feat(vllm): support packed 4D KV cache layouts - #252
Merged
Conversation
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
approved these changes
Aug 10, 2026
linhu-nv
left a comment
Collaborator
There was a problem hiding this comment.
LGTM. thanks for the great work
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)
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
支持 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 只描述存储中是否存在独立的 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,用于清理统计快照、累加计数器、重新计算命中率
并导出归约后的统计值。
验证情况:
和 transfer trace。
LAYERFIRST/BLOCKFIRST 布局、部分 layer 传输,以及双 GPU TP/layerwise
路径。
baseline 和 vectored 路径、部分 block 传输、MHA/MLA 布局,以及
layerwise SSD 到 CPU staging。
KV 数据和量化 scale 的 uint8 字节完全一致。