Prhub

#53781 [3/N] HiSparse: host-resident sparse-MLA decode hot-buffering

原始 PR 作者 MatthewBonanni 合并时间 2026-09-12 12:36 文件变更 124 提交数 138 评论 107 代码增减 +14373 / -1017

执行摘要

HiSparse 新增 host-resident 稀疏 MLA 解码层,GPU 压力下 spill 至主机内存

长上下文稀疏 MLA decode 的 KV 占用远超 GPU 显存容量,原方案 #46326 把 KV 常驻 host,但牺牲了设备端命中率与显存利用率。PR body 明确说明设计动机:"Unlike the original always-host-resident design in #46326, this implementation keeps as much KV as possible in the normal device cache and spills pages to pinned host memory only when GPU capacity must be reclaimed." 目标是让稀疏注意力的 indexer 缓存保持设备驻留、正常前缀可缓存,同时把稀疏 MLA 源 KV 的容量边界扩展到 host 内存,从而支持 GLM 等模型超长上下文(20k/32k/60k)的 decode 场景。

值得精读,尤其建议关注四点:一是 HiSparseCoordinator 的 dirty/clean/unpinned 状态机与 transition_watermark 水位设计,这是 write-back host 层正确性核心;二是连接器如何把每步路由状态塞进标准 build_connector_meta 管道并设置 has_sync_kv_loads,是理解 vLLM 连接器契约的活教材;三是 HiSparseMLAIndexGroup.convert_decode_logical_to_physical_topk 对多 token decode(投机验证)的逐步 swap_in 与输出槽位重排;四是 review 中关于 connector 契约、CUDA graph replay 状态重置的争论,能帮助理解 vLLM 调度/执行两阶段架构的边界。管理层应关注 njhill 提出的接口改动拆分问题,后续 PR 宜把连接器契约变更独立成小步。

讨论亮点

Review 中最有价值的交锋集中在三处:

  1. 连接器契约设计之争。NickLucche 明确反对 HiSparse 绕过标准接口:"the HiSparse connector currently bypasses the connector contract instead of implementing it. It is auto-injected, exposes custom hooks (apply_scheduler_output, finish_forward, set_request_state_indices, etc.) to ActiveKVConnector",并给出应遵循 build_connector_meta 标准管道、按角色实例化 scheduler/worker 的示例。作者随后多轮重构(commit d01e49708、fd61a3531)把命令改走标准 connector metadata 通道。

  2. 合入节奏争议。njhill 在合入前提出尖锐质疑:"I feel that this was rushed in too quickly given how invasive it is. In particular the nontrivial changes to the connector interface and corresponding changes to model runner etc. Could we not have at least separated the kv connector changes into a follow-on PR"。LucasWilkinson 最终以赶 v0.30 为由同意合入:"in the interest of getting this in v0.30 i think we can land this and do further cleanup as followups",该接口层面遗留问题被记 TODO。

  3. 实测暴露的正确性问题。samsja 报告高并发(1024 请求、64/GPU)下 GPU KV 填满后出现 KV 破坏,怀疑 "full CUDA graph replay skips the Python setup for eager host mirroring: each step clears the mirror metadata, replay leaves it unset";MatthewBonanni 确认是 blog 相关 rework 引入并已修复。LCAIZJ 在 H20 上发现三个阻塞问题:tokens_per_state 未实现导致启动 NotImplementedError、首次 forward 崩溃 "HiSparse GPU slots have no scheduler mapping envelope",以及 GLM 打包缓存池导致 NIXL region 几何失配。

实现拆解

