# 带缓存的注意力计算
def attention_with_cache(q, k, v, k_cache, v_cache):
    # 只计算新token的K和V，复用缓存的K和V
    new_k = torch.cat([k_cache, k], dim=1)
    new_v = torch.cat([v_cache, v], dim=1)
    
    scores = torch.matmul(q, new_k.transpose(-2, -1))
    attention_weights = torch.softmax(scores, dim=-1)
    output = torch.matmul(attention_weights, new_v)
    return output, new_k, new_v