remove smem implementation because const vals and bi-search is enough

Author: zigzagcai

csrc/selective_scan/selective_scan_bwd_kernel.cuh

@@ -143,14 +143,8 @@ void selective_scan_bwd_kernel(SSMParamsBwd params) {
     float dD_val = 0;
     float ddelta_bias_val = 0;
 
-    // Load cu_seqlens into shared memory
     const int cu_seqlens_size = params.cu_seqlens_size;
-    long *cu_seqlens = reinterpret_cast<long *>(params.cu_seqlens_ptr);
-    __shared__ long smem_cu_seqlens[1024];  // Adjust size as needed
-    for (int i = threadIdx.x; i < cu_seqlens_size; i += blockDim.x) {
-        smem_cu_seqlens[i] = cu_seqlens[i];
-    }
-    __syncthreads();
+    const long *cu_seqlens = reinterpret_cast<long *>(params.cu_seqlens_ptr);
 
     constexpr int kChunkSize = kNThreads * kNItems;
     u += (params.n_chunks - 1) * kChunkSize;
@@ -267,10 +261,10 @@ void selective_scan_bwd_kernel(SSMParamsBwd params) {
                     int idx = threadIdx.x * kNItems + i + chunk * kChunkSize;
                     while (left <= right) {
                         int mid = (left + right) >> 1;
-                        if (smem_cu_seqlens[mid] == idx) {
+                        if (cu_seqlens[mid] == idx) {
                             delta_a_exp = 0.f;
                             break;
-                        } else if (smem_cu_seqlens[mid] < idx) {
+                        } else if (cu_seqlens[mid] < idx) {
                             left = mid + 1;
                         } else {
                             right = mid - 1;
@@ -372,11 +366,11 @@ void selective_scan_bwd_kernel(SSMParamsBwd params) {
                     int idx = threadIdx.x * kNItems + i + chunk * kChunkSize;
                     while (left <= right) {
                         int mid = (left + right) >> 1;
-                        if (smem_cu_seqlens[mid] == idx) {
+                        if (cu_seqlens[mid] == idx) {
                             delta_a_exp.real_ = 0.f;
                             delta_a_exp.imag_ = 0.f;
                             break;
-                        } else if (smem_cu_seqlens[mid] < idx) {
+                        } else if (cu_seqlens[mid] < idx) {
                             left = mid + 1;
                         } else {
                             right = mid - 1;

csrc/selective_scan/selective_scan_fwd_kernel.cuh

@@ -113,15 +113,8 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {
     input_t *Cvar = reinterpret_cast<input_t *>(params.C_ptr) + batch_id * params.C_batch_stride + group_id * params.C_group_stride;
     scan_t *x = reinterpret_cast<scan_t *>(params.x_ptr) + (batch_id * params.dim + dim_id * kNRows) * params.n_chunks * params.dstate;
 
-    // Load cu_seqlens into shared memory
     const int cu_seqlens_size = params.cu_seqlens_size;
-    long *cu_seqlens = reinterpret_cast<long *>(params.cu_seqlens_ptr);
-    __shared__ long smem_cu_seqlens[1024];  // Adjust size as needed
-    for (int i = threadIdx.x; i < cu_seqlens_size; i += blockDim.x) {
-        smem_cu_seqlens[i] = cu_seqlens[i];
-    }
-    __syncthreads();
-
+    const long *cu_seqlens = reinterpret_cast<long *>(params.cu_seqlens_ptr);
 
     float D_val[kNRows] = {0};
     if (params.D_ptr != nullptr) {
@@ -237,10 +230,10 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {
                         int idx = threadIdx.x * kNItems + i + chunk * kChunkSize;
                         while (left <= right) {
                             int mid = (left + right) >> 1;
-                            if (smem_cu_seqlens[mid] == idx) {
+                            if (cu_seqlens[mid] == idx) {
                                 thread_data[i].x = 0.f;
                                 break;
-                            } else if (smem_cu_seqlens[mid] < idx) {
+                            } else if (cu_seqlens[mid] < idx) {
                                 left = mid + 1;
                             } else {
                                 right = mid - 1;
@@ -264,11 +257,11 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {
                         int idx = threadIdx.x * kNItems + i + chunk * kChunkSize;
                         while (left <= right) {
                             int mid = (left + right) >> 1;
-                            if (smem_cu_seqlens[mid] == idx) {
+                            if (cu_seqlens[mid] == idx) {
                                 thread_data[i].x = 0.f;
                                 thread_data[i].y = 0.f;
                                 break;
-                            } else if (smem_cu_seqlens[mid] < idx) {
+                            } else if (cu_seqlens[mid] < idx) {
                                 left = mid + 1;
                             } else {
                                 right = mid - 1;