ggml-webgpu: improve flash_attn_vec for quantized KV at long contexts (#25956 )

* improve fa of quantized kv cache * Fix some bugs and some comments. * fix v type check and some comments * Fix build error caused by rebasing * editorconfig checking pass
mtmd: add lanczos resize method [no release] (#26341 )
2026-07-31 16:48:01 +02:00 · 2026-07-31 09:08:40 +03:00 · 2026-07-30 21:59:49 +02:00 · 2026-07-30 21:40:38 +02:00 · 2026-07-30 13:11:32 -05:00 · 2026-07-31 01:14:16 +08:00
15 changed files with 484 additions and 188 deletions
@@ -469,6 +469,8 @@ static bool ggml_backend_cpu_device_supports_op(ggml_backend_dev_t dev, const st
            return (src0->type == GGML_TYPE_F32 ||
                    ((src0->type == GGML_TYPE_F16 || ggml_is_quantized(src0->type)) && src0->ne[2] == src1->ne[2] && src0->ne[3] == src1->ne[3])) &&
                src1->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32;
+        case GGML_OP_CONV_2D:
+            return ggml_is_contiguous(op->src[0]);
        default:
            return true;
    }
@@ -126,6 +126,7 @@ void ggml_cuda_op_conv2d(ggml_backend_cuda_context & ctx, ggml_tensor * dst) {
    const float *       X_D    = (const float *) input->data;
    float *             Y_D    = (float *) dst->data;

+    GGML_ASSERT(ggml_is_contiguous(input));
    GGML_ASSERT(ggml_is_contiguous(kernel));
    GGML_ASSERT(kernel->type == GGML_TYPE_F16 || kernel->type == GGML_TYPE_F32);

@@ -5105,7 +5105,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
        case GGML_OP_IM2COL:
        case GGML_OP_IM2COL_3D:
        case GGML_OP_CONV_2D:
-            return true;
+            return (ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1]));
        case GGML_OP_CONV_2D_DW:
            return op->src[0]->type == GGML_TYPE_F32;
        case GGML_OP_CONV_TRANSPOSE_2D:
@@ -1481,6 +1481,11 @@ struct vk_op_binary_push_constants {
    float param1; float param2; int32_t param3;
 };

+// Distinct type with the same layout so concat can overload tensor offset initialization.
+struct vk_op_concat_push_constants : vk_op_binary_push_constants {};
+static_assert(sizeof(vk_op_concat_push_constants) == sizeof(vk_op_binary_push_constants));
+static_assert(std::is_standard_layout_v<vk_op_concat_push_constants>);
+
 struct vk_op_multi_add_push_constants {
    // shape for dst
    uint32_t ne20; uint32_t ne21; uint32_t ne22; uint32_t ne23;
@@ -2246,6 +2251,40 @@ static uint32_t get_misalign_bytes(const ggml_backend_vk_context * ctx, const gg
    return ((vk_tensor_offset(t) + t->view_offs) & (ctx->device->properties.limits.minStorageBufferOffsetAlignment - 1));;
 }

+static uint32_t ggml_vk_concat_unit_size(ggml_type type) {
+    const uint32_t type_size = ggml_type_size(type);
+
+    if (!ggml_is_quantized(type)) {
+        return type_size;
+    }
+
+    // Use the widest existing concat shader that evenly divides a quant block.
+    if (type_size % 8 == 0) {
+        return 8;
+    }
+    if (type_size % 4 == 0) {
+        return 4;
+    }
+    if (type_size % 2 == 0) {
+        return 2;
+    }
+    return 1;
+}
+
+static bool ggml_vk_concat_supported(const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst) {
+    if (src0->type != src1->type || src0->type != dst->type) {
+        return false;
+    }
+
+    if (!ggml_is_quantized(src0->type)) {
+        const size_t type_size = ggml_type_size(src0->type);
+        return type_size == 1 || type_size == 2 || type_size == 4 || type_size == 8;
+    }
+
+    // Quantized tensor rows are block-aligned when created.
+    return ggml_is_contiguous_rows(src0) && ggml_is_contiguous_rows(src1) && ggml_is_contiguous_rows(dst);
+}
+
 template <typename T> void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, T &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, const ggml_tensor * src3, ggml_tensor * dst) {
    GGML_UNUSED(p);
    GGML_UNUSED(src0);
@@ -10896,14 +10935,10 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const
        }
        return nullptr;
    case GGML_OP_CONCAT: {
-        if (src0->type != src1->type || src0->type != dst->type) {
+        if (!ggml_vk_concat_supported(src0, src1, dst)) {
            return nullptr;
        }
-        if (ggml_blck_size(src0->type) != 1) {
-            return nullptr;
-        }
-        const size_t type_size = ggml_type_size(src0->type);
-        switch (type_size) {
+        switch (ggml_vk_concat_unit_size(src0->type)) {
        case 1:
            return ctx->device->pipeline_concat_i8;
        case 2:
@@ -11595,6 +11630,18 @@ template <> void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk
    GGML_UNUSED(src3);
 }

+template <> void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk_op_concat_push_constants &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, const ggml_tensor * src3, ggml_tensor * dst) {
+    const uint32_t unit_size = ggml_vk_concat_unit_size(dst->type);
+    const uint32_t a_offset = get_misalign_bytes(ctx, src0) / unit_size;
+    const uint32_t b_offset = get_misalign_bytes(ctx, src1) / unit_size;
+    const uint32_t d_offset = get_misalign_bytes(ctx, dst) / unit_size;
+
+    p.misalign_offsets = (a_offset << 16) | (b_offset << 8) | d_offset;
+
+    GGML_UNUSED(src2);
+    GGML_UNUSED(src3);
+}
+
 template <> void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk_op_upscale_push_constants &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, const ggml_tensor * src3, ggml_tensor * dst) {
    const uint32_t a_offset = get_misalign_bytes(ctx, src0) / ggml_type_size(src0->type);
    const uint32_t d_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type);
@@ -11630,7 +11677,7 @@ static void ggml_vk_op_f32(ggml_backend_vk_context * ctx, vk_context& subctx, co
    }
    std::cerr << "), (" << dst << ", name=" << dst->name << ", type=" << dst->type << ", ne0=" << dst->ne[0] << ", ne1=" << dst->ne[1] << ", ne2=" << dst->ne[2] << ", ne3=" << dst->ne[3] << ", nb0=" << dst->nb[0] << ", nb1=" << dst->nb[1] << ", nb2=" << dst->nb[2] << ", nb3=" << dst->nb[3];
    std::cerr << "), " << ggml_op_name(op) << ")");
-    GGML_ASSERT(op == GGML_OP_GET_ROWS || op == GGML_OP_CPY || (!ggml_is_quantized(src0->type) && (src1 == nullptr || !ggml_is_quantized(src1->type))));  // NOLINT
+    GGML_ASSERT(op == GGML_OP_GET_ROWS || op == GGML_OP_CPY || op == GGML_OP_CONCAT || (!ggml_is_quantized(src0->type) && (src1 == nullptr || !ggml_is_quantized(src1->type))));  // NOLINT
    GGML_ASSERT(dst->buffer != nullptr);
    const uint64_t ne00 = src0->ne[0];
    const uint64_t ne01 = src0->ne[1];
@@ -11885,6 +11932,9 @@ static void ggml_vk_op_f32(ggml_backend_vk_context * ctx, vk_context& subctx, co
                    ne *= ggml_type_size(src0->type) / 2;
                }
            }
+            if (op == GGML_OP_CONCAT && ggml_is_quantized(dst->type)) {
+                ne = ne / ggml_blck_size(dst->type) * ggml_type_size(dst->type) / ggml_vk_concat_unit_size(dst->type);
+            }
            // copy_to_quant has block size of 32, and each thread does QUANT_K elements.
            // Splitting into 512x512xZ wouldn't work well since each workgroup does 1024 elements.
            // So divide by block size here before splitting into 512x512 groups.
