MLA(Multi-Head Latent Attention)

# 从 HuggingFace config 读取 MLA 维度。
# qk_nope_head_dim: 不参与 RoPE 的 Q/K 维度。
# qk_rope_head_dim: 参与 RoPE 的 Q/K 维度。
# v_head_dim: value head dim。
# kv_lora_rank: compressed KV latent rank。
qk_nope_head_dim = getattr(config, "qk_nope_head_dim", 0)
qk_rope_head_dim = getattr(config, "qk_rope_head_dim", 0)
v_head_dim = getattr(config, "v_head_dim", 0)
kv_lora_rank = getattr(config, "kv_lora_rank", 0)

# DeepSeek v1 或没有 MLA 维度时，走普通 MHA。
use_mha = config.model_type == "deepseek" or all(
    dim == 0 for dim in (qk_nope_head_dim, qk_rope_head_dim)
)

if use_mha:
    # 普通 MHA：传统 Q/K/V cache。
    attn_cls = DeepseekAttention
elif model_config.use_mla:
    # 优化 MLA：latent KV cache + MLA backend。
    attn_cls = DeepseekV2MLAAttention
else:
    # 数学上仍是 MLA，但会展开成普通 K/V 后走通用 Attention。
    # 这是禁用 MLA 优化时的回退路径。
    attn_cls = DeepseekV2Attention

