执行摘要

• 一句话：CUDA图捕获后动态调整KV缓存池大小，减少内存浪费
• 推荐动作：推荐精读本PR，涉及底层CUDA VMM实现和KV池动态调整的设计权衡。对于大规模推理服务，可评估启用以提升内存效率。注意当前opt-in状态和HND槽计算争议，建议在非默认路径下充分验证后再默认开启。

功能与动机

在--mem-fraction-static启发式中，KV池大小在CUDA图捕获前确定，图形内存的占用是猜测值，在大配置下偏差可达~10 GB/GPU，导致部分配置浪费内存，部分配置面临OOM。PR body明确指出：'the KV pool is sized before CUDA-graph capture, so graph memory is a heuristic guess baked into mem_fraction_static. On large configs the guess can be off by ~10 GB/GPU; the padded activation reserve silently absorbs the difference, wasting memory on some configs and risking OOM on others.'因此需要一种机制，在捕获图后根据实际空闲内存精确确定KV池大小。

实现拆解

1. 新增CUDA虚拟内存管理arena (kv_vmm_backing.py)：创建KvVmmArena类，通过CUDA VMM接口预留大块虚拟地址空间（默认256 GiB），不分配物理内存。提供bump分配器（内联C++桩），保证指针按granularity对齐，支持torch.cuda.CUDAPluggableAllocator。每个实例加载独立.so，避免多进程符号冲突。
2. 修改mem_fraction计算路径 (server_args.py)：新增post_capture_kv_sizing_planned()方法，当环境变量启用且满足条件（CUDA、非MLA、非memory_saver等）时，计算mem_fraction_static跳过图占用和激活预留，仅保留基础元数据和并行开销。
3. KV池关联VMM arena (memory_pool.py)：在KVCache及其子类中新增post_capture_active标志和_post_capture_owner引用。启用时_create_buffers使用KvVmmArena创建张量（仅虚拟地址）。新增KvBufferDesc辅助类精确描述缓冲区字节跨度，支持动态计算备份大小。
4. 新增图捕获后调整池大小的方法 (model_runner_kv_cache_mixin.py)：在ModelRunnerKVCacheMixin中添加post_capture_resize_kv_pool()，在CUDA图捕获完成后调用。测量当前空闲GPU内存，减去必要的headroom，计算KV池可用最大token数，调用_finalize_backing物理备份VMA范围。同时修改_profile_available_bytes，在启用时正确处理mamba预捕获预留。
5. 配套修改分配器和SWA池：更新allocator/base.py和allocator/swa.py添加resize方法；swa_memory_pool.py添加finalize_backing协调子池备份。
6. 端到端测试 (test_post_capture_kv_sizing.py)：通过环境变量启动服务器，断言日志含有'Post-capture KV sizing'、API返回max_total_num_tokens>0、GSM8K准确率≥0.80。注册到base-b CI。

关键文件：

• python/sglang/srt/mem_cache/kv_vmm_backing.py（模块 VMM 层；类别 source；类型 dependency-wiring；符号 _driver, _check, align_up, query_granularity）: 新增文件，实现CUDA虚拟内存管理arena，是post-capture sizing的核心基础。
• python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py（模块 KV编排；类别 source；类型 data-contract；符号 post_capture_kv_active, post_capture_resize_kv_pool, _resolve_memory_pool_config, _config_from_budget）: 核心修改文件，添加post_capture_resize_kv_pool和post_capture_kv_active属性，是后调整的编排中心。
• python/sglang/srt/mem_cache/memory_pool.py（模块 内存池；类别 source；类型 core-logic；符号 KvBufferDesc, init, _rows, reserved_span_bytes）: 修改了KVCache基类和子类，添加post_capture_active标志和VMM备份支持，新增KvBufferDesc辅助类。
• python/sglang/srt/server_args.py（模块 配置；类别 source；类型 configuration；符号 post_capture_kv_sizing_planned, mamba_pre_capture_reserve_mb, reserve_for_graph_mb, reserve_for_deepep_a2a_mb）: 修改mem_fraction计算路径，新增post_capture_kv_sizing_planned方法决定何时跳过图预留。
• test/registered/mem_cache/test_post_capture_kv_sizing.py（模块 测试；类别 test；类型 test-coverage；符号 TestPostCaptureKVSizing, setUpClass, tearDownClass, _server_logs）: 新增端到端测试，验证post-capture sizing路径执行、API响应和GSM8K准确率。

关键符号：KvVmmArena.init, post_capture_resize_kv_pool, _profile_available_bytes, post_capture_kv_sizing_planned, KvBufferDesc, finalize_backing, KVCache._create_buffers, SWAKVPool.finalize_backing

