File tree Expand file tree Collapse file tree 2 files changed +4
-0
lines changed Expand file tree Collapse file tree 2 files changed +4
-0
lines changed Original file line number Diff line number Diff line change @@ -253,6 +253,7 @@ void selective_scan_bwd_kernel(SSMParamsBwd params) {
253
253
while (left <= right) {
254
254
if (cu_seqlens[(left + right) >> 1 ] == threadIdx .x * kNItems + i + chunk * kChunkSize ) {
255
255
delta_a_exp = 0 .f ;
256
+ break ;
256
257
} else if (cu_seqlens[(left + right) >> 1 ] < threadIdx .x * kNItems + i + chunk * kChunkSize ) {
257
258
left = ((left + right) >> 1 ) + 1 ;
258
259
} else {
@@ -356,6 +357,7 @@ void selective_scan_bwd_kernel(SSMParamsBwd params) {
356
357
if (cu_seqlens[(left + right) >> 1 ] == threadIdx .x * kNItems + i + chunk * kChunkSize ) {
357
358
delta_a_exp.real_ = 0 .f ;
358
359
delta_a_exp.imag_ = 0 .f ;
360
+ break ;
359
361
} else if (cu_seqlens[(left + right) >> 1 ] < threadIdx .x * kNItems + i + chunk * kChunkSize ) {
360
362
left = ((left + right) >> 1 ) + 1 ;
361
363
} else {
Original file line number Diff line number Diff line change @@ -223,6 +223,7 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {
223
223
while (left <= right) {
224
224
if (cu_seqlens[(left + right) >> 1 ] == threadIdx .x * kNItems + i + chunk * kChunkSize ) {
225
225
thread_data[i].x = 0 .f ;
226
+ break ;
226
227
} else if (cu_seqlens[(left + right) >> 1 ] < threadIdx .x * kNItems + i + chunk * kChunkSize ) {
227
228
left = ((left + right) >> 1 ) + 1 ;
228
229
} else {
@@ -248,6 +249,7 @@ void selective_scan_fwd_kernel(SSMParamsBase params) {
248
249
if (cu_seqlens[(left + right) >> 1 ] == threadIdx .x * kNItems + i + chunk * kChunkSize ) {
249
250
thread_data[i].x = 0 .f ;
250
251
thread_data[i].y = 0 .f ;
252
+ break ;
251
253
} else if (cu_seqlens[(left + right) >> 1 ] < threadIdx .x * kNItems + i + chunk * kChunkSize ) {
252
254
left = ((left + right) >> 1 ) + 1 ;
253
255
} else {
You can’t perform that action at this time.
0 commit comments