系列01：内存管理——当你在管理 ByteBuf 时，你已经在管理 KV Cache 了

这篇讲什么

这是个系列的第一篇。我选"内存管理"作为开篇，因为：

你最熟——Netty ByteBuf、JVM 堆、堆外内存，你每天都在打交道
映射最直接——ByteBuf 池化 = KV Cache Block 管理 = GPU 显存池
代码最少——核心逻辑 100 行就能说清原理

读完这篇，你会理解：

ByteBuf 的引用计数为什么等于 vLLM Block 的 refcount
JVM 的 G1 Region 为什么等于 PagedAttention 的 Block
GPU 显存池为什么是你的 PooledByteBufAllocator 在另一个地址空间的翻版
并且，你可以亲手写一个"Java 版显存管理器"

一、问题本质

内存管理要解决的只有三个问题：

1. 分配——如何快速找到一块可用的内存
2. 回收——用完后如何归还，确保不漏
3. 碎片——如何避免小块空隙越积越多

你在 Java 后端看到过这些方案的演进：

JVM 堆管理：
  Serial GC  → 标记-整理（压缩碎片）
  CMS        → 标记-清除（不压缩，碎片多）
  G1         → Region 分区（固定大小，消除碎片）
  ZGC        → 染色指针 + 读屏障（几乎无停顿）

Netty 堆外内存管理：
  UnpooledByteBufAllocator  → 每次分配新 DirectBuffer
  PooledByteBufAllocator    → 预分配 Arena + Chunk + Page 分三级管理

AI Infra 的显存管理（正在发生和你一模一样的事情）：
  PyTorch 默认（CUDACachingAllocator）→ 有缓存但碎片严重
  vLLM PagedAttention → Region 分块（和 G1 一样的思路）

这不是类比，这是同一个问题在不同约束下的工程演化。

二、Netty 怎么管内存

先看你和 Netty 的交互——这段代码你肯定写过：

// 最常见的写法
ByteBuf buf = PooledByteBufAllocator.DEFAULT.buffer(1024);
// ...
buf.release();

但 PooledByteBufAllocator.DEFAULT.buffer(1024) 背后发生了什么？让我们拆开看。

2.1 三层架构

// 简化后的 Netty PooledByteBufAllocator 核心结构

class PooledByteBufAllocator {
    // Arena = 管理区，每个 Arena 内部有独立的锁
    // 默认 CPU 核数 * 2 个 Arena，减少锁竞争
    final PoolArena[] arenas;

    ByteBuf buffer(int capacity) {
        // 1. 先从 ThreadLocal 缓存拿（零锁）
        PoolThreadCache cache = threadCache.get();
        ByteBuf buf = cache.allocate(capacity);
        if (buf != null) return buf;
        
        // 2. ThreadLocal 没有，从 Arena 分配（上锁）
        PoolArena arena = arenas[randomIndex()];
        return arena.allocate(capacity);
    }
}

class PoolArena {
    // 按大小分层管理，类似 G1 的 Region
    PoolChunkList tinyList;   // < 512 字节
    PoolChunkList smallList;  // < 8K
    PoolChunkList normalList; // < 16M
    // huge: >16M，不缓存，直接分配
    
    ByteBuf allocate(int capacity) {
        // 找到对应大小的 PoolChunk
        // 在 Chunk 里分配一个 page 或 subpage
        // 返回封装后的 ByteBuf（引用计数=1）
    }
}

class PoolChunk {
    // 每块 16MB，分成 2048 个 page（每个 8KB）
    // 用二叉树管理 page 的分配状态
    byte[] memoryMap;  // 2048 个节点的二叉树
    
    long allocate(int size) {
        // 在二叉树中搜索空闲 page
        // 标记为已用
        // 返回偏移量
    }
}

2.2 关键洞察

Netty 做的选择：

