@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / a1c8f3796c
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 15, 19:45 UTC

dynamic distpatch of fp8 kernels (#14245)

Signed-off-by: Jeff Daily <jeff.daily@amd.com>

Jeff Daily committed Mar 11, 2025 at 07:54 UTC a1c8f3796c89f11c48b3e70e66feafd080273fa5
25 files changed +292 -159
benchmarks/kernels/benchmark_moe.py
+1 -2
@@ -18,8 +18,7 @@ from vllm.model_executor.layers.fused_moe.fused_moe import *
18 from vllm.platforms import current_platform
19 from vllm.utils import FlexibleArgumentParser
20
21 -FP8_DTYPE = torch.float8_e4m3fnuz if current_platform.is_rocm(
22 -) else torch.float8_e4m3fn
21 +FP8_DTYPE = current_platform.fp8_dtype()
22
23
24 class BenchmarkConfig(TypedDict):
csrc/dispatch_utils.h
+26 -6
@@ -6,6 +6,11 @@
6
7 #include <torch/all.h>
8
9 +// Need a special dispatch case macro since we will nest the FP8 dispatch.
10 +// Instead of the usual 'scalar_t', this names the dispatched type 'fp8_t'.
11 +#define AT_DISPATCH_FP8_CASE(enum_type, ...) \
12 + AT_PRIVATE_CASE_TYPE_USING_HINT(enum_type, fp8_t, __VA_ARGS__)
13 +
14 #define VLLM_DISPATCH_CASE_FLOATING_TYPES(...) \
15 AT_DISPATCH_CASE(at::ScalarType::Float, __VA_ARGS__) \
16 AT_DISPATCH_CASE(at::ScalarType::Half, __VA_ARGS__) \
@@ -14,17 +19,32 @@
19 #define VLLM_DISPATCH_FLOATING_TYPES(TYPE, NAME, ...) \
20 AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FLOATING_TYPES(__VA_ARGS__))
21
17 -// TODO(luka/varun): use FP8_TYPE macro after refactoring
18 -#ifndef USE_ROCM
19 - #define VLLM_DISPATCH_CASE_QUANT_TYPES(...) \
20 - AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \
21 - AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__)
22 -#else
22 +// ROCm devices might use either fn or fnuz, so set up dispatch table for both.
23 +// A host-based check at runtime will create a preferred FP8 type for ROCm
24 +// such that the correct kernel is dispatched.
25 +#ifdef USE_ROCM
26 + #define VLLM_DISPATCH_CASE_FP8_TYPES(...) \
27 + AT_DISPATCH_FP8_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \
28 + AT_DISPATCH_FP8_CASE(at::ScalarType::Float8_e4m3fnuz, __VA_ARGS__)
29 +
30 #define VLLM_DISPATCH_CASE_QUANT_TYPES(...) \
31 + AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \
32 AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fnuz, __VA_ARGS__) \
33 AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__)
34 +#else
35 + #define VLLM_DISPATCH_CASE_FP8_TYPES(...) \
36 + AT_DISPATCH_FP8_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__)
37 +
38 + #define VLLM_DISPATCH_CASE_QUANT_TYPES(...) \
39 + AT_DISPATCH_CASE(at::ScalarType::Float8_e4m3fn, __VA_ARGS__) \
40 + AT_DISPATCH_CASE(at::ScalarType::Char, __VA_ARGS__)
41 #endif
42
43 +// When using this dispatch macro, the type is 'fp8_t' not 'scalar_t'.
44 +// See AT_DISPATCH_FP8_CASE above.
45 +#define VLLM_DISPATCH_FP8_TYPES(TYPE, NAME, ...) \
46 + AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_FP8_TYPES(__VA_ARGS__))
47 +
48 #define VLLM_DISPATCH_QUANT_TYPES(TYPE, NAME, ...) \
49 AT_DISPATCH_SWITCH(TYPE, NAME, VLLM_DISPATCH_CASE_QUANT_TYPES(__VA_ARGS__))
50
csrc/layernorm_quant_kernels.cu
+33 -25
@@ -21,9 +21,9 @@
21 namespace vllm {
22
23 // TODO(woosuk): Further optimize this kernel.
24 -template <typename scalar_t>
24 +template <typename scalar_t, typename fp8_type>
25 __global__ void rms_norm_static_fp8_quant_kernel(
26 - FP8_TYPE* __restrict__ out, // [..., hidden_size]
26 + fp8_type* __restrict__ out, // [..., hidden_size]
27 const scalar_t* __restrict__ input, // [..., hidden_size]
28 const scalar_t* __restrict__ weight, // [hidden_size]
29 const float* __restrict__ scale, // [1]
@@ -52,7 +52,7 @@ __global__ void rms_norm_static_fp8_quant_kernel(
52 float x = (float)input[blockIdx.x * hidden_size + idx];
53 float const out_norm = ((scalar_t)(x * s_variance)) * weight[idx];
54 out[blockIdx.x * hidden_size + idx] =
55 - scaled_fp8_conversion<true>(out_norm, scale_inv);
55 + scaled_fp8_conversion<true, fp8_type>(out_norm, scale_inv);
56 }
57 }
58
@@ -60,10 +60,10 @@ __global__ void rms_norm_static_fp8_quant_kernel(
60 Additional optimizations we can make in this case are
61 packed and vectorized operations, which help with the
62 memory latency bottleneck. */
63 -template <typename scalar_t, int width>
63 +template <typename scalar_t, int width, typename fp8_type>
64 __global__ std::enable_if_t<(width > 0) && _typeConvert<scalar_t>::exists>
65 fused_add_rms_norm_static_fp8_quant_kernel(
66 - FP8_TYPE* __restrict__ out, // [..., hidden_size]
66 + fp8_type* __restrict__ out, // [..., hidden_size]
67 scalar_t* __restrict__ input, // [..., hidden_size]
68 scalar_t* __restrict__ residual, // [..., hidden_size]
69 const scalar_t* __restrict__ weight, // [hidden_size]
@@ -114,7 +114,7 @@ fused_add_rms_norm_static_fp8_quant_kernel(
114 #pragma unroll
115 for (int i = 0; i < width; ++i) {
116 out[id * width + i] =
117 - scaled_fp8_conversion<true>(float(temp.data[i]), scale_inv);
117 + scaled_fp8_conversion<true, fp8_type>(float(temp.data[i]), scale_inv);
118 }
119 }
120 }
@@ -122,10 +122,10 @@ fused_add_rms_norm_static_fp8_quant_kernel(
122 /* Generic fused_add_rms_norm_kernel
123 The width field is not used here but necessary for other specializations.
124 */
125 -template <typename scalar_t, int width>
125 +template <typename scalar_t, int width, typename fp8_type>
126 __global__ std::enable_if_t<(width == 0) || !_typeConvert<scalar_t>::exists>
127 fused_add_rms_norm_static_fp8_quant_kernel(
128 - FP8_TYPE* __restrict__ out, // [..., hidden_size]
128 + fp8_type* __restrict__ out, // [..., hidden_size]
129 scalar_t* __restrict__ input, // [..., hidden_size]
130 scalar_t* __restrict__ residual, // [..., hidden_size]
131 const scalar_t* __restrict__ weight, // [hidden_size]
@@ -158,7 +158,7 @@ fused_add_rms_norm_static_fp8_quant_kernel(
158 float x = (float)residual[blockIdx.x * hidden_size + idx];
159 float const out_norm = ((scalar_t)(x * s_variance)) * weight[idx];
160 out[blockIdx.x * hidden_size + idx] =
161 - scaled_fp8_conversion<true>(out_norm, scale_inv);
161 + scaled_fp8_conversion<true, fp8_type>(out_norm, scale_inv);
162 }
163 }
164
@@ -176,25 +176,33 @@ void rms_norm_static_fp8_quant(torch::Tensor& out, // [..., hidden_size]
176 dim3 block(std::min(hidden_size, 1024));
177 const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
178 const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
179 - VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "rms_norm_kernel", [&] {
180 - vllm::rms_norm_static_fp8_quant_kernel<scalar_t>
181 - <<<grid, block, 0, stream>>>(
182 - out.data_ptr<FP8_TYPE>(), input.data_ptr<scalar_t>(),
183 - weight.data_ptr<scalar_t>(), scale.data_ptr<float>(), epsilon,
184 - num_tokens, hidden_size);
185 - });
179 + VLLM_DISPATCH_FLOATING_TYPES(
180 + input.scalar_type(), "rms_norm_kernel_scalar_type", [&] {
181 + VLLM_DISPATCH_FP8_TYPES(
182 + out.scalar_type(), "rms_norm_kernel_fp8_type", [&] {
183 + vllm::rms_norm_static_fp8_quant_kernel<scalar_t, fp8_t>
184 + <<<grid, block, 0, stream>>>(
185 + out.data_ptr<fp8_t>(), input.data_ptr<scalar_t>(),
186 + weight.data_ptr<scalar_t>(), scale.data_ptr<float>(),
187 + epsilon, num_tokens, hidden_size);
188 + });
189 + });
190 }
191
188 -#define LAUNCH_FUSED_ADD_RMS_NORM(width) \
189 - VLLM_DISPATCH_FLOATING_TYPES( \
190 - input.scalar_type(), "fused_add_rms_norm_kernel", [&] { \
191 - vllm::fused_add_rms_norm_static_fp8_quant_kernel<scalar_t, width> \
192 - <<<grid, block, 0, stream>>>( \
193 - out.data_ptr<FP8_TYPE>(), input.data_ptr<scalar_t>(), \
194 - residual.data_ptr<scalar_t>(), weight.data_ptr<scalar_t>(), \
195 - scale.data_ptr<float>(), epsilon, num_tokens, hidden_size); \
192 +#define LAUNCH_FUSED_ADD_RMS_NORM(width) \
193 + VLLM_DISPATCH_FLOATING_TYPES( \
194 + input.scalar_type(), "fused_add_rms_norm_kernel_scalar_type", [&] { \
195 + VLLM_DISPATCH_FP8_TYPES( \
196 + out.scalar_type(), "fused_add_rms_norm_kernel_fp8_type", [&] { \
197 + vllm::fused_add_rms_norm_static_fp8_quant_kernel<scalar_t, \
198 + width, fp8_t> \
199 + <<<grid, block, 0, stream>>>( \
200 + out.data_ptr<fp8_t>(), input.data_ptr<scalar_t>(), \
201 + residual.data_ptr<scalar_t>(), \
202 + weight.data_ptr<scalar_t>(), scale.data_ptr<float>(), \
203 + epsilon, num_tokens, hidden_size); \
204 + }); \
205 });
197 -
206 void fused_add_rms_norm_static_fp8_quant(
207 torch::Tensor& out, // [..., hidden_size],
208 torch::Tensor& input, // [..., hidden_size]
csrc/quantization/fp8/amd/quant_utils.cuh
+22
@@ -13,6 +13,28 @@ namespace vllm {
13 namespace fp8 {
14 #ifdef ENABLE_FP8
15
16 +// Use hardware cvt instruction for fp8 on rocm
17 +template <typename fp8_type>
18 +__device__ __forceinline__ fp8_type cvt_c10(float const r) {
19 + return {};
20 +}
21 +
22 +template <>
23 +__device__ __forceinline__ c10::Float8_e4m3fn cvt_c10(float const r) {
24 + return c10::Float8_e4m3fn(
25 + __hip_cvt_float_to_fp8(r, __hip_fp8_e4m3::__default_saturation,
26 + __hip_fp8_e4m3::__default_interpret),
27 + c10::Float8_e4m3fn::from_bits());
28 +}
29 +
30 +template <>
31 +__device__ __forceinline__ c10::Float8_e4m3fnuz cvt_c10(float const r) {
32 + return c10::Float8_e4m3fnuz(
33 + __hip_cvt_float_to_fp8(r, __hip_fp8_e4m3_fnuz::__default_saturation,
34 + __hip_fp8_e4m3_fnuz::__default_interpret),
35 + c10::Float8_e4m3fnuz::from_bits());
36 +}
37 +
38 template <typename Tout, typename Tin>
39 __inline__ __device__ Tout vec_conversion(const Tin& x) {
40 return x;
csrc/quantization/fp8/common.cu
+43 -25
@@ -11,8 +11,8 @@
11
12 namespace vllm {
13
14 -template <typename scalar_t>
15 -__global__ void scaled_fp8_quant_kernel(FP8_TYPE* __restrict__ out,
14 +template <typename scalar_t, typename fp8_type>
15 +__global__ void scaled_fp8_quant_kernel(fp8_type* __restrict__ out,
16 const scalar_t* __restrict__ input,
17 const float* __restrict__ scale,
18 int64_t num_elems) {
@@ -25,12 +25,13 @@ __global__ void scaled_fp8_quant_kernel(FP8_TYPE* __restrict__ out,
25 out, input, inverted_scale, num_elems, tid, blockDim.x * gridDim.x);
26 }
27
28 -template <typename scalar_t>
28 +template <typename scalar_t, typename fp8_type>
29 __global__ void dynamic_per_token_scaled_fp8_quant_kernel(
30 - FP8_TYPE* __restrict__ out, float* __restrict__ scale,
30 + fp8_type* __restrict__ out, float* __restrict__ scale,
31 scalar_t const* __restrict__ input, float const* __restrict__ scale_ub,
32 const int hidden_size) {
33 - float const min_scaling_factor = 1.0f / (FP8_E4M3_MAX * 512.f);
33 + float const min_scaling_factor =
34 + 1.0f / (fp8_e4m3_adjusted_max_v<fp8_type> * 512.f);
35
36 int const tid = threadIdx.x;
37 int const token_idx = blockIdx.x;
@@ -38,7 +39,7 @@ __global__ void dynamic_per_token_scaled_fp8_quant_kernel(
39 // Use int64 to avoid overflowing an int32 when calculating this offset
40 int64_t offset = static_cast<int64_t>(token_idx) * hidden_size;
41 scalar_t const* __restrict__ token_input = &input[offset];
41 - FP8_TYPE* __restrict__ token_output = &out[offset];
42 + fp8_type* __restrict__ token_output = &out[offset];
43
44 // For vectorization, token_input and token_output pointers need to be
45 // aligned at 8-byte and 4-byte addresses respectively.
@@ -66,7 +67,8 @@ __global__ void dynamic_per_token_scaled_fp8_quant_kernel(
67 token_scale = block_absmax_val_maybe;
68 }
69 // token scale computation
69 - token_scale = max(token_scale / FP8_E4M3_MAX, min_scaling_factor);
70 + token_scale = max(token_scale / fp8_e4m3_adjusted_max_v<fp8_type>,
71 + min_scaling_factor);
72 scale[token_idx] = token_scale;
73 }
74 __syncthreads();
@@ -77,7 +79,7 @@ __global__ void dynamic_per_token_scaled_fp8_quant_kernel(
79 token_output, token_input, token_scale, hidden_size, tid, blockDim.x);
80 } else {
81 for (int i = tid; i < hidden_size; i += blockDim.x) {
80 - token_output[i] = scaled_fp8_conversion<false>(
82 + token_output[i] = scaled_fp8_conversion<false, fp8_type>(
83 static_cast<float>(token_input[i]), token_scale);
84 }
85 }
@@ -96,10 +98,14 @@ void static_scaled_fp8_quant(torch::Tensor& out, // [..., d]
98 const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
99 const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
100 VLLM_DISPATCH_FLOATING_TYPES(
99 - input.scalar_type(), "scaled_fp8_quant_kernel", [&] {
100 - vllm::scaled_fp8_quant_kernel<scalar_t><<<grid, block, 0, stream>>>(
101 - out.data_ptr<FP8_TYPE>(), input.data_ptr<scalar_t>(),
102 - scale.data_ptr<float>(), num_elems);
101 + input.scalar_type(), "scaled_fp8_quant_kernel_scalar_type", [&] {
102 + VLLM_DISPATCH_FP8_TYPES(
103 + out.scalar_type(), "scaled_fp8_quant_kernel_fp8_type", [&] {
104 + vllm::scaled_fp8_quant_kernel<scalar_t, fp8_t>
105 + <<<grid, block, 0, stream>>>(
106 + out.data_ptr<fp8_t>(), input.data_ptr<scalar_t>(),
107 + scale.data_ptr<float>(), num_elems);
108 + });
109 });
110 }
111
@@ -114,12 +120,18 @@ void dynamic_scaled_fp8_quant(torch::Tensor& out, // [..., d]
120 const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
121 const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
122 VLLM_DISPATCH_FLOATING_TYPES(
117 - input.scalar_type(), "scaled_fp8_quant_kernel", [&] {
118 - vllm::segmented_max_reduction<scalar_t><<<grid, block, 0, stream>>>(
119 - scale.data_ptr<float>(), input.data_ptr<scalar_t>(), num_elems);
120 - vllm::scaled_fp8_quant_kernel<scalar_t><<<grid, block, 0, stream>>>(
121 - out.data_ptr<FP8_TYPE>(), input.data_ptr<scalar_t>(),
122 - scale.data_ptr<float>(), num_elems);
123 + input.scalar_type(), "scaled_fp8_quant_kernel_scalar_type", [&] {
124 + VLLM_DISPATCH_FP8_TYPES(
125 + out.scalar_type(), "scaled_fp8_quant_kernel_fp8_type", [&] {
126 + vllm::segmented_max_reduction<scalar_t, fp8_t>
127 + <<<grid, block, 0, stream>>>(scale.data_ptr<float>(),
128 + input.data_ptr<scalar_t>(),
129 + num_elems);
130 + vllm::scaled_fp8_quant_kernel<scalar_t, fp8_t>
131 + <<<grid, block, 0, stream>>>(
132 + out.data_ptr<fp8_t>(), input.data_ptr<scalar_t>(),
133 + scale.data_ptr<float>(), num_elems);
134 + });
135 });
136 }
137
@@ -138,12 +150,18 @@ void dynamic_per_token_scaled_fp8_quant(
150 const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
151 const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
152 VLLM_DISPATCH_FLOATING_TYPES(
141 - input.scalar_type(), "dynamic_per_token_scaled_fp8_quant_kernel", [&] {
142 - vllm::dynamic_per_token_scaled_fp8_quant_kernel<scalar_t>
143 - <<<grid, block, 0, stream>>>(
144 - out.data_ptr<FP8_TYPE>(), scales.data_ptr<float>(),
145 - input.data_ptr<scalar_t>(),
146 - scale_ub.has_value() ? scale_ub->data_ptr<float>() : nullptr,
147 - hidden_size);
153 + input.scalar_type(),
154 + "dynamic_per_token_scaled_fp8_quant_kernel_scalar_type", [&] {
155 + VLLM_DISPATCH_FP8_TYPES(
156 + out.scalar_type(),
157 + "dynamic_per_token_scaled_fp8_quant_kernel_fp8_type", [&] {
158 + vllm::dynamic_per_token_scaled_fp8_quant_kernel<scalar_t, fp8_t>
159 + <<<grid, block, 0, stream>>>(
160 + out.data_ptr<fp8_t>(), scales.data_ptr<float>(),
161 + input.data_ptr<scalar_t>(),
162 + scale_ub.has_value() ? scale_ub->data_ptr<float>()
163 + : nullptr,
164 + hidden_size);
165 + });
166 });
167 }
csrc/quantization/fp8/common.cuh
+59 -27
@@ -7,18 +7,52 @@
7
8 #ifndef USE_ROCM
9 #include <c10/util/Float8_e4m3fn.h>
10 -using FP8_TYPE = c10::Float8_e4m3fn;
11 -C10_HOST_DEVICE constexpr auto FP8_E4M3_MAX =
12 - std::numeric_limits<FP8_TYPE>::max();
10 + #define MAYBE_HOST_DEVICE C10_HOST_DEVICE
11 #else
12 + #include <ATen/hip/HIPContext.h>
13 + #include <c10/util/Float8_e4m3fn.h>
14 #include <c10/util/Float8_e4m3fnuz.h>
15 #include "amd/quant_utils.cuh"
16 -using FP8_TYPE = c10::Float8_e4m3fnuz;
17 -// Using the default max value from pytorch (240.0) will cause accuracy
18 -// issue when running dynamic quantization. Here use 224.0f for rocm.
19 -constexpr auto FP8_E4M3_MAX = 224.0f;
16 + // ROCm doesn't seem to need C10_HOST_DEVICE for static constexpr
17 + #define MAYBE_HOST_DEVICE
18 +#endif
19 +
20 +// Determines the preferred FP8 type for the current platform.
21 +// Note that for CUDA this just returns true,
22 +// but on ROCm it will check device props.
23 +static bool is_fp8_ocp() {
24 +#ifndef USE_ROCM
25 + return true;
26 +#else
27 + auto dprops = at::cuda::getCurrentDeviceProperties();
28 + std::string device_arch = dprops->gcnArchName;
29 + size_t substring = device_arch.find("gfx94");
30 + return substring == std::string::npos;
31 #endif
21 -constexpr static auto kFp8Type = c10::CppTypeToScalarType<FP8_TYPE>::value;
32 +}
33 +
34 +template <typename T>
35 +struct fp8_e4m3_adjusted_max;
36 +
37 +template <>
38 +struct fp8_e4m3_adjusted_max<c10::Float8_e4m3fn> {
39 + static constexpr c10::Float8_e4m3fn val() {
40 + return std::numeric_limits<c10::Float8_e4m3fn>::max();
41 + }
42 +};
43 +
44 +// Using the default max value from pytorch (240.0 0x7F) will cause accuracy
45 +// issues when running dynamic quantization. Here use 224.0 0x7E for rocm.
46 +template <>
47 +struct fp8_e4m3_adjusted_max<c10::Float8_e4m3fnuz> {
48 + static constexpr c10::Float8_e4m3fnuz val() {
49 + return c10::Float8_e4m3fnuz(0x7E, c10::Float8_e4m3fnuz::from_bits());
50 + }
51 +};
52 +
53 +template <typename T>
54 +MAYBE_HOST_DEVICE static constexpr T fp8_e4m3_adjusted_max_v =
55 + fp8_e4m3_adjusted_max<T>::val();
56
57 namespace vllm {
58
@@ -32,8 +66,8 @@ __device__ __forceinline__ float atomicMaxFloat(float* addr, float value) {
66 return old;
67 }
68
35 -template <bool is_scale_inverted>
36 -__device__ __forceinline__ FP8_TYPE scaled_fp8_conversion(float const val,
69 +template <bool is_scale_inverted, typename fp8_type>
70 +__device__ __forceinline__ fp8_type scaled_fp8_conversion(float const val,
71 float const scale) {
72 float x = 0.0f;
73 if constexpr (is_scale_inverted) {
@@ -42,15 +76,13 @@ __device__ __forceinline__ FP8_TYPE scaled_fp8_conversion(float const val,
76 x = val / scale;
77 }
78
45 - float r = fmax(-FP8_E4M3_MAX, fmin(x, FP8_E4M3_MAX));
79 + float r = fmax(-fp8_e4m3_adjusted_max_v<fp8_type>,
80 + fmin(x, fp8_e4m3_adjusted_max_v<fp8_type>));
81 #ifndef USE_ROCM
47 - return static_cast<c10::Float8_e4m3fn>(r);
82 + return static_cast<fp8_type>(r);
83 #else
84 // Use hardware cvt instruction for fp8 on rocm
50 - return c10::Float8_e4m3fnuz(
51 - __hip_cvt_float_to_fp8(r, fp8::fp8_type::__default_saturation,
52 - fp8::fp8_type::__default_interpret),
53 - c10::Float8_e4m3fnuz::from_bits());
85 + return fp8::cvt_c10<fp8_type>(r);
86 #endif
87 }
88
@@ -60,7 +92,7 @@ __device__ __forceinline__ FP8_TYPE scaled_fp8_conversion(float const val,
92 // So to get the right answer, *scale needs to be initialized to
93 // a value <= 0.0 and we need to wait for all thread blocks to
94 // finish before consuming *scale.
63 -template <typename scalar_t>
95 +template <typename scalar_t, typename fp8_type>
96 __global__ void segmented_max_reduction(float* __restrict__ scale,
97 const scalar_t* __restrict__ input,
98 int64_t num_elems) {
@@ -91,7 +123,7 @@ __global__ void segmented_max_reduction(float* __restrict__ scale,
123 // Finally, since cache[0] contains the maximum for this thread block,
124 // atomically write the max to the target location
125 if (threadIdx.x == 0) {
94 - atomicMaxFloat(scale, cache[0] / FP8_E4M3_MAX);
126 + atomicMaxFloat(scale, cache[0] / fp8_e4m3_adjusted_max_v<fp8_type>);
127 }
128 }
129
@@ -123,13 +155,13 @@ __device__ float thread_max_vec(scalar_t const* __restrict__ input,
155 return absmax_val;
156 }
157
126 -template <typename scalar_t, bool is_scale_inverted>
127 -__device__ void scaled_fp8_conversion_vec(FP8_TYPE* __restrict__ out,
158 +template <typename scalar_t, bool is_scale_inverted, typename fp8_type>
159 +__device__ void scaled_fp8_conversion_vec(fp8_type* __restrict__ out,
160 scalar_t const* __restrict__ input,
161 float const scale,
162 int64_t const num_elems,
163 int const tid, int const step) {
132 - using float8x4_t = q8x4_t<FP8_TYPE>;
164 + using float8x4_t = q8x4_t<fp8_type>;
165 // Vectorized input/output to better utilize memory bandwidth.
166 auto const* vectorized_in = reinterpret_cast<vec4_t<scalar_t> const*>(input);
167 auto* vectorized_out = reinterpret_cast<float8x4_t*>(out);
@@ -141,22 +173,22 @@ __device__ void scaled_fp8_conversion_vec(FP8_TYPE* __restrict__ out,
173 vec4_t<scalar_t> in_vec = vectorized_in[i];
174 float8x4_t out_vec;
175
144 - out_vec.x = scaled_fp8_conversion<is_scale_inverted>(
176 + out_vec.x = scaled_fp8_conversion<is_scale_inverted, fp8_type>(
177 static_cast<float>(in_vec.x), scale);
146 - out_vec.y = scaled_fp8_conversion<is_scale_inverted>(
178 + out_vec.y = scaled_fp8_conversion<is_scale_inverted, fp8_type>(
179 static_cast<float>(in_vec.y), scale);
148 - out_vec.z = scaled_fp8_conversion<is_scale_inverted>(
180 + out_vec.z = scaled_fp8_conversion<is_scale_inverted, fp8_type>(
181 static_cast<float>(in_vec.z), scale);
150 - out_vec.w = scaled_fp8_conversion<is_scale_inverted>(
182 + out_vec.w = scaled_fp8_conversion<is_scale_inverted, fp8_type>(
183 static_cast<float>(in_vec.w), scale);
184 vectorized_out[i] = out_vec;
185 }
186
187 // Handle the remaining elements if num_elems is not divisible by 4
188 for (int64_t i = num_vec_elems * 4 + tid; i < num_elems; i += step) {
157 - out[i] = scaled_fp8_conversion<is_scale_inverted>(
189 + out[i] = scaled_fp8_conversion<is_scale_inverted, fp8_type>(
190 static_cast<float>(input[i]), scale);
191 }
192 }
193
162 -} // namespace vllm
\ No newline at end of file
194 +} // namespace vllm
csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu
+3
@@ -144,6 +144,9 @@ void rms_norm_dynamic_per_token_quant(
144 torch::Tensor& scales, // [num_tokens]
145 double const var_epsilon, // Variance epsilon used in norm calculation
146 std::optional<at::Tensor> scale_ub, std::optional<at::Tensor> residual) {
147 + static c10::ScalarType kFp8Type = is_fp8_ocp()
148 + ? c10::ScalarType::Float8_e4m3fn
149 + : c10::ScalarType::Float8_e4m3fnuz;
150 TORCH_CHECK(out.dtype() == kFp8Type || out.dtype() == torch::kInt8);
151 TORCH_CHECK(out.is_contiguous() && input.is_contiguous());
152
csrc/quantization/fused_kernels/quant_conversions.cuh
+11 -8
@@ -31,9 +31,11 @@ static __device__ __forceinline__ int8_t float_to_int8_rn(float const x) {
31 #endif
32 }
33
34 -static __device__ __forceinline__ FP8_TYPE float_to_fp8(float const x) {
35 - float const r = fmax(-FP8_E4M3_MAX, fmin(x, FP8_E4M3_MAX));
36 - return static_cast<FP8_TYPE>(r);
34 +template <typename fp8_type>
35 +static __device__ __forceinline__ fp8_type float_to_fp8(float const x) {
36 + float const r = fmax(-fp8_e4m3_adjusted_max_v<fp8_type>,
37 + fmin(x, fp8_e4m3_adjusted_max_v<fp8_type>));
38 + return static_cast<fp8_type>(r);
39 }
40
41 template <typename quant_type_t, bool is_scale_inverted, typename enable = void>
@@ -54,15 +56,16 @@ struct ScaledQuant<
56 };
57
58 template <typename quant_type_t, bool is_scale_inverted>
57 -struct ScaledQuant<
58 - quant_type_t, is_scale_inverted,
59 - typename std::enable_if_t<std::is_same_v<quant_type_t, FP8_TYPE>>> {
59 +struct ScaledQuant<quant_type_t, is_scale_inverted,
60 + typename std::enable_if_t<
61 + std::is_same_v<quant_type_t, c10::Float8_e4m3fn> ||
62 + std::is_same_v<quant_type_t, c10::Float8_e4m3fnuz>>> {
63 static __device__ __forceinline__ quant_type_t quant_fn(float const x,
64 float const scale) {
65 if constexpr (is_scale_inverted) {
63 - return float_to_fp8(x * scale);
66 + return float_to_fp8<quant_type_t>(x * scale);
67 } else {
65 - return float_to_fp8(x / scale);
68 + return float_to_fp8<quant_type_t>(x / scale);
69 }
70 }
71 };
csrc/quantization/vectorization.cuh
-1
@@ -4,7 +4,6 @@
4 */
5
6 // Include both AMD and NVIDIA fp8 types to avoid circular import
7 -// TODO(luka/varun) use FP8_TYPE instead after refactoring
7 #include <c10/util/Float8_e4m3fnuz.h>
8 #include <c10/util/Float8_e4m3fn.h>
9
tests/kernels/quant_utils.py
+1 -2
@@ -9,8 +9,7 @@ from vllm.platforms import current_platform
9 # Using the default value (240.0) from pytorch will cause accuracy
10 # issue on dynamic quantization models. Here use 224.0 for rocm.
11 ROCM_FP8_MAX = 224.0
12 -FP8_DTYPE = torch.float8_e4m3fnuz if current_platform.is_rocm() \
13 - else torch.float8_e4m3fn
12 +FP8_DTYPE = current_platform.fp8_dtype()
13
14
15 def as_float32_tensor(x: Union[float, torch.tensor]) -> torch.tensor:
tests/kernels/test_triton_scaled_mm.py
+2 -5
@@ -32,11 +32,8 @@ def scaled_mm_torch(a: torch.Tensor,
32
33 def get_8bit_types():
34 types = [torch.int8]
35 - supports_fp8 = current_platform.has_device_capability(89)
36 - if current_platform.is_rocm() and supports_fp8:
37 - types.append(torch.float8_e4m3fnuz)
38 - elif current_platform.is_cuda() and supports_fp8:
39 - types.append(torch.float8_e4m3fn)
35 + if current_platform.supports_fp8():
36 + types.append(current_platform.fp8_dtype())
37 return types
38
39
tests/quantization/test_fp8.py
+3 -6
@@ -103,8 +103,7 @@ def test_load_fp16_model(vllm_runner, kv_cache_dtype: str, force_marlin: bool,
103 assert attn._v_scale == 1.0
104
105 if current_platform.is_cuda():
106 - if current_platform.has_device_capability(
107 - 89) and not force_marlin:
106 + if current_platform.supports_fp8() and not force_marlin:
107 # For GPUs with hardware support, we keep weights in fp8
108 assert fc1.weight.dtype == torch.float8_e4m3fn
109 else:
@@ -112,11 +111,9 @@ def test_load_fp16_model(vllm_runner, kv_cache_dtype: str, force_marlin: bool,
111 # for weight-only quantization using Marlin kernels
112 assert fc1.weight.dtype == torch.int32
113 elif current_platform.is_rocm():
115 - # Only MI300 and above support quantization='fp8'
116 - if current_platform.has_device_capability(
117 - 94) and not force_marlin:
114 + if current_platform.supports_fp8() and not force_marlin:
115 # For GPUs with hardware support, we keep weights in fp8
119 - assert fc1.weight.dtype == torch.float8_e4m3fnuz
116 + assert fc1.weight.dtype == current_platform.fp8_dtype()
117 else: # unsupported ROCm platform
118 pytest.skip(
119 "Skip `test_load_fp16_model`. "
vllm/_custom_ops.py
+17 -18
@@ -478,16 +478,16 @@ def cutlass_scaled_mm(a: torch.Tensor,
478 out_dtype: torch.dtype,
479 bias: Optional[torch.Tensor] = None) -> torch.Tensor:
480 """
481 - `cutlass_scaled_mm` implements a fused version of
481 + `cutlass_scaled_mm` implements a fused version of
482 `output = torch.mm((scale_a * a), (scale_b * b)).to(out_dtype)`
483 - where scale_a * a and scale_b * b are implemented using numpy-style
484 - broadcasting.
485 -
486 - In order to support blockwise scaling like found in DeepSeek V3 we also
487 - support extended "group" broadcast rules. We extend the numpy-style
488 - broadcasting rules with the following rule:
489 - "if the extent of a dimension in the source shape is between 1 and
490 - corresponding extent in the target shape we repeat each element along
483 + where scale_a * a and scale_b * b are implemented using numpy-style
484 + broadcasting.
485 +
486 + In order to support blockwise scaling like found in DeepSeek V3 we also
487 + support extended "group" broadcast rules. We extend the numpy-style
488 + broadcasting rules with the following rule:
489 + "if the extent of a dimension in the source shape is between 1 and
490 + corresponding extent in the target shape we repeat each element along
491 that dimension src_shape[dim] // target_shape[dim] times consecutively"
492 example if we have:
493 a = [[1, 2], and target_shape = (2, 4)
@@ -564,7 +564,7 @@ def cutlass_sparse_compress(a: torch.Tensor) \
564 with Cutlass sparse kernels.
565
566 Args:
567 - a (torch.Tensor):
567 + a (torch.Tensor):
568 The input tensor to be compressed. Must have one of the following data types:
569 - `torch.int8`
570 - `torch.float8_e4m3fn`
@@ -572,7 +572,7 @@ def cutlass_sparse_compress(a: torch.Tensor) \
572 - `torch.float16`
573
574 Returns:
575 - tuple[torch.Tensor, torch.Tensor]:
575 + tuple[torch.Tensor, torch.Tensor]:
576 A tuple containing:
577 - `a_nzs` (torch.Tensor): A tensor containing non-zero elements of `a`.
578 - `a_meta` (torch.Tensor): A tensor containing metadata for the sparse representation.
@@ -875,9 +875,8 @@ def scaled_fp8_quant(
875 # This code assumes batch_dim and num_tokens are flattened
876 assert (input.ndim == 2)
877 shape: Union[tuple[int, int], torch.Size] = input.shape
878 - # For rocm, the output fp8 dtype is torch.float_e3m3fnuz
879 - out_dtype: torch.dtype = torch.float8_e4m3fnuz \
880 - if current_platform.is_rocm() else torch.float8_e4m3fn
878 + # For ROCm on MI300, the output fp8 dtype is torch.float_e3m3fnuz
879 + out_dtype: torch.dtype = current_platform.fp8_dtype()
880 if num_token_padding:
881 shape = (max(num_token_padding, input.shape[0]), shape[1])
882 output = torch.empty(shape, device=input.device, dtype=out_dtype)
@@ -908,7 +907,7 @@ def allspark_repack_weight(
907 has_zp: bool = False
908 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
909 """
911 - Rearrange qweight, scale, and zero_point(if asymmetric) to n32k16 format
910 + Rearrange qweight, scale, and zero_point(if asymmetric) to n32k16 format
911 for Ampere W8A16 Fused Gemm kernel
912
913 Args:
@@ -917,10 +916,10 @@ def allspark_repack_weight(
916 zero_point: fp16/bf16 weight zero_point tensor, 1 x n format.
917 Must be provided for asymmetric quantization.
918 has_zp: if use symmetric quantization, has_zp = False.
920 - if use asymmetric quantization, has_zp = True.
921 -
919 + if use asymmetric quantization, has_zp = True.
920 +
921 Returns:
923 - tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]] :
922 + tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]] :
923 rearranged weight, scale, and optionally zero_point.
924 """
925 K = qweight.shape[0]
vllm/attention/backends/mla/common.py
+3 -3
@@ -226,7 +226,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
226 CompressedTensorsW8A8Fp8)
227 from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
228 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
229 - Fp8LinearGenericOp, current_platform_fp8_dtype, is_fp8)
229 + Fp8LinearGenericOp, is_fp8)
230 from vllm.model_executor.layers.quantization.utils.quant_utils import (
231 scaled_quantize)
232 from vllm.model_executor.layers.rotary_embedding import (
@@ -1238,7 +1238,7 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
1238 W_Q_UK, W_Q_UK_scales = scaled_quantize(
1239 W_Q_UK,
1240 self.reqaunt_weight_group_shape,
1241 - quant_dtype=current_platform_fp8_dtype)
1241 + quant_dtype=current_platform.fp8_dtype())
1242 # For FP8 save the transpose so we can use
1243 # `apply_w8a8_block_fp8_linear` directly
1244 self.W_Q_UK = W_Q_UK.T.contiguous()
@@ -1255,7 +1255,7 @@ class MLACommonImpl(MLAAttentionImpl[T], Generic[T]):
1255 W_UV_O, W_UV_O_scales = scaled_quantize(
1256 W_UV_O,
1257 self.reqaunt_weight_group_shape,
1258 - quant_dtype=current_platform_fp8_dtype)
1258 + quant_dtype=current_platform.fp8_dtype())
1259 # For FP8 save the transpose so we can use
1260 # `apply_w8a8_block_fp8_linear` directly
1261 self.W_UV_O = W_UV_O.T.contiguous()
vllm/model_executor/layers/quantization/compressed_tensors/compressed_tensors_moe.py
+1 -2
@@ -158,8 +158,7 @@ class CompressedTensorsW8A8Fp8MoEMethod(CompressedTensorsMoEMethod):
158 layer.w2_input_scale = torch.nn.Parameter(
159 layer.w2_input_scale.max(), requires_grad=False)
160
161 - # If rocm, normalize the weights and scales to e4m3fnuz
162 - if current_platform.is_rocm():
161 + if current_platform.is_fp8_fnuz():
162 # Normalize the weights and scales
163 w13_weight, w13_weight_scale, w13_input_scale = \
164 normalize_e4m3fn_to_e4m3fnuz(
vllm/model_executor/layers/quantization/compressed_tensors/schemes/compressed_tensors_w8a8_fp8.py
+2 -2
@@ -42,7 +42,7 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):
42 logical_widths=layer.logical_widths,
43 )
44
45 - if current_platform.is_rocm():
45 + if current_platform.is_fp8_fnuz():
46 input_scale = getattr(layer, 'input_scale', None)
47
48 weight, max_w_scale, input_scale = normalize_e4m3fn_to_e4m3fnuz(
@@ -60,7 +60,7 @@ class CompressedTensorsW8A8Fp8(CompressedTensorsScheme):
60 elif self.strategy == QuantizationStrategy.CHANNEL:
61 weight = layer.weight
62
63 - if current_platform.is_rocm():
63 + if current_platform.is_fp8_fnuz():
64 input_scale = getattr(layer, 'input_scale', None)
65
66 weight, weight_scale, input_scale = \
vllm/model_executor/layers/quantization/fbgemm_fp8.py
+1 -1
@@ -127,7 +127,7 @@ class FBGEMMFp8LinearMethod(LinearMethodBase):
127
128 weight = layer.weight
129
130 - if current_platform.is_rocm():
130 + if current_platform.is_fp8_fnuz():
131 weight, weight_scale, input_scale = \
132 normalize_e4m3fn_to_e4m3fnuz(
133 weight=weight,
vllm/model_executor/layers/quantization/fp8.py
+5 -9
@@ -270,7 +270,7 @@ class Fp8LinearMethod(LinearMethodBase):
270 # TODO(rob): refactor block quant into separate class.
271 if self.block_quant:
272 assert self.quant_config.activation_scheme == "dynamic"
273 - if current_platform.is_rocm():
273 + if current_platform.is_fp8_fnuz():
274 weight, weight_scale_inv, _ = \
275 normalize_e4m3fn_to_e4m3fnuz(
276 weight=layer.weight,
@@ -327,8 +327,7 @@ class Fp8LinearMethod(LinearMethodBase):
327 weight = layer.weight
328 weight_scale = layer.weight_scale
329
330 - # If rocm, use float8_e4m3fnuz.
331 - if current_platform.is_rocm():
330 + if current_platform.is_fp8_fnuz():
331 weight, weight_scale, input_scale = \
332 normalize_e4m3fn_to_e4m3fnuz(
333 weight=weight,
@@ -533,7 +532,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
532 # TODO (rob): refactor block quant into separate class.
533 if self.block_quant:
534 assert self.quant_config.activation_scheme == "dynamic"
536 - if current_platform.is_rocm():
535 + if current_platform.is_fp8_fnuz():
536 w13_weight, w13_weight_scale_inv, w13_input_scale = \
537 normalize_e4m3fn_to_e4m3fnuz(
538 layer.w13_weight, layer.w13_weight_scale_inv,
@@ -559,9 +558,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
558
559 # If checkpoint is fp16, quantize in place.
560 if not self.quant_config.is_checkpoint_fp8_serialized:
562 - # If rocm, use float8_e4m3fnuz as dtype
563 - fp8_dtype = torch.float8_e4m3fnuz \
564 - if current_platform.is_rocm() else torch.float8_e4m3fn
561 + fp8_dtype = current_platform.fp8_dtype()
562 w13_weight = torch.empty_like(layer.w13_weight.data,
563 dtype=fp8_dtype)
564 w2_weight = torch.empty_like(layer.w2_weight.data, dtype=fp8_dtype)
@@ -608,8 +605,7 @@ class Fp8MoEMethod(FusedMoEMethodBase):
605 layer.w13_input_scale.max(), requires_grad=False)
606 layer.w2_input_scale = torch.nn.Parameter(
607 layer.w2_input_scale.max(), requires_grad=False)
611 - # If rocm, normalize the weights and scales to e4m3fnuz
612 - if current_platform.is_rocm():
608 + if current_platform.is_fp8_fnuz():
609 # Normalize the weights and scales
610 w13_weight, w13_weight_scale, w13_input_scale = \
611 normalize_e4m3fn_to_e4m3fnuz(
vllm/model_executor/layers/quantization/quark/quark_moe.py
+1 -2
@@ -142,8 +142,7 @@ class QuarkW8A8Fp8MoEMethod(QuarkMoEMethod):
142 layer.w2_input_scale = torch.nn.Parameter(
143 layer.w2_input_scale.max(), requires_grad=False)
144
145 - # If rocm, normalize the weights and scales to e4m3fnuz
146 - if current_platform.is_rocm():
145 + if current_platform.is_fp8_fnuz():
146 # Normalize the weights and scales
147 w13_weight, w13_weight_scale, w13_input_scale = \
148 normalize_e4m3fn_to_e4m3fnuz(
vllm/model_executor/layers/quantization/quark/schemes/quark_w8a8_fp8.py
+2 -2
@@ -39,7 +39,7 @@ class QuarkW8A8Fp8(QuarkScheme):
39 logical_widths=layer.logical_widths,
40 )
41
42 - if current_platform.is_rocm():
42 + if current_platform.is_fp8_fnuz():
43 weight, max_w_scale, input_scale = normalize_e4m3fn_to_e4m3fnuz(
44 weight=weight,
45 weight_scale=max_w_scale,
@@ -55,7 +55,7 @@ class QuarkW8A8Fp8(QuarkScheme):
55 elif self.qscheme == "per_channel":
56 weight = layer.weight
57
58 - if current_platform.is_rocm():
58 + if current_platform.is_fp8_fnuz():
59 weight, weight_scale, input_scale = \
60 normalize_e4m3fn_to_e4m3fnuz(
61 weight=weight,
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+2 -10
@@ -22,10 +22,6 @@ from vllm.utils import direct_register_custom_op
22
23 logger = init_logger(__name__)
24
25 -current_platform_fp8_dtype = (torch.float8_e4m3fnuz
26 - if current_platform.is_rocm() else
27 - torch.float8_e4m3fn)
28 -
25
26 def is_fp8(x: Union[torch.dtype, torch.Tensor]) -> bool:
27 if isinstance(x, torch.Tensor):
@@ -165,9 +161,7 @@ def input_to_float8(
161 ) -> Tuple[torch.Tensor, torch.Tensor]:
162 """This function quantizes input values to float8 values "
163 "with tensor-wise quantization."""
168 - if dtype is None:
169 - dtype = (torch.float8_e4m3fnuz
170 - if current_platform.is_rocm() else torch.float8_e4m3fn)
164 + dtype = current_platform.fp8_dtype() if dtype is None else dtype
165 finfo = torch.finfo(dtype)
166 min_val, max_val = x.aminmax()
167 amax = torch.maximum(min_val.abs(), max_val.abs()).clamp(min=1e-12)
@@ -311,9 +305,7 @@ def per_token_group_quant_fp8(
305 Tuple[torch.Tensor, torch.Tensor]: The quantized tensor and the
306 scaling factor for quantization.
307 """
314 - if dtype is None:
315 - dtype = (torch.float8_e4m3fnuz
316 - if current_platform.is_rocm() else torch.float8_e4m3fn)
308 + dtype = current_platform.fp8_dtype() if dtype is None else dtype
309 assert (x.shape[-1] % group_size == 0), (
310 f"the last dimension of `x` {x.shape[-1]} must be divisible "
311 f"by `group_size` {group_size}")
vllm/platforms/cuda.py
+4
@@ -293,6 +293,10 @@ class CudaPlatformBase(Platform):
293 def get_device_communicator_cls(cls) -> str:
294 return "vllm.distributed.device_communicators.cuda_communicator.CudaCommunicator" # noqa
295
296 + @classmethod
297 + def supports_fp8(cls) -> bool:
298 + return cls.has_device_capability(89)
299 +
300
301 # NVML utils
302 # Note that NVML is not affected by `CUDA_VISIBLE_DEVICES`,
vllm/platforms/interface.py
+30
@@ -330,6 +330,36 @@ class Platform:
330 """
331 return "vllm.distributed.device_communicators.base_device_communicator.DeviceCommunicatorBase" # noqa
332
333 + @classmethod
334 + def supports_fp8(cls) -> bool:
335 + """
336 + Returns whether the current platform supports FP8 types.
337 + """
338 + return False
339 +
340 + @classmethod
341 + def is_fp8_fnuz(cls) -> bool:
342 + """
343 + Returns whether the preferred FP8 type is FNUZ on the current platform.
344 +
345 + There are two representations of FP8, OCP FP8 and FNUZ FP8.
346 + The OCP specification can be found at https://tinyurl.com/b7jvwpft.
347 + The FNUZ specification can be found at https://tinyurl.com/5n6hwwu5.
348 +
349 + AMD's MI300 and MI325 have native hardware support for FNUZ. All other
350 + hardware has converged on the OCP FP8 standard.
351 + """
352 + return False
353 +
354 + @classmethod
355 + def fp8_dtype(cls) -> torch.dtype:
356 + """
357 + Returns the preferred FP8 type on the current platform.
358 +
359 + See the documentation for is_fp8_fnuz for details.
360 + """
361 + return torch.float8_e4m3fn
362 +
363 @classmethod
364 def use_all_gather(cls) -> bool:
365 """
vllm/platforms/rocm.py
+17
@@ -231,3 +231,20 @@ class RocmPlatform(Platform):
231 @classmethod
232 def get_device_communicator_cls(cls) -> str:
233 return "vllm.distributed.device_communicators.cuda_communicator.CudaCommunicator" # noqa
234 +
235 + @classmethod
236 + def supports_fp8(cls) -> bool:
237 + gcn_arch = torch.cuda.get_device_properties(0).gcnArchName
238 + return any(gfx in gcn_arch for gfx in ['gfx94', 'gfx95', 'gfx12'])
239 +
240 + @classmethod
241 + def is_fp8_fnuz(cls) -> bool:
242 + # only device 0 is checked, this assumes MI300 platforms are homogeneous
243 + return 'gfx94' in torch.cuda.get_device_properties(0).gcnArchName
244 +
245 + @classmethod
246 + def fp8_dtype(cls) -> torch.dtype:
247 + if cls.is_fp8_fnuz():
248 + return torch.float8_e4m3fnuz
249 + else:
250 + return torch.float8_e4m3fn
vllm/v1/attention/backends/mla/common.py
+3 -3
@@ -219,7 +219,7 @@ from vllm.model_executor.layers.quantization.compressed_tensors.schemes import (
219 CompressedTensorsW8A8Fp8)
220 from vllm.model_executor.layers.quantization.fp8 import Fp8LinearMethod
221 from vllm.model_executor.layers.quantization.utils.fp8_utils import (
222 - Fp8LinearGenericOp, current_platform_fp8_dtype, is_fp8)
222 + Fp8LinearGenericOp, is_fp8)
223 from vllm.model_executor.layers.quantization.utils.quant_utils import (
224 scaled_quantize)
225 from vllm.model_executor.layers.rotary_embedding import RotaryEmbedding
@@ -826,7 +826,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
826 W_Q_UK, W_Q_UK_scales = scaled_quantize(
827 W_Q_UK,
828 self.reqaunt_weight_group_shape,
829 - quant_dtype=current_platform_fp8_dtype)
829 + quant_dtype=current_platform.fp8_dtype())
830 # For FP8 save the transpose so we can use
831 # `apply_w8a8_block_fp8_linear` directly
832 self.W_Q_UK = W_Q_UK.T.contiguous()
@@ -843,7 +843,7 @@ class MLACommonImpl(MLAAttentionImpl[M], Generic[M]):
843 W_UV_O, W_UV_O_scales = scaled_quantize(
844 W_UV_O,
845 self.reqaunt_weight_group_shape,
846 - quant_dtype=current_platform_fp8_dtype)
846 + quant_dtype=current_platform.fp8_dtype())
847 # For FP8 save the transpose so we can use
848 # `apply_w8a8_block_fp8_linear` directly
849 self.W_UV_O = W_UV_O.T.contiguous()