[Perf][Kernel] Vectorized flat abs-max for dynamic per-tensor FP8 quantization (#58194)
Signed-off-by: Monishver Chandrasekaran <monishverchandrasekaran@gmail.com> Co-authored-by: Claude Fable 5.1 <noreply@anthropic.com> Co-authored-by: Misha Goin <mgoin64@gmail.com>
Monishver committed
Sep 25, 2026 at 11:19 UTC
38cc054fc2269941cd35ecdfd198aacbb2eabfb8
2 files changed
+147
-33
csrc/libtorch_stable/quantization/w8a8/fp8/common.cu
+117
-33
@@ -1,11 +1,60 @@
1
#include "../../../../quantization/w8a8/fp8/common.cuh"
2
+#include "../../../../cuda_compat.h"
3
#include "../../../dispatch_utils.h"
4
#include "../../../cub_helpers.h"
5
#include "../../vectorization_utils.cuh"
6
#include "../../../torch_utils.h"
7
#include <torch/csrc/stable/macros.h>
8
+
9
+#include <algorithm>
10
+
11
namespace vllm {
12
13
+namespace {
14
+
15
+// 16 fp8 outputs form one 16-byte store; the matching 16-element input load
16
+// is 32 B for 16-bit types and 64 B for fp32.
17
+constexpr int kFp8QuantVecSize = 16;
18
+constexpr int kFp8QuantBlockSize = 256;
19
+
20
+template <typename T>
21
+using fp8_quant_vec_t = vec_n_t<T, kFp8QuantVecSize>;
22
+
23
+template <typename T>
24
+bool ptr_vec_aligned(const T* ptr) {
25
+ return reinterpret_cast<uintptr_t>(ptr) % alignof(fp8_quant_vec_t<T>) == 0;
26
+}
27
+
28
+__device__ __forceinline__ float warp_reduce_max(float v) {
29
+#pragma unroll
30
+ for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) {
31
+ v = fmaxf(v, VLLM_SHFL_XOR_SYNC(v, offset));
32
+ }
33
+ return v;
34
+}
35
+
36
+// Block-wide max, returned to every thread. Warp partials go through one
37
+// shared-memory exchange, so the whole reduction costs a single
38
+// __syncthreads. blockDim.x must be a multiple of WARP_SIZE and at most
39
+// kFp8QuantBlockSize.
40
+__device__ __forceinline__ float block_reduce_max(float v, float* warp_max) {
41
+ const int lane = threadIdx.x % WARP_SIZE;
42
+ const int warp = threadIdx.x / WARP_SIZE;
43
+ const int num_warps = blockDim.x / WARP_SIZE;
44
+ v = warp_reduce_max(v);
45
+ if (lane == 0) {
46
+ warp_max[warp] = v;
47
+ }
48
+ __syncthreads();
49
+ v = warp_max[0];
50
+ for (int w = 1; w < num_warps; ++w) {
51
+ v = fmaxf(v, warp_max[w]);
52
+ }
53
+ return v;
54
+}
55
+
56
+} // namespace
57
+
58
// STRIDE_I_ZERO: true if scale_stride_i == 0 (per-tensor or per-channel)
59
// STRIDE_J_ZERO: true if scale_stride_j == 0 (per-tensor or per-token)
60
template <typename scalar_t, typename fp8_type, bool STRIDE_I_ZERO,
@@ -75,42 +124,59 @@ __global__ void scaled_fp8_quant_kernel_strided_group_shape(
124
}
125
}
126
127
+// Per-tensor abs-max of a contiguous tensor viewed as a flat array of
128
+// 16-element vectors: grid-stride loop, one atomic per block.
129
template <typename scalar_t, typename fp8_type>
79
-__global__ void segmented_max_reduction_strided(
80
- float* __restrict__ scale, const scalar_t* __restrict__ input,
81
- int hidden_size, int64_t in_row_stride, int64_t num_tokens) {
82
- __shared__ float cache[256];
83
- const int tid = threadIdx.x;
84
- int64_t token_idx = blockIdx.x;
130
+__global__ void __launch_bounds__(kFp8QuantBlockSize)
131
+ absmax_reduction_flat(float* __restrict__ scale,
132
+ const scalar_t* __restrict__ input,
133
+ int64_t num_vecs) {
134
+ using in_vec_t = fp8_quant_vec_t<scalar_t>;
135
+ const auto* v_in = reinterpret_cast<const in_vec_t*>(input);
136
+ const int64_t stride = static_cast<int64_t>(gridDim.x) * blockDim.x;
137
86
- // one block per token. Guard in case gridDim.x > num_tokens.
138
+ float thread_max = 0.0f;
139
+#pragma unroll 4
140
+ for (int64_t v = static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x;
141
+ v < num_vecs; v += stride) {
142
+ const in_vec_t vec = v_in[v];
143
+#pragma unroll
144
+ for (int i = 0; i < kFp8QuantVecSize; ++i) {
145
+ thread_max = fmaxf(thread_max, fabsf(static_cast<float>(vec.val[i])));
146
+ }
147
+ }
148
+
149
+ __shared__ float warp_max[kFp8QuantBlockSize / 32];
150
+ const float block_max = block_reduce_max(thread_max, warp_max);
151
+ if (threadIdx.x == 0) {
152
+ atomicMaxFloat(scale, block_max / quant_type_max_v<fp8_type>);
153
+ }
154
+}
155
+
156
+// Per-tensor abs-max for strided or unaligned rows: one block per token.
157
+template <typename scalar_t, typename fp8_type>
158
+__global__ void __launch_bounds__(kFp8QuantBlockSize)
159
+ segmented_max_reduction_strided(float* __restrict__ scale,
160
+ const scalar_t* __restrict__ input,
161
+ int hidden_size, int64_t in_row_stride,
162
+ int64_t num_tokens) {
163
+ const int64_t token_idx = blockIdx.x;
164
if (token_idx >= num_tokens) {
165
return;
166
}
90
-
167
const scalar_t* row_ptr = input + token_idx * in_row_stride;
168
93
- // each thread scans elements of the row in a strided fashion.
169
float thread_max = 0.0f;
95
- for (int e = tid; e < hidden_size; e += blockDim.x) {
96
- float v = fabsf(static_cast<float>(row_ptr[e]));
97
- thread_max = fmaxf(thread_max, v);
98
- }
99
-
100
- cache[tid] = thread_max;
101
- __syncthreads();
102
-
103
- // parallel reduction to find row max.
104
- for (int offset = blockDim.x / 2; offset > 0; offset >>= 1) {
105
- if (tid < offset) {
106
- cache[tid] = fmaxf(cache[tid], cache[tid + offset]);
107
- }
108
- __syncthreads();
109
- }
170
+ vectorize_read_with_alignment<kFp8QuantVecSize>(
171
+ row_ptr, hidden_size, threadIdx.x, blockDim.x,
172
+ [&] __device__(const scalar_t& v) {
173
+ thread_max = fmaxf(thread_max, fabsf(static_cast<float>(v)));
174
+ });
175
111
- // thread 0 updates global scale (per-tensor) atomically.
112
- if (tid == 0) {
113
- atomicMaxFloat(scale, cache[0] / quant_type_max_v<fp8_type>);
176
+ __shared__ float warp_max[kFp8QuantBlockSize / 32];
177
+ const float block_max = block_reduce_max(thread_max, warp_max);
178
+ if (threadIdx.x == 0) {
179
+ atomicMaxFloat(scale, block_max / quant_type_max_v<fp8_type>);
180
}
181
}
182
@@ -345,7 +411,7 @@ void dynamic_scaled_fp8_quant(torch::stable::Tensor& out, // [..., d]
411
412
const int hidden_size = input.size(-1);
413
const int num_tokens = input.numel() / hidden_size;
348
- const int block_size = 256;
414
+ const int block_size = vllm::kFp8QuantBlockSize;
415
dim3 grid(num_tokens);
416
dim3 block(block_size);
417
@@ -364,11 +430,29 @@ void dynamic_scaled_fp8_quant(torch::stable::Tensor& out, // [..., d]
430
input.scalar_type(), "scaled_fp8_quant_kernel_scalar_type", [&] {
431
VLLM_STABLE_DISPATCH_FP8_TYPES(
432
out.scalar_type(), "scaled_fp8_quant_kernel_fp8_type", [&] {
367
- vllm::segmented_max_reduction_strided<scalar_t, fp8_t>
368
- <<<grid, block, 0, stream>>>(
369
- scale.mutable_data_ptr<float>(),
370
- input.const_data_ptr<scalar_t>(), hidden_size,
371
- in_row_stride, static_cast<int64_t>(num_tokens));
433
+ const scalar_t* in_ptr = input.const_data_ptr<scalar_t>();
434
+ const bool flat = in_row_stride == hidden_size &&
435
+ hidden_size % vllm::kFp8QuantVecSize == 0 &&
436
+ vllm::ptr_vec_aligned(in_ptr);
437
+ if (flat) {
438
+ const int64_t num_vecs = input.numel() / vllm::kFp8QuantVecSize;
439
+ const int64_t max_blocks =
440
+ static_cast<int64_t>(
441
+ get_device_prop()->multiProcessorCount) *
442
+ 8;
443
+ const int64_t num_blocks = std::max<int64_t>(
444
+ 1,
445
+ std::min<int64_t>(
446
+ max_blocks, (num_vecs + block_size - 1) / block_size));
447
+ vllm::absmax_reduction_flat<scalar_t, fp8_t>
448
+ <<<num_blocks, block, 0, stream>>>(
449
+ scale.mutable_data_ptr<float>(), in_ptr, num_vecs);
450
+ } else {
451
+ vllm::segmented_max_reduction_strided<scalar_t, fp8_t>
452
+ <<<grid, block, 0, stream>>>(
453
+ scale.mutable_data_ptr<float>(), in_ptr, hidden_size,
454
+ in_row_stride, static_cast<int64_t>(num_tokens));
455
+ }
456
457
vllm::scaled_fp8_quant_kernel_strided_dynamic<scalar_t, fp8_t>
458
<<<grid, block, 0, stream>>>(out.mutable_data_ptr<fp8_t>(),
tests/kernels/quantization/test_fp8_quant.py
+30
@@ -108,6 +108,36 @@ def test_dynamic_per_tensor_fp8_quant(
108
opcheck_fp8_quant(ops_out, x)
109
110
111
+@pytest.mark.parametrize("pad", [8, 16])
112
+@pytest.mark.parametrize("dtype", DTYPES)
113
+@pytest.mark.parametrize("seed", SEEDS)
114
+@torch.inference_mode()
115
+def test_dynamic_per_tensor_fp8_quant_strided_input(
116
+ pad: int, dtype: torch.dtype, seed: int
117
+) -> None:
118
+ """Non-contiguous rows take the strided abs-max kernel; contiguous rows take
119
+ the flat one. Both must produce exactly the same scale and bytes.
120
+
121
+ pad=8 (16 B for bf16) leaves rows misaligned for 32 B vector loads, pad=16
122
+ keeps them aligned but non-contiguous.
123
+ """
124
+ set_random_seed(seed)
125
+ num_tokens, hidden_size = 64, 4096
126
+
127
+ full = torch.rand(num_tokens, hidden_size + pad, dtype=dtype, device="cuda")
128
+ x = full[:, :hidden_size]
129
+ assert not x.is_contiguous()
130
+
131
+ ops_out, ops_scale = ops.scaled_fp8_quant(x)
132
+ flat_out, flat_scale = ops.scaled_fp8_quant(x.contiguous())
133
+ assert torch.equal(ops_scale, flat_scale)
134
+ assert torch.equal(ops_out.view(torch.uint8), flat_out.view(torch.uint8))
135
+
136
+ ref_out, ref_scale = ref_dynamic_per_tensor_fp8_quant(x)
137
+ assert torch.equal(ref_scale, ops_scale)
138
+ assert torch.equal(ref_out.view(torch.uint8), ops_out.view(torch.uint8))
139
+
140
+
141
# Regression test for a case with large activations where an int32 index cannot
142
# represent the number of elements.
143
@torch.inference_mode()