预分配大块（Chunk=16MB）——减少系统调用（mmap/malloc 的成本很高）
按大小分层（Tiny/Small/Normal）——减少内部碎片
ThreadLocal 缓存（PoolThreadCache）——分配时不加锁
引用计数（refCnt）——允许多个 ByteBuf 共享同一块内存的不同部分

这些优化，每一个都在 AI Infra 里有对应。

三、vLLM 怎么管显存

3.1 先看问题

大模型推理时，每生成一个 token，就需要缓存这个 token 在所有 Transformer 层中的 Key 和 Value 向量。

一个请求生成 1024 个 token：
  KV Cache = 1024 × 层数 × (key_dim + value_dim) × 精度
  70B 模型，32 层，每层 128 维，FP16：
  = 1024 × 32 × 256 × 2 bytes = 16MB 显存

并发 100 个请求，每个请求 2048 token：
  = 100 × 2048 × 32 × 256 × 2 = 3.2GB 显存

如果每个请求按最大长度预分配连续显存——大部分被浪费了（大部分请求的实际生成长度远小于 max_length）。

这就是 vLLM 要解决的问题。

3.2 vLLM 的方案：PagedAttention

核心思路： 把虚拟内存的分页思想搬到 GPU 显存管理上。

# 这是 vLLM block_manager.py 的核心逻辑，我做了大幅度简化但保留结构

class BlockManager:
    """你的 PooledByteBufAllocator 在 GPU 版本。"""
    
    def __init__(self, total_blocks: int, block_size: int = 16):
        # 预分配所有 block（size=16 tokens）
        # 就像 Netty 的 PoolChunk 预分配 16MB
        self.total_blocks = total_blocks
        
        # 空闲列表（类比 PoolChunk 的空闲 page 二叉树）
        self.free_blocks = list(range(total_blocks))
        
        # 引用计数（类比 ByteBuf.refCnt）
        self.refcounts = [0] * total_blocks
    
    def allocate(self) -> int:
        """分配一个 block。O(1)。"""
        if not self.free_blocks:
            raise RuntimeError("GPU OOM")
        block = self.free_blocks.pop()
        self.refcounts[block] = 1
        return block
    
    def free(self, block: int):
        """释放一个 block。
        和 ByteBuf.release() 完全一样的引用计数语义。"""
        self.refcounts[block] -= 1
        if self.refcounts[block] == 0:
            self.free_blocks.append(block)
    
    def fork(self, block: int) -> int:
        """Copy-on-Write 复制。
        参考 Netty 的 retainedDuplicate()。"""
        self.refcounts[block] += 1  # 共享，不拷贝
        return block                 # 直到写入时才真正复制


class SequenceGroup:
    """一个请求组，维护自己的 "页表"。
    类比：一个进程的虚拟地址空间。"""
    
    def __init__(self, block_manager: BlockManager):
        self.block_manager = block_manager
        self.block_table = []  # 逻辑 block → 物理 block 的映射
        # 这就是操作系统的页表！
    
    def append_token(self):
        """追加一个 token 到当前请求。"""
        # 逻辑上的操作：在当前序列末尾追加
        # 物理上的操作：
        #   如果当前 block 没满 → 使用当前 block
        #   如果当前 block 满了 → 分配一个新 block，加入 block_table
        if self._current_block_is_full():
            new_block = self.block_manager.allocate()
            self.block_table.append(new_block)

3.3 映射表

Netty PooledByteBuf	vLLM BlockManager	作用
`PoolArena`	整个 GPU 显存池	管理域
`PoolChunk`（16MB）	预分配固定大小的 block 集合	大块内存
`Page`（8KB）	Block（16 tokens ≈ 16KB-64KB）	最小分配单元
`PoolThreadCache`	每个 sequence 的 block_table	ThreadLocal 缓存
`subpage`	小于 block 的分配（少见）	小对象分配
`refCnt`	`refcounts[block]`	共享跟踪
`retainedDuplicate()`	`fork()` + COW	Copy-on-Write
`capacity()`	`block_size * len(block_table)`	逻辑容量
碎片化	页表机制消除外部碎片	解决的核心问题

