Triton Kernel Writing

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.
• Find reasonable defaults for compile-time knobs such as BLOCK_SIZE, or use a small, legible heuristic when workloads need different choices. Use triton.autotune only when tuning is critical to performance, such as for a matrix multiplication. Otherwise prioritize simple code and fast startup.
• Avoid unintended runtime JIT compilation. Put unimportant runtime integer scalars in do_not_specialize, especially values that alternate between 0 and 1. Pointer alignment also affects specialization: an offset tensor view can trigger a new variant after warmup used an aligned allocation. For scalar-only pointers whose alignment provides no useful optimization, use @triton.jit(do_not_specialize_on_alignment=["end_ptr"]). Retain alignment specialization when vectorized accesses benefit from it. See the Triton JIT options.
• 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).
• Use int64 for offset arithmetic when an index can exceed 32-bit range, especially for KV-cache addressing. Cast operands before multiplication or addition so an intermediate does not overflow in 32-bit arithmetic.
• A [num_tokens, num_heads] grid can be a good low-latency mapping for decode, but it can be very slow for prefill. If the kernel serves prefill, consider tiling tokens or otherwise increasing the work and locality per program.

Validation

• Check correctness at boundary shapes and at sizes that exercise masks and large offsets. Include offset tensor views, not only fresh allocations. Warm aligned buffers, then exercise offset views with vLLM's --jit-monitor-mode error --jit-monitor-verbose to catch missing variants.
• Choose accumulation and intermediate dtypes explicitly. Test numerically difficult inputs, not only random, well-scaled tensors.
• Use $kernel-microbenchmark for benchmark construction, measurement, and interpretation.
• Benchmark a sweep of num_tokens covering decode and representative prefill workloads. 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.
• Include compilation or autotuning overhead when evaluating startup behavior; report steady-state kernel performance separately.