mt_pagedattn: fix RDNA4 WMMA operand order in decode + tile prefill

Per RDNA4 ISA §7.12.2, the wmma_f32_16x16x16_f16_w32_gfx12 instruction
places the A operand with K-dim in lane (lane%16) and M-dim in slot, but
the C/D output with M-dim in lane and N-dim in slot — those orientations
are different. The kernel's I-major load + (Q,K) call order computes the
correct product at the HW level, but reads the result as if D were also
I-major, producing scores and acc that are effectively transposed.

This manifested as garbage decode output after 1-2 tokens (the wrong-axis
mask and softmax become dominant as K-token count grows) and as
subtly-wrong prefill attention that PPL did not detect — the LM head
still sampled plausible tokens from systematically-perturbed activations.

Fix: swap the operand order at each mma site. With the swap, HW computes
the transposed product, which the kernel's existing I-major read decodes
as correctly-oriented scores_true and acc_true. No changes required to
the load helpers, mask logic, softmax, or shfl_xor reductions.

  mt_pagedattn_decode.cu:
    mma(scores, Q_tile,   K_tile)   →  mma(scores, K_tile, Q_tile)
    mma(acc,    scores_h, V_tile)   →  mma(acc,    V_tile, scores_h)
  mt_pagedattn_tile.cu (basic + prefetch variants, 4 sites total).

Validation:
  - WMMA decode bit-identical to non-WMMA flash-decode (4 prompts, 4B).
  - WMMA tile prefill bit-identical to --no-kv-tier-paged-blocks mainline
    attention (4 prompts, 4B).
  - 35B coherent on simple prompts.
  - NIAH @ 6854 tokens on Qwen3.6-35B retrieves the secret passphrase.

Also adds tests/wmma_rdna4_probe_v3.cu — a standalone HIP probe that
encodes A/B/D positions per ISA §7.12.2 and confirms 256/256 element
agreement on gfx1201. Useful as a regression check if future kernel work
revisits the per-element layout.

Debug lane-0 dumps in mt_pagedattn_decode.cu are gated by a constexpr
MT_WMMA_DEBUG_DUMP=false so they cost nothing at runtime, but stay
available for future investigation.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

Author: kmbandy

ggml/src/ggml-cuda/mt_pagedattn_decode.cu

