From a3c83ff2fd050bc6392260d6e84ee1150f238f26 Mon Sep 17 00:00:00 2001 From: Yongye Zhu Date: Thu, 30 Apr 2026 21:09:55 -0400 Subject: [PATCH] Faster per-token fp8 group quant packed kernel for blackwell (#41326) Signed-off-by: Yongye Zhu Co-authored-by: Roger Wang --- .../w8a8/fp8/per_token_group_quant.cu | 258 ++++++++++++------ .../w8a8/per_token_group_quant_8bit.h | 10 + .../test_per_token_group_quant.py | 187 ++++++++++++- 3 files changed, 366 insertions(+), 89 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 8b0356815b..5029e39037 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 @@ -236,17 +236,41 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input, #undef LAUNCH_KERNEL } -template -__global__ void per_token_group_quant_8bit_packed_kernel( +// Register-resident fast path for group_size==128. +// +// Each thread holds 16 source elements (32 B = uint4 x 2) in registers across +// the absmax reduce -> scale compute -> quantize pipeline. No shared memory. +// UE8M0 scale extracted via bit math (bit-exact with exp2f(ceilf(log2f))). +// +// Loads two contiguous uint4s (16 B + 16 B = 32 B) per thread; on Blackwell +// nvcc fuses these into a single 256-bit LDG.E.256. +// +// Constraints: GROUP_SIZE % (THREADS_PER_GROUP * VEC_SIZE) == 0; for +// THREADS_PER_GROUP=8 and bf16/fp16 (VEC_SIZE=16), this means GROUP_SIZE=128. +template +__global__ void per_token_group_quant_8bit_packed_register_kernel( const T* __restrict__ input, void* __restrict__ output_q, - unsigned int* __restrict__ output_s_packed, const int group_size, - 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, + unsigned int* __restrict__ output_s_packed, const int64_t num_groups_padded, + const int groups_per_block, const int padded_groups_per_row, + const int groups_per_row, const int mn, const int output_q_mn_extent, + const int tma_aligned_mn, const int64_t 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; - const int lane_id = threadIdx.x % threads_per_group; + static_assert(GROUP_SIZE == 128, "fast path supports GROUP_SIZE==128"); + constexpr int THREADS_PER_GROUP = 8; + constexpr int VEC_SIZE = 32 / sizeof(T); // 16 for bf16/fp16 + static_assert(GROUP_SIZE == THREADS_PER_GROUP * VEC_SIZE, + "GROUP_SIZE must equal THREADS_PER_GROUP * VEC_SIZE"); + // Each group's 8 threads must live in a single warp octet so the + // 0xffu << (threadIdx.x & 24u) shuffle mask selects exactly the lanes + // that share a group. Requires 32 % THREADS_PER_GROUP == 0 and the host + // to launch num_threads as a multiple of THREADS_PER_GROUP (which it does + // via num_threads = groups_per_block * THREADS_PER_GROUP). + static_assert(32 % THREADS_PER_GROUP == 0, + "THREADS_PER_GROUP must divide warp size for the shuffle " + "mask to be valid"); + + const int 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; @@ -254,141 +278,207 @@ __global__ void per_token_group_quant_8bit_packed_kernel( return; } - // 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); - - // 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; - - // compute scale for valid groups - float y_s = 0.f; + // Load 16 input elements (32 B) into registers as two adjacent uint4 + // loads. nvcc keeps these as 2x LDG.E.128 on sm_100; the per-thread cost + // is dominated by HBM bandwidth at large MN, so a fused 256-bit load via + // inline PTX gave no measurable speedup. + // alignas(16) is required so the uint4* reinterpret_cast below is + // well-defined for T == bf16/fp16 (default alignof is 2). + alignas(16) T regs[VEC_SIZE]; + float local_absmax = eps; 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); + input + static_cast(mn_idx) * groups_per_row * GROUP_SIZE + + sf_k_idx * GROUP_SIZE + lane_id * VEC_SIZE; + uint4* dst = reinterpret_cast(®s[0]); + const uint4* src = reinterpret_cast(group_input); + dst[0] = src[0]; + dst[1] = src[1]; +#pragma unroll + for (int i = 0; i < VEC_SIZE; ++i) { + float v = fabsf(static_cast(regs[i])); + local_absmax = fmaxf(local_absmax, v); + } } - // pack 4 scales into a uint32 exponent + // 8-lane subgroup shuffle reduce (octet of the warp). The mask selects the + // 8 lanes within the warp that share a group. + unsigned mask = 0xffu << (threadIdx.x & 24u); + local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 4)); + local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 2)); + local_absmax = fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 1)); + + float y_s = local_absmax / max_8bit; + y_s = fmaxf(y_s, 1e-10f); + uint32_t bits = __float_as_uint(y_s); + uint32_t exp_bits = (bits >> 23) & 0xffu; + uint32_t mant_bits = bits & 0x7fffffu; + uint8_t exp_byte = + static_cast(exp_bits + (mant_bits != 0u ? 1u : 0u)); + + // Lane 0 writes the packed scale byte. if (lane_id == 0) { - // 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 uint8_t exponent = static_cast((bits >> 23u) & 0xffu); - reinterpret_cast(output_s_packed)[out_idx * 4 + pos] = exponent; + reinterpret_cast(output_s_packed)[out_idx * 4 + pos] = exp_byte; } 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(); - - 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); + // For padded mn rows that fall within output_q's allocated extent, write + // a uint4 of zeros to keep the buffer clean for downstream TMA loads. + // Skip writes for sf_k padding (those positions don't exist in output_q). + if (!is_valid_group) { + if (sf_k_idx < groups_per_row && mn_idx >= mn && + mn_idx < output_q_mn_extent) { + DST_DTYPE* group_output = + static_cast(output_q) + + static_cast(mn_idx) * groups_per_row * GROUP_SIZE + + sf_k_idx * GROUP_SIZE + lane_id * VEC_SIZE; + *reinterpret_cast(group_output) = make_uint4(0, 0, 0, 0); + } + return; } + + // Reconstruct y_s as a power-of-2 float and use its reciprocal. + float y_s_q = __uint_as_float(static_cast(exp_byte) << 23); + float inv_y = 1.0f / y_s_q; + + // Quantize and pack into 16 fp8/int8 bytes (= uint4). VEC_SIZE==16 so we + // fill four 32-bit words, four bytes each. + uint32_t packed_lo = 0; + uint32_t packed_lo_hi = 0; + uint32_t packed_hi_lo = 0; + uint32_t packed_hi = 0; +#pragma unroll + for (int i = 0; i < VEC_SIZE; ++i) { + float q = + fminf(fmaxf(static_cast(regs[i]) * inv_y, min_8bit), max_8bit); + DST_DTYPE qb = DST_DTYPE(q); + uint8_t byte = *reinterpret_cast(&qb); + const int shift = (i & 3) * 8; + if (i < 4) { + packed_lo |= static_cast(byte) << shift; + } else if (i < 8) { + packed_lo_hi |= static_cast(byte) << shift; + } else if (i < 12) { + packed_hi_lo |= static_cast(byte) << shift; + } else { + packed_hi |= static_cast(byte) << shift; + } + } + + uint4 packed_out = + make_uint4(packed_lo, packed_lo_hi, packed_hi_lo, packed_hi); + DST_DTYPE* group_output = + static_cast(output_q) + + static_cast(mn_idx) * groups_per_row * GROUP_SIZE + + sf_k_idx * GROUP_SIZE + lane_id * VEC_SIZE; + *reinterpret_cast(group_output) = packed_out; } +// Public entry point: register-resident packed quant kernel. +// Constraints: group_size == 128 and bf16/fp16 input. void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input, torch::stable::Tensor& output_q, torch::stable::Tensor& output_s_packed, int64_t group_size, double eps, double min_8bit, double max_8bit) { + STD_TORCH_CHECK(group_size == 128, + "per_token_group_quant_8bit_packed only supports " + "group_size==128, got ", + group_size, "."); + const auto in_dtype = input.scalar_type(); + STD_TORCH_CHECK( + in_dtype == torch::headeronly::ScalarType::Half || + in_dtype == torch::headeronly::ScalarType::BFloat16, + "per_token_group_quant_8bit_packed only supports bf16/fp16 input."); + STD_TORCH_CHECK(input.is_contiguous()); STD_TORCH_CHECK(output_q.is_contiguous()); const int64_t k = input.size(-1); - STD_TORCH_CHECK(k % group_size == 0, "Last dimension (", k, - ") must be divisible by group_size (", group_size, ")."); + STD_TORCH_CHECK(k % group_size == 0, "input last dim k=", k, + " is not divisible by group_size=", group_size, "."); const int64_t mn = input.numel() / k; const int64_t groups_per_row = k / group_size; - - STD_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; + // output_q may be allocated with extra padded mn rows (e.g., + // (tma_aligned_mn, k)) so the kernel can zero-fill them in-line and the + // caller can use torch.empty instead of torch.zeros. The grid only covers + // up to tma_aligned_mn, so we cap the extent there. + const int64_t output_q_mn_actual = output_q.numel() / k; + STD_TORCH_CHECK(output_q_mn_actual >= mn, + "output_q must have at least mn rows; got ", + output_q_mn_actual, " rows for mn=", mn, "."); + const int64_t output_q_mn_extent = + output_q_mn_actual < tma_aligned_mn ? output_q_mn_actual : tma_aligned_mn; + STD_TORCH_CHECK( output_s_packed.scalar_type() == torch::headeronly::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. + "output_s_packed must be int32 for UE8M0-packed scales."); STD_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), ", ", + "]; 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 strides must be [1, ", tma_aligned_mn, + "]; got [", output_s_packed.stride(0), ", ", output_s_packed.stride(1), "]."); cudaStream_t stream = get_current_cuda_stream(); - constexpr int THREADS_PER_GROUP = 16; - - // 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). + constexpr int THREADS_PER_GROUP = 8; 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_padded / groups_per_block; + const int64_t num_blocks = num_groups_padded / groups_per_block; const int num_threads = groups_per_block * THREADS_PER_GROUP; + // CUDA caps grid.x at 2^31 - 1; this fits any realistic shape but guard + // against pathological inputs. + STD_TORCH_CHECK(num_blocks <= static_cast(INT32_MAX), + "per_token_group_quant_8bit_packed grid too large: ", + num_blocks, " blocks (max ", INT32_MAX, ")."); -#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)); \ +#define LAUNCH_REG_KERNEL(T, DST_DTYPE) \ + do { \ + dim3 grid(static_cast(num_blocks)); \ + dim3 block(num_threads); \ + per_token_group_quant_8bit_packed_register_kernel \ + <<>>( \ + static_cast(input.data_ptr()), output_q.data_ptr(), \ + reinterpret_cast(output_s_packed.data_ptr()), \ + num_groups_padded, groups_per_block, \ + static_cast(padded_groups_per_row), \ + static_cast(groups_per_row), static_cast(mn), \ + static_cast(output_q_mn_extent), \ + static_cast(tma_aligned_mn), num_scale_elems, \ + static_cast(eps), static_cast(min_8bit), \ + static_cast(max_8bit)); \ } while (0) - VLLM_STABLE_DISPATCH_FLOATING_TYPES( - input.scalar_type(), "per_token_group_quant_8bit_packed", ([&] { + VLLM_STABLE_DISPATCH_HALF_TYPES( + input.scalar_type(), "per_token_group_quant_8bit_packed_register", ([&] { if (dst_type == torch::headeronly::ScalarType::Float8_e4m3fn) { - LAUNCH_PACKED_KERNEL(scalar_t, __nv_fp8_e4m3); + LAUNCH_REG_KERNEL(scalar_t, __nv_fp8_e4m3); } else if (dst_type == torch::headeronly::ScalarType::Char) { - LAUNCH_PACKED_KERNEL(scalar_t, int8_t); + LAUNCH_REG_KERNEL(scalar_t, int8_t); } else { STD_TORCH_CHECK( false, @@ -397,7 +487,7 @@ void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input, } })); -#undef LAUNCH_PACKED_KERNEL +#undef LAUNCH_REG_KERNEL } void per_token_group_quant_fp8(const torch::stable::Tensor& input, diff --git a/csrc/libtorch_stable/quantization/w8a8/per_token_group_quant_8bit.h b/csrc/libtorch_stable/quantization/w8a8/per_token_group_quant_8bit.h index d67fd2b336..6630c0dece 100644 --- a/csrc/libtorch_stable/quantization/w8a8/per_token_group_quant_8bit.h +++ b/csrc/libtorch_stable/quantization/w8a8/per_token_group_quant_8bit.h @@ -8,3 +8,13 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input, torch::stable::Tensor& output_s, int64_t group_size, double eps, double min_8bit, double max_8bit, bool scale_ue8m0 = false); + +// Public op: register-resident packed quant for the DeepGEMM Blackwell path. +// Restricted to group_size == 128 and bf16/fp16 input; other configurations +// raise STD_TORCH_CHECK. The legacy shared-memory fallback was removed because +// no production caller (deep_gemm_moe / input_quant_fp8) uses other shapes. +void per_token_group_quant_8bit_packed(const torch::stable::Tensor& input, + torch::stable::Tensor& output_q, + torch::stable::Tensor& output_s_packed, + int64_t group_size, double eps, + double min_8bit, double max_8bit); diff --git a/tests/kernels/quantization/test_per_token_group_quant.py b/tests/kernels/quantization/test_per_token_group_quant.py index 5447a43ffb..4089e9bc46 100644 --- a/tests/kernels/quantization/test_per_token_group_quant.py +++ b/tests/kernels/quantization/test_per_token_group_quant.py @@ -73,11 +73,6 @@ def test_per_token_group_quant_fp8( # 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]) @@ -161,6 +156,188 @@ def test_per_token_group_quant_fp8_packed( ) +@pytest.mark.skipif( + not current_platform.is_cuda(), reason="DeepGEMM not available on this platform" +) +def test_per_token_group_quant_fp8_packed_all_zero(): + """All-zero input must produce well-defined UE8M0 scale bytes via the eps + floor in the kernel's UE8M0 path. Locks down the all-zero behavior before + optimization. + + The CUDA kernel computes: + y_s = eps / fp8_max + y_s = exp2(ceil(log2(fmax(y_s, 1e-10)))) + For all-zero input, eps/fp8_max < 1e-10, so the inner fmax clamps back to + 1e-10, giving exp2(ceil(log2(1e-10))) = exp2(-33) => UE8M0 byte 0x5E (94). + """ + + device = "cuda" + num_tokens, hidden_dim, group_size = 4, 7168, 128 + x = torch.zeros((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) + + out_q, out_s_packed = fp8_utils.per_token_group_quant_fp8_packed_for_deepgemm( + x, + group_size=group_size, + use_ue8m0=True, + ) + + # Quantized values must be all zero. + assert torch.equal( + out_q.view(torch.uint8), + torch.zeros_like(out_q, dtype=torch.uint8), + ), "All-zero input should produce all-zero FP8 output" + + # UE8M0 byte produced by the kernel for all-zero input. + # The kernel's inner fmax(y_s, 1e-10) clamps eps/fp8_max back to 1e-10. + # 1e-10 as float32 has biased exponent 0x5D and a non-zero mantissa, so + # the kernel's bit-twiddle (exp_bits + (mant_bits != 0)) rounds up to + # 0x5E. This matches exp2(ceil(log2(1e-10))) = exp2(-33). + expected_exp_byte = 0x5E + + 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 + + # All valid scale slots must contain the expected packed value. + # Padding slots must be zero. + actual = torch.as_strided(out_s_packed, (num_scale_elems,), (1,)).cpu() + + 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] |= expected_exp_byte << (pos * 8) + + assert torch.equal(actual, expected), "All-zero scale bytes mismatch" + + +@pytest.mark.skipif( + not current_platform.is_cuda(), reason="DeepGEMM not available on this platform" +) +def test_per_token_group_quant_fp8_packed_mantissa_rounds_up(): + """Inputs whose absmax/max_8bit produces a non-power-of-2 force the + mantissa-rounding-up branch (exp_byte += 1). Locks down this behavior + before optimization.""" + + device = "cuda" + num_tokens, hidden_dim, group_size = 4, 7168, 128 + + # Build a tensor whose per-group absmax = 1.5 * fp8_max * 2^k for various k. + # fp8_max = torch.finfo(torch.float8_e4m3fn).max = 448.0. + # Then absmax/fp8_max = 1.5 * 2^k -> non-zero mantissa, triggers ceil + # rounding to 2^(k+1). Use k=0 for simplicity; the bf16 representation of + # 1.5*448=672.0 is exact. + x = torch.full( + (num_tokens, hidden_dim), + 672.0, + device=device, + dtype=torch.bfloat16, + ) + + out_q, out_s_packed = fp8_utils.per_token_group_quant_fp8_packed_for_deepgemm( + x, + group_size=group_size, + use_ue8m0=True, + ) + + 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, + ) + + assert torch.equal(out_q, ref_q), "Quantized output mismatch" + + 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 + + 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), "Scale bytes mismatch" + + +@pytest.mark.parametrize( + "num_tokens,hidden_dim", + [ + (1, 7168), # mn padded 1 -> 4 + (2, 7168), # mn padded 2 -> 4 + (3, 7168), # mn padded 3 -> 4 + (5, 7168), # mn padded 5 -> 8 + (127, 7168), # mn padded 127 -> 128 + (253, 640), # both mn and groups padded + (1, 384), # extreme: 1 group, 1 mn row -> both axes padded + ], +) +@pytest.mark.skipif( + not current_platform.is_cuda(), reason="DeepGEMM not available on this platform" +) +def test_per_token_group_quant_fp8_packed_zero_fills_padded_output_q( + num_tokens, hidden_dim +): + """When output_q is allocated with shape (tma_aligned_mn, k) instead of + (mn, k), the kernel must overwrite the padded mn rows with zeros so + callers can use ``torch.empty`` instead of ``torch.zeros``.""" + + device = "cuda" + group_size = 128 + 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 + + fp8_dtype = torch.float8_e4m3fn + finfo = torch.finfo(fp8_dtype) + # Allocate output_q with the padded mn extent and pre-fill with 0xFF + # so the kernel cannot rely on a clean buffer. + out_q = torch.empty((tma_aligned_mn, hidden_dim), device=device, dtype=fp8_dtype) + out_q.view(torch.uint8).fill_(0xFF) + + out_s_packed = torch.empty_strided( + (mn, k_num_packed), + (1, tma_aligned_mn), + device=device, + dtype=torch.int32, + ) + + torch.ops._C.per_token_group_fp8_quant_packed( + x, out_q, out_s_packed, group_size, 1e-10, finfo.min, finfo.max + ) + + # Live rows must match the Triton reference. + with patch("vllm.platforms.current_platform.is_cuda", return_value=False): + ref_q, _ = fp8_utils.per_token_group_quant_fp8(x, group_size, use_ue8m0=True) + assert torch.equal(out_q[:mn], ref_q), "Live region mismatch" + + # Padded rows must be all-zero; without this, downstream TMA loads would + # see uninitialised data. + if tma_aligned_mn > mn: + padded_bytes = out_q[mn:tma_aligned_mn].view(torch.uint8) + assert padded_bytes.eq(0).all(), ( + f"Padded rows [{mn}, {tma_aligned_mn}) not zeroed; " + f"{padded_bytes.ne(0).sum().item()} non-zero bytes" + ) + + @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")