A high-throughput and memory-efficient inference and serving engine for LLMs
[Kernel][Helion][1/N] Add Helion kernel for rms_norm_dynamic_per_token_quant (#34432)
Signed-off-by: Sean Chen <seachen@redhat.com> Co-authored-by: Yanan Cao <gmagogsfm@gmail.com>
Xiaohong (Sean) Chen committed
Jun 17, 2026 at 11:03 UTC
46f74e144b33b01684674e3c808babb19835d17b
4 files changed
+6750
tests/kernels/helion/test_rms_norm_dynamic_per_token_quant.py
new
+207
@@ -0,0 +1,207 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for the rms_norm_dynamic_per_token_quant helion kernel
4
+
5
+Run `pytest tests/kernels/helion/test_rms_norm_dynamic_per_token_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 vllm.kernels.helion.case_key import CaseKey
16
+from vllm.kernels.helion.config_manager import ConfigManager
17
+from vllm.kernels.helion.ops.rms_norm_dynamic_per_token_quant import (
18
+ _pick_cache,
19
+ baseline,
20
+ pick_config,
21
+ rms_norm_dynamic_per_token_quant,
22
+)
23
+from vllm.platforms import current_platform
24
+from vllm.utils.import_utils import has_helion
25
+from vllm.utils.torch_utils import set_random_seed
26
+
27
+if not has_helion():
28
+ pytest.skip(
29
+ "Helion is not installed. Install with: pip install vllm[helion]",
30
+ allow_module_level=True,
31
+ )
32
+
33
+
34
+def _generate_fake_input(num_tokens: int, hidden_size: int) -> tuple[Any, ...]:
35
+ with FakeTensorMode():
36
+ input = torch.randn(
37
+ num_tokens, hidden_size, device="cuda", dtype=torch.bfloat16
38
+ )
39
+ result = torch.empty(
40
+ input.shape, device=input.device, dtype=current_platform.fp8_dtype()
41
+ )
42
+ scale = torch.empty((num_tokens, 1), device=input.device, dtype=torch.float32)
43
+ scale_ub = torch.mean(input).to(torch.float32)
44
+ residual = torch.randn_like(input)
45
+ weight = torch.normal(
46
+ mean=1.0,
47
+ std=1.0,
48
+ size=(hidden_size,),
49
+ dtype=input.dtype,
50
+ device=input.device,
51
+ )
52
+ epsilon = 1e-6
53
+ args = (result, input, weight, scale, epsilon, scale_ub, residual)
54
+ return args
55
+
56
+
57
+@pytest.fixture(autouse=True)
58
+def reset_config_manager_singleton():
59
+ ConfigManager.reset_instance()
60
+ ConfigManager()
61
+ yield
62
+ ConfigManager.reset_instance()
63
+
64
+
65
+class TestRmsNormDynamicPerTokenQuantConfigPicker:
66
+ def setup_method(self):
67
+ _pick_cache.clear()
68
+
69
+ def test_config_picker_exact_match(self):
70
+ config_keys = [
71
+ CaseKey({"hidden_size": 2048, "num_tokens": 16}),
72
+ CaseKey({"hidden_size": 4096, "num_tokens": 16}),
73
+ ]
74
+
75
+ args = _generate_fake_input(16, 4096)
76
+ selected_key = pick_config(args, config_keys)
77
+ assert selected_key == CaseKey({"hidden_size": 4096, "num_tokens": 16})
78
+
79
+ def test_config_picker_closest_match(self):
80
+ config_keys = [
81
+ CaseKey({"hidden_size": 2048, "num_tokens": 16}),
82
+ CaseKey({"hidden_size": 2048, "num_tokens": 32}),
83
+ CaseKey({"hidden_size": 4096, "num_tokens": 16}),
84
+ CaseKey({"hidden_size": 4096, "num_tokens": 32}),
85
+ ]
86
+
87
+ args = _generate_fake_input(20, 3000)
88
+ selected_key = pick_config(args, config_keys)
89
+ assert selected_key == CaseKey({"hidden_size": 2048, "num_tokens": 32})
90
+
91
+ def test_config_picker_no_configs(self):
92
+ config_keys: list[dict] = []
93
+
94
+ args = _generate_fake_input(16, 4096)
95
+ selected_key = pick_config(args, config_keys)
96
+ assert selected_key is None
97
+
98
+ def test_config_picker_fallback_to_largest(self):
99
+ config_keys = [
100
+ CaseKey({"hidden_size": 2048, "num_tokens": 16}),
101
+ CaseKey({"hidden_size": 4096, "num_tokens": 16}),
102
+ ]
103
+
104
+ args = _generate_fake_input(32, 8192)
105
+ selected_key = pick_config(args, config_keys)
106
+ assert selected_key == CaseKey({"hidden_size": 4096, "num_tokens": 16})
107
+
108
+
109
+DTYPES = [torch.bfloat16, torch.float]
110
+QUANT_DTYPES = [torch.int8, current_platform.fp8_dtype()]
111
+VEC_HIDDEN_SIZES = [1024, 1025, 1027, 1029]
112
+# Avoid combinatorial explosion with full Cartesian product
113
+NUM_TOKENS_HIDDEN_SIZES = [
114
+ *[(1, i) for i in [1, 64, *VEC_HIDDEN_SIZES, 5120, 5137]],
115
+ *[(2048, i) for i in [1, 64, *VEC_HIDDEN_SIZES, 5137]],
116
+ *[(4096, i) for i in [1, 64, 5137]],
117
+]
118
+
119
+ADD_RESIDUAL = [False, True]
120
+SCALE_UBS = [True, False]
121
+SEEDS = [0]
122
+
123
+EPS = 1e-6
124
+
125
+
126
+class TestRmsNormDynamicPerTokenQuantCorrectness:
127
+ @pytest.mark.parametrize("num_tokens, hidden_size", NUM_TOKENS_HIDDEN_SIZES)
128
+ @pytest.mark.parametrize("add_residual", ADD_RESIDUAL)
129
+ @pytest.mark.parametrize("has_scale_ub", SCALE_UBS)
130
+ @pytest.mark.parametrize("dtype", DTYPES)
131
+ @pytest.mark.parametrize("quant_dtype", QUANT_DTYPES)
132
+ @pytest.mark.parametrize("seed", SEEDS)
133
+ def test_rms_norm_dynamic_per_token_quant(
134
+ self,
135
+ num_tokens: int,
136
+ hidden_size: int,
137
+ add_residual: bool,
138
+ has_scale_ub: bool,
139
+ dtype: torch.dtype,
140
+ quant_dtype: torch.dtype,
141
+ seed: int,
142
+ ) -> None:
143
+ skip_if_platform_unsupported("rms_norm_dynamic_per_token_quant")
144
+
145
+ set_random_seed(seed)
146
+
147
+ if has_scale_ub and quant_dtype != current_platform.fp8_dtype():
148
+ # skip
149
+ return
150
+
151
+ scale = 1 / (hidden_size)
152
+ x = torch.randn(num_tokens, hidden_size, dtype=dtype, device="cuda") * scale
153
+ weight = torch.normal(
154
+ mean=1.0, std=1.0, size=(hidden_size,), dtype=dtype, device=x.device
155
+ )
156
+ residual = torch.randn_like(x) * scale if add_residual else None
157
+ scale_ub = (
158
+ torch.mean(x).to(dtype=torch.float32, device="cuda")
159
+ if has_scale_ub
160
+ else None
161
+ )
162
+
163
+ ref_out = torch.empty(x.shape, device=x.device, dtype=quant_dtype)
164
+ ref_scales = torch.empty((x.shape[0], 1), device=x.device, dtype=torch.float32)
165
+ ref_residual = residual.clone() if residual is not None else None
166
+ baseline(ref_out, x, weight, ref_scales, EPS, scale_ub, ref_residual)
167
+
168
+ ops_out = torch.empty(x.shape, device=x.device, dtype=quant_dtype)
169
+ ops_scales = torch.empty((x.shape[0], 1), device=x.device, dtype=torch.float32)
170
+ ops_residual = residual.clone() if residual is not None else None
171
+ rms_norm_dynamic_per_token_quant(
172
+ ops_out, x, weight, ops_scales, EPS, scale_ub, ops_residual
173
+ )
174
+
175
+ torch.testing.assert_close(ref_scales, ops_scales)
176
+ # allow 1 ULP difference
177
+ assert (
178
+ ref_out.view(torch.uint8).to(torch.int16)
179
+ - ops_out.view(torch.uint8).to(torch.int16)
180
+ ).abs().max() <= 1
181
+
182
+ if add_residual:
183
+ torch.testing.assert_close(ref_residual, ops_residual)
184
+
185
+
186
+class TestRmsNormDynamicPerTokenQuantIntegration:
187
+ def test_kernel_registration_integration(self):
188
+ from vllm.kernels.helion.register import get_registered_kernels
189
+
190
+ registered_kernels = get_registered_kernels()
191
+ assert "rms_norm_dynamic_per_token_quant" in registered_kernels
192
+
193
+ kernel_wrapper = registered_kernels["rms_norm_dynamic_per_token_quant"]
194
+ assert kernel_wrapper.op_name == "rms_norm_dynamic_per_token_quant"
195
+ assert kernel_wrapper._config_picker is not None
196
+ assert kernel_wrapper._mutates_args == ["result", "scale", "residual"]
197
+
198
+ def test_fake_impl_functionality(self):
199
+ skip_if_platform_unsupported("rms_norm_dynamic_per_token_quant")
200
+ from vllm.kernels.helion.register import get_registered_kernels
201
+
202
+ registered_kernels = get_registered_kernels()
203
+ kernel_wrapper = registered_kernels["rms_norm_dynamic_per_token_quant"]
204
+ fake_impl = kernel_wrapper._fake_impl
205
+
206
+ args = _generate_fake_input(16, 4096)
207
+ assert fake_impl(*args) is None
vllm/kernels/helion/configs/rms_norm_dynamic_per_token_quant/nvidia_b200.json
new
+2647
@@ -0,0 +1,2647 @@
1
+[
2
+ {
3
+ "key": {
4
+ "hidden_size": 2048,
5
+ "num_tokens": 1
6
+ },
7
+ "config": {
8
+ "block_sizes": [
9
+ 2048,
10
+ 2048,
11
+ 2048
12
+ ],
13
+ "range_unroll_factors": [
14
+ 0,
15
+ 2,
16
+ 4,
17
+ 3
18
+ ],
19
+ "range_warp_specializes": [],
20
+ "range_num_stages": [],
21
+ "range_multi_buffers": [
22
+ null,
23
+ false,
24
+ false,
25
+ false
26
+ ],
27
+ "range_flattens": [
28
+ null,
29
+ false,
30
+ false,
31
+ true
32
+ ],
33
+ "load_eviction_policies": [
34
+ "last",
35
+ "last",
36
+ "last",
37
+ "last",
38
+ "",
39
+ "first",
40
+ "last",
41
+ "",
42
+ ""
43
+ ],
44
+ "num_warps": 16,
45
+ "num_stages": 3,
46
+ "indexing": [
47
+ "tensor_descriptor",
48
+ "tensor_descriptor",
49
+ "tensor_descriptor",
50
+ "pointer",
51
+ "pointer",
52
+ "pointer",
53
+ "pointer",
54
+ "tensor_descriptor",
55
+ "tensor_descriptor",
56
+ "tensor_descriptor",
57
+ "pointer",
58
+ "pointer"
59
+ ],
60
+ "atomic_indexing": [],
61
+ "pid_type": "flat"
62
+ }
63
+ },
64
+ {
65
+ "key": {
66
+ "hidden_size": 4096,
67
+ "num_tokens": 1
68
+ },
69
+ "config": {
70
+ "block_sizes": [
71
+ 4096,
72
+ 4096,
73
+ 4096
74
+ ],
75
+ "range_unroll_factors": [
76
+ 2,
77
+ 3,
78
+ 1,
79
+ 0
80
+ ],
81
+ "range_warp_specializes": [],
82
+ "range_multi_buffers": [
83
+ true,
84
+ null,
85
+ false,
86
+ null
87
+ ],
88
+ "range_flattens": [
89
+ null,
90
+ null,
91
+ false,
92
+ true
93
+ ],
94
+ "load_eviction_policies": [
95
+ "last",
96
+ "",
97
+ "last",
98
+ "",
99
+ "last",
100
+ "last",
101
+ "",
102
+ "first",
103
+ "last"
104
+ ],
105
+ "num_warps": 16,
106
+ "num_stages": 7,
107
+ "indexing": [
108
+ "pointer",
109
+ "pointer",
110
+ "tensor_descriptor",
111
+ "tensor_descriptor",
112
+ "pointer",
113
+ "pointer",
114
+ "tensor_descriptor",
115
+ "pointer",
116
+ "pointer",
117
+ "tensor_descriptor",
118
+ "pointer",
119
+ "pointer"
120
+ ],
121
+ "atomic_indexing": [],
122
+ "pid_type": "persistent_interleaved",
123
+ "num_sm_multiplier": 1,
124
+ "maxnreg": 128
125
+ }
126
+ },
127
+ {
128
+ "key": {
129
+ "hidden_size": 5120,
130
+ "num_tokens": 1
131
+ },
132
+ "config": {
133
+ "block_sizes": [
134
+ 8192,
135
+ 8192,
136
+ 2048
137
+ ],
138
+ "range_unroll_factors": [
139
+ 0,
140
+ 2,
141
+ 3,
142
+ 2
143
+ ],
144
+ "range_warp_specializes": [],
145
+ "range_num_stages": [],
146
+ "range_multi_buffers": [
147
+ null,
148
+ false,
149
+ false,
150
+ false
151
+ ],
152
+ "range_flattens": [
153
+ null,
154
+ false,
155
+ true,
156
+ false
157
+ ],
158
+ "load_eviction_policies": [
159
+ "last",
160
+ "",
161
+ "last",
162
+ "",
163
+ "",
164
+ "last",
165
+ "last",
166
+ "first",
167
+ "first"
168
+ ],
169
+ "num_warps": 32,
170
+ "num_stages": 6,
171
+ "indexing": [
172
+ "tensor_descriptor",
173
+ "pointer",
174
+ "pointer",
175
+ "pointer",
176
+ "tensor_descriptor",
177
+ "pointer",
178
+ "pointer",
179
+ "tensor_descriptor",
180
+ "pointer",
181
+ "tensor_descriptor",
182
+ "pointer",
183
+ "pointer"
184
+ ],
185
+ "atomic_indexing": [],
186
+ "pid_type": "flat"
187
+ }
188
+ },
189
+ {
190
+ "key": {
191
+ "hidden_size": 2048,
192
+ "num_tokens": 2
193
+ },
194
+ "config": {
195
+ "block_sizes": [
196
+ 4096,
197
+ 4096,
198
+ 2048
199
+ ],
200
+ "range_unroll_factors": [
201
+ 0,
202
+ 3,
203
+ 0,
204
+ 4
205
+ ],
206
+ "range_warp_specializes": [],
207
+ "range_num_stages": [],
208
+ "range_multi_buffers": [
209
+ null,
210
+ false,
211
+ false,
212
+ null
213
+ ],
214
+ "range_flattens": [
215
+ null,
216
+ true,
217
+ null,
218
+ false
219
+ ],
220
+ "load_eviction_policies": [
221
+ "",
222
+ "",
223
+ "",
224
+ "",
225
+ "last",
226
+ "",
227
+ "first",
228
+ "first",
229
+ ""
230
+ ],
231
+ "num_warps": 8,
232
+ "num_stages": 7,
233
+ "indexing": [
234
+ "tensor_descriptor",
235
+ "pointer",
236
+ "tensor_descriptor",
237
+ "tensor_descriptor",
238
+ "pointer",
239
+ "tensor_descriptor",
240
+ "pointer",
241
+ "tensor_descriptor",
242
+ "pointer",
243
+ "tensor_descriptor",
244
+ "tensor_descriptor",
245
+ "pointer"
246
+ ],
247
+ "atomic_indexing": [],
248
+ "pid_type": "flat"
249
+ }
250
+ },
251
+ {
252
+ "key": {
253
+ "hidden_size": 4096,
254
+ "num_tokens": 2
255
+ },
256
+ "config": {
257
+ "block_sizes": [
258
+ 4096,
259
+ 4096,
260
+ 4096
261
+ ],
262
+ "range_unroll_factors": [
263
+ 0,
264
+ 4,
265
+ 2,
266
+ 4
267
+ ],
268
+ "range_warp_specializes": [],
269
+ "range_num_stages": [],
270
+ "range_multi_buffers": [
271
+ null,
272
+ null,
273
+ false,
274
+ null
275
+ ],
276
+ "range_flattens": [
277
+ null,
278
+ false,
279
+ true,
280
+ true
281
+ ],
282
+ "load_eviction_policies": [
283
+ "",
284
+ "first",
285
+ "",
286
+ "first",
287
+ "last",
288
+ "",
289
+ "last",
290
+ "",
291
+ "last"
292
+ ],
293
+ "num_warps": 16,
294
+ "num_stages": 1,
295
+ "indexing": [
296
+ "tensor_descriptor",
297
+ "tensor_descriptor",
298
+ "tensor_descriptor",
299
+ "pointer",
300
+ "tensor_descriptor",
301
+ "tensor_descriptor",
302
+ "tensor_descriptor",
303
+ "pointer",
304
+ "tensor_descriptor",
305
+ "pointer",
306
+ "tensor_descriptor",
307
+ "tensor_descriptor"
308
+ ],
309
+ "atomic_indexing": [],
310
+ "pid_type": "flat"
311
+ }
312
+ },
313
+ {
314
+ "key": {
315
+ "hidden_size": 5120,
316
+ "num_tokens": 2
317
+ },
318
+ "config": {
319
+ "block_sizes": [
320
+ 8192,
321
+ 8192,
322
+ 2048
323
+ ],
324
+ "range_unroll_factors": [
325
+ 0,
326
+ 4,
327
+ 2,
328
+ 3
329
+ ],
330
+ "range_warp_specializes": [],
331
+ "range_num_stages": [],
332
+ "range_multi_buffers": [
333
+ null,
334
+ null,
335
+ true,
336
+ true
337
+ ],
338
+ "range_flattens": [
339
+ null,
340
+ null,
341
+ null,
342
+ true
343
+ ],
344
+ "load_eviction_policies": [
345
+ "",
346
+ "",
347
+ "",
348
+ "",
349
+ "last",
350
+ "last",
351
+ "",
352
+ "first",
353
+ "last"
354
+ ],
355
+ "num_warps": 8,
356
+ "num_stages": 1,
357
+ "indexing": [
358
+ "pointer",
359
+ "tensor_descriptor",
360
+ "tensor_descriptor",
361
+ "pointer",
362
+ "pointer",
363
+ "tensor_descriptor",
364
+ "tensor_descriptor",
365
+ "pointer",
366
+ "tensor_descriptor",
367
+ "pointer",
368
+ "pointer",
369
+ "tensor_descriptor"
370
+ ],
371
+ "atomic_indexing": [],
372
+ "pid_type": "flat"
373
+ }
374
+ },
375
+ {
376
+ "key": {
377
+ "hidden_size": 2048,
378
+ "num_tokens": 4
379
+ },
380
+ "config": {
381
+ "block_sizes": [
382
+ 2048,
383
+ 2048,
384
+ 2048
385
+ ],
386
+ "range_unroll_factors": [
387
+ 0,
388
+ 1,
389
+ 2,
390
+ 0
391
+ ],
392
+ "range_warp_specializes": [
393
+ null,
394
+ null,
395
+ null,
396
+ null
397
+ ],
398
+ "range_num_stages": [],
399
+ "range_multi_buffers": [
400
+ null,
401
+ null,
402
+ false,
403
+ true
404
+ ],
405
+ "range_flattens": [
406
+ null,
407
+ null,
408
+ true,
409
+ null
410
+ ],
411
+ "load_eviction_policies": [
412
+ "",
413
+ "first",
414
+ "first",
415
+ "",
416
+ "last",
417
+ "",
418
+ "",
419
+ "last",
420
+ ""
421
+ ],
422
+ "num_warps": 8,
423
+ "num_stages": 3,
424
+ "indexing": [
425
+ "pointer",
426
+ "tensor_descriptor",
427
+ "tensor_descriptor",
428
+ "tensor_descriptor",
429
+ "tensor_descriptor",
430
+ "pointer",
431
+ "tensor_descriptor",
432
+ "tensor_descriptor",
433
+ "pointer",
434
+ "pointer",
435
+ "pointer",
436
+ "tensor_descriptor"
437
+ ],
438
+ "atomic_indexing": [],
439
+ "pid_type": "flat"
440
+ }
441
+ },
442
+ {
443
+ "key": {
444
+ "hidden_size": 4096,
445
+ "num_tokens": 4
446
+ },
447
+ "config": {
448
+ "block_sizes": [
449
+ 4096,
450
+ 4096,
451
+ 2048
452
+ ],
453
+ "range_unroll_factors": [
454
+ 0,
455
+ 3,
456
+ 0,
457
+ 4
458
+ ],
459
+ "range_warp_specializes": [],
460
+ "range_num_stages": [],
461
+ "range_multi_buffers": [
462
+ null,
463
+ false,
464
+ false,
465
+ null
466
+ ],
467
+ "range_flattens": [
468
+ null,
469
+ true,
470
+ null,
471
+ false
472
+ ],
473
+ "load_eviction_policies": [
474
+ "",
475
+ "",
476
+ "",
477
+ "",
478
+ "last",
479
+ "",
480
+ "first",
481
+ "first",
482
+ ""
483
+ ],
484
+ "num_warps": 8,
485
+ "num_stages": 7,
486
+ "indexing": [
487
+ "tensor_descriptor",
488
+ "pointer",
489
+ "tensor_descriptor",
490
+ "tensor_descriptor",
491
+ "pointer",
492
+ "tensor_descriptor",
493
+ "pointer",
494
+ "tensor_descriptor",
495
+ "pointer",
496
+ "tensor_descriptor",
497
+ "tensor_descriptor",
498
+ "pointer"
499
+ ],
500
+ "atomic_indexing": [],
501
+ "pid_type": "flat"
502
+ }
503
+ },
504
+ {
505
+ "key": {
506
+ "hidden_size": 5120,
507
+ "num_tokens": 4
508
+ },
509
+ "config": {
510
+ "block_sizes": [
511
+ 8192,
512
+ 8192,
513
+ 2048
514
+ ],
515
+ "range_unroll_factors": [
516
+ 0,
517
+ 4,
518
+ 2,
519
+ 3
520
+ ],
521
+ "range_warp_specializes": [],
522
+ "range_num_stages": [],
523
+ "range_multi_buffers": [
524
+ null,
525
+ null,
526
+ true,
527
+ true
528
+ ],
529
+ "range_flattens": [
530
+ null,
531
+ null,
532
+ null,
533
+ true
534
+ ],
535
+ "load_eviction_policies": [
536
+ "",
537
+ "",
538
+ "",
539
+ "",
540
+ "last",
541
+ "last",
542
+ "",
543
+ "first",
544
+ "last"
545
+ ],
546
+ "num_warps": 8,
547
+ "num_stages": 1,
548
+ "indexing": [
549
+ "pointer",
550
+ "tensor_descriptor",
551
+ "tensor_descriptor",
552
+ "pointer",
553
+ "pointer",
554
+ "tensor_descriptor",
555
+ "tensor_descriptor",
556
+ "pointer",
557
+ "tensor_descriptor",
558
+ "pointer",
559
+ "pointer",
560
+ "tensor_descriptor"
561
+ ],
562
+ "atomic_indexing": [],
563
+ "pid_type": "flat"
564
+ }
565
+ },
566
+ {
567
+ "key": {
568
+ "hidden_size": 2048,
569
+ "num_tokens": 8
570
+ },
571
+ "config": {
572
+ "block_sizes": [
573
+ 2048,
574
+ 2048,
575
+ 2048
576
+ ],
577
+ "range_unroll_factors": [
578
+ 0,
579
+ 3,
580
+ 3,
581
+ 1
582
+ ],
583
+ "range_warp_specializes": [],
584
+ "range_num_stages": [],
585
+ "range_multi_buffers": [
586
+ null,
587
+ false,
588
+ true,
589
+ null
590
+ ],
591
+ "range_flattens": [
592
+ null,
593
+ false,
594
+ false,
595
+ true
596
+ ],
597
+ "load_eviction_policies": [
598
+ "last",
599
+ "",
600
+ "last",
601
+ "",
602
+ "last",
603
+ "first",
604
+ "first",
605
+ "",
606
+ ""
607
+ ],
608
+ "num_warps": 8,
609
+ "num_stages": 4,
610
+ "indexing": [
611
+ "tensor_descriptor",
612
+ "pointer",
613
+ "tensor_descriptor",
614
+ "tensor_descriptor",
615
+ "pointer",
616
+ "tensor_descriptor",
617
+ "tensor_descriptor",
618
+ "pointer",
619
+ "tensor_descriptor",
620
+ "tensor_descriptor",
621
+ "tensor_descriptor",
622
+ "tensor_descriptor"
623
+ ],
624
+ "atomic_indexing": [],
625
+ "pid_type": "flat"
626
+ }
627
+ },
628
+ {
629
+ "key": {
630
+ "hidden_size": 4096,
631
+ "num_tokens": 8
632
+ },
633
+ "config": {
634
+ "block_sizes": [
635
+ 4096,
636
+ 4096,
637
+ 4096
638
+ ],
639
+ "range_unroll_factors": [
640
+ 0,
641
+ 1,
642
+ 2,
643
+ 1
644
+ ],
645
+ "range_warp_specializes": [],
646
+ "range_num_stages": [],
647
+ "range_multi_buffers": [
648
+ null,
649
+ null,
650
+ null,
651
+ null
652
+ ],
653
+ "range_flattens": [
654
+ null,
655
+ null,
656
+ null,
657
+ null
658
+ ],
659
+ "load_eviction_policies": [
660
+ "first",
661
+ "",
662
+ "first",
663
+ "",
664
+ "last",
665
+ "first",
666
+ "",
667
+ "",
668
+ "first"
669
+ ],
670
+ "num_warps": 16,
671
+ "num_stages": 7,
672
+ "indexing": [
673
+ "pointer",
674
+ "pointer",
675
+ "pointer",
676
+ "pointer",
677
+ "pointer",
678
+ "pointer",
679
+ "tensor_descriptor",
680
+ "pointer",
681
+ "tensor_descriptor",
682
+ "pointer",
683
+ "tensor_descriptor",
684
+ "tensor_descriptor"
685
+ ],
686
+ "atomic_indexing": [],
687
+ "pid_type": "flat"
688
+ }
689
+ },
690
+ {
691
+ "key": {
692
+ "hidden_size": 5120,
693
+ "num_tokens": 8
694
+ },
695
+ "config": {
696
+ "block_sizes": [
697
+ 8192,
698
+ 8192,
699
+ 2048
700
+ ],
701
+ "range_unroll_factors": [
702
+ 0,
703
+ 4,
704
+ 2,
705
+ 3
706
+ ],
707
+ "range_warp_specializes": [],
708
+ "range_num_stages": [],
709
+ "range_multi_buffers": [
710
+ null,
711
+ null,
712
+ true,
713
+ true
714
+ ],
715
+ "range_flattens": [
716
+ null,
717
+ null,
718
+ null,
719
+ true
720
+ ],
721
+ "load_eviction_policies": [
722
+ "",
723
+ "",
724
+ "",
725
+ "",
726
+ "last",
727
+ "last",
728
+ "",
729
+ "first",
730
+ "last"
731
+ ],
732
+ "num_warps": 8,
733
+ "num_stages": 1,
734
+ "indexing": [
735
+ "pointer",
736
+ "tensor_descriptor",
737
+ "tensor_descriptor",
738
+ "pointer",
739
+ "pointer",
740
+ "tensor_descriptor",
741
+ "tensor_descriptor",
742
+ "pointer",
743
+ "tensor_descriptor",
744
+ "pointer",
745
+ "pointer",
746
+ "tensor_descriptor"
747
+ ],
748
+ "atomic_indexing": [],
749
+ "pid_type": "flat"
750
+ }
751
+ },
752
+ {
753
+ "key": {
754
+ "hidden_size": 2048,
755
+ "num_tokens": 16
756
+ },
757
+ "config": {
758
+ "block_sizes": [
759
+ 2048,
760
+ 2048,
761
+ 2048
762
+ ],
763
+ "range_unroll_factors": [
764
+ 0,
765
+ 4,
766
+ 3,
767
+ 4
768
+ ],
769
+ "range_warp_specializes": [],
770
+ "range_num_stages": [],
771
+ "range_multi_buffers": [
772
+ null,
773
+ true,
774
+ true,
775
+ true
776
+ ],
777
+ "range_flattens": [
778
+ null,
779
+ false,
780
+ false,
781
+ null
782
+ ],
783
+ "load_eviction_policies": [
784
+ "",
785
+ "",
786
+ "",
787
+ "",
788
+ "last",
789
+ "",
790
+ "",
791
+ "",
792
+ "first"
793
+ ],
794
+ "num_warps": 8,
795
+ "num_stages": 1,
796
+ "indexing": [
797
+ "tensor_descriptor",
798
+ "tensor_descriptor",
799
+ "tensor_descriptor",
800
+ "tensor_descriptor",
801
+ "tensor_descriptor",
802
+ "pointer",
803
+ "tensor_descriptor",
804
+ "tensor_descriptor",
805
+ "pointer",
806
+ "tensor_descriptor",
807
+ "tensor_descriptor",
808
+ "pointer"
809
+ ],
810
+ "atomic_indexing": [],
811
+ "pid_type": "flat"
812
+ }
813
+ },
814
+ {
815
+ "key": {
816
+ "hidden_size": 4096,
817
+ "num_tokens": 16
818
+ },
819
+ "config": {
820
+ "block_sizes": [
821
+ 4096,
822
+ 4096,
823
+ 4096
824
+ ],
825
+ "range_unroll_factors": [
826
+ 0,
827
+ 2,
828
+ 0,
829
+ 3
830
+ ],
831
+ "range_warp_specializes": [],
832
+ "range_num_stages": [],
833
+ "range_multi_buffers": [
834
+ null,
835
+ false,
836
+ null,
837
+ null
838
+ ],
839
+ "range_flattens": [
840
+ null,
841
+ true,
842
+ true,
843
+ false
844
+ ],
845
+ "load_eviction_policies": [
846
+ "last",
847
+ "last",
848
+ "last",
849
+ "last",
850
+ "last",
851
+ "last",
852
+ "",
853
+ "",
854
+ ""
855
+ ],
856
+ "num_warps": 16,
857
+ "num_stages": 2,
858
+ "indexing": [
859
+ "pointer",
860
+ "pointer",
861
+ "pointer",
862
+ "tensor_descriptor",
863
+ "pointer",
864
+ "pointer",
865
+ "pointer",
866
+ "pointer",
867
+ "tensor_descriptor",
868
+ "tensor_descriptor",
869
+ "tensor_descriptor",
870
+ "pointer"
871
+ ],
872
+ "atomic_indexing": [],
873
+ "pid_type": "flat"
874
+ }
875
+ },
876
+ {
877
+ "key": {
878
+ "hidden_size": 5120,
879
+ "num_tokens": 16
880
+ },
881
+ "config": {
882
+ "block_sizes": [
883
+ 8192,
884
+ 8192,
885
+ 2048
886
+ ],
887
+ "range_unroll_factors": [
888
+ 0,
889
+ 4,
890
+ 4,
891
+ 1
892
+ ],
893
+ "range_warp_specializes": [
894
+ null,
895
+ false,
896
+ null,
897
+ false
898
+ ],
899
+ "range_num_stages": [],
900
+ "range_multi_buffers": [
901
+ null,
902
+ false,
903
+ false,
904
+ false
905
+ ],
906
+ "range_flattens": [
907
+ null,
908
+ true,
909
+ false,
910
+ true
911
+ ],
912
+ "load_eviction_policies": [
913
+ "last",
914
+ "",
915
+ "last",
916
+ "first",
917
+ "first",
918
+ "",
919
+ "last",
920
+ "",
921
+ "last"
922
+ ],
923
+ "num_warps": 8,
924
+ "num_stages": 7,
925
+ "indexing": [
926
+ "pointer",
927
+ "pointer",
928
+ "tensor_descriptor",
929
+ "pointer",
930
+ "pointer",
931
+ "pointer",
932
+ "pointer",
933
+ "tensor_descriptor",
934
+ "pointer",
935
+ "pointer",
936
+ "pointer",
937
+ "pointer"
938
+ ],
939
+ "atomic_indexing": [],
940
+ "pid_type": "flat"
941
+ }
942
+ },
943
+ {
944
+ "key": {
945
+ "hidden_size": 2048,
946
+ "num_tokens": 32
947
+ },
948
+ "config": {
949
+ "block_sizes": [
950
+ 2048,
951
+ 2048,
952
+ 2048
953
+ ],
954
+ "range_unroll_factors": [
955
+ 0,
956
+ 4,
957
+ 3,
958
+ 4
959
+ ],
960
+ "range_warp_specializes": [],
961
+ "range_num_stages": [],
962
+ "range_multi_buffers": [
963
+ null,
964
+ true,
965
+ true,
966
+ true
967
+ ],
968
+ "range_flattens": [
969
+ null,
970
+ false,
971
+ false,
972
+ null
973
+ ],
974
+ "load_eviction_policies": [
975
+ "",
976
+ "",
977
+ "",
978
+ "",
979
+ "last",
980
+ "",
981
+ "",
982
+ "",
983
+ "first"
984
+ ],
985
+ "num_warps": 8,
986
+ "num_stages": 1,
987
+ "indexing": [
988
+ "tensor_descriptor",
989
+ "tensor_descriptor",
990
+ "tensor_descriptor",
991
+ "tensor_descriptor",
992
+ "tensor_descriptor",
993
+ "pointer",
994
+ "tensor_descriptor",
995
+ "tensor_descriptor",
996
+ "pointer",
997
+ "tensor_descriptor",
998
+ "tensor_descriptor",
999
+ "pointer"
1000
+ ],
1001
+ "atomic_indexing": [],
1002
+ "pid_type": "flat"
1003
+ }
1004
+ },
1005
+ {
1006
+ "key": {
1007
+ "hidden_size": 4096,
1008
+ "num_tokens": 32
1009
+ },
1010
+ "config": {
1011
+ "block_sizes": [
1012
+ 4096,
1013
+ 4096,
1014
+ 4096
1015
+ ],
1016
+ "range_unroll_factors": [
1017
+ 0,
1018
+ 1,
1019
+ 2,
1020
+ 1
1021
+ ],
1022
+ "range_warp_specializes": [],
1023
+ "range_num_stages": [],
1024
+ "range_multi_buffers": [
1025
+ null,
1026
+ null,
1027
+ null,
1028
+ null
1029
+ ],
1030
+ "range_flattens": [
1031
+ null,
1032
+ null,
1033
+ null,
1034
+ null
1035
+ ],
1036
+ "load_eviction_policies": [
1037
+ "first",
1038
+ "",
1039
+ "first",
1040
+ "",
1041
+ "last",
1042
+ "first",
1043
+ "",
1044
+ "",
1045
+ "first"
1046
+ ],
1047
+ "num_warps": 16,
1048
+ "num_stages": 7,
1049
+ "indexing": [
1050
+ "pointer",
1051
+ "pointer",
1052
+ "pointer",
1053
+ "pointer",
1054
+ "pointer",
1055
+ "pointer",
1056
+ "tensor_descriptor",
1057
+ "pointer",
1058
+ "tensor_descriptor",
1059
+ "pointer",
1060
+ "tensor_descriptor",
1061
+ "tensor_descriptor"
1062
+ ],
1063
+ "atomic_indexing": [],
1064
+ "pid_type": "flat"
1065
+ }
1066
+ },
1067
+ {
1068
+ "key": {
1069
+ "hidden_size": 5120,
1070
+ "num_tokens": 32
1071
+ },
1072
+ "config": {
1073
+ "block_sizes": [
1074
+ 8192,
1075
+ 8192,
1076
+ 8192
1077
+ ],
1078
+ "range_unroll_factors": [
1079
+ 0,
1080
+ 2,
1081
+ 4,
1082
+ 3
1083
+ ],
1084
+ "range_warp_specializes": [],
1085
+ "range_num_stages": [],
1086
+ "range_multi_buffers": [
1087
+ null,
1088
+ false,
1089
+ null,
1090
+ null
1091
+ ],
1092
+ "range_flattens": [
1093
+ null,
1094
+ null,
1095
+ true,
1096
+ null
1097
+ ],
1098
+ "load_eviction_policies": [
1099
+ "last",
1100
+ "last",
1101
+ "last",
1102
+ "last",
1103
+ "first",
1104
+ "first",
1105
+ "",
1106
+ "",
1107
+ "last"
1108
+ ],
1109
+ "num_warps": 32,
1110
+ "num_stages": 2,
1111
+ "indexing": [
1112
+ "tensor_descriptor",
1113
+ "tensor_descriptor",
1114
+ "tensor_descriptor",
1115
+ "tensor_descriptor",
1116
+ "tensor_descriptor",
1117
+ "tensor_descriptor",
1118
+ "pointer",
1119
+ "pointer",
1120
+ "pointer",
1121
+ "tensor_descriptor",
1122
+ "pointer",
1123
+ "pointer"
1124
+ ],
1125
+ "atomic_indexing": [],
1126
+ "pid_type": "flat"
1127
+ }
1128
+ },
1129
+ {
1130
+ "key": {
1131
+ "hidden_size": 2048,
1132
+ "num_tokens": 64
1133
+ },
1134
+ "config": {
1135
+ "block_sizes": [
1136
+ 2048,
1137
+ 2048,
1138
+ 2048
1139
+ ],
1140
+ "range_unroll_factors": [
1141
+ 0,
1142
+ 4,
1143
+ 3,
1144
+ 4
1145
+ ],
1146
+ "range_warp_specializes": [],
1147
+ "range_num_stages": [],
1148
+ "range_multi_buffers": [
1149
+ null,
1150
+ true,
1151
+ true,
1152
+ true
1153
+ ],
1154
+ "range_flattens": [
1155
+ null,
1156
+ false,
1157
+ false,
1158
+ null
1159
+ ],
1160
+ "load_eviction_policies": [
1161
+ "",
1162
+ "",
1163
+ "",
1164
+ "",
1165
+ "last",
1166
+ "",
1167
+ "",
1168
+ "",
1169
+ "first"
1170
+ ],
1171
+ "num_warps": 8,
1172
+ "num_stages": 1,
1173
+ "indexing": [
1174
+ "tensor_descriptor",
1175
+ "tensor_descriptor",
1176
+ "tensor_descriptor",
1177
+ "tensor_descriptor",
1178
+ "tensor_descriptor",
1179
+ "pointer",
1180
+ "tensor_descriptor",
1181
+ "tensor_descriptor",
1182
+ "pointer",
1183
+ "tensor_descriptor",
1184
+ "tensor_descriptor",
1185
+ "pointer"
1186
+ ],
1187
+ "atomic_indexing": [],
1188
+ "pid_type": "flat"
1189
+ }
1190
+ },
1191
+ {
1192
+ "key": {
1193
+ "hidden_size": 4096,
1194
+ "num_tokens": 64
1195
+ },
1196
+ "config": {
1197
+ "block_sizes": [
1198
+ 4096,
1199
+ 4096,
1200
+ 4096
1201
+ ],
1202
+ "range_unroll_factors": [
1203
+ 0,
1204
+ 2,
1205
+ 2,
1206
+ 1
1207
+ ],
1208
+ "range_warp_specializes": [],
1209
+ "range_num_stages": [],
1210
+ "range_multi_buffers": [
1211
+ null,
1212
+ null,
1213
+ true,
1214
+ null
1215
+ ],
1216
+ "range_flattens": [
1217
+ null,
1218
+ false,
1219
+ null,
1220
+ false
1221
+ ],
1222
+ "load_eviction_policies": [
1223
+ "last",
1224
+ "",
1225
+ "last",
1226
+ "",
1227
+ "first",
1228
+ "first",
1229
+ "",
1230
+ "first",
1231
+ ""
1232
+ ],
1233
+ "num_warps": 16,
1234
+ "num_stages": 2,
1235
+ "indexing": [
1236
+ "tensor_descriptor",
1237
+ "tensor_descriptor",
1238
+ "pointer",
1239
+ "tensor_descriptor",
1240
+ "pointer",
1241
+ "tensor_descriptor",
1242
+ "pointer",
1243
+ "tensor_descriptor",
1244
+ "tensor_descriptor",
1245
+ "pointer",
1246
+ "pointer",
1247
+ "tensor_descriptor"
1248
+ ],
1249
+ "atomic_indexing": [],
1250
+ "pid_type": "flat"
1251
+ }
1252
+ },
1253
+ {
1254
+ "key": {
1255
+ "hidden_size": 5120,
1256
+ "num_tokens": 64
1257
+ },
1258
+ "config": {
1259
+ "block_sizes": [
1260
+ 8192,
1261
+ 8192,
1262
+ 2048
1263
+ ],
1264
+ "range_unroll_factors": [
1265
+ 0,
1266
+ 4,
1267
+ 3,
1268
+ 2
1269
+ ],
1270
+ "range_warp_specializes": [],
1271
+ "range_num_stages": [],
1272
+ "range_multi_buffers": [
1273
+ null,
1274
+ false,
1275
+ null,
1276
+ false
1277
+ ],
1278
+ "range_flattens": [
1279
+ null,
1280
+ true,
1281
+ true,
1282
+ null
1283
+ ],
1284
+ "load_eviction_policies": [
1285
+ "",
1286
+ "",
1287
+ "",
1288
+ "",
1289
+ "last",
1290
+ "",
1291
+ "first",
1292
+ "first",
1293
+ ""
1294
+ ],
1295
+ "num_warps": 8,
1296
+ "num_stages": 1,
1297
+ "indexing": [
1298
+ "pointer",
1299
+ "pointer",
1300
+ "pointer",
1301
+ "pointer",
1302
+ "tensor_descriptor",
1303
+ "pointer",
1304
+ "tensor_descriptor",
1305
+ "pointer",
1306
+ "tensor_descriptor",
1307
+ "pointer",
1308
+ "pointer",
1309
+ "pointer"
1310
+ ],
1311
+ "atomic_indexing": [],
1312
+ "pid_type": "flat"
1313
+ }
1314
+ },
1315
+ {
1316
+ "key": {
1317
+ "hidden_size": 2048,
1318
+ "num_tokens": 128
1319
+ },
1320
+ "config": {
1321
+ "block_sizes": [
1322
+ 2048,
1323
+ 2048,
1324
+ 2048
1325
+ ],
1326
+ "range_unroll_factors": [
1327
+ 0,
1328
+ 2,
1329
+ 4,
1330
+ 3
1331
+ ],
1332
+ "range_warp_specializes": [],
1333
+ "range_num_stages": [],
1334
+ "range_multi_buffers": [
1335
+ null,
1336
+ false,
1337
+ false,
1338
+ false
1339
+ ],
1340
+ "range_flattens": [
1341
+ null,
1342
+ false,
1343
+ false,
1344
+ true
1345
+ ],
1346
+ "load_eviction_policies": [
1347
+ "last",
1348
+ "last",
1349
+ "last",
1350
+ "last",
1351
+ "",
1352
+ "first",
1353
+ "last",
1354
+ "",
1355
+ ""
1356
+ ],
1357
+ "num_warps": 16,
1358
+ "num_stages": 3,
1359
+ "indexing": [
1360
+ "tensor_descriptor",
1361
+ "tensor_descriptor",
1362
+ "tensor_descriptor",
1363
+ "pointer",
1364
+ "pointer",
1365
+ "pointer",
1366
+ "pointer",
1367
+ "tensor_descriptor",
1368
+ "tensor_descriptor",
1369
+ "tensor_descriptor",
1370
+ "pointer",
1371
+ "pointer"
1372
+ ],
1373
+ "atomic_indexing": [],
1374
+ "pid_type": "flat"
1375
+ }
1376
+ },
1377
+ {
1378
+ "key": {
1379
+ "hidden_size": 4096,
1380
+ "num_tokens": 128
1381
+ },
1382
+ "config": {
1383
+ "block_sizes": [
1384
+ 4096,
1385
+ 4096,
1386
+ 4096
1387
+ ],
1388
+ "range_unroll_factors": [
1389
+ 0,
1390
+ 2,
1391
+ 0,
1392
+ 3
1393
+ ],
1394
+ "range_warp_specializes": [],
1395
+ "range_num_stages": [],
1396
+ "range_multi_buffers": [
1397
+ null,
1398
+ false,
1399
+ null,
1400
+ null
1401
+ ],
1402
+ "range_flattens": [
1403
+ null,
1404
+ true,
1405
+ true,
1406
+ false
1407
+ ],
1408
+ "load_eviction_policies": [
1409
+ "last",
1410
+ "last",
1411
+ "last",
1412
+ "last",
1413
+ "last",
1414
+ "last",
1415
+ "",
1416
+ "",
1417
+ ""
1418
+ ],
1419
+ "num_warps": 16,
1420
+ "num_stages": 2,
1421
+ "indexing": [
1422
+ "pointer",
1423
+ "pointer",
1424
+ "pointer",
1425
+ "tensor_descriptor",
1426
+ "pointer",
1427
+ "pointer",
1428
+ "pointer",
1429
+ "pointer",
1430
+ "tensor_descriptor",
1431
+ "tensor_descriptor",
1432
+ "tensor_descriptor",
1433
+ "pointer"
1434
+ ],
1435
+ "atomic_indexing": [],
1436
+ "pid_type": "flat"
1437
+ }
1438
+ },
1439
+ {
1440
+ "key": {
1441
+ "hidden_size": 5120,
1442
+ "num_tokens": 128
1443
+ },
1444
+ "config": {
1445
+ "block_sizes": [
1446
+ 8192,
1447
+ 8192,
1448
+ 2048
1449
+ ],
1450
+ "range_unroll_factors": [
1451
+ 0,
1452
+ 4,
1453
+ 4,
1454
+ 1
1455
+ ],
1456
+ "range_warp_specializes": [
1457
+ null,
1458
+ false,
1459
+ null,
1460
+ false
1461
+ ],
1462
+ "range_num_stages": [],
1463
+ "range_multi_buffers": [
1464
+ null,
1465
+ false,
1466
+ false,
1467
+ false
1468
+ ],
1469
+ "range_flattens": [
1470
+ null,
1471
+ true,
1472
+ false,
1473
+ true
1474
+ ],
1475
+ "load_eviction_policies": [
1476
+ "last",
1477
+ "",
1478
+ "last",
1479
+ "first",
1480
+ "first",
1481
+ "",
1482
+ "last",
1483
+ "",
1484
+ "last"
1485
+ ],
1486
+ "num_warps": 8,
1487
+ "num_stages": 7,
1488
+ "indexing": [
1489
+ "pointer",
1490
+ "pointer",
1491
+ "tensor_descriptor",
1492
+ "pointer",
1493
+ "pointer",
1494
+ "pointer",
1495
+ "pointer",
1496
+ "tensor_descriptor",
1497
+ "pointer",
1498
+ "pointer",
1499
+ "pointer",
1500
+ "pointer"
1501
+ ],
1502
+ "atomic_indexing": [],
1503
+ "pid_type": "flat"
1504
+ }
1505
+ },
1506
+ {
1507
+ "key": {
1508
+ "hidden_size": 2048,
1509
+ "num_tokens": 256
1510
+ },
1511
+ "config": {
1512
+ "block_sizes": [
1513
+ 2048,
1514
+ 2048,
1515
+ 2048
1516
+ ],
1517
+ "range_unroll_factors": [
1518
+ 0,
1519
+ 4,
1520
+ 3,
1521
+ 4
1522
+ ],
1523
+ "range_warp_specializes": [],
1524
+ "range_num_stages": [],
1525
+ "range_multi_buffers": [
1526
+ null,
1527
+ true,
1528
+ true,
1529
+ true
1530
+ ],
1531
+ "range_flattens": [
1532
+ null,
1533
+ false,
1534
+ false,
1535
+ null
1536
+ ],
1537
+ "load_eviction_policies": [
1538
+ "",
1539
+ "",
1540
+ "",
1541
+ "",
1542
+ "last",
1543
+ "",
1544
+ "",
1545
+ "",
1546
+ "first"
1547
+ ],
1548
+ "num_warps": 8,
1549
+ "num_stages": 1,
1550
+ "indexing": [
1551
+ "tensor_descriptor",
1552
+ "tensor_descriptor",
1553
+ "tensor_descriptor",
1554
+ "tensor_descriptor",
1555
+ "tensor_descriptor",
1556
+ "pointer",
1557
+ "tensor_descriptor",
1558
+ "tensor_descriptor",
1559
+ "pointer",
1560
+ "tensor_descriptor",
1561
+ "tensor_descriptor",
1562
+ "pointer"
1563
+ ],
1564
+ "atomic_indexing": [],
1565
+ "pid_type": "flat"
1566
+ }
1567
+ },
1568
+ {
1569
+ "key": {
1570
+ "hidden_size": 4096,
1571
+ "num_tokens": 256
1572
+ },
1573
+ "config": {
1574
+ "block_sizes": [
1575
+ 4096,
1576
+ 4096,
1577
+ 2048
1578
+ ],
1579
+ "range_unroll_factors": [
1580
+ 0,
1581
+ 3,
1582
+ 0,
1583
+ 4
1584
+ ],
1585
+ "range_warp_specializes": [],
1586
+ "range_num_stages": [],
1587
+ "range_multi_buffers": [
1588
+ null,
1589
+ false,
1590
+ false,
1591
+ null
1592
+ ],
1593
+ "range_flattens": [
1594
+ null,
1595
+ true,
1596
+ null,
1597
+ false
1598
+ ],
1599
+ "load_eviction_policies": [
1600
+ "",
1601
+ "",
1602
+ "",
1603
+ "",
1604
+ "last",
1605
+ "",
1606
+ "first",
1607
+ "first",
1608
+ ""
1609
+ ],
1610
+ "num_warps": 8,
1611
+ "num_stages": 7,
1612
+ "indexing": [
1613
+ "tensor_descriptor",
1614
+ "pointer",
1615
+ "tensor_descriptor",
1616
+ "tensor_descriptor",
1617
+ "pointer",
1618
+ "tensor_descriptor",
1619
+ "pointer",
1620
+ "tensor_descriptor",
1621
+ "pointer",
1622
+ "tensor_descriptor",
1623
+ "tensor_descriptor",
1624
+ "pointer"
1625
+ ],
1626
+ "atomic_indexing": [],
1627
+ "pid_type": "flat"
1628
+ }
1629
+ },
1630
+ {
1631
+ "key": {
1632
+ "hidden_size": 5120,
1633
+ "num_tokens": 256
1634
+ },
1635
+ "config": {
1636
+ "block_sizes": [
1637
+ 8192,
1638
+ 8192,
1639
+ 2048
1640
+ ],
1641
+ "range_unroll_factors": [
1642
+ 0,
1643
+ 4,
1644
+ 2,
1645
+ 3
1646
+ ],
1647
+ "range_warp_specializes": [],
1648
+ "range_num_stages": [],
1649
+ "range_multi_buffers": [
1650
+ null,
1651
+ null,
1652
+ true,
1653
+ true
1654
+ ],
1655
+ "range_flattens": [
1656
+ null,
1657
+ null,
1658
+ null,
1659
+ true
1660
+ ],
1661
+ "load_eviction_policies": [
1662
+ "",
1663
+ "",
1664
+ "",
1665
+ "",
1666
+ "last",
1667
+ "last",
1668
+ "",
1669
+ "first",
1670
+ "last"
1671
+ ],
1672
+ "num_warps": 8,
1673
+ "num_stages": 1,
1674
+ "indexing": [
1675
+ "pointer",
1676
+ "tensor_descriptor",
1677
+ "tensor_descriptor",
1678
+ "pointer",
1679
+ "pointer",
1680
+ "tensor_descriptor",
1681
+ "tensor_descriptor",
1682
+ "pointer",
1683
+ "tensor_descriptor",
1684
+ "pointer",
1685
+ "pointer",
1686
+ "tensor_descriptor"
1687
+ ],
1688
+ "atomic_indexing": [],
1689
+ "pid_type": "flat"
1690
+ }
1691
+ },
1692
+ {
1693
+ "key": {
1694
+ "hidden_size": 2048,
1695
+ "num_tokens": 512
1696
+ },
1697
+ "config": {
1698
+ "block_sizes": [
1699
+ 2048,
1700
+ 2048,
1701
+ 2048
1702
+ ],
1703
+ "range_unroll_factors": [
1704
+ 0,
1705
+ 0,
1706
+ 0,
1707
+ 1
1708
+ ],
1709
+ "range_warp_specializes": [],
1710
+ "range_num_stages": [],
1711
+ "range_multi_buffers": [
1712
+ null,
1713
+ false,
1714
+ null,
1715
+ false
1716
+ ],
1717
+ "range_flattens": [
1718
+ null,
1719
+ null,
1720
+ null,
1721
+ true
1722
+ ],
1723
+ "load_eviction_policies": [
1724
+ "last",
1725
+ "",
1726
+ "last",
1727
+ "",
1728
+ "first",
1729
+ "",
1730
+ "first",
1731
+ "first",
1732
+ ""
1733
+ ],
1734
+ "num_warps": 8,
1735
+ "num_stages": 8,
1736
+ "indexing": [
1737
+ "tensor_descriptor",
1738
+ "tensor_descriptor",
1739
+ "tensor_descriptor",
1740
+ "tensor_descriptor",
1741
+ "pointer",
1742
+ "pointer",
1743
+ "pointer",
1744
+ "tensor_descriptor",
1745
+ "tensor_descriptor",
1746
+ "tensor_descriptor",
1747
+ "tensor_descriptor",
1748
+ "tensor_descriptor"
1749
+ ],
1750
+ "atomic_indexing": [],
1751
+ "pid_type": "flat"
1752
+ }
1753
+ },
1754
+ {
1755
+ "key": {
1756
+ "hidden_size": 4096,
1757
+ "num_tokens": 512
1758
+ },
1759
+ "config": {
1760
+ "block_sizes": [
1761
+ 4096,
1762
+ 4096,
1763
+ 2048
1764
+ ],
1765
+ "range_unroll_factors": [
1766
+ 0,
1767
+ 3,
1768
+ 2,
1769
+ 2
1770
+ ],
1771
+ "range_warp_specializes": [],
1772
+ "range_num_stages": [],
1773
+ "range_multi_buffers": [
1774
+ null,
1775
+ true,
1776
+ null,
1777
+ true
1778
+ ],
1779
+ "range_flattens": [
1780
+ null,
1781
+ false,
1782
+ null,
1783
+ null
1784
+ ],
1785
+ "load_eviction_policies": [
1786
+ "",
1787
+ "",
1788
+ "",
1789
+ "",
1790
+ "",
1791
+ "last",
1792
+ "first",
1793
+ "",
1794
+ ""
1795
+ ],
1796
+ "num_warps": 8,
1797
+ "num_stages": 1,
1798
+ "indexing": [
1799
+ "pointer",
1800
+ "pointer",
1801
+ "tensor_descriptor",
1802
+ "tensor_descriptor",
1803
+ "tensor_descriptor",
1804
+ "pointer",
1805
+ "tensor_descriptor",
1806
+ "pointer",
1807
+ "pointer",
1808
+ "tensor_descriptor",
1809
+ "pointer",
1810
+ "tensor_descriptor"
1811
+ ],
1812
+ "atomic_indexing": [],
1813
+ "pid_type": "flat"
1814
+ }
1815
+ },
1816
+ {
1817
+ "key": {
1818
+ "hidden_size": 5120,
1819
+ "num_tokens": 512
1820
+ },
1821
+ "config": {
1822
+ "block_sizes": [
1823
+ 8192,
1824
+ 8192,
1825
+ 1024
1826
+ ],
1827
+ "range_unroll_factors": [
1828
+ 0,
1829
+ 3,
1830
+ 2,
1831
+ 4
1832
+ ],
1833
+ "range_warp_specializes": [],
1834
+ "range_num_stages": [],
1835
+ "range_multi_buffers": [
1836
+ null,
1837
+ true,
1838
+ null,
1839
+ null
1840
+ ],
1841
+ "range_flattens": [
1842
+ null,
1843
+ null,
1844
+ true,
1845
+ null
1846
+ ],
1847
+ "load_eviction_policies": [
1848
+ "first",
1849
+ "",
1850
+ "last",
1851
+ "last",
1852
+ "",
1853
+ "last",
1854
+ "first",
1855
+ "first",
1856
+ "last"
1857
+ ],
1858
+ "num_warps": 4,
1859
+ "num_stages": 1,
1860
+ "indexing": [
1861
+ "pointer",
1862
+ "tensor_descriptor",
1863
+ "pointer",
1864
+ "tensor_descriptor",
1865
+ "pointer",
1866
+ "tensor_descriptor",
1867
+ "pointer",
1868
+ "pointer",
1869
+ "tensor_descriptor",
1870
+ "tensor_descriptor",
1871
+ "tensor_descriptor",
1872
+ "tensor_descriptor"
1873
+ ],
1874
+ "atomic_indexing": [],
1875
+ "pid_type": "flat"
1876
+ }
1877
+ },
1878
+ {
1879
+ "key": {
1880
+ "hidden_size": 2048,
1881
+ "num_tokens": 1024
1882
+ },
1883
+ "config": {
1884
+ "block_sizes": [
1885
+ 2048,
1886
+ 2048,
1887
+ 1024
1888
+ ],
1889
+ "range_unroll_factors": [
1890
+ 0,
1891
+ 3,
1892
+ 0,
1893
+ 0
1894
+ ],
1895
+ "range_warp_specializes": [
1896
+ null,
1897
+ null,
1898
+ null,
1899
+ true
1900
+ ],
1901
+ "range_num_stages": [],
1902
+ "range_multi_buffers": [
1903
+ null,
1904
+ false,
1905
+ null,
1906
+ true
1907
+ ],
1908
+ "range_flattens": [
1909
+ null,
1910
+ true,
1911
+ true,
1912
+ null
1913
+ ],
1914
+ "load_eviction_policies": [
1915
+ "first",
1916
+ "last",
1917
+ "first",
1918
+ "last",
1919
+ "",
1920
+ "last",
1921
+ "",
1922
+ "",
1923
+ "last"
1924
+ ],
1925
+ "num_warps": 1,
1926
+ "num_stages": 8,
1927
+ "indexing": [
1928
+ "tensor_descriptor",
1929
+ "pointer",
1930
+ "pointer",
1931
+ "tensor_descriptor",
1932
+ "pointer",
1933
+ "pointer",
1934
+ "pointer",
1935
+ "pointer",
1936
+ "tensor_descriptor",
1937
+ "tensor_descriptor",
1938
+ "pointer",
1939
+ "pointer"
1940
+ ],
1941
+ "atomic_indexing": [],
1942
+ "pid_type": "flat"
1943
+ }
1944
+ },
1945
+ {
1946
+ "key": {
1947
+ "hidden_size": 4096,
1948
+ "num_tokens": 1024
1949
+ },
1950
+ "config": {
1951
+ "block_sizes": [
1952
+ 4096,
1953
+ 4096,
1954
+ 2048
1955
+ ],
1956
+ "range_unroll_factors": [
1957
+ 0,
1958
+ 4,
1959
+ 3,
1960
+ 2
1961
+ ],
1962
+ "range_warp_specializes": [
1963
+ null,
1964
+ null,
1965
+ false,
1966
+ null
1967
+ ],
1968
+ "range_num_stages": [],
1969
+ "range_multi_buffers": [
1970
+ null,
1971
+ false,
1972
+ null,
1973
+ null
1974
+ ],
1975
+ "range_flattens": [
1976
+ null,
1977
+ false,
1978
+ null,
1979
+ true
1980
+ ],
1981
+ "load_eviction_policies": [
1982
+ "",
1983
+ "last",
1984
+ "",
1985
+ "last",
1986
+ "",
1987
+ "last",
1988
+ "",
1989
+ "",
1990
+ ""
1991
+ ],
1992
+ "num_warps": 8,
1993
+ "num_stages": 8,
1994
+ "indexing": [
1995
+ "pointer",
1996
+ "pointer",
1997
+ "pointer",
1998
+ "tensor_descriptor",
1999
+ "tensor_descriptor",
2000
+ "pointer",
2001
+ "pointer",
2002
+ "tensor_descriptor",
2003
+ "tensor_descriptor",
2004
+ "tensor_descriptor",
2005
+ "tensor_descriptor",
2006
+ "tensor_descriptor"
2007
+ ],
2008
+ "atomic_indexing": [],
2009
+ "pid_type": "flat"
2010
+ }
2011
+ },
2012
+ {
2013
+ "key": {
2014
+ "hidden_size": 5120,
2015
+ "num_tokens": 1024
2016
+ },
2017
+ "config": {
2018
+ "block_sizes": [
2019
+ 8192,
2020
+ 8192,
2021
+ 2048
2022
+ ],
2023
+ "range_unroll_factors": [
2024
+ 0,
2025
+ 4,
2026
+ 2,
2027
+ 3
2028
+ ],
2029
+ "range_warp_specializes": [],
2030
+ "range_num_stages": [],
2031
+ "range_multi_buffers": [
2032
+ null,
2033
+ null,
2034
+ true,
2035
+ true
2036
+ ],
2037
+ "range_flattens": [
2038
+ null,
2039
+ null,
2040
+ null,
2041
+ true
2042
+ ],
2043
+ "load_eviction_policies": [
2044
+ "",
2045
+ "",
2046
+ "",
2047
+ "",
2048
+ "last",
2049
+ "last",
2050
+ "",
2051
+ "first",
2052
+ "last"
2053
+ ],
2054
+ "num_warps": 8,
2055
+ "num_stages": 1,
2056
+ "indexing": [
2057
+ "pointer",
2058
+ "tensor_descriptor",
2059
+ "tensor_descriptor",
2060
+ "pointer",
2061
+ "pointer",
2062
+ "tensor_descriptor",
2063
+ "tensor_descriptor",
2064
+ "pointer",
2065
+ "tensor_descriptor",
2066
+ "pointer",
2067
+ "pointer",
2068
+ "tensor_descriptor"
2069
+ ],
2070
+ "atomic_indexing": [],
2071
+ "pid_type": "flat"
2072
+ }
2073
+ },
2074
+ {
2075
+ "key": {
2076
+ "hidden_size": 2048,
2077
+ "num_tokens": 2048
2078
+ },
2079
+ "config": {
2080
+ "block_sizes": [
2081
+ 2048,
2082
+ 2048,
2083
+ 1024
2084
+ ],
2085
+ "range_unroll_factors": [
2086
+ 0,
2087
+ 0,
2088
+ 4,
2089
+ 2
2090
+ ],
2091
+ "range_warp_specializes": [],
2092
+ "range_num_stages": [],
2093
+ "range_multi_buffers": [
2094
+ null,
2095
+ true,
2096
+ false,
2097
+ false
2098
+ ],
2099
+ "range_flattens": [
2100
+ null,
2101
+ null,
2102
+ true,
2103
+ null
2104
+ ],
2105
+ "load_eviction_policies": [
2106
+ "first",
2107
+ "",
2108
+ "first",
2109
+ "",
2110
+ "last",
2111
+ "last",
2112
+ "first",
2113
+ "first",
2114
+ "last"
2115
+ ],
2116
+ "num_warps": 4,
2117
+ "num_stages": 1,
2118
+ "indexing": [
2119
+ "tensor_descriptor",
2120
+ "tensor_descriptor",
2121
+ "pointer",
2122
+ "pointer",
2123
+ "tensor_descriptor",
2124
+ "pointer",
2125
+ "tensor_descriptor",
2126
+ "pointer",
2127
+ "tensor_descriptor",
2128
+ "pointer",
2129
+ "tensor_descriptor",
2130
+ "pointer"
2131
+ ],
2132
+ "atomic_indexing": [],
2133
+ "pid_type": "flat"
2134
+ }
2135
+ },
2136
+ {
2137
+ "key": {
2138
+ "hidden_size": 4096,
2139
+ "num_tokens": 2048
2140
+ },
2141
+ "config": {
2142
+ "block_sizes": [
2143
+ 4096,
2144
+ 4096,
2145
+ 2048
2146
+ ],
2147
+ "range_unroll_factors": [
2148
+ 0,
2149
+ 3,
2150
+ 2,
2151
+ 2
2152
+ ],
2153
+ "range_warp_specializes": [],
2154
+ "range_num_stages": [],
2155
+ "range_multi_buffers": [
2156
+ null,
2157
+ true,
2158
+ null,
2159
+ true
2160
+ ],
2161
+ "range_flattens": [
2162
+ null,
2163
+ false,
2164
+ null,
2165
+ null
2166
+ ],
2167
+ "load_eviction_policies": [
2168
+ "",
2169
+ "",
2170
+ "",
2171
+ "",
2172
+ "",
2173
+ "last",
2174
+ "first",
2175
+ "",
2176
+ ""
2177
+ ],
2178
+ "num_warps": 8,
2179
+ "num_stages": 1,
2180
+ "indexing": [
2181
+ "pointer",
2182
+ "pointer",
2183
+ "tensor_descriptor",
2184
+ "tensor_descriptor",
2185
+ "tensor_descriptor",
2186
+ "pointer",
2187
+ "tensor_descriptor",
2188
+ "pointer",
2189
+ "pointer",
2190
+ "tensor_descriptor",
2191
+ "pointer",
2192
+ "tensor_descriptor"
2193
+ ],
2194
+ "atomic_indexing": [],
2195
+ "pid_type": "flat"
2196
+ }
2197
+ },
2198
+ {
2199
+ "key": {
2200
+ "hidden_size": 5120,
2201
+ "num_tokens": 2048
2202
+ },
2203
+ "config": {
2204
+ "block_sizes": [
2205
+ 8192,
2206
+ 8192,
2207
+ 1024
2208
+ ],
2209
+ "range_unroll_factors": [
2210
+ 0,
2211
+ 3,
2212
+ 2,
2213
+ 4
2214
+ ],
2215
+ "range_warp_specializes": [],
2216
+ "range_num_stages": [],
2217
+ "range_multi_buffers": [
2218
+ null,
2219
+ true,
2220
+ null,
2221
+ null
2222
+ ],
2223
+ "range_flattens": [
2224
+ null,
2225
+ null,
2226
+ true,
2227
+ null
2228
+ ],
2229
+ "load_eviction_policies": [
2230
+ "first",
2231
+ "",
2232
+ "last",
2233
+ "last",
2234
+ "",
2235
+ "last",
2236
+ "first",
2237
+ "first",
2238
+ "last"
2239
+ ],
2240
+ "num_warps": 4,
2241
+ "num_stages": 1,
2242
+ "indexing": [
2243
+ "pointer",
2244
+ "tensor_descriptor",
2245
+ "pointer",
2246
+ "tensor_descriptor",
2247
+ "pointer",
2248
+ "tensor_descriptor",
2249
+ "pointer",
2250
+ "pointer",
2251
+ "tensor_descriptor",
2252
+ "tensor_descriptor",
2253
+ "tensor_descriptor",
2254
+ "tensor_descriptor"
2255
+ ],
2256
+ "atomic_indexing": [],
2257
+ "pid_type": "flat"
2258
+ }
2259
+ },
2260
+ {
2261
+ "key": {
2262
+ "hidden_size": 2048,
2263
+ "num_tokens": 4096
2264
+ },
2265
+ "config": {
2266
+ "block_sizes": [
2267
+ 2048,
2268
+ 2048,
2269
+ 1024
2270
+ ],
2271
+ "range_unroll_factors": [
2272
+ 0,
2273
+ 1,
2274
+ 4,
2275
+ 2
2276
+ ],
2277
+ "range_warp_specializes": [
2278
+ null,
2279
+ false,
2280
+ false,
2281
+ null
2282
+ ],
2283
+ "range_num_stages": [],
2284
+ "range_multi_buffers": [
2285
+ null,
2286
+ null,
2287
+ null,
2288
+ true
2289
+ ],
2290
+ "range_flattens": [
2291
+ null,
2292
+ true,
2293
+ true,
2294
+ null
2295
+ ],
2296
+ "load_eviction_policies": [
2297
+ "",
2298
+ "first",
2299
+ "",
2300
+ "last",
2301
+ "last",
2302
+ "",
2303
+ "",
2304
+ "",
2305
+ "last"
2306
+ ],
2307
+ "num_warps": 4,
2308
+ "num_stages": 1,
2309
+ "indexing": [
2310
+ "tensor_descriptor",
2311
+ "pointer",
2312
+ "tensor_descriptor",
2313
+ "pointer",
2314
+ "tensor_descriptor",
2315
+ "tensor_descriptor",
2316
+ "pointer",
2317
+ "tensor_descriptor",
2318
+ "tensor_descriptor",
2319
+ "tensor_descriptor",
2320
+ "tensor_descriptor",
2321
+ "tensor_descriptor"
2322
+ ],
2323
+ "atomic_indexing": [],
2324
+ "pid_type": "flat"
2325
+ }
2326
+ },
2327
+ {
2328
+ "key": {
2329
+ "hidden_size": 4096,
2330
+ "num_tokens": 4096
2331
+ },
2332
+ "config": {
2333
+ "block_sizes": [
2334
+ 4096,
2335
+ 4096,
2336
+ 2048
2337
+ ],
2338
+ "range_unroll_factors": [
2339
+ 0,
2340
+ 4,
2341
+ 3,
2342
+ 2
2343
+ ],
2344
+ "range_warp_specializes": [
2345
+ null,
2346
+ null,
2347
+ false,
2348
+ null
2349
+ ],
2350
+ "range_num_stages": [],
2351
+ "range_multi_buffers": [
2352
+ null,
2353
+ false,
2354
+ null,
2355
+ null
2356
+ ],
2357
+ "range_flattens": [
2358
+ null,
2359
+ false,
2360
+ null,
2361
+ true
2362
+ ],
2363
+ "load_eviction_policies": [
2364
+ "",
2365
+ "last",
2366
+ "",
2367
+ "last",
2368
+ "",
2369
+ "last",
2370
+ "",
2371
+ "",
2372
+ ""
2373
+ ],
2374
+ "num_warps": 8,
2375
+ "num_stages": 8,
2376
+ "indexing": [
2377
+ "pointer",
2378
+ "pointer",
2379
+ "pointer",
2380
+ "tensor_descriptor",
2381
+ "tensor_descriptor",
2382
+ "pointer",
2383
+ "pointer",
2384
+ "tensor_descriptor",
2385
+ "tensor_descriptor",
2386
+ "tensor_descriptor",
2387
+ "tensor_descriptor",
2388
+ "tensor_descriptor"
2389
+ ],
2390
+ "atomic_indexing": [],
2391
+ "pid_type": "flat"
2392
+ }
2393
+ },
2394
+ {
2395
+ "key": {
2396
+ "hidden_size": 5120,
2397
+ "num_tokens": 4096
2398
+ },
2399
+ "config": {
2400
+ "block_sizes": [
2401
+ 8192,
2402
+ 8192,
2403
+ 2048
2404
+ ],
2405
+ "range_unroll_factors": [
2406
+ 0,
2407
+ 4,
2408
+ 3,
2409
+ 2
2410
+ ],
2411
+ "range_warp_specializes": [],
2412
+ "range_num_stages": [],
2413
+ "range_multi_buffers": [
2414
+ null,
2415
+ false,
2416
+ null,
2417
+ false
2418
+ ],
2419
+ "range_flattens": [
2420
+ null,
2421
+ true,
2422
+ true,
2423
+ null
2424
+ ],
2425
+ "load_eviction_policies": [
2426
+ "",
2427
+ "",
2428
+ "",
2429
+ "",
2430
+ "last",
2431
+ "",
2432
+ "first",
2433
+ "first",
2434
+ ""
2435
+ ],
2436
+ "num_warps": 8,
2437
+ "num_stages": 1,
2438
+ "indexing": [
2439
+ "pointer",
2440
+ "pointer",
2441
+ "pointer",
2442
+ "pointer",
2443
+ "tensor_descriptor",
2444
+ "pointer",
2445
+ "tensor_descriptor",
2446
+ "pointer",
2447
+ "tensor_descriptor",
2448
+ "pointer",
2449
+ "pointer",
2450
+ "pointer"
2451
+ ],
2452
+ "atomic_indexing": [],
2453
+ "pid_type": "flat"
2454
+ }
2455
+ },
2456
+ {
2457
+ "key": {
2458
+ "hidden_size": 2048,
2459
+ "num_tokens": 8192
2460
+ },
2461
+ "config": {
2462
+ "block_sizes": [
2463
+ 2048,
2464
+ 2048,
2465
+ 1024
2466
+ ],
2467
+ "range_unroll_factors": [
2468
+ 0,
2469
+ 0,
2470
+ 4,
2471
+ 2
2472
+ ],
2473
+ "range_warp_specializes": [],
2474
+ "range_num_stages": [],
2475
+ "range_multi_buffers": [
2476
+ null,
2477
+ true,
2478
+ false,
2479
+ false
2480
+ ],
2481
+ "range_flattens": [
2482
+ null,
2483
+ null,
2484
+ true,
2485
+ null
2486
+ ],
2487
+ "load_eviction_policies": [
2488
+ "first",
2489
+ "",
2490
+ "first",
2491
+ "",
2492
+ "last",
2493
+ "last",
2494
+ "first",
2495
+ "first",
2496
+ "last"
2497
+ ],
2498
+ "num_warps": 4,
2499
+ "num_stages": 1,
2500
+ "indexing": [
2501
+ "tensor_descriptor",
2502
+ "tensor_descriptor",
2503
+ "pointer",
2504
+ "pointer",
2505
+ "tensor_descriptor",
2506
+ "pointer",
2507
+ "tensor_descriptor",
2508
+ "pointer",
2509
+ "tensor_descriptor",
2510
+ "pointer",
2511
+ "tensor_descriptor",
2512
+ "pointer"
2513
+ ],
2514
+ "atomic_indexing": [],
2515
+ "pid_type": "flat"
2516
+ }
2517
+ },
2518
+ {
2519
+ "key": {
2520
+ "hidden_size": 4096,
2521
+ "num_tokens": 8192
2522
+ },
2523
+ "config": {
2524
+ "block_sizes": [
2525
+ 4096,
2526
+ 4096,
2527
+ 2048
2528
+ ],
2529
+ "range_unroll_factors": [
2530
+ 0,
2531
+ 4,
2532
+ 3,
2533
+ 2
2534
+ ],
2535
+ "range_warp_specializes": [
2536
+ null,
2537
+ null,
2538
+ false,
2539
+ null
2540
+ ],
2541
+ "range_num_stages": [],
2542
+ "range_multi_buffers": [
2543
+ null,
2544
+ false,
2545
+ null,
2546
+ null
2547
+ ],
2548
+ "range_flattens": [
2549
+ null,
2550
+ false,
2551
+ null,
2552
+ true
2553
+ ],
2554
+ "load_eviction_policies": [
2555
+ "",
2556
+ "last",
2557
+ "",
2558
+ "last",
2559
+ "",
2560
+ "last",
2561
+ "",
2562
+ "",
2563
+ ""
2564
+ ],
2565
+ "num_warps": 8,
2566
+ "num_stages": 8,
2567
+ "indexing": [
2568
+ "pointer",
2569
+ "pointer",
2570
+ "pointer",
2571
+ "tensor_descriptor",
2572
+ "tensor_descriptor",
2573
+ "pointer",
2574
+ "pointer",
2575
+ "tensor_descriptor",
2576
+ "tensor_descriptor",
2577
+ "tensor_descriptor",
2578
+ "tensor_descriptor",
2579
+ "tensor_descriptor"
2580
+ ],
2581
+ "atomic_indexing": [],
2582
+ "pid_type": "flat"
2583
+ }
2584
+ },
2585
+ {
2586
+ "key": {
2587
+ "hidden_size": 5120,
2588
+ "num_tokens": 8192
2589
+ },
2590
+ "config": {
2591
+ "block_sizes": [
2592
+ 8192,
2593
+ 8192,
2594
+ 2048
2595
+ ],
2596
+ "range_unroll_factors": [
2597
+ 0,
2598
+ 4,
2599
+ 3,
2600
+ 2
2601
+ ],
2602
+ "range_warp_specializes": [],
2603
+ "range_num_stages": [],
2604
+ "range_multi_buffers": [
2605
+ null,
2606
+ false,
2607
+ null,
2608
+ false
2609
+ ],
2610
+ "range_flattens": [
2611
+ null,
2612
+ true,
2613
+ true,
2614
+ null
2615
+ ],
2616
+ "load_eviction_policies": [
2617
+ "",
2618
+ "",
2619
+ "",
2620
+ "",
2621
+ "last",
2622
+ "",
2623
+ "first",
2624
+ "first",
2625
+ ""
2626
+ ],
2627
+ "num_warps": 8,
2628
+ "num_stages": 1,
2629
+ "indexing": [
2630
+ "pointer",
2631
+ "pointer",
2632
+ "pointer",
2633
+ "pointer",
2634
+ "tensor_descriptor",
2635
+ "pointer",
2636
+ "tensor_descriptor",
2637
+ "pointer",
2638
+ "tensor_descriptor",
2639
+ "pointer",
2640
+ "pointer",
2641
+ "pointer"
2642
+ ],
2643
+ "atomic_indexing": [],
2644
+ "pid_type": "flat"
2645
+ }
2646
+ }
2647
+]
\ No newline at end of file
vllm/kernels/helion/configs/rms_norm_dynamic_per_token_quant/nvidia_h100.json
new
+3663
@@ -0,0 +1,3663 @@
1
+[
2
+ {
3
+ "key": {
4
+ "hidden_size": 2048,
5
+ "num_tokens": 1
6
+ },
7
+ "config": {
8
+ "block_sizes": [
9
+ 2048,
10
+ 2048,
11
+ 2048
12
+ ],
13
+ "range_unroll_factors": [
14
+ 0,
15
+ 0,
16
+ 0,
17
+ 1
18
+ ],
19
+ "range_warp_specializes": [],
20
+ "range_num_stages": [],
21
+ "range_multi_buffers": [
22
+ null,
23
+ false,
24
+ null,
25
+ false
26
+ ],
27
+ "range_flattens": [
28
+ null,
29
+ null,
30
+ null,
31
+ true
32
+ ],
33
+ "load_eviction_policies": [
34
+ "last",
35
+ "",
36
+ "last",
37
+ "",
38
+ "first",
39
+ "",
40
+ "first",
41
+ "first",
42
+ ""
43
+ ],
44
+ "num_warps": 8,
45
+ "num_stages": 8,
46
+ "indexing": [
47
+ "tensor_descriptor",
48
+ "tensor_descriptor",
49
+ "tensor_descriptor",
50
+ "tensor_descriptor",
51
+ "pointer",
52
+ "pointer",
53
+ "pointer",
54
+ "tensor_descriptor",
55
+ "tensor_descriptor",
56
+ "tensor_descriptor",
57
+ "tensor_descriptor",
58
+ "tensor_descriptor"
59
+ ],
60
+ "atomic_indexing": [],
61
+ "pid_type": "flat"
62
+ }
63
+ },
64
+ {
65
+ "key": {
66
+ "hidden_size": 4096,
67
+ "num_tokens": 1
68
+ },
69
+ "config": {
70
+ "block_sizes": [
71
+ 4096,
72
+ 4096,
73
+ 4096
74
+ ],
75
+ "range_unroll_factors": [
76
+ 2,
77
+ 3,
78
+ 1,
79
+ 0
80
+ ],
81
+ "range_warp_specializes": [],
82
+ "range_multi_buffers": [
83
+ true,
84
+ null,
85
+ false,
86
+ null
87
+ ],
88
+ "range_flattens": [
89
+ null,
90
+ null,
91
+ false,
92
+ true
93
+ ],
94
+ "load_eviction_policies": [
95
+ "last",
96
+ "",
97
+ "last",
98
+ "",
99
+ "last",
100
+ "last",
101
+ "",
102
+ "first",
103
+ "last"
104
+ ],
105
+ "num_warps": 16,
106
+ "num_stages": 7,
107
+ "indexing": [
108
+ "pointer",
109
+ "pointer",
110
+ "tensor_descriptor",
111
+ "tensor_descriptor",
112
+ "pointer",
113
+ "pointer",
114
+ "tensor_descriptor",
115
+ "pointer",
116
+ "pointer",
117
+ "tensor_descriptor",
118
+ "pointer",
119
+ "pointer"
120
+ ],
121
+ "atomic_indexing": [],
122
+ "pid_type": "persistent_interleaved",
123
+ "num_sm_multiplier": 1,
124
+ "maxnreg": 128
125
+ }
126
+ },
127
+ {
128
+ "key": {
129
+ "hidden_size": 8192,
130
+ "num_tokens": 1
131
+ },
132
+ "config": {
133
+ "block_sizes": [
134
+ 8192,
135
+ 8192,
136
+ 2048
137
+ ],
138
+ "range_unroll_factors": [
139
+ 0,
140
+ 2,
141
+ 3,
142
+ 2
143
+ ],
144
+ "range_warp_specializes": [],
145
+ "range_num_stages": [],
146
+ "range_multi_buffers": [
147
+ null,
148
+ false,
149
+ false,
150
+ false
151
+ ],
152
+ "range_flattens": [
153
+ null,
154
+ false,
155
+ true,
156
+ false
157
+ ],
158
+ "load_eviction_policies": [
159
+ "last",
160
+ "",
161
+ "last",
162
+ "",
163
+ "",
164
+ "last",
165
+ "last",
166
+ "first",
167
+ "first"
168
+ ],
169
+ "num_warps": 32,
170
+ "num_stages": 6,
171
+ "indexing": [
172
+ "tensor_descriptor",
173
+ "pointer",
174
+ "pointer",
175
+ "pointer",
176
+ "tensor_descriptor",
177
+ "pointer",
178
+ "pointer",
179
+ "tensor_descriptor",
180
+ "pointer",
181
+ "tensor_descriptor",
182
+ "pointer",
183
+ "pointer"
184
+ ],
185
+ "atomic_indexing": [],
186
+ "pid_type": "flat"
187
+ }
188
+ },
189
+ {
190
+ "key": {
191
+ "hidden_size": 2048,
192
+ "num_tokens": 2
193
+ },
194
+ "config": {
195
+ "block_sizes": [
196
+ 2048,
197
+ 2048,
198
+ 2048
199
+ ],
200
+ "range_unroll_factors": [
201
+ 0,
202
+ 3,
203
+ 3,
204
+ 1
205
+ ],
206
+ "range_warp_specializes": [],
207
+ "range_num_stages": [],
208
+ "range_multi_buffers": [
209
+ null,
210
+ false,
211
+ true,
212
+ null
213
+ ],
214
+ "range_flattens": [
215
+ null,
216
+ false,
217
+ false,
218
+ true
219
+ ],
220
+ "load_eviction_policies": [
221
+ "last",
222
+ "",
223
+ "last",
224
+ "",
225
+ "last",
226
+ "first",
227
+ "first",
228
+ "",
229
+ ""
230
+ ],
231
+ "num_warps": 8,
232
+ "num_stages": 4,
233
+ "indexing": [
234
+ "tensor_descriptor",
235
+ "pointer",
236
+ "tensor_descriptor",
237
+ "tensor_descriptor",
238
+ "pointer",
239
+ "tensor_descriptor",
240
+ "tensor_descriptor",
241
+ "pointer",
242
+ "tensor_descriptor",
243
+ "tensor_descriptor",
244
+ "tensor_descriptor",
245
+ "tensor_descriptor"
246
+ ],
247
+ "atomic_indexing": [],
248
+ "pid_type": "flat"
249
+ }
250
+ },
251
+ {
252
+ "key": {
253
+ "hidden_size": 4096,
254
+ "num_tokens": 2
255
+ },
256
+ "config": {
257
+ "block_sizes": [
258
+ 4096,
259
+ 4096,
260
+ 4096
261
+ ],
262
+ "range_unroll_factors": [
263
+ 0,
264
+ 4,
265
+ 4,
266
+ 4
267
+ ],
268
+ "range_warp_specializes": [],
269
+ "range_num_stages": [],
270
+ "range_multi_buffers": [
271
+ null,
272
+ false,
273
+ true,
274
+ false
275
+ ],
276
+ "range_flattens": [
277
+ null,
278
+ true,
279
+ true,
280
+ false
281
+ ],
282
+ "load_eviction_policies": [
283
+ "last",
284
+ "",
285
+ "last",
286
+ "",
287
+ "last",
288
+ "first",
289
+ "first",
290
+ "first",
291
+ ""
292
+ ],
293
+ "num_warps": 16,
294
+ "num_stages": 3,
295
+ "indexing": [
296
+ "tensor_descriptor",
297
+ "tensor_descriptor",
298
+ "tensor_descriptor",
299
+ "tensor_descriptor",
300
+ "tensor_descriptor",
301
+ "pointer",
302
+ "pointer",
303
+ "pointer",
304
+ "pointer",
305
+ "pointer",
306
+ "pointer",
307
+ "pointer"
308
+ ],
309
+ "atomic_indexing": [],
310
+ "pid_type": "flat"
311
+ }
312
+ },
313
+ {
314
+ "key": {
315
+ "hidden_size": 8192,
316
+ "num_tokens": 2
317
+ },
318
+ "config": {
319
+ "block_sizes": [
320
+ 8192,
321
+ 8192,
322
+ 8192
323
+ ],
324
+ "range_unroll_factors": [
325
+ 0,
326
+ 3,
327
+ 2,
328
+ 2
329
+ ],
330
+ "range_warp_specializes": [],
331
+ "range_num_stages": [],
332
+ "range_multi_buffers": [
333
+ null,
334
+ null,
335
+ true,
336
+ null
337
+ ],
338
+ "range_flattens": [
339
+ null,
340
+ true,
341
+ false,
342
+ true
343
+ ],
344
+ "load_eviction_policies": [
345
+ "last",
346
+ "first",
347
+ "last",
348
+ "first",
349
+ "first",
350
+ "",
351
+ "last",
352
+ "",
353
+ "last"
354
+ ],
355
+ "num_warps": 32,
356
+ "num_stages": 8,
357
+ "indexing": [
358
+ "pointer",
359
+ "pointer",
360
+ "tensor_descriptor",
361
+ "pointer",
362
+ "pointer",
363
+ "tensor_descriptor",
364
+ "pointer",
365
+ "tensor_descriptor",
366
+ "tensor_descriptor",
367
+ "tensor_descriptor",
368
+ "pointer",
369
+ "pointer"
370
+ ],
371
+ "atomic_indexing": [],
372
+ "pid_type": "flat"
373
+ }
374
+ },
375
+ {
376
+ "key": {
377
+ "hidden_size": 2048,
378
+ "num_tokens": 4
379
+ },
380
+ "config": {
381
+ "block_sizes": [
382
+ 2048,
383
+ 2048,
384
+ 2048
385
+ ],
386
+ "range_unroll_factors": [
387
+ 0,
388
+ 3,
389
+ 2,
390
+ 1
391
+ ],
392
+ "range_warp_specializes": [],
393
+ "range_num_stages": [],
394
+ "range_multi_buffers": [
395
+ null,
396
+ false,
397
+ false,
398
+ null
399
+ ],
400
+ "range_flattens": [
401
+ null,
402
+ true,
403
+ null,
404
+ false
405
+ ],
406
+ "load_eviction_policies": [
407
+ "",
408
+ "",
409
+ "",
410
+ "",
411
+ "last",
412
+ "first",
413
+ "last",
414
+ "first",
415
+ "last"
416
+ ],
417
+ "num_warps": 8,
418
+ "num_stages": 4,
419
+ "indexing": [
420
+ "pointer",
421
+ "tensor_descriptor",
422
+ "pointer",
423
+ "pointer",
424
+ "pointer",
425
+ "tensor_descriptor",
426
+ "tensor_descriptor",
427
+ "pointer",
428
+ "tensor_descriptor",
429
+ "pointer",
430
+ "pointer",
431
+ "pointer"
432
+ ],
433
+ "atomic_indexing": [],
434
+ "pid_type": "flat"
435
+ }
436
+ },
437
+ {
438
+ "key": {
439
+ "hidden_size": 4096,
440
+ "num_tokens": 4
441
+ },
442
+ "config": {
443
+ "block_sizes": [
444
+ 4096,
445
+ 4096,
446
+ 4096
447
+ ],
448
+ "range_unroll_factors": [
449
+ 0,
450
+ 4,
451
+ 2,
452
+ 4
453
+ ],
454
+ "range_warp_specializes": [],
455
+ "range_num_stages": [],
456
+ "range_multi_buffers": [
457
+ null,
458
+ null,
459
+ false,
460
+ null
461
+ ],
462
+ "range_flattens": [
463
+ null,
464
+ false,
465
+ true,
466
+ true
467
+ ],
468
+ "load_eviction_policies": [
469
+ "",
470
+ "first",
471
+ "",
472
+ "first",
473
+ "last",
474
+ "",
475
+ "last",
476
+ "",
477
+ "last"
478
+ ],
479
+ "num_warps": 16,
480
+ "num_stages": 1,
481
+ "indexing": [
482
+ "tensor_descriptor",
483
+ "tensor_descriptor",
484
+ "tensor_descriptor",
485
+ "pointer",
486
+ "tensor_descriptor",
487
+ "tensor_descriptor",
488
+ "tensor_descriptor",
489
+ "pointer",
490
+ "tensor_descriptor",
491
+ "pointer",
492
+ "tensor_descriptor",
493
+ "tensor_descriptor"
494
+ ],
495
+ "atomic_indexing": [],
496
+ "pid_type": "flat"
497
+ }
498
+ },
499
+ {
500
+ "key": {
501
+ "hidden_size": 8192,
502
+ "num_tokens": 4
503
+ },
504
+ "config": {
505
+ "block_sizes": [
506
+ 8192,
507
+ 8192,
508
+ 4096
509
+ ],
510
+ "range_unroll_factors": [
511
+ 0,
512
+ 1,
513
+ 3,
514
+ 2
515
+ ],
516
+ "range_warp_specializes": [],
517
+ "range_num_stages": [],
518
+ "range_multi_buffers": [
519
+ null,
520
+ null,
521
+ true,
522
+ null
523
+ ],
524
+ "range_flattens": [
525
+ null,
526
+ false,
527
+ true,
528
+ false
529
+ ],
530
+ "load_eviction_policies": [
531
+ "",
532
+ "",
533
+ "",
534
+ "",
535
+ "",
536
+ "last",
537
+ "first",
538
+ "",
539
+ ""
540
+ ],
541
+ "num_warps": 16,
542
+ "num_stages": 1,
543
+ "indexing": [
544
+ "tensor_descriptor",
545
+ "tensor_descriptor",
546
+ "pointer",
547
+ "tensor_descriptor",
548
+ "pointer",
549
+ "tensor_descriptor",
550
+ "pointer",
551
+ "tensor_descriptor",
552
+ "tensor_descriptor",
553
+ "pointer",
554
+ "pointer",
555
+ "tensor_descriptor"
556
+ ],
557
+ "atomic_indexing": [],
558
+ "pid_type": "flat"
559
+ }
560
+ },
561
+ {
562
+ "key": {
563
+ "hidden_size": 2048,
564
+ "num_tokens": 8
565
+ },
566
+ "config": {
567
+ "block_sizes": [
568
+ 2048,
569
+ 2048,
570
+ 2048
571
+ ],
572
+ "range_unroll_factors": [
573
+ 0,
574
+ 4,
575
+ 3,
576
+ 4
577
+ ],
578
+ "range_warp_specializes": [],
579
+ "range_num_stages": [],
580
+ "range_multi_buffers": [
581
+ null,
582
+ true,
583
+ true,
584
+ true
585
+ ],
586
+ "range_flattens": [
587
+ null,
588
+ false,
589
+ false,
590
+ null
591
+ ],
592
+ "load_eviction_policies": [
593
+ "",
594
+ "",
595
+ "",
596
+ "",
597
+ "last",
598
+ "",
599
+ "",
600
+ "",
601
+ "first"
602
+ ],
603
+ "num_warps": 8,
604
+ "num_stages": 1,
605
+ "indexing": [
606
+ "tensor_descriptor",
607
+ "tensor_descriptor",
608
+ "tensor_descriptor",
609
+ "tensor_descriptor",
610
+ "tensor_descriptor",
611
+ "pointer",
612
+ "tensor_descriptor",
613
+ "tensor_descriptor",
614
+ "pointer",
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": 8
627
+ },
628
+ "config": {
629
+ "block_sizes": [
630
+ 4096,
631
+ 4096,
632
+ 2048
633
+ ],
634
+ "range_unroll_factors": [
635
+ 0,
636
+ 0,
637
+ 1,
638
+ 4
639
+ ],
640
+ "range_warp_specializes": [],
641
+ "range_num_stages": [],
642
+ "range_multi_buffers": [
643
+ null,
644
+ null,
645
+ null,
646
+ false
647
+ ],
648
+ "range_flattens": [
649
+ null,
650
+ null,
651
+ false,
652
+ true
653
+ ],
654
+ "load_eviction_policies": [
655
+ "first",
656
+ "",
657
+ "last",
658
+ "",
659
+ "first",
660
+ "last",
661
+ "last",
662
+ "last",
663
+ "first"
664
+ ],
665
+ "num_warps": 16,
666
+ "num_stages": 1,
667
+ "indexing": [
668
+ "tensor_descriptor",
669
+ "pointer",
670
+ "pointer",
671
+ "pointer",
672
+ "tensor_descriptor",
673
+ "pointer",
674
+ "tensor_descriptor",
675
+ "pointer",
676
+ "tensor_descriptor",
677
+ "tensor_descriptor",
678
+ "pointer",
679
+ "pointer"
680
+ ],
681
+ "atomic_indexing": [],
682
+ "pid_type": "flat"
683
+ }
684
+ },
685
+ {
686
+ "key": {
687
+ "hidden_size": 8192,
688
+ "num_tokens": 8
689
+ },
690
+ "config": {
691
+ "block_sizes": [
692
+ 8192,
693
+ 8192,
694
+ 8192
695
+ ],
696
+ "range_unroll_factors": [
697
+ 0,
698
+ 2,
699
+ 0,
700
+ 1
701
+ ],
702
+ "range_warp_specializes": [],
703
+ "range_num_stages": [],
704
+ "range_multi_buffers": [
705
+ null,
706
+ true,
707
+ null,
708
+ false
709
+ ],
710
+ "range_flattens": [
711
+ null,
712
+ null,
713
+ false,
714
+ null
715
+ ],
716
+ "load_eviction_policies": [
717
+ "last",
718
+ "",
719
+ "last",
720
+ "",
721
+ "",
722
+ "",
723
+ "first",
724
+ "first",
725
+ ""
726
+ ],
727
+ "num_warps": 32,
728
+ "num_stages": 7,
729
+ "indexing": [
730
+ "pointer",
731
+ "pointer",
732
+ "tensor_descriptor",
733
+ "pointer",
734
+ "tensor_descriptor",
735
+ "pointer",
736
+ "tensor_descriptor",
737
+ "tensor_descriptor",
738
+ "pointer",
739
+ "tensor_descriptor",
740
+ "pointer",
741
+ "pointer"
742
+ ],
743
+ "atomic_indexing": [],
744
+ "pid_type": "flat"
745
+ }
746
+ },
747
+ {
748
+ "key": {
749
+ "hidden_size": 2048,
750
+ "num_tokens": 16
751
+ },
752
+ "config": {
753
+ "block_sizes": [
754
+ 2048,
755
+ 2048,
756
+ 2048
757
+ ],
758
+ "range_unroll_factors": [
759
+ 0,
760
+ 3,
761
+ 2,
762
+ 4
763
+ ],
764
+ "range_warp_specializes": [],
765
+ "range_num_stages": [],
766
+ "range_multi_buffers": [
767
+ null,
768
+ null,
769
+ null,
770
+ true
771
+ ],
772
+ "range_flattens": [
773
+ null,
774
+ true,
775
+ true,
776
+ false
777
+ ],
778
+ "load_eviction_policies": [
779
+ "first",
780
+ "first",
781
+ "first",
782
+ "",
783
+ "",
784
+ "last",
785
+ "",
786
+ "",
787
+ "first"
788
+ ],
789
+ "num_warps": 16,
790
+ "num_stages": 4,
791
+ "indexing": [
792
+ "tensor_descriptor",
793
+ "tensor_descriptor",
794
+ "tensor_descriptor",
795
+ "pointer",
796
+ "tensor_descriptor",
797
+ "tensor_descriptor",
798
+ "tensor_descriptor",
799
+ "pointer",
800
+ "pointer",
801
+ "pointer",
802
+ "tensor_descriptor",
803
+ "pointer"
804
+ ],
805
+ "atomic_indexing": [],
806
+ "pid_type": "flat"
807
+ }
808
+ },
809
+ {
810
+ "key": {
811
+ "hidden_size": 4096,
812
+ "num_tokens": 16
813
+ },
814
+ "config": {
815
+ "block_sizes": [
816
+ 4096,
817
+ 4096,
818
+ 4096
819
+ ],
820
+ "range_unroll_factors": [
821
+ 0,
822
+ 2,
823
+ 4,
824
+ 3
825
+ ],
826
+ "range_warp_specializes": [],
827
+ "range_num_stages": [],
828
+ "range_multi_buffers": [
829
+ null,
830
+ true,
831
+ false,
832
+ true
833
+ ],
834
+ "range_flattens": [
835
+ null,
836
+ null,
837
+ false,
838
+ true
839
+ ],
840
+ "load_eviction_policies": [
841
+ "",
842
+ "",
843
+ "first",
844
+ "",
845
+ "last",
846
+ "last",
847
+ "",
848
+ "first",
849
+ ""
850
+ ],
851
+ "num_warps": 16,
852
+ "num_stages": 4,
853
+ "indexing": [
854
+ "pointer",
855
+ "tensor_descriptor",
856
+ "pointer",
857
+ "pointer",
858
+ "tensor_descriptor",
859
+ "pointer",
860
+ "tensor_descriptor",
861
+ "tensor_descriptor",
862
+ "pointer",
863
+ "pointer",
864
+ "tensor_descriptor",
865
+ "pointer"
866
+ ],
867
+ "atomic_indexing": [],
868
+ "pid_type": "flat"
869
+ }
870
+ },
871
+ {
872
+ "key": {
873
+ "hidden_size": 8192,
874
+ "num_tokens": 16
875
+ },
876
+ "config": {
877
+ "block_sizes": [
878
+ 8192,
879
+ 8192,
880
+ 8192
881
+ ],
882
+ "range_unroll_factors": [
883
+ 0,
884
+ 2,
885
+ 4,
886
+ 3
887
+ ],
888
+ "range_warp_specializes": [],
889
+ "range_num_stages": [],
890
+ "range_multi_buffers": [
891
+ null,
892
+ false,
893
+ null,
894
+ null
895
+ ],
896
+ "range_flattens": [
897
+ null,
898
+ null,
899
+ true,
900
+ null
901
+ ],
902
+ "load_eviction_policies": [
903
+ "last",
904
+ "last",
905
+ "last",
906
+ "last",
907
+ "first",
908
+ "first",
909
+ "",
910
+ "",
911
+ "last"
912
+ ],
913
+ "num_warps": 32,
914
+ "num_stages": 2,
915
+ "indexing": [
916
+ "tensor_descriptor",
917
+ "tensor_descriptor",
918
+ "tensor_descriptor",
919
+ "tensor_descriptor",
920
+ "tensor_descriptor",
921
+ "tensor_descriptor",
922
+ "pointer",
923
+ "pointer",
924
+ "pointer",
925
+ "tensor_descriptor",
926
+ "pointer",
927
+ "pointer"
928
+ ],
929
+ "atomic_indexing": [],
930
+ "pid_type": "flat"
931
+ }
932
+ },
933
+ {
934
+ "key": {
935
+ "hidden_size": 2048,
936
+ "num_tokens": 32
937
+ },
938
+ "config": {
939
+ "block_sizes": [
940
+ 2048,
941
+ 2048,
942
+ 2048
943
+ ],
944
+ "range_unroll_factors": [
945
+ 0,
946
+ 2,
947
+ 3,
948
+ 2
949
+ ],
950
+ "range_warp_specializes": [],
951
+ "range_num_stages": [],
952
+ "range_multi_buffers": [
953
+ null,
954
+ null,
955
+ null,
956
+ null
957
+ ],
958
+ "range_flattens": [
959
+ null,
960
+ true,
961
+ true,
962
+ false
963
+ ],
964
+ "load_eviction_policies": [
965
+ "last",
966
+ "",
967
+ "last",
968
+ "last",
969
+ "last",
970
+ "last",
971
+ "last",
972
+ "first",
973
+ "first"
974
+ ],
975
+ "num_warps": 16,
976
+ "num_stages": 1,
977
+ "indexing": [
978
+ "tensor_descriptor",
979
+ "pointer",
980
+ "tensor_descriptor",
981
+ "pointer",
982
+ "pointer",
983
+ "tensor_descriptor",
984
+ "tensor_descriptor",
985
+ "pointer",
986
+ "tensor_descriptor",
987
+ "tensor_descriptor",
988
+ "tensor_descriptor",
989
+ "tensor_descriptor"
990
+ ],
991
+ "atomic_indexing": [],
992
+ "pid_type": "flat"
993
+ }
994
+ },
995
+ {
996
+ "key": {
997
+ "hidden_size": 4096,
998
+ "num_tokens": 32
999
+ },
1000
+ "config": {
1001
+ "block_sizes": [
1002
+ 4096,
1003
+ 4096,
1004
+ 4096
1005
+ ],
1006
+ "range_unroll_factors": [
1007
+ 0,
1008
+ 2,
1009
+ 2,
1010
+ 1
1011
+ ],
1012
+ "range_warp_specializes": [],
1013
+ "range_num_stages": [],
1014
+ "range_multi_buffers": [
1015
+ null,
1016
+ null,
1017
+ true,
1018
+ null
1019
+ ],
1020
+ "range_flattens": [
1021
+ null,
1022
+ false,
1023
+ null,
1024
+ false
1025
+ ],
1026
+ "load_eviction_policies": [
1027
+ "last",
1028
+ "",
1029
+ "last",
1030
+ "",
1031
+ "first",
1032
+ "first",
1033
+ "",
1034
+ "first",
1035
+ ""
1036
+ ],
1037
+ "num_warps": 16,
1038
+ "num_stages": 2,
1039
+ "indexing": [
1040
+ "tensor_descriptor",
1041
+ "tensor_descriptor",
1042
+ "pointer",
1043
+ "tensor_descriptor",
1044
+ "pointer",
1045
+ "tensor_descriptor",
1046
+ "pointer",
1047
+ "tensor_descriptor",
1048
+ "tensor_descriptor",
1049
+ "pointer",
1050
+ "pointer",
1051
+ "tensor_descriptor"
1052
+ ],
1053
+ "atomic_indexing": [],
1054
+ "pid_type": "flat"
1055
+ }
1056
+ },
1057
+ {
1058
+ "key": {
1059
+ "hidden_size": 8192,
1060
+ "num_tokens": 32
1061
+ },
1062
+ "config": {
1063
+ "block_sizes": [
1064
+ 8192,
1065
+ 8192,
1066
+ 8192
1067
+ ],
1068
+ "range_unroll_factors": [
1069
+ 0,
1070
+ 3,
1071
+ 3,
1072
+ 3
1073
+ ],
1074
+ "range_warp_specializes": [],
1075
+ "range_num_stages": [],
1076
+ "range_multi_buffers": [
1077
+ null,
1078
+ null,
1079
+ false,
1080
+ false
1081
+ ],
1082
+ "range_flattens": [
1083
+ null,
1084
+ true,
1085
+ null,
1086
+ true
1087
+ ],
1088
+ "load_eviction_policies": [
1089
+ "first",
1090
+ "last",
1091
+ "first",
1092
+ "last",
1093
+ "",
1094
+ "",
1095
+ "last",
1096
+ "",
1097
+ ""
1098
+ ],
1099
+ "num_warps": 32,
1100
+ "num_stages": 1,
1101
+ "indexing": [
1102
+ "pointer",
1103
+ "tensor_descriptor",
1104
+ "tensor_descriptor",
1105
+ "pointer",
1106
+ "tensor_descriptor",
1107
+ "tensor_descriptor",
1108
+ "pointer",
1109
+ "tensor_descriptor",
1110
+ "pointer",
1111
+ "tensor_descriptor",
1112
+ "pointer",
1113
+ "pointer"
1114
+ ],
1115
+ "atomic_indexing": [],
1116
+ "pid_type": "flat"
1117
+ }
1118
+ },
1119
+ {
1120
+ "key": {
1121
+ "hidden_size": 2048,
1122
+ "num_tokens": 64
1123
+ },
1124
+ "config": {
1125
+ "block_sizes": [
1126
+ 2048,
1127
+ 2048,
1128
+ 2048
1129
+ ],
1130
+ "range_unroll_factors": [
1131
+ 0,
1132
+ 2,
1133
+ 4,
1134
+ 3
1135
+ ],
1136
+ "range_warp_specializes": [],
1137
+ "range_num_stages": [],
1138
+ "range_multi_buffers": [
1139
+ null,
1140
+ false,
1141
+ false,
1142
+ false
1143
+ ],
1144
+ "range_flattens": [
1145
+ null,
1146
+ false,
1147
+ false,
1148
+ true
1149
+ ],
1150
+ "load_eviction_policies": [
1151
+ "last",
1152
+ "last",
1153
+ "last",
1154
+ "last",
1155
+ "",
1156
+ "first",
1157
+ "last",
1158
+ "",
1159
+ ""
1160
+ ],
1161
+ "num_warps": 16,
1162
+ "num_stages": 3,
1163
+ "indexing": [
1164
+ "tensor_descriptor",
1165
+ "tensor_descriptor",
1166
+ "tensor_descriptor",
1167
+ "pointer",
1168
+ "pointer",
1169
+ "pointer",
1170
+ "pointer",
1171
+ "tensor_descriptor",
1172
+ "tensor_descriptor",
1173
+ "tensor_descriptor",
1174
+ "pointer",
1175
+ "pointer"
1176
+ ],
1177
+ "atomic_indexing": [],
1178
+ "pid_type": "flat"
1179
+ }
1180
+ },
1181
+ {
1182
+ "key": {
1183
+ "hidden_size": 4096,
1184
+ "num_tokens": 64
1185
+ },
1186
+ "config": {
1187
+ "block_sizes": [
1188
+ 4096,
1189
+ 4096,
1190
+ 4096
1191
+ ],
1192
+ "range_unroll_factors": [
1193
+ 0,
1194
+ 2,
1195
+ 0,
1196
+ 3
1197
+ ],
1198
+ "range_warp_specializes": [],
1199
+ "range_num_stages": [],
1200
+ "range_multi_buffers": [
1201
+ null,
1202
+ false,
1203
+ null,
1204
+ null
1205
+ ],
1206
+ "range_flattens": [
1207
+ null,
1208
+ true,
1209
+ true,
1210
+ false
1211
+ ],
1212
+ "load_eviction_policies": [
1213
+ "last",
1214
+ "last",
1215
+ "last",
1216
+ "last",
1217
+ "last",
1218
+ "last",
1219
+ "",
1220
+ "",
1221
+ ""
1222
+ ],
1223
+ "num_warps": 16,
1224
+ "num_stages": 2,
1225
+ "indexing": [
1226
+ "pointer",
1227
+ "pointer",
1228
+ "pointer",
1229
+ "tensor_descriptor",
1230
+ "pointer",
1231
+ "pointer",
1232
+ "pointer",
1233
+ "pointer",
1234
+ "tensor_descriptor",
1235
+ "tensor_descriptor",
1236
+ "tensor_descriptor",
1237
+ "pointer"
1238
+ ],
1239
+ "atomic_indexing": [],
1240
+ "pid_type": "flat"
1241
+ }
1242
+ },
1243
+ {
1244
+ "key": {
1245
+ "hidden_size": 8192,
1246
+ "num_tokens": 64
1247
+ },
1248
+ "config": {
1249
+ "block_sizes": [
1250
+ 8192,
1251
+ 8192,
1252
+ 8192
1253
+ ],
1254
+ "range_unroll_factors": [
1255
+ 0,
1256
+ 2,
1257
+ 2,
1258
+ 0
1259
+ ],
1260
+ "range_warp_specializes": [],
1261
+ "range_num_stages": [],
1262
+ "range_multi_buffers": [
1263
+ null,
1264
+ false,
1265
+ true,
1266
+ null
1267
+ ],
1268
+ "range_flattens": [
1269
+ null,
1270
+ false,
1271
+ true,
1272
+ null
1273
+ ],
1274
+ "load_eviction_policies": [
1275
+ "last",
1276
+ "",
1277
+ "last",
1278
+ "",
1279
+ "",
1280
+ "first",
1281
+ "",
1282
+ "first",
1283
+ ""
1284
+ ],
1285
+ "num_warps": 32,
1286
+ "num_stages": 6,
1287
+ "indexing": [
1288
+ "tensor_descriptor",
1289
+ "pointer",
1290
+ "pointer",
1291
+ "tensor_descriptor",
1292
+ "pointer",
1293
+ "tensor_descriptor",
1294
+ "pointer",
1295
+ "tensor_descriptor",
1296
+ "pointer",
1297
+ "pointer",
1298
+ "pointer",
1299
+ "pointer"
1300
+ ],
1301
+ "atomic_indexing": [],
1302
+ "pid_type": "flat"
1303
+ }
1304
+ },
1305
+ {
1306
+ "key": {
1307
+ "hidden_size": 2048,
1308
+ "num_tokens": 128
1309
+ },
1310
+ "config": {
1311
+ "block_sizes": [
1312
+ 2048,
1313
+ 2048,
1314
+ 2048
1315
+ ],
1316
+ "range_unroll_factors": [
1317
+ 0,
1318
+ 4,
1319
+ 3,
1320
+ 4
1321
+ ],
1322
+ "range_warp_specializes": [],
1323
+ "range_num_stages": [],
1324
+ "range_multi_buffers": [
1325
+ null,
1326
+ true,
1327
+ true,
1328
+ true
1329
+ ],
1330
+ "range_flattens": [
1331
+ null,
1332
+ false,
1333
+ false,
1334
+ null
1335
+ ],
1336
+ "load_eviction_policies": [
1337
+ "",
1338
+ "",
1339
+ "",
1340
+ "",
1341
+ "last",
1342
+ "",
1343
+ "",
1344
+ "",
1345
+ "first"
1346
+ ],
1347
+ "num_warps": 8,
1348
+ "num_stages": 1,
1349
+ "indexing": [
1350
+ "tensor_descriptor",
1351
+ "tensor_descriptor",
1352
+ "tensor_descriptor",
1353
+ "tensor_descriptor",
1354
+ "tensor_descriptor",
1355
+ "pointer",
1356
+ "tensor_descriptor",
1357
+ "tensor_descriptor",
1358
+ "pointer",
1359
+ "tensor_descriptor",
1360
+ "tensor_descriptor",
1361
+ "pointer"
1362
+ ],
1363
+ "atomic_indexing": [],
1364
+ "pid_type": "flat"
1365
+ }
1366
+ },
1367
+ {
1368
+ "key": {
1369
+ "hidden_size": 4096,
1370
+ "num_tokens": 128
1371
+ },
1372
+ "config": {
1373
+ "block_sizes": [
1374
+ 4096,
1375
+ 4096,
1376
+ 4096
1377
+ ],
1378
+ "range_unroll_factors": [
1379
+ 0,
1380
+ 1,
1381
+ 2,
1382
+ 1
1383
+ ],
1384
+ "range_warp_specializes": [],
1385
+ "range_num_stages": [],
1386
+ "range_multi_buffers": [
1387
+ null,
1388
+ null,
1389
+ null,
1390
+ null
1391
+ ],
1392
+ "range_flattens": [
1393
+ null,
1394
+ null,
1395
+ null,
1396
+ null
1397
+ ],
1398
+ "load_eviction_policies": [
1399
+ "first",
1400
+ "",
1401
+ "first",
1402
+ "",
1403
+ "last",
1404
+ "first",
1405
+ "",
1406
+ "",
1407
+ "first"
1408
+ ],
1409
+ "num_warps": 16,
1410
+ "num_stages": 7,
1411
+ "indexing": [
1412
+ "pointer",
1413
+ "pointer",
1414
+ "pointer",
1415
+ "pointer",
1416
+ "pointer",
1417
+ "pointer",
1418
+ "tensor_descriptor",
1419
+ "pointer",
1420
+ "tensor_descriptor",
1421
+ "pointer",
1422
+ "tensor_descriptor",
1423
+ "tensor_descriptor"
1424
+ ],
1425
+ "atomic_indexing": [],
1426
+ "pid_type": "flat"
1427
+ }
1428
+ },
1429
+ {
1430
+ "key": {
1431
+ "hidden_size": 8192,
1432
+ "num_tokens": 128
1433
+ },
1434
+ "config": {
1435
+ "block_sizes": [
1436
+ 8192,
1437
+ 8192,
1438
+ 4096
1439
+ ],
1440
+ "range_unroll_factors": [
1441
+ 0,
1442
+ 2,
1443
+ 3,
1444
+ 2
1445
+ ],
1446
+ "range_warp_specializes": [],
1447
+ "range_num_stages": [],
1448
+ "range_multi_buffers": [
1449
+ null,
1450
+ false,
1451
+ false,
1452
+ false
1453
+ ],
1454
+ "range_flattens": [
1455
+ null,
1456
+ true,
1457
+ false,
1458
+ true
1459
+ ],
1460
+ "load_eviction_policies": [
1461
+ "last",
1462
+ "",
1463
+ "last",
1464
+ "",
1465
+ "last",
1466
+ "last",
1467
+ "",
1468
+ "",
1469
+ "last"
1470
+ ],
1471
+ "num_warps": 16,
1472
+ "num_stages": 7,
1473
+ "indexing": [
1474
+ "tensor_descriptor",
1475
+ "tensor_descriptor",
1476
+ "tensor_descriptor",
1477
+ "tensor_descriptor",
1478
+ "pointer",
1479
+ "tensor_descriptor",
1480
+ "pointer",
1481
+ "tensor_descriptor",
1482
+ "pointer",
1483
+ "tensor_descriptor",
1484
+ "tensor_descriptor",
1485
+ "pointer"
1486
+ ],
1487
+ "atomic_indexing": [],
1488
+ "pid_type": "flat"
1489
+ }
1490
+ },
1491
+ {
1492
+ "key": {
1493
+ "hidden_size": 2048,
1494
+ "num_tokens": 256
1495
+ },
1496
+ "config": {
1497
+ "block_sizes": [
1498
+ 2048,
1499
+ 2048,
1500
+ 2048
1501
+ ],
1502
+ "range_unroll_factors": [
1503
+ 0,
1504
+ 1,
1505
+ 2,
1506
+ 4
1507
+ ],
1508
+ "range_warp_specializes": [],
1509
+ "range_num_stages": [],
1510
+ "range_multi_buffers": [
1511
+ null,
1512
+ null,
1513
+ true,
1514
+ true
1515
+ ],
1516
+ "range_flattens": [
1517
+ null,
1518
+ null,
1519
+ true,
1520
+ false
1521
+ ],
1522
+ "load_eviction_policies": [
1523
+ "",
1524
+ "",
1525
+ "first",
1526
+ "",
1527
+ "",
1528
+ "",
1529
+ "last",
1530
+ "",
1531
+ ""
1532
+ ],
1533
+ "num_warps": 8,
1534
+ "num_stages": 8,
1535
+ "indexing": [
1536
+ "tensor_descriptor",
1537
+ "pointer",
1538
+ "tensor_descriptor",
1539
+ "pointer",
1540
+ "pointer",
1541
+ "tensor_descriptor",
1542
+ "tensor_descriptor",
1543
+ "tensor_descriptor",
1544
+ "pointer",
1545
+ "pointer",
1546
+ "tensor_descriptor",
1547
+ "tensor_descriptor"
1548
+ ],
1549
+ "atomic_indexing": [],
1550
+ "pid_type": "flat"
1551
+ }
1552
+ },
1553
+ {
1554
+ "key": {
1555
+ "hidden_size": 4096,
1556
+ "num_tokens": 256
1557
+ },
1558
+ "config": {
1559
+ "block_sizes": [
1560
+ 4096,
1561
+ 4096,
1562
+ 4096
1563
+ ],
1564
+ "range_unroll_factors": [
1565
+ 0,
1566
+ 0,
1567
+ 1,
1568
+ 3
1569
+ ],
1570
+ "range_warp_specializes": [],
1571
+ "range_num_stages": [],
1572
+ "range_multi_buffers": [
1573
+ null,
1574
+ true,
1575
+ null,
1576
+ null
1577
+ ],
1578
+ "range_flattens": [
1579
+ null,
1580
+ null,
1581
+ true,
1582
+ true
1583
+ ],
1584
+ "load_eviction_policies": [
1585
+ "first",
1586
+ "",
1587
+ "first",
1588
+ "",
1589
+ "",
1590
+ "last",
1591
+ "last",
1592
+ "last",
1593
+ "last"
1594
+ ],
1595
+ "num_warps": 16,
1596
+ "num_stages": 1,
1597
+ "indexing": [
1598
+ "tensor_descriptor",
1599
+ "pointer",
1600
+ "pointer",
1601
+ "pointer",
1602
+ "tensor_descriptor",
1603
+ "tensor_descriptor",
1604
+ "pointer",
1605
+ "pointer",
1606
+ "tensor_descriptor",
1607
+ "pointer",
1608
+ "tensor_descriptor",
1609
+ "pointer"
1610
+ ],
1611
+ "atomic_indexing": [],
1612
+ "pid_type": "flat"
1613
+ }
1614
+ },
1615
+ {
1616
+ "key": {
1617
+ "hidden_size": 8192,
1618
+ "num_tokens": 256
1619
+ },
1620
+ "config": {
1621
+ "block_sizes": [
1622
+ 8192,
1623
+ 8192,
1624
+ 4096
1625
+ ],
1626
+ "range_unroll_factors": [
1627
+ 0,
1628
+ 1,
1629
+ 3,
1630
+ 2
1631
+ ],
1632
+ "range_warp_specializes": [],
1633
+ "range_num_stages": [],
1634
+ "range_multi_buffers": [
1635
+ null,
1636
+ null,
1637
+ true,
1638
+ null
1639
+ ],
1640
+ "range_flattens": [
1641
+ null,
1642
+ false,
1643
+ true,
1644
+ false
1645
+ ],
1646
+ "load_eviction_policies": [
1647
+ "",
1648
+ "",
1649
+ "",
1650
+ "",
1651
+ "",
1652
+ "last",
1653
+ "first",
1654
+ "",
1655
+ ""
1656
+ ],
1657
+ "num_warps": 16,
1658
+ "num_stages": 1,
1659
+ "indexing": [
1660
+ "tensor_descriptor",
1661
+ "tensor_descriptor",
1662
+ "pointer",
1663
+ "tensor_descriptor",
1664
+ "pointer",
1665
+ "tensor_descriptor",
1666
+ "pointer",
1667
+ "tensor_descriptor",
1668
+ "tensor_descriptor",
1669
+ "pointer",
1670
+ "pointer",
1671
+ "tensor_descriptor"
1672
+ ],
1673
+ "atomic_indexing": [],
1674
+ "pid_type": "flat"
1675
+ }
1676
+ },
1677
+ {
1678
+ "key": {
1679
+ "hidden_size": 2048,
1680
+ "num_tokens": 512
1681
+ },
1682
+ "config": {
1683
+ "block_sizes": [
1684
+ 2048,
1685
+ 2048,
1686
+ 2048
1687
+ ],
1688
+ "range_unroll_factors": [
1689
+ 0,
1690
+ 3,
1691
+ 3,
1692
+ 2
1693
+ ],
1694
+ "range_warp_specializes": [],
1695
+ "range_num_stages": [],
1696
+ "range_multi_buffers": [
1697
+ null,
1698
+ false,
1699
+ null,
1700
+ null
1701
+ ],
1702
+ "range_flattens": [
1703
+ null,
1704
+ false,
1705
+ false,
1706
+ false
1707
+ ],
1708
+ "load_eviction_policies": [
1709
+ "last",
1710
+ "",
1711
+ "last",
1712
+ "",
1713
+ "",
1714
+ "first",
1715
+ "first",
1716
+ "last",
1717
+ "last"
1718
+ ],
1719
+ "num_warps": 8,
1720
+ "num_stages": 6,
1721
+ "indexing": [
1722
+ "tensor_descriptor",
1723
+ "pointer",
1724
+ "pointer",
1725
+ "tensor_descriptor",
1726
+ "pointer",
1727
+ "tensor_descriptor",
1728
+ "pointer",
1729
+ "tensor_descriptor",
1730
+ "tensor_descriptor",
1731
+ "pointer",
1732
+ "pointer",
1733
+ "pointer"
1734
+ ],
1735
+ "atomic_indexing": [],
1736
+ "pid_type": "flat"
1737
+ }
1738
+ },
1739
+ {
1740
+ "key": {
1741
+ "hidden_size": 4096,
1742
+ "num_tokens": 512
1743
+ },
1744
+ "config": {
1745
+ "block_sizes": [
1746
+ 4096,
1747
+ 4096,
1748
+ 2048
1749
+ ],
1750
+ "range_unroll_factors": [
1751
+ 0,
1752
+ 3,
1753
+ 2,
1754
+ 2
1755
+ ],
1756
+ "range_warp_specializes": [],
1757
+ "range_num_stages": [],
1758
+ "range_multi_buffers": [
1759
+ null,
1760
+ true,
1761
+ null,
1762
+ true
1763
+ ],
1764
+ "range_flattens": [
1765
+ null,
1766
+ false,
1767
+ null,
1768
+ null
1769
+ ],
1770
+ "load_eviction_policies": [
1771
+ "",
1772
+ "",
1773
+ "",
1774
+ "",
1775
+ "",
1776
+ "last",
1777
+ "first",
1778
+ "",
1779
+ ""
1780
+ ],
1781
+ "num_warps": 8,
1782
+ "num_stages": 1,
1783
+ "indexing": [
1784
+ "pointer",
1785
+ "pointer",
1786
+ "tensor_descriptor",
1787
+ "tensor_descriptor",
1788
+ "tensor_descriptor",
1789
+ "pointer",
1790
+ "tensor_descriptor",
1791
+ "pointer",
1792
+ "pointer",
1793
+ "tensor_descriptor",
1794
+ "pointer",
1795
+ "tensor_descriptor"
1796
+ ],
1797
+ "atomic_indexing": [],
1798
+ "pid_type": "flat"
1799
+ }
1800
+ },
1801
+ {
1802
+ "key": {
1803
+ "hidden_size": 8192,
1804
+ "num_tokens": 512
1805
+ },
1806
+ "config": {
1807
+ "block_sizes": [
1808
+ 8192,
1809
+ 8192,
1810
+ 2048
1811
+ ],
1812
+ "range_unroll_factors": [
1813
+ 0,
1814
+ 4,
1815
+ 3,
1816
+ 2
1817
+ ],
1818
+ "range_warp_specializes": [],
1819
+ "range_num_stages": [],
1820
+ "range_multi_buffers": [
1821
+ null,
1822
+ false,
1823
+ null,
1824
+ false
1825
+ ],
1826
+ "range_flattens": [
1827
+ null,
1828
+ true,
1829
+ true,
1830
+ null
1831
+ ],
1832
+ "load_eviction_policies": [
1833
+ "",
1834
+ "",
1835
+ "",
1836
+ "",
1837
+ "last",
1838
+ "",
1839
+ "first",
1840
+ "first",
1841
+ ""
1842
+ ],
1843
+ "num_warps": 8,
1844
+ "num_stages": 1,
1845
+ "indexing": [
1846
+ "pointer",
1847
+ "pointer",
1848
+ "pointer",
1849
+ "pointer",
1850
+ "tensor_descriptor",
1851
+ "pointer",
1852
+ "tensor_descriptor",
1853
+ "pointer",
1854
+ "tensor_descriptor",
1855
+ "pointer",
1856
+ "pointer",
1857
+ "pointer"
1858
+ ],
1859
+ "atomic_indexing": [],
1860
+ "pid_type": "flat"
1861
+ }
1862
+ },
1863
+ {
1864
+ "key": {
1865
+ "hidden_size": 2048,
1866
+ "num_tokens": 1024
1867
+ },
1868
+ "config": {
1869
+ "block_sizes": [
1870
+ 2048,
1871
+ 2048,
1872
+ 1024
1873
+ ],
1874
+ "range_unroll_factors": [
1875
+ 0,
1876
+ 0,
1877
+ 4,
1878
+ 2
1879
+ ],
1880
+ "range_warp_specializes": [],
1881
+ "range_num_stages": [],
1882
+ "range_multi_buffers": [
1883
+ null,
1884
+ true,
1885
+ false,
1886
+ false
1887
+ ],
1888
+ "range_flattens": [
1889
+ null,
1890
+ null,
1891
+ true,
1892
+ null
1893
+ ],
1894
+ "load_eviction_policies": [
1895
+ "first",
1896
+ "",
1897
+ "first",
1898
+ "",
1899
+ "last",
1900
+ "last",
1901
+ "first",
1902
+ "first",
1903
+ "last"
1904
+ ],
1905
+ "num_warps": 4,
1906
+ "num_stages": 1,
1907
+ "indexing": [
1908
+ "tensor_descriptor",
1909
+ "tensor_descriptor",
1910
+ "pointer",
1911
+ "pointer",
1912
+ "tensor_descriptor",
1913
+ "pointer",
1914
+ "tensor_descriptor",
1915
+ "pointer",
1916
+ "tensor_descriptor",
1917
+ "pointer",
1918
+ "tensor_descriptor",
1919
+ "pointer"
1920
+ ],
1921
+ "atomic_indexing": [],
1922
+ "pid_type": "flat"
1923
+ }
1924
+ },
1925
+ {
1926
+ "key": {
1927
+ "hidden_size": 4096,
1928
+ "num_tokens": 1024
1929
+ },
1930
+ "config": {
1931
+ "block_sizes": [
1932
+ 4096,
1933
+ 4096,
1934
+ 2048
1935
+ ],
1936
+ "range_unroll_factors": [
1937
+ 0,
1938
+ 3,
1939
+ 0,
1940
+ 4
1941
+ ],
1942
+ "range_warp_specializes": [],
1943
+ "range_num_stages": [],
1944
+ "range_multi_buffers": [
1945
+ null,
1946
+ false,
1947
+ false,
1948
+ null
1949
+ ],
1950
+ "range_flattens": [
1951
+ null,
1952
+ true,
1953
+ null,
1954
+ false
1955
+ ],
1956
+ "load_eviction_policies": [
1957
+ "",
1958
+ "",
1959
+ "",
1960
+ "",
1961
+ "last",
1962
+ "",
1963
+ "first",
1964
+ "first",
1965
+ ""
1966
+ ],
1967
+ "num_warps": 8,
1968
+ "num_stages": 7,
1969
+ "indexing": [
1970
+ "tensor_descriptor",
1971
+ "pointer",
1972
+ "tensor_descriptor",
1973
+ "tensor_descriptor",
1974
+ "pointer",
1975
+ "tensor_descriptor",
1976
+ "pointer",
1977
+ "tensor_descriptor",
1978
+ "pointer",
1979
+ "tensor_descriptor",
1980
+ "tensor_descriptor",
1981
+ "pointer"
1982
+ ],
1983
+ "atomic_indexing": [],
1984
+ "pid_type": "flat"
1985
+ }
1986
+ },
1987
+ {
1988
+ "key": {
1989
+ "hidden_size": 8192,
1990
+ "num_tokens": 1024
1991
+ },
1992
+ "config": {
1993
+ "block_sizes": [
1994
+ 8192,
1995
+ 8192,
1996
+ 4096
1997
+ ],
1998
+ "range_unroll_factors": [
1999
+ 0,
2000
+ 1,
2001
+ 3,
2002
+ 2
2003
+ ],
2004
+ "range_warp_specializes": [],
2005
+ "range_num_stages": [],
2006
+ "range_multi_buffers": [
2007
+ null,
2008
+ null,
2009
+ true,
2010
+ null
2011
+ ],
2012
+ "range_flattens": [
2013
+ null,
2014
+ false,
2015
+ true,
2016
+ false
2017
+ ],
2018
+ "load_eviction_policies": [
2019
+ "",
2020
+ "",
2021
+ "",
2022
+ "",
2023
+ "",
2024
+ "last",
2025
+ "first",
2026
+ "",
2027
+ ""
2028
+ ],
2029
+ "num_warps": 16,
2030
+ "num_stages": 1,
2031
+ "indexing": [
2032
+ "tensor_descriptor",
2033
+ "tensor_descriptor",
2034
+ "pointer",
2035
+ "tensor_descriptor",
2036
+ "pointer",
2037
+ "tensor_descriptor",
2038
+ "pointer",
2039
+ "tensor_descriptor",
2040
+ "tensor_descriptor",
2041
+ "pointer",
2042
+ "pointer",
2043
+ "tensor_descriptor"
2044
+ ],
2045
+ "atomic_indexing": [],
2046
+ "pid_type": "flat"
2047
+ }
2048
+ },
2049
+ {
2050
+ "key": {
2051
+ "hidden_size": 2048,
2052
+ "num_tokens": 2048
2053
+ },
2054
+ "config": {
2055
+ "block_sizes": [
2056
+ 2048,
2057
+ 2048,
2058
+ 1024
2059
+ ],
2060
+ "range_unroll_factors": [
2061
+ 0,
2062
+ 3,
2063
+ 0,
2064
+ 0
2065
+ ],
2066
+ "range_warp_specializes": [],
2067
+ "range_num_stages": [],
2068
+ "range_multi_buffers": [
2069
+ null,
2070
+ false,
2071
+ null,
2072
+ null
2073
+ ],
2074
+ "range_flattens": [
2075
+ null,
2076
+ null,
2077
+ false,
2078
+ null
2079
+ ],
2080
+ "load_eviction_policies": [
2081
+ "",
2082
+ "",
2083
+ "first",
2084
+ "",
2085
+ "",
2086
+ "",
2087
+ "first",
2088
+ "",
2089
+ ""
2090
+ ],
2091
+ "num_warps": 4,
2092
+ "num_stages": 4,
2093
+ "indexing": [
2094
+ "tensor_descriptor",
2095
+ "pointer",
2096
+ "pointer",
2097
+ "tensor_descriptor",
2098
+ "pointer",
2099
+ "tensor_descriptor",
2100
+ "tensor_descriptor",
2101
+ "pointer",
2102
+ "tensor_descriptor",
2103
+ "tensor_descriptor",
2104
+ "tensor_descriptor",
2105
+ "tensor_descriptor"
2106
+ ],
2107
+ "atomic_indexing": [],
2108
+ "pid_type": "flat"
2109
+ }
2110
+ },
2111
+ {
2112
+ "key": {
2113
+ "hidden_size": 4096,
2114
+ "num_tokens": 2048
2115
+ },
2116
+ "config": {
2117
+ "block_sizes": [
2118
+ 4096,
2119
+ 4096,
2120
+ 2048
2121
+ ],
2122
+ "range_unroll_factors": [
2123
+ 0,
2124
+ 3,
2125
+ 2,
2126
+ 2
2127
+ ],
2128
+ "range_warp_specializes": [],
2129
+ "range_num_stages": [],
2130
+ "range_multi_buffers": [
2131
+ null,
2132
+ true,
2133
+ null,
2134
+ true
2135
+ ],
2136
+ "range_flattens": [
2137
+ null,
2138
+ false,
2139
+ null,
2140
+ null
2141
+ ],
2142
+ "load_eviction_policies": [
2143
+ "",
2144
+ "",
2145
+ "",
2146
+ "",
2147
+ "",
2148
+ "last",
2149
+ "first",
2150
+ "",
2151
+ ""
2152
+ ],
2153
+ "num_warps": 8,
2154
+ "num_stages": 1,
2155
+ "indexing": [
2156
+ "pointer",
2157
+ "pointer",
2158
+ "tensor_descriptor",
2159
+ "tensor_descriptor",
2160
+ "tensor_descriptor",
2161
+ "pointer",
2162
+ "tensor_descriptor",
2163
+ "pointer",
2164
+ "pointer",
2165
+ "tensor_descriptor",
2166
+ "pointer",
2167
+ "tensor_descriptor"
2168
+ ],
2169
+ "atomic_indexing": [],
2170
+ "pid_type": "flat"
2171
+ }
2172
+ },
2173
+ {
2174
+ "key": {
2175
+ "hidden_size": 8192,
2176
+ "num_tokens": 2048
2177
+ },
2178
+ "config": {
2179
+ "block_sizes": [
2180
+ 8192,
2181
+ 8192,
2182
+ 4096
2183
+ ],
2184
+ "range_unroll_factors": [
2185
+ 0,
2186
+ 1,
2187
+ 3,
2188
+ 2
2189
+ ],
2190
+ "range_warp_specializes": [],
2191
+ "range_num_stages": [],
2192
+ "range_multi_buffers": [
2193
+ null,
2194
+ null,
2195
+ true,
2196
+ null
2197
+ ],
2198
+ "range_flattens": [
2199
+ null,
2200
+ false,
2201
+ true,
2202
+ false
2203
+ ],
2204
+ "load_eviction_policies": [
2205
+ "",
2206
+ "",
2207
+ "",
2208
+ "",
2209
+ "",
2210
+ "last",
2211
+ "first",
2212
+ "",
2213
+ ""
2214
+ ],
2215
+ "num_warps": 16,
2216
+ "num_stages": 1,
2217
+ "indexing": [
2218
+ "tensor_descriptor",
2219
+ "tensor_descriptor",
2220
+ "pointer",
2221
+ "tensor_descriptor",
2222
+ "pointer",
2223
+ "tensor_descriptor",
2224
+ "pointer",
2225
+ "tensor_descriptor",
2226
+ "tensor_descriptor",
2227
+ "pointer",
2228
+ "pointer",
2229
+ "tensor_descriptor"
2230
+ ],
2231
+ "atomic_indexing": [],
2232
+ "pid_type": "flat"
2233
+ }
2234
+ },
2235
+ {
2236
+ "key": {
2237
+ "hidden_size": 2048,
2238
+ "num_tokens": 4096
2239
+ },
2240
+ "config": {
2241
+ "block_sizes": [
2242
+ 2048,
2243
+ 2048,
2244
+ 1024
2245
+ ],
2246
+ "range_unroll_factors": [
2247
+ 0,
2248
+ 4,
2249
+ 1,
2250
+ 2
2251
+ ],
2252
+ "range_warp_specializes": [],
2253
+ "range_num_stages": [],
2254
+ "range_multi_buffers": [
2255
+ null,
2256
+ false,
2257
+ null,
2258
+ null
2259
+ ],
2260
+ "range_flattens": [
2261
+ null,
2262
+ null,
2263
+ true,
2264
+ true
2265
+ ],
2266
+ "load_eviction_policies": [
2267
+ "",
2268
+ "",
2269
+ "",
2270
+ "last",
2271
+ "",
2272
+ "",
2273
+ "first",
2274
+ "first",
2275
+ "last"
2276
+ ],
2277
+ "num_warps": 4,
2278
+ "num_stages": 1,
2279
+ "indexing": [
2280
+ "pointer",
2281
+ "pointer",
2282
+ "pointer",
2283
+ "pointer",
2284
+ "pointer",
2285
+ "tensor_descriptor",
2286
+ "tensor_descriptor",
2287
+ "tensor_descriptor",
2288
+ "tensor_descriptor",
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": 4096
2301
+ },
2302
+ "config": {
2303
+ "block_sizes": [
2304
+ 4096,
2305
+ 4096,
2306
+ 4096
2307
+ ],
2308
+ "range_unroll_factors": [
2309
+ 0,
2310
+ 2,
2311
+ 4,
2312
+ 2
2313
+ ],
2314
+ "range_warp_specializes": [],
2315
+ "range_num_stages": [],
2316
+ "range_multi_buffers": [
2317
+ null,
2318
+ false,
2319
+ null,
2320
+ null
2321
+ ],
2322
+ "range_flattens": [
2323
+ null,
2324
+ false,
2325
+ null,
2326
+ null
2327
+ ],
2328
+ "load_eviction_policies": [
2329
+ "last",
2330
+ "",
2331
+ "",
2332
+ "last",
2333
+ "first",
2334
+ "",
2335
+ "first",
2336
+ "",
2337
+ "last"
2338
+ ],
2339
+ "num_warps": 16,
2340
+ "num_stages": 8,
2341
+ "indexing": [
2342
+ "pointer",
2343
+ "tensor_descriptor",
2344
+ "pointer",
2345
+ "pointer",
2346
+ "tensor_descriptor",
2347
+ "tensor_descriptor",
2348
+ "pointer",
2349
+ "pointer",
2350
+ "tensor_descriptor",
2351
+ "tensor_descriptor",
2352
+ "tensor_descriptor",
2353
+ "tensor_descriptor"
2354
+ ],
2355
+ "atomic_indexing": [],
2356
+ "pid_type": "flat"
2357
+ }
2358
+ },
2359
+ {
2360
+ "key": {
2361
+ "hidden_size": 8192,
2362
+ "num_tokens": 4096
2363
+ },
2364
+ "config": {
2365
+ "block_sizes": [
2366
+ 8192,
2367
+ 4096,
2368
+ 4096
2369
+ ],
2370
+ "range_unroll_factors": [
2371
+ 1,
2372
+ 1,
2373
+ 2,
2374
+ 1
2375
+ ],
2376
+ "range_warp_specializes": [],
2377
+ "range_multi_buffers": [
2378
+ null,
2379
+ false,
2380
+ false,
2381
+ true
2382
+ ],
2383
+ "range_flattens": [
2384
+ true,
2385
+ null,
2386
+ null,
2387
+ true
2388
+ ],
2389
+ "load_eviction_policies": [
2390
+ "",
2391
+ "last",
2392
+ "",
2393
+ "first",
2394
+ "first",
2395
+ "last",
2396
+ "first",
2397
+ "first",
2398
+ ""
2399
+ ],
2400
+ "num_warps": 16,
2401
+ "num_stages": 8,
2402
+ "indexing": [
2403
+ "pointer",
2404
+ "tensor_descriptor",
2405
+ "pointer",
2406
+ "pointer",
2407
+ "pointer",
2408
+ "tensor_descriptor",
2409
+ "tensor_descriptor",
2410
+ "pointer",
2411
+ "tensor_descriptor",
2412
+ "pointer",
2413
+ "tensor_descriptor",
2414
+ "tensor_descriptor"
2415
+ ],
2416
+ "atomic_indexing": [],
2417
+ "pid_type": "persistent_interleaved",
2418
+ "num_sm_multiplier": 32,
2419
+ "maxnreg": 128
2420
+ }
2421
+ },
2422
+ {
2423
+ "key": {
2424
+ "hidden_size": 2048,
2425
+ "num_tokens": 8192
2426
+ },
2427
+ "config": {
2428
+ "block_sizes": [
2429
+ 2048,
2430
+ 2048,
2431
+ 1024
2432
+ ],
2433
+ "range_unroll_factors": [
2434
+ 0,
2435
+ 4,
2436
+ 1,
2437
+ 2
2438
+ ],
2439
+ "range_warp_specializes": [],
2440
+ "range_num_stages": [],
2441
+ "range_multi_buffers": [
2442
+ null,
2443
+ false,
2444
+ null,
2445
+ null
2446
+ ],
2447
+ "range_flattens": [
2448
+ null,
2449
+ null,
2450
+ true,
2451
+ true
2452
+ ],
2453
+ "load_eviction_policies": [
2454
+ "",
2455
+ "",
2456
+ "",
2457
+ "last",
2458
+ "",
2459
+ "",
2460
+ "first",
2461
+ "first",
2462
+ "last"
2463
+ ],
2464
+ "num_warps": 4,
2465
+ "num_stages": 1,
2466
+ "indexing": [
2467
+ "pointer",
2468
+ "pointer",
2469
+ "pointer",
2470
+ "pointer",
2471
+ "pointer",
2472
+ "tensor_descriptor",
2473
+ "tensor_descriptor",
2474
+ "tensor_descriptor",
2475
+ "tensor_descriptor",
2476
+ "tensor_descriptor",
2477
+ "pointer",
2478
+ "pointer"
2479
+ ],
2480
+ "atomic_indexing": [],
2481
+ "pid_type": "flat"
2482
+ }
2483
+ },
2484
+ {
2485
+ "key": {
2486
+ "hidden_size": 4096,
2487
+ "num_tokens": 8192
2488
+ },
2489
+ "config": {
2490
+ "block_sizes": [
2491
+ 4096,
2492
+ 4096,
2493
+ 2048
2494
+ ],
2495
+ "range_unroll_factors": [
2496
+ 2,
2497
+ 2,
2498
+ 3,
2499
+ 2
2500
+ ],
2501
+ "range_warp_specializes": [],
2502
+ "range_multi_buffers": [
2503
+ true,
2504
+ false,
2505
+ false,
2506
+ null
2507
+ ],
2508
+ "range_flattens": [
2509
+ false,
2510
+ null,
2511
+ false,
2512
+ false
2513
+ ],
2514
+ "load_eviction_policies": [
2515
+ "last",
2516
+ "last",
2517
+ "",
2518
+ "last",
2519
+ "last",
2520
+ "",
2521
+ "",
2522
+ "",
2523
+ ""
2524
+ ],
2525
+ "num_warps": 8,
2526
+ "num_stages": 7,
2527
+ "indexing": [
2528
+ "pointer",
2529
+ "pointer",
2530
+ "pointer",
2531
+ "pointer",
2532
+ "tensor_descriptor",
2533
+ "pointer",
2534
+ "tensor_descriptor",
2535
+ "tensor_descriptor",
2536
+ "tensor_descriptor",
2537
+ "pointer",
2538
+ "tensor_descriptor",
2539
+ "tensor_descriptor"
2540
+ ],
2541
+ "atomic_indexing": [],
2542
+ "pid_type": "persistent_interleaved",
2543
+ "num_sm_multiplier": 64,
2544
+ "maxnreg": 256
2545
+ }
2546
+ },
2547
+ {
2548
+ "key": {
2549
+ "hidden_size": 8192,
2550
+ "num_tokens": 8192
2551
+ },
2552
+ "config": {
2553
+ "block_sizes": [
2554
+ 8192,
2555
+ 8192,
2556
+ 4096
2557
+ ],
2558
+ "range_unroll_factors": [
2559
+ 0,
2560
+ 3,
2561
+ 1,
2562
+ 2
2563
+ ],
2564
+ "range_warp_specializes": [],
2565
+ "range_num_stages": [],
2566
+ "range_multi_buffers": [
2567
+ null,
2568
+ false,
2569
+ true,
2570
+ null
2571
+ ],
2572
+ "range_flattens": [
2573
+ null,
2574
+ true,
2575
+ true,
2576
+ null
2577
+ ],
2578
+ "load_eviction_policies": [
2579
+ "last",
2580
+ "last",
2581
+ "",
2582
+ "last",
2583
+ "last",
2584
+ "",
2585
+ "last",
2586
+ "first",
2587
+ ""
2588
+ ],
2589
+ "num_warps": 16,
2590
+ "num_stages": 7,
2591
+ "indexing": [
2592
+ "tensor_descriptor",
2593
+ "pointer",
2594
+ "pointer",
2595
+ "tensor_descriptor",
2596
+ "tensor_descriptor",
2597
+ "pointer",
2598
+ "pointer",
2599
+ "pointer",
2600
+ "pointer",
2601
+ "tensor_descriptor",
2602
+ "tensor_descriptor",
2603
+ "pointer"
2604
+ ],
2605
+ "atomic_indexing": [],
2606
+ "pid_type": "flat"
2607
+ }
2608
+ },
2609
+ {
2610
+ "key": {
2611
+ "hidden_size": 5120,
2612
+ "num_tokens": 1
2613
+ },
2614
+ "config": {
2615
+ "block_sizes": [
2616
+ 8192,
2617
+ 8192,
2618
+ 2048
2619
+ ],
2620
+ "range_unroll_factors": [
2621
+ 0,
2622
+ 2,
2623
+ 3,
2624
+ 2
2625
+ ],
2626
+ "range_warp_specializes": [],
2627
+ "range_num_stages": [],
2628
+ "range_multi_buffers": [
2629
+ null,
2630
+ false,
2631
+ false,
2632
+ false
2633
+ ],
2634
+ "range_flattens": [
2635
+ null,
2636
+ false,
2637
+ true,
2638
+ false
2639
+ ],
2640
+ "load_eviction_policies": [
2641
+ "last",
2642
+ "",
2643
+ "last",
2644
+ "",
2645
+ "",
2646
+ "last",
2647
+ "last",
2648
+ "first",
2649
+ "first"
2650
+ ],
2651
+ "num_warps": 32,
2652
+ "num_stages": 6,
2653
+ "indexing": [
2654
+ "tensor_descriptor",
2655
+ "pointer",
2656
+ "pointer",
2657
+ "pointer",
2658
+ "tensor_descriptor",
2659
+ "pointer",
2660
+ "pointer",
2661
+ "tensor_descriptor",
2662
+ "pointer",
2663
+ "tensor_descriptor",
2664
+ "pointer",
2665
+ "pointer"
2666
+ ],
2667
+ "atomic_indexing": [],
2668
+ "pid_type": "flat"
2669
+ }
2670
+ },
2671
+ {
2672
+ "key": {
2673
+ "hidden_size": 5120,
2674
+ "num_tokens": 2
2675
+ },
2676
+ "config": {
2677
+ "block_sizes": [
2678
+ 8192,
2679
+ 8192,
2680
+ 2048
2681
+ ],
2682
+ "range_unroll_factors": [
2683
+ 0,
2684
+ 4,
2685
+ 2,
2686
+ 3
2687
+ ],
2688
+ "range_warp_specializes": [],
2689
+ "range_num_stages": [],
2690
+ "range_multi_buffers": [
2691
+ null,
2692
+ null,
2693
+ true,
2694
+ true
2695
+ ],
2696
+ "range_flattens": [
2697
+ null,
2698
+ null,
2699
+ null,
2700
+ true
2701
+ ],
2702
+ "load_eviction_policies": [
2703
+ "",
2704
+ "",
2705
+ "",
2706
+ "",
2707
+ "last",
2708
+ "last",
2709
+ "",
2710
+ "first",
2711
+ "last"
2712
+ ],
2713
+ "num_warps": 8,
2714
+ "num_stages": 1,
2715
+ "indexing": [
2716
+ "pointer",
2717
+ "tensor_descriptor",
2718
+ "tensor_descriptor",
2719
+ "pointer",
2720
+ "pointer",
2721
+ "tensor_descriptor",
2722
+ "tensor_descriptor",
2723
+ "pointer",
2724
+ "tensor_descriptor",
2725
+ "pointer",
2726
+ "pointer",
2727
+ "tensor_descriptor"
2728
+ ],
2729
+ "atomic_indexing": [],
2730
+ "pid_type": "flat"
2731
+ }
2732
+ },
2733
+ {
2734
+ "key": {
2735
+ "hidden_size": 5120,
2736
+ "num_tokens": 4
2737
+ },
2738
+ "config": {
2739
+ "block_sizes": [
2740
+ 8192,
2741
+ 8192,
2742
+ 2048
2743
+ ],
2744
+ "range_unroll_factors": [
2745
+ 0,
2746
+ 4,
2747
+ 2,
2748
+ 3
2749
+ ],
2750
+ "range_warp_specializes": [],
2751
+ "range_num_stages": [],
2752
+ "range_multi_buffers": [
2753
+ null,
2754
+ null,
2755
+ true,
2756
+ true
2757
+ ],
2758
+ "range_flattens": [
2759
+ null,
2760
+ null,
2761
+ null,
2762
+ true
2763
+ ],
2764
+ "load_eviction_policies": [
2765
+ "",
2766
+ "",
2767
+ "",
2768
+ "",
2769
+ "last",
2770
+ "last",
2771
+ "",
2772
+ "first",
2773
+ "last"
2774
+ ],
2775
+ "num_warps": 8,
2776
+ "num_stages": 1,
2777
+ "indexing": [
2778
+ "pointer",
2779
+ "tensor_descriptor",
2780
+ "tensor_descriptor",
2781
+ "pointer",
2782
+ "pointer",
2783
+ "tensor_descriptor",
2784
+ "tensor_descriptor",
2785
+ "pointer",
2786
+ "tensor_descriptor",
2787
+ "pointer",
2788
+ "pointer",
2789
+ "tensor_descriptor"
2790
+ ],
2791
+ "atomic_indexing": [],
2792
+ "pid_type": "flat"
2793
+ }
2794
+ },
2795
+ {
2796
+ "key": {
2797
+ "hidden_size": 5120,
2798
+ "num_tokens": 8
2799
+ },
2800
+ "config": {
2801
+ "block_sizes": [
2802
+ 8192,
2803
+ 8192,
2804
+ 2048
2805
+ ],
2806
+ "range_unroll_factors": [
2807
+ 0,
2808
+ 4,
2809
+ 2,
2810
+ 3
2811
+ ],
2812
+ "range_warp_specializes": [],
2813
+ "range_num_stages": [],
2814
+ "range_multi_buffers": [
2815
+ null,
2816
+ null,
2817
+ true,
2818
+ true
2819
+ ],
2820
+ "range_flattens": [
2821
+ null,
2822
+ null,
2823
+ null,
2824
+ true
2825
+ ],
2826
+ "load_eviction_policies": [
2827
+ "",
2828
+ "",
2829
+ "",
2830
+ "",
2831
+ "last",
2832
+ "last",
2833
+ "",
2834
+ "first",
2835
+ "last"
2836
+ ],
2837
+ "num_warps": 8,
2838
+ "num_stages": 1,
2839
+ "indexing": [
2840
+ "pointer",
2841
+ "tensor_descriptor",
2842
+ "tensor_descriptor",
2843
+ "pointer",
2844
+ "pointer",
2845
+ "tensor_descriptor",
2846
+ "tensor_descriptor",
2847
+ "pointer",
2848
+ "tensor_descriptor",
2849
+ "pointer",
2850
+ "pointer",
2851
+ "tensor_descriptor"
2852
+ ],
2853
+ "atomic_indexing": [],
2854
+ "pid_type": "flat"
2855
+ }
2856
+ },
2857
+ {
2858
+ "key": {
2859
+ "hidden_size": 5120,
2860
+ "num_tokens": 16
2861
+ },
2862
+ "config": {
2863
+ "block_sizes": [
2864
+ 8192,
2865
+ 8192,
2866
+ 2048
2867
+ ],
2868
+ "range_unroll_factors": [
2869
+ 0,
2870
+ 3,
2871
+ 2,
2872
+ 2
2873
+ ],
2874
+ "range_warp_specializes": [],
2875
+ "range_num_stages": [],
2876
+ "range_multi_buffers": [
2877
+ null,
2878
+ true,
2879
+ false,
2880
+ false
2881
+ ],
2882
+ "range_flattens": [
2883
+ null,
2884
+ true,
2885
+ true,
2886
+ null
2887
+ ],
2888
+ "load_eviction_policies": [
2889
+ "",
2890
+ "",
2891
+ "",
2892
+ "",
2893
+ "",
2894
+ "",
2895
+ "first",
2896
+ "first",
2897
+ "first"
2898
+ ],
2899
+ "num_warps": 8,
2900
+ "num_stages": 1,
2901
+ "indexing": [
2902
+ "pointer",
2903
+ "tensor_descriptor",
2904
+ "tensor_descriptor",
2905
+ "tensor_descriptor",
2906
+ "tensor_descriptor",
2907
+ "pointer",
2908
+ "tensor_descriptor",
2909
+ "tensor_descriptor",
2910
+ "tensor_descriptor",
2911
+ "tensor_descriptor",
2912
+ "tensor_descriptor",
2913
+ "pointer"
2914
+ ],
2915
+ "atomic_indexing": [],
2916
+ "pid_type": "flat"
2917
+ }
2918
+ },
2919
+ {
2920
+ "key": {
2921
+ "hidden_size": 5120,
2922
+ "num_tokens": 32
2923
+ },
2924
+ "config": {
2925
+ "block_sizes": [
2926
+ 8192,
2927
+ 8192,
2928
+ 2048
2929
+ ],
2930
+ "range_unroll_factors": [
2931
+ 0,
2932
+ 4,
2933
+ 3,
2934
+ 2
2935
+ ],
2936
+ "range_warp_specializes": [],
2937
+ "range_num_stages": [],
2938
+ "range_multi_buffers": [
2939
+ null,
2940
+ false,
2941
+ null,
2942
+ false
2943
+ ],
2944
+ "range_flattens": [
2945
+ null,
2946
+ true,
2947
+ true,
2948
+ null
2949
+ ],
2950
+ "load_eviction_policies": [
2951
+ "",
2952
+ "",
2953
+ "",
2954
+ "",
2955
+ "last",
2956
+ "",
2957
+ "first",
2958
+ "first",
2959
+ ""
2960
+ ],
2961
+ "num_warps": 8,
2962
+ "num_stages": 1,
2963
+ "indexing": [
2964
+ "pointer",
2965
+ "pointer",
2966
+ "pointer",
2967
+ "pointer",
2968
+ "tensor_descriptor",
2969
+ "pointer",
2970
+ "tensor_descriptor",
2971
+ "pointer",
2972
+ "tensor_descriptor",
2973
+ "pointer",
2974
+ "pointer",
2975
+ "pointer"
2976
+ ],
2977
+ "atomic_indexing": [],
2978
+ "pid_type": "flat"
2979
+ }
2980
+ },
2981
+ {
2982
+ "key": {
2983
+ "hidden_size": 5120,
2984
+ "num_tokens": 64
2985
+ },
2986
+ "config": {
2987
+ "block_sizes": [
2988
+ 8192,
2989
+ 8192,
2990
+ 2048
2991
+ ],
2992
+ "range_unroll_factors": [
2993
+ 0,
2994
+ 4,
2995
+ 2,
2996
+ 3
2997
+ ],
2998
+ "range_warp_specializes": [],
2999
+ "range_num_stages": [],
3000
+ "range_multi_buffers": [
3001
+ null,
3002
+ null,
3003
+ true,
3004
+ true
3005
+ ],
3006
+ "range_flattens": [
3007
+ null,
3008
+ null,
3009
+ null,
3010
+ true
3011
+ ],
3012
+ "load_eviction_policies": [
3013
+ "",
3014
+ "",
3015
+ "",
3016
+ "",
3017
+ "last",
3018
+ "last",
3019
+ "",
3020
+ "first",
3021
+ "last"
3022
+ ],
3023
+ "num_warps": 8,
3024
+ "num_stages": 1,
3025
+ "indexing": [
3026
+ "pointer",
3027
+ "tensor_descriptor",
3028
+ "tensor_descriptor",
3029
+ "pointer",
3030
+ "pointer",
3031
+ "tensor_descriptor",
3032
+ "tensor_descriptor",
3033
+ "pointer",
3034
+ "tensor_descriptor",
3035
+ "pointer",
3036
+ "pointer",
3037
+ "tensor_descriptor"
3038
+ ],
3039
+ "atomic_indexing": [],
3040
+ "pid_type": "flat"
3041
+ }
3042
+ },
3043
+ {
3044
+ "key": {
3045
+ "hidden_size": 5120,
3046
+ "num_tokens": 128
3047
+ },
3048
+ "config": {
3049
+ "block_sizes": [
3050
+ 8192,
3051
+ 8192,
3052
+ 2048
3053
+ ],
3054
+ "range_unroll_factors": [
3055
+ 0,
3056
+ 4,
3057
+ 2,
3058
+ 3
3059
+ ],
3060
+ "range_warp_specializes": [],
3061
+ "range_num_stages": [],
3062
+ "range_multi_buffers": [
3063
+ null,
3064
+ null,
3065
+ true,
3066
+ true
3067
+ ],
3068
+ "range_flattens": [
3069
+ null,
3070
+ null,
3071
+ null,
3072
+ true
3073
+ ],
3074
+ "load_eviction_policies": [
3075
+ "",
3076
+ "",
3077
+ "",
3078
+ "",
3079
+ "last",
3080
+ "last",
3081
+ "",
3082
+ "first",
3083
+ "last"
3084
+ ],
3085
+ "num_warps": 8,
3086
+ "num_stages": 1,
3087
+ "indexing": [
3088
+ "pointer",
3089
+ "tensor_descriptor",
3090
+ "tensor_descriptor",
3091
+ "pointer",
3092
+ "pointer",
3093
+ "tensor_descriptor",
3094
+ "tensor_descriptor",
3095
+ "pointer",
3096
+ "tensor_descriptor",
3097
+ "pointer",
3098
+ "pointer",
3099
+ "tensor_descriptor"
3100
+ ],
3101
+ "atomic_indexing": [],
3102
+ "pid_type": "flat"
3103
+ }
3104
+ },
3105
+ {
3106
+ "key": {
3107
+ "hidden_size": 5120,
3108
+ "num_tokens": 256
3109
+ },
3110
+ "config": {
3111
+ "block_sizes": [
3112
+ 8192,
3113
+ 8192,
3114
+ 2048
3115
+ ],
3116
+ "range_unroll_factors": [
3117
+ 0,
3118
+ 3,
3119
+ 2,
3120
+ 2
3121
+ ],
3122
+ "range_warp_specializes": [],
3123
+ "range_num_stages": [],
3124
+ "range_multi_buffers": [
3125
+ null,
3126
+ true,
3127
+ false,
3128
+ false
3129
+ ],
3130
+ "range_flattens": [
3131
+ null,
3132
+ true,
3133
+ true,
3134
+ null
3135
+ ],
3136
+ "load_eviction_policies": [
3137
+ "",
3138
+ "",
3139
+ "",
3140
+ "",
3141
+ "",
3142
+ "",
3143
+ "first",
3144
+ "first",
3145
+ "first"
3146
+ ],
3147
+ "num_warps": 8,
3148
+ "num_stages": 1,
3149
+ "indexing": [
3150
+ "pointer",
3151
+ "tensor_descriptor",
3152
+ "tensor_descriptor",
3153
+ "tensor_descriptor",
3154
+ "tensor_descriptor",
3155
+ "pointer",
3156
+ "tensor_descriptor",
3157
+ "tensor_descriptor",
3158
+ "tensor_descriptor",
3159
+ "tensor_descriptor",
3160
+ "tensor_descriptor",
3161
+ "pointer"
3162
+ ],
3163
+ "atomic_indexing": [],
3164
+ "pid_type": "flat"
3165
+ }
3166
+ },
3167
+ {
3168
+ "key": {
3169
+ "hidden_size": 5120,
3170
+ "num_tokens": 512
3171
+ },
3172
+ "config": {
3173
+ "block_sizes": [
3174
+ 8192,
3175
+ 8192,
3176
+ 2048
3177
+ ],
3178
+ "range_unroll_factors": [
3179
+ 0,
3180
+ 3,
3181
+ 2,
3182
+ 2
3183
+ ],
3184
+ "range_warp_specializes": [],
3185
+ "range_num_stages": [],
3186
+ "range_multi_buffers": [
3187
+ null,
3188
+ true,
3189
+ false,
3190
+ false
3191
+ ],
3192
+ "range_flattens": [
3193
+ null,
3194
+ true,
3195
+ true,
3196
+ null
3197
+ ],
3198
+ "load_eviction_policies": [
3199
+ "",
3200
+ "",
3201
+ "",
3202
+ "",
3203
+ "",
3204
+ "",
3205
+ "first",
3206
+ "first",
3207
+ "first"
3208
+ ],
3209
+ "num_warps": 8,
3210
+ "num_stages": 1,
3211
+ "indexing": [
3212
+ "pointer",
3213
+ "tensor_descriptor",
3214
+ "tensor_descriptor",
3215
+ "tensor_descriptor",
3216
+ "tensor_descriptor",
3217
+ "pointer",
3218
+ "tensor_descriptor",
3219
+ "tensor_descriptor",
3220
+ "tensor_descriptor",
3221
+ "tensor_descriptor",
3222
+ "tensor_descriptor",
3223
+ "pointer"
3224
+ ],
3225
+ "atomic_indexing": [],
3226
+ "pid_type": "flat"
3227
+ }
3228
+ },
3229
+ {
3230
+ "key": {
3231
+ "hidden_size": 5120,
3232
+ "num_tokens": 1024
3233
+ },
3234
+ "config": {
3235
+ "block_sizes": [
3236
+ 8192,
3237
+ 8192,
3238
+ 1024
3239
+ ],
3240
+ "range_unroll_factors": [
3241
+ 0,
3242
+ 3,
3243
+ 2,
3244
+ 4
3245
+ ],
3246
+ "range_warp_specializes": [],
3247
+ "range_num_stages": [],
3248
+ "range_multi_buffers": [
3249
+ null,
3250
+ true,
3251
+ null,
3252
+ null
3253
+ ],
3254
+ "range_flattens": [
3255
+ null,
3256
+ null,
3257
+ true,
3258
+ null
3259
+ ],
3260
+ "load_eviction_policies": [
3261
+ "first",
3262
+ "",
3263
+ "last",
3264
+ "last",
3265
+ "",
3266
+ "last",
3267
+ "first",
3268
+ "first",
3269
+ "last"
3270
+ ],
3271
+ "num_warps": 4,
3272
+ "num_stages": 1,
3273
+ "indexing": [
3274
+ "pointer",
3275
+ "tensor_descriptor",
3276
+ "pointer",
3277
+ "tensor_descriptor",
3278
+ "pointer",
3279
+ "tensor_descriptor",
3280
+ "pointer",
3281
+ "pointer",
3282
+ "tensor_descriptor",
3283
+ "tensor_descriptor",
3284
+ "tensor_descriptor",
3285
+ "tensor_descriptor"
3286
+ ],
3287
+ "atomic_indexing": [],
3288
+ "pid_type": "flat"
3289
+ }
3290
+ },
3291
+ {
3292
+ "key": {
3293
+ "hidden_size": 5120,
3294
+ "num_tokens": 2048
3295
+ },
3296
+ "config": {
3297
+ "block_sizes": [
3298
+ 8192,
3299
+ 8192,
3300
+ 2048
3301
+ ],
3302
+ "range_unroll_factors": [
3303
+ 0,
3304
+ 4,
3305
+ 3,
3306
+ 2
3307
+ ],
3308
+ "range_warp_specializes": [],
3309
+ "range_num_stages": [],
3310
+ "range_multi_buffers": [
3311
+ null,
3312
+ false,
3313
+ null,
3314
+ false
3315
+ ],
3316
+ "range_flattens": [
3317
+ null,
3318
+ true,
3319
+ true,
3320
+ null
3321
+ ],
3322
+ "load_eviction_policies": [
3323
+ "",
3324
+ "",
3325
+ "",
3326
+ "",
3327
+ "last",
3328
+ "",
3329
+ "first",
3330
+ "first",
3331
+ ""
3332
+ ],
3333
+ "num_warps": 8,
3334
+ "num_stages": 1,
3335
+ "indexing": [
3336
+ "pointer",
3337
+ "pointer",
3338
+ "pointer",
3339
+ "pointer",
3340
+ "tensor_descriptor",
3341
+ "pointer",
3342
+ "tensor_descriptor",
3343
+ "pointer",
3344
+ "tensor_descriptor",
3345
+ "pointer",
3346
+ "pointer",
3347
+ "pointer"
3348
+ ],
3349
+ "atomic_indexing": [],
3350
+ "pid_type": "flat"
3351
+ }
3352
+ },
3353
+ {
3354
+ "key": {
3355
+ "hidden_size": 5120,
3356
+ "num_tokens": 4096
3357
+ },
3358
+ "config": {
3359
+ "block_sizes": [
3360
+ 8192,
3361
+ 8192,
3362
+ 2048
3363
+ ],
3364
+ "range_unroll_factors": [
3365
+ 0,
3366
+ 4,
3367
+ 3,
3368
+ 2
3369
+ ],
3370
+ "range_warp_specializes": [],
3371
+ "range_num_stages": [],
3372
+ "range_multi_buffers": [
3373
+ null,
3374
+ false,
3375
+ null,
3376
+ false
3377
+ ],
3378
+ "range_flattens": [
3379
+ null,
3380
+ true,
3381
+ true,
3382
+ null
3383
+ ],
3384
+ "load_eviction_policies": [
3385
+ "",
3386
+ "",
3387
+ "",
3388
+ "",
3389
+ "last",
3390
+ "",
3391
+ "first",
3392
+ "first",
3393
+ ""
3394
+ ],
3395
+ "num_warps": 8,
3396
+ "num_stages": 1,
3397
+ "indexing": [
3398
+ "pointer",
3399
+ "pointer",
3400
+ "pointer",
3401
+ "pointer",
3402
+ "tensor_descriptor",
3403
+ "pointer",
3404
+ "tensor_descriptor",
3405
+ "pointer",
3406
+ "tensor_descriptor",
3407
+ "pointer",
3408
+ "pointer",
3409
+ "pointer"
3410
+ ],
3411
+ "atomic_indexing": [],
3412
+ "pid_type": "flat"
3413
+ }
3414
+ },
3415
+ {
3416
+ "key": {
3417
+ "hidden_size": 5120,
3418
+ "num_tokens": 8192
3419
+ },
3420
+ "config": {
3421
+ "block_sizes": [
3422
+ 8192,
3423
+ 8192,
3424
+ 2048
3425
+ ],
3426
+ "range_unroll_factors": [
3427
+ 0,
3428
+ 4,
3429
+ 3,
3430
+ 2
3431
+ ],
3432
+ "range_warp_specializes": [],
3433
+ "range_num_stages": [],
3434
+ "range_multi_buffers": [
3435
+ null,
3436
+ false,
3437
+ null,
3438
+ false
3439
+ ],
3440
+ "range_flattens": [
3441
+ null,
3442
+ true,
3443
+ true,
3444
+ null
3445
+ ],
3446
+ "load_eviction_policies": [
3447
+ "",
3448
+ "",
3449
+ "",
3450
+ "",
3451
+ "last",
3452
+ "",
3453
+ "first",
3454
+ "first",
3455
+ ""
3456
+ ],
3457
+ "num_warps": 8,
3458
+ "num_stages": 1,
3459
+ "indexing": [
3460
+ "pointer",
3461
+ "pointer",
3462
+ "pointer",
3463
+ "pointer",
3464
+ "tensor_descriptor",
3465
+ "pointer",
3466
+ "tensor_descriptor",
3467
+ "pointer",
3468
+ "tensor_descriptor",
3469
+ "pointer",
3470
+ "pointer",
3471
+ "pointer"
3472
+ ],
3473
+ "atomic_indexing": [],
3474
+ "pid_type": "flat"
3475
+ }
3476
+ },
3477
+ {
3478
+ "key": {
3479
+ "hidden_size": 5120,
3480
+ "num_tokens": 16384
3481
+ },
3482
+ "config": {
3483
+ "block_sizes": [
3484
+ 8192,
3485
+ 8192,
3486
+ 2048
3487
+ ],
3488
+ "range_unroll_factors": [
3489
+ 0,
3490
+ 0,
3491
+ 2,
3492
+ 4
3493
+ ],
3494
+ "range_warp_specializes": [],
3495
+ "range_num_stages": [],
3496
+ "range_multi_buffers": [
3497
+ null,
3498
+ true,
3499
+ null,
3500
+ true
3501
+ ],
3502
+ "range_flattens": [
3503
+ null,
3504
+ true,
3505
+ null,
3506
+ true
3507
+ ],
3508
+ "load_eviction_policies": [
3509
+ "last",
3510
+ "",
3511
+ "last",
3512
+ "last",
3513
+ "last",
3514
+ "last",
3515
+ "first",
3516
+ "last",
3517
+ ""
3518
+ ],
3519
+ "num_warps": 8,
3520
+ "num_stages": 3,
3521
+ "indexing": [
3522
+ "pointer",
3523
+ "pointer",
3524
+ "pointer",
3525
+ "pointer",
3526
+ "pointer",
3527
+ "pointer",
3528
+ "tensor_descriptor",
3529
+ "tensor_descriptor",
3530
+ "tensor_descriptor",
3531
+ "tensor_descriptor",
3532
+ "pointer",
3533
+ "pointer"
3534
+ ],
3535
+ "atomic_indexing": [],
3536
+ "pid_type": "flat"
3537
+ }
3538
+ },
3539
+ {
3540
+ "key": {
3541
+ "hidden_size": 2048,
3542
+ "num_tokens": 16384
3543
+ },
3544
+ "config": {
3545
+ "block_sizes": [
3546
+ 2048,
3547
+ 2048,
3548
+ 2048
3549
+ ],
3550
+ "range_unroll_factors": [
3551
+ 0,
3552
+ 2,
3553
+ 4,
3554
+ 4
3555
+ ],
3556
+ "range_warp_specializes": [],
3557
+ "range_num_stages": [],
3558
+ "range_multi_buffers": [
3559
+ null,
3560
+ false,
3561
+ null,
3562
+ false
3563
+ ],
3564
+ "range_flattens": [
3565
+ null,
3566
+ true,
3567
+ null,
3568
+ true
3569
+ ],
3570
+ "load_eviction_policies": [
3571
+ "",
3572
+ "",
3573
+ "",
3574
+ "",
3575
+ "last",
3576
+ "",
3577
+ "first",
3578
+ "",
3579
+ "first"
3580
+ ],
3581
+ "num_warps": 8,
3582
+ "num_stages": 6,
3583
+ "indexing": [
3584
+ "pointer",
3585
+ "tensor_descriptor",
3586
+ "pointer",
3587
+ "pointer",
3588
+ "pointer",
3589
+ "tensor_descriptor",
3590
+ "pointer",
3591
+ "tensor_descriptor",
3592
+ "pointer",
3593
+ "tensor_descriptor",
3594
+ "tensor_descriptor",
3595
+ "tensor_descriptor"
3596
+ ],
3597
+ "atomic_indexing": [],
3598
+ "pid_type": "flat"
3599
+ }
3600
+ },
3601
+ {
3602
+ "key": {
3603
+ "hidden_size": 4096,
3604
+ "num_tokens": 16384
3605
+ },
3606
+ "config": {
3607
+ "block_sizes": [
3608
+ 4096,
3609
+ 4096,
3610
+ 4096
3611
+ ],
3612
+ "range_unroll_factors": [
3613
+ 0,
3614
+ 4,
3615
+ 4,
3616
+ 2
3617
+ ],
3618
+ "range_warp_specializes": [],
3619
+ "range_num_stages": [],
3620
+ "range_multi_buffers": [
3621
+ null,
3622
+ null,
3623
+ true,
3624
+ true
3625
+ ],
3626
+ "range_flattens": [
3627
+ null,
3628
+ true,
3629
+ null,
3630
+ false
3631
+ ],
3632
+ "load_eviction_policies": [
3633
+ "last",
3634
+ "last",
3635
+ "last",
3636
+ "",
3637
+ "last",
3638
+ "first",
3639
+ "first",
3640
+ "",
3641
+ "last"
3642
+ ],
3643
+ "num_warps": 16,
3644
+ "num_stages": 1,
3645
+ "indexing": [
3646
+ "tensor_descriptor",
3647
+ "tensor_descriptor",
3648
+ "pointer",
3649
+ "tensor_descriptor",
3650
+ "pointer",
3651
+ "pointer",
3652
+ "pointer",
3653
+ "pointer",
3654
+ "pointer",
3655
+ "pointer",
3656
+ "tensor_descriptor",
3657
+ "pointer"
3658
+ ],
3659
+ "atomic_indexing": [],
3660
+ "pid_type": "flat"
3661
+ }
3662
+ }
3663
+]
\ No newline at end of file
vllm/kernels/helion/ops/rms_norm_dynamic_per_token_quant.py
new
+233
@@ -0,0 +1,233 @@
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.utils import (
11
+ get_fp8_dtype,
12
+ get_int8_min_max,
13
+ get_int8_min_scaling_factor,
14
+)
15
+from vllm.logger import init_logger
16
+from vllm.model_executor.layers.quantization.utils.quant_utils import (
17
+ get_fp8_min_max,
18
+)
19
+from vllm.platforms import current_platform
20
+from vllm.utils.import_utils import has_helion
21
+
22
+if not has_helion():
23
+ raise ImportError(
24
+ "Helion kernel requires helion to be installed. "
25
+ "Install it with: pip install helion"
26
+ )
27
+
28
+import helion
29
+import helion.language as hl
30
+
31
+from vllm.kernels.helion.register import register_kernel
32
+
33
+logger = init_logger(__name__)
34
+
35
+
36
+def generate_inputs() -> dict[CaseKey, tuple[Any, ...]]:
37
+ # TODO(xiaohongchen1991): it is difficult for kernel author to cover all
38
+ # input property combination. Currently, dtypes are fixed. We need
39
+ # optimization to bucket/skip some combinations
40
+ num_tokens_list = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192]
41
+ hidden_size_list = [2048, 4096, 5120]
42
+ in_dtype: torch.dtype = torch.bfloat16
43
+ out_dtype: torch.dtype = current_platform.fp8_dtype()
44
+ scale_dtype: torch.dtype = torch.float32
45
+ inputs = {}
46
+
47
+ for num_tokens, hidden_size in product(num_tokens_list, hidden_size_list):
48
+ input = torch.randn(num_tokens, hidden_size, device="cuda", dtype=in_dtype)
49
+ result = torch.empty(input.shape, device=input.device, dtype=out_dtype)
50
+ scale = torch.empty((num_tokens, 1), device=input.device, dtype=scale_dtype)
51
+ scale_ub = torch.mean(input).to(scale_dtype)
52
+ residual = torch.randn_like(input)
53
+ weight = torch.normal(
54
+ mean=1.0,
55
+ std=1.0,
56
+ size=(hidden_size,),
57
+ dtype=input.dtype,
58
+ device=input.device,
59
+ )
60
+ epsilon = 1e-6
61
+
62
+ config_key = CaseKey({"hidden_size": hidden_size, "num_tokens": num_tokens})
63
+ inputs[config_key] = (result, input, weight, scale, epsilon, scale_ub, residual)
64
+
65
+ return inputs
66
+
67
+
68
+_pick_cache: dict[tuple[int, int], CaseKey | None] = {}
69
+
70
+
71
+def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey | None:
72
+ """Pick the best pre-tuned config for the given input shape.
73
+
74
+ Selection strategy:
75
+ 1. Find the closest hidden_size among available configs
76
+ (exact match preferred).
77
+ 2. Among the num_tokens values tuned for that hidden_size, pick
78
+ the smallest num_tokens >= the input's num_tokens. If the input is
79
+ larger than all available num_tokens, fall back to the largest.
80
+ """
81
+
82
+ if not config_keys:
83
+ return None
84
+
85
+ _, input, *_ = args
86
+ num_tokens, hidden_size = input.shape
87
+
88
+ cache_key = (num_tokens, hidden_size)
89
+ cached = _pick_cache.get(cache_key)
90
+ if cached is not None:
91
+ return cached
92
+
93
+ configs: dict[int, list[int]] = {}
94
+ for key in config_keys:
95
+ if key.is_default():
96
+ continue
97
+ configs.setdefault(key["hidden_size"], []).append(key["num_tokens"])
98
+
99
+ if not configs:
100
+ return None
101
+
102
+ best_hidden_size = min(configs, key=lambda s: abs(s - hidden_size))
103
+ available_num_tokens = sorted(configs[best_hidden_size])
104
+ best_num_tokens = next(
105
+ (n for n in available_num_tokens if n >= num_tokens), available_num_tokens[-1]
106
+ )
107
+
108
+ result = CaseKey({"hidden_size": best_hidden_size, "num_tokens": best_num_tokens})
109
+ _pick_cache[cache_key] = result
110
+ return result
111
+
112
+
113
+def fake_impl(
114
+ result: torch.Tensor, # [num_tokens, hidden_size]
115
+ input: torch.Tensor, # [num_tokens, hidden_size]
116
+ weight: torch.Tensor, # [hidden_size]
117
+ scale: torch.Tensor, # [num_tokens, 1]
118
+ epsilon: float,
119
+ scale_ub: torch.Tensor | None = None, # []
120
+ residual: torch.Tensor | None = None, # [num_tokens, hidden_size]
121
+) -> None:
122
+ return
123
+
124
+
125
+def baseline(
126
+ result: torch.Tensor, # [num_tokens, hidden_size]
127
+ input: torch.Tensor, # [num_tokens, hidden_size]
128
+ weight: torch.Tensor, # [num_tokens]
129
+ scale: torch.Tensor, # [num_tokens, 1]
130
+ epsilon: float,
131
+ scale_ub: torch.Tensor | None = None, # []
132
+ residual: torch.Tensor | None = None, # [num_tokens, hidden_size]
133
+) -> None:
134
+ torch.ops._C.rms_norm_dynamic_per_token_quant(
135
+ result, input, weight, scale, epsilon, scale_ub, residual
136
+ )
137
+
138
+
139
+# Overwrite autotune_baseline_atol and autotune_baseline_rtol
140
+# if too many configs failed due to baseline check during autotuning
141
+@register_kernel(
142
+ mutates_args=["result", "scale", "residual"],
143
+ config_picker=pick_config,
144
+ input_generator=generate_inputs,
145
+ fake_impl=fake_impl,
146
+ helion_settings=helion.Settings(
147
+ autotune_baseline_fn=baseline,
148
+ ignore_warnings=[helion.exc.TensorOperationInWrapper],
149
+ ),
150
+) # type: ignore[misc]
151
+def rms_norm_dynamic_per_token_quant(
152
+ result: torch.Tensor, # [num_tokens, hidden_size]
153
+ input: torch.Tensor, # [num_tokens, hidden_size]
154
+ weight: torch.Tensor, # [hidden_size]
155
+ scale: torch.Tensor, # [num_tokens, 1]
156
+ epsilon: float,
157
+ scale_ub: torch.Tensor | None = None, # []
158
+ residual: torch.Tensor | None = None, # [num_tokens, hidden_size]
159
+) -> None:
160
+ # This code assumes batch_dim and num_tokens are flattened
161
+ assert input.ndim == 2
162
+ num_tokens, hidden_size = input.shape
163
+ hl.specialize(hidden_size)
164
+
165
+ fp8_dtype = get_fp8_dtype()
166
+ assert result.dtype in [fp8_dtype, torch.int8]
167
+ assert result.is_contiguous() and input.is_contiguous()
168
+
169
+ if scale_ub is not None:
170
+ assert result.dtype == fp8_dtype
171
+ assert scale_ub.dtype == torch.float32
172
+
173
+ assert input.dtype == weight.dtype
174
+ assert scale.shape[0] == num_tokens
175
+ assert scale.dtype == torch.float32
176
+
177
+ if residual is not None:
178
+ assert residual.dtype == input.dtype
179
+
180
+ quant_dtype = result.dtype
181
+ qtype_traits_min: int | float
182
+ qtype_traits_max: int | float
183
+ if quant_dtype == torch.int8:
184
+ qtype_traits_min, qtype_traits_max = get_int8_min_max()
185
+ min_scaling_factor = get_int8_min_scaling_factor()
186
+ else:
187
+ qtype_traits_min, qtype_traits_max = get_fp8_min_max()
188
+ min_scaling_factor = 1.0 / (qtype_traits_max * 512.0)
189
+
190
+ qtype_max = float(qtype_traits_max)
191
+
192
+ for tile_m in hl.tile(num_tokens, block_size=1):
193
+ rms = hl.zeros([tile_m], dtype=torch.float32)
194
+ for tile_n in hl.tile(hidden_size):
195
+ x_blk = input[tile_m, tile_n].to(torch.float32)
196
+ if residual is not None:
197
+ x_blk = x_blk + residual[tile_m, tile_n]
198
+ rms = rms + x_blk.pow(2).sum(dim=-1)
199
+
200
+ rms = torch.rsqrt(rms * (1.0 / hidden_size) + epsilon)
201
+ s_blk = hl.zeros([tile_m], dtype=torch.float32)
202
+
203
+ for tile_n in hl.tile(hidden_size):
204
+ x_blk = input[tile_m, tile_n].to(torch.float32)
205
+ if residual is not None:
206
+ x_blk = x_blk + residual[tile_m, tile_n]
207
+ x_blk = (x_blk * rms[:, None]).to(input.dtype) * weight[None, tile_n]
208
+ tmp_blk = torch.amax(torch.abs(x_blk), dim=-1).to(torch.float32)
209
+ s_blk = torch.maximum(s_blk, tmp_blk)
210
+
211
+ if scale_ub is not None:
212
+ scale_ub_s = hl.load(scale_ub, [])
213
+ s_blk = s_blk.clamp(max=scale_ub_s)
214
+ s_blk = s_blk * (1.0 / qtype_max)
215
+ s_blk = s_blk.clamp(min=min_scaling_factor)
216
+ scale[tile_m, 0] = s_blk
217
+
218
+ for tile_n in hl.tile(hidden_size):
219
+ x_blk = input[tile_m, tile_n].to(torch.float32)
220
+ if residual is not None:
221
+ x_blk = x_blk + residual[tile_m, tile_n]
222
+ residual[tile_m, tile_n] = x_blk.to(residual.dtype)
223
+ x_blk = (x_blk * rms[:, None]).to(input.dtype) * weight[None, tile_n]
224
+ if quant_dtype == torch.int8:
225
+ s_inv_blk = 1.0 / s_blk[:, None]
226
+ y_blk = x_blk * s_inv_blk
227
+ y_blk = y_blk.round()
228
+ else:
229
+ y_blk = x_blk / s_blk[:, None]
230
+
231
+ result[tile_m, tile_n] = y_blk.clamp(qtype_traits_min, qtype_traits_max).to(
232
+ result.dtype
233
+ )