DeepseekV4Attention

DeepseekV4Attention
DeepseekV4MultiHeadLatentAttentionWrapper

需要kv_cache的类

DeepseekV4IndexerCache

class DeepseekV4Indexer(nn.Module):
    def __init__():
        assert cache_config is not None, "Deepseek V4 indexer requires cache_config"
        # NOTE(yifan): FP8 indxer cache use the same layout as V3.2:
        # head_dim bytes = 128 fp8 + 4 fp32 scale = 132.
        # For FP4 indexer cache, we still allocate the same amount of memory as FP8,
        # but only use the first half of the memory.
        k_cache_head_dim = self.head_dim + self.head_dim // self.quant_block_size * 4
        self.k_cache = DeepseekV4IndexerCache(
            head_dim=k_cache_head_dim,
            dtype=torch.uint8,
            prefix=f"{prefix}.k_cache",
            cache_config=cache_config,
            compress_ratio=self.compress_ratio,
        )
        self.compressor = DeepseekCompressor(
            vllm_config=vllm_config,
            compress_ratio=self.compress_ratio,
            hidden_size=hidden_size,
            head_dim=self.head_dim,
            rotate=True,
            prefix=f"{prefix}.compressor",
            k_cache_prefix=self.k_cache.prefix,
            use_fp4_cache=self.use_fp4_kv,
        )

        self.indexer_op = SparseAttnIndexer(
            self.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,
            skip_k_cache_insert=True,
            use_fp4_cache=self.use_fp4_kv,
        )

DeepseekV4IndexerCache.kv_cache
在DeepseekV4Indexer中，compressor的压缩的kv entry用于SparseAttnIndexer计算topk索引。 压缩的kv entry存储到DeepseekV4IndexerCache.k_cache。

DeepseekV4SWACache

class DeepseekV4SWACache(torch.nn.Module, AttentionLayerBase):
    def __init__(
        self,
        head_dim: int,
        window_size: int,
        dtype: torch.dtype,
        prefix: str,
        cache_config: CacheConfig,
    ):
        super().__init__()
        self.kv_cache = torch.tensor([])
        self.head_dim = head_dim
        self.window_size = window_size
        self.prefix = prefix
        self.cache_config = cache_config
        self.dtype = dtype
        compilation_config = get_current_vllm_config().compilation_config
        if prefix in compilation_config.static_forward_context:
            raise ValueError(f"Duplicate layer name: {prefix}")
        compilation_config.static_forward_context[prefix] = self

DeepseekV4SWACache.kv_cache

DeepseekV4MLAAttention

class DeepseekV4MLAAttention(nn.Module, AttentionLayerBase):
    def __init__():
        super().__init__()
        self.impl_cls = _select_v4_sparse_impl()
        self.backend_cls = self.impl_cls.backend_cls
        self.num_heads = num_heads
        self.num_kv_heads = 1
        self.head_dim = head_dim
        self.scale = scale
        
        self.kv_cache_dtype = kv_cache_dtype

        # Register with compilation context for metadata lookup
        compilation_config = vllm_config.compilation_config
        if prefix and prefix in compilation_config.static_forward_context:
            raise ValueError(f"Duplicate layer name: {prefix}")
        if prefix:
            compilation_config.static_forward_context[prefix] = self

        self.kv_cache = torch.tensor([])

DeepseekV4MLAAttention.kv_cache

在CSA (Compressed Sparse Attention)和HCA (Heavily Compressed Attention)场景，DeepseekCompressor压缩后的kv entry存储到 DeepseekV4MLAAttention.kv_cache。

CompressorStateCache

class CompressorStateCache(torch.nn.Module, AttentionLayerBase):
    def __init__(
        self,
        state_dim: int,
        dtype: torch.dtype,
        compress_ratio: int,
        prefix: str,
    ):
        super().__init__()
        self.state_dim = state_dim
        self.dtype = dtype
        self.prefix = prefix
        self.kv_cache = torch.tensor([])
        compilation_config = get_current_vllm_config().compilation_config
        if prefix in compilation_config.static_forward_context:
            raise ValueError(f"Duplicate layer name: {prefix}")
        compilation_config.static_forward_context[prefix] = self

