| 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, |
| 61 | bool STRIDE_J_ZERO> |
| 62 | __global__ void scaled_fp8_quant_kernel_strided_group_shape( |
| 63 | fp8_type* __restrict__ out, const scalar_t* __restrict__ input, |
| 64 | const float* __restrict__ scale, int hidden_size, int64_t in_row_stride, |
| 65 | int64_t out_row_stride, int group_m, int group_n, int64_t scale_stride_i, |
| 66 | int64_t scale_stride_j) { |
| 67 | const int64_t token_idx = blockIdx.x; |
| 68 | const int tid = threadIdx.x; |
| 69 | |
| 70 | const scalar_t* token_in = input + token_idx * in_row_stride; |
| 71 | fp8_type* token_out = out + token_idx * out_row_stride; |
| 72 | |
| 73 | // Precompute row-level base offset for scale access (compile-time eliminated |
| 74 | // when STRIDE_I_ZERO) |
| 75 | const int64_t scale_row_base = |
| 76 | STRIDE_I_ZERO ? 0 |
| 77 | : static_cast<int>(token_idx) / group_m * scale_stride_i; |
| 78 | |
| 79 | auto get_inv_scale = [&](int gj) { |
| 80 | return 1.0f / scale[scale_row_base + gj * scale_stride_j]; |
| 81 | }; |
| 82 | |
| 83 | int cached_gj = -1; |
| 84 | float cached_inv_scale = 0.0f; |
| 85 | auto get_inv_scale_cached = [&](int gj) { |
| 86 | if (gj != cached_gj) { |
| 87 | cached_inv_scale = 1.0f / scale[scale_row_base + gj * scale_stride_j]; |
| 88 | cached_gj = gj; |
| 89 | } |
| 90 | return cached_inv_scale; |
| 91 | }; |
| 92 | |
| 93 | constexpr int VEC_SIZE = 16; // FP8 so vectorize to 128 bits |
| 94 | auto scaled_fp8_conversion_vectorized = [&](const scalar_t* in, fp8_type* out, |
| 95 | int size, float inv_scale) { |
| 96 | vectorize_with_alignment<VEC_SIZE>( |
| 97 | in, out, size, tid, blockDim.x, |
| 98 | [=] __device__(fp8_type & dst, const scalar_t& src) { |
| 99 | dst = scaled_fp8_conversion<true, fp8_type>(static_cast<float>(src), |
| 100 | inv_scale); |
| 101 | }); |
| 102 | }; |
| 103 | |
| 104 | if (STRIDE_J_ZERO && hidden_size % VEC_SIZE == 0) { |
| 105 | // Per-tensor or per-token: single scale per row, vectorize full row |
| 106 | scaled_fp8_conversion_vectorized(token_in, token_out, hidden_size, |
| 107 | get_inv_scale(0)); |
| 108 | } else if (group_n % VEC_SIZE == 0) { |
| 109 | // Multiple column groups with vectorization |
| 110 | const int num_groups_n = hidden_size / group_n; |
| 111 | |
| 112 | for (int gj = 0; gj < num_groups_n; gj++) { |
| 113 | scaled_fp8_conversion_vectorized(token_in + gj * group_n, |
| 114 | token_out + gj * group_n, group_n, |
| 115 | get_inv_scale(gj)); |
| 116 | } |
| 117 | } else { |
| 118 | // Scalar path for small column groups (group_n < VEC_SIZE) |
| 119 | for (int n = tid; n < hidden_size; n += blockDim.x) { |
| 120 | const int gj = n / group_n; |
| 121 | token_out[n] = scaled_fp8_conversion<true, fp8_type>( |
| 122 | static_cast<float>(token_in[n]), get_inv_scale_cached(gj)); |
| 123 | } |
| 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> |
| 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 | |
| 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 | } |
| 167 | const scalar_t* row_ptr = input + token_idx * in_row_stride; |
| 168 | |
| 169 | float thread_max = 0.0f; |
| 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 | |
| 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 | |
| 183 | template <typename scalar_t, typename fp8_type> |
| 184 | __global__ void scaled_fp8_quant_kernel_strided_dynamic( |
| 185 | fp8_type* __restrict__ out, const scalar_t* __restrict__ input, |
| 186 | const float* __restrict__ scale, int hidden_size, int64_t in_row_stride, |
| 187 | int64_t out_row_stride) { |
| 188 | const int64_t token_idx = blockIdx.x; |
| 189 | const int tid = threadIdx.x; |
| 190 | |
| 191 | const scalar_t* token_in = input + token_idx * in_row_stride; |
| 192 | fp8_type* token_out = out + token_idx * out_row_stride; |
| 193 | |
| 194 | const float reciprocal_scale = 1.0f / (*scale); |
| 195 | vectorize_with_alignment<16>( |
| 196 | token_in, token_out, hidden_size, tid, blockDim.x, |
| 197 | [=] __device__(fp8_type & dst, const scalar_t& src) { |
| 198 | dst = scaled_fp8_conversion<true, fp8_type>(static_cast<float>(src), |
| 199 | reciprocal_scale); |
| 200 | }); |
| 201 | } |
| 202 | |
| 203 | template <typename scalar_t, typename fp8_type> |
| 204 | __global__ void dynamic_per_token_scaled_fp8_quant_kernel_strided( |
| 205 | fp8_type* __restrict__ out, float* __restrict__ scale, |
| 206 | const scalar_t* __restrict__ input, const float* __restrict__ scale_ub, |
| 207 | int hidden_size, int64_t in_row_stride, int64_t out_row_stride) { |
| 208 | const int64_t token_idx = blockIdx.x; |
| 209 | const int tid = threadIdx.x; |
| 210 | |
| 211 | // Use int64 to avoid overflowing an int32 when calculating this offset |
| 212 | int64_t in_offset = static_cast<int64_t>(token_idx) * in_row_stride; |
| 213 | int64_t out_offset = static_cast<int64_t>(token_idx) * out_row_stride; |
| 214 | const scalar_t* token_in = input + in_offset; |
| 215 | fp8_type* token_out = out + out_offset; |
| 216 | |
| 217 | // 1) per-token absmax |
| 218 | float absmax_val = 0.f; |
| 219 | vectorize_read_with_alignment<16>( |
| 220 | token_in, hidden_size, tid, blockDim.x, [&] __device__(scalar_t v) { |
| 221 | absmax_val = fmaxf(absmax_val, fabsf(static_cast<float>(v))); |
| 222 | }); |
| 223 | |
| 224 | using BlockReduce = cub::BlockReduce<float, 256>; |
| 225 | __shared__ typename BlockReduce::TempStorage tmp; |
| 226 | const float block_max = |
| 227 | BlockReduce(tmp).Reduce(absmax_val, CubMaxOp{}, blockDim.x); |
| 228 | |
| 229 | __shared__ float token_scale; |
| 230 | if (tid == 0) { |
| 231 | token_scale = scale_ub ? fminf(block_max, *scale_ub) : block_max; |
| 232 | token_scale = fmaxf(token_scale / quant_type_max_v<fp8_type>, |
| 233 | min_scaling_factor<fp8_type>::val()); |
| 234 | scale[token_idx] = token_scale; |
| 235 | } |
| 236 | __syncthreads(); |
| 237 | |
| 238 | // 2) quantize |
| 239 | vectorize_with_alignment<16>( |
| 240 | token_in, token_out, hidden_size, tid, blockDim.x, |
| 241 | [=] __device__(fp8_type & dst, const scalar_t& src) { |
| 242 | dst = scaled_fp8_conversion<false, fp8_type>(static_cast<float>(src), |
| 243 | token_scale); |
| 244 | }); |
| 245 | } |
| 246 | |
| 247 | } // namespace vllm |
| 248 | |
| 249 | void static_scaled_fp8_quant( |
| 250 | torch::stable::Tensor& out, // [..., d] |
| 251 | torch::stable::Tensor const& input, // [..., d] |
| 252 | torch::stable::Tensor const& scale, // various shapes |
| 253 | std::optional<torch::headeronly::IntHeaderOnlyArrayRef> |
| 254 | opt_group_shape) // optional explicit [group_m, group_n] |
| 255 | { |
| 256 | STD_TORCH_CHECK(input.stride(-1) == 1, |
| 257 | "last dimension of input must be contiguous"); |
| 258 | STD_TORCH_CHECK(out.stride(-1) == 1, |
| 259 | "last dimension of output must be contiguous"); |
| 260 | |
| 261 | const int hidden_size = input.size(-1); // N (columns) |
| 262 | const int num_tokens = input.numel() / hidden_size; // M (rows) |
| 263 | |
| 264 | // Determine group_m, group_n, and scale strides from scale shape |
| 265 | // Scale indexing: scale[gi * scale_stride_j + gj * scale_stride_i] |
| 266 | // where gi = m / group_m, gj = n / group_n |
| 267 | int group_m, group_n; |
| 268 | int64_t scale_stride_i, scale_stride_j; |
| 269 | |
| 270 | if (scale.dim() == 0 || scale.numel() == 1) { |
| 271 | // Per-tensor: one scale for the entire tensor |
| 272 | group_m = num_tokens; |
| 273 | group_n = hidden_size; |
| 274 | scale_stride_i = 0; |
| 275 | scale_stride_j = 0; |
| 276 | } else if (scale.dim() == 1) { |
| 277 | // 1D scale: require explicit group_shape to disambiguate per-channel vs |
| 278 | // per-token (avoids edge case where num_tokens == hidden_size) |
| 279 | STD_TORCH_CHECK( |
| 280 | opt_group_shape.has_value(), |
| 281 | "1D scale requires explicit group_shape to disambiguate " |
| 282 | "per-channel vs per-token quantization. " |
| 283 | "Use group_shape=(-1, 1) for per-channel or group_shape=(1, " |
| 284 | "-1) for per-token."); |
| 285 | STD_TORCH_CHECK(opt_group_shape->size() == 2, |
| 286 | "group_shape must have exactly 2 elements, got ", |
| 287 | opt_group_shape->size()); |
| 288 | |
| 289 | const auto opt_group_m = (*opt_group_shape)[0]; |
| 290 | const auto opt_group_n = (*opt_group_shape)[1]; |
| 291 | group_m = opt_group_m == -1 ? num_tokens : static_cast<int>(opt_group_m); |
| 292 | group_n = opt_group_n == -1 ? hidden_size : static_cast<int>(opt_group_n); |
| 293 | |
| 294 | // Validate the explicit group shape matches the 1D scale |
| 295 | const int64_t scale_len = scale.numel(); |
| 296 | const int64_t expected_scale_m = num_tokens / group_m; |
| 297 | const int64_t expected_scale_n = hidden_size / group_n; |
| 298 | const int64_t expected_scale_numel = expected_scale_m * expected_scale_n; |
| 299 | |
| 300 | STD_TORCH_CHECK(scale_len == expected_scale_numel, "1D scale length (", |
| 301 | scale_len, ") does not match expected size (", |
| 302 | expected_scale_numel, ") for group_shape (", opt_group_m, |
| 303 | ", ", opt_group_n, ") with input shape (", num_tokens, ", ", |
| 304 | hidden_size, ")"); |
| 305 | |
| 306 | // For 1D scale, determine strides based on which dim is trivial |
| 307 | // Scale indexing: scale[gi * scale_stride_i + gj * scale_stride_j] |
| 308 | // where gi = m / group_m (row group), gj = n / group_n (col group) |
| 309 | if (expected_scale_m == 1) { |
| 310 | // Per-channel style: one scale in M dim, scale varies along N |
| 311 | // gi = 0 always, gj varies, so stride_1 traverses the scale |
| 312 | scale_stride_i = 0; |
| 313 | scale_stride_j = scale.stride(0); |
| 314 | } else if (expected_scale_n == 1) { |
| 315 | // Per-token style: one scale in N dim, scale varies along M |
| 316 | // gj = 0 always, gi varies, so stride_0 traverses the scale |
| 317 | scale_stride_i = scale.stride(0); |
| 318 | scale_stride_j = 0; |
| 319 | } else { |
| 320 | STD_TORCH_CHECK( |
| 321 | false, |
| 322 | "1D scale can only be used when one of the scale dimensions is 1. " |
| 323 | "For 2D group scaling, use a 2D scale tensor."); |
| 324 | } |
| 325 | } else if (scale.dim() == 2) { |
| 326 | // 2D scale: infer group sizes from scale dimensions (or use explicit if |
| 327 | // provided) |
| 328 | const int64_t scale_size_0 = scale.size(0); |
| 329 | const int64_t scale_size_1 = scale.size(1); |
| 330 | |
| 331 | STD_TORCH_CHECK(num_tokens % scale_size_0 == 0, "num_tokens (", num_tokens, |
| 332 | ") must be divisible by scale.size(0) (", scale_size_0, |
| 333 | ")"); |
| 334 | STD_TORCH_CHECK(hidden_size % scale_size_1 == 0, "hidden_size (", |
| 335 | hidden_size, ") must be divisible by scale.size(1) (", |
| 336 | scale_size_1, ")"); |
| 337 | |
| 338 | // Infer from 2D scale shape |
| 339 | int inferred_group_m = num_tokens / scale_size_0; |
| 340 | int inferred_group_n = hidden_size / scale_size_1; |
| 341 | |
| 342 | // Use explicit if provided, otherwise use inferred |
| 343 | if (opt_group_shape.has_value()) { |
| 344 | STD_TORCH_CHECK(opt_group_shape->size() == 2, |
| 345 | "group_shape must have exactly 2 elements, got ", |
| 346 | opt_group_shape->size()); |
| 347 | const auto opt_group_m = (*opt_group_shape)[0]; |
| 348 | const auto opt_group_n = (*opt_group_shape)[1]; |
| 349 | group_m = opt_group_m == -1 ? num_tokens : static_cast<int>(opt_group_m); |
| 350 | group_n = opt_group_n == -1 ? hidden_size : static_cast<int>(opt_group_n); |
| 351 | |
| 352 | // Validate explicit matches inferred |
| 353 | STD_TORCH_CHECK( |
| 354 | group_m == inferred_group_m && group_n == inferred_group_n, |
| 355 | "Explicit group_shape (", opt_group_m, ", ", opt_group_n, |
| 356 | ") does not match inferred group shape (", inferred_group_m, ", ", |
| 357 | inferred_group_n, ") from 2D scale tensor shape (", scale_size_0, |
| 358 | ", ", scale_size_1, ")"); |
| 359 | } else { |
| 360 | group_m = inferred_group_m; |
| 361 | group_n = inferred_group_n; |
| 362 | } |
| 363 | |
| 364 | scale_stride_i = scale.stride(0); |
| 365 | scale_stride_j = scale.stride(1); |
| 366 | } else { |
| 367 | STD_TORCH_CHECK(false, "scale must be 0D, 1D, or 2D tensor, but got ", |
| 368 | scale.dim(), "D"); |
| 369 | } |
| 370 | |
| 371 | const int block_size = 256; |
| 372 | dim3 grid(num_tokens); |
| 373 | dim3 block(block_size); |
| 374 | |
| 375 | const int64_t in_row_stride = input.stride(-2); |
| 376 | const int64_t out_row_stride = out.stride(-2); |
| 377 | |
| 378 | const torch::stable::accelerator::DeviceGuard device_guard( |
| 379 | input.get_device_index()); |
| 380 | const cudaStream_t stream = get_current_cuda_stream(); |
| 381 | |
| 382 | // Dispatch to template-specialized kernel based on stride pattern |
| 383 | VLLM_STABLE_DISPATCH_FLOATING_TYPES( |
| 384 | input.scalar_type(), "scaled_fp8_quant_kernel_scalar_type", [&] { |
| 385 | VLLM_STABLE_DISPATCH_FP8_TYPES( |
| 386 | out.scalar_type(), "scaled_fp8_quant_kernel_fp8_type", [&] { |
| 387 | VLLM_STABLE_DISPATCH_BOOL(scale_stride_i == 0, S0_ZERO, [&] { |
| 388 | VLLM_STABLE_DISPATCH_BOOL(scale_stride_j == 0, S1_ZERO, [&] { |
| 389 | vllm::scaled_fp8_quant_kernel_strided_group_shape< |
| 390 | scalar_t, fp8_t, S0_ZERO, S1_ZERO> |
| 391 | <<<grid, block, 0, stream>>>( |
| 392 | out.mutable_data_ptr<fp8_t>(), |
| 393 | input.const_data_ptr<scalar_t>(), |
| 394 | scale.const_data_ptr<float>(), hidden_size, |
| 395 | in_row_stride, out_row_stride, group_m, group_n, |
| 396 | scale_stride_i, scale_stride_j); |
| 397 | }); |
| 398 | }); |
| 399 | }); |
| 400 | }); |
| 401 | } |
| 402 | |
| 403 | void dynamic_scaled_fp8_quant(torch::stable::Tensor& out, // [..., d] |
| 404 | torch::stable::Tensor const& input, // [..., d] |
| 405 | torch::stable::Tensor& scale) // [1] |
| 406 | { |
| 407 | STD_TORCH_CHECK(input.stride(-1) == 1, |
| 408 | "last dimension of input must be contiguous"); |
| 409 | STD_TORCH_CHECK(out.stride(-1) == 1, |
| 410 | "last dimension of output must be contiguous"); |
| 411 | |
| 412 | const int hidden_size = input.size(-1); |
| 413 | const int num_tokens = input.numel() / hidden_size; |
| 414 | const int block_size = vllm::kFp8QuantBlockSize; |
| 415 | dim3 grid(num_tokens); |
| 416 | dim3 block(block_size); |
| 417 | |
| 418 | const int64_t in_row_stride = input.stride(-2); |
| 419 | const int64_t out_row_stride = out.stride(-2); |
| 420 | |
| 421 | const torch::stable::accelerator::DeviceGuard device_guard( |
| 422 | input.get_device_index()); |
| 423 | const cudaStream_t stream = get_current_cuda_stream(); |
| 424 | |
| 425 | // scale tensor should be initialised to <=0 before reduction |
| 426 | STD_CUDA_CHECK(cudaMemsetAsync(scale.mutable_data_ptr<float>(), 0, |
| 427 | sizeof(float), stream)); |
| 428 | |
| 429 | VLLM_STABLE_DISPATCH_FLOATING_TYPES( |
| 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", [&] { |
| 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>(), |
| 459 | input.const_data_ptr<scalar_t>(), |
| 460 | scale.const_data_ptr<float>(), |
| 461 | hidden_size, in_row_stride, |
| 462 | out_row_stride); |
| 463 | }); |
| 464 | }); |
| 465 | } |
| 466 | |
| 467 | void dynamic_per_token_scaled_fp8_quant( |
| 468 | torch::stable::Tensor& out, // [..., d] |
| 469 | torch::stable::Tensor const& input, // [..., d] |
| 470 | torch::stable::Tensor& scales, |
| 471 | std::optional<torch::stable::Tensor> const& scale_ub) { |
| 472 | STD_TORCH_CHECK(input.stride(-1) == 1, |
| 473 | "last dimension of input must be contiguous"); |
| 474 | STD_TORCH_CHECK(out.stride(-1) == 1, |
| 475 | "last dimension of output must be contiguous"); |
| 476 | |
| 477 | const int hidden_size = input.size(-1); |
| 478 | const int num_tokens = input.numel() / hidden_size; |
| 479 | const int block_size = 256; |
| 480 | dim3 grid(num_tokens); |
| 481 | dim3 block(std::min(hidden_size, block_size)); |
| 482 | |
| 483 | const int64_t in_row_stride = input.stride(-2); |
| 484 | const int64_t out_row_stride = out.stride(-2); |
| 485 | |
| 486 | const torch::stable::accelerator::DeviceGuard device_guard( |
| 487 | input.get_device_index()); |
| 488 | const cudaStream_t stream = get_current_cuda_stream(); |
| 489 | VLLM_STABLE_DISPATCH_FLOATING_TYPES( |
| 490 | input.scalar_type(), |
| 491 | "dynamic_per_token_scaled_fp8_quant_kernel_scalar_type", [&] { |
| 492 | VLLM_STABLE_DISPATCH_FP8_TYPES( |
| 493 | out.scalar_type(), |
| 494 | "dynamic_per_token_scaled_fp8_quant_kernel_fp8_type", [&] { |
| 495 | vllm::dynamic_per_token_scaled_fp8_quant_kernel_strided<scalar_t, |
| 496 | fp8_t> |
| 497 | <<<grid, block, 0, stream>>>( |
| 498 | out.mutable_data_ptr<fp8_t>(), |
| 499 | scales.mutable_data_ptr<float>(), |
| 500 | input.const_data_ptr<scalar_t>(), |
| 501 | scale_ub.has_value() ? scale_ub->const_data_ptr<float>() |
| 502 | : nullptr, |
| 503 | hidden_size, in_row_stride, out_row_stride); |
| 504 | }); |
| 505 | }); |
| 506 | } |