Module neuralops.attentions.triton_cosformer
Triton-Accelerated Cosformer Linear Attention Kernel.
Implements the Cosformer algorithm (Zheng et al., 2022) with hardware-fused OpenAI Triton kernels: 1. In-Register Cosine & Sine Positional Reweighting: - Directly computes cos(pi * i / 2M) and sin(pi * i / 2M) inside GPU registers (tl.cos, tl.sin). - Uses fixed max_length (M = max_position_embeddings) from model config for consistent invariant geometry. - Completely bypasses intermediate DRAM weight buffer allocations.
-
Fused Chunked Causal Cosformer Attention (_triton_chunked_cosformer_fwd):
- Intra-chunk attention score calculation in SRAM tiles.
- Dual inter-chunk state tracking (S_cos, S_sin, Z_cos, Z_sin) in FP32 registers.
- True O(N * d^2) linear complexity with causal lower-triangular masking.
-
Bidirectional Cosformer Kernel (_triton_bidirectional_cosformer_fwd):
- Global parallel reduction over sequence length for encoder backbones (BERT, RoBERTa).
-
Full Autograd Integration (_TritonCosformerAutogradFunction):
- Uses explicit torch.autograd.grad(out, [q_req, k_req, v_req], grad_out) without double-backward warnings.
- Ensures complete differentiability for end-to-end training and recovery fine-tuning.
-
Pure O(N) Fallback for CPU / Mac MPS:
- Evaluates causal attention using chunked streaming prefix sums (chunk_size=64).
- Strictly eliminates O(N^2) memory footprint on non-CUDA platforms.
-
Production-Ready Weight Loading & RoPE:
- Full copy_weights supporting GPT-2 (fused Conv1D c_attn), RoBERTa/BERT, and flat linear layers.
- Complete Rotary Position Embeddings (RoPE) and Grouped Query Attention (GQA) support.
Functions
def apply_rotary_pos_emb(q: torch.Tensor,
k: torch.Tensor,
cos: torch.Tensor,
sin: torch.Tensor,
position_ids: torch.Tensor | None = None) ‑> tuple[torch.Tensor, torch.Tensor]-
Expand source code
def apply_rotary_pos_emb( q: torch.Tensor, k: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor, position_ids: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: if cos.dim() == 2: cos = cos.unsqueeze(0).unsqueeze(0) sin = sin.unsqueeze(0).unsqueeze(0) elif cos.dim() == 3: cos = cos.unsqueeze(1) sin = sin.unsqueeze(1) rotary_dim = cos.shape[-1] q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:] k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:] q_embed = (q_rot * cos) + (_rotate_half(q_rot) * sin) k_embed = (k_rot * cos) + (_rotate_half(k_rot) * sin) if q_pass.shape[-1] > 0: return torch.cat([q_embed, q_pass], dim=-1), torch.cat([k_embed, k_pass], dim=-1) return q_embed, k_embed def inject_triton_cosformer(model: nn.Module = None, rank: int = 64, copy_weights: bool = True, **kwargs: Any) ‑> torch.nn.modules.module.Module | None-
Expand source code
@ATTENTION_REGISTRY.register("inject_triton_cosformer") @ATTENTION_REGISTRY.register("triton_cosformer_inject") def inject_triton_cosformer( model: nn.Module = None, rank: int = 64, copy_weights: bool = True, **kwargs: Any, ) -> nn.Module | None: """Safe factory function to inject TritonCosformerAttention into a model.""" from neuralops.core.attention_utils import inject_attention return inject_attention( model, "triton_cosformer", rank=rank, copy_weights=copy_weights, **kwargs )Safe factory function to inject TritonCosformerAttention into a model.
def triton_cosformer_attention(q: torch.Tensor,
k: torch.Tensor,
v: torch.Tensor,
is_decoder: bool = False,
eps: float = 1e-06,
max_length: int = 2048) ‑> torch.Tensor-
Expand source code
def triton_cosformer_attention( q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, is_decoder: bool = False, eps: float = 1e-6, max_length: int = 2048, ) -> torch.Tensor: """Convenience autograd functional wrapper for Triton Cosformer.""" return _TritonCosformerAutogradFunction.apply(q, k, v, is_decoder, eps, max_length)Convenience autograd functional wrapper for Triton Cosformer.
Classes
class TritonCosformerAttention (config: Any = None,
embed_dim: int | None = None,
num_heads: int | None = None,
num_kv_heads: int | None = None,
is_decoder: bool = False,
eps: float = 1e-06,
pretrained_attn: nn.Module | None = None,
**kwargs: Any)-
Expand source code
@ATTENTION_REGISTRY.register("triton_cosformer", supports_causal=True, supports_bidirectional=True) class TritonCosformerAttention(BaseAttention): r""" Triton-Accelerated Cosformer Linear Attention with Full Autograd Differentiability, RoPE support, Pretrained Weight Loading, and Masked Attention. """ triton_kernel = "linear_attn" def __init__( self, config: Any = None, embed_dim: int | None = None, num_heads: int | None = None, num_kv_heads: int | None = None, is_decoder: bool = False, eps: float = 1e-6, pretrained_attn: nn.Module | None = None, **kwargs: Any, ): super().__init__() self.config = config self.is_decoder = bool( is_decoder or getattr(config, "is_decoder", False) or getattr(config, "is_causal", False) ) self.eps = eps self.triton_kwargs = { "phi": "elu+1", "eps": self.eps, "cos_weight": None, "sin_weight": None, } self.embed_dim = ( embed_dim or kwargs.get("hidden_size") or getattr(config, "hidden_size", 768) ) self.num_heads = ( num_heads or kwargs.get("num_attention_heads") or getattr(config, "num_attention_heads", 12) ) self.num_kv_heads = ( num_kv_heads or kwargs.get("num_kv_heads") or getattr(config, "num_key_value_heads", self.num_heads) ) self.head_dim = kwargs.get("head_dim") or (self.embed_dim // self.num_heads) # BUG 5 FIX: fixed max_length from config self.max_length = getattr(config, "max_position_embeddings", 2048) or 2048 # Linear projections self.q_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True) self.kv_dim = self.num_kv_heads * self.head_dim self.k_proj = nn.Linear(self.embed_dim, self.kv_dim, bias=True) self.v_proj = nn.Linear(self.embed_dim, self.kv_dim, bias=True) self.out_proj = nn.Linear(self.embed_dim, self.embed_dim, bias=True) for m in (self.q_proj, self.k_proj, self.v_proj, self.out_proj): nn.init.xavier_uniform_(m.weight) if m.bias is not None: nn.init.zeros_(m.bias) # BUG 3 FIX: Support pretrained attention weight copying during initialization if pretrained_attn is not None: self.copy_weights(pretrained_attn) def copy_weights(self, src: Any) -> None: """Transfer pretrained projection weights from GPT-2, RoBERTa/BERT, or standard linear layers.""" if src is None: return getattr(src, "__class__", type(src)).__name__ # 1. GPT-2: Fused c_attn (Conv1D) + c_proj if hasattr(src, "c_attn") and src.c_attn is not None: c_attn_mod = src.c_attn base_c = getattr(c_attn_mod, "base_layer", c_attn_mod) c_w = getattr(base_c, "weight", getattr(c_attn_mod, "weight", None)) if c_w is not None: c_w = c_w.data.clone() # If Conv1D or shape is [in, out], transpose to standard PyTorch [out, in] if ( c_attn_mod.__class__.__name__ == "Conv1D" or base_c.__class__.__name__ == "Conv1D" or hasattr(c_attn_mod, "nf") or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1]) ): c_w = c_w.t().contiguous() split_size = c_w.size(0) // 3 bias_mod = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None)) c_b = bias_mod.data.clone() if bias_mod is not None else None self.q_proj.weight.data.copy_(c_w[:split_size, :]) self.k_proj.weight.data.copy_(c_w[split_size : 2 * split_size, :]) self.v_proj.weight.data.copy_(c_w[2 * split_size :, :]) if c_b is not None: if self.q_proj.bias is not None: self.q_proj.bias.data.copy_(c_b[:split_size]) if self.k_proj.bias is not None: self.k_proj.bias.data.copy_(c_b[split_size : 2 * split_size]) if self.v_proj.bias is not None: self.v_proj.bias.data.copy_(c_b[2 * split_size :]) if hasattr(src, "c_proj") and src.c_proj is not None: c_proj_mod = src.c_proj base_o = getattr(c_proj_mod, "base_layer", c_proj_mod) o_w = getattr(base_o, "weight", getattr(c_proj_mod, "weight", None)) if o_w is not None: o_w = o_w.data.clone() if ( c_proj_mod.__class__.__name__ == "Conv1D" or base_o.__class__.__name__ == "Conv1D" or hasattr(c_proj_mod, "nf") or (o_w.ndim == 2 and o_w.shape[0] < o_w.shape[1]) ): o_w = o_w.t().contiguous() self.out_proj.weight.data.copy_(o_w) bias_o = getattr(base_o, "bias", getattr(c_proj_mod, "bias", None)) if bias_o is not None and self.out_proj.bias is not None: self.out_proj.bias.data.copy_(bias_o.data) # 2. RoBERTa / BERT nested self-attention elif hasattr(src, "self") and hasattr(src.self, "query"): s = src.self mapping = [("q_proj", "query"), ("k_proj", "key"), ("v_proj", "value")] for target_name, sub_name in mapping: target_mod = getattr(self, target_name) src_mod = getattr(s, sub_name, None) base_src = getattr(src_mod, "base_layer", src_mod) if base_src is not None and hasattr(base_src, "weight"): target_mod.weight.data.copy_(base_src.weight.data) if base_src.bias is not None and target_mod.bias is not None: target_mod.bias.data.copy_(base_src.bias.data) if hasattr(src, "output") and hasattr(src.output, "dense"): out_mod = src.output.dense base_out = getattr(out_mod, "base_layer", out_mod) self.out_proj.weight.data.copy_(base_out.weight.data) if base_out.bias is not None and self.out_proj.bias is not None: self.out_proj.bias.data.copy_(base_out.bias.data) # 3. Standard flat linear projections (DistilBERT, Qwen, LLaMA) else: proj_candidates = { "q_proj": ["q_proj", "query", "q_lin", "W_q", "wq"], "k_proj": ["k_proj", "key", "k_lin", "W_k", "wk"], "v_proj": ["v_proj", "value", "v_lin", "W_v", "wv"], "out_proj": ["out_proj", "dense", "out_lin", "c_proj", "o_proj", "wo"], } for target_name, candidates in proj_candidates.items(): target_mod = getattr(self, target_name) for cand in candidates: if hasattr(src, cand): src_mod = getattr(src, cand) base_mod = getattr(src_mod, "base_layer", src_mod) if base_mod is not None and hasattr(base_mod, "weight"): w = base_mod.weight.data if ( w.ndim == 2 and w.shape != target_mod.weight.shape and w.t().shape == target_mod.weight.shape ): w = w.t() target_mod.weight.data.copy_(w) if ( getattr(base_mod, "bias", None) is not None and target_mod.bias is not None ): target_mod.bias.data.copy_(base_mod.bias.data) break # Alias for API uniformity copy_weights_from_pretrained = copy_weights def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, position_ids: torch.Tensor | None = None, past_key_value: Any | None = None, *args: Any, **kwargs: Any, ) -> tuple[torch.Tensor, Any | None]: B, N, C = hidden_states.shape # RoPE from kwargs if not explicit if position_embeddings is None: position_embeddings = kwargs.get("position_embeddings") # 1. Project Q, K, V q = ( self.q_proj(hidden_states) .view(B, N, self.num_heads, self.head_dim) .transpose(1, 2) .contiguous() ) k = ( self.k_proj(hidden_states) .view(B, N, self.num_kv_heads, self.head_dim) .transpose(1, 2) .contiguous() ) v = ( self.v_proj(hidden_states) .view(B, N, self.num_kv_heads, self.head_dim) .transpose(1, 2) .contiguous() ) # 2. BUG 4 FIX: Apply Rotary Position Embeddings (RoPE) if provided if position_embeddings is not None: cos, sin = position_embeddings q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids=position_ids) # 3. GQA expansion if num_kv_heads < num_heads if self.num_kv_heads != self.num_heads: k = repeat_kv(k, self.num_heads // self.num_kv_heads) v = repeat_kv(v, self.num_heads // self.num_kv_heads) # 4. BUG 2 FIX: Mask handling for padding tokens on K and V if attention_mask is not None: if attention_mask.dim() == 2: # Shape (B, N) -> (B, 1, N, 1) mask = attention_mask[:, None, :, None].to(q.dtype) if (mask < 0).any(): mask = (mask > -1000.0).to(q.dtype) k = k * mask v = v * mask elif attention_mask.dim() == 4: if attention_mask.shape[-2] == 1: mask = attention_mask.squeeze(1).squeeze(1)[:, None, :, None].to(q.dtype) else: mask = attention_mask[:, 0, -1, :][:, None, :, None].to(q.dtype) if (mask < 0).any(): mask = (mask > -1000.0).to(q.dtype) k = k * mask v = v * mask # 5. Differentiable forward pass using cosformer's own autograd function. # NOTE: We call _TritonCosformerAutogradFunction.apply() directly instead # of routing through triton_dispatch → linear_attn.fwd, because: # - linear_attn.fwd early-exits to PyTorch fallback when cos_weight is set # - That path was broken (ctx.save_for_backward not called before return) # - This module has its own complete Triton + PyTorch fallback built in. out = _TritonCosformerAutogradFunction.apply( q, k, v, self.is_decoder, self.eps, self.max_length ) out = out.transpose(1, 2).contiguous().view(B, N, C) out = self.out_proj(out) return out, NoneTriton-Accelerated Cosformer Linear Attention with Full Autograd Differentiability, RoPE support, Pretrained Weight Loading, and Masked Attention.
Initialize internal Module state, shared by both nn.Module and ScriptModule.
Ancestors
- BaseAttention
- torch.nn.modules.module.Module
Class variables
var triton_kernel-
The type of the None singleton.
Methods
def copy_weights(self, src: Any) ‑> None-
Expand source code
def copy_weights(self, src: Any) -> None: """Transfer pretrained projection weights from GPT-2, RoBERTa/BERT, or standard linear layers.""" if src is None: return getattr(src, "__class__", type(src)).__name__ # 1. GPT-2: Fused c_attn (Conv1D) + c_proj if hasattr(src, "c_attn") and src.c_attn is not None: c_attn_mod = src.c_attn base_c = getattr(c_attn_mod, "base_layer", c_attn_mod) c_w = getattr(base_c, "weight", getattr(c_attn_mod, "weight", None)) if c_w is not None: c_w = c_w.data.clone() # If Conv1D or shape is [in, out], transpose to standard PyTorch [out, in] if ( c_attn_mod.__class__.__name__ == "Conv1D" or base_c.__class__.__name__ == "Conv1D" or hasattr(c_attn_mod, "nf") or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1]) ): c_w = c_w.t().contiguous() split_size = c_w.size(0) // 3 bias_mod = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None)) c_b = bias_mod.data.clone() if bias_mod is not None else None self.q_proj.weight.data.copy_(c_w[:split_size, :]) self.k_proj.weight.data.copy_(c_w[split_size : 2 * split_size, :]) self.v_proj.weight.data.copy_(c_w[2 * split_size :, :]) if c_b is not None: if self.q_proj.bias is not None: self.q_proj.bias.data.copy_(c_b[:split_size]) if self.k_proj.bias is not None: self.k_proj.bias.data.copy_(c_b[split_size : 2 * split_size]) if self.v_proj.bias is not None: self.v_proj.bias.data.copy_(c_b[2 * split_size :]) if hasattr(src, "c_proj") and src.c_proj is not None: c_proj_mod = src.c_proj base_o = getattr(c_proj_mod, "base_layer", c_proj_mod) o_w = getattr(base_o, "weight", getattr(c_proj_mod, "weight", None)) if o_w is not None: o_w = o_w.data.clone() if ( c_proj_mod.__class__.__name__ == "Conv1D" or base_o.__class__.__name__ == "Conv1D" or hasattr(c_proj_mod, "nf") or (o_w.ndim == 2 and o_w.shape[0] < o_w.shape[1]) ): o_w = o_w.t().contiguous() self.out_proj.weight.data.copy_(o_w) bias_o = getattr(base_o, "bias", getattr(c_proj_mod, "bias", None)) if bias_o is not None and self.out_proj.bias is not None: self.out_proj.bias.data.copy_(bias_o.data) # 2. RoBERTa / BERT nested self-attention elif hasattr(src, "self") and hasattr(src.self, "query"): s = src.self mapping = [("q_proj", "query"), ("k_proj", "key"), ("v_proj", "value")] for target_name, sub_name in mapping: target_mod = getattr(self, target_name) src_mod = getattr(s, sub_name, None) base_src = getattr(src_mod, "base_layer", src_mod) if base_src is not None and hasattr(base_src, "weight"): target_mod.weight.data.copy_(base_src.weight.data) if base_src.bias is not None and target_mod.bias is not None: target_mod.bias.data.copy_(base_src.bias.data) if hasattr(src, "output") and hasattr(src.output, "dense"): out_mod = src.output.dense base_out = getattr(out_mod, "base_layer", out_mod) self.out_proj.weight.data.copy_(base_out.weight.data) if base_out.bias is not None and self.out_proj.bias is not None: self.out_proj.bias.data.copy_(base_out.bias.data) # 3. Standard flat linear projections (DistilBERT, Qwen, LLaMA) else: proj_candidates = { "q_proj": ["q_proj", "query", "q_lin", "W_q", "wq"], "k_proj": ["k_proj", "key", "k_lin", "W_k", "wk"], "v_proj": ["v_proj", "value", "v_lin", "W_v", "wv"], "out_proj": ["out_proj", "dense", "out_lin", "c_proj", "o_proj", "wo"], } for target_name, candidates in proj_candidates.items(): target_mod = getattr(self, target_name) for cand in candidates: if hasattr(src, cand): src_mod = getattr(src, cand) base_mod = getattr(src_mod, "base_layer", src_mod) if base_mod is not None and hasattr(base_mod, "weight"): w = base_mod.weight.data if ( w.ndim == 2 and w.shape != target_mod.weight.shape and w.t().shape == target_mod.weight.shape ): w = w.t() target_mod.weight.data.copy_(w) if ( getattr(base_mod, "bias", None) is not None and target_mod.bias is not None ): target_mod.bias.data.copy_(base_mod.bias.data) breakTransfer pretrained projection weights from GPT-2, RoBERTa/BERT, or standard linear layers.
def copy_weights_from_pretrained(self, src: Any) ‑> None-
Expand source code
def copy_weights(self, src: Any) -> None: """Transfer pretrained projection weights from GPT-2, RoBERTa/BERT, or standard linear layers.""" if src is None: return getattr(src, "__class__", type(src)).__name__ # 1. GPT-2: Fused c_attn (Conv1D) + c_proj if hasattr(src, "c_attn") and src.c_attn is not None: c_attn_mod = src.c_attn base_c = getattr(c_attn_mod, "base_layer", c_attn_mod) c_w = getattr(base_c, "weight", getattr(c_attn_mod, "weight", None)) if c_w is not None: c_w = c_w.data.clone() # If Conv1D or shape is [in, out], transpose to standard PyTorch [out, in] if ( c_attn_mod.__class__.__name__ == "Conv1D" or base_c.__class__.__name__ == "Conv1D" or hasattr(c_attn_mod, "nf") or (c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1]) ): c_w = c_w.t().contiguous() split_size = c_w.size(0) // 3 bias_mod = getattr(base_c, "bias", getattr(c_attn_mod, "bias", None)) c_b = bias_mod.data.clone() if bias_mod is not None else None self.q_proj.weight.data.copy_(c_w[:split_size, :]) self.k_proj.weight.data.copy_(c_w[split_size : 2 * split_size, :]) self.v_proj.weight.data.copy_(c_w[2 * split_size :, :]) if c_b is not None: if self.q_proj.bias is not None: self.q_proj.bias.data.copy_(c_b[:split_size]) if self.k_proj.bias is not None: self.k_proj.bias.data.copy_(c_b[split_size : 2 * split_size]) if self.v_proj.bias is not None: self.v_proj.bias.data.copy_(c_b[2 * split_size :]) if hasattr(src, "c_proj") and src.c_proj is not None: c_proj_mod = src.c_proj base_o = getattr(c_proj_mod, "base_layer", c_proj_mod) o_w = getattr(base_o, "weight", getattr(c_proj_mod, "weight", None)) if o_w is not None: o_w = o_w.data.clone() if ( c_proj_mod.__class__.__name__ == "Conv1D" or base_o.__class__.__name__ == "Conv1D" or hasattr(c_proj_mod, "nf") or (o_w.ndim == 2 and o_w.shape[0] < o_w.shape[1]) ): o_w = o_w.t().contiguous() self.out_proj.weight.data.copy_(o_w) bias_o = getattr(base_o, "bias", getattr(c_proj_mod, "bias", None)) if bias_o is not None and self.out_proj.bias is not None: self.out_proj.bias.data.copy_(bias_o.data) # 2. RoBERTa / BERT nested self-attention elif hasattr(src, "self") and hasattr(src.self, "query"): s = src.self mapping = [("q_proj", "query"), ("k_proj", "key"), ("v_proj", "value")] for target_name, sub_name in mapping: target_mod = getattr(self, target_name) src_mod = getattr(s, sub_name, None) base_src = getattr(src_mod, "base_layer", src_mod) if base_src is not None and hasattr(base_src, "weight"): target_mod.weight.data.copy_(base_src.weight.data) if base_src.bias is not None and target_mod.bias is not None: target_mod.bias.data.copy_(base_src.bias.data) if hasattr(src, "output") and hasattr(src.output, "dense"): out_mod = src.output.dense base_out = getattr(out_mod, "base_layer", out_mod) self.out_proj.weight.data.copy_(base_out.weight.data) if base_out.bias is not None and self.out_proj.bias is not None: self.out_proj.bias.data.copy_(base_out.bias.data) # 3. Standard flat linear projections (DistilBERT, Qwen, LLaMA) else: proj_candidates = { "q_proj": ["q_proj", "query", "q_lin", "W_q", "wq"], "k_proj": ["k_proj", "key", "k_lin", "W_k", "wk"], "v_proj": ["v_proj", "value", "v_lin", "W_v", "wv"], "out_proj": ["out_proj", "dense", "out_lin", "c_proj", "o_proj", "wo"], } for target_name, candidates in proj_candidates.items(): target_mod = getattr(self, target_name) for cand in candidates: if hasattr(src, cand): src_mod = getattr(src, cand) base_mod = getattr(src_mod, "base_layer", src_mod) if base_mod is not None and hasattr(base_mod, "weight"): w = base_mod.weight.data if ( w.ndim == 2 and w.shape != target_mod.weight.shape and w.t().shape == target_mod.weight.shape ): w = w.t() target_mod.weight.data.copy_(w) if ( getattr(base_mod, "bias", None) is not None and target_mod.bias is not None ): target_mod.bias.data.copy_(base_mod.bias.data) breakTransfer pretrained projection weights from GPT-2, RoBERTa/BERT, or standard linear layers.
Inherited members