CompressorStateCache.kv_cache

kv cache 的tensor分配分析

ref: vllm分析(四)——kv cache的初始化

DeepseekCompressor

DeepseekCompressor：负载kv块的压缩，用于Heavily Compressed Attention (HSA)，Compressed Sparse Attention (CSA) 和 DeepseekV4Indexer。

class DeepseekCompressor(nn.Module):
    """DeepSeek V4 KV/score compressor.

    Owns the linear / norm / state-cache / ape state and the shared forward
    prologue (kv/score split, save_partial_states launch). The
    compress → norm → RoPE → store step is dispatched to a triton kernel
    (``compress_norm_rope_store_triton``) by default, except for the NVIDIA
    head_dim=128 indexer path which uses the cutedsl kernel
    (``compress_norm_rope_store_cutedsl``) for better performance.
    """

    def __init__(
        self,
        vllm_config: VllmConfig,
        compress_ratio: int,
        hidden_size: int,
        head_dim: int,
        rotate: bool = False,
        prefix: str = "",
        k_cache_prefix="",
        use_fp4_cache: bool = False,
    ):
        self.overlap = compress_ratio == 4
        self.coff = 1 + self.overlap

        state_dtype = torch.float32
        self.ape = nn.Parameter(
            torch.empty(
                (compress_ratio, self.coff * self.head_dim),
                dtype=state_dtype,
                device=self.device,
            ),
            requires_grad=False,
        )

        self.fused_wkv_wgate = MergedColumnParallelLinear(
            self.hidden_size,
            [self.coff * self.head_dim, self.coff * self.head_dim],
            bias=False,
            return_bias=False,
            quant_config=None,
            disable_tp=True,
            prefix=f"{prefix}.fused_wkv_wgate",
        )
        self.norm = RMSNorm(self.head_dim, self.rms_norm_eps)

        self.state_cache = CompressorStateCache(
            state_dim=2 * self.coff * self.head_dim,  # kv_state + score_state
            dtype=state_dtype,
            compress_ratio=compress_ratio,
            prefix=f"{prefix}.state_cache",
        )

    def forward(
        self,
        # [num_tokens, 2 * self.coff * self.head_dim]
        kv_score: torch.Tensor,
        # [num_tokens]
        positions: torch.Tensor,
        rotary_emb,
    ) -> None:
        # Each of shape [num_tokens, coff * self.head_dim]
        # input bf16, output are fp32
        kv, score = kv_score.split(
            [self.coff * self.head_dim, self.coff * self.head_dim], dim=-1
        )

        # Get the metadata and handle dummy profiling run.
        attn_metadata = get_forward_context().attn_metadata
        if not isinstance(attn_metadata, dict):
            return

        state_metadata = cast(
            CompressorMetadata, attn_metadata[self.state_cache.prefix]
        )
        token_to_req_indices = state_metadata.token_to_req_indices
        slot_mapping = state_metadata.slot_mapping
        num_actual = slot_mapping.shape[0]
        block_table = state_metadata.block_table
        block_size = state_metadata.block_size

        # [num_blocks, block_size, kv_dim+score_dim], where kv_dim == score_dim
        state_cache = self.state_cache.kv_cache
        # kv_state stored in first half, score_state stored in second half
        state_width = state_cache.shape[-1] // 2
        pdl_kwargs = (
            {}
            if current_platform.is_rocm() or current_platform.is_xpu()
            else {"launch_pdl": False}
        )

        # Store the KV and score (with fused APE addition) in the state.
        # NOTE: PDL is disabled — both this kernel and the compress kernels
        # below depend on preceding kernel outputs (kv/score from the cublas
        # GEMM; state_cache from this kernel) but neither emits/waits on PDL
        # grid dependency primitives, so launch_pdl=True caused a
        # read-after-write race and non-deterministic output.
        save_partial_states(
            kv=kv,
            score=score,
            ape=self.ape,
            positions=positions,
            state_cache=state_cache,
            slot_mapping=slot_mapping,
            block_size=block_size,
            state_width=state_width,
            compress_ratio=self.compress_ratio,
            pdl_kwargs=pdl_kwargs,
        )

        # Fused: compress → RMSNorm → RoPE → FP8 quant → KV cache write.
        # RoPE requirements (kernel applies forward GPT-J style rotation):
        # - is_neox_style=False (interleaved pairs, NOT split-half)
        # - cos_sin_cache layout: [max_pos, rope_head_dim] with first half cos,
        #   second half sin (per-pair, length rope_head_dim // 2 each)
        # - applied to LAST rope_head_dim elements of head_dim
        # - position used: (positions // compress_ratio) * compress_ratio
        cos_sin_cache = rotary_emb.cos_sin_cache
        k_cache_metadata = cast(Any, attn_metadata[self.k_cache_prefix])
        kv_cache = self._static_forward_context[self.k_cache_prefix].kv_cache

        if current_platform.is_cuda():
            # NVIDIA GPUs.
            if self.head_dim == 512:
                from .nvidia.ops import compress_norm_rope_store_cutedsl

                # Main compressor path.
                # Use a cutedsl kernel for better performance.
                compress_norm_rope_store_fn = compress_norm_rope_store_cutedsl
            else:
                # Indexer path (head_dim == 128).
                # Use a triton kernel.
                compress_norm_rope_store_fn = compress_norm_rope_store_triton
        else:
            # AMD GPUs.
            # Always use a triton kernel.
            compress_norm_rope_store_fn = compress_norm_rope_store_triton

        compress_norm_rope_store_fn(
            state_cache=state_cache,
            num_actual=num_actual,
            token_to_req_indices=token_to_req_indices,
            positions=positions,
            slot_mapping=slot_mapping,
            block_table=block_table,
            block_size=block_size,
            state_width=state_width,
            cos_sin_cache=cos_sin_cache,
            kv_cache=kv_cache,
            k_cache_metadata=k_cache_metadata,
            pdl_kwargs=pdl_kwargs,
            head_dim=self.head_dim,
            rope_head_dim=self.rope_head_dim,
            compress_ratio=self.compress_ratio,
            overlap=self.overlap,
            use_fp4_cache=self.use_fp4_cache,
            rms_norm_weight=self.norm.weight,
            rms_norm_eps=self.rms_norm_eps,
            quant_block=self._quant_block,
            token_stride=self._token_stride,
            scale_dim=self._scale_dim,
        )

