From deecd2ac4a5231aa1d1e2aa1b7426bbe84aae9bf Mon Sep 17 00:00:00 2001 From: yewentao256 Date: Wed, 26 Nov 2025 13:19:08 -0800 Subject: [PATCH 1/3] deepgemm fused layout kernel Signed-off-by: yewentao256 --- csrc/ops.h | 8 + .../w8a8/fp8/per_token_group_quant.cu | 193 ++++++++++++++++++ csrc/torch_bindings.cpp | 9 + .../layers/fused_moe/deep_gemm_moe.py | 9 +- .../layers/quantization/utils/fp8_utils.py | 93 ++++++++- 5 files changed, 307 insertions(+), 5 deletions(-) diff --git a/csrc/ops.h b/csrc/ops.h index 4bb7857b1503..d302f0491326 100644 --- a/csrc/ops.h +++ b/csrc/ops.h @@ -299,6 +299,14 @@ void per_token_group_quant_int8(const torch::Tensor& input, torch::Tensor& output_q, torch::Tensor& output_s, int64_t group_size, double eps, double int8_min, double int8_max); + +// Fused activation quantisation + DeepGEMM-compatible UE8M0-packed scales. +void per_token_group_quant_8bit_packed(const torch::Tensor& input, + torch::Tensor& output_q, + torch::Tensor& output_s_packed, + int64_t group_size, double eps, + double min_8bit, double max_8bit); + #endif void static_scaled_int8_quant(torch::Tensor& out, torch::Tensor const& input, diff --git a/csrc/quantization/w8a8/fp8/per_token_group_quant.cu b/csrc/quantization/w8a8/fp8/per_token_group_quant.cu index e3ab0676b254..16a75d60f584 100644 --- a/csrc/quantization/w8a8/fp8/per_token_group_quant.cu +++ b/csrc/quantization/w8a8/fp8/per_token_group_quant.cu @@ -206,6 +206,199 @@ void per_token_group_quant_8bit(const torch::Tensor& input, #undef LAUNCH_KERNEL } +template +__global__ void per_token_group_quant_8bit_packed_kernel( + const T* __restrict__ input, void* __restrict__ output_q, + unsigned int* __restrict__ output_s_packed, const int group_size, + const int num_groups, const int groups_per_block, const int groups_per_row, + const int mn, const int tma_aligned_mn, const float eps, + const float min_8bit, const float max_8bit) { + const int threads_per_group = 16; + const int64_t local_group_id = threadIdx.x / threads_per_group; + const int lane_id = threadIdx.x % threads_per_group; + + const int64_t block_group_id = blockIdx.x * groups_per_block; + const int64_t global_group_id = block_group_id + local_group_id; + if (global_group_id >= num_groups) { + return; + } + + const int64_t block_group_offset = global_group_id * group_size; + + float local_absmax = eps; + + const T* group_input = input + block_group_offset; + DST_DTYPE* group_output = + static_cast(output_q) + block_group_offset; + + // shared memory to cache each group's data to avoid double DRAM reads. + extern __shared__ __align__(16) char smem_raw[]; + T* smem = reinterpret_cast(smem_raw); + T* smem_group = smem + local_group_id * group_size; + + constexpr int vec_size = 16 / sizeof(T); + using vec_t = vllm::vec_n_t; + + // copy global -> shared & compute absmax + auto scalar_op_cache = [&] __device__(T & dst, const T& src) { + float abs_v = fabsf(static_cast(src)); + local_absmax = fmaxf(local_absmax, abs_v); + dst = src; + }; + + vllm::vectorize_with_alignment( + group_input, // in + smem_group, // out (shared) + group_size, // elements per group + lane_id, // thread id + threads_per_group, // stride in group + scalar_op_cache); // scalar handler + + local_absmax = GroupReduceMax(local_absmax); + + float y_s = local_absmax / max_8bit; + y_s = exp2f(ceilf(log2f(fmaxf(fabsf(y_s), 1e-10f)))); + + // pack 4 scales into a uint32 + if (lane_id == 0) { + // map flat group id to 2D indices (mn_idx, sf_k_idx) + const int sf_k_idx = static_cast(global_group_id % groups_per_row); + const int mn_idx = static_cast(global_group_id / groups_per_row); + + if (mn_idx < mn) { + // each uint32 in output_s_packed stores 4 packed scales + const int sf_k_pack_idx = sf_k_idx / 4; + const int pos = sf_k_idx % 4; + + // reinterpret the UE8M0 scale y_s as IEEE bits and extract its 8-bit + // exponent + unsigned int bits = __float_as_uint(y_s); + unsigned int contrib = 0; + if (pos == 0) { + contrib = (bits >> 23u); + } else if (pos == 1) { + contrib = (bits >> 15u); + } else if (pos == 2) { + contrib = (bits >> 7u); + } else { + contrib = (bits << 1u); + } + + const int out_idx = sf_k_pack_idx * tma_aligned_mn + mn_idx; + // atomically OR 8-bit exponent into the packed scales buffer + atomicOr(output_s_packed + out_idx, contrib); + } + } + + __syncthreads(); + + // quantize shared -> global 8-bit + auto scalar_op_quant = [&] __device__(DST_DTYPE & dst, const T& src) { + float q = fminf(fmaxf(static_cast(src) / y_s, min_8bit), max_8bit); + dst = DST_DTYPE(q); + }; + + vllm::vectorize_with_alignment( + smem_group, // in (shared) + group_output, // out (global quant tensor) + group_size, // elements + lane_id, // tid + threads_per_group, // stride + scalar_op_quant); // scalar handler +} + +void per_token_group_quant_8bit_packed(const torch::Tensor& input, + torch::Tensor& output_q, + torch::Tensor& output_s_packed, + int64_t group_size, double eps, + double min_8bit, double max_8bit) { + TORCH_CHECK(input.is_contiguous()); + TORCH_CHECK(output_q.is_contiguous()); + + const int64_t k = input.size(-1); + TORCH_CHECK(k % group_size == 0, "Last dimension (", k, + ") must be divisible by group_size (", group_size, ")."); + + const int64_t mn = input.numel() / k; + const int64_t groups_per_row = k / group_size; + const int64_t num_groups = mn * groups_per_row; + + TORCH_CHECK(output_s_packed.dim() == 2, + "output_s_packed must be 2D, got dim=", output_s_packed.dim(), + "."); + + const int64_t k_num_packed_sfk = (groups_per_row + 3) / 4; + const int64_t tma_aligned_mn = ((mn + 3) / 4) * 4; + + TORCH_CHECK(output_s_packed.scalar_type() == at::ScalarType::Int, + "output_s_packed must have dtype int32 for UE8M0-packed scales."); + // DeepGEMM expects SFA scales in MN-major form with shape + // [mn, ceil_div(K, 128 * 4)] and TMA-aligned stride on the last + // dimension. + TORCH_CHECK(output_s_packed.size(0) == mn && + output_s_packed.size(1) == k_num_packed_sfk, + "output_s_packed shape must be [", mn, ", ", k_num_packed_sfk, + "], but got [", output_s_packed.size(0), ", ", + output_s_packed.size(1), "]."); + + cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + constexpr int THREADS_PER_GROUP = 16; + + int groups_per_block = 1; + + if (num_groups % 16 == 0) { + groups_per_block = 16; + } else if (num_groups % 8 == 0) { + groups_per_block = 8; + } else if (num_groups % 4 == 0) { + groups_per_block = 4; + } else if (num_groups % 2 == 0) { + groups_per_block = 2; + } + + auto dst_type = output_q.scalar_type(); + const int num_blocks = num_groups / groups_per_block; + const int num_threads = groups_per_block * THREADS_PER_GROUP; + + // zero-initialize packed scales, since we use atomicOr to accumulate + // exponents from different groups. + output_s_packed.zero_(); + +#define LAUNCH_PACKED_KERNEL(T, DST_DTYPE) \ + do { \ + dim3 grid(num_blocks); \ + dim3 block(num_threads); \ + size_t smem_bytes = \ + static_cast(groups_per_block) * group_size * sizeof(T); \ + per_token_group_quant_8bit_packed_kernel \ + <<>>( \ + static_cast(input.data_ptr()), output_q.data_ptr(), \ + reinterpret_cast(output_s_packed.data_ptr()), \ + static_cast(group_size), static_cast(num_groups), \ + groups_per_block, static_cast(groups_per_row), \ + static_cast(mn), static_cast(tma_aligned_mn), \ + static_cast(eps), static_cast(min_8bit), \ + static_cast(max_8bit)); \ + } while (0) + + VLLM_DISPATCH_FLOATING_TYPES( + input.scalar_type(), "per_token_group_quant_8bit_packed", ([&] { + if (dst_type == at::ScalarType::Float8_e4m3fn) { + LAUNCH_PACKED_KERNEL(scalar_t, __nv_fp8_e4m3); + } else if (dst_type == at::ScalarType::Char) { + LAUNCH_PACKED_KERNEL(scalar_t, int8_t); + } else { + TORCH_CHECK( + false, + "per_token_group_quant_8bit_packed only supports FP8/INT8 " + "outputs."); + } + })); + +#undef LAUNCH_PACKED_KERNEL +} + void per_token_group_quant_fp8(const torch::Tensor& input, torch::Tensor& output_q, torch::Tensor& output_s, int64_t group_size, double eps, double fp8_min, diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index e9c96bb8b56c..ca8ad37b77b5 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -610,6 +610,15 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) { ops.impl("per_token_group_fp8_quant", torch::kCUDA, &per_token_group_quant_fp8); + // Compute per-token-group 8-bit quantized tensor and UE8M0-packed, + // TMA-aligned scales for DeepGEMM. + ops.def( + "per_token_group_fp8_quant_packed(Tensor input, Tensor! output_q, " + "Tensor! output_s_packed, int group_size, float eps, float fp8_min, " + "float fp8_max) -> ()"); + ops.impl("per_token_group_fp8_quant_packed", torch::kCUDA, + &per_token_group_quant_8bit_packed); + // Compute per-token-group INT8 quantized tensor and scaling factor. ops.def( "per_token_group_quant_int8(Tensor input, Tensor! output_q, Tensor! " diff --git a/vllm/model_executor/layers/fused_moe/deep_gemm_moe.py b/vllm/model_executor/layers/fused_moe/deep_gemm_moe.py index 86cdd25f2c87..36ba3a035b07 100644 --- a/vllm/model_executor/layers/fused_moe/deep_gemm_moe.py +++ b/vllm/model_executor/layers/fused_moe/deep_gemm_moe.py @@ -24,7 +24,7 @@ ) from vllm.model_executor.layers.fused_moe.utils import _resize_cache from vllm.model_executor.layers.quantization.utils.fp8_utils import ( - per_token_group_quant_fp8, + per_token_group_quant_fp8_packed_for_deepgemm, ) from vllm.utils.deep_gemm import ( get_mk_alignment_for_contiguous_layout, @@ -287,10 +287,11 @@ def apply( self.activation(activation, act_out, mm1_out.view(-1, N)) a2q_scale: torch.Tensor | None = None - a2q, a2q_scale = per_token_group_quant_fp8( - act_out, self.block_shape[1], column_major_scales=True, out_q=quant_out + a2q, a2q_scale = per_token_group_quant_fp8_packed_for_deepgemm( + act_out, + self.block_shape[1], + out_q=quant_out, ) - m_grouped_fp8_gemm_nt_contiguous( (a2q, a2q_scale), (w2, self.w2_scale), mm2_out, expert_ids ) diff --git a/vllm/model_executor/layers/quantization/utils/fp8_utils.py b/vllm/model_executor/layers/quantization/utils/fp8_utils.py index ae63b4a76726..cddc64a92d08 100644 --- a/vllm/model_executor/layers/quantization/utils/fp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/fp8_utils.py @@ -269,7 +269,11 @@ def _run_deepgemm( weight_scale: torch.Tensor, ) -> torch.Tensor: assert self.deepgemm_input_quant_op is not None - q_input, input_scale = self.deepgemm_input_quant_op(input_2d) + q_input, input_scale = per_token_group_quant_fp8_packed_for_deepgemm( + input_2d, + group_size=self.act_quant_group_shape.col, + use_ue8m0=True, + ) output = torch.empty( (q_input.shape[0], weight.shape[0]), dtype=torch.bfloat16, @@ -658,6 +662,93 @@ def per_token_group_quant_fp8( return x_q, x_s +def per_token_group_quant_fp8_packed_for_deepgemm( + x: torch.Tensor, + group_size: int, + eps: float = 1e-10, + use_ue8m0: bool | None = None, + out_q: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """FP8 per-token-group quantization for DeepGEMM. + + This helper: + - Quantizes activations to FP8 (same scheme as `per_token_group_quant_fp8`) + - Writes per-group scales directly into UE8M0-packed, TMA-aligned layout + expected by DeepGEMM kernels (see `smxx_layout.cuh` in DeepGEMM). + + Returns: + (x_q, x_s_packed) + x_q: FP8 activations, same shape as `x`. + x_s_packed: Int32 tensor with logical shape + [mn, ceil(num_groups_per_row / 4)], laid out with + TMA-aligned stride along the packed-K dimension: + - mn = x.numel() // x.shape[-1] + - num_groups_per_row = x.shape[-1] // group_size + """ + if use_ue8m0 is None: + use_ue8m0 = is_deep_gemm_e8m0_used() + # For DeepGEMM UE8M0-packed layout we *require* UE8M0 scales. + assert use_ue8m0, ( + "per_token_group_quant_fp8_packed_for_deepgemm requires UE8M0 scales." + ) + + dtype = current_platform.fp8_dtype() + assert x.shape[-1] % group_size == 0, ( + f"the last dimension of `x` {x.shape[-1]} must be divisible " + f"by `group_size` {group_size}" + ) + assert x.stride(-1) == 1, "`x` groups must be contiguous" + + finfo = torch.finfo(dtype) + fp8_min = finfo.min + fp8_max = finfo.max + + # For the experimental DeepGEMM path we always operate on + # contiguous inputs/outputs; allocate our own FP8 buffer. + + # Compute DeepGEMM-style packed scale tensor shape. + hidden_dim = x.shape[-1] + mn = x.numel() // hidden_dim + num_groups_per_row = hidden_dim // group_size + k_num_packed_sf_k = (num_groups_per_row + 3) // 4 + tma_aligned_mn = ((mn + 3) // 4) * 4 + + # Allocate with the same logical shape and stride as DeepGEMM's + # `get_mn_major_tma_aligned_packed_ue8m0_tensor`, i.e. shape + # [mn, packed_sf_k] and stride [1, tma_aligned_mn]. + x_s_packed = torch.empty_strided( + (mn, k_num_packed_sf_k), + (1, tma_aligned_mn), + device=x.device, + dtype=torch.int32, + ) + + # CUDA kernel path only (DeepGEMM + E8M0 is CUDA-specific). + assert current_platform.is_cuda(), ( + "per_token_group_quant_fp8_packed_for_deepgemm is only valid on CUDA " + "platforms using DeepGEMM." + ) + + # Ensure a contiguous view for the CUDA kernel, independent of how + # `x` is represented inside TorchDynamo / Inductor graphs. + x_contig = x.contiguous() + x_q_local = torch.empty_like(x_contig, device=x.device, dtype=dtype) + + torch.ops._C.per_token_group_fp8_quant_packed( + x_contig, + x_q_local, + x_s_packed, + group_size, + eps, + fp8_min, + fp8_max, + ) + + # Return a tensor with the original logical shape. + x_q = x_q_local.view_as(x) + return x_q, x_s_packed + + @triton.jit def _w8a8_triton_block_scaled_mm( # Pointers to inputs and output From 85dcdf8b57e3fd8c2d316c10a9a850d3898c30a1 Mon Sep 17 00:00:00 2001 From: yewentao256 Date: Wed, 26 Nov 2025 14:16:15 -0800 Subject: [PATCH 2/3] update doc and comments Signed-off-by: yewentao256 --- .../w8a8/fp8/per_token_group_quant.cu | 18 +++----- .../layers/quantization/utils/fp8_utils.py | 41 +++++++------------ 2 files changed, 19 insertions(+), 40 deletions(-) diff --git a/csrc/quantization/w8a8/fp8/per_token_group_quant.cu b/csrc/quantization/w8a8/fp8/per_token_group_quant.cu index 16a75d60f584..f9ac874c4373 100644 --- a/csrc/quantization/w8a8/fp8/per_token_group_quant.cu +++ b/csrc/quantization/w8a8/fp8/per_token_group_quant.cu @@ -270,19 +270,11 @@ __global__ void per_token_group_quant_8bit_packed_kernel( const int sf_k_pack_idx = sf_k_idx / 4; const int pos = sf_k_idx % 4; - // reinterpret the UE8M0 scale y_s as IEEE bits and extract its 8-bit - // exponent - unsigned int bits = __float_as_uint(y_s); - unsigned int contrib = 0; - if (pos == 0) { - contrib = (bits >> 23u); - } else if (pos == 1) { - contrib = (bits >> 15u); - } else if (pos == 2) { - contrib = (bits >> 7u); - } else { - contrib = (bits << 1u); - } + // reinterpret the UE8M0 scale y_s as IEEE bits, extract the 8-bit + // exponent, and place it into the correct byte of the 32-bit word. + const unsigned int bits = __float_as_uint(y_s); + const unsigned int exponent = (bits >> 23u) & 0xffu; + const unsigned int contrib = exponent << (pos * 8u); const int out_idx = sf_k_pack_idx * tma_aligned_mn + mn_idx; // atomically OR 8-bit exponent into the packed scales buffer diff --git a/vllm/model_executor/layers/quantization/utils/fp8_utils.py b/vllm/model_executor/layers/quantization/utils/fp8_utils.py index cddc64a92d08..3dc9dc8241f2 100644 --- a/vllm/model_executor/layers/quantization/utils/fp8_utils.py +++ b/vllm/model_executor/layers/quantization/utils/fp8_utils.py @@ -671,23 +671,16 @@ def per_token_group_quant_fp8_packed_for_deepgemm( ) -> tuple[torch.Tensor, torch.Tensor]: """FP8 per-token-group quantization for DeepGEMM. - This helper: - - Quantizes activations to FP8 (same scheme as `per_token_group_quant_fp8`) - - Writes per-group scales directly into UE8M0-packed, TMA-aligned layout - expected by DeepGEMM kernels (see `smxx_layout.cuh` in DeepGEMM). - Returns: (x_q, x_s_packed) - x_q: FP8 activations, same shape as `x`. - x_s_packed: Int32 tensor with logical shape - [mn, ceil(num_groups_per_row / 4)], laid out with - TMA-aligned stride along the packed-K dimension: - - mn = x.numel() // x.shape[-1] - - num_groups_per_row = x.shape[-1] // group_size + x_q: FP8 activations, same shape as `x`. + x_s_packed: Int32 tensor with logical shape + [mn, ceil(num_groups_per_row / 4)], laid out with + TMA-aligned stride along the packed-K dimension """ if use_ue8m0 is None: use_ue8m0 = is_deep_gemm_e8m0_used() - # For DeepGEMM UE8M0-packed layout we *require* UE8M0 scales. + # for DeepGEMM UE8M0-packed layout we *require* UE8M0 scales. assert use_ue8m0, ( "per_token_group_quant_fp8_packed_for_deepgemm requires UE8M0 scales." ) @@ -700,22 +693,15 @@ def per_token_group_quant_fp8_packed_for_deepgemm( assert x.stride(-1) == 1, "`x` groups must be contiguous" finfo = torch.finfo(dtype) - fp8_min = finfo.min - fp8_max = finfo.max + fp8_min, fp8_max = finfo.min, finfo.max - # For the experimental DeepGEMM path we always operate on - # contiguous inputs/outputs; allocate our own FP8 buffer. - - # Compute DeepGEMM-style packed scale tensor shape. + # compute DeepGEMM-style packed scale tensor shape. hidden_dim = x.shape[-1] mn = x.numel() // hidden_dim num_groups_per_row = hidden_dim // group_size k_num_packed_sf_k = (num_groups_per_row + 3) // 4 tma_aligned_mn = ((mn + 3) // 4) * 4 - # Allocate with the same logical shape and stride as DeepGEMM's - # `get_mn_major_tma_aligned_packed_ue8m0_tensor`, i.e. shape - # [mn, packed_sf_k] and stride [1, tma_aligned_mn]. x_s_packed = torch.empty_strided( (mn, k_num_packed_sf_k), (1, tma_aligned_mn), @@ -729,13 +715,14 @@ def per_token_group_quant_fp8_packed_for_deepgemm( "platforms using DeepGEMM." ) - # Ensure a contiguous view for the CUDA kernel, independent of how - # `x` is represented inside TorchDynamo / Inductor graphs. - x_contig = x.contiguous() - x_q_local = torch.empty_like(x_contig, device=x.device, dtype=dtype) + x_contiguous = x.contiguous() + if out_q is not None: + x_q_local = out_q + else: + x_q_local = torch.empty_like(x_contiguous, device=x.device, dtype=dtype) torch.ops._C.per_token_group_fp8_quant_packed( - x_contig, + x_contiguous, x_q_local, x_s_packed, group_size, @@ -744,7 +731,7 @@ def per_token_group_quant_fp8_packed_for_deepgemm( fp8_max, ) - # Return a tensor with the original logical shape. + # return a tensor with the original logical shape. x_q = x_q_local.view_as(x) return x_q, x_s_packed From 8e690ec6438154c57352a27140559062ec53fd34 Mon Sep 17 00:00:00 2001 From: yewentao256 Date: Tue, 2 Dec 2025 01:05:16 +0000 Subject: [PATCH 3/3] fix deepgemm hopper Signed-off-by: yewentao256 --- .../layers/fused_moe/deep_gemm_moe.py | 29 +++++++++++++++---- 1 file changed, 23 insertions(+), 6 deletions(-) diff --git a/vllm/model_executor/layers/fused_moe/deep_gemm_moe.py b/vllm/model_executor/layers/fused_moe/deep_gemm_moe.py index 36ba3a035b07..801b357aec1a 100644 --- a/vllm/model_executor/layers/fused_moe/deep_gemm_moe.py +++ b/vllm/model_executor/layers/fused_moe/deep_gemm_moe.py @@ -24,9 +24,11 @@ ) from vllm.model_executor.layers.fused_moe.utils import _resize_cache from vllm.model_executor.layers.quantization.utils.fp8_utils import ( + per_token_group_quant_fp8, per_token_group_quant_fp8_packed_for_deepgemm, ) from vllm.utils.deep_gemm import ( + DeepGemmQuantScaleFMT, get_mk_alignment_for_contiguous_layout, m_grouped_fp8_gemm_nt_contiguous, ) @@ -286,12 +288,27 @@ def apply( self.activation(activation, act_out, mm1_out.view(-1, N)) - a2q_scale: torch.Tensor | None = None - a2q, a2q_scale = per_token_group_quant_fp8_packed_for_deepgemm( - act_out, - self.block_shape[1], - out_q=quant_out, - ) + # quantize activations for the second GEMM. + block_k = self.block_shape[1] + scale_fmt = DeepGemmQuantScaleFMT.from_oracle() + + if scale_fmt == DeepGemmQuantScaleFMT.UE8M0: + # Blackwell: packed UE8M0 int32 scales. + a2q, a2q_scale = per_token_group_quant_fp8_packed_for_deepgemm( + act_out, + block_k, + out_q=quant_out, + ) + else: + # Hopper / non-E8M0: float32 scales, with optional E8M0-style ceil. + use_ue8m0 = scale_fmt == DeepGemmQuantScaleFMT.FLOAT32_CEIL_UE8M0 + a2q, a2q_scale = per_token_group_quant_fp8( + act_out, + block_k, + out_q=quant_out, + use_ue8m0=use_ue8m0, + ) + m_grouped_fp8_gemm_nt_contiguous( (a2q, a2q_scale), (w2, self.w2_scale), mm2_out, expert_ids )