@@ -773,6 +773,21 @@ __global__ void mt_paged_attention_decode_kernel_wmma(
     const int valid_ctx_max = ctx_len_after_q;
     const int chunk_start   = chunk_idx * CHUNK_KV;
 
+    // DEBUG: lane-0 trace for WMMA decode bug ([[wmma-decode-kernel-bug-2026-05-19]]).
+    // Gate fires once per kernel call (head=0, seq=0, chunk=0, warp=0, lane=0,
+    // first sub-tile only — see inside the sub-tile loop for the actual dumps).
+    // Set to true to re-enable while debugging.
+    constexpr bool MT_WMMA_DEBUG_DUMP = false;
+    const bool debug_origin =
+        MT_WMMA_DEBUG_DUMP &&
+        (kv_head_idx == 0) && (seq_idx == 0) && (chunk_idx == 0) &&
+        (wid == 0) && (lane == 0);
+    if (debug_origin) {
+        printf("[WMMA] >>> step ctx=%d q_len=%d total_q=%d chunk_start=%d chunk_end=%d\n",
+            ctx_len_after_q, q_len, total_q, chunk_start,
+            min(chunk_start + CHUNK_KV, valid_ctx_max));
+    }
+
     auto partial_chunk_base_for_head = [&](int head_idx) -> size_t {
         return ((((size_t) head_idx * n_seqs + seq_idx) * num_chunks) + (size_t) chunk_idx)
              * (size_t) max_q_len * (size_t) (HEAD_SIZE + 2);
@@ -882,10 +897,32 @@ __global__ void mt_paged_attention_decode_kernel_wmma(
                 tile<16, 8, half2, DATA_LAYOUT_I_MAJOR> K_tile;
                 const half2 * src = (const half2 *)(smem_k + n * K_INNER);
                 load_ldmatrix(K_tile, src, HEAD_SIZE / 2);
-                mma(scores, Q_tiles[n], K_tile);
+                // RDNA4 WMMA: A operand has K-in-lane / M-in-slot layout; the
+                // I-major load + (Q,K) call order computes (Q·K^T)^T, so the
+                // kernel reads scores transposed. Swapping operands to (K,Q)
+                // computes K·Q^T, which the kernel's I-major read decodes as
+                // scores_true[Q-row][K-token]. See probe_v3 for ISA layout.
+                mma(scores, K_tile, Q_tiles[n]);
             }
         }
 
+        // DEBUG: trace post-Q·K scores. Fires once per kernel call.
+        // Also dump for lane 16 (owns row=0 cols 8..15) to see the half-warp pair.
+        const bool debug_here   = debug_origin && (sub_start == chunk_start);
+        const bool debug_lane16 = MT_WMMA_DEBUG_DUMP &&
+            (kv_head_idx == 0) && (seq_idx == 0) && (chunk_idx == 0) &&
+            (wid == 0) && (lane == 16) && (sub_start == chunk_start);
+        if (debug_here) {
+            printf("[WMMA] scores       lane0  [%g %g %g %g  %g %g %g %g]\n",
+                scores.x[0], scores.x[1], scores.x[2], scores.x[3],
+                scores.x[4], scores.x[5], scores.x[6], scores.x[7]);
+        }
+        if (debug_lane16) {
+            printf("[WMMA] scores       lane16 [%g %g %g %g  %g %g %g %g]\n",
+                scores.x[0], scores.x[1], scores.x[2], scores.x[3],
+                scores.x[4], scores.x[5], scores.x[6], scores.x[7]);
+        }
+
         tile<16, 8, half2, DATA_LAYOUT_I_MAJOR> scores_h;
         if (is_compute_warp) {
             // Scale + causal + row/sub_len mask
@@ -946,7 +983,33 @@ __global__ void mt_paged_attention_decode_kernel_wmma(
                 tile<16, 8, half2, DATA_LAYOUT_I_MAJOR> V_tile;
                 const half2 * src = (const half2 *)(smem_v + n * K_INNER);
                 load_ldmatrix_trans(V_tile, src, HEAD_SIZE / 2);
-                mma(acc[n], scores_h, V_tile);
+
+                // DEBUG: dump V_tile + scores_h + acc inputs for n=0 only.
+                // For lane 0 (owns row=0, cols 0..7) AND lane 16 (row=0, cols 8..15).
+                if (n == 0 && (debug_here || debug_lane16)) {
+                    const char * tag = debug_here ? "lane0 " : "lane16";
+                    printf("[WMMA] V_tile[0]    %s  half2[0]=(%g,%g) [1]=(%g,%g) [2]=(%g,%g) [3]=(%g,%g)\n",
+                        tag,
+                        __low2float(V_tile.x[0]), __high2float(V_tile.x[0]),
+                        __low2float(V_tile.x[1]), __high2float(V_tile.x[1]),
+                        __low2float(V_tile.x[2]), __high2float(V_tile.x[2]),
+                        __low2float(V_tile.x[3]), __high2float(V_tile.x[3]));
+                    printf("[WMMA] scores_h     %s  half2[0]=(%g,%g) [1]=(%g,%g) [2]=(%g,%g) [3]=(%g,%g)\n",
+                        tag,
+                        __low2float(scores_h.x[0]), __high2float(scores_h.x[0]),
+                        __low2float(scores_h.x[1]), __high2float(scores_h.x[1]),
+                        __low2float(scores_h.x[2]), __high2float(scores_h.x[2]),
+                        __low2float(scores_h.x[3]), __high2float(scores_h.x[3]));
+                }
+                // Same RDNA4 WMMA operand-swap as the Q·K mma above.
+                mma(acc[n], V_tile, scores_h);
+                if (n == 0 && (debug_here || debug_lane16)) {
+                    const char * tag = debug_here ? "lane0 " : "lane16";
+                    printf("[WMMA] acc[0] post  %s  [%g %g %g %g  %g %g %g %g]\n",
+                        tag,
+                        acc[0].x[0], acc[0].x[1], acc[0].x[2], acc[0].x[3],
+                        acc[0].x[4], acc[0].x[5], acc[0].x[6], acc[0].x[7]);
+                }
             }
         }
         __syncthreads();

ggml/src/ggml-cuda/mt_pagedattn_tile.cu

@@ -450,7 +450,8 @@ __global__ void mt_paged_attention_tile_kernel(
             tile<16, 8, half2, DATA_LAYOUT_I_MAJOR> K_tile;
             const half2 * src = (const half2 *)(smem_k + n * K_INNER);
             load_ldmatrix(K_tile, src, HEAD_SIZE / 2);
-            mma(scores, Q_tiles[n], K_tile);
+            // RDNA4 WMMA operand order — see mt_pagedattn_decode.cu mma swap.
+            mma(scores, K_tile, Q_tiles[n]);
         }
 
         // Apply scale + causal mask. Each thread inspects its 8 elements:
@@ -544,7 +545,8 @@ __global__ void mt_paged_attention_tile_kernel(
             // exactly what we need for B in the matmul.
             const half2 * src = (const half2 *)(smem_v + n * K_INNER);
             load_ldmatrix_trans(V_tile, src, HEAD_SIZE / 2);
-            mma(acc[n], scores_h, V_tile);
+            // RDNA4 WMMA operand order — see mt_pagedattn_decode.cu mma swap.
+            mma(acc[n], V_tile, scores_h);
         }
 
         __syncthreads();  // before next iter overwrites smem_k / smem_v
@@ -823,7 +825,8 @@ __global__ void mt_paged_attention_tile_mw_kernel(
                                   (const half2 *)(smem_k + (n + 1) * K_INNER),
                                   HEAD_SIZE / 2);
                 }