CompressorStateCache
CompressorStateCache的kv_cache空间用到了滑动窗口机制。
SlidingWindowMLASpec到SlidingWindowManager。
滑动窗口滚动时：SlidingWindowManager 会持续计算当前窗口的有效范围。当部分 block 因超出窗口范围被判定为“不再需要”时，便会释放对应的块，remove_skipped_blocks。
save_partial_states
compress_norm_rope_store_cutedsl
compress_norm_rope_store_triton

CSA 层 Compressor的工作原理

kv entry计算过程动态图，c4a attention illustation
head_dim =512。

$$C_a=H\cdot W_a^{KV},C_b=H\cdot W_b^{KV}$$
$$Z_a=H\cdot W_a^Z,Z_b=H\cdot W_b^Z$$
$$C_{\text{Comp}}={\text{Softmax}}_{\text{row}}{\left(\left[\begin{array}{c}{Z}_a+B_a\\ Z_b+B_b\end{array}\right]\right)}^{\top }\odot \left[\begin{array}{c}{C}_a\\ C_b\end{array}\right]$$

wkv包含权重$W_a^{KV}$和$W_b^{KV}$, wgate包含权重$W_a^Z$和$W_b^Z$。

fuzed_wkv_wgagte融合了：$W_a^{KV}$, $W_b^{KV}$, $W_a^Z$, $W_b^Z$。

代码和图片中的ape融合了公式中$B_a$和$B_b$。
在上面的图片中，对于CSA，fuzed_wkv_wgagte，经过split操作后，一个token产生2个 kv block，分别代表$C_a$, $C_b$，2个score block。针对一个token，数据的排列方式为：kv block0，kv block1，score block0，score block1。为了方便展示，step4有两行方块。step5中展示的压缩计算过程，图片将kv state 和 score state 措置摆放。

save_partial_states 将新生成的kv, score存入state_cache.kv_cache 。save_partial_states 在 compress_ratio=4, coff=2 时的状态存储：

