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. - 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](https://triton-lang.org/main/python-api/generated/triton.jit.html). - 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.