# xformers 기반의 메모리 효율적인 인과적 어텐션 구현 def memory_efficient_causal_attention(q, k, v): # 필요한 어텐션 점수만 계산 (마스킹된 점수는 계산하지 않음) batch_size, seq_len, num_heads, head_dim = q.shape # 출력 텐서 초기화 output = torch.zeros_like(q) # 누적 소프트맥스 상태 추적 softmax_scale = 1.0 / math.sqrt(head_dim) m_prev = torch.zeros((batch_size, num_heads, 1), device=q.device) # 최대값 저장 p_sum = torch.zeros((batch_size, num_heads, 1), device=q.device) # 합 저장 o_acc = torch.zeros((batch_size, num_heads, head_dim), device=q.device) # 가중합 저장 # 순차적 처리로 메모리 사용 최소화 for i in range(seq_len): # 현재 위치의 쿼리 qi = q[:, i:i+1] # [B, 1, H, D] # 현재 위치까지의 키와 값만 처리 (인과성) ki = k[:, :i+1] # [B, i+1, H, D] vi = v[:, :i+1] # [B, i+1, H, D] # 어텐션 점수 계산 (스케일링 포함) si = torch.matmul(qi, ki.transpose(-1, -2)) * softmax_scale # [B, 1, H, i+1] # Rabe & Staats 방식: 소프트맥스 최적화 mi = torch.max(si, dim=-1, keepdim=True)[0] m_curr = torch.maximum(m_prev, mi) # 이전 누적값 스케일 조정 scale_factor = torch.exp(m_prev - m_curr) p_sum = p_sum * scale_factor o_acc = o_acc * scale_factor.unsqueeze(-1) # 새로운 어텐션 가중치 계산 pi = torch.exp(si - m_curr) p_sum = p_sum + pi.sum(dim=-1, keepdim=True) # 출력 누적 (중간 어텐션 행렬 저장 없이) o_acc = o_acc + torch.matmul(pi, vi).squeeze(1) # 현재 위치 출력 계산 output[:, i] = o_acc / p_sum # 상태 업데이트 m_prev = m_curr return output