IndexCache：DeepSeek 稀疏注意力中的跨层索引复用

class DeepseekV32ForCausalLM(VerifyAndUpdateConfig):
    @classmethod
    def verify_and_update_config(cls, vllm_config: "VllmConfig") -> None:
        hf_config = vllm_config.model_config.hf_config

        # index_topk 是当前实现识别 DeepSeek V3.2 DSA 的关键配置字段。
        is_v32 = hasattr(hf_config, "index_topk")
        assert is_v32

        cache_config = vllm_config.cache_config
        if cache_config.cache_dtype == "bfloat16":
            # 将显式 BF16 改成 auto，后续由 sparse MLA 后端确定实际布局。
            cache_config.cache_dtype = "auto"
            logger.info("Using bfloat16 kv-cache for DeepSeekV3.2")

class DeepseekV32IndexerCache(torch.nn.Module, AttentionLayerBase):
    def __init__(self, head_dim, dtype, prefix, cache_config):
        super().__init__()

        # 初始化时只是占位；真实分页 cache 由 vLLM KV cache 管理器分配并绑定。
        self.kv_cache = torch.tensor([])
        self.head_dim = head_dim
        self.prefix = prefix
        self.cache_config = cache_config
        self.dtype = dtype

        # 注册为 attention layer，使调度器能够为它建立 cache 和 metadata。
        compilation_config = get_current_vllm_config().compilation_config
        compilation_config.static_forward_context[prefix] = self

    def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
        return MLAAttentionSpec(
            block_size=self.cache_config.block_size,
            num_kv_heads=1,      # 64 个 Indexer Query heads 共享一份历史 K。
            head_size=self.head_dim,
            dtype=self.dtype,
        )

    def get_attn_backend(self) -> AttentionBackend:
        # 后端负责物理布局、slot mapping、block table 和 prefill/decode metadata。
        return DeepseekV32IndexerBackend

class DeepseekV32IndexerBackend(AttentionBackend):
    @staticmethod
    def get_name() -> str:
        return "DEEPSEEK_V32_INDEXER"

    @staticmethod
    def get_kv_cache_shape(
        num_blocks,
        block_size,
        num_kv_heads,
        head_size,
        cache_dtype_str="auto",
    ):
        assert num_kv_heads == 1

        # 每个分页位置只保存一份共享 K，不额外增加 V 维度。
        return (num_blocks, block_size, head_size)

class Indexer(nn.Module):
    def __init__(self, vllm_config, config, hidden_size, q_lora_rank,
                 quant_config, cache_config, topk_indices_buffer, prefix,
                 is_inplace_rope=False):
        super().__init__()

        self.topk_tokens = config.index_topk
        self.n_head = config.index_n_heads       # 通常为 64 个 Indexer Query heads。
        self.head_dim = config.index_head_dim    # 通常为 128。
        self.rope_dim = config.qk_rope_head_dim  # RoPE 子空间通常为 64。

        # 从主 MLA 的 query latent 生成多头 Indexer Query。
        self.wq_b = ReplicatedLinear(
            q_lora_rank,
            self.head_dim * self.n_head,
            bias=False,
            quant_config=quant_config,
            prefix=f"{prefix}.wq_b",
        )

        # 一次 GEMM 同时生成共享 K 和各 Query head 的聚合权重。
        # Fused wk + weights_proj: single GEMM producing [head_dim + n_head].
        self.wk_weights_proj = MergedColumnParallelLinear(
            hidden_size,
            [self.head_dim, self.n_head],
            bias=False,
            quant_config=None,
            disable_tp=True,
            prefix=f"{prefix}.wk_weights_proj",
        )

        self.k_norm = LayerNorm(self.head_dim, eps=1e-6)
        self.topk_indices_buffer = topk_indices_buffer

# FP8 naive cache：128 bytes K + 4 bytes FP32 scale。
self.k_cache = DeepseekV32IndexerCache(
    head_dim=self.head_dim + self.head_dim // self.quant_block_size * 4,
    dtype=torch.uint8,
    prefix=f"{prefix}.k_cache",
    cache_config=cache_config,
)

self.indexer_op = SparseAttnIndexer(
    self.k_cache,               # 历史检索 K 的分页 cache。
    self.quant_block_size,
    self.scale_fmt,
    self.topk_tokens,
    self.head_dim,
    self.max_model_len,
    self.max_total_seq_len,
    self.topk_indices_buffer,   # 跨层共享的是这个 top-k 结果 buffer。
)

