From 1a946ec7459f08bb6d8a05ca449586dd091d4fb3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Sigbj=C3=B8rn=20Skj=C3=A6ret?= Date: Thu, 27 Aug 2026 15:27:17 +0200 Subject: [PATCH 1/3] pr2wt : use ssh/https remote in worktree depending on base (#27800) --- scripts/pr2wt.sh | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/scripts/pr2wt.sh b/scripts/pr2wt.sh index ae03a888cb..ce0327eadb 100755 --- a/scripts/pr2wt.sh +++ b/scripts/pr2wt.sh @@ -48,7 +48,11 @@ echo "org/repo: $org_repo" meta=$(curl -sSLf -H "Accept: application/vnd.github+json" "https://api.github.com/repos/$org_repo/pulls/$PR") -url_remote=$(echo "$meta" | jq -r '.head.repo.clone_url') +if [[ $url_origin =~ ^git@ ]]; then + url_remote=$(echo "$meta" | jq -r '.head.repo.ssh_url') +else + url_remote=$(echo "$meta" | jq -r '.head.repo.clone_url') +fi head_ref=$(echo "$meta" | jq -r '.head.ref') echo "url: $url_remote" From cb300598d5f90189cb69d2702f4930aaf99d32a2 Mon Sep 17 00:00:00 2001 From: Shobhit Date: Thu, 27 Aug 2026 21:34:42 +0800 Subject: [PATCH 2/3] Feature: Added LIGHTNING_INDEXER support for Deepseek V4 ops on Vulkan Backend (#27453) * vulkan: add LIGHTNING_INDEXER op * vulkan: updated lightning_indexer.comp and ggml-vulkan.cpp with 128-lane dot-product reduction moved from a shared-memory tree to subgroupAdd. * vulkan: cleanup; Skip bounds checks * vulkan: cleanup FA_K_ONLY * Revert "vulkan: cleanup FA_K_ONLY" This reverts commit fdcbdd91511945d6878d9070b5447e4a34dce010. * vulkan: restore interleaved K/V buffer ordering * vulkan: Remove FA_K_ONLY * vulkan: Revert flash_attn_dequant * vulkan: Revert tests in backend-ops.cpp --- ggml/src/ggml-vulkan/ggml-vulkan.cpp | 156 +++++++++++++++++- .../ggml-vulkan/vulkan-shaders/fa_types.glsl | 55 ++++++ .../vulkan-shaders/flash_attn_base.glsl | 51 +----- .../vulkan-shaders/lightning_indexer.comp | 151 +++++++++++++++++ .../vulkan-shaders/vulkan-shaders-gen.cpp | 6 + 5 files changed, 365 insertions(+), 54 deletions(-) create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/fa_types.glsl create mode 100644 ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer.comp diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 8108e94c16..72e844aebf 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -767,6 +767,21 @@ static constexpr std::initializer_list> rms_norm_mul_rope_vie { 4, 0, 3 }, // set_rows->src[0] == view }; +static constexpr std::array lightning_indexer_k_types = { + GGML_TYPE_F32, + GGML_TYPE_F16, + GGML_TYPE_BF16, + GGML_TYPE_Q8_0, + GGML_TYPE_Q5_1, + GGML_TYPE_Q5_0, + GGML_TYPE_Q4_1, + GGML_TYPE_Q4_0, + GGML_TYPE_IQ4_NL, +}; + +static bool ggml_vk_lightning_indexer_k_type_supported(ggml_type type) { + return std::find(lightning_indexer_k_types.begin(), lightning_indexer_k_types.end(), type) != lightning_indexer_k_types.end(); +} struct vk_device_struct { std::recursive_mutex mutex; @@ -1068,6 +1083,7 @@ struct vk_device_struct { vk_pipeline pipeline_rwkv_wkv6_f32; vk_pipeline pipeline_rwkv_wkv7_f32; vk_pipeline pipeline_gated_linear_attn_f32; + vk_pipeline pipeline_lightning_indexer_f32[GGML_TYPE_COUNT]; // [size_idx][kda] where size_idx: 0=d16, 1=d32, 2=d64, 3=d128 vk_pipeline pipeline_gated_delta_net[4][2]; vk_pipeline pipeline_ssm_scan_f32_d128; @@ -1848,6 +1864,26 @@ struct vk_op_gated_linear_attn_push_constants { uint32_t H; float scale; }; +struct vk_op_lightning_indexer_push_constants { + uint32_t n_kv; + uint32_t n_heads; + uint32_t n_tokens; + uint32_t n_streams; + uint32_t n_masks; + uint32_t dispatch_x; + uint32_t q_nb1; + uint32_t q_nb2; + uint32_t q_nb3; + uint32_t k_nb2; + uint32_t k_nb3; + uint32_t w_nb1; + uint32_t w_nb3; + uint32_t m_nb1; + uint32_t m_nb3; + uint32_t d_nb1; + uint32_t d_nb3; +}; +static_assert(sizeof(vk_op_lightning_indexer_push_constants) <= 128); struct vk_op_gated_delta_net_push_constants { uint32_t H; uint32_t n_tokens; @@ -3904,11 +3940,16 @@ static vk_fa_pipeline_state get_fa_pipeline_state(const vk_device& device, const return vk_fa_pipeline_state{hsk, hsv, params.block_rows, params.block_cols, params.d_split, params.row_split, params.shmem_staging, params.path, params.workgroup_size, subgroup_size, aligned, f32acc, flags, params.limit_occupancy_shmem, k_type, v_type}; } +// Bytes per buffer block for the FaBlockBytesK/V spec constants. F32 is fed as +// a vec4 "block" of 4 floats, everything else uses its ggml block size. +static uint32_t fa_block_bytes(ggml_type t) { + if (t == GGML_TYPE_F32) { + return 16u; + } + return (uint32_t) ggml_type_size(t); +} + static std::vector get_fa_spec_constants(const vk_fa_pipeline_state& state) { - const auto fa_block_bytes = [](ggml_type t) -> uint32_t { - if (t == GGML_TYPE_F32) return 16u; - return (uint32_t) ggml_type_size(t); - }; return { /* 0 WorkGroupSize */ state.workgroup_size, /* 1 Br */ state.Br, @@ -5847,6 +5888,17 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_gated_linear_attn_f32, "gated_linear_attn_f32", gated_linear_attn_f32_len, gated_linear_attn_f32_data, "main", 6, sizeof(vk_op_gated_linear_attn_push_constants), {1, 1, 1}, {}, 1); + { + const bool li_subgroup = device->subgroup_arithmetic && device->subgroup_require_full_support; + const size_t li_len = li_subgroup ? lightning_indexer_subgroup_f32_len : lightning_indexer_f32_len; + const void * li_data = li_subgroup ? (const void *)lightning_indexer_subgroup_f32_data : (const void *)lightning_indexer_f32_data; + + for (ggml_type k_type : lightning_indexer_k_types) { + const std::string name = "lightning_indexer_" + std::string(ggml_type_name(k_type)) + "_k_f32"; + ggml_vk_create_pipeline(device, device->pipeline_lightning_indexer_f32[k_type], name.c_str(), li_len, li_data, "main", 5, sizeof(vk_op_lightning_indexer_push_constants), {1, 1, 1}, {(uint32_t)k_type, fa_block_bytes(k_type), device->subgroup_size}, 1, true, li_subgroup); + } + } + { const uint32_t gdn_sizes[] = {16, 32, 64, 128}; const char * gdn_names[][2] = { @@ -11697,6 +11749,12 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const return ctx->device->pipeline_gated_linear_attn_f32; } return nullptr; + case GGML_OP_LIGHTNING_INDEXER: + // only the k type selects a pipeline, the other types are fixed by ggml_lightning_indexer() + if (ggml_vk_lightning_indexer_k_type_supported(src1->type)) { + return ctx->device->pipeline_lightning_indexer_f32[src1->type]; + } + return nullptr; case GGML_OP_GATED_DELTA_NET: if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { const uint32_t S_v = dst->src[2]->ne[0]; @@ -12772,6 +12830,55 @@ static void ggml_vk_gated_linear_attn(ggml_backend_vk_context * ctx, vk_context& pc, { (uint32_t)(n_seqs * n_heads), 1, 1 }); } +static void ggml_vk_lightning_indexer(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) { + const ggml_tensor * q = dst->src[0]; + const ggml_tensor * k = dst->src[1]; + const ggml_tensor * w = dst->src[2]; + const ggml_tensor * m = dst->src[3]; + + vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, q, k, w, dst, dst->op); + GGML_ASSERT(pipeline != nullptr); + + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + + const uint32_t n_kv = k->ne[2]; + const uint32_t n_heads = q->ne[1]; + const uint32_t n_tokens = q->ne[2]; + const uint32_t n_streams = q->ne[3]; + const uint32_t n_masks = m->ne[3]; + + const uint32_t n_outputs = (uint32_t)(dst->ne[0] * dst->ne[1] * dst->ne[3]); + const uint32_t dispatch_x = std::min(n_outputs, ctx->device->properties.limits.maxComputeWorkGroupCount[0]); + const uint32_t dispatch_y = CEIL_DIV(n_outputs, dispatch_x); + + // q, w and dst are f32 and m is f16, so their strides are passed in elements; + // k may be quantized, so its strides stay in bytes + const uint32_t q_nb1 = q->nb[1] / sizeof(float); + const uint32_t q_nb2 = q->nb[2] / sizeof(float); + const uint32_t q_nb3 = q->nb[3] / sizeof(float); + const uint32_t k_nb2 = k->nb[2]; + const uint32_t k_nb3 = k->nb[3]; + const uint32_t w_nb1 = w->nb[1] / sizeof(float); + const uint32_t w_nb3 = w->nb[3] / sizeof(float); + const uint32_t m_nb1 = m->nb[1] / sizeof(ggml_fp16_t); + const uint32_t m_nb3 = m->nb[3] / sizeof(ggml_fp16_t); + const uint32_t d_nb1 = dst->nb[1] / sizeof(float); + const uint32_t d_nb3 = dst->nb[3] / sizeof(float); + + const vk_op_lightning_indexer_push_constants pc = { + n_kv, n_heads, n_tokens, n_streams, n_masks, dispatch_x, + q_nb1, q_nb2, q_nb3, + k_nb2, k_nb3, + w_nb1, w_nb3, + m_nb1, m_nb3, + d_nb1, d_nb3, + }; + + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + {ggml_vk_tensor_subbuffer(ctx, q), ggml_vk_tensor_subbuffer(ctx, k), ggml_vk_tensor_subbuffer(ctx, w), ggml_vk_tensor_subbuffer(ctx, m), ggml_vk_tensor_subbuffer(ctx, dst)}, + pc, {dispatch_x, dispatch_y, 1}); +} + static void ggml_vk_gated_delta_net(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) { const ggml_tensor * src_q = dst->src[0]; const ggml_tensor * src_v = dst->src[2]; @@ -15898,6 +16005,11 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr break; + case GGML_OP_LIGHTNING_INDEXER: + ggml_vk_lightning_indexer(ctx, compute_ctx, node); + + break; + case GGML_OP_GATED_DELTA_NET: ggml_vk_gated_delta_net(ctx, compute_ctx, node); @@ -18676,6 +18788,40 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm case GGML_OP_GATED_LINEAR_ATTN: // the shader block size is hardcoded to head_size 64 return op->src[0]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32 && op->src[0]->ne[0] == 64; + case GGML_OP_LIGHTNING_INDEXER: + { + const ggml_tensor * q = op->src[0]; + const ggml_tensor * k = op->src[1]; + const ggml_tensor * w = op->src[2]; + const ggml_tensor * m = op->src[3]; + + // the q/w/m types and the shape relationships between q, k, w, m and dst + // are already asserted in ggml_lightning_indexer() + if (!ggml_vk_lightning_indexer_k_type_supported(k->type) || !device->fp16) { + return false; + } + + // the shader block size is hardcoded to head size 128 + if (q->ne[0] != 128) { + return false; + } + + // the shader indexes the buffers by element stride, and is dispatched + // without allow_misalign + for (const ggml_tensor * t : {q, k, w, m, op}) { + if (t->nb[0] != ggml_type_size(t->type) || + (vk_tensor_offset(t) + t->view_offs) % device->properties.limits.minStorageBufferOffsetAlignment != 0) { + return false; + } + // the strides get scaled down from bytes, so the division must be exact + for (int i = 1; i < GGML_MAX_DIMS; ++i) { + if (t->nb[i] % ggml_type_size(t->type) != 0) { + return false; + } + } + } + return true; + } case GGML_OP_GATED_DELTA_NET: { const uint32_t S_v = op->src[2]->ne[0]; @@ -19685,6 +19831,8 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph * const float * op_params = (const float *)tensor->op_params; tensor_clone = ggml_gated_linear_attn(ggml_ctx, src_clone[0], src_clone[1], src_clone[2], src_clone[3], src_clone[4], op_params[0]); + } else if (tensor->op == GGML_OP_LIGHTNING_INDEXER) { + tensor_clone = ggml_lightning_indexer(ggml_ctx, src_clone[0], src_clone[1], src_clone[2], src_clone[3]); } else if (tensor->op == GGML_OP_GATED_DELTA_NET) { tensor_clone = ggml_gated_delta_net(ggml_ctx, src_clone[0], src_clone[1], src_clone[2], src_clone[3], src_clone[4], src_clone[5], diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/fa_types.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/fa_types.glsl new file mode 100644 index 0000000000..6f414ded12 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/fa_types.glsl @@ -0,0 +1,55 @@ +#if !defined(GGML_FA_TYPES_COMP) +#define GGML_FA_TYPES_COMP + +// FaTypeK / FaTypeV spec constant values. These mirror enum ggml_type so the +// host can pass the type directly. Keep in sync with ggml.h. +#define FA_TYPE_F32 0u +#define FA_TYPE_F16 1u +#define FA_TYPE_Q4_0 2u +#define FA_TYPE_Q4_1 3u +#define FA_TYPE_Q5_0 6u +#define FA_TYPE_Q5_1 7u +#define FA_TYPE_Q8_0 8u +#define FA_TYPE_IQ4_NL 20u +#define FA_TYPE_BF16 30u + +// Number of matrix elements per buffer block, derived from the K/V type spec +// constant. F32 is treated as a vec4 "block" of 4 floats. F16 uses block size 1 +// and bypasses the dequant path entirely. Quants follow their ggml block sizes. +uint fa_block_elems(uint ty) { + switch (ty) { + case FA_TYPE_F32: return 4u; + case FA_TYPE_F16: return 1u; + case FA_TYPE_Q4_0: return uint(QUANT_K_Q4_0); + case FA_TYPE_Q4_1: return uint(QUANT_K_Q4_1); + case FA_TYPE_Q5_0: return uint(QUANT_K_Q5_0); + case FA_TYPE_Q5_1: return uint(QUANT_K_Q5_1); + case FA_TYPE_Q8_0: return uint(QUANT_K_Q8_0); + case FA_TYPE_IQ4_NL: return uint(QUANT_K_IQ4_NL); + case FA_TYPE_BF16: return 1u; + default: return 1u; + } +} + +// QUANT_R_MMQ for FA-eligible K types. Q4_*/Q5_* store two nibbles per byte +// (R==2); Q8_0 stores one byte per element (R==1). Used to derive the number +// of int32s per 32-element block on the MMQ K path: ints_per_block == 8 / R. +uint fa_quant_r_mmq(uint ty) { + switch (ty) { + case FA_TYPE_Q4_0: return uint(QUANT_R_Q4_0); + case FA_TYPE_Q4_1: return uint(QUANT_R_Q4_1); + case FA_TYPE_Q5_0: return uint(QUANT_R_Q5_0); + case FA_TYPE_Q5_1: return uint(QUANT_R_Q5_1); + case FA_TYPE_Q8_0: return uint(QUANT_R_Q8_0); + default: return 1u; + } +} + +bool fa_type_needs_shmem(uint ty) { + switch (ty) { + case FA_TYPE_IQ4_NL: return true; + default: return false; + } +} + +#endif // !defined(GGML_FA_TYPES_COMP) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl index 3c64f91dad..0ce4503a88 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl @@ -88,17 +88,7 @@ layout (binding = 6) readonly buffer MO {uint32_t data_mask_opt[];}; #define BINDING_IDX_K 0 #define BINDING_IDX_V 1 -// FaTypeK / FaTypeV spec constant values. These mirror enum ggml_type so the -// host can pass the type directly. Keep in sync with ggml.h. -#define FA_TYPE_F32 0u -#define FA_TYPE_F16 1u -#define FA_TYPE_Q4_0 2u -#define FA_TYPE_Q4_1 3u -#define FA_TYPE_Q5_0 6u -#define FA_TYPE_Q5_1 7u -#define FA_TYPE_Q8_0 8u -#define FA_TYPE_IQ4_NL 20u -#define FA_TYPE_BF16 30u +#include "fa_types.glsl" #if defined(BFLOAT16) #define O_TYPE float @@ -108,45 +98,6 @@ layout (binding = 6) readonly buffer MO {uint32_t data_mask_opt[];}; #define O_TYPEV4 FLOAT_TYPEV4 #endif -// Number of matrix elements per buffer block, derived from the K/V type spec -// constant. F32 is treated as a vec4 "block" of 4 floats. F16 uses block size 1 -// and bypasses the dequant path entirely. Quants follow their ggml block sizes. -uint fa_block_elems(uint ty) { - switch (ty) { - case FA_TYPE_F32: return 4u; - case FA_TYPE_F16: return 1u; - case FA_TYPE_Q4_0: return uint(QUANT_K_Q4_0); - case FA_TYPE_Q4_1: return uint(QUANT_K_Q4_1); - case FA_TYPE_Q5_0: return uint(QUANT_K_Q5_0); - case FA_TYPE_Q5_1: return uint(QUANT_K_Q5_1); - case FA_TYPE_Q8_0: return uint(QUANT_K_Q8_0); - case FA_TYPE_IQ4_NL: return uint(QUANT_K_IQ4_NL); - case FA_TYPE_BF16: return 1u; - default: return 1u; - } -} - -// QUANT_R_MMQ for FA-eligible K types. Q4_*/Q5_* store two nibbles per byte -// (R==2); Q8_0 stores one byte per element (R==1). Used to derive the number -// of int32s per 32-element block on the MMQ K path: ints_per_block == 8 / R. -uint fa_quant_r_mmq(uint ty) { - switch (ty) { - case FA_TYPE_Q4_0: return uint(QUANT_R_Q4_0); - case FA_TYPE_Q4_1: return uint(QUANT_R_Q4_1); - case FA_TYPE_Q5_0: return uint(QUANT_R_Q5_0); - case FA_TYPE_Q5_1: return uint(QUANT_R_Q5_1); - case FA_TYPE_Q8_0: return uint(QUANT_R_Q8_0); - default: return 1u; - } -} - -bool fa_type_needs_shmem(uint ty) { - switch (ty) { - case FA_TYPE_IQ4_NL: return true; - default: return false; - } -} - // These can't be `const` globals because GLSL forbids function calls in global // const initializers, even when the spec constants would let the driver fold // them. Macros expand at the use site and fold after specialization. diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer.comp b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer.comp new file mode 100644 index 0000000000..ba76ec72ca --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer.comp @@ -0,0 +1,151 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_KHR_shader_subgroup_basic : enable +#if USE_SUBGROUP_ADD +#extension GL_KHR_shader_subgroup_arithmetic : enable +#endif + +#define BINDING_IDX_K 0u + +#include "types.glsl" +#include "fa_types.glsl" +#define FaTypeV FA_TYPE_F32 + +layout(constant_id = 0) const uint FaTypeK = FA_TYPE_F32; +layout(constant_id = 1) const uint FaBlockBytesK = 4; +layout(constant_id = 2) const uint SUBGROUP_SIZE = 32; + +#include "flash_attn_dequant.glsl" + +// one workgroup computes one output element, one invocation per head element +#define HEAD_SIZE 128 + +layout(local_size_x = HEAD_SIZE, local_size_y = 1, local_size_z = 1) in; + +layout(binding = 0) readonly buffer QBuf { float q[]; }; +layout(binding = 1) readonly buffer KBufF16 { float16_t k_f16[]; }; +layout(binding = 1) readonly buffer KBufF32 { float k_f32[]; }; +layout(binding = 1) readonly buffer KBufBF16 { uint16_t k_bf16[]; }; +layout(binding = 2) readonly buffer WBuf { float weights[]; }; +layout(binding = 3) readonly buffer MBuf { float16_t mask[]; }; +layout(binding = 4) writeonly buffer DstBuf { float dst[]; }; + +layout(push_constant) uniform PushConstants { + uint n_kv; + uint n_heads; + uint n_tokens; + uint n_streams; + uint n_masks; + uint dispatch_x; + uint q_nb1; + uint q_nb2; + uint q_nb3; + uint k_nb2; + uint k_nb3; + uint w_nb1; + uint w_nb3; + uint m_nb1; + uint m_nb3; + uint d_nb1; + uint d_nb3; +}; + +shared float k_row[HEAD_SIZE]; + +#if USE_SUBGROUP_ADD +shared float sg_partials[HEAD_SIZE / SUBGROUP_SIZE]; +#else +shared float partials[HEAD_SIZE]; +#endif + +void main() { + const uint tid = gl_LocalInvocationID.x; + const uint output_idx = gl_WorkGroupID.y * dispatch_x + gl_WorkGroupID.x; + const uint n_outputs = n_kv * n_tokens * n_streams; + + if (fa_type_needs_shmem(FaTypeK)) { + init_iq_shmem(gl_WorkGroupSize); + } + + if (output_idx >= n_outputs) { + return; + } + + const uint ik = output_idx % n_kv; + const uint ts = output_idx / n_kv; + const uint t = ts % n_tokens; + const uint s = ts / n_tokens; + const uint k_offset = ik * k_nb2 + s * k_nb3; + + // k strides come in as bytes, so scale them down to the view being indexed + const uint k_block_elems = fa_block_elems(FaTypeK); + const uint k_elem_bytes = FaBlockBytesK / k_block_elems; + + if (FaTypeK == FA_TYPE_F16) { + k_row[tid] = float(k_f16[k_offset / k_elem_bytes + tid]); + } else if (FaTypeK == FA_TYPE_F32) { + k_row[tid] = k_f32[k_offset / k_elem_bytes + tid]; + } else if (FaTypeK == FA_TYPE_BF16) { + k_row[tid] = bf16_to_fp32(uint(k_bf16[k_offset / k_elem_bytes + tid])); + } else if (4 * tid < HEAD_SIZE) { + const uint coord = 4 * tid; + const uint ib = coord / k_block_elems; + const uint iqs = coord % k_block_elems; + const vec4 values = dequantize4(ib, iqs, k_offset / FaBlockBytesK, BINDING_IDX_K); + k_row[coord + 0] = values.x; + k_row[coord + 1] = values.y; + k_row[coord + 2] = values.z; + k_row[coord + 3] = values.w; + } + barrier(); + + const float k_val = k_row[tid]; + + float score = 0.0; + for (uint h = 0; h < n_heads; ++h) { + const float prod = q[h * q_nb1 + t * q_nb2 + s * q_nb3 + tid] * k_val; + +#if USE_SUBGROUP_ADD + const float sg_sum = subgroupAdd(prod); + if (gl_SubgroupInvocationID == 0) { + sg_partials[gl_SubgroupID] = sg_sum; + } + barrier(); + + if (tid == 0) { + float sum = 0.0; + [[unroll]] for (uint i = 0; i < HEAD_SIZE / SUBGROUP_SIZE; ++i) { + sum += sg_partials[i]; + } + score += max(sum, 0.0) * weights[h + t * w_nb1 + s * w_nb3]; + } + // the reads above must complete before the next iteration overwrites sg_partials + barrier(); +#else + partials[tid] = prod; + barrier(); + + [[unroll]] for (uint stride = HEAD_SIZE / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + partials[tid] += partials[tid + stride]; + } + barrier(); + } + + if (tid == 0) { + score += max(partials[0], 0.0) * weights[h + t * w_nb1 + s * w_nb3]; + } + // the read of partials[0] above must complete before the next iteration + // overwrites partials[tid] + barrier(); +#endif + } + + if (tid == 0) { + const uint mask_offset = ik + t * m_nb1 + (s % n_masks) * m_nb3; + dst[ik + t * d_nb1 + s * d_nb3] = score + float(mask[mask_offset]); + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp index dbb99782cf..0da943da95 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -1069,6 +1069,12 @@ void process_shaders() { string_to_spv("gated_linear_attn_f32", "gla.comp", merge_maps(base_dict, {{"A_TYPE", "float"}})); + // Compile IQ4_NL support in so its shared LUT is available when K uses it. + // K quant type is selected at runtime via the FaTypeK spec constant. + std::map li_dict = {{"FLOAT_TYPE", "float"}, {"FLOAT_TYPEV4", "vec4"}, {"DATA_A_IQ4_NL", "1"}}; + string_to_spv("lightning_indexer_f32", "lightning_indexer.comp", li_dict); + string_to_spv("lightning_indexer_subgroup_f32", "lightning_indexer.comp", merge_maps(li_dict, {{"USE_SUBGROUP_ADD", "1"}})); + string_to_spv("rwkv_wkv7_f32", "wkv7.comp", merge_maps(base_dict, {{"A_TYPE", "float"}})); string_to_spv("gated_delta_net_f32", "gated_delta_net.comp", merge_maps(base_dict, {{"FLOAT_TYPE", "float"}, {"USE_SUBGROUP_ADD", "1"}, {"USE_SUBGROUP_CLUSTERED", "1"}})); From 732707dff265be513347b02911e1dd18f7e1f386 Mon Sep 17 00:00:00 2001 From: Xuan-Son Nguyen Date: Thu, 27 Aug 2026 18:31:13 +0200 Subject: [PATCH 3/3] quantize: cap working memory size to avoid loading big tensors onto RAM (#27795) --- include/llama.h | 1 + src/llama-model-loader.cpp | 25 +++---- src/llama-model-loader.h | 5 +- src/llama-quant.cpp | 143 ++++++++++++++++++++---------------- tools/quantize/quantize.cpp | 17 ++++- 5 files changed, 112 insertions(+), 79 deletions(-) diff --git a/include/llama.h b/include/llama.h index b8020a7fce..49a758db26 100644 --- a/include/llama.h +++ b/include/llama.h @@ -445,6 +445,7 @@ extern "C" { const struct llama_model_kv_override * kv_overrides; // pointer to kv overrides const struct llama_model_tensor_override * tt_overrides; // pointer to tensor overrides const int32_t * prune_layers; // pointer to layer indices to prune + size_t max_buf_size; // max bytes of tensor rows kept in memory at once, 0 = default (8 GiB) } llama_model_quantize_params; typedef struct llama_logit_bias { diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index 2f1a09a19a..36df0e3590 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -1418,27 +1418,26 @@ void llama_model_loader::unmap_weight(const llama_tensor_weight & w) const { mappings.at(w.idx)->unmap_fragment(w.offs, w.offs + ggml_nbytes(w.tensor)); } -void llama_model_loader::load_data_for(struct ggml_tensor * cur) const { - const auto & w = require_weight(ggml_get_name(cur)); +const void * llama_model_loader::load_data_range(const llama_tensor_weight & w, size_t offs, size_t size, void * buf) const { + GGML_ASSERT(offs + size <= ggml_nbytes(w.tensor)); + + const void * data = buf; if (use_mmap) { - const auto & mapping = mappings.at(w.idx); - if (cur->data == nullptr) { - cur->data = (uint8_t *)mapping->addr() + w.offs; - } else { - memcpy(cur->data, (uint8_t *)mapping->addr() + w.offs, ggml_nbytes(cur)); - } + data = (const uint8_t *) mappings.at(w.idx)->addr() + w.offs + offs; } else { - GGML_ASSERT(cur->data != nullptr); + GGML_ASSERT(buf != nullptr); GGML_ASSERT(w.idx < files.size()); const auto & file = files.at(w.idx); - file->seek(w.offs, SEEK_SET); - file->read_raw(cur->data, ggml_nbytes(cur)); + file->seek(w.offs + offs, SEEK_SET); + file->read_raw(buf, size); } - if (check_tensors && !ggml_validate_row_data(cur->type, cur->data, ggml_nbytes(cur))) { - throw std::runtime_error(format("tensor '%s' has invalid data", ggml_get_name(cur))); + if (check_tensors && !ggml_validate_row_data(w.tensor->type, data, size)) { + throw std::runtime_error(format("tensor '%s' has invalid data", ggml_get_name(w.tensor))); } + + return data; } bool llama_model_loader::load_all_data( diff --git a/src/llama-model-loader.h b/src/llama-model-loader.h index ec5b692463..407260e990 100644 --- a/src/llama-model-loader.h +++ b/src/llama-model-loader.h @@ -204,8 +204,9 @@ struct llama_model_loader { // release a weight's mmap pages void unmap_weight(const llama_tensor_weight & w) const; - // for backwards compatibility, does not support ggml-backend - void load_data_for(struct ggml_tensor * cur) const; + // read a byte range of a weight's data + // with mmap, returns a pointer into the mapping, otherwise reads into buf and returns buf + const void * load_data_range(const llama_tensor_weight & w, size_t offs, size_t size, void * buf) const; // Returns false if cancelled by progress_callback bool load_all_data( diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp index 20252815d5..49e1dd0c75 100644 --- a/src/llama-quant.cpp +++ b/src/llama-quant.cpp @@ -38,6 +38,9 @@ enum class tensor_category { OTHER }; +// max amount of tensor data kept in memory while quantizing a single tensor +static const size_t LLAMA_QUANT_MAX_BUF_SIZE = 8ull*1024*1024*1024; + static void zeros(std::ofstream & file, size_t n) { char zero = 0; for (size_t i = 0; i < n; ++i) { @@ -211,31 +214,26 @@ struct tensor_metadata { // static void llama_tensor_dequantize_impl( - ggml_tensor * tensor, std::vector> & output, std::vector & workers, + ggml_type type, const void * data, float * f32_output, std::vector & workers, const size_t nelements, const int nthread ) { - if (output.size() < nelements) { - output.resize(nelements); - } - float * f32_output = (float *) output.data(); - - const ggml_type_traits * qtype = ggml_get_type_traits(tensor->type); - if (ggml_is_quantized(tensor->type)) { + const ggml_type_traits * qtype = ggml_get_type_traits(type); + if (ggml_is_quantized(type)) { if (qtype->to_float == NULL) { - throw std::runtime_error(format("type %s unsupported for integer quantization: no dequantization available", ggml_type_name(tensor->type))); + throw std::runtime_error(format("type %s unsupported for integer quantization: no dequantization available", ggml_type_name(type))); } - } else if (tensor->type != GGML_TYPE_F16 && - tensor->type != GGML_TYPE_BF16) { - throw std::runtime_error(format("cannot dequantize/convert tensor type %s", ggml_type_name(tensor->type))); + } else if (type != GGML_TYPE_F16 && + type != GGML_TYPE_BF16) { + throw std::runtime_error(format("cannot dequantize/convert tensor type %s", ggml_type_name(type))); } if (nthread < 2) { - if (tensor->type == GGML_TYPE_F16) { - ggml_fp16_to_fp32_row((ggml_fp16_t *)tensor->data, f32_output, nelements); - } else if (tensor->type == GGML_TYPE_BF16) { - ggml_bf16_to_fp32_row((ggml_bf16_t *)tensor->data, f32_output, nelements); - } else if (ggml_is_quantized(tensor->type)) { - qtype->to_float(tensor->data, f32_output, nelements); + if (type == GGML_TYPE_F16) { + ggml_fp16_to_fp32_row((const ggml_fp16_t *)data, f32_output, nelements); + } else if (type == GGML_TYPE_BF16) { + ggml_bf16_to_fp32_row((const ggml_bf16_t *)data, f32_output, nelements); + } else if (ggml_is_quantized(type)) { + qtype->to_float(data, f32_output, nelements); } else { GGML_ABORT("fatal error"); // unreachable } @@ -243,14 +241,14 @@ static void llama_tensor_dequantize_impl( } size_t block_size; - if (tensor->type == GGML_TYPE_F16 || - tensor->type == GGML_TYPE_BF16) { + if (type == GGML_TYPE_F16 || + type == GGML_TYPE_BF16) { block_size = 1; } else { - block_size = (size_t)ggml_blck_size(tensor->type); + block_size = (size_t)ggml_blck_size(type); } - size_t block_size_bytes = ggml_type_size(tensor->type); + size_t block_size_bytes = ggml_type_size(type); GGML_ASSERT(nelements % block_size == 0); size_t nblocks = nelements / block_size; @@ -265,16 +263,16 @@ static void llama_tensor_dequantize_impl( size_t thr_elems = thr_blocks * block_size; // number of elements for this thread size_t thr_block_bytes = thr_blocks * block_size_bytes; // number of input bytes for this thread - auto compute = [qtype] (ggml_type typ, uint8_t * inbuf, float * outbuf, int nels) { + auto compute = [qtype] (ggml_type typ, const uint8_t * inbuf, float * outbuf, int nels) { if (typ == GGML_TYPE_F16) { - ggml_fp16_to_fp32_row((ggml_fp16_t *)inbuf, outbuf, nels); + ggml_fp16_to_fp32_row((const ggml_fp16_t *)inbuf, outbuf, nels); } else if (typ == GGML_TYPE_BF16) { - ggml_bf16_to_fp32_row((ggml_bf16_t *)inbuf, outbuf, nels); + ggml_bf16_to_fp32_row((const ggml_bf16_t *)inbuf, outbuf, nels); } else { qtype->to_float(inbuf, outbuf, nels); } }; - workers.emplace_back(compute, tensor->type, (uint8_t *) tensor->data + in_buff_offs, f32_output + out_buff_offs, thr_elems); + workers.emplace_back(compute, type, (const uint8_t *) data + in_buff_offs, f32_output + out_buff_offs, thr_elems); in_buff_offs += thr_block_bytes; out_buff_offs += thr_elems; } @@ -1093,6 +1091,8 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: std::vector> work; std::vector> f32_conv_buf; + const size_t max_buf_size = params->max_buf_size ? params->max_buf_size : LLAMA_QUANT_MAX_BUF_SIZE; + int cur_split = -1; std::ofstream fout; auto close_ofstream = [&]() { @@ -1143,15 +1143,13 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: const size_t tensor_size = ggml_nbytes(tensor); - if (!params->dry_run) { - if (!ml.use_mmap) { - if (read_data.size() < tensor_size) { - read_data.resize(tensor_size); - } - tensor->data = read_data.data(); + // read a byte range of the current tensor + auto load_range = [&](size_t offs, size_t size) -> const void * { + if (!ml.use_mmap && read_data.size() < size) { + read_data.resize(size); } - ml.load_data_for(tensor); - } + return ml.load_data_range(weight, offs, size, read_data.data()); + }; LLAMA_LOG_INFO("[%4d/%4d] %-36s - [%s], type = %6s, ", ++idx, ml.n_tensors, @@ -1166,7 +1164,6 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: // in then there's nothing to do. bool quantize = cur_type != new_type; - void * new_data; size_t new_size; if (params->dry_run) { @@ -1190,12 +1187,18 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: } else { // no --dry-run, perform quantization if (!quantize) { - new_data = tensor->data; new_size = tensor_size; LLAMA_LOG_INFO("size = %8.3f MiB\n", tensor_size/1024.0/1024.0); - } else { - const int64_t nelements = ggml_nelements(tensor); + // copy in slabs of whole rows, so that each slab can be validated + const size_t row_size = ggml_row_size(tensor->type, tensor->ne[0]); + const size_t slab_size = std::max(row_size, (max_buf_size/row_size)*row_size); + + for (size_t offs = 0; offs < tensor_size; offs += slab_size) { + const size_t size = std::min(slab_size, tensor_size - offs); + fout.write((const char *) load_range(offs, size), size); + } + } else { const float * imatrix = nullptr; if (imatrix_data) { auto it = imatrix_data->find(tm.remapped_imatrix_name); @@ -1227,43 +1230,60 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: throw std::runtime_error(format("Missing importance matrix for tensor %s in a very low-bit quantization", tensor->name)); } - float * f32_data; - - if (tensor->type == GGML_TYPE_F32) { - f32_data = (float *) tensor->data; - } else if (ggml_is_quantized(tensor->type) && !params->allow_requantize) { + if (ggml_is_quantized(tensor->type) && !params->allow_requantize) { throw std::runtime_error(format("requantizing from type %s is disabled", ggml_type_name(tensor->type))); - } else { - llama_tensor_dequantize_impl(tensor, f32_conv_buf, workers, nelements, nthread); - f32_data = (float *) f32_conv_buf.data(); } LLAMA_LOG_INFO("converting to %s .. ", ggml_type_name(new_type)); fflush(stdout); - if (work.size() < (size_t)nelements * 4) { - work.resize(nelements * 4); // upper bound on size - } - new_data = work.data(); - const int64_t n_per_row = tensor->ne[0]; const int64_t nrows = tensor->ne[1]; + const size_t row_size_src = ggml_row_size(tensor->type, n_per_row); + const size_t row_size_dst = ggml_row_size(new_type, n_per_row); + + // process the rows in slabs, so that the buffers stay below max_buf_size + const size_t bytes_per_row = row_size_src + row_size_dst + (tensor->type == GGML_TYPE_F32 ? 0 : n_per_row*sizeof(float)); + const int64_t nrows_slab = std::max(1, std::min(nrows, max_buf_size/bytes_per_row)); + static const int64_t min_chunk_size = 32 * 512; const int64_t chunk_size = (n_per_row >= min_chunk_size ? n_per_row : n_per_row * ((min_chunk_size + n_per_row - 1)/n_per_row)); - const int64_t nelements_matrix = tensor->ne[0] * tensor->ne[1]; - const int64_t nchunk = (nelements_matrix + chunk_size - 1)/chunk_size; - const int64_t nthread_use = nthread > 1 ? std::max((int64_t)1, std::min((int64_t)nthread, nchunk)) : 1; - // quantize each expert separately since they have different importance matrices new_size = 0; for (int64_t i03 = 0; i03 < tensor->ne[2]; ++i03) { - const float * f32_data_03 = f32_data + i03 * nelements_matrix; - void * new_data_03 = (char *)new_data + ggml_row_size(new_type, n_per_row) * i03 * nrows; const float * imatrix_03 = imatrix ? imatrix + i03 * n_per_row : nullptr; - new_size += llama_tensor_quantize_impl(new_type, f32_data_03, new_data_03, chunk_size, nrows, n_per_row, imatrix_03, workers, nthread_use); + for (int64_t ir = 0; ir < nrows; ir += nrows_slab) { + const int64_t nrows_cur = std::min(nrows_slab, nrows - ir); + const int64_t nelements_cur = nrows_cur * n_per_row; + + const void * src = load_range((i03*nrows + ir)*row_size_src, nrows_cur*row_size_src); + + const float * f32_data; + if (tensor->type == GGML_TYPE_F32) { + f32_data = (const float *) src; + } else { + if (f32_conv_buf.size() < (size_t) nelements_cur) { + f32_conv_buf.resize(nelements_cur); + } + llama_tensor_dequantize_impl(tensor->type, src, (float *) f32_conv_buf.data(), workers, nelements_cur, nthread); + f32_data = (const float *) f32_conv_buf.data(); + } + + if (work.size() < nrows_cur*row_size_dst) { + work.resize(nrows_cur*row_size_dst); + } + + const int64_t nchunk = (nelements_cur + chunk_size - 1)/chunk_size; + const int64_t nthread_use = nthread > 1 ? std::max((int64_t)1, std::min((int64_t)nthread, nchunk)) : 1; + + const size_t size_cur = llama_tensor_quantize_impl(new_type, f32_data, work.data(), chunk_size, nrows_cur, n_per_row, imatrix_03, workers, nthread_use); + + fout.write((const char *) work.data(), size_cur); + new_size += size_cur; + } } LLAMA_LOG_INFO("size = %8.2f MiB -> %8.2f MiB\n", tensor_size/1024.0/1024.0, new_size/1024.0/1024.0); } @@ -1273,10 +1293,8 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: // update the gguf metadata as we go gguf_set_tensor_type(ctx_outs[cur_split].get(), metadata[i].name.c_str(), new_type); GGML_ASSERT(gguf_get_tensor_size(ctx_outs[cur_split].get(), gguf_find_tensor(ctx_outs[cur_split].get(), metadata[i].name.c_str())) == new_size); - gguf_set_tensor_data(ctx_outs[cur_split].get(), metadata[i].name.c_str(), new_data); - // write tensor data + padding - fout.write((const char *) new_data, new_size); + // tensor data is already written, add the padding zeros(fout, GGML_PAD(new_size, align) - new_size); // unmap the tensor to free memory @@ -1323,7 +1341,8 @@ llama_model_quantize_params llama_model_quantize_default_params() { /*.imatrix =*/ nullptr, /*.kv_overrides =*/ nullptr, /*.tensor_type =*/ nullptr, - /*.prune_layers =*/ nullptr + /*.prune_layers =*/ nullptr, + /*.max_buf_size =*/ LLAMA_QUANT_MAX_BUF_SIZE }; return result; diff --git a/tools/quantize/quantize.cpp b/tools/quantize/quantize.cpp index 8d03c8fcd4..38950036cd 100644 --- a/tools/quantize/quantize.cpp +++ b/tools/quantize/quantize.cpp @@ -122,7 +122,7 @@ static bool try_parse_ftype(const std::string & ftype_str_in, llama_ftype & ftyp static void usage(const char * executable) { printf("usage: %s [--help] [--allow-requantize] [--leave-output-tensor] [--pure] [--imatrix] [--include-weights]\n", executable); printf(" [--exclude-weights] [--output-tensor-type] [--token-embedding-type] [--tensor-type] [--tensor-type-file]\n"); - printf(" [--prune-layers] [--keep-split] [--override-kv] [--dry-run]\n"); + printf(" [--prune-layers] [--keep-split] [--override-kv] [--dry-run] [--max-buffer-size]\n"); printf(" model-f32.gguf [model-quant.gguf] type [nthreads]\n\n"); printf(" --allow-requantize\n"); printf(" allow requantizing tensors that have already been quantized\n"); @@ -161,7 +161,10 @@ static void usage(const char * executable) { printf(" WARNING: this is an advanced option, use with care.\n"); printf(" --dry-run\n"); printf(" calculate and show the final quantization size without performing quantization\n"); - printf(" example: llama-quantize --dry-run model-f32.gguf Q4_K\n\n"); + printf(" example: llama-quantize --dry-run model-f32.gguf Q4_K\n"); + printf(" --max-buffer-size MiB\n"); + printf(" max amount of tensor rows kept in memory while quantizing one tensor (default: 8192)\n"); + printf(" lower it to quantize models with very large tensors on a machine with little RAM\n\n"); printf("note: --include-weights and --exclude-weights cannot be used together\n\n"); printf("-----------------------------------------------------------------------------\n"); printf(" allowed quantization types\n"); @@ -467,6 +470,16 @@ int llama_quantize(int argc, char ** argv) { } } else if (strcmp(argv[arg_idx], "--keep-split") == 0) { params.keep_split = true; + } else if (strcmp(argv[arg_idx], "--max-buffer-size") == 0) { + if (arg_idx == argc-1) { + usage(argv[0]); + } + const int mib = atoi(argv[++arg_idx]); + if (mib <= 0) { + fprintf(stderr, "%s: invalid --max-buffer-size '%s'\n", __func__, argv[arg_idx]); + return 1; + } + params.max_buf_size = (size_t) mib * 1024 * 1024; } else { usage(argv[0]); }