关键源码片段

python/sglang/srt/mem_cache/kv_vmm_backing.py

新增文件，实现CUDA虚拟内存管理arena，是post-capture sizing的核心基础。

# 文件 python/sglang/srt/mem_cache/kv_vmm_backing.py ( 新增 )
​
from __future__ import annotations
import ctypes, logging, os, tempfile
from math import prod
from typing import TYPE_CHECKING, List, Optional, Sequence
import torch, torch.utils.cpp_extension
from torch.cuda.memory import CUDAPluggableAllocator
​
logger = logging.getLogger(__name__)
_drv = None
​
def _driver():
    """延迟加载 CUDA driver bindings."""
    global _drv
    if _drv is None:
        from cuda.bindings import driver
        _drv = driver
    return _drv
​
def _check(result, label: str):
    """包装 CUDA API 调用，失败时抛出 RuntimeError."""
    drv = _driver()
    err = result[0] if isinstance(result, tuple) else result
    if err != drv.CUresult.CUDA_SUCCESS:
        raise RuntimeError(f"{label} failed: {err}")
    return result[1] if isinstance(result, tuple) and len(result) > 1 else None
​
def _stub_source(sfx: str) -> str:
    """
    生成 C++ 桩代码，实现线程安全的 bump 分配器。
    符号按后缀区分，每个 arena 实例加载独立 .so，避免多进程/多实例冲突。
    """
    return f"""
#include <cstddef>
#include <cstdint>
#include <mutex>
extern "C" {{
static uintptr_t g_base = 0;
static size_t g_cursor = 0;
static size_t g_reserved = 0;
static size_t g_align = 512;
static std::mutex g_mu;
void kvarena_set_base_{sfx}(uintptr_t b) {{ g_base=b; g_cursor=0; }}
void kvarena_set_reserved_{sfx}(size_t r) {{ g_reserved=r; }}
void* kvarena_malloc_{sfx}(size_t size, int device, void* stream) {{
    size_t need = g_cursor + align_up(size, g_align);
    if (need > g_reserved) return 0;
    void* p = reinterpret_cast<void*>(g_base + g_cursor);
    g_cursor = need;
    return p;
}}
void kvarena_free_{sfx}(void* ptr, size_t size, int device, void* stream) {{}}
}}
"""
​
class KvVmmArena:
    """单个设备上的 CUDA 虚拟内存预留，对外暴露为 torch.cuda.MemPool."""
    _instance_count = 0
​
    def __init__(self, device_id: int, reserve_bytes: int = 256 * 1024**3):
        self.device_id = int(device_id)
        # 使用 PID 避免多进程编译 .so 冲突
        self._sfx = f"{KvVmmArena._instance_count}_{os.getpid()}"
        KvVmmArena._instance_count += 1
        drv = _driver()
        with torch.cuda.device(self.device_id):
            _check(drv.cuInit(0), "cuInit")
            self._prop = drv.CUmemAllocationProp()
            self._prop.type = drv.CUmemAllocationType.CU_MEM_ALLOCATION_TYPE_PINNED
            self._prop.location.type = drv.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
            self._prop.location.id = self.device_id
            self.granularity = query_granularity(self.device_id)
            self._access = drv.CUmemAccessDesc()
            self._access.location.type = drv.CUmemLocationType.CU_MEM_LOCATION_TYPE_DEVICE
            self._access.location.id = self.device_id
            self._access.flags = drv.CUmemAccess_flags.CU_MEM_ACCESS_FLAGS_PROT_READWRITE
            self.reserved = self._align(reserve_bytes)
            # 预留虚拟地址空间
            self.base = int(
                _check(
                    drv.cuMemAddressReserve(self.reserved, self.granularity, 0, 0),
                    "cuMemAddressReserve"
                )
            )
            # 编译 C++ 桩并设置基址和上限 ( 省略具体编译步骤 )

python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py

核心修改文件，添加post_capture_resize_kv_pool和post_capture_kv_active属性，是后调整的编排中心。

# 文件 python/sglang/srt/model_executor/model_runner_kv_cache_mixin.py ( 修改 )
​
class ModelRunnerKVCacheMixin:
​
    @property
    def post_capture_kv_active(self: ModelRunner) -> bool:
        """判断当前是否启用 post-capture KV sizing (在 draft 上禁用)."""
        return (
            self.server_args.post_capture_kv_sizing_planned()
            and current_platform.is_cuda()
            and not self.is_draft_worker
        )
