Module neuralops.attentions.fused_svd_kernel
Fused SVD-Linear Layer with SRAM-Resident RMSNorm & Dual-Projection.
RESEARCH FIX: Eliminates GPU High-Bandwidth Memory (HBM) roundtrip latency by fusing: 1. Down-projection: Z = X @ V^T (hidden_dim -> rank) 2. In-SRAM RMSNorm: Z_norm = Z * rsqrt(mean(Z^2) + eps) * weight 3. Up-projection: Y = Z_norm @ U^T + bias (rank -> out_dim)
The intermediate activation matrix Z \in R^{B x L x r} is kept strictly inside GPU Shared Memory (SRAM) registers and never materialized to DRAM/HBM. Includes Triton JIT accelerated kernel with CUDA tensor-core blocking and PyTorch autograd function with fallback.
Functions
def fused_svd_kernel_factory(config: Any = None, **kwargs: Any) ‑> torch.nn.modules.module.Module-
Expand source code
def fused_svd_kernel_factory(config: Any = None, **kwargs: Any) -> nn.Module: """Factory creating FusedSVDKernelAttention with fused SVD kernel execution.""" return FusedSVDKernelAttention(config=config, **kwargs)Factory creating FusedSVDKernelAttention with fused SVD kernel execution.
def fused_svd_linear(x: torch.Tensor,
V: torch.Tensor,
U: torch.Tensor,
norm_weight: torch.Tensor | None = None,
bias: torch.Tensor | None = None,
eps: float = 1e-06,
use_rms_norm: bool = False) ‑> torch.Tensor-
Expand source code
def fused_svd_linear( x: torch.Tensor, V: torch.Tensor, U: torch.Tensor, norm_weight: torch.Tensor | None = None, bias: torch.Tensor | None = None, eps: float = 1e-6, use_rms_norm: bool = False, ) -> torch.Tensor: """ Functional interface for Fused SVD-Linear with optional in-SRAM RMSNorm. Args: x: Input tensor of shape (..., in_features). V: Down-projection matrix of shape (rank, in_features). U: Up-projection matrix of shape (out_features, rank). norm_weight: Optional RMSNorm affine scale vector of shape (rank,). bias: Optional bias vector of shape (out_features,). eps: Small epsilon for RMSNorm numerical stability. use_rms_norm: Whether to apply in-SRAM RMSNorm on intermediate rank activations. Returns: Output tensor of shape (..., out_features). """ return _FusedSVDLinearAutogradFunction.apply(x, V, U, norm_weight, bias, eps, use_rms_norm)Functional interface for Fused SVD-Linear with optional in-SRAM RMSNorm.
- Args
- -----=
x- Input tensor of shape (…, in_features).
V- Down-projection matrix of shape (rank, in_features).
U- Up-projection matrix of shape (out_features, rank).
norm_weight- Optional RMSNorm affine scale vector of shape (rank,).
bias- Optional bias vector of shape (out_features,).
eps- Small epsilon for RMSNorm numerical stability.
use_rms_norm- Whether to apply in-SRAM RMSNorm on intermediate rank activations.
Returns -----= Output tensor of shape (…, out_features).
def inject_fused_svd_kernel(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) ‑> Any | None-
Expand source code
@ATTENTION_REGISTRY.register("fused_svd_kernel_inject") @ATTENTION_REGISTRY.register("inject_fused_svd_kernel") def inject_fused_svd_kernel( 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, ) -> Any | None: """Safe injection hook for fused SVD kernel attention.""" target = layer if layer is not None else model if target is None: return None cfg = model_config or getattr(target, "config", None) # 1. Whole model traversal support if ( old_attn is None and target is not None and ( hasattr(target, "transformer") or hasattr(target, "distilbert") or hasattr(target, "roberta") or hasattr(target, "bert") or (hasattr(target, "model") and hasattr(target.model, "layers")) or (hasattr(target, "encoder") and hasattr(target.encoder, "layer")) ) and not ( hasattr(target, "q_lin") or hasattr(target, "q_proj") or hasattr(target, "c_attn") or hasattr(target, "query") ) ): try: 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, "bert") and hasattr(target.bert, "encoder"): layer_list = target.bert.encoder.layer elif hasattr(target, "transformer") and hasattr(target.transformer, "h"): layer_list = target.transformer.h elif hasattr(target, "transformer") and hasattr(target.transformer, "layer"): layer_list = target.transformer.layer 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 if layer_list is not None: tot = len(layer_list) for idx, lyr in enumerate(layer_list): inject_fused_svd_kernel( layer=lyr, model_config=cfg, comp_config=comp_config, layer_idx=idx, total_layers=tot, **kwargs, ) return target except Exception as e: logger.error("inject_fused_svd_kernel whole-model traversal failed: %s", e) return None # 2. Named modules whole-model fallback if ( old_attn is None and target is not None and hasattr(target, "named_modules") and not ( hasattr(target, "attention") or hasattr(target, "self") or hasattr(target, "attn") or hasattr(target, "q_proj") or hasattr(target, "q_lin") or hasattr(target, "query") or hasattr(target, "c_attn") ) ): try: replaced = 0 for name, module in list(target.named_modules()): if ( hasattr(module, "q_lin") or hasattr(module, "query") or hasattr(module, "q_proj") or hasattr(module, "c_attn") ): parent_name = name.rsplit(".", 1)[0] if "." in name else "" child_name = name.rsplit(".", 1)[-1] if "." in name else name parent = target.get_submodule(parent_name) if parent_name else target new_attn = inject_fused_svd_kernel( layer=None, model_config=cfg, comp_config=comp_config, old_attn=module, **kwargs, ) if new_attn is not None: setattr(parent, child_name, new_attn) replaced += 1 if replaced > 0: return target except Exception as e: logger.error("inject_fused_svd_kernel named_modules replacement failed: %s", e) return None try: src_attn = ( old_attn if old_attn is not None else ( getattr( layer, "attention", getattr(layer, "self", getattr(layer, "attn", None)), ) if layer is not None else None ) ) new_attn = FusedSVDKernelAttention( config=cfg, old_attn=src_attn, **kwargs, ) if src_attn is not None: new_attn.copy_weights(src_attn) if layer is not None: # Transfer new_attn to the target layer's device and dtype target_param = next(layer.parameters(), None) if target_param is not None: new_attn.to(device=target_param.device, dtype=target_param.dtype) 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 elif hasattr(layer, "self"): layer.self = new_attn new_attn.is_nested_encoder = True elif hasattr(layer, "attn"): layer.attn = new_attn return new_attn except Exception as e: logger.error("inject_fused_svd_kernel failed: %s", e) return NoneSafe injection hook for fused SVD kernel attention.
Classes
class FusedSVDKernelAttention (config: Any | None = None,
rank: int | None = None,
old_attn: nn.Module | None = None,
use_fused_kernel: bool = True,
use_rms_norm: bool = False,
dropout: float = 0.1,
**kwargs: Any)-
Expand source code
@ATTENTION_REGISTRY.register( "fused_svd_kernel", supports_causal=True, supports_bidirectional=True, supports_gqa=True, ) @ATTENTION_REGISTRY.register( "fused_svd_kernel_attention", supports_causal=True, supports_bidirectional=True, supports_gqa=True, ) class FusedSVDKernelAttention(BaseAttention): r"""Fused SVD-Kernel Linear Attention Module. Production-hardened across DistilBERT, RoBERTa/BERT nested, GPT-2 Conv1D, and Qwen/LLaMA. """ supports_gqa: bool = True supports_causal: bool = True supports_bidirectional: bool = True is_custom_attention: bool = True def __init__( self, config: Any | None = None, rank: int | None = None, old_attn: nn.Module | None = None, use_fused_kernel: bool = True, use_rms_norm: bool = False, dropout: float = 0.1, **kwargs: Any, ) -> None: h_dim = int( kwargs.get("hidden_size") or _cfg_get(config, "hidden_size", None) or _cfg_get(config, "dim", None) or _cfg_get(config, "d_model", None) or _cfg_get(config, "n_embd", None) or getattr( old_attn if old_attn is not None else kwargs.get("pretrained_attn"), "embed_dim", None, ) or getattr( old_attn if old_attn is not None else kwargs.get("pretrained_attn"), "hidden_size", None, ) or 768 ) self.hidden_size: int = h_dim actual_rank = rank or _cfg_get( config, "rank", _cfg_get(config, "svd_rank", self.hidden_size) ) super().__init__(config, rank=actual_rank) src_attn = old_attn if old_attn is not None else kwargs.get("pretrained_attn") self.num_heads: int = int( kwargs.get("num_heads") or kwargs.get("num_attention_heads") or _cfg_get(config, "num_attention_heads", None) or _cfg_get(config, "n_heads", None) or _cfg_get(config, "num_heads", None) or _cfg_get(config, "n_head", None) or getattr(src_attn, "num_attention_heads", None) or getattr(src_attn, "n_heads", None) or getattr(src_attn, "num_heads", None) or (self.hidden_size // 64 if self.hidden_size % 64 == 0 else 12) ) self.head_dim: int = int( kwargs.get("head_dim") or _cfg_get( config, "head_dim", self.hidden_size // self.num_heads if self.num_heads > 0 else 64, ) ) cfg_kv = _cfg_get(config, "num_key_value_heads", _cfg_get(config, "num_kv_heads", None)) if cfg_kv is not None: self.num_kv_heads = int(cfg_kv) elif kwargs.get("num_kv_heads") is not None: self.num_kv_heads = int(kwargs["num_kv_heads"]) else: self.num_kv_heads = self.num_heads self.num_kv_heads = max(1, min(self.num_kv_heads, self.num_heads)) self.kv_repeat: int = self.num_heads // self.num_kv_heads self.kv_head_size: int = self.num_kv_heads * self.head_dim self.q_head_size: int = self.num_heads * self.head_dim self.all_head_size: int = self.q_head_size self.rank: int = actual_rank self.use_fused_kernel: bool = use_fused_kernel self.use_rms_norm: bool = use_rms_norm is_dec = ( bool(_cfg_get(config, "is_decoder", False)) or any( str(_cfg_get(config, "model_type", "")).lower().startswith(p) for p in ( "gpt", "qwen", "llama", "mistral", "gemma", "phi", "opt", "bloom", "falcon", "olmo", "mpt", "stablelm", ) ) or (src_attn is not None and hasattr(src_attn, "c_attn")) or bool(kwargs.get("is_causal", False)) or bool(kwargs.get("is_decoder", False)) ) self.is_decoder: bool = is_dec self.is_causal: bool = is_dec self.is_nested_encoder: bool = False self.dropout_rate: float = float( _cfg_get( config, "attention_probs_dropout_prob", _cfg_get(config, "dropout", dropout), ) ) self.dropout = nn.Dropout(self.dropout_rate) self.q_proj = FusedSVDLinear( self.hidden_size, self.q_head_size, rank=self.rank, bias=True, use_rms_norm=self.use_rms_norm, ) self.k_proj = FusedSVDLinear( self.hidden_size, self.kv_head_size, rank=self.rank, bias=True, use_rms_norm=self.use_rms_norm, ) self.v_proj = FusedSVDLinear( self.hidden_size, self.kv_head_size, rank=self.rank, bias=True, use_rms_norm=self.use_rms_norm, ) self.out_proj = FusedSVDLinear( self.q_head_size, self.hidden_size, rank=self.rank, bias=True, use_rms_norm=self.use_rms_norm, ) if src_attn is not None: self.copy_weights(src_attn) @property def query(self) -> FusedSVDLinear: return self.q_proj @property def key(self) -> FusedSVDLinear: return self.k_proj @property def value(self) -> FusedSVDLinear: return self.v_proj def copy_weights(self, old_attn: Any) -> None: """Multi-architecture weight copying supporting GPT-2 Conv1D, RoBERTa nested, and flat projections.""" if old_attn is None: return with torch.no_grad(): w_q, b_q = None, None w_k, b_k = None, None w_v, b_v = None, None w_o, b_o = None, None src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__ has_out = ( hasattr(old_attn, "out_lin") or hasattr(old_attn, "dense") or hasattr(old_attn, "out_proj") or hasattr(old_attn, "c_proj") or hasattr(old_attn, "o_proj") ) if ( hasattr(old_attn, "out_lin") or hasattr(old_attn, "c_proj") or hasattr(old_attn, "o_proj") ): self.is_nested_encoder = False elif ( ( hasattr(old_attn, "self") and hasattr(old_attn.self, "query") and hasattr(old_attn, "output") and hasattr(old_attn.output, "dense") ) or src_cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (hasattr(old_attn, "query") and not has_out) ): self.is_nested_encoder = True else: self.is_nested_encoder = False if hasattr(old_attn, "self") and hasattr(old_attn, "output"): self.is_nested_encoder = True if ( hasattr(old_attn, "output") and hasattr(old_attn.output, "dense") and hasattr(old_attn.output.dense, "weight") ): w_o = old_attn.output.dense.weight.data b_o = getattr(old_attn.output.dense, "bias", None) b_o = b_o.data if b_o is not None else None src = old_attn while hasattr(src, "self") and not any( hasattr(src, attr) for attr in ("query", "key", "value", "q_proj", "k_proj", "v_proj") ): self.is_nested_encoder = True src = src.self # 1. GPT-2 Conv1D fused c_attn & c_proj if hasattr(src, "c_attn") and src.c_attn is not None: self.is_decoder = True self.is_causal = True mod_c = src.c_attn base_c = getattr(mod_c, "base_layer", mod_c) w = getattr(base_c, "weight", getattr(mod_c, "weight", None)) if w is not None: w = w.data.clone() cls_c = getattr(mod_c, "__class__", type(mod_c)).__name__ base_cls_c = getattr(base_c, "__class__", type(base_c)).__name__ if ( "Conv1D" in cls_c or "Conv1D" in base_cls_c or hasattr(mod_c, "nf") or hasattr(base_c, "nf") or (w.ndim == 2 and w.shape[0] < w.shape[1]) ): w = w.t().contiguous() q_dim = self.q_head_size k_dim = self.kv_head_size v_dim = self.kv_head_size if w.shape[0] == q_dim + k_dim + v_dim: w_q = w[:q_dim, :] w_k = w[q_dim : q_dim + k_dim, :] w_v = w[q_dim + k_dim :, :] else: split_size = w.shape[0] // 3 w_q = w[:split_size, :] w_k = w[split_size : 2 * split_size, :] w_v = w[2 * split_size :, :] b = getattr(base_c, "bias", getattr(mod_c, "bias", None)) if b is not None: b_data = b.data.clone() if b_data.shape[0] == q_dim + k_dim + v_dim: b_q = b_data[:q_dim] b_k = b_data[q_dim : q_dim + k_dim] b_v = b_data[q_dim + k_dim :] else: b_q = b_data[:split_size] b_k = b_data[split_size : 2 * split_size] b_v = b_data[2 * split_size :] if hasattr(src, "c_proj") and src.c_proj is not None: mod_o = src.c_proj base_o = getattr(mod_o, "base_layer", mod_o) out_w = getattr(base_o, "weight", getattr(mod_o, "weight", None)) if out_w is not None: out_w = out_w.data.clone() cls_o = getattr(mod_o, "__class__", type(mod_o)).__name__ base_cls_o = getattr(base_o, "__class__", type(base_o)).__name__ if ( "Conv1D" in cls_o or "Conv1D" in base_cls_o or hasattr(mod_o, "nf") or hasattr(base_o, "nf") or out_w.ndim == 2 and out_w.shape[0] < out_w.shape[1] ): out_w = out_w.t().contiguous() w_o = out_w out_b = getattr(base_o, "bias", getattr(mod_o, "bias", None)) b_o = out_b.data.clone() if out_b is not None else None # 2. DistilBERT flat projections (q_lin, k_lin, v_lin, out_lin) elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"): w_q = src.q_lin.weight.data.clone() b_q = ( src.q_lin.bias.data.clone() if getattr(src.q_lin, "bias", None) is not None else None ) w_k = src.k_lin.weight.data.clone() b_k = ( src.k_lin.bias.data.clone() if getattr(src.k_lin, "bias", None) is not None else None ) w_v = src.v_lin.weight.data.clone() b_v = ( src.v_lin.bias.data.clone() if getattr(src.v_lin, "bias", None) is not None else None ) if ( hasattr(src, "out_lin") and getattr(src, "out_lin", None) is not None and hasattr(src.out_lin, "weight") ): w_o = src.out_lin.weight.data.clone() b_o = ( src.out_lin.bias.data.clone() if getattr(src.out_lin, "bias", None) is not None else None ) # 3. Flat projection splitting (single in_proj_weight or qkv_proj / Wqkv / qkv) elif ( ( hasattr(src, "in_proj_weight") and getattr(src, "in_proj_weight", None) is not None ) or ( hasattr(src, "qkv_proj") and getattr(src, "qkv_proj", None) is not None and hasattr(src.qkv_proj, "weight") ) or ( hasattr(src, "qkv") and getattr(src, "qkv", None) is not None and hasattr(src.qkv, "weight") ) or ( hasattr(src, "Wqkv") and getattr(src, "Wqkv", None) is not None and hasattr(src.Wqkv, "weight") ) ): qkv_mod = getattr(src, "qkv_proj", getattr(src, "qkv", getattr(src, "Wqkv", None))) if qkv_mod is not None: in_w = qkv_mod.weight.data.clone() in_b = ( qkv_mod.bias.data.clone() if getattr(qkv_mod, "bias", None) is not None else None ) else: in_w = src.in_proj_weight.data.clone() in_b = getattr(src, "in_proj_bias", None) in_b = in_b.data.clone() if in_b is not None else None if in_w.ndim == 2 and in_w.shape[0] < in_w.shape[1]: in_w = in_w.t().contiguous() q_dim = self.q_head_size k_dim = self.kv_head_size v_dim = self.kv_head_size if in_w.shape[0] == q_dim + k_dim + v_dim: w_q = in_w[:q_dim, :] w_k = in_w[q_dim : q_dim + k_dim, :] w_v = in_w[q_dim + k_dim :, :] else: split_sz = in_w.shape[0] // 3 w_q = in_w[:split_sz, :] w_k = in_w[split_sz : 2 * split_sz, :] w_v = in_w[2 * split_sz :, :] if in_b is not None: if in_b.shape[0] == q_dim + k_dim + v_dim: b_q = in_b[:q_dim] b_k = in_b[q_dim : q_dim + k_dim] b_v = in_b[q_dim + k_dim :] else: split_sz = in_b.shape[0] // 3 b_q = in_b[:split_sz] b_k = in_b[split_sz : 2 * split_sz] b_v = in_b[2 * split_sz :] for out_name in ("out_proj", "o_proj", "dense", "out_lin"): if ( hasattr(src, out_name) and getattr(src, out_name, None) is not None and hasattr(getattr(src, out_name), "weight") ): m_out = getattr(src, out_name) w_o = m_out.weight.data.clone() b_o = ( m_out.bias.data.clone() if getattr(m_out, "bias", None) is not None else None ) break # 4. RoBERTa / BERT separate query, key, value elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"): w_q = src.query.weight.data.clone() b_q = ( src.query.bias.data.clone() if getattr(src.query, "bias", None) is not None else None ) w_k = src.key.weight.data.clone() b_k = ( src.key.bias.data.clone() if getattr(src.key, "bias", None) is not None else None ) w_v = src.value.weight.data.clone() b_v = ( src.value.bias.data.clone() if getattr(src.value, "bias", None) is not None else None ) if w_o is None: for attr in ("dense", "out_proj", "o_proj"): if ( hasattr(src, attr) and getattr(src, attr) is not None and hasattr(getattr(src, attr), "weight") ): m = getattr(src, attr) w_o = m.weight.data.clone() b_o = ( m.bias.data.clone() if getattr(m, "bias", None) is not None else None ) break # 5. Qwen / LLaMA (q_proj, k_proj, v_proj, o_proj) elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"): w_q = src.q_proj.weight.data.clone() b_q = ( src.q_proj.bias.data.clone() if getattr(src.q_proj, "bias", None) is not None else None ) w_k = src.k_proj.weight.data.clone() b_k = ( src.k_proj.bias.data.clone() if getattr(src.k_proj, "bias", None) is not None else None ) w_v = src.v_proj.weight.data.clone() b_v = ( src.v_proj.bias.data.clone() if getattr(src.v_proj, "bias", None) is not None else None ) if ( hasattr(src, "o_proj") and getattr(src, "o_proj", None) is not None and hasattr(src.o_proj, "weight") ): w_o = src.o_proj.weight.data.clone() b_o = ( src.o_proj.bias.data.clone() if getattr(src.o_proj, "bias", None) is not None else None ) # 6. Fallback: UniversalAdapter if w_q is None or w_k is None or w_v is None: try: from neuralops.core.adapters import UniversalAdapter adapter = UniversalAdapter(old_attn) extracted = adapter.extract_qkv_out() if extracted: w_q = extracted.get("q_weight") b_q = extracted.get("q_bias") w_k = extracted.get("k_weight") b_k = extracted.get("k_bias") w_v = extracted.get("v_weight") b_v = extracted.get("v_bias") if w_o is None: w_o = extracted.get("out_weight") b_o = extracted.get("out_bias") except Exception as ex: logger.debug("UniversalAdapter fallback failed: %s", ex) missing = [] if w_q is None: missing.append("q") if w_k is None: missing.append("k") if w_v is None: missing.append("v") if missing: logger.warning( f"Failed to extract required projections {missing} from {src_cls_name}" ) # Frobenius norm verification >= 1e-4 on copied weights for name, w in [("w_q", w_q), ("w_k", w_k), ("w_v", w_v)]: if w is not None and torch.linalg.norm(w.float()).item() < 1e-4: logger.warning(f"Copied weight {name} failed Frobenius norm check (< 1e-4)") if w_o is not None and not self.is_nested_encoder: if torch.linalg.norm(w_o.float()).item() < 1e-4: logger.warning("Copied weight w_o failed Frobenius norm check (< 1e-4)") # SVD decomposition into low-rank factors self.q_proj.decompose_from_weight(w_q, b_q) self.k_proj.decompose_from_weight(w_k, b_k) self.v_proj.decompose_from_weight(w_v, b_v) if w_o is not None: self.out_proj.decompose_from_weight(w_o, b_o) elif self.is_nested_encoder: with torch.no_grad(): U = self.out_proj.U.data V = self.out_proj.V.data U.zero_() V.zero_() d = min(U.shape) U[:d, :d].fill_diagonal_(1.0) d2 = min(V.shape) V[:d2, :d2].fill_diagonal_(1.0) if self.out_proj.bias is not None: self.out_proj.bias.data.zero_() def forward( self, hidden_states: torch.Tensor | None = None, 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, position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None, *args: Any, **kwargs: Any, ) -> tuple[torch.Tensor, torch.Tensor | None]: if encoder_hidden_states is not None: raise NotImplementedError("Cross-attention is not supported in FusedSVDKernelAttention") x = hidden_states if hidden_states is not None else kwargs.get("query", kwargs.get("x")) if x is None and len(args) > 0: x = args[0] if attention_mask is None: attention_mask = kwargs.get("mask", kwargs.get("attn_mask")) B, L, _ = x.shape q = self.q_proj(x).view(B, L, self.num_heads, self.head_dim).transpose(1, 2).contiguous() k = self.k_proj(x).view(B, L, self.num_kv_heads, self.head_dim).transpose(1, 2).contiguous() v = self.v_proj(x).view(B, L, self.num_kv_heads, self.head_dim).transpose(1, 2).contiguous() if position_embeddings is not None: cos, sin = position_embeddings q, k = apply_rotary_pos_emb(q, k, cos, sin) if self.kv_repeat > 1: k = repeat_kv(k, self.kv_repeat) v = repeat_kv(v, self.kv_repeat) scale = 1.0 / math.sqrt(self.head_dim) # Memory-efficient: use SDPA to avoid OOM on large seqs try: _qs = q.half() if q.dtype == torch.float32 else q _ks = k.half() if k.dtype == torch.float32 else k _vs = v.half() if v.dtype == torch.float32 else v _is_causal = ( getattr(self, "is_decoder", False) or getattr(self, "is_causal", False) ) and L > 1 _sdpa_out = torch.nn.functional.scaled_dot_product_attention( _qs, _ks, _vs, is_causal=_is_causal, dropout_p=self.dropout.p if self.training else 0.0, ) _sdpa_out = ( _sdpa_out.to(x.dtype).transpose(1, 2).contiguous().view(B, L, self.all_head_size) ) _sdpa_out = torch.nan_to_num(_sdpa_out, nan=0.0, posinf=1e4, neginf=-1e4) return ( self.out_proj(_sdpa_out) if not getattr(self, "is_nested_encoder", False) else _sdpa_out ), None except Exception: pass scores = torch.matmul(q.float(), k.float().transpose(-2, -1).contiguous()) * scale if self.is_decoder or self.is_causal: causal_mask = torch.tril(torch.ones(L, L, device=x.device, dtype=torch.bool)) scores = scores.masked_fill(~causal_mask.unsqueeze(0).unsqueeze(0), -1e4) if attention_mask is not None: if attention_mask.dim() == 2: valid_mask = get_token_valid_mask(attention_mask, B, L, x.device) if valid_mask is None: valid_mask = torch.ones(B, L, device=x.device, dtype=torch.bool) scores = scores.masked_fill(~valid_mask.unsqueeze(1).unsqueeze(2), -1e4) elif attention_mask.dim() == 4: scores = scores + attention_mask.clamp(min=-1e4) probs = F.softmax(scores, dim=-1, dtype=torch.float32).to(v.dtype) probs = self.dropout(probs) if head_mask is not None: probs = probs * head_mask context = torch.matmul(probs, v).transpose(1, 2).contiguous().view(B, L, self.q_head_size) if self.is_nested_encoder: out = context else: out = self.out_proj(context) return out, (probs if output_attentions else None)Fused SVD-Kernel Linear Attention Module. Production-hardened across DistilBERT, RoBERTa/BERT nested, GPT-2 Conv1D, and Qwen/LLaMA.
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.
var supports_bidirectional : bool-
The type of the None singleton.
var supports_causal : bool-
The type of the None singleton.
var supports_gqa : bool-
The type of the None singleton.
Instance variables
prop key : FusedSVDLinear-
Expand source code
@property def key(self) -> FusedSVDLinear: return self.k_proj prop query : FusedSVDLinear-
Expand source code
@property def query(self) -> FusedSVDLinear: return self.q_proj prop value : FusedSVDLinear-
Expand source code
@property def value(self) -> FusedSVDLinear: return self.v_proj
Methods
def copy_weights(self, old_attn: Any) ‑> None-
Expand source code
def copy_weights(self, old_attn: Any) -> None: """Multi-architecture weight copying supporting GPT-2 Conv1D, RoBERTa nested, and flat projections.""" if old_attn is None: return with torch.no_grad(): w_q, b_q = None, None w_k, b_k = None, None w_v, b_v = None, None w_o, b_o = None, None src_cls_name = getattr(old_attn, "__class__", type(old_attn)).__name__ has_out = ( hasattr(old_attn, "out_lin") or hasattr(old_attn, "dense") or hasattr(old_attn, "out_proj") or hasattr(old_attn, "c_proj") or hasattr(old_attn, "o_proj") ) if ( hasattr(old_attn, "out_lin") or hasattr(old_attn, "c_proj") or hasattr(old_attn, "o_proj") ): self.is_nested_encoder = False elif ( ( hasattr(old_attn, "self") and hasattr(old_attn.self, "query") and hasattr(old_attn, "output") and hasattr(old_attn.output, "dense") ) or src_cls_name in ("BertSelfAttention", "RobertaSelfAttention") or (hasattr(old_attn, "query") and not has_out) ): self.is_nested_encoder = True else: self.is_nested_encoder = False if hasattr(old_attn, "self") and hasattr(old_attn, "output"): self.is_nested_encoder = True if ( hasattr(old_attn, "output") and hasattr(old_attn.output, "dense") and hasattr(old_attn.output.dense, "weight") ): w_o = old_attn.output.dense.weight.data b_o = getattr(old_attn.output.dense, "bias", None) b_o = b_o.data if b_o is not None else None src = old_attn while hasattr(src, "self") and not any( hasattr(src, attr) for attr in ("query", "key", "value", "q_proj", "k_proj", "v_proj") ): self.is_nested_encoder = True src = src.self # 1. GPT-2 Conv1D fused c_attn & c_proj if hasattr(src, "c_attn") and src.c_attn is not None: self.is_decoder = True self.is_causal = True mod_c = src.c_attn base_c = getattr(mod_c, "base_layer", mod_c) w = getattr(base_c, "weight", getattr(mod_c, "weight", None)) if w is not None: w = w.data.clone() cls_c = getattr(mod_c, "__class__", type(mod_c)).__name__ base_cls_c = getattr(base_c, "__class__", type(base_c)).__name__ if ( "Conv1D" in cls_c or "Conv1D" in base_cls_c or hasattr(mod_c, "nf") or hasattr(base_c, "nf") or (w.ndim == 2 and w.shape[0] < w.shape[1]) ): w = w.t().contiguous() q_dim = self.q_head_size k_dim = self.kv_head_size v_dim = self.kv_head_size if w.shape[0] == q_dim + k_dim + v_dim: w_q = w[:q_dim, :] w_k = w[q_dim : q_dim + k_dim, :] w_v = w[q_dim + k_dim :, :] else: split_size = w.shape[0] // 3 w_q = w[:split_size, :] w_k = w[split_size : 2 * split_size, :] w_v = w[2 * split_size :, :] b = getattr(base_c, "bias", getattr(mod_c, "bias", None)) if b is not None: b_data = b.data.clone() if b_data.shape[0] == q_dim + k_dim + v_dim: b_q = b_data[:q_dim] b_k = b_data[q_dim : q_dim + k_dim] b_v = b_data[q_dim + k_dim :] else: b_q = b_data[:split_size] b_k = b_data[split_size : 2 * split_size] b_v = b_data[2 * split_size :] if hasattr(src, "c_proj") and src.c_proj is not None: mod_o = src.c_proj base_o = getattr(mod_o, "base_layer", mod_o) out_w = getattr(base_o, "weight", getattr(mod_o, "weight", None)) if out_w is not None: out_w = out_w.data.clone() cls_o = getattr(mod_o, "__class__", type(mod_o)).__name__ base_cls_o = getattr(base_o, "__class__", type(base_o)).__name__ if ( "Conv1D" in cls_o or "Conv1D" in base_cls_o or hasattr(mod_o, "nf") or hasattr(base_o, "nf") or out_w.ndim == 2 and out_w.shape[0] < out_w.shape[1] ): out_w = out_w.t().contiguous() w_o = out_w out_b = getattr(base_o, "bias", getattr(mod_o, "bias", None)) b_o = out_b.data.clone() if out_b is not None else None # 2. DistilBERT flat projections (q_lin, k_lin, v_lin, out_lin) elif hasattr(src, "q_lin") and hasattr(src, "k_lin") and hasattr(src, "v_lin"): w_q = src.q_lin.weight.data.clone() b_q = ( src.q_lin.bias.data.clone() if getattr(src.q_lin, "bias", None) is not None else None ) w_k = src.k_lin.weight.data.clone() b_k = ( src.k_lin.bias.data.clone() if getattr(src.k_lin, "bias", None) is not None else None ) w_v = src.v_lin.weight.data.clone() b_v = ( src.v_lin.bias.data.clone() if getattr(src.v_lin, "bias", None) is not None else None ) if ( hasattr(src, "out_lin") and getattr(src, "out_lin", None) is not None and hasattr(src.out_lin, "weight") ): w_o = src.out_lin.weight.data.clone() b_o = ( src.out_lin.bias.data.clone() if getattr(src.out_lin, "bias", None) is not None else None ) # 3. Flat projection splitting (single in_proj_weight or qkv_proj / Wqkv / qkv) elif ( ( hasattr(src, "in_proj_weight") and getattr(src, "in_proj_weight", None) is not None ) or ( hasattr(src, "qkv_proj") and getattr(src, "qkv_proj", None) is not None and hasattr(src.qkv_proj, "weight") ) or ( hasattr(src, "qkv") and getattr(src, "qkv", None) is not None and hasattr(src.qkv, "weight") ) or ( hasattr(src, "Wqkv") and getattr(src, "Wqkv", None) is not None and hasattr(src.Wqkv, "weight") ) ): qkv_mod = getattr(src, "qkv_proj", getattr(src, "qkv", getattr(src, "Wqkv", None))) if qkv_mod is not None: in_w = qkv_mod.weight.data.clone() in_b = ( qkv_mod.bias.data.clone() if getattr(qkv_mod, "bias", None) is not None else None ) else: in_w = src.in_proj_weight.data.clone() in_b = getattr(src, "in_proj_bias", None) in_b = in_b.data.clone() if in_b is not None else None if in_w.ndim == 2 and in_w.shape[0] < in_w.shape[1]: in_w = in_w.t().contiguous() q_dim = self.q_head_size k_dim = self.kv_head_size v_dim = self.kv_head_size if in_w.shape[0] == q_dim + k_dim + v_dim: w_q = in_w[:q_dim, :] w_k = in_w[q_dim : q_dim + k_dim, :] w_v = in_w[q_dim + k_dim :, :] else: split_sz = in_w.shape[0] // 3 w_q = in_w[:split_sz, :] w_k = in_w[split_sz : 2 * split_sz, :] w_v = in_w[2 * split_sz :, :] if in_b is not None: if in_b.shape[0] == q_dim + k_dim + v_dim: b_q = in_b[:q_dim] b_k = in_b[q_dim : q_dim + k_dim] b_v = in_b[q_dim + k_dim :] else: split_sz = in_b.shape[0] // 3 b_q = in_b[:split_sz] b_k = in_b[split_sz : 2 * split_sz] b_v = in_b[2 * split_sz :] for out_name in ("out_proj", "o_proj", "dense", "out_lin"): if ( hasattr(src, out_name) and getattr(src, out_name, None) is not None and hasattr(getattr(src, out_name), "weight") ): m_out = getattr(src, out_name) w_o = m_out.weight.data.clone() b_o = ( m_out.bias.data.clone() if getattr(m_out, "bias", None) is not None else None ) break # 4. RoBERTa / BERT separate query, key, value elif hasattr(src, "query") and hasattr(src, "key") and hasattr(src, "value"): w_q = src.query.weight.data.clone() b_q = ( src.query.bias.data.clone() if getattr(src.query, "bias", None) is not None else None ) w_k = src.key.weight.data.clone() b_k = ( src.key.bias.data.clone() if getattr(src.key, "bias", None) is not None else None ) w_v = src.value.weight.data.clone() b_v = ( src.value.bias.data.clone() if getattr(src.value, "bias", None) is not None else None ) if w_o is None: for attr in ("dense", "out_proj", "o_proj"): if ( hasattr(src, attr) and getattr(src, attr) is not None and hasattr(getattr(src, attr), "weight") ): m = getattr(src, attr) w_o = m.weight.data.clone() b_o = ( m.bias.data.clone() if getattr(m, "bias", None) is not None else None ) break # 5. Qwen / LLaMA (q_proj, k_proj, v_proj, o_proj) elif hasattr(src, "q_proj") and hasattr(src, "k_proj") and hasattr(src, "v_proj"): w_q = src.q_proj.weight.data.clone() b_q = ( src.q_proj.bias.data.clone() if getattr(src.q_proj, "bias", None) is not None else None ) w_k = src.k_proj.weight.data.clone() b_k = ( src.k_proj.bias.data.clone() if getattr(src.k_proj, "bias", None) is not None else None ) w_v = src.v_proj.weight.data.clone() b_v = ( src.v_proj.bias.data.clone() if getattr(src.v_proj, "bias", None) is not None else None ) if ( hasattr(src, "o_proj") and getattr(src, "o_proj", None) is not None and hasattr(src.o_proj, "weight") ): w_o = src.o_proj.weight.data.clone() b_o = ( src.o_proj.bias.data.clone() if getattr(src.o_proj, "bias", None) is not None else None ) # 6. Fallback: UniversalAdapter if w_q is None or w_k is None or w_v is None: try: from neuralops.core.adapters import UniversalAdapter adapter = UniversalAdapter(old_attn) extracted = adapter.extract_qkv_out() if extracted: w_q = extracted.get("q_weight") b_q = extracted.get("q_bias") w_k = extracted.get("k_weight") b_k = extracted.get("k_bias") w_v = extracted.get("v_weight") b_v = extracted.get("v_bias") if w_o is None: w_o = extracted.get("out_weight") b_o = extracted.get("out_bias") except Exception as ex: logger.debug("UniversalAdapter fallback failed: %s", ex) missing = [] if w_q is None: missing.append("q") if w_k is None: missing.append("k") if w_v is None: missing.append("v") if missing: logger.warning( f"Failed to extract required projections {missing} from {src_cls_name}" ) # Frobenius norm verification >= 1e-4 on copied weights for name, w in [("w_q", w_q), ("w_k", w_k), ("w_v", w_v)]: if w is not None and torch.linalg.norm(w.float()).item() < 1e-4: logger.warning(f"Copied weight {name} failed Frobenius norm check (< 1e-4)") if w_o is not None and not self.is_nested_encoder: if torch.linalg.norm(w_o.float()).item() < 1e-4: logger.warning("Copied weight w_o failed Frobenius norm check (< 1e-4)") # SVD decomposition into low-rank factors self.q_proj.decompose_from_weight(w_q, b_q) self.k_proj.decompose_from_weight(w_k, b_k) self.v_proj.decompose_from_weight(w_v, b_v) if w_o is not None: self.out_proj.decompose_from_weight(w_o, b_o) elif self.is_nested_encoder: with torch.no_grad(): U = self.out_proj.U.data V = self.out_proj.V.data U.zero_() V.zero_() d = min(U.shape) U[:d, :d].fill_diagonal_(1.0) d2 = min(V.shape) V[:d2, :d2].fill_diagonal_(1.0) if self.out_proj.bias is not None: self.out_proj.bias.data.zero_()Multi-architecture weight copying supporting GPT-2 Conv1D, RoBERTa nested, and flat projections.
Inherited members
class FusedSVDLinear (in_features: int,
out_features: int,
rank: int = 32,
bias: bool = True,
use_rms_norm: bool = False,
eps: float = 1e-06,
dtype: torch.dtype = torch.float32,
device: torch.device | None = None)-
Expand source code
class FusedSVDLinear(nn.Module): """ RESEARCH SOTA: Fused SVD-Linear Module. Combines low-rank projection $Z = x V^T$, in-SRAM RMSNorm, and up-projection $Y = Z U^T + b$ into a single kernel invocation, eliminating HBM roundtrips. """ def __init__( self, in_features: int, out_features: int, rank: int = 32, bias: bool = True, use_rms_norm: bool = False, eps: float = 1e-6, dtype: torch.dtype = torch.float32, device: torch.device | None = None, ) -> None: super().__init__() self.in_features = in_features self.out_features = out_features self.rank = rank self.eps = eps self.use_rms_norm = use_rms_norm # Low-rank factorized parameters self.V = nn.Parameter(torch.empty(rank, in_features, dtype=dtype, device=device)) self.U = nn.Parameter(torch.empty(out_features, rank, dtype=dtype, device=device)) if use_rms_norm: self.norm_weight = nn.Parameter(torch.ones(rank, dtype=dtype, device=device)) else: self.register_parameter("norm_weight", None) if bias: self.bias = nn.Parameter(torch.empty(out_features, dtype=dtype, device=device)) else: self.register_parameter("bias", None) self.reset_parameters() def reset_parameters(self) -> None: # Symmetrical square-root initialization nn.init.orthogonal_(self.U) nn.init.orthogonal_(self.V) if self.bias is not None: nn.init.zeros_(self.bias) if self.norm_weight is not None: nn.init.ones_(self.norm_weight) @property def weight(self) -> torch.Tensor: """Reconstructed effective weight matrix for compatibility.""" return torch.matmul(self.U, self.V) def forward(self, x: torch.Tensor) -> torch.Tensor: # RESEARCH FIX: Dispatches fused kernel executing down-projection, RMSNorm, and up-projection return fused_svd_linear( x=x, V=self.V, U=self.U, norm_weight=self.norm_weight, bias=self.bias, eps=self.eps, use_rms_norm=self.use_rms_norm, ) def set_low_rank_factors(self, U_factor: torch.Tensor, V_factor: torch.Tensor) -> None: """Set precomputed U and V low-rank factors.""" with torch.no_grad(): self.U.copy_(U_factor) self.V.copy_(V_factor) def decompose_from_weight(self, weight: torch.Tensor, bias: torch.Tensor | None = None) -> None: """Decompose dense weight matrix into symmetric SVD factors.""" with torch.no_grad(): W = weight.float() U, S, Vt = torch.linalg.svd(W, full_matrices=False) r = min(self.rank, len(S)) sqrt_s = torch.sqrt(S[:r]) self.U.data.zero_() self.V.data.zero_() self.U.data[:, :r] = (U[:, :r] * sqrt_s).to(self.U.dtype) self.V.data[:r, :] = (sqrt_s.unsqueeze(1) * Vt[:r, :]).to(self.V.dtype) if bias is not None and self.bias is not None: self.bias.data.copy_(bias) def get_orthogonality_loss(self) -> torch.Tensor: """ # RESEARCH FIX: Computes Stiefel manifold penalty to prevent singular subspace collapse. """ u_gram = torch.matmul(self.U.T, self.U) v_gram = torch.matmul(self.V, self.V.T) eye_r = torch.eye(self.rank, device=self.U.device, dtype=self.U.dtype) loss_u = torch.norm(u_gram - eye_r, p="fro") ** 2 loss_v = torch.norm(v_gram - eye_r, p="fro") ** 2 return (loss_u + loss_v) / (self.rank * self.rank) def extra_repr(self) -> str: return ( f"in_features={self.in_features}, out_features={self.out_features}, " f"rank={self.rank}, rms_norm={self.use_rms_norm}, bias={self.bias is not None}" )RESEARCH SOTA: Fused SVD-Linear Module. Combines low-rank projection $Z = x V^T$, in-SRAM RMSNorm, and up-projection $Y = Z U^T + b$ into a single kernel invocation, eliminating HBM roundtrips.
Initialize internal Module state, shared by both nn.Module and ScriptModule.
Ancestors
- torch.nn.modules.module.Module
Instance variables
prop weight : torch.Tensor-
Expand source code
@property def weight(self) -> torch.Tensor: """Reconstructed effective weight matrix for compatibility.""" return torch.matmul(self.U, self.V)Reconstructed effective weight matrix for compatibility.
Methods
def decompose_from_weight(self, weight: torch.Tensor, bias: torch.Tensor | None = None) ‑> None-
Expand source code
def decompose_from_weight(self, weight: torch.Tensor, bias: torch.Tensor | None = None) -> None: """Decompose dense weight matrix into symmetric SVD factors.""" with torch.no_grad(): W = weight.float() U, S, Vt = torch.linalg.svd(W, full_matrices=False) r = min(self.rank, len(S)) sqrt_s = torch.sqrt(S[:r]) self.U.data.zero_() self.V.data.zero_() self.U.data[:, :r] = (U[:, :r] * sqrt_s).to(self.U.dtype) self.V.data[:r, :] = (sqrt_s.unsqueeze(1) * Vt[:r, :]).to(self.V.dtype) if bias is not None and self.bias is not None: self.bias.data.copy_(bias)Decompose dense weight matrix into symmetric SVD factors.
def extra_repr(self) ‑> str-
Expand source code
def extra_repr(self) -> str: return ( f"in_features={self.in_features}, out_features={self.out_features}, " f"rank={self.rank}, rms_norm={self.use_rms_norm}, bias={self.bias is not None}" )Return the extra representation of the module.
To print customized extra information, you should re-implement this method in your own modules. Both single-line and multi-line strings are acceptable.
def forward(self, x: torch.Tensor) ‑> torch.Tensor-
Expand source code
def forward(self, x: torch.Tensor) -> torch.Tensor: # RESEARCH FIX: Dispatches fused kernel executing down-projection, RMSNorm, and up-projection return fused_svd_linear( x=x, V=self.V, U=self.U, norm_weight=self.norm_weight, bias=self.bias, eps=self.eps, use_rms_norm=self.use_rms_norm, )Define 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 get_orthogonality_loss(self) ‑> torch.Tensor-
Expand source code
def get_orthogonality_loss(self) -> torch.Tensor: """ # RESEARCH FIX: Computes Stiefel manifold penalty to prevent singular subspace collapse. """ u_gram = torch.matmul(self.U.T, self.U) v_gram = torch.matmul(self.V, self.V.T) eye_r = torch.eye(self.rank, device=self.U.device, dtype=self.U.dtype) loss_u = torch.norm(u_gram - eye_r, p="fro") ** 2 loss_v = torch.norm(v_gram - eye_r, p="fro") ** 2 return (loss_u + loss_v) / (self.rank * self.rank)RESEARCH FIX: Computes Stiefel manifold penalty to prevent singular subspace collapse.
def reset_parameters(self) ‑> None-
Expand source code
def reset_parameters(self) -> None: # Symmetrical square-root initialization nn.init.orthogonal_(self.U) nn.init.orthogonal_(self.V) if self.bias is not None: nn.init.zeros_(self.bias) if self.norm_weight is not None: nn.init.ones_(self.norm_weight) def set_low_rank_factors(self, U_factor: torch.Tensor, V_factor: torch.Tensor) ‑> None-
Expand source code
def set_low_rank_factors(self, U_factor: torch.Tensor, V_factor: torch.Tensor) -> None: """Set precomputed U and V low-rank factors.""" with torch.no_grad(): self.U.copy_(U_factor) self.V.copy_(V_factor)Set precomputed U and V low-rank factors.
class _FusedSVDLinearAutogradFunction (*args, **kwargs)-
Expand source code
class _FusedSVDLinearAutogradFunction(torch.autograd.Function): is_custom_attention: bool = True """ Autograd Function dispatching Triton JIT fused kernel on CUDA GPUs with exact analytical backward gradient computation. """ @staticmethod def forward( ctx: Any, x: torch.Tensor, V: torch.Tensor, U: torch.Tensor, norm_weight: torch.Tensor | None = None, bias: torch.Tensor | None = None, eps: float = 1e-6, use_rms_norm: bool = False, ) -> torch.Tensor: orig_shape = x.shape x_2d = x.reshape(-1, orig_shape[-1]).contiguous() M, K_in = x_2d.shape R, _ = V.shape K_out, _ = U.shape # Use Triton kernel on CUDA when available and rank is power-of-2 up to 128 use_triton = ( HAS_TRITON and x.is_cuda and V.is_cuda and U.is_cuda and (R & (R - 1) == 0) and (16 <= R <= 128) and x_2d.dtype in [torch.float16, torch.bfloat16, torch.float32] ) if use_triton: try: y_2d = torch.empty((M, K_out), device=x.device, dtype=x.dtype) BLOCK_M = 32 BLOCK_K_IN = ( 64 if K_in >= 64 and (K_in % 64 == 0 or K_in > 64) else min(64, max(16, K_in)) ) BLOCK_K_IN = 1 << (BLOCK_K_IN - 1).bit_length() BLOCK_K_OUT = ( 64 if K_out >= 64 and (K_out % 64 == 0 or K_out > 64) else min(64, max(16, K_out)) ) BLOCK_K_OUT = 1 << (BLOCK_K_OUT - 1).bit_length() grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(K_out, BLOCK_K_OUT)) _fused_svd_linear_fwd_kernel[grid]( x_2d, V, U, norm_weight if (norm_weight is not None and use_rms_norm) else x_2d, bias if bias is not None else x_2d, y_2d, x_2d.stride(0), x_2d.stride(1), V.stride(0), V.stride(1), U.stride(0), U.stride(1), y_2d.stride(0), y_2d.stride(1), M, K_in=K_in, R=R, K_out=K_out, eps=eps, USE_RMS_NORM=use_rms_norm, HAS_BIAS=bias is not None, HAS_NORM_WEIGHT=norm_weight is not None and use_rms_norm, BLOCK_M=BLOCK_M, BLOCK_K_IN=BLOCK_K_IN, BLOCK_K_OUT=BLOCK_K_OUT, ) # RESEARCH FIX: Exact analytical intermediate reconstruction for backward caching with torch.no_grad(): z = F.linear(x_2d, V.to(x_2d.dtype)) if use_rms_norm: if z.dtype in (torch.float16, torch.bfloat16): rms = torch.rsqrt( torch.mean(z.float() ** 2, dim=-1, keepdim=True) + eps ).to(z.dtype) else: rms = torch.rsqrt(torch.mean(z**2, dim=-1, keepdim=True) + eps) z_hat = z * rms z_norm = z_hat * norm_weight if norm_weight is not None else z_hat else: rms = torch.empty(0, device=x.device, dtype=x.dtype) z_hat = z z_norm = z ctx.save_for_backward(x_2d, V, U, norm_weight, bias, z, rms, z_hat, z_norm) ctx.eps = eps ctx.use_rms_norm = use_rms_norm return y_2d.view(*orig_shape[:-1], K_out) except Exception: pass # Fall back to PyTorch native reference below # PyTorch Native Reference Implementation (Fallback / CPU / MPS) # RESEARCH FIX: Numerically identical reference implementation for CPU/MPS z = F.linear(x_2d, V.to(x_2d.dtype)) # [M, R] if use_rms_norm: if z.dtype in (torch.float16, torch.bfloat16): rms = torch.rsqrt(torch.mean(z.float() ** 2, dim=-1, keepdim=True) + eps).to( z.dtype ) else: rms = torch.rsqrt(torch.mean(z**2, dim=-1, keepdim=True) + eps) z_hat = z * rms nw = norm_weight.to(z.dtype) if norm_weight is not None else None z_norm = z_hat * nw if nw is not None else z_hat else: rms = torch.empty(0, device=x.device, dtype=x.dtype) z_hat = z z_norm = z b = bias.to(z_norm.dtype) if bias is not None else None y_2d = F.linear(z_norm, U.to(z_norm.dtype), b) ctx.save_for_backward(x_2d, V, U, norm_weight, bias, z, rms, z_hat, z_norm) ctx.eps = eps ctx.use_rms_norm = use_rms_norm return y_2d.view(*orig_shape[:-1], K_out) @staticmethod def backward(ctx: Any, grad_output: torch.Tensor) -> tuple[torch.Tensor | None, ...]: x_2d, V, U, norm_weight, bias, z, rms, z_hat, z_norm = ctx.saved_tensors use_rms_norm = ctx.use_rms_norm grad_output_2d = grad_output.reshape(-1, U.shape[0]).contiguous() dtype = grad_output_2d.dtype device = grad_output_2d.device z_norm_c = z_norm.to(dtype=dtype, device=device) U_c = U.to(dtype=dtype, device=device) V_c = V.to(dtype=dtype, device=device) x_2d_c = x_2d.to(dtype=dtype, device=device) z_hat_c = z_hat.to(dtype=dtype, device=device) rms_c = rms.to(dtype=dtype, device=device) # 1. Gradient with respect to bias and U grad_bias = grad_output_2d.sum(dim=0).to(bias.dtype) if bias is not None else None grad_U = torch.matmul(grad_output_2d.T, z_norm_c).to(U.dtype) # [K_out, R] # 2. Gradient backpropagated to z_norm grad_z_norm = torch.matmul(grad_output_2d, U_c) # [M, R] if use_rms_norm: if norm_weight is not None: nw_c = norm_weight.to(dtype=dtype, device=device) grad_norm_weight = (grad_z_norm * z_hat_c).sum(dim=0).to(norm_weight.dtype) grad_z_hat = grad_z_norm * nw_c else: grad_norm_weight = None grad_z_hat = grad_z_norm grad_z = rms_c * ( grad_z_hat - z_hat_c * torch.mean(grad_z_hat * z_hat_c, dim=-1, keepdim=True) ) else: grad_norm_weight = None grad_z = grad_z_norm # 5. Gradient with respect to V and input x grad_V = torch.matmul(grad_z.T, x_2d_c).to(V.dtype) # [R, K_in] grad_x = ( torch.matmul(grad_z, V_c).view(*grad_output.shape[:-1], x_2d.shape[-1]).to(x_2d.dtype) ) return grad_x, grad_V, grad_U, grad_norm_weight, grad_bias, None, NoneBase class to create custom
autograd.Function.To create a custom
autograd.Function, subclass this class and implement the :meth:forwardand :meth:backwardstatic methods. Then, to use your custom op in the forward pass, call the class methodapply. Do not call :meth:forwarddirectly.To ensure correctness and best performance, make sure you are calling the correct methods on
ctxand validating your backward function using :func:torch.autograd.gradcheck.See :ref:
extending-autogradfor more details on how to use this class.Examples::
>>> # xdoctest: +REQUIRES(env:TORCH_DOCTEST_AUTOGRAD) >>> class Exp(Function): >>> @staticmethod >>> def forward(ctx, i): >>> result = i.exp() >>> ctx.save_for_backward(result) >>> return result >>> >>> @staticmethod >>> def backward(ctx, grad_output): >>> result, = ctx.saved_tensors >>> return grad_output * result >>> >>> # Use it by calling the apply method: >>> # xdoctest: +SKIP >>> output = Exp.apply(input)Ancestors
- torch.autograd.function.Function
- torch.autograd.function._SingleLevelFunction
- torch._C._FunctionBase
- torch.autograd.function.FunctionCtx
- torch.autograd.function._HookMixin
Class variables
var is_custom_attention : bool-
Autograd Function dispatching Triton JIT fused kernel on CUDA GPUs with exact analytical backward gradient computation.
Static methods
def backward(ctx: Any, grad_output: torch.Tensor) ‑> tuple[torch.Tensor | None, ...]-
Expand source code
@staticmethod def backward(ctx: Any, grad_output: torch.Tensor) -> tuple[torch.Tensor | None, ...]: x_2d, V, U, norm_weight, bias, z, rms, z_hat, z_norm = ctx.saved_tensors use_rms_norm = ctx.use_rms_norm grad_output_2d = grad_output.reshape(-1, U.shape[0]).contiguous() dtype = grad_output_2d.dtype device = grad_output_2d.device z_norm_c = z_norm.to(dtype=dtype, device=device) U_c = U.to(dtype=dtype, device=device) V_c = V.to(dtype=dtype, device=device) x_2d_c = x_2d.to(dtype=dtype, device=device) z_hat_c = z_hat.to(dtype=dtype, device=device) rms_c = rms.to(dtype=dtype, device=device) # 1. Gradient with respect to bias and U grad_bias = grad_output_2d.sum(dim=0).to(bias.dtype) if bias is not None else None grad_U = torch.matmul(grad_output_2d.T, z_norm_c).to(U.dtype) # [K_out, R] # 2. Gradient backpropagated to z_norm grad_z_norm = torch.matmul(grad_output_2d, U_c) # [M, R] if use_rms_norm: if norm_weight is not None: nw_c = norm_weight.to(dtype=dtype, device=device) grad_norm_weight = (grad_z_norm * z_hat_c).sum(dim=0).to(norm_weight.dtype) grad_z_hat = grad_z_norm * nw_c else: grad_norm_weight = None grad_z_hat = grad_z_norm grad_z = rms_c * ( grad_z_hat - z_hat_c * torch.mean(grad_z_hat * z_hat_c, dim=-1, keepdim=True) ) else: grad_norm_weight = None grad_z = grad_z_norm # 5. Gradient with respect to V and input x grad_V = torch.matmul(grad_z.T, x_2d_c).to(V.dtype) # [R, K_in] grad_x = ( torch.matmul(grad_z, V_c).view(*grad_output.shape[:-1], x_2d.shape[-1]).to(x_2d.dtype) ) return grad_x, grad_V, grad_U, grad_norm_weight, grad_bias, None, NoneDefine a formula for differentiating the operation with backward mode automatic differentiation.
This function is to be overridden by all subclasses. (Defining this function is equivalent to defining the
vjpfunction.)It must accept a context :attr:
ctxas the first argument, followed by as many outputs as the :func:forwardreturned (None will be passed in for non tensor outputs of the forward function), and it should return as many tensors, as there were inputs to :func:forward. Each argument is the gradient w.r.t the given output, and each returned value should be the gradient w.r.t. the corresponding input. If an input is not a Tensor or is a Tensor not requiring grads, you can just pass None as a gradient for that input.The context can be used to retrieve tensors saved during the forward pass. It also has an attribute :attr:
ctx.needs_input_gradas a tuple of booleans representing whether each input needs gradient. E.g., :func:backwardwill havectx.needs_input_grad[0] = Trueif the first input to :func:forwardneeds gradient computed w.r.t. the output. def forward(ctx: Any,
x: torch.Tensor,
V: torch.Tensor,
U: torch.Tensor,
norm_weight: torch.Tensor | None = None,
bias: torch.Tensor | None = None,
eps: float = 1e-06,
use_rms_norm: bool = False) ‑> torch.Tensor-
Expand source code
@staticmethod def forward( ctx: Any, x: torch.Tensor, V: torch.Tensor, U: torch.Tensor, norm_weight: torch.Tensor | None = None, bias: torch.Tensor | None = None, eps: float = 1e-6, use_rms_norm: bool = False, ) -> torch.Tensor: orig_shape = x.shape x_2d = x.reshape(-1, orig_shape[-1]).contiguous() M, K_in = x_2d.shape R, _ = V.shape K_out, _ = U.shape # Use Triton kernel on CUDA when available and rank is power-of-2 up to 128 use_triton = ( HAS_TRITON and x.is_cuda and V.is_cuda and U.is_cuda and (R & (R - 1) == 0) and (16 <= R <= 128) and x_2d.dtype in [torch.float16, torch.bfloat16, torch.float32] ) if use_triton: try: y_2d = torch.empty((M, K_out), device=x.device, dtype=x.dtype) BLOCK_M = 32 BLOCK_K_IN = ( 64 if K_in >= 64 and (K_in % 64 == 0 or K_in > 64) else min(64, max(16, K_in)) ) BLOCK_K_IN = 1 << (BLOCK_K_IN - 1).bit_length() BLOCK_K_OUT = ( 64 if K_out >= 64 and (K_out % 64 == 0 or K_out > 64) else min(64, max(16, K_out)) ) BLOCK_K_OUT = 1 << (BLOCK_K_OUT - 1).bit_length() grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(K_out, BLOCK_K_OUT)) _fused_svd_linear_fwd_kernel[grid]( x_2d, V, U, norm_weight if (norm_weight is not None and use_rms_norm) else x_2d, bias if bias is not None else x_2d, y_2d, x_2d.stride(0), x_2d.stride(1), V.stride(0), V.stride(1), U.stride(0), U.stride(1), y_2d.stride(0), y_2d.stride(1), M, K_in=K_in, R=R, K_out=K_out, eps=eps, USE_RMS_NORM=use_rms_norm, HAS_BIAS=bias is not None, HAS_NORM_WEIGHT=norm_weight is not None and use_rms_norm, BLOCK_M=BLOCK_M, BLOCK_K_IN=BLOCK_K_IN, BLOCK_K_OUT=BLOCK_K_OUT, ) # RESEARCH FIX: Exact analytical intermediate reconstruction for backward caching with torch.no_grad(): z = F.linear(x_2d, V.to(x_2d.dtype)) if use_rms_norm: if z.dtype in (torch.float16, torch.bfloat16): rms = torch.rsqrt( torch.mean(z.float() ** 2, dim=-1, keepdim=True) + eps ).to(z.dtype) else: rms = torch.rsqrt(torch.mean(z**2, dim=-1, keepdim=True) + eps) z_hat = z * rms z_norm = z_hat * norm_weight if norm_weight is not None else z_hat else: rms = torch.empty(0, device=x.device, dtype=x.dtype) z_hat = z z_norm = z ctx.save_for_backward(x_2d, V, U, norm_weight, bias, z, rms, z_hat, z_norm) ctx.eps = eps ctx.use_rms_norm = use_rms_norm return y_2d.view(*orig_shape[:-1], K_out) except Exception: pass # Fall back to PyTorch native reference below # PyTorch Native Reference Implementation (Fallback / CPU / MPS) # RESEARCH FIX: Numerically identical reference implementation for CPU/MPS z = F.linear(x_2d, V.to(x_2d.dtype)) # [M, R] if use_rms_norm: if z.dtype in (torch.float16, torch.bfloat16): rms = torch.rsqrt(torch.mean(z.float() ** 2, dim=-1, keepdim=True) + eps).to( z.dtype ) else: rms = torch.rsqrt(torch.mean(z**2, dim=-1, keepdim=True) + eps) z_hat = z * rms nw = norm_weight.to(z.dtype) if norm_weight is not None else None z_norm = z_hat * nw if nw is not None else z_hat else: rms = torch.empty(0, device=x.device, dtype=x.dtype) z_hat = z z_norm = z b = bias.to(z_norm.dtype) if bias is not None else None y_2d = F.linear(z_norm, U.to(z_norm.dtype), b) ctx.save_for_backward(x_2d, V, U, norm_weight, bias, z, rms, z_hat, z_norm) ctx.eps = eps ctx.use_rms_norm = use_rms_norm return y_2d.view(*orig_shape[:-1], K_out)Define the forward of the custom autograd Function.
This function is to be overridden by all subclasses. There are two ways to define forward:
Usage 1 (Combined forward and ctx)::
@staticmethod def forward(ctx: Any, *args: Any, **kwargs: Any) -> Any: pass- It must accept a context ctx as the first argument, followed by any number of arguments (tensors or other types).
- See :ref:
combining-forward-contextfor more details
Usage 2 (Separate forward and ctx)::
@staticmethod def forward(*args: Any, **kwargs: Any) -> Any: pass @staticmethod def setup_context(ctx: Any, inputs: Tuple[Any, ...], output: Any) -> None: pass- The forward no longer accepts a ctx argument.
- Instead, you must also override the :meth:
torch.autograd.Function.setup_contextstaticmethod to handle setting up thectxobject.outputis the output of the forward,inputsare a Tuple of inputs to the forward. - See :ref:
extending-autogradfor more details
The context can be used to store arbitrary data that can be then retrieved during the backward pass. Tensors should not be stored directly on
ctx(though this is not currently enforced for backward compatibility). Instead, tensors should be saved either with :func:ctx.save_for_backwardif they are intended to be used inbackward(equivalently,vjp) or :func:ctx.save_for_forwardif they are intended to be used for injvp.