perf(tree_attn): Add lazy attention mask creation for tree training

areal/engine/fsdp_engine.py

@@ -70,7 +70,11 @@
     gather_packed_tree_logprobs_entropy,
     merge_packed_tree_results,
 )
-from areal.models.tree_attn.module import BLOCK_SIZE, patch_fsdp_for_tree_training
+from areal.models.tree_attn.module import (
+    BLOCK_SIZE,
+    build_block_mask_from_trie,
+    patch_fsdp_for_tree_training,
+)
 from areal.models.tree_attn.tree import TrieNode, build_packed_tree_batch
 from areal.platforms import current_platform
 from areal.utils import (
@@ -541,10 +545,26 @@ def forward_backward_batch(
         for mb_item in mb_list:
             inputs, ctx = self._prepare_mb_inputs(mb_item)
 
+            # Lazily create block mask for tree training just before forward
+            if self.enable_tree_training and ctx.trie_node is not None:
+                padded_size = mb_item.padded_to_length
+                if padded_size is None:
+                    raise ValueError(
+                        "padded_size must be set for tree training with FSDP."
+                    )
+                block_mask = build_block_mask_from_trie(
+                    ctx.trie_node, padded_size, self.device
+                )
+                inputs["block_mask"] = block_mask
+
             with trace_scope("fsdp_engine.forward"):
                 outputs = self.model(**inputs)
             logits = outputs.logits.squeeze(0)
 
+            # Release block mask memory after forward pass
+            if self.enable_tree_training and "block_mask" in inputs:
+                del inputs["block_mask"]
+
             ctx_dict = ctx.to_dict()
             loss = process_output_fn(logits, ctx_dict)
 

areal/engine/megatron_engine.py

@@ -57,7 +57,11 @@
     gather_packed_tree_logprobs_entropy,
     merge_packed_tree_results,
 )
-from areal.models.tree_attn.module import BLOCK_SIZE, patch_bridge_for_tree_training
+from areal.models.tree_attn.module import (
+    BLOCK_SIZE,
+    build_block_mask_from_trie,
+    patch_bridge_for_tree_training,
+)
 from areal.models.tree_attn.tree import build_packed_tree_batch
 from areal.platforms import current_platform
 from areal.utils import logging, name_resolve, names, perf_tracer, stats_tracker
@@ -565,8 +569,23 @@ def forward_step(batch_iter, model):
             mb_input: MicroBatchItem = next(batch_iter)
 
             cu_seqlens = mb_input.padded_mb.get("cu_seqlens", None)
+
+            # Lazily create block mask for tree training just before forward
+            if self.enable_tree_training:
+                trie_node = mb_input.padded_mb.get("trie_node", None)
+                padded_size = mb_input.padded_to_length
+                if trie_node is not None and padded_size is not None:
+                    block_mask = build_block_mask_from_trie(
+                        trie_node, padded_size, self.device
+                    )
+                    mb_input.padded_mb["block_mask"] = block_mask
+
             output = packed_context_parallel_forward(model, mb_input.padded_mb)
 
+            # Release block mask memory after forward pass
+            if self.enable_tree_training and "block_mask" in mb_input.padded_mb:
+                del mb_input.padded_mb["block_mask"]
+
             def _process_output(input_, output_):
                 loss = process_output_fn(output_, input_)
                 if loss is None:

areal/models/tree_attn/module.py

@@ -4,6 +4,7 @@
     patch_fsdp_for_tree_training,
     restore_patch_fsdp_for_tree_training,
 )
+from areal.models.tree_attn.tree import build_block_mask_from_trie
 
 # Conditionally import Megatron functionality
 try:
@@ -22,6 +23,7 @@
     "create_block_mask_from_dense",
     "patch_fsdp_for_tree_training",
     "restore_patch_fsdp_for_tree_training",
+    "build_block_mask_from_trie",
     # Megatron exports (may be None if Megatron not installed)
     "PytorchFlexAttention",
     "patch_bridge_for_tree_training",

areal/models/tree_attn/tree.py

@@ -12,6 +12,7 @@
 
 import torch
 import torch.distributed as dist
+from torch.nn.attention.flex_attention import BlockMask
 
 from areal.api.cli_args import MicroBatchSpec
 from areal.models.tree_attn.constants import BLOCK_SIZE
@@ -370,12 +371,6 @@ def build_packed_tree_batch(
                 mask_template.device,
             )
 
