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

[Performance] Fused blockwise quant RMS norm (#27883)

Signed-off-by: ElizaWszola <ewszola@redhat.com> Signed-off-by: yewentao256 <zhyanwentao@126.com> Co-authored-by: yewentao256 <zhyanwentao@126.com>

ElizaWszola committed Dec 7, 2025 at 17:38 UTC af0444bf40b7db2f3fb9fe1508d25ceba24cac87
14 files changed +946 -154
benchmarks/fused_kernels/layernorm_rms_benchmarks.py
+86 -3
@@ -14,6 +14,9 @@ from tqdm import tqdm
14
15 import vllm._custom_ops as ops
16 from vllm.model_executor.layers.layernorm import RMSNorm
17 +from vllm.model_executor.layers.quantization.utils.fp8_utils import (
18 + per_token_group_quant_fp8,
19 +)
20
21
22 @dataclass
@@ -22,6 +25,7 @@ class bench_params_t:
25 hidden_size: int
26 add_residual: bool
27 dtype: torch.dtype
28 + group_size: list[int]
29
30 def description(self):
31 return (
@@ -29,6 +33,7 @@ class bench_params_t:
33 f"x D {self.hidden_size} "
34 f"x R {self.add_residual} "
35 f"x DT {self.dtype}"
36 + f"x GS {self.group_size}"
37 )
38
39
@@ -38,10 +43,11 @@ def get_bench_params() -> list[bench_params_t]:
43 HIDDEN_SIZES = list(range(1024, 8129, 1024))
44 ADD_RESIDUAL = [True, False]
45 DTYPES = [torch.bfloat16, torch.float]
46 + GROUP_SIZES = [[1, 64], [1, 128]]
47
42 - combinations = product(NUM_TOKENS, HIDDEN_SIZES, ADD_RESIDUAL, DTYPES)
48 + combinations = product(NUM_TOKENS, HIDDEN_SIZES, ADD_RESIDUAL, DTYPES, GROUP_SIZES)
49 bench_params = list(
44 - map(lambda x: bench_params_t(x[0], x[1], x[2], x[3]), combinations)
50 + map(lambda x: bench_params_t(x[0], x[1], x[2], x[3], x[4]), combinations)
51 )
52 return bench_params
53
@@ -52,6 +58,7 @@ def unfused_int8_impl(
58 x: torch.Tensor,
59 residual: torch.Tensor | None,
60 quant_dtype: torch.dtype,
61 + group_size: list[int],
62 ):
63 # Norm
64 torch_out = None
@@ -69,6 +76,7 @@ def unfused_fp8_impl(
76 x: torch.Tensor,
77 residual: torch.Tensor | None,
78 quant_dtype: torch.dtype,
79 + group_size: list[int],
80 ):
81 # Norm
82 torch_out = None
@@ -81,23 +89,63 @@ def unfused_fp8_impl(
89 torch_out, _ = ops.scaled_fp8_quant(torch_out)
90
91
92 +def unfused_groupwise_fp8_impl(
93 + rms_norm_layer: RMSNorm,
94 + x: torch.Tensor,
95 + residual: torch.Tensor | None,
96 + quant_dtype: torch.dtype,
97 + group_size: list[int],
98 +):
99 + # Norm
100 + torch_out = None
101 + if residual is None:
102 + torch_out = rms_norm_layer.forward_cuda(x, residual)
103 + else:
104 + torch_out, _ = rms_norm_layer.forward_cuda(x, residual)
105 +
106 + # Quant
107 + torch_out, _ = per_token_group_quant_fp8(
108 + torch_out, group_size=group_size[1], use_ue8m0=False
109 + )
110 +
111 +
112 def fused_impl(
113 rms_norm_layer: RMSNorm, # this stores the weights
114 x: torch.Tensor,
115 residual: torch.Tensor | None,
116 quant_dtype: torch.dtype,
117 + group_size: list[int],
118 ):
119 out, _ = ops.rms_norm_dynamic_per_token_quant(
120 x, rms_norm_layer.weight, 1e-6, quant_dtype, residual=residual
121 )
122
123
124 +def fused_groupwise_impl(
125 + rms_norm_layer: RMSNorm, # this stores the weights
126 + x: torch.Tensor,
127 + residual: torch.Tensor | None,
128 + quant_dtype: torch.dtype,
129 + group_size: list[int],
130 +):
131 + out, _ = ops.rms_norm_per_block_quant(
132 + x,
133 + rms_norm_layer.weight,
134 + 1e-6,
135 + quant_dtype,
136 + group_size,
137 + residual=residual,
138 + is_scale_transposed=True,
139 + )
140 +
141 +
142 # Bench functions
143 def bench_fn(
144 rms_norm_layer: RMSNorm,
145 x: torch.Tensor,
146 residual: torch.Tensor,
147 quant_dtype: torch.dtype,
148 + group_size: list[int],
149 label: str,
150 sub_label: str,
151 fn: Callable,
@@ -110,10 +158,11 @@ def bench_fn(
158 "x": x,
159 "residual": residual,
160 "quant_dtype": quant_dtype,
161 + "group_size": group_size,
162 "fn": fn,
163 }
164 return TBenchmark.Timer(
116 - stmt="fn(rms_norm_layer, x, residual, quant_dtype)",
165 + stmt="fn(rms_norm_layer, x, residual, quant_dtype, group_size)",
166 globals=globals,
167 label=label,
168 sub_label=sub_label,
@@ -147,6 +196,7 @@ def bench(params: bench_params_t, label: str, sub_label: str) -> Iterable[TMeasu
196 x,
197 residual,
198 torch.int8,
199 + params.group_size,
200 label,
201 sub_label,
202 unfused_int8_impl,
@@ -161,6 +211,7 @@ def bench(params: bench_params_t, label: str, sub_label: str) -> Iterable[TMeasu
211 x,
212 residual,
213 torch.float8_e4m3fn,
214 + params.group_size,
215 label,
216 sub_label,
217 unfused_fp8_impl,
@@ -175,6 +226,7 @@ def bench(params: bench_params_t, label: str, sub_label: str) -> Iterable[TMeasu
226 x,
227 residual,
228 torch.int8,
229 + params.group_size,
230 label,
231 sub_label,
232 fused_impl,
@@ -189,6 +241,7 @@ def bench(params: bench_params_t, label: str, sub_label: str) -> Iterable[TMeasu
241 x,
242 residual,
243 torch.float8_e4m3fn,
244 + params.group_size,
245 label,
246 sub_label,
247 fused_impl,
@@ -196,6 +249,36 @@ def bench(params: bench_params_t, label: str, sub_label: str) -> Iterable[TMeasu
249 )
250 )
251
252 + # unfused groupwise fp8 impl.
253 + timers.append(
254 + bench_fn(
255 + layer,
256 + x,
257 + residual,
258 + torch.float8_e4m3fn,
259 + params.group_size,
260 + label,
261 + sub_label,
262 + unfused_groupwise_fp8_impl,
263 + "unfused_groupwise_fp8_impl",
264 + )
265 + )
266 +
267 + # fused groupwise fp8 impl.
268 + timers.append(
269 + bench_fn(
270 + layer,
271 + x,
272 + residual,
273 + torch.float8_e4m3fn,
274 + params.group_size,
275 + label,
276 + sub_label,
277 + fused_groupwise_impl,
278 + "fused_groupwise_fp8_impl",
279 + )
280 + )
281 +
282 print_timers(timers)
283
284 return timers
csrc/dispatch_utils.h
+18
@@ -118,6 +118,24 @@
118 } \
119 }
120
121 +#define VLLM_DISPATCH_BOOL(expr, const_expr, ...) \
122 + if (expr) { \
123 + constexpr bool const_expr = true; \
124 + __VA_ARGS__(); \
125 + } else { \
126 + constexpr bool const_expr = false; \
127 + __VA_ARGS__(); \
128 + }
129 +
130 +#define VLLM_DISPATCH_GROUP_SIZE(group_size, const_group_size, ...) \
131 + if (group_size == 128) { \
132 + constexpr int const_group_size = 128; \
133 + __VA_ARGS__(); \
134 + } else if (group_size == 64) { \
135 + constexpr int const_group_size = 64; \
136 + __VA_ARGS__(); \
137 + }
138 +
139 #define VLLM_DISPATCH_RANK234(NUM_DIMS, ...) \
140 switch (NUM_DIMS) { \
141 case 2: { \
csrc/ops.h
+7
@@ -128,6 +128,13 @@ void rms_norm_dynamic_per_token_quant(torch::Tensor& out,
128 std::optional<torch::Tensor> scale_ub,
129 std::optional<torch::Tensor> residual);
130
131 +void rms_norm_per_block_quant(torch::Tensor& out, torch::Tensor const& input,
132 + torch::Tensor const& weight,
133 + torch::Tensor& scales, double const epsilon,
134 + std::optional<torch::Tensor> scale_ub,
135 + std::optional<torch::Tensor> residual,
136 + int64_t group_size, bool is_scale_transposed);
137 +
138 void rotary_embedding(torch::Tensor& positions, torch::Tensor& query,
139 std::optional<torch::Tensor> key, int64_t head_size,
140 torch::Tensor& cos_sin_cache, bool is_neox);
csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu
+124 -20
@@ -31,14 +31,15 @@ __device__ void rms_norm_dynamic_per_token_quant_vec(
31
32 // RMS Norm + Quant
33 if constexpr (std::is_same_v<scalar_out_t, int8_t>) {
34 + token_scale = 1.0f / token_scale;
35 vllm::vectorized::norm_and_quant<scalar_t, scalar_out_t, true,
36 has_residual>(
36 - out, input, weight, rms, 1.0f / token_scale, hidden_size, residual);
37 + out, input, weight, rms, &token_scale, hidden_size, residual);
38 } else {
39 // FP8 - Do not invert token_scale for exact match with FBGemm
40 vllm::vectorized::norm_and_quant<scalar_t, scalar_out_t, false,
41 has_residual>(
41 - out, input, weight, rms, token_scale, hidden_size, residual);
42 + out, input, weight, rms, &token_scale, hidden_size, residual);
43 }
44 }
45
@@ -75,14 +76,52 @@ __global__ void rms_norm_dynamic_per_token_quant_kernel(
76
77 // RMS Norm + Quant
78 if constexpr (std::is_same_v<scalar_out_t, int8_t>) {
79 + token_scale = 1.0f / token_scale;
80 vllm::norm_and_quant<scalar_t, scalar_out_t, true, has_residual>(
79 - out, input, weight, rms, 1.0f / token_scale, hidden_size, residual);
81 + out, input, weight, rms, &token_scale, hidden_size, residual);
82 } else {
83 // FP8 - Do not invert s_token_scale for exact match with FBGemm
84 vllm::norm_and_quant<scalar_t, scalar_out_t, false, has_residual>(
83 - out, input, weight, rms, token_scale, hidden_size, residual);
85 + out, input, weight, rms, &token_scale, hidden_size, residual);
86 }
87 }
88 +
89 +// RMS norm + quant kernel
90 +template <typename scalar_t, typename scalar_out_t, bool has_residual = false,
91 + bool is_scale_transposed = false, int32_t group_size = 0>
92 +__global__ void rms_norm_per_block_quant_kernel(
93 + scalar_out_t* __restrict__ out, // [..., hidden_size]
94 + float* __restrict__ scales, // [num_tokens, hidden_size / group_size]
95 + // or
96 + // [hidden_size / group_size, num_tokens]
97 + scalar_t const* __restrict__ input, // [..., hidden_size]
98 + scalar_t const* __restrict__ weight, // [hidden_size]
99 + float const* scale_ub, float const var_epsilon, int32_t const hidden_size,
100 + scalar_t* __restrict__ residual = nullptr) {
101 + float rms;
102 + // Compute RMS
103 + // Always able to vectorize due to constraints on hidden_size
104 + vllm::vectorized::compute_rms<scalar_t, has_residual>(
105 + &rms, input, hidden_size, var_epsilon, residual);
106 +
107 + // Compute Scale
108 + // Always able to vectorize due to constraints on hidden_size and group_size
109 + vllm::vectorized::compute_dynamic_per_token_scales<
110 + scalar_t, scalar_out_t, has_residual, is_scale_transposed, group_size>(
111 + nullptr, scales, input, weight, rms, scale_ub, hidden_size, residual);
112 +
113 + // RMS Norm + Quant
114 + // Always able to vectorize due to constraints on hidden_size
115 + // For int8, don't invert token_scale here: do it inside the norm_and_quant
116 + // kernel. We do it because particular elements of token_scale can be shared
117 + // between multiple threads, so this way, we avoid extra synchronization
118 + // overhead.
119 + vllm::vectorized::norm_and_quant<
120 + scalar_t, scalar_out_t, std::is_same_v<scalar_out_t, int8_t>,
121 + has_residual, is_scale_transposed, group_size>(
122 + out, input, weight, rms, scales, hidden_size, residual);
123 +}
124 +
125 } // namespace vllm
126
127 // Residual add + RMS norm + dynamic per token
@@ -103,30 +142,19 @@ void rms_norm_dynamic_per_token_quant_dispatch(
142 const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
143 const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
144
106 - if (residual.has_value()) {
145 + VLLM_DISPATCH_BOOL(residual.has_value(), has_residual, [&] {
146 VLLM_DISPATCH_QUANT_TYPES(
147 out.scalar_type(), "rms_norm_dynamic_per_token_quant_kernel", [&] {
148 vllm::rms_norm_dynamic_per_token_quant_kernel<scalar_in_t, scalar_t,
110 - true>
149 + has_residual>
150 <<<grid, block, 0, stream>>>(
151 out.data_ptr<scalar_t>(), scales.data_ptr<float>(),
152 input.data_ptr<scalar_in_t>(), weight.data_ptr<scalar_in_t>(),
153 scale_ub.has_value() ? scale_ub->data_ptr<float>() : nullptr,
115 - var_epsilon, hidden_size, residual->data_ptr<scalar_in_t>());
154 + var_epsilon, hidden_size,
155 + has_residual ? residual->data_ptr<scalar_in_t>() : nullptr);
156 });
117 -
118 - } else {
119 - VLLM_DISPATCH_QUANT_TYPES(
120 - out.scalar_type(), "rms_norm_dynamic_per_token_quant_kernel", [&] {
121 - vllm::rms_norm_dynamic_per_token_quant_kernel<scalar_in_t, scalar_t,
122 - false>
123 - <<<grid, block, 0, stream>>>(
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,
127 - var_epsilon, hidden_size, nullptr);
128 - });
129 - }
157 + });
158 }
159
160 void rms_norm_dynamic_per_token_quant(
@@ -157,3 +185,79 @@ void rms_norm_dynamic_per_token_quant(
185 out, input, weight, scales, var_epsilon, scale_ub, residual);
186 });
187 }
188 +
189 +// Residual add + RMS norm + dynamic per token
190 +void rms_norm_per_block_quant_dispatch(
191 + torch::Tensor& out, // [..., hidden_size]
192 + torch::Tensor const& input, // [..., hidden_size]
193 + torch::Tensor const& weight, // [hidden_size]
194 + torch::Tensor& scales, // [num_tokens, hidden_size / group_size] or
195 + // [hidden_size / group_size, num_tokens]
196 + int32_t group_size,
197 + double const var_epsilon, // Variance epsilon used in norm calculation
198 + std::optional<at::Tensor> const& scale_ub,
199 + std::optional<at::Tensor>& residual, bool is_scale_transposed) {
200 + int32_t hidden_size = input.size(-1);
201 + auto num_tokens = input.numel() / hidden_size;
202 +
203 + dim3 grid(num_tokens);
204 + const int max_block_size = (num_tokens <= 256) ? 512 : 256;
205 + dim3 block(std::min(hidden_size, max_block_size));
206 + const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
207 + const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
208 +
209 + VLLM_DISPATCH_FLOATING_TYPES(
210 + input.scalar_type(), "rms_norm_per_block_quant_fp_dispatch", [&] {
211 + using scalar_in_t = scalar_t;
212 + VLLM_DISPATCH_GROUP_SIZE(group_size, gs, [&] {
213 + VLLM_DISPATCH_BOOL(residual.has_value(), has_residual, [&] {
214 + VLLM_DISPATCH_BOOL(is_scale_transposed, transpose_scale, [&] {
215 + VLLM_DISPATCH_QUANT_TYPES(
216 + out.scalar_type(), "rms_norm_per_block_quant_kernel", [&] {
217 + vllm::rms_norm_per_block_quant_kernel<scalar_in_t, scalar_t,
218 + has_residual,
219 + transpose_scale, gs>
220 + <<<grid, block, 0, stream>>>(
221 + out.data_ptr<scalar_t>(), scales.data_ptr<float>(),
222 + input.data_ptr<scalar_in_t>(),
223 + weight.data_ptr<scalar_in_t>(),
224 + scale_ub.has_value() ? scale_ub->data_ptr<float>()
225 + : nullptr,
226 + var_epsilon, hidden_size,
227 + has_residual ? residual->data_ptr<scalar_in_t>()
228 + : nullptr);
229 + });
230 + });
231 + });
232 + });
233 + });
234 +}
235 +
236 +void rms_norm_per_block_quant(torch::Tensor& out, torch::Tensor const& input,
237 + torch::Tensor const& weight,
238 + torch::Tensor& scales, double const var_epsilon,
239 + std::optional<torch::Tensor> scale_ub,
240 + std::optional<torch::Tensor> residual,
241 + int64_t group_size, bool is_scale_transposed) {
242 + static c10::ScalarType kFp8Type = is_fp8_ocp()
243 + ? c10::ScalarType::Float8_e4m3fn
244 + : c10::ScalarType::Float8_e4m3fnuz;
245 + TORCH_CHECK(out.dtype() == kFp8Type || out.dtype() == torch::kInt8);
246 + TORCH_CHECK(out.is_contiguous() && input.is_contiguous());
247 +
248 + if (scale_ub.has_value()) {
249 + TORCH_CHECK(out.dtype() == kFp8Type);
250 + }
251 + TORCH_CHECK(weight.dtype() == input.dtype());
252 + TORCH_CHECK(scales.dtype() == torch::kFloat32);
253 + if (residual) {
254 + TORCH_CHECK(residual->scalar_type() == input.scalar_type());
255 + }
256 +
257 + TORCH_CHECK(group_size == 128 || group_size == 64,
258 + "Unsupported group size: ", group_size);
259 +
260 + rms_norm_per_block_quant_dispatch(out, input, weight, scales, group_size,
261 + var_epsilon, scale_ub, residual,
262 + is_scale_transposed);
263 +}
\ No newline at end of file
csrc/quantization/fused_kernels/layernorm_utils.cuh
+308 -94
@@ -9,6 +9,7 @@
9 #include "quant_conversions.cuh"
10
11 #include "../../cub_helpers.h"
12 +#include "../../cuda_compat.h"
13
14 namespace vllm {
15
@@ -43,62 +44,150 @@ __device__ void compute_rms(float* rms, scalar_t const* __restrict__ input,
44 *rms = s_rms;
45 }
46
46 -template <typename scalar_t, typename scalar_out_t, bool has_residual = false>
47 +__device__ float warpReduceMaxSpecialized(volatile float* val, int64_t tid,
48 + int64_t thread_in_warp,
49 + int64_t reduced_elems) {
50 + static_assert(WARP_SIZE == 32 || WARP_SIZE == 64);
51 + if constexpr (WARP_SIZE == 64) {
52 + if (thread_in_warp + 64 < reduced_elems)
53 + val[tid] = fmaxf(val[tid], val[tid + 64]);
54 + }
55 + if (thread_in_warp + 32 < reduced_elems)
56 + val[tid] = fmaxf(val[tid], val[tid + 32]);
57 + if (thread_in_warp + 16 < reduced_elems)
58 + val[tid] = fmaxf(val[tid], val[tid + 16]);
59 + if (thread_in_warp + 8 < reduced_elems)
60 + val[tid] = fmaxf(val[tid], val[tid + 8]);
61 + if (thread_in_warp + 4 < reduced_elems)
62 + val[tid] = fmaxf(val[tid], val[tid + 4]);
63 + if (thread_in_warp + 2 < reduced_elems)
64 + val[tid] = fmaxf(val[tid], val[tid + 2]);
65 + if (thread_in_warp + 1 < reduced_elems)
66 + val[tid] = fmaxf(val[tid], val[tid + 1]);
67 + return val[tid];
68 +}
69 +
70 +template <typename scalar_t, typename scalar_out_t, bool has_residual = false,
71 + bool is_scale_transposed = false>
72 __device__ void compute_dynamic_per_token_scales(
73 float* __restrict__ token_scale, float* __restrict__ all_token_scales,
74 scalar_t const* __restrict__ input, scalar_t const* __restrict__ weight,
75 float const rms, float const* __restrict__ scale_ub,
51 - int32_t const hidden_size,
52 - scalar_t const* __restrict__ residual = nullptr) {
53 - int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
54 - ;
55 - constexpr scalar_out_t qmax{quant_type_max_v<scalar_out_t>};
56 -
76 + int32_t const hidden_size, scalar_t const* __restrict__ residual = nullptr,
77 + int32_t const group_size = 0) {
78 float block_absmax_val_maybe = 0.0f;
58 - for (auto i = threadIdx.x; i < hidden_size; i += blockDim.x) {
59 - float x = static_cast<float>(input[token_offset + i]);
60 - if constexpr (has_residual) {
61 - x += static_cast<float>(residual[token_offset + i]);
79 + constexpr scalar_out_t qmax{quant_type_max_v<scalar_out_t>};
80 + __syncthreads();
81 + if (group_size > 0) {
82 + __shared__ float s_max_vals[1024];
83 + int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
84 + int64_t num_groups = hidden_size / group_size;
85 + int64_t const threads_per_group = blockDim.x / num_groups;
86 + int64_t const thread_in_group = threadIdx.x % threads_per_group;
87 + int64_t const group_offset = threadIdx.x / threads_per_group * group_size;
88 + int64_t const thread_offset = group_offset + thread_in_group;
89 + int64_t const thread_end =
90 + min(group_offset + group_size, static_cast<int64_t>(hidden_size));
91 + for (auto i = thread_offset; i < thread_end; i += threads_per_group) {
92 + float x = static_cast<float>(input[token_offset + i]);
93 + if constexpr (has_residual) {
94 + x += static_cast<float>(residual[token_offset + i]);
95 + }
96 + x = static_cast<float>(static_cast<scalar_t>(x * rms) * weight[i]);
97 + block_absmax_val_maybe = fmaxf(block_absmax_val_maybe, fabsf(x));
98 }
99 + s_max_vals[threadIdx.x] = block_absmax_val_maybe;
100 + __syncthreads();
101 +
102 + int64_t const warp_size = WARP_SIZE;
103 + int64_t const num_warps = blockDim.x / warp_size;
104 + int64_t const warp_id = threadIdx.x / warp_size;
105 + int64_t const thread_in_warp = threadIdx.x % warp_size;
106 + int64_t const groups_per_warp = (num_groups + num_warps - 1) / num_warps;
107 + for (auto i = 0; i < groups_per_warp; ++i) {
108 + int64_t const group_id = i * num_warps + warp_id;
109 + if (group_id < num_groups) {
110 + int64_t warp_start = group_id * threads_per_group;
111 + int64_t const start = warp_start + thread_in_warp;
112 + int64_t const warp_end = min(warp_start + threads_per_group,
113 + static_cast<int64_t>(hidden_size));
114 + for (auto j = start; j + warp_size < warp_end; j += warp_size) {
115 + s_max_vals[start] =
116 + fmaxf(s_max_vals[start], s_max_vals[j + warp_size]);
117 + }
118 + warpReduceMaxSpecialized(s_max_vals, start, thread_in_warp,
119 + min(warp_end - warp_start, warp_size));
120 + }
121 + }
122 + __syncthreads();
123 +
124 + if (thread_in_group == 0 && thread_offset < thread_end) {
125 + block_absmax_val_maybe = s_max_vals[threadIdx.x];
126 + float scale = 0.0f;
127 + if (scale_ub) {
128 + scale = min(block_absmax_val_maybe, *scale_ub);
129 + } else {
130 + scale = block_absmax_val_maybe;
131 + }
132 + // token scale computation
133 + scale = max(scale / qmax, min_scaling_factor<scalar_out_t>::val());
134 + // Global output store
135 + if constexpr (is_scale_transposed) {
136 + all_token_scales[(threadIdx.x / threads_per_group) * gridDim.x +
137 + blockIdx.x] = scale;
138 + } else {
139 + all_token_scales[blockIdx.x * num_groups +
140 + threadIdx.x / threads_per_group] = scale;
141 + }
142 + }
143 + __syncthreads();
144 + } else {
145 + int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
146 +
147 + for (auto i = threadIdx.x; i < hidden_size; i += blockDim.x) {
148 + float x = static_cast<float>(input[token_offset + i]);
149 + if constexpr (has_residual) {
150 + x += static_cast<float>(residual[token_offset + i]);
151 + }
152
64 - x = static_cast<float>(static_cast<scalar_t>(x * rms) * weight[i]);
65 - block_absmax_val_maybe = fmaxf(block_absmax_val_maybe, fabsf(x));
66 - }
67 -
68 - using BlockReduce = cub::BlockReduce<float, 1024>;
69 - __shared__ typename BlockReduce::TempStorage reduceStore;
70 - block_absmax_val_maybe =
71 - BlockReduce(reduceStore)
72 - .Reduce(block_absmax_val_maybe, CubMaxOp{}, blockDim.x);
73 -
74 - __shared__ float s_token_scale;
75 - if (threadIdx.x == 0) {
76 - float scale = 0.0f;
77 - if (scale_ub) {
78 - scale = min(block_absmax_val_maybe, *scale_ub);
79 - } else {
80 - scale = block_absmax_val_maybe;
153 + x = static_cast<float>(static_cast<scalar_t>(x * rms) * weight[i]);
154 + block_absmax_val_maybe = fmaxf(block_absmax_val_maybe, fabsf(x));
155 }
82 - // token scale computation
83 - scale = max(scale / qmax, min_scaling_factor<scalar_out_t>::val());
84 - s_token_scale = scale; // Shared memory store
85 - all_token_scales[blockIdx.x] = scale; // Global output store
86 - }
87 - __syncthreads();
156 + using BlockReduce = cub::BlockReduce<float, 1024>;
157 + __shared__ typename BlockReduce::TempStorage reduceStore;
158 + block_absmax_val_maybe =
159 + BlockReduce(reduceStore)
160 + .Reduce(block_absmax_val_maybe, CubMaxOp{}, blockDim.x);
161 +
162 + __shared__ float s_token_scale;
163 + if (threadIdx.x == 0) {
164 + float scale = 0.0f;
165 + if (scale_ub) {
166 + scale = min(block_absmax_val_maybe, *scale_ub);
167 + } else {
168 + scale = block_absmax_val_maybe;
169 + }
170 + // token scale computation
171 + scale = max(scale / qmax, min_scaling_factor<scalar_out_t>::val());
172 + s_token_scale = scale; // Shared memory store
173 + all_token_scales[blockIdx.x] = scale; // Global output store
174 + }
175 + __syncthreads();
176
89 - *token_scale = s_token_scale;
177 + *token_scale = s_token_scale;
178 + }
179 }
180
181 template <typename scalar_t, typename scalar_out_t, bool is_scale_inverted,
93 - bool has_residual = false>
182 + bool has_residual = false, bool is_scale_transposed = false>
183 __device__ void norm_and_quant(scalar_out_t* __restrict__ output,
184 scalar_t const* __restrict__ input,
185 scalar_t const* __restrict__ weight,
97 - float const rms, float const scale,
186 + float const rms, float* const scale,
187 int32_t const hidden_size,
99 - scalar_t* __restrict__ residual = nullptr) {
188 + scalar_t* __restrict__ residual = nullptr,
189 + int32_t const group_size = 0) {
190 int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
101 - ;
191
192 for (auto i = threadIdx.x; i < hidden_size; i += blockDim.x) {
193 float x = static_cast<float>(input[token_offset + i]);
@@ -109,8 +198,21 @@ __device__ void norm_and_quant(scalar_out_t* __restrict__ output,
198 // Norm
199 x = static_cast<float>(static_cast<scalar_t>(x * rms) * weight[i]);
200 // Quant
201 + // If groupwise is_scale_inverted is true, so we invert the scale here.
202 + int64_t scale_idx = 0;
203 + if (group_size > 0) {
204 + if constexpr (is_scale_transposed) {
205 + scale_idx = (i / group_size) * gridDim.x + blockIdx.x;
206 + } else {
207 + scale_idx = blockIdx.x * (hidden_size / group_size) + i / group_size;
208 + }
209 + }
210 + auto scale_val =
211 + (group_size > 0
212 + ? (is_scale_inverted ? 1.0f / scale[scale_idx] : scale[scale_idx])
213 + : *scale);
214 output[token_offset + i] =
113 - ScaledQuant<scalar_out_t, is_scale_inverted>::quant_fn(x, scale);
215 + ScaledQuant<scalar_out_t, is_scale_inverted>::quant_fn(x, scale_val);
216 }
217 }
218
@@ -178,95 +280,191 @@ __device__ void compute_rms(float* rms, scalar_t const* __restrict__ input,
280
281 // Vectorized version of vllm::compute_dynamic_per_token_scales
282 // hidden_size must be a multiple of 4
181 -template <typename scalar_t, typename scalar_out_t, bool has_residual = false>
283 +template <typename scalar_t, typename scalar_out_t, bool has_residual = false,
284 + bool is_scale_transposed = false, int32_t group_size = 0>
285 __device__ void compute_dynamic_per_token_scales(
286 float* __restrict__ token_scale, float* __restrict__ all_token_scales,
287 scalar_t const* __restrict__ input, scalar_t const* __restrict__ weight,
288 float const rms, float const* __restrict__ scale_ub,
289 int32_t const hidden_size,
290 scalar_t const* __restrict__ residual = nullptr) {
188 - int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
189 - ;
190 -
191 - // Vectorized input/weight/residual to better utilize memory bandwidth.
192 - vec4_t<scalar_t> const* vec_input =
193 - reinterpret_cast<vec4_t<scalar_t> const*>(&input[token_offset]);
194 - vec4_t<scalar_t> const* vec_weight =
195 - reinterpret_cast<vec4_t<scalar_t> const*>(weight);
196 - vec4_t<scalar_t> const* vec_residual = nullptr;
197 - if constexpr (has_residual) {
198 - vec_residual =
199 - reinterpret_cast<vec4_t<scalar_t> const*>(&residual[token_offset]);
200 - }
201 -
291 constexpr scalar_out_t qmax{quant_type_max_v<scalar_out_t>};
292
293 const int VEC_SIZE = 4;
205 - int32_t const num_vec_elems = hidden_size >> 2;
294 float block_absmax_val_maybe = 0.0f;
295
296 + // Vectorized input/weight/residual to better utilize memory bandwidth.
297 + vec4_t<scalar_t> const* vec_input = nullptr;
298 + vec4_t<scalar_t> const* vec_weight = nullptr;
299 + vec4_t<scalar_t> const* vec_residual = nullptr;
300 +
301 + if constexpr (group_size > 0) {
302 + __shared__ float s_max_vals[1024];
303 +
304 + int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
305 + int64_t const num_groups = hidden_size / group_size;
306 + int64_t const threads_per_group = blockDim.x / num_groups;
307 + int64_t const thread_in_group = threadIdx.x % threads_per_group;
308 + int64_t const group_offset =
309 + threadIdx.x / threads_per_group * (group_size >> 2);
310 + int64_t const thread_offset = group_offset + thread_in_group;
311 + int64_t const thread_end = min(group_offset + (group_size >> 2),
312 + static_cast<int64_t>(hidden_size >> 2));
313 + vec_input = reinterpret_cast<vec4_t<scalar_t> const*>(&input[token_offset]);
314 + vec_weight = reinterpret_cast<vec4_t<scalar_t> const*>(weight);
315 + if constexpr (has_residual) {
316 + vec_residual =
317 + reinterpret_cast<vec4_t<scalar_t> const*>(&residual[token_offset]);
318 + }
319 + int32_t const num_vec_elems = thread_end;
320 +
321 #pragma unroll 4
209 - for (auto i = threadIdx.x; i < num_vec_elems; i += blockDim.x) {
210 - vec4_t<scalar_t> in = vec_input[i];
211 - vec4_t<scalar_t> const w = vec_weight[i];
322 + for (auto i = thread_offset; i < num_vec_elems; i += threads_per_group) {
323 + vec4_t<scalar_t> in = vec_input[i];
324 + vec4_t<scalar_t> const w = vec_weight[i];
325
213 - vec4_t<float> x;
326 + vec4_t<float> x;
327 #pragma unroll
215 - for (int j = 0; j < VEC_SIZE; ++j) {
216 - x.val[j] = static_cast<float>(in.val[j]);
328 + for (int j = 0; j < VEC_SIZE; ++j) {
329 + x.val[j] = static_cast<float>(in.val[j]);
330 + }
331 +
332 + if constexpr (has_residual) {
333 + vec4_t<scalar_t> r = vec_residual[i];
334 +#pragma unroll
335 + for (int j = 0; j < VEC_SIZE; ++j) {
336 + x.val[j] += static_cast<float>(r.val[j]);
337 + }
338 + }
339 +
340 +#pragma unroll
341 + for (int j = 0; j < VEC_SIZE; ++j) {
342 + block_absmax_val_maybe =
343 + fmaxf(block_absmax_val_maybe,
344 + fabs(static_cast<scalar_t>(x.val[j] * rms) * w.val[j]));
345 + }
346 }
347
348 + s_max_vals[threadIdx.x] = block_absmax_val_maybe;
349 + __syncthreads();
350 +
351 + int64_t const warp_size = WARP_SIZE;
352 + int64_t const num_warps = blockDim.x / warp_size;
353 + int64_t const warp_id = threadIdx.x / warp_size;
354 + int64_t const thread_in_warp = threadIdx.x % warp_size;
355 + int64_t const groups_per_warp = (num_groups + num_warps - 1) / num_warps;
356 + for (auto i = 0; i < groups_per_warp; ++i) {
357 + int64_t const group_id = i * num_warps + warp_id;
358 + if (group_id < num_groups) {
359 + int64_t warp_start = group_id * threads_per_group;
360 + int64_t const start = warp_start + thread_in_warp;
361 + int64_t const warp_end = min(warp_start + threads_per_group,
362 + static_cast<int64_t>(hidden_size));
363 + for (auto j = start; j + warp_size < warp_end; j += warp_size) {
364 + s_max_vals[start] =
365 + fmaxf(s_max_vals[start], s_max_vals[j + warp_size]);
366 + }
367 + warpReduceMaxSpecialized(s_max_vals, start, thread_in_warp,
368 + min(warp_end - warp_start, warp_size));
369 + }
370 + }
371 + __syncthreads();
372 +
373 + if (thread_in_group == 0 && thread_offset < thread_end) {
374 + block_absmax_val_maybe = s_max_vals[threadIdx.x];
375 + float scale = 0.0f;
376 + if (scale_ub) {
377 + scale = min(block_absmax_val_maybe, *scale_ub);
378 + } else {
379 + scale = block_absmax_val_maybe;
380 + }
381 + // token scale computation
382 + scale = max(scale / qmax, min_scaling_factor<scalar_out_t>::val());
383 + // Global output store
384 + if constexpr (is_scale_transposed) {
385 + all_token_scales[(threadIdx.x / threads_per_group) * gridDim.x +
386 + blockIdx.x] = scale;
387 + } else {
388 + all_token_scales[blockIdx.x * num_groups +
389 + threadIdx.x / threads_per_group] = scale;
390 + }
391 + }
392 + __syncthreads();
393 +
394 + } else {
395 + int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
396 + vec_input = reinterpret_cast<vec4_t<scalar_t> const*>(&input[token_offset]);
397 + vec_weight = reinterpret_cast<vec4_t<scalar_t> const*>(weight);
398 if constexpr (has_residual) {
220 - vec4_t<scalar_t> r = vec_residual[i];
399 + vec_residual =
400 + reinterpret_cast<vec4_t<scalar_t> const*>(&residual[token_offset]);
401 + }
402 +
403 + int32_t const num_vec_elems = (hidden_size >> 2);
404 +
405 +#pragma unroll 4
406 + for (auto i = threadIdx.x; i < num_vec_elems; i += blockDim.x) {
407 + vec4_t<scalar_t> in = vec_input[i];
408 + vec4_t<scalar_t> const w = vec_weight[i];
409 +
410 + vec4_t<float> x;
411 #pragma unroll
412 for (int j = 0; j < VEC_SIZE; ++j) {
223 - x.val[j] += static_cast<float>(r.val[j]);
413 + x.val[j] = static_cast<float>(in.val[j]);
414 }
225 - }
415
416 + if constexpr (has_residual) {
417 + vec4_t<scalar_t> r = vec_residual[i];
418 #pragma unroll
228 - for (int j = 0; j < VEC_SIZE; ++j) {
229 - block_absmax_val_maybe =
230 - fmaxf(block_absmax_val_maybe,
231 - fabs(static_cast<scalar_t>(x.val[j] * rms) * w.val[j]));
232 - }
233 - }
419 + for (int j = 0; j < VEC_SIZE; ++j) {
420 + x.val[j] += static_cast<float>(r.val[j]);
421 + }
422 + }
423
235 - using BlockReduce = cub::BlockReduce<float, 1024>;
236 - __shared__ typename BlockReduce::TempStorage reduceStore;
237 - block_absmax_val_maybe =
238 - BlockReduce(reduceStore)
239 - .Reduce(block_absmax_val_maybe, CubMaxOp{}, blockDim.x);
424 +#pragma unroll
425 + for (int j = 0; j < VEC_SIZE; ++j) {
426 + block_absmax_val_maybe =
427 + fmaxf(block_absmax_val_maybe,
428 + fabs(static_cast<scalar_t>(x.val[j] * rms) * w.val[j]));
429 + }
430 + }
431
241 - __shared__ float s_token_scale;
242 - if (threadIdx.x == 0) {
243 - float scale = 0.0f;
244 - if (scale_ub) {
245 - scale = min(block_absmax_val_maybe, *scale_ub);
246 - } else {
247 - scale = block_absmax_val_maybe;
432 + using BlockReduce = cub::BlockReduce<float, 1024>;
433 + __shared__ typename BlockReduce::TempStorage reduceStore;
434 + block_absmax_val_maybe =
435 + BlockReduce(reduceStore)
436 + .Reduce(block_absmax_val_maybe, CubMaxOp{}, blockDim.x);
437 +
438 + __shared__ float s_token_scale;
439 + if (threadIdx.x == 0) {
440 + float scale = 0.0f;
441 + if (scale_ub) {
442 + scale = min(block_absmax_val_maybe, *scale_ub);
443 + } else {
444 + scale = block_absmax_val_maybe;
445 + }
446 + // token scale computation
447 + scale = max(scale / qmax, min_scaling_factor<scalar_out_t>::val());
448 + s_token_scale = scale; // shared memory store
449 + all_token_scales[blockIdx.x] = scale; // global output store
450 }
249 - // token scale computation
250 - scale = max(scale / qmax, min_scaling_factor<scalar_out_t>::val());
251 - s_token_scale = scale; // shared memory store
252 - all_token_scales[blockIdx.x] = scale; // global output store
253 - }
254 - __syncthreads();
451 + __syncthreads();
452
256 - *token_scale = s_token_scale;
453 + *token_scale = s_token_scale;
454 + }
455 }
456
457 // hidden_size must be a multiple of 4
458 template <typename scalar_t, typename scalar_out_t, bool is_scale_inverted,
261 - bool has_residual = false>
459 + bool has_residual = false, bool is_scale_transposed = false,
460 + int32_t group_size = 0>
461 __device__ void norm_and_quant(scalar_out_t* __restrict__ output,
462 scalar_t const* __restrict__ input,
463 scalar_t const* __restrict__ weight,
265 - float const rms, float const scale,
464 + float const rms, float* const scale,
465 int32_t const hidden_size,
466 scalar_t* __restrict__ residual = nullptr) {
467 int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
269 - ;
468
469 // Vectorized input/output/weight/residual to better utilize memory bandwidth.
470 vec4_t<scalar_t> const* vec_input =
@@ -311,10 +509,26 @@ __device__ void norm_and_quant(scalar_out_t* __restrict__ output,
509 }
510
511 q8x4_t<scalar_out_t> out;
512 +
513 + float scale_val;
514 +
515 + if constexpr (group_size > 0) {
516 + int64_t const num_groups = hidden_size / group_size;
517 + int64_t scale_idx = 0;
518 + if constexpr (is_scale_transposed) {
519 + scale_idx = (i * VEC_SIZE / group_size) * gridDim.x + blockIdx.x;
520 + } else {
521 + scale_idx = blockIdx.x * num_groups + i * VEC_SIZE / group_size;
522 + }
523 + scale_val =
524 + is_scale_inverted ? 1.0f / scale[scale_idx] : scale[scale_idx];
525 + } else {
526 + scale_val = *scale;
527 + }
528 #pragma unroll
529 for (int j = 0; j < VEC_SIZE; ++j) {
530 out.val[j] = ScaledQuant<scalar_out_t, is_scale_inverted>::quant_fn(
317 - static_cast<scalar_t>(x.val[j] * rms) * w.val[j], scale);
531 + static_cast<scalar_t>(x.val[j] * rms) * w.val[j], scale_val);
532 }
533 vec_output[i] = out;
534 }
csrc/torch_bindings.cpp
+8
@@ -215,6 +215,14 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
215 ops.impl("rms_norm_dynamic_per_token_quant", torch::kCUDA,
216 &rms_norm_dynamic_per_token_quant);
217
218 + // Fused Layernorm + Block quant kernels
219 + ops.def(
220 + "rms_norm_per_block_quant(Tensor! result, Tensor input, "
221 + "Tensor weight, Tensor! scale, float epsilon, "
222 + "Tensor? scale_ub, Tensor!? residual, int group_size, "
223 + "bool is_scale_transposed) -> ()");
224 + ops.impl("rms_norm_per_block_quant", torch::kCUDA, &rms_norm_per_block_quant);
225 +
226 // Rotary embedding
227 // Apply GPT-NeoX or GPT-J style rotary embedding to query and key.
228 ops.def(
tests/compile/test_fusion.py
+55 -14
@@ -18,6 +18,9 @@ from vllm.config import (
18 VllmConfig,
19 )
20 from vllm.model_executor.layers.layernorm import RMSNorm
21 +from vllm.model_executor.layers.quantization.utils.fp8_utils import (
22 + W8A8BlockFp8LinearOp,
23 +)
24 from vllm.model_executor.layers.quantization.utils.quant_utils import (
25 GroupShape,
26 QuantKey,
@@ -25,10 +28,12 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
28 )
29 from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
30 Fp8LinearOp,
31 + cutlass_block_fp8_supported,
32 cutlass_fp8_supported,
33 maybe_create_device_identity,
34 )
35 from vllm.platforms import current_platform
36 +from vllm.utils.deep_gemm import is_deep_gemm_supported
37
38 from ..utils import override_cutlass_fp8_supported
39 from .backend import TestBackend
@@ -44,7 +49,7 @@ class TestModel(torch.nn.Module):
49 self,
50 hidden_size: int,
51 eps: float,
47 - static: bool,
52 + group_shape: GroupShape,
53 cuda_force_torch: bool,
54 *args,
55 **kwargs,
@@ -52,8 +57,17 @@ class TestModel(torch.nn.Module):
57 super().__init__(*args, **kwargs)
58 self.cuda_force_torch = cuda_force_torch
59 self.norm = [RMSNorm(hidden_size, eps) for _ in range(4)]
55 - self.wscale = [torch.rand(1, dtype=torch.float32) for _ in range(3)]
56 - group_shape = GroupShape.PER_TENSOR if static else GroupShape.PER_TOKEN
60 + if group_shape.is_per_group():
61 + self.wscale = [
62 + torch.rand(
63 + (hidden_size // group_shape[1], hidden_size // group_shape[1]),
64 + dtype=torch.float32,
65 + )
66 + for _ in range(3)
67 + ]
68 + else:
69 + self.wscale = [torch.rand(1, dtype=torch.float32) for _ in range(3)]
70 + static = group_shape == GroupShape.PER_TENSOR
71 quant_scale = ScaleDesc(torch.float32, static, group_shape)
72 self.quant_key = QuantKey(dtype=FP8_DTYPE, scale=quant_scale, symmetric=True)
73 if static:
@@ -61,18 +75,29 @@ class TestModel(torch.nn.Module):
75 else:
76 self.scale = [None for _ in range(3)]
77 self.w = [
64 - torch.rand(hidden_size, hidden_size).to(dtype=FP8_DTYPE).t()
65 - for _ in range(3)
78 + torch.rand(hidden_size, hidden_size).to(dtype=FP8_DTYPE) for _ in range(3)
79 ]
80 + if not group_shape.is_per_group():
81 + self.w = [self.w[0].t() for _ in range(3)]
82
68 - with override_cutlass_fp8_supported(not cuda_force_torch):
69 - self.fp8_linear = Fp8LinearOp(
70 - act_quant_static=static,
83 + if group_shape.is_per_group():
84 + self.fp8_linear = W8A8BlockFp8LinearOp(
85 + weight_group_shape=GroupShape(group_shape[1], group_shape[1]),
86 act_quant_group_shape=group_shape,
87 + cutlass_block_fp8_supported=cutlass_block_fp8_supported(),
88 + use_aiter_and_is_supported=False,
89 )
90 + self.enable_quant_fp8_custom_op = self.fp8_linear.input_quant_op.enabled()
91 + else:
92 + with override_cutlass_fp8_supported(not cuda_force_torch):
93 + self.fp8_linear = Fp8LinearOp(
94 + act_quant_static=static,
95 + act_quant_group_shape=group_shape,
96 + )
97 + self.enable_quant_fp8_custom_op = self.fp8_linear.quant_fp8.enabled()
98
99 self.enable_rms_norm_custom_op = self.norm[0].enabled()
75 - self.enable_quant_fp8_custom_op = self.fp8_linear.quant_fp8.enabled()
100 + self.group_shape = group_shape
101
102 def forward(self, x):
103 # avoid having graph input be an arg to a pattern directly
@@ -119,11 +144,19 @@ class TestModel(torch.nn.Module):
144 )
145
146
147 +GROUP_SHAPES = [
148 + GroupShape.PER_TOKEN,
149 + GroupShape.PER_TENSOR,
150 + GroupShape(1, 128),
151 + GroupShape(1, 64),
152 +]
153 +
154 +
155 @pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
123 -@pytest.mark.parametrize("hidden_size", [64])
156 +@pytest.mark.parametrize("hidden_size", [256])
157 @pytest.mark.parametrize("num_tokens", [257])
158 @pytest.mark.parametrize("eps", [1e-5, 1e-6])
126 -@pytest.mark.parametrize("static", [True, False])
159 +@pytest.mark.parametrize("group_shape", GROUP_SHAPES)
160 @pytest.mark.parametrize("enable_rms_norm_custom_op", [True, False])
161 @pytest.mark.parametrize("enable_quant_fp8_custom_op", [True, False])
162 # cuda_force_torch used to test torch code path on platforms that
@@ -139,7 +172,7 @@ def test_fusion_rmsnorm_quant(
172 hidden_size,
173 num_tokens,
174 eps,
142 - static,
175 + group_shape,
176 enable_rms_norm_custom_op,
177 enable_quant_fp8_custom_op,
178 cuda_force_torch,
@@ -149,6 +182,15 @@ def test_fusion_rmsnorm_quant(
182 torch.manual_seed(1)
183 maybe_create_device_identity() # needed for certain non-cutlass fp8 paths
184
185 + if not enable_quant_fp8_custom_op and group_shape.is_per_group():
186 + pytest.skip("Unsupported unwrapped quant fp8 op for blockwise quantization")
187 +
188 + # Skip test for 64-bit group shape when running with cutlass or deepgemm
189 + if group_shape == GroupShape(1, 64) and (
190 + cutlass_block_fp8_supported() or is_deep_gemm_supported()
191 + ):
192 + pytest.skip("Unsupported group shape 64 for CUTLASS/DeepGemm")
193 +
194 custom_ops = []
195 if enable_rms_norm_custom_op:
196 custom_ops.append("+rms_norm")
@@ -172,8 +214,7 @@ def test_fusion_rmsnorm_quant(
214
215 backend = TestBackend(noop_pass, fusion_pass, cleanup_pass)
216 backend2 = TestBackend(noop_pass, cleanup_pass)
175 - model = TestModel(hidden_size, eps, static, cuda_force_torch)
176 -
217 + model = TestModel(hidden_size, eps, group_shape, cuda_force_torch)
218 # First dimension dynamic
219 x = torch.rand(num_tokens, hidden_size)
220 torch._dynamo.mark_dynamic(x, 0)
tests/kernels/core/test_fused_quant_layernorm.py
+63 -19
@@ -8,6 +8,12 @@ import torch
8 import vllm._custom_ops as ops
9 from tests.kernels.utils import opcheck
10 from vllm.model_executor.layers.layernorm import RMSNorm
11 +from vllm.model_executor.layers.quantization.utils.fp8_utils import (
12 + per_token_group_quant_fp8,
13 +)
14 +from vllm.model_executor.layers.quantization.utils.int8_utils import (
15 + per_token_group_quant_int8,
16 +)
17
18 DTYPES = [torch.bfloat16, torch.float]
19 QUANT_DTYPES = [torch.int8, torch.float8_e4m3fn]
@@ -21,6 +27,7 @@ NUM_TOKENS_HIDDEN_SIZES = [
27
28 ADD_RESIDUAL = [False, True]
29 SCALE_UBS = [True, False]
30 +GROUP_SIZES = [None, [1, 64], [1, 128]]
31 SEEDS = [0]
32 CUDA_DEVICES = [f"cuda:{i}" for i in range(1 if torch.cuda.device_count() == 1 else 2)]
33
@@ -45,12 +52,13 @@ def ref_rms_norm(
52 return out, residual
53
54
48 -def ref_dynamic_per_token_quant(
55 +def ref_dynamic_per_token_or_block_quant(
56 rms_norm_layer: RMSNorm,
57 x: torch.Tensor,
58 quant_dtype: torch.dtype,
59 residual: torch.Tensor | None,
60 scale_ub: torch.Tensor | None,
61 + group_size: list[int] | None,
62 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
63 if scale_ub is not None:
64 assert quant_dtype == torch.float8_e4m3fn
@@ -59,13 +67,24 @@ def ref_dynamic_per_token_quant(
67 torch_out, residual = ref_rms_norm(rms_norm_layer, x, residual)
68
69 # Quant
62 - if quant_dtype == torch.float8_e4m3fn:
63 - torch_out, scales = ops.scaled_fp8_quant(
64 - torch_out, scale_ub=scale_ub, use_per_token_if_dynamic=True
65 - )
70 + if group_size is not None:
71 + if quant_dtype == torch.float8_e4m3fn:
72 + torch_out, scales = per_token_group_quant_fp8(
73 + torch_out, group_size=group_size[1], use_ue8m0=False
74 + )
75 + else:
76 + assert quant_dtype == torch.int8
77 + torch_out, scales = per_token_group_quant_int8(
78 + torch_out, group_size=group_size[1]
79 + )
80 else:
67 - assert quant_dtype == torch.int8
68 - torch_out, scales, _ = ops.scaled_int8_quant(torch_out)
81 + if quant_dtype == torch.float8_e4m3fn:
82 + torch_out, scales = ops.scaled_fp8_quant(
83 + torch_out, scale_ub=scale_ub, use_per_token_if_dynamic=True
84 + )
85 + else:
86 + assert quant_dtype == torch.int8
87 + torch_out, scales, _ = ops.scaled_int8_quant(torch_out)
88
89 return torch_out, scales, residual
90
@@ -76,24 +95,32 @@ def ref_impl(
95 quant_dtype: torch.dtype,
96 residual: torch.Tensor | None,
97 scale_ub: torch.Tensor | None,
98 + group_size: list[int] | None,
99 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
80 - return ref_dynamic_per_token_quant(
81 - rms_norm_layer, x, quant_dtype, residual, scale_ub
100 + return ref_dynamic_per_token_or_block_quant(
101 + rms_norm_layer, x, quant_dtype, residual, scale_ub, group_size
102 )
103
104
85 -def ops_dynamic_per_token_quant(
105 +def ops_dynamic_per_token_or_block_quant(
106 weight: torch.Tensor,
107 x: torch.Tensor,
108 quant_dtype: torch.dtype,
109 residual: torch.Tensor | None,
110 scale_ub: torch.Tensor | None,
111 + group_size: list[int] | None,
112 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
113 if residual is not None:
114 residual = residual.clone()
94 - out, scales = ops.rms_norm_dynamic_per_token_quant(
95 - x, weight, EPS, quant_dtype, scale_ub, residual
96 - )
115 + if group_size is not None:
116 + out, scales = ops.rms_norm_per_block_quant(
117 + x, weight, EPS, quant_dtype, group_size, scale_ub, residual, True
118 + )
119 + scales = scales.contiguous()
120 + else:
121 + out, scales = ops.rms_norm_dynamic_per_token_quant(
122 + x, weight, EPS, quant_dtype, scale_ub, residual
123 + )
124 return out, scales, residual
125
126
@@ -103,8 +130,11 @@ def ops_impl(
130 quant_dtype: torch.dtype,
131 residual: torch.Tensor | None,
132 scale_ub: torch.Tensor | None,
133 + group_size: list[int] | None,
134 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor | None]:
107 - return ops_dynamic_per_token_quant(weight, x, quant_dtype, residual, scale_ub)
135 + return ops_dynamic_per_token_or_block_quant(
136 + weight, x, quant_dtype, residual, scale_ub, group_size
137 + )
138
139
140 @pytest.mark.parametrize("num_tokens, hidden_size", NUM_TOKENS_HIDDEN_SIZES)
@@ -112,6 +142,7 @@ def ops_impl(
142 @pytest.mark.parametrize("has_scale_ub", SCALE_UBS)
143 @pytest.mark.parametrize("dtype", DTYPES)
144 @pytest.mark.parametrize("quant_dtype", QUANT_DTYPES)
145 +@pytest.mark.parametrize("group_size", GROUP_SIZES)
146 @pytest.mark.parametrize("seed", SEEDS)
147 @pytest.mark.parametrize("device", CUDA_DEVICES)
148 @torch.inference_mode()
@@ -122,6 +153,7 @@ def test_rms_norm(
153 has_scale_ub: bool,
154 dtype: torch.dtype,
155 quant_dtype: torch.dtype,
156 + group_size: list[int] | None,
157 seed: int,
158 device: str,
159 ) -> None:
@@ -130,6 +162,14 @@ def test_rms_norm(
162 torch.cuda.manual_seed(seed)
163 torch.set_default_device(device)
164
165 + if group_size is not None and hidden_size % group_size[1] != 0:
166 + # skip
167 + return
168 +
169 + if group_size is not None and has_scale_ub:
170 + # blockwise baseline doesn't support scale_ub
171 + return
172 +
173 if has_scale_ub and quant_dtype != torch.float8_e4m3fn:
174 # skip
175 return
@@ -150,10 +190,10 @@ def test_rms_norm(
190 scale_ub = None
191
192 ref_out, ref_scales, ref_residual = ref_impl(
153 - layer, x, quant_dtype, residual, scale_ub
193 + layer, x, quant_dtype, residual, scale_ub, group_size
194 )
195 ops_out, ops_scales, ops_residual = ops_impl(
156 - layer.weight, x, quant_dtype, residual, scale_ub
196 + layer.weight, x, quant_dtype, residual, scale_ub, group_size
197 )
198
199 assert ref_out.dtype == quant_dtype
@@ -166,11 +206,15 @@ def test_rms_norm(
206 assert torch.allclose(ref_scales, ops_scales)
207 a = ref_out.to(dtype=torch.float32)
208 b = ops_out.to(dtype=torch.float32)
169 - ok = torch.allclose(a, b)
209 + ok = torch.allclose(a, b, atol=1e-6)
210 if not ok:
211 # fallback: compare dequantized values with relaxed tolerance
172 - a_deq = a * ref_scales.view(-1, 1)
173 - b_deq = b * ops_scales.view(-1, 1)
212 + if group_size is None:
213 + a_deq = a * ref_scales.view(-1, 1)
214 + b_deq = b * ops_scales.view(-1, 1)
215 + else:
216 + a_deq = a * ref_scales.repeat_interleave(group_size[1], dim=1)
217 + b_deq = b * ops_scales.repeat_interleave(group_size[1], dim=1)
218 # NOTE: It is possible that some future test cases trigger this
219 # max diff due to precision issues. If such an error is
220 # encountered, it's recommended to inspect the differences between
vllm/_custom_ops.py
+40
@@ -436,6 +436,46 @@ def rms_norm_dynamic_per_token_quant(
436 return output, scales
437
438
439 +# fused quant layer norm ops blocked
440 +def rms_norm_per_block_quant(
441 + input: torch.Tensor,
442 + weight: torch.Tensor,
443 + epsilon: float,
444 + quant_dtype: torch.dtype,
445 + group_size: list[int],
446 + scale_ub: torch.Tensor | None = None,
447 + residual: torch.Tensor | None = None,
448 + is_scale_transposed: bool = False,
449 +) -> tuple[torch.Tensor, torch.Tensor]:
450 + assert len(group_size) == 2
451 + output = torch.empty_like(input, dtype=quant_dtype)
452 + if is_scale_transposed:
453 + scales = torch.empty(
454 + (input.shape[-1] // group_size[1], input.numel() // input.shape[-1]),
455 + device=input.device,
456 + dtype=torch.float32,
457 + ).transpose(0, 1)
458 + else:
459 + scales = torch.empty(
460 + (input.numel() // input.shape[-1], input.shape[-1] // group_size[1]),
461 + device=input.device,
462 + dtype=torch.float32,
463 + )
464 +
465 + torch.ops._C.rms_norm_per_block_quant(
466 + output,
467 + input,
468 + weight,
469 + scales,
470 + epsilon,
471 + scale_ub,
472 + residual,
473 + group_size[1],
474 + is_scale_transposed,
475 + )
476 + return output, scales
477 +
478 +
479 # quantization ops
480 # awq
481 def awq_dequantize(
vllm/compilation/fusion.py
+171 -1
@@ -15,13 +15,22 @@ from vllm.model_executor.layers.quantization.utils.quant_utils import (
15 GroupShape,
16 QuantKey,
17 ScaleDesc,
18 + kFp8Dynamic64Sym,
19 + kFp8Dynamic128Sym,
20 kFp8DynamicTensorSym,
21 kFp8DynamicTokenSym,
22 kFp8StaticTensorSym,
23 kNvfp4Quant,
24 kStaticTensorScale,
25 )
26 +from vllm.model_executor.layers.quantization.utils.w8a8_utils import (
27 + cutlass_block_fp8_supported,
28 +)
29 from vllm.platforms import current_platform
30 +from vllm.utils.deep_gemm import (
31 + is_deep_gemm_e8m0_used,
32 + should_use_deepgemm_for_fp8_linear_for_nk,
33 +)
34
35 from .inductor_pass import enable_fake_mode
36 from .matcher_utils import MatcherFusedAddRMSNorm, MatcherQuantFP8, MatcherRMSNorm
@@ -58,6 +67,9 @@ QUANT_OPS: dict[QuantKey, OpOverload] = {
67 }
68 if current_platform.is_cuda() and hasattr(torch.ops._C, "scaled_fp4_quant"):
69 QUANT_OPS[kNvfp4Quant] = torch.ops._C.scaled_fp4_quant.default
70 +if current_platform.is_cuda():
71 + QUANT_OPS[kFp8Dynamic128Sym] = torch.ops._C.per_token_group_fp8_quant.default # noqa: E501
72 + QUANT_OPS[kFp8Dynamic64Sym] = torch.ops._C.per_token_group_fp8_quant.default # noqa: E501
73
74
75 class FusedRMSQuantKey(NamedTuple):
@@ -90,6 +102,18 @@ FUSED_OPS: dict[FusedRMSQuantKey, OpOverload] = {
102 FusedRMSQuantKey(
103 kFp8DynamicTokenSym, True
104 ): torch.ops._C.rms_norm_dynamic_per_token_quant.default, # noqa: E501
105 + FusedRMSQuantKey(
106 + kFp8Dynamic128Sym, False
107 + ): torch.ops._C.rms_norm_per_block_quant.default, # noqa: E501
108 + FusedRMSQuantKey(
109 + kFp8Dynamic128Sym, True
110 + ): torch.ops._C.rms_norm_per_block_quant.default, # noqa: E501
111 + FusedRMSQuantKey(
112 + kFp8Dynamic64Sym, False
113 + ): torch.ops._C.rms_norm_per_block_quant.default, # noqa: E501
114 + FusedRMSQuantKey(
115 + kFp8Dynamic64Sym, True
116 + ): torch.ops._C.rms_norm_per_block_quant.default, # noqa: E501
117 }
118
119
@@ -100,6 +124,15 @@ class RMSNormQuantPattern:
124 config = get_current_vllm_config()
125 self.model_dtype = config.model_config.dtype if config.model_config else None
126
127 + # groupwise FP8 linear uses col major scales if deepgemm and cutlass
128 + using_deepgemm = should_use_deepgemm_for_fp8_linear_for_nk(
129 + self.model_dtype,
130 + config.model_config.hf_config.intermediate_size,
131 + config.model_config.hf_config.hidden_size,
132 + )
133 + use_col_major_scales = using_deepgemm or cutlass_block_fp8_supported()
134 + use_e8m0 = is_deep_gemm_e8m0_used() if using_deepgemm else False
135 +
136 assert key in FUSED_OPS, f"unsupported fused rmsnorm+quant op for {key}"
137 self.FUSED_OP = FUSED_OPS[key]
138
@@ -108,7 +141,9 @@ class RMSNormQuantPattern:
141 if not key.fused_add
142 else MatcherFusedAddRMSNorm(epsilon)
143 )
111 - self.quant_matcher = MatcherQuantFP8(key.quant)
144 + self.quant_matcher = MatcherQuantFP8(
145 + key.quant, use_col_major_scales=use_col_major_scales, use_e8m0=use_e8m0
146 + )
147
148
149 class RMSNormStaticQuantPattern(RMSNormQuantPattern):
@@ -218,6 +253,120 @@ class FusedAddRMSNormStaticQuantPattern(RMSNormQuantPattern):
253 )
254
255
256 +class FusedAddRMSNormGroupQuantPattern(RMSNormQuantPattern):
257 + def __init__(
258 + self,
259 + epsilon: float,
260 + quant_dtype: torch.dtype,
261 + group_shape: GroupShape,
262 + symmetric=True,
263 + ):
264 + scale = ScaleDesc(torch.float32, False, group_shape)
265 + key = FusedRMSQuantKey(
266 + fused_add=True,
267 + quant=QuantKey(dtype=quant_dtype, scale=scale, symmetric=symmetric),
268 + )
269 + self.group_shape = group_shape
270 + super().__init__(epsilon, key)
271 +
272 + def register(self, pm_pass: PatternMatcherPass):
273 + def pattern(input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor):
274 + result_rms, residual = self.rmsnorm_matcher(input, weight, residual)
275 + result, scale = self.quant_matcher(result_rms)
276 + return result, residual, scale
277 +
278 + def replacement(
279 + input: torch.Tensor, weight: torch.Tensor, residual: torch.Tensor
280 + ):
281 + # In case we're matching native rms-norm, conversions might be
282 + # optimized out. We convert here just to be safe.
283 + input = input.to(dtype=self.model_dtype)
284 +
285 + result = torch.empty_like(input, dtype=self.quant_dtype)
286 + scale = self.quant_matcher.make_scale(
287 + input, transposed=self.quant_matcher.use_col_major_scales
288 + )
289 + at = auto_functionalized(
290 + self.FUSED_OP,
291 + result=result,
292 + input=input,
293 + weight=weight,
294 + scale=scale,
295 + epsilon=self.epsilon,
296 + scale_ub=None,
297 + residual=residual,
298 + group_size=self.group_shape[1],
299 + is_scale_transposed=self.quant_matcher.use_col_major_scales,
300 + )
301 +
302 + # result, residual, scale
303 + return at[1], at[3], at[2]
304 +
305 + pm.register_replacement(
306 + pattern,
307 + replacement,
308 + self.rmsnorm_matcher.inputs(),
309 + pm.fwd_only,
310 + pm_pass,
311 + )
312 +
313 +
314 +class RMSNormGroupQuantPattern(RMSNormQuantPattern):
315 + def __init__(
316 + self,
317 + epsilon: float,
318 + quant_dtype: torch.dtype,
319 + group_shape: GroupShape,
320 + symmetric=True,
321 + ):
322 + scale = ScaleDesc(torch.float32, False, group_shape)
323 + key = FusedRMSQuantKey(
324 + fused_add=False,
325 + quant=QuantKey(dtype=quant_dtype, scale=scale, symmetric=symmetric),
326 + )
327 + self.group_shape = group_shape
328 + super().__init__(epsilon, key)
329 +
330 + def register(self, pm_pass: PatternMatcherPass):
331 + def pattern(input: torch.Tensor, weight: torch.Tensor):
332 + result_rms = self.rmsnorm_matcher(input, weight)
333 + result, scale = self.quant_matcher(result_rms)
334 + return result, scale
335 +
336 + def replacement(input: torch.Tensor, weight: torch.Tensor):
337 + # In case we're matching native rms-norm, conversions might be
338 + # optimized out. We convert here just to be safe.
339 + input = input.to(dtype=self.model_dtype)
340 +
341 + result = torch.empty_like(input, dtype=self.quant_dtype)
342 + scale = self.quant_matcher.make_scale(
343 + input, transposed=self.quant_matcher.use_col_major_scales
344 + )
345 + at = auto_functionalized(
346 + self.FUSED_OP,
347 + result=result,
348 + input=input,
349 + weight=weight,
350 + scale=scale,
351 + epsilon=self.epsilon,
352 + scale_ub=None,
353 + residual=None,
354 + group_size=self.group_shape[1],
355 + is_scale_transposed=self.quant_matcher.use_col_major_scales,
356 + )
357 +
358 + # result, scale
359 + return at[1], at[2]
360 +
361 + pm.register_replacement(
362 + pattern,
363 + replacement,
364 + self.rmsnorm_matcher.inputs(),
365 + pm.fwd_only,
366 + pm_pass,
367 + )
368 +
369 +
370 class RMSNormDynamicQuantPattern(RMSNormQuantPattern):
371 def __init__(
372 self,
@@ -340,6 +489,25 @@ class RMSNormQuantFusionPass(VllmPatternMatcherPass):
489 # Make sure fused add patterns are before simple rms norm,
490 # as the latter is a subset of the former in torch ops
491 for epsilon in [1e-5, 1e-6]:
492 + # Fuse fused_add_rms_norm + fp8 group quant
493 + FusedAddRMSNormGroupQuantPattern(
494 + epsilon, FP8_DTYPE, group_shape=GroupShape(1, 128)
495 + ).register(self.patterns)
496 +
497 + # Fuse rms_norm + fp8 group quant
498 + RMSNormGroupQuantPattern(
499 + epsilon, FP8_DTYPE, group_shape=GroupShape(1, 128)
500 + ).register(self.patterns)
501 +
502 + FusedAddRMSNormGroupQuantPattern(
503 + epsilon, FP8_DTYPE, group_shape=GroupShape(1, 64)
504 + ).register(self.patterns)
505 +
506 + # Fuse rms_norm + fp8 group quant
507 + RMSNormGroupQuantPattern(
508 + epsilon, FP8_DTYPE, group_shape=GroupShape(1, 64)
509 + ).register(self.patterns)
510 +
511 # Fuse fused_add_rms_norm + static fp8 quant
512 FusedAddRMSNormStaticQuantPattern(epsilon, FP8_DTYPE).register(
513 self.patterns
@@ -366,9 +534,11 @@ class RMSNormQuantFusionPass(VllmPatternMatcherPass):
534 def uuid(self) -> Any:
535 return self.hash_source(
536 self,
537 + RMSNormGroupQuantPattern,
538 RMSNormQuantPattern,
539 RMSNormStaticQuantPattern,
540 RMSNormDynamicQuantPattern,
541 FusedAddRMSNormStaticQuantPattern,
542 FusedAddRMSNormDynamicQuantPattern,
543 + FusedAddRMSNormGroupQuantPattern,
544 )
vllm/compilation/matcher_utils.py
+42 -2
@@ -13,6 +13,8 @@ from vllm.model_executor.layers.quantization.input_quant_fp8 import QuantFP8
13 from vllm.model_executor.layers.quantization.utils.quant_utils import (
14 QuantKey,
15 _normalize_quant_group_shape,
16 + kFp8Dynamic64Sym,
17 + kFp8Dynamic128Sym,
18 kFp8DynamicTensorSym,
19 kFp8DynamicTokenSym,
20 kFp8StaticTensorSym,
@@ -35,6 +37,10 @@ QUANT_OPS: dict[QuantKey, OpOverload] = {
37 if current_platform.is_cuda() and hasattr(torch.ops._C, "scaled_fp4_quant"):
38 QUANT_OPS[kNvfp4Quant] = torch.ops._C.scaled_fp4_quant.default # noqa: E501
39
40 +if current_platform.is_cuda():
41 + QUANT_OPS[kFp8Dynamic128Sym] = torch.ops._C.per_token_group_fp8_quant.default # noqa: E501
42 + QUANT_OPS[kFp8Dynamic64Sym] = torch.ops._C.per_token_group_fp8_quant.default # noqa: E501
43 +
44 SILU_MUL_OP = torch.ops._C.silu_and_mul.default
45
46
@@ -224,12 +230,20 @@ class MatcherFusedAddRMSNorm(MatcherCustomOp):
230
231
232 class MatcherQuantFP8(MatcherCustomOp):
227 - def __init__(self, quant_key: QuantKey, enabled: bool | None = None):
233 + def __init__(
234 + self,
235 + quant_key: QuantKey,
236 + enabled: bool | None = None,
237 + use_col_major_scales: bool = False,
238 + use_e8m0: bool = False,
239 + ):
240 if enabled is None:
241 enabled = QuantFP8.enabled()
242
243 super().__init__(enabled)
244 self.quant_key = quant_key
245 + self.use_col_major_scales = use_col_major_scales
246 + self.use_e8m0 = use_e8m0
247 assert quant_key in QUANT_OPS, f"unsupported quantization scheme {quant_key}"
248 self.QUANT_OP = QUANT_OPS[quant_key]
249
@@ -248,6 +262,27 @@ class MatcherQuantFP8(MatcherCustomOp):
262 input.shape, device=input.device, dtype=self.quant_key.dtype
263 )
264
265 + if self.quant_key.scale.group_shape.is_per_group():
266 + assert scale is None
267 + scale = self.make_scale(input, transposed=self.use_col_major_scales)
268 +
269 + finfo = torch.finfo(self.quant_key.dtype)
270 + fp8_min = finfo.min
271 + fp8_max = finfo.max
272 +
273 + _, result, scale = auto_functionalized(
274 + self.QUANT_OP,
275 + input=input,
276 + output_q=result,
277 + output_s=scale,
278 + group_size=self.quant_key.scale.group_shape[1],
279 + eps=1e-10,
280 + fp8_min=fp8_min,
281 + fp8_max=fp8_max,
282 + scale_ue8m0=self.use_e8m0,
283 + )
284 + return result, scale
285 +
286 if self.quant_key.scale.static:
287 assert scale is not None
288 _, result = auto_functionalized(
@@ -269,7 +304,7 @@ class MatcherQuantFP8(MatcherCustomOp):
304 ) -> tuple[torch.Tensor, torch.Tensor]:
305 return self.quant_fp8(input, scale)
306
272 - def make_scale(self, input: torch.Tensor):
307 + def make_scale(self, input: torch.Tensor, transposed: bool = False):
308 normalized_group_shape = _normalize_quant_group_shape(
309 input, self.quant_key.scale.group_shape
310 )
@@ -277,6 +312,11 @@ class MatcherQuantFP8(MatcherCustomOp):
312 input.shape[0] // normalized_group_shape[0],
313 input.shape[1] // normalized_group_shape[1],
314 )
315 + if transposed:
316 + scale_shape = tuple(reversed(scale_shape))
317 + return torch.empty(
318 + scale_shape, device=input.device, dtype=torch.float32
319 + ).permute(-1, -2)
320
321 return torch.empty(scale_shape, device=input.device, dtype=torch.float32)
322
vllm/model_executor/layers/quantization/utils/fp8_utils.py
+1 -1
@@ -733,7 +733,7 @@ def per_token_group_quant_fp8(
733 assert out_q is None or out_q.shape == x.shape
734 x_q = out_q
735 if x_q is None:
736 - x_q = torch.empty_like(x, device=x.device, dtype=dtype)
736 + x_q = torch.empty(x.shape, device=x.device, dtype=dtype)
737
738 # Allocate the scale tensor in either row- or column-major format.
739 if column_major_scales:
vllm/model_executor/layers/quantization/utils/quant_utils.py
+6
@@ -115,6 +115,12 @@ kFp8DynamicTokenSym = QuantKey(FP8_DTYPE, kDynamicTokenScale, symmetric=True)
115 kNvfp4GroupScale = ScaleDesc(FP8_DTYPE, False, GroupShape(1, 16))
116 kNvfp4Quant = QuantKey(FP4_DTYPE, scale=kNvfp4GroupScale, scale2=kStaticTensorScale)
117
118 +kDynamic128Scale = ScaleDesc(torch.float32, False, GroupShape(1, 128))
119 +kFp8Dynamic128Sym = QuantKey(FP8_DTYPE, kDynamic128Scale, symmetric=True)
120 +
121 +kDynamic64Scale = ScaleDesc(torch.float32, False, GroupShape(1, 64))
122 +kFp8Dynamic64Sym = QuantKey(FP8_DTYPE, kDynamic64Scale, symmetric=True)
123 +
124
125 # Normalize the group_shape to the full extent for any dims that are -1
126 def _normalize_quant_group_shape(x: torch.Tensor, group_shape: GroupShape):
vllm/utils/deep_gemm.py
+17
@@ -381,6 +381,22 @@ def should_use_deepgemm_for_fp8_linear(
381 )
382
383
384 +def should_use_deepgemm_for_fp8_linear_for_nk(
385 + output_dtype: torch.dtype,
386 + shape0: int,
387 + shape1: int,
388 + supports_deep_gemm: bool | None = None,
389 +):
390 + if supports_deep_gemm is None:
391 + supports_deep_gemm = is_deep_gemm_supported()
392 + return (
393 + supports_deep_gemm
394 + and output_dtype == torch.bfloat16
395 + and shape0 % 128 == 0
396 + and shape1 % 128 == 0
397 + )
398 +
399 +
400 __all__ = [
401 "calc_diff",
402 "fp8_gemm_nt",
@@ -394,6 +410,7 @@ __all__ = [
410 "is_deep_gemm_supported",
411 "get_num_sms",
412 "should_use_deepgemm_for_fp8_linear",
413 + "should_use_deepgemm_for_fp8_linear_for_nk",
414 "get_col_major_tma_aligned_tensor",
415 "get_mk_alignment_for_contiguous_layout",
416 ]