diff --git a/areal/engine/fsdp_engine.py b/areal/engine/fsdp_engine.py index e172bd1085..c6f89b6828 100644 --- a/areal/engine/fsdp_engine.py +++ b/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) diff --git a/areal/engine/megatron_engine.py b/areal/engine/megatron_engine.py index 5cd9afc79f..1bd5b792c3 100644 --- a/areal/engine/megatron_engine.py +++ b/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_) diff --git a/areal/models/tree_attn/constants.py b/areal/models/tree_attn/constants.py index 1deaf43d9c..bcdd16bf5e 100644 --- a/areal/models/tree_attn/constants.py +++ b/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." + ) diff --git a/areal/models/tree_attn/module.py b/areal/models/tree_attn/module.py index 21f5df3793..e6ae284259 100644 --- a/areal/models/tree_attn/module.py +++ b/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", diff --git a/areal/models/tree_attn/module_fsdp.py b/areal/models/tree_attn/module_fsdp.py index b1817ef9bc..34ae2117f9 100644 --- a/areal/models/tree_attn/module_fsdp.py +++ b/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 diff --git a/areal/models/tree_attn/module_megatron.py b/areal/models/tree_attn/module_megatron.py index 667d88aa2c..6d2515f801 100644 --- a/areal/models/tree_attn/module_megatron.py +++ b/areal/models/tree_attn/module_megatron.py @@ -15,7 +15,13 @@ ) from torch.nn.attention.flex_attention import BlockMask +from areal.models.tree_attn.constants import USE_TRITON_TREE_ATTN from areal.models.tree_attn.module_fsdp import _flex_attention +from areal.models.tree_attn.triton_kernel import ( + TRITON_AVAILABLE, + TreeAttentionData, + tree_attention, +) from areal.utils import logging logger = logging.getLogger(__name__) @@ -57,7 +63,7 @@ def forward( query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, - attention_mask: BlockMask, + attention_mask: BlockMask | TreeAttentionData, attn_mask_type: AttnMaskType, attention_bias: torch.Tensor = None, packed_seq_params: PackedSeqParams = None, @@ -65,46 +71,82 @@ def forward( # query: [S, B, H, D] in which B should be 1 in current tree training implementation # key: [S, B, H, D] # value: [S, B, H, D] - # attention_mask: BlockMask (pre-created) + # attention_mask: BlockMask | TreeAttentionData (pre-created) # attention_mask_type: arbitrary - if attention_bias is not None: - raise NotImplementedError( - "PytorchFlexAttention does not support attention_bias yet." - ) - if packed_seq_params is not None: - raise NotImplementedError( - "PytorchFlexAttention does not support packed sequences yet." + # Check for Triton path + if ( + USE_TRITON_TREE_ATTN + and isinstance(attention_mask, TreeAttentionData) + and TRITON_AVAILABLE + ): + # Triton path + if attention_bias is not None: + raise NotImplementedError("Triton path does not support attention_bias") + if packed_seq_params is not None: + raise NotImplementedError( + "Triton path does not support packed sequences" + ) + + # Input: [S, B, H, D], Triton expects [B, H, N, D] + query = query.permute(1, 2, 0, 3).contiguous() + key = key.permute(1, 2, 0, 3).contiguous() + value = value.permute(1, 2, 0, 3).contiguous() + + output = tree_attention( + query, + key, + value, + attention_mask.packed_mask, + attention_mask.kv_indices, + attention_mask.kv_offsets, + attention_mask.q_indices, + attention_mask.q_offsets, + sm_scale=self.softmax_scale, ) - if not isinstance(attention_mask, BlockMask): - raise ValueError( - "PytorchFlexAttention requires a pre-created BlockMask. " - "Use create_block_mask_from_dense() during data preparation." + # Output: [B, H, S, D], convert to [S, B, H*D] + output = output.permute(2, 0, 1, 3).contiguous() + output = output.view(output.shape[0], output.shape[1], -1) + return output + else: + # Existing flex_attention path + if attention_bias is not None: + raise NotImplementedError( + "PytorchFlexAttention does not support attention_bias yet." + ) + if packed_seq_params is not None: + raise NotImplementedError( + "PytorchFlexAttention does not support packed sequences yet." + ) + if not isinstance(attention_mask, BlockMask): + raise ValueError( + "PytorchFlexAttention requires a pre-created BlockMask for flex_attention path. " + "Use create_block_mask_from_dense() during data preparation." + ) + + # query, key, value shape: [S, B, H, D] -> [B, H, S, D] + query = query.permute(1, 2, 0, 3) + key = key.permute(1, 2, 0, 3) + value = value.permute(1, 2, 0, 3) + enable_gqa = query.shape[1] != key.shape[1] + + output = _flex_attention( + query, + key, + value, + block_mask=attention_mask, + score_mod=None, + scale=self.softmax_scale, + enable_gqa=enable_gqa, ) - # query, key, value shape: [S, B, H, D] -> [B, H, S, D] - query = query.permute(1, 2, 0, 3) - key = key.permute(1, 2, 0, 3) - value = value.permute(1, 2, 0, 3) - enable_gqa = query.shape[1] != key.shape[1] - - output = _flex_attention( - query, - key, - value, - block_mask=attention_mask, - score_mod=None, - scale=self.softmax_scale, - enable_gqa=enable_gqa, - ) - - # output shape: [B, H, S, D] -> [S, B, H, D] -> [S, B, H*D] - output = ( - output.permute(2, 0, 1, 3) - .contiguous() - .view(output.shape[2], output.shape[0], -1) - ) - return output + # output shape: [B, H, S, D] -> [S, B, H, D] -> [S, B, H*D] + output = ( + output.permute(2, 0, 1, 3) + .contiguous() + .view(output.shape[2], output.shape[0], -1) + ) + return output @contextmanager diff --git a/areal/models/tree_attn/tree.py b/areal/models/tree_attn/tree.py index e9499e23a1..d3134f43d3 100644 --- a/areal/models/tree_attn/tree.py +++ b/areal/models/tree_attn/tree.py @@ -17,6 +17,10 @@ from areal.api.cli_args import MicroBatchSpec from areal.models.tree_attn.constants import BLOCK_SIZE from areal.models.tree_attn.module_fsdp import create_block_mask_from_dense +from areal.models.tree_attn.triton_kernel import ( + TreeAttentionData, + precompute_tree_attention_data, +) from areal.utils import logging, stats_tracker from areal.utils.data import MicroBatchList from areal.utils.perf_tracer import trace_perf, trace_scope @@ -220,6 +224,42 @@ def _compress_chain( return trie_root +def trie_to_parent_array(trie: TrieNode, max_tokens: int) -> torch.Tensor: + """Build a parent array from TrieNode structure. + + The parent array `fa` is a 1D tensor where `fa[i]` is the index of the + parent token of token `i`. For root tokens, the parent index is -1. + In this tree structure, all tokens within a compressed trie node + share the same parent, which is the last token of the parent node. + + Args: + trie: The root TrieNode. + max_tokens: Maximum number of tokens (length of the output tensor). + + Returns: + torch.Tensor: Parent array of shape (1, max_tokens) with dtype int32. + """ + fa = torch.full((1, max_tokens), -1, dtype=torch.int32) + + if not trie.nodes: # Empty/dummy trie + return fa + + for node in trie.nodes: + parent_end_pos = -1 + if node.ancestors: # Has parent + parent_node = node.ancestors[-1] # Last ancestor is parent + parent_end_pos = parent_node.end_idx + + # First token in node attends to parent; internal tokens attend to previous token + if node.start_idx >= 0 and node.start_idx < max_tokens: + fa[0, node.start_idx] = parent_end_pos + for pos in range(node.start_idx + 1, node.end_idx + 1): + if pos < max_tokens: + fa[0, pos] = pos - 1 + + return fa + + # ============================================================================= # Public API # ============================================================================= @@ -392,7 +432,6 @@ def build_packed_tree_batch( non_packable_keys, ) - # Build micro-batch dict without block_mask (will be created lazily in forward) mb = { "input_ids": input_ids, "position_ids": position_ids, @@ -652,3 +691,49 @@ def build_block_mask_from_trie( del attention_mask return block_mask + + +@trace_perf("tree_attn.build_triton_attn_data_from_trie") +def build_triton_attn_data_from_trie( + trie: TrieNode, + padded_size: int, +) -> TreeAttentionData: + """Lazily build Triton tree attention data from a trie node. + + This function builds the parent array from the trie structure and + precomputes packed masks and sparse block indices for Triton kernels. + + Parameters + ---------- + trie : TrieNode + The root trie node containing the tree structure. + padded_size : int + The padded sequence length. + + Returns + ------- + TreeAttentionData + The precomputed Triton attention data. + """ + # Handle dummy trie (empty tree for DP synchronization) + if not trie.all_sequence_ids: + num_words = (padded_size + 63) >> 6 + num_q_blocks = (padded_size + 128 - 1) // 128 + num_kv_blocks = num_words + packed_mask = torch.zeros((1, padded_size, num_words), dtype=torch.int64) + kv_indices = torch.zeros((0,), dtype=torch.int32) + kv_offsets = torch.zeros((1, num_q_blocks + 1), dtype=torch.int32) + q_indices = torch.zeros((0,), dtype=torch.int32) + q_offsets = torch.zeros((1, num_kv_blocks + 1), dtype=torch.int32) + return TreeAttentionData( + packed_mask=packed_mask, + kv_indices=kv_indices, + kv_offsets=kv_offsets, + q_indices=q_indices, + q_offsets=q_offsets, + ) + + with trace_scope("tree_attn.precompute_triton_data"): + fa = trie_to_parent_array(trie, padded_size) + triton_attn_data = precompute_tree_attention_data(fa) + return triton_attn_data diff --git a/areal/models/tree_attn/triton_kernel.py b/areal/models/tree_attn/triton_kernel.py new file mode 100644 index 0000000000..276c342476 --- /dev/null +++ b/areal/models/tree_attn/triton_kernel.py @@ -0,0 +1,1039 @@ +"""Triton-based tree attention kernel with sparse block iteration.""" + +from dataclasses import dataclass + +import torch + +try: + import triton + import triton.language as tl + + TRITON_AVAILABLE = True +except ImportError: + TRITON_AVAILABLE = False + + +@dataclass +class TreeAttentionData: + packed_mask: torch.Tensor # (B, N, ceil(N/64)) + kv_indices: torch.Tensor # 1D sparse indices + kv_offsets: torch.Tensor # (B, num_q_blocks + 1) + q_indices: torch.Tensor # 1D sparse indices + q_offsets: torch.Tensor # (B, num_kv_blocks + 1) + + +def compute_packed_mask(fa: torch.Tensor) -> torch.Tensor: + """ + Compute packed ancestor mask directly from parent array. + + Uses DP: for each token, inherit parent's ancestor bits and set self bit. + Avoids materializing the full N×N boolean mask. + + Args: + fa: Parent array of shape (B, N). fa[b, i] is parent of token i in batch b, + or -1 for roots. + + Returns: + packed: (B, N, ceil(N/64)) int64 tensor where bit j of packed[b, i, j//64] + is set if token j is an ancestor of token i in batch b. + """ + B, N = fa.shape + num_words = (N + 63) >> 6 + packed = torch.zeros(B, N, num_words, dtype=torch.int64) + + for b in range(B): + for i in range(N): + parent = fa[b, i].item() + if parent >= 0: + parent_word = parent >> 6 + packed[b, i, : parent_word + 1] |= packed[b, parent, : parent_word + 1] + bit_val = torch.tensor(1, dtype=torch.int64) << (i & 63) + packed[b, i, i >> 6] |= bit_val + + return packed + + +def compute_kv_indices_packed( + packed_mask: torch.Tensor, + BLOCK_M: int, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Compute packed KV block indices for each Q block from bit-packed mask. + + Args: + packed_mask: (B, N, num_words) int64 tensor + BLOCK_M: Query block size + + Returns: + kv_indices: 1D tensor of KV block indices (all batches concatenated) + kv_offsets: (B, num_q_blocks + 1) tensor of absolute offsets into kv_indices + """ + B, size_q, num_kv_blocks = packed_mask.shape + num_q_blocks = (size_q + BLOCK_M - 1) // BLOCK_M + + all_indices = [] + offsets = torch.zeros(B, num_q_blocks + 1, dtype=torch.int32) + + for b in range(B): + for qb in range(num_q_blocks): + offsets[b, qb] = len(all_indices) + q_start = qb * BLOCK_M + q_end = min((qb + 1) * BLOCK_M, size_q) + + for kvb in range(num_kv_blocks): + if packed_mask[b, q_start:q_end, kvb].any(): + all_indices.append(kvb) + + offsets[b, num_q_blocks] = len(all_indices) + + kv_indices = torch.tensor(all_indices, dtype=torch.int32) + + return kv_indices, offsets + + +def compute_q_indices_packed( + packed_mask: torch.Tensor, + BLOCK_M: int, + BLOCK_N: int = 64, +) -> tuple[torch.Tensor, torch.Tensor]: + """ + Compute Q block indices for each KV block (reverse mapping for backward dK/dV). + + For each KV block, returns the list of Q blocks that attend to it. + + Args: + packed_mask: (B, N, num_words) int64 tensor + BLOCK_M: Query block size + BLOCK_N: KV block size (default 64) + + Returns: + q_indices: 1D tensor of Q block indices (all batches concatenated) + q_offsets: (B, num_kv_blocks + 1) tensor of absolute offsets into q_indices + """ + B, size_q, num_kv_blocks = packed_mask.shape + num_q_blocks = (size_q + BLOCK_M - 1) // BLOCK_M + + all_indices = [] + offsets = torch.zeros(B, num_kv_blocks + 1, dtype=torch.int32) + + for b in range(B): + for kvb in range(num_kv_blocks): + offsets[b, kvb] = len(all_indices) + + for qb in range(num_q_blocks): + q_start = qb * BLOCK_M + q_end = min((qb + 1) * BLOCK_M, size_q) + + if packed_mask[b, q_start:q_end, kvb].any(): + all_indices.append(qb) + + offsets[b, num_kv_blocks] = len(all_indices) + + q_indices = torch.tensor(all_indices, dtype=torch.int32) + + return q_indices, offsets + + +def precompute_tree_attention_data( + fa: torch.Tensor, + BLOCK_M: int = 128, + BLOCK_N: int = 64, +) -> TreeAttentionData: + """ + Precompute packed mask, KV indices, and Q indices from parent array. + + Args: + fa: Parent array of shape (B, N). fa[b, i] is parent of token i in batch b, + or -1 for roots. + BLOCK_M: Query block size (default 128) + BLOCK_N: KV block size (default 64) + + Returns: + TreeAttentionData object containing all precomputed tensors. + """ + packed_mask = compute_packed_mask(fa) + kv_indices, kv_offsets = compute_kv_indices_packed(packed_mask, BLOCK_M) + q_indices, q_offsets = compute_q_indices_packed(packed_mask, BLOCK_M, BLOCK_N) + return TreeAttentionData( + packed_mask=packed_mask, + kv_indices=kv_indices, + kv_offsets=kv_offsets, + q_indices=q_indices, + q_offsets=q_offsets, + ) + + +if TRITON_AVAILABLE: + + @triton.autotune( + configs=[ + triton.Config({"BLOCK_M": 128, "BLOCK_N": 64}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_M": 128, "BLOCK_N": 64}, num_warps=4, num_stages=3), + triton.Config({"BLOCK_M": 128, "BLOCK_N": 64}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_M": 128, "BLOCK_N": 64}, num_warps=8, num_stages=3), + ], + key=["N", "HEAD_DIM"], + ) + @triton.jit + def _tree_attn_fwd_triton( + Q, + K, + V, + o, + LSE, + packed_mask, + kv_indices, + kv_offsets, + sm_scale, + stride_qb, + stride_qh, + stride_qn, + stride_qd, + stride_kb, + stride_kh, + stride_kn, + stride_kd, + stride_vb, + stride_vh, + stride_vn, + stride_vd, + stride_ob, + stride_oh, + stride_on, + stride_od, + stride_lse_b, + stride_lse_h, + stride_lse_n, + stride_mask_b, + stride_mask_q, + stride_mask_k, + stride_off_b, + stride_off_q, + B: tl.constexpr, + H: tl.constexpr, + N: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + GQA_GROUP_SIZE: tl.constexpr, + ): + """ + Tree attention forward kernel with dense bit-packed mask. + + Grid: (cdiv(N, BLOCK_M), B * H) + """ + pid_m = tl.program_id(0) + pid_bh = tl.program_id(1) + off_b = pid_bh // H + off_h = pid_bh % H + # Handle GQA: map query head to key/value head + off_h_kv = off_h // GQA_GROUP_SIZE + + qo_offset = off_b * stride_qb + off_h * stride_qh + kv_offset = off_b * stride_kb + off_h_kv * stride_kh + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_d = tl.arange(0, HEAD_DIM) + + # Load Q block + q_ptrs = ( + Q + qo_offset + offs_m[:, None] * stride_qn + offs_d[None, :] * stride_qd + ) + q = tl.load(q_ptrs, mask=offs_m[:, None] < N, other=0.0) + + # Get KV block range for this Q block (batch-indexed) + kv_start_idx = tl.load(kv_offsets + off_b * stride_off_b + pid_m * stride_off_q) + kv_end_idx = tl.load( + kv_offsets + off_b * stride_off_b + (pid_m + 1) * stride_off_q + ) + num_kv = kv_end_idx - kv_start_idx + + # Initialize online softmax accumulators + m_i = tl.full([BLOCK_M], float("-inf"), dtype=tl.float32) + l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32) + + qk_scale = sm_scale * 1.44269504 + + # Iterate over valid KV blocks only + for i in range(num_kv): + kv_block_idx = tl.load(kv_indices + kv_start_idx + i) + n_start = kv_block_idx * BLOCK_N + offs_n = n_start + tl.arange(0, BLOCK_N) + + # Load K, V blocks + k_ptrs = ( + K + + kv_offset + + offs_n[:, None] * stride_kn + + offs_d[None, :] * stride_kd + ) + k = tl.load(k_ptrs, mask=offs_n[:, None] < N, other=0.0) + + v_ptrs = ( + V + + kv_offset + + offs_n[:, None] * stride_vn + + offs_d[None, :] * stride_vd + ) + v = tl.load(v_ptrs, mask=offs_n[:, None] < N, other=0.0) + + # Compute QK^T + qk = tl.dot(q, tl.trans(k)) + + # Load and unpack mask block (BLOCK_N=64, one int64 word per block) + mask_ptrs = ( + packed_mask + + off_b * stride_mask_b + + offs_m * stride_mask_q + + kv_block_idx * stride_mask_k + ) + mask_word = tl.load(mask_ptrs, mask=offs_m < N, other=0) # (BLOCK_M,) + + # Unpack 64 bits + bit_indices = tl.arange(0, BLOCK_N) + mask = ((mask_word[:, None] >> bit_indices[None, :]) & 1) != 0 + + # Apply bounds mask + valid_mask = (offs_m[:, None] < N) & (offs_n[None, :] < N) + mask = mask & valid_mask + + # Apply mask to attention scores + # Avoid -inf since -inf - (-inf) = NaN + qk = tl.where(mask, qk * qk_scale, -16384.0) + + # Online softmax update + m_ij = tl.max(qk, axis=1) + m_new = tl.maximum(m_i, m_ij) + + alpha = tl.math.exp2(m_i - m_new) + acc = acc * alpha[:, None] + l_i = l_i * alpha + + p = tl.math.exp2(qk - m_new[:, None]) + l_ij = tl.sum(p, axis=1) + l_i = l_i + l_ij + + p = p.to(v.dtype) + acc = acc + tl.dot(p, v) + + m_i = m_new + + # Normalize + acc = acc / l_i[:, None] + + # Store output + o_ptrs = ( + o + qo_offset + offs_m[:, None] * stride_on + offs_d[None, :] * stride_od + ) + tl.store(o_ptrs, acc.to(o.dtype.element_ty), mask=offs_m[:, None] < N) + + # Store LSE for backward: lse = m + log2(l) + lse = m_i + tl.math.log2(l_i) + lse_ptrs = ( + LSE + off_b * stride_lse_b + off_h * stride_lse_h + offs_m * stride_lse_n + ) + tl.store(lse_ptrs, lse, mask=offs_m < N) + + def tree_attention_triton( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + packed_mask: torch.Tensor, + kv_indices: torch.Tensor, + kv_offsets: torch.Tensor, + sm_scale: float | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Tree attention forward using precomputed dense bit-packed mask with block skipping. + + Returns output and LSE for backward pass. + """ + B, H, N, D = query.shape + # Handle GQA + assert key.shape[0] == B and key.shape[2] == N and key.shape[3] == D + assert value.shape == key.shape + H_kv = key.shape[1] + assert H % H_kv == 0, f"Query heads {H} must be divisible by KV heads {H_kv}" + gqa_group_size = H // H_kv + + BLOCK_M = 128 + # BLOCK_N is fixed to 64 in the kernel + + if sm_scale is None: + sm_scale = 1.0 / (D**0.5) + + packed_mask = packed_mask.to(query.device).contiguous() + kv_indices = kv_indices.to(query.device).contiguous() + kv_offsets = kv_offsets.to(query.device).contiguous() + + query = query.contiguous() + key = key.contiguous() + value = value.contiguous() + + output = torch.empty_like(query) + lse = torch.empty(B, H, N, device=query.device, dtype=torch.float32) + + grid = (triton.cdiv(N, BLOCK_M), B * H) + + _tree_attn_fwd_triton[grid]( + query, + key, + value, + output, + lse, + packed_mask, + kv_indices, + kv_offsets, + sm_scale, + query.stride(0), + query.stride(1), + query.stride(2), + query.stride(3), + key.stride(0), + key.stride(1), + key.stride(2), + key.stride(3), + value.stride(0), + value.stride(1), + value.stride(2), + value.stride(3), + output.stride(0), + output.stride(1), + output.stride(2), + output.stride(3), + lse.stride(0), + lse.stride(1), + lse.stride(2), + packed_mask.stride(0), + packed_mask.stride(1), + packed_mask.stride(2), + kv_offsets.stride(0), + kv_offsets.stride(1), + B=B, + H=H, + N=N, + HEAD_DIM=D, + GQA_GROUP_SIZE=gqa_group_size, + ) + + return output, lse + + @triton.jit + def _tree_attn_bwd_preprocess( + o, + DO, + Delta, + stride_ob, + stride_oh, + stride_on, + stride_od, + N: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_M: tl.constexpr, + ): + pid_m = tl.program_id(0) + pid_bh = tl.program_id(1) + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_d = tl.arange(0, HEAD_DIM) + + o_ptrs = ( + o + + pid_bh * stride_oh + + offs_m[:, None] * stride_on + + offs_d[None, :] * stride_od + ) + do_ptrs = ( + DO + + pid_bh * stride_oh + + offs_m[:, None] * stride_on + + offs_d[None, :] * stride_od + ) + + o = tl.load(o_ptrs, mask=offs_m[:, None] < N, other=0.0).to(tl.float32) + do = tl.load(do_ptrs, mask=offs_m[:, None] < N, other=0.0).to(tl.float32) + + delta = tl.sum(o * do, axis=1) + + delta_ptrs = Delta + pid_bh * N + offs_m + tl.store(delta_ptrs, delta, mask=offs_m < N) + + @triton.autotune( + configs=[ + triton.Config({"BLOCK_M": 128, "BLOCK_N": 64}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_M": 128, "BLOCK_N": 64}, num_warps=4, num_stages=3), + triton.Config({"BLOCK_M": 128, "BLOCK_N": 64}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_M": 128, "BLOCK_N": 64}, num_warps=8, num_stages=3), + ], + key=["N", "HEAD_DIM"], + ) + @triton.jit + def _tree_attn_bwd_dq( + Q, + K, + V, + DO, + DQ, + LSE, + Delta, + packed_mask, + kv_indices, + kv_offsets, + sm_scale, + stride_qb, + stride_qh, + stride_qn, + stride_qd, + stride_kb, + stride_kh, + stride_kn, + stride_kd, + stride_vb, + stride_vh, + stride_vn, + stride_vd, + stride_dob, + stride_doh, + stride_don, + stride_dod, + stride_dqb, + stride_dqh, + stride_dqn, + stride_dqd, + stride_lse_b, + stride_lse_h, + stride_lse_n, + stride_mask_b, + stride_mask_q, + stride_mask_k, + stride_off_b, + stride_off_q, + B: tl.constexpr, + H: tl.constexpr, + N: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + GQA_GROUP_SIZE: tl.constexpr, + ): + pid_m = tl.program_id(0) + pid_bh = tl.program_id(1) + off_b = pid_bh // H + off_h = pid_bh % H + + qo_offset = off_b * stride_qb + off_h * stride_qh + # Handle GQA mapping + off_h_kv = off_h // GQA_GROUP_SIZE + kv_offset = off_b * stride_kb + off_h_kv * stride_kh + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_d = tl.arange(0, HEAD_DIM) + + q = tl.load( + Q + qo_offset + offs_m[:, None] * stride_qn + offs_d[None, :] * stride_qd, + mask=offs_m[:, None] < N, + other=0.0, + ) + do = tl.load( + DO + + off_b * stride_dob + + off_h * stride_doh + + offs_m[:, None] * stride_don + + offs_d[None, :] * stride_dod, + mask=offs_m[:, None] < N, + other=0.0, + ) + lse = tl.load( + LSE + off_b * stride_lse_b + off_h * stride_lse_h + offs_m * stride_lse_n, + mask=offs_m < N, + other=0.0, + ) + delta = tl.load(Delta + pid_bh * N + offs_m, mask=offs_m < N, other=0.0) + + kv_start = tl.load(kv_offsets + off_b * stride_off_b + pid_m * stride_off_q) + kv_end = tl.load(kv_offsets + off_b * stride_off_b + (pid_m + 1) * stride_off_q) + num_kv = kv_end - kv_start + + dq = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32) + qk_scale = sm_scale * 1.44269504 + + for i in range(num_kv): + kv_block_idx = tl.load(kv_indices + kv_start + i) + offs_n = kv_block_idx * BLOCK_N + tl.arange(0, BLOCK_N) + + k = tl.load( + K + + kv_offset + + offs_n[:, None] * stride_kn + + offs_d[None, :] * stride_kd, + mask=offs_n[:, None] < N, + other=0.0, + ) + v = tl.load( + V + + kv_offset + + offs_n[:, None] * stride_vn + + offs_d[None, :] * stride_vd, + mask=offs_n[:, None] < N, + other=0.0, + ) + + mask_word = tl.load( + packed_mask + + off_b * stride_mask_b + + offs_m * stride_mask_q + + kv_block_idx * stride_mask_k, + mask=offs_m < N, + other=0, + ) + bit_indices = tl.arange(0, BLOCK_N) + mask = ((mask_word[:, None] >> bit_indices[None, :]) & 1) != 0 + mask = mask & ((offs_m[:, None] < N) & (offs_n[None, :] < N)) + + qk = tl.dot(q, tl.trans(k)) + qk = tl.where(mask, qk * qk_scale, -16384.0) + p = tl.math.exp2(qk - lse[:, None]) + p = tl.where(mask, p, 0.0) + + dp = tl.dot(do.to(v.dtype), tl.trans(v)).to(tl.float32) + ds = p * (dp - delta[:, None]) + ds = tl.where(mask, ds, 0.0) + + dq += tl.dot(ds.to(k.dtype), k).to(tl.float32) + + dq *= sm_scale + tl.store( + DQ + + off_b * stride_dqb + + off_h * stride_dqh + + offs_m[:, None] * stride_dqn + + offs_d[None, :] * stride_dqd, + dq.to(DQ.dtype.element_ty), + mask=offs_m[:, None] < N, + ) + + @triton.autotune( + configs=[ + triton.Config({"BLOCK_M": 128, "BLOCK_N": 64}, num_warps=4, num_stages=2), + triton.Config({"BLOCK_M": 128, "BLOCK_N": 64}, num_warps=4, num_stages=3), + triton.Config({"BLOCK_M": 128, "BLOCK_N": 64}, num_warps=8, num_stages=2), + triton.Config({"BLOCK_M": 128, "BLOCK_N": 64}, num_warps=8, num_stages=3), + ], + key=["N", "HEAD_DIM"], + ) + @triton.jit + def _tree_attn_bwd_dkdv( + Q, + K, + V, + DO, + DK, + DV, + LSE, + Delta, + packed_mask, + sm_scale, + stride_qb, + stride_qh, + stride_qn, + stride_qd, + stride_kb, + stride_kh, + stride_kn, + stride_kd, + stride_vb, + stride_vh, + stride_vn, + stride_vd, + stride_dob, + stride_doh, + stride_don, + stride_dod, + stride_dkb, + stride_dkh, + stride_dkn, + stride_dkd, + stride_dvb, + stride_dvh, + stride_dvn, + stride_dvd, + stride_lse_b, + stride_lse_h, + stride_lse_n, + stride_mask_b, + stride_mask_q, + stride_mask_k, + B: tl.constexpr, + H: tl.constexpr, + H_KV: tl.constexpr, + N: tl.constexpr, + HEAD_DIM: tl.constexpr, + BLOCK_M: tl.constexpr, + BLOCK_N: tl.constexpr, + NUM_Q_BLOCKS: tl.constexpr, + GQA_GROUP_SIZE: tl.constexpr, + ): + pid_n = tl.program_id(0) + pid_bh_kv = tl.program_id(1) + off_b = pid_bh_kv // H_KV + off_h_kv = pid_bh_kv % H_KV + + kv_offset = off_b * stride_kb + off_h_kv * stride_kh + + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + offs_d = tl.arange(0, HEAD_DIM) + + k = tl.load( + K + kv_offset + offs_n[:, None] * stride_kn + offs_d[None, :] * stride_kd, + mask=offs_n[:, None] < N, + other=0.0, + ) + v = tl.load( + V + kv_offset + offs_n[:, None] * stride_vn + offs_d[None, :] * stride_vd, + mask=offs_n[:, None] < N, + other=0.0, + ) + + dk = tl.zeros([BLOCK_N, HEAD_DIM], dtype=tl.float32) + dv = tl.zeros([BLOCK_N, HEAD_DIM], dtype=tl.float32) + qk_scale = sm_scale * 1.44269504 + + # Loop over Q heads in the group + for g in range(GQA_GROUP_SIZE): + off_h = off_h_kv * GQA_GROUP_SIZE + g + qo_offset = off_b * stride_qb + off_h * stride_qh + + # Loop over Q blocks + for q_block_idx in range(NUM_Q_BLOCKS): + offs_m = q_block_idx * BLOCK_M + tl.arange(0, BLOCK_M) + + mask_word = tl.load( + packed_mask + + off_b * stride_mask_b + + offs_m * stride_mask_q + + pid_n * stride_mask_k, + mask=offs_m < N, + other=0, + ) + has_nonzero = tl.sum(mask_word) != 0 + if has_nonzero: + q = tl.load( + Q + + qo_offset + + offs_m[:, None] * stride_qn + + offs_d[None, :] * stride_qd, + mask=offs_m[:, None] < N, + other=0.0, + ) + do = tl.load( + DO + + off_b * stride_dob + + off_h * stride_doh + + offs_m[:, None] * stride_don + + offs_d[None, :] * stride_dod, + mask=offs_m[:, None] < N, + other=0.0, + ) + lse = tl.load( + LSE + + off_b * stride_lse_b + + off_h * stride_lse_h + + offs_m * stride_lse_n, + mask=offs_m < N, + other=0.0, + ) + delta = tl.load( + Delta + (off_b * H + off_h) * N + offs_m, + mask=offs_m < N, + other=0.0, + ) + + bit_indices = tl.arange(0, BLOCK_N) + mask = ((mask_word[:, None] >> bit_indices[None, :]) & 1) != 0 + mask = mask & ((offs_m[:, None] < N) & (offs_n[None, :] < N)) + + qk = tl.dot(q, tl.trans(k)) + qk = tl.where(mask, qk * qk_scale, -16384.0) + p = tl.math.exp2(qk - lse[:, None]) + p = tl.where(mask, p, 0.0) + + dv += tl.dot(tl.trans(p.to(do.dtype)), do).to(tl.float32) + + dp = tl.dot(do.to(v.dtype), tl.trans(v)).to(tl.float32) + ds = p * (dp - delta[:, None]) + ds = tl.where(mask, ds, 0.0) + + dk += tl.dot(tl.trans(ds.to(q.dtype)), q).to(tl.float32) + + dk *= sm_scale + tl.store( + DK + + off_b * stride_dkb + + off_h_kv * stride_dkh + + offs_n[:, None] * stride_dkn + + offs_d[None, :] * stride_dkd, + dk.to(DK.dtype.element_ty), + mask=offs_n[:, None] < N, + ) + tl.store( + DV + + off_b * stride_dvb + + off_h_kv * stride_dvh + + offs_n[:, None] * stride_dvn + + offs_d[None, :] * stride_dvd, + dv.to(DV.dtype.element_ty), + mask=offs_n[:, None] < N, + ) + + def tree_attention_backward( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + output: torch.Tensor, + lse: torch.Tensor, + dout: torch.Tensor, + packed_mask: torch.Tensor, + kv_indices: torch.Tensor, + kv_offsets: torch.Tensor, + q_indices: torch.Tensor, + q_offsets: torch.Tensor, + sm_scale: float, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + B, H, N, D = query.shape + H_kv = key.shape[1] + gqa_group_size = H // H_kv + BLOCK_M = 128 + BLOCK_N = 64 + + dout = dout.contiguous() + dq = torch.zeros_like(query) + dk = torch.zeros_like(key) + dv = torch.zeros_like(value) + delta = torch.empty(B * H, N, device=query.device, dtype=torch.float32) + + pre_grid = (triton.cdiv(N, BLOCK_M), B * H) + _tree_attn_bwd_preprocess[pre_grid]( + output, + dout, + delta, + output.stride(0) * output.stride(1), + output.stride(1), + output.stride(2), + output.stride(3), + N=N, + HEAD_DIM=D, + BLOCK_M=BLOCK_M, + num_warps=4, + ) + + num_kv_blocks = triton.cdiv(N, BLOCK_N) + num_q_blocks = triton.cdiv(N, BLOCK_M) + + dq_grid = (num_q_blocks, B * H) + _tree_attn_bwd_dq[dq_grid]( + query, + key, + value, + dout, + dq, + lse, + delta, + packed_mask, + kv_indices, + kv_offsets, + sm_scale, + query.stride(0), + query.stride(1), + query.stride(2), + query.stride(3), + key.stride(0), + key.stride(1), + key.stride(2), + key.stride(3), + value.stride(0), + value.stride(1), + value.stride(2), + value.stride(3), + dout.stride(0), + dout.stride(1), + dout.stride(2), + dout.stride(3), + dq.stride(0), + dq.stride(1), + dq.stride(2), + dq.stride(3), + lse.stride(0), + lse.stride(1), + lse.stride(2), + packed_mask.stride(0), + packed_mask.stride(1), + packed_mask.stride(2), + kv_offsets.stride(0), + kv_offsets.stride(1), + B=B, + H=H, + N=N, + HEAD_DIM=D, + GQA_GROUP_SIZE=gqa_group_size, + ) + + dkdv_grid = (num_kv_blocks, B * H_kv) + _tree_attn_bwd_dkdv[dkdv_grid]( + query, + key, + value, + dout, + dk, + dv, + lse, + delta, + packed_mask, + sm_scale, + query.stride(0), + query.stride(1), + query.stride(2), + query.stride(3), + key.stride(0), + key.stride(1), + key.stride(2), + key.stride(3), + value.stride(0), + value.stride(1), + value.stride(2), + value.stride(3), + dout.stride(0), + dout.stride(1), + dout.stride(2), + dout.stride(3), + dk.stride(0), + dk.stride(1), + dk.stride(2), + dk.stride(3), + dv.stride(0), + dv.stride(1), + dv.stride(2), + dv.stride(3), + lse.stride(0), + lse.stride(1), + lse.stride(2), + packed_mask.stride(0), + packed_mask.stride(1), + packed_mask.stride(2), + B=B, + H=H, + H_KV=H_kv, + N=N, + HEAD_DIM=D, + NUM_Q_BLOCKS=num_q_blocks, + GQA_GROUP_SIZE=gqa_group_size, + ) + + return dq, dk, dv + + class TreeAttentionFunction(torch.autograd.Function): + @staticmethod + def forward( + ctx, + query, + key, + value, + packed_mask, + kv_indices, + kv_offsets, + q_indices, + q_offsets, + sm_scale, + ): + output, lse = tree_attention_triton( + query, key, value, packed_mask, kv_indices, kv_offsets, sm_scale + ) + ctx.save_for_backward( + query, + key, + value, + output, + lse, + packed_mask, + kv_indices, + kv_offsets, + q_indices, + q_offsets, + ) + ctx.sm_scale = sm_scale + return output + + @staticmethod + def backward(ctx, dout): + ( + query, + key, + value, + output, + lse, + packed_mask, + kv_indices, + kv_offsets, + q_indices, + q_offsets, + ) = ctx.saved_tensors + dq, dk, dv = tree_attention_backward( + query, + key, + value, + output, + lse, + dout, + packed_mask, + kv_indices, + kv_offsets, + q_indices, + q_offsets, + ctx.sm_scale, + ) + return dq, dk, dv, None, None, None, None, None, None + + def tree_attention( + query: torch.Tensor, + key: torch.Tensor, + value: torch.Tensor, + packed_mask: torch.Tensor, + kv_indices: torch.Tensor, + kv_offsets: torch.Tensor, + q_indices: torch.Tensor, + q_offsets: torch.Tensor, + sm_scale: float | None = None, + ) -> torch.Tensor: + """ + Tree attention with backward support. + + Args: + query, key, value: (B, H, N, D) tensors + packed_mask, kv_indices, kv_offsets, q_indices, q_offsets: from precompute_tree_attention_data + sm_scale: softmax scale (default: 1/sqrt(D)) + """ + if sm_scale is None: + sm_scale = 1.0 / (query.shape[-1] ** 0.5) + packed_mask = packed_mask.to(query.device).contiguous() + kv_indices = kv_indices.to(query.device).contiguous() + kv_offsets = kv_offsets.to(query.device).contiguous() + q_indices = q_indices.to(query.device).contiguous() + q_offsets = q_offsets.to(query.device).contiguous() + return TreeAttentionFunction.apply( + query.contiguous(), + key.contiguous(), + value.contiguous(), + packed_mask, + kv_indices, + kv_offsets, + q_indices, + q_offsets, + sm_scale, + ) + +else: + + def tree_attention(*args, **kwargs): + raise ImportError( + "Triton is not available. Please install triton to use tree_attention." + ) diff --git a/areal/tests/test_tree_training.py b/areal/tests/test_tree_training.py index 3c8a5fd263..6a16c4bd89 100644 --- a/areal/tests/test_tree_training.py +++ b/areal/tests/test_tree_training.py @@ -3,6 +3,7 @@ import pytest import torch +import areal from areal.api.alloc_mode import AllocationMode from areal.api.cli_args import ( MicroBatchSpec, @@ -14,6 +15,7 @@ from areal.engine.megatron_engine import MegatronEngine from areal.infra.platforms import current_platform from areal.models.tree_attn.module import restore_patch_fsdp_for_tree_training +from areal.models.tree_attn.triton_kernel import TRITON_AVAILABLE from areal.tests.utils import get_model_path from areal.utils import logging @@ -199,8 +201,18 @@ def _create_engine( @pytest.mark.parametrize("engine_type", ["megatron", "fsdp"]) -def test_tree_training_forward(engine_type): +@pytest.mark.parametrize("tree_attn_backend", ["flex", "triton"]) +def test_tree_training_forward(engine_type, tree_attn_backend): """Test tree training forward pass produces correct logprobs.""" + if tree_attn_backend == "triton" and not TRITON_AVAILABLE: + pytest.skip("Triton is not available") + + use_triton = tree_attn_backend == "triton" + + # Patch constants to force specific backend + areal.models.tree_attn.tree.USE_TRITON_TREE_ATTN = use_triton + areal.models.tree_attn.module_fsdp.USE_TRITON_TREE_ATTN = use_triton + areal.models.tree_attn.module_megatron.USE_TRITON_TREE_ATTN = use_triton # Create baseline engine inputs = mock_tree_input() baseline_engine = _create_engine(engine_type, port="7777") @@ -274,8 +286,13 @@ def test_tree_training_forward(engine_type): @pytest.mark.parametrize("engine_type", ["megatron", "fsdp"]) -def test_tree_training_forward_backward(engine_type): +@pytest.mark.parametrize("tree_attn_backend", ["flex", "triton"]) +def test_tree_training_forward_backward(engine_type, tree_attn_backend): """Test tree training forward-backward pass produces correct gradients.""" + if tree_attn_backend == "triton" and not TRITON_AVAILABLE: + pytest.skip("Triton is not available") + + use_triton = tree_attn_backend == "triton" def loss_fn(logprobs, entropy, input_data, **kwargs): return logprobs.mean() @@ -283,6 +300,9 @@ def loss_fn(logprobs, entropy, input_data, **kwargs): def loss_weight_fn(input_data): return input_data["loss_mask"].count_nonzero() + areal.models.tree_attn.tree.USE_TRITON_TREE_ATTN = use_triton + areal.models.tree_attn.module_fsdp.USE_TRITON_TREE_ATTN = use_triton + areal.models.tree_attn.module_megatron.USE_TRITON_TREE_ATTN = use_triton inputs = mock_tree_input() # Create baseline engine baseline_engine = _create_engine(engine_type, port="7777") diff --git a/areal/utils/mcore/packed_context_parallel.py b/areal/utils/mcore/packed_context_parallel.py index adb3c81ab6..9381e289e6 100644 --- a/areal/utils/mcore/packed_context_parallel.py +++ b/areal/utils/mcore/packed_context_parallel.py @@ -128,10 +128,11 @@ def packed_context_parallel_forward( position_ids = input_["position_ids"] cu_seqlens = input_.get("cu_seqlens", None) block_mask = input_.get("block_mask", None) + triton_attn_data = input_.get("triton_attn_data", None) packed_seq_params = None if cu_seqlens is not None: - if block_mask is not None: + if block_mask is not None or triton_attn_data is not None: raise ValueError( "Attention mask should be None when using packed sequences." ) @@ -140,10 +141,13 @@ def packed_context_parallel_forward( ) input_ids = input_ids.contiguous() + # Pass triton_attn_data as attention_mask if present (for Triton tree attention) + attention_mask = triton_attn_data if triton_attn_data is not None else block_mask + try: output = model( input_ids=input_ids, - attention_mask=block_mask, + attention_mask=attention_mask, position_ids=position_ids, packed_seq_params=packed_seq_params, )