MAX 流水线 KV 缓存管理详解从内存规划、分页缓存到跨设备传输引擎【免费下载链接】mojoThe Modular Platform (includes MAX Mojo)项目地址: https://gitcode.com/GitHub_Trending/mo/mojo导读max.pipelines.kv_cache是 Modular MAX 平台中负责大语言模型 KVKey-Value缓存全生命周期管理的核心 Python 模块它既决定缓存能占多少显存、按什么页大小组织也负责前缀命中、淘汰换入换出、跨设备传输等运行时行为。本文以 max/python/docs/pipelines.kv_cache.rst 的 API 结构为主线结合仓库源码逐一剖析其内存规划MemoryPlanner、配置体系KVCacheConfig / KVConnectorConfig、分页缓存管理器PagedKVCacheManager、连接器KVConnector与传输引擎KVTransferEngine的设计与用法。读完本文你将能理解 MAX 推理服务中 KV 缓存的显存预算从何而来、如何通过命令行参数精确调控缓存精度与分层离载offload、以及分布式场景下 KV 数据如何跨副本流动。模块的公开入口定义在 max/python/max/pipelines/kv_cache/init.py对外导出KVCacheConfig、KVConnectorConfig、MemoryPlanner、PagedMemoryPlanner、PagedKVCacheManager、KVTransferEngine等全部核心类型并额外导出了max.nn.kv_cache中的KVCacheGroupId。下文按原文档的六大部分逐一展开。一、内存规划KV 缓存的显存预算从哪来自回归文本生成模型的标准推理路径中KV 缓存是最大的动态显存消耗项。max.pipelines.kv_cache的Memory planning部分提供了一套可插拔的预算估算体系实现在 max/python/max/pipelines/kv_cache/memory_planner.py。1.1 结构化协议ModelConfig 与 ModelConfigWithKVCache这两个类并非具体实现而是typing.Protocol带runtime_checkable装饰器可被isinstance运行时检查ModelConfigmemory_planner.py仅要求暴露devices属性模型运行所在的设备列表是最小的模型配置结构契约ModelConfigWithKVCachememory_planner.py在ModelConfig基础上追加get_kv_params() - KVCacheParamInterface使规划器能以类型化接口而非getattr获取缓存参数。这种协议化设计让规划器只依赖配置对象的最小表面不绑定具体架构从而可以独立于整个 pipeline 栈使用。1.2 基类 MemoryPlannerMemoryPlannermemory_planner.py仅由一个ModelConfig构造而非完整PipelineConfig并为所有估算方法提供默认实现子类按需覆写。其核心方法包括estimate_weights_size(pipeline_config)估算模型权重占用字节数默认委托给pipeline_config.model.weights_size()estimate_activation_memory(pipeline_config, huggingface_config)估算权重之外的激活内存临时中间张量默认返回0需要大临时缓冲的架构如 MLA 上投影、专家并行路由缓冲在此覆写infer_max_batch_size(pipeline_config, devices, weights_size)当用户未显式设置max_batch_size时推断架构相关默认值默认返回None交给框架级推断estimate_signal_buffer_memory(...)估算跨设备 P2P 集合通信所需的信号缓冲内存默认单设备返回0对于无条件执行 allreduce 的模型如使用VocabParallelEmbedding的 VLM可通过类级_always_signal_buffers True让单设备也预留Signals.NUM_BYTES。1.3 标准实现 PagedMemoryPlannerPagedMemoryPlannermemory_planner.py是使用分页 KV 缓存的文本生成模型的标准规划器。构造时要求配置实现ModelConfigWithKVCache否则抛出TypeError。对需要预留固定显存空间的架构典型如 VLM 要为视觉处理留出余量源码明确建议不要手写自定义子类而是用工厂方法with_activation_reservation直接生成预配置子类memory_plannerPagedMemoryPlanner.with_activation_reservation( 15 * 1024**3 # 15 GiB )该方法memory_planner.py会动态创建名为PagedMemoryPlanner(15GiB)的子类使其estimate_activation_memory返回固定预留值对于无条件 allreduce 的模型还可追加always_signal_buffersTrue让单 GPU 场景也预留信号缓冲memory_plannerPagedMemoryPlanner.with_activation_reservation( 15 * 1024**3, always_signal_buffersTrue )二、配置体系KVCacheConfig 与 KVConnectorConfig配置层实现在 max/python/max/pipelines/kv_cache/config.py两个 Pydantic 配置类均继承自max.config.ConfigFileModel可直接从命令行、YAML/JSON 配置文件加载。2.1 KVCacheConfig分页缓存的核心参数KVCacheConfigconfig.py的字段即max serve/max generate等命令的--kv-cache-*系列参数来源字段及默认值如下表字段默认值含义kv_cache_page_size128分页 KV 缓存中单个 page 容纳的 token 数enable_prefix_cachingTrue是否启用前缀缓存enable_dp_cross_replica_prefix_copyTrue数据并行DP场景下是否允许把另一副本 GPU 上的前缀缓存块通过设备到设备D2D拷贝搬来命中关闭后跨副本复用只能走共享主机/磁盘层连接器或重算kv_connector_config默认null连接器KV 缓存连接器配置见下文 2.2device_memory_utilization0.9进程应占用的可用显存比例剩余头寸留给 KV 缓存kv_cache_workspace total_free_memory * device_memory_utilization - model_weights_sizekv_cache_formatNone显式覆盖缓存数据类型可选float32、bfloat16、float8_e4m3fnindexer_kv_cache_formatNone独立覆盖 MiniMax 稀疏索引器IndexK缓存 dtype可选bfloat16、float8_e4m3fn默认保持 BF16使--kv-cache-formatfloat8_e4m3fn仍表示主 GQA FP8 索引器 BF16state_pool_dtypeNone覆盖混合模型循环状态池SSM/线性注意力卷积与循环状态的存储 dtype可选bfloat16、float32FP32 使推测解码严格跟随非推测状态轨迹kv_cache_hash_algoahash64缓存块身份哈希算法可选ahash64快速非密码学、sha256256 位密码学哈希、sha256_64SHA-256 链截断到 64 位以兼容协议kv_cache_hash_seedNone可选的 64 位十六进制字符串32 字节集群级种子省略时sha256/sha256_64启动时随机生成ahash64不生成避免影响既有部署其中device_memory_utilization直接出现在max serve的官方示例中见 max/python/docs/cli/serve.rstmax serve \ --model google/gemma-3-12b-it \ --devices gpu:0 \ --max-batch-size 8 \ --device-memory-utilization 0.9KVCacheConfig.to_params()config.py负责把该配置转换为max.nn.kv_cache.cache_params.KVCacheParams并依据is_mla标志自动选择子类MLA 模型生成MLAKVCacheParams此时必须提供num_q_heads否则生成MHAKVCacheParams。架构可通过page_size参数传入其内核强制的页大小下限而不必修改共享配置kv_cache_hash_seed则通过resolve_kv_hash_seed辅助函数解析。2.2 KVConnectorConfig外部缓存层配置KVConnectorConfigconfig.py采用类型与设置同体的设计连接器类型与其参数作为一个对象整体配置例如--kv-connector-config {type: rust_tiered}公共字段是强类型的连接器特有字段通过extraallow透传经model_extra访问。字段如下字段默认值含义typenull连接器类型见下方KVConnectorType枚举默认无外部缓存host_offload_max_gbNone主机内存离载上限GiB未设置时按设备页池的 1.5 倍自动估算disk_offload_dirNone磁盘离载目录tiered类型必填disk_offload_max_gbNone磁盘离载上限GiB 0未设置时为设备页池 2 倍0表示去掉磁盘层纯主机连接器num_disk_workers32磁盘 I/O 工作线程数 0负载下越高排空磁盘操作队列越快约 32 后收益递减block_store_endpointNone同机 dKV 服务端点支持 IPCipc:///path或 TCPtcp://host:portdkv类型必填连接器类型KVConnectorType定义于 max/python/max/nn/kv_cache/cache_params.py四个取值语义明确null无设备外后备存储页仅驻留设备tiered把被逐出的页分层放到主机内存与磁盘已废弃作为rust_tiered的向后兼容别名rust_tiered唯一的主机/磁盘分层实现由 Rustkv_tier_connector扩展支撑拷贝与磁盘 I/O 运行在 Rust 线程无 GIL 争用并通过异步传输句柄与 GPU 计算重叠要求enable_prefix_caching、host_offload_max_gb、disk_offload_dir且仅支持 CUDA/HIP 设备dkv通过分布式 KV 块存储路由页要求block_store_endpoint。值得注意的实现细节源码注释明确说明远程 dKV 端点在运行时由 Orchestrator 的每请求dkv_cache_hint发现而非静态配置——连接器在 Rust 侧解析 hint 并自行拨号每个命名的实例。三、分页缓存管理器PagedKVCacheManager 及其配套3.1 PagedKVCacheManager核心管理器PagedKVCacheManagermax/python/max/pipelines/kv_cache/paged_kv_cache/cache_manager.py是支持数据并行DP与张量并行TP的分页 KV 缓存管理器其 docstring 提供了完整的端到端用法示例cache_manager.pyimport numpy as np from max.driver import CPU from max.dtype import DType from max.engine import InferenceSession from max.graph import DeviceRef from max.nn.kv_cache import MHAKVCacheParams from max.pipelines.context import TextContext, TokenBuffer from max.pipelines.kv_cache import PagedKVCacheManager from max.pipelines.modeling.types import RequestID params MHAKVCacheParams( dtypeDType.float32, n_kv_heads8, head_dim128, num_layers2, page_size128, devices[DeviceRef.CPU()], ) kv_manager PagedKVCacheManager( paramsparams, sessionInferenceSession(devices[CPU()]), total_num_pages8, max_batch_size4, ) def make_context() - TextContext: tokens np.array([1, 2, 3, 4], dtypenp.int64) return TextContext( request_idRequestID(), max_length1000, tokensTokenBuffer(tokens), ) ctx1 make_context() ctx2 make_context() # Allocate metadata for requests in batch kv_manager.claim(ctx1) kv_manager.claim(ctx2) # Allocate blocks for these requests kv_manager.alloc(ctx1) kv_manager.alloc(ctx2) # Get KVCache inputs to feed to graph kv_cache_inputs kv_manager.runtime_inputs([[ctx1, ctx2]]) # Run model... # Update requests with newly generated tokens ctx1.update(42) ctx2.update(42) # Commit newly written blocks to prefix cache kv_manager.step(ctx1)该示例完整呈现了管理器的请求生命周期claim为批内请求分配元数据→alloc分配块→runtime_inputs取图输入→ 前向 →update追加新 token→step把新写块提交到前缀缓存。在设备侧管理器维护每设备一张 LUTlookup table表与缓存长度缓冲LUT 形状为(max_batch_size, padded_lut_cols(max_total_num_pages))、uint32类型内维做了填充以保证 SIMDpopulate在任意合法first_lut_idx之后都能安全加载至多 16 个连续uint32而不越界cache_manager.py。3.2 DummyKVCache测试与禁用场景的无操作实现DummyKVCachemax/python/max/pipelines/kv_cache/paged_kv_cache/dummy_cache_manager.py继承PagedKVCacheManager是缓存被禁用或处于测试时的空操作实现claim/step/release均为 no-opalloc返回一个已完成传输CompletedTransfer.load()contains对任何请求都返回Trueblock_count/host_byte_count/disk_byte_count返回恒定的1 空闲/1 总计占位值指标为空。3.3 计数快照BlockCount 与 ByteCount两者定义于 max/python/max/pipelines/kv_cache/kv_connector.pyBlockCountL35-L67设备G0块池占用快照字段free/total派生出used、used_pct、free_pct属性total 0时百分比返回0而非除零ByteCountL70-L103外部缓存层主机/磁盘的字节占用快照。用字节而非块是因为连接器的主机/磁盘层是操作员按字节host_offload_max_gb、disk_offload_max_gb设置的预算且其块宽度不必与设备一致——MLA 复制单元在主机上只存一份、加载时广播回去。3.4 InsufficientBlocksError分配失败的显式信号InsufficientBlocksErrormax/python/max/pipelines/kv_cache/paged_kv_cache/block_utils.py是空闲块不足无法满足分配时抛出的异常。它不只是运行时错误在 registry.py 的load_kv_manager中会预先校验单请求在max_seq_len下必须能放入设备块池否则请求无法被抢占无物可逐出运行时必然以InsufficientBlocksError崩溃——因此选择在启动阶段直接失败而非留到运行时。四、KV 连接器外部缓存层的加载/卸载协议KVConnector协议kv_connector.py是管理器拥有设备张量、块分配与设备侧前缀缓存连接器负责外部层主机内存、磁盘、dKV 等读写的分工契约。关键方法load(block_ids, block_hashes, replica_idx, hint)把外部缓存数据加载进设备块H2D 拷贝返回KVConnectorTransferoffload(block_ids, block_hashes, replica_idx)把设备块卸载到外部缓存D2H 拷贝touch(block_hashes, replica_idx)刷新外部层对被设备G0侧前缀缓存命中所服务块的 LRU 新鲜度——尽力而为、即发即忘避免设备上仍热门的块被外部层误逐出契约要求从真正根节点开始、按序列顺序传完整集合count_cached_prefix(block_hashes)按前缀顺序统计外部层连续驻留的前导块数返回(num_host_blocks, num_disk_blocks)必须是严格只读操作wait_for_loads()/wait_for_offloads()前向前后对已发布传输的排序/落定当前仅 dKV 连接器使用其余为 no-op。跨该协议的所有块哈希均为规范字节形式ahash64系含sha256_64为 8 字节大端完整 SHA-256 为 32 字节。4.1 传输句柄模型KVConnectorTransfer 与 CompletedTransferKVConnectorTransferkv_connector.py是连接器传输加载/卸载的句柄其核心价值是让管理器把传输与 GPU 计算重叠传输持有其触及的g0设备块 id管理器在is_complete()返回True前保持这些块被钉住随后恰好解除一次。is_complete必须是廉价、无副作用的轮询原子量或cudaEventQuery式检查。协议支持两种完成模型同步 / 流有序连接器dKV拷贝在转发流上或 GPU 有序地排在其前直接返回CompletedTransferkv_connector.pyis_complete立即为True管理器立刻提交复用的前缀、请求不脱离批次异步连接器Rustrust_tiered拷贝在独立拷贝引擎上执行句柄的is_complete仅在拷贝落地后翻转管理器钉住块、推迟提交并暂时把请求隔离出批次GPU 在此期间运行其他就绪工作。五、传输引擎KVTransferEngine 与跨副本数据流动5.1 传输引擎的角色KVTransferEnginemax/python/max/pipelines/kv_cache/paged_kv_cache/transfer_engine.py是支持 DP 与 TP 的 KV 缓存传输引擎基于 NIXLNexus Interconnect eXchange Layer构建是底层TransferEngine之上的薄封装新增能力包括张量形状与设备类型校验、从产者编写的 NIXL group 构建、按组的 TP 复制replicated_per_group、以及from_paged_kv_cache()便捷构造器。其构造接收按 DP 副本索引的外层列表memory[replica][group]每个条目是一个KVCacheMemory对应一个逻辑(child, kind)张量携带其全部 TP 分片视图来自KVCacheBuffer.to_memory()。构造函数会强校验结构不变量transfer_engine.py每个副本必须至少含一个 group、所有副本 group 数量一致、所有 group 的total_num_pages一致。文档特别强调单个 TransferEngine 不是线程安全的它面向 MAX 的单线程调度器使用仅在不同线程/进程之间彼此通信。5.2 元数据与请求数据结构KVTransferEngineMetadatatransfer_engine.py在仅含传输信息的TransferEngineMetadata之上扩展 KV 缓存/拓扑字段——total_num_pages每个张量的总页数、bytes_per_page每页字节数、bytes_per_group每个 NIXL group 的每页字节数首项为主 group后续对应推测解码中的草稿 KV 等额外 group、replicated_per_group每 group 的 TP 复制标志True表示 MLA 式跨 TP 分片完全相同复制False表示分片。该对象可安全跨线程/进程传输是线上唯一的复制信息源。TransferReqDatatransfer_engine.py基于msgspec.Struct的传输请求元数据字段包括源/目标引擎名src_name/dst_name、传输名、每 TP 分片一个的transfer_ids、源/目标索引长度可与transfer_ids不同、副本索引、is_read是否为目标发起的 READ 拉取、tp_shard_count与local_shards_used等同样可安全跨线程/进程传输。5.3 环境变量与端口工具传输引擎读取两个关键环境变量族MODULAR_NIXL_TRANSFER_BACKENDNIXL 传输后端类型默认ucx由统一校验器validate_nixl_backend检查见 max/python/max/pipelines/kv_cache/_nixl_backend.py 中的NIXL_BACKEND_ENV_VARdKV 连接器读取同一变量但为必填无自动回退。UCCL 相关变量若未设置UCCL_SOCKET_IFNAME/NCCL_SOCKET_IFNAME会默认把 UCCL 的带外引导网卡钉到主机的默认路由 NIC控制网络避免自动选择落在仅承载 RoCE 的网卡上导致对端拨号挂起同时默认设置UCCL_P2P_TRANSPORTrdma与UCCL_P2P_DISABLE_IPC1规避同机intranode收发对在 IPC 传输下的死锁transfer_engine.py。available_port(start_port8000, end_port9000, max_attempts100)transfer_engine.py在指定闭区间内随机探测可用 TCP 端口绑定SO_REUSEADDR以避免 TIME_WAIT 问题max_attempts次内未找到则抛RuntimeError。六、工厂函数与工具函数6.1 load_kv_manager按参数装载缓存管理器load_kv_manager(params, max_batch_size, max_seq_len, session, available_cache_memory, is_di_enabled, model_name)registry.py是缓存管理器的统一装载入口接收KVCacheParams单缓存或MultiKVCacheParams多缓存返回的单个管理器原生处理全部缓存。其装载逻辑层层把关若params是MagicMock测试桩直接返回MagicMock编译专用模式虚拟设备模式下返回Mock避免 GPU 显存分配available_cache_memory/max_batch_size未在内存估算阶段设置、或max_batch_size 0时抛ValueError页大小校验必须是128的倍数且至少128registry.py依据_use_jenga_kv_cache决策走 Jenga 还是传统路径。6.2 Jenga 缓存与回退开关Jenga 是仓库中新一代的 KV 缓存管理器JengaKVCacheManager见 max/python/max/pipelines/kv_cache/paged_kv_cache/jenga_cache_manager.py其核心思想是用单一可互换内存板memory slab按 Jenga 两级 huge/little 页分配策略划分使多种 KV 类型共享同一物理内存。当前 Jenga 仅在模型名匹配临时白名单含llama、gemma、gpt-oss、olmo、deepseek、kimi等子串时启用registry.py且存在三类回退到传统缓存的场景用户显式设置MODULAR_USE_LEGACY_KV_CACHE1、启用了分离式推理DI与 Jenga 不兼容、或启用了 dKV 连接器亦不兼容此时源码提示设置MODULAR_USE_LEGACY_KV_CACHE1。源码注释明确标注这是Jenga 切换期的临时开关过渡完成后删除。传统路径下load_kv_manager调用compute_num_device_blocks计算设备块池总页数并传入require_max_seq_len_fitsTrue强制执行单请求最长序列必须放得下的启动期校验registry.py随后构造PagedKVCacheManager。6.3 cache_dtype_for_encoding量化编码到缓存 dtype 的映射cache_dtype_for_encoding(quantization_encoding, kv_cache_format)config.py决定 KV 缓存的存储 dtype规则为显式kv_cache_format覆盖优先否则按权重量化编码推导全部未设置时回退float32。显式覆盖映射_KV_CACHE_FORMAT_TO_DTYPEconfig.pykv_cache_formatdtypefloat32DType.float32bfloat16DType.bfloat16float8_e4m3fnDType.float8_e4m3fn量化编码推导映射_ENCODING_TO_KV_CACHE_DTYPEconfig.py——量化权重格式通常让缓存保持计算 dtypebf16/f32而非权重 dtype量化编码KV 缓存 dtypefloat32float32float16float16bfloat16bfloat16float8_e4m3fnbfloat16float6_e2m3fnbfloat16float4_e2m1fnx2bfloat16q4_kfloat32q4_0float32q6_kfloat32gptqbfloat16无法识别的覆盖串或编码会抛出带明确支持值列表的ValueError。该映射的存在意味着例如部署 FP8 量化权重模型时KV 缓存默认以 BF16 存储以保证精度只有显式传--kv-cache-format float8_e4m3fn才会切换到 FP8 缓存这一取舍在架构层也有印证见 max/python/max/pipelines/architectures/qwen3_5/qwen3_5.py 与 max/python/max/pipelines/architectures/laguna/model_config.py 的相关注释。七、与基准测试工具的联动KV 缓存配置与max benchmark之间存在一个容易踩坑的联动点基准工具的--kv-block-size参数默认128用于每轮缓存驻留指标文档明确建议其与服务器的--kv-cache-page-size保持一致否则驻留指标会失真虽然不匹配本身不影响基准运行见 max/python/docs/cli/benchmark.rst。这也从侧面印证了kv_cache_page_size是贯穿服务端与压测工具的统一概念。总结max.pipelines.kv_cache构成了一套完整的三层 KV 缓存管理栈规划层memory_planner.py通过ModelConfig协议解耦PagedMemoryPlanner提供标准预算估算with_activation_reservation满足 VLM 等特殊预留需求管理层paged_kv_cache/PagedKVCacheManager负责分页分配与前缀缓存KVConnector协议把主机/磁盘/dKV 外部层抽象为可重叠的加载/卸载传输Jenga 管理器则在切换期提供共享内存板的新实现传输层transfer_engine.py基于 NIXL 的KVTransferEngine承载 DP/TP 场景下的 KV 数据流动配合available_port、UCCL 环境变量等工具完成跨副本协调。无论是通过--device-memory-utilization调控显存头寸、用--kv-connector-config {type: rust_tiered}开启分层离载、还是以--kv-cache-format切换缓存精度这套模块都为 MAX 推理服务提供了可预测、可观测、可调优的缓存基础。【免费下载链接】mojoThe Modular Platform (includes MAX Mojo)项目地址: https://gitcode.com/GitHub_Trending/mo/mojo创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考 MAX 流水线 KV 缓存管理详解：从内存规划、分页缓存到跨设备传输引擎
MAX 流水线 KV 缓存管理详解从内存规划、分页缓存到跨设备传输引擎【免费下载链接】mojoThe Modular Platform (includes MAX Mojo)项目地址: https://gitcode.com/GitHub_Trending/mo/mojo导读max.pipelines.kv_cache是 Modular MAX 平台中负责大语言模型 KVKey-Value缓存全生命周期管理的核心 Python 模块它既决定缓存能占多少显存、按什么页大小组织也负责前缀命中、淘汰换入换出、跨设备传输等运行时行为。本文以 max/python/docs/pipelines.kv_cache.rst 的 API 结构为主线结合仓库源码逐一剖析其内存规划MemoryPlanner、配置体系KVCacheConfig / KVConnectorConfig、分页缓存管理器PagedKVCacheManager、连接器KVConnector与传输引擎KVTransferEngine的设计与用法。读完本文你将能理解 MAX 推理服务中 KV 缓存的显存预算从何而来、如何通过命令行参数精确调控缓存精度与分层离载offload、以及分布式场景下 KV 数据如何跨副本流动。模块的公开入口定义在 max/python/max/pipelines/kv_cache/init.py对外导出KVCacheConfig、KVConnectorConfig、MemoryPlanner、PagedMemoryPlanner、PagedKVCacheManager、KVTransferEngine等全部核心类型并额外导出了max.nn.kv_cache中的KVCacheGroupId。下文按原文档的六大部分逐一展开。一、内存规划KV 缓存的显存预算从哪来自回归文本生成模型的标准推理路径中KV 缓存是最大的动态显存消耗项。max.pipelines.kv_cache的Memory planning部分提供了一套可插拔的预算估算体系实现在 max/python/max/pipelines/kv_cache/memory_planner.py。1.1 结构化协议ModelConfig 与 ModelConfigWithKVCache这两个类并非具体实现而是typing.Protocol带runtime_checkable装饰器可被isinstance运行时检查ModelConfigmemory_planner.py仅要求暴露devices属性模型运行所在的设备列表是最小的模型配置结构契约ModelConfigWithKVCachememory_planner.py在ModelConfig基础上追加get_kv_params() - KVCacheParamInterface使规划器能以类型化接口而非getattr获取缓存参数。这种协议化设计让规划器只依赖配置对象的最小表面不绑定具体架构从而可以独立于整个 pipeline 栈使用。1.2 基类 MemoryPlannerMemoryPlannermemory_planner.py仅由一个ModelConfig构造而非完整PipelineConfig并为所有估算方法提供默认实现子类按需覆写。其核心方法包括estimate_weights_size(pipeline_config)估算模型权重占用字节数默认委托给pipeline_config.model.weights_size()estimate_activation_memory(pipeline_config, huggingface_config)估算权重之外的激活内存临时中间张量默认返回0需要大临时缓冲的架构如 MLA 上投影、专家并行路由缓冲在此覆写infer_max_batch_size(pipeline_config, devices, weights_size)当用户未显式设置max_batch_size时推断架构相关默认值默认返回None交给框架级推断estimate_signal_buffer_memory(...)估算跨设备 P2P 集合通信所需的信号缓冲内存默认单设备返回0对于无条件执行 allreduce 的模型如使用VocabParallelEmbedding的 VLM可通过类级_always_signal_buffers True让单设备也预留Signals.NUM_BYTES。1.3 标准实现 PagedMemoryPlannerPagedMemoryPlannermemory_planner.py是使用分页 KV 缓存的文本生成模型的标准规划器。构造时要求配置实现ModelConfigWithKVCache否则抛出TypeError。对需要预留固定显存空间的架构典型如 VLM 要为视觉处理留出余量源码明确建议不要手写自定义子类而是用工厂方法with_activation_reservation直接生成预配置子类memory_plannerPagedMemoryPlanner.with_activation_reservation( 15 * 1024**3 # 15 GiB )该方法memory_planner.py会动态创建名为PagedMemoryPlanner(15GiB)的子类使其estimate_activation_memory返回固定预留值对于无条件 allreduce 的模型还可追加always_signal_buffersTrue让单 GPU 场景也预留信号缓冲memory_plannerPagedMemoryPlanner.with_activation_reservation( 15 * 1024**3, always_signal_buffersTrue )二、配置体系KVCacheConfig 与 KVConnectorConfig配置层实现在 max/python/max/pipelines/kv_cache/config.py两个 Pydantic 配置类均继承自max.config.ConfigFileModel可直接从命令行、YAML/JSON 配置文件加载。2.1 KVCacheConfig分页缓存的核心参数KVCacheConfigconfig.py的字段即max serve/max generate等命令的--kv-cache-*系列参数来源字段及默认值如下表字段默认值含义kv_cache_page_size128分页 KV 缓存中单个 page 容纳的 token 数enable_prefix_cachingTrue是否启用前缀缓存enable_dp_cross_replica_prefix_copyTrue数据并行DP场景下是否允许把另一副本 GPU 上的前缀缓存块通过设备到设备D2D拷贝搬来命中关闭后跨副本复用只能走共享主机/磁盘层连接器或重算kv_connector_config默认null连接器KV 缓存连接器配置见下文 2.2device_memory_utilization0.9进程应占用的可用显存比例剩余头寸留给 KV 缓存kv_cache_workspace total_free_memory * device_memory_utilization - model_weights_sizekv_cache_formatNone显式覆盖缓存数据类型可选float32、bfloat16、float8_e4m3fnindexer_kv_cache_formatNone独立覆盖 MiniMax 稀疏索引器IndexK缓存 dtype可选bfloat16、float8_e4m3fn默认保持 BF16使--kv-cache-formatfloat8_e4m3fn仍表示主 GQA FP8 索引器 BF16state_pool_dtypeNone覆盖混合模型循环状态池SSM/线性注意力卷积与循环状态的存储 dtype可选bfloat16、float32FP32 使推测解码严格跟随非推测状态轨迹kv_cache_hash_algoahash64缓存块身份哈希算法可选ahash64快速非密码学、sha256256 位密码学哈希、sha256_64SHA-256 链截断到 64 位以兼容协议kv_cache_hash_seedNone可选的 64 位十六进制字符串32 字节集群级种子省略时sha256/sha256_64启动时随机生成ahash64不生成避免影响既有部署其中device_memory_utilization直接出现在max serve的官方示例中见 max/python/docs/cli/serve.rstmax serve \ --model google/gemma-3-12b-it \ --devices gpu:0 \ --max-batch-size 8 \ --device-memory-utilization 0.9KVCacheConfig.to_params()config.py负责把该配置转换为max.nn.kv_cache.cache_params.KVCacheParams并依据is_mla标志自动选择子类MLA 模型生成MLAKVCacheParams此时必须提供num_q_heads否则生成MHAKVCacheParams。架构可通过page_size参数传入其内核强制的页大小下限而不必修改共享配置kv_cache_hash_seed则通过resolve_kv_hash_seed辅助函数解析。2.2 KVConnectorConfig外部缓存层配置KVConnectorConfigconfig.py采用类型与设置同体的设计连接器类型与其参数作为一个对象整体配置例如--kv-connector-config {type: rust_tiered}公共字段是强类型的连接器特有字段通过extraallow透传经model_extra访问。字段如下字段默认值含义typenull连接器类型见下方KVConnectorType枚举默认无外部缓存host_offload_max_gbNone主机内存离载上限GiB未设置时按设备页池的 1.5 倍自动估算disk_offload_dirNone磁盘离载目录tiered类型必填disk_offload_max_gbNone磁盘离载上限GiB 0未设置时为设备页池 2 倍0表示去掉磁盘层纯主机连接器num_disk_workers32磁盘 I/O 工作线程数 0负载下越高排空磁盘操作队列越快约 32 后收益递减block_store_endpointNone同机 dKV 服务端点支持 IPCipc:///path或 TCPtcp://host:portdkv类型必填连接器类型KVConnectorType定义于 max/python/max/nn/kv_cache/cache_params.py四个取值语义明确null无设备外后备存储页仅驻留设备tiered把被逐出的页分层放到主机内存与磁盘已废弃作为rust_tiered的向后兼容别名rust_tiered唯一的主机/磁盘分层实现由 Rustkv_tier_connector扩展支撑拷贝与磁盘 I/O 运行在 Rust 线程无 GIL 争用并通过异步传输句柄与 GPU 计算重叠要求enable_prefix_caching、host_offload_max_gb、disk_offload_dir且仅支持 CUDA/HIP 设备dkv通过分布式 KV 块存储路由页要求block_store_endpoint。值得注意的实现细节源码注释明确说明远程 dKV 端点在运行时由 Orchestrator 的每请求dkv_cache_hint发现而非静态配置——连接器在 Rust 侧解析 hint 并自行拨号每个命名的实例。三、分页缓存管理器PagedKVCacheManager 及其配套3.1 PagedKVCacheManager核心管理器PagedKVCacheManagermax/python/max/pipelines/kv_cache/paged_kv_cache/cache_manager.py是支持数据并行DP与张量并行TP的分页 KV 缓存管理器其 docstring 提供了完整的端到端用法示例cache_manager.pyimport numpy as np from max.driver import CPU from max.dtype import DType from max.engine import InferenceSession from max.graph import DeviceRef from max.nn.kv_cache import MHAKVCacheParams from max.pipelines.context import TextContext, TokenBuffer from max.pipelines.kv_cache import PagedKVCacheManager from max.pipelines.modeling.types import RequestID params MHAKVCacheParams( dtypeDType.float32, n_kv_heads8, head_dim128, num_layers2, page_size128, devices[DeviceRef.CPU()], ) kv_manager PagedKVCacheManager( paramsparams, sessionInferenceSession(devices[CPU()]), total_num_pages8, max_batch_size4, ) def make_context() - TextContext: tokens np.array([1, 2, 3, 4], dtypenp.int64) return TextContext( request_idRequestID(), max_length1000, tokensTokenBuffer(tokens), ) ctx1 make_context() ctx2 make_context() # Allocate metadata for requests in batch kv_manager.claim(ctx1) kv_manager.claim(ctx2) # Allocate blocks for these requests kv_manager.alloc(ctx1) kv_manager.alloc(ctx2) # Get KVCache inputs to feed to graph kv_cache_inputs kv_manager.runtime_inputs([[ctx1, ctx2]]) # Run model... # Update requests with newly generated tokens ctx1.update(42) ctx2.update(42) # Commit newly written blocks to prefix cache kv_manager.step(ctx1)该示例完整呈现了管理器的请求生命周期claim为批内请求分配元数据→alloc分配块→runtime_inputs取图输入→ 前向 →update追加新 token→step把新写块提交到前缀缓存。在设备侧管理器维护每设备一张 LUTlookup table表与缓存长度缓冲LUT 形状为(max_batch_size, padded_lut_cols(max_total_num_pages))、uint32类型内维做了填充以保证 SIMDpopulate在任意合法first_lut_idx之后都能安全加载至多 16 个连续uint32而不越界cache_manager.py。3.2 DummyKVCache测试与禁用场景的无操作实现DummyKVCachemax/python/max/pipelines/kv_cache/paged_kv_cache/dummy_cache_manager.py继承PagedKVCacheManager是缓存被禁用或处于测试时的空操作实现claim/step/release均为 no-opalloc返回一个已完成传输CompletedTransfer.load()contains对任何请求都返回Trueblock_count/host_byte_count/disk_byte_count返回恒定的1 空闲/1 总计占位值指标为空。3.3 计数快照BlockCount 与 ByteCount两者定义于 max/python/max/pipelines/kv_cache/kv_connector.pyBlockCountL35-L67设备G0块池占用快照字段free/total派生出used、used_pct、free_pct属性total 0时百分比返回0而非除零ByteCountL70-L103外部缓存层主机/磁盘的字节占用快照。用字节而非块是因为连接器的主机/磁盘层是操作员按字节host_offload_max_gb、disk_offload_max_gb设置的预算且其块宽度不必与设备一致——MLA 复制单元在主机上只存一份、加载时广播回去。3.4 InsufficientBlocksError分配失败的显式信号InsufficientBlocksErrormax/python/max/pipelines/kv_cache/paged_kv_cache/block_utils.py是空闲块不足无法满足分配时抛出的异常。它不只是运行时错误在 registry.py 的load_kv_manager中会预先校验单请求在max_seq_len下必须能放入设备块池否则请求无法被抢占无物可逐出运行时必然以InsufficientBlocksError崩溃——因此选择在启动阶段直接失败而非留到运行时。四、KV 连接器外部缓存层的加载/卸载协议KVConnector协议kv_connector.py是管理器拥有设备张量、块分配与设备侧前缀缓存连接器负责外部层主机内存、磁盘、dKV 等读写的分工契约。关键方法load(block_ids, block_hashes, replica_idx, hint)把外部缓存数据加载进设备块H2D 拷贝返回KVConnectorTransferoffload(block_ids, block_hashes, replica_idx)把设备块卸载到外部缓存D2H 拷贝touch(block_hashes, replica_idx)刷新外部层对被设备G0侧前缀缓存命中所服务块的 LRU 新鲜度——尽力而为、即发即忘避免设备上仍热门的块被外部层误逐出契约要求从真正根节点开始、按序列顺序传完整集合count_cached_prefix(block_hashes)按前缀顺序统计外部层连续驻留的前导块数返回(num_host_blocks, num_disk_blocks)必须是严格只读操作wait_for_loads()/wait_for_offloads()前向前后对已发布传输的排序/落定当前仅 dKV 连接器使用其余为 no-op。跨该协议的所有块哈希均为规范字节形式ahash64系含sha256_64为 8 字节大端完整 SHA-256 为 32 字节。4.1 传输句柄模型KVConnectorTransfer 与 CompletedTransferKVConnectorTransferkv_connector.py是连接器传输加载/卸载的句柄其核心价值是让管理器把传输与 GPU 计算重叠传输持有其触及的g0设备块 id管理器在is_complete()返回True前保持这些块被钉住随后恰好解除一次。is_complete必须是廉价、无副作用的轮询原子量或cudaEventQuery式检查。协议支持两种完成模型同步 / 流有序连接器dKV拷贝在转发流上或 GPU 有序地排在其前直接返回CompletedTransferkv_connector.pyis_complete立即为True管理器立刻提交复用的前缀、请求不脱离批次异步连接器Rustrust_tiered拷贝在独立拷贝引擎上执行句柄的is_complete仅在拷贝落地后翻转管理器钉住块、推迟提交并暂时把请求隔离出批次GPU 在此期间运行其他就绪工作。五、传输引擎KVTransferEngine 与跨副本数据流动5.1 传输引擎的角色KVTransferEnginemax/python/max/pipelines/kv_cache/paged_kv_cache/transfer_engine.py是支持 DP 与 TP 的 KV 缓存传输引擎基于 NIXLNexus Interconnect eXchange Layer构建是底层TransferEngine之上的薄封装新增能力包括张量形状与设备类型校验、从产者编写的 NIXL group 构建、按组的 TP 复制replicated_per_group、以及from_paged_kv_cache()便捷构造器。其构造接收按 DP 副本索引的外层列表memory[replica][group]每个条目是一个KVCacheMemory对应一个逻辑(child, kind)张量携带其全部 TP 分片视图来自KVCacheBuffer.to_memory()。构造函数会强校验结构不变量transfer_engine.py每个副本必须至少含一个 group、所有副本 group 数量一致、所有 group 的total_num_pages一致。文档特别强调单个 TransferEngine 不是线程安全的它面向 MAX 的单线程调度器使用仅在不同线程/进程之间彼此通信。5.2 元数据与请求数据结构KVTransferEngineMetadatatransfer_engine.py在仅含传输信息的TransferEngineMetadata之上扩展 KV 缓存/拓扑字段——total_num_pages每个张量的总页数、bytes_per_page每页字节数、bytes_per_group每个 NIXL group 的每页字节数首项为主 group后续对应推测解码中的草稿 KV 等额外 group、replicated_per_group每 group 的 TP 复制标志True表示 MLA 式跨 TP 分片完全相同复制False表示分片。该对象可安全跨线程/进程传输是线上唯一的复制信息源。TransferReqDatatransfer_engine.py基于msgspec.Struct的传输请求元数据字段包括源/目标引擎名src_name/dst_name、传输名、每 TP 分片一个的transfer_ids、源/目标索引长度可与transfer_ids不同、副本索引、is_read是否为目标发起的 READ 拉取、tp_shard_count与local_shards_used等同样可安全跨线程/进程传输。5.3 环境变量与端口工具传输引擎读取两个关键环境变量族MODULAR_NIXL_TRANSFER_BACKENDNIXL 传输后端类型默认ucx由统一校验器validate_nixl_backend检查见 max/python/max/pipelines/kv_cache/_nixl_backend.py 中的NIXL_BACKEND_ENV_VARdKV 连接器读取同一变量但为必填无自动回退。UCCL 相关变量若未设置UCCL_SOCKET_IFNAME/NCCL_SOCKET_IFNAME会默认把 UCCL 的带外引导网卡钉到主机的默认路由 NIC控制网络避免自动选择落在仅承载 RoCE 的网卡上导致对端拨号挂起同时默认设置UCCL_P2P_TRANSPORTrdma与UCCL_P2P_DISABLE_IPC1规避同机intranode收发对在 IPC 传输下的死锁transfer_engine.py。available_port(start_port8000, end_port9000, max_attempts100)transfer_engine.py在指定闭区间内随机探测可用 TCP 端口绑定SO_REUSEADDR以避免 TIME_WAIT 问题max_attempts次内未找到则抛RuntimeError。六、工厂函数与工具函数6.1 load_kv_manager按参数装载缓存管理器load_kv_manager(params, max_batch_size, max_seq_len, session, available_cache_memory, is_di_enabled, model_name)registry.py是缓存管理器的统一装载入口接收KVCacheParams单缓存或MultiKVCacheParams多缓存返回的单个管理器原生处理全部缓存。其装载逻辑层层把关若params是MagicMock测试桩直接返回MagicMock编译专用模式虚拟设备模式下返回Mock避免 GPU 显存分配available_cache_memory/max_batch_size未在内存估算阶段设置、或max_batch_size 0时抛ValueError页大小校验必须是128的倍数且至少128registry.py依据_use_jenga_kv_cache决策走 Jenga 还是传统路径。6.2 Jenga 缓存与回退开关Jenga 是仓库中新一代的 KV 缓存管理器JengaKVCacheManager见 max/python/max/pipelines/kv_cache/paged_kv_cache/jenga_cache_manager.py其核心思想是用单一可互换内存板memory slab按 Jenga 两级 huge/little 页分配策略划分使多种 KV 类型共享同一物理内存。当前 Jenga 仅在模型名匹配临时白名单含llama、gemma、gpt-oss、olmo、deepseek、kimi等子串时启用registry.py且存在三类回退到传统缓存的场景用户显式设置MODULAR_USE_LEGACY_KV_CACHE1、启用了分离式推理DI与 Jenga 不兼容、或启用了 dKV 连接器亦不兼容此时源码提示设置MODULAR_USE_LEGACY_KV_CACHE1。源码注释明确标注这是Jenga 切换期的临时开关过渡完成后删除。传统路径下load_kv_manager调用compute_num_device_blocks计算设备块池总页数并传入require_max_seq_len_fitsTrue强制执行单请求最长序列必须放得下的启动期校验registry.py随后构造PagedKVCacheManager。6.3 cache_dtype_for_encoding量化编码到缓存 dtype 的映射cache_dtype_for_encoding(quantization_encoding, kv_cache_format)config.py决定 KV 缓存的存储 dtype规则为显式kv_cache_format覆盖优先否则按权重量化编码推导全部未设置时回退float32。显式覆盖映射_KV_CACHE_FORMAT_TO_DTYPEconfig.pykv_cache_formatdtypefloat32DType.float32bfloat16DType.bfloat16float8_e4m3fnDType.float8_e4m3fn量化编码推导映射_ENCODING_TO_KV_CACHE_DTYPEconfig.py——量化权重格式通常让缓存保持计算 dtypebf16/f32而非权重 dtype量化编码KV 缓存 dtypefloat32float32float16float16bfloat16bfloat16float8_e4m3fnbfloat16float6_e2m3fnbfloat16float4_e2m1fnx2bfloat16q4_kfloat32q4_0float32q6_kfloat32gptqbfloat16无法识别的覆盖串或编码会抛出带明确支持值列表的ValueError。该映射的存在意味着例如部署 FP8 量化权重模型时KV 缓存默认以 BF16 存储以保证精度只有显式传--kv-cache-format float8_e4m3fn才会切换到 FP8 缓存这一取舍在架构层也有印证见 max/python/max/pipelines/architectures/qwen3_5/qwen3_5.py 与 max/python/max/pipelines/architectures/laguna/model_config.py 的相关注释。七、与基准测试工具的联动KV 缓存配置与max benchmark之间存在一个容易踩坑的联动点基准工具的--kv-block-size参数默认128用于每轮缓存驻留指标文档明确建议其与服务器的--kv-cache-page-size保持一致否则驻留指标会失真虽然不匹配本身不影响基准运行见 max/python/docs/cli/benchmark.rst。这也从侧面印证了kv_cache_page_size是贯穿服务端与压测工具的统一概念。总结max.pipelines.kv_cache构成了一套完整的三层 KV 缓存管理栈规划层memory_planner.py通过ModelConfig协议解耦PagedMemoryPlanner提供标准预算估算with_activation_reservation满足 VLM 等特殊预留需求管理层paged_kv_cache/PagedKVCacheManager负责分页分配与前缀缓存KVConnector协议把主机/磁盘/dKV 外部层抽象为可重叠的加载/卸载传输Jenga 管理器则在切换期提供共享内存板的新实现传输层transfer_engine.py基于 NIXL 的KVTransferEngine承载 DP/TP 场景下的 KV 数据流动配合available_port、UCCL 环境变量等工具完成跨副本协调。无论是通过--device-memory-utilization调控显存头寸、用--kv-connector-config {type: rust_tiered}开启分层离载、还是以--kv-cache-format切换缓存精度这套模块都为 MAX 推理服务提供了可预测、可观测、可调优的缓存基础。【免费下载链接】mojoThe Modular Platform (includes MAX Mojo)项目地址: https://gitcode.com/GitHub_Trending/mo/mojo创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考