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