def forward(self, hidden_states, qr, positions, rotary_emb) -> torch.Tensor:
    # qr 是主 MLA 产生的 query latent q_c。
    q, _ = self.wq_b(qr)
    q = q.view(-1, self.n_head, self.head_dim)

    q_pe, q_nope = torch.split(
        q,
        [self.rope_dim, self.head_dim - self.rope_dim],
        dim=-1,
    )

    # 一次融合投影得到共享 K 和每个 Query head 的权重。
    kw, _ = self.wk_weights_proj(hidden_states)
    k = kw[:, : self.head_dim]
    weights = kw[:, self.head_dim :]
    k = self.k_norm(k)

    k_pe, k_nope = torch.split(
        k,
        [self.rope_dim, self.head_dim - self.rope_dim],
        dim=-1,
    )

    # 只旋转 PE 子空间；K 增加一个共享 KV-head 维度以匹配 MQA。
    q_pe, k_pe = rotary_emb(positions, q_pe, k_pe.unsqueeze(1))
    q = torch.cat([q_pe, q_nope], dim=-1)
    k = torch.cat([k_pe.squeeze(-2), k_nope], dim=-1)

    # Q 在 Python 层量化；K 的量化与 paged-cache 写入在底层融合完成。
    q_fp8, q_scale = per_token_group_quant_fp8(
        q.view(-1, self.head_dim),
        self.quant_block_size,
        column_major_scales=False,
        use_ue8m0=True,
    )
    q_fp8 = q_fp8.view(-1, self.n_head, self.head_dim)
    q_scale = q_scale.view(-1, self.n_head, 1)

    # 将 Q scale、点积缩放和多头归一化折叠进 weights。
    weights = (
        weights.unsqueeze(-1)
        * q_scale
        * self.softmax_scale
        * self.n_head**-0.5
    ).squeeze(-1)

    return self.indexer_op(hidden_states, q_fp8, k, weights)

return self.indexer_op(
    hidden_states,
    q_fp8,     # 当前 token 的多头量化检索 Query。
    k,         # 当前 token 的共享检索 K，待量化并写入分页 cache。
    weights,   # 聚合多个 Query-head 分数的动态权重。
)

@CustomOp.register("sparse_attn_indexer")
class SparseAttnIndexer(CustomOp):
    def forward_cuda(self, hidden_states, q_quant, k, weights):
        # V3.2 FP8 路径中 q_quant 是 Tensor，q_scale 已折叠进 weights。
        q_values, q_scale = self._split_quantized_q(q_quant)

        return torch.ops.vllm.sparse_attn_indexer(
            hidden_states,
            self.k_cache_prefix,
            self.kv_cache,
            q_values,
            q_scale,
            k,
            weights,
            self.quant_block_size,
            self.scale_fmt,
            self.topk_tokens,
            self.head_dim,
            self.max_model_len,
            self.total_seq_lens,
            self.topk_indices_buffer,
            self.skip_k_cache_insert,
            self.use_fp4_cache,
        )

attn_metadata = get_forward_context().attn_metadata
metadata = attn_metadata[k_cache_prefix]
slot_mapping = metadata.slot_mapping

if not skip_k_cache_insert:
    # V3.2 在这里量化当前 K 并写入本层 Indexer paged cache。
    ops.indexer_k_quant_and_cache(
        k,
        kv_cache,
        slot_mapping,
        quant_block_size,
        scale_fmt,
    )

# -1 用于标记当前 batch 尚未产生或无效的 top-k 位置。
topk_indices_buffer[: hidden_states.shape[0]] = -1

@dataclass
class DeepseekV32IndexerMetadata:
    seq_lens: torch.Tensor
    max_seq_len: int
    slot_mapping: torch.Tensor

    num_decodes: int
    num_decode_tokens: int
    num_prefills: int
    num_prefill_tokens: int

    decode: DeepSeekV32IndexerDecodeMetadata | None = None
    prefill: DeepseekV32IndexerPrefillMetadata | None = None

index_topk_freq = getattr(config, "index_topk_freq", 1)
index_topk_pattern = getattr(config, "index_topk_pattern", None)
index_skip_topk_offset = getattr(config, "index_skip_topk_offset", 2)

if index_topk_pattern is None:
    skip_topk = (
        max(layer_id - index_skip_topk_offset + 1, 0) % index_topk_freq != 0
    )
elif 0 <= layer_id < len(index_topk_pattern):
    skip_topk = index_topk_pattern[layer_id] == "S"

if self.is_v32 and (not skip_topk or is_mtp_layer):
    self.indexer = Indexer(...)
else:
    # Shared backbone 层没有自己的 Indexer，也没有自己的 Indexer K cache。
    self.indexer = None

self.mla_attn = MultiHeadLatentAttentionWrapper(
    ...,
    indexer=self.indexer,
    topk_indices_buffer=topk_indices_buffer,
    skip_topk=skip_topk,
)

self.is_v32 = hasattr(config, "index_topk")

if self.is_v32:
    topk_indices_buffer = torch.empty(
        vllm_config.scheduler_config.max_num_batched_tokens,
        config.index_topk,
        dtype=torch.int32,
        device=self.device,
    )
else:
    topk_indices_buffer = None

# 同一个 tensor 被传给所有 decoder layers。
self.layers = make_layers(
    config.num_hidden_layers,
    lambda prefix: DeepseekV2DecoderLayer(
        vllm_config=vllm_config,
        prefix=prefix,
        topk_indices_buffer=topk_indices_buffer,
    ),
)

# Full 层：运行 Indexer，更新本层 K cache，并刷新共享 top-k buffer。
if self.indexer and self.is_sparse and not self.skip_topk:
    self.indexer(hidden_states, q_c, positions, self.indexer_rope_emb)

# Full 和 Shared 层都会运行 sparse MLA。
# Shared 层没有刷新 buffer，因此会消费最近前置 Full 层留下的 top-k。
attn_out = self.mla_attn(
    q,
    kv_c_normed,
    k_pe,
    output_shape=(hidden_states.shape[0], self.num_heads * self.v_head_dim),
)