Skip to content

feat(vllm): support packed 4D KV cache layouts - #252

Merged
linhu-nv merged 2 commits into
taco-project:mainfrom
xuhao77:fix_vllm_4d_packed
Aug 10, 2026
Merged

feat(vllm): support packed 4D KV cache layouts#252
linhu-nv merged 2 commits into
taco-project:mainfrom
xuhao77:fix_vllm_4d_packed

Conversation

@xuhao77

@xuhao77 xuhao77 commented Aug 6, 2026

Copy link
Copy Markdown
Contributor

支持 packed 4D KV cache 布局

vLLM #44455 将非 MLA KV cache 从具有独立 K/V 维度的五维布局,
调整为将 K 和 V 打包到内容维度中的四维布局。

新版vllm KV 布局,#44455 之后:
Standard packed K/V: (num_blocks, num_kv_heads, block_size, 2 * head_size).
FlashInfer NVFP4: (num_blocks, 2 * num_kv_heads, block_size, packed_head_size).

最后一维包含 FP4 数据和 FP8 scale
packed_head_size = head_size // 2 + head_size // 16.

旧版的vllm KV布局:
LAYERBLOCK (num_blocks, 2, block_size, num_kv_heads, head_size).
LAYERFIRST (2, num_blocks, block_size, num_kv_heads, head_size).
FlashInfer NVFP4 (num_blocks, 2, block_size, num_kv_heads, packed_head_size).

FlexKV 此前复用了 is_mla,使其同时表示“cache 只有一个 KV region”和
“模型使用 MLA”。这会导致 packed MHA/GQA cache 在容量计算、TP head
分片和数据传输路径中错误地继承 MLA 行为。

增加 packed_kv,并通过 use_mla/is_mla 和 packed_kv 推导 single_kv_region:

  single_kv_region = use_mla/is_mla || packed_kv

single_kv_region 只描述存储中是否存在独立的 K、V region,从而决定
kv_dim 是 1 还是 2。is_mla 继续表示模型是否真正使用 MLA,并控制
head 复制、MLA D2H 模式及其他 MLA 专属行为。

主要改动:

  • 物理布局识别与配置传递:
    从 vLLM attention backend 查询实际 cache shape,识别 MLA 3D、
    packed 4D 和旧版 K/V 分离 5D 布局,并提取每个 TP rank 的物理
    head 数量和内容宽度。按照 TP size 汇总非 MLA cache 的存储 head,
    将识别出的 packed_kv 状态继续传递到 ModelConfig,以及 GPU、CPU、
    SSD 和 remote 的 KVCacheLayout。容量计算、存储分配、兼容性检查、
    日志、trace 和 replay 均使用最终得到的物理布局。

  • KV region 语义拆分与接口重命名:
    通过 use_mla 和 packed_kv 推导 ModelConfig.kv_dim,并通过 is_mla
    和 packed_kv 提供 KVCacheLayout.single_kv_region。对于原先名为
    is_mla、实际却只控制 K/V 循环次数或偏移计算的参数,在 CUDA binding、
    SSD、remote、非 TP GDS、NIXL offset helper 和 transfer trace 中统一
    重命名为 single_kv_region。

    对同时需要 KV region 数量和真实模型类型的路径,继续保留这两个参数。
    TP、TP GDS 和 layerwise transfer 同时接收 is_mla 和 packed_kv,并在
    内部推导 single_kv_region。将 CE binding 原来的 is_mla 参数重命名为
    single_kv_region,同时增加 ce_is_mla,使 KV region 遍历和 MLA 专属
    CE 行为相互独立。

  • vLLM GPU cache 注册与布局校验:
    在 vLLM 完成 cache 分配后,根据实际 tensor 维度和 K/V 轴位置构造
    GPU KVCacheLayout。对于 packed 4D cache,在向 FlexKV 注册前校验
    tensor 是否使用 token-major 的 NHD stride。对于当前 transfer kernel
    无法正确寻址的 HND 布局,直接拒绝注册,避免静默损坏 KV 数据。

  • GPU 与 CPU 之间的数据传输:
    在 KVCacheLayout 中通过 is_mla 和 packed_kv 推导 single_kv_region,
    Python worker 使用它计算 region 数量和传输 stride,并通过 binding
    将其传递到 CUDA、CE、TP 和 layerwise 实现。根据 single_kv_region
    遍历一个或两个物理 KV region,同时继续使用 is_mla 控制 TP head
    复制、MLA D2H 模式和 MLA 专属 CE 路径。这样 packed MHA/GQA 可以
    合并 K/V region,同时保留正常的 TP head 分片行为。

  • SSD、remote 和非 TP GDS 数据传输:
    将 packed_kv 继续传递到 CPU、SSD 和 remote 布局,并推导相同的
    single_kv_region。worker 将其传递到 SSD、PCFS 和非 TP GDS 实现,
    使各路径按照物理布局计算 K/V 循环和 chunk offset,只传输实际存在
    的 KV region。

  • TP GDS 数据传输:
    同时向 TP GDS 传递 is_mla 和 packed_kv,并在线程组内部推导
    single_kv_region。使用 single_kv_region 控制 GDS kernel 遍历一个
    还是两个 KV region,同时继续使用 is_mla 选择每个 rank 的 SSD
    offset。MLA rank 访问相同的复制 region,而普通和 packed MHA/GQA
    rank 使用 ssd_tp_stride 访问各自的 head 分片。

  • 非连续 tensor 的 NIXL 数据传输:
    使用 single_kv_region 计算 packed NHD tensor 的 KV chunk offset。
    不再使用可能为非连续逻辑视图生成连续临时副本的 reshape 展平方式,
    而是基于底层 storage、storage offset 和布局 stride 构造字节视图。
    NIXL 现在直接读写原始 KV cache allocation。

  • nvcomp ANS 安全降级:
    初始化压缩时检查 KVCacheLayout.packed_kv。对于 packed cache,禁用
    尚未将 single-KV-region 语义与 MLA size-table 语义分离的 ANS 路径,
    回退到未压缩传输,避免错误寻址 KV region 或 TP size table。

  • vLLM connector 统计聚合:
    适配 vLLM #43877 引入的 scheduler 侧统计流程。vLLM 现在会先处理
    worker connector 输出,再收集 scheduler 侧统计;通过 is_empty 判断
    scheduler 统计是否有效;在 worker 和 scheduler 统计同时存在时进行
    聚合;同时也支持只输出 scheduler 侧统计。

    FlexKV 此前直接实例化 KVConnectorStats,但当前 vLLM 已将其作为接口
    基类,reset、aggregate、reduce 和 is_empty 均没有默认实现。增加
    _FlexKVConnectorStats,用于清理统计快照、累加计数器、重新计算命中率
    并导出归约后的统计值。

