add feature / fix bug: I fixed the kNRows feature in forward

.gitignore

@@ -1,4 +1,5 @@
 *__pycache__/
 *.egg-info/
 build/
-**.so
+*.so
+*.whl

csrc/selective_scan/cus/selective_scan_bwd.cu

@@ -0,0 +1,11 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+// Split into multiple files to compile in paralell
+
+#include "../selective_scan_bwd_kernel.cuh"
+
+template void selective_scan_bwd_cuda<1, float, float>(SSMParamsBwd &params, cudaStream_t stream);
+template void selective_scan_bwd_cuda<1, at::Half, float>(SSMParamsBwd &params, cudaStream_t stream);
+template void selective_scan_bwd_cuda<1, at::BFloat16, float>(SSMParamsBwd &params, cudaStream_t stream);

csrc/selective_scan/cus/selective_scan_bwd2.cu

@@ -0,0 +1,11 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+// Split into multiple files to compile in paralell
+
+#include "../selective_scan_bwd_kernel.cuh"
+
+template void selective_scan_bwd_cuda<2, float, float>(SSMParamsBwd &params, cudaStream_t stream);
+template void selective_scan_bwd_cuda<2, at::Half, float>(SSMParamsBwd &params, cudaStream_t stream);
+template void selective_scan_bwd_cuda<2, at::BFloat16, float>(SSMParamsBwd &params, cudaStream_t stream);

csrc/selective_scan/cus/selective_scan_bwd2_complex.cu

@@ -0,0 +1,11 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+// Split into multiple files to compile in paralell
+
+#include "../selective_scan_bwd_kernel.cuh"
+
+template void selective_scan_bwd_cuda<2, float, complex_t>(SSMParamsBwd &params, cudaStream_t stream);
+template void selective_scan_bwd_cuda<2, at::Half, complex_t>(SSMParamsBwd &params, cudaStream_t stream);
+template void selective_scan_bwd_cuda<2, at::BFloat16, complex_t>(SSMParamsBwd &params, cudaStream_t stream);

csrc/selective_scan/cus/selective_scan_bwd3.cu

@@ -0,0 +1,11 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+// Split into multiple files to compile in paralell
+
+#include "../selective_scan_bwd_kernel.cuh"
+
+template void selective_scan_bwd_cuda<3, float, float>(SSMParamsBwd &params, cudaStream_t stream);
+template void selective_scan_bwd_cuda<3, at::Half, float>(SSMParamsBwd &params, cudaStream_t stream);
+template void selective_scan_bwd_cuda<3, at::BFloat16, float>(SSMParamsBwd &params, cudaStream_t stream);

csrc/selective_scan/cus/selective_scan_bwd3_complex.cu

@@ -0,0 +1,11 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+// Split into multiple files to compile in paralell
+
+#include "../selective_scan_bwd_kernel.cuh"
+
+template void selective_scan_bwd_cuda<3, float, complex_t>(SSMParamsBwd &params, cudaStream_t stream);
+template void selective_scan_bwd_cuda<3, at::Half, complex_t>(SSMParamsBwd &params, cudaStream_t stream);
+template void selective_scan_bwd_cuda<3, at::BFloat16, complex_t>(SSMParamsBwd &params, cudaStream_t stream);

csrc/selective_scan/cus/selective_scan_bwd4.cu

@@ -0,0 +1,11 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+// Split into multiple files to compile in paralell
+
+#include "../selective_scan_bwd_kernel.cuh"
+
+template void selective_scan_bwd_cuda<4, float, float>(SSMParamsBwd &params, cudaStream_t stream);
+template void selective_scan_bwd_cuda<4, at::Half, float>(SSMParamsBwd &params, cudaStream_t stream);
+template void selective_scan_bwd_cuda<4, at::BFloat16, float>(SSMParamsBwd &params, cudaStream_t stream);

csrc/selective_scan/cus/selective_scan_bwd4_complex.cu

@@ -0,0 +1,11 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+// Split into multiple files to compile in paralell
+
+#include "../selective_scan_bwd_kernel.cuh"
+
+template void selective_scan_bwd_cuda<4, float, complex_t>(SSMParamsBwd &params, cudaStream_t stream);
+template void selective_scan_bwd_cuda<4, at::Half, complex_t>(SSMParamsBwd &params, cudaStream_t stream);
+template void selective_scan_bwd_cuda<4, at::BFloat16, complex_t>(SSMParamsBwd &params, cudaStream_t stream);

