[Kernel][Helion][1/N] Add Helion kernel for silu_and_mul_per_block_quant (#43994)
Signed-off-by: Sean Chen <seachen@redhat.com> Co-authored-by: Andreas Karatzas <akaratza@amd.com> Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>
Xiaohong (Sean) Chen committed
Jul 6, 2026 at 12:19 UTC
9fde043f5452b2aa49c90ddeb4848e0834aa32eb
4 files changed
+4782
tests/kernels/helion/test_silu_and_mul_per_block_quant.py
new
+224
@@ -0,0 +1,224 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for the silu_and_mul_per_block_quant helion kernel
4
+Run `pytest tests/kernels/helion/test_silu_and_mul_per_block_quant.py`.
5
+"""
6
+
7
+from typing import Any
8
+
9
+import pytest
10
+import torch
11
+from torch._subclasses.fake_tensor import FakeTensorMode
12
+
13
+from tests.kernels.helion.utils import skip_if_platform_unsupported
14
+from tests.kernels.quant_utils import FP8_DTYPE
15
+from vllm.kernels.helion.case_key import CaseKey
16
+from vllm.kernels.helion.config_manager import ConfigManager
17
+from vllm.kernels.helion.ops.silu_and_mul_per_block_quant import (
18
+ _pick_cache,
19
+ baseline,
20
+ pick_config,
21
+ silu_and_mul_per_block_quant,
22
+)
23
+from vllm.platforms import current_platform
24
+from vllm.utils.import_utils import has_helion
25
+from vllm.utils.torch_utils import set_random_seed
26
+
27
+if not has_helion():
28
+ pytest.skip(
29
+ "Helion is not installed. Install with: pip install vllm[helion]",
30
+ allow_module_level=True,
31
+ )
32
+
33
+
34
+def _generate_fake_input(
35
+ num_tokens: int, intermediate_size: int, group_size: int
36
+) -> tuple[Any, ...]:
37
+ with FakeTensorMode():
38
+ in_dtype: torch.dtype = torch.bfloat16
39
+ out_dtype: torch.dtype = current_platform.fp8_dtype()
40
+ scale_dtype: torch.dtype = torch.float32
41
+ input = torch.randn(
42
+ num_tokens, 2 * intermediate_size, device="cuda", dtype=in_dtype
43
+ )
44
+ result = torch.empty(
45
+ num_tokens, intermediate_size, device=input.device, dtype=out_dtype
46
+ )
47
+ scale = torch.empty(
48
+ (num_tokens, intermediate_size // group_size),
49
+ device=input.device,
50
+ dtype=scale_dtype,
51
+ )
52
+ scale_ub = torch.mean(input).to(scale_dtype)
53
+ args = (
54
+ result,
55
+ input,
56
+ scale,
57
+ group_size,
58
+ scale_ub,
59
+ False,
60
+ )
61
+ return args
62
+
63
+
64
+class TestSiluAndMulPerBlockQuantConfigPicker:
65
+ def setup_method(self):
66
+ _pick_cache.clear()
67
+
68
+ def test_config_picker_exact_match(self):
69
+ config_keys = [
70
+ CaseKey({"intermediate_size": 2048, "group_size": 64, "num_tokens": 16}),
71
+ CaseKey({"intermediate_size": 4096, "group_size": 128, "num_tokens": 16}),
72
+ ]
73
+
74
+ args = _generate_fake_input(16, 4096, 128)
75
+ selected_key = pick_config(args, config_keys)
76
+ assert selected_key == CaseKey(
77
+ {"intermediate_size": 4096, "group_size": 128, "num_tokens": 16}
78
+ )
79
+
80
+ def test_config_picker_closest_match(self):
81
+ config_keys = [
82
+ CaseKey({"intermediate_size": 2048, "group_size": 64, "num_tokens": 16}),
83
+ CaseKey({"intermediate_size": 2048, "group_size": 64, "num_tokens": 32}),
84
+ CaseKey({"intermediate_size": 2048, "group_size": 128, "num_tokens": 16}),
85
+ CaseKey({"intermediate_size": 2048, "group_size": 128, "num_tokens": 32}),
86
+ CaseKey({"intermediate_size": 4096, "group_size": 64, "num_tokens": 16}),
87
+ CaseKey({"intermediate_size": 4096, "group_size": 64, "num_tokens": 32}),
88
+ CaseKey({"intermediate_size": 4096, "group_size": 128, "num_tokens": 16}),
89
+ CaseKey({"intermediate_size": 4096, "group_size": 128, "num_tokens": 32}),
90
+ ]
91
+
92
+ args = _generate_fake_input(20, 3000, 70)
93
+ selected_key = pick_config(args, config_keys)
94
+ assert selected_key == CaseKey(
95
+ {"intermediate_size": 2048, "group_size": 64, "num_tokens": 32}
96
+ )
97
+
98
+ def test_config_picker_no_configs(self):
99
+ config_keys: list[dict] = []
100
+
101
+ args = _generate_fake_input(16, 4096, 128)
102
+ selected_key = pick_config(args, config_keys)
103
+ assert selected_key is None
104
+
105
+ def test_config_picker_fallback_to_largest(self):
106
+ config_keys = [
107
+ CaseKey({"intermediate_size": 2048, "group_size": 64, "num_tokens": 16}),
108
+ CaseKey({"intermediate_size": 2048, "group_size": 64, "num_tokens": 32}),
109
+ CaseKey({"intermediate_size": 2048, "group_size": 128, "num_tokens": 16}),
110
+ CaseKey({"intermediate_size": 2048, "group_size": 128, "num_tokens": 32}),
111
+ CaseKey({"intermediate_size": 4096, "group_size": 64, "num_tokens": 16}),
112
+ CaseKey({"intermediate_size": 4096, "group_size": 64, "num_tokens": 32}),
113
+ CaseKey({"intermediate_size": 4096, "group_size": 128, "num_tokens": 16}),
114
+ CaseKey({"intermediate_size": 4096, "group_size": 128, "num_tokens": 32}),
115
+ ]
116
+
117
+ args = _generate_fake_input(64, 8192, 256)
118
+ selected_key = pick_config(args, config_keys)
119
+ assert selected_key == CaseKey(
120
+ {"intermediate_size": 4096, "group_size": 128, "num_tokens": 32}
121
+ )
122
+
123
+
124
+@pytest.fixture(autouse=True)
125
+def reset_config_manager_singleton():
126
+ ConfigManager.reset_instance()
127
+ ConfigManager()
128
+ yield
129
+ ConfigManager.reset_instance()
130
+
131
+
132
+class TestSiluAndMulPerBlockQuantCorrectness:
133
+ @pytest.mark.parametrize("num_tokens", [1, 7, 4096])
134
+ @pytest.mark.parametrize("hidden_size", [1024, 2048, 5120])
135
+ @pytest.mark.parametrize("group_size", [64, 128])
136
+ @pytest.mark.parametrize("is_scale_transposed", [False, True])
137
+ @pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16])
138
+ @pytest.mark.parametrize("quant_dtype", [current_platform.fp8_dtype(), torch.int8])
139
+ @pytest.mark.parametrize("has_scale_ub", [True, False])
140
+ @pytest.mark.parametrize("seed", [0])
141
+ def test_silu_and_mul_per_block_quant(
142
+ self,
143
+ num_tokens: int,
144
+ hidden_size: int,
145
+ group_size: int,
146
+ is_scale_transposed: bool,
147
+ dtype: torch.dtype,
148
+ quant_dtype: torch.dtype,
149
+ has_scale_ub: bool,
150
+ seed: int,
151
+ ) -> None:
152
+ skip_if_platform_unsupported("silu_and_mul_per_block_quant")
153
+ set_random_seed(seed)
154
+
155
+ if hidden_size % group_size != 0:
156
+ return
157
+
158
+ if has_scale_ub and quant_dtype != FP8_DTYPE:
159
+ # skip
160
+ return
161
+
162
+ scale = 1 / hidden_size
163
+ x = torch.randn(num_tokens, 2 * hidden_size, dtype=dtype, device="cuda") * scale
164
+
165
+ if has_scale_ub:
166
+ act = torch.nn.functional.silu(x[:, :hidden_size]) * x[:, hidden_size:]
167
+ act_abs = act.abs().float()
168
+ scale_ub = 0.5 * (act_abs.mean() + act_abs.amax())
169
+ else:
170
+ scale_ub = None
171
+
172
+ ref_out = torch.empty(num_tokens, hidden_size, device="cuda", dtype=quant_dtype)
173
+
174
+ if is_scale_transposed:
175
+ ref_scales = torch.empty(
176
+ (hidden_size // group_size, x.shape[0]),
177
+ device="cuda",
178
+ dtype=torch.float32,
179
+ ).t()
180
+ else:
181
+ ref_scales = torch.empty(
182
+ (x.shape[0], hidden_size // group_size),
183
+ device="cuda",
184
+ dtype=torch.float32,
185
+ )
186
+
187
+ ops_out = ref_out.clone()
188
+ ops_scales = ref_scales.clone()
189
+
190
+ baseline(ref_out, x, ref_scales, group_size, scale_ub, is_scale_transposed)
191
+ silu_and_mul_per_block_quant(
192
+ ops_out, x, ops_scales, group_size, scale_ub, is_scale_transposed
193
+ )
194
+
195
+ torch.testing.assert_close(ref_scales, ops_scales)
196
+ # allow 1 ULP difference
197
+ assert (
198
+ ref_out.view(torch.uint8).to(torch.int16)
199
+ - ops_out.view(torch.uint8).to(torch.int16)
200
+ ).abs().max() <= 1
201
+
202
+
203
+class TestSiluAndMulPerBlockQuantIntegration:
204
+ def test_kernel_registration_integration(self):
205
+ from vllm.kernels.helion.register import get_registered_kernels
206
+
207
+ registered_kernels = get_registered_kernels()
208
+ assert "silu_and_mul_per_block_quant" in registered_kernels
209
+
210
+ kernel_wrapper = registered_kernels["silu_and_mul_per_block_quant"]
211
+ assert kernel_wrapper.op_name == "silu_and_mul_per_block_quant"
212
+ assert kernel_wrapper._config_picker is not None
213
+ assert kernel_wrapper._mutates_args == ["out", "scales"]
214
+
215
+ def test_fake_impl_functionality(self):
216
+ skip_if_platform_unsupported("silu_and_mul_per_block_quant")
217
+ from vllm.kernels.helion.register import get_registered_kernels
218
+
219
+ registered_kernels = get_registered_kernels()
220
+ kernel_wrapper = registered_kernels["silu_and_mul_per_block_quant"]
221
+ fake_impl = kernel_wrapper._fake_impl
222
+
223
+ args = _generate_fake_input(16, 4096, 128)
224
+ assert fake_impl(*args) is None
vllm/kernels/helion/configs/silu_and_mul_per_block_quant/nvidia_b200.json
new
+2099
@@ -0,0 +1,2099 @@
1
+[
2
+ {
3
+ "key": {
4
+ "intermediate_size": 6144,
5
+ "group_size": 128,
6
+ "num_tokens": 1
7
+ },
8
+ "config": {
9
+ "block_sizes": [
10
+ 4
11
+ ],
12
+ "loop_orders": [
13
+ [
14
+ 1,
15
+ 2,
16
+ 0
17
+ ]
18
+ ],
19
+ "l2_groupings": [
20
+ 32
21
+ ],
22
+ "range_unroll_factors": [
23
+ 0
24
+ ],
25
+ "range_warp_specializes": [],
26
+ "range_num_stages": [],
27
+ "range_multi_buffers": [
28
+ null
29
+ ],
30
+ "range_flattens": [
31
+ null
32
+ ],
33
+ "load_eviction_policies": [
34
+ "first",
35
+ "first",
36
+ ""
37
+ ],
38
+ "num_warps": 2,
39
+ "num_stages": 3,
40
+ "indexing": [
41
+ "tensor_descriptor",
42
+ "pointer",
43
+ "tensor_descriptor",
44
+ "tensor_descriptor",
45
+ "tensor_descriptor"
46
+ ],
47
+ "atomic_indexing": [],
48
+ "pid_type": "flat"
49
+ }
50
+ },
51
+ {
52
+ "key": {
53
+ "intermediate_size": 12288,
54
+ "group_size": 128,
55
+ "num_tokens": 1
56
+ },
57
+ "config": {
58
+ "block_sizes": [
59
+ 4
60
+ ],
61
+ "loop_orders": [
62
+ [
63
+ 0,
64
+ 1,
65
+ 2
66
+ ]
67
+ ],
68
+ "l2_groupings": [
69
+ 16
70
+ ],
71
+ "range_unroll_factors": [
72
+ 0
73
+ ],
74
+ "range_warp_specializes": [],
75
+ "range_num_stages": [],
76
+ "range_multi_buffers": [
77
+ null
78
+ ],
79
+ "range_flattens": [
80
+ null
81
+ ],
82
+ "load_eviction_policies": [
83
+ "",
84
+ "",
85
+ "first"
86
+ ],
87
+ "num_warps": 4,
88
+ "num_stages": 8,
89
+ "indexing": [
90
+ "pointer",
91
+ "tensor_descriptor",
92
+ "tensor_descriptor",
93
+ "pointer",
94
+ "tensor_descriptor"
95
+ ],
96
+ "atomic_indexing": [],
97
+ "pid_type": "flat"
98
+ }
99
+ },
100
+ {
101
+ "key": {
102
+ "intermediate_size": 25600,
103
+ "group_size": 128,
104
+ "num_tokens": 1
105
+ },
106
+ "config": {
107
+ "block_sizes": [
108
+ 4
109
+ ],
110
+ "loop_orders": [
111
+ [
112
+ 1,
113
+ 0,
114
+ 2
115
+ ]
116
+ ],
117
+ "l2_groupings": [
118
+ 32
119
+ ],
120
+ "range_unroll_factors": [
121
+ 4
122
+ ],
123
+ "range_warp_specializes": [
124
+ null
125
+ ],
126
+ "range_multi_buffers": [
127
+ false
128
+ ],
129
+ "range_flattens": [
130
+ null
131
+ ],
132
+ "load_eviction_policies": [
133
+ "last",
134
+ "",
135
+ ""
136
+ ],
137
+ "num_warps": 2,
138
+ "num_stages": 5,
139
+ "indexing": [
140
+ "tensor_descriptor",
141
+ "tensor_descriptor",
142
+ "tensor_descriptor",
143
+ "tensor_descriptor",
144
+ "tensor_descriptor"
145
+ ],
146
+ "atomic_indexing": [],
147
+ "pid_type": "persistent_interleaved",
148
+ "num_sm_multiplier": 4,
149
+ "maxnreg": 128
150
+ }
151
+ },
152
+ {
153
+ "key": {
154
+ "intermediate_size": 6144,
155
+ "group_size": 128,
156
+ "num_tokens": 2
157
+ },
158
+ "config": {
159
+ "block_sizes": [
160
+ 4
161
+ ],
162
+ "loop_orders": [
163
+ [
164
+ 1,
165
+ 2,
166
+ 0
167
+ ]
168
+ ],
169
+ "l2_groupings": [
170
+ 32
171
+ ],
172
+ "range_unroll_factors": [
173
+ 0
174
+ ],
175
+ "range_warp_specializes": [],
176
+ "range_num_stages": [],
177
+ "range_multi_buffers": [
178
+ null
179
+ ],
180
+ "range_flattens": [
181
+ null
182
+ ],
183
+ "load_eviction_policies": [
184
+ "first",
185
+ "first",
186
+ ""
187
+ ],
188
+ "num_warps": 2,
189
+ "num_stages": 3,
190
+ "indexing": [
191
+ "tensor_descriptor",
192
+ "pointer",
193
+ "tensor_descriptor",
194
+ "tensor_descriptor",
195
+ "tensor_descriptor"
196
+ ],
197
+ "atomic_indexing": [],
198
+ "pid_type": "flat"
199
+ }
200
+ },
201
+ {
202
+ "key": {
203
+ "intermediate_size": 12288,
204
+ "group_size": 128,
205
+ "num_tokens": 2
206
+ },
207
+ "config": {
208
+ "block_sizes": [
209
+ 4
210
+ ],
211
+ "loop_orders": [
212
+ [
213
+ 0,
214
+ 1,
215
+ 2
216
+ ]
217
+ ],
218
+ "l2_groupings": [
219
+ 16
220
+ ],
221
+ "range_unroll_factors": [
222
+ 0
223
+ ],
224
+ "range_warp_specializes": [],
225
+ "range_num_stages": [],
226
+ "range_multi_buffers": [
227
+ null
228
+ ],
229
+ "range_flattens": [
230
+ null
231
+ ],
232
+ "load_eviction_policies": [
233
+ "",
234
+ "",
235
+ "first"
236
+ ],
237
+ "num_warps": 4,
238
+ "num_stages": 8,
239
+ "indexing": [
240
+ "pointer",
241
+ "tensor_descriptor",
242
+ "tensor_descriptor",
243
+ "pointer",
244
+ "tensor_descriptor"
245
+ ],
246
+ "atomic_indexing": [],
247
+ "pid_type": "flat"
248
+ }
249
+ },
250
+ {
251
+ "key": {
252
+ "intermediate_size": 25600,
253
+ "group_size": 128,
254
+ "num_tokens": 2
255
+ },
256
+ "config": {
257
+ "block_sizes": [
258
+ 8
259
+ ],
260
+ "loop_orders": [
261
+ [
262
+ 2,
263
+ 1,
264
+ 0
265
+ ]
266
+ ],
267
+ "l2_groupings": [
268
+ 8
269
+ ],
270
+ "range_unroll_factors": [
271
+ 0
272
+ ],
273
+ "range_warp_specializes": [],
274
+ "range_num_stages": [],
275
+ "range_multi_buffers": [
276
+ null
277
+ ],
278
+ "range_flattens": [
279
+ null
280
+ ],
281
+ "load_eviction_policies": [
282
+ "first",
283
+ "first",
284
+ ""
285
+ ],
286
+ "num_warps": 4,
287
+ "num_stages": 4,
288
+ "indexing": [
289
+ "pointer",
290
+ "tensor_descriptor",
291
+ "pointer",
292
+ "pointer",
293
+ "tensor_descriptor"
294
+ ],
295
+ "atomic_indexing": [],
296
+ "pid_type": "flat"
297
+ }
298
+ },
299
+ {
300
+ "key": {
301
+ "intermediate_size": 6144,
302
+ "group_size": 128,
303
+ "num_tokens": 4
304
+ },
305
+ "config": {
306
+ "block_sizes": [
307
+ 4
308
+ ],
309
+ "loop_orders": [
310
+ [
311
+ 0,
312
+ 1,
313
+ 2
314
+ ]
315
+ ],
316
+ "l2_groupings": [
317
+ 16
318
+ ],
319
+ "range_unroll_factors": [
320
+ 0
321
+ ],
322
+ "range_warp_specializes": [],
323
+ "range_num_stages": [],
324
+ "range_multi_buffers": [
325
+ null
326
+ ],
327
+ "range_flattens": [
328
+ null
329
+ ],
330
+ "load_eviction_policies": [
331
+ "",
332
+ "",
333
+ "first"
334
+ ],
335
+ "num_warps": 4,
336
+ "num_stages": 8,
337
+ "indexing": [
338
+ "pointer",
339
+ "tensor_descriptor",
340
+ "tensor_descriptor",
341
+ "pointer",
342
+ "tensor_descriptor"
343
+ ],
344
+ "atomic_indexing": [],
345
+ "pid_type": "flat"
346
+ }
347
+ },
348
+ {
349
+ "key": {
350
+ "intermediate_size": 12288,
351
+ "group_size": 128,
352
+ "num_tokens": 4
353
+ },
354
+ "config": {
355
+ "block_sizes": [
356
+ 4
357
+ ],
358
+ "loop_orders": [
359
+ [
360
+ 2,
361
+ 1,
362
+ 0
363
+ ]
364
+ ],
365
+ "l2_groupings": [
366
+ 8
367
+ ],
368
+ "range_unroll_factors": [
369
+ 0
370
+ ],
371
+ "range_warp_specializes": [],
372
+ "range_num_stages": [],
373
+ "range_multi_buffers": [
374
+ null
375
+ ],
376
+ "range_flattens": [
377
+ null
378
+ ],
379
+ "load_eviction_policies": [
380
+ "first",
381
+ "last",
382
+ ""
383
+ ],
384
+ "num_warps": 4,
385
+ "num_stages": 6,
386
+ "indexing": [
387
+ "tensor_descriptor",
388
+ "pointer",
389
+ "pointer",
390
+ "tensor_descriptor",
391
+ "pointer"
392
+ ],
393
+ "atomic_indexing": [],
394
+ "pid_type": "flat"
395
+ }
396
+ },
397
+ {
398
+ "key": {
399
+ "intermediate_size": 25600,
400
+ "group_size": 128,
401
+ "num_tokens": 4
402
+ },
403
+ "config": {
404
+ "block_sizes": [
405
+ 4
406
+ ],
407
+ "loop_orders": [
408
+ [
409
+ 0,
410
+ 1,
411
+ 2
412
+ ]
413
+ ],
414
+ "l2_groupings": [
415
+ 16
416
+ ],
417
+ "range_unroll_factors": [
418
+ 0
419
+ ],
420
+ "range_warp_specializes": [],
421
+ "range_num_stages": [],
422
+ "range_multi_buffers": [
423
+ null
424
+ ],
425
+ "range_flattens": [
426
+ null
427
+ ],
428
+ "load_eviction_policies": [
429
+ "",
430
+ "",
431
+ "first"
432
+ ],
433
+ "num_warps": 4,
434
+ "num_stages": 8,
435
+ "indexing": [
436
+ "pointer",
437
+ "tensor_descriptor",
438
+ "tensor_descriptor",
439
+ "pointer",
440
+ "tensor_descriptor"
441
+ ],
442
+ "atomic_indexing": [],
443
+ "pid_type": "flat"
444
+ }
445
+ },
446
+ {
447
+ "key": {
448
+ "intermediate_size": 6144,
449
+ "group_size": 128,
450
+ "num_tokens": 8
451
+ },
452
+ "config": {
453
+ "block_sizes": [
454
+ 4
455
+ ],
456
+ "loop_orders": [
457
+ [
458
+ 0,
459
+ 1,
460
+ 2
461
+ ]
462
+ ],
463
+ "l2_groupings": [
464
+ 16
465
+ ],
466
+ "range_unroll_factors": [
467
+ 0
468
+ ],
469
+ "range_warp_specializes": [],
470
+ "range_num_stages": [],
471
+ "range_multi_buffers": [
472
+ null
473
+ ],
474
+ "range_flattens": [
475
+ null
476
+ ],
477
+ "load_eviction_policies": [
478
+ "",
479
+ "",
480
+ "first"
481
+ ],
482
+ "num_warps": 4,
483
+ "num_stages": 8,
484
+ "indexing": [
485
+ "pointer",
486
+ "tensor_descriptor",
487
+ "tensor_descriptor",
488
+ "pointer",
489
+ "tensor_descriptor"
490
+ ],
491
+ "atomic_indexing": [],
492
+ "pid_type": "flat"
493
+ }
494
+ },
495
+ {
496
+ "key": {
497
+ "intermediate_size": 12288,
498
+ "group_size": 128,
499
+ "num_tokens": 8
500
+ },
501
+ "config": {
502
+ "block_sizes": [
503
+ 8
504
+ ],
505
+ "loop_orders": [
506
+ [
507
+ 2,
508
+ 1,
509
+ 0
510
+ ]
511
+ ],
512
+ "l2_groupings": [
513
+ 16
514
+ ],
515
+ "range_unroll_factors": [
516
+ 0
517
+ ],
518
+ "range_warp_specializes": [],
519
+ "range_num_stages": [],
520
+ "range_multi_buffers": [
521
+ null
522
+ ],
523
+ "range_flattens": [
524
+ null
525
+ ],
526
+ "load_eviction_policies": [
527
+ "last",
528
+ "first",
529
+ "first"
530
+ ],
531
+ "num_warps": 4,
532
+ "num_stages": 2,
533
+ "indexing": [
534
+ "tensor_descriptor",
535
+ "tensor_descriptor",
536
+ "pointer",
537
+ "pointer",
538
+ "tensor_descriptor"
539
+ ],
540
+ "atomic_indexing": [],
541
+ "pid_type": "flat"
542
+ }
543
+ },
544
+ {
545
+ "key": {
546
+ "intermediate_size": 25600,
547
+ "group_size": 128,
548
+ "num_tokens": 8
549
+ },
550
+ "config": {
551
+ "block_sizes": [
552
+ 4
553
+ ],
554
+ "loop_orders": [
555
+ [
556
+ 2,
557
+ 1,
558
+ 0
559
+ ]
560
+ ],
561
+ "l2_groupings": [
562
+ 8
563
+ ],
564
+ "range_unroll_factors": [
565
+ 0
566
+ ],
567
+ "range_warp_specializes": [],
568
+ "range_num_stages": [],
569
+ "range_multi_buffers": [
570
+ null
571
+ ],
572
+ "range_flattens": [
573
+ null
574
+ ],
575
+ "load_eviction_policies": [
576
+ "first",
577
+ "last",
578
+ ""
579
+ ],
580
+ "num_warps": 4,
581
+ "num_stages": 6,
582
+ "indexing": [
583
+ "tensor_descriptor",
584
+ "pointer",
585
+ "pointer",
586
+ "tensor_descriptor",
587
+ "pointer"
588
+ ],
589
+ "atomic_indexing": [],
590
+ "pid_type": "flat"
591
+ }
592
+ },
593
+ {
594
+ "key": {
595
+ "intermediate_size": 6144,
596
+ "group_size": 128,
597
+ "num_tokens": 16
598
+ },
599
+ "config": {
600
+ "block_sizes": [
601
+ 4
602
+ ],
603
+ "loop_orders": [
604
+ [
605
+ 0,
606
+ 1,
607
+ 2
608
+ ]
609
+ ],
610
+ "l2_groupings": [
611
+ 16
612
+ ],
613
+ "range_unroll_factors": [
614
+ 0
615
+ ],
616
+ "range_warp_specializes": [],
617
+ "range_num_stages": [],
618
+ "range_multi_buffers": [
619
+ null
620
+ ],
621
+ "range_flattens": [
622
+ null
623
+ ],
624
+ "load_eviction_policies": [
625
+ "",
626
+ "",
627
+ "first"
628
+ ],
629
+ "num_warps": 4,
630
+ "num_stages": 8,
631
+ "indexing": [
632
+ "pointer",
633
+ "tensor_descriptor",
634
+ "tensor_descriptor",
635
+ "pointer",
636
+ "tensor_descriptor"
637
+ ],
638
+ "atomic_indexing": [],
639
+ "pid_type": "flat"
640
+ }
641
+ },
642
+ {
643
+ "key": {
644
+ "intermediate_size": 12288,
645
+ "group_size": 128,
646
+ "num_tokens": 16
647
+ },
648
+ "config": {
649
+ "block_sizes": [
650
+ 4
651
+ ],
652
+ "loop_orders": [
653
+ [
654
+ 2,
655
+ 1,
656
+ 0
657
+ ]
658
+ ],
659
+ "l2_groupings": [
660
+ 8
661
+ ],
662
+ "range_unroll_factors": [
663
+ 0
664
+ ],
665
+ "range_warp_specializes": [],
666
+ "range_num_stages": [],
667
+ "range_multi_buffers": [
668
+ null
669
+ ],
670
+ "range_flattens": [
671
+ null
672
+ ],
673
+ "load_eviction_policies": [
674
+ "first",
675
+ "last",
676
+ ""
677
+ ],
678
+ "num_warps": 4,
679
+ "num_stages": 6,
680
+ "indexing": [
681
+ "tensor_descriptor",
682
+ "pointer",
683
+ "pointer",
684
+ "tensor_descriptor",
685
+ "pointer"
686
+ ],
687
+ "atomic_indexing": [],
688
+ "pid_type": "flat"
689
+ }
690
+ },
691
+ {
692
+ "key": {
693
+ "intermediate_size": 25600,
694
+ "group_size": 128,
695
+ "num_tokens": 16
696
+ },
697
+ "config": {
698
+ "block_sizes": [
699
+ 4
700
+ ],
701
+ "loop_orders": [
702
+ [
703
+ 2,
704
+ 1,
705
+ 0
706
+ ]
707
+ ],
708
+ "l2_groupings": [
709
+ 8
710
+ ],
711
+ "range_unroll_factors": [
712
+ 0
713
+ ],
714
+ "range_warp_specializes": [],
715
+ "range_num_stages": [],
716
+ "range_multi_buffers": [
717
+ null
718
+ ],
719
+ "range_flattens": [
720
+ null
721
+ ],
722
+ "load_eviction_policies": [
723
+ "first",
724
+ "last",
725
+ ""
726
+ ],
727
+ "num_warps": 4,
728
+ "num_stages": 6,
729
+ "indexing": [
730
+ "tensor_descriptor",
731
+ "pointer",
732
+ "pointer",
733
+ "tensor_descriptor",
734
+ "pointer"
735
+ ],
736
+ "atomic_indexing": [],
737
+ "pid_type": "flat"
738
+ }
739
+ },
740
+ {
741
+ "key": {
742
+ "intermediate_size": 6144,
743
+ "group_size": 128,
744
+ "num_tokens": 32
745
+ },
746
+ "config": {
747
+ "block_sizes": [
748
+ 4
749
+ ],
750
+ "loop_orders": [
751
+ [
752
+ 2,
753
+ 1,
754
+ 0
755
+ ]
756
+ ],
757
+ "l2_groupings": [
758
+ 8
759
+ ],
760
+ "range_unroll_factors": [
761
+ 0
762
+ ],
763
+ "range_warp_specializes": [],
764
+ "range_num_stages": [],
765
+ "range_multi_buffers": [
766
+ null
767
+ ],
768
+ "range_flattens": [
769
+ null
770
+ ],
771
+ "load_eviction_policies": [
772
+ "first",
773
+ "last",
774
+ ""
775
+ ],
776
+ "num_warps": 4,
777
+ "num_stages": 6,
778
+ "indexing": [
779
+ "tensor_descriptor",
780
+ "pointer",
781
+ "pointer",
782
+ "tensor_descriptor",
783
+ "pointer"
784
+ ],
785
+ "atomic_indexing": [],
786
+ "pid_type": "flat"
787
+ }
788
+ },
789
+ {
790
+ "key": {
791
+ "intermediate_size": 12288,
792
+ "group_size": 128,
793
+ "num_tokens": 32
794
+ },
795
+ "config": {
796
+ "block_sizes": [
797
+ 4
798
+ ],
799
+ "loop_orders": [
800
+ [
801
+ 1,
802
+ 2,
803
+ 0
804
+ ]
805
+ ],
806
+ "l2_groupings": [
807
+ 32
808
+ ],
809
+ "range_unroll_factors": [
810
+ 0
811
+ ],
812
+ "range_warp_specializes": [
813
+ null
814
+ ],
815
+ "range_num_stages": [],
816
+ "range_multi_buffers": [
817
+ null
818
+ ],
819
+ "range_flattens": [
820
+ null
821
+ ],
822
+ "load_eviction_policies": [
823
+ "",
824
+ "first",
825
+ ""
826
+ ],
827
+ "num_warps": 4,
828
+ "num_stages": 8,
829
+ "indexing": [
830
+ "pointer",
831
+ "tensor_descriptor",
832
+ "tensor_descriptor",
833
+ "pointer",
834
+ "pointer"
835
+ ],
836
+ "atomic_indexing": [],
837
+ "pid_type": "flat"
838
+ }
839
+ },
840
+ {
841
+ "key": {
842
+ "intermediate_size": 25600,
843
+ "group_size": 128,
844
+ "num_tokens": 32
845
+ },
846
+ "config": {
847
+ "block_sizes": [
848
+ 8
849
+ ],
850
+ "loop_orders": [
851
+ [
852
+ 2,
853
+ 1,
854
+ 0
855
+ ]
856
+ ],
857
+ "l2_groupings": [
858
+ 16
859
+ ],
860
+ "range_unroll_factors": [
861
+ 0
862
+ ],
863
+ "range_warp_specializes": [
864
+ null
865
+ ],
866
+ "range_num_stages": [],
867
+ "range_multi_buffers": [
868
+ null
869
+ ],
870
+ "range_flattens": [
871
+ null
872
+ ],
873
+ "load_eviction_policies": [
874
+ "last",
875
+ "last",
876
+ "first"
877
+ ],
878
+ "num_warps": 4,
879
+ "num_stages": 5,
880
+ "indexing": [
881
+ "tensor_descriptor",
882
+ "tensor_descriptor",
883
+ "pointer",
884
+ "tensor_descriptor",
885
+ "tensor_descriptor"
886
+ ],
887
+ "atomic_indexing": [],
888
+ "pid_type": "flat"
889
+ }
890
+ },
891
+ {
892
+ "key": {
893
+ "intermediate_size": 6144,
894
+ "group_size": 128,
895
+ "num_tokens": 64
896
+ },
897
+ "config": {
898
+ "block_sizes": [
899
+ 8
900
+ ],
901
+ "loop_orders": [
902
+ [
903
+ 2,
904
+ 1,
905
+ 0
906
+ ]
907
+ ],
908
+ "l2_groupings": [
909
+ 16
910
+ ],
911
+ "range_unroll_factors": [
912
+ 0
913
+ ],
914
+ "range_warp_specializes": [],
915
+ "range_num_stages": [],
916
+ "range_multi_buffers": [
917
+ null
918
+ ],
919
+ "range_flattens": [
920
+ null
921
+ ],
922
+ "load_eviction_policies": [
923
+ "last",
924
+ "first",
925
+ "first"
926
+ ],
927
+ "num_warps": 4,
928
+ "num_stages": 2,
929
+ "indexing": [
930
+ "tensor_descriptor",
931
+ "tensor_descriptor",
932
+ "pointer",
933
+ "pointer",
934
+ "tensor_descriptor"
935
+ ],
936
+ "atomic_indexing": [],
937
+ "pid_type": "flat"
938
+ }
939
+ },
940
+ {
941
+ "key": {
942
+ "intermediate_size": 12288,
943
+ "group_size": 128,
944
+ "num_tokens": 64
945
+ },
946
+ "config": {
947
+ "block_sizes": [
948
+ 4
949
+ ],
950
+ "loop_orders": [
951
+ [
952
+ 2,
953
+ 1,
954
+ 0
955
+ ]
956
+ ],
957
+ "l2_groupings": [
958
+ 8
959
+ ],
960
+ "range_unroll_factors": [
961
+ 0
962
+ ],
963
+ "range_warp_specializes": [],
964
+ "range_num_stages": [],
965
+ "range_multi_buffers": [
966
+ null
967
+ ],
968
+ "range_flattens": [
969
+ null
970
+ ],
971
+ "load_eviction_policies": [
972
+ "first",
973
+ "last",
974
+ ""
975
+ ],
976
+ "num_warps": 4,
977
+ "num_stages": 6,
978
+ "indexing": [
979
+ "tensor_descriptor",
980
+ "pointer",
981
+ "pointer",
982
+ "tensor_descriptor",
983
+ "pointer"
984
+ ],
985
+ "atomic_indexing": [],
986
+ "pid_type": "flat"
987
+ }
988
+ },
989
+ {
990
+ "key": {
991
+ "intermediate_size": 25600,
992
+ "group_size": 128,
993
+ "num_tokens": 64
994
+ },
995
+ "config": {
996
+ "block_sizes": [
997
+ 4
998
+ ],
999
+ "loop_orders": [
1000
+ [
1001
+ 1,
1002
+ 2,
1003
+ 0
1004
+ ]
1005
+ ],
1006
+ "l2_groupings": [
1007
+ 32
1008
+ ],
1009
+ "range_unroll_factors": [
1010
+ 0
1011
+ ],
1012
+ "range_warp_specializes": [],
1013
+ "range_num_stages": [],
1014
+ "range_multi_buffers": [
1015
+ null
1016
+ ],
1017
+ "range_flattens": [
1018
+ null
1019
+ ],
1020
+ "load_eviction_policies": [
1021
+ "first",
1022
+ "first",
1023
+ ""
1024
+ ],
1025
+ "num_warps": 2,
1026
+ "num_stages": 3,
1027
+ "indexing": [
1028
+ "tensor_descriptor",
1029
+ "pointer",
1030
+ "tensor_descriptor",
1031
+ "tensor_descriptor",
1032
+ "tensor_descriptor"
1033
+ ],
1034
+ "atomic_indexing": [],
1035
+ "pid_type": "flat"
1036
+ }
1037
+ },
1038
+ {
1039
+ "key": {
1040
+ "intermediate_size": 6144,
1041
+ "group_size": 128,
1042
+ "num_tokens": 128
1043
+ },
1044
+ "config": {
1045
+ "block_sizes": [
1046
+ 4
1047
+ ],
1048
+ "loop_orders": [
1049
+ [
1050
+ 2,
1051
+ 1,
1052
+ 0
1053
+ ]
1054
+ ],
1055
+ "l2_groupings": [
1056
+ 8
1057
+ ],
1058
+ "range_unroll_factors": [
1059
+ 0
1060
+ ],
1061
+ "range_warp_specializes": [],
1062
+ "range_num_stages": [],
1063
+ "range_multi_buffers": [
1064
+ null
1065
+ ],
1066
+ "range_flattens": [
1067
+ null
1068
+ ],
1069
+ "load_eviction_policies": [
1070
+ "first",
1071
+ "last",
1072
+ ""
1073
+ ],
1074
+ "num_warps": 4,
1075
+ "num_stages": 6,
1076
+ "indexing": [
1077
+ "tensor_descriptor",
1078
+ "pointer",
1079
+ "pointer",
1080
+ "tensor_descriptor",
1081
+ "pointer"
1082
+ ],
1083
+ "atomic_indexing": [],
1084
+ "pid_type": "flat"
1085
+ }
1086
+ },
1087
+ {
1088
+ "key": {
1089
+ "intermediate_size": 12288,
1090
+ "group_size": 128,
1091
+ "num_tokens": 128
1092
+ },
1093
+ "config": {
1094
+ "block_sizes": [
1095
+ 4
1096
+ ],
1097
+ "loop_orders": [
1098
+ [
1099
+ 1,
1100
+ 2,
1101
+ 0
1102
+ ]
1103
+ ],
1104
+ "l2_groupings": [
1105
+ 32
1106
+ ],
1107
+ "range_unroll_factors": [
1108
+ 0
1109
+ ],
1110
+ "range_warp_specializes": [],
1111
+ "range_num_stages": [],
1112
+ "range_multi_buffers": [
1113
+ null
1114
+ ],
1115
+ "range_flattens": [
1116
+ null
1117
+ ],
1118
+ "load_eviction_policies": [
1119
+ "first",
1120
+ "first",
1121
+ ""
1122
+ ],
1123
+ "num_warps": 2,
1124
+ "num_stages": 3,
1125
+ "indexing": [
1126
+ "tensor_descriptor",
1127
+ "pointer",
1128
+ "tensor_descriptor",
1129
+ "tensor_descriptor",
1130
+ "tensor_descriptor"
1131
+ ],
1132
+ "atomic_indexing": [],
1133
+ "pid_type": "flat"
1134
+ }
1135
+ },
1136
+ {
1137
+ "key": {
1138
+ "intermediate_size": 25600,
1139
+ "group_size": 128,
1140
+ "num_tokens": 128
1141
+ },
1142
+ "config": {
1143
+ "block_sizes": [
1144
+ 8
1145
+ ],
1146
+ "loop_orders": [
1147
+ [
1148
+ 2,
1149
+ 1,
1150
+ 0
1151
+ ]
1152
+ ],
1153
+ "l2_groupings": [
1154
+ 16
1155
+ ],
1156
+ "range_unroll_factors": [
1157
+ 0
1158
+ ],
1159
+ "range_warp_specializes": [],
1160
+ "range_num_stages": [],
1161
+ "range_multi_buffers": [
1162
+ null
1163
+ ],
1164
+ "range_flattens": [
1165
+ null
1166
+ ],
1167
+ "load_eviction_policies": [
1168
+ "last",
1169
+ "first",
1170
+ "first"
1171
+ ],
1172
+ "num_warps": 4,
1173
+ "num_stages": 2,
1174
+ "indexing": [
1175
+ "tensor_descriptor",
1176
+ "tensor_descriptor",
1177
+ "pointer",
1178
+ "pointer",
1179
+ "tensor_descriptor"
1180
+ ],
1181
+ "atomic_indexing": [],
1182
+ "pid_type": "flat"
1183
+ }
1184
+ },
1185
+ {
1186
+ "key": {
1187
+ "intermediate_size": 6144,
1188
+ "group_size": 128,
1189
+ "num_tokens": 256
1190
+ },
1191
+ "config": {
1192
+ "block_sizes": [
1193
+ 4
1194
+ ],
1195
+ "loop_orders": [
1196
+ [
1197
+ 2,
1198
+ 1,
1199
+ 0
1200
+ ]
1201
+ ],
1202
+ "l2_groupings": [
1203
+ 8
1204
+ ],
1205
+ "range_unroll_factors": [
1206
+ 0
1207
+ ],
1208
+ "range_warp_specializes": [],
1209
+ "range_num_stages": [],
1210
+ "range_multi_buffers": [
1211
+ null
1212
+ ],
1213
+ "range_flattens": [
1214
+ null
1215
+ ],
1216
+ "load_eviction_policies": [
1217
+ "first",
1218
+ "last",
1219
+ ""
1220
+ ],
1221
+ "num_warps": 4,
1222
+ "num_stages": 6,
1223
+ "indexing": [
1224
+ "tensor_descriptor",
1225
+ "pointer",
1226
+ "pointer",
1227
+ "tensor_descriptor",
1228
+ "pointer"
1229
+ ],
1230
+ "atomic_indexing": [],
1231
+ "pid_type": "flat"
1232
+ }
1233
+ },
1234
+ {
1235
+ "key": {
1236
+ "intermediate_size": 12288,
1237
+ "group_size": 128,
1238
+ "num_tokens": 256
1239
+ },
1240
+ "config": {
1241
+ "block_sizes": [
1242
+ 16
1243
+ ],
1244
+ "loop_orders": [
1245
+ [
1246
+ 2,
1247
+ 1,
1248
+ 0
1249
+ ]
1250
+ ],
1251
+ "l2_groupings": [
1252
+ 1
1253
+ ],
1254
+ "range_unroll_factors": [
1255
+ 0
1256
+ ],
1257
+ "range_warp_specializes": [
1258
+ null
1259
+ ],
1260
+ "range_num_stages": [],
1261
+ "range_multi_buffers": [
1262
+ null
1263
+ ],
1264
+ "range_flattens": [
1265
+ null
1266
+ ],
1267
+ "load_eviction_policies": [
1268
+ "first",
1269
+ "last",
1270
+ "last"
1271
+ ],
1272
+ "num_warps": 4,
1273
+ "num_stages": 2,
1274
+ "indexing": [
1275
+ "tensor_descriptor",
1276
+ "tensor_descriptor",
1277
+ "pointer",
1278
+ "tensor_descriptor",
1279
+ "pointer"
1280
+ ],
1281
+ "atomic_indexing": [],
1282
+ "pid_type": "flat"
1283
+ }
1284
+ },
1285
+ {
1286
+ "key": {
1287
+ "intermediate_size": 25600,
1288
+ "group_size": 128,
1289
+ "num_tokens": 256
1290
+ },
1291
+ "config": {
1292
+ "block_sizes": [
1293
+ 4
1294
+ ],
1295
+ "loop_orders": [
1296
+ [
1297
+ 2,
1298
+ 1,
1299
+ 0
1300
+ ]
1301
+ ],
1302
+ "l2_groupings": [
1303
+ 1
1304
+ ],
1305
+ "range_unroll_factors": [
1306
+ 0
1307
+ ],
1308
+ "range_warp_specializes": [
1309
+ null
1310
+ ],
1311
+ "range_num_stages": [],
1312
+ "range_multi_buffers": [
1313
+ null
1314
+ ],
1315
+ "range_flattens": [
1316
+ null
1317
+ ],
1318
+ "load_eviction_policies": [
1319
+ "",
1320
+ "first",
1321
+ ""
1322
+ ],
1323
+ "num_warps": 1,
1324
+ "num_stages": 2,
1325
+ "indexing": [
1326
+ "pointer",
1327
+ "pointer",
1328
+ "pointer",
1329
+ "tensor_descriptor",
1330
+ "pointer"
1331
+ ],
1332
+ "atomic_indexing": [],
1333
+ "pid_type": "flat"
1334
+ }
1335
+ },
1336
+ {
1337
+ "key": {
1338
+ "intermediate_size": 6144,
1339
+ "group_size": 128,
1340
+ "num_tokens": 512
1341
+ },
1342
+ "config": {
1343
+ "block_sizes": [
1344
+ 4
1345
+ ],
1346
+ "loop_orders": [
1347
+ [
1348
+ 1,
1349
+ 2,
1350
+ 0
1351
+ ]
1352
+ ],
1353
+ "l2_groupings": [
1354
+ 8
1355
+ ],
1356
+ "range_unroll_factors": [
1357
+ 0
1358
+ ],
1359
+ "range_warp_specializes": [
1360
+ null
1361
+ ],
1362
+ "range_num_stages": [],
1363
+ "range_multi_buffers": [
1364
+ null
1365
+ ],
1366
+ "range_flattens": [
1367
+ null
1368
+ ],
1369
+ "load_eviction_policies": [
1370
+ "",
1371
+ "first",
1372
+ "first"
1373
+ ],
1374
+ "num_warps": 1,
1375
+ "num_stages": 4,
1376
+ "indexing": [
1377
+ "tensor_descriptor",
1378
+ "pointer",
1379
+ "tensor_descriptor",
1380
+ "pointer",
1381
+ "pointer"
1382
+ ],
1383
+ "atomic_indexing": [],
1384
+ "pid_type": "flat"
1385
+ }
1386
+ },
1387
+ {
1388
+ "key": {
1389
+ "intermediate_size": 12288,
1390
+ "group_size": 128,
1391
+ "num_tokens": 512
1392
+ },
1393
+ "config": {
1394
+ "block_sizes": [
1395
+ 4
1396
+ ],
1397
+ "loop_orders": [
1398
+ [
1399
+ 2,
1400
+ 1,
1401
+ 0
1402
+ ]
1403
+ ],
1404
+ "l2_groupings": [
1405
+ 8
1406
+ ],
1407
+ "range_unroll_factors": [
1408
+ 0
1409
+ ],
1410
+ "range_warp_specializes": [
1411
+ null
1412
+ ],
1413
+ "range_num_stages": [],
1414
+ "range_multi_buffers": [
1415
+ null
1416
+ ],
1417
+ "range_flattens": [
1418
+ null
1419
+ ],
1420
+ "load_eviction_policies": [
1421
+ "last",
1422
+ "",
1423
+ "last"
1424
+ ],
1425
+ "num_warps": 1,
1426
+ "num_stages": 5,
1427
+ "indexing": [
1428
+ "tensor_descriptor",
1429
+ "tensor_descriptor",
1430
+ "pointer",
1431
+ "pointer",
1432
+ "pointer"
1433
+ ],
1434
+ "atomic_indexing": [],
1435
+ "pid_type": "flat"
1436
+ }
1437
+ },
1438
+ {
1439
+ "key": {
1440
+ "intermediate_size": 25600,
1441
+ "group_size": 128,
1442
+ "num_tokens": 512
1443
+ },
1444
+ "config": {
1445
+ "block_sizes": [
1446
+ 8
1447
+ ],
1448
+ "loop_orders": [
1449
+ [
1450
+ 2,
1451
+ 1,
1452
+ 0
1453
+ ]
1454
+ ],
1455
+ "l2_groupings": [
1456
+ 8
1457
+ ],
1458
+ "range_unroll_factors": [
1459
+ 0
1460
+ ],
1461
+ "range_warp_specializes": [
1462
+ null
1463
+ ],
1464
+ "range_num_stages": [],
1465
+ "range_multi_buffers": [
1466
+ null
1467
+ ],
1468
+ "range_flattens": [
1469
+ null
1470
+ ],
1471
+ "load_eviction_policies": [
1472
+ "last",
1473
+ "",
1474
+ "last"
1475
+ ],
1476
+ "num_warps": 2,
1477
+ "num_stages": 2,
1478
+ "indexing": [
1479
+ "tensor_descriptor",
1480
+ "tensor_descriptor",
1481
+ "tensor_descriptor",
1482
+ "pointer",
1483
+ "pointer"
1484
+ ],
1485
+ "atomic_indexing": [],
1486
+ "pid_type": "flat"
1487
+ }
1488
+ },
1489
+ {
1490
+ "key": {
1491
+ "intermediate_size": 6144,
1492
+ "group_size": 128,
1493
+ "num_tokens": 1024
1494
+ },
1495
+ "config": {
1496
+ "block_sizes": [
1497
+ 4
1498
+ ],
1499
+ "loop_orders": [
1500
+ [
1501
+ 2,
1502
+ 1,
1503
+ 0
1504
+ ]
1505
+ ],
1506
+ "l2_groupings": [
1507
+ 8
1508
+ ],
1509
+ "range_unroll_factors": [
1510
+ 0
1511
+ ],
1512
+ "range_warp_specializes": [
1513
+ null
1514
+ ],
1515
+ "range_num_stages": [],
1516
+ "range_multi_buffers": [
1517
+ null
1518
+ ],
1519
+ "range_flattens": [
1520
+ null
1521
+ ],
1522
+ "load_eviction_policies": [
1523
+ "last",
1524
+ "",
1525
+ "last"
1526
+ ],
1527
+ "num_warps": 1,
1528
+ "num_stages": 5,
1529
+ "indexing": [
1530
+ "tensor_descriptor",
1531
+ "tensor_descriptor",
1532
+ "pointer",
1533
+ "pointer",
1534
+ "pointer"
1535
+ ],
1536
+ "atomic_indexing": [],
1537
+ "pid_type": "flat"
1538
+ }
1539
+ },
1540
+ {
1541
+ "key": {
1542
+ "intermediate_size": 12288,
1543
+ "group_size": 128,
1544
+ "num_tokens": 1024
1545
+ },
1546
+ "config": {
1547
+ "block_sizes": [
1548
+ 4
1549
+ ],
1550
+ "loop_orders": [
1551
+ [
1552
+ 1,
1553
+ 2,
1554
+ 0
1555
+ ]
1556
+ ],
1557
+ "l2_groupings": [
1558
+ 32
1559
+ ],
1560
+ "range_unroll_factors": [
1561
+ 0
1562
+ ],
1563
+ "range_warp_specializes": [],
1564
+ "range_num_stages": [],
1565
+ "range_multi_buffers": [
1566
+ null
1567
+ ],
1568
+ "range_flattens": [
1569
+ null
1570
+ ],
1571
+ "load_eviction_policies": [
1572
+ "first",
1573
+ "",
1574
+ "last"
1575
+ ],
1576
+ "num_warps": 1,
1577
+ "num_stages": 1,
1578
+ "indexing": [
1579
+ "tensor_descriptor",
1580
+ "tensor_descriptor",
1581
+ "tensor_descriptor",
1582
+ "tensor_descriptor",
1583
+ "pointer"
1584
+ ],
1585
+ "atomic_indexing": [],
1586
+ "pid_type": "flat"
1587
+ }
1588
+ },
1589
+ {
1590
+ "key": {
1591
+ "intermediate_size": 25600,
1592
+ "group_size": 128,
1593
+ "num_tokens": 1024
1594
+ },
1595
+ "config": {
1596
+ "block_sizes": [
1597
+ 8
1598
+ ],
1599
+ "loop_orders": [
1600
+ [
1601
+ 2,
1602
+ 1,
1603
+ 0
1604
+ ]
1605
+ ],
1606
+ "l2_groupings": [
1607
+ 1
1608
+ ],
1609
+ "range_unroll_factors": [
1610
+ 0
1611
+ ],
1612
+ "range_warp_specializes": [
1613
+ null
1614
+ ],
1615
+ "range_num_stages": [],
1616
+ "range_multi_buffers": [
1617
+ null
1618
+ ],
1619
+ "range_flattens": [
1620
+ null
1621
+ ],
1622
+ "load_eviction_policies": [
1623
+ "last",
1624
+ "last",
1625
+ "last"
1626
+ ],
1627
+ "num_warps": 2,
1628
+ "num_stages": 1,
1629
+ "indexing": [
1630
+ "pointer",
1631
+ "pointer",
1632
+ "pointer",
1633
+ "tensor_descriptor",
1634
+ "tensor_descriptor"
1635
+ ],
1636
+ "atomic_indexing": [],
1637
+ "pid_type": "flat"
1638
+ }
1639
+ },
1640
+ {
1641
+ "key": {
1642
+ "intermediate_size": 6144,
1643
+ "group_size": 128,
1644
+ "num_tokens": 2048
1645
+ },
1646
+ "config": {
1647
+ "block_sizes": [
1648
+ 8
1649
+ ],
1650
+ "loop_orders": [
1651
+ [
1652
+ 2,
1653
+ 1,
1654
+ 0
1655
+ ]
1656
+ ],
1657
+ "l2_groupings": [
1658
+ 8
1659
+ ],
1660
+ "range_unroll_factors": [
1661
+ 0
1662
+ ],
1663
+ "range_warp_specializes": [
1664
+ null
1665
+ ],
1666
+ "range_num_stages": [],
1667
+ "range_multi_buffers": [
1668
+ null
1669
+ ],
1670
+ "range_flattens": [
1671
+ null
1672
+ ],
1673
+ "load_eviction_policies": [
1674
+ "last",
1675
+ "",
1676
+ "last"
1677
+ ],
1678
+ "num_warps": 2,
1679
+ "num_stages": 2,
1680
+ "indexing": [
1681
+ "tensor_descriptor",
1682
+ "tensor_descriptor",
1683
+ "tensor_descriptor",
1684
+ "pointer",
1685
+ "pointer"
1686
+ ],
1687
+ "atomic_indexing": [],
1688
+ "pid_type": "flat"
1689
+ }
1690
+ },
1691
+ {
1692
+ "key": {
1693
+ "intermediate_size": 12288,
1694
+ "group_size": 128,
1695
+ "num_tokens": 2048
1696
+ },
1697
+ "config": {
1698
+ "block_sizes": [
1699
+ 16
1700
+ ],
1701
+ "loop_orders": [
1702
+ [
1703
+ 2,
1704
+ 1,
1705
+ 0
1706
+ ]
1707
+ ],
1708
+ "l2_groupings": [
1709
+ 32
1710
+ ],
1711
+ "range_unroll_factors": [
1712
+ 0
1713
+ ],
1714
+ "range_warp_specializes": [
1715
+ null
1716
+ ],
1717
+ "range_num_stages": [],
1718
+ "range_multi_buffers": [
1719
+ null
1720
+ ],
1721
+ "range_flattens": [
1722
+ null
1723
+ ],
1724
+ "load_eviction_policies": [
1725
+ "first",
1726
+ "",
1727
+ ""
1728
+ ],
1729
+ "num_warps": 4,
1730
+ "num_stages": 4,
1731
+ "indexing": [
1732
+ "pointer",
1733
+ "pointer",
1734
+ "tensor_descriptor",
1735
+ "tensor_descriptor",
1736
+ "pointer"
1737
+ ],
1738
+ "atomic_indexing": [],
1739
+ "pid_type": "flat"
1740
+ }
1741
+ },
1742
+ {
1743
+ "key": {
1744
+ "intermediate_size": 25600,
1745
+ "group_size": 128,
1746
+ "num_tokens": 2048
1747
+ },
1748
+ "config": {
1749
+ "block_sizes": [
1750
+ 8
1751
+ ],
1752
+ "loop_orders": [
1753
+ [
1754
+ 2,
1755
+ 1,
1756
+ 0
1757
+ ]
1758
+ ],
1759
+ "l2_groupings": [
1760
+ 8
1761
+ ],
1762
+ "range_unroll_factors": [
1763
+ 0
1764
+ ],
1765
+ "range_warp_specializes": [
1766
+ null
1767
+ ],
1768
+ "range_num_stages": [],
1769
+ "range_multi_buffers": [
1770
+ null
1771
+ ],
1772
+ "range_flattens": [
1773
+ null
1774
+ ],
1775
+ "load_eviction_policies": [
1776
+ "last",
1777
+ "last",
1778
+ "last"
1779
+ ],
1780
+ "num_warps": 2,
1781
+ "num_stages": 5,
1782
+ "indexing": [
1783
+ "pointer",
1784
+ "pointer",
1785
+ "pointer",
1786
+ "tensor_descriptor",
1787
+ "pointer"
1788
+ ],
1789
+ "atomic_indexing": [],
1790
+ "pid_type": "flat"
1791
+ }
1792
+ },
1793
+ {
1794
+ "key": {
1795
+ "intermediate_size": 6144,
1796
+ "group_size": 128,
1797
+ "num_tokens": 4096
1798
+ },
1799
+ "config": {
1800
+ "block_sizes": [
1801
+ 16
1802
+ ],
1803
+ "loop_orders": [
1804
+ [
1805
+ 2,
1806
+ 1,
1807
+ 0
1808
+ ]
1809
+ ],
1810
+ "l2_groupings": [
1811
+ 1
1812
+ ],
1813
+ "range_unroll_factors": [
1814
+ 0
1815
+ ],
1816
+ "range_warp_specializes": [
1817
+ null
1818
+ ],
1819
+ "range_num_stages": [],
1820
+ "range_multi_buffers": [
1821
+ null
1822
+ ],
1823
+ "range_flattens": [
1824
+ null
1825
+ ],
1826
+ "load_eviction_policies": [
1827
+ "first",
1828
+ "last",
1829
+ "last"
1830
+ ],
1831
+ "num_warps": 4,
1832
+ "num_stages": 2,
1833
+ "indexing": [
1834
+ "tensor_descriptor",
1835
+ "tensor_descriptor",
1836
+ "pointer",
1837
+ "tensor_descriptor",
1838
+ "pointer"
1839
+ ],
1840
+ "atomic_indexing": [],
1841
+ "pid_type": "flat"
1842
+ }
1843
+ },
1844
+ {
1845
+ "key": {
1846
+ "intermediate_size": 12288,
1847
+ "group_size": 128,
1848
+ "num_tokens": 4096
1849
+ },
1850
+ "config": {
1851
+ "block_sizes": [
1852
+ 8
1853
+ ],
1854
+ "loop_orders": [
1855
+ [
1856
+ 2,
1857
+ 1,
1858
+ 0
1859
+ ]
1860
+ ],
1861
+ "l2_groupings": [
1862
+ 1
1863
+ ],
1864
+ "range_unroll_factors": [
1865
+ 0
1866
+ ],
1867
+ "range_warp_specializes": [
1868
+ null
1869
+ ],
1870
+ "range_num_stages": [],
1871
+ "range_multi_buffers": [
1872
+ null
1873
+ ],
1874
+ "range_flattens": [
1875
+ null
1876
+ ],
1877
+ "load_eviction_policies": [
1878
+ "last",
1879
+ "last",
1880
+ "last"
1881
+ ],
1882
+ "num_warps": 2,
1883
+ "num_stages": 1,
1884
+ "indexing": [
1885
+ "pointer",
1886
+ "pointer",
1887
+ "pointer",
1888
+ "tensor_descriptor",
1889
+ "tensor_descriptor"
1890
+ ],
1891
+ "atomic_indexing": [],
1892
+ "pid_type": "flat"
1893
+ }
1894
+ },
1895
+ {
1896
+ "key": {
1897
+ "intermediate_size": 25600,
1898
+ "group_size": 128,
1899
+ "num_tokens": 4096
1900
+ },
1901
+ "config": {
1902
+ "block_sizes": [
1903
+ 8
1904
+ ],
1905
+ "loop_orders": [
1906
+ [
1907
+ 2,
1908
+ 1,
1909
+ 0
1910
+ ]
1911
+ ],
1912
+ "l2_groupings": [
1913
+ 4
1914
+ ],
1915
+ "range_unroll_factors": [
1916
+ 0
1917
+ ],
1918
+ "range_warp_specializes": [
1919
+ null
1920
+ ],
1921
+ "range_num_stages": [],
1922
+ "range_multi_buffers": [
1923
+ null
1924
+ ],
1925
+ "range_flattens": [
1926
+ null
1927
+ ],
1928
+ "load_eviction_policies": [
1929
+ "",
1930
+ "last",
1931
+ "last"
1932
+ ],
1933
+ "num_warps": 2,
1934
+ "num_stages": 4,
1935
+ "indexing": [
1936
+ "pointer",
1937
+ "pointer",
1938
+ "pointer",
1939
+ "pointer",
1940
+ "pointer"
1941
+ ],
1942
+ "atomic_indexing": [],
1943
+ "pid_type": "flat"
1944
+ }
1945
+ },
1946
+ {
1947
+ "key": {
1948
+ "intermediate_size": 6144,
1949
+ "group_size": 128,
1950
+ "num_tokens": 8192
1951
+ },
1952
+ "config": {
1953
+ "block_sizes": [
1954
+ 8
1955
+ ],
1956
+ "loop_orders": [
1957
+ [
1958
+ 2,
1959
+ 1,
1960
+ 0
1961
+ ]
1962
+ ],
1963
+ "l2_groupings": [
1964
+ 8
1965
+ ],
1966
+ "range_unroll_factors": [
1967
+ 0
1968
+ ],
1969
+ "range_warp_specializes": [
1970
+ null
1971
+ ],
1972
+ "range_num_stages": [],
1973
+ "range_multi_buffers": [
1974
+ null
1975
+ ],
1976
+ "range_flattens": [
1977
+ null
1978
+ ],
1979
+ "load_eviction_policies": [
1980
+ "last",
1981
+ "last",
1982
+ "last"
1983
+ ],
1984
+ "num_warps": 2,
1985
+ "num_stages": 5,
1986
+ "indexing": [
1987
+ "pointer",
1988
+ "pointer",
1989
+ "pointer",
1990
+ "tensor_descriptor",
1991
+ "pointer"
1992
+ ],
1993
+ "atomic_indexing": [],
1994
+ "pid_type": "flat"
1995
+ }
1996
+ },
1997
+ {
1998
+ "key": {
1999
+ "intermediate_size": 12288,
2000
+ "group_size": 128,
2001
+ "num_tokens": 8192
2002
+ },
2003
+ "config": {
2004
+ "block_sizes": [
2005
+ 8
2006
+ ],
2007
+ "loop_orders": [
2008
+ [
2009
+ 2,
2010
+ 1,
2011
+ 0
2012
+ ]
2013
+ ],
2014
+ "l2_groupings": [
2015
+ 1
2016
+ ],
2017
+ "range_unroll_factors": [
2018
+ 0
2019
+ ],
2020
+ "range_warp_specializes": [
2021
+ null
2022
+ ],
2023
+ "range_num_stages": [],
2024
+ "range_multi_buffers": [
2025
+ null
2026
+ ],
2027
+ "range_flattens": [
2028
+ null
2029
+ ],
2030
+ "load_eviction_policies": [
2031
+ "last",
2032
+ "last",
2033
+ "last"
2034
+ ],
2035
+ "num_warps": 2,
2036
+ "num_stages": 1,
2037
+ "indexing": [
2038
+ "pointer",
2039
+ "pointer",
2040
+ "pointer",
2041
+ "tensor_descriptor",
2042
+ "tensor_descriptor"
2043
+ ],
2044
+ "atomic_indexing": [],
2045
+ "pid_type": "flat"
2046
+ }
2047
+ },
2048
+ {
2049
+ "key": {
2050
+ "intermediate_size": 25600,
2051
+ "group_size": 128,
2052
+ "num_tokens": 8192
2053
+ },
2054
+ "config": {
2055
+ "block_sizes": [
2056
+ 8
2057
+ ],
2058
+ "loop_orders": [
2059
+ [
2060
+ 2,
2061
+ 1,
2062
+ 0
2063
+ ]
2064
+ ],
2065
+ "l2_groupings": [
2066
+ 1
2067
+ ],
2068
+ "range_unroll_factors": [
2069
+ 0
2070
+ ],
2071
+ "range_warp_specializes": [
2072
+ null
2073
+ ],
2074
+ "range_num_stages": [],
2075
+ "range_multi_buffers": [
2076
+ null
2077
+ ],
2078
+ "range_flattens": [
2079
+ null
2080
+ ],
2081
+ "load_eviction_policies": [
2082
+ "last",
2083
+ "last",
2084
+ "last"
2085
+ ],
2086
+ "num_warps": 2,
2087
+ "num_stages": 1,
2088
+ "indexing": [
2089
+ "pointer",
2090
+ "pointer",
2091
+ "pointer",
2092
+ "tensor_descriptor",
2093
+ "tensor_descriptor"
2094
+ ],
2095
+ "atomic_indexing": [],
2096
+ "pid_type": "flat"
2097
+ }
2098
+ }
2099
+]
\ No newline at end of file
vllm/kernels/helion/configs/silu_and_mul_per_block_quant/nvidia_h100.json
new
+2207
@@ -0,0 +1,2207 @@
1
+[
2
+ {
3
+ "key": {
4
+ "intermediate_size": 6144,
5
+ "group_size": 128,
6
+ "num_tokens": 1
7
+ },
8
+ "config": {
9
+ "block_sizes": [
10
+ 4
11
+ ],
12
+ "loop_orders": [
13
+ [
14
+ 2,
15
+ 1,
16
+ 0
17
+ ]
18
+ ],
19
+ "l2_groupings": [
20
+ 4
21
+ ],
22
+ "range_unroll_factors": [
23
+ 0
24
+ ],
25
+ "range_warp_specializes": [],
26
+ "range_num_stages": [],
27
+ "range_multi_buffers": [
28
+ null
29
+ ],
30
+ "range_flattens": [
31
+ null
32
+ ],
33
+ "load_eviction_policies": [
34
+ "",
35
+ "last",
36
+ "last"
37
+ ],
38
+ "num_warps": 4,
39
+ "num_stages": 1,
40
+ "indexing": [
41
+ "pointer",
42
+ "pointer",
43
+ "pointer",
44
+ "tensor_descriptor",
45
+ "tensor_descriptor"
46
+ ],
47
+ "atomic_indexing": [],
48
+ "pid_type": "flat"
49
+ }
50
+ },
51
+ {
52
+ "key": {
53
+ "intermediate_size": 12288,
54
+ "group_size": 128,
55
+ "num_tokens": 1
56
+ },
57
+ "config": {
58
+ "block_sizes": [
59
+ 4
60
+ ],
61
+ "loop_orders": [
62
+ [
63
+ 0,
64
+ 1,
65
+ 2
66
+ ]
67
+ ],
68
+ "l2_groupings": [
69
+ 16
70
+ ],
71
+ "range_unroll_factors": [
72
+ 0
73
+ ],
74
+ "range_warp_specializes": [],
75
+ "range_num_stages": [],
76
+ "range_multi_buffers": [
77
+ null
78
+ ],
79
+ "range_flattens": [
80
+ null
81
+ ],
82
+ "load_eviction_policies": [
83
+ "",
84
+ "",
85
+ "first"
86
+ ],
87
+ "num_warps": 4,
88
+ "num_stages": 8,
89
+ "indexing": [
90
+ "pointer",
91
+ "tensor_descriptor",
92
+ "tensor_descriptor",
93
+ "pointer",
94
+ "tensor_descriptor"
95
+ ],
96
+ "atomic_indexing": [],
97
+ "pid_type": "flat"
98
+ }
99
+ },
100
+ {
101
+ "key": {
102
+ "intermediate_size": 25600,
103
+ "group_size": 128,
104
+ "num_tokens": 1
105
+ },
106
+ "config": {
107
+ "block_sizes": [
108
+ 4
109
+ ],
110
+ "loop_orders": [
111
+ [
112
+ 2,
113
+ 1,
114
+ 0
115
+ ]
116
+ ],
117
+ "l2_groupings": [
118
+ 4
119
+ ],
120
+ "range_unroll_factors": [
121
+ 0
122
+ ],
123
+ "range_warp_specializes": [],
124
+ "range_num_stages": [],
125
+ "range_multi_buffers": [
126
+ null
127
+ ],
128
+ "range_flattens": [
129
+ null
130
+ ],
131
+ "load_eviction_policies": [
132
+ "",
133
+ "last",
134
+ "last"
135
+ ],
136
+ "num_warps": 4,
137
+ "num_stages": 1,
138
+ "indexing": [
139
+ "pointer",
140
+ "pointer",
141
+ "pointer",
142
+ "tensor_descriptor",
143
+ "tensor_descriptor"
144
+ ],
145
+ "atomic_indexing": [],
146
+ "pid_type": "flat"
147
+ }
148
+ },
149
+ {
150
+ "key": {
151
+ "intermediate_size": 6144,
152
+ "group_size": 128,
153
+ "num_tokens": 2
154
+ },
155
+ "config": {
156
+ "block_sizes": [
157
+ 4
158
+ ],
159
+ "loop_orders": [
160
+ [
161
+ 2,
162
+ 1,
163
+ 0
164
+ ]
165
+ ],
166
+ "l2_groupings": [
167
+ 64
168
+ ],
169
+ "range_unroll_factors": [
170
+ 0
171
+ ],
172
+ "range_warp_specializes": [],
173
+ "range_num_stages": [],
174
+ "range_multi_buffers": [
175
+ null
176
+ ],
177
+ "range_flattens": [
178
+ null
179
+ ],
180
+ "load_eviction_policies": [
181
+ "last",
182
+ "first",
183
+ "first"
184
+ ],
185
+ "num_warps": 4,
186
+ "num_stages": 8,
187
+ "indexing": [
188
+ "pointer",
189
+ "pointer",
190
+ "tensor_descriptor",
191
+ "pointer",
192
+ "tensor_descriptor"
193
+ ],
194
+ "atomic_indexing": [],
195
+ "pid_type": "flat"
196
+ }
197
+ },
198
+ {
199
+ "key": {
200
+ "intermediate_size": 12288,
201
+ "group_size": 128,
202
+ "num_tokens": 2
203
+ },
204
+ "config": {
205
+ "block_sizes": [
206
+ 4
207
+ ],
208
+ "loop_orders": [
209
+ [
210
+ 2,
211
+ 1,
212
+ 0
213
+ ]
214
+ ],
215
+ "l2_groupings": [
216
+ 64
217
+ ],
218
+ "range_unroll_factors": [
219
+ 0
220
+ ],
221
+ "range_warp_specializes": [],
222
+ "range_num_stages": [],
223
+ "range_multi_buffers": [
224
+ null
225
+ ],
226
+ "range_flattens": [
227
+ null
228
+ ],
229
+ "load_eviction_policies": [
230
+ "last",
231
+ "first",
232
+ "first"
233
+ ],
234
+ "num_warps": 4,
235
+ "num_stages": 8,
236
+ "indexing": [
237
+ "pointer",
238
+ "pointer",
239
+ "tensor_descriptor",
240
+ "pointer",
241
+ "tensor_descriptor"
242
+ ],
243
+ "atomic_indexing": [],
244
+ "pid_type": "flat"
245
+ }
246
+ },
247
+ {
248
+ "key": {
249
+ "intermediate_size": 25600,
250
+ "group_size": 128,
251
+ "num_tokens": 2
252
+ },
253
+ "config": {
254
+ "block_sizes": [
255
+ 4
256
+ ],
257
+ "loop_orders": [
258
+ [
259
+ 2,
260
+ 1,
261
+ 0
262
+ ]
263
+ ],
264
+ "l2_groupings": [
265
+ 64
266
+ ],
267
+ "range_unroll_factors": [
268
+ 0
269
+ ],
270
+ "range_warp_specializes": [],
271
+ "range_num_stages": [],
272
+ "range_multi_buffers": [
273
+ null
274
+ ],
275
+ "range_flattens": [
276
+ null
277
+ ],
278
+ "load_eviction_policies": [
279
+ "last",
280
+ "first",
281
+ "first"
282
+ ],
283
+ "num_warps": 4,
284
+ "num_stages": 8,
285
+ "indexing": [
286
+ "pointer",
287
+ "pointer",
288
+ "tensor_descriptor",
289
+ "pointer",
290
+ "tensor_descriptor"
291
+ ],
292
+ "atomic_indexing": [],
293
+ "pid_type": "flat"
294
+ }
295
+ },
296
+ {
297
+ "key": {
298
+ "intermediate_size": 6144,
299
+ "group_size": 128,
300
+ "num_tokens": 4
301
+ },
302
+ "config": {
303
+ "block_sizes": [
304
+ 4
305
+ ],
306
+ "loop_orders": [
307
+ [
308
+ 2,
309
+ 1,
310
+ 0
311
+ ]
312
+ ],
313
+ "l2_groupings": [
314
+ 2
315
+ ],
316
+ "range_unroll_factors": [
317
+ 0
318
+ ],
319
+ "range_warp_specializes": [],
320
+ "range_num_stages": [],
321
+ "range_multi_buffers": [
322
+ null
323
+ ],
324
+ "range_flattens": [
325
+ null
326
+ ],
327
+ "load_eviction_policies": [
328
+ "",
329
+ "",
330
+ ""
331
+ ],
332
+ "num_warps": 4,
333
+ "num_stages": 5,
334
+ "indexing": [
335
+ "pointer",
336
+ "pointer",
337
+ "tensor_descriptor",
338
+ "tensor_descriptor",
339
+ "tensor_descriptor"
340
+ ],
341
+ "atomic_indexing": [],
342
+ "pid_type": "flat"
343
+ }
344
+ },
345
+ {
346
+ "key": {
347
+ "intermediate_size": 12288,
348
+ "group_size": 128,
349
+ "num_tokens": 4
350
+ },
351
+ "config": {
352
+ "block_sizes": [
353
+ 4
354
+ ],
355
+ "loop_orders": [
356
+ [
357
+ 2,
358
+ 1,
359
+ 0
360
+ ]
361
+ ],
362
+ "l2_groupings": [
363
+ 4
364
+ ],
365
+ "range_unroll_factors": [
366
+ 0
367
+ ],
368
+ "range_warp_specializes": [],
369
+ "range_num_stages": [],
370
+ "range_multi_buffers": [
371
+ null
372
+ ],
373
+ "range_flattens": [
374
+ null
375
+ ],
376
+ "load_eviction_policies": [
377
+ "",
378
+ "last",
379
+ "last"
380
+ ],
381
+ "num_warps": 4,
382
+ "num_stages": 1,
383
+ "indexing": [
384
+ "pointer",
385
+ "pointer",
386
+ "pointer",
387
+ "tensor_descriptor",
388
+ "tensor_descriptor"
389
+ ],
390
+ "atomic_indexing": [],
391
+ "pid_type": "flat"
392
+ }
393
+ },
394
+ {
395
+ "key": {
396
+ "intermediate_size": 25600,
397
+ "group_size": 128,
398
+ "num_tokens": 4
399
+ },
400
+ "config": {
401
+ "block_sizes": [
402
+ 8
403
+ ],
404
+ "loop_orders": [
405
+ [
406
+ 2,
407
+ 1,
408
+ 0
409
+ ]
410
+ ],
411
+ "l2_groupings": [
412
+ 2
413
+ ],
414
+ "range_unroll_factors": [
415
+ 0
416
+ ],
417
+ "range_warp_specializes": [],
418
+ "range_num_stages": [],
419
+ "range_multi_buffers": [
420
+ null
421
+ ],
422
+ "range_flattens": [
423
+ null
424
+ ],
425
+ "load_eviction_policies": [
426
+ "first",
427
+ "first",
428
+ ""
429
+ ],
430
+ "num_warps": 8,
431
+ "num_stages": 4,
432
+ "indexing": [
433
+ "pointer",
434
+ "tensor_descriptor",
435
+ "tensor_descriptor",
436
+ "pointer",
437
+ "pointer"
438
+ ],
439
+ "atomic_indexing": [],
440
+ "pid_type": "flat"
441
+ }
442
+ },
443
+ {
444
+ "key": {
445
+ "intermediate_size": 6144,
446
+ "group_size": 128,
447
+ "num_tokens": 8
448
+ },
449
+ "config": {
450
+ "block_sizes": [
451
+ 4
452
+ ],
453
+ "loop_orders": [
454
+ [
455
+ 1,
456
+ 2,
457
+ 0
458
+ ]
459
+ ],
460
+ "l2_groupings": [
461
+ 2
462
+ ],
463
+ "range_unroll_factors": [
464
+ 0
465
+ ],
466
+ "range_warp_specializes": [],
467
+ "range_num_stages": [],
468
+ "range_multi_buffers": [
469
+ null
470
+ ],
471
+ "range_flattens": [
472
+ null
473
+ ],
474
+ "load_eviction_policies": [
475
+ "first",
476
+ "",
477
+ "last"
478
+ ],
479
+ "num_warps": 4,
480
+ "num_stages": 3,
481
+ "indexing": [
482
+ "tensor_descriptor",
483
+ "tensor_descriptor",
484
+ "tensor_descriptor",
485
+ "tensor_descriptor",
486
+ "pointer"
487
+ ],
488
+ "atomic_indexing": [],
489
+ "pid_type": "flat"
490
+ }
491
+ },
492
+ {
493
+ "key": {
494
+ "intermediate_size": 12288,
495
+ "group_size": 128,
496
+ "num_tokens": 8
497
+ },
498
+ "config": {
499
+ "block_sizes": [
500
+ 8
501
+ ],
502
+ "loop_orders": [
503
+ [
504
+ 2,
505
+ 1,
506
+ 0
507
+ ]
508
+ ],
509
+ "l2_groupings": [
510
+ 32
511
+ ],
512
+ "range_unroll_factors": [
513
+ 0
514
+ ],
515
+ "range_warp_specializes": [],
516
+ "range_num_stages": [],
517
+ "range_multi_buffers": [
518
+ null
519
+ ],
520
+ "range_flattens": [
521
+ null
522
+ ],
523
+ "load_eviction_policies": [
524
+ "",
525
+ "first",
526
+ ""
527
+ ],
528
+ "num_warps": 8,
529
+ "num_stages": 3,
530
+ "indexing": [
531
+ "tensor_descriptor",
532
+ "pointer",
533
+ "tensor_descriptor",
534
+ "pointer",
535
+ "tensor_descriptor"
536
+ ],
537
+ "atomic_indexing": [],
538
+ "pid_type": "flat"
539
+ }
540
+ },
541
+ {
542
+ "key": {
543
+ "intermediate_size": 25600,
544
+ "group_size": 128,
545
+ "num_tokens": 8
546
+ },
547
+ "config": {
548
+ "block_sizes": [
549
+ 4
550
+ ],
551
+ "loop_orders": [
552
+ [
553
+ 2,
554
+ 1,
555
+ 0
556
+ ]
557
+ ],
558
+ "l2_groupings": [
559
+ 64
560
+ ],
561
+ "range_unroll_factors": [
562
+ 0
563
+ ],
564
+ "range_warp_specializes": [],
565
+ "range_num_stages": [],
566
+ "range_multi_buffers": [
567
+ null
568
+ ],
569
+ "range_flattens": [
570
+ null
571
+ ],
572
+ "load_eviction_policies": [
573
+ "last",
574
+ "first",
575
+ "first"
576
+ ],
577
+ "num_warps": 4,
578
+ "num_stages": 8,
579
+ "indexing": [
580
+ "pointer",
581
+ "pointer",
582
+ "tensor_descriptor",
583
+ "pointer",
584
+ "tensor_descriptor"
585
+ ],
586
+ "atomic_indexing": [],
587
+ "pid_type": "flat"
588
+ }
589
+ },
590
+ {
591
+ "key": {
592
+ "intermediate_size": 6144,
593
+ "group_size": 128,
594
+ "num_tokens": 16
595
+ },
596
+ "config": {
597
+ "block_sizes": [
598
+ 4
599
+ ],
600
+ "loop_orders": [
601
+ [
602
+ 2,
603
+ 1,
604
+ 0
605
+ ]
606
+ ],
607
+ "l2_groupings": [
608
+ 64
609
+ ],
610
+ "range_unroll_factors": [
611
+ 0
612
+ ],
613
+ "range_warp_specializes": [],
614
+ "range_num_stages": [],
615
+ "range_multi_buffers": [
616
+ null
617
+ ],
618
+ "range_flattens": [
619
+ null
620
+ ],
621
+ "load_eviction_policies": [
622
+ "last",
623
+ "first",
624
+ "first"
625
+ ],
626
+ "num_warps": 4,
627
+ "num_stages": 8,
628
+ "indexing": [
629
+ "pointer",
630
+ "pointer",
631
+ "tensor_descriptor",
632
+ "pointer",
633
+ "tensor_descriptor"
634
+ ],
635
+ "atomic_indexing": [],
636
+ "pid_type": "flat"
637
+ }
638
+ },
639
+ {
640
+ "key": {
641
+ "intermediate_size": 12288,
642
+ "group_size": 128,
643
+ "num_tokens": 16
644
+ },
645
+ "config": {
646
+ "block_sizes": [
647
+ 4
648
+ ],
649
+ "loop_orders": [
650
+ [
651
+ 2,
652
+ 1,
653
+ 0
654
+ ]
655
+ ],
656
+ "l2_groupings": [
657
+ 8
658
+ ],
659
+ "range_unroll_factors": [
660
+ 0
661
+ ],
662
+ "range_warp_specializes": [],
663
+ "range_num_stages": [],
664
+ "range_multi_buffers": [
665
+ null
666
+ ],
667
+ "range_flattens": [
668
+ null
669
+ ],
670
+ "load_eviction_policies": [
671
+ "first",
672
+ "last",
673
+ ""
674
+ ],
675
+ "num_warps": 4,
676
+ "num_stages": 6,
677
+ "indexing": [
678
+ "tensor_descriptor",
679
+ "pointer",
680
+ "pointer",
681
+ "tensor_descriptor",
682
+ "pointer"
683
+ ],
684
+ "atomic_indexing": [],
685
+ "pid_type": "flat"
686
+ }
687
+ },
688
+ {
689
+ "key": {
690
+ "intermediate_size": 25600,
691
+ "group_size": 128,
692
+ "num_tokens": 16
693
+ },
694
+ "config": {
695
+ "block_sizes": [
696
+ 8
697
+ ],
698
+ "loop_orders": [
699
+ [
700
+ 2,
701
+ 1,
702
+ 0
703
+ ]
704
+ ],
705
+ "l2_groupings": [
706
+ 16
707
+ ],
708
+ "range_unroll_factors": [
709
+ 0
710
+ ],
711
+ "range_warp_specializes": [],
712
+ "range_num_stages": [],
713
+ "range_multi_buffers": [
714
+ null
715
+ ],
716
+ "range_flattens": [
717
+ null
718
+ ],
719
+ "load_eviction_policies": [
720
+ "last",
721
+ "first",
722
+ "first"
723
+ ],
724
+ "num_warps": 4,
725
+ "num_stages": 2,
726
+ "indexing": [
727
+ "tensor_descriptor",
728
+ "tensor_descriptor",
729
+ "pointer",
730
+ "pointer",
731
+ "tensor_descriptor"
732
+ ],
733
+ "atomic_indexing": [],
734
+ "pid_type": "flat"
735
+ }
736
+ },
737
+ {
738
+ "key": {
739
+ "intermediate_size": 6144,
740
+ "group_size": 128,
741
+ "num_tokens": 32
742
+ },
743
+ "config": {
744
+ "block_sizes": [
745
+ 4
746
+ ],
747
+ "loop_orders": [
748
+ [
749
+ 1,
750
+ 2,
751
+ 0
752
+ ]
753
+ ],
754
+ "l2_groupings": [
755
+ 64
756
+ ],
757
+ "range_unroll_factors": [
758
+ 0
759
+ ],
760
+ "range_warp_specializes": [],
761
+ "range_num_stages": [],
762
+ "range_multi_buffers": [
763
+ null
764
+ ],
765
+ "range_flattens": [
766
+ null
767
+ ],
768
+ "load_eviction_policies": [
769
+ "last",
770
+ "first",
771
+ ""
772
+ ],
773
+ "num_warps": 4,
774
+ "num_stages": 5,
775
+ "indexing": [
776
+ "pointer",
777
+ "pointer",
778
+ "pointer",
779
+ "tensor_descriptor",
780
+ "pointer"
781
+ ],
782
+ "atomic_indexing": [],
783
+ "pid_type": "flat"
784
+ }
785
+ },
786
+ {
787
+ "key": {
788
+ "intermediate_size": 12288,
789
+ "group_size": 128,
790
+ "num_tokens": 32
791
+ },
792
+ "config": {
793
+ "block_sizes": [
794
+ 4
795
+ ],
796
+ "loop_orders": [
797
+ [
798
+ 2,
799
+ 1,
800
+ 0
801
+ ]
802
+ ],
803
+ "l2_groupings": [
804
+ 16
805
+ ],
806
+ "range_unroll_factors": [
807
+ 0
808
+ ],
809
+ "range_warp_specializes": [],
810
+ "range_num_stages": [],
811
+ "range_multi_buffers": [
812
+ null
813
+ ],
814
+ "range_flattens": [
815
+ null
816
+ ],
817
+ "load_eviction_policies": [
818
+ "",
819
+ "first",
820
+ "last"
821
+ ],
822
+ "num_warps": 4,
823
+ "num_stages": 3,
824
+ "indexing": [
825
+ "pointer",
826
+ "tensor_descriptor",
827
+ "tensor_descriptor",
828
+ "pointer",
829
+ "pointer"
830
+ ],
831
+ "atomic_indexing": [],
832
+ "pid_type": "flat"
833
+ }
834
+ },
835
+ {
836
+ "key": {
837
+ "intermediate_size": 25600,
838
+ "group_size": 128,
839
+ "num_tokens": 32
840
+ },
841
+ "config": {
842
+ "block_sizes": [
843
+ 4
844
+ ],
845
+ "loop_orders": [
846
+ [
847
+ 2,
848
+ 1,
849
+ 0
850
+ ]
851
+ ],
852
+ "l2_groupings": [
853
+ 16
854
+ ],
855
+ "range_unroll_factors": [
856
+ 0
857
+ ],
858
+ "range_warp_specializes": [],
859
+ "range_num_stages": [],
860
+ "range_multi_buffers": [
861
+ null
862
+ ],
863
+ "range_flattens": [
864
+ null
865
+ ],
866
+ "load_eviction_policies": [
867
+ "",
868
+ "first",
869
+ "last"
870
+ ],
871
+ "num_warps": 4,
872
+ "num_stages": 3,
873
+ "indexing": [
874
+ "pointer",
875
+ "tensor_descriptor",
876
+ "tensor_descriptor",
877
+ "pointer",
878
+ "pointer"
879
+ ],
880
+ "atomic_indexing": [],
881
+ "pid_type": "flat"
882
+ }
883
+ },
884
+ {
885
+ "key": {
886
+ "intermediate_size": 6144,
887
+ "group_size": 128,
888
+ "num_tokens": 64
889
+ },
890
+ "config": {
891
+ "block_sizes": [
892
+ 4
893
+ ],
894
+ "loop_orders": [
895
+ [
896
+ 2,
897
+ 1,
898
+ 0
899
+ ]
900
+ ],
901
+ "l2_groupings": [
902
+ 32
903
+ ],
904
+ "range_unroll_factors": [
905
+ 0
906
+ ],
907
+ "range_warp_specializes": [],
908
+ "range_num_stages": [],
909
+ "range_multi_buffers": [
910
+ null
911
+ ],
912
+ "range_flattens": [
913
+ null
914
+ ],
915
+ "load_eviction_policies": [
916
+ "last",
917
+ "first",
918
+ ""
919
+ ],
920
+ "num_warps": 4,
921
+ "num_stages": 2,
922
+ "indexing": [
923
+ "tensor_descriptor",
924
+ "tensor_descriptor",
925
+ "pointer",
926
+ "pointer",
927
+ "pointer"
928
+ ],
929
+ "atomic_indexing": [],
930
+ "pid_type": "flat"
931
+ }
932
+ },
933
+ {
934
+ "key": {
935
+ "intermediate_size": 12288,
936
+ "group_size": 128,
937
+ "num_tokens": 64
938
+ },
939
+ "config": {
940
+ "block_sizes": [
941
+ 8
942
+ ],
943
+ "loop_orders": [
944
+ [
945
+ 2,
946
+ 1,
947
+ 0
948
+ ]
949
+ ],
950
+ "l2_groupings": [
951
+ 8
952
+ ],
953
+ "range_unroll_factors": [
954
+ 0
955
+ ],
956
+ "range_warp_specializes": [],
957
+ "range_num_stages": [],
958
+ "range_multi_buffers": [
959
+ null
960
+ ],
961
+ "range_flattens": [
962
+ null
963
+ ],
964
+ "load_eviction_policies": [
965
+ "first",
966
+ "first",
967
+ ""
968
+ ],
969
+ "num_warps": 4,
970
+ "num_stages": 4,
971
+ "indexing": [
972
+ "pointer",
973
+ "tensor_descriptor",
974
+ "pointer",
975
+ "pointer",
976
+ "tensor_descriptor"
977
+ ],
978
+ "atomic_indexing": [],
979
+ "pid_type": "flat"
980
+ }
981
+ },
982
+ {
983
+ "key": {
984
+ "intermediate_size": 25600,
985
+ "group_size": 128,
986
+ "num_tokens": 64
987
+ },
988
+ "config": {
989
+ "block_sizes": [
990
+ 8
991
+ ],
992
+ "loop_orders": [
993
+ [
994
+ 2,
995
+ 1,
996
+ 0
997
+ ]
998
+ ],
999
+ "l2_groupings": [
1000
+ 16
1001
+ ],
1002
+ "range_unroll_factors": [
1003
+ 0
1004
+ ],
1005
+ "range_warp_specializes": [],
1006
+ "range_num_stages": [],
1007
+ "range_multi_buffers": [
1008
+ null
1009
+ ],
1010
+ "range_flattens": [
1011
+ null
1012
+ ],
1013
+ "load_eviction_policies": [
1014
+ "first",
1015
+ "last",
1016
+ "last"
1017
+ ],
1018
+ "num_warps": 4,
1019
+ "num_stages": 3,
1020
+ "indexing": [
1021
+ "pointer",
1022
+ "pointer",
1023
+ "pointer",
1024
+ "tensor_descriptor",
1025
+ "tensor_descriptor"
1026
+ ],
1027
+ "atomic_indexing": [],
1028
+ "pid_type": "flat"
1029
+ }
1030
+ },
1031
+ {
1032
+ "key": {
1033
+ "intermediate_size": 6144,
1034
+ "group_size": 128,
1035
+ "num_tokens": 128
1036
+ },
1037
+ "config": {
1038
+ "block_sizes": [
1039
+ 8
1040
+ ],
1041
+ "loop_orders": [
1042
+ [
1043
+ 1,
1044
+ 2,
1045
+ 0
1046
+ ]
1047
+ ],
1048
+ "l2_groupings": [
1049
+ 32
1050
+ ],
1051
+ "range_unroll_factors": [
1052
+ 0
1053
+ ],
1054
+ "range_warp_specializes": [],
1055
+ "range_num_stages": [],
1056
+ "range_multi_buffers": [
1057
+ null
1058
+ ],
1059
+ "range_flattens": [
1060
+ null
1061
+ ],
1062
+ "load_eviction_policies": [
1063
+ "last",
1064
+ "last",
1065
+ "first"
1066
+ ],
1067
+ "num_warps": 4,
1068
+ "num_stages": 7,
1069
+ "indexing": [
1070
+ "pointer",
1071
+ "tensor_descriptor",
1072
+ "tensor_descriptor",
1073
+ "pointer",
1074
+ "pointer"
1075
+ ],
1076
+ "atomic_indexing": [],
1077
+ "pid_type": "flat"
1078
+ }
1079
+ },
1080
+ {
1081
+ "key": {
1082
+ "intermediate_size": 12288,
1083
+ "group_size": 128,
1084
+ "num_tokens": 128
1085
+ },
1086
+ "config": {
1087
+ "block_sizes": [
1088
+ 8
1089
+ ],
1090
+ "loop_orders": [
1091
+ [
1092
+ 2,
1093
+ 1,
1094
+ 0
1095
+ ]
1096
+ ],
1097
+ "l2_groupings": [
1098
+ 16
1099
+ ],
1100
+ "range_unroll_factors": [
1101
+ 0
1102
+ ],
1103
+ "range_warp_specializes": [],
1104
+ "range_num_stages": [],
1105
+ "range_multi_buffers": [
1106
+ null
1107
+ ],
1108
+ "range_flattens": [
1109
+ null
1110
+ ],
1111
+ "load_eviction_policies": [
1112
+ "last",
1113
+ "first",
1114
+ "first"
1115
+ ],
1116
+ "num_warps": 4,
1117
+ "num_stages": 2,
1118
+ "indexing": [
1119
+ "tensor_descriptor",
1120
+ "tensor_descriptor",
1121
+ "pointer",
1122
+ "pointer",
1123
+ "tensor_descriptor"
1124
+ ],
1125
+ "atomic_indexing": [],
1126
+ "pid_type": "flat"
1127
+ }
1128
+ },
1129
+ {
1130
+ "key": {
1131
+ "intermediate_size": 25600,
1132
+ "group_size": 128,
1133
+ "num_tokens": 128
1134
+ },
1135
+ "config": {
1136
+ "block_sizes": [
1137
+ 8
1138
+ ],
1139
+ "loop_orders": [
1140
+ [
1141
+ 2,
1142
+ 1,
1143
+ 0
1144
+ ]
1145
+ ],
1146
+ "l2_groupings": [
1147
+ 2
1148
+ ],
1149
+ "range_unroll_factors": [
1150
+ 0
1151
+ ],
1152
+ "range_warp_specializes": [],
1153
+ "range_num_stages": [],
1154
+ "range_multi_buffers": [
1155
+ null
1156
+ ],
1157
+ "range_flattens": [
1158
+ null
1159
+ ],
1160
+ "load_eviction_policies": [
1161
+ "first",
1162
+ "first",
1163
+ ""
1164
+ ],
1165
+ "num_warps": 4,
1166
+ "num_stages": 2,
1167
+ "indexing": [
1168
+ "tensor_descriptor",
1169
+ "tensor_descriptor",
1170
+ "pointer",
1171
+ "tensor_descriptor",
1172
+ "pointer"
1173
+ ],
1174
+ "atomic_indexing": [],
1175
+ "pid_type": "flat"
1176
+ }
1177
+ },
1178
+ {
1179
+ "key": {
1180
+ "intermediate_size": 6144,
1181
+ "group_size": 128,
1182
+ "num_tokens": 256
1183
+ },
1184
+ "config": {
1185
+ "block_sizes": [
1186
+ 8
1187
+ ],
1188
+ "loop_orders": [
1189
+ [
1190
+ 0,
1191
+ 2,
1192
+ 1
1193
+ ]
1194
+ ],
1195
+ "l2_groupings": [
1196
+ 1
1197
+ ],
1198
+ "range_unroll_factors": [
1199
+ 0
1200
+ ],
1201
+ "range_warp_specializes": [],
1202
+ "range_num_stages": [],
1203
+ "range_multi_buffers": [
1204
+ null
1205
+ ],
1206
+ "range_flattens": [
1207
+ null
1208
+ ],
1209
+ "load_eviction_policies": [
1210
+ "first",
1211
+ "first",
1212
+ ""
1213
+ ],
1214
+ "num_warps": 4,
1215
+ "num_stages": 7,
1216
+ "indexing": [
1217
+ "pointer",
1218
+ "pointer",
1219
+ "tensor_descriptor",
1220
+ "pointer",
1221
+ "pointer"
1222
+ ],
1223
+ "atomic_indexing": [],
1224
+ "pid_type": "flat"
1225
+ }
1226
+ },
1227
+ {
1228
+ "key": {
1229
+ "intermediate_size": 12288,
1230
+ "group_size": 128,
1231
+ "num_tokens": 256
1232
+ },
1233
+ "config": {
1234
+ "block_sizes": [
1235
+ 8
1236
+ ],
1237
+ "loop_orders": [
1238
+ [
1239
+ 2,
1240
+ 1,
1241
+ 0
1242
+ ]
1243
+ ],
1244
+ "l2_groupings": [
1245
+ 2
1246
+ ],
1247
+ "range_unroll_factors": [
1248
+ 0
1249
+ ],
1250
+ "range_warp_specializes": [],
1251
+ "range_num_stages": [],
1252
+ "range_multi_buffers": [
1253
+ null
1254
+ ],
1255
+ "range_flattens": [
1256
+ null
1257
+ ],
1258
+ "load_eviction_policies": [
1259
+ "first",
1260
+ "",
1261
+ ""
1262
+ ],
1263
+ "num_warps": 4,
1264
+ "num_stages": 2,
1265
+ "indexing": [
1266
+ "tensor_descriptor",
1267
+ "tensor_descriptor",
1268
+ "tensor_descriptor",
1269
+ "pointer",
1270
+ "tensor_descriptor"
1271
+ ],
1272
+ "atomic_indexing": [],
1273
+ "pid_type": "flat"
1274
+ }
1275
+ },
1276
+ {
1277
+ "key": {
1278
+ "intermediate_size": 25600,
1279
+ "group_size": 128,
1280
+ "num_tokens": 256
1281
+ },
1282
+ "config": {
1283
+ "block_sizes": [
1284
+ 8
1285
+ ],
1286
+ "loop_orders": [
1287
+ [
1288
+ 2,
1289
+ 1,
1290
+ 0
1291
+ ]
1292
+ ],
1293
+ "l2_groupings": [
1294
+ 8
1295
+ ],
1296
+ "range_unroll_factors": [
1297
+ 0
1298
+ ],
1299
+ "range_warp_specializes": [],
1300
+ "range_num_stages": [],
1301
+ "range_multi_buffers": [
1302
+ null
1303
+ ],
1304
+ "range_flattens": [
1305
+ null
1306
+ ],
1307
+ "load_eviction_policies": [
1308
+ "first",
1309
+ "first",
1310
+ ""
1311
+ ],
1312
+ "num_warps": 4,
1313
+ "num_stages": 4,
1314
+ "indexing": [
1315
+ "tensor_descriptor",
1316
+ "pointer",
1317
+ "pointer",
1318
+ "pointer",
1319
+ "pointer"
1320
+ ],
1321
+ "atomic_indexing": [],
1322
+ "pid_type": "flat"
1323
+ }
1324
+ },
1325
+ {
1326
+ "key": {
1327
+ "intermediate_size": 6144,
1328
+ "group_size": 128,
1329
+ "num_tokens": 512
1330
+ },
1331
+ "config": {
1332
+ "block_sizes": [
1333
+ 8
1334
+ ],
1335
+ "loop_orders": [
1336
+ [
1337
+ 2,
1338
+ 1,
1339
+ 0
1340
+ ]
1341
+ ],
1342
+ "l2_groupings": [
1343
+ 64
1344
+ ],
1345
+ "range_unroll_factors": [
1346
+ 0
1347
+ ],
1348
+ "range_warp_specializes": [],
1349
+ "range_num_stages": [],
1350
+ "range_multi_buffers": [
1351
+ null
1352
+ ],
1353
+ "range_flattens": [
1354
+ null
1355
+ ],
1356
+ "load_eviction_policies": [
1357
+ "",
1358
+ "",
1359
+ ""
1360
+ ],
1361
+ "num_warps": 4,
1362
+ "num_stages": 6,
1363
+ "indexing": [
1364
+ "tensor_descriptor",
1365
+ "tensor_descriptor",
1366
+ "pointer",
1367
+ "tensor_descriptor",
1368
+ "pointer"
1369
+ ],
1370
+ "atomic_indexing": [],
1371
+ "pid_type": "flat"
1372
+ }
1373
+ },
1374
+ {
1375
+ "key": {
1376
+ "intermediate_size": 12288,
1377
+ "group_size": 128,
1378
+ "num_tokens": 512
1379
+ },
1380
+ "config": {
1381
+ "block_sizes": [
1382
+ 4
1383
+ ],
1384
+ "loop_orders": [
1385
+ [
1386
+ 2,
1387
+ 1,
1388
+ 0
1389
+ ]
1390
+ ],
1391
+ "l2_groupings": [
1392
+ 4
1393
+ ],
1394
+ "range_unroll_factors": [
1395
+ 0
1396
+ ],
1397
+ "range_warp_specializes": [],
1398
+ "range_num_stages": [],
1399
+ "range_multi_buffers": [
1400
+ null
1401
+ ],
1402
+ "range_flattens": [
1403
+ null
1404
+ ],
1405
+ "load_eviction_policies": [
1406
+ "first",
1407
+ "first",
1408
+ ""
1409
+ ],
1410
+ "num_warps": 2,
1411
+ "num_stages": 4,
1412
+ "indexing": [
1413
+ "pointer",
1414
+ "pointer",
1415
+ "tensor_descriptor",
1416
+ "pointer",
1417
+ "tensor_descriptor"
1418
+ ],
1419
+ "atomic_indexing": [],
1420
+ "pid_type": "flat"
1421
+ }
1422
+ },
1423
+ {
1424
+ "key": {
1425
+ "intermediate_size": 25600,
1426
+ "group_size": 128,
1427
+ "num_tokens": 512
1428
+ },
1429
+ "config": {
1430
+ "block_sizes": [
1431
+ 8
1432
+ ],
1433
+ "loop_orders": [
1434
+ [
1435
+ 1,
1436
+ 2,
1437
+ 0
1438
+ ]
1439
+ ],
1440
+ "l2_groupings": [
1441
+ 1
1442
+ ],
1443
+ "range_unroll_factors": [
1444
+ 0
1445
+ ],
1446
+ "range_warp_specializes": [],
1447
+ "range_num_stages": [],
1448
+ "range_multi_buffers": [
1449
+ null
1450
+ ],
1451
+ "range_flattens": [
1452
+ null
1453
+ ],
1454
+ "load_eviction_policies": [
1455
+ "first",
1456
+ "first",
1457
+ ""
1458
+ ],
1459
+ "num_warps": 4,
1460
+ "num_stages": 7,
1461
+ "indexing": [
1462
+ "tensor_descriptor",
1463
+ "pointer",
1464
+ "pointer",
1465
+ "tensor_descriptor",
1466
+ "pointer"
1467
+ ],
1468
+ "atomic_indexing": [],
1469
+ "pid_type": "flat"
1470
+ }
1471
+ },
1472
+ {
1473
+ "key": {
1474
+ "intermediate_size": 6144,
1475
+ "group_size": 128,
1476
+ "num_tokens": 1024
1477
+ },
1478
+ "config": {
1479
+ "block_sizes": [
1480
+ 4
1481
+ ],
1482
+ "loop_orders": [
1483
+ [
1484
+ 1,
1485
+ 2,
1486
+ 0
1487
+ ]
1488
+ ],
1489
+ "l2_groupings": [
1490
+ 32
1491
+ ],
1492
+ "range_unroll_factors": [
1493
+ 0
1494
+ ],
1495
+ "range_warp_specializes": [],
1496
+ "range_num_stages": [],
1497
+ "range_multi_buffers": [
1498
+ null
1499
+ ],
1500
+ "range_flattens": [
1501
+ null
1502
+ ],
1503
+ "load_eviction_policies": [
1504
+ "first",
1505
+ "first",
1506
+ ""
1507
+ ],
1508
+ "num_warps": 2,
1509
+ "num_stages": 3,
1510
+ "indexing": [
1511
+ "tensor_descriptor",
1512
+ "pointer",
1513
+ "tensor_descriptor",
1514
+ "tensor_descriptor",
1515
+ "tensor_descriptor"
1516
+ ],
1517
+ "atomic_indexing": [],
1518
+ "pid_type": "flat"
1519
+ }
1520
+ },
1521
+ {
1522
+ "key": {
1523
+ "intermediate_size": 12288,
1524
+ "group_size": 128,
1525
+ "num_tokens": 1024
1526
+ },
1527
+ "config": {
1528
+ "block_sizes": [
1529
+ 4
1530
+ ],
1531
+ "loop_orders": [
1532
+ [
1533
+ 1,
1534
+ 2,
1535
+ 0
1536
+ ]
1537
+ ],
1538
+ "l2_groupings": [
1539
+ 32
1540
+ ],
1541
+ "range_unroll_factors": [
1542
+ 0
1543
+ ],
1544
+ "range_warp_specializes": [],
1545
+ "range_num_stages": [],
1546
+ "range_multi_buffers": [
1547
+ null
1548
+ ],
1549
+ "range_flattens": [
1550
+ null
1551
+ ],
1552
+ "load_eviction_policies": [
1553
+ "first",
1554
+ "first",
1555
+ ""
1556
+ ],
1557
+ "num_warps": 2,
1558
+ "num_stages": 3,
1559
+ "indexing": [
1560
+ "tensor_descriptor",
1561
+ "pointer",
1562
+ "tensor_descriptor",
1563
+ "tensor_descriptor",
1564
+ "tensor_descriptor"
1565
+ ],
1566
+ "atomic_indexing": [],
1567
+ "pid_type": "flat"
1568
+ }
1569
+ },
1570
+ {
1571
+ "key": {
1572
+ "intermediate_size": 25600,
1573
+ "group_size": 128,
1574
+ "num_tokens": 1024
1575
+ },
1576
+ "config": {
1577
+ "block_sizes": [
1578
+ 8
1579
+ ],
1580
+ "loop_orders": [
1581
+ [
1582
+ 1,
1583
+ 2,
1584
+ 0
1585
+ ]
1586
+ ],
1587
+ "l2_groupings": [
1588
+ 1
1589
+ ],
1590
+ "range_unroll_factors": [
1591
+ 0
1592
+ ],
1593
+ "range_warp_specializes": [],
1594
+ "range_num_stages": [],
1595
+ "range_multi_buffers": [
1596
+ null
1597
+ ],
1598
+ "range_flattens": [
1599
+ null
1600
+ ],
1601
+ "load_eviction_policies": [
1602
+ "first",
1603
+ "first",
1604
+ ""
1605
+ ],
1606
+ "num_warps": 4,
1607
+ "num_stages": 7,
1608
+ "indexing": [
1609
+ "tensor_descriptor",
1610
+ "pointer",
1611
+ "pointer",
1612
+ "tensor_descriptor",
1613
+ "pointer"
1614
+ ],
1615
+ "atomic_indexing": [],
1616
+ "pid_type": "flat"
1617
+ }
1618
+ },
1619
+ {
1620
+ "key": {
1621
+ "intermediate_size": 6144,
1622
+ "group_size": 128,
1623
+ "num_tokens": 2048
1624
+ },
1625
+ "config": {
1626
+ "block_sizes": [
1627
+ 4
1628
+ ],
1629
+ "loop_orders": [
1630
+ [
1631
+ 1,
1632
+ 2,
1633
+ 0
1634
+ ]
1635
+ ],
1636
+ "l2_groupings": [
1637
+ 32
1638
+ ],
1639
+ "range_unroll_factors": [
1640
+ 0
1641
+ ],
1642
+ "range_warp_specializes": [],
1643
+ "range_num_stages": [],
1644
+ "range_multi_buffers": [
1645
+ null
1646
+ ],
1647
+ "range_flattens": [
1648
+ null
1649
+ ],
1650
+ "load_eviction_policies": [
1651
+ "first",
1652
+ "first",
1653
+ ""
1654
+ ],
1655
+ "num_warps": 2,
1656
+ "num_stages": 3,
1657
+ "indexing": [
1658
+ "tensor_descriptor",
1659
+ "pointer",
1660
+ "tensor_descriptor",
1661
+ "tensor_descriptor",
1662
+ "tensor_descriptor"
1663
+ ],
1664
+ "atomic_indexing": [],
1665
+ "pid_type": "flat"
1666
+ }
1667
+ },
1668
+ {
1669
+ "key": {
1670
+ "intermediate_size": 12288,
1671
+ "group_size": 128,
1672
+ "num_tokens": 2048
1673
+ },
1674
+ "config": {
1675
+ "block_sizes": [
1676
+ 8
1677
+ ],
1678
+ "loop_orders": [
1679
+ [
1680
+ 2,
1681
+ 1,
1682
+ 0
1683
+ ]
1684
+ ],
1685
+ "l2_groupings": [
1686
+ 2
1687
+ ],
1688
+ "range_unroll_factors": [
1689
+ 0
1690
+ ],
1691
+ "range_warp_specializes": [],
1692
+ "range_num_stages": [],
1693
+ "range_multi_buffers": [
1694
+ null
1695
+ ],
1696
+ "range_flattens": [
1697
+ null
1698
+ ],
1699
+ "load_eviction_policies": [
1700
+ "first",
1701
+ "first",
1702
+ ""
1703
+ ],
1704
+ "num_warps": 4,
1705
+ "num_stages": 2,
1706
+ "indexing": [
1707
+ "tensor_descriptor",
1708
+ "tensor_descriptor",
1709
+ "pointer",
1710
+ "tensor_descriptor",
1711
+ "pointer"
1712
+ ],
1713
+ "atomic_indexing": [],
1714
+ "pid_type": "flat"
1715
+ }
1716
+ },
1717
+ {
1718
+ "key": {
1719
+ "intermediate_size": 25600,
1720
+ "group_size": 128,
1721
+ "num_tokens": 2048
1722
+ },
1723
+ "config": {
1724
+ "block_sizes": [
1725
+ 8
1726
+ ],
1727
+ "loop_orders": [
1728
+ [
1729
+ 2,
1730
+ 1,
1731
+ 0
1732
+ ]
1733
+ ],
1734
+ "l2_groupings": [
1735
+ 8
1736
+ ],
1737
+ "range_unroll_factors": [
1738
+ 0
1739
+ ],
1740
+ "range_warp_specializes": [],
1741
+ "range_num_stages": [],
1742
+ "range_multi_buffers": [
1743
+ null
1744
+ ],
1745
+ "range_flattens": [
1746
+ null
1747
+ ],
1748
+ "load_eviction_policies": [
1749
+ "first",
1750
+ "first",
1751
+ ""
1752
+ ],
1753
+ "num_warps": 4,
1754
+ "num_stages": 4,
1755
+ "indexing": [
1756
+ "tensor_descriptor",
1757
+ "pointer",
1758
+ "pointer",
1759
+ "pointer",
1760
+ "pointer"
1761
+ ],
1762
+ "atomic_indexing": [],
1763
+ "pid_type": "flat"
1764
+ }
1765
+ },
1766
+ {
1767
+ "key": {
1768
+ "intermediate_size": 6144,
1769
+ "group_size": 128,
1770
+ "num_tokens": 4096
1771
+ },
1772
+ "config": {
1773
+ "block_sizes": [
1774
+ 16
1775
+ ],
1776
+ "loop_orders": [
1777
+ [
1778
+ 2,
1779
+ 1,
1780
+ 0
1781
+ ]
1782
+ ],
1783
+ "l2_groupings": [
1784
+ 4
1785
+ ],
1786
+ "range_unroll_factors": [
1787
+ 0
1788
+ ],
1789
+ "range_warp_specializes": [],
1790
+ "range_num_stages": [],
1791
+ "range_multi_buffers": [
1792
+ null
1793
+ ],
1794
+ "range_flattens": [
1795
+ null
1796
+ ],
1797
+ "load_eviction_policies": [
1798
+ "first",
1799
+ "first",
1800
+ "first"
1801
+ ],
1802
+ "num_warps": 8,
1803
+ "num_stages": 1,
1804
+ "indexing": [
1805
+ "pointer",
1806
+ "tensor_descriptor",
1807
+ "tensor_descriptor",
1808
+ "pointer",
1809
+ "tensor_descriptor"
1810
+ ],
1811
+ "atomic_indexing": [],
1812
+ "pid_type": "flat"
1813
+ }
1814
+ },
1815
+ {
1816
+ "key": {
1817
+ "intermediate_size": 12288,
1818
+ "group_size": 128,
1819
+ "num_tokens": 4096
1820
+ },
1821
+ "config": {
1822
+ "block_sizes": [
1823
+ 8
1824
+ ],
1825
+ "loop_orders": [
1826
+ [
1827
+ 2,
1828
+ 1,
1829
+ 0
1830
+ ]
1831
+ ],
1832
+ "l2_groupings": [
1833
+ 8
1834
+ ],
1835
+ "range_unroll_factors": [
1836
+ 0
1837
+ ],
1838
+ "range_warp_specializes": [],
1839
+ "range_num_stages": [],
1840
+ "range_multi_buffers": [
1841
+ null
1842
+ ],
1843
+ "range_flattens": [
1844
+ null
1845
+ ],
1846
+ "load_eviction_policies": [
1847
+ "first",
1848
+ "first",
1849
+ ""
1850
+ ],
1851
+ "num_warps": 4,
1852
+ "num_stages": 4,
1853
+ "indexing": [
1854
+ "pointer",
1855
+ "tensor_descriptor",
1856
+ "pointer",
1857
+ "pointer",
1858
+ "tensor_descriptor"
1859
+ ],
1860
+ "atomic_indexing": [],
1861
+ "pid_type": "flat"
1862
+ }
1863
+ },
1864
+ {
1865
+ "key": {
1866
+ "intermediate_size": 25600,
1867
+ "group_size": 128,
1868
+ "num_tokens": 4096
1869
+ },
1870
+ "config": {
1871
+ "block_sizes": [
1872
+ 8
1873
+ ],
1874
+ "loop_orders": [
1875
+ [
1876
+ 2,
1877
+ 1,
1878
+ 0
1879
+ ]
1880
+ ],
1881
+ "l2_groupings": [
1882
+ 2
1883
+ ],
1884
+ "range_unroll_factors": [
1885
+ 0
1886
+ ],
1887
+ "range_warp_specializes": [],
1888
+ "range_num_stages": [],
1889
+ "range_multi_buffers": [
1890
+ null
1891
+ ],
1892
+ "range_flattens": [
1893
+ null
1894
+ ],
1895
+ "load_eviction_policies": [
1896
+ "first",
1897
+ "",
1898
+ ""
1899
+ ],
1900
+ "num_warps": 4,
1901
+ "num_stages": 2,
1902
+ "indexing": [
1903
+ "tensor_descriptor",
1904
+ "tensor_descriptor",
1905
+ "tensor_descriptor",
1906
+ "pointer",
1907
+ "tensor_descriptor"
1908
+ ],
1909
+ "atomic_indexing": [],
1910
+ "pid_type": "flat"
1911
+ }
1912
+ },
1913
+ {
1914
+ "key": {
1915
+ "intermediate_size": 6144,
1916
+ "group_size": 128,
1917
+ "num_tokens": 8192
1918
+ },
1919
+ "config": {
1920
+ "block_sizes": [
1921
+ 16
1922
+ ],
1923
+ "loop_orders": [
1924
+ [
1925
+ 2,
1926
+ 1,
1927
+ 0
1928
+ ]
1929
+ ],
1930
+ "l2_groupings": [
1931
+ 4
1932
+ ],
1933
+ "range_unroll_factors": [
1934
+ 0
1935
+ ],
1936
+ "range_warp_specializes": [],
1937
+ "range_num_stages": [],
1938
+ "range_multi_buffers": [
1939
+ null
1940
+ ],
1941
+ "range_flattens": [
1942
+ null
1943
+ ],
1944
+ "load_eviction_policies": [
1945
+ "first",
1946
+ "first",
1947
+ "first"
1948
+ ],
1949
+ "num_warps": 8,
1950
+ "num_stages": 1,
1951
+ "indexing": [
1952
+ "pointer",
1953
+ "tensor_descriptor",
1954
+ "tensor_descriptor",
1955
+ "pointer",
1956
+ "tensor_descriptor"
1957
+ ],
1958
+ "atomic_indexing": [],
1959
+ "pid_type": "flat"
1960
+ }
1961
+ },
1962
+ {
1963
+ "key": {
1964
+ "intermediate_size": 12288,
1965
+ "group_size": 128,
1966
+ "num_tokens": 8192
1967
+ },
1968
+ "config": {
1969
+ "block_sizes": [
1970
+ 8
1971
+ ],
1972
+ "loop_orders": [
1973
+ [
1974
+ 2,
1975
+ 1,
1976
+ 0
1977
+ ]
1978
+ ],
1979
+ "l2_groupings": [
1980
+ 1
1981
+ ],
1982
+ "range_unroll_factors": [
1983
+ 0
1984
+ ],
1985
+ "range_warp_specializes": [],
1986
+ "range_num_stages": [],
1987
+ "range_multi_buffers": [
1988
+ null
1989
+ ],
1990
+ "range_flattens": [
1991
+ null
1992
+ ],
1993
+ "load_eviction_policies": [
1994
+ "first",
1995
+ "",
1996
+ ""
1997
+ ],
1998
+ "num_warps": 4,
1999
+ "num_stages": 5,
2000
+ "indexing": [
2001
+ "pointer",
2002
+ "pointer",
2003
+ "tensor_descriptor",
2004
+ "tensor_descriptor",
2005
+ "pointer"
2006
+ ],
2007
+ "atomic_indexing": [],
2008
+ "pid_type": "flat"
2009
+ }
2010
+ },
2011
+ {
2012
+ "key": {
2013
+ "intermediate_size": 25600,
2014
+ "group_size": 128,
2015
+ "num_tokens": 8192
2016
+ },
2017
+ "config": {
2018
+ "block_sizes": [
2019
+ 8
2020
+ ],
2021
+ "loop_orders": [
2022
+ [
2023
+ 2,
2024
+ 1,
2025
+ 0
2026
+ ]
2027
+ ],
2028
+ "l2_groupings": [
2029
+ 8
2030
+ ],
2031
+ "range_unroll_factors": [
2032
+ 0
2033
+ ],
2034
+ "range_warp_specializes": [],
2035
+ "range_num_stages": [],
2036
+ "range_multi_buffers": [
2037
+ null
2038
+ ],
2039
+ "range_flattens": [
2040
+ null
2041
+ ],
2042
+ "load_eviction_policies": [
2043
+ "first",
2044
+ "",
2045
+ ""
2046
+ ],
2047
+ "num_warps": 4,
2048
+ "num_stages": 3,
2049
+ "indexing": [
2050
+ "tensor_descriptor",
2051
+ "pointer",
2052
+ "pointer",
2053
+ "pointer",
2054
+ "pointer"
2055
+ ],
2056
+ "atomic_indexing": [],
2057
+ "pid_type": "flat"
2058
+ }
2059
+ },
2060
+ {
2061
+ "key": {
2062
+ "intermediate_size": 6144,
2063
+ "group_size": 128,
2064
+ "num_tokens": 16384
2065
+ },
2066
+ "config": {
2067
+ "block_sizes": [
2068
+ 4
2069
+ ],
2070
+ "loop_orders": [
2071
+ [
2072
+ 1,
2073
+ 2,
2074
+ 0
2075
+ ]
2076
+ ],
2077
+ "l2_groupings": [
2078
+ 64
2079
+ ],
2080
+ "range_unroll_factors": [
2081
+ 0
2082
+ ],
2083
+ "range_warp_specializes": [],
2084
+ "range_num_stages": [],
2085
+ "range_multi_buffers": [
2086
+ null
2087
+ ],
2088
+ "range_flattens": [
2089
+ null
2090
+ ],
2091
+ "load_eviction_policies": [
2092
+ "",
2093
+ "first",
2094
+ "last"
2095
+ ],
2096
+ "num_warps": 1,
2097
+ "num_stages": 5,
2098
+ "indexing": [
2099
+ "tensor_descriptor",
2100
+ "tensor_descriptor",
2101
+ "pointer",
2102
+ "pointer",
2103
+ "pointer"
2104
+ ],
2105
+ "atomic_indexing": [],
2106
+ "pid_type": "flat"
2107
+ }
2108
+ },
2109
+ {
2110
+ "key": {
2111
+ "intermediate_size": 12288,
2112
+ "group_size": 128,
2113
+ "num_tokens": 16384
2114
+ },
2115
+ "config": {
2116
+ "block_sizes": [
2117
+ 4
2118
+ ],
2119
+ "loop_orders": [
2120
+ [
2121
+ 1,
2122
+ 2,
2123
+ 0
2124
+ ]
2125
+ ],
2126
+ "l2_groupings": [
2127
+ 32
2128
+ ],
2129
+ "range_unroll_factors": [
2130
+ 0
2131
+ ],
2132
+ "range_warp_specializes": [],
2133
+ "range_num_stages": [],
2134
+ "range_multi_buffers": [
2135
+ null
2136
+ ],
2137
+ "range_flattens": [
2138
+ null
2139
+ ],
2140
+ "load_eviction_policies": [
2141
+ "first",
2142
+ "",
2143
+ "last"
2144
+ ],
2145
+ "num_warps": 1,
2146
+ "num_stages": 1,
2147
+ "indexing": [
2148
+ "tensor_descriptor",
2149
+ "tensor_descriptor",
2150
+ "tensor_descriptor",
2151
+ "tensor_descriptor",
2152
+ "pointer"
2153
+ ],
2154
+ "atomic_indexing": [],
2155
+ "pid_type": "flat"
2156
+ }
2157
+ },
2158
+ {
2159
+ "key": {
2160
+ "intermediate_size": 25600,
2161
+ "group_size": 128,
2162
+ "num_tokens": 16384
2163
+ },
2164
+ "config": {
2165
+ "block_sizes": [
2166
+ 16
2167
+ ],
2168
+ "loop_orders": [
2169
+ [
2170
+ 1,
2171
+ 2,
2172
+ 0
2173
+ ]
2174
+ ],
2175
+ "l2_groupings": [
2176
+ 16
2177
+ ],
2178
+ "range_unroll_factors": [
2179
+ 0
2180
+ ],
2181
+ "range_warp_specializes": [],
2182
+ "range_num_stages": [],
2183
+ "range_multi_buffers": [
2184
+ null
2185
+ ],
2186
+ "range_flattens": [
2187
+ null
2188
+ ],
2189
+ "load_eviction_policies": [
2190
+ "last",
2191
+ "",
2192
+ ""
2193
+ ],
2194
+ "num_warps": 4,
2195
+ "num_stages": 2,
2196
+ "indexing": [
2197
+ "tensor_descriptor",
2198
+ "pointer",
2199
+ "pointer",
2200
+ "tensor_descriptor",
2201
+ "tensor_descriptor"
2202
+ ],
2203
+ "atomic_indexing": [],
2204
+ "pid_type": "flat"
2205
+ }
2206
+ }
2207
+]
\ No newline at end of file
vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py
new
+252
@@ -0,0 +1,252 @@
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 input
38
+ # property combination. Currently, dtypes are fixed. We need optimization to
39
+ # bucket/skip some combinations
40
+ num_tokens_list = [1, 2, 4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048, 4096, 8192]
41
+ intermediate_size_list = [6144, 12288, 25600]
42
+
43
+ in_dtype: torch.dtype = torch.bfloat16
44
+ out_dtype: torch.dtype = current_platform.fp8_dtype()
45
+ scale_dtype: torch.dtype = torch.float32
46
+ group_size_list = [128]
47
+ inputs = {}
48
+ for intermediate_size, group_size, num_tokens in product(
49
+ intermediate_size_list, group_size_list, num_tokens_list
50
+ ):
51
+ input = torch.randn(
52
+ num_tokens, 2 * intermediate_size, device="cuda", dtype=in_dtype
53
+ )
54
+ result = torch.empty(
55
+ num_tokens, intermediate_size, device=input.device, dtype=out_dtype
56
+ )
57
+ scale = torch.empty(
58
+ (num_tokens, intermediate_size // group_size),
59
+ device=input.device,
60
+ dtype=scale_dtype,
61
+ )
62
+ scale_ub = torch.mean(input).to(scale_dtype)
63
+
64
+ config_key = CaseKey(
65
+ {
66
+ "intermediate_size": intermediate_size,
67
+ "group_size": group_size,
68
+ "num_tokens": num_tokens,
69
+ }
70
+ )
71
+ inputs[config_key] = (result, input, scale, group_size, scale_ub, False)
72
+
73
+ return inputs
74
+
75
+
76
+_pick_cache: dict[tuple[int, int, int], CaseKey | None] = {}
77
+
78
+
79
+def pick_config(args: tuple[Any, ...], config_keys: list[CaseKey]) -> CaseKey | None:
80
+ """Pick the best pre-tuned config for the given input shape.
81
+
82
+ Selection strategy:
83
+ 1. Find the closest intermediate_size among available configs
84
+ (exact match preferred).
85
+ 2. Find the closest group_size among available configs
86
+ (exact match preferred).
87
+ 3. Among the num_tokens values tuned for that intermediate_size and group_size,
88
+ pick the smallest num_tokens >= the input's num_tokens. If the input is
89
+ larger than all available num_tokens, fall back to the largest.
90
+ """
91
+
92
+ if not config_keys:
93
+ return None
94
+
95
+ result, _, _, group_size, *_ = args
96
+ num_tokens, intermediate_size = result.shape
97
+
98
+ cache_key = (num_tokens, group_size, intermediate_size)
99
+ cached = _pick_cache.get(cache_key)
100
+ if cached is not None:
101
+ return cached
102
+
103
+ configs: dict[int, dict[int, list[int]]] = {}
104
+ for key in config_keys:
105
+ if key.is_default():
106
+ continue
107
+ configs.setdefault(key["intermediate_size"], {}).setdefault(
108
+ key["group_size"], []
109
+ ).append(key["num_tokens"])
110
+
111
+ if not configs:
112
+ return None
113
+
114
+ best_intermediate_size = min(configs, key=lambda s: abs(s - intermediate_size))
115
+ best_group_size = min(
116
+ configs[best_intermediate_size], key=lambda s: abs(s - group_size)
117
+ )
118
+ available_num_tokens = sorted(configs[best_intermediate_size][best_group_size])
119
+ best_num_tokens = next(
120
+ (n for n in available_num_tokens if n >= num_tokens), available_num_tokens[-1]
121
+ )
122
+
123
+ result = CaseKey(
124
+ {
125
+ "intermediate_size": best_intermediate_size,
126
+ "group_size": best_group_size,
127
+ "num_tokens": best_num_tokens,
128
+ }
129
+ )
130
+ _pick_cache[cache_key] = result
131
+ return result
132
+
133
+
134
+def fake_impl(
135
+ out: torch.Tensor, # [num_tokens, intermediate_size]
136
+ input: torch.Tensor, # [num_tokens, 2 * intermediate_size]
137
+ scales: torch.Tensor, # [num_tokens, groups_per_row]
138
+ group_size: int,
139
+ scale_ub: torch.Tensor | None = None, # scalar tensor
140
+ is_scale_transposed: bool = False,
141
+) -> None:
142
+ return
143
+
144
+
145
+def baseline(
146
+ out: torch.Tensor, # [num_tokens, intermediate_size]
147
+ input: torch.Tensor, # [num_tokens, 2 * intermediate_size]
148
+ scales: torch.Tensor, # [num_tokens, groups_per_row]
149
+ group_size: int,
150
+ scale_ub: torch.Tensor | None = None, # scalar tensor
151
+ is_scale_transposed: bool = False,
152
+) -> None:
153
+ torch.ops._C.silu_and_mul_per_block_quant(
154
+ out, input, scales, group_size, scale_ub, is_scale_transposed
155
+ )
156
+
157
+
158
+@register_kernel(
159
+ mutates_args=["out", "scales"],
160
+ config_picker=pick_config,
161
+ input_generator=generate_inputs,
162
+ fake_impl=fake_impl,
163
+ helion_settings=helion.Settings(
164
+ autotune_baseline_fn=baseline,
165
+ ignore_warnings=[helion.exc.TensorOperationInWrapper],
166
+ ),
167
+) # type: ignore[misc]
168
+def silu_and_mul_per_block_quant(
169
+ out: torch.Tensor, # [num_tokens, intermediate_size]
170
+ input: torch.Tensor, # [num_tokens, 2 * intermediate_size]
171
+ scales: torch.Tensor, # [num_tokens, groups_per_row]
172
+ group_size: int,
173
+ scale_ub: torch.Tensor | None = None, # scalar tensor
174
+ is_scale_transposed: bool = False, # dummy
175
+) -> None:
176
+ # This code assumes batch_dim and num_tokens are flattened
177
+ assert input.ndim == 2
178
+ num_tokens, two_intermediate_size = input.shape
179
+ hl.specialize(two_intermediate_size)
180
+
181
+ assert two_intermediate_size % 2 == 0
182
+ intermediate_size = two_intermediate_size // 2
183
+
184
+ assert out.shape[0] == num_tokens
185
+ assert out.shape[1] == intermediate_size
186
+ fp8_dtype = get_fp8_dtype()
187
+ assert out.dtype in [fp8_dtype, torch.int8]
188
+
189
+ if scale_ub is not None:
190
+ assert out.dtype == fp8_dtype
191
+ assert scale_ub.dtype == torch.float32
192
+
193
+ assert scales.ndim == 2 and scales.dtype == torch.float32
194
+
195
+ assert scales.shape[0] == num_tokens
196
+ groups_per_row = scales.shape[1]
197
+ hl.specialize(groups_per_row)
198
+ assert (
199
+ intermediate_size % group_size == 0
200
+ and intermediate_size // group_size == groups_per_row
201
+ )
202
+
203
+ assert group_size in [64, 128]
204
+ hl.specialize(group_size)
205
+
206
+ assert input.stride()[-1] == 1
207
+ assert out.stride()[-1] == 1
208
+
209
+ quant_dtype = out.dtype
210
+ qtype_traits_min: int | float
211
+ qtype_traits_max: int | float
212
+ if quant_dtype == torch.int8:
213
+ qtype_traits_min, qtype_traits_max = get_int8_min_max()
214
+ min_scaling_factor = get_int8_min_scaling_factor()
215
+ else:
216
+ qtype_traits_min, qtype_traits_max = get_fp8_min_max()
217
+ min_scaling_factor = 1.0 / (qtype_traits_max * 512.0)
218
+
219
+ qtype_max = float(qtype_traits_max)
220
+
221
+ input = input.view(num_tokens, -1, group_size)
222
+ out = out.view(num_tokens, -1, group_size)
223
+
224
+ for tile_m, tile_gn, tile_n in hl.tile(
225
+ [num_tokens, groups_per_row, group_size], block_size=[1, None, group_size]
226
+ ):
227
+ x_a_blk = input[tile_m, tile_gn, tile_n].to(torch.float32)
228
+ x_b_blk = hl.load(
229
+ input,
230
+ [tile_m, tile_gn.index + groups_per_row, tile_n],
231
+ extra_mask=(tile_gn.index + groups_per_row < 2 * groups_per_row)[
232
+ None, :, None
233
+ ],
234
+ ).to(torch.float32)
235
+ x_blk = x_a_blk * torch.sigmoid(x_a_blk) * x_b_blk
236
+ s_blk = torch.amax(torch.abs(x_blk), dim=-1).to(torch.float32)
237
+
238
+ if scale_ub is not None:
239
+ scale_ub_s = hl.load(scale_ub, [])
240
+ s_blk = s_blk.clamp(max=scale_ub_s)
241
+ s_blk = s_blk * (1.0 / qtype_max)
242
+ s_blk = s_blk.clamp(min=min_scaling_factor)
243
+
244
+ scales[tile_m, tile_gn] = s_blk
245
+ if quant_dtype == torch.int8:
246
+ y_blk = (x_blk * (1.0 / s_blk[:, :, None])).round()
247
+ else:
248
+ y_blk = x_blk / s_blk[:, :, None]
249
+
250
+ out[tile_m, tile_gn, tile_n] = y_blk.clamp(
251
+ qtype_traits_min, qtype_traits_max
252
+ ).to(out.dtype)