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 6, 20:30 UTC
[torch.compile] Add support for non-contiguous fused RMSNorm + group quant (#36551)
Signed-off-by: Luka Govedič <lgovedic@redhat.com> Signed-off-by: Luka Govedič <ProExpertProg@users.noreply.github.com> Co-authored-by: Copilot <198982749+Copilot@users.noreply.github.com> Co-authored-by: ProExpertProg <11367180+ProExpertProg@users.noreply.github.com>
Luka Govedič committed
Mar 11, 2026 at 13:56 UTC
9556af87d5d5a38128db0d09eeb7f2fe16f16589
9 files changed
+219
-87
.buildkite/test_areas/compile.yaml
+9
-9
@@ -101,8 +101,8 @@ steps:
101
- nvidia-smi
102
# Run all models and attn backends but only Inductor partition and native custom ops
103
- pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and not +rms_norm and not +quant_fp8"
104
- # Qwen requires +quant_fp8 as -quant_fp8 rms+quant fusion is not supported
105
- - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and not +rms_norm and +quant_fp8 and qwen3"
104
+ # Qwen/Deepseek requires +quant_fp8 as -quant_fp8 rms+quant fusion is not supported
105
+ - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and not +rms_norm and +quant_fp8 and (qwen3 or deepseek)"
106
107
- label: Fusion E2E Config Sweep (H100)
108
timeout_in_minutes: 30
@@ -132,9 +132,9 @@ steps:
132
commands:
133
- nvidia-smi
134
# Run all models but only FLASHINFER, Inductor partition and native custom ops
135
- # Qwen requires +quant_fp8 as -quant_fp8 rms+quant fusion is not supported
135
+ # Qwen/Deepseek requires +quant_fp8 as -quant_fp8 rms+quant fusion is not supported
136
# Run just llama3 (fp8 & fp4) for all config combinations (only inductor partition)
137
- - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and (FLASHINFER and not +rms_norm and (not +quant_fp8 or +quant_fp8 and qwen3) or llama-3)"
137
+ - pytest -v -s tests/compile/fusions_e2e/test_tp1_quant.py -k "inductor_partition and (FLASHINFER and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek)) or llama-3)"
138
139
- label: Fusion E2E TP2 Quick (H100)
140
timeout_in_minutes: 20
@@ -150,8 +150,8 @@ steps:
150
commands:
151
- nvidia-smi
152
# Run all models and attn backends but only Inductor partition and native custom ops
153
- - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "inductor_partition and not +rms_norm and not +quant_fp8"
154
- - pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "inductor_partition and not +rms_norm and not +quant_fp8"
153
+ - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))"
154
+ - pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))"
155
156
- label: Fusion E2E TP2 AR-RMS Config Sweep (H100)
157
timeout_in_minutes: 40
@@ -205,7 +205,7 @@ steps:
205
commands:
206
- nvidia-smi
207
# Run all models but only FLASHINFER, Inductor partition and native custom ops
208
- # include qwen with +quant_fp8 as -quant_fp8 rms+quant fusion is not supported
208
+ # include qwen/deepseek with +quant_fp8 as -quant_fp8 rms+quant fusion is not supported
209
# for ar-rms-quant-fp4, also sweep llama3
210
- - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "(FLASHINFER and inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and qwen3)) or Llama-3.1-8B-Instruct-FP4"
211
- - pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "FLASHINFER and inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and qwen3)"
210
+ - pytest -v -s tests/compile/fusions_e2e/test_tp2_ar_rms.py -k "(FLASHINFER and inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))) or Llama-3.1-8B-Instruct-FP4"
211
+ - pytest -v -s tests/compile/fusions_e2e/test_tp2_async_tp.py -k "FLASHINFER and inductor_partition and not +rms_norm and (not +quant_fp8 or +quant_fp8 and (qwen3 or deepseek))"
csrc/quantization/fused_kernels/fused_layernorm_dynamic_per_token_quant.cu
+45
-24
@@ -15,31 +15,33 @@ __device__ void rms_norm_dynamic_per_token_quant_vec(
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, int32_t const hidden_size,
18
- scalar_t* __restrict__ residual = nullptr) {
18
+ int32_t const input_stride, scalar_t* __restrict__ residual = nullptr) {
19
float rms = 0.0f;
20
float token_scale = 0.0f;
21
22
// Compute rms
23
vllm::vectorized::compute_rms<scalar_t, has_residual>(
24
- &rms, input, hidden_size, var_epsilon, residual);
24
+ &rms, input, hidden_size, input_stride, var_epsilon, residual);
25
26
// Compute scale
27
vllm::vectorized::compute_dynamic_per_token_scales<scalar_t, scalar_out_t,
28
has_residual>(
29
&token_scale, scales, input, weight, rms, scale_ub, hidden_size,
30
- residual);
30
+ input_stride, residual);
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>(
37
- out, input, weight, rms, &token_scale, hidden_size, residual);
36
+ has_residual>(out, input, weight, rms,
37
+ &token_scale, hidden_size,
38
+ input_stride, residual);
39
} else {
40
// FP8 - Do not invert token_scale for exact match with FBGemm
41
vllm::vectorized::norm_and_quant<scalar_t, scalar_out_t, false,
41
- has_residual>(
42
- out, input, weight, rms, &token_scale, hidden_size, residual);
42
+ has_residual>(out, input, weight, rms,
43
+ &token_scale, hidden_size,
44
+ input_stride, residual);
45
}
46
}
47
@@ -51,38 +53,40 @@ __global__ void rms_norm_dynamic_per_token_quant_kernel(
53
scalar_t const* __restrict__ input, // [..., hidden_size]
54
scalar_t const* __restrict__ weight, // [hidden_size]
55
float const* scale_ub, float const var_epsilon, int32_t const hidden_size,
54
- scalar_t* __restrict__ residual = nullptr) {
56
+ int32_t const input_stride, scalar_t* __restrict__ residual = nullptr) {
57
// For vectorization, token_input and token_output pointers need to be
58
// aligned at 8-byte and 4-byte addresses respectively.
57
- bool const can_vectorize = hidden_size % 4 == 0;
59
+ bool const can_vectorize = hidden_size % 4 == 0 and input_stride % 4 == 0;
60
61
if (can_vectorize) {
62
return rms_norm_dynamic_per_token_quant_vec<scalar_t, scalar_out_t,
63
has_residual>(
64
out, scales, input, weight, scale_ub, var_epsilon, hidden_size,
63
- residual);
65
+ input_stride, residual);
66
}
67
68
float rms = 0.0f;
69
float token_scale = 0.0f;
70
71
// Compute RMS
70
- vllm::compute_rms<scalar_t, has_residual>(&rms, input, hidden_size,
71
- var_epsilon, residual);
72
+ vllm::compute_rms<scalar_t, has_residual>(
73
+ &rms, input, hidden_size, input_stride, var_epsilon, residual);
74
// Compute Scale
75
vllm::compute_dynamic_per_token_scales<scalar_t, scalar_out_t, has_residual>(
76
&token_scale, scales, input, weight, rms, scale_ub, hidden_size,
75
- residual);
77
+ input_stride, residual);
78
79
// RMS Norm + Quant
80
if constexpr (std::is_same_v<scalar_out_t, int8_t>) {
81
token_scale = 1.0f / token_scale;
82
vllm::norm_and_quant<scalar_t, scalar_out_t, true, has_residual>(
81
- out, input, weight, rms, &token_scale, hidden_size, residual);
83
+ out, input, weight, rms, &token_scale, hidden_size, input_stride,
84
+ residual);
85
} else {
86
// FP8 - Do not invert s_token_scale for exact match with FBGemm
87
vllm::norm_and_quant<scalar_t, scalar_out_t, false, has_residual>(
85
- out, input, weight, rms, &token_scale, hidden_size, residual);
88
+ out, input, weight, rms, &token_scale, hidden_size, input_stride,
89
+ residual);
90
}
91
}
92
@@ -97,19 +101,20 @@ __global__ void rms_norm_per_block_quant_kernel(
101
scalar_t const* __restrict__ input, // [..., hidden_size]
102
scalar_t const* __restrict__ weight, // [hidden_size]
103
float const* scale_ub, float const var_epsilon, int32_t const hidden_size,
100
- scalar_t* __restrict__ residual = nullptr, int64_t outer_scale_stride = 1) {
104
+ int32_t const input_stride, scalar_t* __restrict__ residual = nullptr,
105
+ int64_t outer_scale_stride = 1) {
106
float rms;
107
// Compute RMS
108
// Always able to vectorize due to constraints on hidden_size
109
vllm::vectorized::compute_rms<scalar_t, has_residual>(
105
- &rms, input, hidden_size, var_epsilon, residual);
110
+ &rms, input, hidden_size, input_stride, var_epsilon, residual);
111
112
// Compute Scale
113
// Always able to vectorize due to constraints on hidden_size and group_size
114
vllm::vectorized::compute_dynamic_per_token_scales<
115
scalar_t, scalar_out_t, has_residual, is_scale_transposed, group_size>(
111
- nullptr, scales, input, weight, rms, scale_ub, hidden_size, residual,
112
- outer_scale_stride);
116
+ nullptr, scales, input, weight, rms, scale_ub, hidden_size, input_stride,
117
+ residual, outer_scale_stride);
118
119
// RMS Norm + Quant
120
// Always able to vectorize due to constraints on hidden_size
@@ -120,7 +125,7 @@ __global__ void rms_norm_per_block_quant_kernel(
125
vllm::vectorized::norm_and_quant<
126
scalar_t, scalar_out_t, std::is_same_v<scalar_out_t, int8_t>,
127
has_residual, is_scale_transposed, group_size>(
123
- out, input, weight, rms, scales, hidden_size, residual,
128
+ out, input, weight, rms, scales, hidden_size, input_stride, residual,
129
outer_scale_stride);
130
}
131
@@ -137,6 +142,7 @@ void rms_norm_dynamic_per_token_quant_dispatch(
142
std::optional<at::Tensor> const& scale_ub,
143
std::optional<at::Tensor>& residual) {
144
int32_t hidden_size = input.size(-1);
145
+ int32_t input_stride = input.view({-1, hidden_size}).stride(0);
146
auto num_tokens = input.numel() / hidden_size;
147
148
dim3 grid(num_tokens);
@@ -153,7 +159,7 @@ void rms_norm_dynamic_per_token_quant_dispatch(
159
out.data_ptr<scalar_t>(), scales.data_ptr<float>(),
160
input.data_ptr<scalar_in_t>(), weight.data_ptr<scalar_in_t>(),
161
scale_ub.has_value() ? scale_ub->data_ptr<float>() : nullptr,
156
- var_epsilon, hidden_size,
162
+ var_epsilon, hidden_size, input_stride,
163
has_residual ? residual->data_ptr<scalar_in_t>() : nullptr);
164
});
165
});
@@ -170,7 +176,9 @@ void rms_norm_dynamic_per_token_quant(
176
? c10::ScalarType::Float8_e4m3fn
177
: c10::ScalarType::Float8_e4m3fnuz;
178
TORCH_CHECK(out.dtype() == kFp8Type || out.dtype() == torch::kInt8);
173
- TORCH_CHECK(out.is_contiguous() && input.is_contiguous());
179
+ TORCH_CHECK(out.is_contiguous());
180
+ TORCH_CHECK(input.stride(-1) == 1,
181
+ "Input must be contiguous in the last dimension");
182
183
if (scale_ub.has_value()) {
184
TORCH_CHECK(out.dtype() == kFp8Type);
@@ -179,6 +187,7 @@ void rms_norm_dynamic_per_token_quant(
187
TORCH_CHECK(scales.dtype() == torch::kFloat32);
188
if (residual) {
189
TORCH_CHECK(residual->scalar_type() == input.scalar_type());
190
+ TORCH_CHECK(residual->is_contiguous());
191
}
192
193
VLLM_DISPATCH_FLOATING_TYPES(
@@ -200,6 +209,15 @@ void rms_norm_per_block_quant_dispatch(
209
std::optional<at::Tensor> const& scale_ub,
210
std::optional<at::Tensor>& residual, bool is_scale_transposed) {
211
int32_t hidden_size = input.size(-1);
212
+ int32_t input_stride = input.view({-1, hidden_size}).stride(0);
213
+
214
+ TORCH_CHECK(hidden_size % 4 == 0,
215
+ "Hidden size must be divisible by 4 for vectorized access");
216
+ TORCH_CHECK(input_stride % 4 == 0,
217
+ "Input stride must be divisible by 4 for vectorized access");
218
+ TORCH_CHECK(group_size % 4 == 0,
219
+ "Group size must be divisible by 4 for vectorized access");
220
+
221
auto num_tokens = input.numel() / hidden_size;
222
223
dim3 grid(num_tokens);
@@ -225,7 +243,7 @@ void rms_norm_per_block_quant_dispatch(
243
weight.data_ptr<scalar_in_t>(),
244
scale_ub.has_value() ? scale_ub->data_ptr<float>()
245
: nullptr,
228
- var_epsilon, hidden_size,
246
+ var_epsilon, hidden_size, input_stride,
247
has_residual ? residual->data_ptr<scalar_in_t>()
248
: nullptr,
249
scales.stride(1));
@@ -246,7 +264,9 @@ void rms_norm_per_block_quant(torch::Tensor& out, torch::Tensor const& input,
264
? c10::ScalarType::Float8_e4m3fn
265
: c10::ScalarType::Float8_e4m3fnuz;
266
TORCH_CHECK(out.dtype() == kFp8Type || out.dtype() == torch::kInt8);
249
- TORCH_CHECK(out.is_contiguous() && input.is_contiguous());
267
+ TORCH_CHECK(out.is_contiguous());
268
+ TORCH_CHECK(input.stride(-1) == 1,
269
+ "Input must be contiguous in the last dimension");
270
271
if (scale_ub.has_value()) {
272
TORCH_CHECK(out.dtype() == kFp8Type);
@@ -255,6 +275,7 @@ void rms_norm_per_block_quant(torch::Tensor& out, torch::Tensor const& input,
275
TORCH_CHECK(scales.dtype() == torch::kFloat32);
276
if (residual) {
277
TORCH_CHECK(residual->scalar_type() == input.scalar_type());
278
+ TORCH_CHECK(residual->is_contiguous());
279
}
280
281
TORCH_CHECK(group_size == 128 || group_size == 64,
csrc/quantization/fused_kernels/layernorm_utils.cuh
+44
-27
@@ -16,14 +16,17 @@ namespace vllm {
16
// has_residual must be true, if residual is not a nullptr
17
template <typename scalar_t, bool has_residual = false>
18
__device__ void compute_rms(float* rms, scalar_t const* __restrict__ input,
19
- int32_t const hidden_size, float const epsilon,
19
+ int32_t const hidden_size,
20
+ int32_t const input_stride, float const epsilon,
21
scalar_t const* __restrict__ residual = nullptr) {
22
+ int64_t const input_token_offset =
23
+ blockIdx.x * static_cast<int64_t>(input_stride);
24
int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
25
// sum of squares
26
float ss = 0.0f;
27
28
for (auto i = threadIdx.x; i < hidden_size; i += blockDim.x) {
26
- float x = static_cast<float>(input[token_offset + i]);
29
+ float x = static_cast<float>(input[input_token_offset + i]);
30
if constexpr (has_residual) {
31
x += static_cast<float>(residual[token_offset + i]);
32
}
@@ -73,15 +76,20 @@ __device__ void compute_dynamic_per_token_scales(
76
float* __restrict__ token_scale, float* __restrict__ all_token_scales,
77
scalar_t const* __restrict__ input, scalar_t const* __restrict__ weight,
78
float const rms, float const* __restrict__ scale_ub,
76
- int32_t const hidden_size, scalar_t const* __restrict__ residual = nullptr,
79
+ int32_t const hidden_size, int32_t const input_stride,
80
+ scalar_t const* __restrict__ residual = nullptr,
81
int32_t const group_size = 0, int64_t outer_scale_stride = 1) {
82
float block_absmax_val_maybe = 0.0f;
83
constexpr scalar_out_t qmax{quant_type_max_v<scalar_out_t>};
84
__syncthreads();
85
+
86
+ int64_t const input_token_offset =
87
+ blockIdx.x * static_cast<int64_t>(input_stride);
88
+ int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
89
+
90
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);
91
int64_t num_groups = hidden_size / group_size;
92
+ __shared__ float s_max_vals[1024];
93
int64_t const threads_per_group = blockDim.x / num_groups;
94
int64_t const thread_in_group = threadIdx.x % threads_per_group;
95
int64_t const group_offset = threadIdx.x / threads_per_group * group_size;
@@ -89,7 +97,7 @@ __device__ void compute_dynamic_per_token_scales(
97
int64_t const thread_end =
98
min(group_offset + group_size, static_cast<int64_t>(hidden_size));
99
for (auto i = thread_offset; i < thread_end; i += threads_per_group) {
92
- float x = static_cast<float>(input[token_offset + i]);
100
+ float x = static_cast<float>(input[input_token_offset + i]);
101
if constexpr (has_residual) {
102
x += static_cast<float>(residual[token_offset + i]);
103
}
@@ -144,10 +152,8 @@ __device__ void compute_dynamic_per_token_scales(
152
}
153
__syncthreads();
154
} else {
147
- int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
148
-
155
for (auto i = threadIdx.x; i < hidden_size; i += blockDim.x) {
150
- float x = static_cast<float>(input[token_offset + i]);
156
+ float x = static_cast<float>(input[input_token_offset + i]);
157
if constexpr (has_residual) {
158
x += static_cast<float>(residual[token_offset + i]);
159
}
@@ -185,12 +191,15 @@ template <typename scalar_t, typename scalar_out_t, bool is_scale_inverted,
191
__device__ void norm_and_quant(
192
scalar_out_t* __restrict__ output, scalar_t const* __restrict__ input,
193
scalar_t const* __restrict__ weight, float const rms, float* const scale,
188
- int32_t const hidden_size, scalar_t* __restrict__ residual = nullptr,
189
- int32_t const group_size = 0, int64_t outer_scale_stride = 1) {
194
+ int32_t const hidden_size, int32_t const input_stride,
195
+ scalar_t* __restrict__ residual = nullptr, int32_t const group_size = 0,
196
+ int64_t outer_scale_stride = 1) {
197
+ int64_t const input_token_offset =
198
+ blockIdx.x * static_cast<int64_t>(input_stride);
199
int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
200
201
for (auto i = threadIdx.x; i < hidden_size; i += blockDim.x) {
193
- float x = static_cast<float>(input[token_offset + i]);
202
+ float x = static_cast<float>(input[input_token_offset + i]);
203
if constexpr (has_residual) {
204
x += static_cast<float>(residual[token_offset + i]);
205
residual[token_offset + i] = static_cast<scalar_t>(x);
@@ -224,13 +233,16 @@ namespace vectorized {
233
// hidden_size must be a multiple of 4
234
template <typename scalar_t, bool has_residual = false>
235
__device__ void compute_rms(float* rms, scalar_t const* __restrict__ input,
227
- int32_t const hidden_size, float const epsilon,
236
+ int32_t const hidden_size,
237
+ int32_t const input_stride, float const epsilon,
238
scalar_t const* __restrict__ residual = nullptr) {
239
+ int64_t const input_token_offset =
240
+ blockIdx.x * static_cast<int64_t>(input_stride);
241
int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
242
243
// Vectorized input/output to better utilize memory bandwidth.
244
vec4_t<scalar_t> const* vec_input =
233
- reinterpret_cast<vec4_t<scalar_t> const*>(&input[token_offset]);
245
+ reinterpret_cast<vec4_t<scalar_t> const*>(&input[input_token_offset]);
246
vec4_t<scalar_t> const* vec_residual = nullptr;
247
if constexpr (has_residual) {
248
vec_residual =
@@ -288,7 +300,8 @@ __device__ void compute_dynamic_per_token_scales(
300
float* __restrict__ token_scale, float* __restrict__ all_token_scales,
301
scalar_t const* __restrict__ input, scalar_t const* __restrict__ weight,
302
float const rms, float const* __restrict__ scale_ub,
291
- int32_t const hidden_size, scalar_t const* __restrict__ residual = nullptr,
303
+ int32_t const hidden_size, int32_t const input_stride,
304
+ scalar_t const* __restrict__ residual = nullptr,
305
int64_t outer_scale_stride = 1) {
306
constexpr scalar_out_t qmax{quant_type_max_v<scalar_out_t>};
307
@@ -300,10 +313,13 @@ __device__ void compute_dynamic_per_token_scales(
313
vec4_t<scalar_t> const* vec_weight = nullptr;
314
vec4_t<scalar_t> const* vec_residual = nullptr;
315
316
+ int64_t const input_token_offset =
317
+ blockIdx.x * static_cast<int64_t>(input_stride);
318
+ int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
319
+
320
if constexpr (group_size > 0) {
321
__shared__ float s_max_vals[1024];
322
306
- int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
323
int64_t const num_groups = hidden_size / group_size;
324
int64_t const threads_per_group = blockDim.x / num_groups;
325
int64_t const thread_in_group = threadIdx.x % threads_per_group;
@@ -312,7 +328,8 @@ __device__ void compute_dynamic_per_token_scales(
328
int64_t const thread_offset = group_offset + thread_in_group;
329
int64_t const thread_end = min(group_offset + (group_size >> 2),
330
static_cast<int64_t>(hidden_size >> 2));
315
- vec_input = reinterpret_cast<vec4_t<scalar_t> const*>(&input[token_offset]);
331
+ vec_input =
332
+ reinterpret_cast<vec4_t<scalar_t> const*>(&input[input_token_offset]);
333
vec_weight = reinterpret_cast<vec4_t<scalar_t> const*>(weight);
334
if constexpr (has_residual) {
335
vec_residual =
@@ -396,8 +413,8 @@ __device__ void compute_dynamic_per_token_scales(
413
__syncthreads();
414
415
} else {
399
- int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
400
- vec_input = reinterpret_cast<vec4_t<scalar_t> const*>(&input[token_offset]);
416
+ vec_input =
417
+ reinterpret_cast<vec4_t<scalar_t> const*>(&input[input_token_offset]);
418
vec_weight = reinterpret_cast<vec4_t<scalar_t> const*>(weight);
419
if constexpr (has_residual) {
420
vec_residual =
@@ -462,18 +479,18 @@ __device__ void compute_dynamic_per_token_scales(
479
template <typename scalar_t, typename scalar_out_t, bool is_scale_inverted,
480
bool has_residual = false, bool is_scale_transposed = false,
481
int32_t group_size = 0>
465
-__device__ void norm_and_quant(scalar_out_t* __restrict__ output,
466
- scalar_t const* __restrict__ input,
467
- scalar_t const* __restrict__ weight,
468
- float const rms, float* const scale,
469
- int32_t const hidden_size,
470
- scalar_t* __restrict__ residual = nullptr,
471
- int64_t outer_scale_stride = 1) {
482
+__device__ void norm_and_quant(
483
+ scalar_out_t* __restrict__ output, scalar_t const* __restrict__ input,
484
+ scalar_t const* __restrict__ weight, float const rms, float* const scale,
485
+ int32_t const hidden_size, int32_t const input_stride,
486
+ scalar_t* __restrict__ residual = nullptr, int64_t outer_scale_stride = 1) {
487
+ int64_t const input_token_offset =
488
+ blockIdx.x * static_cast<int64_t>(input_stride);
489
int64_t const token_offset = blockIdx.x * static_cast<int64_t>(hidden_size);
490
491
// Vectorized input/output/weight/residual to better utilize memory bandwidth.
492
vec4_t<scalar_t> const* vec_input =
476
- reinterpret_cast<vec4_t<scalar_t> const*>(&input[token_offset]);
493
+ reinterpret_cast<vec4_t<scalar_t> const*>(&input[input_token_offset]);
494
vec4_t<scalar_t> const* vec_weight =
495
reinterpret_cast<vec4_t<scalar_t> const*>(weight);
496
q8x4_t<scalar_out_t>* vec_output =
tests/compile/fusions_e2e/conftest.py
+10
@@ -72,6 +72,16 @@ def run_e2e_fusion_test(monkeypatch, caplog_mp_spawn):
72
73
rocm_aiter_ops.refresh_env_variables()
74
75
+ # Filter here to reduce code duplication
76
+ requires_mla = "deepseek" in model_name.lower()
77
+ is_mla = "mla" in attn_backend.backend.name.lower()
78
+
79
+ if requires_mla != is_mla:
80
+ pytest.skip(
81
+ f"Incompatible model '{model_name}' and "
82
+ f"attention backend '{attn_backend.backend.name}'"
83
+ )
84
+
85
# Disable, compile cache to make sure custom passes run.
86
# Otherwise, we can't verify fusion happened through the logs.
87
monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
tests/compile/fusions_e2e/models.py
+36
@@ -44,6 +44,20 @@ ROCM_AITER_UNIFIED_ATTN = pytest.param(
44
),
45
)
46
47
+FLASHINFER_MLA_ATTN = pytest.param(
48
+ AttentionBackendCase(backend=AttentionBackendEnum.FLASHINFER_MLA),
49
+ id="FLASHINFER_MLA",
50
+ marks=pytest.mark.skipif(
51
+ not is_blackwell() or not has_flashinfer(),
52
+ reason="FI backend requires Blackwell and FlashInfer",
53
+ ),
54
+)
55
+
56
+TRITON_MLA_ATTN = pytest.param(
57
+ AttentionBackendCase(backend=AttentionBackendEnum.TRITON_MLA),
58
+ id="TRITON_MLA",
59
+)
60
+
61
# Models
62
llama3_8b = ModelFusionInfo(
63
model_name="meta-llama/Llama-3.1-8B-Instruct",
@@ -126,3 +140,25 @@ qwen3_a3b_fp8 = ModelFusionInfo(
140
async_tp=n_layers * 2,
141
),
142
)
143
+
144
+deepseek_v3_fp8 = ModelFusionInfo(
145
+ model_name="deepseek-ai/DeepSeek-V3",
146
+ matches=lambda n_layers: Matches(
147
+ # 3 per dense layer (first 3):
148
+ # - input_rms + qkv_proj
149
+ # - q_a_layernorm + q_b_proj (inside MLA wrapper)
150
+ # - post_attn_layernorm + MLP
151
+ # 2 per MoE layer (remaining) due to MoE wrapping
152
+ rms_quant_fusion=n_layers * 2 + min(3, n_layers), # add for 3 dense layers
153
+ # TODO silu+block quant
154
+ # act_quant_fusion=min(3, n_layers), # dense layers only
155
+ act_quant_fusion=0,
156
+ # MLA attn + quant not supported yet:
157
+ # https://github.com/vllm-project/vllm/issues/35792
158
+ attn_quant_fusion=0,
159
+ ar_rms_fusion=n_layers * 2 + 1,
160
+ # TODO
161
+ # sequence_parallel= n_layers * 2 + 1,
162
+ # async_tp=n_layers * 2,
163
+ ),
164
+)
tests/compile/fusions_e2e/test_tp1_quant.py
+13
-8
@@ -17,9 +17,12 @@ from .common import (
17
)
18
from .models import (
19
FLASHINFER_ATTN,
20
+ FLASHINFER_MLA_ATTN,
21
ROCM_AITER_UNIFIED_ATTN,
22
ROCM_ATTN,
23
TRITON_ATTN,
24
+ TRITON_MLA_ATTN,
25
+ deepseek_v3_fp8,
26
llama3_8b_fp4,
27
llama3_8b_fp8,
28
llama4_scout_fp4,
@@ -33,6 +36,9 @@ from .models import (
36
[
37
(*llama3_8b_fp8, False),
38
(*qwen3_a3b_fp8, False),
39
+ (*qwen3_a3b_fp8, True),
40
+ (*deepseek_v3_fp8, False),
41
+ (*deepseek_v3_fp8, True),
42
pytest.param(
43
*llama4_scout_fp8,
44
False,
@@ -41,13 +47,6 @@ from .models import (
47
reason="Llama4 Scout FP8 only supported on CUDA",
48
),
49
),
44
- pytest.param(
45
- *qwen3_a3b_fp8,
46
- True,
47
- marks=pytest.mark.skipif(
48
- not current_platform.is_cuda(), reason="DeepGemm only supported on CUDA"
49
- ),
50
- ),
50
],
51
)
52
@pytest.mark.parametrize(
@@ -57,6 +56,8 @@ from .models import (
56
FLASHINFER_ATTN,
57
ROCM_ATTN,
58
ROCM_AITER_UNIFIED_ATTN,
59
+ FLASHINFER_MLA_ATTN,
60
+ TRITON_MLA_ATTN,
61
],
62
)
63
@pytest.mark.parametrize("n_layers", [6])
@@ -75,6 +76,9 @@ def test_tp1_fp8_fusions(
76
run_e2e_fusion_test,
77
monkeypatch,
78
):
79
+ if use_deepgemm and not current_platform.is_cuda():
80
+ pytest.skip("DeepGemm only supported on CUDA")
81
+
82
if use_deepgemm and is_flashinfer_fp8_blockscale_gemm_supported():
83
# Flashinfer block FP8 GEMM has internal quantization, so it can't
84
# be fused with other ops.
@@ -86,7 +90,8 @@ def test_tp1_fp8_fusions(
90
91
matches = matches_fn(n_layers)
92
89
- if "qwen" in model_name.lower() and "-quant_fp8" in custom_ops:
93
+ block_fp8 = "qwen" in model_name.lower() or "deepseek" in model_name.lower()
94
+ if block_fp8 and "-quant_fp8" in custom_ops:
95
# This is why config forces +quant_fp8 by default
96
pytest.skip("native QuantFP8 matching not supported for group quant")
97
tests/compile/fusions_e2e/test_tp2_ar_rms.py
+9
-4
@@ -17,7 +17,9 @@ from .common import (
17
)
18
from .models import (
19
FLASHINFER_ATTN,
20
+ FLASHINFER_MLA_ATTN,
21
TRITON_ATTN,
22
+ deepseek_v3_fp8,
23
llama3_8b,
24
llama3_8b_fp4,
25
llama3_8b_fp8,
@@ -33,10 +35,12 @@ pytestmark = pytest.mark.skipif(not current_platform.is_cuda(), reason="Only tes
35
@multi_gpu_test(num_gpus=2)
36
@pytest.mark.parametrize(
37
"model_name, matches_fn, model_kwargs, hf_overrides",
36
- # qwen3-fp8 should still fuse AR+rms even though group quant is not yet supported
37
- [llama3_8b_fp8, llama4_scout_fp8, qwen3_a3b_fp8],
38
+ # qwen3 & dsv3 should still fuse AR+rms even though group quant is not yet supported
39
+ [llama3_8b_fp8, llama4_scout_fp8, qwen3_a3b_fp8, deepseek_v3_fp8],
40
+)
41
+@pytest.mark.parametrize(
42
+ "attn_backend", [TRITON_ATTN, FLASHINFER_ATTN, FLASHINFER_MLA_ATTN]
43
)
39
-@pytest.mark.parametrize("attn_backend", [TRITON_ATTN, FLASHINFER_ATTN])
44
@pytest.mark.parametrize("n_layers", [4])
45
@pytest.mark.parametrize("custom_ops", custom_ops_combos("quant_fp8", "rms_norm"))
46
@pytest.mark.parametrize("inductor_graph_partition", INDUCTOR_GRAPH_PARTITION)
@@ -54,7 +58,8 @@ def test_tp2_ar_rms_fp8_fusions(
58
):
59
matches = matches_fn(n_layers)
60
57
- if "qwen" in model_name.lower() and "-quant_fp8" in custom_ops:
61
+ block_fp8 = "qwen" in model_name.lower() or "deepseek" in model_name.lower()
62
+ if block_fp8 and "-quant_fp8" in custom_ops:
63
# This is why config forces +quant_fp8 by default
64
pytest.skip("native QuantFP8 matching not supported for group quant")
65
tests/kernels/core/test_fused_quant_layernorm.py
+51
-13
@@ -162,6 +162,7 @@ def ops_impl(
162
)
163
@pytest.mark.parametrize("seed", SEEDS)
164
@pytest.mark.parametrize("device", CUDA_DEVICES)
165
+@pytest.mark.parametrize("strided_input", [False, True])
166
@torch.inference_mode()
167
def test_rms_norm(
168
default_vllm_config,
@@ -175,6 +176,7 @@ def test_rms_norm(
176
tma_alignment: int,
177
seed: int,
178
device: str,
179
+ strided_input: bool,
180
) -> None:
181
torch.random.manual_seed(seed)
182
if torch.cuda.is_available():
@@ -184,17 +186,17 @@ def test_rms_norm(
186
187
if group_size is not None and hidden_size % group_size[1] != 0:
188
# skip
187
- return
189
+ pytest.skip("Skip non-divisible group sizes")
190
191
if group_size is not None and has_scale_ub:
192
# blockwise baseline doesn't support scale_ub
191
- return
193
+ pytest.skip("scale_ub not supported for blockwise/group quantization")
194
195
if (
196
group_size is None or quant_dtype != current_platform.fp8_dtype()
197
) and tma_alignment != 0:
198
# TMA alignment is only supported for groupwise fp8 kernels
197
- return
199
+ pytest.skip("tma alignment not supported for per-token or int8 quantization")
200
201
if (
202
group_size is not None
@@ -202,21 +204,36 @@ def test_rms_norm(
204
and hidden_size // group_size[1] % tma_alignment == 0
205
):
206
# Skip tests where TMA alignment doesn't create extra padding to save time
205
- return
207
+ pytest.skip("Skip TMA alignment cases where no extra padding is added")
208
209
if has_scale_ub and quant_dtype != current_platform.fp8_dtype():
210
# skip
209
- return
211
+ pytest.skip("scale_ub only supported for fp8 quantization")
212
213
layer = RMSNorm(hidden_size, EPS).to(dtype=dtype)
214
215
# Make weights
216
layer.weight.data.normal_(mean=1.0, std=0.1)
217
216
- # Make inputs
218
+ # Make inputs: use a wider tensor and slice to create a non-contiguous
219
+ # (strided) input when strided_input=True. The last dimension stride
220
+ # remains 1, which the kernel requires.
221
scale = 1 / (hidden_size)
218
- x = torch.randn(num_tokens, hidden_size, dtype=dtype) * scale
219
- residual = torch.randn_like(x) * scale if add_residual else None
222
+ last_dim = 2 * hidden_size if strided_input else hidden_size
223
+ x = torch.randn(num_tokens, last_dim, dtype=dtype) * scale
224
+ x = x[:, :hidden_size]
225
+
226
+ # dim 1 gets special-cased
227
+ x_is_strided = strided_input and num_tokens != 1
228
+ # check that the input is strided iff we expect it to be
229
+ assert x.is_contiguous() != x_is_strided
230
+
231
+ # Residual must still be contiguous
232
+ residual = (
233
+ torch.randn(num_tokens, hidden_size, dtype=dtype) * scale
234
+ if add_residual
235
+ else None
236
+ )
237
if has_scale_ub:
238
rms_x, _ = ref_rms_norm(layer, x, residual)
239
scale_ub = torch.mean(rms_x).to(dtype=torch.float32, device="cuda")
@@ -260,12 +277,33 @@ def test_rms_norm(
277
if add_residual:
278
assert torch.allclose(ref_residual, ops_residual)
279
263
- output = torch.empty_like(x, dtype=quant_dtype)
280
+ output = torch.empty(x.shape, dtype=quant_dtype, device=x.device)
281
scales = torch.empty(
282
(x.numel() // x.shape[-1], 1), device=x.device, dtype=torch.float32
283
)
284
268
- opcheck(
269
- torch.ops._C.rms_norm_dynamic_per_token_quant,
270
- (output, x, layer.weight, scales, 1e-5, scale_ub, residual),
271
- )
285
+ if group_size is None:
286
+ opcheck(
287
+ torch.ops._C.rms_norm_dynamic_per_token_quant,
288
+ (output, x, layer.weight, scales, 1e-5, scale_ub, residual),
289
+ )
290
+ else:
291
+ # TODO(luka/eliza) opcheck is broken?
292
+ # Somehow the cloned args are getting mutated in-place,
293
+ # which causes the opcheck to fail.
294
+ # https://github.com/vllm-project/vllm/issues/36688
295
+ return
296
+ opcheck(
297
+ torch.ops._C.rms_norm_per_block_quant,
298
+ (
299
+ output,
300
+ x,
301
+ layer.weight,
302
+ scales,
303
+ 1e-5,
304
+ scale_ub,
305
+ residual,
306
+ group_size[1],
307
+ True, # is_scale_transposed
308
+ ),
309
+ )
vllm/_custom_ops.py
+2
-2
@@ -427,7 +427,7 @@ def rms_norm_dynamic_per_token_quant(
427
scale_ub: torch.Tensor | None = None,
428
residual: torch.Tensor | None = None,
429
) -> tuple[torch.Tensor, torch.Tensor]:
430
- output = torch.empty_like(input, dtype=quant_dtype)
430
+ output = torch.empty(input.shape, dtype=quant_dtype, device=input.device)
431
scales = torch.empty(
432
(input.numel() // input.shape[-1], 1), device=input.device, dtype=torch.float32
433
)
@@ -451,7 +451,7 @@ def rms_norm_per_block_quant(
451
tma_alignment: int = 0,
452
) -> tuple[torch.Tensor, torch.Tensor]:
453
assert len(group_size) == 2
454
- output = torch.empty_like(input, dtype=quant_dtype)
454
+ output = torch.empty(input.shape, dtype=quant_dtype, device=input.device)
455
if is_scale_transposed:
456
if tma_alignment == 0:
457
scales = torch.empty(