From 689e9ef05cd27add25c3c704adef3de8e8842cca Mon Sep 17 00:00:00 2001
From: Xiaoyu Zhang <35585791+BBuf@users.noreply.github.com>
Date: Wed, 25 Mar 2026 18:46:21 +0800
Subject: [PATCH] [Diffusion] Add AKO4ALL kernel optimization skill (#21323)

---
 .../sglang-diffusion-ako4all-kernel/SKILL.md  | 130 ++++++++++++++++++
 .../references/ako-loop.md                    |  54 ++++++++
 .../scripts/ensure_ako4all_clean.sh           |  84 +++++++++++
 3 files changed, 268 insertions(+)
 create mode 100644 python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-ako4all-kernel/SKILL.md
 create mode 100644 python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-ako4all-kernel/references/ako-loop.md
 create mode 100755 python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-ako4all-kernel/scripts/ensure_ako4all_clean.sh

diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-ako4all-kernel/SKILL.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-ako4all-kernel/SKILL.md
new file mode 100644
index 000000000..ca6ed3c61
--- /dev/null
+++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-ako4all-kernel/SKILL.md
@@ -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.
diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-ako4all-kernel/references/ako-loop.md b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-ako4all-kernel/references/ako-loop.md
new file mode 100644
index 000000000..557af9dba
--- /dev/null
+++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-ako4all-kernel/references/ako-loop.md
@@ -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
diff --git a/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-ako4all-kernel/scripts/ensure_ako4all_clean.sh b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-ako4all-kernel/scripts/ensure_ako4all_clean.sh
new file mode 100755
index 000000000..055ad56ec
--- /dev/null
+++ b/python/sglang/multimodal_gen/.claude/skills/sglang-diffusion-ako4all-kernel/scripts/ensure_ako4all_clean.sh
@@ -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}"