[attention] Make the FlexAttention mask type configurable

tests/unit_tests/cpu/test_flex_mask_type.py

@@ -0,0 +1,119 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+
+import torch
+from torch.nn.attention.flex_attention import and_masks
+from torchtitan.models.common.attention import (
+    BaseAttention,
+    FlexInnerAttention,
+    get_causal_mask_mod,
+    get_document_mask_mod,
+    get_flex_causal_mask_mods,
+)
+from torchtitan.models.common.decoder import Decoder
+from torchtitan.models.gpt_oss import build_model_config as build_gpt_oss_config
+from torchtitan.models.gpt_oss.model import GptOssModel
+from torchtitan.models.muse_glimmer import (
+    build_model_config as build_muse_glimmer_config,
+)
+from torchtitan.models.muse_glimmer.model import MuseGlimmerModel
+
+PACKED_POSITIONS = torch.tensor(
+    [0, 1, 2, 0, 1, 0, 1, 2, 0, 1, 0, 1, 2, 3, 0, 1], dtype=torch.int32
+)
+SINGLE_DOC_POSITIONS = torch.arange(16, dtype=torch.int32)
+BLOCK_SIZE = 4
+
+
+def _dense(mask_mod, seq_len: int) -> torch.Tensor:
+    b, h = torch.tensor(0), torch.tensor(0)
+    return torch.tensor(
+        [
+            [
+                bool(mask_mod(b, h, torch.tensor(q), torch.tensor(kv)))
+                for kv in range(seq_len)
+            ]
+            for q in range(seq_len)
+        ]
+    )
+
+
+def _decoder_masks(flex_kwargs: dict) -> dict:
+    attn_config = BaseAttention.Config(
+        n_heads=1,
+        inner_attention=FlexInnerAttention.Config(block_size=BLOCK_SIZE, **flex_kwargs),
+    )
+    decoder = Decoder.__new__(Decoder)
+    return {
+        "mask": decoder._create_flex_attention_mask_for_document(
+            PACKED_POSITIONS, attn_config
+        )
+    }
+
+
+def _model_masks(build_config, model_cls):
+    def masks(flex_kwargs: dict) -> dict:
+        config = build_config("debugmodel", seq_len=16, attn_backend="flex")
+        for layer in config.layers:
+            inner = layer.attention.inner_attention
+            assert isinstance(inner, FlexInnerAttention.Config)
+            inner.block_size = BLOCK_SIZE
+            for key, value in flex_kwargs.items():
+                setattr(inner, key, value)
+        model = model_cls.__new__(model_cls)
+        model.config = config
+        return model.get_attention_masks(PACKED_POSITIONS)
+
+    return masks
+
+
+class TestFlexMaskType(unittest.TestCase):
+    def test_mask_mods_match_mask_type(self):
+        for positions in (PACKED_POSITIONS, SINGLE_DOC_POSITIONS):
+            seq_len = positions.shape[0]
+            causal = _dense(get_causal_mask_mod(), seq_len)
+            same_doc = _dense(get_document_mask_mod(positions), seq_len)
+            for mask_type, expected in (
+                ("block_causal", causal & same_doc),
+                ("causal", causal),
+            ):
+                with self.subTest(
+                    single_doc=positions is SINGLE_DOC_POSITIONS, mask_type=mask_type
+                ):
+                    mask_mods = get_flex_causal_mask_mods(
+                        positions, FlexInnerAttention.Config(mask_type=mask_type)
+                    )
+                    torch.testing.assert_close(
+                        _dense(and_masks(*mask_mods), seq_len), expected
+                    )
+
+    def test_mask_builders_follow_mask_type(self):
+        num_blocks = PACKED_POSITIONS.shape[0] // BLOCK_SIZE
+        causal_blocks = torch.ones(num_blocks, num_blocks, dtype=torch.bool).tril()
+        builders = {
+            "decoder": _decoder_masks,
+            "gpt_oss": _model_masks(build_gpt_oss_config, GptOssModel),
+            "muse_glimmer": _model_masks(build_muse_glimmer_config, MuseGlimmerModel),
+        }
+        for name, build_masks in builders.items():
+            default_masks = build_masks({})
+            causal_masks = build_masks({"mask_type": "causal"})
+            for key, causal_mask in causal_masks.items():
+                with self.subTest(builder=name, mask=key):
+                    torch.testing.assert_close(
+                        causal_mask.to_dense()[0, 0].bool(), causal_blocks
+                    )
+                    # Default (block_causal) drops blocks in earlier documents.
+                    self.assertLess(
+                        default_masks[key].to_dense().sum().item(),
+                        causal_mask.to_dense().sum().item(),
+                    )
+
+
+if __name__ == "__main__":
+    unittest.main()

torchtitan/models/common/__init__.py

@@ -22,6 +22,7 @@
     get_document_mask_mod,
     get_efficient_causal_mask_mod_for_packed_document,
     get_fixed_block_mask_mod,
+    get_flex_causal_mask_mods,
     get_sliding_window_mask_mod,
     GQAttention,
     InnerAttention,
@@ -68,6 +69,7 @@
     "get_document_mask_mod",
     "get_efficient_causal_mask_mod_for_packed_document",
     "get_fixed_block_mask_mod",
+    "get_flex_causal_mask_mods",
     "get_sliding_window_mask_mod",
     "GQAttention",
     "GroupNorm",

torchtitan/models/common/attention.py

@@ -14,7 +14,7 @@
 
 from collections.abc import Callable, Mapping
 from dataclasses import dataclass, field
