@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 9239bf718e
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 15, 17:15 UTC

[Kernel] CUTLASS grouped gemm fp8 MoE kernel (#13972)

Signed-off-by: ElizaWszola <eliza@neuralmagic.com> Signed-off-by: ElizaWszola <ewszola@redhat.com> Co-authored-by: Lucas Wilkinson <wilkinson.lucas@gmail.com>

ElizaWszola committed Mar 27, 2025 at 01:54 UTC 9239bf718e5ebb5ab871ac8ed09fb80ed02fa82b
22 files changed +2317 -15
CMakeLists.txt
+27
@@ -461,6 +461,33 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
461 set(FP4_ARCHS)
462 endif()
463
464 + #
465 + # CUTLASS MoE kernels
466 +
467 + # The MoE kernel cutlass_moe_mm requires CUDA 12.3 or later (and only works
468 + # on Hopper). get_cutlass_moe_mm_data should only be compiled if it's possible
469 + # to compile MoE kernels that use its output.
470 + cuda_archs_loose_intersection(SCALED_MM_ARCHS "9.0a;" "${CUDA_ARCHS}")
471 + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND SCALED_MM_ARCHS)
472 + set(SRCS "csrc/quantization/cutlass_w8a8/moe/grouped_mm_c3x.cu"
473 + "csrc/quantization/cutlass_w8a8/moe/moe_data.cu")
474 + set_gencode_flags_for_srcs(
475 + SRCS "${SRCS}"
476 + CUDA_ARCHS "${SCALED_MM_ARCHS}")
477 + list(APPEND VLLM_EXT_SRC "${SRCS}")
478 + list(APPEND VLLM_GPU_FLAGS "-DENABLE_CUTLASS_MOE_SM90=1")
479 + message(STATUS "Building grouped_mm_c3x for archs: ${SCALED_MM_ARCHS}")
480 + else()
481 + if (NOT ${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND SCALED_MM_ARCHS)
482 + message(STATUS "Not building grouped_mm_c3x kernels as CUDA Compiler version is "
483 + "not >= 12.3, we recommend upgrading to CUDA 12.3 or later "
484 + "if you intend on running FP8 quantized MoE models on Hopper.")
485 + else()
486 + message(STATUS "Not building grouped_mm_c3x as no compatible archs found "
487 + "in CUDA target architectures")
488 + endif()
489 + endif()
490 +
491 #
492 # Machete kernels
493
benchmarks/kernels/benchmark_grouped_gemm_cutlass.py new
+340
@@ -0,0 +1,340 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +
3 +import torch
4 +import torch.utils.benchmark as benchmark
5 +from benchmark_shapes import WEIGHT_SHAPES_MOE
6 +
7 +from vllm import _custom_ops as ops
8 +from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
9 +from vllm.model_executor.layers.fused_moe.fused_moe import (cutlass_moe_fp8,
10 + fused_experts,
11 + fused_topk)
12 +from vllm.utils import FlexibleArgumentParser
13 +
14 +DEFAULT_MODELS = [
15 + "nm-testing/Mixtral-8x7B-Instruct-v0.1", "nm-testing/deepseekv2-lite",
16 + "ibm-granite/granite-3.0-1b-a400m", "ibm-granite/granite-3.0-3b-a800m"
17 +]
18 +DEFAULT_BATCH_SIZES = [1, 4, 8, 16, 32, 64, 128, 256, 512]
19 +DEFAULT_TP_SIZES = [1]
20 +
21 +PER_ACT_TOKEN_OPTS = [False]
22 +PER_OUT_CH_OPTS = [False]
23 +
24 +
25 +def to_fp8(tensor: torch.Tensor):
26 + finfo = torch.finfo(torch.float8_e4m3fn)
27 + return torch.round(tensor.clamp(
28 + min=finfo.min, max=finfo.max)).to(dtype=torch.float8_e4m3fn)
29 +
30 +
31 +def bench_run(results: list[benchmark.Measurement], model: str,
32 + num_experts: int, topk: int, per_act_token: bool,
33 + per_out_ch: bool, mkn: tuple[int, int, int]):
34 + label = "Quant Matmul"
35 +
36 + sub_label = (
37 + "{}, num_experts={}, topk={}, per_act_token={} per_out_ch={}, "
38 + "MKN=({})".format(model, num_experts, topk, per_act_token, per_out_ch,
39 + mkn))
40 +
41 + print(f"Testing: {sub_label}")
42 +
43 + (m, k, n) = mkn
44 +
45 + dtype = torch.half
46 +
47 + a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
48 + w1 = torch.randn((num_experts, 2 * n, k), device="cuda", dtype=dtype) / 10
49 + w2 = torch.randn((num_experts, k, n), device="cuda", dtype=dtype) / 10
50 +
51 + _, a_scale = ops.scaled_fp8_quant(a)
52 +
53 + w1_q = torch.empty((num_experts, 2 * n, k),
54 + device="cuda",
55 + dtype=torch.float8_e4m3fn)
56 + w2_q = torch.empty((num_experts, k, n),
57 + device="cuda",
58 + dtype=torch.float8_e4m3fn)
59 + w1_scale = torch.empty((num_experts, 1, 1),
60 + device="cuda",
61 + dtype=torch.float32)
62 + w2_scale = torch.empty((num_experts, 1, 1),
63 + device="cuda",
64 + dtype=torch.float32)
65 +
66 + ab_strides1 = torch.full((num_experts, ),
67 + k,
68 + device="cuda",
69 + dtype=torch.int64)
70 + c_strides1 = torch.full((num_experts, ),
71 + 2 * n,
72 + device="cuda",
73 + dtype=torch.int64)
74 + ab_strides2 = torch.full((num_experts, ),
75 + n,
76 + device="cuda",
77 + dtype=torch.int64)
78 + c_strides2 = torch.full((num_experts, ),
79 + k,
80 + device="cuda",
81 + dtype=torch.int64)
82 +
83 + for expert in range(num_experts):
84 + w1_q[expert], w1_scale[expert] = ops.scaled_fp8_quant(w1[expert])
85 + w2_q[expert], w2_scale[expert] = ops.scaled_fp8_quant(w2[expert])
86 + w1_q_notransp = w1_q.clone()
87 + w2_q_notransp = w2_q.clone()
88 + w1_q = w1_q.transpose(1, 2)
89 + w2_q = w2_q.transpose(1, 2)
90 +
91 + score = torch.randn((m, num_experts), device="cuda", dtype=dtype)
92 +
93 + topk_weights, topk_ids = fused_topk(a, score, topk, renormalize=False)
94 +
95 + def run_triton_moe(a: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor,
96 + topk_weights: torch.Tensor, topk_ids: torch.Tensor,
97 + w1_scale: torch.Tensor, w2_scale: torch.Tensor,
98 + a_scale: torch.Tensor, num_repeats: int):
99 + for _ in range(num_repeats):
100 + fused_experts(a,
101 + w1,
102 + w2,
103 + topk_weights,
104 + topk_ids,
105 + use_fp8_w8a8=True,
106 + w1_scale=w1_scale,
107 + w2_scale=w2_scale,
108 + a1_scale=a_scale)
109 +
110 + def run_cutlass_moe(a: torch.Tensor, a_scale: torch.Tensor,
111 + w1: torch.Tensor, w2: torch.Tensor,
112 + w1_scale: torch.Tensor, w2_scale: torch.Tensor,
113 + topk_weights: torch.Tensor, topk_ids: torch.Tensor,
114 + ab_strides1: torch.Tensor, c_strides1: torch.Tensor,
115 + ab_strides2: torch.Tensor, c_strides2: torch.Tensor,
116 + num_repeats: int):
117 + for _ in range(num_repeats):
118 + cutlass_moe_fp8(a,
119 + w1,
120 + w2,
121 + w1_scale,
122 + w2_scale,
123 + topk_weights,
124 + topk_ids,
125 + ab_strides1,
126 + c_strides1,
127 + ab_strides2,
128 + c_strides2,
129 + a1_scale=a_scale)
130 +
131 + def run_cutlass_from_graph(
132 + a: torch.Tensor, a_scale: torch.Tensor, w1_q: torch.Tensor,
133 + w2_q: torch.Tensor, w1_scale: torch.Tensor, w2_scale: torch.Tensor,
134 + topk_weights: torch.Tensor, topk_ids: torch.Tensor,
135 + ab_strides1: torch.Tensor, c_strides1: torch.Tensor,
136 + ab_strides2: torch.Tensor, c_strides2: torch.Tensor):
137 + with set_current_vllm_config(
138 + VllmConfig(parallel_config=ParallelConfig(
139 + pipeline_parallel_size=1))):
140 + return cutlass_moe_fp8(a,
141 + w1_q,
142 + w2_q,
143 + w1_scale,
144 + w2_scale,
145 + topk_weights,
146 + topk_ids,
147 + ab_strides1,
148 + c_strides1,
149 + ab_strides2,
150 + c_strides2,
151 + a1_scale=a_scale)
152 +
153 + def run_triton_from_graph(a: torch.Tensor, w1: torch.Tensor,
154 + w2: torch.Tensor, topk_weights: torch.Tensor,
155 + topk_ids: torch.Tensor, w1_scale: torch.Tensor,
156 + w2_scale: torch.Tensor, a_scale: torch.Tensor):
157 + with set_current_vllm_config(
158 + VllmConfig(parallel_config=ParallelConfig(
159 + pipeline_parallel_size=1))):
160 + return fused_experts(a,
161 + w1,
162 + w2,
163 + topk_weights,
164 + topk_ids,
165 + use_fp8_w8a8=True,
166 + w1_scale=w1_scale,
167 + w2_scale=w2_scale,
168 + a1_scale=a_scale)
169 +
170 + def replay_graph(graph, num_repeats):
171 + for _ in range(num_repeats):
172 + graph.replay()
173 + torch.cuda.synchronize()
174 +
175 + cutlass_stream = torch.cuda.Stream()
176 + cutlass_graph = torch.cuda.CUDAGraph()
177 + with torch.cuda.graph(cutlass_graph, stream=cutlass_stream):
178 + run_cutlass_from_graph(a, a_scale, w1_q, w2_q, w1_scale, w2_scale,
179 + topk_weights, topk_ids, ab_strides1, c_strides1,
180 + ab_strides2, c_strides2)
181 + torch.cuda.synchronize()
182 +
183 + triton_stream = torch.cuda.Stream()
184 + triton_graph = torch.cuda.CUDAGraph()
185 + with torch.cuda.graph(triton_graph, stream=triton_stream):
186 + run_triton_from_graph(a, w1_q_notransp, w2_q_notransp, topk_weights,
187 + topk_ids, w1_scale, w2_scale, a_scale)
188 + torch.cuda.synchronize()
189 +
190 + min_run_time = 5
191 + num_warmup = 5
192 + num_runs = 25
193 +
194 + globals = {
195 + # Baseline params
196 + "w1": w1,
197 + "w2": w2,
198 + "score": score,
199 + "topk": topk,
200 + "w1_q_notransp": w1_q_notransp,
201 + "w2_q_notransp": w2_q_notransp,
202 + # Cutlass params
203 + "a_scale": a_scale,
204 + "w1_q": w1_q,
205 + "w2_q": w2_q,
206 + "w1_scale": w1_scale,
207 + "w2_scale": w2_scale,
208 + "ab_strides1": ab_strides1,
209 + "c_strides1": c_strides1,
210 + "ab_strides2": ab_strides2,
211 + "c_strides2": c_strides2,
212 + # cuda graph params
213 + "cutlass_graph": cutlass_graph,
214 + "triton_graph": triton_graph,
215 + # Gen params
216 + "a": a,
217 + "topk_weights": topk_weights,
218 + "topk_ids": topk_ids,
219 + "num_runs": num_runs,
220 + # Kernels
221 + "run_triton_moe": run_triton_moe,
222 + "run_cutlass_moe": run_cutlass_moe,
223 + "replay_graph": replay_graph,
224 + }
225 +
226 + # Warmup
227 + run_triton_moe(a, w1_q_notransp, w2_q_notransp, topk_weights, topk_ids,
228 + w1_scale, w2_scale, a_scale, num_warmup)
229 +
230 + results.append(
231 + benchmark.Timer(
232 + stmt=
233 + "run_triton_moe(a, w1_q_notransp, w2_q_notransp, topk_weights, topk_ids, w1_scale, w2_scale, a_scale, num_runs)", # noqa: E501
234 + globals=globals,
235 + label=label,
236 + sub_label=sub_label,
237 + description="triton_moe",
238 + ).blocked_autorange(min_run_time=min_run_time))
239 +
240 + # Warmup
241 + replay_graph(triton_graph, num_warmup)
242 +
243 + results.append(
244 + benchmark.Timer(
245 + stmt="replay_graph(triton_graph, num_runs)",
246 + globals=globals,
247 + label=label,
248 + sub_label=sub_label,
249 + description="triton_moe_cuda_graphs",
250 + ).blocked_autorange(min_run_time=min_run_time))
251 +
252 + # Warmup
253 + run_cutlass_moe(a, a_scale, w1_q, w2_q, w1_scale, w2_scale, topk_weights,
254 + topk_ids, ab_strides1, c_strides1, ab_strides2, c_strides2,
255 + num_warmup)
256 +
257 + results.append(
258 + benchmark.Timer(
259 + stmt=
260 + "run_cutlass_moe(a, a_scale, w1_q, w2_q, w1_scale, w2_scale, topk_weights, topk_ids, ab_strides1, c_strides1, ab_strides2, c_strides2, num_runs)", # noqa: E501
261 + globals=globals,
262 + label=label,
263 + sub_label=sub_label,
264 + description="grouped_gemm_moe",
265 + ).blocked_autorange(min_run_time=min_run_time))
266 +
267 + # Warmup
268 + replay_graph(cutlass_graph, num_warmup)
269 +
270 + results.append(
271 + benchmark.Timer(
272 + stmt="replay_graph(cutlass_graph, num_runs)",
273 + globals=globals,
274 + label=label,
275 + sub_label=sub_label,
276 + description="grouped_gemm_moe_cuda_graphs",
277 + ).blocked_autorange(min_run_time=min_run_time))
278 +
279 +
280 +def main(args):
281 + print("Benchmarking models:")
282 + for i, model in enumerate(args.models):
283 + print(f"[{i}] {model}")
284 +
285 + results: list[benchmark.Measurement] = []
286 +
287 + for model in args.models:
288 + for tp in args.tp_sizes:
289 + for layer in WEIGHT_SHAPES_MOE[model]:
290 + num_experts = layer[0]
291 + topk = layer[1]
292 + size_k = layer[2]
293 + size_n = layer[3] // tp
294 +
295 + if len(args.limit_k) > 0 and size_k not in args.limit_k:
296 + continue
297 +
298 + if len(args.limit_n) > 0 and size_n not in args.limit_n:
299 + continue
300 +
301 + for per_act_token in PER_ACT_TOKEN_OPTS:
302 + for per_out_ch in PER_OUT_CH_OPTS:
303 + for size_m in DEFAULT_BATCH_SIZES:
304 + mkn = (size_m, size_k, size_n)
305 + bench_run(results, model, num_experts, topk,
306 + per_act_token, per_out_ch, mkn)
307 +
308 + compare = benchmark.Compare(results)
309 + compare.print()
310 +
311 +
312 +if __name__ == "__main__":
313 + parser = FlexibleArgumentParser(
314 + description="Benchmark Marlin across specified models/shapes/batches")
315 + parser.add_argument(
316 + "--models",
317 + nargs="+",
318 + type=str,
319 + default=DEFAULT_MODELS,
320 + choices=WEIGHT_SHAPES_MOE.keys(),
321 + )
322 + parser.add_argument("--tp-sizes",
323 + nargs="+",
324 + type=int,
325 + default=DEFAULT_TP_SIZES)
326 + parser.add_argument("--batch-sizes",
327 + nargs="+",
328 + type=int,
329 + default=DEFAULT_BATCH_SIZES)
330 + parser.add_argument("--limit-k", nargs="+", type=int, default=[])
331 + parser.add_argument("--limit-n", nargs="+", type=int, default=[])
332 + parser.add_argument("--limit-num-groups", nargs="+", type=int, default=[])
333 + parser.add_argument("--limit-per-act-token",
334 + nargs="+",
335 + type=int,
336 + default=[])
337 + parser.add_argument("--limit-per-out-ch", nargs="+", type=int, default=[])
338 +
339 + args = parser.parse_args()
340 + main(args)
benchmarks/kernels/benchmark_shapes.py
+16
@@ -75,3 +75,19 @@ WEIGHT_SHAPES = {
75 [7168, 8192],
76 ],
77 }
78 +
79 +WEIGHT_SHAPES_MOE = {
80 + "nm-testing/Mixtral-8x7B-Instruct-v0.1": [
81 + [8, 2, 4096, 28672],
82 + [8, 2, 14336, 4096],
83 + ],
84 + "nm-testing/deepseekv2-lite": [
85 + [64, 6, 2048, 1408],
86 + ],
87 + "ibm-granite/granite-3.0-1b-a400m": [
88 + [32, 8, 1024, 1024],
89 + ],
90 + "ibm-granite/granite-3.0-3b-a800m": [
91 + [40, 8, 1024, 1536],
92 + ],
93 +}
csrc/cutlass_extensions/common.hpp
+11 -1
@@ -48,4 +48,14 @@ struct enable_sm90_or_later : Kernel {
48 Kernel::operator()(std::forward<Args>(args)...);
49 #endif
50 }
51 -};
\ No newline at end of file
51 +};
52 +
53 +template <typename Kernel>
54 +struct enable_sm90_only : Kernel {
55 + template <typename... Args>
56 + CUTLASS_DEVICE void operator()(Args&&... args) {
57 +#if defined __CUDA_ARCH__ && __CUDA_ARCH__ == 900
58 + Kernel::operator()(std::forward<Args>(args)...);
59 +#endif
60 + }
61 +};
csrc/cutlass_extensions/epilogue/broadcast_load_epilogue_array_c3x.hpp new
+457
@@ -0,0 +1,457 @@
1 +/***************************************************************************************************
2 + * Copyright (c) 2023 - 2024 NVIDIA CORPORATION & AFFILIATES. All rights
3 + *reserved. SPDX-License-Identifier: BSD-3-Clause
4 + *
5 + * Redistribution and use in source and binary forms, with or without
6 + * modification, are permitted provided that the following conditions are met:
7 + *
8 + * 1. Redistributions of source code must retain the above copyright notice,
9 + *this list of conditions and the following disclaimer.
10 + *
11 + * 2. Redistributions in binary form must reproduce the above copyright notice,
12 + * this list of conditions and the following disclaimer in the documentation
13 + * and/or other materials provided with the distribution.
14 + *
15 + * 3. Neither the name of the copyright holder nor the names of its
16 + * contributors may be used to endorse or promote products derived from
17 + * this software without specific prior written permission.
18 + *
19 + * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
20 + * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
21 + * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
22 + *ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE
23 + *LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
24 + *CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
25 + *SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
26 + *INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
27 + *CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
28 + *ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
29 + *POSSIBILITY OF SUCH DAMAGE.
30 + *
31 + **************************************************************************************************/
32 +
33 +//
34 +// This file is a modified excerpt of
35 +// include/cutlass/epilogue/fusion/sm90_visitor_load_tma_warpspecialized.hpp
36 +// from https://github.com/NVIDIA/cutlass v3.5.0
37 +// It has been modified to support either row/column or scalar broadcasting
38 +// where the tensor being loaded from is always passed in via a device pointer.
39 +// This lets one compiled kernel handle all cases of per-tensor or
40 +// per-channel/per-token quantization.
41 +//
42 +// This interface also allows the scales to be passed in as tensors that
43 +// consistently reside on the device, which avoids an issue with a previous
44 +// implementation where scalars needed to be on the CPU since they
45 +// were passed in via float values. This created a potential performance hazard
46 +// if scales were initially on the device, and caused torch.compile graphs
47 +// breaks when moving scales to the CPU.
48 +//
49 +#pragma once
50 +
51 +// Turn off clang-format for the entire file to keep it close to upstream
52 +// clang-format off
53 +
54 +#include "cutlass/cutlass.h"
55 +#include "cutlass/arch/barrier.h"
56 +
57 +#include "cute/tensor.hpp"
58 +#include "cutlass/epilogue/fusion/sm90_visitor_tma_warpspecialized.hpp"
59 +
60 +namespace cutlass::epilogue::fusion {
61 +
62 +using namespace cute;
63 +using namespace detail;
64 +
65 +// Row vector broadcast
66 +template<
67 + int Stages,
68 + class CtaTileShapeMNK,
69 + class Element,
70 + class StrideMNL = Stride<_0,_1,_0>,
71 + int Alignment = 128 / sizeof_bits_v<Element>
72 +>
73 +struct Sm90RowOrScalarBroadcastArray {
74 + static_assert(Stages == 0, "Row broadcast doesn't support smem usage");
75 + static_assert(is_static_v<decltype(take<0,2>(StrideMNL{}))>); // batch stride can be dynamic or static
76 + static_assert(take<0,2>(StrideMNL{}) == Stride<_0,_1>{});
77 +
78 + struct SharedStorage {
79 + array_aligned<Element, size<1>(CtaTileShapeMNK{})> smem;
80 + };
81 +
82 + // This struct has been modified to have a bool indicating that ptr_row is a
83 + // scalar that must be broadcast, instead of containing a scalar that is
84 + // valid if ptr_row is null.
85 + struct Arguments {
86 + const Element* const* ptr_row_array = nullptr;
87 + bool row_broadcast = true;
88 + StrideMNL dRow = {};
89 + };
90 +
91 + using Params = Arguments;
92 +
93 + template <class ProblemShape>
94 + static constexpr Params
95 + to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
96 + return args;
97 + }
98 +
99 + template <class ProblemShape>
100 + static bool
101 + can_implement(ProblemShape const& problem_shape, Arguments const& args) {
102 + return true;
103 + }
104 +
105 + template <class ProblemShape>
106 + static size_t
107 + get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) {
108 + return 0;
109 + }
110 +
111 + template <class ProblemShape>
112 + static cutlass::Status
113 + initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream,
114 + CudaHostAdapter* cuda_adapter = nullptr) {
115 + return cutlass::Status::kSuccess;
116 + }
117 +
118 + CUTLASS_HOST_DEVICE
119 + Sm90RowOrScalarBroadcastArray() { }
120 +
121 + CUTLASS_HOST_DEVICE
122 + Sm90RowOrScalarBroadcastArray(Params const& params, SharedStorage const& shared_storage)
123 + : params(params)
124 + , smem(const_cast<Element*>(shared_storage.smem.data())) { }
125 +
126 + Params params;
127 + Element *smem = nullptr;
128 +
129 + CUTLASS_DEVICE bool
130 + is_producer_load_needed() const {
131 + return false;
132 + }
133 +
134 + CUTLASS_DEVICE bool
135 + is_C_load_needed() const {
136 + return false;
137 + }
138 +
139 + CUTLASS_DEVICE bool
140 + is_zero() const {
141 + return (!params.row_broadcast && *(params.ptr_row_array[group]) == Element(0));
142 + }
143 +
144 + template <class... Args>
145 + CUTLASS_DEVICE auto
146 + get_producer_load_callbacks(ProducerLoadArgs<Args...> const& args) {
147 + return EmptyProducerLoadCallbacks{};
148 + }
149 +
150 + template <class GS_GTensor, class GS_STensor, class GS_CTensor, class Tiled_G2S, class SR_STensor, class SR_RTensor, class CTensor, class ThrResidue, class ThrNum>
151 + struct ConsumerStoreCallbacks : EmptyConsumerStoreCallbacks {
152 + CUTLASS_DEVICE
153 + ConsumerStoreCallbacks(
154 + GS_GTensor tGS_gRow_, GS_STensor tGS_sRow_,
155 + GS_CTensor tGS_cRow_, Tiled_G2S tiled_g2s_,
156 + SR_STensor tSR_sRow_, SR_RTensor tSR_rRow_,
157 + CTensor tCcRow_, ThrResidue residue_tCcRow_, ThrNum thr_num_,
158 + int group, Params const& params_)
159 + : tGS_gRow(tGS_gRow_)
160 + , tGS_sRow(tGS_sRow_)
161 + , tGS_cRow(tGS_cRow_)
162 + , tiled_G2S(tiled_g2s_)
163 + , tSR_sRow(tSR_sRow_)
164 + , tSR_rRow(tSR_rRow_)
165 + , tCcRow(tCcRow_)
166 + , residue_tCcRow(residue_tCcRow_)
167 + , group(group)
168 + , params(params_) {}
169 +
170 + GS_GTensor tGS_gRow; // (CPY,CPY_M,CPY_N)
171 + GS_STensor tGS_sRow; // (CPY,CPY_M,CPY_N)
172 + GS_CTensor tGS_cRow; // (CPY,CPY_M,CPY_N)
173 + Tiled_G2S tiled_G2S;
174 +
175 + SR_STensor tSR_sRow; // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
176 + SR_RTensor tSR_rRow; // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
177 +
178 + CTensor tCcRow; // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
179 + ThrResidue residue_tCcRow; // (m, n)
180 + ThrNum thr_num;
181 + int group;
182 + Params const& params;
183 +
184 + CUTLASS_DEVICE void
185 + begin() {
186 + if (!params.row_broadcast) {
187 + fill(tSR_rRow, *(params.ptr_row_array[group]));
188 + return;
189 + }
190 +
191 + auto synchronize = [&] () { cutlass::arch::NamedBarrier::sync(thr_num, cutlass::arch::ReservedNamedBarriers::EpilogueBarrier); };
192 + Tensor tGS_gRow_flt = filter_zeros(tGS_gRow);
193 + Tensor tGS_sRow_flt = filter_zeros(tGS_sRow);
194 + Tensor tGS_cRow_flt = make_tensor(tGS_cRow.data(), make_layout(tGS_gRow_flt.shape(), tGS_cRow.stride()));
195 +
196 + for (int i = 0; i < size(tGS_gRow_flt); ++i) {
197 + if (get<1>(tGS_cRow_flt(i)) >= size<1>(CtaTileShapeMNK{})) {
198 + continue; // OOB of SMEM,
199 + }
200 + if (elem_less(tGS_cRow_flt(i), make_coord(get<0>(residue_tCcRow), get<1>(residue_tCcRow)))) {
201 + tGS_sRow_flt(i) = tGS_gRow_flt(i);
202 + }
203 + else {
204 + tGS_sRow_flt(i) = Element(0); // Set to Zero when OOB so LDS could be issue without any preds.
205 + }
206 + }
207 + synchronize();
208 + }
209 +
210 + CUTLASS_DEVICE void
211 + begin_loop(int epi_m, int epi_n) {
212 + if (epi_m == 0) { // Assumes M-major subtile loop
213 + if (!params.row_broadcast) return; // Do not issue LDS when row is scalar
214 + Tensor tSR_sRow_flt = filter_zeros(tSR_sRow(_,_,_,epi_m,epi_n));
215 + Tensor tSR_rRow_flt = filter_zeros(tSR_rRow);
216 + copy(tSR_sRow_flt, tSR_rRow_flt);
217 + }
218 + }
219 +
220 + template <typename ElementAccumulator, int FragmentSize>
221 + CUTLASS_DEVICE Array<Element, FragmentSize>
222 + visit(Array<ElementAccumulator, FragmentSize> const& frg_acc, int epi_v, int epi_m, int epi_n) {
223 + Array<Element, FragmentSize> frg_row;
224 +
225 + CUTLASS_PRAGMA_UNROLL
226 + for (int i = 0; i < FragmentSize; ++i) {
227 + frg_row[i] = tSR_rRow(epi_v * FragmentSize + i);
228 + }
229 +
230 + return frg_row;
231 + }
232 + };
233 +
234 + template <
235 + bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
236 + class... Args
237 + >
238 + CUTLASS_DEVICE auto
239 + get_consumer_store_callbacks(ConsumerStoreArgs<Args...> const& args) {
240 + auto [M, N, K, L] = args.problem_shape_mnkl;
241 + auto [m, n, k, l] = args.tile_coord_mnkl;
242 + using ThreadCount = decltype(size(args.tiled_copy));
243 +
244 + Tensor mRow = make_tensor(make_gmem_ptr(params.ptr_row_array[l]), make_shape(M,N,1), params.dRow);
245 + Tensor gRow = local_tile(mRow(_,_,l), take<0,2>(args.tile_shape_mnk), make_coord(m, n)); // (CTA_M, CTA_N)
246 + Tensor sRow = make_tensor(make_smem_ptr(smem),
247 + make_shape(size<0>(CtaTileShapeMNK{}), size<1>(CtaTileShapeMNK{})), make_shape(_0{}, _1{})); // (CTA_M, CTA_N)
248 + //// G2S: Gmem to Smem
249 + auto tiled_g2s = make_tiled_copy(Copy_Atom<DefaultCopy, Element>{},
250 + Layout< Shape<_1, ThreadCount>,
251 + Stride<_0, _1>>{},
252 + Layout<_1>{});
253 + auto thr_g2s = tiled_g2s.get_slice(args.thread_idx);
254 + Tensor tGS_gRow = thr_g2s.partition_S(gRow);
255 + Tensor tGS_sRow = thr_g2s.partition_D(sRow);
256 +
257 + //// G2S: Coord
258 + auto cRow = make_identity_tensor(make_shape(size<0>(CtaTileShapeMNK{}), size<1>(CtaTileShapeMNK{})));
259 + Tensor tGS_cRow = thr_g2s.partition_S(cRow);
260 +
261 + //// S2R: Smem to Reg
262 + Tensor tSR_sRow = sm90_partition_for_epilogue<ReferenceSrc>(sRow, args.epi_tile, args.tiled_copy, args.thread_idx);
263 + Tensor tSR_rRow = make_tensor_like(take<0,3>(tSR_sRow)); // (CPY,CPY_M,CPY_N)
264 +
265 + return ConsumerStoreCallbacks<decltype(tGS_gRow), decltype(tGS_sRow), decltype(tGS_cRow), decltype(tiled_g2s), decltype(tSR_sRow), decltype(tSR_rRow), decltype(args.tCcD), decltype(args.residue_cD), ThreadCount>(
266 + tGS_gRow,
267 + tGS_sRow,
268 + tGS_cRow, tiled_g2s,
269 + tSR_sRow,
270 + tSR_rRow,
271 + args.tCcD,
272 + args.residue_cD,
273 + ThreadCount{},
274 + l,
275 + params);
276 + }
277 +};
278 +
279 +/////////////////////////////////////////////////////////////////////////////////////////////////
280 +
281 +// Column vector broadcast
282 +template<
283 + int Stages,
284 + class CtaTileShapeMNK,
285 + class Element,
286 + class StrideMNL = Stride<_1,_0,_0>,
287 + int Alignment = 128 / sizeof_bits_v<Element>
288 +>
289 +struct Sm90ColOrScalarBroadcastArray {
290 + static_assert(Stages == 0, "Column broadcast doesn't support smem usage yet");
291 + static_assert(Alignment * sizeof_bits_v<Element> % 128 == 0, "sub-16B alignment not supported yet");
292 + static_assert(
293 + (cute::is_same_v<StrideMNL, Stride<_1,_0, _0>>) || // col vector broadcast, e.g. per-row alpha/bias
294 + (cute::is_same_v<StrideMNL, Stride<_1,_0,int>>)); // batched col vector broadcast, e.g. batched per-row bias
295 +
296 + // Accumulator distributes col elements evenly amongst threads so we can just directly load from gmem
297 + struct SharedStorage { };
298 +
299 + // This struct has been modified to have a bool indicating that ptr_col is a
300 + // scalar that must be broadcast, instead of containing a scalar that is
301 + // valid if ptr_col is null.
302 + struct Arguments {
303 + const Element* const* ptr_col_array = nullptr;
304 + bool col_broadcast = true;
305 + StrideMNL dCol = {};
306 + };
307 +
308 + using Params = Arguments;
309 +
310 + template <class ProblemShape>
311 + static constexpr Params
312 + to_underlying_arguments(ProblemShape const& problem_shape, Arguments const& args, void* workspace) {
313 + return args;
314 + }
315 +
316 + template <class ProblemShape>
317 + static bool
318 + can_implement(ProblemShape const& problem_shape, Arguments const& args) {
319 + return true;
320 + }
321 +
322 + template <class ProblemShape>
323 + static size_t
324 + get_workspace_size(ProblemShape const& problem_shape, Arguments const& args) {
325 + return 0;
326 + }
327 +
328 + template <class ProblemShape>
329 + static cutlass::Status
330 + initialize_workspace(ProblemShape const& problem_shape, Arguments const& args, void* workspace, cudaStream_t stream,
331 + CudaHostAdapter* cuda_adapter = nullptr) {
332 + return cutlass::Status::kSuccess;
333 + }
334 +
335 + CUTLASS_DEVICE bool
336 + is_producer_load_needed() const {
337 + return false;
338 + }
339 +
340 + CUTLASS_DEVICE bool
341 + is_C_load_needed() const {
342 + return false;
343 + }
344 +
345 + CUTLASS_DEVICE bool
346 + is_zero() const {
347 + return (!params.col_broadcast && *(params.ptr_col_array[group]) == Element(0));
348 + }
349 +
350 + CUTLASS_HOST_DEVICE
351 + Sm90ColOrScalarBroadcastArray() { }
352 +
353 + CUTLASS_HOST_DEVICE
354 + Sm90ColOrScalarBroadcastArray(Params const& params, SharedStorage const& shared_storage)
355 + : params(params) { }
356 +
357 + Params params;
358 +
359 + template <class... Args>
360 + CUTLASS_DEVICE auto
361 + get_producer_load_callbacks(ProducerLoadArgs<Args...> const& args) {
362 + return EmptyProducerLoadCallbacks{};
363 + }
364 +
365 + template<class GTensor, class RTensor, class CTensor, class ProblemShape>
366 + struct ConsumerStoreCallbacks : EmptyConsumerStoreCallbacks {
367 + CUTLASS_DEVICE
368 + ConsumerStoreCallbacks(
369 + GTensor&& tCgCol,
370 + RTensor&& tCrCol,
371 + CTensor&& tCcCol,
372 + ProblemShape problem_shape,
373 + int group,
374 + Params const& params
375 + ):
376 + tCgCol(cute::forward<GTensor>(tCgCol)),
377 + tCrCol(cute::forward<RTensor>(tCrCol)),
378 + tCcCol(cute::forward<CTensor>(tCcCol)),
379 + m(get<0>(problem_shape)),
380 + group(group),
381 + params(params) {}
382 +
383 + GTensor tCgCol; // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
384 + RTensor tCrCol;
385 + CTensor tCcCol; // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
386 + Params const& params;
387 + int m;
388 + int group;
389 +
390 + CUTLASS_DEVICE void
391 + begin() {
392 + Tensor pred = make_tensor<bool>(shape(tCgCol));
393 + CUTLASS_PRAGMA_UNROLL
394 + for (int i = 0; i < size(pred); ++i) {
395 + pred(i) = get<0>(tCcCol(i)) < m;
396 + }
397 +
398 + if (!params.col_broadcast) {
399 + fill(tCrCol, *(params.ptr_col_array[group]));
400 + return;
401 + }
402 +
403 + // Filter so we don't issue redundant copies over stride-0 modes
404 + // (only works if 0-strides are in same location, which is by construction)
405 + copy_if(pred, filter(tCgCol), filter(tCrCol));
406 + }
407 +
408 + template <typename ElementAccumulator, int FragmentSize>
409 + CUTLASS_DEVICE Array<Element, FragmentSize>
410 + visit(Array<ElementAccumulator, FragmentSize> const& frg_acc, int epi_v, int epi_m, int epi_n) {
411 + Array<Element, FragmentSize> frg_col;
412 + Tensor tCrCol_mn = tCrCol(_,_,_,epi_m,epi_n);
413 +
414 + CUTLASS_PRAGMA_UNROLL
415 + for (int i = 0; i < FragmentSize; ++i) {
416 + frg_col[i] = tCrCol_mn(epi_v * FragmentSize + i);
417 + }
418 +
419 + return frg_col;
420 + }
421 +
422 + };
423 +
424 + template <
425 + bool ReferenceSrc, // do register tensors reference the src or dst layout of the tiled copy
426 + class... Args
427 + >
428 + CUTLASS_DEVICE auto
429 + get_consumer_store_callbacks(ConsumerStoreArgs<Args...> const& args) {
430 +
431 + auto [M, N, K, L] = args.problem_shape_mnkl;
432 + auto [m, n, k, l] = args.tile_coord_mnkl;
433 +
434 + Tensor mCol = make_tensor(make_gmem_ptr(params.ptr_col_array[l]), make_shape(M,N,1), params.dCol);
435 + Tensor tCgCol = sm90_partition_for_epilogue<ReferenceSrc>( // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
436 + mCol, args.tile_shape_mnk, args.tile_coord_mnkl, args.epi_tile, args.tiled_copy, args.thread_idx);
437 + Tensor tCrCol = make_tensor_like(tCgCol); // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
438 +
439 + // Generate an identity tensor matching the shape of the global tensor and
440 + // partition the same way, this will be used to generate the predicate
441 + // tensor for loading
442 + Tensor cCol = make_identity_tensor(mCol.shape());
443 + Tensor tCcCol = sm90_partition_for_epilogue<ReferenceSrc>( // (CPY,CPY_M,CPY_N,EPI_M,EPI_N)
444 + cCol, args.tile_shape_mnk, args.tile_coord_mnkl, args.epi_tile, args.tiled_copy, args.thread_idx);
445 +
446 + return ConsumerStoreCallbacks(
447 + cute::move(tCgCol),
448 + cute::move(tCrCol),
449 + cute::move(tCcCol),
450 + args.problem_shape_mnkl,
451 + l,
452 + params
453 + );
454 + }
455 +};
456 +
457 +}
csrc/cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp
+66
@@ -1,6 +1,7 @@
1 #pragma once
2
3 #include "cutlass_extensions/epilogue/broadcast_load_epilogue_c3x.hpp"
4 +#include "cutlass_extensions/epilogue/broadcast_load_epilogue_array_c3x.hpp"
5
6 /*
7 This file defines custom epilogues for fusing channel scales, token scales,
@@ -69,6 +70,16 @@ struct ScaledEpilogueBase {
70 0 /*Stages*/, TileShape, T, T, Stride<Int<0>, Int<1>, Int<0>>,
71 128 / sizeof_bits_v<T>, EnableNullPtr>;
72
73 + template <typename T>
74 + using ColOrScalarLoadArray =
75 + cutlass::epilogue::fusion::Sm90ColOrScalarBroadcastArray<
76 + 0 /*Stages*/, TileShape, T, Stride<Int<1>, Int<0>, Int<0>>>;
77 +
78 + template <typename T>
79 + using RowOrScalarLoadArray =
80 + cutlass::epilogue::fusion::Sm90RowOrScalarBroadcastArray<
81 + 0 /*Stages*/, TileShape, T, Stride<Int<0>, Int<1>, Int<0>>>;
82 +
83 // This utility function constructs the arguments for the load descriptors
84 // from a tensor. It can handle both row and column, as well as row/column or
85 // scalar cases.
@@ -96,6 +107,14 @@ struct ScaledEpilogueBase {
107 std::is_same_v<Descriptor, RowLoad<T, true>>);
108 return Arguments{data_ptr};
109 }
110 +
111 + template <typename Descriptor, typename T>
112 + static auto args_from_tensor(const T* const* data_ptr, bool do_broadcast) {
113 + using Arguments = typename Descriptor::Arguments;
114 + static_assert(std::is_same_v<Descriptor, ColOrScalarLoadArray<T>> ||
115 + std::is_same_v<Descriptor, RowOrScalarLoadArray<T>>);
116 + return Arguments{data_ptr, do_broadcast};
117 + }
118 };
119
120 /*
@@ -381,4 +400,51 @@ struct ScaledEpilogueBiasAzpToken
400 }
401 };
402
403 +/*
404 + This epilogue works like ScaledEpilogue, but ScaleA and ScaleB are pointers
405 + to arrays containing different scales used in group gemm. The number of
406 + pointers in ScaleA and the number of pointers in ScaleB are equal to the
407 + group size.
408 +*/
409 +template <typename ElementAcc, typename ElementD, typename EpilogueDescriptor>
410 +struct ScaledEpilogueArray
411 + : private ScaledEpilogueBase<ElementAcc, ElementD, EpilogueDescriptor> {
412 + private:
413 + using SUPER = ScaledEpilogueBase<ElementAcc, ElementD, EpilogueDescriptor>;
414 + using Accum = typename SUPER::Accum;
415 + using ScaleA = typename SUPER::template ColOrScalarLoadArray<float>;
416 + using ScaleB = typename SUPER::template RowOrScalarLoadArray<float>;
417 +
418 + using Compute0 = cutlass::epilogue::fusion::Sm90Compute<
419 + cutlass::multiplies, float, float,
420 + cutlass::FloatRoundStyle::round_to_nearest>;
421 +
422 + using EVTCompute0 =
423 + cutlass::epilogue::fusion::Sm90EVT<Compute0, ScaleB, Accum>;
424 +
425 + using Compute1 = cutlass::epilogue::fusion::Sm90Compute<
426 + cutlass::multiplies, ElementD, float,
427 + cutlass::FloatRoundStyle::round_to_nearest>;
428 +
429 + public:
430 + using EVTCompute =
431 + cutlass::epilogue::fusion::Sm90EVT<Compute1, ScaleA, EVTCompute0>;
432 + using ArgumentType = typename EVTCompute::Arguments;
433 +
434 + using ScaleAArray = typename SUPER::template ColOrScalarLoadArray<float>;
435 + using ScaleBArray = typename SUPER::template RowOrScalarLoadArray<float>;
436 +
437 + static ArgumentType prepare_args(float const* const* a_scales_ptr,
438 + float const* const* b_scales_ptr,
439 + bool a_col_broadcast, bool b_row_broadcast) {
440 + auto a_args = SUPER::template args_from_tensor<ScaleAArray, float>(
441 + a_scales_ptr, a_col_broadcast);
442 + auto b_args = SUPER::template args_from_tensor<ScaleBArray, float>(
443 + b_scales_ptr, b_row_broadcast);
444 +
445 + typename EVTCompute0::Arguments evt0_args{b_args, {}, {}};
446 + return ArgumentType{a_args, evt0_args, {}};
447 + }
448 +};
449 +
450 }; // namespace vllm::c3x
csrc/ops.h
+14
@@ -164,6 +164,7 @@ int64_t ggml_moe_get_block_size(int64_t type);
164 bool cutlass_scaled_mm_supports_fp4(int64_t cuda_device_capability);
165 bool cutlass_scaled_mm_supports_fp8(int64_t cuda_device_capability);
166 bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability);
167 +bool cutlass_group_gemm_supported(int64_t cuda_device_capability);
168
169 void cutlass_scaled_fp4_mm(torch::Tensor& D, torch::Tensor const& A,
170 torch::Tensor const& B, torch::Tensor const& A_sf,
@@ -175,6 +176,19 @@ void cutlass_scaled_mm(torch::Tensor& out, torch::Tensor const& a,
176 torch::Tensor const& b_scales,
177 std::optional<torch::Tensor> const& bias);
178
179 +void cutlass_moe_mm(
180 + torch::Tensor& out_tensors, torch::Tensor const& a_tensors,
181 + torch::Tensor const& b_tensors, torch::Tensor const& a_scales,
182 + torch::Tensor const& b_scales, torch::Tensor const& expert_offsets,
183 + torch::Tensor const& problem_sizes, torch::Tensor const& a_strides,
184 + torch::Tensor const& b_strides, torch::Tensor const& c_strides);
185 +
186 +void get_cutlass_moe_mm_data(
187 + const torch::Tensor& topk_ids, torch::Tensor& expert_offsets,
188 + torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
189 + torch::Tensor& input_permutation, torch::Tensor& output_permutation,
190 + const int64_t num_experts, const int64_t n, const int64_t k);
191 +
192 void cutlass_scaled_mm_azp(torch::Tensor& out, torch::Tensor const& a,
193 torch::Tensor const& b,
194 torch::Tensor const& a_scales,
csrc/quantization/cutlass_w8a8/moe/get_group_starts.cuh new
+80
@@ -0,0 +1,80 @@
1 +#pragma once
2 +
3 +#include <cuda.h>
4 +#include <torch/all.h>
5 +#include <c10/cuda/CUDAStream.h>
6 +
7 +#include "core/scalar_type.hpp"
8 +#include "cutlass/bfloat16.h"
9 +#include "cutlass/float8.h"
10 +
11 +template <typename ElementAB, typename ElementC, typename ElementAccumulator>
12 +__global__ void get_group_gemm_starts(
13 + int32_t* expert_offsets, ElementAB** a_offsets, ElementAB** b_offsets,
14 + ElementC** out_offsets, ElementAccumulator** a_scales_offsets,
15 + ElementAccumulator** b_scales_offsets, ElementAB* a_base_as_int,
16 + ElementAB* b_base_as_int, ElementC* out_base_as_int,
17 + ElementAccumulator* a_scales_base_as_int,
18 + ElementAccumulator* b_scales_base_as_int, int64_t n, int64_t k,
19 + bool per_act_token, bool per_out_ch) {
20 + int expert_id = threadIdx.x;
21 +
22 + int64_t expert_offset = expert_offsets[expert_id];
23 +
24 + a_offsets[expert_id] = a_base_as_int + expert_offset * k;
25 + b_offsets[expert_id] = b_base_as_int + expert_id * k * n;
26 + out_offsets[expert_id] = out_base_as_int + expert_offset * n;
27 + a_scales_offsets[expert_id] =
28 + a_scales_base_as_int + (per_act_token ? expert_offset : 0);
29 + b_scales_offsets[expert_id] =
30 + b_scales_base_as_int + (per_out_ch ? n * expert_id : expert_id);
31 +}
32 +
33 +#define __CALL_GET_STARTS_KERNEL(TENSOR_C_TYPE, C_TYPE) \
34 + else if (out_tensors.dtype() == TENSOR_C_TYPE) { \
35 + get_group_gemm_starts<cutlass::float_e4m3_t, C_TYPE, float> \
36 + <<<1, num_experts, 0, stream>>>( \
37 + static_cast<int32_t*>(expert_offsets.data_ptr()), \
38 + static_cast<cutlass::float_e4m3_t**>(a_ptrs.data_ptr()), \
39 + static_cast<cutlass::float_e4m3_t**>(b_ptrs.data_ptr()), \
40 + static_cast<C_TYPE**>(out_ptrs.data_ptr()), \
41 + static_cast<float**>(a_scales_ptrs.data_ptr()), \
42 + static_cast<float**>(b_scales_ptrs.data_ptr()), \
43 + static_cast<cutlass::float_e4m3_t*>(a_tensors.data_ptr()), \
44 + static_cast<cutlass::float_e4m3_t*>(b_tensors.data_ptr()), \
45 + static_cast<C_TYPE*>(out_tensors.data_ptr()), \
46 + static_cast<float*>(a_scales.data_ptr()), \
47 + static_cast<float*>(b_scales.data_ptr()), out_tensors.size(1), \
48 + a_tensors.size(1), per_act_token, per_out_ch); \
49 + }
50 +
51 +namespace {
52 +
53 +void run_get_group_gemm_starts(
54 + torch::Tensor const& expert_offsets, torch::Tensor& a_ptrs,
55 + torch::Tensor& b_ptrs, torch::Tensor& out_ptrs,
56 + torch::Tensor& a_scales_ptrs, torch::Tensor& b_scales_ptrs,
57 + torch::Tensor const& a_tensors, torch::Tensor const& b_tensors,
58 + torch::Tensor& out_tensors, torch::Tensor const& a_scales,
59 + torch::Tensor const& b_scales) {
60 + TORCH_CHECK(a_tensors.dtype() == torch::kFloat8_e4m3fn);
61 + TORCH_CHECK(b_tensors.dtype() == torch::kFloat8_e4m3fn);
62 + TORCH_CHECK(a_scales.dtype() == torch::kFloat32);
63 + TORCH_CHECK(b_scales.dtype() == torch::kFloat32);
64 +
65 + int num_experts = static_cast<int>(expert_offsets.size(0));
66 + bool per_act_token = a_scales.numel() != 1;
67 + bool per_out_ch = b_scales.numel() != num_experts;
68 +
69 + auto stream = at::cuda::getCurrentCUDAStream(a_tensors.device().index());
70 +
71 + if (false) {
72 + }
73 + __CALL_GET_STARTS_KERNEL(torch::kBFloat16, cutlass::bfloat16_t)
74 + __CALL_GET_STARTS_KERNEL(torch::kFloat16, half)
75 + else {
76 + TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)");
77 + }
78 +}
79 +
80 +} // namespace
\ No newline at end of file
csrc/quantization/cutlass_w8a8/moe/grouped_mm_c3x.cu new
+160
@@ -0,0 +1,160 @@
1 +#include <cudaTypedefs.h>
2 +
3 +#include <c10/cuda/CUDAGuard.h>
4 +#include <torch/all.h>
5 +
6 +#include "cutlass/cutlass.h"
7 +#include "grouped_mm_c3x.cuh"
8 +
9 +using namespace cute;
10 +
11 +namespace {
12 +
13 +template <typename InType, typename OutType,
14 + template <typename, typename, typename> typename Epilogue>
15 +struct sm90_fp8_config_default {
16 + // M in (16, inf)
17 + static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
18 + using KernelSchedule =
19 + cutlass::gemm::KernelPtrArrayTmaWarpSpecializedPingpongFP8FastAccum;
20 + using EpilogueSchedule =
21 + cutlass::epilogue::PtrArrayTmaWarpSpecializedPingpong;
22 + using TileShape = cute::Shape<cute::_64, cute::_256, cute::_128>;
23 + using ClusterShape = cute::Shape<cute::_1, cute::_2, cute::_1>;
24 +
25 + using Cutlass3xGemm =
26 + cutlass_3x_group_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
27 + KernelSchedule, EpilogueSchedule>;
28 +};
29 +
30 +template <typename InType, typename OutType,
31 + template <typename, typename, typename> typename Epilogue>
32 +struct sm90_fp8_config_M16 {
33 + // M in [1, 16]
34 + static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
35 + using KernelSchedule =
36 + cutlass::gemm::KernelPtrArrayTmaWarpSpecializedPingpongFP8FastAccum;
37 + using EpilogueSchedule =
38 + cutlass::epilogue::PtrArrayTmaWarpSpecializedPingpong;
39 + using TileShape = cute::Shape<cute::_64, cute::_64, cute::_128>;
40 + using ClusterShape = cute::Shape<cute::_1, cute::_4, cute::_1>;
41 +
42 + using Cutlass3xGemm =
43 + cutlass_3x_group_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
44 + KernelSchedule, EpilogueSchedule>;
45 +};
46 +
47 +template <typename InType, typename OutType,
48 + template <typename, typename, typename> typename Epilogue>
49 +struct sm90_fp8_config_K8192 {
50 + // K in [8192, inf)
51 + static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
52 + using KernelSchedule =
53 + cutlass::gemm::KernelPtrArrayTmaWarpSpecializedPingpongFP8FastAccum;
54 + using EpilogueSchedule =
55 + cutlass::epilogue::PtrArrayTmaWarpSpecializedPingpong;
56 + using TileShape = cute::Shape<cute::_128, cute::_128, cute::_128>;
57 + using ClusterShape = cute::Shape<cute::_1, cute::_8, cute::_1>;
58 +
59 + using Cutlass3xGemm =
60 + cutlass_3x_group_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
61 + KernelSchedule, EpilogueSchedule>;
62 +};
63 +
64 +template <typename InType, typename OutType,
65 + template <typename, typename, typename> typename Epilogue>
66 +struct sm90_fp8_config_N8192 {
67 + // N in [8192, inf)
68 + static_assert(std::is_same<InType, cutlass::float_e4m3_t>());
69 + using KernelSchedule =
70 + cutlass::gemm::KernelPtrArrayTmaWarpSpecializedPingpongFP8FastAccum;
71 + using EpilogueSchedule =
72 + cutlass::epilogue::PtrArrayTmaWarpSpecializedPingpong;
73 + using TileShape = cute::Shape<cute::_64, cute::_128, cute::_256>;
74 + using ClusterShape = cute::Shape<cute::_1, cute::_8, cute::_1>;
75 +
76 + using Cutlass3xGemm =
77 + cutlass_3x_group_gemm<InType, OutType, Epilogue, TileShape, ClusterShape,
78 + KernelSchedule, EpilogueSchedule>;
79 +};
80 +
81 +template <typename InType, typename OutType>
82 +void run_cutlass_moe_mm_sm90(
83 + torch::Tensor& out_tensors, torch::Tensor const& a_tensors,
84 + torch::Tensor const& b_tensors, torch::Tensor const& a_scales,
85 + torch::Tensor const& b_scales, torch::Tensor const& expert_offsets,
86 + torch::Tensor const& problem_sizes, torch::Tensor const& a_strides,
87 + torch::Tensor const& b_strides, torch::Tensor const& c_strides) {
88 + TORCH_CHECK(a_tensors.size(0) > 0, "No input A tensors provided.");
89 + TORCH_CHECK(b_tensors.size(0) > 0, "No input B tensors provided.");
90 + TORCH_CHECK(out_tensors.size(0) > 0, "No output tensors provided.");
91 +
92 + TORCH_CHECK(a_tensors.dtype() == torch::kFloat8_e4m3fn,
93 + "A tensors must be of type float8_e4m3fn.");
94 + TORCH_CHECK(b_tensors.dtype() == torch::kFloat8_e4m3fn,
95 + "B tensors must be of type float8_e4m3fn.");
96 +
97 + TORCH_CHECK(a_tensors.dtype() == torch::kFloat8_e4m3fn);
98 + TORCH_CHECK(b_tensors.dtype() == torch::kFloat8_e4m3fn);
99 +
100 + using Cutlass3xGemmN8192 = typename sm90_fp8_config_N8192<
101 + InType, OutType, vllm::c3x::ScaledEpilogueArray>::Cutlass3xGemm;
102 + using Cutlass3xGemmK8192 = typename sm90_fp8_config_K8192<
103 + InType, OutType, vllm::c3x::ScaledEpilogueArray>::Cutlass3xGemm;
104 + using Cutlass3xGemmM16 = typename sm90_fp8_config_M16<
105 + InType, OutType, vllm::c3x::ScaledEpilogueArray>::Cutlass3xGemm;
106 + using Cutlass3xGemmDefault = typename sm90_fp8_config_default<
107 + InType, OutType, vllm::c3x::ScaledEpilogueArray>::Cutlass3xGemm;
108 +
109 + uint32_t const m = a_tensors.size(0);
110 + uint32_t const n = out_tensors.size(1);
111 + uint32_t const k = a_tensors.size(1);
112 +
113 + if (n >= 8192) {
114 + cutlass_group_gemm_caller<Cutlass3xGemmN8192>(
115 + out_tensors, a_tensors, b_tensors, a_scales, b_scales, expert_offsets,
116 + problem_sizes, a_strides, b_strides, c_strides);
117 + } else if (k >= 8192) {
118 + cutlass_group_gemm_caller<Cutlass3xGemmK8192>(
119 + out_tensors, a_tensors, b_tensors, a_scales, b_scales, expert_offsets,
120 + problem_sizes, a_strides, b_strides, c_strides);
121 + } else if (m <= 16) {
122 + cutlass_group_gemm_caller<Cutlass3xGemmM16>(
123 + out_tensors, a_tensors, b_tensors, a_scales, b_scales, expert_offsets,
124 + problem_sizes, a_strides, b_strides, c_strides);
125 + } else {
126 + cutlass_group_gemm_caller<Cutlass3xGemmDefault>(
127 + out_tensors, a_tensors, b_tensors, a_scales, b_scales, expert_offsets,
128 + problem_sizes, a_strides, b_strides, c_strides);
129 + }
130 +}
131 +
132 +void dispatch_moe_mm_sm90(
133 + torch::Tensor& out_tensors, torch::Tensor const& a_tensors,
134 + torch::Tensor const& b_tensors, torch::Tensor const& a_scales,
135 + torch::Tensor const& b_scales, torch::Tensor const& expert_offsets,
136 + torch::Tensor const& problem_sizes, torch::Tensor const& a_strides,
137 + torch::Tensor const& b_strides, torch::Tensor const& c_strides) {
138 + if (out_tensors.dtype() == torch::kBFloat16) {
139 + run_cutlass_moe_mm_sm90<cutlass::float_e4m3_t, cutlass::bfloat16_t>(
140 + out_tensors, a_tensors, b_tensors, a_scales, b_scales, expert_offsets,
141 + problem_sizes, a_strides, b_strides, c_strides);
142 + } else {
143 + run_cutlass_moe_mm_sm90<cutlass::float_e4m3_t, cutlass::half_t>(
144 + out_tensors, a_tensors, b_tensors, a_scales, b_scales, expert_offsets,
145 + problem_sizes, a_strides, b_strides, c_strides);
146 + }
147 +}
148 +
149 +} // namespace
150 +
151 +void cutlass_moe_mm_sm90(
152 + torch::Tensor& out_tensors, torch::Tensor const& a_tensors,
153 + torch::Tensor const& b_tensors, torch::Tensor const& a_scales,
154 + torch::Tensor const& b_scales, torch::Tensor const& expert_offsets,
155 + torch::Tensor const& problem_sizes, torch::Tensor const& a_strides,
156 + torch::Tensor const& b_strides, torch::Tensor const& c_strides) {
157 + dispatch_moe_mm_sm90(out_tensors, a_tensors, b_tensors, a_scales, b_scales,
158 + expert_offsets, problem_sizes, a_strides, b_strides,
159 + c_strides);
160 +}
csrc/quantization/cutlass_w8a8/moe/grouped_mm_c3x.cuh new
+149
@@ -0,0 +1,149 @@
1 +#pragma once
2 +
3 +#include "cutlass/cutlass.h"
4 +
5 +#include "cutlass/gemm/collective/collective_builder.hpp"
6 +#include "cutlass/epilogue/collective/collective_builder.hpp"
7 +#include "cutlass/gemm/device/gemm_universal_adapter.h"
8 +
9 +#include "cutlass_extensions/epilogue/scaled_mm_epilogues_c3x.hpp"
10 +#include "cutlass_extensions/common.hpp"
11 +#include "get_group_starts.cuh"
12 +
13 +using namespace cute;
14 +
15 +namespace {
16 +
17 +using ProblemShape =
18 + cutlass::gemm::GroupProblemShape<cute::Shape<int, int, int>>;
19 +
20 +using ElementAccumulator = float;
21 +using ArchTag = cutlass::arch::Sm90;
22 +using OperatorClass = cutlass::arch::OpClassTensorOp;
23 +
24 +using LayoutA = cutlass::layout::RowMajor;
25 +using LayoutB = cutlass::layout::ColumnMajor;
26 +using LayoutC = cutlass::layout::RowMajor;
27 +
28 +template <typename ElementAB_, typename ElementC_,
29 + template <typename, typename, typename> typename Epilogue_,
30 + typename TileShape, typename ClusterShape, typename KernelSchedule,
31 + typename EpilogueSchedule>
32 +struct cutlass_3x_group_gemm {
33 + using ElementAB = ElementAB_;
34 + using ElementC = void;
35 + using ElementD = ElementC_;
36 + using ElementAccumulator = float;
37 +
38 + using Epilogue = Epilogue_<ElementAccumulator, ElementD, TileShape>;
39 +
40 + using StrideC =
41 + cute::remove_pointer_t<cute::Stride<int64_t, cute::Int<1>, cute::Int<0>>>;
42 +
43 + static constexpr int AlignmentAB =
44 + 128 / cutlass::sizeof_bits<ElementAB>::value;
45 + static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementD>::value;
46 +
47 + using EVTCompute = typename Epilogue::EVTCompute;
48 +
49 + using CollectiveEpilogue =
50 + typename cutlass::epilogue::collective::CollectiveBuilder<
51 + ArchTag, OperatorClass, TileShape, ClusterShape,
52 + cutlass::epilogue::collective::EpilogueTileAuto, ElementAccumulator,
53 + ElementAccumulator, ElementC, LayoutC*, AlignmentC, ElementD,
54 + LayoutC*, AlignmentC, EpilogueSchedule, EVTCompute>::CollectiveOp;
55 +
56 + static constexpr size_t CEStorageSize =
57 + sizeof(typename CollectiveEpilogue::SharedStorage);
58 + using Stages = typename cutlass::gemm::collective::StageCountAutoCarveout<
59 + static_cast<int>(CEStorageSize)>;
60 +
61 + using CollectiveMainloop =
62 + typename cutlass::gemm::collective::CollectiveBuilder<
63 + ArchTag, OperatorClass, ElementAB, LayoutA*, AlignmentAB, ElementAB,
64 + LayoutB*, AlignmentAB, ElementAccumulator, TileShape, ClusterShape,
65 + Stages, KernelSchedule>::CollectiveOp;
66 +
67 + using KernelType = enable_sm90_only<cutlass::gemm::kernel::GemmUniversal<
68 + ProblemShape, CollectiveMainloop, CollectiveEpilogue>>;
69 +
70 + struct GemmKernel : public KernelType {};
71 +};
72 +
73 +template <typename Gemm>
74 +void cutlass_group_gemm_caller(
75 + torch::Tensor& out_tensors, torch::Tensor const& a_tensors,
76 + torch::Tensor const& b_tensors, torch::Tensor const& a_scales,
77 + torch::Tensor const& b_scales, torch::Tensor const& expert_offsets,
78 + torch::Tensor const& problem_sizes, torch::Tensor const& a_strides,
79 + torch::Tensor const& b_strides, torch::Tensor const& c_strides) {
80 + using ElementAB = typename Gemm::ElementAB;
81 + using ElementD = typename Gemm::ElementD;
82 +
83 + int num_experts = static_cast<int>(expert_offsets.size(0));
84 + int k_size = a_tensors.size(1);
85 + int n_size = out_tensors.size(1);
86 +
87 + bool per_act_token = a_scales.numel() != 1;
88 + bool per_out_ch = b_scales.numel() != num_experts;
89 +
90 + auto stream = at::cuda::getCurrentCUDAStream(a_tensors.device().index());
91 +
92 + auto options_int =
93 + torch::TensorOptions().dtype(torch::kInt64).device(a_tensors.device());
94 +
95 + torch::Tensor a_ptrs = torch::empty(num_experts, options_int);
96 + torch::Tensor b_ptrs = torch::empty(num_experts, options_int);
97 + torch::Tensor out_ptrs = torch::empty(num_experts, options_int);
98 + torch::Tensor a_scales_ptrs = torch::empty(num_experts, options_int);
99 + torch::Tensor b_scales_ptrs = torch::empty(num_experts, options_int);
100 +
101 + run_get_group_gemm_starts(expert_offsets, a_ptrs, b_ptrs, out_ptrs,
102 + a_scales_ptrs, b_scales_ptrs, a_tensors, b_tensors,
103 + out_tensors, a_scales, b_scales);
104 +
105 + using GemmKernel = typename Gemm::GemmKernel;
106 + using StrideA = Stride<int64_t, Int<1>, Int<0>>;
107 + using StrideB = Stride<int64_t, Int<1>, Int<0>>;
108 + using StrideC = typename GemmKernel::InternalStrideC;
109 +
110 + ProblemShape::UnderlyingProblemShape* problem_sizes_as_shapes =
111 + static_cast<ProblemShape::UnderlyingProblemShape*>(
112 + problem_sizes.data_ptr());
113 + ProblemShape prob_shape{num_experts, problem_sizes_as_shapes, nullptr};
114 +
115 + typename GemmKernel::MainloopArguments mainloop_args{
116 + static_cast<const ElementAB**>(a_ptrs.data_ptr()),
117 + static_cast<StrideA*>(a_strides.data_ptr()),
118 + static_cast<const ElementAB**>(b_ptrs.data_ptr()),
119 + static_cast<StrideB*>(b_strides.data_ptr())};
120 +
121 + // Currently, we are only able to do broadcast on either all or none a_scales
122 + // and on either all or none b_scales
123 + typename GemmKernel::EpilogueArguments epilogue_args{
124 + Gemm::Epilogue::prepare_args(
125 + static_cast<const ElementAccumulator**>(a_scales_ptrs.data_ptr()),
126 + static_cast<const ElementAccumulator**>(b_scales_ptrs.data_ptr()),
127 + per_act_token, per_out_ch),
128 + nullptr, static_cast<StrideC*>(c_strides.data_ptr()),
129 + static_cast<ElementD**>(out_ptrs.data_ptr()),
130 + static_cast<StrideC*>(c_strides.data_ptr())};
131 +
132 + typename GemmKernel::Arguments args{
133 + cutlass::gemm::GemmUniversalMode::kGrouped, prob_shape, mainloop_args,
134 + epilogue_args};
135 +
136 + using GemmOp = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
137 + GemmOp gemm_op;
138 + CUTLASS_CHECK(gemm_op.can_implement(args));
139 +
140 + size_t workspace_size = gemm_op.get_workspace_size(args);
141 + auto const workspace_options =
142 + torch::TensorOptions().dtype(torch::kUInt8).device(a_tensors.device());
143 + auto workspace = torch::empty(workspace_size, workspace_options);
144 +
145 + cutlass::Status status = gemm_op.run(args, workspace.data_ptr(), stream);
146 + CUTLASS_CHECK(status);
147 +}
148 +
149 +} // namespace
csrc/quantization/cutlass_w8a8/moe/moe_data.cu new
+90
@@ -0,0 +1,90 @@
1 +#include <cudaTypedefs.h>
2 +
3 +#include <c10/cuda/CUDAGuard.h>
4 +#include <torch/all.h>
5 +
6 +#include <iostream>
7 +
8 +constexpr uint64_t THREADS_PER_EXPERT = 512;
9 +
10 +__global__ void compute_problem_sizes(const int* __restrict__ topk_ids,
11 + int32_t* problem_sizes1,
12 + int32_t* problem_sizes2,
13 + int32_t* atomic_buffer,
14 + const int topk_length, const int n,
15 + const int k) {
16 + int expert_id = blockIdx.x;
17 +
18 + int occurrences = 0;
19 + for (int i = threadIdx.x; i < topk_length; i += THREADS_PER_EXPERT) {
20 + occurrences += (topk_ids[i] == expert_id);
21 + }
22 + atomicAdd(&atomic_buffer[expert_id], occurrences);
23 + __syncthreads();
24 +
25 + if (threadIdx.x == 0) {
26 + int final_occurrences = atomic_buffer[expert_id];
27 + problem_sizes1[expert_id * 3] = final_occurrences;
28 + problem_sizes1[expert_id * 3 + 1] = 2 * n;
29 + problem_sizes1[expert_id * 3 + 2] = k;
30 + problem_sizes2[expert_id * 3] = final_occurrences;
31 + problem_sizes2[expert_id * 3 + 1] = k;
32 + problem_sizes2[expert_id * 3 + 2] = n;
33 + }
34 +}
35 +
36 +__global__ void compute_expert_offsets(
37 + const int32_t* __restrict__ problem_sizes1, int32_t* expert_offsets,
38 + int32_t* atomic_buffer, const int num_experts) {
39 + int32_t tot_offset = 0;
40 + expert_offsets[0] = 0;
41 + for (int i = 0; i < num_experts; ++i) {
42 + atomic_buffer[i] = tot_offset;
43 + tot_offset += problem_sizes1[i * 3];
44 + expert_offsets[i + 1] = tot_offset;
45 + }
46 +}
47 +
48 +__global__ void compute_arg_sorts(const int* __restrict__ topk_ids,
49 + int32_t* input_permutation,
50 + int32_t* output_permutation,
51 + int32_t* atomic_buffer, const int topk_length,
52 + const int topk) {
53 + int expert_id = blockIdx.x;
54 +
55 + for (int i = threadIdx.x; i < topk_length; i += THREADS_PER_EXPERT) {
56 + if (topk_ids[i] == expert_id) {
57 + int start = atomicAdd(&atomic_buffer[expert_id], 1);
58 + input_permutation[start] = i / topk;
59 + output_permutation[i] = start;
60 + }
61 + }
62 +}
63 +
64 +void get_cutlass_moe_mm_data_caller(
65 + const torch::Tensor& topk_ids, torch::Tensor& expert_offsets,
66 + torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
67 + torch::Tensor& input_permutation, torch::Tensor& output_permutation,
68 + const int64_t num_experts, const int64_t n, const int64_t k) {
69 + auto stream = at::cuda::getCurrentCUDAStream(topk_ids.device().index());
70 + auto options_int32 =
71 + torch::TensorOptions().dtype(torch::kInt32).device(topk_ids.device());
72 + torch::Tensor atomic_buffer = torch::zeros(num_experts, options_int32);
73 +
74 + int num_threads = min(THREADS_PER_EXPERT, topk_ids.numel());
75 + compute_problem_sizes<<<num_experts, num_threads, 0, stream>>>(
76 + static_cast<const int32_t*>(topk_ids.data_ptr()),
77 + static_cast<int32_t*>(problem_sizes1.data_ptr()),
78 + static_cast<int32_t*>(problem_sizes2.data_ptr()),
79 + static_cast<int32_t*>(atomic_buffer.data_ptr()), topk_ids.numel(), n, k);
80 + compute_expert_offsets<<<1, 1, 0, stream>>>(
81 + static_cast<const int32_t*>(problem_sizes1.data_ptr()),
82 + static_cast<int32_t*>(expert_offsets.data_ptr()),
83 + static_cast<int32_t*>(atomic_buffer.data_ptr()), num_experts);
84 + compute_arg_sorts<<<num_experts, num_threads, 0, stream>>>(
85 + static_cast<const int32_t*>(topk_ids.data_ptr()),
86 + static_cast<int32_t*>(input_permutation.data_ptr()),
87 + static_cast<int32_t*>(output_permutation.data_ptr()),
88 + static_cast<int32_t*>(atomic_buffer.data_ptr()), topk_ids.numel(),
89 + topk_ids.size(1));
90 +}
csrc/quantization/cutlass_w8a8/scaled_mm_entry.cu
+67
@@ -29,6 +29,20 @@ void cutlass_scaled_mm_sm90(torch::Tensor& c, torch::Tensor const& a,
29 torch::Tensor const& a_scales,
30 torch::Tensor const& b_scales,
31 std::optional<torch::Tensor> const& bias);
32 +
33 +void cutlass_moe_mm_sm90(
34 + torch::Tensor& out_tensors, torch::Tensor const& a_tensors,
35 + torch::Tensor const& b_tensors, torch::Tensor const& a_scales,
36 + torch::Tensor const& b_scales, torch::Tensor const& expert_offsets,
37 + torch::Tensor const& problem_sizes, torch::Tensor const& a_strides,
38 + torch::Tensor const& b_strides, torch::Tensor const& c_strides);
39 +
40 +void get_cutlass_moe_mm_data_caller(
41 + const torch::Tensor& topk_ids, torch::Tensor& expert_offsets,
42 + torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
43 + torch::Tensor& input_permutation, torch::Tensor& output_permutation,
44 + const int64_t num_experts, const int64_t n, const int64_t k);
45 +
46 #endif
47
48 #if defined ENABLE_SCALED_MM_SM100 && ENABLE_SCALED_MM_SM100
@@ -102,6 +116,19 @@ bool cutlass_scaled_mm_supports_block_fp8(int64_t cuda_device_capability) {
116 return false;
117 }
118
119 +bool cutlass_group_gemm_supported(int64_t cuda_device_capability) {
120 + // CUTLASS groped FP8 kernels need at least CUDA 12.3
121 + // and SM90 (Hopper)
122 +
123 +#if defined CUDA_VERSION
124 + if (cuda_device_capability == 90) {
125 + return CUDA_VERSION >= 12030;
126 + }
127 +#endif
128 +
129 + return false;
130 +}
131 +
132 void cutlass_scaled_mm(torch::Tensor& c, torch::Tensor const& a,
133 torch::Tensor const& b, torch::Tensor const& a_scales,
134 torch::Tensor const& b_scales,
@@ -168,6 +195,46 @@ void cutlass_scaled_mm(torch::Tensor& c, torch::Tensor const& a,
195 version_num);
196 }
197
198 +void cutlass_moe_mm(
199 + torch::Tensor& out_tensors, torch::Tensor const& a_tensors,
200 + torch::Tensor const& b_tensors, torch::Tensor const& a_scales,
201 + torch::Tensor const& b_scales, torch::Tensor const& expert_offsets,
202 + torch::Tensor const& problem_sizes, torch::Tensor const& a_strides,
203 + torch::Tensor const& b_strides, torch::Tensor const& c_strides) {
204 + int32_t version_num = get_sm_version_num();
205 +#if defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90
206 + cutlass_moe_mm_sm90(out_tensors, a_tensors, b_tensors, a_scales, b_scales,
207 + expert_offsets, problem_sizes, a_strides, b_strides,
208 + c_strides);
209 + return;
210 +#endif
211 + TORCH_CHECK_NOT_IMPLEMENTED(
212 + false,
213 + "No compiled cutlass_scaled_mm for CUDA device capability: ", version_num,
214 + ". Required capability: 90");
215 +}
216 +
217 +void get_cutlass_moe_mm_data(
218 + const torch::Tensor& topk_ids, torch::Tensor& expert_offsets,
219 + torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
220 + torch::Tensor& input_permutation, torch::Tensor& output_permutation,
221 + const int64_t num_experts, const int64_t n, const int64_t k) {
222 + // This function currently gets compiled only if we have a valid cutlass moe
223 + // mm to run it for.
224 + int32_t version_num = get_sm_version_num();
225 +#if defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90
226 + get_cutlass_moe_mm_data_caller(topk_ids, expert_offsets, problem_sizes1,
227 + problem_sizes2, input_permutation,
228 + output_permutation, num_experts, n, k);
229 + return;
230 +#endif
231 + TORCH_CHECK_NOT_IMPLEMENTED(
232 + false,
233 + "No compiled get_cutlass_moe_mm_data: no cutlass_scaled_mm kernel for "
234 + "CUDA device capability: ",
235 + version_num, ". Required capability: 90");
236 +}
237 +
238 void cutlass_scaled_mm_azp(torch::Tensor& c, torch::Tensor const& a,
239 torch::Tensor const& b,
240 torch::Tensor const& a_scales,
csrc/torch_bindings.cpp
+29
@@ -365,6 +365,35 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
365 ops.def("cutlass_scaled_mm_supports_fp8(int cuda_device_capability) -> bool");
366 ops.impl("cutlass_scaled_mm_supports_fp8", &cutlass_scaled_mm_supports_fp8);
367
368 + // Check if cutlass grouped gemm is supported for CUDA devices of the given
369 + // capability
370 + ops.def("cutlass_group_gemm_supported(int cuda_device_capability) -> bool");
371 + ops.impl("cutlass_group_gemm_supported", &cutlass_group_gemm_supported);
372 +
373 + // CUTLASS w8a8 grouped GEMM
374 + ops.def(
375 + "cutlass_moe_mm(Tensor! out_tensors, Tensor a_tensors, Tensor b_tensors, "
376 + " Tensor a_scales, Tensor b_scales, Tensor expert_offsets, "
377 + " Tensor problem_sizes, Tensor a_strides, "
378 + " Tensor b_strides, Tensor c_strides) -> ()",
379 + {stride_tag});
380 + ops.impl("cutlass_moe_mm", torch::kCUDA, &cutlass_moe_mm);
381 +
382 + // A function that computes data required to run fused MoE with w8a8 grouped
383 + // GEMM. It takes topk_ids as an input, and computes expert_offsets
384 + // (token start indices of each expert). In addition to this, it computes
385 + // problem sizes for each expert's multiplication used by the two mms called
386 + // from fused MoE operation, and arrays with permutations required to shuffle
387 + // and de-shuffle the input/output of the fused operation.
388 + ops.def(
389 + "get_cutlass_moe_mm_data(Tensor topk_ids, Tensor! expert_offsets, "
390 + " Tensor! problem_sizes1, Tensor! problem_sizes2, "
391 + " Tensor! input_permutation, "
392 + " Tensor! output_permutation, int num_experts, "
393 + " int n, int k) -> ()",
394 + {stride_tag});
395 + ops.impl("get_cutlass_moe_mm_data", torch::kCUDA, &get_cutlass_moe_mm_data);
396 +
397 // Check if cutlass scaled_mm supports block quantization (used by DeepSeekV3)
398 ops.def(
399 "cutlass_scaled_mm_supports_block_fp8(int cuda_device_capability) -> "
tests/kernels/test_cutlass.py
+134
@@ -3,6 +3,7 @@
3
4 Run `pytest tests/kernels/test_cutlass.py`.
5 """
6 +import random
7
8 import pytest
9 import torch
@@ -507,3 +508,136 @@ def test_cutlass_cuda_graph(per_act_token: bool, per_out_ch: bool):
508
509 def test_cutlass_support_opcheck():
510 opcheck(torch.ops._C.cutlass_scaled_mm_supports_fp8, (capability, ))
511 +
512 +
513 +@pytest.mark.parametrize("num_experts", [8, 64])
514 +@pytest.mark.parametrize("per_act_token", [True, False])
515 +@pytest.mark.parametrize("per_out_ch", [True, False])
516 +@pytest.mark.parametrize("use_bias", [False])
517 +@pytest.mark.skipif(
518 + (lambda x: x is None or not ops.cutlass_group_gemm_supported(x.to_int()))(
519 + current_platform.get_device_capability()),
520 + reason="Grouped gemm is not supported on this GPU type.")
521 +def test_cutlass_fp8_group_gemm(num_experts: int, per_act_token: bool,
522 + per_out_ch: bool, use_bias: bool):
523 +
524 + # Device and dtype setup
525 + device = "cuda"
526 + out_dtype = torch.half
527 +
528 + # Create separate A, B, C tensors for each group
529 + a_tensors = []
530 + b_tensors = []
531 + a_scales_tensors = []
532 + b_scales_tensors = []
533 + baseline_tensors = []
534 +
535 + expert_offsets = torch.zeros((num_experts + 1),
536 + device=device,
537 + dtype=torch.int32)
538 +
539 + problem_sizes = torch.zeros((num_experts, 3),
540 + device=device,
541 + dtype=torch.int32)
542 +
543 + if not per_act_token:
544 + one_scale_a = torch.randn((1, 1), device=device, dtype=torch.float32)
545 +
546 + alignment = 16 # 128 // 8
547 + # For variation, each group has dimensions
548 + n_g = alignment * random.randint(1, 64)
549 + k_g = alignment * random.randint(1, 64)
550 + for g in range(num_experts):
551 + m_g = alignment * random.randint(1, 64)
552 +
553 + expert_offsets[g + 1] = expert_offsets[g] + m_g
554 + problem_sizes[g][0] = m_g
555 + problem_sizes[g][1] = n_g
556 + problem_sizes[g][2] = k_g
557 +
558 + m_a_scales = m_g if per_act_token else 1
559 + n_b_scales = n_g if per_out_ch else 1
560 +
561 + print("shape:", m_g, n_g, k_g)
562 +
563 + # Create group-specific A and B (FP8) and output (FP16/FP32)
564 + a_g = to_fp8(torch.randn((m_g, k_g), device=device))
565 + b_g = to_fp8(torch.randn((n_g, k_g), device=device).t())
566 + a_tensors.append(a_g)
567 + b_tensors.append(b_g)
568 +
569 + # Set up A/B scales
570 + scale_b = torch.randn((1, n_b_scales),
571 + device=device,
572 + dtype=torch.float32)
573 + b_scales_tensors.append(scale_b)
574 +
575 + if per_act_token:
576 + scale_a = torch.randn((m_a_scales, 1),
577 + device=device,
578 + dtype=torch.float32)
579 + a_scales_tensors.append(scale_a)
580 + else:
581 + scale_a = one_scale_a
582 +
583 + # Compute baseline result for this group
584 + baseline_g = baseline_scaled_mm(a_g, b_g, scale_a, scale_b, out_dtype,
585 + None)
586 + baseline_tensors.append(baseline_g)
587 +
588 + a_tensors_stacked = torch.empty((expert_offsets[num_experts], k_g),
589 + device=device,
590 + dtype=torch.float8_e4m3fn)
591 + b_tensors_stacked = torch.empty((num_experts, n_g, k_g),
592 + device=device,
593 + dtype=torch.float8_e4m3fn)
594 +
595 + for g in range(num_experts):
596 + a_tensors_stacked[expert_offsets[g]:expert_offsets[g +
597 + 1]] = a_tensors[g]
598 + b_tensors_stacked[g] = b_tensors[g].t()
599 + b_tensors_stacked = b_tensors_stacked.transpose(1, 2)
600 +
601 + if per_act_token:
602 + a_scales_tensors_stacked = torch.empty(
603 + (expert_offsets[num_experts], 1),
604 + device=device,
605 + dtype=torch.float32)
606 + for g in range(num_experts):
607 + a_scales_tensors_stacked[
608 + expert_offsets[g]:expert_offsets[g + 1]] = a_scales_tensors[g]
609 + else:
610 + a_scales_tensors_stacked = one_scale_a
611 +
612 + b_scales_tensors_stacked = torch.empty((num_experts, n_b_scales),
613 + device=device,
614 + dtype=torch.float32)
615 + for g in range(num_experts):
616 + b_scales_tensors_stacked[g] = b_scales_tensors[g]
617 +
618 + out_tensors_stacked = torch.zeros((expert_offsets[num_experts], n_g),
619 + device=device,
620 + dtype=out_dtype)
621 +
622 + ab_strides = torch.full((num_experts, ),
623 + a_tensors_stacked.stride(0),
624 + device="cuda",
625 + dtype=torch.int64)
626 + c_strides = torch.full((num_experts, ),
627 + out_tensors_stacked.stride(0),
628 + device="cuda",
629 + dtype=torch.int64)
630 +
631 + ops.cutlass_moe_mm(out_tensors_stacked, a_tensors_stacked,
632 + b_tensors_stacked, a_scales_tensors_stacked,
633 + b_scales_tensors_stacked, expert_offsets[:-1],
634 + problem_sizes, ab_strides, ab_strides, c_strides)
635 +
636 + # Validate each group's result against the baseline
637 + for g in range(num_experts):
638 + baseline = baseline_tensors[g]
639 + c = out_tensors_stacked[expert_offsets[g]:expert_offsets[g + 1]]
640 + print(baseline)
641 + print(c)
642 + print("*")
643 + torch.testing.assert_close(c, baseline, rtol=1e-2, atol=5e-4)
tests/kernels/test_cutlass_moe.py new
+244
@@ -0,0 +1,244 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +import pytest
3 +import torch
4 +
5 +from vllm import _custom_ops as ops
6 +from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
7 +from vllm.model_executor.layers.fused_moe.fused_moe import (cutlass_moe_fp8,
8 + fused_experts,
9 + fused_topk)
10 +from vllm.platforms import current_platform
11 +
12 +NUM_EXPERTS = [40, 64]
13 +TOP_KS = [6, 8]
14 +
15 +
16 +def run(a: torch.Tensor, a_scale: torch.Tensor, w1_q: torch.Tensor,
17 + w2_q: torch.Tensor, w1_scale: torch.Tensor, w2_scale: torch.Tensor,
18 + topk_weights: torch.Tensor, topk_ids: torch.Tensor,
19 + ab_strides1: torch.Tensor, c_strides1: torch.Tensor,
20 + ab_strides2: torch.Tensor, c_strides2: torch.Tensor):
21 + with set_current_vllm_config(
22 + VllmConfig(parallel_config=ParallelConfig(
23 + pipeline_parallel_size=1))):
24 + return cutlass_moe_fp8(a,
25 + w1_q,
26 + w2_q,
27 + w1_scale,
28 + w2_scale,
29 + topk_weights,
30 + topk_ids,
31 + ab_strides1,
32 + c_strides1,
33 + ab_strides2,
34 + c_strides2,
35 + a1_scale=a_scale)
36 +
37 +
38 +@pytest.mark.parametrize("m", [2, 64, 224])
39 +@pytest.mark.parametrize("n", [1024, 3072])
40 +@pytest.mark.parametrize("k", [1024, 1536])
41 +@pytest.mark.parametrize("e", NUM_EXPERTS)
42 +@pytest.mark.parametrize("topk", TOP_KS)
43 +@pytest.mark.parametrize("per_act_token", [True, False])
44 +@pytest.mark.parametrize("per_out_ch", [True, False])
45 +@pytest.mark.skipif(
46 + (lambda x: x is None or not ops.cutlass_group_gemm_supported(x.to_int()))(
47 + current_platform.get_device_capability()),
48 + reason="Grouped gemm is not supported on this GPU type.")
49 +def test_cutlass_moe_no_graph(
50 + m: int,
51 + n: int,
52 + k: int,
53 + e: int,
54 + topk: int,
55 + per_act_token: bool,
56 + per_out_ch: bool,
57 +):
58 + current_platform.seed_everything(7)
59 + with set_current_vllm_config(
60 + VllmConfig(parallel_config=ParallelConfig(
61 + pipeline_parallel_size=1))):
62 +
63 + dtype = torch.half
64 +
65 + a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
66 + w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
67 + w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
68 +
69 + # Get the right scale for tests.
70 + _, a_scale1 = ops.scaled_fp8_quant(
71 + a, use_per_token_if_dynamic=per_act_token)
72 + a_q, _ = ops.scaled_fp8_quant(a,
73 + a_scale1,
74 + use_per_token_if_dynamic=per_act_token)
75 +
76 + a_d = a_q.float().mul(a_scale1).to(dtype)
77 +
78 + n_b_scales = 2 * n if per_out_ch else 1
79 + k_b_scales = k if per_out_ch else 1
80 +
81 + w1_q = torch.empty((e, 2 * n, k),
82 + device="cuda",
83 + dtype=torch.float8_e4m3fn)
84 + w2_q = torch.empty((e, k, n), device="cuda", dtype=torch.float8_e4m3fn)
85 + w1_scale = torch.empty((e, n_b_scales, 1),
86 + device="cuda",
87 + dtype=torch.float32)
88 + w2_scale = torch.empty((e, k_b_scales, 1),
89 + device="cuda",
90 + dtype=torch.float32)
91 +
92 + ab_strides1 = torch.full((e, ), k, device="cuda", dtype=torch.int64)
93 + c_strides1 = torch.full((e, ), 2 * n, device="cuda", dtype=torch.int64)
94 + ab_strides2 = torch.full((e, ), n, device="cuda", dtype=torch.int64)
95 + c_strides2 = torch.full((e, ), k, device="cuda", dtype=torch.int64)
96 +
97 + for expert in range(e):
98 + w1_q[expert], w1_scale[expert] = ops.scaled_fp8_quant(
99 + w1[expert], use_per_token_if_dynamic=per_out_ch)
100 + w2_q[expert], w2_scale[expert] = ops.scaled_fp8_quant(
101 + w2[expert], use_per_token_if_dynamic=per_out_ch)
102 + w1_q = w1_q.transpose(1, 2)
103 + w2_q = w2_q.transpose(1, 2)
104 +
105 + ab_strides1 = torch.full((e, ), k, device="cuda", dtype=torch.int64)
106 + c_strides1 = torch.full((e, ), 2 * n, device="cuda", dtype=torch.int64)
107 + ab_strides2 = torch.full((e, ), n, device="cuda", dtype=torch.int64)
108 + c_strides2 = torch.full((e, ), k, device="cuda", dtype=torch.int64)
109 +
110 + w1_d = torch.empty_like(w1)
111 + w2_d = torch.empty_like(w2)
112 + for expert in range(e):
113 + w1_d[expert] = (w1_q[expert].t().float() * w1_scale[expert]).half()
114 + w2_d[expert] = (w2_q[expert].t().float() * w2_scale[expert]).half()
115 +
116 + score = torch.randn((m, e), device="cuda", dtype=dtype)
117 + topk_weights, topk_ids = fused_topk(a, score, topk, renormalize=False)
118 +
119 + triton_output = fused_experts(a_d, w1_d, w2_d, topk_weights, topk_ids)
120 +
121 + cutlass_output = cutlass_moe_fp8(a,
122 + w1_q,
123 + w2_q,
124 + w1_scale,
125 + w2_scale,
126 + topk_weights,
127 + topk_ids,
128 + ab_strides1,
129 + c_strides1,
130 + ab_strides2,
131 + c_strides2,
132 + a1_scale=a_scale1)
133 +
134 + print(triton_output)
135 + print(cutlass_output)
136 + print("*")
137 +
138 + torch.testing.assert_close(triton_output,
139 + cutlass_output,
140 + atol=5e-2,
141 + rtol=1e-2)
142 +
143 +
144 +@pytest.mark.parametrize("m", [2, 64, 224])
145 +@pytest.mark.parametrize("n", [1024, 3072])
146 +@pytest.mark.parametrize("k", [1024, 1536])
147 +@pytest.mark.parametrize("e", NUM_EXPERTS)
148 +@pytest.mark.parametrize("topk", TOP_KS)
149 +@pytest.mark.parametrize("per_act_token", [True, False])
150 +@pytest.mark.parametrize("per_out_ch", [True, False])
151 +@pytest.mark.skipif(
152 + (lambda x: x is None or not ops.cutlass_group_gemm_supported(x.to_int()))(
153 + current_platform.get_device_capability()),
154 + reason="Grouped gemm is not supported on this GPU type.")
155 +def test_cutlass_moe_cuda_graph(
156 + m: int,
157 + n: int,
158 + k: int,
159 + e: int,
160 + topk: int,
161 + per_act_token: bool,
162 + per_out_ch: bool,
163 +):
164 + current_platform.seed_everything(7)
165 + with set_current_vllm_config(
166 + VllmConfig(parallel_config=ParallelConfig(
167 + pipeline_parallel_size=1))):
168 +
169 + dtype = torch.half
170 +
171 + a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
172 + w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
173 + w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
174 +
175 + # Get the right scale for tests.
176 + _, a_scale1 = ops.scaled_fp8_quant(
177 + a, use_per_token_if_dynamic=per_act_token)
178 + a_q, _ = ops.scaled_fp8_quant(a,
179 + a_scale1,
180 + use_per_token_if_dynamic=per_act_token)
181 +
182 + a_d = a_q.float().mul(a_scale1).to(dtype)
183 +
184 + n_b_scales = 2 * n if per_out_ch else 1
185 + k_b_scales = k if per_out_ch else 1
186 +
187 + w1_q = torch.empty((e, 2 * n, k),
188 + device="cuda",
189 + dtype=torch.float8_e4m3fn)
190 + w2_q = torch.empty((e, k, n), device="cuda", dtype=torch.float8_e4m3fn)
191 + w1_scale = torch.empty((e, n_b_scales, 1),
192 + device="cuda",
193 + dtype=torch.float32)
194 + w2_scale = torch.empty((e, k_b_scales, 1),
195 + device="cuda",
196 + dtype=torch.float32)
197 +
198 + ab_strides1 = torch.full((e, ), k, device="cuda", dtype=torch.int64)
199 + c_strides1 = torch.full((e, ), 2 * n, device="cuda", dtype=torch.int64)
200 + ab_strides2 = torch.full((e, ), n, device="cuda", dtype=torch.int64)
201 + c_strides2 = torch.full((e, ), k, device="cuda", dtype=torch.int64)
202 +
203 + for expert in range(e):
204 + w1_q[expert], w1_scale[expert] = ops.scaled_fp8_quant(
205 + w1[expert], use_per_token_if_dynamic=per_out_ch)
206 + w2_q[expert], w2_scale[expert] = ops.scaled_fp8_quant(
207 + w2[expert], use_per_token_if_dynamic=per_out_ch)
208 + w1_q = w1_q.transpose(1, 2)
209 + w2_q = w2_q.transpose(1, 2)
210 +
211 + ab_strides1 = torch.full((e, ), k, device="cuda", dtype=torch.int64)
212 + c_strides1 = torch.full((e, ), 2 * n, device="cuda", dtype=torch.int64)
213 + ab_strides2 = torch.full((e, ), n, device="cuda", dtype=torch.int64)
214 + c_strides2 = torch.full((e, ), k, device="cuda", dtype=torch.int64)
215 +
216 + w1_d = torch.empty_like(w1)
217 + w2_d = torch.empty_like(w2)
218 + for expert in range(e):
219 + w1_d[expert] = (w1_q[expert].t().float() * w1_scale[expert]).half()
220 + w2_d[expert] = (w2_q[expert].t().float() * w2_scale[expert]).half()
221 +
222 + score = torch.randn((m, e), device="cuda", dtype=dtype)
223 + topk_weights, topk_ids = fused_topk(a, score, topk, renormalize=False)
224 +
225 + triton_output = fused_experts(a_d, w1_d, w2_d, topk_weights, topk_ids)
226 +
227 + stream = torch.cuda.Stream()
228 + graph = torch.cuda.CUDAGraph()
229 + with torch.cuda.graph(graph, stream=stream):
230 + cutlass_output = run(a, a_scale1, w1_q, w2_q, w1_scale, w2_scale,
231 + topk_weights, topk_ids, ab_strides1,
232 + c_strides1, ab_strides2, c_strides2)
233 + torch.cuda.synchronize()
234 + graph.replay()
235 + torch.cuda.synchronize()
236 +
237 + print(triton_output)
238 + print(cutlass_output)
239 + print("*")
240 +
241 + torch.testing.assert_close(triton_output,
242 + cutlass_output,
243 + atol=9e-2,
244 + rtol=1e-2)
vllm/_custom_ops.py
+53
@@ -587,6 +587,9 @@ def cutlass_sparse_scaled_mm_supported(cuda_device_capability: int) -> bool:
587 cuda_device_capability)
588
589
590 +def cutlass_group_gemm_supported(cuda_device_capability: int) -> bool:
591 + return torch.ops._C.cutlass_group_gemm_supported(cuda_device_capability)
592 +
593 def cutlass_sparse_compress(a: torch.Tensor) \
594 -> tuple[torch.Tensor, torch.Tensor]:
595 """
@@ -677,6 +680,56 @@ def cutlass_scaled_sparse_mm(
680 return out
681
682
683 +def get_cutlass_moe_mm_data(
684 + topk_ids: torch.Tensor, expert_offsets: torch.Tensor,
685 + problem_sizes1: torch.Tensor, problem_sizes2: torch.Tensor,
686 + input_permutation: torch.Tensor, output_permutation: torch.Tensor,
687 + num_experts: int, n: int, k: int):
688 + """
689 + Prepare data necessary to perform CUTLASS grouped matrix multiplications
690 + used in CUTLASS-based fused MoE.
691 +
692 + The function takes in topk_ids (token-expert mapping) and uses it to
693 + compute:
694 + - expert_offsets: Indices that mark at which token index each expert begins
695 + its computation after the input is sorted with
696 + input_permutation. The number of tokens computed with
697 + expert E is expert_offsets[E + 1] - expert_offsets[E]
698 + - problem_sizes1, problem_sizes2: MxNxK sizes of each expert's
699 + multiplication in two grouped MMs used in
700 + the fused MoE operation.
701 + - input_permutation: Permutation that must be used to shuffle the input
702 + before executing the MMs.
703 + - output_permutation: Permutation that must be used to shuffle the output
704 + after executing the MMs.
705 + """
706 + torch.ops._C.get_cutlass_moe_mm_data(topk_ids, expert_offsets,
707 + problem_sizes1, problem_sizes2,
708 + input_permutation, output_permutation,
709 + num_experts, n, k)
710 +
711 +
712 +def cutlass_moe_mm(out_tensors: torch.Tensor, a_tensors: torch.Tensor,
713 + b_tensors: torch.Tensor, a_scales: torch.Tensor,
714 + b_scales: torch.Tensor, expert_offsets: torch.Tensor,
715 + problem_sizes: torch.Tensor, a_strides: torch.Tensor,
716 + b_strides: torch.Tensor, c_strides: torch.Tensor):
717 + """
718 + A single grouped matrix multiplication used in CUTLASS-based fused MoE.
719 + The function executes fp8-quantized OUT = AB matrix multiplication.
720 +
721 + - expert_offsets: Indices that mark at which token index each expert begins
722 + its computation. The number of tokens computed with
723 + expert E is expert_offsets[E + 1] - expert_offsets[E]
724 + - problem_sizes: MxNxK sizes of each expert's multiplication in two grouped
725 + MMs used in the fused MoE operation.
726 + - a/b/c_strides: The data strides passed to grouped matrix multiplication.
727 + """
728 + torch.ops._C.cutlass_moe_mm(out_tensors, a_tensors, b_tensors, a_scales,
729 + b_scales, expert_offsets, problem_sizes,
730 + a_strides, b_strides, c_strides)
731 +
732 +
733 # aqlm
734 def aqlm_gemm(input: torch.Tensor, codes: torch.Tensor,
735 codebooks: torch.Tensor, scales: torch.Tensor,
vllm/model_executor/layers/fused_moe/__init__.py
+3 -2
@@ -36,8 +36,8 @@ if HAS_TRITON:
36 import vllm.model_executor.layers.fused_moe.fused_marlin_moe # noqa
37 import vllm.model_executor.layers.fused_moe.fused_moe # noqa
38 from vllm.model_executor.layers.fused_moe.fused_moe import (
39 - fused_experts, fused_moe, fused_topk, get_config_file_name,
40 - grouped_topk)
39 + cutlass_moe_fp8, fused_experts, fused_moe, fused_topk,
40 + get_config_file_name, grouped_topk)
41
42 __all__ += [
43 "fused_moe",
@@ -45,4 +45,5 @@ if HAS_TRITON:
45 "fused_experts",
46 "get_config_file_name",
47 "grouped_topk",
48 + "cutlass_moe_fp8",
49 ]
vllm/model_executor/layers/fused_moe/fused_moe.py
+137
@@ -1623,3 +1623,140 @@ def fused_moe(
1623 a1_scale=a1_scale,
1624 a2_scale=a2_scale,
1625 block_shape=block_shape)
1626 +
1627 +
1628 +#TODO make the grouped gemm kernel consistent with scaled gemm kernel
1629 +def cutlass_moe_fp8(
1630 + a: torch.Tensor,
1631 + w1_q: torch.Tensor,
1632 + w2_q: torch.Tensor,
1633 + w1_scale: torch.Tensor,
1634 + w2_scale: torch.Tensor,
1635 + topk_weights: torch.Tensor,
1636 + topk_ids: torch.Tensor,
1637 + ab_strides1: torch.Tensor,
1638 + c_strides1: torch.Tensor,
1639 + ab_strides2: torch.Tensor,
1640 + c_strides2: torch.Tensor,
1641 + a1_scale: Optional[torch.Tensor] = None,
1642 + a2_scale: Optional[torch.Tensor] = None,
1643 + out_dtype: torch.dtype = torch.half,
1644 +) -> torch.Tensor:
1645 + """
1646 + This function computes a a8w8-quantized Mixture of Experts (MoE) layer
1647 + using two sets of quantized weights, w1_q and w2_q, and top-k gating
1648 + mechanism. The matrix multiplications are implemented with CUTLASS
1649 + grouped gemm.
1650 +
1651 + Parameters:
1652 + - a (torch.Tensor): The input tensor to the MoE layer.
1653 + Shape: [M, K]
1654 + - w1_q (torch.Tensor): The first set of fp8-quantized expert weights.
1655 + Shape: [num_experts, K, 2N] (the weights are passed transposed)
1656 + - w2_q (torch.Tensor): The second set of fp8-quantized expert weights.
1657 + Shape: [num_experts, N, K] (the weights are passed transposed)
1658 + - w1_scale (torch.Tensor): The fp32 scale to dequantize w1_q.
1659 + Shape: [num_experts] or [num_experts, 2N]
1660 + - w2_scale (torch.Tensor): The fp32 scale to dequantize w2_q.
1661 + Shape: [num_experts] or [num_experts, K]
1662 + - gating_output (torch.Tensor): The output of the gating operation
1663 + (before softmax).
1664 + - topk_weights (torch.Tensor): The weights of each token->expert mapping.
1665 + - ab_strides1 (torch.Tensor): The input and weights strides of the first
1666 + grouped gemm.
1667 + - c_strides1 (torch.Tensor): The output strides of the first grouped gemm.
1668 + - ab_strides2 (torch.Tensor): The input and weights strides of the second
1669 + grouped gemm.
1670 + - c_strides2 (torch.Tensor): The output strides of the second grouped gemm.
1671 + - a1_scale (Optional[torch.Tensor]): The optional fp32 scale to quantize a.
1672 + Shape: scalar or [M]
1673 + - a2_scale (Optional[torch.Tensor]): The optional fp32 scale to
1674 + quantize the intermediate result between the gemms.
1675 + Shape: scalar or [M]
1676 + - out_dtype (torch.Tensor): The output tensor type.
1677 +
1678 + Returns:
1679 + - torch.Tensor: The fp16 output tensor after applying the MoE layer.
1680 + """
1681 +
1682 + assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
1683 + assert w1_q.dtype == torch.float8_e4m3fn
1684 + assert w2_q.dtype == torch.float8_e4m3fn
1685 + assert a.shape[1] == w1_q.shape[1], "Hidden size mismatch w1"
1686 + assert w1_q.shape[2] == w2_q.shape[1] * 2, "Hidden size mismatch w2"
1687 + assert w1_q.shape[0] == w2_q.shape[0], "Expert number mismatch"
1688 + assert a1_scale is None or a1_scale.dim(
1689 + ) == 0 or a1_scale.shape[0] == 1 or a1_scale.shape[0] == a.shape[
1690 + 0], "Input scale shape mismatch"
1691 + assert w1_scale.dim() == 1 or w1_scale.shape[1] == 1 or w1_scale.shape[
1692 + 1] == w1_q.shape[2], "W1 scale shape mismatch"
1693 + assert w2_scale.dim() == 1 or w2_scale.shape[1] == 1 or w2_scale.shape[
1694 + 1] == w2_q.shape[2], "W2 scale shape mismatch"
1695 + assert w1_q.shape[0] == w2_q.shape[0], "Weights expert number mismatch"
1696 + assert w1_q.shape[0] == w1_scale.shape[
1697 + 0], "w1 scales expert number mismatch"
1698 + assert w1_q.shape[0] == w2_scale.shape[
1699 + 0], "w2 scales expert number mismatch"
1700 + assert a2_scale is None or a1_scale is None or a2_scale.shape == a1_scale.shape, "Intermediate scale shape mismatch" # noqa: E501
1701 + assert ab_strides1.shape[0] == w1_q.shape[
1702 + 0], "AB Strides 1 expert number mismatch"
1703 + assert c_strides1.shape[0] == w1_q.shape[
1704 + 0], "C Strides 1 expert number mismatch"
1705 + assert ab_strides2.shape[0] == w2_q.shape[
1706 + 0], "AB Strides 2 expert number mismatch"
1707 + assert c_strides2.shape[0] == w2_q.shape[
1708 + 0], "C Strides 2 expert number mismatch"
1709 + assert out_dtype in [torch.half, torch.bfloat16], "Invalid output dtype"
1710 +
1711 + num_experts = w1_q.size(0)
1712 + m = a.size(0)
1713 + k = w1_q.size(1)
1714 + n = w2_q.size(1)
1715 +
1716 + topk = topk_ids.size(1)
1717 + per_act_token = a1_scale.numel() != 1 if a1_scale is not None else (
1718 + a2_scale.numel() != 1 if a2_scale is not None else False)
1719 +
1720 + a_q, a1_scale = ops.scaled_fp8_quant(
1721 + a, a1_scale, use_per_token_if_dynamic=per_act_token)
1722 + device = a_q.device
1723 +
1724 + expert_offsets = torch.empty((num_experts + 1),
1725 + dtype=torch.int32,
1726 + device=device)
1727 + problem_sizes1 = torch.empty((num_experts, 3),
1728 + dtype=torch.int32,
1729 + device=device)
1730 + problem_sizes2 = torch.empty((num_experts, 3),
1731 + dtype=torch.int32,
1732 + device=device)
1733 +
1734 + a_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device)
1735 + c_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device)
1736 +
1737 + ops.get_cutlass_moe_mm_data(topk_ids, expert_offsets, problem_sizes1,
1738 + problem_sizes2, a_map, c_map, num_experts, n,
1739 + k)
1740 +
1741 + rep_a_q = a_q.view(dtype=torch.uint8)[a_map].view(dtype=a_q.dtype)
1742 + rep_a1_scales = a1_scale[a_map] if per_act_token else a1_scale
1743 +
1744 + c1 = torch.empty((m * topk, n * 2), device=device, dtype=out_dtype)
1745 + c2 = torch.empty((m * topk, k), device=device, dtype=out_dtype)
1746 +
1747 + ops.cutlass_moe_mm(c1, rep_a_q, w1_q, rep_a1_scales, w1_scale,
1748 + expert_offsets[:-1], problem_sizes1, ab_strides1,
1749 + ab_strides1, c_strides1)
1750 +
1751 + intermediate = torch.empty((m * topk, n), device=device, dtype=out_dtype)
1752 + torch.ops._C.silu_and_mul(intermediate, c1)
1753 +
1754 + intemediate_q, a2_scale = ops.scaled_fp8_quant(
1755 + intermediate, a2_scale, use_per_token_if_dynamic=per_act_token)
1756 +
1757 + ops.cutlass_moe_mm(c2, intemediate_q, w2_q, a2_scale, w2_scale,
1758 + expert_offsets[:-1], problem_sizes2, ab_strides2,
1759 + ab_strides2, c_strides2)
1760 +
1761 + return (c2[c_map].view(m, topk, k) *
1762 + topk_weights.view(m, topk, 1).to(out_dtype)).sum(dim=1)
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors.py
+23 -8
@@ -96,7 +96,8 @@ class CompressedTensorsConfig(QuantizationConfig):
96 if isinstance(layer, Attention):
97 return CompressedTensorsKVCacheMethod(self)
98 if isinstance(layer, FusedMoE):
99 - return CompressedTensorsMoEMethod.get_moe_method(self)
99 + return CompressedTensorsMoEMethod.get_moe_method(
100 + self, layer.activation, layer.expert_map)
101 return None
102
103 @classmethod
@@ -191,17 +192,26 @@ class CompressedTensorsConfig(QuantizationConfig):
192
193 def _check_scheme_supported(self,
194 min_capability: int,
194 - error: bool = True) -> bool:
195 + error: bool = True,
196 + match_exact: bool = False) -> bool:
197 capability_tuple = current_platform.get_device_capability()
198
199 if capability_tuple is not None:
200 capability = capability_tuple.to_int()
199 - supported = capability >= min_capability
200 - if error and not supported:
201 - raise RuntimeError(
202 - "Quantization scheme is not supported for ",
203 - f"the current GPU. Min capability: {min_capability}. ",
204 - f"Current capability: {capability}.")
201 + if match_exact:
202 + supported = capability == min_capability
203 + if error and not supported:
204 + raise RuntimeError(
205 + "Quantization scheme is not supported for ",
206 + "the current GPU. Required capability: ",
207 + f"{min_capability}. Current capability: {capability}.")
208 + else:
209 + supported = capability >= min_capability
210 + if error and not supported:
211 + raise RuntimeError(
212 + "Quantization scheme is not supported for ",
213 + f"the current GPU. Min capability: {min_capability}. ",
214 + f"Current capability: {capability}.")
215 return supported
216 else:
217 return False
@@ -262,6 +272,11 @@ class CompressedTensorsConfig(QuantizationConfig):
272 input_quant.strategy == QuantizationStrategy.TENSOR)
273 return is_symmetric_activation and is_per_tensor_activation
274
275 + def _is_fp8_w8a8_sm90(self, weight_quant: BaseModel,
276 + input_quant: BaseModel) -> bool:
277 + return (self._check_scheme_supported(90, error=False, match_exact=True)
278 + and self._is_fp8_w8a8(weight_quant, input_quant))
279 +
280 def _is_fp8_w8a16(self, weight_quant: BaseModel,
281 input_quant: BaseModel) -> bool:
282 # Confirm weights quantized.
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+201 -1
@@ -31,6 +31,7 @@ class GPTQMarlinState(Enum):
31
32 __all__ = [
33 "CompressedTensorsMoEMethod", "CompressedTensorsW8A8Fp8MoEMethod",
34 + "CompressedTensorsW8A8Fp8MoECutlassMethod",
35 "CompressedTensorsWNA16MoEMethod"
36 ]
37
@@ -39,7 +40,9 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
40
41 @staticmethod
42 def get_moe_method(
42 - quant_config: "CompressedTensorsConfig" # type: ignore # noqa E501
43 + quant_config: "CompressedTensorsConfig", # type: ignore # noqa E501
44 + activation: str,
45 + expert_map: Optional[torch.Tensor],
46 ) -> "CompressedTensorsMoEMethod":
47 # TODO: @dsikka: refactor this to use schemes as other kernels
48 # are supported + check if the layer is being ignored.
@@ -49,6 +52,9 @@ class CompressedTensorsMoEMethod(FusedMoEMethodBase):
52
53 if quant_config._is_wNa16_group_channel(weight_quant, input_quant):
54 return CompressedTensorsWNA16MoEMethod(quant_config)
55 + elif (quant_config._is_fp8_w8a8_sm90(weight_quant, input_quant)
56 + and activation == "silu" and expert_map is None):
57 + return CompressedTensorsW8A8Fp8MoECutlassMethod(quant_config)
58 elif quant_config._is_fp8_w8a8(weight_quant, input_quant):
59 return CompressedTensorsW8A8Fp8MoEMethod(quant_config)
60 else:
@@ -250,6 +256,200 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
256 a2_scale=layer.w2_input_scale)
257
258
259 +class CompressedTensorsW8A8Fp8MoECutlassMethod(CompressedTensorsMoEMethod):
260 +
261 + def __init__(
262 + self,
263 + quant_config: "CompressedTensorsConfig" # type: ignore # noqa E501
264 + ):
265 + self.quant_config = quant_config
266 + self.weight_quant = self.quant_config.target_scheme_map["Linear"].get(
267 + "weights")
268 + self.input_quant = self.quant_config.target_scheme_map["Linear"].get(
269 + "input_activations")
270 +
271 + if not (self.weight_quant.strategy == QuantizationStrategy.TENSOR
272 + and self.input_quant.strategy == QuantizationStrategy.TENSOR):
273 + raise ValueError(
274 + "For FP8 Fused MoE layers, only per-tensor scales "
275 + "for weights and activations are supported. Found "
276 + f"{self.weight_quant}, {self.input_quant}")
277 +
278 + self.static_input_scales = not self.input_quant.dynamic
279 +
280 + def create_weights(self, layer: torch.nn.Module, num_experts: int,
281 + hidden_size: int, intermediate_size_per_partition: int,
282 + params_dtype: torch.dtype, **extra_weight_attrs):
283 +
284 + params_dtype = torch.float8_e4m3fn
285 +
286 + # WEIGHTS
287 + w13_weight = torch.nn.Parameter(torch.empty(
288 + num_experts,
289 + 2 * intermediate_size_per_partition,
290 + hidden_size,
291 + dtype=params_dtype),
292 + requires_grad=False)
293 + layer.register_parameter("w13_weight", w13_weight)
294 + set_weight_attrs(w13_weight, extra_weight_attrs)
295 +
296 + w2_weight = torch.nn.Parameter(torch.empty(
297 + num_experts,
298 + hidden_size,
299 + intermediate_size_per_partition,
300 + dtype=params_dtype),
301 + requires_grad=False)
302 + layer.register_parameter("w2_weight", w2_weight)
303 + set_weight_attrs(w2_weight, extra_weight_attrs)
304 +
305 + # WEIGHT_SCALES
306 + # Allocate 2 scales for w1 and w3 respectively.
307 + # They will be combined to a single scale after weight loading.
308 + w13_weight_scale = torch.nn.Parameter(torch.ones(num_experts,
309 + 2,
310 + dtype=torch.float32),
311 + requires_grad=False)
312 + layer.register_parameter("w13_weight_scale", w13_weight_scale)
313 +
314 + w2_weight_scale = torch.nn.Parameter(torch.ones(num_experts,
315 + dtype=torch.float32),
316 + requires_grad=False)
317 + layer.register_parameter("w2_weight_scale", w2_weight_scale)
318 + # Add the quantization method used (per tensor/grouped/channel)
319 + # to ensure the weight scales are loaded in properly
320 + extra_weight_attrs.update(
321 + {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value})
322 + set_weight_attrs(w13_weight_scale, extra_weight_attrs)
323 + set_weight_attrs(w2_weight_scale, extra_weight_attrs)
324 +
325 + # INPUT_SCALES
326 + if self.static_input_scales:
327 + w13_input_scale = torch.nn.Parameter(torch.ones(
328 + num_experts, dtype=torch.float32),
329 + requires_grad=False)
330 + layer.register_parameter("w13_input_scale", w13_input_scale)
331 + set_weight_attrs(w13_input_scale, extra_weight_attrs)
332 +
333 + w2_input_scale = torch.nn.Parameter(torch.ones(
334 + num_experts, dtype=torch.float32),
335 + requires_grad=False)
336 + layer.register_parameter("w2_input_scale", w2_input_scale)
337 + set_weight_attrs(w2_input_scale, extra_weight_attrs)
338 + else:
339 + layer.w13_input_scale = None
340 + layer.w2_input_scale = None
341 +
342 + device = w13_weight.device
343 + # TODO strides can be shared across multiple layers
344 + self.ab_strides1 = torch.full((num_experts, ),
345 + hidden_size,
346 + device=device,
347 + dtype=torch.int64)
348 + self.c_strides1 = torch.full((num_experts, ),
349 + 2 * intermediate_size_per_partition,
350 + device=device,
351 + dtype=torch.int64)
352 + self.ab_strides2 = torch.full((num_experts, ),
353 + intermediate_size_per_partition,
354 + device=device,
355 + dtype=torch.int64)
356 + self.c_strides2 = torch.full((num_experts, ),
357 + hidden_size,
358 + device=device,
359 + dtype=torch.int64)
360 +
361 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
362 + # Fp8 moe kernels require a single activation scale.
363 + # We take the max of all the scales in case they differ.
364 + if self.static_input_scales:
365 + if (layer.w13_input_scale is None or layer.w2_input_scale is None):
366 + raise ValueError(
367 + "QuantConfig has static quantization, but found "
368 + "activation scales are None.")
369 + if (not all_close_1d(layer.w13_input_scale)
370 + or not all_close_1d(layer.w2_input_scale)):
371 + logger.warning_once(
372 + "Found input_scales that are not equal for "
373 + "fp8 MoE layer. Using the maximum across experts "
374 + "for each layer.")
375 + layer.w13_input_scale = torch.nn.Parameter(
376 + layer.w13_input_scale.max(), requires_grad=False)
377 + layer.w2_input_scale = torch.nn.Parameter(
378 + layer.w2_input_scale.max(), requires_grad=False)
379 +
380 + # Fp8 moe kernel needs single weight scale for w13 per expert.
381 + # We take the max then dequant and requant each expert.
382 + assert layer.w13_weight_scale is not None
383 + shard_size = layer.intermediate_size_per_partition
384 + max_w13_scales = layer.w13_weight_scale.max(dim=1).values
385 + for expert_id in range(layer.local_num_experts):
386 + start = 0
387 + for shard_id in range(2):
388 + dq_weight = per_tensor_dequantize(
389 + layer.w13_weight[expert_id][start:start + shard_size, :],
390 + layer.w13_weight_scale[expert_id][shard_id])
391 + layer.w13_weight[expert_id][
392 + start:start + shard_size, :], _ = ops.scaled_fp8_quant(
393 + dq_weight, max_w13_scales[expert_id])
394 + start += shard_size
395 +
396 + layer.w13_weight_scale = torch.nn.Parameter(max_w13_scales,
397 + requires_grad=False)
398 +
399 + def apply(
400 + self,
401 + layer: torch.nn.Module,
402 + x: torch.Tensor,
403 + router_logits: torch.Tensor,
404 + top_k: int,
405 + renormalize: bool,
406 + use_grouped_topk: bool = False,
407 + topk_group: Optional[int] = None,
408 + num_expert_group: Optional[int] = None,
409 + global_num_experts: int = -1,
410 + expert_map: Optional[torch.Tensor] = None,
411 + custom_routing_function: Optional[Callable] = None,
412 + scoring_func: str = "softmax",
413 + e_score_correction_bias: Optional[torch.Tensor] = None,
414 + activation: str = "silu",
415 + ) -> torch.Tensor:
416 +
417 + assert activation == "silu"
418 + assert global_num_experts == layer.w13_weight.shape[0]
419 + assert expert_map is None
420 +
421 + topk_weights, topk_ids = FusedMoE.select_experts(
422 + hidden_states=x,
423 + router_logits=router_logits,
424 + use_grouped_topk=use_grouped_topk,
425 + top_k=top_k,
426 + renormalize=renormalize,
427 + topk_group=topk_group,
428 + num_expert_group=num_expert_group,
429 + custom_routing_function=custom_routing_function,
430 + scoring_func=scoring_func,
431 + e_score_correction_bias=e_score_correction_bias)
432 +
433 + from vllm.model_executor.layers.fused_moe import cutlass_moe_fp8
434 +
435 + return cutlass_moe_fp8(
436 + x,
437 + layer.w13_weight.transpose(1, 2),
438 + layer.w2_weight.transpose(1, 2),
439 + layer.w13_weight_scale,
440 + layer.w2_weight_scale,
441 + topk_weights,
442 + topk_ids,
443 + self.ab_strides1,
444 + self.c_strides1,
445 + self.ab_strides2,
446 + self.c_strides2,
447 + a1_scale=layer.w13_input_scale,
448 + a2_scale=layer.w2_input_scale,
449 + out_dtype=x.dtype,
450 + )
451 +
452 +
453 class CompressedTensorsWNA16MoEMethod(CompressedTensorsMoEMethod):
454
455 def __init__(
vllm/model_executor/layers/quantization/utils/w8a8_utils.py
+10
@@ -50,6 +50,16 @@ def cutlass_block_fp8_supported() -> bool:
50 return ops.cutlass_scaled_mm_supports_block_fp8(capability)
51
52
53 +def cutlass_group_gemm_supported() -> bool:
54 + if not current_platform.is_cuda():
55 + return False
56 +
57 + capability_tuple = current_platform.get_device_capability()
58 + capability = -1 if capability_tuple is None else capability_tuple.to_int()
59 +
60 + return ops.cutlass_group_gemm_supported(capability)
61 +
62 +
63 CUTLASS_FP8_SUPPORTED = cutlass_fp8_supported()
64 CUTLASS_BLOCK_FP8_SUPPORTED = cutlass_block_fp8_supported()
65
vllm/utils.py
+6 -3
@@ -1568,18 +1568,21 @@ class ClassRegistry(UserDict[Type[T], _V]):
1568 return any(cls in self.data for cls in key.mro())
1569
1570
1571 -def weak_ref_tensor(tensor: torch.Tensor) -> torch.Tensor:
1571 +def weak_ref_tensor(tensor: Any) -> Any:
1572 """
1573 Create a weak reference to a tensor.
1574 The new tensor will share the same data as the original tensor,
1575 but will not keep the original tensor alive.
1576 """
1577 - return torch.ops._C.weak_ref_tensor(tensor)
1577 + if isinstance(tensor, torch.Tensor):
1578 + return torch.ops._C.weak_ref_tensor(tensor)
1579 + else:
1580 + return tensor
1581
1582
1583 def weak_ref_tensors(
1584 tensors: Union[torch.Tensor, list[torch.Tensor], tuple[torch.Tensor]]
1582 -) -> Union[torch.Tensor, list[torch.Tensor], tuple[torch.Tensor]]:
1585 +) -> Union[torch.Tensor, list[Any], tuple[Any], Any]:
1586 """
1587 Convenience function to create weak references to tensors,
1588 for single tensor, list of tensors or tuple of tensors.