手写 mini-sglang 阶段四 · 更快、更大

自定义 CUDA kernel¶

mini-sglang 自己写的 CUDA kernel 只有三个：把新 K/V 写进 KV 池（store_cache）、词表并行的嵌入查表（indexing）、以及基数树里比较两段 token 的 fast_compare_key（C++，在 CPU 上）。前两个是纯粹的"按下标搬运整行数据"，逻辑简单，但每层每步都要执行，值得写成高效的 kernel。这一章写出它们的 CUDA 版本，用两个版本的 nvcc 编译，并在 CUDA 手册的 CPU 模拟器上运行自检。

自测：能答上来就可以跳过本章

写 KV 缓存的 kernel，瓶颈在计算还是访存？理论上最快能多快？
为什么"一个 warp 负责一行"，而不是"一个线程负责一行"？
用 uint4 拷贝有什么好处？什么时候不能用？
PyTorch 的 k_cache[idx] = k 在 GPU 上也能工作，为什么还要自己写？

自测参考答案（先自己答，再展开对照）

访存：它只是按下标把一行行 K、V 从一个地方搬到另一个地方，没有计算。理论上最快就是读写这些字节的时间 ÷ 显存带宽。
一行有几百到几千个元素：一个 warp 的 32 个线程读同一行里相邻的元素，访问是合并的，也能用满带宽；一个线程负责一行时，相邻线程访问不同的行，完全不合并，而且行数不多时线程数太少。
一条指令搬 16 字节，访存指令数减少到 1/4 甚至更少，在途的数据更多。前提是地址按 16 字节对齐、每行的字节数是 16 的倍数，否则要退回更小的类型或者单独处理尾部。
PyTorch 的高级索引会启动通用的 kernel，有额外的开销（下标检查、元数据处理），还可能引入同步；它也不能和别的操作融合。自己写的 kernel 专门为这个形状优化，还能放进 CUDA Graph、和嵌入查表等操作共用同一套写法。

本章要写的文件：kernel/csrc/kv_kernels.cuh、kernel/csrc/ext.cu、kernel/cuda_ext.py，以及 kernel/__init__.py 中的分派；自检程序 tests/cuda/test_kv_kernels.cu。

到这一章为止的文件：76 个（新建 1，修改 1，提前引入 0）

minisgl/
├── __init__.py
├── core.py
├── env.py
├── __main__.py
├── shell.py
├── attention/  （6 个文件，不变）
├── distributed/  （3 个文件，不变）
├── engine/  （5 个文件，不变）
├── kernel/
│   ├── __init__.py  ← 本章修改
│   ├── torch_ops.py
│   ├── cuda_ext.py  ← 本章新建
├── kvcache/  （5 个文件，不变）
├── layers/  （8 个文件，不变）
├── llm/  （2 个文件，不变）
├── message/  （5 个文件，不变）
├── models/  （7 个文件，不变）
├── moe/  （2 个文件，不变）
├── scheduler/  （9 个文件，不变）
├── server/  （4 个文件，不变）
├── tokenizer/  （4 个文件，不变）
├── utils/  （8 个文件，不变）

这一步的 main：examples/ch19_kernels.py——它只用到上面这些文件；python tools/steps.py check 会逐章搭出这棵树、跑这个 main。

按下标搬运整行¶

写 KV 缓存要做的事是：对本轮每个新 token i，把 k[i]（一行 头数 × head_dim 个元素）拷到 k_cache[out_loc[i]]，v 同理。嵌入查表是反过来：out[i] = weight[ids[i]]。两者都没有任何计算，瓶颈完全是显存带宽——理想情况下，搬运的字节数除以显存带宽就是耗时。

minisgl/kernel/csrc/kv_kernels.cuh

// mini-sglang 的两个访存 kernel：把新算出的 K/V 按位置写进 KV 池，以及（词表并行的）嵌入查表。
// 两者都是纯粹的"按下标搬运整行数据"，瓶颈只有显存带宽：
//   * 一个 warp 负责一行（一个 token），32 个线程协作拷贝；
//   * 按向量类型 V 拷贝（行字节数是 16 的倍数时用 uint4，一次 16 字节），访存是合并的。
#pragma once
#include <cstdint>

