From cf3263d571dfa21e82787356b9506cbf80b68ad6 Mon Sep 17 00:00:00 2001
From: Woosuk Kwon <woosuk@inferact.ai>
Date: Wed, 2 Sep 2026 15:29:16 -0700
Subject: [PATCH] [Agents] Add Triton kernel-writing skill (#55019)

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>
Co-authored-by: OpenAI Codex <codex@openai.com>
---
 .agents/skills/triton-kernel-writing/SKILL.md | 62 +++++++++++++++++++
 .../triton-kernel-writing/agents/openai.yaml  |  4 ++
 2 files changed, 66 insertions(+)
 create mode 100644 .agents/skills/triton-kernel-writing/SKILL.md
 create mode 100644 .agents/skills/triton-kernel-writing/agents/openai.yaml

diff --git a/.agents/skills/triton-kernel-writing/SKILL.md b/.agents/skills/triton-kernel-writing/SKILL.md
new file mode 100644
index 0000000000..6dd6a8520b
--- /dev/null
+++ b/.agents/skills/triton-kernel-writing/SKILL.md
@@ -0,0 +1,62 @@
+---
+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.
diff --git a/.agents/skills/triton-kernel-writing/agents/openai.yaml b/.agents/skills/triton-kernel-writing/agents/openai.yaml
new file mode 100644
index 0000000000..259f16a135
--- /dev/null
+++ b/.agents/skills/triton-kernel-writing/agents/openai.yaml
@@ -0,0 +1,4 @@
+interface:
+  display_name: "Triton Kernel Writing"
+  short_description: "Write robust, practical Triton kernels for vLLM."
+  default_prompt: "Use $triton-kernel-writing to implement or review this Triton kernel."