Merge pull request #104 from ROCm/page-group-update

Fix mha_varlen_fwd num_split and change ck interface
ROCm · Nov 28, 2024 · 58f6f67 · 58f6f67
2 parents 0538b43 + d2dff5b
commit 58f6f67
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 2 deletions.
diff --git a/csrc/composable_kernel b/csrc/composable_kernel
diff --git a/csrc/flash_attn_ck/mha_varlen_fwd.cpp b/csrc/flash_attn_ck/mha_varlen_fwd.cpp
@@ -217,6 +217,7 @@ fmha_fwd_splitkv_args get_ck_fmha_varlen_fwd_splitkv_args(bool has_lse,
         args.page_block_size = 0;
     }
 
+    args.is_gappy = false;
     args.cache_batch_idx = nullptr;
 
     args.seqstart_q_ptr = seqlens_q.data_ptr();
@@ -442,7 +443,7 @@ mha_varlen_fwd(at::Tensor &q,                   // total_q x num_heads x head_si
         if (return_dropout_randval) {p.zero_();}
     }
 
-    int num_splits = 1;
+    int num_splits = 0;
     num_splits = flash::override_num_splits_if_necessary(batch_size, num_heads, max_seqlen_q, head_size, 0, num_splits);
     TORCH_CHECK(num_splits > 0, "num_splits should greater than 0");
     TORCH_CHECK(num_splits <= 128, "num_splits greater than 128 is not supported");