┌─────────────────────────────────────┬─────────────────────────────────────┐
│            kv_state                 │          score_state                │
│         (STATE_WIDTH 个元素)         │         (STATE_WIDTH 个元素)         │
├─────────────────┬───────────────────┼─────────────────┬───────────────────┤
│   block0        │   block1          │   block0        │   block1          │
│ (head_dim 个)    │ (head_dim 个)      │ (head_dim 个)    │ (head_dim 个)      │
└─────────────────┴───────────────────┴─────────────────┴───────────────────┘
offset: 0       head_dim            STATE_WIDTH   STATE_WIDTH+head_dim

假设 token 位置：0,1,2,3,4,5,6,7,8,9,10,11,…
压缩边界在 3,7,11,…（即 position+1 是 4 的倍数）。
在压缩边界（(position+1) % 4 == 0）时，会触发一次压缩，其窗口包含 (1+overlap)*compress_ratio = 8 个 token。

压缩边界	窗口位置（实际位置）	使用的块（每个 token 贡献的块）
pos=3	[-4,-3,-2,-1,0,1,2,3]	token -4…-1: 块0（旧组） token 0…3: 块1（新组）
pos=7	[0,1,2,3,4,5,6,7]	token 0…3: 块0（旧组）token 4…7: 块1（新组）

针对csa，N卡调用 compress_norm_rope_store_cutedsl。amd gpu上 compress_norm_rope_store_triton调用 _fused_kv_compress_norm_rope_insert_sparse_attn
计算过程：

[state_cache] ──┐
[score]        ──┼─ softmax + weighted sum → compressed KV
[block_table]  ──┘
                        ↓
                  [RMSNorm]
                        ↓
            ┌───────────┴───────────┐
            ↓                       ↓
        nope (448)               rope (64)
            ↓                       ↓
      FP8 量化块            RoPE (GPT-J style)
            ↓                       ↓
  [FP8 data] [scale]          [bf16 data]
            └───────────┬───────────┘
                        ↓
                 [k_cache]  (逐 token 布局)

输出维度：

Cache block layout:
[0, bs*576):       token data (448 fp8 + 128 bf16 each)
[bs*576, +bs*8):   uint8 UE8M0 scales (7 real + 1 pad each)

前 448 字节：uint8 类型的 FP8 (E4M3) 数据（nope 部分）
后 128 字节：bfloat16 类型的 RoPE 数据（rope 部分）
紧接着 kv_cache_block_size * 8 字节：每个 token 的 scale 因子（7 个有效 uint8 + 1 个填充），用于反量化 FP8。

HCA Compressor的工作流程示意

博客[2]的流程图，说明了HCA的压缩过程。
HCA将128个kv条目压缩为1个kv entry（compress_ratio=128）。
head_dim =512。

kv entry计算过程不需要重叠。
$$C_a=H\cdot W_a^{KV}$$
$$Z_a=H\cdot W_a^Z$$
$$C_{\text{Comp}}={\text{Softmax}}_{\text{row}}(Z_a+B_a)\odot C_a$$

DeepseekV4Indexer

DeepseekV4Indexer

qr ──→ wq_b ──→ Q ──→ fused_indexer_q_rope_quant ──→ (q_quant, weights)
                              ↑                               │
                         indexer_weights                      │
                              positions                        │
compressed_kv_score ──→ compressor ──→ k ─────────────────────┤
                              │                               │
                              └──→ self.k_cache (写入)        │
                                                               ▼
hidden_states ─────────────────────────────────────→ SparseAttnIndexer
                                                               │
                                                               ▼
                                                          输出张量

fused_indexer_q_rope_quant 和 DeepseekCompressor并行运算。maybe_execute_in_parallel使用不同的stream。
SparseAttnIndexer 对每个 query 执行 top‑k 选择（topk_tokens），将选出的索引写入 topk_indices_buffer（供外部使用）。
DeepseekV4Indexer计算示意，图片来源。

token level Compressor 计算过程：CSA 层 Compressor的工作原理。
SparseAttnIndexer计算示意，图片来源

