HBM 索引架构设计

vLLM v1 中 Scheduler(CPU) ↔ Worker(GPU) 的索引契约：为什么 HBM 索引发到 worker、为什么有两种 SlotMappingMode、DeepSeek V4 的三通道索引设计。

1 · 术语澄清

2 · 为什么要把 HBM 索引发到 worker 模块

一句话：Scheduler 只识逻辑块号，Attention kernel 只识物理 HBM 地址；两者之间的翻译必须在 worker 里完成。

2.1 Scheduler 做什么

Scheduler 通过 KVCacheManager.allocate_slots() 从 free-list 拿出逻辑块号，装进 SchedulerOutput.new_block_ids：

# vllm/v1/core/kv_cache_manager.py:283
def allocate_slots(self, request, num_new_tokens, ...) -> KVCacheBlocks | None:
    # 从 free-list 拿出若干逻辑块号 [block_3, block_7, ...]
    # 返回给 Scheduler → 装进 SchedulerOutput.new_block_ids

这是职责分离的核心：调度策略 和 设备内存布局 是两个独立变化维度。

2.2 Worker 做什么

Worker 收到 block_ids 后，在 GPU HBM 上构建 kernel 直接消费的索引：

# vllm/v1/worker/gpu_model_runner.py:1444-1445
# 1. 收到逻辑 block_ids，append 到 block_table 的 CPU 镜像
self.input_batch.block_table.append_row(new_block_ids, req_index)

# gpu_model_runner.py:1950
# 2. commit：CPU → GPU 异步拷贝（HBM 索引就位）
self.input_batch.block_table.commit_block_table(num_reqs)

# gpu_model_runner.py:2168
# 3. compute_slot_mapping：用 Triton kernel 把 block_table 翻译成
#    每个 token 的物理 slot 索引（HBM 地址的"偏移编号"）
self.input_batch.block_table.compute_slot_mapping(
    num_reqs, query_start_loc, positions)

2.3 物理公式：虚实地址转换

_compute_slot_mapping_kernel（Triton）执行 OS 页表同构的翻译：

# vllm/v1/worker/block_table.py:348-411
@triton.jit
def _compute_slot_mapping_kernel(...):
    # 对每个 token：
    virtual_block_indices = pos // virtual_block_size
    block_indices = virtual_block_indices * BLOCKS_PER_KV_BLOCK
                       + local_block_offsets // block_size
    block_numbers = block_table[req_idx, block_indices]   # 查页表
    slot_offsets = local_block_offsets % block_size
    slot_ids = block_numbers * block_size + slot_offsets  # 物理 slot

2.4 为什么不能在 Scheduler 算好再传

2.5 Scheduler ↔ Worker 契约图

Scheduler (CPU)                              Worker (GPU)
─────────────────                            ─────────────
block_ids: [3, 7, 12]  ──────append────→  block_table.cpu[row] = [3,7,12]
                                         │
                                         ↓ commit_block_table (H2D)
                                         block_table.gpu  ← 这就是"HBM 索引"
                                         │
                                         ↓ compute_slot_mapping (Triton)
                                         slot_mapping.gpu ← kernel 直接消费

3 · 两种 SlotMappingMode 的语义与分发

3.1 定义

# vllm/v1/worker/block_table.py:19-21
class SlotMappingMode(Enum):
    TOKEN_TO_KV_SLOT = "token_to_kv_slot"   # 标准：每 token → 1 slot
    NONE = "none"                            # Mamba/GDN：不用 slot_mapping

3.2 分发逻辑

# vllm/v1/worker/gpu_model_runner.py:7122-7125
if kv_cache_spec_kind == KVCacheSpecKind.MAMBA:
    slot_mapping_modes.append(SlotMappingMode.NONE)   # 循环状态，不写 paged KV
else:
    slot_mapping_modes.append(SlotMappingMode.TOKEN_TO_KV_SLOT)

3.3 为什么要有 NONE

