From f550ad8bf7a3983a16aafd25ad18d8b48f4d5218 Mon Sep 17 00:00:00 2001 From: chuan932 Date: Fri, 25 Sep 2026 23:11:42 +0800 Subject: [PATCH] [Kernel][Perf] Register-resident path for per-token-group 8-bit quant (#55330) Signed-off-by: chao.huan Co-authored-by: Misha Goin --- .../benchmark_per_token_group_quant.py | 13 +- .../w8a8/fp8/per_token_group_quant.cu | 219 +++++++++++++++++- .../test_per_token_group_quant.py | 69 +++++- 3 files changed, 295 insertions(+), 6 deletions(-) diff --git a/benchmarks/kernels/benchmark_per_token_group_quant.py b/benchmarks/kernels/benchmark_per_token_group_quant.py index fda3d2a652..afc31104f7 100644 --- a/benchmarks/kernels/benchmark_per_token_group_quant.py +++ b/benchmarks/kernels/benchmark_per_token_group_quant.py @@ -117,7 +117,18 @@ if __name__ == "__main__": args = parse_args() warmup_iters, bench_iters = args.warmup_iters, args.bench_iters - shapes = [(32, 128), (64, 256), (16, 512)] + shapes = [ + (32, 128), + (64, 256), + (16, 512), + # Serving-sized activations, where this op is memory-bound and its + # schedule dominates: a prefill chunk and a decode batch at two + # common hidden sizes. + (1024, 6144), + (8192, 6144), + (1024, 7168), + (8192, 7168), + ] group_sizes = [64, 128] dtypes = ["fp8", "int8"] if args.dtype == "both" else [args.dtype] 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 cfd1fab1ad..a0aee0cd92 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 @@ -169,6 +169,156 @@ __global__ void per_token_group_quant_8bit_kernel( #endif } +// Register-resident variant for the unpacked float32 scale layout, mirroring +// per_token_group_quant_8bit_packed_register_kernel below, which float32-scale +// callers cannot reach because it is hardwired to UE8M0 exponents packed +// 4-per-int32 with TMA-aligned strides. +// +// A group_size-128 group of 16-bit input is 256 B, so it fits in registers and +// the smem staging plus the __syncthreads() that the kernel above needs between +// its absmax and quantize passes are both avoidable. +// +// Requires GROUP_SIZE == THREADS_PER_GROUP * VEC_SIZE, i.e. group_size 128 with +// 16-bit input; CanUseRegisterPath enforces that. +template +__global__ void per_token_group_quant_8bit_register_kernel( + const T* __restrict__ input, void* __restrict__ output_q, + float* __restrict__ output_s, const int64_t num_groups, const float eps, + const float min_8bit, const float max_8bit, const int scale_num_rows, + const int scale_stride) { + constexpr int THREADS_PER_GROUP = 8; + constexpr int VEC_SIZE = 32 / sizeof(T); + static_assert(32 % THREADS_PER_GROUP == 0, + "THREADS_PER_GROUP must divide the warp size so the shuffle " + "mask stays inside one group"); + static_assert(sizeof(DST_DTYPE) == 1, + "the uint4 store packs one byte per element"); + + // CanUseRegisterPath only routes 16-bit inputs here, but + // VLLM_STABLE_DISPATCH_FLOATING_TYPES instantiates every floating type, so + // gate the body on the shape this schedule can actually serve instead of a + // static_assert. + if constexpr (GROUP_SIZE != THREADS_PER_GROUP * VEC_SIZE) { + return; + } else { + const int64_t group_id = + (static_cast(blockIdx.x) * blockDim.x + threadIdx.x) / + THREADS_PER_GROUP; + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaGridDependencySynchronize(); +#endif + + if (group_id >= num_groups) { +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaTriggerProgrammaticLaunchCompletion(); +#endif + return; + } + + const int lane_id = threadIdx.x % THREADS_PER_GROUP; + + // alignas(16) is required for the uint4 reinterpret_cast: the default + // alignof for 16-bit T is 2. + alignas(16) T regs[VEC_SIZE]; + { + const T* group_input = input + group_id * 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]; + } + + float local_absmax = eps; +#pragma unroll + for (int i = 0; i < VEC_SIZE; ++i) { + local_absmax = fmaxf(local_absmax, fabsf(static_cast(regs[i]))); + } + + // Reduce within the 8-lane subgroup that shares this group. +#ifdef USE_ROCM + const int lane_in_wave = threadIdx.x % warpSize; + const unsigned long long mask = 0xFFull << (lane_in_wave & ~7); + local_absmax = + fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 4, 8)); + local_absmax = + fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 2, 8)); + local_absmax = + fmaxf(local_absmax, __shfl_xor_sync(mask, local_absmax, 1, 8)); +#else + const 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)); +#endif + + float y_s = local_absmax / max_8bit; + if constexpr (SCALE_UE8M0) { + // Bit-exact with exp2f(ceilf(log2f(fmaxf(fabsf(y_s), 1e-10f)))): bump the + // exponent whenever the mantissa is non-zero, then clear the mantissa. + y_s = fmaxf(fabsf(y_s), 1e-10f); + const uint32_t bits = __float_as_uint(y_s); + const uint32_t exp_bits = (bits >> 23) & 0xffu; + const uint32_t mant_bits = bits & 0x7fffffu; + y_s = __uint_as_float((exp_bits + (mant_bits != 0u ? 1u : 0u)) << 23); + } + + if (lane_id == 0) { + if constexpr (IS_COLUMN_MAJOR) { + const int64_t row_idx = group_id / scale_num_rows; + const int64_t col_idx = group_id % scale_num_rows; + output_s[col_idx * scale_stride + row_idx] = y_s; + } else { + output_s[group_id] = y_s; + } + } + + // A UE8M0 scale is an exact power of two, so multiplying by the reciprocal + // matches dividing. Otherwise the reciprocal costs an ulp and can move + // values that land exactly on an 8-bit boundary, so divide instead to stay + // bit-exact with the smem kernel. + const float inv_y = 1.0f / y_s; + uint32_t packed[VEC_SIZE / 4] = {}; +#pragma unroll + for (int i = 0; i < VEC_SIZE; ++i) { + const float v = static_cast(regs[i]); + const float scaled = SCALE_UE8M0 ? (v * inv_y) : (v / y_s); + const float q = fminf(fmaxf(scaled, min_8bit), max_8bit); + DST_DTYPE qb = DST_DTYPE(q); + const uint8_t byte = *reinterpret_cast(&qb); + packed[i >> 2] |= static_cast(byte) << ((i & 3) * 8); + } + + DST_DTYPE* group_output = static_cast(output_q) + + group_id * GROUP_SIZE + lane_id * VEC_SIZE; + *reinterpret_cast(group_output) = + make_uint4(packed[0], packed[1], packed[2], packed[3]); + +#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900)) + cudaTriggerProgrammaticLaunchCompletion(); +#endif + } +} + +// The register path needs group_size 128 with 16-bit input so that one group is +// exactly THREADS_PER_GROUP(8) x VEC_SIZE(16) elements, and it loads/stores +// through uint4, so both buffers must be 16 B aligned. +inline bool CanUseRegisterPath(const torch::stable::Tensor& input, + const torch::stable::Tensor& output_q, + int64_t group_size) { + const auto dtype = input.scalar_type(); + if (group_size != 128) return false; + if (dtype != torch::headeronly::ScalarType::Half && + dtype != torch::headeronly::ScalarType::BFloat16) { + return false; + } + if (!input.is_contiguous() || !output_q.is_contiguous()) return false; + const auto in_addr = reinterpret_cast(input.data_ptr()); + const auto out_addr = reinterpret_cast(output_q.data_ptr()); + return (in_addr % 16 == 0) && (out_addr % 16 == 0); +} + inline int GetGroupsPerBlock(int64_t num_groups) { if (num_groups % 16 == 0) { return 16; @@ -235,6 +385,15 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input, const int scale_num_rows = output_s.size(1); const int scale_stride = output_s.stride(1); + const bool use_register_path = + CanUseRegisterPath(input, output_q, group_size); + constexpr int REG_THREADS_PER_GROUP = 8; + constexpr int REG_BLOCK_THREADS = 256; + const int64_t reg_groups_per_block = + REG_BLOCK_THREADS / REG_THREADS_PER_GROUP; + const int reg_num_blocks = static_cast( + (num_groups + reg_groups_per_block - 1) / reg_groups_per_block); + #ifndef USE_ROCM #define LAUNCH_KERNEL_IMPL(T, DST_DTYPE, COL_MAJOR, UE8M0, CHECK_BOUNDS, \ SMEM_BYTES) \ @@ -281,21 +440,72 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input, } \ } while (0) +#ifndef USE_ROCM + #define LAUNCH_REGISTER_INST(T, DST_DTYPE, COL_MAJOR, UE8M0) \ + do { \ + cudaLaunchConfig_t config = {}; \ + config.gridDim = dim3(reg_num_blocks); \ + config.blockDim = dim3(REG_BLOCK_THREADS); \ + config.dynamicSmemBytes = 0; \ + config.stream = stream; \ + cudaLaunchAttribute attrs[1]; \ + attrs[0].id = cudaLaunchAttributeProgrammaticStreamSerialization; \ + attrs[0].val.programmaticStreamSerializationAllowed = 1; \ + config.numAttrs = 1; \ + config.attrs = attrs; \ + cudaLaunchKernelEx( \ + &config, \ + per_token_group_quant_8bit_register_kernel, \ + static_cast(input.data_ptr()), output_q.data_ptr(), \ + static_cast(output_s.data_ptr()), \ + static_cast(num_groups), (float)eps, (float)min_8bit, \ + (float)max_8bit, scale_num_rows, scale_stride); \ + } while (0) +#else + #define LAUNCH_REGISTER_INST(T, DST_DTYPE, COL_MAJOR, UE8M0) \ + do { \ + per_token_group_quant_8bit_register_kernel \ + <<>>( \ + static_cast(input.data_ptr()), output_q.data_ptr(), \ + static_cast(output_s.data_ptr()), \ + static_cast(num_groups), (float)eps, (float)min_8bit, \ + (float)max_8bit, scale_num_rows, scale_stride); \ + } while (0) +#endif + #define LAUNCH_KERNEL(T, DST_DTYPE) \ do { \ size_t smem_bytes = \ static_cast(groups_per_block) * group_size * sizeof(T); \ if (is_column_major) { \ if (scale_ue8m0) { \ - LAUNCH_KERNEL_INST(T, DST_DTYPE, true, true, smem_bytes); \ + if (use_register_path) { \ + LAUNCH_REGISTER_INST(T, DST_DTYPE, true, true); \ + } else { \ + LAUNCH_KERNEL_INST(T, DST_DTYPE, true, true, smem_bytes); \ + } \ } else { \ - LAUNCH_KERNEL_INST(T, DST_DTYPE, true, false, smem_bytes); \ + if (use_register_path) { \ + LAUNCH_REGISTER_INST(T, DST_DTYPE, true, false); \ + } else { \ + LAUNCH_KERNEL_INST(T, DST_DTYPE, true, false, smem_bytes); \ + } \ } \ } else { \ if (scale_ue8m0) { \ - LAUNCH_KERNEL_INST(T, DST_DTYPE, false, true, smem_bytes); \ + if (use_register_path) { \ + LAUNCH_REGISTER_INST(T, DST_DTYPE, false, true); \ + } else { \ + LAUNCH_KERNEL_INST(T, DST_DTYPE, false, true, smem_bytes); \ + } \ } else { \ - LAUNCH_KERNEL_INST(T, DST_DTYPE, false, false, smem_bytes); \ + if (use_register_path) { \ + LAUNCH_REGISTER_INST(T, DST_DTYPE, false, false); \ + } else { \ + LAUNCH_KERNEL_INST(T, DST_DTYPE, false, false, smem_bytes); \ + } \ } \ } \ } while (0) @@ -314,6 +524,7 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input, #undef LAUNCH_KERNEL #undef LAUNCH_KERNEL_INST #undef LAUNCH_KERNEL_IMPL +#undef LAUNCH_REGISTER_INST } // Register-resident fast path for group_size==128. diff --git a/tests/kernels/quantization/test_per_token_group_quant.py b/tests/kernels/quantization/test_per_token_group_quant.py index 6d3c91022e..a9ac45da03 100644 --- a/tests/kernels/quantization/test_per_token_group_quant.py +++ b/tests/kernels/quantization/test_per_token_group_quant.py @@ -26,6 +26,8 @@ from vllm.platforms import current_platform (16, 512), (3, 2688), (3, 5376), + (333, 128), + (512, 6144), ], ) @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32]) @@ -39,7 +41,7 @@ from vllm.platforms import current_platform ) def test_per_token_group_quant_fp8( shape, - dtype, + dtype: torch.dtype, column_major: bool, tma_aligned: bool, scale_ue8m0: bool, @@ -78,6 +80,50 @@ def test_per_token_group_quant_fp8( assert torch.allclose(scale, ref_s, atol=0.01, rtol=0.01) +def _misaligned_clone(x: torch.Tensor) -> torch.Tensor: + """Contiguous copy of ``x`` whose ``data_ptr`` is not 16B aligned.""" + numel = x.numel() + base = torch.empty(numel + 16 // x.element_size(), device=x.device, dtype=x.dtype) + view = base[1 : 1 + numel] + if view.data_ptr() % 16 == 0: + pytest.skip("allocator did not yield a misaligned offset") + return view.view(x.shape).copy_(x) + + +@pytest.mark.parametrize("shape", [(1, 128), (333, 128), (1024, 6144)]) +@pytest.mark.parametrize("column_major", [False, True]) +@pytest.mark.parametrize("scale_ue8m0", [False, True]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="register-resident schedule is CUDA/ROCm only", +) +def test_per_token_group_quant_fp8_register_matches_smem( + shape, column_major: bool, scale_ue8m0: bool, dtype: torch.dtype +): + """group_size=128 with 16-bit input takes a register-resident schedule that + needs 16B-aligned buffers; a misaligned input falls back to the + shared-memory schedule. The two must agree bit-for-bit, which the + tolerance-based comparison against Triton above cannot establish. + """ + device = current_platform.device_type + torch.manual_seed(42) + x = torch.randn(shape, device=device, dtype=dtype) * 8 + + reg_q, reg_s = fp8_utils.per_token_group_quant_fp8( + x, 128, column_major_scales=column_major, use_ue8m0=scale_ue8m0 + ) + smem_q, smem_s = fp8_utils.per_token_group_quant_fp8( + _misaligned_clone(x), + 128, + column_major_scales=column_major, + use_ue8m0=scale_ue8m0, + ) + + assert torch.equal(reg_q.view(torch.uint8), smem_q.view(torch.uint8)) + assert torch.equal(reg_s, smem_s) + + @pytest.mark.parametrize( "num_tokens,hidden_dim,group_size", [ @@ -472,3 +518,24 @@ def test_per_token_group_quant_int8(shape, dtype, group_size: int): assert torch.allclose(out_q.float(), ref_q.float(), atol=0.15, rtol=0.15) assert torch.allclose(scale, ref_s, atol=0.01, rtol=0.01) + + +@pytest.mark.parametrize("shape", [(1, 128), (333, 128), (1024, 6144)]) +@pytest.mark.parametrize("dtype", [torch.bfloat16, torch.float16]) +@pytest.mark.skipif( + not current_platform.is_cuda_alike(), + reason="register-resident schedule is CUDA/ROCm only", +) +def test_per_token_group_quant_int8_register_matches_smem(shape, dtype: torch.dtype): + """int8 output shares the register-resident schedule with fp8, so it needs + the same bit-exactness guard against the shared-memory schedule. + """ + device = current_platform.device_type + torch.manual_seed(42) + x = torch.randn(shape, device=device, dtype=dtype) * 8 + + reg_q, reg_s = int8_utils.per_token_group_quant_int8(x, 128) + smem_q, smem_s = int8_utils.per_token_group_quant_int8(_misaligned_clone(x), 128) + + assert torch.equal(reg_q, smem_q) + assert torch.equal(reg_s, smem_s)