namespace minisgl {

// k_cache/v_cache：[池中 token 数, row]，行步长 cache_stride；k/v：[n, row]，行步长 input_stride。
// 所有长度和步长都以 V 为单位。
template <typename V>
__global__ void store_kv_kernel(V* __restrict__ k_cache, V* __restrict__ v_cache,
                                const int32_t* __restrict__ indices, const V* __restrict__ k,
                                const V* __restrict__ v, int64_t n, int64_t row,
                                int64_t cache_stride, int64_t input_stride) {
  const int64_t warp = (static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x) / 32;
  const int lane = threadIdx.x % 32;
  if (warp >= n) return;
  const int64_t dst = static_cast<int64_t>(indices[warp]) * cache_stride;
  const int64_t src = warp * input_stride;
  for (int64_t i = lane; i < row; i += 32) {
    k_cache[dst + i] = k[src + i];
    v_cache[dst + i] = v[src + i];
  }
}

// out[t] = weight[ids[t] - start]，若 ids[t] 不在 [start, start + length) 内则填 0（词表并行）。
template <typename V>
__global__ void embedding_kernel(V* __restrict__ out, const V* __restrict__ weight,
                                 const int32_t* __restrict__ ids, int64_t n, int64_t row,
                                 int64_t start, int64_t length) {
  const int64_t warp = (static_cast<int64_t>(blockIdx.x) * blockDim.x + threadIdx.x) / 32;
  const int lane = threadIdx.x % 32;
  if (warp >= n) return;
  const int64_t local = static_cast<int64_t>(ids[warp]) - start;
  const bool valid = local >= 0 && local < length;
  for (int64_t i = lane; i < row; i += 32) {
    out[warp * row + i] = valid ? weight[local * row + i] : V{};
  }
}

}  // namespace minisgl

设计要点（原理见 CUDA 手册的内存层次与访存优化）：

一个 warp 负责一行：Qwen3-0.6B 的一行 K 是 8 × 128 × 2 = 2 KB。一个线程搬 2 KB 太慢，而一个 warp 的 32 个线程各搬相邻的 16 字节，一次就能搬 512 字节，访存是合并的；
向量化：按 uint4（16 字节）为单位拷贝，一条指令搬 16 字节。要求行字节数和指针都是 16 的倍数；否则退回按元素大小（2 或 4 字节）拷贝。选择由主机端根据形状和对齐决定，kernel 本身按类型 V 模板化；
词表并行的掩码：不在本 rank 词表范围内的 token，整行写 0（V{}），配合之后的 all-reduce。

PyTorch 的 k_cache[idx] = k 在 GPU 上当然也能工作，但它是一个通用的 index_put kernel，要处理任意形状和步长；专用 kernel 按固定的行大小生成代码、向量化拷贝，K、V 在同一个 kernel 里完成，更接近带宽上限。官方还加了 PDL（Programmatic Dependent Launch，Hopper 起支持）的开关，让它与前后的 kernel 部分重叠。

接进 PyTorch¶

minisgl/kernel/csrc/ext.cu

// PyTorch 扩展：把 kv_kernels.cuh 里的 kernel 包装成 Python 可调用的函数。
// 由 minisgl/kernel/cuda_ext.py 在第一次使用时即时编译（torch.utils.cpp_extension.load）。
// CUDA stream 由 Python 端传入（torch.cuda.current_stream().cuda_stream），这样只依赖 torch 的
// CPU 头文件和 CUDA runtime，不需要 c10/cuda 的头文件。
#include <cuda_runtime.h>
#include <torch/extension.h>

#include "kv_kernels.cuh"

namespace {

constexpr int kThreads = 128;  // 每个 block 4 个 warp

bool aligned16(const at::Tensor& t, int64_t row_bytes, int64_t stride_bytes) {
  return row_bytes % 16 == 0 && stride_bytes % 16 == 0 &&
         reinterpret_cast<uintptr_t>(t.data_ptr()) % 16 == 0;
}

template <typename V>
void launch_store(at::Tensor& kc, at::Tensor& vc, const at::Tensor& idx, const at::Tensor& k,
                  const at::Tensor& v, int64_t n, int64_t row_bytes, cudaStream_t stream) {
  const int64_t es = kc.element_size();
  const int64_t blocks = (n * 32 + kThreads - 1) / kThreads;
  minisgl::store_kv_kernel<V><<<blocks, kThreads, 0, stream>>>(
      static_cast<V*>(kc.data_ptr()), static_cast<V*>(vc.data_ptr()), idx.data_ptr<int32_t>(),
      static_cast<const V*>(k.data_ptr()), static_cast<const V*>(v.data_ptr()), n,
      row_bytes / sizeof(V), kc.stride(0) * es / sizeof(V), k.stride(0) * es / sizeof(V));
}

}  // namespace

