@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
md 68 lines 3.79 KB
Rendered Raw
1 ---
2 name: triton-kernel-writing
3 description: Write or review Triton kernels for vLLM, with practical guidance for generated-code inspection, launch grids, indexing, specialization, tuning, and representative performance validation.
4 ---
5
6 # Triton Kernel Writing
7
8 ## Implementation
9
10 - Follow the official
11 [Triton semantics](https://triton-lang.org/main/python-api/triton-semantics.html).
12 Check it when behavior may differ from Python or NumPy, especially type
13 promotion, integer division and modulo, casts, broadcasting, and variable
14 scoping.
15 - Use the Triton kernel generated by `torch.compile` as a possible
16 implementation to inspect. Print Inductor's generated code with
17 `TORCH_LOGS="output_code" .venv/bin/python <script>` or enable
18 `torch._logging.set_logs(output_code=True)` before the compiled function
19 runs. Treat generated code as a reference, not as proof of correctness or
20 optimality.
21 - Find reasonable defaults for compile-time knobs such as `BLOCK_SIZE`, or use
22 a small, legible heuristic when workloads need different choices. Use
23 `triton.autotune` only when tuning is critical to performance, such as for a
24 matrix multiplication. Otherwise prioritize simple code and fast startup.
25 - Avoid unintended runtime JIT compilation. Put unimportant runtime integer
26 scalars in `do_not_specialize`, especially values that alternate between 0
27 and 1. Pointer alignment also affects specialization: an offset tensor view
28 can trigger a new variant after warmup used an aligned allocation. For
29 scalar-only pointers whose alignment provides no useful optimization, use
30 `@triton.jit(do_not_specialize_on_alignment=["end_ptr"])`. Retain alignment
31 specialization when vectorized accesses benefit from it. See the
32 [Triton JIT options](https://triton-lang.org/main/python-api/generated/triton.jit.html).
33 - The Triton compiler does not guarantee safe ordering when a kernel writes to
34 a pointer and subsequently reads from the same pointer. This pattern must
35 have a `tl.debug_barrier()` between the write and read. The barrier
36 synchronizes threads in the block; it does not synchronize separate program
37 instances.
38
39 ## Launch and Indexing
40
41 - `grid[1]` and `grid[2]` must be at most 65,535. Choose or flatten the grid
42 order so those dimensions cannot exceed the limit for supported shapes.
43 For example, `num_tokens` is commonly 8K or 16K, but users may configure 32K
44 or more. If `num_tokens` is a grid dimension, it is safe to put it in
45 `grid[0]` (or tile it).
46 - Use `int64` for offset arithmetic when an index can exceed 32-bit range,
47 especially for KV-cache addressing. Cast operands before multiplication or
48 addition so an intermediate does not overflow in 32-bit arithmetic.
49 - A `[num_tokens, num_heads]` grid can be a good low-latency mapping for decode,
50 but it can be very slow for prefill. If the kernel serves prefill, consider
51 tiling tokens or otherwise increasing the work and locality per program.
52
53 ## Validation
54
55 - Check correctness at boundary shapes and at sizes that exercise masks and
56 large offsets. Include offset tensor views, not only fresh allocations.
57 Warm aligned buffers, then exercise offset views with vLLM's
58 `--jit-monitor-mode error --jit-monitor-verbose` to catch missing variants.
59 - Choose accumulation and intermediate dtypes explicitly. Test numerically
60 difficult inputs, not only random, well-scaled tensors.
61 - Use `$kernel-microbenchmark` for benchmark construction, measurement, and
62 interpretation.
63 - Benchmark a sweep of `num_tokens` covering decode and representative prefill
64 workloads. Include relevant head counts and dimensions when they affect the
65 launch shape, and do not select an implementation or tuning heuristic from a
66 single setup.
67 - Include compilation or autotuning overhead when evaluating startup behavior;
68 report steady-state kernel performance separately.