实现按以下 5 步拆解:

  1. KV cache 布局规划(vllm/v1/hisparse/layout.py)。create_hisparse_layout 把稀疏 MLA 的 cache spec 拆成四类:host-resident source group(HISPARSE_SOURCE)、设备驻留 indexer group(HISPARSE_INDEXER)、按 indexer 页大小分组的 resident groups(HiSparseResidentSpec)与每请求固定容量的 hot groups(HiSparseHotSpec)。_partition_hisparse_specs 校验 source/indexer 都存在且拒绝 DeepSeek V4 压缩几何;get_hisparse_host_pool_bytes 从 hisparse_host_pool_gib 导出 host 预算。

  2. 调度器端协调(vllm/v1/hisparse/coordinator.py)。HiSparseCoordinator 持有 host 块分配、host 前缀身份与驻留状态机:_HiSparseRequestState 里每页经历 dirty(仅 GPU 副本)→ clean(host 副本完成)→ unpinned(分配引用释放、池可复用)三态;_ACTIVE_TAIL_PAGES = 2 保证在飞行步写入的页不会被提前交出;transition_watermark = max(hot_cost, num_blocks // 10) 控制请求开始从 host 读取的触发水位;spill 通过 build_offload_command 下发,且 host 页只在拷贝完成后才发布前缀哈希(_PendingPublication)。

  3. 连接器契约与 worker 数据搬运(vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py + worker.py)。HiSparseConnectorScheduler.build_connector_meta 每步把 offload 命令、host CoW 拷贝、row mirrors 和源块 ID 打包进 HiSparseConnectorMetadata,并设置 scheduler_output.has_sync_kv_loads = True 强制前向开始前完成加载;HiSparseConnectorWorker 通过 register_kv_caches 生命周期初始化,使用 _DMADescriptors 批量 DMA、_select_written_row_mirrors 按 slot 映射 coalesce 已写行、跨 TP 广播 host 写事件做到单写入者(_create_hisparse_host_events)。

  4. 稀疏 MLA attention 集成(vllm/v1/attention/backends/mla/index_group.py、sparse_mla_attention.py、flashinfer/flashmla/sparse 各后端)。HiSparseMLAIndexGroup 为每层创建 host 背书的 HiSparseCacheHandle,decode 路径经 convert_logical_to_physical_topk / convert_decode_logical_to_physical_topk 把 indexer 的逻辑 top-k 位置逐个 swap_in 到 hot buffer,再交给各后端 MQA kernel;prefill 侧通过 build_hisparse_prefill_staging_plan 做固定容量 staged gathering,避免动态输出破坏 CUDA graph。flashinfer_mla_sparse.py 新增 autotune_hisparse_decode 直接以合法 decode 包络做 autotune,并把 warmup 与模型 dummy run 解耦。

  5. P/D 与 NIXL 集成及测试配套。NIXL 仍是 P/D 传输层,decoder 在准入时选择落地策略:前缀能装进 GPU 池则写 resident 页,否则 bulk 稀疏 MLA KV 落到注册的 host 池;nixl/base_worker.py 为此引入混合内存类型 region 注册、按 block size 的 descriptor 缓存与失败接收延迟上报。测试配套包括 tests/v1/e2e/general/test_hisparse.py(spill 与前缀恢复)、tests/v1/core/test_prefix_caching.py(+1096 行 DMA row mirrors、异步投机镜像、indexer offload 补全等)、tests/v1/kv_connector/unit/ 下多个 NIXL/offloading 组合用例。

文件 模块 状态 重要度
vllm/v1/hisparse/runtime.py 核心运行时 added 9.18
vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py 连接器层 added 9.08
vllm/distributed/kv_transfer/kv_connector/v1/hisparse/worker.py 数据搬运 added 9.08
vllm/v1/hisparse/coordinator.py 调度协调 added 9.08
vllm/v1/attention/backends/mla/index_group.py 索引分组 added 8.99
vllm/v1/hisparse/binding.py 缓存绑定 added 8.89
vllm/v1/hisparse/layout.py 缓存布局 added 8.89
vllm/v1/core/single_type_kv_cache_manager.py 缓存管理 modified 8.96
vllm/distributed/kv_transfer/kv_connector/v1/nixl/base_worker.py 传输层 modified 8.51
tests/v1/e2e/general/test_hisparse.py 端到端测试 added 7.73
tests/v1/core/test_prefix_caching.py 前缀缓存测试 modified 7.73

关键符号

HiSparseConnectorScheduler.build_connector_meta HiSparseConnectorScheduler.update_connector_output HiSparseConnectorWorker.register_kv_caches HiSparseConnectorWorker.initialize HiSparseConnectorWorker._select_written_row_mirrors HiSparseCoordinator.commit_computed_blocks HiSparseCoordinator.take_host_block_copies HiSparseCoordinator.build_offload_command HiSparseMLAIndexGroup.convert_decode_logical_to_physical_topk HiSparseMLAIndexGroup.register_layer HiSparseSourceManager.allocate_external_computed_blocks HiSparseSourceManager.cache_blocks create_hisparse_layout bind_hisparse_kv_caches resolve_hisparse_block_size is_hisparse_decode_batch check_hisparse_host_memory autotune_hisparse_decode NixlBaseConnectorWorker._report_failed_recv NixlBaseConnectorWorker.get_transfer_results

关键源码片段

vllm/distributed/kv_transfer/kv_connector/v1/hisparse/connector.py core-logic

连接器契约边界核心:把 HiSparse 调度器侧 coordinator 状态(offload 命令、host CoW 拷贝、row mirrors、源块 ID)打包成标准 HiSparseConnectorMetadata,并通过完成计数回执驱动 spill 与 host 拷贝生命周期。

class HiSparseConnectorScheduler:
    """调度器侧连接器:把 coordinator 的每步状态打包成标准连接器元数据。"""
​
    def __init__(self, *, async_speculative: bool, draft_kv_lookahead: int = 0) -> None:
        self.coordinator: HiSparseCoordinator | None = None
        self.async_speculative = async_speculative
        self.draft_kv_lookahead = draft_kv_lookahead
        self.requests: dict[str, Request] = {}
​
    def build_connector_meta(self, scheduler_output: SchedulerOutput) -> HiSparseConnectorMetadata:
        assert self.coordinator is not None
        # HiSparse 每一步都会重建 worker 侧路由状态,因此加载必须在前向开始前
        # 启动,不能像普通 offloading 连接器那样推迟到前向之后。
        scheduler_output.has_sync_kv_loads = True
        scheduler_output.block_table_updates = (
            self.coordinator.take_block_table_updates() or None
        )
        command = self.coordinator.build_offload_command()
        host_block_copies = self.coordinator.take_host_block_copies()
        source_group_id = self.coordinator.host_group_id
        assert source_group_id is not None
        # 收集 host 源组中本次调度新 requests 与缓存命中请求占用的块 ID,
        # worker 据此知道哪些 host 行可以作为本次镜像的源。
        source_block_ids = [
            block_id
            for request in scheduler_output.scheduled_new_reqs
            for block_id in request.block_ids[source_group_id]
        ]
        for new_block_ids in scheduler_output.scheduled_cached_reqs.new_block_ids:
            if new_block_ids is not None:
                source_block_ids.extend(new_block_ids[source_group_id])
        num_computed_tokens = {
            request.req_id: request.num_computed_tokens
            for request in scheduler_output.scheduled_new_reqs
        }
        num_computed_tokens.update(
            zip(
                scheduler_output.scheduled_cached_reqs.req_ids,
                scheduler_output.scheduled_cached_reqs.num_computed_tokens,
            )
        )
        scheduled_requests = tuple(
            (
                request_id,
                num_computed_tokens[request_id],
                scheduled_count,
            )
            for request_id, scheduled_count in scheduler_output.num_scheduled_tokens.items()
        )
        # row_mirrors 描述每请求本步需要镜像到 host 的行区间;异步投机场景下
        # 起点回退到未解析的 placeholder 位置,确保镜像覆盖 draft 与实际
        # target 都可能写入的行。
        row_mirrors = {}
        for request_id, scheduled_start, scheduled_count in scheduled_requests:
            mirror_start = scheduled_start
            if self.async_speculative:
                mirror_start = max(
                    0,
                    scheduled_start - self.requests[request_id].num_output_placeholders,
                )
            row_mirrors[request_id] = self.coordinator.build_row_mirrors(
                (
                    (
                        request_id,
                        mirror_start,
                        scheduled_count + scheduled_start - mirror_start + self.draft_kv_lookahead,
                    ),
                ),
            )
        return HiSparseConnectorMetadata(
            command,
            host_block_copies,
            tuple(source_block_ids),
            row_mirrors,
            self.coordinator.all_context_pages_resident(scheduled_requests),
        )
vllm/v1/hisparse/coordinator.py core-logic

调度器端正确性核心,+752 行:页面三态(dirty/clean/unpinned)状态机、spill 命令生成、host 前缀延迟发布、GPU shadow 副本采纳与 `_ACTIVE_TAIL_PAGES` 飞行窗口保护。

@dataclass
class _HiSparseRequestState:
    """单个请求的驻留状态机。
​
    resident 页依次经历三态:dirty(只有 GPU 副本)、clean(host 副本完成,
    记录在 ``valid_pages``)、unpinned(clean 且分配引用已释放,池可复用,
    记录在 ``unpinned_pages``)。``pinned_clean`` 只存在于请求可以从 host
    读取之前,此时 clean 页仍需持有引用。
    """
​
    valid_pages: set[int] = field(default_factory=set)
    ready_prefix_pages: int = 0
    pending_pages: dict[int, int] = field(default_factory=dict)
    publication: _PendingPublication | None = None
    copies_recorded_blocks: int = 0
    pinned_clean: set[int] = field(default_factory=set)
    unpinned_pages: set[int] = field(default_factory=set)
​
​
class HiSparseCoordinator:
    """调度器侧 host 分配、发布、spill 与 GPU 驻留管理。
​
    GPU resident 页是 host 层的 write-back 缓存:host 副本持久化且所有者
    可读(已分配或已请求 hot 区)后,通过 ``BlockPool.unpin_blocks`` 释放
    引用——块仍留在块表里可读,但池将其计为空闲并可能再分发,此时所有者
    页被置空并重新发布块表。
    """
​
    def __init__(
        self,
        kv_cache_config: KVCacheConfig,
        managers: tuple[SingleTypeKVCacheManager, ...],
        max_model_len: int,
    ) -> None:
        # 收集 resident / hot 管理器并回填 coordinator 引用,供后续页回收回调。
        self.resident_managers = ...
        self.hot_managers = ...
        self.host_manager: HiSparseSourceManager | None = None
        self.host_group_id: int | None = None
        self.max_spill_pages = max_model_len // resident_block_size
        self.gpu_pool = self.resident_managers[0].block_pool
        # 共享池水位低于 transition_watermark 时,请求开始从 host 读取,
        # 这样新准入能找到可驱逐页而不是被钉住的页。
        hot_cost = sum(manager.blocks_per_request for manager in hot_managers)
        self.transition_watermark = max(hot_cost, kv_cache_config.num_blocks // 10)
        # host 块哈希 -> GPU 副本页:让 host 前缀命中能回到 GPU-resident,
        # 形成第一层免拷贝回收。
        self.copies: dict[BlockHashWithGroupId, tuple[KVCacheBlock, ...]] = {}
        # 未 pin 的 GPU 块 -> 仍在读它的 (request, page) 对。
        self._owners: dict[int, set[tuple[str, int]]] = {}
​
    def commit_computed_blocks(self, request_id: str, num_host_pages: int) -> None:
        """记录请求命中的 host 前缀页,并尝试采纳对应的 GPU shadow 副本。"""
        if not self.resident_managers or self.host_manager is None:
            return
        host_blocks = self.host_manager.req_to_blocks.get(request_id, ())
        num_host_pages = min(num_host_pages, len(host_blocks))
        if num_host_pages == 0:
            return
        state = self._get_request_state(request_id)
        state.valid_pages.update(range(num_host_pages))
        state.ready_prefix_pages = max(num_host_pages, state.ready_prefix_pages)
        # 命中长度完全 reconciled 后才做逐页散列采纳,避免影响前缀命中本身。
        self._adopt_copies(request_id, state, host_blocks[:num_host_pages])
        state.copies_recorded_blocks = max(state.copies_recorded_blocks, num_host_pages)

评论区精华

HiSparse 连接器绕过标准连接器契约 设计

NickLucche 指出 HiSparseConnector 自动注入、向 ActiveKVConnector 暴露 apply_scheduler_output/finish_forward/set_request_state_indices 等自定义钩子,"bypasses the connector contract instead of implementing it",并给出按角色实例化 scheduler/worker、使用 build_connector_meta 标准管道的示例。

结论:作者多轮重构(commit d01e49708、fd61a35318)把命令与镜像数据改走标准 connector metadata 通道,移除 ActiveKVConnector 特化绑定。 · 已解决

PR 体量过大、连接器接口改动应拆分 设计

njhill 合入前表示:"I feel that this was rushed in too quickly given how invasive it is. In particular the nontrivial changes to the connector interface and corresponding changes to model runner etc. Could we not have at least separated the kv connector changes into a follow-on PR"。

结论:LucasWilkinson 以赶 v0.30 为由同意合入并承诺 followup:"We can land this and do further cleanup as followups";接口层遗留被记 TODO。 · 待处理

全量 CUDA graph replay 跳过 eager host mirror 状态导致 KV 破坏 正确性

samsja 报告 1024 并发、GPU KV 填满后开始 host 读取时出现 KV 破坏,怀疑 "Full CUDA graph replay skips the Python setup for eager host mirroring: each step clears the mirror metadata, replay leaves it unset, and the spill path assumes the host copy is current"。

结论:MatthewBonanni 确认是 blog 相关 rework 引入并已修复;samsja 另提了 PR#56158 作为替代修复。 · 已解决

NIXL 失败接收只无效化第 0 组的块 正确性

coderabbitai 指出 _report_failed_recv 从 meta.local_block_ids[0] 构建 invalid 集合,而原实现会 union 所有 group;多 group 非 HMA 布局下 1..N 组块仍标记有效,可能被复用为已计算前缀,且会破坏 test_nixl_connector.py 的 ([1,2,3],[7,8]) 断言。

结论:建议改为遍历 meta.local_block_ids 全部 group 取并集;PR 合入前未见作者明确修复确认。 · 待处理

get_transfer_results 缺少 PCP producer 过滤 正确性

coderabbitai 发现 get_transfer_results 直接转发 worker 结果,缺少 get_finished 对非零 pcp_rank producer 清空 finished_sending 的过滤,post_forward 改走该路径后可能导致 KVOutputAggregator 过早判定发送完成。

结论:建议在 kv_role == "kv_producer" 且 pcp_rank > 0 时清空 results.finished_sending 后返回;未确认合入前修复。 · 待处理

SlotMappingMode 枚举是否为无关重构 style

LucasWilkinson 在 model_runner.py 看到把 slot_mapping_enabled 布尔改为 SlotMappingMode 枚举,认为 "This seems like a refactor unrelated to this PR. Can we undo this?"。

结论:作者后续将其归为 cruft 并移除无关改动,保留 HiSparseHotSpec/CircularBufferSpec 走 SlotMappingMode.NONE 的必要分支。 · 已解决

adaptive verification 的 max_decode_tokens 是否影响非 HiSparse 成本曲线 正确性

LucasWilkinson 质疑 batches_to_profile 用 min(max_num_batched_tokens, max_decode_tokens) 会 disable 非 HiSparse 的 tail sizes 成本曲线;MatthewBonanni 解释这是 HiSparse warmup 崩溃的修复。

结论:随后改为仅在 HiSparse 场景限制,并恢复完整 adaptive verification profiling(commit 5fcf84d8e、7aa100d0e)。 · 已解决

tokens_per_state 未实现导致引擎启动崩溃 正确性

LCAIZJ 实测 GLM-5.2 上 resolve_kv_cache_block_sizes 对 HiSparseHotSpec/HiSparseResidentSpec 求 tokens_per_state 时抛 NotImplementedError(两者继承 prefix_cacheable=True 但未实现该属性),并给出 prefix_cacheable=False 的 workaround。

结论:社区实测反馈;PR 正文与 review 均未见针对该问题的明确修复提交,属于合入前已知边界。 · 待处理

build_row_mirrors 中单个请求无块表会丢弃全批镜像 正确性

coderabbitai 指出 host_blocks 为 None 时 build_row_mirrors 直接 return (),会让同批其他请求已构建的 mirrors 全部作废,且镜像窗口只能前移,跳过会导致 host 行永久 stale;应改为 continue 跳过受影响请求。

结论:建议将 return () 改为 continue;未确认合入前修复。 · 待处理

num_gpu_blocks_override 不应作用于 host pool 正确性

coderabbitai 发现 override 同时赋给 host_num_blocks 与 gpu_num_blocks,丢弃 host_pool_gib 预算,override 超过 host 预算时 host 张量会请求超出 operator 分配的 pinned 内存。

结论:建议 override 只作用于 gpu_num_blocks,host 数保持由 host_budget 推导;未确认合入前修复。 · 待处理

warmup 应避免 dummy runs 并覆盖 Triton compile key 测试

LopezCastroRoberto 建议用 TritonWarmupTensor 调 triton_convert_req_index_to_global_index.warmup 而非真实 dummy run,并提醒新调用要与 #50175 的 compile-key 覆盖对齐;作者确认并采用 fake-tensor warmup 方案。

结论:采纳建议,sparse_mla_triton_warmup 使用 TritonWarmupTensor 预编译 index 转换内核。 · 已解决

风险与影响

  1. 连接器接口契约被拓宽:vllm/distributed/kv_transfer/kv_connector/v1/base.py 新增 finish_forward/reset_capture_state/get_transfer_results/bind_kv_cache_manager 等抽象方法,ActiveKVConnector 与 MultiConnector 同步改动。任何第三方连接器实现都会受影响,njhill 与 NickLucche 均认为应与 HiSparse 本体拆分。
  2. NIXL 失败接收路径正确性:coderabbitai 指出 _report_failed_recv 只无效化 meta.local_block_ids[0],多 group 布局下 1..N 组的块仍被标记有效,可能被调度器当作已计算前缀复用;同时 get_transfer_results 缺少 get_finished 已有的 PCP producer 过滤,非零 pcp_rank 的 finished_sending 会导致聚合器过早判定发送完成。
  3. CUDA graph replay 状态残留:samsja 实测的 host mirror 元数据未重置问题说明,任何在 eager segment 之外维护的每步状态都可能被 breakable/full graph replay 跳过,属于该类系统的高发风险面。
  4. HiSparse 与既有连接器组合的回归面:nixl/base_worker.py 的 block-size 解析从单一 cache_config.block_size 改为按 transfer group 取 LCM,offloading/offloading scheduler 的 group 索引语义也做了调整(group_idx=group_id),非 HiSparse 场景存在行为漂移风险;mooncake store 的 region 注册分支亦被改动。
  5. 配置边界:coderabbitai 发现 num_gpu_blocks_override 同时覆盖 host pool 块数,导致 operator 显式钉住 GPU 容量时 host 池被错误缩放;auto-fit 探测只传 kv_cache_groups[0] 会低估非 MLA 组的 GPU stride,可能启动期内存估算失败。

影响范围是系统级的:调度器(scheduler.py 增加 HiSparse coordinator 绑定)、KV cache 管理器(single_type_kv_cache_manager.py 新增 HiSparseSourceManager/HiSparseHotManager/HiSparseResidentManager)、attention 元数据契约(sparse_mla_attention.py 的 SparseMLAPrefillMetadata.host_staging_plan、SparseMLACommonMetadata.decode_max_query_len)、多个 FlashInfer/FlashMLA 稀疏后端、NIXL 与 mooncake 连接器以及 Model Runner V2 的 SlotMappingMode 与 skip_attn 均被改动。对用户而言,启用 hisparse_config + HiSparseConnector 后可在不换硬件的情况下把长上下文 decode 的 KV 容量扩展到 host 内存(B300 上 GLM 20k/32k/60k 均有收益),但配置组合校验严格(不支持 DeepSeek V4、要求 HMA 布局 BLHNC);对团队而言,这是 mrv2 与 kv-connector 两条主线的交叉重构,后续 HiSparse 系列 PR(#53779/#53780/#53263/#53782/#53592)均依赖本 PR 的接口基座。合入前 138 个提交、多轮主线 rebase 与反复 CI 修复也说明维护成本较高。

连接器接口契约变更 CUDA graph replay 状态残留 高并发 KV 破坏实测 NIXL 失败接收路径缺陷 跨模块核心架构变更 超大 PR review 覆盖不足

关联 Issue

未识别关联 Issue

当前没有检测到明确关联的 Issue 链接,后续同步到相关引用后会出现在这里。

完整报告

参与讨论