-from typing import Any, ClassVar, NamedTuple
+from typing import Any, ClassVar, Literal, NamedTuple
 
 import spmd_types as spmd
 import torch
@@ -63,6 +63,7 @@
     "get_document_mask_mod",
     "get_efficient_causal_mask_mod_for_packed_document",
     "get_fixed_block_mask_mod",
+    "get_flex_causal_mask_mods",
     "get_sliding_window_mask_mod",
     "local_head_split",
 ]
@@ -258,6 +259,11 @@ class FlexInnerAttention(InnerAttention):
     class Config(InnerAttention.Config):
         block_size: int | tuple[int, int] = _DEFAULT_SPARSE_BLOCK_SIZE
         kernel_options: dict = field(default_factory=dict)
+        mask_type: Literal["block_causal", "causal"] = "block_causal"
+        """Causal mask for the flex path. ``"block_causal"``: causal within each
+        packed document (positions reset to 0 at document starts). ``"causal"``:
+        plain causal over the whole sequence, which does not depend on
+        ``positions`` and is only correct when each sequence holds one document."""
 
     inductor_configs: ClassVar[dict[str, bool]] = {
         "wrap_inductor_compiled_regions": True,
@@ -499,6 +505,17 @@ def document_mask(
     return document_mask
 
 
+def get_flex_causal_mask_mods(
+    positions: torch.Tensor,
+    config: FlexInnerAttention.Config,
+) -> list[_mask_mod_signature]:
+    """Return the causal mask_mods selected by ``config.mask_type``."""
+    mask_mods = [get_causal_mask_mod()]
+    if config.mask_type == "block_causal":
+        mask_mods.append(get_efficient_causal_mask_mod_for_packed_document(positions))
+    return mask_mods
+
+
 def get_efficient_causal_mask_mod_for_packed_document(
     positions: torch.Tensor,
 ) -> _mask_mod_signature:

torchtitan/models/common/decoder.py

@@ -23,8 +23,7 @@
     create_attention_mask,
     create_varlen_metadata_for_document,
     FlexInnerAttention,
-    get_causal_mask_mod,
-    get_efficient_causal_mask_mod_for_packed_document,
+    get_flex_causal_mask_mods,
     VarlenInnerAttention,
 )
 from torchtitan.models.common.decoder_sharding import decoder_input_sharding
@@ -292,14 +291,12 @@ def _create_flex_attention_mask_for_document(
         positions: torch.Tensor,
         attn_config: BaseAttention.Config,
     ) -> BlockMask:
-        """Build the standard causal + packed-document flex-attention mask."""
+        """Build the causal flex-attention mask selected by ``mask_type``."""
+        assert isinstance(attn_config.inner_attention, FlexInnerAttention.Config)
         return self._create_flex_attention_mask(
             positions,
             attn_config,
-            [
-                get_causal_mask_mod(),
-                get_efficient_causal_mask_mod_for_packed_document(positions),
-            ],
+            get_flex_causal_mask_mods(positions, attn_config.inner_attention),
         )
 
     def preprocess_inputs(

torchtitan/models/gpt_oss/model.py

@@ -23,8 +23,7 @@
     BaseAttention,
     create_varlen_metadata_for_document,
     FlexInnerAttention,
-    get_causal_mask_mod,
-    get_efficient_causal_mask_mod_for_packed_document,
+    get_flex_causal_mask_mods,
     get_sliding_window_mask_mod,
     QKVLinear,
     VarlenInnerAttention,
@@ -292,12 +291,9 @@ def get_attention_masks(
                 max_context_length=max_context_length,
             )
         elif isinstance(inner_attn, FlexInnerAttention.Config):
-            base_mask_mods = [
-                get_causal_mask_mod(),
-                get_efficient_causal_mask_mod_for_packed_document(positions),
-            ]
-            # Full-attention (causal + document) mask, used by layers without a
-            # sliding window.
+            base_mask_mods = get_flex_causal_mask_mods(positions, inner_attn)
+            # Full-attention (causal) mask, used by layers without a sliding
+            # window.
             masks: dict[str, BlockMask] = {
                 "basic_mask": self._create_flex_attention_mask(
                     positions, attn_cfg, base_mask_mods

torchtitan/models/muse_glimmer/model.py

@@ -31,8 +31,7 @@
     create_attention_mask,
     create_varlen_metadata_for_document,
     FlexInnerAttention,
-    get_causal_mask_mod,
-    get_efficient_causal_mask_mod_for_packed_document,
+    get_flex_causal_mask_mods,
     get_sliding_window_mask_mod,
     GQAttention,
     VarlenInnerAttention,
@@ -627,14 +626,10 @@ def get_attention_masks(
                 f"sliding-window masks, got {type(inner_attn).__name__}"
             )
 
-        # Language models always use block-causal (per-document) masking: the
-        # dataloaders emit per-document positions, and the efficient packed-doc
-        # mask ANDed with the causal mask yields same-document causal attention.
+        # ``mask_type="block_causal"`` (default) yields same-document causal
+        # attention from the per-document positions the dataloaders emit.
         seq_len = positions.shape[0]
-        base_mods = [
-            get_causal_mask_mod(),
-            get_efficient_causal_mask_mod_for_packed_document(positions),
-        ]
+        base_mods = get_flex_causal_mask_mods(positions, inner_attn)
 
         # Match the base Decoder mask-building so the configured block size and
         # batch-invariance handling are honored.