[Agents] Add Triton kernel-writing skill (#55019)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: OpenAI Codex <codex@openai.com>
Woosuk Kwon committed
Sep 2, 2026 at 15:29 UTC
cf3263d571dfa21e82787356b9506cbf80b68ad6
2 files changed
+66
.agents/skills/triton-kernel-writing/SKILL.md
new
+62
@@ -0,0 +1,62 @@
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
+- Be careful to avoid unintended runtime JIT compilation. For example, put
26
+ unimportant runtime integer scalars in `do_not_specialize`, especially those
27
+ that may alternate between values such as 0 and 1, which can produce
28
+ different specialization keys.
29
+- The Triton compiler does not guarantee safe ordering when a kernel writes to
30
+ a pointer and subsequently reads from the same pointer. This pattern must
31
+ have a `tl.debug_barrier()` between the write and read. The barrier
32
+ synchronizes threads in the block; it does not synchronize separate program
33
+ instances.
34
+
35
+## Launch and Indexing
36
+
37
+- `grid[1]` and `grid[2]` must be at most 65,535. Choose or flatten the grid
38
+ order so those dimensions cannot exceed the limit for supported shapes.
39
+ For example, `num_tokens` is commonly 8K or 16K, but users may configure 32K
40
+ or more. If `num_tokens` is a grid dimension, it is safe to put it in
41
+ `grid[0]` (or tile it).
42
+- Use `int64` for offset arithmetic when an index can exceed 32-bit range,
43
+ especially for KV-cache addressing. Cast operands before multiplication or
44
+ addition so an intermediate does not overflow in 32-bit arithmetic.
45
+- A `[num_tokens, num_heads]` grid can be a good low-latency mapping for decode,
46
+ but it can be very slow for prefill. If the kernel serves prefill, consider
47
+ tiling tokens or otherwise increasing the work and locality per program.
48
+
49
+## Validation
50
+
51
+- Check correctness at boundary shapes and at sizes that exercise masks and
52
+ large offsets.
53
+- Choose accumulation and intermediate dtypes explicitly. Test numerically
54
+ difficult inputs, not only random, well-scaled tensors.
55
+- Use `$kernel-microbenchmark` for benchmark construction, measurement, and
56
+ interpretation.
57
+- Benchmark a sweep of `num_tokens` covering decode and representative prefill
58
+ workloads. Include relevant head counts and dimensions when they affect the
59
+ launch shape, and do not select an implementation or tuning heuristic from a
60
+ single setup.
61
+- Include compilation or autotuning overhead when evaluating startup behavior;
62
+ report steady-state kernel performance separately.
.agents/skills/triton-kernel-writing/agents/openai.yaml
new
+4
@@ -0,0 +1,4 @@
1
+interface:
2
+ display_name: "Triton Kernel Writing"
3
+ short_description: "Write robust, practical Triton kernels for vLLM."
4
+ default_prompt: "Use $triton-kernel-writing to implement or review this Triton kernel."