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 14, 05:30 UTC
[Model] New model support for Motif-1-Tiny (#23414)
Signed-off-by: ca1207 <ca1207zzz@gmail.com> Signed-off-by: TaehyunKim <73943231+ca1207@users.noreply.github.com> Co-authored-by: WyldeCat <skan1543@gmail.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com>
TaehyunKim committed
Sep 11, 2025 at 15:29 UTC
9bd831f5012b0ad1b4e3b7c1af511b19558d7626
13 files changed
+871
-4
benchmarks/kernels/benchmark_polynorm.py
new
+155
@@ -0,0 +1,155 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+import itertools
5
+
6
+import torch
7
+
8
+from vllm import _custom_ops as vllm_ops
9
+from vllm.triton_utils import triton
10
+
11
+
12
+def polynorm_naive(
13
+ x: torch.Tensor,
14
+ weight: torch.Tensor,
15
+ bias: torch.Tensor,
16
+ eps: float = 1e-6,
17
+):
18
+ orig_shape = x.shape
19
+ x = x.view(-1, x.shape[-1])
20
+
21
+ def norm(x, eps: float):
22
+ return x / torch.sqrt(x.pow(2).mean(-1, keepdim=True) + eps)
23
+
24
+ x = x.float()
25
+ return (
26
+ (
27
+ weight[0] * norm(x**3, eps)
28
+ + weight[1] * norm(x**2, eps)
29
+ + weight[2] * norm(x, eps)
30
+ + bias
31
+ )
32
+ .to(weight.dtype)
33
+ .view(orig_shape)
34
+ )
35
+
36
+
37
+def polynorm_vllm(
38
+ x: torch.Tensor,
39
+ weight: torch.Tensor,
40
+ bias: torch.Tensor,
41
+ eps: float = 1e-6,
42
+):
43
+ orig_shape = x.shape
44
+ x = x.view(-1, x.shape[-1])
45
+
46
+ out = torch.empty_like(x)
47
+ vllm_ops.poly_norm(out, x, weight, bias, eps)
48
+ output = out
49
+
50
+ output = output.view(orig_shape)
51
+ return output
52
+
53
+
54
+def calculate_diff(batch_size, seq_len, hidden_dim):
55
+ dtype = torch.bfloat16
56
+ x = torch.randn(batch_size, seq_len, hidden_dim, dtype=dtype, device="cuda")
57
+ weight = torch.ones(3, dtype=dtype, device="cuda")
58
+ bias = torch.ones(1, dtype=dtype, device="cuda")
59
+
60
+ output_naive = polynorm_naive(x, weight, bias)
61
+ output_vllm = polynorm_vllm(x, weight, bias)
62
+
63
+ if torch.allclose(output_naive, output_vllm, atol=1e-2, rtol=1e-2):
64
+ print("✅ All implementations match")
65
+ else:
66
+ print("❌ Implementations differ")
67
+
68
+
69
+batch_size_range = [2**i for i in range(0, 7, 2)]
70
+seq_length_range = [2**i for i in range(6, 11, 1)]
71
+dim_range = [2048, 4096]
72
+configs = list(itertools.product(dim_range, batch_size_range, seq_length_range))
73
+
74
+
75
+def get_benchmark():
76
+ @triton.testing.perf_report(
77
+ triton.testing.Benchmark(
78
+ x_names=["dim", "batch_size", "seq_len"],
79
+ x_vals=[list(_) for _ in configs],
80
+ line_arg="provider",
81
+ line_vals=["naive", "vllm"],
82
+ line_names=["Naive", "vLLM"],
83
+ styles=[("blue", "-"), ("red", "-")],
84
+ ylabel="us",
85
+ plot_name="polynorm-perf",
86
+ args={},
87
+ )
88
+ )
89
+ def benchmark(dim, batch_size, seq_len, provider):
90
+ dtype = torch.bfloat16
91
+ hidden_dim = dim * 4
92
+
93
+ x = torch.randn(batch_size, seq_len, hidden_dim, dtype=dtype, device="cuda")
94
+ weight = torch.ones(3, dtype=dtype, device="cuda")
95
+ bias = torch.ones(1, dtype=dtype, device="cuda")
96
+
97
+ quantiles = [0.5, 0.2, 0.8]
98
+
99
+ if provider == "naive":
100
+ ms, min_ms, max_ms = triton.testing.do_bench(
101
+ lambda: polynorm_naive(x, weight, bias),
102
+ quantiles=quantiles,
103
+ )
104
+ else:
105
+ ms, min_ms, max_ms = triton.testing.do_bench(
106
+ lambda: polynorm_vllm(x, weight, bias),
107
+ quantiles=quantiles,
108
+ )
109
+
110
+ return 1000 * ms, 1000 * max_ms, 1000 * min_ms
111
+
112
+ return benchmark
113
+
114
+
115
+if __name__ == "__main__":
116
+ import argparse
117
+
118
+ parser = argparse.ArgumentParser()
119
+ parser.add_argument(
120
+ "--batch-size",
121
+ type=int,
122
+ default=4,
123
+ help="Batch size",
124
+ )
125
+ parser.add_argument(
126
+ "--seq-len",
127
+ type=int,
128
+ default=128,
129
+ help="Sequence length",
130
+ )
131
+ parser.add_argument(
132
+ "--hidden-dim",
133
+ type=int,
134
+ default=8192,
135
+ help="Intermediate size of MLP",
136
+ )
137
+ parser.add_argument(
138
+ "--save-path",
139
+ type=str,
140
+ default="./configs/polnorm/",
141
+ help="Path to save polnorm benchmark results",
142
+ )
143
+
144
+ args = parser.parse_args()
145
+
146
+ # Run correctness test
147
+ calculate_diff(
148
+ batch_size=args.batch_size,
149
+ seq_len=args.seq_len,
150
+ hidden_dim=args.hidden_dim,
151
+ )
152
+
153
+ benchmark = get_benchmark()
154
+ # Run performance benchmark
155
+ benchmark.run(print_data=True, save_path=args.save_path)
csrc/layernorm_kernels.cu
+251
@@ -140,6 +140,211 @@ fused_add_rms_norm_kernel(
140
}
141
}
142
143
+/* Function specialization in the case of FP16/BF16 tensors.
144
+ Additional optimizations we can make in this case are
145
+ packed and vectorized operations, which help with the
146
+ memory latency bottleneck.
147
+
148
+ _f16VecPN struct extends _f16Vec to add operations specifically required for
149
+ polynomial normalization (poly norm).
150
+ The original _f16Vec does not include the sum-of-powers computation or
151
+ in-place polynomial normalization logic. */
152
+template <typename scalar_t, int width>
153
+struct alignas(16) _f16VecPN : _f16Vec<scalar_t, width> {
154
+ using Base = _f16Vec<scalar_t, width>;
155
+ using Converter = typename Base::Converter;
156
+ using T1 = typename Base::T1;
157
+ using T2 = typename Base::T2;
158
+ using Base::data;
159
+
160
+ __device__ auto sum_pows() const {
161
+ float s2 = 0.0f, s4 = 0.0f, s6 = 0.0f;
162
+
163
+#pragma unroll
164
+ for (int i = 0; i < width; i += 2) {
165
+ float2 z = Converter::convert(T2{data[i], data[i + 1]});
166
+ float x2 = z.x * z.x;
167
+ float x4 = x2 * x2;
168
+ float x6 = x4 * x2;
169
+
170
+ float y2 = z.y * z.y;
171
+ float y4 = y2 * y2;
172
+ float y6 = y4 * y2;
173
+
174
+ s2 += x2 + y2;
175
+ s4 += x4 + y4;
176
+ s6 += x6 + y6;
177
+ }
178
+ return std::make_tuple(s2, s4, s6);
179
+ }
180
+
181
+ __device__ void poly_norm_inplace(const float w2_inv_std,
182
+ const float w1_inv_std2,
183
+ const float w0_inv_std3, const float bias) {
184
+#pragma unroll
185
+ for (int i = 0; i < width; i += 2) {
186
+ float2 z = Converter::convert(T2{data[i], data[i + 1]});
187
+
188
+ float x2 = z.x * z.x;
189
+ float x3 = x2 * z.x;
190
+ z.x = w2_inv_std * z.x + w1_inv_std2 * x2 + w0_inv_std3 * x3 + bias;
191
+
192
+ float y2 = z.y * z.y;
193
+ float y3 = y2 * z.y;
194
+ z.y = w2_inv_std * z.y + w1_inv_std2 * y2 + w0_inv_std3 * y3 + bias;
195
+
196
+ auto out = Converter::convert(z);
197
+ data[i] = out.x;
198
+ data[i + 1] = out.y;
199
+ }
200
+ }
201
+};
202
+
203
+template <typename scalar_t, int width>
204
+__global__ std::enable_if_t<(width > 0) && _typeConvert<scalar_t>::exists>
205
+poly_norm_kernel(scalar_t* __restrict__ out, // [..., hidden_size]
206
+ const scalar_t* __restrict__ input, // [..., hidden_size]
207
+ const scalar_t* __restrict__ weight, // [3]
208
+ const scalar_t* __restrict__ bias, // [1]
209
+ const float epsilon, const int hidden_size) {
210
+ // Sanity checks on our vector struct and type-punned pointer arithmetic
211
+ static_assert(std::is_pod_v<_f16VecPN<scalar_t, width>>);
212
+ static_assert(sizeof(_f16VecPN<scalar_t, width>) == sizeof(scalar_t) * width);
213
+
214
+ /* These and the argument pointers are all declared `restrict` as they are
215
+ not aliased in practice. Argument pointers should not be dereferenced
216
+ in this kernel as that would be undefined behavior */
217
+ auto* __restrict__ input_v =
218
+ reinterpret_cast<const _f16VecPN<scalar_t, width>*>(input);
219
+ const int vec_hidden_size = hidden_size / width;
220
+ float variance = 0.0f;
221
+ float variance2 = 0.0f;
222
+ float variance3 = 0.0f;
223
+
224
+ for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
225
+ int id = blockIdx.x * vec_hidden_size + idx;
226
+ _f16VecPN<scalar_t, width> temp = input_v[id];
227
+ auto [x2, x4, x6] = temp.sum_pows();
228
+
229
+ variance += x2;
230
+ variance2 += x4;
231
+ variance3 += x6;
232
+ }
233
+
234
+ float3 thread_variances = make_float3(variance, variance2, variance3);
235
+
236
+ struct SumOp {
237
+ __device__ float3 operator()(const float3& a, const float3& b) const {
238
+ return make_float3(a.x + b.x, a.y + b.y, a.z + b.z);
239
+ }
240
+ };
241
+
242
+ using BlockReduce = cub::BlockReduce<float3, 1024>;
243
+ __shared__ typename BlockReduce::TempStorage reduceStore;
244
+ float3 block_variances =
245
+ BlockReduce(reduceStore).Reduce(thread_variances, SumOp{}, blockDim.x);
246
+
247
+ variance = block_variances.x;
248
+ variance2 = block_variances.y;
249
+ variance3 = block_variances.z;
250
+
251
+ __shared__ float s_w2_inv_std;
252
+ __shared__ float s_w1_inv_std2;
253
+ __shared__ float s_w0_inv_std3;
254
+ __shared__ float s_bias;
255
+
256
+ if (threadIdx.x == 0) {
257
+ float w0 = (float)weight[0];
258
+ float w1 = (float)weight[1];
259
+ float w2 = (float)weight[2];
260
+ s_bias = (float)bias[0];
261
+
262
+ s_w2_inv_std = w2 * rsqrtf(variance / hidden_size + epsilon);
263
+ s_w1_inv_std2 = w1 * rsqrtf(variance2 / hidden_size + epsilon);
264
+ s_w0_inv_std3 = w0 * rsqrtf(variance3 / hidden_size + epsilon);
265
+ }
266
+ __syncthreads();
267
+
268
+ auto* __restrict__ out_v = reinterpret_cast<_f16VecPN<scalar_t, width>*>(out);
269
+
270
+ for (int idx = threadIdx.x; idx < vec_hidden_size; idx += blockDim.x) {
271
+ int id = blockIdx.x * vec_hidden_size + idx;
272
+ _f16VecPN<scalar_t, width> temp = input_v[id];
273
+ temp.poly_norm_inplace(s_w2_inv_std, s_w1_inv_std2, s_w0_inv_std3, s_bias);
274
+ out_v[id] = temp;
275
+ }
276
+}
277
+
278
+/* Generic poly_norm_kernel
279
+ The width field is not used here but necessary for other specializations.
280
+ */
281
+template <typename scalar_t, int width>
282
+__global__ std::enable_if_t<(width == 0) || !_typeConvert<scalar_t>::exists>
283
+poly_norm_kernel(scalar_t* __restrict__ out, // [..., hidden_size]
284
+ const scalar_t* __restrict__ input, // [..., hidden_size]
285
+ const scalar_t* __restrict__ weight, // [3]
286
+ const scalar_t* __restrict__ bias, // [1]
287
+ const float epsilon, const int hidden_size) {
288
+ float variance = 0.0f;
289
+ float variance2 = 0.0f;
290
+ float variance3 = 0.0f;
291
+
292
+ for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
293
+ float x = (float)input[blockIdx.x * hidden_size + idx];
294
+ float x2 = x * x;
295
+ float x4 = x2 * x2;
296
+ float x6 = x4 * x2;
297
+
298
+ variance += x2;
299
+ variance2 += x4;
300
+ variance3 += x6;
301
+ }
302
+
303
+ float3 thread_variances = make_float3(variance, variance2, variance3);
304
+
305
+ struct SumOp {
306
+ __device__ float3 operator()(const float3& a, const float3& b) const {
307
+ return make_float3(a.x + b.x, a.y + b.y, a.z + b.z);
308
+ }
309
+ };
310
+
311
+ using BlockReduce = cub::BlockReduce<float3, 1024>;
312
+ __shared__ typename BlockReduce::TempStorage reduceStore;
313
+ float3 block_variances =
314
+ BlockReduce(reduceStore).Reduce(thread_variances, SumOp{}, blockDim.x);
315
+
316
+ variance = block_variances.x;
317
+ variance2 = block_variances.y;
318
+ variance3 = block_variances.z;
319
+
320
+ __shared__ float s_w2_inv_std;
321
+ __shared__ float s_w1_inv_std2;
322
+ __shared__ float s_w0_inv_std3;
323
+ __shared__ float s_bias;
324
+
325
+ if (threadIdx.x == 0) {
326
+ float w0 = (float)weight[0];
327
+ float w1 = (float)weight[1];
328
+ float w2 = (float)weight[2];
329
+ s_bias = (float)bias[0];
330
+
331
+ s_w2_inv_std = w2 * rsqrtf(variance / hidden_size + epsilon);
332
+ s_w1_inv_std2 = w1 * rsqrtf(variance2 / hidden_size + epsilon);
333
+ s_w0_inv_std3 = w0 * rsqrtf(variance3 / hidden_size + epsilon);
334
+ }
335
+ __syncthreads();
336
+
337
+ for (int idx = threadIdx.x; idx < hidden_size; idx += blockDim.x) {
338
+ float x = (float)input[blockIdx.x * hidden_size + idx];
339
+ float x2 = x * x;
340
+ float x3 = x2 * x;
341
+
342
+ out[blockIdx.x * hidden_size + idx] =
343
+ (scalar_t)(x * s_w2_inv_std + x2 * s_w1_inv_std2 + x3 * s_w0_inv_std3 +
344
+ s_bias);
345
+ }
346
+}
347
+
348
} // namespace vllm
349
350
void rms_norm(torch::Tensor& out, // [..., hidden_size]
@@ -219,3 +424,49 @@ void fused_add_rms_norm(torch::Tensor& input, // [..., hidden_size]
424
LAUNCH_FUSED_ADD_RMS_NORM(0);
425
}
426
}
427
+
428
+#define LAUNCH_FUSED_POLY_NORM(width) \
429
+ VLLM_DISPATCH_FLOATING_TYPES(input.scalar_type(), "poly_norm_kernel", [&] { \
430
+ vllm::poly_norm_kernel<scalar_t, width><<<grid, block, 0, stream>>>( \
431
+ out.data_ptr<scalar_t>(), input.data_ptr<scalar_t>(), \
432
+ weight.data_ptr<scalar_t>(), bias.data_ptr<scalar_t>(), epsilon, \
433
+ hidden_size); \
434
+ });
435
+
436
+void poly_norm(torch::Tensor& out, // [..., hidden_size]
437
+ torch::Tensor& input, // [..., hidden_size]
438
+ torch::Tensor& weight, // [3]
439
+ torch::Tensor& bias, // [1]
440
+ double epsilon) {
441
+ TORCH_CHECK(out.is_contiguous());
442
+ TORCH_CHECK(input.is_contiguous());
443
+ TORCH_CHECK(out.data_ptr() != input.data_ptr());
444
+
445
+ int hidden_size = input.size(-1);
446
+ int num_tokens = input.numel() / hidden_size;
447
+
448
+ dim3 grid(num_tokens);
449
+ /* This kernel is memory-latency bound in many scenarios.
450
+ When num_tokens is large, a smaller block size allows
451
+ for increased block occupancy on CUs and better latency
452
+ hiding on global mem ops. */
453
+ const int max_block_size = (num_tokens < 256) ? 1024 : 256;
454
+ dim3 block(std::min(hidden_size, max_block_size));
455
+ const at::cuda::OptionalCUDAGuard device_guard(device_of(input));
456
+ const cudaStream_t stream = at::cuda::getCurrentCUDAStream();
457
+ /*If the tensor types are FP16/BF16, try to use the optimized kernel
458
+ with packed + vectorized ops.
459
+ Max optimization is achieved with a width-8 vector of FP16/BF16s
460
+ since we can load at most 128 bits at once in a global memory op.
461
+ However, this requires each tensor's data to be aligned to 16
462
+ bytes.
463
+ */
464
+ auto inp_ptr = reinterpret_cast<std::uintptr_t>(input.data_ptr());
465
+ auto out_ptr = reinterpret_cast<std::uintptr_t>(out.data_ptr());
466
+ bool ptrs_are_aligned = inp_ptr % 16 == 0 && out_ptr % 16 == 0;
467
+ if (ptrs_are_aligned && hidden_size % 8 == 0) {
468
+ LAUNCH_FUSED_POLY_NORM(8);
469
+ } else {
470
+ LAUNCH_FUSED_POLY_NORM(0);
471
+ }
472
+}
csrc/ops.h
+4
-1
@@ -92,6 +92,9 @@ void rms_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight,
92
void fused_add_rms_norm(torch::Tensor& input, torch::Tensor& residual,
93
torch::Tensor& weight, double epsilon);
94
95
+void poly_norm(torch::Tensor& out, torch::Tensor& input, torch::Tensor& weight,
96
+ torch::Tensor& bias, double epsilon);
97
+
98
void apply_repetition_penalties_(torch::Tensor& logits,
99
const torch::Tensor& prompt_mask,
100
const torch::Tensor& output_mask,
@@ -353,4 +356,4 @@ void qr_open_handles(fptr_t _fa, const std::vector<torch::Tensor>& handles);
356
void qr_all_reduce(fptr_t _fa, torch::Tensor& inp, torch::Tensor& out,
357
int64_t quant_level, bool cast_bf2half = false);
358
int64_t qr_max_size();
356
-#endif
\ No newline at end of file
359
+#endif
csrc/torch_bindings.cpp
+6
@@ -168,6 +168,12 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
168
"float epsilon) -> ()");
169
ops.impl("fused_add_rms_norm", torch::kCUDA, &fused_add_rms_norm);
170
171
+ // Polynomial Normalization.
172
+ ops.def(
173
+ "poly_norm(Tensor! out, Tensor input, Tensor weight, Tensor bias, float "
174
+ "epsilon) -> ()");
175
+ ops.impl("poly_norm", torch::kCUDA, &poly_norm);
176
+
177
// Apply repetition penalties to logits in-place
178
ops.def(
179
"apply_repetition_penalties_(Tensor! logits, Tensor prompt_mask, "
docs/models/supported_models.md
+1
@@ -383,6 +383,7 @@ th {
383
| `MiniCPM3ForCausalLM` | MiniCPM3 | `openbmb/MiniCPM3-4B`, etc. | ✅︎ | ✅︎ | ✅︎ |
384
| `MistralForCausalLM` | Mistral, Mistral-Instruct | `mistralai/Mistral-7B-v0.1`, `mistralai/Mistral-7B-Instruct-v0.1`, etc. | ✅︎ | ✅︎ | ✅︎ |
385
| `MixtralForCausalLM` | Mixtral-8x7B, Mixtral-8x7B-Instruct | `mistralai/Mixtral-8x7B-v0.1`, `mistralai/Mixtral-8x7B-Instruct-v0.1`, `mistral-community/Mixtral-8x22B-v0.1`, etc. | ✅︎ | ✅︎ | ✅︎ |
386
+| `MotifForCausalLM` | Motif-1-Tiny | `Motif-Technologies/Motif-2.6B`, `Motif-Technologies/Motif-2.6b-v1.1-LC`, etc. | ✅︎ | ✅︎ | |
387
| `MPTForCausalLM` | MPT, MPT-Instruct, MPT-Chat, MPT-StoryWriter | `mosaicml/mpt-7b`, `mosaicml/mpt-7b-storywriter`, `mosaicml/mpt-30b`, etc. | | ✅︎ | ✅︎ |
388
| `NemotronForCausalLM` | Nemotron-3, Nemotron-4, Minitron | `nvidia/Minitron-8B-Base`, `mgoin/Nemotron-4-340B-Base-hf-FP8`, etc. | ✅︎ | ✅︎ | ✅︎ |
389
| `NemotronHForCausalLM` | Nemotron-H | `nvidia/Nemotron-H-8B-Base-8K`, `nvidia/Nemotron-H-47B-Base-8K`, `nvidia/Nemotron-H-56B-Base-8K`, etc. | ✅︎ | ✅︎ | ✅︎ |
tests/kernels/core/test_layernorm.py
+32
-1
@@ -6,7 +6,7 @@ import torch
6
7
from tests.kernels.quant_utils import FP8_DTYPE
8
from tests.kernels.utils import opcheck
9
-from vllm.model_executor.layers.layernorm import RMSNorm
9
+from vllm.model_executor.layers.layernorm import PolyNorm, RMSNorm
10
from vllm.platforms import current_platform
11
12
DTYPES = [torch.half, torch.bfloat16, torch.float]
@@ -70,6 +70,37 @@ def test_rms_norm(
70
(out, x, layer.weight.data, layer.variance_epsilon))
71
72
73
+@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
74
+@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
75
+@pytest.mark.parametrize("dtype", DTYPES)
76
+@pytest.mark.parametrize("seed", SEEDS)
77
+@pytest.mark.parametrize("device", CUDA_DEVICES)
78
+@torch.inference_mode()
79
+def test_poly_norm(
80
+ num_tokens: int,
81
+ hidden_size: int,
82
+ dtype: torch.dtype,
83
+ seed: int,
84
+ device: str,
85
+) -> None:
86
+ current_platform.seed_everything(seed)
87
+ torch.set_default_device(device)
88
+ layer = PolyNorm().to(dtype=dtype)
89
+ layer.weight.data.normal_(mean=1.0, std=0.1)
90
+ layer.bias.data.normal_(mean=1.0, std=0.1)
91
+ scale = 1 / (2 * hidden_size)
92
+ x = torch.randn(num_tokens, hidden_size, dtype=dtype)
93
+ x *= scale
94
+
95
+ ref_out = layer.forward_native(x)
96
+ out = layer(x)
97
+ torch.testing.assert_close(out, ref_out, atol=1e-2, rtol=1e-2)
98
+
99
+ opcheck(
100
+ torch.ops._C.poly_norm,
101
+ (out, x, layer.weight.data, layer.bias.data, layer.variance_epsilon))
102
+
103
+
104
@pytest.mark.parametrize("num_tokens", NUM_TOKENS)
105
@pytest.mark.parametrize("hidden_size", HIDDEN_SIZES)
106
@pytest.mark.parametrize("add_residual", ADD_RESIDUAL)
tests/models/registry.py
+3
@@ -291,6 +291,9 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {
291
"MistralForCausalLM": _HfExamplesInfo("mistralai/Mistral-7B-Instruct-v0.1"),
292
"MixtralForCausalLM": _HfExamplesInfo("mistralai/Mixtral-8x7B-Instruct-v0.1", # noqa: E501
293
{"tiny": "TitanML/tiny-mixtral"}), # noqa: E501
294
+ "MotifForCausalLM": _HfExamplesInfo("Motif-Technologies/Motif-2.6B",
295
+ trust_remote_code=True,
296
+ v0_only=True),
297
"MptForCausalLM": _HfExamplesInfo("mpt", is_available_online=False),
298
"MPTForCausalLM": _HfExamplesInfo("mosaicml/mpt-7b"),
299
"NemotronForCausalLM": _HfExamplesInfo("nvidia/Minitron-8B-Base"),
tests/models/test_initialization.py
+3
-2
@@ -63,8 +63,9 @@ def can_initialize(model_arch: str, monkeypatch: pytest.MonkeyPatch,
63
_initialize_kv_caches_v1), monkeypatch.context() as m):
64
if model_info.v0_only:
65
m.setenv("VLLM_USE_V1", "0")
66
- if model_arch == "Phi4FlashForCausalLM":
67
- # Phi4FlashForCausalLM only supports DIFFERENTIAL_FLASH_ATTN backend
66
+ if model_arch in ("Phi4FlashForCausalLM", "MotifForCausalLM"):
67
+ # Phi4FlashForCausalLM and MotifForCausalLM
68
+ # only supports DIFFERENTIAL_FLASH_ATTN backend
69
m.setenv("VLLM_ATTENTION_BACKEND", "DIFFERENTIAL_FLASH_ATTN")
70
if model_arch == "GptOssForCausalLM":
71
# FIXME: A hack to bypass FA3 assertion because our CI's L4 GPU
vllm/_custom_ops.py
+8
@@ -280,6 +280,13 @@ def fused_add_rms_norm(input: torch.Tensor, residual: torch.Tensor,
280
torch.ops._C.fused_add_rms_norm(input, residual, weight, epsilon)
281
282
283
+def poly_norm(out: torch.Tensor, input: torch.Tensor, weight: torch.Tensor,
284
+ bias: torch.Tensor, epsilon: float) -> None:
285
+ # TODO: Remove this contiguous call when the kernel is updated to support non-contiguous input
286
+ input_contiguous = input.contiguous()
287
+ torch.ops._C.poly_norm(out, input_contiguous, weight, bias, epsilon)
288
+
289
+
290
def apply_repetition_penalties_torch(
291
logits: torch.Tensor, prompt_mask: torch.Tensor,
292
output_mask: torch.Tensor, repetition_penalties: torch.Tensor) -> None:
@@ -710,6 +717,7 @@ def cutlass_sparse_scaled_mm_supported(cuda_device_capability: int) -> bool:
717
def cutlass_group_gemm_supported(cuda_device_capability: int) -> bool:
718
return torch.ops._C.cutlass_group_gemm_supported(cuda_device_capability)
719
720
+
721
def cutlass_sparse_compress(a: torch.Tensor) \
722
-> tuple[torch.Tensor, torch.Tensor]:
723
"""
vllm/attention/backends/differential_flash_attn.py
+3
@@ -734,6 +734,7 @@ class DifferentialFlashAttentionImpl(AttentionImpl):
734
window_size=self.sliding_window,
735
alibi_slopes=self.alibi_slopes,
736
softcap=self.logits_soft_cap,
737
+ fa_version=self.vllm_flash_attn_version,
738
)
739
assert prefill_output.shape == output[:
740
num_prefill_tokens].shape
@@ -755,6 +756,7 @@ class DifferentialFlashAttentionImpl(AttentionImpl):
756
window_size=self.sliding_window,
757
alibi_slopes=self.alibi_slopes,
758
softcap=self.logits_soft_cap,
759
+ fa_version=self.vllm_flash_attn_version,
760
).squeeze(1)
761
except Exception as e:
762
logger.error("Error in PagedAttention.forward_decode: %s",
@@ -787,6 +789,7 @@ class DifferentialFlashAttentionImpl(AttentionImpl):
789
window_size=self.sliding_window,
790
alibi_slopes=self.alibi_slopes,
791
softcap=self.logits_soft_cap,
792
+ fa_version=self.vllm_flash_attn_version,
793
).squeeze(1)
794
return output
795
vllm/model_executor/layers/layernorm.py
+59
@@ -43,6 +43,20 @@ def fused_add_rms_norm(
43
return x, residual
44
45
46
+def poly_norm(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor,
47
+ variance_epsilon: float) -> torch.Tensor:
48
+ from vllm import _custom_ops as ops
49
+ out = torch.empty_like(x)
50
+ ops.poly_norm(
51
+ out,
52
+ x,
53
+ weight,
54
+ bias,
55
+ variance_epsilon,
56
+ )
57
+ return out
58
+
59
+
60
def rocm_aiter_rms_norm_impl(x: torch.Tensor, weight: torch.Tensor,
61
variance_epsilon: float) -> torch.Tensor:
62
import aiter as rocm_aiter
@@ -320,3 +334,48 @@ class GemmaRMSNorm(CustomOp):
334
self.forward_static)
335
self._is_compiled = True
336
return self.forward_native(x, residual)
337
+
338
+
339
+@CustomOp.register("poly_norm")
340
+class PolyNorm(CustomOp):
341
+ """Polynomial normalization.
342
+
343
+ Computes x -> w_0 * RMSNorm(x^3) + w_1 * RMSNorm(x^2) + w_2 * RMSNorm(x) + b
344
+ where w_n is the learned weight and b is the bias.
345
+ Refer to https://arxiv.org/html/2411.03884v1
346
+ """
347
+
348
+ def __init__(
349
+ self,
350
+ eps: float = 1e-6,
351
+ ) -> None:
352
+ super().__init__()
353
+ self.weight = torch.nn.Parameter(torch.ones(3) / 3)
354
+ self.bias = torch.nn.Parameter(torch.zeros(1))
355
+ self.variance_epsilon = eps
356
+
357
+ def _norm(self, x):
358
+ return x / torch.sqrt(
359
+ x.pow(2).mean(-1, keepdim=True) + self.variance_epsilon)
360
+
361
+ def forward_native(
362
+ self,
363
+ x: torch.Tensor,
364
+ ) -> torch.Tensor:
365
+ """PyTorch-native implementation equivalent to forward().
366
+
367
+ Refer to https://github.com/BryceZhuo/PolyCom?tab=readme-ov-file/README.md
368
+ """
369
+
370
+ orig_dtype = x.dtype
371
+ x_float = x.to(torch.float32)
372
+ output = (self.weight[0] * self._norm(x_float**3) +
373
+ self.weight[1] * self._norm(x_float**2) +
374
+ self.weight[2] * self._norm(x_float) + self.bias)
375
+ return output.to(orig_dtype)
376
+
377
+ def forward_cuda(
378
+ self,
379
+ x: torch.Tensor,
380
+ ) -> torch.Tensor:
381
+ return poly_norm(x, self.weight, self.bias, self.variance_epsilon)
vllm/model_executor/models/motif.py
new
+345
@@ -0,0 +1,345 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+# Adapted from
5
+# https://huggingface.co/Motif-Technologies/Motif-2.6B/blob/main/modeling_motif.py
6
+# Copyright (c) Alibaba Cloud.
7
+# LICENSE: https://huggingface.co/Motif-Technologies/Motif-2.6B/blob/main/LICENSE
8
+"""Inference-only Motif model compatible with HuggingFace weights."""
9
+import math
10
+from typing import Any, Optional
11
+
12
+import torch
13
+from torch import nn
14
+from transformers import PretrainedConfig
15
+
16
+from vllm.attention import Attention, AttentionType
17
+from vllm.attention.selector import _Backend
18
+from vllm.config import CacheConfig, VllmConfig
19
+from vllm.distributed import get_tensor_model_parallel_world_size
20
+from vllm.model_executor.layers.layernorm import PolyNorm, RMSNorm
21
+from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
22
+ QKVParallelLinear,
23
+ RowParallelLinear)
24
+from vllm.model_executor.layers.quantization import QuantizationConfig
25
+from vllm.model_executor.layers.rotary_embedding import get_rope
26
+from vllm.model_executor.models.llama import LlamaForCausalLM
27
+
28
+from .adapters import as_seq_cls_model
29
+from .interfaces import SupportsV0Only
30
+from .utils import extract_layer_index
31
+
32
+
33
+class MotifMLP(nn.Module):
34
+ """MLP for the language component of the Motif model, which contains a
35
+ MergedColumnParallelLinear merging 2 outputs via PolyNorm activation."""
36
+
37
+ def __init__(
38
+ self,
39
+ hidden_size: int,
40
+ intermediate_size: int,
41
+ hidden_act: str = "poly_norm",
42
+ quant_config: Optional[QuantizationConfig] = None,
43
+ bias: bool = False,
44
+ prefix: str = "",
45
+ reduce_results: bool = True,
46
+ ):
47
+ super().__init__()
48
+ self.gate_up_proj = MergedColumnParallelLinear(
49
+ input_size=hidden_size,
50
+ output_sizes=[intermediate_size] * 2,
51
+ bias=bias,
52
+ quant_config=quant_config,
53
+ prefix=f"{prefix}.gate_up_proj",
54
+ )
55
+ self.down_proj = RowParallelLinear(
56
+ input_size=intermediate_size,
57
+ output_size=hidden_size,
58
+ bias=bias,
59
+ quant_config=quant_config,
60
+ reduce_results=reduce_results,
61
+ prefix=f"{prefix}.down_proj",
62
+ )
63
+ if hidden_act != "poly_norm":
64
+ raise NotImplementedError(f"Unsupported activation: {hidden_act}. "
65
+ "Only poly_norm is supported for now.")
66
+ self.act_fn = PolyNorm()
67
+ self.intermediate_size = intermediate_size
68
+ tp_size = get_tensor_model_parallel_world_size()
69
+ if hidden_act == "poly_norm" and tp_size > 1:
70
+ raise NotImplementedError(
71
+ "Tensor parallelism for poly_norm is not supported yet. "
72
+ "Support will be added in the future.")
73
+
74
+ def forward(self, x):
75
+ x, _ = self.gate_up_proj(x)
76
+ x = self.act_fn(
77
+ x[..., :self.intermediate_size]) * x[..., self.intermediate_size:]
78
+ x, _ = self.down_proj(x)
79
+ return x
80
+
81
+
82
+class MotifAttention(nn.Module):
83
+
84
+ def __init__(
85
+ self,
86
+ config: PretrainedConfig,
87
+ hidden_size: int,
88
+ num_heads: int,
89
+ num_kv_heads: int,
90
+ rope_theta: float = 10000,
91
+ rope_scaling: Optional[dict[str, Any]] = None,
92
+ max_position_embeddings: int = 8192,
93
+ quant_config: Optional[QuantizationConfig] = None,
94
+ bias: bool = False,
95
+ bias_o_proj: bool = False,
96
+ cache_config: Optional[CacheConfig] = None,
97
+ prefix: str = "",
98
+ attn_type: str = AttentionType.DECODER,
99
+ ) -> None:
100
+ super().__init__()
101
+ layer_idx = extract_layer_index(prefix)
102
+ self.hidden_size = hidden_size
103
+ tp_size = get_tensor_model_parallel_world_size()
104
+ self.total_num_heads = num_heads
105
+ assert self.total_num_heads % tp_size == 0
106
+ self.num_heads = self.total_num_heads // tp_size
107
+ self.total_num_kv_heads = num_kv_heads
108
+ if self.total_num_kv_heads >= tp_size:
109
+ # Number of KV heads is greater than TP size, so we partition
110
+ # the KV heads across multiple tensor parallel GPUs.
111
+ assert self.total_num_kv_heads % tp_size == 0
112
+ else:
113
+ # Number of KV heads is less than TP size, so we replicate
114
+ # the KV heads across multiple tensor parallel GPUs.
115
+ assert tp_size % self.total_num_kv_heads == 0
116
+ self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
117
+ # MistralConfig has an optional head_dim introduced by Mistral-Nemo
118
+ head_dim = getattr(config, "head_dim", None)
119
+ if head_dim is None:
120
+ head_dim = self.hidden_size // self.total_num_heads
121
+ self.head_dim = head_dim
122
+ # Phi models introduced a partial_rotary_factor parameter in the config
123
+ self.partial_rotary_factor = getattr(config, "partial_rotary_factor",
124
+ 1)
125
+ self.q_size = self.num_heads * self.head_dim
126
+ self.kv_size = self.num_kv_heads * self.head_dim
127
+ self.scaling = self.head_dim**-0.5
128
+ self.rope_theta = rope_theta
129
+ self.max_position_embeddings = max_position_embeddings
130
+
131
+ assert self.num_heads % 2 == 0, 'num_heads should be even'
132
+ assert self.num_kv_heads % 2 == 0, 'num_heads should be even'
133
+
134
+ self.qkv_proj = QKVParallelLinear(
135
+ hidden_size=hidden_size,
136
+ head_size=self.head_dim,
137
+ total_num_heads=self.total_num_heads,
138
+ total_num_kv_heads=self.total_num_kv_heads,
139
+ bias=bias,
140
+ quant_config=quant_config,
141
+ prefix=f"{prefix}.qkv_proj",
142
+ )
143
+
144
+ self.o_proj = RowParallelLinear(
145
+ input_size=self.total_num_heads * self.head_dim,
146
+ output_size=hidden_size,
147
+ bias=bias_o_proj,
148
+ quant_config=quant_config,
149
+ prefix=f"{prefix}.o_proj",
150
+ )
151
+
152
+ self._init_rotary_emb(config,
153
+ rope_scaling=rope_scaling,
154
+ quant_config=quant_config)
155
+ sliding_window = None
156
+
157
+ self.lambda_init = self.lambda_init_fn(layer_idx)
158
+ self.lambda_q1 = nn.Parameter(
159
+ torch.zeros(self.head_dim, dtype=torch.float32).normal_(mean=0,
160
+ std=0.1))
161
+ self.lambda_k1 = nn.Parameter(
162
+ torch.zeros(self.head_dim, dtype=torch.float32).normal_(mean=0,
163
+ std=0.1))
164
+ self.lambda_q2 = nn.Parameter(
165
+ torch.zeros(self.head_dim, dtype=torch.float32).normal_(mean=0,
166
+ std=0.1))
167
+ self.lambda_k2 = nn.Parameter(
168
+ torch.zeros(self.head_dim, dtype=torch.float32).normal_(mean=0,
169
+ std=0.1))
170
+ self.subln = RMSNorm(2 * self.head_dim, eps=config.attn_rms_norm_eps)
171
+
172
+ params = {
173
+ 'differential_flash_attention_config': {
174
+ 'lambda_init': self.lambda_init,
175
+ 'lambda_q1': self.lambda_q1,
176
+ 'lambda_k1': self.lambda_k1,
177
+ 'lambda_q2': self.lambda_q2,
178
+ 'lambda_k2': self.lambda_k2,
179
+ "subln": self.subln,
180
+ }
181
+ }
182
+
183
+ diff_attn_err_msg = (
184
+ 'Set VLLM_ATTENTION_BACKEND="DIFFERENTIAL_FLASH_ATTN" '
185
+ 'to enable Differential Flash Attention.')
186
+ try:
187
+ self.attn = Attention(
188
+ self.num_heads,
189
+ self.head_dim,
190
+ self.scaling,
191
+ num_kv_heads=self.num_kv_heads,
192
+ cache_config=cache_config,
193
+ quant_config=quant_config,
194
+ per_layer_sliding_window=sliding_window,
195
+ attn_type=attn_type,
196
+ prefix=f"{prefix}.attn",
197
+ **params,
198
+ )
199
+ except TypeError as e:
200
+ raise ValueError(diff_attn_err_msg) from e
201
+ assert (self.attn.backend == _Backend.DIFFERENTIAL_FLASH_ATTN
202
+ ), diff_attn_err_msg
203
+
204
+ def lambda_init_fn(self, depth):
205
+ return 0.8 - 0.6 * math.exp(-0.3 * (depth - 1))
206
+
207
+ def forward(
208
+ self,
209
+ positions: torch.Tensor,
210
+ hidden_states: torch.Tensor,
211
+ ) -> torch.Tensor:
212
+ qkv, _ = self.qkv_proj(hidden_states)
213
+ q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
214
+ q, k = self.rotary_emb(positions, q, k)
215
+ attn_output = self.attn(q, k, v)
216
+ output, _ = self.o_proj(attn_output)
217
+ return output
218
+
219
+ def _init_rotary_emb(self, config: PretrainedConfig,
220
+ rope_scaling: Optional[dict[str, Any]],
221
+ quant_config: Optional[QuantizationConfig]) -> None:
222
+ is_neox_style = True
223
+ is_gguf = quant_config and quant_config.get_name() == "gguf"
224
+ if is_gguf and config.model_type == "llama":
225
+ is_neox_style = False
226
+
227
+ self.rotary_emb = get_rope(
228
+ self.head_dim,
229
+ rotary_dim=self.head_dim,
230
+ max_position=self.max_position_embeddings,
231
+ base=self.rope_theta,
232
+ rope_scaling=rope_scaling,
233
+ is_neox_style=is_neox_style,
234
+ partial_rotary_factor=self.partial_rotary_factor,
235
+ )
236
+
237
+
238
+class MotifDecoderLayer(nn.Module):
239
+
240
+ def __init__(
241
+ self,
242
+ config: PretrainedConfig,
243
+ cache_config: Optional[CacheConfig] = None,
244
+ quant_config: Optional[QuantizationConfig] = None,
245
+ prefix: str = "",
246
+ ) -> None:
247
+ super().__init__()
248
+ self.hidden_size = config.hidden_size
249
+ rope_theta = getattr(config, "rope_theta", 10000)
250
+ rope_scaling = getattr(config, "rope_scaling", None)
251
+ if rope_scaling is not None and getattr(
252
+ config, "original_max_position_embeddings", None):
253
+ rope_scaling["original_max_position_embeddings"] = (
254
+ config.original_max_position_embeddings)
255
+ max_position_embeddings = getattr(config, "max_position_embeddings",
256
+ 8192)
257
+ attention_bias = getattr(config, "attention_bias", False) or getattr(
258
+ config, "use_bias", False)
259
+ bias_o_proj = attention_bias
260
+ if hasattr(config, 'qkv_bias'):
261
+ attention_bias = config.qkv_bias
262
+
263
+ # By default, Motif uses causal attention as it is a decoder-only model.
264
+ # You can override the HF config with `is_causal=False` to enable
265
+ # bidirectional attention, which is used in some embedding models
266
+ # (e.g. parasail-ai/GritLM-7B-vllm)
267
+ if getattr(config, "is_causal", True):
268
+ attn_type = AttentionType.DECODER
269
+ else:
270
+ attn_type = AttentionType.ENCODER_ONLY
271
+
272
+ self.self_attn = MotifAttention(
273
+ config=config,
274
+ hidden_size=self.hidden_size,
275
+ num_heads=config.num_attention_heads,
276
+ num_kv_heads=getattr(config, "num_key_value_heads",
277
+ config.num_attention_heads),
278
+ rope_theta=rope_theta,
279
+ rope_scaling=rope_scaling,
280
+ max_position_embeddings=max_position_embeddings,
281
+ quant_config=quant_config,
282
+ bias=attention_bias,
283
+ bias_o_proj=bias_o_proj,
284
+ cache_config=cache_config,
285
+ prefix=f"{prefix}.self_attn",
286
+ attn_type=attn_type,
287
+ )
288
+ self.mlp = MotifMLP(
289
+ hidden_size=self.hidden_size,
290
+ intermediate_size=config.intermediate_size,
291
+ hidden_act=config.hidden_act,
292
+ quant_config=quant_config,
293
+ bias=getattr(config, "use_bias", False),
294
+ prefix=f"{prefix}.mlp",
295
+ )
296
+ self.input_layernorm = RMSNorm(config.hidden_size,
297
+ eps=config.rms_norm_eps)
298
+ self.post_attention_layernorm = RMSNorm(config.hidden_size,
299
+ eps=config.rms_norm_eps)
300
+
301
+ def forward(
302
+ self,
303
+ positions: torch.Tensor,
304
+ hidden_states: torch.Tensor,
305
+ residual: Optional[torch.Tensor],
306
+ ) -> tuple[torch.Tensor, torch.Tensor]:
307
+ # Self Attention
308
+ if residual is None:
309
+ residual = hidden_states
310
+ hidden_states = self.input_layernorm(hidden_states)
311
+ else:
312
+ hidden_states, residual = self.input_layernorm(
313
+ hidden_states, residual)
314
+ hidden_states = self.self_attn(positions=positions,
315
+ hidden_states=hidden_states)
316
+
317
+ # Fully Connected
318
+ hidden_states, residual = self.post_attention_layernorm(
319
+ hidden_states, residual)
320
+ hidden_states = self.mlp(hidden_states)
321
+ return hidden_states, residual
322
+
323
+
324
+# Motif model uses differential attention
325
+# Only supported in v0 (no chunked prefill support)
326
+class MotifForCausalLM(LlamaForCausalLM, SupportsV0Only):
327
+
328
+ def __init__(self,
329
+ *,
330
+ vllm_config: VllmConfig,
331
+ prefix: str = "",
332
+ layer_type: type[nn.Module] = MotifDecoderLayer):
333
+
334
+ # Prefix caching and chunked prefill is not supported for this model.
335
+ assert not vllm_config.cache_config.enable_prefix_caching, \
336
+ "Motif currently does not support prefix caching"
337
+ assert not vllm_config.scheduler_config.chunked_prefill_enabled, \
338
+ "Motif currently does not support chunked prefill"
339
+
340
+ super().__init__(vllm_config=vllm_config,
341
+ prefix=prefix,
342
+ layer_type=layer_type)
343
+
344
+
345
+MotifForSequenceClassification = as_seq_cls_model(MotifForCausalLM)
vllm/model_executor/models/registry.py
+1
@@ -110,6 +110,7 @@ _TEXT_GENERATION_MODELS = {
110
"MiniCPM3ForCausalLM": ("minicpm3", "MiniCPM3ForCausalLM"),
111
"MistralForCausalLM": ("llama", "LlamaForCausalLM"),
112
"MixtralForCausalLM": ("mixtral", "MixtralForCausalLM"),
113
+ "MotifForCausalLM": ("motif", "MotifForCausalLM"),
114
# transformers's mpt class has lower case
115
"MptForCausalLM": ("mpt", "MPTForCausalLM"),
116
"MPTForCausalLM": ("mpt", "MPTForCausalLM"),