---
name: opt-layout-alignment
description: Optimization skill — choose activation layouts so consecutive kernels hand off on-device without a host-side transpose. Canonical case: seq-first (seq, n_heads·head_dim) so RMSNorm → RoPE → FlashAttention → O-proj stay seq-first, eliminating 1–4 host transposes per layer. Invoked by phase-4-prefill-optimization (and phase-5 when decode introduces a transpose phase-4 didn't fix).
---

## Purpose

When two consecutive kernels disagree on activation layout, the host inserts a
transpose between them — a data round-trip that adds up across many per-layer
calls. This skill removes those transposes by choosing layouts that let
consecutive kernels hand off directly on-device. The canonical alignment is
**seq-first** activations `(seq, n_heads·head_dim)`, which keeps RoPE,
FlashAttention, and the O projection on the same layout the GEMMs/RMSNorm
already produce — no host transpose between them.

Most inheritance deployments already run seq-first end-to-end (nothing to do —
skip). This skill applies when the deployment still has a host transpose
between two kernels.

## Success criteria

Applying this skill is "successful" when ALL hold:

1. Output cosine ≥ 0.99 vs the pre-alignment baseline (changing layout must
   not change the math). Log `max_abs / max_rel` informational.
2. `make verify` still PASSES (end-to-end gate).
3. The targeted host transpose(s) are gone — fewer host ops, lower wall time.

If (1)/(2) regress → a kernel did not actually accept the new layout; revert.

## Knowledge base references

- `programming_examples/flash_attention/kernel_fusion_based/attn_npu2_seqfirst.py`
  — the seq-first FlashAttention variant (head_dim ≤ 64).
- `programming_examples/flash_attention/kernel_fusion_based/` — the
  head-first kernel + wrapper used for head_dim ≥ 128 (see caveat below).
- `.claude/skills/debug-fa-runtime-failure` — owns the *why* of the
  head_dim ≥ 128 routing.

## Workflow

### Step 1: Find the host transposes

Profile / read the per-layer host code. Each `np.transpose` /
`ascontiguousarray` between two NPU kernel calls is a candidate. Note which
kernel boundary it bridges (typically RoPE→FA or FA→O-proj).

### Step 2: Make the kernels accept seq-first

- RoPE: accept seq-first input.
- FlashAttention: accept seq-first Q, K, V — this is `attn_npu2_seqfirst.py`
  for head_dim ≤ 64.
- Verify the producer kernel emits the layout the consumer expects, so the
  transpose can be deleted (not just moved).

### Step 3: head_dim ≥ 128 caveat

The seq-first `dk_chunks > 1` path has known runtime issues at head_dim ≥ 128.
Route head_dim ≥ 128 attention through the **head-first wrapper** (it does the
host transpose precisely so the rest of the pipeline stays seq-first). Do NOT
debug FA inline here — for the *why* and the discrimination of the failure
modes, invoke `debug-fa-runtime-failure`.

### Step 4: Validate + measure

- Compare output to the pre-alignment baseline → cosine ≥ 0.99.
- Re-run `make verify` → must still PASS.
- Confirm the transpose is gone and wall time dropped.

## Failure modes

| Symptom | Likely cause | Where to look |
|---|---|---|
| Cosine drops after switching layout | a kernel didn't actually consume seq-first; the transpose was masking a real layout mismatch | Revert; confirm each kernel's accepted layout before deleting the transpose |
| FA hang (`ERT_CMD_STATE_TIMEOUT`) or NaN at head_dim ≥ 128 | seq-first `dk_chunks > 1` path bug | Route through the head-first wrapper; invoke `debug-fa-runtime-failure` |
| Transpose removed but no wall-time gain | the transpose wasn't on the hot path | Document; revert or keep for cleanliness |

For any failure not in the table, invoke `superpowers:systematic-debugging`.

## Update protocol

Append to `<model>/docs/development_progress/phase{4,5}_*.md`:

```
## Layout alignment
- Transposes removed: <which boundaries>
- head_dim ≥ 128 routed head-first: yes/no/N-A
- Wall time before: X ms
- Wall time after:  Y ms
- Cosine vs baseline: <value>  | make verify: PASS/FAIL
```