Function attention_forward

tests/dump_phase7_reference.py:219–258

(
    q_in: torch.Tensor,
    k_in: torch.Tensor,
    v_in: torch.Tensor,
    prefix: str,
    weights: Dict[str, torch.Tensor],
    num_heads: int,
    rope_cis: torch.Tensor | None = None,
    repeat_freqs_k: bool = False,
    num_k_exclude_rope: int = 0,
)

def attention_forward(
    q_in: torch.Tensor,
    k_in: torch.Tensor,
    v_in: torch.Tensor,
    prefix: str,
    weights: Dict[str, torch.Tensor],
    num_heads: int,
    rope_cis: torch.Tensor | None = None,
    repeat_freqs_k: bool = False,
    num_k_exclude_rope: int = 0,
) -> torch.Tensor:
    q = F.linear(q_in, weights[prefix + ".q_proj.weight"].float(), weights[prefix + ".q_proj.bias"].float())
    k = F.linear(k_in, weights[prefix + ".k_proj.weight"].float(), weights[prefix + ".k_proj.bias"].float())
    v = F.linear(v_in, weights[prefix + ".v_proj.weight"].float(), weights[prefix + ".v_proj.bias"].float())

    bsz, nq, dim = q.shape
    nk = k.shape[1]
    head_dim = dim // num_heads

    qh = q.view(bsz, nq, num_heads, head_dim).permute(0, 2, 1, 3)
    kh = k.view(bsz, nk, num_heads, head_dim).permute(0, 2, 1, 3)
    vh = v.view(bsz, nk, num_heads, head_dim).permute(0, 2, 1, 3)

    if rope_cis is not None:
        nk_rope = kh.shape[2] - num_k_exclude_rope
        if nk_rope > 0:
            qh, kh_rope = apply_rotary_enc(
                qh,
                kh[:, :, :nk_rope, :],
                rope_cis,
                repeat_freqs_k=repeat_freqs_k,
            )
            if nk_rope < kh.shape[2]:
                kh = torch.cat([kh_rope, kh[:, :, nk_rope:, :]], dim=2)
            else:
                kh = kh_rope

    out = F.scaled_dot_product_attention(qh, kh, vh)
    out = out.permute(0, 2, 1, 3).reshape(bsz, nq, dim)
    return F.linear(out, weights[prefix + ".out_proj.weight"].float(), weights[prefix + ".out_proj.bias"].float())


def mlp_forward(x: torch.Tensor, weights: Dict[str, torch.Tensor], prefix: str, num_layers: int) -> torch.Tensor: