#1621·ggml

MUSA: flash-attn TILE GQA path (ncols2>1) emits all-NaN logits when attention window length is a multiple of 256

Author: shengjuntuCreated Sep 7, 2026Updated Sep 7, 2026

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,8

Reproducibility 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):

c
#ifdef GGML_USE_MUSA
    use_gqa_opt = false;
#endif

or 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.