..

mosaicml · Dec 4, 2024 · f6c66e8 · f6c66e8
1 parent fef3a5d
commit f6c66e8
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/llmfoundry/models/layers/attention.py b/llmfoundry/models/layers/attention.py
@@ -1034,6 +1034,7 @@ def __init__(
         reuse_kv_layer_idx: Optional[int] = None,
         attn_logit_softcapping: Optional[float] = None,
         kv_dim: Optional[int] = None,
+        flex_attn_mod_list: Optional[list[dict[str, Any]]] = None,
     ):
         super().__init__(
             d_model=d_model,
@@ -1055,6 +1056,7 @@ def __init__(
             reuse_kv_layer_idx=reuse_kv_layer_idx,
             attn_logit_softcapping=attn_logit_softcapping,
             kv_dim=kv_dim,
+            flex_attn_mod_list=flex_attn_mod_list,
         )
 
 
@@ -1085,6 +1087,7 @@ def __init__(
         reuse_kv_layer_idx: Optional[int] = None,
         attn_logit_softcapping: Optional[float] = None,
         kv_dim: Optional[int] = None,
+        flex_attn_mod_list: Optional[list[dict[str, Any]]] = None,
     ):
         super().__init__(
             d_model=d_model,
@@ -1106,6 +1109,7 @@ def __init__(
             reuse_kv_layer_idx=reuse_kv_layer_idx,
             attn_logit_softcapping=attn_logit_softcapping,
             kv_dim=kv_dim,
+            flex_attn_mod_list=flex_attn_mod_list,
         )