Allow causal scheduler with mask_mod (#155)

* Allow causal scheduler with mask_mod

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

* Pass leading dims for aux tensors

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

* Fix

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

---------

Signed-off-by: Matthew Bonanni <mbonanni@redhat.com>

Author: MatthewBonanni

csrc/composable_kernel

@@ -103,13 +103,14 @@ def to_cute_tensor(t, assumed_align=16, leading_dim=-1, fully_dynamic=False, ena
     return tensor.mark_layout_dynamic(leading_dim=leading_dim)
 
 
-def to_cute_aux_tensor(t, enable_tvm_ffi=True):
+def to_cute_aux_tensor(t, leading_dim=None, enable_tvm_ffi=True):
     """Convert torch tensor to cute tensor for TVM FFI, tailored to FlexAttention aux tensors.
     This allows the user to specify alignment and leading dimension for aux tensors used in
     custom score_mod callables.
     """
     assumed_align: int = getattr(t, "__assumed_align__", None)
-    leading_dim: int = getattr(t, "__leading_dim__", None)
+    if leading_dim is None:
+        leading_dim = getattr(t, "__leading_dim__", None)
     fully_dynamic: bool = leading_dim is None
 
     return to_cute_tensor(
@@ -121,14 +122,28 @@ def to_cute_aux_tensor(t, enable_tvm_ffi=True):
     )
 
 
-def get_aux_tensor_metadata(aux_tensors):
+def get_aux_tensor_metadata(aux_tensors, aux_tensor_leading_dims=None):
+    if aux_tensor_leading_dims is not None:
+        assert len(aux_tensor_leading_dims) == len(aux_tensors)
+
     return tuple(
         (
             getattr(t, "__assumed_align__", 0),
-            getattr(t, "__leading_dim__", -1),
-            hasattr(t, "__leading_dim__"),
+            (
+                leading_dim
+                if leading_dim is not None
+                else getattr(t, "__leading_dim__", -1)
+            ),
+            leading_dim is not None or hasattr(t, "__leading_dim__"),
+        )
+        for t, leading_dim in zip(
+            aux_tensors,
+            (
+                aux_tensor_leading_dims
+                if aux_tensor_leading_dims is not None
+                else [None] * len(aux_tensors)
+            ),
         )
-        for t in aux_tensors
     )
 
 

flash_attn/cute/cute_dsl_utils.py

@@ -345,8 +345,6 @@ def _resolve_causal_local_window(causal, window_size_left, window_size_right, ma
 
     Returns (causal, local, window_size_left, window_size_right).
     """
-    if mask_mod is not None:
-        return False, False, window_size_left, window_size_right
     if causal:
         window_size_right = 0
     if window_size_left is not None and window_size_right is not None and window_size_left + window_size_right < 0:
@@ -396,6 +394,7 @@ def _flash_attn_fwd(
     out: Optional[torch.Tensor] = None,
     lse: Optional[torch.Tensor] = None,
     aux_tensors: Optional[list[torch.Tensor]] = None,
+    aux_tensor_leading_dims: Optional[list[int]] = None,
     aux_scalars: Optional[tuple] = None,
     q_descale: Optional[torch.Tensor] = None,
     k_descale: Optional[torch.Tensor] = None,
@@ -768,8 +767,11 @@ def _flash_attn_fwd(
             q_stage=q_stage,
         )
     if aux_tensors is not None:
-        aux_tensor_metadata = get_aux_tensor_metadata(aux_tensors)
+        aux_tensor_metadata = get_aux_tensor_metadata(
+            aux_tensors, aux_tensor_leading_dims
+        )
     else:
+        assert aux_tensor_leading_dims is None
         aux_tensor_metadata = None
     aux_scalar_metadata = tuple(type(s) for s in aux_scalars) if aux_scalars is not None else None
 
@@ -931,7 +933,15 @@ def _flash_attn_fwd(
         cute_aux_tensors = None
         aux_tensor_metadata = None
         if aux_tensors is not None:
-            cute_aux_tensors = [to_cute_aux_tensor(buf) for buf in aux_tensors]
+            aux_tensor_leading_dims = (
+                aux_tensor_leading_dims
+                if aux_tensor_leading_dims is not None
+                else [None] * len(aux_tensors)
+            )
+            cute_aux_tensors = [
+                to_cute_aux_tensor(buf, leading_dim)
+                for buf, leading_dim in zip(aux_tensors, aux_tensor_leading_dims)
+            ]
 
         qv_tensor = to_cute_tensor(qv)
         gather_kv_indices_tensor = to_cute_tensor(gather_kv_indices)

flash_attn/cute/interface.py

@@ -731,6 +731,41 @@ def apply_mask_sm100(
                 )
 
         else:  # Causal or local
+            if const_expr(mask_mod is not None):
+                assert vec_size % 32 == 0 or 32 % vec_size == 0, (
+                    "vec_size must divide 32 or be a multiple of 32"
+                )
+                if const_expr(vec_size == 1):
+                    self.apply_mask_mod_sm100_scalar(
+                        acc_S,
+                        tScS_t2r,
+                        m_block,
+                        n_block,
+                        mask_seqlen,
+                        mask_mod,
+                        batch_idx,
+                        head_idx,
+                        aux_data,
+                        fastdiv_mods,
+                        head_divmod,
+                        check_q_boundary,
+                    )
+                else:
+                    self.apply_mask_mod_sm100_vector(
+                        acc_S,
+                        tScS_t2r,
+                        m_block,
+                        n_block,
+                        mask_seqlen,
+                        mask_mod,
+                        batch_idx,
+                        head_idx,
+                        vec_size,
+                        aux_data,
+                        fastdiv_mods,
+                        head_divmod,
+                        check_q_boundary,
+                    )
             causal_row_offset = self.seqlen_k - n_block * self.tile_n - self.seqlen_q
             row_idx = tScS_t2r[0][0] + m_block * self.tile_m
             if const_expr(self.qhead_per_kvhead_packgqa != 1):