# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import math from typing import Any, Optional, Tuple, Union import torch from megatron.core import ModelParallelConfig, parallel_state from torch import nn from torch.distributed import _functional_collectives as funcol from transformer_engine.pytorch.attention import _SplitAlongDim, apply_rotary_pos_emb, check_set_window_size from transformer_engine.pytorch.constants import AttnBiasTypes from transformer_engine.pytorch.float8_tensor import Float8Tensor from transformer_engine.pytorch.module.linear import Linear as LinearTE from transformer_engine.pytorch.module.rmsnorm import RMSNorm as RMSNormTE from cosmos_predict1.autoregressive.modules.embedding import RotaryPositionEmbedding from cosmos_predict1.autoregressive.modules.linear import ColumnParallelLinear, RowParallelLinear from cosmos_predict1.autoregressive.modules.normalization import create_norm from cosmos_predict1.autoregressive.utils.parallel import AllReduceBWDRMSNormTE class GQA(nn.Module): """ Grouped Query Attention (GQA) with KV cache (only supported for inference). """ def __init__( self, n_heads: int, n_kv_heads: Union[int, None], dim: int, max_batch_size: int, max_seq_len: int, context_dim: Optional[int] = None, inference: bool = True, flash_attn: bool = True, use_qk_normalization: bool = False, norm_type: str = "rmsnorm", norm_eps: float = 1e-5, attention_dropout: float = 0.0, set_parallel_mode: Optional[bool] = False, model_parallel: Optional[ModelParallelConfig] = None, attention_tp: Optional[bool] = False, causal_mask: Optional[bool] = True, head_dim: Optional[int] = None, fuse_qkv: bool = False, precision: str = "bfloat16", attention_type: str = "self", ): """ Initializes the GQA module. Args: n_heads (int): The number of attention heads. n_kv_heads (int, optional): The number of key-value attention heads. None defaults to n_heads. dim (int): The dimensionality of the input and output. max_batch_size (int): The maximum batch size. max_seq_len (int): The maximum sequence length. context_dim (int, optional): The dimensionality of the context for cross-attn. Defaults to None. inference (bool, optional): Whether the model is in inference mode. Defaults to True. flash_attn (bool, optional): Whether to use Flash attention. Defaults to True. use_qk_normalization (bool, optional): Whether to apply QK normalization. Defaults to False. norm_type (str, optional): The type of normalization layer. Defaults to "rmsnorm". norm_eps (float, optional): The epsilon value for normalization. Defaults to 1e-5. attention_dropout (float, optional): Dropout rate for attention. Defaults to 0.0. tp_group (int, optional): The tensor parallel group. set_parallel_mode (bool, optional): Whether to set parallel mode which enables parallel linear. Defaults to False. model_parallel (ModelParallelConfig, optional): The Megatron model parallel configuration. attention_tp (bool, optional): Whether to use tensor parallelism for attention layers. Defaults to False. causal_mask (bool, optional): Whether to use causal mask. Defaults to True. head_dim (int, optional): The dimensionality of each attention head. If None, defaults to dim // n_heads. fuse_qkv (bool, optional): Whether to fuse QKV projections. Defaults to False. precision (str, optional): The precision of the model. Defaults to "bfloat16". attention_type (str, optional): The type of attention. Defaults to "self". """ super().__init__() assert attention_type in ["self", "cross", "full"], f"Invalid attention type: {attention_type}" self.attention_type = attention_type self.model_parallel = model_parallel if self.model_parallel and self.model_parallel.tensor_model_parallel_size > 1 and attention_tp: self.tp_size = self.model_parallel.tensor_model_parallel_size else: self.tp_size = 1 context_dim = dim if context_dim is None else context_dim self.dim = dim self.context_dim = context_dim self.n_kv_heads = n_heads if n_kv_heads is None else n_kv_heads self.n_local_kv_heads = self.n_kv_heads // self.tp_size self.n_local_heads = n_heads // self.tp_size self.n_rep = self.n_local_heads // self.n_local_kv_heads self.head_dim = dim // n_heads if head_dim is None else head_dim assert flash_attn, "Flash attention is required." self.attention_dropout = attention_dropout self.causal_mask = causal_mask self.fuse_qkv = fuse_qkv self.precision = precision if fuse_qkv: assert context_dim == dim, f"Fuse QKV requires context_dim ({context_dim}) to be equal to dim ({dim})" self.total_head_dim = (n_heads + 2 * self.n_kv_heads) * self.head_dim self.total_local_head_dim = (self.n_local_heads + 2 * self.n_local_kv_heads) * self.head_dim if set_parallel_mode and attention_tp and not inference: kwargs = {"bias": False, "init_method": lambda x: x, "config": self.model_parallel} # Using column and row parallel linear layers if fuse_qkv: self.wqkv = ColumnParallelLinear(dim, self.total_head_dim, **kwargs) else: self.wq = ColumnParallelLinear(dim, n_heads * self.head_dim, **kwargs) self.wk = ColumnParallelLinear(context_dim, self.n_kv_heads * self.head_dim, **kwargs) self.wv = ColumnParallelLinear(context_dim, self.n_kv_heads * self.head_dim, **kwargs) # Linear layer for output projection self.wo = RowParallelLinear( n_heads * self.head_dim, dim, input_is_parallel=True, skip_bias_add=True, **kwargs ) else: # Linear layers for query, key, and value projections if fuse_qkv: self.wqkv = nn.Linear(dim, self.total_local_head_dim, bias=False) else: self.wq = nn.Linear(dim, self.n_local_heads * self.head_dim, bias=False) self.wk = nn.Linear(context_dim, self.n_local_kv_heads * self.head_dim, bias=False) self.wv = nn.Linear(context_dim, self.n_local_kv_heads * self.head_dim, bias=False) self.wo = nn.Linear(self.n_local_heads * self.head_dim, dim, bias=False) self.max_batch_size = max_batch_size self.max_seq_len = max_seq_len if inference and self.attention_type == "self": # Cache for key and value tensors self.init_kv_cache() # QK normalization layers if use_qk_normalization: assert n_heads % self.tp_size == 0, "n_heads must be divisible by tensor_model_parallel_size" assert self.n_kv_heads % self.tp_size == 0, "n_kv_heads must be divisible by tensor_model_parallel_size" self.q_norm = create_norm(norm_type, dim=self.head_dim, eps=norm_eps) self.k_norm = create_norm(norm_type, dim=self.head_dim, eps=norm_eps) else: self.q_norm = nn.Identity() self.k_norm = nn.Identity() self.use_qk_normalization = use_qk_normalization self.inference = inference if fuse_qkv: # Register hook to load fused QKV weights self._register_load_state_dict_pre_hook(self.load_hook) self.to(dtype=getattr(torch, self.precision)) def load_hook(self, state_dict, prefix, *args): if prefix + "wq.weight" in state_dict: wq = state_dict.pop(prefix + "wq.weight") wk = state_dict.pop(prefix + "wk.weight") wv = state_dict.pop(prefix + "wv.weight") state_dict[prefix + "wqkv.weight"] = torch.cat([wq, wk, wv]) def init_kv_cache(self, dtype=None): cache_shape = (self.max_batch_size, self.n_local_kv_heads, self.max_seq_len, self.head_dim) if dtype is None: dtype = getattr(torch, self.precision) if self.attention_type == "self": self.cache_k = torch.zeros(cache_shape, dtype=dtype).cuda() self.cache_v = torch.zeros(cache_shape, dtype=dtype).cuda() def set_inference_flag(self, flag): self.inference = flag if flag and self.attention_type == "self": if self.cache_k is None or self.cache_v is None: self.init_kv_cache() def forward( self, x: torch.Tensor, rope: RotaryPositionEmbedding, input_pos: torch.Tensor, mask: Optional[torch.Tensor] = None, context: Optional[torch.Tensor] = None, ): """ Forward pass of GQA. Args: x: The input tensor of shape (batch_size, seq_len, dim). rope: The rotary positional embedding module. input_pos: The starting position of the current sequence. mask: The attention mask tensor. context: The context tensor of shape (batch_size, context_len, dim). Returns: The output tensor after applying GQA. """ bsz, seqlen, _ = x.shape # Use one single module to handle both self-attn and cross-attn context = x if context is None else context context_len = seqlen if context is None else context.shape[1] if self.fuse_qkv: q_size = self.n_local_heads * self.head_dim kv_size = self.n_local_kv_heads * self.head_dim xq, xk, xv = self.wqkv(x).split([q_size, kv_size, kv_size], dim=-1) else: # Compute query, key, and value projections xq = self.wq(x) xk, xv = self.wk(context), self.wv(context) # Reshape projections xq = xq.view(bsz, seqlen, self.n_local_heads, self.head_dim) xk = xk.view(bsz, context_len, self.n_local_kv_heads, self.head_dim) xv = xv.view(bsz, context_len, self.n_local_kv_heads, self.head_dim) # QK normalization if self.use_qk_normalization: xq = self.q_norm(xq) xk = self.k_norm(xk) # Apply rotary positional embeddings to queries and keys # Only apply RoPE to self-attention! if self.attention_type in ["self", "full"]: xq, xk = rope(xq, xk, input_pos, seqlen) xq, xk, xv = map(lambda x: x.transpose(1, 2), (xq, xk, xv)) # xq: (bs, n_local_heads, seqlen, head_dim) # xk: (bs, n_kv_heads, cache_len + context_len, head_dim) # xv: (bs, n_kv_heads, cache_len + context_len, head_dim) if self.inference and self.attention_type == "self": # Update cache with current key and value tensors assert input_pos is not None self.cache_k[:bsz, :, input_pos] = xk self.cache_v[:bsz, :, input_pos] = xv keys, values = ( self.cache_k[:bsz, :, :], self.cache_v[:bsz, :, :], ) else: keys, values = xk, xv # Repeat keys and values if necessary keys = keys.repeat_interleave(self.n_rep, dim=1) # (bs, n_local_heads, cache_len + context_len, head_dim) values = values.repeat_interleave(self.n_rep, dim=1) # (bs, n_local_heads, cache_len + context_len, head_dim) if self.attention_type == "self" and self.causal_mask: # During inference, `is_causal` should be set to False when KV cache is pre-computed and used, # since the masking is handled outside this attention module. # During training, `is_causal` should be set to None to use the default behavior of FlashAttention. is_causal = False if self.inference else None else: # This is used for full-attention transformer (e.g., ViT) # also for the cross-attn, it's always full-attn w/o causal is_causal = False output = scaled_dot_product_attention( xq, keys, values, head_dim=self.head_dim, mask=mask, is_causal=is_causal, dropout_p=self.attention_dropout if self.training else 0.0, ) output = output.view(bsz, seqlen, -1) output = self.wo(output) if self.inference and self.tp_size > 1: output = funcol.all_reduce(output, "sum", group=parallel_state.get_tensor_model_parallel_group()) return output def init_weights(self, init_std: float): """ Initializes the weights of all modules. """ if self.fuse_qkv: nn.init.trunc_normal_(self.wqkv.weight, mean=0.0, std=0.02) else: for linear in (self.wq, self.wk, self.wv): nn.init.trunc_normal_(linear.weight, mean=0.0, std=0.02) nn.init.trunc_normal_(self.wo.weight, mean=0.0, std=init_std) if self.use_qk_normalization: torch.nn.init.ones_(self.q_norm.weight) torch.nn.init.ones_(self.k_norm.weight) def scaled_dot_product_attention( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, head_dim: int, mask: Optional[torch.Tensor] = None, is_causal: Optional[bool] = None, dropout_p: float = 0.0, ) -> torch.Tensor: """ PyTorch's native implementation of Flash Attention 2. If `is_causal` is given, then the causal attention mask is applied accordingly: - If `is_causal` is True, the standard upper-left causal attention masking is applied. - If `is_causal` is False, no attention mask is applied, unless an explicit mask tensor is provided (i.e., `mask is not None`). If `is_causal` is not given (i.e., `is_causal is None`), then the attention mask is applied based on the provided mask tensor: - If no explicit attention mask is given (i.e., `mask is None`), `is_causal` is set to True, leading to the standard upper-left causal attention masking. - If an attention mask is given (i.e., `mask is not None`), the provided mask is used, and `is_causal` is set to False. Args: q (torch.Tensor): Query tensor k (torch.Tensor): Key tensor v (torch.Tensor): Value tensor head_dim (int): Dimension of each attention head mask (Optional[torch.Tensor], optional): Attention mask. Defaults to None. is_causal (Optional[bool], optional): Whether to apply causal attention mask. Defaults to None. dropout_p (float, optional): Dropout rate. Defaults to 0.0. Returns: torch.Tensor: Output tensor after applying scaled dot-product attention """ scale = 1.0 / math.sqrt(head_dim) if is_causal is None: is_causal = mask is None y = torch.nn.functional.scaled_dot_product_attention( q, k, v, attn_mask=mask, dropout_p=dropout_p, scale=scale, is_causal=is_causal, ) return y.transpose(1, 2).contiguous() def enable_different_context_dim_in_te_ca( te_mha_module, context_dim, args, ): """ Hijacks the MultiheadAttention (MHA) module from TransformerEngine (TE) to use a different context-dim for KV calculation. """ self = te_mha_module common_gemm_kwargs = { "fuse_wgrad_accumulation": args["fuse_wgrad_accumulation"], "tp_group": self.tp_group, "tp_size": self.tp_size, "get_rng_state_tracker": self.get_rng_state_tracker, "sequence_parallel": self.sequence_parallel, "params_dtype": self.params_dtype, } self.key_value = LinearTE( context_dim, 2 * self.hidden_size_kv, init_method=None, bias=args["bias"], return_bias=False, parallel_mode="column" if args["set_parallel_mode"] else None, parameters_split=("key", "value") if not args["fuse_qkv_params"] else None, **common_gemm_kwargs, ) def enable_qk_normalization_in_te_mha( te_mha_module, norm_eps: float, is_self_attn: bool = True, ): """ Hijacks the MultiheadAttention (MHA) module from TransformerEngine (TE) to use our `te_mha_forward_with_qk_norm`. The `te_mha_forward_with_qk_norm` function is just a copy of the TE MHA's forward function (source code at https://github.com/NVIDIA/TransformerEngine/blob/main/transformer_engine/pytorch/attention.py) with the addition of several lines of code for the QK normalization operations. """ self = te_mha_module # First, we add the QK norm layers (RMSNorm class) to the TE's MHA module in advance for our custom forward function. if is_self_attn: common_kwargs = dict( eps=norm_eps, device=self.layernorm_qkv.layer_norm_weight.device, sequence_parallel=self.layernorm_qkv.sequence_parallel, params_dtype=self.layernorm_qkv.layer_norm_weight.dtype, zero_centered_gamma=self.layernorm_qkv.zero_centered_gamma, ) else: common_kwargs = dict( eps=norm_eps, device=self.layernorm_query.query_weight.device, sequence_parallel=self.layernorm_query.sequence_parallel, params_dtype=self.layernorm_query.query_weight.dtype, zero_centered_gamma=self.layernorm_query.zero_centered_gamma, ) if parallel_state.model_parallel_is_initialized() and parallel_state.get_tensor_model_parallel_world_size() > 1: tp_group = parallel_state.get_tensor_model_parallel_group() self.q_norm = AllReduceBWDRMSNormTE( self.hidden_size_per_attention_head, process_group=tp_group, **common_kwargs ) self.k_norm = AllReduceBWDRMSNormTE( self.hidden_size_per_attention_head, process_group=tp_group, **common_kwargs ) else: self.q_norm = RMSNormTE(self.hidden_size_per_attention_head, **common_kwargs) self.k_norm = RMSNormTE(self.hidden_size_per_attention_head, **common_kwargs) # Second, we define the custom forward function for the TE's MHA module, with the QK normalization operations. def te_mha_forward_with_qk_norm( hidden_states: torch.Tensor, attention_mask: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None, encoder_output: Optional[torch.Tensor] = None, attn_mask_type: Optional[str] = None, window_size: Optional[Tuple[int, int]] = None, is_first_microbatch: Optional[bool] = None, checkpoint_core_attention: bool = False, inference_params: Optional[Any] = None, rotary_pos_emb: Optional[Union[torch.Tensor, Tuple[torch.Tensor, torch.Tensor]]] = None, core_attention_bias_type: str = "no_bias", core_attention_bias: Optional[torch.Tensor] = None, alibi_slopes: Optional[torch.Tensor] = None, cu_seqlens_q: Optional[torch.Tensor] = None, cu_seqlens_kv: Optional[torch.Tensor] = None, max_seqlen_q: Optional[int] = None, max_seqlen_kv: Optional[int] = None, fast_zero_fill: bool = True, ) -> Tuple[Union[torch.Tensor, None], ...]: """ Forward propagation for MultiheadAttention layer. """ # hidden_states: [sq, b, h] if attn_mask_type is None: attn_mask_type = self.attn_mask_type if window_size is None: window_size = self.window_size window_size = check_set_window_size(attn_mask_type, window_size) if "padding" in attn_mask_type and attention_mask is not None: for mask in attention_mask: assert mask.dtype == torch.bool, "Attention mask must be in boolean type!" assert ( core_attention_bias_type in AttnBiasTypes ), f"core_attention_bias_type {core_attention_bias_type} is not supported!" # ================================================= # Pre-allocate memory for key-values for inference # ================================================= if inference_params and self.layer_number is not None: if self.layer_number not in inference_params.key_value_memory_dict: inf_max_seq_len = inference_params.max_sequence_length inf_max_batch_size = inference_params.max_batch_size inference_key_memory = self._allocate_memory(inf_max_seq_len, inf_max_batch_size, hidden_states.dtype) inference_value_memory = self._allocate_memory(inf_max_seq_len, inf_max_batch_size, hidden_states.dtype) inference_params.key_value_memory_dict[self.layer_number] = ( inference_key_memory, inference_value_memory, ) else: ( inference_key_memory, inference_value_memory, ) = inference_params.key_value_memory_dict[self.layer_number] # ====================== # Query, Key, and Value # ====================== # fp8_mha = FP8GlobalStateManager.is_fp8_enabled() and FP8GlobalStateManager.get_fp8_recipe().fp8_mha # fp8_kwargs = {"fp8_output": fp8_mha and rotary_pos_emb is None} fp8_kwargs = {} layernorm_output = None if self.attention_type == "self": # Attention heads [sq, b, h] --> [sq, b, ng * (np/ng + 2) * hn] layernorm_qkv_outputs = self.layernorm_qkv( hidden_states, is_first_microbatch=is_first_microbatch, **fp8_kwargs ) mixed_x_layer = layernorm_qkv_outputs num_queries_per_key_value = self.num_attention_heads_per_partition // self.num_gqa_groups_per_partition # [sq, b, ng * (np/ng + 2) * hn] --> [sq, b, (np/ng + 2), ng, hn] new_tensor_shape = mixed_x_layer.size()[:-1] + ( (num_queries_per_key_value + 2), self.num_gqa_groups_per_partition, self.hidden_size_per_attention_head, ) # split along third last dimension split_dim = -3 mixed_x_layer = mixed_x_layer.view(*new_tensor_shape) # [sq, b, (np/ng + 2), ng, hn] # --> [sq, b, np/ng, np, hn], [sq, b, 1, ng, hn], [sq, b, 1, ng, hn] query_layer, key_layer, value_layer = _SplitAlongDim.apply( mixed_x_layer, split_dim, (num_queries_per_key_value, 1, 1) ) # query: -> [sq, b, np, hn] # key, value: -> [sq, b, ng, hn] query_layer, key_layer, value_layer = ( x.reshape(x.size(0), x.size(1), -1, self.hidden_size_per_attention_head) for x in (query_layer, key_layer, value_layer) ) elif self.attention_type == "cross": # Attention heads [sk, b, h] --> [sk, b, (ng * 2 * hn)] mixed_kv_layer = self.key_value(encoder_output, is_first_microbatch=is_first_microbatch, **fp8_kwargs) # [sq, b, (ng * 2 * hn)] --> [sq, b, 2 * ng, hn] new_tensor_shape = mixed_kv_layer.size()[:-1] + ( 2 * self.num_gqa_groups_per_partition, self.hidden_size_per_attention_head, ) # split along second last dimension split_dim = -2 mixed_kv_layer = mixed_kv_layer.view(*new_tensor_shape) # mixed_kv_layer --> 2 [sk, b, ng, hn] key_layer, value_layer = _SplitAlongDim.apply( mixed_kv_layer, split_dim, mixed_kv_layer.shape[split_dim] // 2, ) key_layer, value_layer = ( x.reshape( x.size(0), x.size(1), -1, self.hidden_size_per_attention_head, ) for x in (key_layer, value_layer) ) # Attention head [sq, b, h] --> [sq, b, hp] layernorm_query_outputs = self.layernorm_query( hidden_states, is_first_microbatch=is_first_microbatch, **fp8_kwargs ) query_layer = layernorm_query_outputs # [sq, b, hp] --> [sq, b, np, hn] new_tensor_shape = query_layer.size()[:-1] + ( self.num_attention_heads_per_partition, self.hidden_size_per_attention_head, ) query_layer = query_layer.view(*new_tensor_shape) # ====================================================== # Apply QK normalization (RMSNorm) # ====================================================== # Must use torch.reshape to flatten the tensor, otherwise an error will be triggered in TE's RMSNorm module. query_layer = self.q_norm(query_layer.reshape(-1, self.hidden_size_per_attention_head)).view(query_layer.shape) key_layer = self.k_norm(key_layer.reshape(-1, self.hidden_size_per_attention_head)).view(key_layer.shape) # ====================================================== # Apply relative positional encoding (rotary embedding) # ====================================================== if rotary_pos_emb is not None: assert not isinstance(query_layer, Float8Tensor) and not isinstance( key_layer, Float8Tensor ), "RoPE is not supported for Float8Tensors!" # duplicate the pos_emb for self attention if not isinstance(rotary_pos_emb, tuple): rotary_pos_emb = (rotary_pos_emb,) * 2 q_pos_emb, k_pos_emb = rotary_pos_emb # adjust key and value for inference if inference_params is not None: if self.qkv_format == "sbhd": sequence_length = key_layer.size(0) elif self.qkv_format == "bshd": sequence_length = key_layer.size(1) else: raise ValueError(f"QKV format {self.qkv_format} not supported for KV caching.") sequence_start = inference_params.sequence_len_offset sequence_end = sequence_start + sequence_length q_pos_emb = q_pos_emb[sequence_start:sequence_end, ...] k_pos_emb = k_pos_emb[sequence_start:sequence_end, ...] query_layer = apply_rotary_pos_emb(query_layer, q_pos_emb, self.qkv_format, fused=True) key_layer = apply_rotary_pos_emb(key_layer, k_pos_emb, self.qkv_format, fused=True) # =========================== # Core attention computation # =========================== context_layer = self.core_attention( query_layer, key_layer, value_layer, qkv_format=self.qkv_format, cu_seqlens_q=cu_seqlens_q, cu_seqlens_kv=cu_seqlens_kv, max_seqlen_q=max_seqlen_q, max_seqlen_kv=max_seqlen_kv, attention_mask=attention_mask, attn_mask_type=attn_mask_type, window_size=window_size, checkpoint_core_attention=checkpoint_core_attention, core_attention_bias_type=core_attention_bias_type, core_attention_bias=core_attention_bias, alibi_slopes=alibi_slopes, fast_zero_fill=fast_zero_fill, inference_params=inference_params, ) # =================== # Output. [sq, b, h] # =================== projection_output = self.proj( context_layer, is_first_microbatch=is_first_microbatch, ) if self.return_bias: attention_output, attention_bias = projection_output else: attention_output, attention_bias = projection_output, None outputs = (attention_output,) if self.return_bias: outputs += (attention_bias,) if self.input_layernorm and self.return_layernorm_output: outputs += (layernorm_output,) return outputs if len(outputs) > 1 else outputs[0] # Finally, we replace the forward method of given TE's MHA module with our custom forward function. self.forward = te_mha_forward_with_qk_norm def create_group_causal_attn_mask( num_temporal_groups: int, num_query_per_group: int, num_key_per_group: int, mode: str = "causal" ) -> torch.Tensor: """ Creates a group-based attention mask for scaled dot-product attention with two modes: 'causal' and 'group_diagonal'. Parameters: - num_temporal_groups (int): The number of temporal groups (e.g., frames in a video sequence). - num_query_per_group (int): The number of query tokens per temporal group. (e.g., latent tokens in a frame, H x W). - num_key_per_group (int): The number of key tokens per temporal group. (e.g., action tokens per frame). - mode (str): The mode of the attention mask. Options are: - 'causal': Query tokens can attend to key tokens from the same or previous temporal groups. - 'group_diagonal': Query tokens can attend only to key tokens from the same temporal group. Returns: - attn_mask (torch.Tensor): A boolean tensor of shape (L, S), where: - L = num_temporal_groups * num_query_per_group (total number of query tokens) - S = num_temporal_groups * num_key_per_group (total number of key tokens) The mask indicates where attention is allowed (True) and disallowed (False). Example: Input: num_temporal_groups = 3 num_query_per_group = 4 num_key_per_group = 2 Output: Causal Mask Shape: torch.Size([12, 6]) Group Diagonal Mask Shape: torch.Size([12, 6]) if mode='causal': tensor([[ True, True, False, False, False, False], [ True, True, False, False, False, False], [ True, True, False, False, False, False], [ True, True, False, False, False, False], [ True, True, True, True, False, False], [ True, True, True, True, False, False], [ True, True, True, True, False, False], [ True, True, True, True, False, False], [ True, True, True, True, True, True], [ True, True, True, True, True, True], [ True, True, True, True, True, True], [ True, True, True, True, True, True]]) if mode='group_diagonal': tensor([[ True, True, False, False, False, False], [ True, True, False, False, False, False], [ True, True, False, False, False, False], [ True, True, False, False, False, False], [False, False, True, True, False, False], [False, False, True, True, False, False], [False, False, True, True, False, False], [False, False, True, True, False, False], [False, False, False, False, True, True], [False, False, False, False, True, True], [False, False, False, False, True, True], [False, False, False, False, True, True]]) """ assert mode in ["causal", "group_diagonal"], f"Mode {mode} must be 'causal' or 'group_diagonal'" # Total number of query and key tokens total_num_query_tokens = num_temporal_groups * num_query_per_group # Total number of query tokens (L) total_num_key_tokens = num_temporal_groups * num_key_per_group # Total number of key tokens (S) # Generate time indices for query and key tokens (shape: [L] and [S]) query_time_indices = torch.arange(num_temporal_groups).repeat_interleave(num_query_per_group) # Shape: [L] key_time_indices = torch.arange(num_temporal_groups).repeat_interleave(num_key_per_group) # Shape: [S] # Expand dimensions to compute outer comparison query_time_indices = query_time_indices.unsqueeze(1) # Shape: [L, 1] key_time_indices = key_time_indices.unsqueeze(0) # Shape: [1, S] if mode == "causal": # Causal Mode: Query can attend to keys where key_time <= query_time attn_mask = query_time_indices >= key_time_indices # Shape: [L, S] elif mode == "group_diagonal": # Group Diagonal Mode: Query can attend only to keys where key_time == query_time attn_mask = query_time_indices == key_time_indices # Shape: [L, S] assert attn_mask.shape == (total_num_query_tokens, total_num_key_tokens), "Attention mask shape mismatch" return attn_mask