Skip to content

SKILL.md

Latest commit

 

History

History
114 lines (75 loc) · 6.25 KB

SKILL.md

File metadata and controls

114 lines (75 loc) · 6.25 KB
Raw
Error in user YAML: (<unknown>): mapping values are not allowed in this context at line 3 column 31
---
name: kernel-forge
description: Profile and optimize ROCm/HIP kernels in aiter. Use when benchmarking operators (PA, GEMM, MoE, MHA, Norm, etc.), analyzing roofline bottlenecks, tuning launch configs, or improving HBM bandwidth utilization.
argument-hint: <op-name> [impl: asm|ck|triton|hip] [goal: throughput|latency|bandwidth]
disable-model-invocation: true
effort: high
---

ROCm Kernel Performance Optimization — aiter

Target: $ARGUMENTS

aiter project root: ../aiter (resolved from this skill's working dir, typically /data/raid0/sijing/aiter).

Execution environment

  • File operations (Read / Edit / Write of source under aiter/) run on the host — do not wrap them in any container exec.
  • Run / build / profile / benchmark commands must execute inside the runtime environment the caller specifies (e.g. a docker container with ROCm + GPUs). Ask for it if not provided; do not hardcode a container name.

Phase 1 — Triage: which implementation are we optimizing?

aiter ships up to four implementations per operator. They have very different optimization handles, so you must pick one before going deeper.

Class Where it lives Optimization handle
asm csrc/py_itfs_cu/asm_<op>*.cu + *_configs.hpp Pre-built kernel variant selection; rarely edit the asm itself
ck csrc/ck_<op>*/, csrc/cktile_<op>*/ gen_instances.py codegen + aiter/configs/*_tuned_*.csv
triton aiter/ops/triton/**/*<op>*.py @triton.autotune config space, tl.load cache modifiers
hip csrc/kernels/<op>*.cu (pre-built) or csrc/cpp_itfs/<op>/*.cuh + *.jinja (JIT) Full-source HIP/C++: algorithm, sync structure, intrinsics — most malleable class. The two sub-paths share optimization patterns but have different rebuild flows — see playbooks/hip.md.

Detection (run from aiter root)

op=<op-name>
ls csrc/py_itfs_cu/asm_*${op}*.cu  2>/dev/null         # asm
ls -d csrc/ck_*${op}* csrc/cktile_*${op}* 2>/dev/null  # ck
find aiter/ops/triton -name "*${op}*.py" 2>/dev/null   # triton
ls csrc/kernels/*${op}*.cu 2>/dev/null                 # hip (pre-built)
ls -d csrc/cpp_itfs/*${op}* csrc/cpp_itfs/${op} 2>/dev/null  # hip (JIT)

A common miss: a Python wrapper says from csrc.cpp_itfs.<op> import <op> and you find no .cu under csrc/kernels/ — that means it's the JIT sub-path. The kernel source is csrc/cpp_itfs/<op>/*.cuh driven by a *.jinja template; rebuild and cache layout differ from the pre-built sub-path.

If multiple implementations exist and the user did not pass impl:, list them and ask which to optimize. The Python wrapper in aiter/ops/<op>.py (or aiter/<op>.py) usually shows which dispatcher path is used by default.

Check kernels/ for prior work on this op — ls kernels/<op>/ (run from this skill's dir). If the directory exists, it contains the original kernel pulled from aiter, any optimized versions (v1_*.cu, v2_*.cu, …), benchmark scripts, baseline data, and notes.md with past optimization decisions and failed experiments. Reading these before starting can save a full optimization cycle.

Bottleneck classification

Before going deeper, classify the operator's bottleneck. Three categories — not two:

  • Compute-bound — high arithmetic intensity, sits right of the roofline ridge. Tune tile sizes, ILP, num_stages.
  • Memory-bound — low AI, sits left of the ridge. Tune coalescing, vectorization, occupancy.
  • Sync-bound — dominated by __threadfence / atomicInc-spin / __syncthreads / kernel-launch overhead. Often invisible on a roofline plot — both VALU and MemUnit can be idle. Most common in multi-pass / multi-block-cooperative kernels (radix select, top-k, custom all-reduce, fused norm+quant). Detection: see shared/profiling.md.

Phase 2 — Measure baseline

Follow shared/profiling.md to get baseline throughput/latency/bandwidth and compare against the roofline ceiling in reference.md. Record numbers before changing any code.

Phase 3 — Optimize per implementation

Load the matching playbook and follow it:

  • asm → playbooks/asm.md
  • ck → playbooks/ck.md
  • triton → playbooks/triton.md
  • hip → playbooks/hip.md

Each playbook covers the bottleneck → fix mapping specific to that implementation, including which files to edit, which configs to regenerate, and what to rebuild.

Phase 4 — Validate

Follow shared/validation.md: correctness first (numerical tolerance vs. reference), then re-measure, then produce the before/after report.

Never report a perf win without a passing correctness check on the same build.

Phase 5 — Record

Mandatory. Before declaring the optimization done, update kernels/<op>/notes.md with the optimization record. If the file doesn't exist, create it. Each record should capture:

  • Version tag (e.g. V1, V2) and which layer was applied (structural / sync / instruction / resource)
  • What changed — 1 paragraph summary
  • Result table — before/after latency, correctness PASS/FAIL
  • Claim vs measured — if there was a prior estimate, compare it to actual
  • Failed / not tried — changes attempted that didn't work, and why

This ensures:

  • The next session (model or human) can grep "have we tried X on op Y?" instead of re-discovering.
  • The result table survives even if the source change is reverted, refactored, or merged into a larger PR.
  • Patterns accumulate — repeated wins on similar shapes turn into playbook updates.

Also save each optimized version as kernels/<op>/v{N}_{tag}.cu so the code history is preserved alongside the notes. Update kernels/README.md index if adding a new operator.

Files in this skill

  • SKILL.md — this file (entry + triage)
  • reference.md — GPU peak FLOPS / BW, roofline ridge points, rocprof counter cheat-sheet
  • playbooks/{asm,ck,triton,hip}.md — per-implementation optimization guides
  • shared/profiling.md — baseline measurement + rocprof recipes
  • shared/validation.md — correctness check + final report template
  • kernels/ — per-operator directory with original kernel, optimized versions, benchmarks, baseline data, and notes; read before optimizing a new op to see what has been tried