diff --git a/common/arg.cpp b/common/arg.cpp index 64450d388..53c207bbc 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -353,6 +353,8 @@ static std::string get_default_local_path(const std::string & url) { common_models_handler common_models_handler_init(const common_params & params, llama_example curr_ex) { common_download_hf_plan plan; + common_download_hf_plan plan_spec; + common_download_hf_plan plan_voc; common_download_opts opts; const bool spec_type_draft_mtp = std::find(params.speculative.types.begin(), @@ -378,7 +380,15 @@ common_models_handler common_models_handler_init(const common_params & params, l plan = common_download_get_hf_plan(params.model, opts); } - return common_models_handler{plan, opts}; + if (!params.speculative.draft.mparams.hf_repo.empty()) { + plan_spec = common_download_get_hf_plan(params.speculative.draft.mparams, opts); + } + + if (!params.vocoder.model.hf_repo.empty()) { + plan_voc = common_download_get_hf_plan(params.vocoder.model, opts); + } + + return common_models_handler{plan, plan_spec, plan_voc, opts}; } bool common_models_handler_is_preset_repo(const common_models_handler & handler) { @@ -426,7 +436,9 @@ static std::vector build_url_tasks(const common_params_mod void common_models_handler_apply(common_models_handler & handler, common_params & params, common_download_callback * callback) { std::vector tasks; - auto & plan = handler.plan; + auto & plan = handler.plan; + auto & plan_spec = handler.plan_spec; + auto & plan_voc = handler.plan_voc; auto opts = handler.opts; // copy opts.callback = callback; @@ -485,19 +497,22 @@ void common_models_handler_apply(common_models_handler & handler, common_params } // handle hf_plan tasks - if (!plan.model_files.empty()) { - for (size_t i = 0; i < plan.model_files.size(); ++i) { - auto & model_file = plan.model_files[i]; + auto add_tasks = [&opts, &tasks](const hf_cache::hf_files & model_files, common_params_model & model) { + for (size_t i = 0; i < model_files.size(); ++i) { + auto & model_file = model_files[i]; bool is_first = (i == 0); tasks.emplace_back(model_file, opts, [&, is_first]() { if (is_first) { // only use first part as model path - params.model.path = hf_cache::finalize_file(model_file); + model.path = hf_cache::finalize_file(model_file); } else { hf_cache::finalize_file(model_file); } }); } + }; + if (!plan.model_files.empty()) { + add_tasks(plan.model_files, params.model); } if (!plan.mmproj.local_path.empty()) { tasks.emplace_back(plan.mmproj, opts, [&]() { @@ -523,9 +538,31 @@ void common_models_handler_apply(common_models_handler & handler, common_params }); } + // handle plan_spec (e.g. --spec-draft-hf) + if (!plan_spec.model_files.empty()) { + add_tasks(plan_spec.model_files, params.speculative.draft.mparams); + } + + // handle vocoder plan (e.g. --hf-repo-v) + if (!plan_voc.model_files.empty()) { + add_tasks(plan_voc.model_files, params.vocoder.model); + } + // run all tasks in parallel if (!params.offline) { - common_download_run_tasks(tasks); + // if duplicated files are found, only download once (but still call on_done for each task) + std::unordered_map unique_tasks; + for (auto & task : tasks) { + auto it = unique_tasks.find(task.local_path); + if (it == unique_tasks.end()) { + unique_tasks[task.local_path] = &task; + } + } + std::vector unique_tasks_vec; + for (auto & pair : unique_tasks) { + unique_tasks_vec.push_back(*pair.second); + } + common_download_run_tasks(unique_tasks_vec); } // download successful, update params with the downloaded paths @@ -3712,6 +3749,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex "draft model for speculative decoding (default: unused)", [](common_params & params, const std::string & value) { params.speculative.draft.mparams.path = value; + params.speculative.draft.mparams.hf_file = value; // will be used if --spec-draft-hf is set } ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_MODEL")); add_opt(common_arg( diff --git a/common/arg.h b/common/arg.h index 508e33d29..54a38b9cc 100644 --- a/common/arg.h +++ b/common/arg.h @@ -133,6 +133,8 @@ void common_params_add_preset_options(std::vector & args); struct common_models_handler { common_download_hf_plan plan; + common_download_hf_plan plan_spec; + common_download_hf_plan plan_voc; common_download_opts opts; }; diff --git a/conversion/mamba.py b/conversion/mamba.py index be0e36a29..43d559ffb 100644 --- a/conversion/mamba.py +++ b/conversion/mamba.py @@ -114,7 +114,8 @@ class Mamba2Model(TextModel): hparams["text_config"] = hparams["llm_config"] super().__init__(dir_model, *args, hparams=hparams, **kwargs) self.d_model = self.find_hparam(["hidden_size", "d_model", "dim"]) - self.d_inner = self.find_hparam(["mamba_d_ssm", "intermediate_size", "d_inner"], optional=True) or 2 * self.d_model + self.expand = self.find_hparam(["mamba_expand", "expand"], optional=True) or 2 + self.d_inner = self.find_hparam(["mamba_d_ssm", "intermediate_size", "d_inner"], optional=True) or self.expand * self.d_model self.n_group = self.find_hparam(["n_groups"], optional=True) or 1 def set_vocab(self): @@ -144,11 +145,9 @@ class Mamba2Model(TextModel): rms_norm_eps = self.find_hparam(["layer_norm_epsilon", "rms_norm_eps"], optional=True) or 1e-5 - # Fail early for models which don't have a block expansion factor of 2 - # TODO: does this really matter? # skip the assertion for FalconH1 Model if self.model_arch != gguf.MODEL_ARCH.FALCON_H1: - assert self.d_inner == 2 * self.d_model + assert self.d_inner == self.expand * self.d_model assert self.d_inner % head_dim == 0 self.gguf_writer.add_context_length(2**20) # arbitrary value; for those who use the default diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index c865b07e1..bac27241a 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -1558,6 +1558,8 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s int split_backend_id = split->backend_id; ggml_backend_t split_backend = sched->backends[split_backend_id]; + ggml_backend_synchronize(split_backend); + // copy the input tensors to the split backend for (int input_id = 0; input_id < split->n_inputs; input_id++) { ggml_backend_t input_backend = ggml_backend_sched_get_tensor_backend(sched, split->inputs[input_id]); @@ -1568,15 +1570,15 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s // inputs from the user must be copied immediately to prevent the user overwriting the data before the copy is done if (sched->events[split_backend_id][sched->cur_copy] != NULL) { ggml_backend_event_synchronize(sched->events[split_backend_id][sched->cur_copy]); - } else { + } else if (!split_backend->iface.cpy_tensor_async) { ggml_backend_synchronize(split_backend); } - ggml_backend_tensor_copy(input, input_cpy); + ggml_backend_tensor_copy_async(input_backend, split_backend, input, input_cpy); } else { // wait for the split backend to finish using the input before overwriting it if (sched->events[split_backend_id][sched->cur_copy] != NULL) { ggml_backend_event_wait(split_backend, sched->events[split_backend_id][sched->cur_copy]); - } else { + } else if (!split_backend->iface.cpy_tensor_async) { ggml_backend_synchronize(split_backend); } @@ -1681,6 +1683,8 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s } } + ggml_backend_synchronize(split_backend); + if (!sched->callback_eval) { enum ggml_status ec = ggml_backend_graph_compute_async(split_backend, &split->graph); if (ec != GGML_STATUS_SUCCESS) { diff --git a/ggml/src/ggml-cpu/vec.cpp b/ggml/src/ggml-cpu/vec.cpp index 67b6b05ca..ff2b636df 100644 --- a/ggml/src/ggml-cpu/vec.cpp +++ b/ggml/src/ggml-cpu/vec.cpp @@ -75,12 +75,12 @@ void ggml_vec_dot_f32(int n, float * GGML_RESTRICT s, size_t bs, const float * G ay1 = GGML_F32_VEC_LOAD(y + i); sum1 = GGML_F32_VEC_FMA(sum1, ax1, ay1); } - // maximum number of leftover elements will be less that ggml_f32_epr. Apply predicated svmad on available elements only + // maximum number of leftover elements will be less that ggml_f32_epr. Apply predicated svmla on available elements only if (np2 < n) { svbool_t pg = svwhilelt_b32(np2, n); ax1 = svld1_f32(pg, x + np2); ay1 = svld1_f32(pg, y + np2); - sum1 = svmad_f32_m(pg, ax1, ay1, sum1); + sum1 = svmla_f32_m(pg, sum1, ax1, ay1); } // reduce sum1,sum2 to sum1 GGML_F32_VEC_REDUCE(sumf, sum1, sum2, sum3, sum4, sum5, sum6, sum7, sum8); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 42a7e1c52..4c14a97de 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -3204,11 +3204,24 @@ static bool ggml_backend_cuda_cpy_tensor_async(ggml_backend_t backend_src, ggml_ ggml_backend_buffer_t buf_src = src->view_src ? src->view_src->buffer : src->buffer; ggml_backend_buffer_t buf_dst = dst->view_src ? dst->view_src->buffer : dst->buffer; - if (!ggml_backend_is_cuda(backend_src) || !ggml_backend_is_cuda(backend_dst)) { + // Enables async copies from CPU to CUDA, instead of only CUDA-to-CUDA + // Excluding this path for HIP and MUSA as a precaution. + // According to the summary in https://github.com/ggml-org/llama.cpp/pull/20793#issuecomment-4275794315, this change is not beneficial for hip anyways. + // Additionally, there is a lot of anectodal evidence that hip/musa stream behavior might not always 1:1 match CUDA behavior. + // e.g. https://github.com/ROCm/rocm-systems/issues/5109 + // It thus makes sense to exclude this path for HIP and MUSA. This PR was not aimed these backends, the majority of testing happened on CUDA. + // This can be revisited in the future if enabling copy_from_host benefits hip/MUSA, and if the PR author can extensively test on these backends. +#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA) + const bool copy_from_host = false; +#else + const bool copy_from_host = ggml_backend_buffer_is_host(buf_src) && ggml_backend_dev_type(backend_src->device) == GGML_BACKEND_DEVICE_TYPE_CPU; +#endif + + if (!(copy_from_host || ggml_backend_is_cuda(backend_src)) || !ggml_backend_is_cuda(backend_dst)) { return false; } - if (!ggml_backend_buffer_is_cuda(buf_src) || !ggml_backend_buffer_is_cuda(buf_dst)) { + if (!(copy_from_host || ggml_backend_buffer_is_cuda(buf_src)) || !ggml_backend_buffer_is_cuda(buf_dst)) { return false; } @@ -3219,14 +3232,17 @@ static bool ggml_backend_cuda_cpy_tensor_async(ggml_backend_t backend_src, ggml_ ggml_backend_cuda_buffer_context * buf_ctx_src = (ggml_backend_cuda_buffer_context *) buf_src->context; ggml_backend_cuda_buffer_context * buf_ctx_dst = (ggml_backend_cuda_buffer_context *) buf_dst->context; - if (cuda_ctx_src->device != buf_ctx_src->device || cuda_ctx_dst->device != buf_ctx_dst->device) { + if ((copy_from_host && cuda_ctx_dst->device != buf_ctx_dst->device) || + !copy_from_host && (cuda_ctx_src->device != buf_ctx_src->device || cuda_ctx_dst->device != buf_ctx_dst->device)) { #ifndef NDEBUG GGML_LOG_DEBUG("%s: backend and buffer devices do not match\n", __func__); #endif // NDEBUG return false; } - if (backend_src != backend_dst) { + if (copy_from_host) { + CUDA_CHECK(cudaMemcpyAsync(dst->data, src->data, ggml_nbytes(dst), cudaMemcpyHostToDevice, cuda_ctx_dst->stream())); + } else if (backend_src != backend_dst) { // copy on src stream if (cuda_ctx_src->device == cuda_ctx_dst->device) { CUDA_CHECK(cudaMemcpyAsync(dst->data, src->data, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, cuda_ctx_src->stream())); diff --git a/ggml/src/ggml-cuda/out-prod.cu b/ggml/src/ggml-cuda/out-prod.cu index 499903d09..46b9f3a67 100644 --- a/ggml/src/ggml-cuda/out-prod.cu +++ b/ggml/src/ggml-cuda/out-prod.cu @@ -2,6 +2,28 @@ #include +static __global__ void k_compute_out_prod_ptrs( + const float * src0_d, const float * src1_d, float * dst_d, + const float ** ptrs_a, const float ** ptrs_b, float ** ptrs_c, + const int64_t ne2, const int64_t ne3, + const int64_t dps2, const int64_t dps3, + const size_t s02, const size_t s03, + const size_t s12, const size_t s13, + const size_t s2, const size_t s3) { + const int64_t i2 = blockIdx.x*blockDim.x + threadIdx.x; + const int64_t i3 = blockIdx.y*blockDim.y + threadIdx.y; + + if (i2 >= ne2 || i3 >= ne3) { + return; + } + + const int64_t idx = i3*ne2 + i2; + + ptrs_a[idx] = src0_d + (i3/dps3)*s03 + (i2/dps2)*s02; + ptrs_b[idx] = src1_d + i3 *s13 + i2 *s12; + ptrs_c[idx] = dst_d + i3 *s3 + i2 *s2; +} + void ggml_cuda_out_prod(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; @@ -67,18 +89,39 @@ void ggml_cuda_out_prod(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { &beta, dst_d + i3 *s3, ldc, s2, batch_count)); } + } else if (ne2 > 1 || ne3 > 1) { + // dps2 > 1 (src0 broadcast along dim 2 with non-uniform stride) or multiple GEMMs + // along dim 3: compute per-GEMM pointers on the device and use a single batched GEMM. + GGML_ASSERT(ne3 > 0); + GGML_ASSERT(ne2 <= (int64_t) std::numeric_limits::max() / ne3); + const int batch_count = (int) (ne2 * ne3); + + ggml_cuda_pool_alloc ptrs_a(ctx.pool(), batch_count); + ggml_cuda_pool_alloc ptrs_b(ctx.pool(), batch_count); + ggml_cuda_pool_alloc< float *> ptrs_c(ctx.pool(), batch_count); + + const dim3 block_dims(16, 16); + const dim3 grid_dims((ne2 + block_dims.x - 1)/block_dims.x, (ne3 + block_dims.y - 1)/block_dims.y); + k_compute_out_prod_ptrs<<>>( + src0_d, src1_d, dst_d, + ptrs_a.get(), ptrs_b.get(), ptrs_c.get(), + ne2, ne3, dps2, dps3, s02, s03, s12, s13, s2, s3); + CUDA_CHECK(cudaGetLastError()); + + CUBLAS_CHECK( + cublasSgemmBatched(handle, CUBLAS_OP_N, src1_cublas_op, + ne0, ne1, ne01, + &alpha, ptrs_a.get(), lda, + ptrs_b.get(), ldb, + &beta, ptrs_c.get(), ldc, + batch_count)); } else { - // Fallback: ne2 == 1 (no batching benefit) or dps2 > 1 (src0 broadcast along dim 2 - // with non-uniform stride; would need cublasSgemmBatched with pointer arrays). - for (int64_t i3 = 0; i3 < ne3; ++i3) { - for (int64_t i2 = 0; i2 < ne2; ++i2) { - CUBLAS_CHECK( - cublasSgemm(handle, CUBLAS_OP_N, src1_cublas_op, - ne0, ne1, ne01, - &alpha, src0_d + (i3/dps3)*s03 + (i2/dps2)*s02, lda, - src1_d + i3 *s13 + i2 *s12, ldb, - &beta, dst_d + i3 *s3 + i2 *s2, ldc)); - } - } + // ne2 == 1 && ne3 == 1: single GEMM + CUBLAS_CHECK( + cublasSgemm(handle, CUBLAS_OP_N, src1_cublas_op, + ne0, ne1, ne01, + &alpha, src0_d, lda, + src1_d, ldb, + &beta, dst_d, ldc)); } } diff --git a/ggml/src/ggml-cuda/vendors/hip.h b/ggml/src/ggml-cuda/vendors/hip.h index a6115cd80..d01f1533a 100644 --- a/ggml/src/ggml-cuda/vendors/hip.h +++ b/ggml/src/ggml-cuda/vendors/hip.h @@ -48,6 +48,7 @@ #define cublasSetMathMode(handle, mode) CUBLAS_STATUS_SUCCESS #define cublasSetStream hipblasSetStream #define cublasSgemm hipblasSgemm +#define cublasSgemmBatched hipblasSgemmBatched #define cublasSgemmStridedBatched hipblasSgemmStridedBatched #define cublasStatus_t hipblasStatus_t #define cublasOperation_t hipblasOperation_t diff --git a/ggml/src/ggml-cuda/vendors/musa.h b/ggml/src/ggml-cuda/vendors/musa.h index 99e8fa370..6d725c7ec 100644 --- a/ggml/src/ggml-cuda/vendors/musa.h +++ b/ggml/src/ggml-cuda/vendors/musa.h @@ -32,6 +32,7 @@ #define cublasSetMathMode mublasSetMathMode #define cublasSetStream mublasSetStream #define cublasSgemm mublasSgemm +#define cublasSgemmBatched mublasSgemmBatched #define cublasSgemmStridedBatched mublasSgemmStridedBatched #define cublasStatus_t mublasStatus_t #define cublasOperation_t mublasOperation_t diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 1d893b822..3b7abe0c7 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -314,6 +314,7 @@ enum vk_device_architecture { AMD_RDNA1, AMD_RDNA2, AMD_RDNA3, + INTEL_XE1, INTEL_XE2, NVIDIA_PRE_TURING, NVIDIA_TURING, @@ -371,21 +372,26 @@ static vk_device_architecture get_device_architecture(const vk::PhysicalDevice& const std::vector ext_props = device.enumerateDeviceExtensionProperties(); bool subgroup_size_control = false; + bool integer_dot_product = false; for (const auto& properties : ext_props) { if (strcmp("VK_EXT_subgroup_size_control", properties.extensionName) == 0) { subgroup_size_control = true; + } else if (strcmp("VK_KHR_shader_integer_dot_product", properties.extensionName) == 0) { + integer_dot_product = true; } } - if (!subgroup_size_control) { + if (!subgroup_size_control || !integer_dot_product) { return vk_device_architecture::OTHER; } vk::PhysicalDeviceProperties2 props2; vk::PhysicalDeviceSubgroupSizeControlPropertiesEXT subgroup_size_control_props; + vk::PhysicalDeviceShaderIntegerDotProductPropertiesKHR integer_dot_props; props2.pNext = &subgroup_size_control_props; + subgroup_size_control_props.pNext = &integer_dot_props; device.getProperties2(&props2); if (subgroup_size_control_props.minSubgroupSize == 16) { @@ -394,6 +400,9 @@ static vk_device_architecture get_device_architecture(const vk::PhysicalDevice& // https://www.intel.com/content/www/us/en/content-details/824434/2024-intel-tech-tour-xe2-and-lunar-lake-s-gpu.html // https://www.intel.com/content/www/us/en/docs/oneapi/optimization-guide-gpu/2025-0/intel-xe-gpu-architecture.html return vk_device_architecture::INTEL_XE2; + } else if (subgroup_size_control_props.minSubgroupSize == 8 && + integer_dot_product && integer_dot_props.integerDotProduct4x8BitPackedSignedAccelerated) { + return vk_device_architecture::INTEL_XE1; } } else if (props.vendorID == VK_VENDOR_ID_NVIDIA) { const std::vector ext_props = device.enumerateDeviceExtensionProperties(); @@ -3843,7 +3852,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { l_warptile = { 256, 128, 128, 16, subgroup_size_8, 64, 2, tm_m, tn_m, tk_m, subgroup_size_8 }; l_warptile_mmq = l_warptile_mmq_int = { 256, 128, 128, 32, subgroup_size_8, 64, 2, tm_m, tn_m, tk_m, subgroup_size_8 }; l_warptile_mmq_int_k = { 256, 128, 128, 32, subgroup_size_16, 64, 1, 4, 2, 1, subgroup_size_16 }; - } else if (device->vendor_id == VK_VENDOR_ID_INTEL && device->coopmat_support && device->architecture == INTEL_XE2) { + } else if (device->vendor_id == VK_VENDOR_ID_INTEL && device->coopmat_support) { // Xe2/Xe3 with coopmat enabled - warptile performance tuning l_warptile = { 512, 128, 128, 16, subgroup_size_8, 32, 2, tm_m, tn_m, tk_m, subgroup_size_8 }; l_warptile_mmq = { 512, 128, 128, 32, subgroup_size_8, 32, 2, tm_m, tn_m, tk_m, subgroup_size_8 }; @@ -4716,7 +4725,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { } uint32_t rm_iq = 2 * rm_kq; - const bool use_subgroups = device->subgroup_arithmetic && device->architecture != vk_device_architecture::AMD_GCN; + const bool use_subgroups = device->subgroup_arithmetic; // Ensure a subgroup size >= 16 is available const bool use_subgroups16 = use_subgroups && subgroup_min_size_16; @@ -6383,9 +6392,8 @@ static vk_device ggml_vk_get_device(size_t idx) { break; case VK_VENDOR_ID_INTEL: { // Current Windows driver does not expose BF16 support. - // We only want to use l_warptile if coopmat is available and is Xe2+ - const bool xe2_with_coopmat = device->coopmat_support && device->architecture == INTEL_XE2; - const bool use_l_warptile = (i == GGML_TYPE_BF16) ? (device->coopmat_bf16_support && xe2_with_coopmat) : xe2_with_coopmat; + // We only want to use l_warptile if coopmat is available + const bool use_l_warptile = (i == GGML_TYPE_BF16) ? (device->coopmat_bf16_support && device->coopmat_support) : device->coopmat_support; device->mul_mat_l[i] = use_l_warptile; device->mul_mat_id_l[i] = use_l_warptile; device->mul_mat_m[i] = true; @@ -17922,9 +17930,9 @@ static bool ggml_vk_device_is_supported(const vk::PhysicalDevice & vkdev) { static bool ggml_vk_khr_cooperative_matrix_support(const vk::PhysicalDeviceProperties& props, const vk::PhysicalDeviceDriverProperties& driver_props, vk_device_architecture arch) { switch (props.vendorID) { case VK_VENDOR_ID_INTEL: - // Only allowing Xe2 GPU at the moment since Xe2 GPU can gain significant performance boost, - // while some older hardware (ex. Arc A770) has performance regressions - return arch == vk_device_architecture::INTEL_XE2; + // Only allowing Xe2/Xe3 GPU and integrated Xe GPUs at the moment since older hardware (ex. Arc A770) has performance regressions. + return (arch == vk_device_architecture::INTEL_XE2) || + (arch == vk_device_architecture::INTEL_XE1 && props.deviceType == vk::PhysicalDeviceType::eIntegratedGpu && driver_props.driverID == vk::DriverId::eIntelProprietaryWindows); case VK_VENDOR_ID_AMD: if (driver_props.driverID == vk::DriverId::eAmdProprietary || driver_props.driverID == vk::DriverId::eAmdOpenSource) { // Workaround for AMD proprietary driver reporting support on all GPUs @@ -17972,6 +17980,8 @@ static uint32_t ggml_vk_intel_shader_core_count(const vk::PhysicalDevice& vkdev) case 0xE20B: // B580 case 0xE211: // Pro B60 return 20; + case 0xB080: // PTL Xe3 LPG 2x6 (12 subslices) + return 12; default: return 0; } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/conv2d_mm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/conv2d_mm.comp index 1428ef68d..99400098b 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/conv2d_mm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/conv2d_mm.comp @@ -158,7 +158,7 @@ const uint32_t Csh_stride = BS_NPQ; #ifdef COOPMAT const uint32_t Csh_len = BS_K * Csh_stride; #else -const uint32_t Csh_len = csh_store != 0 ? BS_K * Csh_stride : 1; +const uint32_t Csh_len = csh_store != 0 ? BS_K * Csh_stride : 8; // 8 to workaround compiler bug #endif shared SHMEM_TYPE Csh[Csh_len]; // K x NPQ #endif diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/conv3d_mm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/conv3d_mm.comp index a9712eb3a..f66f299f6 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/conv3d_mm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/conv3d_mm.comp @@ -144,7 +144,7 @@ const uint32_t Csh_stride = BS_NPQ; #ifdef COOPMAT const uint32_t Csh_len = BS_K * Csh_stride; #else -const uint32_t Csh_len = csh_store != 0 ? BS_K * Csh_stride : 1; +const uint32_t Csh_len = csh_store != 0 ? BS_K * Csh_stride : 8; // 8 to workaround compiler bug #endif shared SHMEM_TYPE Csh[Csh_len]; // K x NPQ #endif diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vecq.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vecq.comp index fd84c3c91..7bbee577f 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vecq.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vecq.comp @@ -28,13 +28,10 @@ vec2 cache_b_ds; #include "mul_mat_vecq_funcs.glsl" -void iter(inout FLOAT_TYPE temp[NUM_COLS][NUM_ROWS], const uint first_row, const uint num_rows, const uint tid, const uint i) { +void iter(inout FLOAT_TYPE temp[NUM_COLS][NUM_ROWS], const uint first_row, const uint num_rows, const uint col, const uint b_qs_idx) { [[unroll]] for (uint j = 0; j < NUM_COLS; ++j) { - const uint col = i*BLOCK_SIZE + tid*K_PER_ITER; - // Preload data_b block const uint b_block_idx = (j*p.batch_stride_b + col) / QUANT_K_Q8_1 + b_offset; - const uint b_qs_idx = tid % (32 / K_PER_ITER); const uint b_block_idx_outer = b_block_idx / 4; const uint b_block_idx_inner = b_block_idx % 4; cache_b_ds = vec2(data_b[b_block_idx_outer].ds[b_block_idx_inner]); @@ -91,35 +88,35 @@ void compute_outputs(const uint32_t first_row, const uint32_t num_rows) { } } - uint num_iters = p.ncols / (K_PER_ITER * BLOCK_SIZE); - if (num_iters * K_PER_ITER * BLOCK_SIZE + K_PER_ITER*tid < p.ncols) { + const uint col_stride = K_PER_ITER * BLOCK_SIZE; + uint num_iters = p.ncols / col_stride; + if (num_iters * col_stride + K_PER_ITER * tid < p.ncols) { num_iters++; } - int unroll_count = 4; - uint unrolled_iters = num_iters & ~(unroll_count - 1); - uint i = 0; - while (i < unrolled_iters) { + const uint b_qs_idx = tid % (32 / K_PER_ITER); + uint col = tid * K_PER_ITER; + while (num_iters >= 4) { // Manually partially unroll the loop - [[unroll]] for (uint k = 0; k < unroll_count; ++k) { - iter(temp, first_row, num_rows, tid, i*K_PER_ITER); - i++; + [[unroll]] for (uint k = 0; k < 4; ++k) { + iter(temp, first_row, num_rows, col, b_qs_idx); + col += col_stride; } + + num_iters -= 4; } - unroll_count = 2; - unrolled_iters = num_iters & ~(unroll_count - 1); - - while (i < unrolled_iters) { + if (num_iters >= 2) { // Manually partially unroll the loop - [[unroll]] for (uint k = 0; k < unroll_count; ++k) { - iter(temp, first_row, num_rows, tid, i*K_PER_ITER); - i++; - } + iter(temp, first_row, num_rows, col, b_qs_idx); + col += col_stride; + iter(temp, first_row, num_rows, col, b_qs_idx); + col += col_stride; + num_iters -= 2; } - while (i < num_iters) { - iter(temp, first_row, num_rows, tid, i*K_PER_ITER); - i++; + + if (num_iters > 0) { + iter(temp, first_row, num_rows, col, b_qs_idx); } reduce_result(temp, d_offset, first_row, num_rows, tid); diff --git a/src/models/mamba-base.cpp b/src/models/mamba-base.cpp index c37f29c48..fd3fe3f03 100644 --- a/src/models/mamba-base.cpp +++ b/src/models/mamba-base.cpp @@ -169,7 +169,6 @@ ggml_tensor * llm_build_mamba_base::build_mamba2_layer(llm_graph_input_rs * inp, GGML_ASSERT(ubatch.equal_seqs()); GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs); GGML_ASSERT(d_inner % n_head == 0); - GGML_ASSERT(d_inner % d_state == 0); GGML_ASSERT(d_inner % n_group == 0); ggml_tensor * conv_states_all = mctx_cur->get_r_l(il); diff --git a/src/models/mamba2.cpp b/src/models/mamba2.cpp index c5951cf0f..d5c167cf0 100644 --- a/src/models/mamba2.cpp +++ b/src/models/mamba2.cpp @@ -39,10 +39,11 @@ void llama_model_mamba2::load_arch_tensors(llama_model_loader &) { const int64_t d_inner = hparams.ssm_d_inner; const int64_t d_state = hparams.ssm_d_state; const int64_t n_group = hparams.ssm_n_group; - const int64_t d_in_proj = 2*d_inner + 2*n_group*d_state + n_head; + const int64_t dt_rank = hparams.ssm_dt_rank; + + const int64_t conv_dim = d_inner + 2 * n_group * d_state; + const int64_t d_in_proj = d_inner + conv_dim + dt_rank; - // only an expansion factor of 2 is supported for now - GGML_ASSERT(2 * n_embd == d_inner); tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); @@ -68,11 +69,11 @@ void llama_model_mamba2::load_arch_tensors(llama_model_loader &) { layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", i), {d_conv, d_inner + 2*n_group*d_state}, 0); layer.ssm_conv1d_b = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "bias", i), {d_inner + 2*n_group*d_state}, 0); - layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {n_head}, 0); + layer.ssm_dt_b = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", i), {dt_rank}, 0); // no "weight" suffix for these - layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {1, n_head}, 0); - layer.ssm_d = create_tensor(tn(LLM_TENSOR_SSM_D, i), {1, n_head}, 0); + layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A, i), {1, dt_rank}, 0); + layer.ssm_d = create_tensor(tn(LLM_TENSOR_SSM_D, i), {1, dt_rank}, 0); layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", i), {d_inner / n_group, n_group}, 0); diff --git a/tools/mtmd/clip-model.h b/tools/mtmd/clip-model.h index f86702eba..46be39a64 100644 --- a/tools/mtmd/clip-model.h +++ b/tools/mtmd/clip-model.h @@ -55,8 +55,7 @@ struct clip_hparams { int32_t n_head = 0; int32_t n_head_kv = 0; int32_t n_layer = 0; - // idefics3 - int32_t n_merge = 0; // number of patch merges **per-side** + int32_t n_merge = 1; // number of patch merges **per-side** // for preprocessor int32_t image_longest_edge = 0; @@ -135,8 +134,7 @@ struct clip_hparams { int32_t custom_image_max_tokens = -1; void set_limit_image_tokens(int n_tokens_min, int n_tokens_max) { - const int cur_merge = n_merge == 0 ? 1 : n_merge; - const int patch_area = patch_size * patch_size * cur_merge * cur_merge; + const int patch_area = patch_size * patch_size * n_merge * n_merge; image_min_pixels = (custom_image_min_tokens > 0 ? custom_image_min_tokens : n_tokens_min) * patch_area; image_max_pixels = (custom_image_max_tokens > 0 ? custom_image_max_tokens : n_tokens_max) * patch_area; warmup_image_size = static_cast(std::sqrt(image_max_pixels)); @@ -145,8 +143,7 @@ struct clip_hparams { void set_warmup_n_tokens(int n_tokens) { int n_tok_per_side = static_cast(std::sqrt(n_tokens)); GGML_ASSERT(n_tok_per_side * n_tok_per_side == n_tokens && "n_tokens must be n*n"); - const int cur_merge = n_merge == 0 ? 1 : n_merge; - warmup_image_size = n_tok_per_side * patch_size * cur_merge; + warmup_image_size = n_tok_per_side * patch_size * n_merge; // TODO: support warmup size for custom token numbers } // sam vit deepseek-ocr diff --git a/tools/mtmd/clip.cpp b/tools/mtmd/clip.cpp index 68683bd62..606c03a6d 100644 --- a/tools/mtmd/clip.cpp +++ b/tools/mtmd/clip.cpp @@ -1276,6 +1276,9 @@ struct clip_model_loader { { std::vector pinpoints; get_arr_int(KEY_IMAGE_GRID_PINPOINTS, pinpoints, false); + if (pinpoints.size() % 2 != 0) { + throw std::runtime_error(string_format("%s: image_grid_pinpoints must have an even number of elements, got %zu\n", __func__, pinpoints.size())); + } if (!pinpoints.empty()) { for (size_t i = 0; i < pinpoints.size(); i += 2) { hparams.image_res_candidates.push_back({ @@ -1323,15 +1326,16 @@ struct clip_model_loader { } if (is_vision) { - int idx_mean = gguf_find_key(ctx_gguf.get(), KEY_IMAGE_MEAN); - int idx_std = gguf_find_key(ctx_gguf.get(), KEY_IMAGE_STD); - GGML_ASSERT(idx_mean >= 0 && "image_mean not found"); - GGML_ASSERT(idx_std >= 0 && "image_std not found"); - const float * mean_data = (const float *) gguf_get_arr_data(ctx_gguf.get(), idx_mean); - const float * std_data = (const float *) gguf_get_arr_data(ctx_gguf.get(), idx_std); + std::vector image_mean; + std::vector image_std; + get_arr_f32(KEY_IMAGE_MEAN, image_mean, false); + get_arr_f32(KEY_IMAGE_STD , image_std, false); + if (image_mean.size() < 3 || image_std.size() < 3) { + throw std::runtime_error(string_format("%s: image_mean/image_std arrays must have at least 3 elements, got %zu and %zu\n", __func__, image_mean.size(), image_std.size())); + } for (int i = 0; i < 3; ++i) { - hparams.image_mean[i] = mean_data[i]; - hparams.image_std[i] = std_data[i]; + hparams.image_mean[i] = image_mean[i]; + hparams.image_std[i] = image_std[i]; } } @@ -1762,8 +1766,8 @@ struct clip_model_loader { if (hparams.image_size > 65536) { throw std::runtime_error(string_format("%s: image_size (%d) is too large (max 65536)\n", __func__, hparams.image_size)); } - if (hparams.patch_size <= 0) { - throw std::runtime_error(string_format("%s: patch_size (%d) must be greater than 0\n", __func__, hparams.patch_size)); + if (hparams.patch_size <= 0 || hparams.patch_size >= 65536) { + throw std::runtime_error(string_format("%s: patch_size (%d) must be positive and less than 65536\n", __func__, hparams.patch_size)); } if (hparams.n_embd <= 0) { throw std::runtime_error(string_format("%s: n_embd (%d) must be greater than 0\n", __func__, hparams.n_embd)); @@ -1771,6 +1775,9 @@ struct clip_model_loader { if (hparams.image_max_pixels < hparams.image_min_pixels) { throw std::runtime_error(string_format("%s: image_max_pixels (%d) is less than image_min_pixels (%d)\n", __func__, hparams.image_max_pixels, hparams.image_min_pixels)); } + if (hparams.n_merge < 0 || hparams.n_merge >= 65536) { + throw std::runtime_error(string_format("%s: n_merge (%d) must be greater than 0 and less than 65536\n", __func__, hparams.n_merge)); + } } LOG_INF("%s: projector: %s\n", __func__, proj_type.c_str()); @@ -3148,6 +3155,29 @@ struct clip_model_loader { output = gguf_get_val_f32(ctx_gguf.get(), i); } + void get_arr_f32(const std::string & key, std::vector & output, bool required = true) const { + const int i = gguf_find_key(ctx_gguf.get(), key.c_str()); + if (i < 0) { + if (required) { + throw std::runtime_error("Key not found: " + key); + } + return; + } + const auto type = gguf_get_arr_type(ctx_gguf.get(), i); + if (type != GGUF_TYPE_FLOAT32) { + throw std::runtime_error(string_format("%s: array '%s' has type %d, expected %d (GGUF_TYPE_FLOAT32)\n", __func__, key.c_str(), type, GGUF_TYPE_FLOAT32)); + } + const size_t n = gguf_get_arr_n(ctx_gguf.get(), i); + if (n > (size_t) std::numeric_limits::max()) { + throw std::runtime_error(string_format("%s: array '%s' is too large (%zu elements)\n", __func__, key.c_str(), n)); + } + output.resize(n); + const float * values = (const float *)gguf_get_arr_data(ctx_gguf.get(), i); + for (size_t j = 0; j < n; ++j) { + output[j] = values[j]; + } + } + void get_string(const std::string & key, std::string & output, bool required = true) const { const int i = gguf_find_key(ctx_gguf.get(), key.c_str()); if (i < 0) { @@ -3167,11 +3197,18 @@ struct clip_model_loader { } return; } - int n = gguf_get_arr_n(ctx_gguf.get(), i); + const auto type = gguf_get_arr_type(ctx_gguf.get(), i); + if (type != GGUF_TYPE_INT32) { + throw std::runtime_error(string_format("%s: array '%s' has type %d, expected %d (GGUF_TYPE_INT32)\n", __func__, key.c_str(), type, GGUF_TYPE_INT32)); + } + const size_t n = gguf_get_arr_n(ctx_gguf.get(), i); + if (n > (size_t) std::numeric_limits::max()) { + throw std::runtime_error(string_format("%s: array '%s' is too large (%zu elements)\n", __func__, key.c_str(), n)); + } output.resize(n); const int32_t * values = (const int32_t *)gguf_get_arr_data(ctx_gguf.get(), i); - for (int i = 0; i < n; ++i) { - output[i] = values[i]; + for (size_t j = 0; j < n; ++j) { + output[j] = values[j]; } } @@ -3445,8 +3482,8 @@ int clip_n_output_tokens(const clip_ctx * ctx, const clip_image_f32 * img) { { // dynamic size int n_merge = ctx->model.hparams.n_merge; - int n_patches_x = img->nx() / patch_size / (n_merge > 0 ? n_merge : 1); - int n_patches_y = img->ny() / patch_size / (n_merge > 0 ? n_merge : 1); + int n_patches_x = img->nx() / patch_size / n_merge; + int n_patches_y = img->ny() / patch_size / n_merge; if (ctx->model.token_embd_img_break) { n_patches = n_patches_y * n_patches_x + n_patches_y - 1; // + one [IMG_BREAK] per row, except the last row } else { diff --git a/tools/mtmd/models/pixtral.cpp b/tools/mtmd/models/pixtral.cpp index d6d037b69..edfae0825 100644 --- a/tools/mtmd/models/pixtral.cpp +++ b/tools/mtmd/models/pixtral.cpp @@ -63,8 +63,8 @@ ggml_cgraph * clip_graph_pixtral::build() { // and then concatenate the [IMG_BREAK] token to the end of each row, aka n_patches_per_row dimension // after the concatenation, we have a tensor with shape [n_embd, n_patches_per_row + 1, n_rows] - const int p_y = n_merge > 0 ? n_patches_y / n_merge : n_patches_y; - const int p_x = n_merge > 0 ? n_patches_x / n_merge : n_patches_x; + const int p_y = n_patches_y / n_merge; + const int p_x = n_patches_x / n_merge; const int p_total = p_x * p_y; const int n_embd_text = cur->ne[0]; const int n_tokens_output = p_total + p_y - 1; // one [IMG_BREAK] per row, except the last row diff --git a/tools/mtmd/mtmd-image.cpp b/tools/mtmd/mtmd-image.cpp index 57afb542d..01d9b4517 100644 --- a/tools/mtmd/mtmd-image.cpp +++ b/tools/mtmd/mtmd-image.cpp @@ -628,7 +628,7 @@ mtmd_image_preproc_out mtmd_image_preprocessor_llava_uhd::preprocess(const clip_ mtmd_image_preprocessor_llava_uhd::slice_instructions mtmd_image_preprocessor_llava_uhd::get_slice_instructions(const clip_image_size & original_size) { mtmd_image_preprocessor_llava_uhd::slice_instructions res; // align slices by patch_size * n_merge so an integer number of merger output tokens fits per slice - const int n_merge = hparams.n_merge > 0 ? hparams.n_merge : 1; + const int n_merge = hparams.n_merge; const int patch_size = hparams.patch_size * n_merge; const int slice_size = hparams.image_size; const int original_width = original_size.width; @@ -894,7 +894,7 @@ mtmd_image_preproc_out mtmd_image_preprocessor_dyn_size::preprocess(const clip_i clip_image_u8 resized_image; const clip_image_size original_size = img.get_size(); // the original pixtral model doesn't have n_merge - const int cur_merge = hparams.n_merge == 0 ? 1 : hparams.n_merge; + const int cur_merge = hparams.n_merge; const clip_image_size target_size = img_tool::calc_size_preserved_ratio( original_size, hparams.patch_size * cur_merge, diff --git a/tools/server/README-dev.md b/tools/server/README-dev.md index 5959745e4..dfc9004de 100644 --- a/tools/server/README-dev.md +++ b/tools/server/README-dev.md @@ -57,6 +57,7 @@ The core architecture consists of the following components: - `server_tokens`: Unified representation of token sequences (supports both text and multimodal tokens); used by `server_task` and `server_slot`. - `server_prompt_checkpoint`: For recurrent (e.g., RWKV) and SWA models, stores snapshots of KV cache state. Enables reuse when subsequent requests share the same prompt prefix, saving redundant computation. - `server_models`: Standalone component for managing multiple backend instances (used in router mode). It is completely independent of `server_context`. +- `stream_session_manager`: Process wide owner of resumable SSE stream sessions (`g_stream_sessions`), keyed by conversation id. Backs the replay buffer that lets a client reattach to a generation after an HTTP disconnect. See the "Resumable streaming" section below. ```mermaid graph TD @@ -117,6 +118,58 @@ Here is an example trace of an API request for text completion: - As the response is stateless, `server_res_generator` calls `response->update()` to update the response with the current state. - `server_res_generator` then calls `response->to_json()` and passes the response to the HTTP layer. +### Resumable streaming (SSE replay buffer) + +By default a streaming generation is bound to its HTTP socket: when the socket drops (refresh, tab close, mobile background, transient network) the generation aborts and the live stream is lost. This feature keeps the generation running server side and lets a client reattach. + +It is opt in via the `X-Conversation-Id` header on `POST /v1/chat/completions`. Without the header the OAI strict path is unchanged. The conversation id is the only identity end to end (server map key, client localStorage key, route path), with an optional `::model` suffix for direct routing in router mode. + +The feature lives entirely in `server-stream.{h,cpp}` and rests on three types: + +- `stream_session`: a bounded ring buffer (4 MiB cap, oldest bytes drop first) plus a condvar. `append` pushes raw SSE bytes, `read_from` drains from any offset and blocks for live bytes or finalize, `finalize` wakes readers, `cancel` stops the producer. One conv maps to at most one live session. +- `stream_session_manager` (`g_stream_sessions`): owns all sessions keyed by conv id, enforces the one conv one session invariant via `create_or_replace`, and runs a GC thread that drops completed sessions past their TTL. +- `stream_pipe_producer` / `stream_pipe_consumer`: the write and read ends. The producer owns the session lifetime and finalizes it on destruction; the consumer is read only and never finalizes, so a reader detaching cannot kill a running generation. + +Producer side: `server_res_generator` attaches a producer pipe when the header is present. The HTTP content provider mirrors every chunk into the ring before writing it to the socket. While a pipe is attached, `stream_aware_should_stop` ignores peer disconnect, so a dropped socket does not stop generation: only an explicit `DELETE` does. When the peer leaves early, `on_complete` calls `close()`, which drains the rest of the generation into the ring on the http worker. + +Lifetime safety: the producer pipe holds a shared `alive` flag also captured by the session cancel hook. `~server_res_generator` calls `cleanup()` to clear that hook while the reader is still alive, so a `cancel` arriving during teardown can never call `stop()` on a freed response. This ordering is the most fragile part of the feature: finalizing or destroying the producer before `cleanup()` runs reintroduces a use after free. + +Consumer side: `GET /v1/stream/?from=N` opens a `text/event-stream` that replays buffered bytes from offset `N` and blocks for live bytes, so the browser reattaches like a fresh EventSource. An offset below the dropped prefix returns 400. + +Routes: + +- `GET /v1/stream/:conv_id?from=N`: replay or live reattach. +- `POST /v1/streams/lookup` with `{"conversation_ids": [...]}`: returns session status only for ids the caller already owns. There is no listing route, so live sessions cannot be enumerated (an earlier `GET /v1/streams` was removed for exactly this reason). +- `DELETE /v1/stream/:conv_id`: explicit Stop, idempotent (`evict_and_cancel`). + +Router mode binds the same paths to proxy handlers. A `conv_id -> child` map (`conv_models`), populated when a POST is routed, resolves the owning child in one lookup with no polling. The lookup groups ids per child; GET and DELETE proxy straight to the owner. This loopback REST hop is expected to move to a websocket IPC later, swapping only the transport. + +Lifecycle: `g_stream_sessions.start_gc()` runs in main after common init, `stop_gc()` runs first in `clean_up()` and finalizes every live session so no reader hangs. Reader blocking and the post drop drain both run on httplib worker threads, which block on a condvar rather than spin. + +| Constant | Value | Role | +| --- | --- | --- | +| `STREAM_SESSION_TTL_SECONDS` | 300 | retention of a completed session before GC | +| `STREAM_SESSION_MAX_BYTES` | 4 MiB | ring cap per session | +| `STREAM_SESSION_GC_INTERVAL_SECONDS` | 60 | GC tick | +| `STREAM_READ_WAKE_INTERVAL_MS` | 200 | read_from wake to recheck should_stop | +| `STREAM_LOOKUP_TIMEOUT_MS` | 250 | router to child loopback budget | + +```mermaid +graph TD + Client -- "POST + X-Conversation-Id" --> RG[server_res_generator] + RG -- attach --> Prod[stream_pipe_producer] + Prod -- "write, drain on peer drop" --> Sess + subgraph g_stream_sessions + Sess[stream_session: ring buffer, 4 MiB] + GC[GC thread] -- drop after TTL --> Sess + end + Sess -- read_from offset --> Cons[stream_pipe_consumer] + Cons -- "GET /v1/stream/:id?from=N" --> Client + DEL[DELETE /v1/stream/:id] -- evict_and_cancel --> Sess +``` + +The diagram shows the buffer touch points. The live wire (chunks streamed to the original client during a normal generation) is the producer's default output, described under "Producer side" above. + ### Testing `llama-server` includes an automated test suite based on `pytest`. @@ -223,6 +276,7 @@ The flow for downloading a new model: - Speculative decoding: https://github.com/ggml-org/llama.cpp/pull/17808 and rework in https://github.com/ggml-org/llama.cpp/pull/17808 - INI presets: https://github.com/ggml-org/llama.cpp/pull/17859 (+ refactoring: https://github.com/ggml-org/llama.cpp/pull/18169) - Sleeping mode: https://github.com/ggml-org/llama.cpp/pull/18228 +- Resumable streaming (SSE replay buffer): https://github.com/ggml-org/llama.cpp/pull/23226 diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 39b7eb218..5c33a418f 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -5,6 +5,7 @@ #include "server-task.h" #include "server-queue.h" #include "server-schema.h" +#include "server-stream.h" #include "build-info.h" #include "common.h" @@ -4022,6 +4023,15 @@ struct server_res_generator : server_http_res { queue_tasks.wait_until_no_sleep(); } } + ~server_res_generator() override { + // cleanup() must run while rd is still alive (rd is destroyed after this body returns) + if (spipe) { + spipe->cleanup(); + } + } + void stop() override { + rd.stop(); + } void ok(const json & response_data) { status = 200; data = safe_json_to_str(response_data); @@ -4210,8 +4220,10 @@ std::unique_ptr server_routes::handle_completions_impl( } }; + auto effective_should_stop = stream_aware_should_stop(res_this, req.should_stop); + try { - if (req.should_stop()) { + if (effective_should_stop()) { SRV_DBG("%s", "stopping streaming due to should_stop condition\n"); return false; // should_stop condition met } @@ -4245,8 +4257,8 @@ std::unique_ptr server_routes::handle_completions_impl( // receive subsequent results bool timeout = false; int64_t start_time = ggml_time_ms(); - auto result = rd.next([&timeout, &req, &start_time, ¶ms]() { - if (req.should_stop()) { + auto result = rd.next([&timeout, &start_time, ¶ms, &effective_should_stop]() { + if (effective_should_stop()) { return true; // should_stop condition met } else if (params.sse_ping_interval > 0 && ggml_time_ms() - start_time > (int64_t)params.sse_ping_interval * 1000) { timeout = true; @@ -4264,7 +4276,7 @@ std::unique_ptr server_routes::handle_completions_impl( if (result == nullptr) { SRV_DBG("%s", "stopping streaming due to should_stop condition\n"); - GGML_ASSERT(req.should_stop()); + GGML_ASSERT(effective_should_stop()); return false; // should_stop condition met } @@ -4302,6 +4314,10 @@ std::unique_ptr server_routes::handle_completions_impl( }; } + // attach a producer pipe to the response when X-Conversation-Id is present. + // the pipe mirrors SSE chunks into the ring buffer and wires up the cancel hook. + stream_session_attach_pipe(*res, req.headers); + return res; } diff --git a/tools/server/server-http.cpp b/tools/server/server-http.cpp index 4f2abab00..82f34edac 100644 --- a/tools/server/server-http.cpp +++ b/tools/server/server-http.cpp @@ -1,5 +1,6 @@ #include "common.h" #include "server-http.h" +#include "server-stream.h" #include "server-common.h" #include "ui.h" @@ -456,13 +457,40 @@ static void set_headers(httplib::Response & res, const std::map int { + if (c >= '0' && c <= '9') return c - '0'; + if (c >= 'a' && c <= 'f') return c - 'a' + 10; + if (c >= 'A' && c <= 'F') return c - 'A' + 10; + return -1; + }; + int hi = hex(in[i + 1]); + int lo = hex(in[i + 2]); + if (hi >= 0 && lo >= 0) { + out.push_back(char((hi << 4) | lo)); + i += 2; + continue; + } + } + out.push_back(in[i]); + } + return out; +} + static std::map get_params(const httplib::Request & req) { std::map params; for (const auto & [key, value] : req.params) { params[key] = value; } for (const auto & [key, value] : req.path_params) { - params[key] = value; + params[key] = decode_path_component(value); } return params; } @@ -497,26 +525,41 @@ static void process_handler_response(server_http_req_ptr && request, server_http set_headers(res, response->headers); const std::string content_type = response->content_type; // convert to shared_ptr as both chunked_content_provider() and on_complete() need to use it - std::shared_ptr q_ptr = std::move(request); - std::shared_ptr r_ptr = std::move(response); - const auto chunked_content_provider = [response = r_ptr](size_t, const httplib::DataSink & sink) -> bool { + std::shared_ptr q_ptr = std::move(request); + std::shared_ptr r_ptr = std::move(response); + + const auto chunked_content_provider = [response = r_ptr](size_t, httplib::DataSink & sink) -> bool { std::string chunk; const bool has_next = response->next(chunk); if (!chunk.empty()) { + // mirror into the ring buffer first, the session must reflect every SSE chunk + // whether or not the wire write below succeeds + if (response->spipe) { + response->spipe->write(chunk.data(), chunk.size()); + } if (!sink.write(chunk.data(), chunk.size())) { + // peer is gone, stop the wire path here return false; } SRV_DBG("http: streamed chunk: %s\n", chunk.c_str()); } if (!has_next) { + // producer reached its natural end on the wire, a later close() skips the drain + if (response->spipe) { + response->spipe->done(); + } sink.done(); SRV_DBG("%s", "http: stream ended\n"); } return has_next; }; const auto on_complete = [request = q_ptr, response = r_ptr](bool) mutable { - response.reset(); // trigger the destruction of the response object - request.reset(); // trigger the destruction of the request object + // on a dropped peer, close() drains the rest of the generation into the ring buffer + if (response->spipe) { + response->spipe->close(); + } + response.reset(); // spipe destructor finalizes the session if attached + request.reset(); }; res.set_chunked_content_provider(content_type, chunked_content_provider, on_complete); } else { diff --git a/tools/server/server-http.h b/tools/server/server-http.h index 6b4a4b87a..350813183 100644 --- a/tools/server/server-http.h +++ b/tools/server/server-http.h @@ -3,6 +3,7 @@ #include #include #include +#include #include #include #include @@ -10,6 +11,7 @@ #include struct common_params; +struct stream_pipe_producer; // defined in server-stream.h // generator-like API for HTTP response generation // this object response with one of the 2 modes: @@ -23,12 +25,20 @@ struct server_http_res { std::string data; std::map headers; - // TODO: move this to a virtual function once we have proper polymorphism support + // if set, the stream survives a client disconnect: the producer pipe keeps draining into the + // ring buffer and finalizes the session on destruction, so no explicit on_stream_end is needed. + // shared_ptr (not unique_ptr) so the forward-declared type is safe to delete here. + std::shared_ptr spipe; + std::function next = nullptr; bool is_stream() const { return next != nullptr; } + // called when the session is cancelled (e.g. DELETE /v1/stream/). + // server_res_generator overrides this to stop its reader; the default is a no-op. + virtual void stop() {} + virtual ~server_http_res() = default; }; diff --git a/tools/server/server-models.cpp b/tools/server/server-models.cpp index bb2f43a10..0380f98a3 100644 --- a/tools/server/server-models.cpp +++ b/tools/server/server-models.cpp @@ -1,12 +1,14 @@ #include "server-common.h" #include "server-models.h" #include "server-context.h" +#include "server-stream.h" #include "build-info.h" #include "preset.h" #include "download.h" #include // TODO: remove this once we use HTTP client from download.h +#include #include #include @@ -92,6 +94,9 @@ struct server_subproc { } }; +// short loopback budget for the resumable stream router to child JSON calls (probe, lookup, +// delete). distinct from params.timeout_read/write which only applies to the generation proxy +static constexpr int STREAM_LOOKUP_TIMEOUT_MS = 250; static std::filesystem::path get_server_exec_path() { #if defined(_WIN32) @@ -1580,6 +1585,45 @@ static bool is_autoload(const common_params & params, const server_http_req & re } } +// percent encode one query or path component, covers reserved chars without pulling in +// httplib::detail. used by the stream routes to forward conversation_id to children safely +static std::string encode_qs(const std::string & in) { + std::string out; + out.reserve(in.size() * 3); + for (unsigned char c : in) { + bool safe = (c >= 'A' && c <= 'Z') || (c >= 'a' && c <= 'z') || (c >= '0' && c <= '9') + || c == '-' || c == '_' || c == '.' || c == '~'; + if (safe) { + out.push_back(char(c)); + } else { + char buf[4]; + std::snprintf(buf, sizeof(buf), "%%%02X", c); + out.append(buf, 3); + } + } + return out; +} + +// resolve the child that owns a conversation's stream session via the conv_id -> model map +// populated when the POST was routed. single map lookup then a meta lookup, no polling, no +// parsing of the conv id. returns nullopt when nothing maps, the caller answers not found and +// the client recovers +static std::optional resolve_child_for_conv( + server_models & models, const std::string & conversation_id) { + if (conversation_id.empty()) { + return std::nullopt; + } + auto tracked = models.conv_models.lookup(conversation_id); + if (!tracked.has_value()) { + return std::nullopt; + } + auto meta = models.get_meta(*tracked); + if (meta.has_value() && meta->is_ready()) { + return meta; + } + return std::nullopt; +} + void server_models_routes::init_routes() { this->get_router_props = [this](const server_http_req & req) { std::string name = req.get_param("model"); @@ -1628,6 +1672,12 @@ void server_models_routes::init_routes() { if (!router_validate_model(name, models, autoload, error_res)) { return error_res; } + // remember which child serves this conversation so the stream routes can route straight + // to it without polling, keyed on the exact conv id from the header + std::string conv_id = stream_conv_id_from_headers(req.headers); + if (!conv_id.empty()) { + models.conv_models.remember(conv_id, name); + } return models.proxy_request(req, method, name, true); // update last usage for POST request only }; @@ -1819,6 +1869,128 @@ void server_models_routes::init_routes() { res_ok(res, {{"success", true}}); return res; }; + + this->router_stream_get = [this](const server_http_req & req) { + // GET /v1/stream/?from=N. resolve the owning child from the conv_id -> model + // map, 404 when nothing maps + auto res = std::make_unique(); + std::string conv_id = req.get_param("conv_id"); + if (conv_id.empty()) { + res_err(res, format_error_response("Missing conversation id in path", ERROR_TYPE_INVALID_REQUEST)); + return res; + } + std::optional owner = resolve_child_for_conv(models, conv_id); + if (!owner.has_value()) { + res_err(res, format_error_response("Stream not found or expired", ERROR_TYPE_NOT_FOUND)); + return res; + } + std::string from = req.get_param("from"); + std::string child_path = "/v1/stream/" + encode_qs(conv_id); + if (!from.empty()) { + child_path += "?from=" + from; + } + SRV_INF("proxying stream resume to model %s on port %d, path=%s\n", + owner->name.c_str(), owner->port, child_path.c_str()); + auto proxy = std::make_unique( + "GET", + "http", + CHILD_ADDR, + owner->port, + child_path, + req.headers, + req.body, + req.files, + req.should_stop, + params.timeout_read, + params.timeout_write); + return std::unique_ptr(std::move(proxy)); + }; + + this->router_streams_lookup = [this](const server_http_req & req) { + // POST /v1/streams/lookup. resolve each requested conv id to its owning child via the + // map, group the ids per child, and query only the children that actually own some of + // them instead of fanning out to every ready child. a child only answers for the ids + // it owns, never lists anything else + auto res = std::make_unique(); + std::vector requested; + try { + json body = json::parse(req.body); + if (body.contains("conversation_ids") && body["conversation_ids"].is_array()) { + for (const auto & v : body["conversation_ids"]) { + if (v.is_string() && !v.get().empty()) { + requested.push_back(v.get()); + } + } + } + } catch (const std::exception &) { + res_ok(res, json::array()); + return res; + } + + // group requested ids by the child port that owns them, drop ids that map to nothing + std::unordered_map per_child; + for (const auto & cid : requested) { + auto owner = resolve_child_for_conv(models, cid); + if (!owner.has_value()) { + continue; + } + per_child[owner->port].push_back(cid); + } + + json aggregated = json::array(); + for (auto & [port, ids] : per_child) { + json child_body = {{"conversation_ids", ids}}; + httplib::Client cli(CHILD_ADDR, port); + cli.set_connection_timeout(0, STREAM_LOOKUP_TIMEOUT_MS * 1000); + cli.set_read_timeout(0, STREAM_LOOKUP_TIMEOUT_MS * 1000); + cli.set_write_timeout(0, STREAM_LOOKUP_TIMEOUT_MS * 1000); + auto resp = cli.Post("/v1/streams/lookup", child_body.dump(), "application/json"); + if (!resp || resp->status != 200) { + continue; + } + try { + json child_arr = json::parse(resp->body); + if (!child_arr.is_array()) { + continue; + } + for (auto & entry : child_arr) { + if (entry.is_object()) { + aggregated.push_back(entry); + } + } + } catch (const std::exception &) { + continue; + } + } + res_ok(res, aggregated); + return res; + }; + + this->router_stream_delete = [this](const server_http_req & req) { + // DELETE /v1/stream/. resolve the owning child via the map and forward only to + // it, evict_and_cancel is idempotent on the child + auto res = std::make_unique(); + std::string conv_id = req.get_param("conv_id"); + if (conv_id.empty()) { + res_err(res, format_error_response("Missing conversation id in path", ERROR_TYPE_INVALID_REQUEST)); + return res; + } + std::string child_path = "/v1/stream/" + encode_qs(conv_id); + auto owner = resolve_child_for_conv(models, conv_id); + if (owner.has_value()) { + httplib::Client cli(CHILD_ADDR, owner->port); + cli.set_connection_timeout(0, STREAM_LOOKUP_TIMEOUT_MS * 1000); + cli.set_read_timeout(0, STREAM_LOOKUP_TIMEOUT_MS * 1000); + cli.set_write_timeout(0, STREAM_LOOKUP_TIMEOUT_MS * 1000); + auto resp = cli.Delete(child_path.c_str()); + (void) resp; // best effort, 404 and network errors are equivalent to no op + } + // drop the tracking entry, the session is being torn down + models.conv_models.forget(conv_id); + res->status = 204; + res->content_type = "application/json"; + return res; + }; } diff --git a/tools/server/server-models.h b/tools/server/server-models.h index 9ed4aeead..62bed8725 100644 --- a/tools/server/server-models.h +++ b/tools/server/server-models.h @@ -11,7 +11,10 @@ #include #include #include +#include #include +#include +#include /** * state diagram: @@ -126,6 +129,44 @@ private: // if true, the next get_meta() will trigger a reload of model list bool need_reload = false; + // conv_id -> model name that currently serves its stream session, lets the resumable stream + // routes go straight to the owning child instead of polling every one. populated when + // proxy_request forwards a POST carrying an X-Conversation-Id. best effort: a stale entry just + // makes the child answer not found and the client recovers. owns its lock, one mutex per struct + struct conv_model_tracker { + void remember(const std::string & conv_id, const std::string & model) { + if (conv_id.empty() || model.empty()) { + return; + } + std::lock_guard lock(mu); + map[conv_id] = model; + } + + std::optional lookup(const std::string & conv_id) { + if (conv_id.empty()) { + return std::nullopt; + } + std::lock_guard lock(mu); + auto it = map.find(conv_id); + if (it == map.end()) { + return std::nullopt; + } + return it->second; + } + + void forget(const std::string & conv_id) { + if (conv_id.empty()) { + return; + } + std::lock_guard lock(mu); + map.erase(conv_id); + } + + private: + std::mutex mu; + std::unordered_map map; + }; + common_preset_context ctx_preset; common_params base_params; @@ -145,6 +186,9 @@ private: void notify_sse(const std::string & event, const std::string & model_id, const json & data = nullptr); public: + // conv_id -> model tracker for the resumable stream routes, owns its lock + conv_model_tracker conv_models; + server_models(const common_params & params, int argc, char ** argv); server_response sse; // for real-time updates via SSE endpoint @@ -268,6 +312,12 @@ struct server_models_routes { server_http_context::handler_t get_router_models_sse; server_http_context::handler_t post_router_models; server_http_context::handler_t del_router_models; + + // router side handlers for the resumable streaming routes. each resolves the child that owns + // a conversation through the conv_id -> model map, no probing or fan out + server_http_context::handler_t router_stream_get; + server_http_context::handler_t router_streams_lookup; + server_http_context::handler_t router_stream_delete; }; /** diff --git a/tools/server/server-stream.cpp b/tools/server/server-stream.cpp new file mode 100644 index 000000000..757c36ad2 --- /dev/null +++ b/tools/server/server-stream.cpp @@ -0,0 +1,569 @@ +#include "server-stream.h" +#include "server-common.h" +#include "server-http.h" +#include "server-queue.h" + +#include +#include +#include + +namespace { +constexpr int64_t STREAM_SESSION_TTL_SECONDS = 300; +constexpr size_t STREAM_SESSION_MAX_BYTES = 4 * 1024 * 1024; +constexpr int64_t STREAM_SESSION_GC_INTERVAL_SECONDS = 60; +constexpr int64_t STREAM_READ_WAKE_INTERVAL_MS = 200; + +// returns unix time in seconds +int64_t now_seconds() { + return std::chrono::duration_cast( + std::chrono::system_clock::now().time_since_epoch() + ).count(); +} +} + +stream_session::stream_session(std::string conversation_id_, size_t max_bytes_) + : conversation_id(std::move(conversation_id_)) + , started_ts(now_seconds()) + , prefix_dropped(0) + , cap_bytes(max_bytes_) + , done(false) + , cancelled(false) + , completed_ts(0) { + buffer.reserve(64 * 1024); +} + +bool stream_session::append(const char * data, size_t len) { + if (len == 0) { + return true; + } + { + std::lock_guard lock(mu); + if (done.load(std::memory_order_relaxed)) { + return false; + } + if (len >= cap_bytes) { + // single chunk bigger than the cap, keep only the tail that fits + size_t skip = len - cap_bytes; + prefix_dropped += buffer.size() + skip; + buffer.clear(); + buffer.insert(buffer.end(), data + skip, data + len); + } else { + size_t needed = buffer.size() + len; + if (needed > cap_bytes) { + size_t to_drop = needed - cap_bytes; + buffer.erase(buffer.begin(), buffer.begin() + to_drop); + prefix_dropped += to_drop; + } + buffer.insert(buffer.end(), data, data + len); + } + } + cv.notify_all(); + return true; +} + +void stream_session::finalize() { + bool was_done = done.exchange(true, std::memory_order_acq_rel); + if (was_done) { + return; + } + completed_ts.store(now_seconds(), std::memory_order_release); + cv.notify_all(); +} + +stream_read_status stream_session::read_from(size_t offset, + const std::function & sink, + const std::function & should_stop) { + std::unique_lock lock(mu); + while (true) { + if (should_stop && should_stop()) { + return stream_read_status::OK; + } + if (offset < prefix_dropped) { + return stream_read_status::OFFSET_LOST; + } + size_t logical_end = prefix_dropped + buffer.size(); + if (offset < logical_end) { + size_t local_off = offset - prefix_dropped; + size_t n = buffer.size() - local_off; + // copy the available chunk under the lock, release before calling the sink + std::vector chunk(buffer.begin() + local_off, buffer.begin() + local_off + n); + offset += n; + lock.unlock(); + bool keep_going = sink(chunk.data(), chunk.size()); + if (!keep_going) { + return stream_read_status::OK; + } + lock.lock(); + continue; + } + if (done.load(std::memory_order_acquire)) { + return stream_read_status::OK; + } + // wait for new bytes, finalize, or a periodic wake to re check should_stop + cv.wait_for(lock, std::chrono::milliseconds(STREAM_READ_WAKE_INTERVAL_MS)); + } +} + +bool stream_session::is_done() const { + return done.load(std::memory_order_acquire); +} + +size_t stream_session::total_size() const { + std::lock_guard lock(mu); + return prefix_dropped + buffer.size(); +} + +size_t stream_session::dropped_prefix() const { + std::lock_guard lock(mu); + return prefix_dropped; +} + +int64_t stream_session::completed_at() const { + return completed_ts.load(std::memory_order_acquire); +} + +void stream_session::set_stop_producer(std::function fn) { + std::lock_guard lock(mu); + stop_producer = std::move(fn); +} + +void stream_session::cancel() { + // flip cancelled first so the producer-side stream_aware_should_stop can break out of the + // recv() wait even if remove_waiting_task_ids does not notify the condvar (the cancel task + // posted by rd.stop() will eventually notify, but we do not want to depend on that timing) + cancelled.store(true, std::memory_order_release); + // copy the hook under the lock then invoke outside, the producer side may grab queue locks + // and we do not want to hold our mu across that path + std::function fn; + { + std::lock_guard lock(mu); + fn = stop_producer; + } + if (fn) { + fn(); + } +} + +bool stream_session::is_cancelled() const { + return cancelled.load(std::memory_order_acquire); +} + +stream_session_manager::stream_session_manager() + : running(false) { +} + +stream_session_manager::~stream_session_manager() { + stop_gc(); +} + +stream_session_ptr stream_session_manager::create_or_replace(const std::string & conversation_id) { + // evict any previous session on the same conv, this guarantees the invariant + // "one conv = at most one live session" and propagates cancel to its producer + stream_session_ptr previous; + auto fresh = std::make_shared(conversation_id, STREAM_SESSION_MAX_BYTES); + { + std::unique_lock lock(map_mu); + auto it = sessions.find(conversation_id); + if (it != sessions.end()) { + previous = it->second; + it->second = fresh; + } else { + sessions.emplace(conversation_id, fresh); + } + } + if (previous) { + previous->cancel(); + previous->finalize(); + } + return fresh; +} + +stream_session_ptr stream_session_manager::get(const std::string & conversation_id) { + std::shared_lock lock(map_mu); + auto it = sessions.find(conversation_id); + if (it == sessions.end()) { + return nullptr; + } + return it->second; +} + +std::vector stream_session_manager::list_all() const { + std::vector out; + std::shared_lock lock(map_mu); + out.reserve(sessions.size()); + for (auto & kv : sessions) { + out.push_back(kv.second); + } + return out; +} + +void stream_session_manager::evict(const std::string & conversation_id) { + stream_session_ptr s; + { + std::unique_lock lock(map_mu); + auto it = sessions.find(conversation_id); + if (it == sessions.end()) { + return; + } + s = it->second; + sessions.erase(it); + } + // finalize outside the map lock so any pending readers wake up and exit + s->finalize(); +} + +void stream_session_manager::evict_and_cancel(const std::string & conversation_id) { + stream_session_ptr s; + { + std::unique_lock lock(map_mu); + auto it = sessions.find(conversation_id); + if (it == sessions.end()) { + return; + } + s = it->second; + sessions.erase(it); + } + // signal the producer side first so the inference is cancelled at the queue level, + // then finalize, which wakes any pending HTTP reader and lets the drain exit naturally + s->cancel(); + s->finalize(); +} + +void stream_session_manager::start_gc() { + if (running.exchange(true)) { + return; + } + gc_thread = std::thread([this] { gc_loop(); }); +} + +void stream_session_manager::stop_gc() { + bool was_running = running.exchange(false); + if (was_running) { + { + std::lock_guard lock(gc_wake_mu); + } + gc_wake_cv.notify_all(); + if (gc_thread.joinable()) { + gc_thread.join(); + } + } + // finalize all live sessions so no reader ever hangs + std::vector snapshot; + { + std::unique_lock lock(map_mu); + snapshot.reserve(sessions.size()); + for (auto & kv : sessions) { + snapshot.push_back(kv.second); + } + sessions.clear(); + } + for (auto & s : snapshot) { + s->finalize(); + } +} + +void stream_session_manager::gc_loop() { + while (running.load(std::memory_order_acquire)) { + { + std::unique_lock lock(gc_wake_mu); + gc_wake_cv.wait_for(lock, + std::chrono::seconds(STREAM_SESSION_GC_INTERVAL_SECONDS), + [this] { return !running.load(std::memory_order_acquire); }); + } + if (!running.load(std::memory_order_acquire)) { + return; + } + int64_t cutoff = now_seconds() - STREAM_SESSION_TTL_SECONDS; + std::vector to_drop; + { + std::unique_lock lock(map_mu); + for (auto it = sessions.begin(); it != sessions.end(); ) { + int64_t completed = it->second->completed_at(); + if (completed != 0 && completed <= cutoff) { + to_drop.push_back(it->second); + it = sessions.erase(it); + } else { + ++it; + } + } + } + // finalize outside the map lock, idempotent if the session was already done + for (auto & s : to_drop) { + s->finalize(); + } + } +} + +// process wide manager, lifecycle controlled by llama-server main() via start_gc/stop_gc +stream_session_manager g_stream_sessions; + +// stream_pipe --------------------------------------------------------------------------------- + +stream_pipe::stream_pipe(stream_session_ptr session) + : session_(std::move(session)) { +} + +bool stream_pipe::is_cancelled() const { + return session_->is_cancelled(); +} + +// stream_pipe_producer + +stream_pipe_producer::stream_pipe_producer(stream_session_ptr session) + : stream_pipe(std::move(session)) { +} + +stream_pipe_producer::~stream_pipe_producer() { + cleanup(); + session_->finalize(); +} + +void stream_pipe_producer::cleanup() { + if (!alive_) { + return; + } + alive_->store(false, std::memory_order_release); + session_->set_stop_producer(nullptr); + alive_.reset(); +} + +bool stream_pipe_producer::write(const char * data, size_t len) { + return session_->append(data, len); +} + +void stream_pipe_producer::done() { + done_ = true; +} + +void stream_pipe_producer::close() { + // httplib bails its content provider the moment is_peer_alive() goes false, so pump the rest + // of the generation into the ring buffer here. a DELETE flips is_cancelled and cuts it short + if (done_ || session_->is_cancelled()) { + SRV_INF("stream_pipe close: skip drain (done=%d cancelled=%d) conv=%s\n", + done_ ? 1 : 0, session_->is_cancelled() ? 1 : 0, session_->conversation_id.c_str()); + return; + } + SRV_INF("stream_pipe close: draining conv=%s\n", session_->conversation_id.c_str()); + size_t drained = 0; + std::string chunk; + while (true) { + chunk.clear(); + bool has_next = res_->next(chunk); + if (!chunk.empty()) { + write(chunk.data(), chunk.size()); + drained += chunk.size(); + } + if (!has_next) { + break; + } + } + SRV_INF("stream_pipe close: drain ended conv=%s bytes=%zu\n", session_->conversation_id.c_str(), drained); +} + +std::shared_ptr stream_pipe_producer::create(stream_session_ptr session, + server_http_res & res) { + auto alive = std::make_shared>(true); + auto * res_ptr = &res; + session->set_stop_producer([alive, res_ptr]() { + if (alive->load(std::memory_order_acquire)) { + res_ptr->stop(); + } + }); + auto pipe = std::shared_ptr(new stream_pipe_producer(std::move(session))); + pipe->alive_ = std::move(alive); + pipe->res_ = res_ptr; + return pipe; +} + +// stream_pipe_consumer + +stream_pipe_consumer::stream_pipe_consumer(stream_session_ptr session) + : stream_pipe(std::move(session)) { +} + +stream_read_status stream_pipe_consumer::read(size_t & offset, + const std::function & sink, + const std::function & should_stop) { + return session_->read_from(offset, sink, should_stop); +} + +std::shared_ptr stream_pipe_consumer::create(stream_session_ptr session) { + return std::shared_ptr(new stream_pipe_consumer(std::move(session))); +} + +// helper, builds the standard error response and assigns it to a brand new http_res +static server_http_res_ptr make_error_response(int status, const std::string & message, error_type type) { + auto res = std::make_unique(); + json err = format_error_response(message, type); + res->status = json_value(err, "code", status); + res->content_type = "application/json; charset=utf-8"; + res->data = safe_json_to_str({{"error", err}}); + return res; +} + +server_http_context::handler_t make_stream_get_handler() { + return [](const server_http_req & req) -> server_http_res_ptr { + // GET /v1/stream/?from=N replays the SSE bytes already buffered for the + // session, blocks for more bytes when the session is still running, returns when + // the session is finalized. the body is streamed back as text/event-stream so the + // browser EventSource can attach to it like a fresh request + std::string conv_id = req.get_param("conv_id"); + if (conv_id.empty()) { + return make_error_response(400, "Missing conversation id in path", ERROR_TYPE_INVALID_REQUEST); + } + auto session = g_stream_sessions.get(conv_id); + if (!session) { + return make_error_response(404, "Stream not found or expired", ERROR_TYPE_NOT_FOUND); + } + size_t from = 0; + std::string from_str = req.get_param("from"); + if (!from_str.empty()) { + try { + from = static_cast(std::stoull(from_str)); + } catch (const std::exception &) { + return make_error_response(400, "Invalid 'from' offset", ERROR_TYPE_INVALID_REQUEST); + } + } + if (from < session->dropped_prefix()) { + return make_error_response(400, "Stream offset lost, please restart", ERROR_TYPE_INVALID_REQUEST); + } + auto res = std::make_unique(); + res->status = 200; + res->content_type = "text/event-stream"; + // the next closure reads from the ring buffer at the requested offset, blocks until + // bytes arrive or the session finalizes. exit each call after draining the available + // chunk so set_chunked_content_provider gets a chance to flush to the socket + auto offset_ptr = std::make_shared(from); + // consumer pipe: read-only, does not finalize the session on destruction + auto pipe = stream_pipe_consumer::create(session); + res->next = [pipe, offset_ptr, &req](std::string & output) -> bool { + bool got_any = false; + pipe->read(*offset_ptr, + [&](const char * d, size_t n) { + output.append(d, n); + *offset_ptr += n; + got_any = true; + return false; + }, + req.should_stop); + return got_any; + }; + return res; + }; +} + +server_http_context::handler_t make_streams_lookup_handler() { + return [](const server_http_req & req) -> server_http_res_ptr { + // POST /v1/streams/lookup with body {"conversation_ids": ["X", "Y", ...]} returns the + // matching sessions, only for ids the caller already knows. each id matches the exact key + // and any "::" variant, so one lookup covers every per model session for a conv + std::vector requested; + try { + json body = json::parse(req.body); + if (body.contains("conversation_ids") && body["conversation_ids"].is_array()) { + for (const auto & v : body["conversation_ids"]) { + if (v.is_string()) { + std::string id = v.get(); + if (!id.empty()) { + requested.push_back(std::move(id)); + } + } + } + } + } catch (const std::exception & e) { + auto res = std::make_unique(); + res->status = 400; + res->content_type = "application/json; charset=utf-8"; + res->data = safe_json_to_str({{"error", {{"message", std::string("invalid body: ") + e.what()}, + {"type", "invalid_request_error"}}}}); + return res; + } + + std::vector sessions; + if (!requested.empty()) { + auto all = g_stream_sessions.list_all(); + for (const auto & rid : requested) { + const std::string with_sep = rid + "::"; + for (auto & s : all) { + if (s->conversation_id == rid || + s->conversation_id.compare(0, with_sep.size(), with_sep) == 0) { + sessions.push_back(s); + } + } + } + } + + json arr = json::array(); + for (auto & s : sessions) { + arr.push_back({ + {"conversation_id", s->conversation_id}, + {"is_done", s->is_done()}, + {"total_bytes", s->total_size()}, + {"started_at", s->started_ts}, + {"completed_at", s->completed_at()}, + }); + } + auto res = std::make_unique(); + res->status = 200; + res->content_type = "application/json; charset=utf-8"; + res->data = safe_json_to_str(arr); + return res; + }; +} + +server_http_context::handler_t make_stream_delete_handler() { + return [](const server_http_req & req) -> server_http_res_ptr { + // DELETE /v1/stream/ is the explicit user Stop, cancels the producer hook + // wired by handle_completions_impl and evicts the buffer. idempotent, a session that + // already finalized or was never created returns 204 either way + std::string conv_id = req.get_param("conv_id"); + if (conv_id.empty()) { + return make_error_response(400, "Missing conversation id in path", ERROR_TYPE_INVALID_REQUEST); + } + SRV_INF("DELETE /v1/stream/%s -> evict_and_cancel\n", conv_id.c_str()); + g_stream_sessions.evict_and_cancel(conv_id); + auto res = std::make_unique(); + res->status = 204; + res->content_type = "application/json"; + return res; + }; +} + +std::string stream_conv_id_from_headers(const std::map & headers) { + // case-insensitive scan for x-conversation-id + static constexpr char target[] = "x-conversation-id"; + static constexpr size_t target_len = sizeof(target) - 1; + for (const auto & [hk, hv] : headers) { + if (hk.size() != target_len) continue; + bool match = true; + for (size_t i = 0; i < target_len; ++i) { + char c = hk[i]; + if (c >= 'A' && c <= 'Z') c = char(c + 32); + if (c != target[i]) { match = false; break; } + } + if (match) { + return hv; + } + } + return std::string(); +} + +void stream_session_attach_pipe(server_http_res & res, const std::map & headers) { + std::string conversation_id = stream_conv_id_from_headers(headers); + SRV_INF("stream_session_attach_pipe: conv_id=%s (empty=%d)\n", + conversation_id.c_str(), conversation_id.empty() ? 1 : 0); + if (conversation_id.empty()) { + return; + } + auto session = g_stream_sessions.create_or_replace(conversation_id); + res.spipe = stream_pipe_producer::create(session, res); +} + +std::function stream_aware_should_stop(server_http_res * res, std::function fallback) { + return [res, fallback = std::move(fallback)]() -> bool { + if (res->spipe) { + return res->spipe->is_cancelled(); + } + return fallback(); + }; +} diff --git a/tools/server/server-stream.h b/tools/server/server-stream.h new file mode 100644 index 000000000..ff363bb4c --- /dev/null +++ b/tools/server/server-stream.h @@ -0,0 +1,203 @@ +#pragma once + +#include "server-http.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +enum class stream_read_status { + OK, + OFFSET_LOST, +}; + +// streaming buffer for one generation, survives HTTP disconnect. the producer appends raw SSE +// bytes, readers drain from any offset via read_from and block until more bytes or finalize. +// keyed by conversation_id: one conv = at most one live session +struct stream_session { + std::string conversation_id; + int64_t started_ts; // unix seconds at construction, used by /v1/streams listing + + stream_session(std::string conversation_id_, size_t max_bytes_); + stream_session(const stream_session &) = delete; + stream_session & operator=(const stream_session &) = delete; + + // append raw bytes, drops from the front if the cap is reached. + // returns false if the session is already finalized + bool append(const char * data, size_t len); + + // mark the session as complete, wakes all pending readers + void finalize(); + + // drain bytes from offset, calling sink for each chunk. blocks until more + // bytes arrive or finalize is called. returns OK on clean exit, OFFSET_LOST + // if offset falls below the dropped prefix + stream_read_status read_from(size_t offset, + const std::function & sink, + const std::function & should_stop); + + bool is_done() const; + bool is_cancelled() const; + size_t total_size() const; // bytes that ever entered the session + size_t dropped_prefix() const; // bytes evicted from the front due to cap + int64_t completed_at() const; // 0 while alive, unix seconds after finalize + + // attach the producer stop hook used to cancel its reader, pass an empty function to detach + void set_stop_producer(std::function fn); + + // signal the producer to abort its inference asap via the stop hook, idempotent + void cancel(); + +private: + mutable std::mutex mu; + std::condition_variable cv; + std::vector buffer; + size_t prefix_dropped; + size_t cap_bytes; + std::atomic done; + std::atomic cancelled; + std::atomic completed_ts; + std::function stop_producer; // protected by mu +}; + +using stream_session_ptr = std::shared_ptr; + +// one end of a stream_session pipe. the base holds the session and the shared query, the +// producer and consumer ends derive from it. virtual dtor so each end runs its own teardown: +// the producer finalizes the session, the consumer leaves it untouched +struct stream_pipe { + virtual ~stream_pipe() = default; + + // true if the session was cancelled (e.g. via DELETE /v1/stream/) + bool is_cancelled() const; + +protected: + explicit stream_pipe(stream_session_ptr session); + + stream_session_ptr session_; +}; + +// producer end: writes chunks into the ring buffer and owns the session lifetime, finalizing it +// on destruction. +// +// lifetime safety: holds a shared_ptr> alive also captured by the session's +// stop_producer hook. cleanup() sets alive=false and clears the hook; it must run while the +// response the hook calls stop() on is still alive. ~server_res_generator() does this explicitly. +struct stream_pipe_producer : stream_pipe { + ~stream_pipe_producer() override; + + // append raw bytes to the session's ring buffer, returns false if already finalized + bool write(const char * data, size_t len); + + // mark the natural end on the wire so a later close() is a no-op + void done(); + + // on a peer drop, pump the response next() into the ring buffer until done. runs on the http + // worker from on_complete, no-op after done() or cancel + void close(); + + // disarm the stop hook and drop the alive guard, must run while the response the hook + // references is still alive. idempotent, the destructor calls it too + void cleanup(); + + // res.stop() is invoked when the session is cancelled, the alive guard ensures stop() is not + // called after cleanup() has run + static std::shared_ptr create(stream_session_ptr session, server_http_res & res); + +private: + explicit stream_pipe_producer(stream_session_ptr session); + + bool done_ = false; + std::shared_ptr> alive_; + server_http_res * res_ = nullptr; +}; + +// consumer end: read-only replay of the ring buffer, the destructor does not finalize the session +struct stream_pipe_consumer : stream_pipe { + // drain bytes from offset, calling sink for each available chunk. blocks until more data + // arrives or the session finalizes. should_stop is polled, returns OFFSET_LOST if offset + // fell below the dropped prefix + stream_read_status read(size_t & offset, + const std::function & sink, + const std::function & should_stop); + + static std::shared_ptr create(stream_session_ptr session); + +private: + explicit stream_pipe_consumer(stream_session_ptr session); +}; + +// owns all live sessions, runs a periodic GC to evict expired ones. +// the map is keyed by conversation_id, so the invariant "one conv = at most one +// live session" is enforced at the type level +class stream_session_manager { +public: + stream_session_manager(); + ~stream_session_manager(); + + stream_session_manager(const stream_session_manager &) = delete; + stream_session_manager & operator=(const stream_session_manager &) = delete; + + // install a new session for this conversation, evicting and cancelling any previous one. + // the conversation_id must be non empty, the caller is responsible for that check. + // returns the new session + stream_session_ptr create_or_replace(const std::string & conversation_id); + + // lookup, returns null if unknown or already evicted + stream_session_ptr get(const std::string & conversation_id); + + // list every live or recently completed session, used by GET /v1/streams without filter + std::vector list_all() const; + + // remove from the map and finalize, wakes any pending readers + void evict(const std::string & conversation_id); + + // signal the producer to cancel asap then evict, used by the explicit user Stop path + void evict_and_cancel(const std::string & conversation_id); + + void start_gc(); + void stop_gc(); + +private: + void gc_loop(); + + mutable std::shared_mutex map_mu; + std::unordered_map sessions; // key: conversation_id + std::thread gc_thread; + std::atomic running; + std::mutex gc_wake_mu; + std::condition_variable gc_wake_cv; +}; + +// process wide manager, linked by both llama-server and llama-cli. llama-server main() drives +// start_gc/stop_gc, llama-cli leaves it idle. the dtor calls stop_gc() unconditionally so exit +// is safe whether or not the GC thread ran +extern stream_session_manager g_stream_sessions; + +// route handler factories operating on g_stream_sessions, wired under /v1/stream/* by server.cpp. +// keeps the resumable stream surface confined to server-stream +server_http_context::handler_t make_stream_get_handler(); +server_http_context::handler_t make_streams_lookup_handler(); +server_http_context::handler_t make_stream_delete_handler(); + +// extract the X-Conversation-Id header value (case-insensitive), empty when absent. exposed so +// the router can track which child serves a forwarded POST +std::string stream_conv_id_from_headers(const std::map & headers); + +// on an X-Conversation-Id header, create or replace the session and attach a producer pipe to +// res. no-op when absent, called from the server_res_generator constructor +void stream_session_attach_pipe(server_http_res & res, const std::map & headers); + +// should_stop closure that ignores peer disconnect when a pipe is attached, so only an explicit +// DELETE stops the producer and generation keeps flowing into the ring buffer. without a pipe it +// delegates to fallback, the legacy non-resumable flow +std::function stream_aware_should_stop(server_http_res * res, std::function fallback); diff --git a/tools/server/server.cpp b/tools/server/server.cpp index 0a1947faf..1bbc99d89 100644 --- a/tools/server/server.cpp +++ b/tools/server/server.cpp @@ -2,6 +2,7 @@ #include "server-http.h" #include "server-models.h" #include "server-cors-proxy.h" +#include "server-stream.h" #include "server-tools.h" #include "arg.h" @@ -82,6 +83,10 @@ int llama_server(int argc, char ** argv) { common_init(); + // start the stream session manager GC right after common init, before any HTTP route can + // touch it. lifecycle is symmetric, stop_gc() runs in clean_up() before backend free + g_stream_sessions.start_gc(); + if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_SERVER)) { return 1; } @@ -239,6 +244,29 @@ int llama_server(int argc, char ** argv) { ctx_http.get ("/slots", ex_wrapper(routes.get_slots)); ctx_http.post("/slots/:id_slot", ex_wrapper(routes.post_slots)); + // resumable streaming, the conversation_id is the session identity end to end. router and + // child wire different handlers under the same paths: a child binds the local g_stream_sessions + // backed factories, the router binds proxies that resolve the owning child through the + // conv_id -> model map + server_http_context::handler_t stream_get_h; + server_http_context::handler_t streams_lookup_h; + server_http_context::handler_t stream_delete_h; + if (is_router_server) { + stream_get_h = models_routes->router_stream_get; + streams_lookup_h = models_routes->router_streams_lookup; + stream_delete_h = models_routes->router_stream_delete; + } else { + stream_get_h = make_stream_get_handler(); + streams_lookup_h = make_streams_lookup_handler(); + stream_delete_h = make_stream_delete_handler(); + } + ctx_http.get ("/v1/stream/:conv_id", ex_wrapper(stream_get_h)); + // POST /v1/streams/lookup with body {"conversation_ids": [...]}. you can only ask for ids + // you already own (the WebUI passes the convs visible in its sidebar). the server never + // lists ids it has not been asked about, so a random caller cannot enumerate live sessions + ctx_http.post("/v1/streams/lookup", ex_wrapper(streams_lookup_h)); + ctx_http.del ("/v1/stream/:conv_id", ex_wrapper(stream_delete_h)); + // Google Cloud Platform (Vertex AI) compat ctx_http.register_gcp_compat(); @@ -314,6 +342,8 @@ int llama_server(int argc, char ** argv) { clean_up = [&models_routes]() { SRV_INF("%s: cleaning up before exit...\n", __func__); + // stop the session GC first, it finalizes live sessions and wakes pending readers + g_stream_sessions.stop_gc(); if (models_routes.has_value()) { models_routes->stopping.store(true); // maybe redundant, but just to be safe models_routes->models.unload_all(); @@ -340,6 +370,8 @@ int llama_server(int argc, char ** argv) { // setup clean up function, to be called before exit clean_up = [&ctx_http, &ctx_server]() { SRV_INF("%s: cleaning up before exit...\n", __func__); + // stop the session GC first, it finalizes live sessions and wakes pending readers + g_stream_sessions.stop_gc(); ctx_http.stop(); ctx_server.terminate(); llama_backend_free(); diff --git a/tools/ui/src/lib/components/app/chat/ChatAttachments/ChatAttachmentsList/ChatAttachmentsListItem/ChatAttachmentsListItemMcpPrompt.svelte b/tools/ui/src/lib/components/app/chat/ChatAttachments/ChatAttachmentsList/ChatAttachmentsListItem/ChatAttachmentsListItemMcpPrompt.svelte index 636e93f22..c55dfdec7 100644 --- a/tools/ui/src/lib/components/app/chat/ChatAttachments/ChatAttachmentsList/ChatAttachmentsListItem/ChatAttachmentsListItemMcpPrompt.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatAttachments/ChatAttachmentsList/ChatAttachmentsListItem/ChatAttachmentsListItemMcpPrompt.svelte @@ -33,7 +33,7 @@ {#if !readonly && onRemove}
onRemove?.()} />
diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUserPending.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUserPending.svelte index 4be582b39..5c2913202 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUserPending.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageUser/ChatMessageUserPending.svelte @@ -56,7 +56,7 @@
diff --git a/tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreen.svelte b/tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreen.svelte index 18635ba39..d2fba5a93 100644 --- a/tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreen.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreen.svelte @@ -5,6 +5,7 @@ ChatMessages, ChatScreenDragOverlay, ChatScreenProcessingInfo, + ChatScreenStreamResumeStatus, ServerLoadingSplash, ChatScreenServerError } from '$lib/components/app'; @@ -281,6 +282,10 @@ + {#if page.params.id} + + {/if} +
{#if (isMobile.current ? mobileScrollDownHint || isMobileUserScrolledUp : autoScroll.userScrolledUp) && page.url.hash.includes(ROUTES.CHAT) && page.params.id} + import { chatStore } from '$lib/stores/chat.svelte'; + import { StreamConnectionState } from '$lib/enums'; + import { Loader2 } from '@lucide/svelte'; + + let state = $derived(chatStore.streamConnectionState); + + +{#if state === StreamConnectionState.RESUMING} +
+ + Reconnecting to the stream... +
+{/if} diff --git a/tools/ui/src/lib/components/app/chat/index.ts b/tools/ui/src/lib/components/app/chat/index.ts index 517f24d74..d7004d3ae 100644 --- a/tools/ui/src/lib/components/app/chat/index.ts +++ b/tools/ui/src/lib/components/app/chat/index.ts @@ -683,3 +683,11 @@ export { default as ChatScreenProcessingInfo } from './ChatScreen/ChatScreenProc * Rendered inside ChatScreen when `serverError` store has a value. */ export { default as ChatScreenServerError } from './ChatScreen/ChatScreenServerError.svelte'; + +/** + * Stream resume status indicator. Shows a small "Reconnecting to the stream..." + * banner with a spinner while `chatStore.streamConnectionState` is `resuming`, + * i.e. after a dropped connection is reattaching to the live SSE replay buffer. + * Renders nothing otherwise. Shown inside ChatScreen only on an active conversation route. + */ +export { default as ChatScreenStreamResumeStatus } from './ChatScreen/ChatScreenStreamResumeStatus.svelte'; diff --git a/tools/ui/src/lib/components/app/navigation/SidebarNavigation/SidebarNavigationConversationItem.svelte b/tools/ui/src/lib/components/app/navigation/SidebarNavigation/SidebarNavigationConversationItem.svelte index b1c2b78f6..2c1b9adf2 100644 --- a/tools/ui/src/lib/components/app/navigation/SidebarNavigation/SidebarNavigationConversationItem.svelte +++ b/tools/ui/src/lib/components/app/navigation/SidebarNavigation/SidebarNavigationConversationItem.svelte @@ -39,7 +39,6 @@ depth = 0 }: Props = $props(); - let renderActionsDropdown = $state(false); let dropdownOpen = $state(false); let isLoading = $derived(getAllLoadingChats().includes(conversation.id)); @@ -71,26 +70,10 @@ } } - function handleMouseLeave() { - if (!dropdownOpen) { - renderActionsDropdown = false; - } - } - - function handleMouseOver() { - renderActionsDropdown = true; - } - function handleSelect() { onSelect?.(conversation.id); } - $effect(() => { - if (!dropdownOpen) { - renderActionsDropdown = false; - } - }); - onMount(() => { document.addEventListener('edit-active-conversation', handleGlobalEditEvent as EventListener); @@ -103,23 +86,19 @@ }); - -
{#if depth > 0} @@ -130,7 +109,7 @@ @@ -146,18 +125,15 @@ {#if isLoading} -
e.key === 'Enter' && handleStop(e)} - role="button" - tabindex="0" aria-label="Stop generation" >
+ @@ -169,52 +145,50 @@
- {#if renderActionsDropdown} -
- { - e.stopPropagation(); - handleTogglePin(); - } - }, - { - icon: Pencil, - label: 'Edit', - onclick: handleEdit, - shortcut: ['shift', 'cmd', 'e'] - }, - { - icon: Download, - label: 'Export', - onclick: (e: Event) => { - e.stopPropagation(); - conversationsStore.downloadConversation(conversation.id); - }, - shortcut: ['shift', 'cmd', 's'] - }, - { - icon: Trash2, - label: 'Delete', - onclick: handleDelete, - variant: 'destructive', - shortcut: ['shift', 'cmd', 'd'], - separator: true +
+ { + e.stopPropagation(); + handleTogglePin(); } - ]} - /> -
- {/if} - + }, + { + icon: Pencil, + label: 'Edit', + onclick: handleEdit, + shortcut: ['shift', 'cmd', 'e'] + }, + { + icon: Download, + label: 'Export', + onclick: (e: Event) => { + e.stopPropagation(); + conversationsStore.downloadConversation(conversation.id); + }, + shortcut: ['shift', 'cmd', 's'] + }, + { + icon: Trash2, + label: 'Delete', + onclick: handleDelete, + variant: 'destructive', + shortcut: ['shift', 'cmd', 'd'], + separator: true + } + ]} + /> +
+