From 915dc6d38c9b7545745f7ef71d7f6bd4103103bb Mon Sep 17 00:00:00 2001 From: Niklas Wenzel Date: Thu, 27 Aug 2026 11:53:08 +0200 Subject: [PATCH 1/8] metal : fix memory leaks due to missing autoreleasepools (#27758) --- ggml/src/ggml-metal/ggml-metal-context.m | 174 +++++----- ggml/src/ggml-metal/ggml-metal-device.m | 412 ++++++++++++----------- 2 files changed, 296 insertions(+), 290 deletions(-) diff --git a/ggml/src/ggml-metal/ggml-metal-context.m b/ggml/src/ggml-metal/ggml-metal-context.m index 32d97cd5d0..1227ed39a0 100644 --- a/ggml/src/ggml-metal/ggml-metal-context.m +++ b/ggml/src/ggml-metal/ggml-metal-context.m @@ -84,106 +84,108 @@ struct ggml_metal { ggml_metal_t ggml_metal_init(ggml_metal_device_t dev) { GGML_LOG_INFO("%s: allocating\n", __func__); + @autoreleasepool { #if TARGET_OS_OSX && !GGML_METAL_NDEBUG - // Show all the Metal device instances in the system - NSArray * devices = MTLCopyAllDevices(); - for (id device in devices) { - GGML_LOG_INFO("%s: found device: %s\n", __func__, [[device name] UTF8String]); - } - [devices release]; // since it was created by a *Copy* C method + // Show all the Metal device instances in the system + NSArray * devices = MTLCopyAllDevices(); + for (id device in devices) { + GGML_LOG_INFO("%s: found device: %s\n", __func__, [[device name] UTF8String]); + } + [devices release]; // since it was created by a *Copy* C method #endif - // init context - ggml_metal_t res = calloc(1, sizeof(struct ggml_metal)); + // init context + ggml_metal_t res = calloc(1, sizeof(struct ggml_metal)); - id device = ggml_metal_device_get_obj(dev); + id device = ggml_metal_device_get_obj(dev); - GGML_LOG_INFO("%s: picking default device: %s\n", __func__, [[device name] UTF8String]); - - // TODO: would it be better to have one queue for the backend and one queue for the device? - // the graph encoders and async ops would use the backend queue while the sync ops would use the device queue? - //res->queue = [device newCommandQueue]; [TAG_QUEUE_PER_BACKEND] - id queue = ggml_metal_device_get_queue(dev); - if (queue == nil) { - GGML_LOG_ERROR("%s: error: failed to create command queue\n", __func__); - return NULL; - } - - res->dev = dev; - res->lib = ggml_metal_device_get_library(dev); - if (res->lib == NULL) { - GGML_LOG_WARN("%s: the device does not have a precompiled Metal library - this is unexpected\n", __func__); - GGML_LOG_WARN("%s: will try to compile it on the fly\n", __func__); - - res->lib = ggml_metal_library_init(dev); - if (res->lib == NULL) { - GGML_LOG_ERROR("%s: error: failed to initialize the Metal library\n", __func__); - - free(res); + GGML_LOG_INFO("%s: picking default device: %s\n", __func__, [[device name] UTF8String]); + // TODO: would it be better to have one queue for the backend and one queue for the device? + // the graph encoders and async ops would use the backend queue while the sync ops would use the device queue? + //res->queue = [device newCommandQueue]; [TAG_QUEUE_PER_BACKEND] + id queue = ggml_metal_device_get_queue(dev); + if (queue == nil) { + GGML_LOG_ERROR("%s: error: failed to create command queue\n", __func__); return NULL; } - } - res->ev_cpy = ggml_metal_device_event_init(dev); + res->dev = dev; + res->lib = ggml_metal_device_get_library(dev); + if (res->lib == NULL) { + GGML_LOG_WARN("%s: the device does not have a precompiled Metal library - this is unexpected\n", __func__); + GGML_LOG_WARN("%s: will try to compile it on the fly\n", __func__); - const struct ggml_metal_device_props * props_dev = ggml_metal_device_get_props(dev); + res->lib = ggml_metal_library_init(dev); + if (res->lib == NULL) { + GGML_LOG_ERROR("%s: error: failed to initialize the Metal library\n", __func__); - snprintf(res->name, sizeof(res->name), "%s", props_dev->name); + free(res); - res->d_queue = dispatch_queue_create("ggml-metal", DISPATCH_QUEUE_CONCURRENT); - - res->use_fusion = getenv("GGML_METAL_FUSION_DISABLE") == nil; - res->use_concurrency = getenv("GGML_METAL_CONCURRENCY_DISABLE") == nil; - - { - const char * val = getenv("GGML_METAL_GRAPH_DEBUG"); - res->debug_graph = val ? atoi(val) : 0; - } - - { - const char * val = getenv("GGML_METAL_FUSION_DEBUG"); - res->debug_fusion = val ? atoi(val) : 0; - } - - res->use_graph_optimize = true; - - if (getenv("GGML_METAL_GRAPH_OPTIMIZE_DISABLE") != NULL) { - res->use_graph_optimize = false; - } - - memset(res->fuse_cnt, 0, sizeof(res->fuse_cnt)); - - GGML_LOG_INFO("%s: use fusion = %s\n", __func__, res->use_fusion ? "true" : "false"); - GGML_LOG_INFO("%s: use concurrency = %s\n", __func__, res->use_concurrency ? "true" : "false"); - GGML_LOG_INFO("%s: use graph optimize = %s\n", __func__, res->use_graph_optimize ? "true" : "false"); - - res->capture_compute = 0; - res->capture_started = false; - res->capture_scope = nil; - - { - const char * val = getenv("GGML_METAL_CAPTURE_COMPUTE"); - if (val) { - res->capture_compute = atoi(val); + return NULL; + } } + + res->ev_cpy = ggml_metal_device_event_init(dev); + + const struct ggml_metal_device_props * props_dev = ggml_metal_device_get_props(dev); + + snprintf(res->name, sizeof(res->name), "%s", props_dev->name); + + res->d_queue = dispatch_queue_create("ggml-metal", DISPATCH_QUEUE_CONCURRENT); + + res->use_fusion = getenv("GGML_METAL_FUSION_DISABLE") == nil; + res->use_concurrency = getenv("GGML_METAL_CONCURRENCY_DISABLE") == nil; + + { + const char * val = getenv("GGML_METAL_GRAPH_DEBUG"); + res->debug_graph = val ? atoi(val) : 0; + } + + { + const char * val = getenv("GGML_METAL_FUSION_DEBUG"); + res->debug_fusion = val ? atoi(val) : 0; + } + + res->use_graph_optimize = true; + + if (getenv("GGML_METAL_GRAPH_OPTIMIZE_DISABLE") != NULL) { + res->use_graph_optimize = false; + } + + memset(res->fuse_cnt, 0, sizeof(res->fuse_cnt)); + + GGML_LOG_INFO("%s: use fusion = %s\n", __func__, res->use_fusion ? "true" : "false"); + GGML_LOG_INFO("%s: use concurrency = %s\n", __func__, res->use_concurrency ? "true" : "false"); + GGML_LOG_INFO("%s: use graph optimize = %s\n", __func__, res->use_graph_optimize ? "true" : "false"); + + res->capture_compute = 0; + res->capture_started = false; + res->capture_scope = nil; + + { + const char * val = getenv("GGML_METAL_CAPTURE_COMPUTE"); + if (val) { + res->capture_compute = atoi(val); + } + } + + res->has_error = false; + + res->gf = nil; + res->encode_async = nil; + for (int i = 0; i < GGML_METAL_MAX_COMMAND_BUFFERS; ++i) { + res->cmd_bufs[i].obj = nil; + } + + res->cmd_bufs_ext = [[NSMutableArray alloc] init]; + + res->cmd_buf_last = nil; + + res->pipelines_ext = ggml_metal_pipelines_init(); + + return res; } - - res->has_error = false; - - res->gf = nil; - res->encode_async = nil; - for (int i = 0; i < GGML_METAL_MAX_COMMAND_BUFFERS; ++i) { - res->cmd_bufs[i].obj = nil; - } - - res->cmd_bufs_ext = [[NSMutableArray alloc] init]; - - res->cmd_buf_last = nil; - - res->pipelines_ext = ggml_metal_pipelines_init(); - - return res; } void ggml_metal_free(ggml_metal_t ctx) { diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index 19c57820e8..41ce90dc8a 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -778,7 +778,9 @@ void ggml_metal_encoder_free(ggml_metal_encoder_t encoder) { } void ggml_metal_encoder_debug_group_push(ggml_metal_encoder_t encoder, const char * name) { - [encoder->obj pushDebugGroup:[NSString stringWithCString:name encoding:NSUTF8StringEncoding]]; + @autoreleasepool { + [encoder->obj pushDebugGroup:[NSString stringWithCString:name encoding:NSUTF8StringEncoding]]; + } } void ggml_metal_encoder_debug_group_pop (ggml_metal_encoder_t encoder) { @@ -1023,249 +1025,251 @@ ggml_metal_device_t ggml_metal_device_init(int device, int n_devices) { assert(dev != NULL); - if (dev->mtl_device == nil) { - dev->mtl_device = MTLCreateSystemDefaultDevice(); + @autoreleasepool { + if (dev->mtl_device == nil) { + dev->mtl_device = MTLCreateSystemDefaultDevice(); - if (dev->mtl_device) { - dev->mtl_queue = [dev->mtl_device newCommandQueue]; - if (dev->mtl_queue == nil) { - GGML_LOG_ERROR("%s: error: failed to create command queue\n", __func__); - } + if (dev->mtl_device) { + dev->mtl_queue = [dev->mtl_device newCommandQueue]; + if (dev->mtl_queue == nil) { + GGML_LOG_ERROR("%s: error: failed to create command queue\n", __func__); + } - dev->addr_virt = 0x000000400ULL; + dev->addr_virt = 0x000000400ULL; - dev->props.device = device; + dev->props.device = device; - // the Metal backend uses the system default device as the single physical device; - // additional (virtual) devices are emulated on top of it via GGML_METAL_DEVICES - dev->props.device_phys = 0; - dev->props.device_virt = device; + // the Metal backend uses the system default device as the single physical device; + // additional (virtual) devices are emulated on top of it via GGML_METAL_DEVICES + dev->props.device_phys = 0; + dev->props.device_virt = device; - dev->props.has_simdgroup_reduction = [dev->mtl_device supportsFamily:MTLGPUFamilyApple7]; - dev->props.has_simdgroup_reduction |= [dev->mtl_device supportsFamily:MTLGPUFamilyMetal3_GGML]; + dev->props.has_simdgroup_reduction = [dev->mtl_device supportsFamily:MTLGPUFamilyApple7]; + dev->props.has_simdgroup_reduction |= [dev->mtl_device supportsFamily:MTLGPUFamilyMetal3_GGML]; - dev->props.has_simdgroup_mm = [dev->mtl_device supportsFamily:MTLGPUFamilyApple7]; - dev->props.has_unified_memory = dev->mtl_device.hasUnifiedMemory; + dev->props.has_simdgroup_mm = [dev->mtl_device supportsFamily:MTLGPUFamilyApple7]; + dev->props.has_unified_memory = dev->mtl_device.hasUnifiedMemory; - dev->props.has_bfloat = [dev->mtl_device supportsFamily:MTLGPUFamilyMetal3_GGML]; - dev->props.has_bfloat |= [dev->mtl_device supportsFamily:MTLGPUFamilyApple6]; - if (getenv("GGML_METAL_BF16_DISABLE") != NULL) { - dev->props.has_bfloat = false; - } + dev->props.has_bfloat = [dev->mtl_device supportsFamily:MTLGPUFamilyMetal3_GGML]; + dev->props.has_bfloat |= [dev->mtl_device supportsFamily:MTLGPUFamilyApple6]; + if (getenv("GGML_METAL_BF16_DISABLE") != NULL) { + dev->props.has_bfloat = false; + } - dev->props.has_tensor = [dev->mtl_device supportsFamily:MTLGPUFamilyMetal4_GGML]; - if (getenv("GGML_METAL_TENSOR_DISABLE") != NULL) { - dev->props.has_tensor = false; - } - - // note: disable the tensor API by default for old chips because with the current implementation it is not useful - // - M2 Ultra: ~5% slower - // - M4, M4 Max: no significant difference - // - // TODO: try to update the tensor API kernels to at least match the simdgroup performance - if (getenv("GGML_METAL_TENSOR_ENABLE") == NULL && - ![[dev->mtl_device name] containsString:@"M5"] && - ![[dev->mtl_device name] containsString:@"M6"] && - ![[dev->mtl_device name] containsString:@"A19"] && - ![[dev->mtl_device name] containsString:@"A20"]) { - GGML_LOG_INFO("%s: tensor API disabled for pre-M5 and pre-A19 devices\n", __func__); - dev->props.has_tensor = false; - } - - // double-check that the tensor API compiles - if (dev->props.has_tensor) { - const char * src_tensor_f16 = "\n" - "#include \n" - "#include \n" - "#include \n" - " \n" - "using namespace metal; \n" - "using namespace mpp::tensor_ops; \n" - " \n" - "kernel void dummy_kernel( \n" - " tensor> A [[buffer(0)]], \n" - " tensor> B [[buffer(1)]], \n" - " device float * C [[buffer(2)]], \n" - " uint2 tgid [[threadgroup_position_in_grid]]) \n" - "{ \n" - " auto tA = A.slice(0, (int)tgid.y); \n" - " auto tB = B.slice((int)tgid.x, 0); \n" - " \n" - " matmul2d< \n" - " matmul2d_descriptor(16, 16, dynamic_extent), \n" - " execution_simdgroups<4>> mm; \n" - " \n" - " auto cT = mm.get_destination_cooperative_tensor(); \n" - " \n" - " auto sA = tA.slice(0, 0); \n" - " auto sB = tB.slice(0, 0); \n" - " mm.run(sB, sA, cT); \n" - " \n" - " auto tC = tensor, tensor_inline>(C, dextents(16, 16)); \n" - " \n" - " cT.store(tC); \n" - "}"; - - GGML_LOG_INFO("%s: testing tensor API for f16 support\n", __func__); - ggml_metal_library_t lib = ggml_metal_library_init_from_source(dev, src_tensor_f16, false); - if (lib == NULL) { - GGML_LOG_WARN("%s: - the tensor API is not supported in this environment - disabling\n", __func__); + dev->props.has_tensor = [dev->mtl_device supportsFamily:MTLGPUFamilyMetal4_GGML]; + if (getenv("GGML_METAL_TENSOR_DISABLE") != NULL) { dev->props.has_tensor = false; - } else { - struct ggml_metal_pipeline_with_params ppl = ggml_metal_library_compile_pipeline(lib, "dummy_kernel", "dummy_kernel", nil); - if (!ppl.pipeline) { + } + + // note: disable the tensor API by default for old chips because with the current implementation it is not useful + // - M2 Ultra: ~5% slower + // - M4, M4 Max: no significant difference + // + // TODO: try to update the tensor API kernels to at least match the simdgroup performance + if (getenv("GGML_METAL_TENSOR_ENABLE") == NULL && + ![[dev->mtl_device name] containsString:@"M5"] && + ![[dev->mtl_device name] containsString:@"M6"] && + ![[dev->mtl_device name] containsString:@"A19"] && + ![[dev->mtl_device name] containsString:@"A20"]) { + GGML_LOG_INFO("%s: tensor API disabled for pre-M5 and pre-A19 devices\n", __func__); + dev->props.has_tensor = false; + } + + // double-check that the tensor API compiles + if (dev->props.has_tensor) { + const char * src_tensor_f16 = "\n" + "#include \n" + "#include \n" + "#include \n" + " \n" + "using namespace metal; \n" + "using namespace mpp::tensor_ops; \n" + " \n" + "kernel void dummy_kernel( \n" + " tensor> A [[buffer(0)]], \n" + " tensor> B [[buffer(1)]], \n" + " device float * C [[buffer(2)]], \n" + " uint2 tgid [[threadgroup_position_in_grid]]) \n" + "{ \n" + " auto tA = A.slice(0, (int)tgid.y); \n" + " auto tB = B.slice((int)tgid.x, 0); \n" + " \n" + " matmul2d< \n" + " matmul2d_descriptor(16, 16, dynamic_extent), \n" + " execution_simdgroups<4>> mm; \n" + " \n" + " auto cT = mm.get_destination_cooperative_tensor(); \n" + " \n" + " auto sA = tA.slice(0, 0); \n" + " auto sB = tB.slice(0, 0); \n" + " mm.run(sB, sA, cT); \n" + " \n" + " auto tC = tensor, tensor_inline>(C, dextents(16, 16)); \n" + " \n" + " cT.store(tC); \n" + "}"; + + GGML_LOG_INFO("%s: testing tensor API for f16 support\n", __func__); + ggml_metal_library_t lib = ggml_metal_library_init_from_source(dev, src_tensor_f16, false); + if (lib == NULL) { GGML_LOG_WARN("%s: - the tensor API is not supported in this environment - disabling\n", __func__); dev->props.has_tensor = false; + } else { + struct ggml_metal_pipeline_with_params ppl = ggml_metal_library_compile_pipeline(lib, "dummy_kernel", "dummy_kernel", nil); + if (!ppl.pipeline) { + GGML_LOG_WARN("%s: - the tensor API is not supported in this environment - disabling\n", __func__); + dev->props.has_tensor = false; + } + + ggml_metal_library_free(lib); } - - ggml_metal_library_free(lib); } - } - // try to compile a dummy kernel to determine if the tensor API is supported for bfloat - if (dev->props.has_tensor && dev->props.has_bfloat) { - const char * src_tensor_bf16 = "\n" - "#include \n" - "#include \n" - "#include \n" - " \n" - "using namespace metal; \n" - "using namespace mpp::tensor_ops; \n" - " \n" - "kernel void dummy_kernel( \n" - " tensor> A [[buffer(0)]], \n" - " tensor> B [[buffer(1)]], \n" - " device float * C [[buffer(2)]], \n" - " uint2 tgid [[threadgroup_position_in_grid]]) \n" - "{ \n" - " auto tA = A.slice(0, (int)tgid.y); \n" - " auto tB = B.slice((int)tgid.x, 0); \n" - " \n" - " matmul2d< \n" - " matmul2d_descriptor(16, 16, dynamic_extent), \n" - " execution_simdgroups<4>> mm; \n" - " \n" - " auto cT = mm.get_destination_cooperative_tensor(); \n" - " \n" - " auto sA = tA.slice(0, 0); \n" - " auto sB = tB.slice(0, 0); \n" - " mm.run(sB, sA, cT); \n" - " \n" - " auto tC = tensor, tensor_inline>(C, dextents(16, 16)); \n" - " \n" - " cT.store(tC); \n" - "}"; + // try to compile a dummy kernel to determine if the tensor API is supported for bfloat + if (dev->props.has_tensor && dev->props.has_bfloat) { + const char * src_tensor_bf16 = "\n" + "#include \n" + "#include \n" + "#include \n" + " \n" + "using namespace metal; \n" + "using namespace mpp::tensor_ops; \n" + " \n" + "kernel void dummy_kernel( \n" + " tensor> A [[buffer(0)]], \n" + " tensor> B [[buffer(1)]], \n" + " device float * C [[buffer(2)]], \n" + " uint2 tgid [[threadgroup_position_in_grid]]) \n" + "{ \n" + " auto tA = A.slice(0, (int)tgid.y); \n" + " auto tB = B.slice((int)tgid.x, 0); \n" + " \n" + " matmul2d< \n" + " matmul2d_descriptor(16, 16, dynamic_extent), \n" + " execution_simdgroups<4>> mm; \n" + " \n" + " auto cT = mm.get_destination_cooperative_tensor(); \n" + " \n" + " auto sA = tA.slice(0, 0); \n" + " auto sB = tB.slice(0, 0); \n" + " mm.run(sB, sA, cT); \n" + " \n" + " auto tC = tensor, tensor_inline>(C, dextents(16, 16)); \n" + " \n" + " cT.store(tC); \n" + "}"; - GGML_LOG_INFO("%s: testing tensor API for bfloat support\n", __func__); - ggml_metal_library_t lib = ggml_metal_library_init_from_source(dev, src_tensor_bf16, false); - if (lib == NULL) { - GGML_LOG_WARN("%s: - the tensor API does not support bfloat - disabling bfloat support\n", __func__); - dev->props.has_bfloat = false; - } else { - struct ggml_metal_pipeline_with_params ppl = ggml_metal_library_compile_pipeline(lib, "dummy_kernel", "dummy_kernel", nil); - if (!ppl.pipeline) { + GGML_LOG_INFO("%s: testing tensor API for bfloat support\n", __func__); + ggml_metal_library_t lib = ggml_metal_library_init_from_source(dev, src_tensor_bf16, false); + if (lib == NULL) { GGML_LOG_WARN("%s: - the tensor API does not support bfloat - disabling bfloat support\n", __func__); dev->props.has_bfloat = false; + } else { + struct ggml_metal_pipeline_with_params ppl = ggml_metal_library_compile_pipeline(lib, "dummy_kernel", "dummy_kernel", nil); + if (!ppl.pipeline) { + GGML_LOG_WARN("%s: - the tensor API does not support bfloat - disabling bfloat support\n", __func__); + dev->props.has_bfloat = false; + } + + ggml_metal_library_free(lib); } - - ggml_metal_library_free(lib); } - } - dev->props.use_residency_sets = true; + dev->props.use_residency_sets = true; #if defined(GGML_METAL_HAS_RESIDENCY_SETS) - dev->props.use_residency_sets = getenv("GGML_METAL_NO_RESIDENCY") == nil; + dev->props.use_residency_sets = getenv("GGML_METAL_NO_RESIDENCY") == nil; #endif - dev->props.use_shared_buffers = dev->props.has_unified_memory; + dev->props.use_shared_buffers = dev->props.has_unified_memory; #if TARGET_OS_OSX - // In case of eGPU, shared memory may be preferable. - dev->props.use_shared_buffers |= [dev->mtl_device location] == MTLDeviceLocationExternal; + // In case of eGPU, shared memory may be preferable. + dev->props.use_shared_buffers |= [dev->mtl_device location] == MTLDeviceLocationExternal; #endif - if (getenv("GGML_METAL_SHARED_BUFFERS_DISABLE") != NULL) { - dev->props.use_shared_buffers = false; - } - if (getenv("GGML_METAL_SHARED_BUFFERS_ENABLE") != NULL) { - dev->props.use_shared_buffers = true; - } + if (getenv("GGML_METAL_SHARED_BUFFERS_DISABLE") != NULL) { + dev->props.use_shared_buffers = false; + } + if (getenv("GGML_METAL_SHARED_BUFFERS_ENABLE") != NULL) { + dev->props.use_shared_buffers = true; + } - dev->props.supports_gpu_family_apple7 = [dev->mtl_device supportsFamily:MTLGPUFamilyApple7]; + dev->props.supports_gpu_family_apple7 = [dev->mtl_device supportsFamily:MTLGPUFamilyApple7]; - dev->props.device_id = ggml_metal_device_id_parse([[dev->mtl_device name] UTF8String]); + dev->props.device_id = ggml_metal_device_id_parse([[dev->mtl_device name] UTF8String]); - dev->props.op_offload_min_batch_size = getenv("GGML_OP_OFFLOAD_MIN_BATCH") ? atoi(getenv("GGML_OP_OFFLOAD_MIN_BATCH")) : 32; + dev->props.op_offload_min_batch_size = getenv("GGML_OP_OFFLOAD_MIN_BATCH") ? atoi(getenv("GGML_OP_OFFLOAD_MIN_BATCH")) : 32; - dev->props.max_buffer_size = dev->mtl_device.maxBufferLength; - dev->props.max_theadgroup_memory_size = dev->mtl_device.maxThreadgroupMemoryLength; - if (@available(macOS 10.12, iOS 16.0, *)) { - dev->props.max_working_set_size = dev->mtl_device.recommendedMaxWorkingSetSize; - } else { - dev->props.max_working_set_size = dev->mtl_device.maxBufferLength; - } + dev->props.max_buffer_size = dev->mtl_device.maxBufferLength; + dev->props.max_theadgroup_memory_size = dev->mtl_device.maxThreadgroupMemoryLength; + if (@available(macOS 10.12, iOS 16.0, *)) { + dev->props.max_working_set_size = dev->mtl_device.recommendedMaxWorkingSetSize; + } else { + dev->props.max_working_set_size = dev->mtl_device.maxBufferLength; + } - snprintf(dev->props.name, sizeof(dev->props.name), "%s%d", "MTL", device); - const char * gpu_name = [[dev->mtl_device name] UTF8String]; - if (n_devices > 1) { - snprintf(dev->props.desc, sizeof(dev->props.desc), "%s (dev p%d/v%d)", - gpu_name, dev->props.device_phys, dev->props.device_virt); - } else { - snprintf(dev->props.desc, sizeof(dev->props.desc), "%s", gpu_name); - } + snprintf(dev->props.name, sizeof(dev->props.name), "%s%d", "MTL", device); + const char * gpu_name = [[dev->mtl_device name] UTF8String]; + if (n_devices > 1) { + snprintf(dev->props.desc, sizeof(dev->props.desc), "%s (dev p%d/v%d)", + gpu_name, dev->props.device_phys, dev->props.device_virt); + } else { + snprintf(dev->props.desc, sizeof(dev->props.desc), "%s", gpu_name); + } - dev->library = ggml_metal_library_init(dev); - if (!dev->library) { - GGML_LOG_ERROR("%s: error: failed to create library\n", __func__); - } + dev->library = ggml_metal_library_init(dev); + if (!dev->library) { + GGML_LOG_ERROR("%s: error: failed to create library\n", __func__); + } - if (dev->props.use_residency_sets) { - dev->rsets = ggml_metal_rsets_init(dev); - } else { - dev->rsets = nil; - } + if (dev->props.use_residency_sets) { + dev->rsets = ggml_metal_rsets_init(dev); + } else { + dev->rsets = nil; + } - // print MTL GPU family: - GGML_LOG_INFO("%s: GPU name: %s (%s)\n", __func__, dev->props.name, dev->props.desc); + // print MTL GPU family: + GGML_LOG_INFO("%s: GPU name: %s (%s)\n", __func__, dev->props.name, dev->props.desc); - // determine max supported GPU family - // https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf - // https://developer.apple.com/metal/Metal-Feature-Set-Tables.pdf - { - for (int i = MTLGPUFamilyApple1 + 20; i >= MTLGPUFamilyApple1; --i) { - if ([dev->mtl_device supportsFamily:i]) { - dev->props.gpu_family = i - (int) MTLGPUFamilyApple1 + 1; - GGML_LOG_INFO("%s: GPU family: MTLGPUFamilyApple%d (%d)\n", __func__, dev->props.gpu_family, i); - break; + // determine max supported GPU family + // https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf + // https://developer.apple.com/metal/Metal-Feature-Set-Tables.pdf + { + for (int i = MTLGPUFamilyApple1 + 20; i >= MTLGPUFamilyApple1; --i) { + if ([dev->mtl_device supportsFamily:i]) { + dev->props.gpu_family = i - (int) MTLGPUFamilyApple1 + 1; + GGML_LOG_INFO("%s: GPU family: MTLGPUFamilyApple%d (%d)\n", __func__, dev->props.gpu_family, i); + break; + } + } + + for (int i = MTLGPUFamilyCommon1 + 5; i >= MTLGPUFamilyCommon1; --i) { + if ([dev->mtl_device supportsFamily:i]) { + GGML_LOG_INFO("%s: GPU family: MTLGPUFamilyCommon%d (%d)\n", __func__, i - (int) MTLGPUFamilyCommon1 + 1, i); + break; + } + } + + for (int i = MTLGPUFamilyMetal3_GGML + 5; i >= MTLGPUFamilyMetal3_GGML; --i) { + if ([dev->mtl_device supportsFamily:i]) { + GGML_LOG_INFO("%s: GPU family: MTLGPUFamilyMetal%d (%d)\n", __func__, i - (int) MTLGPUFamilyMetal3_GGML + 3, i); + break; + } } } - for (int i = MTLGPUFamilyCommon1 + 5; i >= MTLGPUFamilyCommon1; --i) { - if ([dev->mtl_device supportsFamily:i]) { - GGML_LOG_INFO("%s: GPU family: MTLGPUFamilyCommon%d (%d)\n", __func__, i - (int) MTLGPUFamilyCommon1 + 1, i); - break; - } - } - - for (int i = MTLGPUFamilyMetal3_GGML + 5; i >= MTLGPUFamilyMetal3_GGML; --i) { - if ([dev->mtl_device supportsFamily:i]) { - GGML_LOG_INFO("%s: GPU family: MTLGPUFamilyMetal%d (%d)\n", __func__, i - (int) MTLGPUFamilyMetal3_GGML + 3, i); - break; - } - } - } - - GGML_LOG_INFO("%s: simdgroup reduction = %s\n", __func__, dev->props.has_simdgroup_reduction ? "true" : "false"); - GGML_LOG_INFO("%s: simdgroup matrix mul. = %s\n", __func__, dev->props.has_simdgroup_mm ? "true" : "false"); - GGML_LOG_INFO("%s: has unified memory = %s\n", __func__, dev->props.has_unified_memory ? "true" : "false"); - GGML_LOG_INFO("%s: has bfloat = %s\n", __func__, dev->props.has_bfloat ? "true" : "false"); - GGML_LOG_INFO("%s: has tensor = %s\n", __func__, dev->props.has_tensor ? "true" : "false"); - GGML_LOG_INFO("%s: use residency sets = %s\n", __func__, dev->props.use_residency_sets ? "true" : "false"); - GGML_LOG_INFO("%s: use shared buffers = %s\n", __func__, dev->props.use_shared_buffers ? "true" : "false"); + GGML_LOG_INFO("%s: simdgroup reduction = %s\n", __func__, dev->props.has_simdgroup_reduction ? "true" : "false"); + GGML_LOG_INFO("%s: simdgroup matrix mul. = %s\n", __func__, dev->props.has_simdgroup_mm ? "true" : "false"); + GGML_LOG_INFO("%s: has unified memory = %s\n", __func__, dev->props.has_unified_memory ? "true" : "false"); + GGML_LOG_INFO("%s: has bfloat = %s\n", __func__, dev->props.has_bfloat ? "true" : "false"); + GGML_LOG_INFO("%s: has tensor = %s\n", __func__, dev->props.has_tensor ? "true" : "false"); + GGML_LOG_INFO("%s: use residency sets = %s\n", __func__, dev->props.use_residency_sets ? "true" : "false"); + GGML_LOG_INFO("%s: use shared buffers = %s\n", __func__, dev->props.use_shared_buffers ? "true" : "false"); #if TARGET_OS_OSX || (TARGET_OS_IOS && __clang_major__ >= 15) - if (@available(macOS 10.12, iOS 16.0, *)) { - GGML_LOG_INFO("%s: recommendedMaxWorkingSetSize = %8.2f MB\n", __func__, dev->props.max_working_set_size / 1e6); - } + if (@available(macOS 10.12, iOS 16.0, *)) { + GGML_LOG_INFO("%s: recommendedMaxWorkingSetSize = %8.2f MB\n", __func__, dev->props.max_working_set_size / 1e6); + } #endif + } } } From f29551215b71778566afef659c8d4b672d7d49ef Mon Sep 17 00:00:00 2001 From: Xuan-Son Nguyen Date: Thu, 27 Aug 2026 12:11:12 +0200 Subject: [PATCH 2/8] args: add --video-* CLI arguments (#24318) * args: add --video-* CLI arguments * gen docs * nits * add mtmd_helper_init_opt --- common/arg.cpp | 21 ++++++++++++++++++ common/common.h | 5 +++++ tools/cli/README.md | 3 +++ tools/mtmd/mtmd-cli.cpp | 11 +++++++++- tools/mtmd/mtmd-helper.cpp | 28 ++++++++++++++++-------- tools/mtmd/mtmd-helper.h | 38 ++++++++++++++++++++++++--------- tools/server/README.md | 3 +++ tools/server/server-common.cpp | 28 ++++++++++++++---------- tools/server/server-common.h | 14 +++++++++--- tools/server/server-context.cpp | 31 ++++++++++++++++----------- tools/server/server-context.h | 2 +- tools/tts/tts.cpp | 2 +- 12 files changed, 138 insertions(+), 48 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index aad8266ed7..dd92e4720b 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -2644,6 +2644,27 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.mtmd_batch_max_tokens = value; } ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_MTMD_BATCH_MAX_TOKENS")); + add_opt(common_arg( + {"--video-fps"}, "N", + string_format("target video frame rate (default: %.1f)", params.video_fps), + [](common_params & params, const std::string & value) { + params.video_fps = std::stof(value); + } + ).set_examples(mmproj_examples).set_env("LLAMA_ARG_VIDEO_FPS")); + add_opt(common_arg( + {"--video-timestamp-interval"}, "N", + string_format("interval in milliseconds between text timestamps (default: %" PRId64 ")", params.video_timestamp_interval_ms), + [](common_params & params, int value) { + params.video_timestamp_interval_ms = value; + } + ).set_examples(mmproj_examples).set_env("LLAMA_ARG_VIDEO_TIMESTAMP_INTERVAL")); + add_opt(common_arg( + {"--video-ffmpeg-dir"}, "DIR", + "path to the directory containing ffmpeg and ffprobe (default: search in PATH)", + [](common_params & params, const std::string & value) { + params.video_ffmpeg_bin_dir = value; + } + ).set_examples(mmproj_examples).set_env("LLAMA_ARG_VIDEO_FFMPEG_DIR")); if (params.is_gen_docs || llama_supports_rpc()) { add_opt(common_arg( {"--rpc"}, "SERVERS", diff --git a/common/common.h b/common/common.h index 51518f343a..9593e10d9a 100644 --- a/common/common.h +++ b/common/common.h @@ -590,6 +590,11 @@ struct common_params { int image_max_tokens = -1; int mtmd_batch_max_tokens = 1024; + // for video input + float video_fps = 4.0f; + int64_t video_timestamp_interval_ms = 5000; + std::string video_ffmpeg_bin_dir = ""; + // finetune struct lr_opt lr; enum ggml_opt_optimizer_type optimizer = GGML_OPT_OPTIMIZER_TYPE_ADAMW; diff --git a/tools/cli/README.md b/tools/cli/README.md index c9cbacafcd..0fba70a90e 100644 --- a/tools/cli/README.md +++ b/tools/cli/README.md @@ -166,6 +166,9 @@ | `--image, --audio, --video FILE` | path to an image, audio, or video file. use with multimodal models, use comma-separated values for multiple files | | `--image-min-tokens N` | minimum number of tokens each image can take, only used by vision models with dynamic resolution (default: read from model)
(env: LLAMA_ARG_IMAGE_MIN_TOKENS) | | `--image-max-tokens N` | maximum number of tokens each image can take, only used by vision models with dynamic resolution (default: read from model)
(env: LLAMA_ARG_IMAGE_MAX_TOKENS) | +| `--video-fps N` | target video frame rate (default: 4.0)
(env: LLAMA_ARG_VIDEO_FPS) | +| `--video-timestamp-interval N` | interval in milliseconds between text timestamps (default: 5000)
(env: LLAMA_ARG_VIDEO_TIMESTAMP_INTERVAL) | +| `--video-ffmpeg-dir DIR` | path to the directory containing ffmpeg and ffprobe (default: search in PATH)
(env: LLAMA_ARG_VIDEO_FFMPEG_DIR) | | `-o, --output, --output-file FNAME` | output file (default: '') | | `--chat-template-kwargs STRING` | sets additional params for the json template parser, must be a valid json object string, e.g. '{"key1":"value1","key2":"value2"}'
(env: LLAMA_ARG_CHAT_TEMPLATE_KWARGS) | | `--jinja, --no-jinja` | whether to use jinja template engine for chat (default: enabled)
(env: LLAMA_ARG_JINJA) | diff --git a/tools/mtmd/mtmd-cli.cpp b/tools/mtmd/mtmd-cli.cpp index f6c787fdb6..97678c6b2e 100644 --- a/tools/mtmd/mtmd-cli.cpp +++ b/tools/mtmd/mtmd-cli.cpp @@ -87,6 +87,9 @@ struct mtmd_cli_context { mtmd::bitmaps bitmaps; std::vector videos; + mtmd_helper_init_opt init_opt = mtmd_helper_init_opt_default(); + std::string video_ffmpeg_bin_dir; + mtmd::batch_ptr mbatch; // chat template @@ -170,6 +173,12 @@ struct mtmd_cli_context { LOG_ERR("Failed to load vision model from %s\n", clip_path); exit(1); } + + video_ffmpeg_bin_dir = params.video_ffmpeg_bin_dir; + init_opt.video_params.fps_target = params.video_fps; + init_opt.video_params.timestamp_interval_ms = params.video_timestamp_interval_ms; + init_opt.video_params.ffmpeg_bin_dir = video_ffmpeg_bin_dir.empty() + ? nullptr : video_ffmpeg_bin_dir.c_str(); } bool check_antiprompt(const llama_tokens & generated_tokens) { @@ -184,7 +193,7 @@ struct mtmd_cli_context { } bool load_media(const std::string & fname) { - auto res = mtmd_helper_bitmap_init_from_file(ctx_vision.get(), fname.c_str(), false); + auto res = mtmd_helper_bitmap_init_from_file(ctx_vision.get(), fname.c_str(), false, init_opt); if (!res.bitmap) { return false; } diff --git a/tools/mtmd/mtmd-helper.cpp b/tools/mtmd/mtmd-helper.cpp index f1defb6477..77f9d58fe2 100644 --- a/tools/mtmd/mtmd-helper.cpp +++ b/tools/mtmd/mtmd-helper.cpp @@ -369,14 +369,18 @@ static bool is_webp_file(const unsigned char * buf, size_t len) { } #ifdef MTMD_VIDEO -static mtmd_bitmap * decode_webp_with_ffmpeg(mtmd_context * mctx, const unsigned char * buf, size_t len, bool placeholder); +static mtmd_bitmap * decode_webp_with_ffmpeg(mtmd_context * mctx, const unsigned char * buf, size_t len, bool placeholder, + const mtmd_helper_video_init_params & params); #endif -mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_buf(mtmd_context * ctx, const unsigned char * buf, size_t len, bool placeholder) { +mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_buf(mtmd_context * ctx, const unsigned char * buf, size_t len, bool placeholder, + mtmd_helper_init_opt opt) { // calculate the hash if needed std::string id; mtmd_bitmap * result = nullptr; + GGML_UNUSED(opt); // only used by video code paths + if (!placeholder) { // use sha256 to prevent cache poisoning id = hash_sha256_hex(buf, len); @@ -414,7 +418,7 @@ mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_buf(mtmd_context * ctx, #ifdef MTMD_VIDEO // stb_image does not support webp; decode it with ffmpeg as a single frame if (!result && is_webp_file(buf, len)) { - result = decode_webp_with_ffmpeg(ctx, buf, len, placeholder); + result = decode_webp_with_ffmpeg(ctx, buf, len, placeholder, opt.video_params); if (!result) { LOG_ERR("%s: failed to decode webp buffer\n", __func__); return {nullptr, nullptr}; @@ -427,8 +431,7 @@ mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_buf(mtmd_context * ctx, // 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); + auto video_ctx = mtmd_helper_video_init_from_buf(ctx, buf, len, opt.video_params); if (!video_ctx) { LOG_ERR("%s: failed to decode buffer as either image/audio/video\n", __func__); return {nullptr, nullptr}; @@ -456,7 +459,8 @@ mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_buf(mtmd_context * ctx, return {nullptr, nullptr}; } -mtmd_helper_bitmap_wrapper 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, + mtmd_helper_init_opt opt) { #ifdef _WIN32 int wlen = MultiByteToWideChar(CP_UTF8, 0, fname, -1, NULL, 0); if (!wlen) { @@ -497,7 +501,7 @@ mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_file(mtmd_context * ctx, return {nullptr, nullptr}; } - return mtmd_helper_bitmap_init_from_buf(ctx, buf.data(), buf.size(), placeholder); + return mtmd_helper_bitmap_init_from_buf(ctx, buf.data(), buf.size(), placeholder, opt); } bool mtmd_helper_support_video(mtmd_context * ctx) { @@ -855,6 +859,12 @@ mtmd_helper_video_init_params mtmd_helper_video_init_params_default() { }; } +mtmd_helper_init_opt mtmd_helper_init_opt_default() { + return { + /* video_params */ mtmd_helper_video_init_params_default(), + }; +} + 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 @@ -876,8 +886,8 @@ static std::string video_resolve_bin(const char * bin_dir, const char * name) { } #ifdef MTMD_VIDEO -static mtmd_bitmap * decode_webp_with_ffmpeg(mtmd_context * mctx, const unsigned char * buf, size_t len, bool placeholder) { - auto params = mtmd_helper_video_init_params_default(); +static mtmd_bitmap * decode_webp_with_ffmpeg(mtmd_context * mctx, const unsigned char * buf, size_t len, bool placeholder, + const mtmd_helper_video_init_params & params) { mtmd_helper_video vctx; vctx.mctx = mctx; vctx.input_buf.assign(buf, buf + len); diff --git a/tools/mtmd/mtmd-helper.h b/tools/mtmd/mtmd-helper.h index 58dfb15250..772e0f091b 100644 --- a/tools/mtmd/mtmd-helper.h +++ b/tools/mtmd/mtmd-helper.h @@ -23,6 +23,23 @@ extern "C" { struct mtmd_helper_video; typedef struct mtmd_helper_video mtmd_helper_video; +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); + +// opt for mtmd_helper_bitmap_init_from_*() +struct mtmd_helper_init_opt { + struct mtmd_helper_video_init_params video_params; +}; +typedef struct mtmd_helper_init_opt mtmd_helper_init_opt; + +MTMD_API struct mtmd_helper_init_opt mtmd_helper_init_opt_default(void); + // 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 @@ -40,7 +57,11 @@ struct mtmd_helper_bitmap_wrapper { // it calls mtmd_helper_bitmap_init_from_buf() internally // returns nullptr on failure // this function is thread-safe -MTMD_API struct mtmd_helper_bitmap_wrapper 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, + struct mtmd_helper_init_opt opt); // helper function to construct a mtmd_bitmap from a buffer containing a file // supported formats: @@ -53,7 +74,11 @@ MTMD_API struct mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_file(mtm // - output bitmap will have SHA-256 hash (hex string) as the ID // returns nullptr on failure // this function is thread-safe -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); +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, + struct mtmd_helper_init_opt opt); // 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); @@ -124,14 +149,7 @@ struct mtmd_helper_video_info { 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); +// note: mtmd_helper_video_init_params is defined at the top, as it is part of mtmd_helper_init_opt // returns NULL on failure (ffprobe not found, file unreadable, etc.) MTMD_API mtmd_helper_video * mtmd_helper_video_init( diff --git a/tools/server/README.md b/tools/server/README.md index 93736c3edf..f49cdb2729 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -182,6 +182,9 @@ For the full list of features, please refer to [server's changelog](https://gith | `--image-min-tokens N` | minimum number of tokens each image can take, only used by vision models with dynamic resolution (default: read from model)
(env: LLAMA_ARG_IMAGE_MIN_TOKENS) | | `--image-max-tokens N` | maximum number of tokens each image can take, only used by vision models with dynamic resolution (default: read from model)
(env: LLAMA_ARG_IMAGE_MAX_TOKENS) | | `--mtmd-batch-max-tokens N` | maximum number of image tokens per batch when encoding images (default: 1024)
(env: LLAMA_ARG_MTMD_BATCH_MAX_TOKENS) | +| `--video-fps N` | target video frame rate (default: 4.0)
(env: LLAMA_ARG_VIDEO_FPS) | +| `--video-timestamp-interval N` | interval in milliseconds between text timestamps (default: 5000)
(env: LLAMA_ARG_VIDEO_TIMESTAMP_INTERVAL) | +| `--video-ffmpeg-dir DIR` | path to the directory containing ffmpeg and ffprobe (default: search in PATH)
(env: LLAMA_ARG_VIDEO_FFMPEG_DIR) | | `-a, --alias STRING` | set model name aliases, comma-separated (to be used by API)
(env: LLAMA_ARG_ALIAS) | | `--tags STRING` | set model tags, comma-separated (informational, not used for routing)
(env: LLAMA_ARG_TAGS) | | `--embd-normalize N` | normalisation for embeddings (default: 2) (-1=none, 0=max absolute int16, 1=taxicab, 2=euclidean, >2=p-norm) | diff --git a/tools/server/server-common.cpp b/tools/server/server-common.cpp index 7997d4016a..c30955e89f 100644 --- a/tools/server/server-common.cpp +++ b/tools/server/server-common.cpp @@ -910,12 +910,17 @@ size_t validate_utf8(const std::string& text) { return len; } -server_tokens process_mtmd_prompt(mtmd_context * mctx, const std::string & prompt, const std::vector & files, bool is_placeholder) { +server_tokens process_mtmd_prompt( + mtmd_context * mctx, + const std::string & prompt, + const std::vector & files, + const mtmd_helper_init_opt & init_opt, + bool is_placeholder) { // these will be freed upon going out of scope mtmd::bitmaps bitmaps; std::vector videos; for (auto & file : files) { - auto out = mtmd_helper_bitmap_init_from_buf(mctx, file.data(), file.size(), is_placeholder); + auto out = mtmd_helper_bitmap_init_from_buf(mctx, file.data(), file.size(), is_placeholder, init_opt); if (!out.bitmap) { throw std::runtime_error("Failed to load image or audio file"); } @@ -956,7 +961,7 @@ server_tokens process_mtmd_prompt(mtmd_context * mctx, const std::string & promp * - "prompt": [12, 34, "string", 56, 78] * - "prompt": { "prompt_string": "string", "multimodal_data": [ "base64" ] } */ -static server_tokens tokenize_input_subprompt(const llama_vocab * vocab, mtmd_context * mctx, const json & json_prompt, bool add_special, bool parse_special) { +static server_tokens tokenize_input_subprompt(const llama_vocab * vocab, mtmd_context * mctx, const json & json_prompt, bool add_special, bool parse_special, const mtmd_helper_init_opt & init_opt) { constexpr char JSON_STRING_PROMPT_KEY[] = "prompt_string"; constexpr char JSON_MTMD_DATA_KEY[] = "multimodal_data"; const bool has_mtmd = mctx != nullptr; @@ -979,7 +984,7 @@ static server_tokens tokenize_input_subprompt(const llama_vocab * vocab, mtmd_co for (const auto & entry : json_prompt.at(JSON_MTMD_DATA_KEY)) { files.push_back(base64_decode(entry)); } - return process_mtmd_prompt(mctx, json_prompt.at(JSON_STRING_PROMPT_KEY), files); + return process_mtmd_prompt(mctx, json_prompt.at(JSON_STRING_PROMPT_KEY), files, init_opt); } else { // Not multimodal, but contains a subobject. llama_tokens tmp = tokenize_mixed(vocab, json_prompt.at(JSON_STRING_PROMPT_KEY), add_special, parse_special); @@ -990,15 +995,15 @@ static server_tokens tokenize_input_subprompt(const llama_vocab * vocab, mtmd_co } } -std::vector tokenize_input_prompts(const llama_vocab * vocab, mtmd_context * mctx, const json & json_prompt, bool add_special, bool parse_special) { +std::vector tokenize_input_prompts(const llama_vocab * vocab, mtmd_context * mctx, const json & json_prompt, bool add_special, bool parse_special, const mtmd_helper_init_opt & init_opt) { std::vector result; if (json_prompt.is_array() && !json_is_array_and_contains_numbers(json_prompt)) { result.reserve(json_prompt.size()); for (const auto & p : json_prompt) { - result.push_back(tokenize_input_subprompt(vocab, mctx, p,add_special, parse_special)); + result.push_back(tokenize_input_subprompt(vocab, mctx, p, add_special, parse_special, init_opt)); } } else { - result.push_back(tokenize_input_subprompt(vocab, mctx, json_prompt, add_special, parse_special)); + result.push_back(tokenize_input_subprompt(vocab, mctx, json_prompt, add_special, parse_special, init_opt)); } if (result.empty()) { throw std::runtime_error("\"prompt\" must not be empty"); @@ -1787,7 +1792,8 @@ server_tokens format_prompt_rerank( const struct llama_vocab * vocab, mtmd_context * mctx, const std::string & query, - const std::string & doc) { + const std::string & doc, + const mtmd_helper_init_opt & init_opt) { server_tokens result = {}; const char * rerank_prompt = llama_model_chat_template(model, "rerank"); @@ -1796,12 +1802,12 @@ server_tokens format_prompt_rerank( std::string prompt = rerank_prompt; string_replace_all(prompt, "{query}" , query); string_replace_all(prompt, "{document}", doc ); - server_tokens tokens = tokenize_input_subprompt(vocab, mctx, prompt, false, true); + server_tokens tokens = tokenize_input_subprompt(vocab, mctx, prompt, false, true, init_opt); result.push_back(tokens); } else { // Get EOS token - use SEP token as fallback if EOS is not available - server_tokens query_tokens = tokenize_input_subprompt(vocab, mctx, query, false, false); - server_tokens doc_tokens = tokenize_input_subprompt(vocab, mctx, doc, false, false); + server_tokens query_tokens = tokenize_input_subprompt(vocab, mctx, query, false, false, init_opt); + server_tokens doc_tokens = tokenize_input_subprompt(vocab, mctx, doc, false, false, init_opt); llama_token eos_token = llama_vocab_eos(vocab); if (eos_token == LLAMA_TOKEN_NULL) { eos_token = llama_vocab_sep(vocab); diff --git a/tools/server/server-common.h b/tools/server/server-common.h index f8ea82ef4c..6c681a2cf5 100644 --- a/tools/server/server-common.h +++ b/tools/server/server-common.h @@ -5,6 +5,7 @@ #include "llama.h" #include "chat.h" #include "mtmd.h" +#include "mtmd-helper.h" #include "json.h" @@ -269,7 +270,12 @@ size_t validate_utf8(const std::string& text); // process mtmd prompt, return the server_tokens containing both text tokens and media chunks // if is_placeholder is true, the media chunk will be treated as placeholder for counting tokens; the output tokens are not usable for actual inference (e.g. for submitting a task to server_queue) -server_tokens process_mtmd_prompt(mtmd_context * mctx, const std::string & prompt, const std::vector & files, bool is_placeholder = false); +server_tokens process_mtmd_prompt( + mtmd_context * mctx, + const std::string & prompt, + const std::vector & files, + const mtmd_helper_init_opt & init_opt, + bool is_placeholder = false); /** * break the input "prompt" object into multiple prompt if needed, then tokenize them @@ -289,7 +295,8 @@ std::vector tokenize_input_prompts( mtmd_context * mctx, const json & json_prompt, bool add_special, - bool parse_special); + bool parse_special, + const mtmd_helper_init_opt & init_opt); // // OAI utils @@ -538,7 +545,8 @@ server_tokens format_prompt_rerank( const struct llama_vocab * vocab, mtmd_context * mctx, const std::string & query, - const std::string & doc); + const std::string & doc, + const mtmd_helper_init_opt & init_opt); // simple implementation of a pipe // used for streaming data between threads diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index a9edbd7be8..9fdfcae563 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -794,6 +794,8 @@ public: llama_model * model_tgt = nullptr; mtmd_context * mctx = nullptr; + // note: video_params.ffmpeg_bin_dir points into params_base, which outlives this struct + mtmd_helper_init_opt init_opt = mtmd_helper_init_opt_default(); const llama_vocab * vocab = nullptr; server_queue queue_tasks; @@ -1118,6 +1120,11 @@ private: } SRV_INF("loaded multimodal model, '%s'\n", mmproj_path.c_str()); + init_opt.video_params.fps_target = params_base.video_fps; + init_opt.video_params.timestamp_interval_ms = params_base.video_timestamp_interval_ms; + init_opt.video_params.ffmpeg_bin_dir = params_base.video_ffmpeg_bin_dir.empty() + ? nullptr : params_base.video_ffmpeg_bin_dir.c_str(); + if (params_base.ctx_shift) { params_base.ctx_shift = false; SRV_WRN("%s\n", "ctx_shift is not supported by multimodal, it will be disabled"); @@ -2134,9 +2141,9 @@ private: try { auto & prompt = task.cli_prompt; if (mctx != nullptr) { - task.tokens = process_mtmd_prompt(mctx, prompt, task.cli_files); + task.tokens = process_mtmd_prompt(mctx, prompt, task.cli_files, init_opt); } else { - task.tokens = std::move(tokenize_input_prompts(vocab, mctx, prompt, true, true)[0]); + task.tokens = std::move(tokenize_input_prompts(vocab, mctx, prompt, true, true, init_opt)[0]); } task.cli_prompt.clear(); task.cli_files.clear(); @@ -4165,10 +4172,10 @@ std::unique_ptr server_routes::handle_completions_impl( if (res_type != TASK_RESPONSE_TYPE_NONE && ctx_server.mctx != nullptr) { // This is the case used by OAI compatible chat path with MTMD. TODO It can be moved to the path below. - inputs.push_back(process_mtmd_prompt(ctx_server.mctx, prompt.get(), files)); + inputs.push_back(process_mtmd_prompt(ctx_server.mctx, prompt.get(), files, ctx_server.init_opt)); } else { // Everything else, including multimodal completions. - inputs = tokenize_input_prompts(ctx_server.vocab, ctx_server.mctx, prompt, true, true); + inputs = tokenize_input_prompts(ctx_server.vocab, ctx_server.mctx, prompt, true, true, ctx_server.init_opt); } // tasks.reserve(inputs.size()); // TODO: this is inaccurate due to child tasks @@ -4752,7 +4759,7 @@ void server_routes::init_routes() { data["input_extra"] = input_extra; // default to empty array if it's not exist std::string prompt = json_value(data, "prompt", std::string()); - std::vector tokenized_prompts = tokenize_input_prompts(ctx_server.vocab, ctx_server.mctx, prompt, false, true); + std::vector tokenized_prompts = tokenize_input_prompts(ctx_server.vocab, ctx_server.mctx, prompt, false, true, ctx_server.init_opt); SRV_DBG("creating infill tasks, n_prompts = %d\n", (int) tokenized_prompts.size()); data["prompt"] = format_prompt_infill( ctx_server.vocab, @@ -4816,7 +4823,7 @@ void server_routes::init_routes() { }; this->post_chat_completions_tok = [this](const server_http_req & req) { - return handle_count_tokens(ctx_server.vocab, ctx_server.mctx, req, TASK_RESPONSE_TYPE_OAI_CHAT); + return handle_count_tokens(ctx_server.vocab, ctx_server.mctx, ctx_server.init_opt, req, TASK_RESPONSE_TYPE_OAI_CHAT); }; this->post_control = [this](const server_http_req & req) { @@ -4875,7 +4882,7 @@ void server_routes::init_routes() { }; this->post_responses_tok_oai = [this](const server_http_req & req) { - return handle_count_tokens(ctx_server.vocab, ctx_server.mctx, req, TASK_RESPONSE_TYPE_OAI_RESP); + return handle_count_tokens(ctx_server.vocab, ctx_server.mctx, ctx_server.init_opt, req, TASK_RESPONSE_TYPE_OAI_RESP); }; this->post_transcriptions_oai = [this](const server_http_req & req) { @@ -4925,7 +4932,7 @@ void server_routes::init_routes() { }; this->post_anthropic_count_tokens = [this](const server_http_req & req) { - return handle_count_tokens(ctx_server.vocab, ctx_server.mctx, req, TASK_RESPONSE_TYPE_ANTHROPIC); + return handle_count_tokens(ctx_server.vocab, ctx_server.mctx, ctx_server.init_opt, req, TASK_RESPONSE_TYPE_ANTHROPIC); }; // same with handle_chat_completions, but without inference part @@ -5058,7 +5065,7 @@ void server_routes::init_routes() { std::vector tasks; tasks.reserve(documents.size()); for (size_t i = 0; i < documents.size(); i++) { - auto tmp = format_prompt_rerank(ctx_server.model_tgt, ctx_server.vocab, ctx_server.mctx, query, documents[i]); + auto tmp = format_prompt_rerank(ctx_server.model_tgt, ctx_server.vocab, ctx_server.mctx, query, documents[i], ctx_server.init_opt); server_task task = server_task(SERVER_TASK_TYPE_RERANK); task.id = rd.get_new_id(); task.tokens = std::move(tmp); @@ -5296,7 +5303,7 @@ std::unique_ptr server_routes::handle_embeddings_impl(cons } } - auto tokenized_prompts = tokenize_input_prompts(ctx_server.vocab, ctx_server.mctx, prompt, true, true); + auto tokenized_prompts = tokenize_input_prompts(ctx_server.vocab, ctx_server.mctx, prompt, true, true, ctx_server.init_opt); for (const auto & tokens : tokenized_prompts) { // this check is necessary for models that do not add BOS token to the input if (tokens.empty()) { @@ -5357,7 +5364,7 @@ std::unique_ptr server_routes::handle_embeddings_impl(cons return res; } -std::unique_ptr server_routes::handle_count_tokens(const llama_vocab * vocab, mtmd_context * mctx, const server_http_req & req, task_response_type res_type) { +std::unique_ptr server_routes::handle_count_tokens(const llama_vocab * vocab, mtmd_context * mctx, const mtmd_helper_init_opt & init_opt, const server_http_req & req, task_response_type res_type) { auto res = create_response(); std::vector files; json body = json::parse(req.body); @@ -5395,7 +5402,7 @@ std::unique_ptr server_routes::handle_count_tokens(const l if (!prompt.is_string()) { throw std::runtime_error("for mtmd, input prompt must be a string."); } - n_tokens = process_mtmd_prompt(mctx, prompt.get(), files, true).size(); + n_tokens = process_mtmd_prompt(mctx, prompt.get(), files, init_opt, true).size(); } else { n_tokens = tokenize_mixed(vocab, prompt, true, true).size(); } diff --git a/tools/server/server-context.h b/tools/server/server-context.h index 5d464b8e8c..0acbbffa9e 100644 --- a/tools/server/server-context.h +++ b/tools/server/server-context.h @@ -169,7 +169,7 @@ private: std::unique_ptr handle_slots_restore(const server_http_req & req, int id_slot); std::unique_ptr handle_slots_erase(const server_http_req &, int id_slot); std::unique_ptr handle_embeddings_impl(const server_http_req & req, task_response_type res_type); - std::unique_ptr handle_count_tokens(const llama_vocab * vocab, mtmd_context * mctx, const server_http_req & req, task_response_type res_type); + std::unique_ptr handle_count_tokens(const llama_vocab * vocab, mtmd_context * mctx, const mtmd_helper_init_opt & init_opt, const server_http_req & req, task_response_type res_type); // using unique_ptr to allow late initialization of const std::unique_ptr meta; diff --git a/tools/tts/tts.cpp b/tools/tts/tts.cpp index 368123baf5..6fd1936324 100644 --- a/tools/tts/tts.cpp +++ b/tools/tts/tts.cpp @@ -103,7 +103,7 @@ int main(int argc, char ** argv) { mtmd::bitmap_ptr speaker_bitmap; if (!params.tts_speaker_file.empty()) { - auto wrapper = mtmd_helper_bitmap_init_from_file(mctx.get(), params.tts_speaker_file.c_str(), false); + auto wrapper = mtmd_helper_bitmap_init_from_file(mctx.get(), params.tts_speaker_file.c_str(), false, mtmd_helper_init_opt_default()); if (!wrapper.bitmap) { LOG_ERR("failed to load speaker file %s\n", params.tts_speaker_file.c_str()); return 1; From deae5ee133a3c4c56fbd46c17c8c2103af3bd643 Mon Sep 17 00:00:00 2001 From: fairydreaming <166155368+fairydreaming@users.noreply.github.com> Date: Thu, 27 Aug 2026 12:27:52 +0200 Subject: [PATCH 3/8] model : simplify MiniMax-01 graph (#27790) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: StanisÅ‚aw Szymczyk --- src/models/minimax-01.cpp | 65 +++++++++------------------------------ 1 file changed, 14 insertions(+), 51 deletions(-) diff --git a/src/models/minimax-01.cpp b/src/models/minimax-01.cpp index a6ccee1917..f14626b2c7 100644 --- a/src/models/minimax-01.cpp +++ b/src/models/minimax-01.cpp @@ -174,11 +174,9 @@ public: bool can_reuse(const llm_graph_params & params) override { bool res = true; - if (params.ubatch.n_seq_tokens > 1) { - res &= ( inp_q_decay && inp_q_decay->ne[2] == params.ubatch.n_seq_tokens); - res &= ( inp_k_decay && inp_k_decay->ne[2] == params.ubatch.n_seq_tokens); - res &= (inp_diag_decay && inp_diag_decay->ne[1] == params.ubatch.n_seq_tokens); - } + res &= ( inp_q_decay && inp_q_decay->ne[2] == params.ubatch.n_seq_tokens); + res &= ( inp_k_decay && inp_k_decay->ne[2] == params.ubatch.n_seq_tokens); + res &= (inp_diag_decay && inp_diag_decay->ne[1] == params.ubatch.n_seq_tokens); return res; } @@ -223,19 +221,17 @@ llama_model_minimax_01::graph::graph(const llama_model & model, const llm_graph_ ggml_set_input(inp->inp_slopes); cb(inp->inp_slopes, "slopes", -1); - if (n_seq_tokens != 1) { - inp->inp_q_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, n_head, n_seq_tokens, n_seqs); - ggml_set_input(inp->inp_q_decay); - cb(inp->inp_q_decay, "q_decay_exp", -1); + inp->inp_q_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, n_head, n_seq_tokens, n_seqs); + ggml_set_input(inp->inp_q_decay); + cb(inp->inp_q_decay, "q_decay_exp", -1); - inp->inp_k_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, n_head, n_seq_tokens, n_seqs); - ggml_set_input(inp->inp_k_decay); - cb(inp->inp_k_decay, "k_decay_exp", -1); + inp->inp_k_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, n_head, n_seq_tokens, n_seqs); + ggml_set_input(inp->inp_k_decay); + cb(inp->inp_k_decay, "k_decay_exp", -1); - inp->inp_diag_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_seq_tokens, n_seq_tokens, n_head, n_seqs); - ggml_set_input(inp->inp_diag_decay); - cb(inp->inp_diag_decay, "diag_decay_exp", -1); - } + inp->inp_diag_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_seq_tokens, n_seq_tokens, n_head, n_seqs); + ggml_set_input(inp->inp_diag_decay); + cb(inp->inp_diag_decay, "diag_decay_exp", -1); la = (llm_graph_input_la *) res->add_input(std::move(inp)); @@ -319,41 +315,8 @@ llama_model_minimax_01::graph::graph(const llama_model & model, const llm_graph_ ggml_tensor * qkv = nullptr; ggml_tensor * kv_new = nullptr; - - if (n_seq_tokens == 1) { - // lightning attention - optimized single token case for TG - - ggml_tensor * slopes_neg = ggml_scale(ctx0, slope_rate, -1.0); - cb(slopes_neg, "slopes_neg", il); - - ggml_tensor * ratio = ggml_exp(ctx0, slopes_neg); - cb(ratio, "ratio", il); - - ggml_tensor * ratio_3d = ggml_reshape_3d(ctx0, ratio, 1, 1, n_head); - cb(ratio_3d, "ratio3d", il); - - ggml_tensor * v_trans = ggml_cont(ctx0, ggml_permute(ctx0, Vcur, 1, 2, 0, 3)); - cb(v_trans, "v_trans", il); - - ggml_tensor * k_trans = ggml_cont(ctx0, ggml_permute(ctx0, Kcur, 1, 2, 0, 3)); - cb(k_trans, "k_trans", il); - - ggml_tensor * kv_cur = ggml_mul_mat(ctx0, k_trans, v_trans); - cb(kv_cur, "kv_cur", il); - - ggml_tensor * kv_old_s = ggml_mul(ctx0, kv_old, ratio_3d); - cb(kv_old_s, "kv_old_s", il); - - kv_new = ggml_add(ctx0, kv_old_s, kv_cur); - cb(kv_new, "kv_new", il); - - ggml_tensor * q_trans = ggml_permute(ctx0, Qcur, 0, 2, 1, 3); - cb(q_trans, "q_trans", il); - - qkv = ggml_mul_mat(ctx0, kv_new, q_trans); - cb(qkv, "qkv", il); - } else if(n_seq_tokens > 1) { - // lightning attention - general multi token case for PP + { + // lightning attention ggml_tensor * q_decay_exp = la->inp_q_decay; ggml_tensor * k_decay_exp = la->inp_k_decay; From 2bb9bddafad44ecbb50889644ca47537ec11841b Mon Sep 17 00:00:00 2001 From: Gaurav Garg Date: Thu, 27 Aug 2026 16:23:42 +0530 Subject: [PATCH 4/8] spec: Add benchmark-only synthetic speculative acceptance options (#27711) * Add benchmark-only synthetic speculative acceptance to llama-server and llama-cli * Address review comments * Address review comments * Add some comments in the code --- common/arg.cpp | 32 ++++ common/common.h | 7 + common/speculative.cpp | 157 ++++++++++++++++++-- common/speculative.h | 9 ++ docs/speculative.md | 9 ++ tests/test-arg-parser.cpp | 77 ++++++++++ tools/cli/README.md | 2 + tools/server/README.md | 2 + tools/server/server-context.cpp | 70 ++++++++- tools/server/tests/unit/test_speculative.py | 72 +++++++++ tools/server/tests/utils.py | 7 + 11 files changed, 427 insertions(+), 17 deletions(-) diff --git a/common/arg.cpp b/common/arg.cpp index dd92e4720b..403cc2b781 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -4132,6 +4132,38 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.speculative.draft.n_min = value; } ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_LOOKUP, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_N_MIN")); + add_opt(common_arg( + {"--spec-synth-len"}, "L", + "target mean synthetic acceptance length, including the target token (benchmarking only)", + [](common_params & params, const std::string & value) { + const std::string text = string_strip(value); + size_t pos = 0; + const double length = std::stod(text, &pos); + if (pos != text.size() || length == -1.0) { + throw std::invalid_argument("invalid value"); + } + params.speculative.synth_len = length; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_SYNTH_LEN")); + add_opt(common_arg( + {"--spec-synth-rates"}, "P0,P1,...", + "comma-separated unconditional per-position synthetic acceptance probabilities (benchmarking only)", + [](common_params & params, const std::string & value) { + const auto values = string_split(value, ','); + std::vector rates; + rates.reserve(values.size()); + for (const auto & raw : values) { + const std::string text = string_strip(raw); + size_t pos = 0; + const double rate = std::stod(text, &pos); + if (pos != text.size()) { + throw std::invalid_argument("invalid value"); + } + rates.push_back(rate); + } + params.speculative.synth_rates = std::move(rates); + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_SYNTH_RATES")); add_opt(common_arg( {"--spec-draft-p-split", "--draft-p-split"}, "P", diff --git a/common/common.h b/common/common.h index 9593e10d9a..1cfb01e1ef 100644 --- a/common/common.h +++ b/common/common.h @@ -370,6 +370,9 @@ struct common_params_speculative_ngram_cache { struct common_params_speculative { std::vector types = { COMMON_SPECULATIVE_TYPE_NONE }; + double synth_len = -1.0; + std::vector synth_rates; + // used by Simple, MTP, Eagle3, etc. - all methods that require some kind of draft model common_params_speculative_draft draft; @@ -384,6 +387,10 @@ struct common_params_speculative { return !draft.mparams.empty(); } + bool has_synth() const { + return synth_len != -1.0 || !synth_rates.empty(); + } + uint32_t need_n_rs_seq() const { bool needs_rs_seq = std::any_of(types.begin(), types.end(), [&](auto t) { return t == COMMON_SPECULATIVE_TYPE_DRAFT_MTP || t == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 || t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || t == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK; diff --git a/common/speculative.cpp b/common/speculative.cpp index 4eef2212e7..393a73bc39 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -14,6 +14,7 @@ #include #include +#include #include #include #include @@ -138,6 +139,7 @@ struct common_speculative_impl { const common_speculative_type type; uint32_t n_seq; + int32_t n_max; // maximum draft length after implementation-specific limits size_t n_call_begin = 0; // number of times this implementation was called for refresh. size_t n_call_draft = 0; // number of times this implementation was called for generation. @@ -157,7 +159,7 @@ struct common_speculative_impl { int64_t t_draft_us = 0; // total time spent in generating drafts in this implementation in microseconds. int64_t t_accept_us = 0; // total time spent in accumulation of this implementation in microseconds. - common_speculative_impl(common_speculative_type type, uint32_t n_seq) : type(type), n_seq(n_seq) {} + common_speculative_impl(common_speculative_type type, uint32_t n_seq, int32_t n_max) : type(type), n_seq(n_seq), n_max(n_max) {} virtual ~common_speculative_impl() = default; @@ -182,7 +184,7 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { std::vector smpls; common_speculative_impl_draft_simple(const common_params_speculative & params, uint32_t n_seq) - : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE, n_seq) + : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE, n_seq, params.draft.n_max) , params(params.draft) { auto * ctx_dft = this->params.ctx_dft; @@ -452,7 +454,7 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl { std::vector g_embd_buf; common_speculative_impl_draft_eagle3(const common_params_speculative & params, uint32_t n_seq) - : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3, n_seq) + : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3, n_seq, params.draft.n_max) , params(params.draft) { SPC_TRC("%s", "adding speculative implementation 'draft-eagle3'\n"); @@ -937,7 +939,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { common_speculative_impl_draft_dflash(const common_params_speculative & params, uint32_t n_seq, common_speculative_type type = COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH) - : common_speculative_impl(type, n_seq) + : common_speculative_impl(type, n_seq, params.draft.n_max) , params(params.draft) , is_dspark(type == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK) { @@ -983,6 +985,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { this->params.n_max = std::min(this->params.n_max, n_draft_max); this->params.n_min = std::min(this->params.n_min, n_draft_max); } + this->n_max = this->params.n_max; batch = llama_batch_init(llama_n_batch(ctx_dft), 0, n_seq); batch_inject = llama_batch_init(llama_n_batch(ctx_dft), n_embd_dec, n_seq); @@ -1315,7 +1318,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { std::vector> chain_h; common_speculative_impl_draft_mtp(const common_params_speculative & params, uint32_t n_seq) - : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_MTP, n_seq) + : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_MTP, n_seq, params.draft.n_max) , params(params.draft) { auto * ctx_tgt = this->params.ctx_tgt; @@ -1382,6 +1385,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { c.reserve((size_t) (this->params.n_max + 1) * n_embd); } } + this->n_max = this->params.n_max; pending_h.assign(n_seq, std::vector(n_embd, 0.0f)); @@ -1726,7 +1730,7 @@ struct common_speculative_impl_ngram_simple : public common_speculative_impl { common_speculative_impl_ngram_simple( const common_params_speculative & params, uint32_t n_seq, common_ngram_simple_config config) - : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE, n_seq) + : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE, n_seq, params.ngram_simple.size_m) , params(params.ngram_simple) , config(config) { @@ -1770,7 +1774,7 @@ struct common_speculative_impl_ngram_map_k : public common_speculative_impl { const common_ngram_map & config, uint32_t n_seq) : common_speculative_impl(config.key_only ? COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K - : COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V, n_seq) + : COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V, n_seq, config.size_value) { for (uint32_t i = 0; i < n_seq; i++) { this->config.push_back(config); @@ -1841,7 +1845,7 @@ struct common_speculative_impl_ngram_mod : public common_speculative_impl { common_speculative_impl_ngram_mod( const common_params_speculative & params, uint32_t n_seq) - : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_MOD, n_seq) + : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_MOD, n_seq, params.ngram_mod.n_max) , params(params.ngram_mod) , mod(params.ngram_mod.n_match, 4*1024*1024) , verbose(std::getenv("LLAMA_TRACE") != nullptr) { @@ -2017,7 +2021,7 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl { const std::string & path_dynamic, bool save_dynamic, bool save_static) - : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_CACHE, n_seq) + : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_CACHE, n_seq, n_draft) , params(params.ngram_cache) , n_draft(n_draft) , save_dynamic(save_dynamic) @@ -2138,6 +2142,8 @@ struct common_speculative { // which implementaion was used for a given seq_id std::vector impl_last; + + std::vector synth_probs; }; static common_ngram_map get_common_ngram_map( @@ -2316,6 +2322,101 @@ int32_t common_speculative_n_max(const common_params_speculative * spec) { return n_max; } +int32_t common_speculative_n_max(const common_speculative * spec) { + int32_t n_max = 0; + + if (spec == nullptr) { + return n_max; + } + + for (const auto & impl : spec->impls) { + n_max = std::max(n_max, std::max(0, impl->n_max)); + } + + return n_max; +} + +std::vector common_speculative_synth_rates_resolve(const common_params_speculative * spec, int32_t n_max) { + const bool has_length = spec->synth_len != -1.0; + const bool has_rates = !spec->synth_rates.empty(); + + if (!has_length && !has_rates) { + return {}; + } + if (has_length && has_rates) { + throw std::invalid_argument("synthetic acceptance length and rates are mutually exclusive"); + } + + if (n_max <= 0) { + throw std::invalid_argument("synthetic acceptance requires at least one speculative token"); + } + + if (has_rates) { + const auto & rates = spec->synth_rates; + if (rates.size() != (size_t) n_max) { + throw std::invalid_argument(string_format( + "synthetic acceptance rates must contain %d values, got %zu", n_max, rates.size())); + } + + for (size_t i = 0; i < rates.size(); ++i) { + if (!std::isfinite(rates[i]) || rates[i] < 0.0 || rates[i] > 1.0) { + throw std::invalid_argument("synthetic acceptance rates must be finite and within [0, 1]"); + } + if (i > 0 && rates[i] > rates[i - 1]) { + throw std::invalid_argument("synthetic acceptance rates must be monotonically non-increasing"); + } + } + + return rates; + } + + const double length = spec->synth_len; + const double length_max = (double) n_max + 1.0; + if (!std::isfinite(length) || length < 1.0 || length > length_max) { + throw std::invalid_argument(string_format( + "synthetic acceptance length must be finite and within [1, %.0f]", length_max)); + } + + double p = 0.0; + if (length == length_max) { + p = 1.0; + } else if (length > 1.0) { + double p_min = 0.0; + double p_max = 1.0; + for (int i = 0; i < 32; ++i) { + const double p_mid = 0.5 * (p_min + p_max); + double sum = 0.0; + double term = p_mid; + for (int32_t j = 0; j < n_max; ++j) { + sum += term; + term *= p_mid; + } + + if (sum < length - 1.0) { + p_min = p_mid; + } else { + p_max = p_mid; + } + } + p = 0.5 * (p_min + p_max); + } + + std::vector rates; + rates.reserve(n_max); + double rate = p; + for (int32_t i = 0; i < n_max; ++i) { + rates.push_back(rate); + rate *= p; + } + + return rates; +} + +const std::vector & common_speculative_get_synth_probs(const common_speculative * spec) { + GGML_ASSERT(spec); + return spec->synth_probs; +} + common_params common_base_params_to_speculative(const common_params & params) { const bool has_draft = params.speculative.has_dft(); @@ -2568,13 +2669,39 @@ common_speculative * common_speculative_init(common_params_speculative & params, return nullptr; } - auto * result = new common_speculative { - /* .dparams = */ common_speculative_draft_params_vec(n_seq), - /* .impls = */ std::move(impls), - /* .impl_last = */ std::vector(n_seq, nullptr) - }; + common_speculative_ptr result(new common_speculative { + /* .dparams = */ common_speculative_draft_params_vec(n_seq), + /* .impls = */ std::move(impls), + /* .impl_last = */ std::vector(n_seq, nullptr), + /* .synth_probs = */ {}, + }); - return result; + const int32_t n_max_configured = common_speculative_n_max(¶ms); + const int32_t n_max_effective = common_speculative_n_max(result.get()); + const auto rates = common_speculative_synth_rates_resolve(¶ms, n_max_effective); + + std::vector rates_str; + rates_str.reserve(rates.size()); + result->synth_probs.reserve(rates.size()); + double rate_prev = 1.0; + double acceptance_length = 1.0; + for (const double rate : rates) { + result->synth_probs.push_back(rate_prev > 0.0 ? rate / rate_prev : 0.0); + rates_str.push_back(string_format("%.6g", rate)); + rate_prev = rate; + acceptance_length += rate; + } + if (!result->synth_probs.empty()) { + SPC_WRN("%s", "synthetic speculative acceptance is enabled for benchmarking; generated output is not valid\n"); + if (n_max_effective != n_max_configured) { + SPC_WRN("synthetic acceptance draft limit was reduced from %d to %d by the initialized speculative implementations\n", + n_max_configured, n_max_effective); + } + SPC_INF("synthetic acceptance: n_max = %zu, mean length = %.6f, rates = [%s]\n", + rates.size(), acceptance_length, string_join(rates_str, ", ").c_str()); + } + + return result.release(); } void common_speculative_free(common_speculative * spec) { diff --git a/common/speculative.h b/common/speculative.h index 12ae31b7de..22505891f7 100644 --- a/common/speculative.h +++ b/common/speculative.h @@ -26,6 +26,15 @@ std::string common_speculative_type_to_str(enum common_speculative_type type); // return the max number of draft tokens based on the speculative parameters int32_t common_speculative_n_max(const common_params_speculative * spec); +// return the max number of draft tokens from the initialized implementations +int32_t common_speculative_n_max(const common_speculative * spec); + +// validate and resolve the unconditional synthetic acceptance rates +std::vector common_speculative_synth_rates_resolve(const common_params_speculative * spec, int32_t n_max); + +// return the conditional synthetic acceptance probabilities +const std::vector & common_speculative_get_synth_probs(const common_speculative * spec); + common_params common_base_params_to_speculative(const common_params & params); struct common_speculative_output_limits { diff --git a/docs/speculative.md b/docs/speculative.md index 0f9f8a3d97..ffb1e34c7f 100644 --- a/docs/speculative.md +++ b/docs/speculative.md @@ -212,6 +212,15 @@ Use `--backend-sampling` to run supported target-model samplers on the model bac Unsupported samplers and device layouts fall back to CPU sampling. Tensor split mode does not support backend sampling. A fixed seed produces repeatable random draws, but stochastic CPU and backend sampling can still select different tokens because floating-point operations can differ between implementations and devices. Use greedy sampling when exact output matching is required. +### Synthetic Acceptance + +`llama-server` and `llama-cli` can replace normal speculative verification with synthetic decisions for benchmarking. The generated output is not valid model output because accepted draft tokens do not have to match the target model. + +Use exactly one of these options: + +- `--spec-synth-rates P0,P1,...` sets unconditional per-position acceptance probabilities. Entry `i` is the probability that the first `i+1` draft tokens are all accepted. The number of entries must match the effective maximum draft length. Values must be finite, within `[0, 1]`, and monotonically non-increasing. +- `--spec-synth-len L` sets the target mean acceptance length, including the target token. For `K` maximum draft tokens, `L` must be within `[1, K+1]`. The server finds a constant conditional probability `p` such that `p + p^2 + ... + p^K = L - 1`, then uses unconditional rates `[p, p^2, ..., p^K]`. + ### General Speculative Parameters ``` diff --git a/tests/test-arg-parser.cpp b/tests/test-arg-parser.cpp index ba58f852eb..e0907631ab 100644 --- a/tests/test-arg-parser.cpp +++ b/tests/test-arg-parser.cpp @@ -4,6 +4,7 @@ #include "llama.h" #include "speculative.h" +#include #include #include #include @@ -34,6 +35,62 @@ static void test(void) { std::numeric_limits::max(), std::numeric_limits::max()); + { + common_params_speculative spec; + spec.synth_len = 3.4; + + auto assert_invalid = [](const common_params_speculative & value, int32_t n_max) { + try { + common_speculative_synth_rates_resolve(&value, n_max); + assert(false); + } catch (const std::invalid_argument &) { + } + }; + + const auto rates = common_speculative_synth_rates_resolve(&spec, 4); + assert(rates.size() == 4); + assert(std::abs(rates[0] - 0.80581) < 1e-5); + assert(std::abs(rates[1] - 0.64933) < 1e-5); + assert(std::abs(rates[2] - 0.52323) < 1e-5); + assert(std::abs(rates[3] - 0.42163) < 1e-5); + assert(std::abs(1.0 + rates[0] + rates[1] + rates[2] + rates[3] - 3.4) < 1e-8); + + spec.synth_len = 1.0; + assert(common_speculative_synth_rates_resolve(&spec, 4) == std::vector({0.0, 0.0, 0.0, 0.0})); + + spec.synth_len = 5.0; + assert(common_speculative_synth_rates_resolve(&spec, 4) == std::vector({1.0, 1.0, 1.0, 1.0})); + + spec.synth_len = 5.1; + assert_invalid(spec, 4); + + spec.synth_len = std::numeric_limits::quiet_NaN(); + assert_invalid(spec, 4); + + spec.synth_len = 0.0; + assert_invalid(spec, 4); + + spec.synth_len = -1.0; + spec.synth_rates = {0.8, 0.6, 0.4}; + assert_invalid(spec, 4); + + spec.synth_rates = {0.8, 0.6, 0.4, 0.2}; + assert(common_speculative_synth_rates_resolve(&spec, 4) == spec.synth_rates); + + spec.synth_rates = {0.8, 0.9, 0.4, 0.2}; + assert_invalid(spec, 4); + + spec.synth_rates = {0.8, std::numeric_limits::quiet_NaN(), 0.4, 0.2}; + assert_invalid(spec, 4); + + spec.synth_rates = {0.8, 0.6, 0.4, -0.2}; + assert_invalid(spec, 4); + + spec.synth_rates = {0.8, 0.6, 0.4, 0.2}; + spec.synth_len = 3.0; + assert_invalid(spec, 4); + } + { common_params base; base.n_parallel = 4; @@ -197,6 +254,26 @@ static void test(void) { assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), params, LLAMA_EXAMPLE_SPECULATIVE)); assert(params.speculative.draft.n_max == 123); + { + common_params synth_params; + argv = {"binary_name", "--spec-synth-len", "3.4"}; + assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), synth_params, LLAMA_EXAMPLE_SERVER)); + assert(synth_params.speculative.synth_len == 3.4); + } + + { + common_params synth_params; + argv = {"binary_name", "--spec-synth-rates", "0.8,0.6,0.2"}; + assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), synth_params, LLAMA_EXAMPLE_SERVER)); + assert(synth_params.speculative.synth_rates == std::vector({0.8, 0.6, 0.2})); + } + + { + common_params synth_params; + argv = {"binary_name", "--spec-synth-len", "3.4x"}; + assert(false == common_params_parse(argv.size(), list_str_to_char(argv).data(), synth_params, LLAMA_EXAMPLE_SERVER)); + } + argv = {"binary_name", "-lm", "none"}; assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), params, LLAMA_EXAMPLE_COMMON)); assert(params.load_mode == LLAMA_LOAD_MODE_NONE); diff --git a/tools/cli/README.md b/tools/cli/README.md index 0fba70a90e..e663cfa3b2 100644 --- a/tools/cli/README.md +++ b/tools/cli/README.md @@ -200,6 +200,8 @@ | `--spec-draft-n-cpu-moe, --spec-draft-ncmoe, -ncmoed, --n-cpu-moe-draft N` | keep the Mixture of Experts (MoE) weights of the first N layers in the CPU for the draft model
(env: LLAMA_ARG_SPEC_DRAFT_N_CPU_MOE) | | `--spec-draft-n-max N` | number of tokens to draft for speculative decoding (default: 3)
(env: LLAMA_ARG_SPEC_DRAFT_N_MAX) | | `--spec-draft-n-min N` | minimum number of draft tokens to use for speculative decoding (default: 0)
(env: LLAMA_ARG_SPEC_DRAFT_N_MIN) | +| `--spec-synth-len L` | target mean synthetic acceptance length, including the target token (benchmarking only)
(env: LLAMA_ARG_SPEC_SYNTH_LEN) | +| `--spec-synth-rates P0,P1,...` | comma-separated unconditional per-position synthetic acceptance probabilities (benchmarking only)
(env: LLAMA_ARG_SPEC_SYNTH_RATES) | | `--spec-draft-p-split, --draft-p-split P` | speculative decoding split probability (default: 0.10)
(env: LLAMA_ARG_SPEC_DRAFT_P_SPLIT) | | `--spec-draft-p-min, --draft-p-min P` | minimum speculative decoding probability (greedy) (default: 0.00)
(env: LLAMA_ARG_SPEC_DRAFT_P_MIN) | | `--spec-draft-backend-sampling, --no-spec-draft-backend-sampling` | offload draft sampling to the backend (default: enabled)
(env: LLAMA_ARG_SPEC_DRAFT_BACKEND_SAMPLING) | diff --git a/tools/server/README.md b/tools/server/README.md index f49cdb2729..6fd27f1389 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -259,6 +259,8 @@ For the full list of features, please refer to [server's changelog](https://gith | `--spec-draft-n-cpu-moe, --spec-draft-ncmoe, -ncmoed, --n-cpu-moe-draft N` | keep the Mixture of Experts (MoE) weights of the first N layers in the CPU for the draft model
(env: LLAMA_ARG_SPEC_DRAFT_N_CPU_MOE) | | `--spec-draft-n-max N` | number of tokens to draft for speculative decoding (default: 3)
(env: LLAMA_ARG_SPEC_DRAFT_N_MAX) | | `--spec-draft-n-min N` | minimum number of draft tokens to use for speculative decoding (default: 0)
(env: LLAMA_ARG_SPEC_DRAFT_N_MIN) | +| `--spec-synth-len L` | target mean synthetic acceptance length, including the target token (benchmarking only)
(env: LLAMA_ARG_SPEC_SYNTH_LEN) | +| `--spec-synth-rates P0,P1,...` | comma-separated unconditional per-position synthetic acceptance probabilities (benchmarking only)
(env: LLAMA_ARG_SPEC_SYNTH_RATES) | | `--spec-draft-p-split, --draft-p-split P` | speculative decoding split probability (default: 0.10)
(env: LLAMA_ARG_SPEC_DRAFT_P_SPLIT) | | `--spec-draft-p-min, --draft-p-min P` | minimum speculative decoding probability (greedy) (default: 0.00)
(env: LLAMA_ARG_SPEC_DRAFT_P_MIN) | | `--spec-draft-backend-sampling, --no-spec-draft-backend-sampling` | offload draft sampling to the backend (default: enabled)
(env: LLAMA_ARG_SPEC_DRAFT_BACKEND_SAMPLING) | diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 9fdfcae563..e6c991f7cf 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -23,6 +23,7 @@ #include #include #include +#include #include #include @@ -51,6 +52,50 @@ static common_speculative_output_limits server_output_limits(const common_params return result; } +// synthetic draft verification for benchmarking - accept draft tokens at random instead of by match with the target +// on replay the draft was already accepted before a context checkpoint restore, so repeat the same decisions +static std::vector server_sample_and_accept_synth( + common_sampler * smpl, + llama_context * ctx, + const std::vector & idxs, + const llama_tokens & draft, + const std::vector & synth_probs, + std::mt19937 & rng, + bool is_replay) { + GGML_ASSERT(idxs.size() == draft.size() + 1); + GGML_ASSERT(synth_probs.size() >= draft.size()); + + std::vector result; + result.reserve(idxs.size()); + + const llama_vocab * vocab = llama_model_get_vocab(llama_get_model(ctx)); + std::uniform_real_distribution dist(0.0, 1.0); + for (size_t i = 0; i < draft.size(); ++i) { + const llama_token id = common_sampler_sample(smpl, ctx, idxs[i]); + const bool accept = is_replay || dist(rng) < synth_probs[i]; + // do not accept a drafted EOG token - it would end the generation early + // on replay the last token is from the target and can be EOG, so skip this check + if (accept && (is_replay || !llama_vocab_is_eog(vocab, draft[i]))) { + // synthetic draft tokens do not advance grammar or reasoning state + // the last replay token is from the target and must advance both + const bool is_replay_target = is_replay && i + 1 == draft.size(); + common_sampler_accept(smpl, draft[i], is_replay_target); + result.push_back(draft[i]); + continue; + } + + common_sampler_accept(smpl, id, true); + result.push_back(id); + return result; + } + + const llama_token id = common_sampler_sample(smpl, ctx, idxs[draft.size()]); + common_sampler_accept(smpl, id, true); + result.push_back(id); + + return result; +} + // state diagram: https://github.com/ggml-org/llama.cpp/pull/9283 enum slot_state { SLOT_STATE_IDLE, @@ -211,6 +256,7 @@ struct server_slot { std::vector spec_i_batch; common_prompt_checkpoint spec_ckpt; bool spec_is_replay = false; + std::mt19937 spec_synth_rng; // TODO: move members that belong to the task (such as `generated_text`, `has_new_line`) to task_results_state // see https://github.com/ggml-org/llama.cpp/pull/18283#issuecomment-3710175837 @@ -1194,6 +1240,9 @@ private: spec.reset(common_speculative_init(params_base.speculative, params_base.n_parallel)); } catch (const std::exception & e) { SRV_ERR("failed to initialize speculative decoding context: %s\n", e.what()); + if (params_base.speculative.has_synth()) { + return false; + } } } @@ -1209,6 +1258,11 @@ private: model_dft = nullptr; } + if (!spec && params_base.speculative.has_synth()) { + SRV_ERR("%s", "synthetic acceptance requires an initialized speculative decoding context\n"); + return false; + } + for (int i = 0; i < params_base.n_parallel; i++) { server_slot & slot = slots[i]; @@ -1717,6 +1771,13 @@ private: SLT_TRC(slot, "sampler chain: %s\n", common_sampler_print(slot.smpl.get()).c_str()); SLT_TRC(slot, "sampler params: \n%s\n", task.params.sampling.print().c_str()); + + if (spec && !common_speculative_get_synth_probs(spec.get()).empty()) { + const uint32_t seed = task.params.sampling.seed == LLAMA_DEFAULT_SEED + ? std::random_device{}() + : task.params.sampling.seed; + slot.spec_synth_rng.seed(seed); + } } else { slot.smpl.reset(); } @@ -3802,7 +3863,12 @@ private: common_sampler_ptr smpl_save(common_sampler_clone(slot.smpl.get())); GGML_ASSERT(slot.spec_i_batch.size() == n_draft + 1); - auto accepted = common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft); + const auto & synth_probs = common_speculative_get_synth_probs(spec.get()); + auto accepted = synth_probs.empty() + ? common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft) + : server_sample_and_accept_synth( + slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft, + synth_probs, slot.spec_synth_rng, slot.spec_is_replay); slot.spec_i_batch.clear(); GGML_ASSERT(accepted.size() >= 1); @@ -3868,7 +3934,7 @@ private: auto & n_accepted_per_pos = slot.n_accepted_per_pos; if (n_accepted_per_pos.empty()) { - n_accepted_per_pos.resize(common_speculative_n_max(¶ms_base.speculative), 0); + n_accepted_per_pos.resize(common_speculative_n_max(spec.get()), 0); } for (size_t i = 0; i < n_accepted && i < n_accepted_per_pos.size(); ++i) { n_accepted_per_pos[i]++; diff --git a/tools/server/tests/unit/test_speculative.py b/tools/server/tests/unit/test_speculative.py index 5837195006..22b523954e 100644 --- a/tools/server/tests/unit/test_speculative.py +++ b/tools/server/tests/unit/test_speculative.py @@ -52,6 +52,18 @@ def test_with_and_without_draft(): assert tokens_no_draft == tokens_draft + server.stop() + create_server() + assert server.spec_draft_n_max is not None + server.spec_synth_rates = [0.0] * server.spec_draft_n_max + server.start() + res = server.make_request("POST", "/completion", data=request) + + assert res.status_code == 200 + assert res.body["timings"]["draft_n"] > 0 + assert res.body["timings"]["draft_n_accepted"] == 0 + assert res.body["tokens"] == tokens_no_draft + def test_different_draft_min_draft_max(): global server @@ -80,6 +92,66 @@ def test_different_draft_min_draft_max(): last_content = res.body["content"] +def test_synth_is_deterministic(): + global server + assert server.spec_draft_n_max is not None + server.spec_synth_rates = [0.75 ** (i + 1) for i in range(server.spec_draft_n_max)] + server.start() + + request = { + "prompt": "I believe the meaning of life is", + "temperature": 0.2, + "top_k": 5, + "seed": 4242, + "n_predict": 32, + } + responses = [server.make_request("POST", "/completion", data=request) for _ in range(2)] + + for res in responses: + assert res.status_code == 200 + assert res.body["timings"]["draft_n"] > 0 + assert responses[0].body["timings"]["draft_n"] == responses[1].body["timings"]["draft_n"] + assert responses[0].body["timings"]["draft_n_accepted"] == responses[1].body["timings"]["draft_n_accepted"] + + +def test_synth_ignores_target_tokens(): + global server + assert server.spec_draft_n_max is not None + server.spec_synth_rates = [1.0] * server.spec_draft_n_max + server.start() + + res = server.make_request("POST", "/completion", data={ + "prompt": "I believe the meaning of life is", + "temperature": 0.0, + "seed": 4242, + "n_predict": 32, + }) + + assert res.status_code == 200 + assert res.body["timings"]["draft_n"] > 0 + assert res.body["timings"]["draft_n_accepted"] == res.body["timings"]["draft_n"] + + res = server.make_request("POST", "/completion", data={ + "prompt": "I believe the meaning of life is", + "temperature": 0.0, + "seed": 4242, + "n_predict": 6, + "grammar": 'root ::= "a"{5,5}', + }) + assert res.status_code == 200, res.body + + res = server.make_request("POST", "/completion", data={ + "prompt": "Respond with only: OK", + "temperature": 0.0, + "seed": 4242, + "n_predict": 64, + "ignore_eos": True, + }) + assert res.status_code == 200, res.body + assert res.body["tokens_predicted"] == 64 + assert res.body["stop_type"] == "limit" + + def test_slot_ctx_not_exceeded(): global server server.n_ctx = 256 diff --git a/tools/server/tests/utils.py b/tools/server/tests/utils.py index a0d2dfa3c5..5a4f31a53e 100644 --- a/tools/server/tests/utils.py +++ b/tools/server/tests/utils.py @@ -99,6 +99,8 @@ class ServerProcess: spec_type: str | None = None spec_draft_n_min: int | None = None spec_draft_n_max: int | None = None + spec_synth_len: float | None = None + spec_synth_rates: List[float] | None = None no_ui: bool | None = None jinja: bool | None = None reasoning_format: Literal['deepseek', 'none', 'nothink'] | None = None @@ -245,6 +247,11 @@ class ServerProcess: server_args.extend(["--spec-draft-n-max", self.spec_draft_n_max]) if self.spec_draft_n_min: server_args.extend(["--spec-draft-n-min", self.spec_draft_n_min]) + if self.spec_synth_len is not None: + server_args.extend(["--spec-synth-len", self.spec_synth_len]) + if self.spec_synth_rates is not None: + rates = ",".join(str(rate) for rate in self.spec_synth_rates) + server_args.extend(["--spec-synth-rates", rates]) if self.no_ui: server_args.append("--no-ui") if self.no_models_autoload: From fe235f434368dfbdf3091a6cea92c00a98d2cee4 Mon Sep 17 00:00:00 2001 From: Aleksander Grygier Date: Thu, 27 Aug 2026 13:08:01 +0200 Subject: [PATCH 5/8] ui: Replace per-conversation MCP overrides with per-conversation tool policy (#27745) * ui: replace per-conversation MCP overrides with per-conversation tool policy MCP server enabled state is now global (server.enabled); per-conversation control moves to disabled tool keys and categories seeded into each new conversation. Aligns the add sheet with the dropdown options and flattens MCP tool groups in the tools submenu. Assisted-by: pi * ui: keep tool policy migration running when defaults parse fails A corrupt disabledToolKeys localStorage entry no longer aborts the migration; it falls through with empty defaults so legacy MCP server overrides still get converted. Assisted-by: pi * ui: fall back to global defaults when agentic flow has no tool policy Passing empty disabled sets bypassed the global defaults and could enable tools for callers that do not pass a policy yet. Assisted-by: pi * ui: align preferences section headers with their methods The Reasoning Effort and Working Directory headers sat above tool policy methods; move them above setCwd and setReasoningEffort. Also clarify the disabled tools JSDoc: existing rows with an unset field have an empty policy, defaults apply only when there is no active conversation. Assisted-by: pi * ui: gate MCP server avatars on conversation tool policy Servers whose tools are disabled for the current conversation (MCP category or server-scoped key) no longer show as enabled for the chat. Assisted-by: pi * ui: drop unused MCP category toggle from tools panel hook Per-conversation MCP control is server-granular; no component renders a whole-category toggle, so remove the dead API. Assisted-by: pi * ui: skip MCP init when flow policy disables the MCP category Resolve the effective tool policy before deciding whether to initialize MCP so flows that will not send any MCP tools skip the init work. Callers without a policy keep falling back to global defaults. Assisted-by: pi * chore: format * ui: restore reasoning section in mobile add sheet The sheet rewrite dropped it; the desktop dropdown still has it. MCP Prompts and Resources stay out of the sheet on purpose. Assisted-by: pi * ui: clear MCP server group key in enableAllToolsForServer The group key disables every tool of the server regardless of per-tool keys, so re-enabling a server from Settings did nothing while it was set. Assisted-by: pi * ui: skip MCP init when no policy-enabled server remains Extends the category-level check: the flow also skips MCP init when every globally-enabled server has its server-scoped group key disabled in the tool policy. Assisted-by: pi * ui: make Settings tools tab edit defaults with category toggles Adds per-category checkboxes and a caption stating the tab applies to new conversations; tool picks inside a chat only affect that chat. Assisted-by: pi * ui: gate cwd picker and mention picker on effective tool policy Both checked the global disabled set directly, so a conversation that disabled file_search still showed search as available. Assisted-by: pi * ui: clean up tool key helpers and store docs Documents getEnabledToolsForLLM properly, unstacks the JSDoc at isEntryEnabled, makes setToolEnabled persist like setCategoryEnabled (toggleTool now delegates to it), and routes the serverId-less MCP branch of toolKey through getMcpServerToolsKey so both key formats come from one place. Preferences banner comments become plain comments so they no longer read as class member docs. Assisted-by: pi * ui: indeterminate group checkboxes and inert grayed rows A category that is on with nothing enabled under it now shows the mixed checkbox state instead of a checked box next to 0/N. Rows grayed out by a disabled parent no longer stay clickable behind opacity. Assisted-by: pi * ui: gate MCP prompt and resource capabilities on tool policy hasPromptsCapability and hasResourcesCapability accept an optional set of usable server ids; ChatFormActions resolves it from global enablement minus the active conversation's policy. Restores the per-chat gating the old mcpServerOverrides provided; callers without arguments keep global behavior. Assisted-by: pi * ui: remove unmounted MCP submenu component Never rendered anywhere; its entries are duplicates (prompts and resources live in the attachment menu, servers in the add menu and sheet) that would need capability wiring maintained for nothing. Assisted-by: pi * ui: fix model information dialog width on all screen sizes The dialog sets container-type: inline-size, so auto width ignores its contents and collapses to padding. Give it an explicit viewport width on mobile and cap at 60rem on desktop. Assisted-by: pi * ui: scroll wide chat template in model information dialog Long unbreakable Jinja tokens blew out the table and dialog width; the block now scrolls horizontally instead of stretching. Assisted-by: pi * ui: use fixed table layout in model information dialog Auto table layout sizes columns to content min-content, so the chat template's long lines kept inflating the dialog despite the scroll wrapper. Fixed layout pins the first column and gives the value column a definite width the wrapper can scroll within. min-w-0 on the grid item guards the same path on the grid side. Assisted-by: pi * ui: make model information dialog full-screen on mobile Matches the settings dialog pattern: full viewport below md, calc-sized and capped at 60rem on desktop. Assisted-by: pi * ui: stack chat template row in model information dialog Label above the block in a single full-width cell, so the template gets the whole table width and its horizontal scroll is usable on narrow screens. Assisted-by: pi * ui: scroll model information header with the content The base dialog header is sticky; this dialog overrides it to relative so the title and description scroll away with the body. relative keeps the header as the close button's containing block. Assisted-by: pi * ui: replace literal comment text in sheet group snippet A // line inside the Svelte snippet rendered as visible text; use an HTML comment. Assisted-by: pi * ui: let indeterminate state win over checked in group checkboxes The checkbox indicator snippet renders the check icon whenever checked, so the mixed state never showed. Pass the checked prop as false while indeterminate. Assisted-by: pi * ui: initialize only policy-enabled MCP servers for a flow ensureInitialized accepts an optional server id set; the agentic flow passes the servers its tool policy leaves usable, so servers disabled for the conversation no longer get connected. Callers without arguments keep the global behavior. Assisted-by: pi * ui: derive group checkbox state in useToolsPanel Moves the mixed-state derivation out of the submenu and sheet snippets into one getGroupCheckState accessor; the snippets just consume checked and indeterminate. Assisted-by: pi * ui: gate /prompt command on the conversation tool policy The slash command's availability now follows the same rule as the agentic flow instead of the global capability check, so it disables itself when the conversation's policy leaves no usable MCP server. Assisted-by: pi * ui: remove dead MCP prompt menu trigger chain The /prompt slash command is the surviving trigger; the menu-button path (onMcpPromptClick, hasMcpPromptsSupport, showMcpPromptButton, the MCP_PROMPT attachment item and its unrendered item arrays) has no consumer left. Message display for inserted prompts is untouched. Assisted-by: pi * ui: render dash for mixed-state group checkboxes The accessor refactor dropped the checked-and-not-indeterminate guard, so the category-on flag won and the dash never showed. The tooltip keeps using the raw parent flag since clicking a mixed group still disables it. Assisted-by: pi * ui: fix group checkbox sticking checked after disable Clicking a mixed-state group box let bits-ui optimistically flip its internal checked flag; the derived checked prop did not change across the transition (both mixed and off map to checked=false), so Svelte never applied the settled value and the check icon stuck while the count already read 0/7. Pass the parent flag as checked and the mix as indeterminate, so every group toggle changes checked; render the dash on top of a checked box for the mixed state. Assisted-by: pi * fix: UI for Model Information dialog * ui: keep MCP connections stable across policy switches ensureInitialized folds the policy into its config signature, so alternating two conversations with different policies tore down and reconnected every server with health checks included. Tool collection already filters by the flow policy, so initialize every settings-enabled server instead and never pass a policy into the MCP config. The duplicated policy-server check becomes one accessor on ConversationPreferences. Assisted-by: pi * ui: remove dead MCP resources menu trigger chain Same shape as the earlier prompt trigger cleanup: nothing renders the MCP resources menu button, and the only live entry into resource browsing is Settings > MCP Servers plus the attachment resource picker. Drop onMcpResourcesClick, hasMcpResourcesSupport, MCP_RESOURCES_CLICK, the AttachmentItemVisibleWhen enum and hasResourcesCapability; the resources display, browser and picker components are untouched. Assisted-by: pi --- .../app/chat/ChatForm/ChatForm.svelte | 24 +- .../ChatFormActionAddDropdown.svelte | 4 - .../ChatFormActionAddMcpSubmenu.svelte | 51 ---- .../ChatFormActionAddSheet.svelte | 206 ++++--------- .../ChatFormActionAddToolsSubmenu.svelte | 186 ++++++------ .../ChatFormActions/ChatFormActions.svelte | 30 +- .../ChatFormCurrentWorkingDirectory.svelte | 9 +- .../ChatFormPickerMcpPrompts.svelte | 5 +- .../ChatFormPickerMention.svelte | 15 +- .../ChatMessages/ChatMessageEditForm.svelte | 1 - .../app/chat/ChatScreen/ChatScreenForm.svelte | 1 - tools/ui/src/lib/components/app/chat/index.ts | 13 - .../dialogs/DialogMcpResourcesBrowser.svelte | 5 +- .../app/dialogs/DialogMcpServerAddNew.svelte | 4 +- .../app/dialogs/DialogModelInformation.svelte | 169 ++++++++--- .../app/mcp/McpActiveServersAvatars.svelte | 8 +- .../SettingsChat/SettingsChatToolsTab.svelte | 15 + .../app/settings/SettingsMcpServers.svelte | 10 +- .../components/ui/checkbox/checkbox.svelte | 6 +- .../constants/attachment-menu.constants.ts | 40 +-- .../ui/src/lib/constants/storage.constants.ts | 3 + tools/ui/src/lib/enums/attachment.enums.ts | 12 - tools/ui/src/lib/enums/index.ts | 3 +- .../lib/hooks/use-attachment-menu.svelte.ts | 16 +- .../src/lib/hooks/use-tools-panel.svelte.ts | 81 +++-- .../ui/src/lib/services/migration.service.ts | 60 +++- .../ui/src/lib/stores/agentic/index.svelte.ts | 24 +- tools/ui/src/lib/stores/chat/index.svelte.ts | 9 +- .../lib/stores/conversations/index.svelte.ts | 11 +- .../conversations/preferences.svelte.ts | 284 +++++++++++------- tools/ui/src/lib/stores/mcp/index.svelte.ts | 150 ++------- tools/ui/src/lib/stores/tools.svelte.ts | 152 +++++++--- tools/ui/src/lib/types/agentic.d.ts | 10 +- tools/ui/src/lib/types/chat.d.ts | 7 - tools/ui/src/lib/types/database.d.ts | 12 +- .../tests/unit/mcp-override-fallback.test.ts | 151 ---------- 36 files changed, 838 insertions(+), 949 deletions(-) delete mode 100644 tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddMcpSubmenu.svelte delete mode 100644 tools/ui/tests/unit/mcp-override-fallback.test.ts diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatForm.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatForm.svelte index 40819d6f1a..893f8077dd 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatForm.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatForm.svelte @@ -23,7 +23,8 @@ FileExtensionText, KeyboardKey, MimeTypeText, - SpecialFileType + SpecialFileType, + ToolSource } from '$lib/enums'; import { useChatFormPickers } from '$lib/hooks/use-chat-form-pickers.svelte'; import { @@ -73,7 +74,6 @@ disabled?: boolean; isLoading?: boolean; placeholder?: string; - showMcpPromptButton?: boolean; showAddButton?: boolean; showModelSelector?: boolean; @@ -103,7 +103,6 @@ onValueChange, placeholder = 'Type a message...', showAddButton = true, - showMcpPromptButton = false, showModelSelector = true, uploadedFiles = $bindable([]), value = $bindable('') @@ -152,9 +151,18 @@ getServerHome: () => toolsStore.serverHome ?? null, getShowModelSelector: () => showModelSelector, getValue: () => value, - hasCwdTools: () => toolsStore.hasEnabledCwdTools, - hasPrompts: () => - mcpStore.hasPromptsCapability(conversationsStore.preferences.getAllMcpServerOverrides()), + hasCwdTools: () => conversationsStore.preferences.hasEnabledCwdTools(), + // policy-aware, same rule as the agentic flow: MCP category on and at + // least one globally-enabled server whose group key is not disabled + hasPrompts: () => { + const prefs = conversationsStore.preferences; + + if (!prefs.isCategoryEnabled(ToolSource.MCP)) return false; + + return mcpStore + .getServers() + .some((s) => s.enabled && prefs.isServerToolsEnabled(s.id) && s.url.trim()); + }, openModelSelector: () => chatFormActionsRef?.openModelSelector(), setCaretOffset: (offset) => inputRef?.setCaretOffset(offset), setValue: (v) => { @@ -620,8 +628,6 @@ isReasoning={chatStore.isReasoning} {isRecording} onFileUpload={handleFileUpload} - onMcpPromptClick={showMcpPromptButton ? () => pickers.openPromptPicker() : undefined} - onMcpResourcesClick={() => (isResourceDialogOpen = true)} onMcpSettingsClick={() => (isMcpServersDialogOpen = true)} onMicClick={handleMicClick} {onStop} @@ -635,7 +641,7 @@ - {#if toolsStore.hasEnabledCwdTools} + {#if conversationsStore.preferences.hasEnabledCwdTools()} ({ hasAudioModality: chatFormActions.hasAudioModality, - hasMcpPromptsSupport: chatFormActions.hasMcpPromptsSupport, - hasMcpResourcesSupport: chatFormActions.hasMcpResourcesSupport, hasVideoModality: chatFormActions.hasVideoModality, hasVisionModality: chatFormActions.hasVisionModality }), () => ({ onFileUpload: chatFormActions.onFileUpload, - onMcpPromptClick: chatFormActions.onMcpPromptClick, - onMcpResourcesClick: chatFormActions.onMcpResourcesClick, onSystemPromptClick: chatFormActions.onSystemPromptClick }), () => { diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddMcpSubmenu.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddMcpSubmenu.svelte deleted file mode 100644 index 07439afd61..0000000000 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddMcpSubmenu.svelte +++ /dev/null @@ -1,51 +0,0 @@ - - - - - - - MCP - - - - - - - Servers - - - {#if chatFormActions.hasMcpPromptsSupport} - - - - Prompts - - {/if} - - {#if chatFormActions.hasMcpResourcesSupport} - - - - Resources - - {/if} - - diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddSheet.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddSheet.svelte index 2f69dc96de..acd0f4d211 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddSheet.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddSheet.svelte @@ -1,18 +1,18 @@
@@ -194,80 +186,15 @@ - (mcpExpanded = open)} open={mcpExpanded}> - - {#if mcpExpanded} - - {:else} - - {/if} + - {/each} - - {#if mcpServers.length === 0} -
- No MCP servers configured -
- {/if} -
- - + System Message + {#if toolsPanel.totalToolCount > 0} (toolsExpanded = open)} open={toolsExpanded}> @@ -289,40 +216,12 @@
- {#each toolsPanel.activeGroups as group (group.key)} - {@const checked = toolsPanel.isGroupChecked(group)} - {@const enabledCount = toolsPanel.getEnabledToolCount(group)} - {@const favicon = toolsPanel.getFavicon(group)} + {#each toolsPanel.categoryGroups as group (group.key)} + {@render sheetGroupRow(group)} + {/each} - + {#each toolsPanel.mcpGroups as group (group.key)} + {@render sheetGroupRow(group)} {/each}
@@ -331,38 +230,55 @@ - - {#if chatFormActions.hasMcpPromptsSupport} - - {/if} - - {#if chatFormActions.hasMcpResourcesSupport} - - {/if} + +{#snippet sheetGroupRow(group: ToolGroup)} + {@const checkState = toolsPanel.getGroupCheckState(group)} + {@const enabledCount = toolsPanel.getEnabledToolCount(group)} + {@const favicon = toolsPanel.getFavicon(group)} + {@const groupDisabled = toolsPanel.isGroupDisabled(group)} + + +{/snippet} diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddToolsSubmenu.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddToolsSubmenu.svelte index 40fed27c70..f495441714 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddToolsSubmenu.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddToolsSubmenu.svelte @@ -7,6 +7,7 @@ import { CLI_FLAGS, ICON_CLASS_DEFAULT } from '$lib/constants'; import { useToolsPanel } from '$lib/hooks/use-tools-panel.svelte'; import { mcpStore, toolsStore } from '$lib/stores'; + import type { ToolGroup } from '$lib/types'; const toolsPanel = useToolsPanel(); const hasMcpServersAvailable = $derived(mcpStore.getServers().length > 0); @@ -62,95 +63,108 @@ {/if} {:else}
- {#each toolsPanel.activeGroups as group (group.key)} - {@const isExpanded = toolsPanel.expandedGroups.has(group.key)} - {@const checked = toolsPanel.isGroupChecked(group)} - {@const favicon = toolsPanel.getFavicon(group)} + {#each toolsPanel.categoryGroups as group (group.key)} + {@render groupRow(group)} + {/each} - toolsPanel.toggleGroupExpanded(group.key)} - open={isExpanded} - > -
- - {#if isExpanded} - - {:else} - - {/if} - - - {#if favicon} - { - (e.currentTarget as HTMLImageElement).style.display = 'none'; - }} - src={favicon} - /> - {/if} - - {group.label} - - - - {toolsPanel.getEnabledToolCount(group)}/{group.tools.length} - - - - - - {#snippet child({ props })} - toolsPanel.toggleGroupByKey(group.key)} - /> - {/snippet} - - - -

- {checked ? 'Disable' : 'Enable'} - {group.tools.length} tool{group.tools.length !== 1 ? 's' : ''} -

-
-
-
- - -
- {#each group.tools as entry (entry.key)} - {@const enabled = toolsStore.isToolEnabled(entry.key)} - - {/each} -
-
-
+ {#each toolsPanel.mcpGroups as group (group.key)} + {@render groupRow(group)} {/each}
{/if} + +{#snippet groupRow(group: ToolGroup)} + {@const isExpanded = toolsPanel.expandedGroups.has(group.key)} + {@const checkState = toolsPanel.getGroupCheckState(group)} + {@const favicon = toolsPanel.getFavicon(group)} + {@const groupDisabled = toolsPanel.isGroupDisabled(group)} + + toolsPanel.toggleGroupExpanded(group.key)} + open={isExpanded} + > +
+ + {#if isExpanded} + + {:else} + + {/if} + + + {#if favicon} + { + (e.currentTarget as HTMLImageElement).style.display = 'none'; + }} + src={favicon} + /> + {/if} + + {group.label} + + + + {toolsPanel.getEnabledToolCount(group)}/{group.tools.length} + + + + + + {#snippet child({ props })} + toolsPanel.toggleGroupByKey(group.key)} + /> + {/snippet} + + + +

+ {checkState.checked ? 'Disable' : 'Enable'} + {group.tools.length} tool{group.tools.length !== 1 ? 's' : ''} +

+
+
+
+ + +
+ {#each group.tools as entry (entry.key)} + {@const enabled = toolsPanel.isToolEnabled(entry)} + {@const parentDisabled = toolsPanel.isToolParentDisabled(entry)} + + {/each} +
+
+
+{/snippet} diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActions.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActions.svelte index f1aa743693..395f2cfbe1 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActions.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActions.svelte @@ -13,7 +13,7 @@ import { setChatFormActionsContext } from '$lib/contexts'; import { FileTypeCategory, MessageRole } from '$lib/enums'; import { ChatService } from '$lib/services'; - import { chatStore, conversationsStore, mcpStore, settingsStore } from '$lib/stores'; + import { chatStore, conversationsStore, settingsStore } from '$lib/stores'; import { getFileTypeCategory } from '$lib/utils'; interface Props { @@ -31,8 +31,6 @@ onMicClick?: () => void; onStop?: () => void; onSystemPromptClick?: () => void; - onMcpPromptClick?: () => void; - onMcpResourcesClick?: () => void; onMcpSettingsClick?: () => void; } @@ -45,8 +43,6 @@ isReasoning = false, isRecording = false, onFileUpload, - onMcpPromptClick, - onMcpResourcesClick, onMcpSettingsClick, onMicClick, onStop, @@ -58,18 +54,6 @@ let currentConfig = $derived(settingsStore.config); - let hasMcpPromptsSupport = $derived.by(() => { - const perChatOverrides = conversationsStore.preferences.getAllMcpServerOverrides(); - - return mcpStore.hasPromptsCapability(perChatOverrides); - }); - - let hasMcpResourcesSupport = $derived.by(() => { - const perChatOverrides = conversationsStore.preferences.getAllMcpServerOverrides(); - - return mcpStore.hasResourcesCapability(perChatOverrides); - }); - let hasAudioModality = $state(false); let hasVideoModality = $state(false); let hasVisionModality = $state(false); @@ -142,12 +126,6 @@ get hasAudioModality() { return hasAudioModality; }, - get hasMcpPromptsSupport() { - return hasMcpPromptsSupport; - }, - get hasMcpResourcesSupport() { - return hasMcpResourcesSupport; - }, get hasVideoModality() { return hasVideoModality; }, @@ -157,12 +135,6 @@ get onFileUpload() { return onFileUpload; }, - get onMcpPromptClick() { - return onMcpPromptClick; - }, - get onMcpResourcesClick() { - return onMcpResourcesClick; - }, get onMcpSettingsClick() { return onMcpSettingsClick; }, diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormCurrentWorkingDirectory/ChatFormCurrentWorkingDirectory.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormCurrentWorkingDirectory/ChatFormCurrentWorkingDirectory.svelte index 307d0e702e..ecd1698467 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormCurrentWorkingDirectory/ChatFormCurrentWorkingDirectory.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormCurrentWorkingDirectory/ChatFormCurrentWorkingDirectory.svelte @@ -5,12 +5,12 @@ import SearchInput from '$lib/components/app/forms/SearchInput.svelte'; import * as Popover from '$lib/components/ui/popover'; import { DEFAULT_MOBILE_BREAKPOINT, HOME_TILDE, SEARCH, UI_DATA_ATTRS } from '$lib/constants'; - import { BuiltInTool, GlobSearchType, KeyboardKey } from '$lib/enums'; + import { BuiltInTool, GlobSearchType, KeyboardKey, ToolSource } from '$lib/enums'; import { useDebouncedSearch } from '$lib/hooks/use-debounced-search.svelte'; import { usePickerNavigation } from '$lib/hooks/use-picker-navigation.svelte'; import { useScrollActiveRow } from '$lib/hooks/use-scroll-active-row.svelte'; import { ToolsService } from '$lib/services/tools.service'; - import { toolsStore } from '$lib/stores'; + import { conversationsStore, toolsStore } from '$lib/stores'; import type { GlobEntry } from '$lib/types'; import { abbreviateHome, @@ -63,8 +63,11 @@ // unavailable instead of firing searches that would only fail. Browse is // hidden too: it resolves the picked folder name through the same tool. const fileSearchKey = $derived(toolsStore.getPermissionKey(BuiltInTool.SERVER_FILE_GLOB_SEARCH)); + // effective policy: the active conversation's tool policy, or global defaults const fileSearchEnabled = $derived( - fileSearchKey !== null && toolsStore.isToolEnabled(fileSearchKey) + fileSearchKey !== null && + conversationsStore.preferences.isToolEnabled(fileSearchKey) && + conversationsStore.preferences.isCategoryEnabled(ToolSource.SERVER) ); const searchUnavailableMessage = $derived( fileSearchKey === null diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormPickers/ChatFormPickerMcpPrompts/ChatFormPickerMcpPrompts.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormPickers/ChatFormPickerMcpPrompts/ChatFormPickerMcpPrompts.svelte index 353d6e7baf..f6a3ee1347 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormPickers/ChatFormPickerMcpPrompts/ChatFormPickerMcpPrompts.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormPickers/ChatFormPickerMcpPrompts/ChatFormPickerMcpPrompts.svelte @@ -9,7 +9,7 @@ } from '$lib/components/app/chat'; import Badge from '$lib/components/ui/badge/badge.svelte'; import { KeyboardKey } from '$lib/enums'; - import { conversationsStore, mcpStore } from '$lib/stores'; + import { mcpStore } from '$lib/stores'; import type { GetPromptResult, MCPPromptInfo, MCPServerSettingsEntry } from '$lib/types'; import { debounce, uuid } from '$lib/utils'; import { SvelteMap } from 'svelte/reactivity'; @@ -87,8 +87,7 @@ isLoading = true; try { - const perChatOverrides = conversationsStore.preferences.getAllMcpServerOverrides(); - const initialized = await mcpStore.ensureInitialized(perChatOverrides); + const initialized = await mcpStore.ensureInitialized(); if (!initialized) { prompts = []; diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormPickers/ChatFormPickerMention.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormPickers/ChatFormPickerMention.svelte index 5a5c8320f9..d09a3d3cb8 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormPickers/ChatFormPickerMention.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormPickers/ChatFormPickerMention.svelte @@ -5,10 +5,16 @@ import * as Popover from '$lib/components/ui/popover'; import * as Tooltip from '$lib/components/ui/tooltip'; import { FILE_GLOB_SEARCH_PICKERS, HOME_TILDE, SEARCH } from '$lib/constants'; - import { BuiltInTool, FileMentionEntryType, GlobSearchType, KeyboardKey } from '$lib/enums'; + import { + BuiltInTool, + FileMentionEntryType, + GlobSearchType, + KeyboardKey, + ToolSource + } from '$lib/enums'; import { useDebouncedSearch } from '$lib/hooks/use-debounced-search.svelte'; import { usePickerNavigation } from '$lib/hooks/use-picker-navigation.svelte'; - import { deviceStore, settingsStore, toolsStore } from '$lib/stores'; + import { conversationsStore, deviceStore, settingsStore, toolsStore } from '$lib/stores'; import type { FileMentionEntry, GlobEntryResult } from '$lib/types'; import { abbreviateHome, runGlobSearchWithChildren } from '$lib/utils'; @@ -52,8 +58,11 @@ // --tools) or the user disabled it, the picker still opens but explains // why instead of firing searches that would only fail. const fileSearchKey = $derived(toolsStore.getPermissionKey(BuiltInTool.SERVER_FILE_GLOB_SEARCH)); + // effective policy: the active conversation's tool policy, or global defaults const fileSearchEnabled = $derived( - fileSearchKey !== null && toolsStore.isToolEnabled(fileSearchKey) + fileSearchKey !== null && + conversationsStore.preferences.isToolEnabled(fileSearchKey) && + conversationsStore.preferences.isCategoryEnabled(ToolSource.SERVER) ); let searchResults = $state([]); diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageEditForm.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageEditForm.svelte index 41d79387b6..6e30cebec6 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageEditForm.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageEditForm.svelte @@ -111,7 +111,6 @@ onValueChange={editCtx.setContent} placeholder="Edit your message..." showAddButton={editCtx.messageRole === MessageRole.USER} - showMcpPromptButton showModelSelector={editCtx.messageRole === MessageRole.USER} value={editCtx.editedContent} /> diff --git a/tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreenForm.svelte b/tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreenForm.svelte index 9825b4b90b..962b6774fc 100644 --- a/tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreenForm.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreenForm.svelte @@ -160,6 +160,5 @@ onSubmit={handleSubmit} onSystemPromptClick={handleSystemPromptClick} onUploadedFileRemove={handleUploadedFileRemove} - showMcpPromptButton /> diff --git a/tools/ui/src/lib/components/app/chat/index.ts b/tools/ui/src/lib/components/app/chat/index.ts index 61ec242e90..d7d7745df1 100644 --- a/tools/ui/src/lib/components/app/chat/index.ts +++ b/tools/ui/src/lib/components/app/chat/index.ts @@ -220,19 +220,6 @@ export { default as ChatFormActionModels } from './ChatForm/ChatFormActions/Chat */ export { default as ChatFormActionAddToolsSubmenu } from './ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddToolsSubmenu.svelte'; -/** - * Dropdown submenu for MCP prompts and resources in the chat form. - * - * Shows an "MCP" sub-menu item with entries for MCP Prompts and MCP - * Resources. Only visible when the server supports them. - * - * @example - * ```svelte - * - * ``` - */ -export { default as ChatFormActionAddMcpSubmenu } from './ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddMcpSubmenu.svelte'; - /** * Dropdown submenu for selecting reasoning effort level. * diff --git a/tools/ui/src/lib/components/app/dialogs/DialogMcpResourcesBrowser.svelte b/tools/ui/src/lib/components/app/dialogs/DialogMcpResourcesBrowser.svelte index 6dfcdb856c..82a07477d1 100644 --- a/tools/ui/src/lib/components/app/dialogs/DialogMcpResourcesBrowser.svelte +++ b/tools/ui/src/lib/components/app/dialogs/DialogMcpResourcesBrowser.svelte @@ -8,7 +8,7 @@ import { Button } from '$lib/components/ui/button'; import * as Dialog from '$lib/components/ui/dialog'; import { ICON_CLASS_DEFAULT } from '$lib/constants'; - import { conversationsStore, mcpStore } from '$lib/stores'; + import { mcpStore } from '$lib/stores'; import type { MCPResourceContent, MCPResourceInfo, MCPResourceTemplateInfo } from '$lib/types'; import { getResourceDisplayName } from '$lib/utils'; import { SvelteSet } from 'svelte/reactivity'; @@ -48,8 +48,7 @@ }); async function loadResources() { - const perChatOverrides = conversationsStore.preferences.getAllMcpServerOverrides(); - const initialized = await mcpStore.ensureInitialized(perChatOverrides); + const initialized = await mcpStore.ensureInitialized(); if (initialized) { await mcpStore.fetchAllResources(); diff --git a/tools/ui/src/lib/components/app/dialogs/DialogMcpServerAddNew.svelte b/tools/ui/src/lib/components/app/dialogs/DialogMcpServerAddNew.svelte index fab45aa970..bc28a754c6 100644 --- a/tools/ui/src/lib/components/app/dialogs/DialogMcpServerAddNew.svelte +++ b/tools/ui/src/lib/components/app/dialogs/DialogMcpServerAddNew.svelte @@ -10,7 +10,7 @@ RECOMMENDED_MCP_SERVERS } from '$lib/constants'; import { BooleanString, HealthCheckStatus } from '$lib/enums'; - import { conversationsStore, mcpStore } from '$lib/stores'; + import { mcpStore } from '$lib/stores'; import { canonicalizeServerUrl, parseHeadersToArray, uuid } from '$lib/utils'; interface Props { @@ -234,8 +234,6 @@ useProxy: newServerUseProxy }); - conversationsStore.preferences.setMcpServerOverride(newServerId, true); - handleOpenChange(false); } diff --git a/tools/ui/src/lib/components/app/dialogs/DialogModelInformation.svelte b/tools/ui/src/lib/components/app/dialogs/DialogModelInformation.svelte index 811c24d6b7..e200c004e5 100644 --- a/tools/ui/src/lib/components/app/dialogs/DialogModelInformation.svelte +++ b/tools/ui/src/lib/components/app/dialogs/DialogModelInformation.svelte @@ -76,22 +76,19 @@ - - + + + - - Model Information +
+
+ Model Information - Current model details and capabilities - + Current model details and capabilities +
-
{#if isLoadingModels || isLoadingRouterProps}
Loading model information...
@@ -100,17 +97,15 @@ {@const modelMeta = firstModel.meta} {#if serverProps} - + +