验证情况:

  • 通过 23 个单元测试,覆盖配置、布局、NIXL storage alias、vLLM adapter
    和 transfer trace。
  • 通过 11 个 packed 4D GPU 回环测试,覆盖 CUDA kernel、CE、
    LAYERFIRST/BLOCKFIRST 布局、部分 layer 传输,以及双 GPU TP/layerwise
    路径。
  • 通过 176 个 SSD transfer 回归测试,覆盖 thread 和 io_uring I/O、
    baseline 和 vectored 路径、部分 block 传输、MHA/MLA 布局,以及
    layerwise SSD 到 CPU staging。
  • 通过 4 个具有代表性的 MHA、MLA、CE 和 layerwise 回归用例。
  • 通过 2 个 DSV4 multi-group 和 SWA layerwise 回环测试。
  • 通过当前 FlashInfer 4D NVFP4 布局的 GPU→CPU→GPU 端到端测试,
    KV 数据和量化 scale 的 uint8 字节完全一致。
  • 通过 git diff 检查和 Python compileall 检查。

xuhao77 added 2 commits August 6, 2026 09:22
vLLM #44455 changes non-MLA KV caches from five-dimensional layouts with a
separate K/V axis to four-dimensional layouts that pack K and V into the
content dimension. FlashInfer NVFP4 additionally uses an expanded physical
head axis to store K/V data and quantization scales.

FlexKV previously overloaded is_mla to mean both "the cache has one KV region"
and "the model uses MLA". This caused packed MHA/GQA caches to incorrectly
inherit MLA behavior in capacity calculations, TP head sharding, and transfer
paths.

Add packed_kv and derive single_kv_region from use_mla/is_mla and packed_kv:

    single_kv_region = use_mla/is_mla || packed_kv

single_kv_region only describes whether storage contains separate K and V
regions and therefore whether kv_dim is one or two. is_mla continues to
represent the actual MLA model type and controls head replication, MLA D2H
modes, and other MLA-specific behavior.

Main changes:

- Physical layout discovery and configuration propagation:
  Query the vLLM attention backend for the actual cache shape, recognize MLA
  3D, packed 4D, and legacy split-K/V 5D layouts, then extract the physical
  head count and content width for each TP rank. Aggregate non-MLA storage
  heads across the TP size and propagate the detected packed_kv state through
  ModelConfig and GPU, CPU, SSD, and remote KVCacheLayout instances. Use the
  resulting physical layout for capacity calculations, storage allocation,
  compatibility checks, logging, tracing, and replay.

- KV-region semantics and API renaming:
  Derive ModelConfig.kv_dim from use_mla and packed_kv, and expose
  KVCacheLayout.single_kv_region from is_mla and packed_kv. Rename transfer
  parameters that were called is_mla but only controlled K/V loop counts or
  offsets to single_kv_region across CUDA bindings, SSD, remote, non-TP GDS,
  NIXL offset helpers, and transfer tracing.

  Keep both concepts in paths that need the KV-region count and the true model
  type. TP, TP GDS, and layerwise transfers receive is_mla and packed_kv and
  derive single_kv_region internally. Rename the CE binding's former is_mla
  argument to single_kv_region and add ce_is_mla so KV-region traversal and
  MLA-specific CE behavior remain independent.

