5.1 从朴素 Attention 到分块执行

Q tile → 加载 K/V tile → S=QKᵀ×scale + mask
       → 行最大值 m_new → p=exp(S-m_new)
       → l_new=exp(m_old-m_new)×l_old+sum(p)
       → acc_new=exp(m_old-m_new)×acc_old+pV
       → 下一 tile；最终 O=acc/l

import torch

def online_attention(q, k, v, block=64):
    # q:[Q,D], k/v:[K,D]；教学版，先用 float32 验证
    scale = q.shape[-1] ** -0.5
    out = []
    for qs in range(0, q.shape[0], block):
        qb = q[qs:qs + block].float()
        m = torch.full((qb.shape[0], 1), -float('inf'), device=q.device)
        l = torch.zeros_like(m)
        acc = torch.zeros((qb.shape[0], v.shape[-1]), device=q.device)
        for ks in range(0, k.shape[0], block):
            score = qb @ k[ks:ks + block].float().T * scale
            # causal: 全局 q 位置与 k 位置比较；空 tile 不参加更新
            qi = torch.arange(qs, qs + qb.shape[0], device=q.device)[:, None]
            ki = torch.arange(ks, ks + score.shape[1], device=q.device)[None, :]
            score = score.masked_fill(ki > qi, -float('inf'))
            new_m = torch.maximum(m, score.max(-1, keepdim=True).values)
            safe_m = torch.where(torch.isfinite(new_m), new_m, torch.zeros_like(new_m))
            alpha = torch.where(torch.isfinite(m), torch.exp(m - safe_m), 0)
            p = torch.exp(score - safe_m)
            acc = acc * alpha + p @ v[ks:ks + block].float()
            l = l * alpha + p.sum(-1, keepdim=True)
            m = new_m
        out.append(acc / l.clamp_min(1e-20))
    return torch.cat(out, dim=0)