DeepSeek – Neuronix Lab

本文深入拆解 DeepSeek-V3 的多头潜在注意力 (Multi-Head Latent Attention, MLA) 机制，推导其如何将 KV Cache 压缩 93.3% 突破长上下文显存墙；同时详解其无辅助损失 MoE 负载均衡设计与 Multi-Token 预测加速策略。

#CUDA #DeepSeek #KV 缓存

mla_attention_kernel.pypython

import torch
import torch.nn as nn
import torch.nn.functional as F

class MultiHeadLatentAttention(nn.Module):
    def __init__(self, dim=7168, num_heads=128, head_dim=128, kv_lora_rank=512, qk_rope_dim=64):
        super().__init__()
        self.num_heads = num_heads
        self.head_dim = head_dim
        self.kv_lora_rank = kv_lora_rank
        self.qk_rope_dim = qk_rope_dim

Dr. Evelyn Vance (方舟元) 2025-02-14

阅读

DeepSeek-V3/R1 架构深度全景解析：MLA 注意力压缩与 Multi-Token 预测内核实战