[Kernel][Helion] Add Helion kernel benchmark script (#48512)
Signed-off-by: Sean Chen <seachen@redhat.com>
Xiaohong (Sean) Chen committed
Jul 15, 2026 at 11:43 UTC
6036bf110a9f4d313c38c257ddcf5b8ba4210383
6 files changed
+574
-28
scripts/benchmark_helion_kernels.py
new
+457
@@ -0,0 +1,457 @@
1
+#!/usr/bin/env python3
2
+# SPDX-License-Identifier: Apache-2.0
3
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
4
+"""
5
+Benchmark a registered Helion kernel against a baseline.
6
+
7
+For each input case produced by the kernel's registered input generator, this
8
+measures the latency of the Helion kernel and a chosen baseline, then reports
9
+the speedup.
10
+
11
+Two baselines are supported (``--baseline``):
12
+
13
+- ``autotune`` (default): the kernel's autotuning baseline
14
+ (``helion_settings.autotune_baseline_fn``), wrapped in ``torch.compile``
15
+ (inductor). This is the native-torch reference used by kernel autotuning
16
+ and correctness unit tests.
17
+- ``cuda``: the corresponding hand-written CUDA op (``torch.ops._C.*``). The
18
+ mapping from Helion kernel name to CUDA op lives in ``CUDA_BASELINE_OPS``
19
+ below. Every Helion kernel shares the same argument interface as its CUDA
20
+ counterpart, so inputs are forwarded verbatim.
21
+
22
+Usage:
23
+ # List available kernels
24
+ python scripts/benchmark_helion_kernels.py --list
25
+
26
+ # Benchmark a kernel against the autotune baseline (default)
27
+ python scripts/benchmark_helion_kernels.py --kernel per_token_group_fp8_quant
28
+
29
+ # Benchmark against the CUDA baseline
30
+ python scripts/benchmark_helion_kernels.py --kernel per_token_group_fp8_quant \\
31
+ --baseline cuda
32
+
33
+ # Disable CUDA graph capture and save results
34
+ python scripts/benchmark_helion_kernels.py --kernel per_token_group_fp8_quant \\
35
+ --no-cudagraph --output results.json
36
+"""
37
+
38
+import argparse
39
+import copy
40
+import gc
41
+import json
42
+import statistics
43
+import sys
44
+from collections.abc import Callable
45
+from dataclasses import asdict, dataclass
46
+
47
+import torch
48
+
49
+from vllm.triton_utils import triton
50
+
51
+try:
52
+ from vllm.benchmarks.lib.utils import default_vllm_config
53
+ from vllm.kernels.helion import get_kernel_by_name, get_registered_kernels
54
+ from vllm.kernels.helion.ops import import_all_kernels
55
+ from vllm.logger import init_logger
56
+ from vllm.utils.import_utils import has_helion
57
+except ImportError as e:
58
+ print(f"Error importing vLLM: {e}")
59
+ print("Please ensure vLLM is installed and in your Python path")
60
+ sys.exit(1)
61
+
62
+logger = init_logger("vllm.scripts.benchmark_helion_kernels")
63
+
64
+
65
+# Maps a Helion kernel name to the CUDA op (attribute on ``torch.ops._C``) that
66
+# implements the same operation. Helion kernels share the CUDA op's argument
67
+# interface, so the kernel's input tuple is forwarded verbatim. Add an entry
68
+# here when introducing a new kernel whose baseline should be the CUDA op.
69
+CUDA_BASELINE_OPS: dict[str, str] = {
70
+ "dynamic_per_token_scaled_fp8_quant": "dynamic_per_token_scaled_fp8_quant",
71
+ "fused_qk_norm_rope": "fused_qk_norm_rope",
72
+ "per_token_group_fp8_quant": "per_token_group_fp8_quant",
73
+ "rms_norm_dynamic_per_token_quant": "rms_norm_dynamic_per_token_quant",
74
+ "rms_norm_per_block_quant": "rms_norm_per_block_quant",
75
+ "silu_and_mul_per_block_quant": "silu_and_mul_per_block_quant",
76
+ "scaled_mm": "cutlass_scaled_mm",
77
+}
78
+
79
+# torch.compile options for the torch baseline, mirroring how these kernels are
80
+# compiled inside vLLM.
81
+_TORCH_COMPILE_OPTIONS: dict[str, bool] = {
82
+ "enable_auto_functionalized_v2": False,
83
+ "size_asserts": False,
84
+ "alignment_asserts": False,
85
+ "scalar_asserts": False,
86
+ "combo_kernels": True,
87
+ "benchmark_combo_kernel": True,
88
+}
89
+
90
+
91
+@dataclass
92
+class Row:
93
+ case: str
94
+ baseline_ms: float
95
+ kernel_ms: float
96
+ speedup_x: float
97
+
98
+
99
+def print_table(rows: list[Row]) -> None:
100
+ headers = ["case", "baseline_ms", "kernel_ms", "speedup(x)"]
101
+
102
+ data = [
103
+ [
104
+ r.case,
105
+ f"{r.baseline_ms:.3f}",
106
+ f"{r.kernel_ms:.3f}",
107
+ f"{r.speedup_x:.3f}",
108
+ ]
109
+ for r in rows
110
+ ]
111
+
112
+ cols = list(zip(*([headers] + data)))
113
+ widths = [max(len(cell) for cell in col) for col in cols]
114
+
115
+ def fmt(row: list[str]) -> str:
116
+ return " | ".join(cell.ljust(w) for cell, w in zip(row, widths))
117
+
118
+ print(fmt(headers))
119
+ print("-+-".join("-" * w for w in widths))
120
+ for row in data:
121
+ print(fmt(row))
122
+
123
+
124
+def list_kernels() -> None:
125
+ kernels = get_registered_kernels()
126
+
127
+ if not kernels:
128
+ print("No Helion kernels found in registry.")
129
+ return
130
+
131
+ print("Available Helion kernels:")
132
+ print("=" * 50)
133
+ for name in sorted(kernels.keys()):
134
+ cuda = CUDA_BASELINE_OPS.get(name)
135
+ suffix = "" if cuda else " (no CUDA baseline mapping)"
136
+ print(f" {name}{suffix}")
137
+ print(f"\nTotal: {len(kernels)} kernels")
138
+
139
+
140
+def check_requirements() -> bool:
141
+ if not torch.accelerator.is_available():
142
+ logger.error("CUDA is not available. Helion benchmarking requires GPU.")
143
+ return False
144
+ if not has_helion():
145
+ logger.error("Helion is not installed. Please install Helion package.")
146
+ return False
147
+ return True
148
+
149
+
150
+def make_cuda_baseline(kernel_name: str) -> Callable:
151
+ """Return a callable invoking the CUDA op mapped to ``kernel_name``.
152
+
153
+ The Helion kernel and its CUDA op share the same argument interface, so the
154
+ input tuple is forwarded verbatim.
155
+ """
156
+ cuda_op_name = CUDA_BASELINE_OPS.get(kernel_name)
157
+ if cuda_op_name is None:
158
+ logger.error(
159
+ "No CUDA baseline mapping for kernel '%s'. Add an entry to "
160
+ "CUDA_BASELINE_OPS in %s (mapping the kernel name to its "
161
+ "torch.ops._C.<op> name), or benchmark with --baseline torch.",
162
+ kernel_name,
163
+ __file__,
164
+ )
165
+ sys.exit(1)
166
+
167
+ cuda_op = getattr(torch.ops._C, cuda_op_name, None)
168
+ if cuda_op is None:
169
+ logger.error(
170
+ "torch.ops._C.%s is not available. Ensure the vLLM C extension is "
171
+ "built and loaded.",
172
+ cuda_op_name,
173
+ )
174
+ sys.exit(1)
175
+
176
+ return cuda_op
177
+
178
+
179
+def make_autotune_baseline(kernel_name: str) -> Callable:
180
+ """Return the kernel's autotune baseline wrapped in ``torch.compile``.
181
+
182
+ The baseline is the native-torch reference the kernel is tuned against,
183
+ registered via ``helion_settings.autotune_baseline_fn``.
184
+ """
185
+ wrapper = get_kernel_by_name(kernel_name)
186
+ settings = wrapper.helion_settings
187
+ baseline_fn = getattr(settings, "autotune_baseline_fn", None)
188
+ if baseline_fn is None:
189
+ logger.error(
190
+ "Kernel '%s' has no autotune_baseline_fn in its helion_settings, so "
191
+ "the 'autotune' baseline is unavailable. Register one via "
192
+ "register_kernel(..., helion_settings=helion.Settings("
193
+ "autotune_baseline_fn=...)), or benchmark with --baseline cuda.",
194
+ kernel_name,
195
+ )
196
+ sys.exit(1)
197
+
198
+ return torch.compile(
199
+ baseline_fn,
200
+ fullgraph=True,
201
+ dynamic=False,
202
+ backend="inductor",
203
+ options=_TORCH_COMPILE_OPTIONS,
204
+ )
205
+
206
+
207
+def cleanup_gpu_resources() -> None:
208
+ try:
209
+ torch.accelerator.empty_cache()
210
+ gc.collect()
211
+ if hasattr(torch, "_dynamo"):
212
+ torch._dynamo.reset()
213
+ torch.accelerator.synchronize()
214
+ except Exception as e:
215
+ logger.warning("Failed to cleanup GPU resources: %s", e)
216
+
217
+
218
+_REDUCERS: dict[str, Callable[[list[float]], float]] = {
219
+ "min": min,
220
+ "max": max,
221
+ "mean": statistics.fmean,
222
+ "median": statistics.median,
223
+}
224
+
225
+
226
+def _reduce(times: list[float], return_mode: str) -> float:
227
+ return _REDUCERS[return_mode](times)
228
+
229
+
230
+def do_bench_cudagraph_l2_clear(
231
+ fn: Callable, rep: int = 100, return_mode: str = "mean"
232
+) -> float:
233
+ """CUDA-graph benchmark that flushes the L2 cache before every call.
234
+
235
+ ``triton.testing.do_bench_cudagraph`` captures back-to-back kernel launches
236
+ with a warm L2 cache, which over-estimates performance for memory-bound
237
+ kernels. This clears L2 (via triton's benchmark cache buffer) before each
238
+ call and subtracts the isolated cache-clear cost from the measured time.
239
+
240
+ Adapted from tritonbench's ``_do_bench_cudagraph_with_cache_clear`` using
241
+ only triton/torch primitives so no extra dependency is introduced.
242
+ """
243
+ cache = triton.runtime.driver.active.get_empty_cache_for_benchmark()
244
+ clear_cache = cache.zero_
245
+
246
+ s = torch.Stream()
247
+ with s:
248
+ clear_cache()
249
+ fn()
250
+
251
+ start_event = torch.Event(enable_timing=True)
252
+ end_event = torch.Event(enable_timing=True)
253
+ start_event.record()
254
+ for _ in range(5):
255
+ clear_cache()
256
+ fn()
257
+ end_event.record()
258
+ torch.accelerator.synchronize()
259
+ estimate_ms = start_event.elapsed_time(end_event) / 5
260
+ n_repeat = 1000 if estimate_ms == 0 else max(1, int(rep / estimate_ms))
261
+
262
+ graph = torch.cuda.CUDAGraph()
263
+ with torch.cuda.graph(graph):
264
+ for _ in range(n_repeat):
265
+ clear_cache()
266
+ fn()
267
+
268
+ clear_graph = torch.cuda.CUDAGraph()
269
+ with torch.cuda.graph(clear_graph):
270
+ for _ in range(n_repeat):
271
+ clear_cache()
272
+ torch.accelerator.synchronize()
273
+
274
+ kernel_times = []
275
+ for _ in range(10):
276
+ start_event = torch.Event(enable_timing=True)
277
+ end_event = torch.Event(enable_timing=True)
278
+ start_event.record()
279
+ clear_graph.replay()
280
+ end_event.record()
281
+ torch.accelerator.synchronize()
282
+ clear_ms = start_event.elapsed_time(end_event) / n_repeat
283
+
284
+ start_event = torch.Event(enable_timing=True)
285
+ end_event = torch.Event(enable_timing=True)
286
+ start_event.record()
287
+ graph.replay()
288
+ end_event.record()
289
+ torch.accelerator.synchronize()
290
+ total_ms = start_event.elapsed_time(end_event) / n_repeat
291
+
292
+ kernel_times.append(total_ms - clear_ms)
293
+
294
+ return _reduce(kernel_times, return_mode)
295
+
296
+
297
+@torch.inference_mode()
298
+def benchmark(
299
+ kernel_name: str,
300
+ baseline_fn: Callable,
301
+ repeat: int,
302
+ cudagraph: bool,
303
+ return_mode: str,
304
+) -> list[Row]:
305
+ kernel = get_kernel_by_name(kernel_name)
306
+ # do_bench already flushes L2 per call; do_bench_cudagraph does not, so use
307
+ # the cache-clearing variant to avoid warm-L2 over-estimates.
308
+ benchmark_fn = do_bench_cudagraph_l2_clear if cudagraph else triton.testing.do_bench
309
+
310
+ inputs_dict = kernel.get_inputs()
311
+ rows: list[Row] = []
312
+
313
+ for key, inputs in inputs_dict.items():
314
+ logger.info("Benchmarking case %s", key)
315
+
316
+ # Kernels may mutate their inputs in place; give each side its own copy.
317
+ kernel_inputs = copy.deepcopy(inputs)
318
+ baseline_inputs = copy.deepcopy(inputs)
319
+
320
+ kernel_latency = benchmark_fn(
321
+ lambda kernel_inputs=kernel_inputs: kernel(*kernel_inputs),
322
+ rep=repeat,
323
+ return_mode=return_mode,
324
+ )
325
+ baseline_latency = benchmark_fn(
326
+ lambda baseline_inputs=baseline_inputs: baseline_fn(*baseline_inputs),
327
+ rep=repeat,
328
+ return_mode=return_mode,
329
+ )
330
+
331
+ rows.append(
332
+ Row(
333
+ case=str(key),
334
+ baseline_ms=baseline_latency,
335
+ kernel_ms=kernel_latency,
336
+ speedup_x=baseline_latency / kernel_latency,
337
+ )
338
+ )
339
+ cleanup_gpu_resources()
340
+
341
+ return rows
342
+
343
+
344
+def main() -> None:
345
+ parser = argparse.ArgumentParser(
346
+ description="Benchmark a Helion kernel against a baseline",
347
+ formatter_class=argparse.RawDescriptionHelpFormatter,
348
+ epilog=__doc__.split("Usage:")[1] if "Usage:" in __doc__ else "",
349
+ )
350
+ parser.add_argument(
351
+ "--list",
352
+ action="store_true",
353
+ help="List available Helion kernels and exit",
354
+ )
355
+ parser.add_argument(
356
+ "--kernel",
357
+ type=str,
358
+ help="Name of the single Helion kernel to benchmark",
359
+ )
360
+ parser.add_argument(
361
+ "--repeat",
362
+ type=int,
363
+ default=100,
364
+ help="Number of benchmark repetitions (default: 100)",
365
+ )
366
+ parser.add_argument(
367
+ "--no-cudagraph",
368
+ dest="cudagraph",
369
+ action="store_false",
370
+ help="Disable CUDA graph mode (enabled by default)",
371
+ )
372
+ parser.add_argument(
373
+ "--baseline",
374
+ choices=["cuda", "autotune"],
375
+ default="autotune",
376
+ help=(
377
+ "Baseline to compare against: 'autotune' uses the kernel's "
378
+ "autotune_baseline_fn under torch.compile; 'cuda' uses the mapped "
379
+ "torch.ops._C op (default: autotune)"
380
+ ),
381
+ )
382
+ parser.add_argument(
383
+ "--return-mode",
384
+ choices=["min", "max", "mean", "median"],
385
+ default="mean",
386
+ help="Statistic to report from the benchmark samples (default: mean)",
387
+ )
388
+ parser.add_argument(
389
+ "--output",
390
+ type=str,
391
+ help="Path to save benchmark results as JSON (default: log only)",
392
+ )
393
+
394
+ args = parser.parse_args()
395
+
396
+ import_all_kernels()
397
+
398
+ if args.list:
399
+ list_kernels()
400
+ return
401
+
402
+ if not args.kernel:
403
+ parser.error("--kernel is required (or use --list to see available kernels)")
404
+
405
+ kernels = get_registered_kernels()
406
+ if args.kernel not in kernels:
407
+ logger.error("Kernel '%s' not found in registry.", args.kernel)
408
+ logger.error("Available kernels: %s", sorted(kernels.keys()))
409
+ sys.exit(1)
410
+
411
+ wrapper = kernels[args.kernel]
412
+ if wrapper._disabled:
413
+ logger.error(
414
+ "Kernel '%s' is disabled: %s",
415
+ args.kernel,
416
+ wrapper._disabled_reason,
417
+ )
418
+ sys.exit(1)
419
+
420
+ if not check_requirements():
421
+ sys.exit(1)
422
+
423
+ with default_vllm_config():
424
+ if args.baseline == "cuda":
425
+ baseline_fn = make_cuda_baseline(args.kernel)
426
+ else:
427
+ baseline_fn = make_autotune_baseline(args.kernel)
428
+
429
+ rows = benchmark(
430
+ args.kernel,
431
+ baseline_fn,
432
+ args.repeat,
433
+ args.cudagraph,
434
+ args.return_mode,
435
+ )
436
+
437
+ print_table(rows)
438
+
439
+ if args.output:
440
+ with open(args.output, "w") as f:
441
+ json.dump(
442
+ {
443
+ "kernel": args.kernel,
444
+ "baseline": args.baseline,
445
+ "cudagraph": args.cudagraph,
446
+ "repeat": args.repeat,
447
+ "return_mode": args.return_mode,
448
+ "results": [asdict(r) for r in rows],
449
+ },
450
+ f,
451
+ indent=2,
452
+ )
453
+ logger.info("Saved results to %s", args.output)
454
+
455
+
456
+if __name__ == "__main__":
457
+ main()
vllm/kernels/helion/ops/dynamic_per_token_scaled_fp8_quant.py
+12
-1
@@ -109,7 +109,18 @@ def baseline(
109
scale: torch.Tensor, # [num_tokens, 1]
110
scale_ub: torch.Tensor | None = None, # scalar tensor
111
) -> None:
112
- torch.ops._C.dynamic_per_token_scaled_fp8_quant(result, input, scale, scale_ub)
112
+ fp8_min, fp8_max = get_fp8_min_max()
113
+ min_scaling_factor = 1.0 / (fp8_max * 512.0)
114
+
115
+ x = input.to(torch.float32)
116
+ s = torch.amax(torch.abs(x), dim=-1, keepdim=True)
117
+ if scale_ub is not None:
118
+ s = s.clamp(max=scale_ub)
119
+ s = (s * (1.0 / fp8_max)).clamp(min=min_scaling_factor)
120
+ y = (x / s).clamp(fp8_min, fp8_max)
121
+
122
+ scale.copy_(s)
123
+ result.copy_(y.to(result.dtype))
124
125
126
# Overwrite autotune_baseline_atol and autotune_baseline_rtol
vllm/kernels/helion/ops/per_token_group_fp8_quant.py
+11
-12
@@ -164,18 +164,17 @@ def baseline(
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
- )
167
+ num_tokens, hidden_size = input.shape
168
+ groups_per_row = hidden_size // group_size
169
+
170
+ x = input.view(num_tokens, groups_per_row, group_size).to(torch.float32)
171
+ s = torch.clamp(torch.amax(torch.abs(x), dim=-1), min=eps) / fp8_max
172
+ if scale_ue8m0:
173
+ s = torch.exp2(torch.ceil(torch.log2(s)))
174
+ y = torch.clamp(x / s[:, :, None], fp8_min, fp8_max)
175
+
176
+ output_s.copy_(s)
177
+ output_q.copy_(y.view(num_tokens, hidden_size).to(output_q.dtype))
178
179
180
@register_kernel(
vllm/kernels/helion/ops/rms_norm_dynamic_per_token_quant.py
+31
-3
@@ -131,9 +131,37 @@ def baseline(
131
scale_ub: torch.Tensor | None = None, # []
132
residual: torch.Tensor | None = None, # [num_tokens, hidden_size]
133
) -> None:
134
- torch.ops._C.rms_norm_dynamic_per_token_quant(
135
- result, input, weight, scale, epsilon, scale_ub, residual
136
- )
134
+ _, hidden_size = input.shape
135
+ quant_dtype = result.dtype
136
+ qtype_min: int | float
137
+ qtype_max: int | float
138
+
139
+ if quant_dtype == torch.int8:
140
+ qtype_min, qtype_max = get_int8_min_max()
141
+ min_scaling_factor = get_int8_min_scaling_factor()
142
+ else:
143
+ qtype_min, qtype_max = get_fp8_min_max()
144
+ min_scaling_factor = 1.0 / (qtype_max * 512.0)
145
+
146
+ x = input.to(torch.float32)
147
+ if residual is not None:
148
+ x = x + residual
149
+ residual.copy_(x.to(residual.dtype))
150
+
151
+ rms = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + epsilon)
152
+ x_norm = (x * rms).to(input.dtype) * weight
153
+
154
+ s = torch.amax(torch.abs(x_norm), dim=-1, keepdim=True).to(torch.float32)
155
+ if scale_ub is not None:
156
+ s = s.clamp(max=scale_ub)
157
+ s = (s * (1.0 / qtype_max)).clamp(min=min_scaling_factor)
158
+
159
+ y = x_norm / s
160
+ if quant_dtype == torch.int8:
161
+ y = y.round()
162
+
163
+ scale.copy_(s)
164
+ result.copy_(y.clamp(qtype_min, qtype_max).to(result.dtype))
165
166
167
# Overwrite autotune_baseline_atol and autotune_baseline_rtol
vllm/kernels/helion/ops/rms_norm_per_block_quant.py
+34
-10
@@ -169,16 +169,40 @@ def baseline(
169
group_size: int,
170
is_scale_transposed: bool,
171
) -> None:
172
- torch.ops._C.rms_norm_per_block_quant(
173
- result,
174
- input,
175
- weight,
176
- scale,
177
- epsilon,
178
- scale_ub,
179
- residual,
180
- group_size,
181
- is_scale_transposed,
172
+ num_tokens, hidden_size = input.shape
173
+ groups_per_row = hidden_size // group_size
174
+ quant_dtype = result.dtype
175
+ qtype_min: int | float
176
+ qtype_max: int | float
177
+
178
+ if quant_dtype == torch.int8:
179
+ qtype_min, qtype_max = get_int8_min_max()
180
+ min_scaling_factor = get_int8_min_scaling_factor()
181
+ else:
182
+ qtype_min, qtype_max = get_fp8_min_max()
183
+ min_scaling_factor = 1.0 / (qtype_max * 512.0)
184
+
185
+ x = input.to(torch.float32)
186
+ if residual is not None:
187
+ x = x + residual
188
+ residual.copy_(x.to(residual.dtype))
189
+
190
+ rms = torch.rsqrt(x.pow(2).mean(dim=-1, keepdim=True) + epsilon)
191
+ x_norm = (x * rms).to(input.dtype) * weight
192
+ x_grouped = x_norm.view(num_tokens, groups_per_row, group_size).to(torch.float32)
193
+
194
+ s = torch.amax(torch.abs(x_grouped), dim=-1).to(torch.float32)
195
+ if scale_ub is not None:
196
+ s = s.clamp(max=scale_ub)
197
+ s = (s * (1.0 / qtype_max)).clamp(min=min_scaling_factor)
198
+
199
+ y = x_grouped / s[:, :, None]
200
+ if quant_dtype == torch.int8:
201
+ y = y.round()
202
+
203
+ scale.copy_(s)
204
+ result.copy_(
205
+ y.clamp(qtype_min, qtype_max).view(num_tokens, hidden_size).to(result.dtype)
206
)
207
208
vllm/kernels/helion/ops/silu_and_mul_per_block_quant.py
+29
-2
@@ -13,6 +13,7 @@ from vllm.kernels.helion.utils import (
13
get_int8_min_scaling_factor,
14
)
15
from vllm.logger import init_logger
16
+from vllm.model_executor.layers.activation import SiluAndMul
17
from vllm.model_executor.layers.quantization.utils.quant_utils import (
18
get_fp8_min_max,
19
)
@@ -150,8 +151,34 @@ def baseline(
151
scale_ub: torch.Tensor | None = None, # scalar tensor
152
is_scale_transposed: bool = False,
153
) -> None:
153
- torch.ops._C.silu_and_mul_per_block_quant(
154
- out, input, scales, group_size, scale_ub, is_scale_transposed
154
+ num_tokens, intermediate_size = out.shape
155
+ groups_per_row = intermediate_size // group_size
156
+ quant_dtype = out.dtype
157
+ qtype_min: int | float
158
+ qtype_max: int | float
159
+
160
+ if quant_dtype == torch.int8:
161
+ qtype_min, qtype_max = get_int8_min_max()
162
+ min_scaling_factor = get_int8_min_scaling_factor()
163
+ else:
164
+ qtype_min, qtype_max = get_fp8_min_max()
165
+ min_scaling_factor = 1.0 / (qtype_max * 512.0)
166
+
167
+ act = SiluAndMul.forward_native(input.to(torch.float32))
168
+ x_grouped = act.view(num_tokens, groups_per_row, group_size)
169
+
170
+ s = torch.amax(torch.abs(x_grouped), dim=-1)
171
+ if scale_ub is not None:
172
+ s = s.clamp(max=scale_ub)
173
+ s = (s * (1.0 / qtype_max)).clamp(min=min_scaling_factor)
174
+
175
+ y = x_grouped / s[:, :, None]
176
+ if quant_dtype == torch.int8:
177
+ y = y.round()
178
+
179
+ scales.copy_(s)
180
+ out.copy_(
181
+ y.clamp(qtype_min, qtype_max).view(num_tokens, intermediate_size).to(out.dtype)
182
)
183
184