Ring Attention：分布式长上下文

第0步: GPU_i 用本地 Q_i 和 K_i 计算注意力，结果存为 partial_O_i
       同时，GPU_i 将 K_i, V_i 发送给 GPU_{(i+1) % N}

第1步: GPU_i 接收 K_{(i-1) % N}, V_{(i-1) % N}
       用 Q_i 和接收的 K 计算注意力，累加到 partial_O_i
       同时发送 K_{(i-1) % N}, V_{(i-1) % N} 给 GPU_{(i+1) % N}

... (重复N-1步)

第N-1步: GPU_i 已与所有K/V块计算过注意力
         partial_O_i 即为最终结果

def ring_attention_forward(Q_chunks, K_chunks, V_chunks, rank, world_size):
    """
    Q_chunks, K_chunks, V_chunks: 本地GPU的Q/K/V
    rank: 当前GPU编号
    world_size: GPU数量
    """
    # 本地Q和输出
    Q_local = Q_chunks[rank]
    partial_O = torch.zeros_like(Q_local)
    
    # 当前持有的K/V（初始为本地K/V）
    K_current = K_chunks[rank].clone()
    V_current = V_chunks[rank].clone()
    
    for step in range(world_size):
        # 当前K/V的全局块索引
        kv_rank = (rank - step) % world_size
        
        # 计算Q_local与K_current的注意力
        # 需要处理因果掩码：只有kv_rank对应的token在Q_local之前才能计算
        attn_out = blockwise_attention(Q_local, K_current, V_current, 
                                        q_start=rank * block_size,
                                        kv_start=kv_rank * block_size)
        partial_O += attn_out
        
        if step < world_size - 1:
            # 异步发送当前K/V给下一个GPU
            send_req = dist.isend(K_current, dst=(rank + 1) % world_size)
            send_req_v = dist.isend(V_current, dst=(rank + 1) % world_size)
            
            # 异步接收上一个GPU的K/V
            K_next = torch.empty_like(K_current)
            V_next = torch.empty_like(V_current)
            recv_req = dist.irecv(K_next, src=(rank - 1) % world_size)
            recv_req_v = dist.irecv(V_next, src=(rank - 1) % world_size)
            
            # 等待通信完成
            send_req.wait()
            send_req_v.wait()
            recv_req.wait()
            recv_req_v.wait()
            
            K_current = K_next
            V_current = V_next
    
    return partial_O

def ring_attention_with_causal(Q_local, K_current, V_current, 
                                q_rank, kv_rank, block_size):
    if kv_rank > q_rank:
        # KV全在Q之后，完全不可见
        return torch.zeros_like(Q_local)
    elif kv_rank < q_rank:
        # KV全在Q之前，完全可见
        return standard_attention(Q_local, K_current, V_current)
    else:
        # 同一块，需要因果掩码
        mask = torch.tril(torch.ones(block_size, block_size))
        return masked_attention(Q_local, K_current, V_current, mask)

def ring_flash_attention(Q, K, V, rank, world_size):
    """Ring Attention + Flash Attention"""
    block_size = Q.shape[1]  # 本地序列长度
    
    # Flash Attention的在线softmax状态
    lse = torch.full((Q.shape[0], Q.shape[1], 1), float('-inf'))  # log-sum-exp
    O = torch.zeros_like(Q)
    
    K_current = K.clone()
    V_current = V.clone()
    
    for step in range(world_size):
        kv_rank = (rank - step) % world_size
        
        if kv_rank <= rank:
            # 使用Flash Attention计算块级注意力
            # 返回更新后的O和lse
            O, lse = flash_attention_block(Q, K_current, V_current, 
                                            O, lse, q_rank, kv_rank)
        
        # 环形通信
        if step < world_size - 1:
            K_current, V_current = ring_send_recv(K_current, V_current, 
                                                    rank, world_size)
    
    # 最终归一化
    O = O / torch.exp(lse)
    return O

block_size × d_model × N × 2 ≈ block_size × total_seq_len × d_model