From 07d386507023373e1475673ae1becf58df440777 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Wed, 2 Sep 2026 19:43:01 +0000 Subject: [PATCH 1/4] [Agents] Import Triton kernel writing skill Co-authored-by: OpenAI Codex Signed-off-by: Woosuk Kwon --- .../skills/kernel-triton-writing/ORIGIN.md | 17 + .agents/skills/kernel-triton-writing/SKILL.md | 362 +++++++++++++++++ .../references/api-core.md | 331 +++++++++++++++ .../references/api-language.md | 284 +++++++++++++ .../references/concepts-semantics.md | 199 +++++++++ .../references/operator-routing.md | 125 ++++++ .../references/patterns-advanced.md | 320 +++++++++++++++ .../references/patterns-basic.md | 238 +++++++++++ .../references/patterns-fusion.md | 349 ++++++++++++++++ .../references/patterns-gemm.md | 292 ++++++++++++++ .../references/troubleshooting.md | 282 +++++++++++++ .../kernel-triton-writing/scripts/__init__.py | 17 + .../scripts/benchmark_kernel.py | 308 ++++++++++++++ .../scripts/verify_kernel.py | 381 ++++++++++++++++++ 14 files changed, 3505 insertions(+) create mode 100644 .agents/skills/kernel-triton-writing/ORIGIN.md create mode 100644 .agents/skills/kernel-triton-writing/SKILL.md create mode 100644 .agents/skills/kernel-triton-writing/references/api-core.md create mode 100644 .agents/skills/kernel-triton-writing/references/api-language.md create mode 100644 .agents/skills/kernel-triton-writing/references/concepts-semantics.md create mode 100644 .agents/skills/kernel-triton-writing/references/operator-routing.md create mode 100644 .agents/skills/kernel-triton-writing/references/patterns-advanced.md create mode 100644 .agents/skills/kernel-triton-writing/references/patterns-basic.md create mode 100644 .agents/skills/kernel-triton-writing/references/patterns-fusion.md create mode 100644 .agents/skills/kernel-triton-writing/references/patterns-gemm.md create mode 100644 .agents/skills/kernel-triton-writing/references/troubleshooting.md create mode 100644 .agents/skills/kernel-triton-writing/scripts/__init__.py create mode 100644 .agents/skills/kernel-triton-writing/scripts/benchmark_kernel.py create mode 100644 .agents/skills/kernel-triton-writing/scripts/verify_kernel.py diff --git a/.agents/skills/kernel-triton-writing/ORIGIN.md b/.agents/skills/kernel-triton-writing/ORIGIN.md new file mode 100644 index 000000000000..a809bf8c9641 --- /dev/null +++ b/.agents/skills/kernel-triton-writing/ORIGIN.md @@ -0,0 +1,17 @@ +# Source and License + +This skill was copied from NVIDIA's TensorRT-LLM repository: + +- Source: `https://github.com/NVIDIA/TensorRT-LLM/tree/main/.claude/skills/kernel-triton-writing` +- Snapshot commit: `395985c025c8d1cf5aa842bc752b337ba88721b6` +- Copyright: Copyright (c) 2011-2026 NVIDIA CORPORATION & AFFILIATES. + All rights reserved. +- License: Apache License 2.0 + +The NVIDIA copyright and Apache-2.0 SPDX notices are preserved in the copied +reference and script files. The vLLM copy adds explicit source comments, +removes unsupported skill metadata, replaces unsafe cache-removal examples, +keeps upstream Markdown tables/code blocks with targeted lint suppressions, +and adapts paths and Python commands to vLLM's `.venv/bin/python` and `uv` +workflow. The benchmark helper uses vLLM's accelerator-neutral synchronization +API. diff --git a/.agents/skills/kernel-triton-writing/SKILL.md b/.agents/skills/kernel-triton-writing/SKILL.md new file mode 100644 index 000000000000..2b5b6b124a2d --- /dev/null +++ b/.agents/skills/kernel-triton-writing/SKILL.md @@ -0,0 +1,362 @@ +--- +name: kernel-triton-writing +description: > + ONLY for OpenAI Triton (@triton.jit) kernel development. NEVER use for + CUDA C++ kernels, TileIR, or profiling tools (ncu, nsys). + The user's request must involve Triton explicitly. Covers Triton-specific + patterns: fused elementwise, reductions (softmax, LayerNorm, RMSNorm), + tiled GEMM with triton.autotune, and flash attention. Workflow: + design, write, verify (with fast-path for explicit requests). +license: Apache-2.0 +metadata: + author: NVIDIA Corporation + source: https://github.com/NVIDIA/TensorRT-LLM + source_commit: 395985c025c8d1cf5aa842bc752b337ba88721b6 +--- + +# Triton Kernel Writing + + + + + +## Principles + +### Correctness First + +1. Never benchmark before verification passes. +2. Always mask loads and stores for non-divisible shapes. +3. Include `kernel_fn`, `reference_fn`, and `get_inputs()` exports for companion scripts. +4. Always run `scripts/verify_kernel.py` to validate against the reference. + +### FP16/BF16 Precision Rules (LOW FREEDOM -- follow exactly) + +Transcendental functions (`tl.exp`, `tl.log`, `tl.math.erf`, `tl.math.tanh`) require fp32 inputs. + +```python +# WRONG -- compilation error or wrong results with fp16/bf16: +result = tl.exp(x) + +# CORRECT -- cast to fp32, compute, cast back: +x_fp32 = x.to(tl.float32) +result = tl.exp(x_fp32).to(x.dtype) +``` + +Rule: any math function beyond basic arithmetic (+, -, *, /) requires fp32 cast in, original dtype cast out. + +Additional precision constraints: + +- `tl.sigmoid()` is unavailable in some Triton versions. Use `1.0 / (1.0 + tl.exp(-x_fp32))`. +- Always cast back to `x.dtype` before `tl.store` -- mismatches cause "Type mismatch, store Float32 to Float16". +- Unlike PyTorch, Triton does NOT auto-promote fp16/bf16 to fp32 for accumulation. Always use `tl.float32` accumulators for `tl.dot`. +- **TF32 for matmul:** On Ampere+/Hopper, `tl.dot` uses TF32 by default for fp32 inputs (same as `torch.mm`). Do NOT add `input_precision="ieee"` — it is 3-8x slower. TF32 is the correct default. If verification fails due to TF32 precision (~0.01-0.1 abs diff), ensure `reference_fn` also uses TF32 (plain `torch.mm`, no `allow_tf32=False`). + +### CPU-GPU Sync Avoidance (LOW FREEDOM) + +Never call `.item()` in kernel wrappers. It forces a CPU-GPU sync (~50-100us per call). + +| Pitfall | Fix | +|---------|-----| +| `tensor.item()` for seed | `x.data_ptr() % (2**31)` | +| `torch.randint(...).item()` | Use tensor metadata for pseudo-random seed | +| Allocating output every call | Accept pre-allocated output as parameter | +| Python loops calling kernel | Batch operations | + +### C Integer Division Semantics (CRITICAL) + +Triton uses **C semantics** (round toward zero) for `//` and `%`, NOT Python semantics (round toward negative infinity). This only matters when operands can be negative. + +| Expression | Python | Triton/C | +|------------|--------|----------| +| `-7 // 2` | `-4` | `-3` | +| `-7 % 2` | `1` | `-1` | + +**Safe pattern:** Ensure all index/offset values are non-negative. If negative values are possible, use `(idx % BLOCK + BLOCK) % BLOCK`. + +See [references/concepts-semantics.md](references/concepts-semantics.md) for full rules and scalar-only exception. + +### Kernel Design Mental Model + +- **Parallelization axis:** Element-wise kernels parallelize over flattened elements. Row-wise kernels (LayerNorm, softmax) parallelize over rows. Matmul kernels tile in 2D (M, N). +- **Block size:** Power-of-2 only (256, 512, 1024, 2048). Start with 1024 for H100, 512 for V100. +- **Memory coalescing:** Adjacent threads must access adjacent memory addresses. The compiler handles this automatically from block-level pointer arithmetic. +- **Grid:** Use `triton.cdiv(n_elements, BLOCK_SIZE)`. With autotune, grid must be a lambda: `lambda meta: (triton.cdiv(n, meta['BLOCK_SIZE']),)`. +- **Decorator order:** `@triton.autotune` (outermost) -> `@triton.heuristics` -> `@triton.jit` (innermost). +- **`reset_to_zero`:** Required for autotune on kernels that accumulate (e.g., matmul output). Without it, later configs see leftover values from earlier trials. + +## Workflow + +**Fast path:** If the user explicitly requests a Triton kernel (e.g., "Write a Triton kernel for X", "Implement softmax in Triton"), start at **Phase 2**. Only use Phase 0-1 when the request is ambiguous about whether Triton is appropriate. + +### Phase 0: Route the Operator (only for ambiguous requests) + +Skip this phase if the user explicitly asks for a Triton kernel. Only use when the request is ambiguous (e.g., "optimize this operation"). + +Triton wins when 2+ operations can share registers instead of writing/reading global memory. Quick rules: + +| Pattern | Decision | +|---------|----------| +| Single element-wise op (`relu`, `sigmoid`) | SKIP — PyTorch already optimal | +| Standalone matmul | SKIP — cuBLAS is optimized | +| Standard attention | SKIP — Use FlashAttention | +| Element-wise chain (2+ ops), reduction, matmul + epilogue | USE TRITON | + +If SKIP, recommend the alternative and STOP. See [references/operator-routing.md](references/operator-routing.md) for edge cases. + +### Phase 1: Analyze the Operator (only for ambiguous requests) + +From the user's request, identify: (1) operation type, (2) parallelization strategy, (3) input shapes and dtypes. + +### Phase 2: Design the Kernel + +Pick the skeleton below that matches your operation. **These skeletons are sufficient for element-wise, reduction, matmul, and fusion kernels — do NOT read reference files for these common patterns.** Only consult `references/` when implementing uncommon patterns (grouped GEMM, TMA, extern functions) or debugging issues. + +**Element-wise skeleton** (GELU, dropout, fused ops on flat tensors): + +```python +@triton.jit +def kernel(x_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(x_ptr + offsets, mask=mask) + # ... compute ... + tl.store(out_ptr + offsets, result, mask=mask) +``` + +**Row-wise skeleton** (softmax, LayerNorm, RMSNorm — one program per row): + +```python +@triton.jit +def kernel(x_ptr, out_ptr, n_cols, BLOCK_SIZE: tl.constexpr): + row_idx = tl.program_id(0) + col_offsets = tl.arange(0, BLOCK_SIZE) + mask = col_offsets < n_cols + x = tl.load(x_ptr + row_idx * n_cols + col_offsets, mask=mask, other=0.0) + # ... reduce / normalize ... + tl.store(out_ptr + row_idx * n_cols + col_offsets, result, mask=mask) +``` + +**Tiled matmul skeleton** (GEMM with 2D tiling, grouped ordering, and autotune): + +```python +@triton.autotune( + configs=[ + triton.Config({'BLOCK_M': 128, 'BLOCK_N': 256, 'BLOCK_K': 64, 'GROUP_SIZE_M': 8}, num_warps=8, num_stages=3), + triton.Config({'BLOCK_M': 64, 'BLOCK_N': 256, 'BLOCK_K': 32, 'GROUP_SIZE_M': 8}, num_warps=4, num_stages=4), + triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 32, 'GROUP_SIZE_M': 8}, num_warps=4, num_stages=4), + triton.Config({'BLOCK_M': 256, 'BLOCK_N': 64, 'BLOCK_K': 32, 'GROUP_SIZE_M': 8}, num_warps=4, num_stages=4), + ], + key=['M', 'N', 'K'], +) +@triton.jit +def matmul_kernel( + a_ptr, b_ptr, c_ptr, M, N, K, + stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn, + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, +): + pid = tl.program_id(0) + num_m_blocks = tl.cdiv(M, BLOCK_M) + num_n_blocks = tl.cdiv(N, BLOCK_N) + # Grouped ordering for L2 cache locality + num_pid_in_group = GROUP_SIZE_M * num_n_blocks + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_m_blocks - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + offs_k = tl.arange(0, BLOCK_K) + + a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak + b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + + for k in range(0, tl.cdiv(K, BLOCK_K)): + a_mask = (offs_m[:, None] < M) & (offs_k[None, :] < K) + b_mask = (offs_k[:, None] < K) & (offs_n[None, :] < N) + a = tl.load(a_ptrs, mask=a_mask, other=0.0) + b = tl.load(b_ptrs, mask=b_mask, other=0.0) + acc += tl.dot(a, b) + a_ptrs += BLOCK_K * stride_ak + b_ptrs += BLOCK_K * stride_bk + offs_k += BLOCK_K + + c_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N) + c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn + tl.store(c_ptrs, acc.to(c_ptr.dtype.element_ty), mask=c_mask) +``` + +### Phase 3: Write the Kernel + +Create an output directory, then write the kernel file to `{output_dir}/kernel.py`. + +The kernel file MUST include: + +- `@triton.jit` decorated kernel function +- `@triton.autotune` for production kernels (see [references/api-core.md](references/api-core.md)) +- Python wrapper function (descriptive name for external import) +- **Fixed contract exports** (companion scripts rely on these exact names): + - `kernel_fn` — alias to the wrapper function + - `reference_fn(*args)` — PyTorch reference with identical signature + - `get_inputs()` — returns `list` of fresh CUDA tensors for testing/benchmarking + +Concise example (fused GELU + dropout): + +```python +import triton +import triton.language as tl +import torch + +@triton.autotune( + configs=[ + triton.Config({'BLOCK_SIZE': 1024}, num_warps=4), + triton.Config({'BLOCK_SIZE': 2048}, num_warps=8), + ], + key=['n_elements'], +) +@triton.jit +def fused_gelu_dropout_kernel( + x_ptr, out_ptr, n_elements, p, seed, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + + x = tl.load(x_ptr + offsets, mask=mask) + x_fp32 = x.to(tl.float32) + x = (0.5 * x_fp32 * (1.0 + tl.math.erf(x_fp32 * 0.7071067811865476))).to(x.dtype) + + random = tl.rand(seed, offsets) + x = tl.where(random > p, x / (1.0 - p), 0.0) + + tl.store(out_ptr + offsets, x, mask=mask) + + +def fused_gelu_dropout_triton(x: torch.Tensor, p: float = 0.1) -> torch.Tensor: + n_elements = x.numel() + out = torch.empty_like(x) + grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),) + seed = (x.data_ptr() % (2**31)) ^ n_elements # sync-free seed + fused_gelu_dropout_kernel[grid](x, out, n_elements, p, seed) + return out + + +# --- Fixed contract (companion scripts rely on these names) --- +kernel_fn = fused_gelu_dropout_triton + +def reference_fn(x, p=0.1): + torch.manual_seed((x.data_ptr() % (2**31)) ^ x.numel()) + return torch.nn.functional.dropout( + torch.nn.functional.gelu(x), p, training=True + ) + +def get_inputs(): + return [torch.randn(128 * 1024 * 1024, device="cuda")] +``` + +For more patterns (SiLU+mul, RMSNorm, linear+GELU, add+LayerNorm), see [references/patterns-fusion.md](references/patterns-fusion.md). For GEMM patterns, see [references/patterns-gemm.md](references/patterns-gemm.md). + +### Phase 4: Verify Correctness + +Run the companion verification script: + +```bash +.venv/bin/python \ + .agents/skills/kernel-triton-writing/scripts/verify_kernel.py \ + {output_dir}/kernel.py \ + --rtol 1e-3 --atol 1e-3 +``` + +Output: + +```json +{"correct": true, "max_abs_diff": 1.2e-7, "max_rel_diff": 3.4e-6, "details": "..."} +``` + +**Stop if `correct: false`.** Fix the kernel before benchmarking. + +**Tolerance guide:** + +| Dtype | rtol | atol | Notes | +|-------|------|------|-------| +| float16 | 1e-3 | 1e-3 | | +| bfloat16 | 1e-2 | 1e-2 | | +| float32 | 1e-5 | 1e-5 | Element-wise ops | +| float32 (matmul) | 1e-2 | 1e-1 | TF32 accumulation order differs between Triton tiles and cuBLAS | + +### Phase 5: Benchmark Performance (optional) + +Only benchmark if the user explicitly requests performance numbers. Skip this phase for correctness-focused requests. + +```bash +.venv/bin/python \ + .agents/skills/kernel-triton-writing/scripts/benchmark_kernel.py \ + {output_dir}/kernel.py +``` + +Output: + +```json +{"kernel_time_ms": 0.45, "reference_time_ms": 1.23, "speedup": 2.73, "warmup_iters": 10, "benchmark_iters": 40} +``` + +## References (consult only when stuck) + +The skeletons and principles above cover element-wise, reduction, matmul, and fusion kernels. **Do NOT read reference files for these common patterns.** + +Only consult `references/` when: + +- Implementing **uncommon patterns** (grouped GEMM, TMA, persistent matmul, extern functions) +- **Debugging** a compile error or incorrect result not covered by the error table below +- Needing **API details** for an unfamiliar `tl.*` operation + +**How to search:** Grep for your keyword across `references/`. Read only the file Grep points to. + +| File | When to use | +|---|---| +| `references/api-core.md` | Unfamiliar `triton.autotune` / `triton.Config` options | +| `references/api-language.md` | Unfamiliar `tl.*` operations | +| `references/patterns-gemm.md` | Grouped GEMM, persistent matmul, TMA, MX formats | +| `references/patterns-advanced.md` | Flash attention details, backward passes, libdevice | +| `references/troubleshooting.md` | Debug ops, interpreter mode, env vars | + +## Error Handling and Troubleshooting + +### Common Errors + +| Error / Symptom | Cause | Fix | +|---------|-------|-----| +| "Type mismatch, store Float32 to Float16" | Missing `.to(x.dtype)` before store | Cast fp32 result back | +| `BLOCK_SIZE is not a constexpr` | Block size passed as runtime value | Add `: tl.constexpr` annotation | +| `shape mismatch` in binary op | Tensor shapes don't broadcast | Check with `tl.static_print`; use `[:, None]` / `[None, :]` | +| Large diffs everywhere | Wrong dtype in `tl.load` | Check load dtype matches input | +| Matmul 3-8x slower than expected | `input_precision="ieee"` on `tl.dot` | Remove it; use TF32 default. Ensure `reference_fn` also uses TF32 | +| Matmul ~0.01-0.1 abs diff vs reference | TF32 vs IEEE mismatch | Use same precision in both kernel and reference (TF32 for both) | +| Diffs at boundaries | Missing mask | Add mask to all load/store ops | +| Random diffs | Race condition | Check atomics and ordering | +| NaN/Inf | Division by zero or fp16 overflow | Guard with epsilon; use `tl.float32` accumulator | +| `grid must be a tuple` | Grid lambda returns int, not tuple | Return `(value,)` with trailing comma | +| `expected constexpr` in `tl.arange` | Non-constexpr argument | Both args of `tl.arange(start, end)` must be constexpr | +| `triton.OutOfResources` | Register/shared memory pressure | Reduce BLOCK_SIZE or `num_stages` | +| Kernel not updating after edit | Stale compilation cache | Move the confirmed Triton cache directory aside and retry | +| Mismatched results vs PyTorch | C integer division semantics | Triton uses truncation; see `references/concepts-semantics.md` | + +For extended error table, interpreter mode issues, and environment variables, see [references/troubleshooting.md](references/troubleshooting.md). + +### When to Abort + +Stop and report failure if: + +1. **Not a good fit** -- Pure matmul or complex control flow (Phase 0 should catch this). +2. **Verification fails after 3 attempts** -- Numerical issues too severe to fix. +3. **No speedup** -- Reference is already well-optimized (cuBLAS, cuDNN). +4. **Hardware mismatch** -- Target GPU not available for testing. diff --git a/.agents/skills/kernel-triton-writing/references/api-core.md b/.agents/skills/kernel-triton-writing/references/api-core.md new file mode 100644 index 000000000000..ed3c16685af8 --- /dev/null +++ b/.agents/skills/kernel-triton-writing/references/api-core.md @@ -0,0 +1,331 @@ + + + + +# Triton Core API Reference + +## triton.jit + +Decorator that JIT-compiles a function into a GPU kernel using the Triton compiler. + +### Signature + +```python +@triton.jit # simple form — no parens needed +@triton.jit(do_not_specialize=None, do_not_specialize_on_alignment=None, + debug=None, noinline=None, repr=None, launch_metadata=None) +``` + +| Param | Type | Purpose | +|-------|------|---------| +| `do_not_specialize` | `Iterable[int\|str]\|None` | Args to skip value-specialization (by index or name) | +| `do_not_specialize_on_alignment` | `Iterable[int\|str]\|None` | Args to skip alignment-specialization | +| `debug` | `bool\|None` | Enable interpreter mode / debug prints | +| `noinline` | `bool\|None` | Prevent inlining when called from another jit'd function | + +### Implicit Pointer Conversion + +Objects with both `.data_ptr()` and `.dtype` (e.g., PyTorch tensors) are auto-converted +to device pointers. You never call `.data_ptr()` yourself in the launch call: + +```python +@triton.jit +def add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr): + pid = tl.program_id(0) + offs = pid * BLOCK + tl.arange(0, BLOCK) + mask = offs < n + tl.store(out_ptr + offs, tl.load(x_ptr + offs, mask=mask) + tl.load(y_ptr + offs, mask=mask), mask=mask) + +# Launch — pass tensors directly, NOT x.data_ptr() +add_kernel[(grid,)](x, y, out, x.numel(), BLOCK=1024) +``` + +### Specialization Rules + +Triton recompiles a kernel when argument properties change. For each argument: + +| Arg type | Specialized on | Effect | +|----------|---------------|--------| +| Pointer (tensor) | 16-byte alignment of `.data_ptr()` | Enables vectorized loads/stores | +| Integer scalar | Whether value == 1 | Dead-code elimination for guards | +| Integer scalar | Whether value is divisible by 16 | Enables optimized indexing | +| `tl.constexpr` | Exact value | Baked into compiled code as literal | + +**Gotcha:** Each unique specialization signature triggers a full recompile. If a size arg +oscillates between aligned/unaligned values, you get 2 cached versions (fine). But if you +pass truly random integers, use `do_not_specialize` to avoid cache explosion: + +```python +@triton.jit(do_not_specialize=["stride_x"]) +def my_kernel(x_ptr, stride_x, BLOCK: tl.constexpr): + ... +``` + +### constexpr Parameters + +Annotate with `tl.constexpr` to make a param a compile-time constant. Required for +values used in `tl.arange()`, `tl.zeros()`, tensor shapes, and `tl.static_assert`. +Each distinct value triggers recompilation. + +```python +@triton.jit +def kernel(x_ptr, N: tl.constexpr, BLOCK_SIZE: tl.constexpr): + ... +``` + +--- + +## triton.autotune + +Decorator that benchmarks multiple `triton.Config`s and caches the fastest per key. + +### Signature + +```python +@triton.autotune( + configs: list[triton.Config], + key: list[str], + prune_configs_by: dict | None = None, + reset_to_zero: list[str] | None = None, + restore_value: list[str] | None = None, + warmup: int = 25, + rep: int = 100, + use_cuda_graph: bool = False, +) +``` + +| Param | Purpose | +|-------|---------| +| `configs` | List of `triton.Config` objects to benchmark | +| `key` | Arg names whose values form the cache key (e.g., `["M", "N", "K"]`) | +| `prune_configs_by` | Dict with `early_config_prune`, `perf_model`, `top_k` to reduce search space | +| `reset_to_zero` | Arg names zeroed before each config trial (for accumulator correctness) | +| `restore_value` | Arg names restored to original value after each trial | +| `warmup` | Warmup time in ms per config (default 25) | +| `rep` | Benchmark time in ms per config (default 100) | + +### Example + +```python +@triton.autotune( + configs=[ + triton.Config({"BLOCK_M": 128, "BLOCK_N": 128}, num_warps=4, num_stages=3), + triton.Config({"BLOCK_M": 64, "BLOCK_N": 256}, num_warps=8, num_stages=3), + triton.Config({"BLOCK_M": 256, "BLOCK_N": 64}, num_warps=4, num_stages=4), + ], + key=["M", "N", "K"], + reset_to_zero=["c_ptr"], # zero output buffer between trials +) +@triton.jit +def matmul_kernel(a_ptr, b_ptr, c_ptr, M, N, K, + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr): + ... +``` + +### Debugging Autotuning + +```bash +TRITON_PRINT_AUTOTUNING=1 .venv/bin/python my_script.py +``` + +Prints the winning config and total tuning time per kernel to stdout. + +### prune_configs_by Details + +```python +def early_prune(configs, named_args, **kwargs): + """Drop configs that exceed shared memory or have bad aspect ratios.""" + M, N = named_args["M"], named_args["N"] + return [c for c in configs if c.kwargs["BLOCK_M"] <= M and c.kwargs["BLOCK_N"] <= N] + +@triton.autotune( + configs=[...], + key=["M", "N"], + prune_configs_by={"early_config_prune": early_prune}, +) +``` + +Fields: `early_config_prune(configs, named_args, **kwargs) -> list[Config]`, +`perf_model(named_args, config, **kwargs) -> float` (estimated time), +`top_k` (int, keep only top_k configs after perf_model ranking). + +### Gotchas + +- `reset_to_zero` is critical for kernels that accumulate (e.g., matmul output). + Without it, later configs see leftover values from earlier trials. +- Autotuning happens on first call with each unique key combination. Subsequent calls + with the same key values use the cached winner. +- Decorator order: `@triton.autotune` must be the outermost, then `@triton.heuristics` + (if used), then `@triton.jit` innermost. + +--- + +## triton.Config + +Represents one candidate configuration for `triton.autotune`. + +### Signature + +```python +triton.Config( + kwargs: dict[str, Any], + num_warps: int = 4, + num_stages: int = 3, + num_ctas: int = 1, + maxnreg: int | None = None, + pre_hook: Callable | None = None, +) +``` + +| Param | Default | Purpose | +|-------|---------|---------| +| `kwargs` | (required) | Dict mapping `tl.constexpr` param names to values | +| `num_warps` | 4 | Threads per block = `num_warps * 32` | +| `num_stages` | 3 | Software pipelining depth for global loads | +| `num_ctas` | 1 | Cooperative thread arrays (multi-CTA kernels, Hopper+) | +| `maxnreg` | None | Max registers per thread (trades occupancy vs spilling) | +| `pre_hook` | None | `fn(args: dict)` called before kernel launch | + +### Example with pre_hook + +```python +def zero_output(args): + """Zero the output tensor before the kernel runs.""" + args["c_ptr"].zero_() + +triton.Config( + {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32}, + num_warps=4, + num_stages=3, + pre_hook=zero_output, +) +``` + +### Tuning Guidance + +| Parameter | Small tiles / low occupancy | Large tiles / high throughput | +|-----------|----------------------------|-------------------------------| +| `num_warps` | 2-4 | 8-16 | +| `num_stages` | 2 (less shared mem) | 3-5 (hide global latency) | +| `maxnreg` | None (let compiler decide) | 128-255 (force occupancy) | + +**Gotcha:** `num_stages > 1` requires shared memory for buffering. Large tiles + +many stages can exceed shared memory limits, causing silent fallback or launch failure. + +### GPU-Specific Config Guidelines + +**H100 (Hopper):** HBM3, 168 SMs, large shared memory. + +- Prefer larger blocks (1024-4096), more warps (8-16), `num_stages=4+`. + +**A100 (Ampere):** Balanced config. + +- Block sizes 512-2048, `num_stages=3` typically optimal. + +**V100 (Volta):** Less shared memory. + +- Smaller blocks (256-1024), fewer stages (2), warps 4-8. + +--- + +## triton.heuristics + +Decorator that computes meta-parameters from kernel arguments at launch time, +avoiding the cost of autotuning for values that can be derived deterministically. + +### Signature + +```python +@triton.heuristics(values: dict[str, Callable]) +``` + +`values` maps constexpr parameter names to functions. Each function receives +the kernel's named arguments as a dict and returns the computed value. + +### Example + +```python +@triton.heuristics( + values={ + "BLOCK_SIZE": lambda args: triton.next_power_of_2(args["n_cols"]), + "num_warps": lambda args: 4 if args["n_cols"] <= 1024 else 8, + } +) +@triton.jit +def softmax_kernel(x_ptr, out_ptr, n_cols, + BLOCK_SIZE: tl.constexpr): + ... + +# Launch — BLOCK_SIZE is computed automatically, not passed +softmax_kernel[(n_rows,)](x, out, n_cols) +``` + +### Combined with autotune + +```python +@triton.autotune( + configs=[ + triton.Config({"BLOCK_M": 64}, num_warps=4), + triton.Config({"BLOCK_M": 128}, num_warps=8), + ], + key=["M", "N"], +) +@triton.heuristics( + values={"BLOCK_N": lambda args: triton.next_power_of_2(args["N"])} +) +@triton.jit +def kernel(x_ptr, M, N, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr): + ... +``` + +**Required order (outermost to innermost):** `@autotune` -> `@heuristics` -> `@jit` + +### triton.next_power_of_2 + +```python +triton.next_power_of_2(n) # 7 -> 8, 8 -> 8, 1000 -> 1024 +``` + +Host-side utility. Common pattern: derive BLOCK_SIZE from a problem dimension +so the block covers the full row/column in one pass. + +### Gotchas + +- Heuristic functions run on the host (CPU) at every kernel launch, not on GPU. +- `@heuristics` must come AFTER `@autotune` but BEFORE `@jit` in decorator stack. +- Values computed by heuristics override any same-named values in `triton.Config.kwargs`. +- Returning non-power-of-2 for a `BLOCK_*` param is valid but usually suboptimal. + +--- + +## Decorator Stacking Summary + +``` +@triton.autotune(...) # outermost — optional +@triton.heuristics(...) # middle — optional +@triton.jit # innermost — required +def kernel(...): +``` + +| Combo | Use case | +|-------|----------| +| `@jit` only | Fixed config, simplest kernels | +| `@autotune` + `@jit` | Search over tile sizes and hardware params | +| `@heuristics` + `@jit` | Derive config from args, no search needed | +| `@autotune` + `@heuristics` + `@jit` | Search some params, derive others | diff --git a/.agents/skills/kernel-triton-writing/references/api-language.md b/.agents/skills/kernel-triton-writing/references/api-language.md new file mode 100644 index 000000000000..02ccb1061b45 --- /dev/null +++ b/.agents/skills/kernel-triton-writing/references/api-language.md @@ -0,0 +1,284 @@ + + + + +# Triton Language API (`triton.language` / `tl`) + +## Programming Model + +| Function | Signature | Notes | +|---|---|---| +| `program_id` | `program_id(axis)` | Returns ID of current program instance along `axis` (0, 1, or 2) | +| `num_programs` | `num_programs(axis)` | Returns number of program instances along `axis` | +| `tensor` | N-D array type | Block-structured; all ops are implicitly vectorized over the block | +| `tensor_descriptor` | Returned by `make_tensor_descriptor` | Opaque handle; backed by TMA on supported NVIDIA GPUs | + +## Creation Operations + +| Function | Signature | Notes | +|---|---|---| +| `arange` | `arange(start, end)` | Half-open `[start, end)`, returns 1-D int32 tensor | +| `full` | `full(shape, value, dtype)` | Broadcast scalar `value` to `shape` | +| `zeros` | `zeros(shape, dtype)` | Shorthand for `full(shape, 0, dtype)` | +| `zeros_like` | `zeros_like(x)` | Zeros with same shape/dtype as `x` | +| `cat` | `cat(x, y, can_reorder=False)` | Concatenate along dim 0; `can_reorder` allows compiler flexibility | +| `cast` | `cast(x, dtype, fp_downcast_rounding="rtne")` | Type conversion; rounding: `"rtne"` (default) or `"rtz"` | + +## Memory Operations (Pointer-based) + +### `tl.load` -- most-used memory op + +```python +load(pointer, mask=None, other=None, boundary_check=(), + padding_option='', cache_modifier='', eviction_policy='', volatile=False) +``` + +**Key semantics:** + +- `mask`: block of `int1`. Where False, returns `other` (default 0). Required for out-of-bounds safety. +- `other`: fallback value where mask is False. Must match dtype. +- `boundary_check`: tuple of dims for block-pointer bounds checking (mutually exclusive with `mask`). +- `padding_option`: `"zero"` or `"nan"` (only with `boundary_check`). +- `cache_modifier`: `""`, `".cg"`, `".cs"`, `".ca"`, `".wb"`, `".wt"`. +- `eviction_policy`: `""`, `"evict_first"`, `"evict_last"`. + +```python +# Typical masked load pattern +offs = pid * BLOCK + tl.arange(0, BLOCK) +mask = offs < n_elements +x = tl.load(ptr + offs, mask=mask, other=0.0) +``` + +### `tl.store` + +```python +store(pointer, value, mask=None, boundary_check=(), + cache_modifier='', eviction_policy='') +``` + +Same mask semantics as load. Where mask is False, store is skipped (no side effect). + +```python +tl.store(out_ptr + offs, result, mask=mask) +``` + +## Memory Operations (Block Pointer) + +| Function | Signature | Notes | +|---|---|---| +| `make_block_ptr` | `(base, shape, strides, offsets, block_shape, order)` | Structured pointer; `order` controls memory layout (e.g., `(1,0)` for col-major) | +| `advance` | `advance(block_ptr, offsets)` | Returns NEW ptr (no mutation); `offsets` is tuple by dim | + +```python +a_ptr = tl.make_block_ptr(a, (M, K), (stride_am, stride_ak), (pid_m * BM, 0), (BM, BK), order=(1, 0)) +a_ptr = tl.advance(a_ptr, (0, BK)) # advance K dimension +a = tl.load(a_ptr, boundary_check=(0, 1)) +``` + +## Memory Operations (Tensor Descriptor / TMA) + +```python +make_tensor_descriptor(base, shape, strides, block_shape, padding_option="zero") +# base must be 16-byte aligned. Supports 2-5D tensors. +# On NVIDIA GPUs with TMA, uses hardware TMA descriptor. +``` + +| Function | Signature | Notes | +|---|---|---| +| `tensor_descriptor.load` | `.load(offsets, boundary_check=True)` | Load block at `offsets` from descriptor | +| `tensor_descriptor.store` | `.store(offsets, value)` | Store block at `offsets` | + +## Linear Algebra + +### `tl.dot` + +```python +dot(input, other, acc=None, input_precision="tf32", max_num_imprecise_acc=None, out_dtype=float32) +``` + +- Both operands must be 2-D or 3-D (batched matmul). Inner dims must match (min 16). +- `input` dtype: int8, float8_e5m2, float8_e4m3fn, float16, bfloat16, float32. +- `input_precision`: `"tf32"` (default, NVIDIA), `"tf32x3"`, `"ieee"`. +- `acc`: accumulator tensor; if provided, result is added to it. + +### `tl.dot_scaled` (Microscaling / MX formats) + +```python +dot_scaled(lhs, lhs_scale, lhs_format, rhs, rhs_scale, rhs_format, + acc=None, out_dtype=float32) +``` + +- Formats: `"e2m1"`, `"e4m3"`, `"e5m2"`, `"bf16"`, `"fp16"`. +- Scales are e8m0 (uint8 tensors), shape `[M, K//group_size]`. + +## Math Operations + +| Function | Signature | Notes | +|---|---|---| +| `abs` | `abs(x)` | Elementwise absolute value | +| `cdiv` | `cdiv(x, div)` | Ceiling division: `(x + div - 1) // div` | +| `ceil` | `ceil(x)` | Ceiling (float) | +| `floor` | `floor(x)` | Floor (float) | +| `exp` | `exp(x)` | Base-e exponential | +| `exp2` | `exp2(x)` | Base-2 exponential | +| `log` | `log(x)` | Natural logarithm | +| `log2` | `log2(x)` | Base-2 logarithm | +| `cos` | `cos(x)` | Cosine | +| `sin` | `sin(x)` | Sine | +| `sqrt` | `sqrt(x)` | Square root | +| `rsqrt` | `rsqrt(x)` | Reciprocal square root: `1/sqrt(x)` | +| `sigmoid` | `sigmoid(x)` | `1 / (1 + exp(-x))` | +| `softmax` | `softmax(x, axis)` | Numerically-stable softmax along `axis` | +| `umulhi` | `umulhi(x, y)` | Upper 32 bits of `x * y` (uint32) | +| `fdiv` | `fdiv(x, y, ieee_rounding=False)` | Floating-point division | +| `fma` | `fma(x, y, z)` | Fused multiply-add: `x * y + z` | +| `clamp` | `clamp(x, min, max)` | Clamp to range `[min, max]` | +| `minimum` | `minimum(x, y)` | Elementwise min (propagates NaN) | +| `maximum` | `maximum(x, y)` | Elementwise max (propagates NaN) | + +## Where (Critical for Masking) + +```python +where(condition, x, y) +``` + +Returns elements from `x` where `condition` is True, else from `y`. Both `x` and `y` are broadcast to `condition`'s shape. This is the primary tool for conditional logic in Triton. + +```python +# Causal mask in attention +mask = offs_m[:, None] >= offs_n[None, :] +attn = tl.where(mask, attn, float("-inf")) +``` + +## Reduction Operations + +All reductions: `fn(input, axis=None, keep_dims=False)`. When `axis=None`, reduces all dims. + +| Function | Signature | Notes | +|---|---|---| +| `max` | `max(input, axis, keep_dims=False)` | Maximum along axis | +| `min` | `min(input, axis, keep_dims=False)` | Minimum along axis | +| `argmax` | `argmax(input, axis)` | Index of max along axis | +| `argmin` | `argmin(input, axis)` | Index of min along axis | +| `sum` | `sum(input, axis, keep_dims=False, dtype=None)` | Sum; int/bool auto-upcast to int32, float to float32 | +| `xor_sum` | `xor_sum(input, axis)` | XOR reduction along axis | +| `reduce` | `reduce(input, axis, combine_fn, keep_dims=False)` | Generic reduction with user-defined `combine_fn(a, b) -> c` | + +```python +# Reduction pattern: online softmax +row_max = tl.max(row, axis=1, keep_dims=True) +row = tl.exp(row - row_max) +row_sum = tl.sum(row, axis=1, keep_dims=True) +``` + +## Scan and Sort Operations + +| Function | Signature | Notes | +|---|---|---| +| `associative_scan` | `associative_scan(input, axis, combine_fn, reverse=False)` | Prefix scan with user-defined associative `combine_fn` | +| `cumsum` | `cumsum(input, axis, dtype=None)` | Cumulative sum (specialization of scan) | +| `cumprod` | `cumprod(input, axis, dtype=None)` | Cumulative product | +| `histogram` | `histogram(input, num_bins)` | Counts per bin; input values are bin indices | +| `sort` | `sort(input, axis=-1, descending=False, stable=True)` | Sort along axis | +| `topk` | `topk(input, k, axis=-1, descending=True)` | Top-k values along axis | +| `gather` | `gather(input, indices, axis)` | Gather elements along axis using indices | + +## Atomic Operations + +All atomics: `atomic_*(pointer, val, mask=None, sem="acq_rel", scope="gpu")`. + +- `sem`: `"acquire"`, `"release"`, `"acq_rel"` (default), `"relaxed"`. +- `scope`: `"gpu"` (default), `"cta"` (thread block), `"sys"` (system). + +| Function | Signature | Notes | +|---|---|---| +| `atomic_add` | `(ptr, val, mask=None, sem, scope)` | Atomic add; returns old value | +| `atomic_max` | `(ptr, val, mask=None, sem, scope)` | Atomic max; returns old value | +| `atomic_min` | `(ptr, val, mask=None, sem, scope)` | Atomic min; returns old value | +| `atomic_and` | `(ptr, val, mask=None, sem, scope)` | Atomic bitwise AND | +| `atomic_or` | `(ptr, val, mask=None, sem, scope)` | Atomic bitwise OR | +| `atomic_xor` | `(ptr, val, mask=None, sem, scope)` | Atomic bitwise XOR | +| `atomic_xchg` | `(ptr, val, mask=None, sem, scope)` | Atomic exchange; returns old value | +| `atomic_cas` | `(ptr, cmp, val, sem, scope)` | Compare-and-swap: if `*ptr == cmp`, set to `val`; returns old value | + +## Random Number Generation (Philox PRNG) + +| Function | Signature | Notes | +|---|---|---| +| `randint4x` | `randint4x(seed, offset)` | Returns 4 blocks of int32; fastest for multiple streams | +| `randint` | `randint(seed, offset, n_rounds=6)` | Single block of random int32 | +| `rand` | `rand(seed, offset, n_rounds=6)` | Uniform float32 in `[0, 1)` | +| `randn` | `randn(seed, offset, n_rounds=6)` | Normal distribution (float32) | + +`seed`: scalar int32. `offset`: block of int32 (determines which element gets which random value). + +## Iterators + +| Function | Signature | Notes | +|---|---|---| +| `range` | `range(start, stop, step=1)` | Dynamic loop; bounds can be runtime values | +| `static_range` | `static_range(start, stop, step=1)` | Fully unrolled at compile time; bounds must be `constexpr` | + +## Debug Operations + +| Function | Signature | Notes | +|---|---|---| +| `static_print` | `static_print(*args)` | Print at **compile time**; same interface as Python `print` | +| `static_assert` | `static_assert(cond, msg="")` | Assert at **compile time** | +| `device_print` | `device_print(prefix, *args)` | Print at **runtime** on device; first arg must be string, rest are scalars/tensors | +| `device_assert` | `device_assert(cond, msg="")` | Assert at **runtime**; requires `TRITON_DEBUG=1` env var | + +## Compiler Hints + +| Function | Signature | Notes | +|---|---|---| +| `assume` | `assume(cond)` | Hint to backend for address calculation optimization | +| `max_contiguous` | `max_contiguous(input, values)` | Declare max contiguous extent per dim; enables coalesced access | +| `max_constancy` | `max_constancy(input, values)` | Declare max constant extent per dim | +| `multiple_of` | `multiple_of(input, values)` | Declare that values are multiples of given constants | +| `debug_barrier` | `debug_barrier()` | Thread barrier (debugging only; not for correctness) | + +## Shape Manipulation + +| Function | Signature | Notes | +|---|---|---| +| `broadcast` | `broadcast(x, y)` | Broadcast `x` and `y` to compatible shape (returns both) | +| `broadcast_to` | `broadcast_to(x, shape)` | Broadcast `x` to explicit `shape` | +| `expand_dims` | `expand_dims(x, axis)` | Insert length-1 dim at `axis` | +| `reshape` | `reshape(x, shape)` | Reshape (total elements must match) | +| `view` | `view(x, shape)` | Like reshape; bitcast semantics | +| `trans` | `trans(x, *dims)` | Transpose; default swaps last two dims | +| `permute` | `permute(x, *dims)` | Reorder dims; `permute(x, 2, 1, 0)` or `permute(x, (2,1,0))` | +| `ravel` | `ravel(x)` | Flatten to 1-D | +| `split` | `split(x)` | Split first dim into separate tensors | +| `join` | `join(x, y)` | Concatenate along a new innermost dim | +| `interleave` | `interleave(x, y)` | Interleave elements from `x` and `y` | + +## Inline Assembly + +```python +inline_asm_elementwise(asm, constraints, args, dtype, is_pure, pack) +``` + +- `asm`: PTX/ASM string with `$0`, `$1`, ... placeholders. +- `constraints`: register constraint string (e.g., `"=r,r"` for one int output, one int input). +- `args`: list of input tensors. +- `dtype`: output dtype (or tuple for multi-output). +- `is_pure`: True if no side effects (enables CSE). +- `pack`: number of elements per register (typically 1). diff --git a/.agents/skills/kernel-triton-writing/references/concepts-semantics.md b/.agents/skills/kernel-triton-writing/references/concepts-semantics.md new file mode 100644 index 000000000000..c74081b45b2a --- /dev/null +++ b/.agents/skills/kernel-triton-writing/references/concepts-semantics.md @@ -0,0 +1,199 @@ + + + + +# Triton Concepts and Semantics + +## Programming Model — Block-Based Execution + +Triton programs operate on **blocks** (tiles) of data, not individual scalar threads. +Each kernel instance (called a "program") processes an entire block of elements at once. + +| Concept | CUDA | Triton | +|---------|------|--------| +| Execution unit | Single scalar thread | Program operating on a block | +| Memory coalescing | Manual (stride patterns) | Automatic (compiler) | +| Shared memory | Manual (`__shared__`, sync) | Automatic (compiler) | +| Vectorization | Manual (float4, etc.) | Automatic (compiler) | +| Tensor core usage | Manual (wmma/mma) | Automatic (compiler) | +| Thread synchronization | Manual (`__syncthreads`) | Not needed | + +### Launch Grid and Program IDs + +```python +@triton.jit +def kernel(X_ptr, Y_ptr, N, BLOCK_SIZE: tl.constexpr): + # Each program instance gets a unique ID along each grid axis + pid = tl.program_id(axis=0) # which block of data this program handles + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < N + x = tl.load(X_ptr + offsets, mask=mask) + tl.store(Y_ptr + offsets, x * 2, mask=mask) + +# Launch with a 1D grid: one program per block of data +grid = lambda meta: (triton.cdiv(N, meta['BLOCK_SIZE']),) +kernel[grid](x_ptr, y_ptr, N, BLOCK_SIZE=1024) +``` + +### Key Takeaway + +The programmer thinks in blocks. The compiler decides how to map blocks to warps, +how to stage data through shared memory, and when to use tensor cores. This is the +core design tradeoff: less control, but far less boilerplate and fewer correctness bugs. + +--- + +## Type Promotion Rules + +Triton applies automatic type promotion for binary ops and `tl.where` (last two args). + +### Promotion Hierarchy + +``` +{bool} < {int8, int16, int32, int64, uint8, uint16, uint32, uint64} < {fp8, fp16, bf16, fp32, fp64} + ^ ^ (integral types) ^ (floating types) + kind 0 kind 1 kind 2 +``` + +### Rules Applied in Order + +| Priority | Rule | Example | +|----------|------|---------| +| 1 | **Cross-kind**: lower kind promotes to higher kind's dtype | `(int32, bf16)` -> `bf16` | +| 2 | **Same-kind widening**: narrower promotes to wider | `(fp16, fp32)` -> `fp32` | +| 3 | **Same-width float tie**: bf16 and fp16 both promote to fp16 | `(fp16, bf16)` -> `fp16` | +| 4 | **Same-width sign tie**: promote to unsigned | `(int32, uint32)` -> `uint32` | + +### Scalar-Tensor Interaction + +When a Python scalar interacts with a Triton tensor: + +| Scalar Type | Tensor Type | Result | +|-------------|-------------|--------| +| Python `int` | Any int tensor | Tensor's dtype (no widening) | +| Python `int` | Any float tensor | Tensor's dtype | +| Python `float` | Any float tensor | Tensor's dtype | +| Python `float` | Any int tensor | `fp64` (float is higher kind) | + +### Gotchas: Type Promotion + +| Gotcha | Detail | +|--------|--------| +| `int32 + bf16` -> `bf16` | Integer silently truncated to bf16 precision (only ~3 decimal digits) | +| `fp16 + bf16` -> `fp16` | bf16 promotes to fp16, NOT fp32. May lose bf16 range | +| `int8 + uint8` -> `uint8` | Signed values become unsigned, wrapping negative values | +| Cross-kind hides widening | `(int64, fp16)` -> `fp16`, losing 64-bit integer precision | +| No implicit fp32 promotion | Unlike PyTorch, Triton does NOT auto-promote fp16/bf16 to fp32 for accumulation | + +--- + +## Broadcasting Rules + +Triton broadcasting follows NumPy conventions with one key constraint: +tensors are at most 2D in practice (block pointers may extend this). + +### Rules + +1. **Left-pad with ones**: If tensors have different numbers of dimensions, the + shorter shape is padded on the left with 1s. +2. **Dimension-1 expansion**: Dimensions of size 1 are stretched to match the + corresponding dimension of the other tensor. +3. **Incompatible = error**: If dimensions differ and neither is 1, it is a compile error. + +### Example: Row-Column Broadcast + +```python +# Create a row vector (1, N) and column vector (M, 1) +row = tl.arange(0, N)[None, :] # shape: (1, N) +col = tl.arange(0, M)[:, None] # shape: (M, 1) + +# Broadcast produces (M, N) — outer product pattern +result = row + col # shape: (M, N) +``` + +### Common Broadcasting Patterns + +| Pattern | Shape A | Shape B | Result Shape | Use Case | +|---------|---------|---------|-------------|----------| +| Row + Col | `(1, N)` | `(M, 1)` | `(M, N)` | 2D index grids, outer products | +| Scalar + Block | `()` | `(M, N)` | `(M, N)` | Add bias, scale | +| Row mask | `(1, N)` | `(M, N)` | `(M, N)` | Column-wise masking | + +### Gotcha: Broadcasting + +| Gotcha | Detail | +|--------|--------| +| No implicit unsqueeze | You must explicitly reshape with `[:, None]` or `[None, :]` | +| 1D + 1D does NOT broadcast | Two 1D tensors of different length are an error, not broadcast | +| Mask must broadcast to data | `tl.load(ptr, mask=mask)` — mask shape must broadcast to ptr block shape | + +--- + +## Integer Division and Modulus — C Semantics + +**CRITICAL**: Triton uses **C semantics** (round toward zero), NOT Python semantics +(round toward negative infinity). This is the most common source of subtle bugs +when porting Python logic to Triton kernels. + +### Comparison Table + +| Expression | Python Result | Triton Result | Why | +|------------|--------------|---------------|-----| +| `-7 // 2` | `-4` | `-3` | Python: floor division. Triton/C: truncation toward zero | +| `-7 % 2` | `1` | `-1` | Follows from division: `a == (a // b) * b + (a % b)` | +| `7 // -2` | `-4` | `-3` | Same: truncation vs floor | +| `7 % -2` | `-1` | `1` | Remainder keeps dividend sign in C | +| `-7 // -2` | `3` | `3` | Both agree when signs match (positive quotient) | + +### The Identity + +Both C and Python satisfy: `a == (a // b) * b + (a % b)` + +But they disagree on which direction to round the quotient, which changes the remainder. + +### Exception: Scalar-Only Computations + +When **all inputs are Python scalars** (not Triton tensors), division and modulus +follow **Python semantics**. This only applies to compile-time constant folding. + +```python +@triton.jit +def kernel(X_ptr, N, BLOCK: tl.constexpr): + # Python semantics — both are Python scalars at compile time + blocks_per_row = (-7) // 2 # = -4 (Python floor division) + + # C semantics — pid is a Triton value + pid = tl.program_id(0) + row = pid // N # truncation toward zero + col = pid % N # C remainder +``` + +### Gotcha: Safe Patterns for Negative Values + +| Unsafe Pattern | Problem | Safe Alternative | +|----------------|---------|------------------| +| `(-offset) // stride` | C truncation gives wrong block | `-(offset // stride)` or use unsigned | +| `idx % BLOCK` for negative idx | Negative remainder | Ensure idx is non-negative, or add `+ BLOCK) % BLOCK` | +| Porting Python `divmod` logic | Both `//` and `%` differ | Rewrite with explicit floor: `q = (a - (a % b + b) % b) // b` | + +### When It Matters + +This only causes bugs when **operands can be negative**. If all values are +non-negative (which is common for pointer offsets and indices), C and Python +semantics agree. Guard against negative values explicitly when in doubt. diff --git a/.agents/skills/kernel-triton-writing/references/operator-routing.md b/.agents/skills/kernel-triton-writing/references/operator-routing.md new file mode 100644 index 000000000000..9bfd3d3b7c13 --- /dev/null +++ b/.agents/skills/kernel-triton-writing/references/operator-routing.md @@ -0,0 +1,125 @@ + + + + +# Operator Routing Decision Reference + +Detailed decision rules for determining whether an operator should be implemented +as a custom Triton kernel or handled by existing libraries. + +## Decision Procedure + +Follow these rules in order. Stop at the first match. + +1. **Single element-wise op** (e.g., `relu(x)`, `sigmoid(x)`) -- SKIP. PyTorch + already optimal, no fusion benefit. +2. **Standalone matmul** (e.g., `torch.matmul(a, b)`) -- SKIP. cuBLAS is highly + optimized and hard to beat. +3. **Standard attention** (e.g., `F.scaled_dot_product_attention`) -- SKIP. Use + FlashAttention. +4. **Element-wise chain (2+ ops)** (e.g., `gelu(dropout(x))`, `silu(x) * y`) -- + USE TRITON. Fuse memory-bound ops into compute-bound kernel. +5. **Reduction op** (e.g., LayerNorm, RMSNorm, Softmax) -- USE TRITON. Custom + single-pass implementation beats generic PyTorch decomposition. +6. **Matmul + element-wise epilogue** (e.g., `matmul(a, b) + bias`, + `matmul + gelu`) -- USE TRITON. Epilogue fusion avoids memory round-trip. +7. **Matmul + reduction** (e.g., `matmul -> softmax`, `matmul -> layernorm`) -- + USE TRITON. Common transformer pattern with clear fusion benefit. +8. **Custom attention variant** -- Check FlashAttention support first. Only use + Triton if the variant is unsupported. +9. **Sparse operations** -- Triton can help, but evaluate specialized libraries + (cuSPARSE, Triton block-sparse) first. +10. **Very small tensors** -- Launch overhead may dominate. Benchmark before + committing. +11. **Default** -- Analyze operator code and shapes, then decide. + +## Output Format + +Report the routing decision as: + +```markdown +## Routing Decision: [OPERATOR_NAME] + +**Decision:** USE TRITON | SKIP TRITON | EVALUATE FURTHER + +**Pattern:** [e.g., Element-wise chain, Reduction, Matmul+epilogue] + +**Rationale:** [Why -- reference fusion benefit or lack thereof] + +**Next Steps:** +- [USE TRITON] Proceed to Phase 1 (Analyze the Operator) +- [SKIP] Recommend alternative (cuBLAS, FlashAttention, PyTorch) +- [EVALUATE] Profile operator, analyze shapes, then re-decide +``` + +## Examples + +### Fused GELU + Dropout + +```python +def fused_op(x, p=0.1): + return F.dropout(F.gelu(x), p=p) +``` + +**Decision:** USE TRITON | **Pattern:** Element-wise chain (2 ops) +Fusing eliminates one intermediate tensor write+read (~2x memory traffic reduction). + +### Simple ReLU + +```python +def simple_relu(x): + return F.relu(x) +``` + +**Decision:** SKIP TRITON | **Pattern:** Single element-wise op +No fusion benefit. PyTorch ReLU is already a single memory-bound kernel. + +### RMSNorm + +```python +def rmsnorm(x, weight, eps=1e-6): + rms = torch.sqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + eps) + return x / rms * weight +``` + +**Decision:** USE TRITON | **Pattern:** Reduction op +Triton fuses square, mean, sqrt, divide, multiply in a single pass over the data. + +### Linear + GELU + +```python +def linear_gelu(x, weight, bias): + return F.gelu(F.linear(x, weight, bias)) +``` + +**Decision:** USE TRITON | **Pattern:** Matmul + element-wise epilogue +Fusing GELU into the matmul epilogue avoids an extra full tensor read+write. + +## Edge Cases + +- **Dynamic shapes or data-dependent branching** -- Triton requires static grid + dimensions at launch. If shapes change per-sample, fall back to PyTorch eager + or `torch.compile`. +- **Operators already in `torch.compile` fusion groups** -- Check whether + `torch.compile` already fuses the pattern before writing a manual kernel. + A manual Triton kernel is only justified if it measurably outperforms the + compiler-generated version. +- **Mixed precision boundaries** -- Triton handles dtype casting well, but verify + that the fused kernel preserves numerical behavior (especially around + loss scaling and FP16/BF16 reductions). diff --git a/.agents/skills/kernel-triton-writing/references/patterns-advanced.md b/.agents/skills/kernel-triton-writing/references/patterns-advanced.md new file mode 100644 index 000000000000..635f10897c3c --- /dev/null +++ b/.agents/skills/kernel-triton-writing/references/patterns-advanced.md @@ -0,0 +1,320 @@ + + + + +# Advanced Triton Patterns + +Source tutorials: + +- [05-layer-norm](https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html) +- [06-fused-attention](https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html) +- [07-extern-functions](https://triton-lang.org/main/getting-started/tutorials/07-extern-functions.html) + +## Layer Normalization + +Layer norm normalizes across the hidden dimension: `y = (x - mean) / sqrt(var + eps) * w + b`. +Each program instance processes one row of the input (one token). The hidden dimension +is tiled into blocks so arbitrary sizes are supported. + +### Forward Kernel (Complete) + +```python +@triton.jit +def _layer_norm_fwd_fused( + X, # input pointer, shape (M, N) + Y, # output pointer, shape (M, N) + W, # weight pointer, shape (N,) + B, # bias pointer, shape (N,) + Mean, # mean pointer, shape (M,) — written for backward + Rstd, # rstd pointer, shape (M,) — written for backward + stride, # row stride of X and Y + N, # number of columns (hidden size) + eps: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + # Each program handles one row + row = tl.program_id(0) + Y += row * stride + X += row * stride + + # --- Compute mean --- + _mean = tl.zeros([BLOCK_SIZE], dtype=tl.float32) + for off in range(0, N, BLOCK_SIZE): + cols = off + tl.arange(0, BLOCK_SIZE) + a = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) + _mean += a + mean = tl.sum(_mean, axis=0) / N + + # --- Compute variance --- + _var = tl.zeros([BLOCK_SIZE], dtype=tl.float32) + for off in range(0, N, BLOCK_SIZE): + cols = off + tl.arange(0, BLOCK_SIZE) + x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) + x = tl.where(cols < N, x - mean, 0.0) + _var += x * x + var = tl.sum(_var, axis=0) / N + rstd = 1 / tl.sqrt(var + eps) + + # Store mean and rstd for backward + tl.store(Mean + row, mean) + tl.store(Rstd + row, rstd) + + # --- Normalize and apply affine transform --- + for off in range(0, N, BLOCK_SIZE): + cols = off + tl.arange(0, BLOCK_SIZE) + mask = cols < N + w = tl.load(W + cols, mask=mask) + b = tl.load(B + cols, mask=mask) + x = tl.load(X + cols, mask=mask, other=0.0).to(tl.float32) + x_hat = (x - mean) * rstd + y = x_hat * w + b + tl.store(Y + cols, y, mask=mask) +``` + +**Key pattern:** Two-pass reduction (mean then variance) using block-wise tiling. +Each loop iteration processes `BLOCK_SIZE` elements with masking for the tail. +Mean and rstd are saved for the backward pass. + +### Backward Kernel: Atomic Lock Pattern + +The backward pass computes `dw` and `db` which require reducing across all rows +(all programs contribute). Triton uses a spin-lock pattern with `atomic_cas` for +mutual exclusion: + +```python +@triton.jit +def _layer_norm_bwd_dwdb( + DW, DB, # output accumulators, shape (N,) + DWEIGHT, DBIAS, # partial sums buffer, shape (GROUP_SIZE_M, N) + Lock, # lock array, shape (1,) — int32 + ... + GROUP_SIZE_M: tl.constexpr, +): + row_block_id = tl.program_id(0) + # Each group of rows accumulates partials then atomically adds to DW/DB + + # --- Compute partial dw, db for this row group --- + # (loop over assigned rows, accumulate _dw and _db) + + # --- Acquire lock --- + lock_id = tl.program_id(1) # column block index + Lock += lock_id + Count = Lock + tl.num_programs(1) # second half stores count + + while tl.atomic_cas(Lock, 0, 1) == 1: # spin until we get 0->1 + pass + count = tl.load(Count) # how many groups have accumulated so far + + if count == 0: + # First group: just store + tl.store(DWEIGHT + cols, _dw, mask=mask) + tl.store(DBIAS + cols, _db, mask=mask) + else: + # Subsequent groups: accumulate + _dw += tl.load(DWEIGHT + cols, mask=mask) + _db += tl.load(DBIAS + cols, mask=mask) + tl.store(DWEIGHT + cols, _dw, mask=mask) + tl.store(DBIAS + cols, _db, mask=mask) + + if count == GROUP_SIZE_M - 1: + # Last group: write final result + tl.store(DW + cols, _dw, mask=mask) + tl.store(DB + cols, _db, mask=mask) + + # --- Release lock and increment count --- + tl.atomic_xchg(Lock, 0) # release: set lock back to 0 + tl.store(Count, count + 1) # must store AFTER release for correctness + tl.debug_barrier() # ensure memory operations are visible +``` + +**Gotchas:** + +- `atomic_cas(Lock, 0, 1)` returns the old value; spin while it returns 1 (already held). +- `atomic_xchg(Lock, 0)` unconditionally sets to 0 (release). Do NOT use `atomic_cas` for release. +- The count update (`tl.store(Count, count + 1)`) must happen after the lock is released. +- `tl.debug_barrier()` forces memory ordering visibility across programs. +- Lock array must be zero-initialized before each backward call. +- This pattern is needed because Triton has no native cross-program reduction for non-atomic dtypes. + +## Fused Attention + +Implements Flash Attention v2: fused Q*K^T softmax and V accumulation in a single +kernel, avoiding materializing the full N x N attention matrix. + +### Online Softmax Algorithm + +The key insight is computing softmax in a single streaming pass using running +statistics. For each block of K/V columns processed: + +``` +# For each new block j of keys: +qk = Q_block @ K_block_j^T # [BLOCK_M, BLOCK_N] +m_ij = max(qk, axis=1) # new block max +m_i_new = max(m_i, m_ij) # update running max +alpha = exp(m_i - m_i_new) # correction factor for old accumulators +p = exp(qk - m_i_new[:, None]) # stable softmax numerator +l_i = alpha * l_i + sum(p, axis=1) # update running denominator +acc = alpha[:, None] * acc + p @ V_j # rescale old acc + new contribution +m_i = m_i_new # commit new max +# After all blocks: +acc = acc / l_i[:, None] # final normalization +``` + +**Why this works:** Each time a new block raises the running max, all previous +accumulations are rescaled by `exp(old_max - new_max)`, maintaining numerical +equivalence to the two-pass softmax. + +### Multi-Stage Processing (Causal Masking) + +The STAGE parameter controls masking behavior in the inner loop: + +| STAGE | Behavior | When used | +|-------|----------|-----------| +| 1 | Off-band: skip blocks entirely below diagonal | Causal, early blocks | +| 3 | On-band: apply causal mask within block | Causal, diagonal blocks | +| 2 | No masking | Non-causal attention | + +```python +# Causal masking within a block (STAGE == 3): +if STAGE == 3: + # Current query rows: [start_m, start_m + BLOCK_M) + # Current key cols: [start_n, start_n + BLOCK_N) + offs_m = start_m + tl.arange(0, BLOCK_M) + offs_n = start_n + tl.arange(0, BLOCK_N) + causal_mask = offs_m[:, None] >= offs_n[None, :] + qk = tl.where(causal_mask, qk, float("-inf")) +``` + +**Gotcha:** When `STAGE == 1`, blocks where all keys are below the diagonal are +skipped entirely (the inner loop `start_n` begins past the diagonal). This is a +major performance win for causal attention on long sequences. + +### Kernel Structure Skeleton + +```python +@triton.jit +def _attn_fwd( + Q, K, V, sm_scale, + M, # log-sum-exp for backward, shape (batch, nheads, seqlen) + Out, + stride_qz, stride_qh, stride_qm, stride_qk, # Q strides + # ... K, V, Out strides ... + Z, H, N_CTX, + BLOCK_M: tl.constexpr, # query block size (e.g. 128) + BLOCK_N: tl.constexpr, # key block size (e.g. 64) + HEAD_DIM: tl.constexpr, # head dimension (e.g. 64) + STAGE: tl.constexpr, +): + start_m = tl.program_id(0) # which query block + off_hz = tl.program_id(1) # batch * head index + + # Initialize pointers for Q[start_m], K, V + # Load Q block into registers (stays resident) + q = tl.load(Q_block_ptr) # [BLOCK_M, HEAD_DIM] + + # Accumulator in float32 + acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32) + m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") + l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0 + + # --- Inner loop over K/V blocks --- + for start_n in range(lo, hi, BLOCK_N): + k = tl.load(K_block_ptr) # [BLOCK_N, HEAD_DIM] + v = tl.load(V_block_ptr) # [BLOCK_N, HEAD_DIM] + + qk = tl.dot(q, tl.trans(k)) * sm_scale # [BLOCK_M, BLOCK_N] + + # Apply causal mask if STAGE == 3 + # Online softmax update (m_i, l_i, acc) as shown above + + K_block_ptr = tl.advance(K_block_ptr, (BLOCK_N, 0)) + V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0)) + + # Final normalization + acc = acc / l_i[:, None] + # Store lse = m_i + log(l_i) for backward + tl.store(Out_block_ptr, acc.to(Out.type.element_ty)) +``` + +### TensorDescriptor (Hopper+) + +On Hopper/Blackwell, `TensorDescriptor` enables TMA (Tensor Memory Accelerator): + +```python +desc_q = TensorDescriptor(Q_ptr, shape=[N_CTX, HEAD_DIM], + strides=[stride_qm, stride_qk], + block_shape=[BLOCK_M, HEAD_DIM]) +q = desc_q.load([start_m * BLOCK_M, 0]) # replaces pointer arithmetic +``` + +### Warp Specialization and Performance + +On Blackwell (sm_100+), warp specialization lets different warp groups play +producer/consumer roles, overlapping loads with compute via `tl.async_task`. +Flash Attention in Triton reaches ~165 TFLOPS (fp16, H100). Causal masking +with the STAGE optimization roughly halves unnecessary work. + +## External Functions + +Triton can call functions from external device libraries (libdevice for CUDA, +ROCm device libs for HIP) for math operations not built into the language. + +### Basic Usage + +```python +@triton.jit +def asin_kernel( + x_ptr, y_ptr, + n_elements, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(axis=0) + offset = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offset < n_elements + x = tl.load(x_ptr + offset, mask=mask) + # Call libdevice asin — dispatches based on dtype (fp32 or fp64) + y = tl.extra.cuda.libdevice.asin(x) + tl.store(y_ptr + offset, y, mask=mask) +``` + +Type dispatch is automatic: `libdevice.asin` calls `__nv_asinf` for float32 +and `__nv_asin` for float64 under the hood. + +### Custom Library Paths + +Pass external libraries explicitly via `extern_libs` at compile time: + +```python +grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),) +asin_kernel[grid](x, y, n_elements, BLOCK_SIZE=1024, + extern_libs={"libdevice": "/path/to/libdevice.10.bc"}) +``` + +### Backend Detection + +| Backend | Library file | Namespace | +|---------|-------------|-----------| +| CUDA | `libdevice.10.bc` | `tl.extra.cuda.libdevice.*` | +| HIP | `ocml.bc` / `ockl.bc` | Functions mapped through HIP backend | + +Common functions: `asin`, `acos`, `atan`, `exp`, `log`, `pow`, `sqrt`, `rsqrt`, +`fma`, `cbrt`, `erf`, `erfc`, `ceil`, `floor`, `round`. + +**Gotcha:** `extern_libs` must point to the `.bc` bitcode file (typically +`/usr/local/cuda/nvvm/libdevice/libdevice.10.bc`). Missing file = compile-time linker error. diff --git a/.agents/skills/kernel-triton-writing/references/patterns-basic.md b/.agents/skills/kernel-triton-writing/references/patterns-basic.md new file mode 100644 index 000000000000..9435f2f90aa2 --- /dev/null +++ b/.agents/skills/kernel-triton-writing/references/patterns-basic.md @@ -0,0 +1,238 @@ + + + + +# Triton Basic Kernel Patterns + +Reusable patterns extracted from the official Triton tutorials. +Each section contains a complete kernel, its launch wrapper, and annotations. + +--- + +## Vector Addition + +The simplest Triton pattern: 1D parallel map over contiguous data. + +### Kernel + +```python +import torch +import triton +import triton.language as tl + +@triton.jit +def add_kernel( + x_ptr, y_ptr, output_ptr, + n_elements, + BLOCK_SIZE: tl.constexpr, # compile-time constant: controls tile width +): + # Each program instance owns one tile of BLOCK_SIZE elements. + pid = tl.program_id(axis=0) + # Compute the start offset for this program's tile, then the per-lane offsets. + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + # Guard against out-of-bounds access on the final, possibly partial tile. + mask = offsets < n_elements + # Load inputs from DRAM — masked lanes get a safe default (0.0). + x = tl.load(x_ptr + offsets, mask=mask) + y = tl.load(y_ptr + offsets, mask=mask) + output = x + y + # Write result back — only masked lanes write. + tl.store(output_ptr + offsets, output, mask=mask) +``` + +### Launch Wrapper + +```python +def add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + output = torch.empty_like(x) + assert x.is_cuda and y.is_cuda and output.is_cuda + n_elements = output.numel() + # Grid is a callable: Triton passes meta-parameters (incl. BLOCK_SIZE) at launch. + grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) + add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024) + return output +``` + +### Benchmark Pattern + +```python +@triton.testing.perf_report( + triton.testing.Benchmark( + x_names=["size"], + x_vals=[2**i for i in range(12, 28, 1)], + x_log=True, + line_arg="provider", + line_vals=["triton", "torch"], + line_names=["Triton", "Torch"], + ylabel="GB/s", + plot_name="vector-add-performance", + args={}, + ) +) +def benchmark(size, provider): + x = torch.rand(size, device="cuda", dtype=torch.float32) + y = torch.rand(size, device="cuda", dtype=torch.float32) + quantiles = [0.5, 0.2, 0.8] + if provider == "torch": + ms, min_ms, max_ms = triton.testing.do_bench(lambda: x + y, quantiles=quantiles) + if provider == "triton": + ms, min_ms, max_ms = triton.testing.do_bench(lambda: add(x, y), quantiles=quantiles) + gbps = lambda ms: 3 * x.numel() * x.element_size() * 1e-9 / (ms * 1e-3) + return gbps(ms), gbps(max_ms), gbps(min_ms) + +benchmark.run(print_data=True, show_plots=True) +``` + +### Key Takeaways + +- **Offset pattern:** `pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)` — the universal 1D tiling idiom. +- **Masking:** Always guard the last tile with `offsets < n_elements`. +- **Grid as callable:** `lambda meta: (triton.cdiv(n, meta["BLOCK_SIZE"]),)` lets autotune vary BLOCK_SIZE. +- **Pointer arithmetic:** Triton pointers support `ptr + offset_tensor` for vectorized addressing. + +--- + +## Fused Softmax + +Row-wise softmax fused into a single kernel. Fusion reduces DRAM traffic from +`5*M*N + 2*M` bytes (naive: read 3x for max/exp/sum, write 2x) down to `M*N` +read + `M*N` write by keeping intermediate results in SRAM. + +### Kernel + +```python +@triton.jit +def softmax_kernel( + output_ptr, input_ptr, + input_row_stride, output_row_stride, + n_rows, n_cols, + BLOCK_SIZE: tl.constexpr, # must be >= n_cols (padded to power-of-2) + num_stages: tl.constexpr, # software pipelining depth +): + # Persistent-kernel style: each program processes multiple rows, strided. + row_start = tl.program_id(0) + row_step = tl.num_programs(0) + for row_idx in tl.range(row_start, n_rows, row_step, num_stages=num_stages): + # Compute pointers for this row. + row_start_ptr = input_ptr + row_idx * input_row_stride + col_offsets = tl.arange(0, BLOCK_SIZE) + input_ptrs = row_start_ptr + col_offsets + # Mask: BLOCK_SIZE is rounded up to power-of-2, so some lanes are OOB. + mask = col_offsets < n_cols + # Load row; OOB lanes get -inf so they don't affect max/sum. + row = tl.load(input_ptrs, mask=mask, other=-float("inf")) + # --- Numerical stability: subtract row-max before exp --- + row_minus_max = row - tl.max(row, axis=0) + numerator = tl.exp(row_minus_max) + denominator = tl.sum(numerator, axis=0) + softmax_output = numerator / denominator + # Store result. + output_row_start_ptr = output_ptr + row_idx * output_row_stride + output_ptrs = output_row_start_ptr + col_offsets + tl.store(output_ptrs, softmax_output, mask=mask) +``` + +### Launch Wrapper + +```python +def softmax(x: torch.Tensor) -> torch.Tensor: + n_rows, n_cols = x.shape + # BLOCK_SIZE must cover the full row — round up to power-of-2. + BLOCK_SIZE = triton.next_power_of_2(n_cols) + # Heuristic: use more warps for wider rows. + num_warps = 4 if BLOCK_SIZE <= 2048 else 8 + # Persistent kernel: launch fewer programs than rows for large inputs. + # Each SM can run ~4 programs concurrently (occupancy dependent). + num_stages = 4 if BLOCK_SIZE > 2048 else 2 + y = torch.empty_like(x) + # Grid: one dimension, capped by number of rows. + num_programs = min(n_rows, 1024) # cap to avoid over-subscription + softmax_kernel[(num_programs, 1, 1)]( + y, x, + x.stride(0), y.stride(0), + n_rows, n_cols, + BLOCK_SIZE=BLOCK_SIZE, + num_stages=num_stages, + num_warps=num_warps, + ) + return y +``` + +### Key Takeaways + +- **Numerical stability:** Always `row - tl.max(row, axis=0)` before `tl.exp`. +- **Power-of-2 padding:** `BLOCK_SIZE = triton.next_power_of_2(n_cols)` with `-inf` masking for OOB lanes. +- **Persistent kernel:** `tl.range(start, end, step, num_stages=...)` loops over multiple rows per + program, improving occupancy and enabling software pipelining. +- **Fusion benefit:** One kernel replaces three separate passes (max, exp/sum, divide), keeping + all intermediates in registers/SRAM instead of round-tripping through DRAM. + +--- + +## Low-Memory Dropout + +Traditional dropout stores a full-size bit mask. This pattern stores only an `int32` seed and +recomputes the mask on-the-fly via Triton's built-in PRNG. The same seed + offsets produce +identical random values, so forward and backward passes see the same mask without storing it. + +### Kernel + +```python +@triton.jit +def _seeded_dropout( + x_ptr, output_ptr, + n_elements, + p, # dropout probability (float, 0 to 1) + seed, # int32 seed — the ONLY state needed to reproduce the mask + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(axis=0) + block_start = pid * BLOCK_SIZE + offsets = block_start + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + x = tl.load(x_ptr + offsets, mask=mask) + # tl.rand: deterministic PRNG — given the same (seed, offsets), produces + # the same uniform float32 values in [0, 1). No global state needed. + random = tl.rand(seed, offsets) + x_keep = random > p + # Scale kept elements by 1/(1-p) so expected value is unchanged (inverted dropout). + # Dropped elements become 0.0. + output = tl.where(x_keep, x / (1 - p), 0.0) + tl.store(output_ptr + offsets, output, mask=mask) +``` + +### Launch Wrapper + +```python +def seeded_dropout(x: torch.Tensor, p: float, seed: int) -> torch.Tensor: + output = torch.empty_like(x) + assert x.is_contiguous() + n_elements = x.numel() + grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) + _seeded_dropout[grid](x, output, n_elements, p, seed, BLOCK_SIZE=1024) + return output +``` + +### Key Takeaways + +- **Memory savings:** State = 1 `int32` seed, not an `(N,)` bool tensor. +- **Deterministic PRNG:** `tl.rand(seed, offsets)` is pure-functional — same inputs, same outputs. +- **Inverted dropout:** `x / (1 - p)` scales at train time so inference needs no adjustment. +- **`tl.where` pattern:** `tl.where(cond, val_true, val_false)` is the standard Triton conditional — + works element-wise on block tensors, compiles to predicated instructions (no branch divergence). diff --git a/.agents/skills/kernel-triton-writing/references/patterns-fusion.md b/.agents/skills/kernel-triton-writing/references/patterns-fusion.md new file mode 100644 index 000000000000..ac8d6db0912d --- /dev/null +++ b/.agents/skills/kernel-triton-writing/references/patterns-fusion.md @@ -0,0 +1,349 @@ + + + + +# Deep Learning Fusion Patterns + +Ready-to-use Triton kernel patterns for common DL operator fusions. +Each pattern includes autotune configs, kernel, wrapper, and expected speedups. + +For foundational patterns (vector add, softmax, dropout), see `patterns-basic.md`. +For LayerNorm with backward, fused attention, and extern functions, see `patterns-advanced.md`. + +--- + +## GELU + Dropout + +**Use when:** Transformer FFN layers with dropout. +**Expected speedup:** 1.8-2.2x vs separate ops. + +```python +@triton.autotune( + configs=[ + triton.Config({'BLOCK_SIZE': 1024}, num_warps=4), + triton.Config({'BLOCK_SIZE': 2048}, num_warps=8), + ], + key=['n_elements'], +) +@triton.jit +def fused_gelu_dropout_kernel( + x_ptr, out_ptr, n_elements, p, seed, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + + x = tl.load(x_ptr + offsets, mask=mask) + + # GELU (exact): cast to fp32 for erf, then cast back + x_fp32 = x.to(tl.float32) + x_gelu = 0.5 * x_fp32 * (1.0 + tl.math.erf(x_fp32 * 0.7071067811865476)) + x = x_gelu.to(x.dtype) + + # Dropout + random = tl.rand(seed, offsets) + x = tl.where(random > p, x / (1 - p), 0.0) + + tl.store(out_ptr + offsets, x, mask=mask) + + +def fused_gelu_dropout(x: torch.Tensor, p: float = 0.1, training: bool = True) -> torch.Tensor: + if not training or p == 0.0: + return torch.nn.functional.gelu(x) + n_elements = x.numel() + out = torch.empty_like(x) + grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),) + seed = (x.data_ptr() % (2**31)) ^ n_elements + fused_gelu_dropout_kernel[grid](x, out, n_elements, p, seed) + return out +``` + +--- + +## SiLU + Multiply (SwiGLU) + +**Use when:** LLaMA-style FFN with SwiGLU activation. +**Expected speedup:** 1.5-2x vs `F.silu(gate) * x`. + +```python +@triton.autotune( + configs=[ + triton.Config({'BLOCK_SIZE': 1024}, num_warps=4), + triton.Config({'BLOCK_SIZE': 2048}, num_warps=8), + ], + key=['n_elements'], +) +@triton.jit +def fused_silu_mul_kernel( + x_ptr, gate_ptr, out_ptr, n_elements, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + + x = tl.load(x_ptr + offsets, mask=mask) + gate = tl.load(gate_ptr + offsets, mask=mask) + + # SiLU(gate) * x = gate * sigmoid(gate) * x + silu_gate = gate * tl.sigmoid(gate) + out = silu_gate * x + + tl.store(out_ptr + offsets, out, mask=mask) + + +def fused_silu_mul(x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: + assert x.shape == gate.shape + n_elements = x.numel() + out = torch.empty_like(x) + grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),) + fused_silu_mul_kernel[grid](x, gate, out, n_elements) + return out +``` + +--- + +## Residual Add + Activation + +**Use when:** Adding residual connection with activation. +**Expected speedup:** 1.4-1.8x vs `F.gelu(x + residual)`. + +```python +@triton.autotune( + configs=[ + triton.Config({'BLOCK_SIZE': 1024}, num_warps=4), + triton.Config({'BLOCK_SIZE': 2048}, num_warps=8), + ], + key=['n_elements'], +) +@triton.jit +def fused_residual_gelu_kernel( + x_ptr, residual_ptr, out_ptr, n_elements, + BLOCK_SIZE: tl.constexpr, +): + pid = tl.program_id(0) + offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + mask = offsets < n_elements + + x = tl.load(x_ptr + offsets, mask=mask) + residual = tl.load(residual_ptr + offsets, mask=mask) + x = x + residual + + # GELU (exact) + x_fp32 = x.to(tl.float32) + x = (0.5 * x_fp32 * (1.0 + tl.math.erf(x_fp32 * 0.7071067811865476))).to(x.dtype) + + tl.store(out_ptr + offsets, x, mask=mask) + + +def fused_residual_gelu(x: torch.Tensor, residual: torch.Tensor) -> torch.Tensor: + n_elements = x.numel() + out = torch.empty_like(x) + grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),) + fused_residual_gelu_kernel[grid](x, residual, out, n_elements) + return out +``` + +--- + +## RMSNorm + +**Use when:** LLaMA-style normalization (no mean subtraction). +**Expected speedup:** 1.4-2x vs naive PyTorch RMSNorm. + +```python +@triton.autotune( + configs=[ + triton.Config({'BLOCK_SIZE': 1024}, num_warps=8), + triton.Config({'BLOCK_SIZE': 2048}, num_warps=8), + triton.Config({'BLOCK_SIZE': 4096}, num_warps=16), + ], + key=['n_cols'], +) +@triton.jit +def rmsnorm_kernel( + x_ptr, out_ptr, weight_ptr, + n_rows, n_cols, eps, + BLOCK_SIZE: tl.constexpr, +): + row_idx = tl.program_id(0) + col_offsets = tl.arange(0, BLOCK_SIZE) + mask = col_offsets < n_cols + + row_start = row_idx * n_cols + x = tl.load(x_ptr + row_start + col_offsets, mask=mask, other=0.0) + + # Compute RMS + x_sq = x * x + rms = tl.sqrt(tl.sum(x_sq, axis=0) / n_cols + eps) + + # Normalize and scale + x_norm = x / rms + weight = tl.load(weight_ptr + col_offsets, mask=mask, other=1.0) + out = x_norm * weight + + tl.store(out_ptr + row_start + col_offsets, out, mask=mask) + + +def triton_rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: + assert x.is_contiguous() + shape = x.shape + x = x.view(-1, shape[-1]) + n_rows, n_cols = x.shape + out = torch.empty_like(x) + grid = (n_rows,) + rmsnorm_kernel[grid](x, out, weight, n_rows, n_cols, eps) + return out.view(shape) +``` + +--- + +## Linear + GELU (Matmul + Epilogue) + +**Use when:** Transformer FFN first linear with activation. +**Expected speedup:** 1.3-1.6x vs `F.gelu(F.linear(x, weight, bias))`. + +```python +@triton.autotune( + configs=[ + triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64, 'BLOCK_K': 32}, num_stages=3, num_warps=4), + triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'BLOCK_K': 32}, num_stages=3, num_warps=4), + triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 32}, num_stages=3, num_warps=8), + ], + key=['M', 'N', 'K'], +) +@triton.jit +def linear_gelu_kernel( + x_ptr, weight_ptr, bias_ptr, out_ptr, + M, N, K, + stride_xm, stride_xk, stride_wk, stride_wn, + BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, +): + pid_m = tl.program_id(0) + pid_n = tl.program_id(1) + + offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) + offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) + offs_k = tl.arange(0, BLOCK_K) + + acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) + for k in range(0, K, BLOCK_K): + k_offs = k + offs_k + x_ptrs = x_ptr + offs_m[:, None] * stride_xm + k_offs[None, :] * stride_xk + x_mask = (offs_m[:, None] < M) & (k_offs[None, :] < K) + x = tl.load(x_ptrs, mask=x_mask, other=0.0) + + w_ptrs = weight_ptr + k_offs[:, None] * stride_wk + offs_n[None, :] * stride_wn + w_mask = (k_offs[:, None] < K) & (offs_n[None, :] < N) + w = tl.load(w_ptrs, mask=w_mask, other=0.0) + acc += tl.dot(x, w) + + # Add bias + fused GELU epilogue + bias = tl.load(bias_ptr + offs_n, mask=offs_n < N, other=0.0) + acc = acc + bias[None, :] + acc = 0.5 * acc * (1.0 + tl.math.erf(acc * 0.7071067811865476)) + + out_ptrs = out_ptr + offs_m[:, None] * N + offs_n[None, :] + out_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N) + tl.store(out_ptrs, acc.to(tl.float16), mask=out_mask) + + +def linear_gelu(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor) -> torch.Tensor: + assert x.is_contiguous() and weight.is_contiguous() + M, K = x.shape + K2, N = weight.shape + assert K == K2 + out = torch.empty((M, N), device=x.device, dtype=x.dtype) + grid = lambda meta: (triton.cdiv(M, meta['BLOCK_M']), triton.cdiv(N, meta['BLOCK_N'])) + linear_gelu_kernel[grid](x, weight, bias, out, M, N, K, x.stride(0), x.stride(1), weight.stride(0), weight.stride(1)) + return out +``` + +--- + +## Fused Add + LayerNorm + +**Use when:** Post-attention residual add + normalization. +**Expected speedup:** 1.5-2x vs `F.layer_norm(x + residual, ...)`. + +```python +@triton.autotune( + configs=[ + triton.Config({'BLOCK_SIZE': 1024}, num_warps=8), + triton.Config({'BLOCK_SIZE': 2048}, num_warps=8), + triton.Config({'BLOCK_SIZE': 4096}, num_warps=16), + ], + key=['n_cols'], +) +@triton.jit +def fused_add_layernorm_kernel( + x_ptr, residual_ptr, out_ptr, weight_ptr, bias_ptr, + n_rows, n_cols, eps, + BLOCK_SIZE: tl.constexpr, +): + row_idx = tl.program_id(0) + col_offsets = tl.arange(0, BLOCK_SIZE) + mask = col_offsets < n_cols + row_start = row_idx * n_cols + + # Load and add + x = tl.load(x_ptr + row_start + col_offsets, mask=mask, other=0.0) + residual = tl.load(residual_ptr + row_start + col_offsets, mask=mask, other=0.0) + x = x + residual + + # LayerNorm + mean = tl.sum(x, axis=0) / n_cols + x_centered = x - mean + var = tl.sum(x_centered * x_centered, axis=0) / n_cols + x_norm = x_centered / tl.sqrt(var + eps) + + weight = tl.load(weight_ptr + col_offsets, mask=mask, other=1.0) + bias = tl.load(bias_ptr + col_offsets, mask=mask, other=0.0) + out = x_norm * weight + bias + + tl.store(out_ptr + row_start + col_offsets, out, mask=mask) + + +def fused_add_layernorm( + x: torch.Tensor, residual: torch.Tensor, + weight: torch.Tensor, bias: torch.Tensor, eps: float = 1e-5, +) -> torch.Tensor: + assert x.is_contiguous() and residual.is_contiguous() + shape = x.shape + x = x.view(-1, shape[-1]) + residual = residual.view(-1, shape[-1]) + n_rows, n_cols = x.shape + out = torch.empty_like(x) + grid = (n_rows,) + fused_add_layernorm_kernel[grid](x, residual, out, weight, bias, n_rows, n_cols, eps) + return out.view(shape) +``` + +--- + +## Pattern Selection Guide + +| Use Case | Pattern | Expected Speedup | +|----------|---------|------------------| +| FFN activation + dropout | GELU + Dropout | 1.8-2.2x | +| LLaMA FFN gate | SiLU + Multiply | 1.5-2x | +| LLaMA norm | RMSNorm | 1.4-2x | +| FFN with activation | Linear + GELU | 1.3-1.6x | +| Post-attention | Add + LayerNorm | 1.5-2x | diff --git a/.agents/skills/kernel-triton-writing/references/patterns-gemm.md b/.agents/skills/kernel-triton-writing/references/patterns-gemm.md new file mode 100644 index 000000000000..62d2b6d5be0e --- /dev/null +++ b/.agents/skills/kernel-triton-writing/references/patterns-gemm.md @@ -0,0 +1,292 @@ + + + + +# Triton GEMM Patterns + +Reusable matrix multiplication patterns from Triton tutorials 03, 08, 09, 10. + +## Matrix Multiplication + +Block-tiled GEMM with L2 cache optimization. The workhorse pattern for dense matmul. + +### Autotune Configs + +| Config | BLOCK_M | BLOCK_N | BLOCK_K | num_stages | num_warps | +|--------|---------|---------|---------|------------|-----------| +| 1 | 128 | 256 | 64 | 3 | 8 | +| 2 | 64 | 256 | 32 | 4 | 4 | +| 3 | 128 | 128 | 32 | 4 | 4 | +| 4 | 128 | 64 | 32 | 4 | 4 | +| 5 | 64 | 128 | 32 | 4 | 4 | +| 6 | 128 | 32 | 32 | 4 | 4 | +| 7 | 64 | 32 | 32 | 5 | 2 | +| 8 | 32 | 64 | 32 | 5 | 2 | + +All configs use `GROUP_SIZE_M=8`. `key=["M", "N", "K"]` triggers re-autotuning on shape change. + +### Complete Kernel + +```python +@triton.autotune( + configs=[ + triton.Config({"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 8}, num_stages=3, num_warps=8), + triton.Config({"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 8}, num_stages=4, num_warps=4), + triton.Config({"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 8}, num_stages=4, num_warps=4), + triton.Config({"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 8}, num_stages=4, num_warps=4), + triton.Config({"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 8}, num_stages=4, num_warps=4), + triton.Config({"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 8}, num_stages=4, num_warps=4), + triton.Config({"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 8}, num_stages=5, num_warps=2), + triton.Config({"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 8}, num_stages=5, num_warps=2), + ], + key=["M", "N", "K"], +) +@triton.jit +def matmul_kernel( + a_ptr, b_ptr, c_ptr, + M, N, K, + stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn, + BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, + GROUP_SIZE_M: tl.constexpr, +): + """C = A @ B. A is (M,K), B is (K,N), C is (M,N).""" + pid = tl.program_id(axis=0) + num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) + num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) + + # L2 cache optimization: super-grouping — nearby pids share B columns + num_pid_in_group = GROUP_SIZE_M * num_pid_n + group_id = pid // num_pid_in_group + first_pid_m = group_id * GROUP_SIZE_M + group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) + pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) + pid_n = (pid % num_pid_in_group) // group_size_m + + # Multi-dimensional pointer arithmetic via broadcasting + # 1D offset vectors + strides -> 2D block of pointers + offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M + offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N + offs_k = tl.arange(0, BLOCK_SIZE_K) + a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) + b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) + + # Accumulate along K + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0) + b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0) + accumulator = tl.dot(a, b, accumulator) + a_ptrs += BLOCK_SIZE_K * stride_ak + b_ptrs += BLOCK_SIZE_K * stride_bk + c = accumulator.to(tl.float16) + + # Store with boundary masking + offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) + offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) + c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] + c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) + tl.store(c_ptrs, c, mask=c_mask) +``` + +### Launch Wrapper + +```python +def matmul(a, b): + assert a.shape[1] == b.shape[0], "Incompatible dimensions" + assert a.is_contiguous(), "Matrix A must be contiguous" + M, K = a.shape + K, N = b.shape + c = torch.empty((M, N), device=a.device, dtype=torch.float16) + grid = lambda META: (triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),) + matmul_kernel[grid]( + a, b, c, M, N, K, + a.stride(0), a.stride(1), b.stride(0), b.stride(1), c.stride(0), c.stride(1), + ) + return c +``` + +### Key Patterns + +- **Pointer broadcasting:** `offs_row[:, None] * stride_row + offs_col[None, :] * stride_col` creates 2D pointer block from 1D offsets. +- **tl.dot(a, b, acc):** Accumulates `a @ b` into `acc`. Always use float32 accumulator. +- **Super-grouping:** `GROUP_SIZE_M` controls how many M-tiles share N-tiles, improving L2 hit rate on B. Typically 8. +- **Boundary masking:** `% M`/`% N` wraps OOB offsets to valid addresses (loads still masked). K-dim uses explicit `mask=offs_k < remaining`. + +## Grouped GEMM + +Batched independent matmuls in a single persistent kernel. Use case: mixture-of-experts (MoE). + +### Core Pattern + +```python +@triton.jit +def grouped_matmul_kernel( + a_ptrs, b_ptrs, c_ptrs, # device arrays of per-group pointers + m_sizes, n_sizes, k_sizes, # per-group dimensions + lds_a, lds_b, lds_c, # per-group leading dimensions + group_offsets, # cumulative tile count per group + num_tiles, + NUM_SM: tl.constexpr, BLOCK_SIZE_M: tl.constexpr, + BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, +): + tile_idx = tl.program_id(0) + # Persistent: each SM strides across all tiles + for tile_idx in tl.range(tile_idx, num_tiles, NUM_SM, num_stages=0): + # Binary-search group_offsets to find which group owns this tile + # Compute (pid_m, pid_n) within that group + # Standard matmul accumulation loop for group's (M,K) x (K,N) + ... + +# Launch: one program per SM +NUM_SM = torch.cuda.get_device_properties("cuda").multi_processor_count +grouped_matmul_kernel[(NUM_SM,)]( + a_ptrs, b_ptrs, c_ptrs, m_sizes, n_sizes, k_sizes, + lds_a, lds_b, lds_c, group_offsets, total_tiles, + NUM_SM=NUM_SM, BLOCK_SIZE_M=128, BLOCK_SIZE_N=128, BLOCK_SIZE_K=32, +) +``` + +### TMA Variant (Hopper+) + +```python +# Device-side TMA descriptors — shape varies per group +desc_a = tl.make_tensor_descriptor( + a_group_ptr, shape=[M_g, K_g], strides=[lda, 1], + block_shape=[BLOCK_SIZE_M, BLOCK_SIZE_K], +) +desc_b = tl.make_tensor_descriptor( + b_group_ptr, shape=[K_g, N_g], strides=[ldb, 1], + block_shape=[BLOCK_SIZE_K, BLOCK_SIZE_N], +) +a = desc_a.load([pid_m * BLOCK_SIZE_M, k * BLOCK_SIZE_K]) +b = desc_b.load([k * BLOCK_SIZE_K, pid_n * BLOCK_SIZE_N]) +``` + +### Key Patterns + +- **Persistent kernel:** Grid = NUM_SM. Each program loops via `tl.range(pid, total, NUM_SM)`. +- **Device-side scheduling:** Binary search on cumulative tile offsets maps flat tile_id to (group, tile_m, tile_n). +- **`tl.make_tensor_descriptor`:** Creates TMA descriptors on-device (needed because shape changes per group). +- **`num_stages=0`:** Outer loop has no pipelining; inner K loop is pipelined. + +## Persistent Matmul + +TMA descriptors and warp specialization for Hopper/Blackwell. Three progressive variants. + +### Variant 1: Persistent with Pointer Arithmetic + +Same as basic GEMM but with `tl.range(start_pid, num_tiles, NUM_SM)` outer loop +and grid = `(NUM_SM,)`. See Grouped GEMM pattern above for the persistent loop structure. + +### Variant 2: TMA Descriptors + +```python +# Host-side: create TMA descriptors before launch +from triton.tools.experimental_descriptor import create_2d_tma_descriptor +desc_a = create_2d_tma_descriptor(a_ptr, M, K, BLOCK_SIZE_M, BLOCK_SIZE_K, a.element_size()) +desc_b = create_2d_tma_descriptor(b_ptr, K, N, BLOCK_SIZE_K, BLOCK_SIZE_N, b.element_size()) + +# Kernel: load via hardware TMA unit (no manual pointer math) +@triton.jit +def matmul_kernel_tma(desc_a, desc_b, c_ptr, M, N, K, ...): + # Inside K-loop: + a = tl._experimental_descriptor_load( + desc_a, [pid_m * BLOCK_SIZE_M, k * BLOCK_SIZE_K], + [BLOCK_SIZE_M, BLOCK_SIZE_K], tl.float16) + b = tl._experimental_descriptor_load( + desc_b, [k * BLOCK_SIZE_K, pid_n * BLOCK_SIZE_N], + [BLOCK_SIZE_K, BLOCK_SIZE_N], tl.float16) + accumulator = tl.dot(a, b, accumulator) +``` + +### Variant 3: Warp Specialization + +```python +# Warps split into producers (TMA loads) and consumers (tl.dot compute). +# Compiler manages producer/consumer synchronization automatically. +matmul_kernel_warp_spec[(NUM_SM,)]( + desc_a, desc_b, c, M, N, K, ..., + BLOCK_SIZE_M=128, BLOCK_SIZE_N=256, BLOCK_SIZE_K=64, + num_stages=4, num_warps=8, + num_consumer_groups=2, # warp groups for compute + num_buffers_warp_spec=4, # pipeline depth for producer/consumer overlap +) +``` + +### FP8 Support (compute capability >= 9.0) + +```python +a = tl.load(a_ptrs, mask=..., other=0.0).to(tl.float8e5m2) # or tl.float8e4m3fn +b = tl.load(b_ptrs, mask=..., other=0.0).to(tl.float8e5m2) +accumulator = tl.dot(a, b, accumulator) # acc stays float32 +``` + +### Key Patterns + +- **tl.range(start, end, step):** Persistent loop with SM-count stride. +- **TMA descriptors:** Host-side `create_2d_tma_descriptor` or device-side `tl.make_tensor_descriptor`. Offloads address gen to hardware. +- **Warp specialization:** `num_consumer_groups` + `num_buffers_warp_spec` split warps into memory producers and compute consumers. +- **Epilogue subtiling:** Slice accumulator along N for the store phase to cut register pressure. + +## Block-Scaled Matmul + +Per-block scale factors for microscaling (MX) formats. +Requires 5th-gen Tensor Cores (compute capability >= 10.0, Blackwell+). + +### Supported Formats + +| Format | Element Type | Scale Block Size | Platform | +|--------|-------------|------------------|----------| +| mxfp8 | float8 (e5m2 / e4m3) | 32 elements | NVIDIA + AMD | +| mxfp4 | float4 (e2m1) | 32 elements | NVIDIA + AMD | +| nvfp4 | float4 (e2m1) | 16 elements | NVIDIA only | + +### Kernel with tl.dot_scaled + +```python +@triton.jit +def matmul_kernel_block_scaled( + a_ptr, b_ptr, c_ptr, a_scale_ptr, b_scale_ptr, + M, N, K, + stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn, + stride_a_scale_m, stride_a_scale_k, stride_b_scale_n, stride_b_scale_k, + BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, + BLOCK_SIZE_K: tl.constexpr, GROUP_SIZE_M: tl.constexpr, +): + pid = tl.program_id(axis=0) + # ... super-grouping (same as basic GEMM) ... + + accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) + for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): + a = tl.load(a_ptrs, mask=..., other=0.0) + b = tl.load(b_ptrs, mask=..., other=0.0) + a_scale = tl.load(a_scale_ptrs) # [BLOCK_M, BLOCK_K // 32] + b_scale = tl.load(b_scale_ptrs) # [BLOCK_N, BLOCK_K // 32] + # Hardware-accelerated scaled dot product + accumulator = tl.dot_scaled(a, a_scale, "e4m3", b, b_scale, "e4m3", accumulator) + # advance pointers ... + tl.store(c_ptrs, accumulator.to(tl.float16), mask=c_mask) +``` + +### Key Patterns + +- **tl.dot_scaled(a, a_scale, a_fmt, b, b_scale, b_fmt, acc):** Single instruction applying per-block scales during matmul. Replaces manual dequant-then-multiply. +- **Format strings:** `"e4m3"`, `"e5m2"`, `"e2m1"` passed to `tl.dot_scaled`. +- **Preshuffling:** Always preprocess scales into vendor-specific layout before kernel launch. +- **Hardware:** NVIDIA CC >= 10.0 (Blackwell, PTX 8.7+). AMD CDNA3+ (MI300X). diff --git a/.agents/skills/kernel-triton-writing/references/troubleshooting.md b/.agents/skills/kernel-triton-writing/references/troubleshooting.md new file mode 100644 index 000000000000..c1269f24aa14 --- /dev/null +++ b/.agents/skills/kernel-triton-writing/references/troubleshooting.md @@ -0,0 +1,282 @@ + + + + +# Triton Troubleshooting, Debugging, and Benchmarking + +## Debug Operations — Compile-Time + +### static_print — Inspect Types and Constants at Compile Time + +Prints values during kernel compilation (not at runtime). Use to verify +constexpr values, tensor shapes, and dtypes. + +```python +@triton.jit +def kernel(X_ptr, N, BLOCK_SIZE: tl.constexpr): + tl.static_print("BLOCK_SIZE", BLOCK_SIZE) # prints the constexpr value + x = tl.load(X_ptr + tl.arange(0, BLOCK_SIZE)) + tl.static_print("x dtype", x.dtype) # prints the tensor dtype + tl.static_print("x shape", x.shape) # prints the tensor shape +``` + +Output appears in stderr during compilation (not on device): + +``` +BLOCK_SIZE 1024 +x dtype float32 +x shape (1024,) +``` + +### static_assert — Compile-Time Invariant Checks + +Fails compilation if condition is false. Use for constexpr guards. + +```python +@triton.jit +def kernel(X_ptr, BLOCK_SIZE: tl.constexpr): + tl.static_assert(BLOCK_SIZE % 32 == 0, "BLOCK_SIZE must be multiple of 32") + tl.static_assert(BLOCK_SIZE <= 4096, "BLOCK_SIZE too large") +``` + +| Function | Runs When | Requires TRITON_DEBUG | Use For | +|----------|-----------|----------------------|---------| +| `tl.static_print(label, value)` | Compilation | No | Inspecting types, shapes, constexpr values | +| `tl.static_assert(cond, msg)` | Compilation | No | Enforcing constexpr constraints | + +--- + +## Debug Operations — Runtime (On-Device) + +### device_print — Print Tensor Values from GPU + +Prints values at runtime from every active thread. Produces large output +on multi-element blocks; use masks or conditions to limit output. + +```python +@triton.jit +def kernel(X_ptr, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(0) + offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + x = tl.load(X_ptr + offs) + + # Print from all programs — very verbose + tl.device_print("x", x) + + # Print from only program 0 — much less output + if pid == 0: + tl.device_print("x[0]", x) +``` + +### device_assert — Runtime Assertions (Requires TRITON_DEBUG=1) + +Only executes when `TRITON_DEBUG=1` is set. Silent otherwise. + +```python +@triton.jit +def kernel(X_ptr, N, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(0) + offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + tl.device_assert(offs < N, "out-of-bounds access") + x = tl.load(X_ptr + offs) +``` + +```bash +# Enable device_assert checks +TRITON_DEBUG=1 .venv/bin/python my_kernel.py +``` + +| Function | Runs When | Requires TRITON_DEBUG | Use For | +|----------|-----------|----------------------|---------| +| `tl.device_print(label, value)` | Runtime (GPU) | No | Inspecting tensor values on device | +| `tl.device_assert(cond, msg)` | Runtime (GPU) | **Yes** (`=1`) | Bounds checks, NaN guards, invariants | + +### Gotcha: device_assert Does Nothing Without TRITON_DEBUG + +If you add `tl.device_assert` and your kernel still silently produces wrong results, +check that `TRITON_DEBUG=1` is exported **before** the kernel is compiled/cached. + +--- + +## Interpreter Mode — CPU Step-Through Debugging + +Setting `TRITON_INTERPRET=1` runs all Triton kernels on the CPU using NumPy, +bypassing GPU compilation entirely. This enables standard Python debugging. + +### Basic Usage + +```bash +TRITON_INTERPRET=1 .venv/bin/python my_kernel.py +``` + +### Debugging with pdb + +```bash +TRITON_INTERPRET=1 .venv/bin/python -m pdb my_kernel.py +``` + +Set breakpoints inside `@triton.jit` functions — they execute as normal Python +in interpreter mode. + +```python +@triton.jit +def kernel(X_ptr, Y_ptr, BLOCK_SIZE: tl.constexpr): + pid = tl.program_id(0) + offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + x = tl.load(X_ptr + offs) + # In interpreter mode, you can set a pdb breakpoint here + # and inspect x as a numpy array + import pdb; pdb.set_trace() # works only with TRITON_INTERPRET=1 + y = x * 2 + tl.store(Y_ptr + offs, y) +``` + +### Interpreter Mode Limitations + +| Limitation | Detail | +|------------|--------| +| No bfloat16 support | NumPy lacks native bf16; operations may error or use fp32 fallback | +| No indirect memory access | Gather/scatter patterns may not work correctly | +| No GPU-specific behavior | Race conditions, warp-level ops not simulated | +| Performance | Orders of magnitude slower than GPU — use small inputs only | +| Caching | Set `TRITON_INTERPRET=1` before any kernel is compiled/cached | +| atomic_add with fp16 | Known issue — may raise `ValueError('unsupported data type')` | + +--- + +## Third-Party Debug Tools + +| Tool | Vendor | Purpose | Usage | +|------|--------|---------|-------| +| `compute-sanitizer` | NVIDIA | Memory access checker (out-of-bounds, races) | `compute-sanitizer .venv/bin/python my_kernel.py` | +| `compute-sanitizer --tool memcheck` | NVIDIA | Detailed memory error reports | `compute-sanitizer --tool memcheck .venv/bin/python my_kernel.py` | +| `compute-sanitizer --tool racecheck` | NVIDIA | Shared memory race detection | `compute-sanitizer --tool racecheck .venv/bin/python my_kernel.py` | +| AddressSanitizer | AMD (ROCm) | Memory error detection on AMD GPUs | Compile with ASan flags | +| `triton-viz` | Community | Visual trace of memory access patterns | `uv pip install triton-viz` | + +### compute-sanitizer Example + +```bash +# Check for out-of-bounds memory access +compute-sanitizer --tool memcheck .venv/bin/python my_kernel.py + +# Check for shared memory race conditions +compute-sanitizer --tool racecheck .venv/bin/python my_kernel.py +``` + +--- + +## Benchmarking — triton.testing + +### do_bench — Micro-Benchmark a Function + +```python +import triton + +ms = triton.testing.do_bench(lambda: my_kernel[grid](x, y, N, BLOCK_SIZE=1024)) +print(f"{ms:.3f} ms") +``` + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `fn` | `Callable` | required | Zero-arg function to benchmark (use lambda) | +| `warmup` | `int` | `25` | Warmup time in milliseconds | +| `rep` | `int` | `100` | Repetition time in milliseconds | +| `grad_to_none` | `torch.Tensor` | `None` | Reset this tensor's gradient to None each iteration | +| `quantiles` | `list[float]` | `None` | Percentiles to return (e.g., `[0.2, 0.5, 0.8]`) | +| `return_mode` | `str` | `"mean"` | `"min"`, `"max"`, `"mean"`, `"median"`, or `"all"` | + +### Benchmark Class + perf_report — Parameterized Benchmarks + +```python +import triton +from triton.testing import Benchmark, perf_report + +@perf_report( + Benchmark( + x_names=["N"], # argument to vary + x_vals=[2**i for i in range(10, 25)], # values for N + line_arg="provider", # line grouping + line_vals=["triton", "torch"], # line values + line_names=["Triton", "PyTorch"], # legend labels + plot_name="vector-add-performance", # plot filename + args={}, # fixed args + xlabel="Vector Size (N)", + ylabel="GB/s", + x_log=True, + ) +) +def benchmark(N, provider): + x = torch.randn(N, device='cuda', dtype=torch.float32) + y = torch.randn(N, device='cuda', dtype=torch.float32) + output = torch.empty_like(x) + if provider == "triton": + ms = triton.testing.do_bench(lambda: my_kernel[grid](x, y, output, N, BLOCK_SIZE=1024)) + else: + ms = triton.testing.do_bench(lambda: x + y) + gbps = 3 * x.numel() * x.element_size() / ms * 1e-6 # 3 = 2 reads + 1 write + return gbps + +# Run and save plot +benchmark.run(show_plots=True, save_path="./benchmarks/") +``` + +### Correctness Testing — torch.testing.assert_close + +Triton does not ship its own `assert_close`. Use PyTorch: + +```python +import torch +torch.testing.assert_close(triton_output, torch_reference, atol=1e-2, rtol=1e-2) +``` + +For fp16/bf16 kernels, use relaxed tolerances (`atol=1e-1, rtol=1e-1`). + +--- + +## Common Errors Table + +| Error / Symptom | Cause | Fix | +|-----------------|-------|-----| +| `shape mismatch` in binary op | Tensor shapes do not broadcast | Check shapes with `tl.static_print`; add `[:, None]` or `[None, :]` | +| `BLOCK_SIZE is not a constexpr` | Block size passed as runtime value | Add `: tl.constexpr` annotation to the parameter | +| `mask dimensions do not match` | Mask shape incompatible with load/store block | Ensure mask broadcasts to the pointer offset shape | +| OOM during autotuning | Too many `@triton.autotune` configs | Reduce config list; avoid combinatorial explosion of BLOCK_M/N/K | +| `device_assert` has no effect | `TRITON_DEBUG` not set to `1` | Export `TRITON_DEBUG=1` before running | +| Silent wrong results | Off-by-one in pointer arithmetic | Use `tl.device_print` to inspect offsets; test with `TRITON_INTERPRET=1` | +| `incompatible types` in store | Computed dtype does not match output pointer dtype | Cast explicitly: `tl.store(ptr, val.to(tl.float16))` | +| Kernel not updating after edit | Triton cache serving stale binary | Move the confirmed Triton cache directory aside and retry | +| `ValueError: unsupported data type` in interpreter | bf16 or fp8 used with `TRITON_INTERPRET=1` | Use fp16 or fp32 for interpreter debugging | +| `grid must be a tuple` | Lambda grid returns int, not tuple | Return `(value,)` with trailing comma | +| NaN output, correct logic | fp16 overflow in accumulator | Use `tl.float32` for accumulation, cast on store | +| `expected constexpr` in `tl.arange` | Non-constexpr argument to arange | Both args of `tl.arange(start, end)` must be constexpr | +| Mismatched results vs PyTorch | C integer division semantics | See concepts-semantics.md: Triton uses truncation, not floor division | +| `triton.OutOfResources` | Register/shared memory pressure | Reduce BLOCK_SIZE or number of live variables | + +--- + +## Environment Variables Reference + +| Variable | Value | Effect | +|----------|-------|--------| +| `TRITON_DEBUG` | `1` | Enable `device_assert`, extra runtime checks | +| `TRITON_INTERPRET` | `1` | Run kernels on CPU via NumPy (no GPU) | +| `TRITON_CACHE_DIR` | path | Override default cache directory (`~/.triton/cache/`) | +| `MLIR_ENABLE_DUMP` | `1` | Dump MLIR intermediate representations | +| `TRITON_PRINT_AUTOTUNING` | `1` | Print autotuning results to stderr | diff --git a/.agents/skills/kernel-triton-writing/scripts/__init__.py b/.agents/skills/kernel-triton-writing/scripts/__init__.py new file mode 100644 index 000000000000..650a104621d7 --- /dev/null +++ b/.agents/skills/kernel-triton-writing/scripts/__init__.py @@ -0,0 +1,17 @@ +# SPDX-FileCopyrightText: Copyright (c) 2011-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileComment: Copied from NVIDIA TensorRT-LLM at commit +# 395985c025c8d1cf5aa842bc752b337ba88721b6. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. diff --git a/.agents/skills/kernel-triton-writing/scripts/benchmark_kernel.py b/.agents/skills/kernel-triton-writing/scripts/benchmark_kernel.py new file mode 100644 index 000000000000..59235d8d560e --- /dev/null +++ b/.agents/skills/kernel-triton-writing/scripts/benchmark_kernel.py @@ -0,0 +1,308 @@ +#!/usr/bin/env -S uv run --python .venv/bin/python + +# SPDX-FileCopyrightText: Copyright (c) 2011-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileComment: Copied from NVIDIA TensorRT-LLM at commit +# 395985c025c8d1cf5aa842bc752b337ba88721b6. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Benchmark a Triton kernel using the fixed-name contract. + +Standalone script -- only Python stdlib required (torch/triton needed at +runtime for GPU benchmarking, but not for --mock mode). +Outputs structured JSON to stdout. + +Contract: + The kernel file must export: + - ``kernel_fn``: callable -- the Triton kernel wrapper + - ``reference_fn``: callable -- reference implementation (optional) + - ``get_inputs()``: returns a list of CUDA tensors + +Usage: + .venv/bin/python benchmark_kernel.py kernel.py [--warmup 10] [--iters 40] \ + [--timeout 120] [--mock] +""" + +from __future__ import annotations + +import argparse +import json +import os +import subprocess +import sys +import tempfile +import textwrap + +# --------------------------------------------------------------------------- +# Benchmark harness generation +# --------------------------------------------------------------------------- + + +def _build_benchmark_script( + kernel_path: str, + warmup: int, + iters: int, +) -> str: + """Generate a temporary benchmark harness script. + + The harness: + 1. Imports the kernel module using the fixed-name contract + 2. Calls ``get_inputs()`` for shared test tensors + 3. Benchmarks ``kernel_fn`` with CUDA events + 4. If ``reference_fn`` exists, benchmarks it too + 5. Prints a structured ``BENCHMARK:`` line to stdout + """ + abs_kernel_path = os.path.abspath(kernel_path) + kernel_dir = os.path.dirname(abs_kernel_path) + + return textwrap.dedent(f"""\ + import sys + import os + import importlib.util + import torch + + sys.path.insert(0, {kernel_dir!r}) + + # Import the kernel module + _spec = importlib.util.spec_from_file_location( + "kernel_module", {abs_kernel_path!r} + ) + _mod = importlib.util.module_from_spec(_spec) + _spec.loader.exec_module(_mod) + + # Validate fixed-name contract + for _attr in ("kernel_fn", "get_inputs"): + if not hasattr(_mod, _attr): + print(f"ERROR:kernel file missing required export: {{_attr}}") + sys.exit(1) + + def gpu_benchmark(fn, args, warmup, iters): + for _ in range(warmup): + fn(*args) + torch.accelerator.synchronize() + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + for _ in range(iters): + fn(*args) + end.record() + torch.accelerator.synchronize() + return start.elapsed_time(end) / iters + + # Get inputs + inputs = _mod.get_inputs() + if not isinstance(inputs, (list, tuple)): + print("ERROR:get_inputs() must return a list or tuple of tensors") + sys.exit(1) + + # Clone inputs for each benchmark to avoid in-place mutation + def _clone(x): + if isinstance(x, torch.Tensor): + return x.clone() + return x + + # Benchmark kernel_fn + kern_inputs = [_clone(t) for t in inputs] + kernel_ms = gpu_benchmark( + _mod.kernel_fn, kern_inputs, + warmup={warmup}, iters={iters}, + ) + + # Benchmark reference_fn if available + if hasattr(_mod, "reference_fn"): + ref_inputs = [_clone(t) for t in inputs] + ref_ms = gpu_benchmark( + _mod.reference_fn, ref_inputs, + warmup={warmup}, iters={iters}, + ) + speedup = ref_ms / kernel_ms if kernel_ms > 0 else 0 + print( + f"BENCHMARK:kernel_ms={{kernel_ms:.6f}}," + f"ref_ms={{ref_ms:.6f}}," + f"speedup={{speedup:.4f}}" + ) + else: + print(f"BENCHMARK:kernel_ms={{kernel_ms:.6f}}") + """) + + +# --------------------------------------------------------------------------- +# Core benchmark function +# --------------------------------------------------------------------------- + + +def benchmark_kernel( + kernel_path: str, + warmup: int = 10, + iters: int = 40, + timeout: int = 120, +) -> dict: + """Benchmark kernel performance using the fixed-name contract. + + Args: + kernel_path: Path to Python file exporting ``kernel_fn``, + optionally ``reference_fn``, and ``get_inputs()``. + warmup: Number of warmup iterations. + iters: Number of measured iterations. + timeout: Execution timeout in seconds. + + Returns: + Dict with keys ``kernel_time_ms``, ``reference_time_ms``, + ``speedup``, ``warmup_iters``, ``benchmark_iters``. + """ + if not os.path.exists(kernel_path): + print(f"Kernel file not found: {kernel_path}", file=sys.stderr) + sys.exit(1) + + script = _build_benchmark_script(kernel_path, warmup, iters) + + with tempfile.NamedTemporaryFile( + mode="w", suffix=".py", delete=False + ) as script_file: + script_file.write(script) + script_path = script_file.name + + working_dir = os.path.dirname(os.path.abspath(kernel_path)) + + try: + result = subprocess.run( + [sys.executable, script_path], + capture_output=True, + text=True, + timeout=timeout, + cwd=working_dir, + ) + + output = result.stdout + result.stderr + + if "ERROR:" in output: + error_msg = output.split("ERROR:")[1].strip().split("\n")[0] + print(f"Benchmark error: {error_msg}", file=sys.stderr) + sys.exit(1) + + if "BENCHMARK:" not in output: + print(f"Benchmark failed:\n{output[:1000]}", file=sys.stderr) + sys.exit(1) + + result_line = [line for line in output.split("\n") if "BENCHMARK:" in line][0] + parts = result_line.split("BENCHMARK:")[1].split(",") + try: + parsed = {kv.split("=")[0]: float(kv.split("=")[1]) for kv in parts} + except (IndexError, ValueError) as exc: + return { + "kernel_time_ms": None, + "reference_time_ms": None, + "speedup": None, + "warmup_iters": warmup, + "benchmark_iters": iters, + "error": f"Failed to parse BENCHMARK line: {exc}", + } + + result_dict: dict = { + "kernel_time_ms": parsed["kernel_ms"], + "warmup_iters": warmup, + "benchmark_iters": iters, + } + + if "ref_ms" in parsed: + result_dict["reference_time_ms"] = parsed["ref_ms"] + result_dict["speedup"] = parsed["speedup"] + else: + result_dict["reference_time_ms"] = None + result_dict["speedup"] = None + + return result_dict + + except subprocess.TimeoutExpired: + print(f"Benchmark timed out after {timeout} seconds", file=sys.stderr) + sys.exit(1) + finally: + if os.path.exists(script_path): + os.unlink(script_path) + + +# --------------------------------------------------------------------------- +# Mock data +# --------------------------------------------------------------------------- + + +def _mock_data(warmup: int = 10, iters: int = 40) -> dict: + """Return realistic mock benchmark data for testing.""" + return { + "kernel_time_ms": 0.45, + "reference_time_ms": 1.23, + "speedup": 2.73, + "warmup_iters": warmup, + "benchmark_iters": iters, + } + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main() -> None: + """Entry point for CLI invocation.""" + parser = argparse.ArgumentParser( + description="Benchmark a Triton kernel using the fixed-name contract." + ) + parser.add_argument( + "kernel_path", + nargs="?", + help="Path to Python file exporting kernel_fn, reference_fn, get_inputs().", + ) + parser.add_argument( + "--warmup", + type=int, + default=10, + help="Warmup iterations (default: 10).", + ) + parser.add_argument( + "--iters", + type=int, + default=40, + help="Measured iterations (default: 40).", + ) + parser.add_argument( + "--timeout", + type=int, + default=120, + help="Execution timeout in seconds (default: 120).", + ) + parser.add_argument( + "--mock", + action="store_true", + help="Return mock data for testing (no GPU required).", + ) + args = parser.parse_args() + + if args.mock: + data = _mock_data(warmup=args.warmup, iters=args.iters) + elif args.kernel_path: + data = benchmark_kernel( + kernel_path=args.kernel_path, + warmup=args.warmup, + iters=args.iters, + timeout=args.timeout, + ) + else: + parser.error("Either --mock or kernel_path is required.") + + json.dump(data, sys.stdout, indent=2) + print() + + +if __name__ == "__main__": + main() diff --git a/.agents/skills/kernel-triton-writing/scripts/verify_kernel.py b/.agents/skills/kernel-triton-writing/scripts/verify_kernel.py new file mode 100644 index 000000000000..b96ad596c05d --- /dev/null +++ b/.agents/skills/kernel-triton-writing/scripts/verify_kernel.py @@ -0,0 +1,381 @@ +#!/usr/bin/env -S uv run --python .venv/bin/python + +# SPDX-FileCopyrightText: Copyright (c) 2011-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# SPDX-FileComment: Copied from NVIDIA TensorRT-LLM at commit +# 395985c025c8d1cf5aa842bc752b337ba88721b6. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Verify a Triton kernel against its reference using the fixed-name contract. + +Standalone script -- only Python stdlib required (torch/triton needed at +runtime for GPU verification, but not for --mock mode). +Outputs structured JSON to stdout. + +Contract: + The kernel file must export: + - ``kernel_fn``: callable -- the Triton kernel wrapper + - ``reference_fn``: callable -- reference implementation (same signature) + - ``get_inputs()``: returns a list of CUDA tensors + +Usage: + .venv/bin/python verify_kernel.py kernel.py [--rtol 1e-3] [--atol 1e-3] \ + [--timeout 60] [--mock] +""" + +from __future__ import annotations + +import argparse +import json +import os +import subprocess +import sys +import tempfile +import textwrap + +# --------------------------------------------------------------------------- +# Verification harness generation +# --------------------------------------------------------------------------- + + +def _build_verification_script( + kernel_path: str, + rtol: float, + atol: float, +) -> str: + """Generate a temporary verification harness script. + + The harness: + 1. Imports the kernel module using the fixed-name contract + 2. Calls ``get_inputs()`` for shared test tensors + 3. Runs ``reference_fn`` and ``kernel_fn`` with the same inputs + 4. Recursively compares outputs (handles tuples, lists, dicts, tensors, + scalars) + 5. Prints a structured ``RESULT:`` line to stdout + """ + abs_kernel_path = os.path.abspath(kernel_path) + kernel_dir = os.path.dirname(abs_kernel_path) + + return textwrap.dedent(f"""\ + import sys + import importlib.util + import math + + sys.path.insert(0, {kernel_dir!r}) + + # Import the kernel module + _spec = importlib.util.spec_from_file_location( + "kernel_module", {abs_kernel_path!r} + ) + _mod = importlib.util.module_from_spec(_spec) + _spec.loader.exec_module(_mod) + + # Validate fixed-name contract + for _attr in ("kernel_fn", "reference_fn", "get_inputs"): + if not hasattr(_mod, _attr): + print(f"ERROR:kernel file missing required export: {{_attr}}") + sys.exit(1) + + import torch + + # Get shared inputs + inputs = _mod.get_inputs() + if not isinstance(inputs, (list, tuple)): + print("ERROR:get_inputs() must return a list or tuple of tensors") + sys.exit(1) + + # Clone inputs for each call to avoid in-place mutation issues + def _clone(x): + if isinstance(x, torch.Tensor): + return x.clone() + return x + + ref_inputs = [_clone(t) for t in inputs] + kern_inputs = [_clone(t) for t in inputs] + + # Run both implementations + ref_out = _mod.reference_fn(*ref_inputs) + kern_out = _mod.kernel_fn(*kern_inputs) + + # Recursive comparison + _global_max_abs = 0.0 + _global_max_rel = 0.0 + _all_correct = True + _mismatches = [] + + + def _compare(ref, kern, path="root"): + global _global_max_abs, _global_max_rel, _all_correct + if isinstance(ref, torch.Tensor) and isinstance(kern, torch.Tensor): + abs_diff = (kern.float() - ref.float()).abs() + max_abs = abs_diff.max().item() + ref_abs = ref.float().abs() + safe_ref = torch.where( + ref_abs > 0, ref_abs, torch.ones_like(ref_abs) + ) + max_rel = (abs_diff / safe_ref).max().item() + _global_max_abs = max(_global_max_abs, max_abs) + _global_max_rel = max(_global_max_rel, max_rel) + if not torch.allclose( + kern.float(), ref.float(), rtol={rtol}, atol={atol} + ): + _all_correct = False + _mismatches.append( + f"Tensor mismatch at {{path}}: " + f"max_abs={{max_abs:.2e}}" + ) + elif isinstance(ref, dict) and isinstance(kern, dict): + for k in ref: + if k not in kern: + _all_correct = False + _mismatches.append( + f"Missing key at {{path}}: {{k!r}}" + ) + return + _compare(ref[k], kern[k], path + f"[{{k!r}}]") + elif isinstance(ref, (list, tuple)) and isinstance(kern, (list, tuple)): + if len(ref) != len(kern): + _all_correct = False + _mismatches.append( + f"Length mismatch at {{path}}: " + f"{{len(ref)}} vs {{len(kern)}}" + ) + return + for i, (r, k) in enumerate(zip(ref, kern)): + _compare(r, k, path + f"[{{i}}]") + elif isinstance(ref, (int, float)) and isinstance(kern, (int, float)): + abs_d = abs(kern - ref) + safe_r = abs(ref) if abs(ref) > 0 else 1.0 + rel_d = abs_d / safe_r + _global_max_abs = max(_global_max_abs, abs_d) + _global_max_rel = max(_global_max_rel, rel_d) + if abs_d > {atol} + {rtol} * safe_r: + _all_correct = False + _mismatches.append( + f"Scalar mismatch at {{path}}: " + f"{{kern}} vs {{ref}}" + ) + else: + # Non-comparable types + if ref != kern: + _all_correct = False + _mismatches.append( + f"Value mismatch at {{path}}: " + f"{{kern!r}} vs {{ref!r}}" + ) + + + try: + _compare(ref_out, kern_out) + _result_parts = [ + f"passed={{_all_correct}}", + f"max_abs={{_global_max_abs}}", + f"max_rel={{_global_max_rel}}", + ] + print("RESULT:" + ",".join(_result_parts)) + for _m in _mismatches: + print(f"MISMATCH:{{_m}}", file=sys.stderr) + except Exception as e: + print(f"ERROR:{{e}}") + sys.exit(1) + """) + + +# --------------------------------------------------------------------------- +# Core verification function +# --------------------------------------------------------------------------- + + +def verify_kernel( + kernel_path: str, + rtol: float = 1e-3, + atol: float = 1e-3, + timeout: int = 60, +) -> dict: + """Verify kernel correctness using the fixed-name contract. + + Args: + kernel_path: Path to Python file exporting ``kernel_fn``, + ``reference_fn``, and ``get_inputs()``. + rtol: Relative tolerance. + atol: Absolute tolerance. + timeout: Execution timeout in seconds. + + Returns: + Dict with keys ``correct``, ``max_abs_diff``, ``max_rel_diff``, + ``details``. + """ + if not os.path.exists(kernel_path): + print(f"Kernel file not found: {kernel_path}", file=sys.stderr) + sys.exit(1) + + script = _build_verification_script(kernel_path, rtol, atol) + + with tempfile.NamedTemporaryFile( + mode="w", suffix=".py", delete=False + ) as script_file: + script_file.write(script) + script_path = script_file.name + + working_dir = os.path.dirname(os.path.abspath(kernel_path)) + + try: + result = subprocess.run( + [sys.executable, script_path], + capture_output=True, + text=True, + timeout=timeout, + cwd=working_dir, + ) + + output = result.stdout + result.stderr + + if "RESULT:" in output: + try: + result_line = [ + line for line in output.split("\n") if "RESULT:" in line + ][0] + parts_str = result_line.split("RESULT:")[1] + parts = parts_str.split(",") + passed = "True" in parts[0] + max_abs = float(parts[1].split("=")[1]) + max_rel = float(parts[2].split("=")[1]) + except (IndexError, ValueError) as exc: + return { + "correct": False, + "max_abs_diff": float("inf"), + "max_rel_diff": float("inf"), + "details": ( + f"Failed to parse RESULT line: {exc}. " + f"Raw output: {output[:500]}" + ), + } + + if passed: + details = ( + f"All outputs match within tolerance (rtol={rtol}, atol={atol})" + ) + else: + details = ( + f"Outputs differ beyond tolerance" + f" (max_abs={max_abs:.2e}, rtol={rtol}, atol={atol})" + ) + + return { + "correct": passed, + "max_abs_diff": max_abs, + "max_rel_diff": max_rel, + "details": details, + } + elif "ERROR:" in output: + error_msg = output.split("ERROR:")[1].strip().split("\n")[0] + return { + "correct": False, + "max_abs_diff": float("inf"), + "max_rel_diff": float("inf"), + "details": f"Verification error: {error_msg}", + } + else: + return { + "correct": False, + "max_abs_diff": float("inf"), + "max_rel_diff": float("inf"), + "details": f"Unexpected output: {output[:500]}", + } + + except subprocess.TimeoutExpired: + return { + "correct": False, + "max_abs_diff": float("inf"), + "max_rel_diff": float("inf"), + "details": f"Verification timed out after {timeout} seconds", + } + finally: + if os.path.exists(script_path): + os.unlink(script_path) + + +# --------------------------------------------------------------------------- +# Mock data +# --------------------------------------------------------------------------- + + +def _mock_data(rtol: float = 1e-3, atol: float = 1e-3) -> dict: + """Return realistic mock verification data for testing.""" + return { + "correct": True, + "max_abs_diff": 1.2e-7, + "max_rel_diff": 3.4e-6, + "details": f"All outputs match within tolerance (rtol={rtol}, atol={atol})", + } + + +# --------------------------------------------------------------------------- +# Main +# --------------------------------------------------------------------------- + + +def main() -> None: + """Entry point for CLI invocation.""" + parser = argparse.ArgumentParser( + description="Verify Triton kernel correctness using fixed-name contract." + ) + parser.add_argument( + "kernel_path", + nargs="?", + help="Path to Python file exporting kernel_fn, reference_fn, get_inputs().", + ) + parser.add_argument( + "--rtol", + type=float, + default=1e-3, + help="Relative tolerance (default: 1e-3).", + ) + parser.add_argument( + "--atol", + type=float, + default=1e-3, + help="Absolute tolerance (default: 1e-3).", + ) + parser.add_argument( + "--timeout", + type=int, + default=60, + help="Execution timeout in seconds (default: 60).", + ) + parser.add_argument( + "--mock", + action="store_true", + help="Return mock data for testing (no GPU required).", + ) + args = parser.parse_args() + + if args.mock: + data = _mock_data(rtol=args.rtol, atol=args.atol) + elif args.kernel_path: + data = verify_kernel( + kernel_path=args.kernel_path, + rtol=args.rtol, + atol=args.atol, + timeout=args.timeout, + ) + else: + parser.error("Either --mock or kernel_path is required.") + + json.dump(data, sys.stdout, indent=2) + print() + + +if __name__ == "__main__": + main() From e0617ddaae12d1374e67d8982c6899b9cbf24c18 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Wed, 2 Sep 2026 20:57:32 +0000 Subject: [PATCH 2/4] [Agents] Use vLLM kernel validation workflows Co-authored-by: OpenAI Codex Signed-off-by: Woosuk Kwon --- .../skills/kernel-triton-writing/ORIGIN.md | 15 +- .agents/skills/kernel-triton-writing/SKILL.md | 69 +--- .../kernel-triton-writing/scripts/__init__.py | 17 - .../scripts/benchmark_kernel.py | 308 -------------- .../scripts/verify_kernel.py | 381 ------------------ 5 files changed, 30 insertions(+), 760 deletions(-) delete mode 100644 .agents/skills/kernel-triton-writing/scripts/__init__.py delete mode 100644 .agents/skills/kernel-triton-writing/scripts/benchmark_kernel.py delete mode 100644 .agents/skills/kernel-triton-writing/scripts/verify_kernel.py diff --git a/.agents/skills/kernel-triton-writing/ORIGIN.md b/.agents/skills/kernel-triton-writing/ORIGIN.md index a809bf8c9641..d7a718fea90e 100644 --- a/.agents/skills/kernel-triton-writing/ORIGIN.md +++ b/.agents/skills/kernel-triton-writing/ORIGIN.md @@ -9,9 +9,12 @@ This skill was copied from NVIDIA's TensorRT-LLM repository: - License: Apache License 2.0 The NVIDIA copyright and Apache-2.0 SPDX notices are preserved in the copied -reference and script files. The vLLM copy adds explicit source comments, -removes unsupported skill metadata, replaces unsafe cache-removal examples, -keeps upstream Markdown tables/code blocks with targeted lint suppressions, -and adapts paths and Python commands to vLLM's `.venv/bin/python` and `uv` -workflow. The benchmark helper uses vLLM's accelerator-neutral synchronization -API. +reference files. The vLLM copy adds explicit source comments, removes +unsupported skill metadata, replaces unsafe cache-removal examples, and keeps +upstream Markdown tables/code blocks with targeted lint suppressions. + +The upstream standalone verification and benchmark scripts are omitted. vLLM +uses its existing parametrized kernel pytest suites for correctness and the +`kernel-microbenchmark` skill and `benchmarks/kernels/` for performance work. +The associated fixed-name export contract and workflow sections are adapted to +those vLLM conventions. diff --git a/.agents/skills/kernel-triton-writing/SKILL.md b/.agents/skills/kernel-triton-writing/SKILL.md index 2b5b6b124a2d..fc6b6925e46a 100644 --- a/.agents/skills/kernel-triton-writing/SKILL.md +++ b/.agents/skills/kernel-triton-writing/SKILL.md @@ -6,7 +6,7 @@ description: > The user's request must involve Triton explicitly. Covers Triton-specific patterns: fused elementwise, reductions (softmax, LayerNorm, RMSNorm), tiled GEMM with triton.autotune, and flash attention. Workflow: - design, write, verify (with fast-path for explicit requests). + route, design, implement, and validate in vLLM. license: Apache-2.0 metadata: author: NVIDIA Corporation @@ -30,8 +30,10 @@ licensed under Apache-2.0. See ORIGIN.md for the source and local changes. 1. Never benchmark before verification passes. 2. Always mask loads and stores for non-divisible shapes. -3. Include `kernel_fn`, `reference_fn`, and `get_inputs()` exports for companion scripts. -4. Always run `scripts/verify_kernel.py` to validate against the reference. +3. Validate the public wrapper against a PyTorch reference in the nearest + existing pytest suite. +4. Cover the relevant shapes, dtypes, strides, and boundary conditions with + explicit tolerances. ### FP16/BF16 Precision Rules (LOW FREEDOM -- follow exactly) @@ -53,7 +55,7 @@ Additional precision constraints: - `tl.sigmoid()` is unavailable in some Triton versions. Use `1.0 / (1.0 + tl.exp(-x_fp32))`. - Always cast back to `x.dtype` before `tl.store` -- mismatches cause "Type mismatch, store Float32 to Float16". - Unlike PyTorch, Triton does NOT auto-promote fp16/bf16 to fp32 for accumulation. Always use `tl.float32` accumulators for `tl.dot`. -- **TF32 for matmul:** On Ampere+/Hopper, `tl.dot` uses TF32 by default for fp32 inputs (same as `torch.mm`). Do NOT add `input_precision="ieee"` — it is 3-8x slower. TF32 is the correct default. If verification fails due to TF32 precision (~0.01-0.1 abs diff), ensure `reference_fn` also uses TF32 (plain `torch.mm`, no `allow_tf32=False`). +- **TF32 for matmul:** On Ampere+/Hopper, `tl.dot` uses TF32 by default for fp32 inputs (same as `torch.mm`). Do NOT add `input_precision="ieee"` — it is 3-8x slower. TF32 is the correct default. If verification fails due to TF32 precision (~0.01-0.1 abs diff), ensure the PyTorch reference also uses TF32 (plain `torch.mm`, no `allow_tf32=False`). ### CPU-GPU Sync Avoidance (LOW FREEDOM) @@ -196,17 +198,12 @@ def matmul_kernel( ### Phase 3: Write the Kernel -Create an output directory, then write the kernel file to `{output_dir}/kernel.py`. - -The kernel file MUST include: +Implement the kernel in the appropriate vLLM module and match nearby public +interfaces and style. The implementation should include: - `@triton.jit` decorated kernel function - `@triton.autotune` for production kernels (see [references/api-core.md](references/api-core.md)) - Python wrapper function (descriptive name for external import) -- **Fixed contract exports** (companion scripts rely on these exact names): - - `kernel_fn` — alias to the wrapper function - - `reference_fn(*args)` — PyTorch reference with identical signature - - `get_inputs()` — returns `list` of fresh CUDA tensors for testing/benchmarking Concise example (fused GELU + dropout): @@ -248,41 +245,25 @@ def fused_gelu_dropout_triton(x: torch.Tensor, p: float = 0.1) -> torch.Tensor: seed = (x.data_ptr() % (2**31)) ^ n_elements # sync-free seed fused_gelu_dropout_kernel[grid](x, out, n_elements, p, seed) return out - - -# --- Fixed contract (companion scripts rely on these names) --- -kernel_fn = fused_gelu_dropout_triton - -def reference_fn(x, p=0.1): - torch.manual_seed((x.data_ptr() % (2**31)) ^ x.numel()) - return torch.nn.functional.dropout( - torch.nn.functional.gelu(x), p, training=True - ) - -def get_inputs(): - return [torch.randn(128 * 1024 * 1024, device="cuda")] ``` For more patterns (SiLU+mul, RMSNorm, linear+GELU, add+LayerNorm), see [references/patterns-fusion.md](references/patterns-fusion.md). For GEMM patterns, see [references/patterns-gemm.md](references/patterns-gemm.md). ### Phase 4: Verify Correctness -Run the companion verification script: - -```bash -.venv/bin/python \ - .agents/skills/kernel-triton-writing/scripts/verify_kernel.py \ - {output_dir}/kernel.py \ - --rtol 1e-3 --atol 1e-3 -``` +Extend the nearest existing pytest suite under `tests/kernels/`. Compare the +public wrapper against a PyTorch reference and use the smallest cases that +cover the intended shapes, dtypes, strides, non-divisible boundaries, and any +in-place behavior. Keep tolerances explicit. -Output: +Run the focused test file: -```json -{"correct": true, "max_abs_diff": 1.2e-7, "max_rel_diff": 3.4e-6, "details": "..."} +```bash +.venv/bin/python -m pytest {test_path} -v ``` -**Stop if `correct: false`.** Fix the kernel before benchmarking. +Stop and fix the kernel if correctness fails. Do not use benchmark agreement +as a substitute for a pytest regression test. **Tolerance guide:** @@ -297,17 +278,9 @@ Output: Only benchmark if the user explicitly requests performance numbers. Skip this phase for correctness-focused requests. -```bash -.venv/bin/python \ - .agents/skills/kernel-triton-writing/scripts/benchmark_kernel.py \ - {output_dir}/kernel.py -``` - -Output: - -```json -{"kernel_time_ms": 0.45, "reference_time_ms": 1.23, "speedup": 2.73, "warmup_iters": 10, "benchmark_iters": 40} -``` +Use `$kernel-microbenchmark` for benchmark design and interpretation. Put +persistent kernel performance work in `benchmarks/kernels/` and follow its +CUPTI timing, cold-L2, CUDA graph, throughput, and reproducibility guidance. ## References (consult only when stuck) @@ -339,7 +312,7 @@ Only consult `references/` when: | `BLOCK_SIZE is not a constexpr` | Block size passed as runtime value | Add `: tl.constexpr` annotation | | `shape mismatch` in binary op | Tensor shapes don't broadcast | Check with `tl.static_print`; use `[:, None]` / `[None, :]` | | Large diffs everywhere | Wrong dtype in `tl.load` | Check load dtype matches input | -| Matmul 3-8x slower than expected | `input_precision="ieee"` on `tl.dot` | Remove it; use TF32 default. Ensure `reference_fn` also uses TF32 | +| Matmul 3-8x slower than expected | `input_precision="ieee"` on `tl.dot` | Remove it; use TF32 default. Ensure the PyTorch reference also uses TF32 | | Matmul ~0.01-0.1 abs diff vs reference | TF32 vs IEEE mismatch | Use same precision in both kernel and reference (TF32 for both) | | Diffs at boundaries | Missing mask | Add mask to all load/store ops | | Random diffs | Race condition | Check atomics and ordering | diff --git a/.agents/skills/kernel-triton-writing/scripts/__init__.py b/.agents/skills/kernel-triton-writing/scripts/__init__.py deleted file mode 100644 index 650a104621d7..000000000000 --- a/.agents/skills/kernel-triton-writing/scripts/__init__.py +++ /dev/null @@ -1,17 +0,0 @@ -# SPDX-FileCopyrightText: Copyright (c) 2011-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# SPDX-FileComment: Copied from NVIDIA TensorRT-LLM at commit -# 395985c025c8d1cf5aa842bc752b337ba88721b6. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. diff --git a/.agents/skills/kernel-triton-writing/scripts/benchmark_kernel.py b/.agents/skills/kernel-triton-writing/scripts/benchmark_kernel.py deleted file mode 100644 index 59235d8d560e..000000000000 --- a/.agents/skills/kernel-triton-writing/scripts/benchmark_kernel.py +++ /dev/null @@ -1,308 +0,0 @@ -#!/usr/bin/env -S uv run --python .venv/bin/python - -# SPDX-FileCopyrightText: Copyright (c) 2011-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# SPDX-FileComment: Copied from NVIDIA TensorRT-LLM at commit -# 395985c025c8d1cf5aa842bc752b337ba88721b6. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Benchmark a Triton kernel using the fixed-name contract. - -Standalone script -- only Python stdlib required (torch/triton needed at -runtime for GPU benchmarking, but not for --mock mode). -Outputs structured JSON to stdout. - -Contract: - The kernel file must export: - - ``kernel_fn``: callable -- the Triton kernel wrapper - - ``reference_fn``: callable -- reference implementation (optional) - - ``get_inputs()``: returns a list of CUDA tensors - -Usage: - .venv/bin/python benchmark_kernel.py kernel.py [--warmup 10] [--iters 40] \ - [--timeout 120] [--mock] -""" - -from __future__ import annotations - -import argparse -import json -import os -import subprocess -import sys -import tempfile -import textwrap - -# --------------------------------------------------------------------------- -# Benchmark harness generation -# --------------------------------------------------------------------------- - - -def _build_benchmark_script( - kernel_path: str, - warmup: int, - iters: int, -) -> str: - """Generate a temporary benchmark harness script. - - The harness: - 1. Imports the kernel module using the fixed-name contract - 2. Calls ``get_inputs()`` for shared test tensors - 3. Benchmarks ``kernel_fn`` with CUDA events - 4. If ``reference_fn`` exists, benchmarks it too - 5. Prints a structured ``BENCHMARK:`` line to stdout - """ - abs_kernel_path = os.path.abspath(kernel_path) - kernel_dir = os.path.dirname(abs_kernel_path) - - return textwrap.dedent(f"""\ - import sys - import os - import importlib.util - import torch - - sys.path.insert(0, {kernel_dir!r}) - - # Import the kernel module - _spec = importlib.util.spec_from_file_location( - "kernel_module", {abs_kernel_path!r} - ) - _mod = importlib.util.module_from_spec(_spec) - _spec.loader.exec_module(_mod) - - # Validate fixed-name contract - for _attr in ("kernel_fn", "get_inputs"): - if not hasattr(_mod, _attr): - print(f"ERROR:kernel file missing required export: {{_attr}}") - sys.exit(1) - - def gpu_benchmark(fn, args, warmup, iters): - for _ in range(warmup): - fn(*args) - torch.accelerator.synchronize() - start = torch.cuda.Event(enable_timing=True) - end = torch.cuda.Event(enable_timing=True) - start.record() - for _ in range(iters): - fn(*args) - end.record() - torch.accelerator.synchronize() - return start.elapsed_time(end) / iters - - # Get inputs - inputs = _mod.get_inputs() - if not isinstance(inputs, (list, tuple)): - print("ERROR:get_inputs() must return a list or tuple of tensors") - sys.exit(1) - - # Clone inputs for each benchmark to avoid in-place mutation - def _clone(x): - if isinstance(x, torch.Tensor): - return x.clone() - return x - - # Benchmark kernel_fn - kern_inputs = [_clone(t) for t in inputs] - kernel_ms = gpu_benchmark( - _mod.kernel_fn, kern_inputs, - warmup={warmup}, iters={iters}, - ) - - # Benchmark reference_fn if available - if hasattr(_mod, "reference_fn"): - ref_inputs = [_clone(t) for t in inputs] - ref_ms = gpu_benchmark( - _mod.reference_fn, ref_inputs, - warmup={warmup}, iters={iters}, - ) - speedup = ref_ms / kernel_ms if kernel_ms > 0 else 0 - print( - f"BENCHMARK:kernel_ms={{kernel_ms:.6f}}," - f"ref_ms={{ref_ms:.6f}}," - f"speedup={{speedup:.4f}}" - ) - else: - print(f"BENCHMARK:kernel_ms={{kernel_ms:.6f}}") - """) - - -# --------------------------------------------------------------------------- -# Core benchmark function -# --------------------------------------------------------------------------- - - -def benchmark_kernel( - kernel_path: str, - warmup: int = 10, - iters: int = 40, - timeout: int = 120, -) -> dict: - """Benchmark kernel performance using the fixed-name contract. - - Args: - kernel_path: Path to Python file exporting ``kernel_fn``, - optionally ``reference_fn``, and ``get_inputs()``. - warmup: Number of warmup iterations. - iters: Number of measured iterations. - timeout: Execution timeout in seconds. - - Returns: - Dict with keys ``kernel_time_ms``, ``reference_time_ms``, - ``speedup``, ``warmup_iters``, ``benchmark_iters``. - """ - if not os.path.exists(kernel_path): - print(f"Kernel file not found: {kernel_path}", file=sys.stderr) - sys.exit(1) - - script = _build_benchmark_script(kernel_path, warmup, iters) - - with tempfile.NamedTemporaryFile( - mode="w", suffix=".py", delete=False - ) as script_file: - script_file.write(script) - script_path = script_file.name - - working_dir = os.path.dirname(os.path.abspath(kernel_path)) - - try: - result = subprocess.run( - [sys.executable, script_path], - capture_output=True, - text=True, - timeout=timeout, - cwd=working_dir, - ) - - output = result.stdout + result.stderr - - if "ERROR:" in output: - error_msg = output.split("ERROR:")[1].strip().split("\n")[0] - print(f"Benchmark error: {error_msg}", file=sys.stderr) - sys.exit(1) - - if "BENCHMARK:" not in output: - print(f"Benchmark failed:\n{output[:1000]}", file=sys.stderr) - sys.exit(1) - - result_line = [line for line in output.split("\n") if "BENCHMARK:" in line][0] - parts = result_line.split("BENCHMARK:")[1].split(",") - try: - parsed = {kv.split("=")[0]: float(kv.split("=")[1]) for kv in parts} - except (IndexError, ValueError) as exc: - return { - "kernel_time_ms": None, - "reference_time_ms": None, - "speedup": None, - "warmup_iters": warmup, - "benchmark_iters": iters, - "error": f"Failed to parse BENCHMARK line: {exc}", - } - - result_dict: dict = { - "kernel_time_ms": parsed["kernel_ms"], - "warmup_iters": warmup, - "benchmark_iters": iters, - } - - if "ref_ms" in parsed: - result_dict["reference_time_ms"] = parsed["ref_ms"] - result_dict["speedup"] = parsed["speedup"] - else: - result_dict["reference_time_ms"] = None - result_dict["speedup"] = None - - return result_dict - - except subprocess.TimeoutExpired: - print(f"Benchmark timed out after {timeout} seconds", file=sys.stderr) - sys.exit(1) - finally: - if os.path.exists(script_path): - os.unlink(script_path) - - -# --------------------------------------------------------------------------- -# Mock data -# --------------------------------------------------------------------------- - - -def _mock_data(warmup: int = 10, iters: int = 40) -> dict: - """Return realistic mock benchmark data for testing.""" - return { - "kernel_time_ms": 0.45, - "reference_time_ms": 1.23, - "speedup": 2.73, - "warmup_iters": warmup, - "benchmark_iters": iters, - } - - -# --------------------------------------------------------------------------- -# Main -# --------------------------------------------------------------------------- - - -def main() -> None: - """Entry point for CLI invocation.""" - parser = argparse.ArgumentParser( - description="Benchmark a Triton kernel using the fixed-name contract." - ) - parser.add_argument( - "kernel_path", - nargs="?", - help="Path to Python file exporting kernel_fn, reference_fn, get_inputs().", - ) - parser.add_argument( - "--warmup", - type=int, - default=10, - help="Warmup iterations (default: 10).", - ) - parser.add_argument( - "--iters", - type=int, - default=40, - help="Measured iterations (default: 40).", - ) - parser.add_argument( - "--timeout", - type=int, - default=120, - help="Execution timeout in seconds (default: 120).", - ) - parser.add_argument( - "--mock", - action="store_true", - help="Return mock data for testing (no GPU required).", - ) - args = parser.parse_args() - - if args.mock: - data = _mock_data(warmup=args.warmup, iters=args.iters) - elif args.kernel_path: - data = benchmark_kernel( - kernel_path=args.kernel_path, - warmup=args.warmup, - iters=args.iters, - timeout=args.timeout, - ) - else: - parser.error("Either --mock or kernel_path is required.") - - json.dump(data, sys.stdout, indent=2) - print() - - -if __name__ == "__main__": - main() diff --git a/.agents/skills/kernel-triton-writing/scripts/verify_kernel.py b/.agents/skills/kernel-triton-writing/scripts/verify_kernel.py deleted file mode 100644 index b96ad596c05d..000000000000 --- a/.agents/skills/kernel-triton-writing/scripts/verify_kernel.py +++ /dev/null @@ -1,381 +0,0 @@ -#!/usr/bin/env -S uv run --python .venv/bin/python - -# SPDX-FileCopyrightText: Copyright (c) 2011-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. -# SPDX-License-Identifier: Apache-2.0 -# SPDX-FileCopyrightText: Copyright contributors to the vLLM project -# SPDX-FileComment: Copied from NVIDIA TensorRT-LLM at commit -# 395985c025c8d1cf5aa842bc752b337ba88721b6. -# -# Licensed under the Apache License, Version 2.0 (the "License"); -# you may not use this file except in compliance with the License. -# You may obtain a copy of the License at -# -# http://www.apache.org/licenses/LICENSE-2.0 -# -# Unless required by applicable law or agreed to in writing, software -# distributed under the License is distributed on an "AS IS" BASIS, -# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. -# See the License for the specific language governing permissions and -# limitations under the License. -"""Verify a Triton kernel against its reference using the fixed-name contract. - -Standalone script -- only Python stdlib required (torch/triton needed at -runtime for GPU verification, but not for --mock mode). -Outputs structured JSON to stdout. - -Contract: - The kernel file must export: - - ``kernel_fn``: callable -- the Triton kernel wrapper - - ``reference_fn``: callable -- reference implementation (same signature) - - ``get_inputs()``: returns a list of CUDA tensors - -Usage: - .venv/bin/python verify_kernel.py kernel.py [--rtol 1e-3] [--atol 1e-3] \ - [--timeout 60] [--mock] -""" - -from __future__ import annotations - -import argparse -import json -import os -import subprocess -import sys -import tempfile -import textwrap - -# --------------------------------------------------------------------------- -# Verification harness generation -# --------------------------------------------------------------------------- - - -def _build_verification_script( - kernel_path: str, - rtol: float, - atol: float, -) -> str: - """Generate a temporary verification harness script. - - The harness: - 1. Imports the kernel module using the fixed-name contract - 2. Calls ``get_inputs()`` for shared test tensors - 3. Runs ``reference_fn`` and ``kernel_fn`` with the same inputs - 4. Recursively compares outputs (handles tuples, lists, dicts, tensors, - scalars) - 5. Prints a structured ``RESULT:`` line to stdout - """ - abs_kernel_path = os.path.abspath(kernel_path) - kernel_dir = os.path.dirname(abs_kernel_path) - - return textwrap.dedent(f"""\ - import sys - import importlib.util - import math - - sys.path.insert(0, {kernel_dir!r}) - - # Import the kernel module - _spec = importlib.util.spec_from_file_location( - "kernel_module", {abs_kernel_path!r} - ) - _mod = importlib.util.module_from_spec(_spec) - _spec.loader.exec_module(_mod) - - # Validate fixed-name contract - for _attr in ("kernel_fn", "reference_fn", "get_inputs"): - if not hasattr(_mod, _attr): - print(f"ERROR:kernel file missing required export: {{_attr}}") - sys.exit(1) - - import torch - - # Get shared inputs - inputs = _mod.get_inputs() - if not isinstance(inputs, (list, tuple)): - print("ERROR:get_inputs() must return a list or tuple of tensors") - sys.exit(1) - - # Clone inputs for each call to avoid in-place mutation issues - def _clone(x): - if isinstance(x, torch.Tensor): - return x.clone() - return x - - ref_inputs = [_clone(t) for t in inputs] - kern_inputs = [_clone(t) for t in inputs] - - # Run both implementations - ref_out = _mod.reference_fn(*ref_inputs) - kern_out = _mod.kernel_fn(*kern_inputs) - - # Recursive comparison - _global_max_abs = 0.0 - _global_max_rel = 0.0 - _all_correct = True - _mismatches = [] - - - def _compare(ref, kern, path="root"): - global _global_max_abs, _global_max_rel, _all_correct - if isinstance(ref, torch.Tensor) and isinstance(kern, torch.Tensor): - abs_diff = (kern.float() - ref.float()).abs() - max_abs = abs_diff.max().item() - ref_abs = ref.float().abs() - safe_ref = torch.where( - ref_abs > 0, ref_abs, torch.ones_like(ref_abs) - ) - max_rel = (abs_diff / safe_ref).max().item() - _global_max_abs = max(_global_max_abs, max_abs) - _global_max_rel = max(_global_max_rel, max_rel) - if not torch.allclose( - kern.float(), ref.float(), rtol={rtol}, atol={atol} - ): - _all_correct = False - _mismatches.append( - f"Tensor mismatch at {{path}}: " - f"max_abs={{max_abs:.2e}}" - ) - elif isinstance(ref, dict) and isinstance(kern, dict): - for k in ref: - if k not in kern: - _all_correct = False - _mismatches.append( - f"Missing key at {{path}}: {{k!r}}" - ) - return - _compare(ref[k], kern[k], path + f"[{{k!r}}]") - elif isinstance(ref, (list, tuple)) and isinstance(kern, (list, tuple)): - if len(ref) != len(kern): - _all_correct = False - _mismatches.append( - f"Length mismatch at {{path}}: " - f"{{len(ref)}} vs {{len(kern)}}" - ) - return - for i, (r, k) in enumerate(zip(ref, kern)): - _compare(r, k, path + f"[{{i}}]") - elif isinstance(ref, (int, float)) and isinstance(kern, (int, float)): - abs_d = abs(kern - ref) - safe_r = abs(ref) if abs(ref) > 0 else 1.0 - rel_d = abs_d / safe_r - _global_max_abs = max(_global_max_abs, abs_d) - _global_max_rel = max(_global_max_rel, rel_d) - if abs_d > {atol} + {rtol} * safe_r: - _all_correct = False - _mismatches.append( - f"Scalar mismatch at {{path}}: " - f"{{kern}} vs {{ref}}" - ) - else: - # Non-comparable types - if ref != kern: - _all_correct = False - _mismatches.append( - f"Value mismatch at {{path}}: " - f"{{kern!r}} vs {{ref!r}}" - ) - - - try: - _compare(ref_out, kern_out) - _result_parts = [ - f"passed={{_all_correct}}", - f"max_abs={{_global_max_abs}}", - f"max_rel={{_global_max_rel}}", - ] - print("RESULT:" + ",".join(_result_parts)) - for _m in _mismatches: - print(f"MISMATCH:{{_m}}", file=sys.stderr) - except Exception as e: - print(f"ERROR:{{e}}") - sys.exit(1) - """) - - -# --------------------------------------------------------------------------- -# Core verification function -# --------------------------------------------------------------------------- - - -def verify_kernel( - kernel_path: str, - rtol: float = 1e-3, - atol: float = 1e-3, - timeout: int = 60, -) -> dict: - """Verify kernel correctness using the fixed-name contract. - - Args: - kernel_path: Path to Python file exporting ``kernel_fn``, - ``reference_fn``, and ``get_inputs()``. - rtol: Relative tolerance. - atol: Absolute tolerance. - timeout: Execution timeout in seconds. - - Returns: - Dict with keys ``correct``, ``max_abs_diff``, ``max_rel_diff``, - ``details``. - """ - if not os.path.exists(kernel_path): - print(f"Kernel file not found: {kernel_path}", file=sys.stderr) - sys.exit(1) - - script = _build_verification_script(kernel_path, rtol, atol) - - with tempfile.NamedTemporaryFile( - mode="w", suffix=".py", delete=False - ) as script_file: - script_file.write(script) - script_path = script_file.name - - working_dir = os.path.dirname(os.path.abspath(kernel_path)) - - try: - result = subprocess.run( - [sys.executable, script_path], - capture_output=True, - text=True, - timeout=timeout, - cwd=working_dir, - ) - - output = result.stdout + result.stderr - - if "RESULT:" in output: - try: - result_line = [ - line for line in output.split("\n") if "RESULT:" in line - ][0] - parts_str = result_line.split("RESULT:")[1] - parts = parts_str.split(",") - passed = "True" in parts[0] - max_abs = float(parts[1].split("=")[1]) - max_rel = float(parts[2].split("=")[1]) - except (IndexError, ValueError) as exc: - return { - "correct": False, - "max_abs_diff": float("inf"), - "max_rel_diff": float("inf"), - "details": ( - f"Failed to parse RESULT line: {exc}. " - f"Raw output: {output[:500]}" - ), - } - - if passed: - details = ( - f"All outputs match within tolerance (rtol={rtol}, atol={atol})" - ) - else: - details = ( - f"Outputs differ beyond tolerance" - f" (max_abs={max_abs:.2e}, rtol={rtol}, atol={atol})" - ) - - return { - "correct": passed, - "max_abs_diff": max_abs, - "max_rel_diff": max_rel, - "details": details, - } - elif "ERROR:" in output: - error_msg = output.split("ERROR:")[1].strip().split("\n")[0] - return { - "correct": False, - "max_abs_diff": float("inf"), - "max_rel_diff": float("inf"), - "details": f"Verification error: {error_msg}", - } - else: - return { - "correct": False, - "max_abs_diff": float("inf"), - "max_rel_diff": float("inf"), - "details": f"Unexpected output: {output[:500]}", - } - - except subprocess.TimeoutExpired: - return { - "correct": False, - "max_abs_diff": float("inf"), - "max_rel_diff": float("inf"), - "details": f"Verification timed out after {timeout} seconds", - } - finally: - if os.path.exists(script_path): - os.unlink(script_path) - - -# --------------------------------------------------------------------------- -# Mock data -# --------------------------------------------------------------------------- - - -def _mock_data(rtol: float = 1e-3, atol: float = 1e-3) -> dict: - """Return realistic mock verification data for testing.""" - return { - "correct": True, - "max_abs_diff": 1.2e-7, - "max_rel_diff": 3.4e-6, - "details": f"All outputs match within tolerance (rtol={rtol}, atol={atol})", - } - - -# --------------------------------------------------------------------------- -# Main -# --------------------------------------------------------------------------- - - -def main() -> None: - """Entry point for CLI invocation.""" - parser = argparse.ArgumentParser( - description="Verify Triton kernel correctness using fixed-name contract." - ) - parser.add_argument( - "kernel_path", - nargs="?", - help="Path to Python file exporting kernel_fn, reference_fn, get_inputs().", - ) - parser.add_argument( - "--rtol", - type=float, - default=1e-3, - help="Relative tolerance (default: 1e-3).", - ) - parser.add_argument( - "--atol", - type=float, - default=1e-3, - help="Absolute tolerance (default: 1e-3).", - ) - parser.add_argument( - "--timeout", - type=int, - default=60, - help="Execution timeout in seconds (default: 60).", - ) - parser.add_argument( - "--mock", - action="store_true", - help="Return mock data for testing (no GPU required).", - ) - args = parser.parse_args() - - if args.mock: - data = _mock_data(rtol=args.rtol, atol=args.atol) - elif args.kernel_path: - data = verify_kernel( - kernel_path=args.kernel_path, - rtol=args.rtol, - atol=args.atol, - timeout=args.timeout, - ) - else: - parser.error("Either --mock or kernel_path is required.") - - json.dump(data, sys.stdout, indent=2) - print() - - -if __name__ == "__main__": - main() From 71344ae8bc2f555cb2ab39cfc0dcb88a1140fb29 Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Wed, 2 Sep 2026 21:09:58 +0000 Subject: [PATCH 3/4] [Agents] Remove unsupported Triton guidance Co-authored-by: OpenAI Codex Signed-off-by: Woosuk Kwon --- .../skills/kernel-triton-writing/ORIGIN.md | 13 +- .agents/skills/kernel-triton-writing/SKILL.md | 413 ++++++------------ .../references/api-core.md | 331 -------------- .../references/api-language.md | 284 ------------ .../references/concepts-semantics.md | 199 --------- .../references/operator-routing.md | 125 ------ .../references/patterns-advanced.md | 320 -------------- .../references/patterns-basic.md | 238 ---------- .../references/patterns-fusion.md | 349 --------------- .../references/patterns-gemm.md | 292 ------------- .../references/semantics.md | 66 +++ .../references/troubleshooting.md | 319 +++----------- 12 files changed, 254 insertions(+), 2695 deletions(-) delete mode 100644 .agents/skills/kernel-triton-writing/references/api-core.md delete mode 100644 .agents/skills/kernel-triton-writing/references/api-language.md delete mode 100644 .agents/skills/kernel-triton-writing/references/concepts-semantics.md delete mode 100644 .agents/skills/kernel-triton-writing/references/operator-routing.md delete mode 100644 .agents/skills/kernel-triton-writing/references/patterns-advanced.md delete mode 100644 .agents/skills/kernel-triton-writing/references/patterns-basic.md delete mode 100644 .agents/skills/kernel-triton-writing/references/patterns-fusion.md delete mode 100644 .agents/skills/kernel-triton-writing/references/patterns-gemm.md create mode 100644 .agents/skills/kernel-triton-writing/references/semantics.md diff --git a/.agents/skills/kernel-triton-writing/ORIGIN.md b/.agents/skills/kernel-triton-writing/ORIGIN.md index d7a718fea90e..8a36bfab46b5 100644 --- a/.agents/skills/kernel-triton-writing/ORIGIN.md +++ b/.agents/skills/kernel-triton-writing/ORIGIN.md @@ -8,13 +8,16 @@ This skill was copied from NVIDIA's TensorRT-LLM repository: All rights reserved. - License: Apache License 2.0 -The NVIDIA copyright and Apache-2.0 SPDX notices are preserved in the copied -reference files. The vLLM copy adds explicit source comments, removes -unsupported skill metadata, replaces unsafe cache-removal examples, and keeps -upstream Markdown tables/code blocks with targeted lint suppressions. +The NVIDIA copyright and Apache-2.0 SPDX notices are preserved in the adapted +reference files. The vLLM copy adds explicit source comments and removes +unsupported skill metadata. The upstream standalone verification and benchmark scripts are omitted. vLLM uses its existing parametrized kernel pytest suites for correctness and the `kernel-microbenchmark` skill and `benchmarks/kernels/` for performance work. The associated fixed-name export contract and workflow sections are adapted to -those vLLM conventions. +those vLLM conventions. Copied API catalogs, fixed tuning recipes, performance +claims, and incomplete kernel examples were removed after review because they +duplicated versioned Triton documentation or were not generally supportable. +The remaining guidance directs contributors to current official Triton +documentation and device-specific measurement. diff --git a/.agents/skills/kernel-triton-writing/SKILL.md b/.agents/skills/kernel-triton-writing/SKILL.md index fc6b6925e46a..8b318de9fbac 100644 --- a/.agents/skills/kernel-triton-writing/SKILL.md +++ b/.agents/skills/kernel-triton-writing/SKILL.md @@ -2,11 +2,9 @@ name: kernel-triton-writing description: > ONLY for OpenAI Triton (@triton.jit) kernel development. NEVER use for - CUDA C++ kernels, TileIR, or profiling tools (ncu, nsys). - The user's request must involve Triton explicitly. Covers Triton-specific - patterns: fused elementwise, reductions (softmax, LayerNorm, RMSNorm), - tiled GEMM with triton.autotune, and flash attention. Workflow: - route, design, implement, and validate in vLLM. + CUDA C++ kernels, TileIR, or profiling tools such as ncu or nsys. Use when + the request explicitly involves implementing, reviewing, or debugging a + Triton kernel in vLLM. license: Apache-2.0 metadata: author: NVIDIA Corporation @@ -17,319 +15,160 @@ metadata: # Triton Kernel Writing - - -## Principles - -### Correctness First - -1. Never benchmark before verification passes. -2. Always mask loads and stores for non-divisible shapes. -3. Validate the public wrapper against a PyTorch reference in the nearest - existing pytest suite. -4. Cover the relevant shapes, dtypes, strides, and boundary conditions with - explicit tolerances. - -### FP16/BF16 Precision Rules (LOW FREEDOM -- follow exactly) - -Transcendental functions (`tl.exp`, `tl.log`, `tl.math.erf`, `tl.math.tanh`) require fp32 inputs. - -```python -# WRONG -- compilation error or wrong results with fp16/bf16: -result = tl.exp(x) - -# CORRECT -- cast to fp32, compute, cast back: -x_fp32 = x.to(tl.float32) -result = tl.exp(x_fp32).to(x.dtype) -``` - -Rule: any math function beyond basic arithmetic (+, -, *, /) requires fp32 cast in, original dtype cast out. - -Additional precision constraints: - -- `tl.sigmoid()` is unavailable in some Triton versions. Use `1.0 / (1.0 + tl.exp(-x_fp32))`. -- Always cast back to `x.dtype` before `tl.store` -- mismatches cause "Type mismatch, store Float32 to Float16". -- Unlike PyTorch, Triton does NOT auto-promote fp16/bf16 to fp32 for accumulation. Always use `tl.float32` accumulators for `tl.dot`. -- **TF32 for matmul:** On Ampere+/Hopper, `tl.dot` uses TF32 by default for fp32 inputs (same as `torch.mm`). Do NOT add `input_precision="ieee"` — it is 3-8x slower. TF32 is the correct default. If verification fails due to TF32 precision (~0.01-0.1 abs diff), ensure the PyTorch reference also uses TF32 (plain `torch.mm`, no `allow_tf32=False`). - -### CPU-GPU Sync Avoidance (LOW FREEDOM) - -Never call `.item()` in kernel wrappers. It forces a CPU-GPU sync (~50-100us per call). - -| Pitfall | Fix | -|---------|-----| -| `tensor.item()` for seed | `x.data_ptr() % (2**31)` | -| `torch.randint(...).item()` | Use tensor metadata for pseudo-random seed | -| Allocating output every call | Accept pre-allocated output as parameter | -| Python loops calling kernel | Batch operations | - -### C Integer Division Semantics (CRITICAL) - -Triton uses **C semantics** (round toward zero) for `//` and `%`, NOT Python semantics (round toward negative infinity). This only matters when operands can be negative. - -| Expression | Python | Triton/C | -|------------|--------|----------| -| `-7 // 2` | `-4` | `-3` | -| `-7 % 2` | `1` | `-1` | - -**Safe pattern:** Ensure all index/offset values are non-negative. If negative values are possible, use `(idx % BLOCK + BLOCK) % BLOCK`. - -See [references/concepts-semantics.md](references/concepts-semantics.md) for full rules and scalar-only exception. - -### Kernel Design Mental Model - -- **Parallelization axis:** Element-wise kernels parallelize over flattened elements. Row-wise kernels (LayerNorm, softmax) parallelize over rows. Matmul kernels tile in 2D (M, N). -- **Block size:** Power-of-2 only (256, 512, 1024, 2048). Start with 1024 for H100, 512 for V100. -- **Memory coalescing:** Adjacent threads must access adjacent memory addresses. The compiler handles this automatically from block-level pointer arithmetic. -- **Grid:** Use `triton.cdiv(n_elements, BLOCK_SIZE)`. With autotune, grid must be a lambda: `lambda meta: (triton.cdiv(n, meta['BLOCK_SIZE']),)`. -- **Decorator order:** `@triton.autotune` (outermost) -> `@triton.heuristics` -> `@triton.jit` (innermost). -- **`reset_to_zero`:** Required for autotune on kernels that accumulate (e.g., matmul output). Without it, later configs see leftover values from earlier trials. - -## Workflow - -**Fast path:** If the user explicitly requests a Triton kernel (e.g., "Write a Triton kernel for X", "Implement softmax in Triton"), start at **Phase 2**. Only use Phase 0-1 when the request is ambiguous about whether Triton is appropriate. - -### Phase 0: Route the Operator (only for ambiguous requests) - -Skip this phase if the user explicitly asks for a Triton kernel. Only use when the request is ambiguous (e.g., "optimize this operation"). - -Triton wins when 2+ operations can share registers instead of writing/reading global memory. Quick rules: - -| Pattern | Decision | -|---------|----------| -| Single element-wise op (`relu`, `sigmoid`) | SKIP — PyTorch already optimal | -| Standalone matmul | SKIP — cuBLAS is optimized | -| Standard attention | SKIP — Use FlashAttention | -| Element-wise chain (2+ ops), reduction, matmul + epilogue | USE TRITON | - -If SKIP, recommend the alternative and STOP. See [references/operator-routing.md](references/operator-routing.md) for edge cases. - -### Phase 1: Analyze the Operator (only for ambiguous requests) - -From the user's request, identify: (1) operation type, (2) parallelization strategy, (3) input shapes and dtypes. - -### Phase 2: Design the Kernel - -Pick the skeleton below that matches your operation. **These skeletons are sufficient for element-wise, reduction, matmul, and fusion kernels — do NOT read reference files for these common patterns.** Only consult `references/` when implementing uncommon patterns (grouped GEMM, TMA, extern functions) or debugging issues. - -**Element-wise skeleton** (GELU, dropout, fused ops on flat tensors): +Use this workflow for OpenAI Triton (`@triton.jit`) work in vLLM. Use the +`kernel-microbenchmark` skill as well when performance measurement or generated +code inspection is part of the task. + +## 1. Confirm the fit + +If Triton was explicitly requested, honor that choice. Otherwise, first inspect +nearby vLLM implementations and decide whether Triton is appropriate. Compare it +with existing vLLM operators, PyTorch compilation, and maintained vendor or +third-party kernels. Fusion potential alone does not guarantee a speedup, and a +standalone operation is not automatically a poor Triton candidate. + +Record the intended devices, dtypes, layouts, shape distribution, numerical +contract, and whether compilation or autotuning latency matters. Unless support +follows an existing vLLM compatibility contract, do not claim backend or device +support without relevant test coverage. + +## 2. Design around the contract + +- Define which program owns each output or whether an atomic update is required. + Make pointer arithmetic and strides explicit; do not assume inputs are + contiguous unless the public contract does. +- Mask every potentially out-of-bounds load and store. Select masked-load + values that are neutral for the operation, such as zero for a sum or negative + infinity for a floating-point maximum. +- Remember that `tl.where` evaluates both branches. Use load/store masks when a + branch must prevent a memory access. +- Choose accumulator and intermediate dtypes from the algorithm's numerical + requirements. Promotion is operation-specific: for example, reductions and + `tl.dot` have their own accumulation rules. Do not apply a blanket rule that + every math function requires fp32. +- `tl.store` converts values to the pointer element type. Cast explicitly when + it documents a deliberate rounding point, not because every store requires + one. +- Keep index calculations non-negative when possible. Triton integer division + and remainder can differ from Python for negative tensor operands; consult + [semantics.md](references/semantics.md) when porting signed index math. +- Treat block sizes, warp counts, stage counts, and launch order as tuning + choices, not GPU-family rules. Constraints on a specific operation, such as + `tl.arange`, do not imply that every meta-parameter must be a power of two. +- Avoid device-to-host scalar extraction such as `.item()` in a hot wrapper + when the source is a device tensor. It can synchronize the host and device. + Do not replace a real random seed with a pointer-derived value; preserve the + operator's RNG and determinism contract. + +Use a tuple launch grid when it is static. Use a callable grid only when it must +depend on compile-time meta-parameters, including autotuned values. + +## 3. Implement in vLLM + +Match the nearest vLLM module's public interface, dispatch, platform guards, +device handling, and style. Prefer extending an existing implementation over +creating a parallel abstraction. + +A basic one-dimensional kernel has this shape: ```python @triton.jit def kernel(x_ptr, out_ptr, n_elements, BLOCK_SIZE: tl.constexpr): - pid = tl.program_id(0) - offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) + offsets = tl.program_id(0) * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) mask = offsets < n_elements x = tl.load(x_ptr + offsets, mask=mask) - # ... compute ... + result = x # Replace with the operation. tl.store(out_ptr + offsets, result, mask=mask) ``` -**Row-wise skeleton** (softmax, LayerNorm, RMSNorm — one program per row): - -```python -@triton.jit -def kernel(x_ptr, out_ptr, n_cols, BLOCK_SIZE: tl.constexpr): - row_idx = tl.program_id(0) - col_offsets = tl.arange(0, BLOCK_SIZE) - mask = col_offsets < n_cols - x = tl.load(x_ptr + row_idx * n_cols + col_offsets, mask=mask, other=0.0) - # ... reduce / normalize ... - tl.store(out_ptr + row_idx * n_cols + col_offsets, result, mask=mask) -``` - -**Tiled matmul skeleton** (GEMM with 2D tiling, grouped ordering, and autotune): - -```python -@triton.autotune( - configs=[ - triton.Config({'BLOCK_M': 128, 'BLOCK_N': 256, 'BLOCK_K': 64, 'GROUP_SIZE_M': 8}, num_warps=8, num_stages=3), - triton.Config({'BLOCK_M': 64, 'BLOCK_N': 256, 'BLOCK_K': 32, 'GROUP_SIZE_M': 8}, num_warps=4, num_stages=4), - triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 32, 'GROUP_SIZE_M': 8}, num_warps=4, num_stages=4), - triton.Config({'BLOCK_M': 256, 'BLOCK_N': 64, 'BLOCK_K': 32, 'GROUP_SIZE_M': 8}, num_warps=4, num_stages=4), - ], - key=['M', 'N', 'K'], -) -@triton.jit -def matmul_kernel( - a_ptr, b_ptr, c_ptr, M, N, K, - stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn, - BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, - GROUP_SIZE_M: tl.constexpr, -): - pid = tl.program_id(0) - num_m_blocks = tl.cdiv(M, BLOCK_M) - num_n_blocks = tl.cdiv(N, BLOCK_N) - # Grouped ordering for L2 cache locality - num_pid_in_group = GROUP_SIZE_M * num_n_blocks - group_id = pid // num_pid_in_group - first_pid_m = group_id * GROUP_SIZE_M - group_size_m = min(num_m_blocks - first_pid_m, GROUP_SIZE_M) - pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) - pid_n = (pid % num_pid_in_group) // group_size_m - - offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - offs_k = tl.arange(0, BLOCK_K) - - a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak - b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - - for k in range(0, tl.cdiv(K, BLOCK_K)): - a_mask = (offs_m[:, None] < M) & (offs_k[None, :] < K) - b_mask = (offs_k[:, None] < K) & (offs_n[None, :] < N) - a = tl.load(a_ptrs, mask=a_mask, other=0.0) - b = tl.load(b_ptrs, mask=b_mask, other=0.0) - acc += tl.dot(a, b) - a_ptrs += BLOCK_K * stride_ak - b_ptrs += BLOCK_K * stride_bk - offs_k += BLOCK_K - - c_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N) - c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn - tl.store(c_ptrs, acc.to(c_ptr.dtype.element_ty), mask=c_mask) -``` - -### Phase 3: Write the Kernel - -Implement the kernel in the appropriate vLLM module and match nearby public -interfaces and style. The implementation should include: - -- `@triton.jit` decorated kernel function -- `@triton.autotune` for production kernels (see [references/api-core.md](references/api-core.md)) -- Python wrapper function (descriptive name for external import) +This is a structural example, not a recommended block size or complete public +wrapper. For GEMM, attention, persistent kernels, tensor descriptors, or other +specialized designs, start from a current official Triton tutorial and adapt it +to the installed Triton version and vLLM conventions. Do not copy experimental +APIs without checking that vLLM's supported Triton versions expose them. -Concise example (fused GELU + dropout): +### Autotuning -```python -import triton -import triton.language as tl -import torch - -@triton.autotune( - configs=[ - triton.Config({'BLOCK_SIZE': 1024}, num_warps=4), - triton.Config({'BLOCK_SIZE': 2048}, num_warps=8), - ], - key=['n_elements'], -) -@triton.jit -def fused_gelu_dropout_kernel( - x_ptr, out_ptr, n_elements, p, seed, - BLOCK_SIZE: tl.constexpr, -): - pid = tl.program_id(0) - offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - mask = offsets < n_elements - - x = tl.load(x_ptr + offsets, mask=mask) - x_fp32 = x.to(tl.float32) - x = (0.5 * x_fp32 * (1.0 + tl.math.erf(x_fp32 * 0.7071067811865476))).to(x.dtype) +Use `triton.autotune` only when its runtime cost and cache behavior fit the +deployment path. Fixed configurations, heuristics, or an existing vLLM tuning +mechanism may be preferable. - random = tl.rand(seed, offsets) - x = tl.where(random > p, x / (1.0 - p), 0.0) +When autotuning a kernel that mutates a buffer, ensure every candidate sees the +same initial state. Use the installed Triton version's `reset_to_zero`, +`restore_value`, or hooks as appropriate. A normal matmul that overwrites its +output does not need `reset_to_zero` merely because it uses an accumulator +internally. - tl.store(out_ptr + offsets, x, mask=mask) +Do not encode generic H100/A100/V100 recipes. SKU resources, shapes, dtypes, +compiler versions, and register pressure all affect the best configuration. +Measure representative production shapes on each supported target. +## 4. Verify correctness -def fused_gelu_dropout_triton(x: torch.Tensor, p: float = 0.1) -> torch.Tensor: - n_elements = x.numel() - out = torch.empty_like(x) - grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),) - seed = (x.data_ptr() % (2**31)) ^ n_elements # sync-free seed - fused_gelu_dropout_kernel[grid](x, out, n_elements, p, seed) - return out -``` +Extend the nearest existing pytest suite, normally under `tests/kernels/`. +Before writing tests, identify the public behavior, failure mode, and smallest +test level that catches it. -For more patterns (SiLU+mul, RMSNorm, linear+GELU, add+LayerNorm), see [references/patterns-fusion.md](references/patterns-fusion.md). For GEMM patterns, see [references/patterns-gemm.md](references/patterns-gemm.md). +Cover the dimensions relevant to the contract: -### Phase 4: Verify Correctness +- empty or minimum supported sizes and non-divisible tile boundaries; +- representative production shapes, including awkward dimensions; +- supported dtypes and layouts, including non-contiguous inputs if promised; +- aliasing, in-place behavior, RNG state, and determinism when applicable; +- numerical edge cases such as large magnitudes, zeros, infinities, or NaNs + when the operator defines behavior for them. -Extend the nearest existing pytest suite under `tests/kernels/`. Compare the -public wrapper against a PyTorch reference and use the smallest cases that -cover the intended shapes, dtypes, strides, non-divisible boundaries, and any -in-place behavior. Keep tolerances explicit. +Compare the public wrapper with an independent reference. Derive tolerances +from the dtype, operation, reduction depth, and documented precision mode; do +not use a universal tolerance table. For matmul-like operations, configure the +reference and Triton kernel to use comparable input and accumulation precision. -Run the focused test file: +Run the focused suite through the repository environment: ```bash -.venv/bin/python -m pytest {test_path} -v +.venv/bin/python -m pytest tests/path/to/test_file.py -v ``` -Stop and fix the kernel if correctness fails. Do not use benchmark agreement -as a substitute for a pytest regression test. - -**Tolerance guide:** - -| Dtype | rtol | atol | Notes | -|-------|------|------|-------| -| float16 | 1e-3 | 1e-3 | | -| bfloat16 | 1e-2 | 1e-2 | | -| float32 | 1e-5 | 1e-5 | Element-wise ops | -| float32 (matmul) | 1e-2 | 1e-1 | TF32 accumulation order differs between Triton tiles and cuBLAS | - -### Phase 5: Benchmark Performance (optional) - -Only benchmark if the user explicitly requests performance numbers. Skip this phase for correctness-focused requests. - -Use `$kernel-microbenchmark` for benchmark design and interpretation. Put -persistent kernel performance work in `benchmarks/kernels/` and follow its -CUPTI timing, cold-L2, CUDA graph, throughput, and reproducibility guidance. - -## References (consult only when stuck) - -The skeletons and principles above cover element-wise, reduction, matmul, and fusion kernels. **Do NOT read reference files for these common patterns.** - -Only consult `references/` when: - -- Implementing **uncommon patterns** (grouped GEMM, TMA, persistent matmul, extern functions) -- **Debugging** a compile error or incorrect result not covered by the error table below -- Needing **API details** for an unfamiliar `tl.*` operation +Do not use benchmark agreement as a substitute for a correctness test. -**How to search:** Grep for your keyword across `references/`. Read only the file Grep points to. +## 5. Measure only after correctness passes -| File | When to use | -|---|---| -| `references/api-core.md` | Unfamiliar `triton.autotune` / `triton.Config` options | -| `references/api-language.md` | Unfamiliar `tl.*` operations | -| `references/patterns-gemm.md` | Grouped GEMM, persistent matmul, TMA, MX formats | -| `references/patterns-advanced.md` | Flash attention details, backward passes, libdevice | -| `references/troubleshooting.md` | Debug ops, interpreter mode, env vars | +When performance is in scope, follow `$kernel-microbenchmark`. Put durable +kernel benchmarks under `benchmarks/kernels/` and report distributions, +representative shapes, hardware, software versions, and benchmark conditions. +Compare end-to-end cost when compilation, autotuning, allocations, or wrapper +overhead can affect the user-visible result. -## Error Handling and Troubleshooting +Treat a slowdown or regression as evidence to investigate, not proof that the +reference is optimal or Triton is unsuitable. Inspect generated code and +resource use when the benchmark warrants it. -### Common Errors +## 6. Debug systematically -| Error / Symptom | Cause | Fix | -|---------|-------|-----| -| "Type mismatch, store Float32 to Float16" | Missing `.to(x.dtype)` before store | Cast fp32 result back | -| `BLOCK_SIZE is not a constexpr` | Block size passed as runtime value | Add `: tl.constexpr` annotation | -| `shape mismatch` in binary op | Tensor shapes don't broadcast | Check with `tl.static_print`; use `[:, None]` / `[None, :]` | -| Large diffs everywhere | Wrong dtype in `tl.load` | Check load dtype matches input | -| Matmul 3-8x slower than expected | `input_precision="ieee"` on `tl.dot` | Remove it; use TF32 default. Ensure the PyTorch reference also uses TF32 | -| Matmul ~0.01-0.1 abs diff vs reference | TF32 vs IEEE mismatch | Use same precision in both kernel and reference (TF32 for both) | -| Diffs at boundaries | Missing mask | Add mask to all load/store ops | -| Random diffs | Race condition | Check atomics and ordering | -| NaN/Inf | Division by zero or fp16 overflow | Guard with epsilon; use `tl.float32` accumulator | -| `grid must be a tuple` | Grid lambda returns int, not tuple | Return `(value,)` with trailing comma | -| `expected constexpr` in `tl.arange` | Non-constexpr argument | Both args of `tl.arange(start, end)` must be constexpr | -| `triton.OutOfResources` | Register/shared memory pressure | Reduce BLOCK_SIZE or `num_stages` | -| Kernel not updating after edit | Stale compilation cache | Move the confirmed Triton cache directory aside and retry | -| Mismatched results vs PyTorch | C integer division semantics | Triton uses truncation; see `references/concepts-semantics.md` | +Reduce failures to a small shape, separate compilation failures from numerical +errors and memory faults, and use the tools in +[troubleshooting.md](references/troubleshooting.md). Never delete a broad or +unresolved cache path. If cache invalidation is justified, first resolve and +confirm the exact Triton cache directory, then move that directory aside so it +can be restored. -For extended error table, interpreter mode issues, and environment variables, see [references/troubleshooting.md](references/troubleshooting.md). +## Current authoritative references -### When to Abort +Check these before relying on signatures, backend support, or experimental +features: -Stop and report failure if: +- [Triton language API](https://triton-lang.org/main/python-api/triton.language.html) +- [Official Triton tutorials](https://triton-lang.org/main/getting-started/tutorials/) +- [Triton debugging guide](https://triton-lang.org/main/programming-guide/chapter-3/debugging.html) +- [vLLM kernel benchmarks](../../../benchmarks/kernels/) -1. **Not a good fit** -- Pure matmul or complex control flow (Phase 0 should catch this). -2. **Verification fails after 3 attempts** -- Numerical issues too severe to fix. -3. **No speedup** -- Reference is already well-optimized (cuBLAS, cuDNN). -4. **Hardware mismatch** -- Target GPU not available for testing. +Consult [semantics.md](references/semantics.md) for the few semantic hazards +worth keeping local. Consult [troubleshooting.md](references/troubleshooting.md) +for a compact debugging checklist. Prefer current official documentation and +the installed API over copied signature catalogs. diff --git a/.agents/skills/kernel-triton-writing/references/api-core.md b/.agents/skills/kernel-triton-writing/references/api-core.md deleted file mode 100644 index ed3c16685af8..000000000000 --- a/.agents/skills/kernel-triton-writing/references/api-core.md +++ /dev/null @@ -1,331 +0,0 @@ - - - - -# Triton Core API Reference - -## triton.jit - -Decorator that JIT-compiles a function into a GPU kernel using the Triton compiler. - -### Signature - -```python -@triton.jit # simple form — no parens needed -@triton.jit(do_not_specialize=None, do_not_specialize_on_alignment=None, - debug=None, noinline=None, repr=None, launch_metadata=None) -``` - -| Param | Type | Purpose | -|-------|------|---------| -| `do_not_specialize` | `Iterable[int\|str]\|None` | Args to skip value-specialization (by index or name) | -| `do_not_specialize_on_alignment` | `Iterable[int\|str]\|None` | Args to skip alignment-specialization | -| `debug` | `bool\|None` | Enable interpreter mode / debug prints | -| `noinline` | `bool\|None` | Prevent inlining when called from another jit'd function | - -### Implicit Pointer Conversion - -Objects with both `.data_ptr()` and `.dtype` (e.g., PyTorch tensors) are auto-converted -to device pointers. You never call `.data_ptr()` yourself in the launch call: - -```python -@triton.jit -def add_kernel(x_ptr, y_ptr, out_ptr, n, BLOCK: tl.constexpr): - pid = tl.program_id(0) - offs = pid * BLOCK + tl.arange(0, BLOCK) - mask = offs < n - tl.store(out_ptr + offs, tl.load(x_ptr + offs, mask=mask) + tl.load(y_ptr + offs, mask=mask), mask=mask) - -# Launch — pass tensors directly, NOT x.data_ptr() -add_kernel[(grid,)](x, y, out, x.numel(), BLOCK=1024) -``` - -### Specialization Rules - -Triton recompiles a kernel when argument properties change. For each argument: - -| Arg type | Specialized on | Effect | -|----------|---------------|--------| -| Pointer (tensor) | 16-byte alignment of `.data_ptr()` | Enables vectorized loads/stores | -| Integer scalar | Whether value == 1 | Dead-code elimination for guards | -| Integer scalar | Whether value is divisible by 16 | Enables optimized indexing | -| `tl.constexpr` | Exact value | Baked into compiled code as literal | - -**Gotcha:** Each unique specialization signature triggers a full recompile. If a size arg -oscillates between aligned/unaligned values, you get 2 cached versions (fine). But if you -pass truly random integers, use `do_not_specialize` to avoid cache explosion: - -```python -@triton.jit(do_not_specialize=["stride_x"]) -def my_kernel(x_ptr, stride_x, BLOCK: tl.constexpr): - ... -``` - -### constexpr Parameters - -Annotate with `tl.constexpr` to make a param a compile-time constant. Required for -values used in `tl.arange()`, `tl.zeros()`, tensor shapes, and `tl.static_assert`. -Each distinct value triggers recompilation. - -```python -@triton.jit -def kernel(x_ptr, N: tl.constexpr, BLOCK_SIZE: tl.constexpr): - ... -``` - ---- - -## triton.autotune - -Decorator that benchmarks multiple `triton.Config`s and caches the fastest per key. - -### Signature - -```python -@triton.autotune( - configs: list[triton.Config], - key: list[str], - prune_configs_by: dict | None = None, - reset_to_zero: list[str] | None = None, - restore_value: list[str] | None = None, - warmup: int = 25, - rep: int = 100, - use_cuda_graph: bool = False, -) -``` - -| Param | Purpose | -|-------|---------| -| `configs` | List of `triton.Config` objects to benchmark | -| `key` | Arg names whose values form the cache key (e.g., `["M", "N", "K"]`) | -| `prune_configs_by` | Dict with `early_config_prune`, `perf_model`, `top_k` to reduce search space | -| `reset_to_zero` | Arg names zeroed before each config trial (for accumulator correctness) | -| `restore_value` | Arg names restored to original value after each trial | -| `warmup` | Warmup time in ms per config (default 25) | -| `rep` | Benchmark time in ms per config (default 100) | - -### Example - -```python -@triton.autotune( - configs=[ - triton.Config({"BLOCK_M": 128, "BLOCK_N": 128}, num_warps=4, num_stages=3), - triton.Config({"BLOCK_M": 64, "BLOCK_N": 256}, num_warps=8, num_stages=3), - triton.Config({"BLOCK_M": 256, "BLOCK_N": 64}, num_warps=4, num_stages=4), - ], - key=["M", "N", "K"], - reset_to_zero=["c_ptr"], # zero output buffer between trials -) -@triton.jit -def matmul_kernel(a_ptr, b_ptr, c_ptr, M, N, K, - BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr): - ... -``` - -### Debugging Autotuning - -```bash -TRITON_PRINT_AUTOTUNING=1 .venv/bin/python my_script.py -``` - -Prints the winning config and total tuning time per kernel to stdout. - -### prune_configs_by Details - -```python -def early_prune(configs, named_args, **kwargs): - """Drop configs that exceed shared memory or have bad aspect ratios.""" - M, N = named_args["M"], named_args["N"] - return [c for c in configs if c.kwargs["BLOCK_M"] <= M and c.kwargs["BLOCK_N"] <= N] - -@triton.autotune( - configs=[...], - key=["M", "N"], - prune_configs_by={"early_config_prune": early_prune}, -) -``` - -Fields: `early_config_prune(configs, named_args, **kwargs) -> list[Config]`, -`perf_model(named_args, config, **kwargs) -> float` (estimated time), -`top_k` (int, keep only top_k configs after perf_model ranking). - -### Gotchas - -- `reset_to_zero` is critical for kernels that accumulate (e.g., matmul output). - Without it, later configs see leftover values from earlier trials. -- Autotuning happens on first call with each unique key combination. Subsequent calls - with the same key values use the cached winner. -- Decorator order: `@triton.autotune` must be the outermost, then `@triton.heuristics` - (if used), then `@triton.jit` innermost. - ---- - -## triton.Config - -Represents one candidate configuration for `triton.autotune`. - -### Signature - -```python -triton.Config( - kwargs: dict[str, Any], - num_warps: int = 4, - num_stages: int = 3, - num_ctas: int = 1, - maxnreg: int | None = None, - pre_hook: Callable | None = None, -) -``` - -| Param | Default | Purpose | -|-------|---------|---------| -| `kwargs` | (required) | Dict mapping `tl.constexpr` param names to values | -| `num_warps` | 4 | Threads per block = `num_warps * 32` | -| `num_stages` | 3 | Software pipelining depth for global loads | -| `num_ctas` | 1 | Cooperative thread arrays (multi-CTA kernels, Hopper+) | -| `maxnreg` | None | Max registers per thread (trades occupancy vs spilling) | -| `pre_hook` | None | `fn(args: dict)` called before kernel launch | - -### Example with pre_hook - -```python -def zero_output(args): - """Zero the output tensor before the kernel runs.""" - args["c_ptr"].zero_() - -triton.Config( - {"BLOCK_M": 128, "BLOCK_N": 128, "BLOCK_K": 32}, - num_warps=4, - num_stages=3, - pre_hook=zero_output, -) -``` - -### Tuning Guidance - -| Parameter | Small tiles / low occupancy | Large tiles / high throughput | -|-----------|----------------------------|-------------------------------| -| `num_warps` | 2-4 | 8-16 | -| `num_stages` | 2 (less shared mem) | 3-5 (hide global latency) | -| `maxnreg` | None (let compiler decide) | 128-255 (force occupancy) | - -**Gotcha:** `num_stages > 1` requires shared memory for buffering. Large tiles + -many stages can exceed shared memory limits, causing silent fallback or launch failure. - -### GPU-Specific Config Guidelines - -**H100 (Hopper):** HBM3, 168 SMs, large shared memory. - -- Prefer larger blocks (1024-4096), more warps (8-16), `num_stages=4+`. - -**A100 (Ampere):** Balanced config. - -- Block sizes 512-2048, `num_stages=3` typically optimal. - -**V100 (Volta):** Less shared memory. - -- Smaller blocks (256-1024), fewer stages (2), warps 4-8. - ---- - -## triton.heuristics - -Decorator that computes meta-parameters from kernel arguments at launch time, -avoiding the cost of autotuning for values that can be derived deterministically. - -### Signature - -```python -@triton.heuristics(values: dict[str, Callable]) -``` - -`values` maps constexpr parameter names to functions. Each function receives -the kernel's named arguments as a dict and returns the computed value. - -### Example - -```python -@triton.heuristics( - values={ - "BLOCK_SIZE": lambda args: triton.next_power_of_2(args["n_cols"]), - "num_warps": lambda args: 4 if args["n_cols"] <= 1024 else 8, - } -) -@triton.jit -def softmax_kernel(x_ptr, out_ptr, n_cols, - BLOCK_SIZE: tl.constexpr): - ... - -# Launch — BLOCK_SIZE is computed automatically, not passed -softmax_kernel[(n_rows,)](x, out, n_cols) -``` - -### Combined with autotune - -```python -@triton.autotune( - configs=[ - triton.Config({"BLOCK_M": 64}, num_warps=4), - triton.Config({"BLOCK_M": 128}, num_warps=8), - ], - key=["M", "N"], -) -@triton.heuristics( - values={"BLOCK_N": lambda args: triton.next_power_of_2(args["N"])} -) -@triton.jit -def kernel(x_ptr, M, N, BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr): - ... -``` - -**Required order (outermost to innermost):** `@autotune` -> `@heuristics` -> `@jit` - -### triton.next_power_of_2 - -```python -triton.next_power_of_2(n) # 7 -> 8, 8 -> 8, 1000 -> 1024 -``` - -Host-side utility. Common pattern: derive BLOCK_SIZE from a problem dimension -so the block covers the full row/column in one pass. - -### Gotchas - -- Heuristic functions run on the host (CPU) at every kernel launch, not on GPU. -- `@heuristics` must come AFTER `@autotune` but BEFORE `@jit` in decorator stack. -- Values computed by heuristics override any same-named values in `triton.Config.kwargs`. -- Returning non-power-of-2 for a `BLOCK_*` param is valid but usually suboptimal. - ---- - -## Decorator Stacking Summary - -``` -@triton.autotune(...) # outermost — optional -@triton.heuristics(...) # middle — optional -@triton.jit # innermost — required -def kernel(...): -``` - -| Combo | Use case | -|-------|----------| -| `@jit` only | Fixed config, simplest kernels | -| `@autotune` + `@jit` | Search over tile sizes and hardware params | -| `@heuristics` + `@jit` | Derive config from args, no search needed | -| `@autotune` + `@heuristics` + `@jit` | Search some params, derive others | diff --git a/.agents/skills/kernel-triton-writing/references/api-language.md b/.agents/skills/kernel-triton-writing/references/api-language.md deleted file mode 100644 index 02ccb1061b45..000000000000 --- a/.agents/skills/kernel-triton-writing/references/api-language.md +++ /dev/null @@ -1,284 +0,0 @@ - - - - -# Triton Language API (`triton.language` / `tl`) - -## Programming Model - -| Function | Signature | Notes | -|---|---|---| -| `program_id` | `program_id(axis)` | Returns ID of current program instance along `axis` (0, 1, or 2) | -| `num_programs` | `num_programs(axis)` | Returns number of program instances along `axis` | -| `tensor` | N-D array type | Block-structured; all ops are implicitly vectorized over the block | -| `tensor_descriptor` | Returned by `make_tensor_descriptor` | Opaque handle; backed by TMA on supported NVIDIA GPUs | - -## Creation Operations - -| Function | Signature | Notes | -|---|---|---| -| `arange` | `arange(start, end)` | Half-open `[start, end)`, returns 1-D int32 tensor | -| `full` | `full(shape, value, dtype)` | Broadcast scalar `value` to `shape` | -| `zeros` | `zeros(shape, dtype)` | Shorthand for `full(shape, 0, dtype)` | -| `zeros_like` | `zeros_like(x)` | Zeros with same shape/dtype as `x` | -| `cat` | `cat(x, y, can_reorder=False)` | Concatenate along dim 0; `can_reorder` allows compiler flexibility | -| `cast` | `cast(x, dtype, fp_downcast_rounding="rtne")` | Type conversion; rounding: `"rtne"` (default) or `"rtz"` | - -## Memory Operations (Pointer-based) - -### `tl.load` -- most-used memory op - -```python -load(pointer, mask=None, other=None, boundary_check=(), - padding_option='', cache_modifier='', eviction_policy='', volatile=False) -``` - -**Key semantics:** - -- `mask`: block of `int1`. Where False, returns `other` (default 0). Required for out-of-bounds safety. -- `other`: fallback value where mask is False. Must match dtype. -- `boundary_check`: tuple of dims for block-pointer bounds checking (mutually exclusive with `mask`). -- `padding_option`: `"zero"` or `"nan"` (only with `boundary_check`). -- `cache_modifier`: `""`, `".cg"`, `".cs"`, `".ca"`, `".wb"`, `".wt"`. -- `eviction_policy`: `""`, `"evict_first"`, `"evict_last"`. - -```python -# Typical masked load pattern -offs = pid * BLOCK + tl.arange(0, BLOCK) -mask = offs < n_elements -x = tl.load(ptr + offs, mask=mask, other=0.0) -``` - -### `tl.store` - -```python -store(pointer, value, mask=None, boundary_check=(), - cache_modifier='', eviction_policy='') -``` - -Same mask semantics as load. Where mask is False, store is skipped (no side effect). - -```python -tl.store(out_ptr + offs, result, mask=mask) -``` - -## Memory Operations (Block Pointer) - -| Function | Signature | Notes | -|---|---|---| -| `make_block_ptr` | `(base, shape, strides, offsets, block_shape, order)` | Structured pointer; `order` controls memory layout (e.g., `(1,0)` for col-major) | -| `advance` | `advance(block_ptr, offsets)` | Returns NEW ptr (no mutation); `offsets` is tuple by dim | - -```python -a_ptr = tl.make_block_ptr(a, (M, K), (stride_am, stride_ak), (pid_m * BM, 0), (BM, BK), order=(1, 0)) -a_ptr = tl.advance(a_ptr, (0, BK)) # advance K dimension -a = tl.load(a_ptr, boundary_check=(0, 1)) -``` - -## Memory Operations (Tensor Descriptor / TMA) - -```python -make_tensor_descriptor(base, shape, strides, block_shape, padding_option="zero") -# base must be 16-byte aligned. Supports 2-5D tensors. -# On NVIDIA GPUs with TMA, uses hardware TMA descriptor. -``` - -| Function | Signature | Notes | -|---|---|---| -| `tensor_descriptor.load` | `.load(offsets, boundary_check=True)` | Load block at `offsets` from descriptor | -| `tensor_descriptor.store` | `.store(offsets, value)` | Store block at `offsets` | - -## Linear Algebra - -### `tl.dot` - -```python -dot(input, other, acc=None, input_precision="tf32", max_num_imprecise_acc=None, out_dtype=float32) -``` - -- Both operands must be 2-D or 3-D (batched matmul). Inner dims must match (min 16). -- `input` dtype: int8, float8_e5m2, float8_e4m3fn, float16, bfloat16, float32. -- `input_precision`: `"tf32"` (default, NVIDIA), `"tf32x3"`, `"ieee"`. -- `acc`: accumulator tensor; if provided, result is added to it. - -### `tl.dot_scaled` (Microscaling / MX formats) - -```python -dot_scaled(lhs, lhs_scale, lhs_format, rhs, rhs_scale, rhs_format, - acc=None, out_dtype=float32) -``` - -- Formats: `"e2m1"`, `"e4m3"`, `"e5m2"`, `"bf16"`, `"fp16"`. -- Scales are e8m0 (uint8 tensors), shape `[M, K//group_size]`. - -## Math Operations - -| Function | Signature | Notes | -|---|---|---| -| `abs` | `abs(x)` | Elementwise absolute value | -| `cdiv` | `cdiv(x, div)` | Ceiling division: `(x + div - 1) // div` | -| `ceil` | `ceil(x)` | Ceiling (float) | -| `floor` | `floor(x)` | Floor (float) | -| `exp` | `exp(x)` | Base-e exponential | -| `exp2` | `exp2(x)` | Base-2 exponential | -| `log` | `log(x)` | Natural logarithm | -| `log2` | `log2(x)` | Base-2 logarithm | -| `cos` | `cos(x)` | Cosine | -| `sin` | `sin(x)` | Sine | -| `sqrt` | `sqrt(x)` | Square root | -| `rsqrt` | `rsqrt(x)` | Reciprocal square root: `1/sqrt(x)` | -| `sigmoid` | `sigmoid(x)` | `1 / (1 + exp(-x))` | -| `softmax` | `softmax(x, axis)` | Numerically-stable softmax along `axis` | -| `umulhi` | `umulhi(x, y)` | Upper 32 bits of `x * y` (uint32) | -| `fdiv` | `fdiv(x, y, ieee_rounding=False)` | Floating-point division | -| `fma` | `fma(x, y, z)` | Fused multiply-add: `x * y + z` | -| `clamp` | `clamp(x, min, max)` | Clamp to range `[min, max]` | -| `minimum` | `minimum(x, y)` | Elementwise min (propagates NaN) | -| `maximum` | `maximum(x, y)` | Elementwise max (propagates NaN) | - -## Where (Critical for Masking) - -```python -where(condition, x, y) -``` - -Returns elements from `x` where `condition` is True, else from `y`. Both `x` and `y` are broadcast to `condition`'s shape. This is the primary tool for conditional logic in Triton. - -```python -# Causal mask in attention -mask = offs_m[:, None] >= offs_n[None, :] -attn = tl.where(mask, attn, float("-inf")) -``` - -## Reduction Operations - -All reductions: `fn(input, axis=None, keep_dims=False)`. When `axis=None`, reduces all dims. - -| Function | Signature | Notes | -|---|---|---| -| `max` | `max(input, axis, keep_dims=False)` | Maximum along axis | -| `min` | `min(input, axis, keep_dims=False)` | Minimum along axis | -| `argmax` | `argmax(input, axis)` | Index of max along axis | -| `argmin` | `argmin(input, axis)` | Index of min along axis | -| `sum` | `sum(input, axis, keep_dims=False, dtype=None)` | Sum; int/bool auto-upcast to int32, float to float32 | -| `xor_sum` | `xor_sum(input, axis)` | XOR reduction along axis | -| `reduce` | `reduce(input, axis, combine_fn, keep_dims=False)` | Generic reduction with user-defined `combine_fn(a, b) -> c` | - -```python -# Reduction pattern: online softmax -row_max = tl.max(row, axis=1, keep_dims=True) -row = tl.exp(row - row_max) -row_sum = tl.sum(row, axis=1, keep_dims=True) -``` - -## Scan and Sort Operations - -| Function | Signature | Notes | -|---|---|---| -| `associative_scan` | `associative_scan(input, axis, combine_fn, reverse=False)` | Prefix scan with user-defined associative `combine_fn` | -| `cumsum` | `cumsum(input, axis, dtype=None)` | Cumulative sum (specialization of scan) | -| `cumprod` | `cumprod(input, axis, dtype=None)` | Cumulative product | -| `histogram` | `histogram(input, num_bins)` | Counts per bin; input values are bin indices | -| `sort` | `sort(input, axis=-1, descending=False, stable=True)` | Sort along axis | -| `topk` | `topk(input, k, axis=-1, descending=True)` | Top-k values along axis | -| `gather` | `gather(input, indices, axis)` | Gather elements along axis using indices | - -## Atomic Operations - -All atomics: `atomic_*(pointer, val, mask=None, sem="acq_rel", scope="gpu")`. - -- `sem`: `"acquire"`, `"release"`, `"acq_rel"` (default), `"relaxed"`. -- `scope`: `"gpu"` (default), `"cta"` (thread block), `"sys"` (system). - -| Function | Signature | Notes | -|---|---|---| -| `atomic_add` | `(ptr, val, mask=None, sem, scope)` | Atomic add; returns old value | -| `atomic_max` | `(ptr, val, mask=None, sem, scope)` | Atomic max; returns old value | -| `atomic_min` | `(ptr, val, mask=None, sem, scope)` | Atomic min; returns old value | -| `atomic_and` | `(ptr, val, mask=None, sem, scope)` | Atomic bitwise AND | -| `atomic_or` | `(ptr, val, mask=None, sem, scope)` | Atomic bitwise OR | -| `atomic_xor` | `(ptr, val, mask=None, sem, scope)` | Atomic bitwise XOR | -| `atomic_xchg` | `(ptr, val, mask=None, sem, scope)` | Atomic exchange; returns old value | -| `atomic_cas` | `(ptr, cmp, val, sem, scope)` | Compare-and-swap: if `*ptr == cmp`, set to `val`; returns old value | - -## Random Number Generation (Philox PRNG) - -| Function | Signature | Notes | -|---|---|---| -| `randint4x` | `randint4x(seed, offset)` | Returns 4 blocks of int32; fastest for multiple streams | -| `randint` | `randint(seed, offset, n_rounds=6)` | Single block of random int32 | -| `rand` | `rand(seed, offset, n_rounds=6)` | Uniform float32 in `[0, 1)` | -| `randn` | `randn(seed, offset, n_rounds=6)` | Normal distribution (float32) | - -`seed`: scalar int32. `offset`: block of int32 (determines which element gets which random value). - -## Iterators - -| Function | Signature | Notes | -|---|---|---| -| `range` | `range(start, stop, step=1)` | Dynamic loop; bounds can be runtime values | -| `static_range` | `static_range(start, stop, step=1)` | Fully unrolled at compile time; bounds must be `constexpr` | - -## Debug Operations - -| Function | Signature | Notes | -|---|---|---| -| `static_print` | `static_print(*args)` | Print at **compile time**; same interface as Python `print` | -| `static_assert` | `static_assert(cond, msg="")` | Assert at **compile time** | -| `device_print` | `device_print(prefix, *args)` | Print at **runtime** on device; first arg must be string, rest are scalars/tensors | -| `device_assert` | `device_assert(cond, msg="")` | Assert at **runtime**; requires `TRITON_DEBUG=1` env var | - -## Compiler Hints - -| Function | Signature | Notes | -|---|---|---| -| `assume` | `assume(cond)` | Hint to backend for address calculation optimization | -| `max_contiguous` | `max_contiguous(input, values)` | Declare max contiguous extent per dim; enables coalesced access | -| `max_constancy` | `max_constancy(input, values)` | Declare max constant extent per dim | -| `multiple_of` | `multiple_of(input, values)` | Declare that values are multiples of given constants | -| `debug_barrier` | `debug_barrier()` | Thread barrier (debugging only; not for correctness) | - -## Shape Manipulation - -| Function | Signature | Notes | -|---|---|---| -| `broadcast` | `broadcast(x, y)` | Broadcast `x` and `y` to compatible shape (returns both) | -| `broadcast_to` | `broadcast_to(x, shape)` | Broadcast `x` to explicit `shape` | -| `expand_dims` | `expand_dims(x, axis)` | Insert length-1 dim at `axis` | -| `reshape` | `reshape(x, shape)` | Reshape (total elements must match) | -| `view` | `view(x, shape)` | Like reshape; bitcast semantics | -| `trans` | `trans(x, *dims)` | Transpose; default swaps last two dims | -| `permute` | `permute(x, *dims)` | Reorder dims; `permute(x, 2, 1, 0)` or `permute(x, (2,1,0))` | -| `ravel` | `ravel(x)` | Flatten to 1-D | -| `split` | `split(x)` | Split first dim into separate tensors | -| `join` | `join(x, y)` | Concatenate along a new innermost dim | -| `interleave` | `interleave(x, y)` | Interleave elements from `x` and `y` | - -## Inline Assembly - -```python -inline_asm_elementwise(asm, constraints, args, dtype, is_pure, pack) -``` - -- `asm`: PTX/ASM string with `$0`, `$1`, ... placeholders. -- `constraints`: register constraint string (e.g., `"=r,r"` for one int output, one int input). -- `args`: list of input tensors. -- `dtype`: output dtype (or tuple for multi-output). -- `is_pure`: True if no side effects (enables CSE). -- `pack`: number of elements per register (typically 1). diff --git a/.agents/skills/kernel-triton-writing/references/concepts-semantics.md b/.agents/skills/kernel-triton-writing/references/concepts-semantics.md deleted file mode 100644 index c74081b45b2a..000000000000 --- a/.agents/skills/kernel-triton-writing/references/concepts-semantics.md +++ /dev/null @@ -1,199 +0,0 @@ - - - - -# Triton Concepts and Semantics - -## Programming Model — Block-Based Execution - -Triton programs operate on **blocks** (tiles) of data, not individual scalar threads. -Each kernel instance (called a "program") processes an entire block of elements at once. - -| Concept | CUDA | Triton | -|---------|------|--------| -| Execution unit | Single scalar thread | Program operating on a block | -| Memory coalescing | Manual (stride patterns) | Automatic (compiler) | -| Shared memory | Manual (`__shared__`, sync) | Automatic (compiler) | -| Vectorization | Manual (float4, etc.) | Automatic (compiler) | -| Tensor core usage | Manual (wmma/mma) | Automatic (compiler) | -| Thread synchronization | Manual (`__syncthreads`) | Not needed | - -### Launch Grid and Program IDs - -```python -@triton.jit -def kernel(X_ptr, Y_ptr, N, BLOCK_SIZE: tl.constexpr): - # Each program instance gets a unique ID along each grid axis - pid = tl.program_id(axis=0) # which block of data this program handles - offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - mask = offsets < N - x = tl.load(X_ptr + offsets, mask=mask) - tl.store(Y_ptr + offsets, x * 2, mask=mask) - -# Launch with a 1D grid: one program per block of data -grid = lambda meta: (triton.cdiv(N, meta['BLOCK_SIZE']),) -kernel[grid](x_ptr, y_ptr, N, BLOCK_SIZE=1024) -``` - -### Key Takeaway - -The programmer thinks in blocks. The compiler decides how to map blocks to warps, -how to stage data through shared memory, and when to use tensor cores. This is the -core design tradeoff: less control, but far less boilerplate and fewer correctness bugs. - ---- - -## Type Promotion Rules - -Triton applies automatic type promotion for binary ops and `tl.where` (last two args). - -### Promotion Hierarchy - -``` -{bool} < {int8, int16, int32, int64, uint8, uint16, uint32, uint64} < {fp8, fp16, bf16, fp32, fp64} - ^ ^ (integral types) ^ (floating types) - kind 0 kind 1 kind 2 -``` - -### Rules Applied in Order - -| Priority | Rule | Example | -|----------|------|---------| -| 1 | **Cross-kind**: lower kind promotes to higher kind's dtype | `(int32, bf16)` -> `bf16` | -| 2 | **Same-kind widening**: narrower promotes to wider | `(fp16, fp32)` -> `fp32` | -| 3 | **Same-width float tie**: bf16 and fp16 both promote to fp16 | `(fp16, bf16)` -> `fp16` | -| 4 | **Same-width sign tie**: promote to unsigned | `(int32, uint32)` -> `uint32` | - -### Scalar-Tensor Interaction - -When a Python scalar interacts with a Triton tensor: - -| Scalar Type | Tensor Type | Result | -|-------------|-------------|--------| -| Python `int` | Any int tensor | Tensor's dtype (no widening) | -| Python `int` | Any float tensor | Tensor's dtype | -| Python `float` | Any float tensor | Tensor's dtype | -| Python `float` | Any int tensor | `fp64` (float is higher kind) | - -### Gotchas: Type Promotion - -| Gotcha | Detail | -|--------|--------| -| `int32 + bf16` -> `bf16` | Integer silently truncated to bf16 precision (only ~3 decimal digits) | -| `fp16 + bf16` -> `fp16` | bf16 promotes to fp16, NOT fp32. May lose bf16 range | -| `int8 + uint8` -> `uint8` | Signed values become unsigned, wrapping negative values | -| Cross-kind hides widening | `(int64, fp16)` -> `fp16`, losing 64-bit integer precision | -| No implicit fp32 promotion | Unlike PyTorch, Triton does NOT auto-promote fp16/bf16 to fp32 for accumulation | - ---- - -## Broadcasting Rules - -Triton broadcasting follows NumPy conventions with one key constraint: -tensors are at most 2D in practice (block pointers may extend this). - -### Rules - -1. **Left-pad with ones**: If tensors have different numbers of dimensions, the - shorter shape is padded on the left with 1s. -2. **Dimension-1 expansion**: Dimensions of size 1 are stretched to match the - corresponding dimension of the other tensor. -3. **Incompatible = error**: If dimensions differ and neither is 1, it is a compile error. - -### Example: Row-Column Broadcast - -```python -# Create a row vector (1, N) and column vector (M, 1) -row = tl.arange(0, N)[None, :] # shape: (1, N) -col = tl.arange(0, M)[:, None] # shape: (M, 1) - -# Broadcast produces (M, N) — outer product pattern -result = row + col # shape: (M, N) -``` - -### Common Broadcasting Patterns - -| Pattern | Shape A | Shape B | Result Shape | Use Case | -|---------|---------|---------|-------------|----------| -| Row + Col | `(1, N)` | `(M, 1)` | `(M, N)` | 2D index grids, outer products | -| Scalar + Block | `()` | `(M, N)` | `(M, N)` | Add bias, scale | -| Row mask | `(1, N)` | `(M, N)` | `(M, N)` | Column-wise masking | - -### Gotcha: Broadcasting - -| Gotcha | Detail | -|--------|--------| -| No implicit unsqueeze | You must explicitly reshape with `[:, None]` or `[None, :]` | -| 1D + 1D does NOT broadcast | Two 1D tensors of different length are an error, not broadcast | -| Mask must broadcast to data | `tl.load(ptr, mask=mask)` — mask shape must broadcast to ptr block shape | - ---- - -## Integer Division and Modulus — C Semantics - -**CRITICAL**: Triton uses **C semantics** (round toward zero), NOT Python semantics -(round toward negative infinity). This is the most common source of subtle bugs -when porting Python logic to Triton kernels. - -### Comparison Table - -| Expression | Python Result | Triton Result | Why | -|------------|--------------|---------------|-----| -| `-7 // 2` | `-4` | `-3` | Python: floor division. Triton/C: truncation toward zero | -| `-7 % 2` | `1` | `-1` | Follows from division: `a == (a // b) * b + (a % b)` | -| `7 // -2` | `-4` | `-3` | Same: truncation vs floor | -| `7 % -2` | `-1` | `1` | Remainder keeps dividend sign in C | -| `-7 // -2` | `3` | `3` | Both agree when signs match (positive quotient) | - -### The Identity - -Both C and Python satisfy: `a == (a // b) * b + (a % b)` - -But they disagree on which direction to round the quotient, which changes the remainder. - -### Exception: Scalar-Only Computations - -When **all inputs are Python scalars** (not Triton tensors), division and modulus -follow **Python semantics**. This only applies to compile-time constant folding. - -```python -@triton.jit -def kernel(X_ptr, N, BLOCK: tl.constexpr): - # Python semantics — both are Python scalars at compile time - blocks_per_row = (-7) // 2 # = -4 (Python floor division) - - # C semantics — pid is a Triton value - pid = tl.program_id(0) - row = pid // N # truncation toward zero - col = pid % N # C remainder -``` - -### Gotcha: Safe Patterns for Negative Values - -| Unsafe Pattern | Problem | Safe Alternative | -|----------------|---------|------------------| -| `(-offset) // stride` | C truncation gives wrong block | `-(offset // stride)` or use unsigned | -| `idx % BLOCK` for negative idx | Negative remainder | Ensure idx is non-negative, or add `+ BLOCK) % BLOCK` | -| Porting Python `divmod` logic | Both `//` and `%` differ | Rewrite with explicit floor: `q = (a - (a % b + b) % b) // b` | - -### When It Matters - -This only causes bugs when **operands can be negative**. If all values are -non-negative (which is common for pointer offsets and indices), C and Python -semantics agree. Guard against negative values explicitly when in doubt. diff --git a/.agents/skills/kernel-triton-writing/references/operator-routing.md b/.agents/skills/kernel-triton-writing/references/operator-routing.md deleted file mode 100644 index 9bfd3d3b7c13..000000000000 --- a/.agents/skills/kernel-triton-writing/references/operator-routing.md +++ /dev/null @@ -1,125 +0,0 @@ - - - - -# Operator Routing Decision Reference - -Detailed decision rules for determining whether an operator should be implemented -as a custom Triton kernel or handled by existing libraries. - -## Decision Procedure - -Follow these rules in order. Stop at the first match. - -1. **Single element-wise op** (e.g., `relu(x)`, `sigmoid(x)`) -- SKIP. PyTorch - already optimal, no fusion benefit. -2. **Standalone matmul** (e.g., `torch.matmul(a, b)`) -- SKIP. cuBLAS is highly - optimized and hard to beat. -3. **Standard attention** (e.g., `F.scaled_dot_product_attention`) -- SKIP. Use - FlashAttention. -4. **Element-wise chain (2+ ops)** (e.g., `gelu(dropout(x))`, `silu(x) * y`) -- - USE TRITON. Fuse memory-bound ops into compute-bound kernel. -5. **Reduction op** (e.g., LayerNorm, RMSNorm, Softmax) -- USE TRITON. Custom - single-pass implementation beats generic PyTorch decomposition. -6. **Matmul + element-wise epilogue** (e.g., `matmul(a, b) + bias`, - `matmul + gelu`) -- USE TRITON. Epilogue fusion avoids memory round-trip. -7. **Matmul + reduction** (e.g., `matmul -> softmax`, `matmul -> layernorm`) -- - USE TRITON. Common transformer pattern with clear fusion benefit. -8. **Custom attention variant** -- Check FlashAttention support first. Only use - Triton if the variant is unsupported. -9. **Sparse operations** -- Triton can help, but evaluate specialized libraries - (cuSPARSE, Triton block-sparse) first. -10. **Very small tensors** -- Launch overhead may dominate. Benchmark before - committing. -11. **Default** -- Analyze operator code and shapes, then decide. - -## Output Format - -Report the routing decision as: - -```markdown -## Routing Decision: [OPERATOR_NAME] - -**Decision:** USE TRITON | SKIP TRITON | EVALUATE FURTHER - -**Pattern:** [e.g., Element-wise chain, Reduction, Matmul+epilogue] - -**Rationale:** [Why -- reference fusion benefit or lack thereof] - -**Next Steps:** -- [USE TRITON] Proceed to Phase 1 (Analyze the Operator) -- [SKIP] Recommend alternative (cuBLAS, FlashAttention, PyTorch) -- [EVALUATE] Profile operator, analyze shapes, then re-decide -``` - -## Examples - -### Fused GELU + Dropout - -```python -def fused_op(x, p=0.1): - return F.dropout(F.gelu(x), p=p) -``` - -**Decision:** USE TRITON | **Pattern:** Element-wise chain (2 ops) -Fusing eliminates one intermediate tensor write+read (~2x memory traffic reduction). - -### Simple ReLU - -```python -def simple_relu(x): - return F.relu(x) -``` - -**Decision:** SKIP TRITON | **Pattern:** Single element-wise op -No fusion benefit. PyTorch ReLU is already a single memory-bound kernel. - -### RMSNorm - -```python -def rmsnorm(x, weight, eps=1e-6): - rms = torch.sqrt(torch.mean(x ** 2, dim=-1, keepdim=True) + eps) - return x / rms * weight -``` - -**Decision:** USE TRITON | **Pattern:** Reduction op -Triton fuses square, mean, sqrt, divide, multiply in a single pass over the data. - -### Linear + GELU - -```python -def linear_gelu(x, weight, bias): - return F.gelu(F.linear(x, weight, bias)) -``` - -**Decision:** USE TRITON | **Pattern:** Matmul + element-wise epilogue -Fusing GELU into the matmul epilogue avoids an extra full tensor read+write. - -## Edge Cases - -- **Dynamic shapes or data-dependent branching** -- Triton requires static grid - dimensions at launch. If shapes change per-sample, fall back to PyTorch eager - or `torch.compile`. -- **Operators already in `torch.compile` fusion groups** -- Check whether - `torch.compile` already fuses the pattern before writing a manual kernel. - A manual Triton kernel is only justified if it measurably outperforms the - compiler-generated version. -- **Mixed precision boundaries** -- Triton handles dtype casting well, but verify - that the fused kernel preserves numerical behavior (especially around - loss scaling and FP16/BF16 reductions). diff --git a/.agents/skills/kernel-triton-writing/references/patterns-advanced.md b/.agents/skills/kernel-triton-writing/references/patterns-advanced.md deleted file mode 100644 index 635f10897c3c..000000000000 --- a/.agents/skills/kernel-triton-writing/references/patterns-advanced.md +++ /dev/null @@ -1,320 +0,0 @@ - - - - -# Advanced Triton Patterns - -Source tutorials: - -- [05-layer-norm](https://triton-lang.org/main/getting-started/tutorials/05-layer-norm.html) -- [06-fused-attention](https://triton-lang.org/main/getting-started/tutorials/06-fused-attention.html) -- [07-extern-functions](https://triton-lang.org/main/getting-started/tutorials/07-extern-functions.html) - -## Layer Normalization - -Layer norm normalizes across the hidden dimension: `y = (x - mean) / sqrt(var + eps) * w + b`. -Each program instance processes one row of the input (one token). The hidden dimension -is tiled into blocks so arbitrary sizes are supported. - -### Forward Kernel (Complete) - -```python -@triton.jit -def _layer_norm_fwd_fused( - X, # input pointer, shape (M, N) - Y, # output pointer, shape (M, N) - W, # weight pointer, shape (N,) - B, # bias pointer, shape (N,) - Mean, # mean pointer, shape (M,) — written for backward - Rstd, # rstd pointer, shape (M,) — written for backward - stride, # row stride of X and Y - N, # number of columns (hidden size) - eps: tl.constexpr, - BLOCK_SIZE: tl.constexpr, -): - # Each program handles one row - row = tl.program_id(0) - Y += row * stride - X += row * stride - - # --- Compute mean --- - _mean = tl.zeros([BLOCK_SIZE], dtype=tl.float32) - for off in range(0, N, BLOCK_SIZE): - cols = off + tl.arange(0, BLOCK_SIZE) - a = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) - _mean += a - mean = tl.sum(_mean, axis=0) / N - - # --- Compute variance --- - _var = tl.zeros([BLOCK_SIZE], dtype=tl.float32) - for off in range(0, N, BLOCK_SIZE): - cols = off + tl.arange(0, BLOCK_SIZE) - x = tl.load(X + cols, mask=cols < N, other=0.0).to(tl.float32) - x = tl.where(cols < N, x - mean, 0.0) - _var += x * x - var = tl.sum(_var, axis=0) / N - rstd = 1 / tl.sqrt(var + eps) - - # Store mean and rstd for backward - tl.store(Mean + row, mean) - tl.store(Rstd + row, rstd) - - # --- Normalize and apply affine transform --- - for off in range(0, N, BLOCK_SIZE): - cols = off + tl.arange(0, BLOCK_SIZE) - mask = cols < N - w = tl.load(W + cols, mask=mask) - b = tl.load(B + cols, mask=mask) - x = tl.load(X + cols, mask=mask, other=0.0).to(tl.float32) - x_hat = (x - mean) * rstd - y = x_hat * w + b - tl.store(Y + cols, y, mask=mask) -``` - -**Key pattern:** Two-pass reduction (mean then variance) using block-wise tiling. -Each loop iteration processes `BLOCK_SIZE` elements with masking for the tail. -Mean and rstd are saved for the backward pass. - -### Backward Kernel: Atomic Lock Pattern - -The backward pass computes `dw` and `db` which require reducing across all rows -(all programs contribute). Triton uses a spin-lock pattern with `atomic_cas` for -mutual exclusion: - -```python -@triton.jit -def _layer_norm_bwd_dwdb( - DW, DB, # output accumulators, shape (N,) - DWEIGHT, DBIAS, # partial sums buffer, shape (GROUP_SIZE_M, N) - Lock, # lock array, shape (1,) — int32 - ... - GROUP_SIZE_M: tl.constexpr, -): - row_block_id = tl.program_id(0) - # Each group of rows accumulates partials then atomically adds to DW/DB - - # --- Compute partial dw, db for this row group --- - # (loop over assigned rows, accumulate _dw and _db) - - # --- Acquire lock --- - lock_id = tl.program_id(1) # column block index - Lock += lock_id - Count = Lock + tl.num_programs(1) # second half stores count - - while tl.atomic_cas(Lock, 0, 1) == 1: # spin until we get 0->1 - pass - count = tl.load(Count) # how many groups have accumulated so far - - if count == 0: - # First group: just store - tl.store(DWEIGHT + cols, _dw, mask=mask) - tl.store(DBIAS + cols, _db, mask=mask) - else: - # Subsequent groups: accumulate - _dw += tl.load(DWEIGHT + cols, mask=mask) - _db += tl.load(DBIAS + cols, mask=mask) - tl.store(DWEIGHT + cols, _dw, mask=mask) - tl.store(DBIAS + cols, _db, mask=mask) - - if count == GROUP_SIZE_M - 1: - # Last group: write final result - tl.store(DW + cols, _dw, mask=mask) - tl.store(DB + cols, _db, mask=mask) - - # --- Release lock and increment count --- - tl.atomic_xchg(Lock, 0) # release: set lock back to 0 - tl.store(Count, count + 1) # must store AFTER release for correctness - tl.debug_barrier() # ensure memory operations are visible -``` - -**Gotchas:** - -- `atomic_cas(Lock, 0, 1)` returns the old value; spin while it returns 1 (already held). -- `atomic_xchg(Lock, 0)` unconditionally sets to 0 (release). Do NOT use `atomic_cas` for release. -- The count update (`tl.store(Count, count + 1)`) must happen after the lock is released. -- `tl.debug_barrier()` forces memory ordering visibility across programs. -- Lock array must be zero-initialized before each backward call. -- This pattern is needed because Triton has no native cross-program reduction for non-atomic dtypes. - -## Fused Attention - -Implements Flash Attention v2: fused Q*K^T softmax and V accumulation in a single -kernel, avoiding materializing the full N x N attention matrix. - -### Online Softmax Algorithm - -The key insight is computing softmax in a single streaming pass using running -statistics. For each block of K/V columns processed: - -``` -# For each new block j of keys: -qk = Q_block @ K_block_j^T # [BLOCK_M, BLOCK_N] -m_ij = max(qk, axis=1) # new block max -m_i_new = max(m_i, m_ij) # update running max -alpha = exp(m_i - m_i_new) # correction factor for old accumulators -p = exp(qk - m_i_new[:, None]) # stable softmax numerator -l_i = alpha * l_i + sum(p, axis=1) # update running denominator -acc = alpha[:, None] * acc + p @ V_j # rescale old acc + new contribution -m_i = m_i_new # commit new max -# After all blocks: -acc = acc / l_i[:, None] # final normalization -``` - -**Why this works:** Each time a new block raises the running max, all previous -accumulations are rescaled by `exp(old_max - new_max)`, maintaining numerical -equivalence to the two-pass softmax. - -### Multi-Stage Processing (Causal Masking) - -The STAGE parameter controls masking behavior in the inner loop: - -| STAGE | Behavior | When used | -|-------|----------|-----------| -| 1 | Off-band: skip blocks entirely below diagonal | Causal, early blocks | -| 3 | On-band: apply causal mask within block | Causal, diagonal blocks | -| 2 | No masking | Non-causal attention | - -```python -# Causal masking within a block (STAGE == 3): -if STAGE == 3: - # Current query rows: [start_m, start_m + BLOCK_M) - # Current key cols: [start_n, start_n + BLOCK_N) - offs_m = start_m + tl.arange(0, BLOCK_M) - offs_n = start_n + tl.arange(0, BLOCK_N) - causal_mask = offs_m[:, None] >= offs_n[None, :] - qk = tl.where(causal_mask, qk, float("-inf")) -``` - -**Gotcha:** When `STAGE == 1`, blocks where all keys are below the diagonal are -skipped entirely (the inner loop `start_n` begins past the diagonal). This is a -major performance win for causal attention on long sequences. - -### Kernel Structure Skeleton - -```python -@triton.jit -def _attn_fwd( - Q, K, V, sm_scale, - M, # log-sum-exp for backward, shape (batch, nheads, seqlen) - Out, - stride_qz, stride_qh, stride_qm, stride_qk, # Q strides - # ... K, V, Out strides ... - Z, H, N_CTX, - BLOCK_M: tl.constexpr, # query block size (e.g. 128) - BLOCK_N: tl.constexpr, # key block size (e.g. 64) - HEAD_DIM: tl.constexpr, # head dimension (e.g. 64) - STAGE: tl.constexpr, -): - start_m = tl.program_id(0) # which query block - off_hz = tl.program_id(1) # batch * head index - - # Initialize pointers for Q[start_m], K, V - # Load Q block into registers (stays resident) - q = tl.load(Q_block_ptr) # [BLOCK_M, HEAD_DIM] - - # Accumulator in float32 - acc = tl.zeros([BLOCK_M, HEAD_DIM], dtype=tl.float32) - m_i = tl.zeros([BLOCK_M], dtype=tl.float32) - float("inf") - l_i = tl.zeros([BLOCK_M], dtype=tl.float32) + 1.0 - - # --- Inner loop over K/V blocks --- - for start_n in range(lo, hi, BLOCK_N): - k = tl.load(K_block_ptr) # [BLOCK_N, HEAD_DIM] - v = tl.load(V_block_ptr) # [BLOCK_N, HEAD_DIM] - - qk = tl.dot(q, tl.trans(k)) * sm_scale # [BLOCK_M, BLOCK_N] - - # Apply causal mask if STAGE == 3 - # Online softmax update (m_i, l_i, acc) as shown above - - K_block_ptr = tl.advance(K_block_ptr, (BLOCK_N, 0)) - V_block_ptr = tl.advance(V_block_ptr, (BLOCK_N, 0)) - - # Final normalization - acc = acc / l_i[:, None] - # Store lse = m_i + log(l_i) for backward - tl.store(Out_block_ptr, acc.to(Out.type.element_ty)) -``` - -### TensorDescriptor (Hopper+) - -On Hopper/Blackwell, `TensorDescriptor` enables TMA (Tensor Memory Accelerator): - -```python -desc_q = TensorDescriptor(Q_ptr, shape=[N_CTX, HEAD_DIM], - strides=[stride_qm, stride_qk], - block_shape=[BLOCK_M, HEAD_DIM]) -q = desc_q.load([start_m * BLOCK_M, 0]) # replaces pointer arithmetic -``` - -### Warp Specialization and Performance - -On Blackwell (sm_100+), warp specialization lets different warp groups play -producer/consumer roles, overlapping loads with compute via `tl.async_task`. -Flash Attention in Triton reaches ~165 TFLOPS (fp16, H100). Causal masking -with the STAGE optimization roughly halves unnecessary work. - -## External Functions - -Triton can call functions from external device libraries (libdevice for CUDA, -ROCm device libs for HIP) for math operations not built into the language. - -### Basic Usage - -```python -@triton.jit -def asin_kernel( - x_ptr, y_ptr, - n_elements, - BLOCK_SIZE: tl.constexpr, -): - pid = tl.program_id(axis=0) - offset = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - mask = offset < n_elements - x = tl.load(x_ptr + offset, mask=mask) - # Call libdevice asin — dispatches based on dtype (fp32 or fp64) - y = tl.extra.cuda.libdevice.asin(x) - tl.store(y_ptr + offset, y, mask=mask) -``` - -Type dispatch is automatic: `libdevice.asin` calls `__nv_asinf` for float32 -and `__nv_asin` for float64 under the hood. - -### Custom Library Paths - -Pass external libraries explicitly via `extern_libs` at compile time: - -```python -grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),) -asin_kernel[grid](x, y, n_elements, BLOCK_SIZE=1024, - extern_libs={"libdevice": "/path/to/libdevice.10.bc"}) -``` - -### Backend Detection - -| Backend | Library file | Namespace | -|---------|-------------|-----------| -| CUDA | `libdevice.10.bc` | `tl.extra.cuda.libdevice.*` | -| HIP | `ocml.bc` / `ockl.bc` | Functions mapped through HIP backend | - -Common functions: `asin`, `acos`, `atan`, `exp`, `log`, `pow`, `sqrt`, `rsqrt`, -`fma`, `cbrt`, `erf`, `erfc`, `ceil`, `floor`, `round`. - -**Gotcha:** `extern_libs` must point to the `.bc` bitcode file (typically -`/usr/local/cuda/nvvm/libdevice/libdevice.10.bc`). Missing file = compile-time linker error. diff --git a/.agents/skills/kernel-triton-writing/references/patterns-basic.md b/.agents/skills/kernel-triton-writing/references/patterns-basic.md deleted file mode 100644 index 9435f2f90aa2..000000000000 --- a/.agents/skills/kernel-triton-writing/references/patterns-basic.md +++ /dev/null @@ -1,238 +0,0 @@ - - - - -# Triton Basic Kernel Patterns - -Reusable patterns extracted from the official Triton tutorials. -Each section contains a complete kernel, its launch wrapper, and annotations. - ---- - -## Vector Addition - -The simplest Triton pattern: 1D parallel map over contiguous data. - -### Kernel - -```python -import torch -import triton -import triton.language as tl - -@triton.jit -def add_kernel( - x_ptr, y_ptr, output_ptr, - n_elements, - BLOCK_SIZE: tl.constexpr, # compile-time constant: controls tile width -): - # Each program instance owns one tile of BLOCK_SIZE elements. - pid = tl.program_id(axis=0) - # Compute the start offset for this program's tile, then the per-lane offsets. - offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - # Guard against out-of-bounds access on the final, possibly partial tile. - mask = offsets < n_elements - # Load inputs from DRAM — masked lanes get a safe default (0.0). - x = tl.load(x_ptr + offsets, mask=mask) - y = tl.load(y_ptr + offsets, mask=mask) - output = x + y - # Write result back — only masked lanes write. - tl.store(output_ptr + offsets, output, mask=mask) -``` - -### Launch Wrapper - -```python -def add(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: - output = torch.empty_like(x) - assert x.is_cuda and y.is_cuda and output.is_cuda - n_elements = output.numel() - # Grid is a callable: Triton passes meta-parameters (incl. BLOCK_SIZE) at launch. - grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) - add_kernel[grid](x, y, output, n_elements, BLOCK_SIZE=1024) - return output -``` - -### Benchmark Pattern - -```python -@triton.testing.perf_report( - triton.testing.Benchmark( - x_names=["size"], - x_vals=[2**i for i in range(12, 28, 1)], - x_log=True, - line_arg="provider", - line_vals=["triton", "torch"], - line_names=["Triton", "Torch"], - ylabel="GB/s", - plot_name="vector-add-performance", - args={}, - ) -) -def benchmark(size, provider): - x = torch.rand(size, device="cuda", dtype=torch.float32) - y = torch.rand(size, device="cuda", dtype=torch.float32) - quantiles = [0.5, 0.2, 0.8] - if provider == "torch": - ms, min_ms, max_ms = triton.testing.do_bench(lambda: x + y, quantiles=quantiles) - if provider == "triton": - ms, min_ms, max_ms = triton.testing.do_bench(lambda: add(x, y), quantiles=quantiles) - gbps = lambda ms: 3 * x.numel() * x.element_size() * 1e-9 / (ms * 1e-3) - return gbps(ms), gbps(max_ms), gbps(min_ms) - -benchmark.run(print_data=True, show_plots=True) -``` - -### Key Takeaways - -- **Offset pattern:** `pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)` — the universal 1D tiling idiom. -- **Masking:** Always guard the last tile with `offsets < n_elements`. -- **Grid as callable:** `lambda meta: (triton.cdiv(n, meta["BLOCK_SIZE"]),)` lets autotune vary BLOCK_SIZE. -- **Pointer arithmetic:** Triton pointers support `ptr + offset_tensor` for vectorized addressing. - ---- - -## Fused Softmax - -Row-wise softmax fused into a single kernel. Fusion reduces DRAM traffic from -`5*M*N + 2*M` bytes (naive: read 3x for max/exp/sum, write 2x) down to `M*N` -read + `M*N` write by keeping intermediate results in SRAM. - -### Kernel - -```python -@triton.jit -def softmax_kernel( - output_ptr, input_ptr, - input_row_stride, output_row_stride, - n_rows, n_cols, - BLOCK_SIZE: tl.constexpr, # must be >= n_cols (padded to power-of-2) - num_stages: tl.constexpr, # software pipelining depth -): - # Persistent-kernel style: each program processes multiple rows, strided. - row_start = tl.program_id(0) - row_step = tl.num_programs(0) - for row_idx in tl.range(row_start, n_rows, row_step, num_stages=num_stages): - # Compute pointers for this row. - row_start_ptr = input_ptr + row_idx * input_row_stride - col_offsets = tl.arange(0, BLOCK_SIZE) - input_ptrs = row_start_ptr + col_offsets - # Mask: BLOCK_SIZE is rounded up to power-of-2, so some lanes are OOB. - mask = col_offsets < n_cols - # Load row; OOB lanes get -inf so they don't affect max/sum. - row = tl.load(input_ptrs, mask=mask, other=-float("inf")) - # --- Numerical stability: subtract row-max before exp --- - row_minus_max = row - tl.max(row, axis=0) - numerator = tl.exp(row_minus_max) - denominator = tl.sum(numerator, axis=0) - softmax_output = numerator / denominator - # Store result. - output_row_start_ptr = output_ptr + row_idx * output_row_stride - output_ptrs = output_row_start_ptr + col_offsets - tl.store(output_ptrs, softmax_output, mask=mask) -``` - -### Launch Wrapper - -```python -def softmax(x: torch.Tensor) -> torch.Tensor: - n_rows, n_cols = x.shape - # BLOCK_SIZE must cover the full row — round up to power-of-2. - BLOCK_SIZE = triton.next_power_of_2(n_cols) - # Heuristic: use more warps for wider rows. - num_warps = 4 if BLOCK_SIZE <= 2048 else 8 - # Persistent kernel: launch fewer programs than rows for large inputs. - # Each SM can run ~4 programs concurrently (occupancy dependent). - num_stages = 4 if BLOCK_SIZE > 2048 else 2 - y = torch.empty_like(x) - # Grid: one dimension, capped by number of rows. - num_programs = min(n_rows, 1024) # cap to avoid over-subscription - softmax_kernel[(num_programs, 1, 1)]( - y, x, - x.stride(0), y.stride(0), - n_rows, n_cols, - BLOCK_SIZE=BLOCK_SIZE, - num_stages=num_stages, - num_warps=num_warps, - ) - return y -``` - -### Key Takeaways - -- **Numerical stability:** Always `row - tl.max(row, axis=0)` before `tl.exp`. -- **Power-of-2 padding:** `BLOCK_SIZE = triton.next_power_of_2(n_cols)` with `-inf` masking for OOB lanes. -- **Persistent kernel:** `tl.range(start, end, step, num_stages=...)` loops over multiple rows per - program, improving occupancy and enabling software pipelining. -- **Fusion benefit:** One kernel replaces three separate passes (max, exp/sum, divide), keeping - all intermediates in registers/SRAM instead of round-tripping through DRAM. - ---- - -## Low-Memory Dropout - -Traditional dropout stores a full-size bit mask. This pattern stores only an `int32` seed and -recomputes the mask on-the-fly via Triton's built-in PRNG. The same seed + offsets produce -identical random values, so forward and backward passes see the same mask without storing it. - -### Kernel - -```python -@triton.jit -def _seeded_dropout( - x_ptr, output_ptr, - n_elements, - p, # dropout probability (float, 0 to 1) - seed, # int32 seed — the ONLY state needed to reproduce the mask - BLOCK_SIZE: tl.constexpr, -): - pid = tl.program_id(axis=0) - block_start = pid * BLOCK_SIZE - offsets = block_start + tl.arange(0, BLOCK_SIZE) - mask = offsets < n_elements - x = tl.load(x_ptr + offsets, mask=mask) - # tl.rand: deterministic PRNG — given the same (seed, offsets), produces - # the same uniform float32 values in [0, 1). No global state needed. - random = tl.rand(seed, offsets) - x_keep = random > p - # Scale kept elements by 1/(1-p) so expected value is unchanged (inverted dropout). - # Dropped elements become 0.0. - output = tl.where(x_keep, x / (1 - p), 0.0) - tl.store(output_ptr + offsets, output, mask=mask) -``` - -### Launch Wrapper - -```python -def seeded_dropout(x: torch.Tensor, p: float, seed: int) -> torch.Tensor: - output = torch.empty_like(x) - assert x.is_contiguous() - n_elements = x.numel() - grid = lambda meta: (triton.cdiv(n_elements, meta["BLOCK_SIZE"]),) - _seeded_dropout[grid](x, output, n_elements, p, seed, BLOCK_SIZE=1024) - return output -``` - -### Key Takeaways - -- **Memory savings:** State = 1 `int32` seed, not an `(N,)` bool tensor. -- **Deterministic PRNG:** `tl.rand(seed, offsets)` is pure-functional — same inputs, same outputs. -- **Inverted dropout:** `x / (1 - p)` scales at train time so inference needs no adjustment. -- **`tl.where` pattern:** `tl.where(cond, val_true, val_false)` is the standard Triton conditional — - works element-wise on block tensors, compiles to predicated instructions (no branch divergence). diff --git a/.agents/skills/kernel-triton-writing/references/patterns-fusion.md b/.agents/skills/kernel-triton-writing/references/patterns-fusion.md deleted file mode 100644 index ac8d6db0912d..000000000000 --- a/.agents/skills/kernel-triton-writing/references/patterns-fusion.md +++ /dev/null @@ -1,349 +0,0 @@ - - - - -# Deep Learning Fusion Patterns - -Ready-to-use Triton kernel patterns for common DL operator fusions. -Each pattern includes autotune configs, kernel, wrapper, and expected speedups. - -For foundational patterns (vector add, softmax, dropout), see `patterns-basic.md`. -For LayerNorm with backward, fused attention, and extern functions, see `patterns-advanced.md`. - ---- - -## GELU + Dropout - -**Use when:** Transformer FFN layers with dropout. -**Expected speedup:** 1.8-2.2x vs separate ops. - -```python -@triton.autotune( - configs=[ - triton.Config({'BLOCK_SIZE': 1024}, num_warps=4), - triton.Config({'BLOCK_SIZE': 2048}, num_warps=8), - ], - key=['n_elements'], -) -@triton.jit -def fused_gelu_dropout_kernel( - x_ptr, out_ptr, n_elements, p, seed, - BLOCK_SIZE: tl.constexpr, -): - pid = tl.program_id(0) - offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - mask = offsets < n_elements - - x = tl.load(x_ptr + offsets, mask=mask) - - # GELU (exact): cast to fp32 for erf, then cast back - x_fp32 = x.to(tl.float32) - x_gelu = 0.5 * x_fp32 * (1.0 + tl.math.erf(x_fp32 * 0.7071067811865476)) - x = x_gelu.to(x.dtype) - - # Dropout - random = tl.rand(seed, offsets) - x = tl.where(random > p, x / (1 - p), 0.0) - - tl.store(out_ptr + offsets, x, mask=mask) - - -def fused_gelu_dropout(x: torch.Tensor, p: float = 0.1, training: bool = True) -> torch.Tensor: - if not training or p == 0.0: - return torch.nn.functional.gelu(x) - n_elements = x.numel() - out = torch.empty_like(x) - grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),) - seed = (x.data_ptr() % (2**31)) ^ n_elements - fused_gelu_dropout_kernel[grid](x, out, n_elements, p, seed) - return out -``` - ---- - -## SiLU + Multiply (SwiGLU) - -**Use when:** LLaMA-style FFN with SwiGLU activation. -**Expected speedup:** 1.5-2x vs `F.silu(gate) * x`. - -```python -@triton.autotune( - configs=[ - triton.Config({'BLOCK_SIZE': 1024}, num_warps=4), - triton.Config({'BLOCK_SIZE': 2048}, num_warps=8), - ], - key=['n_elements'], -) -@triton.jit -def fused_silu_mul_kernel( - x_ptr, gate_ptr, out_ptr, n_elements, - BLOCK_SIZE: tl.constexpr, -): - pid = tl.program_id(0) - offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - mask = offsets < n_elements - - x = tl.load(x_ptr + offsets, mask=mask) - gate = tl.load(gate_ptr + offsets, mask=mask) - - # SiLU(gate) * x = gate * sigmoid(gate) * x - silu_gate = gate * tl.sigmoid(gate) - out = silu_gate * x - - tl.store(out_ptr + offsets, out, mask=mask) - - -def fused_silu_mul(x: torch.Tensor, gate: torch.Tensor) -> torch.Tensor: - assert x.shape == gate.shape - n_elements = x.numel() - out = torch.empty_like(x) - grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),) - fused_silu_mul_kernel[grid](x, gate, out, n_elements) - return out -``` - ---- - -## Residual Add + Activation - -**Use when:** Adding residual connection with activation. -**Expected speedup:** 1.4-1.8x vs `F.gelu(x + residual)`. - -```python -@triton.autotune( - configs=[ - triton.Config({'BLOCK_SIZE': 1024}, num_warps=4), - triton.Config({'BLOCK_SIZE': 2048}, num_warps=8), - ], - key=['n_elements'], -) -@triton.jit -def fused_residual_gelu_kernel( - x_ptr, residual_ptr, out_ptr, n_elements, - BLOCK_SIZE: tl.constexpr, -): - pid = tl.program_id(0) - offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - mask = offsets < n_elements - - x = tl.load(x_ptr + offsets, mask=mask) - residual = tl.load(residual_ptr + offsets, mask=mask) - x = x + residual - - # GELU (exact) - x_fp32 = x.to(tl.float32) - x = (0.5 * x_fp32 * (1.0 + tl.math.erf(x_fp32 * 0.7071067811865476))).to(x.dtype) - - tl.store(out_ptr + offsets, x, mask=mask) - - -def fused_residual_gelu(x: torch.Tensor, residual: torch.Tensor) -> torch.Tensor: - n_elements = x.numel() - out = torch.empty_like(x) - grid = lambda meta: (triton.cdiv(n_elements, meta['BLOCK_SIZE']),) - fused_residual_gelu_kernel[grid](x, residual, out, n_elements) - return out -``` - ---- - -## RMSNorm - -**Use when:** LLaMA-style normalization (no mean subtraction). -**Expected speedup:** 1.4-2x vs naive PyTorch RMSNorm. - -```python -@triton.autotune( - configs=[ - triton.Config({'BLOCK_SIZE': 1024}, num_warps=8), - triton.Config({'BLOCK_SIZE': 2048}, num_warps=8), - triton.Config({'BLOCK_SIZE': 4096}, num_warps=16), - ], - key=['n_cols'], -) -@triton.jit -def rmsnorm_kernel( - x_ptr, out_ptr, weight_ptr, - n_rows, n_cols, eps, - BLOCK_SIZE: tl.constexpr, -): - row_idx = tl.program_id(0) - col_offsets = tl.arange(0, BLOCK_SIZE) - mask = col_offsets < n_cols - - row_start = row_idx * n_cols - x = tl.load(x_ptr + row_start + col_offsets, mask=mask, other=0.0) - - # Compute RMS - x_sq = x * x - rms = tl.sqrt(tl.sum(x_sq, axis=0) / n_cols + eps) - - # Normalize and scale - x_norm = x / rms - weight = tl.load(weight_ptr + col_offsets, mask=mask, other=1.0) - out = x_norm * weight - - tl.store(out_ptr + row_start + col_offsets, out, mask=mask) - - -def triton_rmsnorm(x: torch.Tensor, weight: torch.Tensor, eps: float = 1e-6) -> torch.Tensor: - assert x.is_contiguous() - shape = x.shape - x = x.view(-1, shape[-1]) - n_rows, n_cols = x.shape - out = torch.empty_like(x) - grid = (n_rows,) - rmsnorm_kernel[grid](x, out, weight, n_rows, n_cols, eps) - return out.view(shape) -``` - ---- - -## Linear + GELU (Matmul + Epilogue) - -**Use when:** Transformer FFN first linear with activation. -**Expected speedup:** 1.3-1.6x vs `F.gelu(F.linear(x, weight, bias))`. - -```python -@triton.autotune( - configs=[ - triton.Config({'BLOCK_M': 64, 'BLOCK_N': 64, 'BLOCK_K': 32}, num_stages=3, num_warps=4), - triton.Config({'BLOCK_M': 128, 'BLOCK_N': 64, 'BLOCK_K': 32}, num_stages=3, num_warps=4), - triton.Config({'BLOCK_M': 128, 'BLOCK_N': 128, 'BLOCK_K': 32}, num_stages=3, num_warps=8), - ], - key=['M', 'N', 'K'], -) -@triton.jit -def linear_gelu_kernel( - x_ptr, weight_ptr, bias_ptr, out_ptr, - M, N, K, - stride_xm, stride_xk, stride_wk, stride_wn, - BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, BLOCK_K: tl.constexpr, -): - pid_m = tl.program_id(0) - pid_n = tl.program_id(1) - - offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M) - offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N) - offs_k = tl.arange(0, BLOCK_K) - - acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32) - for k in range(0, K, BLOCK_K): - k_offs = k + offs_k - x_ptrs = x_ptr + offs_m[:, None] * stride_xm + k_offs[None, :] * stride_xk - x_mask = (offs_m[:, None] < M) & (k_offs[None, :] < K) - x = tl.load(x_ptrs, mask=x_mask, other=0.0) - - w_ptrs = weight_ptr + k_offs[:, None] * stride_wk + offs_n[None, :] * stride_wn - w_mask = (k_offs[:, None] < K) & (offs_n[None, :] < N) - w = tl.load(w_ptrs, mask=w_mask, other=0.0) - acc += tl.dot(x, w) - - # Add bias + fused GELU epilogue - bias = tl.load(bias_ptr + offs_n, mask=offs_n < N, other=0.0) - acc = acc + bias[None, :] - acc = 0.5 * acc * (1.0 + tl.math.erf(acc * 0.7071067811865476)) - - out_ptrs = out_ptr + offs_m[:, None] * N + offs_n[None, :] - out_mask = (offs_m[:, None] < M) & (offs_n[None, :] < N) - tl.store(out_ptrs, acc.to(tl.float16), mask=out_mask) - - -def linear_gelu(x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor) -> torch.Tensor: - assert x.is_contiguous() and weight.is_contiguous() - M, K = x.shape - K2, N = weight.shape - assert K == K2 - out = torch.empty((M, N), device=x.device, dtype=x.dtype) - grid = lambda meta: (triton.cdiv(M, meta['BLOCK_M']), triton.cdiv(N, meta['BLOCK_N'])) - linear_gelu_kernel[grid](x, weight, bias, out, M, N, K, x.stride(0), x.stride(1), weight.stride(0), weight.stride(1)) - return out -``` - ---- - -## Fused Add + LayerNorm - -**Use when:** Post-attention residual add + normalization. -**Expected speedup:** 1.5-2x vs `F.layer_norm(x + residual, ...)`. - -```python -@triton.autotune( - configs=[ - triton.Config({'BLOCK_SIZE': 1024}, num_warps=8), - triton.Config({'BLOCK_SIZE': 2048}, num_warps=8), - triton.Config({'BLOCK_SIZE': 4096}, num_warps=16), - ], - key=['n_cols'], -) -@triton.jit -def fused_add_layernorm_kernel( - x_ptr, residual_ptr, out_ptr, weight_ptr, bias_ptr, - n_rows, n_cols, eps, - BLOCK_SIZE: tl.constexpr, -): - row_idx = tl.program_id(0) - col_offsets = tl.arange(0, BLOCK_SIZE) - mask = col_offsets < n_cols - row_start = row_idx * n_cols - - # Load and add - x = tl.load(x_ptr + row_start + col_offsets, mask=mask, other=0.0) - residual = tl.load(residual_ptr + row_start + col_offsets, mask=mask, other=0.0) - x = x + residual - - # LayerNorm - mean = tl.sum(x, axis=0) / n_cols - x_centered = x - mean - var = tl.sum(x_centered * x_centered, axis=0) / n_cols - x_norm = x_centered / tl.sqrt(var + eps) - - weight = tl.load(weight_ptr + col_offsets, mask=mask, other=1.0) - bias = tl.load(bias_ptr + col_offsets, mask=mask, other=0.0) - out = x_norm * weight + bias - - tl.store(out_ptr + row_start + col_offsets, out, mask=mask) - - -def fused_add_layernorm( - x: torch.Tensor, residual: torch.Tensor, - weight: torch.Tensor, bias: torch.Tensor, eps: float = 1e-5, -) -> torch.Tensor: - assert x.is_contiguous() and residual.is_contiguous() - shape = x.shape - x = x.view(-1, shape[-1]) - residual = residual.view(-1, shape[-1]) - n_rows, n_cols = x.shape - out = torch.empty_like(x) - grid = (n_rows,) - fused_add_layernorm_kernel[grid](x, residual, out, weight, bias, n_rows, n_cols, eps) - return out.view(shape) -``` - ---- - -## Pattern Selection Guide - -| Use Case | Pattern | Expected Speedup | -|----------|---------|------------------| -| FFN activation + dropout | GELU + Dropout | 1.8-2.2x | -| LLaMA FFN gate | SiLU + Multiply | 1.5-2x | -| LLaMA norm | RMSNorm | 1.4-2x | -| FFN with activation | Linear + GELU | 1.3-1.6x | -| Post-attention | Add + LayerNorm | 1.5-2x | diff --git a/.agents/skills/kernel-triton-writing/references/patterns-gemm.md b/.agents/skills/kernel-triton-writing/references/patterns-gemm.md deleted file mode 100644 index 62d2b6d5be0e..000000000000 --- a/.agents/skills/kernel-triton-writing/references/patterns-gemm.md +++ /dev/null @@ -1,292 +0,0 @@ - - - - -# Triton GEMM Patterns - -Reusable matrix multiplication patterns from Triton tutorials 03, 08, 09, 10. - -## Matrix Multiplication - -Block-tiled GEMM with L2 cache optimization. The workhorse pattern for dense matmul. - -### Autotune Configs - -| Config | BLOCK_M | BLOCK_N | BLOCK_K | num_stages | num_warps | -|--------|---------|---------|---------|------------|-----------| -| 1 | 128 | 256 | 64 | 3 | 8 | -| 2 | 64 | 256 | 32 | 4 | 4 | -| 3 | 128 | 128 | 32 | 4 | 4 | -| 4 | 128 | 64 | 32 | 4 | 4 | -| 5 | 64 | 128 | 32 | 4 | 4 | -| 6 | 128 | 32 | 32 | 4 | 4 | -| 7 | 64 | 32 | 32 | 5 | 2 | -| 8 | 32 | 64 | 32 | 5 | 2 | - -All configs use `GROUP_SIZE_M=8`. `key=["M", "N", "K"]` triggers re-autotuning on shape change. - -### Complete Kernel - -```python -@triton.autotune( - configs=[ - triton.Config({"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 64, "GROUP_SIZE_M": 8}, num_stages=3, num_warps=8), - triton.Config({"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 256, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 8}, num_stages=4, num_warps=4), - triton.Config({"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 8}, num_stages=4, num_warps=4), - triton.Config({"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 8}, num_stages=4, num_warps=4), - triton.Config({"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 128, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 8}, num_stages=4, num_warps=4), - triton.Config({"BLOCK_SIZE_M": 128, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 8}, num_stages=4, num_warps=4), - triton.Config({"BLOCK_SIZE_M": 64, "BLOCK_SIZE_N": 32, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 8}, num_stages=5, num_warps=2), - triton.Config({"BLOCK_SIZE_M": 32, "BLOCK_SIZE_N": 64, "BLOCK_SIZE_K": 32, "GROUP_SIZE_M": 8}, num_stages=5, num_warps=2), - ], - key=["M", "N", "K"], -) -@triton.jit -def matmul_kernel( - a_ptr, b_ptr, c_ptr, - M, N, K, - stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn, - BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, - GROUP_SIZE_M: tl.constexpr, -): - """C = A @ B. A is (M,K), B is (K,N), C is (M,N).""" - pid = tl.program_id(axis=0) - num_pid_m = tl.cdiv(M, BLOCK_SIZE_M) - num_pid_n = tl.cdiv(N, BLOCK_SIZE_N) - - # L2 cache optimization: super-grouping — nearby pids share B columns - num_pid_in_group = GROUP_SIZE_M * num_pid_n - group_id = pid // num_pid_in_group - first_pid_m = group_id * GROUP_SIZE_M - group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M) - pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m) - pid_n = (pid % num_pid_in_group) // group_size_m - - # Multi-dimensional pointer arithmetic via broadcasting - # 1D offset vectors + strides -> 2D block of pointers - offs_am = (pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)) % M - offs_bn = (pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)) % N - offs_k = tl.arange(0, BLOCK_SIZE_K) - a_ptrs = a_ptr + (offs_am[:, None] * stride_am + offs_k[None, :] * stride_ak) - b_ptrs = b_ptr + (offs_k[:, None] * stride_bk + offs_bn[None, :] * stride_bn) - - # Accumulate along K - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) - for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): - a = tl.load(a_ptrs, mask=offs_k[None, :] < K - k * BLOCK_SIZE_K, other=0.0) - b = tl.load(b_ptrs, mask=offs_k[:, None] < K - k * BLOCK_SIZE_K, other=0.0) - accumulator = tl.dot(a, b, accumulator) - a_ptrs += BLOCK_SIZE_K * stride_ak - b_ptrs += BLOCK_SIZE_K * stride_bk - c = accumulator.to(tl.float16) - - # Store with boundary masking - offs_cm = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M) - offs_cn = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N) - c_ptrs = c_ptr + stride_cm * offs_cm[:, None] + stride_cn * offs_cn[None, :] - c_mask = (offs_cm[:, None] < M) & (offs_cn[None, :] < N) - tl.store(c_ptrs, c, mask=c_mask) -``` - -### Launch Wrapper - -```python -def matmul(a, b): - assert a.shape[1] == b.shape[0], "Incompatible dimensions" - assert a.is_contiguous(), "Matrix A must be contiguous" - M, K = a.shape - K, N = b.shape - c = torch.empty((M, N), device=a.device, dtype=torch.float16) - grid = lambda META: (triton.cdiv(M, META["BLOCK_SIZE_M"]) * triton.cdiv(N, META["BLOCK_SIZE_N"]),) - matmul_kernel[grid]( - a, b, c, M, N, K, - a.stride(0), a.stride(1), b.stride(0), b.stride(1), c.stride(0), c.stride(1), - ) - return c -``` - -### Key Patterns - -- **Pointer broadcasting:** `offs_row[:, None] * stride_row + offs_col[None, :] * stride_col` creates 2D pointer block from 1D offsets. -- **tl.dot(a, b, acc):** Accumulates `a @ b` into `acc`. Always use float32 accumulator. -- **Super-grouping:** `GROUP_SIZE_M` controls how many M-tiles share N-tiles, improving L2 hit rate on B. Typically 8. -- **Boundary masking:** `% M`/`% N` wraps OOB offsets to valid addresses (loads still masked). K-dim uses explicit `mask=offs_k < remaining`. - -## Grouped GEMM - -Batched independent matmuls in a single persistent kernel. Use case: mixture-of-experts (MoE). - -### Core Pattern - -```python -@triton.jit -def grouped_matmul_kernel( - a_ptrs, b_ptrs, c_ptrs, # device arrays of per-group pointers - m_sizes, n_sizes, k_sizes, # per-group dimensions - lds_a, lds_b, lds_c, # per-group leading dimensions - group_offsets, # cumulative tile count per group - num_tiles, - NUM_SM: tl.constexpr, BLOCK_SIZE_M: tl.constexpr, - BLOCK_SIZE_N: tl.constexpr, BLOCK_SIZE_K: tl.constexpr, -): - tile_idx = tl.program_id(0) - # Persistent: each SM strides across all tiles - for tile_idx in tl.range(tile_idx, num_tiles, NUM_SM, num_stages=0): - # Binary-search group_offsets to find which group owns this tile - # Compute (pid_m, pid_n) within that group - # Standard matmul accumulation loop for group's (M,K) x (K,N) - ... - -# Launch: one program per SM -NUM_SM = torch.cuda.get_device_properties("cuda").multi_processor_count -grouped_matmul_kernel[(NUM_SM,)]( - a_ptrs, b_ptrs, c_ptrs, m_sizes, n_sizes, k_sizes, - lds_a, lds_b, lds_c, group_offsets, total_tiles, - NUM_SM=NUM_SM, BLOCK_SIZE_M=128, BLOCK_SIZE_N=128, BLOCK_SIZE_K=32, -) -``` - -### TMA Variant (Hopper+) - -```python -# Device-side TMA descriptors — shape varies per group -desc_a = tl.make_tensor_descriptor( - a_group_ptr, shape=[M_g, K_g], strides=[lda, 1], - block_shape=[BLOCK_SIZE_M, BLOCK_SIZE_K], -) -desc_b = tl.make_tensor_descriptor( - b_group_ptr, shape=[K_g, N_g], strides=[ldb, 1], - block_shape=[BLOCK_SIZE_K, BLOCK_SIZE_N], -) -a = desc_a.load([pid_m * BLOCK_SIZE_M, k * BLOCK_SIZE_K]) -b = desc_b.load([k * BLOCK_SIZE_K, pid_n * BLOCK_SIZE_N]) -``` - -### Key Patterns - -- **Persistent kernel:** Grid = NUM_SM. Each program loops via `tl.range(pid, total, NUM_SM)`. -- **Device-side scheduling:** Binary search on cumulative tile offsets maps flat tile_id to (group, tile_m, tile_n). -- **`tl.make_tensor_descriptor`:** Creates TMA descriptors on-device (needed because shape changes per group). -- **`num_stages=0`:** Outer loop has no pipelining; inner K loop is pipelined. - -## Persistent Matmul - -TMA descriptors and warp specialization for Hopper/Blackwell. Three progressive variants. - -### Variant 1: Persistent with Pointer Arithmetic - -Same as basic GEMM but with `tl.range(start_pid, num_tiles, NUM_SM)` outer loop -and grid = `(NUM_SM,)`. See Grouped GEMM pattern above for the persistent loop structure. - -### Variant 2: TMA Descriptors - -```python -# Host-side: create TMA descriptors before launch -from triton.tools.experimental_descriptor import create_2d_tma_descriptor -desc_a = create_2d_tma_descriptor(a_ptr, M, K, BLOCK_SIZE_M, BLOCK_SIZE_K, a.element_size()) -desc_b = create_2d_tma_descriptor(b_ptr, K, N, BLOCK_SIZE_K, BLOCK_SIZE_N, b.element_size()) - -# Kernel: load via hardware TMA unit (no manual pointer math) -@triton.jit -def matmul_kernel_tma(desc_a, desc_b, c_ptr, M, N, K, ...): - # Inside K-loop: - a = tl._experimental_descriptor_load( - desc_a, [pid_m * BLOCK_SIZE_M, k * BLOCK_SIZE_K], - [BLOCK_SIZE_M, BLOCK_SIZE_K], tl.float16) - b = tl._experimental_descriptor_load( - desc_b, [k * BLOCK_SIZE_K, pid_n * BLOCK_SIZE_N], - [BLOCK_SIZE_K, BLOCK_SIZE_N], tl.float16) - accumulator = tl.dot(a, b, accumulator) -``` - -### Variant 3: Warp Specialization - -```python -# Warps split into producers (TMA loads) and consumers (tl.dot compute). -# Compiler manages producer/consumer synchronization automatically. -matmul_kernel_warp_spec[(NUM_SM,)]( - desc_a, desc_b, c, M, N, K, ..., - BLOCK_SIZE_M=128, BLOCK_SIZE_N=256, BLOCK_SIZE_K=64, - num_stages=4, num_warps=8, - num_consumer_groups=2, # warp groups for compute - num_buffers_warp_spec=4, # pipeline depth for producer/consumer overlap -) -``` - -### FP8 Support (compute capability >= 9.0) - -```python -a = tl.load(a_ptrs, mask=..., other=0.0).to(tl.float8e5m2) # or tl.float8e4m3fn -b = tl.load(b_ptrs, mask=..., other=0.0).to(tl.float8e5m2) -accumulator = tl.dot(a, b, accumulator) # acc stays float32 -``` - -### Key Patterns - -- **tl.range(start, end, step):** Persistent loop with SM-count stride. -- **TMA descriptors:** Host-side `create_2d_tma_descriptor` or device-side `tl.make_tensor_descriptor`. Offloads address gen to hardware. -- **Warp specialization:** `num_consumer_groups` + `num_buffers_warp_spec` split warps into memory producers and compute consumers. -- **Epilogue subtiling:** Slice accumulator along N for the store phase to cut register pressure. - -## Block-Scaled Matmul - -Per-block scale factors for microscaling (MX) formats. -Requires 5th-gen Tensor Cores (compute capability >= 10.0, Blackwell+). - -### Supported Formats - -| Format | Element Type | Scale Block Size | Platform | -|--------|-------------|------------------|----------| -| mxfp8 | float8 (e5m2 / e4m3) | 32 elements | NVIDIA + AMD | -| mxfp4 | float4 (e2m1) | 32 elements | NVIDIA + AMD | -| nvfp4 | float4 (e2m1) | 16 elements | NVIDIA only | - -### Kernel with tl.dot_scaled - -```python -@triton.jit -def matmul_kernel_block_scaled( - a_ptr, b_ptr, c_ptr, a_scale_ptr, b_scale_ptr, - M, N, K, - stride_am, stride_ak, stride_bk, stride_bn, stride_cm, stride_cn, - stride_a_scale_m, stride_a_scale_k, stride_b_scale_n, stride_b_scale_k, - BLOCK_SIZE_M: tl.constexpr, BLOCK_SIZE_N: tl.constexpr, - BLOCK_SIZE_K: tl.constexpr, GROUP_SIZE_M: tl.constexpr, -): - pid = tl.program_id(axis=0) - # ... super-grouping (same as basic GEMM) ... - - accumulator = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32) - for k in range(0, tl.cdiv(K, BLOCK_SIZE_K)): - a = tl.load(a_ptrs, mask=..., other=0.0) - b = tl.load(b_ptrs, mask=..., other=0.0) - a_scale = tl.load(a_scale_ptrs) # [BLOCK_M, BLOCK_K // 32] - b_scale = tl.load(b_scale_ptrs) # [BLOCK_N, BLOCK_K // 32] - # Hardware-accelerated scaled dot product - accumulator = tl.dot_scaled(a, a_scale, "e4m3", b, b_scale, "e4m3", accumulator) - # advance pointers ... - tl.store(c_ptrs, accumulator.to(tl.float16), mask=c_mask) -``` - -### Key Patterns - -- **tl.dot_scaled(a, a_scale, a_fmt, b, b_scale, b_fmt, acc):** Single instruction applying per-block scales during matmul. Replaces manual dequant-then-multiply. -- **Format strings:** `"e4m3"`, `"e5m2"`, `"e2m1"` passed to `tl.dot_scaled`. -- **Preshuffling:** Always preprocess scales into vendor-specific layout before kernel launch. -- **Hardware:** NVIDIA CC >= 10.0 (Blackwell, PTX 8.7+). AMD CDNA3+ (MI300X). diff --git a/.agents/skills/kernel-triton-writing/references/semantics.md b/.agents/skills/kernel-triton-writing/references/semantics.md new file mode 100644 index 000000000000..bef898a2a04d --- /dev/null +++ b/.agents/skills/kernel-triton-writing/references/semantics.md @@ -0,0 +1,66 @@ + + +# Triton Semantics That Affect Correctness + +Use the installed Triton API and the +[official language semantics](https://triton-lang.org/main/python-api/triton-semantics.html) +as the authority. These reminders highlight common porting hazards. + +## Programs, shapes, and memory + +A Triton program operates on blocks of values. The programmer still controls +the launch grid, block shapes, pointer arithmetic, masks, and access pattern; +the compiler does not make an arbitrary layout coalesced or race-free. + +Broadcasting follows documented tensor-shape rules. Make dimensions explicit +with operations such as `[:, None]` and `[None, :]`, and ensure masks broadcast +to the corresponding pointer block. Do not assume Triton tensors are limited +to two dimensions. + +Masked loads require an `other` value when masked lanes can participate in +later computation. Choose a value that is neutral for the operation. Masked +stores are still required at output boundaries. + +`tl.where` evaluates both branches. It selects values; it does not guard an +otherwise-invalid load or store. + +## Numeric behavior + +Use the documented semantics of each operation rather than one global +promotion rule. In particular, reductions, dot products, transcendental +functions, and stores can have different conversion or precision behavior. +Specify accumulator or input precision when the operator contract requires it, +and make the reference use a comparable precision mode. + +Python scalars and Triton tensors do not always promote like PyTorch tensors. +If promotion affects range or precision, cast deliberately and cover the case +with a focused test. + +## Signed integer division + +For integer division and remainder involving Triton tensor operands, Triton +uses C-style truncation toward zero. Python's `//` instead rounds toward +negative infinity. The difference matters only when an operand can be negative. +Keep offsets non-negative when possible, or implement and test the intended +floor-division or modulo operation explicitly. + +For example, with a positive divisor, normalizing a remainder can be expressed +as `(value % divisor + divisor) % divisor`. Verify signed edge cases against the +intended public semantics rather than assuming translated Python expressions +behave identically. diff --git a/.agents/skills/kernel-triton-writing/references/troubleshooting.md b/.agents/skills/kernel-triton-writing/references/troubleshooting.md index c1269f24aa14..c539f8adf31b 100644 --- a/.agents/skills/kernel-triton-writing/references/troubleshooting.md +++ b/.agents/skills/kernel-triton-writing/references/troubleshooting.md @@ -1,7 +1,7 @@ - - -# Triton Troubleshooting, Debugging, and Benchmarking - -## Debug Operations — Compile-Time - -### static_print — Inspect Types and Constants at Compile Time - -Prints values during kernel compilation (not at runtime). Use to verify -constexpr values, tensor shapes, and dtypes. - -```python -@triton.jit -def kernel(X_ptr, N, BLOCK_SIZE: tl.constexpr): - tl.static_print("BLOCK_SIZE", BLOCK_SIZE) # prints the constexpr value - x = tl.load(X_ptr + tl.arange(0, BLOCK_SIZE)) - tl.static_print("x dtype", x.dtype) # prints the tensor dtype - tl.static_print("x shape", x.shape) # prints the tensor shape -``` - -Output appears in stderr during compilation (not on device): - -``` -BLOCK_SIZE 1024 -x dtype float32 -x shape (1024,) -``` - -### static_assert — Compile-Time Invariant Checks - -Fails compilation if condition is false. Use for constexpr guards. - -```python -@triton.jit -def kernel(X_ptr, BLOCK_SIZE: tl.constexpr): - tl.static_assert(BLOCK_SIZE % 32 == 0, "BLOCK_SIZE must be multiple of 32") - tl.static_assert(BLOCK_SIZE <= 4096, "BLOCK_SIZE too large") -``` - -| Function | Runs When | Requires TRITON_DEBUG | Use For | -|----------|-----------|----------------------|---------| -| `tl.static_print(label, value)` | Compilation | No | Inspecting types, shapes, constexpr values | -| `tl.static_assert(cond, msg)` | Compilation | No | Enforcing constexpr constraints | - ---- - -## Debug Operations — Runtime (On-Device) - -### device_print — Print Tensor Values from GPU - -Prints values at runtime from every active thread. Produces large output -on multi-element blocks; use masks or conditions to limit output. - -```python -@triton.jit -def kernel(X_ptr, BLOCK_SIZE: tl.constexpr): - pid = tl.program_id(0) - offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - x = tl.load(X_ptr + offs) - - # Print from all programs — very verbose - tl.device_print("x", x) - - # Print from only program 0 — much less output - if pid == 0: - tl.device_print("x[0]", x) -``` - -### device_assert — Runtime Assertions (Requires TRITON_DEBUG=1) - -Only executes when `TRITON_DEBUG=1` is set. Silent otherwise. - -```python -@triton.jit -def kernel(X_ptr, N, BLOCK_SIZE: tl.constexpr): - pid = tl.program_id(0) - offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - tl.device_assert(offs < N, "out-of-bounds access") - x = tl.load(X_ptr + offs) -``` - -```bash -# Enable device_assert checks -TRITON_DEBUG=1 .venv/bin/python my_kernel.py -``` - -| Function | Runs When | Requires TRITON_DEBUG | Use For | -|----------|-----------|----------------------|---------| -| `tl.device_print(label, value)` | Runtime (GPU) | No | Inspecting tensor values on device | -| `tl.device_assert(cond, msg)` | Runtime (GPU) | **Yes** (`=1`) | Bounds checks, NaN guards, invariants | - -### Gotcha: device_assert Does Nothing Without TRITON_DEBUG - -If you add `tl.device_assert` and your kernel still silently produces wrong results, -check that `TRITON_DEBUG=1` is exported **before** the kernel is compiled/cached. - ---- - -## Interpreter Mode — CPU Step-Through Debugging - -Setting `TRITON_INTERPRET=1` runs all Triton kernels on the CPU using NumPy, -bypassing GPU compilation entirely. This enables standard Python debugging. - -### Basic Usage - -```bash -TRITON_INTERPRET=1 .venv/bin/python my_kernel.py -``` - -### Debugging with pdb - -```bash -TRITON_INTERPRET=1 .venv/bin/python -m pdb my_kernel.py -``` - -Set breakpoints inside `@triton.jit` functions — they execute as normal Python -in interpreter mode. - -```python -@triton.jit -def kernel(X_ptr, Y_ptr, BLOCK_SIZE: tl.constexpr): - pid = tl.program_id(0) - offs = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE) - x = tl.load(X_ptr + offs) - # In interpreter mode, you can set a pdb breakpoint here - # and inspect x as a numpy array - import pdb; pdb.set_trace() # works only with TRITON_INTERPRET=1 - y = x * 2 - tl.store(Y_ptr + offs, y) -``` - -### Interpreter Mode Limitations - -| Limitation | Detail | -|------------|--------| -| No bfloat16 support | NumPy lacks native bf16; operations may error or use fp32 fallback | -| No indirect memory access | Gather/scatter patterns may not work correctly | -| No GPU-specific behavior | Race conditions, warp-level ops not simulated | -| Performance | Orders of magnitude slower than GPU — use small inputs only | -| Caching | Set `TRITON_INTERPRET=1` before any kernel is compiled/cached | -| atomic_add with fp16 | Known issue — may raise `ValueError('unsupported data type')` | - ---- - -## Third-Party Debug Tools - -| Tool | Vendor | Purpose | Usage | -|------|--------|---------|-------| -| `compute-sanitizer` | NVIDIA | Memory access checker (out-of-bounds, races) | `compute-sanitizer .venv/bin/python my_kernel.py` | -| `compute-sanitizer --tool memcheck` | NVIDIA | Detailed memory error reports | `compute-sanitizer --tool memcheck .venv/bin/python my_kernel.py` | -| `compute-sanitizer --tool racecheck` | NVIDIA | Shared memory race detection | `compute-sanitizer --tool racecheck .venv/bin/python my_kernel.py` | -| AddressSanitizer | AMD (ROCm) | Memory error detection on AMD GPUs | Compile with ASan flags | -| `triton-viz` | Community | Visual trace of memory access patterns | `uv pip install triton-viz` | - -### compute-sanitizer Example - -```bash -# Check for out-of-bounds memory access -compute-sanitizer --tool memcheck .venv/bin/python my_kernel.py - -# Check for shared memory race conditions -compute-sanitizer --tool racecheck .venv/bin/python my_kernel.py -``` - ---- - -## Benchmarking — triton.testing - -### do_bench — Micro-Benchmark a Function - -```python -import triton - -ms = triton.testing.do_bench(lambda: my_kernel[grid](x, y, N, BLOCK_SIZE=1024)) -print(f"{ms:.3f} ms") -``` - -| Parameter | Type | Default | Description | -|-----------|------|---------|-------------| -| `fn` | `Callable` | required | Zero-arg function to benchmark (use lambda) | -| `warmup` | `int` | `25` | Warmup time in milliseconds | -| `rep` | `int` | `100` | Repetition time in milliseconds | -| `grad_to_none` | `torch.Tensor` | `None` | Reset this tensor's gradient to None each iteration | -| `quantiles` | `list[float]` | `None` | Percentiles to return (e.g., `[0.2, 0.5, 0.8]`) | -| `return_mode` | `str` | `"mean"` | `"min"`, `"max"`, `"mean"`, `"median"`, or `"all"` | - -### Benchmark Class + perf_report — Parameterized Benchmarks - -```python -import triton -from triton.testing import Benchmark, perf_report - -@perf_report( - Benchmark( - x_names=["N"], # argument to vary - x_vals=[2**i for i in range(10, 25)], # values for N - line_arg="provider", # line grouping - line_vals=["triton", "torch"], # line values - line_names=["Triton", "PyTorch"], # legend labels - plot_name="vector-add-performance", # plot filename - args={}, # fixed args - xlabel="Vector Size (N)", - ylabel="GB/s", - x_log=True, - ) -) -def benchmark(N, provider): - x = torch.randn(N, device='cuda', dtype=torch.float32) - y = torch.randn(N, device='cuda', dtype=torch.float32) - output = torch.empty_like(x) - if provider == "triton": - ms = triton.testing.do_bench(lambda: my_kernel[grid](x, y, output, N, BLOCK_SIZE=1024)) - else: - ms = triton.testing.do_bench(lambda: x + y) - gbps = 3 * x.numel() * x.element_size() / ms * 1e-6 # 3 = 2 reads + 1 write - return gbps - -# Run and save plot -benchmark.run(show_plots=True, save_path="./benchmarks/") -``` - -### Correctness Testing — torch.testing.assert_close - -Triton does not ship its own `assert_close`. Use PyTorch: - -```python -import torch -torch.testing.assert_close(triton_output, torch_reference, atol=1e-2, rtol=1e-2) -``` - -For fp16/bf16 kernels, use relaxed tolerances (`atol=1e-1, rtol=1e-1`). - ---- - -## Common Errors Table - -| Error / Symptom | Cause | Fix | -|-----------------|-------|-----| -| `shape mismatch` in binary op | Tensor shapes do not broadcast | Check shapes with `tl.static_print`; add `[:, None]` or `[None, :]` | -| `BLOCK_SIZE is not a constexpr` | Block size passed as runtime value | Add `: tl.constexpr` annotation to the parameter | -| `mask dimensions do not match` | Mask shape incompatible with load/store block | Ensure mask broadcasts to the pointer offset shape | -| OOM during autotuning | Too many `@triton.autotune` configs | Reduce config list; avoid combinatorial explosion of BLOCK_M/N/K | -| `device_assert` has no effect | `TRITON_DEBUG` not set to `1` | Export `TRITON_DEBUG=1` before running | -| Silent wrong results | Off-by-one in pointer arithmetic | Use `tl.device_print` to inspect offsets; test with `TRITON_INTERPRET=1` | -| `incompatible types` in store | Computed dtype does not match output pointer dtype | Cast explicitly: `tl.store(ptr, val.to(tl.float16))` | -| Kernel not updating after edit | Triton cache serving stale binary | Move the confirmed Triton cache directory aside and retry | -| `ValueError: unsupported data type` in interpreter | bf16 or fp8 used with `TRITON_INTERPRET=1` | Use fp16 or fp32 for interpreter debugging | -| `grid must be a tuple` | Lambda grid returns int, not tuple | Return `(value,)` with trailing comma | -| NaN output, correct logic | fp16 overflow in accumulator | Use `tl.float32` for accumulation, cast on store | -| `expected constexpr` in `tl.arange` | Non-constexpr argument to arange | Both args of `tl.arange(start, end)` must be constexpr | -| Mismatched results vs PyTorch | C integer division semantics | See concepts-semantics.md: Triton uses truncation, not floor division | -| `triton.OutOfResources` | Register/shared memory pressure | Reduce BLOCK_SIZE or number of live variables | - ---- - -## Environment Variables Reference - -| Variable | Value | Effect | -|----------|-------|--------| -| `TRITON_DEBUG` | `1` | Enable `device_assert`, extra runtime checks | -| `TRITON_INTERPRET` | `1` | Run kernels on CPU via NumPy (no GPU) | -| `TRITON_CACHE_DIR` | path | Override default cache directory (`~/.triton/cache/`) | -| `MLIR_ENABLE_DUMP` | `1` | Dump MLIR intermediate representations | -| `TRITON_PRINT_AUTOTUNING` | `1` | Print autotuning results to stderr | +# Triton Troubleshooting + +Consult the current +[Triton debugging guide](https://triton-lang.org/main/programming-guide/chapter-3/debugging.html) +before relying on environment variables or interpreter limitations, which can +change between Triton versions. + +## Triage order + +1. Reproduce with the smallest failing shape and a deterministic input. +2. Determine whether the failure occurs during Python wrapping, Triton + compilation, launch, memory access, or numerical comparison. +3. Compare pointer offsets, strides, block shapes, masks, dtypes, and precision + modes with the reference contract. +4. Test boundary tiles and one full tile separately. +5. Add the smallest regression test that reproduces the failure before + broadening the shape matrix. + +## Built-in tools + +- `tl.static_print` and `tl.static_assert` inspect or validate compile-time + values. +- `tl.device_print` inspects runtime values. Restrict the printed programs and + lanes to keep output usable. +- `tl.device_assert` can check runtime invariants when enabled as documented by + the installed Triton version. +- `TRITON_INTERPRET=1` can help with supported operations, but interpreter + behavior is not a substitute for running on the target GPU. Check current + documented limitations before drawing conclusions from it. + +For memory faults on NVIDIA GPUs, run a focused reproducer under +`compute-sanitizer --tool memcheck`. Use backend-appropriate tooling on other +platforms. Sanitizer success does not establish numerical correctness or the +absence of logical races. + +## Symptom checklist + +| Symptom | Inspect | +| --- | --- | +| Boundary-only differences | Load/store masks, neutral masked-load values, and final partial tiles | +| Shape or broadcast error | Block shapes and explicit singleton dimensions | +| NaN or infinity | Input domain, masked values, division guards, intermediate dtype, and overflow | +| Nondeterministic differences | Aliasing, atomics, cross-program ownership, and RNG state | +| Matmul-like precision mismatch | Input precision, accumulator dtype, reference backend settings, and reduction order | +| Resource exhaustion | Tile sizes, live values, warp count, pipeline stages, and compiler diagnostics | +| Unexpected recompilation | Specialization keys, meta-parameters, shapes, strides, and cache configuration | +| Unexpected stale result | Confirm the loaded source and cache path before moving the exact cache directory aside | + +Do not prescribe a universal numerical tolerance. Derive it from the operation, +dtype, reduction depth, and supported precision contract. + +Performance diagnosis belongs to `$kernel-microbenchmark`, which covers timing, +generated-code inspection, multi-GPU comparisons, and speed-of-light checks. From 7459213637477372ea65fd53c2542df5f324084c Mon Sep 17 00:00:00 2001 From: Woosuk Kwon Date: Wed, 2 Sep 2026 21:16:11 +0000 Subject: [PATCH 4/4] [Agents] Update Triton skill attribution Co-authored-by: OpenAI Codex Signed-off-by: Woosuk Kwon --- .../skills/kernel-triton-writing/ORIGIN.md | 24 +++++++++---------- .agents/skills/kernel-triton-writing/SKILL.md | 8 +++---- .../references/semantics.md | 16 ++----------- .../references/troubleshooting.md | 16 ++----------- 4 files changed, 18 insertions(+), 46 deletions(-) diff --git a/.agents/skills/kernel-triton-writing/ORIGIN.md b/.agents/skills/kernel-triton-writing/ORIGIN.md index 8a36bfab46b5..ddc6d7e6ab6f 100644 --- a/.agents/skills/kernel-triton-writing/ORIGIN.md +++ b/.agents/skills/kernel-triton-writing/ORIGIN.md @@ -1,23 +1,21 @@ -# Source and License +# Provenance -This skill was copied from NVIDIA's TensorRT-LLM repository: +The initial version of this skill was copied from NVIDIA's TensorRT-LLM +repository: - Source: `https://github.com/NVIDIA/TensorRT-LLM/tree/main/.claude/skills/kernel-triton-writing` - Snapshot commit: `395985c025c8d1cf5aa842bc752b337ba88721b6` -- Copyright: Copyright (c) 2011-2026 NVIDIA CORPORATION & AFFILIATES. - All rights reserved. -- License: Apache License 2.0 +- Upstream license: Apache License 2.0 -The NVIDIA copyright and Apache-2.0 SPDX notices are preserved in the adapted -reference files. The vLLM copy adds explicit source comments and removes -unsupported skill metadata. +The content has since been substantially rewritten for vLLM. The source and +snapshot commit remain here to record the history of the initial import. The upstream standalone verification and benchmark scripts are omitted. vLLM uses its existing parametrized kernel pytest suites for correctness and the `kernel-microbenchmark` skill and `benchmarks/kernels/` for performance work. The associated fixed-name export contract and workflow sections are adapted to -those vLLM conventions. Copied API catalogs, fixed tuning recipes, performance -claims, and incomplete kernel examples were removed after review because they -duplicated versioned Triton documentation or were not generally supportable. -The remaining guidance directs contributors to current official Triton -documentation and device-specific measurement. +those vLLM conventions. The copied API catalogs, fixed tuning recipes, +performance claims, and incomplete kernel examples were removed after review +because they duplicated versioned Triton documentation or were not generally +supportable. The remaining guidance directs contributors to current official +Triton documentation and device-specific measurement. diff --git a/.agents/skills/kernel-triton-writing/SKILL.md b/.agents/skills/kernel-triton-writing/SKILL.md index 8b318de9fbac..9388e3ab91e9 100644 --- a/.agents/skills/kernel-triton-writing/SKILL.md +++ b/.agents/skills/kernel-triton-writing/SKILL.md @@ -7,7 +7,6 @@ description: > Triton kernel in vLLM. license: Apache-2.0 metadata: - author: NVIDIA Corporation source: https://github.com/NVIDIA/TensorRT-LLM source_commit: 395985c025c8d1cf5aa842bc752b337ba88721b6 --- @@ -15,10 +14,9 @@ metadata: # Triton Kernel Writing Use this workflow for OpenAI Triton (`@triton.jit`) work in vLLM. Use the diff --git a/.agents/skills/kernel-triton-writing/references/semantics.md b/.agents/skills/kernel-triton-writing/references/semantics.md index bef898a2a04d..4b6e1cd77921 100644 --- a/.agents/skills/kernel-triton-writing/references/semantics.md +++ b/.agents/skills/kernel-triton-writing/references/semantics.md @@ -1,19 +1,7 @@ # Triton Semantics That Affect Correctness diff --git a/.agents/skills/kernel-triton-writing/references/troubleshooting.md b/.agents/skills/kernel-triton-writing/references/troubleshooting.md index c539f8adf31b..be9759f1a1c9 100644 --- a/.agents/skills/kernel-triton-writing/references/troubleshooting.md +++ b/.agents/skills/kernel-triton-writing/references/troubleshooting.md @@ -1,19 +1,7 @@ # Triton Troubleshooting