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

[Kernel][Perf][Quantization] Fix odd-row performance cliff in per-token-group quantization (#56478)

Browse Source 
Signed-off-by: Ayushman Singh <ayush1399@users.noreply.github.com>
Co-authored-by: Ayushman Singh <ayush1399@users.noreply.github.com>
This commit is contained in:
Ayushman Singh
2026-09-12 20:27:08 -07:00
committed by GitHub
co-authored by Ayushman Singh
parent 533dbe424b
commit c2ea9f1a5f
2 changed files with 131 additions and 64 deletions
Download Patch File Download Diff File Expand all files Collapse all files
csrc/libtorch_stable/quantization/w8a8/fp8/per_token_group_quant.cu
+88 -58
View File
@@ -96,7 +96,8 @@ __device__ __forceinline__ void QuantizeGroup(
}
template <typename T, typename DST_DTYPE, bool IS_COLUMN_MAJOR = false,
bool SCALE_UE8M0 = false, typename scale_packed_t = float>
bool SCALE_UE8M0 = false, typename scale_packed_t = float,
bool CHECK_BOUNDS = false>
__global__ void per_token_group_quant_8bit_kernel(
const T* __restrict__ input, void* __restrict__ output_q,
scale_packed_t* __restrict__ output_s, const int group_size,
@@ -114,49 +115,54 @@ __global__ void per_token_group_quant_8bit_kernel(
using scale_element_t = float;
static_assert(sizeof(scale_packed_t) % sizeof(scale_element_t) == 0);
const T* group_input = input + block_group_offset;
DST_DTYPE* group_output =
static_cast<DST_DTYPE*>(output_q) + block_group_offset;
scale_element_t* scale_output;
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
cudaGridDependencySynchronize();
#endif
if constexpr (IS_COLUMN_MAJOR) {
const int num_elems_per_pack =
static_cast<int>(sizeof(scale_packed_t) / sizeof(scale_element_t));
const int scale_num_rows_element = scale_num_rows * num_elems_per_pack;
const int row_idx = global_group_id / scale_num_rows_element;
const int col_idx_raw = global_group_id % scale_num_rows_element;
const int col_idx = col_idx_raw / num_elems_per_pack;
const int pack_idx = col_idx_raw % num_elems_per_pack;
scale_output = reinterpret_cast<scale_element_t*>(output_s) +
(col_idx * scale_stride * num_elems_per_pack +
row_idx * num_elems_per_pack + pack_idx);
} else {
scale_output = output_s + global_group_id;
}
// shared memory to cache each group's data to avoid double DRAM reads.
extern __shared__ __align__(16) char smem_raw[];
T* smem = reinterpret_cast<T*>(smem_raw);
T* smem_group = smem + local_group_id * group_size;
const bool valid_group = !CHECK_BOUNDS || global_group_id < num_groups;
float y_s = 1.0f;
if (valid_group) {
const T* group_input = input + block_group_offset;
scale_element_t* scale_output;
const float y_s = ComputeGroupScale<T, SCALE_UE8M0>(
group_input, smem_group, group_size, lane_id, threads_per_group, eps,
max_8bit);
if constexpr (IS_COLUMN_MAJOR) {
const int num_elems_per_pack =
static_cast<int>(sizeof(scale_packed_t) / sizeof(scale_element_t));
const int scale_num_rows_element = scale_num_rows * num_elems_per_pack;
const int row_idx = global_group_id / scale_num_rows_element;
const int col_idx_raw = global_group_id % scale_num_rows_element;
const int col_idx = col_idx_raw / num_elems_per_pack;
const int pack_idx = col_idx_raw % num_elems_per_pack;
scale_output = reinterpret_cast<scale_element_t*>(output_s) +
(col_idx * scale_stride * num_elems_per_pack +
row_idx * num_elems_per_pack + pack_idx);
} else {
scale_output = output_s + global_group_id;
}
scale_element_t y_s_quant = y_s;
y_s = ComputeGroupScale<T, SCALE_UE8M0>(group_input, smem_group, group_size,
lane_id, threads_per_group, eps,
max_8bit);
if (lane_id == 0) {
*scale_output = y_s_quant;
scale_element_t y_s_quant = y_s;
if (lane_id == 0) {
*scale_output = y_s_quant;
}
}
__syncthreads();
QuantizeGroup<T, DST_DTYPE>(smem_group, group_output, group_size, lane_id,
threads_per_group, y_s, min_8bit, max_8bit);
if (valid_group) {
DST_DTYPE* group_output =
static_cast<DST_DTYPE*>(output_q) + block_group_offset;
QuantizeGroup<T, DST_DTYPE>(smem_group, group_output, group_size, lane_id,
threads_per_group, y_s, min_8bit, max_8bit);
}
#if (defined(__CUDA_ARCH__) && (__CUDA_ARCH__ >= 900))
cudaTriggerProgrammaticLaunchCompletion();
@@ -209,10 +215,20 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
constexpr int THREADS_PER_GROUP = 16;
const int groups_per_block = GetGroupsPerBlock(num_groups);
if (num_groups == 0) {
return;
}
STD_TORCH_CHECK(input.size(-1) % group_size == 0);
const int64_t groups_per_row = input.size(-1) / group_size;
const int64_t num_rows = num_groups / groups_per_row;
// Odd row counts inherit the next even row count's launch geometry.
const int64_t policy_num_groups =
(num_rows + (num_rows & 1)) * groups_per_row;
const int groups_per_block = GetGroupsPerBlock(policy_num_groups);
const bool needs_bounds_check = num_groups % groups_per_block != 0;
auto dst_type = output_q.scalar_type();
const int num_blocks = num_groups / groups_per_block;
const int num_blocks = (num_groups + groups_per_block - 1) / groups_per_block;
const int num_threads = groups_per_block * THREADS_PER_GROUP;
const bool is_column_major = output_s.stride(0) < output_s.stride(1);
@@ -220,38 +236,51 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
const int scale_stride = output_s.stride(1);
#ifndef USE_ROCM
#define LAUNCH_KERNEL_INST(T, DST_DTYPE, COL_MAJOR, UE8M0, SMEM_BYTES) \
do { \
cudaLaunchConfig_t config = {}; \
config.gridDim = dim3(num_blocks); \
config.blockDim = dim3(num_threads); \
config.dynamicSmemBytes = (SMEM_BYTES); \
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_kernel<T, DST_DTYPE, COL_MAJOR, UE8M0>, \
static_cast<T*>(input.data_ptr()), output_q.data_ptr(), \
static_cast<float*>(output_s.data_ptr()), group_size, num_groups, \
groups_per_block, (float)eps, (float)min_8bit, (float)max_8bit, \
scale_num_rows, scale_stride); \
#define LAUNCH_KERNEL_IMPL(T, DST_DTYPE, COL_MAJOR, UE8M0, CHECK_BOUNDS, \
SMEM_BYTES) \
do { \
cudaLaunchConfig_t config = {}; \
config.gridDim = dim3(num_blocks); \
config.blockDim = dim3(num_threads); \
config.dynamicSmemBytes = (SMEM_BYTES); \
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_kernel<T, DST_DTYPE, COL_MAJOR, UE8M0, \
float, CHECK_BOUNDS>, \
static_cast<T*>(input.data_ptr()), output_q.data_ptr(), \
static_cast<float*>(output_s.data_ptr()), group_size, num_groups, \
groups_per_block, (float)eps, (float)min_8bit, (float)max_8bit, \
scale_num_rows, scale_stride); \
} while (0)
#else
#define LAUNCH_KERNEL_INST(T, DST_DTYPE, COL_MAJOR, UE8M0, SMEM_BYTES) \
do { \
per_token_group_quant_8bit_kernel<T, DST_DTYPE, COL_MAJOR, UE8M0> \
<<<dim3(num_blocks), dim3(num_threads), (SMEM_BYTES), stream>>>( \
static_cast<T*>(input.data_ptr()), output_q.data_ptr(), \
static_cast<float*>(output_s.data_ptr()), group_size, \
num_groups, groups_per_block, (float)eps, (float)min_8bit, \
(float)max_8bit, scale_num_rows, scale_stride); \
#define LAUNCH_KERNEL_IMPL(T, DST_DTYPE, COL_MAJOR, UE8M0, CHECK_BOUNDS, \
SMEM_BYTES) \
do { \
per_token_group_quant_8bit_kernel<T, DST_DTYPE, COL_MAJOR, UE8M0, float, \
CHECK_BOUNDS> \
<<<dim3(num_blocks), dim3(num_threads), (SMEM_BYTES), stream>>>( \
static_cast<T*>(input.data_ptr()), output_q.data_ptr(), \
static_cast<float*>(output_s.data_ptr()), group_size, \
num_groups, groups_per_block, (float)eps, (float)min_8bit, \
(float)max_8bit, scale_num_rows, scale_stride); \
} while (0)
#endif
#define LAUNCH_KERNEL_INST(T, DST_DTYPE, COL_MAJOR, UE8M0, SMEM_BYTES) \
do { \
if (needs_bounds_check) { \
LAUNCH_KERNEL_IMPL(T, DST_DTYPE, COL_MAJOR, UE8M0, true, SMEM_BYTES); \
} else { \
LAUNCH_KERNEL_IMPL(T, DST_DTYPE, COL_MAJOR, UE8M0, false, SMEM_BYTES); \
} \
} while (0)
#define LAUNCH_KERNEL(T, DST_DTYPE) \
do { \
size_t smem_bytes = \
@@ -284,6 +313,7 @@ void per_token_group_quant_8bit(const torch::stable::Tensor& input,
#undef LAUNCH_KERNEL
#undef LAUNCH_KERNEL_INST
#undef LAUNCH_KERNEL_IMPL
}
// Register-resident fast path for group_size==128.
tests/kernels/quantization/test_per_token_group_quant.py
+43 -6
View File
@@ -11,8 +11,24 @@ from vllm.platforms import current_platform
@pytest.mark.parametrize(
"shape", [(31, 128), (32, 128), (63, 256), (64, 256), (16, 512)]
"shape",
[
(0, 128),
(1, 128),
(3, 384),
(5, 384),
(7, 384),
(15, 384),
(31, 128),
(32, 128),
(63, 256),
(64, 256),
(16, 512),
(3, 2688),
(3, 5376),
],
)
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
@pytest.mark.parametrize("column_major", [False, True])
@pytest.mark.parametrize("tma_aligned", [False, True])
@pytest.mark.parametrize("scale_ue8m0", [False, True])
@@ -22,14 +38,20 @@ from vllm.platforms import current_platform
reason="Only test on CUDA/ROCm/XPU.",
)
def test_per_token_group_quant_fp8(
shape, column_major: bool, tma_aligned: bool, scale_ue8m0: bool, group_size: int
shape,
dtype,
column_major: bool,
tma_aligned: bool,
scale_ue8m0: bool,
group_size: int,
):
"""Tail groups preserve FP8 values and scales in every supported layout."""
device = current_platform.device_type
torch.manual_seed(42)
num_tokens, hidden_dim = shape
x = torch.randn((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) * 8
x = torch.randn((num_tokens, hidden_dim), device=device, dtype=dtype) * 8
# native kernel path
out_q, scale = fp8_utils.per_token_group_quant_fp8(
@@ -411,19 +433,34 @@ def test_per_token_group_quant_fp8_packed_large_mn():
assert torch.equal(out_s_packed.cpu(), expected.cpu()), "Packed scale mismatch"
@pytest.mark.parametrize("shape", [(32, 128), (64, 256), (16, 512)])
@pytest.mark.parametrize(
"shape",
[
(0, 128),
(1, 128),
(3, 384),
(5, 384),
(7, 384),
(15, 384),
(32, 128),
(64, 256),
(16, 512),
],
)
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
@pytest.mark.parametrize("group_size", [64, 128])
@pytest.mark.skipif(
not (current_platform.is_cuda_alike() or current_platform.is_xpu()),
reason="Only test on CUDA/ROCm/XPU.",
)
def test_per_token_group_quant_int8(shape, group_size: int):
def test_per_token_group_quant_int8(shape, dtype, group_size: int):
"""The shared launch policy also handles INT8 tail groups."""
device = current_platform.device_type
torch.manual_seed(42)
num_tokens, hidden_dim = shape
x = torch.randn((num_tokens, hidden_dim), device=device, dtype=torch.bfloat16) * 8
x = torch.randn((num_tokens, hidden_dim), device=device, dtype=dtype) * 8
# cuda path
out_q, scale = int8_utils.per_token_group_quant_int8(