class DeepseekV2MLAAttention(nn.Module):
    """
    核心组件:
    - q_lora_rank: Q 的低秩维度
    - kv_lora_rank: KV 的低秩维度
    - fused_qkv_a_proj: 融合的 QKV 投影
    - q_a_layernorm, q_b_proj: Q 的低秩分解
    - kv_a_layernorm, kv_b_proj: KV 的低秩分解
    - o_proj: 输出投影
    """

    def __init__(self, ...):
        ...
        if self.q_lora_rank is not None:
            # DeepSeek V3/R1 常见路径：把 q_a_proj 和 kv_a_proj 融合成一个 GEMM。
            # 输出维度是 [q_lora_rank, kv_lora_rank + qk_rope_head_dim]。
            self.fused_qkv_a_proj = DeepSeekV2FusedQkvAProjLinear(
                proj_input_size,
                [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
                quant_config=quant_config,
                prefix=f"{prefix}.fused_qkv_a_proj",
            )
        else:
            # 没有 q_lora_rank 时，KV latent 单独投影。
            self.kv_a_proj_with_mqa = ReplicatedLinear(
                proj_input_size,
                self.kv_lora_rank + self.qk_rope_head_dim,
                bias=False,
                quant_config=quant_config,
                prefix=f"{prefix}.kv_a_proj_with_mqa",
            )

        # KV latent 先 RMSNorm，再通过 kv_b_proj 展开成每个 head 的 k_nope + v。
        self.kv_a_layernorm = RMSNorm(self.kv_lora_rank, eps=config.rms_norm_eps)
        self.kv_b_proj = ColumnParallelLinear(
            self.kv_lora_rank,
            self.num_heads * (self.qk_nope_head_dim + self.v_head_dim),
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.kv_b_proj",
        )

        # 输出投影。注意 MLA attention 内部先得到 [heads * v_head_dim]，
        # 最后仍需 o_proj 回到 hidden_size。
        self.o_proj = RowParallelLinear(
            self.num_heads * self.v_head_dim,
            self.hidden_size,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.o_proj",
        )

cache entry dim = kv_lora_rank + qk_rope_head_dim
q head dim      = qk_nope_head_dim + qk_rope_head_dim
value head dim  = v_head_dim

kv_lora_rank = 512
qk_rope_head_dim = 64
cache entry dim = 576

@PluggableLayer.register("multi_head_latent_attention")
class MultiHeadLatentAttentionWrapper(PluggableLayer):
    def forward(self, positions, hidden_states, llama_4_scaling=None):
        q_c = None
        kv_lora = None

        if self.q_lora_rank is not None:
            # 一次 fused projection 同时得到 Q 的低秩表示 q_c，
            # 以及 KV cache 所需的 kv_lora。
            qkv_lora = self.fused_qkv_a_proj(hidden_states)[0]

            # q_c:      [tokens, q_lora_rank]
            # kv_lora:  [tokens, kv_lora_rank + qk_rope_head_dim]
            q_c, kv_lora = qkv_lora.split(
                [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
                dim=-1,
            )

            # Q 低秩表示归一化后上投影为完整 Q。
            q_c = self.q_a_layernorm(q_c)
            q = self.q_b_proj(q_c)[0]
        else:
            # 非 q_lora 路径：Q 和 KV latent 分开算。
            kv_lora = self.kv_a_proj_with_mqa(hidden_states)[0]
            q = self.q_proj(hidden_states)[0]

        # kv_c 是 compressed KV latent。
        # k_pe 是 decoupled rope key。
        kv_c, k_pe = kv_lora.split([self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)

        # 重要：写入 KV cache 的不是原始 kv_c，而是 RMSNorm 后的 kv_c_normed。
        kv_c_normed = self.kv_a_layernorm(kv_c)

        # Q reshape 成 [tokens, local_heads, qk_head_dim]。
        q = q.view(-1, self.num_heads, self.qk_head_dim)

        # k_pe 是 MQA 形态：所有 head 共享一份 rope key，所以 head 维是 1。
        k_pe = k_pe.unsqueeze(1)

        if self.rotary_emb is not None:
            # 只对 Q 的 rope 部分和 K 的 rope 部分做 RoPE。
            # Q 的 nope 部分不参与 RoPE。
            q[..., self.qk_nope_head_dim :], k_pe = self.rotary_emb(
                positions, q[..., self.qk_nope_head_dim :], k_pe
            )

        # sparse MLA 时，额外算 top-k indices，供 sparse backend 选择历史 token。
        if self.indexer and self.is_sparse:
            _topk_indices = self.indexer(
                hidden_states, q_c, positions, self.indexer_rope_emb
            )

        # 进入通用 MLA attention 层。
        attn_out = self.mla_attn(
            q,
            kv_c_normed,
            k_pe,
            output_shape=(hidden_states.shape[0], self.num_heads * self.v_head_dim),
        )

        # attention 输出仍需 o_proj。
        return self.o_proj(attn_out)[0]

class MLAAttention(nn.Module, AttentionLayerBase):
    def __init__(...):
        ...
        self.num_heads = num_heads
        self.scale = scale
        self.qk_nope_head_dim = qk_nope_head_dim
        self.qk_rope_head_dim = qk_rope_head_dim
        self.v_head_dim = v_head_dim
        self.q_lora_rank = q_lora_rank
        self.kv_lora_rank = kv_lora_rank
        self.kv_b_proj = kv_b_proj

        # MLA cache 每个 token 只有一个向量：
        # [kv_lora_rank, qk_rope_head_dim]
        # 所以 head_size 不是 qk_head_dim，也不是 v_head_dim。
        self.head_size = kv_lora_rank + qk_rope_head_dim

        # MLA 的 KV head 固定为 1，因为 latent KV 是 MQA 风格共享的。
        self.num_kv_heads = 1
        self.qk_head_dim = self.qk_nope_head_dim + self.qk_rope_head_dim

        # use_mla=True 是 backend 选择的关键。
        # use_sparse=True 时会选择 sparse MLA backend。
        self.attn_backend = get_attn_backend(
            self.head_size,
            dtype,
            kv_cache_dtype,
            use_mla=True,
            use_sparse=use_sparse,
            num_heads=self.num_heads,
        )

        # backend impl 接收 MLA 专用维度和 kv_b_proj。
        # kv_b_proj 后续用于 prefill 展开 K/V，也用于权重后处理拆 W_UK/W_UV。
        self.impl = impl_cls(
            num_heads=self.num_heads,
            head_size=self.head_size,
            scale=self.scale,
            num_kv_heads=1,
            kv_cache_dtype=self.kv_cache_dtype,
            attn_type=AttentionType.DECODER,
            q_lora_rank=self.q_lora_rank,
            kv_lora_rank=self.kv_lora_rank,
            qk_nope_head_dim=self.qk_nope_head_dim,
            qk_rope_head_dim=self.qk_rope_head_dim,
            qk_head_dim=self.qk_nope_head_dim + self.qk_rope_head_dim,
            v_head_dim=self.v_head_dim,
            kv_b_proj=kv_b_proj,
            indexer=indexer,
        )

class MLAAttention(nn.Module, AttentionLayerBase):
    def forward(self, q, kv_c_normed, k_pe, output_shape=None):
        if self.calculate_kv_scales:
            # FP8 / quantized KV cache 场景可动态计算 q/k/v scale。
            torch.ops.vllm.maybe_calc_kv_scales(
                q,
                kv_c_normed,
                k_pe,
                _encode_layer_name(self.layer_name),
            )

        if self.use_direct_call:
            forward_context = get_forward_context()
            attn_metadata = forward_context.attn_metadata

            # 多 layer metadata 以 dict 按 layer_name 存储。
            if isinstance(attn_metadata, dict):
                attn_metadata = attn_metadata[self.layer_name]

            self_kv_cache = self.kv_cache
            slot_mapping = forward_context.slot_mapping

            # slot_mapping 由 scheduler/KV manager 生成，决定当前 token 写到哪个 cache slot。
            # 区分不同的实现（例如Triton。FlashAttn, ...)
            self.impl.do_kv_cache_update(
                kv_c_normed,
                k_pe,
                self_kv_cache,
                slot_mapping.get(self.layer_name),
                self.kv_cache_dtype,
                self._k_scale,
            )

            # cache 已更新，随后执行 attention。
            output = torch.empty(output_shape, dtype=q.dtype, device=q.device)
            self.forward_impl(
                q,
                kv_c_normed,
                k_pe,
                self_kv_cache,
                attn_metadata,
                output=output,
            )
            return output

def do_kv_cache_update(
    self,
    kv_c_normed,
    k_pe,
    kv_cache,
    slot_mapping,
    kv_cache_dtype,
    k_scale,
):
    if kv_cache.numel() == 0:
        return

    # concat_and_cache_mla 会把两段拼起来：
    #   kv_c_normed: [tokens, kv_lora_rank]
    #   k_pe:        [tokens, qk_rope_head_dim]
    # 写入：
    #   kv_cache[slot] = [kv_c_normed, k_pe]
    ops.concat_and_cache_mla(
        kv_c_normed,
        k_pe.squeeze(1),
        kv_cache,
        slot_mapping.flatten(),
        kv_cache_dtype=kv_cache_dtype,
        scale=k_scale,
    )

fp8_attention = is_quantized_kv_cache(self.kv_cache_dtype)
num_actual_toks = attn_metadata.num_actual_tokens

# CUDA graph 会 padding 输入输出，这里只截取真实 token。
output = output[:num_actual_toks, ...]
q = q[:num_actual_toks, ...]
k_c_normed = k_c_normed[:num_actual_toks, ...]
k_pe = k_pe[:num_actual_toks, ...]

# sparse MLA backend 只支持 MQA/decode-style attention。
is_sparse_impl = isinstance(self.impl, SparseMLAAttentionImpl)

if is_sparse_impl:
    # sparse 情况下所有 token 都走 MQA。
    num_mqa_tokens = q.size(0)
    num_mha_tokens = 0
else:
    # 普通 MLA 由 metadata 告诉当前 batch 中多少 token 是 decode。
    num_mqa_tokens = attn_metadata.num_decode_tokens
    num_mha_tokens = q.size(0) - num_mqa_tokens

if num_mha_tokens > 0:
    # prefill 部分：compute-friendly MHA。
    self.impl.forward_mha(
        q[num_mqa_tokens:],
        k_c_normed[num_mqa_tokens:],
        k_pe[num_mqa_tokens:],
        kv_cache,
        attn_metadata,
        self._k_scale,
        output=output[num_mqa_tokens:],
    )

if num_mqa_tokens > 0:
    # decode 部分：data-movement-friendly MQA。
    mqa_q = q[:num_mqa_tokens]
    mqa_output_slice = output[:num_mqa_tokens]

    # q_nope: [B, heads, qk_nope_head_dim]
    # q_pe:   [B, heads, qk_rope_head_dim]
    mqa_q_nope, mqa_q_pe = mqa_q.split(
        [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1
    )

    # 将 q_nope 投到 latent KV 空间：
    # [heads, B, P] x [heads, P, Lkv] -> [heads, B, Lkv]
    mqa_q_nope = mqa_q_nope.transpose(0, 1)
    torch.bmm(mqa_q_nope, self.W_UK_T, out=mqa_ql_nope)
    mqa_ql_nope = mqa_ql_nope.transpose(0, 1)

    # backend 的 MQA 输入是 latent query + rope query。
    mqa_q = (mqa_ql_nope, mqa_q_pe)

    # backend 从 paged latent KV cache 中读取历史 KV。
    attn_out, lse = self.impl.forward_mqa(mqa_q, kv_cache, attn_metadata, self)

    # MQA 输出仍在 latent value 空间，需要通过 W_UV 投回 v_head_dim。
    self._v_up_proj(attn_out, out=mqa_output_slice)

def forward_mha(self, q, kv_c_normed, k_pe, kv_cache, attn_metadata, k_scale, output):
    prefill_metadata = attn_metadata.prefill

    # 当前 prefill chunk 的新 token，需要从 latent KV 展开出 k_nope 和 v。
    kv_nope = self.kv_b_proj(kv_c_normed)[0].view(
        -1,
        self.num_heads,
        self.qk_nope_head_dim + self.v_head_dim,
    )
    k_nope, v = kv_nope.split([self.qk_nope_head_dim, self.v_head_dim], dim=-1)

    # 完整 K = [k_nope, k_pe]。
    k = self._concat_k_nope_k_pe(k_nope, k_pe)

    # 当前新 token 之间做 causal attention。
    output_prefill = self._run_prefill_new_tokens(
        prefill=prefill_metadata,
        q=q,
        k=k,
        v=v,
        return_softmax_lse=has_context,
    )

    if has_context:
        # chunked prefill 有历史 context。
        # 历史 context 已在 KV cache 中，不能直接用当前 kv_c_normed。
        context_output, context_lse = self._compute_prefill_context(
            q,
            kv_cache,
            attn_metadata,
            k_scale,
        )

        # 当前 suffix 和历史 context 是两次 attention，需要按 LSE 稳定合并。
        merge_attn_states(
            output=output.view(-1, self.num_heads, self.v_head_dim),
            prefix_output=context_output,
            prefix_lse=context_lse,
            suffix_output=suffix_output,
            suffix_lse=suffix_lse,
            prefill_tokens_with_context=(
                prefill_metadata.chunked_context.prefill_tokens_with_context
            ),
        )
    else:
        # 没有历史 context，直接写输出。
        output_prefill = output_prefill[..., : v.shape[-1]].flatten(start_dim=-2)
        output.copy_(output_prefill)

def forward_mqa(self, q, kv_c_and_k_pe_cache, attn_metadata, layer):
    assert kv_c_and_k_pe_cache.numel() > 0
    assert attn_metadata.decode is not None

    # 通用层可能传入 tuple: (ql_nope, q_pe)。
    # Triton backend 这里拼成完整 latent-space query。
    if type(q) is tuple:
        q = torch.cat(q, dim=-1)

    B = q.shape[0]
    q_num_heads = q.shape[1]

    # decode attention 的输出不是 v_head_dim，而是 kv_lora_rank。
    # 后续由通用层 _v_up_proj 投到 v_head_dim。
    o = torch.zeros(
        B,
        q_num_heads,
        self.kv_lora_rank,
        dtype=q.dtype,
        device=q.device,
    )

    # 根据 max_seq_len 和 SM 数估计把 KV 分成几段并行处理。
    ideal_splits = max(1, attn_metadata.max_seq_len // min_work_per_split)
    ideal_splits = triton.next_power_of_2(ideal_splits)
    num_kv_splits = min(ideal_splits, max_splits)

    # cache 原始 shape 是 [num_blocks, block_size, Lkv + R]。
    # 这里增加一个 head 维，方便复用 decode kernel 形态。
    kv_c_and_k_pe_cache = kv_c_and_k_pe_cache.unsqueeze(2)

    # latent value 使用 kv_c 部分。
    kv_c_cache = kv_c_and_k_pe_cache[..., : self.kv_lora_rank]

    decode_attention_fwd(
        q,
        kv_c_and_k_pe_cache,
        kv_c_cache,
        o,
        lse,
        attn_metadata.decode.block_table,
        attn_metadata.decode.seq_lens,
        attn_logits,
        num_kv_splits,
        self.scale,
        PAGE_SIZE,
        k_scale=layer._k_scale,
        v_scale=layer._k_scale,
        is_mla=True,
    )

    return o, lse

def get_kv_cache_spec(self, vllm_config):
    kv_cache_dtype = kv_cache_dtype_str_to_dtype(
        self.kv_cache_dtype,
        vllm_config.model_config,
    )

    return MLAAttentionSpec(
        block_size=vllm_config.cache_config.block_size,

        # MLA latent KV 是 MQA 风格，所以 KV head 固定为 1。
        num_kv_heads=1,

        # 每个 token 存 [kv_c_normed, k_pe]。
        head_size=self.head_size,

        dtype=kv_cache_dtype,
        cache_dtype_str=vllm_config.cache_config.cache_dtype,
    )

@dataclass(frozen=True, kw_only=True)
class MLAAttentionSpec(FullAttentionSpec):
    cache_dtype_str: str | None = None

    @property
    def real_page_size_bytes(self) -> int:
        if self.cache_dtype_str == "fp8_ds_mla":
            # DeepSeek/FlashMLA sparse 的特殊格式：
            # 每 token 固定 656 bytes。
            return self.block_size * 656

        # 普通 MLA 页大小：
        # block_size * 1 * (kv_lora_rank + qk_rope_head_dim) * dtype_size。
        # 注意这里没有传统 attention 的 K/V 两份乘 2。
        return (
            self.block_size
            * self.num_kv_heads
            * self.head_size
            * get_dtype_size(self.dtype)
        )

# FullAttentionSpec.real_page_size_bytes
return (
    2
    * self.block_size
    * self.num_kv_heads
    * self.head_size
    * get_dtype_size(self.dtype)
)

def get_kv_cache_shape(
    num_blocks,
    block_size,
    num_kv_heads,  # MLA 假设为 1。
    head_size,
    cache_dtype_str="auto",
):
    # MLA cache 逻辑 shape。
    # 不是 [2, num_blocks, block_size, heads, head_dim]。
    return (num_blocks, block_size, head_size)

kv_cache.shape = [num_blocks, block_size, 576]
576 = 512 kv_lora_rank + 64 qk_rope_head_dim

if cache_dtype_str == "fp8_ds_mla":
    # DeepSeek 特化 packed cache 格式。
    return (num_blocks, block_size, 656)
else:
    return (num_blocks, block_size, head_size)