[Diffusion] Add AKO4ALL kernel optimization skill (#21323)
This commit is contained in:
Xiaoyu Zhang
@@ -0,0 +1,130 @@
|
||||
---
|
||||
name: sglang-diffusion-ako4all-kernel
|
||||
description: Use when optimizing an existing SGLang diffusion kernel with AKO4ALL, including AKO4ALL repo hygiene, custom microbench setup, ncu-guided iteration, and end-to-end denoise validation. Also use when a sibling AKO4ALL repo must be cloned or refreshed before starting kernel tuning work.
|
||||
---
|
||||
|
||||
# SGLang Diffusion AKO4ALL Kernel
|
||||
|
||||
Use this skill to run the full AKO4ALL-based optimization loop for an existing SGLang diffusion kernel.
|
||||
It packages the workflow we used for diffusion Triton and JIT kernel tuning: bootstrap a custom AKO harness, benchmark and profile the kernel, iterate with `ncu`, port the best version back to `sglang`, then validate with targeted tests and model-level denoise runs.
|
||||
|
||||
This skill assumes a sibling repo layout like:
|
||||
|
||||
```text
|
||||
<base-dir>/
|
||||
├── sglang/
|
||||
└── AKO4ALL/
|
||||
```
|
||||
|
||||
If `AKO4ALL/` is missing under the current base directory, clone it first.
|
||||
|
||||
## Use This Skill When
|
||||
|
||||
- tuning an existing diffusion Triton, CUDA JIT, CuTeDSL, or runtime-integrated kernel in `sglang`
|
||||
- creating a custom AKO4ALL harness for a real diffusion kernel instead of using the default benchmark tasks
|
||||
- validating that a kernel-level win transfers to Qwen, FLUX, Wan, Hunyuan, MOVA, or other diffusion denoise latency
|
||||
- preparing PR artifacts such as microbench tables, `ncu` before/after data, and proof image outputs
|
||||
|
||||
Do not use this skill when adding a brand-new kernel from scratch with no existing SGLang integration.
|
||||
For that, start from the sibling kernel-authoring skills first:
|
||||
|
||||
- Triton: [../sglang-diffusion-triton-kernel/SKILL.md](../sglang-diffusion-triton-kernel/SKILL.md)
|
||||
- CUDA JIT: [../sglang-diffusion-cuda-kernel/SKILL.md](../sglang-diffusion-cuda-kernel/SKILL.md)
|
||||
- Denoise benchmark/profile: [../sglang-diffusion-benchmark-profile/SKILL.md](../sglang-diffusion-benchmark-profile/SKILL.md)
|
||||
|
||||
## Mandatory AKO4ALL Preflight
|
||||
|
||||
Before any AKO work:
|
||||
|
||||
1. Run `scripts/ensure_ako4all_clean.sh [base-dir]`.
|
||||
2. If `<base-dir>/AKO4ALL` does not exist, the script clones it.
|
||||
3. Do not continue unless `AKO4ALL` is:
|
||||
- on the upstream default branch, usually `main`
|
||||
- fully clean with no tracked or untracked local changes
|
||||
- exactly synced to `upstream/<default-branch>`
|
||||
4. If the script reports local commits, divergence, or a dirty worktree, stop and clean or re-clone the repo before continuing.
|
||||
|
||||
The script creates an `upstream` remote automatically when missing.
|
||||
By default it uses the existing `origin` URL, or `AKO4ALL_URL` if you need to override the clone source.
|
||||
|
||||
## Workflow
|
||||
|
||||
### 1. Scope the Kernel
|
||||
|
||||
- Identify the exact kernel entry point and runtime call sites in `sglang`.
|
||||
- Record the target shapes, dtypes, model families, and whether the kernel is on a hot path.
|
||||
- Reuse existing unit tests and benchmark entry points when they already exist.
|
||||
|
||||
If the implementation work is primarily Triton or CUDA authoring, read only the relevant sibling skill from the list above.
|
||||
|
||||
### 2. Bootstrap the AKO Harness
|
||||
|
||||
Inside the clean `AKO4ALL` repo:
|
||||
|
||||
- read `TASK.md` and `HINTS.md`
|
||||
- create a custom harness instead of relying on the stock benchmark tasks
|
||||
- mirror the real SGLang kernel into:
|
||||
- `input/reference.py`
|
||||
- `input/<kernel>.py`
|
||||
- `solution/<kernel>.py`
|
||||
- `bench/bench_<kernel>.py`
|
||||
- keep a short context note in `context/` when the kernel has model-specific shape assumptions or perf conclusions
|
||||
|
||||
The custom benchmark should:
|
||||
|
||||
- cover representative diffusion shapes
|
||||
- check correctness against the reference kernel
|
||||
- report aggregate runtime plus per-shape results when useful
|
||||
|
||||
### 3. Establish the Baseline
|
||||
|
||||
- run the AKO custom microbench before changing the kernel
|
||||
- capture one representative `ncu` baseline on the hottest meaningful shape
|
||||
- note whether the bottleneck looks like registers, occupancy, instruction count, launch config, or memory latency
|
||||
|
||||
### 4. Iterate in AKO4ALL
|
||||
|
||||
- change one idea at a time
|
||||
- rerun the microbench after every change
|
||||
- update `ITERATIONS.md` with hypothesis, result, and next step
|
||||
- prefer simple, explainable wins over clever rewrites that do not transfer
|
||||
|
||||
After 3 consecutive no-improvement or regression iterations:
|
||||
|
||||
- rerun `ncu`
|
||||
- re-read `ITERATIONS.md`
|
||||
- change direction instead of continuing blind sweeps
|
||||
|
||||
### 5. Port the Best Version Back to SGLang
|
||||
|
||||
- apply the best candidate to the real `sglang` kernel file
|
||||
- run import or syntax checks and targeted tests first
|
||||
- keep the AKO `solution/` version aligned with the main-tree version you actually want to keep
|
||||
|
||||
### 6. Validate on Real Models
|
||||
|
||||
- use the benchmark/profile skill for denoise perf dumps and before/after comparison
|
||||
- prefer exact local snapshot validation when testing local edits on a GPU box
|
||||
- run targeted kernel tests first
|
||||
- run model-level denoise benchmarks with perf dumps
|
||||
- compare baseline vs optimized runs with `compare_perf.py`
|
||||
- if the PR needs proof that generation still works, save one real model output image
|
||||
|
||||
### 7. Prepare PR Artifacts
|
||||
|
||||
At minimum, keep:
|
||||
|
||||
- one microbench table
|
||||
- one denoise-stage table
|
||||
- one end-to-end table
|
||||
- one `ncu` before/after pair on the most representative kernel shape
|
||||
- one generated image when the kernel affects production inference
|
||||
|
||||
See [references/ako-loop.md](references/ako-loop.md) for the checklist and common stop rules.
|
||||
|
||||
## Operating Rules
|
||||
|
||||
- Treat AKO4ALL repo hygiene as a gate, not a suggestion.
|
||||
- Prefer exact local snapshot validation over hand-wavy “remote tree is close enough”.
|
||||
- Keep model-level validation honest: if microbench improves but denoise does not, do not keep the AKO-only variant in the main code path.
|
||||
- When writing conclusions, explain the win in terms of measurable causes such as lower registers per thread, higher occupancy, fewer executed instructions, or better scheduler eligibility.
|
||||
+54
@@ -0,0 +1,54 @@
|
||||
# AKO Loop Checklist
|
||||
|
||||
Use this checklist after `scripts/ensure_ako4all_clean.sh` succeeds.
|
||||
|
||||
## Minimum Repo Layout
|
||||
|
||||
Inside `AKO4ALL/`, prefer these files for a diffusion kernel task:
|
||||
|
||||
- `input/reference.py`
|
||||
- `input/<kernel>.py`
|
||||
- `solution/<kernel>.py`
|
||||
- `bench/bench_<kernel>.py`
|
||||
- `context/<kernel>_notes.md`
|
||||
|
||||
## Baseline Checklist
|
||||
|
||||
- Reproduce the current SGLang kernel exactly in AKO first.
|
||||
- Run the custom microbench before making edits.
|
||||
- Record one representative `ncu` report on a real hot shape.
|
||||
- Note the baseline bottleneck in plain language.
|
||||
|
||||
## Iteration Discipline
|
||||
|
||||
- One optimization idea per iteration.
|
||||
- Re-benchmark after every code change.
|
||||
- Log the result in `ITERATIONS.md`.
|
||||
- Keep the best candidate easy to identify.
|
||||
|
||||
Stop a direction early when:
|
||||
|
||||
- 3 consecutive iterations do not beat the best runtime
|
||||
- correctness gets fragile
|
||||
- AKO-only gains stop transferring to real denoise runs
|
||||
|
||||
## Real Validation Gate
|
||||
|
||||
Before calling a kernel "done", validate all of:
|
||||
|
||||
- syntax or import checks
|
||||
- targeted unit test or regression test
|
||||
- kernel or op-level benchmark
|
||||
- model-level denoise benchmark with perf dumps
|
||||
- one generated image if the PR needs production proof
|
||||
|
||||
## PR Artifact Checklist
|
||||
|
||||
Prepare these artifacts:
|
||||
|
||||
- microbench table
|
||||
- denoise-stage table
|
||||
- end-to-end table
|
||||
- one `ncu` before or after pair
|
||||
- one short explanation of why the kernel got faster
|
||||
- one generated output image when applicable
|
||||
+84
@@ -0,0 +1,84 @@
|
||||
#!/usr/bin/env bash
|
||||
set -euo pipefail
|
||||
|
||||
BASE_DIR="${1:-$PWD}"
|
||||
BASE_DIR="$(cd "$BASE_DIR" && pwd)"
|
||||
AKO_DIR="${BASE_DIR}/AKO4ALL"
|
||||
CANONICAL_UPSTREAM_URL="https://github.com/TongmingLAIC/AKO4ALL.git"
|
||||
UPSTREAM_URL="${AKO4ALL_UPSTREAM_URL:-$CANONICAL_UPSTREAM_URL}"
|
||||
CLONE_URL="${AKO4ALL_URL:-$UPSTREAM_URL}"
|
||||
|
||||
say() {
|
||||
printf '[ako4all] %s\n' "$*"
|
||||
}
|
||||
|
||||
fail() {
|
||||
printf '[ako4all] ERROR: %s\n' "$*" >&2
|
||||
exit 1
|
||||
}
|
||||
|
||||
if [[ ! -d "$AKO_DIR/.git" ]]; then
|
||||
say "AKO4ALL not found under ${BASE_DIR}; cloning ${CLONE_URL}"
|
||||
git clone "$CLONE_URL" "$AKO_DIR"
|
||||
fi
|
||||
|
||||
cd "$AKO_DIR"
|
||||
|
||||
if ! git remote get-url origin >/dev/null 2>&1; then
|
||||
fail "AKO4ALL exists but has no origin remote."
|
||||
fi
|
||||
|
||||
if ! git remote get-url upstream >/dev/null 2>&1; then
|
||||
say "Adding missing upstream remote -> ${UPSTREAM_URL}"
|
||||
git remote add upstream "$UPSTREAM_URL"
|
||||
fi
|
||||
|
||||
git fetch upstream --prune
|
||||
git remote set-head upstream -a >/dev/null 2>&1 || true
|
||||
|
||||
default_branch="${AKO4ALL_BRANCH:-}"
|
||||
if [[ -z "$default_branch" ]]; then
|
||||
if upstream_head="$(git symbolic-ref --quiet --short refs/remotes/upstream/HEAD 2>/dev/null)"; then
|
||||
default_branch="${upstream_head#upstream/}"
|
||||
else
|
||||
default_branch="main"
|
||||
fi
|
||||
fi
|
||||
|
||||
if [[ -n "$(git status --porcelain)" ]]; then
|
||||
fail "AKO4ALL worktree is dirty. Clean all local changes before using this skill."
|
||||
fi
|
||||
|
||||
if git show-ref --verify --quiet "refs/heads/${default_branch}"; then
|
||||
git switch "$default_branch" >/dev/null
|
||||
else
|
||||
git switch -c "$default_branch" --track "upstream/${default_branch}" >/dev/null
|
||||
fi
|
||||
|
||||
git fetch upstream --prune
|
||||
|
||||
local_head="$(git rev-parse HEAD)"
|
||||
upstream_head="$(git rev-parse "upstream/${default_branch}")"
|
||||
|
||||
if [[ "$local_head" != "$upstream_head" ]]; then
|
||||
if git merge-base --is-ancestor "$local_head" "$upstream_head"; then
|
||||
say "Fast-forwarding ${default_branch} to upstream/${default_branch}"
|
||||
git merge --ff-only "upstream/${default_branch}" >/dev/null
|
||||
else
|
||||
fail "Local ${default_branch} diverges from upstream/${default_branch}. Reset or re-clone AKO4ALL before continuing."
|
||||
fi
|
||||
fi
|
||||
|
||||
if [[ -n "$(git status --porcelain)" ]]; then
|
||||
fail "AKO4ALL became dirty after sync; stop and inspect the repo."
|
||||
fi
|
||||
|
||||
final_head="$(git rev-parse HEAD)"
|
||||
expected_head="$(git rev-parse "upstream/${default_branch}")"
|
||||
if [[ "$final_head" != "$expected_head" ]]; then
|
||||
fail "AKO4ALL is not exactly at upstream/${default_branch}."
|
||||
fi
|
||||
|
||||
say "Ready: ${AKO_DIR}"
|
||||
say "Branch: ${default_branch}"
|
||||
say "Commit: ${final_head}"
|
||||
Reference in New Issue
Block a user