-        # Create block mask from dense mask
-        with trace_scope("tree_attn.create_block_mask"):
-            block_mask = create_block_mask_from_dense(
-                attention_mask, padded_size, mask_template.device
-            )
-
         # Compute position_ids (needs dense attention_mask)
         with trace_scope("tree_attn.get_position_ids"):
             position_ids = get_packed_tree_position_ids(
@@ -384,6 +379,7 @@ def build_packed_tree_batch(
             )
 
         # Release dense attention mask memory after position_ids are computed
+        # Block mask will be lazily created in forward_backward_batch
         del attention_mask
 
         # Pack extra data
@@ -396,10 +392,9 @@ def build_packed_tree_batch(
                 non_packable_keys,
             )
 
-        # Build micro-batch dict with block_mask
+        # Build micro-batch dict without block_mask (will be created lazily in forward)
         mb = {
             "input_ids": input_ids,
-            "block_mask": block_mask,
             "position_ids": position_ids,
             "trie_node": trie,
             **extra_data,
@@ -611,3 +606,49 @@ def get_packed_tree_position_ids(
         position_ids = torch.clamp_min(ancestor_counts - 1, 0)
 
     return position_ids.unsqueeze(0)
+
+
+@trace_perf("tree_attn.build_block_mask_from_trie")
+def build_block_mask_from_trie(
+    trie: TrieNode,
+    padded_size: int,
+    device: torch.device,
+) -> BlockMask:
+    """Lazily build a block mask from a trie node.
+
+    This function builds the dense attention mask from the trie structure and
+    converts it to a block mask for use with flex attention. It should be called
+    just before the forward pass to minimize memory usage.
+
+    Parameters
+    ----------
+    trie : TrieNode
+        The root trie node containing the tree structure.
+    padded_size : int
+        The padded sequence length.
+    device : torch.device
+        Device to create the block mask on.
+
+    Returns
+    -------
+    BlockMask
+        The created block mask for use with flex_attention.
+    """
+    # Handle dummy trie (empty tree for DP synchronization)
+    if not trie.all_sequence_ids:
+        # Create a minimal valid block mask for empty trees
+        dummy_mask = torch.zeros(
+            (padded_size, padded_size), dtype=torch.bool, device=device
+        )
+        return create_block_mask_from_dense(dummy_mask, padded_size, device)
+
+    with trace_scope("tree_attn.build_attention_mask"):
+        attention_mask = _build_attention_mask(trie, padded_size, device)
+
+    with trace_scope("tree_attn.create_block_mask"):
+        block_mask = create_block_mask_from_dense(attention_mask, padded_size, device)
+
+    # Release dense attention mask memory
+    del attention_mask
+
+    return block_mask

areal/utils/data.py

@@ -372,12 +372,14 @@ class MicroBatchItem(NamedTuple):
         padded_mb: Padded micro-batch dict (for model forward)
         padding_length: Batch-level padding added to this micro-batch
         old_cu_seqlens: Original cu_seqlens before sequence alignment (or None)
+        padded_to_length: The padded sequence length for this micro-batch (or None)
     """
 
     orig_mb: dict[str, Any]
     padded_mb: dict[str, Any]
     padding_length: int
     old_cu_seqlens: torch.Tensor | None
+    padded_to_length: int | None = None
 
 
 @dataclass
@@ -421,18 +423,23 @@ def __iter__(self) -> Iterator[MicroBatchItem]:
                 - padded_mb: Padded micro-batch dict (for model forward)
                 - padding_length: Batch-level padding added to this micro-batch
                 - old_cu_seqlens: Original cu_seqlens before sequence alignment (or None)
+                - padded_to_length: The padded sequence length for this micro-batch (or None)
         """
         if self.padded_mbs is None:
             raise ValueError("padded_mbs is None. Call pad_mb_list first.")
         for i in range(len(self.mbs)):
             old_cu_seqlens = (
                 self.old_cu_seqlens_list[i] if self.old_cu_seqlens_list else None
             )
+            padded_to_length = (
+                self.padded_to_lengths[i] if self.padded_to_lengths else None
+            )
             yield MicroBatchItem(
                 orig_mb=self.mbs[i],
                 padded_mb=self.padded_mbs[i],
                 padding_length=self.padding_lengths[i],
                 old_cu_seqlens=old_cu_seqlens,
+                padded_to_length=padded_to_length,
             )
 
     def to(self, *args, **kwargs):