import math from typing import Optional, Literal, Tuple import torch import torch.nn.functional as F from torch import nn from .LMConfig import LMConfig block_size = 128 attn_impl: Literal["naive", "absorb"] = "absorb" class Linear(nn.Module): """ 这是deepseek的线性层实现，暂时没有使用 Custom linear layer with support for quantized weights and optional bias. Args: in_features (int): Number of input features. out_features (int): Number of output features. bias (bool): Whether to include a bias term. Defaults to False. dtype (optional): Data type for the layer. Defaults to `torch.bfloat16`. """ dtype = torch.bfloat16 def __init__(self, in_features: int, out_features: int, bias: bool = False, dtype = None): super().__init__() self.in_features = in_features self.out_features = out_features self.weight = nn.Parameter(torch.empty(out_features, in_features, dtype=dtype or Linear.dtype)) nn.init.kaiming_normal_(self.weight, mode='fan_out', nonlinearity='linear') self.register_parameter("scale", None) if bias: self.bias = nn.Parameter(torch.empty(out_features)) else: self.register_parameter("bias", None) def forward(self, x: torch.Tensor) -> torch.Tensor: """ Forward pass for the custom linear layer. Args: x (torch.Tensor): Input tensor. Returns: torch.Tensor: Transformed tensor after linear computation. """ return F.linear(x, self.weight, self.bias) def apply_rotary_emb(x: torch.Tensor, freqs_cis: torch.Tensor) -> torch.Tensor: """ Applies rotary positional embeddings to the input tensor. Args: x (torch.Tensor): Input tensor with positional embeddings to be applied. freqs_cis (torch.Tensor): Precomputed complex exponential values for positional embeddings. Returns: torch.Tensor: Tensor with rotary embeddings applied. """ dtype = x.dtype x = torch.view_as_complex(x.float().view(*x.shape[:-1], -1, 2)) freqs_cis = freqs_cis.view(1, x.size(1), 1, x.size(-1)) y = torch.view_as_real(x * freqs_cis).flatten(3) return y.to(dtype) class MLA(nn.Module): """ Multi-Headed Attention Layer (MLA). Attributes: dim (int): Dimensionality of the input features. n_heads (int): Number of attention heads. n_local_heads (int): Number of local attention heads for distributed systems. q_lora_rank (int): Rank for low-rank query projection. kv_lora_rank (int): Rank for low-rank key/value projection. qk_nope_head_dim (int): Dimensionality of non-positional query/key projections. qk_rope_head_dim (int): Dimensionality of rotary-positional query/key projections. qk_head_dim (int): Total dimensionality of query/key projections. v_head_dim (int): Dimensionality of value projections. softmax_scale (float): Scaling factor for softmax in attention computation. """ def __init__(self, args: LMConfig): super().__init__() self.dim = args.dim self.n_heads = args.n_heads self.n_local_heads = args.n_heads self.q_lora_rank = args.q_lora_rank self.kv_lora_rank = args.kv_lora_rank self.qk_nope_head_dim = args.qk_nope_head_dim self.qk_rope_head_dim = args.qk_rope_head_dim self.qk_head_dim = args.qk_nope_head_dim + args.qk_rope_head_dim self.v_head_dim = args.v_head_dim dtype = args.torch_dtype if self.q_lora_rank == 0: self.wq = nn.Linear(self.dim, self.n_heads * self.qk_head_dim, bias=False, dtype=dtype) else: self.wq_a = nn.Linear(self.dim, self.q_lora_rank, bias=False, dtype=dtype) self.q_norm = nn.RMSNorm(self.q_lora_rank, dtype=dtype) self.wq_b = nn.Linear(self.q_lora_rank, self.n_heads * self.qk_head_dim, bias=False, dtype=dtype) self.wkv_a = nn.Linear(self.dim, self.kv_lora_rank + self.qk_rope_head_dim, bias=False, dtype=dtype) self.kv_norm = nn.RMSNorm(self.kv_lora_rank, dtype=dtype) self.wkv_b = nn.Linear(self.kv_lora_rank, self.n_heads * (self.qk_nope_head_dim + self.v_head_dim), bias=False, dtype=dtype) self.wo = nn.Linear(self.n_heads * self.v_head_dim, self.dim, bias=False, dtype=dtype) self.softmax_scale = self.qk_head_dim ** -0.5 if args.max_seq_len > args.original_seq_len: mscale = 0.1 * args.mscale * math.log(args.rope_factor) + 1.0 self.softmax_scale = self.softmax_scale * mscale * mscale self.use_cache = args.use_cache if self.use_cache: if attn_impl == "naive": self.register_buffer("k_cache", torch.zeros(args.max_batch_size, args.max_seq_len, self.n_local_heads, self.qk_head_dim, dtype=dtype), persistent=False) self.register_buffer("v_cache", torch.zeros(args.max_batch_size, args.max_seq_len, self.n_local_heads, self.v_head_dim, dtype=dtype), persistent=False) else: self.register_buffer("kv_cache", torch.zeros(args.max_batch_size, args.max_seq_len, self.kv_lora_rank, dtype=dtype), persistent=False) self.register_buffer("pe_cache", torch.zeros(args.max_batch_size, args.max_seq_len, self.qk_rope_head_dim, dtype=dtype), persistent=False) def forward(self, x: torch.Tensor, start_pos: int, freqs_cis: torch.Tensor, past_key_value: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, mask: Optional[torch.Tensor] = None): """ Forward pass for the Multi-Headed Attention Layer (MLA). Args: x (torch.Tensor): Input tensor of shape (batch_size, seq_len, dim). start_pos (int): Starting position in the sequence for caching. freqs_cis (torch.Tensor): Precomputed complex exponential values for rotary embeddings. mask (Optional[torch.Tensor]): Mask tensor to exclude certain positions from attention. Returns: torch.Tensor: Output tensor with the same shape as the input. """ bsz, seqlen, _ = x.size() end_pos = start_pos + seqlen if self.q_lora_rank == 0: q = self.wq(x) else: q = self.wq_b(self.q_norm(self.wq_a(x))) q = q.view(bsz, seqlen, self.n_local_heads, self.qk_head_dim) q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1) q_pe = apply_rotary_emb(q_pe, freqs_cis) kv = self.wkv_a(x) kv, k_pe = torch.split(kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1) k_pe = apply_rotary_emb(k_pe.unsqueeze(2), freqs_cis) v = None if attn_impl == "naive": q = torch.cat([q_nope, q_pe], dim=-1) kv = self.wkv_b(self.kv_norm(kv)) kv = kv.view(bsz, seqlen, self.n_local_heads, self.qk_nope_head_dim + self.v_head_dim) k_nope, v = torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1) k = torch.cat([k_nope, k_pe.expand(-1, -1, self.n_local_heads, -1)], dim=-1) if self.use_cache: self.k_cache[:bsz, start_pos:end_pos] = k self.v_cache[:bsz, start_pos:end_pos] = v k = self.k_cache[:bsz, :end_pos] v = self.k_cache[:bsz, :end_pos] scores = torch.einsum("bshd,bthd->bsht", q, k) * self.softmax_scale else: # wkv_b = self.wkv_b.weight if self.wkv_b.scale is None else weight_dequant(self.wkv_b.weight, self.wkv_b.scale, block_size) wkv_b = self.wkv_b.weight wkv_b = wkv_b.view(self.n_local_heads, -1, self.kv_lora_rank) q_nope = torch.einsum("bshd,hdc->bshc", q_nope, wkv_b[:, :self.qk_nope_head_dim]) kv = self.kv_norm(kv) k_pe = k_pe.squeeze(2) if self.use_cache: self.kv_cache[:bsz, start_pos:end_pos] = kv self.pe_cache[:bsz, start_pos:end_pos] = k_pe kv = self.kv_cache[:bsz, :end_pos] k_pe = self.pe_cache[:bsz, :end_pos] scores = (torch.einsum("bshc,btc->bsht", q_nope, kv) + torch.einsum("bshr,btr->bsht", q_pe, k_pe)) * self.softmax_scale if mask is not None: scores += mask.unsqueeze(1) scores = scores.softmax(dim=-1, dtype=torch.float32).type_as(x) if attn_impl == "naive": x = torch.einsum("bsht,bthd->bshd", scores, v) else: x = torch.einsum("bsht,btc->bshc", scores, kv) x = torch.einsum("bshc,hdc->bshd", x, wkv_b[:, -self.v_head_dim:]) x = self.wo(x.flatten(2)) return x, None