SparseAttnIndexer用于为每个 query token 选出最相关的 key-value tokens。
对于第 t个 query token与历史上的每个 token 的$k_s^I$计算相关性得分：
$$I_{t,s}=\sum _{j=1}^{H_I}{w}_{t,j}^I\cdot \text{ReLU}\left({\mathbf{q}}_{t,j}^I\cdot {\mathbf{k}}_s^I\right)$$

• $H_I$：indexer 头数（固定为 64）
• $q_{t,j}^I$: 第t个 Token 在第j个索引头（Indexer Head）的 Query 向量。
• $k_s^I$: 第 s个 Token 的 Key 向量。该 Key 向量只有一个，被所有 64 个索引头共享。图片中的indexer wk后的箭头，方框为虚线形式。
• $w_{t,j}^I$：表示第 j个头的重要性。
• $ReLU$：激活函数，具备高吞吐量的计算优势。

Attention计算过程

DeepseekV4MultiHeadLatentAttentionWrapper.forward
deepseek_v4_attention
DeepseekV4MultiHeadLatentAttentionWrapper.attention_impl

[输入] hidden_states, positions
         │
         ▼
╔════════════════════════════════════════════════════════════╗
║ 阶段1：并行 GEMM 投影 (attn_gemm_parallel_execute)         ║
╠════════════════════════════════════════════════════════════╣
║  default stream: fused_wqa_wkv → qr_kv (主投影)           ║
║          ↓                                                ║
║  aux_streams (最多3个，可选并行):                         ║
║    - compressor_kv_score (如果 compress_ratio > 1)        ║
║    - indexer_weights_proj (如果 indexer exists)           ║
║    - indexer_compressor_kv_score (如果 indexer exists)    ║
║                                                           ║
║  同步: start_event 广播 → aux等待 → done_events等待       ║
╚════════════════════════════════════════════════════════════╝
         │
         ▼
   qr_kv, kv_score, indexer_kv_score, indexer_weights
         │
         ▼
╔════════════════════════════════════════════════════════════╗
║ 阶段2：RMSNorm 归一化 (fused_q_kv_rmsnorm)                 ║
╠════════════════════════════════════════════════════════════╣
║  qr, kv = split(qr_kv, [q_lora_rank, head_dim])           ║
║  qr ← RMSNorm(qr, q_norm.weight)                          ║
║  kv ← RMSNorm(kv, kv_norm.weight)                         ║
╚════════════════════════════════════════════════════════════╝
         │
         ▼
╔════════════════════════════════════════════════════════════╗
║ 阶段3：Q/KV 变换 + 缓存写入 (带多流分支)                    ║
╠════════════════════════════════════════════════════════════╣
║ ┌─────────────────────────────────────────────────────────┐
║ │ 分支 A: 存在 indexer (压缩器必然存在)                    │
║ │   default: wq_b_kv_insert                               │
║ │     → wq_b(qr) → [n_heads, head_dim]                   │
║ │     → (融合kernel):        │
║ │         - Q: per-head RMSNorm + RoPE + 填充至padded_heads│
║ │         - KV: RoPE + FP8量化 + 写入SWA缓存              │
║ │   aux0: indexer.forward (含其内部wq_b+量化+稀疏索引)     │
║ │   aux1: compressor.forward (压缩kv_score写索引器K缓存)   │
║ │   同步: default等待aux0,aux1完成 → 返回 q_padded         │
║ ├─────────────────────────────────────────────────────────┤
║ │ 分支 B: 仅存在压缩器 (无 indexer)                        │
║ │   default: wq_b_kv_insert (同上)                        │
║ │   aux: compressor.forward                               │
║ │   同步: maybe_execute_in_parallel                       │
║ ├─────────────────────────────────────────────────────────┤
║ │ 分支 C: 无压缩器 & 无 indexer (纯SWA)                   │
║ │   顺序执行 wq_b_kv_insert                               │
║ └─────────────────────────────────────────────────────────┘
╚════════════════════════════════════════════════════════════╝
         │
         ▼
   q_padded ( [num_tokens, padded_heads, head_dim] )
   kv (原始, 实际读缓存)
         │
         ▼
