A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 15, 02:15 UTC
[Kernel] Triton implementation of causal-conv1d for Mamba-based models (#18218)
Signed-off-by: Tuan M. Hoang-Trong <tmhoangt@us.ibm.com> Co-authored-by: Tuan M. Hoang-Trong <tmhoangt@us.ibm.com> Co-authored-by: Tyler Michael Smith <tysmith@redhat.com> Co-authored-by: Tyler Michael Smith <tyler@neuralmagic.com>
Tuan, Hoang-Trong committed
Jul 9, 2025 at 15:53 UTC
47043eb6787bc8f62e4e746d7a9ad2451aaf7983
15 files changed
+1120
-1145
CMakeLists.txt
-1
@@ -232,7 +232,6 @@ endif()
232
233
set(VLLM_EXT_SRC
234
"csrc/mamba/mamba_ssm/selective_scan_fwd.cu"
235
- "csrc/mamba/causal_conv1d/causal_conv1d.cu"
235
"csrc/cache_kernels.cu"
236
"csrc/attention/paged_attention_v1.cu"
237
"csrc/attention/paged_attention_v2.cu"
csrc/mamba/causal_conv1d/causal_conv1d.cu
deleted
-656
@@ -1,656 +0,0 @@
1
-// clang-format off
2
-// adapted from https://github.com/Dao-AILab/causal-conv1d/blob/main/csrc/causal_conv1d_fwd.cu
3
-// and https://github.com/Dao-AILab/causal-conv1d/blob/main/csrc/causal_conv1d_update.cu
4
-#include <torch/all.h>
5
-#include <ATen/cuda/CUDAContext.h>
6
-#include <c10/cuda/CUDAGuard.h>
7
-
8
-#include "causal_conv1d.h"
9
-#include <c10/util/BFloat16.h>
10
-#include <c10/util/Half.h>
11
-#include <c10/cuda/CUDAException.h> // For C10_CUDA_CHECK and C10_CUDA_KERNEL_LAUNCH_CHECK
12
-
13
-#include <cub/block/block_load.cuh>
14
-#include <cub/block/block_store.cuh>
15
-
16
-#ifdef USE_ROCM
17
- namespace cub = hipcub;
18
-#endif
19
-
20
-#include "static_switch.h"
21
-
22
-
23
-
24
-#define CHECK_SHAPE(x, ...) TORCH_CHECK(x.sizes() == torch::IntArrayRef({__VA_ARGS__}), #x " must have shape (" #__VA_ARGS__ ")")
25
-
26
-#define DISPATCH_WTYPE_ITYPE_FLOAT_AND_HALF_AND_BF16(ITYPE, NAME, ...) \
27
- if (ITYPE == at::ScalarType::Half) { \
28
- using input_t = at::Half; \
29
- using weight_t = at::Half; \
30
- __VA_ARGS__(); \
31
- } else if (ITYPE == at::ScalarType::BFloat16) { \
32
- using input_t = at::BFloat16; \
33
- using weight_t = at::BFloat16; \
34
- __VA_ARGS__(); \
35
- } else if (ITYPE == at::ScalarType::Float) { \
36
- using input_t = float; \
37
- using weight_t = float; \
38
- __VA_ARGS__(); \
39
- } else { \
40
- AT_ERROR(#NAME, " not implemented for input type '", toString(ITYPE), "'"); \
41
- }
42
-
43
-
44
-template<typename input_t, typename weight_t>
45
-void causal_conv1d_fwd_cuda(ConvParamsBase ¶ms, cudaStream_t stream);
46
-
47
-template<typename input_t, typename weight_t>
48
-void causal_conv1d_update_cuda(ConvParamsBase ¶ms, cudaStream_t stream);
49
-
50
-void set_conv_params_fwd(ConvParamsBase ¶ms,
51
- // sizes
52
- const size_t batch,
53
- const size_t dim,
54
- const size_t seqlen,
55
- const size_t width,
56
- // device pointers
57
- const at::Tensor x,
58
- const at::Tensor weight,
59
- const at::Tensor out,
60
- const std::optional<at::Tensor>& bias,
61
- bool silu_activation,
62
- int64_t pad_slot_id,
63
- const std::optional<at::Tensor>& query_start_loc = std::nullopt,
64
- const std::optional<at::Tensor>& cache_indices = std::nullopt,
65
- const std::optional<at::Tensor>& has_initial_state = std::nullopt) {
66
-
67
- // Reset the parameters
68
- memset(¶ms, 0, sizeof(params));
69
-
70
- params.batch = batch;
71
- params.dim = dim;
72
- params.seqlen = seqlen;
73
- params.width = width;
74
- params.pad_slot_id = pad_slot_id;
75
-
76
- params.silu_activation = silu_activation;
77
-
78
- // Set the pointers and strides.
79
- params.x_ptr = x.data_ptr();
80
- params.weight_ptr = weight.data_ptr();
81
- params.bias_ptr = bias.has_value() ? bias.value().data_ptr() : nullptr;
82
- params.out_ptr = out.data_ptr();
83
- // All stride are in elements, not bytes.
84
- params.query_start_loc_ptr = query_start_loc.has_value() ? query_start_loc.value().data_ptr() : nullptr;
85
- params.cache_indices_ptr = cache_indices.has_value() ? cache_indices.value().data_ptr() : nullptr;
86
- params.has_initial_state_ptr = has_initial_state.has_value() ? has_initial_state.value().data_ptr() : nullptr;
87
- const bool varlen = params.query_start_loc_ptr != nullptr;
88
- params.x_batch_stride = x.stride(varlen ? 1 : 0);
89
- params.x_c_stride = x.stride(varlen ? 0 : 1);
90
- params.x_l_stride = x.stride(varlen ? 1 : -1);
91
- params.weight_c_stride = weight.stride(0);
92
- params.weight_width_stride = weight.stride(1);
93
- params.out_batch_stride = out.stride(varlen ? 1 : 0);
94
- params.out_c_stride = out.stride(varlen ? 0 : 1);
95
- params.out_l_stride = out.stride(varlen ? 1 : -1);
96
-}
97
-
98
-
99
-void causal_conv1d_fwd(const at::Tensor &x, const at::Tensor &weight,
100
- const std::optional<at::Tensor> &bias_,
101
- const std::optional<at::Tensor> &conv_states,
102
- const std::optional<at::Tensor> &query_start_loc,
103
- const std::optional<at::Tensor> &cache_indices,
104
- const std::optional<at::Tensor> &has_initial_state,
105
- bool silu_activation,
106
- // used to identify padding entries if cache_indices provided
107
- // in case of padding, the kernel will return early
108
- int64_t pad_slot_id) {
109
- auto input_type = x.scalar_type();
110
- auto weight_type = weight.scalar_type();
111
- TORCH_CHECK(input_type == at::ScalarType::Float || input_type == at::ScalarType::Half || input_type == at::ScalarType::BFloat16);
112
- TORCH_CHECK(weight_type == at::ScalarType::Float || weight_type == at::ScalarType::Half || weight_type == at::ScalarType::BFloat16);
113
-
114
- TORCH_CHECK(x.is_cuda());
115
- TORCH_CHECK(weight.is_cuda());
116
-
117
- const bool varlen = query_start_loc.has_value() ? true : false;
118
- const auto sizes = x.sizes();
119
- const int batch_size = varlen ? query_start_loc.value().sizes()[0] - 1 : sizes[0];
120
- const int dim = varlen ? sizes[0] : sizes[1];
121
- const int seqlen = varlen ? sizes[1] : sizes[2];
122
- const int width = weight.size(-1);
123
- if (varlen){
124
- CHECK_SHAPE(x, dim, seqlen);
125
- }
126
- else {
127
- CHECK_SHAPE(x, batch_size, dim, seqlen);
128
- }
129
- CHECK_SHAPE(weight, dim, width);
130
-
131
-
132
-
133
- if (bias_.has_value()) {
134
- auto bias = bias_.value();
135
- TORCH_CHECK(bias.scalar_type() == weight_type);
136
- TORCH_CHECK(bias.is_cuda());
137
- TORCH_CHECK(bias.stride(-1) == 1);
138
- CHECK_SHAPE(bias, dim);
139
- }
140
-
141
-
142
- if (has_initial_state.has_value()) {
143
- auto has_initial_state_ = has_initial_state.value();
144
- TORCH_CHECK(has_initial_state_.scalar_type() == at::ScalarType::Bool);
145
- TORCH_CHECK(has_initial_state_.is_cuda());
146
- CHECK_SHAPE(has_initial_state_, batch_size);
147
- }
148
-
149
-
150
- if (query_start_loc.has_value()) {
151
- auto query_start_loc_ = query_start_loc.value();
152
- TORCH_CHECK(query_start_loc_.scalar_type() == at::ScalarType::Int);
153
- TORCH_CHECK(query_start_loc_.is_cuda());
154
- }
155
-
156
-
157
- if (cache_indices.has_value()) {
158
- auto cache_indices_ = cache_indices.value();
159
- TORCH_CHECK(cache_indices_.scalar_type() == at::ScalarType::Int);
160
- TORCH_CHECK(cache_indices_.is_cuda());
161
- CHECK_SHAPE(cache_indices_, batch_size);
162
- }
163
-
164
- at::Tensor out = x;
165
-
166
- ConvParamsBase params;
167
- set_conv_params_fwd(params, batch_size, dim, seqlen, width, x, weight, out,
168
- bias_,
169
- silu_activation,
170
- pad_slot_id,
171
- query_start_loc,
172
- cache_indices,
173
- has_initial_state
174
- );
175
-
176
- if (conv_states.has_value()) {
177
- auto conv_states_ = conv_states.value();
178
- TORCH_CHECK(conv_states_.scalar_type() == input_type);
179
- TORCH_CHECK(conv_states_.is_cuda());
180
- params.conv_states_ptr = conv_states_.data_ptr();
181
- params.conv_states_batch_stride = conv_states_.stride(0);
182
- params.conv_states_c_stride = conv_states_.stride(1);
183
- params.conv_states_l_stride = conv_states_.stride(2);
184
- } else {
185
- params.conv_states_ptr = nullptr;
186
- }
187
-
188
- const at::cuda::OptionalCUDAGuard device_guard(device_of(x));
189
- auto stream = at::cuda::getCurrentCUDAStream().stream();
190
- DISPATCH_WTYPE_ITYPE_FLOAT_AND_HALF_AND_BF16(x.scalar_type(), "causal_conv1d_fwd", [&] {
191
- causal_conv1d_fwd_cuda<input_t, weight_t>(params, stream);
192
- });
193
-}
194
-
195
-
196
-void causal_conv1d_update(const at::Tensor &x,
197
- const at::Tensor &conv_state,
198
- const at::Tensor &weight,
199
- const std::optional<at::Tensor> &bias_,
200
- bool silu_activation,
201
- const std::optional<at::Tensor> &cache_seqlens_,
202
- const std::optional<at::Tensor> &conv_state_indices_,
203
- // used to identify padding entries if cache_indices provided
204
- // in case of padding, the kernel will return early
205
- int64_t pad_slot_id) {
206
- auto input_type = x.scalar_type();
207
- auto weight_type = weight.scalar_type();
208
- TORCH_CHECK(input_type == at::ScalarType::Float || input_type == at::ScalarType::Half || input_type == at::ScalarType::BFloat16);
209
- TORCH_CHECK(weight_type == at::ScalarType::Float || weight_type == at::ScalarType::Half || weight_type == at::ScalarType::BFloat16);
210
- TORCH_CHECK(weight_type == input_type, "weight type must equal to input type, other variations are disabled due to binary size limitations");
211
- TORCH_CHECK(conv_state.scalar_type() == input_type);
212
-
213
- TORCH_CHECK(x.is_cuda());
214
- TORCH_CHECK(conv_state.is_cuda());
215
- TORCH_CHECK(weight.is_cuda());
216
-
217
- const auto sizes = x.sizes();
218
- const int batch_size = sizes[0];
219
- const int dim = sizes[1];
220
- const int seqlen = sizes[2];
221
- const int width = weight.size(-1);
222
- const int conv_state_len = conv_state.size(2);
223
- TORCH_CHECK(conv_state_len >= width - 1);
224
-
225
- CHECK_SHAPE(x, batch_size, dim, seqlen);
226
- CHECK_SHAPE(weight, dim, width);
227
-
228
- TORCH_CHECK(width >= 2 && width <= 4, "causal_conv1d only supports width between 2 and 4");
229
-
230
- if (bias_.has_value()) {
231
- auto bias = bias_.value();
232
- TORCH_CHECK(bias.scalar_type() == weight_type);
233
- TORCH_CHECK(bias.is_cuda());
234
- TORCH_CHECK(bias.stride(-1) == 1);
235
- CHECK_SHAPE(bias, dim);
236
- }
237
-
238
- at::Tensor out = x;
239
-
240
- ConvParamsBase params;
241
- set_conv_params_fwd(params, batch_size, dim, seqlen, width, x, weight, out,
242
- bias_,
243
- silu_activation,
244
- pad_slot_id);
245
- params.conv_state_ptr = conv_state.data_ptr();
246
- params.conv_state_len = conv_state_len;
247
- // All stride are in elements, not bytes.
248
- params.conv_state_batch_stride = conv_state.stride(0);
249
- params.conv_state_c_stride = conv_state.stride(1);
250
- params.conv_state_l_stride = conv_state.stride(2);
251
-
252
- if (cache_seqlens_.has_value()) {
253
- auto cache_seqlens = cache_seqlens_.value();
254
- TORCH_CHECK(cache_seqlens.scalar_type() == torch::kInt32);
255
- TORCH_CHECK(cache_seqlens.is_cuda());
256
- TORCH_CHECK(cache_seqlens.stride(-1) == 1);
257
- CHECK_SHAPE(cache_seqlens, batch_size);
258
- params.cache_seqlens = cache_seqlens.data_ptr<int32_t>();
259
- } else {
260
- params.cache_seqlens = nullptr;
261
- }
262
-
263
- if (conv_state_indices_.has_value()) {
264
- auto conv_state_indices = conv_state_indices_.value();
265
- TORCH_CHECK(conv_state_indices.scalar_type() == torch::kInt32)
266
- TORCH_CHECK(conv_state_indices.is_cuda());
267
- TORCH_CHECK(conv_state_indices.stride(0) == 1)
268
- CHECK_SHAPE(conv_state_indices, batch_size);
269
-
270
- int conv_state_entries = conv_state.size(0);
271
- CHECK_SHAPE(conv_state, conv_state_entries, dim, conv_state_len);
272
-
273
- params.conv_state_indices_ptr = conv_state_indices.data_ptr<int32_t>();
274
- } else {
275
- CHECK_SHAPE(conv_state, batch_size, dim, conv_state_len);
276
- params.conv_state_indices_ptr = nullptr;
277
- }
278
-
279
- const at::cuda::OptionalCUDAGuard device_guard(device_of(x));
280
- auto stream = at::cuda::getCurrentCUDAStream().stream();
281
- DISPATCH_WTYPE_ITYPE_FLOAT_AND_HALF_AND_BF16(x.scalar_type(), "causal_conv1d_update", [&] {
282
- causal_conv1d_update_cuda<input_t, weight_t>(params, stream);
283
- });
284
-}
285
-
286
-template<int kNThreads_, int kWidth_, bool kIsVecLoad_, typename input_t_, typename weight_t_>
287
-struct Causal_conv1d_fwd_kernel_traits {
288
- using input_t = input_t_;
289
- using weight_t = weight_t_;
290
- static constexpr int kNThreads = kNThreads_;
291
- static constexpr int kWidth = kWidth_;
292
- static constexpr int kNBytes = sizeof(input_t);
293
- static_assert(kNBytes == 2 || kNBytes == 4);
294
- static constexpr int kNElts = kNBytes == 4 ? 4 : 8;
295
- static_assert(kWidth <= kNElts);
296
- static constexpr bool kIsVecLoad = kIsVecLoad_;
297
- using vec_t = typename BytesToType<kNBytes * kNElts>::Type;
298
- using BlockLoadT = cub::BlockLoad<input_t, kNThreads, kNElts, cub::BLOCK_LOAD_WARP_TRANSPOSE>;
299
- using BlockLoadVecT = cub::BlockLoad<vec_t, kNThreads, 1, cub::BLOCK_LOAD_DIRECT>;
300
- using BlockStoreT = cub::BlockStore<input_t, kNThreads, kNElts, cub::BLOCK_STORE_WARP_TRANSPOSE>;
301
- using BlockStoreVecT = cub::BlockStore<vec_t, kNThreads, 1, cub::BLOCK_STORE_DIRECT>;
302
- static constexpr int kSmemIOSize = kIsVecLoad
303
- ? 0
304
- : custom_max({sizeof(typename BlockLoadT::TempStorage), sizeof(typename BlockStoreT::TempStorage)});
305
- static constexpr int kSmemExchangeSize = kNThreads * kNBytes * kNElts;
306
- static constexpr int kSmemSize = kSmemIOSize + kSmemExchangeSize;
307
-};
308
-
309
-template<typename Ktraits>
310
-__global__ __launch_bounds__(Ktraits::kNThreads)
311
-void causal_conv1d_fwd_kernel(ConvParamsBase params) {
312
- constexpr int kWidth = Ktraits::kWidth;
313
- constexpr int kNThreads = Ktraits::kNThreads;
314
- constexpr int kNElts = Ktraits::kNElts;
315
- constexpr bool kIsVecLoad = Ktraits::kIsVecLoad;
316
- using input_t = typename Ktraits::input_t;
317
- using vec_t = typename Ktraits::vec_t;
318
- using weight_t = typename Ktraits::weight_t;
319
-
320
- // Shared memory.
321
- extern __shared__ char smem_[];
322
- auto& smem_load = reinterpret_cast<typename Ktraits::BlockLoadT::TempStorage&>(smem_);
323
- auto& smem_load_vec = reinterpret_cast<typename Ktraits::BlockLoadVecT::TempStorage&>(smem_);
324
- auto& smem_store = reinterpret_cast<typename Ktraits::BlockStoreT::TempStorage&>(smem_);
325
- auto& smem_store_vec = reinterpret_cast<typename Ktraits::BlockStoreVecT::TempStorage&>(smem_);
326
- vec_t *smem_exchange = reinterpret_cast<vec_t *>(smem_ + Ktraits::kSmemIOSize);
327
-
328
- const bool kVarlen = params.query_start_loc_ptr != nullptr;
329
- const int tidx = threadIdx.x;
330
- const int batch_id = blockIdx.x;
331
- const int channel_id = blockIdx.y;
332
- const int *query_start_loc = kVarlen ? reinterpret_cast<int *>(params.query_start_loc_ptr) : nullptr;
333
- const int sequence_start_index = kVarlen ? query_start_loc[batch_id] : batch_id;
334
- const int seqlen = kVarlen ? query_start_loc[batch_id + 1] - sequence_start_index : params.seqlen;
335
-
336
- input_t *x = reinterpret_cast<input_t *>(params.x_ptr) + sequence_start_index * params.x_batch_stride
337
- + channel_id * params.x_c_stride;
338
- weight_t *weight = reinterpret_cast<weight_t *>(params.weight_ptr) + channel_id * params.weight_c_stride;
339
- input_t *out = reinterpret_cast<input_t *>(params.out_ptr) + sequence_start_index * params.out_batch_stride
340
- + channel_id * params.out_c_stride;
341
- float bias_val = params.bias_ptr == nullptr ? 0.f : float(reinterpret_cast<weight_t *>(params.bias_ptr)[channel_id]);
342
-
343
- bool has_initial_state = params.has_initial_state_ptr == nullptr ? false
344
- : reinterpret_cast<bool *>(params.has_initial_state_ptr)[batch_id];
345
-
346
- int* cache_indices = params.cache_indices_ptr == nullptr ? nullptr
347
- : reinterpret_cast<int *>(params.cache_indices_ptr);
348
- int cache_index = cache_indices == nullptr ? batch_id : cache_indices[batch_id];
349
- // cache_index == params.pad_slot_id is defined as padding, so we exit early
350
- if (cache_index == params.pad_slot_id){
351
- return;
352
- }
353
- input_t *conv_states = params.conv_states_ptr == nullptr ? nullptr
354
- : reinterpret_cast<input_t *>(params.conv_states_ptr) + cache_index * params.conv_states_batch_stride + channel_id * params.conv_states_c_stride;
355
-
356
- // Thread 0 will load the last elements of the previous chunk, so we initialize those to 0.
357
- if (tidx == 0) {
358
- input_t initial_state[kNElts] = {0};
359
- if (has_initial_state) {
360
- #pragma unroll
361
- for (int w = 0; w < kWidth - 1; ++w){ initial_state[kNElts - 1 - (kWidth - 2) + w ] = conv_states[w]; }
362
- }
363
- smem_exchange[kNThreads - 1] = reinterpret_cast<vec_t *>(initial_state)[0];
364
- }
365
-
366
- float weight_vals[kWidth];
367
- #pragma unroll
368
- for (int i = 0; i < kWidth; ++i) { weight_vals[i] = float(weight[i * params.weight_width_stride]); }
369
-
370
- constexpr int kChunkSize = kNThreads * kNElts;
371
- const int n_chunks = (seqlen + kChunkSize - 1) / kChunkSize;
372
- for (int chunk = 0; chunk < n_chunks; ++chunk) {
373
- input_t x_vals_load[2 * kNElts] = {0};
374
- if constexpr(kIsVecLoad) {
375
- typename Ktraits::BlockLoadVecT(smem_load_vec).Load(reinterpret_cast<vec_t*>(x), *reinterpret_cast<vec_t (*)[1]>(&x_vals_load[kNElts]), (seqlen - chunk * kChunkSize) / kNElts);
376
- } else {
377
- __syncthreads();
378
- typename Ktraits::BlockLoadT(smem_load).Load(x, *reinterpret_cast<input_t (*)[kNElts]>(&x_vals_load[kNElts]), seqlen - chunk * kChunkSize);
379
- }
380
- x += kChunkSize;
381
- __syncthreads();
382
- // Thread kNThreads - 1 don't write yet, so that thread 0 can read
383
- // the last elements of the previous chunk.
384
- if (tidx < kNThreads - 1) { smem_exchange[tidx] = reinterpret_cast<vec_t *>(x_vals_load)[1]; }
385
- __syncthreads();
386
- reinterpret_cast<vec_t *>(x_vals_load)[0] = smem_exchange[tidx > 0 ? tidx - 1 : kNThreads - 1];
387
- __syncthreads();
388
- // Now thread kNThreads - 1 can write the last elements of the current chunk.
389
- if (tidx == kNThreads - 1) { smem_exchange[tidx] = reinterpret_cast<vec_t *>(x_vals_load)[1]; }
390
-
391
- float x_vals[2 * kNElts];
392
- #pragma unroll
393
- for (int i = 0; i < 2 * kNElts; ++i) { x_vals[i] = float(x_vals_load[i]); }
394
-
395
- float out_vals[kNElts];
396
- #pragma unroll
397
- for (int i = 0; i < kNElts; ++i) {
398
- out_vals[i] = bias_val;
399
- #pragma unroll
400
- for (int w = 0; w < kWidth; ++w) {
401
- out_vals[i] += weight_vals[w] * x_vals[kNElts + i - (kWidth - w - 1)];
402
- }
403
- }
404
-
405
- if (params.silu_activation) {
406
- #pragma unroll
407
- for (int i = 0; i < kNElts; ++i) {
408
- out_vals[i] = out_vals[i] / (1 + expf(-out_vals[i]));
409
- }
410
- }
411
-
412
- input_t out_vals_store[kNElts];
413
- #pragma unroll
414
- for (int i = 0; i < kNElts; ++i) { out_vals_store[i] = out_vals[i]; }
415
- if constexpr(kIsVecLoad) {
416
- typename Ktraits::BlockStoreVecT(smem_store_vec).Store(reinterpret_cast<vec_t*>(out), reinterpret_cast<vec_t (&)[1]>(out_vals_store), (seqlen - chunk * kChunkSize) / kNElts);
417
- } else {
418
- typename Ktraits::BlockStoreT(smem_store).Store(out, out_vals_store, seqlen - chunk * kChunkSize);
419
- }
420
- out += kChunkSize;
421
-
422
- int final_state_position = ((seqlen - (kWidth - 1)) - (n_chunks - 1) * kChunkSize);
423
- // in case the final state is separated between the last "smem_exchange" and
424
- // and the one before it (chunk = n_chunks - 1 and chunk = n_chunks - 2),
425
- // (which occurs when `final_state_position` is a non-positive index)
426
- // we load the correct data from smem_exchange from both chunks, the last chunk iteration and the one before it
427
- if (conv_states != nullptr && final_state_position < 0 && seqlen > kWidth){
428
- input_t vals_load[kNElts] = {0};
429
- if ((chunk == n_chunks - 2) && (tidx == kNThreads - 1)){
430
- // chunk = n_chunks - 2, a segment of the final state sits in the last index
431
- reinterpret_cast<vec_t *>(vals_load)[0] = smem_exchange[kNThreads - 1];
432
- #pragma unroll
433
- for (int w = 0; w < -final_state_position; ++w){
434
- conv_states[w] = vals_load[kNElts + final_state_position + w];
435
- }
436
- }
437
- if ((chunk == n_chunks - 1) && tidx == 0){
438
- // chunk = n_chunks - 1, the second segment of the final state first positions
439
- reinterpret_cast<vec_t *>(vals_load)[0] = smem_exchange[0];
440
- for (int w = -final_state_position; w < kWidth - 1; ++w){
441
- conv_states[w] = vals_load[w + final_state_position];
442
- }
443
- return;
444
- }
445
- }
446
- }
447
- // Final state is stored in the smem_exchange last token slot,
448
- // in case seqlen < kWidth, we would need to take the final state from the
449
- // initial state which is stored in conv_states
450
- // in case seqlen > kWidth, we would need to load the last kWidth - 1 data
451
- // and load it into conv_state accordingly
452
- int last_thread = ((seqlen - (kWidth - 1)) - (n_chunks - 1) * kChunkSize) / kNElts;
453
- if (conv_states != nullptr && tidx == last_thread) {
454
- input_t x_vals_load[kNElts * 2] = {0};
455
- // in case we are on the first kWidth tokens
456
- if (last_thread == 0 && seqlen < kWidth){
457
- // Need to take the initial state
458
- reinterpret_cast<vec_t *>(x_vals_load)[0] = smem_exchange[0];
459
- const int offset = seqlen - (kWidth - 1);
460
- #pragma unroll
461
- for (int w = 0; w < kWidth - 1; ++w){
462
- // pad the existing state
463
- if ((w - seqlen) >= 0 && has_initial_state) { conv_states[w - seqlen] = conv_states[w]; }
464
- else if ((w - seqlen) >= 0 && !has_initial_state) { conv_states[w - seqlen] = input_t(0.0f); }
465
- }
466
- #pragma unroll
467
- for (int w = 0; w < kWidth - 1; ++w){
468
- if (offset + w >= 0)
469
- conv_states[w] = x_vals_load[offset + w ];
470
- }
471
- }
472
- else {
473
- // in case the final state is in between the threads data
474
- const int offset = ((seqlen - (kWidth - 1)) % (kNElts));
475
- if ((offset + kWidth - 2) >= kNElts && (last_thread + 1 < kNThreads)){
476
- // In case last_thread == kNThreads - 1, accessing last_thread + 1 will result in a
477
- // illegal access error on H100.
478
- // Therefore, we access last_thread + 1, only if the final state data sits there
479
- reinterpret_cast<vec_t *>(x_vals_load)[1] = smem_exchange[last_thread + 1];
480
- }
481
- reinterpret_cast<vec_t *>(x_vals_load)[0] = smem_exchange[last_thread];
482
- #pragma unroll
483
- for (int w = 0; w < kWidth - 1; ++w){
484
- conv_states[w] = x_vals_load[offset + w ];
485
- }
486
- }
487
-
488
- }
489
-}
490
-
491
-
492
-template<int kNThreads, int kWidth, typename input_t, typename weight_t>
493
-void causal_conv1d_fwd_launch(ConvParamsBase ¶ms, cudaStream_t stream) {
494
- static constexpr int kNElts = sizeof(input_t) == 4 ? 4 : 8;
495
- const bool kVarlen = params.query_start_loc_ptr != nullptr;
496
- BOOL_SWITCH(params.seqlen % kNElts == 0 && !kVarlen, kIsVecLoad, [&] {
497
- using Ktraits = Causal_conv1d_fwd_kernel_traits<kNThreads, kWidth, kIsVecLoad, input_t, weight_t>;
498
- constexpr int kSmemSize = Ktraits::kSmemSize;
499
- dim3 grid(params.batch, params.dim);
500
-
501
- auto kernel = &causal_conv1d_fwd_kernel<Ktraits>;
502
-
503
- if (kSmemSize >= 48 * 1024) {
504
- C10_CUDA_CHECK(cudaFuncSetAttribute(
505
- (void *) kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, kSmemSize));
506
- std::cerr << "Warning (causal_conv1d fwd launch): attempting to set maxDynamicSharedMemorySize on an AMD GPU which is currently a non-op (in ROCm versions <= 6.1). This might lead to undefined behavior. \n" << std::endl;
507
- }
508
- kernel<<<grid, Ktraits::kNThreads, kSmemSize, stream>>>(params);
509
-
510
- C10_CUDA_KERNEL_LAUNCH_CHECK();
511
- });
512
-}
513
-
514
-template<typename input_t, typename weight_t>
515
-void causal_conv1d_fwd_cuda(ConvParamsBase ¶ms, cudaStream_t stream) {
516
- if (params.width == 2) {
517
- causal_conv1d_fwd_launch<128, 2, input_t, weight_t>(params, stream);
518
- } else if (params.width == 3) {
519
- causal_conv1d_fwd_launch<128, 3, input_t, weight_t>(params, stream);
520
- } else if (params.width == 4) {
521
- causal_conv1d_fwd_launch<128, 4, input_t, weight_t>(params, stream);
522
- }
523
-}
524
-
525
-
526
-template void causal_conv1d_fwd_cuda<float, float>(ConvParamsBase ¶ms, cudaStream_t stream);
527
-template void causal_conv1d_fwd_cuda<at::Half, at::Half>(ConvParamsBase ¶ms, cudaStream_t stream);
528
-template void causal_conv1d_fwd_cuda<at::BFloat16, at::BFloat16>(ConvParamsBase ¶ms, cudaStream_t stream);
529
-
530
-
531
-
532
-
533
-template<int kNThreads_, int kWidth_, typename input_t_, typename weight_t_>
534
-struct Causal_conv1d_update_kernel_traits {
535
- using input_t = input_t_;
536
- using weight_t = weight_t_;
537
- static constexpr int kNThreads = kNThreads_;
538
- static constexpr int kWidth = kWidth_;
539
- static constexpr int kNBytes = sizeof(input_t);
540
- static_assert(kNBytes == 2 || kNBytes == 4);
541
-};
542
-
543
-template<typename Ktraits, bool kIsCircularBuffer>
544
-__global__ __launch_bounds__(Ktraits::kNThreads)
545
-void causal_conv1d_update_kernel(ConvParamsBase params) {
546
- constexpr int kWidth = Ktraits::kWidth;
547
- constexpr int kNThreads = Ktraits::kNThreads;
548
- using input_t = typename Ktraits::input_t;
549
- using weight_t = typename Ktraits::weight_t;
550
-
551
- const int tidx = threadIdx.x;
552
- const int batch_id = blockIdx.x;
553
- const int channel_id = blockIdx.y * kNThreads + tidx;
554
- if (channel_id >= params.dim) return;
555
-
556
- input_t *x = reinterpret_cast<input_t *>(params.x_ptr) + batch_id * params.x_batch_stride
557
- + channel_id * params.x_c_stride;
558
-
559
- // If params.conv_state_batch_indices is set, then the conv state is gathered from the conv state tensor
560
- // along the batch axis. Otherwise, the conv state coordinate is the same as the batch id.
561
- const int conv_state_batch_coord = params.conv_state_indices_ptr == nullptr
562
- ? batch_id
563
- : params.conv_state_indices_ptr[batch_id];
564
- // conv_state_batch_coord == params.pad_slot_id is defined as padding so we exit early
565
- if (conv_state_batch_coord == params.pad_slot_id){
566
- return;
567
- }
568
- input_t *conv_state = reinterpret_cast<input_t *>(params.conv_state_ptr)
569
- + conv_state_batch_coord * params.conv_state_batch_stride
570
- + channel_id * params.conv_state_c_stride;
571
-
572
- weight_t *weight = reinterpret_cast<weight_t *>(params.weight_ptr) + channel_id * params.weight_c_stride;
573
- input_t *out = reinterpret_cast<input_t *>(params.out_ptr) + batch_id * params.out_batch_stride
574
- + channel_id * params.out_c_stride;
575
- float bias_val = params.bias_ptr == nullptr ? 0.f : float(reinterpret_cast<weight_t *>(params.bias_ptr)[channel_id]);
576
-
577
- int state_len = params.conv_state_len;
578
- int advance_len = params.seqlen;
579
- int cache_seqlen = kIsCircularBuffer ? params.cache_seqlens[batch_id] % state_len : 0;
580
- int update_idx = cache_seqlen - (kWidth - 1);
581
- update_idx = update_idx < 0 ? update_idx + state_len : update_idx;
582
-
583
- float weight_vals[kWidth] = {0};
584
- #pragma unroll
585
- for (int i = 0; i < kWidth; ++i) { weight_vals[i] = float(weight[i * params.weight_width_stride]); }
586
-
587
- float x_vals[kWidth] = {0};
588
- if constexpr (!kIsCircularBuffer) {
589
- #pragma unroll 2
590
- for (int i = 0; i < state_len - advance_len - (kWidth - 1); ++i) {
591
- conv_state[i * params.conv_state_l_stride] = conv_state[(i + advance_len) * params.conv_state_l_stride];
592
- }
593
- #pragma unroll
594
- for (int i = 0; i < kWidth - 1; ++i) {
595
- input_t state_val = conv_state[(state_len - (kWidth - 1) + i) * params.conv_state_l_stride];
596
- if (i < advance_len + (kWidth - 1) && state_len - advance_len - (kWidth - 1) + i >= 0) {
597
- conv_state[(state_len - advance_len - (kWidth - 1) + i) * params.conv_state_l_stride] = state_val;
598
- }
599
- x_vals[i] = float(state_val);
600
- }
601
- } else {
602
- #pragma unroll
603
- for (int i = 0; i < kWidth - 1; ++i, update_idx = update_idx + 1 >= state_len ? update_idx + 1 - state_len : update_idx + 1) {
604
- input_t state_val = conv_state[update_idx * params.conv_state_l_stride];
605
- x_vals[i] = float(state_val);
606
- }
607
- }
608
- #pragma unroll 2
609
- for (int i = 0; i < params.seqlen; ++i) {
610
- input_t x_val = x[i * params.x_l_stride];
611
- if constexpr (!kIsCircularBuffer) {
612
- if (i < advance_len && state_len - advance_len + i >= 0) {
613
- conv_state[(state_len - advance_len + i) * params.conv_state_l_stride] = x_val;
614
- }
615
- } else {
616
- conv_state[update_idx * params.conv_state_l_stride] = x_val;
617
- ++update_idx;
618
- update_idx = update_idx >= state_len ? update_idx - state_len : update_idx;
619
- }
620
- x_vals[kWidth - 1] = float(x_val);
621
- float out_val = bias_val;
622
- #pragma unroll
623
- for (int j = 0; j < kWidth; ++j) { out_val += weight_vals[j] * x_vals[j]; }
624
- if (params.silu_activation) { out_val = out_val / (1 + expf(-out_val)); }
625
- out[i * params.out_l_stride] = input_t(out_val);
626
- // Shift the input buffer by 1
627
- #pragma unroll
628
- for (int i = 0; i < kWidth - 1; ++i) { x_vals[i] = x_vals[i + 1]; }
629
- }
630
-}
631
-
632
-template<int kNThreads, int kWidth, typename input_t, typename weight_t>
633
-void causal_conv1d_update_launch(ConvParamsBase ¶ms, cudaStream_t stream) {
634
- using Ktraits = Causal_conv1d_update_kernel_traits<kNThreads, kWidth, input_t, weight_t>;
635
- dim3 grid(params.batch, (params.dim + kNThreads - 1) / kNThreads);
636
- auto kernel = params.cache_seqlens == nullptr
637
- ? &causal_conv1d_update_kernel<Ktraits, false>
638
- : &causal_conv1d_update_kernel<Ktraits, true>;
639
- kernel<<<grid, Ktraits::kNThreads, 0, stream>>>(params);
640
- C10_CUDA_KERNEL_LAUNCH_CHECK();
641
-}
642
-
643
-template<typename input_t, typename weight_t>
644
-void causal_conv1d_update_cuda(ConvParamsBase ¶ms, cudaStream_t stream) {
645
- if (params.width == 2) {
646
- causal_conv1d_update_launch<64, 2, input_t, weight_t>(params, stream);
647
- } else if (params.width == 3) {
648
- causal_conv1d_update_launch<64, 3, input_t, weight_t>(params, stream);
649
- } else if (params.width == 4) {
650
- causal_conv1d_update_launch<64, 4, input_t, weight_t>(params, stream);
651
- }
652
-}
653
-
654
-template void causal_conv1d_update_cuda<float, float>(ConvParamsBase ¶ms, cudaStream_t stream);
655
-template void causal_conv1d_update_cuda<at::Half, at::Half>(ConvParamsBase ¶ms, cudaStream_t stream);
656
-template void causal_conv1d_update_cuda<at::BFloat16, at::BFloat16>(ConvParamsBase ¶ms, cudaStream_t stream);
csrc/mamba/causal_conv1d/causal_conv1d.h
deleted
-159
@@ -1,159 +0,0 @@
1
-/******************************************************************************
2
- * Copyright (c) 2024, Tri Dao.
3
- ******************************************************************************/
4
-// clang-format off
5
-// adapted from https://github.com/Dao-AILab/causal-conv1d/blob/main/csrc/causal_conv1d.h
6
-#pragma once
7
-
8
-#include <cuda_bf16.h>
9
-#include <cuda_fp16.h>
10
-////////////////////////////////////////////////////////////////////////////////////////////////////
11
-
12
-struct ConvParamsBase {
13
- using index_t = uint32_t;
14
-
15
- int batch, dim, seqlen, width;
16
- int64_t pad_slot_id;
17
- bool silu_activation;
18
-
19
- index_t x_batch_stride;
20
- index_t x_c_stride;
21
- index_t x_l_stride;
22
- index_t weight_c_stride;
23
- index_t weight_width_stride;
24
- index_t out_batch_stride;
25
- index_t out_c_stride;
26
- index_t out_l_stride;
27
-
28
- int conv_state_len;
29
- index_t conv_state_batch_stride;
30
- index_t conv_state_c_stride;
31
- index_t conv_state_l_stride;
32
-
33
- // Common data pointers.
34
- void *__restrict__ x_ptr;
35
- void *__restrict__ weight_ptr;
36
- void *__restrict__ bias_ptr;
37
- void *__restrict__ out_ptr;
38
-
39
- void *__restrict__ conv_state_ptr;
40
- void *__restrict__ query_start_loc_ptr;
41
- void *__restrict__ has_initial_state_ptr;
42
- void *__restrict__ cache_indices_ptr;
43
- int32_t *__restrict__ cache_seqlens;
44
-
45
- // For the continuous batching case. Makes it so that the mamba state for
46
- // the current batch doesn't need to be a contiguous tensor.
47
- int32_t *__restrict__ conv_state_indices_ptr;
48
-
49
- void *__restrict__ seq_idx_ptr;
50
-
51
- // No __restrict__ since initial_states could be the same as final_states.
52
- void * initial_states_ptr;
53
- index_t initial_states_batch_stride;
54
- index_t initial_states_l_stride;
55
- index_t initial_states_c_stride;
56
-
57
- void * final_states_ptr;
58
- index_t final_states_batch_stride;
59
- index_t final_states_l_stride;
60
- index_t final_states_c_stride;
61
-
62
- void * conv_states_ptr;
63
- index_t conv_states_batch_stride;
64
- index_t conv_states_l_stride;
65
- index_t conv_states_c_stride;
66
-};
67
-
68
-
69
-#ifndef USE_ROCM
70
- #include <cuda_bf16.h>
71
-
72
- template<typename T>
73
- __device__ inline T shuffle_xor(T val, int offset) {
74
- return __shfl_xor_sync(uint32_t(-1), val, offset);
75
- }
76
-
77
- constexpr size_t custom_max(std::initializer_list<size_t> ilist)
78
- {
79
- return std::max(ilist);
80
- }
81
-
82
- template<typename T>
83
- constexpr T constexpr_min(T a, T b) {
84
- return std::min(a, b);
85
- }
86
-
87
-#else
88
- #include <hip/hip_bf16.h>
89
-
90
- template<typename T>
91
- __device__ inline T shuffle_xor(T val, int offset) {
92
- return __shfl_xor(val, offset);
93
- }
94
- constexpr size_t custom_max(std::initializer_list<size_t> ilist)
95
- {
96
- return *std::max_element(ilist.begin(), ilist.end());
97
- }
98
-
99
- template<typename T>
100
- constexpr T constexpr_min(T a, T b) {
101
- return a < b ? a : b;
102
- }
103
-#endif
104
-
105
-////////////////////////////////////////////////////////////////////////////////////////////////////
106
-
107
-template<int BYTES> struct BytesToType {};
108
-
109
-template<> struct BytesToType<16> {
110
- using Type = uint4;
111
- static_assert(sizeof(Type) == 16);
112
-};
113
-
114
-template<> struct BytesToType<8> {
115
- using Type = uint64_t;
116
- static_assert(sizeof(Type) == 8);
117
-};
118
-
119
-template<> struct BytesToType<4> {
120
- using Type = uint32_t;
121
- static_assert(sizeof(Type) == 4);
122
-};
123
-
124
-template<> struct BytesToType<2> {
125
- using Type = uint16_t;
126
- static_assert(sizeof(Type) == 2);
127
-};
128
-
129
-template<> struct BytesToType<1> {
130
- using Type = uint8_t;
131
- static_assert(sizeof(Type) == 1);
132
-};
133
-
134
-////////////////////////////////////////////////////////////////////////////////////////////////////
135
-
136
-template<typename T>
137
-struct SumOp {
138
-__device__ inline T operator()(T const & x, T const & y) { return x + y; }
139
-};
140
-
141
-template<int THREADS>
142
-struct Allreduce {
143
- static_assert(THREADS == 32 || THREADS == 16 || THREADS == 8 || THREADS == 4);
144
- template<typename T, typename Operator>
145
- static __device__ inline T run(T x, Operator &op) {
146
- constexpr int OFFSET = THREADS / 2;
147
- x = op(x, __shfl_xor_sync(uint32_t(-1), x, OFFSET));
148
- return Allreduce<OFFSET>::run(x, op);
149
- }
150
-};
151
-
152
-template<>
153
-struct Allreduce<2> {
154
-template<typename T, typename Operator>
155
-static __device__ inline T run(T x, Operator &op) {
156
- x = op(x, __shfl_xor_sync(uint32_t(-1), x, 1));
157
- return x;
158
-}
159
-};
csrc/mamba/causal_conv1d/static_switch.h
deleted
-28
@@ -1,28 +0,0 @@
1
-// Inspired by
2
-// https://github.com/NVIDIA/DALI/blob/main/include/dali/core/static_switch.h
3
-// and https://github.com/pytorch/pytorch/blob/master/aten/src/ATen/Dispatch.h
4
-// clang-format off
5
-// adapted from https://github.com/Dao-AILab/causal-conv1d/blob/main/csrc/static_switch.h
6
-
7
-#pragma once
8
-
9
-/// @param COND - a boolean expression to switch by
10
-/// @param CONST_NAME - a name given for the constexpr bool variable.
11
-/// @param ... - code to execute for true and false
12
-///
13
-/// Usage:
14
-/// ```
15
-/// BOOL_SWITCH(flag, BoolConst, [&] {
16
-/// some_function<BoolConst>(...);
17
-/// });
18
-/// ```
19
-#define BOOL_SWITCH(COND, CONST_NAME, ...) \
20
- [&] { \
21
- if (COND) { \
22
- static constexpr bool CONST_NAME = true; \
23
- return __VA_ARGS__(); \
24
- } else { \
25
- static constexpr bool CONST_NAME = false; \
26
- return __VA_ARGS__(); \
27
- } \
28
- }()
csrc/ops.h
-16
@@ -326,22 +326,6 @@ void selective_scan_fwd(const torch::Tensor& u, const torch::Tensor& delta,
326
const std::optional<torch::Tensor>& has_initial_state,
327
const torch::Tensor& ssm_states, int64_t pad_slot_id);
328
329
-void causal_conv1d_update(const at::Tensor& x, const at::Tensor& conv_state,
330
- const at::Tensor& weight,
331
- const std::optional<at::Tensor>& bias_,
332
- bool silu_activation,
333
- const std::optional<at::Tensor>& cache_seqlens_,
334
- const std::optional<at::Tensor>& conv_state_indices_,
335
- int64_t pad_slot_id);
336
-
337
-void causal_conv1d_fwd(const at::Tensor& x, const at::Tensor& weight,
338
- const std::optional<at::Tensor>& bias_,
339
- const std::optional<at::Tensor>& conv_states,
340
- const std::optional<at::Tensor>& query_start_loc,
341
- const std::optional<at::Tensor>& cache_indices,
342
- const std::optional<at::Tensor>& has_initial_state,
343
- bool silu_activation, int64_t pad_slot_id);
344
-
329
using fptr_t = int64_t;
330
fptr_t init_custom_ar(const std::vector<int64_t>& fake_ipc_ptrs,
331
torch::Tensor& rank_data, int64_t rank,
csrc/torch_bindings.cpp
-22
@@ -594,28 +594,6 @@ TORCH_LIBRARY_EXPAND(TORCH_EXTENSION_NAME, ops) {
594
"int pad_slot_id) -> ()");
595
ops.impl("selective_scan_fwd", torch::kCUDA, &selective_scan_fwd);
596
597
- ops.def(
598
- "causal_conv1d_update(Tensor! x,"
599
- "Tensor! conv_state,"
600
- "Tensor! weight,"
601
- "Tensor? bias_,"
602
- "bool silu_activation,"
603
- "Tensor? cache_seqlens_,"
604
- "Tensor? conv_state_indices,"
605
- "int pad_slot_id) -> ()");
606
- ops.impl("causal_conv1d_update", torch::kCUDA, &causal_conv1d_update);
607
-
608
- ops.def(
609
- "causal_conv1d_fwd(Tensor! x, Tensor! weight,"
610
- "Tensor? bias_,"
611
- "Tensor!? conv_states,"
612
- "Tensor? query_start_loc,"
613
- "Tensor? cache_indices,"
614
- "Tensor? has_initial_state,"
615
- "bool silu_activation,"
616
- "int pad_slot_id) -> ()");
617
- ops.impl("causal_conv1d_fwd", torch::kCUDA, &causal_conv1d_fwd);
618
-
597
#ifndef USE_ROCM
598
// reorder weight for AllSpark Ampere W8A16 Fused Gemm kernel
599
ops.def(
tests/kernels/mamba/test_causal_conv1d.py
+44
-114
@@ -6,9 +6,8 @@ from typing import Optional
6
import pytest
7
import torch
8
import torch.nn.functional as F
9
+from einops import rearrange
10
10
-from tests.kernels.utils import opcheck
11
-from vllm import _custom_ops as ops # noqa: F401
11
from vllm.attention.backends.utils import PAD_SLOT_ID
12
from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
13
causal_conv1d_fn, causal_conv1d_update)
@@ -144,79 +143,6 @@ def causal_conv1d_opcheck_fn(x: torch.Tensor,
143
x = x.contiguous()
144
bias = bias.contiguous() if bias is not None else None
145
147
- opcheck(torch.ops._C.causal_conv1d_fwd,
148
- (x, weight, bias, conv_states, cu_seq_len, cache_indices,
149
- has_initial_state, activation in ["silu", "swish"], pad_slot_id))
150
-
151
-
152
-@pytest.mark.parametrize("itype", [torch.bfloat16, torch.float])
153
-@pytest.mark.parametrize("silu_activation", [True])
154
-@pytest.mark.parametrize("has_bias", [True])
155
-@pytest.mark.parametrize("has_initial_state", [True, False])
156
-@pytest.mark.parametrize("width", [4])
157
-@pytest.mark.parametrize(
158
- 'seqlen', [1, 8, 16, 32, 64, 128, 256, 512, 784, 1024, 1025, 2048, 4096])
159
-@pytest.mark.parametrize('dim', [64])
160
-@pytest.mark.parametrize('batch', [1])
161
-def test_causal_conv1d(batch, dim, seqlen, width, has_bias, silu_activation,
162
- has_initial_state, itype):
163
- device = "cuda"
164
- rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3)
165
- if itype == torch.bfloat16:
166
- rtol, atol = 1e-2, 5e-2
167
- # set seed
168
- current_platform.seed_everything(0)
169
- x = torch.randn(batch, dim, seqlen, device=device,
170
- dtype=itype).contiguous()
171
-
172
- weight = torch.randn(dim, width, device=device, dtype=itype)
173
- bias = torch.randn(dim, device=device, dtype=itype) if has_bias else None
174
- if has_initial_state:
175
- initial_states = torch.randn(batch,
176
- dim,
177
- width - 1,
178
- device=device,
179
- dtype=itype)
180
- has_initial_state_tensor = torch.ones(batch,
181
- dtype=torch.bool,
182
- device=x.device)
183
- else:
184
- initial_states = None
185
- has_initial_state_tensor = None
186
- x_ref = x.clone()
187
- weight_ref = weight.clone()
188
- bias_ref = bias.clone() if bias is not None else None
189
- initial_states_ref = initial_states.clone(
190
- ) if initial_states is not None else None
191
- activation = None if not silu_activation else "silu"
192
- out = causal_conv1d_fn(x,
193
- weight,
194
- bias,
195
- activation=activation,
196
- conv_states=initial_states,
197
- has_initial_state=has_initial_state_tensor)
198
- out_ref, final_states_ref = causal_conv1d_ref(
199
- x_ref,
200
- weight_ref,
201
- bias_ref,
202
- initial_states=initial_states_ref,
203
- return_final_states=True,
204
- activation=activation)
205
- if has_initial_state:
206
- assert initial_states is not None and final_states_ref is not None
207
- assert torch.allclose(initial_states,
208
- final_states_ref,
209
- rtol=rtol,
210
- atol=atol)
211
- assert torch.allclose(out, out_ref, rtol=rtol, atol=atol)
212
-
213
- causal_conv1d_opcheck_fn(x,
214
- weight,
215
- bias,
216
- activation=activation,
217
- conv_states=initial_states,
218
- has_initial_state=has_initial_state_tensor)
219
-
146
147
@pytest.mark.parametrize("itype", [torch.bfloat16])
148
@pytest.mark.parametrize("silu_activation", [False, True])
@@ -255,22 +181,19 @@ def test_causal_conv1d_update(dim, width, seqlen, has_bias, silu_activation,
181
assert torch.equal(conv_state, conv_state_ref)
182
assert torch.allclose(out, out_ref, rtol=rtol, atol=atol)
183
258
- opcheck(torch.ops._C.causal_conv1d_update,
259
- (x, conv_state, weight, bias, activation
260
- in ["silu", "swish"], None, None, PAD_SLOT_ID))
261
-
184
185
@pytest.mark.parametrize("itype",
186
[torch.float32, torch.float16, torch.bfloat16])
187
@pytest.mark.parametrize("silu_activation", [False, True])
188
@pytest.mark.parametrize("has_bias", [False, True])
267
-@pytest.mark.parametrize("seqlen", [1, 4, 5])
268
-@pytest.mark.parametrize("width", [2, 3, 4])
269
-@pytest.mark.parametrize("dim", [2048, 2048 + 16, 4096])
189
+@pytest.mark.parametrize("seqlen", [1, 3])
190
+@pytest.mark.parametrize("width", [3, 4])
191
+@pytest.mark.parametrize("dim", [2048 + 16, 4096])
192
# tests correctness in case subset of the sequences are padded
193
@pytest.mark.parametrize("with_padding", [True, False])
272
-def test_causal_conv1d_update_with_batch_gather(with_padding, dim, width,
273
- seqlen, has_bias,
194
+@pytest.mark.parametrize("batch_size", [3])
195
+def test_causal_conv1d_update_with_batch_gather(batch_size, with_padding, dim,
196
+ width, seqlen, has_bias,
197
silu_activation, itype):
198
device = "cuda"
199
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3)
@@ -280,12 +203,15 @@ def test_causal_conv1d_update_with_batch_gather(with_padding, dim, width,
203
# set seed
204
current_platform.seed_everything(0)
205
283
- batch_size = 3
206
padding = 5 if with_padding else 0
207
padded_batch_size = batch_size + padding
208
+ # total_entries = number of cache line
209
total_entries = 10 * batch_size
210
288
- x = torch.randn(padded_batch_size, dim, 1, device=device, dtype=itype)
211
+ # x will be (batch, dim, seqlen) with contiguous along dim-axis
212
+ x = torch.randn(padded_batch_size, seqlen, dim, device=device,
213
+ dtype=itype).transpose(1, 2)
214
+
215
x_ref = x.clone()
216
217
conv_state_indices = torch.randperm(total_entries)[:batch_size].to(
@@ -300,17 +226,22 @@ def test_causal_conv1d_update_with_batch_gather(with_padding, dim, width,
226
[PAD_SLOT_ID] * padding, dtype=torch.int32, device=device)
227
],
228
dim=0)
229
+
230
+ # conv_state will be (cache_lines, dim, state_len)
231
+ # with contiguous along dim-axis
232
conv_state = torch.randn(total_entries,
304
- dim,
233
width - 1,
234
+ dim,
235
device=device,
307
- dtype=itype)
236
+ dtype=itype).transpose(1, 2)
237
+
238
conv_state_for_padding_test = conv_state.clone()
239
240
weight = torch.randn(dim, width, device=device, dtype=itype)
241
bias = torch.randn(dim, device=device, dtype=itype) if has_bias else None
242
conv_state_ref = conv_state[conv_state_indices, :].detach().clone()
243
activation = None if not silu_activation else "silu"
244
+
245
out = causal_conv1d_update(x,
246
conv_state,
247
weight,
@@ -325,26 +256,21 @@ def test_causal_conv1d_update_with_batch_gather(with_padding, dim, width,
256
activation=activation)
257
258
assert torch.equal(conv_state[conv_state_indices, :], conv_state_ref)
328
- assert torch.allclose(out[:batch_size], out_ref, rtol=rtol, atol=atol)
259
assert torch.equal(conv_state[unused_states_bool],
260
conv_state_for_padding_test[unused_states_bool])
331
-
332
- opcheck(torch.ops._C.causal_conv1d_update,
333
- (x, conv_state, weight, bias, activation
334
- in ["silu", "swish"], None, padded_state_indices, PAD_SLOT_ID))
261
+ assert torch.allclose(out[:batch_size], out_ref, rtol=rtol, atol=atol)
262
263
264
@pytest.mark.parametrize("itype", [torch.bfloat16])
265
@pytest.mark.parametrize("silu_activation", [True])
266
@pytest.mark.parametrize("has_bias", [True])
267
@pytest.mark.parametrize("width", [4])
341
-@pytest.mark.parametrize(
342
- 'seqlen', [8, 16, 32, 64, 128, 256, 512, 784, 1024, 2048, 2049, 4096])
268
+@pytest.mark.parametrize('seqlen', [8, 30, 249, 2049, 4096])
269
@pytest.mark.parametrize('dim', [64, 4096])
344
-# tests correctness in case subset of the sequences are padded
270
@pytest.mark.parametrize('with_padding', [True, False])
346
-def test_causal_conv1d_varlen(with_padding, dim, seqlen, width, has_bias,
347
- silu_activation, itype):
271
+@pytest.mark.parametrize('batch', [4, 10])
272
+def test_causal_conv1d_varlen(batch, with_padding, dim, seqlen, width,
273
+ has_bias, silu_activation, itype):
274
device = "cuda"
275
torch.cuda.empty_cache()
276
rtol, atol = (3e-4, 1e-3) if itype == torch.float32 else (3e-3, 5e-3)
@@ -353,14 +279,13 @@ def test_causal_conv1d_varlen(with_padding, dim, seqlen, width, has_bias,
279
# set seed
280
current_platform.seed_everything(0)
281
seqlens = []
356
- batch_size = 4
357
- if seqlen < 10:
358
- batch_size = 1
282
+ batch_size = batch
283
padding = 3 if with_padding else 0
284
padded_batch_size = batch_size + padding
285
nsplits = padded_batch_size - 1
286
287
eos_pos = torch.randperm(seqlen - 1)[:nsplits].sort().values
288
+
289
seqlens.append(
290
torch.diff(
291
torch.cat(
@@ -373,19 +298,22 @@ def test_causal_conv1d_varlen(with_padding, dim, seqlen, width, has_bias,
298
cumsum = torch.cumsum(torch.tensor(seqlens[0]), dim=0).to(torch.int32)
299
cumsum = torch.concat([torch.tensor([0], dtype=torch.int32), cumsum],
300
dim=0)
376
- x = torch.randn(1, 4096 + dim + 64, seqlen, device=device,
377
- dtype=itype)[:, 4096:4096 + dim, :]
301
+ x = rearrange(
302
+ torch.randn(1, seqlen, 4096 + dim + 64, device=device, dtype=itype),
303
+ "b s d -> b d s")[:, 4096:4096 + dim, :]
304
+
305
weight = torch.randn(dim, width, device=device, dtype=itype)
306
+
307
bias = torch.randn(dim, device=device, dtype=itype) if has_bias else None
308
x_ref = x.clone()
309
weight_ref = weight.clone()
310
bias_ref = bias.clone() if bias is not None else None
311
activation = None if not silu_activation else "silu"
312
final_states = torch.randn(total_entries,
385
- dim,
313
width - 1,
314
+ dim,
315
device=x.device,
388
- dtype=x.dtype)
316
+ dtype=x.dtype).transpose(1, 2)
317
final_states_ref = final_states.clone()
318
has_initial_states = torch.randint(0,
319
2, (cumsum.shape[0] - 1, ),
@@ -400,10 +328,16 @@ def test_causal_conv1d_varlen(with_padding, dim, seqlen, width, has_bias,
328
[PAD_SLOT_ID] * padding, dtype=torch.int32, device=device),
329
],
330
dim=-1)
331
+ out = causal_conv1d_fn(x.squeeze(0),
332
+ weight,
333
+ bias=bias,
334
+ conv_states=final_states,
335
+ query_start_loc=cumsum.cuda(),
336
+ cache_indices=padded_state_indices,
337
+ has_initial_state=has_initial_states,
338
+ activation=activation,
339
+ pad_slot_id=PAD_SLOT_ID)
340
404
- out = causal_conv1d_fn(x.squeeze(0), weight, bias, cumsum.cuda(),
405
- padded_state_indices, has_initial_states,
406
- final_states, activation, PAD_SLOT_ID)
341
out_ref = []
342
out_ref_b = []
343
@@ -426,13 +360,9 @@ def test_causal_conv1d_varlen(with_padding, dim, seqlen, width, has_bias,
360
out_ref.append(torch.cat([t[0] for t in out_ref_b], dim=2))
361
out_ref_tensor = torch.cat(out_ref, dim=0)
362
429
- unpadded_out = out[:, :out_ref_tensor.shape[-1]]
430
- assert torch.allclose(unpadded_out, out_ref_tensor, rtol=rtol, atol=atol)
363
assert torch.allclose(final_states[state_indices],
364
final_states_ref[state_indices],
365
rtol=rtol,
366
atol=atol)
435
-
436
- causal_conv1d_opcheck_fn(x.squeeze(0), weight, bias, cumsum.cuda(),
437
- padded_state_indices, has_initial_states,
438
- final_states, activation)
367
+ unpadded_out = out[:, :out_ref_tensor.shape[-1]]
368
+ assert torch.allclose(unpadded_out, out_ref_tensor, rtol=rtol, atol=atol)
tests/kernels/mamba/test_mamba_ssm_ssd.py
+2
-2
@@ -6,11 +6,11 @@ import torch
6
import torch.nn.functional as F
7
from einops import rearrange, repeat
8
9
-from vllm.model_executor.layers.mamba.mamba2_metadata import (
10
- _query_start_loc_to_chunk_indices_offsets)
9
from vllm.model_executor.layers.mamba.ops.ssd_combined import (
10
mamba_chunk_scan_combined)
11
from vllm.platforms import current_platform
12
+from vllm.v1.attention.backends.mamba_attn import (
13
+ _query_start_loc_to_chunk_indices_offsets)
14
15
# Added by the IBM Team, 2024
16
vllm/_custom_ops.py
+5
-29
@@ -963,17 +963,17 @@ def cutlass_fp4_moe_mm(a_tensors: torch.Tensor, b_tensors: torch.Tensor,
963
expert_offsets: torch.Tensor, sf_offsets: torch.Tensor,
964
out_dtype: torch.dtype, device: torch.device):
965
"""
966
- An FP4 Blockscaled Group Gemm that takes in a_tensors, b_tensors and runs
966
+ An FP4 Blockscaled Group Gemm that takes in a_tensors, b_tensors and runs
967
the gemms for each combination based on the specified problem sizes.
968
969
This is used as the MoE gemm during NVFP4 Quantized FusedMoE forward.
970
- a/b_tensors: the NVFP4 a_ptrs and b_ptrs tensors which are quantized
971
input and expert weights.
972
- a_/b_scales: The blockscales in FP8-E4M3 precision
973
- - expert_offsets/sf_offsets: Indices that mark at which token index
974
- each expert begins its computation. The number of tokens
975
- computed with expert E is expert_offsets[E + 1] -
976
- expert_offsets[E] And the sf_size per expert is
973
+ - expert_offsets/sf_offsets: Indices that mark at which token index
974
+ each expert begins its computation. The number of tokens
975
+ computed with expert E is expert_offsets[E + 1] -
976
+ expert_offsets[E] And the sf_size per expert is
977
sf_offset[E+1] - sf_offset[E]
978
- problem_sizes: MxNxK sizes of each expert's multiplication in two grouped
979
MMs used in the fused MoE operation.
@@ -1464,30 +1464,6 @@ def ggml_moe_get_block_size(quant_type: int) -> int:
1464
1465
1466
# mamba
1467
-def causal_conv1d_fwd(x: torch.Tensor, weight: torch.Tensor,
1468
- bias_: Optional[torch.Tensor],
1469
- conv_states: Optional[torch.Tensor],
1470
- query_start_loc: Optional[torch.Tensor],
1471
- cache_indices: Optional[torch.Tensor],
1472
- has_initial_state: Optional[torch.Tensor],
1473
- silu_activation: bool, pad_slot_id: int):
1474
- torch.ops._C.causal_conv1d_fwd(x, weight, bias_, conv_states,
1475
- query_start_loc, cache_indices,
1476
- has_initial_state, silu_activation,
1477
- pad_slot_id)
1478
-
1479
-
1480
-def causal_conv1d_update(x: torch.Tensor, conv_state: torch.Tensor,
1481
- weight: torch.Tensor, bias_: Optional[torch.Tensor],
1482
- silu_activation: bool,
1483
- cache_seqlens: Optional[torch.Tensor],
1484
- conv_state_indices: Optional[torch.Tensor],
1485
- pad_slot_id: int):
1486
- torch.ops._C.causal_conv1d_update(x, conv_state, weight, bias_,
1487
- silu_activation, cache_seqlens,
1488
- conv_state_indices, pad_slot_id)
1489
-
1490
-
1467
def selective_scan_fwd(u: torch.Tensor, delta: torch.Tensor, A: torch.Tensor,
1468
B: torch.Tensor, C: torch.Tensor,
1469
D_: Optional[torch.Tensor], z_: Optional[torch.Tensor],
vllm/model_executor/layers/mamba/mamba2_metadata.py
+103
-42
@@ -1,14 +1,18 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
-import math
3
from dataclasses import dataclass
4
+from typing import Optional, Union
5
6
+import numpy as np
7
import torch
8
9
from vllm.attention.backends.abstract import AttentionMetadata
10
from vllm.attention.backends.placeholder_attn import (
11
PlaceholderAttentionMetadata)
12
+from vllm.attention.backends.utils import PAD_SLOT_ID
13
from vllm.platforms import current_platform
14
+from vllm.v1.attention.backends.mamba_attn import (
15
+ Mamba2AttentionMetadata, _query_start_loc_to_chunk_indices_offsets)
16
17
18
@dataclass
@@ -21,6 +25,29 @@ class Mamba2Metadata:
25
seq_idx: torch.Tensor
26
chunk_indices: torch.Tensor
27
chunk_offsets: torch.Tensor
28
+ """
29
+ With continuous batching layout of `x` in vLLM, to enable a Triton program
30
+ to handle a request in parallel, two supporting tensors are used
31
+ (batch_ptr, token_chunk_offset_ptr)
32
+ BLOCK_M = the # tokens to be handled by a Triton program
33
+ (can be customized for different hardware)
34
+
35
+ nums_dict:
36
+ tracks the data associated with a given value of BLOCK_M
37
+ BLOCK_M = #tokens handled by a Triton program
38
+ cu_seqlen: total tokens per batch
39
+ (used as flag to update other data at each new input)
40
+ batch_ptr: tracks batch-id handled by the Triton program
41
+ token_chunk_offset_ptr: tracks token group_idx handled by the Triton program
42
+ (Triton implementation of causal_conv1d handles parallelism in 3-axes
43
+ - feature-axis
44
+ - batch-axis
45
+ - sequence-axis)
46
+ """
47
+ nums_dict: Optional[dict] = None
48
+ cu_seqlen: Optional[int] = None
49
+ batch_ptr: Optional[torch.tensor] = None
50
+ token_chunk_offset_ptr: Optional[torch.tensor] = None
51
52
53
def get_platform_metadata_classes() -> tuple[type[AttentionMetadata], ...]:
@@ -38,45 +65,10 @@ def get_platform_metadata_classes() -> tuple[type[AttentionMetadata], ...]:
65
f"Unsupported platform for Mamba2: {current_platform.device_type}")
66
67
41
-def _query_start_loc_to_chunk_indices_offsets(query_start_loc: torch.Tensor,
42
- chunk_size: int,
43
- total_seqlens: int):
44
-
45
- cu_seqlens = query_start_loc[1:] # remove prepended 0
46
-
47
- # outputs will have length expansion of chunks that do not divide
48
- # chunk_size
49
- N = math.ceil(total_seqlens / chunk_size) + (cu_seqlens[:-1] % chunk_size
50
- > 0).sum()
51
- chunk_indices = torch.arange(N,
52
- dtype=torch.int,
53
- device=query_start_loc.device)
54
- chunk_offsets = torch.zeros((N, ),
55
- dtype=torch.int,
56
- device=query_start_loc.device)
57
-
58
- p = 0 # num of insertions
59
- for s, e in zip(cu_seqlens[:-1], cu_seqlens[1:]):
60
-
61
- # if does not divide chunk_size, then there is one chunk insertion
62
- p += (s % chunk_size > 0)
63
-
64
- # get the dimensions
65
- # - the + 1 for _e is to shift the boundary by one chunk
66
- # - this shifting is not needed if chunk_size divides e
67
- _s, _e = s // chunk_size + p, e // chunk_size + p + (e % chunk_size
68
- > 0)
69
-
70
- # adjust inidces and offsets
71
- chunk_indices[_s:_e] -= p
72
- chunk_offsets[_s] = s % chunk_size
73
-
74
- return chunk_indices, chunk_offsets
75
-
76
-
68
def prepare_mamba2_metadata(
69
chunk_size: int,
70
attn_metadata: AttentionMetadata,
71
+ mamba2_metadata=None,
72
) -> Mamba2Metadata:
73
74
# compute number of prefill and decode requests
@@ -96,12 +88,12 @@ def prepare_mamba2_metadata(
88
attn_metadata_instances = get_platform_metadata_classes()
89
if (isinstance(attn_metadata, attn_metadata_instances)
90
and attn_metadata.context_lens_tensor is not None):
99
- has_initial_states = \
100
- attn_metadata.context_lens_tensor[:num_prefills] > 0 #[batch,]
101
- # precompute flag to avoid device syncs in mamba2 layer forwards
91
+ # precompute flag to avoid device syncs later in mamba2 layer
92
+ # forwards
93
# prep is only needed for mamba2 ssd prefill processing
103
- prep_initial_states = torch.any(has_initial_states).item()
104
-
94
+ has_initial_states = attn_metadata.context_lens_tensor > 0
95
+ prep_initial_states = torch.any(
96
+ has_initial_states[:num_prefills]).item()
97
query_start_loc = attn_metadata.query_start_loc[:num_prefills + 1]
98
seq_idx = torch.repeat_interleave(torch.arange(
99
num_prefills, dtype=torch.int32, device=query_start_loc.device),
@@ -117,9 +109,78 @@ def prepare_mamba2_metadata(
109
_query_start_loc_to_chunk_indices_offsets(
110
query_start_loc, chunk_size, num_prefill_tokens)
111
112
+ if mamba2_metadata is not None:
113
+ mamba2_metadata.has_initial_states = has_initial_states
114
+ mamba2_metadata.prep_initial_states = prep_initial_states
115
+ mamba2_metadata.chunk_size = chunk_size
116
+ mamba2_metadata.seq_idx = seq_idx
117
+ mamba2_metadata.chunk_indices = chunk_indices
118
+ mamba2_metadata.chunk_offsets = chunk_offsets
119
+ # We use 1 reset flag:
120
+ # * mamba2_metadata.cu_seqlen is None
121
+ # update config specific to (each input)
122
+ # (become available at first layer, e.g. conv_weights)
123
+ mamba2_metadata.cu_seqlen = None # suppose to be updated at each input
124
+
125
+ return mamba2_metadata
126
return Mamba2Metadata(has_initial_states=has_initial_states,
127
prep_initial_states=prep_initial_states,
128
chunk_size=chunk_size,
129
seq_idx=seq_idx,
130
chunk_indices=chunk_indices,
131
chunk_offsets=chunk_offsets)
132
+
133
+
134
+def update_metadata(x: torch.Tensor, query_start_loc: torch.Tensor,
135
+ mamba2_metadata: Union[Mamba2Metadata,
136
+ Mamba2AttentionMetadata]):
137
+ """
138
+ this is triggered upon handling a new input at the first layer
139
+ """
140
+ dim, cu_seqlen = x.shape
141
+ mamba2_metadata.cu_seqlen = cu_seqlen
142
+ seqlens = np.diff(query_start_loc.to('cpu'))
143
+ nums_dict = {} # type: ignore
144
+ for BLOCK_M in [8]: # cover all BLOCK_M values
145
+ nums = -(-seqlens // BLOCK_M)
146
+ nums_dict[BLOCK_M] = {}
147
+ nums_dict[BLOCK_M]['nums'] = nums
148
+ nums_dict[BLOCK_M]['tot'] = nums.sum().item()
149
+ mlist = torch.from_numpy(np.repeat(np.arange(len(nums)), nums))
150
+ nums_dict[BLOCK_M]['mlist'] = mlist
151
+ mlist_len = len(nums_dict[BLOCK_M]['mlist'])
152
+ nums_dict[BLOCK_M]['mlist_len'] = mlist_len
153
+ MAX_NUM_PROGRAMS = max(1024, mlist_len) * 2
154
+ offsetlist = [] # type: ignore
155
+ for idx, num in enumerate(nums):
156
+ offsetlist.extend(range(num))
157
+ offsetlist = torch.tensor(offsetlist, dtype=torch.int32)
158
+ nums_dict[BLOCK_M]['offsetlist'] = offsetlist
159
+
160
+ if mamba2_metadata.batch_ptr is None:
161
+ # Update default value after class definition
162
+ #mamba2_metadata.MAX_NUM_PROGRAMS *= 2
163
+ mamba2_metadata.batch_ptr = torch.full((MAX_NUM_PROGRAMS, ),
164
+ PAD_SLOT_ID,
165
+ dtype=torch.int32,
166
+ device='cuda')
167
+ mamba2_metadata.token_chunk_offset_ptr = torch.full(
168
+ (MAX_NUM_PROGRAMS, ),
169
+ PAD_SLOT_ID,
170
+ dtype=torch.int32,
171
+ device='cuda')
172
+ else:
173
+ if mamba2_metadata.batch_ptr.nelement() < MAX_NUM_PROGRAMS:
174
+ mamba2_metadata.batch_ptr.resize_(MAX_NUM_PROGRAMS).fill_(
175
+ PAD_SLOT_ID)
176
+ mamba2_metadata.token_chunk_offset_ptr.resize_( # type: ignore
177
+ MAX_NUM_PROGRAMS).fill_(PAD_SLOT_ID)
178
+
179
+ mamba2_metadata.batch_ptr[0:mlist_len].copy_(mlist)
180
+ mamba2_metadata.token_chunk_offset_ptr[ # type: ignore
181
+ 0:mlist_len].copy_(offsetlist)
182
+ nums_dict[BLOCK_M]['batch_ptr'] = mamba2_metadata.batch_ptr
183
+ nums_dict[BLOCK_M]['token_chunk_offset_ptr'] = (
184
+ mamba2_metadata.token_chunk_offset_ptr) # type: ignore
185
+ mamba2_metadata.nums_dict = nums_dict
186
+ return mamba2_metadata
vllm/model_executor/layers/mamba/mamba_mixer.py
+1
-1
@@ -159,7 +159,7 @@ class MambaMixer(CustomOp):
159
hidden_states = causal_conv1d_fn(
160
hidden_states,
161
conv_weights,
162
- self.conv1d.bias,
162
+ bias=self.conv1d.bias,
163
activation=self.activation,
164
conv_states=mamba_cache_params.conv_state,
165
has_initial_state=attn_metadata.context_lens_tensor > 0,
vllm/model_executor/layers/mamba/mamba_mixer2.py
+17
-9
@@ -17,7 +17,8 @@ from vllm.forward_context import get_forward_context
17
from vllm.model_executor.custom_op import CustomOp
18
from vllm.model_executor.layers.linear import (ColumnParallelLinear,
19
RowParallelLinear)
20
-from vllm.model_executor.layers.mamba.mamba2_metadata import Mamba2Metadata
20
+from vllm.model_executor.layers.mamba.mamba2_metadata import (Mamba2Metadata,
21
+ update_metadata)
22
from vllm.model_executor.layers.mamba.ops.causal_conv1d import (
23
causal_conv1d_fn, causal_conv1d_update)
24
from vllm.model_executor.layers.mamba.ops.mamba_ssm import (
@@ -161,9 +162,9 @@ def mamba_v2_sharded_weight_loader(
162
tp_size: int,
163
tp_rank: int,
164
) -> LoaderFunction:
164
- """Create a weight loader for mamba v2. This ensures that the projections
165
- are correctly sharded so that they can be split into x, B, C. It also
166
- ensures that all the groups corresponding to a head shard is placed
165
+ """Create a weight loader for mamba v2. This ensures that the projections
166
+ are correctly sharded so that they can be split into x, B, C. It also
167
+ ensures that all the groups corresponding to a head shard is placed
168
together with it.
169
"""
170
@@ -458,9 +459,11 @@ class MambaMixer2(CustomOp):
459
if attn_metadata is not None:
460
assert isinstance(attn_metadata, dict)
461
attn_metadata = attn_metadata[self.prefix]
462
+ mamba2_metadata = attn_metadata
463
assert isinstance(attn_metadata, Mamba2AttentionMetadata)
464
self_kv_cache = self.kv_cache[forward_context.virtual_engine]
463
- conv_state = self_kv_cache[0]
465
+ # conv_state = (..., dim, width-1) yet contiguous along 'dim'
466
+ conv_state = self_kv_cache[0].transpose(-1, -2)
467
ssm_state = self_kv_cache[1]
468
state_indices_tensor = attn_metadata.state_indices_tensor
469
has_initial_states_p = attn_metadata.has_initial_states
@@ -531,6 +534,7 @@ class MambaMixer2(CustomOp):
534
# NOTE: V0 put prefill before decode, v1 puts decode before prefill
535
# Separate prefill and decode by splitting varlen input
536
# Split along token dimension
537
+ # NOTE: V0 put prefill before decode, v1 puts decode before prefill
538
if envs.VLLM_USE_V1:
539
hidden_states_B_C_d, hidden_states_B_C_p = torch.split(
540
hidden_states_B_C,
@@ -579,8 +583,13 @@ class MambaMixer2(CustomOp):
583
# 2. Convolution sequence transformation
584
# - "cache_indices" updates the conv_state cache in positions
585
# pointed to by "state_indices_tensor"
586
+ x = hidden_states_B_C_p.transpose(
587
+ 0, 1) # this is the form that causal-conv see
588
+ if mamba2_metadata.cu_seqlen is None:
589
+ mamba2_metadata = update_metadata(
590
+ x, attn_metadata.query_start_loc, mamba2_metadata)
591
hidden_states_B_C_p = causal_conv1d_fn(
583
- hidden_states_B_C_p.transpose(0, 1),
592
+ x,
593
conv_weights,
594
self.conv1d.bias,
595
activation=self.activation,
@@ -590,8 +599,6 @@ class MambaMixer2(CustomOp):
599
query_start_loc=query_start_loc_p).transpose(
600
0, 1)[:num_prefill_tokens]
601
593
- # TODO: Why is this needed?
594
- hidden_states_B_C_p = hidden_states_B_C_p.contiguous()
602
hidden_states_p, B_p, C_p = split_hidden_states_B_C_fn(
603
hidden_states_B_C_p)
604
@@ -715,9 +722,10 @@ class MambaMixer2(CustomOp):
722
# - heads and n_groups are TP-ed
723
conv_dim = (self.intermediate_size +
724
2 * n_groups * self.ssm_state_size)
725
+ # contiguous along 'dim' axis
726
conv_state_shape = (
719
- divide(conv_dim, world_size),
727
self.conv_kernel_size - 1,
728
+ divide(conv_dim, world_size),
729
)
730
731
# These are not TP-ed as they depend on A, dt_bias, D
vllm/model_executor/layers/mamba/ops/causal_conv1d.py
+902
-61
@@ -4,102 +4,943 @@
4
# Copyright (c) 2024, Tri Dao.
5
# Adapted from https://github.com/Dao-AILab/causal-conv1d/blob/main/causal_conv1d/causal_conv1d_interface.py
6
7
-from typing import Optional
7
+from typing import Optional, Union
8
9
+import numpy as np
10
import torch
11
+import triton
12
+import triton.language as tl
13
11
-from vllm import _custom_ops as ops
14
from vllm.attention.backends.utils import PAD_SLOT_ID
15
16
15
-def causal_conv1d_fn(x: torch.Tensor,
16
- weight: torch.Tensor,
17
- bias: Optional[torch.Tensor] = None,
18
- query_start_loc: Optional[torch.Tensor] = None,
19
- cache_indices: Optional[torch.Tensor] = None,
20
- has_initial_state: Optional[torch.Tensor] = None,
21
- conv_states: Optional[torch.Tensor] = None,
22
- activation: Optional[str] = "silu",
23
- pad_slot_id: int = PAD_SLOT_ID):
24
- """
25
- x: (batch, dim, seqlen) or (dim,cu_seq_len) for varlen
17
+@triton.jit()
18
+def _causal_conv1d_fwd_kernel( # continuous batching
19
+ # Pointers to matrices
20
+ x_ptr, # (dim, cu_seqlen) holding `batch` of actual sequences + padded sequences
21
+ w_ptr, # (dim, width)
22
+ bias_ptr,
23
+ initial_states_ptr, # conv_states_ptr
24
+ cache_indices_ptr, # conv_state_indices_ptr
25
+ has_initial_states_ptr,
26
+ query_start_loc_ptr,
27
+ batch_ptr,
28
+ token_chunk_offset_ptr,
29
+ o_ptr, # (dim, seqlen) - actually pointing to x_ptr
30
+ # Matrix dimensions
31
+ batch: tl.int32, # actually padded_batch
32
+ dim: tl.constexpr,
33
+ seqlen: tl.int32, # cu_seqlen
34
+ num_cache_lines: tl.constexpr, # added to support vLLM larger cache lines
35
+ # Strides
36
+ stride_x_seq: tl.constexpr, # stride to get to next sequence,
37
+ stride_x_dim: tl.constexpr, # stride to get to next feature-value,
38
+ stride_x_token: tl.
39
+ constexpr, # stride to get to next token (same feature-index, same sequence-index)
40
+ stride_w_dim: tl.constexpr, # stride to get to next dim-axis value
41
+ stride_w_width: tl.constexpr, # stride to get to next width-axis value
42
+ stride_istate_seq: tl.constexpr,
43
+ stride_istate_dim: tl.constexpr,
44
+ stride_istate_token: tl.constexpr,
45
+ stride_o_seq: tl.constexpr,
46
+ stride_o_dim: tl.constexpr,
47
+ stride_o_token: tl.constexpr,
48
+ # others
49
+ pad_slot_id: tl.constexpr,
50
+ # Meta-parameters
51
+ HAS_BIAS: tl.constexpr,
52
+ KERNEL_WIDTH: tl.constexpr,
53
+ SILU_ACTIVATION: tl.constexpr,
54
+ HAS_INITIAL_STATES: tl.constexpr,
55
+ HAS_CACHE: tl.constexpr,
56
+ IS_CONTINUOUS_BATCHING: tl.constexpr,
57
+ USE_PAD_SLOT: tl.constexpr,
58
+ NP2_STATELEN: tl.constexpr,
59
+ DECODE_SEQLEN: tl.constexpr,
60
+ BLOCK_M: tl.constexpr,
61
+ BLOCK_N: tl.constexpr,
62
+):
63
+ conv_states_ptr = initial_states_ptr
64
+ conv_state_indices_ptr = cache_indices_ptr
65
+ stride_conv_state_seq = stride_istate_seq
66
+ stride_conv_state_dim = stride_istate_dim
67
+ stride_conv_state_tok = stride_istate_token
68
+ state_len = KERNEL_WIDTH - 1 # can be passed via argument if it's not the same as this value
69
+
70
+ # one program handles one chunk in a single sequence
71
+ # rather than mixing sequences - to make updating initial_states across sequences efficiently
72
+
73
+ # single-sequence id
74
+ idx_seq = tl.load(batch_ptr + tl.program_id(0))
75
+ chunk_offset = tl.load(token_chunk_offset_ptr + tl.program_id(0))
76
+
77
+ # BLOCK_N elements along the feature-dimension (channel)
78
+ idx_feats = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N)
79
+
80
+ if idx_seq == pad_slot_id:
81
+ return
82
+
83
+ sequence_start_index = tl.load(query_start_loc_ptr + idx_seq)
84
+ sequence_end_index = tl.load(query_start_loc_ptr + idx_seq + 1)
85
+ # find the actual sequence length
86
+ seqlen = sequence_end_index - sequence_start_index
87
+
88
+ token_offset = BLOCK_M * chunk_offset
89
+ segment_len = min(BLOCK_M, seqlen - token_offset)
90
+
91
+ # base of the sequence
92
+ x_base = x_ptr + sequence_start_index * stride_x_token + idx_feats * stride_x_dim # [BLOCK_N,]
93
+
94
+ if IS_CONTINUOUS_BATCHING:
95
+ # cache_idx
96
+ conv_state_batch_coord = tl.load(conv_state_indices_ptr + idx_seq)
97
+ else:
98
+ # cache_idx
99
+ conv_state_batch_coord = idx_seq
100
+ if USE_PAD_SLOT: # noqa
101
+ if conv_state_batch_coord == pad_slot_id:
102
+ # not processing as this is not the actual sequence
103
+ return
104
+ conv_states_base = (conv_states_ptr +
105
+ (conv_state_batch_coord * stride_conv_state_seq) +
106
+ (idx_feats * stride_conv_state_dim)) # [BLOCK_N,]
107
+
108
+ w_base = w_ptr + (idx_feats * stride_w_dim) # [BLOCK_N,]
109
+
110
+ # Does 2 things:
111
+ # 1. READ prior-block init-state data - [done by every Triton programs]
112
+ # 2. update conv_state with new data [only by the Triton program handles chunk_offset=0]
113
+ if chunk_offset == 0:
114
+ # read from conv_states
115
+ load_init_state = False
116
+ if HAS_INITIAL_STATES: # the new HAS_INITIAL_STATES
117
+ load_init_state = tl.load(has_initial_states_ptr + idx_seq).to(
118
+ tl.int1)
119
+ if load_init_state:
120
+ # load from conv_states
121
+ prior_tokens = conv_states_base + (state_len -
122
+ 1) * stride_conv_state_tok
123
+ mask_w = idx_feats < dim
124
+ if KERNEL_WIDTH == 2:
125
+ conv_states_ptrs = prior_tokens # [BLOCK_N]
126
+ col0 = tl.load(conv_states_ptrs, mask_w, 0.0)
127
+ if KERNEL_WIDTH == 3:
128
+ conv_states_ptrs = prior_tokens # [BLOCK_N]
129
+ col1 = tl.load(conv_states_ptrs, mask_w, 0.0)
130
+ conv_states_ptrs = prior_tokens - 1 * stride_conv_state_tok # [BLOCK_N]
131
+ col0 = tl.load(conv_states_ptrs, mask_w, 0.0)
132
+ if KERNEL_WIDTH == 4:
133
+ conv_states_ptrs = prior_tokens # [BLOCK_N]
134
+ col2 = tl.load(conv_states_ptrs, mask_w, 0.0)
135
+ conv_states_ptrs = prior_tokens - 1 * stride_conv_state_tok # [BLOCK_N]
136
+ col1 = tl.load(conv_states_ptrs, mask_w, 0.0)
137
+ conv_states_ptrs = prior_tokens - 2 * stride_conv_state_tok # [BLOCK_N]
138
+ col0 = tl.load(conv_states_ptrs, mask_w, 0.0)
139
+ if KERNEL_WIDTH == 5:
140
+ conv_states_ptrs = prior_tokens # [BLOCK_N]
141
+ col3 = tl.load(conv_states_ptrs, mask_w, 0.0)
142
+ conv_states_ptrs = prior_tokens - 1 * stride_conv_state_tok # [BLOCK_N]
143
+ col2 = tl.load(conv_states_ptrs, mask_w, 0.0)
144
+ conv_states_ptrs = prior_tokens - 2 * stride_conv_state_tok # [BLOCK_N]
145
+ col1 = tl.load(conv_states_ptrs, mask_w, 0.0)
146
+ conv_states_ptrs = prior_tokens - 3 * stride_conv_state_tok # [BLOCK_N]
147
+ col0 = tl.load(conv_states_ptrs, mask_w, 0.0)
148
+ else:
149
+ # prior-tokens are zeros
150
+ if KERNEL_WIDTH >= 2: # STRATEGY1
151
+ # first chunk and does not have prior-token, so just set to 0
152
+ col0 = tl.zeros((BLOCK_N, ), dtype=x_ptr.dtype.element_ty)
153
+ if KERNEL_WIDTH >= 3: # STRATEGY1
154
+ col1 = tl.zeros((BLOCK_N, ), dtype=x_ptr.dtype.element_ty)
155
+ if KERNEL_WIDTH >= 4: # STRATEGY1
156
+ col2 = tl.zeros((BLOCK_N, ), dtype=x_ptr.dtype.element_ty)
157
+ if KERNEL_WIDTH >= 5: # STRATEGY1
158
+ col3 = tl.zeros((BLOCK_N, ), dtype=x_ptr.dtype.element_ty)
159
+
160
+ # STEP 2:
161
+ # here prepare data for updating conv_state
162
+ if state_len <= seqlen: # SMALL_CACHE=True (only move part of 'x' into conv_state cache)
163
+ # just read from 'x'
164
+ # copy 'x' data to conv_state
165
+ # load only 'x' data (and set 0 before 'x' if seqlen < state_len)
166
+ idx_tokens_last = (seqlen - state_len) + tl.arange(
167
+ 0, NP2_STATELEN) # [BLOCK_M]
168
+ x_ptrs = x_ptr + (
169
+ (sequence_start_index + idx_tokens_last) *
170
+ stride_x_token)[:, None] + (
171
+ idx_feats * stride_x_dim)[None, :] # [BLOCK_M,BLOCK_N,]
172
+ mask_x = ((idx_tokens_last >= 0)[:, None] &
173
+ (idx_tokens_last < seqlen)[:, None] &
174
+ (idx_feats < dim)[None, :]
175
+ ) # token-index # token-index # feature-index
176
+ loaded_x = tl.load(x_ptrs, mask_x, 0.0)
177
+ new_conv_state = tl.load(x_ptrs, mask_x, 0.0)
178
+ idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M]
179
+ conv_states_ptrs_target = conv_states_base[None, :] + (
180
+ idx_tokens_conv * stride_conv_state_tok)[:, None]
181
+
182
+ mask = (idx_tokens_conv < state_len)[:, None] & (idx_feats
183
+ < dim)[None, :]
184
+ tl.debug_barrier() # NOTE: use this due to bug in Triton compiler
185
+ tl.store(conv_states_ptrs_target, new_conv_state, mask)
186
+
187
+ else:
188
+ if load_init_state:
189
+ # update conv_state by shifting left, i.e. take last few cols from conv_state + cols from 'x'
190
+ idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M]
191
+
192
+ conv_states_ptrs_source = (
193
+ conv_states_ptr +
194
+ (conv_state_batch_coord * stride_conv_state_seq) +
195
+ (idx_feats * stride_conv_state_dim)[None, :] +
196
+ ((idx_tokens_conv + seqlen) * stride_conv_state_tok)[:,
197
+ None]
198
+ ) # [BLOCK_M, BLOCK_N]
199
+ mask = ((conv_state_batch_coord < num_cache_lines)
200
+ & ((idx_tokens_conv + seqlen) < state_len)[:, None]
201
+ & (idx_feats < dim)[None, :])
202
+ conv_state = tl.load(conv_states_ptrs_source, mask, other=0.0)
203
+
204
+ VAL = state_len - seqlen
205
+
206
+ x_ptrs = x_base[None, :] + (
207
+ (idx_tokens_conv - VAL) *
208
+ stride_x_token)[:, None] # [BLOCK_M, BLOCK_N]
209
+
210
+ mask_x = ((idx_tokens_conv - VAL >= 0)[:, None] &
211
+ (idx_tokens_conv - VAL < seqlen)[:, None] &
212
+ (idx_feats < dim)[None, :]
213
+ ) # token-index # token-index # feature-index
214
+ loaded_x = tl.load(x_ptrs, mask_x, 0.0)
215
+
216
+ tl.debug_barrier(
217
+ ) # need this due to the bug in tl.where not enforcing this when data is the result of another tl.load
218
+ new_conv_state = tl.where(
219
+ mask, conv_state, loaded_x
220
+ ) # BUG in 'tl.where' which requires a barrier before this
221
+ conv_states_ptrs_target = conv_states_base + (
222
+ idx_tokens_conv *
223
+ stride_conv_state_tok)[:, None] # [BLOCK_M, BLOCK_N]
224
+ mask = (idx_tokens_conv
225
+ < state_len)[:, None] & (idx_feats < dim)[None, :]
226
+ tl.store(conv_states_ptrs_target, new_conv_state, mask)
227
+ else: # load_init_state == False
228
+ # update conv_state by shifting left, BUT
229
+ # set cols prior to 'x' as zeros + cols from 'x'
230
+ idx_tokens_conv = tl.arange(0, NP2_STATELEN) # [BLOCK_M]
231
+
232
+ VAL = state_len - seqlen
233
+
234
+ x_ptrs = x_base[None, :] + (
235
+ (idx_tokens_conv - VAL) *
236
+ stride_x_token)[:, None] # [BLOCK_M, BLOCK_N]
237
+
238
+ mask_x = ((idx_tokens_conv - VAL >= 0)[:, None] &
239
+ (idx_tokens_conv - VAL < seqlen)[:, None] &
240
+ (idx_feats < dim)[None, :]
241
+ ) # token-index # token-index # feature-index
242
+ new_conv_state = tl.load(x_ptrs, mask_x, 0.0)
243
+
244
+ conv_states_ptrs_target = conv_states_base + (
245
+ idx_tokens_conv *
246
+ stride_conv_state_tok)[:, None] # [BLOCK_M, BLOCK_N]
247
+ mask = (idx_tokens_conv
248
+ < state_len)[:, None] & (idx_feats < dim)[None, :]
249
+ tl.store(conv_states_ptrs_target, new_conv_state, mask)
250
+
251
+ else: # chunk_offset > 0
252
+ # read prior-token data from `x`
253
+ load_init_state = True
254
+ prior_tokens = x_base + (token_offset - 1) * stride_x_token
255
+ mask_w = idx_feats < dim
256
+ if KERNEL_WIDTH == 2:
257
+ conv_states_ptrs = prior_tokens # [BLOCK_N]
258
+ col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier='.ca')
259
+ if KERNEL_WIDTH == 3:
260
+ conv_states_ptrs = prior_tokens # [BLOCK_N]
261
+ col1 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier='.ca')
262
+ conv_states_ptrs = prior_tokens - 1 * stride_x_token # [BLOCK_N]
263
+ col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier='.ca')
264
+ if KERNEL_WIDTH == 4:
265
+ conv_states_ptrs = prior_tokens # [BLOCK_N]
266
+ col2 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier='.ca')
267
+ conv_states_ptrs = prior_tokens - 1 * stride_x_token # [BLOCK_N]
268
+ col1 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier='.ca')
269
+ conv_states_ptrs = prior_tokens - 2 * stride_x_token # [BLOCK_N]
270
+ col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier='.ca')
271
+ if KERNEL_WIDTH == 5:
272
+ # ruff: noqa: F841
273
+ conv_states_ptrs = prior_tokens # [BLOCK_N]
274
+ col3 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier='.ca')
275
+ conv_states_ptrs = prior_tokens - 1 * stride_x_token # [BLOCK_N]
276
+ col2 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier='.ca')
277
+ conv_states_ptrs = prior_tokens - 2 * stride_x_token # [BLOCK_N]
278
+ col1 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier='.ca')
279
+ conv_states_ptrs = prior_tokens - 3 * stride_x_token # [BLOCK_N]
280
+ col0 = tl.load(conv_states_ptrs, mask_w, 0.0, cache_modifier='.ca')
281
+
282
+ if HAS_BIAS:
283
+ bias = bias_ptr + idx_feats
284
+ mask_bias = idx_feats < dim
285
+ acc_preload = tl.load(bias, mask=mask_bias,
286
+ other=0.0).to(tl.float32) # [BLOCK_N]
287
+ else:
288
+ acc_preload = tl.zeros((BLOCK_N, ), dtype=tl.float32)
289
+
290
+ x_base_1d = x_base + token_offset * stride_x_token # starting of chunk
291
+
292
+ # PRE-LOAD WEIGHTS
293
+ mask_w = idx_feats < dim
294
+ if KERNEL_WIDTH >= 2:
295
+ w_ptrs = w_base + (0 * stride_w_width) # [BLOCK_N] tensor
296
+ w_col0 = tl.load(w_ptrs, mask_w, other=0.0)
297
+ w_ptrs = w_base + (1 * stride_w_width) # [BLOCK_N] tensor
298
+ w_col1 = tl.load(w_ptrs, mask_w, other=0.0)
299
+ if KERNEL_WIDTH >= 3:
300
+ w_ptrs = w_base + (2 * stride_w_width) # [BLOCK_N] tensor
301
+ w_col2 = tl.load(w_ptrs, mask_w, other=0.0)
302
+ if KERNEL_WIDTH >= 4:
303
+ w_ptrs = w_base + (3 * stride_w_width) # [BLOCK_N] tensor
304
+ w_col3 = tl.load(w_ptrs, mask_w, other=0.0)
305
+ mask_x_1d = idx_feats < dim
306
+ for idx_token in range(segment_len):
307
+ acc = acc_preload
308
+
309
+ matrix_w = w_col0
310
+ matrix_x = col0
311
+ for j in tl.static_range(KERNEL_WIDTH):
312
+
313
+ if KERNEL_WIDTH == 2:
314
+ if j == 1: # KERNEL_WIDTH-1:
315
+ matrix_w = w_col1
316
+ x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N]
317
+ matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d)
318
+ elif KERNEL_WIDTH == 3:
319
+ if j == 1:
320
+ matrix_w = w_col1
321
+ matrix_x = col1
322
+ elif j == 2:
323
+ matrix_w = w_col2
324
+ x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N]
325
+ matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d)
326
+ elif KERNEL_WIDTH == 4:
327
+ if j == 1:
328
+ matrix_w = w_col1
329
+ matrix_x = col1
330
+ elif j == 2:
331
+ matrix_w = w_col2
332
+ matrix_x = col2
333
+ elif j == 3:
334
+ matrix_w = w_col3
335
+ x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N]
336
+ matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d)
337
+
338
+ acc += matrix_x * matrix_w # [BLOCK_N]
339
+
340
+ if KERNEL_WIDTH == 2:
341
+ col0 = matrix_x
342
+ elif KERNEL_WIDTH == 3:
343
+ col0 = col1
344
+ col1 = matrix_x
345
+ elif KERNEL_WIDTH == 4:
346
+ col0 = col1
347
+ col1 = col2
348
+ col2 = matrix_x
349
+
350
+ if SILU_ACTIVATION:
351
+ acc = acc / (1 + tl.exp(-acc))
352
+ mask_1d = (idx_token < segment_len) & (
353
+ idx_feats < dim) # token-index # feature-index
354
+ o_ptrs = o_ptr + (sequence_start_index + token_offset + idx_token
355
+ ) * stride_o_token + (idx_feats * stride_o_dim)
356
+
357
+ tl.store(o_ptrs, acc, mask=mask_1d)
358
+
359
+
360
+def causal_conv1d_fn(
361
+ x: torch.Tensor,
362
+ weight: torch.Tensor,
363
+ bias: Union[torch.Tensor, None],
364
+ conv_states: torch.Tensor,
365
+ query_start_loc: torch.Tensor,
366
+ cache_indices: Optional[torch.Tensor] = None,
367
+ has_initial_state: Optional[torch.Tensor] = None,
368
+ activation: Optional[str] = "silu",
369
+ pad_slot_id: int = PAD_SLOT_ID,
370
+ metadata=None,
371
+ validate_data=False,
372
+):
373
+ """support varlen + continuous batching when x is 2D tensor
374
+
375
+ x: (dim,cu_seq_len)
376
+ cu_seq_len = total tokens of all seqs in that batch
377
sequences are concatenated from left to right for varlen
378
weight: (dim, width)
28
- bias: (dim,)
379
+ conv_states: (...,dim,width - 1) itype
380
+ updated inplace if provided
381
+ [it use `cache_indices` to get the index to the cache of conv_state for that sequence
382
+
383
+ conv_state[cache_indices[i]] for seq-i - to be used as initial_state when has_initial_state[i] = True
384
+ and after that conv_state[cache_indices[i]] need to be shift-left and updated with values from 'x'
385
+ ]
386
query_start_loc: (batch + 1) int32
387
The cumulative sequence lengths of the sequences in
388
the batch, used to index into sequence. prepended by 0.
32
- for example: query_start_loc = torch.Tensor([0,10,16,17]),
389
+ if
390
+ x = [5, 1, 1, 1] <- continuous batching (batch=4)
391
+ then
392
+ query_start_loc = [0, 5, 6, 7, 8] <- the starting index of the next sequence; while the last value is
393
+ the ending index of the last sequence
394
+ [length(query_start_loc)-1 == batch]
395
+ for example: query_start_loc = torch.Tensor([0,10,16,17]),
396
x.shape=(dim,17)
397
cache_indices: (batch) int32
35
- indicates the corresponding state index,
398
+ indicates the corresponding state index,
399
like so: conv_state = conv_states[cache_indices[batch_id]]
400
has_initial_state: (batch) bool
38
- indicates whether should the kernel take the current state as initial
401
+ indicates whether should the kernel take the current state as initial
402
state for the calculations
40
- conv_states: (...,dim,width - 1) itype
41
- updated inplace if provided
42
- activation: either None or "silu" or "swish"
403
+ [single boolean for each sequence in the batch: True or False]
404
+ bias: (dim,)
405
+ activation: either None or "silu" or "swish" or True
406
pad_slot_id: int
44
- if cache_indices is passed, lets the kernel identify padded
45
- entries that will not be processed,
46
- for example: cache_indices = [pad_slot_id, 1, 20, pad_slot_id]
47
- in this case, the kernel will not process entries at
48
- indices 0 and 3
49
-
407
+ if cache_indices is passed, lets the kernel identify padded
408
+ entries that will not be processed,
409
+ for example: cache_indices = [pad_slot_id, 1, 20, pad_slot_id]
410
+ in this case, the kernel will not process entries at
411
+ indices 0 and 3
412
51
- out: (batch, dim, seqlen)
413
+ out: same shape as `x`
414
"""
53
- if activation not in [None, "silu", "swish"]:
54
- raise NotImplementedError("activation must be None, silu, or swish")
55
- if x.stride(-1) != 1:
56
- x = x.contiguous()
57
- bias = bias.contiguous() if bias is not None else None
58
-
59
- ops.causal_conv1d_fwd(x, weight, bias, conv_states, query_start_loc,
60
- cache_indices, has_initial_state, activation
61
- in ["silu", "swish"], pad_slot_id)
62
- return x
63
-
64
-
65
-def causal_conv1d_update(x: torch.Tensor,
66
- conv_state: torch.Tensor,
67
- weight: torch.Tensor,
68
- bias: Optional[torch.Tensor] = None,
69
- activation: Optional[str] = None,
70
- cache_seqlens: Optional[torch.Tensor] = None,
71
- conv_state_indices: Optional[torch.Tensor] = None,
72
- pad_slot_id: int = PAD_SLOT_ID):
415
+ if isinstance(activation, bool) and activation:
416
+ activation = "silu"
417
+
418
+ args = None
419
+ out = torch.zeros_like(x)
420
+ if metadata is not None:
421
+ cu_seqlen = metadata.cu_seqlen
422
+ nums_dict = metadata.nums_dict
423
+ #x = metadata.x
424
+ args = nums_dict
425
+ batch_ptr = metadata.batch_ptr
426
+ token_chunk_offset_ptr = metadata.token_chunk_offset_ptr
427
+ else:
428
+ seqlens = np.diff(query_start_loc.to('cpu'))
429
+ args = seqlens
430
+ MAX_NUM_PROGRAMS = 1024
431
+
432
+ batch_ptr = torch.full(
433
+ (MAX_NUM_PROGRAMS, ),
434
+ PAD_SLOT_ID,
435
+ dtype=torch.int32,
436
+ device=x.device
437
+ ) # tracking which seq-idx the Triton program is handling
438
+ token_chunk_offset_ptr = torch.full(
439
+ (MAX_NUM_PROGRAMS, ),
440
+ PAD_SLOT_ID,
441
+ dtype=torch.int32,
442
+ device=x.device
443
+ ) # tracking BLOCK_M-based index in the sequence the Triton program is handling
444
+
445
+ is_channel_last = (x.stride(0) == 1) & (x.stride(1) > 1)
446
+ dim, cu_seqlen = x.shape
447
+ _, width = weight.shape
448
+ state_len = width - 1
449
+ np2_statelen = triton.next_power_of_2(state_len)
450
+
451
+ padded_batch = query_start_loc.size(0) - 1
452
+ stride_x_seq = 0
453
+ stride_x_dim = x.stride(0)
454
+ stride_x_token = x.stride(1)
455
+ stride_w_dim = weight.stride(0)
456
+ stride_w_width = weight.stride(1)
457
+ stride_istate_seq = 0
458
+ stride_istate_dim = 0
459
+ stride_istate_token = 0
460
+ num_cache_lines = 0
461
+ if conv_states is not None:
462
+ # extensions to support vLLM:
463
+ # 1. conv_states is used to replaced initial_states
464
+ # 2. conv_states serve as a cache with num cache lines can be larger than batch size
465
+ # 3. mapping from sequence x[idx] to a cache line at index as specified via cache_indices[idx]
466
+ # 4. computation can be skipped if cache_indices[idx] == pad_slot_id
467
+ num_cache_lines = conv_states.size(0)
468
+ assert (num_cache_lines, dim, width - 1) == conv_states.shape
469
+ stride_istate_seq = conv_states.stride(0)
470
+ stride_istate_dim = conv_states.stride(1)
471
+ stride_istate_token = conv_states.stride(2)
472
+ assert stride_istate_dim == 1
473
+ if out.dim() == 2:
474
+ stride_o_seq = 0
475
+ stride_o_dim = out.stride(0)
476
+ stride_o_token = out.stride(1)
477
+ else:
478
+ stride_o_seq = out.stride(0)
479
+ stride_o_dim = out.stride(1)
480
+ stride_o_token = out.stride(2)
481
+
482
+ if validate_data:
483
+ assert x.dim() == 2
484
+ assert query_start_loc is not None
485
+ assert query_start_loc.dim() == 1
486
+ assert x.stride(0) == 1 or x.stride(1) == 1
487
+ if bias is not None:
488
+ assert bias.dim() == 1
489
+ assert dim == bias.size(0)
490
+ if cache_indices is not None:
491
+ assert cache_indices.dim() == 1
492
+ assert padded_batch == cache_indices.size(0)
493
+ if has_initial_state is not None:
494
+ assert has_initial_state.size() == (padded_batch, )
495
+ assert conv_states is not None, "ERROR: `has_initial_state` is used, which needs also `conv_states`"
496
+ assert weight.stride(1) == 1
497
+ assert (dim, width) == weight.shape
498
+ assert is_channel_last, "Need to run in channel-last layout"
499
+
500
+ if metadata is None:
501
+
502
+ def num_program(META, seqlens):
503
+ tot = 0
504
+
505
+ mlist = []
506
+ offsetlist = [] # type: ignore
507
+
508
+ nums = -(-seqlens // META["BLOCK_M"])
509
+
510
+ tot = nums.sum().item()
511
+ mlist = np.repeat(np.arange(len(nums)), nums)
512
+ for idx, num in enumerate(nums):
513
+ offsetlist.extend(
514
+ range(num)
515
+ ) # chunk-idx if a sequence is split into multiple chunks
516
+
517
+ if META["batch_ptr"].nelement() < len(mlist):
518
+ newlen = len(mlist) + 1
519
+ META["batch_ptr"].resize_(newlen).fill_(PAD_SLOT_ID)
520
+ META["token_chunk_offset_ptr"].resize_(newlen).fill_(
521
+ PAD_SLOT_ID)
522
+
523
+ if META["batch_ptr"].nelement() >= len(mlist):
524
+ META["batch_ptr"][0:len(mlist)].copy_(
525
+ torch.from_numpy(np.array(mlist)))
526
+ META["token_chunk_offset_ptr"][0:len(mlist)].copy_(
527
+ torch.from_numpy(np.array(offsetlist)))
528
+
529
+ META["batch_ptr"] = META["batch_ptr"].to(META["x_ptr"].device)
530
+ META["token_chunk_offset_ptr"] = META["token_chunk_offset_ptr"].to(
531
+ META["x_ptr"].device)
532
+ return tot
533
+ else:
534
+
535
+ def num_program(META, nums_dict):
536
+ tot = nums_dict[META["BLOCK_M"]]['tot']
537
+
538
+ mlist = nums_dict[META["BLOCK_M"]]['mlist']
539
+ mlist_len = nums_dict[META["BLOCK_M"]]['mlist_len']
540
+
541
+ offsetlist = nums_dict[META["BLOCK_M"]]['offsetlist']
542
+
543
+ if nums_dict[META["BLOCK_M"]]["batch_ptr"] is not None:
544
+ META["batch_ptr"] = nums_dict[META["BLOCK_M"]]["batch_ptr"]
545
+ META["token_chunk_offset_ptr"] = nums_dict[
546
+ META["BLOCK_M"]]["token_chunk_offset_ptr"]
547
+ else:
548
+ if META["batch_ptr"].nelement() < mlist_len:
549
+ newlen = mlist_len + 1
550
+ META["batch_ptr"].resize_(newlen).fill_(PAD_SLOT_ID)
551
+ META["token_chunk_offset_ptr"].resize_(newlen).fill_(
552
+ PAD_SLOT_ID)
553
+
554
+ if META["batch_ptr"].nelement() >= mlist_len:
555
+ META["batch_ptr"][0:mlist_len].copy_(mlist)
556
+ META["token_chunk_offset_ptr"][0:mlist_len].copy_(
557
+ offsetlist)
558
+ return tot
559
+
560
+ def grid(META):
561
+ return (
562
+ num_program(META, args),
563
+ triton.cdiv(dim, META["BLOCK_N"]),
564
+ )
565
+
566
+ if batch_ptr.device != x.device:
567
+ batch_ptr = batch_ptr.to(x.device)
568
+ token_chunk_offset_ptr = token_chunk_offset_ptr.to(x.device)
569
+
570
+ _causal_conv1d_fwd_kernel[grid](
571
+ # Pointers to matrices
572
+ x,
573
+ weight,
574
+ bias,
575
+ conv_states,
576
+ cache_indices,
577
+ has_initial_state,
578
+ query_start_loc,
579
+ batch_ptr,
580
+ token_chunk_offset_ptr,
581
+ out,
582
+ # Matrix dimensions
583
+ padded_batch,
584
+ dim,
585
+ cu_seqlen,
586
+ num_cache_lines,
587
+ # stride
588
+ stride_x_seq,
589
+ stride_x_dim,
590
+ stride_x_token,
591
+ stride_w_dim,
592
+ stride_w_width,
593
+ stride_istate_seq,
594
+ stride_istate_dim,
595
+ stride_istate_token,
596
+ stride_o_seq,
597
+ stride_o_dim,
598
+ stride_o_token,
599
+ # others
600
+ pad_slot_id,
601
+ # META
602
+ HAS_BIAS=bias is not None,
603
+ KERNEL_WIDTH=width,
604
+ SILU_ACTIVATION=activation in ["silu", "swish"],
605
+ HAS_INITIAL_STATES=has_initial_state is not None,
606
+ HAS_CACHE=conv_states is not None,
607
+ IS_CONTINUOUS_BATCHING=cache_indices is not None,
608
+ USE_PAD_SLOT=pad_slot_id is not None,
609
+ NP2_STATELEN=np2_statelen,
610
+ DECODE_SEQLEN=1,
611
+ #launch_cooperative_grid=True
612
+ BLOCK_M=8,
613
+ BLOCK_N=256,
614
+ num_stages=2,
615
+ )
616
+ return out
617
+
618
+
619
+@triton.jit()
620
+def _causal_conv1d_update_kernel(
621
+ # Pointers to matrices
622
+ x_ptr, # (batch, dim, seqlen)
623
+ w_ptr, # (dim, width)
624
+ bias_ptr,
625
+ conv_state_ptr,
626
+ cache_seqlens_ptr, # circular buffer
627
+ conv_state_indices_ptr,
628
+ o_ptr, # (batch, dim, seqlen)
629
+ # Matrix dimensions
630
+ batch: int,
631
+ dim: tl.constexpr,
632
+ seqlen: tl.constexpr,
633
+ state_len: tl.constexpr,
634
+ num_cache_lines: tl.constexpr, # added to support vLLM larger cache lines
635
+ # Strides
636
+ stride_x_seq: tl.constexpr,
637
+ stride_x_dim: tl.constexpr,
638
+ stride_x_token: tl.constexpr,
639
+ stride_w_dim: tl.constexpr,
640
+ stride_w_width: tl.constexpr,
641
+ stride_conv_state_seq: tl.constexpr,
642
+ stride_conv_state_dim: tl.constexpr,
643
+ stride_conv_state_tok: tl.constexpr,
644
+ stride_o_seq: tl.constexpr,
645
+ stride_o_dim: tl.constexpr,
646
+ stride_o_token: tl.constexpr,
647
+ # others
648
+ pad_slot_id: tl.constexpr,
649
+ # Meta-parameters
650
+ HAS_BIAS: tl.constexpr,
651
+ KERNEL_WIDTH: tl.constexpr,
652
+ SILU_ACTIVATION: tl.constexpr,
653
+ IS_CONTINUOUS_BATCHING: tl.constexpr,
654
+ NP2_STATELEN: tl.constexpr,
655
+ USE_PAD_SLOT: tl.constexpr,
656
+ BLOCK_N: tl.constexpr,
657
+):
658
+ # ruff: noqa: E501
659
+ idx_seq = tl.program_id(0)
660
+ if idx_seq >= batch:
661
+ return
662
+
663
+ # [BLOCK_N,] elements along the feature-dimension (channel)
664
+ idx_feats = tl.program_id(1) * BLOCK_N + tl.arange(0, BLOCK_N)
665
+
666
+ if IS_CONTINUOUS_BATCHING:
667
+ # mask = idx_seq < batch
668
+ conv_state_batch_coord = tl.load(conv_state_indices_ptr + idx_seq)
669
+ else:
670
+ conv_state_batch_coord = idx_seq
671
+ if USE_PAD_SLOT: # noqa
672
+ if conv_state_batch_coord == pad_slot_id:
673
+ # not processing as this is not the actual sequence
674
+ return
675
+
676
+ # STEP 1: READ init_state data
677
+ conv_states_base = (conv_state_ptr +
678
+ (conv_state_batch_coord * stride_conv_state_seq) +
679
+ (idx_feats * stride_conv_state_dim))
680
+ mask_w = idx_feats < dim
681
+
682
+ prior_tokens = conv_states_base
683
+ if KERNEL_WIDTH >= 2:
684
+ conv_states_ptrs = prior_tokens # [BLOCK_N]
685
+ col0 = tl.load(conv_states_ptrs, mask_w, 0.0)
686
+ if KERNEL_WIDTH >= 3:
687
+ conv_states_ptrs = prior_tokens + 1 * stride_conv_state_tok # [BLOCK_N]
688
+ col1 = tl.load(conv_states_ptrs, mask_w, 0.0)
689
+ if KERNEL_WIDTH >= 4:
690
+ conv_states_ptrs = prior_tokens + 2 * stride_conv_state_tok # [BLOCK_N]
691
+ col2 = tl.load(conv_states_ptrs, mask_w, 0.0)
692
+ if KERNEL_WIDTH == 5:
693
+ conv_states_ptrs = prior_tokens + 3 * stride_conv_state_tok # [BLOCK_N]
694
+ col3 = tl.load(conv_states_ptrs, mask_w, 0.0)
695
+
696
+ # STEP 2: assume state_len > seqlen
697
+ idx_tokens = tl.arange(0, NP2_STATELEN) # [BLOCK_M]
698
+
699
+ conv_state_ptrs_source = (
700
+ conv_state_ptr + (conv_state_batch_coord * stride_conv_state_seq) +
701
+ (idx_feats * stride_conv_state_dim)[None, :] +
702
+ ((idx_tokens + seqlen) * stride_conv_state_tok)[:, None]
703
+ ) # [BLOCK_M, BLOCK_N]
704
+ mask = ((conv_state_batch_coord < num_cache_lines)
705
+ & ((idx_tokens + seqlen) < state_len)[:, None]
706
+ & (idx_feats < dim)[None, :])
707
+ conv_state = tl.load(conv_state_ptrs_source, mask, other=0.0)
708
+
709
+ VAL = state_len - seqlen
710
+ x_base = x_ptr + (idx_seq * stride_x_seq) + (idx_feats * stride_x_dim
711
+ ) # [BLOCK_N]
712
+
713
+ x_ptrs = x_base[None, :] + (
714
+ (idx_tokens - VAL) * stride_x_token)[:, None] # [BLOCK_M, BLOCK_N]
715
+
716
+ mask_x = ((idx_tokens - VAL >= 0)[:, None] &
717
+ (idx_tokens - VAL < seqlen)[:, None] & (idx_feats < dim)[None, :]
718
+ ) # token-index # token-index # feature-index
719
+ loaded_x = tl.load(x_ptrs, mask_x, 0.0)
720
+ tl.debug_barrier()
721
+
722
+ new_conv_state = tl.where(mask, conv_state, loaded_x)
723
+
724
+ conv_state_base = (conv_state_ptr +
725
+ (conv_state_batch_coord * stride_conv_state_seq) +
726
+ (idx_feats * stride_conv_state_dim)) # [BLOCK_N,]
727
+ conv_state_ptrs_target = conv_state_base + (
728
+ idx_tokens * stride_conv_state_tok)[:, None] # [BLOCK_M, BLOCK_N]
729
+ mask = (idx_tokens < state_len)[:, None] & (idx_feats < dim)[None, :]
730
+ tl.store(conv_state_ptrs_target, new_conv_state, mask)
731
+
732
+ # STEP 3: init accumulator
733
+ if HAS_BIAS:
734
+ bias = bias_ptr + idx_feats
735
+ mask_bias = idx_feats < dim
736
+ acc_preload = tl.load(bias, mask=mask_bias,
737
+ other=0.0).to(tl.float32) # [BLOCK_N]
738
+ else:
739
+ acc_preload = tl.zeros((BLOCK_N, ), dtype=tl.float32)
740
+
741
+ # STEP 4:
742
+ # PRE-LOAD WEIGHTS
743
+ # first kernel column, configured for weights to handle BLOCK_N features in range
744
+ w_base = w_ptr + (idx_feats * stride_w_dim) # [BLOCK_N,]
745
+ mask_w = idx_feats < dim
746
+ if KERNEL_WIDTH >= 2:
747
+ w_ptrs = w_base + (0 * stride_w_width) # [BLOCK_N] tensor
748
+ w_col0 = tl.load(w_ptrs, mask_w, other=0.0)
749
+ w_ptrs = w_base + (1 * stride_w_width) # [BLOCK_N] tensor
750
+ w_col1 = tl.load(w_ptrs, mask_w, other=0.0)
751
+ if KERNEL_WIDTH >= 3:
752
+ w_ptrs = w_base + (2 * stride_w_width) # [BLOCK_N] tensor
753
+ w_col2 = tl.load(w_ptrs, mask_w, other=0.0)
754
+ if KERNEL_WIDTH >= 4:
755
+ w_ptrs = w_base + (3 * stride_w_width) # [BLOCK_N] tensor
756
+ w_col3 = tl.load(w_ptrs, mask_w, other=0.0)
757
+
758
+ x_base_1d = x_base # starting of chunk [BLOCK_N]
759
+ mask_x_1d = idx_feats < dim
760
+
761
+ # STEP 5: compute each token
762
+ for idx_token in tl.static_range(seqlen):
763
+ acc = acc_preload
764
+
765
+ matrix_w = w_col0
766
+ matrix_x = col0
767
+ for j in tl.static_range(KERNEL_WIDTH):
768
+ if KERNEL_WIDTH == 2:
769
+ if j == 1: # KERNEL_WIDTH-1:
770
+ matrix_w = w_col1
771
+ x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N]
772
+ matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d)
773
+ elif KERNEL_WIDTH == 3:
774
+ if j == 1:
775
+ matrix_w = w_col1
776
+ matrix_x = col1
777
+ elif j == 2:
778
+ matrix_w = w_col2
779
+ x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N]
780
+ matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d)
781
+ elif KERNEL_WIDTH == 4:
782
+ if j == 1:
783
+ matrix_w = w_col1
784
+ matrix_x = col1
785
+ elif j == 2:
786
+ matrix_w = w_col2
787
+ matrix_x = col2
788
+ elif j == 3:
789
+ matrix_w = w_col3
790
+ x_ptrs_1d = x_base_1d + idx_token * stride_x_token # [BLOCK_N]
791
+ matrix_x = tl.load(x_ptrs_1d, mask=mask_x_1d)
792
+
793
+ acc += matrix_x * matrix_w # [BLOCK_N]
794
+
795
+ if KERNEL_WIDTH == 2:
796
+ col0 = matrix_x
797
+ elif KERNEL_WIDTH == 3:
798
+ col0 = col1
799
+ col1 = matrix_x
800
+ elif KERNEL_WIDTH == 4:
801
+ col0 = col1
802
+ col1 = col2
803
+ col2 = matrix_x
804
+
805
+ if SILU_ACTIVATION:
806
+ acc = acc / (1 + tl.exp(-acc))
807
+ mask_1d = (idx_token < seqlen) & (idx_feats < dim
808
+ ) # token-index # feature-index
809
+ o_ptrs = o_ptr + (
810
+ idx_seq) * stride_o_seq + idx_token * stride_o_token + (
811
+ idx_feats * stride_o_dim)
812
+
813
+ tl.store(o_ptrs, acc, mask=mask_1d)
814
+
815
+
816
+def causal_conv1d_update(
817
+ x: torch.Tensor,
818
+ conv_state: torch.Tensor,
819
+ weight: torch.Tensor,
820
+ bias: Optional[torch.Tensor] = None,
821
+ activation: Union[bool, str, None] = None,
822
+ cache_seqlens: Optional[torch.Tensor] = None,
823
+ conv_state_indices: Optional[torch.Tensor] = None,
824
+ pad_slot_id: int = PAD_SLOT_ID,
825
+ metadata=None,
826
+ validate_data=False,
827
+):
828
"""
829
x: (batch, dim) or (batch, dim, seqlen)
75
- conv_state: (batch, dim, state_len), where state_len >= width - 1
830
+ [shape=2: single token prediction]
831
+ [shape=3: single or multiple tokens prediction]
832
+ conv_state: (..., dim, state_len), where state_len >= width - 1
833
weight: (dim, width)
834
bias: (dim,)
835
cache_seqlens: (batch,), dtype int32.
836
If not None, the conv_state is treated as a circular buffer.
80
- The conv_state will be updated by copying x to the conv_state
837
+ The conv_state will be updated by copying x to the conv_state
838
starting at the index
839
@cache_seqlens % state_len.
840
conv_state_indices: (batch,), dtype int32
84
- If not None, the conv_state is a larger tensor along the batch dim,
841
+ If not None, the conv_state is a larger tensor along the batch dim,
842
and we are selecting the batch coords specified by conv_state_indices.
843
Useful for a continuous batching scenario.
844
pad_slot_id: int
88
- if cache_indices is passed, lets the kernel identify padded
89
- entries that will not be processed,
90
- for example: cache_indices = [pad_slot_id, 1 ,20 ,pad_slot_id]
91
- in this case, the kernel will not process entries at
845
+ if cache_indices is passed, lets the kernel identify padded
846
+ entries that will not be processed,
847
+ for example: cache_indices = [pad_slot_id, 1 ,20 ,pad_slot_id]
848
+ in this case, the kernel will not process entries at
849
indices 0 and 3
850
out: (batch, dim) or (batch, dim, seqlen)
851
"""
95
- if activation not in [None, "silu", "swish"]:
96
- raise NotImplementedError("activation must be None, silu, or swish")
97
- activation_val = activation in ["silu", "swish"]
852
+ if validate_data:
853
+ assert cache_seqlens is None # not implemented yet - ok for vLLM
854
+ assert pad_slot_id is not None
855
+ assert x.stride(1) == 1
856
+ if isinstance(activation, bool):
857
+ activation = "silu" if activation is True else None
858
+ elif activation is not None:
859
+ assert activation in ["silu", "swish"]
860
unsqueeze = x.dim() == 2
861
if unsqueeze:
862
+ # make it (batch, dim, seqlen) with seqlen == 1
863
x = x.unsqueeze(-1)
101
- ops.causal_conv1d_update(x, conv_state, weight, bias, activation_val,
102
- cache_seqlens, conv_state_indices, pad_slot_id)
864
+ batch, dim, seqlen = x.shape
865
+ _, width = weight.shape
866
+ # conv_state: (..., dim, state_len), where state_len >= width - 1
867
+ num_cache_lines, _, state_len = conv_state.size()
868
+
869
+ if validate_data:
870
+ assert dim == weight.size(0)
871
+ assert conv_state.stride(
872
+ -2
873
+ ) == 1, f"ERROR: expect contiguous along feat-dim of conv_state (currently stride={conv_state.stride()})"
874
+ assert state_len >= width - 1
875
+ # when above happens, we don't shift-left to keep any records in conv_state
876
+ assert dim == conv_state.size(1)
877
+ if conv_state_indices is None:
878
+ assert conv_state.size(0) >= batch
879
+ else:
880
+ assert (batch, ) == conv_state_indices.shape
881
+
882
+ assert num_cache_lines >= batch
883
+ assert weight.stride(1) == 1 # Need this
884
+ assert cache_seqlens is None # not needed for vLLM - circular buffer
885
+
886
+ # adopt the strategy in vLLM that overwrite on 'x' directly, rather than creating a new tensor 'o'
887
+ out = x
888
+ stride_w_dim, stride_w_width = weight.stride()
889
+
890
+ stride_x_seq, stride_x_dim, stride_x_token = x.stride(
891
+ ) # X (batch, dim, seqlen)
892
+
893
+ stride_o_seq, stride_o_dim, stride_o_token = out.stride()
894
+
895
+ stride_istate_seq, stride_istate_dim, stride_istate_token = conv_state.stride(
896
+ )
897
+ state_len = width - 1
898
+ np2_statelen = triton.next_power_of_2(state_len)
899
+
900
+ def grid(META):
901
+ return (
902
+ batch,
903
+ triton.cdiv(dim, META["BLOCK_N"]),
904
+ )
905
+
906
+ _causal_conv1d_update_kernel[grid](
907
+ # Pointers to matrices
908
+ x,
909
+ weight,
910
+ bias,
911
+ conv_state,
912
+ cache_seqlens,
913
+ conv_state_indices,
914
+ out,
915
+ # Matrix dimensions
916
+ batch,
917
+ dim,
918
+ seqlen,
919
+ state_len,
920
+ num_cache_lines,
921
+ # stride
922
+ stride_x_seq,
923
+ stride_x_dim,
924
+ stride_x_token,
925
+ stride_w_dim,
926
+ stride_w_width,
927
+ stride_istate_seq,
928
+ stride_istate_dim,
929
+ stride_istate_token,
930
+ stride_o_seq,
931
+ stride_o_dim,
932
+ stride_o_token,
933
+ # others
934
+ pad_slot_id,
935
+ # META
936
+ HAS_BIAS=bias is not None,
937
+ KERNEL_WIDTH=width,
938
+ SILU_ACTIVATION=activation in ["silu", "swish"],
939
+ IS_CONTINUOUS_BATCHING=conv_state_indices is not None,
940
+ NP2_STATELEN=np2_statelen,
941
+ USE_PAD_SLOT=pad_slot_id is not None,
942
+ BLOCK_N=256,
943
+ )
944
if unsqueeze:
104
- x = x.squeeze(-1)
105
- return x
945
+ out = out.squeeze(-1)
946
+ return out
vllm/model_executor/models/mamba_cache.py
+4
-2
@@ -36,10 +36,12 @@ class MambaCacheManager(ConstantSizeCache):
36
# Initialize parent class
37
super().__init__(max_batch_size)
38
39
+ # assume conv_state = (dim, state_len)
40
+ assert conv_state_shape[0] > conv_state_shape[1]
41
conv_state = torch.empty(size=(num_mamba_layers, max_batch_size) +
40
- conv_state_shape,
42
+ (conv_state_shape[1], conv_state_shape[0]),
43
dtype=dtype,
42
- device="cuda")
44
+ device="cuda").transpose(-1, -2)
45
temporal_state = torch.empty(size=(num_mamba_layers, max_batch_size) +
46
temporal_state_shape,
47
dtype=dtype,
vllm/v1/attention/backends/mamba_attn.py
+42
-3
@@ -1,14 +1,13 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import math
4
from dataclasses import dataclass
4
-from typing import TYPE_CHECKING
5
+from typing import TYPE_CHECKING, Optional
6
7
import torch
8
9
from vllm.attention.backends.abstract import AttentionBackend
10
from vllm.config import VllmConfig, get_layers_from_vllm_config
10
-from vllm.model_executor.layers.mamba.mamba2_metadata import (
11
- _query_start_loc_to_chunk_indices_offsets)
11
from vllm.v1.attention.backends.utils import (AttentionMetadataBuilder,
12
CommonAttentionMetadata)
13
from vllm.v1.kv_cache_interface import MambaSpec
@@ -29,6 +28,42 @@ def get_mamba2_chunk_size(vllm_config: VllmConfig) -> int:
28
return chunk_sizes.pop()
29
30
31
+def _query_start_loc_to_chunk_indices_offsets(query_start_loc: torch.Tensor,
32
+ chunk_size: int,
33
+ total_seqlens: int):
34
+
35
+ cu_seqlens = query_start_loc[1:] # remove prepended 0
36
+
37
+ # outputs will have length expansion of chunks that do not divide
38
+ # chunk_size
39
+ N = math.ceil(total_seqlens / chunk_size) + (cu_seqlens[:-1] % chunk_size
40
+ > 0).sum()
41
+ chunk_indices = torch.arange(N,
42
+ dtype=torch.int,
43
+ device=query_start_loc.device)
44
+ chunk_offsets = torch.zeros((N, ),
45
+ dtype=torch.int,
46
+ device=query_start_loc.device)
47
+
48
+ p = 0 # num of insertions
49
+ for s, e in zip(cu_seqlens[:-1], cu_seqlens[1:]):
50
+
51
+ # if does not divide chunk_size, then there is one chunk insertion
52
+ p += (s % chunk_size > 0)
53
+
54
+ # get the dimensions
55
+ # - the + 1 for _e is to shift the boundary by one chunk
56
+ # - this shifting is not needed if chunk_size divides e
57
+ _s, _e = s // chunk_size + p, e // chunk_size + p + (e % chunk_size
58
+ > 0)
59
+
60
+ # adjust indices and offsets
61
+ chunk_indices[_s:_e] -= p
62
+ chunk_offsets[_s] = s % chunk_size
63
+
64
+ return chunk_indices, chunk_offsets
65
+
66
+
67
class Mamba2AttentionBackend(AttentionBackend):
68
69
@staticmethod
@@ -53,6 +88,10 @@ class Mamba2AttentionMetadata:
88
chunk_offsets: torch.Tensor
89
90
state_indices_tensor: torch.Tensor # shape: [batch,]
91
+ nums_dict: Optional[dict] = None
92
+ cu_seqlen: Optional[int] = None
93
+ batch_ptr: Optional[torch.tensor] = None
94
+ token_chunk_offset_ptr: Optional[torch.tensor] = None
95
96
97
class Mamba2AttentionMetadataBuilder(