---
name: triton-kernel-writing
description: Write or review Triton kernels for vLLM, with practical guidance for generated-code inspection, launch grids, indexing, specialization, tuning, and representative performance validation.
---

# 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.