@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 19bee6d12d
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 12, 02:15 UTC

[Performance][DP/EP] Add silu_mul_per_token_group_quant_fp8_colmajor kernel (#29470)

Signed-off-by: Varun Sundar Rabindranath <vsundarr@redhat.com> Co-authored-by: Varun Sundar Rabindranath <vsundarr@redhat.com> Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com>

Varun Sundar Rabindranath committed Dec 3, 2025 at 13:04 UTC 19bee6d12d985c231b16374c99836376fc0c5706
4 files changed +496 -81
benchmarks/kernels/benchmark_2d_silu_mul_fp8_quant.py new
+244
@@ -0,0 +1,244 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from dataclasses import dataclass
5 +from enum import Enum
6 +from itertools import product
7 +from typing import Any
8 +
9 +import torch
10 +import torch.utils.benchmark as TBenchmark
11 +from torch.utils.benchmark import Measurement as TMeasurement
12 +
13 +from vllm.model_executor.layers.quantization.utils.fp8_utils import (
14 + _per_token_group_quant_fp8_colmajor,
15 + silu_mul_per_token_group_quant_fp8_colmajor,
16 +)
17 +from vllm.triton_utils import triton
18 +from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used
19 +
20 +from .utils import ArgPool, Bench, CudaGraphBenchParams
21 +
22 +GROUP_SIZE = 128
23 +FLOAT8_T = torch.float8_e4m3fn
24 +
25 +
26 +def print_timers(timers: list[TMeasurement], cuda_graph_nops: int):
27 + print(
28 + f"Note : The timings reported above is for {cuda_graph_nops} "
29 + "consecutive invocations of the benchmarking functions. "
30 + f"Please divide by {cuda_graph_nops} for single invocation "
31 + "timings."
32 + )
33 + compare = TBenchmark.Compare(timers)
34 + compare.print()
35 +
36 +
37 +class ImplType(Enum):
38 + SILU_MUL_PER_TOKEN_GROUP_QUANT_FP8_COLMAJOR = 1
39 + REFERENCE = 2
40 +
41 + def get_impl(self):
42 + if self == ImplType.SILU_MUL_PER_TOKEN_GROUP_QUANT_FP8_COLMAJOR:
43 + return silu_mul_per_token_group_quant_fp8_colmajor
44 + elif self == ImplType.REFERENCE:
45 + return reference
46 + raise ValueError(f"Unrecognized ImplType {self}")
47 +
48 +
49 +@dataclass
50 +class BenchmarkTensors:
51 + input: torch.Tensor
52 + output: torch.Tensor
53 +
54 + # Reference act output tensor
55 + ref_act_out: torch.Tensor
56 + ref_quant_out: torch.Tensor
57 +
58 + @staticmethod
59 + def make(T: int, N: int) -> "BenchmarkTensors":
60 + assert T % GROUP_SIZE == 0
61 + assert N % (GROUP_SIZE * 2) == 0
62 +
63 + input = torch.rand((T, N), dtype=torch.bfloat16, device="cuda")
64 +
65 + # silu_mul_per_token_group_quant_fp8_colmajor output.
66 + output = torch.rand((T, N // 2), dtype=torch.bfloat16, device="cuda").to(
67 + FLOAT8_T
68 + )
69 +
70 + # reference output.
71 + ref_act_out = torch.empty((T, N // 2), dtype=torch.bfloat16, device="cuda")
72 + ref_quant_out = torch.empty(
73 + (T, N // 2), dtype=torch.bfloat16, device="cuda"
74 + ).to(FLOAT8_T)
75 +
76 + return BenchmarkTensors(
77 + input=input,
78 + output=output,
79 + ref_act_out=ref_act_out,
80 + ref_quant_out=ref_quant_out,
81 + )
82 +
83 + @property
84 + def T(self):
85 + return self.input.size(0)
86 +
87 + @property
88 + def N(self):
89 + return self.input.size(1)
90 +
91 + def make_impl_kwargs(self, impl_type: ImplType) -> dict[str, Any]:
92 + if impl_type == ImplType.SILU_MUL_PER_TOKEN_GROUP_QUANT_FP8_COLMAJOR:
93 + return {
94 + "input": self.input,
95 + "output": self.output,
96 + "use_ue8m0": is_deep_gemm_e8m0_used(),
97 + }
98 + elif impl_type == ImplType.REFERENCE:
99 + return {
100 + "input": self.input,
101 + "act_out": self.ref_act_out,
102 + "quant_out": self.ref_quant_out,
103 + "use_ue8m0": is_deep_gemm_e8m0_used(),
104 + }
105 + raise ValueError(f"Unrecognized impl_type {impl_type}")
106 +
107 +
108 +def reference_quant(x: torch.Tensor, quant_out: torch.Tensor, use_ue8m0: bool):
109 + """
110 + Reference triton quant kernel from,
111 + vllm.model_executor.layers.quantization.utils.fp8_utils
112 + """
113 + assert quant_out.size() == x.size()
114 + # Allocate the scale tensor column-major format.
115 + shape = (x.shape[-1] // GROUP_SIZE,) + x.shape[:-1]
116 + x_q = quant_out
117 + x_s = torch.empty(shape, device=x.device, dtype=torch.float32).permute(-1, -2)
118 +
119 + M = x.numel() // GROUP_SIZE
120 + N = GROUP_SIZE
121 + BLOCK = triton.next_power_of_2(N)
122 + # heuristics for number of warps
123 + num_warps = min(max(BLOCK // 256, 1), 8)
124 + num_stages = 1
125 +
126 + finfo = torch.finfo(FLOAT8_T)
127 + fp8_min = finfo.min
128 + fp8_max = finfo.max
129 +
130 + _per_token_group_quant_fp8_colmajor[(M,)](
131 + x,
132 + x_q,
133 + x_s,
134 + GROUP_SIZE,
135 + x.shape[1],
136 + x.stride(0),
137 + x_s.stride(1),
138 + eps=1e-10,
139 + fp8_min=fp8_min,
140 + fp8_max=fp8_max,
141 + use_ue8m0=use_ue8m0,
142 + BLOCK=BLOCK,
143 + num_warps=num_warps,
144 + num_stages=num_stages,
145 + )
146 + return x_q, x_s
147 +
148 +
149 +def reference(
150 + input: torch.Tensor,
151 + act_out: torch.Tensor,
152 + quant_out: torch.Tensor,
153 + use_ue8m0: bool,
154 +) -> tuple[torch.Tensor, torch.Tensor]:
155 + torch.ops._C.silu_and_mul(act_out, input)
156 + return reference_quant(act_out, quant_out, use_ue8m0)
157 +
158 +
159 +def bench_impl(
160 + bench_tensors: list[BenchmarkTensors], impl_type: ImplType
161 +) -> TMeasurement:
162 + T = bench_tensors[0].T
163 + N = bench_tensors[0].N
164 +
165 + arg_pool_size = len(bench_tensors)
166 + kwargs_list = [bt.make_impl_kwargs(impl_type) for bt in bench_tensors]
167 +
168 + # warmup
169 + for kwargs in kwargs_list:
170 + impl_type.get_impl()(**kwargs)
171 + torch.cuda.synchronize()
172 +
173 + # Merge into a single kwargs and qualify arguments as ArgPool
174 + kwargs = {k: ArgPool([]) for k in kwargs_list[0]}
175 + for _kwargs in kwargs_list:
176 + for k, v in _kwargs.items():
177 + kwargs[k].values.append(v)
178 +
179 + cuda_graph_params = None
180 + cuda_graph_params = CudaGraphBenchParams(arg_pool_size)
181 + timer = None
182 + with Bench(
183 + cuda_graph_params,
184 + "silu-mul-quant",
185 + f"num_tokens={T}, N={N}",
186 + impl_type.name,
187 + impl_type.get_impl(),
188 + **kwargs,
189 + ) as bench:
190 + timer = bench.run()
191 + return timer
192 +
193 +
194 +def test_correctness(T: int, N: int):
195 + print(f"Testing num_tokens={T}, N={N} ...")
196 +
197 + bench_tensor = BenchmarkTensors.make(T, N)
198 +
199 + def output_from_impl(impl: ImplType) -> tuple[torch.Tensor, torch.Tensor]:
200 + return impl.get_impl()(**bench_tensor.make_impl_kwargs(impl))
201 +
202 + # reference output
203 + ref_out_q, ref_out_s = output_from_impl(ImplType.REFERENCE)
204 +
205 + # test ouptut
206 + out_q, out_s = output_from_impl(
207 + ImplType.SILU_MUL_PER_TOKEN_GROUP_QUANT_FP8_COLMAJOR
208 + )
209 +
210 + torch.testing.assert_close(ref_out_q.to(torch.float32), out_q.to(torch.float32))
211 + torch.testing.assert_close(ref_out_s, out_s)
212 +
213 +
214 +def run(Ts: list[int], Ns: list[int], arg_pool_size: int) -> list[TMeasurement]:
215 + timers = []
216 + for N, T in product(Ns, Ts):
217 + test_correctness(T, N)
218 +
219 + bench_tensors: list[BenchmarkTensors] = [
220 + BenchmarkTensors.make(T, N) for _ in range(arg_pool_size)
221 + ]
222 +
223 + silu_mul_quant_timer = bench_impl(
224 + bench_tensors, ImplType.SILU_MUL_PER_TOKEN_GROUP_QUANT_FP8_COLMAJOR
225 + )
226 + timers.append(silu_mul_quant_timer)
227 + reference_timer = bench_impl(bench_tensors, ImplType.REFERENCE)
228 + timers.append(reference_timer)
229 +
230 + print_timers(
231 + [silu_mul_quant_timer, reference_timer], cuda_graph_nops=arg_pool_size
232 + )
233 +
234 + print_timers(timers, cuda_graph_nops=arg_pool_size)
235 +
236 + return timers
237 +
238 +
239 +if __name__ == "__main__":
240 + T = [128 * i for i in range(1, 16)] + [2048 * i for i in range(1, 65)]
241 + N = [2048, 4096, 8192]
242 +
243 + print(f"T = {T}, N = {N}")
244 + run(T, N, arg_pool_size=8)
tests/kernels/moe/test_silu_mul_per_token_group_quant_fp8_colmajor.py new
+86
@@ -0,0 +1,86 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +import torch
6 +
7 +from vllm.model_executor.layers.quantization.utils.fp8_utils import (
8 + _per_token_group_quant_fp8_colmajor,
9 + silu_mul_per_token_group_quant_fp8_colmajor,
10 +)
11 +from vllm.platforms import current_platform
12 +from vllm.triton_utils import triton
13 +from vllm.utils.deep_gemm import is_deep_gemm_e8m0_used
14 +
15 +FLOAT8_DTYPE = torch.float8_e4m3fn
16 +GROUP_SIZE = 128
17 +
18 +
19 +def reference_quant(x: torch.Tensor, use_ue8m0: bool):
20 + """
21 + Reference triton quant kernel from,
22 + vllm.model_executor.layers.quantization.utils.fp8_utils
23 + """
24 +
25 + x_q = torch.empty_like(x, device=x.device, dtype=FLOAT8_DTYPE)
26 +
27 + # Allocate the scale tensor in column-major format.
28 + shape = (x.shape[-1] // GROUP_SIZE,) + x.shape[:-1]
29 + x_s = torch.empty(shape, device=x.device, dtype=torch.float32).permute(-1, -2)
30 +
31 + M = x.numel() // GROUP_SIZE
32 + N = GROUP_SIZE
33 + BLOCK = triton.next_power_of_2(N)
34 + # heuristics for number of warps
35 + num_warps = min(max(BLOCK // 256, 1), 8)
36 + num_stages = 1
37 +
38 + finfo = torch.finfo(FLOAT8_DTYPE)
39 + fp8_min = finfo.min
40 + fp8_max = finfo.max
41 +
42 + _per_token_group_quant_fp8_colmajor[(M,)](
43 + x,
44 + x_q,
45 + x_s,
46 + GROUP_SIZE,
47 + x.shape[1],
48 + x.stride(0),
49 + x_s.stride(1),
50 + eps=1e-10,
51 + fp8_min=fp8_min,
52 + fp8_max=fp8_max,
53 + use_ue8m0=use_ue8m0,
54 + BLOCK=BLOCK,
55 + num_warps=num_warps,
56 + num_stages=num_stages,
57 + )
58 + return x_q, x_s
59 +
60 +
61 +def reference(x: torch.Tensor, use_ue8m0: bool) -> tuple[torch.Tensor, torch.Tensor]:
62 + T, N = x.size()
63 + ref_act_out = torch.empty((T, N // 2), dtype=torch.bfloat16, device="cuda")
64 + torch.ops._C.silu_and_mul(ref_act_out, x)
65 + return reference_quant(ref_act_out, use_ue8m0)
66 +
67 +
68 +@pytest.mark.parametrize("T", [128, 256, 512])
69 +@pytest.mark.parametrize("N", [128 * 2, 256 * 2, 768 * 2, 2048 * 2, 7168 * 2])
70 +def test_silu_mul_fp8_quant_deep_gemm(T: int, N: int):
71 + current_platform.seed_everything(42)
72 +
73 + input = torch.rand((T, N), dtype=torch.bfloat16, device="cuda")
74 +
75 + use_ue8m0 = is_deep_gemm_e8m0_used()
76 +
77 + # Test
78 + output, output_scales = silu_mul_per_token_group_quant_fp8_colmajor(
79 + input, use_ue8m0=use_ue8m0
80 + )
81 +
82 + # Reference
83 + ref_output, ref_output_scales = reference(input, use_ue8m0)
84 +
85 + torch.testing.assert_close(output.to(torch.float32), ref_output.to(torch.float32))
86 + torch.testing.assert_close(output_scales, ref_output_scales)
vllm/model_executor/layers/fused_moe/deep_gemm_moe.py
+33 -81
@@ -2,9 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import torch
5 -from tqdm import tqdm
5
7 -import vllm.envs as env
6 import vllm.model_executor.layers.fused_moe.modular_kernel as mk
7 from vllm.logger import init_logger
8 from vllm.model_executor.layers.fused_moe.config import (
@@ -25,12 +23,12 @@ from vllm.model_executor.layers.fused_moe.topk_weight_and_reduce import (
23 from vllm.model_executor.layers.fused_moe.utils import _resize_cache
24 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
25 per_token_group_quant_fp8,
26 + silu_mul_per_token_group_quant_fp8_colmajor,
27 )
28 from vllm.utils.deep_gemm import (
29 get_mk_alignment_for_contiguous_layout,
30 m_grouped_fp8_gemm_nt_contiguous,
31 )
33 -from vllm.utils.func_utils import run_once
32 from vllm.utils.import_utils import has_deep_gemm
33
34 logger = init_logger(__name__)
@@ -108,70 +106,6 @@ def _valid_deep_gemm(
106 return True
107
108
111 -@run_once
112 -def warmup_deepgemm_gg_contiguous_kernels(
113 - w1: torch.Tensor,
114 - w2: torch.Tensor,
115 - w1_scale: torch.Tensor,
116 - w2_scale: torch.Tensor,
117 - num_topk: int,
118 -):
119 - """
120 - DeepGemm JITs the grouped-gemm kernels. The JIT'ing happens based on the
121 - input tensor shapes. In this function, we construct all possible input
122 - tensor shapes so all the kernels are JIT'ed and cached.
123 - Note that this warmup is expected to happen during the model profile
124 - call and not during actual model inference.
125 - """
126 -
127 - assert w1.size(0) == w2.size(0), "w1 and w2 must have the same number of experts"
128 -
129 - block_m = get_mk_alignment_for_contiguous_layout()[0]
130 - num_experts = w1.size(0)
131 - device = w1.device
132 -
133 - # This is the maximum GroupedGemm M size that we expect to run
134 - # the grouped_gemm with.
135 - MAX_M = compute_aligned_M(
136 - env.VLLM_FUSED_MOE_CHUNK_SIZE,
137 - num_topk,
138 - num_experts,
139 - block_m,
140 - expert_tokens_meta=None,
141 - )
142 - # Distribute expert-ids evenly.
143 - MAX_BLOCKS = MAX_M // block_m
144 - expert_ids_block = torch.randint(
145 - low=0, high=num_experts, size=(MAX_BLOCKS,), device=device, dtype=torch.int32
146 - )
147 - expert_ids = torch.repeat_interleave(expert_ids_block, block_m, dim=0)
148 -
149 - def _warmup(w: torch.Tensor, w_scale: torch.Tensor):
150 - _, n, k = w.size()
151 - a1q = torch.empty((MAX_M, k), device=device).to(torch.float8_e4m3fn)
152 - a1q_scales = torch.empty(
153 - (MAX_M, k // block_m), device=device, dtype=torch.float32
154 - )
155 - out = torch.empty((MAX_M, n), device=device, dtype=torch.bfloat16)
156 -
157 - pbar = tqdm(
158 - total=MAX_BLOCKS, desc=f"DeepGemmExperts GEMM warmup (MAX_M={MAX_M})"
159 - )
160 - num_tokens = MAX_M
161 - while num_tokens > 0:
162 - m_grouped_fp8_gemm_nt_contiguous(
163 - (a1q[:num_tokens], a1q_scales[:num_tokens]),
164 - (w, w_scale),
165 - out[:num_tokens],
166 - expert_ids[:num_tokens],
167 - )
168 - pbar.update(1)
169 - num_tokens = num_tokens - block_m
170 -
171 - _warmup(w1, w1_scale)
172 - _warmup(w2, w2_scale)
173 -
174 -
109 class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
110 def __init__(self, quant_config: FusedMoEQuantConfig):
111 super().__init__(quant_config)
@@ -215,11 +149,32 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
149 )
150 assert M_sum % block_m == 0
151
218 - workspace1 = (M_sum, N)
219 - workspace2 = (M_sum, max(N // 2, K))
152 + workspace1 = (M_sum, max(N // 2, K))
153 + workspace2 = (M_sum, max(N, K))
154 output = (M, K)
155 return (workspace1, workspace2, output)
156
157 + def _act_mul_quant(
158 + self, input: torch.Tensor, output: torch.Tensor, activation: str
159 + ) -> tuple[torch.Tensor, torch.Tensor]:
160 + if activation == "silu":
161 + return silu_mul_per_token_group_quant_fp8_colmajor(
162 + input=input, output=output
163 + )
164 + else:
165 + # This is a fallback path. If we find ourselves using any activation other
166 + # than silu, we should add that activation to
167 + # silu_mul_per_token_group_quant_fp8_colmajor kernel as it is much faster.
168 + M_sum, N = input.size()
169 + act_out = torch.empty(
170 + (M_sum, N // 2), dtype=input.dtype, device=input.device
171 + )
172 + self.activation(activation, act_out, input)
173 + assert self.block_shape is not None
174 + return per_token_group_quant_fp8(
175 + act_out, self.block_shape[1], column_major_scales=True, out_q=output
176 + )
177 +
178 def apply(
179 self,
180 output: torch.Tensor,
@@ -261,14 +216,9 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
216 expert_tokens_meta=expert_tokens_meta,
217 )
218
264 - a1q_perm = _resize_cache(workspace2.view(dtype=torch.float8_e4m3fn), (M_sum, K))
265 - mm1_out = _resize_cache(workspace13, (M_sum, N))
266 - act_out = _resize_cache(workspace2, (M_sum, N // 2))
267 - quant_out = _resize_cache(
268 - workspace13.view(dtype=torch.float8_e4m3fn), (M_sum, N // 2)
219 + a1q_perm = _resize_cache(
220 + workspace13.view(dtype=torch.float8_e4m3fn), (M_sum, K)
221 )
270 - mm2_out = _resize_cache(workspace2, (M_sum, K))
271 -
222 a1q, a1q_scale, expert_ids, inv_perm = deepgemm_moe_permute(
223 aq=a1q,
224 aq_scale=a1q_scale,
@@ -280,17 +230,19 @@ class DeepGemmExperts(mk.FusedMoEPermuteExpertsUnpermute):
230 )
231 assert a1q.size(0) == M_sum
232
233 + mm1_out = _resize_cache(workspace2, (M_sum, N))
234 m_grouped_fp8_gemm_nt_contiguous(
235 (a1q, a1q_scale), (w1, self.w1_scale), mm1_out, expert_ids
236 )
237
287 - self.activation(activation, act_out, mm1_out.view(-1, N))
288 -
289 - a2q_scale: torch.Tensor | None = None
290 - a2q, a2q_scale = per_token_group_quant_fp8(
291 - act_out, self.block_shape[1], column_major_scales=True, out_q=quant_out
238 + quant_out = _resize_cache(
239 + workspace13.view(dtype=torch.float8_e4m3fn), (M_sum, N // 2)
240 + )
241 + a2q, a2q_scale = self._act_mul_quant(
242 + input=mm1_out.view(-1, N), output=quant_out, activation=activation
243 )
244
245 + mm2_out = _resize_cache(workspace2, (M_sum, K))
246 m_grouped_fp8_gemm_nt_contiguous(
247 (a2q, a2q_scale), (w2, self.w2_scale), mm2_out, expert_ids
248 )
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+133
@@ -492,6 +492,139 @@ def _per_token_group_quant_fp8(
492 tl.store(y_s_ptr, y_s)
493
494
495 +@triton.jit
496 +def _silu_mul_per_token_group_quant_fp8_colmajor(
497 + y_ptr, # [M, N]
498 + y_q_ptr, # [M, N // 2]
499 + y_s_ptr, # [M, (N // 2) // GROUP_SIZE]
500 + M, # num tokens
501 + N, # intermediate size
502 + # Stride
503 + y_s_col_stride: tl.int64,
504 + # Information for float8
505 + eps,
506 + fp8_min,
507 + fp8_max,
508 + use_ue8m0: tl.constexpr,
509 + # Meta-parameters
510 + GROUP_SIZE: tl.constexpr,
511 + BLOCK_M: tl.constexpr,
512 + BLOCK_N: tl.constexpr,
513 +):
514 + # TODO(varun) : Add expert_ids so we may early-exit no-op thread blocks.
515 + """
516 + Each thread block (BLOCK_N) computes [BLOCK_M, GROUP_SIZE] act-mul outputs. Then
517 + the thread block quantizes the [BLOCK_M, GROUP_SIZE] block of values and fills
518 + the outputs tensors at the right positions.
519 + """
520 +
521 + pid_m = tl.program_id(0)
522 + pid_n = tl.program_id(1)
523 + N_2 = N // 2
524 +
525 + m_offset = pid_m * BLOCK_M
526 + n_offset = pid_n * BLOCK_N
527 + if m_offset >= M:
528 + return
529 +
530 + offs_n = tl.arange(0, BLOCK_N).to(tl.int64)
531 + offs_m = tl.arange(0, BLOCK_M).to(tl.int64)
532 +
533 + base_y_ptr = y_ptr + m_offset * N + n_offset
534 +
535 + act_in_ptrs = base_y_ptr + offs_m[:, None] * N + offs_n[None, :]
536 +
537 + act_in = tl.load(act_in_ptrs)
538 + mul_in = tl.load(act_in_ptrs + N_2)
539 +
540 + # silu & mul
541 + act_in = act_in.to(tl.float32)
542 + one_f32 = tl.cast(1, tl.float32)
543 + silu_out = (act_in / (one_f32 + tl.exp(-act_in))).to(y_ptr.dtype.element_ty)
544 + y = (silu_out * mul_in).to(tl.float32)
545 +
546 + # quant
547 + _absmax = tl.maximum(tl.max(tl.abs(y), axis=1), eps)
548 + scale_raw = _absmax / fp8_max
549 + y_s = tl.math.exp2(tl.ceil(tl.log2(scale_raw))) if use_ue8m0 else scale_raw
550 + y_s = tl.reshape(y_s, (BLOCK_M, 1))
551 + y_q = tl.clamp(y / y_s, fp8_min, fp8_max).to(y_q_ptr.dtype.element_ty)
552 +
553 + # store y_q
554 + base_y_q_ptr = y_q_ptr + m_offset * N_2 + n_offset
555 + y_q_ptrs = base_y_q_ptr + offs_m[:, None] * N_2 + offs_n[None, :]
556 + tl.store(y_q_ptrs, y_q)
557 +
558 + # store y_s
559 + group_id = n_offset // GROUP_SIZE
560 + base_y_s_ptr = y_s_ptr + group_id * y_s_col_stride + m_offset
561 + y_s_ptrs = base_y_s_ptr + offs_m
562 + y_s = tl.reshape(y_s, (BLOCK_M,))
563 + tl.store(y_s_ptrs, y_s)
564 +
565 +
566 +def silu_mul_per_token_group_quant_fp8_colmajor(
567 + input: torch.Tensor, # [M, N]
568 + output: torch.Tensor | None = None, # [M, N // 2]
569 + use_ue8m0: bool | None = None,
570 + eps: float = 1e-10,
571 +):
572 + """
573 + silu+mul + block-fp8 quant with group size 128.
574 + """
575 + GROUP_SIZE = 128
576 + assert input.ndim == 2
577 + if output is not None:
578 + assert output.ndim == 2
579 + assert input.size(0) % GROUP_SIZE == 0
580 + assert input.size(1) % (GROUP_SIZE * 2) == 0
581 +
582 + if use_ue8m0 is None:
583 + use_ue8m0 = is_deep_gemm_e8m0_used()
584 +
585 + M, N = input.size()
586 + N_2 = N // 2
587 +
588 + if output is None:
589 + output = torch.empty((M, N_2), dtype=torch.float8_e4m3fn, device=input.device)
590 +
591 + output_scales = torch.empty(
592 + ((N_2 // GROUP_SIZE), M), dtype=torch.float32, device=input.device
593 + ).transpose(0, 1)
594 +
595 + BLOCK_M = 8
596 + BLOCK_N = GROUP_SIZE
597 + assert M % BLOCK_M == 0
598 + assert N_2 % BLOCK_N == 0
599 +
600 + finfo = torch.finfo(torch.float8_e4m3fn)
601 + fp8_min = finfo.min
602 + fp8_max = finfo.max
603 +
604 + # Force even division so we can avoid edgecases within the kernel.
605 + assert M % BLOCK_M == 0
606 + assert N_2 % BLOCK_N == 0
607 + grid = (M // BLOCK_M, N_2 // BLOCK_N)
608 +
609 + _silu_mul_per_token_group_quant_fp8_colmajor[grid](
610 + input,
611 + output,
612 + output_scales,
613 + M,
614 + N,
615 + output_scales.stride(-1),
616 + eps,
617 + fp8_min,
618 + fp8_max,
619 + use_ue8m0,
620 + GROUP_SIZE,
621 + BLOCK_M,
622 + BLOCK_N,
623 + )
624 +
625 + return output, output_scales
626 +
627 +
628 @triton.jit
629 def _per_token_group_quant_fp8_colmajor(
630 # Pointers to inputs and output