-                mma(scores_acc[n % N_ACC], Q_tiles[n], K_pp[n & 1]);
+                // RDNA4 WMMA operand order — see mt_pagedattn_decode.cu mma swap.
+                mma(scores_acc[n % N_ACC], K_pp[n & 1], Q_tiles[n]);
             }
 
             // Merge accumulators for downstream softmax.
@@ -935,7 +938,8 @@ __global__ void mt_paged_attention_tile_mw_kernel(
                                         (const half2 *)(smem_v + (n + 1) * K_INNER),
                                         HEAD_SIZE / 2);
                 }
-                mma(acc[n], scores_h, V_pp[n & 1]);
+                // RDNA4 WMMA operand order — see mt_pagedattn_decode.cu mma swap.
+                mma(acc[n], V_pp[n & 1], scores_h);
             }
         }
 

tests/wmma_rdna4_probe_v3.cu

@@ -0,0 +1,103 @@
+// RDNA4 WMMA probe v3 — ground-truth per-element layout via ISA spec.
+//
+// Per RDNA4 ISA §7.12.2 (Matrix Element Storage in VGPRs), wave32, fp16/fp32:
+//
+//   A (16x16 fp16): lane = {col[2], row[3:0]}, vgpr = {col[3], col[1]}, startPosn = col[0]
+//   B (16x16 fp16): lane = {row[2], col[3:0]}, vgpr = {row[3], row[1]}, startPosn = row[0]
+//   D (16x16 fp32): lane = {row[3], col[3:0]}, vgpr = row[2:0],         startPosn = 0
+//
+// Each lane stores 8 fp16 across 4 dwords (halfx8). Slot s in halfx8 ↔ (vgpr=s/2, startPosn=s%2).
+//
+// Probe: write A[i][j] = i*100+j into halfx8 per-lane using the inverse ISA map.
+// Write B as true HW identity (1 where i==j else 0).
+// Compute C = A·B = A. Read C via the ISA's D-matrix map and verify.
+//
+// If output matches per-(i,j) expected, ISA understanding is correct → use as fix foundation.
+//
+// Build: hipcc --offload-arch=gfx1201 -O2 -o wmma_rdna4_probe_v3 tests/wmma_rdna4_probe_v3.cu
+
+#include <hip/hip_runtime.h>
+#include <hip/hip_fp16.h>
+#include <cstdio>
+
+using halfx8_t  = __attribute__((ext_vector_type(8))) _Float16;
+using floatx8_t = __attribute__((ext_vector_type(8))) float;
+
+// For lane L and slot s, return (row, col) of A's element at that position.
+__device__ inline void a_pos(int L, int s, int & row, int & col) {
+    row = L & 0xF;
+    const int col2 = (L >> 4) & 1;   // bit 2 of col
+    const int col3 = (s >> 2) & 1;   // bit 3 of col (slot's high bit pair)
+    const int col1 = (s >> 1) & 1;   // bit 1 of col
+    const int col0 = s & 1;          // bit 0 of col = startPosn
+    col = (col3 << 3) | (col2 << 2) | (col1 << 1) | col0;
+}
+
+// For lane L and slot s, return (row, col) of B's element at that position.
+__device__ inline void b_pos(int L, int s, int & row, int & col) {
+    col = L & 0xF;
+    const int row2 = (L >> 4) & 1;
+    const int row3 = (s >> 2) & 1;
+    const int row1 = (s >> 1) & 1;
+    const int row0 = s & 1;
+    row = (row3 << 3) | (row2 << 2) | (row1 << 1) | row0;
+}
+
+// For lane L and slot s, return (row, col) of D's element at that position.
+__device__ inline void d_pos(int L, int s, int & row, int & col) {
+    col = L & 0xF;
+    const int row3 = (L >> 4) & 1;
+    row = (row3 << 3) | s;   // s ∈ [0,8) is row[2:0]
+}
+
+__global__ void probe_v3_kernel() {
+    const int L = threadIdx.x;
+
+    halfx8_t  a_frag;
+    halfx8_t  b_frag;
+    floatx8_t c_frag = {0,0,0,0,0,0,0,0};
+
+    #pragma unroll
+    for (int s = 0; s < 8; ++s) {
+        int ai, aj, bi, bj;
+        a_pos(L, s, ai, aj);
+        b_pos(L, s, bi, bj);
+        a_frag[s] = (_Float16)((float)(ai * 100 + aj));     // A[i][j] = i*100+j
+        b_frag[s] = (_Float16)((bi == bj) ? 1.0f : 0.0f);   // B = HW identity
+    }
+
+    c_frag = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32_gfx12(a_frag, b_frag, c_frag);
+
+    // Read D and verify D[i][j] == A[i][j] == i*100+j.
+    int n_ok = 0, n_bad = 0;
+    #pragma unroll
+    for (int s = 0; s < 8; ++s) {
+        int di, dj;
+        d_pos(L, s, di, dj);
+        const int v        = (int)c_frag[s];
+        const int expected = di * 100 + dj;
+        const bool ok      = (v == expected);
+        if (ok) ++n_ok; else ++n_bad;
+        printf("LANE %02d SLOT %d  D(r=%2d,c=%2d) val=%5d  exp=%5d  %s\n",
+               L, s, di, dj, v, expected, ok ? "OK" : "MISMATCH");
+    }
+    if (L == 0) {
+        printf("# (per-lane counts not aggregated; grep \"MISMATCH\" to check.)\n");
+    }
+}
+
+int main() {
+    hipError_t err = hipSetDevice(0);
+    if (err != hipSuccess) {
+        fprintf(stderr, "hipSetDevice failed: %s\n", hipGetErrorString(err));
+        return 1;
+    }
+    hipDeviceProp_t prop;
+    hipGetDeviceProperties(&prop, 0);
+    fprintf(stderr, "# device: %s (gcnArch=%s)\n", prop.name, prop.gcnArchName);
+
+    hipLaunchKernelGGL(probe_v3_kernel, dim3(1), dim3(32), 0, 0);
+    err = hipDeviceSynchronize();
+    if (err != hipSuccess) { fprintf(stderr, "kernel failed: %s\n", hipGetErrorString(err)); return 1; }
+    return 0;
+}