LESSON 07 · 2026.04.19 · T4

PyTorch Custom Op — wiring CUDA in ~50 lines

Register Lesson 6's Flash Attention kernel as torch.ops.mylib.flash_attention. From this moment on, vLLM's csrc/*.cu stops looking foreign.

The core ~5 lines

In a single file (330 lines, mostly the kernel body), the pivot into the PyTorch world is these five lines:

check_qkv(Q, "Q");                           // shape/device/dtype/contig
auto O = torch::empty({N, d}, Q.options());  // PyTorch allocates
auto stream = at::cuda::getCurrentCUDAStream();
my_kernel<<<grid, block, 0, stream>>>(
    Q.data_ptr<float>(), ..., O.data_ptr<float>());
return O;

That replaces Lesson 6's hundreds of lines of main() (CLI parsing, malloc, CSV output, CPU ref check) entirely. This is the production pattern.

Accuracy

3–5× FP32 machine epsilon (~1.2e-7). SDPA and our kernel both sit inside FP32 rounding limits.

Speed (vs F.scaled_dot_product_attention, T4, d=64)

0.33–0.55× of SDPA. The gap narrows with N (both are bound by N², converging toward a similar regime). The number is where cuDNN's tuning level shows up clearly.

Lesson 1 · stream-aware launches are non-negotiable

Skip at::cuda::getCurrentCUDAStream() as the 4th launch argument and you go to the default stream. PyTorch might be using a different stream → silent race condition. No crash, non-deterministic output. vLLM's PagedAttention launch follows exactly this pattern.

Lesson 2 · the dispatcher gives you the CPU guard for free

TORCH_LIBRARY_IMPL(mylib, CUDA, m) {
  m.impl("flash_attention", &flash_attention_forward);
}

Register only under the CUDA backend. When a CPU tensor arrives, the dispatcher blocks it upfront:

RuntimeError: Could not run 'mylib::flash_attention' with arguments from the 'CPU' backend

Autograd, dtype promotion, device dispatch all wired into one pipe — a clean part of PyTorch's op system.

Lesson 3 · the 3× gap is a line and reality

• cuDNN has register allocation and instruction scheduling hand-tuned by NVIDIA engineers
• Uses warp-specialized FP32 paths (WMMA is FP16-only)
• Smem layouts that avoid bank conflicts
• At our Br=64 / d=64, register pressure may be causing spills

Our place isn't SOTA chasing. It's filling operators the library is missing (novel ops, custom sparsity).

Trap log

• import mylib_ext → PyInit_mylib_ext not defined. Fix: add an empty PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {}.
• nvcc not found. Set CUDA_HOME=/usr/local/cuda.
• CUDA 12.9 toolkit vs PyTorch cu121 wheel — runs under minor-version compat. For production, keep them aligned.

Lessons 1–7 stack

Python model (vLLM, my service)
        │
        │   torch.ops.mylib.flash_attention(q,k,v)    ← layer broken by lesson 07
        ▼
torch dispatcher
        │
        ▼
C++ host wrapper (tensor → raw ptr, stream)
        │
        ▼
CUDA kernel (flash_attention_v1)                     ← lesson 06
        │
        ▼
Warp / thread (shuffle, tiled mma)                   ← lessons 03, 04, 05
        │
        ▼
Memory hierarchy (HBM↔L2↔smem↔reg)                   ← lessons 01, 02

That's CUDA Phase 1 in the bag. Triton starts next lesson.