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

[Perf][Kernel] Fused SiLU+Mul+Quant kernel for NVFP4 cutlass_moe (#31832)

Signed-off-by: mgoin <mgoin64@gmail.com> Signed-off-by: Michael Goin <mgoin64@gmail.com>

Michael Goin committed Jan 9, 2026 at 09:40 UTC 34cd32fe302f63270ac90c3ac8ec68b561aac3d7
8 files changed +278 -95
csrc/ops.h
+6
@@ -301,6 +301,12 @@ void scaled_fp4_experts_quant(
301 torch::Tensor const& input_offset_by_experts,
302 torch::Tensor const& output_scale_offset_by_experts);
303
304 +void silu_and_mul_scaled_fp4_experts_quant(
305 + torch::Tensor& output, torch::Tensor& output_scale,
306 + torch::Tensor const& input, torch::Tensor const& input_global_scale,
307 + torch::Tensor const& input_offset_by_experts,
308 + torch::Tensor const& output_scale_offset_by_experts);
309 +
310 void per_token_group_quant_fp8(const torch::Tensor& input,
311 torch::Tensor& output_q, torch::Tensor& output_s,
312 int64_t group_size, double eps, double fp8_min,
csrc/quantization/fp4/activation_nvfp4_quant_fusion_kernels.cu
-31
@@ -31,37 +31,6 @@
31
32 namespace vllm {
33
34 -// silu in float32
35 -__device__ __forceinline__ float silu(float x) {
36 - return __fdividef(x, (1.f + __expf(-x)));
37 -}
38 -
39 -__device__ __forceinline__ float2 silu2(float2 x) {
40 - return make_float2(silu(x.x), silu(x.y));
41 -}
42 -
43 -template <class Type>
44 -__inline__ __device__ PackedVec<Type> compute_silu_mul(PackedVec<Type>& vec,
45 - PackedVec<Type>& vec2) {
46 - PackedVec<Type> result;
47 - using packed_type = typename TypeConverter<Type>::Type;
48 -
49 -#pragma unroll
50 - for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; ++i) {
51 - // silu_mul in float32
52 - if constexpr (std::is_same_v<Type, half>) {
53 - float2 silu_vec = silu2(__half22float2(vec.elts[i]));
54 - result.elts[i] =
55 - __float22half2_rn(__fmul2_rn(silu_vec, __half22float2(vec2.elts[i])));
56 - } else {
57 - float2 silu_vec = silu2(__bfloat1622float2(vec.elts[i]));
58 - result.elts[i] = __float22bfloat162_rn(
59 - __fmul2_rn(silu_vec, __bfloat1622float2(vec2.elts[i])));
60 - }
61 - }
62 - return result;
63 -}
64 -
34 // Use UE4M3 by default.
35 template <class Type, bool UE8M0_SF = false>
36 __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
csrc/quantization/fp4/nvfp4_experts_quant.cu
+136 -57
@@ -31,8 +31,12 @@
31
32 namespace vllm {
33
34 +// NVFP4 quantization kernel for experts (low-latency path).
35 +// When FUSE_SILU_MUL=true, expects input with gate||up layout and fuses
36 +// SiLU(gate)*up before quantization.
37 // Use UE4M3 by default.
35 -template <class Type, bool UE8M0_SF = false, bool SMALL_NUM_EXPERTS = false>
38 +template <class Type, bool FUSE_SILU_MUL = false, bool UE8M0_SF = false,
39 + bool SMALL_NUM_EXPERTS = false>
40 __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
41 cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
42 float const* SFScale, uint32_t* out, uint32_t* SFout,
@@ -50,6 +54,8 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
54
55 int tid = blockIdx.x * blockDim.x + threadIdx.x;
56 int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
57 + // When fusing SiLU+Mul, input has gate || up layout (doubled width)
58 + int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow;
59
60 // Each global thread processes one element
61 for (int globalIdx = tid; globalIdx < numRows * colsPerRow;
@@ -58,13 +64,6 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
64 int rowIdx = globalIdx / colsPerRow;
65 int colIdx = globalIdx % colsPerRow;
66
61 - int64_t inOffset = rowIdx * colsPerRow + colIdx;
62 - PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
63 - // Get the output tensor offset.
64 - // Same as inOffset because 8 elements are packed into one uint32_t.
65 - int64_t outOffset = inOffset;
66 - auto& out_pos = out[outOffset];
67 -
67 // Find index within the experts using different strategies based on expert
68 // count
69 int rowIdx_in_expert = 0;
@@ -111,6 +110,23 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
110 }
111 }
112
113 + // Load input and optionally apply fused SiLU+Mul
114 + int64_t inOffset = rowIdx * inColsPerRow + colIdx;
115 + PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
116 + PackedVec quant_input;
117 + if constexpr (FUSE_SILU_MUL) {
118 + PackedVec in_vec_up =
119 + reinterpret_cast<PackedVec const*>(in)[inOffset + colsPerRow];
120 + quant_input = compute_silu_mul(in_vec, in_vec_up);
121 + } else {
122 + quant_input = in_vec;
123 + }
124 +
125 + // Get the output tensor offset.
126 + // Same as inOffset because 8 elements are packed into one uint32_t.
127 + int64_t outOffset = rowIdx * colsPerRow + colIdx;
128 + auto& out_pos = out[outOffset];
129 +
130 // Get the global scaling factor, which will be applied to the SF.
131 // Note SFScale is the same as next GEMM's alpha, which is
132 // (448.f / (Alpha_A / 6.f)).
@@ -124,12 +140,16 @@ __global__ void __launch_bounds__(512, VLLM_BLOCKS_PER_SM(512))
140 CVT_FP4_NUM_THREADS_PER_SF>(
141 rowIdx_in_expert, colIdx, numKTiles, SFout_in_expert);
142
127 - out_pos = cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(in_vec, SFScaleVal, sf_out);
143 + out_pos =
144 + cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(quant_input, SFScaleVal, sf_out);
145 }
146 }
147
131 -// Kernel for LARGE_M_TOPK = true (large m_topk optimized version)
132 -template <class Type, bool UE8M0_SF = false, bool SMALL_NUM_EXPERTS = false>
148 +// NVFP4 quantization kernel for LARGE_M_TOPK = true (large m_topk optimized
149 +// version). When FUSE_SILU_MUL=true, expects input with gate||up layout and
150 +// fuses SiLU(gate)*up before quantization.
151 +template <class Type, bool FUSE_SILU_MUL = false, bool UE8M0_SF = false,
152 + bool SMALL_NUM_EXPERTS = false>
153 __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
154 cvt_fp16_to_fp4(int32_t numRows, int32_t numCols, Type const* in,
155 float const* SFScale, uint32_t* out, uint32_t* SFout,
@@ -167,6 +187,8 @@ __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
187
188 int tid = blockIdx.x * blockDim.x + threadIdx.x;
189 int colsPerRow = numCols / CVT_FP4_ELTS_PER_THREAD;
190 + // When fusing SiLU+Mul, input has gate || up layout (doubled width)
191 + int inColsPerRow = FUSE_SILU_MUL ? colsPerRow * 2 : colsPerRow;
192
193 // Each global thread processes one element
194 for (int globalIdx = tid; globalIdx < numRows * colsPerRow;
@@ -175,11 +197,6 @@ __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
197 int rowIdx = globalIdx / colsPerRow;
198 int colIdx = globalIdx % colsPerRow;
199
178 - int64_t inOffset = rowIdx * colsPerRow + colIdx;
179 - PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
180 - int64_t outOffset = inOffset;
181 - auto& out_pos = out[outOffset];
182 -
200 // Find expert using binary search for better performance with large m_topk
201 int rowIdx_in_expert = 0;
202 int expert_idx = 0;
@@ -204,6 +221,21 @@ __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
221 }
222 }
223
224 + // Load input and optionally apply fused SiLU+Mul
225 + int64_t inOffset = rowIdx * inColsPerRow + colIdx;
226 + PackedVec in_vec = reinterpret_cast<PackedVec const*>(in)[inOffset];
227 + PackedVec quant_input;
228 + if constexpr (FUSE_SILU_MUL) {
229 + PackedVec in_vec_up =
230 + reinterpret_cast<PackedVec const*>(in)[inOffset + colsPerRow];
231 + quant_input = compute_silu_mul(in_vec, in_vec_up);
232 + } else {
233 + quant_input = in_vec;
234 + }
235 +
236 + int64_t outOffset = rowIdx * colsPerRow + colIdx;
237 + auto& out_pos = out[outOffset];
238 +
239 float const SFScaleVal = SFScale == nullptr ? 1.0f : SFScale[expert_idx];
240
241 uint32_t* SFout_in_expert =
@@ -214,11 +246,12 @@ __global__ void __launch_bounds__(1024, VLLM_BLOCKS_PER_SM(1024))
246 CVT_FP4_NUM_THREADS_PER_SF>(
247 rowIdx_in_expert, colIdx, numKTiles, SFout_in_expert);
248
217 - out_pos = cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(in_vec, SFScaleVal, sf_out);
249 + out_pos =
250 + cvt_warp_fp16_to_fp4<Type, UE8M0_SF>(quant_input, SFScaleVal, sf_out);
251 }
252 }
253
221 -template <typename T>
254 +template <typename T, bool FUSE_SILU_MUL = false>
255 void quant_impl(void* output, void* output_scale, void* input,
256 void* input_global_scale, void* input_offset_by_experts,
257 void* output_scale_offset_by_experts, int m_topk, int k,
@@ -246,7 +279,7 @@ void quant_impl(void* output, void* output_scale, void* input,
279 if (blockRepeat > 1) {
280 size_t shared_mem_size = (n_experts + 1) * sizeof(uint32_t);
281 if (n_experts >= 4) {
249 - cvt_fp16_to_fp4<T, false, false>
282 + cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false, false>
283 <<<grid, block, shared_mem_size, stream>>>(
284 m_topk, k, reinterpret_cast<T*>(input),
285 reinterpret_cast<float*>(input_global_scale),
@@ -256,34 +289,37 @@ void quant_impl(void* output, void* output_scale, void* input,
289 reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
290 n_experts);
291 } else {
259 - cvt_fp16_to_fp4<T, false, true><<<grid, block, shared_mem_size, stream>>>(
260 - m_topk, k, reinterpret_cast<T*>(input),
261 - reinterpret_cast<float*>(input_global_scale),
262 - reinterpret_cast<uint32_t*>(output),
263 - reinterpret_cast<uint32_t*>(output_scale),
264 - reinterpret_cast<uint32_t*>(input_offset_by_experts),
265 - reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
266 - n_experts);
292 + cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false, true>
293 + <<<grid, block, shared_mem_size, stream>>>(
294 + m_topk, k, reinterpret_cast<T*>(input),
295 + reinterpret_cast<float*>(input_global_scale),
296 + reinterpret_cast<uint32_t*>(output),
297 + reinterpret_cast<uint32_t*>(output_scale),
298 + reinterpret_cast<uint32_t*>(input_offset_by_experts),
299 + reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
300 + n_experts);
301 }
302 } else {
303 if (n_experts >= 16) {
270 - cvt_fp16_to_fp4<T, false, false><<<grid, block, 0, stream>>>(
271 - m_topk, k, reinterpret_cast<T*>(input),
272 - reinterpret_cast<float*>(input_global_scale),
273 - reinterpret_cast<uint32_t*>(output),
274 - reinterpret_cast<uint32_t*>(output_scale),
275 - reinterpret_cast<uint32_t*>(input_offset_by_experts),
276 - reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
277 - n_experts, /* bool low_latency */ true);
304 + cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false, false>
305 + <<<grid, block, 0, stream>>>(
306 + m_topk, k, reinterpret_cast<T*>(input),
307 + reinterpret_cast<float*>(input_global_scale),
308 + reinterpret_cast<uint32_t*>(output),
309 + reinterpret_cast<uint32_t*>(output_scale),
310 + reinterpret_cast<uint32_t*>(input_offset_by_experts),
311 + reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
312 + n_experts, /* bool low_latency */ true);
313 } else {
279 - cvt_fp16_to_fp4<T, false, true><<<grid, block, 0, stream>>>(
280 - m_topk, k, reinterpret_cast<T*>(input),
281 - reinterpret_cast<float*>(input_global_scale),
282 - reinterpret_cast<uint32_t*>(output),
283 - reinterpret_cast<uint32_t*>(output_scale),
284 - reinterpret_cast<uint32_t*>(input_offset_by_experts),
285 - reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
286 - n_experts, /* bool low_latency */ true);
314 + cvt_fp16_to_fp4<T, FUSE_SILU_MUL, false, true>
315 + <<<grid, block, 0, stream>>>(
316 + m_topk, k, reinterpret_cast<T*>(input),
317 + reinterpret_cast<float*>(input_global_scale),
318 + reinterpret_cast<uint32_t*>(output),
319 + reinterpret_cast<uint32_t*>(output_scale),
320 + reinterpret_cast<uint32_t*>(input_offset_by_experts),
321 + reinterpret_cast<uint32_t*>(output_scale_offset_by_experts),
322 + n_experts, /* bool low_latency */ true);
323 }
324 }
325 }
@@ -304,19 +340,19 @@ constexpr auto FLOAT = at::ScalarType::Float;
340 constexpr auto INT = at::ScalarType::Int;
341 constexpr auto UINT8 = at::ScalarType::Byte;
342
307 -void scaled_fp4_experts_quant_sm1xxa(
308 - torch::Tensor& output, torch::Tensor& output_scale,
343 +// Common validation for fp4 experts quantization entry points.
344 +static void validate_fp4_experts_quant_inputs(
345 + torch::Tensor const& output, torch::Tensor const& output_scale,
346 torch::Tensor const& input, torch::Tensor const& input_global_scale,
347 torch::Tensor const& input_offset_by_experts,
311 - torch::Tensor const& output_scale_offset_by_experts) {
312 - CHECK_INPUT(output, "output must be a CUDA tensor");
313 - CHECK_INPUT(output_scale, "output_scale must be a CUDA tensor");
314 - CHECK_INPUT(input, "input must be a CUDA tensor");
315 - CHECK_INPUT(input_global_scale, "input_global_scale must be a CUDA tensor");
316 - CHECK_INPUT(input_offset_by_experts,
317 - "input_offset_by_experts must be a CUDA tensor");
318 - CHECK_INPUT(output_scale_offset_by_experts,
319 - "output_scale_offset_by_experts must be a CUDA tensor");
348 + torch::Tensor const& output_scale_offset_by_experts, int64_t m_topk,
349 + int64_t k) {
350 + CHECK_INPUT(output, "output");
351 + CHECK_INPUT(output_scale, "output_scale");
352 + CHECK_INPUT(input, "input");
353 + CHECK_INPUT(input_global_scale, "input_global_scale");
354 + CHECK_INPUT(input_offset_by_experts, "input_offset_by_experts");
355 + CHECK_INPUT(output_scale_offset_by_experts, "output_scale_offset_by_experts");
356
357 TORCH_CHECK(output.dim() == 2);
358 TORCH_CHECK(output_scale.dim() == 2);
@@ -335,8 +371,6 @@ void scaled_fp4_experts_quant_sm1xxa(
371 TORCH_CHECK(output_scale.scalar_type() == INT);
372
373 const int BLOCK_SIZE = 16;
338 - auto m_topk = input.size(0);
339 - auto k = input.size(1);
374 TORCH_CHECK(k % BLOCK_SIZE == 0, "k must be a multiple of 16");
375 auto n_experts = input_global_scale.size(0);
376 TORCH_CHECK(input_offset_by_experts.size(0) == n_experts + 1);
@@ -348,7 +382,21 @@ void scaled_fp4_experts_quant_sm1xxa(
382 int padded_k = (scales_k + (4 - 1)) / 4 * 4;
383 // 4 means 4 fp8 values are packed into one int32
384 TORCH_CHECK(output_scale.size(1) * 4 == padded_k);
385 +}
386 +
387 +void scaled_fp4_experts_quant_sm1xxa(
388 + torch::Tensor& output, torch::Tensor& output_scale,
389 + torch::Tensor const& input, torch::Tensor const& input_global_scale,
390 + torch::Tensor const& input_offset_by_experts,
391 + torch::Tensor const& output_scale_offset_by_experts) {
392 + auto m_topk = input.size(0);
393 + auto k = input.size(1);
394 +
395 + validate_fp4_experts_quant_inputs(output, output_scale, input,
396 + input_global_scale, input_offset_by_experts,
397 + output_scale_offset_by_experts, m_topk, k);
398
399 + auto n_experts = input_global_scale.size(0);
400 const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
401 const cudaStream_t stream =
402 at::cuda::getCurrentCUDAStream(input.get_device());
@@ -356,7 +404,38 @@ void scaled_fp4_experts_quant_sm1xxa(
404 VLLM_DISPATCH_HALF_TYPES(
405 input.scalar_type(), "nvfp4_experts_quant_kernel", [&] {
406 using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
359 - vllm::quant_impl<cuda_type>(
407 + vllm::quant_impl<cuda_type, /*FUSE_SILU_MUL=*/false>(
408 + output.data_ptr(), output_scale.data_ptr(), input.data_ptr(),
409 + input_global_scale.data_ptr(), input_offset_by_experts.data_ptr(),
410 + output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts,
411 + stream);
412 + });
413 +}
414 +
415 +void silu_and_mul_scaled_fp4_experts_quant_sm1xxa(
416 + torch::Tensor& output, torch::Tensor& output_scale,
417 + torch::Tensor const& input, torch::Tensor const& input_global_scale,
418 + torch::Tensor const& input_offset_by_experts,
419 + torch::Tensor const& output_scale_offset_by_experts) {
420 + auto m_topk = input.size(0);
421 + // Input has gate || up layout, so k = input.size(1) / 2
422 + auto k_times_2 = input.size(1);
423 + TORCH_CHECK(k_times_2 % 2 == 0, "input width must be even (gate || up)");
424 + auto k = k_times_2 / 2;
425 +
426 + validate_fp4_experts_quant_inputs(output, output_scale, input,
427 + input_global_scale, input_offset_by_experts,
428 + output_scale_offset_by_experts, m_topk, k);
429 +
430 + auto n_experts = input_global_scale.size(0);
431 + const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
432 + const cudaStream_t stream =
433 + at::cuda::getCurrentCUDAStream(input.get_device());
434 +
435 + VLLM_DISPATCH_HALF_TYPES(
436 + input.scalar_type(), "silu_mul_nvfp4_experts_quant_kernel", [&] {
437 + using cuda_type = vllm::CUDATypeConverter<scalar_t>::Type;
438 + vllm::quant_impl<cuda_type, /*FUSE_SILU_MUL=*/true>(
439 output.data_ptr(), output_scale.data_ptr(), input.data_ptr(),
440 input_global_scale.data_ptr(), input_offset_by_experts.data_ptr(),
441 output_scale_offset_by_experts.data_ptr(), m_topk, k, n_experts,
csrc/quantization/fp4/nvfp4_quant_entry.cu
+24
@@ -41,6 +41,15 @@ void silu_and_mul_nvfp4_quant_sm1xxa(torch::Tensor& output,
41 torch::Tensor& input_sf);
42 #endif
43
44 +#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
45 + (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
46 +void silu_and_mul_scaled_fp4_experts_quant_sm1xxa(
47 + torch::Tensor& output, torch::Tensor& output_scale,
48 + torch::Tensor const& input, torch::Tensor const& input_global_scale,
49 + torch::Tensor const& input_offset_by_experts,
50 + torch::Tensor const& output_scale_offset_by_experts);
51 +#endif
52 +
53 void scaled_fp4_quant(torch::Tensor& output, torch::Tensor const& input,
54 torch::Tensor& output_sf, torch::Tensor const& input_sf) {
55 #if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
@@ -74,3 +83,18 @@ void silu_and_mul_nvfp4_quant(torch::Tensor& output, torch::Tensor& output_sf,
83 TORCH_CHECK_NOT_IMPLEMENTED(
84 false, "No compiled silu_and_mul nvfp4 quantization kernel");
85 }
86 +
87 +void silu_and_mul_scaled_fp4_experts_quant(
88 + torch::Tensor& output, torch::Tensor& output_scale,
89 + torch::Tensor const& input, torch::Tensor const& input_global_scale,
90 + torch::Tensor const& input_offset_by_experts,
91 + torch::Tensor const& output_scale_offset_by_experts) {
92 +#if (defined(ENABLE_NVFP4_SM100) && ENABLE_NVFP4_SM100) || \
93 + (defined(ENABLE_NVFP4_SM120) && ENABLE_NVFP4_SM120)
94 + return silu_and_mul_scaled_fp4_experts_quant_sm1xxa(
95 + output, output_scale, input, input_global_scale, input_offset_by_experts,
96 + output_scale_offset_by_experts);
97 +#endif
98 + TORCH_CHECK_NOT_IMPLEMENTED(
99 + false, "No compiled silu_and_mul nvfp4 experts quantization kernel");
100 +}
csrc/quantization/fp4/nvfp4_utils.cuh
+30
@@ -239,4 +239,34 @@ __device__ uint32_t cvt_warp_fp16_to_fp4(PackedVec<Type>& vec, float SFScaleVal,
239 return e2m1Vec;
240 }
241
242 +// silu in float32
243 +__device__ __forceinline__ float silu(float x) {
244 + return __fdividef(x, (1.f + __expf(-x)));
245 +}
246 +
247 +__device__ __forceinline__ float2 silu2(float2 x) {
248 + return make_float2(silu(x.x), silu(x.y));
249 +}
250 +
251 +template <class Type>
252 +__inline__ __device__ PackedVec<Type> compute_silu_mul(
253 + const PackedVec<Type>& x_vec, const PackedVec<Type>& y_vec) {
254 + PackedVec<Type> result;
255 +
256 +#pragma unroll
257 + for (int i = 0; i < CVT_FP4_ELTS_PER_THREAD / 2; ++i) {
258 + // silu_mul in float32
259 + if constexpr (std::is_same_v<Type, half>) {
260 + float2 silu_vec = silu2(__half22float2(x_vec.elts[i]));
261 + result.elts[i] = __float22half2_rn(
262 + __fmul2_rn(silu_vec, __half22float2(y_vec.elts[i])));
263 + } else {
264 + float2 silu_vec = silu2(__bfloat1622float2(x_vec.elts[i]));
265 + result.elts[i] = __float22bfloat162_rn(
266 + __fmul2_rn(silu_vec, __bfloat1622float2(y_vec.elts[i])));
267 + }
268 + }
269 + return result;
270 +}
271 +
272 } // namespace vllm
csrc/torch_bindings.cpp
+9
@@ -558,6 +558,15 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
558 "Tensor output_scale_offset_by_experts) -> ()");
559 ops.impl("scaled_fp4_experts_quant", torch::kCUDA, &scaled_fp4_experts_quant);
560
561 + // Fused SiLU+Mul+NVFP4 experts quantization.
562 + ops.def(
563 + "silu_and_mul_scaled_fp4_experts_quant(Tensor! output, Tensor! "
564 + "output_scale,"
565 + "Tensor input, Tensor input_global_scale, Tensor input_offset_by_experts,"
566 + "Tensor output_scale_offset_by_experts) -> ()");
567 + ops.impl("silu_and_mul_scaled_fp4_experts_quant", torch::kCUDA,
568 + &silu_and_mul_scaled_fp4_experts_quant);
569 +
570 // Check if cutlass_scaled_mm_fp4 is supported for CUDA devices
571 // of the given capability
572 ops.def("cutlass_scaled_mm_supports_fp4(int cuda_device_capability) -> bool");
vllm/_custom_ops.py
+68 -3
@@ -1606,15 +1606,15 @@ def scaled_fp4_experts_quant(
1606 topk: int,
1607 ) -> tuple[torch.Tensor, torch.Tensor]:
1608 """
1609 - Quantize input tensor to FP4 and return quantized tensor and scale, for
1609 + Quantize input tensor to NVFP4 and return quantized tensor and scale, for
1610 packed MoE Inputs.
1611 Args:
1612 - input_tensor: The input tensor to be quantized to FP4
1612 + input_tensor: The input tensor to be quantized to NVFP4
1613 input_global_scale: A scalar scaling factor for the entire tensor.
1614 expert_offsets: The expert offsets tensor
1615 blockscale_offsets: The blockscale offsets tensor
1616 Outputs:
1617 - output: The quantized tensor in FP4
1617 + output: The quantized tensor in NVFP4
1618 output_scales: The blockscale tensor in FP8-E4M3
1619 """
1620 assert not current_platform.is_rocm()
@@ -1660,6 +1660,71 @@ def scaled_fp4_experts_quant(
1660 return output, output_scales
1661
1662
1663 +def silu_and_mul_scaled_fp4_experts_quant(
1664 + input_tensor: torch.Tensor,
1665 + input_global_scale: torch.Tensor,
1666 + expert_offsets: torch.Tensor,
1667 + blockscale_offsets: torch.Tensor,
1668 + topk: int,
1669 +) -> tuple[torch.Tensor, torch.Tensor]:
1670 + """
1671 + Fused SiLU+Mul+NVFP4 quantization for MoE intermediate activations.
1672 +
1673 + Args:
1674 + input_tensor: The input tensor with gate || up layout [m_topk, k*2]
1675 + input_global_scale: A per-expert scaling factor [n_experts]
1676 + expert_offsets: The expert offsets tensor [n_experts+1]
1677 + blockscale_offsets: The blockscale offsets tensor [n_experts+1]
1678 + topk: Number of top-k experts selected
1679 + Outputs:
1680 + output: The quantized tensor in NVFP4 [m_topk, k/2]
1681 + output_scales: The blockscale tensor in FP8-E4M3
1682 + """
1683 + assert not current_platform.is_rocm()
1684 + assert input_tensor.ndim == 2, (
1685 + f"input.ndim needs to be == 2, but got {input_tensor.ndim}."
1686 + )
1687 +
1688 + # Control the maximum number of tokens per expert supported by the
1689 + # NVFP4 MoE Expert Quantization. This is used to prevent the kernel
1690 + # from running out of memory. This value can also be increased to support
1691 + # larger models.
1692 + MAX_TOKENS_PER_EXPERT = envs.VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE
1693 + m_numtopk, k_times_2 = input_tensor.shape
1694 + assert k_times_2 % 2 == 0, "input width must be even (gate || up layout)"
1695 + k = k_times_2 // 2
1696 +
1697 + assert m_numtopk <= MAX_TOKENS_PER_EXPERT * topk, (
1698 + f"m_numtopk must be less than MAX_TOKENS_PER_EXPERT("
1699 + f"{MAX_TOKENS_PER_EXPERT})"
1700 + f" for cutlass_moe_fp4, observed m_numtopk = {m_numtopk}. Use"
1701 + f" VLLM_MAX_TOKENS_PER_EXPERT_FP4_MOE to set this value."
1702 + )
1703 + scales_k = k // 16
1704 + padded_k = (scales_k + (4 - 1)) // 4
1705 +
1706 + # output is uint8 and packed fp4 values
1707 + output = torch.empty(
1708 + m_numtopk, k // 2, device=input_tensor.device, dtype=torch.uint8
1709 + )
1710 + output_scales = torch.empty(
1711 + MAX_TOKENS_PER_EXPERT * topk,
1712 + padded_k,
1713 + dtype=torch.int32,
1714 + device=input_tensor.device,
1715 + )
1716 + torch.ops._C.silu_and_mul_scaled_fp4_experts_quant(
1717 + output,
1718 + output_scales,
1719 + input_tensor,
1720 + input_global_scale,
1721 + expert_offsets,
1722 + blockscale_offsets,
1723 + )
1724 + output_scales = output_scales.view(torch.float8_e4m3fn)
1725 + return output, output_scales
1726 +
1727 +
1728 # fp8
1729 def scaled_fp8_quant(
1730 input: torch.Tensor,
vllm/model_executor/layers/fused_moe/cutlass_moe.py
+5 -4
@@ -549,7 +549,8 @@ def run_cutlass_moe_fp4(
549 num_topk,
550 )
551 c1 = _resize_cache(workspace13, (m * topk, n * 2))
552 - c2 = _resize_cache(workspace2, (m * topk, n))
552 + # Note: c2 workspace is no longer needed since SiLU is fused with quantization.
553 + # c3 reuses workspace13 after c1 is consumed.
554 c3 = _resize_cache(workspace13, (m * topk, k))
555 ops.cutlass_fp4_moe_mm(
556 c1,
@@ -563,9 +564,9 @@ def run_cutlass_moe_fp4(
564 blockscale_offsets[:-1],
565 )
566 del rep_a_fp4, rep_a_blockscale
566 - torch.ops._C.silu_and_mul(c2, c1)
567 - int_fp4, int_blockscale = ops.scaled_fp4_experts_quant(
568 - c2, a2_gscale, expert_offsets, blockscale_offsets, num_topk
567 + # Fused SiLU+Mul+NVFP4 quantization
568 + int_fp4, int_blockscale = ops.silu_and_mul_scaled_fp4_experts_quant(
569 + c1, a2_gscale, expert_offsets, blockscale_offsets, num_topk
570 )
571
572 ops.cutlass_fp4_moe_mm(