server : fix swa checkpoint invalidation logic (#30666)

* server : fix swa checkpoint invalidation logic

Assisted-by: Claude

* add testcase that fails with > n_past, but works with > n_prefix

.github/workflows/build-ibm.yml

@@ -20,7 +20,11 @@ on:
     paths: [
       '.github/workflows/build-ibm.yml',
       'ggml/src/ggml-cpu/**',
-      'ggml/src/ggml-zdnn/**'
+      'ggml/src/ggml-zdnn/**',
+      'ggml/src/ggml-quants.c',
+      'src/llama-quant.cpp',
+      'gguf-py/gguf/scripts/gguf_convert_endian.py',
+      'tests/test-quantize-be.sh'
     ]
 
 cache-mode: none
@@ -105,6 +109,11 @@ jobs:
           wget https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories260K-be.gguf
           ./bin/llama-completion -m stories260K-be.gguf -p "One day, Lily met a Shoggoth" -n 500 -c 256
 
+      - name: Test quantize (s390x)
+        id: quantize_test_s390x
+        run: |
+          ./tests/test-quantize-be.sh ./build/bin/llama-quantize ./build/bin/llama-perplexity 2>&1
+
   ubuntu-26-zdnn-s390x:
     name: ubuntu-26-zdnn-s390x
     runs-on: ubuntu-24.04-s390x

.pi/gg/APPEND_SYSTEM.md

@@ -0,0 +1,6 @@
+
+Project-specific instructions:
+
+## Coding
+
+- In `test-backend-ops.cpp`, do not mention specific backends (e.g. Metal, CUDA) in comments

.pi/gg/SYSTEM.md

@@ -1,44 +0,0 @@
-You are a coding agent. Here are some very important rules that you must follow:
-
-General:
-- Be very precise and concise when writing code, comments, explanations, etc.
-- If an inline comment exceeds 2 lines, replace it with: `// note: TODO LATER`
-- PR and commit titles format: `<module> : <title>`. Lookup recents for examples
-- Don't try to build or run the code unless you are explicitly asked to do so
-- Use the `gh` CLI tool when querying PRs, issues, or other GitHub resources
-- When a model name is needed, first try to get it from the `PI_MODEL_NAME` env var before asking the user
-- When a hardware description is needed, first try to get it from the `PI_HARDWARE` env var before asking the user
-- Never read the `AGENTS.md` file
-
-Coding:
-- Never write inline comments in the code unless the user tells you to
-- When in doubt, always refer to the CONTRIBUTING.md file of the project
-- In `test-backend-ops.cpp`, do not mention specific backends (e.g. Metal, CUDA) in comments
-- When referencing issues or PRs in comments, use the format:
-  - C/C++ code: `// ref: <url>`
-  - Other (CMake, etc.): `# ref: <url>`
-
-Pull requests (PRs):
-- New branch names are prefixed with "gg/"
-- Before opening a pull request, ask the user to confirm the description
-- Don't explicitly wrap lines in the PR description (each paragraph and bullet is a single line)
-- When creating a pull request, look for the repository's PR template and follow it
-- For the AI usage disclosure section, write "YES", then these 4 bullets:
-  - Harness: `pi:llama.cpp`
-  - Hardware: `<hardware>`
-  - Model: `<model>`, linked to its Hugging Face page - the link text is the plain model name
-  - Stats: `<stats>`
-- Run `~/.pi/stats.sh "$PI_SESSION_FILE"` to get `<stats>` and write it as
-  `◷ <time> · ↑ <input tokens> · ↓ <output tokens> · ≡ <messages> · ⌨ <user messages>`,
-  e.g. `◷ 12.3m · ↑ 120k · ↓ 45k · ≡ 87 · ⌨ 6`
-- Write these symbols as literal unicode in the PR description, never `:shortcodes:`, and keep
-  commit trailers ASCII
-- If `PI_MODEL_NAME` or `PI_HARDWARE` is not set, ask the user and use their answer for `<model>` or `<hardware>`
-- Always create the pull requests in draft mode
-- Never reply to review comments or post comments on issues/PRs without explicit permission from the user
-
-Commits:
-- On every commit that you make, include an `Assisted-by: pi:llama.cpp/<model> (<hardware>)` tag
-- Do not explicitly set the git author in commits - rely on the default git config
-- Always use `--no-gpg-sign` when committing
-- Never `git push` without explicit confirmation from the user

conversion/base.py

@@ -1625,7 +1625,7 @@ class TextModel(ModelBase):
             self.gguf_writer.add_head_count_kv(n_head_kv)
             logger.info(f"gguf: key-value head count = {n_head_kv}")
 
-        if self.hparams.get("is_causal") is False:
+        if self.hparams.get("is_causal") is False or self.hparams.get("use_bidirectional_attention") is True:
             self.gguf_writer.add_causal_attention(False)
             logger.info("gguf: causal attention = False")
 

conversion/qwen.py

@@ -253,6 +253,11 @@ class Qwen3Model(Qwen2Model):
                     yield from super().modify_tensors(data_torch, name, bid)
                 return
 
+        # output embedding projection (voyage-4-nano)
+        if name == "linear.weight":
+            self.gguf_writer.add_embedding_length_out(data_torch.shape[0])
+            name = "dense_2_out.weight"
+
         yield from super().modify_tensors(data_torch, name, bid)
 
 

ggml/src/ggml-cann/ggml-cann.cpp

@@ -2455,9 +2455,10 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev, const ggml_ten
 #ifdef ASCEND_310P
                         // Q4 && Q8 per group is not support on 310p device
                         return false;
-#endif
+#else
                         // only support contiguous for quantized types.
                         return ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1]);
+#endif
                     default:
                         return false;
                 }
@@ -2472,9 +2473,10 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev, const ggml_ten
 #ifdef ASCEND_310P
                     // Q4 && Q8 per group is not support on 310p device
                     return false;
-#endif
+#else
                     // only support contiguous for quantized types.
                     return ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1]);
+#endif
                 default:
                     return false;
             }
@@ -2641,7 +2643,7 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev, const ggml_ten
 #ifdef ASCEND_310P
                 // Ger is not supported on 310p device
                 return false;
-#endif
+#else
                 switch (op->src[0]->type) {
                     case GGML_TYPE_F16:
                     case GGML_TYPE_F32:
@@ -2649,6 +2651,7 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev, const ggml_ten
                     default:
                         return false;
                 }
+#endif
             }
         case GGML_OP_CONV_TRANSPOSE_1D:
             return true;
@@ -2667,7 +2670,7 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev, const ggml_ten
 #ifdef ASCEND_310P
                 // FA not support on 310p device
                 return false;
-#endif
+#else
                 // derived from [ggml-cuda.cu]
                 if (op->src[1]->type != GGML_TYPE_F16 || op->src[2]->type != GGML_TYPE_F16) {
                     return false;
@@ -2693,6 +2696,7 @@ static bool ggml_backend_cann_supports_op(ggml_backend_dev_t dev, const ggml_ten
                     return false;
                 }
                 return true;
+#endif
             }
         case GGML_OP_SSM_CONV:
             return true;

ggml/src/ggml-cpu/vec.cpp

