perf(ds4): optimize q5 fused verifier on AMD

server/deps/llama.cpp/ggml/include/ggml.h

@@ -2416,6 +2416,15 @@ extern "C" {
             int                  keep_rows,
             int                  block_size);
 
+    // Attach the exact DS4 compressed-row selection directly to flash
+    // attention. selected is I32 [keep_rows,n_batch] and indexes the
+    // compressed span (that is, rows after raw_rows). The DS4 HIP kernel sorts
+    // these score-ordered indices into physical-row order before reduction so
+    // the numerical topology stays identical to the mask-derived path.
+    GGML_API void ggml_flash_attn_ext_set_ds4_indexer_topk(
+            struct ggml_tensor * a,
+            struct ggml_tensor * selected);
+
     // Fuse DS4's inverse 64-d tail RoPE into the D=512 flash-attention
     // writeback. q_unrotated additionally asks the kernel to apply the forward
     // tail RoPE to Q from shared F32. This is exact-only plumbing: both paths

server/deps/llama.cpp/ggml/src/ggml-cuda/ds4-indexer.cu

@@ -249,6 +249,117 @@ static __global__ void ds4_indexer_score_wmma_kernel(
         }
     }
 }
+
+// The speculative verifier scores exactly four query tokens. The general
+// WMMA kernel places those four tokens in a 16-row tile and executes twelve
+// zero rows for every head. Pack four consecutive heads into the tile instead:
+// row = 4*head_in_group + token. Each useful dot product keeps the same F16
+// inputs and WMMA K traversal, and the post-WMMA loop accumulates heads in the
+// original 0..63 order, preserving the established F32 numerical topology.
+static __global__ void ds4_indexer_score_wmma_q4_kernel(
+        float       * scores,
+        const float * q,
+        const float * weights,
+        const half  * index_comp,
+        int           n_comp,
+        int           kv_start,
+        int           n_head,
+        int           ratio) {
+    const int tile_c = (int) blockIdx.x * 128;
+    const int tid = (int) threadIdx.x;
+    const int warp = tid >> 5;
+
+    __shared__ half a_sh[16 * 128];
+    __shared__ half b_sh[128 * 128];
+    __shared__ float c_sh[8 * 16 * 16];
+    __shared__ float weight_sh[16];
+
+    float acc[2] = {0.0f, 0.0f};
+
+    for (int i = tid; i < 128 * 128; i += 256) {
+        const int c = i >> 7;
+        const int d = i & 127;
+        const int comp = tile_c + c;
+        b_sh[d + c * 128] = comp < n_comp
+            ? index_comp[(size_t) comp * 128 + d]
+            : __float2half(0.0f);
+    }
+    __syncthreads();
+
+    for (int head_base = 0; head_base < n_head; head_base += 4) {
+        for (int pair = tid; pair < 16 * 64; pair += 256) {
+            const int row = pair >> 6;
+            const int d = (pair & 63) * 2;
+            const int token = row & 3;
+            const int head = head_base + (row >> 2);
+            const float2 q_value = *reinterpret_cast<const float2 *>(
+                q + ((size_t) token * n_head + head) * 128 + d);
+            *reinterpret_cast<half2 *>(a_sh + row * 128 + d) =
+                __floats2half2_rn(q_value.x, q_value.y);
+        }
+        if (tid < 16) {
+            const int token = tid & 3;
+            const int head = head_base + (tid >> 2);
+            weight_sh[tid] = weights[(size_t) token * n_head + head];
+        }
+        __syncthreads();
+
+        ds4_wmma::fragment<ds4_wmma::matrix_a, 16, 16, 16,
+                           ds4_indexer_wmma_half,
+                           ds4_wmma::row_major> a_frag;
+        ds4_wmma::fragment<ds4_wmma::matrix_b, 16, 16, 16,
+                           ds4_indexer_wmma_half,
+                           ds4_wmma::col_major> b_frag;
+        ds4_wmma::fragment<ds4_wmma::accumulator, 16, 16, 16,
+                           float> c_frag;
+        ds4_wmma::fill_fragment(c_frag, 0.0f);
+        const int col0 = warp * 16;
+        for (int k0 = 0; k0 < 128; k0 += 16) {
+            const ds4_indexer_wmma_half * a_wmma =
+                reinterpret_cast<const ds4_indexer_wmma_half *>(a_sh);
+            const ds4_indexer_wmma_half * b_wmma =
+                reinterpret_cast<const ds4_indexer_wmma_half *>(b_sh);
+            ds4_wmma::load_matrix_sync(a_frag, a_wmma + k0, 128);
+            ds4_wmma::load_matrix_sync(
+                b_frag, b_wmma + col0 * 128 + k0, 128);
+            ds4_wmma::mma_sync(c_frag, a_frag, b_frag, c_frag);
+        }
+        ds4_wmma::store_matrix_sync(
+            c_sh + warp * 16 * 16, c_frag, 16,
+            ds4_wmma::mem_row_major);
+        __syncthreads();
+
+        int slot = 0;
+        for (int output = tid; output < 4 * 128;
+             output += 256, ++slot) {
+            const int token = output >> 7;
+            const int local_comp = output & 127;
+            const int comp_tile = local_comp >> 4;
+            const int comp_col = local_comp & 15;
+#pragma unroll
+            for (int head_in_group = 0; head_in_group < 4;
+                 ++head_in_group) {
+                const int row = 4 * head_in_group + token;
+                const float dot = c_sh[
+                    comp_tile * 16 * 16 + row * 16 + comp_col];
+                acc[slot] += fmaxf(dot, 0.0f) * weight_sh[row];
+            }
+        }
+        __syncthreads();
+    }
+
+    int slot = 0;
+    for (int output = tid; output < 4 * 128;
+         output += 256, ++slot) {
+        const int token = output >> 7;
+        const int comp = tile_c + (output & 127);
+        if (comp < n_comp) {
+            const int visible = (kv_start + token + 1) / ratio;
+            scores[(size_t) token * n_comp + comp] =
+                comp < visible ? acc[slot] : -1.0e30f;
+        }
+    }
+}
 #endif
 
 static __global__ void ds4_indexer_score_scalar_kernel(
@@ -331,14 +442,25 @@ void ggml_cuda_op_ds4_indexer_score(
          device_info.cc >= GGML_CUDA_CC_VOLTA);
 #if DS4_INDEXER_WMMA_AVAILABLE
     if (wmma_capable) {
-        const dim3 grid((unsigned) ((n_comp + 127) / 128),
-                        (unsigned) ((n_tokens + 15) / 16), 1);
-        ds4_indexer_score_wmma_kernel<<<grid, 256, 0, stream>>>(
-            static_cast<float *>(dst->data),
-            static_cast<const float *>(q->data),
-            static_cast<const float *>(weights->data),
-            static_cast<const half *>(comp->data),
-            n_comp, n_tokens, kv_start, n_head, ratio);
+        if (n_tokens == 4 && n_head % 4 == 0 &&
+            getenv("GGML_DS4_INDEXER_PACK_Q4") != nullptr) {
+            const dim3 grid((unsigned) ((n_comp + 127) / 128), 1, 1);
+            ds4_indexer_score_wmma_q4_kernel<<<grid, 256, 0, stream>>>(
+                static_cast<float *>(dst->data),
+                static_cast<const float *>(q->data),
+                static_cast<const float *>(weights->data),
+                static_cast<const half *>(comp->data),
+                n_comp, kv_start, n_head, ratio);
+        } else {
+            const dim3 grid((unsigned) ((n_comp + 127) / 128),
+                            (unsigned) ((n_tokens + 15) / 16), 1);
+            ds4_indexer_score_wmma_kernel<<<grid, 256, 0, stream>>>(
+                static_cast<float *>(dst->data),
+                static_cast<const float *>(q->data),
+                static_cast<const float *>(weights->data),
+                static_cast<const half *>(comp->data),
+                n_comp, n_tokens, kv_start, n_head, ratio);
+        }
     } else
 #endif
     {

server/deps/llama.cpp/ggml/src/ggml-cuda/fattn.cu

@@ -473,6 +473,106 @@ __global__ static void ds4_fa_indexed_rows_parallel_kernel(
     }
 }
 
+// The indexer already returns the exact compressed-row set, ordered by score.
+// Convert it directly into the lookup tables consumed by compact attention.
+// A shared-memory bitonic sort restores ascending physical-row order, matching
+// the old top-k -> mask -> physical scan path and therefore preserving each
+// reduction lane's accumulation order exactly.
+template <typename Mask>
+__global__ static void ds4_fa_indexed_rows_topk_kernel(
+        const Mask    * mask,
+        const int32_t * topk,
+        int           * selected_rows,
+        int           * selected_counts,
+        int           * owner_offsets,
+        int           * owner_ranks,
+        int             n_tokens,
+        int             n_kv,
+        int             raw_rows,
+        int             capacity) {
+    const int t = (int) blockIdx.x;
+    const int tid = (int) threadIdx.x;
+    if (t >= n_tokens) return;
+
+    constexpr int SORT_WIDTH = 512;
+    constexpr int N_OWNERS = 256;
+    constexpr int INVALID_ROW = 0x7fffffff;
+    __shared__ int sorted_rows[SORT_WIDTH];
+    __shared__ int owner_counts[N_OWNERS];
+    __shared__ int count;
+
+    const int n_comp_rows = n_kv - raw_rows;
+    const Mask * token_mask = mask + (size_t) t * n_kv;
+    const int32_t * token_topk = topk + (size_t) t * capacity;
+    int * token_rows = selected_rows + (size_t) t * capacity;
+    int * token_owner_offsets = owner_offsets + (size_t) t * (N_OWNERS + 1);
+    int * token_owner_ranks = owner_ranks + (size_t) t * capacity;
+
+    int row = INVALID_ROW;
+    if (tid < capacity) {
+        const int comp = token_topk[tid];
+        const int physical = raw_rows + comp;
+        if (comp >= 0 && comp < n_comp_rows &&
+            ds4_fa_load<Mask, Mask>(token_mask + physical) > -1.0e20f) {
+            row = physical;
+        }
+    }
+    sorted_rows[tid] = row;
+    if (tid < N_OWNERS) owner_counts[tid] = 0;
+    __syncthreads();
+
+    for (int width = 2; width <= SORT_WIDTH; width <<= 1) {
+        for (int stride = width >> 1; stride > 0; stride >>= 1) {
+            const int peer = tid ^ stride;
+            if (peer > tid) {
+                const int lhs = sorted_rows[tid];
+                const int rhs = sorted_rows[peer];
+                const bool ascending = (tid & width) == 0;
+                if ((lhs > rhs) == ascending) {
+                    sorted_rows[tid] = rhs;
+                    sorted_rows[peer] = lhs;
+                }
+            }
+            __syncthreads();
+        }
+    }
+
+    if (tid == 0) {
+        int valid = 0;
+        while (valid < capacity && sorted_rows[valid] != INVALID_ROW) {
+            ++valid;
+        }
+        count = valid;
+        selected_counts[t] = valid;
+    }
+    __syncthreads();
+
+    if (tid < count) {
+        token_rows[tid] = sorted_rows[tid];
+        atomicAdd(owner_counts + (sorted_rows[tid] & (N_OWNERS - 1)), 1);
+    }
+    __syncthreads();
+
+    if (tid == 0) {
+        int prefix = 0;
+        for (int owner = 0; owner < N_OWNERS; ++owner) {
+            token_owner_offsets[owner] = prefix;
+            prefix += owner_counts[owner];
+        }
+        token_owner_offsets[N_OWNERS] = prefix;
+    }
+    __syncthreads();
+
+    if (tid < N_OWNERS) {
+        int write = token_owner_offsets[tid];
+        for (int rank = 0; rank < count; ++rank) {
+            if ((token_rows[rank] & (N_OWNERS - 1)) == tid) {
+                token_owner_ranks[write++] = rank;
+            }
+        }
+    }
+}
+
 template <typename KV, typename Mask>
 __global__ static void ds4_flash_attn_d512_shared_kv_kernel(
         float       * dst,
@@ -1618,6 +1718,7 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32_supported(const ggml_tensor * dst)
     const ggml_tensor * V = dst->src[2];
     const ggml_tensor * mask = dst->src[3];
     const ggml_tensor * sinks = dst->src[4];
+    const ggml_tensor * indexer_topk = dst->src[5];
     const bool kv_f32 = K && V && K->type == GGML_TYPE_F32 &&
                         V->type == GGML_TYPE_F32;
     const bool kv_f16 = K && V && K->type == GGML_TYPE_F16 &&
@@ -1676,6 +1777,15 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32_supported(const ggml_tensor * dst)
     const int raw_window = (int) (ds4_layout >> 16);
     const int sparse_block_size = (int) (ds4_layout & 0xffffu);
     const int rope_flags = ggml_get_op_params_i32(dst, 7);
+    if (indexer_topk &&
+        (sparse_keep_rows >= 0 || -sparse_keep_rows > 512 ||
+         indexer_topk->type != GGML_TYPE_I32 ||
+         indexer_topk->ne[0] != -sparse_keep_rows ||
+         indexer_topk->ne[1] != Q->ne[1] ||
+         indexer_topk->ne[2] != 1 || indexer_topk->ne[3] != 1 ||
+         !ggml_is_contiguous(indexer_topk))) {
+        return false;
+    }
     if (raw_rows < 0 || raw_rows > n_kv ||
         (ds4_layout != 0 && (raw_window <= 0 || sparse_block_size <= 0)) ||
         sparse_keep_rows == INT_MIN ||
@@ -1699,6 +1809,7 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32(
     const ggml_tensor * V = dst->src[2];
     const ggml_tensor * mask = dst->src[3];
     const ggml_tensor * sinks = dst->src[4];
+    const ggml_tensor * indexer_topk = dst->src[5];
     const bool kv_f32 = K->type == GGML_TYPE_F32;
     const bool kv_f16 = K->type == GGML_TYPE_F16;
     const int n_tokens = (int) Q->ne[1];
@@ -1856,7 +1967,14 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32(
             const bool parallel_index_scan = n_comp_rows > 512 &&
                 getenv("GGML_DS4_FA_SERIAL_INDEX_SCAN") == nullptr;
             if (mask->type == GGML_TYPE_F16) {
-                if (parallel_index_scan) {
+                if (indexer_topk) {
+                    ds4_fa_indexed_rows_topk_kernel<half><<<n_tokens, 512, 0, stream>>>(
+                        (const half *) mask->data,
+                        (const int32_t *) indexer_topk->data,
+                        indexed_rows, indexed_counts,
+                        indexed_owner_offsets, indexed_owner_ranks,
+                        n_tokens, n_kv, raw_rows, indexed_capacity);
+                } else if (parallel_index_scan) {
                     ds4_fa_indexed_rows_parallel_kernel<half><<<n_tokens, 256, 0, stream>>>(
                         (const half *) mask->data, indexed_rows, indexed_counts,
                         indexed_owner_offsets, indexed_owner_ranks,
@@ -1870,7 +1988,14 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32(
                         indexed_capacity);
                 }
             } else {
-                if (parallel_index_scan) {
+                if (indexer_topk) {
+                    ds4_fa_indexed_rows_topk_kernel<float><<<n_tokens, 512, 0, stream>>>(
+                        (const float *) mask->data,
+                        (const int32_t *) indexer_topk->data,
+                        indexed_rows, indexed_counts,
+                        indexed_owner_offsets, indexed_owner_ranks,
+                        n_tokens, n_kv, raw_rows, indexed_capacity);
+                } else if (parallel_index_scan) {
                     ds4_fa_indexed_rows_parallel_kernel<float><<<n_tokens, 256, 0, stream>>>(
                         (const float *) mask->data, indexed_rows, indexed_counts,
                         indexed_owner_offsets, indexed_owner_ranks,