╔════════════════════════════════════════════════════════════╗
║ 阶段4：稀疏注意力计算 (mla_attn)                           ║
╠════════════════════════════════════════════════════════════╣
║  backend: FlashMLASparseBackend (NVIDIA) / ROCm AITER     ║
║  输入: q_padded, kv, positions                            ║
║  缓存: SWA缓存(FP8) + 索引器提供的稀疏索引(若存在)         ║
║  输出: out [num_tokens, padded_heads, head_dim] (预分配)   ║
║  注意: 仅使用SWA缓存中的KV，忽略输入的kv tensor            ║
╚════════════════════════════════════════════════════════════╝
         │
         ▼
[输出] out (后续经过逆RoPE + FP8 einsum + wo_b 得到最终结果)

_select_v4_sparse_impl

def _select_v4_sparse_impl() -> "type[DeepseekV4SparseMLAAttentionImpl]":
    """Pick the platform-specific V4 sparse MLA impl class. Sole platform check."""
    if current_platform.is_rocm():
        from vllm.models.deepseek_v4.amd.rocm import (
            DeepseekV4ROCMAiterMLASparseImpl,
        )

        return DeepseekV4ROCMAiterMLASparseImpl
    from vllm.models.deepseek_v4.nvidia.flashmla import (
        DeepseekV4FlashMLASparseImpl,
    )

    return DeepseekV4FlashMLASparseImpl

结合vllm代码，画出csa计算过程示意图，图片参考了:图片源。

wqa和wkv的权重融合为 fused_wqa_wkv
q rmsnorm和 kv rmsnorm 融合为 fused_q_kv_rmsnorm。
w_qb的投影计算，self.wq_b(qr).view(-1, self.n_local_heads, self.head_dim)
qnorm rope 和 kv rope quant insert 融合为 _fused_qnorm_rope_kv_insert。
inv RoPE + FP8 einsum + wo_b 对应的代码位置

forward_mqa

DeepseekV4FlashMLAAttention.forward_mqa

forward_mqa
├─ 若是预热 dummy run → 分配 workspace，return
├─ 获取 swa_metadata / flashmla_metadata
├─ 根据 compress_ratio 判断是否 swa_only
├─ 若存在 prefill
│   └─ _forward_prefill
│       ├─ 分块
│       ├─ 收集压缩 KV（非 swa_only）和 SWA KV
│       ├─ 合并 topk 与 SWA 索引
│       └─ 调用 flash_mla_sparse_fwd
└─ 若存在 decode
    └─ _forward_decode
        ├─ 准备 topk 索引（实时计算或预取）
        ├─ 获取对应 tile_sched_metadata
        └─ 调用 flash_mla_with_kvcache（同时传入 SWA 和全局缓存）

_forward_prefill 执行流程