// k_cache/v_cache：[池中 token 数, H * D]（视图）；k/v：[n, H * D]；indices：int32 [n]
void store_kv(at::Tensor k_cache, at::Tensor v_cache, at::Tensor indices, at::Tensor k,
              at::Tensor v, int64_t stream_ptr) {
  auto stream = reinterpret_cast<cudaStream_t>(stream_ptr);
  TORCH_CHECK(indices.scalar_type() == at::kInt, "indices must be int32");
  TORCH_CHECK(k_cache.stride(1) == 1 && k.stride(1) == 1, "rows must be contiguous");
  const int64_t n = indices.size(0);
  const int64_t row_bytes = k.size(1) * k.element_size();
  if (n == 0) return;
  const bool vec = aligned16(k_cache, row_bytes, k_cache.stride(0) * k_cache.element_size()) &&
                   aligned16(k, row_bytes, k.stride(0) * k.element_size());
  if (vec) {  // 快路径：每个线程一次搬 16 字节
    launch_store<uint4>(k_cache, v_cache, indices, k, v, n, row_bytes, stream);
  } else if (k.element_size() == 4) {
    launch_store<uint32_t>(k_cache, v_cache, indices, k, v, n, row_bytes, stream);
  } else {
    TORCH_CHECK(k.element_size() == 2, "unsupported element size");
    launch_store<uint16_t>(k_cache, v_cache, indices, k, v, n, row_bytes, stream);
  }
}

at::Tensor indexing(at::Tensor weight, at::Tensor ids, int64_t start, int64_t length,
                    int64_t stream_ptr) {
  auto stream = reinterpret_cast<cudaStream_t>(stream_ptr);
  TORCH_CHECK(ids.scalar_type() == at::kInt && weight.is_contiguous());
  const int64_t n = ids.size(0);
  auto out = at::empty({n, weight.size(1)}, weight.options());
  const int64_t row_bytes = weight.size(1) * weight.element_size();
  TORCH_CHECK(row_bytes % 16 == 0, "embedding row must be a multiple of 16 bytes");
  const int64_t blocks = (n * 32 + kThreads - 1) / kThreads;
  if (n > 0) {
    minisgl::embedding_kernel<uint4><<<blocks, kThreads, 0, stream>>>(
        static_cast<uint4*>(out.data_ptr()), static_cast<const uint4*>(weight.data_ptr()),
        ids.data_ptr<int32_t>(), n, row_bytes / 16, start, length);
  }
  return out;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("store_kv", &store_kv);
  m.def("indexing", &indexing);
}

扩展由 torch.utils.cpp_extension.load 在第一次使用时即时编译。CUDA stream 由 Python 端传入（torch.cuda.current_stream().cuda_stream），这样扩展只依赖 PyTorch 的 CPU 头文件和 CUDA runtime——也让我们在没有 CUDA 版 PyTorch 的开发机上能够编译检查它。

minisgl/kernel/__init__.py · store_cache

def store_cache(k_cache, v_cache, indices, k, v) -> None:
    if (ext := _cuda_ext(k)) is not None:  # 第 19 章的自定义 CUDA kernel
        n = k.shape[0]
        ext.store_kv(k_cache.view(k_cache.shape[0], -1), v_cache.view(v_cache.shape[0], -1),
                     indices.int(), k.reshape(n, -1), v.reshape(n, -1),
                     torch.cuda.current_stream().cuda_stream)
        return
    torch_ops.store_cache(k_cache, v_cache, indices, k, v)

GPU 上且扩展编译成功时用自定义 kernel，否则（CPU、没有 nvcc）用 PyTorch 参考实现。

与官方的差异：tvm-ffi 与 torch 扩展

