iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[Kernel][Perf] Register-resident path for per-token-group 8-bit quant (#55330)

Browse Source 
Signed-off-by: chao.huan <chao.huan@nio.com>
Co-authored-by: Misha Goin <mgoin64@gmail.com>
This commit is contained in:
chuan932
2026-09-25 08:11:42 -07:00
committed by GitHub
co-authored by Misha Goin
parent bb35e2587b
commit f550ad8bf7
3 changed files with 295 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu
+215 -4
View File
@@ -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 <typename T, typename DST_DTYPE, bool IS_COLUMN_MAJOR = false,
bool SCALE_UE8M0 = false, int GROUP_SIZE = 128>
__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<int64_t>(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<uint4*>(&regs[0]);
const uint4* src = reinterpret_cast<const uint4*>(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<float>(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<float>(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<uint8_t*>(&qb);
packed[i >> 2] |= static_cast<uint32_t>(byte) << ((i & 3) * 8);
}
DST_DTYPE* group_output = static_cast<DST_DTYPE*>(output_q) +
group_id * GROUP_SIZE + lane_id * VEC_SIZE;
*reinterpret_cast<uint4*>(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<uintptr_t>(input.data_ptr());
const auto out_addr = reinterpret_cast<uintptr_t>(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<int>(
(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<T, DST_DTYPE, COL_MAJOR, \
UE8M0>, \
static_cast<T*>(input.data_ptr()), output_q.data_ptr(), \
static_cast<float*>(output_s.data_ptr()), \
static_cast<int64_t>(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<T, DST_DTYPE, COL_MAJOR, \
UE8M0> \
<<<dim3(reg_num_blocks), dim3(REG_BLOCK_THREADS), 0, stream>>>( \
static_cast<T*>(input.data_ptr()), output_q.data_ptr(), \
static_cast<float*>(output_s.data_ptr()), \
static_cast<int64_t>(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<size_t>(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.