Rebased and changed calls as needed

Author: AdityaKane2001

csrc/include/natten/cuda/flash_fna/flash_kernel/flash.h

@@ -22,7 +22,7 @@ struct NA_params {
   NADim stride;
   NADim dilation;
 
-  int num_heads_actual;
+  int batch_size_actual;
 
 };
 

csrc/include/natten/cuda/flash_fna/flash_kernel/flash_bwd_launch_template.h

@@ -156,7 +156,7 @@ void run_flash_bwd(Flash_fna_bwd_params<NADim> &params, cudaStream_t stream) {
         params.window_size,
         params.stride,
         params.dilation,
-        params.num_heads_actual
+        params.batch_size_actual
     };
     // The case work with GQA is ugly but idk how to fix it.
     typename CollectiveEpilogue::Arguments epilogue_args {

csrc/include/natten/cuda/flash_fna/flash_kernel/flash_fwd_launch_template.h

@@ -93,7 +93,7 @@ void run_flash_fwd(Flash_fna_fwd_params<NADim> &params, cudaStream_t stream) {
         {params.v_descale_batch_stride, params.v_descale_head_stride},
         params.num_splits,
         // NA Args
-        params.qkv_shape, params.q_shape, params.kv_shape, params.window_size, params.stride, params.dilation, params.num_heads_actual
+        params.qkv_shape, params.q_shape, params.kv_shape, params.window_size, params.stride, params.dilation, params.batch_size_actual
     };
     typename CollectiveEpilogue::Arguments epilogue_args {
         static_cast<ElementOut*>(params.o_ptr),

csrc/include/natten/cuda/flash_fna/flash_kernel/mainloop_bwd_sm80.hpp

@@ -323,7 +323,7 @@ struct CollectiveMainloopBwdSm80 {
         NADim window_size;
         NADim stride;
         NADim dilation;
-        int num_heads_actual;
+        int batch_size_actual;
     };
 
     // Device side kernel params
@@ -365,7 +365,7 @@ struct CollectiveMainloopBwdSm80 {
         NADim window_right;
         NADim stride;
         NADim dilation;
-        int num_heads_actual;
+        int batch_size_actual;
         bool is_fully_block_sparse;
         bool has_q_padding;
         bool requires_qkv_fixup;
@@ -420,7 +420,7 @@ struct CollectiveMainloopBwdSm80 {
                 // !Has_softcap ? 0.f : args.softmax_scale / args.softcap_val,
                 args.num_batch, args.dq_semaphore,
                 args.qkv_shape, args.q_shape, args.kv_shape, args.window_size, window_left, window_right,
-                args.stride, args.dilation, args.num_heads_actual,
+                args.stride, args.dilation, args.batch_size_actual,
                 is_fully_block_sparse_, has_q_padding_, requires_qkv_fixup_, is_dilated_
                 // args.cu_seqlens_q, args.cu_seqlens_k, args.seqused_q, args.seqused_k
         };
@@ -445,13 +445,12 @@ struct CollectiveMainloopBwdSm80 {
             // params.cu_seqlens_q, params.cu_seqlens_k, params.seqused_q, params.seqused_k
         };
 
-        int head_idx = bidh;
         auto qkv_shape = params.qkv_shape;
         bool is_fully_block_sparse = params.is_fully_block_sparse;
         bool has_q_padding = params.has_q_padding;
 
         if (params.requires_qkv_fixup) {
-          qkv_shape = correct_qkv_shape(params.qkv_shape, head_idx, params.dilation, params.num_heads_actual);
+          qkv_shape = correct_qkv_shape(params.qkv_shape, bidb, params.dilation, params.batch_size_actual);
           is_fully_block_sparse = fully_block_sparse<Causal>(
               qkv_shape,
               params.window_size,

csrc/include/natten/cuda/flash_fna/flash_kernel/mainloop_fwd_sm80.hpp

@@ -196,7 +196,7 @@ struct CollectiveMainloopFwdSm80 {
         NADim window_size;
         NADim stride;
         NADim dilation;
-        int num_heads_actual;
+        int batch_size_actual;
     };
 
     // Device side kernel params
@@ -227,7 +227,7 @@ struct CollectiveMainloopFwdSm80 {
         NADim window_right;
         NADim stride;
         NADim dilation;
-        int num_heads_actual;
+        int batch_size_actual;
         bool is_fully_block_sparse;
         bool has_kv_padding;
         bool requires_qkv_fixup;
@@ -277,7 +277,7 @@ struct CollectiveMainloopFwdSm80 {
                 args.stride_q_descale, args.stride_k_descale, args.stride_v_descale,
                 1 /* args.num_splits */,
                 args.qkv_shape, args.q_shape, args.kv_shape, args.window_size, window_left, window_right,
-                args.stride, args.dilation, args.num_heads_actual,
+                args.stride, args.dilation, args.batch_size_actual,
                 is_fully_block_sparse, has_kv_padding, requires_qkv_fixup, is_dilated_
                };
     }
@@ -304,12 +304,11 @@ struct CollectiveMainloopFwdSm80 {
         int const split_idx = get<3>(block_coord);
         int const bidh_kv = !PackGQA ? params.qhead_per_khead_divmod.divide(bidh) : bidh;
 
-        int head_idx = bidh;
         auto qkv_shape = params.qkv_shape;
         bool is_fully_block_sparse = params.is_fully_block_sparse;
         bool has_kv_padding = params.has_kv_padding;
         if (params.requires_qkv_fixup) {
-          qkv_shape = correct_qkv_shape(params.qkv_shape, head_idx, params.dilation, params.num_heads_actual);
+          qkv_shape = correct_qkv_shape(params.qkv_shape, bidb, params.dilation, params.batch_size_actual);
           is_fully_block_sparse = fully_block_sparse<Causal>(
               qkv_shape,
               params.window_size,

csrc/include/natten/cuda/flash_fna/flash_kernel/na_utils.h

@@ -434,11 +434,11 @@ CUTLASS_DEVICE auto correct_qkv_shape(
     NADim const& qkv_shape, // this is pre-padding, pre-token permute, just
                                // the original shape of the sequence mode in
                                // the self attention
-    int head_idx,
+    int batch_idx,
     NADim const& dilation,
-    int num_heads_actual) {
+    int batch_size_actual) {
 
-  auto dilation_group_idx = head_idx / num_heads_actual;
+  auto dilation_group_idx = batch_idx % product(dilation);
   auto dilation_group_crd = idx2crd(dilation_group_idx, dilation);
 
   return correct_qkv_shape_wrt_dilation(

csrc/include/natten/cuda/flash_fna/flash_kernel/param_utils.h

@@ -86,7 +86,7 @@ Flash_fna_fwd_params<NADim> set_flash_fna_fwd_params(
   params.window_size = window_size;
   params.stride = stride;
   params.dilation = dilation;
-  params.num_heads_actual = H / product(dilation);
+  params.batch_size_actual = B / product(dilation);
 
   params.is_bf16 = query.scalar_type() == torch::kBFloat16;
   params.is_e4m3 = false;
@@ -283,7 +283,7 @@ Flash_fna_bwd_params<NADim> set_flash_fna_bwd_params(
   params.window_size = window_size;
   params.stride = stride;
   params.dilation = dilation;
-  params.num_heads_actual = H / product(dilation);
+  params.batch_size_actual = B / product(dilation);
 
   params.q_ptr = static_cast<void*>(query.data_ptr());
   params.k_ptr = static_cast<void*>(key.data_ptr());

setup.py

@@ -63,11 +63,9 @@
 AUTOGEN_POLICY = AUTOGEN_POLICY if AUTOGEN_POLICY != "" else "default"
 
 tmp_dir = tempfile.TemporaryDirectory()
-print(f"***************** {tmp_dir=}")
 NATTEN_BUILD_DIR = os.getenv("NATTEN_BUILD_DIR", tmp_dir.name)
 if not os.path.isdir(NATTEN_BUILD_DIR):
     NATTEN_BUILD_DIR = tmp_dir.name
-print(f"***************** {NATTEN_BUILD_DIR=}")
 
 DEFAULT_N_WORKERS = max(1, (multiprocessing.cpu_count() // 4))
 try:

src/natten/_libnatten/__init__.py

@@ -122,6 +122,8 @@
     "blackwell_na3d_backward",
     "blackwell_na3d_forward",
     "compute_delta",
+    "flash_fmha_backward",
+    "flash_fmha_forward",
     "fmha_backward",
     "fmha_forward",
     "hopper_fmha_backward",
@@ -132,6 +134,12 @@
     "hopper_na2d_forward",
     "hopper_na3d_backward",
     "hopper_na3d_forward",
+    "flash_na1d_backward",
+    "flash_na1d_forward",
+    "flash_na2d_backward",
+    "flash_na2d_forward",
+    "flash_na3d_backward",
+    "flash_na3d_forward",
     "na1d_backward",
     "na1d_forward",
     "na2d_backward",