官方用 tvm-ffi 做 JIT 编译和 Python 绑定：kernel 的行大小、线程数等作为 C++ 模板参数，按实际形状即时生成专门的版本（kernel/utils.py 的 load_jit）。我们用更常见的 PyTorch 扩展，行大小作为运行时参数。kernel 的思路（一个 warp 一行、向量化拷贝、词表掩码）相同。

在没有 GPU 的机器上验证¶

自检程序用随机数据分别测试两个 kernel 的向量化路径和逐元素路径，与 CPU 上的参考结果逐字节比较：

tests/cuda/test_kv_kernels.cu

// nvcc -std=c++17 -O2 -I../../python/minisgl/kernel/csrc test_kv_kernels.cu -o test_kv_kernels
// 独立的自检程序：用随机数据分别测试两个 kernel 的向量化路径（uint4）和逐元素路径，与 CPU 结果逐位比较。
#include <cuda_runtime.h>

#include <cstdint>
#include <cstdio>
#include <cstdlib>
#include <cstring>
#include <vector>

#include "kv_kernels.cuh"

#define CHECK(x) do { cudaError_t e = (x); if (e != cudaSuccess) { \
  printf("CUDA error %s at line %d\n", cudaGetErrorString(e), __LINE__); exit(1); } } while (0)

template <typename V>
bool test_store(int pool_tokens, int n, int row_bytes) {
  const int row = row_bytes / sizeof(V);
  std::vector<V> kc(pool_tokens * row), vc(pool_tokens * row), k(n * row), v(n * row);
  std::vector<int32_t> idx(n);
  std::vector<int> perm(pool_tokens);
  for (int i = 0; i < pool_tokens; ++i) perm[i] = i;
  for (int i = pool_tokens - 1; i > 0; --i) std::swap(perm[i], perm[rand() % (i + 1)]);
  for (int i = 0; i < n; ++i) idx[i] = perm[i];  // 互不相同的目标位置
  auto fill = [](std::vector<V>& a) { memset(a.data(), 0, a.size() * sizeof(V));
    for (size_t i = 0; i < a.size() * sizeof(V); ++i) reinterpret_cast<uint8_t*>(a.data())[i] = rand(); };
  fill(kc); fill(vc); fill(k); fill(v);
  V *dkc, *dvc, *dk, *dv; int32_t* didx;
  CHECK(cudaMalloc(&dkc, kc.size() * sizeof(V))); CHECK(cudaMalloc(&dvc, vc.size() * sizeof(V)));
  CHECK(cudaMalloc(&dk, k.size() * sizeof(V))); CHECK(cudaMalloc(&dv, v.size() * sizeof(V)));
  CHECK(cudaMalloc(&didx, n * sizeof(int32_t)));
  CHECK(cudaMemcpy(dkc, kc.data(), kc.size() * sizeof(V), cudaMemcpyHostToDevice));
  CHECK(cudaMemcpy(dvc, vc.data(), vc.size() * sizeof(V), cudaMemcpyHostToDevice));
  CHECK(cudaMemcpy(dk, k.data(), k.size() * sizeof(V), cudaMemcpyHostToDevice));
  CHECK(cudaMemcpy(dv, v.data(), v.size() * sizeof(V), cudaMemcpyHostToDevice));
  CHECK(cudaMemcpy(didx, idx.data(), n * sizeof(int32_t), cudaMemcpyHostToDevice));
  const int threads = 128, blocks = (n * 32 + threads - 1) / threads;
  minisgl::store_kv_kernel<V><<<blocks, threads>>>(dkc, dvc, didx, dk, dv, n, row, row, row);
  CHECK(cudaGetLastError());
  std::vector<V> out_k(kc.size()), out_v(vc.size());
  CHECK(cudaMemcpy(out_k.data(), dkc, kc.size() * sizeof(V), cudaMemcpyDeviceToHost));
  CHECK(cudaMemcpy(out_v.data(), dvc, vc.size() * sizeof(V), cudaMemcpyDeviceToHost));
  for (int i = 0; i < n; ++i) {  // CPU 参考结果
    memcpy(&kc[idx[i] * row], &k[i * row], row * sizeof(V));
    memcpy(&vc[idx[i] * row], &v[i * row], row * sizeof(V));
  }
  bool ok = memcmp(out_k.data(), kc.data(), kc.size() * sizeof(V)) == 0 &&
            memcmp(out_v.data(), vc.data(), vc.size() * sizeof(V)) == 0;
  cudaFree(dkc); cudaFree(dvc); cudaFree(dk); cudaFree(dv); cudaFree(didx);
  return ok;
}

