Intro-ops/docs/en/phase1-kernel-writing.md

Phase 1 Training Objective: Writing Kernels

Goal

This phase has one primary task:

Write kernel files for the existing operators.

All other code (descriptor lifecycle, Python bindings, tests, build system) is already implemented. Your only job is to fill in the computation logic.

Default requirements:

1. Priority: get the kernel correct and running.
2. You do NOT need to pursue peak performance initially.
3. Complete the TODOs in kernel.cuh / kernel.py first, then consider aggressive optimizations.

If you've completed the basic goals and want to push performance further, you can modify the wrapper layer (e.g., .cu files for the NVIDIA backend), adjust launch policy, switch between kernels, or add more aggressive specializations. These are advanced topics and not required for Phase 1.

---

Files to Write

Each operator has two backends, corresponding to two kernel files:

Operator	NVIDIA kernel	TileLang kernel
copy	ops/copy/nvidia/kernel.cuh	ops/copy/tilelang/kernel.py
vector_add	ops/vector_add/nvidia/kernel.cuh	ops/vector_add/tilelang/kernel.py
reduce_sum	ops/reduce_sum/nvidia/kernel.cuh	ops/reduce_sum/tilelang/kernel.py
softmax	ops/softmax/nvidia/kernel.cuh	ops/softmax/tilelang/kernel.py

Recommended order (easy to hard): copy → vector_add → reduce_sum → softmax.

---

NVIDIA Kernel (.cuh files)

What to Write

A __global__ function placed under the operator's namespace. The function handles only computation logic, taking raw pointers as input/output. It does NOT involve any descriptor or public API.

By default, you do NOT need to modify the .cu files in the same directory this phase. The repository already provides compatibility-focused launch boilerplate. This phase only requires you to make kernel.cuh correct and runnable.

If you've completed the basic goal and want to push performance, treat the .cu file as an advanced optimization layer: adjust thread counts there, switch kernel variants, or add specialized paths. These are outside the scope of Phase 1.

Difficulty Progression

copy / vector_add: Element-wise operations using the grid-stride loop pattern. Each thread handles several independent elements with no inter-thread communication.

reduce_sum: Row-wise reduction. One block per row. Threads within the block first independently accumulate their assigned columns, then perform a tree reduction via shared memory. __syncthreads() is required for synchronization.

softmax: The basic version uses a three-pass flow: first pass finds the row max (numerical stability), second pass computes exp(x - max) sum and writes intermediate results to output, third pass divides by the sum. Each pass requires intra-block synchronization.

Concepts to Understand

• grid-stride loop: Why this pattern handles arbitrarily-sized tensors
• shared memory reduction: What each step of tree reduction does, and why __syncthreads() is necessary
• softmax minus max: Why computing exp(x) directly causes problems, and why subtracting the row max does not change the result

---

TileLang Kernel (.py files)

What to Write

A Python function decorated with @tilelang.jit that describes tile-level computation using the TileLang DSL. TileLang compiles it into real CUDA kernels.

As with NVIDIA, you only need to complete the TODOs in kernel.py this phase — no need to modify outer adapter code.

Difficulty Progression

copy: Use the built-in T.copy to move data between two tiles. No manual loop needed.

vector_add: Use a T.Parallel loop to compute element-wise within a tile. Understand the difference between T.Parallel and T.Serial.

reduce_sum: Block-wise accumulation required. The outer loop uses T.Serial to traverse column-direction blocks sequentially (because state accumulates), and the inner loop uses T.reduce_sum to reduce over a fragment.

softmax: Two-pass scan using the online softmax algorithm. The first pass maintains a rolling log-sum-exp state; the second pass normalizes with the final LSE. Uses exp2 / log2 instead of exp / log.

Concepts to Understand

• T.Parallel vs T.Serial: When can loop iterations run in parallel, and when must they be sequential
• T.alloc_fragment: A tile-level local buffer, corresponding to registers or shared memory
• T.copy: Moves a chunk of global memory into a fragment — not element-by-element assignment
• online softmax: Why a single pass can produce the correct normalization, and the rolling update logic for log-sum-exp
• exp2 / log2: Why TileLang uses these instead of exp / log

---

Verification

After writing each kernel, verify with the corresponding test:

# NVIDIA kernel (requires rebuild)
bash scripts/build_nvidia.sh build
PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia pytest tests/op_tests/test_<op>.py -v --backend nvidia

# TileLang kernel (no build needed)
PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia pytest tests/op_tests/test_<op>.py -v --backend tilelang

# Both correctness + benchmark
PYTHONPATH=python:. CAMP_BUILD_DIR=build-nvidia python tests/run_ops.py --op <op> --backend nvidia --mode all

All four operators passing both backends = Phase 1 complete.