优化pa_opt的编译速度

csrc/attention/attention_kernels_opt.cu

@@ -861,7 +861,7 @@ void paged_attention_v1_launcher(
           dim3 block(NUM_THREADS);
           const at::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(device_of(query));
           const hipStream_t stream = at::hip::getCurrentHIPStreamMasqueradingAsCUDA();
-          LAUNCH_PAGED_ATTENTION_V1(HEAD_SIZE);
+          if constexpr(BLOCK_SIZE==16 && IS_BLOCK_SPARSE==false && sizeof(T)==2){LAUNCH_PAGED_ATTENTION_V1(HEAD_SIZE);}
           });
         });
       });
@@ -1006,7 +1006,7 @@ void paged_attention_v2_launcher(
         dim3 block(NUM_THREADS);
         const at::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(device_of(query));
         const hipStream_t stream = at::hip::getCurrentHIPStreamMasqueradingAsCUDA();
-        LAUNCH_PAGED_ATTENTION_V2(HEAD_SIZE);
+        if constexpr(BLOCK_SIZE==16 && IS_BLOCK_SPARSE==false && sizeof(T)==2){LAUNCH_PAGED_ATTENTION_V2(HEAD_SIZE);}
         });
       });
     });

csrc/attention/attention_with_mask_kernels_opt.cu

@@ -883,7 +883,7 @@ void paged_attention_v1_launcher(
           dim3 block(NUM_THREADS);
           const at::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(device_of(query));
           const hipStream_t stream = at::hip::getCurrentHIPStreamMasqueradingAsCUDA();
-          LAUNCH_PAGED_ATTENTION_V1(HEAD_SIZE);
+          if constexpr(BLOCK_SIZE==16 && IS_BLOCK_SPARSE==false && sizeof(T)==2){LAUNCH_PAGED_ATTENTION_V1(HEAD_SIZE);}
           });
         });
       });
@@ -1041,7 +1041,7 @@ void paged_attention_v2_launcher(
         dim3 block(NUM_THREADS);
         const at::hip::OptionalHIPGuardMasqueradingAsCUDA device_guard(device_of(query));
         const hipStream_t stream = at::hip::getCurrentHIPStreamMasqueradingAsCUDA();
-        LAUNCH_PAGED_ATTENTION_V2(HEAD_SIZE);
+        if constexpr(BLOCK_SIZE==16 && IS_BLOCK_SPARSE==false && sizeof(T)==2){LAUNCH_PAGED_ATTENTION_V2(HEAD_SIZE);}
         });
       });
     });