非连续 KV Cache

1. 概述

vLLM 通常将逻辑 KV-cache 块暴露为常规张量视图。在昇腾上，混合模型可以从一个公共的页对齐 cache 池中分配 Attention 和 Mamba cache 组。一个物理页可能包含对齐填充，不同的 cache 视图可能使用不同的区域或底层张量，因此将每个 cache 都视为稠密 K/V 张量可能会浪费内存或产生错误的 stride。

该特性在保持逻辑块表和块 ID 不变的同时，在现有物理分配之上构建带 stride 的张量视图。它在 Model Runner V1 中实现，适用于混合 Attention + Mamba 模型以及受支持的纯 GQA Attention 模型。

2. 目标与非目标

该实现旨在：

• 保留 vLLM 的逻辑块和块表语义；
• 避免将填充页复制到稠密 K/V 分配中；
• 向昇腾算子暴露正确的 shape、stride 和 storage-offset 元数据；
• 在一个分配方案中支持独立的 Attention 和 Mamba cache 组；
• 通过临时调试日志使布局可观测。

它不改变 token 调度、块表编号或逻辑 Attention/Mamba cache 规格。

3. 物理布局与逻辑布局

逻辑块是调度器寻址的单位。物理页是分配器预留的字节范围。当页为混合 cache 组进行对齐时，它们的大小可能不同。

对于 Attention cache，原始分配被视为块主序（block-major）张量。当物理页包含填充时，视图使用基于完整页大小的块 stride，而不是未填充的逻辑张量大小：

logical view:  [block, kv_head, token, head_dim]
physical page: [logical elements][alignment padding]
block stride:  physical page size / element size

该视图在设计上就是非连续的。即使 K 和 V 共享同一底层分配，它们仍保持为独立的逻辑视图。

3.1 连续兼容布局

以下布局是当实现需要所有状态张量保持连续时使用的旧版兼容路径：

tensor 1: [KV padding][conv state]...
tensor 2: [K blocks][SSM state]...
tensor 3: [V blocks][Mamba padding]...

此图不得被解释为非连续布局。它描述的是当填充被显式实体化时物理缓冲区如何被拆分。

3.2 非连续布局

非连续路径不要求将 K、V、卷积状态和 SSM 状态拼接为一个逻辑张量。相反，每个 cache 组通过各自的视图暴露：

Attention raw allocation  -> K view and V view
Mamba raw allocation      -> conv-state view and SSM-state view

这些视图可以具有不同的形状、步长和存储偏移量。它们使用相同的逻辑块编号，并且在缓存池需要时，使用相同的物理页步长。实现不得假设 Attention K 区域、Mamba SSM 区域和 Attention V 区域在一个原始张量内始终相邻或交错。

两个层是否共享底层张量由 cache 分配器和层映射决定。共享的底层张量通过现有的共享布局路径解释；否则每个层获得各自的带 stride 视图。

3.3 布局对比

方面	连续兼容路径	非连续路径
K/V 表示	稠密或显式切片的张量	独立的带 stride K/V 视图
Mamba 状态表示	显式拆分的填充区域	按物理页 stride 推进的视图
填充	实体化为张量区域	通过 stride 或 storage offset 跳过
数据移动	可能需要稠密暂存视图	设计为避免整页复制
主要不变式	保持每个状态张量连续	保留物理分配和正确的 stride

这两条路径使用相同的逻辑块 ID，但它们的物理解释不同。对连续兼容布局的更改不得用作关于非连续布局的依据。

3.4 示例：块 stride

假设一个逻辑 Attention 块包含 48 个 float16 元素，而混合页对齐到 128 字节：

logical payload = 48 elements * 2 bytes = 96 bytes
physical page   = 128 bytes
block stride    = 128 / 2 = 64 elements

因此下一个逻辑块从当前块之后 64 个元素处开始，而不是 48 个元素处。这 16 个元素的间隙不属于逻辑视图。

对于合并的 K/V 分配，简化的双通道视图可以表示为：

raw allocation: [K lane][V lane]
logical shape:  (2, blocks, block_tokens, kv_heads, head_dim)

第一个通道寻址 K，第二个通道寻址 V。通道 stride 和块 stride 均源自物理分配；它们不能仅从逻辑张量大小推断。

4. Cache 构建流程

NPUModelRunner._reshape_kv_cache_tensors 执行以下步骤：

1. 读取每个 cache 组的逻辑 KVCacheSpec。
2. 为该组选择算子块大小。
3. 识别原始分配是独立的 K/V 元组、合并张量还是 Mamba 状态分配。
4. 使用 view、as_strided 或现有的昇腾布局辅助函数构建类型化视图，而不复制底层存储。
5. 保留 worker 使用的逻辑层到 cache 的映射。