_forward_prefill
│
├─ 1. 解析元数据
│     - swa_only = (attn_metadata is None)
│     - num_prefills, num_prefill_tokens, num_decodes
│     - seq_lens, gather_lens
│     - query_start_loc_cpu, query_start_loc
│     - prefill_token_base = query_start_loc_cpu[num_decodes]
│
├─ 2. 确定压缩池大小与 topk 索引
│     ├─ 若非 swa_only:
│     │    - 若 compress_ratio == 4 → topk_indices 从 self.topk_indices_buffer 切片
│     │    - 若 compress_ratio == 128 → topk_indices = attn_metadata.c128a_prefill_topk_indices
│     │    - top_k = topk_indices.shape[-1]
│     │    - N = ceil(max_model_len / compress_ratio)
│     └─ 若 swa_only:
│          - topk_indices 占位（不使用），top_k = 0，N = 0
│
├─ 3. 分配临时 KV 缓冲区
│     - M = N + window_size + max_num_batched_tokens
│     - chunk_size_const = PREFILL_CHUNK_SIZE
│     - num_chunks = ceil(num_prefills / chunk_size_const)
│     - kv = workspace_manager.get_simultaneous((chunk_size_const, M, head_dim), bf16)[0]
│
├─ 4. 分块循环（chunk_idx = 0 到 num_chunks-1）
│     │
│     ├─ 4.1 确定当前 chunk 范围
│     │      chunk_start, chunk_end, chunk_size
│     │
│     ├─ 4.2 若非 swa_only：
│     │      dequantize_and_gather_k_cache
│     │      - 从 compressed_k_cache 解量化并收集到 kv[:chunk_size, 0:N]
│     │      - 使用压缩后的 block_table（偏移 num_decodes）
│     │
│     ├─ 4.3 收集 SWA KV
│     │      dequantize_and_gather_k_cache
│     │      - 从 swa_k_cache 收集到 kv[:chunk_size, N : N+max_gather_lens]
│     │      - 使用 swa_metadata.block_table（偏移 num_decodes）
│     │
│     ├─ 4.4 计算当前 chunk 在 q/output 中的 token 范围
│     │      query_start = query_start_loc_cpu[num_decodes+chunk_start] - prefill_token_base
│     │      query_end   = query_start_loc_cpu[num_decodes+chunk_end]   - prefill_token_base
│     │
│     ├─ 4.5 合并索引
│     │      combined_indices, combined_lens = combine_topk_swa_indices(
│     │          topk_indices[query_start:query_end],
│     │          query_start_loc[chunk_start:chunk_end+1],
│     │          seq_lens[chunk_start:chunk_end],
│     │          gather_lens[chunk_start:chunk_end],
│     │          window_size, compress_ratio, top_k, M, N
│     │      )
│     │
│     └─ 4.6 执行稀疏注意力
│            flash_mla_sparse_fwd(
│                q = q[query_start:query_end],
│                kv = kv.view(-1, 1, head_dim),
│                indices = combined_indices.unsqueeze(1),
│                topk_length = combined_lens,
│                out = output[query_start:query_end]
│            )
│
└─ 5. 结束（所有 chunk 处理完毕）

第一次调用dequantize_and_gather_k_cache, offset=0, kv[batch_idx, offset + i, :] = kv[batch_idx, i, :]，其中 i 从 0 到 gather_len-1（压缩后的 token 索引）。
第一次调用dequantize_and_gather_k_cache, offset=N，写入位置：kv[batch_idx, N + i, :]，其中 i 从 0 到 gather_len-1（原始 token 索引）。
_forward_decode 执行流程

_forward_decode
│
├─ 1. 获取 decode 相关元数据
│     - num_decodes, num_decode_tokens
│     - swa_indices, swa_lens
│
├─ 2. 若非 SWA-only 模式，准备全局 topk 索引
│     ├─ 根据 compress_ratio 分支：
│     │   - ratio=4: 调用 compute_global_topk_indices_and_lens 实时计算
│     │   - ratio=128: 直接使用预计算的 c128a_global_decode_topk_indices
│     └─ 得到 topk_indices, topk_lens
│
├─ 3. 准备缓存张量（添加一个维度以匹配内核期望）
│     - swa_cache = swa_cache_layer.kv_cache.unsqueeze(-2)
│     - 若有全局缓存，kv_cache = kv_cache.unsqueeze(-2)
│
├─ 4. 根据 compress_ratio 选择对应的 tile_scheduler_metadata
│     - ≤1  → tile_sched_swaonly
│     - =4  → tile_sched_c4a
│     - =128→ tile_sched_c128a
│
└─ 5. 调用 flash_mla_with_kvcache 内核
       - 传入 SWA 缓存和（可选）全局压缩缓存
       - 分别传入 SWA 索引和全局 topk 索引
       - 使用 tile scheduler 自动分块调度
       - 输出直接写入 output

flash_mla_with_kvcache 的 block_table=None。全局压缩缓存的索引已经通过 compute_global_topk_indices_and_lens 映射为物理 slot ID。

reference

[1] DeepSeek V4 in vLLM: Efficient Long-context Attention
[2] 手撕 DeepSeek-V4 (3): HCA
[3] DeepSeek v4 Compressor kv cache压缩模块
[4] DeepSeek V4-vLLM预览
[5] 图解DeepSeek V4：详细计算流程解析
[6] Deepseek-V4模型结构与源码解析
[7] Deepseek FlashMLA解析
[8] flashMLA中sparse_attn_fwd_kernel分析
[9] FlashMLA 深度解析：FP8 KV Cache 与 DSA 稀疏注意力实现原理
[10] 深度解析FlashMLA: 一文读懂大模型加速新利器