Fix local attention (#102)

* fix local attention

Signed-off-by: Lucas Wilkinson <[email protected]>

* use tot_seqlen_k

Signed-off-by: Lucas Wilkinson <[email protected]>

* more stable fix

Signed-off-by: Lucas Wilkinson <[email protected]>

* better error message

Signed-off-by: Lucas Wilkinson <[email protected]>

* make consistent with og dcp pr

Signed-off-by: Lucas Wilkinson <[email protected]>

---------

Signed-off-by: Lucas Wilkinson <[email protected]>

Author: LucasWilkinson

hopper/block.h

@@ -40,11 +40,12 @@ struct BlockMN {
             // If local, blocking (m_idx_max - m_idx_min + window_size_right + window_size_left)  
             // when cp is not enabled, tot_seqlen_k is equal to seqlen_k, and cp_world_size is 1.
             // cp_world_size is guaranteed to be greater than 0
-            n_block_max = std::min(n_block_max,
-                                    cute::ceil_div(
-                                    cute::ceil_div(m_idx_max + seqlen_info.tot_seqlen_k - seqlen_q + window_size_right - seqlen_info.cp_rank,
-                                                  seqlen_info.cp_world_size),
-                                    kBlockN));
+            int tot_seqlen_k = (Is_local) ? seqlen_k : seqlen_info.tot_seqlen_k;
+            int n_token_max = m_idx_max + tot_seqlen_k - seqlen_q + window_size_right;
+            if (seqlen_info.cp_world_size > 1 && !Is_local) {
+                n_token_max = cute::ceil_div(n_token_max - seqlen_info.cp_rank, seqlen_info.cp_world_size);
+            }
+            n_block_max = std::min(n_block_max, cute::ceil_div(n_token_max, kBlockN));
         }
         // Now, only adjust n_block_min if split
         int n_block_min = 0;

hopper/flash_api.cpp

@@ -1170,6 +1170,9 @@ mha_fwd(at::Tensor &q,   // (b, s_q, h, d) or (total_q, h, d) if there is cu_seq
     TORCH_CHECK(cp_world_size > 0, "cp_world_size must be positive, required by downstream unified code path. Use 1 if CP is not enabled.");
     TORCH_CHECK(cp_world_size != 1 || cp_rank == 0, "When context parallelism is disabled, cp_rank must be zero");
     TORCH_CHECK(cp_world_size == 1 || cp_tot_seqused_k_.has_value(), "cp_tot_seqused_k_ must be provided when context parallelism is enabled.");
+    TORCH_CHECK(!(params.is_local && cp_world_size > 1), 
+        "Local attention (sliding window) is not currently supported with context parallelism (cp_world_size > 1)."
+        "Requires proper n_offset handling in block boundary calculations in mainloop and block.h");
 
     #ifdef FLASHATTENTION_DISABLE_LOCAL
     TORCH_CHECK(!params.is_local, "This flash attention build does not support local attention.");