Triton Kernel Writing

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.
• Be careful to avoid unintended runtime JIT compilation. For example, put unimportant runtime integer scalars in do_not_specialize, especially those that may alternate between values such as 0 and 1, which can produce different specialization keys.
• 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.
• 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.