fix skipahead compile error on older cuda versions

facebookresearch · Dec 20, 2022 · b9337bd · b9337bd
1 parent d1a2103
commit b9337bd
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/xformers/components/attention/csrc/cuda/mem_eff_attention/kernel_forward.h b/xformers/components/attention/csrc/cuda/mem_eff_attention/kernel_forward.h
@@ -764,9 +764,10 @@ struct AttentionKernel {
         if (thread_i < problem_size_0_m && thread_start_j < problem_size_0_n) {
           curandStatePhilox4_32_10_t curand_state = curand_state_init;
           skipahead(
-              (query_start + thread_i) * p.num_keys + (iter_key_start + thread_start_j),
+              static_cast<unsigned long long>(
+                  (query_start + thread_i) * p.num_keys +
+                  (iter_key_start + thread_start_j)),
               &curand_state);
-
           const float dropout_scale = 1.0 / (1.0 - p.dropout_prob);
 
           // apply dropout scaling to elements this thread is responsible for,