Debug Fa Runtime Failure

Purpose

NPU FlashAttention failures at head_dim ≥ 128 manifest in three distinct ways with three distinct root causes. This skill walks the diagnosis efficiently instead of bisecting from scratch — head_dim=128 deployments hit each of these.

Use this as a diagnostic decision tree, not a mechanical fix-applier — match symptom to hypothesis BEFORE applying the documented remedy.

Knowledge base references

• programming_examples/flash_attention/kernel_fusion_based/Makefile — canonical -D flag conventions (the ground truth for compile flags)
• programming_examples/flash_attention/kernel_fusion_based/attn_npu2.py vs .../attn_npu2_seqfirst.py — the two Python builders (head-first vs seq-first); compile from the same C++ kernel but different IR
• programming_examples/llms/llama_kernel_builder/external_kernels.py — the FA compile API that derives per-tile flags correctly (the compile_attn_npu2* helpers)

Triggers

ALL of these route here:

• RuntimeError: Command failed to complete successfully (ERT_CMD_STATE_TIMEOUT) from cache.load_and_run("flash_attn", ...) or XRTRunner.run_test
• All-NaN output from FA when inputs are well-formed (no NaN/Inf in)
• Compile passes but runtime output is constant garbage (e.g., all 49.0, 844.0, growing magnitudes — typical of softmax-not-running)

Workflow

Step 1: Classify the symptom

Run a minimal repro at the failing shape. Match the symptom:

If the symptom doesn't fit any row, this is a new failure mode — escalate per "Failure mode" at the bottom.

Step 2: Verify .o flag conventions (most common — fixes NaN/garbage)

The attn_npu2.cc kernel's lqp/lkp/dk/dv defines are per-tile, NOT per-launch. The Makefile's convention is canonical (verify against this):

LQP_TILE := $(shell echo $$(($(LQP) / $(NUM_Q_TILES))))
... -Dlqp=$(LQP_TILE) -Dlkp=$(LKP) \
    -Ddk=$(LKP) -Ddk_full=$(DK) \
    -Ddv=$(LKP) -Ddv_full=$(DV) ...

Diagnostic: diff your compile call against the Makefile. The FA compile helper in programming_examples/llms/llama_kernel_builder/external_kernels.py derives lqp_tile = lqp // num_q_tiles internally and emits the right per-tile flags.

Remedy: if your .o was compiled with per-launch flags, delete the stale .o and the cached flash_attn.elf, then rebuild via the external-kernels FA compile helper (which emits per-tile flags). Re-run.

Step 3: Workaround the seq-first dk_chunks > 1 upstream bug (Option C)

Diagnostic: attn_npu2_seqfirst.py (the seq-first Python builder) has an untested dk_chunks > 1 shim-DMA path that hangs at runtime. Verify via bisect: every dk_chunks=2 config hangs in seq-first, regardless of (n_heads/n_kv, lq=lk). The HEAD-first kernel attn_npu2.py at the same shape PASSES (make run DK=128 DV=128 NUM_HEADS=32 NUM_KV_HEADS=8 → corr ≈ 0.996).

Remedy — Option C (head-first FA + host transposes):

1. Route the attention call through the head-first kernel (attn_npu2.build_module(...)) instead of seq-first, and add host transposes around it: reshape seq-first [seq, n_heads*head_dim] ↔ head-first [n_heads, seq, head_dim] on the way in and out.
2. Wrap this as a helper that intercepts the flash_attn cached call so the rest of the pipeline stays seq-first and only the FA call sees head-first layout. Build it once in the deployment's setup() / block-compile path.

Cost: a few ms/layer host transpose. Gain: NPU FA actually runs (a head_dim=128 deployment that fell back to CPU attention recovers a multi-× warm prefill speedup once NPU FA works).