相关的实现要点包括：

_allocate_kv_cache_tensors()
    Allocate the raw backing tensor for each cache layer or group.

_reshape_kv_cache_tensors()
    Select the Attention/Mamba path and construct typed cache views.

_adjust_kv_layout()
    Build Mamba views whose block dimension advances by the physical page size.

gather_ssm_states()
    Gather indexed strided state rows into a kernel-friendly tensor.

scatter_ssm_states_()
    Write updated kernel state rows back to the original allocation.

调度器继续基于逻辑块 ID 运行。这些辅助函数仅改变 worker 将逻辑 ID 映射到物理地址的方式。

当启用 use_hybrid_blocks 时，混合 Attention 使用通用的混合分配约定。纯 GQA Attention 可以使用块主序非连续视图，而无需在同一模型中包含 Mamba 层。

5. Mamba 状态访问

Mamba 卷积和 SSM 状态的物理步长可能大于其逻辑行大小。因此，Mamba 状态辅助函数会将请求的行收集到对内核友好的张量中，并将更新散射回原始的带步长分配。

收集/散射路径会保留：

• 请求到状态的索引映射；
• has_initial_state 掩码；
• 物理页之间的存储填充；
• 原始缓存分配的就地更新。

状态索引测试同时覆盖零填充和带填充的物理行。

因此，访问流程为：

block table / state indices
          |
          v
strided physical state
          |
          +--> gather --> FLA or Mamba kernel
          |                  |
          +<-- scatter <-----+

6. FLA 算子集成

标准 GDN 预处理和循环路径使用 FLA AscendC 接口：

from fla_npu.ops.ascendc import (
    causal_conv1d_fn,
    causal_conv1d_update,
    recurrent_gated_delta_rule,
)

FLA 测试使用独立的 PyTorch 黄金实现，并放置在 vLLM-Ascend e2e 和 ops conftest 树之外。这可以防止无关的 Triton/自定义算子初始化在同一进程中改变 FLA 运行时状态。

7. 可观测性

在布局仍处于验证阶段时，Model Runner 会发出带有以下前缀的调试消息：

[non-contiguous-kv-cache]

Attention 日志包含模式（hybrid 或 gqa）、形状、步长以及 K/V 连续性。Mamba 日志包含页大小、状态形状、步长、存储偏移和连续性。这些日志是临时的观测点，在布局实现被认为稳定后可以移除。

8. 兼容性与限制

• 该功能针对 Model Runner V1 缓存构建实现。
• 逻辑块表和调度器可见的块 ID 保持不变。
• 确切的物理布局取决于模型的 Attention 和 Mamba 缓存规格以及所选的内核块大小。
• 稀疏、压缩、卸载或连接器特定的缓存路径必须保留各自的布局契约；不应在没有专门验证的情况下假定它们接受带步长的视图。
• 图捕获、PD 分离和 KV 连接器需要针对每个受支持的配置进行端到端验证。

当某条路径使用连接器、稀疏缓存、压缩缓存或卸载后端时，其注册和传输契约可能会施加单独的物理布局。不能仅因为模型是混合或 GQA 模型就为该路径启用非连续视图；连接器特定的路径需要其自身的验证。

9. 测试

该 PR 涵盖以下测试类别：

• 针对混合和纯 GQA 缓存视图的 Model Runner 单元测试；
• Mamba 状态收集/散射测试，包括物理填充；
• FLA 因果卷积和循环 GDN 数值测试；
• 使用 FLA 算子 mock 的 GDN 逐层 KV 连接器测试；
• 确认混合和 GQA 非连续布局的调试日志断言。

预期的验证矩阵为：

领域	用例
缓存构建	混合 Attention + Mamba；纯 GQA；带填充的页；组合 K/V 分配
状态访问	零填充；非零填充；索引收集；就地散射
FLA 算子	causal_conv1d_fn；causal_conv1d_update；recurrent_gated_delta_rule
执行模式	即时执行；启用布局后的图捕获
调度	预填充；解码；在模型路径启用时的 MTP/投机元数据
分布式路径	仅在传输布局验证后才使用 PD 分离和 KV 连接器

单元测试验证形状、步长、存储偏移和调度元数据。 FLA 单算子数值测试维护在 FLA 仓库中。 昇腾 nightly 状态索引测试将本地 Triton 辅助函数与 PyTorch 参考进行比较，包括未选中的行和物理页填充。 对于声称支持的每种图、连接器和分布式配置，仍然需要端到端模型测试。

代表性的 NPU 命令为：

pytest -q tests/e2e/nightly/single_node/ops/singlecard_ops/triton/test_mamba_state_index.py
pytest -q tests/ut/worker/test_model_runner_v1.py