triton-kernel-writing

# Triton Kernel Writing ## Implementation - Follow the official [Triton semantics](https://triton-lang.org/main/python-api/triton-semantics.html). 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.