- vLLM GPU cache registration and layout validation:
  Build the GPU KVCacheLayout from the actual tensor rank and K/V-axis
  position after vLLM allocates the cache. For packed 4D caches, validate that
  the tensor uses token-major NHD strides before registering it with FlexKV.
  Reject HND layouts that the current transfer kernels cannot address
  correctly instead of silently corrupting KV data.

- GPU-to-CPU and CPU-to-GPU transfers:
  Derive single_kv_region from is_mla and packed_kv in KVCacheLayout, use it in
  Python workers to calculate region counts and transfer strides, and pass it
  through the bindings to CUDA, CE, TP, and layerwise implementations. Iterate
  over one or two physical KV regions based on single_kv_region while retaining
  is_mla for TP head replication, MLA D2H modes, and MLA-specific CE paths.
  Packed MHA/GQA therefore merges the K/V region without losing normal TP head
  sharding.

- SSD, remote, and non-TP GDS transfers:
  Propagate packed_kv through CPU, SSD, and remote layouts and derive the same
  single_kv_region value. Pass it from the workers into SSD, PCFS, and non-TP
  GDS implementations so each path calculates K/V loops and chunk offsets from
  the physical layout and transfers only the regions that actually exist.

- TP GDS transfers:
  Pass both is_mla and packed_kv into TP GDS and derive single_kv_region inside
  the thread group. Use single_kv_region to control whether GDS kernels traverse
  one or two KV regions, while is_mla continues to select the per-rank SSD
  offset. MLA ranks access the same replicated region, whereas regular and
  packed MHA/GQA ranks use ssd_tp_stride to access their own head shards.

- NIXL transfers for non-contiguous tensors:
  Use single_kv_region to calculate KV chunk offsets for packed NHD tensors.
  Replace reshape-based flattening, which can materialize a contiguous
  temporary for non-contiguous logical views, with byte views constructed from
  the underlying storage, storage offset, and layout strides. NIXL now reads
  and writes the original KV cache allocation directly.

- Safe nvcomp ANS fallback:
  Check KVCacheLayout.packed_kv when compression is initialized. Disable the
  ANS path, which does not yet separate single-KV-region semantics from MLA
  size-table semantics, and fall back to uncompressed transfers for packed
  caches instead of addressing their regions or TP size tables incorrectly.

- vLLM connector statistics aggregation:
  Adapt to the scheduler-side statistics flow introduced by vLLM #43877. vLLM
  now processes worker connector output before collecting scheduler-side
  statistics, calls is_empty to determine whether scheduler statistics are
  valid, aggregates worker and scheduler statistics when both exist, and can
  emit scheduler-only statistics.

  FlexKV previously instantiated KVConnectorStats directly, but current vLLM
  treats it as an interface base whose reset, aggregate, reduce, and is_empty
  methods have no default implementation. Add _FlexKVConnectorStats to clear
  snapshots, accumulate counters, recompute hit ratios, and export reduced
  values. Keep the worker-side sentinel stats so aggregation is delegated to
  the real scheduler-side snapshot and the final data reaches SchedulerStats
  and the logging pipeline.

Validation:

- Passed 23 unit tests covering configuration, layouts, NIXL storage aliasing,
  the vLLM adapter, and transfer tracing.
- Passed 11 packed 4D GPU round-trip tests covering CUDA kernels, CE,
  LAYERFIRST/BLOCKFIRST layouts, partial-layer transfers, and two-GPU
  TP/layerwise paths.
- Passed 176 SSD transfer regression tests covering thread and io_uring I/O,
  baseline and vectored paths, partial-block transfers, MHA/MLA layouts, and
  layerwise SSD-to-CPU staging.
- Passed four representative MHA, MLA, CE, and layerwise regression cases.
- Passed two DSV4 multi-group and SWA layerwise round-trip tests.
- Passed the end-to-end GPU-to-CPU-to-GPU test for the current FlashInfer 4D
  NVFP4 layout with byte-exact uint8 KV data and quantization scales.
- Passed git diff checks and Python compileall checks.
… Layout are incorporated into test_kvmanager.py.

@linhu-nv linhu-nv left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

LGTM. thanks for the great work

@linhu-nv
linhu-nv merged commit a79ddd5 into taco-project:main Aug 10, 2026
3 of 4 checks passed
zhjc1124 added a commit to zhjc1124/FlexKV that referenced this pull request Aug 11, 2026
IndexerCacheConfig was removed in PR taco-project#252/taco-project#255 — indexer data is now
packed inside KV blocks via layer_groups, not a separate pool.
PoolKind.INDEXER and enable_pool_specs() also removed.

Changes:
- Remove IndexerCacheConfig import and indexer= kwarg
- _make_cache_config() no longer takes with_indexer
- Rename joint-match tests to KV-only match tests
- Replace PoolKind.INDEXER with PoolKind.KV/SWA
- Remove enable_pool_specs tests (method no longer exists)
@xuhao77
xuhao77 deleted the fix_vllm_4d_packed branch August 13, 2026 01:41
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants