feat(tree_attn): implement Triton kernel

areal/engine/fsdp_engine.py

@@ -73,7 +73,10 @@
 )
 from areal.models.tree_attn.module import (
     BLOCK_SIZE,
+    TRITON_AVAILABLE,
+    USE_TRITON_TREE_ATTN,
     build_block_mask_from_trie,
+    build_triton_attn_data_from_trie,
     patch_fsdp_for_tree_training,
 )
 from areal.models.tree_attn.tree import TrieNode, build_packed_tree_batch
@@ -545,25 +548,34 @@ 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
+            # Lazily create tree attention metadata 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
+                if USE_TRITON_TREE_ATTN and TRITON_AVAILABLE:
+                    triton_attn_data = build_triton_attn_data_from_trie(
+                        ctx.trie_node, padded_size
+                    )
+                    inputs["triton_attn_data"] = triton_attn_data
+                else:
+                    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"]
+            # Release tree attention metadata after forward pass
+            if self.enable_tree_training:
+                if "block_mask" in inputs:
+                    del inputs["block_mask"]
+                if "triton_attn_data" in inputs:
+                    del inputs["triton_attn_data"]
 
             ctx_dict = ctx.to_dict()
             loss = process_output_fn(logits, ctx_dict)

areal/engine/megatron_engine.py

@@ -60,7 +60,10 @@
 )
 from areal.models.tree_attn.module import (
     BLOCK_SIZE,
+    TRITON_AVAILABLE,
+    USE_TRITON_TREE_ATTN,
     build_block_mask_from_trie,
+    build_triton_attn_data_from_trie,
     patch_bridge_for_tree_training,
 )
 from areal.models.tree_attn.tree import build_packed_tree_batch
@@ -570,21 +573,32 @@ def forward_step(batch_iter, model):
 
             cu_seqlens = mb_input.padded_mb.get("cu_seqlens", None)
 
-            # Lazily create block mask for tree training just before forward
+            # Lazily create tree attention metadata 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
+                    if USE_TRITON_TREE_ATTN and TRITON_AVAILABLE:
+                        triton_attn_data = build_triton_attn_data_from_trie(
+                            trie_node, padded_size
+                        )
+                        mb_input.padded_mb["triton_attn_data"] = triton_attn_data
+                    else:
+                        block_mask = build_block_mask_from_trie(
+                            trie_node,
+                            padded_size,
+                            mb_input.padded_mb["input_ids"].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"]
+            # Release tree attention metadata after forward pass
+            if self.enable_tree_training:
+                if "block_mask" in mb_input.padded_mb:
+                    del mb_input.padded_mb["block_mask"]
+                if "triton_attn_data" in mb_input.padded_mb:
+                    del mb_input.padded_mb["triton_attn_data"]
 
             def _process_output(input_, output_):
                 loss = process_output_fn(output_, input_)

areal/models/tree_attn/constants.py

@@ -2,4 +2,15 @@
 
 import os
 
+from areal.utils.logging import getLogger
+
+logger = getLogger(__name__)
+
 BLOCK_SIZE = int(os.environ.get("AREAL_FLEX_ATTENTION_BLOCK_SIZE", "128"))
+USE_TRITON_TREE_ATTN = int(os.environ.get("AREAL_USE_TRITON_TREE_ATTN", "0")) == 1
+
+if USE_TRITON_TREE_ATTN:
+    logger.warning(
+        "Triton tree attention kernel is only an experimental feature "
+        "that requires further practical RL experiment testing."
+    )

areal/models/tree_attn/module.py

@@ -1,10 +1,25 @@
-from areal.models.tree_attn.constants import BLOCK_SIZE
+from areal.models.tree_attn.constants import BLOCK_SIZE, USE_TRITON_TREE_ATTN
 from areal.models.tree_attn.module_fsdp import (
     create_block_mask_from_dense,
     patch_fsdp_for_tree_training,
     restore_patch_fsdp_for_tree_training,
 )
-from areal.models.tree_attn.tree import build_block_mask_from_trie
+from areal.models.tree_attn.tree import (
+    build_block_mask_from_trie,
+    build_triton_attn_data_from_trie,
+)
+
+# Conditionally import Triton functionality
+try:
+    from areal.models.tree_attn.triton_kernel import (
+        TRITON_AVAILABLE,
+        TreeAttentionData,
+        tree_attention,
+    )
+except ImportError:
+    TRITON_AVAILABLE = False
+    TreeAttentionData = None
+    tree_attention = None
 
 # Conditionally import Megatron functionality
 try:
@@ -19,11 +34,17 @@
 __all__ = [
     # Shared constants
     "BLOCK_SIZE",
+    "USE_TRITON_TREE_ATTN",
     # FSDP/common exports
     "create_block_mask_from_dense",
     "patch_fsdp_for_tree_training",
     "restore_patch_fsdp_for_tree_training",
     "build_block_mask_from_trie",
+    "build_triton_attn_data_from_trie",
+    # Triton exports (may be None if Triton not installed)
+    "TRITON_AVAILABLE",
+    "TreeAttentionData",
+    "tree_attention",
     # Megatron exports (may be None if Megatron not installed)
     "PytorchFlexAttention",
     "patch_bridge_for_tree_training",

areal/models/tree_attn/module_fsdp.py

@@ -7,7 +7,8 @@
     flex_attention,
 )
 
-from areal.models.tree_attn.constants import BLOCK_SIZE
+from areal.models.tree_attn.constants import BLOCK_SIZE, USE_TRITON_TREE_ATTN
+from areal.models.tree_attn.triton_kernel import TRITON_AVAILABLE, tree_attention
 from areal.utils import logging
 
 logger = logging.getLogger(__name__)
@@ -88,33 +89,57 @@ def _tree_attn_fwd_func(
     *args,
     **kwargs,
 ):
-    # Require pre-created block_mask
-    block_mask = kwargs.get("block_mask", None)
-    if block_mask is None or not isinstance(block_mask, BlockMask):
-        raise ValueError(
-            "_tree_attn_fwd_func requires a pre-created BlockMask in kwargs['block_mask']. "
-            "Use create_block_mask_from_dense() during data preparation."
+    # Check for Triton path
+    triton_attn_data = kwargs.get("triton_attn_data", None)
+
+    if USE_TRITON_TREE_ATTN and triton_attn_data is not None and TRITON_AVAILABLE:
+        # [B, S, H, D] -> [B, H, S, D]
+        query = query.permute(0, 2, 1, 3).contiguous()
+        key = key.permute(0, 2, 1, 3).contiguous()
+        value = value.permute(0, 2, 1, 3).contiguous()
+
+        output = tree_attention(
+            query,
+            key,
+            value,
+            triton_attn_data.packed_mask,
+            triton_attn_data.kv_indices,
+            triton_attn_data.kv_offsets,
+            triton_attn_data.q_indices,
+            triton_attn_data.q_offsets,
+            sm_scale=softmax_scale,
         )
-
-    # [B, S, H, D] -> [B, H, S, D]
-    query = query.permute(0, 2, 1, 3).contiguous()
-    key = key.permute(0, 2, 1, 3).contiguous()
-    value = value.permute(0, 2, 1, 3).contiguous()
-
-    enable_gqa = query.shape[1] != key.shape[1]
-
-    output = _flex_attention(
-        query,
-        key,
-        value,
-        block_mask=block_mask,
-        score_mod=None,
-        scale=softmax_scale,
-        enable_gqa=enable_gqa,
-    )
-    # [B, H, S, D] -> [B, S, H, D]
-    output = output.permute(0, 2, 1, 3).contiguous()
-    return output
+        # [B, H, S, D] -> [B, S, H, D]
+        output = output.permute(0, 2, 1, 3).contiguous()
+        return output
+    else:
+        # Require pre-created block_mask
+        block_mask = kwargs.get("block_mask", None)
+        if block_mask is None or not isinstance(block_mask, BlockMask):
+            raise ValueError(
+                "_tree_attn_fwd_func requires a pre-created BlockMask in kwargs['block_mask']. "
+                "Use create_block_mask_from_dense() during data preparation."
+            )
+
+        # [B, S, H, D] -> [B, H, S, D]
+        query = query.permute(0, 2, 1, 3).contiguous()
+        key = key.permute(0, 2, 1, 3).contiguous()
+        value = value.permute(0, 2, 1, 3).contiguous()
+
+        enable_gqa = query.shape[1] != key.shape[1]
+
+        output = _flex_attention(
+            query,
+            key,
+            value,
+            block_mask=block_mask,
+            score_mod=None,
+            scale=softmax_scale,
+            enable_gqa=enable_gqa,
+        )
+        # [B, H, S, D] -> [B, S, H, D]
+        output = output.permute(0, 2, 1, 3).contiguous()
+        return output
 
 
 ORIGINAL_FLASH_ATTENTION_FORWARD = None