上篇
deepseek v4的kv cache类型
python
# https://github.com/vllm-project/vllm/blob/v0.27.1/vllm/models/deepseek_v4/attention.py#L655
class DeepseekV4Attention(nn.Module, AttentionLayerBase, ABC):
def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec | None:
if (
self.compress_ratio <= 1
): # SWA part. Allocated separately as DeepseekV4SWACache.
return None
# fp8_ds_mla is a UE8M0 block-scaled uint8 layout and needs 576B
# alignment; plain bf16 / per-tensor fp8 rows use natural element-size
# pages.
uses_fp8_ds_mla_layout = self.kv_cache_dtype == "fp8_ds_mla"
return MLAAttentionSpec(
block_size=vllm_config.cache_config.block_size,
num_kv_heads=1,
head_size=self.head_dim,
dtype=torch.uint8 if uses_fp8_ds_mla_layout else self.kv_cache_torch_dtype,
compress_ratio=self.compress_ratio,
cache_dtype_str=self.kv_cache_dtype,
alignment=576 if uses_fp8_ds_mla_layout else 512,
model_version="deepseek_v4",
kv_quant_mode=get_kv_quant_mode(self.kv_cache_dtype),
)
# https://github.com/vllm-project/vllm/blob/v0.27.1/vllm/models/deepseek_v4/attention.py#L698
class DeepseekV4IndexerCache(torch.nn.Module, AttentionLayerBase):
def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
# head_dim already carries the fp8 scale padding
# compress_ratio=1 for V3.2, >1 for DeepseekV4; both use the same cache layout.
uses_fp8_ds_mla_layout = vllm_config.cache_config.cache_dtype == "fp8_ds_mla"
return MLAAttentionSpec(
block_size=self.cache_config.block_size,
num_kv_heads=1,
head_size=self.head_dim,
dtype=self.dtype,
compress_ratio=self.compress_ratio,
# 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577).
alignment=576 if uses_fp8_ds_mla_layout else 512,
)
# https://github.com/vllm-project/vllm/blob/v0.27.1/vllm/v1/attention/backends/mla/sparse_swa.py#L87
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
# Block size is constrained by tensor sharing between SWA and C4A KV blocks.
# Since both block types share the same physical tensor, they must use the
# same page size. The C4A KV block shape [256//4, head_dim] = [64, head_dim]
# determines the SWA block size of 64 tokens per block.
# TODO(yifan): make SWA block size automatically determined and configurable.
self.block_size = 64
def get_kv_cache_spec(self, vllm_config: VllmConfig) -> KVCacheSpec:
# fp8_ds_mla's UE8M0 paged layout needs 576B alignment; contiguous
# bf16/fp8 cache uses the natural element-size page.
uses_fp8_ds_mla_layout = self.cache_config.cache_dtype == "fp8_ds_mla"
return SlidingWindowMLASpec(
block_size=self.block_size,
num_kv_heads=1,
head_size=self.head_dim,
dtype=self.dtype,
sliding_window=self.window_size,
cache_dtype_str=self.cache_config.cache_dtype,
# 576B for FlashMLA packing; 512B for FlashInfer sparse (#44577).
alignment=576 if uses_fp8_ds_mla_layout else 512,
model_version="deepseek_v4",
kv_quant_mode=get_kv_quant_mode(self.cache_config.cache_dtype),
)
python
class DeepseekV4SparseMLABackend(AttentionBackend):
"""DeepSeek-V4 sparse-MLA backend base.
Subclasses ``AttentionBackend`` directly (not the V3.2
``FlashMLASparseBackend``): DeepSeek-V4 runs its own attention layer
(``DeepseekV4Attention``), so it does not reuse the V3.2 builder or impl, and
only needs to declare its own metadata builder, KV-cache layout, and the
sparse-MLA capability flags.
"""
supported_dtypes: ClassVar[list[torch.dtype]] = [torch.bfloat16]
supported_kv_cache_dtypes: ClassVar[list[CacheDType]] = [
"auto",
"fp8_ds_mla",
"fp8", # alias for fp8_ds_mla
]
@staticmethod
def get_supported_kernel_block_sizes() -> list[int | MultipleOf]:
return [256]
@classmethod
def get_supported_head_sizes(cls) -> list[int]:
# DeepSeek V4 layout: 448 NoPE + 64 RoPE = 512.
return [512]
@staticmethod
def get_kv_cache_shape(
num_blocks: int,
block_size: int,
num_kv_heads: int,
head_size: int,
cache_dtype_str: str = "auto",
) -> tuple[int, ...]:
if cache_dtype_str == "fp8_ds_mla":
# DeepseekV4 main MLA: 584B per token (448 NoPE + 128 RoPE + 8 fp8 scale).
# head_size passed in is the semantic head_dim (512).
return (num_blocks, block_size, 584)
else:
return (num_blocks, block_size, head_size)page_size_bytes 的计算
page_size_bytes实际上是按照alignment对齐后的字节数。后续分析先忽略这一点。
MLAAttentionSpec.real_page_size_bytes
python
def _apply_alignment_padding(spec: MLAAttentionSpec | SlidingWindowMLASpec):
if spec.alignment is None:
return
actual_page_size = spec.real_page_size_bytes
padded_page_size = round_up(actual_page_size, spec.alignment)
if padded_page_size != actual_page_size:
object.__setattr__(spec, "page_size_padded", padded_page_size)
class MLAAttentionSpec(FullAttentionSpec):
def __post_init__(self):
super().__post_init__()
_apply_alignment_padding(self)
@property
def real_page_size_bytes(self) -> int:
if self.cache_dtype_str == "fp8_ds_mla":
if self.model_version == "deepseek_v4":
# DeepseekV4: 448B NoPE + 128B RoPE + 8B fp8 scale = 584B per token.
# head_size stays semantic (512); bytes are determined here.
return self.storage_block_size * 584
# V3.2 main MLA: 656-byte custom layout (kv_lora_rank=512 +
# qk_rope_head_dim=64, head_size=576). See flashmla_sparse.py.
return self.block_size * 656
if self.kv_quant_mode == KVQuantMode.INT4_PER_TOKEN_HEAD:
head_dim = self.head_size // 2
else:
head_dim = self.head_size
return (
self.storage_block_size
* self.num_kv_heads
* head_dim
* get_dtype_size(self.dtype)
)
# https://github.com/vllm-project/vllm/blob/v0.27.1/vllm/v1/kv_cache_interface.py#L631
class SlidingWindowMLASpec(SlidingWindowSpec):
@property
def real_page_size_bytes(self) -> int:
if self.model_version == "deepseek_v4" and self.cache_dtype_str == "fp8_ds_mla":
# DeepseekV4 FlashMLA: 448B NoPE + 128B RoPE + 8B fp8 scale = 584B
# per token. FlashInfer's contiguous bf16/fp8 cache falls through to
# the element-size formula below.
return self.storage_block_size * 584
assert self.model_version in (None, "deepseek_v4"), (
f"Unsupported model version: {self.model_version}"
)
return (
self.storage_block_size
* self.num_kv_heads
* self.head_size
* get_dtype_size(self.dtype)
)DeepSeek V4 in vLLM: Efficient Long-context Attention
Different layers compress at different rates (1/4 for c4a, 1/128 for c128a, 1/1 for SWA). we fix the logical block at 256 native token positions for every compressed layer. A c4a block then physically holds 256 / 4 = 64 compressed entries, and a c128a block holds 256 / 128 = 2.
CSA 和HCA 逻辑 block_size = 256, storage_block_size = block_size/compress_ratio.
SWA的 block_size = 64。
| Group | compress_ratio | storage_block_size | 每token大小 | real_page_size_bytes |
|---|---|---|---|---|
| CSA | 4 | 256/4 = 64 | 584B | 64 × 584 = 37,376B |
| HCA | 128 | 256/128 = 2 | 584B | 2 × 584 = 1,168B |
| SWA | 1 | 64 | 584B | 64 × 584 = 37,376B |
num_blocks 的计算
text
get_kv_cache_configs()
├── get_kv_cache_groups()
│ └── group_and_unify_kv_cache_specs() # DeepseekV4 特殊处理
│ └── _get_kv_cache_groups_uniform_groups()
├── _project_kv_cache_groups_to_worker()
└── get_kv_cache_config_from_groups() # 循环调用
└── _use_packed_kv_cache_config() # 判断是否使用 packed 布局
└── _get_kv_cache_config_packed()
└── _get_packed_kv_cache_layout()
python
def get_kv_cache_configs(
vllm_config: VllmConfig,
kv_cache_specs: list[dict[str, KVCacheSpec]],
available_memory: list[int],
) -> list[KVCacheConfig]:get_kv_cache_config_from_groups
python
def get_kv_cache_config_from_groups(
vllm_config: VllmConfig,
kv_cache_groups: list[KVCacheGroupSpec],
available_memory: int,
) -> KVCacheConfig:
"""
Generate the KV cache configuration from the KV cache groups and spec
of each layer.
Args:
vllm_config: The global VllmConfig
kv_cache_groups: The KV cache groups
available_memory: Memory available for KV cache in bytes
Returns:
The generated KVCacheConfig
"""
if len(kv_cache_groups) == 0:
# Attention free models do not have KV cache.
# Return num_blocks=1 as BlockPool always needs a null_block.
return KVCacheConfig(
num_blocks=1,
kv_cache_tensors=[],
kv_cache_groups=kv_cache_groups,
)
# Determine how model runners should initialize the KV cache tensors.
if len(kv_cache_groups) == 1 and isinstance(
kv_cache_groups[0].kv_cache_spec, UniformTypeKVCacheSpecs
):
# Special case: all layers have the same type of KV cache but with
# different hidden sizes. Allocate different amount of memory for each
# layer based on its hidden size.
num_blocks = (
available_memory // kv_cache_groups[0].kv_cache_spec.page_size_bytes
)
num_blocks = may_override_num_blocks(vllm_config, num_blocks)
per_layer_specs = kv_cache_groups[0].kv_cache_spec.kv_cache_specs
kv_cache_tensors = [
KVCacheTensor(
size=per_layer_specs[layer_name].page_size_bytes * num_blocks,
shared_by=[layer_name],
)
for layer_name in kv_cache_groups[0].layer_names
]
elif _use_packed_kv_cache_config(vllm_config, kv_cache_groups):
# DeepSeek V4 uses the packed layout by default. Other multi-group
# layouts can opt in with --enable-cross-layers.
num_blocks, kv_cache_tensors = _get_kv_cache_config_packed(
vllm_config, kv_cache_groups, available_memory
)
else:
# General case:
# We will have group_size memory pools, each is shared by one layer from
# each group. As layers of different groups have different block table,
# they will use different parts of the shared Tensor.
# The memory layout for 3 groups (full.0, full.1), (sw.0, sw.2),
# (sw.1, padding) will be: (group_size = 2)
# full.0, sw.0, sw.1: share a Tensor with size=available_memory//2
# full.1, sw.2: share another Tensor with size=available_memory//2
group_size = max(len(group.layer_names) for group in kv_cache_groups)
page_size = get_uniform_page_size(
[group.kv_cache_spec for group in kv_cache_groups]
)
assert group_size > 0, "group_size must be greater than 0"
num_blocks = get_num_blocks(
vllm_config, group_size, available_memory, page_size
)
kv_cache_tensors = []
for i in range(group_size):
shared_by = []
for j in range(len(kv_cache_groups)):
if i < len(kv_cache_groups[j].layer_names):
shared_by.append(kv_cache_groups[j].layer_names[i])
kv_cache_tensors.append(
KVCacheTensor(size=page_size * num_blocks, shared_by=shared_by)
)
return KVCacheConfig(
num_blocks=num_blocks,
kv_cache_tensors=kv_cache_tensors,
kv_cache_groups=kv_cache_groups,
)
python
def _get_kv_cache_config_packed(
vllm_config: VllmConfig,
kv_cache_groups: list[KVCacheGroupSpec],
available_memory: int,
) -> tuple[int, list[KVCacheTensor]]:
"""Plan a packed per-block KV cache tensor layout.
Cache groups use dense, overlapping layouts within one block slab. Each
emitted tensor aliases the same physical backing allocation.
"""
block_stride, layers_by_offset = _get_packed_kv_cache_layout(kv_cache_groups)
num_blocks = available_memory // block_stride
num_blocks = may_override_num_blocks(vllm_config, num_blocks)
total_size = block_stride * num_blocks
kv_cache_tensors: list[KVCacheTensor] = []
for byte_offset in sorted(layers_by_offset):
kv_cache_tensors.append(
KVCacheTensor(
size=total_size,
shared_by=layers_by_offset[byte_offset],
offset=byte_offset,
block_stride=block_stride,
)
)
return num_blocks, kv_cache_tensors针对packed kv cache kayout，_allocate_kv_cache_tensors只会分配一块物理内存packed_backing。
block_stride的计算：_get_packed_kv_cache_layout
python
def _get_packed_kv_cache_layout(
kv_cache_groups: list[KVCacheGroupSpec],
) -> tuple[int, dict[int, list[str]]]:
"""Lay out each cache group densely in one shared block slab.
A block ID is owned by one cache group at a time, so layouts from different
groups may overlap. Layers within a group remain disjoint.
"""
layers_by_offset: dict[int, list[str]] = defaultdict(list)
block_stride = 0
for group in kv_cache_groups:
spec = group.kv_cache_spec
byte_offset = 0
for layer_name in group.layer_names:
if isinstance(spec, UniformTypeKVCacheSpecs):
page_size = spec.kv_cache_specs[layer_name].page_size_bytes
else:
page_size = spec.page_size_bytes
layers_by_offset[byte_offset].append(layer_name)
byte_offset += page_size
block_stride = max(block_stride, byte_offset)
assert block_stride > 0
return block_stride, layers_by_offset计算 block_stride
python
CSA Group: 30层 × 37,376B = 1,121,280B
HCA Group: 31层 × 1,168B = 36,208B
SWA Group: 61层 × 37,376B = 2,279,936B ← 最大
block_stride = max(1,121,280, 36,208, 2,279,936) = 2,279,936B内存分配时候，按照最大block_stride分配。一个block可以覆盖所有group的内存需求。同一时刻，一个 block 只会被一个 Group 使用。
num_blocks的计算：
num_blocks = available_memory // block_strideblock_stride = 2,279,936 字节 (约 2.28 MB)
假设 available_memory = 10 GB (即 10,737,418,240 字节)
则 num_blocks = 10,737,418,240 // 2,279,936 ≈ 4710 个块
内存布局：
text
物理内存（total_size = block_stride × num_blocks）:
Block 0 (block_stride B) Block 1 (block_stride B) ... Block N (block_stride B)
┌──────────────────────┐ ┌──────────────────────┐ ┌──────────────────────┐
│ Layer0 (offset=0) │ │ Layer0 (offset=0) │ │ Layer0 (offset=0) │
│ Layer1 (offset=100B) │ │ Layer1 (offset=100B) │ │ Layer1 (offset=100B) │
│ ... │ │ ... │ │ ... │
│ LayerN (offset=...) │ │ LayerN (offset=...) │ │ LayerN (offset=...) │
└──────────────────────┘ └──────────────────────┘ └──────────────────────┘layers_by_offset 构造过程模拟
Group 1: SWA Group (61层)
python
byte_offset = 0
for i in range(61):
layer_name = f"swa_{i}"
page_size = 37,376B
layers_by_offset[byte_offset].append(layer_name) # 0: ['swa_0']
byte_offset += 37,376遍历完成后：
text
offset 0: ['swa_0']
offset 37,376: ['swa_1']
offset 74,752: ['swa_2']
offset 112,128: ['swa_3']
...
offset : 2,242,560['swa_60'] # 60 * 37,376
byte_offset = 61 * 37,376 = 2,279,936BGroup 2: CSA Group (30层)
python
byte_offset = 0
for i in range(30):
layer_name = f"csa_{i}"
page_size = 37,376B
layers_by_offset[byte_offset].append(layer_name)
byte_offset += 37,376遍历完成后：
text
offset 0: ['swa_0', 'csa_0'] # 与swa_0重叠
offset 37,376: ['swa_1', 'csa_1']
offset 74,752: ['swa_2', 'csa_2']
offset 112,128: ['swa_3', 'csa_3']
...
offset 1,084,416: ['swa_29', 'csa_29'] # 29 * 37,376
byte_offset = 30 * 37,376 = 1,121,280BGroup 3: HCA Group (31层)
python
byte_offset = 0
for i in range(31):
layer_name = f"hca_{i}"
page_size = 1,168B
layers_by_offset[byte_offset].append(layer_name)
byte_offset += 1,168遍历完成后：
text
offset 0: ['swa_0', 'csa_0', 'hca_0'] # 三者重叠
offset 1,168: ['hca_1']
offset 2,336: ['hca_2']
offset 3,504: ['hca_3']
...
offset 35,040: ['hca_30'] # 30 * 1,168
byte_offset = 31 * 1,168 = 36,208B最终的 layers_by_offset
text
layers_by_offset = {
# === 偏移 0 处，三组的第一层完全重叠 ===
0: ['swa_0', 'csa_0', 'hca_0'],
# === 偏移 1,168 到 35,040：HCA 组各层占据的间隙 ===
1,168: ['hca_1'],
2,336: ['hca_2'],
3,504: ['hca_3'],
# ... (中间按 1,168B 步长递增)
35,040: ['hca_30'],
# === 偏移 37,376 开始，SWA 和 CSA 层按 37,376B 步长对齐 ===
37,376: ['swa_1', 'csa_1'],
74,752: ['swa_2', 'csa_2'],
112,128: ['swa_3', 'csa_3'],
# ... (中间按 37,376B 步长递增)
1,084,416: ['swa_29', 'csa_29'], # 29 * 37,376
# === 仅属于 SWA 组的剩余层 (CSA 只有30层) ===
1,121,280: ['swa_30'],
1,158,656: ['swa_31'],
1,196,032: ['swa_32'],
# ... (中间按 37,376B 步长递增)
2,242,560: ['swa_60'], # 60 * 37,376
}kv cache 内存分配
python
# https://github.com/vllm-project/vllm/blob/v0.27.1/vllm/v1/worker/gpu_model_runner.py#L7541
class GPUModelRunner(
LoRAModelRunnerMixin, KVConnectorModelRunnerMixin, ECConnectorModelRunnerMixin
):
def initialize_kv_cache_tensors(
self, kv_cache_config: KVCacheConfig, kernel_block_sizes: list[int]
) -> dict[str, torch.Tensor]:
"""
Initialize the memory buffer for KV cache.
Args:
kv_cache_config: The KV cache config
kernel_block_sizes: The kernel block sizes for each KV cache group.
Returns:
Dict[str, torch.Tensor]: A map between layer names to their
corresponding memory buffer for KV cache.
"""
# Try creating KV caches optimized for kv-connector transfers
cache_dtype = self.cache_config.cache_dtype
if self.use_uniform_kv_cache(self.attn_groups):
kv_caches, cross_layers_kv_cache, attn_backend = (
self.allocate_uniform_kv_caches(
kv_cache_config,
self.attn_groups,
cache_dtype,
self.device,
kernel_block_sizes,
)
)
self.cross_layers_kv_cache = cross_layers_kv_cache
self.cross_layers_attn_backend = attn_backend
else:
# Fallback to the general case
# Initialize the memory buffer for KV cache
kv_cache_raw_tensors = self._allocate_kv_cache_tensors(kv_cache_config)
# Change the memory buffer to the desired shape
kv_caches = self._reshape_kv_cache_tensors(
kv_cache_raw_tensors, kernel_block_sizes
)
# Set up cross-layer KV cache sharing
for layer_name, target_layer_name in self.shared_kv_cache_layers.items():
logger.debug("%s reuses KV cache of %s", layer_name, target_layer_name)
kv_caches[layer_name] = kv_caches[target_layer_name]
num_attn_module = (
2 if self.model_config.hf_config.model_type == "longcat_flash" else 1
)
bind_kv_cache(
kv_caches,
self.compilation_config.static_forward_context,
self.kv_caches,
num_attn_module,
)
return kv_caches
# https://github.com/vllm-project/vllm/blob/v0.27.1/vllm/v1/worker/gpu_model_runner.py#L7312
def _allocate_kv_cache_tensors(
self, kv_cache_config: KVCacheConfig
) -> dict[str, torch.Tensor]:
"""
Initializes the KV cache buffer with the correct size. The buffer needs
to be reshaped to the desired shape before being used by the models.
Args:
kv_cache_config: The KV cache config
Returns:
dict[str, torch.Tensor]: A map between layer names to their
corresponding memory buffer for KV cache.
"""
kv_cache_raw_tensors: dict[str, torch.Tensor] = {}
packed_backing: torch.Tensor | None = None
for kv_cache_tensor in kv_cache_config.kv_cache_tensors:
if kv_cache_tensor.block_stride > 0:
# Allocate once; all packed tensors alias the same backing.
if packed_backing is None:
packed_backing = torch.zeros(
kv_cache_tensor.size,
dtype=torch.int8,
device=self.device,
)
tensor = packed_backing
else:
tensor = torch.zeros(
kv_cache_tensor.size, dtype=torch.int8, device=self.device
)
for layer_name in kv_cache_tensor.shared_by:
kv_cache_raw_tensors[layer_name] = tensor
layer_names = set()
for group in kv_cache_config.kv_cache_groups:
for layer_name in group.layer_names:
if layer_name in self.runner_only_attn_layers:
continue
layer_names.add(layer_name)
assert layer_names == set(kv_cache_raw_tensors.keys()), (
"Some layers are not correctly initialized"
)
return kv_cache_raw_tensors
# https://github.com/vllm-project/vllm/blob/v0.27.1/vllm/v1/worker/gpu_model_runner.py#L7364
def _reshape_kv_cache_tensors(
self,
kv_cache_raw_tensors: dict[str, torch.Tensor],
kernel_block_sizes: list[int],
) -> dict[str, torch.Tensor]:
"""
Reshape the KV cache tensors to the desired shape and dtype.
Args:
kv_cache_raw_tensors: The KV cache buffer of each layer, with
correct size but uninitialized shape.
kernel_block_sizes: The kernel block sizes for each KV cache group.
Returns:
Dict[str, torch.Tensor]: A map between layer names to their
corresponding memory buffer for KV cache.
"""
kv_caches: dict[str, torch.Tensor] = {}
has_attn, has_mamba = False, False
# Map layer names to (offset, block_stride) within the packed
# backing tensor so we can create strided views per layer.
layer_packing: dict[str, tuple[int, int]] = {}
for kv_tensor in self.kv_cache_config.kv_cache_tensors:
if kv_tensor.block_stride > 0:
for ln in kv_tensor.shared_by:
layer_packing[ln] = (kv_tensor.offset, kv_tensor.block_stride)
for group in self._kv_cache_spec_attn_group_iterator():
kv_cache_spec = group.kv_cache_spec
attn_backend = group.backend
if group.kv_cache_group_id == len(kernel_block_sizes):
# There may be a last group for layers without kv cache.
continue
kernel_block_size = kernel_block_sizes[group.kv_cache_group_id]
for layer_name in group.layer_names:
if layer_name in self.runner_only_attn_layers:
continue
raw_tensor = kv_cache_raw_tensors[layer_name]
packing = layer_packing.get(layer_name)
if packing is not None:
_, blk_stride = packing
num_blocks = raw_tensor.numel() // blk_stride
else:
assert raw_tensor.numel() % kv_cache_spec.page_size_bytes == 0
num_blocks = raw_tensor.numel() // kv_cache_spec.page_size_bytes
if isinstance(kv_cache_spec, AttentionSpec):
has_attn = True
num_blocks_per_kv_block = (
kv_cache_spec.block_size // kernel_block_size
)
kernel_num_blocks = num_blocks * num_blocks_per_kv_block
# For MLA with compression, storage_block_size != block_size
if kv_cache_spec.storage_block_size != kv_cache_spec.block_size:
shape_block_size = kv_cache_spec.storage_block_size
else:
shape_block_size = kernel_block_size
# Skipped layers (--kv-cache-dtype-skip-layers) need
# the unquantized shape.
layer_cache_dtype_str = (
"auto"
if kv_cache_spec.kv_quant_mode == KVQuantMode.NONE
else getattr(
kv_cache_spec,
"cache_dtype_str",
None,
)
or self.cache_config.cache_dtype
)
kv_cache_shape = attn_backend.get_kv_cache_shape(
kernel_num_blocks,
shape_block_size,
kv_cache_spec.num_kv_heads,
kv_cache_spec.head_size,
cache_dtype_str=layer_cache_dtype_str,
)
try:
kv_cache_stride_order = attn_backend.get_kv_cache_stride_order()
assert len(kv_cache_stride_order) == len(kv_cache_shape)
except (AttributeError, NotImplementedError):
kv_cache_stride_order = tuple(range(len(kv_cache_shape)))
raw_tensor = kv_cache_raw_tensors[layer_name]
kv_caches[layer_name] = _reshape_attention_kv_cache(
raw_tensor,
kv_cache_spec,
kv_cache_shape,
kv_cache_stride_order,
kernel_num_blocks,
packing,
)
elif isinstance(kv_cache_spec, MambaSpec):
has_mamba = True
raw_tensor = kv_cache_raw_tensors[layer_name]
page_size_bytes = kv_cache_spec.page_size_bytes
# Hold a single contiguous [num_blocks, 1, 1, page_size_bytes]
# int8 page view per layer; the layer's bind_kv_cache unpacks
# each block's bytes into its conv/ssm state views. Keeping
# one tensor per layer lets the KV connector register it
# without special-casing Mamba.
kv_caches[layer_name] = raw_tensor[
: num_blocks * page_size_bytes
].view(num_blocks, 1, 1, page_size_bytes)
else:
raise NotImplementedError
# Reconcile divergent KV layouts to blocks-first. Triggered by hybrid
# attention/mamba models, and by encoder-decoder models whose shared
# decoder/cross-attention allocation mixes K/V-first and blocks-first
# backends (see _has_mixed_attention_kv_layout).
if has_attn and (
has_mamba or self._has_mixed_attention_kv_layout(kernel_block_sizes)
):
self._update_hybrid_attention_mamba_layout(kv_caches, kernel_block_sizes)
return kv_caches
# https://github.com/vllm-project/vllm/blob/v0.27.1/vllm/v1/worker/gpu/attn_utils.py#L211
def _reshape_attention_kv_cache(
kv_raw_tensor: torch.Tensor,
kv_cache_spec: AttentionSpec,
kv_cache_shape: tuple[int, ...],
kv_cache_stride_order: tuple[int, ...],
num_blocks: int,
packing: tuple[int, int] | None,
page_aligned_blocks: bool = False,
) -> torch.Tensor:
permuted_kv_cache_shape = tuple(kv_cache_shape[i] for i in kv_cache_stride_order)
inv_order = [
kv_cache_stride_order.index(i) for i in range(len(kv_cache_stride_order))
]
dtype = kv_cache_spec.dtype
if packing is not None:
offset, block_stride = packing
assert inv_order[0] == 0
page_bytes = prod(kv_cache_shape[1:]) * get_dtype_size(dtype)
kv_cache = (
kv_raw_tensor.view(-1, block_stride)[:, offset : offset + page_bytes]
.view(dtype)
.view(permuted_kv_cache_shape)
)
elif kv_cache_spec.page_size_padded is not None:
# Use a strided view to skip the padding between physical pages.
#
# Only num-blocks-first layouts are supported (the block dimension is
# dim 0 of the unpermuted shape). kv-first layouts such as ROCm's
# ``(2, num_blocks, ...)`` are intentionally not supported here. For a
# num-blocks-first layout the only stride that must change is the block
# stride: every other (contiguous) stride already steps within the
# unpadded region of a page, so no further adjustment is needed.
assert kv_cache_shape[0] == num_blocks, (
"Padded KV pages require a num-blocks-first KV cache layout (got "
f"shape {kv_cache_shape} with num_blocks={num_blocks}); "
"kv-first layouts are not supported."
)
dtype_size = get_dtype_size(kv_cache_spec.dtype)
page_stride = kv_cache_spec.page_size_bytes // dtype_size
num_blocks_dim = inv_order[0]
strides = list(torch.empty(permuted_kv_cache_shape, device="meta").stride())
strides[num_blocks_dim] = page_stride
kv_cache = torch.as_strided(
kv_raw_tensor.view(dtype),
size=permuted_kv_cache_shape,
stride=tuple(strides),
)
elif page_aligned_blocks:
# A KV-first layout such as ROCm's ``(2, num_blocks, ...)`` puts block
# ``b``'s K and V in two far-apart halves of the allocation, so block
# ``b`` does not cover page ``b``. Mamba layers sharing the allocation
# do address their state by page, so the two would resolve the same
# bytes. Build the view page-first instead, then swap the dims back.
assert kv_cache_shape[1] == num_blocks and kv_cache_stride_order == tuple(
range(len(kv_cache_shape))
), (
"Page-aligned KV blocks expect a default-strided (kv, num_blocks, "
f"...) layout, got shape {kv_cache_shape} with stride order "
f"{kv_cache_stride_order} and num_blocks={num_blocks}."
)
kv_cache = (
kv_raw_tensor.view(dtype)
.view(num_blocks, kv_cache_shape[0], *kv_cache_shape[2:])
.transpose(0, 1)
)
else:
# No padding --- safe to use a contiguous view.
kv_cache = kv_raw_tensor.view(dtype).view(permuted_kv_cache_shape)
return kv_cache.permute(*inv_order)DeepseekV4SparseMLABackend.get_kv_cache_shape
python
class DeepseekV4FlashMLABackend(AttentionBackend):
@staticmethod
def get_kv_cache_shape(
num_blocks: int,
block_size: int,
num_kv_heads: int,
head_size: int,
cache_dtype_str: str = "auto",
) -> tuple[int, ...]:
if cache_dtype_str == "fp8_ds_mla":
# DeepseekV4 main MLA: 584B per token (448 NoPE + 128 RoPE + 8 fp8 scale).
# head_size passed in is the semantic head_dim (512).
return (num_blocks, block_size, 584)
else:
return (num_blocks, block_size, head_size)针对packing情形，_allocate_kv_cache_tensors只分配一次packed_backing tensor。函数返回字典kv_cache_raw_tensors，值均为packed_backing。_reshape_kv_cache_tensors 将 KV 缓存内存块（kv_cache_raw_tensors）重塑为模型各注意力层所需具体形状，求取切片。
_reshape_kv_cache_tensors
在函数开始时，会遍历 kv_cache_config.kv_cache_tensors，为所有参与打包的层构建一个映射字典 layer_packing，记录每一层的 (offset, block_stride)。对于上面的例子：
python
layer_packing = {
'swa_0': (0, 2_279_936),
'csa_0': (0, 2_279_936),
'hca_0': (0, 2_279_936),
'swa_1': (37_376, 2_279_936),
'csa_1': (37_376, 2_279_936),
'hca_1': (1_168, 2_279_936),
# ...
}_reshape_attention_kv_cache 为每一层创建视图。
python
page_bytes = prod(kv_cache_shape[1:]) * dtype_size # 计算单页字节数
kv_cache = (raw_tensor.view(-1, block_stride) # 视图1: [总字节数/block_stride, block_stride]
[:, offset : offset + page_bytes] # 视图2: [num_blocks, page_bytes]，切片出该层的区域
.view(dtype) # 视图3: 转为目标数据类型 (如 fp8)
.view(permuted_kv_cache_shape)) # 视图4: 重塑为后端要求的形状以SWA_1层为例，shape_block_size = kv_cache_spec.storage_block_size =64。
DeepseekV4SparseMLABackend.get_kv_cache_shape 返回 (num_blocks, 64, 584)。
page_bytes = page_size_bytes = 64 x 584 = 37,376B针对SWA_1层，_reshape_attention_kv_cache返回的tensor切片(layer_view)就是下图中红色方框中的存储空间。
SWA_1层可用的存储空间（37,376 字节/块）不连续，而是被分割成 num_blocks 个碎片，分别存储在 packed_backing 的 Block 0、Block 1...Block N 中对应的 offset 处。
图片中的数字编号只是示意数据块个数，不是slot id的编号。SWA0，SWA1...SWA60对应的存储空间，slot id 各自独立编号。
kv cache的绑定
AttentionLayerBase.bind_kv_cache
python
class AttentionLayerBase(ABC):
"""
Base class for attention-like layers (Attention, Mamba, etc.)
that support the v1 engine.
This provides a common interface for getting attention backends
from different layer types.
"""
impl: "AttentionImpl"
supports_dcp: bool = True
def bind_kv_cache(self, kv_cache: torch.Tensor) -> None:
"""Bind the allocated KV cache tensor to this layer.
The default stores the cache view as-is; subclasses (e.g. Mamba)
override this to unpack the raw buffer into per-state views.
"""
self.kv_cache = kv_cache当 forward_contextlayer_name.bind_kv_cache(kv_cache) 被调用时：
- forward_contextlayer_name 获取到模型中的对应注意力层实例。
- 调用该实例的 bind_kv_cache 方法，并将 kv_cache（即 _reshape_kv_cache_tensors 返回的张量视图）作为参数传入。
- 该方法执行 self.kv_cache = kv_cache，将这个张量视图赋值给该层实例的 kv_cache 属性。
kv cache 寻址过程
通用寻址公式
标准 KV 缓存shape为 (num_blocks, block_size, num_heads, head_size)。
一个 slot_id 通常编码了块索引（block_id）和块内偏移（token_offset），即 slot_id = block_id * block_size + token_offset。对于给定的 slot_id 和需要访问的注意力头索引 head_idx（范围 0 到 num_heads-1），其 KV 数据在张量中的物理地址计算方式为：
物理地址 = kv_cache 的基地址
+ block_id * kv_cache.stride(0) # 跳到目标物理块
+ token_offset * kv_cache.stride(1) # 跳到块内目标 token 位置
+ head_idx * kv_cache.stride(2) # 跳到目标 head
+ head_offset * kv_cache.stride(3) # (可选) 跳到 head 内的具体位置- kv_cache.stride(0) = 每个物理块的总字节数（block_size * num_heads * head_size * dtype_size）。
- kv_cache.stride(1) = 每个 token 的总字节数（num_heads * head_size * dtype_size）。
- kv_cache.stride(2) = 单个头的字节数（head_size * dtype_size）。
- kv_cache.stride(3) = dtype_size（例如 bfloat16 为 2，uint8 为 1）。
head_offset：当需要访问 head_size 中的特定元素时使用，通常在注意力计算中，内核会顺序读取整个 head_size。
swa cache的寻址
DeepseekV4FlashMLAAttention._forward_decode
python
class DeepseekV4FlashMLAAttention(DeepseekV4Attention):
def _forward_decode(
self,
q: torch.Tensor,
kv_cache: torch.Tensor | None, # Only used when compress_ratio > 1
swa_metadata: "DeepseekSparseSWAMetadata",
attn_metadata: DeepseekV4FlashMLAMetadata | None,
swa_only: bool,
output: torch.Tensor,
) -> None:
# Prepare SWA cache (num_blocks, swa_block_size, 1, head_bytes)
# Use unsqueeze to preserve strides (handles padded blocks correctly)
swa_cache = self.swa_cache_layer.kv_cache.unsqueeze(-2)
# Reshape KV cache to (num_blocks, block_size, 1, head_bytes)
if kv_cache is not None:
kv_cache = kv_cache.unsqueeze(-2)self.swa_cache_layer.kv_cache 就是通过 bind_kv_cache 绑定到 SWA 层的张量视图，其形状为 (num_blocks, 64, 584)。
执行 unsqueeze(-2) 后，swa_cache 的形状变为：(num_blocks, 64, 1, 584)。
在标准的 MHA 中，KV 缓存通常有 4 个维度：(num_blocks, block_size, num_kv_heads, head_size)。对于 DeepSeek V4 的 MLA 或 SWA，num_kv_heads 为 1。添加这个维度可以让 SWA 缓存与内核期望的通用 4D 布局对齐。
pytorch 为swa_cache张量维护一个"地图"，即它的 stride（步长）属性。这个地图精确地记录了如何在物理内存中定位每个逻辑元素。例如：
- swa_cache.stride(0): 要跳到下一个 block，需要在内存中前进多少字节（block_stride）。
- swa_cache.stride(1): 在同一个 block 内，要跳到下一个 Token，需要前进多少字节（通常是 584 字节）。
- swa_cache.stride(2): num_heads 维度（大小为1）的步长。
- swa_cache.stride(3)：head_size 维度（584字节）的步长，这里通常是1。
针对swa_cache，根据slot id 定位物理地址：
slot_id：block_id = slot_id // 64, token_offset = slot_id % 64
text
物理地址 = swa_cache 的基地址
+ block_id * swa_cache.stride(0)
+ token_offset * swa_cache.stride(1)
+ 0 * swa_cache.stride(2)需要注意，SWA0，SWA1...SWA60对应的存储空间，slot id 各自独立编号。
slot id 计算
_compute_swa_indices_and_lens_kernel 根据当前批次中每个 Token 的位置，计算出它需要在 SWA 缓存中关注的 Token 的物理槽位（slot）索引和有效窗口长度。
text
def _compute_swa_indices_and_lens_kernel(
swa_indices_ptr,
swa_indices_stride,
swa_lens_ptr,
window_size,
query_start_loc_ptr,
seq_lens_ptr,
token_to_req_indices_ptr,
is_valid_token_ptr,
block_table_ptr,
block_table_stride,
block_size,
token_offset,
TRITON_BLOCK_SIZE: tl.constexpr,
):
计算过程:
1. 根据 pos 计算窗口 start_pos 和 end_pos
2. 对窗口内每个 offset:
pos_offset = start_pos + offset
block_indices = pos_offset // block_size
block_number = block_table[block_indices]
block_offset = pos_offset % block_size
slot_id = block_number * block_size + block_offset SWA 缓存：其 block_size 固定为 64，对应 DeepseekV4SWACache 的 storage_block_size。
_compute_global_topk_indices_and_lens_kernel 为全局稀疏索引（Top-K）构建 slot_id 索引。
cpp
def _compute_global_topk_indices_and_lens_kernel(
global_topk_indices_ptr,
global_topk_indices_stride,
topk_lens_ptr,
topk_indices_ptr,
topk_indices_stride,
topk,
token_to_req_indices_ptr,
block_table_ptr,
block_table_stride,
block_size,
is_valid_token_ptr,
TRITON_BLOCK_SIZE: tl.constexpr,
)
计算过程:
1. 直接从 topk_indices_buffer 读取 local_idx
2. 对每个 local_idx (>=0):
block_indices = local_idx // block_size
block_number = block_table[block_indices]
block_offset = local_idx % block_size
slot_id = block_number * block_size + block_offsetCSA/HCA 缓存：它们管理的是压缩后的 KV 缓存。传入_compute_global_topk_indices_and_lens_kernel的block_size, 代码位置：
text
block_size = attn_metadata.block_size // self.compress_ratio例如，CSA 的 compress_ratio=4，则其 block_size = 256 / 4 = 64；HCA 的 compress_ratio=128，则其 block_size = 256 / 128 = 2。
block id 的分配
每个 SingleTypeKVCacheManager 独立管理各自的 KV 缓存组（如 SWA 组、CSA 组），但它们共享同一个 BlockPool 实例，并从该实例中申请全局唯一的 block_id。
text
KVCacheManager.allocate_slots
│
├── 1. 计算需求 & 准入检查
│ ├── num_tokens_main_model = total_computed_tokens + num_new_tokens
│ ├── 如果 full_sequence_must_fit:
│ │ ├── num_blocks_to_allocate = coordinator.get_num_blocks_to_allocate(...)
│ │ ├── 如果 num_blocks_to_allocate + watermark_blocks > 空闲块数 → 返回 None
│ └── (实际分配时的检查见后)
│
├── 2. 协调层: coordinator.get_num_blocks_to_allocate
│ └── for each manager in single_type_managers:
│ └── manager.get_num_blocks_to_allocate(...) # 计算该组需求
│
├── 3. 第二次准入检查 (实际分配前)
│ ├── available_blocks = block_pool.get_num_free_blocks() - reserved_blocks
│ ├── 如果 num_blocks_to_allocate + watermark_blocks > available_blocks → 返回 None
│
├── 4. 处理新命中的前缀缓存块 (可选)
│ └── coordinator.allocate_new_computed_blocks(...)
│
├── 5. 执行分配: coordinator.allocate_new_blocks
│ └── for each manager in single_type_managers:
│ └── SingleTypeKVCacheManager.allocate_new_blocks
│ ├── ① 处理部分命中 (CoW)
│ │ ├── 如果 request_id 在 _partial_hit_reqs 中:
│ │ │ ├── cow_block = block_pool.get_new_blocks(1)[0] ← 第一次调用 get_new_blocks
│ │ │ └── 记录 cow_block.block_id 到 new_block_ids
│ │ └── (该块用于写时复制)
│ │
│ ├── ② 计算常规新块需求
│ │ ├── num_required_blocks = cdiv(num_tokens, self.block_size)
│ │ ├── num_new_blocks = num_required_blocks - len(req_blocks)
│ │ └── 如果 num_new_blocks <= 0 → 跳过，只返回 CoW 块
│ │
│ └── ③ 从 BlockPool 获取新块
│ ├── new_blocks = block_pool.get_new_blocks(num_new_blocks) ← 第二次调用 get_new_blocks
│ ├── req_blocks.extend(new_blocks)
│ ├── 如果 _record_new_block_ids:
│ │ └── new_block_ids.extend(b.block_id for b in new_blocks)
│ └── 返回 cow_blocks + new_blocks
│
└── 6. 返回结果
└── 返回 KVCacheBlocks(new_blocks) # 包含所有组的新块列表KVCacheCoordinator.allocate_new_blocks