common.cu 11.1 KB
Newer Older
1
#include <ATen/cuda/CUDAContext.h>
2
#include <torch/all.h>
3
4
5
6
7
8
9
#include <c10/cuda/CUDAGuard.h>

#include <cmath>

#include "cuda_compat.h"
#include "dispatch_utils.h"

10
11
12
13
14
15
16
#ifndef USE_ROCM
  #include <cub/util_type.cuh>
  #include <cub/cub.cuh>
#else
  #include <hipcub/util_type.hpp>
  #include <hipcub/hipcub.hpp>
#endif
17

18
19
20
21
22
23
24
25
26
27
28
29
#ifndef USE_ROCM
using FP8_TYPE = c10::Float8_e4m3fn;
C10_HOST_DEVICE constexpr auto FP8_E4M3_MAX =
    std::numeric_limits<FP8_TYPE>::max();
#else
  #include "amd/hip_float8.h"
using FP8_TYPE = c10::Float8_e4m3fnuz;
// Using the default max value from pytorch (240.0) will cause accuracy
// issue when running dynamic quantization. Here use 224.0f for rocm.
constexpr auto FP8_E4M3_MAX = 224.0f;
#endif

30
31
32
namespace vllm {

__device__ __forceinline__ float atomicMaxFloat(float* addr, float value) {
33
34
35
36
37
  float old;
  old = (value >= 0)
            ? __int_as_float(atomicMax((int*)addr, __float_as_int(value)))
            : __uint_as_float(
                  atomicMin((unsigned int*)addr, __float_as_uint(value)));
38

39
  return old;
40
41
}

42
template <bool is_scale_inverted>
43
44
__device__ __forceinline__ FP8_TYPE scaled_fp8_conversion(float const val,
                                                          float const scale) {
45
46
47
48
49
50
51
  float x = 0.0f;
  if constexpr (is_scale_inverted) {
    x = val * scale;
  } else {
    x = val / scale;
  }

52
  float r = fmax(-FP8_E4M3_MAX, fmin(x, FP8_E4M3_MAX));
53
#ifndef USE_ROCM
54
  return static_cast<c10::Float8_e4m3fn>(r);
55
56
57
58
59
#else
  // Use hardware cvt instruction for fp8 on rocm
  return c10::Float8_e4m3fnuz(hip_fp8(r).data,
                              c10::Float8_e4m3fnuz::from_bits());
#endif
60
61
}

62
63
64
65
66
67
// Compute the absolute maximum m of the input tensor and store
// m / float8_e4m3::max() in *scale. Each thread block performs a
// reduction tree and the memory in scale is atomically updated.
// So to get the right answer, *scale needs to be initialized to
// a value <= 0.0 and we need to wait for all thread blocks to
// finish before consuming *scale.
68
69
70
71
template <typename scalar_t>
__global__ void segmented_max_reduction(float* __restrict__ scale,
                                        const scalar_t* __restrict__ input,
                                        int64_t num_elems) {
72
  __shared__ float cache[1024];
73
  int64_t i = blockDim.x * blockIdx.x + threadIdx.x;
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90

  // First store maximum for all values processes by
  // the current thread in cache[threadIdx.x]
  scalar_t tmp = 0.0;
  while (i < num_elems) {
    float x = static_cast<float>(input[i]);
    tmp = max(tmp, fabs(x));
    i += blockDim.x * gridDim.x;
  }
  cache[threadIdx.x] = tmp;

  __syncthreads();

  // Now perform parallel reduction within the thread block
  int ib = blockDim.x / 2;
  while (ib != 0) {
    if (threadIdx.x < ib && cache[threadIdx.x + ib] > cache[threadIdx.x]) {
91
      cache[threadIdx.x] = cache[threadIdx.x + ib];
92
93
94
95
96
97
98
    }
    __syncthreads();
    ib /= 2;
  }
  // Finally, since cache[0] contains the maximum for this thread block,
  // atomically write the max to the target location
  if (threadIdx.x == 0) {
99
    atomicMaxFloat(scale, cache[0] / FP8_E4M3_MAX);
100
101
102
  }
}

103
104
105
106
107
108
109
110
111
template <typename scalar_t>
struct __align__(8) vec4_t {
  scalar_t x;
  scalar_t y;
  scalar_t z;
  scalar_t w;
};

typedef struct __align__(4) {
112
113
114
115
  FP8_TYPE x;
  FP8_TYPE y;
  FP8_TYPE z;
  FP8_TYPE w;
116
117
118
}
float8x4_t;

119
template <typename scalar_t>
120
121
122
123
124
125
__device__ float thread_max_vec(scalar_t const* __restrict__ input,
                                int64_t const num_elems, int const tid,
                                int const step) {
  // Vectorized input/output to better utilize memory bandwidth.
  vec4_t<scalar_t> const* vectorized_in =
      reinterpret_cast<vec4_t<scalar_t> const*>(input);
126

127
  int64_t const num_vec_elems = num_elems >> 2;
128
129
130
  float absmax_val = 0.0f;

#pragma unroll 4
131
  for (int64_t i = tid; i < num_vec_elems; i += step) {
132
133
134
135
136
137
    vec4_t<scalar_t> in_vec = vectorized_in[i];
    absmax_val = max(absmax_val, fabs(in_vec.x));
    absmax_val = max(absmax_val, fabs(in_vec.y));
    absmax_val = max(absmax_val, fabs(in_vec.z));
    absmax_val = max(absmax_val, fabs(in_vec.w));
  }
138

139
  // Handle the remaining elements if num_elems is not divisible by 4
140
  for (int64_t i = num_vec_elems * 4 + tid; i < num_elems; i += step) {
141
142
143
144
145
146
    absmax_val = max(absmax_val, fabs(input[i]));
  }

  return absmax_val;
}

147
template <typename scalar_t, bool is_scale_inverted>
148
__device__ void scaled_fp8_conversion_vec(FP8_TYPE* __restrict__ out,
149
                                          scalar_t const* __restrict__ input,
150
                                          float const scale,
151
152
                                          int64_t const num_elems,
                                          int const tid, int const step) {
153
  // Vectorized input/output to better utilize memory bandwidth.
154
155
  vec4_t<scalar_t> const* vectorized_in =
      reinterpret_cast<vec4_t<scalar_t> const*>(input);
156
157
  float8x4_t* vectorized_out = reinterpret_cast<float8x4_t*>(out);

158
  int64_t const num_vec_elems = num_elems >> 2;
159
160

#pragma unroll 4
161
  for (int64_t i = tid; i < num_vec_elems; i += step) {
162
163
164
    vec4_t<scalar_t> in_vec = vectorized_in[i];
    float8x4_t out_vec;

165
166
167
168
169
170
171
172
    out_vec.x = scaled_fp8_conversion<is_scale_inverted>(
        static_cast<float>(in_vec.x), scale);
    out_vec.y = scaled_fp8_conversion<is_scale_inverted>(
        static_cast<float>(in_vec.y), scale);
    out_vec.z = scaled_fp8_conversion<is_scale_inverted>(
        static_cast<float>(in_vec.z), scale);
    out_vec.w = scaled_fp8_conversion<is_scale_inverted>(
        static_cast<float>(in_vec.w), scale);
173
174
175
176
    vectorized_out[i] = out_vec;
  }

  // Handle the remaining elements if num_elems is not divisible by 4
177
  for (int64_t i = num_vec_elems * 4 + tid; i < num_elems; i += step) {
178
179
    out[i] = scaled_fp8_conversion<is_scale_inverted>(
        static_cast<float>(input[i]), scale);
180
181
182
  }
}

183
template <typename scalar_t>
184
__global__ void scaled_fp8_quant_kernel(FP8_TYPE* __restrict__ out,
185
186
187
188
189
190
191
192
                                        const scalar_t* __restrict__ input,
                                        const float* __restrict__ scale,
                                        int64_t num_elems) {
  int tid = blockDim.x * blockIdx.x + threadIdx.x;

  // Invert the scale so that we can use multiplications to avoid expensive
  // division.
  const float inverted_scale = 1.0f / (*scale);
193
194
  scaled_fp8_conversion_vec<scalar_t, true>(
      out, input, inverted_scale, num_elems, tid, blockDim.x * gridDim.x);
195
196
197
198
}

template <typename scalar_t>
__global__ void dynamic_per_token_scaled_fp8_quant_kernel(
199
    FP8_TYPE* __restrict__ out, float* __restrict__ scale,
200
201
202
203
    scalar_t const* __restrict__ input, float const* __restrict__ scale_ub,
    const int hidden_size) {
  float const min_scaling_factor = 1.0f / (FP8_E4M3_MAX * 512.f);

204
205
206
  int const tid = threadIdx.x;
  int const token_idx = blockIdx.x;

207
208
209
210
  // Use int64 to avoid overflowing an int32 when calculating this offset
  int64_t offset = static_cast<int64_t>(token_idx) * hidden_size;
  scalar_t const* __restrict__ token_input = &input[offset];
  FP8_TYPE* __restrict__ token_output = &out[offset];
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225

  // For vectorization, token_input and token_output pointers need to be
  // aligned at 8-byte and 4-byte addresses respectively.
  bool const can_vectorize = hidden_size % 4 == 0;

  float absmax_val = 0.0f;
  if (can_vectorize) {
    absmax_val = thread_max_vec(token_input, hidden_size, tid, blockDim.x);
  } else {
    for (int i = tid; i < hidden_size; i += blockDim.x) {
      float const x = static_cast<float>(token_input[i]);
      absmax_val = max(absmax_val, fabs(x));
    }
  }

226
227
228
229
  using BlockReduce = cub::BlockReduce<float, 1024>;
  __shared__ typename BlockReduce::TempStorage reduceStorage;
  float const block_absmax_val_maybe =
      BlockReduce(reduceStorage).Reduce(absmax_val, cub::Max{}, blockDim.x);
230
  __shared__ float token_scale;
231
  if (tid == 0) {
232
233
234
235
236
237
238
239
    if (scale_ub) {
      token_scale = min(block_absmax_val_maybe, *scale_ub);
    } else {
      token_scale = block_absmax_val_maybe;
    }
    // token scale computation
    token_scale = max(token_scale / FP8_E4M3_MAX, min_scaling_factor);
    scale[token_idx] = token_scale;
240
241
242
  }
  __syncthreads();

243
  // Note that we don't use inverted scales so we can match FBGemm impl.
244
  if (can_vectorize) {
245
246
    scaled_fp8_conversion_vec<scalar_t, false>(
        token_output, token_input, token_scale, hidden_size, tid, blockDim.x);
247
248
  } else {
    for (int i = tid; i < hidden_size; i += blockDim.x) {
249
250
      token_output[i] = scaled_fp8_conversion<false>(
          static_cast<float>(token_input[i]), token_scale);
251
252
253
254
    }
  }
}

255
}  // namespace vllm
256