3.4 碎片化的对比

无 PagedAttention（传统做法——连续分配）：
  [请求A的KV Cache: 连续 64 个 block]
  [请求B: 连续 48 个 block]
  [请求C: 连续 72 个 block]
  [请求A完成 → 释放 64 个 block]
  现在空闲区域是：64 block 大小
  但如果新请求需要 80 个 block → 碎片，只能等 GC

有 PagedAttention（分页管理）：
  [Block 0-9: A, Block 10-19: B, Block 20-29: C, ...]
  A 释放后：Block 0-9 立即可以被任何请求使用
  没有碎片问题，因为所有 block 大小相同
  ← 这就是 G1 用固定大小 Region 解决 CMS 碎片问题的同款思路

四、JVM 级别的对应

4.1 G1 GC 和 PagedAttention

JVM 的 G1 和 vLLM 的 PagedAttention 在抽象层面是同一个架构：

// G1 GC（你已经理解的）
//
// 堆被分成 2048 个 Region（每个 1MB，和分代无关）
// Eden/Old/Survivor 只是这些 Region 上的逻辑标签
//
// 对象分配：
//   在当前线程的 ThreadLocal Buffer (TLAB) 里分配
//   用完了 → 申请新的 Region
//
// Region 可以随时被重新标记为不同代
// 消除了 CMS 的分代碎片问题

// PagedAttention（你正在理解的）
//
// 显存被分成 N 个 Block（每个 fixed size）
// 请求的 KV Cache 是这些 Block 上的逻辑连续
//
// Block 分配：
//   从当前 sequence 的 block_table 里找空位
//   用完了 → 申请新的 Block
//
// Block 可以分配给任何 sequence
// 消除了传统的请求级连续分配碎片问题

代码级别也是同构的：

// G1 分配对象（JVM 内部）：
HeapRegion* hr = g1h->allocator()->new_alloc_region(word_size, ...);
// 从空闲链表取一个 Region，标记为 Eden

// PagedAttention 分配 Block（vLLM）：
block = self.free_blocks.pop()
self.refcounts[block] = 1
self.block_table.append(block)

都是"从预分区域里拿一个固定大小的块"。

4.2 ThreadLocal + 引用计数

Java:
  ThreadLocal<TLAB> → 每个线程在堆里有自己的分配缓冲区

ByteBuf:
  PoolThreadCache → 每个 EventLoop 有自己的 ByteBuf 缓存

vLLM:
  block_table → 每个 sequence 有自己的 block 地址映射

都遵循同一个模式：共享资源 + 按线程/请求分片 + 避免全局锁。

五、4 个代码实验

下面的代码你可以直接复制到 IDE 里跑。每个实验说明你把 Java 知识延伸到了 AI Infra 的哪个点。

实验 1：引用计数分配的 ByteBuf 版

public class RefCountedBlock {
    private static final int TOTAL_BLOCKS = 1024;
    
    private final List<Integer> freeList;
    private final int[] refcounts;
    
    public RefCountedBlock() {
        this.freeList = new ArrayList<>();
        this.refcounts = new int[TOTAL_BLOCKS];
        for (int i = 0; i < TOTAL_BLOCKS; i++) {
            freeList.add(i);
            refcounts[i] = 0;
        }
    }
    
    public synchronized int allocate() {
        if (freeList.isEmpty()) throw new RuntimeException("OOM");
        int block = freeList.remove(freeList.size() - 1);
        refcounts[block] = 1;
        System.out.println("Alloc block " + block + " (free: " + freeList.size() + ")");
        return block;
    }
    
    public synchronized void retain(int block) {
        refcounts[block]++;
        System.out.println("Retain block " + block + " (refcnt: " + refcounts[block] + ")");
    }
    
