| 199 | |
| 200 | |
| 201 | def attention_forward(q: torch.Tensor, |
| 202 | k: torch.Tensor, |
| 203 | v: torch.Tensor, |
| 204 | prefix: str, |
| 205 | weights: Dict[str, torch.Tensor], |
| 206 | num_heads: int) -> torch.Tensor: |
| 207 | q_w = weights[prefix + ".q_proj.weight"].float() |
| 208 | q_b = weights[prefix + ".q_proj.bias"].float() |
| 209 | k_w = weights[prefix + ".k_proj.weight"].float() |
| 210 | k_b = weights[prefix + ".k_proj.bias"].float() |
| 211 | v_w = weights[prefix + ".v_proj.weight"].float() |
| 212 | v_b = weights[prefix + ".v_proj.bias"].float() |
| 213 | out_w = weights[prefix + ".out_proj.weight"].float() |
| 214 | out_b = weights[prefix + ".out_proj.bias"].float() |
| 215 | |
| 216 | q_proj = F.linear(q, q_w, q_b) |
| 217 | k_proj = F.linear(k, k_w, k_b) |
| 218 | v_proj = F.linear(v, v_w, v_b) |
| 219 | |
| 220 | bsz, nq, dim = q_proj.shape |
| 221 | nk = k_proj.shape[1] |
| 222 | head_dim = dim // num_heads |
| 223 | |
| 224 | qh = q_proj.view(bsz, nq, num_heads, head_dim).permute(0, 2, 1, 3) |
| 225 | kh = k_proj.view(bsz, nk, num_heads, head_dim).permute(0, 2, 1, 3) |
| 226 | vh = v_proj.view(bsz, nk, num_heads, head_dim).permute(0, 2, 1, 3) |
| 227 | |
| 228 | out = F.scaled_dot_product_attention(qh, kh, vh) |
| 229 | out = out.permute(0, 2, 1, 3).reshape(bsz, nq, dim) |
| 230 | return F.linear(out, out_w, out_b) |
| 231 | |
| 232 | |
| 233 | def mlp_forward(x: torch.Tensor, |