DeepSeek-V4的CSA（Compressed Sparse Attention）和HCA（Heavily Compressed Attention）机制通过不同的策略优化注意力计算。以下是两种机制的核心Python实现代码示例，结合具体的技术细节进行说明。 ## 1. HCA（Heavily Compressed Attention）机制实现 HCA通过极端压缩和混合精度存储来优化KV Cache，核心是低秩压缩、分块合并和FP4存储。 ```python import torch import torch.nn as nn import torch.nn.functional as F class HeavilyCompressedAttention(nn.Module): """HCA: 极端压缩注意力机制实现""" def __init__(self, embed_dim=8192, latent_dim=512, # 低秩维度 block_size=128, # 分块大小 num_heads=64, fp4_storage=True): # 启用FP4存储 super().__init__() self.embed_dim = embed_dim self.latent_dim = latent_dim self.block_size = block_size self.num_heads = num_heads self.fp4_storage = fp4_storage # 低秩投影矩阵 [ref_1] self.kv_proj = nn.Linear(embed_dim, latent_dim, bias=False) # KV压缩投影 self.q_proj = nn.Linear(embed_dim, embed_dim, bias=False) # Q保持全维度 # 解压缩矩阵（用于从低维恢复高维） self.k_decompress = nn.Linear(latent_dim, embed_dim, bias=False) self.v_decompress = nn.Linear(latent_dim, embed_dim, bias=False) # 分块合并的权重矩阵 self.block_weights = nn.Parameter( torch.randn(block_size, latent_dim) / latent_dim**0.5 ) def compress_to_fp4(self, tensor): """将FP16张量压缩为FP4格式 [ref_1]""" if not self.fp4_storage: return tensor # 模拟FP4量化：将FP16量化为4位 # 实际实现会使用更复杂的量化策略 min_val = tensor.min() max_val = tensor.max() scale = (max_val - min_val) / 15 # 4位有15个可表示的值（排除0） # 量化到4位整数 quantized = ((tensor - min_val) / scale).round().clamp(0, 15) # 压缩存储：每两个4位值打包成一个字节 # 这里简化为返回量化后的整数 return quantized.to(torch.uint8), min_val, scale def decompress_from_fp4(self, quantized, min_val, scale): """从FP4恢复为FP16张量 [ref_1]""" if not self.fp4_storage: return quantized # 反量化 return quantized.float() * scale + min_val def block_merge(self, kv_latent): """分块合并：每128个token合并为一个entry [ref_1]""" B, L, D = kv_latent.shape num_blocks = L // self.block_size # 重塑为块状结构 kv_blocks = kv_latent.view(B, num_blocks, self.block_size, D) # 使用可学习权重进行加权合并 weights = F.softmax(self.block_weights, dim=0) # [block_size, D] weights = weights.unsqueeze(0).unsqueeze(0) # [1, 1, block_size, D] # 加权合并 merged = torch.sum(kv_blocks * weights, dim=2) # [B, num_blocks, D] return merged def forward(self, x, use_cache=False, past_kv=None): """ 前向传播 x: [batch_size, seq_len, embed_dim] past_kv: 缓存的压缩KV """ B, L, D = x.shape # 1. 生成查询向量（保持全维度） Q = self.q_proj(x) # [B, L, D] # 2. 生成Key和Value的低秩表示 [ref_1] kv_latent = self.kv_proj(x) # [B, L, latent_dim] # 3. 分块合并 if use_cache: # 推理时使用缓存 if past_kv is None: # 首次计算，创建缓存 merged_kv = self.block_merge(kv_latent) # FP4压缩存储 [ref_1] if self.fp4_storage: kv_cache = self.compress_to_fp4(merged_kv) else: kv_cache = merged_kv else: # 从缓存恢复 if self.fp4_storage: quantized, min_val, scale = past_kv merged_kv = self.decompress_from_fp4(quantized, min_val, scale) else: merged_kv = past_kv else: # 训练时不使用缓存 merged_kv = self.block_merge(kv_latent) # 4. 从低维恢复高维K和V K_restored = self.k_decompress(merged_kv) # [B, num_blocks, D] V_restored = self.v_decompress(merged_kv) # [B, num_blocks, D] # 5. 多头注意力计算 num_blocks = K_restored.shape[1] # 重塑为多头 Q = Q.view(B, L, self.num_heads, D // self.num_heads).transpose(1, 2) K = K_restored.view(B, num_blocks, self.num_heads, D // self.num_heads).transpose(1, 2) V = V_restored.view(B, num_blocks, self.num_heads, D // self.num_heads).transpose(1, 2) # 注意力计算 attn_scores = torch.matmul(Q, K.transpose(-2, -1)) / (D // self.num_heads)**0.5 attn_weights = F.softmax(attn_scores, dim=-1) attn_output = torch.matmul(attn_weights, V) # 重塑回原始维度 attn_output = attn_output.transpose(1, 2).contiguous().view(B, L, D) # 返回输出和缓存 if use_cache: kv_cache = merged_kv if not self.fp4_storage else self.compress_to_fp4(merged_kv) return attn_output, kv_cache else: return attn_output ``` ## 2. CSA（Compressed Sparse Attention）机制实现 CSA通过压缩和稀疏选择来减少计算量，核心是轻量级索引器和Top-K选择。 ```python class CompressedSparseAttention(nn.Module): """CSA: 压缩稀疏注意力机制实现""" def __init__(self, embed_dim=8192, compressed_dim=512, # 压缩维度 block_size=4, # 合并块大小 top_k=128, # 保留的块数 num_heads=64): super().__init__() self.embed_dim = embed_dim self.compressed_dim = compressed_dim self.block_size = block_size self.top_k = top_k self.num_heads = num_heads # 压缩投影 [ref_1] self.kv_compress = nn.Linear(embed_dim, compressed_dim, bias=False) self.q_proj = nn.Linear(embed_dim, embed_dim, bias=False) # 闪电索引器（轻量级评分网络） self.indexer = nn.Sequential( nn.Linear(compressed_dim, compressed_dim // 2), nn.GELU(), nn.Linear(compressed_dim // 2, 1) # 输出评分 ) # MLA组件（复用低秩注意力） self.mla = MLAAttention(embed_dim, compressed_dim, num_heads) def block_compress(self, x, weights=None): """块压缩：每4个token加权合并 [ref_1]""" B, L, D = x.shape num_blocks = L // self.block_size # 重塑为块状 x_blocks = x.view(B, num_blocks, self.block_size, D) # 动态权重生成（如果未提供） if weights is None: weights = torch.ones(B, num_blocks, self.block_size, 1, device=x.device) weights = F.softmax(weights, dim=2) # 加权合并 compressed = torch.sum(x_blocks * weights, dim=2) # [B, num_blocks, D] return compressed, weights def lightning_indexer(self, query_compressed, key_compressed): """闪电索引器：快速评分 [ref_1]""" B, L, D = query_compressed.shape _, N, _ = key_compressed.shape # 扩展维度以进行批量矩阵乘法 query_expanded = query_compressed.unsqueeze(2) # [B, L, 1, D] key_expanded = key_compressed.unsqueeze(1) # [B, 1, N, D] # 计算初步相关性得分 raw_scores = torch.sum(query_expanded * key_expanded, dim=-1) # [B, L, N] # 通过索引器网络细化得分 # 将query和key拼接作为输入 combined = query_expanded + key_expanded # [B, L, N, D] combined = combined.view(-1, D) refined_scores = self.indexer(combined) # [B*L*N, 1] refined_scores = refined_scores.view(B, L, N) # 综合得分 final_scores = raw_scores * 0.7 + refined_scores * 0.3 return final_scores def forward(self, x, use_cache=False, past_kv=None): """ 前向传播 x: [batch_size, seq_len, embed_dim] """ B, L, D = x.shape # 1. 生成查询向量 Q = self.q_proj(x) # [B, L, D] # 2. 压缩Key和Value [ref_1] kv_compressed = self.kv_compress(x) # [B, L, compressed_dim] # 3. 块压缩：每4个token合并 [ref_1] kv_blocks, block_weights = self.block_compress(kv_compressed) num_blocks = kv_blocks.shape[1] # 通常是 L // 4 # 4. 闪电索引器评分 # 查询也需要压缩以进行评分 q_compressed = self.kv_compress(x) q_blocks, _ = self.block_compress(q_compressed) # 计算所有块的得分 block_scores = self.lightning_indexer(q_blocks, kv_blocks) # [B, num_blocks, num_blocks] # 5. Top-K选择：每个查询选择最相关的K个块 [ref_1] # 取平均得分作为块的全局重要性 global_scores = block_scores.mean(dim=1) # [B, num_blocks] # 选择top-k个块 topk_values, topk_indices = torch.topk(global_scores, k=min(self.top_k, num_blocks), dim=-1) # 6. 收集被选中的块 selected_kv = [] for b in range(B): selected = torch.index_select(kv_blocks[b], 0, topk_indices[b]) selected_kv.append(selected) selected_kv = torch.stack(selected_kv, dim=0) # [B, top_k, compressed_dim] # 7. 对选中的块执行MLA注意力 [ref_1] # 这里简化处理，实际MLA会更复杂 attn_output = self.mla(Q, selected_kv) # 缓存相关信息 if use_cache: cache_data = { 'kv_blocks': kv_blocks, 'block_scores': block_scores, 'topk_indices': topk_indices } return attn_output, cache_data else: return attn_output class MLAAttention(nn.Module): """简化的MLA注意力实现""" def __init__(self, embed_dim, latent_dim, num_heads): super().__init__() self.embed_dim = embed_dim self.latent_dim = latent_dim self.num_heads = num_heads # 低秩投影 self.q_proj = nn.Linear(embed_dim, embed_dim) self.k_proj = nn.Linear(latent_dim, embed_dim // 2) # 部分维度 self.v_proj = nn.Linear(latent_dim, embed_dim // 2) def forward(self, query, kv_compressed): B, L, D = query.shape _, N, _ = kv_compressed.shape # 投影 Q = self.q_proj(query) K = self.k_proj(kv_compressed) V = self.v_proj(kv_compressed) # 重塑为多头 head_dim = D // self.num_heads Q = Q.view(B, L, self.num_heads, head_dim).transpose(1, 2) K = K.view(B, N, self.num_heads, head_dim // 2).transpose(1, 2) V = V.view(B, N, self.num_heads, head_dim // 2).transpose(1, 2) # 注意力计算 attn_scores = torch.matmul(Q, K.transpose(-2, -1)) / (head_dim ** 0.5) attn_weights = F.softmax(attn_scores, dim=-1) attn_output = torch.matmul(attn_weights, V) # 重塑回原始维度 attn_output = attn_output.transpose(1, 2).contiguous() attn_output = attn_output.view(B, L, D // 2) # 补全另一半维度（简化处理） output = torch.cat([attn_output, attn_output], dim=-1) return output ``` ## 3. 混合架构集成示例 在实际的DeepSeek-V4中，CSA和HCA会交替使用： ```python class DeepSeekV4Layer(nn.Module): """DeepSeek-V4的Transformer层，交替使用CSA和HCA""" def __init__(self, layer_id, embed_dim=8192, num_heads=64): super().__init__() self.layer_id = layer_id # 根据层ID决定使用CSA还是HCA [ref_1] # 低层多用CSA，高层多用HCA if layer_id % 3 == 0: # 每3层一个HCA self.attention = HeavilyCompressedAttention( embed_dim=embed_dim, latent_dim=512, block_size=128, num_heads=num_heads, fp4_storage=True ) self.attn_type = "HCA" else: self.attention = CompressedSparseAttention( embed_dim=embed_dim, compressed_dim=512, block_size=4, top_k=128, num_heads=num_heads ) self.attn_type = "CSA" # 前馈网络 self.ffn = nn.Sequential( nn.Linear(embed_dim, embed_dim * 4), nn.GELU(), nn.Linear(embed_dim * 4, embed_dim) ) # 层归一化 self.norm1 = nn.LayerNorm(embed_dim) self.norm2 = nn.LayerNorm(embed_dim) def forward(self, x, attention_mask=None, use_cache=False, past_kv=None): residual = x # 注意力层 x = self.norm1(x) if use_cache: attn_output, new_kv = self.attention(x, use_cache=True, past_kv=past_kv) else: attn_output = self.attention(x) new_kv = None x = residual + attn_output # 前馈层 residual = x x = self.norm2(x) ffn_output = self.ffn(x) x = residual + ffn_output if use_cache: return x, new_kv else: return x class DeepSeekV4Model(nn.Module): """完整的DeepSeek-V4模型""" def __init__(self, num_layers=64, embed_dim=8192, num_heads=64): super().__init__() self.embed_dim = embed_dim # 词嵌入 self.embedding = nn.Embedding(100000, embed_dim) # Transformer层 self.layers = nn.ModuleList([ DeepSeekV4Layer(i, embed_dim, num_heads) for i in range(num_layers) ]) # 输出层 self.output_norm = nn.LayerNorm(embed_dim) self.lm_head = nn.Linear(embed_dim, 100000, bias=False) # KV缓存 self.kv_cache = [None] *