# vllm/v1/worker/block_table.py:153-163
def compute_slot_mapping(self, num_reqs, query_start_loc, positions):
    if self.slot_mapping_mode == SlotMappingMode.NONE:
        # Mamba/GDN groups consume the block table as recurrent state
        # indices and do not use per-token slot mappings.
        return

4 · HBM 索引 vs KV cache host：两层抽象的分裂

这是两个层级的抽象，变化速率不同：

4.1 具体差异

# vllm/v1/worker/block_table.py:58-78
if kernel_block_size == block_size:
    # Standard case: direct mapping
    self.block_size = block_size
    self.blocks_per_kv_block = 1
else:
    # Hybrid case: memory blocks subdivided to match kernel
    self.block_size = kernel_block_size
    self.blocks_per_kv_block = block_size // kernel_block_size

4.2 为什么不让 host 直接产出 slot_mapping

因为 slot_mapping 依赖：

• positions（每个 token 在序列中的绝对位置）—— Scheduler 传下来的
• query_start_loc（batch 内每个 req 的起止）—— worker 在 GPU 上算的
• block_size、compress_ratio、cp_kv_cache_interleave_size 等设备相关参数

5 · DeepSeek V4 的三通道索引设计

DeepSeek V4 是 vLLM 里索引路径最复杂的模型，因为它有三条独立的 KV 索引通道，各自有 block_table / slot_mapping：

DeepSeek V4 的 3 条 KV 索引通道
─────────────────────────────────
① Main MLA cache      — 压缩 KV（compress_ratio=4 或 128），block_size=256
② Indexer cache       — Lightning Indexer 的稀疏 top-k 选择器，独立 block_table
③ SWA cache           — Sliding Window Attention，block_size=64

5.1 Main MLA Cache 的索引

# vllm/models/deepseek_v4/attention.py:605-624
def get_kv_cache_spec(self, vllm_config) -> KVCacheSpec | None:
    if self.compress_ratio <= 1:  # SWA-only 层
        return None               # SWA 走③，不走这里
    return MLAAttentionSpec(
        block_size=...,           # 256
        num_kv_heads=1,
        head_size=self.head_dim,  # 512
        compress_ratio=self.compress_ratio,  # 4 或 128
        alignment=576,            # fp8_ds_mla 专用对齐
    )

关键：storage_block_size = block_size // compress_ratio（kv_cache_interface.py:402-403）：

• compress_ratio=4 时：256-token 的 host block 只存 64 个 KV（每 4 token 压 1 个）
• compress_ratio=128 时：256-token 的 host block 只存 2 个 KV

5.2 压缩 slot_mapping 怎么算

只有满足 (pos+1) % compress_ratio == 0 的 token 才有合法 slot：

# vllm/v1/attention/backends/mla/compressor_utils.py:9-50
@triton.jit
def _compressed_slot_mapping_kernel(...):
    for i in range(0, query_len, TRITON_BLOCK_SIZE):
        pos = start_pos + i + arange(0, TRITON_BLOCK_SIZE)
        is_valid = (pos + 1) % COMPRESS_RATIO == 0
        pos_after_compress = pos // COMPRESS_RATIO

        block_ids = pos_after_compress // block_size
        block_numbers = block_table[req_idx, block_ids]       # 标准 block_table 查表
        slot_ids = block_numbers * block_size + pos_after_compress % block_size
        slot_ids = tl.where(is_valid, slot_ids, PAD_ID)       # 非压缩位填 -1

5.3 Indexer Cache 的索引

Indexer 是 DeepSeek V4 的稀疏注意力选择器：每个 query token 通过它挑出 top-k 个 KV token。

# vllm/models/deepseek_v4/attention.py:668-764
class DeepseekV4Indexer:
    def __init__(self, ...):
        # 独立的 KV cache（128 fp8 + 4 fp32 scale = 132B per token head）
        self.k_cache = DeepseekV4IndexerCache(
            head_dim=head_dim + head_dim//quant_block_size*4,
            compress_ratio=compress_ratio,
        )
        # 独立的 compressor（把 hidden_states 投成 KV score）
        self.compressor = DeepseekCompressor(...)
        # indexer op 内部有自己的 block_table 查表逻辑
        self.indexer_op = SparseAttnIndexer(
            self.k_cache, ..., skip_k_cache_insert=True)