    public synchronized boolean release(int block) {
        refcounts[block]--;
        System.out.println("Release block " + block + " (refcnt: " + refcounts[block] + ")");
        if (refcounts[block] == 0) {
            freeList.add(block);
            System.out.println("  → Block " + block + " returned to pool");
            return true;
        }
        return false;
    }
    
    public static void main(String[] args) {
        var pool = new RefCountedBlock();
        
        // 模拟：请求 A 分配两个 block
        int a1 = pool.allocate();  // refcnt=1
        int a2 = pool.allocate();  // refcnt=1
        
        // 模拟：共享 prefix（Copy-on-Write）
        pool.retain(a1);           // refcnt=2
        
        // 释放请求 A
        pool.release(a1);          // refcnt=1 → 还没释放（有人共享）
        pool.release(a2);          // refcnt=0 → 回收
        pool.release(a1);          // refcnt=0 → 回收
        
        // 输出: 所有 block 都被回收，无泄漏
        System.out.println("Final free blocks: " + pool.freeList.size());
        // → 1024
    }
}

这个实验让你理解 vLLM 的 BlockManager。 你已经在写 Java 版的显存管理器了。

实验 2：页表（Block Table）

public class BlockTable {
    private final RefCountedBlock pool;
    private final List<Integer> blocks = new ArrayList<>();
    
    public BlockTable(RefCountedBlock pool) {
        this.pool = pool;
    }
    
    public void append() {
        blocks.add(pool.allocate());
    }
    
    public int getPhysicalBlock(int logicalIndex) {
        return blocks.get(logicalIndex);
    }
    
    public void free() {
        for (int b : blocks) pool.release(b);
        blocks.clear();
    }
    
    public static void main(String[] args) {
        var pool = new RefCountedBlock();
        var table = new BlockTable(pool);
        
        // 请求 A：需要 5 个 block 的 KV Cache
        for (int i = 0; i < 5; i++) table.append();
        
        // 逻辑 block 2 对应物理 block 2
        System.out.println("Logical block 2 → Physical block " + table.getPhysicalBlock(2));
        
        // 释放
        table.free();
    }
}

这个实验让你理解 PagedAttention 的地址映射。 这是操作系统虚拟内存思想的 GPU 版。

实验 3：碎片化对比

public class FragmentationDemo {
    public static void main(String[] args) {
        // 模拟：连续分配（传统推理框架）
        boolean[] continuous = new boolean[100];
        
        // 分配：请求 A 要 20 块，请求 B 要 30 块
        fill(continuous, 0, 20);   // A
        fill(continuous, 20, 50);  // B
        
        // A 完成释放
        clear(continuous, 0, 20);
        // 空闲：块 0-19
        
        // 新请求 C 需要 25 块
        if (countFree(continuous) >= 25) {
            // 连续分配需要连续的空闲块
            int start = findContinuousFree(continuous, 25);
            System.out.println("连续分配：找到起始位置 " + start);
            // 找不到！因为最大连续只有 20 块
            // → 碎片化
        }
        System.out.println("连续分配：无法分配 25 块（最大连续空闲：" + maxContinuousFree(continuous) + "）");
        
        // === 一次性实验：用固定大小 block 消除碎片 ===
        // BlockManager: 所有 block 一样大，不要求连续
        var pool = new RefCountedBlock();
        var tableA = new BlockTable(pool);
        var tableB = new BlockTable(pool);
        
        for (int i = 0; i < 20; i++) tableA.append();  // A 拿 20 块
        for (int i = 0; i < 30; i++) tableB.append();  // B 拿 30 块
        tableA.free();  // A 释放 20 块 → 回到空闲池
        
        var tableC = new BlockTable(pool);
        for (int i = 0; i < 25; i++) tableC.append();  // C 拿 25 块 → 无碎片，成功！
        System.out.println("分页管理：成功分配 25 块！");
    }
    
