Flash Attention – Tính Toán Attention Theo Chunk (Memory-Efficient)

Mục tiêu kiến thức

Standard attention phải tạo ma trận N × N tốn O(N²) bộ nhớ. Flash Attention tính toán theo chunk nhỏ:

Online Softmax với 2 chunk: m_i = max(m_i-1, max_j S_ij) (max tracker) l_i = e^{m_i-1 - m_i} l_i-1 + ∑_j e^{S_ij - m_i} (sum tracker) O_i = e^{m_i-1 - m_i} l_i-1 O_i-1 + ∑_j e^{S_ij - m_i} V_jl_i

Yêu cầu

def chunked_attention(
    Q: torch.Tensor,  # [N, d]
    K: torch.Tensor,  # [N, d]
    V: torch.Tensor,  # [N, d]
    chunk_size: int = 32
) -> torch.Tensor:
    # Returns attention output [N, d] WITHOUT materializing full N×N attention matrix
    pass

Ràng buộc

• Thời gian chạy tối đa: 10000ms.
• Giới hạn bộ nhớ: 1024MB.
• Dữ liệu đầu vào hợp lệ theo đúng kiểu dữ liệu và miền giá trị được mô tả.