257
258
259
void static_scaled_fp8_quant(torch::Tensor& out,          // [..., d]
                             torch::Tensor const& input,  // [..., d]
                             torch::Tensor const& scale)  // [1]
260
261
262
263
264
265
266
267
{
  int64_t num_tokens = input.numel() / input.size(-1);
  int64_t num_elems = input.numel();
  dim3 grid(num_tokens);
  dim3 block(1024);
  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  VLLM_DISPATCH_FLOATING_TYPES(
268
269
      input.scalar_type(), "scaled_fp8_quant_kernel", [&] {
        vllm::scaled_fp8_quant_kernel<scalar_t><<<grid, block, 0, stream>>>(
270
            out.data_ptr<FP8_TYPE>(), input.data_ptr<scalar_t>(),
271
            scale.data_ptr<float>(), num_elems);
272
273
274
      });
}

275
276
277
void dynamic_scaled_fp8_quant(torch::Tensor& out,          // [..., d]
                              torch::Tensor const& input,  // [..., d]
                              torch::Tensor& scale)        // [1]
278
279
280
281
282
283
284
285
{
  int64_t num_tokens = input.numel() / input.size(-1);
  int64_t num_elems = input.numel();
  dim3 grid(num_tokens);
  dim3 block(1024);
  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  VLLM_DISPATCH_FLOATING_TYPES(
286
287
288
289
      input.scalar_type(), "scaled_fp8_quant_kernel", [&] {
        vllm::segmented_max_reduction<scalar_t><<<grid, block, 0, stream>>>(
            scale.data_ptr<float>(), input.data_ptr<scalar_t>(), num_elems);
        vllm::scaled_fp8_quant_kernel<scalar_t><<<grid, block, 0, stream>>>(
290
            out.data_ptr<FP8_TYPE>(), input.data_ptr<scalar_t>(),
291
            scale.data_ptr<float>(), num_elems);
292
293
      });
}
294

295
296
297
298
void dynamic_per_token_scaled_fp8_quant(
    torch::Tensor& out,          // [..., d]
    torch::Tensor const& input,  // [..., d]
    torch::Tensor& scales, std::optional<at::Tensor> const& scale_ub) {
299
300
301
302
303
304
305
306
307
308
309
310
311
312
  TORCH_CHECK(input.is_contiguous());
  TORCH_CHECK(out.is_contiguous());

  int const hidden_size = input.size(-1);
  int const num_tokens = input.numel() / hidden_size;
  dim3 const grid(num_tokens);
  dim3 const block(std::min(hidden_size, 1024));

  const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
  const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
  VLLM_DISPATCH_FLOATING_TYPES(
      input.scalar_type(), "dynamic_per_token_scaled_fp8_quant_kernel", [&] {
        vllm::dynamic_per_token_scaled_fp8_quant_kernel<scalar_t>
            <<<grid, block, 0, stream>>>(
313
                out.data_ptr<FP8_TYPE>(), scales.data_ptr<float>(),
314
315
316
                input.data_ptr<scalar_t>(),
                scale_ub.has_value() ? scale_ub->data_ptr<float>() : nullptr,
                hidden_size);
317
318
      });
}