Switch llama3 and qwen3 model configs from sdpa/causal to flex/block_causal

With packed datasets, sdpa/causal allows cross-document attention
leakage and uses sequential positions across document boundaries.
flex/block_causal isolates documents in attention and enables
per-document RoPE position IDs.

llama4, deepseek_v3, and gpt_oss already used flex/block_causal.

Author: joecummings

torchtitan/models/llama3/__init__.py

@@ -44,7 +44,8 @@
             ),
             attention=GQAttention.Config(
                 n_heads=16,
-                attn_backend="sdpa",
+                attn_backend="flex",
+                attn_mask_type="block_causal",
                 rope_backend="complex",
             ),
         ),
@@ -57,6 +58,7 @@
             scaling="llama",
         ),
     ),
+    # TODO: now identical to "debugmodel", can be removed
     "debugmodel_flex_attn": Llama3Model.Config(
         dim=256,
         n_layers=6,
@@ -130,7 +132,8 @@
             attention=GQAttention.Config(
                 n_heads=32,
                 n_kv_heads=8,
-                attn_backend="sdpa",
+                attn_backend="flex",
+                attn_mask_type="block_causal",
                 rope_backend="complex",
             ),
         ),
@@ -142,6 +145,7 @@
             scaling="llama",
         ),
     ),
+    # TODO: now identical to "8B", can be removed
     "8B_flex": Llama3Model.Config(
         dim=4096,
         n_layers=32,
@@ -219,7 +223,8 @@
             attention=GQAttention.Config(
                 n_heads=64,
                 n_kv_heads=8,
-                attn_backend="sdpa",
+                attn_backend="flex",
+                attn_mask_type="block_causal",
                 rope_backend="complex",
             ),
         ),
@@ -248,7 +253,8 @@
             attention=GQAttention.Config(
                 n_heads=128,
                 n_kv_heads=8,
-                attn_backend="sdpa",
+                attn_backend="flex",
+                attn_mask_type="block_causal",
                 rope_backend="complex",
             ),
         ),

torchtitan/models/qwen3/__init__.py

@@ -46,7 +46,8 @@
                 head_dim=128,
                 q_norm=RMSNorm.Config(eps=1e-6),
                 k_norm=RMSNorm.Config(eps=1e-6),
-                attn_backend="sdpa",
+                attn_backend="flex",
+                attn_mask_type="block_causal",
                 rope_backend="cos_sin",
             ),
         ),
@@ -57,6 +58,7 @@
             backend="cos_sin",
         ),
     ),
+    # TODO: now identical to "debugmodel", can be removed
     "debugmodel_flex": Qwen3Model.Config(
         vocab_size=2048,
         dim=256,
@@ -76,6 +78,7 @@
                 q_norm=RMSNorm.Config(eps=1e-6),
                 k_norm=RMSNorm.Config(eps=1e-6),
                 attn_backend="flex",
+                attn_mask_type="block_causal",
                 rope_backend="cos_sin",
             ),
         ),
@@ -106,7 +109,8 @@
                 head_dim=128,
                 q_norm=RMSNorm.Config(eps=1e-6),
                 k_norm=RMSNorm.Config(eps=1e-6),
-                attn_backend="sdpa",
+                attn_backend="flex",
+                attn_mask_type="block_causal",
                 rope_backend="cos_sin",
             ),
         ),
@@ -137,7 +141,8 @@
                 head_dim=128,
                 q_norm=RMSNorm.Config(eps=1e-6),
                 k_norm=RMSNorm.Config(eps=1e-6),
-                attn_backend="sdpa",
+                attn_backend="flex",
+                attn_mask_type="block_causal",
                 rope_backend="cos_sin",
             ),
         ),
@@ -168,7 +173,8 @@
                 head_dim=128,
                 q_norm=RMSNorm.Config(eps=1e-6),
                 k_norm=RMSNorm.Config(eps=1e-6),
-                attn_backend="sdpa",
+                attn_backend="flex",
+                attn_mask_type="block_causal",
                 rope_backend="cos_sin",
             ),
         ),
@@ -198,7 +204,8 @@
                 head_dim=128,
                 q_norm=RMSNorm.Config(eps=1e-6),
                 k_norm=RMSNorm.Config(eps=1e-6),
-                attn_backend="sdpa",
+                attn_backend="flex",
+                attn_mask_type="block_causal",
                 rope_backend="cos_sin",
             ),
         ),
@@ -228,7 +235,8 @@
                 head_dim=128,
                 q_norm=RMSNorm.Config(eps=1e-6),
                 k_norm=RMSNorm.Config(eps=1e-6),
-                attn_backend="sdpa",
+                attn_backend="flex",
+                attn_mask_type="block_causal",
                 rope_backend="cos_sin",
             ),
         ),
@@ -258,7 +266,8 @@
                 head_dim=128,
                 q_norm=RMSNorm.Config(eps=1e-6),
                 k_norm=RMSNorm.Config(eps=1e-6),
-                attn_backend="sdpa",
+                attn_backend="flex",
+                attn_mask_type="block_causal",
                 rope_backend="cos_sin",
             ),
         ),
@@ -300,7 +309,8 @@
                 head_dim=128,
                 q_norm=RMSNorm.Config(eps=1e-6),
                 k_norm=RMSNorm.Config(eps=1e-6),
-                attn_backend="sdpa",
+                attn_backend="flex",
+                attn_mask_type="block_causal",
                 rope_backend="cos_sin",
             ),
         ),
@@ -341,7 +351,8 @@
                 head_dim=128,
                 q_norm=RMSNorm.Config(eps=1e-6),
                 k_norm=RMSNorm.Config(eps=1e-6),
-                attn_backend="sdpa",
+                attn_backend="flex",
+                attn_mask_type="block_causal",
                 rope_backend="cos_sin",
             ),
         ),
@@ -382,7 +393,8 @@
                 head_dim=128,
                 q_norm=RMSNorm.Config(eps=1e-6),
                 k_norm=RMSNorm.Config(eps=1e-6),
-                attn_backend="sdpa",
+                attn_backend="flex",
+                attn_mask_type="block_causal",
                 rope_backend="cos_sin",
             ),
         ),