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