only keep hdim128

Signed-off-by: Minmin Sun <[email protected]>

Author: minminsun

csrc/flash_attn/flash_api.cpp

@@ -236,11 +236,11 @@ void run_mha_fwd(Flash_fwd_params &params, cudaStream_t stream, bool force_split
 
 void run_mha_fwd_sparse(Flash_fwd_params &params, cudaStream_t stream, bool force_split_kernel=false) {
     TORCH_CHECK(params.num_splits <= 1 && !force_split_kernel, "run_mha_fwd_sparse does not support splitkv.");
+    TORCH_CHECK(params.d == 128, "run_mha_fwd_sparse only supports headdim=128 for now to keep binary small.");
     FP16_SWITCH(!params.is_bf16, [&] {
-        HEADDIM_SWITCH(params.d, [&] {
-            BOOL_SWITCH(params.is_causal, Is_causal, [&] {
-                run_mha_fwd_sparse_<elem_type, kHeadDim, Is_causal>(params, stream);
-            });
+        constexpr static int kHeadDim = 128;
+        BOOL_SWITCH(params.is_causal, Is_causal, [&] {
+            run_mha_fwd_sparse_<elem_type, kHeadDim, Is_causal>(params, stream);
         });
     });
 }

csrc/flash_attn/src/flash_fwd_sparse_hdim160_bf16_causal_sm80.cu

@@ -270,7 +270,7 @@ def test_varlen_with_paged_kv(
 
 @pytest.mark.parametrize("seq_lens", [(1023, 2049), (1023, 1023), (32, 32), (65, 65), (129, 129)])
 @pytest.mark.parametrize("num_heads", [1, 2, 4])
-@pytest.mark.parametrize("head_size", [64, 128, 256])
+@pytest.mark.parametrize("head_size", [128])
 @pytest.mark.parametrize("dtype", DTYPES)
 @pytest.mark.parametrize("NNZ_S", [1, 2, 3, 7, 15, 32])
 @torch.inference_mode()