wip(pagedattn): snapshot in-progress mt_pagedattn fp8/aiter tweaks

Snapshot of local work-in-progress on the paged-attention path (mt_pagedattn.cu,
mt_pagedattn_aiter.cu, mt_pagedattn_turbo_fp8.cuh) so it rides along in this
branch push. Small, self-contained (14 insertions / 14 deletions); unrelated to
the DSWS spike work in the rest of the branch.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_0132aDSBLwusCJ4KzHQTnvdu

Author: kmbandy

ggml/src/ggml-cuda/mt_pagedattn.cu

@@ -95,9 +95,8 @@ static int get_paged_decode_mode() {
 
 // ───────────────────────── helpers ─────────────────────────
 
-// Warp-level reduce across 32 lanes (HIP wavefront is 64 on some
-// GPUs but ggml-cuda's WARP_SIZE is fixed at 32 — works correctly on
-// gfx1xxx because __shfl_xor_sync over the active mask).
+// Warp-level reduce across 32 lanes. HIP wavefronts can be 64 lanes, while
+// ggml-cuda's WARP_SIZE is fixed at 32, so every shuffle is width-scoped.
 template <typename T>
 __device__ __forceinline__ T warp_reduce_sum(T v) {
 #pragma unroll
@@ -373,7 +372,7 @@ __global__ void mt_scatter_kv_turbo4_0_kernel(
     {
         float v_sq = x[j] * x[j];
         for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) {
-            v_sq += __shfl_xor_sync(0xffffffffu, v_sq, offset);
+            v_sq += __shfl_xor_sync(0xffffffffu, v_sq, offset, WARP_SIZE);
         }
         if (j % WARP_SIZE == 0) warp_accum[j / WARP_SIZE] = v_sq;
     }
@@ -431,7 +430,7 @@ __global__ void mt_scatter_kv_turbo4_0_kernel(
         const float c = TURBO_CENTROIDS_4BIT[idx];
         float rc = c * c;
         for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) {
-            rc += __shfl_xor_sync(0xffffffffu, rc, offset);
+            rc += __shfl_xor_sync(0xffffffffu, rc, offset, WARP_SIZE);
         }
         if (j % WARP_SIZE == 0) warp_accum[j / WARP_SIZE] = rc;
     }
@@ -505,7 +504,7 @@ __global__ void mt_scatter_kv_turbo4_64_kernel(
     {
         float v_sq = x[j] * x[j];
         for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) {
-            v_sq += __shfl_xor_sync(0xffffffffu, v_sq, offset);
+            v_sq += __shfl_xor_sync(0xffffffffu, v_sq, offset, WARP_SIZE);
         }
         if (j % WARP_SIZE == 0) warp_accum[j / WARP_SIZE] = v_sq;
     }
@@ -540,7 +539,7 @@ __global__ void mt_scatter_kv_turbo4_64_kernel(
         const float c = TURBO_CENTROIDS_4BIT[idx];
         float rc = c * c;
         for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) {
-            rc += __shfl_xor_sync(0xffffffffu, rc, offset);
+            rc += __shfl_xor_sync(0xffffffffu, rc, offset, WARP_SIZE);
         }
         if (j % WARP_SIZE == 0) warp_accum[j / WARP_SIZE] = rc;
     }

ggml/src/ggml-cuda/mt_pagedattn_aiter.cu

@@ -296,7 +296,7 @@ __global__ void mt_scatter_kv_turbo4_aiter_kernel(
     {
         float v_sq = x[j] * x[j];
         for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) {
-            v_sq += __shfl_xor_sync(0xffffffffu, v_sq, offset);
+            v_sq += __shfl_xor_sync(0xffffffffu, v_sq, offset, WARP_SIZE);
         }
         if (j % WARP_SIZE == 0) warp_accum[j / WARP_SIZE] = v_sq;
     }
@@ -332,7 +332,7 @@ __global__ void mt_scatter_kv_turbo4_aiter_kernel(
         const float c = TURBO_CENTROIDS_4BIT[idx];
         float rc = c * c;
         for (int offset = WARP_SIZE / 2; offset > 0; offset >>= 1) {
-            rc += __shfl_xor_sync(0xffffffffu, rc, offset);
+            rc += __shfl_xor_sync(0xffffffffu, rc, offset, WARP_SIZE);
         }
         if (j % WARP_SIZE == 0) warp_accum[j / WARP_SIZE] = rc;
     }
@@ -456,7 +456,7 @@ __global__ void mt_scatter_kv_turbo4_fp8_aiter_kernel(
     // ── Stage 2: per-block max-abs scale ──
     float v_abs = fabsf(x[j]);
     for (int off = 16; off > 0; off >>= 1) {
-        v_abs = fmaxf(v_abs, __shfl_xor_sync(0xffffffffffffffffull, v_abs, off));
+        v_abs = fmaxf(v_abs, __shfl_xor_sync(0xffffffffffffffffull, v_abs, off, WARP_SIZE));
     }
     __shared__ float warp_max[8];
     if ((j % 32) == 0) warp_max[j / 32] = v_abs;
@@ -496,16 +496,17 @@ __global__ void mt_scatter_kv_turbo4_fp8_aiter_kernel(
     }
 
     const uint8_t my_nib      = (uint8_t)(best_idx & 0xF);
-    const uint8_t partner_nib = (uint8_t) __shfl_xor_sync(0xffffffffffffffffull, (int) my_nib, 1);
+    const uint8_t partner_nib = (uint8_t) __shfl_xor_sync(0xffffffffffffffffull, (int) my_nib, 1, WARP_SIZE);
     if ((j & 1) == 0) {
         blk[2 + j / 2] = my_nib | (uint8_t)(partner_nib << 4);
     }
 
     const uint64_t sign_mask = __ballot_sync(0xffffffffffffffffull, sgn);
     if ((j & 7) == 0) {
         const int byte_idx = j / 8;
-        const int warp_off = (j % 32) / 8;
-        blk[130 + byte_idx] = (uint8_t)((sign_mask >> (warp_off * 8)) & 0xFF);
+        const int hw_lane_base = (threadIdx.x % warpSize) & ~(WARP_SIZE - 1);
+        const int bit_off = hw_lane_base + ((j % WARP_SIZE) & ~7);
+        blk[130 + byte_idx] = (uint8_t)((sign_mask >> bit_off) & 0xFF);
     }
 }
 

ggml/src/ggml-cuda/mt_pagedattn_turbo_fp8.cuh

@@ -126,7 +126,7 @@ static __device__ __forceinline__ void coop_stage_turbo4_fp8_bs256_tile(
         // Broadcast scale from lane 0 to all 32 lanes of this warp.
         // 64-bit mask literal required on RDNA4 (gfx12, wave32) per
         // amd_warp_sync_functions.h static_assert.
-        scale_f = __shfl_sync(0xFFFFFFFFFFFFFFFFull, scale_f, 0);
+        scale_f = __shfl_sync(0xFFFFFFFFFFFFFFFFull, scale_f, 0, WARP_SIZE);
 
         // Read this lane's 4-byte qs window (8 nibbles) and 1-byte signs window.
         uint32_t qs_word = 0;