@@ -12525,18 +12575,28 @@ static void ggml_vk_opt_step_sgd(ggml_backend_vk_context * ctx, vk_context& subc
 static void ggml_vk_concat(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) {
    int * op_params = (int *)dst->op_params;

-    const uint32_t src0_type_size = ggml_type_size(src0->type);
-    const uint32_t src1_type_size = ggml_type_size(src1->type);
-    const uint32_t dst_type_size = ggml_type_size(dst->type);
+    const uint32_t unit_size = ggml_vk_concat_unit_size(dst->type);
+    const uint32_t units_per_block = ggml_type_size(dst->type) / unit_size;
+    const uint32_t block_size = ggml_blck_size(dst->type);
+    const bool quantized = ggml_is_quantized(dst->type);

-    ggml_vk_op_f32<vk_op_binary_push_constants>(ctx, subctx, src0, src1, nullptr, nullptr, dst, GGML_OP_CONCAT, {
-        (uint32_t)ggml_nelements(dst),
-        (uint32_t)src0->ne[0], (uint32_t)src0->ne[1], (uint32_t)src0->ne[2],(uint32_t)src0->ne[3], (uint32_t)src0->nb[0] / src0_type_size, (uint32_t)src0->nb[1] / src0_type_size, (uint32_t)src0->nb[2] / src0_type_size, (uint32_t)src0->nb[3] / src0_type_size,
-        (uint32_t)src1->ne[0], (uint32_t)src1->ne[1], (uint32_t)src1->ne[2],(uint32_t)src1->ne[3], (uint32_t)src1->nb[0] / src1_type_size, (uint32_t)src1->nb[1] / src1_type_size, (uint32_t)src1->nb[2] / src1_type_size, (uint32_t)src1->nb[3] / src1_type_size,
-        (uint32_t) dst->ne[0], (uint32_t) dst->ne[1], (uint32_t) dst->ne[2],(uint32_t) dst->ne[3], (uint32_t) dst->nb[0] /  dst_type_size, (uint32_t) dst->nb[1] /  dst_type_size, (uint32_t) dst->nb[2] /  dst_type_size, (uint32_t) dst->nb[3] /  dst_type_size,
+    // Address dimension 0 in packed storage units; higher strides may be noncontiguous.
+    const uint32_t ne00 = src0->ne[0] / block_size * units_per_block;
+    const uint32_t ne10 = src1->ne[0] / block_size * units_per_block;
+    const uint32_t ne20 =  dst->ne[0] / block_size * units_per_block;
+    const uint32_t nb00 = quantized ? 1 : src0->nb[0] / unit_size;
+    const uint32_t nb10 = quantized ? 1 : src1->nb[0] / unit_size;
+    const uint32_t nb20 = quantized ? 1 :  dst->nb[0] / unit_size;
+
+    vk_op_concat_push_constants pc {{
+        ne20 * (uint32_t)dst->ne[1] * (uint32_t)dst->ne[2] * (uint32_t)dst->ne[3],
+        ne00, (uint32_t)src0->ne[1], (uint32_t)src0->ne[2],(uint32_t)src0->ne[3], nb00, (uint32_t)src0->nb[1] / unit_size, (uint32_t)src0->nb[2] / unit_size, (uint32_t)src0->nb[3] / unit_size,
+        ne10, (uint32_t)src1->ne[1], (uint32_t)src1->ne[2],(uint32_t)src1->ne[3], nb10, (uint32_t)src1->nb[1] / unit_size, (uint32_t)src1->nb[2] / unit_size, (uint32_t)src1->nb[3] / unit_size,
+        ne20, (uint32_t) dst->ne[1], (uint32_t) dst->ne[2],(uint32_t) dst->ne[3], nb20, (uint32_t) dst->nb[1] / unit_size, (uint32_t) dst->nb[2] / unit_size, (uint32_t) dst->nb[3] / unit_size,
        0,
        0.0f, 0.0f, op_params[0],
-    });
+    }};
+    ggml_vk_op_f32<vk_op_concat_push_constants>(ctx, subctx, src0, src1, nullptr, nullptr, dst, GGML_OP_CONCAT, std::move(pc));
 }

 static void ggml_vk_upscale(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst) {
@@ -17872,12 +17932,7 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
            return op->src[0]->type == op->src[1]->type && op->src[0]->type == op->type &&
                   (op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_I32);
        case GGML_OP_CONCAT: {
-            if (op->src[0]->type != op->src[1]->type || op->src[0]->type != op->type) {
-                return false;
-            }
-            const size_t type_size = ggml_type_size(op->type);
-            return ggml_blck_size(op->type) == 1 &&
-                   (type_size == 1 || type_size == 2 || type_size == 4 || type_size == 8);
+            return ggml_vk_concat_supported(op->src[0], op->src[1], op);
        }
        case GGML_OP_ADD1:
            return (op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32)
@@ -18016,10 +18071,17 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm
        case GGML_OP_CONV_2D:
        case GGML_OP_CONV_TRANSPOSE_2D:
            {
+                const bool transpose = op->op == GGML_OP_CONV_TRANSPOSE_2D;
+                const int64_t cout = !transpose ? op->src[0]->ne[3] : op->src[0]->ne[2];
+                const int64_t cin  = !transpose ? op->src[0]->ne[2] : op->src[0]->ne[3];
+
                // Channel-contiguous format is not supported yet.
                return ((op->src[0]->type == GGML_TYPE_F32 || op->src[0]->type == GGML_TYPE_F16) &&
+                    (op->src[0]->nb[0] == sizeof(float) || op->src[0]->nb[0] == sizeof(ggml_fp16_t) ) &&
                    op->src[1]->type == GGML_TYPE_F32 &&
                    op->type == GGML_TYPE_F32 &&
+                    cout == op->ne[2] &&
+                    cin == op->src[1]->ne[2] &&
                    ggml_is_contiguous(op->src[0]) &&
                    ggml_is_contiguous(op->src[1]) &&
                    ggml_is_contiguous(op));
@@ -591,7 +591,8 @@ struct ggml_webgpu_flash_attn_common_pipeline_key {
    ggml_type dst_type;
    uint32_t  head_dim_qk;
    uint32_t  head_dim_v;
-    bool      kv_direct;
+    bool      k_direct;
+    bool      v_direct;
    bool      kv_overlap;
    bool      has_mask;
    bool      has_sinks;
@@ -600,8 +601,9 @@ struct ggml_webgpu_flash_attn_common_pipeline_key {
    bool operator==(const ggml_webgpu_flash_attn_common_pipeline_key & other) const {
        return q_type == other.q_type && k_type == other.k_type && v_type == other.v_type &&
               dst_type == other.dst_type && head_dim_qk == other.head_dim_qk && head_dim_v == other.head_dim_v &&
-               kv_direct == other.kv_direct && kv_overlap == other.kv_overlap && has_mask == other.has_mask &&
-               has_sinks == other.has_sinks && uses_logit_softcap == other.uses_logit_softcap;
+               k_direct == other.k_direct && v_direct == other.v_direct && kv_overlap == other.kv_overlap &&
+               has_mask == other.has_mask && has_sinks == other.has_sinks &&
+               uses_logit_softcap == other.uses_logit_softcap;
    }
 };

@@ -613,7 +615,8 @@ inline void ggml_webgpu_flash_attn_hash_common_pipeline_key(size_t &
    ggml_webgpu_hash_combine(seed, key.dst_type);
    ggml_webgpu_hash_combine(seed, key.head_dim_qk);
    ggml_webgpu_hash_combine(seed, key.head_dim_v);
-    ggml_webgpu_hash_combine(seed, key.kv_direct);
+    ggml_webgpu_hash_combine(seed, key.k_direct);
+    ggml_webgpu_hash_combine(seed, key.v_direct);
    ggml_webgpu_hash_combine(seed, key.kv_overlap);
    ggml_webgpu_hash_combine(seed, key.has_mask);
    ggml_webgpu_hash_combine(seed, key.has_sinks);
@@ -687,12 +690,13 @@ inline bool ggml_webgpu_flash_attn_float_vec4_aligned(const ggml_tensor * K,
           ggml_webgpu_flash_attn_float_vec4_aligned(V, storage_offset_alignment);
 }

-inline bool ggml_webgpu_flash_attn_kv_direct(const ggml_tensor * Q,
-                                             const ggml_tensor * K,
-                                             const ggml_tensor * V,
-                                             uint32_t            kv_direct_align) {
-    return K->type == GGML_TYPE_F16 && V->type == GGML_TYPE_F16 && (Q->ne[0] % kv_direct_align == 0) &&
-           (K->ne[1] % GGML_WEBGPU_KV_SEQ_PAD == 0);
+inline bool ggml_webgpu_flash_attn_k_direct(const ggml_tensor * Q, const ggml_tensor * K, uint32_t kv_direct_align) {
+    return (K->type == GGML_TYPE_F16 || K->type == GGML_TYPE_Q8_0 || K->type == GGML_TYPE_Q4_0) &&
+           (Q->ne[0] % kv_direct_align == 0) && (K->ne[1] % GGML_WEBGPU_KV_SEQ_PAD == 0);
+}
+
+inline bool ggml_webgpu_flash_attn_v_direct(const ggml_tensor * Q, const ggml_tensor * V, uint32_t kv_direct_align) {
+    return ggml_webgpu_flash_attn_k_direct(Q, V, kv_direct_align);
 }

 inline ggml_webgpu_flash_attn_common_pipeline_key ggml_webgpu_flash_attn_make_common_pipeline_key(
@@ -706,10 +710,11 @@ inline ggml_webgpu_flash_attn_common_pipeline_key ggml_webgpu_flash_attn_make_co
    key.dst_type                                   = context.dst->type;
    key.head_dim_qk                                = (uint32_t) context.src0->ne[0];
    key.head_dim_v                                 = (uint32_t) context.src2->ne[0];
-    key.kv_direct  = ggml_webgpu_flash_attn_kv_direct(context.src0, context.src1, context.src2, kv_direct_align);
-    key.kv_overlap = kv_overlap;
-    key.has_mask   = context.src3 != nullptr;
-    key.has_sinks  = context.src4 != nullptr;
+    key.k_direct           = ggml_webgpu_flash_attn_k_direct(context.src0, context.src1, kv_direct_align);
+    key.v_direct           = ggml_webgpu_flash_attn_v_direct(context.src0, context.src2, kv_direct_align);
+    key.kv_overlap         = kv_overlap;
+    key.has_mask           = context.src3 != nullptr;
+    key.has_sinks          = context.src4 != nullptr;
    key.uses_logit_softcap = ggml_get_op_params_f32(context.dst, 2) != 0.0f;
    return key;
 }
@@ -794,9 +799,13 @@ inline std::vector<std::string> ggml_webgpu_flash_attn_common_defines(
        defines.push_back("LOGIT_SOFTCAP");
        variant += "_lgsc";
    }
-    if (key.kv_direct) {
-        defines.push_back("KV_DIRECT");
-        variant += "_kvdirect";
+    if (key.k_direct) {
+        defines.push_back("K_DIRECT");
+        variant += "_k_direct";
+    }
+    if (key.v_direct) {
+        defines.push_back("V_DIRECT");
+        variant += "_v_direct";
    }
    if (key.kv_overlap) {
        defines.push_back("KV_OVERLAP");
@@ -815,6 +824,12 @@ inline std::vector<std::string> ggml_webgpu_flash_attn_common_defines(

    if (ggml_is_quantized(key.k_type) || ggml_is_quantized(key.v_type)) {
        defines.push_back("U32_DEQUANT_HELPERS");
+        if (ggml_is_quantized(key.k_type)) {
+            defines.push_back("LOADERS_QUANTIZED_K");
+        }
+        if (ggml_is_quantized(key.v_type)) {
+            defines.push_back("LOADERS_QUANTIZED_V");
+        }
    }

    return defines;
@@ -2792,12 +2807,14 @@ class ggml_webgpu_shader_lib {
        ggml_webgpu_flash_attn_pipeline_key key = {};
        key.common                              = ggml_webgpu_flash_attn_make_common_pipeline_key(
            context, decisions.use_sg_matrix ? context.sg_mat_k : 1u, kv_overlap);
-        key.common.kv_direct = decisions.use_sg_matrix && key.common.kv_direct;
-        key.use_sg_matrix    = decisions.use_sg_matrix;
+        key.common.k_direct &= decisions.use_sg_matrix && key.common.k_type == GGML_TYPE_F16;
+        key.common.v_direct &= decisions.use_sg_matrix && key.common.v_type == GGML_TYPE_F16;
+        key.use_sg_matrix = decisions.use_sg_matrix;

        const uint32_t max_kv_tile = ggml_webgpu_flash_attn_max_kv_tile(
            context.wg_mem_limit_bytes, decisions.q_tile, decisions.use_sg_matrix ? context.sg_mat_n : 1u,
-            key.common.head_dim_qk, key.common.head_dim_v, key.common.has_mask, key.common.kv_direct);
+            key.common.head_dim_qk, key.common.head_dim_v, key.common.has_mask,
+            key.common.k_direct || key.common.v_direct);
        GGML_ASSERT(max_kv_tile > 0);

        decisions.kv_tile = decisions.use_sg_matrix ?
@@ -2809,7 +2826,7 @@ class ggml_webgpu_shader_lib {
                std::min(context.max_wg_size, std::max(GGML_WEBGPU_FLASH_ATTN_PREFERRED_WG_SIZE,
                                                       GGML_WEBGPU_FLASH_ATTN_TILE_Q_TILE * context.max_subgroup_size));

-        if (key.common.kv_direct) {
+        if (key.common.k_direct || key.common.v_direct) {
            decisions.kv_tile = std::min(decisions.kv_tile, GGML_WEBGPU_KV_SEQ_PAD);
            while (GGML_WEBGPU_KV_SEQ_PAD % decisions.kv_tile != 0) {
                decisions.kv_tile -= decisions.use_sg_matrix ? context.sg_mat_n : context.min_subgroup_size;
@@ -2856,9 +2873,9 @@ class ggml_webgpu_shader_lib {
        }

        ggml_webgpu_flash_attn_vec_decisions decisions = {};
-        decisions.kv_tile =
-            ggml_webgpu_flash_attn_get_vec_kv_tile(context.wg_mem_limit_bytes, key.common.head_dim_qk,
-                                                   key.common.head_dim_v, key.common.has_mask, key.common.kv_direct);
+        decisions.kv_tile = ggml_webgpu_flash_attn_get_vec_kv_tile(context.wg_mem_limit_bytes, key.common.head_dim_qk,
+                                                                   key.common.head_dim_v, key.common.has_mask,
+                                                                   key.common.k_direct || key.common.v_direct);
        decisions.wg_size = context.max_subgroup_size;

        std::string              variant = "flash_attn_vec";
@@ -2870,12 +2887,10 @@ class ggml_webgpu_shader_lib {
            variant += "_mask_blk";
        }

-        uint32_t d_split = context.min_subgroup_size;
-        if (key.common.k_type == GGML_TYPE_F16 && key.common.v_type == GGML_TYPE_F16) {
-            const uint32_t D     = key.common.head_dim_qk | key.common.head_dim_v;
-            const uint32_t D_lsb = D & (~(D - 1u));
-            d_split              = std::min(std::min(context.min_subgroup_size, 4u), std::max(D_lsb / 4u, 1u));
-        }
+        uint32_t       d_split = context.min_subgroup_size;
+        const uint32_t D       = key.common.head_dim_qk | key.common.head_dim_v;
+        const uint32_t D_lsb   = D & (~(D - 1u));
+        d_split                = std::min(std::min(context.min_subgroup_size, 4u), std::max(D_lsb / 4u, 1u));

        defines.push_back(std::string("D_SPLIT=") + std::to_string(d_split));
        variant += "_dsplit" + std::to_string(d_split);
@@ -3839,7 +3839,8 @@ static size_t ggml_backend_webgpu_buffer_type_get_alloc_size(ggml_backend_buffer
                const auto &        capabilities = ctx->webgpu_global_ctx->capabilities;
                if (ggml_webgpu_flash_attn_use_vec_path(ctx->webgpu_global_ctx, Q, K, V)) {
                    const bool kv_direct =
-                        ggml_webgpu_flash_attn_kv_direct(Q, K, V, GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH);
+                        ggml_webgpu_flash_attn_k_direct(Q, K, GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH) ||
+                        ggml_webgpu_flash_attn_v_direct(Q, V, GGML_WEBGPU_FLASH_ATTN_TILE_KV_VEC_WIDTH);
                    const uint32_t kv_tile = ggml_webgpu_flash_attn_get_vec_kv_tile(
                        capabilities.limits.maxComputeWorkgroupStorageSize, (uint32_t) Q->ne[0], (uint32_t) V->ne[0],
                        mask != nullptr, kv_direct);
@@ -4448,9 +4449,10 @@ static bool ggml_backend_webgpu_device_supports_op(ggml_backend_dev_t dev, const
                const uint32_t q_tile =
                    use_subgroup_matrix ? capabilities.sg_mat_m : GGML_WEBGPU_FLASH_ATTN_TILE_Q_TILE;
                const uint32_t kv_granularity = use_subgroup_matrix ? capabilities.sg_mat_n : 1u;
-                const bool kv_direct = use_subgroup_matrix ?
-                                           ggml_webgpu_flash_attn_kv_direct(src0, src1, src2, capabilities.sg_mat_k) :
-                                           false;
+                const bool     kv_direct = use_subgroup_matrix ?
+                                               ggml_webgpu_flash_attn_k_direct(src0, src1, capabilities.sg_mat_k) ||
+                                                   ggml_webgpu_flash_attn_v_direct(src0, src2, capabilities.sg_mat_k) :
+                                               false;
                const uint32_t max_kv_tile = ggml_webgpu_flash_attn_max_kv_tile(
                    capabilities.limits.maxComputeWorkgroupStorageSize, q_tile, kv_granularity, (uint32_t) src0->ne[0],
                    (uint32_t) src2->ne[0], op->src[3] != nullptr, kv_direct);
@@ -9,6 +9,12 @@ fn get_byte_i32(value: u32, index: u32) -> i32 {
 #endif

 #ifdef U32_DEQUANT_HELPERS
+
+fn f16_from_u16(bits: u32) -> f16 {
+    let packed = unpack2x16float(bits);
+    return f16(packed[0]);
+}
+
 #ifdef DECLARE_BYTE_LOADERS_SRC
 fn load_u16_at_src(byte_offset: u32) -> u32 {
    let word = src[byte_offset / 4u];
@@ -36,7 +42,7 @@ fn load_f16_as_f32_at_src(byte_offset: u32) -> f32 {
    let d_bits = (word >> shift) & 0xFFFFu;
    return unpack2x16float(d_bits)[0];
 }
-#endif
+#endif // DECLARE_BYTE_LOADERS_SRC

 #ifdef DECLARE_BYTE_LOADERS_SRC0
 fn load_u16_at_src0(byte_offset: u32) -> u32 {
@@ -72,8 +78,47 @@ fn load_f16_as_f32_at_src0(byte_offset: u32) -> f32 {
    let d_bits = (word >> shift) & 0xFFFFu;
    return unpack2x16float(d_bits)[0];
 }
-#endif
-#endif
+#endif // DECLARE_BYTE_LOADERS_SRC0
+
+#ifdef LOADERS_QUANTIZED_K
+fn load_k_u16_at(byte_offset: u32) -> u32 {
+    let word = K[byte_offset / 4u];
+    let shift = (byte_offset & 2u) * 8u;
+    return (word >> shift) & 0xFFFFu;
+}
+
+fn load_k_u32_at(byte_offset: u32) -> u32 {
+    let word_idx = byte_offset / 4u;
+    let shift = (byte_offset & 3u) * 8u;
+    let lo = K[word_idx];
+    if (shift == 0u) {
+        return lo;
+    }
+    let hi = K[word_idx + 1u];
+    return (lo >> shift) | (hi << (32u - shift));
+}
+#endif // LOADERS_QUANTIZED_K
+
+#ifdef LOADERS_QUANTIZED_V
+fn load_v_u16_at(byte_offset: u32) -> u32 {
+    let word = V[byte_offset / 4u];
+    let shift = (byte_offset & 2u) * 8u;
+    return (word >> shift) & 0xFFFFu;
+}
+
+fn load_v_u32_at(byte_offset: u32) -> u32 {
+    let word_idx = byte_offset / 4u;
+    let shift = (byte_offset & 3u) * 8u;
+    let lo = V[word_idx];
+    if (shift == 0u) {
+        return lo;
+    }
+    let hi = V[word_idx + 1u];
+    return (lo >> shift) | (hi << (32u - shift));
+}
+#endif // LOADERS_QUANTIZED_V
+
+#endif // U32_DEQUANT_HELPERS



@@ -138,7 +138,7 @@ const FLOAT_MIN: f32 = -1.0e9;
 // The number of Q rows processed per workgroup
 var<workgroup> q_shmem: array<f16, Q_TILE * HEAD_DIM_QK>;

-#ifndef KV_DIRECT
+#if !defined(K_DIRECT) || !defined(V_DIRECT)
 const kv_shmem_size = KV_TILE * max(HEAD_DIM_QK, HEAD_DIM_V);
 // we can reuse the same shmem for K and V since we only need one at a time
 var<workgroup> kv_shmem: array<f16, kv_shmem_size>;
@@ -183,13 +183,12 @@ fn load_kx4(buf: ptr<storage, array<vec4<K_TYPE>>, read_write>, scalar_index: u3
    return (*buf)[scalar_index >> 2u];
 }

-#ifndef KV_DIRECT
+#if !defined(K_DIRECT) || !defined(V_DIRECT)
 #define QUANT_SHMEM kv_shmem
 #define QUANT_OUT_TYPE f16
-#include "quant_inner_loops.tmpl"
 #include "flash_attn_quant_staging.tmpl"

-#if !defined(K_Q4_0) && !defined(K_Q8_0)
+#if !defined(K_DIRECT) && !defined(K_Q4_0) && !defined(K_Q8_0)
 fn load_k_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, k_head_offset: u32) {
    for (var elem_idx = local_x; elem_idx < KV_TILE * HEAD_DIM_QK; elem_idx += WG_SIZE) {
        let k_row = elem_idx / HEAD_DIM_QK;
@@ -204,7 +203,7 @@ fn load_k_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, k_head_offset: u
 }
 #endif

-#if !defined(V_Q4_0) && !defined(V_Q8_0)
+#if !defined(V_DIRECT) && !defined(V_Q4_0) && !defined(V_Q8_0)
 fn load_v_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, v_head_offset: u32) {
    for (var elem_idx = local_x; elem_idx < KV_TILE * HEAD_DIM_V; elem_idx += WG_SIZE) {
        let v_row = elem_idx / HEAD_DIM_V;
@@ -296,7 +295,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
        }

      // load k tile into shared memory
-#ifndef KV_DIRECT
+#ifndef K_DIRECT
      load_k_tile_block(local_id.x, kv_count, kv_tile, k_head_offset);
 #endif

@@ -306,7 +305,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
      // TODO: this loop seems to be the current largest bottleneck
      // this bracket exists to scope the lifetime of variables, reducing register pressure
      {
-#ifdef KV_DIRECT
+#ifdef K_DIRECT
          let k_block_row = kv_tile + subgroup_id * SG_MAT_N;
          var k_global_offset = k_head_offset + k_block_row * params.stride_k1;
 #else
@@ -318,7 +317,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,

              var q_cur = subgroupMatrixLoad<subgroup_matrix_left<f16, SG_MAT_K, SG_MAT_M>>(&q_shmem, 0u, false, HEAD_DIM_QK);

-#ifdef KV_DIRECT
+#ifdef K_DIRECT
              var k_cur = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(&K, k_global_offset + 0u, true, params.stride_k1);
 #else
              var k_cur = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(&kv_shmem, k_block_offset + 0u, true, HEAD_DIM_QK);
@@ -328,7 +327,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
              for (; t + 1u < HEAD_DIM_QK / SG_MAT_K; t += 2u) {
                  let h0 = t * SG_MAT_K;
                  var q0 = subgroupMatrixLoad<subgroup_matrix_left<f16, SG_MAT_K, SG_MAT_M>>(&q_shmem, h0, false, HEAD_DIM_QK);
-#ifdef KV_DIRECT
+#ifdef K_DIRECT
                  var k0 = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(&K, k_global_offset + h0, true, params.stride_k1);
 #else
                  var k0 = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(&kv_shmem, k_block_offset + h0, true, HEAD_DIM_QK);
@@ -339,7 +338,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,

                  let h1 = (t + 1u) * SG_MAT_K;
                  var q1g = subgroupMatrixLoad<subgroup_matrix_left<f16, SG_MAT_K, SG_MAT_M>>(&q_shmem, h1, false, HEAD_DIM_QK);
-#ifdef KV_DIRECT
+#ifdef K_DIRECT
                  var k1g = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(&K, k_global_offset + h1, true, params.stride_k1);
 #else
                  var k1g = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(&kv_shmem, k_block_offset + h1, true, HEAD_DIM_QK);
@@ -353,7 +352,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
              if (t < HEAD_DIM_QK / SG_MAT_K) {
                  let h = t * SG_MAT_K;
                  var qn = subgroupMatrixLoad<subgroup_matrix_left<f16, SG_MAT_K, SG_MAT_M>>(&q_shmem, h, false, HEAD_DIM_QK);
-#ifdef KV_DIRECT
+#ifdef K_DIRECT
                  var kn = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(&K, k_global_offset + h, true, params.stride_k1);
 #else
                  var kn = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(&kv_shmem, k_block_offset + h, true, HEAD_DIM_QK);
@@ -365,7 +364,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,

              acc = subgroupMatrixMultiplyAccumulate(q_cur, k_cur, acc);

-#ifdef KV_DIRECT
+#ifdef K_DIRECT
              k_global_offset += num_subgroups * SG_MAT_N * params.stride_k1;
 #else
              k_block_offset += num_subgroups * SG_MAT_N * HEAD_DIM_QK;
@@ -436,7 +435,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
      }

      // load v tile into shared memory
-#ifndef KV_DIRECT
+#ifndef V_DIRECT
      load_v_tile_block(local_id.x, kv_count, kv_tile, v_head_offset);
 #endif

@@ -464,7 +463,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
                  );

                  // load V submatrix from global or shared memory
-#ifdef KV_DIRECT
+#ifdef V_DIRECT
                  let v_block_row = kv_tile + kv_block * SG_MAT_N;
                  let v_global_offset = v_head_offset + v_block_row * params.stride_v1 + head_dim_block;
                  var v_sg_mat: subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K> = subgroupMatrixLoad<subgroup_matrix_right<f16, SG_MAT_N, SG_MAT_K>>(
@@ -1,3 +1,5 @@
+#include "quant_inner_loops.tmpl"
+
 #define BLOCK_SIZE 32
 #define BLOCKS_K ((HEAD_DIM_QK + BLOCK_SIZE - 1) / BLOCK_SIZE)
 #define BLOCKS_V ((HEAD_DIM_V + BLOCK_SIZE - 1) / BLOCK_SIZE)
@@ -26,49 +28,6 @@
 #define V_BYTES_PER_INNER_LOOP 4u
 #endif

-#if defined(K_Q4_0) || defined(K_Q8_0)
-fn load_k_u16_at(byte_offset: u32) -> u32 {
-    let word = K[byte_offset / 4u];
-    let shift = (byte_offset & 2u) * 8u;
-    return (word >> shift) & 0xFFFFu;
-}
-
-fn load_k_u32_at(byte_offset: u32) -> u32 {
-    let word_idx = byte_offset / 4u;
-    let shift = (byte_offset & 3u) * 8u;
-    let lo = K[word_idx];
-    if (shift == 0u) {
-        return lo;
-    }
-    let hi = K[word_idx + 1u];
-    return (lo >> shift) | (hi << (32u - shift));
-}
-#endif
-
-#if defined(V_Q4_0) || defined(V_Q8_0)
-fn load_v_u16_at(byte_offset: u32) -> u32 {
-    let word = V[byte_offset / 4u];
-    let shift = (byte_offset & 2u) * 8u;
-    return (word >> shift) & 0xFFFFu;
-}
-
-fn load_v_u32_at(byte_offset: u32) -> u32 {
-    let word_idx = byte_offset / 4u;
-    let shift = (byte_offset & 3u) * 8u;
-    let lo = V[word_idx];
-    if (shift == 0u) {
-        return lo;
-    }
-    let hi = V[word_idx + 1u];
-    return (lo >> shift) | (hi << (32u - shift));
-}
-#endif
-
-fn f16_from_u16(bits: u32) -> f16 {
-    let packed = unpack2x16float(bits);
-    return f16(packed[0]);
-}
-
 #if defined(K_Q4_0) || defined(K_Q8_0)
 fn load_k_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, k_head_offset: u32) {
    for (var elem_idx = local_x * K_NQ; elem_idx < kv_count * HEAD_DIM_QK; elem_idx += WG_SIZE * K_NQ) {
@@ -153,7 +153,6 @@ var<workgroup> p_shmem: array<f16, Q_TILE * KV_TILE>;

 #define QUANT_SHMEM kv_shmem
 #define QUANT_OUT_TYPE f16
-#include "quant_inner_loops.tmpl"
 #include "flash_attn_quant_staging.tmpl"

 #if !defined(K_Q4_0) && !defined(K_Q8_0)
@@ -270,7 +269,9 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
            local_scores[slot] = FLOAT_MIN;
        }

-#ifndef KV_DIRECT
+        // The tile path stages K/V in shared memory so each tile can be reused across
+        // Q_TILE query rows. It therefore does not use the direct path.
+#ifndef K_DIRECT
        load_k_tile_block(local_id.x, kv_count, kv_tile, k_head_offset);
 #endif

@@ -333,7 +334,9 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,

        workgroupBarrier();

-#ifndef KV_DIRECT
+        // The tile path stages K/V in shared memory so each tile can be reused across
+        // Q_TILE query rows. It therefore does not use the direct path.
+#ifndef V_DIRECT
        load_v_tile_block(local_id.x, kv_count, kv_tile, v_head_offset);
 #endif

@@ -196,49 +196,35 @@ struct Params {

 // Just a very small float value.
 const FLOAT_MIN: f32 = -1.0e9;
+const kv_shmem_size = KV_TILE * max(HEAD_DIM_QK, HEAD_DIM_V);

 var<workgroup> q_shmem: array<f32, HEAD_DIM_QK>;
-
-#ifndef KV_DIRECT
-const kv_shmem_size = KV_TILE * max(HEAD_DIM_QK, HEAD_DIM_V);
-// we can reuse the same shmem for K and V since we only need one at a time
-var<workgroup> kv_shmem: array<f32, kv_shmem_size>;
-#endif
-
 var<workgroup> o_shmem: array<f32, HEAD_DIM_V>;
+// note that we reuse the same storage for both since we only need one at a time
+var<workgroup> inter_shmem: array<f32, KV_TILE>;

 #ifdef MASK
 // storage for mask values
 var<workgroup> mask_shmem: array<f32, KV_TILE>;
 #endif

-// note that we reuse the same storage for both since we only need one at a time
-var<workgroup> inter_shmem: array<f32, KV_TILE>;
-
-// Storage for row max and exp sum during online softmax
-fn calc_softmax_term(kv_idx: u32, slope: f32, has_bias: bool, apply_mask: bool) -> f32 {
-    var v = select(FLOAT_MIN,
-                   inter_shmem[kv_idx] * params.scale,
-                   kv_idx < KV_TILE);
-#ifdef LOGIT_SOFTCAP
-    v = params.logit_softcap * tanh(v);
+#if defined(K_DIRECT) || defined(V_DIRECT)
+// Shared memory for scale factor (d) in quantized K/V. Multiple threads use the same value,
+// so caching it is more efficient, even on the direct path.
+var<workgroup> d_shmem: array<f32, kv_shmem_size / 32>;
 #endif
-#ifdef MASK
-    if (apply_mask) {
-        var mask_val = select(0.0, mask_shmem[kv_idx], kv_idx < KV_TILE);
-        v += select(mask_val, slope * mask_val, has_bias);
-    }
-#endif
-    return v;
-}

-#ifndef KV_DIRECT
+// K/V shared memory handling
+#if !defined(K_DIRECT) || !defined(V_DIRECT)
+
+// we can reuse the same shmem for K and V since we only need one at a time
+var<workgroup> kv_shmem: array<f32, kv_shmem_size>;
+
 #define QUANT_SHMEM kv_shmem
 #define QUANT_OUT_TYPE f32
-#include "quant_inner_loops.tmpl"
 #include "flash_attn_quant_staging.tmpl"

-#if !defined(K_Q4_0) && !defined(K_Q8_0)
+#if !defined(K_DIRECT) && !defined(K_Q4_0) && !defined(K_Q8_0)
 fn load_k_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, k_head_offset: u32) {
    for (var elem_idx = local_x * 4u; elem_idx < KV_TILE * HEAD_DIM_QK; elem_idx += WG_SIZE * 4u) {
        let k_row = elem_idx / HEAD_DIM_QK;
@@ -256,7 +242,7 @@ fn load_k_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, k_head_offset: u
 }
 #endif

-#if !defined(V_Q4_0) && !defined(V_Q8_0)
+#if !defined(V_DIRECT) && !defined(V_Q4_0) && !defined(V_Q8_0)
 fn load_v_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, v_head_offset: u32) {
    for (var elem_idx = local_x * 4u; elem_idx < KV_TILE * HEAD_DIM_V; elem_idx += WG_SIZE * 4u) {
        let v_row = elem_idx / HEAD_DIM_V;
@@ -273,7 +259,24 @@ fn load_v_tile_block(local_x: u32, kv_count: u32, kv_tile: u32, v_head_offset: u
    }
 }
 #endif
+#endif // !defined(K_DIRECT) || !defined(V_DIRECT)
+
+// Storage for row max and exp sum during online softmax
+fn calc_softmax_term(kv_idx: u32, slope: f32, has_bias: bool, apply_mask: bool) -> f32 {
+    var v = select(FLOAT_MIN,
+                   inter_shmem[kv_idx] * params.scale,
+                   kv_idx < KV_TILE);
+#ifdef LOGIT_SOFTCAP
+    v = params.logit_softcap * tanh(v);
 #endif
+#ifdef MASK
+    if (apply_mask) {
+        var mask_val = select(0.0, mask_shmem[kv_idx], kv_idx < KV_TILE);
+        v += select(mask_val, slope * mask_val, has_bias);
+    }
+#endif
+    return v;
+}

@compute @workgroup_size(WG_SIZE)
 fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
@@ -355,12 +358,31 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
            inter_shmem[elem_idx] = 0.0;
        }

-      // load k tile into shared memory
-#ifndef KV_DIRECT
-      load_k_tile_block(local_id.x, kv_count, kv_tile, k_head_offset);
+#ifdef K_DIRECT
+      // load only the scale factor (d) from each quantized block into shared memory on the direct path.
+#if defined(K_Q8_0)
+        for (var j = local_id.x * 32; j < KV_TILE * HEAD_DIM_QK; j += WG_SIZE * 32) {
+            let kv_row = kv_tile + j / HEAD_DIM_QK;
+            let block_idx = (j % HEAD_DIM_QK) / 32;
+            let block_byte_base = 34 * (k_head_offset + kv_row * params.stride_k1 + block_idx);
+            let d = f32(f16_from_u16(load_k_u16_at(block_byte_base)));
+            d_shmem[j / 32] = d;
+        }
+#elif defined(K_Q4_0)
+        for (var j = local_id.x * 32; j < KV_TILE * HEAD_DIM_QK; j += WG_SIZE * 32) {
+            let kv_row = kv_tile + j / HEAD_DIM_QK;
+            let block_idx = (j % HEAD_DIM_QK) / 32;
+            let block_byte_base = 18 * (k_head_offset + kv_row * params.stride_k1 + block_idx);
+            let d = f32(f16_from_u16(load_k_u16_at(block_byte_base)));
+            d_shmem[j / 32] = d;
+        }
 #endif
+#else
+      // load k tile into shared memory
+      load_k_tile_block(local_id.x, kv_count, kv_tile, k_head_offset);
+#endif // defined(K_DIRECT)

-      workgroupBarrier();
+        workgroupBarrier();

      // accumulate q block * k block into registers across the entire KV tile
      if (!skip_tile) {
@@ -381,9 +403,40 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
                            q_shmem[q_off + 1u],
                            q_shmem[q_off + 2u],
                            q_shmem[q_off + 3u]);
-#ifdef KV_DIRECT
+#ifdef K_DIRECT
+#if defined(K_Q8_0)
+                        let kv_row = kv_tile + kv_idx;
+                        let block_idx = (i * 4u) / 32;
+                        let id_in_block = (i * 4u) % 32;
+                        let block_byte_base = 34 * (k_head_offset + kv_row * params.stride_k1 + block_idx);
+                        let q_byte_base = block_byte_base + 2u;
+                        let d = d_shmem[(kv_idx * HEAD_DIM_QK) / 32 + block_idx];
+                        let q8u4 = load_k_u32_at(q_byte_base + id_in_block);
+                        let kv = vec4<f32>(
+                            d * f32(get_byte_i32(q8u4, 0)),
+                            d * f32(get_byte_i32(q8u4, 1)),
+                            d * f32(get_byte_i32(q8u4, 2)),
+                            d * f32(get_byte_i32(q8u4, 3)),
+                        );
+#elif defined(K_Q4_0)
+                        let kv_row = kv_tile + kv_idx;
+                        let block_idx = (i * 4u) / 32;
+                        let id_in_block = (i * 4u) % 32;
+                        let phase = id_in_block / 16;
+                        let block_byte_base = 18 * (k_head_offset + kv_row * params.stride_k1 + block_idx);
+                        let q_byte_base = block_byte_base + 2u;
+                        let d = d_shmem[(kv_idx * HEAD_DIM_QK) / 32 + block_idx];
+                        let q8u4 = load_k_u32_at(q_byte_base + (id_in_block - phase * 16u));
+                        let kv = vec4<f32>(
+                            d * (f32((get_byte(q8u4, 0) >> (phase * 4u)) & 0xFu) - 8.0),
+                            d * (f32((get_byte(q8u4, 1) >> (phase * 4u)) & 0xFu) - 8.0),
+                            d * (f32((get_byte(q8u4, 2) >> (phase * 4u)) & 0xFu) - 8.0),
+                            d * (f32((get_byte(q8u4, 3) >> (phase * 4u)) & 0xFu) - 8.0),
+                        );
+#else
                        let idx = k_head_offset + (kv_tile + kv_idx) * params.stride_k1 + (i * 4u);
                        let kv = vec4<f32>(K[idx >> 2u]);
+#endif
 #else
                        let idx = kv_idx * HEAD_DIM_QK + (i * 4u);
                        let kv = vec4<f32>(
@@ -391,7 +444,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
                            kv_shmem[idx + 1u],
                            kv_shmem[idx + 2u],
                            kv_shmem[idx + 3u]);
-#endif
+#endif // defined(K_DIRECT)
                        partial_sum += dot(qv, kv);
                    }
                  }
@@ -473,12 +526,32 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
          }
      }

-      // load v tile into shared memory
-#ifndef KV_DIRECT
-      load_v_tile_block(local_id.x, kv_count, kv_tile, v_head_offset);
-#endif

-      workgroupBarrier();
+#ifdef V_DIRECT
+      // load only `d` of quantized block into shared memory in the direct path
+#if defined(V_Q8_0)
+        for (var j = local_id.x * 32; j < KV_TILE * HEAD_DIM_V; j += WG_SIZE * 32) {
+            let v_row = kv_tile + j / HEAD_DIM_V;
+            let block_idx = (j % HEAD_DIM_V) / 32;
+            let block_byte_base = 34 * (v_head_offset + v_row * params.stride_v1 + block_idx);
+            let d = f32(f16_from_u16(load_v_u16_at(block_byte_base)));
+            d_shmem[j / 32] = d;
+        }
+#elif defined(V_Q4_0)
+        for (var j = local_id.x * 32; j < KV_TILE * HEAD_DIM_V; j += WG_SIZE * 32) {
+            let v_row = kv_tile + j / HEAD_DIM_V;
+            let block_idx = (j % HEAD_DIM_V) / 32;
+            let block_byte_base = 18 * (v_head_offset + v_row * params.stride_v1 + block_idx);
+            let d = f32(f16_from_u16(load_v_u16_at(block_byte_base)));
+            d_shmem[j / 32] = d;
+        }
+#endif
+#else
+      // load v tile into shared memory
+      load_v_tile_block(local_id.x, kv_count, kv_tile, v_head_offset);
+#endif // V_DIRECT
+
+        workgroupBarrier();

      if (!skip_tile) {
          // we have P (KV_TILE) in inter_shmem and V (KV_TILE x head_dim_v) in kv_shmem
@@ -501,9 +574,38 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
                      }

                      let p = inter_shmem[kv_idx];
-#ifdef KV_DIRECT
+#ifdef V_DIRECT
+#if defined(V_Q8_0)
+                        let block_idx = (vec_col * 4u) / 32;
+                        let id_in_block = (vec_col * 4u) % 32;
+                        let block_byte_base = 34 * (v_head_offset + v_row * params.stride_v1 + block_idx);
+                        let q_byte_base = block_byte_base + 2u;
+                        let d = d_shmem[(kv_idx * HEAD_DIM_V) / 32 + block_idx];
+                        let q8u4 = load_v_u32_at(q_byte_base + id_in_block);
+                        let v4 = vec4<f32>(
+                            d * f32(get_byte_i32(q8u4, 0)),
+                            d * f32(get_byte_i32(q8u4, 1)),
+                            d * f32(get_byte_i32(q8u4, 2)),
+                            d * f32(get_byte_i32(q8u4, 3)),
+                        );
+#elif defined(V_Q4_0)
+                        let block_idx = (vec_col * 4u) / 32;
+                        let id_in_block = (vec_col * 4u) % 32;
+                        let phase = id_in_block / 16;
+                        let block_byte_base = 18 * (v_head_offset + v_row * params.stride_v1 + block_idx);
+                        let q_byte_base = block_byte_base + 2u;
+                        let d = d_shmem[(kv_idx * HEAD_DIM_V) / 32 + block_idx];
+                        let q8u4 = load_v_u32_at(q_byte_base + (id_in_block - phase * 16u));
+                        let v4 = vec4<f32>(
+                            d * (f32((get_byte(q8u4, 0) >> (phase * 4u)) & 0xFu) - 8.0),
+                            d * (f32((get_byte(q8u4, 1) >> (phase * 4u)) & 0xFu) - 8.0),
+                            d * (f32((get_byte(q8u4, 2) >> (phase * 4u)) & 0xFu) - 8.0),
+                            d * (f32((get_byte(q8u4, 3) >> (phase * 4u)) & 0xFu) - 8.0),
+                        );
+#else
                      let v_idx = v_head_offset + v_row * params.stride_v1 + vec_col * 4u;
                      let v4 = vec4<f32>(V[v_idx >> 2u]);
+#endif
 #else
                      let v_idx = kv_idx * HEAD_DIM_V + vec_col * 4u;
                      let v4 = vec4<f32>(
@@ -511,7 +613,7 @@ fn main(@builtin(workgroup_id) wg_id: vec3<u32>,
                          kv_shmem[v_idx + 1u],
                          kv_shmem[v_idx + 2u],
                          kv_shmem[v_idx + 3u]);
-#endif
+#endif // defined(V_DIRECT)
                      lo += p * v4;
                  }

@@ -8329,7 +8329,11 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
            test_cases.emplace_back(new test_conv_2d(
                { act_case[iwh_idx], act_case[iwh_idx], act_case[Cin_idx], act_case[B_idx] },
                { act_case[kwh_idx], act_case[kwh_idx], act_case[Cin_idx], act_case[Cout_idx] },
-                kernel_type, 1, 1, 0, 0, 1, 1, false));
+                kernel_type, 1, 1, 0, 0, 1, 1, false));  // bool cwhn = false
+            test_cases.emplace_back(new test_conv_2d(
+                { act_case[iwh_idx], act_case[iwh_idx], act_case[Cin_idx], act_case[B_idx] },
+                { act_case[kwh_idx], act_case[kwh_idx], act_case[Cin_idx], act_case[Cout_idx] },
+                kernel_type, 1, 1, 0, 0, 1, 1, true));  // bool cwhn = true
        }
    }
 #endif
@@ -8358,7 +8362,9 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
                                        calc_conv_output_size(H, KH, s1, p1, d1) > 0) {
                                        for (auto kernel_type : {GGML_TYPE_F32, GGML_TYPE_F16}) {
                                            test_cases.emplace_back(new test_conv_2d(
-                                                { W, H, Cin, 2 }, { KW, KH, Cin, Cout }, kernel_type, s0, s1, p0, p1, d0, d1, false));
+                                                { W, H, Cin, 2 }, { KW, KH, Cin, Cout }, kernel_type, s0, s1, p0, p1, d0, d1, false)); // bool cwhn = false
+                                            test_cases.emplace_back(new test_conv_2d(
+                                                { W, H, Cin, 2 }, { KW, KH, Cin, Cout }, kernel_type, s0, s1, p0, p1, d0, d1, true));  // bool cwhn = true
                                        }
                                    }
                                }
@@ -8370,7 +8376,8 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_eval() {
        }
    }
    for (auto kernel_type : {GGML_TYPE_F32, GGML_TYPE_F16}) {
-        test_cases.emplace_back(new test_conv_2d({ 256, 256, 192, 1 }, { 3, 3, 192, 96 }, kernel_type, 1, 1, 1, 1, 1, 1, false));
+        test_cases.emplace_back(new test_conv_2d({ 256, 256, 192, 1 }, { 3, 3, 192, 96 }, kernel_type, 1, 1, 1, 1, 1, 1, false)); // bool cwhn = false
+        test_cases.emplace_back(new test_conv_2d({ 256, 256, 192, 1 }, { 3, 3, 192, 96 }, kernel_type, 1, 1, 1, 1, 1, 1, true));  // bool cwhn = true
    }

    // sycl backend will limit task global_range < MAX_INT
@@ -9757,7 +9764,11 @@ static std::vector<std::unique_ptr<test_case>> make_test_cases_perf() {
            test_cases.emplace_back(new test_conv_2d(
                { act_case[iwh_idx], act_case[iwh_idx], act_case[Cin_idx], act_case[B_idx] },
                { act_case[kwh_idx], act_case[kwh_idx], act_case[Cin_idx], act_case[Cout_idx] },
-                kernel_type, 1, 1, 0, 0, 1, 1, false));
+                kernel_type, 1, 1, 0, 0, 1, 1, false));   // bool cwhn = false
+            test_cases.emplace_back(new test_conv_2d(
+                { act_case[iwh_idx], act_case[iwh_idx], act_case[Cin_idx], act_case[B_idx] },
+                { act_case[kwh_idx], act_case[kwh_idx], act_case[Cin_idx], act_case[Cout_idx] },
+                kernel_type, 1, 1, 0, 0, 1, 1, true));    // bool cwhn = true
        }
    }

@@ -33,7 +33,7 @@ enum resize_algo {
    RESIZE_ALGO_BILINEAR, // stretch to target resolution
    RESIZE_ALGO_BICUBIC, // center-crop when aspect ratio doesn't match
    RESIZE_ALGO_BICUBIC_PILLOW,
-    // RESIZE_ALGO_LANCZOS, // TODO
+    RESIZE_ALGO_LANCZOS,
 };

 // Padding style for img_tool::resize
@@ -68,6 +68,9 @@ struct img_tool {
                case RESIZE_ALGO_BICUBIC_PILLOW:
                    resize_bicubic_pillow(src, dst, target_resolution.width, target_resolution.height);
                    break;
+                case RESIZE_ALGO_LANCZOS:
+                    resize_lanczos_pillow(src, dst, target_resolution.width, target_resolution.height);
+                    break;
                default:
                    throw std::runtime_error("Unsupported resize algorithm");
            }
@@ -97,6 +100,9 @@ struct img_tool {
                case RESIZE_ALGO_BICUBIC_PILLOW:
                    resize_bicubic_pillow(src, resized_image, new_width, new_height);
                    break;
+                case RESIZE_ALGO_LANCZOS:
+                    resize_lanczos_pillow(src, resized_image, new_width, new_height);
+                    break;
                default:
                    throw std::runtime_error("Unsupported resize algorithm");
            }
@@ -337,22 +343,50 @@ private:
        }
    }

-    // Bicubic resize function using Pillow's ImagingResample algorithm
+    // Pillow-compatible separable resampling (Bicubic and Lanczos)
    // Adapted from https://github.com/python-pillow/Pillow/blob/main/src/libImaging/Resample.c
    //
-    // Key Difference with resize_bicubic:
-    // 1. Uses separable filtering: horizontal pass followed by vertical pass
+    // Key properties:
+    // 1. Separable filtering: horizontal pass followed by vertical pass
    // 2. Pre-computes normalized filter coefficients for each output pixel
-    // 3. Applies convolution using fixed-point integer arithmetic for performance
+    // 3. Fixed-point integer arithmetic (22 fractional bits) for speed and determinism
    static bool resize_bicubic_pillow(const clip_image_u8 & img, clip_image_u8 & dst, int target_width, int target_height) {
+        return resize_pillow(img, dst, target_width, target_height, /*use_lanczos=*/false);
+    }
+
+    // Lanczos-3 (support radius 3), matches Pillow's Image.LANCZOS
+    static bool resize_lanczos_pillow(const clip_image_u8 & img, clip_image_u8 & dst, int target_width, int target_height) {
+        return resize_pillow(img, dst, target_width, target_height, /*use_lanczos=*/true);
+    }
+
+    static bool resize_pillow(
+            const clip_image_u8 & img,
+            clip_image_u8 & dst,
+            int target_width,
+            int target_height,
+            bool use_lanczos) {
        // Fixed-point precision: 22 bits = 32 (int32_t) - 8 (uint8_t pixels) - 2 (headroom for accumulation)
        // This allows encoding fractional weights as integers: weight * 2^22
        const int PRECISION_BITS = 32 - 8 - 2;

-        // Bicubic filter function with a = -0.5 (Note that GGML/PyTorch takes a = -0.75)
+        // Resample filter: Lanczos-3 (support [-3, 3]) or bicubic with a = -0.5 (support [-2, 2])
+        // Note: GGML/PyTorch bicubic uses a = -0.75, Pillow uses a = -0.5
        // Returns filter weight for distance x from pixel center
-        // Support: [-2, 2], meaning the filter influences pixels within 2 units of distance
-        auto bicubic_filter = [](double x) -> double {
+        auto resample_filter = [use_lanczos](double x) -> double {
+            if (use_lanczos) {
+                if (-3.0 <= x && x < 3.0) {
+                    auto sinc = [](double v) {
+                        if (v == 0.0) {
+                            return 1.0;
+                        }
+                        const double pi_v = v * 3.141592653589793238462643383279502884;
+                        return std::sin(pi_v) / pi_v;
+                    };
+                    return sinc(x) * sinc(x / 3.0);
+                }
+                return 0.0;
+            }
+
            constexpr double a = -0.5;
            if (x < 0.0) {
                x = -x;
@@ -366,8 +400,8 @@ private:
            return 0.0;  // Zero outside [-2, 2]
        };

-        // Filter support radius: bicubic extends 2 pixels in each direction
-        constexpr double filter_support = 2.0;
+        // Filter support radius: 2 for bicubic, 3 for lanczos
+        const double filter_support = use_lanczos ? 3.0 : 2.0;

        // Clipping function for 8-bit values
        auto clip8 = [](int val) -> uint8_t {
@@ -434,7 +468,7 @@ private:
                // Compute filter weights for each contributing input pixel
                for (x = 0; x < xmax; x++) {
                    // Distance from input pixel center to output pixel center in input space
-                    double w = bicubic_filter((x + xmin - center + 0.5) * ss);
+                    double w = resample_filter((x + xmin - center + 0.5) * ss);
                    pre_weights[xx * ksize + x] = w;
                    ww += w;  // Accumulate for normalization
                }
@@ -463,6 +497,12 @@ private:
            const double fxp_scale = std::ldexp(1.0, PRECISION_BITS); // 1.0 * 2^PRECISION_BITS

            for (int i = 0; i < outSize * ksize; i++) {
+                if (use_lanczos) {
+                    // Pillow adds +/- 0.5 then truncates toward zero; std::round would round twice
+                    const double rounded = pre_weights[i] * fxp_scale + (pre_weights[i] < 0 ? -0.5 : 0.5);
+                    weights[i] = static_cast<int32_t>(rounded);
+                    continue;
+                }
                double tmp_val = pre_weights[i] * fxp_scale;
                if (pre_weights[i] < 0) {
                    tmp_val -= 0.5;
@@ -78,31 +78,41 @@ struct server_batch {
    };
    std::vector<token> tokens;
    int32_t n_tokens_alloc = 0;
+    int32_t n_embd = 0;

    // track if given slot can be batched with slots already in the batch
    server_slot * slot_batched = nullptr;

+    // in embd mode, we temporarily swap out the tokens arr and restore it on clear()
+    bool has_embd = false;
+    llama_token * tokens_ptr = nullptr;
+    std::vector<float> embd;
+
    float  alora_scale       = -1.0f;
    size_t alora_disabled_id = 0;

    server_batch() {
-        batch.token = nullptr; // sentinel: uninitialized batch
+        batch.pos = nullptr; // sentinel: uninitialized batch
    }

    ~server_batch() {
-        if (batch.token != nullptr) {
+        if (batch.pos != nullptr) {
+            clear();
            llama_batch_free(batch);
        }
    }

-    void init(int32_t n_tokens_alloc) {
+    void init(int32_t n_tokens_alloc, int32_t n_embd) {
        this->n_tokens_alloc = n_tokens_alloc;
+        this->n_embd = n_embd;
        batch = llama_batch_init(n_tokens_alloc, 0, 1);
+        tokens_ptr = batch.token;
        tokens.reserve(n_tokens_alloc);
    }

    bool add(int32_t id_slot, llama_token token, llama_pos pos, bool output) {
-        GGML_ASSERT(batch.token != nullptr);
+        GGML_ASSERT(!has_embd); // cannot mix tokens + embd in same batch
+        GGML_ASSERT(batch.pos != nullptr);
        if ((int32_t)tokens.size() >= n_tokens_alloc) {
            return false;
        }
@@ -110,13 +120,30 @@ struct server_batch {
        return true;
    }

+    bool add(int32_t id_slot, const std::vector<float> & embd_in, llama_pos pos, bool output) {
+        GGML_ASSERT(batch.pos != nullptr);
+        if ((int32_t)tokens.size() >= n_tokens_alloc) {
+            return false;
+        }
+        tokens.push_back({ id_slot, LLAMA_TOKEN_NULL, pos, output });
+        has_embd = true;
+        embd.insert(embd.end(), embd_in.begin(), embd_in.end());
+        return true;
+    }
+
    void clear() {
        tokens.clear();
+        embd.clear();
        common_batch_clear(batch);
        slot_batched      = nullptr;
        alora_scale       = -1.0f;
        alora_disabled_id = 0;
        batch_rendered    = false;
+        has_embd          = false;
+        if (batch.token == nullptr) {
+            batch.token = tokens_ptr;
+            batch.embd  = nullptr;
+        }
    }

    int32_t size() const {
@@ -129,25 +156,33 @@ struct server_batch {
    }

    void render() {
-        GGML_ASSERT(batch.token != nullptr);
+        GGML_ASSERT(!batch_rendered);
+        GGML_ASSERT(batch.pos != nullptr);
        common_batch_clear(batch);
        for (int32_t i = 0; i < size(); i++) {
            const auto & t = tokens[i];
            common_batch_add(batch, t.token, t.pos, { t.id_slot }, t.output);
        }
+        if (has_embd) {
+            batch.token = nullptr; // will be restored on clear()
+            batch.embd  = embd.data();
+        }
        batch_rendered = true;
    }

    llama_batch get_view(int32_t off, int32_t n_tokens) const {
-        GGML_ASSERT(batch.token != nullptr);
+        GGML_ASSERT(batch.pos != nullptr);
        GGML_ASSERT(batch_rendered);
        GGML_ASSERT(off >= 0 && off < size());
        GGML_ASSERT(n_tokens > 0 && off + n_tokens <= size());

+        auto * token = batch.token ? batch.token + off          : nullptr;
+        auto * embd  = batch.embd  ? batch.embd  + off * n_embd : nullptr;
+
        llama_batch view = {
            n_tokens,
-            batch.token    + off,
-            nullptr,
+            token,
+            embd,
            batch.pos      + off,
            batch.n_seq_id + off,
            batch.seq_id   + off,
@@ -270,6 +305,10 @@ struct server_slot {

    llama_token sampled; // in speculative mode, this is the last accepted token

+    // for TTS models, this is the embd generated from prev step, decode this to generate next hidden state
+    // corresponding to one token position (size = n_embd)
+    std::vector<float> inp_embd;
+
    // stats
    size_t n_sent_text = 0; // number of sent text character

@@ -378,7 +417,9 @@ struct server_slot {
    bool can_batch_with(server_slot & other_slot) const {
        GGML_ASSERT(task);

-        return task->type == other_slot.task->type && are_lora_equal(lora, other_slot.lora);
+        return task->type == other_slot.task->type
+            && inp_embd.size() == other_slot.inp_embd.size()
+            && are_lora_equal(lora, other_slot.lora);
    }

    bool has_budget(const common_params & global_params) {
@@ -444,7 +485,11 @@ struct server_slot {
            // no speculative decoding
            i_batch = batch.size();

-            add_ok &= batch.add(id, sampled, prompt.tokens.pos_next(), true);
+            if (!inp_embd.empty()) {
+                add_ok &= batch.add(id, inp_embd, prompt.tokens.pos_next(), true);
+            } else {
+                add_ok &= batch.add(id, sampled, prompt.tokens.pos_next(), true);
+            }

            SLT_DBG(*this, "slot decode token, id=%d, n_ctx = %d, n_tokens = %d, truncated = %d\n",
                    sampled, n_ctx, prompt.n_tokens(), truncated);
@@ -1334,7 +1379,8 @@ private:
        // note that n_batch can be > n_ctx (e.g. for non-causal attention models such as BERT where the KV cache is not used)
        {
            const int32_t n_batch = llama_n_batch(ctx_tgt);
-            batch.init(std::max(n_batch, params_base.n_parallel));
+            const int32_t n_embd  = llama_model_n_embd_inp(model_tgt);
+            batch.init(std::max(n_batch, params_base.n_parallel), n_embd);
        }

        if (params_base.cache_ram_mib != 0) {
@@ -3578,6 +3624,15 @@ private:
            n_empty_consecutive = 0;
        }

+        // TODO @ngxson : dft model may have different n_embd than the tgt model, so we check & reject if that's the case
+        // this case is not currently used by any models, but may need to be supported in the future
+        if (spec && batch.has_embd) {
+            if (llama_model_n_embd_inp(model_dft) != llama_model_n_embd_inp(model_tgt)) {
+                SRV_ERR("%s", "unsupported batch.has_embd + spec case\n");
+                throw std::runtime_error("unsupported batch.has_embd + spec case");
+            }
+        }
+
        const int ret = llama_decode(ctx_tgt, batch_view);

        metrics.on_decoded(slots);
Author	SHA1	Message	Date
Masashi YoshimuraandGitHub	8f4646a63e	ggml-webgpu: improve flash_attn_vec for quantized KV at long contexts (#25956 ) * improve fa of quantized kv cache * Fix some bugs and some comments. * fix v type check and some comments * Fix build error caused by rebasing * editorconfig checking pass	2026-07-31 09:08:40 +03:00
Xuan-Son NguyenandGitHub	5f55650a78	mtmd: add lanczos resize method [no release] (#26341 )	2026-07-30 21:59:49 +02:00
Xuan-Son NguyenandGitHub	b4ca032ae3	server: support inp embd to generate next token (#26313 ) * server: support embd for sampled token * fix ~server_batch()	2026-07-30 21:40:38 +02:00
Jeff BolzandGitHub	ea63b4d32e	vulkan: Support quantized concat (#25684 )	2026-07-30 13:11:32 -05:00
pmaybankandGitHub	958d9c0b61	Test support for alternative conv layout (#25617 ) * add bool cwhn = true to conv_2d test cases * add layout check at graph building time * extend layout checks for conv2d.cu kernel * in CPU back-end kernel needs to be stored contiguously to prevent test failures with cwhn=1 * trim white space * do op support check in vulkan backend * fix CI failure and vulkan run-time assert failure by introducing new graph build-time check in ggml_backend_vk_device_supports_op * add additional check in support_op function for Vulkan to fix run-time assert failure	2026-07-31 01:14:16 +08:00