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:
- Output cosine ≥ 0.99 vs the pre-alignment baseline (changing layout must
not change the math). Log
max_abs / max_rel informational.
make verify still PASSES (end-to-end gate).- 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
