From 59b2f7b64007c6eae7d8ec82cbf1f33500d1ad2d Mon Sep 17 00:00:00 2001 From: Wei Zhao <51183510+wzhao18@users.noreply.github.com> Date: Sat, 11 Apr 2026 10:16:51 -0400 Subject: [PATCH] [Perf] Fuse Zero Initializer for FP8 DeepGemm Block Quant Kernel (#39547) Signed-off-by: wzhao18 --- .../w8a8/fp8/per_token_group_quant.cu | 119 ++++++++++-------- .../test_per_token_group_quant.py | 110 ++++++++++++++++ 2 files changed, 180 insertions(+), 49 deletions(-) diff --git a/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu b/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu index 69b6564be7..8b0356815b 100644 --- a/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu +++ b/csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu @@ -240,8 +240,9 @@ 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 int num_groups_padded, const int groups_per_block, + const int padded_groups_per_row, const int groups_per_row, const int mn, + const int tma_aligned_mn, const int num_scale_elems, 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; @@ -249,51 +250,62 @@ __global__ void per_token_group_quant_8bit_packed_kernel( 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) { + if (global_group_id >= num_groups_padded) { return; } - const int64_t block_group_offset = global_group_id * group_size; + // map flat group id to 2D indices (mn_idx, sf_k_idx) + const int sf_k_idx = + static_cast(global_group_id % padded_groups_per_row); + const int mn_idx = static_cast(global_group_id / padded_groups_per_row); - const T* group_input = input + block_group_offset; - DST_DTYPE* group_output = - static_cast(output_q) + block_group_offset; + // whether it is a valid group (not padding) + const bool is_valid_group = (mn_idx < mn) && (sf_k_idx < groups_per_row); // 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; - const float y_s = - ComputeGroupScale(group_input, smem_group, group_size, lane_id, - threads_per_group, eps, max_8bit); - // pack 4 scales into a uint32 + // compute scale for valid groups + float y_s = 0.f; + if (is_valid_group) { + const T* group_input = + input + static_cast(mn_idx) * groups_per_row * group_size + + sf_k_idx * group_size; + y_s = ComputeGroupScale(group_input, smem_group, group_size, + lane_id, threads_per_group, eps, max_8bit); + } + + // pack 4 scales into a uint32 exponent 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; + // 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; + const int out_idx = sf_k_pack_idx * tma_aligned_mn + mn_idx; + if (is_valid_group) { // 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 - atomicOr(output_s_packed + out_idx, contrib); + const uint8_t exponent = static_cast((bits >> 23u) & 0xffu); + reinterpret_cast(output_s_packed)[out_idx * 4 + pos] = exponent; + } else if (out_idx < num_scale_elems) { + // write zero for padding groups if within bounds of output_s_packed + reinterpret_cast(output_s_packed)[out_idx * 4 + pos] = 0; } } __syncthreads(); - QuantizeGroup(smem_group, group_output, group_size, lane_id, - threads_per_group, y_s, min_8bit, max_8bit); + if (is_valid_group) { + DST_DTYPE* group_output = + static_cast(output_q) + + static_cast(mn_idx) * groups_per_row * group_size + + sf_k_idx * group_size; + QuantizeGroup(smem_group, group_output, group_size, lane_id, + threads_per_group, y_s, min_8bit, max_8bit); + } } void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input, @@ -310,7 +322,6 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input, 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; STD_TORCH_CHECK(output_s_packed.dim() == 2, "output_s_packed must be 2D, got dim=", output_s_packed.dim(), @@ -330,36 +341,46 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input, "output_s_packed shape must be [", mn, ", ", k_num_packed_sfk, "], but got [", output_s_packed.size(0), ", ", output_s_packed.size(1), "]."); + // Verify column-major TMA-aligned layout + STD_TORCH_CHECK(output_s_packed.stride(0) == 1 && + output_s_packed.stride(1) == tma_aligned_mn, + "output_s_packed must have strides [1, ", tma_aligned_mn, + "], but got [", output_s_packed.stride(0), ", ", + output_s_packed.stride(1), "]."); cudaStream_t stream = get_current_cuda_stream(); constexpr int THREADS_PER_GROUP = 16; - const int groups_per_block = GetGroupsPerBlock(num_groups); + // Expand the grid to cover MN and K padding so every byte in + // output_s_packed is written (padding bytes get zeroed by the kernel). + const int64_t padded_groups_per_row = k_num_packed_sfk * 4; + const int64_t num_groups_padded = tma_aligned_mn * padded_groups_per_row; + // Number of elements in output_s_packed. + const int64_t num_scale_elems = mn + (k_num_packed_sfk - 1) * tma_aligned_mn; + + const int groups_per_block = GetGroupsPerBlock(num_groups_padded); auto dst_type = output_q.scalar_type(); - const int num_blocks = num_groups / groups_per_block; + const int num_blocks = num_groups_padded / 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. - torch::stable::zero_(output_s_packed); - -#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)); \ +#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_padded), \ + groups_per_block, static_cast(padded_groups_per_row), \ + static_cast(groups_per_row), static_cast(mn), \ + static_cast(tma_aligned_mn), \ + static_cast(num_scale_elems), static_cast(eps), \ + static_cast(min_8bit), static_cast(max_8bit)); \ } while (0) VLLM_STABLE_DISPATCH_FLOATING_TYPES( diff --git a/tests/kernels/quantization/test_per_token_group_quant.py b/tests/kernels/quantization/test_per_token_group_quant.py index e3b934722b..f2edd92765 100644 --- a/tests/kernels/quantization/test_per_token_group_quant.py +++ b/tests/kernels/quantization/test_per_token_group_quant.py @@ -48,6 +48,116 @@ def test_per_token_group_quant_fp8( assert torch.allclose(scale, ref_s, atol=0.01, rtol=0.01) +@pytest.mark.parametrize( + "num_tokens,hidden_dim,group_size", + [ + # No padding: mn=4 (mult of 4), groups_per_row=56 (mult of 4) + (4, 7168, 128), + # MN padding only: mn=1, tma_aligned_mn=4 + (1, 7168, 128), + # MN padding only: mn=3, tma_aligned_mn=4 + (3, 7168, 128), + # K padding only: groups_per_row=5 (5%4=1) + (4, 640, 128), + # K padding only: groups_per_row=6 (6%4=2) + (4, 768, 128), + # Single packed column, no padding: k_num_packed=1, mn%4=0 + (4, 384, 128), + # Both MN and K padding + (1, 384, 128), + (3, 640, 128), + # Larger shapes with no padding + (64, 7168, 128), + (128, 14336, 128), + # Larger shapes with padding + (127, 7168, 128), + (253, 640, 128), + # Non-power-of-2 group size + (4, 768, 96), # 768/96=8 groups, no padding + (3, 768, 96), # 768/96=8 groups, MN padding + (4, 480, 96), # 480/96=5 groups, K padding + (1, 480, 96), # both MN and K padding + ], +) +@pytest.mark.parametrize("poisoned_scales", [False, True]) +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available") +def test_per_token_group_quant_fp8_packed( + num_tokens, hidden_dim, group_size, poisoned_scales +): + """Test the packed DeepGEMM quantization kernel against the Triton + reference (row-major, UE8M0 scales).""" + + device = "cuda" + torch.manual_seed(42) + + x = torch.randn((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) * 8 + + mn = num_tokens + groups_per_row = hidden_dim // group_size + k_num_packed = (groups_per_row + 3) // 4 + tma_aligned_mn = ((mn + 3) // 4) * 4 + num_scale_elems = mn + (k_num_packed - 1) * tma_aligned_mn + + if poisoned_scales: + # Call the kernel with poisoned scale buffer to + # ensure padded indices are correctly zeroed. + fp8_dtype = torch.float8_e4m3fn + finfo = torch.finfo(fp8_dtype) + out_q = torch.empty_like(x, dtype=fp8_dtype) + out_s_packed = torch.empty_strided( + (mn, k_num_packed), + (1, tma_aligned_mn), + device=device, + dtype=torch.int32, + ) + torch.as_strided(out_s_packed, (num_scale_elems,), (1,)).fill_(0x7F7F7F7F) + torch.ops._C.per_token_group_fp8_quant_packed( + x, + out_q, + out_s_packed, + group_size, + 1e-10, + finfo.min, + finfo.max, + ) + else: + out_q, out_s_packed = fp8_utils.per_token_group_quant_fp8_packed_for_deepgemm( + x, + group_size=group_size, + use_ue8m0=True, + ) + + # Triton reference (row-major float32 scales, UE8M0) + with patch("vllm.platforms.current_platform.is_cuda", return_value=False): + ref_q, ref_s = fp8_utils.per_token_group_quant_fp8( + x, + group_size, + use_ue8m0=True, + ) + + # Quantized values must match. + assert torch.equal(out_q, ref_q), "Quantized output mismatch" + + # Verify packed scales (valid exponents + padding zeros). + ref_s_flat = ref_s.reshape(mn, groups_per_row) + ref_exponents = (ref_s_flat.view(torch.int32) >> 23) & 0xFF + + expected = torch.zeros(num_scale_elems, dtype=torch.int32, device="cpu") + for row in range(mn): + for g in range(groups_per_row): + pack_col = g // 4 + pos = g % 4 + idx = pack_col * tma_aligned_mn + row + expected[idx] |= int(ref_exponents[row, g].item()) << (pos * 8) + + actual = torch.as_strided(out_s_packed, (num_scale_elems,), (1,)).cpu() + assert torch.equal(actual, expected), ( + f"Packed scale storage mismatch.\n" + f"First diff at index " + f"{(actual != expected).nonzero(as_tuple=True)[0][0].item()}" + ) + + @pytest.mark.parametrize("shape", [(32, 128), (64, 256), (16, 512)]) @pytest.mark.parametrize("group_size", [64, 128]) @pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA not available")