Skand Tandon
Standalone · Kernel

Crossing the boundary: custom kernels and the C++/Python ABI in vLLM

Python is a productive orchestration language for inference serving, but it is the wrong tool for the critical path of token generation. Every call has to cross the Python/C++ boundary, and that crossing involves an ABI, a dispatcher, and the overhead of launching work on the GPU.

2026-08-20 kernel engine 22 min read

Python is a productive orchestration language for inference serving, but it is the wrong tool for the critical path of token generation. Large-model inference is bound by memory bandwidth and by per-operation latency, and the interpreter cannot meet either constraint. So frameworks like vLLM keep the control flow in Python and push the arithmetic into compiled C++ and CUDA kernels.

That split is not free. Every call has to cross the Python/C++ boundary, and that crossing involves an Application Binary Interface (ABI), a dispatcher, and the overhead of launching work on the GPU. This walks through how vLLM crosses that boundary: how kernels are registered with PyTorch, the ABI constraints that make the boundary fragile, the hardware-level decisions inside the kernels themselves, and how vLLM amortizes the per-call overhead with CUDA graphs.

The important feature of this path is that the bulk data, the multi-gigabyte tensors, never moves. Only metadata and pointers cross the boundary; the launch is asynchronous, so the CPU returns to Python and queues the next op while the GPU is still working on the last one.

The cost of crossing the Python/C++ boundary

When a Python script invokes a PyTorch operation, very little of the work happens in Python. The call crosses into a pre-compiled C++ backend, and the crossing has a fixed cost regardless of how small the tensors are. That cost comes from several places:

None of these is large in isolation. The problem is the multiplier. Generating one token runs a tight sequence of operations, and a decode step can issue hundreds of ops. At a few microseconds of fixed overhead each, the boundary cost alone can rival the time the GPU spends on the math, especially at small batch sizes where each kernel is short. The two structural responses are to do more work per crossing (fused kernels) and to remove the crossings entirely on the hot path (CUDA graphs).

The integration layer: torch.library and the dispatcher

vLLM implements operations that PyTorch does not ship natively: PagedAttention, AWQ/GPTQ dequantization, fused RoPE, fused RMSNorm, and others. It has to make them first-class citizens of the dispatcher so they participate in autograd, torch.compile, and device routing like any built-in op.

The dispatcher is a routing table. Each operation has a schema (a typed signature), and for each dispatch key there can be a different registered implementation. Registration happens in C++ through the torch.library API:

#include <torch/library.h>

// The C++ wrapper that validates inputs and launches the CUDA kernel.
void paged_attention_v1(
    torch::Tensor& out,            // mutated output  (marked Tensor! in the schema)
    torch::Tensor const& query,
    torch::Tensor const& key_cache,
    torch::Tensor const& value_cache,
    torch::Tensor const& block_tables,
    torch::Tensor const& seq_lens,
    int64_t block_size,
    int64_t max_seq_len) {
  TORCH_CHECK(query.is_cuda(), "query must be a CUDA tensor");
  TORCH_CHECK(query.scalar_type() == at::kHalf ||
              query.scalar_type() == at::kBFloat16,
              "query must be fp16 or bf16");
  TORCH_CHECK(out.is_contiguous(), "out must be contiguous");
  // launch_paged_attention_v1(out.data_ptr<...>(), query.data_ptr<...>(), ...);
}

// Declare the schema. The `!` on `Tensor!` records that `out` is mutated.
TORCH_LIBRARY(vllm, m) {
  m.def(
    "paged_attention_v1(Tensor! out, Tensor query, Tensor key_cache, "
    "Tensor value_cache, Tensor block_tables, Tensor seq_lens, "
    "int block_size, int max_seq_len) -> ()");
}

// Bind the implementation to the CUDA dispatch key.
TORCH_LIBRARY_IMPL(vllm, CUDA, m) {
  m.impl("paged_attention_v1", &paged_attention_v1);
}

Two details are worth dwelling on, because they are easy to get wrong:

The ABI boundary itself

An ABI is the machine-level contract for how compiled code interoperates: how arguments are passed in registers or on the stack, how structs are laid out and aligned, how names are encoded into symbols, and how exceptions and destructors behave. The Python/C++ boundary in vLLM is really three nested contracts, and all three have to agree.

The direction of travel is to make this boundary stable. PyTorch 2.10+ exposes an ABI-stable LibTorch API built around torch::stable::Tensor and a C-style shim, so an extension can target one frozen interface and survive across PyTorch versions instead of being rebuilt for each. It trades a little expressiveness for the ability to ship a single wheel, the same bargain that the dispatcher already makes for name mangling, now extended to the type layout.

Hardware-level optimizations in custom CUDA kernels

Once the boundary is crossed and the wrapper launches, the constraints change entirely. The kernel's job is to keep the SMs fed and the memory subsystem busy. A few principles drive vLLM's kernel design, illustrated with PagedAttention and the quantized-dequant kernels.

Memory coalescing and alignment. When a warp (32 threads) accesses global memory, the hardware coalesces the accesses into the fewest possible transactions when the addresses are contiguous and aligned. vLLM lays out data so reads coalesce, and uses vectorized loads to move 128 bits per thread per instruction:

// Each thread loads 8 fp16 values (128 bits) in one instruction.
// Casting to uint4 makes the compiler emit a single 16-byte vectorized load.
const uint4* kv_vec = reinterpret_cast<const uint4*>(key_cache_ptr);
uint4 chunk = kv_vec[thread_offset];   // 1 transaction per aligned warp access

The alignment is a hard precondition, not a nicety: a uint4 load from a non-16-byte-aligned address is undefined behavior on the GPU, and even a merely uncoalesced (but legal) access splits one transaction into several, which is costly on a bandwidth-bound kernel. PagedAttention's block tables are sized and padded precisely so that each KV block starts on an aligned boundary.

Register pressure and occupancy. Registers are the fastest storage on the GPU and the scarcest. Each SM has a fixed register file shared by all resident threads, so a kernel that uses more registers per thread allows fewer thread blocks to be resident at once, which lowers occupancy and the hardware's ability to hide memory latency by switching warps. vLLM kernels use #pragma unroll deliberately rather than reflexively, and keep variable lifetimes short, because unrolling that inflates the live register set can reduce occupancy enough to erase the gain. Block dimensions are chosen to balance registers against shared-memory use for a target architecture such as Ampere or Hopper.

Mental model

The Python/C++ boundary is a toll booth on every op. The toll is small, but a decode step passes through hundreds of booths. Fused kernels are carpooling, and CUDA graphs are a season pass: capture once, ride forever.

Amortizing the overhead: CUDA graphs

CUDA graphs capture a sequence of kernel launches and replay them without CPU involvement. The graph is captured once (the first time the model runs), and every subsequent decode step replays it. This removes the per-call crossing cost entirely from the hot path: no dispatcher, no marshaling, no launch overhead. The CPU hands the GPU a graph handle and the GPU replays the whole sequence.

The tradeoff is that the graph is fixed: the sequence of kernels, their shapes, and their memory addresses are frozen at capture time. This is why vLLM's graph capture happens per shape and why dynamic shapes break graph replay. It is also why the op IR and the dispatcher matter: they give the engine the flexibility to choose what goes into the graph and what stays dynamic.

The Python/C++ boundary is where inference serving's two worlds meet: orchestration and arithmetic. Every microsecond of crossing cost is a microsecond the GPU is not working.

Back to the blog