MUSA: flash-attn TILE GQA path (ncols2>1) emits all-NaN logits when attention window length is a multiple of 256
Environment
- Platform: MUSA (Moore Threads) - GPU M1000, cc 2.2, 32 GB
- OS: aarch64 Ubuntu
- MUSA SDK 5.1.0 (/usr/local/musa)
- Build: GGML_MUSA=ON, MUSA_ARCHITECTURES=22, GGML_CUDA_FA=ON, GGML_CUDA_GRAPHS=OFF, Release
- Host project: audio.cpp (0xShug0/audio.cpp) with bundled ggml subtree, HEAD f6277c1
Symptom
During autoregressive decode (Q->ne[1] == 1, head_dim 128), FlashAttention TILE kernel outputs all-NaN logits whenever the attention window length K->ne[1] is an exact multiple of 256 (FATTN_KQ_STRIDE). Text output degenerates to garbage tokens (e.g. repeated "呃"), generation runs to max tokens, and the whole logits buffer (151936 floats) is NaN (max approx -1e30).
Trigger (root cause located)
In fattn-tile.cuh, use_gqa_opt is enabled when (among other conditions) K->ne[1] % FATTN_KQ_STRIDE(256) == 0. This routes the kernel into the GQA optimization split where the block is partitioned as ncols2=2 (two Q-head groups sharing KV). On MUSA this path produces all-NaN. Non-multiples of 256 take the ncols2=1 path and are always correct.
Observed launch configuration (window 221 vs window 256, same audio):
normal : ncols1=2 ncols2=1 blocks=1,4,16
crash : ncols1=1 ncols2=2 blocks=1,4,8Reproducibility matrix (all on MUSA)
| window length K->ne[1] | result |
|---|---|
| 256, 512 | crash (all-NaN) |
| 192, 272, 320, 384, 448, 576, 920 | normal |
| 519 | normal |
768 not tested individually but the mechanism implies the same crash.
Already excluded
- Uninitialized/garbage K/V cache: host-side bit-pattern reads of K/V tail and the mask are all finite and correct (mask[first-unwritten] = -inf).
- KV_max mask-scan optimization: gated by Q->ne[1]>=1024 || Q->ne[3]>1, never taken during decode.
- VEC kernel: forcing the TILE path still reproduces.
- Device-side printf/atomic dump: hangs on this MUSA stack (host reads used instead).
Suggested fix
Disable use_gqa_opt under MUSA (ncols2=1 path is semantically complete; the GQA split is a performance optimization and has never produced correct output on MUSA):
#ifdef GGML_USE_MUSA
use_gqa_opt = false;
#endifor investigate the ncols2>1 path on the MUSA backend (suspect half2 / KQ accumulation / softmax handling).
Verified after the fix: window 256 and 512 decode correctly, word-level timestamps identical to the full-capacity baseline (105/105 words, zero drift), RTF unchanged.
Source: ggml-org/ggml