bool test_embedding(int vocab, int start, int length, int n, int dim) {
  const int row = dim * 2 / 16;  // bf16 行，按 uint4 搬运
  std::vector<uint4> w(length * row), out(n * row);
  for (auto& x : w) x = uint4{(unsigned)rand(), (unsigned)rand(), (unsigned)rand(), (unsigned)rand()};
  std::vector<int32_t> ids(n);
  for (int i = 0; i < n; ++i) ids[i] = rand() % vocab;
  uint4 *dw, *dout; int32_t* dids;
  CHECK(cudaMalloc(&dw, w.size() * sizeof(uint4))); CHECK(cudaMalloc(&dout, out.size() * sizeof(uint4)));
  CHECK(cudaMalloc(&dids, n * sizeof(int32_t)));
  CHECK(cudaMemcpy(dw, w.data(), w.size() * sizeof(uint4), cudaMemcpyHostToDevice));
  CHECK(cudaMemcpy(dids, ids.data(), n * sizeof(int32_t), cudaMemcpyHostToDevice));
  const int threads = 128, blocks = (n * 32 + threads - 1) / threads;
  minisgl::embedding_kernel<uint4><<<blocks, threads>>>(dout, dw, dids, n, row, start, length);
  CHECK(cudaGetLastError());
  CHECK(cudaMemcpy(out.data(), dout, out.size() * sizeof(uint4), cudaMemcpyDeviceToHost));
  bool ok = true;
  for (int i = 0; i < n && ok; ++i) {
    const int local = ids[i] - start;
    for (int j = 0; j < row; ++j) {
      uint4 e = (local >= 0 && local < length) ? w[local * row + j] : uint4{0, 0, 0, 0};
      uint4 g = out[i * row + j];
      ok &= e.x == g.x && e.y == g.y && e.z == g.z && e.w == g.w;
    }
  }
  cudaFree(dw); cudaFree(dout); cudaFree(dids);
  return ok;
}

int main() {
  srand(42);
  bool all = true;
  auto report = [&](const char* name, bool ok) { printf("%-40s %s\n", name, ok ? "PASS" : "FAIL"); all &= ok; };
  report("store_kv uint4  (8 heads x 128 x bf16)", test_store<uint4>(300, 37, 8 * 128 * 2));
  report("store_kv uint16 (row = 6 bytes)", test_store<uint16_t>(64, 17, 6));
  report("embedding, full vocab", test_embedding(1000, 0, 1000, 23, 64));
  report("embedding, vocab shard [500, 750)", test_embedding(1000, 500, 250, 41, 64));
  return all ? 0 : 1;
}

它用 CUDA 13.4 和 12.9 两个版本的 nvcc 编译，然后在 CUDA 手册的 CPU 模拟器上运行——模拟器为每个 CUDA 线程开一个协程，真实地执行 warp 内的并行逻辑：

examples/ch19_kernels.py

"""第 19 章：用两个版本的 nvcc 编译 kernel，并在 CPU 模拟器上运行自检。"""

import re
import subprocess
import sys
from pathlib import Path

sys.path.insert(0, "tools")
import emu_kernels  # noqa: E402

csrc = Path("python/minisgl/kernel/csrc")
for nvcc in sorted(Path.home().glob("cuda-handbook/.toolkit*/bin/nvcc")):
    version = re.search(r"release ([\d.]+)", subprocess.run([str(nvcc), "--version"], capture_output=True,
                                                            text=True).stdout).group(1)
    r = subprocess.run([str(nvcc), "-std=c++17", "-O2", "-Wno-deprecated-gpu-targets", f"-I{csrc}", "-c",
                        "tests/cuda/test_kv_kernels.cu", "-o", "/dev/null"], capture_output=True, text=True)
    print(f"nvcc {version}: 编译 {'通过' if r.returncode == 0 else '失败'}")
print("CPU 模拟器上运行自检：")
ok = emu_kernels.run(Path("tests/cuda/test_kv_kernels.cu"))
print("全部通过" if ok else "有失败")