    static void fill(boolean[] arr, int from, int to) {
        for (int i = from; i < to; i++) arr[i] = true;
    }
    static void clear(boolean[] arr, int from, int to) {
        for (int i = from; i < to; i++) arr[i] = false;
    }
    static int countFree(boolean[] arr) {
        int c = 0; for (boolean b : arr) if (!b) c++; return c;
    }
    static int findContinuousFree(boolean[] arr, int need) {
        int cur = 0;
        for (int i = 0; i < arr.length; i++) {
            if (!arr[i]) { cur++; if (cur == need) return i - need + 1; }
            else cur = 0;
        }
        return -1;
    }
    static int maxContinuousFree(boolean[] arr) {
        int max = 0, cur = 0;
        for (boolean b : arr) {
            if (!b) { cur++; max = Math.max(max, cur); }
            else cur = 0;
        }
        return max;
    }
}

这个实验让你理解为什么 PagedAttention 能把显存利用率从 40% 提升到 95%。 原因就在这个差异里。

实验 4：前缀共享（Prefix Caching）

public class PrefixCachingDemo {
    // vLLM 中的 Automatic Prefix Caching
    // 核心思想：如果两个请求的前几个 token 相同，共享对应的 KV Cache block
    
    public static void main(String[] args) {
        var pool = new RefCountedBlock();
        
        // 请求 A: "今天天气真不错啊"（5 个 token → 5 个 block）
        var tableA = new BlockTable(pool);
        for (int i = 0; i < 5; i++) tableA.append();
        System.out.println("A 分配了 " + 5 + " 个 block");
        
        // 请求 B: "今天天气真好"（4 个 token，前 3 个和 A 相同）
        // vLLM 检测到 "今天天气" 已经在 A 的 block 0-2 里了
        // block 0-2 共享（refcount++），block 3 新分配
        for (int i = 0; i < 3; i++) {
            pool.retain(tableA.getPhysicalBlock(i));  // 共享 A 的前 3 个 block
        }
        var tableB = new BlockTable(pool);
        // B 只需要额外分配 1 个 block（第 4 个 token）
        tableB.append();
        System.out.println("B 共享了 A 的 3 个 block，新分配了 1 个 block");
        System.out.println("总 block 使用: " + (5 + 1) + " (非共享: " + (5 + 4) + ")");
        System.out.println("节省了 " + ((5+4) - (5+1)) + " 个 block");
    }
}

这个实验让你理解 vLLM 的 Automatic Prefix Caching（APC）。 和 JVM 的 String.intern()——相同内容只存一份——是同一个思路。

六、总结

你已经在做 AI Infra 工程师做的事了。只是名字不一样。

你的世界	AI Infra 的世界
ByteBuf.release()	block_table 的 refcount–
PooledByteBufAllocator	BlockManager
G1 Region	PagedAttention Block
ThreadLocal + TLAB	Per-sequence block_table
Full GC 调优	GPU OOM 排查
String.intern	Automatic Prefix Caching
Heap dump 分析	Block table dump 分析

下一篇预告：事件循环——从 Netty EventLoop 到 GPU Warp Scheduler

你写过无数遍 NioEventLoop.run() 里的 selector.select() → processSelectedKeys() → runAllTasks()。

GPU 的 SM 里跑着完全一样的循环——只是名字变成了 warp scheduler dispatch → execute warp instruction → switch to next warp。

这篇讲什么#

一、问题本质#

二、Netty 怎么管内存#

2.1 三层架构#

2.2 关键洞察#

三、vLLM 怎么管显存#

3.1 先看问题#

3.2 vLLM 的方案：PagedAttention#

3.3 映射表#

3.4 碎片化的对比#

四、JVM 级别的对应#

4.1 G1 GC 和 PagedAttention#

4.2 ThreadLocal + 引用计数#

五、4 个代码实验#

实验 1：引用计数分配的 ByteBuf 版#

实验 2：页表（Block Table）#

实验 3：碎片化对比#

实验 4：前缀共享（Prefix Caching）#

六、总结#