Triton Kernel Writing

Native kernels live in rapid_llm/kernels/ops/<family>/, one module per op family: a module docstring, a public wrapper, and _-prefixed @triton.jit kernels. An implementation module imports Triton at the top level; the dispatch tier above it (dispatcher/, ops/interfaces.py) stays torch-free and kernels/__init__.py re-exports lazily, which is what keeps a CPU-only install importable. Most kernels are reached through dispatch() behind a KernelSpec row; the few the engine calls directly are re-exported by name. Triton is the default kernel path; when it cannot express an op (inline PTX, special intrinsics, fine-grained memory control), add-jit-kernel covers the CUDA C++ JIT alternative and carries the full path-decision table. Correctness lives in tests/kernels/ against tests/reference.py; measurement belongs to benchmarks/kernels/ and the kernel-microbenchmark skill.

Implementation

• Follow the official Triton semantics. Check it when behavior may differ from Python or NumPy, especially type promotion, integer division and modulo, casts, broadcasting, and variable scoping.
• Use the Triton kernel generated by torch.compile as a possible implementation to inspect. Print Inductor's generated code with TORCH_LOGS="output_code" .venv/bin/python <script> or enable torch._logging.set_logs(output_code=True) before the compiled function runs. Treat generated code as a reference, not as proof of correctness or optimality.
• Resolve tile knobs through the repo's tile policy, not triton.autotune. Dense GEMM and MoE launchers call resolve_tiles() from rapid_llm/kernels/ops/tile_policy.py: autotune store first, then a device-tiered heuristic table (TileTier.PRE_HOPPER covers sm86/sm89 — those parts carry roughly half of Hopper's smem, so an H100-measured wide tier spills or fails to compile there; HOPPER_UP is sm90+). The store is filled offline by the scripts under benchmarks/kernels/, not by an @triton.autotune at import time. Pass block_k_multiple when a format's scales are grouped along k (int4 group_size, nvfp4's 16-element blocks) so a tuned BLOCK_K still covers whole groups.
• Prefer simple code and fast startup over tuning. Cold start is a stated repo goal, so a legible heuristic beats a first-call search; where a knob forks per device, hang the fork on the existing sm_version / has_native_fp8 queries rather than adding a mechanism.
• Keep unsupported hardware working: a kernel below its device minimum ships a fallback path or a CapabilityRequirement that filters its row out before the launch. quantization/fp8.py falls back to the torch quantiser below sm89 because Triton cannot emit the e4m3 cast there — fp8 has no native MMA on that hardware to begin with.
• Be careful to avoid unintended runtime JIT compilation. A runtime integer scalar's specialization key records divisibility by 16 and equality to 1, so a value that alternates (0 and 1, or an aligned and a ragged stride) compiles a second variant mid-serving. Put unimportant runtime scalars in do_not_specialize.
• The Triton compiler does not guarantee safe ordering when a kernel writes to a pointer and subsequently reads from the same pointer. This pattern must have a tl.debug_barrier() between the write and read. The barrier synchronizes threads in the block; it does not synchronize separate program instances.

Launch and Indexing

• grid[1] and grid[2] must be at most 65,535. Choose or flatten the grid order so those dimensions cannot exceed the limit for supported shapes. For example, num_tokens is commonly 8K or 16K, but users may configure 32K or more. If num_tokens is a grid dimension, it is safe to put it in grid[0] (or tile it). flashattention2_nopad follows this shape: grid = (cdiv(max_seq_len, BLOCK_M), batchs * n_heads, 1), with the bounded batch-head product on axis 1.
• Cast offset arithmetic to int64 before the multiply when the addressed table's element count can exceed 32-bit range. The native kernels open with tl.program_id(0).to(tl.int64) and keep the cast through the row-stride multiply (swiglu.py, vocab_embedding.py, fused_moe.py) — vocab rows and expert weights both qualify. The KV scatter loads an int32 row id from the allocator (update_kv_buffer) and stays correct only while max_rows * row_stride_elements fits in 2**31; check that bound before enlarging a pool, and do not widen a stride instead.
• A [num_tokens, num_heads] grid can be a good low-latency mapping for decode, but it can be very slow for prefill. This repo keeps the two as separate kernels on purpose: flash_decoding maps one program per (batch, head) row (plus a split-K partition axis for long rows), while flashattention2_nopad tiles query blocks against KV blocks so one program covers many query rows. Do not stretch the decode mapping to serve prefill.
• Take strides as arguments and honor the caller's layout. update_kv_buffer walks a combined [tokens, 2 * heads, dim] buffer with K and V halves; the flash kernels read strided views passed down by modules/attention.py. A .contiguous() inside a wrapper turns a free launch into a bandwidth op — if the kernel truly requires contiguity, assert it and let the caller pay for the copy outside the timed path.
• Small scatter kernels state their launch shape (num_warps=1, num_stages=1 in update_kv_buffer / update_kv_index); tile kernels take theirs from resolve_tiles. Launch parameters are part of the configuration that gets measured — do not leave them at defaults in a perf-sensitive path.

Validation

• Check correctness at boundary shapes and at sizes that exercise masks and large offsets. In tests/kernels/ this reads as parametrized cases (decode-single-token, ragged-count) plus negative cases that pin what must not happen — a scatter must leave unselected rows untouched, and destination order must not be assumed monotonic.
• Choose accumulation and intermediate dtypes explicitly. Test numerically difficult inputs, not only random, well-scaled tensors. Verify a scatter by reading rows back through the destination index, not by comparing whole buffers.
• Use the kernel-microbenchmark skill for benchmark construction, measurement, and interpretation.
• Benchmark a sweep of num_tokens covering decode and representative prefill workloads (bench_kv_write.py sweeps the scatter across both; bench_quant_gemm.py isolates the decode end at m=1). Include relevant head counts and dimensions when they affect the launch shape, and do not select an implementation or tuning heuristic from a single setup.
• Gate a latency number on a passing correctness check: an unverified row is excluded by dispatch ranking, so a number measured on it never serves.
• Include compilation overhead when evaluating startup behavior; report steady-state kernel performance separately.

In-repo examples

• rapid_llm/kernels/ops/kvcache/update_kv_buffer.py — the smallest complete kernel: one program per row, int32 row ids from the allocator, K/V halves in one buffer.
• rapid_llm/kernels/ops/quantization/fp8.py — per-token two-pass quantiser with the sm89 floor and the torch fallback; nvfp4.py and w8a16.py for block scales and the e4m3 bit trick.
• rapid_llm/kernels/ops/activation/swiglu.py, rapid_llm/kernels/ops/embeddings/vocab_embedding.py — int64 openings and mask-guarded negative pointers.
• rapid_llm/kernels/ops/tile_policy.py — resolve_tiles / TileTier, the entry point every new launcher's knobs go through.
• tests/kernels/test_kv_cache_ops.py — how a scatter gets verified; benchmarks/kernels/kv_pool.py + microbench.py — how it gets measured.