Skip to content
mini-sglang from Scratch Stage 4 · Faster and bigger

Custom CUDA kernels

mini-sglang writes only three CUDA kernels of its own: writing new K/V into the KV pool (store_cache), the vocabulary-parallel embedding lookup (indexing), and fast_compare_key for comparing two stretches of tokens in the radix tree (C++, on the CPU). The first two are pure "move whole rows by index" with simple logic, but they run every layer of every step, so they are worth writing as efficient kernels. This chapter writes their CUDA versions, compiles them with two versions of nvcc, and self-checks them on the CUDA handbook's CPU simulator.

Self-test: if you can answer these, skip the chapter

  1. Is the KV-cache write kernel compute-bound or memory-bound? How fast can it be in theory?
  2. Why "one warp per row" rather than "one thread per row"?
  3. What does copying with uint4 buy? When can it not be used?
  4. PyTorch's k_cache[idx] = k works on a GPU too. Why write your own?
Answers (try it yourself first, then expand)
  1. Memory-bound: it only moves rows of K and V from one place to another by index, with no computation. The theoretical best is the bytes read and written divided by memory bandwidth.
  2. A row has hundreds to thousands of elements: a warp's 32 threads read neighbouring elements of the same row, so the accesses coalesce and saturate the bandwidth. With one thread per row, neighbouring threads touch different rows, nothing coalesces, and with few rows there are too few threads.
  3. One instruction moves 16 bytes, cutting the number of memory instructions to a quarter or less and keeping more data in flight. It requires the addresses to be 16-byte aligned and the row's byte count to be a multiple of 16; otherwise fall back to a smaller type or handle the tail separately.
  4. PyTorch's advanced indexing launches a generic kernel with extra overhead (index checks, metadata handling) and can introduce synchronization; it also cannot be fused with anything else. A hand-written kernel is specialized to this shape, can go into a CUDA Graph, and shares one pattern with the embedding lookup.

Files you will write: kernel/csrc/kv_kernels.cuh, kernel/csrc/ext.cu, kernel/cuda_ext.py, and the dispatch in kernel/__init__.py; plus the self-check program tests/cuda/test_kv_kernels.cu.

Files so far: 76 (1 new, 1 modified, 0 brought in early)
minisgl/
├── __init__.py
├── core.py
├── env.py
├── __main__.py
├── shell.py
├── attention/  (6 files, unchanged)
├── distributed/  (3 files, unchanged)
├── engine/  (5 files, unchanged)
├── kernel/
│   ├── __init__.py  ← modified in this chapter
│   ├── torch_ops.py
│   ├── cuda_ext.py  ← new in this chapter
├── kvcache/  (5 files, unchanged)
├── layers/  (8 files, unchanged)
├── llm/  (2 files, unchanged)
├── message/  (5 files, unchanged)
├── models/  (7 files, unchanged)
├── moe/  (2 files, unchanged)
├── scheduler/  (9 files, unchanged)
├── server/  (4 files, unchanged)
├── tokenizer/  (4 files, unchanged)
├── utils/  (8 files, unchanged)
This step's main: examples/ch19_kernels.py — it uses only the files above; python tools/steps.py check rebuilds this tree chapter by chapter and runs it.

Moving whole rows by index

Writing the KV cache means: for each new token i this step, copy k[i] (one row of heads × head_dim elements) to k_cache[out_loc[i]], and the same for v. The embedding lookup is the reverse: out[i] = weight[ids[i]]. Neither computes anything, so the bottleneck is purely memory bandwidth: at best the bytes moved divided by the bandwidth.

minisgl/kernel/csrc/kv_kernels.cuh
// mini-sglang's two memory kernels: writing newly computed K/V into the KV pool by location, and the (vocabulary-parallel) embedding lookup.
// Both are pure "move whole rows by index", bounded only by memory bandwidth:
//   * one warp per row (one token), with its 32 threads copying together;
//   * copying in the vector type V (uint4, 16 bytes at a time, when the row's byte count is a multiple of 16), with coalesced accesses.
#pragma once
#include <cstdint>

namespace minisgl {

// k_cache/v_cache: [tokens in the pool, row] with row stride cache_stride; k/v: [n, row] with row stride input_stride.
// Every length and stride is counted in units of 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], zero when ids[t] falls outside [start, start + length) (vocabulary parallel).
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
The design (the background is in the CUDA handbook's memory hierarchy and access optimization):

  • one warp per row: one row of K for Qwen3-0.6B is 8 × 128 × 2 = 2 KB. One thread moving 2 KB is far too slow, while a warp's 32 threads each moving a neighbouring 16 bytes move 512 bytes at once, with coalesced accesses;
  • vectorization: copy in uint4 units (16 bytes), one instruction per 16 bytes. This needs both the row's byte count and the pointers to be multiples of 16; otherwise it falls back to copying by element size (2 or 4 bytes). The host side chooses based on the shape and alignment, and the kernel itself is templated on the type V;
  • the vocabulary-parallel mask: a token outside this rank's vocabulary slice gets a whole row of zeros (V{}), which pairs with the all-reduce that follows.

PyTorch's k_cache[idx] = k works on a GPU, of course, but it is a generic index_put kernel that has to handle any shape and stride; a dedicated kernel generates code for a fixed row size, vectorizes the copy and does K and V in one kernel, landing much closer to the bandwidth limit. Upstream also adds a switch for PDL (Programmatic Dependent Launch, supported from Hopper on) so it can partly overlap with the kernels around it.

Hooking into PyTorch

minisgl/kernel/csrc/ext.cu
// The PyTorch extension: wraps the kernels in kv_kernels.cuh as functions callable from Python.
// Compiled just in time on first use by minisgl/kernel/cuda_ext.py (torch.utils.cpp_extension.load).
// the CUDA stream is passed in from Python (torch.cuda.current_stream().cuda_stream), so this depends only on torch's
// CPU headers and the CUDA runtime; the c10/cuda headers are not needed.
#include <cuda_runtime.h>
#include <torch/extension.h>

#include "kv_kernels.cuh"

namespace {

constexpr int kThreads = 128;  // 4 warps per block

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: [tokens in the pool, H * D] (a view); 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) {  // the fast path: 16 bytes per thread per move
    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);
}
The extension is compiled just in time on first use by torch.utils.cpp_extension.load. The CUDA stream is passed in from Python (torch.cuda.current_stream().cuda_stream), so the extension depends only on PyTorch's CPU headers and the CUDA runtime, which also lets us compile-check it on a development machine without a CUDA build of 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:  # the custom CUDA kernels of chapter 19
        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)
