implement-kernel

Files (1)

• SKILL.md

SKILL.md (first lines)

# Implement Kernel Optimization

> **PIPELINE_MODE**: When `pipeline_mode: true` — **NO `EnterPlanMode`**, **NO `AskUserQuestion`**, **NO `/validate`**, **NO `/commit`**. Implement the Tier-1 change directly and return. The pipeline orchestrates validation and commit externally.

End-to-end workflow: design → implement → validate → commit a kernel optimization for SGLang on MI355 (ROCm/HIP). **Never commit before validation passes.**

## Usage

```
/implement-kernel <description of what to optimize>
```

If no description is provided, ask with `AskUserQuestion`.

### PIPELINE_MODE (from `/kernel-fusion-pipeline`)

When invoked with `PIPELINE_MODE=1` or CONFIG containing `"pipeline_mode": true`:

- **Tier 1 only** — skip `EnterPlanMode` / `ExitPlanMode` and plan approval; read HIP path and implement immediately.
- Use `CONFIG.worktree` as `$SGLANG_ROOT` if set.
- Do **not** use `AskUserQuestion`.
- Do **not** invoke `/validate` or `/commit` from this skill — the pipeline handles those per worktree.

---

## Step 0: Gather optimization design

### 0a: Detect the active SGLang installation

Before doing anything else, determine which SGLang repo is in use. There may be multiple SGLang directories on the system (e.g., `/sgl-workspace/sglang`, `$HOME/sglang`). **Always use the one that's currently installed in the active Python environment.**

```bash
SGLANG_ROOT=$(python3 -c "import sglang, pathlib; print(pathlib.Path(sglang.__file__).resolve().parents[2])")
echo "Active SGLang root: $SGLANG_ROOT"
```

Use `$SGLANG_ROOT` for all subsequent paths instead of hardcoded paths like `$HOME/sglang` or `/sgl-workspace/sglang`. All file paths in the source reference table below should be prefixed with `$SGLANG_ROOT/python/sglang/srt/` rather than a hardcoded path.

### 0b: If the user provided a description after `/implement-kernel`:

Parse the free text for:
- **What to optimize**: the specific kernel, module, or operation
- **Tier**: infer from the description:
  - Tier 1 (code refactor) — keywords: "switch dispatch", "remove", "fuse existing", "enable flag", "set env var"
  - Tier 2 (Triton kernel) — keywords: "write Triton", "fuse into Triton", "Triton kernel"
  - Tier 3 (aiter C++/HIP) — keywords: "tune CK", "aiter kernel", "HIP kernel", "assembly"

### If no description provided:

Use `AskUserQuestion` with these questions:

1. **What to optimize?** (free text — e.g., "fuse softmax+weighted_sum in compressor decode", "switch Q proj to BF16 for small-M decode")
2. **Which tier?**
   - Tier 1: Code refactor (Python-level, no new kernels) (Recommended)
   - Tier 2: Triton kernel (write/modify Triton kernel)
   - Tier 3: aiter C++/HIP kernel (CK tile, assembly, low-level HIP)

### Identify target files

Use the source file reference table (at the bottom of this document) to identify which files the optimization will touch. If ambiguous, ask the user. (**PIPELINE_MODE**: infer from the Tier-1 description, do not ask.)

---

## Step 1: Design the implementation (plan mode)

**Skip this entire step when PIPELINE_MODE is active and tier is 1** — go directly to Step 2.

Enter plan mode with `EnterPlanMode`.

### 1a. Read the current MI355 (HIP) code path

Read the target files identified in Step 0. Trace the full call chain from the SGLang layer Python code down to the kernel invocation on HIP:

1. Read the SGLang layer file — find the `is_hip()` / `_use_aiter` branch
2. Trace into aiter library at `/sgl-workspace/aiter/` if applicable
3. Identify what operations run as separate kernels

### 1b. Read the B200 (CUDA) reference path

Read the same SGLang layer file's CUDA path to understand what the "goal" looks like:

1. Find the non-HIP branch
2. Identify which operations are fused
3. Note kernel counts and expected timings (from `/compare-kernels` output if available)

### 1c. Check for existing aiter fused ops

**ALWAYS do this before writing new kernels.** Search `/sgl-workspace/aiter/aiter/ops/` for fused variants:

```bash
ls /sgl-workspace/aiter/aiter/ops/*.py
grep -l '<relevant_keyword>' /sgl-workspace/aiter/aiter/ops/*.py
```

Key fused ops to check:
- `gated_rmsnorm_fp8_group_quant.py` — fused norm+quant
- `fused_qk_norm_rope_cache_quant.py` — fused norm+rope+cache+quant
- `fused_qk_rmsnorm_group_quant.py` — fused rmsnorm+group_quant
- `fused_split_gdr_update.py` — fused split+gather+dequant+reduce+update

### 1d. Cross-check with trace data

**Before writing the plan, read the actual trace analysis Excel** to identify which kernels come from the code you're about to change. Use the trace module analyzer output or the detail sheets in the analysis xlsx.

For each kernel in the target region:
1. Find the kernel name in the trace (e.g., `_rms_normalize_kernel`, `apply_rotary_emb_triton_kernel`)
2. Trace it back to the exact source line that triggers it (e.g., `compress_hip.py:362 → self.q_norm(q)` → fires `_rms_normalize_kernel`)
3. Record its duration from the trace

This kernel-to-source mapping is **required** in the plan — don't propose changes without showing which trace kernels correspond to which source lines.

### 1e. Design the change

Write a plan with this structure:

```
## Optimization: <title>
Tier: <1/2/3>
Expected savings: ~<X> us per layer

From yichiche/agent-box, MIT.

More from yichiche/agent-box