perf(ds4): accelerate long context and add PFlash

README.md

@@ -93,6 +93,10 @@ These runs use different prompts, quantizations, and inference policies. They sh
 
 See [Recommended server setups](server/docs/RECOMMENDED_SETUPS.md) for the model and hardware matrix, including single-GPU and mixed-GPU profiles.
 
+The DS4 guide also documents the Strix long-context sparse-verifier profile and
+Qwen3-0.6B PFlash integration. PFlash is lossy prompt compression; keep it off
+for exact-retrieval and matched true-context benchmarks.
+
 ## Client Harnesses
 
 [`harness/`](harness/) runs Lucebox through popular coding clients and checks server compatibility.

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

@@ -2471,6 +2471,13 @@ extern "C" {
             int                  n_ctx_orig,
             bool                 q_unrotated);
 
+    // Optional runtime positions for both fused RoPE directions. I32 [n_query]
+    // replaces kv_start + query_index, allowing a cached graph to advance
+    // without rebuilding its topology or retaining stale position constants.
+    GGML_API void ggml_flash_attn_ext_set_ds4_rope_positions(
+            struct ggml_tensor * a,
+            struct ggml_tensor * positions);
+
     // True when flash_attn_ext carries the DS4 sparse-layout or fused-RoPE
     // contract. Backends must implement that complete contract or reject it.
     GGML_API bool ggml_flash_attn_ext_is_ds4(

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

@@ -351,12 +351,14 @@ static __global__ void ds4_indexer_score_decode_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. The post-WMMA loop still accumulates heads in
-// their original order, preserving the established F32 numerical topology.
-static __global__ void ds4_indexer_score_wmma_q4_kernel(
+// The speculative verifier scores only a few query tokens. The general WMMA
+// kernel places them in a 16-row tile and executes the unused rows for every
+// head. Pack consecutive heads into the tile instead:
+// row = N_TOKENS*head_in_group + token. The post-WMMA loop still accumulates
+// heads in their original order, preserving the established F32 numerical
+// topology. This is the HIP equivalent of the Vulkan small-CM dispatch.
+template<int N_TOKENS>
+static __global__ void ds4_indexer_score_wmma_small_kernel(
         float       * scores,
         const float * q,
         const float * weights,
@@ -375,7 +377,14 @@ static __global__ void ds4_indexer_score_wmma_q4_kernel(
     __shared__ float c_sh[8 * 16 * 16];
     __shared__ float weight_sh[16];
 
-    float acc[2] = {0.0f, 0.0f};
+    static_assert(N_TOKENS >= 2 && N_TOKENS <= 5,
+                  "small-CM kernel is specialized for verifier widths 2..5");
+    constexpr int HEADS_PER_TILE = 16 / N_TOKENS;
+    constexpr int USED_ROWS = HEADS_PER_TILE * N_TOKENS;
+    constexpr int ACC_SLOTS = (N_TOKENS + 1) / 2;
+    float acc[ACC_SLOTS];
+#pragma unroll
+    for (int slot = 0; slot < ACC_SLOTS; ++slot) acc[slot] = 0.0f;
 
     for (int i = tid; i < 128 * 128; i += 256) {
         const int c = i >> 7;
@@ -387,21 +396,30 @@ static __global__ void ds4_indexer_score_wmma_q4_kernel(
     }
     __syncthreads();
 
-    for (int head_base = 0; head_base < n_head; head_base += 4) {
+    for (int head_base = 0; head_base < n_head;
+         head_base += HEADS_PER_TILE) {
         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);
+            half2 value = __float2half2_rn(0.0f);
+            if (row < USED_ROWS) {
+                const int token = row % N_TOKENS;
+                const int head = head_base + row / N_TOKENS;
+                if (head < n_head) {
+                    const float2 q_value =
+                        *reinterpret_cast<const float2 *>(
+                            q + ((size_t) token * n_head + head) * 128 + d);
+                    value = __floats2half2_rn(q_value.x, q_value.y);
+                }
+            }
+            *reinterpret_cast<half2 *>(a_sh + row * 128 + d) = value;
         }
         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];
+            const int token = tid % N_TOKENS;
+            const int head = head_base + tid / N_TOKENS;
+            weight_sh[tid] = tid < USED_ROWS && head < n_head
+                ? weights[(size_t) token * n_head + head]
+                : 0.0f;
         }
         __syncthreads();
 
@@ -431,16 +449,17 @@ static __global__ void ds4_indexer_score_wmma_q4_kernel(
         __syncthreads();
 
         int slot = 0;
-        for (int output = tid; output < 4 * 128;
+        for (int output = tid; output < N_TOKENS * 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;
+            for (int head_in_group = 0;
+                 head_in_group < HEADS_PER_TILE;
                  ++head_in_group) {
-                const int row = 4 * head_in_group + token;
+                const int row = N_TOKENS * 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];
@@ -450,7 +469,7 @@ static __global__ void ds4_indexer_score_wmma_q4_kernel(
     }
 
     int slot = 0;
-    for (int output = tid; output < 4 * 128;
+    for (int output = tid; output < N_TOKENS * 128;
          output += 256, ++slot) {
         const int token = output >> 7;
         const int comp = tile_c + (output & 127);
@@ -553,6 +572,17 @@ void ggml_cuda_op_ds4_indexer_score(
         warp_size == 32 &&
         (!GGML_CUDA_CC_IS_NVIDIA(device_info.cc) ||
          device_info.cc >= GGML_CUDA_CC_VOLTA);
+    const char * packed_small_name = "GGML_DS4_INDEXER_PACK_SMALL";
+    const char * packed_small_env = std::getenv(packed_small_name);
+    if (!packed_small_env) {
+        // Backward-compatible alias for the original q=4-only prototype.
+        packed_small_name = "GGML_DS4_INDEXER_PACK_Q4";
+        packed_small_env = std::getenv(packed_small_name);
+    }
+    const bool use_packed_small = packed_small_env
+        ? ds4_env_flag_enabled(packed_small_name)
+        : GGML_CUDA_CC_IS_RDNA3_5(device_info.cc) ||
+          GGML_CUDA_CC_IS_RDNA4(device_info.cc);
 #if DS4_INDEXER_WMMA_AVAILABLE
     if (wmma_capable && n_tokens == 1) {
         const dim3 grid((unsigned) ((n_comp + 127) / 128), 1, 1);
@@ -564,10 +594,39 @@ void ggml_cuda_op_ds4_indexer_score(
             visibility_mask
                 ? static_cast<const float *>(visibility_mask->data) : nullptr,
             n_comp, kv_start, n_head, ratio);
-    } else if (wmma_capable && n_tokens == 4 && n_head % 4 == 0 &&
-               ds4_env_flag_enabled("GGML_DS4_INDEXER_PACK_Q4")) {
+    } else if (wmma_capable && n_tokens == 2 && use_packed_small) {
+        const dim3 grid((unsigned) ((n_comp + 127) / 128), 1, 1);
+        ds4_indexer_score_wmma_small_kernel<2><<<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),
+            visibility_mask
+                ? static_cast<const float *>(visibility_mask->data) : nullptr,
+            n_comp, kv_start, n_head, ratio);
+    } else if (wmma_capable && n_tokens == 3 && use_packed_small) {
+        const dim3 grid((unsigned) ((n_comp + 127) / 128), 1, 1);
+        ds4_indexer_score_wmma_small_kernel<3><<<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),
+            visibility_mask
+                ? static_cast<const float *>(visibility_mask->data) : nullptr,
+            n_comp, kv_start, n_head, ratio);
+    } else if (wmma_capable && n_tokens == 4 && use_packed_small) {
+        const dim3 grid((unsigned) ((n_comp + 127) / 128), 1, 1);
+        ds4_indexer_score_wmma_small_kernel<4><<<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),
+            visibility_mask
+                ? static_cast<const float *>(visibility_mask->data) : nullptr,
+            n_comp, kv_start, n_head, ratio);
+    } else if (wmma_capable && n_tokens == 5 && use_packed_small) {
         const dim3 grid((unsigned) ((n_comp + 127) / 128), 1, 1);
-        ds4_indexer_score_wmma_q4_kernel<<<grid, 256, 0, stream>>>(
+        ds4_indexer_score_wmma_small_kernel<5><<<grid, 256, 0, stream>>>(
             static_cast<float *>(dst->data),
             static_cast<const float *>(q->data),
             static_cast<const float *>(weights->data),

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

@@ -103,6 +103,7 @@ struct ds4_inverse_rope_params {
     int   enabled;
     int   forward_q_enabled;
     int   kv_start;
+    const int32_t * positions;
     float freq_scale;
     float ext_factor;
     float attn_factor;
@@ -157,7 +158,8 @@ __device__ static __forceinline__ void ds4_inverse_rope_coefficients(
         const ds4_inverse_rope_params & p,
         float & cos_theta, float & sin_theta) {
     ds4_rope_coefficients_at_position(
-        pair, -(p.kv_start + token), p, cos_theta, sin_theta);
+        pair, -(p.positions ? p.positions[token] : p.kv_start + token),
+        p, cos_theta, sin_theta);
 }
 
 // Forward counterpart of ds4_inverse_rope_coefficients. Keep the expressions
@@ -169,7 +171,8 @@ __device__ static __forceinline__ void ds4_forward_rope_coefficients(
         const ds4_inverse_rope_params & p,
         float & cos_theta, float & sin_theta) {
     ds4_rope_coefficients_at_position(
-        pair, p.kv_start + token, p, cos_theta, sin_theta);
+        pair, p.positions ? p.positions[token] : p.kv_start + token,
+        p, cos_theta, sin_theta);
 }
 
 __device__ static __forceinline__ void ds4_apply_inverse_rope_pair(
@@ -478,7 +481,7 @@ __global__ static void ds4_fa_indexed_rows_parallel_kernel(
 // 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>
+template <typename Mask, int SORT_WIDTH>
 __global__ static void ds4_fa_indexed_rows_topk_kernel(
         const Mask    * mask,
         const int32_t * topk,
@@ -494,7 +497,6 @@ __global__ static void ds4_fa_indexed_rows_topk_kernel(
     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];
@@ -573,6 +575,26 @@ __global__ static void ds4_fa_indexed_rows_topk_kernel(
     }
 }
 
+template <typename Mask>
+static void ds4_launch_indexed_rows_topk(
+        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,
+        cudaStream_t stream) {
+    // The learned top-512 stays on its original launch. A batched verifier
+    // appends a small saved-raw suffix and needs the next sorting bucket.
+    if (capacity <= 512) {
+        ds4_fa_indexed_rows_topk_kernel<Mask, 512><<<n_tokens, 512, 0, stream>>>(
+            mask, topk, selected_rows, selected_counts, owner_offsets, owner_ranks,
+            n_tokens, n_kv, raw_rows, capacity);
+    } else {
+        ds4_fa_indexed_rows_topk_kernel<Mask, 1024><<<n_tokens, 1024, 0, stream>>>(
+            mask, topk, selected_rows, selected_counts, owner_offsets, owner_ranks,
+            n_tokens, n_kv, raw_rows, capacity);
+    }
+}
+
 template <typename KV, typename Mask>
 __global__ static void ds4_flash_attn_d512_shared_kv_kernel(
         float       * dst,
@@ -1777,6 +1799,14 @@ 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);
+    const ggml_tensor * rope_positions = dst->src[6];
+    if (rope_positions &&
+        ((rope_flags & 1) == 0 || rope_positions->type != GGML_TYPE_I32 ||
+         rope_positions->ne[0] != n_tokens || rope_positions->ne[1] != 1 ||
+         rope_positions->ne[2] != 1 || rope_positions->ne[3] != 1 ||
+         !ggml_is_contiguous(rope_positions))) {
+        return false;
+    }
     if (sparse_keep_rows == INT_MIN) {
         return false;
     }
@@ -1789,7 +1819,7 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32_supported(const ggml_tensor * dst)
     }
     const int n_comp_rows = n_kv - raw_rows;
     if (indexer_topk &&
-        (sparse_keep_rows >= 0 || -sparse_keep_rows > 512 ||
+        (sparse_keep_rows >= 0 || -sparse_keep_rows > 1024 ||
          -sparse_keep_rows > n_comp_rows ||
          indexer_topk->type != GGML_TYPE_I32 ||
          indexer_topk->ne[0] != -sparse_keep_rows ||
@@ -1868,6 +1898,8 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32(
     inverse_rope.forward_q_enabled = (rope_flags & 2) != 0;
     if (rope_flags != 0) {
         inverse_rope.kv_start = ggml_get_op_params_i32(dst, 8);
+        inverse_rope.positions = dst->src[6]
+            ? static_cast<const int32_t *>(dst->src[6]->data) : nullptr;
         const float freq_base = ggml_get_op_params_f32(dst, 9);
         inverse_rope.freq_scale = ggml_get_op_params_f32(dst, 10);
         inverse_rope.ext_factor = ggml_get_op_params_f32(dst, 11);
@@ -1994,12 +2026,12 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32(
                 getenv("GGML_DS4_FA_SERIAL_INDEX_SCAN") == nullptr;
             if (mask->type == GGML_TYPE_F16) {
                 if (indexer_topk) {
-                    ds4_fa_indexed_rows_topk_kernel<half><<<n_tokens, 512, 0, stream>>>(
+                    ds4_launch_indexed_rows_topk<half>(
                         (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);
+                        n_tokens, n_kv, raw_rows, indexed_capacity, stream);
                 } 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,
@@ -2015,12 +2047,12 @@ static bool ggml_cuda_ds4_flash_attn_d512_f32(
                 }
             } else {
                 if (indexer_topk) {
-                    ds4_fa_indexed_rows_topk_kernel<float><<<n_tokens, 512, 0, stream>>>(
+                    ds4_launch_indexed_rows_topk<float>(
                         (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);
+                        n_tokens, n_kv, raw_rows, indexed_capacity, stream);
                 } 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,

server/deps/llama.cpp/ggml/src/ggml.c

@@ -5668,6 +5668,19 @@ void ggml_flash_attn_ext_set_ds4_inverse_rope(
     ggml_set_op_params_i32(a, 15, n_ctx_orig);
 }
 
+void ggml_flash_attn_ext_set_ds4_rope_positions(
+        struct ggml_tensor * a,
+        struct ggml_tensor * positions) {
+    GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT);
+    GGML_ASSERT((ggml_get_op_params_i32(a, 7) & 1) != 0);
+    GGML_ASSERT(a->src[6] == NULL);
+    GGML_ASSERT(positions && positions->type == GGML_TYPE_I32);
+    GGML_ASSERT(ggml_is_contiguous(positions));
+    GGML_ASSERT(positions->ne[0] == a->src[0]->ne[1]);
+    GGML_ASSERT(positions->ne[1] == 1 && positions->ne[2] == 1 && positions->ne[3] == 1);
+    a->src[6] = positions;
+}
+
 bool ggml_flash_attn_ext_is_ds4(const struct ggml_tensor * a) {
     return a && a->op == GGML_OP_FLASH_ATTN_EXT &&
            (ggml_get_op_params_i32(a, 6) != 0 ||