causal-conv1d 1.5.1__tar.gz → 1.5.2__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- causal_conv1d-1.5.2/MANIFEST.in +3 -0
- {causal_conv1d-1.5.1 → causal_conv1d-1.5.2}/PKG-INFO +1 -1
- {causal_conv1d-1.5.1 → causal_conv1d-1.5.2}/causal_conv1d/__init__.py +1 -1
- {causal_conv1d-1.5.1 → causal_conv1d-1.5.2}/causal_conv1d.egg-info/PKG-INFO +1 -1
- {causal_conv1d-1.5.1 → causal_conv1d-1.5.2}/causal_conv1d.egg-info/SOURCES.txt +1 -0
- causal_conv1d-1.5.2/csrc/causal_conv1d.cpp +466 -0
- causal_conv1d-1.5.1/MANIFEST.in +0 -4
- {causal_conv1d-1.5.1 → causal_conv1d-1.5.2}/AUTHORS +0 -0
- {causal_conv1d-1.5.1 → causal_conv1d-1.5.2}/LICENSE +0 -0
- {causal_conv1d-1.5.1 → causal_conv1d-1.5.2}/README.md +0 -0
- {causal_conv1d-1.5.1 → causal_conv1d-1.5.2}/causal_conv1d/causal_conv1d_interface.py +0 -0
- {causal_conv1d-1.5.1 → causal_conv1d-1.5.2}/causal_conv1d/causal_conv1d_varlen.py +0 -0
- {causal_conv1d-1.5.1 → causal_conv1d-1.5.2}/causal_conv1d/cpp_functions.py +0 -0
- {causal_conv1d-1.5.1 → causal_conv1d-1.5.2}/causal_conv1d.egg-info/dependency_links.txt +0 -0
- {causal_conv1d-1.5.1 → causal_conv1d-1.5.2}/causal_conv1d.egg-info/requires.txt +0 -0
- {causal_conv1d-1.5.1 → causal_conv1d-1.5.2}/causal_conv1d.egg-info/top_level.txt +0 -0
- {causal_conv1d-1.5.1 → causal_conv1d-1.5.2}/csrc/causal_conv1d.h +0 -0
- {causal_conv1d-1.5.1 → causal_conv1d-1.5.2}/csrc/causal_conv1d_bwd.cu +0 -0
- {causal_conv1d-1.5.1 → causal_conv1d-1.5.2}/csrc/causal_conv1d_common.h +0 -0
- {causal_conv1d-1.5.1 → causal_conv1d-1.5.2}/csrc/causal_conv1d_fwd.cu +0 -0
- {causal_conv1d-1.5.1 → causal_conv1d-1.5.2}/csrc/causal_conv1d_update.cu +0 -0
- {causal_conv1d-1.5.1 → causal_conv1d-1.5.2}/csrc/static_switch.h +0 -0
- {causal_conv1d-1.5.1 → causal_conv1d-1.5.2}/pyproject.toml +0 -0
- {causal_conv1d-1.5.1 → causal_conv1d-1.5.2}/setup.cfg +0 -0
- {causal_conv1d-1.5.1 → causal_conv1d-1.5.2}/setup.py +0 -0
@@ -13,6 +13,7 @@ causal_conv1d.egg-info/SOURCES.txt
|
|
13
13
|
causal_conv1d.egg-info/dependency_links.txt
|
14
14
|
causal_conv1d.egg-info/requires.txt
|
15
15
|
causal_conv1d.egg-info/top_level.txt
|
16
|
+
csrc/causal_conv1d.cpp
|
16
17
|
csrc/causal_conv1d.h
|
17
18
|
csrc/causal_conv1d_bwd.cu
|
18
19
|
csrc/causal_conv1d_common.h
|
@@ -0,0 +1,466 @@
|
|
1
|
+
/******************************************************************************
|
2
|
+
* Copyright (c) 2024, Tri Dao.
|
3
|
+
******************************************************************************/
|
4
|
+
|
5
|
+
#include <c10/cuda/CUDAGuard.h>
|
6
|
+
#include <c10/cuda/CUDAStream.h>
|
7
|
+
#include <torch/python.h>
|
8
|
+
#include <vector>
|
9
|
+
|
10
|
+
#include "causal_conv1d.h"
|
11
|
+
|
12
|
+
#define CHECK_SHAPE(x, ...) TORCH_CHECK(x.sizes() == torch::IntArrayRef({__VA_ARGS__}), #x " must have shape (" #__VA_ARGS__ ")")
|
13
|
+
|
14
|
+
#define DISPATCH_ITYPE_FLOAT_AND_HALF_AND_BF16(ITYPE, NAME, ...) \
|
15
|
+
if (ITYPE == at::ScalarType::Half) { \
|
16
|
+
using input_t = at::Half; \
|
17
|
+
__VA_ARGS__(); \
|
18
|
+
} else if (ITYPE == at::ScalarType::BFloat16) { \
|
19
|
+
using input_t = at::BFloat16; \
|
20
|
+
__VA_ARGS__(); \
|
21
|
+
} else if (ITYPE == at::ScalarType::Float) { \
|
22
|
+
using input_t = float; \
|
23
|
+
__VA_ARGS__(); \
|
24
|
+
} else { \
|
25
|
+
AT_ERROR(#NAME, " not implemented for input type '", toString(ITYPE), "'"); \
|
26
|
+
}
|
27
|
+
|
28
|
+
#define DISPATCH_WTYPE_FLOAT_AND_HALF_AND_BF16(WTYPE, NAME, ...) \
|
29
|
+
if (WTYPE == at::ScalarType::Half) { \
|
30
|
+
using weight_t = at::Half; \
|
31
|
+
__VA_ARGS__(); \
|
32
|
+
} else if (WTYPE == at::ScalarType::BFloat16) { \
|
33
|
+
using weight_t = at::BFloat16; \
|
34
|
+
__VA_ARGS__(); \
|
35
|
+
} else if (WTYPE == at::ScalarType::Float) { \
|
36
|
+
using weight_t = float; \
|
37
|
+
__VA_ARGS__(); \
|
38
|
+
} else { \
|
39
|
+
AT_ERROR(#NAME, " not implemented for weight type '", toString(WTYPE), "'"); \
|
40
|
+
}
|
41
|
+
|
42
|
+
template<typename input_t, typename weight_t>
|
43
|
+
void causal_conv1d_fwd_cuda(ConvParamsBase ¶ms, cudaStream_t stream);
|
44
|
+
template <typename input_t, typename weight_t>
|
45
|
+
void causal_conv1d_channellast_fwd_cuda(ConvParamsBase ¶ms, cudaStream_t stream);
|
46
|
+
|
47
|
+
template<typename input_t, typename weight_t>
|
48
|
+
void causal_conv1d_bwd_cuda(ConvParamsBwd ¶ms, cudaStream_t stream);
|
49
|
+
template<typename input_t, typename weight_t>
|
50
|
+
void causal_conv1d_channellast_bwd_cuda(ConvParamsBwd ¶ms, cudaStream_t stream);
|
51
|
+
|
52
|
+
template<typename input_t, typename weight_t>
|
53
|
+
void causal_conv1d_update_cuda(ConvParamsBase ¶ms, cudaStream_t stream);
|
54
|
+
|
55
|
+
void set_conv_params_fwd(ConvParamsBase ¶ms,
|
56
|
+
// sizes
|
57
|
+
const size_t batch,
|
58
|
+
const size_t dim,
|
59
|
+
const size_t seqlen,
|
60
|
+
const size_t width,
|
61
|
+
// device pointers
|
62
|
+
const at::Tensor x,
|
63
|
+
const at::Tensor weight,
|
64
|
+
const at::Tensor out,
|
65
|
+
void* bias_ptr,
|
66
|
+
bool silu_activation) {
|
67
|
+
|
68
|
+
// Reset the parameters
|
69
|
+
memset(¶ms, 0, sizeof(params));
|
70
|
+
|
71
|
+
params.batch = batch;
|
72
|
+
params.dim = dim;
|
73
|
+
params.seqlen = seqlen;
|
74
|
+
params.width = width;
|
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_ptr;
|
82
|
+
params.out_ptr = out.data_ptr();
|
83
|
+
// All stride are in elements, not bytes.
|
84
|
+
params.x_batch_stride = x.stride(0);
|
85
|
+
params.x_c_stride = x.stride(1);
|
86
|
+
params.x_l_stride = x.stride(-1);
|
87
|
+
params.weight_c_stride = weight.stride(0);
|
88
|
+
params.weight_width_stride = weight.stride(1);
|
89
|
+
params.out_batch_stride = out.stride(0);
|
90
|
+
params.out_c_stride = out.stride(1);
|
91
|
+
params.out_l_stride = out.stride(-1);
|
92
|
+
}
|
93
|
+
|
94
|
+
|
95
|
+
void set_conv_params_bwd(ConvParamsBwd ¶ms,
|
96
|
+
// sizes
|
97
|
+
const size_t batch,
|
98
|
+
const size_t dim,
|
99
|
+
const size_t seqlen,
|
100
|
+
const size_t width,
|
101
|
+
// device pointers
|
102
|
+
const at::Tensor x,
|
103
|
+
const at::Tensor weight,
|
104
|
+
void* bias_ptr,
|
105
|
+
const at::Tensor dout,
|
106
|
+
const at::Tensor dx,
|
107
|
+
const at::Tensor dweight,
|
108
|
+
void* dbias_ptr,
|
109
|
+
bool silu_activation) {
|
110
|
+
// Pass in "dout" instead of "out", we're not gonna use "out" at all.
|
111
|
+
set_conv_params_fwd(params, batch, dim, seqlen, width,
|
112
|
+
x, weight, dout, bias_ptr, silu_activation);
|
113
|
+
|
114
|
+
// Set the pointers and strides.
|
115
|
+
params.dout_ptr = dout.data_ptr();
|
116
|
+
params.dx_ptr = dx.data_ptr();
|
117
|
+
params.dweight_ptr = dweight.data_ptr();
|
118
|
+
params.dbias_ptr = dbias_ptr;
|
119
|
+
// All stride are in elements, not bytes.
|
120
|
+
params.dout_batch_stride = dout.stride(0);
|
121
|
+
params.dout_c_stride = dout.stride(1);
|
122
|
+
params.dout_l_stride = dout.stride(2);
|
123
|
+
params.dweight_c_stride = dweight.stride(0);
|
124
|
+
params.dweight_width_stride = dweight.stride(1);
|
125
|
+
params.dx_batch_stride = dx.stride(0);
|
126
|
+
params.dx_c_stride = dx.stride(1);
|
127
|
+
params.dx_l_stride = dx.stride(2);
|
128
|
+
}
|
129
|
+
|
130
|
+
void
|
131
|
+
causal_conv1d_fwd(const at::Tensor &x,
|
132
|
+
const at::Tensor &weight,
|
133
|
+
const c10::optional<at::Tensor> &bias_,
|
134
|
+
const c10::optional<at::Tensor> &seq_idx_,
|
135
|
+
const c10::optional<at::Tensor> &initial_states_,
|
136
|
+
at::Tensor &out,
|
137
|
+
c10::optional<at::Tensor> &final_states_out_,
|
138
|
+
bool silu_activation) {
|
139
|
+
auto input_type = x.scalar_type();
|
140
|
+
auto weight_type = weight.scalar_type();
|
141
|
+
TORCH_CHECK(input_type == at::ScalarType::Float || input_type == at::ScalarType::Half || input_type == at::ScalarType::BFloat16);
|
142
|
+
TORCH_CHECK(weight_type == at::ScalarType::Float || weight_type == at::ScalarType::Half || weight_type == at::ScalarType::BFloat16);
|
143
|
+
|
144
|
+
TORCH_CHECK(x.is_cuda());
|
145
|
+
TORCH_CHECK(weight.is_cuda());
|
146
|
+
|
147
|
+
const auto sizes = x.sizes();
|
148
|
+
const int batch_size = sizes[0];
|
149
|
+
const int dim = sizes[1];
|
150
|
+
const int seqlen = sizes[2];
|
151
|
+
const int width = weight.size(-1);
|
152
|
+
|
153
|
+
CHECK_SHAPE(x, batch_size, dim, seqlen);
|
154
|
+
CHECK_SHAPE(weight, dim, width);
|
155
|
+
|
156
|
+
TORCH_CHECK(x.stride(2) == 1 || x.stride(1) == 1);
|
157
|
+
const bool is_channel_last = x.stride(1) == 1 && x.stride(2) > 1;
|
158
|
+
|
159
|
+
if (is_channel_last) {
|
160
|
+
TORCH_CHECK(dim % 8 == 0, "causal_conv1d only supports channel dimension divisible by 8 for now");
|
161
|
+
TORCH_CHECK(x.stride(2) % 8 == 0 and x.stride(0) % 8 == 0, "causal_conv1d with channel last layout requires strides (x.stride(0) and x.stride(2)) to be multiples of 8");
|
162
|
+
}
|
163
|
+
TORCH_CHECK(width >= 2 && width <= 4, "causal_conv1d only supports width between 2 and 4");
|
164
|
+
|
165
|
+
if (bias_.has_value()) {
|
166
|
+
auto bias = bias_.value();
|
167
|
+
TORCH_CHECK(bias.scalar_type() == weight_type);
|
168
|
+
TORCH_CHECK(bias.is_cuda());
|
169
|
+
TORCH_CHECK(bias.stride(-1) == 1);
|
170
|
+
CHECK_SHAPE(bias, dim);
|
171
|
+
}
|
172
|
+
|
173
|
+
if (seq_idx_.has_value()) {
|
174
|
+
TORCH_CHECK(is_channel_last, "seq_idx is only supported for channel last layout");
|
175
|
+
auto seq_idx = seq_idx_.value();
|
176
|
+
TORCH_CHECK(seq_idx.scalar_type() == torch::kInt32);
|
177
|
+
TORCH_CHECK(seq_idx.is_cuda());
|
178
|
+
TORCH_CHECK(seq_idx.is_contiguous());
|
179
|
+
CHECK_SHAPE(seq_idx, batch_size, seqlen);
|
180
|
+
}
|
181
|
+
|
182
|
+
ConvParamsBase params;
|
183
|
+
set_conv_params_fwd(params, batch_size, dim, seqlen, width, x, weight, out,
|
184
|
+
bias_.has_value() ? bias_.value().data_ptr() : nullptr,
|
185
|
+
silu_activation);
|
186
|
+
|
187
|
+
if (seq_idx_.has_value()) {
|
188
|
+
params.seq_idx_ptr = seq_idx_.value().data_ptr();
|
189
|
+
} else {
|
190
|
+
params.seq_idx_ptr = nullptr;
|
191
|
+
}
|
192
|
+
|
193
|
+
if (initial_states_.has_value()) {
|
194
|
+
TORCH_CHECK(is_channel_last, "initial_states is only supported for channel last layout");
|
195
|
+
auto initial_states = initial_states_.value();
|
196
|
+
TORCH_CHECK(initial_states.scalar_type() == input_type);
|
197
|
+
TORCH_CHECK(initial_states.is_cuda());
|
198
|
+
CHECK_SHAPE(initial_states, batch_size, dim, width - 1);
|
199
|
+
TORCH_CHECK(initial_states.stride(1) == 1);
|
200
|
+
params.initial_states_ptr = initial_states.data_ptr();
|
201
|
+
params.initial_states_batch_stride = initial_states.stride(0);
|
202
|
+
params.initial_states_c_stride = initial_states.stride(1);
|
203
|
+
params.initial_states_l_stride = initial_states.stride(2);
|
204
|
+
} else {
|
205
|
+
params.initial_states_ptr = nullptr;
|
206
|
+
}
|
207
|
+
|
208
|
+
if (final_states_out_.has_value()) {
|
209
|
+
TORCH_CHECK(is_channel_last, "final_states is only supported for channel last layout");
|
210
|
+
auto final_states = final_states_out_.value();
|
211
|
+
TORCH_CHECK(final_states.scalar_type() == input_type);
|
212
|
+
TORCH_CHECK(final_states.is_cuda());
|
213
|
+
CHECK_SHAPE(final_states, batch_size, dim, width - 1);
|
214
|
+
TORCH_CHECK(final_states.stride(1) == 1);
|
215
|
+
params.final_states_ptr = final_states.data_ptr();
|
216
|
+
params.final_states_batch_stride = final_states.stride(0);
|
217
|
+
params.final_states_c_stride = final_states.stride(1);
|
218
|
+
params.final_states_l_stride = final_states.stride(2);
|
219
|
+
} else {
|
220
|
+
params.final_states_ptr = nullptr;
|
221
|
+
}
|
222
|
+
|
223
|
+
// Otherwise the kernel will be launched from cuda:0 device
|
224
|
+
at::cuda::CUDAGuard device_guard{x.device()};
|
225
|
+
auto stream = at::cuda::getCurrentCUDAStream().stream();
|
226
|
+
DISPATCH_ITYPE_FLOAT_AND_HALF_AND_BF16(x.scalar_type(), "causal_conv1d_fwd", [&] {
|
227
|
+
DISPATCH_WTYPE_FLOAT_AND_HALF_AND_BF16(weight.scalar_type(), "causal_conv1d_fwd", [&] {
|
228
|
+
if (!is_channel_last) {
|
229
|
+
causal_conv1d_fwd_cuda<input_t, weight_t>(params, stream);
|
230
|
+
} else {
|
231
|
+
causal_conv1d_channellast_fwd_cuda<input_t, weight_t>(params, stream);
|
232
|
+
}
|
233
|
+
});
|
234
|
+
});
|
235
|
+
}
|
236
|
+
|
237
|
+
void
|
238
|
+
causal_conv1d_bwd(const at::Tensor &x,
|
239
|
+
const at::Tensor &weight,
|
240
|
+
const c10::optional<at::Tensor> &bias_,
|
241
|
+
at::Tensor &dout,
|
242
|
+
const c10::optional<at::Tensor> &seq_idx_,
|
243
|
+
const c10::optional<at::Tensor> &initial_states_,
|
244
|
+
const c10::optional<at::Tensor> &dfinal_states_,
|
245
|
+
at::Tensor &dx,
|
246
|
+
at::Tensor &dweight,
|
247
|
+
c10::optional<at::Tensor> &dbias_,
|
248
|
+
c10::optional<at::Tensor> &dinitial_states_,
|
249
|
+
bool silu_activation) {
|
250
|
+
auto input_type = x.scalar_type();
|
251
|
+
auto weight_type = weight.scalar_type();
|
252
|
+
TORCH_CHECK(input_type == at::ScalarType::Float || input_type == at::ScalarType::Half || input_type == at::ScalarType::BFloat16);
|
253
|
+
TORCH_CHECK(weight_type == at::ScalarType::Float || weight_type == at::ScalarType::Half || weight_type == at::ScalarType::BFloat16);
|
254
|
+
|
255
|
+
TORCH_CHECK(x.is_cuda());
|
256
|
+
TORCH_CHECK(weight.is_cuda());
|
257
|
+
TORCH_CHECK(dout.is_cuda());
|
258
|
+
TORCH_CHECK(bias_.has_value() == dbias_.has_value());
|
259
|
+
|
260
|
+
const auto sizes = x.sizes();
|
261
|
+
const int batch_size = sizes[0];
|
262
|
+
const int dim = sizes[1];
|
263
|
+
const int seqlen = sizes[2];
|
264
|
+
const int width = weight.size(-1);
|
265
|
+
|
266
|
+
TORCH_CHECK(width >= 2 && width <= 4, "causal_conv1d only supports width between 2 and 4");
|
267
|
+
|
268
|
+
CHECK_SHAPE(x, batch_size, dim, seqlen);
|
269
|
+
CHECK_SHAPE(weight, dim, width);
|
270
|
+
CHECK_SHAPE(dout, batch_size, dim, seqlen);
|
271
|
+
|
272
|
+
TORCH_CHECK(x.stride(2) == 1 || x.stride(1) == 1);
|
273
|
+
const bool is_channel_last = x.stride(1) == 1 && x.stride(2) > 1;
|
274
|
+
if (!is_channel_last && dout.stride(2) != 1) { dout = dout.contiguous(); }
|
275
|
+
if (is_channel_last && dout.stride(1) != 1) { dout = dout.transpose(-1, -2).contiguous().transpose(-1, -2); }
|
276
|
+
|
277
|
+
if (is_channel_last) {
|
278
|
+
TORCH_CHECK(dim % 8 == 0, "causal_conv1d only supports channel dimension divisible by 8 for now");
|
279
|
+
TORCH_CHECK(x.stride(2) % 8 == 0 and x.stride(0) % 8 == 0, "causal_conv1d with channel last layout requires strides (x.stride(0) and x.stride(2)) to be multiples of 8");
|
280
|
+
TORCH_CHECK(dout.stride(2) % 8 == 0 and dout.stride(0) % 8 == 0, "causal_conv1d with channel last layout requires strides (dout.stride(0) and dout.stride(2)) to be multiples of 8");
|
281
|
+
}
|
282
|
+
|
283
|
+
if (bias_.has_value()) {
|
284
|
+
auto bias = bias_.value();
|
285
|
+
TORCH_CHECK(bias.scalar_type() == weight_type);
|
286
|
+
TORCH_CHECK(bias.is_cuda());
|
287
|
+
TORCH_CHECK(bias.stride(-1) == 1);
|
288
|
+
CHECK_SHAPE(bias, dim);
|
289
|
+
}
|
290
|
+
|
291
|
+
if (seq_idx_.has_value()) {
|
292
|
+
TORCH_CHECK(is_channel_last, "seq_idx only supported for channel last layout");
|
293
|
+
auto seq_idx = seq_idx_.value();
|
294
|
+
TORCH_CHECK(seq_idx.scalar_type() == torch::kInt32);
|
295
|
+
TORCH_CHECK(seq_idx.is_cuda());
|
296
|
+
TORCH_CHECK(seq_idx.is_contiguous());
|
297
|
+
CHECK_SHAPE(seq_idx, batch_size, seqlen);
|
298
|
+
}
|
299
|
+
|
300
|
+
TORCH_CHECK(dx.scalar_type() == input_type);
|
301
|
+
TORCH_CHECK(dx.is_cuda());
|
302
|
+
CHECK_SHAPE(dx, batch_size, dim, seqlen);
|
303
|
+
if (!is_channel_last) { TORCH_CHECK(dx.stride(2) == 1); }
|
304
|
+
if (is_channel_last) { TORCH_CHECK(dx.stride(1) == 1); }
|
305
|
+
|
306
|
+
// Otherwise the kernel will be launched from cuda:0 device
|
307
|
+
at::cuda::CUDAGuard device_guard{x.device()};
|
308
|
+
|
309
|
+
ConvParamsBwd params;
|
310
|
+
set_conv_params_bwd(params, batch_size, dim, seqlen, width,
|
311
|
+
x, weight, bias_.has_value() ? bias_.value().data_ptr() : nullptr,
|
312
|
+
dout, dx, dweight, bias_.has_value() ? dbias_.value().data_ptr() : nullptr,
|
313
|
+
silu_activation);
|
314
|
+
|
315
|
+
if (seq_idx_.has_value()) {
|
316
|
+
params.seq_idx_ptr = seq_idx_.value().data_ptr();
|
317
|
+
} else {
|
318
|
+
params.seq_idx_ptr = nullptr;
|
319
|
+
}
|
320
|
+
|
321
|
+
if (initial_states_.has_value()) {
|
322
|
+
TORCH_CHECK(is_channel_last, "initial_states is only supported for channel last layout");
|
323
|
+
auto initial_states = initial_states_.value();
|
324
|
+
TORCH_CHECK(initial_states.scalar_type() == input_type);
|
325
|
+
TORCH_CHECK(initial_states.is_cuda());
|
326
|
+
CHECK_SHAPE(initial_states, batch_size, dim, width - 1);
|
327
|
+
TORCH_CHECK(initial_states.stride(1) == 1);
|
328
|
+
params.initial_states_ptr = initial_states.data_ptr();
|
329
|
+
params.initial_states_batch_stride = initial_states.stride(0);
|
330
|
+
params.initial_states_c_stride = initial_states.stride(1);
|
331
|
+
params.initial_states_l_stride = initial_states.stride(2);
|
332
|
+
} else {
|
333
|
+
params.initial_states_ptr = nullptr;
|
334
|
+
}
|
335
|
+
|
336
|
+
if (dfinal_states_.has_value()) {
|
337
|
+
TORCH_CHECK(is_channel_last, "dfinal_states is only supported for channel last layout");
|
338
|
+
auto dfinal_states = dfinal_states_.value();
|
339
|
+
TORCH_CHECK(dfinal_states.scalar_type() == input_type);
|
340
|
+
TORCH_CHECK(dfinal_states.is_cuda());
|
341
|
+
CHECK_SHAPE(dfinal_states, batch_size, dim, width - 1);
|
342
|
+
params.dfinal_states_ptr = dfinal_states.data_ptr();
|
343
|
+
params.dfinal_states_batch_stride = dfinal_states.stride(0);
|
344
|
+
params.dfinal_states_c_stride = dfinal_states.stride(1);
|
345
|
+
params.dfinal_states_l_stride = dfinal_states.stride(2);
|
346
|
+
} else {
|
347
|
+
params.dfinal_states_ptr = nullptr;
|
348
|
+
}
|
349
|
+
|
350
|
+
if (dinitial_states_.has_value()) {
|
351
|
+
at::Tensor dinitial_states = dinitial_states_.value();
|
352
|
+
TORCH_CHECK(dinitial_states.stride(1) == 1);
|
353
|
+
params.dinitial_states_ptr = dinitial_states.data_ptr();
|
354
|
+
params.dinitial_states_batch_stride = dinitial_states.stride(0);
|
355
|
+
params.dinitial_states_c_stride = dinitial_states.stride(1);
|
356
|
+
params.dinitial_states_l_stride = dinitial_states.stride(2);
|
357
|
+
} else {
|
358
|
+
params.dinitial_states_ptr = nullptr;
|
359
|
+
}
|
360
|
+
|
361
|
+
auto stream = at::cuda::getCurrentCUDAStream().stream();
|
362
|
+
DISPATCH_ITYPE_FLOAT_AND_HALF_AND_BF16(x.scalar_type(), "causal_conv1d_bwd", [&] {
|
363
|
+
DISPATCH_WTYPE_FLOAT_AND_HALF_AND_BF16(weight.scalar_type(), "causal_conv1d_bwd", [&] {
|
364
|
+
if (!is_channel_last) {
|
365
|
+
causal_conv1d_bwd_cuda<input_t, weight_t>(params, stream);
|
366
|
+
} else {
|
367
|
+
causal_conv1d_channellast_bwd_cuda<input_t, weight_t>(params, stream);
|
368
|
+
}
|
369
|
+
});
|
370
|
+
});
|
371
|
+
}
|
372
|
+
|
373
|
+
void
|
374
|
+
causal_conv1d_update(const at::Tensor &x,
|
375
|
+
const at::Tensor &conv_state,
|
376
|
+
const at::Tensor &weight,
|
377
|
+
const c10::optional<at::Tensor> &bias_,
|
378
|
+
at::Tensor &out,
|
379
|
+
bool silu_activation,
|
380
|
+
const c10::optional<at::Tensor> &cache_seqlens_,
|
381
|
+
const c10::optional<at::Tensor> &conv_state_indices_
|
382
|
+
) {
|
383
|
+
auto input_type = x.scalar_type();
|
384
|
+
auto weight_type = weight.scalar_type();
|
385
|
+
TORCH_CHECK(input_type == at::ScalarType::Float || input_type == at::ScalarType::Half || input_type == at::ScalarType::BFloat16);
|
386
|
+
TORCH_CHECK(weight_type == at::ScalarType::Float || weight_type == at::ScalarType::Half || weight_type == at::ScalarType::BFloat16);
|
387
|
+
TORCH_CHECK(conv_state.scalar_type() == input_type);
|
388
|
+
|
389
|
+
TORCH_CHECK(x.is_cuda());
|
390
|
+
TORCH_CHECK(conv_state.is_cuda());
|
391
|
+
TORCH_CHECK(weight.is_cuda());
|
392
|
+
|
393
|
+
const auto sizes = x.sizes();
|
394
|
+
const int batch_size = sizes[0];
|
395
|
+
const int dim = sizes[1];
|
396
|
+
const int seqlen = sizes[2];
|
397
|
+
const int width = weight.size(-1);
|
398
|
+
const int conv_state_len = conv_state.size(2);
|
399
|
+
TORCH_CHECK(conv_state_len >= width - 1);
|
400
|
+
|
401
|
+
CHECK_SHAPE(x, batch_size, dim, seqlen);
|
402
|
+
CHECK_SHAPE(weight, dim, width);
|
403
|
+
|
404
|
+
TORCH_CHECK(width >= 2 && width <= 4, "causal_conv1d only supports width between 2 and 4");
|
405
|
+
|
406
|
+
if (bias_.has_value()) {
|
407
|
+
auto bias = bias_.value();
|
408
|
+
TORCH_CHECK(bias.scalar_type() == weight_type);
|
409
|
+
TORCH_CHECK(bias.is_cuda());
|
410
|
+
TORCH_CHECK(bias.stride(-1) == 1);
|
411
|
+
CHECK_SHAPE(bias, dim);
|
412
|
+
}
|
413
|
+
|
414
|
+
ConvParamsBase params;
|
415
|
+
set_conv_params_fwd(params, batch_size, dim, seqlen, width, x, weight, out,
|
416
|
+
bias_.has_value() ? bias_.value().data_ptr() : nullptr,
|
417
|
+
silu_activation);
|
418
|
+
params.conv_state_ptr = conv_state.data_ptr();
|
419
|
+
params.conv_state_len = conv_state_len;
|
420
|
+
// All stride are in elements, not bytes.
|
421
|
+
params.conv_state_batch_stride = conv_state.stride(0);
|
422
|
+
params.conv_state_c_stride = conv_state.stride(1);
|
423
|
+
params.conv_state_l_stride = conv_state.stride(2);
|
424
|
+
|
425
|
+
if (conv_state_indices_.has_value()) {
|
426
|
+
auto conv_state_indices = conv_state_indices_.value();
|
427
|
+
TORCH_CHECK(conv_state_indices.scalar_type() == torch::kInt32)
|
428
|
+
TORCH_CHECK(conv_state_indices.is_cuda());
|
429
|
+
TORCH_CHECK(conv_state_indices.stride(0) == 1)
|
430
|
+
CHECK_SHAPE(conv_state_indices, batch_size);
|
431
|
+
|
432
|
+
int conv_state_entries = conv_state.size(0);
|
433
|
+
CHECK_SHAPE(conv_state, conv_state_entries, dim, conv_state_len);
|
434
|
+
|
435
|
+
params.conv_state_indices_ptr = conv_state_indices.data_ptr<int32_t>();
|
436
|
+
} else {
|
437
|
+
CHECK_SHAPE(conv_state, batch_size, dim, conv_state_len);
|
438
|
+
params.conv_state_indices_ptr = nullptr;
|
439
|
+
}
|
440
|
+
|
441
|
+
if (cache_seqlens_.has_value()) {
|
442
|
+
auto cache_seqlens = cache_seqlens_.value();
|
443
|
+
TORCH_CHECK(cache_seqlens.scalar_type() == torch::kInt32);
|
444
|
+
TORCH_CHECK(cache_seqlens.is_cuda());
|
445
|
+
TORCH_CHECK(cache_seqlens.stride(-1) == 1);
|
446
|
+
CHECK_SHAPE(cache_seqlens, batch_size);
|
447
|
+
params.cache_seqlens = cache_seqlens.data_ptr<int32_t>();
|
448
|
+
} else {
|
449
|
+
params.cache_seqlens = nullptr;
|
450
|
+
}
|
451
|
+
|
452
|
+
// Otherwise the kernel will be launched from cuda:0 device
|
453
|
+
at::cuda::CUDAGuard device_guard{x.device()};
|
454
|
+
auto stream = at::cuda::getCurrentCUDAStream().stream();
|
455
|
+
DISPATCH_ITYPE_FLOAT_AND_HALF_AND_BF16(x.scalar_type(), "causal_conv1d_update", [&] {
|
456
|
+
DISPATCH_WTYPE_FLOAT_AND_HALF_AND_BF16(weight.scalar_type(), "causal_conv1d_update", [&] {
|
457
|
+
causal_conv1d_update_cuda<input_t, weight_t>(params, stream);
|
458
|
+
});
|
459
|
+
});
|
460
|
+
}
|
461
|
+
|
462
|
+
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
|
463
|
+
m.def("causal_conv1d_fwd", &causal_conv1d_fwd, "Causal conv1d forward");
|
464
|
+
m.def("causal_conv1d_bwd", &causal_conv1d_bwd, "Causal conv1d backward");
|
465
|
+
m.def("causal_conv1d_update", &causal_conv1d_update, "Causal conv1d update");
|
466
|
+
}
|
causal_conv1d-1.5.1/MANIFEST.in
DELETED
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|
File without changes
|