输出投影

    self.out_proj = OptimizedTensorParallelLinear(d_model, d_model, num_partitions)
    
    # Dropout层
    self.dropout_layer = nn.Dropout(dropout)

def forward(self, 
            query: torch.Tensor, 
            key: torch.Tensor, 
            value: torch.Tensor, 
            mask: Optional[torch.Tensor] = None) -> torch.Tensor:
    """Flash注意力前向传播"""
    batch_size, seq_len, _ = query.shape
    
    # 查询、键、值投影
    q = self.q_proj(query)  # [batch_size, seq_len, d_model]
    k = self.k_proj(key)    # [batch_size, seq_len, d_model]
    v = self.v_proj(value)  # [batch_size, seq_len, d_model]
    
    # 重塑为多头格式
    q = q.view(batch_size, seq_len, self.nhead, self.d_k).transpose(1, 2)
    k = k.view(batch_size, seq_len, self.nhead, self.d_k).transpose(1, 2)
    v = v.view(batch_size, seq_len, self.nhead, self.d_k).transpose(1, 2)
    
    # Flash注意力计算（分块计算以减少内存使用）
    output = self._flash_attention(q, k, v, mask)
    
    # 重塑回原始格式
    output = output.transpose(1, 2).contiguous()
    output = output.view(batch_size, seq_len, self.d_model)
    
    # 输出投影
    output = self.out_proj(output)
    
    return output

def _flash_attention(self, 
                    q: torch.Tensor, 
                    k: torch.Tensor, 
                    v: torch.Tensor, 
                    mask: Optional[torch.Tensor] = None) -> torch.Tensor:
    """Flash注意力实现"""
    batch_size, seq_len, nhead, d_k = q.shape
    
    # 分块大小
    block_size = 64  # 可以根据GPU内存调整
    
    # 初始化输出
    output = torch.zeros_like(q)
    
    # 分块计算
    for i in range(0, seq_len, block_size):
        for j in range(0, seq_len, block_size):
            # 获取当前块
            q_block = q[:, i:i+block_size, :, :]
            k_block = k[:, j:j+block_size, :, :]
            v_block = v[:, j:j+block_size, :, :]
            
            # 计算当前块的注意力
            scores = torch.matmul(q_block, k_block.transpose(-2, -1)) / math.sqrt(d_k)
            
            # 应用mask
            if mask is not None:
                mask_block = mask[:, i:i+block_size, j:j+block_size]
                scores = scores.masked_fill(mask_block == 0, float('-inf'))
            
            # 计算注意力权重
            attn_weights = torch.softmax(scores, dim=-1)
            attn_weights = self.dropout_layer(attn_weights)
            
            # 应用注意力权重
            output_block = torch.matmul(attn_weights, v_block)
            output[:, i:i+block_size, :, :] += output_block
    
    return output
​

演示优化的注意力机制

def demonstrate_optimized_attention():
"""演示优化的注意力机制"""
print("\n优化的注意力机制演示")
print("-" * 50)

# 参数设置
d_model = 512
nhead = 8
num_partitions = 4
batch_size = 16
seq_len = 128

# 创建输入数据
query = torch.randn(batch_size, seq_len, d_model)
key = torch.randn(batch_size, seq_len, d_model)
value = torch.randn(batch_size, seq_len, d_model)

# 测试标准张量并行注意力
print("\n标准张量并行注意力:")
std_attention = OptimizedTensorParallelAttention(d_model, nhead, num_partitions)
std_output = std_attention(query, key, value)
print(f"输出形状: {std_output.shape}")

# 测试Flash注意力
print("\nFlash注意力:")
flash_attention = FlashAttentionOptimized(d_model, nhead, num_partitions)
flash_output = flash_attention(query, key, value)
print(f"输出形状: {flash_output.shape}")

# 验证精度
print("\n精度验证:")
print(f"标准输出均值: {std_output.mean().item():.6f}")
print(f"Flash输出均值: {flash_output.mean().item():.6f}")
print(f"差异: {torch.norm(std_output - flash_output).item():.6f}")

# 性能对比
print("\n性能对比:")
import time

# 测试标准注意力
start_time = time.time()
for _ in range(10):
    _ = std_attention(query, key, value)
std_time = time.time() - start
​