@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 0c0fdae84f
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, 11:00 UTC

[Hardware/NVIDIA/Kernel] Enable nvidia/DeepSeek-R1-FP4 Model (#16362)

Pavani Majety committed May 9, 2025 at 16:24 UTC 0c0fdae84f1da5e45518aafc7b32e8139055adae
16 files changed +1994 -112
CMakeLists.txt
+5 -2
@@ -288,6 +288,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
288 "csrc/quantization/cutlass_w8a8/scaled_mm_entry.cu"
289 "csrc/quantization/fp4/nvfp4_quant_entry.cu"
290 "csrc/quantization/fp4/nvfp4_scaled_mm_entry.cu"
291 + "csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu"
292 "csrc/sparse/cutlass/sparse_scaled_mm_entry.cu"
293 "csrc/cutlass_extensions/common.cpp"
294 "csrc/attention/mla/cutlass_mla_entry.cu")
@@ -495,7 +496,9 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
496 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER 12.8 AND FP4_ARCHS)
497 set(SRCS
498 "csrc/quantization/fp4/nvfp4_quant_kernels.cu"
498 - "csrc/quantization/fp4/nvfp4_scaled_mm_kernels.cu")
499 + "csrc/quantization/fp4/nvfp4_experts_quant.cu"
500 + "csrc/quantization/fp4/nvfp4_scaled_mm_kernels.cu"
501 + "csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu")
502 set_gencode_flags_for_srcs(
503 SRCS "${SRCS}"
504 CUDA_ARCHS "${FP4_ARCHS}")
@@ -533,7 +536,7 @@ if(VLLM_GPU_LANG STREQUAL "CUDA")
536 # The MoE kernel cutlass_moe_mm requires CUDA 12.3 or later (and only works
537 # on Hopper). get_cutlass_moe_mm_data should only be compiled if it's possible
538 # to compile MoE kernels that use its output.
536 - cuda_archs_loose_intersection(SCALED_MM_ARCHS "9.0a;" "${CUDA_ARCHS}")
539 + cuda_archs_loose_intersection(SCALED_MM_ARCHS "9.0a;10.0a" "${CUDA_ARCHS}")
540 if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.3 AND SCALED_MM_ARCHS)
541 set(SRCS "csrc/quantization/cutlass_w8a8/moe/grouped_mm_c3x.cu"
542 "csrc/quantization/cutlass_w8a8/moe/moe_data.cu")
benchmarks/kernels/benchmark_cutlass_fp4_moe.py new
+408
@@ -0,0 +1,408 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +"""
3 +Benchmark the performance of the cutlass_moe_fp4 kernel vs the triton_moe
4 +kernel. The cutlass_moe_fp4 kernel takes in fp4 quantized weights and 16-bit
5 +activations. The triton_moe kernel takes in fp8 weights(tensor scaled to fp8)
6 +and 16-bit activations.
7 +"""
8 +import nvtx
9 +import torch
10 +import torch.utils.benchmark as benchmark
11 +
12 +from vllm import _custom_ops as ops
13 +from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
14 +from vllm.model_executor.layers.fused_moe.cutlass_moe import cutlass_moe_fp4
15 +from vllm.model_executor.layers.fused_moe.fused_moe import (fused_experts,
16 + fused_topk)
17 +from vllm.scalar_type import scalar_types
18 +from vllm.utils import FlexibleArgumentParser
19 +
20 +WEIGHT_SHAPES_MOE = {
21 + "nvidia/DeepSeek-R1-FP4": [
22 + [256, 8, 2048, 7168],
23 + ],
24 +}
25 +
26 +DEFAULT_MODELS = [
27 + "nvidia/DeepSeek-R1-FP4",
28 +]
29 +
30 +DEFAULT_BATCH_SIZES = [4, 8, 16, 32, 64, 128, 256, 512, 1024, 2048]
31 +DEFAULT_TP_SIZES = [1]
32 +
33 +PER_ACT_TOKEN_OPTS = [False]
34 +PER_OUT_CH_OPTS = [False]
35 +FLOAT4_E2M1_MAX = scalar_types.float4_e2m1f.max()
36 +FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
37 +
38 +
39 +def to_fp8(tensor: torch.Tensor):
40 + finfo = torch.finfo(torch.float8_e4m3fn)
41 + return torch.round(tensor.clamp(
42 + min=finfo.min, max=finfo.max)).to(dtype=torch.float8_e4m3fn)
43 +
44 +
45 +def bench_run(results: list[benchmark.Measurement], model: str,
46 + num_experts: int, topk: int, per_act_token: bool,
47 + per_out_ch: bool, mkn: tuple[int, int, int]):
48 + label = "NVFP4 Blockscaled CUTLASS MOE vs FP8 Tensor Scaled Triton"
49 +
50 + sub_label = (
51 + "{}, num_experts={}, topk={}, per_act_token={} per_out_ch={}, "
52 + "MKN=({})".format(model, num_experts, topk, per_act_token, per_out_ch,
53 + mkn))
54 +
55 + print(f"Testing: {sub_label}")
56 +
57 + (m, k, n) = mkn
58 +
59 + dtype = torch.half
60 + device = "cuda"
61 + a = torch.randn((m, k), device=device, dtype=dtype) / 10
62 + w1 = torch.randn((num_experts, 2 * n, k), device=device, dtype=dtype) / 10
63 + w2 = torch.randn((num_experts, k, n), device=device, dtype=dtype) / 10
64 +
65 + _, a_fp8_scale = ops.scaled_fp8_quant(a)
66 +
67 + w1_fp8q = torch.empty((num_experts, 2 * n, k),
68 + device=device,
69 + dtype=torch.float8_e4m3fn)
70 + w2_fp8q = torch.empty((num_experts, k, n),
71 + device=device,
72 + dtype=torch.float8_e4m3fn)
73 + w1_fp8scale = torch.empty((num_experts, 1, 1),
74 + device=device,
75 + dtype=torch.float32)
76 + w2_fp8scale = torch.empty((num_experts, 1, 1),
77 + device=device,
78 + dtype=torch.float32)
79 +
80 + for expert in range(num_experts):
81 + w1_fp8q[expert], w1_fp8scale[expert] = ops.scaled_fp8_quant(w1[expert])
82 + w2_fp8q[expert], w2_fp8scale[expert] = ops.scaled_fp8_quant(w2[expert])
83 +
84 + w1_fp8q_notransp = w1_fp8q.clone()
85 + w2_fp8q_notransp = w2_fp8q.clone()
86 + w1_fp8q = w1_fp8q.transpose(1, 2)
87 + w2_fp8q = w2_fp8q.transpose(1, 2)
88 +
89 + score = torch.randn((m, num_experts), device=device, dtype=dtype)
90 +
91 + topk_weights, topk_ids = fused_topk(a, score, topk, renormalize=False)
92 +
93 + quant_blocksize = 16
94 + w1_blockscale = torch.empty((num_experts, 2 * n, k // quant_blocksize),
95 + device=device,
96 + dtype=torch.float8_e4m3fn)
97 + w2_blockscale = torch.empty((num_experts, k, n // quant_blocksize),
98 + device=device,
99 + dtype=torch.float8_e4m3fn)
100 +
101 + # n_b_scales = 2 * n if per_out_ch else 1
102 + # k_b_scales = k if per_out_ch else 1
103 + w1_fp4 = torch.empty((num_experts, 2 * n, k // 2),
104 + device=device,
105 + dtype=torch.uint8)
106 + w2_fp4 = torch.empty((num_experts, k, n // 2),
107 + device=device,
108 + dtype=torch.uint8)
109 +
110 + w1_gs = torch.empty((num_experts, ), device=device, dtype=torch.float32)
111 + w2_gs = torch.empty((num_experts, ), device=device, dtype=torch.float32)
112 + a1_gs = torch.ones((num_experts, ), device=device, dtype=torch.float32)
113 + a2_gs = torch.ones((num_experts, ), device=device, dtype=torch.float32)
114 +
115 + for expert in range(num_experts):
116 + w1_e = w1[expert]
117 + w2_e = w2[expert]
118 + w1_amax = torch.abs(w1_e).max().to(torch.float32)
119 + w2_amax = torch.abs(w2_e).max().to(torch.float32)
120 + w1_gs[expert] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w1_amax
121 + w2_gs[expert] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w2_amax
122 +
123 + w1_fp4[expert], w1_blockscale[expert] = ops.scaled_fp4_quant(
124 + w1_e, w1_gs[expert])
125 +
126 + w2_fp4[expert], w2_blockscale[expert] = ops.scaled_fp4_quant(
127 + w2_e, w2_gs[expert])
128 +
129 + def run_triton_moe(a: torch.Tensor, w1: torch.Tensor, w2: torch.Tensor,
130 + topk_weights: torch.Tensor, topk_ids: torch.Tensor,
131 + w1_scale: torch.Tensor, w2_scale: torch.Tensor,
132 + a_fp8_scale: torch.Tensor, num_repeats: int):
133 + for _ in range(num_repeats):
134 + fused_experts(a,
135 + w1,
136 + w2,
137 + topk_weights,
138 + topk_ids,
139 + use_fp8_w8a8=True,
140 + w1_scale=w1_scale,
141 + w2_scale=w2_scale,
142 + a1_scale=a_fp8_scale)
143 +
144 + def run_cutlass_moe_fp4(a: torch.Tensor, w1_fp4: torch.Tensor,
145 + w2_fp4: torch.Tensor, w1_blockscale: torch.Tensor,
146 + w2_blockscale: torch.Tensor, w1_gs: torch.Tensor,
147 + w2_gs: torch.Tensor, a1_gs: torch.Tensor,
148 + a2_gs: torch.Tensor, topk_weights: torch.Tensor,
149 + topk_ids: torch.Tensor, m: int, n: int, k: int,
150 + e: int, device: torch.device, num_repeats: int):
151 + for _ in range(num_repeats):
152 + with nvtx.annotate("cutlass_moe_fp4", color="green"):
153 + cutlass_moe_fp4(a=a,
154 + a1_gscale=a1_gs,
155 + a2_gscale=a2_gs,
156 + w1_fp4=w1_fp4,
157 + w1_blockscale=w1_blockscale,
158 + w1_alphas=w1_gs,
159 + w2_fp4=w2_fp4,
160 + w2_blockscale=w2_blockscale,
161 + w2_alphas=w2_gs,
162 + topk_weights=topk_weights,
163 + topk_ids=topk_ids,
164 + m=m,
165 + n=n,
166 + k=k,
167 + e=num_experts,
168 + device=device)
169 +
170 + def run_cutlass_from_graph(
171 + a: torch.Tensor, a1_gscale: torch.Tensor, w1_fp4: torch.Tensor,
172 + w1_blockscale: torch.Tensor, w1_alphas: torch.Tensor,
173 + a2_gscale: torch.Tensor, w2_fp4: torch.Tensor,
174 + w2_blockscale: torch.Tensor, w2_alphas: torch.Tensor,
175 + topk_weights: torch.Tensor, topk_ids: torch.Tensor, m: int, n: int,
176 + k: int, e: int, device: torch.device):
177 + with set_current_vllm_config(
178 + VllmConfig(parallel_config=ParallelConfig(
179 + pipeline_parallel_size=1))):
180 + return cutlass_moe_fp4(a=a,
181 + a1_gscale=a1_gs,
182 + w1_fp4=w1_fp4,
183 + w1_blockscale=w1_blockscale,
184 + w1_alphas=w1_alphas,
185 + a2_gscale=a2_gs,
186 + w2_fp4=w2_fp4,
187 + w2_blockscale=w2_blockscale,
188 + w2_alphas=w2_alphas,
189 + topk_weights=topk_weights,
190 + topk_ids=topk_ids,
191 + m=m,
192 + n=n,
193 + k=k,
194 + e=num_experts,
195 + device=device)
196 +
197 + def run_triton_from_graph(a: torch.Tensor, w1: torch.Tensor,
198 + w2: torch.Tensor, topk_weights: torch.Tensor,
199 + topk_ids: torch.Tensor, w1_scale: torch.Tensor,
200 + w2_scale: torch.Tensor,
201 + a_fp8_scale: torch.Tensor):
202 + with set_current_vllm_config(
203 + VllmConfig(parallel_config=ParallelConfig(
204 + pipeline_parallel_size=1))):
205 + return fused_experts(a,
206 + w1,
207 + w2,
208 + topk_weights,
209 + topk_ids,
210 + use_fp8_w8a8=True,
211 + w1_scale=w1_scale,
212 + w2_scale=w2_scale,
213 + a1_scale=a_fp8_scale)
214 +
215 + def replay_graph(graph, num_repeats):
216 + for _ in range(num_repeats):
217 + graph.replay()
218 + torch.cuda.synchronize()
219 +
220 + cutlass_stream = torch.cuda.Stream()
221 + cutlass_graph = torch.cuda.CUDAGraph()
222 + with torch.cuda.graph(cutlass_graph, stream=cutlass_stream):
223 + run_cutlass_from_graph(a=a,
224 + a1_gscale=a1_gs,
225 + w1_fp4=w1_fp4,
226 + w1_blockscale=w1_blockscale,
227 + w1_alphas=w1_gs,
228 + a2_gscale=a2_gs,
229 + w2_fp4=w2_fp4,
230 + w2_blockscale=w2_blockscale,
231 + w2_alphas=w2_gs,
232 + topk_weights=topk_weights,
233 + topk_ids=topk_ids,
234 + m=m,
235 + n=n,
236 + k=k,
237 + e=num_experts,
238 + device=device)
239 + torch.cuda.synchronize()
240 +
241 + triton_stream = torch.cuda.Stream()
242 + triton_graph = torch.cuda.CUDAGraph()
243 + with torch.cuda.graph(triton_graph, stream=triton_stream):
244 + run_triton_from_graph(a, w1_fp8q_notransp, w2_fp8q_notransp,
245 + topk_weights, topk_ids, w1_fp8scale, w2_fp8scale,
246 + a_fp8_scale)
247 + torch.cuda.synchronize()
248 +
249 + min_run_time = 5
250 + num_warmup = 5
251 + num_runs = 25
252 +
253 + globals = {
254 + # Baseline params
255 + "w1": w1,
256 + "w2": w2,
257 + "score": score,
258 + "topk": topk,
259 + "w1_fp8q_notransp": w1_fp8q_notransp,
260 + "w2_fp8q_notransp": w2_fp8q_notransp,
261 + "w1_fp8scale": w1_fp8scale,
262 + "w2_fp8scale": w2_fp8scale,
263 + "a_fp8_scale": a_fp8_scale,
264 + # Cutlass params
265 + "a": a,
266 + "a1_gscale": a1_gs,
267 + "w1_fp4": w1_fp4,
268 + "w1_blockscale": w1_blockscale,
269 + "w1_alphas": w1_gs,
270 + "a2_gscale": a2_gs,
271 + "w2_fp4": w2_fp4,
272 + "w2_blockscale": w2_blockscale,
273 + "w2_alphas": w2_gs,
274 + "topk_weights": topk_weights,
275 + "topk_ids": topk_ids,
276 + "m": m,
277 + "n": n,
278 + "k": k,
279 + "e": num_experts,
280 + "device": device,
281 + # cuda graph params
282 + "cutlass_graph": cutlass_graph,
283 + "triton_graph": triton_graph,
284 + # Gen params
285 + "num_runs": num_runs,
286 + # Kernels
287 + "run_triton_moe": run_triton_moe,
288 + "run_cutlass_moe_fp4": run_cutlass_moe_fp4,
289 + "replay_graph": replay_graph,
290 + }
291 +
292 + # Warmup
293 + run_triton_moe(a, w1_fp8q_notransp, w2_fp8q_notransp, topk_weights,
294 + topk_ids, w1_fp8scale, w2_fp8scale, a_fp8_scale, num_warmup)
295 +
296 + results.append(
297 + benchmark.Timer(
298 + stmt=
299 + "run_triton_moe(a, w1_fp8q_notransp, w2_fp8q_notransp, topk_weights, topk_ids, w1_fp8scale, w2_fp8scale, a_fp8_scale, num_runs)", # noqa: E501
300 + globals=globals,
301 + label=label,
302 + sub_label=sub_label,
303 + description="triton_moe",
304 + ).blocked_autorange(min_run_time=min_run_time))
305 +
306 + # Warmup
307 + replay_graph(triton_graph, num_warmup)
308 +
309 + results.append(
310 + benchmark.Timer(
311 + stmt="replay_graph(triton_graph, num_runs)",
312 + globals=globals,
313 + label=label,
314 + sub_label=sub_label,
315 + description="triton_moe_cuda_graphs",
316 + ).blocked_autorange(min_run_time=min_run_time))
317 +
318 + # Warmup
319 +
320 + run_cutlass_moe_fp4(a, w1_fp4, w2_fp4, w1_blockscale, w2_blockscale, w1_gs,
321 + w2_gs, a1_gs, a2_gs, topk_weights, topk_ids, m, n, k,
322 + num_experts, device, num_warmup)
323 +
324 + results.append(
325 + benchmark.Timer(
326 + stmt=
327 + "run_cutlass_moe_fp4(a, w1_fp4, w2_fp4, w1_blockscale, w2_blockscale, w1_alphas, w2_alphas, a1_gscale, a2_gscale, topk_weights, topk_ids, m, n, k, e, device, num_runs)", # noqa: E501
328 + globals=globals,
329 + label=label,
330 + sub_label=sub_label,
331 + description="cutlass_moe_fp4",
332 + ).blocked_autorange(min_run_time=min_run_time))
333 +
334 + # Warmup
335 + replay_graph(cutlass_graph, num_warmup)
336 +
337 + results.append(
338 + benchmark.Timer(
339 + stmt="replay_graph(cutlass_graph, num_runs)",
340 + globals=globals,
341 + label=label,
342 + sub_label=sub_label,
343 + description="cutlass_moe_fp4_cuda_graphs",
344 + ).blocked_autorange(min_run_time=min_run_time))
345 +
346 +
347 +def main(args):
348 + print("Benchmarking models:")
349 + for i, model in enumerate(args.models):
350 + print(f"[{i}] {model}")
351 +
352 + results: list[benchmark.Measurement] = []
353 +
354 + for model in args.models:
355 + for tp in args.tp_sizes:
356 + for layer in WEIGHT_SHAPES_MOE[model]:
357 + num_experts = layer[0]
358 + topk = layer[1]
359 + size_k = layer[2]
360 + size_n = layer[3] // tp
361 +
362 + if len(args.limit_k) > 0 and size_k not in args.limit_k:
363 + continue
364 +
365 + if len(args.limit_n) > 0 and size_n not in args.limit_n:
366 + continue
367 +
368 + for per_act_token in PER_ACT_TOKEN_OPTS:
369 + for per_out_ch in PER_OUT_CH_OPTS:
370 + for size_m in args.batch_sizes:
371 + mkn = (size_m, size_k, size_n)
372 + bench_run(results, model, num_experts, topk,
373 + per_act_token, per_out_ch, mkn)
374 +
375 + compare = benchmark.Compare(results)
376 + compare.print()
377 +
378 +
379 +if __name__ == "__main__":
380 + parser = FlexibleArgumentParser(
381 + description="Benchmark NVFP4 CUTLASS MOE across specified "
382 + "models/shapes/batches")
383 + parser.add_argument(
384 + "--models",
385 + nargs="+",
386 + type=str,
387 + default=DEFAULT_MODELS,
388 + choices=WEIGHT_SHAPES_MOE.keys(),
389 + )
390 + parser.add_argument("--tp-sizes",
391 + nargs="+",
392 + type=int,
393 + default=DEFAULT_TP_SIZES)
394 + parser.add_argument("--batch-sizes",
395 + nargs="+",
396 + type=int,
397 + default=DEFAULT_BATCH_SIZES)
398 + parser.add_argument("--limit-k", nargs="+", type=int, default=[])
399 + parser.add_argument("--limit-n", nargs="+", type=int, default=[])
400 + parser.add_argument("--limit-num-groups", nargs="+", type=int, default=[])
401 + parser.add_argument("--limit-per-act-token",
402 + nargs="+",
403 + type=int,
404 + default=[])
405 + parser.add_argument("--limit-per-out-ch", nargs="+", type=int, default=[])
406 +
407 + args = parser.parse_args()
408 + main(args)
csrc/ops.h
+12
@@ -208,6 +208,12 @@ void cutlass_moe_mm(
208 torch::Tensor const& problem_sizes, torch::Tensor const& a_strides,
209 torch::Tensor const& b_strides, torch::Tensor const& c_strides);
210
211 +void cutlass_fp4_group_mm(
212 + torch::Tensor& output, const torch::Tensor& a, const torch::Tensor& b,
213 + const torch::Tensor& a_blockscale, const torch::Tensor& b_blockscales,
214 + const torch::Tensor& alphas, const torch::Tensor& problem_sizes,
215 + const torch::Tensor& expert_offsets, const torch::Tensor& sf_offsets);
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,
@@ -235,6 +241,12 @@ std::vector<torch::Tensor> cutlass_sparse_compress(torch::Tensor const& a);
241 void scaled_fp4_quant(torch::Tensor& output, torch::Tensor const& input,
242 torch::Tensor& output_scale,
243 torch::Tensor const& input_scale);
244 +
245 +void scaled_fp4_experts_quant(
246 + torch::Tensor& output, torch::Tensor& output_scale,
247 + torch::Tensor const& input, torch::Tensor const& input_global_scale,
248 + torch::Tensor const& input_offset_by_experts,
249 + torch::Tensor const& output_scale_offset_by_experts);
250 #endif
251
252 void static_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input,
csrc/quantization/cutlass_w8a8/scaled_mm_entry.cu
+11 -7
@@ -37,12 +37,6 @@ void cutlass_moe_mm_sm90(
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 -
40 #endif
41
42 #if defined ENABLE_SCALED_MM_SM100 && ENABLE_SCALED_MM_SM100
@@ -53,6 +47,15 @@ void cutlass_scaled_mm_sm100(torch::Tensor& c, torch::Tensor const& a,
47 std::optional<torch::Tensor> const& bias);
48 #endif
49
50 +#if defined(ENABLE_SCALED_MM_SM90) && ENABLE_SCALED_MM_SM90 || \
51 + defined(ENABLE_SCALED_MM_SM100) && ENABLE_SCALED_MM_SM100
52 +void get_cutlass_moe_mm_data_caller(
53 + const torch::Tensor& topk_ids, torch::Tensor& expert_offsets,
54 + torch::Tensor& problem_sizes1, torch::Tensor& problem_sizes2,
55 + torch::Tensor& input_permutation, torch::Tensor& output_permutation,
56 + const int64_t num_experts, const int64_t n, const int64_t k);
57 +#endif
58 +
59 void cutlass_scaled_mm_azp_sm75(torch::Tensor& c, torch::Tensor const& a,
60 torch::Tensor const& b,
61 torch::Tensor const& a_scales,
@@ -224,7 +227,8 @@ void get_cutlass_moe_mm_data(
227 // This function currently gets compiled only if we have a valid cutlass moe
228 // mm to run it for.
229 int32_t version_num = get_sm_version_num();
227 -#if defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90
230 +#if (defined ENABLE_CUTLASS_MOE_SM90 && ENABLE_CUTLASS_MOE_SM90) || \
231 + (defined ENABLE_SCALED_MM_SM100 && ENABLE_SCALED_MM_SM90)
232 get_cutlass_moe_mm_data_caller(topk_ids, expert_offsets, problem_sizes1,
233 problem_sizes2, input_permutation,
234 output_permutation, num_experts, n, k);
csrc/quantization/fp4/nvfp4_blockwise_moe_kernel.cu new
+402
@@ -0,0 +1,402 @@
1 +#include <torch/all.h>
2 +#include <cutlass/arch/arch.h>
3 +
4 +#include <ATen/cuda/CUDAContext.h>
5 +#include <c10/cuda/CUDAGuard.h>
6 +#include <c10/cuda/CUDAStream.h>
7 +
8 +#include "cute/tensor.hpp"
9 +#include "cutlass/tensor_ref.h"
10 +#include "cutlass/epilogue/collective/default_epilogue.hpp"
11 +#include "cutlass/epilogue/thread/linear_combination.h"
12 +#include "cutlass/gemm/dispatch_policy.hpp"
13 +#include "cutlass/gemm/group_array_problem_shape.hpp"
14 +#include "cutlass/gemm/collective/collective_builder.hpp"
15 +#include "cutlass/epilogue/collective/collective_builder.hpp"
16 +#include "cutlass/gemm/device/gemm_universal_adapter.h"
17 +#include "cutlass/gemm/kernel/gemm_universal.hpp"
18 +
19 +#include "cutlass/util/command_line.h"
20 +#include "cutlass/util/distribution.h"
21 +#include "cutlass/util/host_tensor.h"
22 +#include "cutlass/util/packed_stride.hpp"
23 +#include "cutlass/util/tensor_view_io.h"
24 +#include "cutlass/util/reference/device/gemm.h"
25 +#include "cutlass/util/reference/device/tensor_compare.h"
26 +#include "cutlass/util/reference/host/tensor_fill.h"
27 +#include "cutlass/util/reference/host/gett.hpp"
28 +#include "cutlass/util/reference/host/tensor_norm.h"
29 +#include "cutlass/util/reference/host/tensor_compare.h"
30 +#include <cassert>
31 +
32 +using namespace cute;
33 +
34 +template <typename ElementAB, typename ElementC, typename ElementSF,
35 + typename ElementAccumulator, typename LayoutSFA, typename LayoutSFB,
36 + typename ScaleConfig>
37 +__global__ void __get_group_gemm_starts(
38 + ElementAB** a_offsets, ElementAB** b_offsets, ElementC** out_offsets,
39 + ElementSF** a_scales_offsets, ElementSF** b_scales_offsets,
40 + ElementAccumulator** alpha_offsets, LayoutSFA* layout_sfa_base_as_int,
41 + LayoutSFB* layout_sfb_base_as_int, ElementAB* a_base_as_int,
42 + ElementAB* b_base_as_int, ElementC* out_base_as_int,
43 + ElementSF* a_scales_base_as_int, ElementSF* b_scales_base_as_int,
44 + ElementAccumulator* alphas_base_as_int, const int32_t* expert_offsets,
45 + const int32_t* sf_offsets, const int32_t* problem_sizes_as_shapes,
46 + const int K, const int N) {
47 + int64_t expert_id = threadIdx.x;
48 + if (expert_id >= gridDim.x * blockDim.x) {
49 + return;
50 + }
51 + // Originally int32_t but upcasting to int64_t to avoid overflow
52 + // during offset calculations
53 + int64_t expert_offset = static_cast<int64_t>(expert_offsets[expert_id]);
54 + int64_t sf_offset = static_cast<int64_t>(sf_offsets[expert_id]);
55 + // size for block in block scale.
56 + int64_t group_size = 16;
57 + int64_t m = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3]);
58 + int64_t n = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3 + 1]);
59 + int64_t k = static_cast<int64_t>(problem_sizes_as_shapes[expert_id * 3 + 2]);
60 + assert((m >= 0 && n == N && k == K && k % 2 == 0) &&
61 + "unexpected problem sizes");
62 +
63 + int64_t half_k = static_cast<int64_t>(k / 2);
64 + int64_t group_k = static_cast<int64_t>(k / group_size);
65 + // Shape of A as uint8/byte = [M, K // 2]
66 + // Shape of B as uint8/byte = [E, N, K // 2]
67 + a_offsets[expert_id] = a_base_as_int + expert_offset * half_k;
68 +
69 + b_offsets[expert_id] = b_base_as_int + expert_id * n * half_k;
70 + // Shape of C = [M, N]
71 + out_offsets[expert_id] = out_base_as_int + expert_offset * n;
72 + // Shape of a_scale = [sum(sf_sizes), K // group_size]
73 + a_scales_offsets[expert_id] = a_scales_base_as_int + sf_offset * group_k;
74 +
75 + assert((reinterpret_cast<uintptr_t>(a_scales_offsets[expert_id]) % 128) ==
76 + 0 &&
77 + "TMA requires 128-byte alignment");
78 +
79 + // Shape of B scale = [E, N, K // group_size]
80 + b_scales_offsets[expert_id] = b_scales_base_as_int + expert_id * n * group_k;
81 + assert((reinterpret_cast<uintptr_t>(b_scales_offsets[expert_id]) % 128) ==
82 + 0 &&
83 + "TMA requires 128-byte alignment");
84 + // Shape of alpha = [E]
85 + alpha_offsets[expert_id] = alphas_base_as_int + expert_id;
86 +
87 + LayoutSFA* layout_sfa_ptr = layout_sfa_base_as_int + expert_id;
88 + LayoutSFB* layout_sfb_ptr = layout_sfb_base_as_int + expert_id;
89 +
90 + *layout_sfa_ptr = ScaleConfig::tile_atom_to_shape_SFA(cute::make_shape(
91 + static_cast<int>(m), static_cast<int>(n), static_cast<int>(k), 1));
92 + *layout_sfb_ptr = ScaleConfig::tile_atom_to_shape_SFB(cute::make_shape(
93 + static_cast<int>(m), static_cast<int>(n), static_cast<int>(k), 1));
94 +}
95 +
96 +#define __CALL_GET_STARTS_KERNEL_BLOCKSCALE(ELEMENT_AB_TYPE, SF_TYPE, \
97 + TENSOR_C_TYPE, C_TYPE, LayoutSFA, \
98 + LayoutSFB, ScaleConfig) \
99 + else if (out_tensors.dtype() == TENSOR_C_TYPE) { \
100 + __get_group_gemm_starts<ELEMENT_AB_TYPE, C_TYPE, SF_TYPE, float, \
101 + LayoutSFA, LayoutSFB, ScaleConfig> \
102 + <<<1, num_experts, 0, stream>>>( \
103 + static_cast<ELEMENT_AB_TYPE**>(a_starts.data_ptr()), \
104 + static_cast<ELEMENT_AB_TYPE**>(b_starts.data_ptr()), \
105 + static_cast<C_TYPE**>(out_starts.data_ptr()), \
106 + static_cast<SF_TYPE**>(a_scales_starts.data_ptr()), \
107 + static_cast<SF_TYPE**>(b_scales_starts.data_ptr()), \
108 + static_cast<float**>(alpha_starts.data_ptr()), \
109 + reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()), \
110 + reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr()), \
111 + static_cast<ELEMENT_AB_TYPE*>(a_tensors.data_ptr()), \
112 + static_cast<ELEMENT_AB_TYPE*>(b_tensors.data_ptr()), \
113 + static_cast<C_TYPE*>(out_tensors.data_ptr()), \
114 + static_cast<SF_TYPE*>(a_scales.data_ptr()), \
115 + static_cast<SF_TYPE*>(b_scales.data_ptr()), \
116 + static_cast<float*>(alphas.data_ptr()), \
117 + static_cast<int32_t*>(expert_offsets.data_ptr()), \
118 + static_cast<int32_t*>(sf_offsets.data_ptr()), \
119 + static_cast<int32_t*>(problem_sizes.data_ptr()), K, N); \
120 + }
121 +
122 +template <typename LayoutSFA, typename LayoutSFB, typename ScaleConfig>
123 +void run_get_group_gemm_starts(
124 + const torch::Tensor& a_starts, const torch::Tensor& b_starts,
125 + const torch::Tensor& out_starts, const torch::Tensor& a_scales_starts,
126 + const torch::Tensor& b_scales_starts, const torch::Tensor& alpha_starts,
127 + const torch::Tensor& layout_sfa, const torch::Tensor& layout_sfb,
128 + /*these are used for their base addresses*/
129 + torch::Tensor const& a_tensors, torch::Tensor const& b_tensors,
130 + torch::Tensor const& out_tensors, torch::Tensor const& a_scales,
131 + torch::Tensor const& b_scales, torch::Tensor const& alphas,
132 + torch::Tensor const& expert_offsets, torch::Tensor const& sf_offsets,
133 + torch::Tensor const& problem_sizes, int M, int N, int K) {
134 + int num_experts = (int)expert_offsets.size(0);
135 + auto stream = at::cuda::getCurrentCUDAStream(a_tensors.device().index());
136 +
137 + TORCH_CHECK(out_tensors.size(1) == N,
138 + "Output tensor shape doesn't match expected shape");
139 + TORCH_CHECK(K / 2 == b_tensors.size(2),
140 + "b_tensors(dim = 2) and a_tensors(dim = 1) trailing"
141 + " dimension must match");
142 + if (false) {
143 + }
144 + //(ELEMENT_AB_TYPE, BS_TYPE, TENSOR_C_TYPE, C_TYPE, LayoutSFA, LayoutSFB,
145 + // ScaleConfig)
146 + __CALL_GET_STARTS_KERNEL_BLOCKSCALE(
147 + cutlass::float_e2m1_t, cutlass::float_ue4m3_t, torch::kBFloat16,
148 + cutlass::bfloat16_t, LayoutSFA, LayoutSFB, ScaleConfig)
149 + __CALL_GET_STARTS_KERNEL_BLOCKSCALE(cutlass::float_e2m1_t,
150 + cutlass::float_ue4m3_t, torch::kFloat16,
151 + half, LayoutSFA, LayoutSFB, ScaleConfig)
152 + else {
153 + TORCH_CHECK(false, "Invalid output type (must be float16 or bfloat16)");
154 + }
155 +}
156 +
157 +template <typename OutType>
158 +void run_fp4_blockwise_scaled_group_mm(
159 + torch::Tensor& output, const torch::Tensor& a, const torch::Tensor& b,
160 + const torch::Tensor& a_blockscale, const torch::Tensor& b_blockscales,
161 + const torch::Tensor& alphas, const torch::Tensor& problem_sizes,
162 + const torch::Tensor& expert_offsets, const torch::Tensor& sf_offsets, int M,
163 + int N, int K) {
164 + using ProblemShape =
165 + cutlass::gemm::GroupProblemShape<Shape<int32_t, int32_t, int32_t>>;
166 + using ElementType = cutlass::float_e2m1_t;
167 + using ElementSFType = cutlass::float_ue4m3_t;
168 + using ElementA = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
169 + using ElementB = cutlass::nv_float4_t<cutlass::float_e2m1_t>;
170 +
171 + using ElementC = OutType;
172 + using ElementD = ElementC;
173 + using ElementAccumulator = float;
174 + // Layout definitions
175 + using LayoutA = cutlass::layout::RowMajor;
176 + using LayoutB = cutlass::layout::ColumnMajor;
177 + using LayoutC = cutlass::layout::RowMajor;
178 + using LayoutD = LayoutC;
179 +
180 + // Alignment constraints
181 + static constexpr int AlignmentA = 32;
182 + static constexpr int AlignmentB = 32;
183 + static constexpr int AlignmentC = 128 / cutlass::sizeof_bits<ElementC>::value;
184 + static constexpr int AlignmentD = 128 / cutlass::sizeof_bits<ElementD>::value;
185 +
186 + // Architecture definitions
187 + using ArchTag = cutlass::arch::Sm100;
188 + using EpilogueOperatorClass =
189 + cutlass::arch::OpClassTensorOp; // Epilogue Operator class tag
190 + using MainloopOperatorClass =
191 + cutlass::arch::OpClassBlockScaledTensorOp; // Mainloop Operator class tag
192 + using StageCountType =
193 + cutlass::gemm::collective::StageCountAuto; // Stage count maximized based
194 + // on the tile size
195 +
196 + using ClusterShape = Shape<_1, _1, _1>;
197 + struct MMA1SMConfig {
198 + using MmaTileShape = Shape<_128, _128, _128>;
199 + using KernelSchedule = cutlass::gemm::
200 + KernelPtrArrayTmaWarpSpecialized1SmNvf4Sm100; // Kernel to launch
201 + using EpilogueSchedule =
202 + cutlass::epilogue::PtrArrayTmaWarpSpecialized1Sm; // Epilogue to launch
203 + };
204 +
205 + using CollectiveEpilogue =
206 + typename cutlass::epilogue::collective::CollectiveBuilder<
207 + ArchTag, EpilogueOperatorClass, typename MMA1SMConfig::MmaTileShape,
208 + ClusterShape, Shape<_128, _64>, ElementAccumulator,
209 + ElementAccumulator, ElementC, LayoutC*, AlignmentC, ElementD,
210 + LayoutC*, AlignmentD,
211 + typename MMA1SMConfig::EpilogueSchedule>::CollectiveOp;
212 +
213 + using CollectiveMainloop =
214 + typename cutlass::gemm::collective::CollectiveBuilder<
215 + ArchTag, MainloopOperatorClass, ElementA, LayoutA*, AlignmentA,
216 + ElementB, LayoutB*, AlignmentB, ElementAccumulator,
217 + typename MMA1SMConfig::MmaTileShape, ClusterShape,
218 + cutlass::gemm::collective::StageCountAutoCarveout<static_cast<int>(
219 + sizeof(typename CollectiveEpilogue::SharedStorage))>,
220 + typename MMA1SMConfig::KernelSchedule>::CollectiveOp;
221 +
222 + using GemmKernel =
223 + cutlass::gemm::kernel::GemmUniversal<ProblemShape, CollectiveMainloop,
224 + CollectiveEpilogue>;
225 +
226 + using Gemm1SM = cutlass::gemm::device::GemmUniversalAdapter<GemmKernel>;
227 + using Gemm = Gemm1SM;
228 + using StrideA = typename Gemm::GemmKernel::InternalStrideA;
229 + using StrideB = typename Gemm::GemmKernel::InternalStrideB;
230 + using StrideC = typename Gemm::GemmKernel::InternalStrideC;
231 + using StrideD = typename Gemm::GemmKernel::InternalStrideD;
232 +
233 + using LayoutSFA =
234 + typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFA;
235 + using LayoutSFB =
236 + typename Gemm::GemmKernel::CollectiveMainloop::InternalLayoutSFB;
237 + using ScaleConfig =
238 + typename Gemm::GemmKernel::CollectiveMainloop::Sm1xxBlkScaledConfig;
239 +
240 + using UnderlyingProblemShape = ProblemShape::UnderlyingProblemShape;
241 + int num_experts = static_cast<int>(expert_offsets.size(0));
242 + auto options_int =
243 + torch::TensorOptions().dtype(torch::kInt64).device(a.device());
244 +
245 + torch::Tensor a_ptrs = torch::empty(num_experts, options_int);
246 + torch::Tensor b_ptrs = torch::empty(num_experts, options_int);
247 + torch::Tensor out_ptrs = torch::empty(num_experts, options_int);
248 + torch::Tensor a_scales_ptrs = torch::empty(num_experts, options_int);
249 + torch::Tensor b_scales_ptrs = torch::empty(num_experts, options_int);
250 + torch::Tensor alpha_ptrs = torch::empty(num_experts, options_int);
251 + torch::Tensor layout_sfa = torch::empty({num_experts, 5}, options_int);
252 + torch::Tensor layout_sfb = torch::empty({num_experts, 5}, options_int);
253 + torch::Tensor c_strides1 =
254 + torch::full({num_experts}, output.stride(0), options_int);
255 + torch::Tensor a_strides1 =
256 + torch::full({num_experts}, a.stride(0) * 2, options_int);
257 + torch::Tensor b_strides1 =
258 + torch::full({num_experts}, b.stride(1) * 2, options_int);
259 +
260 + run_get_group_gemm_starts<LayoutSFA, LayoutSFB, ScaleConfig>(
261 + a_ptrs, b_ptrs, out_ptrs, a_scales_ptrs, b_scales_ptrs, alpha_ptrs,
262 + layout_sfa, layout_sfb, a, b, output, a_blockscale, b_blockscales, alphas,
263 + expert_offsets, sf_offsets, problem_sizes, M, N, K);
264 +
265 + // Create an instance of the GEMM
266 + Gemm gemm_op;
267 +
268 + // Initialize problem_sizes_as_shapes correctly
269 + UnderlyingProblemShape* problem_sizes_as_shapes =
270 + static_cast<UnderlyingProblemShape*>(problem_sizes.data_ptr());
271 +
272 + // Set the Scheduler info
273 + cutlass::KernelHardwareInfo hw_info;
274 + using RasterOrderOptions = typename cutlass::gemm::kernel::detail::
275 + PersistentTileSchedulerSm100GroupParams<
276 + typename ProblemShape::UnderlyingProblemShape>::RasterOrderOptions;
277 + typename Gemm::GemmKernel::TileSchedulerArguments scheduler;
278 + scheduler.raster_order = RasterOrderOptions::AlongM;
279 + hw_info.device_id = a.get_device();
280 + static std::unordered_map<int, int> cached_sm_counts;
281 + if (cached_sm_counts.find(hw_info.device_id) == cached_sm_counts.end()) {
282 + cached_sm_counts[hw_info.device_id] =
283 + cutlass::KernelHardwareInfo::query_device_multiprocessor_count(
284 + hw_info.device_id);
285 + }
286 + hw_info.sm_count = min(cached_sm_counts[hw_info.device_id], INT_MAX);
287 +
288 + // Mainloop Arguments
289 + typename GemmKernel::MainloopArguments mainloop_args{
290 + static_cast<const ElementType**>(a_ptrs.data_ptr()),
291 + static_cast<StrideA*>(a_strides1.data_ptr()),
292 + static_cast<const ElementType**>(b_ptrs.data_ptr()),
293 + static_cast<StrideB*>(b_strides1.data_ptr()),
294 + static_cast<const ElementSFType**>(a_scales_ptrs.data_ptr()),
295 + reinterpret_cast<LayoutSFA*>(layout_sfa.data_ptr()),
296 + static_cast<const ElementSFType**>(b_scales_ptrs.data_ptr()),
297 + reinterpret_cast<LayoutSFB*>(layout_sfb.data_ptr())};
298 +
299 + // Epilogue Arguments
300 + typename GemmKernel::EpilogueArguments epilogue_args{
301 + {}, // epilogue.thread
302 + nullptr,
303 + static_cast<StrideC*>(c_strides1.data_ptr()),
304 + static_cast<ElementD**>(out_ptrs.data_ptr()),
305 + static_cast<StrideC*>(c_strides1.data_ptr())};
306 + auto& fusion_args = epilogue_args.thread;
307 + fusion_args.alpha_ptr_array =
308 + reinterpret_cast<float**>(alpha_ptrs.data_ptr());
309 + fusion_args.dAlpha = {_0{}, _0{}, 1};
310 +
311 + // Gemm Arguments
312 + typename GemmKernel::Arguments args{
313 + cutlass::gemm::GemmUniversalMode::kGrouped,
314 + {num_experts, problem_sizes_as_shapes, nullptr},
315 + mainloop_args,
316 + epilogue_args,
317 + hw_info,
318 + scheduler};
319 +
320 + size_t workspace_size = Gemm::get_workspace_size(args);
321 + auto const workspace_options =
322 + torch::TensorOptions().dtype(torch::kUInt8).device(a.device());
323 + auto workspace = torch::empty(workspace_size, workspace_options);
324 + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(a.get_device());
325 +
326 + auto can_implement_status = gemm_op.can_implement(args);
327 + TORCH_CHECK(can_implement_status == cutlass::Status::kSuccess,
328 + "Failed to implement GEMM");
329 +
330 + // Run the GEMM
331 + auto status = gemm_op.initialize(args, workspace.data_ptr());
332 + TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to initialize GEMM");
333 +
334 + status = gemm_op.run(args, workspace.data_ptr(), stream);
335 + TORCH_CHECK(status == cutlass::Status::kSuccess, "Failed to run GEMM");
336 +}
337 +
338 +constexpr auto FLOAT4_E2M1X2 = at::ScalarType::Byte;
339 +constexpr auto SF_DTYPE = at::ScalarType::Float8_e4m3fn;
340 +
341 +#define CHECK_TYPE(x, st, m) \
342 + TORCH_CHECK(x.scalar_type() == st, ": Inconsistency of Tensor type:", m)
343 +#define CHECK_TH_CUDA(x, m) \
344 + TORCH_CHECK(x.is_cuda(), m, ": must be a CUDA tensor.")
345 +#define CHECK_CONTIGUOUS(x, m) \
346 + TORCH_CHECK(x.is_contiguous(), m, ": must be contiguous.")
347 +#define CHECK_INPUT(x, st, m) \
348 + CHECK_TH_CUDA(x, m); \
349 + CHECK_CONTIGUOUS(x, m); \
350 + CHECK_TYPE(x, st, m)
351 +
352 +void cutlass_fp4_group_mm(
353 + torch::Tensor& output, const torch::Tensor& a, const torch::Tensor& b,
354 + const torch::Tensor& a_blockscale, const torch::Tensor& b_blockscales,
355 + const torch::Tensor& alphas, const torch::Tensor& problem_sizes,
356 + const torch::Tensor& expert_offsets, const torch::Tensor& sf_offsets) {
357 +#if defined ENABLE_NVFP4 && ENABLE_NVFP4
358 + // Input validation
359 + CHECK_INPUT(a, FLOAT4_E2M1X2, "a");
360 + CHECK_INPUT(b, FLOAT4_E2M1X2, "b");
361 + CHECK_INPUT(a_blockscale, SF_DTYPE, "a_blockscale");
362 + CHECK_INPUT(b_blockscales, SF_DTYPE, "b_blockscales");
363 + CHECK_INPUT(alphas, at::ScalarType::Float, "alphas");
364 +
365 + TORCH_CHECK(a_blockscale.dim() == 2,
366 + "expected a_blockscale to be of shape [num_experts, rounded_m,"
367 + " k // group_size], observed rank: ",
368 + a_blockscale.dim())
369 + TORCH_CHECK(b_blockscales.dim() == 3,
370 + "expected b_blockscale to be of shape: "
371 + " [num_experts, n, k // group_size], observed rank: ",
372 + b_blockscales.dim())
373 + TORCH_CHECK(problem_sizes.dim() == 2, "problem_sizes must be a 2D tensor");
374 + TORCH_CHECK(problem_sizes.size(1) == 3,
375 + "problem_sizes must have the shape (num_experts, 3)");
376 + TORCH_CHECK(problem_sizes.size(0) == expert_offsets.size(0),
377 + "Number of experts in problem_sizes must match expert_offsets");
378 + TORCH_CHECK(problem_sizes.dtype() == torch::kInt32,
379 + "problem_sizes must be int32.");
380 +
381 + int M = static_cast<int>(a.size(0));
382 + int N = static_cast<int>(b.size(1));
383 + int E = static_cast<int>(b.size(0));
384 + int K = static_cast<int>(2 * b.size(2));
385 +
386 + if (output.scalar_type() == torch::kBFloat16) {
387 + run_fp4_blockwise_scaled_group_mm<cutlass::bfloat16_t>(
388 + output, a, b, a_blockscale, b_blockscales, alphas, problem_sizes,
389 + expert_offsets, sf_offsets, M, N, K);
390 + } else {
391 + run_fp4_blockwise_scaled_group_mm<cutlass::half_t>(
392 + output, a, b, a_blockscale, b_blockscales, alphas, problem_sizes,
393 + expert_offsets, sf_offsets, M, N, K);
394 + }
395 +#else
396 + TORCH_CHECK_NOT_IMPLEMENTED(
397 + false,
398 + "No compiled cutlass_fp4_group_mm kernel, vLLM must "
399 + "be compiled with ENABLE_NVFP4 for SM100+ and CUDA "
400 + "12.8 or above.");
401 +#endif
402 +}
csrc/quantization/fp4/nvfp4_experts_quant.cu new
+404
@@ -0,0 +1,404 @@
1 +#include <torch/all.h>
2 +
3 +#include <ATen/cuda/CUDAContext.h>
4 +#include <c10/cuda/CUDAGuard.h>
5 +
6 +#include <cuda_runtime.h>
7 +#include <cuda_fp8.h>
8 +
9 +template <typename T>
10 +struct TypeConverter {
11 + using Type = half2;
12 +}; // keep for generality
13 +
14 +template <>
15 +struct TypeConverter<half2> {
16 + using Type = half;
17 +};
18 +
19 +template <>
20 +struct TypeConverter<half> {
21 + using Type = half2;
22 +};
23 +
24 +template <>
25 +struct TypeConverter<__nv_bfloat162> {
26 + using Type = __nv_bfloat16;
27 +};
28 +
29 +template <>
30 +struct TypeConverter<__nv_bfloat16> {
31 + using Type = __nv_bfloat162;
32 +};
33 +
34 +#define ELTS_PER_THREAD 8
35 +
36 +constexpr int CVT_FP4_ELTS_PER_THREAD = 8;
37 +constexpr int CVT_FP4_SF_VEC_SIZE = 16;
38 +
39 +// Convert 8 float32 values into 8 e2m1 values (represented as one uint32_t).
40 +inline __device__ uint32_t fp32_vec_to_e2m1(float (&array)[8]) {
41 +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
42 + uint32_t val;
43 + asm volatile(
44 + "{\n"
45 + ".reg .b8 byte0;\n"
46 + ".reg .b8 byte1;\n"
47 + ".reg .b8 byte2;\n"
48 + ".reg .b8 byte3;\n"
49 + "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n"
50 + "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n"
51 + "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n"
52 + "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n"
53 + "mov.b32 %0, {byte0, byte1, byte2, byte3};\n"
54 + "}"
55 + : "=r"(val)
56 + : "f"(array[0]), "f"(array[1]), "f"(array[2]), "f"(array[3]),
57 + "f"(array[4]), "f"(array[5]), "f"(array[6]), "f"(array[7]));
58 + return val;
59 +#else
60 + return 0;
61 +#endif
62 +}
63 +
64 +// Convert 4 float2 values into 8 e2m1 values (represented as one uint32_t).
65 +inline __device__ uint32_t fp32_vec_to_e2m1(float2 (&array)[4]) {
66 +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
67 + uint32_t val;
68 + asm volatile(
69 + "{\n"
70 + ".reg .b8 byte0;\n"
71 + ".reg .b8 byte1;\n"
72 + ".reg .b8 byte2;\n"
73 + ".reg .b8 byte3;\n"
74 + "cvt.rn.satfinite.e2m1x2.f32 byte0, %2, %1;\n"
75 + "cvt.rn.satfinite.e2m1x2.f32 byte1, %4, %3;\n"
76 + "cvt.rn.satfinite.e2m1x2.f32 byte2, %6, %5;\n"
77 + "cvt.rn.satfinite.e2m1x2.f32 byte3, %8, %7;\n"
78 + "mov.b32 %0, {byte0, byte1, byte2, byte3};\n"
79 + "}"
80 + : "=r"(val)
81 + : "f"(array[0].x), "f"(array[0].y), "f"(array[1].x), "f"(array[1].y),
82 + "f"(array[2].x), "f"(array[2].y), "f"(array[3].x), "f"(array[3].y));
83 + return val;
84 +#else
85 + return 0;
86 +#endif
87 +}
88 +
89 +// Fast reciprocal.
90 +inline __device__ float reciprocal_approximate_ftz(float a) {
91 + float b;
92 + asm volatile("rcp.approx.ftz.f32 %0, %1;\n" : "=f"(b) : "f"(a));
93 + return b;
94 +}
95 +
96 +template <class SFType, int CVT_FP4_NUM_THREADS_PER_SF>
97 +__device__ uint8_t* cvt_quant_to_fp4_get_sf_out_offset(int rowIdx, int colIdx,
98 + int numCols,
99 + SFType* SFout) {
100 +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
101 + static_assert(CVT_FP4_NUM_THREADS_PER_SF == 1 ||
102 + CVT_FP4_NUM_THREADS_PER_SF == 2);
103 +
104 + // One pair of threads write one SF to global memory.
105 + // TODO: stage through smem for packed STG.32
106 + // is it better than STG.8 from 4 threads ?
107 + if (threadIdx.x % CVT_FP4_NUM_THREADS_PER_SF == 0) {
108 + // SF vector index (16 elements share one SF in the K dimension).
109 + int32_t kIdx = colIdx / CVT_FP4_NUM_THREADS_PER_SF;
110 + int32_t mIdx = rowIdx;
111 +
112 + // SF layout [numMTiles, numKTiles, 32 (mTile), 4 (mTile), 4(kTile)]
113 + // --> index [mTileIdx, kTileIdx, outerMIdx, innerMIdx, innerKIdx]
114 +
115 + int32_t mTileIdx = mIdx / (32 * 4);
116 + // SF vector size 16.
117 + int factor = CVT_FP4_SF_VEC_SIZE * 4;
118 + int32_t numKTiles = (numCols + factor - 1) / factor;
119 + int64_t mTileStride = numKTiles * 32 * 4 * 4;
120 +
121 + int32_t kTileIdx = (kIdx / 4);
122 + int64_t kTileStride = 32 * 4 * 4;
123 +
124 + // M tile layout [32, 4] is column-major.
125 + int32_t outerMIdx = (mIdx % 32);
126 + int64_t outerMStride = 4 * 4;
127 +
128 + int32_t innerMIdx = (mIdx % (32 * 4)) / 32;
129 + int64_t innerMStride = 4;
130 +
131 + int32_t innerKIdx = (kIdx % 4);
132 + int64_t innerKStride = 1;
133 +
134 + // Compute the global offset.
135 + int64_t SFOffset = mTileIdx * mTileStride + kTileIdx * kTileStride +
136 + outerMIdx * outerMStride + innerMIdx * innerMStride +
137 + innerKIdx * innerKStride;
138 +
139 + return reinterpret_cast<uint8_t*>(SFout) + SFOffset;
140 + }
141 +#endif
142 + return nullptr;
143 +}
144 +
145 +// Define a 16 bytes packed data type.
146 +template <class Type>
147 +struct PackedVec {
148 + typename TypeConverter<Type>::Type elts[4];
149 +};
150 +
151 +template <>
152 +struct PackedVec<__nv_fp8_e4m3> {
153 + __nv_fp8x2_e4m3 elts[8];
154 +};
155 +
156 +// Quantizes the provided PackedVec into the uint32_t output
157 +template <class Type, bool UE8M0_SF = false>
158 +__device__ uint32_t cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal,
159 + uint8_t* SFout) {
160 +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
161 + // Get absolute maximum values among the local 8 values.
162 + auto localMax = __habs2(vec.elts[0]);
163 +
164 + // Local maximum value.
165 + #pragma unroll
166 + for (int i = 1; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
167 + localMax = __hmax2(localMax, __habs2(vec.elts[i]));
168 + }
169 +
170 + // Get the absolute maximum among all 16 values (two threads).
171 + localMax = __hmax2(__shfl_xor_sync(uint32_t(-1), localMax, 1), localMax);
172 + // Get the final absolute maximum values.
173 + float vecMax = float(__hmax(localMax.x, localMax.y));
174 +
175 + // Get the SF (max value of the vector / max value of e2m1).
176 + // maximum value of e2m1 = 6.0.
177 + // TODO: use half as compute data type.
178 + float SFValue = SFScaleVal * (vecMax * reciprocal_approximate_ftz(6.0f));
179 + // 8 bits representation of the SF.
180 + uint8_t fp8SFVal;
181 + // Write the SF to global memory (STG.8).
182 + if constexpr (UE8M0_SF) {
183 + // Extract the 8 exponent bits from float32.
184 + // float 32bits = 1 sign bit + 8 exponent bits + 23 mantissa bits.
185 + uint32_t tmp = reinterpret_cast<uint32_t&>(SFValue) >> 23;
186 + fp8SFVal = tmp & 0xff;
187 + // Convert back to fp32.
188 + reinterpret_cast<uint32_t&>(SFValue) = tmp << 23;
189 + } else {
190 + // Here SFValue is always positive, so E4M3 is the same as UE4M3.
191 + __nv_fp8_e4m3 tmp = __nv_fp8_e4m3(SFValue);
192 + reinterpret_cast<__nv_fp8_e4m3&>(fp8SFVal) = tmp;
193 + // Convert back to fp32.
194 + SFValue = float(tmp);
195 + }
196 + // Get the output scale.
197 + // Recipe: final_scale = reciprocal(fp32(fp8(SFValue * SFScaleVal))) *
198 + // reciprocal(SFScaleVal))
199 + float outputScale =
200 + SFValue != 0 ? reciprocal_approximate_ftz(
201 + SFValue * reciprocal_approximate_ftz(SFScaleVal))
202 + : 0.0f;
203 +
204 + if (SFout) {
205 + // Write the SF to global memory (STG.8).
206 + *SFout = fp8SFVal;
207 + }
208 +
209 + // Convert the input to float.
210 + float2 fp2Vals[CVT_FP4_ELTS_PER_THREAD / 2];
211 +
212 + #pragma unroll
213 + for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; i++) {
214 + if constexpr (std::is_same_v<Type, half>) {
215 + fp2Vals[i] = __half22float2(vec.elts[i]);
216 + } else {
217 + fp2Vals[i] = __bfloat1622float2(vec.elts[i]);
218 + }
219 + fp2Vals[i].x *= outputScale;
220 + fp2Vals[i].y *= outputScale;
221 + }
222 +
223 + // Convert to e2m1 values.
224 + uint32_t e2m1Vec = fp32_vec_to_e2m1(fp2Vals);
225 +
226 + // Write the e2m1 values to global memory.
227 + return e2m1Vec;
228 +#else
229 + return 0;
230 +#endif
231 +}
232 +
233 +// Use UE4M3 by default.
234 +template <class Type, bool UE8M0_SF = false>
235 +__global__ void
236 +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
237 +__launch_bounds__(512, 4) cvt_fp16_to_fp4(
238 +#else
239 +cvt_fp16_to_fp4(
240 +#endif
241 + int32_t numRows, int32_t numCols, Type const* in, float const* SFScale,
242 + uint32_t* out, uint32_t* SFout, uint32_t* input_offset_by_experts,
243 + uint32_t* output_scale_offset_by_experts, int n_experts) {
244 +#if defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 1000)
245 + using PackedVec = PackedVec<Type>;
246 + static constexpr int CVT_FP4_NUM_THREADS_PER_SF =
247 + (CVT_FP4_SF_VEC_SIZE / CVT_FP4_ELTS_PER_THREAD);
248 + static_assert(sizeof(PackedVec) == sizeof(Type) * CVT_FP4_ELTS_PER_THREAD,
249 + "Vec size is not matched.");
250 +
251 + // Input tensor row/col loops.
252 + for (int rowIdx = blockIdx.x; rowIdx < numRows; rowIdx += gridDim.x) {
253 + for (int colIdx = threadIdx.x; colIdx < numCols / CVT_FP4_ELTS_PER_THREAD;
254 + colIdx += blockDim.x) {
255 + int64_t inOffset = rowIdx * (numCols / CVT_FP4_ELTS_PER_THREAD) + colIdx;
256 + PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
257 + // Get the output tensor offset.
258 + // Same as inOffset because 8 elements are packed into one uint32_t.
259 + int64_t outOffset = inOffset;
260 + auto& out_pos = out[outOffset];
261 +
262 + // Find index within the experts.
263 + int rowIdx_in_expert = 0;
264 + int expert_idx = 0;
265 + for (int i = 0; i < n_experts; i++) {
266 + if (rowIdx >= input_offset_by_experts[i] &&
267 + rowIdx < input_offset_by_experts[i + 1]) {
268 + rowIdx_in_expert = rowIdx - input_offset_by_experts[i];
269 + expert_idx = i;
270 + break;
271 + }
272 + }
273 +
274 + // Get the global scaling factor, which will be applied to the SF.
275 + // Note SFScale is the same as next GEMM's alpha, which is
276 + // (448.f / (Alpha_A / 6.f)).
277 + float const SFScaleVal = SFScale == nullptr ? 1.0f : SFScale[expert_idx];
278 +
279 + int factor = CVT_FP4_SF_VEC_SIZE * 4;
280 + // The actual output_scales dim is computed from the padded numCols.
281 + int32_t numCols_padded = (numCols + factor - 1) / factor * factor;
282 + int numCols_SFout = numCols_padded / CVT_FP4_SF_VEC_SIZE / 4;
283 + uint32_t* SFout_in_expert =
284 + SFout + output_scale_offset_by_experts[expert_idx] * numCols_SFout;
285 +
286 + auto sf_out =
287 + cvt_quant_to_fp4_get_sf_out_offset<uint32_t,
288 + CVT_FP4_NUM_THREADS_PER_SF>(
289 + rowIdx_in_expert, colIdx, numCols, SFout_in_expert);
290 +
291 + out_pos =
292 + cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(in_vec, SFScaleVal, sf_out);
293 + }
294 + }
295 +#endif
296 +}
297 +
298 +template <typename T>
299 +void quant_impl(void* output, void* output_scale, void* input,
300 + void* input_global_scale, void* input_offset_by_experts,
301 + void* output_scale_offset_by_experts, int m_topk, int k,
302 + int n_experts, cudaStream_t stream) {
303 + // TODO: this multiProcessorCount should be cached.
304 + int device;
305 + cudaGetDevice(&device);
306 + int multiProcessorCount;
307 + cudaDeviceGetAttribute(&multiProcessorCount, cudaDevAttrMultiProcessorCount,
308 + device);
309 +
310 + // Grid, Block size.
311 + // Each thread converts 8 values.
312 + dim3 block(std::min(int(k / ELTS_PER_THREAD), 512));
313 + // Get number of blocks per SM (assume we can fully utilize the SM).
314 + int const numBlocksPerSM = 2048 / block.x;
315 + dim3 grid(std::min(int(m_topk), multiProcessorCount * numBlocksPerSM));
316 +
317 + cvt_fp16_to_fp4<T, false><<<grid, block, 0, stream>>>(
318 + m_topk, k, reinterpret_cast<T*>(input),
319 + reinterpret_cast<float*>(input_global_scale),
320 + reinterpret_cast<uint32_t*>(output),
321 + reinterpret_cast<uint32_t*>(output_scale),
322 + reinterpret_cast<uint32_t*>(input_offset_by_experts),
323 + reinterpret_cast<uint32_t*>(output_scale_offset_by_experts), n_experts);
324 +}
325 +
326 +/*Quantization entry for fp4 experts quantization*/
327 +#define CHECK_TH_CUDA(x, m) TORCH_CHECK(x.is_cuda(), m, "must be a CUDA tensor")
328 +#define CHECK_CONTIGUOUS(x, m) \
329 + TORCH_CHECK(x.is_contiguous(), m, "must be contiguous")
330 +#define CHECK_INPUT(x, m) \
331 + CHECK_TH_CUDA(x, m); \
332 + CHECK_CONTIGUOUS(x, m);
333 +
334 +constexpr auto HALF = at::ScalarType::Half;
335 +constexpr auto BF16 = at::ScalarType::BFloat16;
336 +constexpr auto FLOAT = at::ScalarType::Float;
337 +constexpr auto INT = at::ScalarType::Int;
338 +constexpr auto UINT8 = at::ScalarType::Byte;
339 +
340 +void scaled_fp4_experts_quant_sm100a(
341 + torch::Tensor& output, torch::Tensor& output_scale,
342 + torch::Tensor const& input, torch::Tensor const& input_global_scale,
343 + torch::Tensor const& input_offset_by_experts,
344 + torch::Tensor const& output_scale_offset_by_experts) {
345 + CHECK_INPUT(output, "output must be a CUDA tensor");
346 + CHECK_INPUT(output_scale, "output_scale must be a CUDA tensor");
347 + CHECK_INPUT(input, "input must be a CUDA tensor");
348 + CHECK_INPUT(input_global_scale, "input_global_scale must be a CUDA tensor");
349 + CHECK_INPUT(input_offset_by_experts,
350 + "input_offset_by_experts must be a CUDA tensor");
351 + CHECK_INPUT(output_scale_offset_by_experts,
352 + "output_scale_offset_by_experts must be a CUDA tensor");
353 +
354 + TORCH_CHECK(output.dim() == 2);
355 + TORCH_CHECK(output_scale.dim() == 2);
356 + TORCH_CHECK(input.dim() == 2);
357 + TORCH_CHECK(input_global_scale.dim() == 1);
358 + TORCH_CHECK(input_offset_by_experts.dim() == 1);
359 + TORCH_CHECK(output_scale_offset_by_experts.dim() == 1);
360 +
361 + TORCH_CHECK(input.scalar_type() == HALF || input.scalar_type() == BF16);
362 + TORCH_CHECK(input_global_scale.scalar_type() == FLOAT);
363 + TORCH_CHECK(input_offset_by_experts.scalar_type() == INT);
364 + TORCH_CHECK(output_scale_offset_by_experts.scalar_type() == INT);
365 + // output is uint8 (two nvfp4 values are packed into one uint8)
366 + // output_scale is int32 (four fp8 values are packed into one int32)
367 + TORCH_CHECK(output.scalar_type() == UINT8);
368 + TORCH_CHECK(output_scale.scalar_type() == INT);
369 +
370 + const int BLOCK_SIZE = 16;
371 + auto m_topk = input.size(0);
372 + auto k = input.size(1);
373 + TORCH_CHECK(k % BLOCK_SIZE == 0, "k must be a multiple of 16");
374 + auto n_experts = input_global_scale.size(0);
375 + TORCH_CHECK(input_offset_by_experts.size(0) == n_experts + 1);
376 + TORCH_CHECK(output_scale_offset_by_experts.size(0) == n_experts + 1);
377 + TORCH_CHECK(output.size(0) == m_topk);
378 + TORCH_CHECK(output.size(1) == k / 2);
379 + int scales_k = k / BLOCK_SIZE;
380 + // 4 means the swizzle requirement by nvidia nvfp4.
381 + int padded_k = (scales_k + (4 - 1)) / 4 * 4;
382 + // 4 means 4 fp8 values are packed into one int32
383 + TORCH_CHECK(output_scale.size(1) * 4 == padded_k);
384 +
385 + auto in_dtype = input.dtype();
386 + at::cuda::CUDAGuard device_guard{(char)input.get_device()};
387 + const cudaStream_t stream =
388 + at::cuda::getCurrentCUDAStream(input.get_device());
389 + if (in_dtype == at::ScalarType::Half) {
390 + quant_impl<half>(output.data_ptr(), output_scale.data_ptr(),
391 + input.data_ptr(), input_global_scale.data_ptr(),
392 + input_offset_by_experts.data_ptr(),
393 + output_scale_offset_by_experts.data_ptr(), m_topk, k,
394 + n_experts, stream);
395 + } else if (in_dtype == at::ScalarType::BFloat16) {
396 + quant_impl<__nv_bfloat16>(output.data_ptr(), output_scale.data_ptr(),
397 + input.data_ptr(), input_global_scale.data_ptr(),
398 + input_offset_by_experts.data_ptr(),
399 + output_scale_offset_by_experts.data_ptr(), m_topk,
400 + k, n_experts, stream);
401 + } else {
402 + TORCH_CHECK(false, "Expected input data type to be half or bfloat16");
403 + }
404 +}
\ No newline at end of file
csrc/quantization/fp4/nvfp4_quant_entry.cu
+23 -1
@@ -23,10 +23,32 @@ void scaled_fp4_quant_sm100a(torch::Tensor const& output,
23 torch::Tensor const& input_sf);
24 #endif
25
26 +#if defined ENABLE_NVFP4 && ENABLE_NVFP4
27 +void scaled_fp4_experts_quant_sm100a(
28 + torch::Tensor& output, torch::Tensor& output_scale,
29 + torch::Tensor const& input, torch::Tensor const& input_global_scale,
30 + torch::Tensor const& input_offset_by_experts,
31 + torch::Tensor const& output_scale_offset_by_experts);
32 +#endif
33 +
34 void scaled_fp4_quant(torch::Tensor& output, torch::Tensor const& input,
35 torch::Tensor& output_sf, torch::Tensor const& input_sf) {
36 #if defined ENABLE_NVFP4 && ENABLE_NVFP4
37 return scaled_fp4_quant_sm100a(output, input, output_sf, input_sf);
38 #endif
31 - TORCH_CHECK_NOT_IMPLEMENTED(false, "No compiled nvfp4 quantization");
39 + TORCH_CHECK_NOT_IMPLEMENTED(false, "No compiled nvfp4 quantization kernel");
40 +}
41 +
42 +void scaled_fp4_experts_quant(
43 + torch::Tensor& output, torch::Tensor& output_scale,
44 + torch::Tensor const& input, torch::Tensor const& input_global_scale,
45 + torch::Tensor const& input_offset_by_experts,
46 + torch::Tensor const& output_scale_offset_by_experts) {
47 +#if defined ENABLE_NVFP4 && ENABLE_NVFP4
48 + return scaled_fp4_experts_quant_sm100a(
49 + output, output_scale, input, input_global_scale, input_offset_by_experts,
50 + output_scale_offset_by_experts);
51 +#endif
52 + TORCH_CHECK_NOT_IMPLEMENTED(false,
53 + "No compiled nvfp4 experts quantization kernel");
54 }
csrc/torch_bindings.cpp
+15
@@ -363,6 +363,14 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
363 {stride_tag});
364 ops.impl("cutlass_scaled_fp4_mm", torch::kCUDA, &cutlass_scaled_fp4_mm);
365
366 + // cutlass nvfp4 block scaled group GEMM
367 + ops.def(
368 + "cutlass_fp4_group_mm(Tensor! out, Tensor a, Tensor b,"
369 + " Tensor a_blockscale, Tensor b_blockscales, Tensor alphas,"
370 + " Tensor problem_sizes, Tensor expert_offsets, Tensor sf_offsets) -> ()",
371 + {stride_tag});
372 + ops.impl("cutlass_fp4_group_mm", torch::kCUDA, &cutlass_fp4_group_mm);
373 +
374 // CUTLASS w8a8 GEMM, supporting symmetric per-tensor or per-row/column
375 // quantization, as well as bias
376 ops.def(
@@ -492,6 +500,13 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
500 " Tensor! output_scale, Tensor input_scale) -> ()");
501 ops.impl("scaled_fp4_quant", torch::kCUDA, &scaled_fp4_quant);
502
503 + // Compute NVFP4 experts quantization.
504 + ops.def(
505 + "scaled_fp4_experts_quant(Tensor! output, Tensor! output_scale,"
506 + "Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts,"
507 + "Tensor output_scale_offset_by_experts) -> ()");
508 + ops.impl("scaled_fp4_experts_quant", torch::kCUDA, &scaled_fp4_experts_quant);
509 +
510 // Check if cutlass_scaled_mm_fp4 is supported for CUDA devices
511 // of the given capability
512 ops.def("cutlass_scaled_mm_supports_fp4(int cuda_device_capability) -> bool");
tests/kernels/moe/test_nvfp4_moe.py new
+144
@@ -0,0 +1,144 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +import pytest
3 +import torch
4 +
5 +from tests.kernels.quantization.nvfp4_utils import (FLOAT4_E2M1_MAX,
6 + FLOAT8_E4M3_MAX,
7 + dequantize_nvfp4_to_dtype)
8 +from tests.kernels.utils import torch_moe
9 +from vllm import _custom_ops as ops
10 +from vllm.config import ParallelConfig, VllmConfig, set_current_vllm_config
11 +from vllm.model_executor.layers.fused_moe.cutlass_moe import cutlass_moe_fp4
12 +from vllm.model_executor.layers.fused_moe.fused_moe import fused_topk
13 +from vllm.platforms import current_platform
14 +
15 +if not current_platform.has_device_capability(100):
16 + pytest.skip(reason="Nvfp4 Requires compute capability of 10 or above.",
17 + allow_module_level=True)
18 +
19 +MNK_FACTORS = [
20 + (2, 1024, 1024),
21 + (2, 1024, 1536),
22 + (2, 3072, 1024),
23 + (2, 3072, 1536),
24 + (64, 1024, 1024),
25 + (64, 1024, 1536),
26 + (64, 3072, 1024),
27 + (64, 2048, 1536),
28 + (224, 1024, 1024),
29 + (224, 1024, 1536),
30 +]
31 +
32 +
33 +@pytest.mark.parametrize("m,n,k", MNK_FACTORS)
34 +@pytest.mark.parametrize("e", [40, 64, 256])
35 +@pytest.mark.parametrize("topk", [1, 6, 8])
36 +@pytest.mark.parametrize("dtype", [torch.half, torch.bfloat16])
37 +@torch.inference_mode()
38 +def test_cutlass_fp4_moe_no_graph(m: int, n: int, k: int, e: int, topk: int,
39 + dtype: torch.dtype):
40 + current_platform.seed_everything(7)
41 + with set_current_vllm_config(
42 + VllmConfig(parallel_config=ParallelConfig(
43 + pipeline_parallel_size=1))):
44 +
45 + a = torch.randn((m, k), device="cuda", dtype=dtype) / 10
46 + w1 = torch.randn((e, 2 * n, k), device="cuda", dtype=dtype) / 10
47 + quant_blocksize = 16
48 + round_up = lambda x, y: (x + y - 1) // y * y
49 + sf_w1_2n = round_up(2 * n, 128)
50 + sf_w1_k = round_up(k // quant_blocksize, 4)
51 + w1_blockscale = torch.empty((e, sf_w1_2n, sf_w1_k),
52 + device="cuda",
53 + dtype=torch.float8_e4m3fn)
54 +
55 + w2 = torch.randn((e, k, n), device="cuda", dtype=dtype) / 10
56 + sf_w2_k = round_up(k, 128)
57 + sf_w2_n = round_up(n // quant_blocksize, 4)
58 + w2_blockscale = torch.empty((e, sf_w2_k, sf_w2_n),
59 + device="cuda",
60 + dtype=torch.float8_e4m3fn)
61 +
62 + w1_q = torch.empty((e, 2 * n, k // 2),
63 + device="cuda",
64 + dtype=torch.uint8)
65 + w2_q = torch.empty((e, k, n // 2), device="cuda", dtype=torch.uint8)
66 + w1_gs = torch.empty((e, ), device="cuda", dtype=torch.float32)
67 + w2_gs = torch.empty((e, ), device="cuda", dtype=torch.float32)
68 +
69 + for expert in range(e):
70 + w1_amax = torch.abs(w1).max().to(torch.float32)
71 + w2_amax = torch.abs(w2).max().to(torch.float32)
72 + w1_gs[expert] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w1_amax
73 + w2_gs[expert] = FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX / w2_amax
74 +
75 + w1_q[expert], w1_blockscale[expert] = ops.scaled_fp4_quant(
76 + w1[expert], w1_gs[expert])
77 +
78 + w2_q[expert], w2_blockscale[expert] = ops.scaled_fp4_quant(
79 + w2[expert], w2_gs[expert])
80 +
81 + score = torch.randn((m, e), device="cuda", dtype=dtype)
82 + topk_weights, topk_ids = fused_topk(a, score, topk, renormalize=False)
83 +
84 + a1_gs = torch.ones((e, ), device="cuda", dtype=torch.float32)
85 + a2_gs = torch.ones((e, ), device="cuda", dtype=torch.float32)
86 +
87 + cutlass_output = cutlass_moe_fp4(
88 + a=a,
89 + a1_gscale=a1_gs,
90 + w1_fp4=w1_q,
91 + w1_blockscale=w1_blockscale,
92 + w1_alphas=(1 / w1_gs),
93 + a2_gscale=a2_gs,
94 + w2_fp4=w2_q,
95 + w2_blockscale=w2_blockscale,
96 + w2_alphas=(1 / w2_gs),
97 + topk_weights=topk_weights,
98 + topk_ids=topk_ids,
99 + m=m,
100 + n=n,
101 + k=k,
102 + e=e,
103 + device=a.device,
104 + )
105 +
106 + # Reference check:
107 + a_global_scale = ((FLOAT8_E4M3_MAX * FLOAT4_E2M1_MAX) /
108 + torch.amax(a.flatten(), dim=-1)).to(torch.float32)
109 + a_fp4, a_scale_interleaved = ops.scaled_fp4_quant(a, a_global_scale)
110 + _, m_k = a_fp4.shape
111 + a_in_dtype = dequantize_nvfp4_to_dtype(a_fp4,
112 + a_scale_interleaved,
113 + a_global_scale,
114 + dtype=a.dtype,
115 + device=a.device,
116 + block_size=quant_blocksize)
117 +
118 + w1_d = torch.empty((e, 2 * n, k), device="cuda", dtype=dtype)
119 + w2_d = torch.empty((e, k, n), device="cuda", dtype=dtype)
120 +
121 + for idx in range(0, e):
122 + w1_d[idx] = dequantize_nvfp4_to_dtype(w1_q[idx],
123 + w1_blockscale[idx],
124 + w1_gs[idx],
125 + dtype=w1.dtype,
126 + device=w1.device,
127 + block_size=quant_blocksize)
128 + w2_d[idx] = dequantize_nvfp4_to_dtype(w2_q[idx],
129 + w2_blockscale[idx],
130 + w2_gs[idx],
131 + dtype=w2.dtype,
132 + device=w2.device,
133 + block_size=quant_blocksize)
134 +
135 + torch_output = torch_moe(a_in_dtype, w1_d, w2_d, score, topk, None)
136 +
137 + torch.testing.assert_close(torch_output,
138 + cutlass_output,
139 + atol=1e-1,
140 + rtol=1e-1)
141 +
142 +
143 +if __name__ == "__main__":
144 + test_cutlass_fp4_moe_no_graph((2, 1024, 1024), 40, 1, torch.half)
tests/kernels/quantization/nvfp4_utils.py new
+66
@@ -0,0 +1,66 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +import torch
3 +
4 +from vllm.scalar_type import scalar_types
5 +
6 +FLOAT4_E2M1_MAX = scalar_types.float4_e2m1f.max()
7 +FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
8 +
9 +kE2M1ToFloat = torch.tensor([0., 0.5, 1., 1.5, 2., 3., 4., 6.],
10 + dtype=torch.float32)
11 +
12 +
13 +def convert_swizzled_to_linear(a_sf_swizzled: torch.Tensor, m, k, block_size):
14 + m_tiles = (m + 128 - 1) // 128
15 + f = block_size * 4
16 + k_tiles = (k + f - 1) // f
17 + tmp = torch.reshape(a_sf_swizzled, (1, m_tiles, k_tiles, 32, 4, 4))
18 + tmp = torch.permute(tmp, (0, 1, 4, 3, 2, 5))
19 + out = tmp.reshape(m_tiles * 128, k_tiles * f // block_size)
20 + return out[0:m, 0:k]
21 +
22 +
23 +def dequantize_nvfp4_to_dtype(tensor_fp4,
24 + tensor_sf,
25 + global_scale,
26 + dtype,
27 + device,
28 + block_size=16):
29 + """Dequantize the fp4 tensor back to high precision."""
30 + # Two fp4 values are packed into one uint8.
31 + assert tensor_fp4.dtype == torch.uint8
32 + m, packed_k = tensor_fp4.shape
33 + k = packed_k * 2
34 + tensor_f32 = break_fp4_bytes(tensor_fp4, dtype)
35 + tensor_f32 = tensor_f32.reshape(m, k // block_size, block_size)
36 + tensor_sf = tensor_sf.view(torch.float8_e4m3fn)
37 + tensor_sf = convert_swizzled_to_linear(tensor_sf, m, k, block_size)
38 + tensor_sf_dtype = tensor_sf.to(torch.float32) / global_scale
39 +
40 + # scale the tensor
41 + out = (tensor_f32 * tensor_sf_dtype.unsqueeze(-1)).reshape(m, k)
42 + return out.to(dtype=dtype)
43 +
44 +
45 +def break_fp4_bytes(a, dtype):
46 + assert a.dtype == torch.uint8
47 + m, n = a.shape
48 +
49 + # Vectorized nibble processing
50 + a_flat = a.flatten()
51 + high = (a_flat & 0xF0) >> 4 # Upper nibbles
52 + low = a_flat & 0x0F # Lower nibbles
53 +
54 + # Combine nibbles for batch processing
55 + combined = torch.stack((low, high), dim=1).flatten()
56 +
57 + # Vectorized sign and magnitude extraction
58 + signs = (combined & 0x08).to(torch.bool) # Sign bits
59 + abs_vals = (combined & 0x07).to(torch.long) # Magnitude indices
60 +
61 + # Device-aware lookup and sign application
62 + kE2M1 = kE2M1ToFloat.to(device=a.device)
63 + values = kE2M1[abs_vals] * torch.where(signs, -1.0, 1.0)
64 +
65 + # Reshape to final form
66 + return values.reshape(m, n * 2).to(dtype=dtype)
tests/kernels/quantization/test_nvfp4_scaled_mm.py
+14 -84
@@ -1,10 +1,11 @@
1 # SPDX-License-Identifier: Apache-2.0
2 import pytest
3 import torch
4 +from nvfp4_utils import (FLOAT4_E2M1_MAX, FLOAT8_E4M3_MAX,
5 + dequantize_nvfp4_to_dtype)
6
7 from vllm import _custom_ops as ops
8 from vllm.platforms import current_platform
7 -from vllm.scalar_type import scalar_types
9
10 if not current_platform.has_device_capability(100):
11 pytest.skip(reason="Nvfp4 Requires compute capability of 10 or above.",
@@ -19,95 +20,24 @@ SHAPES.extend(PAD_SHAPES)
20 SEEDS = [42]
21 CUDA_DEVICES = ['cuda:0']
22
22 -FLOAT4_E2M1_MAX = scalar_types.float4_e2m1fn.max()
23 -FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
24 -
25 -kE2M1ToFloatArray = [
26 - 0.,
27 - 0.5,
28 - 1.,
29 - 1.5,
30 - 2.,
31 - 3.,
32 - 4.,
33 - 6.,
34 -]
35 -
36 -
37 -def e2m1_to_fp32(int4_value):
38 - signBit = (int4_value & 0x8)
39 - int4_absValue = int4_value & 0x7
40 - float_result = kE2M1ToFloatArray[int4_absValue]
41 - if (signBit):
42 - float_result = -float_result
43 - return float_result
44 -
45 -
46 -def break_fp4_bytes(a, dtype):
47 - assert (a.dtype == torch.uint8)
48 - m, n = a.shape
49 - a = a.flatten()
50 - # Get upper 4 bits
51 - highHalfByte = (a & 0xF0) >> 4
52 - # Get lower 4 bits
53 - lowHalfByte = a & 0x0F
54 - fH = torch.tensor([e2m1_to_fp32(x) for x in highHalfByte]).to(a.device)
55 - fL = torch.tensor([e2m1_to_fp32(x) for x in lowHalfByte]).to(a.device)
56 - # [0xAB, 0xCD] -> [0xB, 0xA, 0xD, 0xC]
57 - out = torch.stack((fL, fH), dim=-1).reshape(m, n * 2)
58 - return out
59 -
60 -
61 -def convert_swizzled_to_linear(a_sf_swizzled: torch.Tensor, m, k, block_size):
62 - sf_m, sf_k = a_sf_swizzled.shape
63 - m_tiles = (m + 128 - 1) // 128
64 - f = block_size * 4
65 - k_tiles = (k + f - 1) // f
66 - tmp = torch.reshape(a_sf_swizzled, (1, m_tiles, k_tiles, 32, 4, 4))
67 - tmp = torch.permute(tmp, (0, 1, 4, 3, 2, 5))
68 - out = tmp.reshape(m_tiles * 128, k_tiles * f // block_size)
69 - return out[0:m, 0:k]
70 -
71 -
72 -def dequantize_to_dtype(tensor_fp4,
73 - tensor_sf,
74 - global_scale,
75 - dtype,
76 - device,
77 - block_size=16):
78 - """Dequantize the fp4 tensor back to high precision."""
79 - # Two fp4 values are packed into one uint8.
80 - assert tensor_fp4.dtype == torch.uint8
81 - m, packed_k = tensor_fp4.shape
82 - k = packed_k * 2
83 - tensor_f32 = break_fp4_bytes(tensor_fp4, dtype)
84 - tensor_f32 = tensor_f32.reshape(m, k // block_size, block_size)
85 - tensor_sf = tensor_sf.view(torch.float8_e4m3fn)
86 - tensor_sf = convert_swizzled_to_linear(tensor_sf, m, k, block_size)
87 - tensor_sf_dtype = tensor_sf.to(torch.float32) / global_scale
88 -
89 - # scale the tensor
90 - out = (tensor_f32 * tensor_sf_dtype.unsqueeze(-1)).reshape(m, k)
91 - return out
92 -
23
24 def get_ref_results(a_fp4, b_fp4, a_sf, b_sf, a_global_scale, b_global_scale,
25 m, n, dtype, block_size, device):
26 _, m_k = a_fp4.shape
27 _, n_k = b_fp4.shape
28 assert (m_k == n_k)
99 - a_in_dtype = dequantize_to_dtype(a_fp4,
100 - a_sf,
101 - a_global_scale,
102 - dtype=dtype,
103 - device=device,
104 - block_size=block_size)
105 - b_in_dtype = dequantize_to_dtype(b_fp4,
106 - b_sf,
107 - b_global_scale,
108 - dtype=dtype,
109 - device=device,
110 - block_size=block_size)
29 + a_in_dtype = dequantize_nvfp4_to_dtype(a_fp4,
30 + a_sf,
31 + a_global_scale,
32 + dtype=dtype,
33 + device=device,
34 + block_size=block_size)
35 + b_in_dtype = dequantize_nvfp4_to_dtype(b_fp4,
36 + b_sf,
37 + b_global_scale,
38 + dtype=dtype,
39 + device=device,
40 + block_size=block_size)
41 return torch.matmul(a_in_dtype, b_in_dtype.t())
42
43
vllm/_custom_ops.py
+91 -7
@@ -745,10 +745,11 @@ def get_cutlass_moe_mm_data(
745 - output_permutation: Permutation that must be used to shuffle the output
746 after executing the MMs.
747 """
748 - torch.ops._C.get_cutlass_moe_mm_data(topk_ids, expert_offsets,
749 - problem_sizes1, problem_sizes2,
750 - input_permutation, output_permutation,
751 - num_experts, n, k)
748 + return torch.ops._C.get_cutlass_moe_mm_data(topk_ids, expert_offsets,
749 + problem_sizes1, problem_sizes2,
750 + input_permutation,
751 + output_permutation,
752 + num_experts, n, k)
753
754
755 def cutlass_moe_mm(out_tensors: torch.Tensor, a_tensors: torch.Tensor,
@@ -767,9 +768,41 @@ def cutlass_moe_mm(out_tensors: torch.Tensor, a_tensors: torch.Tensor,
768 MMs used in the fused MoE operation.
769 - a/b/c_strides: The data strides passed to grouped matrix multiplication.
770 """
770 - torch.ops._C.cutlass_moe_mm(out_tensors, a_tensors, b_tensors, a_scales,
771 - b_scales, expert_offsets, problem_sizes,
772 - a_strides, b_strides, c_strides)
771 + return torch.ops._C.cutlass_moe_mm(out_tensors, a_tensors, b_tensors,
772 + a_scales, b_scales, expert_offsets,
773 + problem_sizes, a_strides, b_strides,
774 + c_strides)
775 +
776 +
777 +def cutlass_fp4_moe_mm(a_tensors: torch.Tensor, b_tensors: torch.Tensor,
778 + a_scales: torch.Tensor, b_scales: torch.Tensor,
779 + alphas: torch.Tensor, problem_sizes: torch.Tensor,
780 + expert_offsets: torch.Tensor, sf_offsets: torch.Tensor,
781 + out_dtype: torch.dtype, device: torch.device):
782 + """
783 + An FP4 Blockscaled Group Gemm that takes in a_tensors, b_tensors and runs
784 + the gemms for each combination based on the specified problem sizes.
785 +
786 + This is used as the MoE gemm during NVFP4 Quantized FusedMoE forward.
787 + - a/b_tensors: the NVFP4 a_ptrs and b_ptrs tensors which are quantized
788 + input and expert weights.
789 + - a_/b_scales: The blockscales in FP8-E4M3 precision
790 + - expert_offsets/sf_offsets: Indices that mark at which token index
791 + each expert begins its computation. The number of tokens
792 + computed with expert E is expert_offsets[E + 1] -
793 + expert_offsets[E] And the sf_size per expert is
794 + sf_offset[E+1] - sf_offset[E]
795 + - problem_sizes: MxNxK sizes of each expert's multiplication in two grouped
796 + MMs used in the fused MoE operation.
797 + """
798 + m_topk = a_tensors.shape[0]
799 + n = b_tensors.shape[1]
800 + c_shape = (m_topk, n)
801 + c = torch.empty(c_shape, device=device, dtype=out_dtype)
802 + torch.ops._C.cutlass_fp4_group_mm(c, a_tensors, b_tensors, a_scales,
803 + b_scales, alphas, problem_sizes,
804 + expert_offsets, sf_offsets)
805 + return c.to(out_dtype)
806
807
808 # aqlm
@@ -960,6 +993,57 @@ def scaled_fp4_quant(
993 return output, output_scale
994
995
996 +def scaled_fp4_experts_quant(
997 + input_tensor: torch.Tensor,
998 + input_global_scale: torch.Tensor,
999 + expert_offsets: torch.Tensor,
1000 + blockscale_offsets: torch.Tensor,
1001 + topk: int,
1002 + expert_map: Optional[torch.Tensor] = None,
1003 + MAX_TOKENS_PER_EXPERT: int = 163840,
1004 +) -> tuple[torch.Tensor, torch.Tensor]:
1005 + """
1006 + Quantize input tensor to FP4 and return quantized tensor and scale, for
1007 + packed MoE Inputs.
1008 + Args:
1009 + input: The input tensor to be quantized to FP4
1010 + expert_map: The expert map tensor
1011 + input_global_scale: A scalar scaling factor for the entire tensor.
1012 + expert_offsets: The expert offsets tensor
1013 + blockscale_offsets: The blockscale offsets tensor
1014 + Outputs:
1015 + output: The quantized tensor in FP4
1016 + output_scales: The blockscale tensor in FP8-E4M3
1017 + """
1018 + assert not current_platform.is_rocm()
1019 + assert input_tensor.ndim == 2, (
1020 + f'input.ndim needs to be == 2, but got {input_tensor.ndim}.')
1021 +
1022 + input_tensor = input_tensor[
1023 + expert_map] if expert_map is not None else input_tensor
1024 + m_numtopk, k = input_tensor.shape
1025 + assert (m_numtopk <= MAX_TOKENS_PER_EXPERT * topk), (
1026 + f"m_numtopk must be less than MAX_TOKENS_PER_EXPERT * topk for"
1027 + f" scaled_fp4_experts_quant kernel, observed m_numtopk = {m_numtopk}")
1028 + scales_k = k // 16
1029 + padded_k = (scales_k + (4 - 1)) // 4
1030 +
1031 + # output is uint8 and packed fp4 values
1032 + output = torch.empty(m_numtopk,
1033 + k // 2,
1034 + device=input_tensor.device,
1035 + dtype=torch.uint8)
1036 + output_scales = torch.empty(MAX_TOKENS_PER_EXPERT * topk,
1037 + padded_k,
1038 + dtype=torch.int32,
1039 + device=input_tensor.device)
1040 + torch.ops._C.scaled_fp4_experts_quant(output, output_scales, input_tensor,
1041 + input_global_scale, expert_offsets,
1042 + blockscale_offsets)
1043 + output_scales = output_scales.view(torch.float8_e4m3fn)
1044 + return output, output_scales
1045 +
1046 +
1047 # fp8
1048 def scaled_fp8_quant(
1049 input: torch.Tensor,
vllm/model_executor/layers/fused_moe/__init__.py
+2 -1
@@ -36,7 +36,7 @@ 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.cutlass_moe import (
39 - cutlass_moe_fp8)
39 + cutlass_moe_fp4, cutlass_moe_fp8)
40 from vllm.model_executor.layers.fused_moe.fused_moe import (
41 fused_experts, fused_moe, fused_topk, get_config_file_name,
42 grouped_topk)
@@ -48,4 +48,5 @@ if HAS_TRITON:
48 "get_config_file_name",
49 "grouped_topk",
50 "cutlass_moe_fp8",
51 + "cutlass_moe_fp4",
52 ]
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+125 -1
@@ -1,10 +1,11 @@
1 # SPDX-License-Identifier: Apache-2.0
2 -"""Fused MoE kernel."""
2 +""" CUTLASS based Fused MoE kernels."""
3 from typing import Optional
4
5 import torch
6
7 from vllm import _custom_ops as ops
8 +from vllm.scalar_type import scalar_types
9
10
11 #TODO make the grouped gemm kernel consistent with scaled gemm kernel
@@ -178,3 +179,126 @@ def cutlass_moe_fp8(
179 if not apply_router_weight_on_input:
180 c2 = c2 * topk_weights.view(m, topk, 1).to(out_dtype)
181 return c2.sum(dim=1)
182 +
183 +
184 +FLOAT4_E2M1_MAX = scalar_types.float4_e2m1f.max()
185 +FLOAT8_E4M3_MAX = torch.finfo(torch.float8_e4m3fn).max
186 +MAX_TOKENS_PER_EXPERT = 65536
187 +
188 +
189 +def cutlass_moe_fp4(a: torch.Tensor, a1_gscale: torch.Tensor,
190 + w1_fp4: torch.Tensor, w1_blockscale: torch.Tensor,
191 + w1_alphas: torch.Tensor, a2_gscale: torch.Tensor,
192 + w2_fp4: torch.Tensor, w2_blockscale: torch.Tensor,
193 + w2_alphas: torch.Tensor, topk_weights: torch.Tensor,
194 + topk_ids: torch.Tensor, m: int, n: int, k: int, e: int,
195 + device: torch.device):
196 + """
197 + MoE implementation for FP4 Inputs
198 +
199 + # Gemm 1
200 + a: Input tensor: [m, k] (half/bfloat16)
201 + a1_gscale: Activation scale per expert: [e] (float32)
202 + w1(gate up) (not an argument to cutlass_moe_fp4): [e, 2 * n, k]
203 + w1_fp4: [e, 2 * n, k // 2], dtype: torch.uint8 (stacked fp4: E2M1)
204 + (Note: `n` is the up projection output dim, `k` is the input dim in
205 + full precision)
206 + w1_blockscale: [e, 2 * n, k // block_size] (float8_e4m3)
207 + (Block size = 16 for NVFP4)
208 +
209 + # Gemm 2
210 + a2_gscale: Activation scale per expert: [e]
211 + w2(down projection) (not an argument to cutlass_moe_fp4): [e, k, n]
212 + w2_fp4: [e, k, n // 2], dtype: torch.uint8 (stacked E2M1)
213 + w2_blockscale: [e, k, n // block_size], dtype: float8_e4m3
214 +
215 + topk_weights: [m, topk] dtype: float8
216 + topk_ids: [m, topk] dtype: float8
217 +
218 + m, n, k: Unquantized weight shapes, dtype: int
219 + e: number of experts, dtype: int
220 +
221 + assumes that topk < k < n to satisfy - up/down projection expectations.
222 + """
223 + assert topk_weights.shape == topk_ids.shape, "topk shape mismatch"
224 + assert w1_fp4.dtype == torch.uint8, "weight 1 must be uint8"
225 + assert w2_fp4.dtype == torch.uint8, "weight 2 must be uint8"
226 + assert (w1_fp4.ndim == 3 and w2_fp4.ndim == 3 and w1_blockscale.ndim == 3
227 + and w2_blockscale.ndim
228 + == 3), ("All Weights must be of rank 3 for cutlass_moe_fp4")
229 + m_a, k_a = a.shape
230 + e_w1, nx2_w1, half_k_w1 = w1_fp4.shape
231 + e_w2, k_w2, half_n_w2 = w2_fp4.shape
232 +
233 + assert (e_w1 == e_w2 and e_w1 == e), ("Number of experts must match",
234 + " between weights.")
235 + assert (k_a // 2 == half_k_w1
236 + and k == k_w2), ("Hidden size mismatch between a, w1 and w2")
237 + assert (nx2_w1 == n * 2 and half_n_w2 == n // 2), ("mismatch in "
238 + "expected `n`")
239 + assert (m == m_a), "input shape mismatch"
240 + assert 2 * half_k_w1 == k_w2, "Hidden size mismatch w2 and w1"
241 + assert a.dtype in [torch.half, torch.bfloat16], "Invalid input dtype"
242 + assert (topk_weights.shape[0] == m and topk_ids.shape[0]
243 + == m), ("topk must be provided for each row of a")
244 + assert (m <= MAX_TOKENS_PER_EXPERT), (
245 + f"m must be less than MAX_TOKENS_PER_EXPERT({MAX_TOKENS_PER_EXPERT})"
246 + f" for cutlass_moe_fp4, observed m = {m}")
247 + out_dtype = a.dtype
248 + num_topk = topk_ids.shape[1]
249 +
250 + expert_offsets = torch.empty((e + 1), dtype=torch.int32, device=device)
251 + # Problem size: (num_experts, (m,2n,k))
252 + problem_sizes1 = torch.empty((e, 3), dtype=torch.int32, device=device)
253 + # Problem size: (num_experts, (m,n,k))
254 + problem_sizes2 = torch.empty((e, 3), dtype=torch.int32, device=device)
255 +
256 + a_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device)
257 + c_map = torch.empty((topk_ids.numel()), dtype=torch.int32, device=device)
258 +
259 + # problem shapes should have [m, n, k]
260 + # Note that problem sizes are based on logical number of elements.
261 + ops.get_cutlass_moe_mm_data(topk_ids, expert_offsets, problem_sizes1,
262 + problem_sizes2, a_map, c_map, e, n, k)
263 +
264 + tokens_per_expert = problem_sizes1[:, 0]
265 + rounded_tokens_per_expert = (tokens_per_expert + (128 - 1)) // 128 * 128
266 + blockscale_offsets = torch.zeros(e + 1, dtype=torch.int32, device=device)
267 + blockscale_offsets[1:] = torch.cumsum(rounded_tokens_per_expert, dim=0)
268 +
269 + rep_a_fp4, rep_a_blockscale = ops.scaled_fp4_experts_quant(
270 + a,
271 + a1_gscale,
272 + expert_offsets,
273 + blockscale_offsets,
274 + num_topk,
275 + expert_map=a_map,
276 + MAX_TOKENS_PER_EXPERT=MAX_TOKENS_PER_EXPERT)
277 +
278 + c1 = ops.cutlass_fp4_moe_mm(rep_a_fp4, w1_fp4, rep_a_blockscale,
279 + w1_blockscale, w1_alphas, problem_sizes1,
280 + expert_offsets[:-1], blockscale_offsets[:-1],
281 + out_dtype, device)
282 + del rep_a_fp4, rep_a_blockscale
283 + # hidden size dimension is split to one halfpytho sized tensor.
284 + intermediate = torch.empty((m * num_topk, w1_fp4.shape[1] // 2),
285 + device=device,
286 + dtype=out_dtype)
287 +
288 + torch.ops._C.silu_and_mul(intermediate, c1)
289 +
290 + int_fp4, int_blockscale = ops.scaled_fp4_experts_quant(
291 + intermediate,
292 + a2_gscale,
293 + expert_offsets,
294 + blockscale_offsets,
295 + num_topk,
296 + MAX_TOKENS_PER_EXPERT=MAX_TOKENS_PER_EXPERT)
297 +
298 + c2 = ops.cutlass_fp4_moe_mm(int_fp4, w2_fp4, int_blockscale, w2_blockscale,
299 + w2_alphas, problem_sizes2, expert_offsets[:-1],
300 + blockscale_offsets[:-1], out_dtype, device)
301 + del int_fp4, int_blockscale
302 + out = (c2[c_map].view(m, num_topk, k) *
303 + topk_weights.view(m, num_topk, 1).half()).sum(dim=1)
304 + return out.to(dtype=out_dtype)
vllm/model_executor/layers/fused_moe/layer.py
+20 -3
@@ -643,7 +643,7 @@ class FusedMoE(torch.nn.Module):
643 expert_id = self._map_global_expert_id_to_local_expert_id(expert_id)
644 if expert_id == -1:
645 return
646 -
646 + quant_method_name = self.quant_method.__class__.__name__
647 # compressed-tensors checkpoints with packed weights are stored flipped
648 # TODO (mgoin): check self.quant_method.quant_config.quant_format
649 # against known CompressionFormat enum values that have this quality
@@ -697,8 +697,9 @@ class FusedMoE(torch.nn.Module):
697 # this is needed for compressed-tensors only
698 loaded_weight = loaded_weight.to(param.data.device)
699
700 - if param.data[expert_id] != 1 and (param.data[expert_id] -
701 - loaded_weight).abs() > 1e-5:
700 + if ("compressed" in quant_method_name.lower()
701 + and param.data[expert_id] != 1
702 + and (param.data[expert_id] - loaded_weight).abs() > 1e-5):
703 raise ValueError(
704 "input_scales of w1 and w3 of a layer "
705 f"must be equal. But got {param.data[expert_id]} "
@@ -718,6 +719,22 @@ class FusedMoE(torch.nn.Module):
719 tp_rank=self.tp_rank)
720 return
721
722 + if "ModelOpt" in quant_method_name:
723 + if ('weight_scale_2' in weight_name
724 + or 'input_scale' in weight_name):
725 + self._load_per_tensor_weight_scale(shard_id=shard_id,
726 + param=param,
727 + loaded_weight=loaded_weight,
728 + expert_id=expert_id)
729 + elif "weight" in weight_name:
730 + self._load_model_weight_or_group_weight_scale(
731 + shard_id=shard_id,
732 + shard_dim=shard_dim,
733 + loaded_weight=loaded_weight,
734 + expert_data=expert_data,
735 + tp_rank=self.tp_rank)
736 + return
737 +
738 # Case weight scales, zero_points and offset
739 if ("scale" in weight_name or "zero" in weight_name
740 or "offset" in weight_name):
vllm/model_executor/layers/quantization/modelopt.py
+252 -6
@@ -1,6 +1,6 @@
1 # SPDX-License-Identifier: Apache-2.0
2
3 -from typing import Any, Dict, List, Optional, Union
3 +from typing import Any, Callable, Dict, List, Optional, Union
4
5 import torch
6 from torch.nn import Module
@@ -9,6 +9,8 @@ from torch.nn.parameter import Parameter
9 from vllm._custom_ops import (cutlass_scaled_fp4_mm,
10 cutlass_scaled_mm_supports_fp4, scaled_fp4_quant)
11 from vllm.logger import init_logger
12 +from vllm.model_executor.layers.fused_moe.layer import (
13 + FusedMoE, FusedMoEMethodBase, FusedMoeWeightScaleSupported)
14 from vllm.model_executor.layers.linear import (LinearBase, LinearMethodBase,
15 UnquantizedLinearMethod)
16 from vllm.model_executor.layers.quantization import QuantizationMethods
@@ -210,25 +212,37 @@ class ModelOptNvFp4Config(QuantizationConfig):
212 "`hf_quant_config.json` file for your model's "
213 "quant configuration.")
214 is_checkpoint_nvfp4_serialized = ("NVFP4" in quant_method)
213 - kv_cache_quant_algo = quant_config["kv_cache_quant_algo"]
214 - group_size = quant_config["group_size"]
215 - exclude_modules = quant_config["exclude_modules"]
216 - if not (group_size and kv_cache_quant_algo and exclude_modules):
215 + if ("group_size" and "kv_cache_quant_algo"
216 + and "exclude_modules") not in quant_config:
217 raise ValueError("NVFP4 quantization requires group size and "
218 "kv_cache_quant_algo specified in "
219 "hf_quant_config.json")
220 + kv_cache_quant_algo = quant_config["kv_cache_quant_algo"]
221 + group_size = quant_config["group_size"]
222 + exclude_modules = quant_config["exclude_modules"]
223 return cls(is_checkpoint_nvfp4_serialized, kv_cache_quant_algo,
224 exclude_modules, group_size)
225
226 + def is_layer_excluded(self, prefix: str, exclude_modules: List):
227 + import re
228 + for pattern in exclude_modules:
229 + regex_str = pattern.replace('.', r'\.').replace('*', r'.*')
230 + if re.fullmatch(regex_str, prefix):
231 + return True
232 + return False
233 +
234 def get_quant_method(self, layer: torch.nn.Module,
235 prefix: str) -> Optional["QuantizeMethodBase"]:
236 from vllm.attention.layer import Attention # Avoid circular import
237 if isinstance(layer, LinearBase):
227 - if is_layer_skipped(prefix, self.exclude_modules):
238 + if (is_layer_skipped(prefix, self.exclude_modules)
239 + or self.is_layer_excluded(prefix, self.exclude_modules)):
240 return UnquantizedLinearMethod()
241 return ModelOptNvFp4LinearMethod(self)
242 elif isinstance(layer, Attention):
243 return ModelOptFp8KVCacheMethod(self)
244 + elif isinstance(layer, FusedMoE):
245 + return ModelOptNvFp4FusedMoE(self)
246 return None
247
248
@@ -409,3 +423,235 @@ class ModelOptNvFp4LinearMethod(LinearMethodBase):
423 if bias is not None:
424 out = out + bias
425 return out.view(*output_shape)
426 +
427 +
428 +class ModelOptNvFp4FusedMoE(FusedMoEMethodBase):
429 + """
430 + MoE Method for FP4 Quantization.
431 + Args:
432 + quant_config: NVFP4 Quant Config
433 + """
434 +
435 + def __init__(self, quant_config: ModelOptNvFp4Config):
436 + self.quant_config = quant_config
437 +
438 + def create_weights(self, layer: torch.nn.Module, num_experts: int,
439 + hidden_size: int, intermediate_size_per_partition: int,
440 + params_dtype: torch.dtype, **extra_weight_attrs):
441 + if not self.quant_config.is_checkpoint_nvfp4_serialized:
442 + raise ValueError("NVFP4 quantization was selected, "
443 + " dynamic quantization is not supported.")
444 +
445 + layer.quant_config = self.quant_config
446 + weight_dtype = torch.uint8
447 + weight_scale_dtype = torch.float8_e4m3fn
448 + weight_loader = extra_weight_attrs.get("weight_loader")
449 + # GEMM 1
450 + w13_weight = ModelWeightParameter(
451 + data=torch.empty(
452 + num_experts,
453 + 2 * intermediate_size_per_partition,
454 + # 2 fp4 items are packed in the input dimension
455 + hidden_size // 2,
456 + dtype=weight_dtype),
457 + input_dim=1,
458 + output_dim=2,
459 + weight_loader=weight_loader)
460 + layer.register_parameter("w13_weight", w13_weight)
461 +
462 + # GEMM 2
463 + w2_weight = ModelWeightParameter(
464 + data=torch.empty(
465 + num_experts,
466 + hidden_size,
467 + # 2 fp4 items are packed in the input dimension
468 + intermediate_size_per_partition // 2,
469 + dtype=weight_dtype),
470 + input_dim=1,
471 + output_dim=2,
472 + weight_loader=weight_loader)
473 + layer.register_parameter("w2_weight", w2_weight)
474 +
475 + w13_weight_scale = ModelWeightParameter(
476 + data=torch.empty(
477 + num_experts,
478 + 2 * intermediate_size_per_partition,
479 + # 2 fp4 items are packed in the input dimension
480 + hidden_size // self.quant_config.group_size,
481 + dtype=weight_scale_dtype),
482 + input_dim=1,
483 + output_dim=2,
484 + weight_loader=weight_loader)
485 + layer.register_parameter("w13_weight_scale", w13_weight_scale)
486 +
487 + w2_weight_scale = ModelWeightParameter(
488 + data=torch.empty(
489 + num_experts,
490 + hidden_size,
491 + # 2 fp4 items are packed in the input dimension
492 + intermediate_size_per_partition //
493 + self.quant_config.group_size,
494 + dtype=weight_scale_dtype),
495 + input_dim=1,
496 + output_dim=2,
497 + weight_loader=weight_loader)
498 + layer.register_parameter("w2_weight_scale", w2_weight_scale)
499 +
500 + extra_weight_attrs.update(
501 + {"quant_method": FusedMoeWeightScaleSupported.BLOCK.value})
502 +
503 + w13_weight_scale_2 = PerTensorScaleParameter(
504 + data=torch.empty(num_experts, 2, dtype=torch.float32),
505 + weight_loader=weight_loader)
506 + layer.register_parameter("w13_weight_scale_2", w13_weight_scale_2)
507 +
508 + w2_weight_scale_2 = PerTensorScaleParameter(
509 + data=torch.empty(num_experts, dtype=torch.float32),
510 + weight_loader=weight_loader)
511 + layer.register_parameter("w2_weight_scale_2", w2_weight_scale_2)
512 +
513 + extra_weight_attrs.update(
514 + {"quant_method": FusedMoeWeightScaleSupported.TENSOR.value})
515 +
516 + w13_input_scale = PerTensorScaleParameter(data=torch.empty(
517 + num_experts, 2, dtype=torch.float32),
518 + weight_loader=weight_loader)
519 + layer.register_parameter("w13_input_scale", w13_input_scale)
520 +
521 + w2_input_scale = PerTensorScaleParameter(data=torch.empty(
522 + num_experts, dtype=torch.float32),
523 + weight_loader=weight_loader)
524 + layer.register_parameter("w2_input_scale", w2_input_scale)
525 +
526 + def swizzle_blockscale(self, scale: torch.tensor):
527 + assert (scale.dtype == torch.float8_e4m3fn)
528 + # Pad and blockwise interleave weight_scale
529 + scale_ndim = scale.ndim
530 + if scale.ndim == 2:
531 + scale = scale.unsqueeze(0)
532 + assert scale.ndim == 3
533 + B, M, K = scale.shape
534 + round_up_multiple = lambda x, m: (x + m - 1) // m * m
535 + M_padded = round_up_multiple(M, 128)
536 + K_padded = round_up_multiple(K, 4)
537 + padded_scale = torch.zeros((B, M_padded, K_padded), dtype=scale.dtype)
538 + padded_scale[:B, :M, :K] = scale
539 + batches, rows, cols = padded_scale.shape
540 + assert rows % 128 == 0
541 + assert cols % 4 == 0
542 + padded_scale = padded_scale.reshape(batches, rows // 128, 4, 32,
543 + cols // 4, 4)
544 + swizzled_scale = padded_scale.permute((0, 1, 4, 3, 2, 5))
545 + swizzled_scale = swizzled_scale.contiguous().cuda()
546 + return (swizzled_scale.reshape(M, K)
547 + if scale_ndim == 2 else swizzled_scale.reshape(B, M, K))
548 +
549 + def process_weights_after_loading(self, layer: torch.nn.Module) -> None:
550 + # GEMM 1
551 +
552 + assert torch.allclose(
553 + layer.w13_weight_scale_2[:, 0], layer.w13_weight_scale_2[:, 1]), (
554 + "Expected w1_weight_scale_2 to equal w3_weight_scale_2")
555 +
556 + w13_weight_scale_2 = layer.w13_weight_scale_2[:, 0]
557 + layer.w13_weight_scale_2 = Parameter(w13_weight_scale_2,
558 + requires_grad=False)
559 +
560 + w13_input_scale = layer.w13_input_scale.max(dim=1).values.to(
561 + torch.float32)
562 + layer.g1_alphas = Parameter(
563 + (w13_input_scale * w13_weight_scale_2).to(torch.float32),
564 + requires_grad=False)
565 +
566 + assert (layer.w13_weight_scale.shape[2] % 16 == 0), (
567 + "Expected weight_scale.dim(1) to be divisible by 16")
568 + assert (layer.w13_weight_scale.dtype == torch.float8_e4m3fn), (
569 + "Weight Blockscale must be represented as FP8-E4M3")
570 + w13_blockscale_swizzled = self.swizzle_blockscale(
571 + layer.w13_weight_scale)
572 +
573 + layer.w13_blockscale_swizzled = Parameter(w13_blockscale_swizzled,
574 + requires_grad=False)
575 +
576 + # This is for quantization, so we need to invert it.
577 + layer.w13_input_scale_quant = Parameter(
578 + (1 / w13_input_scale).to(torch.float32), requires_grad=False)
579 +
580 + # GEMM 2
581 + layer.g2_alphas = Parameter(
582 + (layer.w2_input_scale * layer.w2_weight_scale_2).to(torch.float32),
583 + requires_grad=False)
584 +
585 + # This is for quantization, so we need to invert it.
586 + layer.w2_input_scale_quant = Parameter(
587 + (1 / layer.w2_input_scale).to(torch.float32), requires_grad=False)
588 +
589 + assert (layer.w2_weight_scale.shape[2] % 16 == 0), (
590 + "Expected weight_scale.dim(1) to be divisible by 16")
591 + assert (layer.w2_weight_scale.dtype == torch.float8_e4m3fn), (
592 + "Weight Blockscale must be represented as FP8-E4M3")
593 + w2_blockscale_swizzled = self.swizzle_blockscale(layer.w2_weight_scale)
594 +
595 + layer.w2_blockscale_swizzled = Parameter(w2_blockscale_swizzled,
596 + requires_grad=False)
597 + return
598 +
599 + def apply(
600 + self,
601 + layer: torch.nn.Module,
602 + x: torch.Tensor,
603 + router_logits: torch.Tensor,
604 + top_k: int,
605 + renormalize: bool,
606 + use_grouped_topk: bool = False,
607 + topk_group: Optional[int] = None,
608 + num_expert_group: Optional[int] = None,
609 + global_num_experts: int = -1,
610 + expert_map: Optional[torch.Tensor] = None,
611 + custom_routing_function: Optional[Callable] = None,
612 + scoring_func: str = "softmax",
613 + e_score_correction_bias: Optional[torch.Tensor] = None,
614 + apply_router_weight_on_input: bool = False,
615 + activation: str = "silu",
616 + ):
617 + assert activation == "silu", "Only SiLU activation is supported."
618 + assert not apply_router_weight_on_input, (
619 + "Router weight on input is not "
620 + "supported for ModelOptNvFp4FusedMoE.")
621 + assert expert_map is None, ("Expert Parallelism /expert_map "
622 + "is currently not supported for "
623 + "ModelOptNvFp4FusedMoE.")
624 +
625 + topk_weights, topk_ids = FusedMoE.select_experts(
626 + hidden_states=x,
627 + router_logits=router_logits,
628 + use_grouped_topk=use_grouped_topk,
629 + top_k=top_k,
630 + renormalize=renormalize,
631 + topk_group=topk_group,
632 + num_expert_group=num_expert_group,
633 + custom_routing_function=custom_routing_function,
634 + scoring_func=scoring_func,
635 + e_score_correction_bias=e_score_correction_bias)
636 +
637 + from vllm.model_executor.layers.fused_moe.cutlass_moe import (
638 + cutlass_moe_fp4)
639 +
640 + # Cutlass moe takes in activations in BF16/Half precision
641 + # and fp4 quantized weights loaded from the checkpoint
642 + return cutlass_moe_fp4(a=x,
643 + w1_fp4=layer.w13_weight,
644 + w1_blockscale=layer.w13_blockscale_swizzled,
645 + w1_alphas=layer.g1_alphas,
646 + w2_fp4=layer.w2_weight,
647 + w2_blockscale=layer.w2_blockscale_swizzled,
648 + w2_alphas=layer.g2_alphas,
649 + topk_weights=topk_weights,
650 + topk_ids=topk_ids,
651 + m=x.shape[0],
652 + n=layer.w2_weight.shape[2] * 2,
653 + k=x.shape[1],
654 + e=layer.w13_weight.shape[0],
655 + a1_gscale=layer.w13_input_scale_quant,
656 + a2_gscale=layer.w2_input_scale_quant,
657 + device=x.device).to(x.dtype)