@@ -401,6 +401,10 @@ void ggml_vec_silu_f32(const int n, float * y, const float * x) {
     for (; i + 3 < n; i += 4) {
         vst1q_f32(y + i, ggml_v_silu(vld1q_f32(x + i)));
     }
+#elif defined(__VXE__) || defined(__VXE2__)
+    for (; i + 3 < n; i += 4) {
+        vec_xst(ggml_v_silu(vec_xl(0, x + i)), 0, y + i);
+    }
 #elif defined(__riscv_v_intrinsic)
     for (int vl; i < n; i += vl) {
         vl = __riscv_vsetvl_e32m2(n - i);
@@ -438,6 +442,10 @@ void ggml_vec_swiglu_f32(const int n, float * y, const float * x, const float *
     for (; i + 3 < n; i += 4) {
         vst1q_f32(y + i, vmulq_f32(ggml_v_silu(vld1q_f32(x + i)), vld1q_f32(g + i)));
     }
+#elif defined(__VXE__) || defined(__VXE2__)
+    for (; i + 3 < n; i += 4) {
+        vec_xst(vec_mul(ggml_v_silu(vec_xl(0, x + i)), vec_xl(0, g + i)), 0, y + i);
+    }
 #elif defined(__riscv_v_intrinsic)
     for (int vl; i < n; i += vl) {
         vl = __riscv_vsetvl_e32m2(n - i);
@@ -581,6 +589,12 @@ ggml_float ggml_vec_soft_max_f32(const int n, float * y, const float * x, float
         vst1q_f32(y + i, val);
         sum += (ggml_float)vaddvq_f32(val);
     }
+#elif defined(__VXE__) || defined(__VXE2__)
+    for (; i + 3 < n; i += 4) {
+        float32x4_t v_val = ggml_v_expf(vec_sub(vec_xl(0, x + i), vec_splats(max)));
+        vec_xst(v_val, 0, y + i);
+        sum += (ggml_float)vec_hsum_f32x4(v_val);
+    }
 #elif defined(__riscv_v_intrinsic)
     vfloat64m1_t vsum = __riscv_vfmv_v_f_f64m1(0, 1);
     for (int avl; i < n; i += avl) {

ggml/src/ggml-cpu/vec.h

@@ -1334,6 +1334,46 @@ inline static __m128 ggml_v_silu(__m128 x) {
     return _mm_div_ps(x, one_plus_exp_neg_x);
 }
 
+#elif defined(__VXE__) || defined(__VXE2__)
+
+// adapted from arm limited optimized routine
+// the maximum error is 1.45358 plus 0.5 ulps
+// numbers above 88.38 will flush to infinity
+// numbers beneath -103.97 will flush to zero
+inline static float32x4_t ggml_v_expf(float32x4_t v_x) {
+    const float32x4_t v_r = vec_splats(0x1.8p23f);
+    const float32x4_t v_z = vec_madd(v_x, vec_splats(0x1.715476p+0f), v_r);
+    const float32x4_t v_n = vec_sub(v_z, v_r);
+    const float32x4_t v_b = vec_nmsub(v_n, vec_splats(0x1.7f7d1cp-20f),
+                                      vec_nmsub(v_n, vec_splats(0x1.62e4p-1f), v_x));
+    const uint32x4_t  v_e = vec_sl((uint32x4_t)v_z, 23);
+    const float32x4_t v_k = (float32x4_t)vec_add(v_e, (uint32x4_t)vec_splats(1.0f));
+    const uint32x4_t  v_c = (uint32x4_t)vec_cmpgt(vec_abs(v_n), vec_splats(126.0f));
+    const float32x4_t v_u = vec_mul(v_b, v_b);
+    const float32x4_t v_j = vec_madd(
+        vec_madd(vec_madd(vec_splats(0x1.0e4020p-7f), v_b, vec_splats(0x1.573e2ep-5f)), v_u,
+                 vec_madd(vec_splats(0x1.555e66p-3f), v_b, vec_splats(0x1.fffdb6p-2f))),
+        v_u, vec_mul(vec_splats(0x1.ffffecp-1f), v_b));
+    if (vec_all_eq(v_c, vec_splats(UINT32_C(0)))) {
+        return vec_madd(v_j, v_k, v_k);
+    }
+    const uint32x4_t  v_d  = vec_and((uint32x4_t)vec_cmple(v_n, vec_splats(0.0f)), vec_splats(UINT32_C(0x82000000)));
+    const float32x4_t v_s1 = (float32x4_t)vec_add(v_d, vec_splats(UINT32_C(0x7f000000)));
+    const float32x4_t v_s2 = (float32x4_t)vec_sub(v_e, v_d);
+    return vec_sel(vec_sel(vec_madd(v_k, v_j, v_k), vec_mul(vec_madd(v_s2, v_j, v_s2), v_s1), v_c),
+                   vec_mul(v_s1, v_s1), (uint32x4_t)vec_cmpgt(vec_abs(v_n), vec_splats(192.0f)));
+}
+
+// computes silu x/(1+exp(-x)) in single precision vector
+inline static float32x4_t ggml_v_silu(float32x4_t v_x) {
+    const float32x4_t v_one = vec_splats(1.0f);
+    const float32x4_t v_zero = vec_splats(0.0f);
+    const float32x4_t v_neg_x = vec_sub(v_zero, v_x);
+    const float32x4_t v_exp_neg_x = ggml_v_expf(v_neg_x);
+    const float32x4_t v_one_plus_exp_neg_x = vec_add(v_one, v_exp_neg_x);
+    return vec_div(v_x, v_one_plus_exp_neg_x);
+}
+
 #elif defined(__riscv_v_intrinsic)
 
 // adapted from arm limited optimized routine
@@ -1386,7 +1426,7 @@ inline static vfloat32m2_t ggml_v_silu_m2(vfloat32m2_t x, int vl) {
     return __riscv_vfdiv_vv_f32m2(x, one_plus_exp_neg_x, vl);
 }
 
-#endif // __ARM_NEON / __AVX2__ / __SSE2__ / __riscv_v_intrinsic
+#endif // __ARM_NEON / __AVX2__ / __SSE2__ / __VXE__ / __riscv_v_intrinsic
 
 inline static void ggml_vec_silu_f16(const int n, ggml_fp16_t * y, const ggml_fp16_t * x) {
     for (int i = 0; i < n; ++i) {

ggml/src/ggml-cuda/fattn-vec.cuh

@@ -317,7 +317,7 @@ static __global__ void flash_attn_ext_vec(
 #endif // V_DOT2_F32_F16_AVAILABLE
         }
 
-        ggml_cuda_syncwarp();
+        ggml_cuda_syncwarp(); // Needed for KQ synchronization.
 
 #pragma unroll
         for (int k0 = 0; k0 < WARP_SIZE; k0 += V_cols_per_iter) {
@@ -374,6 +374,8 @@ static __global__ void flash_attn_ext_vec(
             }
 #endif // V_DOT2_F32_F16_AVAILABLE
         }
+
+        ggml_cuda_syncwarp(); // Needed for KQ synchronization.
     }
 
     if (sinks && blockIdx.y == 0) {

ggml/src/ggml-cuda/gated-delta-net-mma.cu

@@ -4,7 +4,7 @@
 #include <climits>
 #include <mutex>
 
-namespace {
+namespace ggml_cuda_gated_delta_net_mma {
 constexpr int D = 128;
 using bf16      = nv_bfloat16;
 
@@ -417,6 +417,7 @@ static gdn_path select_gdn_mma_path(int device, const ggml_cuda_gdn_mma_args & a
 #endif // GGML_USE_HIP
 }
 }  // namespace
+using namespace ggml_cuda_gated_delta_net_mma;
 
 bool ggml_cuda_gdn_mma_launch(int device, const ggml_cuda_gdn_mma_args & a, cudaStream_t stream) {
     if (!a.state_out) {

ggml/src/ggml-cuda/mmq.cuh

@@ -1441,10 +1441,12 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a
     const int ntiles_dst = ntx * nty * ntzw;
     const int tiles_nwaves = (ntiles_dst + nsm - 1) / nsm;
     const int tiles_efficiency_percent = 100 * ntiles_dst / (nsm*tiles_nwaves);
-    const dim3 block_nums_stream_k(GGML_CUDA_CC_IS_NVIDIA(cc) && tiles_efficiency_percent >= 90 ? ntiles_dst : nsm, 1, 1);
 
     GGML_ASSERT(ntiles_dst * blocks_per_ne00_fd.z < (1 << 30)); // Assert that variable kbc will not overflow.
 
+    const int nwork = ntiles_dst * ((args.ncols_x + config.K_vram - 1) / config.K_vram); // Maximum number of CUDA blocks that can do useful work.
+    const dim3 block_nums_stream_k(GGML_CUDA_CC_IS_NVIDIA(cc) && tiles_efficiency_percent >= 90 ? ntiles_dst : std::min(nsm, nwork), 1, 1);
+
     const bool fixup_needed = ntiles_dst % block_nums_stream_k.x != 0;
 
     ggml_cuda_pool & pool = ctx.pool(id);

ggml/src/ggml-hexagon/ggml-hexagon.cpp

@@ -3472,15 +3472,16 @@ static bool ggml_hexagon_flash_attn_is_hmx_eligible(
         return true;
     }
 
-    // Compute-bound head dim
-    if (DK >= 256) {
-        return true;
-    }
-
     const uint32_t n_head = q->ne[2];
     const uint32_t n_kv_heads = k->ne[2];
     const uint32_t G = n_kv_heads > 0 ? n_head / n_kv_heads : 1;
     const uint32_t S = k->ne[1];
+    const uint32_t total_threads = sess->n_threads * sess->mdev.count;
+
+    // Compute-bound head dim (for DK == 256 decode, HVX is faster when KV heads span all threads)
+    if (DK > 256 || (DK == 256 && (DV > 256 || n_kv_heads < total_threads))) {
+        return true;
+    }
 
     // Tile alignment for 32-row HMX tiles
     const bool is_tile_aligned = (G > 0 && (32 % G == 0));
@@ -3493,7 +3494,9 @@ static bool ggml_hexagon_flash_attn_is_hmx_eligible(
 
     // Context depth crossover
     uint32_t s_cross = 512;
-    if (DK <= 64) {
+    if (DK == 256) {
+        s_cross = 4096;
+    } else if (DK <= 64) {
         s_cross = (sess->n_threads >= 8) ? 2048 : ((sess->n_threads >= 6) ? 768 : 512);
     } else {
         s_cross = (sess->n_threads >= 8) ? 1024 : 512;
@@ -3665,10 +3668,6 @@ static bool ggml_hexagon_supported_flash_attn_ext(const struct ggml_hexagon_sess
         return false;
     }
 
-    if (dst->ne[3] != 1) {
-        return false;
-    }
-
     struct htp_fa_kernel_params kparams;
     if (!ggml_hexagon_precompute_flash_attn_params(sess, op, &kparams)) {
         return false;

ggml/src/ggml-hexagon/htp/flash-attn-ops.c

@@ -64,6 +64,8 @@ struct htp_fa_context {
     __fp16 slopes[512];
 
     uint32_t n_blocks;
+    uint32_t n_active_blocks;
+    bool     shared_active_blocks;
 
     size_t size_q_row_padded;
     size_t size_k_row_padded;
@@ -82,12 +84,13 @@ struct htp_fa_context {
     size_t size_q_block;
     size_t size_vkq_acc;
 
-    uint8_t * spad_q;
-    uint8_t * spad_k;
-    uint8_t * spad_v;
-    uint8_t * spad_m;
-    uint8_t * spad_a;
-    float *   spad_sinks;
+    uint8_t *  spad_q;
+    uint8_t *  spad_k;
+    uint8_t *  spad_v;
+    uint8_t *  spad_m;
+    uint8_t *  spad_a;
+    float *    spad_sinks;
+    uint16_t   active_blocks[FA_MAX_KV_BLOCKS];
 
     const struct htp_tensor * k;
     const struct htp_tensor * v;
@@ -156,6 +159,7 @@ struct hmx_fa_context {
     size_t       row_buf_stride;       // HVX vectors per row buffer (Bc/64)
     size_t       mask_buf_row_stride;  // elements (__fp16) per row in mask buffer
     size_t       q_tile_bytes;
+    size_t       k_tile_bytes;
     size_t       o_tile_bytes;
     size_t       col_vec_bytes;
     size_t       d_tile_bytes;
@@ -163,6 +167,77 @@ struct hmx_fa_context {
     dma_cache_fa m_cache;
 };
 
+// Scan mask rows [q_start, q_start + n_rows_q) and build compact list of KV blocks with any entry != -inf.
+static uint32_t fa_scan_active_blocks(
+        dma_queue *               dma_q,
+        const struct htp_tensor * mask,
+        uint32_t                  q_start,
+        uint32_t                  n_rows_q,
+        uint32_t                  im2,
+        uint32_t                  im3,
+        uint32_t                  nek1,
+        uint32_t                  Bc,
+        uint32_t                  n_kv_blocks,
+        void *                    vtcm_scratch,
+        size_t                    scratch_bytes,
+        uint16_t *                active_blocks) {
+    const size_t row_bytes  = (size_t) nek1 * sizeof(__fp16);
+    const size_t row_stride = hex_round_up(row_bytes, 128);
+
+    if (!mask || n_kv_blocks <= 1 || n_kv_blocks > FA_MAX_KV_BLOCKS ||
+        (size_t) n_rows_q * row_stride > scratch_bytes) {
+        #pragma clang loop vectorize(disable)
+        for (uint32_t b = 0; b < n_kv_blocks; ++b) {
+            active_blocks[b] = (uint16_t) b;
+        }
+        return n_kv_blocks;
+    }
+
+    const dma_addr_t m_src = mask->data + q_start * mask->nb[1] + im2 * mask->nb[2] + im3 * mask->nb[3];
+    dma_queue_push(dma_q, dma_make_data(vtcm_scratch, m_src), row_stride, mask->nb[1], row_bytes, n_rows_q);
+    const uint8_t * m_buf = (const uint8_t *) dma_queue_pop(dma_q).dst;
+
+    const HVX_Vector v_ninf   = Q6_Vh_vsplat_R(0xFC00);
+    uint32_t         n_active = 0;
+
+    for (uint32_t b = 0; b < n_kv_blocks; ++b) {
+        const uint32_t col0 = b * Bc;
+        const uint32_t cols = MIN(Bc, nek1 - col0);
+
+        for (uint32_t r = 0; r < n_rows_q; ++r) {
+            const __fp16 * row_ptr = (const __fp16 *) (m_buf + r * row_stride) + col0;
+            HVX_Vector v_any = Q6_V_vzero();
+            uint32_t c = 0;
+            for (; c + 64 <= cols; c += 64) {
+                const HVX_Vector v_m = *(const HVX_Vector *) (row_ptr + c);
+                v_any = Q6_V_vor_VV(v_any, Q6_V_vxor_VV(v_m, v_ninf));
+            }
+            if (c < cols) {
+                const uint32_t rem = cols - c;
+                const HVX_VectorPred q_rem = Q6_Q_vsetq2_R(rem * sizeof(__fp16));
+                const HVX_Vector v_m = *(const HVX_Vector *) (row_ptr + c);
+                const HVX_Vector v_diff = Q6_V_vand_QV(q_rem, Q6_V_vxor_VV(v_m, v_ninf));
+                v_any = Q6_V_vor_VV(v_any, v_diff);
+            }
+            v_any = Q6_V_vor_VV(v_any, Q6_V_vror_VR(v_any, 64));
+            v_any = Q6_V_vor_VV(v_any, Q6_V_vror_VR(v_any, 32));
+            v_any = Q6_V_vor_VV(v_any, Q6_V_vror_VR(v_any, 16));
+            v_any = Q6_V_vor_VV(v_any, Q6_V_vror_VR(v_any, 8));
+            v_any = Q6_V_vor_VV(v_any, Q6_V_vror_VR(v_any, 4));
+            if (Q6_R_vextract_VR(v_any, 0) != 0) {
+                active_blocks[n_active++] = (uint16_t) b;
+                break;
+            }
+        }
+    }
+
+    if (n_active == 0 && n_kv_blocks > 0) {
+        active_blocks[0] = 0;
+        n_active = 1;
+    }
+    return n_active;
+}
+
 static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * data) {
     struct htp_fa_context * factx = (struct htp_fa_context *) data;
     const struct htp_ops_context * octx = factx->octx;
@@ -236,6 +311,18 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
     dma_cache_dm m_cache;
     dma_cache_dm_init(&m_cache, spad_m, factx->size_m_block, HVX_FA_DMA_CACHE_SIZE);
 
+    uint16_t thread_active_blocks[FA_MAX_KV_BLOCKS];
+    const uint16_t * active_blocks = factx->shared_active_blocks ? factx->active_blocks : thread_active_blocks;
+    uint32_t n_active_blocks       = factx->n_active_blocks;
+    dma_addr_t prev_mp_base        = 0;
+
+    if (!factx->shared_active_blocks && !mask) {
+        #pragma clang loop vectorize(disable)
+        for (uint32_t b = 0; b < factx->n_blocks; ++b) {
+            thread_active_blocks[b] = (uint16_t) b;
+        }
+    }
+
     const size_t size_vkq_acc_single = hex_round_up(DV * sizeof(float), 128);
 
     uint32_t ir = ir0;
@@ -271,6 +358,19 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
             }
         }
 
+        dma_addr_t mp_base = 0;
+        if (mask) {
+            const uint32_t im2 = fastmodulo(iq2, mask->ne[2], &factx->src3_div2);
+            const uint32_t im3 = fastmodulo(iq3, mask->ne[3], &factx->src3_div3);
+            mp_base = mask->data + iq1*mask->nb[1] + im2*mask->nb[2] + im3*mask->nb[3];
+            if (!factx->shared_active_blocks && mp_base != prev_mp_base) {
+                n_active_blocks = fa_scan_active_blocks(
+                    dma_q, mask, iq1, 1, im2, im3, nek1, FLASH_ATTN_BLOCK_SIZE,
+                    factx->n_blocks, spad_k, factx->size_k_block * 2, thread_active_blocks);
+                prev_mp_base = mp_base;
+            }
+        }
+
         uint32_t heads[FA_HVX_G_MAX];
         HVX_Vector slope_vecs[FA_HVX_G_MAX] __attribute__((aligned(128)));
         HVX_Vector S_vec[FA_HVX_G_MAX]      __attribute__((aligned(128)));
@@ -303,26 +403,20 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
             dma_queue_push(dma_q, dma_make_data(q_dst, q_row_ptr), factx->size_q_row_padded, nbq1, size_q_row, 1);
         }
 
-        dma_addr_t mp_base = 0;
-        if (mask) {
-            const uint32_t im2 = fastmodulo(iq2, mask->ne[2], &factx->src3_div2);
-            const uint32_t im3 = fastmodulo(iq3, mask->ne[3], &factx->src3_div3);
-            mp_base = mask->data + iq1*mask->nb[1] + im2*mask->nb[2] + im3*mask->nb[3];
-        }
-
-        // Prefetch first two blocks
-        for (uint32_t ib = 0; ib < MIN(factx->n_blocks, 2); ++ib) {
+        // Prefetch first two active blocks
+        for (uint32_t step = 0; step < MIN(n_active_blocks, 2); ++step) {
+            const uint32_t ib = active_blocks[step];
             const uint32_t ic_start = ib * FLASH_ATTN_BLOCK_SIZE;
             const uint32_t current_block_size = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start);
 
             // K
             const dma_addr_t k_src = k->data + ic_start*nbk1 + ik2*nbk2 + ik3*nbk3;
-            uint8_t * k_dst = spad_k + (ib % 2) * factx->size_k_block;
+            uint8_t * k_dst = spad_k + (step % 2) * factx->size_k_block;
             dma_queue_push(dma_q, dma_make_data(k_dst, k_src), factx->size_k_row_padded, nbk1, size_k_row, current_block_size);
 
             // V
             const dma_addr_t v_src = v->data + ic_start*nbv1 + iv2*nbv2 + iv3*nbv3;
-            uint8_t * v_dst = spad_v + (ib % 2) * factx->size_v_block;
+            uint8_t * v_dst = spad_v + (step % 2) * factx->size_v_block;
             dma_queue_push(dma_q, dma_make_data(v_dst, v_src), factx->size_v_row_padded, nbv1, size_v_row, current_block_size);
 
             // Mask
@@ -348,7 +442,8 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
         const HVX_Vector v_log2e   = hvx_vec_splat_f16(EXP_LOG2E_F);
         const uint32_t stride_v2   = factx->size_v_row_padded * 2;
 
-        for (uint32_t ib = 0; ib < factx->n_blocks; ++ib) {
+        for (uint32_t step = 0; step < n_active_blocks; ++step) {
+            const uint32_t ib = active_blocks[step];
             const uint32_t ic_start = ib * FLASH_ATTN_BLOCK_SIZE;
             const uint32_t current_block_size = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - ic_start);
 
@@ -411,9 +506,9 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
                 HVX_Vector v_max     = Q6_V_lo_W(hvx_vec_f16_to_f32(v_max_f16));
                 htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_QK, head_ir);
 
-                // prefetch K for block ib + 2 after last head finished QK
-                if (g + 1 == G_local && ib + 2 < factx->n_blocks) {
-                    const uint32_t next_ib = ib + 2;
+                // prefetch K for step + 2 after last head finished QK
+                if (g + 1 == G_local && step + 2 < n_active_blocks) {
+                    const uint32_t next_ib = active_blocks[step + 2];
                     const uint32_t next_ic_start = next_ib * FLASH_ATTN_BLOCK_SIZE;
                     const uint32_t next_block_size = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - next_ic_start);
 
@@ -504,6 +599,56 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
 
                         ((HVX_VectorPair *) VKQ32)[0] = vkq0;
                         ((HVX_VectorPair *) VKQ32)[1] = vkq1;
+                    } else if (DV == 256) {
+                        HVX_VectorPair vkq0 = ((const HVX_VectorPair *) VKQ32)[0];
+                        HVX_VectorPair vkq1 = ((const HVX_VectorPair *) VKQ32)[1];
+                        HVX_VectorPair vkq2 = ((const HVX_VectorPair *) VKQ32)[2];
+                        HVX_VectorPair vkq3 = ((const HVX_VectorPair *) VKQ32)[3];
+                        vkq0 = Q6_W_vcombine_VV(
+                            HVX_OP_MUL_F32(Q6_V_hi_W(vkq0), ms_vec),
+                            HVX_OP_MUL_F32(Q6_V_lo_W(vkq0), ms_vec)
+                        );
+                        vkq1 = Q6_W_vcombine_VV(
+                            HVX_OP_MUL_F32(Q6_V_hi_W(vkq1), ms_vec),
+                            HVX_OP_MUL_F32(Q6_V_lo_W(vkq1), ms_vec)
+                        );
+                        vkq2 = Q6_W_vcombine_VV(
+                            HVX_OP_MUL_F32(Q6_V_hi_W(vkq2), ms_vec),
+                            HVX_OP_MUL_F32(Q6_V_lo_W(vkq2), ms_vec)
+                        );
+                        vkq3 = Q6_W_vcombine_VV(
+                            HVX_OP_MUL_F32(Q6_V_hi_W(vkq3), ms_vec),
+                            HVX_OP_MUL_F32(Q6_V_lo_W(vkq3), ms_vec)
+                        );
+
+                        for (uint32_t j = 0; j < current_block_size; j += 2) {
+                            HVX_Vector S0 = hvx_vec_repl_f16(Q6_V_vror_VR(P, j * 2));
+                            const HVX_Vector * vx0 = (const HVX_Vector *) v_ptr;
+                            if (j + 1 == current_block_size) {
+                                vkq0 = hvx_vec_mpyacc_f32_f16(vkq0, Q6_Vh_vshuff_Vh(vx0[0]), S0);
+                                vkq1 = hvx_vec_mpyacc_f32_f16(vkq1, Q6_Vh_vshuff_Vh(vx0[1]), S0);
+                                vkq2 = hvx_vec_mpyacc_f32_f16(vkq2, Q6_Vh_vshuff_Vh(vx0[2]), S0);
+                                vkq3 = hvx_vec_mpyacc_f32_f16(vkq3, Q6_Vh_vshuff_Vh(vx0[3]), S0);
+                                break;
+                            }
+
+                            HVX_Vector S1 = hvx_vec_repl_f16(Q6_V_vror_VR(P, (j + 1) * 2));
+                            const HVX_Vector * vx1 = (const HVX_Vector *) (v_ptr + factx->size_v_row_padded);
+                            vkq0 = hvx_vec_mpyacc_f32_f16(vkq0, Q6_Vh_vshuff_Vh(vx0[0]), S0);
+                            vkq0 = hvx_vec_mpyacc_f32_f16(vkq0, Q6_Vh_vshuff_Vh(vx1[0]), S1);
+                            vkq1 = hvx_vec_mpyacc_f32_f16(vkq1, Q6_Vh_vshuff_Vh(vx0[1]), S0);
+                            vkq1 = hvx_vec_mpyacc_f32_f16(vkq1, Q6_Vh_vshuff_Vh(vx1[1]), S1);
+                            vkq2 = hvx_vec_mpyacc_f32_f16(vkq2, Q6_Vh_vshuff_Vh(vx0[2]), S0);
+                            vkq2 = hvx_vec_mpyacc_f32_f16(vkq2, Q6_Vh_vshuff_Vh(vx1[2]), S1);
+                            vkq3 = hvx_vec_mpyacc_f32_f16(vkq3, Q6_Vh_vshuff_Vh(vx0[3]), S0);
+                            vkq3 = hvx_vec_mpyacc_f32_f16(vkq3, Q6_Vh_vshuff_Vh(vx1[3]), S1);
+                            v_ptr += stride_v2;
+                        }
+
+                        ((HVX_VectorPair *) VKQ32)[0] = vkq0;
+                        ((HVX_VectorPair *) VKQ32)[1] = vkq1;
+                        ((HVX_VectorPair *) VKQ32)[2] = vkq2;
+                        ((HVX_VectorPair *) VKQ32)[3] = vkq3;
                     } else {
                         hvx_scale_vec_f32_aa((uint8_t *) VKQ32, (const uint8_t *) VKQ32, DV, ms_vec);
 
@@ -524,9 +669,9 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
                 }
                 htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_SFM, head_ir);
 
-                // prefetch V and mask for block ib + 2 after last head finished V accumulation
-                if (g + 1 == G_local && ib + 2 < factx->n_blocks) {
-                    const uint32_t next_ib = ib + 2;
+                // prefetch V and mask for step + 2 after last head finished V accumulation
+                if (g + 1 == G_local && step + 2 < n_active_blocks) {
+                    const uint32_t next_ib = active_blocks[step + 2];
                     const uint32_t next_ic_start = next_ib * FLASH_ATTN_BLOCK_SIZE;
                     const uint32_t next_block_size = MIN(FLASH_ATTN_BLOCK_SIZE, nek1 - next_ic_start);
 
@@ -541,7 +686,7 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void *
                     }
                 }
             } // end for g
-        } // end for ib
+        } // end for step
 
         for (uint32_t g = 0; g < G_local; ++g) {
             const uint32_t head_ir = ir + g;
@@ -1997,6 +2142,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
     factx.vtcm_mask_buf       = VTCM_LAYOUT_PTR(__fp16, base, L.off_mask_buf);
     factx.mask_buf_row_stride = L.mask_buf_row_stride;
     factx.q_tile_bytes        = L.q_tile_bytes;
+    factx.k_tile_bytes        = L.k_tile_bytes;
     factx.o_tile_bytes        = L.o_tile_bytes;
     factx.col_vec_bytes       = L.col_vec_bytes;
     factx.d_tile_bytes        = L.d_tile_bytes;
@@ -2049,7 +2195,6 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
     const size_t q_src_stride               = q_transposed ? q->nb[2] : q->nb[1];
     const size_t q_row_bytes_untransposed   = factx.G * factx.DK * qo_element_size;
     const size_t q_row_bytes_trans_factor   = factx.DK * qo_element_size;
-    const uint32_t kv_rows0                 = hex_smin(Bc, nek1);
 
     // ======== Reusable job descriptors for pipeline ========
     hmx_fa_qk_job_t       qk_job;
@@ -2065,34 +2210,51 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
             const size_t   g_br_actual = hex_align_up(n_rows_g, HMX_FP16_TILE_N_ROWS);
             const size_t   n_row_tiles = g_br_actual / HMX_FP16_TILE_N_ROWS;
 
+            uint16_t active_blocks[FA_MAX_KV_BLOCKS];
+            uint32_t n_active_blocks = factx.n_kv_blocks;
+            if (mask && factx.mask_broadcast && factx.n_kv_blocks > 1) {
+                const size_t scratch_bytes = factx.k_tile_bytes * (factx.pipeline ? 2 : 1);
+                n_active_blocks = fa_scan_active_blocks(
+                    dma_q, mask, q_start, n_rows_q, 0, im3, nek1, Bc,
+                    factx.n_kv_blocks, factx.vtcm_k_tiles[0], scratch_bytes, active_blocks);
+            } else {
+                #pragma clang loop vectorize(disable)
+                for (uint32_t b = 0; b < factx.n_kv_blocks; ++b) {
+                    active_blocks[b] = (uint16_t) b;
+                }
+            }
+            const uint32_t b0       = (n_active_blocks > 0) ? active_blocks[0] : 0;
+            const uint32_t b0_start = b0 * Bc;
+            const uint32_t b0_rows  = hex_smin(Bc, nek1 - b0_start);
+
             for (uint32_t kv_head = kv_head_min; kv_head < kv_head_max; ++kv_head) {
                 const uint32_t ik2 = kv_head;
                 const uint32_t ik3 = fastdiv(ib3, &kparams->broadcast_rk3);
                 const uint32_t iv2 = kv_head;
                 const uint32_t iv3 = fastdiv(ib3, &kparams->broadcast_rv3);
 
-                // 1. Push Q and KV DMAs for the very first iteration.
-                // Subsequent iterations are enqueued early at the end of the previous iteration.
-                if (ib3 == 0 && q_start == q_start_min && kv_head == kv_head_min) {
+                // 1. Push Q and KV DMAs for the first kv_head of this q_start block.
+                // Subsequent kv_head iterations are enqueued early at the end of the previous iteration.
+                if (kv_head == kv_head_min) {
                     const dma_addr_t q_ptr = q->data + q_start * q->nb[1] +
                                             (kv_head * factx.G) * q->nb[2] + ib3 * q->nb[3];
                     const size_t q_row_bytes = q_transposed ? n_rows_q * q_row_bytes_trans_factor : q_row_bytes_untransposed;
                     const size_t n_rows      = q_transposed ? factx.G : n_rows_q;
                     dma_queue_push(dma_q, dma_make_data(factx.vtcm_q_dma, q_ptr), q_row_bytes, hex_smax(q_src_stride, q_row_bytes), q_row_bytes, n_rows);
 
-                    if (factx.n_kv_blocks > 0) {
-                        const dma_addr_t k_src = k->data + ik2 * k->nb[2] + ik3 * k->nb[3];
-                        dma_queue_push(dma_q, dma_make_data(factx.vtcm_k_fp16[0], k_src), size_k_row_padded, k->nb[1], size_k_row, kv_rows0);
+                    if (n_active_blocks > 0) {
+                        const dma_addr_t k_src = k->data + b0_start * k->nb[1] + ik2 * k->nb[2] + ik3 * k->nb[3];
+                        dma_queue_push(dma_q, dma_make_data(factx.vtcm_k_fp16[0], k_src), size_k_row_padded, k->nb[1], size_k_row, b0_rows);
 
-                        const dma_addr_t v_src = v->data + iv2 * v->nb[2] + iv3 * v->nb[3];
-                        dma_queue_push(dma_q, dma_make_data(factx.vtcm_v_fp16[0], v_src), size_v_row_padded, v->nb[1], size_v_row, kv_rows0);
+                        const dma_addr_t v_src = v->data + b0_start * v->nb[1] + iv2 * v->nb[2] + iv3 * v->nb[3];
+                        dma_queue_push(dma_q, dma_make_data(factx.vtcm_v_fp16[0], v_src), size_v_row_padded, v->nb[1], size_v_row, b0_rows);
 
                         if (factx.pipeline && mask) {
                             if (__builtin_expect(factx.mask_broadcast, true)) {
-                                const dma_addr_t ms_src = mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + 0;
-                                dma_cache_fa_push(dma_q, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows0 * sizeof(__fp16), n_rows_q);
+                                const dma_addr_t ms_src = mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + b0_start * sizeof(__fp16);
+                                dma_cache_fa_push(dma_q, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], b0_rows * sizeof(__fp16), n_rows_q);
                             } else {
-                                fa_push_mask_dma_gqa(dma_q, mask, q_start, im3, 0, kv_head, G, m_line_bytes, kv_rows0, n_rows_q, &factx);
+                                fa_push_mask_dma_gqa(dma_q, mask, q_start, im3, b0_start, kv_head, G, m_line_bytes, b0_rows, n_rows_q, &factx);
                             }
                         }
                     }
@@ -2124,27 +2286,28 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                     hmx_fa_qk_job_t qk_job[2];
                     hmx_fa_o_update_job_t ou_job[2];
 
-                    // Prefetch block 1 early if there are multiple blocks
-                    if (factx.n_kv_blocks > 1) {
-                        fa_prefetch_block(dma_q, k, v, mask, 1, Bc, size_k_row_padded, size_k_row, size_v_row_padded, size_v_row,
+                    // Prefetch second active block early if there are multiple active blocks
+                    if (n_active_blocks > 1) {
+                        fa_prefetch_block(dma_q, k, v, mask, active_blocks[1], Bc, size_k_row_padded, size_k_row, size_v_row_padded, size_v_row,
                                           ik2, ik3, iv2, iv3, q_start, im3, kv_head, G, m_line_bytes, n_rows_q, nek1, 1, &factx);
                     }
 
                     // Prep and start QK-dot(0)
                     void * curr_k0 = (void *) dma_queue_pop(dma_q).dst;
-                    fa_phase_k_interleave(&factx, kv_rows0, k_src_stride, curr_k0, 0, 0);
+                    fa_phase_k_interleave(&factx, b0_rows, k_src_stride, curr_k0, b0_start, 0);
 
                     qk_job[0].q_tiles        = factx.vtcm_q_tiles;
                     qk_job[0].k_tiles        = factx.vtcm_k_tiles[0];
                     qk_job[0].s_tiles        = factx.vtcm_s_tiles[0];
                     qk_job[0].n_row_tiles    = n_row_tiles;
-                    qk_job[0].n_col_tiles    = hmx_ceil_div(kv_rows0, HMX_FP16_TILE_N_COLS);
+                    qk_job[0].n_col_tiles    = hmx_ceil_div(b0_rows, HMX_FP16_TILE_N_COLS);
                     qk_job[0].n_dot_tiles    = DK_pad / 32;
                     qk_job[0].n_tiles_per_bc = n_tiles_per_bc;
                     qk_job[0].hmx_scales     = factx.vtcm_hmx_scales_qk;
                     hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_qk_dot_worker, &qk_job[0]));
 
-                    for (uint32_t kv_blk = 0; kv_blk < factx.n_kv_blocks; ++kv_blk) {
+                    for (uint32_t step = 0; step < n_active_blocks; ++step) {
+                        const uint32_t kv_blk      = active_blocks[step];
                         const uint32_t kv_start    = kv_blk * Bc;
                         const uint32_t kv_rows     = hex_smin(Bc, nek1 - kv_start);
                         const size_t   n_col_tiles = hmx_ceil_div(kv_rows, HMX_FP16_TILE_N_COLS);
@@ -2164,13 +2327,14 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                             }
                         }
 
-                        // ---- 3. Start HMX O update for block kv_blk - 1 (reads P[1 - buf_idx], V[1 - buf_idx], D) ----
-                        // O update relys on the previous block's P and V tiles.
+                        // ---- 3. Start HMX O update for step - 1 (reads P[1 - buf_idx], V[1 - buf_idx], D) ----
+                        // O update relies on the previous block's P and V tiles.
                         // O update MUST be pushed before the next block's QK-dot: hmx_queue_pop() retires the
                         // oldest descriptor, so push order alone decides which pop waits for which job.
                         // If OU went in after QK(i+1), the pop below would retire QK(i+1) and leave
                         // OU(i-1) in flight into the next iteration, where V-prep overwrites V[prev_buf].
-                        if (kv_blk > 0) {
+                        if (step > 0) {
+                            const uint32_t prev_blk      = active_blocks[step - 1];
                             const size_t prev_buf        = 1 - buf_idx;
                             ou_job[prev_buf].o_curr      = o_tile_curr;
                             ou_job[prev_buf].o_prev      = o_tile_prev;
@@ -2180,16 +2344,17 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                             ou_job[prev_buf].hmx_scales  = factx.vtcm_hmx_scales_id;
                             ou_job[prev_buf].n_row_tiles = n_row_tiles;
                             ou_job[prev_buf].n_col_tiles =
-                                hmx_ceil_div(hex_smin(Bc, nek1 - (kv_blk - 1) * Bc), HMX_FP16_TILE_N_COLS);
+                                hmx_ceil_div(hex_smin(Bc, nek1 - prev_blk * Bc), HMX_FP16_TILE_N_COLS);
                             ou_job[prev_buf].n_row_tiles_g_br = n_row_tiles_g_br;
                             ou_job[prev_buf].n_tiles_per_bc   = n_tiles_per_bc;
                             ou_job[prev_buf].DV               = DV_pad;
                             hmx_queue_push(hmx_q, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job[prev_buf]));
                         }
 
-                        // ---- 4. Pop and run K-prep for next block & push next QK-dot ----
-                        if (kv_blk + 1 < factx.n_kv_blocks) {
-                            const uint32_t next_start = (kv_blk + 1) * Bc;
+                        // ---- 4. Pop and run K-prep for next active block & push next QK-dot ----
+                        if (step + 1 < n_active_blocks) {
+                            const uint32_t next_blk   = active_blocks[step + 1];
+                            const uint32_t next_start = next_blk * Bc;
                             const uint32_t next_rows  = hex_smin(Bc, nek1 - next_start);
                             const size_t   next_buf   = 1 - buf_idx;
 
@@ -2236,15 +2401,15 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                         // Run Softmax on HVX (blocking call)
                         fa_phase_softmax_and_build_d(&factx, &sargs, n_row_tiles, n_row_tiles_g_br);
 
-                        // Wait for HMX O update for block kv_blk - 1 to finish
-                        if (kv_blk > 0) {
+                        // Wait for HMX O update for step - 1 to finish
+                        if (step > 0) {
                             hmx_queue_pop(hmx_q);
                             hex_swap_ptr((void **) &o_tile_curr, (void **) &o_tile_prev);
                         }
 
-                        // Prefetch block kv_blk + 2
-                        if (kv_blk + 2 < factx.n_kv_blocks) {
-                            fa_prefetch_block(dma_q, k, v, mask, kv_blk + 2, Bc, size_k_row_padded, size_k_row, size_v_row_padded, size_v_row,
+                        // Prefetch block step + 2
+                        if (step + 2 < n_active_blocks) {
+                            fa_prefetch_block(dma_q, k, v, mask, active_blocks[step + 2], Bc, size_k_row_padded, size_k_row, size_v_row_padded, size_v_row,
                                               ik2, ik3, iv2, iv3, q_start, im3, kv_head, G, m_line_bytes, n_rows_q, nek1, buf_idx, &factx);
                         }
 
@@ -2252,8 +2417,8 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                     }
 
                     // Epilogue
-                    if (factx.n_kv_blocks > 0) {
-                        const uint32_t last_blk = factx.n_kv_blocks - 1;
+                    if (n_active_blocks > 0) {
+                        const uint32_t last_blk = active_blocks[n_active_blocks - 1];
                         const size_t last_cols  = hmx_ceil_div(hex_smin(Bc, nek1 - last_blk * Bc), HMX_FP16_TILE_N_COLS);
                         ou_job[0].o_curr           = o_tile_curr;
                         ou_job[0].o_prev           = o_tile_prev;
@@ -2279,7 +2444,8 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
 
                 } else {
                     // Fallback path
-                    for (uint32_t kv_blk = 0; kv_blk < factx.n_kv_blocks; ++kv_blk) {
+                    for (uint32_t step = 0; step < n_active_blocks; ++step) {
+                        const uint32_t kv_blk      = active_blocks[step];
                         const uint32_t kv_start    = kv_blk * Bc;
                         const uint32_t kv_rows     = hex_smin(Bc, nek1 - kv_start);
                         const size_t   n_col_tiles = hmx_ceil_div(kv_rows, HMX_FP16_TILE_N_COLS);
@@ -2293,8 +2459,9 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                             }
                         }
 
-                        if (kv_blk + 1 < factx.n_kv_blocks) {
-                            const uint32_t  prefetch_start = (kv_blk + 1) * Bc;
+                        if (step + 1 < n_active_blocks) {
+                            const uint32_t  prefetch_blk   = active_blocks[step + 1];
+                            const uint32_t  prefetch_start = prefetch_blk * Bc;
                             const uint32_t  prefetch_rows  = hex_smin(Bc, nek1 - prefetch_start);
                             const size_t    prefetch_buf   = 1 - buf_idx;
                             const dma_addr_t k_prefetch_src = k->data + prefetch_start * k->nb[1] + ik2 * k->nb[2] + ik3 * k->nb[3];
@@ -2372,7 +2539,7 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                             ou_job.DV               = DV_pad;
 
                             hmx_queue_push(ctx->hmx_queue, hmx_queue_make_desc(hmx_fa_o_update_worker, &ou_job));
-                            if (kv_blk + 1 == factx.n_kv_blocks) {
+                            if (step + 1 == n_active_blocks) {
                                 // Overlapped: run HVX build diag inv L while HMX is busy executing the update
                                 htp_trace_event_start(tr_hvx, HTP_TRACE_EVT_HVX_O_PROC, (uint16_t) q_start);
                                 fa_build_d_diag_inv_l(&factx, n_row_tiles, n_row_tiles_g_br, kv_head, n_rows_g);
@@ -2387,53 +2554,32 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) {
                     }
                 }
 
-                // Enqueue DMAs for the next iteration early so they overlap with O-PROC
-                uint32_t next_kv_head = kv_head + 1;
-                uint32_t next_q_start = q_start;
-                uint32_t next_ib3     = ib3;
-                if (next_kv_head >= kv_head_max) {
-                    next_kv_head = kv_head_min;
-                    next_q_start = q_start + Br;
-                    if (next_q_start >= q_start_max) {
-                        next_q_start = q_start_min;
-                        next_ib3     = ib3 + 1;
-                    }
-                }
-                bool has_next = (next_ib3 < neq3);
+                // Enqueue DMAs for the next kv_head early so they overlap with O-PROC
+                const uint32_t next_kv_head = kv_head + 1;
+                const bool has_next = (next_kv_head < kv_head_max);
 
                 if (has_next) {
-                    const uint32_t next_n_rows_q = hex_smin(Br, neq1 - next_q_start);
-                    const dma_addr_t next_q_ptr = q->data + next_q_start * q->nb[1] + (next_kv_head * factx.G) * q->nb[2] + next_ib3 * q->nb[3];
-                    const size_t next_q_row_bytes = q_transposed ? next_n_rows_q * q_row_bytes_trans_factor : q_row_bytes_untransposed;
-                    const size_t next_n_rows      = q_transposed ? factx.G : next_n_rows_q;
+                    const dma_addr_t next_q_ptr = q->data + q_start * q->nb[1] + (next_kv_head * factx.G) * q->nb[2] + ib3 * q->nb[3];
+                    const size_t next_q_row_bytes = q_transposed ? n_rows_q * q_row_bytes_trans_factor : q_row_bytes_untransposed;
+                    const size_t next_n_rows      = q_transposed ? factx.G : n_rows_q;
                     dma_queue_push(dma_q, dma_make_data(factx.vtcm_q_dma, next_q_ptr), next_q_row_bytes, hex_smax(q_src_stride, next_q_row_bytes), next_q_row_bytes, next_n_rows);
 
-                    if (factx.n_kv_blocks > 0) {
+                    if (n_active_blocks > 0) {
                         const uint32_t next_ik2 = next_kv_head;
                         const uint32_t next_iv2 = next_kv_head;
-                        uint32_t next_ik3 = ik3;
-                        uint32_t next_iv3 = iv3;
-                        if (next_ib3 != ib3) {
-                            next_ik3 = fastdiv(next_ib3, &kparams->broadcast_rk3);
-                            next_iv3 = fastdiv(next_ib3, &kparams->broadcast_rv3);
-                        }
 
-                        const dma_addr_t next_k_src = k->data + next_ik2 * k->nb[2] + next_ik3 * k->nb[3];
-                        dma_queue_push(dma_q, dma_make_data(factx.vtcm_k_fp16[0], next_k_src), size_k_row_padded, k->nb[1], size_k_row, kv_rows0);
+                        const dma_addr_t next_k_src = k->data + b0_start * k->nb[1] + next_ik2 * k->nb[2] + ik3 * k->nb[3];
+                        dma_queue_push(dma_q, dma_make_data(factx.vtcm_k_fp16[0], next_k_src), size_k_row_padded, k->nb[1], size_k_row, b0_rows);
 
-                        const dma_addr_t next_v_src = v->data + next_iv2 * v->nb[2] + next_iv3 * v->nb[3];
-                        dma_queue_push(dma_q, dma_make_data(factx.vtcm_v_fp16[0], next_v_src), size_v_row_padded, v->nb[1], size_v_row, kv_rows0);
+                        const dma_addr_t next_v_src = v->data + b0_start * v->nb[1] + next_iv2 * v->nb[2] + iv3 * v->nb[3];
+                        dma_queue_push(dma_q, dma_make_data(factx.vtcm_v_fp16[0], next_v_src), size_v_row_padded, v->nb[1], size_v_row, b0_rows);
 
                         if (factx.pipeline && mask) {
-                            uint32_t next_im3 = im3;
-                            if (next_ib3 != ib3) {
-                                next_im3 = fastmodulo(next_ib3, mask->ne[3], &factx.src3_div3);
-                            }
                             if (__builtin_expect(factx.mask_broadcast, true)) {
-                                const dma_addr_t ms_src = mask->data + next_q_start * mask->nb[1] + next_im3 * mask->nb[3] + 0;
-                                dma_cache_fa_push(dma_q, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], kv_rows0 * sizeof(__fp16), next_n_rows_q);
+                                const dma_addr_t ms_src = mask->data + q_start * mask->nb[1] + im3 * mask->nb[3] + b0_start * sizeof(__fp16);
+                                dma_cache_fa_push(dma_q, &factx.m_cache, ms_src, m_line_bytes, mask->nb[1], b0_rows * sizeof(__fp16), n_rows_q);
                             } else {
-                                fa_push_mask_dma_gqa(dma_q, mask, next_q_start, next_im3, 0, next_kv_head, G, m_line_bytes, kv_rows0, next_n_rows_q, &factx);
+                                fa_push_mask_dma_gqa(dma_q, mask, q_start, im3, b0_start, next_kv_head, G, m_line_bytes, b0_rows, n_rows_q, &factx);
                             }
                         }
                     }
@@ -2481,7 +2627,7 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
 
     const struct htp_fa_kernel_params * kparams = (const struct htp_fa_kernel_params *) octx->kernel_params;
 
-    if (kparams->kernel_type == HTP_FA_KERNEL_UNSUPPORTED) {
+    if (kparams->kernel_type == HTP_FA_KERNEL_UNSUPPORTED || kparams->n_kv_blocks > FA_MAX_KV_BLOCKS) {
         return HTP_STATUS_NO_SUPPORT;
     }
 
@@ -2596,12 +2742,21 @@ int op_flash_attn_ext(struct htp_ops_context * octx) {
 
     uint8_t * const base = octx->ctx->vtcm_base;
 
-    factx.spad_q     = VTCM_LAYOUT_PTR(uint8_t, base, L.off_q);
-    factx.spad_k     = VTCM_LAYOUT_PTR(uint8_t, base, L.off_k);
-    factx.spad_v     = VTCM_LAYOUT_PTR(uint8_t, base, L.off_v);
-    factx.spad_m     = VTCM_LAYOUT_PTR_OPTIONAL(uint8_t, base, L.off_m, has_mask);
-    factx.spad_a     = VTCM_LAYOUT_PTR(uint8_t, base, L.off_a);
-    factx.spad_sinks = VTCM_LAYOUT_PTR_OPTIONAL(float, base, L.off_sinks, has_sinks);
+    factx.spad_q               = VTCM_LAYOUT_PTR(uint8_t, base, L.off_q);
+    factx.spad_k               = VTCM_LAYOUT_PTR(uint8_t, base, L.off_k);
+    factx.spad_v               = VTCM_LAYOUT_PTR(uint8_t, base, L.off_v);
+    factx.spad_m               = VTCM_LAYOUT_PTR_OPTIONAL(uint8_t, base, L.off_m, has_mask);
+    factx.spad_a               = VTCM_LAYOUT_PTR(uint8_t, base, L.off_a);
+    factx.spad_sinks           = VTCM_LAYOUT_PTR_OPTIONAL(float, base, L.off_sinks, has_sinks);
+    factx.n_active_blocks      = factx.n_blocks;
+    factx.shared_active_blocks = false;
+
+    if (has_mask && neq1 == 1 && neq3 == 1 && mask->ne[2] == 1 && factx.n_blocks > 1 && factx.n_blocks <= FA_MAX_KV_BLOCKS) {
+        factx.n_active_blocks = fa_scan_active_blocks(
+            octx->ctx->dma[0], mask, 0, 1, 0, 0, k->ne[1], FLASH_ATTN_BLOCK_SIZE,
+            factx.n_blocks, factx.spad_k, factx.size_k_block * 2, factx.active_blocks);
+        factx.shared_active_blocks = true;
+    }
 
     if (has_sinks) {
         const struct htp_tensor * sinks = octx->src[4];

ggml/src/ggml-hexagon/htp/flash-attn-ops.h

@@ -126,6 +126,7 @@ struct hmx_fa_vtcm_layout {
     // Region byte sizes reused by the device at runtime (not just for allocation).
     size_t q_tile_bytes;
     size_t o_tile_bytes;
+    size_t k_tile_bytes;
     size_t s_tile_bytes;       // S and P tiles (same size)
     size_t d_tile_bytes;       // d_tiles[0..1] + d_inv_l, allocated back to back
     size_t m_line_bytes;       // one mask row
@@ -141,6 +142,8 @@ struct hmx_fa_vtcm_layout {
     size_t total_bytes;
 };
 
+#define FA_MAX_KV_BLOCKS 1024
+
 // Build the VTCM layout.
 
 static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L,
@@ -224,6 +227,7 @@ static inline void hmx_fa_vtcm_layout_build(struct hmx_fa_vtcm_layout * L,
 
     L->q_tile_bytes        = q_tile_size;
     L->o_tile_bytes        = o_tile_size;
+    L->k_tile_bytes        = k_tile_size;
     L->col_vec_bytes       = col_vec_size;
     L->s_tile_bytes        = s_tile_size;
     // Measured from the actual offsets rather than assumed to be N * d_tile_size, so
@@ -332,8 +336,12 @@ static inline int hmx_fa_find_chunk_size(size_t * Br_out,
 
     // Pipeline constraint: cap Bc so n_kv_blocks >= FA_MIN_KV_BLOCKS.
     // Only relax when kv_len is too short to form enough blocks.
-    const size_t Bc_limit     = can_pipeline ? hex_align_down(kv_len / FA_MIN_KV_BLOCKS, bc_unit) :
-                                               (kv_len >= bc_unit ? hex_align_down(kv_len, bc_unit) : bc_unit);
+    size_t Bc_limit     = can_pipeline ? hex_align_down(kv_len / FA_MIN_KV_BLOCKS, bc_unit) :
+                                         (kv_len >= bc_unit ? hex_align_down(kv_len, bc_unit) : bc_unit);
+    // For small-Q SWA-sized KV caches (qo_len <= 16, kv_len <= 2048), cap Bc at 256 for block-skipping granularity
+    if (qo_len <= 16 && kv_len <= 2048 && Bc_limit > 256) {
+        Bc_limit = 256;
+    }
     // Cost coefficients calibrated from profiling
     const size_t c_q_fixed    = 800;   // per-Q-block: q_load + epilogue o_update + o_norm + o_store
     const size_t c_iter_base  = 200;   // per-KV-iter base (HMX dot/update + DMA)

ggml/src/ggml-metal/ggml-metal-fusion.cpp

@@ -636,6 +636,11 @@ static bool ggml_metal_fusion_check_moe_reduce(
         if (!match.experts->data || !match.weights->data || !match.dst->data) {
             return false;
         }
+
+        // the fused kernel reads experts and weights while writing dst, so dst must not overlap either input
+        if (ggml_metal_fusion_overlap(match.dst, match.experts) || ggml_metal_fusion_overlap(match.dst, match.weights)) {
+            return false;
+        }
     }
 
     return true;

gguf-py/gguf/constants.py

@@ -2929,6 +2929,7 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = {
         MODEL_TENSOR.FFN_GATE,
         MODEL_TENSOR.FFN_DOWN,
         MODEL_TENSOR.FFN_UP,
+        MODEL_TENSOR.DENSE_2_OUT, # voyage-4-nano
     ],
     MODEL_ARCH.QWEN3MOE: [
         MODEL_TENSOR.TOKEN_EMBD,

scripts/pr2wt.sh

@@ -72,19 +72,12 @@ dir=$(basename $(pwd))
 git branch -D pr/$PR 2> /dev/null
 git worktree add -b pr/$PR ../$dir-pr-$PR pr/$PR/$head_ref 2> /dev/null
 
-og_path=$(pwd)
 wt_path=$(cd ../$dir-pr-$PR && pwd)
 
 echo "git worktree created in $wt_path"
 
 cd $wt_path
 
-# pi agent setup in the worktree
-if [[ -f "$og_path/.pi/SYSTEM.md" && ! -f ".pi/SYSTEM.md" ]]; then
-    mkdir -p .pi
-    ln -sfn "$og_path/.pi/SYSTEM.md" .pi/SYSTEM.md
-fi
-
 git branch --set-upstream-to=pr/$PR/$head_ref
 git pull   --ff-only || {
     echo "error: failed to pull pr/$PR"

scripts/sync_vendor.py

@@ -5,7 +5,7 @@ import os
 import sys
 import subprocess
 
-HTTPLIB_VERSION = "refs/tags/v0.61.0"
+HTTPLIB_VERSION = "refs/tags/v0.62.0"
 
 # used by examples/gguf-hash, these repos have no release tag, so we pin a commit
 XXHASH_COMMIT      = "9f465f1ea932d6ad9a26cd77496311ffa544cd68"

scripts/wc2wt.sh

@@ -39,19 +39,12 @@ dir_suffix=$(echo "$BRANCH" | tr '/' '-')
 
 git worktree add "../$dir-$dir_suffix" "$BRANCH" || git worktree add -b "$BRANCH" "../$dir-$dir_suffix" HEAD
 
-og_path=$(pwd)
 wt_path=$(cd "../$dir-$dir_suffix" && pwd)
 
 echo "git worktree created in $wt_path"
 
 cd "$wt_path"
 
-# pi agent setup in the worktree
-if [[ -f "$og_path/.pi/SYSTEM.md" && ! -f ".pi/SYSTEM.md" ]]; then
-    mkdir -p .pi
-    ln -sfn "$og_path/.pi/SYSTEM.md" .pi/SYSTEM.md
-fi
-
 if [[ $# -eq 2 ]]; then
     echo "executing: $2"
     eval "$2"

skills/update-fusion-table/SKILL.md

@@ -0,0 +1,39 @@
+---
+name: update-fusion-table
+description: Regenerate a per-device test-fusion baseline (tests/fusion/MTL.csv, ...) after fusion patterns or the dummy-model set change. Use when test-fusion --check fails.
+---
+
+# Update a fusion baseline
+
+`tests/fusion/<BACKEND>.csv` is the `test-fusion` golden baseline: `arch,moe,mode,label,count`,
+where `count` is how many times a fusion label matched per graph run. See `tests/fusion/README.md`.
+
+Re-record only when the counts moved for the reason the contributor intended. A row failing on
+`nmse_fus` (> 1e-4) is a fused-vs-unfused correctness bug, not baseline drift - fix it, never
+record it away.
+
+## Steps
+
+```sh
+cmake -B build -DCMAKE_BUILD_TYPE=Release -DGGML_METAL=ON       # enable the target backend
+cmake --build build --config Release --target test-llama-archs --target test-fusion -j
+
+rm -rf build-ci-models && mkdir -p build-ci-models              # always regenerate
+./build/bin/test-llama-archs -o build-ci-models
+
+D=MTL0; B=tests/fusion/MTL.csv; T=/tmp/$(basename $B)   # keep the target backend/device in sync
+./build/bin/test-fusion --models build-ci-models --device $D --check $B   # read the FAIL rows
+./build/bin/test-fusion --models build-ci-models --device $D --record $T
+diff $B $T
+cp $T $B
+./build/bin/test-fusion --models build-ci-models --device $D --check $B   # 0 failed, exit 0
+```
+
+- Read the diff before applying it. Every added, removed or changed row must be explainable by the
+  change under review; an unexplained row means the run picked up something unintended.
+- `expected 0` in a FAIL row = no baseline row for that label (new arch, or a fusion newly
+  firing); non-zero `expected` = count changed; `fusion ... is no longer matched` = row stopped
+  firing.
+- `--record` rewrites the whole file, so never record a partial run (`--model FILE`, a filtered
+  models dir) - it drops the rows of every architecture that was not run.
+- `--record` marks every row `ok`; only the `--check` run validates.

src/llama-moe-cache.cpp

@@ -12,19 +12,21 @@
 
 namespace {
 
-// LRU of the experts of a group of layers, the slot of each expert is kept in the slot map of its layer
-struct moe_cache_lru {
+// SLRU of the experts of a group of layers, the slot of each expert is kept in the slot map of its layer
+struct moe_cache_slru {
     int32_t n_expert = 0;
     int32_t n_slots  = 0;
 
     std::vector<int32_t *> slot_map; // [n_layer] data of the slot maps, -1 if the expert is not cached
     std::vector<int32_t>   key_of;   // [n_slots] il*n_expert + expert, -1 if empty
 
-    // doubly linked list of the slots, head is the least recently used
+    // probation and protected LRU lists; hits promote slots to the protected half
     std::vector<int32_t> prev;
     std::vector<int32_t> next;
-    int32_t head = -1;
-    int32_t tail = -1;
+    int32_t head[2] = { -1, -1 };
+    int32_t tail[2] = { -1, -1 };
+    std::vector<bool> is_protected;
+    int32_t n_protected = 0;
 
     std::vector<uint32_t> seen; // [n_expert]
     uint32_t seen_gen = 0;
@@ -41,27 +43,59 @@ struct moe_cache_lru {
             prev[s] = s - 1;
             next[s] = s + 1 < n_slots ? s + 1 : -1;
         }
-        head = 0;
-        tail = n_slots - 1;
+        head[0] = 0;
+        tail[0] = n_slots - 1;
+        is_protected.assign(n_slots, false);
         seen.assign(n_expert, 0);
     }
 
-    // move slot s to the tail (most recently used)
-    void touch(int32_t s) {
-        if (s == tail) {
+    // move slot s to the most recently used end of a segment
+    void move(int32_t s, bool protect) {
+        const bool from = is_protected[s];
+        if (from == protect && s == tail[from]) {
             return;
         }
         if (prev[s] >= 0) {
             next[prev[s]] = next[s];
         } else {
-            head = next[s];
+            head[from] = next[s];
+        }
+        if (next[s] >= 0) {
+            prev[next[s]] = prev[s];
+        } else {
+            tail[from] = prev[s];
         }
-        prev[next[s]] = prev[s];
 
-        prev[s] = tail;
+        prev[s] = tail[protect];
         next[s] = -1;
-        next[tail] = s;
-        tail = s;
+        if (tail[protect] >= 0) {
+            next[tail[protect]] = s;
+        } else {
+            head[protect] = s;
+        }
+        tail[protect] = s;
+        is_protected[s] = protect;
+        n_protected += int(protect) - int(from);
+    }
+
+    void touch(int32_t s) {
+        move(s, true);
+        if (n_protected > n_slots/2) {
+            move(head[1], false);
+        }
+    }
+
+    int32_t victim(int32_t il) const {
+        for (int segment = 0; segment < 2; ++segment) {
+            for (int32_t s = head[segment]; s >= 0; s = next[s]) {
+                const int32_t key = key_of[s];
+                // keep every expert selected by this request, including earlier fills
+                if (key < 0 || key / n_expert != il || seen[key % n_expert] != seen_gen) {
+                    return s;
+                }
+            }
+        }
+        GGML_ABORT("no SLRU victim");
     }
 
     struct fill {
@@ -69,7 +103,7 @@ struct moe_cache_lru {
         int32_t slot;
     };
 
-    // give a slot to each expert selected by ids in layer il, the misses evict the least recently used experts
+    // give a slot to each expert selected by ids in layer il, evicting from probation first
     // returns false if the ids select more distinct experts than there are slots
     bool plan(int32_t il, const int32_t * ids, size_t n_ids, std::vector<fill> & fills, size_t & n_hit) {
         fills.clear();
@@ -93,7 +127,7 @@ struct moe_cache_lru {
 
         int32_t * slots = slot_map[il];
 
-        // hits go to the tail first, so the head can be evicted below
+        // promote hits before choosing victims for misses
         for (int32_t e : uniq) {
             if (slots[e] >= 0) {
                 touch(slots[e]);
@@ -106,13 +140,13 @@ struct moe_cache_lru {
             if (slots[e] >= 0) {
                 continue;
             }
-            const int32_t s = head;
+            const int32_t s = victim(il);
             if (key_of[s] >= 0) {
                 slot_map[key_of[s] / n_expert][key_of[s] % n_expert] = -1;
             }
             key_of[s] = il*n_expert + e;
             slots[e] = s;
-            touch(s);
+            move(s, false);
             fills.push_back({ e, s });
         }
         return true;
@@ -172,7 +206,7 @@ struct llama_moe_cache::impl {
         std::vector<ggml_tensor *> banks; // device storage of all slots, one per expert tensor
         size_t host_bytes = 0;
         int32_t n_slots = 0;
-        moe_cache_lru lru;
+        moe_cache_slru slru;
     };
 
     struct layer {
@@ -208,7 +242,7 @@ struct llama_moe_cache::impl {
     std::unordered_map<const ggml_tensor *, int32_t> layer_of; // slot map -> layer
 
     std::vector<int32_t> ids;
-    std::vector<moe_cache_lru::fill> fills;
+    std::vector<moe_cache_slru::fill> fills;
 
     // slot maps in host memory
     ggml_context_ptr ctx_host;
@@ -323,7 +357,7 @@ struct llama_moe_cache::impl {
                 g.n_slots = 0;
                 continue;
             }
-            g.lru.init(model.layers.size(), n_expert, g.n_slots);
+            g.slru.init(model.layers.size(), n_expert, g.n_slots);
             n_tensors[g.id] += g.ref.size()*(1 + g.layers.size());
             n_tensors_host  += g.layers.size();
         }
@@ -421,7 +455,7 @@ struct llama_moe_cache::impl {
             for (group & g : groups) {
                 for (int32_t il : g.layers) {
                     if (layers[il].slot_map != nullptr) {
-                        g.lru.slot_map[il] = (int32_t *) layers[il].slot_map->data;
+                        g.slru.slot_map[il] = (int32_t *) layers[il].slot_map->data;
                     }
                 }
             }
@@ -489,7 +523,7 @@ struct llama_moe_cache::impl {
         }
 
         // large batches only read the cache, so the experts used in generation stay in it
-        const int32_t * slots = g.lru.slot_map[b.il];
+        const int32_t * slots = g.slru.slot_map[b.il];
         if (slots == nullptr || slots[e] < 0) {
             return 0;
         }
@@ -550,7 +584,7 @@ struct llama_moe_cache::impl {
         ggml_backend_synchronize(backend);
 
         size_t n_hit = 0;
-        if (!g.lru.plan(il, ids.data(), ids.size(), fills, n_hit)) {
+        if (!g.slru.plan(il, ids.data(), ids.size(), fills, n_hit)) {
             GGML_ABORT("the MoE cache is too small for the experts selected in layer %d", il);
         }
 

src/llama-moe-cache.h

@@ -8,7 +8,7 @@
 
 struct llama_model;
 
-// keeps the most recently used experts of host-resident MoE layers in a device buffer
+// keeps recently and repeatedly used experts of host-resident MoE layers in a device buffer
 // each layer has a slot map in host memory: when the scheduler copies it to the device, the copy callback uploads the missing experts
 class llama_moe_cache {
 public:

src/models/qwen3.cpp

@@ -28,6 +28,10 @@ void llama_model_qwen3::load_arch_tensors(llama_model_loader &) {
     // output rerank head
     cls_out = create_tensor(tn(LLM_TENSOR_CLS_OUT, "weight"), {n_embd, hparams.n_cls_out}, TENSOR_NOT_REQUIRED);
 
+    // output embedding projection
+    dense_2_out_layers   = create_tensor(tn(LLM_TENSOR_DENSE_2_OUT, "weight"), {n_embd, hparams.n_embd_out()}, TENSOR_NOT_REQUIRED);
+    dense_2_out_layers_b = create_tensor(tn(LLM_TENSOR_DENSE_2_OUT, "bias"),   {hparams.n_embd_out()        }, TENSOR_NOT_REQUIRED);
+
     for (int i = 0; i < n_layer; ++i) {
         auto & layer = layers[i];
 

tests/fusion/MTL.csv

@@ -20,6 +20,8 @@ bailingmoe3         ,1   ,any     ,GATED_DELTA_NET+CPY         ,      1
 bailingmoe3         ,1   ,any     ,MUL+ADD                     ,      2
 bailingmoe3         ,1   ,any     ,RMS_NORM+MUL                ,      8
 bailingmoe3         ,1   ,any     ,RMS_NORM+SCALE              ,      2
+beam                ,0   ,any     ,MUL+ADD                     ,      6
+beam                ,0   ,any     ,RMS_NORM+MUL                ,     25
 bloom               ,0   ,any     ,NORM+MUL+ADD                ,      6
 chatglm             ,0   ,any     ,RMS_NORM+MUL                ,      5
 codeshell           ,0   ,any     ,NORM+MUL+ADD                ,      5

tests/test-quantize-be.sh

@@ -0,0 +1,178 @@
+#!/usr/bin/env bash
+set -euo pipefail
+
+if [ $# -lt 2 ]; then
+    printf "Usage: %s <llama-quantize> <llama-perplexity> [<scratch-dir>] [--no-check]\n" "$0"
+    printf "  <scratch-dir>  dir for the model and outputs, default /tmp/test-quantize-be\n"
+    printf "  --no-check     print the perplexity of each output without checking it\n"
+    printf "Example: %s ./build/bin/llama-quantize ./build/bin/llama-perplexity\n" "$0"
+    exit 1
+fi
+
+quantize=$1
+perplexity=$2
+shift 2
+
+scratch_dir=/tmp/test-quantize-be
+no_check=
+for arg in "$@"; do
+    case $arg in
+        --no-check) no_check=--no-check ;;
+        *)          scratch_dir=$arg ;;
+    esac
+done
+
+for bin in "$quantize" "$perplexity"; do
+    if [ ! -x "$bin" ]; then
+        echo "error: $bin not found, build it first" >&2
+        exit 1
+    fi
+done
+
+if [ "$(uname -m)" != s390x ]; then
+    echo "error: this test needs a big-endian s390x host" >&2
+    exit 1
+fi
+
+MODEL_URL=https://huggingface.co/ggml-org/models/resolve/499bc8821c6b12b4e53c5bffcb21ec206f212d81/tinyllamas/stories15M.gguf
+MODEL_SHA256=61b50d457809a5194818fd22e6724b456cd7bb9a6264c52c8110684c53f3704a
+F32_SHA256=55174d7a7d46f29fc6be163ccb040330c357aa727b22ec13e2c5d713e5ba5b76
+
+# perplexity ceiling of each big-endian output on the text below, about 10% above the measured value
+# a wrong byte order gives nan or a perplexity in the hundreds
+# F32 is the converted input, all other types come from llama-quantize --pure
+declare -a outputs=(
+    "F32       3.60"
+    "F16       3.60"
+    "BF16      3.59"
+    "Q4_0      4.14"
+    "Q4_1      4.01"
+    "Q5_0      3.82"
+    "Q5_1      3.65"
+    "Q8_0      3.60"
+    "Q2_K      5.24"
+    "Q3_K      4.25"
+    "Q4_K      3.83"
+    "Q5_K      3.67"
+    "Q6_K      3.60"
+    "IQ4_NL    3.86"
+    "IQ4_XS    3.85"
+    "MXFP4_MOE 4.27"
+)
+
+check_sha256() {
+    local name=$1
+    local file=$2
+    local expected=$3
+
+    local actual
+    actual=$(sha256sum "$file" | cut -d ' ' -f 1)
+
+    if [ "$actual" != "$expected" ]; then
+        echo "error: sha256 mismatch for $name: expected '$expected', got '$actual'" >&2
+        exit 1
+    fi
+}
+
+# keep the report readable, only show a command's output when it fails
+run_logged() {
+    local log=$1
+    shift
+
+    if ! "$@" > "$log" 2>&1; then
+        cat "$log" >&2
+        echo "error: command failed: $*" >&2
+        return 1
+    fi
+}
+
+# quantize one type and print its perplexity, print nothing if a step fails
+measure_ppl() {
+    local type=$1
+    local out=$2
+    local log=$scratch_dir/stories15M-be-$type-ppl.log
+
+    if [ "$type" != F32 ]; then
+        run_logged "$scratch_dir/stories15M-be-$type.log" "$quantize" --pure "$f32" "$out" "$type" || return 1
+    fi
+    run_logged "$log" "$perplexity" -m "$out" -f "$text" -c 128 || return 1
+
+    # llama-perplexity prints no final estimate when the result is nan
+    sed -n 's/.*Final estimate: PPL = \([0-9.]*\).*/\1/p' "$log"
+}
+
+mkdir -p "$scratch_dir"
+
+model=$scratch_dir/stories15M.gguf
+if [ ! -f "$model" ]; then
+    wget -q -O "$model" "$MODEL_URL"
+fi
+
+check_sha256 stories15M.gguf "$model" "$MODEL_SHA256"
+
+f32=$scratch_dir/stories15M-be-F32.gguf
+cp "$model" "$f32"
+echo YES | run_logged "$scratch_dir/stories15M-be-F32.log" python3 "$(dirname "$0")/../gguf-py/gguf/scripts/gguf_convert_endian.py" "$f32" big || exit 1
+check_sha256 F32 "$f32" "$F32_SHA256"
+
+# 588 tokens, enough for 4 chunks of 128
+text=$scratch_dir/stories.txt
+cat > "$text" << 'EOF'
+Once upon a time, there was a little girl named Lily. She lived in a small house with her mom and dad. Lily loved to play outside in the garden. One sunny day, she saw a big red ball under a tree. She ran to the ball and picked it up. It was soft and round.
+Lily wanted to play with her friend Tom. She went to his house and knocked on the door. Tom opened the door and smiled. "Hi, Lily! Do you want to play?" he asked. "Yes! Look at my new ball," said Lily. They went to the park together.
+At the park, Lily and Tom threw the ball to each other. They laughed and ran on the green grass. Then the ball went too high. It flew over a fence and into a pond. Lily was sad. "Oh no, my ball is in the water," she said.
+Tom had an idea. He found a long stick near the pond. He held the stick out over the water and pushed the ball to the side. Lily helped him. Soon the ball came close and Lily picked it up. It was wet, but it was not broken.
+"Thank you, Tom," said Lily. "You are a good friend." Tom smiled and said, "You are welcome. Let's play again, but not near the pond." They both laughed and went back to the grass.
+When the sun went down, Lily's mom called her home. Lily said goodbye to Tom and walked home with her ball. She told her mom about the pond and the stick. Her mom gave her a big hug. "I am happy you and Tom helped each other," she said.
+That night, Lily put the ball next to her bed. She thought about the fun day at the park. She closed her eyes and smiled. Lily knew that sharing and helping made every day better.
+The next morning, Lily woke up early. She saw a little bird outside her window. The bird was singing a happy song. Lily wanted to give the bird some food. She asked her dad for some bread. Her dad gave her a small piece, and Lily put it on the window.
+The bird hopped over and ate the bread. Then it sang again, louder than before. Lily clapped her hands. "Thank you for the song, little bird," she said. The bird flew away, but it came back every morning after that.
+Lily and the bird became good friends. Every day she gave it bread, and every day it sang for her. Lily learned that being kind to others, even small birds, can make the world a happy place.
+EOF
+
+rows=()
+n_ok=0
+n_bad=0
+for entry in "${outputs[@]}"; do
+    read -r type max <<< "$entry"
+    out=$scratch_dir/stories15M-be-$type.gguf
+
+    if ! ppl=$(measure_ppl "$type" "$out"); then
+        ppl=-
+        status=FAIL
+    elif [ "$no_check" = --no-check ]; then
+        status=-
+    elif [[ $ppl =~ ^[0-9.]+$ ]] && awk -v p="$ppl" -v m="$max" 'BEGIN { exit !(p <= m) }'; then
+        status=OK
+    else
+        status=FAIL
+    fi
+
+    case $status in
+        OK)   n_ok=$((n_ok + 1)) ;;
+        FAIL) n_bad=$((n_bad + 1)) ;;
+    esac
+
+    rows+=("$type ${ppl:-nan} $max $status")
+    printf "%s: %-10s (%s) done\n" "$(basename "$0")" "$type" "$out"
+done
+
+printf "\n%-10s %10s %10s %s\n" type ppl max status
+for row in "${rows[@]}"; do
+    read -r type ppl max status <<< "$row"
+    case $status in
+        OK)   status="\033[32mOK\033[0m" ;;
+        FAIL) status="\033[31mFAIL\033[0m" ;;
+    esac
+    printf "%-10s %10s %10s %b\n" "$type" "$ppl" "$max" "$status"
+done
+
+if [ "$no_check" != --no-check ]; then
+    printf "summary: %d ok, %d failed\n" "$n_ok" "$n_bad"
+fi
+
+if [ "$n_bad" -ne 0 ]; then
+    printf "\nif the perplexity is expected to change, run with --no-check and update the max values in %s:\n\n" "$0"
+    printf "%s %s %s %s --no-check\n" "$0" "$quantize" "$perplexity" "$scratch_dir"
+    exit 1
+fi

tools/export-lora/export-lora.cpp

@@ -104,6 +104,9 @@ struct file_input {
         auto offset = gguf_get_data_offset(ctx_gguf) + gguf_get_tensor_offset(ctx_gguf, i_tensor_in);
         f_in.seekg(offset);
         f_in.read((char* )buf.data(), len);
+        if (!f_in) {
+            throw std::runtime_error("failed to read data for tensor " + name + ", input file may be truncated");
+        }
     }
 
     ~file_input() {

tools/gguf-split/gguf-split.cpp

@@ -362,6 +362,10 @@ struct split_strategy {
         }
         f_in.seekg(in_offset);
         f_in.read((char *)read_buf.data(), len);
+        if (!f_in) {
+            fprintf(stderr, "%s: failed to read %zu bytes at offset %zu, input file may be truncated\n", __func__, len, in_offset);
+            exit(EXIT_FAILURE);
+        }
         f_out.write((const char *)read_buf.data(), len);
     }
 };

tools/server/server-context.cpp

@@ -3781,6 +3781,9 @@ private:
 
                             llama_pos pos_next = slot.prompt.tokens.pos_next(n_past);
 
+                            // number of cached tokens shared with the new prompt, before any checkpoint is restored
+                            const int n_prefix = n_past;
+
                             // ref: https://github.com/ggml-org/llama.cpp/pull/24110
                             const bool has_new_tokens = (n_past < slot.task->n_tokens());
 
@@ -3891,10 +3894,10 @@ private:
                             }
 
                             {
-                                // erase any checkpoints with pos_max > pos_next
+                                // erase any checkpoints with pos_max > pos_next, or that cover tokens not shared with the new prompt
                                 for (auto it = slot.prompt.checkpoints.begin(); it != slot.prompt.checkpoints.end();) {
                                     const auto & cur = *it;
-                                    if (cur.pos_max > pos_next) {
+                                    if (cur.pos_max > pos_next || cur.n_tokens > n_prefix) {
                                         SLT_TRC(slot, "erased invalidated context checkpoint (pos_min = %d, pos_max = %d, n_tokens = %" PRId64 ", n_swa = %d, pos_next = %d, size = %.3f MiB)\n", cur.pos_min, cur.pos_max, cur.n_tokens, n_swa, pos_next, (float) cur.size() / 1024 / 1024);
                                         it = slot.prompt.checkpoints.erase(it);
                                     } else {

tools/server/tests/unit/test_completion.py

@@ -689,3 +689,81 @@ def test_completion_prompt_cache():
         assert "prompt_n" in timings and timings["prompt_n"] + timings["cache_n"] == n_prompt
         assert "predicted_n" in timings and timings["predicted_n"] == n_predict
         assert "tokens" in res.body and isinstance(res.body["tokens"], list)
+
+
+# a slot that switched to an unrelated prompt must not restore a context checkpoint of the previous one
+def test_completion_checkpoint_after_unrelated_prompt():
+    global server
+    server = ServerPreset.tinygemma3()
+    server.n_slots = 1
+    server.cache_ram = 0
+    server.temperature = 0.0
+    server.start()
+
+    def tokenize(content):
+        res = server.make_request("POST", "/tokenize", data={"content": content})
+        assert res.status_code == 200
+        return res.body["tokens"]
+
+    def top_logprobs(prompt, cache_prompt):
+        res = server.make_request("POST", "/completion", data={
+            "prompt": prompt,
+            "n_predict": 1,
+            "n_probs": 5,
+            "cache_prompt": cache_prompt,
+        })
+        assert res.status_code == 200
+        return [(p["id"], p["logprob"]) for p in res.body["completion_probabilities"][0]["top_logprobs"]]
+
+    tokens = tokenize("Purple elephants quietly compose sonnets about tax law")
+    expected = top_logprobs(tokens[:2], False)
+
+    # 5 tokens: the checkpoint 4 tokens before the end of the prompt holds only its first token
+    top_logprobs(tokenize("The quick brown fox jumps over the lazy dog")[:5], True)
+    top_logprobs(tokens[:1], True)
+    actual = top_logprobs(tokens[:2], True)
+
+    assert [tid for tid, _ in actual] == [tid for tid, _ in expected]
+    for (_, lp), (_, lp_expected) in zip(actual, expected):
+        assert lp == pytest.approx(lp_expected, abs=1e-3)
+
+
+# a restored context checkpoint must be kept for the next request that shares the same prefix
+def test_completion_checkpoint_kept_after_restore():
+    global server
+    server = ServerPreset.tinygemma3()
+    server.n_slots = 1
+    server.n_batch = 16
+    server.n_ubatch = 16
+    server.cache_ram = 0
+    server.temperature = 0.0
+    server.start()
+
+    def tokenize(content):
+        res = server.make_request("POST", "/tokenize", data={"content": content})
+        assert res.status_code == 200
+        return res.body["tokens"]
+
+    def cache_n(prompt):
+        res = server.make_request("POST", "/completion", data={
+            "prompt": prompt,
+            "n_predict": 1,
+            "cache_prompt": True,
+        })
+        assert res.status_code == 200
+        return res.body["timings"]["cache_n"]
+
+    prefix = tokenize("Purple elephants quietly compose sonnets about tax law in the moonlight")
+    suffixes = [tokenize(s) for s in [
+        "The quick brown fox jumps over the lazy dog while the farmer counts his sheep and the sun slowly sets behind the distant hills",
+        "Seven tiny robots repaired the broken lighthouse lens before the storm arrived, then danced on the rocks until the morning tide came in",
+        "Green apples rolled down the hill toward the river, where a curious otter collected them one by one and stacked them on a flat stone",
+    ]]
+
+    # 4 tokens: the checkpoint 4 tokens before the end of the prompt holds exactly the prefix
+    cache_n(prefix + suffixes[0][:4])
+
+    # suffixes longer than n_ubatch + 4: no new checkpoint within the prefix, only the restored one can be reused
+    for suffix in suffixes[1:]:
+        assert len(suffix) > server.n_ubatch + 4
+        assert cache_n(prefix + suffix) == len(prefix) - 1

tools/ui/src/lib/components/app/content/MarkdownContent/markdown-content.css

@@ -139,13 +139,13 @@
 /* Lists */
 .markdown-content :global(ul) {
 	list-style-type: disc;
-	margin-inline-start: 1.5rem;
+	margin-inline-start: 2rem;
 	margin-bottom: 1rem;
 }
 
 .markdown-content :global(ol) {
 	list-style-type: decimal;
-	margin-inline-start: 1.5rem;
+	margin-inline-start: 2rem;
 	margin-bottom: 1rem;
 }
 
@@ -500,6 +500,8 @@ div.markdown-user-content :global(.table-wrapper) {
 	margin: 1.5rem 0;
 	border-radius: 0.5rem;
 	border: 1px solid var(--border);
+	/* a wide table scrolls here, it does not widen the column */
+	max-width: 100%;
 }
 
 .markdown-content :global(.table-wrapper table) {
@@ -530,11 +532,9 @@ div.markdown-user-content :global(.table-wrapper) {
 		padding: 0.375rem 0.5rem;
 	}
 
+	/* no full-bleed: it hid the outer columns */
 	.markdown-content :global(.table-wrapper) {
-		margin: 0.5rem -1rem;
-		border-radius: 0;
-		border-left: none;
-		border-right: none;
+		margin: 1rem 0;
 	}
 }
 

tools/ui/tests/stories/ChatMessage.stories.svelte

@@ -1,4 +1,5 @@
 <script lang="ts" module>
+	import { JOKES_TABLE_MD, WIDE_CELL_TABLE_MD } from './fixtures/mobile-tables.js';
 	import { defineMeta } from '@storybook/addon-svelte-csf';
 	import ChatMessage from '$lib/components/app/chat/ChatMessages/ChatMessage/ChatMessage.svelte';
 	import type { ChatMessageActions } from '$lib/types';
@@ -49,6 +50,18 @@
 		type: 'message'
 	};
 
+	const assistantWithTable: DatabaseMessage = {
+		children: [],
+		content: `${JOKES_TABLE_MD}\n\n${WIDE_CELL_TABLE_MD}`,
+		convId: 'conv-1',
+		id: '7',
+		parent: '1',
+		role: 'assistant',
+		thinking: '',
+		timestamp: Date.now() - 1000 * 60 * 3,
+		type: 'message'
+	};
+
 	const assistantWithReasoning: DatabaseMessage = {
 		children: [],
 		content: "Here's the concise answer, now that I've thought it through carefully for you.",
@@ -124,6 +137,19 @@
 	}}
 />
 
+<Story
+	args={{
+		class: 'max-w-[56rem] w-[calc(100vw-2rem)]',
+		message: assistantWithTable
+	}}
+	name="AssistantWithTable"
+	play={async () => {
+		const { settingsStore } = await import('$lib/stores/settings/index.svelte');
+
+		settingsStore.updateConfig('showRawOutputSwitch', false);
+	}}
+/>
+
 <Story
 	args={{
 		class: 'max-w-[56rem] w-[calc(100vw-2rem)]',

tools/ui/tests/stories/fixtures/mobile-tables.ts

@@ -0,0 +1,27 @@
+// Table-heavy answers, to check the table layout at a phone width.
+export const JOKES_TABLE_MD = String.raw`
+Sure! Here are the jokes in a table:
+
+| No. | Joke |
+|-----|------|
+| 1 | Why don't scientists trust atoms? Because they make up everything. |
+| 2 | I told my wife she was drawing her eyebrows too high. She looked surprised. |
+| 3 | Why did the scarecrow win an award? Because he was outstanding in his field. |
+| 4 | I'm reading a book about anti-gravity. It's impossible to put down. |
+| 5 | Why don't skeletons fight each other? They don't have the guts. |
+| 6 | I used to hate clowns, but they all have their fat-sies. |
+| 7 | What do you call a fake noodle? An impasta. |
+| 8 | Why did the bicycle fall over? Because it was two-tired. |
+| 9 | I would tell you a chemistry joke, but I know I wouldn't get a reaction. |
+| 10 | Why did the golfer bring two pairs of pants? In case he got a hole in one. |
+`;
+
+export const WIDE_CELL_TABLE_MD = String.raw`
+Some endpoints, in a table, for reference:
+
+| Name | Endpoint | Description |
+|------|----------|-------------|
+| Completions | https://huggingface.co/api/models/ggml-org/gpt-oss-20b-GGUF/summary | Long endpoint description |
+| Models | ./v1/models | Lists the models loaded in the server slot |
+| Props | ./props | Server capabilities and default sampling |
+`;

vendor/cpp-httplib/httplib.h

@@ -8,8 +8,8 @@
 #ifndef CPPHTTPLIB_HTTPLIB_H
 #define CPPHTTPLIB_HTTPLIB_H
 
-#define CPPHTTPLIB_VERSION "0.61.0"
-#define CPPHTTPLIB_VERSION_NUM "0x003d00"
+#define CPPHTTPLIB_VERSION "0.62.0"
+#define CPPHTTPLIB_VERSION_NUM "0x003e00"
 
 #ifdef _WIN32
 #if defined(_WIN32_WINNT) && _WIN32_WINNT < 0x0A00
@@ -18,6 +18,14 @@
 #endif
 #endif
 
+#if defined(__SANITIZE_THREAD__)
+#define CPPHTTPLIB_THREAD_SANITIZER
+#elif defined(__has_feature)
+#if __has_feature(thread_sanitizer)
+#define CPPHTTPLIB_THREAD_SANITIZER
+#endif
+#endif
+
 /*
  * Configuration
  */
@@ -122,6 +130,10 @@
 #define CPPHTTPLIB_REDIRECT_MAX_COUNT 20
 #endif
 
+#ifndef CPPHTTPLIB_INTERIM_RESPONSE_MAX_COUNT
+#define CPPHTTPLIB_INTERIM_RESPONSE_MAX_COUNT 100
+#endif
+
 #ifndef CPPHTTPLIB_MULTIPART_FORM_DATA_FILE_MAX_COUNT
 #define CPPHTTPLIB_MULTIPART_FORM_DATA_FILE_MAX_COUNT 1024
 #endif
@@ -130,6 +142,10 @@
 #define CPPHTTPLIB_PAYLOAD_MAX_LENGTH (100 * 1024 * 1024) // 100MB
 #endif
 
+#ifndef CPPHTTPLIB_SSE_MAX_EVENT_LENGTH
+#define CPPHTTPLIB_SSE_MAX_EVENT_LENGTH (16 * 1024 * 1024) // 16MB
+#endif
+
 #ifndef CPPHTTPLIB_FORM_URL_ENCODED_PAYLOAD_MAX_LENGTH
 #define CPPHTTPLIB_FORM_URL_ENCODED_PAYLOAD_MAX_LENGTH 8192
 #endif
@@ -211,6 +227,16 @@
 #define CPPHTTPLIB_MAX_LINE_LENGTH 32768
 #endif
 
+// These bound the parsing of every request, so 0 cannot lift them.
+static_assert(CPPHTTPLIB_REQUEST_URI_MAX_LENGTH > 0,
+              "CPPHTTPLIB_REQUEST_URI_MAX_LENGTH must be positive");
+static_assert(CPPHTTPLIB_HEADER_MAX_LENGTH > 0,
+              "CPPHTTPLIB_HEADER_MAX_LENGTH must be positive");
+static_assert(CPPHTTPLIB_HEADER_MAX_COUNT > 0,
+              "CPPHTTPLIB_HEADER_MAX_COUNT must be positive");
+static_assert(CPPHTTPLIB_MAX_LINE_LENGTH > 0,
+              "CPPHTTPLIB_MAX_LINE_LENGTH must be positive");
+
 #ifndef CPPHTTPLIB_WEBSOCKET_MAX_PAYLOAD_LENGTH
 #define CPPHTTPLIB_WEBSOCKET_MAX_PAYLOAD_LENGTH 16777216
 #endif
@@ -340,7 +366,6 @@ using socklen_t = int;
 #include <netinet/tcp.h>
 #include <poll.h>
 #include <pthread.h>
-#include <sys/mman.h>
 #include <sys/socket.h>
 #include <sys/un.h>
 #include <unistd.h>
@@ -451,6 +476,7 @@ using socket_t = int;
 
 #include <openssl/err.h>
 #include <openssl/evp.h>
+#include <openssl/rand.h>
 #include <openssl/ssl.h>
 #include <openssl/x509v3.h>
 
@@ -581,6 +607,10 @@ namespace ws {
 class WebSocket;
 } // namespace ws
 
+namespace sse {
+class SSEClient;
+} // namespace sse
+
 namespace detail {
 
 /*
@@ -660,6 +690,15 @@ inline bool equal(const std::string &a, const std::string &b) {
          });
 }
 
+// For a literal, without building a std::string from it.
+template <size_t N>
+inline bool equal(const std::string &a, const char (&b)[N]) {
+  return a.size() == N - 1 &&
+         std::equal(a.begin(), a.end(), b, [](char ca, char cb) {
+           return to_lower(ca) == to_lower(cb);
+         });
+}
+
 struct equal_to {
   bool operator()(const std::string &a, const std::string &b) const {
     return equal(a, b);
@@ -668,17 +707,13 @@ struct equal_to {
 
 struct hash {
   size_t operator()(const std::string &key) const {
-    return hash_core(key.data(), key.size(), 0);
-  }
-
-  size_t hash_core(const char *s, size_t l, size_t h) const {
-    return (l == 0) ? h
-                    : hash_core(s + 1, l - 1,
-                                // Unsets the 6 high bits of h, therefore no
-                                // overflow happens
-                                (((std::numeric_limits<size_t>::max)() >> 6) &
-                                 h * 33) ^
-                                    static_cast<unsigned char>(to_lower(*s)));
+    size_t h = 0;
+    for (auto c : key) {
+      // Unsets the 6 high bits of h, therefore no overflow happens
+      h = (((std::numeric_limits<size_t>::max)() >> 6) & h * 33) ^
+          static_cast<unsigned char>(to_lower(c));
+    }
+    return h;
   }
 };
 
@@ -736,6 +771,7 @@ inline from_chars_result<T> from_chars(const char *first, const char *last,
   }
   if (p == last) { return {first, std::errc::invalid_argument}; }
 
+  auto radix = static_cast<T>(base);
   T result = 0;
   for (; p != last; ++p) {
     char c = *p;
@@ -751,10 +787,11 @@ inline from_chars_result<T> from_chars(const char *first, const char *last,
     }
 
     if (digit < 0 || digit >= base) { break; }
-    if (result > ((std::numeric_limits<T>::max)() - digit) / base) {
+    auto d = static_cast<T>(digit);
+    if (result > ((std::numeric_limits<T>::max)() - d) / radix) {
       return {p, std::errc::result_out_of_range};
     }
-    result = result * base + digit;
+    result = result * radix + d;
   }
 
   if (p == first || (negative && p == first + 1)) {
@@ -881,6 +918,12 @@ inline bool parse_url(const std::string &url, UrlComponents &uc) {
   auto has_authority = has_authority_prefix || (!url.empty() && url[0] != '/' &&
                                                 url[0] != '?' && url[0] != '#');
   if (has_authority) {
+    // userinfo is not supported, and a backslash that some parsers take for
+    // '/' would leave them seeing another host than this one.
+    if (url.find_first_of("@\\", pos) < url.find_first_of("/?#", pos)) {
+      return false;
+    }
+
     if (pos < url.size() && url[pos] == '[') {
       auto close = url.find(']', pos);
       if (close == std::string::npos) { return false; }
@@ -1803,6 +1846,7 @@ struct Request {
   bool is_multipart_form_data() const;
 
   // private members...
+  bool body_read_started_ = false;
   bool body_consumed_ = false;
   bool expect_100_continue_pending_ = false;
   size_t redirect_count_ = CPPHTTPLIB_REDIRECT_MAX_COUNT;
@@ -1972,6 +2016,16 @@ public:
     (void)usec;
   }
 
+  // Bounds the reads from now on to `msec` in total, on top of the timeout of
+  // each one. 0 lifts the bound.
+  virtual void set_max_timeout(time_t msec) { (void)msec; }
+
+  // Lets a wait for data end early once `interrupted` returns true, so that
+  // another thread can stop a blocked read on every platform.
+  virtual void set_read_interrupt(std::function<bool()> interrupted) {
+    (void)interrupted;
+  }
+
   // Bytes already pulled off the socket and sitting in this stream's own
   // buffer. Exposing them lets a line reader scan for a terminator in one
   // pass instead of asking for a byte at a time. A stream that does no
@@ -2425,7 +2479,8 @@ private:
   int bind_internal(const std::string &host, int port, int socket_flags);
   bool listen_internal();
 
-  bool routing(Request &req, Response &res, Stream &strm);
+  bool routing(Request &req, Response &res, Stream &strm,
+               bool skip_pre_routing);
   bool handle_file_request(Request &req, Response &res);
   bool check_if_not_modified(const Request &req, Response &res,
                              const std::string &etag, time_t mtime) const;
@@ -2724,6 +2779,7 @@ public:
 
   private:
     friend class ClientImpl;
+    friend class sse::SSEClient;
 
     ssize_t read_with_decompression(char *buf, size_t len);
 
@@ -2941,14 +2997,13 @@ protected:
   void close_socket(Socket &socket);
   void disconnect(bool gracefully);
 
+  std::string make_request_target(const Request &req) const;
   bool process_request(Stream &strm, Request &req, Response &res,
                        bool close_connection, Error &error);
 
   bool write_content_with_provider(Stream &strm, const Request &req,
                                    Error &error) const;
 
-  void copy_settings(const ClientImpl &rhs);
-
   void output_log(const Request &req, const Response &res) const;
   void output_error_log(const Error &err, const Request *req) const;
 
@@ -2960,6 +3015,9 @@ protected:
   Socket socket_;
   mutable std::mutex socket_mutex_;
   std::recursive_mutex request_mutex_;
+  // How many send() calls are running on this client, guarded by
+  // request_mutex_
+  size_t send_depth_ = 0;
 
   // These are all protected under socket_mutex
   size_t socket_requests_in_flight_ = 0;
@@ -3071,6 +3129,9 @@ private:
   virtual bool is_ssl() const;
 
   void transfer_socket_ownership_to_handle(StreamHandle &handle);
+  StreamHandle open_stream_(Request &req);
+  void add_credential_headers(Request &req);
+  std::string proxy_target_prefix() const;
 
 #ifdef CPPHTTPLIB_SSL_ENABLED
 public:
@@ -3100,6 +3161,24 @@ protected:
 #endif
 };
 
+namespace detail {
+
+// Stands in for a client that Client could not make from its arguments, so a
+// request fails rather than going through a null pointer.
+class InvalidClientImpl final : public ClientImpl {
+public:
+  InvalidClientImpl() : ClientImpl(std::string(), 0) {}
+  bool is_valid() const override { return false; }
+
+private:
+  bool create_and_connect_socket(Socket &, Error &error) override {
+    error = Error::Connection;
+    return false;
+  }
+};
+
+} // namespace detail
+
 class Client {
 public:
   // Universal interface
@@ -3364,7 +3443,14 @@ private:
   bool process_and_close_socket(socket_t sock) override;
 
   tls::ctx_t ctx_ = nullptr;
+  // update_certs_pem() replaces the certificates and CAs that handshakes read
+  // from ctx_, so it waits for the handshakes in progress, which their timeout
+  // bounds, and new ones wait while it is pending so that they cannot keep it
+  // waiting.
   std::mutex ctx_mutex_;
+  std::condition_variable ctx_cv_;
+  size_t handshakes_in_progress_ = 0;
+  size_t updates_pending_ = 0;
 
   int last_ssl_error_ = 0;
 };
@@ -3446,11 +3532,12 @@ private:
 
   std::function<SSLVerifierResponse(tls::session_t)> session_verifier_;
 
+  // Like ca_cert_store_set_, keeps what ctx_ cannot report back: the verifier
+  // set_server_certificate_verifier() installed there.
+  tls::VerifyCallback server_certificate_verifier_;
+
 #ifdef CPPHTTPLIB_WINDOWS_AUTOMATIC_ROOT_CERTIFICATES_UPDATE
   bool enable_windows_cert_verification_ = true;
-  // Like ca_cert_store_set_, tracks what ctx_ cannot report back: whether
-  // set_server_certificate_verifier() installed a verifier.
-  bool server_certificate_verifier_set_ = false;
 #endif
 
   friend class ClientImpl;
@@ -3571,7 +3658,7 @@ inline void ClientImpl::set_max_timeout(
     const std::chrono::duration<Rep, Period> &duration) {
   auto msec =
       std::chrono::duration_cast<std::chrono::milliseconds>(duration).count();
-  set_max_timeout(msec);
+  set_max_timeout(static_cast<time_t>(msec));
 }
 
 template <class Rep, class Period>
@@ -3937,6 +4024,8 @@ struct ChunkedDecoder {
   bool parse_trailers_into(Headers &dest, const Headers &src_headers);
 };
 
+// A regular file opened to be served: mapped on Windows, read with pread()
+// elsewhere.
 class mmap {
 public:
   mmap(const char *path);
@@ -3947,18 +4036,20 @@ public:
 
   bool is_open() const;
   size_t size() const;
-  const char *data() const;
+  // Copies up to len bytes from offset into buf and returns how many: fewer,
+  // down to 0, if the file has shrunk since it was opened.
+  size_t read(size_t offset, char *buf, size_t len) const;
 
 private:
 #if defined(_WIN32)
   HANDLE hFile_ = NULL;
   HANDLE hMapping_ = NULL;
+  void *addr_ = nullptr;
+  bool is_open_empty_file = false;
 #else
   int fd_ = -1;
 #endif
   size_t size_ = 0;
-  void *addr_ = nullptr;
-  bool is_open_empty_file = false;
 };
 
 // NOTE: https://www.rfc-editor.org/rfc/rfc9110#section-5
@@ -4003,8 +4094,10 @@ struct MbedTlsContext {
   mbedtls_x509_crt own_cert;
   mbedtls_pk_context own_key;
   bool is_server = false;
+  // Peer certificate required (set_verify_client()); a client defaults to true
   bool verify_client = false;
-  bool has_verify_callback = false;
+  // Set by tls::set_verify_callback().
+  VerifyCallback verify_callback;
 
   MbedTlsContext();
   ~MbedTlsContext();
@@ -4027,8 +4120,10 @@ namespace impl {
 struct WolfSSLContext {
   WOLFSSL_CTX *ctx = nullptr;
   bool is_server = false;
+  // Peer certificate required (set_verify_client()); a client defaults to true
   bool verify_client = false;
-  bool has_verify_callback = false;
+  // Replaced under verify_callback_mutex()
+  std::shared_ptr<VerifyCallback> verify_callback;
   std::string ca_pem_data_; // accumulated PEM for get_ca_names/get_ca_certs
 
   WolfSSLContext();
@@ -4453,7 +4548,8 @@ enum ReadResult : int { Fail = 0, Text = 1, Binary = 2, Timeout = 3 };
 // Result of WebSocketClient::connect(). Truthy only when the WebSocket
 // upgrade handshake fully succeeded. On failure error() identifies the
 // failing layer; status()/headers() expose the server's upgrade response
-// when one was received (status() is -1 otherwise).
+// when one was received (status() is -1 otherwise), or the proxy's reply to
+// CONNECT when error() is ProxyConnection.
 class Result {
 public:
   Result() = default;
@@ -4555,6 +4651,8 @@ private:
   int max_missed_pongs_;
   int unacked_pings_ = 0;
   std::atomic<bool> closed_{false};
+  // Set when the peer stopped answering pings, to end a read() waiting on it
+  std::atomic<bool> pong_timed_out_{false};
   // Set once the caller has bounded read() through set_read_timeout(). Until
   // then the timeout in effect is the compile-time default, and elapsing it
   // is a failure that closes the connection, not a Timeout.
@@ -4612,6 +4710,12 @@ public:
   void set_interface(const std::string &intf);
   void set_hostname_addr_map(std::map<std::string, std::string> addr_map);
 
+  void set_proxy(const std::string &host, int port);
+  void set_proxy_basic_auth(const std::string &username,
+                            const std::string &password);
+  void set_proxy_bearer_token_auth(const std::string &token);
+  void set_no_proxy(const std::vector<std::string> &patterns);
+
 #ifdef CPPHTTPLIB_SSL_ENABLED
   struct PemMemory {
     const char *cert_pem;
@@ -4630,13 +4734,23 @@ public:
   void enable_server_certificate_verification(bool enabled);
   void enable_server_hostname_verification(bool enabled);
   void enable_system_ca(bool enabled);
+  void set_server_certificate_verifier(tls::VerifyCallback verifier);
+  void set_session_verifier(
+      std::function<SSLVerifierResponse(tls::session_t)> verifier);
+  void set_proxy_digest_auth(const std::string &username,
+                             const std::string &password);
 #endif
 
 private:
   void shutdown_and_close();
+  std::pair<time_t, time_t> handshake_timeout() const;
   bool create_stream(std::unique_ptr<Stream> &strm, Error &error,
                      int &ssl_error, uint64_t &ssl_backend_error);
   void prepare_default_headers(Request &req);
+  bool is_proxy_enabled() const;
+  bool create_and_connect_socket(Error &error);
+  bool send_proxy_connect(Request &req, Response &proxy_res, Error &error);
+  bool connect_with_proxy(Response &proxy_res, Error &error);
 
   std::string host_;
   int port_;
@@ -4666,6 +4780,13 @@ private:
   // hostname; only the connection target changes, never the identity.
   std::map<std::string, std::string> addr_map_;
 
+  std::string proxy_host_;
+  int proxy_port_ = -1;
+  std::string proxy_basic_auth_username_;
+  std::string proxy_basic_auth_password_;
+  std::string proxy_bearer_token_auth_token_;
+  std::vector<detail::NoProxyEntry> no_proxy_entries_;
+
 #ifdef CPPHTTPLIB_SSL_ENABLED
   bool is_ssl_ = false;
   tls::ctx_t tls_ctx_ = nullptr;
@@ -4677,6 +4798,9 @@ private:
   SystemCAMode system_ca_mode_ = SystemCAMode::Auto;
   bool server_certificate_verification_ = true;
   bool server_hostname_verification_ = true;
+  std::function<SSLVerifierResponse(tls::session_t)> session_verifier_;
+  std::string proxy_digest_auth_username_;
+  std::string proxy_digest_auth_password_;
 #endif
 };