[Agents] Add Triton kernel-writing skill (#55019)
Signed-off-by: Woosuk Kwon <woosuk@inferact.ai> Co-authored-by: OpenAI Codex <codex@openai.com>
This commit is contained in:
co-authored by
OpenAI Codex
parent
6c6376a09e
commit
cf3263d571
@@ -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.
|
||||
@@ -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."
|
||||
Reference in New Issue
Block a user