运行结果

nvcc 13.4: 编译 通过
nvcc 12.9: 编译 通过
CPU 模拟器上运行自检：
store_kv uint4  (8 heads x 128 x bf16)   PASS
store_kv uint16 (row = 6 bytes)          PASS
embedding, full vocab                    PASS
embedding, vocab shard [500, 750)        PASS
全部通过

PyTorch 扩展 ext.cu 同样用两个版本的 nvcc 编译检查（tools/check.py），但它依赖 CUDA 版的 PyTorch 运行时，只能在 GPU 上执行。

官方实现

写 KV：kernel/csrc/jit/store.cu，Python 端 kernel/store.py · store_cache（第 30～42 行）
查表：kernel/csrc/jit/index.cu，Python 端 kernel/index.py · indexing（第 31～50 行）（行字节数是 2048、1024 的倍数时，用 4 个或 2 个 warp 分担一行）
基数树的比较：kernel/csrc/src/radix.cpp（CPU 上的 std::mismatch）

测试¶

tests/test_ch19_kernels.py · test_cuda_kernels_on_cpu_emulator

def test_cuda_kernels_on_cpu_emulator():
    """CUDA kernel 的自检程序在 CUDA 手册的 CPU 模拟器上运行（没有 GPU 也能验证 kernel 逻辑）。"""
    r = subprocess.run([sys.executable, str(ROOT / "tools" / "emu_kernels.py")],
                       capture_output=True, text=True, timeout=900)
    assert r.returncode == 0, r.stdout + r.stderr
    assert r.stdout.count("PASS") == 4

怎么讲清楚

讲自定义 kernel：写 KV 缓存、嵌入查表都是"按下标搬运整行"，瓶颈是带宽，理论最快是"读写字节数 / 显存带宽"。一个 warp 负责一行，32 个线程读连续的地址、合并访问（一个线程一行会让相邻线程的地址相差一整行）；用 uint4 一次搬 16 字节，减少访存指令，前提是地址 16 字节对齐、行长能被整除。PyTorch 的 k_cache[idx] = k 是通用的 index_put，要处理任意形状和步长、可能有多次 kernel 启动，专用 kernel 按固定行大小生成代码更快，也能被 CUDA Graph 录制。没有 GPU 时用 CPU 模拟器逐字节自检。

练习¶

在 H100（显存带宽约 3.35 TB/s）上，Qwen3-0.6B 一次 decode、批大小 64，每层写 KV 要搬多少字节？理论耗时是多少？这个 kernel 值得用 CUDA Graph 吗？
官方的 indexing 在行很大时用多个 warp 分担一行。为 embedding_kernel 加上这个优化，并在模拟器上验证。
把 store_kv_kernel 改成"一个线程块负责多行、每个 warp 一行"的写法，和现在"每 4 个 warp 一个块"的写法有什么区别？

参考答案

每个 token 的 K、V 各 2 KB，64 个 token 共 256 KB（读 k、v 再写入池子，读写合计约 512 KB），约 0.15 µs——远小于 kernel 的启动开销（几微秒）。所以它必须放进 CUDA Graph，否则启动开销会是实际工作的几十倍。
num_splits 个 warp 负责一行：warp_id / num_splits 是行号，warp_id % num_splits 是这一行里的第几段，每段 row / num_splits 个元素。
现在的写法本来就是每个线程块 128 个线程（4 个 warp），每个 warp 一行；改变每块的 warp 数只影响调度粒度和 occupancy，对这种纯访存的 kernel 影响不大，关键仍是每个 warp 内的合并访问。

小结¶

写 KV 和嵌入查表都是"按下标搬运整行"，瓶颈是带宽：一个 warp 一行、uint4 向量化、合并访问。
PyTorch 扩展即时编译，stream 从 Python 传入；GPU 上用 kernel，其他情况退回参考实现。
kernel 用两个版本的 nvcc 编译，并在 CPU 模拟器上逐字节自检。

本章练习题（1 道）

在浏览器里直接写代码、跑测试；也可以在本地用 python practice/judge.py 判题。

202. 写 KV Cache 的 kernel（中等）

发现错误或有更好的写法？在 GitHub 上编辑此页提交问题