​
    def post_capture_resize_kv_pool(self: ModelRunner) -> None:
        """
        在 CUDA 图捕获完成后调用，根据实际空闲内存重新确定 KV 池大小并物理备份。
        """
        pool = self.token_to_kv_pool
        torch.cuda.synchronize()
        free_gb = get_available_gpu_memory(
            self.device, self.gpu_id,
            distributed=get_world_group().world_size > 1,
            cpu_group=get_world_group().cpu_group,
        )
        headroom_gb = self.pre_model_load_memory * (1 - self.mem_fraction_static)
        decode_cfg = self.server_args.cuda_graph_config.decode
        eager_gap = (
            self.server_args.disaggregation_mode != "prefill"
            and decode_cfg.backend != Backend.DISABLED
            and decode_cfg.max_bs < self.max_running_requests
        )
        if eager_gap or self.mambaish_config is not None:
            headroom_gb = max(
                headroom_gb,
                self.server_args.mamba_pre_capture_reserve_mb(
                    get_device_memory_capacity(self.device)
                ) / 1024,
            )
        pool_gb = max(free_gb - headroom_gb, 0)
        # 将可用 GB 转换为 token 数 ( 具体计算略 )
        max_total_num_tokens = int(pool_gb * 1024**3 / self.pool_token_size_bytes)
        logger.info("Post-capture KV sizing: max_total_num_tokens=%d", max_total_num_tokens)
        if hasattr(pool, "finalize_backing"):
            pool.finalize_backing(max_total_num_tokens)
        self.dynamic_max_total_num_tokens = max_total_num_tokens

评论区精华

• gemini-code-assist[bot]的高优先级问题：
  • cuMemAddressReserve未包_check → 已修复。
  • 多进程库名竞态 → 已添加PID后缀。
  • 单位混淆（字节 vs MB） → 作者澄清无问题。
  • HND分页槽计算错误 → 作者反驳，未完全解决。
  • 释放顺序不安全 → 已调整。
• merrymercy的设计评论：
  • fallback headroom过于宽松 → 后续commit优化为只在必要场景保留。
  • 不要用getattr → 已移除getattr改用显式属性。

• 未解决争议：HND槽计算逻辑仍需进一步验证。

  • cuMemAddressReserve缺少错误检查 (correctness): 已在后续commit中修复，添加_check包装。
  • 多进程环境下库名竞态 (correctness): 已添加os.getpid()到后缀，保证唯一性。
  • get_device_memory_capacity单位混淆 (correctness): 无实际操作问题，但文档需澄清。
  • HND分页KV缓存槽计算错误 (correctness): 未完全达成一致，但代码保持原样。需要进一步验证。
  • 释放顺序不安全 (correctness): 已调整删除顺序，先删张量再释放owner。
  • fallback headroom过于宽松 (design): 后续commit优化了条件，仅在eager gap或mamba场景保留headroom。
  • 不要使用getattr访问post_capture_active (style): 已移除getattr，所有子类显式声明post_capture_active属性。

风险与影响

• 风险：
  • CUDA VMM兼容性：依赖较新驱动和GPU，在旧硬件或非NVIDIA平台上可能启动失败（已有条件检查）。
  • 多进程竞态：虽已加入PID后缀，但多个进程共享临时目录时仍可能冲突（低概率）。
  • 激活headroom逻辑：headroom计算依赖mem_fraction_static准确性，若估算不足可能导致OOM；已有eager decode和mamba保护，但未覆盖全部边缘。
  • HND槽计算争议：若review中问题未被正确修复，可能导致物理内存过度分配或不足，影响大量用户。
  • 测试覆盖不足：仅测试了Llama-3.1-8B单模型，未覆盖MLA、DeepSeek V4、MoE等架构，也未测试多GPU场景。
  • CI负担：测试注册在base-b，每PR执行约4分钟，可能增加CI排队时间。
• 影响：
  • 用户：opt-in特性，默认关闭，对现有用户无影响。启用用户可在大GPU上节省约10 GB内存，或避免OOM。需了解CUDA VMM兼容性。
  • 系统：修改KV池分配核心路径和mem_fraction计算。启用时KV池初始化延迟到图捕获后，物理备份在后期单次完成，可能短暂增大抖动。
  • 团队：新增424行核心代码和测试，需要熟悉CUDA VMM的开发者维护。讨论成果提高了代码质量。
  • 风险标记：CUDA VMM兼容性, 多进程竞态, HND槽计算争议, 激活headroom逻辑, 测试覆盖有限

关联脉络

• 暂无明显关联 PR