On a GPU with the extension compiled it uses the custom kernel; otherwise (CPU, no nvcc) it falls back to the PyTorch reference implementation.

Difference from upstream: tvm-ffi versus torch extensions

Upstream uses tvm-ffi for JIT compilation and Python bindings: the kernel's row size, thread count and so on are C++ template parameters, and a specialized version is generated on the fly for the actual shape (load_jit in kernel/utils.py). We use the more familiar PyTorch extension with the row size as a runtime argument. The kernel's ideas (one warp per row, vectorized copies, the vocabulary mask) are the same.

Verifying on a machine with no GPU

The self-check program tests both kernels' vectorized and element-wise paths on random data and compares byte for byte against a CPU reference:

tests/cuda/test_kv_kernels.cu
// nvcc -std=c++17 -O2 -I../../python/minisgl/kernel/csrc test_kv_kernels.cu -o test_kv_kernels
// A standalone self-check: test both kernels' vectorized (uint4) and element-wise paths on random data, bit for bit against the CPU result.
#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];  // distinct destination locations
  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) {  // the CPU reference result
    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 rows, moved as 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;
}
It is compiled with both CUDA 13.4 and 12.9 nvcc and then run on the CUDA handbook's CPU simulator, which gives every CUDA thread a coroutine and really executes the parallel logic within a warp:

examples/ch19_kernels.py
"""Chapter 19: compile the kernels with two versions of nvcc and self-check them on the CPU simulator."""

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}: compiled {'ok' if r.returncode == 0 else 'FAILED'}")
print("running the self-check on the CPU simulator:")
ok = emu_kernels.run(Path("tests/cuda/test_kv_kernels.cu"))
print("all passed" if ok else "some failed")
What it prints
nvcc 13.4: compiled ok
nvcc 12.9: compiled ok
running the self-check on the CPU simulator:
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
all passed
The PyTorch extension ext.cu is compile-checked with both nvcc versions too (tools/check.py), but it depends on a CUDA build of PyTorch at runtime and can only execute on a GPU.

The official implementation

Tests

tests/test_ch19_kernels.py · test_cuda_kernels_on_cpu_emulator
def test_cuda_kernels_on_cpu_emulator():
    """The CUDA kernels' self-check runs on the CUDA handbook's CPU simulator (kernel logic verified with no GPU)."""
    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

How to explain it

On custom kernels: writing the KV cache and the embedding lookup are both "move whole rows by index", the bottleneck is bandwidth, and the theoretical best is bytes moved divided by memory bandwidth. One warp handles a row so its 32 threads read contiguous addresses and coalesce (one thread per row would put a whole row between neighbouring threads' addresses); uint4 moves 16 bytes at a time and cuts the memory instructions, given 16-byte alignment and a divisible row length. PyTorch's k_cache[idx] = k is a generic index_put that handles any shape and stride and may launch several kernels, while a dedicated kernel generated for a fixed row size is faster and can be captured into a CUDA Graph. Without a GPU, a CPU simulator self-checks it byte for byte.

Exercises

  1. On an H100 (about 3.35 TB/s of memory bandwidth), for one Qwen3-0.6B decode step at batch size 64, how many bytes does the per-layer KV write move? What is the theoretical time? Is this kernel worth putting in a CUDA Graph?
  2. Upstream's indexing has several warps share a row when the row is large. Add that optimization to embedding_kernel and verify it on the simulator.
  3. Rewrite store_kv_kernel as "one block handles several rows, one warp per row". How does that differ from the current "4 warps per block"?
Answers
  1. Each token's K and V are 2 KB each, so 64 tokens are 256 KB (reading k and v and writing the pool makes roughly 512 KB of traffic), about 0.15 µs, which is far below a kernel launch's few microseconds. So it has to go into a CUDA Graph, or the launch overhead would be tens of times the actual work.
  2. Have num_splits warps share a row: warp_id / num_splits is the row and warp_id % num_splits is the segment within it, each segment covering row / num_splits elements.
  3. The current version already uses 128 threads per block (4 warps) with one warp per row; changing the warps per block only affects scheduling granularity and occupancy, which matters little for a purely memory-bound kernel where coalescing within each warp is what counts.

Summary

  • Writing KV and the embedding lookup are both "move whole rows by index" and bandwidth-bound: one warp per row, uint4 vectorization, coalesced accesses.
  • The PyTorch extension compiles just in time with the stream passed in from Python; the kernel is used on a GPU and the reference implementation everywhere else.
  • The kernels are compiled with two nvcc versions and self-checked byte for byte on the CPU simulator.

Exercises for this chapter (1)

Write the code and run the tests in your browser, or judge locally with python practice/judge.py. The exercises are in Chinese for now.

Found a mistake or a better way to say it? Edit this page on GitHubReport an issue