# 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