fix

InternLM · Nov 24, 2023 · 295715d · 295715d
1 parent c07f60f
commit 295715d
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/src/turbomind/models/llama/flash_attention2/flash_fwd_launch_template.h b/src/turbomind/models/llama/flash_attention2/flash_fwd_launch_template.h
@@ -14,7 +14,13 @@
 template<typename Kernel_traits, bool Is_dropout, bool Is_causal, bool Is_even_N, bool Is_even_K, bool Return_softmax>
 __global__ void flash_fwd_kernel(Flash_fwd_params params)
 {
+
+#if __CUDA_ARCH__ >= 800
     flash::compute_attn<Kernel_traits, Is_dropout, Is_causal, Is_even_N, Is_even_K, Return_softmax>(params);
+#else
+// TODO: support flash attention2 on sm<80
+    assert(false);
+#endif
 }
 
 template<typename Kernel_traits, bool Is_dropout, bool Is_causal>