@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / e85829450d
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, 17:15 UTC

[Feature][ROCm]Enable fusion pass for torch.compile on ROCm (#15050)

Signed-off-by: charlifu <charlifu@amd.com>

Charlie Fu committed Mar 31, 2025 at 06:42 UTC e85829450d8016309d71de9f347e2147ee03400a
8 files changed +92 -72
csrc/quantization/fp8/common.cu
+2 -5
@@ -30,9 +30,6 @@ __global__ void dynamic_per_token_scaled_fp8_quant_kernel(
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 =
34 - 1.0f / (fp8_e4m3_adjusted_max_v<fp8_type> * 512.f);
35 -
33 int const tid = threadIdx.x;
34 int const token_idx = blockIdx.x;
35
@@ -67,8 +64,8 @@ __global__ void dynamic_per_token_scaled_fp8_quant_kernel(
64 token_scale = block_absmax_val_maybe;
65 }
66 // token scale computation
70 - token_scale = max(token_scale / fp8_e4m3_adjusted_max_v<fp8_type>,
71 - min_scaling_factor);
67 + token_scale = max(token_scale / quant_type_max_v<fp8_type>,
68 + min_scaling_factor<fp8_type>::val());
69 scale[token_idx] = token_scale;
70 }
71 __syncthreads();
csrc/quantization/fp8/common.cuh
+5 -36
@@ -1,20 +1,12 @@
1 #pragma once
2
3 #include "quantization/vectorization.cuh"
4 +#include "quantization/utils.cuh"
5
6 #include <cmath>
6 -#include <c10/core/ScalarType.h>
7
8 -#ifndef USE_ROCM
9 - #include <c10/util/Float8_e4m3fn.h>
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>
8 +#ifdef USE_ROCM
9 #include "amd/quant_utils.cuh"
16 - // ROCm doesn't seem to need C10_HOST_DEVICE for static constexpr
17 - #define MAYBE_HOST_DEVICE
10 #endif
11
12 // Determines the preferred FP8 type for the current platform.
@@ -31,29 +23,6 @@ static bool is_fp8_ocp() {
23 #endif
24 }
25
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 -
26 namespace vllm {
27
28 __device__ __forceinline__ float atomicMaxFloat(float* addr, float value) {
@@ -76,8 +45,8 @@ __device__ __forceinline__ fp8_type scaled_fp8_conversion(float const val,
45 x = val / scale;
46 }
47
79 - float r = fmax(-fp8_e4m3_adjusted_max_v<fp8_type>,
80 - fmin(x, fp8_e4m3_adjusted_max_v<fp8_type>));
48 + float r =
49 + fmax(-quant_type_max_v<fp8_type>, fmin(x, quant_type_max_v<fp8_type>));
50 #ifndef USE_ROCM
51 return static_cast<fp8_type>(r);
52 #else
@@ -123,7 +92,7 @@ __global__ void segmented_max_reduction(float* __restrict__ scale,
92 // Finally, since cache[0] contains the maximum for this thread block,
93 // atomically write the max to the target location
94 if (threadIdx.x == 0) {
126 - atomicMaxFloat(scale, cache[0] / fp8_e4m3_adjusted_max_v<fp8_type>);
95 + atomicMaxFloat(scale, cache[0] / quant_type_max_v<fp8_type>);
96 }
97 }
98
csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu
+10 -18
@@ -14,8 +14,7 @@ __device__ void rms_norm_dynamic_per_token_quant_vec(
14 float* __restrict__ scales, // [num_tokens]
15 scalar_t const* __restrict__ input, // [..., hidden_size]
16 scalar_t const* __restrict__ weight, // [hidden_size]
17 - float const* scale_ub, float const var_epsilon,
18 - float const min_scaling_factor, int32_t const hidden_size,
17 + float const* scale_ub, float const var_epsilon, int32_t const hidden_size,
18 scalar_t* __restrict__ residual = nullptr) {
19 float rms = 0.0f;
20 float token_scale = 0.0f;
@@ -27,8 +26,8 @@ __device__ void rms_norm_dynamic_per_token_quant_vec(
26 // Compute scale
27 vllm::vectorized::compute_dynamic_per_token_scales<scalar_t, scalar_out_t,
28 has_residual>(
30 - &token_scale, scales, input, weight, rms, scale_ub, min_scaling_factor,
31 - hidden_size, residual);
29 + &token_scale, scales, input, weight, rms, scale_ub, hidden_size,
30 + residual);
31
32 // RMS Norm + Quant
33 if constexpr (std::is_same_v<scalar_out_t, int8_t>) {
@@ -50,8 +49,7 @@ __global__ void rms_norm_dynamic_per_token_quant_kernel(
49 float* __restrict__ scales, // [num_tokens]
50 scalar_t const* __restrict__ input, // [..., hidden_size]
51 scalar_t const* __restrict__ weight, // [hidden_size]
53 - float const* scale_ub, float const var_epsilon,
54 - float const min_scaling_factor, int32_t const hidden_size,
52 + float const* scale_ub, float const var_epsilon, int32_t const hidden_size,
53 scalar_t* __restrict__ residual = nullptr) {
54 // For vectorization, token_input and token_output pointers need to be
55 // aligned at 8-byte and 4-byte addresses respectively.
@@ -60,8 +58,8 @@ __global__ void rms_norm_dynamic_per_token_quant_kernel(
58 if (can_vectorize) {
59 return rms_norm_dynamic_per_token_quant_vec<scalar_t, scalar_out_t,
60 has_residual>(
63 - out, scales, input, weight, scale_ub, var_epsilon, min_scaling_factor,
64 - hidden_size, residual);
61 + out, scales, input, weight, scale_ub, var_epsilon, hidden_size,
62 + residual);
63 }
64
65 float rms = 0.0f;
@@ -72,8 +70,8 @@ __global__ void rms_norm_dynamic_per_token_quant_kernel(
70 var_epsilon, residual);
71 // Compute Scale
72 vllm::compute_dynamic_per_token_scales<scalar_t, scalar_out_t, has_residual>(
75 - &token_scale, scales, input, weight, rms, scale_ub, min_scaling_factor,
76 - hidden_size, residual);
73 + &token_scale, scales, input, weight, rms, scale_ub, hidden_size,
74 + residual);
75
76 // RMS Norm + Quant
77 if constexpr (std::is_same_v<scalar_out_t, int8_t>) {
@@ -105,11 +103,6 @@ void rms_norm_dynamic_per_token_quant_dispatch(
103 const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
104 const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
105
108 - const float min_scaling_factor =
109 - out.dtype() == torch::kInt8
110 - ? std::numeric_limits<float>::epsilon()
111 - : 1.0f / (std::numeric_limits<c10::Float8_e4m3fn>::max() * 512.f);
112 -
106 if (residual.has_value()) {
107 VLLM_DISPATCH_QUANT_TYPES(
108 out.scalar_type(), "rms_norm_dynamic_per_token_quant_kernel", [&] {
@@ -119,8 +112,7 @@ void rms_norm_dynamic_per_token_quant_dispatch(
112 out.data_ptr<scalar_t>(), scales.data_ptr<float>(),
113 input.data_ptr<scalar_in_t>(), weight.data_ptr<scalar_in_t>(),
114 scale_ub.has_value() ? scale_ub->data_ptr<float>() : nullptr,
122 - var_epsilon, min_scaling_factor, hidden_size,
123 - residual->data_ptr<scalar_in_t>());
115 + var_epsilon, hidden_size, residual->data_ptr<scalar_in_t>());
116 });
117
118 } else {
@@ -132,7 +124,7 @@ void rms_norm_dynamic_per_token_quant_dispatch(
124 out.data_ptr<scalar_t>(), scales.data_ptr<float>(),
125 input.data_ptr<scalar_in_t>(), weight.data_ptr<scalar_in_t>(),
126 scale_ub.has_value() ? scale_ub->data_ptr<float>() : nullptr,
135 - var_epsilon, min_scaling_factor, hidden_size, nullptr);
127 + var_epsilon, hidden_size, nullptr);
128 });
129 }
130 }
csrc/quantization/fused_kernels/layernorm_utils.cuh
+7 -6
@@ -5,6 +5,7 @@
5 */
6
7 #include "quantization/vectorization.cuh"
8 +#include "quantization/utils.cuh"
9 #include "quant_conversions.cuh"
10
11 #ifndef USE_ROCM
@@ -51,11 +52,11 @@ __device__ void compute_dynamic_per_token_scales(
52 float* __restrict__ token_scale, float* __restrict__ all_token_scales,
53 scalar_t const* __restrict__ input, scalar_t const* __restrict__ weight,
54 float const rms, float const* __restrict__ scale_ub,
54 - float const min_scaling_factor, int32_t const hidden_size,
55 + int32_t const hidden_size,
56 scalar_t const* __restrict__ residual = nullptr) {
57 int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
58 ;
58 - constexpr scalar_out_t qmax{std::numeric_limits<scalar_out_t>::max()};
59 + constexpr scalar_out_t qmax{quant_type_max_v<scalar_out_t>};
60
61 float block_absmax_val_maybe = 0.0f;
62 for (auto i = threadIdx.x; i < hidden_size; i += blockDim.x) {
@@ -83,7 +84,7 @@ __device__ void compute_dynamic_per_token_scales(
84 scale = block_absmax_val_maybe;
85 }
86 // token scale computation
86 - scale = max(scale / qmax, min_scaling_factor);
87 + scale = max(scale / qmax, min_scaling_factor<scalar_out_t>::val());
88 s_token_scale = scale; // Shared memory store
89 all_token_scales[blockIdx.x] = scale; // Global output store
90 }
@@ -184,7 +185,7 @@ __device__ void compute_dynamic_per_token_scales(
185 float* __restrict__ token_scale, float* __restrict__ all_token_scales,
186 scalar_t const* __restrict__ input, scalar_t const* __restrict__ weight,
187 float const rms, float const* __restrict__ scale_ub,
187 - float const min_scaling_factor, int32_t const hidden_size,
188 + int32_t const hidden_size,
189 scalar_t const* __restrict__ residual = nullptr) {
190 int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
191 ;
@@ -200,7 +201,7 @@ __device__ void compute_dynamic_per_token_scales(
201 reinterpret_cast<vec4_t<scalar_t> const*>(&residual[token_offset]);
202 }
203
203 - constexpr scalar_out_t qmax{std::numeric_limits<scalar_out_t>::max()};
204 + constexpr scalar_out_t qmax{quant_type_max_v<scalar_out_t>};
205
206 int32_t const num_vec_elems = hidden_size >> 2;
207 float block_absmax_val_maybe = 0.0f;
@@ -248,7 +249,7 @@ __device__ void compute_dynamic_per_token_scales(
249 scale = block_absmax_val_maybe;
250 }
251 // token scale computation
251 - scale = max(scale / qmax, min_scaling_factor);
252 + scale = max(scale / qmax, min_scaling_factor<scalar_out_t>::val());
253 s_token_scale = scale; // shared memory store
254 all_token_scales[blockIdx.x] = scale; // global output store
255 }
csrc/quantization/fused_kernels/quant_conversions.cuh
+2 -2
@@ -33,8 +33,8 @@ static __device__ __forceinline__ int8_t float_to_int8_rn(float const x) {
33
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>));
36 + float const r =
37 + fmax(-quant_type_max_v<fp8_type>, fmin(x, quant_type_max_v<fp8_type>));
38 return static_cast<fp8_type>(r);
39 }
40
csrc/quantization/utils.cuh new
+59
@@ -0,0 +1,59 @@
1 +#pragma once
2 +
3 +/**
4 + * Quantization utilities including:
5 + * Adjusted maximum values for qtypes.
6 + * Minimum scaling factors for qtypes.
7 + */
8 +
9 +#include <cmath>
10 +#include <torch/types.h>
11 +
12 +#ifndef USE_ROCM
13 + #include <c10/util/Float8_e4m3fn.h>
14 + #define MAYBE_HOST_DEVICE C10_HOST_DEVICE
15 +#else
16 + #include <ATen/hip/HIPContext.h>
17 + #include <c10/util/Float8_e4m3fn.h>
18 + #include <c10/util/Float8_e4m3fnuz.h>
19 + // ROCm doesn't seem to need C10_HOST_DEVICE for static constexpr
20 + #define MAYBE_HOST_DEVICE
21 +#endif
22 +
23 +template <typename T,
24 + typename = std::enable_if_t<std::is_same_v<T, c10::Float8_e4m3fn> ||
25 + std::is_same_v<T, c10::Float8_e4m3fnuz> ||
26 + std::is_same_v<T, int8_t>>>
27 +struct quant_type_max {
28 + static constexpr T val() { return std::numeric_limits<T>::max(); }
29 +};
30 +
31 +// Using the default max value from pytorch (240.0 0x7F) will cause accuracy
32 +// issues when running dynamic quantization. Here use 224.0 0x7E for rocm.
33 +template <>
34 +struct quant_type_max<c10::Float8_e4m3fnuz> {
35 + static constexpr c10::Float8_e4m3fnuz val() {
36 + return c10::Float8_e4m3fnuz(0x7E, c10::Float8_e4m3fnuz::from_bits());
37 + }
38 +};
39 +
40 +template <typename T>
41 +MAYBE_HOST_DEVICE static constexpr T quant_type_max_v =
42 + quant_type_max<T>::val();
43 +
44 +template <typename T,
45 + typename = std::enable_if_t<std::is_same_v<T, c10::Float8_e4m3fn> ||
46 + std::is_same_v<T, c10::Float8_e4m3fnuz> ||
47 + std::is_same_v<T, int8_t>>>
48 +struct min_scaling_factor {
49 + C10_DEVICE C10_ALWAYS_INLINE static float val() {
50 + return 1.0f / (quant_type_max_v<T> * 512.0f);
51 + }
52 +};
53 +
54 +template <>
55 +struct min_scaling_factor<int8_t> {
56 + C10_DEVICE C10_ALWAYS_INLINE static float val() {
57 + return std::numeric_limits<float>::epsilon();
58 + }
59 +};
\ No newline at end of file
tests/compile/test_fusion.py
+5 -3
@@ -2,7 +2,6 @@
2
3 import pytest
4 import torch
5 -from compressed_tensors.quantization import FP8_DTYPE
5
6 import vllm.envs as envs
7 import vllm.plugins
@@ -14,9 +13,12 @@ from vllm.config import CompilationConfig, CompilationLevel, VllmConfig
13 from vllm.model_executor.layers.layernorm import RMSNorm
14 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
15 CUTLASS_FP8_SUPPORTED, Fp8LinearOp, maybe_create_device_identity)
16 +from vllm.platforms import current_platform
17
18 from .backend import TestBackend
19
20 +FP8_DTYPE = current_platform.fp8_dtype()
21 +
22
23 class TestModel(torch.nn.Module):
24
@@ -59,8 +61,8 @@ class TestModel(torch.nn.Module):
61 @pytest.mark.parametrize("static", [True, False])
62 @pytest.mark.parametrize("cutlass_fp8_enabled",
63 [True, False] if CUTLASS_FP8_SUPPORTED else [False])
62 -@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE != "cuda",
63 - reason="Only test on CUDA")
64 +@pytest.mark.skipif(envs.VLLM_TARGET_DEVICE not in ["cuda", "rocm"],
65 + reason="Only test on CUDA and ROCm")
66 def test_fusion_rmsnorm_quant(dtype, hidden_size, num_tokens, eps, static,
67 cutlass_fp8_enabled):
68 torch.set_default_device("cuda")
vllm/compilation/fusion.py
+2 -2
@@ -4,8 +4,6 @@ from typing import Callable, Dict, List, NamedTuple, Optional, Tuple
4
5 import torch
6 import torch._inductor.pattern_matcher as pm
7 -# TODO(luka) use vllm.utils once #10836 landed
8 -from compressed_tensors.quantization import FP8_DTYPE
7 from torch import fx
8 from torch._higher_order_ops.auto_functionalize import auto_functionalized
9 from torch._inductor.pattern_matcher import PatternMatcherPass
@@ -13,12 +11,14 @@ from torch._ops import OpOverload
11
12 from vllm.config import CompilationConfig
13 from vllm.logger import init_logger
14 +from vllm.platforms import current_platform
15
16 from .fx_utils import find_getitem_maybe
17 from .multi_output_match import MultiOutputMatch
18 from .vllm_inductor_pass import VllmInductorPass
19
20 logger = init_logger(__name__)
21 +FP8_DTYPE = current_platform.fp8_dtype()
22
23
24 def empty_bf16(*args, **kwargs):