Module neuralops.attentions.flash_attention3_fp8
FlashAttention-3 FP8: Low-Precision 8-bit Float Transformer Attention.
Implements FlashAttention-3 utilizing FP8 (e4m3fn) precision format for NVIDIA Hopper (SM 9.0+) Tensor Cores with per-tensor dynamic scaling and de-scaling, with automatic fallback to hardware-accelerated SDPA on pre-Hopper architectures (such as Ada Lovelace SM 8.9 or Ampere SM 8.0) and CPU.
NOTE ON METHODOLOGY: FlashAttention-3 FP8 is a hardware execution acceleration technique, NOT a model compression technique. It does not reduce parameter count or weight storage (unless LoRA adapters are configured). Base model weights are maintained in native precision and activation tensors are dynamically scaled and cast to FP8 during Hopper GPU attention computation.
Hardened for production across DistilBERT, RoBERTa/BERT, GPT-2, and Qwen/LLaMA architectures: - Robust pretrained weight extraction (GPT-2 fused c_attn, RoBERTa nested, Qwen GQA, DistilBERT) - Strict Frobenius norm verification (>= 1e-4) for all projections - Strict validation requiring all mandatory Q/K/V projections (raises RuntimeError if missing) - Grouped Query Attention (GQA) with kv_head_size projections and repeat_kv - Rotary Position Embedding (RoPE) support with partial-coordinate slicing - Strict causal masking for decoder models across FlashAttention and SDPA paths - Nested encoder detection to bypass out_proj in inner self-attention - Parameter-efficient fine-tuning via LoRALinear adapters and frozen base weights - Safe FP8 casting with per-tensor dynamic scaling (max absolute value / 448.0) and de-scaling - Hardware capability gating (Hopper compute capability >= 9.0 check) - Full cross-attention 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]: r"""Applies rotary position embeddings with partial-RoPE dimension slicing.""" if cos is None or sin is None: return q, k 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] head_dim = q.shape[-1] if rotary_dim < head_dim: q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:] k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:] q_rot_embed = (q_rot * cos) + (_rotate_half(q_rot) * sin) k_rot_embed = (k_rot * cos) + (_rotate_half(k_rot) * sin) q = torch.cat([q_rot_embed, q_pass], dim=-1) k = torch.cat([k_rot_embed, k_pass], dim=-1) else: q = (q * cos) + (_rotate_half(q) * sin) k = (k * cos) + (_rotate_half(k) * sin) return q, kApplies rotary position embeddings with partial-RoPE dimension slicing.
def inject_flash3_fp8(layer: Any = None,
model: Any = None,
model_config: Any = None,
comp_config: Any = None,
layer_idx: int = 0,
total_layers: int = 12,
old_attn: nn.Module | None = None,
**kwargs: Any)-
Expand source code
@ATTENTION_REGISTRY.register("flash3_fp8_inject") def inject_flash3_fp8( layer: Any = None, model: Any = None, model_config: Any = None, comp_config: Any = None, layer_idx: int = 0, total_layers: int = 12, old_attn: nn.Module | None = None, **kwargs: Any, ): target = layer if layer is not None else model if target is None: target = kwargs.get("model", kwargs.get("layer")) if target is None: return None cfg = model_config if model_config is not None else getattr(target, "config", None) # Whole-model traversal layer_list = None if hasattr(target, "distilbert") and hasattr(target.distilbert, "transformer"): layer_list = target.distilbert.transformer.layer elif hasattr(target, "roberta") and hasattr(target.roberta, "encoder"): layer_list = target.roberta.encoder.layer elif hasattr(target, "transformer") and hasattr(target.transformer, "h"): layer_list = target.transformer.h elif hasattr(target, "h") and isinstance(target.h, (torch.nn.ModuleList, list)): layer_list = target.h elif hasattr(target, "model") and hasattr(target.model, "layers"): layer_list = target.model.layers elif hasattr(target, "encoder") and hasattr(target.encoder, "layer"): layer_list = target.encoder.layer elif ( hasattr(target, "bert") and hasattr(target.bert, "encoder") and hasattr(target.bert.encoder, "layer") ): layer_list = target.bert.encoder.layer elif hasattr(target, "layers") and isinstance(target.layers, (torch.nn.ModuleList, list)): layer_list = target.layers if layer_list is not None: num_injected = 0 tot = len(layer_list) for idx, lyr in enumerate(layer_list): res = inject_flash3_fp8( layer=lyr, model_config=cfg, comp_config=( comp_config if "comp_config" in locals() or "comp_config" in kwargs else None ), layer_idx=idx, total_layers=tot, **kwargs, ) if res is not None: num_injected += 1 return target if num_injected > 0 else None layer = target r"""Safe injection factory for FlashAttention3FP8.""" src_attn = old_attn if src_attn is None: if hasattr(layer, "attention"): src_attn = layer.attention elif hasattr(layer, "attn"): src_attn = layer.attn elif hasattr(layer, "self_attn"): src_attn = layer.self_attn elif hasattr(layer, "self"): src_attn = layer.self try: new_attn = FlashAttention3FP8( config=model_config, comp_config=comp_config, layer_id=layer_idx, layer_idx=layer_idx, old_attn=src_attn, **kwargs, ) if hasattr(layer, "attention"): if hasattr(layer.attention, "self"): layer.attention.self = new_attn new_attn.is_nested_encoder = True else: layer.attention = new_attn new_attn.is_nested_encoder = False elif hasattr(layer, "self") and hasattr(layer, "output"): layer.self = new_attn new_attn.is_nested_encoder = True elif hasattr(layer, "attn"): layer.attn = new_attn new_attn.is_nested_encoder = False elif hasattr(layer, "self_attn"): layer.self_attn = new_attn new_attn.is_nested_encoder = False return new_attn except Exception as e: logger.error(f"FlashAttention3FP8 injection failed: {e}") return None def repeat_kv(hidden_states: torch.Tensor, n_rep: int) ‑> torch.Tensor-
Expand source code
def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: r"""Expands grouped key/value heads to match query heads.""" if n_rep == 1: return hidden_states batch, num_kv_heads, slen, head_dim = hidden_states.shape hidden_states = hidden_states[:, :, None, :, :].expand( batch, num_kv_heads, n_rep, slen, head_dim ) return hidden_states.reshape(batch, num_kv_heads * n_rep, slen, head_dim)Expands grouped key/value heads to match query heads.
Classes
class FlashAttention3FP8 (config: Any | None = None,
rank: int | None = None,
d_model: int | Any = 768,
n_heads: int = 12,
use_fp8: bool = False,
pretrained_attn: nn.Module | None = None,
old_attn: nn.Module | None = None,
comp_config: Any | None = None,
layer_id: int | None = None,
layer_idx: int | None = None,
**kwargs: Any)-
Expand source code
@ATTENTION_REGISTRY.register( "flash3_fp8", supports_causal=True, supports_bidirectional=True, supports_gqa=True ) @ATTENTION_REGISTRY.register( "flash_attention3_fp8", supports_causal=True, supports_bidirectional=True, supports_gqa=True, ) class FlashAttention3FP8(BaseAttention): r"""FlashAttention-3 FP8 Hardware-Accelerated Multi-Head Attention. Accelerates attention using FP8 (e4m3fn) precision with FlashAttention-3 kernels on Hopper (SM 9.0+) GPUs. Automatically falls back to native PyTorch SDPA on pre-Hopper hardware or CPU. """ is_custom_attention: bool = True def __init__( self, config: Any | None = None, rank: int | None = None, d_model: int | Any = 768, n_heads: int = 12, use_fp8: bool = False, pretrained_attn: nn.Module | None = None, old_attn: nn.Module | None = None, comp_config: Any | None = None, layer_id: int | None = None, layer_idx: int | None = None, **kwargs: Any, ) -> None: source_attn = pretrained_attn if pretrained_attn is not None else old_attn if config is None: class DummyConfig: pass config = DummyConfig() if source_attn is not None: out_layer = getattr( source_attn, "out_lin", getattr( source_attn, "dense", getattr(source_attn, "c_proj", getattr(source_attn, "o_proj", None)), ), ) q_layer = getattr( source_attn, "q_lin", getattr( source_attn, "query", getattr(source_attn, "c_attn", getattr(source_attn, "q_proj", None)), ), ) if out_layer is not None and hasattr(out_layer, "out_features"): config.hidden_size = out_layer.out_features elif q_layer is not None and hasattr(q_layer, "in_features"): config.hidden_size = q_layer.in_features else: config.hidden_size = kwargs.get("hidden_size", kwargs.get("d_model", d_model)) config.num_attention_heads = getattr( source_attn, "n_heads", getattr( source_attn, "num_attention_heads", getattr( source_attn, "num_heads", kwargs.get("num_heads", kwargs.get("num_attention_heads", n_heads)), ), ), ) else: config.hidden_size = kwargs.get("hidden_size", kwargs.get("d_model", d_model)) config.num_attention_heads = kwargs.get( "num_heads", kwargs.get("num_attention_heads", n_heads) ) super().__init__(config, rank) self.num_heads: int = ( getattr( config, "num_attention_heads", getattr( config, "num_heads", kwargs.get("num_heads", kwargs.get("num_attention_heads", n_heads)), ), ) if config is not None else kwargs.get("num_heads", kwargs.get("num_attention_heads", n_heads)) ) self.num_kv_heads: int = kwargs.get( "num_kv_heads", kwargs.get( "num_key_value_heads", ( getattr( config, "num_key_value_heads", getattr(config, "num_kv_heads", self.num_heads), ) if config is not None else self.num_heads ), ), ) hidden_size: int = ( getattr( config, "hidden_size", getattr( config, "d_model", kwargs.get("hidden_size", kwargs.get("d_model", d_model)), ), ) if config is not None else kwargs.get("hidden_size", kwargs.get("d_model", d_model)) ) cfg_head_dim = getattr(config, "head_dim", None) if config is not None else None if cfg_head_dim is None and source_attn is not None: for attr in ("q_proj", "query", "q_lin", "W_q"): mod = getattr(source_attn, attr, None) if mod is not None and hasattr(mod, "weight"): cfg_head_dim = mod.weight.shape[0] // max(1, self.num_heads) break self.head_dim: int = ( int(cfg_head_dim) if cfg_head_dim is not None else kwargs.get("head_dim", hidden_size // max(1, self.num_heads)) ) self.all_head_size: int = self.num_heads * self.head_dim self.kv_head_size: int = self.num_kv_heads * self.head_dim self.kv_repeat: int = max(1, self.num_heads // max(1, self.num_kv_heads)) # Check hardware capability for FP8 (requires Hopper SM 9.0+) self.use_fp8: bool = bool(use_fp8) if self.use_fp8: hopper_capable = False if torch.cuda.is_available() and hasattr(torch.cuda, "get_device_capability"): try: cc = torch.cuda.get_device_capability() hopper_capable = cc[0] >= 9 except Exception: hopper_capable = False if ( not hopper_capable or not FLASH_ATTN_AVAILABLE or not hasattr(torch, "float8_e4m3fn") ): logger.info( "FlashAttention3FP8: Hopper architecture (SM 9.0+), FlashAttention, or torch.float8_e4m3fn " "not available; disabling FP8 and falling back to native SDPA." ) self.use_fp8 = False # Static layer index resolved_layer_id = ( layer_id if layer_id is not None else ( layer_idx if layer_idx is not None else kwargs.get("layer_id", kwargs.get("layer_idx", 0)) ) ) self.layer_idx: int = int(resolved_layer_id) # LoRA adaptation configuration self.lora_rank: int = getattr( config, "lora_rank", getattr(comp_config, "lora_rank", kwargs.get("lora_rank", 16)), ) self.lora_alpha: float = getattr( config, "lora_alpha", getattr(comp_config, "lora_alpha", kwargs.get("lora_alpha", 32.0)), ) self.freeze_base: bool = self.lora_rank > 0 # Decoder & nested encoder detection _mt = str(getattr(config, "model_type", "") or "").lower() if config is not None else "" _decoder_prefixes = ( "gpt", "qwen", "llama", "mistral", "gemma", "phi", "opt", "bloom", "falcon", "olmo", "mpt", "stablelm", ) is_dec = ( ( getattr(config, "is_decoder", False) or getattr(config, "is_causal", False) or any(_mt.startswith(p) for p in _decoder_prefixes) ) if config is not None else False ) if hasattr(config, "architectures") and any( "GPT" in str(a) or "CausalLM" in str(a) for a in (getattr(config, "architectures", []) or []) ): is_dec = True if kwargs.get("is_decoder", False) or kwargs.get("is_causal", False): is_dec = True self.is_decoder: bool = bool(is_dec) self.is_causal: bool = bool(is_dec) is_nested = kwargs.get("is_nested_encoder") if is_nested is None and source_attn is not None: src_cls = source_attn.__class__.__name__ if ( src_cls in [ "BertSelfAttention", "RobertaSelfAttention", "XLMRobertaSelfAttention", "ElectraSelfAttention", "DebertaSelfAttention", ] or (src_cls.endswith("SelfAttention") and "distil" not in src_cls.lower()) or (hasattr(source_attn, "self") and hasattr(source_attn, "output")) ): is_nested = True elif ( hasattr(source_attn, "out_lin") or hasattr(source_attn, "c_proj") or hasattr(source_attn, "o_proj") ): is_nested = False if is_nested is None: is_nested = False self.is_nested_encoder: bool = bool(is_nested) self.gradient_checkpointing: bool = ( getattr(config, "gradient_checkpointing", False) if config is not None else False ) # Multi-head projections equipped with LoRA adapters self.q_proj = LoRALinear( hidden_size, self.all_head_size, bias=True, lora_rank=self.lora_rank, lora_alpha=self.lora_alpha, freeze_base=self.freeze_base, ) self.k_proj = LoRALinear( hidden_size, self.kv_head_size, bias=True, lora_rank=self.lora_rank, lora_alpha=self.lora_alpha, freeze_base=self.freeze_base, ) self.v_proj = LoRALinear( hidden_size, self.kv_head_size, bias=True, lora_rank=self.lora_rank, lora_alpha=self.lora_alpha, freeze_base=self.freeze_base, ) self.out_proj = LoRALinear( self.all_head_size, hidden_size, bias=True, lora_rank=self.lora_rank, lora_alpha=self.lora_alpha, freeze_base=self.freeze_base, ) dropout_p = ( getattr(config, "dropout", getattr(config, "attention_probs_dropout_prob", 0.1)) if config else 0.1 ) self.dropout = nn.Dropout(dropout_p) self.rotary_emb = None self.comp_config: Any | None = comp_config if source_attn is not None: self.copy_weights(source_attn) # Compatibility properties @property def query(self) -> LoRALinear: return self.q_proj @property def key(self) -> LoRALinear: return self.k_proj @property def value(self) -> LoRALinear: return self.v_proj @property def dense(self) -> LoRALinear: return self.out_proj @property def o_proj(self) -> LoRALinear: return self.out_proj @property def q_lin(self) -> LoRALinear: return self.q_proj @property def k_lin(self) -> LoRALinear: return self.k_proj @property def v_lin(self) -> LoRALinear: return self.v_proj @property def out_lin(self) -> LoRALinear: return self.out_proj @property def W_q(self) -> LoRALinear: return self.q_proj @property def W_k(self) -> LoRALinear: return self.k_proj @property def W_v(self) -> LoRALinear: return self.v_proj @property def c_proj(self) -> LoRALinear: return self.out_proj def copy_weights(self, src: Any) -> None: r"""Copy projection weights from pretrained attention module with Frobenius norm verification.""" if src is None: return cls_name = src.__class__.__name__ has_out_lin = hasattr(src, "out_lin") or hasattr(src, "c_proj") or hasattr(src, "o_proj") is_nested = False if has_out_lin: is_nested = False elif ( hasattr(src, "self") and hasattr(src.self, "query") and hasattr(src, "output") and hasattr(src.output, "dense") ) or ( cls_name in [ "BertSelfAttention", "RobertaSelfAttention", "XLMRobertaSelfAttention", "ElectraSelfAttention", "DebertaSelfAttention", ] or (cls_name.endswith("SelfAttention") and "distil" not in cls_name.lower()) ): is_nested = True else: is_nested = False self.is_nested_encoder = bool(is_nested) if hasattr(src, "rotary_emb") and src.rotary_emb is not None: self.rotary_emb = src.rotary_emb extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = { "q": (None, None), "k": (None, None), "v": (None, None), "o": (None, None), } # 1. GPT-2: Fused c_attn + c_proj if hasattr(src, "c_attn") and src.c_attn is not None: c_attn_mod = src.c_attn if hasattr(c_attn_mod, "merge") and not getattr(c_attn_mod, "merged", False): try: c_attn_mod.merge() except Exception as e: logger.warning(f"copy_weights c_attn merge warning: {e}") 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() cls_c = getattr(c_attn_mod, "__class__", type(c_attn_mod)).__name__ base_cls_c = getattr(base_c, "__class__", type(base_c)).__name__ if ( (cls_c == "Conv1D") or (base_cls_c == "Conv1D") or hasattr(c_attn_mod, "nf") or hasattr(base_c, "nf") or ( c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1] and "Linear" not in cls_c and "Linear" not in base_cls_c ) ): 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 extracted["q"] = ( c_w[:split_size, :], c_b[:split_size] if c_b is not None else None, ) extracted["k"] = ( c_w[split_size : 2 * split_size, :], c_b[split_size : 2 * split_size] if c_b is not None else None, ) extracted["v"] = ( c_w[2 * split_size :, :], c_b[2 * split_size :] if c_b is not None else None, ) 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() cls_o = getattr(c_proj_mod, "__class__", type(c_proj_mod)).__name__ base_cls_o = getattr(base_o, "__class__", type(base_o)).__name__ if ( (cls_o == "Conv1D") or (base_cls_o == "Conv1D") or hasattr(c_proj_mod, "nf") or hasattr(base_o, "nf") or ( o_w.ndim == 2 and o_w.shape[0] < o_w.shape[1] and "Linear" not in cls_o and "Linear" not in base_cls_o ) ): o_w = o_w.t().contiguous() bias_o = getattr(base_o, "bias", getattr(c_proj_mod, "bias", None)) extracted["o"] = ( o_w, bias_o.data.clone() if bias_o is not None else None, ) # 2. RoBERTa / BERT nested self-attention elif hasattr(src, "self") and hasattr(src.self, "query"): s = src.self for proj_key, sub_name in [("q", "query"), ("k", "key"), ("v", "value")]: mod = getattr(s, sub_name, None) if mod is not None and hasattr(mod, "weight"): extracted[proj_key] = ( mod.weight.data.clone(), mod.bias.data.clone() if mod.bias is not None else None, ) if hasattr(src, "output") and hasattr(src.output, "dense"): out_mod = src.output.dense extracted["o"] = ( out_mod.weight.data.clone(), out_mod.bias.data.clone() if out_mod.bias is not None else None, ) # 3. Standard linear projections (DistilBERT, Qwen, LLaMA) else: proj_candidates = { "q": ["q_lin", "q_proj", "query", "W_q"], "k": ["k_lin", "k_proj", "key", "W_k"], "v": ["v_lin", "v_proj", "value", "W_v"], "o": ["out_lin", "out_proj", "dense", "o_proj", "proj", "output"], } for proj_key, names in proj_candidates.items(): for name in names: if hasattr(src, name): mod = getattr(src, name) if mod is not None and hasattr(mod, "weight"): extracted[proj_key] = ( mod.weight.data.clone(), mod.bias.data.clone() if mod.bias is not None else None, ) break # Fallback via UniversalAdapter missing = [ k for k, (w, _) in extracted.items() if w is None and (k != "o" or not self.is_nested_encoder) ] if missing: try: from neuralops.core.adapters import UniversalAdapter adapter = UniversalAdapter() for k in missing: w, b = adapter.get_weight_bias(src, k) if w is not None: extracted[k] = (w.clone(), b.clone() if b is not None else None) except Exception as e: logger.debug(f"UniversalAdapter fallback warning: {e}") missing_after = [ k for k, (w, _) in extracted.items() if w is None and (k != "o" or not self.is_nested_encoder) ] if missing_after: raise RuntimeError( f"FlashAttention3FP8 copy_weights: missing mandatory projections {missing_after} " f"from source module {src.__class__.__name__}." ) target_map = { "q": self.q_proj, "k": self.k_proj, "v": self.v_proj, "o": self.out_proj, } with torch.no_grad(): for proj_key, target_mod in target_map.items(): w, b = extracted[proj_key] if w is None: if proj_key == "o" and self.is_nested_encoder: nn.init.eye_(target_mod.weight) if target_mod.bias is not None: nn.init.zeros_(target_mod.bias) continue if w.ndim != 2: w = w.reshape(w.size(0), -1) if ( w.shape != target_mod.weight.shape and target_mod.weight.shape == w.t().contiguous().shape ): w = w.t().contiguous() out_d = min(target_mod.weight.size(0), w.size(0)) in_d = min(target_mod.weight.size(1), w.size(1)) target_mod.weight.data[:out_d, :in_d].copy_( w[:out_d, :in_d].to( device=target_mod.weight.device, dtype=target_mod.weight.dtype ) ) if b is not None and target_mod.bias is not None: min_b = min(target_mod.bias.numel(), b.numel()) target_mod.bias.data[:min_b].copy_( b[:min_b].to(device=target_mod.bias.device, dtype=target_mod.bias.dtype) ) elif target_mod.bias is not None: nn.init.zeros_(target_mod.bias) # Frobenius norm verification (>= 1e-4) w_norm = float(torch.linalg.norm(target_mod.weight.data.float()).item()) if w_norm < 1e-4: logger.warning( f"FlashAttention3FP8 weight copy verification: norm too small ({w_norm:.2e}) for '{proj_key}'" ) if self.freeze_base: target_mod.freeze_base_weights() def forward( self, hidden_states: torch.Tensor, attention_mask: torch.Tensor | None = None, head_mask: torch.Tensor | None = None, encoder_hidden_states: torch.Tensor | None = None, encoder_attention_mask: torch.Tensor | None = None, past_key_value: Any | None = None, output_attentions: bool | None = False, output_hidden_states: bool | None = False, return_dict: bool | None = True, position_ids: torch.Tensor | None = None, position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, *args: Any, **kwargs: Any, ) -> tuple[torch.Tensor, torch.Tensor | None]: B, L, D = hidden_states.shape H, Dh = self.num_heads, self.head_dim device = hidden_states.device dtype = hidden_states.dtype is_cross_attention = encoder_hidden_states is not None kv_input = encoder_hidden_states if is_cross_attention else hidden_states kv_len = kv_input.shape[1] # 1. Projections q = self.q_proj(hidden_states).view(B, L, H, Dh).transpose(1, 2).contiguous() k = ( self.k_proj(kv_input) .view(B, kv_len, self.num_kv_heads, Dh) .transpose(1, 2) .contiguous() ) v = ( self.v_proj(kv_input) .view(B, kv_len, self.num_kv_heads, Dh) .transpose(1, 2) .contiguous() ) # 2. Rotary Position Embeddings (RoPE) if position_embeddings is not None and not is_cross_attention: if isinstance(position_embeddings, tuple): cos, sin = position_embeddings else: cos, sin = position_embeddings, None if sin is not None: q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids) elif self.rotary_emb is not None and not is_cross_attention: if position_ids is not None: cos, sin = self.rotary_emb(v, seq_len=kv_len) q, k = apply_rotary_pos_emb(q, k, cos, sin, position_ids) # 3. Grouped Query Attention (GQA) head expansion if self.num_kv_heads != self.num_heads: k = repeat_kv(k, self.kv_repeat) v = repeat_kv(v, self.kv_repeat) causal_flag = bool(self.is_decoder and not is_cross_attention and kv_len == L) # 4. FlashAttention-3 FP8 hardware acceleration path # Requires Hopper architecture (SM 9.0+), flash_attn, and CUDA tensor out = None if FLASH_ATTN_AVAILABLE and self.use_fp8 and q.is_cuda and hasattr(torch, "float8_e4m3fn"): # Check Hopper compute capability is_hopper = False if hasattr(torch.cuda, "get_device_capability"): try: is_hopper = torch.cuda.get_device_capability(device)[0] >= 9 except Exception: is_hopper = False if is_hopper: try: # FlashAttention expects shape (B, L, H, Dh) q_perm = q.transpose(1, 2).contiguous() k_perm = k.transpose(1, 2).contiguous() v_perm = v.transpose(1, 2).contiguous() # Dynamic per-tensor scaling to FP8 max representable magnitude (448.0 for float8_e4m3fn) fp8_max = 448.0 q_scale = (q_perm.abs().amax() / fp8_max).clamp(min=1e-6) k_scale = (k_perm.abs().amax() / fp8_max).clamp(min=1e-6) v_scale = (v_perm.abs().amax() / fp8_max).clamp(min=1e-6) # Scale and cast to FP8 q_fa = (q_perm / q_scale).to(torch.float8_e4m3fn) k_fa = (k_perm / k_scale).to(torch.float8_e4m3fn) v_fa = (v_perm / v_scale).to(torch.float8_e4m3fn) # Execute FlashAttention kernel out_fp8 = flash_attn_func(q_fa, k_fa, v_fa, causal=causal_flag) # De-scale back: output is mathematically scaled by v_scale out_tensor = (out_fp8.float() * v_scale).to(dtype) # Apply head_mask if provided if head_mask is not None: # out_tensor is (B, L, H, Dh) -> transpose to (B, H, L, Dh) for head masking out_tensor_h = out_tensor.transpose(1, 2) if head_mask.dim() == 1: hm = head_mask.view(1, -1, 1, 1) elif head_mask.dim() == 2: hm = head_mask.view(head_mask.shape[0], head_mask.shape[1], 1, 1) elif head_mask.dim() == 3: hm = head_mask.unsqueeze(-1) else: hm = head_mask out_tensor_h = out_tensor_h * hm out = ( out_tensor_h.transpose(1, 2).contiguous().view(B, L, self.all_head_size) ) else: out = out_tensor.contiguous().view(B, L, self.all_head_size) except Exception as e: logger.warning( f"FlashAttention-3 FP8 execution failed: {e}; falling back to native SDPA." ) out = None # 5. SDPA / PyTorch native fallback path if out is None: eff_mask = encoder_attention_mask if is_cross_attention else attention_mask attn_mask = None if eff_mask is not None: if eff_mask.dim() == 2: min_val = torch.finfo(q.dtype).min if q.dtype.is_floating_point else -1e4 attn_mask = (1.0 - eff_mask.unsqueeze(1).unsqueeze(2).float()) * min_val elif eff_mask.dim() == 3: attn_mask = eff_mask.unsqueeze(1).float() elif eff_mask.dim() == 4: attn_mask = eff_mask.float() if causal_flag and attn_mask is None: out_sdpa = F.scaled_dot_product_attention( q, k, v, attn_mask=None, is_causal=True, dropout_p=self.dropout.p if self.training else 0.0, ) else: if causal_flag: causal_m = torch.triu( torch.full( (L, kv_len), float("-inf"), device=device, dtype=torch.float32, ), diagonal=1, ) attn_mask = ( (attn_mask + causal_m.unsqueeze(0).unsqueeze(0)) if attn_mask is not None else causal_m.unsqueeze(0).unsqueeze(0) ) out_sdpa = F.scaled_dot_product_attention( q, k, v, attn_mask=attn_mask, is_causal=False, dropout_p=self.dropout.p if self.training else 0.0, ) # Head mask handling if head_mask is not None: if head_mask.dim() == 1: head_mask = head_mask.view(1, -1, 1, 1) elif head_mask.dim() == 2: head_mask = head_mask.view(head_mask.shape[0], head_mask.shape[1], 1, 1) elif head_mask.dim() == 3: head_mask = head_mask.unsqueeze(-1) out_sdpa = out_sdpa * head_mask out = out_sdpa.transpose(1, 2).contiguous().view(B, L, self.all_head_size) # 6. Output projection with nested encoder bypass if self.is_nested_encoder: final_out = out else: final_out = self.out_proj(out) return final_out, NoneFlashAttention-3 FP8 Hardware-Accelerated Multi-Head Attention.
Accelerates attention using FP8 (e4m3fn) precision with FlashAttention-3 kernels on Hopper (SM 9.0+) GPUs. Automatically falls back to native PyTorch SDPA on pre-Hopper hardware or CPU.
Initialize internal Module state, shared by both nn.Module and ScriptModule.
Ancestors
- BaseAttention
- torch.nn.modules.module.Module
Class variables
var is_custom_attention : bool-
The type of the None singleton.
Instance variables
prop W_k : LoRALinear-
Expand source code
@property def W_k(self) -> LoRALinear: return self.k_proj prop W_q : LoRALinear-
Expand source code
@property def W_q(self) -> LoRALinear: return self.q_proj prop W_v : LoRALinear-
Expand source code
@property def W_v(self) -> LoRALinear: return self.v_proj prop c_proj : LoRALinear-
Expand source code
@property def c_proj(self) -> LoRALinear: return self.out_proj prop dense : LoRALinear-
Expand source code
@property def dense(self) -> LoRALinear: return self.out_proj prop k_lin : LoRALinear-
Expand source code
@property def k_lin(self) -> LoRALinear: return self.k_proj prop key : LoRALinear-
Expand source code
@property def key(self) -> LoRALinear: return self.k_proj prop o_proj : LoRALinear-
Expand source code
@property def o_proj(self) -> LoRALinear: return self.out_proj prop out_lin : LoRALinear-
Expand source code
@property def out_lin(self) -> LoRALinear: return self.out_proj prop q_lin : LoRALinear-
Expand source code
@property def q_lin(self) -> LoRALinear: return self.q_proj prop query : LoRALinear-
Expand source code
@property def query(self) -> LoRALinear: return self.q_proj prop v_lin : LoRALinear-
Expand source code
@property def v_lin(self) -> LoRALinear: return self.v_proj prop value : LoRALinear-
Expand source code
@property def value(self) -> LoRALinear: return self.v_proj
Methods
def copy_weights(self, src: Any) ‑> None-
Expand source code
def copy_weights(self, src: Any) -> None: r"""Copy projection weights from pretrained attention module with Frobenius norm verification.""" if src is None: return cls_name = src.__class__.__name__ has_out_lin = hasattr(src, "out_lin") or hasattr(src, "c_proj") or hasattr(src, "o_proj") is_nested = False if has_out_lin: is_nested = False elif ( hasattr(src, "self") and hasattr(src.self, "query") and hasattr(src, "output") and hasattr(src.output, "dense") ) or ( cls_name in [ "BertSelfAttention", "RobertaSelfAttention", "XLMRobertaSelfAttention", "ElectraSelfAttention", "DebertaSelfAttention", ] or (cls_name.endswith("SelfAttention") and "distil" not in cls_name.lower()) ): is_nested = True else: is_nested = False self.is_nested_encoder = bool(is_nested) if hasattr(src, "rotary_emb") and src.rotary_emb is not None: self.rotary_emb = src.rotary_emb extracted: dict[str, tuple[torch.Tensor | None, torch.Tensor | None]] = { "q": (None, None), "k": (None, None), "v": (None, None), "o": (None, None), } # 1. GPT-2: Fused c_attn + c_proj if hasattr(src, "c_attn") and src.c_attn is not None: c_attn_mod = src.c_attn if hasattr(c_attn_mod, "merge") and not getattr(c_attn_mod, "merged", False): try: c_attn_mod.merge() except Exception as e: logger.warning(f"copy_weights c_attn merge warning: {e}") 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() cls_c = getattr(c_attn_mod, "__class__", type(c_attn_mod)).__name__ base_cls_c = getattr(base_c, "__class__", type(base_c)).__name__ if ( (cls_c == "Conv1D") or (base_cls_c == "Conv1D") or hasattr(c_attn_mod, "nf") or hasattr(base_c, "nf") or ( c_w.ndim == 2 and c_w.shape[0] < c_w.shape[1] and "Linear" not in cls_c and "Linear" not in base_cls_c ) ): 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 extracted["q"] = ( c_w[:split_size, :], c_b[:split_size] if c_b is not None else None, ) extracted["k"] = ( c_w[split_size : 2 * split_size, :], c_b[split_size : 2 * split_size] if c_b is not None else None, ) extracted["v"] = ( c_w[2 * split_size :, :], c_b[2 * split_size :] if c_b is not None else None, ) 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() cls_o = getattr(c_proj_mod, "__class__", type(c_proj_mod)).__name__ base_cls_o = getattr(base_o, "__class__", type(base_o)).__name__ if ( (cls_o == "Conv1D") or (base_cls_o == "Conv1D") or hasattr(c_proj_mod, "nf") or hasattr(base_o, "nf") or ( o_w.ndim == 2 and o_w.shape[0] < o_w.shape[1] and "Linear" not in cls_o and "Linear" not in base_cls_o ) ): o_w = o_w.t().contiguous() bias_o = getattr(base_o, "bias", getattr(c_proj_mod, "bias", None)) extracted["o"] = ( o_w, bias_o.data.clone() if bias_o is not None else None, ) # 2. RoBERTa / BERT nested self-attention elif hasattr(src, "self") and hasattr(src.self, "query"): s = src.self for proj_key, sub_name in [("q", "query"), ("k", "key"), ("v", "value")]: mod = getattr(s, sub_name, None) if mod is not None and hasattr(mod, "weight"): extracted[proj_key] = ( mod.weight.data.clone(), mod.bias.data.clone() if mod.bias is not None else None, ) if hasattr(src, "output") and hasattr(src.output, "dense"): out_mod = src.output.dense extracted["o"] = ( out_mod.weight.data.clone(), out_mod.bias.data.clone() if out_mod.bias is not None else None, ) # 3. Standard linear projections (DistilBERT, Qwen, LLaMA) else: proj_candidates = { "q": ["q_lin", "q_proj", "query", "W_q"], "k": ["k_lin", "k_proj", "key", "W_k"], "v": ["v_lin", "v_proj", "value", "W_v"], "o": ["out_lin", "out_proj", "dense", "o_proj", "proj", "output"], } for proj_key, names in proj_candidates.items(): for name in names: if hasattr(src, name): mod = getattr(src, name) if mod is not None and hasattr(mod, "weight"): extracted[proj_key] = ( mod.weight.data.clone(), mod.bias.data.clone() if mod.bias is not None else None, ) break # Fallback via UniversalAdapter missing = [ k for k, (w, _) in extracted.items() if w is None and (k != "o" or not self.is_nested_encoder) ] if missing: try: from neuralops.core.adapters import UniversalAdapter adapter = UniversalAdapter() for k in missing: w, b = adapter.get_weight_bias(src, k) if w is not None: extracted[k] = (w.clone(), b.clone() if b is not None else None) except Exception as e: logger.debug(f"UniversalAdapter fallback warning: {e}") missing_after = [ k for k, (w, _) in extracted.items() if w is None and (k != "o" or not self.is_nested_encoder) ] if missing_after: raise RuntimeError( f"FlashAttention3FP8 copy_weights: missing mandatory projections {missing_after} " f"from source module {src.__class__.__name__}." ) target_map = { "q": self.q_proj, "k": self.k_proj, "v": self.v_proj, "o": self.out_proj, } with torch.no_grad(): for proj_key, target_mod in target_map.items(): w, b = extracted[proj_key] if w is None: if proj_key == "o" and self.is_nested_encoder: nn.init.eye_(target_mod.weight) if target_mod.bias is not None: nn.init.zeros_(target_mod.bias) continue if w.ndim != 2: w = w.reshape(w.size(0), -1) if ( w.shape != target_mod.weight.shape and target_mod.weight.shape == w.t().contiguous().shape ): w = w.t().contiguous() out_d = min(target_mod.weight.size(0), w.size(0)) in_d = min(target_mod.weight.size(1), w.size(1)) target_mod.weight.data[:out_d, :in_d].copy_( w[:out_d, :in_d].to( device=target_mod.weight.device, dtype=target_mod.weight.dtype ) ) if b is not None and target_mod.bias is not None: min_b = min(target_mod.bias.numel(), b.numel()) target_mod.bias.data[:min_b].copy_( b[:min_b].to(device=target_mod.bias.device, dtype=target_mod.bias.dtype) ) elif target_mod.bias is not None: nn.init.zeros_(target_mod.bias) # Frobenius norm verification (>= 1e-4) w_norm = float(torch.linalg.norm(target_mod.weight.data.float()).item()) if w_norm < 1e-4: logger.warning( f"FlashAttention3FP8 weight copy verification: norm too small ({w_norm:.2e}) for '{proj_key}'" ) if self.freeze_base: target_mod.freeze_base_weights()Copy projection weights from pretrained attention module with Frobenius norm verification.
Inherited members
class LoRALinear (in_features: int,
out_features: int,
bias: bool = True,
lora_rank: int = 16,
lora_alpha: float = 32.0,
freeze_base: bool = False,
device: torch.device | None = None,
dtype: torch.dtype | None = None)-
Expand source code
class LoRALinear(nn.Module): r"""Linear projection layer with frozen base weights and trainable residual LoRA adapters.""" def __init__( self, in_features: int, out_features: int, bias: bool = True, lora_rank: int = 16, lora_alpha: float = 32.0, freeze_base: bool = False, device: torch.device | None = None, dtype: torch.dtype | None = None, ) -> None: super().__init__() factory_kwargs = {"device": device, "dtype": dtype} self.in_features = in_features self.out_features = out_features self.lora_rank = lora_rank self.lora_alpha = lora_alpha self.lora_scaling = (lora_alpha / lora_rank) if lora_rank > 0 else 1.0 self.freeze_base = freeze_base self.weight = nn.Parameter(torch.empty((out_features, in_features), **factory_kwargs)) if bias: self.bias = nn.Parameter(torch.empty(out_features, **factory_kwargs)) else: self.register_parameter("bias", None) if self.lora_rank > 0: self.lora_A = nn.Parameter(torch.empty((lora_rank, in_features), **factory_kwargs)) self.lora_B = nn.Parameter(torch.empty((out_features, lora_rank), **factory_kwargs)) else: self.register_parameter("lora_A", None) self.register_parameter("lora_B", None) self.reset_parameters() if freeze_base: self.freeze_base_weights() def reset_parameters(self) -> None: nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) if self.bias is not None: fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0 nn.init.uniform_(self.bias, -bound, bound) if self.lora_rank > 0 and self.lora_A is not None: nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) self.lora_A.requires_grad = True if self.lora_rank > 0 and self.lora_B is not None: nn.init.zeros_(self.lora_B) self.lora_B.requires_grad = True def freeze_base_weights(self) -> None: """Freeze base projection parameters so only LoRA parameters adapt.""" self.weight.requires_grad = False if self.bias is not None: self.bias.requires_grad = False if self.lora_rank > 0 and self.lora_A is not None: self.lora_A.requires_grad = True if self.lora_rank > 0 and self.lora_B is not None: self.lora_B.requires_grad = True def unfreeze_base_weights(self) -> None: """Unfreeze base projection parameters.""" self.weight.requires_grad = True if self.bias is not None: self.bias.requires_grad = True def forward(self, x: torch.Tensor) -> torch.Tensor: w = self.weight.to(x.dtype) if self.weight.dtype != x.dtype else self.weight b = ( self.bias.to(x.dtype) if (self.bias is not None and self.bias.dtype != x.dtype) else self.bias ) out = F.linear(x, w, b) if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None: a = self.lora_A.to(x.dtype) if self.lora_A.dtype != x.dtype else self.lora_A bb = self.lora_B.to(x.dtype) if self.lora_B.dtype != x.dtype else self.lora_B lora_out = F.linear(F.linear(x, a), bb) * self.lora_scaling out = out + lora_out return outLinear projection layer with frozen base weights and trainable residual LoRA adapters.
Initialize internal Module state, shared by both nn.Module and ScriptModule.
Ancestors
- torch.nn.modules.module.Module
Methods
def forward(self, x: torch.Tensor) ‑> torch.Tensor-
Expand source code
def forward(self, x: torch.Tensor) -> torch.Tensor: w = self.weight.to(x.dtype) if self.weight.dtype != x.dtype else self.weight b = ( self.bias.to(x.dtype) if (self.bias is not None and self.bias.dtype != x.dtype) else self.bias ) out = F.linear(x, w, b) if self.lora_rank > 0 and self.lora_A is not None and self.lora_B is not None: a = self.lora_A.to(x.dtype) if self.lora_A.dtype != x.dtype else self.lora_A bb = self.lora_B.to(x.dtype) if self.lora_B.dtype != x.dtype else self.lora_B lora_out = F.linear(F.linear(x, a), bb) * self.lora_scaling out = out + lora_out return outDefine the computation performed at every call.
Should be overridden by all subclasses.
Note
Although the recipe for forward pass needs to be defined within this function, one should call the :class:
Moduleinstance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them. def freeze_base_weights(self) ‑> None-
Expand source code
def freeze_base_weights(self) -> None: """Freeze base projection parameters so only LoRA parameters adapt.""" self.weight.requires_grad = False if self.bias is not None: self.bias.requires_grad = False if self.lora_rank > 0 and self.lora_A is not None: self.lora_A.requires_grad = True if self.lora_rank > 0 and self.lora_B is not None: self.lora_B.requires_grad = TrueFreeze base projection parameters so only LoRA parameters adapt.
def reset_parameters(self) ‑> None-
Expand source code
def reset_parameters(self) -> None: nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5)) if self.bias is not None: fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self.weight) bound = 1 / math.sqrt(fan_in) if fan_in > 0 else 0 nn.init.uniform_(self.bias, -bound, bound) if self.lora_rank > 0 and self.lora_A is not None: nn.init.kaiming_uniform_(self.lora_A, a=math.sqrt(5)) self.lora_A.requires_grad = True if self.lora_rank > 0 and self.lora_B is not None: nn.init.zeros_(self.lora_B) self.lora_B.requires_grad = True def unfreeze_base_weights(self) ‑> None-
Expand source code
def unfreeze_base_weights(self) -> None: """Unfreeze base projection parameters.""" self.weight.requires_grad = True if self.bias is not None: self.bias.requires_grad = TrueUnfreeze base projection parameters.