diff --git a/CMakeLists.txt b/CMakeLists.txt index db32112fa..d4c04abb6 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -389,8 +389,6 @@ endif() # Build libraries # -add_subdirectory(vendor/hash) - add_library(ggml ggml/src/ggml.c ggml/include/ggml.h @@ -497,12 +495,22 @@ add_library(common2 otherarch/llmutils.h common/reasoning-budget.cpp common/reasoning-budget.h + vendor/hash/hash.cpp + vendor/hash/hash.h + vendor/hash/sha256/sha256.c + vendor/hash/sha256/sha256.h tools/mtmd/mtmd-audio.cpp tools/mtmd/mtmd-audio.h) -target_include_directories(common2 PUBLIC . ./ggml/include ./ggml/src ./ggml/src/ggml-cpu ./include ./otherarch ./otherarch/tools ./vendor/stb ./vendor/nlohmann ./vendor ./otherarch/sdcpp ./otherarch/sdcpp/thirdparty ./tools ./common) +target_include_directories(common2 PUBLIC . ./ggml/include ./ggml/src ./ggml/src/ggml-cpu ./include ./otherarch ./otherarch/tools ./vendor/stb ./vendor/nlohmann ./vendor ./vendor/hash ./otherarch/sdcpp ./otherarch/sdcpp/thirdparty ./tools ./common) target_compile_features(common2 PUBLIC cxx_std_17) # don't bump -target_link_libraries(common2 PRIVATE ggml vendor-hash ${LLAMA_EXTRA_LIBS}) +target_link_libraries(common2 PRIVATE ggml ${LLAMA_EXTRA_LIBS}) set_target_properties(common2 PROPERTIES POSITION_INDEPENDENT_CODE ON) +if (CMAKE_C_COMPILER_ID STREQUAL "MSVC") + set(KCPP_HASH_NO_WARN_FLAG /w) +else() + set(KCPP_HASH_NO_WARN_FLAG -w) +endif() +set_source_files_properties(vendor/hash/sha256/sha256.c PROPERTIES COMPILE_OPTIONS ${KCPP_HASH_NO_WARN_FLAG}) add_library(sdtype_adapter otherarch/sdcpp/sdtype_adapter.cpp diff --git a/common/arg.cpp b/common/arg.cpp index ac59481b1..4235e303e 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -4659,6 +4659,12 @@ void common_params_add_preset_options(std::vector & args) { [](common_params &, int) { /* unused */ } ).set_env(COMMON_ARG_PRESET_STOP_TIMEOUT).set_preset_only()); + args.push_back(common_arg( + {"dedup-cache-models"}, "0|1", + "in server router mode, hide a cached model from the model list when this preset resolves to the same model file", + [](common_params &, const std::string &) { /* unused */ } + ).set_env(COMMON_ARG_PRESET_DEDUP_CACHE_MODELS).set_preset_only()); + // args.push_back(common_arg( // {"pin"}, // "in server router mode, do not unload this model if models_max is exceeded", diff --git a/common/arg.h b/common/arg.h index 44b9e887c..421bc295f 100644 --- a/common/arg.h +++ b/common/arg.h @@ -11,8 +11,9 @@ #include // pseudo-env variable to identify preset-only arguments -#define COMMON_ARG_PRESET_LOAD_ON_STARTUP "__PRESET_LOAD_ON_STARTUP" -#define COMMON_ARG_PRESET_STOP_TIMEOUT "__PRESET_STOP_TIMEOUT" +#define COMMON_ARG_PRESET_LOAD_ON_STARTUP "__PRESET_LOAD_ON_STARTUP" +#define COMMON_ARG_PRESET_STOP_TIMEOUT "__PRESET_STOP_TIMEOUT" +#define COMMON_ARG_PRESET_DEDUP_CACHE_MODELS "__PRESET_DEDUP_CACHE_MODELS" // // CLI argument parsing diff --git a/common/common.cpp b/common/common.cpp index b0d2cc599..47831f64e 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1784,6 +1784,8 @@ void common_threadpools::init(llama_context * ctx, const common_params & params) struct ggml_threadpool_params tpp = ggml_threadpool_params_from_cpu_params(params.cpuparams); + // each pool needs to match the respective n_threads exactly + // see: https://github.com/ggml-org/llama.cpp/pull/27138#issuecomment-5332307332 if (!ggml_threadpool_params_match(&tpp, &tpp_batch)) { threadpool_batch = ggml_threadpool_new_fn(&tpp_batch); if (!threadpool_batch) { diff --git a/common/download.cpp b/common/download.cpp index 44c6cea42..2509f75ab 100644 --- a/common/download.cpp +++ b/common/download.cpp @@ -989,6 +989,26 @@ std::vector common_list_cached_models() { return result; } +std::string common_download_resolve_path(const std::string & hf_repo_with_tag, const std::string & hf_file) { + auto [repo, tag] = common_download_split_repo_tag(hf_repo_with_tag); + + auto files = hf_cache::get_cached_files(repo); + if (files.empty()) { + return ""; + } + + if (!hf_file.empty()) { + for (const auto & f : files) { + if (f.path == hf_file) { + return f.local_path; + } + } + return ""; + } + + return find_best_model(files, tag).local_path; +} + bool common_download_remove(const std::string & hf_repo_with_tag) { namespace fs = std::filesystem; diff --git a/common/download.h b/common/download.h index 9da595d1f..a6162f7bf 100644 --- a/common/download.h +++ b/common/download.h @@ -86,6 +86,10 @@ std::vector common_download_get_all_parts(const std::string & url); // returns list of cached models std::vector common_list_cached_models(); +// resolve the local cached file path for a HF repo without network access (hf_file, if given, must match exactly) +// returns an empty string if the model is not present in the cache +std::string common_download_resolve_path(const std::string & hf_repo_with_tag, const std::string & hf_file = ""); + // download single file from url to local path // returns status code or -1 on error // skip_etag: if true, don't read/write .etag files (for HF cache where filename is the hash) diff --git a/conversion/__init__.py b/conversion/__init__.py index 3232a1050..5ae6ad819 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -109,6 +109,8 @@ TEXT_MODEL_MAP: dict[str, str] = { "GraniteSwitchForCausalLM": "granite", "GraniteSpeechForConditionalGeneration": "granite", "GraniteSpeechPlusForConditionalGeneration": "granite", + "GraniteSWAForCausalLM": "granite", + "GraniteMoeSWAForCausalLM": "granite", "Grok1ForCausalLM": "grok", "GrokForCausalLM": "grok", "GroveMoeForCausalLM": "grovemoe", diff --git a/conversion/granite.py b/conversion/granite.py index 5f1e3e847..796d37cca 100644 --- a/conversion/granite.py +++ b/conversion/granite.py @@ -74,6 +74,108 @@ class GraniteModel(LlamaModel): return super().filter_tensors(item) +@ModelBase.register("GraniteSWAForCausalLM") +class GraniteSWAModel(GraniteModel): + """Conversion for IBM's GraniteSWAForCausalLM (interleaved sliding window attention)""" + model_arch = gguf.MODEL_ARCH.GRANITE_SWA + + @classmethod + def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: + name, gen = item + + if name.endswith("sinks"): + name += ".weight" + + return super().filter_tensors((name, gen)) + + def set_gguf_parameters(self): + """GraniteSWA uses Granite parameters plus sliding window configuration.""" + super().set_gguf_parameters() + + # Add sliding_window from config + sliding_window = self.hparams.get("sliding_window", 128) + self.gguf_writer.add_sliding_window(sliding_window) + logger.info("gguf: (granite_swa) sliding_window = %s", sliding_window) + + # Derive sliding_window_pattern from layer_types + if layer_types := self.hparams.get("layer_types"): + is_swa = [t == "sliding_attention" for t in layer_types] + self.gguf_writer.add_sliding_window_pattern(is_swa) + logger.info("gguf: (granite_swa) sliding_window_pattern = %d SWA layers / %d total", + sum(is_swa), len(is_swa)) + else: + # Fall back to period-based pattern: i % 4 != 0 + # This matches the transformers default pattern + n_layers = self.block_count + is_swa = [i % 4 != 0 for i in range(n_layers)] + self.gguf_writer.add_sliding_window_pattern(is_swa) + logger.info("gguf: (granite_swa) sliding_window_pattern (inferred) = %d SWA layers / %d total", + sum(is_swa), n_layers) + + # Add rope_pattern from no_rope_layers + if no_rope_layers := self.hparams.get("no_rope_layers"): + # Convert 1/0 to bool (1 = use RoPE, 0 = NoPE) + rope_pattern = [bool(x) for x in no_rope_layers] + self.gguf_writer.add_rope_pattern(rope_pattern) + logger.info("gguf: (granite_swa) rope_pattern = %d RoPE layers / %d total", + sum(rope_pattern), len(rope_pattern)) + + +@ModelBase.register("GraniteMoeSWAForCausalLM") +class GraniteMoeSWAModel(GraniteSWAModel): + """Conversion for IBM's GraniteMoeSWAForCausalLM (unified dense + MoE with iSWA)""" + model_arch = gguf.MODEL_ARCH.GRANITE_SWA + + def set_gguf_parameters(self): + super().set_gguf_parameters() + if shared_intermediate_size := self.hparams.get("shared_intermediate_size"): + self.gguf_writer.add_expert_shared_feed_forward_length(shared_intermediate_size) + logger.info("gguf: (granitemoewa) shared_intermediate_size = %s", shared_intermediate_size) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + """Split merged MoE tensors (gate+up) following standard MoE pattern.""" + + # Handle expert FFN tensors (merged gate+up) - swash format: experts.gate_up_proj + # Kept fused since inference (build_moe_ffn) supports a single gate_up_exps + # tensor for the routed experts. + if name.endswith("block_sparse_moe.experts.gate_up_proj"): + ffn_dim = self.hparams["intermediate_size"] + assert data_torch.shape[-2] == 2 * ffn_dim, f"Merged FFN tensor size must be 2 * intermediate_size, got {data_torch.shape[-2]}" + yield from ModelBase.modify_tensors(self, data_torch, self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE_UP_EXP, bid), bid) + return + + # Handle expert FFN down projection - swash format: experts.down_proj + if name.endswith("block_sparse_moe.experts.down_proj"): + yield from ModelBase.modify_tensors(self, data_torch, self.format_tensor_name(gguf.MODEL_TENSOR.FFN_DOWN_EXP, bid), bid) + return + + # Handle expert FFN tensors (merged gate+up) - standard granite format: input_linear.weight + # Kept fused since inference (build_moe_ffn) supports a single gate_up_exps + # tensor for the routed experts. + if name.endswith("block_sparse_moe.input_linear.weight"): + ffn_dim = self.hparams["intermediate_size"] + assert data_torch.shape[-2] == 2 * ffn_dim, "Merged FFN tensor size must be 2 * intermediate_size" + yield from ModelBase.modify_tensors(self, data_torch, self.format_tensor_name(gguf.MODEL_TENSOR.FFN_GATE_UP_EXP, bid), bid) + return + + # Handle shared expert FFN tensors (if present) - kept fused since + # inference (build_ffn) supports a single ffn_up_shexp tensor with + # LLM_FFN_SWIGLU for the shared expert. + if name.endswith("shared_mlp.input_linear.weight"): + ffn_dim = self.hparams.get("shared_intermediate_size", self.hparams["intermediate_size"]) + assert data_torch.shape[-2] == 2 * ffn_dim, "Merged FFN tensor size must be 2 * shared_intermediate_size" + yield from ModelBase.modify_tensors(self, data_torch, self.format_tensor_name(gguf.MODEL_TENSOR.FFN_UP_SHEXP, bid), bid) + return + + # Handle shared expert output (if present) + if name.endswith("shared_mlp.output_linear.weight"): + yield from ModelBase.modify_tensors(self, data_torch, self.format_tensor_name(gguf.MODEL_TENSOR.FFN_DOWN_SHEXP, bid), bid) + return + + # Pass through to parent for all other tensors (including sinks) + yield from super().modify_tensors(data_torch, name, bid) + + @ModelBase.register("GraniteMoeForCausalLM", "GraniteMoeSharedForCausalLM") @ModelBase.example("ibm-granite/granite-3.1-3b-a800m-instruct") class GraniteMoeModel(GraniteModel): diff --git a/ggml/include/ggml-rpc.h b/ggml/include/ggml-rpc.h index 276aea00e..059e44962 100644 --- a/ggml/include/ggml-rpc.h +++ b/ggml/include/ggml-rpc.h @@ -7,7 +7,7 @@ extern "C" { #endif #define RPC_PROTO_MAJOR_VERSION 5 -#define RPC_PROTO_MINOR_VERSION 0 +#define RPC_PROTO_MINOR_VERSION 1 #define RPC_PROTO_PATCH_VERSION 0 #ifdef __cplusplus diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 604cc6902..40c6e3bd8 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -1999,6 +1999,14 @@ extern "C" { float beta_fast, float beta_slow); + // set the offset dims for RoPE + // a must be GGML_OP_ROPE or GGML_OP_ROPE_BACK + // vision RoPE is not supported + // example: (marking: x = rotated, 0 = unrotated) + // n_embd = 10, n_dims = 4, offset = 2 --> [00xxxx0000] + GGML_API struct ggml_tensor * ggml_rope_set_offset( + struct ggml_tensor * a, + int n_offs); // clamp // in-place, returns view(a) diff --git a/ggml/src/ggml-backend-impl.h b/ggml/src/ggml-backend-impl.h index 9c56ec30c..40cea024c 100644 --- a/ggml/src/ggml-backend-impl.h +++ b/ggml/src/ggml-backend-impl.h @@ -83,6 +83,7 @@ extern "C" { GGML_API ggml_backend_buffer_t ggml_backend_multi_buffer_alloc_buffer(ggml_backend_buffer_t * buffers, size_t n_buffers); GGML_API bool ggml_backend_buffer_is_multi_buffer(ggml_backend_buffer_t buffer); GGML_API void ggml_backend_multi_buffer_set_usage(ggml_backend_buffer_t buffer, enum ggml_backend_buffer_usage usage); + GGML_API void ggml_backend_meta_buffer_set_usage (ggml_backend_buffer_t buffer, enum ggml_backend_buffer_usage usage); // // Backend (meta) diff --git a/ggml/src/ggml-backend-meta.cpp b/ggml/src/ggml-backend-meta.cpp index 7654ea1f3..775ae9926 100644 --- a/ggml/src/ggml-backend-meta.cpp +++ b/ggml/src/ggml-backend-meta.cpp @@ -1118,7 +1118,6 @@ static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state( } static struct ggml_backend_meta_split_state ggml_backend_meta_get_split_state(const struct ggml_tensor * tensor, bool assume_sync) { - GGML_ASSERT(ggml_backend_buffer_is_meta(tensor->buffer)); ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) tensor->buffer->context; return ggml_backend_meta_get_split_state(buf_ctx->get_simple_tensor_container(tensor), tensor, assume_sync); } @@ -1178,7 +1177,15 @@ static enum ggml_status ggml_backend_meta_buffer_init_tensor_impl(ggml_backend_m t_ij->flags = tensor->flags; memcpy(t_ij->op_params, tensor->op_params, sizeof(tensor->op_params)); ggml_set_name(t_ij, tensor->name); + t_ij->buffer = simple_buf; + if (simple_buf) { + // the backend that owns the buffer will set .extra + ggml_backend_buffer_init_tensor(simple_buf, t_ij); + } else { + t_ij->extra = tensor->extra; + } + t_ij->view_src = tensor->view_src; t_ij->view_offs = tensor->view_offs; if (t_ij->view_src != nullptr && ggml_backend_buffer_is_meta(t_ij->view_src->buffer)) { @@ -1209,7 +1216,6 @@ static enum ggml_status ggml_backend_meta_buffer_init_tensor_impl(ggml_backend_m t_ij->data = (char *) ggml_backend_buffer_get_base(simple_buf) + size_t(tensor->data) - size_t(ggml_backend_buffer_get_base(tensor->buffer)); } - t_ij->extra = tensor->extra; for (int i = 0; i < GGML_MAX_SRC; i++) { t_ij->src[i] = tensor->src[i]; if (tensor->src[i] == tensor) { @@ -1502,6 +1508,16 @@ bool ggml_backend_buffer_is_meta(ggml_backend_buffer_t buf) { return buf != nullptr && buf->iface.free_buffer == ggml_backend_meta_buffer_iface.free_buffer; } +void ggml_backend_meta_buffer_set_usage(ggml_backend_buffer_t buffer, enum ggml_backend_buffer_usage usage) { + GGML_ASSERT(ggml_backend_buffer_is_meta(buffer)); + ggml_backend_meta_buffer_context * buf_ctx = (ggml_backend_meta_buffer_context *) buffer->context; + for (size_t i = 0; i < buf_ctx->bufs.size(); i++) { + if (buf_ctx->bufs[i]) { + ggml_backend_buffer_set_usage(buf_ctx->bufs[i].get(), usage); + } + } +} + static ggml_backend_buffer_t ggml_backend_meta_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { const size_t n_simple_bufts = ggml_backend_meta_buft_n_bufts(buft); diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 8e65c448e..1f35c5f86 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -182,6 +182,8 @@ void ggml_backend_buffer_set_usage(ggml_backend_buffer_t buffer, enum ggml_backe // FIXME: add a generic callback to the buffer interface if (ggml_backend_buffer_is_multi_buffer(buffer)) { ggml_backend_multi_buffer_set_usage(buffer, usage); + } else if (ggml_backend_buffer_is_meta(buffer)) { + ggml_backend_meta_buffer_set_usage(buffer, usage); } } diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 001e1ae85..2b5f68444 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -5979,6 +5979,8 @@ static void ggml_compute_forward_rope_flt( memcpy(&beta_slow, (int32_t *) dst->op_params + 10, sizeof(float)); memcpy(§ions, (int32_t *) dst->op_params + 11, sizeof(int)*4); + const int n_offs = ((int32_t *) dst->op_params)[15]; + GGML_TENSOR_UNARY_OP_LOCALS //printf("ne0: %d, ne1: %d, ne2: %d, ne3: %d\n", ne0, ne1, ne2, ne3); @@ -5995,6 +5997,10 @@ static void ggml_compute_forward_rope_flt( GGML_ASSERT(n_dims <= ne0); GGML_ASSERT(n_dims % 2 == 0); + GGML_ASSERT(n_offs >= 0); + GGML_ASSERT(n_offs % 2 == 0); + GGML_ASSERT(n_offs + n_dims <= ne0); + // rows per thread const int dr = (nr + nth - 1)/nth; @@ -6020,6 +6026,7 @@ static void ggml_compute_forward_rope_flt( if (is_vision) { GGML_ASSERT(n_dims == ne0/2); + GGML_ASSERT(n_offs == 0); } const float * freq_factors = NULL; @@ -6068,12 +6075,12 @@ static void ggml_compute_forward_rope_flt( switch (mode) { case GGML_ROPE_TYPE_NORMAL: - rotate_pairs(n_dims, 1, cache, src, dst_data, 1); + rotate_pairs(n_dims, 1, cache, src + n_offs, dst_data + n_offs, 1); break; case GGML_ROPE_TYPE_NEOX: case GGML_ROPE_TYPE_MROPE: case GGML_ROPE_TYPE_IMROPE: - rotate_pairs(n_dims, n_dims/2, cache, src, dst_data); + rotate_pairs(n_dims, n_dims/2, cache, src + n_offs, dst_data + n_offs); break; case GGML_ROPE_TYPE_VISION: rotate_pairs(ne0, n_dims, cache, src, dst_data); @@ -6084,7 +6091,11 @@ static void ggml_compute_forward_rope_flt( if (!is_vision) { // fill the remain channels with data from src tensor - for (int64_t i0 = n_dims; i0 < ne0; i0 += 2) { + for (int64_t i0 = 0; i0 < ne0; i0 += 2) { + if (i0 == n_offs) { + i0 += n_dims - 2; // skip the rotated channels + continue; + } const T * const src = (T *)((char *) src0->data + i3*nb03 + i2*nb02 + i1*nb01 + i0*nb00); T * dst_data = (T *)((char *) dst->data + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0); diff --git a/ggml/src/ggml-cpu/simd-mappings.h b/ggml/src/ggml-cpu/simd-mappings.h index fca5119e1..10ce4bfc5 100644 --- a/ggml/src/ggml-cpu/simd-mappings.h +++ b/ggml/src/ggml-cpu/simd-mappings.h @@ -29,13 +29,15 @@ extern "C" { // FP16 to FP32 conversion // 16-bit float -// on Arm, we use __fp16 +// on Arm, we use __fp16, which requires the IEEE fp16 format: implied on +// AArch64, selected by -mfp16-format=ieee on 32 bit Arm, where the compiler +// may otherwise reject the type // on x86, we use uint16_t // // for old CUDA compilers (<= 11), we use uint16_t: ref https://github.com/ggml-org/llama.cpp/pull/10616 // for MUSA compilers , we use uint16_t: ref https://github.com/ggml-org/llama.cpp/pull/11843 // -#if defined(__ARM_NEON) && !(defined(__CUDACC__) && __CUDACC_VER_MAJOR__ <= 11) && !defined(__MUSACC__) +#if defined(__ARM_NEON) && defined(__ARM_FP16_FORMAT_IEEE) && !(defined(__CUDACC__) && __CUDACC_VER_MAJOR__ <= 11) && !defined(__MUSACC__) #define GGML_CPU_COMPUTE_FP16_TO_FP32(x) neon_compute_fp16_to_fp32(x) #define GGML_CPU_COMPUTE_FP32_TO_FP16(x) neon_compute_fp32_to_fp16(x) @@ -326,7 +328,7 @@ inline static float ggml_lookup_fp16_to_fp32(ggml_fp16_t f) { #define GGML_F16_VEC_REDUCE GGML_F32Cx4_REDUCE #endif -#elif defined(__ARM_NEON) && defined(__ARM_FEATURE_FMA) +#elif defined(__ARM_NEON) && defined(__ARM_FEATURE_FMA) && defined(__ARM_FP16_FORMAT_IEEE) #define GGML_SIMD diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 53a12ee2a..a68544f66 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -2731,6 +2731,12 @@ static bool ggml_cuda_should_fuse_rms_norm_mul_rope(const ggml_tensor * rms_norm return false; } + // ggml_rope_set_offset is not yet supported in the fused kernel + const int n_offs = ((const int32_t *) rope->op_params)[15]; + if (n_offs != 0) { + return false; + } + return true; } diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index 0589e65bd..c99923804 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -4,6 +4,7 @@ #include "vecdotq.cuh" #include +#include typedef float (*vec_dot_q_cuda_t)(const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs); @@ -69,7 +70,8 @@ enum mmvq_parameter_table_id { MMVQ_PARAMETERS_GCN, MMVQ_PARAMETERS_RDNA2, MMVQ_PARAMETERS_RDNA3_0, - MMVQ_PARAMETERS_RDNA4 + MMVQ_PARAMETERS_RDNA4, + MMVQ_PARAMETERS_GB10 }; static constexpr __device__ mmvq_parameter_table_id get_device_table_id() { @@ -83,6 +85,8 @@ static constexpr __device__ mmvq_parameter_table_id get_device_table_id() { return MMVQ_PARAMETERS_GCN; #elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= GGML_CUDA_CC_TURING && __CUDA_ARCH__ < GGML_CUDA_CC_AMPERE return MMVQ_PARAMETERS_TURING; +#elif defined(__CUDA_ARCH__) && __CUDA_ARCH__ == GGML_CUDA_CC_DGX_SPARK + return MMVQ_PARAMETERS_GB10; #else return MMVQ_PARAMETERS_GENERIC; #endif @@ -104,6 +108,9 @@ static __host__ mmvq_parameter_table_id get_device_table_id(int cc) { if (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_TURING && ggml_cuda_highest_compiled_arch(cc) < GGML_CUDA_CC_AMPERE) { return MMVQ_PARAMETERS_TURING; } + if (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) == GGML_CUDA_CC_DGX_SPARK) { + return MMVQ_PARAMETERS_GB10; + } return MMVQ_PARAMETERS_GENERIC; } @@ -351,7 +358,7 @@ static constexpr __device__ int get_mmvq_mmid_max_batch_for_device() { #endif } -static constexpr __host__ __device__ int calc_nwarps(ggml_type type, int ncols_dst, mmvq_parameter_table_id table_id) { +static constexpr __host__ __device__ int calc_nwarps(ggml_type type, int ncols_dst, mmvq_parameter_table_id table_id, bool small_k = false, bool halve_iters = false) { if (table_id == MMVQ_PARAMETERS_GENERIC) { switch (ncols_dst) { case 1: @@ -454,11 +461,32 @@ static constexpr __host__ __device__ int calc_nwarps(ggml_type type, int ncols_d return 1; } } + if (table_id == MMVQ_PARAMETERS_GB10) { + const int generic = calc_nwarps(type, ncols_dst, MMVQ_PARAMETERS_GENERIC); + // Only worth the wider block when it actually retires the K loop in half the trips (Observation) + if (ncols_dst == 1 && !small_k && halve_iters) { + switch (type) { + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_IQ4_NL: + return 2 * generic; + default: + break; + } + } + return generic; + } return 1; } static constexpr __host__ __device__ int calc_rows_per_block(int ncols_dst, int table_id, bool small_k = false, int nwarps = 1) { - if (table_id == MMVQ_PARAMETERS_GENERIC || table_id == MMVQ_PARAMETERS_GCN || table_id == MMVQ_PARAMETERS_TURING) { + if (table_id == MMVQ_PARAMETERS_GENERIC || table_id == MMVQ_PARAMETERS_GCN || table_id == MMVQ_PARAMETERS_TURING || table_id == MMVQ_PARAMETERS_GB10) { switch (ncols_dst) { case 1: return small_k ? nwarps : 1; @@ -477,8 +505,8 @@ static constexpr __host__ __device__ int calc_rows_per_block(int ncols_dst, int return 1; } -template -__launch_bounds__(calc_nwarps(type, ncols_dst, get_device_table_id())*ggml_cuda_get_physical_warp_size(), 1) +template +__launch_bounds__(calc_nwarps(type, ncols_dst, get_device_table_id(), small_k, halve_iters)*ggml_cuda_get_physical_warp_size(), 1) static __global__ void mul_mat_vec_q( const void * vx_ptr, const void * vy_ptr, const int32_t * ids_ptr, const ggml_cuda_mm_fusion_args_device fusion, float * dst_ptr, const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t stride_row_x, const uint32_t stride_col_y, @@ -495,7 +523,7 @@ static __global__ void mul_mat_vec_q( constexpr int qi = ggml_cuda_type_traits::qi; constexpr int vdr = get_vdr_mmvq(type); constexpr mmvq_parameter_table_id table_id = get_device_table_id(); - constexpr int nwarps = calc_nwarps(type, ncols_dst, table_id); + constexpr int nwarps = calc_nwarps(type, ncols_dst, table_id, small_k, halve_iters); constexpr int rows_per_cuda_block = calc_rows_per_block(ncols_dst, table_id, small_k, nwarps); constexpr int warp_size = ggml_cuda_get_physical_warp_size(); @@ -773,8 +801,8 @@ static __global__ void mul_mat_vec_q_moe( template static std::pair calc_launch_params( const int ncols_dst, const int nrows_x, const int nchannels_dst, const int nsamples_or_ntokens, - const int warp_size, const mmvq_parameter_table_id table_id, const bool small_k = false) { - const int nwarps = calc_nwarps(type, ncols_dst, table_id); + const int warp_size, const mmvq_parameter_table_id table_id, const bool small_k = false, const bool halve_iters = false) { + const int nwarps = calc_nwarps(type, ncols_dst, table_id, small_k, halve_iters); const int rpb = calc_rows_per_block(ncols_dst, table_id, small_k, nwarps); const int64_t nblocks = (nrows_x + rpb - 1) / rpb; const dim3 block_nums(nblocks, nchannels_dst, nsamples_or_ntokens); @@ -782,7 +810,7 @@ static std::pair calc_launch_params( return {block_nums, block_dims}; } -template +template static void mul_mat_vec_q_switch_fusion( const void * vx, const void * vy, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst, const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t stride_row_x, const uint32_t stride_col_y, @@ -797,7 +825,7 @@ static void mul_mat_vec_q_switch_fusion( if constexpr (c_ncols_dst == 1) { if (has_fusion) { const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, nbytes_shared, stream); - ggml_cuda_kernel_launch(mul_mat_vec_q, launch_params, + ggml_cuda_kernel_launch(mul_mat_vec_q, launch_params, vx, vy, ids, fusion, dst, ncols_x, nchannels_y, stride_row_x, stride_col_y, stride_col_dst, channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride); @@ -808,7 +836,7 @@ static void mul_mat_vec_q_switch_fusion( GGML_ASSERT(!has_fusion && "fusion only supported for ncols_dst=1"); const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, nbytes_shared, stream); - ggml_cuda_kernel_launch(mul_mat_vec_q, launch_params, + ggml_cuda_kernel_launch(mul_mat_vec_q, launch_params, vx, vy, ids, fusion, dst, ncols_x, nchannels_y, stride_row_x, stride_col_y, stride_col_dst, channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride); @@ -860,16 +888,18 @@ static void mul_mat_vec_q_switch_ncols_dst( const bool has_ids = ids != nullptr; + // How the K loop divides up at the baseline block width, both decisions below use these. + constexpr int qk = ggml_cuda_type_traits::qk; + constexpr int qi = ggml_cuda_type_traits::qi; + constexpr int vdr = get_vdr_mmvq(type); + const int blocks_per_row_x = ncols_x / qk; + const int blocks_per_iter_1warp = vdr * warp_size / qi; + const auto should_use_small_k = [&](int c_ncols_dst) { // When K is small, increase rows_per_block to match nwarps so each warp has more work to do // Trigger when the full thread block covers all K blocks in a single loop iteration and few threads remain idle. - constexpr int qk = ggml_cuda_type_traits::qk; - constexpr int qi = ggml_cuda_type_traits::qi; - constexpr int vdr = get_vdr_mmvq(type); - const int blocks_per_row_x = ncols_x / qk; - const int blocks_per_iter_1warp = vdr * warp_size / qi; - const int nwarps = calc_nwarps(type, c_ncols_dst, table_id); - bool use = nwarps > 1 && blocks_per_row_x < nwarps * blocks_per_iter_1warp; + const int nwarps = calc_nwarps(type, c_ncols_dst, table_id); + bool use = nwarps > 1 && blocks_per_row_x < nwarps * blocks_per_iter_1warp; constexpr std::array iq_slow_turing = { GGML_TYPE_IQ3_XXS, @@ -902,6 +932,28 @@ static void mul_mat_vec_q_switch_ncols_dst( return use; }; + // Whether doubling nwarps pays off on the ncols_dst == 1 path, where K sets the K loop trip count. + const auto should_halve_iters = [&] { + if (table_id != MMVQ_PARAMETERS_GB10) { + return false; + } + + // Expert rows are gathered per token, so a wider block adds reduction work without reuse. + if (has_ids) { + return false; + } + + const int blocks_per_iter = calc_nwarps(type, 1, table_id) * blocks_per_iter_1warp; + const int iters = (blocks_per_row_x + blocks_per_iter - 1) / blocks_per_iter; + const int iters_wide = (blocks_per_row_x + blocks_per_iter * 2 - 1) / (blocks_per_iter * 2); + + // An odd trip count leaves half the wider block idle for its last iteration, that tail is + // only affordable once the loop is long enough to dilute it to an eighth of the work (observation). + const int idle = iters_wide * 2 - iters; + + return idle * 8 <= iters_wide * 2; + }; + if (has_ids && ncols_dst > 1) { // Multi-token MUL_MAT_ID path - dedicated MoE kernel mul_mat_vec_q_moe_launch( @@ -914,26 +966,34 @@ static void mul_mat_vec_q_switch_ncols_dst( switch (ncols_dst) { case 1: { - constexpr int c_ncols_dst = 1; + // static, else MSVC lambda capture breaks the constexpr uses below + static constexpr int c_ncols_dst = 1; - bool use_small_k = should_use_small_k(c_ncols_dst); + // Tag types keep the flags compile-time, so __launch_bounds__ matches what is launched. + const auto launch = [&](auto small_k_tag, auto halve_iters_tag) { + constexpr bool c_small_k = decltype(small_k_tag)::value; + // Types the table does not promote would compile a second, identical kernel. + constexpr bool c_promoted = + calc_nwarps(type, c_ncols_dst, MMVQ_PARAMETERS_GB10, false, true) != + calc_nwarps(type, c_ncols_dst, MMVQ_PARAMETERS_GB10, false, false); - if (use_small_k) { - std::pair dims = calc_launch_params(c_ncols_dst, nrows_x, nchannels_dst, - nsamples_dst, warp_size, table_id, true); - mul_mat_vec_q_switch_fusion( + constexpr bool c_halve_iters = decltype(halve_iters_tag)::value && c_promoted; + + const std::pair dims = calc_launch_params(c_ncols_dst, nrows_x, nchannels_dst, + nsamples_dst, warp_size, table_id, c_small_k, c_halve_iters); + mul_mat_vec_q_switch_fusion( vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, dims.first, dims.second, 0, ids_stride, stream); + }; + + if (should_use_small_k(c_ncols_dst)) { + launch(std::true_type{}, std::false_type{}); + } else if (should_halve_iters()) { + launch(std::false_type{}, std::true_type{}); } else { - std::pair dims = calc_launch_params(c_ncols_dst, nrows_x, nchannels_dst, - nsamples_dst, warp_size, table_id); - mul_mat_vec_q_switch_fusion( - vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, - channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio_fd, - stride_sample_x, stride_sample_y, stride_sample_dst, dims.first, dims.second, 0, ids_stride, - stream); + launch(std::false_type{}, std::false_type{}); } } break; case 2: { diff --git a/ggml/src/ggml-cuda/rope.cu b/ggml/src/ggml-cuda/rope.cu index 504c6b818..e546fb655 100644 --- a/ggml/src/ggml-cuda/rope.cu +++ b/ggml/src/ggml-cuda/rope.cu @@ -53,6 +53,7 @@ static __global__ void rope_norm(const T * x, const int s2, const int s3, const int n_dims, + const int n_offs, const int32_t * pos, const float freq_scale, const float ext_factor, @@ -61,7 +62,8 @@ static __global__ void rope_norm(const T * x, const float theta_scale, const float * freq_factors, const int64_t * row_indices, - const int set_rows_stride) { + const int set_rows_stride, + const bool inplace) { const int i0 = 2*(blockDim.y*blockIdx.y + threadIdx.y); if (i0 >= ne00) { @@ -92,19 +94,24 @@ static __global__ void rope_norm(const T * x, ggml_cuda_memcpy_1<4>(dst + idst, &v); } }; - if (i0 >= n_dims) { + if (i0 < n_offs || i0 >= n_offs + n_dims) { + if (inplace) { + return; + } store_coaelsced(x[ix + 0], x[ix + 1]); return; } - const float theta_base = pos[i2]*powf(theta_scale, i0/2.0f); + const int iw = i0 - n_offs; // relative idx - const float freq_factor = has_ff ? freq_factors[i0/2] : 1.0f; + const float theta_base = pos[i2]*powf(theta_scale, iw/2.0f); + + const float freq_factor = has_ff ? freq_factors[iw/2] : 1.0f; float cos_theta; float sin_theta; - rope_yarn(theta_base/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor, cos_theta, sin_theta); + rope_yarn(theta_base/freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor, cos_theta, sin_theta); const float x0 = x[ix + 0]; const float x1 = x[ix + 1]; @@ -125,6 +132,7 @@ static __global__ void rope_neox(const T * x, const int s2, const int s3, const int n_dims, + const int n_offs, const int32_t * pos, const float freq_scale, const float ext_factor, @@ -133,7 +141,8 @@ static __global__ void rope_neox(const T * x, const float theta_scale, const float * freq_factors, const int64_t * row_indices, - const int set_rows_stride) { + const int set_rows_stride, + const bool inplace) { ggml_cuda_pdl_lc(); const int i0 = 2*(blockDim.y*blockIdx.y + threadIdx.y); @@ -158,27 +167,33 @@ static __global__ void rope_neox(const T * x, idst += row_indices[i2] * set_rows_stride; } - if (i0 >= n_dims) { + if (i0 < n_offs || i0 >= n_offs + n_dims) { + if (inplace) { + return; + } dst[idst + i0 / 2 + 0] = ggml_cuda_cast(x[ix + i0 / 2 + 0]); dst[idst + i0 / 2 + 1] = ggml_cuda_cast(x[ix + i0 / 2 + 1]); return; } - const float theta_base = pos[i2]*powf(theta_scale, i0/2.0f); + const int iw = i0 - n_offs; // relative idx - const float freq_factor = has_ff ? freq_factors[i0/2] : 1.0f; + const float theta_base = pos[i2]*powf(theta_scale, iw/2.0f); + + const float freq_factor = has_ff ? freq_factors[iw/2] : 1.0f; float cos_theta; float sin_theta; - rope_yarn(theta_base/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor, cos_theta, sin_theta); + rope_yarn(theta_base/freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor, cos_theta, sin_theta); - const float x0 = x[ix + 0]; - const float x1 = x[ix + n_dims/2]; + // idst/ix point at channel i0/2; the first channel of the rotated pair is n_offs + iw/2 = i0/2 + n_offs/2 + const float x0 = x[ix + n_offs/2 + 0]; + const float x1 = x[ix + n_offs/2 + n_dims/2]; - dst[idst + 0] = ggml_cuda_cast(x0 * cos_theta - x1 * sin_theta); - dst[idst + n_dims / 2] = ggml_cuda_cast(x0 * sin_theta + x1 * cos_theta); + dst[idst + n_offs/2 + 0] = ggml_cuda_cast(x0 * cos_theta - x1 * sin_theta); + dst[idst + n_offs/2 + n_dims / 2] = ggml_cuda_cast(x0 * sin_theta + x1 * cos_theta); } template @@ -194,6 +209,7 @@ static __global__ void rope_multi(const T * x, const int s2, const int s3, const int n_dims, + const int n_offs, const int32_t * pos, const float freq_scale, const float ext_factor, @@ -202,7 +218,8 @@ static __global__ void rope_multi(const T * x, const float theta_scale, const float * freq_factors, const mrope_sections sections, - const bool is_imrope) { + const bool is_imrope, + const bool inplace) { const int i0 = 2 * (blockDim.y * blockIdx.y + threadIdx.y); if (i0 >= ne00) { @@ -219,52 +236,58 @@ static __global__ void rope_multi(const T * x, const int ix = i0 / 2 + i1 * s01 + i2 * s02 + i3 * s03; ggml_cuda_pdl_sync(); - if (i0 >= n_dims) { + if (i0 < n_offs || i0 >= n_offs + n_dims) { + if (inplace) { + return; + } dst[idst + i0/2 + 0] = x[ix + i0/2 + 0]; dst[idst + i0/2 + 1] = x[ix + i0/2 + 1]; return; } + const int iw = i0 - n_offs; // relative idx + const int sect_dims = sections.v[0] + sections.v[1] + sections.v[2] + sections.v[3]; const int sec_w = sections.v[1] + sections.v[0]; - const int sector = (i0 / 2) % sect_dims; + const int sector = (iw / 2) % sect_dims; float theta_base = 0.0; if (is_imrope) { if (sector % 3 == 1 && sector < 3 * sections.v[1]) { // h - theta_base = pos[i2 + ne02 * 1] * powf(theta_scale, i0 / 2.0f); + theta_base = pos[i2 + ne02 * 1] * powf(theta_scale, iw / 2.0f); } else if (sector % 3 == 2 && sector < 3 * sections.v[2]) { // w - theta_base = pos[i2 + ne02 * 2] * powf(theta_scale, i0 / 2.0f); + theta_base = pos[i2 + ne02 * 2] * powf(theta_scale, iw / 2.0f); } else if (sector % 3 == 0 && sector < 3 * sections.v[0]) { // t - theta_base = pos[i2] * powf(theta_scale, i0 / 2.0f); + theta_base = pos[i2] * powf(theta_scale, iw / 2.0f); } else { - theta_base = pos[i2 + ne02 * 3] * powf(theta_scale, i0 / 2.0f); + theta_base = pos[i2 + ne02 * 3] * powf(theta_scale, iw / 2.0f); } } else { if (sector < sections.v[0]) { - theta_base = pos[i2] * powf(theta_scale, i0 / 2.0f); + theta_base = pos[i2] * powf(theta_scale, iw / 2.0f); } else if (sector >= sections.v[0] && sector < sec_w) { - theta_base = pos[i2 + ne02 * 1] * powf(theta_scale, i0 / 2.0f); + theta_base = pos[i2 + ne02 * 1] * powf(theta_scale, iw / 2.0f); } else if (sector >= sec_w && sector < sec_w + sections.v[2]) { - theta_base = pos[i2 + ne02 * 2] * powf(theta_scale, i0 / 2.0f); + theta_base = pos[i2 + ne02 * 2] * powf(theta_scale, iw / 2.0f); } else if (sector >= sec_w + sections.v[2]) { - theta_base = pos[i2 + ne02 * 3] * powf(theta_scale, i0 / 2.0f); + theta_base = pos[i2 + ne02 * 3] * powf(theta_scale, iw / 2.0f); } } - const float freq_factor = has_ff ? freq_factors[i0/2] : 1.0f; + const float freq_factor = has_ff ? freq_factors[iw/2] : 1.0f; float cos_theta; float sin_theta; - rope_yarn(theta_base/freq_factor, freq_scale, corr_dims, i0, ext_factor, attn_factor, cos_theta, sin_theta); + rope_yarn(theta_base/freq_factor, freq_scale, corr_dims, iw, ext_factor, attn_factor, cos_theta, sin_theta); - const float x0 = x[ix + 0]; - const float x1 = x[ix + n_dims/2]; + // idst/ix point at channel i0/2; the first channel of the rotated pair is n_offs + iw/2 = i0/2 + n_offs/2 + const float x0 = x[ix + n_offs/2 + 0]; + const float x1 = x[ix + n_offs/2 + n_dims/2]; - dst[idst + 0] = x0*cos_theta - x1*sin_theta; - dst[idst + n_dims/2] = x0*sin_theta + x1*cos_theta; + dst[idst + n_offs/2 + 0] = x0*cos_theta - x1*sin_theta; + dst[idst + n_offs/2 + n_dims/2] = x0*sin_theta + x1*cos_theta; } template @@ -344,6 +367,7 @@ static void rope_norm_cuda(const T * x, const int s2, const int s3, const int n_dims, + const int n_offs, const int nr, const int32_t * pos, const float freq_scale, @@ -354,6 +378,7 @@ static void rope_norm_cuda(const T * x, const float * freq_factors, const int64_t * row_indices, const int set_rows_stride, + const bool inplace, cudaStream_t stream) { GGML_ASSERT(ne00 % 2 == 0); const dim3 block_dims(1, CUDA_ROPE_BLOCK_SIZE, 1); @@ -364,12 +389,12 @@ static void rope_norm_cuda(const T * x, if (freq_factors == nullptr) { rope_norm<<>>( - x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, pos, freq_scale, ext_factor, - attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride); + x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, n_offs, pos, freq_scale, ext_factor, + attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride, inplace); } else { rope_norm<<>>( - x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, pos, freq_scale, ext_factor, - attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride); + x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, n_offs, pos, freq_scale, ext_factor, + attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride, inplace); } } @@ -386,6 +411,7 @@ static void rope_neox_cuda(const T * x, const int s2, const int s3, const int n_dims, + const int n_offs, const int nr, const int32_t * pos, const float freq_scale, @@ -396,6 +422,7 @@ static void rope_neox_cuda(const T * x, const float * freq_factors, const int64_t * row_indices, const int set_rows_stride, + const bool inplace, cudaStream_t stream) { GGML_ASSERT(ne00 % 2 == 0); const dim3 block_dims(1, CUDA_ROPE_BLOCK_SIZE, 1); @@ -407,12 +434,12 @@ static void rope_neox_cuda(const T * x, if (freq_factors == nullptr) { ggml_cuda_kernel_launch(rope_neox, launch_params, - x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, pos, freq_scale, ext_factor, - attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride); + x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, n_offs, pos, freq_scale, ext_factor, + attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride, inplace); } else { ggml_cuda_kernel_launch(rope_neox, launch_params, - x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, pos, freq_scale, ext_factor, - attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride); + x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, n_offs, pos, freq_scale, ext_factor, + attn_factor, corr_dims, theta_scale, freq_factors, row_indices, set_rows_stride, inplace); } } @@ -429,6 +456,7 @@ static void rope_multi_cuda(const T * x, const int s2, const int s3, const int n_dims, + const int n_offs, const int nr, const int32_t * pos, const float freq_scale, @@ -439,6 +467,7 @@ static void rope_multi_cuda(const T * x, const float * freq_factors, const mrope_sections sections, const bool is_imrope, + const bool inplace, cudaStream_t stream) { GGML_ASSERT(ne00 % 2 == 0); const dim3 block_dims(1, CUDA_ROPE_BLOCK_SIZE, 1); @@ -450,13 +479,13 @@ static void rope_multi_cuda(const T * x, if (freq_factors == nullptr) { const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream); ggml_cuda_kernel_launch(rope_multi, launch_params, - x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, pos, freq_scale, ext_factor, - attn_factor, corr_dims, theta_scale, freq_factors, sections, is_imrope); + x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, n_offs, pos, freq_scale, ext_factor, + attn_factor, corr_dims, theta_scale, freq_factors, sections, is_imrope, inplace); } else { const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream); ggml_cuda_kernel_launch(rope_multi, launch_params, - x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, pos, freq_scale, ext_factor, - attn_factor, corr_dims, theta_scale, freq_factors, sections, is_imrope); + x, dst, ne00, ne01, ne02, s01, s02, s03, s1, s2, s3, n_dims, n_offs, pos, freq_scale, ext_factor, + attn_factor, corr_dims, theta_scale, freq_factors, sections, is_imrope, inplace); } } @@ -552,8 +581,12 @@ void ggml_cuda_op_rope_impl(ggml_backend_cuda_context & ctx, const int mode = ((int32_t *) dst->op_params)[2]; //const int n_ctx = ((int32_t *) dst->op_params)[3]; const int n_ctx_orig = ((int32_t *) dst->op_params)[4]; + const int n_offs = ((int32_t *) dst->op_params)[15]; mrope_sections sections; + // when dst aliases src0, the channels outside the rotated window already hold the correct data + const bool inplace = dst_d == src0->data; + // RoPE alteration for extended context float freq_base; float freq_scale; @@ -581,6 +614,7 @@ void ggml_cuda_op_rope_impl(ggml_backend_cuda_context & ctx, if (is_vision) { GGML_ASSERT(n_dims == ne00/2); + GGML_ASSERT(n_offs == 0); // offset not supported for vision, as the rotated pairs span the whole row } const int32_t * pos = (const int32_t *) src1_d; @@ -597,31 +631,31 @@ void ggml_cuda_op_rope_impl(ggml_backend_cuda_context & ctx, if (is_neox) { if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F32) { rope_neox_cuda((const float *) src0_d, (float *) dst_d, ne00, ne01, ne02, s01, s02, - s03, s1, s2, s3, n_dims, nr, pos, freq_scale, freq_base, + s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims, freq_factors, row_indices, - set_rows_stride, stream); + set_rows_stride, inplace, stream); } else if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F16) { rope_neox_cuda((const float *) src0_d, (half *) dst_d, ne00, ne01, ne02, s01, s02, - s03, s1, s2, s3, n_dims, nr, pos, freq_scale, freq_base, + s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims, freq_factors, row_indices, - set_rows_stride, stream); + set_rows_stride, inplace, stream); } else if (src0->type == GGML_TYPE_F16 && dst_type == GGML_TYPE_F16) { rope_neox_cuda((const half *) src0_d, (half *) dst_d, ne00, ne01, ne02, s01, s02, - s03, s1, s2, s3, n_dims, nr, pos, freq_scale, freq_base, + s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims, freq_factors, row_indices, - set_rows_stride, stream); + set_rows_stride, inplace, stream); } else { GGML_ABORT("fatal error"); } } else if (is_mrope && !is_vision) { if (src0->type == GGML_TYPE_F32) { rope_multi_cuda((const float *) src0_d, (float *) dst_d, ne00, ne01, ne02, s01, s02, s03, s1, - s2, s3, n_dims, nr, pos, freq_scale, freq_base, ext_factor, attn_factor, - corr_dims, freq_factors, sections, is_imrope, stream); + s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base, ext_factor, attn_factor, + corr_dims, freq_factors, sections, is_imrope, inplace, stream); } else if (src0->type == GGML_TYPE_F16) { rope_multi_cuda((const half *) src0_d, (half *) dst_d, ne00, ne01, ne02, s01, s02, s03, s1, - s2, s3, n_dims, nr, pos, freq_scale, freq_base, ext_factor, attn_factor, - corr_dims, freq_factors, sections, is_imrope, stream); + s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base, ext_factor, attn_factor, + corr_dims, freq_factors, sections, is_imrope, inplace, stream); } else { GGML_ABORT("fatal error"); } @@ -640,19 +674,19 @@ void ggml_cuda_op_rope_impl(ggml_backend_cuda_context & ctx, } else { if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F32) { rope_norm_cuda((const float *) src0_d, (float *) dst_d, ne00, ne01, ne02, s01, s02, - s03, s1, s2, s3, n_dims, nr, pos, freq_scale, freq_base, + s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims, freq_factors, row_indices, - set_rows_stride, stream); + set_rows_stride, inplace, stream); } else if (src0->type == GGML_TYPE_F32 && dst_type == GGML_TYPE_F16) { rope_norm_cuda((const float *) src0_d, (half *) dst_d, ne00, ne01, ne02, s01, s02, - s03, s1, s2, s3, n_dims, nr, pos, freq_scale, freq_base, + s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims, freq_factors, row_indices, - set_rows_stride, stream); + set_rows_stride, inplace, stream); } else if (src0->type == GGML_TYPE_F16 && dst_type == GGML_TYPE_F16) { rope_norm_cuda((const half *) src0_d, (half *) dst_d, ne00, ne01, ne02, s01, s02, - s03, s1, s2, s3, n_dims, nr, pos, freq_scale, freq_base, + s03, s1, s2, s3, n_dims, n_offs, nr, pos, freq_scale, freq_base, ext_factor, attn_factor, corr_dims, freq_factors, row_indices, - set_rows_stride, stream); + set_rows_stride, inplace, stream); } else { GGML_ABORT("fatal error"); } diff --git a/ggml/src/ggml-metal/ggml-metal-impl.h b/ggml/src/ggml-metal/ggml-metal-impl.h index 1f6e8c48b..05ea7470e 100644 --- a/ggml/src/ggml-metal/ggml-metal-impl.h +++ b/ggml/src/ggml-metal/ggml-metal-impl.h @@ -329,6 +329,7 @@ typedef struct { uint64_t nb3; int32_t n_past; int32_t n_dims; + int32_t n_offs; int32_t n_ctx_orig; float freq_base; float freq_scale; @@ -341,6 +342,7 @@ typedef struct { int32_t sect_2; int32_t sect_3; bool src2; + bool inplace; } ggml_metal_kargs_rope; typedef struct { diff --git a/ggml/src/ggml-metal/ggml-metal-ops.cpp b/ggml/src/ggml-metal/ggml-metal-ops.cpp index b7f9b2d0d..d8435e957 100644 --- a/ggml/src/ggml-metal/ggml-metal-ops.cpp +++ b/ggml/src/ggml-metal/ggml-metal-ops.cpp @@ -3884,6 +3884,11 @@ int ggml_metal_op_rope(ggml_metal_op_t ctx, int idx) { const int sect_2 = ((const int32_t *) op->op_params)[13]; const int sect_3 = ((const int32_t *) op->op_params)[14]; + const int n_offs = ((const int32_t *) op->op_params)[15]; + + // when dst aliases src0, the channels outside the rotated window already hold the correct data + const bool inplace = op->data == op->src[0]->data; + ggml_metal_kargs_rope args = { /*.ne00 =*/ ne00, /*.ne01 =*/ ne01, @@ -3903,6 +3908,7 @@ int ggml_metal_op_rope(ggml_metal_op_t ctx, int idx) { /*.nb3 =*/ nb3, /*.n_past =*/ n_past, /*.n_dims =*/ n_dims, + /*.n_offs =*/ n_offs, /*.n_ctx_orig =*/ n_ctx_orig, /*.freq_base =*/ freq_base, /*.freq_scale =*/ freq_scale, @@ -3915,6 +3921,7 @@ int ggml_metal_op_rope(ggml_metal_op_t ctx, int idx) { /* sect_2 =*/ sect_2, /* sect_3 =*/ sect_3, /* src2 =*/ op->src[2] != nullptr, + /* inplace =*/ inplace, }; auto pipeline = ggml_metal_library_get_pipeline_rope(lib, op); diff --git a/ggml/src/ggml-metal/ggml-metal.metal b/ggml/src/ggml-metal/ggml-metal.metal index 243c997fc..0537fa4cf 100644 --- a/ggml/src/ggml-metal/ggml-metal.metal +++ b/ggml/src/ggml-metal/ggml-metal.metal @@ -656,13 +656,13 @@ void dequantize_q5_1_t4(device const block_q5_1 * xb, short il, thread type4 & r template void dequantize_q8_0(device const block_q8_0 *xb, short il, thread type4x4 & reg) { - device const int8_t * qs = ((device const int8_t *)xb->qs); + device const packed_char4 * qs = (device const packed_char4 *) xb->qs; const float d = xb->d; float4x4 reg_f; - for (int i = 0; i < 16; i++) { - reg_f[i/4][i%4] = (qs[i + 16*il] * d); + for (int i = 0; i < 4; ++i) { + reg_f[i] = float4(qs[4*il + i]) * d; } reg = (type4x4) reg_f; @@ -670,12 +670,10 @@ void dequantize_q8_0(device const block_q8_0 *xb, short il, thread type4x4 & reg template void dequantize_q8_0_t4(device const block_q8_0 *xb, short il, thread type4 & reg) { - device const int8_t * qs = ((device const int8_t *)xb->qs); + device const packed_char4 * qs = (device const packed_char4 *) xb->qs; const float d = xb->d; - for (int i = 0; i < 4; i++) { - reg[i] = (qs[4*(il%4) + i + 16*(il/4)] * d); - } + reg = (type4) (float4(qs[il]) * d); } template @@ -4688,14 +4686,15 @@ kernel void kernel_rope_norm( float sin_theta; for (int i0 = 2*tiitg; i0 < args.ne0; i0 += 2*tptg.x) { - if (i0 < args.n_dims) { - const int ic = i0/2; + if (i0 >= args.n_offs && i0 < args.n_offs + args.n_dims) { + const int iw = i0 - args.n_offs; // relative idx + const int ic = iw/2; - const float theta = theta_base * pow(args.freq_base, inv_ndims*i0); + const float theta = theta_base * pow(args.freq_base, inv_ndims*iw); const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f; - rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, i0, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta); + rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, iw, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta); device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + i0*args.nb00); device T * dst_data = (device T *)( dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0); @@ -4706,6 +4705,10 @@ kernel void kernel_rope_norm( dst_data[0] = x0*cos_theta - x1*sin_theta; dst_data[1] = x0*sin_theta + x1*cos_theta; } else { + if (args.inplace) { + continue; + } + device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + i0*args.nb00); device T * dst_data = (device T *)( dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0); @@ -4741,17 +4744,18 @@ kernel void kernel_rope_neox( float sin_theta; for (int i0 = 2*tiitg; i0 < args.ne0; i0 += 2*tptg.x) { - if (i0 < args.n_dims) { - const int ic = i0/2; + if (i0 >= args.n_offs && i0 < args.n_offs + args.n_dims) { + const int iw = i0 - args.n_offs; // relative idx + const int ic = iw/2; - const float theta = theta_base * pow(args.freq_base, inv_ndims*i0); + const float theta = theta_base * pow(args.freq_base, inv_ndims*iw); const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f; - rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, i0, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta); + rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, iw, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta); - device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + ic*args.nb00); - device T * dst_data = (device T *)( dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + ic*args.nb0); + device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + (args.n_offs + ic)*args.nb00); + device T * dst_data = (device T *)( dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + (args.n_offs + ic)*args.nb0); const float x0 = src[0]; const float x1 = src[args.n_dims/2]; @@ -4759,6 +4763,10 @@ kernel void kernel_rope_neox( dst_data[0] = x0*cos_theta - x1*sin_theta; dst_data[args.n_dims/2] = x0*sin_theta + x1*cos_theta; } else { + if (args.inplace) { + continue; + } + device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + i0*args.nb00); device T * dst_data = (device T *)( dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0); @@ -4793,8 +4801,9 @@ kernel void kernel_rope_multi( float sin_theta; for (int i0 = 2*tiitg; i0 < args.ne0; i0 += 2*tptg.x) { - if (i0 < args.n_dims) { - const int ic = i0/2; + if (i0 >= args.n_offs && i0 < args.n_offs + args.n_dims) { + const int iw = i0 - args.n_offs; // relative idx + const int ic = iw/2; // mrope theta calculations // note: the rest is the same as kernel_rope_neox @@ -4827,14 +4836,14 @@ kernel void kernel_rope_multi( } // end of mrope - const float theta = theta_base * pow(args.freq_base, inv_ndims*i0); + const float theta = theta_base * pow(args.freq_base, inv_ndims*iw); const float freq_factor = args.src2 ? ((device const float *) src2)[ic] : 1.0f; - rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, i0, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta); + rope_yarn(theta/freq_factor, args.freq_scale, corr_dims, iw, args.ext_factor, args.attn_factor, &cos_theta, &sin_theta); - device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + ic*args.nb00); - device T * dst_data = (device T *)( dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + ic*args.nb0); + device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + (args.n_offs + ic)*args.nb00); + device T * dst_data = (device T *)( dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + (args.n_offs + ic)*args.nb0); const float x0 = src[0]; const float x1 = src[args.n_dims/2]; @@ -4842,6 +4851,10 @@ kernel void kernel_rope_multi( dst_data[0] = x0*cos_theta - x1*sin_theta; dst_data[args.n_dims/2] = x0*sin_theta + x1*cos_theta; } else { + if (args.inplace) { + continue; + } + device const T * const src = (device T *)(src0 + i3*args.nb03 + i2*args.nb02 + i1*args.nb01 + i0*args.nb00); device T * dst_data = (device T *)( dst + i3*args.nb3 + i2*args.nb2 + i1*args.nb1 + i0*args.nb0); diff --git a/ggml/src/ggml-rpc/ggml-rpc.cpp b/ggml/src/ggml-rpc/ggml-rpc.cpp index f48552f31..c4edd7190 100644 --- a/ggml/src/ggml-rpc/ggml-rpc.cpp +++ b/ggml/src/ggml-rpc/ggml-rpc.cpp @@ -49,7 +49,7 @@ struct rpc_tensor { uint64_t data; char name[RPC_TENSOR_NAME_SIZE]; - char padding[4]; + int32_t use_count; }; static_assert(RPC_TENSOR_NAME_SIZE == 64, "rpc_tensor name size must match the upstream RPC wire ABI"); @@ -450,7 +450,7 @@ static rpc_tensor serialize_tensor(const ggml_tensor * tensor) { // Avoid sending uninitialized data over the wire memset(result.name, 0, sizeof(result.name)); - memset(result.padding, 0, sizeof(result.padding)); + result.use_count = 0; snprintf(result.name, sizeof(result.name), "%s", tensor->name); return result; @@ -678,7 +678,7 @@ static void ggml_backend_rpc_synchronize(ggml_backend_t backend) { // this is no-op because we don't have any async operations } -static void add_tensor(ggml_tensor * tensor, std::vector & tensors, std::unordered_set & visited) { +static void add_tensor(ggml_tensor * tensor, const ggml_cgraph * cgraph, std::vector & tensors, std::unordered_set & visited) { if (tensor == nullptr) { return; } @@ -687,10 +687,15 @@ static void add_tensor(ggml_tensor * tensor, std::vector & tensors, } visited.insert(tensor); for (int i = 0; i < GGML_MAX_SRC; i++) { - add_tensor(tensor->src[i], tensors, visited); + add_tensor(tensor->src[i], cgraph, tensors, visited); } - add_tensor(tensor->view_src, tensors, visited); - tensors.push_back(serialize_tensor(tensor)); + add_tensor(tensor->view_src, cgraph, tensors, visited); + rpc_tensor result = serialize_tensor(tensor); + const size_t hash_pos = ggml_hash_find(&cgraph->visited_hash_set, tensor); + if (hash_pos != GGML_HASHSET_FULL && ggml_bitset_get(cgraph->visited_hash_set.used, hash_pos)) { + result.use_count = cgraph->use_counts[hash_pos]; + } + tensors.push_back(result); } static void serialize_graph(uint32_t device, const ggml_cgraph * cgraph, std::vector & output) { @@ -698,7 +703,7 @@ static void serialize_graph(uint32_t device, const ggml_cgraph * cgraph, std::ve std::vector tensors; std::unordered_set visited; for (uint32_t i = 0; i < n_nodes; i++) { - add_tensor(cgraph->nodes[i], tensors, visited); + add_tensor(cgraph->nodes[i], cgraph, tensors, visited); } // serialization format: // | device (4 bytes) | n_nodes (4 bytes) | nodes (n_nodes * sizeof(uint64_t) | n_tensors (4 bytes) | tensors (n_tensors * sizeof(rpc_tensor)) | @@ -1454,6 +1459,10 @@ bool rpc_server::graph_compute(const std::vector & input) { GGML_LOG_ERROR("[%s] failed to create graph node %d (id=%" PRId64 ")\n", __func__, i, id); return false; } + if (graph->nodes[i] != nullptr) { + const size_t hash_pos = ggml_hash_insert(&graph->visited_hash_set, graph->nodes[i]); + graph->use_counts[hash_pos] = tensor_ptrs.at(id)->use_count; + } } ggml_status status = ggml_backend_graph_compute(backends[device], graph); GGML_ASSERT(status == GGML_STATUS_SUCCESS && "Unsuccessful graph computations are not supported with RPC"); diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index ff86fe2a1..cc44cdd79 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -919,6 +919,7 @@ struct vk_device_struct { vk_pipeline pipeline_quantize_q8_1_x4; vk_pipeline pipeline_dequant[GGML_TYPE_COUNT]; + vk_pipeline pipeline_dequant_transpose[GGML_TYPE_COUNT]; // fused dequant+transpose for FA quant-KV vk_pipeline pipeline_dequant_mul_mat_vec_f32_f32[DMMV_WG_SIZE_COUNT][GGML_TYPE_COUNT][mul_mat_vec_max_cols]; vk_pipeline pipeline_dequant_mul_mat_vec_f16_f32[DMMV_WG_SIZE_COUNT][GGML_TYPE_COUNT][mul_mat_vec_max_cols]; vk_pipeline pipeline_dequant_mul_mat_vec_id_f32[DMMV_WG_SIZE_COUNT][GGML_TYPE_COUNT]; @@ -968,6 +969,7 @@ struct vk_device_struct { vk_pipeline pipeline_cpy_f32_quant[GGML_TYPE_COUNT]; vk_pipeline pipeline_cpy_quant_f32[GGML_TYPE_COUNT]; vk_pipeline pipeline_cpy_transpose_16, pipeline_cpy_transpose_32; + vk_pipeline pipeline_cpy_transpose_02_16, pipeline_cpy_transpose_02_32; // [src0 0=fp32,1=fp16][dst] vk_pipeline pipeline_set_rows_i32[2][GGML_TYPE_COUNT]; vk_pipeline pipeline_set_rows_i64[2][GGML_TYPE_COUNT]; @@ -1650,6 +1652,7 @@ struct vk_op_rope_push_constants { uint32_t rope_mode; uint32_t nrows; uint32_t n_dims; + uint32_t n_offs; float freq_scale; float freq_base; float ext_factor; @@ -3388,10 +3391,10 @@ static void ggml_vk_queue_command_pools_cleanup(vk_device& device) { // Arbitrary frequency to cleanup/reuse command buffers static constexpr uint32_t cleanup_frequency = 10; - if (device->compute_queue->cmd_pool.buffers_in_use() >= cleanup_frequency) { + if (device->compute_queue && device->compute_queue->cmd_pool.buffers_in_use() >= cleanup_frequency) { ggml_vk_command_pool_cleanup(device, device->compute_queue->cmd_pool); } - if (device->transfer_queue->cmd_pool.buffers_in_use() >= cleanup_frequency) { + if (device->transfer_queue && device->transfer_queue->cmd_pool.buffers_in_use() >= cleanup_frequency) { ggml_vk_command_pool_cleanup(device, device->transfer_queue->cmd_pool); } } @@ -5395,6 +5398,7 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q5_0], "dequant_q5_0", dequant_q5_0_len, dequant_q5_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q5_1], "dequant_q5_1", dequant_q5_1_len, dequant_q5_1_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q8_0], "dequant_q8_0", dequant_q8_0_len, dequant_q8_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant_transpose[GGML_TYPE_Q8_0], "dequant_q8_0_transpose", dequant_q8_0_transpose_len, dequant_q8_0_transpose_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q2_K], "dequant_q2_k", dequant_q2_k_len, dequant_q2_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_TQ2_0], "dequant_tq2_0", dequant_tq2_0_len, dequant_tq2_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q3_K], "dequant_q3_k", dequant_q3_k_len, dequant_q3_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); @@ -5531,6 +5535,8 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_cpy_transpose_32, "cpy_transpose_32", cpy_transpose_32_len, cpy_transpose_32_data, "main", 2, sizeof(vk_op_unary_push_constants), {1, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_cpy_transpose_16, "cpy_transpose_16", cpy_transpose_16_len, cpy_transpose_16_data, "main", 2, sizeof(vk_op_unary_push_constants), {1, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_cpy_transpose_02_32, "cpy_transpose_02_32", cpy_transpose_02_32_len, cpy_transpose_02_32_data, "main", 2, sizeof(vk_op_unary_push_constants), {1, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_cpy_transpose_02_16, "cpy_transpose_02_16", cpy_transpose_02_16_len, cpy_transpose_02_16_data, "main", 2, sizeof(vk_op_unary_push_constants), {1, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_quant[GGML_TYPE_Q1_0], "cpy_f32_q1_0", cpy_f32_q1_0_len, cpy_f32_q1_0_data, "main", 2, sizeof(vk_op_unary_push_constants), {32, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_quant[GGML_TYPE_Q2_0], "cpy_f32_q2_0", cpy_f32_q2_0_len, cpy_f32_q2_0_data, "main", 2, sizeof(vk_op_unary_push_constants), {32, 1, 1}, {}, 1); @@ -8964,6 +8970,18 @@ static vk_pipeline ggml_vk_get_cpy_pipeline(ggml_backend_vk_context * ctx, const } } + // Same, for a 0<->2 swap: src dim2 is the innermost dimension. + bool transpose02 = dst && !contig && src->nb[2] == ggml_type_size(to) && + ggml_is_contiguous(dst) && ggml_are_same_shape(dst, src); + + if (transpose02 && src->type == to) { + if (ggml_type_size(to) == 4) { + return ctx->device->pipeline_cpy_transpose_02_32; + } else if (ggml_type_size(to) == 2) { + return ctx->device->pipeline_cpy_transpose_02_16; + } + } + if (src->type == GGML_TYPE_F32 && to == GGML_TYPE_F32) { if (contig) { return ctx->device->pipeline_contig_cpy_f32_f32; @@ -10840,9 +10858,32 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx const bool f32acc = !ctx->device->fp16 || dst->op_params[3] == GGML_PREC_F32 || k->type == GGML_TYPE_BF16; + // dequant K/V once into an f16 scratch, reordered KV layout so FA can read without a stride + auto is_dense_kv_cache = [](const ggml_tensor * t) { + return t->nb[0] == ggml_type_size(t->type) && + t->nb[2] == ggml_row_size(t->type, t->ne[0]) && + t->nb[1] == t->nb[2] * t->ne[2] && + t->nb[3] == t->nb[1] * t->ne[1]; + }; + const bool k_quant = k->type != GGML_TYPE_F16 && k->type != GGML_TYPE_BF16 && k->type != GGML_TYPE_F32; + const bool v_quant = v->type != GGML_TYPE_F16 && v->type != GGML_TYPE_BF16 && v->type != GGML_TYPE_F32; + const bool use_dequant_kv = k_quant && v_quant && neq1 >= 64 && + is_dense_kv_cache(k) && is_dense_kv_cache(v) && + (uint64_t)ggml_nelements(k) * sizeof(ggml_fp16_t) <= ctx->device->properties.limits.maxStorageBufferRange && + (uint64_t)ggml_nelements(v) * sizeof(ggml_fp16_t) <= ctx->device->properties.limits.maxStorageBufferRange && + ctx->device->pipeline_dequant_transpose[k->type] != nullptr && + ctx->device->pipeline_dequant_transpose[v->type] != nullptr && + // coopmat2 path does not benefit from the f16 scratch + !ctx->device->coopmat2 && + // Intel Xe1 regresses, see PR 25494 + (ctx->device->vendor_id != VK_VENDOR_ID_INTEL || + (ctx->device->coopmat_support && ctx->device->architecture != vk_device_architecture::INTEL_XE1)); + const ggml_type k_type_eff = use_dequant_kv ? GGML_TYPE_F16 : k->type; + const ggml_type v_type_eff = use_dequant_kv ? GGML_TYPE_F16 : v->type; + // For scalar/coopmat1 FA, we can use the "large" size to accommodate qga. // For coopmat2 FA, we always use the small size (which is still pretty large for gqa). - vk_fa_tuning_params tuning_params = get_fa_tuning_params(ctx->device, HSK, HSV, 512, KV, k->type, v->type, f32acc); + vk_fa_tuning_params tuning_params = get_fa_tuning_params(ctx->device, HSK, HSV, 512, KV, k_type_eff, v_type_eff, f32acc); const uint32_t max_gqa = std::min(tuning_params.block_rows, 32u); if (N <= 8 && qk_ratio > 1 && qk_ratio <= max_gqa && @@ -10855,7 +10896,7 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx workgroups_y /= gqa_ratio; } - tuning_params = get_fa_tuning_params(ctx->device, HSK, HSV, N, KV, k->type, v->type, f32acc); + tuning_params = get_fa_tuning_params(ctx->device, HSK, HSV, N, KV, k_type_eff, v_type_eff, f32acc); const uint32_t q_stride = (uint32_t)(nbq1 / ggml_type_size(q->type)); uint32_t k_stride = (uint32_t)(nbk1 / ggml_type_size(k->type)); @@ -10869,6 +10910,17 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx v_stride /= 4; } + uint32_t nbk2_eff = (uint32_t)nbk2, nbk3_eff = (uint32_t)nbk3; + uint32_t nbv2_eff = (uint32_t)nbv2, nbv3_eff = (uint32_t)nbv3; + if (use_dequant_kv) { + k_stride = HSK; + v_stride = HSV; + nbk2_eff = (uint32_t)((uint64_t)HSK * KV * sizeof(ggml_fp16_t)); + nbk3_eff = (uint32_t)((uint64_t)HSK * KV * nek2 * sizeof(ggml_fp16_t)); + nbv2_eff = (uint32_t)((uint64_t)HSV * KV * sizeof(ggml_fp16_t)); + nbv3_eff = (uint32_t)((uint64_t)HSV * KV * nev2 * sizeof(ggml_fp16_t)); + } + const uint32_t alignment = tuning_params.block_cols; bool aligned = (KV % alignment) == 0 && // the "aligned" shader variant will forcibly align strides, for performance @@ -10895,7 +10947,7 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx bool use_mask_opt = mask && nem1 >= 32 && nem0 * nem1 > 32768 && nem0 >= tuning_params.block_cols * 16 && (ctx->device->architecture != vk_device_architecture::AMD_GCN || HSK > 256 || HSV > 256); vk_fa_pipeline_state fa_pipeline_state = get_fa_pipeline_state(ctx->device, tuning_params, HSK, HSV, aligned, f32acc, - mask != nullptr, use_mask_opt, logit_softcap != 0, k->type, v->type); + mask != nullptr, use_mask_opt, logit_softcap != 0, k_type_eff, v_type_eff); vk_pipeline pipeline = nullptr; @@ -10999,6 +11051,34 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx vk_subbuffer sinks_buf = sinks ? ggml_vk_tensor_subbuffer(ctx, sinks) : q_buf; vk_subbuffer mask_opt_buf = use_mask_opt ? ggml_vk_subbuffer(ctx, ctx->prealloc_y, 0) : q_buf; + if (use_dequant_kv) { + const uint64_t fp = sizeof(ggml_fp16_t); + const uint64_t k_f16_sz = (uint64_t)ggml_nelements(k) * fp; + const uint64_t v_f16_sz = (uint64_t)ggml_nelements(v) * fp; + if (ctx->prealloc_size_x < k_f16_sz + v_f16_sz) { + ctx->prealloc_size_x = k_f16_sz + v_f16_sz; + ggml_vk_preallocate_buffers(ctx, subctx); + } + vk_pipeline tr_k = ctx->device->pipeline_dequant_transpose[k->type]; + vk_pipeline tr_v = ctx->device->pipeline_dequant_transpose[v->type]; + ggml_pipeline_request_descriptor_sets(ctx, tr_k, 1); + ggml_pipeline_request_descriptor_sets(ctx, tr_v, 1); + if (ctx->prealloc_x_need_sync) { + ggml_vk_sync_buffers(ctx, subctx); + } + vk_subbuffer k_dst = vk_subbuffer{ ctx->prealloc_x, 0, k_f16_sz }; + vk_subbuffer v_dst = vk_subbuffer{ ctx->prealloc_x, k_f16_sz, v_f16_sz }; + const uint32_t k_nel = (uint32_t)ggml_nelements(k); + const uint32_t v_nel = (uint32_t)ggml_nelements(v); + { const std::vector pc = { (uint32_t)HSK, (uint32_t)nek2, (uint32_t)KV, 0, k_nel }; + ggml_vk_dispatch_pipeline(ctx, subctx, tr_k, { k_buf, k_dst }, pc, { k_nel, 1, 1 }); } + { const std::vector pc = { (uint32_t)HSV, (uint32_t)nev2, (uint32_t)KV, 0, v_nel }; + ggml_vk_dispatch_pipeline(ctx, subctx, tr_v, { v_buf, v_dst }, pc, { v_nel, 1, 1 }); } + ggml_vk_sync_buffers(ctx, subctx); + k_buf = k_dst; + v_buf = v_dst; + } + uint32_t mask_n_head_log2 = ((sinks != nullptr) << 24) | n_head_log2; if (use_mask_opt) @@ -11028,8 +11108,8 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx (uint32_t)nev2, (uint32_t)nev3, nem1, nem2, nem3, q_stride, (uint32_t)nbq2, (uint32_t)nbq3, - k_stride, (uint32_t)nbk2, (uint32_t)nbk3, - v_stride, (uint32_t)nbv2, (uint32_t)nbv3, + k_stride, nbk2_eff, nbk3_eff, + v_stride, nbv2_eff, nbv3_eff, scale, max_bias, logit_softcap, mask_n_head_log2, m0, m1, gqa_ratio, split_kv, split_k }; @@ -11071,6 +11151,10 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx {q_buf, k_buf, v_buf, mask_buf, sinks_buf, dst_buf, mask_opt_buf}, pc, { workgroups_x, workgroups_y, workgroups_z }); } + + if (use_dequant_kv) { + ctx->prealloc_x_need_sync = true; + } } static vk_conv_shapes ggml_vk_conv_select_shape(ggml_backend_vk_context * ctx, uint32_t K, uint32_t NPQ) { @@ -12225,7 +12309,16 @@ static void ggml_vk_op_f32(ggml_backend_vk_context * ctx, vk_context& subctx, co elements = { ne, 1, 1 }; } - if (pipeline == ctx->device->pipeline_cpy_transpose_32 || + if (pipeline == ctx->device->pipeline_cpy_transpose_02_32 || + pipeline == ctx->device->pipeline_cpy_transpose_02_16) { + // 32x32 tiles over dims 0 and 2; dim1 and dim3 are the batch + elements[0] = (uint32_t)CEIL_DIV(dst->ne[0], 32); + elements[1] = (uint32_t)CEIL_DIV(dst->ne[2], 32); + elements[2] = (uint32_t)(dst->ne[1]*dst->ne[3]); + elements[0] = std::min(elements[0], ctx->device->properties.limits.maxComputeWorkGroupCount[0]); + elements[1] = std::min(elements[1], ctx->device->properties.limits.maxComputeWorkGroupCount[1]); + elements[2] = std::min(elements[2], ctx->device->properties.limits.maxComputeWorkGroupCount[2]); + } else if (pipeline == ctx->device->pipeline_cpy_transpose_32 || pipeline == ctx->device->pipeline_cpy_transpose_16) { // 32x32 tiles elements[0] = (uint32_t)CEIL_DIV(dst->ne[0], 32); @@ -13153,6 +13246,7 @@ static uint32_t ggml_vk_rms_partials_size(ggml_backend_vk_context * ctx, const g static vk_op_rope_push_constants ggml_vk_make_rope_constants(const ggml_tensor *dst, const ggml_tensor *src0, const bool has_ff, bool backprop, const uint32_t set_rows_stride) { const int n_dims = ((const int32_t *) dst->op_params)[1]; const int mode = ((const int32_t *) dst->op_params)[2]; + const int n_offs = ((const int32_t *) dst->op_params)[15]; // const int n_ctx = ((const int32_t *) dst->op_params)[3]; const int n_ctx_orig = ((const int32_t *) dst->op_params)[4]; const float freq_base = ((const float *) dst->op_params)[5]; @@ -13182,7 +13276,7 @@ static vk_op_rope_push_constants ggml_vk_make_rope_constants(const ggml_tensor * uint32_t nb13 = dst->nb[3] / ggml_type_size(dst->type); vk_op_rope_push_constants rope { - (uint32_t)mode, (uint32_t)ggml_nrows(src0), (uint32_t)n_dims, freq_scale, + (uint32_t)mode, (uint32_t)ggml_nrows(src0), (uint32_t)n_dims, (uint32_t)n_offs, freq_scale, freq_base, ext_factor, attn_factor, {corr_dims[0], corr_dims[1]}, theta_scale, has_ff, { sections[0], sections[1], sections[2], sections[3] }, is_imrope, backprop, set_rows_stride, @@ -19228,6 +19322,10 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph * tensor_clone = ggml_rope_ext_back(ggml_ctx, src_clone[0], src_clone[1], src_clone[2], n_dims, mode, n_ctx_orig_ggml, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); } } + const int n_offs = ((int32_t *) tensor->op_params)[15]; + if (n_offs != 0) { + tensor_clone = ggml_rope_set_offset(tensor_clone, n_offs); + } } else if (tensor->op == GGML_OP_UNARY) { switch (ggml_get_unary_op(tensor)) { case GGML_UNARY_OP_EXP: diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/copy_transpose_02.comp b/ggml/src/ggml-vulkan/vulkan-shaders/copy_transpose_02.comp new file mode 100644 index 000000000..5a3d66dab --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/copy_transpose_02.comp @@ -0,0 +1,61 @@ +#version 450 + +#include "types.glsl" +#include "generic_unary_head.glsl" + +// workgroup does 32x32 tile, but uses 32x8 threads +#define TILE_DIM 32 +layout(local_size_x = 32, local_size_y = 8, local_size_z = 1) in; + +// +1 padding avoids shared-memory bank conflicts on the transposed read +shared uint sh[TILE_DIM][TILE_DIM + 1]; + +void iter(uvec3 wg_id) { + const uint tile_i0 = wg_id.x; // tiles dst ne10 (== src ne00) + const uint tile_i2 = wg_id.y; // tiles dst ne12 (== src ne02) + + const uint tid_col = gl_LocalInvocationID.x; + const uint tid_row = gl_LocalInvocationID.y; + + const uint i1 = wg_id.z % p.ne11; + const uint i3 = wg_id.z / p.ne11; + const uint i01 = i1; + const uint i03 = i3; + + [[unroll]] for (uint y = 0; y < 4; ++y) { + const uint i00 = tile_i0 * TILE_DIM + tid_row + 8 * y; + const uint i02 = tile_i2 * TILE_DIM + tid_col; + if (i00 < p.ne00 && i01 < p.ne01 && i02 < p.ne02 && i03 < p.ne03) { + const uint src_idx = i00 * p.nb00 + i01 * p.nb01 + i02 * p.nb02 + i03 * p.nb03; + sh[tid_row + 8 * y][tid_col] = uint(data_a[get_aoffset() + src_idx]); + } + } + + barrier(); + + [[unroll]] for (uint y = 0; y < 4; ++y) { + const uint i0 = tile_i0 * TILE_DIM + tid_col; + const uint i2 = tile_i2 * TILE_DIM + tid_row + 8 * y; + if (i0 < p.ne10 && i1 < p.ne11 && i2 < p.ne12 && i3 < p.ne13) { + const uint dst_idx = i0 * p.nb10 + i1 * p.nb11 + i2 * p.nb12 + i3 * p.nb13; + data_d[get_doffset() + dst_idx] = D_TYPE(sh[tid_col][tid_row + 8 * y]); + } + } +} + +#define CEIL_DIV(a, b) (((a) + (b) - 1) / (b)) + +void main() { + bool need_barrier = false; + for (uint z = gl_WorkGroupID.z; z < p.ne11 * p.ne13; z += gl_NumWorkGroups.z) { + for (uint y = gl_WorkGroupID.y; y < CEIL_DIV(p.ne12, TILE_DIM); y += gl_NumWorkGroups.y) { + for (uint x = gl_WorkGroupID.x; x < CEIL_DIV(p.ne10, TILE_DIM); x += gl_NumWorkGroups.x) { + if (need_barrier) { + barrier(); + } + need_barrier = true; + iter(uvec3(x, y, z)); + } + } + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q8_0.comp b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q8_0.comp index 10844ddf7..3b3fbbe89 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q8_0.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q8_0.comp @@ -18,7 +18,18 @@ void main() { return; } +#ifdef DEQUANT_TRANSPOSE + // read [HS, NH, KV, NS], write [HS, KV, NH, NS] + const uint HS = p.M, NH = p.K, KVn = p.stride_a; + const uint e0 = ib * 32; + const uint b_idx = (e0 % HS) + + ((e0 / (HS * NH)) % KVn) * HS + + ((e0 / HS) % NH) * (HS * KVn) + + (e0 / (HS * NH * KVn)) * (HS * KVn * NH) + + 16 * il; +#else const uint b_idx = 1024*i + 32*ir + 16*il; +#endif const float d = float(data_a[ib].d); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/rope_funcs.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/rope_funcs.glsl index 033587931..feb55b203 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/rope_funcs.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/rope_funcs.glsl @@ -50,19 +50,21 @@ void rope_norm(const uint i0, const uint i1, const uint i2, const uint i3, rope_ } idst += p.d_offset; - if (i0 >= p.n_dims) { + if (i0 < p.n_offs || i0 >= p.n_offs + p.n_dims) { rope_data_d[idst + 0] = ROPE_D_TYPE(rope_data_a[ix + 0]); rope_data_d[idst + 1] = ROPE_D_TYPE(rope_data_a[ix + 1]); return; } - const float theta_base = rope_data_pos[i2] * pow(p.theta_scale, i0/2.0f); + const uint iw = i0 - p.n_offs; // relative idx - const float freq_factor = p.has_ff != 0 ? rope_data_ff[i0/2] : 1.0f; + const float theta_base = rope_data_pos[i2] * pow(p.theta_scale, iw/2.0f); + + const float freq_factor = p.has_ff != 0 ? rope_data_ff[iw/2] : 1.0f; float cos_theta, sin_theta; - rope_yarn(theta_base / freq_factor, i0, cos_theta, sin_theta, p); + rope_yarn(theta_base / freq_factor, iw, cos_theta, sin_theta, p); const float x0 = float(rope_data_a[ix + 0]); const float x1 = float(rope_data_a[ix + 1]); @@ -87,25 +89,28 @@ void rope_neox(const uint i0, const uint i1, const uint i2, const uint i3, rope_ } idst += p.d_offset; - if (i0 >= p.n_dims) { + if (i0 < p.n_offs || i0 >= p.n_offs + p.n_dims) { rope_data_d[idst + i0/2 + 0] = ROPE_D_TYPE(rope_data_a[ix + i0/2 + 0]); rope_data_d[idst + i0/2 + 1] = ROPE_D_TYPE(rope_data_a[ix + i0/2 + 1]); return; } - const float theta_base = rope_data_pos[i2] * pow(p.theta_scale, i0/2.0f); + const uint iw = i0 - p.n_offs; // relative idx - const float freq_factor = p.has_ff != 0 ? rope_data_ff[i0/2] : 1.0f; + const float theta_base = rope_data_pos[i2] * pow(p.theta_scale, iw/2.0f); + + const float freq_factor = p.has_ff != 0 ? rope_data_ff[iw/2] : 1.0f; float cos_theta, sin_theta; - rope_yarn(theta_base / freq_factor, i0, cos_theta, sin_theta, p); + rope_yarn(theta_base / freq_factor, iw, cos_theta, sin_theta, p); - const float x0 = float(rope_data_a[ix + 0]); - const float x1 = float(rope_data_a[ix + p.n_dims/2]); + // idst/ix point at channel i0/2; the first channel of the rotated pair is p.n_offs + iw/2 = i0/2 + p.n_offs/2 + const float x0 = float(rope_data_a[ix + p.n_offs/2 + 0]); + const float x1 = float(rope_data_a[ix + p.n_offs/2 + p.n_dims/2]); - rope_data_d[idst + 0] = ROPE_D_TYPE(x0*cos_theta - x1*sin_theta); - rope_data_d[idst + p.n_dims/2] = ROPE_D_TYPE(x0*sin_theta + x1*cos_theta); + rope_data_d[idst + p.n_offs/2 + 0] = ROPE_D_TYPE(x0*cos_theta - x1*sin_theta); + rope_data_d[idst + p.n_offs/2 + p.n_dims/2] = ROPE_D_TYPE(x0*sin_theta + x1*cos_theta); } @@ -125,53 +130,56 @@ void rope_multi(const uint i0, const uint i1, const uint i2, const uint i3, rope } idst += p.d_offset; - if (i0 >= p.n_dims) { + if (i0 < p.n_offs || i0 >= p.n_offs + p.n_dims) { rope_data_d[idst + i0/2 + 0] = ROPE_D_TYPE(rope_data_a[ix + i0/2 + 0]); rope_data_d[idst + i0/2 + 1] = ROPE_D_TYPE(rope_data_a[ix + i0/2 + 1]); return; } + const uint iw = i0 - p.n_offs; // relative idx + const int sect_dims = p.sections[0] + p.sections[1] + p.sections[2] + p.sections[3]; const int sec_w = p.sections[1] + p.sections[0]; - const uint sector = (i0 / 2) % sect_dims; + const uint sector = (iw / 2) % sect_dims; float theta_base = 0.0; if (p.is_imrope != 0) { if (sector % 3 == 1 && sector < 3 * p.sections[1]) { - theta_base = rope_data_pos[i2 + p.ne02 * 1]*pow(p.theta_scale, i0/2.0f); + theta_base = rope_data_pos[i2 + p.ne02 * 1]*pow(p.theta_scale, iw/2.0f); } else if (sector % 3 == 2 && sector < 3 * p.sections[2]) { - theta_base = rope_data_pos[i2 + p.ne02 * 2]*pow(p.theta_scale, i0/2.0f); + theta_base = rope_data_pos[i2 + p.ne02 * 2]*pow(p.theta_scale, iw/2.0f); } else if (sector % 3 == 0 && sector < 3 * p.sections[0]) { - theta_base = rope_data_pos[i2]*pow(p.theta_scale, i0/2.0f); + theta_base = rope_data_pos[i2]*pow(p.theta_scale, iw/2.0f); } else { - theta_base = rope_data_pos[i2 + p.ne02 * 3]*pow(p.theta_scale, i0/2.0f); + theta_base = rope_data_pos[i2 + p.ne02 * 3]*pow(p.theta_scale, iw/2.0f); } } else { if (sector < p.sections[0]) { - theta_base = rope_data_pos[i2]*pow(p.theta_scale, i0/2.0f); + theta_base = rope_data_pos[i2]*pow(p.theta_scale, iw/2.0f); } else if (sector >= p.sections[0] && sector < sec_w) { - theta_base = rope_data_pos[i2 + p.ne02 * 1]*pow(p.theta_scale, i0/2.0f); + theta_base = rope_data_pos[i2 + p.ne02 * 1]*pow(p.theta_scale, iw/2.0f); } else if (sector >= sec_w && sector < sec_w + p.sections[2]) { - theta_base = rope_data_pos[i2 + p.ne02 * 2]*pow(p.theta_scale, i0/2.0f); + theta_base = rope_data_pos[i2 + p.ne02 * 2]*pow(p.theta_scale, iw/2.0f); } else if (sector >= sec_w + p.sections[2]) { - theta_base = rope_data_pos[i2 + p.ne02 * 3]*pow(p.theta_scale, i0/2.0f); + theta_base = rope_data_pos[i2 + p.ne02 * 3]*pow(p.theta_scale, iw/2.0f); } } - const float freq_factor = p.has_ff != 0 ? rope_data_ff[i0/2] : 1.0f; + const float freq_factor = p.has_ff != 0 ? rope_data_ff[iw/2] : 1.0f; float cos_theta, sin_theta; - rope_yarn(theta_base / freq_factor, i0, cos_theta, sin_theta, p); + rope_yarn(theta_base / freq_factor, iw, cos_theta, sin_theta, p); - const float x0 = float(rope_data_a[ix + 0]); - const float x1 = float(rope_data_a[ix + p.n_dims/2]); + // idst/ix point at channel i0/2; the first channel of the rotated pair is p.n_offs + iw/2 = i0/2 + p.n_offs/2 + const float x0 = float(rope_data_a[ix + p.n_offs/2 + 0]); + const float x1 = float(rope_data_a[ix + p.n_offs/2 + p.n_dims/2]); - rope_data_d[idst + 0] = ROPE_D_TYPE(x0*cos_theta - x1*sin_theta); - rope_data_d[idst + p.n_dims/2] = ROPE_D_TYPE(x0*sin_theta + x1*cos_theta); + rope_data_d[idst + p.n_offs/2 + 0] = ROPE_D_TYPE(x0*cos_theta - x1*sin_theta); + rope_data_d[idst + p.n_offs/2 + p.n_dims/2] = ROPE_D_TYPE(x0*sin_theta + x1*cos_theta); } void rope_vision(const uint i0, const uint i1, const uint i2, const uint i3, rope_params p) { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/rope_params.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/rope_params.glsl index 3602485b9..b88a73fcc 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/rope_params.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/rope_params.glsl @@ -5,6 +5,7 @@ struct rope_params { uint rope_mode; uint nrows; uint n_dims; + uint n_offs; float freq_scale; float freq_base; float ext_factor; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp index 591ab6602..b40a9d291 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -806,6 +806,10 @@ void process_shaders() { if (tname != "f16" && tname != "bf16") { string_to_spv("dequant_" + tname, "dequant_" + tname + ".comp", merge_maps(base_dict, {{data_a_key, "1"}, {"D_TYPE", "float16_t"}})); } + // Fused dequant+transpose variant for FA quant-KV (per-head-contiguous f16 scratch). + if (tname == "q8_0") { + string_to_spv("dequant_" + tname + "_transpose", "dequant_" + tname + ".comp", merge_maps(base_dict, {{data_a_key, "1"}, {"D_TYPE", "float16_t"}, {"DEQUANT_TRANSPOSE", "1"}})); + } shader = (tname == "f32" || tname == "f16" || tname == "bf16") ? "get_rows.comp" : "get_rows_quant.comp"; @@ -852,6 +856,8 @@ void process_shaders() { string_to_spv("cpy_transpose_16", "copy_transpose.comp", {{"A_TYPE", "uint16_t"}, {"D_TYPE", "uint16_t"}}); string_to_spv("cpy_transpose_32", "copy_transpose.comp", {{"A_TYPE", "uint"}, {"D_TYPE", "uint"}}); + string_to_spv("cpy_transpose_02_16", "copy_transpose_02.comp", {{"A_TYPE", "uint16_t"}, {"D_TYPE", "uint16_t"}}); + string_to_spv("cpy_transpose_02_32", "copy_transpose_02.comp", {{"A_TYPE", "uint"}, {"D_TYPE", "uint"}}); for (std::string t : {"q1_0", "q2_0", "q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "iq4_nl"}) { string_to_spv("cpy_f32_" + t, "copy_to_quant.comp", {{"DATA_A_" + to_uppercase(t), "1"}, {"S_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 8cf8ab342..b43058405 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -4216,7 +4216,7 @@ static struct ggml_tensor * ggml_rope_impl( struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - int32_t params[15] = { /*n_past*/ 0, n_dims, mode, /*n_ctx*/ 0, n_ctx_orig }; + int32_t params[16] = { /*n_past*/ 0, n_dims, mode, /*n_ctx*/ 0, n_ctx_orig }; memcpy(params + 5, &freq_base, sizeof(float)); memcpy(params + 6, &freq_scale, sizeof(float)); memcpy(params + 7, &ext_factor, sizeof(float)); @@ -4228,6 +4228,8 @@ static struct ggml_tensor * ggml_rope_impl( } else { memset(params + 11, 0, sizeof(int32_t) * GGML_MROPE_SECTIONS); } + params[15] = 0; // n_offs, set via ggml_rope_set_offset() + ggml_set_op_params(result, params, sizeof(params)); result->op = GGML_OP_ROPE; @@ -4438,6 +4440,20 @@ struct ggml_tensor * ggml_rope_multi_back( result->op = GGML_OP_ROPE_BACK; return result; } + +struct ggml_tensor * ggml_rope_set_offset( + struct ggml_tensor * a, + int n_offs) { + GGML_ASSERT(a->op == GGML_OP_ROPE || a->op == GGML_OP_ROPE_BACK); + GGML_ASSERT(n_offs >= 0); + + const int32_t mode = ggml_get_op_params_i32(a, 2); + GGML_ASSERT(mode != GGML_ROPE_TYPE_VISION); + + ggml_set_op_params_i32(a, 15, n_offs); + return a; +} + // ggml_clamp struct ggml_tensor * ggml_clamp( diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index d043c9b6e..fad8d1fd8 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -208,6 +208,7 @@ class Keys: SHARED_KV_LAYERS = "{arch}.attention.shared_kv_layers" SLIDING_WINDOW_PATTERN = "{arch}.attention.sliding_window_pattern" TEMPERATURE_SCALE = "{arch}.attention.temperature_scale" + ROPE_PATTERN = "{arch}.attention.rope_pattern" class Indexer: HEAD_COUNT = "{arch}.attention.indexer.head_count" @@ -549,6 +550,7 @@ class MODEL_ARCH(IntEnum): GRANITE_MOE = auto() GRANITE_HYBRID = auto() GRANITE_SWITCH = auto() + GRANITE_SWA = auto() CHAMELEON = auto() WAVTOKENIZER_DEC = auto() PLM = auto() @@ -1265,6 +1267,7 @@ MODEL_ARCH_NAMES: dict[MODEL_ARCH, str] = { MODEL_ARCH.GRANITE_MOE: "granitemoe", MODEL_ARCH.GRANITE_HYBRID: "granitehybrid", MODEL_ARCH.GRANITE_SWITCH: "graniteswitch", + MODEL_ARCH.GRANITE_SWA: "granite_swa", MODEL_ARCH.CHAMELEON: "chameleon", MODEL_ARCH.WAVTOKENIZER_DEC: "wavtokenizer-dec", MODEL_ARCH.PLM: "plm", @@ -4152,6 +4155,31 @@ MODEL_TENSORS: dict[MODEL_ARCH, list[MODEL_TENSOR]] = { MODEL_TENSOR.FFN_DOWN, MODEL_TENSOR.FFN_UP, ], + MODEL_ARCH.GRANITE_SWA: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT_NORM, + MODEL_TENSOR.OUTPUT, + MODEL_TENSOR.ATTN_NORM, + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.ATTN_SINKS, + MODEL_TENSOR.ROPE_FREQS, + MODEL_TENSOR.FFN_NORM, + MODEL_TENSOR.FFN_GATE, + MODEL_TENSOR.FFN_DOWN, + MODEL_TENSOR.FFN_UP, + # MoE (GraniteMoeSWA) + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_GATE_UP_EXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + # Shared expert - gate+up kept fused in FFN_UP_SHEXP (LLM_FFN_SWIGLU) + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + ], MODEL_ARCH.CHAMELEON: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, diff --git a/gguf-py/gguf/gguf_reader.py b/gguf-py/gguf/gguf_reader.py index ea241ada2..bf3e08380 100644 --- a/gguf-py/gguf/gguf_reader.py +++ b/gguf-py/gguf/gguf_reader.py @@ -32,6 +32,10 @@ from gguf.constants import ( GGUFEndian, ) +# limits mirroring ggml/src/gguf.cpp (not part of gguf.h) +GGUF_MAX_STRING_LENGTH = 1024 * 1024 * 1024 +GGUF_MAX_ARRAY_ELEMENTS = 1024 * 1024 * 1024 + logger = logging.getLogger(__name__) READER_SUPPORTED_VERSIONS = [2, GGUF_VERSION] @@ -167,6 +171,10 @@ class GGUFReader: offs += self._push_field(ReaderField(offs, 'GGUF.tensor_count', [temp_counts[:1]], [0], [GGUFValueType.UINT64])) offs += self._push_field(ReaderField(offs, 'GGUF.kv_count', [temp_counts[1:]], [0], [GGUFValueType.UINT64])) tensor_count, kv_count = temp_counts + if tensor_count > GGUF_MAX_ARRAY_ELEMENTS: + raise ValueError(f'Tensor count {tensor_count} exceeds maximum {GGUF_MAX_ARRAY_ELEMENTS}') + if kv_count > GGUF_MAX_ARRAY_ELEMENTS: + raise ValueError(f'KV count {kv_count} exceeds maximum {GGUF_MAX_ARRAY_ELEMENTS}') offs = self._build_fields(offs, kv_count) # Build Tensor Info Fields @@ -217,6 +225,10 @@ class GGUFReader: def _get_str(self, offset: int) -> tuple[npt.NDArray[np.uint64], npt.NDArray[np.uint8]]: slen = self._get(offset, np.uint64) + if int(slen[0]) > GGUF_MAX_STRING_LENGTH: + raise ValueError(f'String length {int(slen[0])} exceeds maximum {GGUF_MAX_STRING_LENGTH}') + if offset + 8 + int(slen[0]) > self.data.nbytes: + raise ValueError(f'String length {int(slen[0])} exceeds remaining file size {self.data.nbytes - offset - 8}') return slen, self._get(offset + 8, np.uint8, slen[0]) def _get_field_parts( @@ -241,6 +253,8 @@ class GGUFReader: raw_itype = self._get(offs, np.uint32) offs += int(raw_itype.nbytes) alen = self._get(offs, np.uint64) + if int(alen[0]) > GGUF_MAX_ARRAY_ELEMENTS: + raise ValueError(f'Array length {int(alen[0])} exceeds maximum {GGUF_MAX_ARRAY_ELEMENTS}') offs += int(alen.nbytes) aparts: list[npt.NDArray[Any]] = [raw_itype, alen] data_idxs: list[int] = [] diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index 9e0914fd8..16ae9f999 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -824,6 +824,9 @@ class GGUFWriter: else: self.add_array(key, value) + def add_rope_pattern(self, value: Sequence[bool]) -> None: + self.add_array(Keys.Attention.ROPE_PATTERN.format(arch=self.arch), value) + def add_dense_features_dims(self, dense:str, in_f:int, out_f:int) -> None: self.add_uint32(Keys.LLM.DENSE_FEAT_IN_SIZE.format(arch=self.arch, dense=dense), in_f) self.add_uint32(Keys.LLM.DENSE_FEAT_OUT_SIZE.format(arch=self.arch, dense=dense), out_f) diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index 3292942b4..a0571ccd3 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -458,6 +458,7 @@ class TensorNameMap: "transformer.decoder_layer.{bid}.router", # Grok "transformer.blocks.{bid}.ffn.router.layer", # dbrx "model.layers.{bid}.block_sparse_moe.router.layer", # granitemoe + "model.layers.{bid}.block_sparse_moe.router", # granite_swa "model.layers.{bid}.feed_forward.router", # llama4 jamba "encoder.layers.{bid}.mlp.router.layer", # nomic-bert-moe "model.layers.{bid}.mlp.router", # openai-moe diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 5b88bde14..955c2d796 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -102,6 +102,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_GRANITE_MOE, "granitemoe" }, { LLM_ARCH_GRANITE_HYBRID, "granitehybrid" }, { LLM_ARCH_GRANITE_SWITCH, "graniteswitch" }, + { LLM_ARCH_GRANITE_SWA, "granite_swa" }, { LLM_ARCH_CHAMELEON, "chameleon" }, { LLM_ARCH_WAVTOKENIZER_DEC, "wavtokenizer-dec" }, { LLM_ARCH_PLM, "plm" }, @@ -261,6 +262,8 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_ATTENTION_RELATIVE_BUCKETS_COUNT, "%s.attention.relative_buckets_count" }, { LLM_KV_ATTENTION_SLIDING_WINDOW, "%s.attention.sliding_window" }, { LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, "%s.attention.sliding_window_pattern" }, + { LLM_KV_ATTENTION_ROPE_PATTERN, "%s.attention.rope_pattern" }, + { LLM_KV_ATTENTION_SCALE, "%s.attention.scale" }, { LLM_KV_ATTENTION_OUTPUT_SCALE, "%s.attention.output_scale" }, { LLM_KV_ATTENTION_VALUE_SCALE, "%s.attention.value_scale" }, diff --git a/src/llama-arch.h b/src/llama-arch.h index 8042120a2..48fe051a9 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -107,6 +107,7 @@ enum llm_arch { LLM_ARCH_GRANITE_MOE, LLM_ARCH_GRANITE_HYBRID, LLM_ARCH_GRANITE_SWITCH, + LLM_ARCH_GRANITE_SWA, LLM_ARCH_CHAMELEON, LLM_ARCH_WAVTOKENIZER_DEC, LLM_ARCH_PLM, @@ -267,6 +268,8 @@ enum llm_kv { LLM_KV_ATTENTION_SLIDING_WINDOW, LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, LLM_KV_ATTENTION_SCALE, + LLM_KV_ATTENTION_ROPE_PATTERN, + LLM_KV_ATTENTION_OUTPUT_SCALE, LLM_KV_ATTENTION_VALUE_SCALE, LLM_KV_ATTENTION_TEMPERATURE_LENGTH, diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index e3f0cf0ed..cbe31134f 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -291,7 +291,11 @@ bool llama_hparams::has_rope(uint32_t il) const { return false; } - return true; + if (il < n_layer_all) { + return rope_pattern[il] != 0; + } + + GGML_ABORT("%s: il (%u) out of bounds (n_layer_all: %u)\n", __func__, il, n_layer_all); } uint32_t llama_hparams::n_layer() const { diff --git a/src/llama-hparams.h b/src/llama-hparams.h index e91ce1cc3..f6af36436 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -144,6 +144,10 @@ struct llama_hparams { std::array rope_sections; + // Per-layer RoPE enable flags (1 = use RoPE, 0 = NoPE) + // by default, all layers use RoPE (controlled by rope_finetuned) + std::array rope_pattern; + // Sliding Window Attention (SWA) llama_swa_type swa_type = LLAMA_SWA_TYPE_NONE; // the size of the sliding window (0 - no SWA) diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index 6481b62e4..7d91a297b 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -1396,6 +1396,11 @@ void llama_model_loader::get_mapping_range(size_t * first, size_t * last, void * } } +void llama_model_loader::unmap_weight(const llama_tensor_weight & w) const { + if (!use_mmap) { return; } + mappings.at(w.idx)->unmap_fragment(w.offs, w.offs + ggml_nbytes(w.tensor)); +} + void llama_model_loader::load_data_for(struct ggml_tensor * cur) const { const auto & w = require_weight(ggml_get_name(cur)); diff --git a/src/llama-model-loader.h b/src/llama-model-loader.h index d6b31c231..e9fe3592d 100644 --- a/src/llama-model-loader.h +++ b/src/llama-model-loader.h @@ -194,6 +194,9 @@ struct llama_model_loader { void get_mapping_range(size_t * first, size_t * last, void ** addr, int idx, ggml_context * ctx) const; + // release a weight's mmap pages + void unmap_weight(const llama_tensor_weight & w) const; + // for backwards compatibility, does not support ggml-backend void load_data_for(struct ggml_tensor * cur) const; diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index be9524d40..b9e0a6009 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -30,6 +30,7 @@ bool llama_model_saver_supports_arch(llm_arch arch) { case LLM_ARCH_MUSE_GLIMMER: case LLM_ARCH_MELLUM: case LLM_ARCH_LAGUNA: + case LLM_ARCH_GRANITE_SWA: return false; default: return true; @@ -272,6 +273,7 @@ void llama_model_saver::add_kv_from_model() { add_kv(LLM_KV_ATTENTION_VALUE_RESIDUAL_MIX_LORA_RANK, hparams.n_lora_value_res_mix); add_kv(LLM_KV_ATTENTION_GATE_LORA_RANK, hparams.n_lora_gate); add_kv(LLM_KV_ATTENTION_RELATIVE_BUCKETS_COUNT, hparams.n_rel_attn_bkts); + add_kv(LLM_KV_ATTENTION_ROPE_PATTERN, hparams.rope_pattern, true); add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa); // add_kv(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, ???); add_kv(LLM_KV_ATTENTION_SCALE, hparams.f_attention_scale); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index e35a683d1..e38c89d03 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -92,6 +92,7 @@ #include "models/gptneox.cpp" #include "models/granite-hybrid.cpp" #include "models/granite-moe.cpp" +#include "models/granite-swa.cpp" #include "models/granite-switch.cpp" #include "models/granite.cpp" #include "models/grok.cpp" @@ -395,6 +396,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_minicpm(params); case LLM_ARCH_GRANITE_HYBRID: return new llama_model_granite_hybrid(params); + case LLM_ARCH_GRANITE_SWA: + return new llama_model_granite_swa(params); case LLM_ARCH_CHAMELEON: return new llama_model_chameleon(params); case LLM_ARCH_WAVTOKENIZER_DEC: @@ -1306,6 +1309,7 @@ void llama_model_base::load_hparams(llama_model_loader & ml) { std::fill(hparams.n_ff_arr.begin(), hparams.n_ff_arr.end(), 0); std::fill(hparams.rope_sections.begin(), hparams.rope_sections.end(), 0); + std::fill(hparams.rope_pattern.begin(), hparams.rope_pattern.end(), 1); std::fill(hparams.is_swa_impl.begin(), hparams.is_swa_impl.end(), 0); std::fill(hparams.is_recr_impl.begin(), hparams.is_recr_impl.end(), llm_arch_is_recurrent(ml.get_arch()) ? 1 : 0); std::fill(hparams.is_indexer_full_impl.begin(), hparams.is_indexer_full_impl.end(), 0); @@ -2788,6 +2792,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_GRANITE_MOE: case LLM_ARCH_GRANITE_HYBRID: case LLM_ARCH_GRANITE_SWITCH: + case LLM_ARCH_GRANITE_SWA: case LLM_ARCH_CHAMELEON: case LLM_ARCH_BAILINGMOE: case LLM_ARCH_BAILINGMOE3: diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp index 20ba3827f..982f2e426 100644 --- a/src/llama-quant.cpp +++ b/src/llama-quant.cpp @@ -1272,7 +1272,7 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: total_size_org += tensor_size; total_size_new += new_size; - // update the gguf meta data as we go + // update the gguf metadata as we go gguf_set_tensor_type(ctx_outs[cur_split].get(), metadata[i].name.c_str(), new_type); GGML_ASSERT(gguf_get_tensor_size(ctx_outs[cur_split].get(), gguf_find_tensor(ctx_outs[cur_split].get(), metadata[i].name.c_str())) == new_size); gguf_set_tensor_data(ctx_outs[cur_split].get(), metadata[i].name.c_str(), new_data); @@ -1280,6 +1280,10 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: // write tensor data + padding fout.write((const char *) new_data, new_size); zeros(fout, GGML_PAD(new_size, align) - new_size); + + // unmap the tensor to free memory + if (ml.use_mmap) { ml.unmap_weight(weight); } + } // no --dry-run } // main loop diff --git a/src/models/deepseek32.cpp b/src/models/deepseek32.cpp index 08555a801..2b82a780c 100644 --- a/src/models/deepseek32.cpp +++ b/src/models/deepseek32.cpp @@ -10,8 +10,6 @@ void llama_model_deepseek32::load_arch_hparams(llama_model_loader & ml) { ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, false); // MoE parameters - ml.get_key(LLM_KV_EXPERT_COUNT, hparams.n_expert); - ml.get_key(LLM_KV_EXPERT_USED_COUNT, hparams.n_expert_used); ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared); ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead, false); ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false); diff --git a/src/models/glm-dsa.cpp b/src/models/glm-dsa.cpp index 803ef7674..93a1448b4 100644 --- a/src/models/glm-dsa.cpp +++ b/src/models/glm-dsa.cpp @@ -32,8 +32,6 @@ void llama_model_glm_dsa::load_arch_hparams(llama_model_loader & ml) { ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, false); // MoE parameters - ml.get_key(LLM_KV_EXPERT_COUNT, hparams.n_expert); - ml.get_key(LLM_KV_EXPERT_USED_COUNT, hparams.n_expert_used); ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared); ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead, false); ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false); diff --git a/src/models/glm4-moe.cpp b/src/models/glm4-moe.cpp index d60e47ddf..8cde66978 100644 --- a/src/models/glm4-moe.cpp +++ b/src/models/glm4-moe.cpp @@ -6,8 +6,6 @@ void llama_model_glm4_moe::load_arch_hparams(llama_model_loader & ml) { ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, false); // MoE parameters - ml.get_key(LLM_KV_EXPERT_COUNT, hparams.n_expert); - ml.get_key(LLM_KV_EXPERT_USED_COUNT, hparams.n_expert_used); ml.get_key(LLM_KV_EXPERT_SHARED_COUNT, hparams.n_expert_shared); ml.get_key(LLM_KV_LEADING_DENSE_BLOCK_COUNT, hparams.n_layer_dense_lead, false); ml.get_key(LLM_KV_EXPERT_WEIGHTS_SCALE, hparams.expert_weights_scale, false); diff --git a/src/models/granite-hybrid.cpp b/src/models/granite-hybrid.cpp index eb23095ae..8a8f7e19f 100644 --- a/src/models/granite-hybrid.cpp +++ b/src/models/granite-hybrid.cpp @@ -16,7 +16,8 @@ void llama_model_granite_hybrid::load_arch_hparams(llama_model_loader & ml) { // Granite uses rope_finetuned as a switch for rope, so default to true bool rope_finetuned = true; ml.get_key(LLM_KV_ROPE_SCALING_FINETUNED, rope_finetuned, false); - hparams.rope_finetuned = rope_finetuned; + hparams.rope_finetuned = rope_finetuned; // needed for round trip save + std::fill(hparams.rope_pattern.begin(), hparams.rope_pattern.end(), rope_finetuned); // A layer is recurrent IFF the n_head_kv value is set to 0 for (uint32_t i = 0; i < hparams.n_layer(); ++i) { @@ -147,7 +148,7 @@ llama_model_granite_hybrid::graph::graph(const llama_model & model, const llm_gr // Positional embeddings populated if rope enabled ggml_tensor * inp_pos = nullptr; - if (hparams.rope_finetuned) { + if (hparams.has_rope(0)) { inp_pos = build_inp_pos(); } @@ -206,8 +207,7 @@ ggml_tensor * llama_model_granite_hybrid::graph::build_attention_layer(ggml_tens const int il) { auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur, n_embd_head, hparams.n_head(il), hparams.n_head_kv(il), il); - const bool use_rope = hparams.rope_finetuned; - if (use_rope) { + if (hparams.has_rope(il)) { ggml_tensor * rope_factors = model.get_rope_factors(cparams, il); Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, rope_factors, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); diff --git a/src/models/granite-moe.cpp b/src/models/granite-moe.cpp index 115263c41..09be49393 100644 --- a/src/models/granite-moe.cpp +++ b/src/models/granite-moe.cpp @@ -7,11 +7,6 @@ void llama_model_granite_moe::load_arch_hparams(llama_model_loader & ml) { ml.get_key(LLM_KV_EMBEDDING_SCALE, hparams.f_embedding_scale, false); ml.get_key(LLM_KV_ATTENTION_SCALE, hparams.f_attention_scale, false); - // Granite uses rope_finetuned as a switch for rope, so default to true - bool rope_finetuned = true; - ml.get_key(LLM_KV_ROPE_SCALING_FINETUNED, rope_finetuned, false); - hparams.rope_finetuned = rope_finetuned; - switch (hparams.n_layer()) { case 32: type = LLM_TYPE_3B; break; case 40: type = LLM_TYPE_3B; break; diff --git a/src/models/granite-swa.cpp b/src/models/granite-swa.cpp new file mode 100644 index 000000000..3aa2b63b2 --- /dev/null +++ b/src/models/granite-swa.cpp @@ -0,0 +1,319 @@ +#include "models.h" + +#include + +void llama_model_granite_swa::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + ml.get_key(LLM_KV_LOGIT_SCALE, hparams.f_logit_scale); + ml.get_key(LLM_KV_RESIDUAL_SCALE, hparams.f_residual_scale, false); + ml.get_key(LLM_KV_EMBEDDING_SCALE, hparams.f_embedding_scale, false); + ml.get_key(LLM_KV_ATTENTION_SCALE, hparams.f_attention_scale, false); + + // MoE expert configuration + ml.get_key(LLM_KV_EXPERT_COUNT, hparams.n_expert, false); + ml.get_key(LLM_KV_EXPERT_USED_COUNT, hparams.n_expert_used, false); + + // iSWA configuration + ml.get_arr(LLM_KV_ATTENTION_SLIDING_WINDOW_PATTERN, hparams.is_swa_impl); + ml.get_key(LLM_KV_ATTENTION_SLIDING_WINDOW, hparams.n_swa); + hparams.swa_type = LLAMA_SWA_TYPE_STANDARD; + + // Granite4 Vision uses array deepstack_mapping + ml.get_arr(LLM_KV_DEEPSTACK_MAPPING, hparams.deepstack_mapping_arr, false); + + // Count the unique deepstack input indices + std::unordered_set unique_deepstack_idxs; + for (const auto val : hparams.deepstack_mapping_arr) { + if (val >= 0) { + unique_deepstack_idxs.insert(val); + } + } + hparams.n_deepstack_layers = unique_deepstack_idxs.size(); + + // Ensure all values are valid (avoid overflow attacks) + for (const auto val : unique_deepstack_idxs) { + if (val > hparams.n_deepstack_layers) { + std::stringstream ss; + ss << "Invalid deepstack index: " << val << " > " << hparams.n_deepstack_layers; + throw std::runtime_error(ss.str()); + } + } + + // Per-layer RoPE pattern (optional) + ml.get_arr(LLM_KV_ATTENTION_ROPE_PATTERN, hparams.rope_pattern, false); + + switch (hparams.n_layer()) { + case 32: type = LLM_TYPE_3B; break; + case 40: type = LLM_TYPE_3B; break; + // Add additional layer/vocab/etc checks here for other model sizes + default: type = LLM_TYPE_UNKNOWN; + } + + // For Granite MoE Shared + ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, /* required */ false); +} + +void llama_model_granite_swa::load_arch_tensors(llama_model_loader &) { + LLAMA_LOAD_LOCALS; + + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); + + // output + output_norm = create_tensor(tn(LLM_TENSOR_OUTPUT_NORM, "weight"), {n_embd}, 0); + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), {n_embd, n_vocab}, TENSOR_NOT_REQUIRED); + + // if output is NULL, init from the input tok embed + if (output == NULL) { + output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, TENSOR_DUPLICATED); + } + + for (int i = 0; i < n_layer; ++i) { + auto & layer = layers[i]; + + layer.attn_norm = create_tensor(tn(LLM_TENSOR_ATTN_NORM, "weight", i), {n_embd}, 0); + + create_tensor_qkv(layer, i, n_embd, n_embd_head_k * n_head, n_embd_k_gqa, n_embd_v_gqa, 0); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", i), {n_embd_head_k * n_head, n_embd}, 0); + + // optional bias tensors + layer.wo_b = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "bias", i), {n_embd}, TENSOR_NOT_REQUIRED); + + // Per-layer attention sinks for iSWA + layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", i), {n_head}, 0); + + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), {n_embd}, 0); + + if (hparams.rope_scaling_type_train == LLAMA_ROPE_SCALING_TYPE_LONGROPE) { + layer.rope_long = create_tensor(tn(LLM_TENSOR_ROPE_FACTORS_LONG, "weight", i), {n_rot/2}, TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0)); + layer.rope_short = create_tensor(tn(LLM_TENSOR_ROPE_FACTORS_SHORT, "weight", i), {n_rot/2}, TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0)); + } + else { + layer.rope_freqs = create_tensor(tn(LLM_TENSOR_ROPE_FREQS, "weight", i), {n_rot/2}, TENSOR_NOT_REQUIRED | (i != 0 ? TENSOR_DUPLICATED : 0)); + } + + if (n_expert == 0) { + layer.ffn_gate = create_tensor(tn(LLM_TENSOR_FFN_GATE, "weight", i), {n_embd, n_ff}, 0); + layer.ffn_down = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "weight", i), { n_ff, n_embd}, 0); + layer.ffn_up = create_tensor(tn(LLM_TENSOR_FFN_UP, "weight", i), {n_embd, n_ff}, 0); + + // optional MLP bias + layer.ffn_gate_b = create_tensor(tn(LLM_TENSOR_FFN_GATE, "bias", i), {n_ff}, TENSOR_NOT_REQUIRED); + layer.ffn_down_b = create_tensor(tn(LLM_TENSOR_FFN_DOWN, "bias", i), {n_embd}, TENSOR_NOT_REQUIRED); + layer.ffn_up_b = create_tensor(tn(LLM_TENSOR_FFN_UP, "bias", i), {n_ff}, TENSOR_NOT_REQUIRED); + } else { + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", i), {n_embd, n_expert}, 0); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", i), { n_ff, n_embd, n_expert}, 0); + create_tensor_gate_up_exps(layer, i, n_embd, n_ff, n_expert, 0); + + // For Granite MoE Shared - gate+up kept fused in ffn_up_shexp (see LLM_FFN_SWIGLU below) + if (hparams.n_ff_shexp > 0) { + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", i), {n_embd, 2*hparams.n_ff_shexp}, 0); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", i), {hparams.n_ff_shexp, n_embd}, 0); + } + } + } +} + +std::unique_ptr llama_model_granite_swa::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*this, params); +} + +llama_model_granite_swa::graph::graph( + const llama_model & model, + const llm_graph_params & params) + : llm_graph_context(params) { + + const int64_t n_embd_head = hparams.n_embd_head_v(); + + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); + GGML_ASSERT(n_embd_head == n_rot); + + ggml_tensor * cur; + ggml_tensor * inpL; + + inpL = build_inp_embd(model.tok_embd); + + // inp_pos - built only if rope enabled + ggml_tensor * inp_pos = build_inp_pos(); + auto * inp_attn = build_attn_inp_kv_iswa(); + + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + for (int il = 0; il < n_layer; ++il) { + + // Granite Vision 4.1 deepstack: inject the projector stream that + // targets decoder layer `il` before the decoder runs. + // NOTE: skip the first deepstack layer since that's inpL + const auto & deepstack_emb_idx = hparams.deepstack_mapping_arr[il]; + if (il > 0 && deepstack_emb_idx >= 0) { + ggml_tensor * ds = ggml_view_2d(ctx0, + res->t_inp_embd, n_embd, n_tokens, + res->t_inp_embd->nb[1], + deepstack_emb_idx * n_embd * sizeof(float)); + inpL = ggml_add(ctx0, inpL, ds); + cb(inpL, "deepstack_in", il); + } + + ggml_tensor * inpSA = inpL; + + // norm + cur = build_norm(inpL, + model.layers[il].attn_norm, NULL, + LLM_NORM_RMS, il); + cb(cur, "attn_norm", il); + + // self-attention + cur = build_attention_layer( + cur, inp_pos, inp_attn, + model, n_embd_head, il); + + if (il == n_layer - 1 && inp_out_ids) { + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + inpSA = ggml_get_rows(ctx0, inpSA, inp_out_ids); + } + // ffn + cur = build_layer_ffn(cur, inpSA, model, il); + + // input for next layer + inpL = cur; + } + cur = inpL; + + cur = build_norm(cur, + model.output_norm, NULL, + LLM_NORM_RMS, -1); + + cb(cur, "result_norm", -1); + res->t_embd = cur; + + // lm_head + cur = build_lora_mm(model.output, cur, model.output_s); + + // For Granite architectures - scale logits + cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_logit_scale); + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} + +ggml_tensor * llama_model_granite_swa::graph::build_attention_layer( + ggml_tensor * cur, + ggml_tensor * inp_pos, + llm_graph_input_attn_kv_iswa * inp_attn, + const llama_model & model, + const int64_t n_embd_head, + const int il) { + + auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur, + n_embd_head, hparams.n_head(il), hparams.n_head_kv(il), il); + + const bool use_rope = hparams.has_rope(il); + if (use_rope) { + ggml_tensor * rope_factors = model.get_rope_factors(cparams, il); + Qcur = ggml_rope_ext( + ctx0, Qcur, inp_pos, rope_factors, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow + ); + + Kcur = ggml_rope_ext( + ctx0, Kcur, inp_pos, rope_factors, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow + ); + } + + cb(Qcur, "Qcur", il); + cb(Kcur, "Kcur", il); + cb(Vcur, "Vcur", il); + + const float kq_scale = hparams.f_attention_scale == 0.0f ? 1.0f/sqrtf(float(n_embd_head)) : hparams.f_attention_scale; + + // Pass layer.attn_sinks to build_attn for sink-based attention modulation + cur = build_attn(inp_attn, + model.layers[il].wo, model.layers[il].wo_b, model.layers[il].wo_s, + Qcur, Kcur, Vcur, nullptr, model.layers[il].attn_sinks, nullptr, kq_scale, il); + cb(cur, "attn_out", il); + return cur; +} + +ggml_tensor * llama_model_granite_swa::graph::build_layer_ffn( + ggml_tensor * cur, + ggml_tensor * inpSA, + const llama_model & model, + const int il) { + + // For Granite architectures - scale residual + if (hparams.f_residual_scale) { + cur = ggml_scale(ctx0, cur, hparams.f_residual_scale); + } + ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpSA); + cb(ffn_inp, "ffn_inp", il); + + // feed-forward network (non-MoE) + if (model.layers[il].ffn_gate_inp == nullptr) { + + cur = build_norm(ffn_inp, + model.layers[il].ffn_norm, NULL, + LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + cur = build_ffn(cur, + model.layers[il].ffn_up, model.layers[il].ffn_up_b, NULL, + model.layers[il].ffn_gate, model.layers[il].ffn_gate_b, NULL, + model.layers[il].ffn_down, model.layers[il].ffn_down_b, NULL, + NULL, + LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(cur, "ffn_out", il); + + } else { + // MoE branch + cur = build_norm(ffn_inp, + model.layers[il].ffn_norm, NULL, + LLM_NORM_RMS, il); + cb(cur, "ffn_norm", il); + + ggml_tensor * moe_out = build_moe_ffn(cur, + model.layers[il].ffn_gate_inp, + model.layers[il].ffn_up_exps, + model.layers[il].ffn_gate_exps, + model.layers[il].ffn_down_exps, + nullptr, + n_expert, n_expert_used, + LLM_FFN_SILU, true, + hparams.expert_weights_scale, + LLAMA_EXPERT_GATING_FUNC_TYPE_SOFTMAX, + il, + nullptr, model.layers[il].ffn_gate_up_exps); + cb(moe_out, "ffn_moe_out", il); + + // For Granite MoE Shared - gate+up kept fused in ffn_up_shexp + if (hparams.n_ff_shexp > 0) { + ggml_tensor * ffn_shexp = build_ffn(cur, + model.layers[il].ffn_up_shexp, NULL, NULL, + NULL, NULL, NULL, + model.layers[il].ffn_down_shexp, NULL, NULL, + NULL, + LLM_FFN_SWIGLU, LLM_FFN_SEQ, il); + cb(ffn_shexp, "ffn_shexp", il); + + cur = ggml_add(ctx0, moe_out, ffn_shexp); + cb(cur, "ffn_out", il); + } else { + cur = moe_out; + } + } + + // For Granite architectures - scale residual + if (hparams.f_residual_scale) { + cur = ggml_scale(ctx0, cur, hparams.f_residual_scale); + } + cur = ggml_add(ctx0, cur, ffn_inp); + cb(cur, "ffn_out", il); + + cur = build_cvec(cur, il); + cb(cur, "l_out", il); + + return cur; +} diff --git a/src/models/granite-switch.cpp b/src/models/granite-switch.cpp index 80f6b86ed..7c9a901c8 100644 --- a/src/models/granite-switch.cpp +++ b/src/models/granite-switch.cpp @@ -11,7 +11,8 @@ void llama_model_granite_switch::load_arch_hparams(llama_model_loader & ml) { bool rope_finetuned = true; ml.get_key(LLM_KV_ROPE_SCALING_FINETUNED, rope_finetuned, false); - hparams.rope_finetuned = rope_finetuned; + hparams.rope_finetuned = rope_finetuned; // needed for round trip save + std::fill(hparams.rope_pattern.begin(), hparams.rope_pattern.end(), rope_finetuned); switch (hparams.n_layer()) { case 40: type = hparams.n_embd == 4096 ? LLM_TYPE_8B : LLM_TYPE_3B; break; @@ -254,7 +255,7 @@ llama_model_granite_switch::graph::graph( cb(inpL, "inp_embd", -1); ggml_tensor * inp_pos = nullptr; - if (hparams.rope_finetuned) { + if (hparams.has_rope(0)) { inp_pos = build_inp_pos(); } auto * inp_attn = build_attn_inp_kv(); @@ -361,7 +362,7 @@ ggml_tensor * llama_model_granite_switch::graph::build_attention_layer( Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); - if (hparams.rope_finetuned) { + if (hparams.has_rope(il)) { ggml_tensor * rope_factors = model.get_rope_factors(cparams, il); Qcur = ggml_rope_ext(ctx0, Qcur, inp_pos, rope_factors, n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, diff --git a/src/models/granite.cpp b/src/models/granite.cpp index 4a75c5ff3..9e9f97e94 100644 --- a/src/models/granite.cpp +++ b/src/models/granite.cpp @@ -33,7 +33,8 @@ void llama_model_granite::load_arch_hparams(llama_model_loader & ml) { // Granite uses rope_finetuned as a switch for rope, so default to true bool rope_finetuned = true; ml.get_key(LLM_KV_ROPE_SCALING_FINETUNED, rope_finetuned, false); - hparams.rope_finetuned = rope_finetuned; + hparams.rope_finetuned = rope_finetuned; // needed for round trip save + std::fill(hparams.rope_pattern.begin(), hparams.rope_pattern.end(), rope_finetuned); switch (hparams.n_layer()) { case 32: type = LLM_TYPE_3B; break; @@ -127,7 +128,7 @@ llama_model_granite::graph::graph( // inp_pos - built only if rope enabled ggml_tensor * inp_pos = nullptr; - if (hparams.rope_finetuned) { + if (hparams.has_rope(0)) { inp_pos = build_inp_pos(); } auto * inp_attn = build_attn_inp_kv(); @@ -203,8 +204,7 @@ ggml_tensor * llama_model_granite::graph::build_attention_layer( auto [Qcur, Kcur, Vcur] = build_qkv(model.layers[il], cur, n_embd_head, hparams.n_head(il), hparams.n_head_kv(il), il); - const bool use_rope = hparams.rope_finetuned; - if (use_rope) { + if (hparams.has_rope(il)) { ggml_tensor * rope_factors = model.get_rope_factors(cparams, il); Qcur = ggml_rope_ext( ctx0, Qcur, inp_pos, rope_factors, diff --git a/src/models/models.h b/src/models/models.h index 180b30a46..1dd30dfd1 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -1719,6 +1719,34 @@ struct llama_model_granite_hybrid : public llama_model_base { }; +struct llama_model_granite_swa : public llama_model_base { + llama_model_granite_swa(const struct llama_model_params & params) : llama_model_base(params) {} + void load_arch_hparams(llama_model_loader & ml) override; + void load_arch_tensors(llama_model_loader & ml) override; + + struct graph : public llm_graph_context { + graph(const llama_model & model, const llm_graph_params & params); + + private: + ggml_tensor * build_attention_layer( + ggml_tensor * cur, + ggml_tensor * inp_pos, + llm_graph_input_attn_kv_iswa * inp_attn, + const llama_model & model, + const int64_t n_embd_head, + const int il); + + ggml_tensor * build_layer_ffn( + ggml_tensor * cur, + ggml_tensor * inpSA, + const llama_model & model, + const int il); + }; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + + struct llama_model_chameleon : public llama_model_base { llama_model_chameleon(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override; diff --git a/src/unicode.cpp b/src/unicode.cpp index b02ecdc93..93996f9dd 100644 --- a/src/unicode.cpp +++ b/src/unicode.cpp @@ -1241,7 +1241,7 @@ std::vector unicode_regex_split(const std::string & text, const std { unicode_cpt_flags::LETTER, "\x41-\x5A\x61-\x7A" }, // A-Za-z { unicode_cpt_flags::PUNCTUATION, "\x21-\x23\x25-\x2A\x2C-\x2F\x3A-\x3B\x3F-\x40\\\x5B-\\\x5D\x5F\\\x7B\\\x7D" }, // !-#%-*,-/:-;?-@\[-\]_\{\} { unicode_cpt_flags::ACCENT_MARK, "" }, // no sub-128 codepoints - { unicode_cpt_flags::SYMBOL, "\\\x24\\\x2B\x3C-\x3E\x5E\x60\\\x7C" }, // $+<=>^`| + { unicode_cpt_flags::SYMBOL, "\\\x24\\\x2B\x3C-\x3E\x5E\x60\\\x7C\\\x7E" }, // $+<=>^`|~ }; // compute collapsed codepoints only if needed by at least one regex diff --git a/tools/mtmd/README-dev.md b/tools/mtmd/README-dev.md index ac43e1b81..e14906823 100644 --- a/tools/mtmd/README-dev.md +++ b/tools/mtmd/README-dev.md @@ -21,7 +21,7 @@ A typical pipeline of the core libmtmd is as follows: - A bitmap (RGB image or PCM audio) is created - Bitmap and the text prompt is provided to `mtmd_tokenize()` that breaks the input into chunks - The tokenizer function first expands a "lazy" bitmap if it finds one. Typically, this is used by video, so that one media token corresponds to one input bitmap - - For models that support "fused" temporal frames like Qwen-VL, the tokenizer tries to merge pair of consecutive frames into one batch + - For models that support "fused" temporal frames like Qwen-VL, the tokenizer tries to merge pair of consecutive frames into one batch. Only bitmaps marked by `mtmd_bitmap_set_mergeable()` are merged - The preprocessor will then be called, which produces a list of chunks - Depending on the model itself, special tokens will be injected to separate image chunks (i.e. llava-uhd-style models) - Multiple bitmaps may be batched together to form a larger `mtmd_batch()` diff --git a/tools/mtmd/clip-impl.h b/tools/mtmd/clip-impl.h index 86c2c1c94..3fcd12233 100644 --- a/tools/mtmd/clip-impl.h +++ b/tools/mtmd/clip-impl.h @@ -860,6 +860,9 @@ static std::ifstream open_ifstream_binary(const std::string & fname) { } #endif +// in test-mtmd-impl, we include woth common.h and this file, and these functions are duplicated +// this is a quick fix to avoid compilation errors +#ifndef DIRECTORY_SEPARATOR static std::string string_format(const char * fmt, ...) { va_list ap; va_list ap2; @@ -918,6 +921,7 @@ inline bool string_ends_with2(std::string_view str, std::string_view suffix) { return str.size() >= suffix.size() && str.compare(str.size() - suffix.size(), suffix.size(), suffix) == 0; } +#endif // // gguf utils diff --git a/tools/mtmd/models/deepseekocr.cpp b/tools/mtmd/models/deepseekocr.cpp index 0ba5a4d2a..b784cdad6 100644 --- a/tools/mtmd/models/deepseekocr.cpp +++ b/tools/mtmd/models/deepseekocr.cpp @@ -88,6 +88,22 @@ static ggml_tensor * get_rel_pos(ggml_context * ctx0, return cur; // [C, k_size, q_size] } +// ggml_conv_2d with the im2col kept in F32: the F16 im2col it emits since #23660 degrades OCR +static ggml_tensor * conv_2d_f32(ggml_context * ctx0, ggml_tensor * a, ggml_tensor * b, + int s0, int s1, int p0, int p1, int d0, int d1) { + const ggml_type im2col_type = a->type == GGML_TYPE_F16 ? GGML_TYPE_F16 : GGML_TYPE_F32; + ggml_tensor * im2col = ggml_im2col(ctx0, a, b, s0, s1, p0, p1, d0, d1, true, im2col_type); // [N, OH, OW, IC * KH * KW] + + ggml_tensor * result = ggml_mul_mat(ctx0, + ggml_reshape_2d(ctx0, im2col, im2col->ne[0], im2col->ne[3] * im2col->ne[2] * im2col->ne[1]), + ggml_reshape_2d(ctx0, a, (a->ne[0] * a->ne[1] * a->ne[2]), a->ne[3])); + + result = ggml_reshape_4d(ctx0, result, im2col->ne[1], im2col->ne[2], im2col->ne[3], a->ne[3]); // [OC, N, OH, OW] + result = ggml_cont(ctx0, ggml_permute(ctx0, result, 0, 1, 3, 2)); // [N, OC, OH, OW] + + return result; +} + ggml_tensor * clip_graph_deepseekocr::build_sam(ggml_tensor * inp_raw) { // Building SAM @@ -101,7 +117,8 @@ ggml_tensor * clip_graph_deepseekocr::build_sam(ggml_tensor * inp_raw) { ggml_tensor * inpL; - inpL = ggml_conv_2d_sk_p0(ctx0, model.patch_embed_proj_w, inp_raw); + inpL = conv_2d_f32(ctx0, model.patch_embed_proj_w, inp_raw, + (int) model.patch_embed_proj_w->ne[0], (int) model.patch_embed_proj_w->ne[1], 0, 0, 1, 1); inpL = ggml_add(ctx0, inpL, ggml_reshape_3d(ctx0, model.patch_embed_proj_b, 1, 1, n_embd)); inpL = ggml_cont(ctx0, ggml_permute(ctx0, inpL, 1, 2, 0, 3)); @@ -229,18 +246,18 @@ ggml_tensor * clip_graph_deepseekocr::build_sam(ggml_tensor * inp_raw) { cur = ggml_cont(ctx0, ggml_permute(ctx0, cur, 2, 0, 1, 3)); - cur = ggml_conv_2d(ctx0, model.neck_0_w, cur, 1, 1, 0, 0, 1, 1); + cur = conv_2d_f32(ctx0, model.neck_0_w, cur, 1, 1, 0, 0, 1, 1); cur = ggml_cont(ctx0, ggml_permute(ctx0, cur, 1, 2, 0, 3)); cur = build_norm(cur, model.neck_1_w, model.neck_1_b, NORM_TYPE_NORMAL, sam_eps, -1); cur = ggml_cont(ctx0, ggml_permute(ctx0, cur, 2, 0, 1, 3)); - cur = ggml_conv_2d(ctx0, model.neck_2_w, cur, 1, 1, 1, 1, 1, 1); + cur = conv_2d_f32(ctx0, model.neck_2_w, cur, 1, 1, 1, 1, 1, 1); cur = ggml_cont(ctx0, ggml_permute(ctx0, cur, 1, 2, 0, 3)); cur = build_norm(cur, model.neck_3_w, model.neck_3_b, NORM_TYPE_NORMAL, sam_eps, -1); cur = ggml_cont(ctx0, ggml_permute(ctx0, cur, 2, 0, 1, 3)); - cur = ggml_conv_2d(ctx0, model.net_2, cur, 2, 2, 1, 1, 1, 1); - cur = ggml_conv_2d(ctx0, model.net_3, cur, 2, 2, 1, 1, 1, 1); + cur = conv_2d_f32(ctx0, model.net_2, cur, 2, 2, 1, 1, 1, 1); + cur = conv_2d_f32(ctx0, model.net_3, cur, 2, 2, 1, 1, 1, 1); cb(cur, "sam_output", -1); ggml_build_forward_expand(gf, cur); diff --git a/tools/mtmd/mtmd-helper.cpp b/tools/mtmd/mtmd-helper.cpp index 3681b9ee0..788089652 100644 --- a/tools/mtmd/mtmd-helper.cpp +++ b/tools/mtmd/mtmd-helper.cpp @@ -728,7 +728,9 @@ struct mtmd_helper_video { LOG_DBG("%s: frame %d read OK\n", __func__, current_frame); current_frame++; - return mtmd_bitmap_init(info.width, info.height, frame_buf.data()); + mtmd_bitmap * frame = mtmd_bitmap_init(info.width, info.height, frame_buf.data()); + mtmd_bitmap_set_mergeable(frame, true); + return frame; } int32_t read_next(mtmd_bitmap ** out_bitmap, char ** out_text) { diff --git a/tools/mtmd/mtmd-image.cpp b/tools/mtmd/mtmd-image.cpp index 769b6efe6..0d9db4f62 100644 --- a/tools/mtmd/mtmd-image.cpp +++ b/tools/mtmd/mtmd-image.cpp @@ -1013,14 +1013,31 @@ mtmd_image_preproc_out mtmd_image_preprocessor_lfm2::preprocess(const clip_image return output; } +bool mtmd_image_preprocessor_lfm2::should_tile( + const clip_hparams & hparams, + const clip_image_size & original_size) { + const int align_size = hparams.patch_size * hparams.n_merge; + + const auto round_by_factor = [align_size](float x) { + // see https://github.com/ggml-org/llama.cpp/pull/27057#discussion_r3796264887 + return static_cast(std::nearbyint(static_cast(x) / align_size)) * align_size; + }; + + const int h_bar = std::max(hparams.patch_size, round_by_factor(original_size.height)); + const int w_bar = std::max(hparams.patch_size, round_by_factor(original_size.width)); + + return static_cast(h_bar) * static_cast(w_bar) > + static_cast(hparams.image_max_pixels) * max_pixels_tolerance; +} + mtmd_image_preprocessor_llava_uhd::slice_instructions mtmd_image_preprocessor_lfm2::get_slice_instructions(const clip_image_size & original_size) { mtmd_image_preprocessor_llava_uhd::slice_instructions inst; const int align_size = hparams.patch_size * hparams.n_merge; inst.overview_size = img_tool::calc_size_preserved_ratio( original_size, { align_size, hparams.image_min_pixels, hparams.image_max_pixels, 0 }); - // tile if either dimension exceeds tile_size with tolerance - const bool needs_tiling = original_size.width > tile_size * max_pixels_tolerance || original_size.height > tile_size * max_pixels_tolerance; + + const bool needs_tiling = should_tile(hparams, original_size); if (!needs_tiling) { inst.refined_size = clip_image_size{0, 0}; diff --git a/tools/mtmd/mtmd-image.h b/tools/mtmd/mtmd-image.h index 40dfea7eb..732e27379 100644 --- a/tools/mtmd/mtmd-image.h +++ b/tools/mtmd/mtmd-image.h @@ -148,6 +148,8 @@ struct mtmd_image_preprocessor_lfm2 : mtmd_image_preprocessor_llava_uhd { mtmd_image_preproc_out preprocess(const clip_image_u8 & img) override; slice_instructions get_slice_instructions(const clip_image_size & original_size) override; + static bool should_tile(const clip_hparams & hparams, const clip_image_size & original_size); + private: clip_image_size find_closest_aspect_ratio( float aspect_ratio, diff --git a/tools/mtmd/mtmd-internal.h b/tools/mtmd/mtmd-internal.h new file mode 100644 index 000000000..067fa88b9 --- /dev/null +++ b/tools/mtmd/mtmd-internal.h @@ -0,0 +1,19 @@ +#pragma once + +#include "mtmd.h" + +#include +#include + +// !!! Internal header, to be used by mtmd and its unit tests only !!! + +#define MTMD_INTERNAL_HEADER + +// bitmap is null for text parts +struct mtmd_input_part { + std::string text; + const mtmd_bitmap * bitmap; +}; + +// [QWEN_VIDEO] merged parts are erased from `parts`, so one group always maps to one part +std::vector> mtmd_group_mergeable_bitmaps(std::vector & parts, int n_merge); diff --git a/tools/mtmd/mtmd.cpp b/tools/mtmd/mtmd.cpp index 4063d28e0..0f5cb8c7a 100644 --- a/tools/mtmd/mtmd.cpp +++ b/tools/mtmd/mtmd.cpp @@ -1,6 +1,7 @@ #include "clip.h" #include "clip-impl.h" #include "mtmd.h" +#include "mtmd-internal.h" #include "mtmd-audio.h" #include "mtmd-image.h" #include "debug/mtmd-debug.h" @@ -149,6 +150,7 @@ struct mtmd_bitmap { uint32_t ny = 0; std::string id; // optional user-defined id, for ex: can be set to image hash, useful for KV cache tracking bool is_audio = false; // true if the bitmap is audio + bool mergeable = false; // [QWEN_VIDEO] set only on frames of the same video // lazy-loaded bitmap mtmd_bitmap_lazy_callback lazy_callback = nullptr; @@ -186,7 +188,9 @@ struct mtmd_bitmap { bool can_merge_with(const mtmd_bitmap & other) const { // [QWEN_VIDEO] can (temporal) merge if both are images with same size - return !is_audio && !other.is_audio && nx == other.nx && ny == other.ny; + return mergeable && other.mergeable + && !is_audio && !other.is_audio + && nx == other.nx && ny == other.ny; } private: @@ -1076,6 +1080,25 @@ void mtmd_free(mtmd_context * ctx) { delete ctx; } +std::vector> mtmd_group_mergeable_bitmaps(std::vector & parts, int n_merge) { + std::vector> output; + for (size_t i = 0; i < parts.size(); i++) { + if (parts[i].bitmap == nullptr) { + continue; // text part + } + const bool has_next = n_merge > 1 && i + 1 < parts.size() && parts[i + 1].bitmap != nullptr; + if (has_next && parts[i].bitmap->can_merge_with(*parts[i + 1].bitmap)) { + LOG_DBG("%s: merging 2 frames at part index %zu and %zu\n", __func__, i, i + 1); + output.push_back({parts[i].bitmap, parts[i + 1].bitmap}); + parts.erase(parts.begin() + i + 1); + continue; + } + LOG_DBG("%s: no merging for part index %zu\n", __func__, i); + output.push_back({parts[i].bitmap}); + } + return output; +} + struct mtmd_tokenizer { mtmd_context * ctx; @@ -1084,10 +1107,7 @@ struct mtmd_tokenizer { bool parse_special; const llama_vocab * vocab; - struct part { - std::string text; - const mtmd_bitmap * bitmap; - }; + using part = mtmd_input_part; std::vector parts; // these will be freed when mtmd_tokenizer finishes std::vector bm_from_lazy; // TODO @ngxson : refactor, free bm_from_lazy progressively @@ -1192,34 +1212,7 @@ struct mtmd_tokenizer { GGML_ASSERT(n_merge_frames <= 2 && "we only support merging maximum 2 images for now; open an issue if this model supports merging more"); } - // Build merged_bitmaps: each entry is a group of 1 or 2 bitmaps. - // For consecutive mergeable bitmap parts, merge them and collapse the second part out of this->parts. - std::vector> merged_bitmaps; - if (n_merge_frames > 1) { - for (size_t i = 0; i < parts.size(); ++i) { - if (parts[i].bitmap == nullptr) { - continue; - } - if (i + 1 < parts.size() && parts[i + 1].bitmap != nullptr) { - const mtmd_bitmap * bm_a = parts[i].bitmap; - const mtmd_bitmap * bm_b = parts[i + 1].bitmap; - if (bm_a->can_merge_with(*bm_b)) { - LOG_DBG("%s: merging 2 frames at part index %zu and %zu\n", __func__, i, i + 1); - merged_bitmaps.push_back({bm_a, bm_b}); - parts.erase(parts.begin() + i + 1); // collapse the second bitmap part - continue; - } - } - LOG_DBG("%s: no merging for part index %zu\n", __func__, i); - merged_bitmaps.push_back({parts[i].bitmap}); - } - } else { - for (const auto & p : parts) { - if (p.bitmap != nullptr) { - merged_bitmaps.push_back({p.bitmap}); - } - } - } + auto merged_bitmaps = mtmd_group_mergeable_bitmaps(parts, n_merge_frames); size_t i_bm = 0; for (const auto & p : parts) { @@ -2200,6 +2193,10 @@ void mtmd_bitmap_set_id(mtmd_bitmap * bitmap, const char * id) { } } +void mtmd_bitmap_set_mergeable(mtmd_bitmap * bitmap, bool mergeable) { + bitmap->mergeable = mergeable; +} + mtmd_bitmap * mtmd_bitmap_init_lazy(mtmd_context * ctx, const char * id, void * user_data, diff --git a/tools/mtmd/mtmd.h b/tools/mtmd/mtmd.h index 78587f3fe..ef4f99c0b 100644 --- a/tools/mtmd/mtmd.h +++ b/tools/mtmd/mtmd.h @@ -154,7 +154,8 @@ MTMD_API const char * mtmd_get_marker(const mtmd_context * ctx); // length of data must be nx * ny * 3 // the data is in RGBRGBRGB... format // note: some video-capable models (i.e. qwen-vl) can merge consecutive bitmaps -// into one chunk, mtmd_tokenize() will automatically handle this +// into one chunk; mtmd_tokenize() handles this, but remember to set +// mtmd_bitmap_set_mergeable(true) for every frame // if bitmap is audio: // length of data must be n_samples * sizeof(float) // the data is in float format (PCM F32) @@ -175,6 +176,8 @@ MTMD_API void mtmd_bitmap_free (mtmd_bitmap * bitmap); // these getters/setters are dedicated functions, so you can for example calculate the hash of the image based on mtmd_bitmap_get_data() MTMD_API const char * mtmd_bitmap_get_id(const mtmd_bitmap * bitmap); MTMD_API void mtmd_bitmap_set_id(mtmd_bitmap * bitmap, const char * id); +// if true, this bitmap can be merged (temporal merge) with an adjacent mergeable bitmap by certain video input models +MTMD_API void mtmd_bitmap_set_mergeable(mtmd_bitmap * bitmap, bool mergeable); // mtmd_bitmap lazy // diff --git a/tools/server/README-dev.md b/tools/server/README-dev.md index 613017acf..0f42b2ee1 100644 --- a/tools/server/README-dev.md +++ b/tools/server/README-dev.md @@ -189,7 +189,7 @@ This endpoint is intended to be used internally by the Web UI and subject to cha Get a list of tools, each tool has these fields: - `tool` (string): the ID name of the tool, to be used in POST call. Example: `read_file` - `display_name` (string): the name to be displayed on UI. Example: `Read file` -- `type` (string): `"builtin"` for a built-in tool, or `"mcp"` for a tool exposed by an MCP server +- `type` (string): `"server"` for a server tool, or `"mcp"` for a tool exposed by an MCP server - `permissions` (object): a mapping string --> boolean that indicates the permission required by this tool. This is useful for the UI to ask the user before calling the tool. For now, the only permission supported is `"write"` - `definition` (object): the OAI-compat definition of this tool @@ -291,6 +291,36 @@ The flow for downloading a new model: - If a stop request comes in, the router asks the child process to stop (same mechanism as running a model in child process) - Otherwise, upon completion, we call `load_models()` to refresh the list of models +### Sleep mode + +Sleep mode was initially introduced in PR [#18228](https://github.com/ggml-org/llama.cpp/pull/18228). The main idea is to have: +- `server_queue` keeping track of the idle timeout +- When the timeout is detected, `server_queue` signals to `server_context_impl` that it should go into sleep +- `server_context_impl` frees all `llama_context` and `mtmd_context` + +Compared to simply exiting the whole process, this approach allows accessing some read-only endpoints during sleep, while also handling wakeup-on-request. Any inference request will wake the server up. + +Call stack on entering sleeping: +- `server_queue::start_loop` (main thread) sees no task for `idle_sleep_ms` --> `sleeping = true` +- `cb0(true)` --> `server_routes::update_cached_responses` + - snapshots `/props`, `/models` and metrics; the model is still alive here +- `cb1(true)` --> `server_context_impl::handle_sleeping_state` + - `callback_state(SERVER_STATE_SLEEPING)` --> reported to router in child mode + - `destroy()` --> frees `llama_context` and `mtmd_context` +- `condition_tasks.wait` until `req_stop_sleeping` + +Call stack on waking up: +- `server_res_generator` constructor (HTTP thread) --> `server_queue::wait_until_no_sleep` + - sets `req_stop_sleeping = true`, then waits until `sleeping == false` +- `server_queue::start_loop` (main thread) wakes up +- `cb1(false)` --> `server_context_impl::handle_sleeping_state` + - `load_model()`, which then emits `callback_state(SERVER_STATE_READY)` +- `cb0(false)` --> `server_routes::update_cached_responses` + - nothing to do, the cache is only read during sleep +- `sleeping = false` --> `notify_all` unblocks the HTTP thread, the request is handled as usual + +Endpoints created with `create_response(true)` (`/health`, `/props`, `/models`, `/metrics`) skip `wait_until_no_sleep`, so they answer from the cached responses instead of waking the server. + ### Notable Related PRs - Initial server implementation: https://github.com/ggml-org/llama.cpp/pull/1443 diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index 842e4203c..21ff78394 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -818,6 +818,14 @@ public: } } + server_metrics get_metrics() const { + return metrics; + } + + void reset_metrics_bucket() { + metrics.reset_bucket(); + } + private: // note: accessing these fields outside of this class is not thread-safe // use server_context methods instead @@ -898,6 +906,10 @@ private: void handle_sleeping_state(bool new_state) { GGML_ASSERT(sleeping != new_state); if (new_state) { + if (callback_state) { + callback_state(SERVER_STATE_SLEEPING, {}); + // note: for sleeping == false, event is emitted by load_model() + } SRV_INF("%s", "server is entering sleeping state\n"); destroy(); } else { @@ -2290,8 +2302,8 @@ private: // returns false to decline the task, it is offered again after the decode is done bool process_single_task(server_task && task, bool is_yielding) { - // while yielding, an encode / decode is running and only accessing metrics is safe - if (is_yielding && task.type != SERVER_TASK_TYPE_METRICS) { + // while yielding, an encode / decode is running and only reading the server state is safe + if (is_yielding && task.type != SERVER_TASK_TYPE_METRICS && task.type != SERVER_TASK_TYPE_SLOT_GET) { SRV_DBG("decoding, decline task, id_task = %d\n", task.id); return false; } @@ -2417,28 +2429,17 @@ private: } break; case SERVER_TASK_TYPE_METRICS: { - json slots_data = json::array(); - - int n_idle_slots = 0; int n_processing_slots = 0; for (server_slot & slot : slots) { - json slot_data = slot.to_json(slots_debug == 0); - if (slot.is_processing()) { n_processing_slots++; - } else { - n_idle_slots++; } - - slots_data.push_back(slot_data); } - SRV_DBG("n_idle_slots = %d, n_processing_slots = %d\n", n_idle_slots, n_processing_slots); + SRV_DBG("n_processing_slots = %d\n", n_processing_slots); auto res = std::make_unique(); res->id = task.id; - res->slots_data = std::move(slots_data); - res->n_idle_slots = n_idle_slots; res->n_processing_slots = n_processing_slots; res->n_tasks_deferred = queue_tasks.queue_tasks_deferred_size(); res->metrics = metrics; @@ -2446,6 +2447,28 @@ private: if (task.metrics_reset_bucket) { metrics.reset_bucket(); } + queue_results.send(std::move(res)); + } break; + case SERVER_TASK_TYPE_SLOT_GET: + { + json slots_data = json::array(); + + int n_idle_slots = 0; + + for (server_slot & slot : slots) { + if (!slot.is_processing()) { + n_idle_slots++; + } + + slots_data.push_back(slot.to_json(slots_debug == 0)); + } + SRV_DBG("n_idle_slots = %d\n", n_idle_slots); + + auto res = std::make_unique(); + res->id = task.id; + res->slots_data = std::move(slots_data); + res->n_idle_slots = n_idle_slots; + queue_results.send(std::move(res)); } break; case SERVER_TASK_TYPE_SLOT_SAVE: @@ -4142,12 +4165,6 @@ struct server_res_generator : server_res_spipe { void server_context::set_state_callback(server_state_callback_t callback) { impl->callback_state = std::move(callback); - impl->queue_tasks.on_sleeping_state([this](bool sleeping) { - if (sleeping) { - impl->callback_state(SERVER_STATE_SLEEPING, {}); - } - // for sleeping == false, event is emitted by load_model() - }); } // @@ -4431,6 +4448,119 @@ server_routes::server_routes(const common_params & params, server_context & ctx_ queue_tasks(ctx_server.impl->queue_tasks), queue_results(ctx_server.impl->queue_results) { init_routes(); + + // note: this must be registered before load_model() + // so that on sleep phase, the callback is called before ctx is destroyed + queue_tasks.on_sleeping_state([this](bool is_sleeping) { + update_cached_responses(is_sleeping); + }); +} + +static json get_res_model_info(const server_context_meta & meta) { + // note: do NOT use ctx_server here, otherwise it's not possible to use this during sleep + + return { + {"id", meta.model_name}, + {"aliases", meta.model_aliases}, + {"tags", meta.model_tags}, + {"object", "model"}, + {"created", std::time(0)}, + {"owned_by", "llamacpp"}, + {"meta", { + {"vocab_type", meta.model_vocab_type}, + {"n_vocab", meta.model_vocab_n_tokens}, + {"n_ctx", meta.slot_n_ctx}, + {"n_ctx_train", meta.model_n_ctx_train}, + {"n_embd", meta.model_n_embd_inp}, + {"n_params", meta.model_n_params}, + {"size", meta.model_size}, + {"ftype", meta.model_ftype}, + }}, + }; +} + +static json get_res_models(const server_context_meta & meta) { + // note: do NOT use ctx_server here, otherwise it's not possible to use this during sleep + + return { + {"models", { + { + {"name", meta.model_name}, + {"model", meta.model_name}, + {"modified_at", ""}, + {"size", ""}, + {"digest", ""}, // dummy value, llama.cpp does not support managing model file's hash + {"type", "model"}, + {"description", ""}, + {"tags", {""}}, + {"capabilities", meta.has_mtmd ? json({"completion","multimodal"}) : json({"completion"})}, + {"parameters", ""}, + {"details", { + {"parent_model", ""}, + {"format", "gguf"}, + {"family", ""}, + {"families", {""}}, + {"parameter_size", ""}, + {"quantization_level", ""} + }} + } + }}, + {"object", "list"}, + {"data", { + get_res_model_info(meta), + }} + }; +} + +static json get_res_props(const server_context_meta & meta, const common_params & params, bool is_sleeping) { + // note: do NOT use ctx_server here, otherwise it's not possible to use this during sleep + + task_params tparams; + tparams.sampling = params.sampling; + json default_generation_settings_for_props = json { + { "params", tparams.to_json(true) }, + { "n_ctx", meta.slot_n_ctx }, + }; + + std::string tmpl_default = common_chat_templates_source(meta.chat_params.tmpls.get(), ""); + std::string tmpl_tools = common_chat_templates_source(meta.chat_params.tmpls.get(), "tool_use"); + + json props = { + { "default_generation_settings", default_generation_settings_for_props }, + { "total_slots", params.n_parallel }, + { "model_alias", meta.model_name }, + { "model_ftype", meta.model_ftype }, + { "model_path", meta.model_path }, + { "modalities", json { + {"vision", meta.has_inp_image}, + {"video", meta.has_inp_video}, + {"audio", meta.has_inp_audio}, + } }, + { "media_marker", get_media_marker() }, + { "endpoint_slots", params.endpoint_slots }, + { "endpoint_props", params.endpoint_props }, + { "endpoint_metrics", params.endpoint_metrics }, + { "ui", params.ui }, + { "ui_settings", meta.json_ui_settings }, + { "chat_template", tmpl_default }, + { "chat_template_caps", meta.chat_template_caps }, + { "bos_token", meta.bos_token_str }, + { "eos_token", meta.eos_token_str }, + { "build_info", meta.build_info }, + { "is_sleeping", is_sleeping }, + { "cors_proxy_enabled", params.ui_mcp_proxy }, + }; + if (params.use_jinja) { + if (!tmpl_tools.empty()) { + props["chat_template_tool_use"] = tmpl_tools; + } + } + + return props; +} + +json server_routes::get_model_info() const { + return get_res_model_info(*meta); } void server_routes::init_routes() { @@ -4451,41 +4581,64 @@ void server_routes::init_routes() { }; this->get_metrics = [this](const server_http_req & req) { - auto res = create_response(); + auto res = create_response(true); if (!params.endpoint_metrics) { res->error(format_error_response("This server does not support metrics endpoint. Start it with `--metrics`", ERROR_TYPE_NOT_SUPPORTED)); return res; } - // request slots data using task queue - { - server_task task(SERVER_TASK_TYPE_METRICS); - task.id = res->rd.get_new_id(); + // render response using cached_metrics + auto use_cached_metrics = [&]() { + std::unique_lock lock(mutex_cache); + res->headers["Process-Start-Time-Unix"] = std::to_string(cached_metrics.t_start); + server_task_result_metrics tmp; + tmp.metrics = cached_metrics; + res->content_type = "text/plain; version=0.0.4"; + res->status = 200; + res->data = tmp.to_metrics(); // the gauges are averaged over the window between two scrapes - task.metrics_reset_bucket = true; - res->rd.post_task(std::move(task), true); // high-priority task + cached_metrics.reset_bucket(); + should_reset_buckets = true; + }; + + if (queue_tasks.is_sleeping()) { + use_cached_metrics(); + + } else { + // request slots data using task queue + { + server_task task(SERVER_TASK_TYPE_METRICS); + task.id = res->rd.get_new_id(); + // the gauges are averaged over the window between two scrapes + task.metrics_reset_bucket = true; + res->rd.post_task(std::move(task), true); // high-priority task + } + + // a task posted right before sleeping is never processed, do not wait for it + auto result = res->rd.next([&]{ + return req.should_stop() || queue_tasks.is_sleeping(); + }); + if (!result) { + if (!req.should_stop()) { + use_cached_metrics(); + } + return res; + } + + if (result->is_error()) { + res->error(result->to_json()); + return res; + } + + auto res_task = dynamic_cast(result.get()); + GGML_ASSERT(res_task != nullptr); + + res->headers["Process-Start-Time-Unix"] = std::to_string(res_task->metrics.t_start); + res->content_type = "text/plain; version=0.0.4"; + res->status = 200; + res->data = res_task->to_metrics(); } - // get the result - auto result = res->rd.next(req.should_stop); - if (!result) { - // connection was closed - GGML_ASSERT(req.should_stop()); - return res; - } - - if (result->is_error()) { - res->error(result->to_json()); - return res; - } - - auto res_task = dynamic_cast(result.get()); - GGML_ASSERT(res_task != nullptr); - - res->headers["Process-Start-Time-Unix"] = std::to_string(res_task->metrics.t_start); - res->content_type = "text/plain; version=0.0.4"; - res->status = 200; - res->data = res_task->to_metrics(); return res; }; @@ -4498,7 +4651,7 @@ void server_routes::init_routes() { // request slots data using task queue { - server_task task(SERVER_TASK_TYPE_METRICS); + server_task task(SERVER_TASK_TYPE_SLOT_GET); task.id = res->rd.get_new_id(); res->rd.post_task(std::move(task), true); // high-priority task } @@ -4516,7 +4669,7 @@ void server_routes::init_routes() { return res; } - auto * res_task = dynamic_cast(result.get()); + auto * res_task = dynamic_cast(result.get()); GGML_ASSERT(res_task != nullptr); // optionally return "fail_on_no_slot" error @@ -4566,53 +4719,13 @@ void server_routes::init_routes() { this->get_props = [this](const server_http_req &) { auto res = create_response(true); - - // this endpoint can be accessed during sleeping - // the next LOC is to avoid someone accidentally use ctx_server - bool ctx_server; // do NOT delete this line - GGML_UNUSED(ctx_server); - - task_params tparams; - tparams.sampling = params.sampling; - json default_generation_settings_for_props = json { - { "params", tparams.to_json(true) }, - { "n_ctx", meta->slot_n_ctx }, - }; - - std::string tmpl_default = common_chat_templates_source(meta->chat_params.tmpls.get(), ""); - std::string tmpl_tools = common_chat_templates_source(meta->chat_params.tmpls.get(), "tool_use"); - - json props = { - { "default_generation_settings", default_generation_settings_for_props }, - { "total_slots", params.n_parallel }, - { "model_alias", meta->model_name }, - { "model_ftype", meta->model_ftype }, - { "model_path", meta->model_path }, - { "modalities", json { - {"vision", meta->has_inp_image}, - {"video", meta->has_inp_video}, - {"audio", meta->has_inp_audio}, - } }, - { "media_marker", get_media_marker() }, - { "endpoint_slots", params.endpoint_slots }, - { "endpoint_props", params.endpoint_props }, - { "endpoint_metrics", params.endpoint_metrics }, - { "ui", params.ui }, - { "ui_settings", meta->json_ui_settings }, - { "chat_template", tmpl_default }, - { "chat_template_caps", meta->chat_template_caps }, - { "bos_token", meta->bos_token_str }, - { "eos_token", meta->eos_token_str }, - { "build_info", meta->build_info }, - { "is_sleeping", queue_tasks.is_sleeping() }, - { "cors_proxy_enabled", params.ui_mcp_proxy }, - }; - if (params.use_jinja) { - if (!tmpl_tools.empty()) { - props["chat_template_tool_use"] = tmpl_tools; - } + // note: do NOT use ctx_server here, this endpoint must be accessible during sleep + if (queue_tasks.is_sleeping()) { + std::unique_lock lock(mutex_cache); + res->ok(cached_props); + } else { + res->ok(get_res_props(*meta, params, false)); } - res->ok(props); return res; }; @@ -4874,42 +4987,13 @@ void server_routes::init_routes() { this->get_models = [this](const server_http_req &) { auto res = create_response(true); - - // this endpoint can be accessed during sleeping - // the next LOC is to avoid someone accidentally use ctx_server - bool ctx_server; // do NOT delete this line - GGML_UNUSED(ctx_server); - - json models = { - {"models", { - { - {"name", meta->model_name}, - {"model", meta->model_name}, - {"modified_at", ""}, - {"size", ""}, - {"digest", ""}, // dummy value, llama.cpp does not support managing model file's hash - {"type", "model"}, - {"description", ""}, - {"tags", {""}}, - {"capabilities", meta->has_mtmd ? json({"completion","multimodal"}) : json({"completion"})}, - {"parameters", ""}, - {"details", { - {"parent_model", ""}, - {"format", "gguf"}, - {"family", ""}, - {"families", {""}}, - {"parameter_size", ""}, - {"quantization_level", ""} - }} - } - }}, - {"object", "list"}, - {"data", { - get_model_info(), - }} - }; - - res->ok(models); + // note: do NOT use ctx_server here, this endpoint must be accessible during sleep + if (queue_tasks.is_sleeping()) { + std::unique_lock lock(mutex_cache); + res->ok(cached_models); + } else { + res->ok(get_res_models(*meta)); + } return res; }; @@ -5119,27 +5203,6 @@ void server_routes::init_routes() { }; } -json server_routes::get_model_info() const { - return json { - {"id", meta->model_name}, - {"aliases", meta->model_aliases}, - {"tags", meta->model_tags}, - {"object", "model"}, - {"created", std::time(0)}, - {"owned_by", "llamacpp"}, - {"meta", { - {"vocab_type", meta->model_vocab_type}, - {"n_vocab", meta->model_vocab_n_tokens}, - {"n_ctx", meta->slot_n_ctx}, - {"n_ctx_train", meta->model_n_ctx_train}, - {"n_embd", meta->model_n_embd_inp}, - {"n_params", meta->model_n_params}, - {"size", meta->model_size}, - {"ftype", meta->model_ftype}, - }}, - }; -} - std::unique_ptr server_routes::handle_slots_save(const server_http_req & req, int id_slot) { auto res = create_response(); const json request_data = json::parse(req.body); @@ -5388,3 +5451,24 @@ std::unique_ptr server_routes::handle_count_tokens(const l res->ok(response); return res; } + +void server_routes::update_cached_responses(bool is_sleeping) { + // caller is task_queue, so ctx_server can be accessed without holding locks + std::unique_lock lock(mutex_cache); + + if (is_sleeping) { + cached_models = get_res_models(*meta); + cached_props = get_res_props(*meta, params, true); + cached_metrics = ctx_server.get_metrics(); + + should_reset_buckets = false; + + SRV_DBG("%s\n", "cached responses updated"); + + } else if (should_reset_buckets) { + // a scrape during sleep already reported these buckets + ctx_server.reset_metrics_bucket(); + + should_reset_buckets = false; + } +} diff --git a/tools/server/server-context.h b/tools/server/server-context.h index f9ab1132b..764df0e08 100644 --- a/tools/server/server-context.h +++ b/tools/server/server-context.h @@ -8,6 +8,7 @@ #include #include +#include #include struct server_context_impl; // private implementation @@ -174,9 +175,19 @@ private: std::unique_ptr meta; const common_params & params; - const server_context_impl & ctx_server; + server_context_impl & ctx_server; server_queue & queue_tasks; server_response & queue_results; std::unique_ptr create_response(bool bypass_sleep = false); + + // cached responses, to be used during sleep + std::mutex mutex_cache; + json cached_models = nullptr; + json cached_props = nullptr; + server_metrics cached_metrics; + // set when a scrape during sleep already reported the throughput buckets + bool should_reset_buckets = false; + // call right before sleep to update the cached responses + void update_cached_responses(bool is_sleeping); }; diff --git a/tools/server/server-http.cpp b/tools/server/server-http.cpp index b11dc09d0..2ec137aa0 100644 --- a/tools/server/server-http.cpp +++ b/tools/server/server-http.cpp @@ -198,8 +198,6 @@ bool server_http_context::init(const common_params & params) { std::unordered_set endpoints { "/health", "/v1/health", - "/models", - "/v1/models", }; endpoints.insert(frontend_paths.begin(), frontend_paths.end()); return endpoints; diff --git a/tools/server/server-models.cpp b/tools/server/server-models.cpp index 93e940951..35b935570 100644 --- a/tools/server/server-models.cpp +++ b/tools/server/server-models.cpp @@ -555,6 +555,40 @@ void server_models::load_models() { return source_map.count(name) ? source_map.at(name) : SERVER_MODEL_SOURCE_PRESET; }; + // hide cache models whose resolved file is already used by a preset with dedup-cache-models enabled + std::set hidden_models; + { + std::set preset_paths; + for (const auto & [name, preset] : custom_presets) { + std::string val; + if (!preset.get_option(COMMON_ARG_PRESET_DEDUP_CACHE_MODELS, val) || !common_arg_utils::is_truthy(val)) { + continue; + } + std::string hf_repo; + if (!preset.get_option("LLAMA_ARG_HF_REPO", hf_repo) || hf_repo.empty()) { + continue; + } + std::string hf_file; + preset.get_option("LLAMA_ARG_HF_FILE", hf_file); + std::string path = common_download_resolve_path(hf_repo, hf_file); + if (!path.empty()) { + preset_paths.insert(path); + } + } + if (!preset_paths.empty()) { + for (const auto & [name, preset] : cached_models) { + if (get_source(name) != SERVER_MODEL_SOURCE_CACHE) { + continue; // merged with another source, not a pure cache entry + } + std::string path = common_download_resolve_path(name); + if (!path.empty() && preset_paths.count(path)) { + SRV_INF("hiding cache model name=%s (deduplicated by a preset)\n", name.c_str()); + hidden_models.insert(name); + } + } + } + } + // Helpers that read `mapping` - must be called while holding the lock. std::unordered_set custom_names; for (const auto & [name, preset] : custom_presets) custom_names.insert(name); @@ -590,6 +624,11 @@ void server_models::load_models() { } } }; + auto apply_hidden = [&]() { + for (auto & [name, inst] : mapping) { + inst.meta.hidden = hidden_models.count(name) > 0; + } + }; // update_args() injects HOST/PORT/ALIAS, so strip them before comparing presets auto preset_options_for_compare = [](common_preset p) { p.unset_option("LLAMA_ARG_HOST"); @@ -630,6 +669,7 @@ void server_models::load_models() { add_model(std::move(meta)); } apply_stop_timeout(); + apply_hidden(); log_available_models(); std::vector models_to_load; @@ -806,6 +846,7 @@ void server_models::load_models() { } apply_stop_timeout(); + apply_hidden(); // clear reload flag before unlocking for autoload - load() blocks on !is_reloading, // so clearing it here (while still locked) prevents a deadlock in the autoload calls below @@ -1025,10 +1066,13 @@ void server_models::load(const std::string & name, const load_options & opts) { char * buffer = vec_buf.data(); if (stdout_file) { while (fgets(buffer, vec_buf.size(), stdout_file) != nullptr) { - LOG("[%5d] %s", port, buffer); std::string str(buffer); if (string_starts_with(buffer, CMD_CHILD_TO_ROUTER_STATE)) { + LOG_DBG("[%5d] %s", port, buffer); // prevent spamming the log this->handle_child_state(name, str); + } else { + // forward log + LOG("[%5d] %s", port, buffer); } } } else { @@ -1926,6 +1970,9 @@ void server_models_routes::init_routes() { auto all_models = models.get_all_meta(); std::time_t t = std::time(0); for (const auto & meta : all_models) { + if (meta.hidden) { + continue; // cache model deduplicated by a preset + } json status { {"value", server_model_status_to_string(meta.status)}, {"args", meta.args}, diff --git a/tools/server/server-models.h b/tools/server/server-models.h index 615acb577..79b231cba 100644 --- a/tools/server/server-models.h +++ b/tools/server/server-models.h @@ -84,6 +84,7 @@ struct server_model_meta { int exit_code = 0; // exit code of the model instance process (only valid if status == FAILED) int stop_timeout = 0; // seconds to wait before force-killing the model instance during shutdown mtmd_caps multimodal; // multimodal capabilities + bool hidden = false; // hidden from GET /models, but still accept if requested bool is_ready() const { return status == SERVER_MODEL_STATUS_LOADED; diff --git a/tools/server/server-queue.cpp b/tools/server/server-queue.cpp index 2bcc9bd8f..78169e9a5 100644 --- a/tools/server/server-queue.cpp +++ b/tools/server/server-queue.cpp @@ -3,6 +3,7 @@ #include "log.h" +#include #include #include @@ -20,6 +21,10 @@ // server_queue // +static bool task_resets_idle_timer(server_task_type type) { + return type != SERVER_TASK_TYPE_METRICS; +} + int server_queue::post(server_task && task, bool front) { std::unique_lock lock(mutex_tasks); GGML_ASSERT(task.id != -1); @@ -27,20 +32,24 @@ int server_queue::post(server_task && task, bool front) { if (task.type == SERVER_TASK_TYPE_CANCEL) { cleanup_pending_task(task.id_target); } - const int task_id = task.id; + const int task_id = task.id; + const bool reset_timer = task_resets_idle_timer(task.type); QUE_DBG("new task, id = %d, front = %d\n", task_id, front); if (front) { queue_tasks.push_front(std::move(task)); } else { queue_tasks.push_back(std::move(task)); } - time_last_task = ggml_time_ms(); + if (reset_timer) { + time_last_task = ggml_time_ms(); + } condition_tasks.notify_one(); return task_id; } int server_queue::post(std::vector && tasks, bool front) { std::unique_lock lock(mutex_tasks); + bool reset_timer = false; for (auto & task : tasks) { if (task.id == -1) { task.id = id++; @@ -49,6 +58,7 @@ int server_queue::post(std::vector && tasks, bool front) { if (task.type == SERVER_TASK_TYPE_CANCEL) { cleanup_pending_task(task.id_target); } + reset_timer |= task_resets_idle_timer(task.type); QUE_DBG("new task, id = %d/%d, front = %d\n", task.id, (int) tasks.size(), front); if (front) { queue_tasks.push_front(std::move(task)); @@ -56,7 +66,9 @@ int server_queue::post(std::vector && tasks, bool front) { queue_tasks.push_back(std::move(task)); } } - time_last_task = ggml_time_ms(); + if (reset_timer) { + time_last_task = ggml_time_ms(); + } condition_tasks.notify_one(); return 0; } @@ -294,11 +306,14 @@ void server_queue::start_loop(int64_t idle_sleep_ms) { QUE_DBG("%s", "update slots\n"); // this will run the main inference process for all slots + const int64_t t_update_slots = ggml_time_ms(); callback_update_slots(); { // update_slots() may take a while to finish, we need to make sure it's not counted as idle + // shift instead of reset, so that non-task_resets_idle_timer tasks do not delay the sleep std::unique_lock lock(mutex_tasks); - time_last_task = ggml_time_ms(); + const int64_t now = ggml_time_ms(); + time_last_task = std::min(now, time_last_task + (now - t_update_slots)); } QUE_DBG("%s", "waiting for new tasks\n"); @@ -312,7 +327,10 @@ void server_queue::start_loop(int64_t idle_sleep_ms) { if (should_sleep()) { QUE_INF("%s", "entering sleeping state\n"); sleeping = true; - callback_sleeping_state(true); + // Call order cb0 -> cb1 -> cb{N} + for (auto & cb : callback_sleeping_state) { + cb(true); + } req_stop_sleeping = false; // wait until we are requested to exit sleeping state condition_tasks.wait(lock, [&]{ @@ -323,7 +341,10 @@ void server_queue::start_loop(int64_t idle_sleep_ms) { } QUE_INF("%s", "exiting sleeping state\n"); req_stop_sleeping = false; - callback_sleeping_state(false); + // Call order cb{N} -> cb1 -> cb0 + for (size_t i = callback_sleeping_state.size(); i > 0; i--) { + callback_sleeping_state[i - 1](false); + } sleeping = false; time_last_task = ggml_time_ms(); condition_tasks.notify_all(); // notify wait_until_no_sleep() diff --git a/tools/server/server-queue.h b/tools/server/server-queue.h index 52d30095c..e17733a74 100644 --- a/tools/server/server-queue.h +++ b/tools/server/server-queue.h @@ -44,7 +44,7 @@ private: // callback functions std::function callback_new_task; std::function callback_update_slots; - std::function callback_sleeping_state; + std::vector> callback_sleeping_state; public: ~server_queue() { worker_stop(); } @@ -86,6 +86,7 @@ public: * * Sleeping procedure (disabled if idle_sleep_ms < 0): * - If there is no task after idle_sleep_ms, enter sleeping state + * note: metrics tasks are processed as usual, but do not reset the idle timer * - Call callback_sleeping_state(true) * - Wait until req_stop_sleeping is set to true * - Call callback_sleeping_state(false) @@ -127,18 +128,12 @@ public: } // Register callback for sleeping state change; multiple callbacks are allowed - // note: when entering sleeping state, the callback is called AFTER sleeping is set to true - // when leaving sleeping state, the callback is called BEFORE sleeping is set to false + // for example: register order cb0, cb1, cb2 + // entering sleep: queue.sleeping = true --> cb0(true) --> cb1(true) --> cb2(true) + // leaving sleep: cb2(false) --> cb1(false) --> cb0(false) --> queue.sleeping = false + // note: caller will hold mutex_tasks while calling the callbacks void on_sleeping_state(std::function callback) { - if (callback_sleeping_state) { - auto prev_callback = std::move(callback_sleeping_state); - callback_sleeping_state = [prev_callback, callback](bool sleeping) { - prev_callback(sleeping); - callback(sleeping); - }; - } else { - callback_sleeping_state = std::move(callback); - } + callback_sleeping_state.push_back(std::move(callback)); } private: diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index 64afbc5ed..258cdcf8f 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -1512,10 +1512,15 @@ json server_task_result_error::to_json() { // // server_task_result_metrics // -json server_task_result_metrics::to_json() { +json server_task_result_slots::to_json() { return slots_data; } +json server_task_result_metrics::to_json() { + // not used, /metrics renders prometheus text via to_metrics() + return json{}; +} + // metrics definition: https://prometheus.io/docs/practices/naming/#metric-names std::string server_task_result_metrics::to_metrics() { const std::vector counters = { diff --git a/tools/server/server-task.h b/tools/server/server-task.h index b6da4d4bd..25ff01512 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -22,6 +22,7 @@ enum server_task_type { SERVER_TASK_TYPE_CONTROL, SERVER_TASK_TYPE_NEXT_RESPONSE, SERVER_TASK_TYPE_METRICS, + SERVER_TASK_TYPE_SLOT_GET, SERVER_TASK_TYPE_SLOT_SAVE, SERVER_TASK_TYPE_SLOT_RESTORE, SERVER_TASK_TYPE_SLOT_ERASE, @@ -489,22 +490,16 @@ struct server_task_result_error : server_task_result { virtual json to_json() override; }; +// used by /metrics API struct server_task_result_metrics : server_task_result { // these are immediate stats, not accumulated (server_metrics is cumulative) - int n_idle_slots; - int n_processing_slots; - int n_tasks_deferred; + int n_processing_slots = 0; + int n_tasks_deferred = 0; server_metrics metrics; - // while we can also use std::vector this requires copying the slot object which can be quite messy - // therefore, we use json to temporarily store the slot.to_json() result - json slots_data = json::array(); - - // used by /slots API virtual json to_json() override; - // used by /metrics API struct metric_item { std::string name; std::string description; @@ -513,6 +508,17 @@ struct server_task_result_metrics : server_task_result { std::string to_metrics(); }; +// used by /slots API +struct server_task_result_slots : server_task_result { + int n_idle_slots = 0; + + // while we can also use std::vector this requires copying the slot object which can be quite messy + // therefore, we use json to temporarily store the slot.to_json() result + json slots_data = json::array(); + + virtual json to_json() override; +}; + struct server_task_result_slot_save_load : server_task_result { std::string filename; bool is_save; // true = save, false = load diff --git a/tools/server/server-tools.cpp b/tools/server/server-tools.cpp index 5e5e60efd..b5c5c078a 100644 --- a/tools/server/server-tools.cpp +++ b/tools/server/server-tools.cpp @@ -2035,7 +2035,7 @@ void server_tools::setup(const std::vector & enabled_tools, } } - // append MCP tools, skipping any that collide with a built-in or another MCP tool of the same "_" name + // append MCP tools, skipping any that collide with a server tool or another MCP tool of the same "_" name if (!mcp_mgr.empty()) { std::unordered_set seen_names; for (auto & t : tools) { diff --git a/tools/server/server-tools.h b/tools/server/server-tools.h index c4509ca80..e7332f2e5 100644 --- a/tools/server/server-tools.h +++ b/tools/server/server-tools.h @@ -18,7 +18,7 @@ struct server_tool { virtual ~server_tool() = default; virtual json get_definition() const = 0; - virtual std::string type() const { return "builtin"; } + virtual std::string type() const { return "server"; } struct stream { server_response & qr; diff --git a/tools/server/server.cpp b/tools/server/server.cpp index 6d1aa4351..01cc6633a 100644 --- a/tools/server/server.cpp +++ b/tools/server/server.cpp @@ -235,8 +235,8 @@ int llama_server(common_params & params, int argc, char ** argv) { ctx_http.get ("/metrics", ex_wrapper(routes.get_metrics)); ctx_http.get ("/props", ex_wrapper(routes.get_props)); ctx_http.post("/props", ex_wrapper(routes.post_props)); - ctx_http.get ("/models", ex_wrapper(routes.get_models)); // public endpoint (no API key check) - ctx_http.get ("/v1/models", ex_wrapper(routes.get_models)); // public endpoint (no API key check) + ctx_http.get ("/models", ex_wrapper(routes.get_models)); + ctx_http.get ("/v1/models", ex_wrapper(routes.get_models)); ctx_http.post("/completion", ex_wrapper(routes.post_completions)); // legacy ctx_http.post("/completions", ex_wrapper(routes.post_completions)); ctx_http.post("/v1/completions", ex_wrapper(routes.post_completions_oai)); @@ -346,7 +346,7 @@ int llama_server(common_params & params, int argc, char ** argv) { ctx_http.get ("/tools", ex_wrapper(tools.handle_get)); ctx_http.post("/tools", ex_wrapper(tools.handle_post)); if (!params.server_tools.empty()) { - warn_names.push_back("built-in tools (experimental)"); + warn_names.push_back("server tools (experimental)"); } if (!params.server_tools_runtime.empty()) { warn_names.push_back("tools runtime (experimental)"); diff --git a/tools/server/tests/unit/test_router.py b/tools/server/tests/unit/test_router.py index 5ab62666c..96eb87978 100644 --- a/tools/server/tests/unit/test_router.py +++ b/tools/server/tests/unit/test_router.py @@ -63,14 +63,16 @@ def test_router_chat_completion_stream(model: str, success: bool): assert content == "" -def _get_model_ids(is_reload: bool) -> set[str]: - res = server.make_request("GET", "/models" + ("?reload=1" if is_reload else "")) +def _get_model_ids(is_reload: bool, headers: dict | None = None) -> set[str]: + res = server.make_request( + "GET", "/models" + ("?reload=1" if is_reload else ""), headers=headers + ) assert res.status_code == 200 return {item["id"] for item in res.body.get("data", [])} -def _get_model_status(model_id: str) -> str: - res = server.make_request("GET", "/models") +def _get_model_status(model_id: str, headers: dict | None = None) -> str: + res = server.make_request("GET", "/models", headers=headers) assert res.status_code == 200 for item in res.body.get("data", []): if item.get("id") == model_id or item.get("model") == model_id: @@ -78,11 +80,11 @@ def _get_model_status(model_id: str) -> str: raise AssertionError(f"Model {model_id} not found in /models response") -def _wait_for_model_status(model_id: str, desired: set[str], timeout: int = 60) -> str: +def _wait_for_model_status(model_id: str, desired: set[str], timeout: int = 60, headers: dict | None = None) -> str: deadline = time.time() + timeout last_status = None while time.time() < deadline: - last_status = _get_model_status(model_id) + last_status = _get_model_status(model_id, headers=headers) if last_status in desired: return last_status time.sleep(0.01) @@ -100,7 +102,7 @@ def _load_model_and_wait( assert load_res.status_code == 200 assert isinstance(load_res.body, dict) assert load_res.body.get("success") is True - _wait_for_model_status(model_id, {"loaded"}, timeout=timeout) + _wait_for_model_status(model_id, {"loaded"}, timeout=timeout, headers=headers) def test_router_unload_model(): @@ -406,6 +408,59 @@ def test_router_reload_models(): os.remove(preset_path) +def test_router_dedup_cache_models(): + """dedup-cache-models hides the cache entry backing a preset from GET /models""" + global server + + preset_path = os.path.join(TMP_DIR, "test_dedup.ini") + cache_id = "ggml-org/test-model-stories260K:F32" + + with open(preset_path, "w") as f: + f.write( + "[model-dedup]\n" + "hf-repo = ggml-org/test-model-stories260K\n" + "dedup-cache-models = 1\n" + ) + + server.models_preset = preset_path + server.start() + + try: + ids = _get_model_ids(is_reload=False) + assert "model-dedup" in ids + assert cache_id not in ids, "cache model should be hidden by dedup" + # other cache models are unaffected + assert "ggml-org/tinygemma3-GGUF:Q8_0" in ids + + # the hidden model is only hidden from the listing, it can still be used + res = server.make_request("POST", "/tokenize", data={"model": cache_id, "content": "hello"}) + assert res.status_code == 200 + + # disabling the flag brings the cache entry back on reload + with open(preset_path, "w") as f: + f.write( + "[model-dedup]\n" + "hf-repo = ggml-org/test-model-stories260K\n" + ) + ids = _get_model_ids(is_reload=True) + assert cache_id in ids + + # the flag also works from the global section + with open(preset_path, "w") as f: + f.write( + "[*]\n" + "dedup-cache-models = 1\n" + "\n" + "[model-dedup]\n" + "hf-repo = ggml-org/test-model-stories260K\n" + ) + ids = _get_model_ids(is_reload=True) + assert "model-dedup" in ids + assert cache_id not in ids, "cache model should be hidden by global dedup" + finally: + os.remove(preset_path) + + def test_router_remote_preset(): global server server.model_hf_repo = "ggml-org/test-preset-ci" diff --git a/tools/server/tests/unit/test_security.py b/tools/server/tests/unit/test_security.py index ac0544575..36fc439f9 100644 --- a/tools/server/tests/unit/test_security.py +++ b/tools/server/tests/unit/test_security.py @@ -15,7 +15,7 @@ def create_server(): server.api_key = TEST_API_KEY -@pytest.mark.parametrize("endpoint", ["/health", "/models"]) +@pytest.mark.parametrize("endpoint", ["/health"]) def test_access_public_endpoint(endpoint: str): global server server.start() diff --git a/tools/server/tests/unit/test_sleep.py b/tools/server/tests/unit/test_sleep.py index 3374165e8..515f7077d 100644 --- a/tools/server/tests/unit/test_sleep.py +++ b/tools/server/tests/unit/test_sleep.py @@ -11,6 +11,35 @@ def create_server(): server = ServerPreset.tinyllama2() +def is_sleeping(server: ServerProcess) -> bool: + res = server.make_request("GET", "/props") + assert res.status_code == 200 + return res.body["is_sleeping"] + + +def wait_for_sleep(server: ServerProcess, timeout: float = 10.0): + start = time.time() + while time.time() - start < timeout: + if is_sleeping(server): + return + time.sleep(0.1) + raise TimeoutError("server did not go to sleep") + + +def fetch_metrics(server: ServerProcess) -> str: + res = server.make_request("GET", "/metrics") + assert res.status_code == 200 + assert isinstance(res.body, str) + return res.body + + +def get_metric(text: str, name: str) -> float: + prefix = f"llamacpp:{name} " + values = [ln for ln in text.splitlines() if ln.startswith(prefix)] + assert len(values) == 1, f"{name} not found in metrics" + return float(values[0][len(prefix):]) + + def test_server_sleep(): global server server.sleep_idle_seconds = 1 @@ -25,6 +54,10 @@ def test_server_sleep(): res = server.make_request("GET", "/props") assert res.status_code == 200 assert res.body["is_sleeping"] == True + res = server.make_request("GET", "/models") + assert res.status_code == 200 + assert len(res.body["data"]) == 1 + assert res.body["data"][0]["id"] == server.model_alias # make a generation request to wake up the server res = server.make_request("POST", "/completion", data={ @@ -37,3 +70,58 @@ def test_server_sleep(): res = server.make_request("GET", "/props") assert res.status_code == 200 assert res.body["is_sleeping"] == False + + +def test_server_sleep_read_only_endpoints(): + global server + server.sleep_idle_seconds = 1 + server.server_metrics = True + server.start() + + res = server.make_request("POST", "/completion", data={ + "n_predict": 4, + "prompt": "Hello", + }) + assert res.status_code == 200 + + # the first scrape resets the throughput buckets, so that the second one reports + # the same zero rates as the snapshot taken on entering sleep + fetch_metrics(server) + metrics_awake = fetch_metrics(server) + assert get_metric(metrics_awake, "tokens_predicted_total") > 0 + + wait_for_sleep(server) + + # during sleep, metrics are served from the snapshot taken right before sleeping + assert fetch_metrics(server) == metrics_awake + + # scraping /metrics must not wake the server up + assert is_sleeping(server) + + +def test_server_sleep_metrics_buckets(): + global server + server.sleep_idle_seconds = 1 + server.server_metrics = True + server.start() + + res = server.make_request("POST", "/completion", data={ + "n_predict": 8, + "prompt": "Hello", + }) + assert res.status_code == 200 + + wait_for_sleep(server) + + # the first scrape reports the throughput of the last generation + assert get_metric(fetch_metrics(server), "predicted_tokens_seconds") > 0 + + # nothing runs while sleeping, so the next scrapes report an empty window + assert get_metric(fetch_metrics(server), "predicted_tokens_seconds") == 0 + assert is_sleeping(server) + + # waking up must not report the buckets again + res = server.make_request("POST", "/tokenize", data={"content": "Hello"}) + assert res.status_code == 200 + assert is_sleeping(server) == False + assert get_metric(fetch_metrics(server), "predicted_tokens_seconds") == 0 diff --git a/tools/ui/eslint.config.js b/tools/ui/eslint.config.js index c65484048..b8bdb216e 100644 --- a/tools/ui/eslint.config.js +++ b/tools/ui/eslint.config.js @@ -61,6 +61,9 @@ export default ts.config( { blankLine: 'always', next: ['return', 'throw', 'break', 'continue'], prev: '*' } ], + // Alphabetical order for enum members + 'perfectionist/sort-enums': ['error', { type: 'natural' }], + 'perfectionist/sort-objects': ['error', { type: 'natural' }], // Alphabetical order for variable declarations and object keys diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatForm.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatForm.svelte index a9f721e47..e937d2730 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatForm.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatForm.svelte @@ -48,6 +48,7 @@ containsFileMentionLink, findCommandToken, findMentionToken, + getConversationModel, isIMEComposing, isOffsetInCodeBlock, parseClipboardContent, @@ -190,31 +191,9 @@ let isRouter = $derived(serverStore.isRouterMode); let conversationModel = $derived( - chatStore.getConversationModel(conversationsStore.activeMessages as DatabaseMessage[]) + getConversationModel(conversationsStore.activeMessages as DatabaseMessage[]) ); - let activeModelId = $derived.by(() => { - const options = modelsStore.models; - - if (!isRouter) { - return options.length > 0 ? options[0].model : null; - } - - const selectedId = modelsStore.selectedModelId; - - if (selectedId) { - const model = options.find((m) => m.id === selectedId); - - if (model) return model.model; - } - - if (conversationModel) { - const model = options.find((m) => m.model === conversationModel); - - if (model) return model.model; - } - - return null; - }); + let activeModelId = $derived(modelsStore.activeModelId); let hasModelSelected = $derived( !isRouter || !!conversationModel || !!modelsStore.selectedModelId diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddToolsSubmenu.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddToolsSubmenu.svelte index 1204390fd..58ae10673 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddToolsSubmenu.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddToolsSubmenu.svelte @@ -35,7 +35,7 @@ Run llama-server with {CLI_FLAGS.TOOLS} flag to enable - Built-in Tools. + Server Tools. diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionsAdd.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionsAdd.svelte index 47bdb47a4..b2581f11e 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionsAdd.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionsAdd.svelte @@ -2,10 +2,10 @@ import ChatFormActionAddButton from './ChatFormActionAddButton.svelte'; import ChatFormActionAddDropdown from './ChatFormActionAddDropdown.svelte'; import ChatFormActionAddSheet from './ChatFormActionAddSheet.svelte'; - import { isMobile } from '$lib/stores'; + import { deviceStore } from '$lib/stores'; -{#if isMobile.current} +{#if deviceStore.isMobile} {#snippet trigger({ disabled, onclick })} diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionModels.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionModels.svelte index 9f163c6d6..690e13dfa 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionModels.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionModels.svelte @@ -1,6 +1,7 @@ -{#if isMobile.current} +{#if deviceStore.isMobile} import ContextGaugeDial from './ContextGaugeDial.svelte'; - import { useContextGauge } from '$lib/hooks/use-context-gauge.svelte'; import { - chatStore, - conversationsStore, gaugeTriggerClick, gaugeTriggerEnter, gaugeTriggerKeydown, gaugeTriggerLeave, gaugeTriggerPointerDown - } from '$lib/stores'; + } from './gauge-popup.svelte'; + import { useContextGauge } from '$lib/hooks/use-context-gauge.svelte'; + import { chatStore, conversationsStore } from '$lib/stores'; import { untrack } from 'svelte'; const gauge = useContextGauge(); diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeDetails.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeDetails.svelte index eaaba69de..1153c70fd 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeDetails.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormContextGauge/ContextGaugeDetails.svelte @@ -1,9 +1,9 @@ diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ToolCallBlock.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ToolCallBlock.svelte index 4b8524fe6..08da39a2a 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ToolCallBlock.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/ToolCallBlock.svelte @@ -14,8 +14,8 @@ import { ICON_CLASS_DEFAULT, ICON_CLASS_SPIN } from '$lib/constants'; import { AgenticSectionType } from '$lib/enums'; import { mcpStore } from '$lib/stores'; - import type { AgenticSection, BuiltinToolUiEntry } from '$lib/types'; - import { getBuiltinToolUi } from '$lib/utils'; + import type { AgenticSection, ToolUiEntry } from '$lib/types'; + import { getToolUi } from '$lib/utils'; import type { Component, Snippet } from 'svelte'; type ToolCallBlockMetaWithError = TMeta & { errorMessage?: string }; @@ -82,7 +82,7 @@ const showSpinner = $derived(isPending || (isStreamingCall && isStreaming) || extraLiveStreaming); const isCodeStreaming = $derived(isStreaming && (isPending || isStreamingCall)); - const toolUi: BuiltinToolUiEntry | null = $derived(getBuiltinToolUi(section.toolName)); + const toolUi: ToolUiEntry | null = $derived(getToolUi(section.toolName)); const toolIcon: Component = $derived( spinIconWhenActive && showSpinner ? Loader2 : (toolUi?.icon ?? Wrench) ); diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/edit-file.ts b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/edit-file.ts index dadb33a49..9ed6f92bc 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/edit-file.ts +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/edit-file.ts @@ -24,7 +24,7 @@ export type EditFileMeta = { }; export function parseEditFileMeta(section: AgenticSection): EditFileMeta | null { - const args = parseToolArgs(BuiltInTool.EDIT_FILE, section, { partial: true }); + const args = parseToolArgs(BuiltInTool.SERVER_EDIT_FILE, section, { partial: true }); if (!args) return null; diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/exec-shell-command.ts b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/exec-shell-command.ts index 496fcdde6..7cf767535 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/exec-shell-command.ts +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/exec-shell-command.ts @@ -14,7 +14,7 @@ export type ExecShellCommandMeta = { }; export function parseExecShellCommandMeta(section: AgenticSection): ExecShellCommandMeta | null { - const args = parseToolArgs(BuiltInTool.EXEC_SHELL_COMMAND, section); + const args = parseToolArgs(BuiltInTool.SERVER_EXEC_SHELL_COMMAND, section); if (!args) return null; diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/file-glob-search.ts b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/file-glob-search.ts index 0acd53d77..237afa599 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/file-glob-search.ts +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/file-glob-search.ts @@ -19,7 +19,7 @@ export type FileGlobSearchMeta = { }; export function parseFileGlobSearchMeta(section: AgenticSection): FileGlobSearchMeta | null { - const args = parseToolArgs(BuiltInTool.FILE_GLOB_SEARCH, section); + const args = parseToolArgs(BuiltInTool.SERVER_FILE_GLOB_SEARCH, section); if (!args) return null; diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/grep-search.ts b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/grep-search.ts index c25d76b1a..90889ff27 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/grep-search.ts +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/grep-search.ts @@ -28,7 +28,7 @@ export type GrepSearchMeta = { }; export function parseGrepSearchMeta(section: AgenticSection): GrepSearchMeta | null { - const args = parseToolArgs(BuiltInTool.GREP_SEARCH, section); + const args = parseToolArgs(BuiltInTool.SERVER_GREP_SEARCH, section); if (!args) return null; diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/read-file.ts b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/read-file.ts index 9ee748ed7..af0f3d925 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/read-file.ts +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/read-file.ts @@ -16,7 +16,7 @@ export type ReadFileMeta = { }; export function parseReadFileMeta(section: AgenticSection): ReadFileMeta | null { - const args = parseToolArgs(BuiltInTool.READ_FILE, section, { partial: true }); + const args = parseToolArgs(BuiltInTool.SERVER_READ_FILE, section, { partial: true }); if (!args) return null; diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/run-javascript.ts b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/run-javascript.ts index a524478ab..440a1f5d6 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/run-javascript.ts +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/run-javascript.ts @@ -16,7 +16,7 @@ export type RunJavascriptMeta = { }; export function parseRunJavascriptMeta(section: AgenticSection): RunJavascriptMeta | null { - const args = parseToolArgs(BuiltInTool.RUN_JAVASCRIPT, section); + const args = parseToolArgs(BuiltInTool.BROWSER_RUN_JAVASCRIPT, section); if (!args) return null; diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/write-file.ts b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/write-file.ts index 53ba38e12..5b9bf9f88 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/write-file.ts +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessage/ChatMessageToolCall/parsers/write-file.ts @@ -20,7 +20,7 @@ export type WriteFileMeta = { }; export function parseWriteFileMeta(section: AgenticSection): WriteFileMeta | null { - const args = parseToolArgs(BuiltInTool.WRITE_FILE, section, { partial: true }); + const args = parseToolArgs(BuiltInTool.SERVER_WRITE_FILE, section, { partial: true }); if (!args) return null; diff --git a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageActions/ChatMessageActionCard/ChatMessageActionCardPermissionRequest.svelte b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageActions/ChatMessageActionCard/ChatMessageActionCardPermissionRequest.svelte index a0cee94f4..e8af94464 100644 --- a/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageActions/ChatMessageActionCard/ChatMessageActionCardPermissionRequest.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatMessages/ChatMessageActions/ChatMessageActionCard/ChatMessageActionCardPermissionRequest.svelte @@ -61,8 +61,8 @@ {:else} {@const source = toolsStore.getToolSource(toolName)} {@const providerName = - source === ToolSource.BUILTIN - ? TOOL_SERVER_LABELS[ToolSource.BUILTIN] + source === ToolSource.SERVER + ? TOOL_SERVER_LABELS[ToolSource.SERVER] : source === ToolSource.CUSTOM ? TOOL_SERVER_LABELS[ToolSource.CUSTOM] : 'MCP Tools'} 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 2b5ca68de..d78de1f85 100644 --- a/tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreen.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreen.svelte @@ -21,8 +21,7 @@ import { chatStore, conversationsStore, - device, - isMobile, + deviceStore, serverStore, settingsStore } from '$lib/stores'; @@ -32,7 +31,7 @@ let { showCenteredEmpty = false } = $props(); let disableAutoScroll = $derived( - Boolean(settingsStore.config.disableAutoScroll) || isMobile.current + Boolean(settingsStore.config.disableAutoScroll) || deviceStore.isMobile ); let isMobileUserScrolledUp = $state(false); let mobileScrollDownHint = $state(false); @@ -52,11 +51,11 @@ let hasPropsError = $derived(!!serverStore.error); let isCurrentConversationLoading = $derived(chatStore.isLoading || chatStore.isStreaming()); let chatFormBottomPosition = $derived.by(() => { - if (!isMobile.current) return '1rem'; + if (!deviceStore.isMobile) return '1rem'; - if (device.isStandalone) return '1.5rem'; + if (deviceStore.isStandalone) return '1.5rem'; - if (device.isIOSSafari) return '0.25rem'; + if (deviceStore.isIOSSafari) return '0.25rem'; return '0.5rem'; }); @@ -84,7 +83,7 @@ }); function handleMobileScroll() { - if (!isMobile.current) return; + if (!deviceStore.isMobile) return; const container = scroll.chatScrollContainer; @@ -184,7 +183,7 @@ } function handleSendLikeScroll() { - if (!isMobile.current) { + if (!deviceStore.isMobile) { autoScroll.enable(); } @@ -197,7 +196,7 @@ '.chat-message:nth-last-child(2) .chat-message-user .chat-message-user-bubble' ) as HTMLElement | null; - if (isMobile.current) { + if (deviceStore.isMobile) { // Keep the last user message bubble just above the input on mobile const bubbleHeight = lastUserBubble?.scrollHeight ?? 0; const baseHeight = container.scrollHeight - innerHeight; @@ -220,7 +219,7 @@ } }, 100); - if (isMobile.current) { + if (deviceStore.isMobile) { autoScroll.setDisabled(disableAutoScroll); mobileScrollDownHint = true; mobileScrollDownHintLockedUntil = Date.now() + 500; @@ -243,7 +242,8 @@ $effect(() => { const shouldDisableAutoScroll = - settingsStore.config.disableAutoScroll || (isMobile.current && isCurrentConversationLoading); + settingsStore.config.disableAutoScroll || + (deviceStore.isMobile && isCurrentConversationLoading); autoScroll.setDisabled(shouldDisableAutoScroll); @@ -266,7 +266,7 @@ autoScroll.enable(); } - if (isMobile.current && isCurrentConversationLoading) { + if (deviceStore.isMobile && isCurrentConversationLoading) { mobileScrollDownHint = true; mobileScrollDownHintLockedUntil = Date.now() + 500; } @@ -318,9 +318,9 @@
- {#if (isMobile.current ? mobileScrollDownHint || isMobileUserScrolledUp : autoScroll.userScrolledUp) && page.url.hash.includes(ROUTES.CHAT) && page.params.id} + {#if (deviceStore.isMobile ? mobileScrollDownHint || isMobileUserScrolledUp : autoScroll.userScrolledUp) && page.url.hash.includes(ROUTES.CHAT) && page.params.id} { mobileScrollDownHint = false; diff --git a/tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreenForm.svelte b/tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreenForm.svelte index 4119b2816..464502980 100644 --- a/tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreenForm.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatScreen/ChatScreenForm.svelte @@ -3,7 +3,7 @@ import { page } from '$app/state'; import { ChatForm } from '$lib/components/app'; import { useDraftMessages } from '$lib/hooks/use-draft-messages.svelte'; - import { isMobile } from '$lib/stores'; + import { deviceStore } from '$lib/stores'; import { onMount } from 'svelte'; interface Props { @@ -120,13 +120,13 @@ } onMount(() => { - if (!isMobile.current) { + if (!deviceStore.isMobile) { setTimeout(focusFormUnlessCaptured, 100); } }); afterNavigate((navigation) => { - if (navigation?.from != null && !isMobile.current) { + if (navigation?.from != null && !deviceStore.isMobile) { setTimeout(focusFormUnlessCaptured, 100); } }); diff --git a/tools/ui/src/lib/components/app/chat/index.ts b/tools/ui/src/lib/components/app/chat/index.ts index 2de8a6ace..34571d53b 100644 --- a/tools/ui/src/lib/components/app/chat/index.ts +++ b/tools/ui/src/lib/components/app/chat/index.ts @@ -278,7 +278,7 @@ export { default as ChatFormInput } from './ChatForm/ChatFormInput/ChatFormInput /** * Working directory selector for agent mode. Renders a chip below the chat * form; clicking it opens a popover with a directory picker backed by the - * server's `file_glob_search` built-in tool (POST /tools). The picked + * server's `file_glob_search` server tool (POST /tools). The picked * directory is exposed via `bind:directory`; changing it records a * synthetic "Set working directory to ..." user message into chat history * and is enforced on tool calls via the `x-tool-cwd` request header. @@ -380,7 +380,7 @@ export { default as ChatFormPickerListItemSkeleton } from './ChatForm/ChatFormPi /** * `@`-triggered file/folder mention picker. Resolves `@` in the chat - * input to a filesystem match via the server's `file_glob_search` built-in + * input to a filesystem match via the server's `file_glob_search` server tool * tool, scoped to the conversation cwd (or server home when unset). * Selection splices a `[name](file:///)` link into the input. */ diff --git a/tools/ui/src/lib/components/app/navigation/SidebarNavigation/SidebarNavigation.svelte b/tools/ui/src/lib/components/app/navigation/SidebarNavigation/SidebarNavigation.svelte index 260ff985b..aa63c2915 100644 --- a/tools/ui/src/lib/components/app/navigation/SidebarNavigation/SidebarNavigation.svelte +++ b/tools/ui/src/lib/components/app/navigation/SidebarNavigation/SidebarNavigation.svelte @@ -14,7 +14,7 @@ import { useKeyboardShortcuts } from '$lib/hooks/use-keyboard-shortcuts.svelte'; import { useMarqueeSelection } from '$lib/hooks/use-marquee-selection.svelte'; import { RouterService } from '$lib/services/router.service'; - import { chatStore, conversationsStore, device, isMobile, settingsStore } from '$lib/stores'; + import { chatStore, conversationsStore, deviceStore, settingsStore } from '$lib/stores'; import { buildConversationTree } from '$lib/utils'; import { circIn } from 'svelte/easing'; import { SvelteSet } from 'svelte/reactivity'; @@ -36,7 +36,7 @@ let logoHovered = $state(false); const isStripExpanded = $derived(isExpandedMode || hoveredTooltip !== null); - const isOnMobile = $derived(isMobile.current); + const isOnMobile = $derived(deviceStore.isMobile); const alwaysShowOnDesktop = $derived(settingsStore.config.alwaysShowSidebarOnDesktop as boolean); $effect(() => { @@ -65,7 +65,7 @@ }); $effect(() => { - if (isMobile.current && page.url.hash.includes(ROUTES.SEARCH)) { + if (deviceStore.isMobile && page.url.hash.includes(ROUTES.SEARCH)) { isExpandedMode = false; } }); @@ -227,7 +227,7 @@ } async function selectConversation(id: string) { - if (isMobile.current) { + if (deviceStore.isMobile) { scheduleMobileCollapse(); } @@ -315,9 +315,9 @@ 'fixed md:sticky top-2 left-2 md:left-0 md:ml-2 md:mt-2 pt-2 z-10 w-[calc(100dvw-1rem)]', 'md:h-[calc(100dvh-1.125rem)]', isExpandedMode && - (device.isStandalone + (deviceStore.isStandalone ? 'h-[calc(100dvh-2rem)]' - : device.isIOSDevice + : deviceStore.isIOSDevice ? 'h-[calc(100dvh-0.5rem)]' : 'h-[calc(100dvh-1rem)]'), 'rounded-3xl md:rounded-2xl', @@ -353,7 +353,7 @@ {#if isOnMobile || (isExpandedMode && !alwaysShowOnDesktop)}
@@ -395,7 +395,7 @@ isSearchModeActive = true; }} onNewChat={() => { - if (isMobile.current) { + if (deviceStore.isMobile) { scheduleMobileCollapse(); } }} diff --git a/tools/ui/src/lib/components/app/navigation/SidebarNavigation/SidebarNavigationActions.svelte b/tools/ui/src/lib/components/app/navigation/SidebarNavigation/SidebarNavigationActions.svelte index 29880f0dd..dc333134b 100644 --- a/tools/ui/src/lib/components/app/navigation/SidebarNavigation/SidebarNavigationActions.svelte +++ b/tools/ui/src/lib/components/app/navigation/SidebarNavigation/SidebarNavigationActions.svelte @@ -12,7 +12,7 @@ SIDEBAR_ACTIONS_ITEMS } from '$lib/constants'; import { TooltipSide } from '$lib/enums'; - import { isMobile } from '$lib/stores'; + import { deviceStore } from '$lib/stores'; import type { Component } from 'svelte'; import { onMount } from 'svelte'; import { circIn } from 'svelte/easing'; @@ -42,7 +42,7 @@ let showIcons = $state(false); let searchInputRef = $state(null); - const isOnMobile = $derived(isMobile.current); + const isOnMobile = $derived(deviceStore.isMobile); $effect(() => { if (isSearchModeActive && searchInputRef) { @@ -107,7 +107,7 @@ > {#each SIDEBAR_ACTIONS_ITEMS as item, i (item.tooltip)} {@const isActive = isItemActive(item)} - {@const isSearchOnMobile = item.icon === Search && isMobile.current} + {@const isSearchOnMobile = item.icon === Search && deviceStore.isMobile} {@const itemHref = isSearchOnMobile ? ROUTES.SEARCH : item.route} {@const itemOnClick = item.route ? () => { @@ -156,7 +156,7 @@