diff --git a/common/arg.cpp b/common/arg.cpp index de9118240..fda1dd52e 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1361,7 +1361,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex add_opt(common_arg( {"--cache-idle-slots"}, {"--no-cache-idle-slots"}, - "save and clear idle slots on new task (default: enabled, requires unified KV and cache-ram)", + "save idle slots to the prompt cache on new task, and clear them when using unified KV (default: enabled, requires cache-ram)", [](common_params & params, bool value) { params.cache_idle_slots = value; } @@ -2222,8 +2222,8 @@ common_params_context common_params_parser_init(common_params & params, llama_ex } ).set_examples(mmproj_examples).set_env("LLAMA_ARG_MMPROJ_OFFLOAD")); add_opt(common_arg( - {"--image", "--audio"}, "FILE", - "path to an image or audio file. use with multimodal models, use comma-separated values for multiple files\n", + {"--image", "--audio", "--video"}, "FILE", + "path to an image, audio, or video file. use with multimodal models, use comma-separated values for multiple files\n", [](common_params & params, const std::string & value) { for (const auto & item : parse_csv_row(value)) { params.image.emplace_back(item); @@ -3334,6 +3334,13 @@ common_params_context common_params_parser_init(common_params & params, llama_ex common_log_set_file(common_log_main(), value.c_str()); } ).set_env("LLAMA_ARG_LOG_FILE")); + add_opt(common_arg( + {"--log-prompts-dir"}, "PATH", + "Log prompts to directory (only used for debugging, default: disabled)", + [](common_params & params, const std::string & value) { + params.path_prompts_log_dir = value; + } + ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); add_opt(common_arg( {"--log-colors"}, "[on|off|auto]", "Set colored logging ('on', 'off', or 'auto', default: 'auto')\n" diff --git a/common/common.h b/common/common.h index 1836947f6..86646ba82 100644 --- a/common/common.h +++ b/common/common.h @@ -490,6 +490,7 @@ struct common_params { std::string input_prefix = ""; // string to prefix user inputs with // NOLINT std::string input_suffix = ""; // string to suffix user inputs with // NOLINT std::string logits_file = ""; // file for saving *all* logits // NOLINT + std::string path_prompts_log_dir = ""; // directory with logged prompts // NOLINT // llama-debug specific options std::string logits_output_dir = "data"; // directory for saving logits output files // NOLINT @@ -572,7 +573,7 @@ struct common_params { struct common_params_model mmproj; bool mmproj_use_gpu = true; // use GPU for multimodal model bool no_mmproj = false; // explicitly disable multimodal model - std::vector image; // path to image file(s) + std::vector image; // path to image file(s) ; TODO: change the name to "media" int image_min_tokens = -1; int image_max_tokens = -1; diff --git a/conversion/gemma.py b/conversion/gemma.py index d8cf8be57..5b4ca5c58 100644 --- a/conversion/gemma.py +++ b/conversion/gemma.py @@ -789,6 +789,16 @@ class Gemma4UnifiedModel(Gemma4Model): class Gemma4AssistantModel(Gemma4Model): model_arch = gguf.MODEL_ARCH.GEMMA4_ASSISTANT + @classmethod + def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: + name, gen = item + + if "masked_embedding" in name: + logger.debug(f"Skipping get tensor {name!r} in safetensors so that convert can end normally.") + return None + + return super().filter_tensors(item) + def set_gguf_parameters(self): super().set_gguf_parameters() self.gguf_writer.add_embedding_length_out(self.hparams["backbone_hidden_size"]) diff --git a/ggml/include/ggml-rpc.h b/ggml/include/ggml-rpc.h index 6fcf5a433..5ad121ae5 100644 --- a/ggml/include/ggml-rpc.h +++ b/ggml/include/ggml-rpc.h @@ -8,10 +8,10 @@ extern "C" { #define RPC_PROTO_MAJOR_VERSION 4 #define RPC_PROTO_MINOR_VERSION 0 -#define RPC_PROTO_PATCH_VERSION 0 +#define RPC_PROTO_PATCH_VERSION 1 #ifdef __cplusplus -static_assert(GGML_OP_COUNT == 96, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION"); +static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT has changed - update RPC_PROTO_PATCH_VERSION"); #endif #define GGML_RPC_MAX_SERVERS 16 diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 489486639..9467cd6b9 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -541,6 +541,7 @@ extern "C" { GGML_OP_IM2COL, GGML_OP_IM2COL_BACK, GGML_OP_IM2COL_3D, + GGML_OP_COL2IM_1D, GGML_OP_CONV_2D, GGML_OP_CONV_3D, GGML_OP_CONV_2D_DW, @@ -2025,6 +2026,16 @@ extern "C" { int d1, // dilation dimension 1 bool is_2D); + // col2im_1d: scatter-add GEMM columns back to 1D signal + // a: [K*OC, T_in] (columns from matmul, K = a->ne[0]/OC) + // result: [T_out, OC] where T_out = (T_in - 1)*s0 + K - 2*p0 + GGML_API struct ggml_tensor * ggml_col2im_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, // columns [K*OC, T_in] + int s0, // stride + int oc, // output channels + int p0); // padding to crop from both sides + GGML_API struct ggml_tensor * ggml_conv_1d( struct ggml_context * ctx, struct ggml_tensor * a, // convolution kernel diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index 0c5c59f63..a3b9f67e9 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -2684,6 +2684,10 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm { ggml_compute_forward_im2col_3d(params, tensor); } break; + case GGML_OP_COL2IM_1D: + { + ggml_compute_forward_col2im_1d(params, tensor); + } break; case GGML_OP_CONV_2D: { ggml_compute_forward_conv_2d(params, tensor); @@ -3156,6 +3160,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { case GGML_OP_CONV_2D: case GGML_OP_CONV_3D: case GGML_OP_CONV_2D_DW: + case GGML_OP_COL2IM_1D: case GGML_OP_CONV_TRANSPOSE_1D: case GGML_OP_CONV_TRANSPOSE_2D: { diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 0686d8bae..2cc163068 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -4008,12 +4008,12 @@ static void ggml_compute_forward_rms_norm_back_f32( // dx := scale(dx, rrms) float * dx = (float *) ((char *) dst->data + i01*nb1 + i02*nb2 + i03*nb3); - // dx[i00] = (x*(-sum_xdz/sum_eps) + dz) / sqrtf(mean_eps) - ggml_vec_cpy_f32 (ne00, dx, x); - // ggml_vec_scale_f32(ne00, dx, -mean_xdz/mean_eps); - ggml_vec_scale_f32(ne00, dx, (float)(-sum_xdz)/sum_eps); - ggml_vec_acc_f32 (ne00, dx, dz); - ggml_vec_scale_f32(ne00, dx, rrms); + // dx[i00] = (dz + x*(-sum_xdz/sum_eps)) * rrms + // note: https://github.com/ggml-org/ggml/issues/1491 + const float scale_x = (float) (-sum_xdz) / sum_eps; + for (int64_t i00 = 0; i00 < ne00; i00++) { + dx[i00] = (dz[i00] + x[i00] * scale_x) * rrms; + } } } } @@ -6730,6 +6730,78 @@ static inline int64_t ggml_wrap_around(int64_t coord, int64_t size) { return (coord + size) % size; // adding size avoids negative number weirdness } +// ggml_compute_forward_col2im_1d +// +// Scatter-add columns [K*OC, T_in] -> signal [T_out, OC] +// where T_out = (T_in - 1)*s + K - 2*p. Gather approach: each output reads ceil(K/s) inputs. +// Parallelized over the time axis so the split stays balanced whatever OC is. +// Supports F32, F16, BF16 input/output (same type), F32 accumulator. + +template +static void ggml_compute_forward_col2im_1d_impl( + const ggml_compute_params * params, + ggml_tensor * dst) { + + const ggml_tensor * src = dst->src[0]; // [K*OC, T_in] + + GGML_ASSERT(ggml_is_contiguous(src)); + GGML_ASSERT(ggml_is_contiguous(dst)); + + const int32_t s0 = ((const int32_t *)(dst->op_params))[0]; + const int32_t OC = ((const int32_t *)(dst->op_params))[1]; + const int32_t p0 = ((const int32_t *)(dst->op_params))[2]; + + const int64_t K_OC = src->ne[0]; + const int64_t T_in = src->ne[1]; + const int64_t K = K_OC / OC; + const int64_t T_out = dst->ne[0]; + + const elem_t * col_data = (const elem_t *) src->data; + elem_t * dst_data = (elem_t *) dst->data; + + const int ith = params->ith; + const int nth = params->nth; + + // Parallelize over the time axis: the split stays balanced whatever OC is, + // down to OC = 1 for mono audio, and threads read disjoint column bands + const int64_t dr = (T_out + nth - 1) / nth; + const int64_t it0 = dr * ith; + const int64_t it1 = it0 + dr < T_out ? it0 + dr : T_out; + + for (int64_t oc = 0; oc < OC; oc++) { + for (int64_t t_out = it0; t_out < it1; t_out++) { + const int64_t t_abs = t_out + p0; // absolute position in uncropped signal + // Gather: find all (t_in, k) where t_in * s + k == t_abs, 0 <= k < K + int64_t t_in_min = (t_abs - K + 1 + s0 - 1) / s0; // ceil((t_abs-K+1)/s) + if (t_in_min < 0) t_in_min = 0; + int64_t t_in_max = t_abs / s0; + if (t_in_max >= T_in) t_in_max = T_in - 1; + + float sum = 0.0f; + for (int64_t t_in = t_in_min; t_in <= t_in_max; t_in++) { + int64_t k = t_abs - t_in * s0; + if (k >= 0 && k < K) { + // col layout: [K*OC, T_in], element (oc*K+k, t_in) + sum += type_conversion_table::to_f32(col_data[(oc * K + k) + t_in * K_OC]); + } + } + // dst layout: [T_out, OC], element (t_out, oc) + dst_data[t_out + oc * T_out] = type_conversion_table::from_f32(sum); + } + } +} + +void ggml_compute_forward_col2im_1d( + const ggml_compute_params * params, + ggml_tensor * dst) { + switch (dst->src[0]->type) { + case GGML_TYPE_F32: ggml_compute_forward_col2im_1d_impl (params, dst); break; + case GGML_TYPE_F16: ggml_compute_forward_col2im_1d_impl(params, dst); break; + case GGML_TYPE_BF16: ggml_compute_forward_col2im_1d_impl(params, dst); break; + default: GGML_ABORT("col2im_1d: unsupported type %d", dst->src[0]->type); + } +} + // ggml_compute_forward_conv_2d diff --git a/ggml/src/ggml-cpu/ops.h b/ggml/src/ggml-cpu/ops.h index 7398e5618..a8e18c716 100644 --- a/ggml/src/ggml-cpu/ops.h +++ b/ggml/src/ggml-cpu/ops.h @@ -68,6 +68,7 @@ void ggml_compute_forward_conv_transpose_1d(const struct ggml_compute_params * p void ggml_compute_forward_im2col(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_im2col_back_f32(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_im2col_3d(const struct ggml_compute_params * params, struct ggml_tensor * dst); +void ggml_compute_forward_col2im_1d(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_conv_2d(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_conv_3d(const struct ggml_compute_params * params, struct ggml_tensor * dst); void ggml_compute_forward_conv_transpose_2d(const struct ggml_compute_params * params, struct ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index bdfbfd2d3..fe44a58da 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -411,7 +411,6 @@ static constexpr __host__ __device__ int calc_nwarps(ggml_type type, int ncols_d case GGML_TYPE_Q5_0: case GGML_TYPE_Q5_1: case GGML_TYPE_Q8_0: - case GGML_TYPE_Q4_K: return 8; case GGML_TYPE_Q6_K: return 2; diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index bf9bcbe09..4101c5593 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -119,6 +119,21 @@ typedef struct VkPhysicalDeviceShaderBfloat16FeaturesKHR { } VkPhysicalDeviceShaderBfloat16FeaturesKHR; #endif +#if !defined(VK_VALVE_shader_mixed_float_dot_product) +#define VK_VALVE_shader_mixed_float_dot_product 1 +#define VK_VALVE_SHADER_MIXED_FLOAT_DOT_PRODUCT_SPEC_VERSION 1 +#define VK_VALVE_SHADER_MIXED_FLOAT_DOT_PRODUCT_EXTENSION_NAME "VK_VALVE_shader_mixed_float_dot_product" +#define VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_MIXED_FLOAT_DOT_PRODUCT_FEATURES_VALVE ((VkStructureType)1000673000) +typedef struct VkPhysicalDeviceShaderMixedFloatDotProductFeaturesVALVE { + VkStructureType sType; + void* pNext; + VkBool32 shaderMixedFloatDotProductFloat16AccFloat32; + VkBool32 shaderMixedFloatDotProductFloat16AccFloat16; + VkBool32 shaderMixedFloatDotProductBFloat16Acc; + VkBool32 shaderMixedFloatDotProductFloat8AccFloat32; +} VkPhysicalDeviceShaderMixedFloatDotProductFeaturesVALVE; +#endif + #define ROUNDUP_POW2(M, N) (((M) + (N) - 1) & ~((N) - 1)) #define CEIL_DIV(M, N) (((M) + (N)-1) / (N)) static bool is_pow2(uint32_t x) { return x > 1 && (x & (x-1)) == 0; } @@ -711,6 +726,8 @@ struct vk_device_struct { bool coopmat2_bf16_support {}; bool coopmat2_decode_vector; + bool dot2_f16 {}; + bool pipeline_executable_properties_support {}; size_t idx; @@ -3383,7 +3400,9 @@ static bool ggml_vk_matmul_shmem_support(const vk_device& device, const std::vec switch (src0_type) { case GGML_TYPE_IQ1_S: case GGML_TYPE_IQ1_M: - lut_size = 2*2048 + 4*2048; + // Regular matmul uses the compact uint16_t IQ1 grid; the expanded + // uint32_t grid is only enabled for the q8_1/int-dot vector path. + lut_size = 2*2048; break; case GGML_TYPE_IQ2_XXS: lut_size = 8*256; @@ -3926,8 +3945,13 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { name = aligned ? "flash_attn_f32_f16_aligned" : "flash_attn_f32_f16"; } else { if (device->fp16) { - if (f32acc) { spv_data = flash_attn_f32_f16_data; spv_size = flash_attn_f32_f16_len; } - else { spv_data = flash_attn_f32_f16_f16acc_data; spv_size = flash_attn_f32_f16_f16acc_len; } + if (device->dot2_f16) { + if (f32acc) { spv_data = flash_attn_f32_f16_dot2_data; spv_size = flash_attn_f32_f16_dot2_len; } + else { spv_data = flash_attn_f32_f16_dot2_f16acc_data; spv_size = flash_attn_f32_f16_dot2_f16acc_len; } + } else { + if (f32acc) { spv_data = flash_attn_f32_f16_data; spv_size = flash_attn_f32_f16_len; } + else { spv_data = flash_attn_f32_f16_f16acc_data; spv_size = flash_attn_f32_f16_f16acc_len; } + } } else { spv_data = flash_attn_f32_f16_fp32_data; spv_size = flash_attn_f32_f16_fp32_len; @@ -4221,7 +4245,23 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { #endif // defined(VK_KHR_cooperative_matrix) && defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT) if (device->fp16) { // Create 6 variants, {s,m,l}x{unaligned,aligned} + // Selects dot2 SPIR-V variant at runtime when device->dot2_f16 is true #define CREATE_MM(TYPE, PIPELINE_NAME, NAMELC, F16ACC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID, REQSUBGROUPSIZE) \ + if (device->mul_mat ## ID ## _l[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->l, #NAMELC #F16ACC "_l", (device->dot2_f16 ? NAMELC ## _dot2 ## F16ACC ## _len : NAMELC ## F16ACC ## _len), (device->dot2_f16 ? NAMELC ## _dot2 ## F16ACC ## _data : NAMELC ## F16ACC ## _data), "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, 1, false, REQSUBGROUPSIZE > 0, REQSUBGROUPSIZE); \ + if (device->mul_mat ## ID ## _m[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->m, #NAMELC #F16ACC "_m", (device->dot2_f16 ? NAMELC ## _dot2 ## F16ACC ## _len : NAMELC ## F16ACC ## _len), (device->dot2_f16 ? NAMELC ## _dot2 ## F16ACC ## _data : NAMELC ## F16ACC ## _data), "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, 1, false, REQSUBGROUPSIZE > 0, REQSUBGROUPSIZE); \ + if (device->mul_mat ## ID ## _s[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->s, #NAMELC #F16ACC "_s", (device->dot2_f16 ? NAMELC ## _dot2 ## F16ACC ## _len : NAMELC ## F16ACC ## _len), (device->dot2_f16 ? NAMELC ## _dot2 ## F16ACC ## _data : NAMELC ## F16ACC ## _data), "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, 1, false, REQSUBGROUPSIZE > 0, REQSUBGROUPSIZE); \ + if (device->mul_mat ## ID ## _l[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_l, #NAMELC #F16ACC "_aligned_l", (device->dot2_f16 ? NAMELC ## _dot2_aligned ## F16ACC ## _len : NAMELC ## _aligned ## F16ACC ## _len), (device->dot2_f16 ? NAMELC ## _dot2_aligned ## F16ACC ## _data : NAMELC ## _aligned ## F16ACC ## _data), "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, l_align, false, REQSUBGROUPSIZE > 0, REQSUBGROUPSIZE); \ + if (device->mul_mat ## ID ## _m[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_m, #NAMELC #F16ACC "_aligned_m", (device->dot2_f16 ? NAMELC ## _dot2_aligned ## F16ACC ## _len : NAMELC ## _aligned ## F16ACC ## _len), (device->dot2_f16 ? NAMELC ## _dot2_aligned ## F16ACC ## _data : NAMELC ## _aligned ## F16ACC ## _data), "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, m_align, false, REQSUBGROUPSIZE > 0, REQSUBGROUPSIZE); \ + if (device->mul_mat ## ID ## _s[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_s, #NAMELC #F16ACC "_aligned_s", (device->dot2_f16 ? NAMELC ## _dot2_aligned ## F16ACC ## _len : NAMELC ## _aligned ## F16ACC ## _len), (device->dot2_f16 ? NAMELC ## _dot2_aligned ## F16ACC ## _data : NAMELC ## _aligned ## F16ACC ## _data), "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, s_align, false, REQSUBGROUPSIZE > 0, REQSUBGROUPSIZE); \ + + // bf16 scalar path promotes to f32, no dot2 variant +#define CREATE_MM_NODOT2(TYPE, PIPELINE_NAME, NAMELC, F16ACC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID, REQSUBGROUPSIZE) \ if (device->mul_mat ## ID ## _l[TYPE]) \ ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->l, #NAMELC #F16ACC "_l", NAMELC ## F16ACC ## _len, NAMELC ## F16ACC ## _data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, 1, false, REQSUBGROUPSIZE > 0, REQSUBGROUPSIZE); \ if (device->mul_mat ## ID ## _m[TYPE]) \ @@ -4256,7 +4296,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { CREATE_MM2(GGML_TYPE_F16, pipeline_matmul_f16, matmul_f16, wg_denoms, warptile, vk_mat_mat_push_constants, 3, , 0); CREATE_MM2(GGML_TYPE_F16, pipeline_matmul_f16_f32, matmul_f16_f32, wg_denoms, warptile, vk_mat_mat_push_constants, 3, , 0); - CREATE_MM(GGML_TYPE_BF16, pipeline_matmul_bf16, matmul_bf16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM_NODOT2(GGML_TYPE_BF16, pipeline_matmul_bf16, matmul_bf16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, , 0); CREATE_MM2(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q1_0], matmul_q1_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); CREATE_MM2(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_0], matmul_q4_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); @@ -4264,7 +4304,6 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { CREATE_MM2(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_0], matmul_q5_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); CREATE_MM2(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_1], matmul_q5_1_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); CREATE_MM2(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q8_0], matmul_q8_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); - CREATE_MM2(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_K], matmul_q2_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); CREATE_MM2(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_K], matmul_q3_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); CREATE_MM2(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_K], matmul_q4_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); @@ -4304,8 +4343,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { CREATE_MM(GGML_TYPE_F32, pipeline_matmul_id_f32, matmul_id_subgroup_f32_f32, , wg_denoms, warptile_id, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size_16); CREATE_MM2(GGML_TYPE_F16, pipeline_matmul_id_f16, matmul_id_subgroup_f16, wg_denoms, warptile_id, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size_16); CREATE_MM2(GGML_TYPE_F16, pipeline_matmul_id_f16_f32, matmul_id_subgroup_f16_f32, wg_denoms, warptile_id, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size_16); - CREATE_MM(GGML_TYPE_BF16, pipeline_matmul_id_bf16, matmul_id_subgroup_bf16, , wg_denoms, warptile_id, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size_16); - + CREATE_MM_NODOT2(GGML_TYPE_BF16, pipeline_matmul_id_bf16, matmul_id_subgroup_bf16, , wg_denoms, warptile_id, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size_16); CREATE_MM2(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q1_0], matmul_id_subgroup_q1_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); CREATE_MM2(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0], matmul_id_subgroup_q4_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); CREATE_MM2(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1], matmul_id_subgroup_q4_1_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); @@ -4350,8 +4388,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { CREATE_MM(GGML_TYPE_F32, pipeline_matmul_id_f32, matmul_id_f32_f32, , wg_denoms, warptile, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); CREATE_MM2(GGML_TYPE_F16, pipeline_matmul_id_f16, matmul_id_f16, wg_denoms, warptile, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); CREATE_MM2(GGML_TYPE_F16, pipeline_matmul_id_f16_f32, matmul_id_f16_f32, wg_denoms, warptile, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); - CREATE_MM(GGML_TYPE_BF16, pipeline_matmul_id_bf16, matmul_id_bf16, , wg_denoms, warptile, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); - + CREATE_MM_NODOT2(GGML_TYPE_BF16, pipeline_matmul_id_bf16, matmul_id_bf16, , wg_denoms, warptile, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); CREATE_MM2(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q1_0], matmul_id_q1_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); CREATE_MM2(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0], matmul_id_q4_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); CREATE_MM2(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1], matmul_id_q4_1_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); @@ -4396,6 +4433,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { #undef CREATE_MM2 #undef CREATE_MMQ #undef CREATE_MM +#undef CREATE_MM_NODOT2 } else { // Create 6 variants, {s,m,l}x{unaligned,aligned} #define CREATE_MM(TYPE, PIPELINE_NAME, NAMELC, F16ACC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID, REQSUBGROUPSIZE) \ @@ -5465,6 +5503,7 @@ static vk_device ggml_vk_get_device(size_t idx) { device->integer_dot_product = false; device->shader_64b_indexing = false; bool bfloat16_support = false; + bool dot2_f16_support = false; for (const auto& properties : ext_props) { if (strcmp("VK_KHR_maintenance4", properties.extensionName) == 0) { @@ -5507,6 +5546,9 @@ static vk_device ggml_vk_get_device(size_t idx) { !getenv("GGML_VK_DISABLE_BFLOAT16")) { bfloat16_support = true; #endif + } else if (strcmp("VK_VALVE_shader_mixed_float_dot_product", properties.extensionName) == 0 && + !getenv("GGML_VK_DISABLE_DOT2")) { + dot2_f16_support = true; } else if (strcmp("VK_KHR_pipeline_executable_properties", properties.extensionName) == 0) { pipeline_executable_properties_support = true; } else if (strcmp("VK_EXT_memory_priority", properties.extensionName) == 0 && @@ -5822,6 +5864,14 @@ static vk_device ggml_vk_get_device(size_t idx) { device_extensions.push_back("VK_KHR_shader_integer_dot_product"); } + VkPhysicalDeviceShaderMixedFloatDotProductFeaturesVALVE dot2_features {}; + dot2_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_MIXED_FLOAT_DOT_PRODUCT_FEATURES_VALVE; + if (dot2_f16_support) { + last_struct->pNext = (VkBaseOutStructure *)&dot2_features; + last_struct = (VkBaseOutStructure *)&dot2_features; + device_extensions.push_back("VK_VALVE_shader_mixed_float_dot_product"); + } + VkPhysicalDevicePipelineExecutablePropertiesFeaturesKHR pep_features {}; pep_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_PIPELINE_EXECUTABLE_PROPERTIES_FEATURES_KHR; if (pipeline_executable_properties_support) { @@ -5856,6 +5906,8 @@ static vk_device ggml_vk_get_device(size_t idx) { device->bf16 = false; #endif + device->dot2_f16 = dot2_f16_support && dot2_features.shaderMixedFloatDotProductFloat16AccFloat32; + device->pipeline_robustness = pl_robustness_features.pipelineRobustness; device->multi_add = vk12_props.shaderRoundingModeRTEFloat16 && @@ -6270,6 +6322,7 @@ static void ggml_vk_print_gpu_info(size_t idx) { bool coopmat2_decode_vector_support = false; bool integer_dot_product = false; bool bfloat16_support = false; + bool dot2_f16_support = false; for (auto properties : ext_props) { if (strcmp("VK_KHR_16bit_storage", properties.extensionName) == 0) { @@ -6299,6 +6352,9 @@ static void ggml_vk_print_gpu_info(size_t idx) { !getenv("GGML_VK_DISABLE_BFLOAT16")) { bfloat16_support = true; #endif + } else if (strcmp("VK_VALVE_shader_mixed_float_dot_product", properties.extensionName) == 0 && + !getenv("GGML_VK_DISABLE_DOT2")) { + dot2_f16_support = true; } } @@ -6389,6 +6445,13 @@ static void ggml_vk_print_gpu_info(size_t idx) { last_struct = (VkBaseOutStructure *)&coopmat2_decode_vector_features; } + VkPhysicalDeviceShaderMixedFloatDotProductFeaturesVALVE dot2_features {}; + dot2_features.sType = VK_STRUCTURE_TYPE_PHYSICAL_DEVICE_SHADER_MIXED_FLOAT_DOT_PRODUCT_FEATURES_VALVE; + if (dot2_f16_support) { + last_struct->pNext = (VkBaseOutStructure *)&dot2_features; + last_struct = (VkBaseOutStructure *)&dot2_features; + } + vkGetPhysicalDeviceFeatures2(physical_device, &device_features2); fp16 = fp16 && vk12_features.shaderFloat16; @@ -6435,9 +6498,12 @@ static void ggml_vk_print_gpu_info(size_t idx) { : coopmat_support ? "KHR_coopmat" : "none"; + bool dot2_f16 = dot2_f16_support && dot2_features.shaderMixedFloatDotProductFloat16AccFloat32; + const char *fp16_str = fp16 ? (dot2_f16 ? "dot2" : "1") : "0"; + std::string device_name = props2.properties.deviceName.data(); - GGML_LOG_DEBUG("ggml_vulkan: %zu = %s (%s) | uma: %d | fp16: %d | bf16: %d | warp size: %zu | shared memory: %d | int dot: %d | matrix cores: %s\n", - idx, device_name.c_str(), driver_props.driverName.data(), uma, fp16, bf16, subgroup_size, + GGML_LOG_DEBUG("ggml_vulkan: %zu = %s (%s) | uma: %d | fp16: %s | bf16: %d | warp size: %zu | shared memory: %d | int dot: %d | matrix cores: %s\n", + idx, device_name.c_str(), driver_props.driverName.data(), uma, fp16_str, bf16, subgroup_size, props2.properties.limits.maxComputeSharedMemorySize, integer_dot_product, matrix_cores.c_str()); if (props2.properties.deviceType == vk::PhysicalDeviceType::eCpu) { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dot_product_funcs.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/dot_product_funcs.glsl new file mode 100644 index 000000000..c474bfe09 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dot_product_funcs.glsl @@ -0,0 +1,27 @@ +#ifdef DOT2_F16 +#extension GL_EXT_spirv_intrinsics : require + +spirv_instruction(extensions = ["SPV_VALVE_mixed_float_dot_product"], + capabilities = [6912], id = 6916) +float v_dot2_f32_f16(f16vec2 a, f16vec2 b, float acc); + +ACC_TYPE dot_product(f16vec4 a, f16vec4 b, ACC_TYPE acc) { + return ACC_TYPE(v_dot2_f32_f16(a.zw, b.zw, v_dot2_f32_f16(a.xy, b.xy, float(acc)))); +} + +ACC_TYPE dot_product(f16vec2 a, f16vec2 b, ACC_TYPE acc) { + return ACC_TYPE(v_dot2_f32_f16(a, b, float(acc))); +} + +#else + +ACC_TYPE dot_product(FLOAT_TYPEV4 a, FLOAT_TYPEV4 b, ACC_TYPE acc) { + return fma(ACC_TYPE(a.x), ACC_TYPE(b.x), fma(ACC_TYPE(a.y), ACC_TYPE(b.y), + fma(ACC_TYPE(a.z), ACC_TYPE(b.z), fma(ACC_TYPE(a.w), ACC_TYPE(b.w), acc)))); +} + +ACC_TYPE dot_product(FLOAT_TYPEV2 a, FLOAT_TYPEV2 b, ACC_TYPE acc) { + return fma(ACC_TYPE(a.x), ACC_TYPE(b.x), fma(ACC_TYPE(a.y), ACC_TYPE(b.y), acc)); +} + +#endif diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp index 6ac095489..91fb07c93 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp @@ -21,6 +21,7 @@ #extension GL_KHR_shader_subgroup_vote : enable #include "types.glsl" +#include "dot_product_funcs.glsl" #include "flash_attn_base.glsl" #include "flash_attn_dequant.glsl" @@ -318,7 +319,7 @@ void main() { K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + (j * Bc + c * cols_per_iter + col_tid) * k_stride / 4 + d * D_split + d_tid]); } [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { - Sf[r][c] += dot(ACC_TYPEV4(Q_cache[r]), ACC_TYPEV4(K_Tf)); + Sf[r][c] = dot_product(Q_cache[r], K_Tf, Sf[r][c]); } } } @@ -341,7 +342,7 @@ void main() { K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + (j * Bc + c * cols_per_iter + col_tid) * k_stride / 4 + d * D_split + d_tid]); } [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { - Sf[r][c] += dot(ACC_TYPEV4(Qf[tile_row(r) * qf_stride + d * D_split + d_tid]), ACC_TYPEV4(K_Tf)); + Sf[r][c] = dot_product(Qf[tile_row(r) * qf_stride + d * D_split + d_tid], K_Tf, Sf[r][c]); } } } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vecq.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vecq.comp index 6fe3e2dc0..fd84c3c91 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vecq.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vecq.comp @@ -4,6 +4,7 @@ #extension GL_EXT_integer_dot_product : require #define MMQ +#define NEEDS_IQ1S_GRID_GPU #define B_TYPE block_q8_1_x4 #include "mul_mat_vec_base.glsl" diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp index 89346e48e..f39410d74 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp @@ -29,6 +29,7 @@ #endif #include "types.glsl" +#include "dot_product_funcs.glsl" #ifndef LOAD_VEC_A #define LOAD_VEC_A 1 @@ -329,15 +330,8 @@ void main() { [[unroll]] for (uint cr = 0; cr < TM / 2; cr++) { // [WNITER][TN][WMITER][TM / 2] -> [wsic][cc][wsir][cr] const uint sums_idx = (wsic * TN + cc) * WMITER * (TM / 2) + wsir * (TM / 2) + cr; - #if defined(DATA_A_F32) || defined(DATA_A_F16) - sums[sums_idx].x = fma(ACC_TYPE(cache_a[wsir * TM + 2 * cr ].x), ACC_TYPE(cache_b.x), fma(ACC_TYPE(cache_a[wsir * TM + 2 * cr ].y), ACC_TYPE(cache_b.y), - fma(ACC_TYPE(cache_a[wsir * TM + 2 * cr ].z), ACC_TYPE(cache_b.z), fma(ACC_TYPE(cache_a[wsir * TM + 2 * cr ].w), ACC_TYPE(cache_b.w), sums[sums_idx].x)))); - sums[sums_idx].y = fma(ACC_TYPE(cache_a[wsir * TM + 2 * cr + 1].x), ACC_TYPE(cache_b.x), fma(ACC_TYPE(cache_a[wsir * TM + 2 * cr + 1].y), ACC_TYPE(cache_b.y), - fma(ACC_TYPE(cache_a[wsir * TM + 2 * cr + 1].z), ACC_TYPE(cache_b.z), fma(ACC_TYPE(cache_a[wsir * TM + 2 * cr + 1].w), ACC_TYPE(cache_b.w), sums[sums_idx].y)))); - #else - sums[sums_idx].x = fma(ACC_TYPE(cache_a[wsir * TM + 2 * cr ].x), ACC_TYPE(cache_b.x), fma(ACC_TYPE(cache_a[wsir * TM + 2 * cr ].y), ACC_TYPE(cache_b.y), sums[sums_idx].x)); - sums[sums_idx].y = fma(ACC_TYPE(cache_a[wsir * TM + 2 * cr + 1].x), ACC_TYPE(cache_b.x), fma(ACC_TYPE(cache_a[wsir * TM + 2 * cr + 1].y), ACC_TYPE(cache_b.y), sums[sums_idx].y)); - #endif + sums[sums_idx].x = dot_product(cache_a[wsir * TM + 2 * cr ], cache_b, sums[sums_idx].x); + sums[sums_idx].y = dot_product(cache_a[wsir * TM + 2 * cr + 1], cache_b, sums[sums_idx].y); } } } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/types.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/types.glsl index f84d6f873..8c6b20c68 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/types.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/types.glsl @@ -598,9 +598,10 @@ const uint[1024] iq1s_grid_const = { 0x55dd55df, 0x55d555d7, 0x5503550c, 0x557f5501, 0x5577557d, 0x55405575, 0x555d555f, 0x55555557 }; +#if defined(NEEDS_IQ1S_GRID_GPU) // Same content as iq1s_grid_const except each 2-bit value is expanded to 4-bit // and has 1 added to it (allows packed values to be extracted with & 0x0F0F0F0F -// and 0xF0F0F0F0). +// and 0xF0F0F0F0). This is only used by the q8_1/int-dot vector path. const uint32_t[2048] iq1s_grid_gpu_const = { 0x00000000, 0x00000002, 0x00000101, 0x00000200, 0x00000202, 0x00010001, 0x00010101, 0x00020000, 0x00020002, 0x00020200, 0x00020202, 0x01000101, 0x01010001, 0x01010100, 0x01010102, 0x01020101, @@ -859,9 +860,12 @@ const uint32_t[2048] iq1s_grid_gpu_const = { 0x20222020, 0x20222022, 0x20222220, 0x20222222, 0x21212021, 0x21212120, 0x21212122, 0x22202020, 0x22202022, 0x22202220, 0x22202222, 0x22212121, 0x22222020, 0x22222022, 0x22222220, 0x22222222, }; +#endif shared uint16_t iq1s_grid[2048]; +#if defined(NEEDS_IQ1S_GRID_GPU) shared uint32_t iq1s_grid_gpu[2048]; +#endif #define NEEDS_INIT_IQ_SHMEM void init_iq_shmem(uvec3 wgsize) @@ -875,12 +879,14 @@ void init_iq_shmem(uvec3 wgsize) iq1s_grid[2*idx+1] = g.y; } } +#if defined(NEEDS_IQ1S_GRID_GPU) [[unroll]] for (uint i = 0; i < iq1s_grid_gpu_const.length(); i += wgsize.x) { uint idx = i + gl_LocalInvocationIndex.x; if (iq1s_grid_gpu_const.length() % wgsize.x == 0 || idx < iq1s_grid_gpu_const.length()) { iq1s_grid_gpu[idx] = iq1s_grid_gpu_const[idx]; } } +#endif barrier(); } #endif 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 d1c9618c5..1d653c336 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -350,7 +350,8 @@ void string_to_spv_func(std::string name, std::string in_path, std::string out_p // disable spirv-opt for coopmat shaders for https://github.com/ggml-org/llama.cpp/issues/10734 // disable spirv-opt for bf16 shaders for https://github.com/ggml-org/llama.cpp/issues/15344 // disable spirv-opt for rope shaders for https://github.com/ggml-org/llama.cpp/issues/16860 - if (!coopmat && name.find("bf16") == std::string::npos && name.find("rope") == std::string::npos) { + // disable spirv-opt for dot2 shaders (spirv-opt doesn't recognize SPV_VALVE_mixed_float_dot_product capability) + if (!coopmat && name.find("bf16") == std::string::npos && name.find("rope") == std::string::npos && name.find("_dot2") == std::string::npos) { cmd.push_back("-O"); } @@ -444,10 +445,11 @@ void string_to_spv(std::string name, const std::string& source, const std::map base_dict; std::string shader_name = "matmul"; @@ -480,6 +482,10 @@ void matmul_shaders(bool fp16, MatMulIdType matmul_id_type, bool coopmat, bool c } #endif + if (dot2) { + base_dict["DOT2_F16"] = "1"; + } + const std::string source_name = coopmat2 ? "mul_mm_cm2.comp" : "mul_mm.comp"; auto const &FLOAT_TYPE = [&](int vec, const std::string &t) -> std::string { @@ -545,11 +551,11 @@ void matmul_shaders(bool fp16, MatMulIdType matmul_id_type, bool coopmat, bool c }; // Shaders with f16 B_TYPE - string_to_spv(shader_name + "_f32_f16", source_name, merge_maps(merge_maps(base_dict, float_type_dict_f16), {{"DATA_A_F32", "1"}, {"B_TYPE", "float16_t"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}, }), fp16, coopmat, coopmat2, f16acc); - string_to_spv(shader_name + "_f32_f16_aligned", source_name, merge_maps(merge_maps(base_dict, float_type_dict_f16), {{"DATA_A_F32", "1"}, {"LOAD_VEC_A", load_vec}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}, {"ALIGNED", "1"}}), fp16, coopmat, coopmat2, f16acc); + string_to_spv(shader_name + "_f32_f16" + dot2_sfx, source_name, merge_maps(merge_maps(base_dict, float_type_dict_f16), {{"DATA_A_F32", "1"}, {"B_TYPE", "float16_t"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}, }), fp16, coopmat, coopmat2, f16acc); + string_to_spv(shader_name + "_f32_f16" + dot2_sfx + "_aligned", source_name, merge_maps(merge_maps(base_dict, float_type_dict_f16), {{"DATA_A_F32", "1"}, {"LOAD_VEC_A", load_vec}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}, {"ALIGNED", "1"}}), fp16, coopmat, coopmat2, f16acc); - string_to_spv(shader_name + "_f16", source_name, merge_maps(merge_maps(base_dict, float_type_dict_f16), {{"DATA_A_F16", "1"}, {"B_TYPE", "float16_t"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc); - string_to_spv(shader_name + "_f16_aligned", source_name, merge_maps(merge_maps(base_dict, float_type_dict_f16), {{"DATA_A_F16", "1"}, {"LOAD_VEC_A", load_vec}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}, {"ALIGNED", "1"}}), fp16, coopmat, coopmat2, f16acc); + string_to_spv(shader_name + "_f16" + dot2_sfx, source_name, merge_maps(merge_maps(base_dict, float_type_dict_f16), {{"DATA_A_F16", "1"}, {"B_TYPE", "float16_t"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc); + string_to_spv(shader_name + "_f16" + dot2_sfx + "_aligned", source_name, merge_maps(merge_maps(base_dict, float_type_dict_f16), {{"DATA_A_F16", "1"}, {"LOAD_VEC_A", load_vec}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}, {"ALIGNED", "1"}}), fp16, coopmat, coopmat2, f16acc); // bf16 { @@ -570,8 +576,10 @@ void matmul_shaders(bool fp16, MatMulIdType matmul_id_type, bool coopmat, bool c if (!(coopmat || coopmat2)) #endif { - string_to_spv(shader_name + "_bf16", source_name, merge_maps(merge_maps(base_dict, float_type_dict_bf16), {{"TO_FLOAT_TYPE", to_float_type}, {"DATA_A_BF16", "1"}, {"B_TYPE", coopmat2 ? "bfloat16_t" : "uint16_t"}, {"B_TYPEV4", "bf16vec4"}, {"D_TYPE", "float"}, {"B_IS_FLOAT", "1"}, {"DATA_B_BF16", "1"}}), fp16, coopmat, coopmat2, f16acc); - string_to_spv(shader_name + "_bf16_aligned", source_name, merge_maps(merge_maps(base_dict, float_type_dict_bf16), {{"TO_FLOAT_TYPE", to_float_type}, {"DATA_A_BF16", "1"}, {"LOAD_VEC_A", load_vec_a}, {"LOAD_VEC_B", "4"}, {"B_TYPE", coopmat2 ? "bfloat16_t" : "u16vec4"}, {"B_TYPEV4", "bf16vec4"}, {"D_TYPE", "float"}, {"B_IS_FLOAT", "1"}, {"DATA_B_BF16", "1"}, {"ALIGNED", "1"}}), fp16, coopmat, coopmat2, f16acc); + if (!dot2) { + string_to_spv(shader_name + "_bf16", source_name, merge_maps(merge_maps(base_dict, float_type_dict_bf16), {{"TO_FLOAT_TYPE", to_float_type}, {"DATA_A_BF16", "1"}, {"B_TYPE", coopmat2 ? "bfloat16_t" : "uint16_t"}, {"B_TYPEV4", "bf16vec4"}, {"D_TYPE", "float"}, {"B_IS_FLOAT", "1"}, {"DATA_B_BF16", "1"}}), fp16, coopmat, coopmat2, f16acc); + string_to_spv(shader_name + "_bf16_aligned", source_name, merge_maps(merge_maps(base_dict, float_type_dict_bf16), {{"TO_FLOAT_TYPE", to_float_type}, {"DATA_A_BF16", "1"}, {"LOAD_VEC_A", load_vec_a}, {"LOAD_VEC_B", "4"}, {"B_TYPE", coopmat2 ? "bfloat16_t" : "u16vec4"}, {"B_TYPEV4", "bf16vec4"}, {"D_TYPE", "float"}, {"B_IS_FLOAT", "1"}, {"DATA_B_BF16", "1"}, {"ALIGNED", "1"}}), fp16, coopmat, coopmat2, f16acc); + } } } @@ -601,18 +609,18 @@ void matmul_shaders(bool fp16, MatMulIdType matmul_id_type, bool coopmat, bool c // don't generate f32 variants for coopmat2 if (!coopmat2) { - string_to_spv(shader_name + "_" + tname + "_f32", source_name, merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a_unaligned}, {"B_TYPE", "float"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc); - string_to_spv(shader_name + "_" + tname + "_f32_aligned", source_name, merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f32}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}, {"ALIGNED", "1"}}), fp16, coopmat, coopmat2, f16acc); + string_to_spv(shader_name + "_" + tname + "_f32" + dot2_sfx, source_name, merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a_unaligned}, {"B_TYPE", "float"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc); + string_to_spv(shader_name + "_" + tname + "_f32" + dot2_sfx + "_aligned", source_name, merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f32}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}, {"ALIGNED", "1"}}), fp16, coopmat, coopmat2, f16acc); } if (tname != "f16" && tname != "f32") { - string_to_spv(shader_name + "_" + tname + "_f16", source_name, merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a_unaligned}, {"B_TYPE", "float16_t"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc); - string_to_spv(shader_name + "_" + tname + "_f16_aligned", source_name, merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}, {"ALIGNED", "1"}}), fp16, coopmat, coopmat2, f16acc); + string_to_spv(shader_name + "_" + tname + "_f16" + dot2_sfx, source_name, merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a_unaligned}, {"B_TYPE", "float16_t"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc); + string_to_spv(shader_name + "_" + tname + "_f16" + dot2_sfx + "_aligned", source_name, merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}, {"ALIGNED", "1"}}), fp16, coopmat, coopmat2, f16acc); } #if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT) - // Integer dot mmq performs better with f32 accumulators - if (!f16acc && !coopmat && !coopmat2 && (is_legacy_quant(tname) || is_k_quant(tname) || tname == "mxfp4")) { + // Integer dot mmq performs better with f32 accumulators (different shader, skip for dot2) + if (!f16acc && !coopmat && !coopmat2 && !dot2 && (is_legacy_quant(tname) || is_k_quant(tname) || tname == "mxfp4")) { string_to_spv(shader_name + "_" + tname + "_q8_1", "mul_mmq.comp", merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"D_TYPE", "float"},}), fp16, coopmat, coopmat2, f16acc); } #endif @@ -630,6 +638,10 @@ void process_shaders() { matmul_shaders(true, matmul_id_type, false, false, false); matmul_shaders(true, matmul_id_type, false, false, true); + // dot2 variants (scalar fp16 only) + matmul_shaders(true, matmul_id_type, false, false, false, true); + matmul_shaders(true, matmul_id_type, false, false, true, true); + if (matmul_id_type != MatMulIdType::DEFAULT) { #if defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT) // Coopmat, fp32acc and fp16acc @@ -677,6 +689,12 @@ void process_shaders() { string_to_spv("flash_attn_f32_f16", "flash_attn.comp", merge_maps(fa_base_dict, {{"Q_TYPE", "float"}, {"D_TYPE", "float"}, {"D_TYPEV4", "vec4"}}), fp16, false, false, f16acc); + + if (fp16) { + string_to_spv("flash_attn_f32_f16_dot2", "flash_attn.comp", + merge_maps(fa_base_dict, {{"Q_TYPE", "float"}, {"D_TYPE", "float"}, {"D_TYPEV4", "vec4"}, {"DOT2_F16", "1"}}), fp16, false, false, f16acc); + } + #if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT) string_to_spv("flash_attn_f32_f16", "flash_attn.comp", merge_maps(fa_base_dict, {{"Q_TYPE", "float"}, {"D_TYPE", "float"}, {"D_TYPEV4", "vec4"}, {"MMQ", "1"}, {"FA_MMQ_MIXED", "1"}}), fp16, false, false, f16acc, "_int8"); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index dc1d5c6f0..389499e5b 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -1047,6 +1047,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "IM2COL", "IM2COL_BACK", "IM2COL_3D", + "COL2IM_1D", "CONV_2D", "CONV_3D", "CONV_2D_DW", @@ -1096,7 +1097,7 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "GLU", }; -static_assert(GGML_OP_COUNT == 96, "GGML_OP_COUNT != 96"); +static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT != 97"); static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "none", @@ -1157,6 +1158,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "im2col(x)", "im2col_back(x)", "im2col_3d(x)", + "col2im_1d(x)", "conv_2d(x)", "conv_3d(x)", "conv_2d_dw(x)", @@ -1206,7 +1208,7 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "glu(x)", }; -static_assert(GGML_OP_COUNT == 96, "GGML_OP_COUNT != 96"); +static_assert(GGML_OP_COUNT == 97, "GGML_OP_COUNT != 97"); static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); @@ -4557,6 +4559,41 @@ struct ggml_tensor * ggml_conv_1d_dw_ph( return ggml_conv_1d_dw(ctx, a, b, s0, a->ne[0] / 2, d0); } +// ggml_col2im_1d + +struct ggml_tensor * ggml_col2im_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int s0, + int oc, + int p0) { + GGML_ASSERT(ggml_is_matrix(a)); + GGML_ASSERT(ggml_is_contiguous(a)); + GGML_ASSERT(a->type == GGML_TYPE_F32 || a->type == GGML_TYPE_F16 || a->type == GGML_TYPE_BF16); + GGML_ASSERT(s0 > 0); + GGML_ASSERT(oc > 0); + GGML_ASSERT(p0 >= 0); + + const int64_t K_OC = a->ne[0]; + const int64_t T_in = a->ne[1]; + const int64_t K = K_OC / oc; + const int64_t T_out = (T_in - 1) * s0 + K - 2 * p0; + + GGML_ASSERT(K_OC == K * oc); // a->ne[0] must be a whole number of oc blocks + GGML_ASSERT(K > 0 && T_out > 0); + + const int64_t ne[4] = { T_out, oc, 1, 1 }; + struct ggml_tensor * result = ggml_new_tensor(ctx, a->type, 2, ne); + + int32_t params[] = { s0, (int32_t)oc, (int32_t)p0 }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_COL2IM_1D; + result->src[0] = a; + + return result; +} + // ggml_conv_transpose_1d static int64_t ggml_calc_conv_transpose_1d_output_size(int64_t ins, int64_t ks, int s, int p, int d) { diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index bd6246137..584594097 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -538,6 +538,8 @@ class VISION_PROJECTOR_TYPE(IntEnum): class MODEL_TENSOR(IntEnum): TOKEN_EMBD = auto() TOKEN_EMBD_NORM = auto() + MASKED_EMBD_CENTROIDS= auto() + MASKED_EMBD_ORDERING = auto() TOKEN_TYPES = auto() POS_EMBD = auto() OUTPUT = auto() @@ -1087,6 +1089,8 @@ TENSOR_NAMES: dict[MODEL_TENSOR, str] = { MODEL_TENSOR.TOKEN_EMBD: "token_embd", MODEL_TENSOR.TOKEN_EMBD_NORM: "token_embd_norm", MODEL_TENSOR.TOKEN_TYPES: "token_types", + MODEL_TENSOR.MASKED_EMBD_CENTROIDS: "masked_embd_centroids", + MODEL_TENSOR.MASKED_EMBD_ORDERING: "masked_embd_ordering", MODEL_TENSOR.POS_EMBD: "position_embd", MODEL_TENSOR.OUTPUT_NORM: "output_norm", MODEL_TENSOR.OUTPUT: "output", @@ -2586,6 +2590,8 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_ARCH.GEMMA4_ASSISTANT: [ MODEL_TENSOR.ROPE_FREQS, MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.MASKED_EMBD_CENTROIDS, + MODEL_TENSOR.MASKED_EMBD_ORDERING, MODEL_TENSOR.OUTPUT_NORM, MODEL_TENSOR.NEXTN_PROJ_PRE, MODEL_TENSOR.NEXTN_PROJ_POST, diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index a9537983d..5f1e28818 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -37,6 +37,14 @@ class TensorNameMap: "model.embed", # talkie ), + # Masked embeddings + MODEL_TENSOR.MASKED_EMBD_CENTROIDS: ( + "masked_embedding.centroids", # gemma-4 E2B/E4B assistants + ), + MODEL_TENSOR.MASKED_EMBD_ORDERING: ( + "masked_embedding.token_ordering", # gemma-4 E2B/E4B assistants + ), + # Token type embeddings MODEL_TENSOR.TOKEN_TYPES: ( "embeddings.token_type_embeddings", # bert nomic-bert diff --git a/gpttype_adapter.cpp b/gpttype_adapter.cpp index cfd616f81..ab15d7ac1 100644 --- a/gpttype_adapter.cpp +++ b/gpttype_adapter.cpp @@ -4491,7 +4491,7 @@ static void PrepareMediaEmbds(const int nctx, const std::vector & media_int std::string media_obj = media_objects[i].b64data; const std::vector media_data_buffer = kcpp_base64_decode(media_obj); mtmd::bitmap bitmap(media_objects[i].is_audio - ? mtmd_helper_bitmap_init_from_buf(mtmd_ctx, media_data_buffer.data(), media_data_buffer.size(),false) + ? mtmd_helper_bitmap_init_from_buf(mtmd_ctx, media_data_buffer.data(), media_data_buffer.size(),false).bitmap : kcpp_mtmd_bitmap_init_image_from_buf(media_data_buffer.data(), media_data_buffer.size(), vision_max_res)); if(!bitmap.ptr) { diff --git a/koboldcpp.py b/koboldcpp.py index 0cd4f3492..a10fe2f39 100644 --- a/koboldcpp.py +++ b/koboldcpp.py @@ -5077,8 +5077,7 @@ class KcppServerRequestHandler(http.server.SimpleHTTPRequestHandler): tc["function"]["arguments"] = json.dumps(tcarg) recvtxt = None currfinishreason = "tool_calls" - if args.debugmode >= 1: - print(f"\nDebug ToolCall Response: {json.dumps(tool_calls)}") + utfprint(f"\nExecute Toolcall: {json.dumps(tool_calls)}",1) modelNameToReturn = friendlymodelname if autoswapmode and textName is not None: diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 6a5d5f8d2..680b5fc64 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -559,6 +559,8 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_INDEXER_PROJ, "blk.%d.indexer.proj" }, { LLM_TENSOR_INDEXER_ATTN_K, "blk.%d.indexer.attn_k" }, { LLM_TENSOR_INDEXER_ATTN_Q_B, "blk.%d.indexer.attn_q_b" }, + { LLM_TENSOR_MASKED_EMBD_CENTROIDS, "masked_embd_centroids" }, + { LLM_TENSOR_MASKED_EMBD_ORDERING, "masked_embd_ordering" }, }; // declare information about the model weight tensors: @@ -783,6 +785,8 @@ static const std::map LLM_TENSOR_INFOS = { // latent projections feed ggml_mul_mat, the buft probe must use MUL_MAT to keep them on GPU {LLM_TENSOR_FFN_LATENT_DOWN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_FFN_LATENT_UP, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_MASKED_EMBD_CENTROIDS, {LLM_TENSOR_LAYER_INPUT, GGML_OP_NONE}}, + {LLM_TENSOR_MASKED_EMBD_ORDERING, {LLM_TENSOR_LAYER_INPUT, GGML_OP_NONE}}, }; LLM_KV::LLM_KV(llm_arch arch, const char * suffix) : arch(arch), suffix(suffix) {} diff --git a/src/llama-arch.h b/src/llama-arch.h index 03b1a265d..b65fce72e 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -566,8 +566,11 @@ enum llm_tensor { LLM_TENSOR_NEXTN_HNORM, LLM_TENSOR_NEXTN_SHARED_HEAD_HEAD, LLM_TENSOR_NEXTN_SHARED_HEAD_NORM, + LLM_TENSOR_MASKED_EMBD_CENTROIDS, + LLM_TENSOR_MASKED_EMBD_ORDERING, }; + enum llm_tensor_layer { LLM_TENSOR_LAYER_INPUT, LLM_TENSOR_LAYER_REPEATING, diff --git a/src/llama-graph.cpp b/src/llama-graph.cpp index 66ea59a95..a37483aa6 100644 --- a/src/llama-graph.cpp +++ b/src/llama-graph.cpp @@ -567,7 +567,10 @@ void llm_graph_input_attn_kv_iswa::set_input(const llama_ubatch * ubatch) { mctx->get_base()->set_input_v_idxs(self_v_idxs, ubatch); } - mctx->get_base()->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn); + // the kq mask guards on its own buffer: shared cells leave idxs unbacked while the mask stays live + if (self_kq_mask && self_kq_mask->buffer) { + mctx->get_base()->set_input_kq_mask(self_kq_mask, ubatch, cparams.causal_attn); + } // swa tensors may not be allocated if there are no SWA attention layers if (self_k_idxs_swa && self_k_idxs_swa->buffer) { @@ -575,7 +578,9 @@ void llm_graph_input_attn_kv_iswa::set_input(const llama_ubatch * ubatch) { mctx->get_swa()->set_input_v_idxs(self_v_idxs_swa, ubatch); } - mctx->get_swa()->set_input_kq_mask(self_kq_mask_swa, ubatch, cparams.causal_attn); + if (self_kq_mask_swa && self_kq_mask_swa->buffer) { + mctx->get_swa()->set_input_kq_mask(self_kq_mask_swa, ubatch, cparams.causal_attn); + } if (self_k_rot) { mctx->get_base()->set_input_k_rot(self_k_rot); @@ -607,7 +612,9 @@ bool llm_graph_input_attn_kv_iswa::can_reuse(const llm_graph_params & params) { //res &= self_v_idxs->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there } - res &= can_reuse_kq_mask(self_kq_mask, mctx->get_base(), params.ubatch, params.cparams); + if (self_kq_mask && self_kq_mask->buffer) { + res &= can_reuse_kq_mask(self_kq_mask, mctx->get_base(), params.ubatch, params.cparams); + } // swa tensors may not be allocated if there are no SWA attention layers if (self_k_idxs_swa && self_k_idxs_swa->buffer) { @@ -615,7 +622,9 @@ bool llm_graph_input_attn_kv_iswa::can_reuse(const llm_graph_params & params) { //res &= self_v_idxs_swa->ne[0] == params.ubatch.n_tokens; // TODO: need to move this to the unified cache and check there } - res &= can_reuse_kq_mask(self_kq_mask_swa, mctx->get_swa(), params.ubatch, params.cparams); + if (self_kq_mask_swa && self_kq_mask_swa->buffer) { + res &= can_reuse_kq_mask(self_kq_mask_swa, mctx->get_swa(), params.ubatch, params.cparams); + } return res; } diff --git a/src/models/gemma4-assistant.cpp b/src/models/gemma4-assistant.cpp index 5b7a25a5a..6378130e7 100644 --- a/src/models/gemma4-assistant.cpp +++ b/src/models/gemma4-assistant.cpp @@ -39,6 +39,9 @@ void llama_model_gemma4_assistant::load_arch_tensors(llama_model_loader &) { output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), { n_embd }, 0); + create_tensor(tn(LLM_TENSOR_MASKED_EMBD_CENTROIDS, "weight"), {}, TENSOR_NOT_REQUIRED); + create_tensor(tn(LLM_TENSOR_MASKED_EMBD_ORDERING), {}, TENSOR_NOT_REQUIRED); + const int64_t n_embd_backbone = hparams.n_embd_inp(); nextn_proj_post = create_tensor(tn(LLM_TENSOR_NEXTN_PROJ_POST, "weight"), { n_embd, n_embd_backbone }, 0); diff --git a/src/models/plamo2.cpp b/src/models/plamo2.cpp index b93cf48bc..0b81513c3 100644 --- a/src/models/plamo2.cpp +++ b/src/models/plamo2.cpp @@ -11,6 +11,10 @@ void llama_model_plamo2::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_SSM_TIME_STEP_RANK, hparams.ssm_dt_rank); ml.get_key(LLM_KV_SSM_GROUP_COUNT, hparams.ssm_n_group); + // Load attention parameters + ml.get_key(LLM_KV_ATTENTION_KEY_LENGTH, hparams.n_embd_head_k_full, false); + ml.get_key(LLM_KV_ATTENTION_VALUE_LENGTH, hparams.n_embd_head_v_full, false); + for (uint32_t i = 0; i < hparams.n_layer(); ++i) { hparams.is_recr_impl[i] = hparams.n_head_kv(i) == 0; } @@ -273,7 +277,7 @@ ggml_tensor * llama_model_plamo2::graph::build_plamo2_mamba_layer(llm_graph_inpu GGML_ASSERT(n_seqs != 0); GGML_ASSERT(ubatch.equal_seqs()); GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs); - GGML_ASSERT(d_inner % n_head == 0); + GGML_ASSERT(d_inner % n_heads == 0); GGML_ASSERT(n_group == 0); ggml_tensor * conv_states_all = mctx_cur->get_r_l(il); diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index d7c6014e4..a6480f2e9 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -362,11 +362,17 @@ ggml_tensor * clip_graph::build_vit( std::function add_pos, const build_vit_opts & opts ) { + // batch dim: inp is [n_embd, n_pos] (B==1) or [n_embd, n_pos, B] (multi-tile encode) + const int64_t B = inp->ne[2]; + if (learned_pos_embd) { inp = ggml_add(ctx0, inp, learned_pos_embd); cb(inp, "pos_embed", -1); } + // flatten batch; unflatten again in attention + inp = ggml_reshape_2d(ctx0, inp, n_embd, n_pos * B); + ggml_tensor * inpL = inp; // pre-layernorm @@ -396,20 +402,24 @@ ggml_tensor * clip_graph::build_vit( cur = ggml_add(ctx0, cur, layer.qkv_b); } - Qcur = ggml_view_3d(ctx0, cur, d_head, n_head, n_pos, - /* nb1 */ ggml_row_size(cur->type, d_head), - /* nb2 */ cur->nb[1], - /* offset */ 0); + // Q/K/V as [d_head, n_head, n_pos, B], the batch stride is cur->nb[1]*n_pos. + Qcur = ggml_view_4d(ctx0, cur, d_head, n_head, n_pos, B, + /* nb1 */ ggml_row_size(cur->type, d_head), + /* nb2 */ cur->nb[1], + /* nb3 */ cur->nb[1] * n_pos, + /* offset */ 0); - Kcur = ggml_view_3d(ctx0, cur, d_head, n_head, n_pos, - /* nb1 */ ggml_row_size(cur->type, d_head), - /* nb2 */ cur->nb[1], - /* offset */ ggml_row_size(cur->type, n_embd)); + Kcur = ggml_view_4d(ctx0, cur, d_head, n_head, n_pos, B, + /* nb1 */ ggml_row_size(cur->type, d_head), + /* nb2 */ cur->nb[1], + /* nb3 */ cur->nb[1] * n_pos, + /* offset */ ggml_row_size(cur->type, n_embd)); - Vcur = ggml_view_3d(ctx0, cur, d_head, n_head, n_pos, - /* nb1 */ ggml_row_size(cur->type, d_head), - /* nb2 */ cur->nb[1], - /* offset */ ggml_row_size(cur->type, 2 * n_embd)); + Vcur = ggml_view_4d(ctx0, cur, d_head, n_head, n_pos, B, + /* nb1 */ ggml_row_size(cur->type, d_head), + /* nb2 */ cur->nb[1], + /* nb3 */ cur->nb[1] * n_pos, + /* offset */ ggml_row_size(cur->type, 2 * n_embd)); if (layer.q_norm) { GGML_ASSERT(layer.q_norm->ne[0] == Qcur->ne[0]); @@ -454,9 +464,9 @@ ggml_tensor * clip_graph::build_vit( } } - Qcur = ggml_reshape_3d(ctx0, Qcur, d_head, n_head, n_pos); - Kcur = ggml_reshape_3d(ctx0, Kcur, d_head, n_head_kv, n_pos); - Vcur = ggml_reshape_3d(ctx0, Vcur, d_head, n_head_kv, n_pos); + Qcur = ggml_reshape_4d(ctx0, Qcur, d_head, n_head, n_pos, B); + Kcur = ggml_reshape_4d(ctx0, Kcur, d_head, n_head_kv, n_pos, B); + Vcur = ggml_reshape_4d(ctx0, Vcur, d_head, n_head_kv, n_pos, B); if (norm_per_head) { if (layer.q_norm) { @@ -486,6 +496,7 @@ ggml_tensor * clip_graph::build_vit( cb(Vcur, "Vcur_normed", il); } + // build_attn returns a flat 2D [n_embd, n_pos*B] cur = build_attn(layer.o_w, layer.o_b, Qcur, Kcur, Vcur, opts.attn_mask, kq_scale, il); cb(cur, "attn_out", il); @@ -557,6 +568,10 @@ ggml_tensor * clip_graph::build_vit( if (model.post_ln_w) { inpL = build_norm(inpL, model.post_ln_w, model.post_ln_b, norm_t, eps, -1); } + + // restore the batch dim + GGML_ASSERT(inpL->ne[1] % B == 0); + inpL = ggml_reshape_3d(ctx0, inpL, n_embd, inpL->ne[1] / B, B); return inpL; } diff --git a/tools/mtmd/mtmd-cli.cpp b/tools/mtmd/mtmd-cli.cpp index bd7f9871c..a3cad7cd0 100644 --- a/tools/mtmd/mtmd-cli.cpp +++ b/tools/mtmd/mtmd-cli.cpp @@ -77,6 +77,7 @@ struct mtmd_cli_context { int n_batch; mtmd::bitmaps bitmaps; + std::vector videos; // chat template common_chat_templates_ptr tmpls; @@ -166,11 +167,14 @@ struct mtmd_cli_context { } bool load_media(const std::string & fname) { - mtmd::bitmap bmp(mtmd_helper_bitmap_init_from_file(ctx_vision.get(), fname.c_str(), false)); - if (!bmp.ptr) { + auto res = mtmd_helper_bitmap_init_from_file(ctx_vision.get(), fname.c_str(), false); + if (!res.bitmap) { return false; } - bitmaps.entries.push_back(std::move(bmp)); + bitmaps.entries.emplace_back(res.bitmap); + if (res.video_ctx) { + videos.emplace_back(res.video_ctx); + } return true; } }; @@ -253,6 +257,7 @@ static int eval_message(mtmd_cli_context & ctx, common_chat_msg & msg) { } ctx.bitmaps.entries.clear(); + ctx.videos.clear(); llama_pos new_n_past; if (mtmd_helper_eval_chunks(ctx.ctx_vision.get(), @@ -373,6 +378,9 @@ int main(int argc, char ** argv) { if (mtmd_support_audio(ctx.ctx_vision.get())) { LOG("\n /audio load an audio"); } + if (mtmd_helper_support_video(ctx.ctx_vision.get())) { + LOG("\n /video load a video"); + } LOG("\n /clear clear the chat history"); LOG("\n /quit or /exit exit the program"); LOG("\n"); @@ -407,14 +415,15 @@ int main(int argc, char ** argv) { g_is_generating = true; bool is_image = line == "/image" || line.find("/image ") == 0; bool is_audio = line == "/audio" || line.find("/audio ") == 0; - if (is_image || is_audio) { + bool is_video = line == "/video" || line.find("/video ") == 0; + if (is_image || is_audio || is_video) { if (line.size() < 8) { LOG_ERR("ERR: Missing media filename\n"); continue; } std::string media_path = line.substr(7); if (ctx.load_media(media_path)) { - LOG("%s %s loaded\n", media_path.c_str(), is_image ? "image" : "audio"); + LOG("%s %s loaded\n", media_path.c_str(), is_image ? "image" : is_audio ? "audio" : "video"); content += mtmd_default_marker(); } // else, error is already printed by libmtmd diff --git a/tools/mtmd/mtmd-helper.cpp b/tools/mtmd/mtmd-helper.cpp index 1c4d93c70..31cbdbf17 100644 --- a/tools/mtmd/mtmd-helper.cpp +++ b/tools/mtmd/mtmd-helper.cpp @@ -37,6 +37,11 @@ #error "mtmd-helper is a public library outside of mtmd. it must not include internal headers" #endif +#ifdef MTMD_VIDEO +#include "sheredom/subprocess.h" +#include +#endif + // // internal logging functions // @@ -80,6 +85,7 @@ struct mtmd_helper_logger { } } g_logger; +#define LOG_DBG(...) g_logger.log(GGML_LOG_LEVEL_DEBUG, __VA_ARGS__) #define LOG_INF(...) g_logger.log(GGML_LOG_LEVEL_INFO, __VA_ARGS__) #define LOG_WRN(...) g_logger.log(GGML_LOG_LEVEL_WARN, __VA_ARGS__) #define LOG_ERR(...) g_logger.log(GGML_LOG_LEVEL_ERROR, __VA_ARGS__) @@ -479,42 +485,94 @@ static bool decode_audio_from_buf(const unsigned char * buf_in, size_t len, int } // namespace audio_helpers -mtmd_bitmap * mtmd_helper_bitmap_init_from_buf(mtmd_context * ctx, const unsigned char * buf, size_t len, bool placeholder) { +// Computes FNV-1a hash of the data +static std::string fnv_hash(const uint8_t * data, size_t len) { + const uint64_t fnv_prime = 0x100000001b3ULL; + uint64_t hash = 0xcbf29ce484222325ULL; + + for (size_t i = 0; i < len; ++i) { + hash ^= data[i]; + hash *= fnv_prime; + } + return std::to_string(hash); +} + +mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_buf(mtmd_context * ctx, const unsigned char * buf, size_t len, bool placeholder) { + // calculate the hash if needed + std::string id; + mtmd_bitmap * result = nullptr; + + if (!placeholder) { + id = fnv_hash(buf, len); + } + if (audio_helpers::is_audio_file((const char *)buf, len)) { std::vector pcmf32; const int sample_rate = mtmd_get_audio_sample_rate(ctx); if (sample_rate < 0) { LOG_ERR("This model does not support audio input\n"); - return nullptr; + return {nullptr, nullptr}; } if (!audio_helpers::decode_audio_from_buf(buf, len, sample_rate, pcmf32)) { LOG_ERR("Unable to read WAV audio file from buffer\n"); - return nullptr; + return {nullptr, nullptr}; } - return mtmd_bitmap_init_from_audio(pcmf32.size(), placeholder ? nullptr : pcmf32.data()); + result = mtmd_bitmap_init_from_audio(pcmf32.size(), placeholder ? nullptr : pcmf32.data()); + mtmd_bitmap_set_id(result, id.empty() ? nullptr : id.c_str()); + return {result, nullptr}; } // otherwise, we assume it's an image - mtmd_bitmap * result = nullptr; - { + if (!result) { int nx, ny, nc; auto * data = stbi_load_from_memory(buf, len, &nx, &ny, &nc, 3); - if (!data) { - LOG_ERR("%s: failed to decode image bytes\n", __func__); - return nullptr; + if (data) { + result = mtmd_bitmap_init(nx, ny, placeholder ? nullptr : data); + mtmd_bitmap_set_id(result, id.empty() ? nullptr : id.c_str()); + stbi_image_free(data); + return {result, nullptr}; } - result = mtmd_bitmap_init(nx, ny, placeholder ? nullptr : data); - stbi_image_free(data); + // otherwise, fallthrough to video decoding (if supported) } - return result; + + // last try: load as video +#ifdef MTMD_VIDEO + if (!result) { + auto params = mtmd_helper_video_init_params_default(); + auto video_ctx = mtmd_helper_video_init_from_buf(ctx, buf, len, params); + if (!video_ctx) { + LOG_ERR("%s: failed to decode buffer as either image/audio/video\n", __func__); + return {nullptr, nullptr}; + } + result = mtmd_bitmap_init_lazy(ctx, + id.empty() ? nullptr : id.c_str(), + video_ctx, + [](size_t, void * user_data, mtmd_bitmap ** out_bitmap, char ** out_text) -> int { + auto * vctx = static_cast(user_data); + char * text = nullptr; + int ret = mtmd_helper_video_read_next(vctx, out_bitmap, &text); + *out_text = text; // heap-allocated by read_next; freed automatically by mtmd + return ret; + }); + return {result, video_ctx}; + } +#else + if (!result) { + LOG_ERR("%s: failed to decode buffer as either image or audio (video support not compiled in)\n", __func__); + return {nullptr, nullptr}; + } +#endif + + // should not reach here + return {nullptr, nullptr}; } -mtmd_bitmap * mtmd_helper_bitmap_init_from_file(mtmd_context * ctx, const char * fname, bool placeholder) { +mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_file(mtmd_context * ctx, const char * fname, bool placeholder) { std::vector buf; FILE * f = fopen(fname, "rb"); if (!f) { LOG_ERR("Unable to open file %s: %s\n", fname, strerror(errno)); - return nullptr; + return {nullptr, nullptr}; } fseek(f, 0, SEEK_END); @@ -523,7 +581,7 @@ mtmd_bitmap * mtmd_helper_bitmap_init_from_file(mtmd_context * ctx, const char * if (file_size < 0) { LOG_ERR("Failed to get file size of %s\n", fname); fclose(f); - return nullptr; + return {nullptr, nullptr}; } buf.resize(file_size); @@ -531,9 +589,444 @@ mtmd_bitmap * mtmd_helper_bitmap_init_from_file(mtmd_context * ctx, const char * fclose(f); if (n_read != (size_t)file_size) { LOG_ERR("Failed to read entire file %s", fname); - return nullptr; + return {nullptr, nullptr}; } return mtmd_helper_bitmap_init_from_buf(ctx, buf.data(), buf.size(), placeholder); } +bool mtmd_helper_support_video(mtmd_context * ctx) { +#ifdef MTMD_VIDEO + return mtmd_support_vision(ctx); +#else + return false; +#endif +} + +// +// Video input helpers +// + +#ifdef MTMD_VIDEO + +struct mtmd_helper_video { + mtmd_context * mctx; + std::string path; + std::vector input_buf; // non-empty when initialized from buffer + std::string ffmpeg_bin; + std::string ffprobe_bin; + float fps_target = 0.0f; + mtmd_helper_video_info info = {}; + + // RAII wrapper for managing subprocess + struct subprocess_handle { + struct subprocess_s proc = {}; + bool alive = false; + std::thread feeder; + + subprocess_handle() = default; + subprocess_handle(const subprocess_handle &) = delete; + subprocess_handle & operator=(const subprocess_handle &) = delete; + ~subprocess_handle() { stop(); } + + void stop() { + if (alive) { + subprocess_terminate(&proc); + } + // join before destroy: feeder holds a FILE* from subprocess_stdin; + // subprocess_destroy closes it, so the thread must finish first + if (feeder.joinable()) { + feeder.join(); + } + if (alive) { + subprocess_destroy(&proc); + alive = false; + } + } + + FILE * stdout_pipe() { + return subprocess_stdout(&proc); + } + + // buf is tied to lifetime of mtmd_helper_video, so it's guaranteed to outlive the feeder thread + void start_feeder(const std::vector & buf) { + feeder = std::thread([this, &buf]() { + FILE * f = subprocess_stdin(&proc); + if (!f) { + return; + } + fwrite(buf.data(), 1, buf.size(), f); + fclose(f); + proc.stdin_file = nullptr; // prevent double-close in subprocess_destroy + }); + } + }; + + subprocess_handle sp; + int32_t current_frame = 0; + + std::string prompt_start = "Video:"; + int32_t timestamp_interval_ms = 5000; // emit a timestamp text every N ms (0 = disabled) + float next_timestamp_ms = 0.0f; // next elapsed-ms threshold at which to emit + + std::vector frame_buf; + std::string pending_text; // text queued to be returned before the next frame + bool start_emitted = false; + + bool is_buf_input() const { + return !input_buf.empty(); + } + + bool probe(float fps_target_arg) { + const char * input_arg = is_buf_input() ? "pipe:0" : path.c_str(); + const char * cmd[] = { + ffprobe_bin.c_str(), + "-v", "quiet", + "-show_entries", "stream=width,height,r_frame_rate,nb_frames,duration", + "-select_streams", "v:0", + "-of", "default=noprint_wrappers=1", + input_arg, + nullptr, + }; + + LOG_DBG("%s: launching:", __func__); + for (size_t i = 0; cmd[i]; i++) { LOG_DBG(" %s", cmd[i]); } + LOG_DBG("\n"); + + subprocess_handle probe_sp; + if (subprocess_create(cmd, + subprocess_option_search_user_path | subprocess_option_inherit_environment, + &probe_sp.proc) != 0) { + LOG_ERR("%s: failed to launch ffprobe\n", __func__); + return false; + } + probe_sp.alive = true; + + if (is_buf_input()) { + probe_sp.start_feeder(input_buf); + } + + uint32_t width = 0; + uint32_t height = 0; + float orig_fps = 0.0f; + float duration = -1.0f; + int32_t n_frames_orig = -1; + char line[256]; + FILE * fp = probe_sp.stdout_pipe(); + + while (fgets(line, sizeof(line), fp)) { + char * eq = strchr(line, '='); + if (!eq) continue; + *eq = '\0'; + const char * key = line; + const char * val = eq + 1; + char * nl = (char *)strchr(val, '\n'); + if (nl) *nl = '\0'; + + if (strcmp(key, "width") == 0) { + width = (uint32_t)atoi(val); + } else if (strcmp(key, "height") == 0) { + height = (uint32_t)atoi(val); + } else if (strcmp(key, "r_frame_rate") == 0) { + orig_fps = parse_rational(val); + } else if (strcmp(key, "nb_frames") == 0 && strcmp(val, "N/A") != 0) { + n_frames_orig = atoi(val); + } else if (strcmp(key, "duration") == 0 && strcmp(val, "N/A") != 0) { + duration = (float)atof(val); + } + } + + probe_sp.stop(); + + if (width == 0 || height == 0 || orig_fps <= 0.0f) { + return false; + } + + if (duration < 0.0f && n_frames_orig > 0) { + duration = (float)n_frames_orig / orig_fps; + } + + fps_target = fps_target_arg > 0.0f ? fps_target_arg : orig_fps; + info.width = width; + info.height = height; + info.fps = fps_target; + LOG_DBG("%s: %ux%u fps=%.2f duration=%.2fs n_frames=%d\n", + __func__, width, height, fps_target, duration, info.n_frames); + info.n_frames = duration > 0.0f ? (int32_t)(duration * fps_target + 0.5f) : -1; + frame_buf.resize((size_t)width * height * 3); + return true; + } + + bool start_ffmpeg(float seek_seconds) { + char seek_buf[64]; + char fps_buf[64]; + + std::vector cmd; + cmd.push_back(ffmpeg_bin.c_str()); + + if (!is_buf_input() && seek_seconds > 0.0f) { + // input-side seek: fast, keyframe-accurate; only valid for seekable file inputs + snprintf(seek_buf, sizeof(seek_buf), "%.6f", seek_seconds); + cmd.push_back("-ss"); + cmd.push_back(seek_buf); + } + + cmd.push_back("-nostdin"); + cmd.push_back("-i"); + // cache:pipe:0 wraps stdin with a seekable in-memory cache, letting ffmpeg seek + // backwards for container headers (e.g. MP4 moov atom at end of file) + cmd.push_back(is_buf_input() ? "cache:pipe:0" : path.c_str()); + + if (seek_seconds > 0.0f && is_buf_input()) { + // output-side seek: frame-accurate but decodes and discards frames up to seek point + snprintf(seek_buf, sizeof(seek_buf), "%.6f", seek_seconds); + cmd.push_back("-ss"); + cmd.push_back(seek_buf); + } + + if (fps_target > 0.0f) { + snprintf(fps_buf, sizeof(fps_buf), "fps=%.6f", fps_target); + cmd.push_back("-vf"); + cmd.push_back(fps_buf); + } + + cmd.push_back("-f"); + cmd.push_back("rawvideo"); + cmd.push_back("-pix_fmt"); + cmd.push_back("rgb24"); + cmd.push_back("pipe:1"); + cmd.push_back("-loglevel"); + cmd.push_back("error"); + cmd.push_back(nullptr); + + LOG_DBG("%s: launching:", __func__); + for (size_t i = 0; cmd[i]; i++) { + LOG_DBG(" %s", cmd[i]); + } + LOG_DBG("\n"); + + int ret = subprocess_create( + cmd.data(), + subprocess_option_search_user_path | subprocess_option_inherit_environment, + &sp.proc); + + sp.alive = (ret == 0); + LOG_DBG("%s: subprocess_create ret=%d proc_alive=%d\n", __func__, ret, (int)sp.alive); + + if (sp.alive && is_buf_input()) { + LOG_DBG("%s: starting feeder thread for %zu-byte buffer\n", __func__, input_buf.size()); + sp.start_feeder(input_buf); + } + + return sp.alive; + } + + void stop_ffmpeg() { + sp.stop(); + } + + mtmd_bitmap * read_next_frame() { + if (!sp.alive) return nullptr; + + FILE * fp = sp.stdout_pipe(); + const size_t frame_size = (size_t)info.width * info.height * 3; + LOG_DBG("%s: reading frame %d, expecting %zu bytes (%ux%u)\n", + __func__, current_frame, frame_size, info.width, info.height); + + size_t total_read = 0; + while (total_read < frame_size) { + size_t n = fread(frame_buf.data() + total_read, 1, frame_size - total_read, fp); + if (n == 0) { + // clean EOF only if no bytes read yet; partial frame is an error + LOG_DBG("%s: fread returned 0 after %zu/%zu bytes (ferror=%d)\n", + __func__, total_read, frame_size, ferror(fp)); + sp.alive = false; + return nullptr; + } + total_read += n; + } + + LOG_DBG("%s: frame %d read OK\n", __func__, current_frame); + current_frame++; + return mtmd_bitmap_init(info.width, info.height, frame_buf.data()); + } + + int32_t read_next(mtmd_bitmap ** out_bitmap, char ** out_text) { + *out_bitmap = nullptr; + *out_text = nullptr; + + if (!pending_text.empty()) { + *out_text = strdup(pending_text.c_str()); + pending_text.clear(); + return *out_text ? 0 : -2; + } + + LOG_DBG("%s: proc_alive=%d start_emitted=%d current_frame=%d\n", + __func__, (int)sp.alive, (int)start_emitted, current_frame); + + if (!sp.alive) { + return (current_frame == 0) ? -2 : -1; + } + + if (!start_emitted) { + start_emitted = true; + if (!prompt_start.empty()) { + *out_text = strdup(prompt_start.c_str()); + return *out_text ? 0 : -2; + } + } + + mtmd_bitmap * frame = read_next_frame(); + if (!frame) return -1; + *out_bitmap = frame; + + if (timestamp_interval_ms > 0) { + // current_frame was already incremented by read_next_frame(); undo for elapsed calc + float elapsed_ms = (float)(current_frame - 1) / info.fps * 1000.0f; + if (elapsed_ms >= next_timestamp_ms) { + char ts_buf[32]; + float elapsed_s = elapsed_ms / 1000.0f; + int minutes = (int)(elapsed_s / 60); + float seconds = elapsed_s - minutes * 60.0f; + snprintf(ts_buf, sizeof(ts_buf), "[%dm%.2fs]", minutes, seconds); + pending_text = ts_buf; + next_timestamp_ms += (float)timestamp_interval_ms; + } + } + + return 0; + } + + static float parse_rational(const char * s) { + int num = 0, den = 1; + if (sscanf(s, "%d/%d", &num, &den) == 2 && den > 0) { + return (float)num / (float)den; + } + float val; + if (sscanf(s, "%f", &val) == 1) { + return val; + } + return 0.0f; + } +}; +#endif + +mtmd_helper_video_init_params mtmd_helper_video_init_params_default() { + return { + /* fps_target */ 4.0f, + /* ffmpeg_bin_dir */ nullptr, + /* timestamp_interval_ms */ 5000, + }; +} + +static std::string video_resolve_bin(const char * bin_dir, const char * name) { + if (!bin_dir || bin_dir[0] == '\0') { + return name; // rely on PATH + } + std::string result = bin_dir; + char last = result.back(); + if (last != '/' && last != '\\') { +#ifdef _WIN32 + result += '\\'; +#else + result += '/'; +#endif + } + result += name; +#ifdef _WIN32 + result += ".exe"; +#endif + return result; +} + +mtmd_helper_video * mtmd_helper_video_init( + mtmd_context * mctx, + const char * path, + mtmd_helper_video_init_params params) { +#ifdef MTMD_VIDEO + auto * ctx = new mtmd_helper_video(); + + ctx->mctx = mctx; + ctx->path = path; + ctx->ffmpeg_bin = video_resolve_bin(params.ffmpeg_bin_dir, "ffmpeg"); + ctx->ffprobe_bin = video_resolve_bin(params.ffmpeg_bin_dir, "ffprobe"); + ctx->timestamp_interval_ms = params.timestamp_interval_ms; + + if (!ctx->probe(params.fps_target)) { + LOG_ERR("%s: ffprobe failed for '%s' (is ffprobe in PATH?)\n", __func__, path); + delete ctx; + return nullptr; + } + + if (!ctx->start_ffmpeg(0.0f)) { + LOG_ERR("%s: failed to start ffmpeg for '%s' (is ffmpeg in PATH?)\n", __func__, path); + delete ctx; + return nullptr; + } + + return ctx; +#else + LOG_ERR("%s: video is not supported in this build (MTMD_VIDEO is set to OFF)\n", __func__); + return nullptr; +#endif +} + +mtmd_helper_video * mtmd_helper_video_init_from_buf( + mtmd_context * mctx, + const unsigned char * buf, size_t len, + mtmd_helper_video_init_params params) { +#ifdef MTMD_VIDEO + auto * ctx = new mtmd_helper_video(); + + ctx->mctx = mctx; + ctx->input_buf.assign(buf, buf + len); + ctx->ffmpeg_bin = video_resolve_bin(params.ffmpeg_bin_dir, "ffmpeg"); + ctx->ffprobe_bin = video_resolve_bin(params.ffmpeg_bin_dir, "ffprobe"); + ctx->timestamp_interval_ms = params.timestamp_interval_ms; + + if (!ctx->probe(params.fps_target)) { + LOG_ERR("%s: ffprobe failed on buffer (is ffprobe in PATH?)\n", __func__); + delete ctx; + return nullptr; + } + + if (!ctx->start_ffmpeg(0.0f)) { + LOG_ERR("%s: failed to start ffmpeg on buffer (is ffmpeg in PATH?)\n", __func__); + delete ctx; + return nullptr; + } + + return ctx; +#else + LOG_ERR("%s: video is not supported in this build (MTMD_VIDEO is set to OFF)\n", __func__); + return nullptr; +#endif +} + +void mtmd_helper_video_free(mtmd_helper_video * ctx) { +#ifdef MTMD_VIDEO + if (!ctx) return; + ctx->stop_ffmpeg(); + delete ctx; +#else + LOG_ERR("%s: video is not supported in this build (MTMD_VIDEO is set to OFF)\n", __func__); +#endif +} + +mtmd_helper_video_info mtmd_helper_video_get_info(const mtmd_helper_video * ctx) { +#ifdef MTMD_VIDEO + return ctx->info; +#else + GGML_ASSERT(false && "video is not supported in this build (MTMD_VIDEO is set to OFF)"); +#endif +} + +int32_t mtmd_helper_video_read_next(mtmd_helper_video * ctx, + mtmd_bitmap ** out_bitmap, char ** out_text) { +#ifdef MTMD_VIDEO + if (!ctx) return -2; + return ctx->read_next(out_bitmap, out_text); +#else + GGML_ASSERT(false && "video is not supported in this build (MTMD_VIDEO is set to OFF)"); +#endif +} diff --git a/tools/mtmd/mtmd-helper.h b/tools/mtmd/mtmd-helper.h index 7eecbb067..164b7c668 100644 --- a/tools/mtmd/mtmd-helper.h +++ b/tools/mtmd/mtmd-helper.h @@ -20,25 +20,39 @@ extern "C" { // BREAKING CHANGES are expected. // +struct mtmd_helper_video; +typedef struct mtmd_helper_video mtmd_helper_video; + // Set callback for all future logging events. // If this is not called, or NULL is supplied, everything is output on stderr. // Note: this also call mtmd_log_set() internally MTMD_API void mtmd_helper_log_set(ggml_log_callback log_callback, void * user_data); +// Returns true if this build includes video support (MTMD_VIDEO was ON at compile time). +MTMD_API bool mtmd_helper_support_video(mtmd_context * ctx); + +struct mtmd_helper_bitmap_wrapper { + mtmd_bitmap * bitmap; + mtmd_helper_video * video_ctx; +}; + // helper function to construct a mtmd_bitmap from a file // it calls mtmd_helper_bitmap_init_from_buf() internally // returns nullptr on failure // this function is thread-safe -MTMD_API mtmd_bitmap * mtmd_helper_bitmap_init_from_file(mtmd_context * ctx, const char * fname, bool placeholder); +MTMD_API struct mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_file(mtmd_context * ctx, const char * fname, bool placeholder); // helper function to construct a mtmd_bitmap from a buffer containing a file // supported formats: // image: formats supported by stb_image: jpg, png, bmp, gif, etc. // audio: formats supported by miniaudio: wav, mp3, flac -// note: audio files will be auto-detected based on magic bytes +// note: +// - for now, video input is only supported via C++ helper functions +// - audio files will be auto-detected based on magic bytes +// - output bitmap will have FNV hash as the ID // returns nullptr on failure // this function is thread-safe -MTMD_API mtmd_bitmap * mtmd_helper_bitmap_init_from_buf(mtmd_context * ctx, const unsigned char * buf, size_t len, bool placeholder); +MTMD_API struct mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_buf(mtmd_context * ctx, const unsigned char * buf, size_t len, bool placeholder); // helper to count the total number of tokens from a list of chunks, useful to keep track of KV cache MTMD_API size_t mtmd_helper_get_n_tokens(const mtmd_input_chunks * chunks); @@ -89,6 +103,56 @@ MTMD_API int32_t mtmd_helper_decode_image_chunk(mtmd_context * ctx, int32_t n_batch, llama_pos * new_n_past); +// +// video input helpers (requires ffmpeg/ffprobe installed on the system) +// the notion of video only exists at the helper level, it is not visible to the core mtmd library +// +// NOTE: this implementation is model-agnostic, it can be used with any vision-capable model +// however, it may not be accurate for some specific models +// (this is expected for now, to keep the implementation simple) +// + +struct mtmd_helper_video_info { + uint32_t width; + uint32_t height; + float fps; // effective fps (fps_target if set, else original video fps) + int32_t n_frames; // estimated total frames at effective fps (-1 if unknown) +}; + +struct mtmd_helper_video_init_params { + float fps_target; // desired output fps; <= 0 means use the video's native fps, defaulted to 4.0f + const char * ffmpeg_bin_dir; // directory containing ffmpeg/ffprobe binaries; NULL means search PATH + int64_t timestamp_interval_ms; // interval for adding timestamp as text chunk (example: "[10m50.5s]"); <= 0 means no timestamp, defaulted to 5000ms + // TODO @ngxson : allow "placeholder" bitmap output for counting tokens +}; + +MTMD_API struct mtmd_helper_video_init_params mtmd_helper_video_init_params_default(void); + +// returns NULL on failure (ffprobe not found, file unreadable, etc.) +MTMD_API mtmd_helper_video * mtmd_helper_video_init( + struct mtmd_context * mctx, + const char * path, + struct mtmd_helper_video_init_params params); + +// Same as mtmd_helper_video_init(), but reads from an in-memory buffer. +// The buffer is copied internally; the caller does not need to keep it alive. +// Note: pipe input is not seekable, so seeking will use output-side seeking +// (ffmpeg decodes and discards frames up to the target position). +MTMD_API mtmd_helper_video * mtmd_helper_video_init_from_buf( + struct mtmd_context * mctx, + const unsigned char * buf, size_t len, + struct mtmd_helper_video_init_params params); +MTMD_API void mtmd_helper_video_free(mtmd_helper_video * ctx); +MTMD_API struct mtmd_helper_video_info mtmd_helper_video_get_info(const mtmd_helper_video * ctx); + +// Read the next item from the video stream; exactly one of out_bitmap or out_text is set per call. +// *out_bitmap - heap-allocated; caller must free with mtmd_bitmap_free() +// *out_text - heap-allocated (always via strdup/malloc); caller must free with free() +// returns 0 on success, -1 on EOF, -2 on error +MTMD_API int32_t mtmd_helper_video_read_next(mtmd_helper_video * ctx, + mtmd_bitmap ** out_bitmap, + char ** out_text); + #ifdef __cplusplus } // extern "C" #endif @@ -97,4 +161,16 @@ MTMD_API int32_t mtmd_helper_decode_image_chunk(mtmd_context * ctx, // C++ wrappers // +#ifdef __cplusplus +namespace mtmd_helper { + +// video-related C++ wrappers +struct mtmd_helper_video_deleter { + void operator()(mtmd_helper_video * val) { mtmd_helper_video_free(val); } +}; +using video_ptr = std::unique_ptr; + +} // namespace mtmd_helper +#endif + #endif diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 39936c64a..f5deed557 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -35,6 +35,10 @@ struct mtmd_bitmap { std::string id; // optional user-defined id, for ex: can be set to image hash, useful for KV cache tracking bool is_audio = false; // true if the bitmap is audio + // lazy-loaded bitmap + mtmd_bitmap_lazy_callback lazy_callback = nullptr; + void * lazy_user_data = nullptr; + mtmd_bitmap(const unsigned char * data, uint32_t nx, uint32_t ny) : nx(nx), ny(ny), is_audio(false) { if (data) { @@ -732,30 +736,111 @@ void mtmd_free(mtmd_context * ctx) { struct mtmd_tokenizer { mtmd_context * ctx; - std::vector bitmaps; std::string input_text; bool add_special; bool parse_special; const llama_vocab * vocab; + struct part { + std::string text; + const mtmd_bitmap * bitmap; + }; + std::vector parts; + // these will be freed when mtmd_tokenizer finishes + std::vector bm_from_lazy; // TODO @ngxson : refactor, free bm_from_lazy progressively + std::vector text_from_lazy; + mtmd_input_chunks cur; uint32_t n_images_added = 0; // 0-based index assigned to the next image chunk + ~mtmd_tokenizer() { + // note: mtmd::bitmap is already RAII + for (auto & str : text_from_lazy) { + free((void *)str); + } + } + mtmd_tokenizer(mtmd_context * ctx, const mtmd_input_text * text, - const mtmd_bitmap ** bitmaps, - size_t n_bitmaps) : ctx(ctx), bitmaps(bitmaps, bitmaps + n_bitmaps) { + const mtmd_bitmap ** bmps, + size_t n_bitmaps) : ctx(ctx) { add_special = text->add_special; parse_special = text->parse_special; input_text = text->text; vocab = ctx->vocab; + + std::vector bitmaps(bmps, bmps + n_bitmaps); + auto parts_str = split_text(input_text, ctx->media_marker); + size_t i_bm = 0; + for (const auto & part : parts_str) { + if (part == ctx->media_marker) { + if (i_bm >= bitmaps.size()) { + throw std::runtime_error(string_format("number of media markers in text (%zu) exceeds number of bitmaps (%zu)", i_bm + 1, bitmaps.size())); + } + parts.push_back({"", bitmaps[i_bm++]}); + } else { + parts.push_back({std::move(part), nullptr}); + } + } + + size_t n_markers = 0; + for (const auto & part : parts) { + if (part.bitmap != nullptr) { + n_markers++; + } + } + if (n_markers != bitmaps.size()) { + throw std::runtime_error(string_format("number of media markers in text (%zu) does not match number of bitmaps (%zu)", n_markers, bitmaps.size())); + } + + expand_lazy_bitmaps(); + } + + void expand_lazy_bitmaps() { + std::vector expanded; + expanded.reserve(parts.size()); + for (auto & p : parts) { + if (p.bitmap != nullptr && p.bitmap->lazy_callback) { + LOG_DBG("%s: expanding lazy bitmap\n", __func__); + for (size_t i = 0;; i++) { + char * out_str = nullptr; + mtmd_bitmap * out_bm = nullptr; + int res = p.bitmap->lazy_callback(i, + p.bitmap->lazy_user_data, + &out_bm, + &out_str); + if (out_bm && out_str) { + throw std::runtime_error(string_format("lazy callback cannot return both bitmap and text")); + } + if (res == 0) { + // OK, append the returned chunk; lazy part is not yet added + if (out_bm) { + auto & ptr = bm_from_lazy.emplace_back(out_bm); // remember to free it later + expanded.push_back({"", ptr.ptr.get()}); + LOG_DBG("%s: lazy callback returned bitmap with dimensions %d x %d\n", __func__, out_bm->nx, out_bm->ny); + } else if (out_str) { + auto & ptr = text_from_lazy.emplace_back(out_str); // remember to free it later + expanded.push_back({ptr, nullptr}); + LOG_DBG("%s: lazy callback returned text: %s\n", __func__, out_str); + } + } else if (res == -1) { + // EOF: lazy part removes itself (not added to expanded) + break; + } else if (res == -2) { + // error + throw std::runtime_error(string_format("lazy callback returned error")); + } + } + } else { + expanded.push_back(std::move(p)); + } + } + parts = std::move(expanded); } int32_t tokenize(mtmd_input_chunks * output) { cur.entries.clear(); - std::vector parts = split_text(input_text, ctx->media_marker); - size_t i_bm = 0; // index of the current bitmap // [QWEN_VIDEO] handle frame merging for models that support it (i.e. qwen-vl) int n_merge_frames = 1; @@ -764,53 +849,50 @@ struct mtmd_tokenizer { GGML_ASSERT(n_merge_frames <= 2 && "we only support merging maximum 2 images for now; open an issue if this model supports merging more"); } + // Build merged_bitmaps: each entry is a group of 1 or 2 bitmaps. + // For consecutive mergeable bitmap parts, merge them and collapse the second part out of this->parts. std::vector> merged_bitmaps; if (n_merge_frames > 1) { - size_t i_bm_scan = 0; for (size_t i = 0; i < parts.size(); ++i) { - if (parts[i] != ctx->media_marker) { + if (parts[i].bitmap == nullptr) { continue; } - if (i + 1 < parts.size() - && parts[i + 1] == ctx->media_marker - && i_bm_scan + 1 < bitmaps.size()) { - const mtmd_bitmap * bm_a = bitmaps[i_bm_scan]; - const mtmd_bitmap * bm_b = bitmaps[i_bm_scan + 1]; + if (i + 1 < parts.size() && parts[i + 1].bitmap != nullptr) { + const mtmd_bitmap * bm_a = parts[i].bitmap; + const mtmd_bitmap * bm_b = parts[i + 1].bitmap; if (bm_a->can_batch_with(*bm_b)) { - LOG_DBG("%s: merging 2 frames at bitmap index %zu and %zu\n", __func__, i_bm_scan, i_bm_scan + 1); + LOG_DBG("%s: merging 2 frames at part index %zu and %zu\n", __func__, i, i + 1); merged_bitmaps.push_back({bm_a, bm_b}); - parts.erase(parts.begin() + i + 1); // remove the second marker - i_bm_scan += 2; + parts.erase(parts.begin() + i + 1); // collapse the second bitmap part continue; } } - LOG_DBG("%s: no merging for bitmap index %zu\n", __func__, i_bm_scan); - merged_bitmaps.push_back({bitmaps[i_bm_scan]}); - ++i_bm_scan; + LOG_DBG("%s: no merging for part index %zu\n", __func__, i); + merged_bitmaps.push_back({parts[i].bitmap}); } } else { - for (size_t i = 0; i < bitmaps.size(); ++i) { - merged_bitmaps.push_back({bitmaps[i]}); + for (const auto & p : parts) { + if (p.bitmap != nullptr) { + merged_bitmaps.push_back({p.bitmap}); + } } } - i_bm = 0; - for (auto & part : parts) { - if (part == ctx->media_marker) { - // this is a marker, we should add the next bitmap + size_t i_bm = 0; + for (const auto & p : parts) { + if (p.bitmap != nullptr) { if (i_bm >= merged_bitmaps.size()) { LOG_ERR("%s: error: number of bitmaps (%zu) does not match number of markers (%zu)\n", __func__, merged_bitmaps.size(), parts.size() - 1); return 1; } - auto & bmps = merged_bitmaps[i_bm++]; + auto bmps = merged_bitmaps[i_bm++]; int32_t res = add_media(bmps); if (res != 0) { return res; } } else { - // this is a text part, we should add it as text - add_text(part, parse_special); + add_text(p.text, parse_special); } } @@ -1236,8 +1318,13 @@ int32_t mtmd_tokenize(mtmd_context * ctx, const mtmd_input_text * text, const mtmd_bitmap ** bitmaps, size_t n_bitmaps) { - mtmd_tokenizer tokenizer(ctx, text, bitmaps, n_bitmaps); - return tokenizer.tokenize(output); + try { + mtmd_tokenizer tokenizer(ctx, text, bitmaps, n_bitmaps); + return tokenizer.tokenize(output); + } catch (const std::exception & e) { + LOG_ERR("%s: error: %s\n", __func__, e.what()); + return 2; + } } int32_t mtmd_encode_chunk(mtmd_context * ctx, const mtmd_input_chunk * chunk) { @@ -1373,6 +1460,10 @@ int mtmd_get_audio_sample_rate(const mtmd_context * ctx) { return clip_get_hparams(ctx->ctx_a)->audio_sample_rate; } +const char * mtmd_get_marker(const mtmd_context * ctx) { + return ctx->media_marker.c_str(); +} + // // public API functions // @@ -1405,10 +1496,16 @@ uint32_t mtmd_bitmap_get_ny(const mtmd_bitmap * bitmap) { } const unsigned char * mtmd_bitmap_get_data(const mtmd_bitmap * bitmap) { + if (bitmap->is_placeholder()) { + return nullptr; + } return bitmap->get_ro_buf().data(); } size_t mtmd_bitmap_get_n_bytes(const mtmd_bitmap * bitmap) { + if (bitmap->is_placeholder()) { + return 0; + } return bitmap->get_ro_buf().size(); } @@ -1428,6 +1525,18 @@ void mtmd_bitmap_set_id(mtmd_bitmap * bitmap, const char * id) { } } +mtmd_bitmap * mtmd_bitmap_init_lazy(mtmd_context * ctx, + const char * id, + void * user_data, + mtmd_bitmap_lazy_callback callback) { + GGML_UNUSED(ctx); // reserved for future use + mtmd_bitmap * bitmap = new mtmd_bitmap(nullptr, 0, 0); + bitmap->lazy_callback = callback; + bitmap->lazy_user_data = user_data; + mtmd_bitmap_set_id(bitmap, id); + return bitmap; +} + void mtmd_bitmap_free(mtmd_bitmap * bitmap) { if (bitmap) { delete bitmap; diff --git a/tools/mtmd/mtmd.h b/tools/mtmd/mtmd.h index a314b11f7..9a09f364f 100644 --- a/tools/mtmd/mtmd.h +++ b/tools/mtmd/mtmd.h @@ -128,6 +128,9 @@ MTMD_API bool mtmd_support_audio(const mtmd_context * ctx); // return -1 if audio is not supported MTMD_API int mtmd_get_audio_sample_rate(const mtmd_context * ctx); +// get the current marker string +MTMD_API const char * mtmd_get_marker(const mtmd_context * ctx); + // mtmd_bitmap // // if bitmap is image: @@ -156,6 +159,34 @@ MTMD_API void mtmd_bitmap_free (mtmd_bitmap * bitmap); MTMD_API const char * mtmd_bitmap_get_id(const mtmd_bitmap * bitmap); MTMD_API void mtmd_bitmap_set_id(mtmd_bitmap * bitmap, const char * id); +// mtmd_bitmap lazy +// +// this is a special bitmap that: +// - does not hold the actual data +// - can be expanded into one or more chunks (either media to text chunks) +// user must provide a callback to fill in the data when mtmd_tokenize() is called +// this is useful for large video inputs: +// - allow reading video frame by frame, without loading the entire video into memory +// - allow tracking the whole video with a single ID (for example, the file hash) + +// set (*out_bitmap) to non-nullptr to emit a bitmap chunk; it will be freed automatically +// set (*out_text) to non-nullptr to emit a text chunk; it must be heap-allocated, null-terminated and will be freed automatically +// either out_bitmap or out_text can be set, but not both +// out_bitmap cannot be another lazy bitmap (no nested lazy allowed) +// return value: +// 0 on success +// -1 on EOF (signal to mtmd_tokenize to move on) +// -2 on error (signal to mtmd_tokenize to abort) +typedef int(* mtmd_bitmap_lazy_callback)( + size_t chunk_idx, + void * user_data, + mtmd_bitmap ** out_bitmap, + char ** out_text); + +MTMD_API mtmd_bitmap * mtmd_bitmap_init_lazy(mtmd_context * ctx, + const char * id, // usually set to file hash + void * user_data, + mtmd_bitmap_lazy_callback callback); // mtmd_input_chunks // diff --git a/tools/server/server-common.cpp b/tools/server/server-common.cpp index dfd286d24..9f3caac8f 100644 --- a/tools/server/server-common.cpp +++ b/tools/server/server-common.cpp @@ -701,29 +701,19 @@ size_t validate_utf8(const std::string& text) { return len; } -// Computes FNV-1a hash of the data -static std::string fnv_hash(const uint8_t * data, size_t len) { - const uint64_t fnv_prime = 0x100000001b3ULL; - uint64_t hash = 0xcbf29ce484222325ULL; - - for (size_t i = 0; i < len; ++i) { - hash ^= data[i]; - hash *= fnv_prime; - } - return std::to_string(hash); -} - server_tokens process_mtmd_prompt(mtmd_context * mctx, const std::string & prompt, const std::vector & files, bool is_placeholder) { + // these will be freed upon going out of scope mtmd::bitmaps bitmaps; + std::vector videos; for (auto & file : files) { - mtmd::bitmap bmp(mtmd_helper_bitmap_init_from_buf(mctx, file.data(), file.size(), is_placeholder)); - if (!bmp.ptr) { + auto out = mtmd_helper_bitmap_init_from_buf(mctx, file.data(), file.size(), is_placeholder); + if (!out.bitmap) { throw std::runtime_error("Failed to load image or audio file"); } - // calculate bitmap hash (for KV caching) - std::string hash = fnv_hash(bmp.data(), bmp.n_bytes()); - bmp.set_id(hash.c_str()); - bitmaps.entries.push_back(std::move(bmp)); + bitmaps.entries.emplace_back(out.bitmap); + if (out.video_ctx) { + videos.emplace_back(out.video_ctx); + } } // process prompt std::vector inputs; @@ -1023,6 +1013,20 @@ json oaicompat_chat_params_parse( p["text"] = get_media_marker(); p.erase("input_audio"); + } else if (type == "input_video") { + if (!opt.allow_video) { + throw std::runtime_error("video input is not supported - hint: if this is unexpected, you may need to provide the mmproj"); + } + + json input_video = json_value(p, "input_video", json::object()); + std::string data = json_value(input_video, "data", std::string()); + auto decoded_data = base64_decode(data); // expected to be base64 encoded + out_files.push_back(decoded_data); + + p["type"] = "media_marker"; + p["text"] = get_media_marker(); + p.erase("input_video"); + } else if (type != "text") { throw std::invalid_argument("unsupported content[].type"); } diff --git a/tools/server/server-common.h b/tools/server/server-common.h index 51b161317..249b97c2f 100644 --- a/tools/server/server-common.h +++ b/tools/server/server-common.h @@ -294,6 +294,7 @@ struct server_chat_params { common_chat_templates_ptr tmpls; bool allow_image; bool allow_audio; + bool allow_video; bool enable_thinking = true; int reasoning_budget = -1; std::string reasoning_budget_message; diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 07759f417..bdfa51718 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -27,6 +27,7 @@ #include #include #include +#include // fix problem with std::min and std::max #if defined(_WIN32) @@ -127,7 +128,11 @@ struct server_slot { server_prompt prompt; - void prompt_save(server_prompt_cache & prompt_cache) const { + bool prompt_save(server_prompt_cache & prompt_cache) const { + if (prompt.tokens.size() == 0) { + return false; + } + GGML_ASSERT(prompt.data.size() == 0); const size_t cur_size_tgt = llama_state_seq_get_size_ext(ctx_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE); @@ -140,13 +145,15 @@ struct server_slot { auto * cur = prompt_cache.alloc(prompt, cur_size_tgt, cur_size_dft); if (cur == nullptr) { - return; + return false; } llama_state_seq_get_data_ext(ctx_tgt, cur->data.main.data(), cur_size_tgt, id, LLAMA_STATE_SEQ_FLAGS_NONE); if (ctx_dft) { llama_state_seq_get_data_ext(ctx_dft, cur->data.drft.data(), cur_size_dft, id, LLAMA_STATE_SEQ_FLAGS_NONE); } + + return true; } bool prompt_load(server_prompt_cache & prompt_cache, const server_tokens & tokens) { @@ -739,17 +746,6 @@ private: llama_batch_free(batch); } - void slot_save_and_clear(server_slot & slot) { - if (slot.prompt.n_tokens() == 0) { - return; - } - SLT_INF(slot, "%s", "saving idle slot to prompt cache\n"); - SLT_DBG(slot, "%s", "__TEST_TAG_CACHE_IDLE_SLOT__\n"); - slot.prompt_save(*prompt_cache); - slot.prompt_clear(false); - prompt_cache->update(); - } - void handle_sleeping_state(bool new_state) { GGML_ASSERT(sleeping != new_state); if (new_state) { @@ -1186,14 +1182,17 @@ private: metrics.init(); if (params_base.cache_idle_slots) { - if (!params_base.kv_unified) { - SRV_WRN("%s", "--cache-idle-slots requires --kv-unified, disabling\n"); - params_base.cache_idle_slots = false; - } else if (params_base.cache_ram_mib == 0) { + if (params_base.cache_ram_mib == 0) { SRV_WRN("%s", "--cache-idle-slots requires --cache-ram, disabling\n"); params_base.cache_idle_slots = false; } else { - SRV_INF("%s", "idle slots will be saved to prompt cache and cleared upon starting a new task\n"); + if (params_base.kv_unified) { + SRV_INF("%s", "idle slots will be saved to prompt cache and cleared upon starting a new task\n"); + } else { + // without a unified KV cache, clearing a slot frees no reusable room, so we only + // publish a RAM-cache copy of idle slots (their KV stays in VRAM) [TAG_IDLE_SLOT_CLEAR] + SRV_INF("%s", "idle slots will be saved to prompt cache upon starting a new task\n"); + } SRV_DBG("%s", "__TEST_TAG_CACHE_IDLE_SLOTS_ENABLED__\n"); } } @@ -1247,6 +1246,7 @@ private: /* tmpls */ std::move(chat_templates), /* allow_image */ mctx ? mtmd_support_vision(mctx) : false, /* allow_audio */ mctx ? mtmd_support_audio (mctx) : false, + /* allow_video */ mctx ? mtmd_helper_support_video(mctx) : false, /* enable_thinking */ enable_thinking, /* reasoning_budget */ params_base.sampling.reasoning_budget_tokens, /* reasoning_budget_msg */ params_base.sampling.reasoning_budget_message, @@ -1356,8 +1356,6 @@ private: } if (ret) { - const auto & tokens = ret->prompt.tokens; - update_cache = update_cache && prompt_cache; // cache prompts only for completion tasks @@ -1368,10 +1366,7 @@ private: const int64_t t_start = ggml_time_us(); - // don't save the slot's state if its context is empty - if (tokens.size() > 0) { - ret->prompt_save(*prompt_cache); - } + ret->prompt_save(*prompt_cache); if (!ret->prompt_load(*prompt_cache, task.tokens)) { ret->prompt_clear(false); @@ -2119,9 +2114,19 @@ private: } if (params_base.cache_idle_slots) { - for (auto & s : slots) { - if (!s.is_processing()) { - slot_save_and_clear(s); + for (auto & slot : slots) { + if (!slot.is_processing()) { + SLT_INF(slot, "%s", "saving idle slot to prompt cache\n"); + + if (slot.prompt_save(*prompt_cache)) { + SLT_DBG(slot, "%s", "__TEST_TAG_CACHE_IDLE_SLOT__\n"); + prompt_cache->update(); + } + + if (params_base.kv_unified) { + // [TAG_IDLE_SLOT_CLEAR] + slot.prompt_clear(false); + } } } } @@ -3586,6 +3591,7 @@ server_context_meta server_context::get_meta() const { /* has_mtmd */ impl->mctx != nullptr, /* has_inp_image */ impl->chat_params.allow_image, /* has_inp_audio */ impl->chat_params.allow_audio, + /* has_inp_video */ impl->chat_params.allow_video, /* json_ui_settings */ impl->json_ui_settings, /* json_webui_settings */ impl->json_webui_settings, // Deprecated /* slot_n_ctx */ impl->get_slot_n_ctx(), @@ -3714,6 +3720,16 @@ std::unique_ptr server_routes::handle_completions_impl( // TODO: this log can become very long, put it behind a flag or think about a more compact format //SRV_DBG("Prompt: %s\n", prompt.is_string() ? prompt.get().c_str() : prompt.dump(2).c_str()); + if (!params.path_prompts_log_dir.empty()) { + const auto file_path = std::filesystem::path(params.path_prompts_log_dir) / string_format("%012" PRId64 ".txt", ggml_time_ms()); + std::ofstream f(file_path); + if (f) { + f << (prompt.is_string() ? prompt.get().c_str() : prompt.dump(2).c_str()); + } else { + SRV_ERR("failed to create %s\n", file_path.string().c_str()); + } + } + // process prompt std::vector inputs; @@ -4183,6 +4199,7 @@ void server_routes::init_routes() { { "model_path", meta->model_path }, { "modalities", json { {"vision", meta->has_inp_image}, + {"video", meta->has_inp_video}, {"audio", meta->has_inp_audio}, } }, { "media_marker", get_media_marker() }, @@ -4976,7 +4993,7 @@ std::unique_ptr server_routes::handle_count_tokens(const l n_tokens = tokenize_mixed(vocab, prompt, true, true).size(); } - json response = {{"input_tokens", static_cast(n_tokens)}}; + json response = {{"input_tokens", static_cast(n_tokens)}}; if (is_oai) { response["object"] = "response.input_tokens"; } diff --git a/tools/server/server-context.h b/tools/server/server-context.h index 72a1f40e0..0e84785af 100644 --- a/tools/server/server-context.h +++ b/tools/server/server-context.h @@ -21,6 +21,7 @@ struct server_context_meta { bool has_mtmd; bool has_inp_image; bool has_inp_audio; + bool has_inp_video; json json_ui_settings; // Primary: new name json json_webui_settings; // Deprecated: use json_ui_settings instead (kept for backward compat) int slot_n_ctx; diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index 842be2ad3..72a4bd076 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -1393,6 +1393,9 @@ json server_task_result_cmpl_final::to_json_anthropic_stream() { // void server_task_result_cmpl_partial::update(task_result_state & state) { is_updated = true; + if (is_begin) { + return; // begin marker only flushes headers, skip parsing + } state.update_chat_msg(content, true, oaicompat_msg_diffs); // Copy current state for use in to_json_*() (reflects state BEFORE this chunk) diff --git a/tools/ui/eslint.config.js b/tools/ui/eslint.config.js index 6b3b1b5c0..b90376b6c 100644 --- a/tools/ui/eslint.config.js +++ b/tools/ui/eslint.config.js @@ -46,7 +46,14 @@ export default ts.config( }, { // Exclude generated build output and Storybook files from ESLint - ignores: ['dist/**', 'build/**', '.svelte-kit/**', 'test-results/**', '.storybook/**/*'] + ignores: [ + 'dist/**', + 'build/**', + '.svelte-kit/**', + 'test-results/**', + '.storybook/**/*', + 'src/lib/services/sandbox-worker.js' + ] }, storybook.configs['flat/recommended'] ); diff --git a/tools/ui/scripts/vite-plugin-llama-cpp-build.ts b/tools/ui/scripts/vite-plugin-llama-cpp-build.ts index 01c714a24..74e3de9ba 100644 --- a/tools/ui/scripts/vite-plugin-llama-cpp-build.ts +++ b/tools/ui/scripts/vite-plugin-llama-cpp-build.ts @@ -46,10 +46,12 @@ export function llamaCppBuildPlugin(): Plugin { content = content.replace(/\r/g, ''); content = GUIDE_FOR_FRONTEND + '\n' + content; - content = content.replace(/\/_app\/immutable\/bundle\.[^"]+\.js/g, './bundle.js'); + + // Keep the Vite hash as a query string so each build busts the browser cache + content = content.replace(/\/_app\/immutable\/bundle\.([^".]+)\.js/g, './bundle.js?$1'); content = content.replace( - /\/_app\/immutable\/assets\/bundle\.[^"]+\.css/g, - './bundle.css' + /\/_app\/immutable\/assets\/bundle\.([^".]+)\.css/g, + './bundle.css?$1' ); content = content.replace(/__sveltekit_[a-z0-9]+/g, '__sveltekit__'); diff --git a/tools/ui/src/app.css b/tools/ui/src/app.css index aa96bac32..9254a96df 100644 --- a/tools/ui/src/app.css +++ b/tools/ui/src/app.css @@ -148,7 +148,6 @@ * { scrollbar-width: thin; scrollbar-color: transparent transparent; - transition: scrollbar-color 0.2s ease; } *:hover { diff --git a/tools/ui/src/lib/constants/index.ts b/tools/ui/src/lib/constants/index.ts index c4334132c..9ab864cd0 100644 --- a/tools/ui/src/lib/constants/index.ts +++ b/tools/ui/src/lib/constants/index.ts @@ -37,6 +37,7 @@ export * from './model-id'; export * from './precision'; export * from './processing-info'; export * from './routes'; +export * from './sandbox'; export * from './settings-keys'; export * from './settings-registry'; export * from './supported-file-types'; diff --git a/tools/ui/src/lib/constants/sandbox.ts b/tools/ui/src/lib/constants/sandbox.ts new file mode 100644 index 000000000..30e769509 --- /dev/null +++ b/tools/ui/src/lib/constants/sandbox.ts @@ -0,0 +1,39 @@ +import { JsonSchemaType, ToolCallType } from '$lib/enums'; +import type { OpenAIToolDefinition } from '$lib/types'; + +export const SANDBOX_TOOL_NAME = 'run_javascript'; + +export const SANDBOX_TIMEOUT_MS_DEFAULT = 10000; + +export const SANDBOX_TIMEOUT_MS_MAX = 30000; + +export const SANDBOX_OUTPUT_MAX_CHARS = 8192; + +export const SANDBOX_EMPTY_OUTPUT = '(no output)'; + +export const SANDBOX_TRUNCATION_NOTICE = '[output truncated]'; + +export const SANDBOX_TOOL_DEFINITION: OpenAIToolDefinition = { + type: ToolCallType.FUNCTION, + function: { + name: SANDBOX_TOOL_NAME, + description: + 'Execute JavaScript in a sandboxed browser worker (no DOM, no page access). ' + + 'Top level await is supported. Use console.log to print intermediate values; ' + + 'a top level return statement is captured as the result.', + parameters: { + type: JsonSchemaType.OBJECT, + properties: { + code: { + type: JsonSchemaType.STRING, + description: 'JavaScript source to execute' + }, + timeout_ms: { + type: JsonSchemaType.NUMBER, + description: `Execution timeout in milliseconds, default ${SANDBOX_TIMEOUT_MS_DEFAULT}, max ${SANDBOX_TIMEOUT_MS_MAX}` + } + }, + required: ['code'] + } + } +}; diff --git a/tools/ui/src/lib/constants/settings-keys.ts b/tools/ui/src/lib/constants/settings-keys.ts index 5fff9f94c..3e19166df 100644 --- a/tools/ui/src/lib/constants/settings-keys.ts +++ b/tools/ui/src/lib/constants/settings-keys.ts @@ -69,6 +69,7 @@ export const SETTINGS_KEYS = { ENABLE_THINKING: 'enableThinking', SHOW_RAW_OUTPUT_SWITCH: 'showRawOutputSwitch', // PY_INTERPRETER_ENABLED: 'pyInterpreterEnabled', + JS_SANDBOX_ENABLED: 'jsSandboxEnabled', CUSTOM_JSON: 'customJson', CUSTOM_CSS: 'customCss' } as const; diff --git a/tools/ui/src/lib/constants/settings-registry.ts b/tools/ui/src/lib/constants/settings-registry.ts index 9246b9703..910e5c2df 100644 --- a/tools/ui/src/lib/constants/settings-registry.ts +++ b/tools/ui/src/lib/constants/settings-registry.ts @@ -690,6 +690,14 @@ const SETTINGS_REGISTRY: Record = { paramType: SyncableParameterType.BOOLEAN } }, + { + key: SETTINGS_KEYS.JS_SANDBOX_ENABLED, + label: 'JavaScript sandbox tool', + help: 'Expose a run_javascript tool to the model. Code runs in a Web Worker inside a sandboxed iframe with an opaque origin, isolated from the WebUI and its API, with a hard timeout.', + defaultValue: false, + type: SettingsFieldType.CHECKBOX, + section: SETTINGS_SECTION_SLUGS.DEVELOPER + }, { key: SETTINGS_KEYS.CUSTOM_JSON, label: 'Custom JSON', diff --git a/tools/ui/src/lib/constants/tools.ts b/tools/ui/src/lib/constants/tools.ts index efc3476cd..4d9385f9f 100644 --- a/tools/ui/src/lib/constants/tools.ts +++ b/tools/ui/src/lib/constants/tools.ts @@ -2,10 +2,12 @@ import { ToolSource } from '$lib/enums/tools.enums'; export const TOOL_GROUP_LABELS = { [ToolSource.BUILTIN]: 'Built-in', - [ToolSource.CUSTOM]: 'JSON Schema' + [ToolSource.CUSTOM]: 'JSON Schema', + [ToolSource.FRONTEND]: 'Browser' } as const; export const TOOL_SERVER_LABELS = { [ToolSource.BUILTIN]: 'Built-in Tools', - [ToolSource.CUSTOM]: 'Custom Tools' + [ToolSource.CUSTOM]: 'Custom Tools', + [ToolSource.FRONTEND]: 'Browser Tools' } as const; diff --git a/tools/ui/src/lib/enums/mcp.enums.ts b/tools/ui/src/lib/enums/mcp.enums.ts index d2c27e1a0..3d9a2070d 100644 --- a/tools/ui/src/lib/enums/mcp.enums.ts +++ b/tools/ui/src/lib/enums/mcp.enums.ts @@ -54,7 +54,9 @@ export enum MCPContentType { * JSON Schema types used in MCP tool definitions */ export enum JsonSchemaType { - OBJECT = 'object' + OBJECT = 'object', + STRING = 'string', + NUMBER = 'number' } /** diff --git a/tools/ui/src/lib/enums/tools.enums.ts b/tools/ui/src/lib/enums/tools.enums.ts index 4b2cdab32..120c1d471 100644 --- a/tools/ui/src/lib/enums/tools.enums.ts +++ b/tools/ui/src/lib/enums/tools.enums.ts @@ -1,7 +1,8 @@ export enum ToolSource { BUILTIN = 'builtin', MCP = 'mcp', - CUSTOM = 'custom' + CUSTOM = 'custom', + FRONTEND = 'frontend' } export enum ToolPermissionDecision { diff --git a/tools/ui/src/lib/services/index.ts b/tools/ui/src/lib/services/index.ts index 475e6419b..386f740b8 100644 --- a/tools/ui/src/lib/services/index.ts +++ b/tools/ui/src/lib/services/index.ts @@ -261,6 +261,26 @@ export { ParameterSyncService } from './parameter-sync.service'; */ export { MCPService } from './mcp.service'; +/** + * **SandboxService** - Frontend JavaScript execution in a browser sandbox + * + * Stateless executor for the run_javascript frontend tool. Model generated + * code runs in a Web Worker spawned inside a sandboxed iframe with an opaque + * origin: no access to the app origin, its storage or its API, and outgoing + * requests carry a null origin. The code never touches a main thread, so the + * parent enforces the timeout by removing the iframe, which terminates the + * worker at the browser level. + * + * **Architecture & Relationships:** + * - **SandboxService** (this class): Stateless sandbox execution + * - **toolsStore**: Exposes the tool definition when the sandbox is enabled + * - **agenticStore**: Dispatches ToolSource.FRONTEND calls here + * + * @see SANDBOX_TOOL_DEFINITION in constants/sandbox.ts - tool schema sent to the LLM + * @see agenticStore in stores/agentic.svelte.ts - tool dispatch + */ +export { SandboxService } from './sandbox.service'; + /** * **RouterService** — Dynamic route URL construction utility * diff --git a/tools/ui/src/lib/services/sandbox-harness.ts b/tools/ui/src/lib/services/sandbox-harness.ts new file mode 100644 index 000000000..27b05e24b --- /dev/null +++ b/tools/ui/src/lib/services/sandbox-harness.ts @@ -0,0 +1,25 @@ +import WORKER_SHIM from './sandbox-worker.js?raw'; + +/** + * Harness loaded as srcdoc into a sandboxed iframe (allow-scripts only). + * The opaque origin is the security boundary: no access to the app origin, + * its storage or its API. The harness spawns a worker so model code never + * runs on a main thread, which makes the parent timeout enforceable by + * removing the iframe. + */ +export const SANDBOX_HARNESS_HTML = ``; diff --git a/tools/ui/src/lib/services/sandbox-worker.js b/tools/ui/src/lib/services/sandbox-worker.js new file mode 100644 index 000000000..689b9211e --- /dev/null +++ b/tools/ui/src/lib/services/sandbox-worker.js @@ -0,0 +1,30 @@ +const logs = []; +const fmt = (value) => { + if (typeof value === 'string') return value; + try { + return JSON.stringify(value); + } catch { + return String(value); + } +}; +const capture = + (level, prefix) => + (...args) => { + logs.push(prefix + args.map(fmt).join(' ')); + }; +console.log = capture('log', ''); +console.info = capture('info', ''); +console.debug = capture('debug', ''); +console.warn = capture('warn', 'warn: '); +console.error = capture('error', 'error: '); +self.onmessage = async (event) => { + const reply = { logs, result: null, error: null }; + try { + const AsyncFunction = Object.getPrototypeOf(async function () {}).constructor; + const value = await new AsyncFunction(event.data.code)(); + if (value !== undefined) reply.result = fmt(value); + } catch (err) { + reply.error = err instanceof Error ? err.stack || err.message : String(err); + } + self.postMessage(reply); +}; diff --git a/tools/ui/src/lib/services/sandbox.service.ts b/tools/ui/src/lib/services/sandbox.service.ts new file mode 100644 index 000000000..ef0e6cf4e --- /dev/null +++ b/tools/ui/src/lib/services/sandbox.service.ts @@ -0,0 +1,112 @@ +import { + NEWLINE_SEPARATOR, + SANDBOX_EMPTY_OUTPUT, + SANDBOX_OUTPUT_MAX_CHARS, + SANDBOX_TIMEOUT_MS_DEFAULT, + SANDBOX_TIMEOUT_MS_MAX, + SANDBOX_TOOL_NAME, + SANDBOX_TRUNCATION_NOTICE +} from '$lib/constants'; +import { SANDBOX_HARNESS_HTML } from './sandbox-harness'; +import type { ToolExecutionResult } from '$lib/types'; + +interface SandboxReply { + logs?: unknown; + result?: unknown; + error?: unknown; +} + +function formatReply(reply: SandboxReply): ToolExecutionResult { + const lines: string[] = []; + + if (Array.isArray(reply.logs)) { + for (const line of reply.logs) lines.push(String(line)); + } + + if (reply.error != null) { + lines.push(`Error: ${String(reply.error)}`); + } else if (reply.result != null) { + lines.push(`=> ${String(reply.result)}`); + } + + let content = lines.join(NEWLINE_SEPARATOR); + if (!content) content = SANDBOX_EMPTY_OUTPUT; + if (content.length > SANDBOX_OUTPUT_MAX_CHARS) { + content = `${content.slice(0, SANDBOX_OUTPUT_MAX_CHARS)}${NEWLINE_SEPARATOR}${SANDBOX_TRUNCATION_NOTICE}`; + } + + return { content, isError: reply.error != null }; +} + +export class SandboxService { + /** + * Execute a frontend sandbox tool call and return its output. + * One disposable iframe per execution, removed on completion, + * timeout or abort. Removing the iframe terminates the worker + * at the browser level, so runaway code cannot outlive it. + */ + static executeTool( + toolName: string, + params: Record, + signal?: AbortSignal + ): Promise { + if (toolName !== SANDBOX_TOOL_NAME) { + return Promise.resolve({ content: `Unknown frontend tool: ${toolName}`, isError: true }); + } + + const code = typeof params.code === 'string' ? params.code : ''; + if (!code) { + return Promise.resolve({ content: 'Missing required parameter: code', isError: true }); + } + + const requested = Number(params.timeout_ms); + const timeoutMs = + Number.isFinite(requested) && requested > 0 + ? Math.min(requested, SANDBOX_TIMEOUT_MS_MAX) + : SANDBOX_TIMEOUT_MS_DEFAULT; + + return new Promise((resolve, reject) => { + const iframe = document.createElement('iframe'); + iframe.setAttribute('sandbox', 'allow-scripts'); + iframe.style.display = 'none'; + iframe.srcdoc = SANDBOX_HARNESS_HTML; + + let settled = false; + + const cleanup = () => { + settled = true; + clearTimeout(timer); + window.removeEventListener('message', onMessage); + signal?.removeEventListener('abort', onAbort); + iframe.remove(); + }; + + const finish = (result: ToolExecutionResult) => { + if (settled) return; + cleanup(); + resolve(result); + }; + + const onAbort = () => { + if (settled) return; + cleanup(); + reject(new DOMException('Sandbox execution aborted', 'AbortError')); + }; + + const onMessage = (event: MessageEvent) => { + if (event.source !== iframe.contentWindow) return; + finish(formatReply((event.data ?? {}) as SandboxReply)); + }; + + const timer = setTimeout( + () => finish({ content: `Execution timed out after ${timeoutMs} ms`, isError: true }), + timeoutMs + ); + + window.addEventListener('message', onMessage); + signal?.addEventListener('abort', onAbort); + iframe.onload = () => iframe.contentWindow?.postMessage({ code }, '*'); + document.body.appendChild(iframe); + }); + } +} diff --git a/tools/ui/src/lib/stores/agentic.svelte.ts b/tools/ui/src/lib/stores/agentic.svelte.ts index 947737d7c..5579cc1e5 100644 --- a/tools/ui/src/lib/stores/agentic.svelte.ts +++ b/tools/ui/src/lib/stores/agentic.svelte.ts @@ -29,6 +29,7 @@ import { permissionsStore } from '$lib/stores/permissions.svelte'; import { ToolSource, ToolPermissionDecision } from '$lib/enums'; import { SvelteMap } from 'svelte/reactivity'; import { ToolsService } from '$lib/services/tools.service'; +import { SandboxService } from '$lib/services/sandbox.service'; import { isAbortError } from '$lib/utils'; import { DEFAULT_AGENTIC_CONFIG, NEWLINE_SEPARATOR } from '$lib/constants'; import { @@ -784,6 +785,13 @@ class AgenticStore { result = executionResult.content; + if (executionResult.isError) toolSuccess = false; + } else if (toolSource === ToolSource.FRONTEND) { + const args = this.parseToolArguments(toolCall.function.arguments); + const executionResult = await SandboxService.executeTool(toolName, args, signal); + + result = executionResult.content; + if (executionResult.isError) toolSuccess = false; } else { const mcpCall: MCPToolCall = { diff --git a/tools/ui/src/lib/stores/tools.svelte.ts b/tools/ui/src/lib/stores/tools.svelte.ts index 82e41f0bf..9f0101a82 100644 --- a/tools/ui/src/lib/stores/tools.svelte.ts +++ b/tools/ui/src/lib/stores/tools.svelte.ts @@ -5,6 +5,7 @@ import { HealthCheckStatus, JsonSchemaType, ToolCallType, ToolSource } from '$li import { config } from '$lib/stores/settings.svelte'; import { DISABLED_TOOL_KEYS_LOCALSTORAGE_KEY, + SANDBOX_TOOL_DEFINITION, TOOL_GROUP_LABELS, TOOL_SERVER_LABELS } from '$lib/constants'; @@ -18,6 +19,8 @@ function toolKey(source: ToolSource, name: string, serverId?: string): string { return serverId ? `mcp-${serverId}:${name}` : `mcp:${name}`; case ToolSource.CUSTOM: return `custom:${name}`; + case ToolSource.FRONTEND: + return `frontend:${name}`; default: return `builtin:${name}`; } @@ -82,6 +85,10 @@ class ToolsStore { return mcpStore.getToolDefinitionsForLLM(); } + get frontendTools(): OpenAIToolDefinition[] { + return config().jsSandboxEnabled ? [SANDBOX_TOOL_DEFINITION] : []; + } + get customTools(): OpenAIToolDefinition[] { const raw = config().customJson; if (!raw || typeof raw !== 'string') return []; @@ -156,6 +163,15 @@ class ToolsStore { push({ source: ToolSource.BUILTIN, key: toolKey(ToolSource.BUILTIN, name), definition: def }); } + for (const def of this.frontendTools) { + const name = def.function.name; + push({ + source: ToolSource.FRONTEND, + key: toolKey(ToolSource.FRONTEND, name), + definition: def + }); + } + for (const { serverId, serverName, definition } of this.mcpEntries()) { const name = definition.function.name; push({ @@ -208,6 +224,8 @@ class ToolsStore { return entry.serverName ?? ''; case ToolSource.CUSTOM: return TOOL_GROUP_LABELS[ToolSource.CUSTOM]; + case ToolSource.FRONTEND: + return TOOL_GROUP_LABELS[ToolSource.FRONTEND]; default: return TOOL_GROUP_LABELS[ToolSource.BUILTIN]; } @@ -237,6 +255,7 @@ class ToolsStore { }; for (const def of this._builtinTools) take(def); + for (const def of this.frontendTools) take(def); for (const def of mcpStore.getToolDefinitionsForLLM()) take(def); for (const def of this.customTools) take(def); @@ -346,6 +365,7 @@ class ToolsStore { if (entry.serverName) return mcpStore.getServerDisplayName(entry.serverName); if (entry.source === ToolSource.BUILTIN) return TOOL_SERVER_LABELS[ToolSource.BUILTIN]; if (entry.source === ToolSource.CUSTOM) return TOOL_SERVER_LABELS[ToolSource.CUSTOM]; + if (entry.source === ToolSource.FRONTEND) return TOOL_SERVER_LABELS[ToolSource.FRONTEND]; return ''; } diff --git a/tools/ui/src/routes/+layout.svelte b/tools/ui/src/routes/+layout.svelte index be474109a..aa023840b 100644 --- a/tools/ui/src/routes/+layout.svelte +++ b/tools/ui/src/routes/+layout.svelte @@ -254,7 +254,7 @@ /> -
+
diff --git a/tools/ui/tsconfig.json b/tools/ui/tsconfig.json index 7c585f4db..51f55971e 100644 --- a/tools/ui/tsconfig.json +++ b/tools/ui/tsconfig.json @@ -24,7 +24,8 @@ "tests/**/*.svelte", ".storybook/**/*.ts", ".storybook/**/*.svelte" - ] + ], + "exclude": ["src/lib/services/sandbox-worker.js"] // Path aliases are handled by https://svelte.dev/docs/kit/configuration#alias // except $lib which is handled by https://svelte.dev/docs/kit/configuration#files //