Indexer 的 slot_mapping 也是独立计算的：

# vllm/v1/attention/backends/mla/indexer.py:579-591
compressed_slot_mapping = slot_mapping
if self.compress_ratio > 1:
    compressed_slot_mapping = get_compressed_slot_mapping(
        num_tokens, query_start_loc, seq_lens, block_table,
        self.kv_cache_spec.storage_block_size,
        self.compress_ratio, out=self.compressed_slot_mapping_buffer)

5.4 SWA Cache 的索引

SWA 存最近 window_size（如 2048）个 token 的全量 KV，用于滑动窗口内精确注意力。

# vllm/v1/attention/backends/mla/sparse_swa.py:51-102
class DeepseekV4SWACache:
    def __init__(self, ...):
        # 强制 block_size = 64（因为 C4A 和 SWA 共享物理 tensor）
        self.block_size = 64

    def get_kv_cache_spec(self, vllm_config):
        return SlidingWindowMLASpec(
            block_size=64,  # 与 main MLA 的 256 不同！
            sliding_window=self.window_size,
        )

SWA 的索引特殊之处：不只在 build 时算 slot_mapping，还要算 decode_swa_indices 和 prefill_swa_indices——每个 token 在窗口内的所有历史 slot 列表：

# vllm/v1/attention/backends/mla/sparse_swa.py:660-716
@triton.jit
def _compute_swa_indices_and_lens_kernel(...):
    # 标准 block_table 查表，但"一对多"：一个 query token → window_size 个 KV slot
    for i in range(0, window_size, TRITON_BLOCK_SIZE):
        pos_offset = start_pos + offset
        block_indices = pos_offset // block_size
        block_numbers = block_table[req_idx, block_indices]
        block_offsets = pos_offset % block_size
        slot_ids = block_numbers * block_size + block_offsets

5.5 三条通道如何在 worker 里协同

# vllm/v1/worker/gpu_model_runner.py:2294-2538
# 1. 通过 kv_cache_gid 索引到不同的 BlockTable 实例
blk_table = self.input_batch.block_table[kv_cache_gid]

# 2. 每个 gid 产出独立的 block_table_tensor + slot_mapping
for kv_cache_gid, kv_cache_group in enumerate(kv_cache_groups):
    if kv_cache_gid > 0:
        cm.block_table_tensor = _get_block_table(kv_cache_gid)
        cm.slot_mapping = slot_mappings[kv_cache_gid]

# 3. 每个 attention backend 的 builder.build() 消费各自的 cm

最终一个 DeepSeek V4 forward 里：

• Main MLA builder（DeepseekV4FlashMLAMetadataBuilder）：用 compressed_slot_mapping
• Indexer builder（DeepseekV32IndexerMetadataBuilder）：用 indexer 自己的 compressed_slot_mapping
• SWA builder（DeepseekSparseSWAMetadataBuilder）：用标准 slot_mapping + SWA indices

5.6 设计过程的"第一性原理"

需求：在 16K-128K 上下文下做全量注意力太贵
  │
  ↓
观察：大部分 KV 对当前 query 无用（稀疏性）
  │
  ├─→ 方案① 压缩（compress_ratio=4/128）—— 只存 1/4 或 1/128 的 KV
  │         → 需要 compressed_slot_mapping（非压缩位不写）
  │
  ├─→ 方案② Lightning Indexer —— 每个 query 先快速打分挑 top-k
  │         → 需要独立 indexer cache + 独立 block_table
  │
  └─→ 方案③ Sliding Window —— 近端 token 全部保留（生成质量关键）
            → 需要独立 SWA cache + 独立 block_size=64

三条通道的 block_table 共享逻辑 block_ids（同一个 SchedulerOutput），
但各自产出独立的 slot_mapping（因为 block_size / compress_ratio 不同）

6 · 第一性原理总结

7 · 关键源文件索引

分析基于 vLLM 2db39c704 · 2026-07-21 · 代码路径: vllm_research/vllm/