csrc/selective_scan/cus/selective_scan_bwd_complex.cu

@@ -0,0 +1,11 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+// Split into multiple files to compile in paralell
+
+#include "../selective_scan_bwd_kernel.cuh"
+
+template void selective_scan_bwd_cuda<1, float, complex_t>(SSMParamsBwd &params, cudaStream_t stream);
+template void selective_scan_bwd_cuda<1, at::Half, complex_t>(SSMParamsBwd &params, cudaStream_t stream);
+template void selective_scan_bwd_cuda<1, at::BFloat16, complex_t>(SSMParamsBwd &params, cudaStream_t stream);

csrc/selective_scan/cus/selective_scan_fwd.cu

@@ -0,0 +1,11 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+// Split into multiple files to compile in paralell
+
+#include "../selective_scan_fwd_kernel.cuh"
+
+template void selective_scan_fwd_cuda<1, float, float>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<1, at::BFloat16, float>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<1, at::Half, float>(SSMParamsBase &params, cudaStream_t stream);

csrc/selective_scan/cus/selective_scan_fwd2.cu

@@ -0,0 +1,11 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+// Split into multiple files to compile in paralell
+
+#include "../selective_scan_fwd_kernel.cuh"
+
+template void selective_scan_fwd_cuda<2, at::BFloat16, float>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<2, at::Half, float>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<2, float, float>(SSMParamsBase &params, cudaStream_t stream);

csrc/selective_scan/cus/selective_scan_fwd2_complex.cu

@@ -0,0 +1,11 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+// Split into multiple files to compile in paralell
+
+#include "../selective_scan_fwd_kernel.cuh"
+
+template void selective_scan_fwd_cuda<2, at::BFloat16, complex_t>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<2, at::Half, complex_t>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<2, float, complex_t>(SSMParamsBase &params, cudaStream_t stream);

csrc/selective_scan/cus/selective_scan_fwd3.cu

@@ -0,0 +1,11 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+// Split into multiple files to compile in paralell
+
+#include "../selective_scan_fwd_kernel.cuh"
+
+template void selective_scan_fwd_cuda<3, at::BFloat16, float>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<3, at::Half, float>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<3, float, float>(SSMParamsBase &params, cudaStream_t stream);

csrc/selective_scan/cus/selective_scan_fwd3_complex.cu

@@ -0,0 +1,11 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+// Split into multiple files to compile in paralell
+
+#include "../selective_scan_fwd_kernel.cuh"
+
+template void selective_scan_fwd_cuda<3, at::BFloat16, complex_t>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<3, at::Half, complex_t>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<3, float, complex_t>(SSMParamsBase &params, cudaStream_t stream);

csrc/selective_scan/cus/selective_scan_fwd4.cu

@@ -0,0 +1,11 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+// Split into multiple files to compile in paralell
+
+#include "../selective_scan_fwd_kernel.cuh"
+
+template void selective_scan_fwd_cuda<4, at::BFloat16, float>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<4, at::Half, float>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<4, float, float>(SSMParamsBase &params, cudaStream_t stream);

csrc/selective_scan/cus/selective_scan_fwd4_complex.cu

@@ -0,0 +1,11 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+// Split into multiple files to compile in paralell
+
+#include "../selective_scan_fwd_kernel.cuh"
+
+template void selective_scan_fwd_cuda<4, at::BFloat16, complex_t>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<4, at::Half, complex_t>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<4, float, complex_t>(SSMParamsBase &params, cudaStream_t stream);

csrc/selective_scan/cus/selective_scan_fwd_complex.cu

@@ -0,0 +1,11 @@
+/******************************************************************************
+ * Copyright (c) 2023, Tri Dao.
+ ******************************************************************************/
+
+// Split into multiple files to compile in paralell
+
+#include "../selective_scan_fwd_kernel.cuh"
+
+template void selective_scan_fwd_cuda<1, at::BFloat16, complex_t>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<1, at::Half, complex_t>(SSMParamsBase &params, cudaStream_t stream);
+template void selective_scan_fwd_cuda<1, float, complex_t>(SSMParamsBase &params, cudaStream_t stream);