diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml index ee1b6e8a5069..98250c860650 100644 --- a/.github/workflows/release.yml +++ b/.github/workflows/release.yml @@ -714,10 +714,10 @@ jobs: with: key: release-windows-2025-vs2026-${{ matrix.arch }}-cpu - # TODO: build only the ggml-hip backend like the other windows backend jobs - # (windows-cuda, windows-sycl), then drop the ui-build dependency + # note: builds only the ggml-hip backend - llama-server is injected from the + # windows-cpu zip during the release "Merge artifacts" step windows-rocm: - needs: [check-release, ui-build] + needs: [check-release] if: ${{ needs.check-release.outputs.should_release == 'true' }} runs-on: windows-2022 @@ -736,11 +736,9 @@ jobs: with: fetch-depth: 0 - - name: Download UI build - uses: actions/download-artifact@v7 - with: - name: llama-ui.zip - path: tools/ui/dist + - name: Install Ninja + run: | + choco install ninja - name: ccache uses: ggml-org/ccache-action@v1.2.21 @@ -795,33 +793,28 @@ jobs: - name: Build run: | - mkdir build - cd build - cmake .. ` - -G "Unix Makefiles" ` + cmake -S . -B build ` + -G "Ninja Multi-Config" ` -DCMAKE_PREFIX_PATH="${env:HIP_PATH}" ` - -DCMAKE_BUILD_TYPE=Release ` -DGGML_BACKEND_DL=ON ` -DGGML_NATIVE=OFF ` - -DGGML_CPU=ON ` - -DGGML_CPU_ALL_VARIANTS=ON ` + -DGGML_CPU=OFF ` -DGGML_HIP=ON ` -DCMAKE_C_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang.exe" ` -DCMAKE_CXX_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang++.exe" ` -DCMAKE_C_FLAGS="-Wno-error=incompatible-pointer-types" ` -DCMAKE_HIP_COMPILER="${env:HIP_PATH}\lib\llvm\bin\clang.exe" ` -DHIP_PATH="${env:HIP_PATH}" ` - -DGGML_HIP_ROCWMMA_FATTN=ON ` -DAMDGPU_TARGETS="${{ matrix.gpu_targets }}" - cmake --build . --config Release --parallel ${env:NUMBER_OF_PROCESSORS} + cmake --build build --config Release --parallel ${env:NUMBER_OF_PROCESSORS} --target ggml-hip - name: Verify HIP backend was built run: | - $hipDll = Get-ChildItem -Path build\bin -Filter "ggml-hip*.dll" -ErrorAction SilentlyContinue + $hipDll = Get-ChildItem -Path build\bin\Release -Filter "ggml-hip*.dll" -ErrorAction SilentlyContinue if (-not $hipDll) { Write-Host "##[error]ggml-hip*.dll was NOT produced. The HIP backend silently failed to build." - Write-Host "Contents of build\bin:" - Get-ChildItem build\bin | Format-Table -AutoSize + Write-Host "Contents of build\bin\Release:" + Get-ChildItem build\bin\Release | Format-Table -AutoSize exit 1 } Write-Host "HIP backend artifact found:" @@ -836,10 +829,40 @@ jobs: $rocmVersionShort = ('${{ matrix.ROCM_VERSION }}'.Split('.')[0..1] -join '.') echo "ROCM_VERSION_SHORT=$rocmVersionShort" >> $env:GITHUB_ENV + - name: Bundle HIP runtime DLLs (amdhip64_7.dll, rocm_kpack.dll, amd_comgr.dll) + run: | + $ErrorActionPreference = "Stop" + # See issue https://github.com/ggml-org/llama.cpp/issues/26929. + # ggml-hip.dll loads amdhip64_7.dll at run time. The Adrenalin driver + # ships an amdhip64_7.dll in System32, which the loader searches before PATH, + # so a matching DLL from PATH cannot win. Copy amdhip64 next to the + # binaries (exe directory is searched before System32) so the correct + # runtime is used. rocm_kpack.dll is amdhip64_7's direct dependency, so + # copy the matching version too. amd_comgr is copied as well to keep it + # in sync with the bundled amdhip64, avoiding a version mismatch with a + # amd_comgr from System32. + # rocblas/hipblaslt kernels resolve fine via PATH and are not copied. + $binPath = (rocm-sdk path --bin).Trim() + if (-not $binPath) { throw "rocm-sdk path --bin returned empty" } + write-host "ROCm bin path: $binPath" + + $patterns = @("amdhip64_7.dll", "rocm_kpack.dll", "amd_comgr.dll") + foreach ($pattern in $patterns) { + $files = Get-ChildItem -Path $binPath -Filter $pattern -ErrorAction SilentlyContinue + if (-not $files) { throw "no match for $pattern in $binPath" } + foreach ($f in $files) { + Copy-Item $f.FullName -Destination build\bin\Release -Force + write-host " copied $($f.Name)" + } + } + - name: Pack artifacts run: | - cp "LICENSE" "build\bin\" - 7z a -snl llama-bin-win-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.zip .\build\bin\* + 7z a -snl llama-bin-win-rocm-${{ env.ROCM_VERSION_SHORT }}-${{ matrix.build }}.zip ` + .\build\bin\Release\ggml-hip.dll ` + .\build\bin\Release\amdhip64_7.dll ` + .\build\bin\Release\rocm_kpack.dll ` + .\build\bin\Release\amd_comgr.dll - name: Upload artifacts uses: actions/upload-artifact@v6 diff --git a/common/arg.cpp b/common/arg.cpp index 86f8610a56d0..e346863e51fd 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -1643,6 +1643,14 @@ common_params_context common_params_parser_init(common_params & params, llama_ex } } ).set_env("LLAMA_ARG_CTX_SIZE")); + add_opt(common_arg( + { "--kv-unified-per-slot" }, "N", + "context limit per parallel slot (default: unset, behavior unchanged).\n" + "when set without -c/--ctx-size, the shared KV pool is sized to n_parallel*N", + [](common_params & params, int value) { + params.kv_unified_per_slot = value; + } + ).set_env("LLAMA_ARG_KV_UNIFIED_PER_SLOT").set_examples({ LLAMA_EXAMPLE_SERVER })); add_opt(common_arg( {"-n", "--predict", "--n-predict"}, "N", string_format( @@ -2644,6 +2652,27 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.mtmd_batch_max_tokens = value; } ).set_examples({LLAMA_EXAMPLE_SERVER}).set_env("LLAMA_ARG_MTMD_BATCH_MAX_TOKENS")); + add_opt(common_arg( + {"--video-fps"}, "N", + string_format("target video frame rate (default: %.1f)", params.video_fps), + [](common_params & params, const std::string & value) { + params.video_fps = std::stof(value); + } + ).set_examples(mmproj_examples).set_env("LLAMA_ARG_VIDEO_FPS")); + add_opt(common_arg( + {"--video-timestamp-interval"}, "N", + string_format("interval in milliseconds between text timestamps (default: %" PRId64 ")", params.video_timestamp_interval_ms), + [](common_params & params, int value) { + params.video_timestamp_interval_ms = value; + } + ).set_examples(mmproj_examples).set_env("LLAMA_ARG_VIDEO_TIMESTAMP_INTERVAL")); + add_opt(common_arg( + {"--video-ffmpeg-dir"}, "DIR", + "path to the directory containing ffmpeg and ffprobe (default: search in PATH)", + [](common_params & params, const std::string & value) { + params.video_ffmpeg_bin_dir = value; + } + ).set_examples(mmproj_examples).set_env("LLAMA_ARG_VIDEO_FFMPEG_DIR")); if (params.is_gen_docs || llama_supports_rpc()) { add_opt(common_arg( {"--rpc"}, "SERVERS", @@ -2699,6 +2728,19 @@ common_params_context common_params_parser_init(common_params & params, llama_ex else { throw std::invalid_argument("invalid value"); } } ).set_env("LLAMA_ARG_LOAD_MODE")); + add_opt(common_arg( + {"--tensor-read-lazy"}, "MODE", + "on-demand reading of certain tensors, for example per-layer embeddings (default: auto)\n" + "- on: read the rows of such tensors from disk on demand instead of keeping them resident (requires mmap)\n" + "- auto: on, but only for tensors larger than 4 GiB\n" + "- off: always keep them resident", + [](common_params & params, const std::string & value) { + /**/ if (value == "on") { params.tensor_read_lazy = LLAMA_TENSOR_READ_LAZY_ON; } + else if (value == "auto") { params.tensor_read_lazy = LLAMA_TENSOR_READ_LAZY_AUTO; } + else if (value == "off") { params.tensor_read_lazy = LLAMA_TENSOR_READ_LAZY_OFF; } + else { throw std::invalid_argument("invalid value"); } + } + ).set_env("LLAMA_ARG_TENSOR_READ_LAZY")); add_opt(common_arg( {"--numa"}, "TYPE", "attempt optimizations that help on some NUMA systems\n" @@ -2750,14 +2792,20 @@ common_params_context common_params_parser_init(common_params & params, llama_ex if (value < 0) { throw std::invalid_argument("invalid value"); } - for (int i = 0; i < value; ++i) { - // keep strings alive and avoid leaking memory by storing them in a static vector - static std::list buft_overrides; - buft_overrides.push_back(llm_ffn_exps_block_regex(i)); - params.tensor_buft_overrides.push_back({buft_overrides.back().c_str(), ggml_backend_cpu_buffer_type()}); - } + llm_add_n_cpu_ffn_overrides(value, LLM_FFN_EXPS_REGEX, params.tensor_buft_overrides); } ).set_env("LLAMA_ARG_N_CPU_MOE")); + add_opt(common_arg( + {"-ncffn", "--n-cpu-ffn"}, "N", + "keep the dense FFN weights of the first N layers in the CPU\n" + "(dense models; for MoE expert weights use --n-cpu-moe)", + [](common_params & params, int value) { + if (value < 0) { + throw std::invalid_argument("invalid value"); + } + llm_add_n_cpu_ffn_overrides(value, LLM_FFN_DENSE_REGEX, params.tensor_buft_overrides); + } + ).set_env("LLAMA_ARG_N_CPU_FFN")); GGML_ASSERT(params.n_gpu_layers < 0); // string_format would need to be extended for a default >= 0 add_opt(common_arg( {"-ngl", "--gpu-layers", "--n-gpu-layers"}, "N", @@ -4084,11 +4132,7 @@ common_params_context common_params_parser_init(common_params & params, llama_ex if (value < 0) { throw std::invalid_argument("invalid value"); } - for (int i = 0; i < value; ++i) { - static std::list buft_overrides_draft; - buft_overrides_draft.push_back(llm_ffn_exps_block_regex(i)); - params.speculative.draft.tensor_buft_overrides.push_back({buft_overrides_draft.back().c_str(), ggml_backend_cpu_buffer_type()}); - } + llm_add_n_cpu_ffn_overrides(value, LLM_FFN_EXPS_REGEX, params.speculative.draft.tensor_buft_overrides); } ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_N_CPU_MOE")); @@ -4109,6 +4153,38 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.speculative.draft.n_min = value; } ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_LOOKUP, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_N_MIN")); + add_opt(common_arg( + {"--spec-synth-len"}, "L", + "target mean synthetic acceptance length, including the target token (benchmarking only)", + [](common_params & params, const std::string & value) { + const std::string text = string_strip(value); + size_t pos = 0; + const double length = std::stod(text, &pos); + if (pos != text.size() || length == -1.0) { + throw std::invalid_argument("invalid value"); + } + params.speculative.synth_len = length; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_SYNTH_LEN")); + add_opt(common_arg( + {"--spec-synth-rates"}, "P0,P1,...", + "comma-separated unconditional per-position synthetic acceptance probabilities (benchmarking only)", + [](common_params & params, const std::string & value) { + const auto values = string_split(value, ','); + std::vector rates; + rates.reserve(values.size()); + for (const auto & raw : values) { + const std::string text = string_strip(raw); + size_t pos = 0; + const double rate = std::stod(text, &pos); + if (pos != text.size()) { + throw std::invalid_argument("invalid value"); + } + rates.push_back(rate); + } + params.speculative.synth_rates = std::move(rates); + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_SYNTH_RATES")); add_opt(common_arg( {"--spec-draft-p-split", "--draft-p-split"}, "P", diff --git a/common/common.cpp b/common/common.cpp index 3d54bd6002d3..347e8e9fc416 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1688,6 +1688,7 @@ struct llama_model_params common_model_params_to_llama(common_params & params) { mparams.main_gpu = params.main_gpu; mparams.split_mode = params.split_mode; mparams.load_mode = params.load_mode; + mparams.tensor_read_lazy = params.tensor_read_lazy; mparams.tensor_split = params.tensor_split; mparams.check_tensors = params.check_tensors; mparams.use_extra_bufts = !params.no_extra_bufts; diff --git a/common/common.h b/common/common.h index de49dac9f63a..a333f702ac1d 100644 --- a/common/common.h +++ b/common/common.h @@ -8,6 +8,7 @@ #include "ggml.h" #include "llama.h" +#include #include #include #include @@ -369,6 +370,9 @@ struct common_params_speculative_ngram_cache { struct common_params_speculative { std::vector types = { COMMON_SPECULATIVE_TYPE_NONE }; + double synth_len = -1.0; + std::vector synth_rates; + // used by Simple, MTP, Eagle3, etc. - all methods that require some kind of draft model common_params_speculative_draft draft; @@ -383,6 +387,10 @@ struct common_params_speculative { return !draft.mparams.empty(); } + bool has_synth() const { + return synth_len != -1.0 || !synth_rates.empty(); + } + uint32_t need_n_rs_seq() const { bool needs_rs_seq = std::any_of(types.begin(), types.end(), [&](auto t) { return t == COMMON_SPECULATIVE_TYPE_DRAFT_MTP || t == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 || t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || t == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK; @@ -475,6 +483,8 @@ struct common_params { enum llama_split_mode split_mode = LLAMA_SPLIT_MODE_LAYER; // how to split the model across GPUs enum llama_load_mode load_mode = LLAMA_LOAD_MODE_AUTO; // how to load the model + enum llama_tensor_read_lazy tensor_read_lazy = LLAMA_TENSOR_READ_LAZY_AUTO; // on-demand reading of tensors marked by the arch + common_cpu_params cpuparams; common_cpu_params cpuparams_batch; @@ -589,6 +599,11 @@ struct common_params { int image_max_tokens = -1; int mtmd_batch_max_tokens = 1024; + // for video input + float video_fps = 4.0f; + int64_t video_timestamp_interval_ms = 5000; + std::string video_ffmpeg_bin_dir = ""; + // finetune struct lr_opt lr; enum ggml_opt_optimizer_type optimizer = GGML_OPT_OPTIMIZER_TYPE_ADAMW; @@ -612,6 +627,7 @@ struct common_params { bool cache_prompt = true; // whether to enable prompt caching bool cache_idle_slots = true; // save and clear idle slots upon starting a new task int32_t n_ctx_checkpoints = 32; // max number of context checkpoints per slot + int32_t kv_unified_per_slot = 0; // max context per parallel slot; 0 = unset int32_t checkpoint_min_step = 8192; // minimum spacing between context checkpoints int32_t cache_ram_mib = 8192; // -1 = no limit, 0 - disable, 1 = 1 MiB, etc. @@ -1108,19 +1124,30 @@ const char * const LLM_KV_SPLIT_TENSORS_COUNT = "split.tensors.count"; } // -// MoE utils +// FFN offload utils // const char * const LLM_FFN_EXPS_REGEX = "\\.ffn_(up|down|gate|gate_up)_(ch|)exps"; -inline std::string llm_ffn_exps_block_regex(int idx) { - return string_format("blk\\.%d%s", idx, LLM_FFN_EXPS_REGEX); +const char * const LLM_FFN_DENSE_REGEX = "\\.ffn_(up|down|gate)\\."; + +inline std::string llm_ffn_block_regex(int idx, const char * ffn_regex) { + return string_format("blk\\.%d%s", idx, ffn_regex); } inline llama_model_tensor_buft_override llm_ffn_exps_cpu_override() { return { LLM_FFN_EXPS_REGEX, ggml_backend_cpu_buffer_type() }; } +inline void llm_add_n_cpu_ffn_overrides(int n, const char * ffn_regex, std::vector & overrides) { + // keep strings alive and avoid leaking memory by storing them in a static list + static std::list buft_override_strings; + for (int i = 0; i < n; ++i) { + buft_override_strings.push_back(llm_ffn_block_regex(i, ffn_regex)); + overrides.push_back({buft_override_strings.back().c_str(), ggml_backend_cpu_buffer_type()}); + } +} + // // training utils // diff --git a/common/speculative.cpp b/common/speculative.cpp index 4eef2212e751..d34d1c9c5950 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -14,6 +14,7 @@ #include #include +#include #include #include #include @@ -138,6 +139,7 @@ struct common_speculative_impl { const common_speculative_type type; uint32_t n_seq; + int32_t n_max; // maximum draft length after implementation-specific limits size_t n_call_begin = 0; // number of times this implementation was called for refresh. size_t n_call_draft = 0; // number of times this implementation was called for generation. @@ -157,7 +159,7 @@ struct common_speculative_impl { int64_t t_draft_us = 0; // total time spent in generating drafts in this implementation in microseconds. int64_t t_accept_us = 0; // total time spent in accumulation of this implementation in microseconds. - common_speculative_impl(common_speculative_type type, uint32_t n_seq) : type(type), n_seq(n_seq) {} + common_speculative_impl(common_speculative_type type, uint32_t n_seq, int32_t n_max) : type(type), n_seq(n_seq), n_max(n_max) {} virtual ~common_speculative_impl() = default; @@ -182,7 +184,7 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { std::vector smpls; common_speculative_impl_draft_simple(const common_params_speculative & params, uint32_t n_seq) - : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE, n_seq) + : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_SIMPLE, n_seq, params.draft.n_max) , params(params.draft) { auto * ctx_dft = this->params.ctx_dft; @@ -452,7 +454,7 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl { std::vector g_embd_buf; common_speculative_impl_draft_eagle3(const common_params_speculative & params, uint32_t n_seq) - : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3, n_seq) + : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3, n_seq, params.draft.n_max) , params(params.draft) { SPC_TRC("%s", "adding speculative implementation 'draft-eagle3'\n"); @@ -923,12 +925,19 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { int32_t block_size = 0; llama_token mask_token_id = 0; + bool is_dflash2 = false; + bool is_mrope = false; + int32_t selector_top_k = 0; + // draft-dspark: the draft carries a Markov head and uses an anchor-first block layout const bool is_dspark; // dspark speculators bool sample_from_anchor = true; + // block-internal attention + bool causal_attn = false; + const int32_t * target_layer_ids = nullptr; // model_dft's extract layer indices uint32_t target_layer_ids_n = 0; @@ -937,7 +946,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { common_speculative_impl_draft_dflash(const common_params_speculative & params, uint32_t n_seq, common_speculative_type type = COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH) - : common_speculative_impl(type, n_seq) + : common_speculative_impl(type, n_seq, params.draft.n_max) , params(params.draft) , is_dspark(type == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK) { @@ -966,9 +975,25 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { if (llama_model_meta_val_str(model_dft, "dflash.sample_from_anchor", buf, sizeof(buf)) >= 0) { sample_from_anchor = std::strcmp(buf, "true") == 0; } + if (llama_model_meta_val_str(model_dft, "dflash.attention.causal", buf, sizeof(buf)) >= 0) { + causal_attn = std::strcmp(buf, "true") == 0; + } } + + selector_top_k = llama_model_dflash_selector_top_k(model_dft); + is_dflash2 = selector_top_k > 0; mask_token_id = llama_vocab_mask(llama_model_get_vocab(model_dft)); + if (is_dspark && this->params.p_min > 0.0f) { + char buf[16] = {}; + const bool has_conf = + llama_model_meta_val_str(model_dft, "dflash.has_confidence_head", buf, sizeof(buf)) < 0 || + std::strcmp(buf, "true") == 0; + if (!has_conf) { + throw std::runtime_error("DSpark draft has no confidence head: please set --spec-draft-p-min 0"); + } + } + LOG_INF("%s: adding speculative implementation '%s'\n", __func__, common_speculative_type_to_str(type).c_str()); LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%.2f\n", __func__, this->params.n_max, this->params.n_min, this->params.p_min); LOG_INF("%s: - block_size=%d, mask_token_id=%d, n_extract=%u, sample_from_anchor=%s\n", __func__, @@ -983,10 +1008,18 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { this->params.n_max = std::min(this->params.n_max, n_draft_max); this->params.n_min = std::min(this->params.n_min, n_draft_max); } + this->n_max = this->params.n_max; batch = llama_batch_init(llama_n_batch(ctx_dft), 0, n_seq); batch_inject = llama_batch_init(llama_n_batch(ctx_dft), n_embd_dec, n_seq); + // embd batches on an M-RoPE draft need 4 position rows per token + is_mrope = llama_model_rope_type(model_dft) == LLAMA_ROPE_TYPE_MROPE; + if (is_mrope) { + free(batch_inject.pos); + batch_inject.pos = (llama_pos *) malloc(sizeof(llama_pos) * 4 * llama_n_batch(ctx_dft)); + } + smpls.resize(n_seq); for (auto & s : smpls) { common_params_sampling sparams; @@ -998,7 +1031,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { // offload draft sampling to the backend backend_chains.assign(n_seq, nullptr); - if (this->params.backend_sampling) { + if (this->params.backend_sampling && !is_dflash2) { for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) { llama_sampler * chain = llama_sampler_chain_init(llama_sampler_chain_default_params()); llama_sampler_chain_add(chain, llama_sampler_init_top_k(10)); @@ -1017,8 +1050,9 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { llama_set_embeddings_layer_inp(ctx_tgt, (uint32_t) target_layer_ids[k], true); } - llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ true); - llama_set_causal_attn(ctx_dft, false); // DFlash needs non-causal attention + // DFlash2 reads its selector lattice from h_nextn and never consumes raw logits. + llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ !is_dflash2); + llama_set_causal_attn(ctx_dft, causal_attn); // DFlash needs non-causal attention unless the model says otherwise } ~common_speculative_impl_draft_dflash() override { @@ -1118,11 +1152,24 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { } // fuse extracted features through DFlash encoder + // M-RoPE drafts read 4 position rows per token from embd batches, so pass them explicitly + std::vector enc_pos; + if (is_mrope) { + enc_pos.resize((size_t) 4 * n_chunk); + for (int32_t i = 0; i < n_chunk; ++i) { + const llama_pos p = batch_in.pos[i_batch_beg[seq_id] + offset + i]; + enc_pos[0 * n_chunk + i] = p; + enc_pos[1 * n_chunk + i] = p; + enc_pos[2 * n_chunk + i] = p; + enc_pos[3 * n_chunk + i] = 0; + } + } + llama_batch enc_batch = { /*.n_tokens =*/ n_chunk, /*.token =*/ nullptr, /*.embd =*/ features_buf.data(), - /*.pos =*/ nullptr, + /*.pos =*/ is_mrope ? enc_pos.data() : nullptr, /*.n_seq_id =*/ nullptr, /*.seq_id =*/ nullptr, /*.logits =*/ nullptr, @@ -1143,7 +1190,13 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { std::memcpy(batch_inject.embd, inp_g, (size_t) n_chunk * n_embd_dec * sizeof(float)); for (int32_t i = 0; i < n_chunk; ++i) { - batch_inject.pos[i] = batch_in.pos[i_batch_beg[seq_id] + offset + i]; + const llama_pos p = batch_in.pos[i_batch_beg[seq_id] + offset + i]; + batch_inject.pos[i] = p; + if (is_mrope) { + batch_inject.pos[1 * n_chunk + i] = p; + batch_inject.pos[2 * n_chunk + i] = p; + batch_inject.pos[3 * n_chunk + i] = 0; + } batch_inject.n_seq_id[i] = 1; batch_inject.seq_id[i][0] = seq_id; batch_inject.logits[i] = false; @@ -1186,7 +1239,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { i_block_beg[seq_id] = batch.n_tokens; n_block [seq_id] = n_block_tokens; for (int32_t i = 0; i < n_block_tokens; ++i) { - common_batch_add(batch, i == 0 ? dp.id_last : mask_token_id, n + i, { seq_id }, true); + common_batch_add(batch, i == 0 ? dp.id_last : mask_token_id, n + i, { seq_id }, !is_dflash2); } } @@ -1214,6 +1267,36 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { auto & result = *dp.result; + if (is_dflash2) { + const float * lattice = llama_get_embeddings_nextn(ctx_dft); + GGML_ASSERT(lattice && "DFlash2 selector produced no lattice"); + + int32_t predecessor = 0; + for (int32_t i = 1; i < n_block_tokens; ++i) { + const float * row = lattice + (size_t) (beg + i) * n_embd_dec; + const float * scores = row + selector_top_k + (size_t) predecessor * selector_top_k; + + predecessor = (int32_t) std::distance(scores, + std::max_element(scores, scores + selector_top_k)); + if (params.p_min > 0.0f) { + // softmax(scores) at the argmax, i.e. 1 / sum(exp(s_k - s_max)) + float sum = 0.0f; + for (int32_t k = 0; k < selector_top_k; ++k) { + sum += std::exp(scores[k] - scores[predecessor]); + } + if (1.0f / sum < params.p_min) { + break; + } + } + result.push_back((llama_token) row[predecessor]); + } + + if (result.size() < (size_t) params.n_min) { + result.clear(); + } + continue; + } + if (is_dspark) { // DSpark: read from the first draft slot, truncate below the confidence threshold const float * conf = params.p_min > 0.0f ? llama_get_embeddings_nextn(ctx_dft) : nullptr; @@ -1315,7 +1398,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { std::vector> chain_h; common_speculative_impl_draft_mtp(const common_params_speculative & params, uint32_t n_seq) - : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_MTP, n_seq) + : common_speculative_impl(COMMON_SPECULATIVE_TYPE_DRAFT_MTP, n_seq, params.draft.n_max) , params(params.draft) { auto * ctx_tgt = this->params.ctx_tgt; @@ -1382,6 +1465,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { c.reserve((size_t) (this->params.n_max + 1) * n_embd); } } + this->n_max = this->params.n_max; pending_h.assign(n_seq, std::vector(n_embd, 0.0f)); @@ -1726,7 +1810,7 @@ struct common_speculative_impl_ngram_simple : public common_speculative_impl { common_speculative_impl_ngram_simple( const common_params_speculative & params, uint32_t n_seq, common_ngram_simple_config config) - : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE, n_seq) + : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_SIMPLE, n_seq, params.ngram_simple.size_m) , params(params.ngram_simple) , config(config) { @@ -1770,7 +1854,7 @@ struct common_speculative_impl_ngram_map_k : public common_speculative_impl { const common_ngram_map & config, uint32_t n_seq) : common_speculative_impl(config.key_only ? COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K - : COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V, n_seq) + : COMMON_SPECULATIVE_TYPE_NGRAM_MAP_K4V, n_seq, config.size_value) { for (uint32_t i = 0; i < n_seq; i++) { this->config.push_back(config); @@ -1841,7 +1925,7 @@ struct common_speculative_impl_ngram_mod : public common_speculative_impl { common_speculative_impl_ngram_mod( const common_params_speculative & params, uint32_t n_seq) - : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_MOD, n_seq) + : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_MOD, n_seq, params.ngram_mod.n_max) , params(params.ngram_mod) , mod(params.ngram_mod.n_match, 4*1024*1024) , verbose(std::getenv("LLAMA_TRACE") != nullptr) { @@ -2017,7 +2101,7 @@ struct common_speculative_impl_ngram_cache : public common_speculative_impl { const std::string & path_dynamic, bool save_dynamic, bool save_static) - : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_CACHE, n_seq) + : common_speculative_impl(COMMON_SPECULATIVE_TYPE_NGRAM_CACHE, n_seq, n_draft) , params(params.ngram_cache) , n_draft(n_draft) , save_dynamic(save_dynamic) @@ -2138,6 +2222,8 @@ struct common_speculative { // which implementaion was used for a given seq_id std::vector impl_last; + + std::vector synth_probs; }; static common_ngram_map get_common_ngram_map( @@ -2316,6 +2402,101 @@ int32_t common_speculative_n_max(const common_params_speculative * spec) { return n_max; } +int32_t common_speculative_n_max(const common_speculative * spec) { + int32_t n_max = 0; + + if (spec == nullptr) { + return n_max; + } + + for (const auto & impl : spec->impls) { + n_max = std::max(n_max, std::max(0, impl->n_max)); + } + + return n_max; +} + +std::vector common_speculative_synth_rates_resolve(const common_params_speculative * spec, int32_t n_max) { + const bool has_length = spec->synth_len != -1.0; + const bool has_rates = !spec->synth_rates.empty(); + + if (!has_length && !has_rates) { + return {}; + } + if (has_length && has_rates) { + throw std::invalid_argument("synthetic acceptance length and rates are mutually exclusive"); + } + + if (n_max <= 0) { + throw std::invalid_argument("synthetic acceptance requires at least one speculative token"); + } + + if (has_rates) { + const auto & rates = spec->synth_rates; + if (rates.size() != (size_t) n_max) { + throw std::invalid_argument(string_format( + "synthetic acceptance rates must contain %d values, got %zu", n_max, rates.size())); + } + + for (size_t i = 0; i < rates.size(); ++i) { + if (!std::isfinite(rates[i]) || rates[i] < 0.0 || rates[i] > 1.0) { + throw std::invalid_argument("synthetic acceptance rates must be finite and within [0, 1]"); + } + if (i > 0 && rates[i] > rates[i - 1]) { + throw std::invalid_argument("synthetic acceptance rates must be monotonically non-increasing"); + } + } + + return rates; + } + + const double length = spec->synth_len; + const double length_max = (double) n_max + 1.0; + if (!std::isfinite(length) || length < 1.0 || length > length_max) { + throw std::invalid_argument(string_format( + "synthetic acceptance length must be finite and within [1, %.0f]", length_max)); + } + + double p = 0.0; + if (length == length_max) { + p = 1.0; + } else if (length > 1.0) { + double p_min = 0.0; + double p_max = 1.0; + for (int i = 0; i < 32; ++i) { + const double p_mid = 0.5 * (p_min + p_max); + double sum = 0.0; + double term = p_mid; + for (int32_t j = 0; j < n_max; ++j) { + sum += term; + term *= p_mid; + } + + if (sum < length - 1.0) { + p_min = p_mid; + } else { + p_max = p_mid; + } + } + p = 0.5 * (p_min + p_max); + } + + std::vector rates; + rates.reserve(n_max); + double rate = p; + for (int32_t i = 0; i < n_max; ++i) { + rates.push_back(rate); + rate *= p; + } + + return rates; +} + +const std::vector & common_speculative_get_synth_probs(const common_speculative * spec) { + GGML_ASSERT(spec); + return spec->synth_probs; +} + common_params common_base_params_to_speculative(const common_params & params) { const bool has_draft = params.speculative.has_dft(); @@ -2568,13 +2749,39 @@ common_speculative * common_speculative_init(common_params_speculative & params, return nullptr; } - auto * result = new common_speculative { - /* .dparams = */ common_speculative_draft_params_vec(n_seq), - /* .impls = */ std::move(impls), - /* .impl_last = */ std::vector(n_seq, nullptr) - }; + common_speculative_ptr result(new common_speculative { + /* .dparams = */ common_speculative_draft_params_vec(n_seq), + /* .impls = */ std::move(impls), + /* .impl_last = */ std::vector(n_seq, nullptr), + /* .synth_probs = */ {}, + }); - return result; + const int32_t n_max_configured = common_speculative_n_max(¶ms); + const int32_t n_max_effective = common_speculative_n_max(result.get()); + const auto rates = common_speculative_synth_rates_resolve(¶ms, n_max_effective); + + std::vector rates_str; + rates_str.reserve(rates.size()); + result->synth_probs.reserve(rates.size()); + double rate_prev = 1.0; + double acceptance_length = 1.0; + for (const double rate : rates) { + result->synth_probs.push_back(rate_prev > 0.0 ? rate / rate_prev : 0.0); + rates_str.push_back(string_format("%.6g", rate)); + rate_prev = rate; + acceptance_length += rate; + } + if (!result->synth_probs.empty()) { + SPC_WRN("%s", "synthetic speculative acceptance is enabled for benchmarking; generated output is not valid\n"); + if (n_max_effective != n_max_configured) { + SPC_WRN("synthetic acceptance draft limit was reduced from %d to %d by the initialized speculative implementations\n", + n_max_configured, n_max_effective); + } + SPC_INF("synthetic acceptance: n_max = %zu, mean length = %.6f, rates = [%s]\n", + rates.size(), acceptance_length, string_join(rates_str, ", ").c_str()); + } + + return result.release(); } void common_speculative_free(common_speculative * spec) { diff --git a/common/speculative.h b/common/speculative.h index 12ae31b7de59..22505891f7ef 100644 --- a/common/speculative.h +++ b/common/speculative.h @@ -26,6 +26,15 @@ std::string common_speculative_type_to_str(enum common_speculative_type type); // return the max number of draft tokens based on the speculative parameters int32_t common_speculative_n_max(const common_params_speculative * spec); +// return the max number of draft tokens from the initialized implementations +int32_t common_speculative_n_max(const common_speculative * spec); + +// validate and resolve the unconditional synthetic acceptance rates +std::vector common_speculative_synth_rates_resolve(const common_params_speculative * spec, int32_t n_max); + +// return the conditional synthetic acceptance probabilities +const std::vector & common_speculative_get_synth_probs(const common_speculative * spec); + common_params common_base_params_to_speculative(const common_params & params); struct common_speculative_output_limits { diff --git a/conversion/__init__.py b/conversion/__init__.py index d80c787f761b..ea521a580d90 100644 --- a/conversion/__init__.py +++ b/conversion/__init__.py @@ -54,6 +54,7 @@ "DeepseekV3ForCausalLM": "deepseek", "DeepseekV32ForCausalLM": "deepseek", "DFlashDraftModel": "qwen", + "DFlash2DraftModel": "qwen", "Qwen3DSparkModel": "qwen", "DSparkDraftModel": "qwen", "DSparkSpeculator": "qwen", @@ -236,6 +237,8 @@ "Qwen3_5ForConditionalGeneration": "qwen", "Qwen3_5MoeForCausalLM": "qwen", "Qwen3_5MoeForConditionalGeneration": "qwen", + "Qwen4ExpForCausalLM": "qwen4exp", + "Qwen4ExpForConditionalGeneration": "qwen4exp", "RND1": "qwen", "RWForCausalLM": "falcon", "RWKV6Qwen2ForCausalLM": "rwkv", @@ -334,6 +337,7 @@ "Qwen3VLMoeForConditionalGeneration": "qwen3vl", "Qwen3_5ForConditionalGeneration": "qwen3vl", "Qwen3_5MoeForConditionalGeneration": "qwen3vl", + "Qwen4ExpForConditionalGeneration": "qwen4exp", "RADIOModel": "nemotron", "Sarashina2VisionForCausalLM": "sarashina2", "SmolVLMForConditionalGeneration": "smolvlm", diff --git a/conversion/base.py b/conversion/base.py index 39ccdc414277..d0257c2b76c4 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -1029,12 +1029,16 @@ def load(): else: raise ValueError(f"Unknown file type: {self.ftype.name}") + # a chunked tensor quantizes as one chunk at a time, while it is written + quantize = data.quantize if isinstance(data, gguf.LazyChunkedTensor) else ( + lambda qtype, d=data: gguf.quants.quantize(d, qtype)) + try: - data = gguf.quants.quantize(data, data_qtype) + data = quantize(data_qtype) except gguf.QuantError as e: logger.warning("%s, %s", e, "falling back to F16") data_qtype = gguf.GGMLQuantizationType.F16 - data = gguf.quants.quantize(data, data_qtype) + data = quantize(data_qtype) shape = gguf.quant_shape_from_byte_shape(data.shape, data_qtype) if data.dtype == np.uint8 else data.shape diff --git a/conversion/nemotron.py b/conversion/nemotron.py index cd0d48c8f0f7..07fbc6531495 100644 --- a/conversion/nemotron.py +++ b/conversion/nemotron.py @@ -302,6 +302,10 @@ def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Ca ) if not keep: return None + # PEFT names adapter tensors using model.layers.*, while Nemotron-H checkpoints + # and the GGUF tensor map use backbone.layers.* + if name.startswith("model.layers.") and ".mixer." in name: + name = name.replace("model.layers.", "backbone.layers.", 1) return super().filter_tensors((name, gen)) def prepare_metadata(self, vocab_only: bool): diff --git a/conversion/qwen.py b/conversion/qwen.py index cdba8a63e9c9..419611896fc3 100644 --- a/conversion/qwen.py +++ b/conversion/qwen.py @@ -639,7 +639,7 @@ class Qwen3_5MoeTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase): model_arch = gguf.MODEL_ARCH.QWEN35MOE -@ModelBase.register("DFlashDraftModel") +@ModelBase.register("DFlashDraftModel", "DFlash2DraftModel") @ModelBase.example("z-lab/Qwen3.5-9B-DFlash") class DFlashModel(Qwen3Model): model_arch = gguf.MODEL_ARCH.DFLASH @@ -678,34 +678,98 @@ def set_vocab(self): def set_gguf_parameters(self): super().set_gguf_parameters() - block_size = self.hparams.get("block_size", 16) - self.gguf_writer.add_block_size(block_size) dflash_config = self.hparams.get("dflash_config", {}) + block_size = dflash_config.get("block_size", self.hparams.get("block_size", 16)) + self.gguf_writer.add_block_size(block_size) + + if "conv_kernel_size" in dflash_config: + self.gguf_writer.add_conv_kernel_size(int(dflash_config["conv_kernel_size"])) + self.gguf_writer.add_conv_group_size(int(dflash_config["conv_group_size"])) + self.gguf_writer.add_selector_rank(int(dflash_config["selector_rank"])) + self.gguf_writer.add_selector_top_k(int(dflash_config["selector_top_k"])) + + output_multiplier = dflash_config.get( + "output_multiplier", self.hparams.get("output_multiplier") + ) + if output_multiplier is not None: + self.gguf_writer.add_logit_scale(float(output_multiplier)) + softcap = dflash_config.get( + "final_logit_softcapping", self.hparams.get("final_logit_softcapping") + ) + if softcap is not None and float(softcap) > 0: + self.gguf_writer.add_final_logit_softcapping(float(softcap)) + embedding_scale = dflash_config.get( + "input_embedding_scale", self.hparams.get("input_embedding_scale") + ) + if embedding_scale is not None: + self.gguf_writer.add_embedding_scale(float(embedding_scale)) target_layer_ids = dflash_config.get("target_layer_ids", []) if target_layer_ids: extract_layer_ids = [i + 1 for i in target_layer_ids] self.gguf_writer.add_target_layers(extract_layer_ids) - use_sliding_window = self.hparams.get("use_sliding_window", False) - sliding_window = self.hparams.get("sliding_window") + use_sliding_window = self.hparams.get("use_sliding_window", False) or dflash_config.get("use_swa", False) + sliding_window = dflash_config.get("swa_window_size") or self.hparams.get("sliding_window") layer_types = self.hparams.get("layer_types") if use_sliding_window and sliding_window and layer_types: is_swa = [lt == "sliding_attention" for lt in layer_types] self.gguf_writer.add_sliding_window(sliding_window) self.gguf_writer.add_sliding_window_pattern(is_swa) + causal = self.hparams.get("is_causal") + if causal is None: + causal = dflash_config.get("causal") + if causal is not None: + self.gguf_writer.add_causal_attention(bool(causal)) + + # M-RoPE target: the draft ropes on the temporal dim only, so write + # degenerate sections [n_rot/2, 0, 0, 0] + if self._target_uses_mrope(): + head_dim = self.hparams.get("head_dim") or self.hparams["hidden_size"] // self.hparams["num_attention_heads"] + self.gguf_writer.add_rope_dimension_sections([head_dim // 2, 0, 0, 0]) + + def _target_uses_mrope(self) -> bool: + if self.target_model_dir is None: + return False + with open(self.target_model_dir / "config.json", "r", encoding="utf-8") as f: + cfg = json.load(f) + cfg = cfg.get("text_config", cfg) + rope = cfg.get("rope_parameters") or cfg.get("rope_scaling") or {} + return "mrope_section" in rope + @classmethod def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: name, gen = item if not name.startswith("model."): name = "model." + name + if "sink" in name and not name.endswith(".weight"): + name += ".weight" return super().filter_tensors((name, gen)) + _ROPE_PERMUTE_SUFFIXES = ( + "self_attn.q_proj.weight", + "self_attn.k_proj.weight", + "self_attn.q_norm.weight", + "self_attn.k_norm.weight", + ) + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: if name == "model.embed_tokens.weight" and not self.hparams.get("has_embed_tokens", True): return + # interleaved-rope checkpoints (rope_is_neox_style = false) -> NeoX layout: per head, even dims first then odd + if not self.hparams.get("rope_is_neox_style", True) and name.endswith(self._ROPE_PERMUTE_SUFFIXES): + head_dim = self.hparams["head_dim"] + shape = data_torch.shape + data_torch = data_torch.reshape(-1, head_dim // 2, 2, *shape[1:]).transpose(1, 2).reshape(shape) + + if name in ( + "model.candidate_selector.predecessor_codebook", + "model.candidate_selector.successor_codebook", + ): + name += ".weight" + yield from super().modify_tensors(data_torch, name, bid) @@ -759,6 +823,10 @@ def set_gguf_parameters(self): super().set_gguf_parameters() self.gguf_writer.add_sample_from_anchor(self._sample_from_anchor) + # confidence head is optional: vanilla-markov exports ship without it + has_conf = any("confidence_head.proj" in name for name in self.model_tensors) + self.gguf_writer.add_has_confidence_head(has_conf) + @classmethod def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: if item[0] == "t2d": # not used at runtime @@ -777,7 +845,7 @@ def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iter self._d2t = data_torch return - if self._n_vocab_draft == self.hparams["vocab_size"] and name.endswith(("embed_tokens.weight", "lm_head.weight")): + if self._n_vocab_draft == self.hparams["vocab_size"] and name.endswith("lm_head.weight"): return # interleaved-rope checkpoints (rope_is_neox_style = false) -> NeoX layout: per head, even dims first then odd diff --git a/conversion/qwen4exp.py b/conversion/qwen4exp.py new file mode 100644 index 000000000000..168796d616b9 --- /dev/null +++ b/conversion/qwen4exp.py @@ -0,0 +1,195 @@ +from __future__ import annotations + +from typing import Iterable, cast + +import torch +from torch import Tensor + +import gguf +import numpy as np + +from .base import ModelBase +from .qwen import _LinearAttentionVReorderBase, _Qwen35MRopeMixin +from .qwen3vl import Qwen3VLVisionModel + + +@ModelBase.register("Qwen4ExpForConditionalGeneration", "Qwen4ExpForCausalLM") +@ModelBase.example("Qwen/Qwen3.8-Flash-Next") +class Qwen4ExpTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase): + """Qwen3.8-Flash-Next. + + Shares the Qwen3.5 gated delta net and interleaved mrope, and adds three things: + hyper-connections in place of every layer norm, QSA sparse attention on the full + attention layers, and PLE n-gram hash embeddings on a single layer. + """ + + model_arch = gguf.MODEL_ARCH.QWEN4EXP + + # the MTP block is a separate draft head; vLLM drops it too + supports_mtp_export = False + no_mtp = True + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + # only the shard names, so the table itself is never held + self._ple_shards: dict[int, str] = {} + self._ple_row_dim: int | None = None + + def _read_hash_constants(self, suffix: str) -> list[int]: + """Read an int64 PLE constant straight from the checkpoint. + + prepare_tensors() casts every non-float dtype to float32 before + modify_tensors() sees it (base.py), which would silently round these + 45-bit multipliers. Reading the lazy tensor here bypasses that. + """ + for name, gen in self.model_tensors.items(): + if name.endswith(suffix): + t = gen() + if t.dtype != torch.int64: + t = t.to(torch.int64) + return [int(x) for x in t.tolist()] + raise ValueError(f"PLE constant {suffix!r} missing from the checkpoint") + + def set_gguf_parameters(self): + super().set_gguf_parameters() + hp = self.hparams + + self.gguf_writer.add_hyper_connection_count(hp["hc_count"]) + self.gguf_writer.add_hyper_connection_low_rank(hp["hc_lowrank"]) + + n_layer = hp["num_hidden_layers"] + self.gguf_writer.add_indexer_head_count(hp["indexer_n_heads"]) + self.gguf_writer.add_indexer_key_length(hp["indexer_head_dim"]) + self.gguf_writer.add_indexer_top_k(hp["indexer_budget"]) + ratio = hp["indexer_compress_ratio"] + layer_types = hp["layer_types"] + self.gguf_writer.add_attention_compress_ratios( + [ratio if layer_types[i] == "full_attention" else 0 for i in range(n_layer)] + ) + + # ple_layer_ids is 1-based in the HF config; empty means no n-gram table, + # so emit no PLE keys rather than optional ones + ple_layers = [i - 1 for i in hp["ple_layer_ids"]] + if not ple_layers: + return + self.gguf_writer.add_ple_layers(ple_layers) + self.gguf_writer.add_ple_ngram_size(hp["ngram_size"]) + self.gguf_writer.add_ple_heads_per_ngram(hp["heads_per_ngram"]) + self.gguf_writer.add_ple_conv_kernel(hp["ple_conv_kernel_size"]) + self.gguf_writer.add_ple_eos_token_id(self._eos_token_id()) + # an image is decoded as an embeddings-only batch, so the graph has no placeholder + # ids to hash; carry the id and let it stand in for those positions + _img = self._image_token_id() + if _img is not None: + self.gguf_writer.add_ple_image_token_id(int(_img)) + if self._ple_row_dim is not None: + self.gguf_writer.add_embedding_length_per_layer_input(self._ple_row_dim) + + self.gguf_writer.add_ple_layer_multipliers( + self._read_hash_constants("ple_embedding.layer_multipliers")) + self.gguf_writer.add_ple_head_offsets( + self._read_hash_constants("ple_embedding.ngram_heads_offsets")) + self.gguf_writer.add_ple_head_vocab_sizes( + self._read_hash_constants("ple_embedding.ngram_heads_vocab_sizes")) + + def _image_token_id(self) -> int | None: + img = self.hparams.get("image_token_id") + return None if img is None else int(img) + + def _eos_token_id(self) -> int: + eos = self.hparams.get("eos_token_id") + if isinstance(eos, list): + # the PLE hash resets n-grams on the primary EOS + return int(eos[-1]) + if eos is None: + raise ValueError("eos_token_id is required: the PLE hash resets its n-grams on it") + return int(eos) + + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: + # int64 hash constants must stay exact; 1-D tensors force F32, so use KV + if name.endswith("ple_embedding.layer_multipliers"): + self._ple_multipliers = [int(x) for x in data_torch.tolist()] + return [] + if name.endswith("ple_embedding.ngram_heads_offsets"): + self._ple_head_offsets = [int(x) for x in data_torch.tolist()] + return [] + if name.endswith("ple_embedding.ngram_heads_vocab_sizes"): + self._ple_head_vocab_sizes = [int(x) for x in data_torch.tolist()] + return [] + + if ".ngram_embedding.shard_" in name: + return self._place_ple_shard(data_torch, name) + + # one projection feeds indexer q and k; split it, as minimax-m3 does + if ".indexer.index_qk_proj.weight" in name: + n_q = self.hparams["indexer_n_heads"] * self.hparams["indexer_head_dim"] + q = data_torch[:n_q] + k = data_torch[n_q:] + return [ + (self.format_tensor_name(gguf.MODEL_TENSOR.INDEXER_Q_PROJ, bid, ".weight"), q), + (self.format_tensor_name(gguf.MODEL_TENSOR.INDEXER_K_PROJ, bid, ".weight"), k), + ] + + # Gemma zero-centred gammas the inherited norm.weight rule misses + if name.endswith((".ple.norm_key.weight", ".ple.norm_query.weight", ".ple.norm_conv.weight", + ".indexer.q_layernorm.weight", ".indexer.k_layernorm.weight")): + return [(self.map_tensor_name(name), data_torch + 1)] + + if name.endswith(".ple.conv1d.weight"): + return [(self.map_tensor_name(name), data_torch.squeeze())] + + return super().modify_tensors(data_torch, name, bid) + + # the shards concatenate into a tensor of well over 100 GB + # use LazyChunkedTensor here, a single shard resident at a time + def _place_ple_shard(self, data_torch: Tensor, name: str) -> Iterable[tuple[str, Tensor]]: + + idx = int(name.rpartition(".shard_")[2].partition(".")[0]) + n_parts = self.hparams["split_ngram_parts"] + + self._ple_shards[idx] = name + self._ple_row_dim = int(data_torch.shape[-1]) + + if len(self._ple_shards) < n_parts: + return [] + + # the checkpoint may yield the shards in any order, the row order is by index + shards = [self._ple_shards[i] for i in sorted(self._ple_shards)] + rows = 0 + for shard in shards: + shape = self.model_tensors[shard]().shape + if int(shape[-1]) != self._ple_row_dim: + raise ValueError( + f"PLE shard {shard} has row dim {int(shape[-1])}, expected {self._ple_row_dim}") + rows += int(shape[0]) + + table = gguf.LazyChunkedTensor( + [self._load_ple_shard(shard) for shard in shards], + shape=(rows, self._ple_row_dim), + dtype=np.float32, + ) + gguf_name = gguf.TENSOR_NAMES[gguf.MODEL_TENSOR.PER_LAYER_TOKEN_EMBD] + return [(gguf_name + ".weight", cast(Tensor, table))] + + def _load_ple_shard(self, name: str): + def load() -> np.ndarray: + from .base import LazyTorchTensor + + # a fresh lazy tensor every call, or to_eager() memoizes every shard + eager = LazyTorchTensor.to_eager(self.model_tensors[name]()) + return eager.to(torch.float32).contiguous().numpy() + return load + + def prepare_tensors(self): + super().prepare_tensors() + n_parts = self.hparams.get("split_ngram_parts", 0) + if self._ple_shards and len(self._ple_shards) != n_parts: + raise ValueError( + f"got {len(self._ple_shards)} PLE embedding shards, expected {n_parts}" + ) + + +@ModelBase.register("Qwen4ExpForConditionalGeneration") +@ModelBase.example("Qwen/Qwen3.8-Flash-Next") +class Qwen4ExpVisionModel(Qwen3VLVisionModel): + """The vision tower is an unmodified Qwen3-VL ViT.""" diff --git a/docs/backend/snapdragon/CMakeUserPresets.json b/docs/backend/snapdragon/CMakeUserPresets.json index 848d735f1c5c..afc73923237a 100644 --- a/docs/backend/snapdragon/CMakeUserPresets.json +++ b/docs/backend/snapdragon/CMakeUserPresets.json @@ -8,7 +8,7 @@ "toolset": { "value": "host=x86_64", "strategy": "external" }, "cacheVariables": { "ANDROID_ABI": "arm64-v8a", - "ANDROID_PLATFORM": "android-31", + "ANDROID_PLATFORM": "android-34", "CMAKE_TOOLCHAIN_FILE": "$env{ANDROID_NDK_ROOT}/build/cmake/android.toolchain.cmake", "CMAKE_C_FLAGS": "-march=armv8.7a+fp16+dotprod+i8mm -fvectorize -ffp-model=fast -fno-finite-math-only -flto -D_GNU_SOURCE", "CMAKE_CXX_FLAGS": "-march=armv8.7a+fp16+dotprod+i8mm -fvectorize -ffp-model=fast -fno-finite-math-only -flto -D_GNU_SOURCE", diff --git a/docs/backend/snapdragon/README.md b/docs/backend/snapdragon/README.md index e9f0e215858c..391c8bf230f0 100644 --- a/docs/backend/snapdragon/README.md +++ b/docs/backend/snapdragon/README.md @@ -2,39 +2,47 @@ ## Setup -### Android +The cross-compilation toolchain images are provided by the +[Qualcomm Snapdragon Toolchain registry](https://github.com/snapdragon-toolchain). +These Docker images include the Android NDK, OpenCL SDK, Hexagon SDK, CMake, and the necessary cross-compilers: -The easiest way to build llama.cpp for a Snapdragon-based Android device is using the toolchain Docker image (see github.com/snapdragon-toolchain). -This image includes Android NDK, OpenCL SDK, Hexagon SDK, CMake, etc. +* **Android toolchain**: `ghcr.io/snapdragon-toolchain/arm64-android:v0.7` +* **Linux toolchain**: `ghcr.io/snapdragon-toolchain/arm64-linux:v0.7` -This method works on Linux, macOS, and Windows. macOS and Windows users should install Docker Desktop. +The unified build utility (`scripts/snapdragon/build.py`) automatically pulls +and orchestrates these containers to perform target compilation. +You only need to ensure that Docker (or Docker Desktop on macOS/Windows) is running on your host machine. +Specific setup, build, and installation details for Linux and Windows on Snapdragon platforms are documented in: +* [Linux on Snapdragon guide](linux.md) +* [Windows on Snapdragon guide](windows.md) -``` -~/src/llama.cpp$ docker run -it -u $(id -u):$(id -g) --volume $(pwd):/workspace --platform linux/amd64 ghcr.io/snapdragon-toolchain/arm64-android:v0.7 -[d]/> cd /workspace -``` +## How to Build + +### Using build.py script (Recommended) -Note: The rest of the **Android** build process assumes that you're running inside the toolchain container. +The easiest way to build llama.cpp is by using the `scripts/snapdragon/build.py` script. It automatically copies the CMake presets, +launches the correct compilation Docker container, builds the libraries and tools, +installs them, and optionally pushes them to your ADB device. -### Windows On Snapdragon +Build and deploy for Android target (accepts `android` or `adb` alias): +``` +$ ./scripts/snapdragon/build.py --target adb --push +``` -Native Windows 11 arm64 builds has the following tools dependencies: -- MS Visual Studio 2026 (Community Edition or Pro) - - MSVC arm64 standard and runtime libraries - - UCRT and Driver Kit -- LLVM core libraries and Clang compiler (winget) -- CMake, Git, Python (winget) -- Hexagon SDK Community Edition 6.6 or later (see windows.md) -- OpenCL SDK 2.3 or later (see windows.md) +Build and deploy for Linux target (accepts `linux` or `lnx` alias): +``` +$ ./scripts/snapdragon/build.py --target linux:user@host --push +``` -Note: The rest of the **Windows** build process assumes that you're running natively in Powershell. -Adapt below build commands accordingly. +### Manual CMake Build -## How to Build +Alternatively, you can build llama.cpp manually by entering the cross-compilation Docker container and running the CMake commands: -Let's build llama.cpp with CPU, OpenCL, and Hexagon backends via CMake presets: +```bash +# Start the cross-compilation container manually: +~/src/llama.cpp$ docker run -it --rm -u $(id -u):$(id -g) --volume $(pwd):/workspace --platform linux/amd64 ghcr.io/snapdragon-toolchain/arm64-android:v0.7 -``` +# Inside the container, build the project using presets: [d]/workspace> cp docs/backend/snapdragon/CMakeUserPresets.json . [d]/workspace> cmake --preset arm64-android-snapdragon-release -B build-snapdragon @@ -68,19 +76,19 @@ Preset CMake variables: To generate an installable "package" simply use cmake --install: ``` -[d]/workspace> cmake --install build-snapdragon --prefix pkg-snapdragon/llama.cpp +[d]/workspace> cmake --install build-snapdragon --prefix pkg-android/llama.cpp -- Install configuration: "Release" --- Installing: /workspace/pkg-snapdragon/llama.cpp/lib/libggml-cpu.so --- Installing: /workspace/pkg-snapdragon/llama.cpp/lib/libggml-opencl.so --- Installing: /workspace/pkg-snapdragon/llama.cpp/lib/libggml-hexagon.so --- Installing: /workspace/pkg-snapdragon/llama.cpp/lib/libggml-htp-v73.so --- Installing: /workspace/pkg-snapdragon/llama.cpp/lib/libggml-htp-v75.so --- Installing: /workspace/pkg-snapdragon/llama.cpp/lib/libggml-htp-v79.so --- Installing: /workspace/pkg-snapdragon/llama.cpp/lib/libggml-htp-v81.so --- Installing: /workspace/pkg-snapdragon/llama.cpp/lib/libggml.so +-- Installing: /workspace/pkg-android/llama.cpp/lib/libggml-cpu.so +-- Installing: /workspace/pkg-android/llama.cpp/lib/libggml-opencl.so +-- Installing: /workspace/pkg-android/llama.cpp/lib/libggml-hexagon.so +-- Installing: /workspace/pkg-android/llama.cpp/lib/libggml-htp-v73.so +-- Installing: /workspace/pkg-android/llama.cpp/lib/libggml-htp-v75.so +-- Installing: /workspace/pkg-android/llama.cpp/lib/libggml-htp-v79.so +-- Installing: /workspace/pkg-android/llama.cpp/lib/libggml-htp-v81.so +-- Installing: /workspace/pkg-android/llama.cpp/lib/libggml.so ... --- Installing: /workspace/pkg-snapdragon/llama.cpp/bin/llama-bench --- Installing: /workspace/pkg-snapdragon/llama.cpp/bin/llama-cli +-- Installing: /workspace/pkg-android/llama.cpp/bin/llama-bench +-- Installing: /workspace/pkg-android/llama.cpp/bin/llama-cli ... ``` @@ -91,14 +99,14 @@ To generate an installable "package" simply use cmake --install: For this step, your device needs to be configured for on-device development. Please see https://developer.android.com/studio/debug/dev-options for details. -Once ADB is enabled, use `adb push` to install `pkg-snapdragon` on the device. +Once ADB is enabled, use `adb push` to install `pkg-android` on the device. **Note that the toolchain Docker image doesn't have ADB and doesn't set up the ADB bridge. Please use native ADB on the host.** ``` -~/src/llama.cpp$ adb push pkg-snapdragon/llama.cpp /data/local/tmp/ -pkg-snapdragon/llama.cpp/bin/: 67 files pushed, 0 skipped. 190.2 MB/s (919095042 bytes in 4.607s) -pkg-snapdragon/llama.cpp/include/: 19 files pushed, 0 skipped. 20.5 MB/s (255173 bytes in 0.012s) -pkg-snapdragon/llama.cpp/lib/: 16 files pushed, 0 skipped. 144.4 MB/s (43801382 bytes in 0.289s) +~/src/llama.cpp$ adb push pkg-android/llama.cpp /data/local/tmp/ +pkg-android/llama.cpp/bin/: 67 files pushed, 0 skipped. 190.2 MB/s (919095042 bytes in 4.607s) +pkg-android/llama.cpp/include/: 19 files pushed, 0 skipped. 20.5 MB/s (255173 bytes in 0.012s) +pkg-android/llama.cpp/lib/: 16 files pushed, 0 skipped. 144.4 MB/s (43801382 bytes in 0.289s) 102 files pushed, 0 skipped. 186.9 MB/s (963151597 bytes in 4.914s) ``` @@ -115,24 +123,44 @@ Llama-3.2-1B-Instruct-Q4_0.gguf: 1 file pushed, 0 skipped. 38.3 MB/s (773025920 ### Windows -All artifacts are already installed in the `pkg-snapdragon` folder. -To run, adapt below instructions to use Powershell scripts in `scripts/snapdragon/windows`. +All artifacts are already installed in the `pkg-wos` folder. +To run, you can use the `scripts/snapdragon/run.py` runner script (see details below). ## How to Run -The easiest way to run llama.cpp cli tools is using provided wrapper scripts that properly set up all required environment variables. +The easiest way to run llama.cpp cli tools is using the provided `scripts/snapdragon/run.py` wrapper script. This script automatically +maps CLI options to environment variables, resolves executable paths, and runs the command locally, via ADB, or remotely via SSH on the +target device. -llama.cpp supports three backends on Snapdragon-based devices: CPU, Adreno GPU (GPUOpenCL), and Hexagon NPU (HTP0-4). -You can select which backend to run the model on using the `D=` variable, which maps to the `--device` option. +llama.cpp supports three backends on Snapdragon-based devices: CPU, Adreno GPU (GPUOpenCL), and Hexagon NPU. +You can select which backend(s) to run the model on using the `--device` option of the tool (or `--devices` option in `run.py`). Hexagon NPU behaves as a "GPU" device when it comes to `-ngl` and other offload-related options. -Here are some examples of running various llama.cpp tools via ADB. +Here are some examples of running various llama.cpp tools. -Simple question for Llama-3.2-1B +Generating a completion with Gemma on Android (relying on default `HTP0:0` device and default thread count `-t 6`): ``` -~/src/llama.cpp$ M=Llama-3.2-1B-Instruct-Q4_0.gguf D=HTP0 ./scripts/snapdragon/adb/run-completion.sh -p "what is the most popular cookie in the world?" +~/src/llama.cpp$ ./scripts/snapdragon/run.py --target adb -- llama-completion -m models/gemma-2-2b-it-Q4_0.gguf -f prompts/sample_prompt_1024.txt --jinja -st +... +ggml-hex: Hexagon backend (experimental) : allocating new registry : ndev 1 +ggml-hex: Hexagon Arch version v79 +ggml-hex: allocating new session: HTP0:0 +... +load_tensors: offloading output layer to GPU +load_tensors: offloaded 27/27 layers to GPU +load_tensors: CPU model buffer size = 300.00 MiB +load_tensors: HTP0:0 model buffer size = 1400.26 MiB +... +llama_perf_context_print: prompt eval time = 320.00 ms / 1024 tokens ( 0.31 ms per token, 3200.00 tokens per second) +llama_perf_context_print: eval time = 2100.00 ms / 100 runs ( 21.00 ms per token, 47.62 tokens per second) +``` + +Simple question for Llama-3.2-1B: + +``` +~/src/llama.cpp$ ./scripts/snapdragon/run.py --target android --devices HTP0 -- llama-cli -m Llama-3.2-1B-Instruct-Q4_0.gguf -p "what is the most popular cookie in the world?" ... ggml-hex: Hexagon backend (experimental) : allocating new registry : ndev 1 ggml-hex: Hexagon Arch version v79 @@ -142,8 +170,7 @@ ggml-hex: new session: HTP0 : session-id 0 domain-id 3 uri file:///libggml-htp-v load_tensors: offloading output layer to GPU load_tensors: offloaded 17/17 layers to GPU load_tensors: CPU model buffer size = 225.49 MiB -load_tensors: HTP0 model buffer size = 0.26 MiB -load_tensors: HTP0-REPACK model buffer size = 504.00 MiB +load_tensors: HTP0 model buffer size = 504.26 MiB ... I hope this helps you understand the world's most popular cookies! [end of text] ... @@ -156,60 +183,25 @@ llama_perf_context_print: graphs reused = 473 llama_memory_breakdown_print: | memory breakdown [MiB] | total free self model context compute unaccounted | llama_memory_breakdown_print: | - HTP0 (Hexagon) | 2048 = 2048 + ( 0 = 0 + 0 + 0) + 0 | llama_memory_breakdown_print: | - Host | 439 = 225 + 136 + 77 | -llama_memory_breakdown_print: | - HTP0-REPACK | 504 = 504 + 0 + 0 | -``` - -Summary request for OLMoE-1B-7B. This is a large model that requires two HTP sessions/devices - -``` -~/src/llama.cpp$ M=OLMoE-1B-7B-0125-Instruct-Q4_0.gguf NDEV=2 D=HTP0,HTP1 ./scripts/snapdragon/adb/run-completion.sh -f surfing.txt -... -ggml-hex: Hexagon backend (experimental) : allocating new registry : ndev 1 -ggml-hex: Hexagon Arch version v81 -ggml-hex: allocating new session: HTP0 -ggml-hex: allocating new session: HTP1 -... -load_tensors: offloading output layer to GPU -load_tensors: offloaded 17/17 layers to GPU -load_tensors: CPU model buffer size = 143.86 MiB -load_tensors: HTP1 model buffer size = 0.23 MiB -load_tensors: HTP1-REPACK model buffer size = 1575.00 MiB -load_tensors: HTP0 model buffer size = 0.28 MiB -load_tensors: HTP0-REPACK model buffer size = 2025.00 MiB -... -llama_context: CPU output buffer size = 0.19 MiB -llama_kv_cache: HTP1 KV buffer size = 238.00 MiB -llama_kv_cache: HTP0 KV buffer size = 306.00 MiB -llama_kv_cache: size = 544.00 MiB ( 8192 cells, 16 layers, 1/1 seqs), K (q8_0): 272.00 MiB, V (q8_0): 272.00 MiB -llama_context: HTP0 compute buffer size = 15.00 MiB -llama_context: HTP1 compute buffer size = 15.00 MiB -llama_context: CPU compute buffer size = 24.56 MiB -... -llama_perf_context_print: prompt eval time = 1730.57 ms / 212 tokens ( 8.16 ms per token, 122.50 tokens per second) -llama_perf_context_print: eval time = 5624.75 ms / 257 runs ( 21.89 ms per token, 45.69 tokens per second) -llama_perf_context_print: total time = 7377.33 ms / 469 tokens -llama_perf_context_print: graphs reused = 255 -llama_memory_breakdown_print: | memory breakdown [MiB] | total free self model context compute unaccounted | -llama_memory_breakdown_print: | - HTP0 (Hexagon) | 2048 = 2048 + ( 0 = 0 + 0 + 0) + 0 | -llama_memory_breakdown_print: | - HTP1 (Hexagon) | 2048 = 2048 + ( 0 = 0 + 0 + 0) + 0 | -llama_memory_breakdown_print: | - Host | 742 = 144 + 544 + 54 | -llama_memory_breakdown_print: | - HTP1-REPACK | 1575 = 1575 + 0 + 0 | -llama_memory_breakdown_print: | - HTP0-REPACK | 2025 = 2025 + 0 + 0 | ``` -Op test for MUL_MAT +Op test for MUL_MAT: ``` -~/src/llama.cpp$ HB=0 ./scripts/snapdragon/adb/run-tool.sh test-backend-ops -b HTP0 -o MUL_MAT +~/src/llama.cpp$ ./scripts/snapdragon/run.py --target adb --hex-hostbuf 0 --devices HTP0:0 -- test-backend-ops -b HTP0:0 -o MUL_MAT ... -Backend 2/3: HTP0 +Backend 2/3: HTP0:0 Device description: Hexagon Device memory: 2048 MB (2048 MB free) MUL_MAT(type_a=q4_0,type_b=f32,m=16,n=1,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK MUL_MAT(type_a=q4_0,type_b=f32,m=16,n=2,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK MUL_MAT(type_a=q4_0,type_b=f32,m=16,n=3,k=256,bs=[1,1],nr=[1,1],per=[0,1,2,3],v=0,o=1): OK +``` + +Llama benchmark: -~/src/llama.cpp-hexagon$ M=Llama-3.2-1B-Instruct-Q4_0.gguf ./scripts/snapdragon/adb/run-bench.sh -p 128 -n 64 +``` +~/src/llama.cpp$ ./scripts/snapdragon/run.py --target adb --devices HTP0 -- llama-bench -p 128 -n 64 -m Llama-3.2-1B-Instruct-Q4_0.gguf ... ggml-hex: Hexagon backend (experimental) : allocating new registry : ndev 1 ggml-hex: Hexagon Arch version v79 @@ -219,15 +211,20 @@ ggml-hex: new session: HTP0 : session-id 0 domain-id 3 uri file:///libggml-htp-v | ---------------| ---------: | -----: | ---------- | --: | ------: | ------: | ---: | ----: | ------------: | | llama 1B Q4_0 | 729.75 MiB | 1.24 B | HTP | 99 | 4 | 128 | 0 | pp128 | 169.42 ± 1.75 | | llama 1B Q4_0 | 729.75 MiB | 1.24 B | HTP | 99 | 4 | 128 | 0 | tg64 | 51.54 ± 1.13 | - -build: 6a8cf8914 (6733) ``` ## Environment variables -- `GGML_HEXAGON_NDEV=1` - Controls the number of devices/sessions to allocate. The default is 1. - Most quantized models under 4B fit into a single session; an 8B model needs two, and a 20B model needs four. +- `GGML_HEXAGON_DEVICES` (default: not set, defaults to HTP0 session) + Controls which NPU devices and sessions to allocate. Can be configured as: + - A single integer `N`: Allocates `N` sessions named `HTP0`, `HTP1`, ..., `HTP` (behaves identically to `GGML_HEXAGON_NDEV=N`). + - A comma-separated list of device names in `HTP:` format (or legacy `HTP` format). For example, `HTP0:0,HTP0:1` creates two virtual + sessions on the first physical NPU (useful for memory limits). `HTP0:0,HTP1:0` allocates one session on each of the two physical NPUs + on a dual-NPU device. + +- `GGML_HEXAGON_NDEV` (deprecated) + Replaced by `GGML_HEXAGON_DEVICES`. Controls the number of virtual sessions to allocate on physical NPU `0`. + Allocates sessions named `HTP0`, `HTP1`, etc. - `GGML_HEXAGON_NHVX=0` Controls the number of HVX hardware threads to use. The default is all (actual number varies depending on the hardware version). @@ -255,26 +252,17 @@ build: 6a8cf8914 (6733) - `2` Extended profile with per-op `usecs`, `cycles` and default PMU counter data - `0x1,...,0x8` Extended profile with per-op `usecs`, `cycles` and custom PMU counter data - The logging output can be either saved into a file for post-processing or it can be piped directly into the post-processing tool to generate the report. + The logging output can be either saved into a file for post-processing or it can be piped directly into the post-processing tool + to generate the report. Examples: - `GGML_HEXAGON_PROFILE=1 llama-completion ... |& ./scripts/snapdragon/ggml-hexagon-profile.py -` - -- `GGML_HEXAGON_OPSTAGE=0x0` - Allows enabling specific stages of the Op processing pipeline: - - - `0x1` Enable Op Queue (i.e., queuing Ops into NPU) - - `0x2` Enable Op Compute (MUL_MAT, etc.) - - Examples: - - `GGML_HEXAGON_OPSTAGE=0x1 llama-completion ...` - Ops are enqueued to the NPU but dma & compute are disabled - `GGML_HEXAGON_OPSTAGE=0x3 llama-completion ...` - Full queuing and processing of Ops (default) + `GGML_HEXAGON_PROFILE=1 ./scripts/snapdragon/run.py --target adb -- llama-cli ... |& ./scripts/snapdragon/ggml-hexagon-profile.py -` - `GGML_HEXAGON_OPFILTER=regex` Allows filtering (disabling) Ops that match the regex pattern: Examples: - `GGML_HEXAGON_OPFILTER="FLASH_ATTN_EXT" llama-completion ...` - Disable Flash Attention on Hexagon (falls back to CPU or GPU) - `GGML_HEXAGON_OPFILTER="ADD\|SUB" llama-completion ...` - Disable ADD and SUB on Hexagon (fall back to CPU or GPU) + `GGML_HEXAGON_OPFILTER="FLASH_ATTN_EXT" ./scripts/snapdragon/run.py --target adb -- llama-cli ...` - Disable Flash Attention on Hexagon (falls back to CPU or GPU) + `GGML_HEXAGON_OPFILTER="ADD\|SUB" ./scripts/snapdragon/run.py --target adb -- llama-cli ...` - Disable ADD and SUB on Hexagon (fall back to CPU or GPU) + diff --git a/docs/backend/snapdragon/developer.md b/docs/backend/snapdragon/developer.md index 9d56638e3d59..d7d9f2a2790d 100644 --- a/docs/backend/snapdragon/developer.md +++ b/docs/backend/snapdragon/developer.md @@ -39,22 +39,21 @@ the repacking. ## Large model handling -Hexagon NPU session (aka Process Domain (PD) in the Hexagon docs) is limited to a memory mapping of around 3.5GB. -In llama.cpp/GGML the Hexagon session is mapped to a single GGML backend device (HTP0, HTP1, etc). +Hexagon NPU sessions (aka Process Domains (PD) in the Hexagon SDK) are limited to a maximum memory mapping window of around 3.5GB. +In llama.cpp/GGML, each Hexagon session is mapped to a single GGML backend device (e.g., `HTP0:0`, `HTP0:1`, etc. when using +`GGML_HEXAGON_DEVICES`, or `HTP0`, `HTP1` in legacy mode). -In order to map models larger than 3.5GB we need to allocate multiple devices and split the model. -For this we're taking advantage of the llama.cpp/GGML multi-GPU layer-splitting support. -Each Hexagon device behaves like a GPU from the offload and model splitting perspective. +To support running models larger than 3.5GB on a single device, the Hexagon backend dynamically maps and unmaps execution buffers +during the graph execution cycle to stay within the Process Domain window. This enables large models to run successfully on a single +NPU device. -Here is an example of running GPT-OSS-20B model on a newer Snapdragon device with 16GB of DDR. +Alternatively, users can choose to use standard llama.cpp/GGML layer-splitting mode to partition and split the model across +multiple Hexagon devices or virtual sessions (which behave like multiple GPUs from the offload and splitting perspective). + +Here is an example of running GPT-OSS-20B model on a Snapdragon device using 4 virtual sessions on a single NPU (physical index 0). ``` -M=gpt-oss-20b-Q4_0.gguf NDEV=4 D=HTP0,HTP1,HTP2,HTP3 P=surfing.txt scripts/snapdragon/adb/run-completion.sh -f surfing.txt -n 32 -... -LD_LIBRARY_PATH=/data/local/tmp/llama.cpp/lib -ADSP_LIBRARY_PATH=/data/local/tmp/llama.cpp/lib -GGML_HEXAGON_NDEV=4 ./bin/llama-cli --load-mode none -m /data/local/tmp/llama.cpp/../gguf/gpt-oss-20b-Q4_0.gguf - -t 4 --ctx-size 8192 --batch-size 128 -ctk q8_0 -ctv q8_0 -fa on -ngl 99 --device HTP0,HTP1,HTP2,HTP3 -no-cnv -f surfing.txt +~/src/llama.cpp$ ./scripts/snapdragon/run.py --target adb --devices HTP0:0,HTP0:1,HTP0:2,HTP0:3 -- llama-cli --load-mode none -m /data/local/tmp/gguf/gpt-oss-20b-Q4_0.gguf -t 4 --ctx-size 8192 --batch-size 128 -ctk q8_0 -ctv q8_0 -fa on -ngl 99 -no-cnv -f surfing.txt ... llama_model_loader: - type f32: 289 tensors llama_model_loader: - type q4_0: 96 tensors @@ -63,33 +62,29 @@ llama_model_loader: - type mxfp4: 72 tensors ... load_tensors: offloaded 25/25 layers to GPU load_tensors: CPU model buffer size = 1182.09 MiB -load_tensors: HTP1 model buffer size = 6.64 MiB -load_tensors: HTP1-REPACK model buffer size = 2505.94 MiB -load_tensors: HTP3 model buffer size = 5.55 MiB -load_tensors: HTP3-REPACK model buffer size = 2088.28 MiB -load_tensors: HTP0 model buffer size = 7.75 MiB -load_tensors: HTP0-REPACK model buffer size = 2923.59 MiB -load_tensors: HTP2 model buffer size = 6.64 MiB -load_tensors: HTP2-REPACK model buffer size = 2505.94 MiB +load_tensors: HTP0:1 model buffer size = 2512.58 MiB +load_tensors: HTP0:3 model buffer size = 2093.83 MiB +load_tensors: HTP0:0 model buffer size = 2931.34 MiB +load_tensors: HTP0:2 model buffer size = 2512.58 MiB ... llama_context: n_ctx_per_seq (8192) < n_ctx_train (131072) -- the full capacity of the model will not be utilized llama_context: CPU output buffer size = 0.77 MiB llama_kv_cache_iswa: creating non-SWA KV cache, size = 8192 cells -llama_kv_cache: HTP1 KV buffer size = 25.50 MiB -llama_kv_cache: HTP3 KV buffer size = 25.50 MiB -llama_kv_cache: HTP0 KV buffer size = 25.50 MiB -llama_kv_cache: HTP2 KV buffer size = 25.50 MiB +llama_kv_cache: HTP0:1 KV buffer size = 25.50 MiB +llama_kv_cache: HTP0:3 KV buffer size = 25.50 MiB +llama_kv_cache: HTP0:0 KV buffer size = 25.50 MiB +llama_kv_cache: HTP0:2 KV buffer size = 25.50 MiB llama_kv_cache: size = 102.00 MiB ( 8192 cells, 12 layers, 1/1 seqs), K (q8_0): 51.00 MiB, V (q8_0): 51.00 MiB llama_kv_cache_iswa: creating SWA KV cache, size = 256 cells -llama_kv_cache: HTP1 KV buffer size = 0.80 MiB -llama_kv_cache: HTP3 KV buffer size = 0.53 MiB -llama_kv_cache: HTP0 KV buffer size = 1.06 MiB -llama_kv_cache: HTP2 KV buffer size = 0.80 MiB +llama_kv_cache: HTP0:1 KV buffer size = 0.80 MiB +llama_kv_cache: HTP0:3 KV buffer size = 0.53 MiB +llama_kv_cache: HTP0:0 KV buffer size = 1.06 MiB +llama_kv_cache: HTP0:2 KV buffer size = 0.80 MiB llama_kv_cache: size = 3.19 MiB ( 256 cells, 12 layers, 1/1 seqs), K (q8_0): 1.59 MiB, V (q8_0): 1.59 MiB -llama_context: HTP0 compute buffer size = 16.06 MiB -llama_context: HTP1 compute buffer size = 16.06 MiB -llama_context: HTP2 compute buffer size = 16.06 MiB -llama_context: HTP3 compute buffer size = 16.06 MiB +llama_context: HTP0:0 compute buffer size = 16.06 MiB +llama_context: HTP0:1 compute buffer size = 16.06 MiB +llama_context: HTP0:2 compute buffer size = 16.06 MiB +llama_context: HTP0:3 compute buffer size = 16.06 MiB llama_context: CPU compute buffer size = 98.19 MiB ... llama_perf_context_print: prompt eval time = 3843.67 ms / 197 tokens ( 19.51 ms per token, 51.25 tokens per second) @@ -97,13 +92,9 @@ llama_perf_context_print: eval time = 1686.13 ms / 31 runs ( 54.3 llama_perf_context_print: total time = 6266.30 ms / 228 tokens llama_perf_context_print: graphs reused = 30 llama_memory_breakdown_print: | memory breakdown [MiB] | total free self model context compute unaccounted | -llama_memory_breakdown_print: | - HTP0 (Hexagon) | 2048 = 2048 + ( 0 = 0 + 0 + 0) + 0 | -llama_memory_breakdown_print: | - HTP1 (Hexagon) | 2048 = 2048 + ( 0 = 0 + 0 + 0) + 0 | -llama_memory_breakdown_print: | - HTP2 (Hexagon) | 2048 = 2048 + ( 0 = 0 + 0 + 0) + 0 | -llama_memory_breakdown_print: | - HTP3 (Hexagon) | 2048 = 2048 + ( 0 = 0 + 0 + 0) + 0 | +llama_memory_breakdown_print: | - HTP0:0 (Hexagon) | 2048 = 2048 + ( 0 = 0 + 0 + 0) + 0 | +llama_memory_breakdown_print: | - HTP0:1 (Hexagon) | 2048 = 2048 + ( 0 = 0 + 0 + 0) + 0 | +llama_memory_breakdown_print: | - HTP0:2 (Hexagon) | 2048 = 2048 + ( 0 = 0 + 0 + 0) + 0 | +llama_memory_breakdown_print: | - HTP0:3 (Hexagon) | 2048 = 2048 + ( 0 = 0 + 0 + 0) + 0 | llama_memory_breakdown_print: | - Host | 1476 = 1208 + 105 + 162 | -llama_memory_breakdown_print: | - HTP1-REPACK | 2505 = 2505 + 0 + 0 | -llama_memory_breakdown_print: | - HTP3-REPACK | 2088 = 2088 + 0 + 0 | -llama_memory_breakdown_print: | - HTP0-REPACK | 2923 = 2923 + 0 + 0 | -llama_memory_breakdown_print: | - HTP2-REPACK | 2505 = 2505 + 0 + 0 | ``` diff --git a/docs/backend/snapdragon/linux.md b/docs/backend/snapdragon/linux.md index 90fdadb6c93f..d4ecd9b1b167 100644 --- a/docs/backend/snapdragon/linux.md +++ b/docs/backend/snapdragon/linux.md @@ -1,25 +1,37 @@ # Snapdragon-based Linux devices -## Docker Setup +The cross-compilation is performed using the Snapdragon Linux Docker toolchain image (see +[github.com/snapdragon-toolchain](https://github.com/snapdragon-toolchain)): -The easiest way to build llama.cpp for a Snapdragon-based Linux device is using the toolchain Docker image (see [github.com/snapdragon-toolchain](https://github.com/snapdragon-toolchain)). -This image includes OpenCL SDK, Hexagon SDK, CMake, and the ARM64 Linux cross-compilation toolchain. +* **Linux toolchain**: `ghcr.io/snapdragon-toolchain/arm64-linux:v0.7` -Cross-compilation is supported on **Linux X86** hosts. The resulting binaries are deployed to and run on the target **Qualcomm Snapdragon ARM64 Linux** device. +The unified build utility (`scripts/snapdragon/build.py`) automatically pulls +and orchestrates this container to perform target compilation. You only need to +ensure that Docker is running on your host machine. + +## How to Build + +### Using build.py script (Recommended) + +The easiest way to build llama.cpp is by using the `scripts/snapdragon/build.py` script. It automatically copies the CMake presets, +launches the correct compilation Docker container, builds the libraries and tools, +installs them, and optionally pushes them to your target device. + +Build and deploy for a Linux target (using SSH deployment alias `lnx` or `linux`): ``` -~/src/llama.cpp$ docker run -it -u $(id -u):$(id -g) --volume $(pwd):/workspace --platform linux/amd64 ghcr.io/snapdragon-toolchain/arm64-linux:v0.1 -[d]/> cd /workspace +$ ./scripts/snapdragon/build.py --target lnx:user@host --push ``` -Note: The rest of the **Linux** build process assumes that you're running inside the toolchain container. +### Manual CMake Build +Alternatively, you can build llama.cpp manually by entering the cross-compilation Docker container and running the CMake commands: -## How to Build +```bash +# Start the cross-compilation container manually: +~/src/llama.cpp$ docker run -it --rm -u $(id -u):$(id -g) --volume $(pwd):/workspace --platform linux/amd64 ghcr.io/snapdragon-toolchain/arm64-linux:v0.7 -Let's build llama.cpp with CPU, OpenCL, and Hexagon backends via CMake presets: - -``` +# Inside the container, build the project using presets: [d]/workspace> cp docs/backend/snapdragon/CMakeUserPresets.json . [d]/workspace> cmake --preset arm64-linux-snapdragon-release -B build-snapdragon @@ -30,17 +42,19 @@ Let's build llama.cpp with CPU, OpenCL, and Hexagon backends via CMake presets: To generate an installable "package" simply use cmake --install, then zip it: ``` -[d]/workspace> cmake --install build-snapdragon --prefix pkg-snapdragon -[d]/workspace> zip -r pkg-snapdragon.zip pkg-snapdragon +[d]/workspace> cmake --install build-snapdragon --prefix pkg-linux +[d]/workspace> zip -r pkg-linux.zip pkg-linux ``` ## How to Install -For this step, you will deploy the built binaries and libraries to the target Linux device. Transfer `pkg-snapdragon.zip` to the target device, then unzip it and set up the environment variables: +For this step, you will deploy the built binaries and libraries to the target +Linux device. Transfer `pkg-linux.zip` to the target device, then unzip it +and set up the environment variables: ``` -$ unzip pkg-snapdragon.zip -$ cd pkg-snapdragon +$ unzip pkg-linux.zip +$ cd pkg-linux $ export LD_LIBRARY_PATH=./lib $ export ADSP_LIBRARY_PATH=./lib ``` @@ -52,7 +66,28 @@ $ wget https://huggingface.co/bartowski/Llama-3.2-3B-Instruct-GGUF/resolve/main/ ``` ## How to Run -Next, since we have setup the environment variables, we can run the llama-cli with the Hexagon backends: +You can run locally on the Snapdragon Linux device: +``` +$ ./scripts/snapdragon/run.py --devices HTP0 -- llama-cli -m Llama-3.2-3B-Instruct-Q4_0.gguf -ngl 99 -p "what is the most popular cookie in the world?" +``` + +Or run remotely from your host development machine using the SSH target option: +``` +$ ./scripts/snapdragon/run.py --target lnx:user@host --devices HTP0 -- llama-cli -m Llama-3.2-3B-Instruct-Q4_0.gguf -ngl 99 -p "what is the most popular cookie in the world?" +``` + +For multi-NPU systems, you can run a tensor split completion command targeting a remote Linux system: +``` +$ ./scripts/snapdragon/run.py --target ubuntu:maxk@192.168.1.87 --device HTP0:0,HTP1:0 -- llama-completion -m models/gemma-2b-it-Q4_0.gguf -f prompts/sample_prompt_1024.txt --jinja -st --split-mode tensor --ctx-size 8192 +``` + +This translates to the following command being executed remotely via SSH: +``` ++ ssh maxk@192.168.1.87 "cd ~/llama.cpp && ulimit -c unlimited && LD_LIBRARY_PATH=./lib ADSP_LIBRARY_PATH=./lib GGML_HEXAGON_DEVICES=HTP0:0,HTP1:0 GGML_HEXAGON_OPPOLL=1 ./bin/llama-completion -m models/gemma-2b-it-Q4_0.gguf -f prompts/sample_prompt_1024.txt --jinja -st --split-mode tensor --ctx-size 8192 -v -n 16 --device HTP0:0,HTP1:0 -ngl 99 --ubatch-size 1024 -fa on -t 6" +``` + +Alternatively, you can run the binary directly on the device: ``` $ ./bin/llama-cli -m Llama-3.2-3B-Instruct-Q4_0.gguf --device HTP0 -ngl 99 -p "what is the most popular cookie in the world?" ``` + diff --git a/docs/backend/snapdragon/windows.md b/docs/backend/snapdragon/windows.md index aa731413c909..3f7d60dd95f8 100644 --- a/docs/backend/snapdragon/windows.md +++ b/docs/backend/snapdragon/windows.md @@ -1,3 +1,18 @@ +# Snapdragon-based Windows devices + +## Tool Dependencies + +Native Windows 11 arm64 builds have the following tool dependencies: +- MS Visual Studio 2026 (Community Edition or Pro) + - MSVC arm64 standard and runtime libraries + - UCRT and Driver Kit +- LLVM core libraries and Clang compiler (winget) +- CMake, Git, Python (winget) +- Hexagon SDK Community Edition 6.6 or later (see below) +- OpenCL SDK 2.3 or later (see below) + +Note: The rest of the **Windows** build process assumes that you're running natively in Powershell. + ## Overview The document covers procedures for installing the latest GPU and NPU drivers, and OpenCL and Hexagon SDKs. @@ -53,7 +68,8 @@ Download the driver from https://softwarecenter.qualcomm.com/catalog/item/Qualcomm_HND -After the automated installation and reboot please make sure that the Hexagon NPU device shows up in the `Device Manager` (under `Neural Processors`). +After the automated installation and reboot please make sure that the Hexagon NPU device shows up in the `Device Manager` +(under `Neural Processors`). If the device is not available you can try installing all components (`qcnspmcdm8380`, `qcnspmcdm8380_ext`) manually. The components are extracted into @@ -130,12 +146,12 @@ However, additional settings are required for generating and signing HTP Ops lib > cmake --preset arm64-windows-snapdragon-release -B build-wos ... -> cmake --install build-wos --prefix pkg-snapdragon +> cmake --install build-wos --prefix pkg-wos ``` Once the build is complete HTP ops libraries will be installed like this ``` -> dir pkg-snapdragon/lib +> dir pkg-wos/lib ... -a---- 1/22/2026 6:01 PM 187656 libggml-htp-v73.so -a---- 1/22/2026 6:01 PM 191752 libggml-htp-v75.so @@ -147,8 +163,8 @@ Once the build is complete HTP ops libraries will be installed like this The .cat file, the signature and proper certificate installation can be verified with ``` -> signtool.exe verify /v /pa .\pkg-snapdragon\lib\libggml-htp.cat -Verifying: .\pkg-snapdragon\lib\libggml-htp.cat +> signtool.exe verify /v /pa .\pkg-wos\lib\libggml-htp.cat +Verifying: .\pkg-wos\lib\libggml-htp.cat Signature Index: 0 (Primary Signature) Hash of file (sha256): 9820C664DA59D5EAE31DBB664127FCDAEF59CDC31502496BC567544EC2F401CF @@ -156,6 +172,6 @@ Hash of file (sha256): 9820C664DA59D5EAE31DBB664127FCDAEF59CDC31502496BC567544EC Signing Certificate Chain: Issued to: GGML.HTP.v1 ... -Successfully verified: .\pkg-snapdragon\lib\libggml-htp.cat +Successfully verified: .\pkg-wos\lib\libggml-htp.cat ... ``` diff --git a/docs/speculative.md b/docs/speculative.md index 0f9f8a3d977a..ffb1e34c7fbf 100644 --- a/docs/speculative.md +++ b/docs/speculative.md @@ -212,6 +212,15 @@ Use `--backend-sampling` to run supported target-model samplers on the model bac Unsupported samplers and device layouts fall back to CPU sampling. Tensor split mode does not support backend sampling. A fixed seed produces repeatable random draws, but stochastic CPU and backend sampling can still select different tokens because floating-point operations can differ between implementations and devices. Use greedy sampling when exact output matching is required. +### Synthetic Acceptance + +`llama-server` and `llama-cli` can replace normal speculative verification with synthetic decisions for benchmarking. The generated output is not valid model output because accepted draft tokens do not have to match the target model. + +Use exactly one of these options: + +- `--spec-synth-rates P0,P1,...` sets unconditional per-position acceptance probabilities. Entry `i` is the probability that the first `i+1` draft tokens are all accepted. The number of entries must match the effective maximum draft length. Values must be finite, within `[0, 1]`, and monotonically non-increasing. +- `--spec-synth-len L` sets the target mean acceptance length, including the target token. For `K` maximum draft tokens, `L` must be within `[1, K+1]`. The server finds a constant conditional probability `p` such that `p + p^2 + ... + p^K = L - 1`, then uses unconditional rates `[p, p^2, ..., p^K]`. + ### General Speculative Parameters ``` diff --git a/ggml/src/ggml-hexagon/ggml-hexagon.cpp b/ggml/src/ggml-hexagon/ggml-hexagon.cpp index e8a5009b381b..53e860755910 100644 --- a/ggml/src/ggml-hexagon/ggml-hexagon.cpp +++ b/ggml/src/ggml-hexagon/ggml-hexagon.cpp @@ -6,6 +6,7 @@ #include #include +#include #include #include #include @@ -18,6 +19,7 @@ #include #include #include +#include #include #ifdef _WIN32 @@ -52,6 +54,8 @@ #include "htp/matmul-ops.h" #include "htp/flash-attn-ops.h" #include "htp/unary-ops.h" +#include "htp/get-rows-ops.h" +#include "htp/set-rows-ops.h" #include "htp_iface.h" #include "htp-drv.h" @@ -59,6 +63,36 @@ using intvec = std::vector; using uintvec = std::vector; using u32vec = std::vector; +#define GGML_HEXAGON_MAX_SESSIONS 16 + +#define GGML_HEXAGON_FENCE_BUFFER_SIZE 8192 +#define GGML_HEXAGON_FENCE_SLOT_SIZE 128 + +struct ggml_hexagon_device_config { + int physical_idx = 0; + int virtual_idx = 0; + std::string name; +}; + +static ggml_hexagon_device_config opt_device_configs[GGML_HEXAGON_MAX_SESSIONS]; + +static int get_domain_id(int physical_idx) { + switch (physical_idx) { + case 0: return 3; // CDSP0 (all devices) + case 1: return 4; // CDSP1 (IQ9, IQ10) + case 2: return 18; // CDSP2 (IQ10) + case 3: return 19; // CDSP3 (IQ10) + default: return CDSP_DOMAIN_ID + physical_idx; + } +} + +static std::string get_domain_name(int physical_idx) { + if (physical_idx == 0) { + return CDSP_DOMAIN_NAME; + } + return std::string("cdsp") + std::to_string(physical_idx); +} + static int opt_arch = 0; // autodetect static size_t opt_ndev = 1; static size_t opt_nhvx = 0; // use all @@ -68,20 +102,19 @@ static size_t opt_mbuf = 1ul * 1024 * 1024 * 1024; // max buffer size static int opt_etm = 0; static int opt_verbose = 0; static int opt_profile = 0; // profiling mode (0-disabled, 1-basic, 2-pmu) -static int opt_hostbuf = 1; // hostbuf ON by default +static bool opt_hostbuf = false; static int opt_mm_select = 3; // 3 = HMX -> Tiled -> Flat -> CPU, 2 = Tiled -> Flat -> CPU, 1 = Flat -> CPU static int opt_fa_select = 2; // 2 = HMX -> HVX -> CPU, 1 = HVX -> CPU, 0 = CPU (unsupported) +static int opt_ar_select = 2; // 2 = fused ALLREDUCE+ADD (DMA, default), 1 = unfused ALLREDUCE (DMA), 0 = fallback to CPY+FENCE // Default PMU events, if profiling with PMU (mode=2) is enabled // See https://docs.qualcomm.com/doc/80-N2040-60/topic/pmu-events.html // https://docs.qualcomm.com/doc/80-N2040-61/topic/hvx-pmu-events.html static u32vec opt_pmu_evt { 0x3, 0x111, 0x100, 0x105, 0x240, 0x256, 0x7D, 0x8C }; -// Enable all stages by default -static int opt_opstage = HTP_OPSTAGE_QUEUE | HTP_OPSTAGE_COMPUTE; static int opt_opbatch = 1024; // max number of ops in a batch -static int opt_opqueue = 16; // max number of pending batches +static int opt_opqueue = 64; // max number of pending batches static int opt_optrace = 0; // trace buffer size per thread (0 means default) static int opt_oppoll = 0; // polling for batch completions static int opt_opfusion = 1; // enable/disable op fusion @@ -121,7 +154,7 @@ static void ggml_hexagon_dump_op_exec(const std::string &sess_name, const htp_op static void ggml_hexagon_dump_op_supp(const std::string &sess_name, const struct ggml_tensor * op, bool supp) { if (!opt_verbose) return; - htp_opformat fmt(htp_opformat(htp_opnode{const_cast(op), {}, HTP_OP_INVALID})); + htp_opformat fmt(htp_opformat(htp_opnode(HTP_OP_INVALID, const_cast(op)))); GGML_LOG_DEBUG("ggml-hex: %s supports-op %s|%s|%s|%s|%s|%s|%s\n", sess_name.c_str(), ggml_op_desc(op), fmt.names, fmt.dims, fmt.types, fmt.strides, fmt.buffs, supp ? "yes" : "no"); } @@ -144,6 +177,7 @@ static const char * htp_event_name(uint16_t id) { case HTP_TRACE_EVT_L2FLUSH: return "L2FLUSH"; case HTP_TRACE_EVT_INIT: return "INIT"; case HTP_TRACE_EVT_BUFF: return "BUFF"; + case HTP_TRACE_EVT_FENCE: return "FENCE"; default: return "UNKNOWN"; } } @@ -205,7 +239,12 @@ static void ggml_hexagon_dump_trace_events(const std::string & sess_name, const } } -// ** +enum ggml_hexagon_tensor_flags { + GGML_HEXAGON_TENSOR_REPACK = (1 << 0), + GGML_HEXAGON_TENSOR_WEIGHT = (1 << 1), + GGML_HEXAGON_TENSOR_FENCE = (1 << 2), + GGML_HEXAGON_TENSOR_FUSEABLE = (1 << 3), +}; static inline bool ggml_hexagon_is_repack_type(enum ggml_type type) { return type == GGML_TYPE_Q4_0 || type == GGML_TYPE_Q4_1 || @@ -227,6 +266,15 @@ static void ggml_hexagon_precompute_matmul_params( struct htp_mm_kernel_params * kparams ); +static void ggml_hexagon_precompute_fused_matmul_add_params( + const struct ggml_hexagon_session * sess, + const struct ggml_tensor * src0, + const struct ggml_tensor * src1, + const struct ggml_tensor * src2, + const struct ggml_tensor * dst, + struct htp_mm_kernel_params * kparams +); + static void ggml_hexagon_precompute_unary_params( const struct ggml_hexagon_session * sess, uint32_t op, @@ -236,25 +284,75 @@ static void ggml_hexagon_precompute_unary_params( struct htp_unary_kernel_params * kparams ); -static void ggml_hexagon_precompute_fused_qkv_params( +static void ggml_hexagon_precompute_get_rows_params( const struct ggml_hexagon_session * sess, const struct ggml_tensor * src0, const struct ggml_tensor * src1, - struct htp_mm_kernel_params * kparams + const struct ggml_tensor * dst, + struct htp_get_rows_kernel_params * kparams ); -static void ggml_hexagon_precompute_fused_ffn_params( +static void ggml_hexagon_precompute_set_rows_params( const struct ggml_hexagon_session * sess, const struct ggml_tensor * src0, const struct ggml_tensor * src1, + const struct ggml_tensor * dst, + struct htp_set_rows_kernel_params * kparams +); + +static void ggml_hexagon_precompute_fused_mmnx_params( + const struct ggml_hexagon_session * sess, + const struct ggml_tensor * src0, + const struct ggml_tensor * src1, + int32_t n_weights, struct htp_mm_kernel_params * kparams ); +static bool ggml_hexagon_precompute_allreduce_params( + const struct ggml_hexagon_session * sess, + const struct ggml_tensor * dst, + uint32_t rank, + uint32_t n_ranks, + bool has_add, + bool is_row_bcast, + struct htp_allreduce_kernel_params * kparams +); + +static bool mm_is_hmx_eligible(const ggml_tensor * t); +static bool is_mergeable_mul_mat(const ggml_tensor * t); +static bool is_mergeable_mul_mat_pair(const ggml_tensor * n1, const ggml_tensor * n2); + // ** backend sessions +struct ggml_hexagon_tensor_extra { + std::vector shadow_buf; + size_t shadow_size { 0 }; + uint32_t flags { 0 }; +}; + +static inline bool ggml_hexagon_tensor_is_fuseable(const struct ggml_tensor * t) { + if (!t || !t->extra) return false; + auto extra = (const struct ggml_hexagon_tensor_extra *) t->extra; + return (extra->flags & GGML_HEXAGON_TENSOR_FUSEABLE) != 0; +} + +struct htp_opnode; + struct ggml_hexagon_opbatch; struct ggml_hexagon_opqueue; -struct htp_opnode; +struct ggml_hexagon_shared_buffer; +struct ggml_hexagon_session; + +struct ggml_backend_hexagon_comm_context { + std::vector backends; + size_t n_backends = 0; + uint32_t fence_seq = 0; +}; + +struct ggml_hexagon_event { + ggml_hexagon_session * sess = nullptr; + uint64_t seq = 0; +}; struct ggml_hexagon_session { std::string name; @@ -264,6 +362,8 @@ struct ggml_hexagon_session { uint32_t domain_id; uint64_t queue_id; int dev_id; + int phys_idx; + int virt_idx; bool valid_session; bool valid_handle; bool valid_queue; @@ -273,20 +373,24 @@ struct ggml_hexagon_session { ggml_hexagon_opbatch* op_batch; ggml_hexagon_opqueue* op_queue; + std::unordered_map> cloned_buffers; + std::unordered_set sync_peers; + ggml_backend_buffer_type buffer_type = {}; - ggml_backend_buffer_type repack_buffer_type = {}; + ggml_backend_buffer_type host_buffer_type = {}; - uint32_t n_threads = 0; - uint32_t n_hvx = 0; - uint32_t n_hmx = 0; - uint64_t vtcm_size = 0; - size_t max_vmem = 0; + uint32_t n_threads = 0; + uint32_t n_hvx = 0; + uint32_t n_hmx = 0; + uint64_t vtcm_size = 0; + size_t max_vmem = 0; size_t max_bufsize = 0; + uint32_t fence_seq; - struct { - uint64_t uid = 0; - std::vector htp_nodes; - } cached_graph; + uint64_t cached_uid = 0; + std::vector cached_nodes; + + mutable std::unordered_set needs_repack; ggml_hexagon_session(int dev_id, ggml_backend_dev_t dev) noexcept(false); ~ggml_hexagon_session() noexcept(true); @@ -297,10 +401,31 @@ struct ggml_hexagon_session { void release() noexcept(true); void enqueue_op(const htp_opnode & node); - void flush(bool all = true); + void enqueue_cpy(const ggml_tensor * src, ggml_tensor * dst, const ggml_tensor * sync_tensor = nullptr, uint32_t fence_seq = 0); + void enqueue_fence(const ggml_tensor * sync_tensor, uint32_t fence_seq = 0); + void enqueue_allreduce(const ggml_tensor * dst, const std::vector & src_tensors, const std::vector & sync_tensors, uint32_t rank, uint32_t n_ranks, uint32_t fence_seq_entry = 0, uint32_t fence_seq_exit = 0); + void flush(bool all = true); void flush_pending(bool all = false); - void flush_batch(); + void flush_batch(size_t min_ops = 1); + + uint64_t record_event(); + void wait_event(uint64_t seq); + + bool clone_buffer(const ggml_hexagon_shared_buffer*); + + void add_sync_peer(ggml_hexagon_session * peer) { + sync_peers.insert(peer); + } + + void flush_sync_peers() { + if (sync_peers.empty()) return; + + for (auto * peer : sync_peers) { + peer->flush_batch(); + } + sync_peers.clear(); + } }; // ** backend buffers @@ -315,26 +440,68 @@ struct ggml_backend_hexagon_buffer_type_context { std::string name; }; +struct ggml_hexagon_rpcmem_block { + uint8_t * base = nullptr; + int fd = -1; + size_t size = 0; + + ggml_hexagon_rpcmem_block(size_t size) { + base = (uint8_t *) rpcmem_alloc2(RPCMEM_HEAP_ID_SYSTEM, RPCMEM_DEFAULT_FLAGS, size); + if (!base) { + throw std::runtime_error("ggml-hex: rpcmem_alloc failed"); + } + fd = rpcmem_to_fd(base); + if (fd < 0) { + rpcmem_free(base); + throw std::runtime_error("ggml-hex: rpcmem_to_fd failed"); + } + this->size = size; + } + + ~ggml_hexagon_rpcmem_block() { + if (base) { + rpcmem_free(base); + } + } +}; + struct ggml_hexagon_shared_buffer { - ggml_hexagon_session * sess; - uint8_t * base; - size_t size; - int fd; - bool mapped; - bool pinned; + ggml_hexagon_session * sess; + std::shared_ptr mem; + std::vector tensor_extra; + uint32_t fence_head = 0; + size_t fences_size = 0; + bool mapped; + bool pinned; + + const char * c_name() const { return sess->c_name(); } + uint8_t * base() const { return mem ? mem->base : nullptr; } + size_t size() const { return mem ? mem->size : 0; } + int fd() const { return mem ? mem->fd : -1; } + + uint8_t * alloc_fence() { + if (fences_size == 0) return nullptr; + int max_slots = fences_size / GGML_HEXAGON_FENCE_SLOT_SIZE; + uint32_t slot = (fence_head++) % max_slots; + + size_t guard_offset = size() - fences_size; + uint8_t * fence_ptr = base() + guard_offset + (size_t)slot * GGML_HEXAGON_FENCE_SLOT_SIZE; + return fence_ptr; + } void mmap() { + if (!this->mem) return; fastrpc_map_flags flags = this->pinned ? FASTRPC_MAP_FD : FASTRPC_MAP_FD_DELAYED; - int err = fastrpc_mmap(sess->domain_id, this->fd, (void *) this->base, 0, this->size, flags); + int err = fastrpc_mmap(sess->domain_id, fd(), (void *) base(), 0, size(), flags); if (err != 0) { GGML_LOG_ERROR("ggml-hex: %s buffer mapping failed : domain_id %d size %zu fd %d error 0x%08x\n", sess->c_name(), - sess->domain_id, this->size, this->fd, (unsigned) err); + sess->domain_id, size(), fd(), (unsigned) err); throw std::runtime_error("ggml-hex: fastrpc_mmap failed (see log for details)"); } HEX_VERBOSE("ggml-hex: %s mapped buffer: base %p size %zu fd %d pinned %u\n", - sess->c_name(), (void *) this->base, this->size, this->fd, pinned); + sess->c_name(), (void *) base(), size(), fd(), pinned); this->mapped = true; } @@ -342,66 +509,69 @@ struct ggml_hexagon_shared_buffer { void unmap() { if (!this->mapped) return; - if (!this->pinned) { + if (!this->pinned && mem) { // HTP might still hold a reference, tell it drop it - htp_iface_munmap(sess->handle, this->fd); + htp_iface_munmap(sess->handle, fd()); } - fastrpc_munmap(sess->domain_id, this->fd, (void *) this->base, this->size); + if (mem) { + fastrpc_munmap(sess->domain_id, fd(), (void *) base(), size()); + } HEX_VERBOSE("ggml-hex: %s unmapped buffer: base %p size %zu fd %d\n", sess->c_name(), - (void *) this->base, size, this->fd); + (void *) base(), size(), fd()); this->mapped = false; - this->fd = -1; } void alloc(size_t size) { - if (this->base) return; - - this->base = (uint8_t *) rpcmem_alloc2(RPCMEM_HEAP_ID_SYSTEM, RPCMEM_DEFAULT_FLAGS, size); - if (!this->base) { - GGML_LOG_ERROR("ggml-hex: %s failed to allocate buffer : size %zu\n", sess->c_name(), size); - throw std::runtime_error("ggml-hex: rpcmem_alloc failed (see log for details)"); - } + if (this->mem) return; - this->fd = rpcmem_to_fd(this->base); - if (this->fd < 0) { - GGML_LOG_ERROR("ggml-hex: %s failed to get FD for buffer %p\n", sess->c_name(), (void *) this->base); - throw std::runtime_error("ggml-hex: rpcmem_to_fd failed (see log for details)"); - } - this->size = size; + this->mem = std::make_shared(size); HEX_VERBOSE("ggml-hex: %s allocated buffer: base %p size %zu fd %d pinned %d\n", sess->c_name(), - (void *) this->base, this->size, this->fd, (int) pinned); + (void *) base(), this->size(), fd(), (int) pinned); mmap(); } void free() { - if (!this->base) return; - unmap(); - rpcmem_free(this->base); - - HEX_VERBOSE("ggml-hex: %s freed buffer: base %p size %zu fd %d\n", sess->c_name(), - (void *) this->base, size, this->fd); + // The memory is freed when the shared_ptr refcount drops to 0. + HEX_VERBOSE("ggml-hex: %s release ref on buffer: base %p size %zu fd %d\n", sess->c_name(), + (void *) base(), size(), fd()); + this->mem = nullptr; + } + + ggml_hexagon_shared_buffer(ggml_hexagon_session * sess, size_t size, bool pinned = false, size_t fence_size = 0) { + this->sess = sess; + this->mapped = false; + this->pinned = pinned; + this->fences_size = fence_size; + + // Size adjustment inside the buffer class + size_t guard_offset = (size + 4095) & ~4095; + size_t total_size = guard_offset; + if (fence_size > 0) { + total_size += 4096 + fence_size; + } - this->base = NULL; + alloc(total_size); } - ggml_hexagon_shared_buffer(ggml_hexagon_session * sess, size_t size, bool pinned = false) { - this->sess = sess; - this->size = 0; - this->base = nullptr; - this->fd = -1; - this->mapped = false; - this->pinned = pinned; - - alloc(size); + // Clone constructor for cross-session mapping + ggml_hexagon_shared_buffer(ggml_hexagon_session * sess, const ggml_hexagon_shared_buffer & other) { + this->sess = sess; + this->mem = other.mem; + this->mapped = false; + this->pinned = other.pinned; + this->fences_size = other.fences_size; } ~ggml_hexagon_shared_buffer() { free(); + for (auto * extra : tensor_extra) { + delete extra; + } } }; @@ -416,18 +586,25 @@ static void ggml_backend_hexagon_buffer_free_buffer(ggml_backend_buffer_t buffer static void * ggml_backend_hexagon_buffer_get_base(ggml_backend_buffer_t buffer) { auto sbuf = static_cast(buffer->context); - return sbuf->base; + return sbuf->base(); } static enum ggml_status ggml_backend_hexagon_buffer_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) { auto sbuf = static_cast(buffer->context); auto sess = sbuf->sess; - HEX_VERBOSE("ggml-hex: %s init-tensor %s : base %p data %p nbytes %zu usage %d\n", sess->c_name(), - tensor->name, (void *) sbuf->base, tensor->data, ggml_nbytes(tensor), (int) buffer->usage); + HEX_VERBOSE("ggml-hex: %s init-tensor %s : base %p data %p nbytes %zu\n", sess->c_name(), + tensor->name, (void *) sbuf->base(), tensor->data, ggml_nbytes(tensor)); + + auto extra = new ggml_hexagon_tensor_extra(); + sbuf->tensor_extra.push_back(extra); - if (tensor->view_src != NULL && tensor->view_offs == 0) { - return GGML_STATUS_SUCCESS; // nothing to do for the view + tensor->extra = extra; + if (ggml_hexagon_is_repack_type(tensor->type)) { + if (sess->needs_repack.count(tensor)) { + extra->flags |= GGML_HEXAGON_TENSOR_REPACK; + sess->needs_repack.erase(tensor); + } } return GGML_STATUS_SUCCESS; @@ -499,7 +676,7 @@ static void pack_mxfp4_quants(block_mxfp4 * x, const uint8_t * qs, unsigned int } // repack q4_0 data into q4_0_tiled tensor -static void repack_q4_0_tiled(ggml_tensor * t, const void * data, size_t size) { +static void repack_q4_0_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) { const block_q4_0 * src_matrix = (const block_q4_0 *) data; int64_t ne0 = t->ne[0]; int64_t ne1 = t->ne[1]; @@ -513,46 +690,49 @@ static void repack_q4_0_tiled(ggml_tensor * t, const void * data, size_t size) { const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_Q4_0; const size_t matrix_size = n_col_tiles * n_k_tiles * tile_size; - for (int i3 = 0; i3 < ne3; i3++) { - for (int i2 = 0; i2 < ne2; i2++) { - const block_q4_0 * src_expert = src_matrix + (i3 * ne2 + i2) * (ne1 * (ne0 / 32)); - uint8_t * matrix_dst = (uint8_t *) t->data + (i3 * ne2 + i2) * matrix_size; + size_t slice_size = ne1 * ggml_row_size(t->type, ne0); + int64_t start_slice = offset / slice_size; + int64_t end_slice = (offset + size + slice_size - 1) / slice_size; + if (end_slice > ne2 * ne3) { + end_slice = ne2 * ne3; + } - for (int ct = 0; ct < n_col_tiles; ct++) { - for (int kt = 0; kt < n_k_tiles; kt++) { - uint8_t * tile_dst = matrix_dst + (ct * n_k_tiles + kt) * tile_size; + for (int64_t slice_idx = start_slice; slice_idx < end_slice; slice_idx++) { + const block_q4_0 * src_slice = src_matrix + (slice_idx - start_slice) * (ne1 * (ne0 / 32)); + uint8_t * matrix_dst = (uint8_t *) t->data + slice_idx * matrix_size; - uint8_t tile_quants[32][32]; - for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - if (r < ne1 && kt < ne0 / 32) { - unpack_q4_0_quants(tile_quants[row], &src_expert[r * (ne0 / 32) + kt], 0); - } else { - memset(tile_quants[row], 8, 32); - } - } + for (int ct = 0; ct < n_col_tiles; ct++) { + for (int kt = 0; kt < n_k_tiles; kt++) { + uint8_t * tile_dst = matrix_dst + (ct * n_k_tiles + kt) * tile_size; - for (int cp = 0; cp < 16; cp++) { - for (int row = 0; row < 32; row++) { - tile_dst[cp * 32 + row] = (tile_quants[row][2 * cp + 1] << 4) | tile_quants[row][2 * cp]; - } + uint8_t tile_quants[32][32]; + for (int row = 0; row < 32; row++) { + int64_t r = ct * 32 + row; + if (r < ne1 && kt < ne0 / 32) { + unpack_q4_0_quants(tile_quants[row], &src_slice[r * (ne0 / 32) + kt], 0); + } else { + memset(tile_quants[row], 8, 32); } + } - ggml_half * scale_dst = (ggml_half *)(tile_dst + 512); + for (int cp = 0; cp < 16; cp++) { for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - scale_dst[row] = (r < ne1 && kt < ne0 / 32) ? src_expert[r * (ne0 / 32) + kt].d : 0; + tile_dst[cp * 32 + row] = (tile_quants[row][2 * cp + 1] << 4) | tile_quants[row][2 * cp]; } } + + ggml_half * scale_dst = (ggml_half *)(tile_dst + 512); + for (int row = 0; row < 32; row++) { + int64_t r = ct * 32 + row; + scale_dst[row] = (r < ne1 && kt < ne0 / 32) ? src_slice[r * (ne0 / 32) + kt].d : 0; + } } } } - - GGML_UNUSED(size); } // repack q4_0_tiled tensor into q4_0 data -static void repack_tiled_q4_0(void * data, const ggml_tensor * t, size_t size) { +static void repack_tiled_q4_0(void * data, const ggml_tensor * t, size_t offset, size_t size) { block_q4_0 * dst_matrix = (block_q4_0 *) data; int64_t ne0 = t->ne[0]; int64_t ne1 = t->ne[1]; @@ -566,48 +746,65 @@ static void repack_tiled_q4_0(void * data, const ggml_tensor * t, size_t size) { const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_Q4_0; const size_t matrix_size = n_col_tiles * n_k_tiles * tile_size; - for (int i3 = 0; i3 < ne3; i3++) { - for (int i2 = 0; i2 < ne2; i2++) { - block_q4_0 * dst_expert = dst_matrix + (i3 * ne2 + i2) * (ne1 * (ne0 / 32)); - const uint8_t * matrix_src = (const uint8_t *) t->data + (i3 * ne2 + i2) * matrix_size; - - for (int ct = 0; ct < n_col_tiles; ct++) { - for (int kt = 0; kt < n_k_tiles; kt++) { - const uint8_t * tile_src = matrix_src + (ct * n_k_tiles + kt) * tile_size; - - uint8_t tile_quants[32][32]; - for (int cp = 0; cp < 16; cp++) { - for (int row = 0; row < 32; row++) { - uint8_t val = tile_src[cp * 32 + row]; - tile_quants[row][2 * cp + 0] = val & 0x0F; - tile_quants[row][2 * cp + 1] = val >> 4; - } - } + size_t slice_size = ne1 * ggml_row_size(t->type, ne0); + size_t row_size_bytes = ggml_row_size(t->type, ne0); + int64_t start_slice = offset / slice_size; + int64_t end_slice = (offset + size + slice_size - 1) / slice_size; + if (end_slice > ne2 * ne3) { + end_slice = ne2 * ne3; + } + + for (int64_t slice_idx = start_slice; slice_idx < end_slice; slice_idx++) { + size_t cur_start_byte = (std::max)(offset, (size_t) slice_idx * slice_size); + size_t cur_end_byte = (std::min)(offset + size, (size_t) (slice_idx + 1) * slice_size); + size_t slice_offset_start = cur_start_byte - (size_t) slice_idx * slice_size; + size_t slice_offset_end = cur_end_byte - (size_t) slice_idx * slice_size; + + int64_t start_row = slice_offset_start / row_size_bytes; + int64_t end_row = (slice_offset_end + row_size_bytes - 1) / row_size_bytes; + end_row = (std::min)(end_row, ne1); + + int start_ct = start_row / 32; + int end_ct = (end_row + 31) / 32; + end_ct = (std::min)(end_ct, n_col_tiles); + + block_q4_0 * dst_slice = dst_matrix + (cur_start_byte - offset) / sizeof(block_q4_0); + const uint8_t * matrix_src = (const uint8_t *) t->data + slice_idx * matrix_size; + for (int ct = start_ct; ct < end_ct; ct++) { + for (int kt = 0; kt < n_k_tiles; kt++) { + const uint8_t * tile_src = matrix_src + (ct * n_k_tiles + kt) * tile_size; + + uint8_t tile_quants[32][32]; + for (int cp = 0; cp < 16; cp++) { for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - if (r < ne1 && kt < ne0 / 32) { - pack_q4_0_quants(&dst_expert[r * (ne0 / 32) + kt], tile_quants[row], 0); - } + uint8_t val = tile_src[cp * 32 + row]; + tile_quants[row][2 * cp + 0] = val & 0x0F; + tile_quants[row][2 * cp + 1] = val >> 4; } + } - const ggml_half * scale_src = (const ggml_half *)(tile_src + 512); - for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - if (r < ne1 && kt < ne0 / 32) { - dst_expert[r * (ne0 / 32) + kt].d = scale_src[row]; - } + for (int row = 0; row < 32; row++) { + int64_t r = ct * 32 + row; + if (r >= start_row && r < end_row && kt < ne0 / 32) { + pack_q4_0_quants(&dst_slice[(r - start_row) * (ne0 / 32) + kt], tile_quants[row], 0); + } + } + + const ggml_half * scale_src = (const ggml_half *)(tile_src + 512); + for (int row = 0; row < 32; row++) { + int64_t r = ct * 32 + row; + if (r >= start_row && r < end_row && kt < ne0 / 32) { + dst_slice[(r - start_row) * (ne0 / 32) + kt].d = scale_src[row]; } } } } } - - GGML_UNUSED(size); } // repack q4_1 data into q4_1_tiled tensor -static void repack_q4_1_tiled(ggml_tensor * t, const void * data, size_t size) { +static void repack_q4_1_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) { const block_q4_1 * src_matrix = (const block_q4_1 *) data; int64_t ne0 = t->ne[0]; int64_t ne1 = t->ne[1]; @@ -621,52 +818,55 @@ static void repack_q4_1_tiled(ggml_tensor * t, const void * data, size_t size) { const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_Q4_1; const size_t matrix_size = n_col_tiles * n_k_tiles * tile_size; - for (int i3 = 0; i3 < ne3; i3++) { - for (int i2 = 0; i2 < ne2; i2++) { - const block_q4_1 * src_expert = src_matrix + (i3 * ne2 + i2) * (ne1 * (ne0 / 32)); - uint8_t * matrix_dst = (uint8_t *) t->data + (i3 * ne2 + i2) * matrix_size; + size_t slice_size = ne1 * ggml_row_size(t->type, ne0); + int64_t start_slice = offset / slice_size; + int64_t end_slice = (offset + size + slice_size - 1) / slice_size; + if (end_slice > ne2 * ne3) { + end_slice = ne2 * ne3; + } - for (int ct = 0; ct < n_col_tiles; ct++) { - for (int kt = 0; kt < n_k_tiles; kt++) { - uint8_t * tile_dst = matrix_dst + (ct * n_k_tiles + kt) * tile_size; + for (int64_t slice_idx = start_slice; slice_idx < end_slice; slice_idx++) { + const block_q4_1 * src_slice = src_matrix + (slice_idx - start_slice) * (ne1 * (ne0 / 32)); + uint8_t * matrix_dst = (uint8_t *) t->data + slice_idx * matrix_size; - uint8_t tile_quants[32][32]; - for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - if (r < ne1 && kt < ne0 / 32) { - unpack_q4_1_quants(tile_quants[row], &src_expert[r * (ne0 / 32) + kt], 0); - } else { - memset(tile_quants[row], 0, 32); - } - } + for (int ct = 0; ct < n_col_tiles; ct++) { + for (int kt = 0; kt < n_k_tiles; kt++) { + uint8_t * tile_dst = matrix_dst + (ct * n_k_tiles + kt) * tile_size; - for (int cp = 0; cp < 16; cp++) { - for (int row = 0; row < 32; row++) { - tile_dst[cp * 32 + row] = (tile_quants[row][2 * cp + 1] << 4) | tile_quants[row][2 * cp]; - } + uint8_t tile_quants[32][32]; + for (int row = 0; row < 32; row++) { + int64_t r = ct * 32 + row; + if (r < ne1 && kt < ne0 / 32) { + unpack_q4_1_quants(tile_quants[row], &src_slice[r * (ne0 / 32) + kt], 0); + } else { + memset(tile_quants[row], 0, 32); } + } - ggml_half * scale_dst = (ggml_half *)(tile_dst + 512); + for (int cp = 0; cp < 16; cp++) { for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - if (r < ne1 && kt < ne0 / 32) { - scale_dst[2 * row + 0] = src_expert[r * (ne0 / 32) + kt].d; - scale_dst[2 * row + 1] = src_expert[r * (ne0 / 32) + kt].m; - } else { - scale_dst[2 * row + 0] = 0; - scale_dst[2 * row + 1] = 0; - } + tile_dst[cp * 32 + row] = (tile_quants[row][2 * cp + 1] << 4) | tile_quants[row][2 * cp]; + } + } + + ggml_half * scale_dst = (ggml_half *)(tile_dst + 512); + for (int row = 0; row < 32; row++) { + int64_t r = ct * 32 + row; + if (r < ne1 && kt < ne0 / 32) { + scale_dst[2 * row + 0] = src_slice[r * (ne0 / 32) + kt].d; + scale_dst[2 * row + 1] = src_slice[r * (ne0 / 32) + kt].m; + } else { + scale_dst[2 * row + 0] = 0; + scale_dst[2 * row + 1] = 0; } } } } } - - GGML_UNUSED(size); } // repack q4_1_tiled tensor into q4_1 data -static void repack_tiled_q4_1(void * data, const ggml_tensor * t, size_t size) { +static void repack_tiled_q4_1(void * data, const ggml_tensor * t, size_t offset, size_t size) { block_q4_1 * dst_matrix = (block_q4_1 *) data; int64_t ne0 = t->ne[0]; int64_t ne1 = t->ne[1]; @@ -680,49 +880,66 @@ static void repack_tiled_q4_1(void * data, const ggml_tensor * t, size_t size) { const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_Q4_1; const size_t matrix_size = n_col_tiles * n_k_tiles * tile_size; - for (int i3 = 0; i3 < ne3; i3++) { - for (int i2 = 0; i2 < ne2; i2++) { - block_q4_1 * dst_expert = dst_matrix + (i3 * ne2 + i2) * (ne1 * (ne0 / 32)); - const uint8_t * matrix_src = (const uint8_t *) t->data + (i3 * ne2 + i2) * matrix_size; - - for (int ct = 0; ct < n_col_tiles; ct++) { - for (int kt = 0; kt < n_k_tiles; kt++) { - const uint8_t * tile_src = matrix_src + (ct * n_k_tiles + kt) * tile_size; - - uint8_t tile_quants[32][32]; - for (int cp = 0; cp < 16; cp++) { - for (int row = 0; row < 32; row++) { - uint8_t val = tile_src[cp * 32 + row]; - tile_quants[row][2 * cp + 0] = val & 0x0F; - tile_quants[row][2 * cp + 1] = val >> 4; - } - } + size_t slice_size = ne1 * ggml_row_size(t->type, ne0); + size_t row_size_bytes = ggml_row_size(t->type, ne0); + int64_t start_slice = offset / slice_size; + int64_t end_slice = (offset + size + slice_size - 1) / slice_size; + if (end_slice > ne2 * ne3) { + end_slice = ne2 * ne3; + } + for (int64_t slice_idx = start_slice; slice_idx < end_slice; slice_idx++) { + size_t cur_start_byte = (std::max)(offset, (size_t) slice_idx * slice_size); + size_t cur_end_byte = (std::min)(offset + size, (size_t) (slice_idx + 1) * slice_size); + size_t slice_offset_start = cur_start_byte - (size_t) slice_idx * slice_size; + size_t slice_offset_end = cur_end_byte - (size_t) slice_idx * slice_size; + + int64_t start_row = slice_offset_start / row_size_bytes; + int64_t end_row = (slice_offset_end + row_size_bytes - 1) / row_size_bytes; + end_row = (std::min)(end_row, ne1); + + int start_ct = start_row / 32; + int end_ct = (end_row + 31) / 32; + end_ct = (std::min)(end_ct, n_col_tiles); + + block_q4_1 * dst_slice = dst_matrix + (cur_start_byte - offset) / sizeof(block_q4_1); + const uint8_t * matrix_src = (const uint8_t *) t->data + slice_idx * matrix_size; + + for (int ct = start_ct; ct < end_ct; ct++) { + for (int kt = 0; kt < n_k_tiles; kt++) { + const uint8_t * tile_src = matrix_src + (ct * n_k_tiles + kt) * tile_size; + + uint8_t tile_quants[32][32]; + for (int cp = 0; cp < 16; cp++) { for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - if (r < ne1 && kt < ne0 / 32) { - pack_q4_1_quants(&dst_expert[r * (ne0 / 32) + kt], tile_quants[row], 0); - } + uint8_t val = tile_src[cp * 32 + row]; + tile_quants[row][2 * cp + 0] = val & 0x0F; + tile_quants[row][2 * cp + 1] = val >> 4; } + } - const ggml_half * scale_src = (const ggml_half *)(tile_src + 512); - for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - if (r < ne1 && kt < ne0 / 32) { - dst_expert[r * (ne0 / 32) + kt].d = scale_src[2 * row]; - dst_expert[r * (ne0 / 32) + kt].m = scale_src[2 * row + 1]; - } + for (int row = 0; row < 32; row++) { + int64_t r = ct * 32 + row; + if (r >= start_row && r < end_row && kt < ne0 / 32) { + pack_q4_1_quants(&dst_slice[(r - start_row) * (ne0 / 32) + kt], tile_quants[row], 0); + } + } + + const ggml_half * scale_src = (const ggml_half *)(tile_src + 512); + for (int row = 0; row < 32; row++) { + int64_t r = ct * 32 + row; + if (r >= start_row && r < end_row && kt < ne0 / 32) { + dst_slice[(r - start_row) * (ne0 / 32) + kt].d = scale_src[2 * row]; + dst_slice[(r - start_row) * (ne0 / 32) + kt].m = scale_src[2 * row + 1]; } } } } } - - GGML_UNUSED(size); } // repack q8_0 data into q8_0_tiled tensor -static void repack_q8_0_tiled(ggml_tensor * t, const void * data, size_t size) { +static void repack_q8_0_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) { const block_q8_0 * src_matrix = (const block_q8_0 *) data; int64_t ne0 = t->ne[0]; int64_t ne1 = t->ne[1]; @@ -736,41 +953,44 @@ static void repack_q8_0_tiled(ggml_tensor * t, const void * data, size_t size) { const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_Q8_0; const size_t matrix_size = n_col_tiles * n_k_tiles * tile_size; - for (int i3 = 0; i3 < ne3; i3++) { - for (int i2 = 0; i2 < ne2; i2++) { - const block_q8_0 * src_expert = src_matrix + (i3 * ne2 + i2) * (ne1 * (ne0 / 32)); - uint8_t * matrix_dst = (uint8_t *) t->data + (i3 * ne2 + i2) * matrix_size; - - for (int ct = 0; ct < n_col_tiles; ct++) { - for (int kt = 0; kt < n_k_tiles; kt++) { - uint8_t * tile_dst = matrix_dst + (ct * n_k_tiles + kt) * tile_size; - - for (int cp = 0; cp < 16; cp++) { - int col0 = cp * 2; - int col1 = col0 + 1; - for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - const block_q8_0 * b = (r < ne1 && kt < ne0 / 32) ? &src_expert[r * (ne0 / 32) + kt] : NULL; - tile_dst[cp * 64 + 2 * row + 0] = b ? b->qs[col0] : 0; - tile_dst[cp * 64 + 2 * row + 1] = b ? b->qs[col1] : 0; - } - } + size_t slice_size = ne1 * ggml_row_size(t->type, ne0); + int64_t start_slice = offset / slice_size; + int64_t end_slice = (offset + size + slice_size - 1) / slice_size; + if (end_slice > ne2 * ne3) { + end_slice = ne2 * ne3; + } + + for (int64_t slice_idx = start_slice; slice_idx < end_slice; slice_idx++) { + const block_q8_0 * src_slice = src_matrix + (slice_idx - start_slice) * (ne1 * (ne0 / 32)); + uint8_t * matrix_dst = (uint8_t *) t->data + slice_idx * matrix_size; - ggml_half * scale_dst = (ggml_half *)(tile_dst + 1024); + for (int ct = 0; ct < n_col_tiles; ct++) { + for (int kt = 0; kt < n_k_tiles; kt++) { + uint8_t * tile_dst = matrix_dst + (ct * n_k_tiles + kt) * tile_size; + + for (int cp = 0; cp < 16; cp++) { + int col0 = cp * 2; + int col1 = col0 + 1; for (int row = 0; row < 32; row++) { int64_t r = ct * 32 + row; - scale_dst[row] = (r < ne1 && kt < ne0 / 32) ? src_expert[r * (ne0 / 32) + kt].d : 0; + const block_q8_0 * b = (r < ne1 && kt < ne0 / 32) ? &src_slice[r * (ne0 / 32) + kt] : NULL; + tile_dst[cp * 64 + 2 * row + 0] = b ? b->qs[col0] : 0; + tile_dst[cp * 64 + 2 * row + 1] = b ? b->qs[col1] : 0; } } + + ggml_half * scale_dst = (ggml_half *)(tile_dst + 1024); + for (int row = 0; row < 32; row++) { + int64_t r = ct * 32 + row; + scale_dst[row] = (r < ne1 && kt < ne0 / 32) ? src_slice[r * (ne0 / 32) + kt].d : 0; + } } } } - - GGML_UNUSED(size); } // repack q8_0_tiled tensor into q8_0 data -static void repack_tiled_q8_0(void * data, const ggml_tensor * t, size_t size) { +static void repack_tiled_q8_0(void * data, const ggml_tensor * t, size_t offset, size_t size) { block_q8_0 * dst_matrix = (block_q8_0 *) data; int64_t ne0 = t->ne[0]; int64_t ne1 = t->ne[1]; @@ -784,45 +1004,62 @@ static void repack_tiled_q8_0(void * data, const ggml_tensor * t, size_t size) { const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_Q8_0; const size_t matrix_size = n_col_tiles * n_k_tiles * tile_size; - for (int i3 = 0; i3 < ne3; i3++) { - for (int i2 = 0; i2 < ne2; i2++) { - block_q8_0 * dst_expert = dst_matrix + (i3 * ne2 + i2) * (ne1 * (ne0 / 32)); - const uint8_t * matrix_src = (const uint8_t *) t->data + (i3 * ne2 + i2) * matrix_size; - - for (int ct = 0; ct < n_col_tiles; ct++) { - for (int kt = 0; kt < n_k_tiles; kt++) { - const uint8_t * tile_src = matrix_src + (ct * n_k_tiles + kt) * tile_size; - - for (int cp = 0; cp < 16; cp++) { - int col0 = cp * 2; - int col1 = col0 + 1; - for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - if (r < ne1 && kt < ne0 / 32) { - block_q8_0 & b = dst_expert[r * (ne0 / 32) + kt]; - b.qs[col0] = tile_src[cp * 64 + 2 * row + 0]; - b.qs[col1] = tile_src[cp * 64 + 2 * row + 1]; - } - } - } + size_t slice_size = ne1 * ggml_row_size(t->type, ne0); + size_t row_size_bytes = ggml_row_size(t->type, ne0); + int64_t start_slice = offset / slice_size; + int64_t end_slice = (offset + size + slice_size - 1) / slice_size; + if (end_slice > ne2 * ne3) { + end_slice = ne2 * ne3; + } + + for (int64_t slice_idx = start_slice; slice_idx < end_slice; slice_idx++) { + size_t cur_start_byte = (std::max)(offset, (size_t) slice_idx * slice_size); + size_t cur_end_byte = (std::min)(offset + size, (size_t) (slice_idx + 1) * slice_size); + size_t slice_offset_start = cur_start_byte - (size_t) slice_idx * slice_size; + size_t slice_offset_end = cur_end_byte - (size_t) slice_idx * slice_size; + + int64_t start_row = slice_offset_start / row_size_bytes; + int64_t end_row = (slice_offset_end + row_size_bytes - 1) / row_size_bytes; + end_row = (std::min)(end_row, ne1); + + int start_ct = start_row / 32; + int end_ct = (end_row + 31) / 32; + end_ct = (std::min)(end_ct, n_col_tiles); + + block_q8_0 * dst_slice = dst_matrix + (cur_start_byte - offset) / sizeof(block_q8_0); + const uint8_t * matrix_src = (const uint8_t *) t->data + slice_idx * matrix_size; - const ggml_half * scale_src = (const ggml_half *)(tile_src + 1024); + for (int ct = start_ct; ct < end_ct; ct++) { + for (int kt = 0; kt < n_k_tiles; kt++) { + const uint8_t * tile_src = matrix_src + (ct * n_k_tiles + kt) * tile_size; + + for (int cp = 0; cp < 16; cp++) { + int col0 = cp * 2; + int col1 = col0 + 1; for (int row = 0; row < 32; row++) { int64_t r = ct * 32 + row; - if (r < ne1 && kt < ne0 / 32) { - dst_expert[r * (ne0 / 32) + kt].d = scale_src[row]; + if (r >= start_row && r < end_row && kt < ne0 / 32) { + block_q8_0 & b = dst_slice[(r - start_row) * (ne0 / 32) + kt]; + b.qs[col0] = tile_src[cp * 64 + 2 * row + 0]; + b.qs[col1] = tile_src[cp * 64 + 2 * row + 1]; } } } + + const ggml_half * scale_src = (const ggml_half *)(tile_src + 1024); + for (int row = 0; row < 32; row++) { + int64_t r = ct * 32 + row; + if (r >= start_row && r < end_row && kt < ne0 / 32) { + dst_slice[(r - start_row) * (ne0 / 32) + kt].d = scale_src[row]; + } + } } } } - - GGML_UNUSED(size); } // repack mxfp4 data into mxfp4_tiled tensor -static void repack_mxfp4_tiled(ggml_tensor * t, const void * data, size_t size) { +static void repack_mxfp4_tiled(ggml_tensor * t, const void * data, size_t offset, size_t size) { const block_mxfp4 * src_matrix = (const block_mxfp4 *) data; int64_t ne0 = t->ne[0]; int64_t ne1 = t->ne[1]; @@ -836,46 +1073,49 @@ static void repack_mxfp4_tiled(ggml_tensor * t, const void * data, size_t size) const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_MXFP4; const size_t matrix_size = n_col_tiles * n_k_tiles * tile_size; - for (int i3 = 0; i3 < ne3; i3++) { - for (int i2 = 0; i2 < ne2; i2++) { - const block_mxfp4 * src_expert = src_matrix + (i3 * ne2 + i2) * (ne1 * (ne0 / 32)); - uint8_t * matrix_dst = (uint8_t *) t->data + (i3 * ne2 + i2) * matrix_size; + size_t slice_size = ne1 * ggml_row_size(t->type, ne0); + int64_t start_slice = offset / slice_size; + int64_t end_slice = (offset + size + slice_size - 1) / slice_size; + if (end_slice > ne2 * ne3) { + end_slice = ne2 * ne3; + } - for (int ct = 0; ct < n_col_tiles; ct++) { - for (int kt = 0; kt < n_k_tiles; kt++) { - uint8_t * tile_dst = matrix_dst + (ct * n_k_tiles + kt) * tile_size; + for (int64_t slice_idx = start_slice; slice_idx < end_slice; slice_idx++) { + const block_mxfp4 * src_slice = src_matrix + (slice_idx - start_slice) * (ne1 * (ne0 / 32)); + uint8_t * matrix_dst = (uint8_t *) t->data + slice_idx * matrix_size; - uint8_t tile_quants[32][32]; - for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - if (r < ne1 && kt < ne0 / 32) { - unpack_mxfp4_quants(tile_quants[row], &src_expert[r * (ne0 / 32) + kt], 0); - } else { - memset(tile_quants[row], 0, 32); - } - } + for (int ct = 0; ct < n_col_tiles; ct++) { + for (int kt = 0; kt < n_k_tiles; kt++) { + uint8_t * tile_dst = matrix_dst + (ct * n_k_tiles + kt) * tile_size; - for (int cp = 0; cp < 16; cp++) { - for (int row = 0; row < 32; row++) { - tile_dst[cp * 32 + row] = (tile_quants[row][2 * cp + 1] << 4) | tile_quants[row][2 * cp]; - } + uint8_t tile_quants[32][32]; + for (int row = 0; row < 32; row++) { + int64_t r = ct * 32 + row; + if (r < ne1 && kt < ne0 / 32) { + unpack_mxfp4_quants(tile_quants[row], &src_slice[r * (ne0 / 32) + kt], 0); + } else { + memset(tile_quants[row], 0, 32); } + } - uint8_t * scale_dst = tile_dst + 512; + for (int cp = 0; cp < 16; cp++) { for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - scale_dst[row] = (r < ne1 && kt < ne0 / 32) ? src_expert[r * (ne0 / 32) + kt].e : 0; + tile_dst[cp * 32 + row] = (tile_quants[row][2 * cp + 1] << 4) | tile_quants[row][2 * cp]; } } + + uint8_t * scale_dst = tile_dst + 512; + for (int row = 0; row < 32; row++) { + int64_t r = ct * 32 + row; + scale_dst[row] = (r < ne1 && kt < ne0 / 32) ? src_slice[r * (ne0 / 32) + kt].e : 0; + } } } } - - GGML_UNUSED(size); } // repack mxfp4_tiled tensor into mxfp4 data -static void repack_tiled_mxfp4(void * data, const ggml_tensor * t, size_t size) { +static void repack_tiled_mxfp4(void * data, const ggml_tensor * t, size_t offset, size_t size) { block_mxfp4 * dst_matrix = (block_mxfp4 *) data; int64_t ne0 = t->ne[0]; int64_t ne1 = t->ne[1]; @@ -889,133 +1129,179 @@ static void repack_tiled_mxfp4(void * data, const ggml_tensor * t, size_t size) const size_t tile_size = HTP_MM_WEIGHT_TILE_SIZE_MXFP4; const size_t matrix_size = n_col_tiles * n_k_tiles * tile_size; - for (int i3 = 0; i3 < ne3; i3++) { - for (int i2 = 0; i2 < ne2; i2++) { - block_mxfp4 * dst_expert = dst_matrix + (i3 * ne2 + i2) * (ne1 * (ne0 / 32)); - const uint8_t * matrix_src = (const uint8_t *) t->data + (i3 * ne2 + i2) * matrix_size; - - for (int ct = 0; ct < n_col_tiles; ct++) { - for (int kt = 0; kt < n_k_tiles; kt++) { - const uint8_t * tile_src = matrix_src + (ct * n_k_tiles + kt) * tile_size; - - uint8_t tile_quants[32][32]; - for (int cp = 0; cp < 16; cp++) { - for (int row = 0; row < 32; row++) { - uint8_t val = tile_src[cp * 32 + row]; - tile_quants[row][2 * cp + 0] = val & 0x0F; - tile_quants[row][2 * cp + 1] = val >> 4; - } - } + size_t slice_size = ne1 * ggml_row_size(t->type, ne0); + size_t row_size_bytes = ggml_row_size(t->type, ne0); + int64_t start_slice = offset / slice_size; + int64_t end_slice = (offset + size + slice_size - 1) / slice_size; + if (end_slice > ne2 * ne3) { + end_slice = ne2 * ne3; + } + + for (int64_t slice_idx = start_slice; slice_idx < end_slice; slice_idx++) { + size_t cur_start_byte = (std::max)(offset, (size_t) slice_idx * slice_size); + size_t cur_end_byte = (std::min)(offset + size, (size_t) (slice_idx + 1) * slice_size); + size_t slice_offset_start = cur_start_byte - (size_t) slice_idx * slice_size; + size_t slice_offset_end = cur_end_byte - (size_t) slice_idx * slice_size; + + int64_t start_row = slice_offset_start / row_size_bytes; + int64_t end_row = (slice_offset_end + row_size_bytes - 1) / row_size_bytes; + end_row = (std::min)(end_row, ne1); + + int start_ct = start_row / 32; + int end_ct = (end_row + 31) / 32; + end_ct = (std::min)(end_ct, n_col_tiles); + block_mxfp4 * dst_slice = dst_matrix + (cur_start_byte - offset) / sizeof(block_mxfp4); + const uint8_t * matrix_src = (const uint8_t *) t->data + slice_idx * matrix_size; + + for (int ct = start_ct; ct < end_ct; ct++) { + for (int kt = 0; kt < n_k_tiles; kt++) { + const uint8_t * tile_src = matrix_src + (ct * n_k_tiles + kt) * tile_size; + + uint8_t tile_quants[32][32]; + for (int cp = 0; cp < 16; cp++) { for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - if (r < ne1 && kt < ne0 / 32) { - pack_mxfp4_quants(&dst_expert[r * (ne0 / 32) + kt], tile_quants[row], 0); - } + uint8_t val = tile_src[cp * 32 + row]; + tile_quants[row][2 * cp + 0] = val & 0x0F; + tile_quants[row][2 * cp + 1] = val >> 4; } + } - const uint8_t * scale_src = tile_src + 512; - for (int row = 0; row < 32; row++) { - int64_t r = ct * 32 + row; - if (r < ne1 && kt < ne0 / 32) { - dst_expert[r * (ne0 / 32) + kt].e = scale_src[row]; - } + for (int row = 0; row < 32; row++) { + int64_t r = ct * 32 + row; + if (r >= start_row && r < end_row && kt < ne0 / 32) { + pack_mxfp4_quants(&dst_slice[(r - start_row) * (ne0 / 32) + kt], tile_quants[row], 0); + } + } + + const uint8_t * scale_src = tile_src + 512; + for (int row = 0; row < 32; row++) { + int64_t r = ct * 32 + row; + if (r >= start_row && r < end_row && kt < ne0 / 32) { + dst_slice[(r - start_row) * (ne0 / 32) + kt].e = scale_src[row]; } } } } } - - GGML_UNUSED(size); } -static void ggml_backend_hexagon_buffer_set_tensor(ggml_backend_buffer_t buffer, - ggml_tensor * tensor, - const void * data, - size_t offset, - size_t size) { - auto sbuf = (ggml_hexagon_shared_buffer *) buffer->context; - auto sess = sbuf->sess; - - HEX_VERBOSE("ggml-hex: %s set-tensor %s : data %p offset %zu size %zu\n", sess->c_name(), tensor->name, data, offset, size); - +static void repack_tensor_tiled(ggml_tensor * tensor, const void * data, size_t size) { switch (tensor->type) { case GGML_TYPE_Q4_0: - GGML_ASSERT(offset == 0); - GGML_ASSERT(offset + size <= ggml_nbytes(tensor)); - repack_q4_0_tiled(tensor, data, size); + repack_q4_0_tiled(tensor, data, 0, size); break; case GGML_TYPE_Q4_1: - GGML_ASSERT(offset == 0); - GGML_ASSERT(offset + size <= ggml_nbytes(tensor)); - repack_q4_1_tiled(tensor, data, size); + repack_q4_1_tiled(tensor, data, 0, size); break; case GGML_TYPE_Q8_0: - GGML_ASSERT(offset == 0); - GGML_ASSERT(offset + size <= ggml_nbytes(tensor)); - repack_q8_0_tiled(tensor, data, size); + repack_q8_0_tiled(tensor, data, 0, size); break; case GGML_TYPE_IQ4_NL: - GGML_ASSERT(offset == 0); - GGML_ASSERT(offset + size <= ggml_nbytes(tensor)); - // IQ4_NL has identical block layout to Q4_0 (ggml_half d + uint8_t qs[16]) - repack_q4_0_tiled(tensor, data, size); + repack_q4_0_tiled(tensor, data, 0, size); break; case GGML_TYPE_MXFP4: - GGML_ASSERT(offset == 0); - GGML_ASSERT(offset + size <= ggml_nbytes(tensor)); - repack_mxfp4_tiled(tensor, data, size); + repack_mxfp4_tiled(tensor, data, 0, size); break; default: - memcpy((char *) tensor->data + offset, data, size); break; } } +static void ggml_backend_hexagon_buffer_set_tensor(ggml_backend_buffer_t buffer, + ggml_tensor * tensor, + const void * data, + size_t offset, + size_t size) { + auto extra = (ggml_hexagon_tensor_extra *) tensor->extra; + auto sbuf = (ggml_hexagon_shared_buffer *) buffer->context; + auto sess = sbuf->sess; + + if (ggml_backend_buffer_get_usage(buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS) { + extra->flags |= GGML_HEXAGON_TENSOR_WEIGHT; + if (ggml_hexagon_is_repack_type(tensor->type)) { + extra->flags |= GGML_HEXAGON_TENSOR_REPACK; + } + } + + HEX_VERBOSE("ggml-hex: %s set-tensor %s : data %p offset %zu size %zu usage %d flags 0x%x\n", + sess->c_name(), tensor->name, data, offset, size, (int) buffer->usage, extra->flags); + + if ((extra->flags & GGML_HEXAGON_TENSOR_REPACK) == 0) { + memcpy((char *) tensor->data + offset, data, size); + return; + } + + if (offset == 0 && size == ggml_nbytes(tensor) && extra->shadow_buf.empty()) { + repack_tensor_tiled(tensor, data, size); + return; + } + + if (extra->shadow_buf.size() < ggml_nbytes(tensor)) { + extra->shadow_buf.resize(ggml_nbytes(tensor)); + } + memcpy(extra->shadow_buf.data() + offset, data, size); + extra->shadow_size += size; + + if (extra->shadow_size >= ggml_nbytes(tensor)) { + repack_tensor_tiled(tensor, extra->shadow_buf.data(), extra->shadow_buf.size()); + extra->shadow_buf.clear(); + extra->shadow_buf.shrink_to_fit(); + extra->shadow_size = 0; + } +} + static void ggml_backend_hexagon_buffer_get_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * tensor, void * data, size_t offset, size_t size) { - auto sbuf = (ggml_hexagon_shared_buffer *) buffer->context; - auto sess = sbuf->sess; + auto extra = (ggml_hexagon_tensor_extra *) tensor->extra; + auto sbuf = (ggml_hexagon_shared_buffer *) buffer->context; + auto sess = sbuf->sess; - HEX_VERBOSE("ggml-hex: %s get-tensor %s : data %p offset %zu size %zu\n", sess->c_name(), tensor->name, data, offset, size); + HEX_VERBOSE("ggml-hex: %s get-tensor %s : data %p offset %zu size %zu usage %d flags 0x%x\n", + sess->c_name(), tensor->name, data, offset, size, (int) buffer->usage, extra->flags); + + if ((extra->flags & GGML_HEXAGON_TENSOR_REPACK) == 0) { + memcpy(data, (const char *) tensor->data + offset, size); + return; + } switch (tensor->type) { case GGML_TYPE_Q4_0: GGML_ASSERT(offset == 0); GGML_ASSERT(offset + size <= ggml_nbytes(tensor)); - repack_tiled_q4_0(data, tensor, size); + repack_tiled_q4_0(data, tensor, offset, size); break; case GGML_TYPE_Q4_1: GGML_ASSERT(offset == 0); GGML_ASSERT(offset + size <= ggml_nbytes(tensor)); - repack_tiled_q4_1(data, tensor, size); + repack_tiled_q4_1(data, tensor, offset, size); break; case GGML_TYPE_Q8_0: GGML_ASSERT(offset == 0); GGML_ASSERT(offset + size <= ggml_nbytes(tensor)); - repack_tiled_q8_0(data, tensor, size); + repack_tiled_q8_0(data, tensor, offset, size); break; case GGML_TYPE_IQ4_NL: GGML_ASSERT(offset == 0); GGML_ASSERT(offset + size <= ggml_nbytes(tensor)); - repack_tiled_q4_0(data, tensor, size); + repack_tiled_q4_0(data, tensor, offset, size); break; case GGML_TYPE_MXFP4: GGML_ASSERT(offset == 0); GGML_ASSERT(offset + size <= ggml_nbytes(tensor)); - repack_tiled_mxfp4(data, tensor, size); + repack_tiled_mxfp4(data, tensor, offset, size); break; default: @@ -1035,11 +1321,121 @@ static bool ggml_backend_hexagon_buffer_cpy_tensor(ggml_backend_buffer_t bu GGML_UNUSED(dst); } +static void ggml_backend_hexagon_buffer_set_tensor_2d(ggml_backend_buffer_t buffer, + ggml_tensor * tensor, + const void * data, + size_t offset, + size_t size, + size_t n_copies, + size_t stride_tensor, + size_t stride_data) { + auto extra = (ggml_hexagon_tensor_extra *) tensor->extra; + auto sbuf = (ggml_hexagon_shared_buffer *) buffer->context; + auto sess = sbuf->sess; + + if (ggml_backend_buffer_get_usage(buffer) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS) { + extra->flags |= GGML_HEXAGON_TENSOR_WEIGHT; + if (ggml_hexagon_is_repack_type(tensor->type)) { + extra->flags |= GGML_HEXAGON_TENSOR_REPACK; + } + } + + HEX_VERBOSE("ggml-hex: %s set-tensor-2d %s : data %p offset %zu size %zu n_copies %zu stride_tensor %zu stride_data %zu usage %d flags 0x%x\n", + sess->c_name(), tensor->name, data, offset, size, n_copies, stride_tensor, stride_data, (int) buffer->usage, extra->flags); + + if ((extra->flags & GGML_HEXAGON_TENSOR_REPACK) == 0) { + for (size_t i = 0; i < n_copies; i++) { + memcpy((uint8_t *) tensor->data + offset + i * stride_tensor, (const uint8_t *) data + i * stride_data, size); + } + return; + } + + if (extra->shadow_buf.size() < ggml_nbytes(tensor)) { + extra->shadow_buf.resize(ggml_nbytes(tensor)); + } + for (size_t i = 0; i < n_copies; i++) { + memcpy(extra->shadow_buf.data() + offset + i * stride_tensor, (const uint8_t *) data + i * stride_data, size); + } + extra->shadow_size += n_copies * size; + + if (extra->shadow_size >= ggml_nbytes(tensor)) { + repack_tensor_tiled(tensor, extra->shadow_buf.data(), extra->shadow_buf.size()); + extra->shadow_buf.clear(); + extra->shadow_buf.shrink_to_fit(); + extra->shadow_size = 0; + } +} + +static void ggml_backend_hexagon_buffer_get_tensor_2d(ggml_backend_buffer_t buffer, + const ggml_tensor * tensor, + void * data, + size_t offset, + size_t size, + size_t n_copies, + size_t stride_tensor, + size_t stride_data) { + auto extra = (ggml_hexagon_tensor_extra *) tensor->extra; + auto sbuf = (ggml_hexagon_shared_buffer *) buffer->context; + auto sess = sbuf->sess; + + HEX_VERBOSE("ggml-hex: %s get-tensor-2d %s : data %p offset %zu size %zu n_copies %zu stride_tensor %zu stride_data %zu usage %d\n", + sess->c_name(), tensor->name, data, offset, size, n_copies, stride_tensor, stride_data, (int) buffer->usage); + + if ((extra->flags & GGML_HEXAGON_TENSOR_REPACK) == 0) { + for (size_t i = 0; i < n_copies; i++) { + memcpy((uint8_t *)data + i * stride_data, (const uint8_t *)tensor->data + offset + i * stride_tensor, size); + } + return; + } + + size_t temp_size = n_copies > 0 ? (n_copies - 1) * stride_tensor + size : 0; + size_t slice_size = tensor->ne[1] * ggml_row_size(tensor->type, tensor->ne[0]); + size_t slice_offset = offset % slice_size; + size_t row_size_bytes = ggml_row_size(tensor->type, tensor->ne[0]); + + GGML_ASSERT((slice_offset % row_size_bytes) == 0 && "offset must be aligned to row boundary"); + GGML_ASSERT((temp_size % row_size_bytes) == 0 && "temp_size must be a multiple of row size"); + GGML_ASSERT((slice_offset / row_size_bytes) % 32 == 0 && "offset must be aligned to tile size (32 rows)"); + GGML_ASSERT((offset + temp_size) <= ggml_nbytes(tensor)); + + std::vector temp_buf(temp_size); + + switch (tensor->type) { + case GGML_TYPE_Q4_0: + repack_tiled_q4_0(temp_buf.data(), tensor, offset, temp_size); + break; + + case GGML_TYPE_Q4_1: + repack_tiled_q4_1(temp_buf.data(), tensor, offset, temp_size); + break; + + case GGML_TYPE_Q8_0: + repack_tiled_q8_0(temp_buf.data(), tensor, offset, temp_size); + break; + + case GGML_TYPE_IQ4_NL: + repack_tiled_q4_0(temp_buf.data(), tensor, offset, temp_size); + break; + + case GGML_TYPE_MXFP4: + repack_tiled_mxfp4(temp_buf.data(), tensor, offset, temp_size); + break; + + default: + memcpy(temp_buf.data(), (const uint8_t *) tensor->data + offset, temp_size); + break; + } + + for (size_t i = 0; i < n_copies; i++) { + memcpy((uint8_t *) data + i * stride_data, temp_buf.data() + i * stride_tensor, size); + } +} + static void ggml_backend_hexagon_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { auto sbuf = (ggml_hexagon_shared_buffer *) buffer->context; auto sess = sbuf->sess; - HEX_VERBOSE("ggml-hex: %s clear-buff base %p size %zu\n", sess->c_name(), (void *) sbuf->base, sbuf->size); - memset(sbuf->base, value, sbuf->size); + HEX_VERBOSE("ggml-hex: %s clear-buff base %p size %zu\n", sess->c_name(), (void *) sbuf->base(), sbuf->size()); + memset(sbuf->base(), value, sbuf->size()); } static ggml_backend_buffer_i ggml_backend_hexagon_buffer_interface = { @@ -1049,6 +1445,40 @@ static ggml_backend_buffer_i ggml_backend_hexagon_buffer_interface = { /* .memset_tensor = */ NULL, /* .set_tensor = */ ggml_backend_hexagon_buffer_set_tensor, /* .get_tensor = */ ggml_backend_hexagon_buffer_get_tensor, + /* .set_tensor_2d = */ ggml_backend_hexagon_buffer_set_tensor_2d, + /* .get_tensor_2d = */ ggml_backend_hexagon_buffer_get_tensor_2d, + /* .cpy_tensor = */ ggml_backend_hexagon_buffer_cpy_tensor, + /* .clear = */ ggml_backend_hexagon_buffer_clear, + /* .reset = */ NULL, +}; + +// ** backend buffer type + +static void ggml_backend_hexagon_host_buffer_set_tensor(ggml_backend_buffer_t buffer, + ggml_tensor * tensor, + const void * data, + size_t offset, + size_t size) { + memcpy((char *) tensor->data + offset, data, size); + GGML_UNUSED(buffer); +} + +static void ggml_backend_hexagon_host_buffer_get_tensor(ggml_backend_buffer_t buffer, + const ggml_tensor * tensor, + void * data, + size_t offset, + size_t size) { + memcpy(data, (const char *) tensor->data + offset, size); + GGML_UNUSED(buffer); +} + +static ggml_backend_buffer_i ggml_backend_hexagon_host_buffer_interface = { + /* .free_buffer = */ ggml_backend_hexagon_buffer_free_buffer, + /* .get_base = */ ggml_backend_hexagon_buffer_get_base, + /* .init_tensor = */ ggml_backend_hexagon_buffer_init_tensor, + /* .memset_tensor = */ NULL, + /* .set_tensor = */ ggml_backend_hexagon_host_buffer_set_tensor, + /* .get_tensor = */ ggml_backend_hexagon_host_buffer_get_tensor, /* .set_tensor_2d = */ NULL, /* .get_tensor_2d = */ NULL, /* .cpy_tensor = */ ggml_backend_hexagon_buffer_cpy_tensor, @@ -1066,24 +1496,22 @@ static ggml_backend_buffer_t ggml_backend_hexagon_buffer_type_alloc_buffer( ggml_backend_buffer_type_t buffer_type, size_t size) { auto sess = static_cast(buffer_type->context)->sess; try { - size += 4 * 1024; // guard page - ggml_hexagon_shared_buffer * sbuf = new ggml_hexagon_shared_buffer(sess, size); + ggml_hexagon_shared_buffer * sbuf = new ggml_hexagon_shared_buffer(sess, size, false, GGML_HEXAGON_FENCE_BUFFER_SIZE); return ggml_backend_buffer_init(buffer_type, ggml_backend_hexagon_buffer_interface, sbuf, size); } catch (const std::exception & exc) { - GGML_LOG_ERROR("ggml-hex: %s failed to allocate buffer context (host): %s\n", sess->c_name(), exc.what()); + GGML_LOG_ERROR("ggml-hex: %s failed to allocate device buffer context: %s\n", sess->c_name(), exc.what()); return nullptr; } } -static ggml_backend_buffer_t ggml_backend_hexagon_repack_buffer_type_alloc_buffer( +static ggml_backend_buffer_t ggml_backend_hexagon_host_buffer_type_alloc_buffer( ggml_backend_buffer_type_t buffer_type, size_t size) { auto sess = static_cast(buffer_type->context)->sess; try { - size += 4 * 1024; // guard page - ggml_hexagon_shared_buffer * sbuf = new ggml_hexagon_shared_buffer(sess, size); - return ggml_backend_buffer_init(buffer_type, ggml_backend_hexagon_buffer_interface, sbuf, size); + ggml_hexagon_shared_buffer * sbuf = new ggml_hexagon_shared_buffer(sess, size, false, GGML_HEXAGON_FENCE_BUFFER_SIZE); + return ggml_backend_buffer_init(buffer_type, ggml_backend_hexagon_host_buffer_interface, sbuf, size); } catch (const std::exception & exc) { - GGML_LOG_ERROR("ggml-hex: %s failed to allocate buffer context (repack): %s\n", sess->c_name(), exc.what()); + GGML_LOG_ERROR("ggml-hex: %s failed to allocate host buffer context: %s\n", sess->c_name(), exc.what()); return nullptr; } } @@ -1094,7 +1522,7 @@ static size_t ggml_backend_hexagon_buffer_type_get_alignment(ggml_backend_buffer } static size_t ggml_backend_hexagon_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const struct ggml_tensor * t) { - if (t->type == GGML_TYPE_Q4_0 || t->type == GGML_TYPE_Q4_1 || t->type == GGML_TYPE_Q8_0 || t->type == GGML_TYPE_IQ4_NL || t->type == GGML_TYPE_MXFP4) { + if (ggml_hexagon_is_repack_type(t->type)) { int64_t ne0 = hex_round_up(t->ne[0], 32); int64_t ne1 = hex_round_up(t->ne[1], 32); int64_t ne2 = t->ne[2]; @@ -1112,14 +1540,12 @@ static size_t ggml_backend_hexagon_buffer_type_get_max_size(ggml_backend_buffer_ } static bool ggml_backend_hexagon_buffer_type_is_host(ggml_backend_buffer_type_t buft) { - return opt_hostbuf; - + return false; GGML_UNUSED(buft); } -static bool ggml_backend_hexagon_repack_buffer_type_is_host(ggml_backend_buffer_type_t buft) { - return false; - +static bool ggml_backend_hexagon_host_buffer_type_is_host(ggml_backend_buffer_type_t buft) { + return true; GGML_UNUSED(buft); } @@ -1132,26 +1558,19 @@ static ggml_backend_buffer_type_i ggml_backend_hexagon_buffer_type_interface = { /* .is_host = */ ggml_backend_hexagon_buffer_type_is_host, }; -static ggml_backend_buffer_type_i ggml_backend_hexagon_repack_buffer_type_interface = { +static ggml_backend_buffer_type_i ggml_backend_hexagon_host_buffer_type_interface = { /* .get_name = */ ggml_backend_hexagon_buffer_type_name, - /* .alloc_buffer = */ ggml_backend_hexagon_repack_buffer_type_alloc_buffer, + /* .alloc_buffer = */ ggml_backend_hexagon_host_buffer_type_alloc_buffer, /* .get_alignment = */ ggml_backend_hexagon_buffer_type_get_alignment, /* .get_max_size = */ ggml_backend_hexagon_buffer_type_get_max_size, /* .get_alloc_size = */ ggml_backend_hexagon_buffer_type_get_alloc_size, - /* .is_host = */ ggml_backend_hexagon_repack_buffer_type_is_host, + /* .is_host = */ ggml_backend_hexagon_host_buffer_type_is_host, }; static bool ggml_backend_buffer_is_hexagon(const struct ggml_backend_buffer * b) { return b->buft->iface.get_alignment == ggml_backend_hexagon_buffer_type_get_alignment; } -static inline bool ggml_backend_buffer_is_hexagon_repack(const struct ggml_backend_buffer * b) { - if (!opt_hostbuf) { - return ggml_backend_buffer_is_hexagon(b); - } - return b->buft->iface.alloc_buffer == ggml_backend_hexagon_repack_buffer_type_alloc_buffer; -} - struct ggml_hexagon_opbatch { ggml_hexagon_session* sess; @@ -1165,8 +1584,6 @@ struct ggml_hexagon_opbatch { std::unordered_map t_map; // tensor ptr to index std::unordered_multimap d_map; // tensor data to index - - unsigned int n_bufs; // num buffers in the batch unsigned int n_tens; // num tensors ... unsigned int n_ops; // num ops ... @@ -1186,6 +1603,7 @@ struct ggml_hexagon_opbatch { b_map.clear(); t_map.clear(); d_map.clear(); + ops.resize(n_ops_max); } ggml_hexagon_opbatch(ggml_hexagon_session *sess, size_t batch_size, size_t max_vmem) { @@ -1218,39 +1636,39 @@ struct ggml_hexagon_opbatch { // add buffer and return its index int add_buffer(ggml_hexagon_shared_buffer * sbuf) { // Lookup by fd - auto it = b_map.find(sbuf->fd); + auto it = b_map.find(sbuf->fd()); if (it != b_map.end()) { return it->second; } // Add new buffer to the batch int bi = n_bufs++; GGML_ASSERT(n_bufs < HTP_OP_MAX_BUFS); - b_map.insert({sbuf->fd, bi}); + b_map.insert({sbuf->fd(), bi}); htp_buf_desc &b = h_bufs[bi]; - b.base = (uint64_t) sbuf->base; - b.fd = sbuf->fd; - b.size = sbuf->size; + b.base = (uint64_t) sbuf->base(); + b.fd = sbuf->fd(); + b.size = sbuf->size(); b_vmem += b.size; - HEX_VERBOSE("ggml-hex: %s add-buffer #%u : fd %d base %p size %zu : vmem %zu\n", sess->c_name(), bi, b.fd, (void*) sbuf->base, (size_t) b.size, b_vmem); + HEX_VERBOSE("ggml-hex: %s add-buffer #%u : fd %d base %p size %zu : vmem %zu\n", sess->c_name(), bi, b.fd, (void*) sbuf->base(), (size_t) b.size, b_vmem); return bi; } - - bool same_shape(const htp_tensor * h, const ggml_tensor * t) const { + auto extra = (ggml_hexagon_tensor_extra *) t->extra; + int64_t ne0 = t->ne[0]; int64_t ne1 = t->ne[1]; - const bool is_repack = ggml_backend_buffer_is_hexagon_repack(t->buffer) && ggml_hexagon_is_repack_type(t->type); + const bool is_repack = (extra->flags & GGML_HEXAGON_TENSOR_REPACK) != 0; if (is_repack) { ne0 = hex_round_up(ne0, 32); ne1 = hex_round_up(ne1, 32); } int64_t nb1 = is_repack ? ggml_row_size(t->type, ne0) : t->nb[1]; - int64_t nb2 = is_repack ? nb1 * ne1 : t->nb[2]; + int64_t nb2 = is_repack ? nb1 * ne1 : t->nb[2]; int64_t nb3 = is_repack ? nb2 * t->ne[2] : t->nb[3]; return (h->type == t->type) && @@ -1260,7 +1678,8 @@ struct ggml_hexagon_opbatch { // add tensor and return its index int add_tensor(const ggml_tensor * t) { - auto sbuf = static_cast(t->buffer->context); + auto extra = (ggml_hexagon_tensor_extra *) t->extra; + auto sbuf = static_cast(t->buffer->context); // First lookup by tensor data auto range = d_map.equal_range(t->data); @@ -1280,7 +1699,7 @@ struct ggml_hexagon_opbatch { t_map.insert({t, ti}); d_map.insert({t->data, ti}); - uint64_t t_offset = (uint8_t *) t->data - sbuf->base; + uint64_t t_offset = (uint8_t *) t->data - sbuf->base(); size_t t_size = ggml_nbytes(t); htp_tensor &h = h_tens[ti]; @@ -1289,7 +1708,7 @@ struct ggml_hexagon_opbatch { h.data = t_offset; h.type = t->type; - const bool is_repack = ggml_backend_buffer_is_hexagon_repack(t->buffer) && ggml_hexagon_is_repack_type(t->type); + const bool is_repack = (extra->flags & GGML_HEXAGON_TENSOR_REPACK) != 0; if (is_repack) { h.ne[0] = hex_round_up(t->ne[0], 32); h.ne[1] = hex_round_up(t->ne[1], 32); @@ -1308,11 +1727,15 @@ struct ggml_hexagon_opbatch { h.nb[0] = t->nb[0]; h.nb[1] = t->nb[1]; h.nb[2] = t->nb[2]; h.nb[3] = t->nb[3]; } - - h.flags = 0; - if (ggml_backend_buffer_get_usage(t->buffer) != GGML_BACKEND_BUFFER_USAGE_WEIGHTS) { - h.flags |= HTP_TENSOR_COMPUTE; + if ((extra->flags & GGML_HEXAGON_TENSOR_WEIGHT) != 0) { + h.flags |= HTP_TENSOR_WEIGHT; + } + if ((extra->flags & GGML_HEXAGON_TENSOR_REPACK) != 0) { + h.flags |= HTP_TENSOR_REPACK; + } + if ((extra->flags & GGML_HEXAGON_TENSOR_FENCE) != 0) { + h.flags |= HTP_TENSOR_FENCE; } HEX_VERBOSE("ggml-hex: %s add-tensor #%u %s : bi %d data %p offset %zu size %zu flags 0x%x : %zu:%zu:%zu:%zu\n", sess->c_name(), @@ -1336,8 +1759,8 @@ struct ggml_hexagon_opbatch { extra_tens++; auto sbuf = static_cast(t->buffer->context); - if (!b_map.count(sbuf->fd)) { - extra_vmem += sbuf->size; + if (!b_map.count(sbuf->fd())) { + extra_vmem += sbuf->size(); extra_bufs += 1; } } @@ -1372,32 +1795,497 @@ struct ggml_hexagon_opbatch { o.opcode = node.opcode; o.flags = 0; - if (!(opt_opstage & HTP_OPSTAGE_COMPUTE)) { - o.flags |= HTP_OPFLAGS_SKIP_COMPUTE; - } + ggml_hexagon_dump_op_exec(sess->c_name(), ops[n], o.flags); + + auto inputs = node.get_inputs(); + for (unsigned int i=0; i < HTP_OP_MAX_INPUTS; i++) { + o.src[i] = (i < inputs.size() && inputs[i]) ? add_tensor(inputs[i]) : 0xffff; + } + + auto outputs = node.get_outputs(); + for (unsigned int i=0; i < HTP_OP_MAX_OUTPUTS; i++) { + o.dst[i] = (i < outputs.size() && outputs[i]) ? add_tensor(outputs[i]) : 0xffff; + } + } + + bool try_fuse_allreduce_add(const htp_opnode & node) { + if (n_ops == 0 || opt_ar_select != 2) return false; + if (node.opcode != HTP_OP_ADD) return false; + + htp_opnode & last_node = ops[n_ops - 1]; + if (last_node.opcode != HTP_OP_ALLREDUCE) return false; + + auto * ar_kparams = (struct htp_allreduce_kernel_params *) last_node.kernel_params; + const uint32_t rank = (uint32_t) ar_kparams->rank; + const ggml_tensor * ar_local = (rank < last_node.inputs.size()) ? last_node.inputs[rank] : nullptr; + const ggml_tensor * add_src0 = node.src0(); + const ggml_tensor * add_src1 = node.src1(); + + if (!add_src0 || !add_src1 || !ar_local) return false; + if (!ggml_hexagon_tensor_is_fuseable(ar_local)) return false; + + const ggml_tensor * res_tensor = nullptr; + if (add_src0 == ar_local || add_src0->data == ar_local->data) { + res_tensor = add_src1; + } else if (add_src1 == ar_local || add_src1->data == ar_local->data) { + res_tensor = add_src0; + } else { + return false; + } + + if (!res_tensor || !res_tensor->data) return false; + + if (ar_local->type != res_tensor->type) return false; + + const bool is_same_shape = (ar_local->ne[0] == res_tensor->ne[0] && ar_local->ne[1] == res_tensor->ne[1] && + ar_local->ne[2] == res_tensor->ne[2] && ar_local->ne[3] == res_tensor->ne[3]); + const bool is_row_bcast = (ar_local->ne[0] == res_tensor->ne[0] && + res_tensor->ne[1] == 1 && res_tensor->ne[2] == 1 && res_tensor->ne[3] == 1); + + if (!is_same_shape && !is_row_bcast) return false; + + if (is_same_shape) { + if (ar_local->nb[1] != res_tensor->nb[1] || ar_local->nb[2] != res_tensor->nb[2] || + ar_local->nb[3] != res_tensor->nb[3]) { + return false; + } + if (ggml_is_contiguous(ar_local) != ggml_is_contiguous(res_tensor)) { + return false; + } + } + if (ggml_is_contiguous(ar_local) != ggml_is_contiguous(node.dst())) { + return false; + } + + struct htp_allreduce_kernel_params new_kparams; + if (!ggml_hexagon_precompute_allreduce_params( + sess, node.dst(), (uint32_t) ar_kparams->rank, (uint32_t) ar_kparams->n_ranks, true, is_row_bcast, &new_kparams + )) { + HEX_VERBOSE("ggml-hex: %s skip ALLREDUCE_ADD fusion: solver failed\n", sess->c_name()); + return false; + } + + size_t extra_bufs = 0, extra_vmem = 0, extra_tens = 0; + auto fit_t = [&](const ggml_tensor * t) { + if (!t) return; + if (!t_map.count(t)) { + extra_tens++; + auto sbuf = static_cast(t->buffer->context); + if (!b_map.count(sbuf->fd())) { + extra_vmem += sbuf->size(); + extra_bufs += 1; + } + } + }; + fit_t(res_tensor); + fit_t(node.dst()); + if ((extra_bufs + n_bufs) > n_bufs_max || (extra_tens + n_tens) > n_tens_max || (extra_vmem + b_vmem) > b_vmem_max) { + return false; + } + + last_node.opcode = HTP_OP_ALLREDUCE_ADD; + last_node.name = "ALLREDUCE+ADD"; + last_node.inputs.push_back(res_tensor); + last_node.outputs.clear(); + last_node.outputs.push_back(node.dst()); + last_node.fused.push_back(node.node); + memcpy(last_node.kernel_params, &new_kparams, sizeof(new_kparams)); + + htp_op_desc & o = h_ops[n_ops - 1]; + o.opcode = HTP_OP_ALLREDUCE_ADD; + memcpy(o.kernel_params, &new_kparams, sizeof(new_kparams)); + + const uint32_t n_ranks = (uint32_t) ar_kparams->n_ranks; + o.src[2 * n_ranks] = add_tensor(res_tensor); + o.dst[0] = add_tensor(node.dst()); + for (uint32_t d = 1; d < HTP_OP_MAX_OUTPUTS; d++) { + o.dst[d] = 0xffff; + } + + HEX_VERBOSE("ggml-hex: %s fused ALLREDUCE+ADD (#%u)\n", sess->c_name(), n_ops - 1); + return true; + } + + bool try_fuse_rms_norm_mul(const htp_opnode & node) { + if (n_ops == 0) return false; + if (node.opcode != HTP_OP_MUL) return false; + + htp_opnode & last_node = ops[n_ops - 1]; + if (last_node.opcode != HTP_OP_RMS_NORM) return false; + + const ggml_tensor * mul_src0 = node.src0(); + const ggml_tensor * mul_src1 = node.src1(); + const ggml_tensor * rms_out = last_node.dst(); + + if (!mul_src0 || !mul_src1 || !rms_out) return false; + if (!ggml_hexagon_tensor_is_fuseable(rms_out)) return false; + + const ggml_tensor * weight = nullptr; + if (mul_src0 == rms_out || mul_src0->data == rms_out->data) { + weight = mul_src1; + } else if (mul_src1 == rms_out || mul_src1->data == rms_out->data) { + weight = mul_src0; + } else { + return false; + } + + if (!weight || !weight->data) return false; + + const ggml_tensor * src0 = last_node.src0(); + if (!src0 || !src0->data) return false; + + if (src0->ne[0] != weight->ne[0] || src0->ne[0] != node.dst()->ne[0]) { + return false; + } + + const bool is_row_bcast = (weight->ne[1] == 1 && weight->ne[2] == 1 && weight->ne[3] == 1); + const bool is_same_shape = (src0->ne[0] == weight->ne[0] && src0->ne[1] == weight->ne[1] && + src0->ne[2] == weight->ne[2] && src0->ne[3] == weight->ne[3]); + if (!is_row_bcast && !is_same_shape) return false; + + if (!ggml_are_same_shape(src0, node.dst())) { + return false; + } + if (ggml_is_contiguous(src0) != ggml_is_contiguous(node.dst())) { + return false; + } + + struct htp_unary_kernel_params new_kparams; + ggml_hexagon_precompute_unary_params( + sess, HTP_OP_RMS_NORM_MUL, src0, weight, node.dst(), &new_kparams + ); + + if ((size_t) new_kparams.vtcm_size > sess->vtcm_size) { + HEX_VERBOSE("ggml-hex: %s skip RMS_NORM_MUL fusion: VTCM needed (%d) > budget (%zu)\n", + sess->c_name(), new_kparams.vtcm_size, sess->vtcm_size); + return false; + } + + size_t extra_bufs = 0, extra_vmem = 0, extra_tens = 0; + auto fit_t = [&](const ggml_tensor * t) { + if (!t) return; + if (!t_map.count(t)) { + extra_tens++; + auto sbuf = static_cast(t->buffer->context); + if (!b_map.count(sbuf->fd())) { + extra_vmem += sbuf->size(); + extra_bufs += 1; + } + } + }; + fit_t(weight); + fit_t(node.dst()); + if ((extra_bufs + n_bufs) > n_bufs_max || (extra_tens + n_tens) > n_tens_max || (extra_vmem + b_vmem) > b_vmem_max) { + return false; + } + + last_node.opcode = HTP_OP_RMS_NORM_MUL; + last_node.name = "RMS_NORM+MUL"; + last_node.inputs.clear(); + last_node.inputs.push_back(src0); + last_node.inputs.push_back(weight); + last_node.outputs.clear(); + last_node.outputs.push_back(node.dst()); + last_node.fused.push_back(node.node); + memcpy(last_node.kernel_params, &new_kparams, sizeof(new_kparams)); + + htp_op_desc & o = h_ops[n_ops - 1]; + o.opcode = HTP_OP_RMS_NORM_MUL; + memcpy(o.kernel_params, &new_kparams, sizeof(new_kparams)); + + o.src[0] = add_tensor(src0); + o.src[1] = add_tensor(weight); + for (uint32_t s = 2; s < HTP_OP_MAX_INPUTS; s++) { + o.src[s] = 0xffff; + } + o.dst[0] = add_tensor(node.dst()); + for (uint32_t d = 1; d < HTP_OP_MAX_OUTPUTS; d++) { + o.dst[d] = 0xffff; + } + + HEX_VERBOSE("ggml-hex: %s fused RMS_NORM+MUL (#%u)\n", sess->c_name(), n_ops - 1); + return true; + } + + bool try_fuse_mul_mat_add(const htp_opnode & node) { + if (n_ops == 0) return false; + if (node.opcode != HTP_OP_ADD) return false; + + htp_opnode & last_node = ops[n_ops - 1]; + if (last_node.opcode != HTP_OP_MUL_MAT) return false; + + const ggml_tensor * add_src0 = node.src0(); + const ggml_tensor * add_src1 = node.src1(); + const ggml_tensor * mm_out = last_node.dst(); + + if (!add_src0 || !add_src1 || !mm_out) return false; + if (!ggml_hexagon_tensor_is_fuseable(mm_out)) return false; + + const ggml_tensor * src2 = nullptr; + if (add_src0 == mm_out || add_src0->data == mm_out->data) { + src2 = add_src1; + } else if (add_src1 == mm_out || add_src1->data == mm_out->data) { + src2 = add_src0; + } else { + return false; + } + + if (!src2 || !src2->data) return false; + + const ggml_tensor * src0 = last_node.src0(); + const ggml_tensor * src1 = last_node.src1(); + if (!src0 || !src1) return false; + + struct htp_mm_kernel_params kparams; + ggml_hexagon_precompute_fused_matmul_add_params(sess, src0, src1, src2, node.dst(), &kparams); + const int src1_nrows = src1->ne[1] * src1->ne[2] * src1->ne[3]; + const bool can_fuse = (kparams.n_hmx > 0) || (src1_nrows == 1); + if (!can_fuse) return false; + + if ((size_t) kparams.vtcm_size > sess->vtcm_size) { + HEX_VERBOSE("ggml-hex: %s skip MUL_MAT_ADD fusion: VTCM needed (%d) > budget (%zu)\n", + sess->c_name(), kparams.vtcm_size, sess->vtcm_size); + return false; + } + + size_t extra_bufs = 0, extra_vmem = 0, extra_tens = 0; + auto fit_t = [&](const ggml_tensor * t) { + if (!t) return; + if (!t_map.count(t)) { + extra_tens++; + auto sbuf = static_cast(t->buffer->context); + if (!b_map.count(sbuf->fd())) { + extra_vmem += sbuf->size(); + extra_bufs += 1; + } + } + }; + fit_t(src2); + fit_t(node.dst()); + if ((extra_bufs + n_bufs) > n_bufs_max || (extra_tens + n_tens) > n_tens_max || (extra_vmem + b_vmem) > b_vmem_max) { + return false; + } + + last_node.opcode = HTP_OP_MUL_MAT_ADD; + last_node.name = "MUL_MAT+ADD"; + last_node.inputs.clear(); + last_node.inputs.push_back(src0); + last_node.inputs.push_back(src1); + last_node.inputs.push_back(src2); + last_node.outputs.clear(); + last_node.outputs.push_back(node.dst()); + last_node.fused.push_back(node.node); + memcpy(last_node.kernel_params, &kparams, sizeof(kparams)); + + htp_op_desc & o = h_ops[n_ops - 1]; + o.opcode = HTP_OP_MUL_MAT_ADD; + memcpy(o.kernel_params, &kparams, sizeof(kparams)); + + o.src[0] = add_tensor(src0); + o.src[1] = add_tensor(src1); + o.src[2] = add_tensor(src2); + for (uint32_t s = 3; s < HTP_OP_MAX_INPUTS; s++) { + o.src[s] = 0xffff; + } + o.dst[0] = add_tensor(node.dst()); + for (uint32_t d = 1; d < HTP_OP_MAX_OUTPUTS; d++) { + o.dst[d] = 0xffff; + } + + HEX_VERBOSE("ggml-hex: %s fused MUL_MAT+ADD (#%u)\n", sess->c_name(), n_ops - 1); + return true; + } + + bool try_fuse_mul_mat_nx(const htp_opnode & node) { + if (n_ops == 0 || node.opcode != HTP_OP_MUL_MAT) return false; + if (!is_mergeable_mul_mat(node.node)) return false; + + const ggml_tensor * w_in = node.src0(); + const ggml_tensor * x_in = node.src1(); + const ggml_tensor * d_in = node.dst(); + if (!w_in || !x_in || !d_in) return false; + + htp_opnode & last_node = ops[n_ops - 1]; + + // Case 1: last_node is already MUL_MAT_NX + if (last_node.opcode == HTP_OP_MUL_MAT_NX) { + const uint32_t curr_n = (uint32_t) last_node.outputs.size(); + if (curr_n >= HTP_OP_MAX_OUTPUTS || curr_n + 1 >= HTP_OP_MAX_INPUTS) { + return false; + } + + const ggml_tensor * w0 = last_node.inputs[0]; + const ggml_tensor * x = last_node.inputs[curr_n]; + + if (x_in != x || w_in->type != w0->type || w_in->ne[0] != w0->ne[0]) { + return false; + } + + struct htp_mm_kernel_params kparams; + ggml_hexagon_precompute_fused_mmnx_params(sess, w0, x, curr_n + 1, &kparams); + if ((size_t) kparams.vtcm_size > sess->vtcm_size) { + HEX_VERBOSE("ggml-hex: %s skip NX fusion: VTCM needed (%d) > budget (%zu)\n", + sess->c_name(), kparams.vtcm_size, sess->vtcm_size); + return false; + } + + size_t extra_bufs = 0, extra_vmem = 0, extra_tens = 0; + auto fit_t = [&](const ggml_tensor * t) { + if (!t) return; + if (!t_map.count(t)) { + extra_tens++; + auto sbuf = static_cast(t->buffer->context); + if (!b_map.count(sbuf->fd())) { + extra_vmem += sbuf->size(); + extra_bufs += 1; + } + } + }; + fit_t(w_in); + fit_t(d_in); + if ((extra_bufs + n_bufs) > n_bufs_max || (extra_tens + n_tens) > n_tens_max || (extra_vmem + b_vmem) > b_vmem_max) { + return false; + } + + last_node.inputs[curr_n] = w_in; + last_node.inputs.push_back(x); + last_node.outputs.push_back(d_in); + last_node.fused.push_back(node.node); + memcpy(last_node.kernel_params, &kparams, sizeof(kparams)); + + htp_op_desc & o = h_ops[n_ops - 1]; + memcpy(o.kernel_params, &kparams, sizeof(kparams)); + + for (uint32_t s = 0; s <= curr_n + 1; s++) { + o.src[s] = add_tensor(last_node.inputs[s]); + } + for (uint32_t s = curr_n + 2; s < HTP_OP_MAX_INPUTS; s++) { + o.src[s] = 0xffff; + } + for (uint32_t d = 0; d <= curr_n; d++) { + o.dst[d] = add_tensor(last_node.outputs[d]); + } + for (uint32_t d = curr_n + 1; d < HTP_OP_MAX_OUTPUTS; d++) { + o.dst[d] = 0xffff; + } + + HEX_VERBOSE("ggml-hex: %s fused MUL_MAT_NX (N=%u, #%u)\n", sess->c_name(), curr_n + 1, n_ops - 1); + return true; + } + + // Case 2: last_node is single MUL_MAT + if (last_node.opcode == HTP_OP_MUL_MAT) { + if (!is_mergeable_mul_mat_pair(last_node.node, node.node)) { + return false; + } + + const ggml_tensor * w0 = last_node.src0(); + const ggml_tensor * x = last_node.src1(); + const ggml_tensor * w1 = node.src0(); + if (!w0 || !x || !w1) return false; - ggml_hexagon_dump_op_exec(sess->c_name(), ops[n], o.flags); + struct htp_mm_kernel_params kparams; + ggml_hexagon_precompute_fused_mmnx_params(sess, w0, x, 2, &kparams); + if ((size_t) kparams.vtcm_size > sess->vtcm_size) { + HEX_VERBOSE("ggml-hex: %s skip NX fusion: VTCM needed (%d) > budget (%zu)\n", + sess->c_name(), kparams.vtcm_size, sess->vtcm_size); + return false; + } - auto inputs = node.get_inputs(); - for (unsigned int i=0; i < HTP_OP_MAX_INPUTS; i++) { - o.src[i] = (i < inputs.size() && inputs[i]) ? add_tensor(inputs[i]) : 0xffff; - } + size_t extra_bufs = 0, extra_vmem = 0, extra_tens = 0; + auto fit_t = [&](const ggml_tensor * t) { + if (!t) return; + if (!t_map.count(t)) { + extra_tens++; + auto sbuf = static_cast(t->buffer->context); + if (!b_map.count(sbuf->fd())) { + extra_vmem += sbuf->size(); + extra_bufs += 1; + } + } + }; + fit_t(w1); + fit_t(node.dst()); + if ((extra_bufs + n_bufs) > n_bufs_max || (extra_tens + n_tens) > n_tens_max || (extra_vmem + b_vmem) > b_vmem_max) { + return false; + } - auto outputs = node.get_outputs(); - for (unsigned int i=0; i < HTP_OP_MAX_OUTPUTS; i++) { - o.dst[i] = (i < outputs.size() && outputs[i]) ? add_tensor(outputs[i]) : 0xffff; + const ggml_tensor * dst_0 = last_node.dst(); + const ggml_tensor * dst_1 = node.dst(); + + last_node.opcode = HTP_OP_MUL_MAT_NX; + last_node.name = "MUL_MAT_NX"; + last_node.inputs.clear(); + last_node.inputs.push_back(w0); + last_node.inputs.push_back(w1); + last_node.inputs.push_back(x); + last_node.outputs.clear(); + last_node.outputs.push_back(dst_0); + last_node.outputs.push_back(dst_1); + last_node.fused.push_back(node.node); + memcpy(last_node.kernel_params, &kparams, sizeof(kparams)); + + htp_op_desc & o = h_ops[n_ops - 1]; + o.opcode = HTP_OP_MUL_MAT_NX; + memcpy(o.kernel_params, &kparams, sizeof(kparams)); + + o.src[0] = add_tensor(w0); + o.src[1] = add_tensor(w1); + o.src[2] = add_tensor(x); + for (uint32_t s = 3; s < HTP_OP_MAX_INPUTS; s++) { + o.src[s] = 0xffff; + } + o.dst[0] = add_tensor(dst_0); + o.dst[1] = add_tensor(dst_1); + for (uint32_t d = 2; d < HTP_OP_MAX_OUTPUTS; d++) { + o.dst[d] = 0xffff; + } + + HEX_VERBOSE("ggml-hex: %s fused MUL_MAT_NX (N=2, #%u)\n", sess->c_name(), n_ops - 1); + return true; } + + return false; } - void finalize_ranges() { +enum ggml_hexagon_fusion_flags { + GGML_HEXAGON_FUSE_ALLREDUCE_ADD = (1 << 1), // 2 + GGML_HEXAGON_FUSE_RMS_NORM_MUL = (1 << 2), // 4 + GGML_HEXAGON_FUSE_MUL_MAT_ADD = (1 << 3), // 8 + GGML_HEXAGON_FUSE_MUL_MAT_NX = (1 << 4), // 16 +}; + +static inline bool ggml_hexagon_is_fusion_enabled(int flag) { + if (opt_opfusion <= 0) return false; + if (opt_opfusion == 1) return true; // 1 enables all + return (opt_opfusion & flag) != 0; +} + + bool try_fuse(const htp_opnode & node) { + if (!opt_opfusion) return false; + if (ggml_hexagon_is_fusion_enabled(GGML_HEXAGON_FUSE_ALLREDUCE_ADD) && try_fuse_allreduce_add(node)) return true; + if (ggml_hexagon_is_fusion_enabled(GGML_HEXAGON_FUSE_RMS_NORM_MUL) && try_fuse_rms_norm_mul(node)) return true; + if (ggml_hexagon_is_fusion_enabled(GGML_HEXAGON_FUSE_MUL_MAT_ADD) && try_fuse_mul_mat_add(node)) return true; + if (ggml_hexagon_is_fusion_enabled(GGML_HEXAGON_FUSE_MUL_MAT_NX) && try_fuse_mul_mat_nx(node)) return true; + return false; } }; +struct ggml_hexagon_registry { + ggml_hexagon_registry(ggml_backend_reg_t reg); + ~ggml_hexagon_registry(); + + ggml_backend_device devices[GGML_HEXAGON_MAX_SESSIONS]; +}; + struct ggml_hexagon_opqueue { // Shared buffer for storing batches ggml_hexagon_shared_buffer *shm_buf; size_t shm_blk_size; + uint64_t req_seq = 0; + uint64_t rsp_seq = 0; + using opvec = std::vector; std::queue done; // completed batch ids @@ -1429,8 +2317,8 @@ struct ggml_hexagon_opqueue { for (unsigned int i = 0; i < depth; i++) { done.push(i); } if (opt_verbose) { - GGML_LOG_INFO("ggml-hex: %s allocated op-queue : batch-size %zu depth %zu shm-size %zu shm-block-size %zu\n", - sess->c_name(), batch_size, depth, shm_buf->size, shm_blk_size); + GGML_LOG_INFO("ggml-hex: %s allocated opqueue : batch-size %zu depth %zu shm-size %zu shm-block-size %zu\n", + sess->c_name(), batch_size, depth, shm_buf->size(), shm_blk_size); } } @@ -1453,8 +2341,9 @@ struct ggml_hexagon_opqueue { req.n_bufs = op_batch->n_bufs; req.n_tensors = op_batch->n_tens; req.n_ops = op_batch->n_ops; + req.seq = ++req_seq; - op_cache[req.id] = op_batch->ops; + op_cache[req.id] = std::move(op_batch->ops); start_usec[req.id] = ggml_time_us(); const size_t b_size = sizeof(htp_buf_desc) * req.n_bufs; @@ -1470,10 +2359,10 @@ struct ggml_hexagon_opqueue { req.n_traces = 0; } - dbuf.ptr = shm_buf->base + (req.id * shm_blk_size); - dbuf.fd = shm_buf->fd; + dbuf.ptr = shm_buf->base() + (req.id * shm_blk_size); + dbuf.fd = shm_buf->fd(); dbuf.flags = DSPQUEUE_BUFFER_FLAG_FLUSH_SENDER | DSPQUEUE_BUFFER_FLAG_INVALIDATE_RECIPIENT; - dbuf.offset = (uint8_t*) dbuf.ptr - (uint8_t*) shm_buf->base; + dbuf.offset = (uint8_t*) dbuf.ptr - (uint8_t*) shm_buf->base(); dbuf.size = b_size + t_size + o_size + p_size + tr_size; GGML_ASSERT(dbuf.size <= shm_blk_size); @@ -1487,7 +2376,7 @@ struct ggml_hexagon_opqueue { memcpy(t_ptr, (void *) op_batch->h_tens.data(), t_size); memcpy(o_ptr, (void *) op_batch->h_ops.data(), o_size); - HEX_VERBOSE("ggml-hex: %s op-queue push batch #%u : n-bufs %u n-tensors %u n-ops %u vmem %zu : b-size %zu t-size %zu o-size %zu m-size %zu\n", + HEX_VERBOSE("ggml-hex: %s opqueue-push batch #%u : n-bufs %u n-tensors %u n-ops %u vmem %zu : b-size %zu t-size %zu o-size %zu m-size %zu\n", shm_buf->sess->c_name(), req.id, req.n_bufs, req.n_tensors, req.n_ops, op_batch->b_vmem, b_size, t_size, o_size, (size_t) dbuf.size); @@ -1530,33 +2419,40 @@ struct ggml_hexagon_opqueue { const size_t m_size = b_size + t_size + o_size + p_size + tr_size; GGML_ASSERT(m_size <= shm_blk_size); - HEX_VERBOSE("ggml-hex: %s op-queue pop batch #%u : n-bufs %u n-tensors %u n-ops %u : m-size %zu b-size %zu t-size %zu o-size %zu\n", + HEX_VERBOSE("ggml-hex: %s opqueue-pop batch #%u : n-bufs %u n-tensors %u n-ops %u : m-size %zu b-size %zu t-size %zu o-size %zu\n", shm_buf->sess->c_name(), rsp.id, rsp.n_bufs, rsp.n_tensors, rsp.n_ops, (size_t) dbuf.size, b_size, t_size, o_size); uint8_t * m_ptr = (uint8_t*) dbuf.ptr; uint8_t * p_ptr = m_ptr + (b_size + t_size + o_size); - if (opt_profile && rsp.n_ops > 0) { + if (rsp.n_ops > 0) { auto & ops = op_cache[rsp.id]; - GGML_ASSERT(rsp.n_ops <= ops.size()); const htp_prof_desc * pd = (const htp_prof_desc *) p_ptr; - const htp_trace_desc * trace_events = nullptr; - if (opt_profile == 3) { trace_events = (const htp_trace_desc *) (p_ptr + p_size); } - ggml_hexagon_dump_batch_prof(shm_buf->sess->name, rsp); + if (opt_profile) { + ggml_hexagon_dump_batch_prof(shm_buf->sess->name, rsp); + } for (uint32_t i = 0; i < rsp.n_ops; i++) { - ggml_hexagon_dump_op_prof(shm_buf->sess->name, ops[i], pd[i]); + if (opt_profile) { + ggml_hexagon_dump_op_prof(shm_buf->sess->name, ops[i], pd[i]); + } } - ggml_hexagon_dump_trace_events(shm_buf->sess->name, rsp, trace_events, n_traces); + if (opt_profile) { + ggml_hexagon_dump_trace_events(shm_buf->sess->name, rsp, trace_events, n_traces); + } + } + + if (rsp.seq > rsp_seq) { + rsp_seq = rsp.seq; } } }; @@ -1601,10 +2497,8 @@ void ggml_hexagon_session::flush_pending(bool all) { } } -void ggml_hexagon_session::flush_batch() { - if (op_batch->empty()) { return; } - - op_batch->finalize_ranges(); +void ggml_hexagon_session::flush_batch(size_t min_ops) { + if (op_batch->n_ops < min_ops) { return; } htp_opbatch_req req {}; dspqueue_buffer dbuf{}; @@ -1625,17 +2519,267 @@ void ggml_hexagon_session::flush_batch() { } } +void ggml_hexagon_session::flush(bool all) { + flush_sync_peers(); + flush_batch(); + flush_pending(all); +} + void ggml_hexagon_session::enqueue_op(const htp_opnode & node) { + for (auto t : node.get_inputs()) { + if (t && t->buffer && ggml_backend_buffer_is_hexagon(t->buffer)) { + if (ggml_backend_hexagon_buffer_get_sess(t->buffer) != this) { + this->clone_buffer(static_cast(t->buffer->context)); + } + } + } + for (auto t : node.get_outputs()) { + if (t && t->buffer && ggml_backend_buffer_is_hexagon(t->buffer)) { + if (ggml_backend_hexagon_buffer_get_sess(t->buffer) != this) { + this->clone_buffer(static_cast(t->buffer->context)); + } + } + } + + if (opt_opfusion && op_batch->try_fuse(node)) { + return; + } + if (!op_batch->fit_op(node)) { flush_batch(); } op_batch->add_op(node); } -// Flush HTP response queue i.e wait for all outstanding requests to complete -void ggml_hexagon_session::flush(bool all) { +void ggml_hexagon_session::enqueue_cpy(const ggml_tensor * src, ggml_tensor * dst, const ggml_tensor * sync_tensor, uint32_t fence_seq) { + htp_opnode cpy_node(HTP_OP_CPY); + + ggml_tensor* node = cpy_node.add_dummy(*dst); + node->op = GGML_OP_CPY; + node->src[0] = const_cast(src); + node->src[1] = sync_tensor ? cpy_node.add_dummy(*sync_tensor) : nullptr; + if (sync_tensor) { + node->op_params[0] = (int32_t) fence_seq; + } + + cpy_node.init(node); + if (sync_tensor) { + cpy_node.name = "CPY+FENCE"; + } + this->enqueue_op(cpy_node); +} + +void ggml_hexagon_session::enqueue_fence(const ggml_tensor * sync_tensor, uint32_t fence_seq) { + htp_opnode sync_node(HTP_OP_FENCE); + + ggml_tensor* node = sync_node.add_dummy(*sync_tensor); + node->op = GGML_OP_NONE; + node->src[0] = node; + node->op_params[0] = (int32_t) fence_seq; + + sync_node.init(node); + sync_node.name = "FENCE"; + this->enqueue_op(sync_node); +} + +static bool ggml_hexagon_precompute_allreduce_params( + const struct ggml_hexagon_session * sess, + const struct ggml_tensor * dst, + uint32_t rank, + uint32_t n_ranks, + bool has_add, + bool is_row_bcast, + struct htp_allreduce_kernel_params * kparams +) { + memset(kparams, 0, sizeof(*kparams)); + kparams->rank = (int32_t) rank; + kparams->n_ranks = (int32_t) n_ranks; + kparams->is_row_bcast = (has_add && is_row_bcast) ? 1 : 0; + + const uint32_t n_bufs = n_ranks + 1 + (has_add ? 1 : 0); + const uint32_t nelem = (uint32_t) ggml_nelements(dst); + const uint32_t elem_size = (dst->type == GGML_TYPE_F16) ? sizeof(ggml_fp16_t) : sizeof(float); + const bool is_contiguous = ggml_is_contiguous(dst); + + const uint32_t ne0 = (uint32_t) dst->ne[0]; + const uint32_t ne1 = (uint32_t) (dst->ne[1] * dst->ne[2] * dst->ne[3]); + kparams->ne0 = (int32_t) ne0; + kparams->ne1 = (int32_t) ne1; + + const bool use_1d = is_contiguous && !(has_add && is_row_bcast && ne1 > 1); + + if (has_add) { + kparams->n_dsts = 1; + if (use_1d) { + kparams->rank_elem_start = 0; + kparams->rank_nelem = (int32_t) nelem; + } else { + kparams->rank_elem_start = 0; + kparams->rank_nelem = (int32_t) ne1; + } + } else { + kparams->n_dsts = (int32_t) n_ranks; + if (use_1d) { + const uint32_t rank_chunk_elems = hex_round_up((nelem + n_ranks - 1) / n_ranks, 128); + const uint32_t rank_elem_start = (std::min)(rank * rank_chunk_elems, nelem); + const uint32_t rank_elem_end = (std::min)(rank_elem_start + rank_chunk_elems, nelem); + const uint32_t rank_nelem = rank_elem_end - rank_elem_start; + kparams->rank_elem_start = (int32_t) rank_elem_start; + kparams->rank_nelem = (int32_t) rank_nelem; + } else { + const uint32_t rank_chunk_rows = (ne1 + n_ranks - 1) / n_ranks; + const uint32_t rank_r0 = (std::min)(rank * rank_chunk_rows, ne1); + const uint32_t rank_r1 = (std::min)(rank_r0 + rank_chunk_rows, ne1); + const uint32_t rank_nrows = rank_r1 - rank_r0; + kparams->rank_elem_start = (int32_t) rank_r0; + kparams->rank_nelem = (int32_t) rank_nrows; + } + } + + if (use_1d) { + const uint32_t rank_nelem = (uint32_t) kparams->rank_nelem; + const uint32_t n_threads = (std::min)((uint32_t) sess->n_threads, (std::max)(1u, rank_nelem / 128)); + kparams->n_threads = n_threads; + + uint32_t block_elems = 65536; + if (block_elems > rank_nelem / n_threads && rank_nelem / n_threads > 128) { + block_elems = hex_round_up(rank_nelem / (n_threads * 2), 128); + } + block_elems = (std::max)(128u, block_elems); + + kparams->block_elems = block_elems; + kparams->vtcm_size_per_thread = 2 * block_elems * elem_size; + kparams->vtcm_size = n_threads * n_bufs * kparams->vtcm_size_per_thread; + + while ((size_t) kparams->vtcm_size > sess->vtcm_size && block_elems > 128) { + const size_t max_bytes_per_buf = sess->vtcm_size / (n_threads * n_bufs * 2); + block_elems = (uint32_t) hex_align_down((size_t) (max_bytes_per_buf / elem_size), 128); + if (block_elems < 128) break; + kparams->block_elems = block_elems; + kparams->vtcm_size_per_thread = 2 * block_elems * elem_size; + kparams->vtcm_size = n_threads * n_bufs * kparams->vtcm_size_per_thread; + } + + if (sess->vtcm_size < (size_t) kparams->vtcm_size || block_elems < 128) { + HEX_VERBOSE("ggml-hex: %s allreduce 1D solver failed to fit VTCM (%d > %zu)\n", + sess->c_name(), kparams->vtcm_size, sess->vtcm_size); + return false; + } + + kparams->elems_per_thread = hex_round_up((rank_nelem + n_threads - 1) / n_threads, block_elems); + kparams->kernel_type = HTP_ALLREDUCE_KERNEL_DMA_1D; + return true; + } else { + const uint32_t rank_nrows = (uint32_t) kparams->rank_nelem; + const uint32_t n_threads = (std::min)((uint32_t) sess->n_threads, (std::max)(1u, rank_nrows)); + kparams->n_threads = n_threads; + + const uint32_t row_bytes = ne0 * elem_size; + const uint32_t row_size_aligned = (uint32_t) hex_align_up(row_bytes, 128); + kparams->row_size_aligned = row_size_aligned; + + const uint32_t nrows_per_thread = (rank_nrows + n_threads - 1) / n_threads; + uint32_t block_rows = (std::min)(128u, nrows_per_thread); + block_rows = (std::max)(1u, block_rows); + kparams->block_elems = block_rows; + + kparams->vtcm_size_per_thread = 2 * (block_rows * row_size_aligned); + kparams->vtcm_size = n_threads * n_bufs * kparams->vtcm_size_per_thread; + + while ((size_t) kparams->vtcm_size > sess->vtcm_size && block_rows > 1) { + const size_t max_rows_per_buf = sess->vtcm_size / (n_threads * n_bufs * 2 * row_size_aligned); + block_rows = (std::max)(1u, (uint32_t) max_rows_per_buf); + kparams->block_elems = block_rows; + kparams->vtcm_size_per_thread = 2 * (block_rows * row_size_aligned); + kparams->vtcm_size = n_threads * n_bufs * kparams->vtcm_size_per_thread; + if (max_rows_per_buf == 0) break; + } + + if (sess->vtcm_size < (size_t) kparams->vtcm_size || block_rows < 1) { + HEX_VERBOSE("ggml-hex: %s allreduce 2D solver failed to fit VTCM (%d > %zu)\n", + sess->c_name(), kparams->vtcm_size, sess->vtcm_size); + return false; + } + + kparams->elems_per_thread = nrows_per_thread; + kparams->kernel_type = HTP_ALLREDUCE_KERNEL_DMA_2D; + return true; + } +} + +void ggml_hexagon_session::enqueue_allreduce( + const ggml_tensor * dst, + const std::vector & src_tensors, + const std::vector & sync_tensors, + uint32_t rank, + uint32_t n_ranks, + uint32_t fence_seq_entry, + uint32_t fence_seq_exit +) { + htp_opnode ar_node(HTP_OP_ALLREDUCE); + + ggml_tensor* node = ar_node.add_dummy(*dst); + node->op = GGML_OP_NONE; + node->op_params[0] = (int32_t) fence_seq_entry; + node->op_params[1] = (int32_t) fence_seq_exit; + + ar_node.init(node); + + ar_node.inputs.clear(); + for (size_t i = 0; i < src_tensors.size(); i++) { + ar_node.inputs.push_back(src_tensors[i]); + } + for (size_t i = 0; i < sync_tensors.size(); i++) { + ar_node.inputs.push_back(ar_node.add_dummy(*sync_tensors[i])); + } + + ar_node.outputs.clear(); + for (size_t i = 0; i < src_tensors.size(); i++) { + ar_node.outputs.push_back(src_tensors[i]); + } + + ggml_hexagon_precompute_allreduce_params( + this, dst, rank, n_ranks, false, false, + (struct htp_allreduce_kernel_params *) ar_node.kernel_params + ); + + ar_node.name = "ALLREDUCE"; + this->enqueue_op(ar_node); +} + +void ggml_hexagon_session::wait_event(uint64_t seq) { + flush_sync_peers(); + HEX_VERBOSE("ggml-hex: %s opqueue-wait start: seq %llu, current rsp-seq %llu, pending %d\n", + this->name.c_str(), (unsigned long long)seq, (unsigned long long)op_queue->rsp_seq, (int)this->op_pending); + while (op_queue->rsp_seq < seq && this->op_pending > 0) { + this->flush_pending(false); + } + HEX_VERBOSE("ggml-hex: %s opqueue-wait end: seq %llu, current rsp-seq %llu, pending %d\n", + this->name.c_str(), (unsigned long long)seq, (unsigned long long)op_queue->rsp_seq, (int)this->op_pending); +} + +uint64_t ggml_hexagon_session::record_event() { flush_batch(); - flush_pending(all); + return op_queue->req_seq; +} + +bool ggml_hexagon_session::clone_buffer(const ggml_hexagon_shared_buffer *sbuf) +{ + if (this->cloned_buffers.find(sbuf->fd()) != this->cloned_buffers.end()) return true; + + HEX_VERBOSE("ggml-hex: %s clone-buffer: %s base %p size %zu fd %d\n", this->name.c_str(), + sbuf->c_name(), sbuf->base(), sbuf->size(), sbuf->fd()); + + auto clone = std::make_unique(this, *sbuf); + try { + clone->mmap(); + } catch (const std::exception & exc) { + GGML_LOG_ERROR("ggml-hex: %s lazy mapping of buffer context failed: %s\n", this->c_name(), exc.what()); + return false; + } + + this->cloned_buffers[sbuf->fd()] = std::move(clone); + return true; } static size_t ggml_hexagon_measure_max_vmem(ggml_hexagon_session *sess) { @@ -1668,37 +2812,44 @@ static size_t ggml_hexagon_measure_max_vmem(ggml_hexagon_session *sess) { } void ggml_hexagon_session::allocate(int dev_id) noexcept(false) { + const auto & config = opt_device_configs[dev_id]; + int phys_idx = config.physical_idx; + int virt_idx = config.virtual_idx; + this->valid_session = false; this->valid_handle = false; this->valid_queue = false; this->valid_iface = false; - this->domain_id = 3; // Default for CDSP, updated after the session is created - this->session_id = 0; // Default for CDSP, updated after the session is created + this->phys_idx = phys_idx; + this->virt_idx = virt_idx; + this->domain_id = get_domain_id(phys_idx); + this->session_id = 0; this->dev_id = dev_id; - this->name = std::string("HTP") + std::to_string(dev_id); - - this->op_pending = 0; + this->name = config.name; + this->op_pending = 0; GGML_LOG_DEBUG("ggml-hex: %s allocating new session\n", this->name.c_str()); domain * my_domain = htpdrv_get_domain(this->domain_id); if (my_domain == NULL) { - GGML_LOG_ERROR("ggml-hex: unable to get domain struct for CDSP\n"); + GGML_LOG_ERROR("ggml-hex: unable to get domain struct for CDSP (domain_id %d)\n", this->domain_id); throw std::runtime_error("ggml-hex: failed to get CDSP domain (see log for details)"); } - // Create new session - if (dev_id != 0) { + std::string dom_name = get_domain_name(phys_idx); + + // Create new session if virtual_idx > 0 + if (virt_idx > 0) { struct remote_rpc_reserve_new_session n; - n.domain_name_len = strlen(CDSP_DOMAIN_NAME); - n.domain_name = const_cast(CDSP_DOMAIN_NAME); + n.domain_name_len = dom_name.size(); + n.domain_name = const_cast(dom_name.c_str()); n.session_name = const_cast(this->name.c_str()); n.session_name_len = this->name.size(); int err = remote_session_control(FASTRPC_RESERVE_NEW_SESSION, (void *) &n, sizeof(n)); if (err != AEE_SUCCESS) { - GGML_LOG_ERROR("ggml-hex: failed to reserve new session %d : error 0x%x\n", dev_id, err); + GGML_LOG_ERROR("ggml-hex: failed to reserve new session %d (physical %d, virtual %d) : error 0x%x\n", dev_id, phys_idx, virt_idx, err); throw std::runtime_error("ggml-hex: remote_session_control(new-sess) failed (see log for details)"); } @@ -1717,8 +2868,8 @@ void ggml_hexagon_session::allocate(int dev_id) noexcept(false) { struct remote_rpc_get_uri u = {}; u.session_id = this->session_id; - u.domain_name = const_cast(CDSP_DOMAIN_NAME); - u.domain_name_len = strlen(CDSP_DOMAIN_NAME); + u.domain_name = const_cast(dom_name.c_str()); + u.domain_name_len = dom_name.size(); u.module_uri = const_cast(htp_uri); u.module_uri_len = strlen(htp_uri); u.uri = session_uri; @@ -1731,7 +2882,7 @@ void ggml_hexagon_session::allocate(int dev_id) noexcept(false) { snprintf(session_uri, htp_URI_domain_len, "%s%s", htp_uri, my_domain->uri); - GGML_LOG_WARN("ggml-hex: failed to get URI for session %d : error 0x%x. Falling back to single session URI: %s\n", dev_id, err, session_uri); + GGML_LOG_WARN("ggml-hex: failed to get URI for session %d (physical %d, virtual %d) : error 0x%x. Falling back to single session URI: %s\n", dev_id, phys_idx, virt_idx, err, session_uri); } } @@ -1899,14 +3050,17 @@ void ggml_hexagon_session::release() noexcept(true) { if (this->valid_handle) { htp_iface_close(this->handle); } + + this->cloned_buffers.clear(); } ggml_hexagon_session::ggml_hexagon_session(int dev_id, ggml_backend_dev_t dev) noexcept(false) { - buffer_type.device = dev; - repack_buffer_type.device = dev; + buffer_type.device = dev; + host_buffer_type.device = dev; op_batch = nullptr; op_queue = nullptr; + fence_seq = ((uintptr_t)this) & 0xFFFF; try { allocate(dev_id); @@ -1914,8 +3068,8 @@ ggml_hexagon_session::ggml_hexagon_session(int dev_id, ggml_backend_dev_t dev) n buffer_type.iface = ggml_backend_hexagon_buffer_type_interface; buffer_type.context = new ggml_backend_hexagon_buffer_type_context(this->name, this); - repack_buffer_type.iface = ggml_backend_hexagon_repack_buffer_type_interface; - repack_buffer_type.context = new ggml_backend_hexagon_buffer_type_context(this->name + "-REPACK", this); + host_buffer_type.iface = ggml_backend_hexagon_host_buffer_type_interface; + host_buffer_type.context = new ggml_backend_hexagon_buffer_type_context(this->name + "-HOST", this); } catch (const std::exception & exc) { release(); throw; @@ -1926,7 +3080,7 @@ ggml_hexagon_session::~ggml_hexagon_session() noexcept(true) { release(); delete static_cast(buffer_type.context); - delete static_cast(repack_buffer_type.context); + delete static_cast(host_buffer_type.context); } // ** backend interface @@ -1946,7 +3100,8 @@ static bool ggml_hexagon_flash_attn_is_hmx_eligible( return false; } - if (k->type != GGML_TYPE_F16 || v->type != GGML_TYPE_F16) { + if ((k->type != GGML_TYPE_F16 && k->type != GGML_TYPE_Q8_0) || + (v->type != GGML_TYPE_F16 && v->type != GGML_TYPE_Q8_0)) { return false; } @@ -2098,8 +3253,10 @@ static bool ggml_hexagon_supported_flash_attn_ext(const struct ggml_hexagon_sess const struct ggml_tensor * src4 = op->src[4]; const struct ggml_tensor * dst = op; - // Check for F16 support only as requested - if ((src0->type != GGML_TYPE_F16 && src0->type != GGML_TYPE_F32) || src1->type != GGML_TYPE_F16 || src2->type != GGML_TYPE_F16) { + // Check for F16/Q8_0 support + if ((src0->type != GGML_TYPE_F16 && src0->type != GGML_TYPE_F32) || + (src1->type != GGML_TYPE_F16 && src1->type != GGML_TYPE_Q8_0) || + (src2->type != GGML_TYPE_F16 && src2->type != GGML_TYPE_Q8_0)) { return false; } @@ -2352,7 +3509,7 @@ static void ggml_hexagon_precompute_hvx_mm_params( for (uint32_t d = max_prefetch; d >= 2; d /= 2) { htp_mm_hvx_vtcm_layout_build( &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, - 0, src0->nb[1], 0, src2_row_size, d, true, false, false + 0, src0->nb[1], 0, src2_row_size, d, true, false ); if (L.total_bytes <= vtcm_budget) { best_n_prefetch = d; @@ -2362,7 +3519,7 @@ static void ggml_hexagon_precompute_hvx_mm_params( if (best_n_prefetch == 2 && L.total_bytes > vtcm_budget) { htp_mm_hvx_vtcm_layout_build( &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, - 0, src0->nb[1], 0, src2_row_size, 2, true, false, false + 0, src0->nb[1], 0, src2_row_size, 2, true, false ); } kparams->n_prefetch = best_n_prefetch; @@ -2386,7 +3543,7 @@ static void ggml_hexagon_precompute_hvx_mm_params( for (uint32_t d = max_prefetch; d >= 2; d /= 2) { htp_mm_hvx_vtcm_layout_build( &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, - dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, d, false, false, false + dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, d, false, false ); if (L.total_bytes <= vtcm_budget) { best_n_prefetch = d; @@ -2396,7 +3553,7 @@ static void ggml_hexagon_precompute_hvx_mm_params( if (best_n_prefetch == 2 && L.total_bytes > vtcm_budget) { htp_mm_hvx_vtcm_layout_build( &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, - dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 2, false, false, false + dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 2, false, false ); } @@ -2420,7 +3577,7 @@ static void ggml_hexagon_precompute_hvx_mm_params( struct htp_mm_hvx_vtcm_layout L; htp_mm_hvx_vtcm_layout_build( &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, - dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false, false + dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false ); kparams->n_prefetch = 16; @@ -2440,7 +3597,7 @@ static void ggml_hexagon_precompute_hvx_mm_params( struct htp_mm_hvx_vtcm_layout L; htp_mm_hvx_vtcm_layout_build( &L, HTP_MM_KERNEL_HVX_F16_F16_VTCM, wtype, ne10, src1_nrows, sess->n_threads, - dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false, false + dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false ); if (!is_batched && !is_permuted && L.total_bytes <= vtcm_budget) { @@ -2460,7 +3617,7 @@ static void ggml_hexagon_precompute_hvx_mm_params( kparams->src1_row_size = src1->nb[1]; htp_mm_hvx_vtcm_layout_build( &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, - dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false, false + dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false ); kparams->vtcm_size = L.total_bytes; kparams->vtcm_src0_size = L.src0_bytes; @@ -2476,7 +3633,7 @@ static void ggml_hexagon_precompute_hvx_mm_params( struct htp_mm_hvx_vtcm_layout L; htp_mm_hvx_vtcm_layout_build( &L, HTP_MM_KERNEL_HVX_F32_F32_VTCM, wtype, ne10, src1_nrows, sess->n_threads, - dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false, false + dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false ); if (!is_batched && !is_permuted && L.total_bytes <= vtcm_budget) { @@ -2492,7 +3649,7 @@ static void ggml_hexagon_precompute_hvx_mm_params( kparams->src1_row_size = src1->nb[1]; htp_mm_hvx_vtcm_layout_build( &L, kparams->kernel_type, wtype, ne10, src1_nrows, sess->n_threads, - dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false, false + dst->nb[1], src0->nb[1], src1->nb[1], src2_row_size, 16, false, false ); kparams->vtcm_size = L.total_bytes; kparams->vtcm_src0_size = L.src0_bytes; @@ -2642,80 +3799,99 @@ static void ggml_hexagon_precompute_unary_params( kparams->div_tpr = init_fastdiv_values(tiles_per_row); } -static void ggml_hexagon_precompute_fused_qkv_params( +static void ggml_hexagon_precompute_get_rows_params( const struct ggml_hexagon_session * sess, - const struct ggml_tensor * src0, // Wk - const struct ggml_tensor * src1, // x - struct htp_mm_kernel_params * kparams + const struct ggml_tensor * src0, + const struct ggml_tensor * src1, + const struct ggml_tensor * dst, + struct htp_get_rows_kernel_params * kparams ) { memset(kparams, 0, sizeof(*kparams)); - const int wtype = src0->type; - const bool is_repack = ggml_hexagon_is_repack_type((ggml_type) wtype); + const uint32_t ne00 = src0->ne[0]; + const uint32_t ne02 = src0->ne[2]; + const uint32_t ne03 = src0->ne[3]; - const int ne10 = src1->ne[0]; - const int src1_nrows = src1->ne[1] * src1->ne[2] * src1->ne[3]; - const size_t src1_row_size = (wtype == GGML_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10); - const size_t src0_row_size = src0->nb[1]; + const uint32_t ne10 = src1->ne[0]; + const uint32_t ne11 = src1->ne[1]; + const uint32_t ne12 = src1->ne[2]; + const uint32_t nr = ne10 * ne11 * ne12; - uint32_t best_n_prefetch = 16; + const size_t nb01 = src0->nb[1]; + const size_t nb1 = dst->nb[1]; - if (is_repack) { - const uint32_t max_prefetch = (src1_nrows > HTP_MM_HMX_MIN_NROWS) ? 2 : 16; - best_n_prefetch = 2; - for (uint32_t d = max_prefetch; d >= 2; d /= 2) { - struct htp_mm_hvx_vtcm_layout L; - htp_mm_hvx_vtcm_layout_build( - &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads, - 0, src0_row_size, src1_row_size, 0, d, false, true, false - ); - if (L.total_bytes <= sess->vtcm_size) { - best_n_prefetch = d; - break; + const bool can_use_dma = (src0->type == dst->type) && (nb01 == nb1); + const bool use_dma = can_use_dma && (ne00 >= 2048); + + kparams->use_dma = use_dma ? 1 : 0; + + uint32_t chunks_per_row = 1; + uint32_t chunk_size = ne00; + uint32_t total_tasks = nr; + + if (use_dma) { + kparams->n_threads = (std::min)((uint32_t)sess->n_threads, nr); + kparams->tasks_per_thread = (nr + kparams->n_threads - 1) / kparams->n_threads; + } else { + if (src0->type == GGML_TYPE_F32 && nr < sess->n_threads) { + const uint32_t min_chunk_size = 1024; + uint32_t max_chunks = ne00 / min_chunk_size; + if (max_chunks == 0) { + max_chunks = 1; } + chunks_per_row = (std::min)((sess->n_threads + nr - 1) / nr, max_chunks); + chunk_size = (ne00 + chunks_per_row - 1) / chunks_per_row; + total_tasks = nr * chunks_per_row; } + kparams->n_threads = (std::min)(total_tasks, (uint32_t)sess->n_threads); + kparams->tasks_per_thread = (total_tasks + kparams->n_threads - 1) / kparams->n_threads; } - struct htp_mm_hvx_vtcm_layout L; - bool try_tiled = (opt_mm_select >= 2); + kparams->chunks_per_row = chunks_per_row; + kparams->chunk_size = chunk_size; + kparams->total_tasks = total_tasks; - // Test tiled first - htp_mm_hvx_vtcm_layout_build( - &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads, - 0, src0_row_size, src1_row_size, 0, best_n_prefetch, false, true, false - ); + kparams->div_ne10 = init_fastdiv_values(ne10); + kparams->div_ne10_ne11 = init_fastdiv_values(ne10 * ne11); + kparams->div_chunks_per_row = init_fastdiv_values(chunks_per_row); + kparams->div_ne02 = init_fastdiv_values(ne02); + kparams->div_ne03 = init_fastdiv_values(ne03); - if (try_tiled && L.total_bytes <= sess->vtcm_size) { - kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW; - kparams->vtcm_src0_size = L.src0_bytes; - kparams->vtcm_src1_size = L.src1_bytes; - kparams->vtcm_src2_size = L.src2_bytes; - kparams->vtcm_src3_size = L.src3_bytes; - kparams->vtcm_dst_size = L.dst_bytes; - kparams->vtcm_size = L.total_bytes; - kparams->n_prefetch = best_n_prefetch; - } else { - kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT; - size_t flat_src1_row_size = (wtype == GGML_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10); + struct htp_get_rows_vtcm_layout vtcm_layout; + htp_get_rows_vtcm_layout_build(&vtcm_layout, src0->type, ne00, kparams->n_threads); + kparams->vtcm_size = vtcm_layout.total_bytes; +} - htp_mm_hvx_vtcm_layout_build( - &L, HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT, wtype, ne10, src1_nrows, sess->n_threads, - 0, src0_row_size, flat_src1_row_size, 0, best_n_prefetch, false, true, false - ); - kparams->vtcm_src0_size = L.src0_bytes; - kparams->vtcm_src1_size = L.src1_bytes; - kparams->vtcm_src2_size = L.src2_bytes; - kparams->vtcm_src3_size = L.src3_bytes; - kparams->vtcm_dst_size = L.dst_bytes; - kparams->vtcm_size = L.total_bytes; - kparams->n_prefetch = best_n_prefetch; - } +static void ggml_hexagon_precompute_set_rows_params( + const struct ggml_hexagon_session * sess, + const struct ggml_tensor * src0, // values + const struct ggml_tensor * src1, // indices + const struct ggml_tensor * dst, // destination + struct htp_set_rows_kernel_params * kparams +) { + memset(kparams, 0, sizeof(*kparams)); + + const uint32_t nr = src0->ne[1]; + + kparams->n_threads = (std::min)((uint32_t)sess->n_threads, nr); + kparams->tasks_per_thread = (nr + kparams->n_threads - 1) / kparams->n_threads; + kparams->total_tasks = nr; + + kparams->div_ne11 = init_fastdiv_values(src1->ne[1]); + kparams->div_ne12 = init_fastdiv_values(src1->ne[2]); + kparams->div_tasks_per_thread = init_fastdiv_values(kparams->tasks_per_thread); + kparams->div_ne02 = init_fastdiv_values(src0->ne[2]); + + struct htp_set_rows_vtcm_layout vtcm_layout; + htp_set_rows_vtcm_layout_build(&vtcm_layout, dst->type, src0->ne[0], kparams->n_threads); + kparams->vtcm_size = vtcm_layout.total_bytes; } -static void ggml_hexagon_precompute_fused_ffn_params( +static void ggml_hexagon_precompute_fused_mmnx_params( const struct ggml_hexagon_session * sess, - const struct ggml_tensor * src0, // Wgate - const struct ggml_tensor * src1, // y + const struct ggml_tensor * src0, // W0 + const struct ggml_tensor * src1, // x + int32_t n_weights, struct htp_mm_kernel_params * kparams ) { memset(kparams, 0, sizeof(*kparams)); @@ -2737,7 +3913,7 @@ static void ggml_hexagon_precompute_fused_ffn_params( struct htp_mm_hvx_vtcm_layout L; htp_mm_hvx_vtcm_layout_build( &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads, - 0, src0_row_size, src1_row_size, 0, d, false, false, true + 0, src0_row_size, src1_row_size, 0, d, false, true ); if (L.total_bytes <= sess->vtcm_size) { best_n_prefetch = d; @@ -2752,34 +3928,42 @@ static void ggml_hexagon_precompute_fused_ffn_params( // Test tiled first htp_mm_hvx_vtcm_layout_build( &L, HTP_MM_KERNEL_HVX_QUANT_ROW, wtype, ne10, src1_nrows, sess->n_threads, - 0, src0_row_size, src1_row_size, 0, best_n_prefetch, false, false, true + 0, src0_row_size, src1_row_size, 0, best_n_prefetch, false, true ); if (try_tiled && L.total_bytes <= sess->vtcm_size) { kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW; kparams->vtcm_src0_size = L.src0_bytes; kparams->vtcm_src1_size = L.src1_bytes; - kparams->vtcm_src2_size = L.src2_bytes; kparams->vtcm_dst_size = L.dst_bytes; kparams->vtcm_size = L.total_bytes; kparams->n_prefetch = best_n_prefetch; + kparams->n_weights = n_weights; } else { kparams->kernel_type = HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT; size_t flat_src1_row_size = (wtype == GGML_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10); htp_mm_hvx_vtcm_layout_build( &L, HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT, wtype, ne10, src1_nrows, sess->n_threads, - 0, src0_row_size, flat_src1_row_size, 0, best_n_prefetch, false, false, true + 0, src0_row_size, flat_src1_row_size, 0, best_n_prefetch, false, true ); kparams->vtcm_src0_size = L.src0_bytes; kparams->vtcm_src1_size = L.src1_bytes; - kparams->vtcm_src2_size = L.src2_bytes; kparams->vtcm_dst_size = L.dst_bytes; kparams->vtcm_size = L.total_bytes; kparams->n_prefetch = best_n_prefetch; + kparams->n_weights = n_weights; } } +static bool ggml_hexagon_tensor_is_host(const struct ggml_hexagon_session * sess, const struct ggml_tensor * t) { + return t && t->buffer && t->buffer->buft == &sess->host_buffer_type; +} + +static bool ggml_hexagon_tensor_is_non_host(const struct ggml_hexagon_session * sess, const struct ggml_tensor * t) { + return t && t->buffer && t->buffer->buft != &sess->host_buffer_type; +} + static bool ggml_hexagon_supported_mul_mat(const struct ggml_hexagon_session * sess, const struct ggml_tensor * dst) { const struct ggml_tensor * src0 = dst->src[0]; const struct ggml_tensor * src1 = dst->src[1]; @@ -2811,9 +3995,8 @@ static bool ggml_hexagon_supported_mul_mat(const struct ggml_hexagon_session * s return false; // no broadcasting (for now) } - // src0 (weights) must be repacked - if (src0->buffer && !ggml_backend_buffer_is_hexagon_repack(src0->buffer)) { - return false; + if (!src0->buffer) { + sess->needs_repack.insert(src0); } break; @@ -2872,9 +4055,8 @@ static bool ggml_hexagon_supported_mul_mat_id(const struct ggml_hexagon_session return false; } - // src0 (weights) must be repacked - if (src0->buffer && !ggml_backend_buffer_is_hexagon_repack(src0->buffer)) { - return false; + if (!src0->buffer) { + sess->needs_repack.insert(src0); } break; @@ -2970,7 +4152,7 @@ static bool ggml_hexagon_supported_unary(const struct ggml_hexagon_session * ses if (dst->type != GGML_TYPE_F32) { return false; } - if (ggml_is_permuted(src0)) { + if (!ggml_is_contiguous_rows(src0)) { return false; } if (!ggml_are_same_shape(src0, dst)) { @@ -3114,7 +4296,11 @@ static bool ggml_hexagon_supported_softmax(const struct ggml_hexagon_session * s static bool ggml_hexagon_supported_set_rows(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) { const struct ggml_tensor * src0 = op->src[0]; // values const struct ggml_tensor * src1 = op->src[1]; // indices - const struct ggml_tensor * dst = op; + const struct ggml_tensor * dst = op->src[2] ? op->src[2] : op; + + if (dst->type == GGML_TYPE_Q8_0 && src0->ne[0] < 32) { + return false; + } if (src0->type != GGML_TYPE_F32) { return false; @@ -3124,7 +4310,7 @@ static bool ggml_hexagon_supported_set_rows(const struct ggml_hexagon_session * return false; } - if (dst->type != GGML_TYPE_F16) { + if (dst->type != GGML_TYPE_F32 && dst->type != GGML_TYPE_F16 && dst->type != GGML_TYPE_Q8_0) { return false; } @@ -3138,7 +4324,11 @@ static bool ggml_hexagon_supported_get_rows(const struct ggml_hexagon_session * const struct ggml_tensor * src1 = op->src[1]; // indices const struct ggml_tensor * dst = op; - if (src0->type != GGML_TYPE_F32) { + if (src0->type != GGML_TYPE_F32 && src0->ne[0] < 32) { + return false; + } + + if (src0->type != GGML_TYPE_F32 && src0->type != GGML_TYPE_F16 && src0->type != GGML_TYPE_Q8_0) { return false; } @@ -3453,6 +4643,7 @@ static htp_op_code op_remap_to_htp(const ggml_tensor * t) { case GGML_OP_CLAMP: return HTP_OP_CLAMP; case GGML_OP_SQR: return HTP_OP_SQR; case GGML_OP_SQRT: return HTP_OP_SQRT; + case GGML_OP_LOG: return HTP_OP_UNARY_LOG; case GGML_OP_SOFT_MAX: return HTP_OP_SOFTMAX; case GGML_OP_SSM_CONV: return HTP_OP_SSM_CONV; case GGML_OP_GATED_DELTA_NET: return HTP_OP_GATED_DELTA_NET; @@ -3476,6 +4667,7 @@ static htp_op_code op_remap_to_htp(const ggml_tensor * t) { case GGML_UNARY_OP_EXP: return HTP_OP_UNARY_EXP; case GGML_UNARY_OP_SOFTPLUS: return HTP_OP_UNARY_SOFTPLUS; case GGML_UNARY_OP_TANH: return HTP_OP_UNARY_TANH; + case GGML_UNARY_OP_ABS: return HTP_OP_UNARY_ABS; default: break; } @@ -3521,133 +4713,22 @@ static bool is_mergeable_mul_mat(const ggml_tensor * t) { if (!t || t->op != GGML_OP_MUL_MAT) return false; if (t->src[1]->type != GGML_TYPE_F32) return false; return ggml_is_quantized(t->src[0]->type) && !mm_is_hmx_eligible(t); -} - -static bool is_mergeable_mul_mat_pair(const ggml_tensor * n1, const ggml_tensor * n2) { - if (!is_mergeable_mul_mat(n1) || !is_mergeable_mul_mat(n2)) { - return false; - } - if (n1->src[1] != n2->src[1]) { - return false; - } - if (n1->src[0]->ne[0] != n2->src[0]->ne[0] || - n1->src[0]->ne[1] != n2->src[0]->ne[1]) { - return false; - } - if (n1->src[0]->type != n2->src[0]->type) { - return false; - } - return true; -} - -static bool is_qkv_mergeable(const ggml_tensor * n_q, const ggml_tensor * n_k, const ggml_tensor * n_v) { - if (!is_mergeable_mul_mat(n_q) || !is_mergeable_mul_mat(n_k) || !is_mergeable_mul_mat(n_v)) { - return false; - } - if (n_q->src[1] != n_k->src[1] || n_q->src[1] != n_v->src[1]) { - return false; - } - if (n_q->src[0]->type != n_k->src[0]->type || n_q->src[0]->type != n_v->src[0]->type) { - return false; - } - if (n_k->src[0]->ne[0] != n_v->src[0]->ne[0] || - n_k->src[0]->ne[1] != n_v->src[0]->ne[1]) { - return false; - } - if (n_q->src[0]->ne[0] != n_k->src[0]->ne[0]) { - return false; - } - return true; -} - -static bool try_fuse_node(const ggml_hexagon_session * sess, const ggml_cgraph * graph, int & i, std::vector & nodes) { - if (!opt_opfusion) { - return false; - } - - ggml_tensor * n = graph->nodes[i]; - ggml_tensor * next_node = (i + 1 < graph->n_nodes) ? graph->nodes[i + 1] : nullptr; - - if (n->op == GGML_OP_RMS_NORM && next_node) { - if (next_node->op == GGML_OP_MUL && op_is_compute(next_node) && ggml_can_fuse(graph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL })) { - htp_opnode node(n, {}, HTP_OP_RMS_NORM_MUL); - node.add_fused(next_node); - - auto inputs = node.get_inputs(); - const struct ggml_tensor * src0 = inputs[0]; - const struct ggml_tensor * src1 = inputs.size() > 1 ? inputs[1] : nullptr; - ggml_hexagon_precompute_unary_params(sess, - node.opcode, src0, src1, node.dst(), - (struct htp_unary_kernel_params *)node.kernel_params - ); - - nodes.push_back(std::move(node)); - i++; // skip the fused MUL node - return true; - } - } - - if (is_mergeable_mul_mat(n)) { - ggml_tensor * n1 = (i + 1 < graph->n_nodes) ? graph->nodes[i + 1] : nullptr; - ggml_tensor * n2 = (i + 2 < graph->n_nodes) ? graph->nodes[i + 2] : nullptr; - if (is_qkv_mergeable(n, n1, n2)) { - struct htp_mm_kernel_params kparams; - ggml_hexagon_precompute_fused_qkv_params(sess, n1->src[0], n1->src[1], &kparams); - if ((size_t)kparams.vtcm_size <= sess->vtcm_size) { - // Reorder to KVQ: K (n1), V (n2), Q (n) - htp_opnode node(n1, {}, HTP_OP_MUL_MAT_QKV); - node.add_fused(n2, true); - node.add_fused(n, true); - memcpy(node.kernel_params, &kparams, sizeof(kparams)); - nodes.push_back(std::move(node)); - i += 2; - return true; - } else { - HEX_VERBOSE("ggml-hex: skip QKV fusion because VTCM needed (%d) > budget (%zu)\n", - kparams.vtcm_size, sess->vtcm_size); - } - } - if (is_mergeable_mul_mat_pair(n, n1)) { - struct htp_mm_kernel_params kparams; - ggml_hexagon_precompute_fused_ffn_params(sess, n->src[0], n->src[1], &kparams); - if ((size_t)kparams.vtcm_size <= sess->vtcm_size) { - htp_opnode node(n, {}, HTP_OP_MUL_MAT_FFN); - node.add_fused(n1, true); - memcpy(node.kernel_params, &kparams, sizeof(kparams)); - nodes.push_back(std::move(node)); - i += 1; - return true; - } else { - HEX_VERBOSE("ggml-hex: skip FFN fusion because VTCM needed (%d) > budget (%zu)\n", - kparams.vtcm_size, sess->vtcm_size); - } - } - } +} - if (n->op == GGML_OP_MUL_MAT && next_node) { - if (next_node->op == GGML_OP_ADD && op_is_compute(next_node) && ggml_can_fuse(graph, i, { GGML_OP_MUL_MAT, GGML_OP_ADD })) { - if (next_node->src[0] == n || next_node->src[1] == n) { - const struct ggml_tensor * src2 = (next_node->src[0] == n) ? next_node->src[1] : next_node->src[0]; - struct htp_mm_kernel_params kparams; - ggml_hexagon_precompute_fused_matmul_add_params(sess, n->src[0], n->src[1], src2, next_node, &kparams); - const int src1_nrows = n->src[1]->ne[1] * n->src[1]->ne[2] * n->src[1]->ne[3]; - const bool can_fuse = (kparams.n_hmx > 0) || (src1_nrows == 1); - if (can_fuse && (size_t)kparams.vtcm_size <= sess->vtcm_size) { - htp_opnode node(n, {}, HTP_OP_MUL_MAT_ADD); - node.add_fused(next_node); - memcpy(node.kernel_params, &kparams, sizeof(kparams)); - nodes.push_back(std::move(node)); - i += 1; - return true; - } else if (can_fuse) { - HEX_VERBOSE("ggml-hex: skip MUL_MAT_ADD fusion because VTCM needed (%d) > budget (%zu)\n", - kparams.vtcm_size, sess->vtcm_size); - } - } - } +static bool is_mergeable_mul_mat_pair(const ggml_tensor * n1, const ggml_tensor * n2) { + if (!is_mergeable_mul_mat(n1) || !is_mergeable_mul_mat(n2)) { + return false; } - - return false; + if (n1->src[1] != n2->src[1]) { + return false; + } + if (n1->src[0]->ne[0] != n2->src[0]->ne[0]) { + return false; + } + if (n1->src[0]->type != n2->src[0]->type) { + return false; + } + return true; } static ggml_status ggml_backend_hexagon_graph_compute(ggml_backend_t backend, ggml_cgraph * graph) { @@ -3659,24 +4740,34 @@ static ggml_status ggml_backend_hexagon_graph_compute(ggml_backend_t backend, gg std::vector computed_nodes; // Check for cache hit - bool cache_hit = (graph->uid != 0 && sess->cached_graph.uid == graph->uid); + bool cache_hit = (graph->uid != 0 && sess->cached_uid == graph->uid); if (cache_hit) { - nodes_ptr = &sess->cached_graph.htp_nodes; + nodes_ptr = &sess->cached_nodes; } else { + // Tag fusable tensors in graph + for (int i = 0; i < graph->n_nodes; i++) { + auto * extra = (ggml_hexagon_tensor_extra *) graph->nodes[i]->extra; + if (!extra) continue; + + if (graph->nodes[i]->op == GGML_OP_RMS_NORM && ggml_can_fuse(graph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL })) { + extra->flags |= GGML_HEXAGON_TENSOR_FUSEABLE; + } else if (graph->nodes[i]->op == GGML_OP_MUL_MAT) { + if ((i + 1 < graph->n_nodes && graph->nodes[i + 1]->op == GGML_OP_ADD && ggml_can_fuse(graph, i, { GGML_OP_MUL_MAT, GGML_OP_ADD })) || + ggml_node_has_n_uses(graph, i, 1)) { + extra->flags |= GGML_HEXAGON_TENSOR_FUSEABLE; + } + } + } + computed_nodes.reserve(graph->n_nodes); - // Fuse and finalize for (int i = 0; i < graph->n_nodes; ++i) { ggml_tensor * n = graph->nodes[i]; if (!op_is_compute(n)) { continue; } - if (try_fuse_node(sess, graph, i, computed_nodes)) { - continue; - } - - htp_opnode node(n, {}, HTP_OP_INVALID); + htp_opnode node(HTP_OP_INVALID, n); node.opcode = op_remap_to_htp(n); if (node.opcode == HTP_OP_MUL_MAT || node.opcode == HTP_OP_MUL_MAT_ID) { ggml_hexagon_precompute_matmul_params(sess, @@ -3696,29 +4787,34 @@ static ggml_status ggml_backend_hexagon_graph_compute(ggml_backend_t backend, gg node.opcode, src0, src1, node.dst(), (struct htp_unary_kernel_params *)node.kernel_params ); + } else if (node.opcode == HTP_OP_GET_ROWS) { + ggml_hexagon_precompute_get_rows_params(sess, + node.node->src[0], node.node->src[1], node.dst(), + (struct htp_get_rows_kernel_params *)node.kernel_params + ); + } else if (node.opcode == HTP_OP_SET_ROWS) { + ggml_hexagon_precompute_set_rows_params(sess, + node.node->src[0], node.node->src[1], node.dst(), + (struct htp_set_rows_kernel_params *)node.kernel_params + ); } computed_nodes.push_back(std::move(node)); } if (graph->uid != 0) { - sess->cached_graph.uid = graph->uid; - sess->cached_graph.htp_nodes = std::move(computed_nodes); - nodes_ptr = &sess->cached_graph.htp_nodes; + sess->cached_uid = graph->uid; + sess->cached_nodes = std::move(computed_nodes); + nodes_ptr = &sess->cached_nodes; } else { nodes_ptr = &computed_nodes; } } // Queue and execute - if (opt_opstage & HTP_OPSTAGE_QUEUE) { - for (const auto & node : *nodes_ptr) { - sess->enqueue_op(node); - } + for (const auto & node : *nodes_ptr) { + sess->enqueue_op(node); } - // Wait until all pending ops complete - sess->flush(); - return GGML_STATUS_SUCCESS; } @@ -3731,6 +4827,106 @@ static void ggml_backend_hexagon_synchronize(ggml_backend_t backend) { sess->flush(); } +enum ggml_hexagon_mem_range_type { + HEXAGON_MEM_RANGE_TYPE_SRC, + HEXAGON_MEM_RANGE_TYPE_DST, +}; + +struct ggml_hexagon_mem_range { + uint64_t pb; + uint64_t p0; + uint64_t p1; + ggml_hexagon_mem_range_type pt; +}; + +struct ggml_hexagon_mem_ranges { + std::vector ranges; + + void reset() { + ranges.clear(); + } + + void add(const ggml_hexagon_mem_range & mr) { + ranges.push_back(mr); + } + + bool check(const ggml_hexagon_mem_range & mr) const { + for (const auto & cmp : ranges) { + if (mr.pb != cmp.pb) { + continue; + } + if (mr.pt == HEXAGON_MEM_RANGE_TYPE_SRC && cmp.pt == HEXAGON_MEM_RANGE_TYPE_SRC) { + continue; + } + if (mr.p0 < cmp.p1 && mr.p1 > cmp.p0) { + return false; + } + } + return true; + } +}; + +static ggml_hexagon_mem_range ggml_hexagon_mem_range_from_tensor(const ggml_tensor * tensor, ggml_hexagon_mem_range_type pt) { + const ggml_tensor * base = tensor->view_src ? tensor->view_src : tensor; + ggml_hexagon_mem_range mr; + if (tensor->buffer) { + mr = { + /*.pb =*/ (uint64_t) tensor->buffer, + /*.p0 =*/ (uint64_t) tensor->data, + /*.p1 =*/ (uint64_t) tensor->data + ggml_backend_buft_get_alloc_size(tensor->buffer->buft, tensor), + /*.pt =*/ pt, + }; + } else { + mr = { + /*.pb =*/ (uint64_t) base, + /*.p0 =*/ 0, + /*.p1 =*/ 1024, + /*.pt =*/ pt, + }; + } + return mr; +} + +static void ggml_hexagon_mem_ranges_add_node(ggml_hexagon_mem_ranges & mrs, const htp_opnode & node) { + if (node.is_empty()) return; + + for (int i = 0; i < GGML_MAX_SRC; i++) { + if (node.node->src[i]) { + mrs.add(ggml_hexagon_mem_range_from_tensor(node.node->src[i], HEXAGON_MEM_RANGE_TYPE_SRC)); + } + } + for (const auto * fused : node.fused) { + for (int i = 0; i < GGML_MAX_SRC; i++) { + if (fused->src[i]) { + mrs.add(ggml_hexagon_mem_range_from_tensor(fused->src[i], HEXAGON_MEM_RANGE_TYPE_SRC)); + } + } + } + mrs.add(ggml_hexagon_mem_range_from_tensor(node.dst(), HEXAGON_MEM_RANGE_TYPE_DST)); +} + +static bool ggml_hexagon_mem_ranges_check_node(const ggml_hexagon_mem_ranges & mrs, const htp_opnode & node) { + if (node.is_empty()) return true; + + for (int i = 0; i < GGML_MAX_SRC; i++) { + if (node.node->src[i]) { + if (!mrs.check(ggml_hexagon_mem_range_from_tensor(node.node->src[i], HEXAGON_MEM_RANGE_TYPE_SRC))) { + return false; + } + } + } + for (const auto * fused : node.fused) { + for (int i = 0; i < GGML_MAX_SRC; i++) { + if (fused->src[i]) { + if (!mrs.check(ggml_hexagon_mem_range_from_tensor(fused->src[i], HEXAGON_MEM_RANGE_TYPE_SRC))) { + return false; + } + } + } + } + return mrs.check(ggml_hexagon_mem_range_from_tensor(node.dst(), HEXAGON_MEM_RANGE_TYPE_DST)); +} + static std::vector ggml_hexagon_graph_optimize_reorder(const std::vector & nodes) { const int n = nodes.size(); @@ -3739,28 +4935,32 @@ static std::vector ggml_hexagon_graph_optimize_reorder(const std::vector used(n, false); - // The main goal here is to stack the MUL_MAT ops with the same src1 input. - // This allows use to reuse dynamically quantized src1 in VTCM. + ggml_hexagon_mem_ranges mrs; - // TODO: the current version might do incorrect reordering in cases where quantized src0 - // input is an output of another Op. + // The main goal here is to stack the MUL_MAT ops with the same src1 input. + // This allows us to reuse dynamically quantized src1 in VTCM. for (int i0 = 0; i0 < n; i0++) { if (used[i0]) { continue; } - res.push_back(i0); - const auto & node0 = nodes[i0]; if (!node0.stackable()) { + res.push_back(i0); + used[i0] = true; continue; } // that many nodes forward to search for stackable nodes that can reuse VTCM constexpr int N_FORWARD = 16; + std::vector stack; + stack.push_back(i0); + + mrs.reset(); + for (int i1 = i0 + 1; i1 < i0 + N_FORWARD && i1 < n; i1++) { if (used[i1]) { continue; @@ -3768,11 +4968,17 @@ static std::vector ggml_hexagon_graph_optimize_reorder(const std::vector nodes; nodes.reserve(gf->n_nodes); - // fuse nodes: - // we don't want to make reorders that break fusing, so we first pack all fusable tensors - // and perform the reorder over the fused nodes. after the reorder is done, we unfuse + // Pack nodes for reordering for (int i = 0; i < n; i++) { - htp_opnode node = { - /*.node =*/gf->nodes[i], - /*.fused =*/{}, - }; + htp_opnode node(HTP_OP_INVALID, gf->nodes[i]); // fuse only ops that start with these operations // can be expanded when needed @@ -3856,22 +5057,200 @@ static void ggml_backend_hexagon_graph_optimize(ggml_backend_t backend, ggml_cgr GGML_UNUSED(backend); } +static bool ggml_hexagon_cpy_tensor_async_phys(ggml_backend_t backend_src, ggml_backend_t backend_dst, const ggml_tensor * src, ggml_tensor * dst) { + auto sess_src = static_cast(backend_src->context); + auto sess_dst = static_cast(backend_dst->context); + auto sbuf_dst = (ggml_hexagon_shared_buffer *) dst->buffer->context; + + if (sess_dst->fence_seq == 0) sess_dst->fence_seq = 1; + uint32_t fence_seq = sess_dst->fence_seq++; + if (sess_dst->fence_seq == 0) sess_dst->fence_seq = 1; + + volatile uint32_t * fence = (volatile uint32_t *) sbuf_dst->alloc_fence(); + + HEX_VERBOSE("ggml-hex: %s cpy-tensor-async %s -> %s size %zu : seq %u\n", + sess_dst->name.c_str(), src->name, dst->name, ggml_nbytes(src), fence_seq); + + // dummy extra (must be static) + static ggml_hexagon_tensor_extra fence_extra { {}, 0, GGML_HEXAGON_TENSOR_FENCE }; + + ggml_tensor fence_tensor {}; + fence_tensor.buffer = dst->buffer; + fence_tensor.extra = &fence_extra; + fence_tensor.data = (void *) fence; + fence_tensor.type = GGML_TYPE_I32; + fence_tensor.ne[0] = 1; + fence_tensor.ne[1] = 1; + fence_tensor.ne[2] = 1; + fence_tensor.ne[3] = 1; + fence_tensor.nb[0] = sizeof(int32_t); + fence_tensor.nb[1] = sizeof(int32_t); + fence_tensor.nb[2] = sizeof(int32_t); + fence_tensor.nb[3] = sizeof(int32_t); + fence_tensor.op = GGML_OP_NONE; + + sess_src->enqueue_cpy(src, dst, &fence_tensor, fence_seq); + sess_dst->enqueue_fence(&fence_tensor, fence_seq); + + sess_dst->add_sync_peer(sess_src); + + return true; +} + +static bool ggml_hexagon_cpy_tensor_async_virt(ggml_backend_t backend_src, ggml_backend_t backend_dst, const ggml_tensor * src, ggml_tensor * dst) { + auto sess_src = static_cast(backend_src->context); + auto sess_dst = static_cast(backend_dst->context); + auto sbuf_dst = (ggml_hexagon_shared_buffer *) dst->buffer->context; + + if (!sess_src->clone_buffer(sbuf_dst)) { return false; } + + HEX_VERBOSE("ggml-hex: %s cpy-tensor-async %s -> %s size %zu\n", + sess_dst->name.c_str(), src->name, dst->name, ggml_nbytes(src)); + + sess_src->enqueue_cpy(src, dst); + sess_src->flush(true); + + return true; +} + +static bool ggml_backend_hexagon_cpy_tensor_async(ggml_backend_t backend_src, ggml_backend_t backend_dst, const ggml_tensor * src, ggml_tensor * dst) { + if (!ggml_backend_is_hexagon(backend_src) || !ggml_backend_is_hexagon(backend_dst)) { + return false; + } + + *(ggml_hexagon_tensor_extra *) dst->extra = *(const ggml_hexagon_tensor_extra *) src->extra; + + auto sess_src = static_cast(backend_src->context); + auto sess_dst = static_cast(backend_dst->context); + + if (sess_src == sess_dst) { + HEX_VERBOSE("ggml-hex: %s cpy-tensor-async %s -> %s size %zu\n", sess_dst->name.c_str(), src->name, dst->name, ggml_nbytes(src)); + sess_src->enqueue_cpy(src, dst); + sess_src->flush_batch(); + return true; + } + + if (sess_src->phys_idx != sess_dst->phys_idx) + return ggml_hexagon_cpy_tensor_async_phys(backend_src, backend_dst, src, dst); + + return ggml_hexagon_cpy_tensor_async_virt(backend_src, backend_dst, src, dst); +} + +static ggml_backend_event_t ggml_backend_hexagon_device_event_new(ggml_backend_dev_t dev) { + ggml_hexagon_event * hex_event = new ggml_hexagon_event(); + HEX_VERBOSE("ggml-hex: %s event-new : event %p\n", ggml_backend_dev_name(dev), (void *)hex_event); + + return new ggml_backend_event { + /* .device = */ dev, + /* .context = */ hex_event, + }; +} + +static void ggml_backend_hexagon_device_event_free(ggml_backend_dev_t dev, ggml_backend_event_t event) { + GGML_UNUSED(dev); + + if (event == nullptr) { + return; + } + + ggml_hexagon_event * hex_event = (ggml_hexagon_event *)event->context; + HEX_VERBOSE("ggml-hex: %s event-free : event %p\n", ggml_backend_dev_name(dev), (void *)hex_event); + delete hex_event; + delete event; +} + +static void ggml_backend_hexagon_device_event_synchronize(ggml_backend_dev_t dev, ggml_backend_event_t event) { + GGML_UNUSED(dev); + + ggml_hexagon_event * hex_event = (ggml_hexagon_event *)event->context; + HEX_VERBOSE("ggml-hex: %s event-synchronize : event %p seq %llu\n", + ggml_backend_dev_name(dev), (void *)hex_event, (unsigned long long)hex_event->seq); + if (hex_event->sess != nullptr) { + hex_event->sess->wait_event(hex_event->seq); + } +} + +static void ggml_backend_hexagon_event_record(ggml_backend_t backend, ggml_backend_event_t event) { + auto sess = static_cast(backend->context); + ggml_hexagon_event * hex_event = (ggml_hexagon_event *)event->context; + + hex_event->sess = sess; + hex_event->seq = sess->record_event(); + HEX_VERBOSE("ggml-hex: %s event-record : event %p seq %llu\n", + sess->c_name(), (void *)hex_event, (unsigned long long)hex_event->seq); +} + +static void ggml_backend_hexagon_event_wait(ggml_backend_t backend, ggml_backend_event_t event) { + GGML_UNUSED(backend); + + ggml_hexagon_event * hex_event = (ggml_hexagon_event *)event->context; + if (hex_event->sess != nullptr) { + HEX_VERBOSE("ggml-hex: %s event-wait : event %p seq %llu\n", + hex_event->sess->c_name(), (void *)hex_event, (unsigned long long)hex_event->seq); + hex_event->sess->wait_event(hex_event->seq); + } +} + +static void ggml_backend_hexagon_set_tensor_async(ggml_backend_t backend, struct ggml_tensor * tensor, const void * data, size_t offset, size_t size) { + auto sess = static_cast(backend->context); + HEX_VERBOSE("ggml-hex: %s set-tensor-async %s : data %p offset %zu size %zu usage %d\n", + sess->c_name(), tensor->name, data, offset, size, tensor->buffer ? (int) tensor->buffer->usage : -1); + ggml_backend_tensor_set(tensor, data, offset, size); +} + +static void ggml_backend_hexagon_get_tensor_async(ggml_backend_t backend, const struct ggml_tensor * tensor, void * data, size_t offset, size_t size) { + auto sess = static_cast(backend->context); + HEX_VERBOSE("ggml-hex: %s get-tensor-async %s : data %p offset %zu size %zu usage %d\n", + sess->c_name(), tensor->name, data, offset, size, tensor->buffer ? (int) tensor->buffer->usage : -1); + sess->flush(true); + ggml_backend_tensor_get(tensor, data, offset, size); +} + +static void ggml_backend_hexagon_set_tensor_2d_async(ggml_backend_t backend, + struct ggml_tensor * tensor, + const void * data, + size_t offset, + size_t size, + size_t n_copies, + size_t stride_tensor, + size_t stride_data) { + auto sess = static_cast(backend->context); + HEX_VERBOSE("ggml-hex: %s set-tensor-2d-async %s : data %p offset %zu size %zu n_copies %zu stride_tensor %zu stride_data %zu usage %d\n", + sess->c_name(), tensor->name, data, offset, size, n_copies, stride_tensor, stride_data, tensor->buffer ? (int) tensor->buffer->usage : -1); + ggml_backend_tensor_set_2d(tensor, data, offset, size, n_copies, stride_tensor, stride_data); +} + +static void ggml_backend_hexagon_get_tensor_2d_async(ggml_backend_t backend, + const struct ggml_tensor * tensor, + void * data, + size_t offset, + size_t size, + size_t n_copies, + size_t stride_tensor, + size_t stride_data) { + auto sess = static_cast(backend->context); + HEX_VERBOSE("ggml-hex: %s get-tensor-2d-async %s : data %p offset %zu size %zu n_copies %zu stride_tensor %zu stride_data %zu usage %d\n", + sess->c_name(), tensor->name, data, offset, size, n_copies, stride_tensor, stride_data, tensor->buffer ? (int) tensor->buffer->usage : -1); + sess->flush(true); + ggml_backend_tensor_get_2d(tensor, data, offset, size, n_copies, stride_tensor, stride_data); +} + static struct ggml_backend_i hexagon_backend_i = { /* .get_name = */ ggml_backend_hexagon_name, /* .free = */ ggml_backend_hexagon_free, - /* .set_tensor_async = */ NULL, - /* .get_tensor_async = */ NULL, - /* .set_tensor_2d_async = */ NULL, - /* .get_tensor_2d_async = */ NULL, - /* .cpy_tensor_async = */ NULL, + /* .set_tensor_async = */ ggml_backend_hexagon_set_tensor_async, + /* .get_tensor_async = */ ggml_backend_hexagon_get_tensor_async, + /* .set_tensor_2d_async = */ ggml_backend_hexagon_set_tensor_2d_async, + /* .get_tensor_2d_async = */ ggml_backend_hexagon_get_tensor_2d_async, + /* .cpy_tensor_async = */ ggml_backend_hexagon_cpy_tensor_async, /* .synchronize = */ ggml_backend_hexagon_synchronize, /* .graph_plan_create = */ NULL, /* .graph_plan_free = */ NULL, /* .graph_plan_update = */ NULL, /* .graph_plan_compute = */ NULL, /* .graph_compute = */ ggml_backend_hexagon_graph_compute, - /* .event_record = */ NULL, - /* .event_wait = */ NULL, + /* .event_record = */ ggml_backend_hexagon_event_record, + /* .event_wait = */ ggml_backend_hexagon_event_wait, /* .graph_optimize = */ ggml_backend_hexagon_graph_optimize, }; @@ -3932,9 +5311,9 @@ static void ggml_backend_hexagon_device_get_props(ggml_backend_dev_t dev, struct ggml_backend_hexagon_device_get_memory(dev, &props->memory_free, &props->memory_total); props->caps = { /* .async = */ true, - /* .host_buffer = */ (bool) opt_hostbuf, + /* .host_buffer = */ false, /* .buffer_from_host_ptr = */ false, - /* .events = */ false, + /* .events = */ true, /* .mmap_support = */ false, }; } @@ -3944,32 +5323,12 @@ static ggml_backend_buffer_type_t ggml_backend_hexagon_device_get_buffer_type(gg return &sess->buffer_type; } -static ggml_backend_buffer_type_t ggml_backend_hexagon_device_get_repack_buffer_type(ggml_backend_dev_t dev) { - auto sess = static_cast(dev->context); - return &sess->repack_buffer_type; -} - -static bool ggml_hexagon_supported_buffer(ggml_hexagon_session *sess, const struct ggml_tensor * t) { - if (t && t->buffer) { - if (ggml_backend_buffer_is_hexagon(t->buffer) == false) return false; // not our buffer - if (ggml_backend_hexagon_buffer_get_sess(t->buffer) != sess) return false; // wrong session - } - return true; -} - -static bool ggml_hexagon_supported_buffers(ggml_hexagon_session *sess, const struct ggml_tensor * t) { - // all srcs & dsts must be mapped to the same session - if (!ggml_hexagon_supported_buffer(sess, t)) { - return false; - } - - for (int i = 0; i < GGML_MAX_SRC; i++) { - if (!ggml_hexagon_supported_buffer(sess, t->src[i])) { - return false; - } +static ggml_backend_buffer_type_t ggml_backend_hexagon_device_get_host_buffer_type(ggml_backend_dev_t dev) { + if (!opt_hostbuf) { + return NULL; } - - return true; + auto sess = static_cast(dev->context); + return &sess->host_buffer_type; } static bool ggml_hexagon_supported_cpy(const struct ggml_hexagon_session * sess, const struct ggml_tensor * op) { @@ -4067,12 +5426,6 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons return false; } - // all srcs & dsts must be mapped to the same session - if (!ggml_hexagon_supported_buffers(sess, op)) { - ggml_hexagon_dump_op_supp(sess->name, op, false); - return false; - } - bool supp = false; switch (op->op) { case GGML_OP_NONE: @@ -4112,6 +5465,7 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons case GGML_OP_SQR: case GGML_OP_SQRT: + case GGML_OP_LOG: supp = ggml_hexagon_supported_unary(sess, op); break; @@ -4130,6 +5484,7 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons case GGML_UNARY_OP_SIGMOID: case GGML_UNARY_OP_SOFTPLUS: case GGML_UNARY_OP_TANH: + case GGML_UNARY_OP_ABS: case GGML_UNARY_OP_SILU: case GGML_UNARY_OP_GELU: case GGML_UNARY_OP_GELU_QUICK: @@ -4233,31 +5588,20 @@ static bool ggml_backend_hexagon_device_supports_op(ggml_backend_dev_t dev, cons } static bool ggml_backend_hexagon_device_supports_buft(ggml_backend_dev_t dev, ggml_backend_buffer_type_t buft) { - if (buft->iface.get_alignment != ggml_backend_hexagon_buffer_type_get_alignment) { - return false; - } - - auto s0 = static_cast(dev->context); - auto s1 = static_cast(buft->context)->sess; - - // Need session/domain-id for buffers to be compatible - bool supp = (s0->session_id == s1->session_id); + auto sess = static_cast(dev->context); - HEX_VERBOSE("ggml-hex: %s device-supports-buft %s (%d)\n", s0->name.c_str(), s1->name.c_str(), (int) supp); + // Technically we can clone hexagon buffers from any session but for some reason the output is garbled with layer-split, + // tensor-split works correctly, so it needs mode debugging and investigation. For now accept only our own buffers. +#if 0 + bool supp = (buft->iface.get_alignment == ggml_backend_hexagon_buffer_type_get_alignment); +#else + bool supp = (buft == &sess->host_buffer_type) || (buft == &sess->buffer_type); +#endif + HEX_VERBOSE("ggml-hex: %s device-supports-buft %s %s\n", sess->name.c_str(), ggml_backend_buft_name(buft), supp ? "yes" : "no"); return supp; } -static ggml_backend_buffer_type_t * ggml_backend_hexagon_device_get_extra_buffers_type(ggml_backend_dev_t dev) { - auto s0 = static_cast(dev->context); - HEX_VERBOSE("ggml-hex: device-get-extra-buft : %s \n", s0->name.c_str()); - - static ggml_backend_buffer_type_t bufts[2]; - bufts[0] = ggml_backend_hexagon_device_get_repack_buffer_type(dev); - bufts[1] = NULL; - return bufts; -} - static const struct ggml_backend_device_i ggml_backend_hexagon_device_i = { /* .get_name = */ ggml_backend_hexagon_device_get_name, /* .get_description = */ ggml_backend_hexagon_device_get_description, @@ -4266,27 +5610,18 @@ static const struct ggml_backend_device_i ggml_backend_hexagon_device_i = { /* .get_props = */ ggml_backend_hexagon_device_get_props, /* .init_backend = */ ggml_backend_hexagon_device_init, /* .get_buffer_type = */ ggml_backend_hexagon_device_get_buffer_type, - /* .get_host_buffer_type = */ NULL, // ggml_backend_hexagon_device_get_host_buffer_type, + /* .get_host_buffer_type = */ ggml_backend_hexagon_device_get_host_buffer_type, /* .buffer_from_host_ptr = */ NULL, // ggml_backend_hexagon_device_buffer_from_ptr, /* .supports_op = */ ggml_backend_hexagon_device_supports_op, /* .supports_buft = */ ggml_backend_hexagon_device_supports_buft, /* .offload_op = */ NULL, // ggml_backend_hexagon_device_offload_op, - /* .event_new = */ NULL, - /* .event_free = */ NULL, - /* .event_synchronize = */ NULL, + /* .event_new = */ ggml_backend_hexagon_device_event_new, + /* .event_free = */ ggml_backend_hexagon_device_event_free, + /* .event_synchronize = */ ggml_backend_hexagon_device_event_synchronize, }; //** backend registry -#define GGML_HEXAGON_MAX_SESSIONS 16 - -struct ggml_hexagon_registry { - ggml_hexagon_registry(ggml_backend_reg_t reg); - ~ggml_hexagon_registry(); - - ggml_backend_device devices[GGML_HEXAGON_MAX_SESSIONS]; -}; - ggml_hexagon_registry::ggml_hexagon_registry(ggml_backend_reg_t reg) { GGML_LOG_INFO("ggml-hex: Hexagon backend (experimental) : allocating new registry : ndev %zu\n", opt_ndev); @@ -4303,6 +5638,7 @@ ggml_hexagon_registry::ggml_hexagon_registry(ggml_backend_reg_t reg) { devices[i].context = nullptr; } } + } ggml_hexagon_registry::~ggml_hexagon_registry() { @@ -4335,14 +5671,129 @@ static ggml_backend_dev_t ggml_backend_hexagon_reg_get_device(ggml_backend_reg_t return &hreg->devices[index]; } -static void * ggml_backend_hexagon_get_proc_address(ggml_backend_reg_t reg, const char * name) { - if (strcmp(name, "ggml_backend_dev_get_extra_bufts") == 0 && opt_hostbuf) { - ggml_backend_dev_get_extra_bufts_t fct = ggml_backend_hexagon_device_get_extra_buffers_type; - return (void *) fct; +// ** communication context for tensor-split allreduce + +static void * ggml_backend_hexagon_comm_init(ggml_backend_t * backends, size_t n_backends) { + if (n_backends < 2 || n_backends > 4) { + return nullptr; } - return NULL; + for (size_t i = 0; i < n_backends; ++i) { + if (!ggml_backend_is_hexagon(backends[i])) { + return nullptr; + } + } + + auto * ctx = new ggml_backend_hexagon_comm_context(); + ctx->backends.assign(backends, backends + n_backends); + ctx->n_backends = n_backends; + ctx->fence_seq = (((uintptr_t) ctx) & 0xFFFF) | 1; + + return ctx; +} + +static void ggml_backend_hexagon_comm_free(void * comm_ctx_v) { + if (!comm_ctx_v) return; + delete static_cast(comm_ctx_v); +} + +static bool ggml_backend_hexagon_comm_allreduce_tensor(void * comm_ctx_v, struct ggml_tensor ** tensors) { + if (opt_ar_select == 0 || !comm_ctx_v) return false; + auto * comm_ctx = static_cast(comm_ctx_v); + const size_t n_backends = comm_ctx->n_backends; + + if (n_backends < 2 || n_backends > 4) return false; + + for (size_t i = 0; i < n_backends; i++) { + if (!tensors[i] || !tensors[i]->buffer || !ggml_backend_buffer_is_hexagon(tensors[i]->buffer)) { + return false; + } + if (tensors[i]->type != tensors[0]->type) { + return false; + } + if (!ggml_is_contiguous(tensors[i])) { + return false; + } + if (ggml_nelements(tensors[i]) != ggml_nelements(tensors[0])) { + return false; + } + } + + if (tensors[0]->type != GGML_TYPE_F16 && tensors[0]->type != GGML_TYPE_F32) { + return false; + } + + for (size_t r = 0; r < n_backends; r++) { + auto sess = static_cast(comm_ctx->backends[r]->context); + struct htp_allreduce_kernel_params kparams; + if (!ggml_hexagon_precompute_allreduce_params(sess, tensors[r], (uint32_t) r, (uint32_t) n_backends, false, false, &kparams)) { + return false; + } + } + + if (comm_ctx->fence_seq == 0) comm_ctx->fence_seq = 1; + uint32_t fence_seq_entry = comm_ctx->fence_seq++; + if (comm_ctx->fence_seq == 0) comm_ctx->fence_seq = 1; + uint32_t fence_seq_exit = comm_ctx->fence_seq++; + if (comm_ctx->fence_seq == 0) comm_ctx->fence_seq = 1; + + volatile uint32_t * fences[GGML_HEXAGON_MAX_SESSIONS]; + for (size_t i = 0; i < n_backends; i++) { + auto sbuf = (ggml_hexagon_shared_buffer *) tensors[i]->buffer->context; + fences[i] = (volatile uint32_t *) sbuf->alloc_fence(); + } + + static ggml_hexagon_tensor_extra fence_extra { {}, 0, GGML_HEXAGON_TENSOR_FENCE }; + ggml_tensor fence_tensors[GGML_HEXAGON_MAX_SESSIONS]; + for (size_t i = 0; i < n_backends; i++) { + fence_tensors[i] = {}; + fence_tensors[i].buffer = tensors[i]->buffer; + fence_tensors[i].extra = &fence_extra; + fence_tensors[i].data = (void *) fences[i]; + fence_tensors[i].type = GGML_TYPE_I32; + fence_tensors[i].ne[0] = 4; + fence_tensors[i].ne[1] = 1; + fence_tensors[i].ne[2] = 1; + fence_tensors[i].ne[3] = 1; + fence_tensors[i].nb[0] = sizeof(int32_t); + fence_tensors[i].nb[1] = sizeof(int32_t); + fence_tensors[i].nb[2] = sizeof(int32_t); + fence_tensors[i].nb[3] = sizeof(int32_t); + fence_tensors[i].op = GGML_OP_NONE; + } + + std::vector data_tensors(n_backends); + std::vector sync_tensors(n_backends); + for (size_t i = 0; i < n_backends; i++) { + data_tensors[i] = tensors[i]; + sync_tensors[i] = &fence_tensors[i]; + } + + for (size_t r = 0; r < n_backends; r++) { + auto sess = static_cast(comm_ctx->backends[r]->context); + sess->enqueue_allreduce(tensors[r], data_tensors, sync_tensors, (uint32_t) r, (uint32_t) n_backends, fence_seq_entry, fence_seq_exit); + for (size_t j = 0; j < n_backends; j++) { + if (r != j) { + sess->add_sync_peer(static_cast(comm_ctx->backends[j]->context)); + } + } + } + + return true; +} + +static void * ggml_backend_hexagon_get_proc_address(ggml_backend_reg_t reg, const char * name) { GGML_UNUSED(reg); + if (strcmp(name, "ggml_backend_comm_init") == 0) { + return (void *) ggml_backend_hexagon_comm_init; + } + if (strcmp(name, "ggml_backend_comm_free") == 0) { + return (void *) ggml_backend_hexagon_comm_free; + } + if (strcmp(name, "ggml_backend_comm_allreduce_tensor") == 0) { + return (void *) ggml_backend_hexagon_comm_allreduce_tensor; + } + return NULL; } template std::vector str_to_vec(const char* str) { @@ -4379,8 +5830,6 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) { "please update hexagon_type to match ggml_type"); const char * str_verbose = getenv("GGML_HEXAGON_VERBOSE"); - const char * str_hostbuf = getenv("GGML_HEXAGON_HOSTBUF"); - const char * str_opstage = getenv("GGML_HEXAGON_OPSTAGE"); const char * str_opbatch = getenv("GGML_HEXAGON_OPBATCH"); const char * str_opqueue = getenv("GGML_HEXAGON_OPQUEUE"); const char * str_oppoll = getenv("GGML_HEXAGON_OPPOLL"); @@ -4389,15 +5838,16 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) { const char * str_profile = getenv("GGML_HEXAGON_PROFILE"); const char * str_etm = getenv("GGML_HEXAGON_ETM"); const char * str_nhvx = getenv("GGML_HEXAGON_NHVX"); - const char * str_use_hmx = getenv("GGML_HEXAGON_USE_HMX"); const char * str_nhmx = getenv("GGML_HEXAGON_NHMX"); const char * str_mm_select = getenv("GGML_HEXAGON_MM_SELECT"); const char * str_fa_select = getenv("GGML_HEXAGON_FA_SELECT"); + const char * str_ar_select = getenv("GGML_HEXAGON_AR_SELECT"); const char * str_ndev = getenv("GGML_HEXAGON_NDEV"); const char * str_arch = getenv("GGML_HEXAGON_ARCH"); const char * str_vmem = getenv("GGML_HEXAGON_VMEM"); const char * str_mbuf = getenv("GGML_HEXAGON_MBUF"); const char * str_optrace = getenv("GGML_HEXAGON_OPTRACE"); + const char * str_hostbuf = getenv("GGML_HEXAGON_HOSTBUF"); // Init Arch first since it affects other defaults if (!str_arch) { @@ -4430,8 +5880,6 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) { opt_opfilter = str_opfilter ? new std::regex(str_opfilter, RE_ICASE) : NULL; opt_verbose = str_verbose ? atoi(str_verbose) : 0; - opt_hostbuf = str_hostbuf ? atoi(str_hostbuf) : opt_hostbuf; - opt_opstage = str_opstage ? strtoul(str_opstage, NULL, 0) : opt_opstage; opt_opbatch = str_opbatch ? strtoul(str_opbatch, NULL, 0) : opt_opbatch; opt_opqueue = str_opqueue ? strtoul(str_opqueue, NULL, 0) : opt_opqueue; opt_optrace = str_optrace ? strtoul(str_optrace, NULL, 0) : (opt_opbatch * 256); @@ -4440,16 +5888,90 @@ static void ggml_hexagon_init(ggml_backend_reg * reg) { opt_profile = str_profile ? atoi(str_profile) : 0; opt_etm = str_etm ? atoi(str_etm) : 0; opt_nhvx = str_nhvx ? strtoul(str_nhvx, NULL, 0) : opt_nhvx; - opt_nhmx = str_nhmx ? atoi(str_nhmx) : (str_use_hmx ? atoi(str_use_hmx) : opt_nhmx); + opt_nhmx = str_nhmx ? atoi(str_nhmx) : opt_nhmx; opt_mm_select = str_mm_select ? atoi(str_mm_select) : opt_mm_select; opt_fa_select = str_fa_select ? atoi(str_fa_select) : opt_fa_select; - opt_ndev = str_ndev ? strtoul(str_ndev, NULL, 0) : opt_ndev; - opt_hostbuf = str_hostbuf ? atoi(str_hostbuf) : opt_hostbuf; + opt_ar_select = str_ar_select ? atoi(str_ar_select) : opt_ar_select; opt_mbuf = str_mbuf ? strtoul(str_mbuf, NULL, 0) * MiB : opt_mbuf; opt_vmem = str_vmem ? strtoul(str_vmem, NULL, 0) * MiB : opt_vmem; + opt_hostbuf = str_hostbuf ? atoi(str_hostbuf) != 0 : opt_hostbuf; + + // Parse device configuration + const char * str_devices = getenv("GGML_HEXAGON_DEVICES"); + if (!str_devices && str_ndev && str_ndev[0] != '\0') { + GGML_LOG_WARN("DEPRECATED: GGML_HEXAGON_NDEV is deprecated. use GGML_HEXAGON_DEVICES instead\n"); + str_devices = str_ndev; + } + + if (str_devices && str_devices[0] != '\0') { + bool is_single_number = true; + for (int i = 0; str_devices[i] != '\0'; i++) { + if (!isdigit((unsigned char)str_devices[i])) { + is_single_number = false; + break; + } + } + if (is_single_number) { + int n = atoi(str_devices); + if (n < 1) n = 1; + if (n > GGML_HEXAGON_MAX_SESSIONS) n = GGML_HEXAGON_MAX_SESSIONS; + opt_ndev = n; + for (size_t i = 0; i < opt_ndev; i++) { + opt_device_configs[i].physical_idx = 0; + opt_device_configs[i].virtual_idx = (int)i; + opt_device_configs[i].name = "HTP" + std::to_string(i); + } + } else { + std::string s_devices(str_devices); + std::stringstream ss(s_devices); + std::string item; + opt_ndev = 0; + while (std::getline(ss, item, ',')) { + size_t start = item.find_first_not_of(" \t\r\n"); + size_t end = item.find_last_not_of(" \t\r\n"); + if (start == std::string::npos) { + continue; + } + item = item.substr(start, end - start + 1); + + if (item.rfind("HTP", 0) == 0) { + std::string rest = item.substr(3); + size_t colon_pos = rest.find(':'); + int phys = 0; + int virt = 0; + try { + if (colon_pos == std::string::npos) { + phys = std::stoi(rest); + virt = 0; + } else { + phys = std::stoi(rest.substr(0, colon_pos)); + virt = std::stoi(rest.substr(colon_pos + 1)); + } + } catch (...) { + GGML_LOG_WARN("ggml-hex: failed to parse device index in '%s'\n", item.c_str()); + continue; + } - if (opt_ndev > GGML_HEXAGON_MAX_SESSIONS) { - opt_ndev = GGML_HEXAGON_MAX_SESSIONS; + if (opt_ndev < GGML_HEXAGON_MAX_SESSIONS) { + opt_device_configs[opt_ndev].physical_idx = phys; + opt_device_configs[opt_ndev].virtual_idx = virt; + opt_device_configs[opt_ndev].name = colon_pos == std::string::npos + ? "HTP" + std::to_string(phys) + : "HTP" + std::to_string(phys) + ":" + std::to_string(virt); + opt_ndev++; + } else { + GGML_LOG_WARN("ggml-hex: max sessions limit reached (%d), ignoring device %s\n", GGML_HEXAGON_MAX_SESSIONS, item.c_str()); + } + } else { + GGML_LOG_WARN("ggml-hex: invalid device name format '%s', must start with HTP\n", item.c_str()); + } + } + } + } else { + opt_ndev = 1; + opt_device_configs[0].physical_idx = 0; + opt_device_configs[0].virtual_idx = 0; + opt_device_configs[0].name = "HTP0"; } #if defined(__ANDROID__) diff --git a/ggml/src/ggml-hexagon/htp-opnode.h b/ggml/src/ggml-hexagon/htp-opnode.h index b0c859dacf9a..741b5e04eb89 100644 --- a/ggml/src/ggml-hexagon/htp-opnode.h +++ b/ggml/src/ggml-hexagon/htp-opnode.h @@ -8,60 +8,107 @@ #include #include #include +#include #include #include "htp-ops.h" #include "htp/matmul-ops.h" #include "htp/flash-attn-ops.h" #include "htp/unary-ops.h" +#include "htp/allreduce-ops.h" struct htp_opnode { - ggml_tensor * node = nullptr; - - std::vector fused; - - htp_op_code opcode = HTP_OP_INVALID; + ggml_tensor * node { nullptr }; + htp_op_code opcode { HTP_OP_INVALID }; + int32_t kernel_params[HTP_OP_MAX_KERN_PARAMS] {0}; + + std::vector fused; + std::vector> dummy; + + std::vector inputs; + std::vector outputs; + std::string name; + + int n_active_src(const ggml_tensor * t) const { + if (!t) return 0; + for (int i = GGML_MAX_SRC - 1; i >= 0; i--) { + if (t->src[i]) { + return i + 1; + } + } + return 0; + } - std::vector extra_dsts; + void init(ggml_tensor * node) { + this->node = node; + if (this->node) { + this->name = ggml_op_desc(this->node); - int32_t kernel_params[HTP_OP_MAX_KERN_PARAMS] = {0}; + // Build inputs (preserving optional nullptrs) + int n_inputs = n_active_src(this->node); + this->inputs.resize(n_inputs, nullptr); + for (int i = 0; i < n_inputs; i++) { + this->inputs[i] = this->node->src[i]; + } - htp_opnode(ggml_tensor * node = nullptr, std::vector fused = {}, htp_op_code opcode = HTP_OP_INVALID, std::vector extra_dsts = {}) - : node(node), fused(std::move(fused)), opcode(opcode), extra_dsts(std::move(extra_dsts)) {} + // Build outputs + this->outputs.push_back(this->dst()); + } + } - ggml_op op() const { - return node->op; + htp_opnode(htp_op_code opcode = HTP_OP_INVALID, ggml_tensor * node = nullptr) : opcode(opcode) { + init(node); } - const ggml_tensor * dst() const { - return fused.empty() ? node : fused.back(); + ggml_op op() const { return node->op; } + const ggml_tensor * src0() const { return node->src[0]; } + const ggml_tensor * src1() const { return node->src[1]; } + const ggml_tensor * dst() const { return outputs.empty() ? node : outputs.back(); } + + ggml_tensor * add_dummy(const ggml_tensor & t) { + dummy.push_back(std::make_shared(t)); + return dummy.back().get(); } void add_fused(ggml_tensor * t, bool extra_dst = false) { fused.push_back(t); + + name += "+"; + name += ggml_op_desc(t); + if (extra_dst) { - extra_dsts.push_back(t); + outputs.push_back(t); + } else { + outputs.clear(); + outputs.push_back(t); } - } - std::vector get_outputs() const { - std::vector res; - if (extra_dsts.empty()) { - res.push_back(dst()); - } else { - res.push_back(node); - for (const auto * x : extra_dsts) { - res.push_back(x); + // Remove the newly fused intermediate output tensor t from inputs (if it was there) + inputs.erase(std::remove(inputs.begin(), inputs.end(), t), inputs.end()); + + // Append new inputs from t, preserving middle nullptrs + int n_inputs = n_active_src(t); + for (int i = 0; i < n_inputs; i++) { + const auto * src = t->src[i]; + if (!src) { + inputs.push_back(nullptr); + } else if (src != node && + std::find(fused.begin(), fused.end(), src) == fused.end() && + std::find(inputs.begin(), inputs.end(), src) == inputs.end()) { + inputs.push_back(src); } } - return res; } - const ggml_tensor * src0() const { - return node->src[0]; + const std::vector & get_inputs() const { + return inputs; } - const ggml_tensor * src1() const { - return node->src[1]; + const std::vector & get_outputs() const { + return outputs; + } + + std::string op_name() const { + return name; } bool is_empty() const { @@ -81,75 +128,6 @@ struct htp_opnode { bool same_input(const htp_opnode& n) const { return n.src1() == this->src1(); } - - std::vector get_inputs() const { - if (fused.empty()) { - int last_non_null = -1; - for (int i = 0; i < GGML_MAX_SRC; i++) { - if (node->src[i]) { - last_non_null = i; - } - } - std::vector inputs(last_non_null + 1, nullptr); - for (int i = 0; i <= last_non_null; i++) { - inputs[i] = node->src[i]; - } - return inputs; - } - - std::vector inputs(GGML_MAX_SRC, nullptr); - std::vector outputs; - outputs.push_back(node); - for (const auto * f : fused) { - outputs.push_back(f); - } - - auto contains = [&](const std::vector & vec, const ggml_tensor * t) { - for (const auto * x : vec) { - if (x == t) return true; - } - return false; - }; - - int count = 0; - auto add_input = [&](const ggml_tensor * t) { - if (t && !contains(outputs, t) && !contains(inputs, t)) { - if (count < (int)inputs.size()) { - inputs[count++] = t; - } else { - inputs.push_back(t); - } - } - }; - - for (int i = 0; i < GGML_MAX_SRC; i++) { - if (node->src[i]) { - add_input(node->src[i]); - } - } - for (const auto * f : fused) { - for (int i = 0; i < GGML_MAX_SRC; i++) { - if (f->src[i]) { - add_input(f->src[i]); - } - } - } - - inputs.resize(count); - return inputs; - } - - std::string op_name() const { - if (fused.empty()) { - return ggml_op_desc(node); - } - std::string name = ggml_op_desc(node); - for (const auto * f : fused) { - name += "+"; - name += ggml_op_desc(f); - } - return name; - } }; struct htp_opformat { @@ -337,8 +315,7 @@ struct htp_opformat { } void format_kernel_params(char * str, size_t max_size, const htp_opnode & node) { if (node.opcode == HTP_OP_MUL_MAT || node.opcode == HTP_OP_MUL_MAT_ID || - node.opcode == HTP_OP_MUL_MAT_QKV || node.opcode == HTP_OP_MUL_MAT_FFN || - node.opcode == HTP_OP_MUL_MAT_ADD) { + node.opcode == HTP_OP_MUL_MAT_NX || node.opcode == HTP_OP_MUL_MAT_ADD) { const auto * kparams = (const struct htp_mm_kernel_params *) node.kernel_params; const char * path = "unknown"; int32_t type = kparams->kernel_type; diff --git a/ggml/src/ggml-hexagon/htp/CMakeLists.txt b/ggml/src/ggml-hexagon/htp/CMakeLists.txt index b00aa2bc94c3..77f3ee39dd3c 100644 --- a/ggml/src/ggml-hexagon/htp/CMakeLists.txt +++ b/ggml/src/ggml-hexagon/htp/CMakeLists.txt @@ -43,6 +43,7 @@ add_library(${HTP_LIB} SHARED pad-ops.c argsort-ops.c im2col-ops.c + allreduce-ops.c ) target_compile_definitions(${HTP_LIB} PRIVATE diff --git a/ggml/src/ggml-hexagon/htp/act-ops.c b/ggml/src/ggml-hexagon/htp/act-ops.c index 9973c088dda7..0a8bf84e382a 100644 --- a/ggml/src/ggml-hexagon/htp/act-ops.c +++ b/ggml/src/ggml-hexagon/htp/act-ops.c @@ -183,6 +183,53 @@ static void swiglu_oai_f32(const float * restrict src0, static const float GELU_COEF_A = 0.044715f; static const float SQRT_2_OVER_PI = 0.79788456080286535587989211986876f; +static inline HVX_Vector hvx_vec_fast_sigmoid_f32_2it(HVX_Vector v) { + v = Q6_Vqf32_vmpy_VsfVsf(v, Q6_V_vsplat_R(FAST_SIGMOID_LOG2F)); + v = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(v), Q6_V_vsplat_R(FAST_SIGMOID_C3)); + + HVX_Vector in_int = hvx_vec_truncate_f32(Q6_Vsf_equals_Vqf32(v)); + HVX_Vector x = Q6_Vqf32_vsub_Vqf32Vsf(v, Q6_Vsf_equals_Vw(in_int)); + HVX_Vector xx = Q6_Vqf32_vmpy_Vqf32Vqf32(x, x); + + HVX_Vector v1 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(xx), Q6_V_vsplat_R(FAST_SIGMOID_C2)); + v1 = Q6_Vqf32_vadd_Vqf32Vsf(v1, Q6_V_vsplat_R(FAST_SIGMOID_LOG2F)); + + HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(x), Q6_V_vsplat_R(FAST_SIGMOID_C1)); + v2 = Q6_Vqf32_vmpy_Vqf32Vqf32(v2, xx); + v2 = Q6_Vqf32_vadd_Vqf32Vqf32(v2, x); + + HVX_Vector v3 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vqf32(v2, v1)); + v3 = Q6_Vw_vaslacc_VwVwR(v3, in_int, 24); + + HVX_Vector v4 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_Vqf32Vqf32(v2, v1)); + HVX_Vector v5 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v3, v4)); + + // Newton-Raphson with 2 iterations + HVX_Vector two_sf = hvx_vec_splat_f32(2.0f); + HVX_Vector i_sf = Q6_Vw_vsub_VwVw(Q6_V_vsplat_R(0x7EEEEBB3), v5); + HVX_Vector r_qf = Q6_Vqf32_vmpy_VsfVsf( + i_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(two_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(i_sf, v5))))); + r_qf = Q6_Vqf32_vmpy_Vqf32Vqf32( + r_qf, Q6_Vqf32_vsub_VsfVsf(two_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(r_qf), v5)))); + HVX_Vector res = Q6_Vsf_equals_Vqf32(r_qf); + + res = Q6_Vqf32_vmpy_VsfVsf(v3, res); + + return Q6_Vsf_equals_Vqf32(res); +} + +static inline HVX_Vector hvx_vec_fast_sigmoid_f32_guard_2it(HVX_Vector v, + HVX_Vector one, + HVX_Vector max_exp, + HVX_Vector min_exp) { + const HVX_VectorPred pred_max = Q6_Q_vcmp_gt_VsfVsf(max_exp, v); + const HVX_VectorPred pred_min = Q6_Q_vcmp_gt_VsfVsf(v, min_exp); + + HVX_Vector out = hvx_vec_fast_sigmoid_f32_2it(v); + out = Q6_V_vmux_QVV(pred_max, out, one); + return Q6_V_vmux_QVV(pred_min, out, Q6_V_vzero()); +} + static inline void hvx_geglu_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { assert((unsigned long) dst % 128 == 0); assert((unsigned long) src0 % 128 == 0); @@ -200,20 +247,13 @@ static inline void hvx_geglu_f32_aa(uint8_t * restrict dst, const uint8_t * rest const HVX_Vector v_coef_a_times_sqrt = hvx_vec_splat_f32(GELU_COEF_A_TIMES_SQRT); const HVX_Vector v_sqrt_2_pi = hvx_vec_splat_f32(SQRT_2_OVER_PI); - const HVX_Vector v_half = hvx_vec_splat_f32(0.5f); const HVX_Vector v_one = hvx_vec_splat_f32(1.0f); - const HVX_Vector v_two = hvx_vec_splat_f32(2.0f); - - // Hoisted fast sigmoid / inverse constants to avoid loop-internal overhead - const HVX_Vector v_log2f = Q6_V_vsplat_R(FAST_SIGMOID_LOG2F); - const HVX_Vector v_c1 = Q6_V_vsplat_R(FAST_SIGMOID_C1); - const HVX_Vector v_c2 = Q6_V_vsplat_R(FAST_SIGMOID_C2); - const HVX_Vector v_inv_aprox = Q6_V_vsplat_R(0x7EEEEBB3); const HVX_Vector v_max_exp = hvx_vec_splat_f32(87.0f); const HVX_Vector v_min_exp = hvx_vec_splat_f32(-87.0f); uint32_t i = 0; + _Pragma("unroll(4)") for (; i < nvec; i++) { HVX_Vector x = vsrc0[i]; HVX_Vector g = vsrc1[i]; @@ -223,56 +263,13 @@ static inline void hvx_geglu_f32_aa(uint8_t * restrict dst, const uint8_t * rest coef = hvx_vec_add_f32_f32(coef, v_sqrt_2_pi); HVX_Vector inner = hvx_vec_mul_f32_f32(x, coef); - // y2 = 2 * inner - HVX_Vector y2 = hvx_vec_mul_f32_f32(inner, v_two); - - // Sigmoid guard check predicates - HVX_VectorPred pred_max = Q6_Q_vcmp_gt_VsfVsf(v_max_exp, y2); - HVX_VectorPred pred_min = Q6_Q_vcmp_gt_VsfVsf(y2, v_min_exp); - - // Fast sigmoid approximation - HVX_Vector v = Q6_Vqf32_vmpy_VsfVsf(y2, v_log2f); - v = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(v), v_half); - - HVX_Vector in_int = hvx_vec_truncate_f32(Q6_Vsf_equals_Vqf32(v)); - HVX_Vector x_sig = Q6_Vqf32_vsub_Vqf32Vsf(v, Q6_Vsf_equals_Vw(in_int)); - HVX_Vector xx_sig = Q6_Vqf32_vmpy_Vqf32Vqf32(x_sig, x_sig); - - HVX_Vector v1 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(xx_sig), v_c2); - v1 = Q6_Vqf32_vadd_Vqf32Vsf(v1, v_log2f); + // y2 = 2 * inner = inner + inner + HVX_Vector y2 = hvx_vec_add_f32_f32(inner, inner); - HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(x_sig), v_c1); - v2 = Q6_Vqf32_vmpy_Vqf32Vqf32(v2, xx_sig); - v2 = Q6_Vqf32_vadd_Vqf32Vqf32(v2, x_sig); - - HVX_Vector v3 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vqf32(v2, v1)); - v3 = Q6_Vw_vaslacc_VwVwR(v3, in_int, 24); - - HVX_Vector v4 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_Vqf32Vqf32(v2, v1)); - HVX_Vector v5 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v3, v4)); - - // Fast division (Newton-Raphson with 2 iterations) - HVX_Vector i_sf = Q6_Vw_vsub_VwVw(v_inv_aprox, v5); - HVX_Vector r_qf = Q6_Vqf32_vmpy_VsfVsf( - i_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v_two, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(i_sf, v5))))); - r_qf = Q6_Vqf32_vmpy_Vqf32Vqf32( - r_qf, Q6_Vqf32_vsub_VsfVsf(v_two, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(r_qf), v5)))); - HVX_Vector res_inv = Q6_Vsf_equals_Vqf32(r_qf); - - HVX_Vector sig2y = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(v3, res_inv)); - - // Sigmoid guards - sig2y = Q6_V_vmux_QVV(pred_max, sig2y, v_one); - sig2y = Q6_V_vmux_QVV(pred_min, sig2y, Q6_V_vzero()); - - // tanh(inner) = 2 * sigmoid(2 * inner) - 1 - HVX_Vector tanh_val = hvx_vec_mul_f32_f32(sig2y, v_two); - tanh_val = hvx_vec_sub_f32_f32(tanh_val, v_one); - - HVX_Vector tanh_plus_one = hvx_vec_add_f32_f32(tanh_val, v_one); - HVX_Vector half_x = hvx_vec_mul_f32_f32(x, v_half); - HVX_Vector gelu_x = hvx_vec_mul_f32_f32(half_x, tanh_plus_one); + // Fast sigmoid approximation (2 iterations) + HVX_Vector sig2y = hvx_vec_fast_sigmoid_f32_guard_2it(y2, v_one, v_max_exp, v_min_exp); + HVX_Vector gelu_x = hvx_vec_mul_f32_f32(x, sig2y); vdst[i] = hvx_vec_mul_f32_f32(gelu_x, g); } @@ -285,50 +282,11 @@ static inline void hvx_geglu_f32_aa(uint8_t * restrict dst, const uint8_t * rest coef = hvx_vec_add_f32_f32(coef, v_sqrt_2_pi); HVX_Vector inner = hvx_vec_mul_f32_f32(x, coef); - HVX_Vector y2 = hvx_vec_mul_f32_f32(inner, v_two); - - HVX_VectorPred pred_max = Q6_Q_vcmp_gt_VsfVsf(v_max_exp, y2); - HVX_VectorPred pred_min = Q6_Q_vcmp_gt_VsfVsf(y2, v_min_exp); - - HVX_Vector v = Q6_Vqf32_vmpy_VsfVsf(y2, v_log2f); - v = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(v), v_half); - - HVX_Vector in_int = hvx_vec_truncate_f32(Q6_Vsf_equals_Vqf32(v)); - HVX_Vector x_sig = Q6_Vqf32_vsub_Vqf32Vsf(v, Q6_Vsf_equals_Vw(in_int)); - HVX_Vector xx_sig = Q6_Vqf32_vmpy_Vqf32Vqf32(x_sig, x_sig); - - HVX_Vector v1 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(xx_sig), v_c2); - v1 = Q6_Vqf32_vadd_Vqf32Vsf(v1, v_log2f); - - HVX_Vector v2 = Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(x_sig), v_c1); - v2 = Q6_Vqf32_vmpy_Vqf32Vqf32(v2, xx_sig); - v2 = Q6_Vqf32_vadd_Vqf32Vqf32(v2, x_sig); - - HVX_Vector v3 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vadd_Vqf32Vqf32(v2, v1)); - v3 = Q6_Vw_vaslacc_VwVwR(v3, in_int, 24); - - HVX_Vector v4 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_Vqf32Vqf32(v2, v1)); - HVX_Vector v5 = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v3, v4)); - - HVX_Vector i_sf = Q6_Vw_vsub_VwVw(v_inv_aprox, v5); - HVX_Vector r_qf = Q6_Vqf32_vmpy_VsfVsf( - i_sf, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vsub_VsfVsf(v_two, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(i_sf, v5))))); - r_qf = Q6_Vqf32_vmpy_Vqf32Vqf32( - r_qf, Q6_Vqf32_vsub_VsfVsf(v_two, Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(Q6_Vsf_equals_Vqf32(r_qf), v5)))); - HVX_Vector res_inv = Q6_Vsf_equals_Vqf32(r_qf); - - HVX_Vector sig2y = Q6_Vsf_equals_Vqf32(Q6_Vqf32_vmpy_VsfVsf(v3, res_inv)); - - sig2y = Q6_V_vmux_QVV(pred_max, sig2y, v_one); - sig2y = Q6_V_vmux_QVV(pred_min, sig2y, Q6_V_vzero()); - - HVX_Vector tanh_val = hvx_vec_mul_f32_f32(sig2y, v_two); - tanh_val = hvx_vec_sub_f32_f32(tanh_val, v_one); + HVX_Vector y2 = hvx_vec_add_f32_f32(inner, inner); - HVX_Vector tanh_plus_one = hvx_vec_add_f32_f32(tanh_val, v_one); - HVX_Vector half_x = hvx_vec_mul_f32_f32(x, v_half); - HVX_Vector gelu_x = hvx_vec_mul_f32_f32(half_x, tanh_plus_one); + HVX_Vector sig2y = hvx_vec_fast_sigmoid_f32_guard_2it(y2, v_one, v_max_exp, v_min_exp); + HVX_Vector gelu_x = hvx_vec_mul_f32_f32(x, sig2y); HVX_Vector res = hvx_vec_mul_f32_f32(gelu_x, g); hvx_vec_store_a((void *) &vdst[i], nloe * sizeof(float), res); } diff --git a/ggml/src/ggml-hexagon/htp/allreduce-ops.c b/ggml/src/ggml-hexagon/htp/allreduce-ops.c new file mode 100644 index 000000000000..d35f685a6dc0 --- /dev/null +++ b/ggml/src/ggml-hexagon/htp/allreduce-ops.c @@ -0,0 +1,398 @@ +#pragma clang diagnostic ignored "-Wunused-variable" +#pragma clang diagnostic ignored "-Wunused-function" +#pragma clang diagnostic ignored "-Wunused-but-set-variable" + +#include +#include +#include +#include +#include + +#define GGML_COMMON_DECL_C +#include "ggml-common.h" +#include "htp-ctx.h" +#include "htp-ops.h" +#include "hvx-utils.h" +#include "htp-tensor.h" +#include "hex-dma.h" +#include "hex-profile.h" +#include "allreduce-ops.h" + +struct htp_allreduce_context { + struct htp_ops_context * octx; + uint32_t n_ranks; + uint32_t n_dsts; + uint32_t nelem; + uint32_t ne0; + uint32_t ne1; + uint32_t row_size_aligned; + uint32_t rank_elem_start; + uint32_t rank_nelem; + uint32_t elems_per_thread; + uint32_t block_elems; + uint32_t vtcm_size_per_thread; + bool is_row_bcast; + uint8_t * src_spad_base[HTP_ALLREDUCE_MAX_RANKS]; + uint8_t * dst_spad_base; + uint8_t * res_spad_base; +}; + +#define DEFINE_ALLREDUCE_THREAD_DMA_1D(SUFFIX, TYPE, HVX_ADD_FN, HAS_ADD) \ +static void allreduce_thread_dma_1d_##SUFFIX(unsigned int nth, unsigned int ith, void * data) { \ + struct htp_allreduce_context * actx = (struct htp_allreduce_context *) data; \ + struct htp_ops_context * octx = actx->octx; \ + \ + const uint32_t n_ranks = actx->n_ranks; \ + const uint32_t n_dsts = actx->n_dsts; \ + const uint32_t block_elems = actx->block_elems; \ + \ + const uint32_t dr = actx->elems_per_thread; \ + const uint32_t ir0 = actx->rank_elem_start + dr * ith; \ + const uint32_t ir1 = MIN(ir0 + dr, actx->rank_elem_start + actx->rank_nelem); \ + if (ir0 >= ir1) return; \ + \ + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ + dma_queue * q = octx->ctx->dma[ith]; \ + \ + uint8_t * src_spad_base[HTP_ALLREDUCE_MAX_RANKS]; \ + for (uint32_t s = 0; s < n_ranks; s++) { \ + src_spad_base[s] = actx->src_spad_base[s] + (ith * actx->vtcm_size_per_thread); \ + } \ + uint8_t * dst_spad_base = actx->dst_spad_base + (ith * actx->vtcm_size_per_thread); \ + uint8_t * res_spad_base = HAS_ADD ? (actx->res_spad_base + (ith * actx->vtcm_size_per_thread)) : NULL; \ + \ + const size_t spad_half = actx->vtcm_size_per_thread / 2; \ + uint32_t ir_prefetch = ir0; \ + int spad_idx = 0; \ + \ + for (int k = 0; k < 2 && ir_prefetch < ir1; k++) { \ + uint32_t cur_elems = MIN(block_elems, ir1 - ir_prefetch); \ + size_t cur_bytes = cur_elems * sizeof(TYPE); \ + uint8_t * d_spad = dst_spad_base + spad_idx * spad_half; \ + for (uint32_t d = 0; d < n_dsts; d++) { \ + uint8_t * d_ddr = (uint8_t *) octx->dsts[d]->data + ir_prefetch * sizeof(TYPE); \ + dma_queue_push(q, dma_make_ptr(d_ddr, d_spad), cur_bytes, cur_bytes, cur_bytes, 0); \ + } \ + for (uint32_t s = 0; s < n_ranks; s++) { \ + uint8_t * s_spad = src_spad_base[s] + spad_idx * spad_half; \ + const uint8_t * s_ddr = (const uint8_t *) octx->src[s]->data + ir_prefetch * sizeof(TYPE); \ + dma_queue_push(q, dma_make_ptr(s_spad, s_ddr), cur_bytes, cur_bytes, cur_bytes, 1); \ + } \ + if (HAS_ADD) { \ + uint8_t * r_spad = res_spad_base + spad_idx * spad_half; \ + const uint8_t * r_ddr = (const uint8_t *) octx->src[2 * n_ranks]->data + ir_prefetch * sizeof(TYPE); \ + dma_queue_push(q, dma_make_ptr(r_spad, r_ddr), cur_bytes, cur_bytes, cur_bytes, 1); \ + } \ + ir_prefetch += cur_elems; \ + spad_idx ^= 1; \ + } \ + \ + for (uint32_t ir = ir0; ir < ir1; ) { \ + uint32_t cur_elems = MIN(block_elems, ir1 - ir); \ + size_t cur_bytes = cur_elems * sizeof(TYPE); \ + uint8_t * d_spad = NULL; \ + for (uint32_t d = 0; d < n_dsts; d++) { \ + d_spad = (uint8_t *) dma_queue_pop(q).src; \ + } \ + uint8_t * s_spad[HTP_ALLREDUCE_MAX_RANKS]; \ + for (uint32_t s = 0; s < n_ranks; s++) { \ + s_spad[s] = (uint8_t *) dma_queue_pop(q).dst; \ + } \ + uint8_t * r_spad = HAS_ADD ? (uint8_t *) dma_queue_pop(q).dst : NULL; \ + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir); \ + HVX_ADD_FN(d_spad, s_spad[0], s_spad[1], cur_elems); \ + for (uint32_t s = 2; s < n_ranks; s++) { \ + HVX_ADD_FN(d_spad, d_spad, s_spad[s], cur_elems); \ + } \ + if (HAS_ADD) { \ + HVX_ADD_FN(d_spad, d_spad, r_spad, cur_elems); \ + } \ + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) ir); \ + for (uint32_t d = 0; d < n_dsts; d++) { \ + uint8_t * d_ddr = (uint8_t *) octx->dsts[d]->data + ir * sizeof(TYPE); \ + dma_queue_push(q, dma_make_ptr(d_ddr, d_spad), cur_bytes, cur_bytes, cur_bytes, 1); \ + } \ + if (ir_prefetch < ir1) { \ + uint32_t next_elems = MIN(block_elems, ir1 - ir_prefetch); \ + size_t next_bytes = next_elems * sizeof(TYPE); \ + for (uint32_t s = 0; s < n_ranks; s++) { \ + const uint8_t * s_next = (const uint8_t *) octx->src[s]->data + ir_prefetch * sizeof(TYPE); \ + dma_queue_push(q, dma_make_ptr(s_spad[s], s_next), next_bytes, next_bytes, next_bytes, 1); \ + } \ + if (HAS_ADD) { \ + const uint8_t * r_next = (const uint8_t *) octx->src[2 * n_ranks]->data + ir_prefetch * sizeof(TYPE); \ + dma_queue_push(q, dma_make_ptr(r_spad, r_next), next_bytes, next_bytes, next_bytes, 1); \ + } \ + ir_prefetch += next_elems; \ + } \ + ir += cur_elems; \ + } \ + dma_queue_flush(q); \ +} + +DEFINE_ALLREDUCE_THREAD_DMA_1D(f16, __fp16, hvx_add_f16_aaa, 0) +DEFINE_ALLREDUCE_THREAD_DMA_1D(f32, float, hvx_add_f32_aaa, 0) +DEFINE_ALLREDUCE_THREAD_DMA_1D(add_f16, __fp16, hvx_add_f16_aaa, 1) +DEFINE_ALLREDUCE_THREAD_DMA_1D(add_f32, float, hvx_add_f32_aaa, 1) + +#define DEFINE_ALLREDUCE_THREAD_DMA_2D(SUFFIX, TYPE, HVX_ADD_FN, HAS_ADD, IS_ROW_BCAST) \ +static void allreduce_thread_dma_2d_##SUFFIX(unsigned int nth, unsigned int ith, void * data) { \ + struct htp_allreduce_context * actx = (struct htp_allreduce_context *) data; \ + struct htp_ops_context * octx = actx->octx; \ + \ + const uint32_t n_ranks = actx->n_ranks; \ + const uint32_t n_dsts = actx->n_dsts; \ + const uint32_t ne0 = actx->ne0; \ + const uint32_t block_rows = actx->block_elems; \ + const uint32_t row_size_aligned = actx->row_size_aligned; \ + const uint32_t row_bytes = ne0 * sizeof(TYPE); \ + \ + const uint32_t dr = actx->elems_per_thread; \ + const uint32_t r0 = actx->rank_elem_start + dr * ith; \ + const uint32_t r1 = MIN(r0 + dr, actx->rank_elem_start + actx->rank_nelem); \ + if (r0 >= r1) return; \ + \ + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ + dma_queue * q = octx->ctx->dma[ith]; \ + \ + uint8_t * src_spad_base[HTP_ALLREDUCE_MAX_RANKS]; \ + for (uint32_t s = 0; s < n_ranks; s++) { \ + src_spad_base[s] = actx->src_spad_base[s] + (ith * actx->vtcm_size_per_thread); \ + } \ + uint8_t * dst_spad_base = actx->dst_spad_base + (ith * actx->vtcm_size_per_thread); \ + uint8_t * res_spad_base = HAS_ADD ? (IS_ROW_BCAST ? actx->res_spad_base : (actx->res_spad_base + (ith * actx->vtcm_size_per_thread))) : NULL; \ + \ + const size_t spad_half = actx->vtcm_size_per_thread / 2; \ + uint32_t r_prefetch = r0; \ + int spad_idx = 0; \ + \ + for (int k = 0; k < 2 && r_prefetch < r1; k++) { \ + uint32_t cur_rows = MIN(block_rows, r1 - r_prefetch); \ + uint8_t * d_spad = dst_spad_base + spad_idx * spad_half; \ + for (uint32_t d = 0; d < n_dsts; d++) { \ + uint8_t * d_ddr = (uint8_t *) octx->dsts[d]->data + r_prefetch * octx->dsts[d]->nb[1]; \ + dma_queue_push(q, dma_make_ptr(d_ddr, d_spad), octx->dsts[d]->nb[1], row_size_aligned, row_bytes, 0); \ + } \ + for (uint32_t s = 0; s < n_ranks; s++) { \ + uint8_t * s_spad = src_spad_base[s] + spad_idx * spad_half; \ + const uint8_t * s_ddr = (const uint8_t *) octx->src[s]->data + r_prefetch * octx->src[s]->nb[1]; \ + dma_queue_push(q, dma_make_ptr(s_spad, s_ddr), row_size_aligned, octx->src[s]->nb[1], row_bytes, cur_rows); \ + } \ + if (HAS_ADD && !IS_ROW_BCAST) { \ + uint8_t * r_spad = res_spad_base + spad_idx * spad_half; \ + const uint8_t * r_ddr = (const uint8_t *) octx->src[2 * n_ranks]->data + r_prefetch * octx->src[2 * n_ranks]->nb[1]; \ + dma_queue_push(q, dma_make_ptr(r_spad, r_ddr), row_size_aligned, octx->src[2 * n_ranks]->nb[1], row_bytes, cur_rows); \ + } \ + r_prefetch += cur_rows; \ + spad_idx ^= 1; \ + } \ + \ + for (uint32_t r = r0; r < r1; ) { \ + uint32_t cur_rows = MIN(block_rows, r1 - r); \ + uint8_t * d_spad = NULL; \ + for (uint32_t d = 0; d < n_dsts; d++) { \ + d_spad = (uint8_t *) dma_queue_pop(q).src; \ + } \ + uint8_t * s_spad[HTP_ALLREDUCE_MAX_RANKS]; \ + for (uint32_t s = 0; s < n_ranks; s++) { \ + s_spad[s] = (uint8_t *) dma_queue_pop(q).dst; \ + } \ + uint8_t * r_spad = (HAS_ADD && !IS_ROW_BCAST) ? (uint8_t *) dma_queue_pop(q).dst : NULL; \ + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) r); \ + for (uint32_t row = 0; row < cur_rows; row++) { \ + uint8_t * d_row = d_spad + row * row_size_aligned; \ + const uint8_t * s0_row = s_spad[0] + row * row_size_aligned; \ + const uint8_t * s1_row = s_spad[1] + row * row_size_aligned; \ + HVX_ADD_FN(d_row, s0_row, s1_row, ne0); \ + for (uint32_t s = 2; s < n_ranks; s++) { \ + const uint8_t * ss_row = s_spad[s] + row * row_size_aligned; \ + HVX_ADD_FN(d_row, d_row, ss_row, ne0); \ + } \ + if (HAS_ADD) { \ + const uint8_t * res_row = IS_ROW_BCAST ? res_spad_base : (r_spad + row * row_size_aligned); \ + HVX_ADD_FN(d_row, d_row, res_row, ne0); \ + } \ + } \ + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, (uint16_t) r); \ + for (uint32_t d = 0; d < n_dsts; d++) { \ + uint8_t * d_ddr = (uint8_t *) octx->dsts[d]->data + r * octx->dsts[d]->nb[1]; \ + dma_queue_push(q, dma_make_ptr(d_ddr, d_spad), octx->dsts[d]->nb[1], row_size_aligned, row_bytes, cur_rows); \ + } \ + if (r_prefetch < r1) { \ + uint32_t next_rows = MIN(block_rows, r1 - r_prefetch); \ + for (uint32_t s = 0; s < n_ranks; s++) { \ + const uint8_t * s_next = (const uint8_t *) octx->src[s]->data + r_prefetch * octx->src[s]->nb[1]; \ + dma_queue_push(q, dma_make_ptr(s_spad[s], s_next), row_size_aligned, octx->src[s]->nb[1], row_bytes, next_rows); \ + } \ + if (HAS_ADD && !IS_ROW_BCAST) { \ + const uint8_t * r_next = (const uint8_t *) octx->src[2 * n_ranks]->data + r_prefetch * octx->src[2 * n_ranks]->nb[1]; \ + dma_queue_push(q, dma_make_ptr(r_spad, r_next), row_size_aligned, octx->src[2 * n_ranks]->nb[1], row_bytes, next_rows); \ + } \ + r_prefetch += next_rows; \ + } \ + r += cur_rows; \ + } \ + dma_queue_flush(q); \ +} + +DEFINE_ALLREDUCE_THREAD_DMA_2D(f16, __fp16, hvx_add_f16_aaa, 0, 0) +DEFINE_ALLREDUCE_THREAD_DMA_2D(f32, float, hvx_add_f32_aaa, 0, 0) +DEFINE_ALLREDUCE_THREAD_DMA_2D(add_f16, __fp16, hvx_add_f16_aaa, 1, 0) +DEFINE_ALLREDUCE_THREAD_DMA_2D(add_f32, float, hvx_add_f32_aaa, 1, 0) +DEFINE_ALLREDUCE_THREAD_DMA_2D(add_bcast_f16, __fp16, hvx_add_f16_aaa, 1, 1) +DEFINE_ALLREDUCE_THREAD_DMA_2D(add_bcast_f32, float, hvx_add_f32_aaa, 1, 1) + +int op_allreduce(struct htp_ops_context * octx) { + const struct htp_allreduce_kernel_params * kparams = (const struct htp_allreduce_kernel_params *) octx->kernel_params; + const struct htp_tensor * dst = octx->dst; + + const uint32_t rank = (uint32_t) kparams->rank; + const uint32_t n_ranks = (uint32_t) kparams->n_ranks; + + if (n_ranks < 2 || n_ranks > HTP_ALLREDUCE_MAX_RANKS || rank >= n_ranks) { + return HTP_STATUS_INVAL_PARAMS; + } + + if (dst->type != HTP_TYPE_F16 && dst->type != HTP_TYPE_F32) { + return HTP_STATUS_NO_SUPPORT; + } + + const uint32_t nelem = dst->ne[0] * dst->ne[1] * dst->ne[2] * dst->ne[3]; + const uint32_t fence_seq_entry = (uint32_t) octx->op_params[0]; + const uint32_t fence_seq_exit = (uint32_t) octx->op_params[1]; + + // 1. Entry Barrier: Synchronize all ranks before reading + struct htp_thread_trace * tr0 = &octx->ctx->trace[0]; + htp_trace_event_start(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_entry); + + const struct htp_tensor * my_sync = octx->src[n_ranks + rank]; + atomic_uint * my_fence = (atomic_uint *) my_sync->data; + + atomic_store(&my_fence[0], fence_seq_entry); + asm volatile ("syncht" : : : "memory"); + Q6_dccleaninva_A((void *) my_fence); + + for (uint32_t j = 0; j < n_ranks; j++) { + if (j == rank) continue; + const struct htp_tensor * peer_sync = octx->src[n_ranks + j]; + atomic_uint * peer_fence = (atomic_uint *) peer_sync->data; + uint64_t spins = 0; + while (1) { + Q6_dccleaninva_A((void *) peer_fence); + uint32_t val = atomic_load(&peer_fence[0]); + if (val == fence_seq_entry || val == fence_seq_exit) { + break; + } + if (++spins > HTP_FENCE_TIMEOUT) { + FARF(ERROR, "ggml-hex: allreduce entry fence-wait TIMEOUT: rank %u waiting on %u (fence %p seq %u)\n", rank, j, peer_fence, fence_seq_entry); + return HTP_STATUS_INTERNAL_ERR; + } + hex_pause(); + } + } + asm volatile ("syncht" : : : "memory"); + + htp_trace_event_stop(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_entry); + + // 2. Multi-threaded Reduction across assigned rank chunk + if (nelem > 0) { + const uint32_t n_threads = (uint32_t) kparams->n_threads; + const uint32_t block_elems = (uint32_t) kparams->block_elems; + const uint32_t elems_per_thread = (uint32_t) kparams->elems_per_thread; + const uint32_t vtcm_size_per_thread = (uint32_t) kparams->vtcm_size_per_thread; + + const bool has_add = (octx->op == HTP_OP_ALLREDUCE_ADD); + + struct htp_allreduce_context actx; + actx.octx = octx; + actx.n_ranks = n_ranks; + actx.n_dsts = (uint32_t) kparams->n_dsts ? (uint32_t) kparams->n_dsts : n_ranks; + actx.nelem = nelem; + actx.ne0 = (uint32_t) kparams->ne0; + actx.ne1 = (uint32_t) kparams->ne1; + actx.row_size_aligned = (uint32_t) kparams->row_size_aligned; + actx.rank_elem_start = (uint32_t) kparams->rank_elem_start; + actx.rank_nelem = (uint32_t) kparams->rank_nelem; + actx.elems_per_thread = elems_per_thread; + actx.block_elems = block_elems; + actx.vtcm_size_per_thread = vtcm_size_per_thread; + actx.is_row_bcast = (kparams->is_row_bcast != 0); + + work_queue_func_t reduce_fun = NULL; + switch (kparams->kernel_type) { + case HTP_ALLREDUCE_KERNEL_DMA_1D: + if (has_add) { + reduce_fun = (dst->type == HTP_TYPE_F16) ? allreduce_thread_dma_1d_add_f16 : allreduce_thread_dma_1d_add_f32; + } else { + reduce_fun = (dst->type == HTP_TYPE_F16) ? allreduce_thread_dma_1d_f16 : allreduce_thread_dma_1d_f32; + } + break; + case HTP_ALLREDUCE_KERNEL_DMA_2D: + if (has_add) { + if (kparams->is_row_bcast) { + reduce_fun = (dst->type == HTP_TYPE_F16) ? allreduce_thread_dma_2d_add_bcast_f16 : allreduce_thread_dma_2d_add_bcast_f32; + } else { + reduce_fun = (dst->type == HTP_TYPE_F16) ? allreduce_thread_dma_2d_add_f16 : allreduce_thread_dma_2d_add_f32; + } + } else { + reduce_fun = (dst->type == HTP_TYPE_F16) ? allreduce_thread_dma_2d_f16 : allreduce_thread_dma_2d_f32; + } + break; + default: + return HTP_STATUS_NO_SUPPORT; + } + + uint8_t * vtcm_ptr = (uint8_t *) octx->ctx->vtcm_base; + for (uint32_t s = 0; s < n_ranks; s++) { + actx.src_spad_base[s] = vtcm_ptr; + vtcm_ptr += n_threads * vtcm_size_per_thread; + } + actx.dst_spad_base = vtcm_ptr; + vtcm_ptr += n_threads * vtcm_size_per_thread; + if (has_add) { + actx.res_spad_base = vtcm_ptr; + vtcm_ptr += (actx.is_row_bcast ? 1 : n_threads) * vtcm_size_per_thread; + } + + if (has_add && actx.is_row_bcast) { + const uint8_t * r_ddr = (const uint8_t *) octx->src[2 * n_ranks]->data; + const uint32_t row_bytes = actx.ne0 * (dst->type == HTP_TYPE_F16 ? sizeof(__fp16) : sizeof(float)); + dma_queue * q = octx->ctx->dma[0]; + dma_queue_push(q, dma_make_ptr(actx.res_spad_base, r_ddr), actx.row_size_aligned, 0, row_bytes, 1); + dma_queue_pop(q); + } + + work_queue_run(octx->ctx->work_queue, reduce_fun, &actx, n_threads); + } + + // 4. Exit Barrier: Synchronize all ranks after writing + htp_trace_event_start(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_exit); + + atomic_store(&my_fence[0], fence_seq_exit); + asm volatile ("syncht" : : : "memory"); + Q6_dccleaninva_A((void *) my_fence); + + for (uint32_t j = 0; j < n_ranks; j++) { + if (j == rank) continue; + const struct htp_tensor * peer_sync = octx->src[n_ranks + j]; + atomic_uint * peer_fence = (atomic_uint *) peer_sync->data; + uint64_t spins = 0; + while (1) { + Q6_dccleaninva_A((void *) peer_fence); + uint32_t val = atomic_load(&peer_fence[0]); + if (val == fence_seq_exit) { + break; + } + if (++spins > HTP_FENCE_TIMEOUT) { + FARF(ERROR, "ggml-hex: allreduce exit fence-wait TIMEOUT: rank %u waiting on %u (fence %p seq %u)\n", rank, j, peer_fence, fence_seq_exit); + return HTP_STATUS_INTERNAL_ERR; + } + hex_pause(); + } + } + asm volatile ("syncht" : : : "memory"); + + htp_trace_event_stop(tr0, HTP_TRACE_EVT_FENCE, (uint16_t) fence_seq_exit); + + return HTP_STATUS_OK; +} diff --git a/ggml/src/ggml-hexagon/htp/allreduce-ops.h b/ggml/src/ggml-hexagon/htp/allreduce-ops.h new file mode 100644 index 000000000000..de447d87e912 --- /dev/null +++ b/ggml/src/ggml-hexagon/htp/allreduce-ops.h @@ -0,0 +1,40 @@ +#ifndef ALLREDUCE_OPS_H +#define ALLREDUCE_OPS_H + +#include + +#define HTP_ALLREDUCE_MAX_RANKS 4 + +#ifdef __cplusplus +extern "C" { +#endif + +enum htp_allreduce_kernel_type { + HTP_ALLREDUCE_KERNEL_UNSUPPORTED = 0, + HTP_ALLREDUCE_KERNEL_DMA_1D, + HTP_ALLREDUCE_KERNEL_DMA_2D, +}; + +struct htp_allreduce_kernel_params { + int32_t rank; + int32_t n_ranks; + int32_t n_threads; + int32_t block_elems; // 1D: block_elems, 2D: block_rows + int32_t elems_per_thread; // 1D: nelem_per_thread, 2D: nrows_per_thread + int32_t vtcm_size_per_thread; + int32_t vtcm_size; + int32_t kernel_type; + int32_t ne0; + int32_t ne1; + int32_t row_size_aligned; + int32_t rank_elem_start; + int32_t rank_nelem; + int32_t n_dsts; + int32_t is_row_bcast; +}; + +#ifdef __cplusplus +} +#endif + +#endif /* ALLREDUCE_OPS_H */ diff --git a/ggml/src/ggml-hexagon/htp/cpy-ops.c b/ggml/src/ggml-hexagon/htp/cpy-ops.c index ae507effa51a..15bc8dc244fb 100644 --- a/ggml/src/ggml-hexagon/htp/cpy-ops.c +++ b/ggml/src/ggml-hexagon/htp/cpy-ops.c @@ -4,6 +4,7 @@ #include #include +#include #include #include @@ -14,6 +15,7 @@ #include "htp-ops.h" #include "htp-ops.h" #include "hvx-utils.h" +#include "htp-tensor.h" struct htp_copy_context { struct htp_ops_context * octx; @@ -78,7 +80,7 @@ static void cpy_thread_##NAME##_sameshape(unsigned int nth, unsigned int ith, vo } \ } -DEFINE_CPY_SAMESHAPE(f32, float, 4) +DEFINE_CPY_SAMESHAPE(f32, float, 4) DEFINE_CPY_SAMESHAPE(f16, __fp16, 2) #define DEFINE_CPY_RESHAPE(NAME, ELEM_TYPE, ELEM_SIZE) \ @@ -179,7 +181,7 @@ static void cpy_thread_##NAME##_reshape(unsigned int nth, unsigned int ith, void } \ } -DEFINE_CPY_RESHAPE(f32, float, 4) +DEFINE_CPY_RESHAPE(f32, float, 4) DEFINE_CPY_RESHAPE(f16, __fp16, 2) static void cpy_thread_f16_f32_sameshape(unsigned int nth, unsigned int ith, void * data) { @@ -232,6 +234,41 @@ static void cpy_thread_f32_f16_sameshape(unsigned int nth, unsigned int ith, voi } } +static inline void cpy_dma_sametype_sameshape( + struct htp_ops_context * octx, + const struct htp_tensor * dst, + const struct htp_tensor * src0, + uint32_t elem_size, + uint32_t ne00, uint32_t ne01, uint32_t ne02, uint32_t ne03, + uint32_t nb01, uint32_t nb02, uint32_t nb03, + uint32_t nb1, uint32_t nb2, uint32_t nb3 +) { + const bool contiguous_outer = + (ne02 == 1 || (nb02 == ne01 * nb01 && nb2 == ne01 * nb1)) && + (ne03 == 1 || (nb03 == ne02 * nb02 && nb3 == ne02 * nb2)); + + dma_queue * q = octx->ctx->dma[0]; + + if (contiguous_outer) { + dma_queue_push(q, dma_make_ptr((void *) dst->data, (const void *) src0->data), nb1, nb01, ne00 * elem_size, ne01 * ne02 * ne03); + dma_queue_pop(q); + return; + } + + for (uint32_t i03 = 0; i03 < ne03; i03++) { + for (uint32_t i02 = 0; i02 < ne02; i02++) { + uint8_t* dst_ptr = (uint8_t*) dst->data + i02*nb2 + i03*nb3; + uint8_t* src0_ptr = (uint8_t*) src0->data + i02*nb02 + i03*nb03; + if (!dma_queue_push(q, dma_make_ptr(dst_ptr, src0_ptr), nb1, nb01, ne00 * elem_size, ne01)) { + dma_queue_flush(q); + dma_queue_push(q, dma_make_ptr(dst_ptr, src0_ptr), nb1, nb01, ne00 * elem_size, ne01); + } + } + } + + dma_queue_flush(q); +} + int op_cpy(struct htp_ops_context * octx) { cpy_preamble; @@ -264,14 +301,11 @@ int op_cpy(struct htp_ops_context * octx) { ct.src0_nrows_per_thread = (nr + n_threads - 1) / n_threads; - worker_callback_t copy_fun; + worker_callback_t copy_fun = NULL; + bool use_dma = false; if (sametype && sameshape) { - if (src0->type == HTP_TYPE_F32) { - copy_fun = cpy_thread_f32_sameshape; - } else { - copy_fun = cpy_thread_f16_sameshape; - } + use_dma = true; } else if (sameshape) { /**/ if (dst->type == HTP_TYPE_F16 && src0->type == HTP_TYPE_F32) copy_fun = cpy_thread_f16_f32_sameshape; @@ -289,7 +323,28 @@ int op_cpy(struct htp_ops_context * octx) { return HTP_STATUS_NO_SUPPORT; } - worker_pool_run_func(octx->ctx->worker_pool, copy_fun, &ct, n_threads); + if (use_dma) { + cpy_dma_sametype_sameshape(octx, dst, src0, ct.src0_type_size, ne00, ne01, ne02, ne03, nb01, nb02, nb03, nb1, nb2, nb3); + } else { + worker_pool_run_func(octx->ctx->worker_pool, copy_fun, &ct, n_threads); + } + + const struct htp_tensor *sync = octx->src[1]; + if (sync) { + if (!use_dma) { + // htp_tensor_flush_all(octx->ctx, octx->dsts, 1); + qurt_mem_cache_clean((qurt_addr_t) 0, 0, QURT_MEM_CACHE_FLUSH_INVALIDATE_ALL, QURT_MEM_DCACHE); + } + + atomic_uint * sync_fence = (atomic_uint *) sync->data; + const uint32_t seq = (uint32_t) octx->op_params[0]; + + atomic_store(&sync_fence[0], seq); + asm volatile ("syncht" : : : "memory"); + Q6_dccleaninva_A((void *) sync_fence); + + FARF(HIGH, "ggml-hex: sync-release : fence %p seq %u\n", sync_fence, seq); + } return HTP_STATUS_OK; } diff --git a/ggml/src/ggml-hexagon/htp/dma-queue.h b/ggml/src/ggml-hexagon/htp/dma-queue.h index 264284bda828..190ca3a9b9e1 100644 --- a/ggml/src/ggml-hexagon/htp/dma-queue.h +++ b/ggml/src/ggml-hexagon/htp/dma-queue.h @@ -244,17 +244,18 @@ static inline dma_ptr dma_queue_pop(dma_queue * q) { return dptr; } - dma_descriptor_2d * desc = &r->desc[r->pop_idx]; + dptr = r->dptr[r->pop_idx]; + + volatile dma_descriptor_2d * desc = &r->desc[r->pop_idx]; // Wait for desc to complete if (!desc->done) { + // FARF(ALWAYS, "dma-poll: idx %u dst %p src %p", r->pop_idx, dptr.dst, dptr.src); while (!desc->done) { dmpoll(); } } - dptr = r->dptr[r->pop_idx]; - htp_trace_event_stop(r->trace, HTP_TRACE_EVT_DMA, r->pop_idx); r->pop_idx = (r->pop_idx + 1) & r->idx_mask; diff --git a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c index 81765629046a..c76b4d3a3ac6 100644 --- a/ggml/src/ggml-hexagon/htp/flash-attn-ops.c +++ b/ggml/src/ggml-hexagon/htp/flash-attn-ops.c @@ -30,6 +30,8 @@ #include "ggml-common.h" #include "htp-ctx.h" #include "htp-ops.h" +#include "htp-tensor.h" +#include "hvx-quant.h" #include "flash-attn-ops.h" #include "hvx-fa-kernels.h" @@ -85,12 +87,17 @@ struct htp_fa_context { uint8_t * spad_m; uint8_t * spad_a; + const struct htp_tensor * k; + const struct htp_tensor * v; + uint64_t t_start; }; struct hmx_fa_context { const struct htp_ops_context * octx; const struct htp_tensor * sinks; // attention sinks (src[4]), NULL if absent + const struct htp_tensor * k; + const struct htp_tensor * v; bool pipeline; // true when n_kv_blocks >= FA_MIN_KV_BLOCKS && n_threads >= 2 uint32_t n_threads; @@ -214,8 +221,8 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * const uint32_t DV = nev0; const size_t size_q_row = DK * ((q->type == HTP_TYPE_F32) ? 4 : 2); - const size_t size_k_row = DK * sizeof(__fp16); - const size_t size_v_row = DV * sizeof(__fp16); + const size_t size_k_row = htp_tensor_get_row_size(k->type, DK); + const size_t size_v_row = htp_tensor_get_row_size(v->type, DV); // Scratchpad buffers for Q, K, V, Mask, and VKQ32 accumulator uint8_t * spad_q = factx->spad_q + factx->size_q_block * ith; @@ -364,6 +371,23 @@ static void flash_attn_ext_f16_thread(unsigned int nth, unsigned int ith, void * uint8_t * v_base = dma_queue_pop(dma).dst; // V __fp16 * m_base = mask ? dma_queue_pop(dma).dst : NULL; // M + if (factx->k->type == HTP_TYPE_Q8_0) { + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, ir); + for (uint32_t r = 0; r < current_block_size; ++r) { + __fp16 * row_k = (__fp16 *)(k_base + r * factx->size_k_row_padded); + hvx_dequantize_row_q8_0_f16(row_k, row_k, DK); + } + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, ir); + } + if (factx->v->type == HTP_TYPE_Q8_0) { + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_V_PREP, ir); + for (uint32_t r = 0; r < current_block_size; ++r) { + __fp16 * row_v = (__fp16 *)(v_base + r * factx->size_v_row_padded); + hvx_dequantize_row_q8_0_f16(row_v, row_v, DV); + } + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_V_PREP, ir); + } + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_QK, ir); // Inner loop processing the block from VTCM @@ -625,6 +649,12 @@ static void fa_k_interleave_thread(unsigned int n, unsigned int i, void * data) struct htp_thread_trace * tr = &factx->octx->ctx->trace[i]; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, (uint16_t) (args->kv_start + start)); + if (factx->k->type == HTP_TYPE_Q8_0) { + for (uint32_t r = start; r < end; ++r) { + __fp16 * row_k = (__fp16 *)((char *)args->curr_k + r * args->src_stride * sizeof(__fp16)); + hvx_dequantize_row_q8_0_f16(row_k, row_k, factx->DK); + } + } hmx_interleave_rows_to_tiles(factx->vtcm_k_tiles[args->buf_idx], (const __fp16 *) args->curr_k, total_rows, factx->DK, args->src_stride, start, end); htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_K_PREP, (uint16_t) (args->kv_start + start)); @@ -673,6 +703,12 @@ static void fa_v_interleave_thread(unsigned int n, unsigned int i, void * data) struct htp_thread_trace * tr = &factx->octx->ctx->trace[i]; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_FA_V_PREP, (uint16_t) (args->kv_start + start)); + if (factx->v->type == HTP_TYPE_Q8_0) { + for (uint32_t r = start; r < end; ++r) { + __fp16 * row_v = (__fp16 *)((char *)args->v_src + r * args->src_stride * sizeof(__fp16)); + hvx_dequantize_row_q8_0_f16(row_v, row_v, factx->DV); + } + } hmx_interleave_cols_to_tiles(v_tiles_dst, (const __fp16 *) args->v_src, total_rows, factx->DV, args->src_stride, (uint32_t) args->n_col_tiles, start, end); htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_FA_V_PREP, (uint16_t) (args->kv_start + start)); @@ -1809,6 +1845,8 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { memset(&factx, 0, sizeof(factx)); factx.octx = octx; factx.sinks = octx->src[4]; // NULL if this op has no attention sinks + factx.k = k; + factx.v = v; factx.n_threads = kparams->n_threads; factx.DK = DK; factx.DV = DV; @@ -1853,10 +1891,10 @@ int hmx_flash_attn_ext(struct htp_ops_context * octx) { // ======== VTCM allocation (GQA-aware) ======== // K/V row sizes drive the DMA descriptors (not the VTCM layout) and are used // throughout the KV loop below. - const size_t size_k_row = DK * sizeof(__fp16); - const size_t size_v_row = DV * sizeof(__fp16); - const size_t size_k_row_padded = hex_round_up(size_k_row, 128); - const size_t size_v_row_padded = hex_round_up(size_v_row, 128); + const size_t size_k_row = htp_tensor_get_row_size(k->type, DK); + const size_t size_v_row = htp_tensor_get_row_size(v->type, DV); + const size_t size_k_row_padded = hex_round_up(DK * sizeof(__fp16), 128); + const size_t size_v_row_padded = hex_round_up(DV * sizeof(__fp16), 128); // Build the VTCM layout once (shared with the host estimator) and place every // scratch buffer at its computed offset. @@ -2348,7 +2386,9 @@ int op_flash_attn_ext(struct htp_ops_context * octx) { const struct htp_tensor * dst = octx->dst; // Check support - if ((q->type != HTP_TYPE_F16 && q->type != HTP_TYPE_F32) || k->type != HTP_TYPE_F16 || v->type != HTP_TYPE_F16) { + if ((q->type != HTP_TYPE_F16 && q->type != HTP_TYPE_F32) || + (k->type != HTP_TYPE_F16 && k->type != HTP_TYPE_Q8_0) || + (v->type != HTP_TYPE_F16 && v->type != HTP_TYPE_Q8_0)) { return HTP_STATUS_NO_SUPPORT; } @@ -2364,6 +2404,8 @@ int op_flash_attn_ext(struct htp_ops_context * octx) { struct htp_fa_context factx; factx.octx = octx; + factx.k = k; + factx.v = v; factx.t_start = HAP_perf_get_qtimer_count(); diff --git a/ggml/src/ggml-hexagon/htp/get-rows-ops.c b/ggml/src/ggml-hexagon/htp/get-rows-ops.c index bf7063e9880a..05769d17f74e 100644 --- a/ggml/src/ggml-hexagon/htp/get-rows-ops.c +++ b/ggml/src/ggml-hexagon/htp/get-rows-ops.c @@ -12,18 +12,17 @@ #include "ggml-common.h" #include "htp-ctx.h" #include "htp-ops.h" -#include "htp-ops.h" +#include "htp-tensor.h" #include "hvx-utils.h" +#include "hvx-quant.h" +#include "get-rows-ops.h" +#include "work-queue.h" struct get_rows_context { struct htp_ops_context * octx; - uint32_t tasks_per_thread; - uint32_t total_tasks; - uint32_t chunks_per_row; - uint32_t chunk_size; - struct fastdiv_values get_rows_div_ne10; - struct fastdiv_values get_rows_div_ne10_ne11; - struct fastdiv_values get_rows_div_chunks_per_row; + const struct htp_get_rows_kernel_params * kparams; + struct htp_get_rows_vtcm_layout vtcm_layout; + uint8_t * vtcm_base; }; #define get_rows_preamble \ @@ -56,102 +55,161 @@ struct get_rows_context { \ const uint32_t nr = ne10 * ne11 * ne12; -static void get_rows_thread_f32_f32_dma(unsigned int nth, unsigned int ith, void *data) { - struct get_rows_context * grctx = (struct get_rows_context *)data; - struct htp_ops_context * octx = grctx->octx; - get_rows_preamble; - - uint64_t qt = HAP_perf_get_qtimer_count(); - - const uint32_t dr = grctx->tasks_per_thread; - const uint32_t ir0 = dr * ith; - if (ir0 >= grctx->total_tasks) { - return; - } - const uint32_t ir1 = MIN(ir0 + dr, grctx->total_tasks); - - const bool is_i32 = (octx->src[1]->type == HTP_TYPE_I32); - - dma_queue * dma_queue = octx->ctx->dma[ith]; - for (uint32_t i = ir0; i < ir1; ++i) { - const uint32_t i12 = fastdiv(i, &grctx->get_rows_div_ne10_ne11); - const uint32_t rem = i - i12 * ne11 * ne10; - const uint32_t i11 = fastdiv(rem, &grctx->get_rows_div_ne10); - const uint32_t i10 = rem - i11 * ne10; - - const uintptr_t src1_addr = octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12; - uint32_t i01 = is_i32 ? *(int32_t *)src1_addr : *(int64_t *)src1_addr; - - if (i01 >= ne01) { - continue; - } - - const uintptr_t src0_ptr = octx->src[0]->data + i01*nb01 + i11*nb02 + i12*nb03; - const uintptr_t dst_ptr = octx->dst->data + i10*nb1 + i11*nb2 + i12*nb3; - - while (!dma_queue_push(dma_queue, dma_make_ptr((void *)dst_ptr, (const void *)src0_ptr), nb1, nb01, ne00 * sizeof(float), 1)) { - dma_queue_pop(dma_queue); - } - } - dma_queue_flush(dma_queue); - - qt = HAP_perf_qtimer_count_to_us(HAP_perf_get_qtimer_count() - qt); - FARF(HIGH, "get-rows-f32-f32-dma %d/%d: %ux%ux%ux%u (%u:%u) x %ux%ux%ux%u -> %ux%ux%ux%u usec %u\n", ith, nth, - ne00, ne01, ne02, ne03, ir0, ir1, ne10, ne11, ne12, ne13, ne0, ne1, ne2, ne3, (unsigned) qt); +#define GET_ROWS_THREAD_ST_FN(IDX_TYPE) \ +static void get_rows_thread_st_##IDX_TYPE(unsigned int nth, unsigned int ith, void *data) { \ + struct get_rows_context * grctx = (struct get_rows_context *)data; \ + struct htp_ops_context * octx = grctx->octx; \ + const struct htp_get_rows_kernel_params * kparams = grctx->kparams; \ + get_rows_preamble; \ + const uint32_t dr = kparams->tasks_per_thread; \ + const uint32_t ir0 = dr * ith; \ + if (ir0 >= kparams->total_tasks) { \ + return; \ + } \ + const uint32_t ir1 = MIN(ir0 + dr, kparams->total_tasks); \ + const uint32_t row_size_bytes = htp_tensor_get_row_size(octx->src[0]->type, ne00); \ + dma_queue * dma_queue = octx->ctx->dma[ith]; \ + for (uint32_t i = ir0; i < ir1; ++i) { \ + const uint32_t i12 = fastdiv(i, &kparams->div_ne10_ne11); \ + const uint32_t rem = i - i12 * ne11 * ne10; \ + const uint32_t i11 = fastdiv(rem, &kparams->div_ne10); \ + const uint32_t i10 = rem - i11 * ne10; \ + const IDX_TYPE * src1_ptr = (const IDX_TYPE *)(octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12); \ + const uint32_t i01 = (uint32_t)*src1_ptr; \ + assert(i01 < ne01); \ + const uint32_t q02 = fastdiv(i11, &kparams->div_ne02); \ + const uint32_t i02 = i11 - q02 * ne02; \ + const uint32_t q03 = fastdiv(i12, &kparams->div_ne03); \ + const uint32_t i03 = i12 - q03 * ne03; \ + const uintptr_t src0_ptr = octx->src[0]->data + i01*nb01 + i02*nb02 + i03*nb03; \ + const uintptr_t dst_ptr = octx->dst->data + i10*nb1 + i11*nb2 + i12*nb3; \ + while (!dma_queue_push(dma_queue, dma_make_ptr((void *)dst_ptr, (const void *)src0_ptr), nb1, nb01, \ + row_size_bytes, 1)) { \ + dma_queue_pop(dma_queue); \ + } \ + } \ + dma_queue_flush(dma_queue); \ } -static void get_rows_thread_f32_f32_hvx(unsigned int nth, unsigned int ith, void *data) { - struct get_rows_context * grctx = (struct get_rows_context *)data; - struct htp_ops_context * octx = grctx->octx; - get_rows_preamble; - - uint64_t qt = HAP_perf_get_qtimer_count(); - - const uint32_t dr = grctx->tasks_per_thread; - const uint32_t ir0 = dr * ith; - if (ir0 >= grctx->total_tasks) { - return; - } - const uint32_t ir1 = MIN(ir0 + dr, grctx->total_tasks); - - const bool is_i32 = (octx->src[1]->type == HTP_TYPE_I32); +GET_ROWS_THREAD_ST_FN(int32_t) +GET_ROWS_THREAD_ST_FN(int64_t) + +#define GET_ROWS_THREAD_DT_FN(TYPE_NAME, SRC0_SIZE_EXPR, IDX_TYPE, COMPUTE_EXPR) \ +static void get_rows_thread_##TYPE_NAME##_##IDX_TYPE(unsigned int nth, unsigned int ith, void *data) { \ + struct get_rows_context * grctx = (struct get_rows_context *)data; \ + struct htp_ops_context * octx = grctx->octx; \ + const struct htp_get_rows_kernel_params * kparams = grctx->kparams; \ + get_rows_preamble; \ + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ + const uint32_t dr = kparams->tasks_per_thread; \ + const uint32_t ir0 = dr * ith; \ + if (ir0 >= kparams->total_tasks) { \ + return; \ + } \ + const uint32_t ir1 = MIN(ir0 + dr, kparams->total_tasks); \ + const uint32_t chunks_per_row = kparams->chunks_per_row; \ + const uint32_t chunk_size = kparams->chunk_size; \ + dma_queue * dma_queue = octx->ctx->dma[ith]; \ + const struct htp_get_rows_vtcm_layout * vtcm_layout = &grctx->vtcm_layout; \ + uint8_t * vtcm_src0 = grctx->vtcm_base + vtcm_layout->off_src0 + ith * vtcm_layout->src0_bytes_per_thread; \ + uint8_t * vtcm_dst = grctx->vtcm_base + vtcm_layout->off_dst + ith * vtcm_layout->dst_bytes_per_thread; \ + for (uint32_t step = 0, spad_idx = 0; step < ir1 - ir0 && spad_idx < 2; ++step, spad_idx++) { \ + const uint32_t i = ir0 + step; \ + const uint32_t row_idx = fastdiv(i, &kparams->div_chunks_per_row); \ + const uint32_t chunk_idx = i - row_idx * chunks_per_row; \ + const uint32_t i12 = fastdiv(row_idx, &kparams->div_ne10_ne11); \ + const uint32_t rem = row_idx - i12 * ne11 * ne10; \ + const uint32_t i11 = fastdiv(rem, &kparams->div_ne10); \ + const uint32_t i10 = rem - i11 * ne10; \ + const IDX_TYPE * src1_ptr = (const IDX_TYPE *)(octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12); \ + const uint32_t i01 = (uint32_t)*src1_ptr; \ + assert(i01 < ne01); \ + const uint32_t q02 = fastdiv(i11, &kparams->div_ne02); \ + const uint32_t i02 = i11 - q02 * ne02; \ + const uint32_t q03 = fastdiv(i12, &kparams->div_ne03); \ + const uint32_t i03 = i12 - q03 * ne03; \ + const uint32_t offset = chunk_idx * chunk_size; \ + const uint32_t cur_elems = (offset < ne00) ? MIN(chunk_size, ne00 - offset) : 0; \ + const uint32_t cur_src0_bytes = SRC0_SIZE_EXPR(cur_elems); \ + const uint32_t cur_dst_bytes = cur_elems * sizeof(float); \ + const uintptr_t src0_ptr = octx->src[0]->data + i01*nb01 + i02*nb02 + i03*nb03 + SRC0_SIZE_EXPR(offset); \ + dma_queue_push(dma_queue, \ + dma_make_ptr((void *)(uintptr_t)octx->dst->data, \ + vtcm_dst + spad_idx * vtcm_layout->dst_spad_half_size), \ + cur_dst_bytes, vtcm_layout->dst_spad_half_size, cur_dst_bytes, 0); \ + dma_queue_push(dma_queue, \ + dma_make_ptr((void *)(vtcm_src0 + spad_idx * vtcm_layout->src0_spad_half_size), \ + (const void *)src0_ptr), \ + vtcm_layout->src0_spad_half_size, cur_src0_bytes, cur_src0_bytes, 1); \ + } \ + for (uint32_t step = 0; step < ir1 - ir0; ++step) { \ + const uint32_t i = ir0 + step; \ + void * dst_spad = (void *) dma_queue_pop(dma_queue).src; \ + void * src_spad = (void *) dma_queue_pop(dma_queue).dst; \ + const uint32_t row_idx = fastdiv(i, &kparams->div_chunks_per_row); \ + const uint32_t chunk_idx = i - row_idx * chunks_per_row; \ + const uint32_t i12 = fastdiv(row_idx, &kparams->div_ne10_ne11); \ + const uint32_t rem = row_idx - i12 * ne11 * ne10; \ + const uint32_t i11 = fastdiv(rem, &kparams->div_ne10); \ + const uint32_t i10 = rem - i11 * ne10; \ + const uint32_t offset = chunk_idx * chunk_size; \ + const uint32_t cur_elems = (offset < ne00) ? MIN(chunk_size, ne00 - offset) : 0; \ + const uint32_t cur_dst_bytes = cur_elems * sizeof(float); \ + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, i); \ + COMPUTE_EXPR; \ + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, i); \ + const uintptr_t dst_ptr = octx->dst->data + i10*nb1 + i11*nb2 + i12*nb3 + offset * sizeof(float); \ + dma_queue_push(dma_queue, \ + dma_make_ptr((void *)dst_ptr, (const void *)dst_spad), \ + cur_dst_bytes, vtcm_layout->dst_spad_half_size, cur_dst_bytes, 1); \ + const uint32_t next_step = step + 2; \ + if (next_step < ir1 - ir0) { \ + const uint32_t pi = ir0 + next_step; \ + const uint32_t prow_idx = fastdiv(pi, &kparams->div_chunks_per_row); \ + const uint32_t pchunk_idx = pi - prow_idx * chunks_per_row; \ + const uint32_t pi12 = fastdiv(prow_idx, &kparams->div_ne10_ne11); \ + const uint32_t prem = prow_idx - pi12 * ne11 * ne10; \ + const uint32_t pi11 = fastdiv(prem, &kparams->div_ne10); \ + const uint32_t pi10 = prem - pi11 * ne10; \ + const IDX_TYPE * psrc1_ptr = (const IDX_TYPE *)(octx->src[1]->data + pi10*nb10 + pi11*nb11 + pi12*nb12); \ + const uint32_t pi01 = (uint32_t)*psrc1_ptr; \ + assert(pi01 < ne01); \ + const uint32_t pq02 = fastdiv(pi11, &kparams->div_ne02); \ + const uint32_t pi02 = pi11 - pq02 * ne02; \ + const uint32_t pq03 = fastdiv(pi12, &kparams->div_ne03); \ + const uint32_t pi03 = pi12 - pq03 * ne03; \ + const uint32_t poffset = pchunk_idx * chunk_size; \ + const uint32_t pcur_elems = (poffset < ne00) ? MIN(chunk_size, ne00 - poffset) : 0; \ + const uint32_t pcur_src0_bytes = SRC0_SIZE_EXPR(pcur_elems); \ + const uintptr_t psrc0_ptr = \ + octx->src[0]->data + pi01*nb01 + pi02*nb02 + pi03*nb03 + SRC0_SIZE_EXPR(poffset); \ + dma_queue_push(dma_queue, \ + dma_make_ptr((void *)src_spad, (const void *)psrc0_ptr), \ + vtcm_layout->src0_spad_half_size, pcur_src0_bytes, pcur_src0_bytes, 1); \ + } \ + } \ + dma_queue_flush(dma_queue); \ +} - const uint32_t chunks_per_row = grctx->chunks_per_row; - const uint32_t chunk_size = grctx->chunk_size; - for (uint32_t i = ir0; i < ir1; ++i) { - const uint32_t row_idx = fastdiv(i, &grctx->get_rows_div_chunks_per_row); - const uint32_t chunk_idx = i - row_idx * chunks_per_row; +#define F32_BYTES(n) ((n) * sizeof(float)) +#define F16_BYTES(n) ((n) * sizeof(__fp16)) +#define Q8_0_BYTES(n) (((n) / 32) * sizeof(block_q8_0)) - const uint32_t i12 = fastdiv(row_idx, &grctx->get_rows_div_ne10_ne11); - const uint32_t rem = row_idx - i12 * ne11 * ne10; - const uint32_t i11 = fastdiv(rem, &grctx->get_rows_div_ne10); - const uint32_t i10 = rem - i11 * ne10; +GET_ROWS_THREAD_DT_FN(f32, F32_BYTES, int32_t, { if (cur_elems > 0) hvx_copy_f32_uu((uint8_t *)dst_spad, (const uint8_t *)src_spad, cur_elems); }) +GET_ROWS_THREAD_DT_FN(f32, F32_BYTES, int64_t, { if (cur_elems > 0) hvx_copy_f32_uu((uint8_t *)dst_spad, (const uint8_t *)src_spad, cur_elems); }) - const uintptr_t src1_addr = octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12; - uint32_t i01 = is_i32 ? *(int32_t *)src1_addr : *(int64_t *)src1_addr; +GET_ROWS_THREAD_DT_FN(f16, F16_BYTES, int32_t, { hvx_dequantize_row_f16_f32((float *)dst_spad, src_spad, ne00); }) +GET_ROWS_THREAD_DT_FN(f16, F16_BYTES, int64_t, { hvx_dequantize_row_f16_f32((float *)dst_spad, src_spad, ne00); }) - if (i01 >= ne01) { - continue; - } - - const uint32_t offset = chunk_idx * chunk_size; - if (offset < ne00) { - const uint32_t copy_size = MIN(chunk_size, ne00 - offset); - const uintptr_t src0_ptr = octx->src[0]->data + i01*nb01 + i11*nb02 + i12*nb03 + offset * sizeof(float); - const uintptr_t dst_ptr = octx->dst->data + i10*nb1 + i11*nb2 + i12*nb3 + offset * sizeof(float); - hvx_copy_f32_uu((uint8_t *)dst_ptr, (const uint8_t *)src0_ptr, copy_size); - } - } - - qt = HAP_perf_qtimer_count_to_us(HAP_perf_get_qtimer_count() - qt); - FARF(HIGH, "get-rows-f32-f32-hvx %d/%d: %ux%ux%ux%u (%u:%u) x %ux%ux%ux%u -> %ux%ux%ux%u usec %u\n", ith, nth, - ne00, ne01, ne02, ne03, ir0, ir1, ne10, ne11, ne12, ne13, ne0, ne1, ne2, ne3, (unsigned) qt); -} +GET_ROWS_THREAD_DT_FN(q8_0, Q8_0_BYTES, int32_t, { hvx_dequantize_row_q8_0_f32((float *)dst_spad, src_spad, ne00); }) +GET_ROWS_THREAD_DT_FN(q8_0, Q8_0_BYTES, int64_t, { hvx_dequantize_row_q8_0_f32((float *)dst_spad, src_spad, ne00); }) int op_get_rows(struct htp_ops_context * octx) { - get_rows_preamble; + const struct htp_get_rows_kernel_params * kparams = (const struct htp_get_rows_kernel_params *) octx->kernel_params; - if (octx->src[0]->type != HTP_TYPE_F32) { + if (octx->src[0]->type != HTP_TYPE_F32 && + octx->src[0]->type != HTP_TYPE_F16 && + octx->src[0]->type != HTP_TYPE_Q8_0) { return HTP_STATUS_NO_SUPPORT; } @@ -167,52 +225,28 @@ int op_get_rows(struct htp_ops_context * octx) { return HTP_STATUS_OK; } - const uint32_t nb00 = octx->src[0]->nb[0]; - const uint32_t nb0 = octx->dst->nb[0]; - - const bool can_use_dma = (nb00 == sizeof(float)) && (nb0 == sizeof(float)); - const bool use_dma = can_use_dma && (ne00 >= 2048); - struct get_rows_context grctx; grctx.octx = octx; - grctx.get_rows_div_ne10 = init_fastdiv_values(octx->src[1]->ne[0]); - grctx.get_rows_div_ne10_ne11 = init_fastdiv_values(octx->src[1]->ne[0] * octx->src[1]->ne[1]); + grctx.kparams = kparams; + grctx.vtcm_base = (uint8_t *)octx->ctx->vtcm_base; - if (use_dma) { - grctx.chunks_per_row = 1; - grctx.chunk_size = ne00; - grctx.total_tasks = nr; - grctx.get_rows_div_chunks_per_row = init_fastdiv_values(1); + const uint32_t ne00 = octx->src[0]->ne[0]; + htp_get_rows_vtcm_layout_build(&grctx.vtcm_layout, octx->src[0]->type, ne00, kparams->n_threads); - const uint32_t n_threads = MIN(nr, octx->n_threads); - grctx.tasks_per_thread = (nr + n_threads - 1) / n_threads; + const bool is_i32 = (octx->src[1]->type == HTP_TYPE_I32); - worker_pool_run_func(octx->ctx->worker_pool, get_rows_thread_f32_f32_dma, &grctx, n_threads); + work_queue_func_t q_func = NULL; + if (kparams->use_dma) { + q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_st_int32_t : get_rows_thread_st_int64_t); } else { - uint32_t chunks_per_row = 1; - uint32_t chunk_size = ne00; - uint32_t total_tasks = nr; - - if (nr < octx->n_threads) { - const uint32_t min_chunk_size = 1024; - uint32_t max_chunks = ne00 / min_chunk_size; - if (max_chunks == 0) { - max_chunks = 1; - } - chunks_per_row = MIN((octx->n_threads + nr - 1) / nr, max_chunks); - chunk_size = (ne00 + chunks_per_row - 1) / chunks_per_row; - total_tasks = nr * chunks_per_row; + switch (octx->src[0]->type) { + case HTP_TYPE_F32: q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_f32_int32_t : get_rows_thread_f32_int64_t); break; + case HTP_TYPE_F16: q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_f16_int32_t : get_rows_thread_f16_int64_t); break; + case HTP_TYPE_Q8_0: q_func = (work_queue_func_t)(is_i32 ? get_rows_thread_q8_0_int32_t : get_rows_thread_q8_0_int64_t); break; + default: return HTP_STATUS_NO_SUPPORT; } - - grctx.chunks_per_row = chunks_per_row; - grctx.chunk_size = chunk_size; - grctx.total_tasks = total_tasks; - grctx.get_rows_div_chunks_per_row = init_fastdiv_values(chunks_per_row); - - const uint32_t n_threads = MIN(total_tasks, octx->n_threads); - grctx.tasks_per_thread = (total_tasks + n_threads - 1) / n_threads; - - worker_pool_run_func(octx->ctx->worker_pool, get_rows_thread_f32_f32_hvx, &grctx, n_threads); } + + work_queue_run(octx->ctx->work_queue, q_func, &grctx, kparams->n_threads); return HTP_STATUS_OK; } diff --git a/ggml/src/ggml-hexagon/htp/get-rows-ops.h b/ggml/src/ggml-hexagon/htp/get-rows-ops.h new file mode 100644 index 000000000000..0e7c2ca8cf0b --- /dev/null +++ b/ggml/src/ggml-hexagon/htp/get-rows-ops.h @@ -0,0 +1,77 @@ +#ifndef HTP_GET_ROWS_OPS_H +#define HTP_GET_ROWS_OPS_H + +#include "hex-fastdiv.h" + +struct htp_get_rows_kernel_params { + int32_t n_threads; + int32_t use_dma; + int32_t chunks_per_row; + int32_t chunk_size; + int32_t total_tasks; + int32_t tasks_per_thread; + int32_t vtcm_size; + + // Fastdiv helpers + struct fastdiv_values div_ne10; + struct fastdiv_values div_ne10_ne11; + struct fastdiv_values div_chunks_per_row; + struct fastdiv_values div_ne02; + struct fastdiv_values div_ne03; +}; + +struct htp_get_rows_vtcm_layout { + size_t total_bytes; + size_t off_src0; + size_t off_dst; + + size_t src0_bytes_per_thread; + size_t dst_bytes_per_thread; + + size_t src0_spad_half_size; + size_t dst_spad_half_size; +}; + +static inline void htp_get_rows_vtcm_layout_build( + struct htp_get_rows_vtcm_layout * vtcm_layout, + int type, + uint32_t ne00, + uint32_t n_threads) { + + uint32_t src0_row_size = 0; + switch (type) { + case 0: // HTP_TYPE_F32 + src0_row_size = ne00 * 4; + break; + case 1: // HTP_TYPE_F16 + src0_row_size = ne00 * 2; + break; + case 8: // HTP_TYPE_Q8_0 + src0_row_size = (ne00 / 32) * 34; + break; + default: + src0_row_size = 0; + break; + } + + size_t src0_row_size_aligned = (src0_row_size + 255) & ~255; + size_t dst_row_size_aligned = (ne00 * sizeof(float) + 255) & ~255; + + vtcm_layout->src0_spad_half_size = src0_row_size_aligned; + vtcm_layout->dst_spad_half_size = dst_row_size_aligned; + + vtcm_layout->src0_bytes_per_thread = src0_row_size_aligned * 2; + vtcm_layout->dst_bytes_per_thread = dst_row_size_aligned * 2; + + vtcm_layout->off_src0 = 0; + vtcm_layout->off_dst = vtcm_layout->off_src0 + vtcm_layout->src0_bytes_per_thread * n_threads; + vtcm_layout->total_bytes = vtcm_layout->off_dst + vtcm_layout->dst_bytes_per_thread * n_threads; +} + +#if defined(__cplusplus) +static_assert(sizeof(struct htp_get_rows_kernel_params) <= 128, "htp_get_rows_kernel_params is too large for kernel_params blob"); +#else +_Static_assert(sizeof(struct htp_get_rows_kernel_params) <= 128, "htp_get_rows_kernel_params is too large for kernel_params blob"); +#endif + +#endif // HTP_GET_ROWS_OPS_H diff --git a/ggml/src/ggml-hexagon/htp/hex-utils.h b/ggml/src/ggml-hexagon/htp/hex-utils.h index 93e87efcb4c4..1b3965030009 100644 --- a/ggml/src/ggml-hexagon/htp/hex-utils.h +++ b/ggml/src/ggml-hexagon/htp/hex-utils.h @@ -39,17 +39,22 @@ static inline void hex_l2fetch_block(const void * addr, size_t size) { #define HEX_L2_LINE_SIZE 128 #define HEX_L2_BLOCK_SIZE (HEX_L2_LINE_SIZE * 4) // flush granularity (lines per loop iteration) +#define HEX_L2_FLUSH_IL_THRESHOLD 1024 // inline flush threshold #define HEX_L2_FLUSH_WQ_THRESHOLD (4 * 1024) #define HEX_L2_FLUSH_ALL_THRESHOLD (4 * 1024 * 1024) static inline void hex_l2flush(void * addr, size_t size) { const uint32_t s = ((uint32_t) addr) & ~(HEX_L2_LINE_SIZE - 1); const uint32_t e = (((uint32_t) addr) + size + HEX_L2_LINE_SIZE - 1) & ~(HEX_L2_LINE_SIZE - 1); - for (uint32_t i = s; i < e; i += HEX_L2_BLOCK_SIZE) { - Q6_dccleaninva_A((void *) i + HEX_L2_LINE_SIZE * 0); - Q6_dccleaninva_A((void *) i + HEX_L2_LINE_SIZE * 1); - Q6_dccleaninva_A((void *) i + HEX_L2_LINE_SIZE * 2); - Q6_dccleaninva_A((void *) i + HEX_L2_LINE_SIZE * 3); + const uint32_t eb = s + ((e - s) & ~(HEX_L2_BLOCK_SIZE - 1)); + for (uint32_t i = s; i < eb; i += HEX_L2_BLOCK_SIZE) { + Q6_dccleaninva_A((void *) (i + HEX_L2_LINE_SIZE * 0)); + Q6_dccleaninva_A((void *) (i + HEX_L2_LINE_SIZE * 1)); + Q6_dccleaninva_A((void *) (i + HEX_L2_LINE_SIZE * 2)); + Q6_dccleaninva_A((void *) (i + HEX_L2_LINE_SIZE * 3)); + } + for (uint32_t i = eb; i < e; i += HEX_L2_LINE_SIZE) { + Q6_dccleaninva_A((void *) i); } } diff --git a/ggml/src/ggml-hexagon/htp/htp-ctx.h b/ggml/src/ggml-hexagon/htp/htp-ctx.h index e0f9a0c40d19..88ecf144b944 100644 --- a/ggml/src/ggml-hexagon/htp/htp-ctx.h +++ b/ggml/src/ggml-hexagon/htp/htp-ctx.h @@ -117,8 +117,7 @@ struct htp_context { int op_matmul(struct htp_ops_context * octx); int op_matmul_id(struct htp_ops_context * octx); -int op_matmul_qkv(struct htp_ops_context * octx); -int op_matmul_ffn(struct htp_ops_context * octx); +int op_matmul_nx(struct htp_ops_context * octx); int op_binary(struct htp_ops_context * octx); int op_unary(struct htp_ops_context * octx); int op_sum_rows(struct htp_ops_context * octx); @@ -141,5 +140,6 @@ int op_solve_tri(struct htp_ops_context * octx); int op_gated_delta_net(struct htp_ops_context * octx); int op_pad(struct htp_ops_context * octx); int op_im2col(struct htp_ops_context * octx); +int op_allreduce(struct htp_ops_context * octx); #endif /* HTP_CTX_H */ diff --git a/ggml/src/ggml-hexagon/htp/htp-ops.h b/ggml/src/ggml-hexagon/htp/htp-ops.h index a138f062aa68..e804844d5996 100644 --- a/ggml/src/ggml-hexagon/htp/htp-ops.h +++ b/ggml/src/ggml-hexagon/htp/htp-ops.h @@ -43,13 +43,6 @@ enum htp_data_type { -// Mask to enable various stages of the Ops. -// Used for debugging and profiling. -enum htp_op_stage { - HTP_OPSTAGE_QUEUE = (1 << 0), // Enable Queueing (ie calls into NPU) - HTP_OPSTAGE_COMPUTE = (1 << 1), // Enable Compute -}; - // Do not reorder first 4 (used as an index) enum htp_op_code { HTP_OP_MUL = 0, @@ -58,8 +51,7 @@ enum htp_op_code { HTP_OP_DIV = 3, HTP_OP_MUL_MAT, HTP_OP_MUL_MAT_ID, - HTP_OP_MUL_MAT_QKV, - HTP_OP_MUL_MAT_FFN, + HTP_OP_MUL_MAT_NX, HTP_OP_MUL_MAT_ADD, HTP_OP_RMS_NORM, HTP_OP_RMS_NORM_MUL, @@ -70,6 +62,8 @@ enum htp_op_code { HTP_OP_UNARY_NEG, HTP_OP_UNARY_SOFTPLUS, HTP_OP_UNARY_TANH, + HTP_OP_UNARY_ABS, + HTP_OP_UNARY_LOG, HTP_OP_GLU_SWIGLU, HTP_OP_GLU_SWIGLU_OAI, HTP_OP_GLU_GEGLU, @@ -99,12 +93,15 @@ enum htp_op_code { HTP_OP_CONCAT, HTP_OP_CLAMP, HTP_OP_IM2COL, + HTP_OP_FENCE, + HTP_OP_ALLREDUCE, + HTP_OP_ALLREDUCE_ADD, HTP_OP_INVALID }; #define HTP_OP_MAX_DIMS 4 // aka GGML_MAX_DIMS -#define HTP_OP_MAX_INPUTS 6 // aka GGML_MAX_SRCS +#define HTP_OP_MAX_INPUTS 10 // aka GGML_MAX_SRCS #define HTP_OP_MAX_OUTPUTS 4 #define HTP_OP_MAX_PARAMS 16 // aka GGML_MAX_OP_PARAMS #define HTP_OP_MAX_KERN_PARAMS 32 @@ -112,13 +109,16 @@ enum htp_op_code { #define HTP_OP_MAX_BUFS 16 #define HTP_OP_MAX_TENSORS 8192 // must stay under 64K (uint16) +#define HTP_FENCE_TIMEOUT (1000000000ULL) + #define HTP_OP_MAX_VMEM_DEFAULT (3355443200u) #define HTP_MMAP_MAX_VMEM (2147483648u) enum htp_tensor_flags { - HTP_TENSOR_COMPUTE = (1U << 0), // Tensor buffer temporal compute data (not weights) - HTP_TENSOR_DIRTY = (1U << 1) // Tensor buffer is dirty and needs to be flushed + HTP_TENSOR_WEIGHT = (1U << 0), // Tensor buffer model weight data (not compute) + HTP_TENSOR_REPACK = (1U << 1), // Tensor is in repacked tiled format + HTP_TENSOR_FENCE = (1U << 2) // Tensor is synchronization fence (explicitly managed) }; // Tensor descriptor @@ -175,6 +175,7 @@ enum htp_trace_event_id { HTP_TRACE_EVT_L2FLUSH = 1, HTP_TRACE_EVT_INIT = 2, HTP_TRACE_EVT_BUFF = 3, + HTP_TRACE_EVT_FENCE = 4, HTP_TRACE_EVT_HVX_COMP = 20, HTP_TRACE_EVT_HVX_A_QUANT = 21, @@ -215,6 +216,7 @@ struct htp_opbatch_req { uint32_t n_ops; // Number of ops uint32_t n_traces; // Number of trace descriptors per thread uint32_t pad; // unused + uint64_t seq; // Sequence number // struct htp_buf_desc bufs[]; -- dspqueue buf 0 // struct htp_tensor tensors[]; -- dspqueue buf 0 // struct htp_op_desc ops[]; -- dspqueue buf 0 @@ -231,6 +233,7 @@ struct htp_opbatch_rsp { uint32_t pad; // align to 8 bytes uint64_t cycles_start; // Start cycle counter uint64_t cycles_stop; // Stop cycle counter + uint64_t seq; // Sequence number // struct htp_prof_desc profs[]; -- dspqueue buf 0 }; diff --git a/ggml/src/ggml-hexagon/htp/htp-tensor.c b/ggml/src/ggml-hexagon/htp/htp-tensor.c index 39436e26dfff..ae377c9221ff 100644 --- a/ggml/src/ggml-hexagon/htp/htp-tensor.c +++ b/ggml/src/ggml-hexagon/htp/htp-tensor.c @@ -79,7 +79,14 @@ void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * co for (uint32_t i = 0; i < n; i++) { const struct htp_tensor * t = tensors[i]; - if (!t) continue; + if (!t || (t->flags & (HTP_TENSOR_WEIGHT | HTP_TENSOR_FENCE))) { + continue; + } + + if (t->size <= HEX_L2_FLUSH_IL_THRESHOLD) { + hex_l2flush((void *) (uintptr_t) t->data, t->size); + continue; + } uint32_t t_start = t->data; uint32_t t_end = t_start + t->size; @@ -242,7 +249,7 @@ void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * co for (uint32_t i = 0; i < n; i++) { const struct htp_tensor * t = tensors[i]; - if (t && (t->flags & HTP_TENSOR_COMPUTE) && is_tensor_dirty(ctx, t)) { + if (t && !(t->flags & (HTP_TENSOR_WEIGHT | HTP_TENSOR_FENCE)) && is_tensor_dirty(ctx, t)) { dirty_tensors[n_dirty++] = t; total_dirty += t->size; } diff --git a/ggml/src/ggml-hexagon/htp/htp-tensor.h b/ggml/src/ggml-hexagon/htp/htp-tensor.h index 2c3fc54c748f..c9cadbae3f23 100644 --- a/ggml/src/ggml-hexagon/htp/htp-tensor.h +++ b/ggml/src/ggml-hexagon/htp/htp-tensor.h @@ -13,6 +13,15 @@ static inline uint32_t * htp_tensor_flags(const struct htp_tensor * t) { return (uint32_t *) &t->flags; } +static inline uint32_t htp_tensor_get_row_size(int type, uint32_t ne00) { + switch (type) { + case HTP_TYPE_F32: return ne00 * 4; + case HTP_TYPE_F16: return ne00 * 2; + case HTP_TYPE_Q8_0: return (ne00 / 32) * 34; + default: return 0; + } +} + struct htp_context; void htp_tensor_flush_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n); void htp_tensor_dirty_all(struct htp_context * ctx, const struct htp_tensor * const * tensors, uint32_t n); diff --git a/ggml/src/ggml-hexagon/htp/hvx-arith.h b/ggml/src/ggml-hexagon/htp/hvx-arith.h index 82e3416970b4..765c3577668e 100644 --- a/ggml/src/ggml-hexagon/htp/hvx-arith.h +++ b/ggml/src/ggml-hexagon/htp/hvx-arith.h @@ -17,9 +17,9 @@ #define hvx_arith_loop_body(dst_type, src0_type, src1_type, elem_size, vec_store, vec_op) \ do { \ - dst_type * restrict vdst = (dst_type *) dst; \ - src0_type * restrict vsrc0 = (src0_type *) src0; \ - src1_type * restrict vsrc1 = (src1_type *) src1; \ + dst_type * vdst = (dst_type *) dst; \ + src0_type * vsrc0 = (src0_type *) src0; \ + src1_type * vsrc1 = (src1_type *) src1; \ \ const uint32_t epv = 128 / (elem_size); \ const uint32_t nvec = n / epv; \ @@ -57,40 +57,40 @@ // Generic macro to define alignment permutations for an op #define DEFINE_HVX_BINARY_OP_VARIANTS(OP_NAME, OP_MACRO, ELEM_TYPE) \ -static inline void OP_NAME##_aaa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ +static inline void OP_NAME##_aaa(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) { \ assert((uintptr_t) dst % 128 == 0); \ assert((uintptr_t) src0 % 128 == 0); \ assert((uintptr_t) src1 % 128 == 0); \ hvx_arith_loop_body(HVX_Vector, HVX_Vector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO); \ } \ -static inline void OP_NAME##_aau(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ +static inline void OP_NAME##_aau(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) { \ assert((uintptr_t) dst % 128 == 0); \ assert((uintptr_t) src0 % 128 == 0); \ hvx_arith_loop_body(HVX_Vector, HVX_Vector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO); \ } \ -static inline void OP_NAME##_aua(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ +static inline void OP_NAME##_aua(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) { \ assert((uintptr_t) dst % 128 == 0); \ assert((uintptr_t) src1 % 128 == 0); \ hvx_arith_loop_body(HVX_Vector, HVX_UVector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO); \ } \ -static inline void OP_NAME##_auu(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ +static inline void OP_NAME##_auu(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) { \ assert((uintptr_t) dst % 128 == 0); \ hvx_arith_loop_body(HVX_Vector, HVX_UVector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_a, OP_MACRO); \ } \ -static inline void OP_NAME##_uaa(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ +static inline void OP_NAME##_uaa(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) { \ assert((uintptr_t) src0 % 128 == 0); \ assert((uintptr_t) src1 % 128 == 0); \ hvx_arith_loop_body(HVX_UVector, HVX_Vector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO); \ } \ -static inline void OP_NAME##_uau(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ +static inline void OP_NAME##_uau(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) { \ assert((uintptr_t) src0 % 128 == 0); \ hvx_arith_loop_body(HVX_UVector, HVX_Vector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO); \ } \ -static inline void OP_NAME##_uua(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ +static inline void OP_NAME##_uua(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) { \ assert((uintptr_t) src1 % 128 == 0); \ hvx_arith_loop_body(HVX_UVector, HVX_UVector, HVX_Vector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO); \ } \ -static inline void OP_NAME##_uuu(uint8_t * restrict dst, const uint8_t * restrict src0, const uint8_t * restrict src1, uint32_t n) { \ +static inline void OP_NAME##_uuu(uint8_t * dst, const uint8_t * src0, const uint8_t * src1, uint32_t n) { \ hvx_arith_loop_body(HVX_UVector, HVX_UVector, HVX_UVector, sizeof(ELEM_TYPE), hvx_vec_store_u, OP_MACRO); \ } \ @@ -358,6 +358,34 @@ static inline void hvx_clamp_scalar_f32(uint8_t * restrict dst, const uint8_t * } } +// +// Abs +// + +static inline void hvx_abs_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + assert((unsigned long) dst % 128 == 0); + assert((unsigned long) src % 128 == 0); + + HVX_Vector * restrict vdst = (HVX_Vector *) dst; + HVX_Vector * restrict vsrc = (HVX_Vector *) src; + + const uint32_t elem_size = sizeof(float); + const uint32_t epv = 128 / elem_size; + const uint32_t nvec = n / epv; + const uint32_t nloe = n % epv; + + uint32_t i = 0; + + _Pragma("unroll(4)") + for (; i < nvec; i++) { + vdst[i] = hvx_vec_abs_f32(vsrc[i]); + } + if (nloe) { + HVX_Vector v = hvx_vec_abs_f32(vsrc[i]); + hvx_vec_store_a((void *) &vdst[i], nloe * elem_size, v); + } +} + // // Square // diff --git a/ggml/src/ggml-hexagon/htp/hvx-log.h b/ggml/src/ggml-hexagon/htp/hvx-log.h index 7013dae785ac..a209f88d555c 100644 --- a/ggml/src/ggml-hexagon/htp/hvx-log.h +++ b/ggml/src/ggml-hexagon/htp/hvx-log.h @@ -62,4 +62,28 @@ static inline HVX_Vector hvx_vec_log_f32(HVX_Vector x) { return hvx_vec_add_f32_f32(term_e, res); } +static inline void hvx_log_f32_aa(uint8_t * restrict dst, const uint8_t * restrict src, uint32_t n) { + assert((unsigned long) dst % 128 == 0); + assert((unsigned long) src % 128 == 0); + + HVX_Vector * restrict vdst = (HVX_Vector *) dst; + HVX_Vector * restrict vsrc = (HVX_Vector *) src; + + const uint32_t elem_size = sizeof(float); + const uint32_t epv = 128 / elem_size; + const uint32_t nvec = n / epv; + const uint32_t nloe = n % epv; + + uint32_t i = 0; + + _Pragma("unroll(4)") + for (; i < nvec; i++) { + vdst[i] = hvx_vec_log_f32(vsrc[i]); + } + if (nloe) { + HVX_Vector v = hvx_vec_log_f32(vsrc[i]); + hvx_vec_store_a((void *) &vdst[i], nloe * elem_size, v); + } +} + #endif /* HVX_LOG_H */ diff --git a/ggml/src/ggml-hexagon/htp/hvx-quant.h b/ggml/src/ggml-hexagon/htp/hvx-quant.h new file mode 100644 index 000000000000..6b172cd63c3a --- /dev/null +++ b/ggml/src/ggml-hexagon/htp/hvx-quant.h @@ -0,0 +1,165 @@ +#ifndef HVX_QUANT_H +#define HVX_QUANT_H + +#include +#include +#include + +#include "hvx-arith.h" +#include "hvx-base.h" +#include "hvx-reduce.h" +#include "hvx-repl.h" +#include "hvx-utils.h" + +#ifndef GGML_COMMON_DECL_C +#define GGML_COMMON_DECL_C +#endif +#include "ggml-common.h" +#include "ggml-impl.h" + +static inline void hvx_quantize_row_q8_0_f32(void * restrict dst_ptr, const float * restrict src_ptr, int n) { + const int nb = n / QK8_0; + block_q8_0 * dst = (block_q8_0 *) dst_ptr; + HVX_Vector zero = Q6_V_vzero(); + + int i = 0; + for (; i + 3 < nb; i += 4) { + HVX_Vector * vx = (HVX_Vector *) (src_ptr + i * QK8_0); + + HVX_Vector vmax0_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[0])); + HVX_Vector vmax1_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[1])); + HVX_Vector vmax2_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[2])); + HVX_Vector vmax3_sf = hvx_vec_reduce_max_f32(hvx_vec_abs_f32(vx[3])); + + HVX_Vector vx0_qf = Q6_Vqf32_vsub_VsfVsf(vx[0], zero); + HVX_Vector vx1_qf = Q6_Vqf32_vsub_VsfVsf(vx[1], zero); + HVX_Vector vx2_qf = Q6_Vqf32_vsub_VsfVsf(vx[2], zero); + HVX_Vector vx3_qf = Q6_Vqf32_vsub_VsfVsf(vx[3], zero); + + HVX_Vector vmax0_qf = Q6_Vqf32_vsub_VsfVsf(vmax0_sf, zero); + HVX_Vector vmax1_qf = Q6_Vqf32_vsub_VsfVsf(vmax1_sf, zero); + HVX_Vector vmax2_qf = Q6_Vqf32_vsub_VsfVsf(vmax2_sf, zero); + HVX_Vector vmax3_qf = Q6_Vqf32_vsub_VsfVsf(vmax3_sf, zero); + + HVX_Vector vmax01_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vmax1_qf, vmax0_qf))); + HVX_Vector vmax23_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vmax3_qf, vmax2_qf))); + + HVX_Vector vx01_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vx1_qf, vx0_qf))); + HVX_Vector vx23_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(vx3_qf, vx2_qf))); + + HVX_Vector vd01_qf16 = Q6_Vqf16_vmpy_VhfVhf(vmax01_hf, Q6_Vh_vsplat_R(0x2008)); // 1.0 / 127.0 + HVX_Vector vd23_qf16 = Q6_Vqf16_vmpy_VhfVhf(vmax23_hf, Q6_Vh_vsplat_R(0x2008)); // 1.0 / 127.0 + HVX_Vector vd01_hf = Q6_Vhf_equals_Vqf16(vd01_qf16); + HVX_Vector vd23_hf = Q6_Vhf_equals_Vqf16(vd23_qf16); + + HVX_Vector vd01_inv_hf = hvx_vec_inverse_f16(vd01_hf); + HVX_Vector vd23_inv_hf = hvx_vec_inverse_f16(vd23_hf); + vx01_hf = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(vx01_hf, vd01_inv_hf)); + vx23_hf = Q6_Vhf_equals_Vqf16(Q6_Vqf16_vmpy_VhfVhf(vx23_hf, vd23_inv_hf)); + + HVX_Vector vx01_i16 = hvx_vec_i16_from_hf_rnd_sat(vx01_hf); + HVX_Vector vx23_i16 = hvx_vec_i16_from_hf_rnd_sat(vx23_hf); + HVX_Vector vx_i8 = Q6_Vb_vpack_VhVh_sat(vx23_i16, vx01_i16); + + hvx_vec_store_u(&dst[i + 0].d, 2, vd01_hf); + hvx_vec_store_u(dst[i + 0].qs, 32, vx_i8); + + hvx_vec_store_u(&dst[i + 1].d, 2, Q6_V_vror_VR(vd01_hf, 64)); + hvx_vec_store_u(dst[i + 1].qs, 32, Q6_V_vror_VR(vx_i8, 32)); + + hvx_vec_store_u(&dst[i + 2].d, 2, vd23_hf); + hvx_vec_store_u(dst[i + 2].qs, 32, Q6_V_vror_VR(vx_i8, 64)); + + hvx_vec_store_u(&dst[i + 3].d, 2, Q6_V_vror_VR(vd23_hf, 64)); + hvx_vec_store_u(dst[i + 3].qs, 32, Q6_V_vror_VR(vx_i8, 96)); + } + + for (; i < nb; i++) { + const float * block_src = src_ptr + i * QK8_0; + HVX_Vector vx = *(const HVX_UVector *) block_src; + HVX_Vector v_abs = hvx_vec_abs_f32(vx); + HVX_Vector v_max = hvx_vec_reduce_max_f32(v_abs); + float amax = hvx_vec_get_f32(v_max); + + const float d = amax / 127.0f; + const float id = d ? (1.0f / d) : 0.0f; + dst[i].d = GGML_FP32_TO_FP16(d); + + HVX_Vector vid = hvx_vec_splat_f32(id); + HVX_Vector v_scaled = hvx_vec_mul_f32_f32(vx, vid); + HVX_Vector v_scaled_qf = Q6_Vqf32_vsub_VsfVsf(v_scaled, zero); + HVX_Vector v_scaled_hf = Q6_Vh_vdeal_Vh(Q6_Vhf_equals_Wqf32(Q6_W_vcombine_VV(zero, v_scaled_qf))); + HVX_Vector v_i16 = hvx_vec_i16_from_hf_rnd_sat(v_scaled_hf); + HVX_Vector v_i8 = Q6_Vb_vpack_VhVh_sat(zero, v_i16); + + hvx_vec_store_u(dst[i].qs, 32, v_i8); + } +} + +static inline void hvx_dequantize_row_q8_0_f32(float * restrict dst_ptr, const void * restrict src_ptr, int n) { + const int nb = n / QK8_0; + const block_q8_0 * src = (const block_q8_0 *) src_ptr; + + for (int i = 0; i < nb; i++) { + HVX_Vector vd_f16 = Q6_Vh_vsplat_R(*(const int16_t *) &src[i].d); + HVX_VectorPair vp_f32 = hvx_vec_f16_to_f32(vd_f16); + HVX_Vector vd = Q6_V_lo_W(vp_f32); + + HVX_Vector vq_i8 = *(const HVX_UVector *) src[i].qs; + + HVX_VectorPair p16 = Q6_Wh_vunpack_Vb(vq_i8); + HVX_Vector v_i16 = Q6_V_lo_W(p16); + HVX_VectorPair p32 = Q6_Ww_vunpack_Vh(v_i16); + HVX_Vector v_i32 = Q6_V_lo_W(p32); + + HVX_Vector v_f32 = Q6_Vsf_equals_Vw(v_i32); + HVX_Vector res = hvx_vec_mul_f32_f32(v_f32, vd); + + float * block_dst = dst_ptr + i * QK8_0; + hvx_vmem(block_dst) = res; + } +} + +static inline void hvx_dequantize_row_q8_0_f16(__fp16 * restrict dst_ptr, const void * restrict src_ptr, int n) { + const int nb = n / QK8_0; + const block_q8_0 * src = (const block_q8_0 *) src_ptr; + + for (int i = nb - 1; i >= 0; i--) { + HVX_Vector vd_f16 = Q6_Vh_vsplat_R(*(const int16_t *) &src[i].d); + HVX_VectorPair vp_f32 = hvx_vec_f16_to_f32(vd_f16); + HVX_Vector vd = Q6_V_lo_W(vp_f32); + + HVX_Vector vq_i8 = *(const HVX_UVector *) src[i].qs; + + HVX_VectorPair p16 = Q6_Wh_vunpack_Vb(vq_i8); + HVX_Vector v_i16 = Q6_V_lo_W(p16); + HVX_VectorPair p32 = Q6_Ww_vunpack_Vh(v_i16); + HVX_Vector v_i32 = Q6_V_lo_W(p32); + + HVX_Vector v_f32 = Q6_Vsf_equals_Vw(v_i32); + HVX_Vector res_f32 = hvx_vec_mul_f32_f32(v_f32, vd); + + HVX_Vector res_f16 = hvx_vec_f32_to_f16(res_f32, Q6_V_vzero()); + + __fp16 * block_dst = dst_ptr + i * QK8_0; + hvx_vec_store_u(block_dst, QK8_0 * sizeof(__fp16), res_f16); + } +} + +static inline void hvx_dequantize_row_f16_f32(float * restrict dst_ptr, const void * restrict src_ptr, int n) { + const int nb = n / 32; + const _Float16 * src = (const _Float16 *) src_ptr; + + for (int i = 0; i < nb; i++) { + HVX_Vector v_f16 = *(const HVX_UVector *) (src + i * 32); + HVX_VectorPair vp_f32 = hvx_vec_f16_to_f32(v_f16); + HVX_Vector res = Q6_V_lo_W(vp_f32); + + float * block_dst = dst_ptr + i * 32; + hvx_vmem(block_dst) = res; + } +} + + + +#endif // HVX_QUANT_H diff --git a/ggml/src/ggml-hexagon/htp/main.c b/ggml/src/ggml-hexagon/htp/main.c index 880e20c99597..fe7d093a81c1 100644 --- a/ggml/src/ggml-hexagon/htp/main.c +++ b/ggml/src/ggml-hexagon/htp/main.c @@ -18,6 +18,7 @@ #include #include #include +#include #include "hex-utils.h" #include "hex-dma.h" @@ -32,6 +33,7 @@ #include "htp_iface.h" #include "work-queue.h" #include "hex-profile.h" +#include "allreduce-ops.h" #define HMX_QUEUE_CAPACITY 16 #define HMX_QUEUE_STACK_SIZE 16384 @@ -46,6 +48,36 @@ struct htp_handle { struct htp_context * ctx; }; +static inline void * htp_mmap(uint32_t fd, uint32_t size) { + void * va = (void *)-1; + for (int retry = 0; retry < 2; retry++) { +#if __HVX_ARCH__ > 73 + va = HAP_mmap2(NULL, size, HAP_PROT_READ | HAP_PROT_WRITE, 0, fd, 0); +#else + if (size > HTP_MMAP_MAX_VMEM) { + FARF(ERROR, "mmap failed : size %u exceeds 2GB limit for HAP_mmap", (uint32_t) size); + abort(); + } + va = HAP_mmap(NULL, size, HAP_PROT_READ | HAP_PROT_WRITE, 0, fd, 0); +#endif + if (va != (void *)-1 && va != NULL) { + return va; + } + if (retry == 0) { + FARF(HIGH, "mmap failed first try (va %p fd %u size %u), retrying...", va, fd, size); + } + } + return NULL; +} + +static inline void htp_munmap(void * va, uint32_t size) { +#if __HVX_ARCH__ > 73 + HAP_munmap2(va, size); +#else + HAP_munmap(va, size); +#endif +} + AEEResult htp_iface_open(const char * uri, remote_handle64 * handle) { (void) uri; struct htp_handle * h = calloc(1, sizeof(*h)); @@ -127,11 +159,7 @@ AEEResult htp_iface_close(remote_handle64 handle) { // release the mmaps (if any) for (uint32_t i=0; immap[i].size) { -#if __HVX_ARCH__ > 73 - HAP_munmap2((void *) ctx->mmap[i].base, ctx->mmap[i].size); -#else - HAP_munmap((void *) ctx->mmap[i].base, ctx->mmap[i].size); -#endif + htp_munmap((void *) ctx->mmap[i].base, ctx->mmap[i].size); ctx->mmap[i].size = 0; ctx->mmap[i].base = NULL; ctx->mmap[i].fd = -1; @@ -175,18 +203,9 @@ AEEResult htp_iface_mmap(remote_handle64 handle, uint32_t fd, uint32_t size) { struct htp_mmap *m = &ctx->mmap[i]; if (!m->size) { FARF(HIGH, "mmap : fd %u size %u", fd, size); -#if __HVX_ARCH__ > 73 - void *va = HAP_mmap2(NULL, size, HAP_PROT_READ | HAP_PROT_WRITE, 0, fd, 0); -#else - if (size > HTP_MMAP_MAX_VMEM) { // HAP_mmap has a size limit of 2GB - FARF(ERROR, "mmap failed : size %u exceeds 2GB limit for HAP_mmap", (uint32_t) size); - abort(); // can't do much else at this point - } - - void *va = HAP_mmap(NULL, size, HAP_PROT_READ | HAP_PROT_WRITE, 0, fd, 0); -#endif - if (va == (void*)-1) { - FARF(ERROR, "mmap failed : va %p fd %u size %u", va, fd, (uint32_t) size); + void *va = htp_mmap(fd, size); + if (va == NULL) { + FARF(ERROR, "mmap failed : fd %u size %u", fd, (uint32_t) size); return AEE_EFAILED; } @@ -212,11 +231,7 @@ AEEResult htp_iface_munmap(remote_handle64 handle, uint32 fd) { struct htp_mmap *m = &ctx->mmap[i]; if (fd < 0 || m->fd == fd) { FARF(HIGH, "unmmap : base %p fd %u size %u", (void*) m->base, m->fd, (uint32_t) m->size); -#if __HVX_ARCH__ > 73 - HAP_munmap2((void *) m->base, m->size); -#else - HAP_munmap((void *) m->base, m->size); -#endif + htp_munmap((void *) m->base, m->size); m->size = 0; m->base = NULL; m->fd = -1; @@ -228,7 +243,7 @@ AEEResult htp_iface_munmap(remote_handle64 handle, uint32 fd) { static void vtcm_acquire(struct htp_context * ctx) { if (!ctx->vtcm_valid) { - int err = HAP_compute_res_acquire_cached(ctx->vtcm_rctx, 1000000u); + int err = HAP_compute_res_acquire_cached(ctx->vtcm_rctx, 10000000u); if (err != 0) { FARF(ERROR, "ggml-hex: failed to acquire VTCM: 0x%08x", (unsigned)err); abort(); @@ -692,8 +707,45 @@ static inline void profile_stop(uint32_t mode, struct profile_data * d) { } } +static int op_fence(struct htp_ops_context * octx) { + struct htp_context *ctx = octx->ctx; + struct htp_thread_trace * tr = &ctx->trace[0]; + const uint32_t seq = (uint32_t) octx->op_params[0]; + + htp_trace_event_start(tr, HTP_TRACE_EVT_FENCE, (uint16_t) seq); + + const struct htp_tensor * sync = octx->src[0]; + atomic_uint * sync_fence = (atomic_uint *) sync->data; + uint64_t spins = 0; + while (1) { + Q6_dccleaninva_A((void *) sync_fence); + asm volatile ("syncht" : : : "memory"); + uint32_t val = atomic_load(&sync_fence[0]); + if ((int32_t)(val - seq) >= 0) { + break; + } + if (++spins > HTP_FENCE_TIMEOUT) { + FARF(ERROR, "ggml-hex: sync-wait TIMEOUT : fence %p spins %llu seq %u\n", sync_fence, spins, seq); + break; + } + hex_pause(); + } + + htp_trace_event_stop(tr, HTP_TRACE_EVT_FENCE, (uint16_t) seq); + + FARF(HIGH, "ggml-hex: sync-done : fence %p spins %llu seq %u\n", sync_fence, spins, seq); + return HTP_STATUS_OK; +} + static int execute_op(struct htp_ops_context * octx) { switch (octx->op) { + case HTP_OP_FENCE: + return op_fence(octx); + + case HTP_OP_ALLREDUCE: + case HTP_OP_ALLREDUCE_ADD: + return op_allreduce(octx); + case HTP_OP_MUL_MAT: case HTP_OP_MUL_MAT_ADD: return op_matmul(octx); @@ -701,11 +753,8 @@ static int execute_op(struct htp_ops_context * octx) { case HTP_OP_MUL_MAT_ID: return op_matmul_id(octx); - case HTP_OP_MUL_MAT_QKV: - return op_matmul_qkv(octx); - - case HTP_OP_MUL_MAT_FFN: - return op_matmul_ffn(octx); + case HTP_OP_MUL_MAT_NX: + return op_matmul_nx(octx); case HTP_OP_MUL: case HTP_OP_ADD: @@ -728,6 +777,8 @@ static int execute_op(struct htp_ops_context * octx) { case HTP_OP_UNARY_NEG: case HTP_OP_UNARY_EXP: case HTP_OP_UNARY_TANH: + case HTP_OP_UNARY_ABS: + case HTP_OP_UNARY_LOG: case HTP_OP_L2_NORM: return op_unary(octx); @@ -818,12 +869,8 @@ static inline bool reuse_buf(struct htp_context *ctx, uint32_t *m_reuse, struct static inline void drop_mmap(struct htp_context *ctx, struct htp_mmap *m) { if (m->size) { - FARF(HIGH, "unmap : fd %u base %p size %u", m->fd, (void*) m->base, (uint32_t) m->size); -#if __HVX_ARCH__ > 73 - HAP_munmap2((void *) m->base, m->size); -#else - HAP_munmap((void *) m->base, m->size); -#endif + FARF(ALWAYS, "unmap : fd %u base %p size %u", m->fd, (void*) m->base, (uint32_t) m->size); + htp_munmap((void *) m->base, m->size); m->size = 0; m->base = 0; m->fd = -1; @@ -837,18 +884,9 @@ static inline void mmap_buf(struct htp_context *ctx, struct htp_buf_desc *b) { for (uint32_t i=0; i < HTP_MAX_MMAPS; i++) { struct htp_mmap *m = &ctx->mmap[i]; if (!m->size) { -#if __HVX_ARCH__ > 73 - void *va = HAP_mmap2(NULL, b->size, HAP_PROT_READ | HAP_PROT_WRITE, 0, b->fd, 0); -#else - if (b->size > HTP_MMAP_MAX_VMEM) { // HAP_mmap has a size limit of 2GB - FARF(ERROR, "mmap failed : size %u exceeds 2GB limit for HAP_mmap", (uint32_t) b->size); - abort(); // can't do much else at this point - } - - void *va = HAP_mmap(NULL, b->size, HAP_PROT_READ | HAP_PROT_WRITE, 0, b->fd, 0); -#endif - if (va == (void*)-1) { - FARF(ERROR, "mmap failed : va %p fd %u size %u", va, b->fd, (uint32_t) b->size); + void *va = htp_mmap(b->fd, b->size); + if (va == NULL) { + FARF(ERROR, "mmap failed : fd %u size %u", b->fd, (uint32_t) b->size); abort(); // can't do much else at this point } @@ -856,10 +894,13 @@ static inline void mmap_buf(struct htp_context *ctx, struct htp_buf_desc *b) { m->fd = b->fd; m->size = b->size; - FARF(HIGH, "mmap : fd %u base %p size %u", m->fd, (void*) m->base, (uint32_t) m->size); + FARF(ALWAYS, "mmap : fd %u base %p size %u", m->fd, (void*) m->base, (uint32_t) m->size); return; } } + + FARF(ERROR, "mmap failed : exceeded mapping capacity limit of %u", HTP_MAX_MMAPS); + abort(); } static void prep_op_bufs(struct htp_context *ctx, struct htp_buf_desc *bufs, uint32_t n_bufs) { @@ -1081,6 +1122,7 @@ static void process_opbatch(struct htp_context * ctx, const struct htp_opbatch_r rsp.usecs = batch_prof.usecs; rsp.cycles_start = batch_prof.cycles_start; rsp.cycles_stop = batch_prof.cycles_stop; + rsp.seq = req->seq; if (ctx->profiler == HTP_PROF_TRACE) { for (int t = 0; t <= HTP_MAX_NTHREADS; t++) { diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.c b/ggml/src/ggml-hexagon/htp/matmul-ops.c index 9d385469ae9f..a6adc0e61fa6 100644 --- a/ggml/src/ggml-hexagon/htp/matmul-ops.c +++ b/ggml/src/ggml-hexagon/htp/matmul-ops.c @@ -64,6 +64,7 @@ typedef struct { struct htp_mm_context { const char * type; struct htp_ops_context * octx; + const struct htp_tensor * act; void (*vec_dot_1x1)(const uint32_t n, float * restrict s0, const void * restrict vx0, @@ -478,7 +479,7 @@ static void hvx_mv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void \ htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ DOT_2X1(ne10, dst_ptr, w_tile, src1_col, valid_rows, NULL); \ - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ \ if (push_ct < ct_end) { \ dma_queue_push(dma_queue, dma_make_ptr((uint8_t *)w_tile, src0_row + push_ct * tile_row_stride), \ @@ -502,150 +503,67 @@ static void hvx_mv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void } \ } -#define MATMUL_QKV_2D_REPACKED_IMPL(SUFFIX, TILE_SIZE, DOT_2X2, DOT_2X1) \ -static void hvx_mm_qkv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void * data) { \ +#define MATMUL_NX_2D_REPACKED_IMPL(SUFFIX, TILE_SIZE, DOT_2X2, DOT_2X1) \ +static void hvx_mm_nx_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void * data) { \ struct htp_mm_context * mmctx = data; \ struct htp_ops_context * octx = mmctx->octx; \ + const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; \ + const uint32_t n_weights = kparams->n_weights; \ \ - const struct htp_tensor * restrict src0 = octx->src[0]; /* Wk */ \ - const struct htp_tensor * restrict src1 = octx->src[1]; /* x */ \ - const struct htp_tensor * restrict src2 = octx->src[2]; /* Wv */ \ - const struct htp_tensor * restrict src3 = octx->src[3]; /* Wq */ \ - const struct htp_tensor * restrict dst_k = octx->dsts[0]; \ - const struct htp_tensor * restrict dst_v = octx->dsts[1]; \ - const struct htp_tensor * restrict dst_q = octx->dsts[2]; \ - \ - const uint32_t ne00 = src0->ne[0]; \ - const uint32_t ne10 = src1->ne[0]; \ - const uint32_t src1_nrows = src1->ne[1] * src1->ne[2] * src1->ne[3]; \ - \ - const size_t dst_k_row_size = dst_k->nb[1]; /* K and V share output width */ \ - const size_t dst_q_row_size = dst_q->nb[1]; /* Q may be wider (GQA) */ \ + const struct htp_tensor * restrict act = octx->src[n_weights]; /* x */ \ + const uint32_t ne10 = act->ne[0]; \ + const uint32_t src1_nrows = act->ne[1] * act->ne[2] * act->ne[3]; \ const size_t src1_stride = mmctx->vtcm_src1_stride; \ \ - uint8_t * restrict vtcm_src0_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; \ - uint8_t * restrict vtcm_src2_ptr = mmctx->vtcm_src2 + mmctx->vtcm_src2_size_per_thread * ith; \ - uint8_t * restrict vtcm_src3_ptr = mmctx->vtcm_src3 + mmctx->vtcm_src3_size_per_thread * ith; \ - uint8_t * restrict src1_data = mmctx->vtcm_src1; \ + uint8_t * restrict vtcm_weight_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; \ + uint8_t * restrict src1_data = mmctx->vtcm_src1; \ \ struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ - \ - const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; \ const uint32_t n_prefetch = kparams->n_prefetch; \ assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0); \ \ - const uint8_t * restrict src0_row = (const uint8_t *) src0->data; \ - const uint8_t * restrict src2_row = (const uint8_t *) src2->data; \ - const uint8_t * restrict src3_row = (const uint8_t *) src3->data; \ - \ const uint32_t tile_size = TILE_SIZE; \ const uint32_t aligned_tile_size = hex_align_up(tile_size, 128); \ - \ - uint32_t n_k_tiles_w = ne00 / 32; \ uint32_t n_k_tiles_a = ne10 / 32; \ - uint32_t tile_row_stride = n_k_tiles_w * tile_size; \ uint32_t tile_row_transfer_size_aligned = n_k_tiles_a * aligned_tile_size; \ \ dma_queue * dma_queue = octx->ctx->dma[ith]; \ \ - /* 1. Process K and V together */ \ - const uint32_t src0_nrows_kv = src0->ne[1] * src0->ne[2] * src0->ne[3]; /* src0 is Wk */ \ - uint32_t src0_nrows_per_thread_kv = (src0_nrows_kv + nth - 1) / nth; \ - src0_nrows_per_thread_kv = hex_round_up(src0_nrows_per_thread_kv, 32); \ - \ - const uint32_t start_row_kv = src0_nrows_per_thread_kv * ith; \ - const uint32_t end_row_kv = MIN(start_row_kv + src0_nrows_per_thread_kv, src0_nrows_kv); \ - \ - uint32_t ct_start_kv = start_row_kv / 32; \ - uint32_t ct_end_kv = (end_row_kv + 31) / 32; \ - \ - uint32_t push_ct = ct_start_kv; \ - if (start_row_kv < end_row_kv) { \ - for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end_kv; d++, push_ct++) { \ - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + d * tile_row_transfer_size_aligned, \ - src0_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr + d * tile_row_transfer_size_aligned, \ - src2_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ - } \ - } \ - \ hvx_mm_run_quant_task(mmctx, ith); \ \ - if (start_row_kv < end_row_kv) { \ - \ - for (uint32_t ct = ct_start_kv; ct < ct_end_kv; ct++) { \ - const uint8_t * w_tile_k = dma_queue_pop(dma_queue).dst; \ - const uint8_t * w_tile_v = dma_queue_pop(dma_queue).dst; \ - \ - int valid_rows = (int)src0->ne[1] - (int)(ct * 32); \ - valid_rows = MIN(32, MAX(0, valid_rows)); \ - \ - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ith); \ - uint32_t ir1 = 0; \ - for (; ir1 + 1 < src1_nrows; ir1 += 2) { \ - const uint8_t * restrict src1_col0 = (const uint8_t *) (src1_data + (ir1+0) * src1_stride); \ - const uint8_t * restrict src1_col1 = (const uint8_t *) (src1_data + (ir1+1) * src1_stride); \ - \ - float * restrict dst_row0_k = (float *) (dst_k->data + ((ir1+0) * dst_k_row_size)); \ - float * restrict dst_row1_k = (float *) (dst_k->data + ((ir1+1) * dst_k_row_size)); \ - float * dst_ptr0_k = &dst_row0_k[ct * 32]; \ - float * dst_ptr1_k = &dst_row1_k[ct * 32]; \ - \ - float * restrict dst_row0_v = (float *) (dst_v->data + ((ir1+0) * dst_k_row_size)); \ - float * restrict dst_row1_v = (float *) (dst_v->data + ((ir1+1) * dst_k_row_size)); \ - float * dst_ptr0_v = &dst_row0_v[ct * 32]; \ - float * dst_ptr1_v = &dst_row1_v[ct * 32]; \ - \ - DOT_2X2(ne10, dst_ptr0_k, dst_ptr1_k, w_tile_k, src1_col0, src1_col1, valid_rows, NULL, NULL); \ - DOT_2X2(ne10, dst_ptr0_v, dst_ptr1_v, w_tile_v, src1_col0, src1_col1, valid_rows, NULL, NULL); \ - } \ + for (uint32_t widx = 0; widx < n_weights; widx++) { \ + const struct htp_tensor * restrict src_w = octx->src[widx]; \ + const struct htp_tensor * restrict dst = octx->dsts[widx]; \ + if (!src_w || !dst) continue; \ \ - for (; ir1 < src1_nrows; ++ir1) { \ - const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + ir1 * src1_stride); \ - \ - float * restrict dst_row_k = (float *) (dst_k->data + (ir1 * dst_k_row_size)); \ - float * dst_ptr_k = &dst_row_k[ct * 32]; \ - \ - float * restrict dst_row_v = (float *) (dst_v->data + (ir1 * dst_k_row_size)); \ - float * dst_ptr_v = &dst_row_v[ct * 32]; \ - \ - DOT_2X1(ne10, dst_ptr_k, w_tile_k, src1_col, valid_rows, NULL); \ - DOT_2X1(ne10, dst_ptr_v, w_tile_v, src1_col, valid_rows, NULL); \ - } \ - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ith); \ + const uint32_t ne00 = src_w->ne[0]; \ + const uint32_t ne01 = src_w->ne[1]; \ + const size_t dst_row_size = dst->nb[1]; \ + const uint8_t * restrict src_w_row = (const uint8_t *) src_w->data; \ \ - if (push_ct < ct_end_kv) { \ - dma_queue_push(dma_queue, dma_make_ptr((uint8_t *)w_tile_k, src0_row + push_ct * tile_row_stride), \ - aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ - dma_queue_push(dma_queue, dma_make_ptr((uint8_t *)w_tile_v, src2_row + push_ct * tile_row_stride), \ - aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ - push_ct++; \ - } \ - } \ - } \ + uint32_t n_k_tiles_w = ne00 / 32; \ + uint32_t tile_row_stride = n_k_tiles_w * tile_size; \ \ - /* 2. Process Q separately */ \ - const uint32_t src0_nrows_q = src3->ne[1] * src3->ne[2] * src3->ne[3]; /* src3 is Wq */ \ - uint32_t src0_nrows_per_thread_q = (src0_nrows_q + nth - 1) / nth; \ - src0_nrows_per_thread_q = hex_round_up(src0_nrows_per_thread_q, 32); \ + const uint32_t src0_nrows = ne01 * src_w->ne[2] * src_w->ne[3]; \ + uint32_t src0_nrows_per_thread = (src0_nrows + nth - 1) / nth; \ + src0_nrows_per_thread = hex_round_up(src0_nrows_per_thread, 32); \ \ - const uint32_t start_row_q = src0_nrows_per_thread_q * ith; \ - const uint32_t end_row_q = MIN(start_row_q + src0_nrows_per_thread_q, src0_nrows_q); \ + const uint32_t start_row = src0_nrows_per_thread * ith; \ + const uint32_t end_row = MIN(start_row + src0_nrows_per_thread, src0_nrows); \ + if (start_row >= end_row) continue; \ \ - if (start_row_q < end_row_q) { \ - uint32_t ct_start_q = start_row_q / 32; \ - uint32_t ct_end_q = (end_row_q + 31) / 32; \ + uint32_t ct_start = start_row / 32; \ + uint32_t ct_end = (end_row + 31) / 32; \ \ - uint32_t push_ct = ct_start_q; \ - for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end_q; d++, push_ct++) { \ - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src3_ptr + d * tile_row_transfer_size_aligned, \ - src3_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ + uint32_t push_ct = ct_start; \ + for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end; d++, push_ct++) { \ + dma_queue_push(dma_queue, dma_make_ptr(vtcm_weight_ptr + d * tile_row_transfer_size_aligned, \ + src_w_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ } \ \ - for (uint32_t ct = ct_start_q; ct < ct_end_q; ct++) { \ - const uint8_t * w_tile_q = dma_queue_pop(dma_queue).dst; \ - \ - int valid_rows = (int)src3->ne[1] - (int)(ct * 32); \ + for (uint32_t ct = ct_start; ct < ct_end; ct++) { \ + const uint8_t * w_tile = dma_queue_pop(dma_queue).dst; \ + int valid_rows = (int)ne01 - (int)(ct * 32); \ valid_rows = MIN(32, MAX(0, valid_rows)); \ \ htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ @@ -654,26 +572,24 @@ static void hvx_mm_qkv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, const uint8_t * restrict src1_col0 = (const uint8_t *) (src1_data + (ir1+0) * src1_stride); \ const uint8_t * restrict src1_col1 = (const uint8_t *) (src1_data + (ir1+1) * src1_stride); \ \ - float * restrict dst_row0_q = (float *) (dst_q->data + ((ir1+0) * dst_q_row_size)); \ - float * restrict dst_row1_q = (float *) (dst_q->data + ((ir1+1) * dst_q_row_size)); \ - float * dst_ptr0_q = &dst_row0_q[ct * 32]; \ - float * dst_ptr1_q = &dst_row1_q[ct * 32]; \ + float * restrict dst_row0 = (float *) (dst->data + ((ir1+0) * dst_row_size)); \ + float * restrict dst_row1 = (float *) (dst->data + ((ir1+1) * dst_row_size)); \ + float * dst_ptr0 = &dst_row0[ct * 32]; \ + float * dst_ptr1 = &dst_row1[ct * 32]; \ \ - DOT_2X2(ne10, dst_ptr0_q, dst_ptr1_q, w_tile_q, src1_col0, src1_col1, valid_rows, NULL, NULL); \ + DOT_2X2(ne10, dst_ptr0, dst_ptr1, w_tile, src1_col0, src1_col1, valid_rows, NULL, NULL); \ } \ \ for (; ir1 < src1_nrows; ++ir1) { \ const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + ir1 * src1_stride); \ - \ - float * restrict dst_row_q = (float *) (dst_q->data + (ir1 * dst_q_row_size)); \ - float * dst_ptr_q = &dst_row_q[ct * 32]; \ - \ - DOT_2X1(ne10, dst_ptr_q, w_tile_q, src1_col, valid_rows, NULL); \ + float * restrict dst_row = (float *) (dst->data + (ir1 * dst_row_size)); \ + float * dst_ptr = &dst_row[ct * 32]; \ + DOT_2X1(ne10, dst_ptr, w_tile, src1_col, valid_rows, NULL); \ } \ htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ \ - if (push_ct < ct_end_q) { \ - dma_queue_push(dma_queue, dma_make_ptr((uint8_t *)w_tile_q, src3_row + push_ct * tile_row_stride), \ + if (push_ct < ct_end) { \ + dma_queue_push(dma_queue, dma_make_ptr((uint8_t *)w_tile, src_w_row + push_ct * tile_row_stride), \ aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ push_ct++; \ } \ @@ -681,121 +597,6 @@ static void hvx_mm_qkv_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, } \ } -#define MATMUL_FFN_2D_REPACKED_IMPL(SUFFIX, TILE_SIZE, DOT_2X2, DOT_2X1) \ -static void hvx_mm_ffn_2d_repacked_##SUFFIX(unsigned int nth, unsigned int ith, void * data) { \ - struct htp_mm_context * mmctx = data; \ - struct htp_ops_context * octx = mmctx->octx; \ - \ - const struct htp_tensor * restrict src0 = octx->src[0]; /* Wgate */ \ - const struct htp_tensor * restrict src1 = octx->src[1]; /* y */ \ - const struct htp_tensor * restrict src2 = octx->src[2]; /* Wup */ \ - const struct htp_tensor * restrict dst_gate = octx->dsts[0]; \ - const struct htp_tensor * restrict dst_up = octx->dsts[1]; \ - \ - const uint32_t ne00 = src0->ne[0]; \ - const uint32_t ne01 = src0->ne[1]; \ - const uint32_t ne10 = src1->ne[0]; \ - const uint32_t src1_nrows = src1->ne[1] * src1->ne[2] * src1->ne[3]; \ - \ - const size_t dst_row_size = dst_gate->nb[1]; \ - const size_t src1_stride = mmctx->vtcm_src1_stride; \ - \ - uint8_t * restrict vtcm_src0_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; \ - uint8_t * restrict vtcm_src2_ptr = mmctx->vtcm_src2 + mmctx->vtcm_src2_size_per_thread * ith; \ - uint8_t * restrict src1_data = mmctx->vtcm_src1; \ - \ - struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ - \ - const uint8_t * restrict src0_row = (const uint8_t *) src0->data; \ - const uint8_t * restrict src2_row = (const uint8_t *) src2->data; \ - \ - const uint32_t tile_size = TILE_SIZE; \ - const uint32_t aligned_tile_size = hex_align_up(tile_size, 128); \ - \ - const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; \ - const uint32_t n_prefetch = kparams->n_prefetch; \ - assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0); \ - \ - uint32_t n_k_tiles_w = ne00 / 32; \ - uint32_t n_k_tiles_a = ne10 / 32; \ - uint32_t tile_row_stride = n_k_tiles_w * tile_size; \ - uint32_t tile_row_transfer_size_aligned = n_k_tiles_a * aligned_tile_size; \ - dma_queue * dma_queue = octx->ctx->dma[ith]; \ - \ - const uint32_t src0_nrows = ne01 * src0->ne[2] * src0->ne[3]; \ - const uint32_t src0_start_row = mmctx->src0_nrows_per_thread * ith; \ - const uint32_t src0_end_row = MIN(src0_start_row + mmctx->src0_nrows_per_thread, src0_nrows); \ - \ - uint32_t ct_start = src0_start_row / 32; \ - uint32_t ct_end = (src0_end_row + 31) / 32; \ - \ - uint32_t push_ct = ct_start; \ - if (src0_start_row < src0_end_row) { \ - for (uint32_t d = 0; d < n_prefetch && push_ct < ct_end; d++, push_ct++) { \ - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + d * tile_row_transfer_size_aligned, \ - src0_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr + d * tile_row_transfer_size_aligned, \ - src2_row + push_ct * tile_row_stride), aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ - } \ - } \ - \ - hvx_mm_run_quant_task(mmctx, ith); \ - \ - if (src0_start_row >= src0_end_row) { \ - return; \ - } \ - \ - for (uint32_t ct = ct_start; ct < ct_end; ct++) { \ - const uint8_t * w_tile_gate = dma_queue_pop(dma_queue).dst; \ - const uint8_t * w_tile_up = dma_queue_pop(dma_queue).dst; \ - \ - int valid_rows = (int)ne01 - (int)(ct * 32); \ - valid_rows = MIN(32, MAX(0, valid_rows)); \ - \ - htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ - uint32_t ir1 = 0; \ - for (; ir1 + 1 < src1_nrows; ir1 += 2) { \ - const uint8_t * restrict src1_col0 = (const uint8_t *) (src1_data + (ir1+0) * src1_stride); \ - const uint8_t * restrict src1_col1 = (const uint8_t *) (src1_data + (ir1+1) * src1_stride); \ - \ - float * restrict dst_row0_gate = (float *) (dst_gate->data + ((ir1+0) * dst_row_size)); \ - float * restrict dst_row1_gate = (float *) (dst_gate->data + ((ir1+1) * dst_row_size)); \ - float * dst_ptr0_gate = &dst_row0_gate[ct * 32]; \ - float * dst_ptr1_gate = &dst_row1_gate[ct * 32]; \ - \ - float * restrict dst_row0_up = (float *) (dst_up->data + ((ir1+0) * dst_row_size)); \ - float * restrict dst_row1_up = (float *) (dst_up->data + ((ir1+1) * dst_row_size)); \ - float * dst_ptr0_up = &dst_row0_up[ct * 32]; \ - float * dst_ptr1_up = &dst_row1_up[ct * 32]; \ - \ - DOT_2X2(ne10, dst_ptr0_gate, dst_ptr1_gate, w_tile_gate, src1_col0, src1_col1, valid_rows, NULL, NULL); \ - DOT_2X2(ne10, dst_ptr0_up, dst_ptr1_up, w_tile_up, src1_col0, src1_col1, valid_rows, NULL, NULL); \ - } \ - \ - for (; ir1 < src1_nrows; ++ir1) { \ - const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + ir1 * src1_stride); \ - \ - float * restrict dst_row_gate = (float *) (dst_gate->data + (ir1 * dst_row_size)); \ - float * dst_ptr_gate = &dst_row_gate[ct * 32]; \ - \ - float * restrict dst_row_up = (float *) (dst_up->data + (ir1 * dst_row_size)); \ - float * dst_ptr_up = &dst_row_up[ct * 32]; \ - \ - DOT_2X1(ne10, dst_ptr_gate, w_tile_gate, src1_col, valid_rows, NULL); \ - DOT_2X1(ne10, dst_ptr_up, w_tile_up, src1_col, valid_rows, NULL); \ - } \ - htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ct); \ - \ - if (push_ct < ct_end) { \ - dma_queue_push(dma_queue, dma_make_ptr((uint8_t *)w_tile_gate, src0_row + push_ct * tile_row_stride), \ - aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ - dma_queue_push(dma_queue, dma_make_ptr((uint8_t *)w_tile_up, src2_row + push_ct * tile_row_stride), \ - aligned_tile_size, tile_size, tile_size, n_k_tiles_a); \ - push_ct++; \ - } \ - } \ -} - MATMUL_2D_REPACKED_IMPL(q4_0, 576, tiled_vec_dot_q4_0_32x2, tiled_vec_dot_q4_0_32x1) MATMUL_2D_REPACKED_IMPL(q4_1, 640, tiled_vec_dot_q4_1_32x2, tiled_vec_dot_q4_1_32x1) MATMUL_2D_REPACKED_IMPL(q8_0, 1088, tiled_vec_dot_q8_0_32x2, tiled_vec_dot_q8_0_32x1) @@ -812,7 +613,7 @@ MATMUL_2D_REPACKED_IMPL(mxfp4_flat, 544, flat_vec_dot_mxfp4_32x2, flat_vec_dot static void name(unsigned int nth, unsigned int ith, void * data) { \ struct htp_mm_context * mmctx = data; \ struct htp_ops_context * octx = mmctx->octx; \ - const struct htp_tensor * src = octx->src[1]; \ + const struct htp_tensor * src = mmctx->act; \ const uint32_t ne0 = src->ne[0]; \ const uint32_t ne1 = src->ne[1]; \ const uint32_t ne2 = src->ne[2]; \ @@ -854,7 +655,7 @@ static void quantize_f32_q8_0_tiled_block(unsigned int nth, unsigned int ith, vo struct htp_thread_trace * tr = &octx->ctx->trace[ith]; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_QUANT, mmctx->quant_ib_first[ith]); - const struct htp_tensor * src = octx->src[1]; + const struct htp_tensor * src = mmctx->act; quantize_f32_q8_0_tiled_block_kernel( (const float *) src->data, @@ -878,7 +679,7 @@ static void quantize_f32_q8_1_tiled_block(unsigned int nth, unsigned int ith, vo struct htp_thread_trace * tr = &octx->ctx->trace[ith]; htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_A_QUANT, mmctx->quant_ib_first[ith]); - const struct htp_tensor * src = octx->src[1]; + const struct htp_tensor * src = mmctx->act; quantize_f32_q8_1_tiled_block_kernel( (const float *) src->data, @@ -909,30 +710,17 @@ MATVEC_2D_REPACKED_IMPL(iq4nl_flat, 576, flat_vec_dot_iq4nl_32x1) MATVEC_2D_REPACKED_IMPL(mxfp4_flat, 544, flat_vec_dot_mxfp4_32x1) -MATMUL_QKV_2D_REPACKED_IMPL(q4_0, 576, tiled_vec_dot_q4_0_32x2, tiled_vec_dot_q4_0_32x1) -MATMUL_QKV_2D_REPACKED_IMPL(q4_1, 640, tiled_vec_dot_q4_1_32x2, tiled_vec_dot_q4_1_32x1) -MATMUL_QKV_2D_REPACKED_IMPL(q8_0, 1088, tiled_vec_dot_q8_0_32x2, tiled_vec_dot_q8_0_32x1) -MATMUL_QKV_2D_REPACKED_IMPL(iq4nl, 576, tiled_vec_dot_iq4nl_32x2, tiled_vec_dot_iq4nl_32x1) -MATMUL_QKV_2D_REPACKED_IMPL(mxfp4, 544, tiled_vec_dot_mxfp4_32x2, tiled_vec_dot_mxfp4_32x1) +MATMUL_NX_2D_REPACKED_IMPL(q4_0, 576, tiled_vec_dot_q4_0_32x2, tiled_vec_dot_q4_0_32x1) +MATMUL_NX_2D_REPACKED_IMPL(q4_1, 640, tiled_vec_dot_q4_1_32x2, tiled_vec_dot_q4_1_32x1) +MATMUL_NX_2D_REPACKED_IMPL(q8_0, 1088, tiled_vec_dot_q8_0_32x2, tiled_vec_dot_q8_0_32x1) +MATMUL_NX_2D_REPACKED_IMPL(iq4nl, 576, tiled_vec_dot_iq4nl_32x2, tiled_vec_dot_iq4nl_32x1) +MATMUL_NX_2D_REPACKED_IMPL(mxfp4, 544, tiled_vec_dot_mxfp4_32x2, tiled_vec_dot_mxfp4_32x1) -MATMUL_QKV_2D_REPACKED_IMPL(q4_0_flat, 576, flat_vec_dot_q4_0_32x2, flat_vec_dot_q4_0_32x1) -MATMUL_QKV_2D_REPACKED_IMPL(q4_1_flat, 640, flat_vec_dot_q4_1_32x2, flat_vec_dot_q4_1_32x1) -MATMUL_QKV_2D_REPACKED_IMPL(q8_0_flat, 1088, flat_vec_dot_q8_0_32x2, flat_vec_dot_q8_0_32x1) -MATMUL_QKV_2D_REPACKED_IMPL(iq4nl_flat, 576, flat_vec_dot_iq4nl_32x2, flat_vec_dot_iq4nl_32x1) -MATMUL_QKV_2D_REPACKED_IMPL(mxfp4_flat, 544, flat_vec_dot_mxfp4_32x2, flat_vec_dot_mxfp4_32x1) - - -MATMUL_FFN_2D_REPACKED_IMPL(q4_0, 576, tiled_vec_dot_q4_0_32x2, tiled_vec_dot_q4_0_32x1) -MATMUL_FFN_2D_REPACKED_IMPL(q4_1, 640, tiled_vec_dot_q4_1_32x2, tiled_vec_dot_q4_1_32x1) -MATMUL_FFN_2D_REPACKED_IMPL(q8_0, 1088, tiled_vec_dot_q8_0_32x2, tiled_vec_dot_q8_0_32x1) -MATMUL_FFN_2D_REPACKED_IMPL(iq4nl, 576, tiled_vec_dot_iq4nl_32x2, tiled_vec_dot_iq4nl_32x1) -MATMUL_FFN_2D_REPACKED_IMPL(mxfp4, 544, tiled_vec_dot_mxfp4_32x2, tiled_vec_dot_mxfp4_32x1) - -MATMUL_FFN_2D_REPACKED_IMPL(q4_0_flat, 576, flat_vec_dot_q4_0_32x2, flat_vec_dot_q4_0_32x1) -MATMUL_FFN_2D_REPACKED_IMPL(q4_1_flat, 640, flat_vec_dot_q4_1_32x2, flat_vec_dot_q4_1_32x1) -MATMUL_FFN_2D_REPACKED_IMPL(q8_0_flat, 1088, flat_vec_dot_q8_0_32x2, flat_vec_dot_q8_0_32x1) -MATMUL_FFN_2D_REPACKED_IMPL(iq4nl_flat, 576, flat_vec_dot_iq4nl_32x2, flat_vec_dot_iq4nl_32x1) -MATMUL_FFN_2D_REPACKED_IMPL(mxfp4_flat, 544, flat_vec_dot_mxfp4_32x2, flat_vec_dot_mxfp4_32x1) +MATMUL_NX_2D_REPACKED_IMPL(q4_0_flat, 576, flat_vec_dot_q4_0_32x2, flat_vec_dot_q4_0_32x1) +MATMUL_NX_2D_REPACKED_IMPL(q4_1_flat, 640, flat_vec_dot_q4_1_32x2, flat_vec_dot_q4_1_32x1) +MATMUL_NX_2D_REPACKED_IMPL(q8_0_flat, 1088, flat_vec_dot_q8_0_32x2, flat_vec_dot_q8_0_32x1) +MATMUL_NX_2D_REPACKED_IMPL(iq4nl_flat, 576, flat_vec_dot_iq4nl_32x2, flat_vec_dot_iq4nl_32x1) +MATMUL_NX_2D_REPACKED_IMPL(mxfp4_flat, 544, flat_vec_dot_mxfp4_32x2, flat_vec_dot_mxfp4_32x1) static void hvx_mm_2d(unsigned int nth, unsigned int ith, void * data) { htp_matmul_preamble; @@ -1353,6 +1141,7 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { struct htp_mm_context mmctx_struct = {0}; struct htp_mm_context * mmctx = &mmctx_struct; mmctx->octx = octx; + mmctx->act = src1; const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; @@ -1528,7 +1317,7 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { struct htp_mm_hvx_vtcm_layout L; htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, ne10, src1_nrows, octx->n_threads, - dst_row_size, src0_row_size, src1_row_size, src2 ? src2->nb[1] : 0, kparams->n_prefetch, false, false, false); + dst_row_size, src0_row_size, src1_row_size, src2 ? src2->nb[1] : 0, kparams->n_prefetch, false, false); if (kparams->kernel_type == HTP_MM_KERNEL_HVX_F16_F16_VTCM || kparams->kernel_type == HTP_MM_KERNEL_HVX_F32_F32_VTCM || @@ -1587,297 +1376,97 @@ static int hvx_mm_matmul(struct htp_ops_context * octx) { return HTP_STATUS_OK; } -static void hvx_mm_qkv_2d(unsigned int nth, unsigned int ith, void * data) { +static void hvx_mm_nx_2d(unsigned int nth, unsigned int ith, void * data) { struct htp_mm_context * mmctx = data; struct htp_ops_context * octx = mmctx->octx; + const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; + const uint32_t n_weights = kparams->n_weights; - const struct htp_tensor * restrict src0 = octx->src[0]; // Wk - const struct htp_tensor * restrict src1 = octx->src[1]; // x - const struct htp_tensor * restrict src2 = octx->src[2]; // Wv - const struct htp_tensor * restrict src3 = octx->src[3]; // Wq - const struct htp_tensor * restrict dst_k = octx->dsts[0]; - const struct htp_tensor * restrict dst_v = octx->dsts[1]; - const struct htp_tensor * restrict dst_q = octx->dsts[2]; - - const uint32_t ne00 = src0->ne[0]; - const uint32_t ne01 = src0->ne[1]; - const uint32_t ne02 = src0->ne[2]; - const uint32_t ne03 = src0->ne[3]; - - const uint32_t ne11 = src1->ne[1]; - const uint32_t ne12 = src1->ne[2]; - const uint32_t ne13 = src1->ne[3]; - - const uint32_t src0_nrows = ne01 * ne02 * ne03; - const uint32_t src1_nrows = ne11 * ne12 * ne13; - - const uint32_t src0_nrows_per_thread = mmctx->src0_nrows_per_thread; - const uint32_t src0_start_row = src0_nrows_per_thread * ith; - const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows); - const uint32_t src0_end_row_x2 = src0_start_row + ((src0_end_row - src0_start_row) & ~1U); - - const size_t dst_k_row_size = dst_k->nb[1]; // K and V share output width - const size_t dst_q_row_size = dst_q->nb[1]; // Q may be wider (GQA) - const size_t src0_row_size = src0->nb[1]; - const size_t src2_row_size = src2->nb[1]; - const size_t src3_row_size = src3->nb[1]; - - const size_t src0_stride = mmctx->vtcm_src0_stride; - const size_t src2_stride = mmctx->vtcm_src2_stride; - const size_t src3_stride = mmctx->vtcm_src3_stride; + const struct htp_tensor * restrict act = octx->src[n_weights]; + const uint32_t src1_nrows = act->ne[1] * act->ne[2] * act->ne[3]; const size_t src1_stride = mmctx->vtcm_src1_stride; uint8_t * restrict vtcm_src0_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; - uint8_t * restrict vtcm_src2_ptr = mmctx->vtcm_src2 + mmctx->vtcm_src2_size_per_thread * ith; - uint8_t * restrict vtcm_src3_ptr = mmctx->vtcm_src3 + mmctx->vtcm_src3_size_per_thread * ith; - uint8_t * restrict src1_data = mmctx->vtcm_src1; + uint8_t * restrict src1_data = mmctx->vtcm_src1; dma_queue * dma_queue = octx->ctx->dma[ith]; - - const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; const uint32_t n_prefetch = kparams->n_prefetch; assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0); const uint32_t prefetch_mask = n_prefetch - 1; - const uint8_t * restrict src0_row = (const uint8_t *) src0->data; - const uint8_t * restrict src2_row = (const uint8_t *) src2->data; - const uint8_t * restrict src3_row = (const uint8_t *) src3->data; - - // Prefill spad with src0, src2, src3 rows - if (src0_start_row < src0_end_row) { - for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) { - const int is0 = (ir0 - src0_start_row); - if (is0 >= (int)n_prefetch) { - break; - } - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size), - src0_stride, src0_row_size, src0_row_size, 2); - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr + is0 * src2_stride, src2_row + ir0 * src2_row_size), - src2_stride, src2_row_size, src2_row_size, 2); - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src3_ptr + is0 * src3_stride, src3_row + ir0 * src3_row_size), - src3_stride, src3_row_size, src3_row_size, 2); - } - } + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; hvx_mm_run_quant_task(mmctx, ith); - if (src0_start_row >= src0_end_row) { - return; - } - - // Process rows - for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) { - const uint8_t * ss0 = dma_queue_pop(dma_queue).dst; - const uint8_t * ss2 = dma_queue_pop(dma_queue).dst; - const uint8_t * ss3 = dma_queue_pop(dma_queue).dst; - - // Process src1 columns in pairs (2×2 tiling) - uint32_t ir1 = 0; - for (; ir1 + 1 < src1_nrows; ir1 += 2) { - const uint8_t * restrict src1_col0 = (const uint8_t *) (src1_data + (ir1+0) * src1_stride); - const uint8_t * restrict src1_col1 = (const uint8_t *) (src1_data + (ir1+1) * src1_stride); - - float * restrict dst_row0_k = (float *) (dst_k->data + ((ir1+0) * dst_k_row_size)); - float * restrict dst_row1_k = (float *) (dst_k->data + ((ir1+1) * dst_k_row_size)); - mmctx->vec_dot_2x2(ne00, &dst_row0_k[ir0], &dst_row1_k[ir0], ss0, ss0 + src0_stride, src1_col0, src1_col1); - - float * restrict dst_row0_v = (float *) (dst_v->data + ((ir1+0) * dst_k_row_size)); - float * restrict dst_row1_v = (float *) (dst_v->data + ((ir1+1) * dst_k_row_size)); - mmctx->vec_dot_2x2(ne00, &dst_row0_v[ir0], &dst_row1_v[ir0], ss2, ss2 + src2_stride, src1_col0, src1_col1); - - float * restrict dst_row0_q = (float *) (dst_q->data + ((ir1+0) * dst_q_row_size)); - float * restrict dst_row1_q = (float *) (dst_q->data + ((ir1+1) * dst_q_row_size)); - mmctx->vec_dot_2x2(ne00, &dst_row0_q[ir0], &dst_row1_q[ir0], ss3, ss3 + src3_stride, src1_col0, src1_col1); - } - - // Handle remaining src1 rows (fallback to 2×1) - for (; ir1 < src1_nrows; ++ir1) { - const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + ir1 * src1_stride); - - float * restrict dst_row_k = (float *) (dst_k->data + (ir1 * dst_k_row_size)); - mmctx->vec_dot_2x1(ne00, &dst_row_k[ir0], ss0, ss0 + src0_stride, src1_col); - - float * restrict dst_row_v = (float *) (dst_v->data + (ir1 * dst_k_row_size)); - mmctx->vec_dot_2x1(ne00, &dst_row_v[ir0], ss2, ss2 + src2_stride, src1_col); - - float * restrict dst_row_q = (float *) (dst_q->data + (ir1 * dst_q_row_size)); - mmctx->vec_dot_2x1(ne00, &dst_row_q[ir0], ss3, ss3 + src3_stride, src1_col); - } - - // Prefetch next (n + vtcm_nrows) rows - const int pr0 = (ir0 + n_prefetch); - const int is0 = (pr0 - src0_start_row) & prefetch_mask; - if (pr0 < src0_end_row_x2) { - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + pr0 * src0_row_size), - src0_stride, src0_row_size, src0_row_size, 2); - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr + is0 * src2_stride, src2_row + pr0 * src2_row_size), - src2_stride, src2_row_size, src2_row_size, 2); - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src3_ptr + is0 * src3_stride, src3_row + pr0 * src3_row_size), - src3_stride, src3_row_size, src3_row_size, 2); - } - } - - // Process last row (if any) - if (src0_end_row != src0_end_row_x2) { - uint32_t ir0 = src0_end_row_x2; - const int is0 = (ir0 - src0_start_row) & prefetch_mask; - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size), - src0_stride, src0_row_size, src0_row_size, 1); - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr + is0 * src2_stride, src2_row + ir0 * src2_row_size), - src2_stride, src2_row_size, src2_row_size, 1); - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src3_ptr + is0 * src3_stride, src3_row + ir0 * src3_row_size), - src3_stride, src3_row_size, src3_row_size, 1); - - const uint8_t * ss0 = dma_queue_pop(dma_queue).dst; - const uint8_t * ss2 = dma_queue_pop(dma_queue).dst; - const uint8_t * ss3 = dma_queue_pop(dma_queue).dst; - - for (uint32_t ir1 = 0; ir1 < src1_nrows; ++ir1) { - const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + ir1 * src1_stride); - - float * restrict dst_row_k = (float *) (dst_k->data + (ir1 * dst_k_row_size)); - mmctx->vec_dot_1x1(ne00, &dst_row_k[ir0], ss0, src1_col); + for (uint32_t widx = 0; widx < n_weights; widx++) { + const struct htp_tensor * restrict src_w = octx->src[widx]; + const struct htp_tensor * restrict dst = octx->dsts[widx]; + if (!src_w || !dst) continue; - float * restrict dst_row_v = (float *) (dst_v->data + (ir1 * dst_k_row_size)); - mmctx->vec_dot_1x1(ne00, &dst_row_v[ir0], ss2, src1_col); + const uint32_t ne00 = src_w->ne[0]; + const uint32_t ne01 = src_w->ne[1]; + const uint32_t src0_nrows = ne01 * src_w->ne[2] * src_w->ne[3]; - float * restrict dst_row_q = (float *) (dst_q->data + (ir1 * dst_q_row_size)); - mmctx->vec_dot_1x1(ne00, &dst_row_q[ir0], ss3, src1_col); - } - } -} - -static void hvx_mm_ffn_2d(unsigned int nth, unsigned int ith, void * data) { - struct htp_mm_context * mmctx = data; - struct htp_ops_context * octx = mmctx->octx; + uint32_t src0_nrows_per_thread = (src0_nrows + nth - 1) / nth; + src0_nrows_per_thread += (src0_nrows_per_thread & 1); - const struct htp_tensor * restrict src0 = octx->src[0]; // Wgate - const struct htp_tensor * restrict src1 = octx->src[1]; // y - const struct htp_tensor * restrict src2 = octx->src[2]; // Wup - const struct htp_tensor * restrict dst_gate = octx->dsts[0]; - const struct htp_tensor * restrict dst_up = octx->dsts[1]; + const uint32_t src0_start_row = src0_nrows_per_thread * ith; + const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows); + const uint32_t src0_end_row_x2 = src0_start_row + ((src0_end_row - src0_start_row) & ~1U); + if (src0_start_row >= src0_end_row) continue; - const uint32_t ne00 = src0->ne[0]; - const uint32_t ne01 = src0->ne[1]; - const uint32_t ne02 = src0->ne[2]; - const uint32_t ne03 = src0->ne[3]; + const size_t dst_row_size = dst->nb[1]; + const size_t src0_row_size = src_w->nb[1]; + const size_t src0_stride = hex_round_up(src0_row_size, 128); - const uint32_t ne11 = src1->ne[1]; - const uint32_t ne12 = src1->ne[2]; - const uint32_t ne13 = src1->ne[3]; + const uint8_t * restrict src0_row = (const uint8_t *) src_w->data; - const uint32_t src0_nrows = ne01 * ne02 * ne03; - const uint32_t src1_nrows = ne11 * ne12 * ne13; - - const uint32_t src0_nrows_per_thread = mmctx->src0_nrows_per_thread; - const uint32_t src0_start_row = src0_nrows_per_thread * ith; - const uint32_t src0_end_row = MIN(src0_start_row + src0_nrows_per_thread, src0_nrows); - const uint32_t src0_end_row_x2 = src0_start_row + ((src0_end_row - src0_start_row) & ~1U); - - const size_t dst_row_size = dst_gate->nb[1]; - const size_t src0_row_size = src0->nb[1]; - const size_t src2_row_size = src2->nb[1]; - - const size_t src0_stride = mmctx->vtcm_src0_stride; - const size_t src2_stride = mmctx->vtcm_src2_stride; - const size_t src1_stride = mmctx->vtcm_src1_stride; - - uint8_t * restrict vtcm_src0_ptr = mmctx->vtcm_src0 + mmctx->vtcm_src0_size_per_thread * ith; - uint8_t * restrict vtcm_src2_ptr = mmctx->vtcm_src2 + mmctx->vtcm_src2_size_per_thread * ith; - uint8_t * restrict src1_data = mmctx->vtcm_src1; - - dma_queue * dma_queue = octx->ctx->dma[ith]; - - const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; - const uint32_t n_prefetch = kparams->n_prefetch; - assert(n_prefetch >= 2 && n_prefetch <= HTP_MM_MAX_PREFETCH && (n_prefetch & (n_prefetch - 1)) == 0); - const uint32_t prefetch_mask = n_prefetch - 1; - - const uint8_t * restrict src0_row = (const uint8_t *) src0->data; - const uint8_t * restrict src2_row = (const uint8_t *) src2->data; - - // Prefill spad with src0, src2 rows - if (src0_start_row < src0_end_row) { for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) { const int is0 = (ir0 - src0_start_row); - if (is0 >= (int)n_prefetch) { - break; - } + if (is0 >= (int)n_prefetch) break; dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size), src0_stride, src0_row_size, src0_row_size, 2); - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr + is0 * src2_stride, src2_row + ir0 * src2_row_size), - src2_stride, src2_row_size, src2_row_size, 2); - } - } - - hvx_mm_run_quant_task(mmctx, ith); - - if (src0_start_row >= src0_end_row) { - return; - } - - // Process rows - for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) { - const uint8_t * ss0 = dma_queue_pop(dma_queue).dst; - const uint8_t * ss2 = dma_queue_pop(dma_queue).dst; - - // Process src1 columns in pairs (2×2 tiling) - uint32_t ir1 = 0; - for (; ir1 + 1 < src1_nrows; ir1 += 2) { - const uint8_t * restrict src1_col0 = (const uint8_t *) (src1_data + (ir1+0) * src1_stride); - const uint8_t * restrict src1_col1 = (const uint8_t *) (src1_data + (ir1+1) * src1_stride); - - float * restrict dst_row0_gate = (float *) (dst_gate->data + ((ir1+0) * dst_row_size)); - float * restrict dst_row1_gate = (float *) (dst_gate->data + ((ir1+1) * dst_row_size)); - mmctx->vec_dot_2x2(ne00, &dst_row0_gate[ir0], &dst_row1_gate[ir0], ss0, ss0 + src0_stride, src1_col0, src1_col1); - - float * restrict dst_row0_up = (float *) (dst_up->data + ((ir1+0) * dst_row_size)); - float * restrict dst_row1_up = (float *) (dst_up->data + ((ir1+1) * dst_row_size)); - mmctx->vec_dot_2x2(ne00, &dst_row0_up[ir0], &dst_row1_up[ir0], ss2, ss2 + src2_stride, src1_col0, src1_col1); } - // Handle remaining src1 rows (fallback to 2×1) - for (; ir1 < src1_nrows; ++ir1) { - const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + ir1 * src1_stride); - - float * restrict dst_row_gate = (float *) (dst_gate->data + (ir1 * dst_row_size)); - mmctx->vec_dot_2x1(ne00, &dst_row_gate[ir0], ss0, ss0 + src0_stride, src1_col); - - float * restrict dst_row_up = (float *) (dst_up->data + (ir1 * dst_row_size)); - mmctx->vec_dot_2x1(ne00, &dst_row_up[ir0], ss2, ss2 + src2_stride, src1_col); - } + for (uint32_t ir0 = src0_start_row; ir0 < src0_end_row_x2; ir0 += 2) { + const uint8_t * ss0 = dma_queue_pop(dma_queue).dst; + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir0); + uint32_t ir1 = 0; + for (; ir1 + 1 < src1_nrows; ir1 += 2) { + const uint8_t * restrict src1_col0 = (const uint8_t *) (src1_data + (ir1+0) * src1_stride); + const uint8_t * restrict src1_col1 = (const uint8_t *) (src1_data + (ir1+1) * src1_stride); + float * restrict dst_row0 = (float *) (dst->data + ((ir1+0) * dst_row_size)); + float * restrict dst_row1 = (float *) (dst->data + ((ir1+1) * dst_row_size)); + mmctx->vec_dot_2x2(ne00, &dst_row0[ir0], &dst_row1[ir0], ss0, ss0 + src0_stride, src1_col0, src1_col1); + } + for (; ir1 < src1_nrows; ++ir1) { + const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + ir1 * src1_stride); + float * restrict dst_row = (float *) (dst->data + (ir1 * dst_row_size)); + mmctx->vec_dot_2x1(ne00, &dst_row[ir0], ss0, ss0 + src0_stride, src1_col); + } + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir0); - // Prefetch next rows - const int pr0 = (ir0 + n_prefetch); - const int is0 = (pr0 - src0_start_row) & prefetch_mask; - if (pr0 < src0_end_row_x2) { - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + pr0 * src0_row_size), - src0_stride, src0_row_size, src0_row_size, 2); - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr + is0 * src2_stride, src2_row + pr0 * src2_row_size), - src2_stride, src2_row_size, src2_row_size, 2); + const int pr0 = (ir0 + n_prefetch); + const int is0 = (pr0 - src0_start_row) & prefetch_mask; + if (pr0 < src0_end_row_x2) { + dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + pr0 * src0_row_size), + src0_stride, src0_row_size, src0_row_size, 2); + } } - } - - // Process last row (if any) - if (src0_end_row != src0_end_row_x2) { - uint32_t ir0 = src0_end_row_x2; - const int is0 = (ir0 - src0_start_row) & prefetch_mask; - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size), - src0_stride, src0_row_size, src0_row_size, 1); - dma_queue_push(dma_queue, dma_make_ptr(vtcm_src2_ptr + is0 * src2_stride, src2_row + ir0 * src2_row_size), - src2_stride, src2_row_size, src2_row_size, 1); - const uint8_t * ss0 = dma_queue_pop(dma_queue).dst; - const uint8_t * ss2 = dma_queue_pop(dma_queue).dst; - - for (uint32_t ir1 = 0; ir1 < src1_nrows; ++ir1) { - const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + ir1 * src1_stride); - - float * restrict dst_row_gate = (float *) (dst_gate->data + (ir1 * dst_row_size)); - mmctx->vec_dot_1x1(ne00, &dst_row_gate[ir0], ss0, src1_col); - - float * restrict dst_row_up = (float *) (dst_up->data + (ir1 * dst_row_size)); - mmctx->vec_dot_1x1(ne00, &dst_row_up[ir0], ss2, src1_col); + if (src0_end_row != src0_end_row_x2) { + uint32_t ir0 = src0_end_row_x2; + const int is0 = (ir0 - src0_start_row) & prefetch_mask; + dma_queue_push(dma_queue, dma_make_ptr(vtcm_src0_ptr + is0 * src0_stride, src0_row + ir0 * src0_row_size), + src0_stride, src0_row_size, src0_row_size, 1); + const uint8_t * ss0 = dma_queue_pop(dma_queue).dst; + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, ir0); + for (uint32_t ir1 = 0; ir1 < src1_nrows; ++ir1) { + const uint8_t * restrict src1_col = (const uint8_t *) (src1_data + ir1 * src1_stride); + float * restrict dst_row = (float *) (dst->data + (ir1 * dst_row_size)); + mmctx->vec_dot_1x1(ne00, &dst_row[ir0], ss0, src1_col); + } + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, ir0); } } } @@ -3485,7 +3074,7 @@ static int hvx_mm_matmul_id( struct htp_mm_hvx_vtcm_layout L; htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, ne10, src1_nrows, octx->n_threads, - 0, src0_row_size, src1_row_size, 0, kparams->n_prefetch, true, false, false); + 0, src0_row_size, src1_row_size, 0, kparams->n_prefetch, true, false); size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes; @@ -3610,6 +3199,7 @@ int op_matmul_id(struct htp_ops_context * octx) { struct htp_mm_context mmctx_struct = {0}; struct htp_mm_context * mmctx = &mmctx_struct; mmctx->octx = octx; + mmctx->act = src1; const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; @@ -3690,18 +3280,15 @@ int op_matmul_id(struct htp_ops_context * octx) { return s; } - -int op_matmul_qkv(struct htp_ops_context * octx) { +int op_matmul_nx(struct htp_ops_context * octx) { struct htp_thread_trace * tr = &octx->ctx->trace[0]; htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0); - const struct htp_tensor * restrict src0 = octx->src[0]; // Wk - const struct htp_tensor * restrict src1 = octx->src[1]; // x - const struct htp_tensor * restrict src2 = octx->src[2]; // Wv - const struct htp_tensor * restrict src3 = octx->src[3]; // Wq - const struct htp_tensor * restrict dst_k = octx->dsts[0]; - const struct htp_tensor * restrict dst_v = octx->dsts[1]; - const struct htp_tensor * restrict dst_q = octx->dsts[2]; + const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; + const uint32_t n_weights = kparams->n_weights; + + const struct htp_tensor * restrict src0 = octx->src[0]; // first weight + const struct htp_tensor * restrict act = octx->src[n_weights]; // activation x bool is_repacked = (src0->type == HTP_TYPE_Q4_0 || src0->type == HTP_TYPE_Q4_1 || src0->type == HTP_TYPE_Q8_0 || src0->type == HTP_TYPE_IQ4_NL || @@ -3710,19 +3297,9 @@ int op_matmul_qkv(struct htp_ops_context * octx) { struct htp_mm_context mmctx_struct = {0}; struct htp_mm_context * mmctx = &mmctx_struct; mmctx->octx = octx; + mmctx->act = act; - const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; - - const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3]; - const uint32_t src1_nrows = src1->ne[1] * src1->ne[2] * src1->ne[3]; - - // Compute src0_nrows_per_thread - mmctx->src0_nrows_per_thread = (src0_nrows + octx->n_threads - 1) / octx->n_threads; - if (is_repacked) { - mmctx->src0_nrows_per_thread = hex_round_up(mmctx->src0_nrows_per_thread, 32); - } else { - mmctx->src0_nrows_per_thread += (mmctx->src0_nrows_per_thread & 1); // round up to even - } + const uint32_t src1_nrows = act->ne[1] * act->ne[2] * act->ne[3]; const size_t src0_row_size = src0->nb[1]; const size_t src0_row_size_padded = hex_round_up(src0_row_size, 128); @@ -3732,7 +3309,7 @@ int op_matmul_qkv(struct htp_ops_context * octx) { } const uint32_t qk = QK_Q8_0_TILED; - const uint32_t nb = (src1->ne[0] + qk - 1) / qk; + const uint32_t nb = (act->ne[0] + qk - 1) / qk; const uint32_t total_nb = src1_nrows * nb; worker_callback_t quant_task_func; @@ -3758,185 +3335,39 @@ int op_matmul_qkv(struct htp_ops_context * octx) { size_t src1_row_size; if (kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT) { - src1_row_size = (src0->type == HTP_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(src1->ne[0]) : htp_mm_q8_0_flat_row_size(src1->ne[0]); + src1_row_size = (src0->type == HTP_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(act->ne[0]) : htp_mm_q8_0_flat_row_size(act->ne[0]); } else { - src1_row_size = (src0->type == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(src1->ne[0]) : htp_mm_q8_0_tiled_row_size(src1->ne[0]); + src1_row_size = (src0->type == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(act->ne[0]) : htp_mm_q8_0_tiled_row_size(act->ne[0]); } struct htp_mm_hvx_vtcm_layout L; - htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, src1->ne[0], src1_nrows, octx->n_threads, - 0, src0_row_size, src1_row_size, 0, kparams->n_prefetch, false, true, false); + htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, act->ne[0], src1_nrows, octx->n_threads, + 0, src0_row_size, src1_row_size, 0, kparams->n_prefetch, false, true); size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes; if (octx->ctx->vtcm_size < vtcm_size) { - FARF(ERROR, "matmul-qkv: current VTCM reservation %zu is too small, needed %zu\n", + FARF(ERROR, "matmul-nx: current VTCM reservation %zu is too small, needed %zu\n", octx->ctx->vtcm_size, vtcm_size); return HTP_STATUS_VTCM_TOO_SMALL; } uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base; - mmctx->vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1); mmctx->vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0); - mmctx->vtcm_src2 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src2); - mmctx->vtcm_src3 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src3); - mmctx->vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst); - - octx->src1_spad.src = NULL; - octx->src0_spad.src = NULL; - octx->src2_spad.src = NULL; - octx->src3_spad.src = NULL; - octx->dst_spad.src = NULL; - - mmctx->vtcm_src0_stride = is_repacked ? 0 : src0_row_size_padded; - mmctx->vtcm_src2_stride = is_repacked ? 0 : src0_row_size_padded; - mmctx->vtcm_src3_stride = is_repacked ? 0 : src0_row_size_padded; - mmctx->vtcm_src1_stride = src1_row_size; - - mmctx->vtcm_src0_size_per_thread = L.src0_bytes / octx->n_threads; - mmctx->vtcm_src1_size_per_thread = L.src1_bytes; - mmctx->vtcm_src2_size_per_thread = L.src2_bytes / octx->n_threads; - mmctx->vtcm_src3_size_per_thread = L.src3_bytes / octx->n_threads; - mmctx->vtcm_dst_size_per_thread = L.dst_bytes / octx->n_threads; - - mmctx->n_quant_rows_per_thread = (src1_nrows + n_quant_tasks - 1) / n_quant_tasks; - mmctx->quant_task_func = quant_task_func; - mmctx->n_quant_tasks = n_quant_tasks; - atomic_init(&mmctx->quant_barrier, n_quant_tasks); - - // Run fused matmul - const uint32_t n_matmul_jobs = octx->n_threads; - worker_callback_t matmul_job_func; - if (is_repacked) { - if (kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT) { - switch (src0->type) { - case HTP_TYPE_Q4_0: matmul_job_func = hvx_mm_qkv_2d_repacked_q4_0_flat; break; - case HTP_TYPE_Q4_1: matmul_job_func = hvx_mm_qkv_2d_repacked_q4_1_flat; break; - case HTP_TYPE_Q8_0: matmul_job_func = hvx_mm_qkv_2d_repacked_q8_0_flat; break; - case HTP_TYPE_IQ4_NL: matmul_job_func = hvx_mm_qkv_2d_repacked_iq4nl_flat; break; - case HTP_TYPE_MXFP4: matmul_job_func = hvx_mm_qkv_2d_repacked_mxfp4_flat; break; - default: return HTP_STATUS_NO_SUPPORT; - } - } else { - switch (src0->type) { - case HTP_TYPE_Q4_0: matmul_job_func = hvx_mm_qkv_2d_repacked_q4_0; break; - case HTP_TYPE_Q4_1: matmul_job_func = hvx_mm_qkv_2d_repacked_q4_1; break; - case HTP_TYPE_Q8_0: matmul_job_func = hvx_mm_qkv_2d_repacked_q8_0; break; - case HTP_TYPE_IQ4_NL: matmul_job_func = hvx_mm_qkv_2d_repacked_iq4nl; break; - case HTP_TYPE_MXFP4: matmul_job_func = hvx_mm_qkv_2d_repacked_mxfp4; break; - default: return HTP_STATUS_NO_SUPPORT; - } - } - } else { - matmul_job_func = hvx_mm_qkv_2d; - } - - htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0); - - worker_pool_run_func(octx->ctx->worker_pool, matmul_job_func, mmctx, n_matmul_jobs); - - return HTP_STATUS_OK; -} - -int op_matmul_ffn(struct htp_ops_context * octx) { - struct htp_thread_trace * tr = &octx->ctx->trace[0]; - htp_trace_event_start(tr, HTP_TRACE_EVT_INIT, 0); - - const struct htp_tensor * restrict src0 = octx->src[0]; // Wgate - const struct htp_tensor * restrict src1 = octx->src[1]; // y - const struct htp_tensor * restrict src2 = octx->src[2]; // Wup - const struct htp_tensor * restrict dst_gate = octx->dsts[0]; - const struct htp_tensor * restrict dst_up = octx->dsts[1]; - - bool is_repacked = (src0->type == HTP_TYPE_Q4_0 || src0->type == HTP_TYPE_Q4_1 || - src0->type == HTP_TYPE_Q8_0 || src0->type == HTP_TYPE_IQ4_NL || - src0->type == HTP_TYPE_MXFP4); - - struct htp_mm_context mmctx_struct = {0}; - struct htp_mm_context * mmctx = &mmctx_struct; - mmctx->octx = octx; - - const struct htp_mm_kernel_params * kparams = (const struct htp_mm_kernel_params *) octx->kernel_params; - - const uint32_t src0_nrows = src0->ne[1] * src0->ne[2] * src0->ne[3]; - const uint32_t src1_nrows = src1->ne[1] * src1->ne[2] * src1->ne[3]; - - // Compute src0_nrows_per_thread - mmctx->src0_nrows_per_thread = (src0_nrows + octx->n_threads - 1) / octx->n_threads; - if (is_repacked) { - mmctx->src0_nrows_per_thread = hex_round_up(mmctx->src0_nrows_per_thread, 32); - } else { - mmctx->src0_nrows_per_thread += (mmctx->src0_nrows_per_thread & 1); // round up to even - } - - const size_t src0_row_size = src0->nb[1]; - const size_t src0_row_size_padded = hex_round_up(src0_row_size, 128); - - if (hvx_mm_init_vec_dot(mmctx, src0->type) != 0) { - return HTP_STATUS_NO_SUPPORT; - } - - const uint32_t qk = QK_Q8_0_TILED; - const uint32_t nb = (src1->ne[0] + qk - 1) / qk; - const uint32_t total_nb = src1_nrows * nb; - - worker_callback_t quant_task_func; - uint32_t n_quant_tasks = 1; - if (kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT) { - n_quant_tasks = MIN(src1_nrows, octx->n_threads); - quant_task_func = (src0->type == HTP_TYPE_Q4_1) ? quantize_f32_q8_1_flat : quantize_f32_q8_0_flat; - } else if (src1_nrows < octx->n_threads) { - n_quant_tasks = MIN(total_nb, octx->n_threads); - quant_task_func = (src0->type == HTP_TYPE_Q4_1) ? quantize_f32_q8_1_tiled_block : quantize_f32_q8_0_tiled_block; - for (uint32_t ith = 0; ith < n_quant_tasks; ++ith) { - uint32_t ib_first = (total_nb * (ith + 0)) / n_quant_tasks; - uint32_t ib_last = (total_nb * (ith + 1)) / n_quant_tasks; - mmctx->quant_ib_first[ith] = ib_first; - mmctx->quant_ib_last[ith] = ib_last; - mmctx->quant_r[ith] = ib_first / nb; - mmctx->quant_c[ith] = ib_first % nb; - } - } else { - n_quant_tasks = MIN(src1_nrows, octx->n_threads); - quant_task_func = (src0->type == HTP_TYPE_Q4_1) ? quantize_f32_q8_1_tiled : quantize_f32_q8_0_tiled; - } - - size_t src1_row_size; - if (kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT) { - src1_row_size = (src0->type == HTP_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(src1->ne[0]) : htp_mm_q8_0_flat_row_size(src1->ne[0]); - } else { - src1_row_size = (src0->type == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(src1->ne[0]) : htp_mm_q8_0_tiled_row_size(src1->ne[0]); - } - - struct htp_mm_hvx_vtcm_layout L; - htp_mm_hvx_vtcm_layout_build(&L, kparams->kernel_type, src0->type, src1->ne[0], src1_nrows, octx->n_threads, - 0, src0_row_size, src1_row_size, 0, kparams->n_prefetch, false, false, true); - - size_t vtcm_size = kparams->vtcm_size > 0 ? (size_t)kparams->vtcm_size : L.total_bytes; - - if (octx->ctx->vtcm_size < vtcm_size) { - FARF(ERROR, "matmul-ffn: current VTCM reservation %zu is too small, needed %zu\n", octx->ctx->vtcm_size, vtcm_size); - return HTP_STATUS_VTCM_TOO_SMALL; - } - - uint8_t * const base = (uint8_t *) octx->ctx->vtcm_base; mmctx->vtcm_src1 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src1); - mmctx->vtcm_src0 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src0); - mmctx->vtcm_src2 = VTCM_LAYOUT_PTR(uint8_t, base, L.off_src2); mmctx->vtcm_dst = VTCM_LAYOUT_PTR(uint8_t, base, L.off_dst); - octx->src1_spad.src = NULL; octx->src0_spad.src = NULL; + octx->src1_spad.src = NULL; octx->src2_spad.src = NULL; + octx->src3_spad.src = NULL; octx->dst_spad.src = NULL; mmctx->vtcm_src0_stride = is_repacked ? 0 : src0_row_size_padded; - mmctx->vtcm_src2_stride = is_repacked ? 0 : src0_row_size_padded; mmctx->vtcm_src1_stride = src1_row_size; mmctx->vtcm_src0_size_per_thread = L.src0_bytes / octx->n_threads; mmctx->vtcm_src1_size_per_thread = L.src1_bytes; - mmctx->vtcm_src2_size_per_thread = L.src2_bytes / octx->n_threads; mmctx->vtcm_dst_size_per_thread = L.dst_bytes / octx->n_threads; mmctx->n_quant_rows_per_thread = (src1_nrows + n_quant_tasks - 1) / n_quant_tasks; @@ -3950,25 +3381,25 @@ int op_matmul_ffn(struct htp_ops_context * octx) { if (is_repacked) { if (kparams->kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT) { switch (src0->type) { - case HTP_TYPE_Q4_0: matmul_job_func = hvx_mm_ffn_2d_repacked_q4_0_flat; break; - case HTP_TYPE_Q4_1: matmul_job_func = hvx_mm_ffn_2d_repacked_q4_1_flat; break; - case HTP_TYPE_Q8_0: matmul_job_func = hvx_mm_ffn_2d_repacked_q8_0_flat; break; - case HTP_TYPE_IQ4_NL: matmul_job_func = hvx_mm_ffn_2d_repacked_iq4nl_flat; break; - case HTP_TYPE_MXFP4: matmul_job_func = hvx_mm_ffn_2d_repacked_mxfp4_flat; break; + case HTP_TYPE_Q4_0: matmul_job_func = hvx_mm_nx_2d_repacked_q4_0_flat; break; + case HTP_TYPE_Q4_1: matmul_job_func = hvx_mm_nx_2d_repacked_q4_1_flat; break; + case HTP_TYPE_Q8_0: matmul_job_func = hvx_mm_nx_2d_repacked_q8_0_flat; break; + case HTP_TYPE_IQ4_NL: matmul_job_func = hvx_mm_nx_2d_repacked_iq4nl_flat; break; + case HTP_TYPE_MXFP4: matmul_job_func = hvx_mm_nx_2d_repacked_mxfp4_flat; break; default: return HTP_STATUS_NO_SUPPORT; } } else { switch (src0->type) { - case HTP_TYPE_Q4_0: matmul_job_func = hvx_mm_ffn_2d_repacked_q4_0; break; - case HTP_TYPE_Q4_1: matmul_job_func = hvx_mm_ffn_2d_repacked_q4_1; break; - case HTP_TYPE_Q8_0: matmul_job_func = hvx_mm_ffn_2d_repacked_q8_0; break; - case HTP_TYPE_IQ4_NL: matmul_job_func = hvx_mm_ffn_2d_repacked_iq4nl; break; - case HTP_TYPE_MXFP4: matmul_job_func = hvx_mm_ffn_2d_repacked_mxfp4; break; + case HTP_TYPE_Q4_0: matmul_job_func = hvx_mm_nx_2d_repacked_q4_0; break; + case HTP_TYPE_Q4_1: matmul_job_func = hvx_mm_nx_2d_repacked_q4_1; break; + case HTP_TYPE_Q8_0: matmul_job_func = hvx_mm_nx_2d_repacked_q8_0; break; + case HTP_TYPE_IQ4_NL: matmul_job_func = hvx_mm_nx_2d_repacked_iq4nl; break; + case HTP_TYPE_MXFP4: matmul_job_func = hvx_mm_nx_2d_repacked_mxfp4; break; default: return HTP_STATUS_NO_SUPPORT; } } } else { - matmul_job_func = hvx_mm_ffn_2d; + matmul_job_func = hvx_mm_nx_2d; } htp_trace_event_stop(tr, HTP_TRACE_EVT_INIT, 0); diff --git a/ggml/src/ggml-hexagon/htp/matmul-ops.h b/ggml/src/ggml-hexagon/htp/matmul-ops.h index 6c393664c6e8..dbc8e359093d 100644 --- a/ggml/src/ggml-hexagon/htp/matmul-ops.h +++ b/ggml/src/ggml-hexagon/htp/matmul-ops.h @@ -88,6 +88,7 @@ struct htp_mm_kernel_params { int32_t vtcm_src2_size; // src2 scratchpad size in VTCM (fused only) int32_t vtcm_src3_size; // src3 scratchpad size in VTCM (fused only) int32_t vtcm_dst_size; // dst scratchpad size in VTCM + int32_t n_weights; // Number of weights for fused NX // Precomputed division values struct fastdiv_values div_ne12_ne1; @@ -463,8 +464,7 @@ static inline void htp_mm_hvx_vtcm_layout_build( size_t src2_row_size, uint32_t n_prefetch, bool is_matmul_id, - bool is_fused_qkv, - bool is_fused_ffn + bool is_fused_nx ) { size_t src0_sz = 0; size_t src1_sz = 0; @@ -476,44 +476,33 @@ static inline void htp_mm_hvx_vtcm_layout_build( wtype == HTP_TYPE_Q8_0 || wtype == HTP_TYPE_IQ4_NL || wtype == HTP_TYPE_MXFP4); - if (is_fused_qkv || is_fused_ffn) { + if (is_fused_nx) { const size_t src0_row_size_padded = hex_round_up(src0_row_size, 128); const size_t quant_scratch_size = hex_round_up(ne10 * sizeof(float), QK_Q8_0_TILED * sizeof(float)) * n_threads; - size_t src0_sz_per_thread = 0; - size_t src2_sz_per_thread = 0; - size_t src3_sz_per_thread = 0; + size_t weight_sz_per_thread = 0; if (is_repack) { uint32_t aligned_tile_size = htp_mm_get_weight_aligned_tile_size(wtype); uint32_t n_k_tiles = hex_round_up(ne10, 32) / 32; uint32_t tile_row_size = n_k_tiles * aligned_tile_size; - src0_sz_per_thread = hex_round_up(n_prefetch * tile_row_size, 128); - src2_sz_per_thread = hex_round_up(n_prefetch * tile_row_size, 128); - if (is_fused_qkv) { - src3_sz_per_thread = hex_round_up(n_prefetch * tile_row_size, 128); - } + weight_sz_per_thread = hex_round_up(n_prefetch * tile_row_size, 128); } else { - src0_sz_per_thread = hex_round_up(n_prefetch * src0_row_size_padded, 128); - src2_sz_per_thread = hex_round_up(n_prefetch * src0_row_size_padded, 128); - if (is_fused_qkv) { - src3_sz_per_thread = hex_round_up(n_prefetch * src0_row_size_padded, 128); - } + weight_sz_per_thread = hex_round_up(n_prefetch * src0_row_size_padded, 128); } - size_t flat_src1_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10); - size_t tiled_src1_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10); + size_t flat_act_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_flat_row_size(ne10) : htp_mm_q8_0_flat_row_size(ne10); + size_t tiled_act_row_size = (wtype == HTP_TYPE_Q4_1) ? htp_mm_q8_1_tiled_row_size(ne10) : htp_mm_q8_0_tiled_row_size(ne10); - if (kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT) { - src1_sz = hex_round_up(flat_src1_row_size * src1_nrows, 128); - } else { - src1_sz = hex_round_up(tiled_src1_row_size * src1_nrows, 128); - } + size_t act_sz = (kernel_type == HTP_MM_KERNEL_HVX_QUANT_ROW_FLAT) + ? hex_round_up(flat_act_row_size * src1_nrows, 128) + : hex_round_up(tiled_act_row_size * src1_nrows, 128); - src0_sz = src0_sz_per_thread * n_threads; - src2_sz = src2_sz_per_thread * n_threads; - src3_sz = src3_sz_per_thread * n_threads; + src0_sz = weight_sz_per_thread * n_threads; // shared single-weight prefetch buffer + src1_sz = act_sz; // quantized activation buffer + src2_sz = 0; + src3_sz = 0; dst_sz = quant_scratch_size; } else if (is_matmul_id) { const size_t src0_row_size_padded = htp_mm_round_up(src0_row_size, 128); @@ -616,8 +605,8 @@ static inline void htp_mm_hvx_vtcm_layout_build( } size_t off = 0; - VTCM_LAYOUT_ALLOC(off, off_src1, src1_sz); VTCM_LAYOUT_ALLOC(off, off_src0, src0_sz); + VTCM_LAYOUT_ALLOC(off, off_src1, src1_sz); VTCM_LAYOUT_ALLOC(off, off_src2, src2_sz); VTCM_LAYOUT_ALLOC(off, off_src3, src3_sz); VTCM_LAYOUT_ALLOC(off, off_dst, dst_sz); diff --git a/ggml/src/ggml-hexagon/htp/set-rows-ops.c b/ggml/src/ggml-hexagon/htp/set-rows-ops.c index 58c54967db09..fa14bf0ef6b6 100644 --- a/ggml/src/ggml-hexagon/htp/set-rows-ops.c +++ b/ggml/src/ggml-hexagon/htp/set-rows-ops.c @@ -8,14 +8,20 @@ #include #include -#include "hex-dma.h" +#include "dma-queue.h" +#include "work-queue.h" #include "hvx-utils.h" +#include "hex-utils.h" +#include "hvx-copy.h" +#include "hvx-quant.h" #define GGML_COMMON_DECL_C #include "ggml-common.h" + #include "htp-ctx.h" #include "htp-ops.h" -#include "htp-ops.h" +#include "htp-tensor.h" +#include "htp/set-rows-ops.h" #define set_rows_preamble \ const uint32_t ne00 = octx->src[0]->ne[0]; \ @@ -47,116 +53,142 @@ \ const uint32_t nr = ne01; -struct htp_set_rows_context { +struct set_rows_context { struct htp_ops_context * octx; - struct fastdiv_values div_ne12; - struct fastdiv_values div_ne11; - uint32_t src0_nrows_per_thread; + const struct htp_set_rows_kernel_params * kparams; + struct htp_set_rows_vtcm_layout vtcm_layout; + uint8_t * vtcm_base; }; -static void set_rows_thread_f32_f32(unsigned int nth, unsigned int ith, void *data) { - struct htp_set_rows_context * srctx = (struct htp_set_rows_context *)data; - struct htp_ops_context * octx = srctx->octx; - - set_rows_preamble; - - uint64_t qt = HAP_perf_get_qtimer_count(); - - // parallelize by rows of src0 - const uint32_t dr = srctx->src0_nrows_per_thread; - const uint32_t ir0 = dr * ith; - if (ir0 >= nr) { - return; - } - const uint32_t ir1 = (ir0 + dr < nr) ? (ir0 + dr) : nr; - - const bool is_i32 = (octx->src[1]->type == HTP_TYPE_I32); - - for (uint32_t i03 = 0; i03 < ne03; ++i03) { - for (uint32_t i02 = 0; i02 < ne02; ++i02) { - for (uint32_t i = ir0; i < ir1; ++i) { - const uint32_t i12 = fastmodulo(i03, ne12, &srctx->div_ne12); - const uint32_t i11 = fastmodulo(i02, ne11, &srctx->div_ne11); - const uint32_t i10 = i; - - const uintptr_t src1_addr = octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12; - - uint32_t i1 = is_i32 ? *(int32_t *)src1_addr : *(int64_t *)src1_addr; - if (i1 >= ne1) { - // ignore invalid indices - continue; - } - - const uintptr_t src0_ptr = octx->src[0]->data + i*nb01 + i02*nb02 + i03*nb03; - const uintptr_t dst_ptr = octx->dst->data + i1*nb1 + i02*nb2 + i03*nb3; - - // copy row - hvx_copy_f32_uu((uint8_t *)dst_ptr, (const uint8_t *)src0_ptr, ne00); - } - } - } - - qt = HAP_perf_qtimer_count_to_us(HAP_perf_get_qtimer_count() - qt); - FARF(HIGH, "set-rows-f32-f32 %d/%d: %ux%ux%ux%u (%u:%u) x %ux%ux%ux%u -> %ux%ux%ux%u usec %u\n", ith, nth, - ne00, ne01, ne02, ne03, ir0, ir1, ne10, ne11, ne12, ne13, ne0, ne1, ne2, ne3, (unsigned) qt); +#define SET_ROWS_THREAD_DMA_FN(TYPE_NAME, IDX_TYPE, COMPUTE_EXPR) \ +static void set_rows_thread_dma_##TYPE_NAME##_##IDX_TYPE(unsigned int nth, unsigned int ith, void *data) { \ + struct set_rows_context * srctx = (struct set_rows_context *)data; \ + struct htp_ops_context * octx = srctx->octx; \ + const struct htp_set_rows_kernel_params * kparams = srctx->kparams; \ + set_rows_preamble; \ + struct htp_thread_trace * tr = &octx->ctx->trace[ith]; \ + const uint32_t dr = kparams->tasks_per_thread; \ + const uint32_t ir0 = dr * ith; \ + if (ir0 >= kparams->total_tasks) { \ + return; \ + } \ + const uint32_t ir1 = MIN(ir0 + dr, kparams->total_tasks); \ + dma_queue * dma_queue = octx->ctx->dma[ith]; \ + const struct htp_set_rows_vtcm_layout * vtcm_layout = &srctx->vtcm_layout; \ + uint8_t * vtcm_src0 = srctx->vtcm_base + vtcm_layout->off_src0 + ith * vtcm_layout->src0_bytes_per_thread; \ + uint8_t * vtcm_dst = srctx->vtcm_base + vtcm_layout->off_dst + ith * vtcm_layout->dst_bytes_per_thread; \ + const uint32_t src0_row_size = ne00 * sizeof(float); \ + const uint32_t dst_row_size = htp_tensor_get_row_size(octx->dst->type, ne00); \ + const uint32_t nrows_per_thread = ir1 - ir0; \ + const uint32_t total_steps = ne03 * ne02 * nrows_per_thread; \ + uint32_t pi_step = 0; \ + uint32_t pi02 = 0; \ + uint32_t pi03 = 0; \ + for (uint32_t step = 0, spad_idx = 0; step < total_steps && spad_idx < 2; ++step, spad_idx++) { \ + uint32_t i = ir0 + pi_step; \ + const uintptr_t src0_ptr = octx->src[0]->data + i*nb01 + pi02*nb02 + pi03*nb03; \ + dma_queue_push(dma_queue, \ + dma_make_ptr((void *)octx->dst->data, \ + vtcm_dst + spad_idx * vtcm_layout->dst_spad_half_size), \ + dst_row_size, vtcm_layout->dst_spad_half_size, dst_row_size, 0); \ + dma_queue_push(dma_queue, \ + dma_make_ptr((void *)(vtcm_src0 + spad_idx * vtcm_layout->src0_spad_half_size), \ + (const void *)src0_ptr), \ + vtcm_layout->src0_spad_half_size, src0_row_size, src0_row_size, 1); \ + pi_step++; \ + if (pi_step == nrows_per_thread) { \ + pi_step = 0; \ + pi02++; \ + if (pi02 == ne02) { \ + pi02 = 0; \ + pi03++; \ + } \ + } \ + } \ + uint32_t ci_step = 0; \ + uint32_t ci02 = 0; \ + uint32_t ci03 = 0; \ + uint32_t ci11_base = 0; \ + uint32_t ci12_base = 0; \ + for (uint32_t step = 0; step < total_steps; ++step) { \ + void * dst_spad = (void *) dma_queue_pop(dma_queue).src; \ + void * src_spad = (void *) dma_queue_pop(dma_queue).dst; \ + uint32_t i = ir0 + ci_step; \ + const uintptr_t src1_addr = octx->src[1]->data + i*nb10 + ci11_base*nb11 + ci12_base*nb12; \ + const IDX_TYPE i1 = *(const IDX_TYPE *)src1_addr; \ + const bool valid_i1 = ((uint64_t)i1 < (uint64_t)ne1); \ + const uint32_t target_i1 = (uint32_t)i1; \ + htp_trace_event_start(tr, HTP_TRACE_EVT_HVX_COMP, step); \ + if (valid_i1) { \ + COMPUTE_EXPR; \ + } \ + htp_trace_event_stop(tr, HTP_TRACE_EVT_HVX_COMP, step); \ + if (valid_i1) { \ + const uintptr_t dst_ptr = octx->dst->data + target_i1*nb1 + ci02*nb2 + ci03*nb3; \ + dma_queue_push(dma_queue, \ + dma_make_ptr((void *)dst_ptr, (const void *)dst_spad), \ + dst_row_size, vtcm_layout->dst_spad_half_size, dst_row_size, 1); \ + } else { \ + dma_queue_push(dma_queue, \ + dma_make_ptr((void *)octx->dst->data, (const void *)dst_spad), \ + dst_row_size, vtcm_layout->dst_spad_half_size, dst_row_size, 0); \ + } \ + const uint32_t next_step = step + 2; \ + if (next_step < total_steps) { \ + uint32_t ni = ir0 + pi_step; \ + const uintptr_t psrc0_ptr = octx->src[0]->data + ni*nb01 + pi02*nb02 + pi03*nb03; \ + dma_queue_push(dma_queue, \ + dma_make_ptr((void *)src_spad, (const void *)psrc0_ptr), \ + vtcm_layout->src0_spad_half_size, src0_row_size, src0_row_size, 1); \ + pi_step++; \ + if (pi_step == nrows_per_thread) { \ + pi_step = 0; \ + pi02++; \ + if (pi02 == ne02) { \ + pi02 = 0; \ + pi03++; \ + } \ + } \ + } \ + ci_step++; \ + if (ci_step == nrows_per_thread) { \ + ci_step = 0; \ + ci02++; \ + ci11_base++; \ + if (ci11_base == ne11) { \ + ci11_base = 0; \ + } \ + if (ci02 == ne02) { \ + ci02 = 0; \ + ci03++; \ + ci12_base++; \ + if (ci12_base == ne12) { \ + ci12_base = 0; \ + } \ + } \ + } \ + } \ + dma_queue_flush(dma_queue); \ } -static void set_rows_thread_f16_f32(unsigned int nth, unsigned int ith, void *data) { - struct htp_set_rows_context * srctx = (struct htp_set_rows_context *)data; - struct htp_ops_context * octx = srctx->octx; - - set_rows_preamble; +SET_ROWS_THREAD_DMA_FN(f32, int32_t, { hvx_copy_f32_uu((uint8_t *)dst_spad, (const uint8_t *)src_spad, ne00); }) +SET_ROWS_THREAD_DMA_FN(f32, int64_t, { hvx_copy_f32_uu((uint8_t *)dst_spad, (const uint8_t *)src_spad, ne00); }) - uint64_t qt = HAP_perf_get_qtimer_count(); +SET_ROWS_THREAD_DMA_FN(f16, int32_t, { hvx_copy_f16_f32_uu((uint8_t *)dst_spad, (const uint8_t *)src_spad, ne00); }) +SET_ROWS_THREAD_DMA_FN(f16, int64_t, { hvx_copy_f16_f32_uu((uint8_t *)dst_spad, (const uint8_t *)src_spad, ne00); }) - // parallelize by rows of src0 - const uint32_t dr = srctx->src0_nrows_per_thread; - const uint32_t ir0 = dr * ith; - if (ir0 >= nr) { - return; - } - const uint32_t ir1 = (ir0 + dr < nr) ? (ir0 + dr) : nr; - - const bool is_i32 = (octx->src[1]->type == HTP_TYPE_I32); - - for (uint32_t i03 = 0; i03 < ne03; ++i03) { - for (uint32_t i02 = 0; i02 < ne02; ++i02) { - for (uint32_t i = ir0; i < ir1; ++i) { - const uint32_t i12 = fastmodulo(i03, ne12, &srctx->div_ne12); - const uint32_t i11 = fastmodulo(i02, ne11, &srctx->div_ne11); - const uint32_t i10 = i; - - const uintptr_t src1_addr = octx->src[1]->data + i10*nb10 + i11*nb11 + i12*nb12; - - uint32_t i1 = is_i32 ? *(int32_t *)src1_addr : *(int64_t *)src1_addr; - if (i1 >= ne1) { - // ignore invalid indices - continue; - } - - const uint8_t* src0_ptr = (const uint8_t *) octx->src[0]->data + i*nb01 + i02*nb02 + i03*nb03; - uint8_t* dst_ptr = (uint8_t *) octx->dst->data + i1*nb1 + i02*nb2 + i03*nb3; - - hvx_copy_f16_f32_uu(dst_ptr, src0_ptr, ne00); - } - } - } - - qt = HAP_perf_qtimer_count_to_us(HAP_perf_get_qtimer_count() - qt); - FARF(HIGH, "set-rows-f16-f32 %d/%d: %ux%ux%ux%u (%u:%u) x %ux%ux%ux%u -> %ux%ux%ux%u usec %u\n", ith, nth, - ne00, ne01, ne02, ne03, ir0, ir1, ne10, ne11, ne12, ne13, ne0, ne1, ne2, ne3, (unsigned) qt); -} +SET_ROWS_THREAD_DMA_FN(q8_0, int32_t, { hvx_quantize_row_q8_0_f32(dst_spad, (const float *)src_spad, ne00); }) +SET_ROWS_THREAD_DMA_FN(q8_0, int64_t, { hvx_quantize_row_q8_0_f32(dst_spad, (const float *)src_spad, ne00); }) int op_set_rows(struct htp_ops_context * octx) { + const struct htp_set_rows_kernel_params * kparams = (const struct htp_set_rows_kernel_params *)octx->kernel_params; set_rows_preamble; - const uint32_t n_threads = MIN(nr, octx->n_threads); - if (octx->src[0]->type != HTP_TYPE_F32) { return HTP_STATUS_NO_SUPPORT; } - if (octx->dst->type != HTP_TYPE_F32 && octx->dst->type != HTP_TYPE_F16) { + if (octx->dst->type != HTP_TYPE_F32 && octx->dst->type != HTP_TYPE_F16 && octx->dst->type != HTP_TYPE_Q8_0) { return HTP_STATUS_NO_SUPPORT; } @@ -164,27 +196,27 @@ int op_set_rows(struct htp_ops_context * octx) { return HTP_STATUS_NO_SUPPORT; } - if (octx->flags & HTP_OPFLAGS_SKIP_COMPUTE) { - return HTP_STATUS_OK; - } + // l2fetch the src1 (indices) tensor in the main thread + hex_l2fetch_block((const void *)octx->src[1]->data, octx->src[1]->ne[3] * octx->src[1]->nb[3]); - struct htp_set_rows_context srctx; + struct set_rows_context srctx; srctx.octx = octx; - srctx.div_ne12 = init_fastdiv_values(ne12); - srctx.div_ne11 = init_fastdiv_values(ne11); - - srctx.src0_nrows_per_thread = (nr + n_threads - 1) / n_threads; - - switch(octx->dst->type) { - case HTP_TYPE_F32: - worker_pool_run_func(octx->ctx->worker_pool, set_rows_thread_f32_f32, &srctx, n_threads); - break; - case HTP_TYPE_F16: - worker_pool_run_func(octx->ctx->worker_pool, set_rows_thread_f16_f32, &srctx, n_threads); - break; - default: - return HTP_STATUS_NO_SUPPORT; + srctx.kparams = kparams; + + htp_set_rows_vtcm_layout_build(&srctx.vtcm_layout, octx->dst->type, ne00, kparams->n_threads); + srctx.vtcm_base = (uint8_t *)octx->ctx->vtcm_base; + + work_queue_func_t q_func = NULL; + const bool is_i32 = (octx->src[1]->type == HTP_TYPE_I32); + + switch (octx->dst->type) { + case HTP_TYPE_F32: q_func = is_i32 ? set_rows_thread_dma_f32_int32_t : set_rows_thread_dma_f32_int64_t; break; + case HTP_TYPE_F16: q_func = is_i32 ? set_rows_thread_dma_f16_int32_t : set_rows_thread_dma_f16_int64_t; break; + case HTP_TYPE_Q8_0: q_func = is_i32 ? set_rows_thread_dma_q8_0_int32_t : set_rows_thread_dma_q8_0_int64_t; break; + default: return HTP_STATUS_NO_SUPPORT; } + work_queue_run(octx->ctx->work_queue, q_func, &srctx, kparams->n_threads); + return HTP_STATUS_OK; } diff --git a/ggml/src/ggml-hexagon/htp/set-rows-ops.h b/ggml/src/ggml-hexagon/htp/set-rows-ops.h new file mode 100644 index 000000000000..5e98d2cb55c5 --- /dev/null +++ b/ggml/src/ggml-hexagon/htp/set-rows-ops.h @@ -0,0 +1,74 @@ +#ifndef HTP_SET_ROWS_OPS_H +#define HTP_SET_ROWS_OPS_H + +#include "hex-fastdiv.h" + +struct htp_set_rows_kernel_params { + int32_t n_threads; + int32_t total_tasks; + int32_t tasks_per_thread; + int32_t vtcm_size; + + // Fastdiv helpers + struct fastdiv_values div_ne11; + struct fastdiv_values div_ne12; + struct fastdiv_values div_tasks_per_thread; + struct fastdiv_values div_ne02; +}; + +struct htp_set_rows_vtcm_layout { + size_t total_bytes; + size_t off_src0; + size_t off_dst; + + size_t src0_bytes_per_thread; + size_t dst_bytes_per_thread; + + size_t src0_spad_half_size; + size_t dst_spad_half_size; +}; + +static inline void htp_set_rows_vtcm_layout_build( + struct htp_set_rows_vtcm_layout * vtcm_layout, + int dst_type, + uint32_t ne00, + uint32_t n_threads) { + + size_t src0_row_size = ne00 * 4; + size_t dst_row_size = 0; + switch (dst_type) { + case 0: // HTP_TYPE_F32 + dst_row_size = ne00 * 4; + break; + case 1: // HTP_TYPE_F16 + dst_row_size = ne00 * 2; + break; + case 8: // HTP_TYPE_Q8_0 + dst_row_size = (ne00 / 32) * 34; + break; + default: + dst_row_size = 0; + break; + } + + size_t src0_row_size_aligned = (src0_row_size + 255) & ~255; + size_t dst_row_size_aligned = (dst_row_size + 255) & ~255; + + vtcm_layout->src0_spad_half_size = src0_row_size_aligned; + vtcm_layout->dst_spad_half_size = dst_row_size_aligned; + + vtcm_layout->src0_bytes_per_thread = src0_row_size_aligned * 2; + vtcm_layout->dst_bytes_per_thread = dst_row_size_aligned * 2; + + vtcm_layout->off_src0 = 0; + vtcm_layout->off_dst = vtcm_layout->off_src0 + vtcm_layout->src0_bytes_per_thread * n_threads; + vtcm_layout->total_bytes = vtcm_layout->off_dst + vtcm_layout->dst_bytes_per_thread * n_threads; +} + +#if defined(__cplusplus) +static_assert(sizeof(struct htp_set_rows_kernel_params) <= 128, "htp_set_rows_kernel_params is too large for kernel_params blob"); +#else +_Static_assert(sizeof(struct htp_set_rows_kernel_params) <= 128, "htp_set_rows_kernel_params is too large for kernel_params blob"); +#endif + +#endif // HTP_SET_ROWS_OPS_H diff --git a/ggml/src/ggml-hexagon/htp/unary-ops.c b/ggml/src/ggml-hexagon/htp/unary-ops.c index b21415a67d64..1a632bf5631e 100644 --- a/ggml/src/ggml-hexagon/htp/unary-ops.c +++ b/ggml/src/ggml-hexagon/htp/unary-ops.c @@ -443,6 +443,34 @@ static void tanh_f32(const float * restrict src, } } +static void abs_f32(const float * restrict src, + float * restrict dst, + const uint32_t num_rows, + const struct htp_unary_context * uctx) { + htp_unary_op_preamble; + + for (uint32_t ir = 0; ir < num_rows; ir++) { + const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned); + uint8_t * restrict dst_local = (uint8_t *)dst + (ir * dst_row_size_aligned); + + hvx_abs_f32_aa(dst_local, src_local, ne0); + } +} + +static void log_f32(const float * restrict src, + float * restrict dst, + const uint32_t num_rows, + const struct htp_unary_context * uctx) { + htp_unary_op_preamble; + + for (uint32_t ir = 0; ir < num_rows; ir++) { + const uint8_t * restrict src_local = (const uint8_t *)src + (ir * src0_row_size_aligned); + uint8_t * restrict dst_local = (uint8_t *)dst + (ir * dst_row_size_aligned); + + hvx_log_f32_aa(dst_local, src_local, ne0); + } +} + #define DEFINE_UNARY_TASK(NAME, IS_RMS_NORM_MUL, IS_TRI, CORE_EXPR) \ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * data) { \ const struct htp_unary_context * uctx = (const struct htp_unary_context *) data; \ @@ -478,6 +506,9 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat const uint32_t nb11 = src1 ? src1->nb[1] : 0; \ const uint32_t nb12 = src1 ? src1->nb[2] : 0; \ const uint32_t nb13 = src1 ? src1->nb[3] : 0; \ + const uint32_t nb11_bc = (src1 && src1->ne[1] > 1) ? nb11 : 0; \ + const uint32_t nb12_bc = (src1 && src1->ne[2] > 1) ? nb12 : 0; \ + const uint32_t nb13_bc = (src1 && src1->ne[3] > 1) ? nb13 : 0; \ const bool src1_contig = src1 ? ((nb12 == (size_t)ne01 * nb11) && (nb13 == (size_t)ne02 * nb12)) : false; \ \ uint8_t * src0_vtcm_data = uctx->vtcm_src0 + (ith * uctx->vtcm_src0_size_per_thread); \ @@ -497,8 +528,12 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat const struct fastdiv_values * div_ne02 = &uctx->kparams->div_ne02; \ const struct fastdiv_values * div_ne012 = &uctx->kparams->div_ne012; \ \ - const uint32_t src0_max_block = src0_contig ? uctx->block : MIN((uint32_t)uctx->block, ne01); \ - const uint32_t dst_max_block = dst_contig ? uctx->block : MIN((uint32_t)uctx->block, ne1); \ + const bool src1_needs_row_clip = (IS_RMS_NORM_MUL) && !uctx->broadcast_weight && !src1_contig; \ + const bool block_src0_contig = src0_contig && !src1_needs_row_clip; \ + const bool block_dst_contig = dst_contig && !src1_needs_row_clip; \ + \ + const uint32_t src0_max_block = block_src0_contig ? uctx->block : MIN((uint32_t)uctx->block, ne01); \ + const uint32_t dst_max_block = block_dst_contig ? uctx->block : MIN((uint32_t)uctx->block, ne1); \ const uint32_t BLOCK = MIN(src0_max_block, dst_max_block); \ if (BLOCK == 0) { \ FARF(ERROR, "unary-f32 : current VTCM reservation %zu is too small, needed at least %zu\n", \ @@ -515,8 +550,8 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat } \ \ for (uint32_t ir = src0_start_row, vtcm_idx = 0; ir < src0_end_row && vtcm_idx < 2; vtcm_idx++) { \ - const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, src0_contig, dst_contig, ne01, \ - div_ne01); \ + const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, block_src0_contig, block_dst_contig, \ + ne01, div_ne01); \ \ dma_queue_push(dma_queue, \ dma_make_ptr(data_dst, dst_vtcm_data + (vtcm_idx * dst_vtcm_half_size)), \ @@ -530,7 +565,7 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat \ if ((IS_RMS_NORM_MUL) && !uctx->broadcast_weight) { \ const size_t src1_off = src1_contig ? (ir * nb11) : \ - unary_row_offset(ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb11, nb12, nb13); \ + unary_row_offset(ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb11_bc, nb12_bc, nb13_bc); \ dma_queue_push(dma_queue, \ dma_make_ptr(src1_vtcm_data + (vtcm_idx * src1_vtcm_half_size), data_src1 + src1_off), \ uctx->src1_row_size_aligned, nb11, uctx->src1_data_row_size, block_size); \ @@ -540,8 +575,8 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat } \ \ for (uint32_t ir = src0_start_row; ir < src0_end_row; ) { \ - const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, src0_contig, dst_contig, ne01, \ - div_ne01); \ + const uint32_t block_size = unary_block_size(ir, src0_end_row, BLOCK, block_src0_contig, block_dst_contig, \ + ne01, div_ne01); \ \ float * dst_vtcm = (float *) dma_queue_pop(dma_queue).src; \ float * src0_vtcm = (float *) dma_queue_pop(dma_queue).dst; \ @@ -562,12 +597,12 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat \ const uint32_t next_ir = ir + block_size; \ if (next_ir < src0_end_row) { \ - const uint32_t next_block_size = unary_block_size(next_ir, src0_end_row, BLOCK, src0_contig, dst_contig,\ - ne01, div_ne01); \ + const uint32_t next_block_size = unary_block_size(next_ir, src0_end_row, BLOCK, block_src0_contig, \ + block_dst_contig, ne01, div_ne01); \ const uint32_t pref_ir = next_ir + next_block_size; \ if (pref_ir < src0_end_row) { \ - const uint32_t pref_block_size = unary_block_size(pref_ir, src0_end_row, BLOCK, src0_contig, \ - dst_contig, ne01, div_ne01); \ + const uint32_t pref_block_size = unary_block_size(pref_ir, src0_end_row, BLOCK, block_src0_contig, \ + block_dst_contig, ne01, div_ne01); \ const size_t src0_pref_off = src0_contig ? (pref_ir * nb01) : \ unary_row_offset(pref_ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb01, nb02, nb03); \ dma_queue_push(dma_queue, \ @@ -576,7 +611,8 @@ static void unary_task_f32_##NAME(unsigned int nth, unsigned int ith, void * dat \ if ((IS_RMS_NORM_MUL) && !uctx->broadcast_weight) { \ const size_t src1_pref_off = src1_contig ? (pref_ir * nb11) : \ - unary_row_offset(pref_ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb11, nb12, nb13); \ + unary_row_offset(pref_ir, ne01, ne02, div_ne01, div_ne02, div_ne012, nb11_bc, nb12_bc, \ + nb13_bc); \ dma_queue_push(dma_queue, \ dma_make_ptr(src1_vtcm, data_src1 + src1_pref_off), \ uctx->src1_row_size_aligned, nb11, uctx->src1_data_row_size, pref_block_size); \ @@ -603,6 +639,8 @@ DEFINE_UNARY_TASK(unary_silu, false, false, silu_f32(src0_vtcm, dst_vtcm, bl DEFINE_UNARY_TASK(unary_gelu, false, false, gelu_f32(src0_vtcm, dst_vtcm, block_size, uctx)) DEFINE_UNARY_TASK(unary_softplus, false, false, softplus_f32(src0_vtcm, dst_vtcm, block_size, uctx)) DEFINE_UNARY_TASK(unary_tanh, false, false, tanh_f32(src0_vtcm, dst_vtcm, block_size, uctx)) +DEFINE_UNARY_TASK(unary_abs, false, false, abs_f32(src0_vtcm, dst_vtcm, block_size, uctx)) +DEFINE_UNARY_TASK(unary_log, false, false, log_f32(src0_vtcm, dst_vtcm, block_size, uctx)) DEFINE_UNARY_TASK(l2_norm, false, false, l2_norm_f32(src0_vtcm, dst_vtcm, block_size, uctx)) DEFINE_UNARY_TASK(tri, false, true, tri_f32(src0_vtcm, dst_vtcm, block_size, ir, uctx)) @@ -850,6 +888,8 @@ DEFINE_UNARY_TILED_TASK(unary_silu, false, tile_silu_f32(dst_vtcm, src_vtcm, DEFINE_UNARY_TILED_TASK(unary_gelu, false, tile_gelu_f32(dst_vtcm, src_vtcm, tw)) DEFINE_UNARY_TILED_TASK(unary_softplus, false, tile_unary_softplus_f32(dst_vtcm, src_vtcm, tw)) DEFINE_UNARY_TILED_TASK(unary_tanh, false, hvx_tanh_f32_aa(dst_vtcm, src_vtcm, tw)) +DEFINE_UNARY_TILED_TASK(unary_abs, false, hvx_abs_f32_aa(dst_vtcm, src_vtcm, tw)) +DEFINE_UNARY_TILED_TASK(unary_log, false, hvx_log_f32_aa(dst_vtcm, src_vtcm, tw)) DEFINE_UNARY_TILED_TASK(tri, true, tri_apply_tile_f32(src_vtcm, dst_vtcm, tw, col, i01, ne0, tri_ttype)) static int execute_op_unary_f32(struct htp_ops_context * octx) { @@ -875,6 +915,8 @@ static int execute_op_unary_f32(struct htp_ops_context * octx) { case HTP_OP_UNARY_GELU: op_type = "gelu-f32"; break; case HTP_OP_UNARY_SOFTPLUS: op_type = "softplus-f32"; break; case HTP_OP_UNARY_TANH: op_type = "tanh-f32"; break; + case HTP_OP_UNARY_ABS: op_type = "abs-f32"; break; + case HTP_OP_UNARY_LOG: op_type = "log-f32"; break; case HTP_OP_L2_NORM: op_type = "l2norm-f32"; break; case HTP_OP_TRI: op_type = "tri-f32"; break; @@ -973,6 +1015,8 @@ static int execute_op_unary_f32(struct htp_ops_context * octx) { case HTP_OP_UNARY_GELU: task_func = unary_task_f32_tiled_unary_gelu; break; case HTP_OP_UNARY_SOFTPLUS: task_func = unary_task_f32_tiled_unary_softplus; break; case HTP_OP_UNARY_TANH: task_func = unary_task_f32_tiled_unary_tanh; break; + case HTP_OP_UNARY_ABS: task_func = unary_task_f32_tiled_unary_abs; break; + case HTP_OP_UNARY_LOG: task_func = unary_task_f32_tiled_unary_log; break; case HTP_OP_TRI: task_func = unary_task_f32_tiled_tri; break; default: break; } @@ -992,6 +1036,8 @@ static int execute_op_unary_f32(struct htp_ops_context * octx) { case HTP_OP_UNARY_GELU: task_func = unary_task_f32_unary_gelu; break; case HTP_OP_UNARY_SOFTPLUS: task_func = unary_task_f32_unary_softplus; break; case HTP_OP_UNARY_TANH: task_func = unary_task_f32_unary_tanh; break; + case HTP_OP_UNARY_ABS: task_func = unary_task_f32_unary_abs; break; + case HTP_OP_UNARY_LOG: task_func = unary_task_f32_unary_log; break; case HTP_OP_L2_NORM: task_func = unary_task_f32_l2_norm; break; case HTP_OP_TRI: task_func = unary_task_f32_tri; break; default: break; diff --git a/ggml/src/ggml-hexagon/htp/unary-ops.h b/ggml/src/ggml-hexagon/htp/unary-ops.h index 1f4c3a5c4d96..458218ff4431 100644 --- a/ggml/src/ggml-hexagon/htp/unary-ops.h +++ b/ggml/src/ggml-hexagon/htp/unary-ops.h @@ -55,6 +55,8 @@ static inline bool htp_op_is_unary(uint32_t opcode) { case HTP_OP_UNARY_GELU: case HTP_OP_UNARY_SOFTPLUS: case HTP_OP_UNARY_TANH: + case HTP_OP_UNARY_ABS: + case HTP_OP_UNARY_LOG: case HTP_OP_L2_NORM: case HTP_OP_TRI: return true; diff --git a/ggml/src/ggml-metal/ggml-metal-context.m b/ggml/src/ggml-metal/ggml-metal-context.m index 32d97cd5d0af..1227ed39a095 100644 --- a/ggml/src/ggml-metal/ggml-metal-context.m +++ b/ggml/src/ggml-metal/ggml-metal-context.m @@ -84,106 +84,108 @@ ggml_metal_t ggml_metal_init(ggml_metal_device_t dev) { GGML_LOG_INFO("%s: allocating\n", __func__); + @autoreleasepool { #if TARGET_OS_OSX && !GGML_METAL_NDEBUG - // Show all the Metal device instances in the system - NSArray * devices = MTLCopyAllDevices(); - for (id device in devices) { - GGML_LOG_INFO("%s: found device: %s\n", __func__, [[device name] UTF8String]); - } - [devices release]; // since it was created by a *Copy* C method + // Show all the Metal device instances in the system + NSArray * devices = MTLCopyAllDevices(); + for (id device in devices) { + GGML_LOG_INFO("%s: found device: %s\n", __func__, [[device name] UTF8String]); + } + [devices release]; // since it was created by a *Copy* C method #endif - // init context - ggml_metal_t res = calloc(1, sizeof(struct ggml_metal)); - - id device = ggml_metal_device_get_obj(dev); + // init context + ggml_metal_t res = calloc(1, sizeof(struct ggml_metal)); - GGML_LOG_INFO("%s: picking default device: %s\n", __func__, [[device name] UTF8String]); + id device = ggml_metal_device_get_obj(dev); - // TODO: would it be better to have one queue for the backend and one queue for the device? - // the graph encoders and async ops would use the backend queue while the sync ops would use the device queue? - //res->queue = [device newCommandQueue]; [TAG_QUEUE_PER_BACKEND] - id queue = ggml_metal_device_get_queue(dev); - if (queue == nil) { - GGML_LOG_ERROR("%s: error: failed to create command queue\n", __func__); - return NULL; - } + GGML_LOG_INFO("%s: picking default device: %s\n", __func__, [[device name] UTF8String]); - res->dev = dev; - res->lib = ggml_metal_device_get_library(dev); - if (res->lib == NULL) { - GGML_LOG_WARN("%s: the device does not have a precompiled Metal library - this is unexpected\n", __func__); - GGML_LOG_WARN("%s: will try to compile it on the fly\n", __func__); + // TODO: would it be better to have one queue for the backend and one queue for the device? + // the graph encoders and async ops would use the backend queue while the sync ops would use the device queue? + //res->queue = [device newCommandQueue]; [TAG_QUEUE_PER_BACKEND] + id queue = ggml_metal_device_get_queue(dev); + if (queue == nil) { + GGML_LOG_ERROR("%s: error: failed to create command queue\n", __func__); + return NULL; + } - res->lib = ggml_metal_library_init(dev); + res->dev = dev; + res->lib = ggml_metal_device_get_library(dev); if (res->lib == NULL) { - GGML_LOG_ERROR("%s: error: failed to initialize the Metal library\n", __func__); + GGML_LOG_WARN("%s: the device does not have a precompiled Metal library - this is unexpected\n", __func__); + GGML_LOG_WARN("%s: will try to compile it on the fly\n", __func__); - free(res); + res->lib = ggml_metal_library_init(dev); + if (res->lib == NULL) { + GGML_LOG_ERROR("%s: error: failed to initialize the Metal library\n", __func__); - return NULL; + free(res); + + return NULL; + } } - } - res->ev_cpy = ggml_metal_device_event_init(dev); + res->ev_cpy = ggml_metal_device_event_init(dev); - const struct ggml_metal_device_props * props_dev = ggml_metal_device_get_props(dev); + const struct ggml_metal_device_props * props_dev = ggml_metal_device_get_props(dev); - snprintf(res->name, sizeof(res->name), "%s", props_dev->name); + snprintf(res->name, sizeof(res->name), "%s", props_dev->name); - res->d_queue = dispatch_queue_create("ggml-metal", DISPATCH_QUEUE_CONCURRENT); + res->d_queue = dispatch_queue_create("ggml-metal", DISPATCH_QUEUE_CONCURRENT); - res->use_fusion = getenv("GGML_METAL_FUSION_DISABLE") == nil; - res->use_concurrency = getenv("GGML_METAL_CONCURRENCY_DISABLE") == nil; + res->use_fusion = getenv("GGML_METAL_FUSION_DISABLE") == nil; + res->use_concurrency = getenv("GGML_METAL_CONCURRENCY_DISABLE") == nil; - { - const char * val = getenv("GGML_METAL_GRAPH_DEBUG"); - res->debug_graph = val ? atoi(val) : 0; - } + { + const char * val = getenv("GGML_METAL_GRAPH_DEBUG"); + res->debug_graph = val ? atoi(val) : 0; + } - { - const char * val = getenv("GGML_METAL_FUSION_DEBUG"); - res->debug_fusion = val ? atoi(val) : 0; - } + { + const char * val = getenv("GGML_METAL_FUSION_DEBUG"); + res->debug_fusion = val ? atoi(val) : 0; + } - res->use_graph_optimize = true; + res->use_graph_optimize = true; - if (getenv("GGML_METAL_GRAPH_OPTIMIZE_DISABLE") != NULL) { - res->use_graph_optimize = false; - } + if (getenv("GGML_METAL_GRAPH_OPTIMIZE_DISABLE") != NULL) { + res->use_graph_optimize = false; + } - memset(res->fuse_cnt, 0, sizeof(res->fuse_cnt)); + memset(res->fuse_cnt, 0, sizeof(res->fuse_cnt)); - GGML_LOG_INFO("%s: use fusion = %s\n", __func__, res->use_fusion ? "true" : "false"); - GGML_LOG_INFO("%s: use concurrency = %s\n", __func__, res->use_concurrency ? "true" : "false"); - GGML_LOG_INFO("%s: use graph optimize = %s\n", __func__, res->use_graph_optimize ? "true" : "false"); + GGML_LOG_INFO("%s: use fusion = %s\n", __func__, res->use_fusion ? "true" : "false"); + GGML_LOG_INFO("%s: use concurrency = %s\n", __func__, res->use_concurrency ? "true" : "false"); + GGML_LOG_INFO("%s: use graph optimize = %s\n", __func__, res->use_graph_optimize ? "true" : "false"); - res->capture_compute = 0; - res->capture_started = false; - res->capture_scope = nil; + res->capture_compute = 0; + res->capture_started = false; + res->capture_scope = nil; - { - const char * val = getenv("GGML_METAL_CAPTURE_COMPUTE"); - if (val) { - res->capture_compute = atoi(val); + { + const char * val = getenv("GGML_METAL_CAPTURE_COMPUTE"); + if (val) { + res->capture_compute = atoi(val); + } } - } - res->has_error = false; + res->has_error = false; - res->gf = nil; - res->encode_async = nil; - for (int i = 0; i < GGML_METAL_MAX_COMMAND_BUFFERS; ++i) { - res->cmd_bufs[i].obj = nil; - } + res->gf = nil; + res->encode_async = nil; + for (int i = 0; i < GGML_METAL_MAX_COMMAND_BUFFERS; ++i) { + res->cmd_bufs[i].obj = nil; + } - res->cmd_bufs_ext = [[NSMutableArray alloc] init]; + res->cmd_bufs_ext = [[NSMutableArray alloc] init]; - res->cmd_buf_last = nil; + res->cmd_buf_last = nil; - res->pipelines_ext = ggml_metal_pipelines_init(); + res->pipelines_ext = ggml_metal_pipelines_init(); - return res; + return res; + } } void ggml_metal_free(ggml_metal_t ctx) { diff --git a/ggml/src/ggml-metal/ggml-metal-device.m b/ggml/src/ggml-metal/ggml-metal-device.m index 19c57820e859..41ce90dc8a93 100644 --- a/ggml/src/ggml-metal/ggml-metal-device.m +++ b/ggml/src/ggml-metal/ggml-metal-device.m @@ -778,7 +778,9 @@ void ggml_metal_encoder_free(ggml_metal_encoder_t encoder) { } void ggml_metal_encoder_debug_group_push(ggml_metal_encoder_t encoder, const char * name) { - [encoder->obj pushDebugGroup:[NSString stringWithCString:name encoding:NSUTF8StringEncoding]]; + @autoreleasepool { + [encoder->obj pushDebugGroup:[NSString stringWithCString:name encoding:NSUTF8StringEncoding]]; + } } void ggml_metal_encoder_debug_group_pop (ggml_metal_encoder_t encoder) { @@ -1023,249 +1025,251 @@ ggml_metal_device_t ggml_metal_device_init(int device, int n_devices) { assert(dev != NULL); - if (dev->mtl_device == nil) { - dev->mtl_device = MTLCreateSystemDefaultDevice(); - - if (dev->mtl_device) { - dev->mtl_queue = [dev->mtl_device newCommandQueue]; - if (dev->mtl_queue == nil) { - GGML_LOG_ERROR("%s: error: failed to create command queue\n", __func__); - } + @autoreleasepool { + if (dev->mtl_device == nil) { + dev->mtl_device = MTLCreateSystemDefaultDevice(); - dev->addr_virt = 0x000000400ULL; + if (dev->mtl_device) { + dev->mtl_queue = [dev->mtl_device newCommandQueue]; + if (dev->mtl_queue == nil) { + GGML_LOG_ERROR("%s: error: failed to create command queue\n", __func__); + } - dev->props.device = device; + dev->addr_virt = 0x000000400ULL; - // the Metal backend uses the system default device as the single physical device; - // additional (virtual) devices are emulated on top of it via GGML_METAL_DEVICES - dev->props.device_phys = 0; - dev->props.device_virt = device; + dev->props.device = device; - dev->props.has_simdgroup_reduction = [dev->mtl_device supportsFamily:MTLGPUFamilyApple7]; - dev->props.has_simdgroup_reduction |= [dev->mtl_device supportsFamily:MTLGPUFamilyMetal3_GGML]; + // the Metal backend uses the system default device as the single physical device; + // additional (virtual) devices are emulated on top of it via GGML_METAL_DEVICES + dev->props.device_phys = 0; + dev->props.device_virt = device; - dev->props.has_simdgroup_mm = [dev->mtl_device supportsFamily:MTLGPUFamilyApple7]; - dev->props.has_unified_memory = dev->mtl_device.hasUnifiedMemory; + dev->props.has_simdgroup_reduction = [dev->mtl_device supportsFamily:MTLGPUFamilyApple7]; + dev->props.has_simdgroup_reduction |= [dev->mtl_device supportsFamily:MTLGPUFamilyMetal3_GGML]; - dev->props.has_bfloat = [dev->mtl_device supportsFamily:MTLGPUFamilyMetal3_GGML]; - dev->props.has_bfloat |= [dev->mtl_device supportsFamily:MTLGPUFamilyApple6]; - if (getenv("GGML_METAL_BF16_DISABLE") != NULL) { - dev->props.has_bfloat = false; - } + dev->props.has_simdgroup_mm = [dev->mtl_device supportsFamily:MTLGPUFamilyApple7]; + dev->props.has_unified_memory = dev->mtl_device.hasUnifiedMemory; - dev->props.has_tensor = [dev->mtl_device supportsFamily:MTLGPUFamilyMetal4_GGML]; - if (getenv("GGML_METAL_TENSOR_DISABLE") != NULL) { - dev->props.has_tensor = false; - } + dev->props.has_bfloat = [dev->mtl_device supportsFamily:MTLGPUFamilyMetal3_GGML]; + dev->props.has_bfloat |= [dev->mtl_device supportsFamily:MTLGPUFamilyApple6]; + if (getenv("GGML_METAL_BF16_DISABLE") != NULL) { + dev->props.has_bfloat = false; + } - // note: disable the tensor API by default for old chips because with the current implementation it is not useful - // - M2 Ultra: ~5% slower - // - M4, M4 Max: no significant difference - // - // TODO: try to update the tensor API kernels to at least match the simdgroup performance - if (getenv("GGML_METAL_TENSOR_ENABLE") == NULL && - ![[dev->mtl_device name] containsString:@"M5"] && - ![[dev->mtl_device name] containsString:@"M6"] && - ![[dev->mtl_device name] containsString:@"A19"] && - ![[dev->mtl_device name] containsString:@"A20"]) { - GGML_LOG_INFO("%s: tensor API disabled for pre-M5 and pre-A19 devices\n", __func__); - dev->props.has_tensor = false; - } + dev->props.has_tensor = [dev->mtl_device supportsFamily:MTLGPUFamilyMetal4_GGML]; + if (getenv("GGML_METAL_TENSOR_DISABLE") != NULL) { + dev->props.has_tensor = false; + } - // double-check that the tensor API compiles - if (dev->props.has_tensor) { - const char * src_tensor_f16 = "\n" - "#include \n" - "#include \n" - "#include \n" - " \n" - "using namespace metal; \n" - "using namespace mpp::tensor_ops; \n" - " \n" - "kernel void dummy_kernel( \n" - " tensor> A [[buffer(0)]], \n" - " tensor> B [[buffer(1)]], \n" - " device float * C [[buffer(2)]], \n" - " uint2 tgid [[threadgroup_position_in_grid]]) \n" - "{ \n" - " auto tA = A.slice(0, (int)tgid.y); \n" - " auto tB = B.slice((int)tgid.x, 0); \n" - " \n" - " matmul2d< \n" - " matmul2d_descriptor(16, 16, dynamic_extent), \n" - " execution_simdgroups<4>> mm; \n" - " \n" - " auto cT = mm.get_destination_cooperative_tensor(); \n" - " \n" - " auto sA = tA.slice(0, 0); \n" - " auto sB = tB.slice(0, 0); \n" - " mm.run(sB, sA, cT); \n" - " \n" - " auto tC = tensor, tensor_inline>(C, dextents(16, 16)); \n" - " \n" - " cT.store(tC); \n" - "}"; - - GGML_LOG_INFO("%s: testing tensor API for f16 support\n", __func__); - ggml_metal_library_t lib = ggml_metal_library_init_from_source(dev, src_tensor_f16, false); - if (lib == NULL) { - GGML_LOG_WARN("%s: - the tensor API is not supported in this environment - disabling\n", __func__); + // note: disable the tensor API by default for old chips because with the current implementation it is not useful + // - M2 Ultra: ~5% slower + // - M4, M4 Max: no significant difference + // + // TODO: try to update the tensor API kernels to at least match the simdgroup performance + if (getenv("GGML_METAL_TENSOR_ENABLE") == NULL && + ![[dev->mtl_device name] containsString:@"M5"] && + ![[dev->mtl_device name] containsString:@"M6"] && + ![[dev->mtl_device name] containsString:@"A19"] && + ![[dev->mtl_device name] containsString:@"A20"]) { + GGML_LOG_INFO("%s: tensor API disabled for pre-M5 and pre-A19 devices\n", __func__); dev->props.has_tensor = false; - } else { - struct ggml_metal_pipeline_with_params ppl = ggml_metal_library_compile_pipeline(lib, "dummy_kernel", "dummy_kernel", nil); - if (!ppl.pipeline) { + } + + // double-check that the tensor API compiles + if (dev->props.has_tensor) { + const char * src_tensor_f16 = "\n" + "#include \n" + "#include \n" + "#include \n" + " \n" + "using namespace metal; \n" + "using namespace mpp::tensor_ops; \n" + " \n" + "kernel void dummy_kernel( \n" + " tensor> A [[buffer(0)]], \n" + " tensor> B [[buffer(1)]], \n" + " device float * C [[buffer(2)]], \n" + " uint2 tgid [[threadgroup_position_in_grid]]) \n" + "{ \n" + " auto tA = A.slice(0, (int)tgid.y); \n" + " auto tB = B.slice((int)tgid.x, 0); \n" + " \n" + " matmul2d< \n" + " matmul2d_descriptor(16, 16, dynamic_extent), \n" + " execution_simdgroups<4>> mm; \n" + " \n" + " auto cT = mm.get_destination_cooperative_tensor(); \n" + " \n" + " auto sA = tA.slice(0, 0); \n" + " auto sB = tB.slice(0, 0); \n" + " mm.run(sB, sA, cT); \n" + " \n" + " auto tC = tensor, tensor_inline>(C, dextents(16, 16)); \n" + " \n" + " cT.store(tC); \n" + "}"; + + GGML_LOG_INFO("%s: testing tensor API for f16 support\n", __func__); + ggml_metal_library_t lib = ggml_metal_library_init_from_source(dev, src_tensor_f16, false); + if (lib == NULL) { GGML_LOG_WARN("%s: - the tensor API is not supported in this environment - disabling\n", __func__); dev->props.has_tensor = false; - } + } else { + struct ggml_metal_pipeline_with_params ppl = ggml_metal_library_compile_pipeline(lib, "dummy_kernel", "dummy_kernel", nil); + if (!ppl.pipeline) { + GGML_LOG_WARN("%s: - the tensor API is not supported in this environment - disabling\n", __func__); + dev->props.has_tensor = false; + } - ggml_metal_library_free(lib); + ggml_metal_library_free(lib); + } } - } - // try to compile a dummy kernel to determine if the tensor API is supported for bfloat - if (dev->props.has_tensor && dev->props.has_bfloat) { - const char * src_tensor_bf16 = "\n" - "#include \n" - "#include \n" - "#include \n" - " \n" - "using namespace metal; \n" - "using namespace mpp::tensor_ops; \n" - " \n" - "kernel void dummy_kernel( \n" - " tensor> A [[buffer(0)]], \n" - " tensor> B [[buffer(1)]], \n" - " device float * C [[buffer(2)]], \n" - " uint2 tgid [[threadgroup_position_in_grid]]) \n" - "{ \n" - " auto tA = A.slice(0, (int)tgid.y); \n" - " auto tB = B.slice((int)tgid.x, 0); \n" - " \n" - " matmul2d< \n" - " matmul2d_descriptor(16, 16, dynamic_extent), \n" - " execution_simdgroups<4>> mm; \n" - " \n" - " auto cT = mm.get_destination_cooperative_tensor(); \n" - " \n" - " auto sA = tA.slice(0, 0); \n" - " auto sB = tB.slice(0, 0); \n" - " mm.run(sB, sA, cT); \n" - " \n" - " auto tC = tensor, tensor_inline>(C, dextents(16, 16)); \n" - " \n" - " cT.store(tC); \n" - "}"; - - GGML_LOG_INFO("%s: testing tensor API for bfloat support\n", __func__); - ggml_metal_library_t lib = ggml_metal_library_init_from_source(dev, src_tensor_bf16, false); - if (lib == NULL) { - GGML_LOG_WARN("%s: - the tensor API does not support bfloat - disabling bfloat support\n", __func__); - dev->props.has_bfloat = false; - } else { - struct ggml_metal_pipeline_with_params ppl = ggml_metal_library_compile_pipeline(lib, "dummy_kernel", "dummy_kernel", nil); - if (!ppl.pipeline) { + // try to compile a dummy kernel to determine if the tensor API is supported for bfloat + if (dev->props.has_tensor && dev->props.has_bfloat) { + const char * src_tensor_bf16 = "\n" + "#include \n" + "#include \n" + "#include \n" + " \n" + "using namespace metal; \n" + "using namespace mpp::tensor_ops; \n" + " \n" + "kernel void dummy_kernel( \n" + " tensor> A [[buffer(0)]], \n" + " tensor> B [[buffer(1)]], \n" + " device float * C [[buffer(2)]], \n" + " uint2 tgid [[threadgroup_position_in_grid]]) \n" + "{ \n" + " auto tA = A.slice(0, (int)tgid.y); \n" + " auto tB = B.slice((int)tgid.x, 0); \n" + " \n" + " matmul2d< \n" + " matmul2d_descriptor(16, 16, dynamic_extent), \n" + " execution_simdgroups<4>> mm; \n" + " \n" + " auto cT = mm.get_destination_cooperative_tensor(); \n" + " \n" + " auto sA = tA.slice(0, 0); \n" + " auto sB = tB.slice(0, 0); \n" + " mm.run(sB, sA, cT); \n" + " \n" + " auto tC = tensor, tensor_inline>(C, dextents(16, 16)); \n" + " \n" + " cT.store(tC); \n" + "}"; + + GGML_LOG_INFO("%s: testing tensor API for bfloat support\n", __func__); + ggml_metal_library_t lib = ggml_metal_library_init_from_source(dev, src_tensor_bf16, false); + if (lib == NULL) { GGML_LOG_WARN("%s: - the tensor API does not support bfloat - disabling bfloat support\n", __func__); dev->props.has_bfloat = false; - } + } else { + struct ggml_metal_pipeline_with_params ppl = ggml_metal_library_compile_pipeline(lib, "dummy_kernel", "dummy_kernel", nil); + if (!ppl.pipeline) { + GGML_LOG_WARN("%s: - the tensor API does not support bfloat - disabling bfloat support\n", __func__); + dev->props.has_bfloat = false; + } - ggml_metal_library_free(lib); + ggml_metal_library_free(lib); + } } - } - dev->props.use_residency_sets = true; + dev->props.use_residency_sets = true; #if defined(GGML_METAL_HAS_RESIDENCY_SETS) - dev->props.use_residency_sets = getenv("GGML_METAL_NO_RESIDENCY") == nil; + dev->props.use_residency_sets = getenv("GGML_METAL_NO_RESIDENCY") == nil; #endif - dev->props.use_shared_buffers = dev->props.has_unified_memory; + dev->props.use_shared_buffers = dev->props.has_unified_memory; #if TARGET_OS_OSX - // In case of eGPU, shared memory may be preferable. - dev->props.use_shared_buffers |= [dev->mtl_device location] == MTLDeviceLocationExternal; + // In case of eGPU, shared memory may be preferable. + dev->props.use_shared_buffers |= [dev->mtl_device location] == MTLDeviceLocationExternal; #endif - if (getenv("GGML_METAL_SHARED_BUFFERS_DISABLE") != NULL) { - dev->props.use_shared_buffers = false; - } - if (getenv("GGML_METAL_SHARED_BUFFERS_ENABLE") != NULL) { - dev->props.use_shared_buffers = true; - } + if (getenv("GGML_METAL_SHARED_BUFFERS_DISABLE") != NULL) { + dev->props.use_shared_buffers = false; + } + if (getenv("GGML_METAL_SHARED_BUFFERS_ENABLE") != NULL) { + dev->props.use_shared_buffers = true; + } - dev->props.supports_gpu_family_apple7 = [dev->mtl_device supportsFamily:MTLGPUFamilyApple7]; + dev->props.supports_gpu_family_apple7 = [dev->mtl_device supportsFamily:MTLGPUFamilyApple7]; - dev->props.device_id = ggml_metal_device_id_parse([[dev->mtl_device name] UTF8String]); + dev->props.device_id = ggml_metal_device_id_parse([[dev->mtl_device name] UTF8String]); - dev->props.op_offload_min_batch_size = getenv("GGML_OP_OFFLOAD_MIN_BATCH") ? atoi(getenv("GGML_OP_OFFLOAD_MIN_BATCH")) : 32; + dev->props.op_offload_min_batch_size = getenv("GGML_OP_OFFLOAD_MIN_BATCH") ? atoi(getenv("GGML_OP_OFFLOAD_MIN_BATCH")) : 32; - dev->props.max_buffer_size = dev->mtl_device.maxBufferLength; - dev->props.max_theadgroup_memory_size = dev->mtl_device.maxThreadgroupMemoryLength; - if (@available(macOS 10.12, iOS 16.0, *)) { - dev->props.max_working_set_size = dev->mtl_device.recommendedMaxWorkingSetSize; - } else { - dev->props.max_working_set_size = dev->mtl_device.maxBufferLength; - } - - snprintf(dev->props.name, sizeof(dev->props.name), "%s%d", "MTL", device); - const char * gpu_name = [[dev->mtl_device name] UTF8String]; - if (n_devices > 1) { - snprintf(dev->props.desc, sizeof(dev->props.desc), "%s (dev p%d/v%d)", - gpu_name, dev->props.device_phys, dev->props.device_virt); - } else { - snprintf(dev->props.desc, sizeof(dev->props.desc), "%s", gpu_name); - } + dev->props.max_buffer_size = dev->mtl_device.maxBufferLength; + dev->props.max_theadgroup_memory_size = dev->mtl_device.maxThreadgroupMemoryLength; + if (@available(macOS 10.12, iOS 16.0, *)) { + dev->props.max_working_set_size = dev->mtl_device.recommendedMaxWorkingSetSize; + } else { + dev->props.max_working_set_size = dev->mtl_device.maxBufferLength; + } - dev->library = ggml_metal_library_init(dev); - if (!dev->library) { - GGML_LOG_ERROR("%s: error: failed to create library\n", __func__); - } + snprintf(dev->props.name, sizeof(dev->props.name), "%s%d", "MTL", device); + const char * gpu_name = [[dev->mtl_device name] UTF8String]; + if (n_devices > 1) { + snprintf(dev->props.desc, sizeof(dev->props.desc), "%s (dev p%d/v%d)", + gpu_name, dev->props.device_phys, dev->props.device_virt); + } else { + snprintf(dev->props.desc, sizeof(dev->props.desc), "%s", gpu_name); + } - if (dev->props.use_residency_sets) { - dev->rsets = ggml_metal_rsets_init(dev); - } else { - dev->rsets = nil; - } + dev->library = ggml_metal_library_init(dev); + if (!dev->library) { + GGML_LOG_ERROR("%s: error: failed to create library\n", __func__); + } - // print MTL GPU family: - GGML_LOG_INFO("%s: GPU name: %s (%s)\n", __func__, dev->props.name, dev->props.desc); + if (dev->props.use_residency_sets) { + dev->rsets = ggml_metal_rsets_init(dev); + } else { + dev->rsets = nil; + } - // determine max supported GPU family - // https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf - // https://developer.apple.com/metal/Metal-Feature-Set-Tables.pdf - { - for (int i = MTLGPUFamilyApple1 + 20; i >= MTLGPUFamilyApple1; --i) { - if ([dev->mtl_device supportsFamily:i]) { - dev->props.gpu_family = i - (int) MTLGPUFamilyApple1 + 1; - GGML_LOG_INFO("%s: GPU family: MTLGPUFamilyApple%d (%d)\n", __func__, dev->props.gpu_family, i); - break; + // print MTL GPU family: + GGML_LOG_INFO("%s: GPU name: %s (%s)\n", __func__, dev->props.name, dev->props.desc); + + // determine max supported GPU family + // https://developer.apple.com/metal/Metal-Shading-Language-Specification.pdf + // https://developer.apple.com/metal/Metal-Feature-Set-Tables.pdf + { + for (int i = MTLGPUFamilyApple1 + 20; i >= MTLGPUFamilyApple1; --i) { + if ([dev->mtl_device supportsFamily:i]) { + dev->props.gpu_family = i - (int) MTLGPUFamilyApple1 + 1; + GGML_LOG_INFO("%s: GPU family: MTLGPUFamilyApple%d (%d)\n", __func__, dev->props.gpu_family, i); + break; + } } - } - for (int i = MTLGPUFamilyCommon1 + 5; i >= MTLGPUFamilyCommon1; --i) { - if ([dev->mtl_device supportsFamily:i]) { - GGML_LOG_INFO("%s: GPU family: MTLGPUFamilyCommon%d (%d)\n", __func__, i - (int) MTLGPUFamilyCommon1 + 1, i); - break; + for (int i = MTLGPUFamilyCommon1 + 5; i >= MTLGPUFamilyCommon1; --i) { + if ([dev->mtl_device supportsFamily:i]) { + GGML_LOG_INFO("%s: GPU family: MTLGPUFamilyCommon%d (%d)\n", __func__, i - (int) MTLGPUFamilyCommon1 + 1, i); + break; + } } - } - for (int i = MTLGPUFamilyMetal3_GGML + 5; i >= MTLGPUFamilyMetal3_GGML; --i) { - if ([dev->mtl_device supportsFamily:i]) { - GGML_LOG_INFO("%s: GPU family: MTLGPUFamilyMetal%d (%d)\n", __func__, i - (int) MTLGPUFamilyMetal3_GGML + 3, i); - break; + for (int i = MTLGPUFamilyMetal3_GGML + 5; i >= MTLGPUFamilyMetal3_GGML; --i) { + if ([dev->mtl_device supportsFamily:i]) { + GGML_LOG_INFO("%s: GPU family: MTLGPUFamilyMetal%d (%d)\n", __func__, i - (int) MTLGPUFamilyMetal3_GGML + 3, i); + break; + } } } - } - GGML_LOG_INFO("%s: simdgroup reduction = %s\n", __func__, dev->props.has_simdgroup_reduction ? "true" : "false"); - GGML_LOG_INFO("%s: simdgroup matrix mul. = %s\n", __func__, dev->props.has_simdgroup_mm ? "true" : "false"); - GGML_LOG_INFO("%s: has unified memory = %s\n", __func__, dev->props.has_unified_memory ? "true" : "false"); - GGML_LOG_INFO("%s: has bfloat = %s\n", __func__, dev->props.has_bfloat ? "true" : "false"); - GGML_LOG_INFO("%s: has tensor = %s\n", __func__, dev->props.has_tensor ? "true" : "false"); - GGML_LOG_INFO("%s: use residency sets = %s\n", __func__, dev->props.use_residency_sets ? "true" : "false"); - GGML_LOG_INFO("%s: use shared buffers = %s\n", __func__, dev->props.use_shared_buffers ? "true" : "false"); + GGML_LOG_INFO("%s: simdgroup reduction = %s\n", __func__, dev->props.has_simdgroup_reduction ? "true" : "false"); + GGML_LOG_INFO("%s: simdgroup matrix mul. = %s\n", __func__, dev->props.has_simdgroup_mm ? "true" : "false"); + GGML_LOG_INFO("%s: has unified memory = %s\n", __func__, dev->props.has_unified_memory ? "true" : "false"); + GGML_LOG_INFO("%s: has bfloat = %s\n", __func__, dev->props.has_bfloat ? "true" : "false"); + GGML_LOG_INFO("%s: has tensor = %s\n", __func__, dev->props.has_tensor ? "true" : "false"); + GGML_LOG_INFO("%s: use residency sets = %s\n", __func__, dev->props.use_residency_sets ? "true" : "false"); + GGML_LOG_INFO("%s: use shared buffers = %s\n", __func__, dev->props.use_shared_buffers ? "true" : "false"); #if TARGET_OS_OSX || (TARGET_OS_IOS && __clang_major__ >= 15) - if (@available(macOS 10.12, iOS 16.0, *)) { - GGML_LOG_INFO("%s: recommendedMaxWorkingSetSize = %8.2f MB\n", __func__, dev->props.max_working_set_size / 1e6); - } + if (@available(macOS 10.12, iOS 16.0, *)) { + GGML_LOG_INFO("%s: recommendedMaxWorkingSetSize = %8.2f MB\n", __func__, dev->props.max_working_set_size / 1e6); + } #endif + } } } diff --git a/ggml/src/ggml-opencl/ggml-opencl.cpp b/ggml/src/ggml-opencl/ggml-opencl.cpp index 64f3325b2a59..6ae83449b082 100644 --- a/ggml/src/ggml-opencl/ggml-opencl.cpp +++ b/ggml/src/ggml-opencl/ggml-opencl.cpp @@ -903,6 +903,8 @@ struct ggml_backend_opencl_context { cl_kernel kernel_gemv_moe_mxfp4_f32_ns_wimg = nullptr; // weight-as-texture MoE decode GEMV cl_kernel kernel_gemm_moe_mxfp4_q8_1_dp4a = nullptr; // dp4a (int8) mxfp4 MoE prefill GEMM cl_kernel kernel_gemm_moe_q4_0_q8_1_dp4a = nullptr; // dp4a (int8) q4_0 MoE prefill GEMM + cl_kernel kernel_gemm_moe_mxfp4_q8_1_dp4a_bin = nullptr; // binary dp4a (int8) mxfp4 MoE prefill GEMM + cl_kernel kernel_gemm_moe_q4_0_q8_1_dp4a_bin = nullptr; // binary dp4a (int8) q4_0 MoE prefill GEMM cl_kernel kernel_moe_reorder_b; cl_kernel kernel_moe_histogram, kernel_moe_scan, kernel_moe_fill, kernel_moe_scatter; cl_kernel kernel_moe_scatter_stable = nullptr; // deterministic slot assignment @@ -4248,6 +4250,24 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) { GGML_LOG_CONT("."); } + // gemm_moe_mxfp4_q8_1_dp4a_bin (dp4a prefill GEMM) + if (backend_ctx->has_integer_dot) { + size_t bin_size = 0; + backend_ctx->kernel_gemm_moe_mxfp4_q8_1_dp4a_bin = nullptr; + + if (use_adreno_bin_kernels(backend_ctx)) { + const char * kernel_bin = (const char *)backend_ctx->get_adreno_bin_kernel("gemm_moe_mxfp4_q8_1_dp4a_ila", &bin_size); + if (kernel_bin && bin_size > 0) { + cl_program prog = + build_program_from_binary(backend_ctx->context, backend_ctx->device, kernel_bin, CL_moe_compile_opts, bin_size); + + CL_CHECK((backend_ctx->kernel_gemm_moe_mxfp4_q8_1_dp4a_bin = clCreateKernel(prog, "kernel_gemm_moe_mxfp4_q8_1_dp4a_ila", &err), err)); + CL_CHECK(clReleaseProgram(prog)); + GGML_LOG_CONT("."); + } + } + } + // gemm_moe_q4_0_q8_1_dp4a (dp4a prefill GEMM) if (backend_ctx->has_integer_dot) { #ifdef GGML_OPENCL_EMBED_KERNELS @@ -4265,6 +4285,24 @@ static void load_cl_kernels(ggml_backend_opencl_context *backend_ctx) { GGML_LOG_CONT("."); } + // gemm_moe_q4_0_q8_1_dp4a_bin (dp4a prefill GEMM) + if (backend_ctx->has_integer_dot) { + size_t bin_size = 0; + backend_ctx->kernel_gemm_moe_q4_0_q8_1_dp4a_bin = nullptr; + + if (use_adreno_bin_kernels(backend_ctx)) { + const char * kernel_bin = (const char *)backend_ctx->get_adreno_bin_kernel("gemm_moe_q4_0_q8_1_dp4a_ila", &bin_size); + if (kernel_bin && bin_size > 0) { + cl_program prog = + build_program_from_binary(backend_ctx->context, backend_ctx->device, kernel_bin, CL_moe_compile_opts, bin_size); + + CL_CHECK((backend_ctx->kernel_gemm_moe_q4_0_q8_1_dp4a_bin = clCreateKernel(prog, "kernel_gemm_moe_q4_0_q8_1_dp4a_ila", &err), err)); + CL_CHECK(clReleaseProgram(prog)); + GGML_LOG_CONT("."); + } + } + } + // gemm_moe_q8_1_dp4a (generic dp4a MoE GEMM; MOE_QT=80 -> q8_0 expert variant) if (backend_ctx->has_integer_dot) { #ifdef GGML_OPENCL_EMBED_KERNELS @@ -21519,7 +21557,9 @@ static void ggml_cl_mul_mat_id(ggml_backend_t backend, const ggml_tensor * src0, // dot prod has to be available use_moe_dp4a = backend_ctx->has_integer_dot && use_moe_dp4a; // bin kernel takes precedence - use_moe_dp4a = use_moe_dp4a && backend_ctx->kernel_gemm_moe_q4_0_f32_ns_bin == nullptr; + if (backend_ctx->kernel_gemm_moe_q4_0_q8_1_dp4a_bin == nullptr) { + use_moe_dp4a = use_moe_dp4a && backend_ctx->kernel_gemm_moe_q4_0_f32_ns_bin == nullptr; + } cl_buffer_region region; region.origin = 0; @@ -21625,6 +21665,10 @@ static void ggml_cl_mul_mat_id(ggml_backend_t backend, const ggml_tensor * src0, // dp4a GEMM cl_kernel dk = backend_ctx->kernel_gemm_moe_q4_0_q8_1_dp4a; + if (backend_ctx->kernel_gemm_moe_q4_0_q8_1_dp4a_bin) { + dk = backend_ctx->kernel_gemm_moe_q4_0_q8_1_dp4a_bin; + } + int aidx = 0; CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_q4_0->q_img)); CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_q4_0->d)); @@ -23463,8 +23507,10 @@ static void ggml_cl_mul_mat_id(ggml_backend_t backend, const ggml_tensor * src0, : (backend_ctx->adreno_gen == ADRENO_GPU_GEN::X2E); // dot prod has to be available use_moe_dp4a = backend_ctx->has_integer_dot && use_moe_dp4a; - // bin kernel takes precedence - use_moe_dp4a = use_moe_dp4a && backend_ctx->kernel_gemm_moe_mxfp4_f32_ns_bin == nullptr; + // bin kernel takes precedence, dp4a bin kernel has higher priority than normal bin kernel + if (backend_ctx->kernel_gemm_moe_mxfp4_q8_1_dp4a_bin == nullptr) { + use_moe_dp4a = use_moe_dp4a && backend_ctx->kernel_gemm_moe_mxfp4_f32_ns_bin == nullptr; + } cl_buffer_region region; region.origin = 0; @@ -23573,6 +23619,10 @@ static void ggml_cl_mul_mat_id(ggml_backend_t backend, const ggml_tensor * src0, // dp4a GEMM cl_kernel dk = backend_ctx->kernel_gemm_moe_mxfp4_q8_1_dp4a; + if (backend_ctx->kernel_gemm_moe_mxfp4_q8_1_dp4a_bin) { + dk = backend_ctx->kernel_gemm_moe_mxfp4_q8_1_dp4a_bin; + } + int aidx = 0; CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_mxfp4->q_img)); CL_CHECK(clSetKernelArg(dk, aidx++, sizeof(cl_mem), &extra0_mxfp4->e)); diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index 8108e94c16c3..72e844aebfd5 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -767,6 +767,21 @@ static constexpr std::initializer_list> rms_norm_mul_rope_vie { 4, 0, 3 }, // set_rows->src[0] == view }; +static constexpr std::array lightning_indexer_k_types = { + GGML_TYPE_F32, + GGML_TYPE_F16, + GGML_TYPE_BF16, + GGML_TYPE_Q8_0, + GGML_TYPE_Q5_1, + GGML_TYPE_Q5_0, + GGML_TYPE_Q4_1, + GGML_TYPE_Q4_0, + GGML_TYPE_IQ4_NL, +}; + +static bool ggml_vk_lightning_indexer_k_type_supported(ggml_type type) { + return std::find(lightning_indexer_k_types.begin(), lightning_indexer_k_types.end(), type) != lightning_indexer_k_types.end(); +} struct vk_device_struct { std::recursive_mutex mutex; @@ -1068,6 +1083,7 @@ struct vk_device_struct { vk_pipeline pipeline_rwkv_wkv6_f32; vk_pipeline pipeline_rwkv_wkv7_f32; vk_pipeline pipeline_gated_linear_attn_f32; + vk_pipeline pipeline_lightning_indexer_f32[GGML_TYPE_COUNT]; // [size_idx][kda] where size_idx: 0=d16, 1=d32, 2=d64, 3=d128 vk_pipeline pipeline_gated_delta_net[4][2]; vk_pipeline pipeline_ssm_scan_f32_d128; @@ -1848,6 +1864,26 @@ struct vk_op_gated_linear_attn_push_constants { uint32_t H; float scale; }; +struct vk_op_lightning_indexer_push_constants { + uint32_t n_kv; + uint32_t n_heads; + uint32_t n_tokens; + uint32_t n_streams; + uint32_t n_masks; + uint32_t dispatch_x; + uint32_t q_nb1; + uint32_t q_nb2; + uint32_t q_nb3; + uint32_t k_nb2; + uint32_t k_nb3; + uint32_t w_nb1; + uint32_t w_nb3; + uint32_t m_nb1; + uint32_t m_nb3; + uint32_t d_nb1; + uint32_t d_nb3; +}; +static_assert(sizeof(vk_op_lightning_indexer_push_constants) <= 128); struct vk_op_gated_delta_net_push_constants { uint32_t H; uint32_t n_tokens; @@ -3904,11 +3940,16 @@ static vk_fa_pipeline_state get_fa_pipeline_state(const vk_device& device, const return vk_fa_pipeline_state{hsk, hsv, params.block_rows, params.block_cols, params.d_split, params.row_split, params.shmem_staging, params.path, params.workgroup_size, subgroup_size, aligned, f32acc, flags, params.limit_occupancy_shmem, k_type, v_type}; } +// Bytes per buffer block for the FaBlockBytesK/V spec constants. F32 is fed as +// a vec4 "block" of 4 floats, everything else uses its ggml block size. +static uint32_t fa_block_bytes(ggml_type t) { + if (t == GGML_TYPE_F32) { + return 16u; + } + return (uint32_t) ggml_type_size(t); +} + static std::vector get_fa_spec_constants(const vk_fa_pipeline_state& state) { - const auto fa_block_bytes = [](ggml_type t) -> uint32_t { - if (t == GGML_TYPE_F32) return 16u; - return (uint32_t) ggml_type_size(t); - }; return { /* 0 WorkGroupSize */ state.workgroup_size, /* 1 Br */ state.Br, @@ -5847,6 +5888,17 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_gated_linear_attn_f32, "gated_linear_attn_f32", gated_linear_attn_f32_len, gated_linear_attn_f32_data, "main", 6, sizeof(vk_op_gated_linear_attn_push_constants), {1, 1, 1}, {}, 1); + { + const bool li_subgroup = device->subgroup_arithmetic && device->subgroup_require_full_support; + const size_t li_len = li_subgroup ? lightning_indexer_subgroup_f32_len : lightning_indexer_f32_len; + const void * li_data = li_subgroup ? (const void *)lightning_indexer_subgroup_f32_data : (const void *)lightning_indexer_f32_data; + + for (ggml_type k_type : lightning_indexer_k_types) { + const std::string name = "lightning_indexer_" + std::string(ggml_type_name(k_type)) + "_k_f32"; + ggml_vk_create_pipeline(device, device->pipeline_lightning_indexer_f32[k_type], name.c_str(), li_len, li_data, "main", 5, sizeof(vk_op_lightning_indexer_push_constants), {1, 1, 1}, {(uint32_t)k_type, fa_block_bytes(k_type), device->subgroup_size}, 1, true, li_subgroup); + } + } + { const uint32_t gdn_sizes[] = {16, 32, 64, 128}; const char * gdn_names[][2] = { @@ -11697,6 +11749,12 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const return ctx->device->pipeline_gated_linear_attn_f32; } return nullptr; + case GGML_OP_LIGHTNING_INDEXER: + // only the k type selects a pipeline, the other types are fixed by ggml_lightning_indexer() + if (ggml_vk_lightning_indexer_k_type_supported(src1->type)) { + return ctx->device->pipeline_lightning_indexer_f32[src1->type]; + } + return nullptr; case GGML_OP_GATED_DELTA_NET: if (src0->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { const uint32_t S_v = dst->src[2]->ne[0]; @@ -12772,6 +12830,55 @@ static void ggml_vk_gated_linear_attn(ggml_backend_vk_context * ctx, vk_context& pc, { (uint32_t)(n_seqs * n_heads), 1, 1 }); } +static void ggml_vk_lightning_indexer(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) { + const ggml_tensor * q = dst->src[0]; + const ggml_tensor * k = dst->src[1]; + const ggml_tensor * w = dst->src[2]; + const ggml_tensor * m = dst->src[3]; + + vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, q, k, w, dst, dst->op); + GGML_ASSERT(pipeline != nullptr); + + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + + const uint32_t n_kv = k->ne[2]; + const uint32_t n_heads = q->ne[1]; + const uint32_t n_tokens = q->ne[2]; + const uint32_t n_streams = q->ne[3]; + const uint32_t n_masks = m->ne[3]; + + const uint32_t n_outputs = (uint32_t)(dst->ne[0] * dst->ne[1] * dst->ne[3]); + const uint32_t dispatch_x = std::min(n_outputs, ctx->device->properties.limits.maxComputeWorkGroupCount[0]); + const uint32_t dispatch_y = CEIL_DIV(n_outputs, dispatch_x); + + // q, w and dst are f32 and m is f16, so their strides are passed in elements; + // k may be quantized, so its strides stay in bytes + const uint32_t q_nb1 = q->nb[1] / sizeof(float); + const uint32_t q_nb2 = q->nb[2] / sizeof(float); + const uint32_t q_nb3 = q->nb[3] / sizeof(float); + const uint32_t k_nb2 = k->nb[2]; + const uint32_t k_nb3 = k->nb[3]; + const uint32_t w_nb1 = w->nb[1] / sizeof(float); + const uint32_t w_nb3 = w->nb[3] / sizeof(float); + const uint32_t m_nb1 = m->nb[1] / sizeof(ggml_fp16_t); + const uint32_t m_nb3 = m->nb[3] / sizeof(ggml_fp16_t); + const uint32_t d_nb1 = dst->nb[1] / sizeof(float); + const uint32_t d_nb3 = dst->nb[3] / sizeof(float); + + const vk_op_lightning_indexer_push_constants pc = { + n_kv, n_heads, n_tokens, n_streams, n_masks, dispatch_x, + q_nb1, q_nb2, q_nb3, + k_nb2, k_nb3, + w_nb1, w_nb3, + m_nb1, m_nb3, + d_nb1, d_nb3, + }; + + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + {ggml_vk_tensor_subbuffer(ctx, q), ggml_vk_tensor_subbuffer(ctx, k), ggml_vk_tensor_subbuffer(ctx, w), ggml_vk_tensor_subbuffer(ctx, m), ggml_vk_tensor_subbuffer(ctx, dst)}, + pc, {dispatch_x, dispatch_y, 1}); +} + static void ggml_vk_gated_delta_net(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) { const ggml_tensor * src_q = dst->src[0]; const ggml_tensor * src_v = dst->src[2]; @@ -15898,6 +16005,11 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr break; + case GGML_OP_LIGHTNING_INDEXER: + ggml_vk_lightning_indexer(ctx, compute_ctx, node); + + break; + case GGML_OP_GATED_DELTA_NET: ggml_vk_gated_delta_net(ctx, compute_ctx, node); @@ -18676,6 +18788,40 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm case GGML_OP_GATED_LINEAR_ATTN: // the shader block size is hardcoded to head_size 64 return op->src[0]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32 && op->src[0]->ne[0] == 64; + case GGML_OP_LIGHTNING_INDEXER: + { + const ggml_tensor * q = op->src[0]; + const ggml_tensor * k = op->src[1]; + const ggml_tensor * w = op->src[2]; + const ggml_tensor * m = op->src[3]; + + // the q/w/m types and the shape relationships between q, k, w, m and dst + // are already asserted in ggml_lightning_indexer() + if (!ggml_vk_lightning_indexer_k_type_supported(k->type) || !device->fp16) { + return false; + } + + // the shader block size is hardcoded to head size 128 + if (q->ne[0] != 128) { + return false; + } + + // the shader indexes the buffers by element stride, and is dispatched + // without allow_misalign + for (const ggml_tensor * t : {q, k, w, m, op}) { + if (t->nb[0] != ggml_type_size(t->type) || + (vk_tensor_offset(t) + t->view_offs) % device->properties.limits.minStorageBufferOffsetAlignment != 0) { + return false; + } + // the strides get scaled down from bytes, so the division must be exact + for (int i = 1; i < GGML_MAX_DIMS; ++i) { + if (t->nb[i] % ggml_type_size(t->type) != 0) { + return false; + } + } + } + return true; + } case GGML_OP_GATED_DELTA_NET: { const uint32_t S_v = op->src[2]->ne[0]; @@ -19685,6 +19831,8 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph * const float * op_params = (const float *)tensor->op_params; tensor_clone = ggml_gated_linear_attn(ggml_ctx, src_clone[0], src_clone[1], src_clone[2], src_clone[3], src_clone[4], op_params[0]); + } else if (tensor->op == GGML_OP_LIGHTNING_INDEXER) { + tensor_clone = ggml_lightning_indexer(ggml_ctx, src_clone[0], src_clone[1], src_clone[2], src_clone[3]); } else if (tensor->op == GGML_OP_GATED_DELTA_NET) { tensor_clone = ggml_gated_delta_net(ggml_ctx, src_clone[0], src_clone[1], src_clone[2], src_clone[3], src_clone[4], src_clone[5], diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/fa_types.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/fa_types.glsl new file mode 100644 index 000000000000..6f414ded1230 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/fa_types.glsl @@ -0,0 +1,55 @@ +#if !defined(GGML_FA_TYPES_COMP) +#define GGML_FA_TYPES_COMP + +// FaTypeK / FaTypeV spec constant values. These mirror enum ggml_type so the +// host can pass the type directly. Keep in sync with ggml.h. +#define FA_TYPE_F32 0u +#define FA_TYPE_F16 1u +#define FA_TYPE_Q4_0 2u +#define FA_TYPE_Q4_1 3u +#define FA_TYPE_Q5_0 6u +#define FA_TYPE_Q5_1 7u +#define FA_TYPE_Q8_0 8u +#define FA_TYPE_IQ4_NL 20u +#define FA_TYPE_BF16 30u + +// Number of matrix elements per buffer block, derived from the K/V type spec +// constant. F32 is treated as a vec4 "block" of 4 floats. F16 uses block size 1 +// and bypasses the dequant path entirely. Quants follow their ggml block sizes. +uint fa_block_elems(uint ty) { + switch (ty) { + case FA_TYPE_F32: return 4u; + case FA_TYPE_F16: return 1u; + case FA_TYPE_Q4_0: return uint(QUANT_K_Q4_0); + case FA_TYPE_Q4_1: return uint(QUANT_K_Q4_1); + case FA_TYPE_Q5_0: return uint(QUANT_K_Q5_0); + case FA_TYPE_Q5_1: return uint(QUANT_K_Q5_1); + case FA_TYPE_Q8_0: return uint(QUANT_K_Q8_0); + case FA_TYPE_IQ4_NL: return uint(QUANT_K_IQ4_NL); + case FA_TYPE_BF16: return 1u; + default: return 1u; + } +} + +// QUANT_R_MMQ for FA-eligible K types. Q4_*/Q5_* store two nibbles per byte +// (R==2); Q8_0 stores one byte per element (R==1). Used to derive the number +// of int32s per 32-element block on the MMQ K path: ints_per_block == 8 / R. +uint fa_quant_r_mmq(uint ty) { + switch (ty) { + case FA_TYPE_Q4_0: return uint(QUANT_R_Q4_0); + case FA_TYPE_Q4_1: return uint(QUANT_R_Q4_1); + case FA_TYPE_Q5_0: return uint(QUANT_R_Q5_0); + case FA_TYPE_Q5_1: return uint(QUANT_R_Q5_1); + case FA_TYPE_Q8_0: return uint(QUANT_R_Q8_0); + default: return 1u; + } +} + +bool fa_type_needs_shmem(uint ty) { + switch (ty) { + case FA_TYPE_IQ4_NL: return true; + default: return false; + } +} + +#endif // !defined(GGML_FA_TYPES_COMP) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl index 3c64f91dad36..0ce4503a8847 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl @@ -88,17 +88,7 @@ layout (binding = 6) readonly buffer MO {uint32_t data_mask_opt[];}; #define BINDING_IDX_K 0 #define BINDING_IDX_V 1 -// FaTypeK / FaTypeV spec constant values. These mirror enum ggml_type so the -// host can pass the type directly. Keep in sync with ggml.h. -#define FA_TYPE_F32 0u -#define FA_TYPE_F16 1u -#define FA_TYPE_Q4_0 2u -#define FA_TYPE_Q4_1 3u -#define FA_TYPE_Q5_0 6u -#define FA_TYPE_Q5_1 7u -#define FA_TYPE_Q8_0 8u -#define FA_TYPE_IQ4_NL 20u -#define FA_TYPE_BF16 30u +#include "fa_types.glsl" #if defined(BFLOAT16) #define O_TYPE float @@ -108,45 +98,6 @@ layout (binding = 6) readonly buffer MO {uint32_t data_mask_opt[];}; #define O_TYPEV4 FLOAT_TYPEV4 #endif -// Number of matrix elements per buffer block, derived from the K/V type spec -// constant. F32 is treated as a vec4 "block" of 4 floats. F16 uses block size 1 -// and bypasses the dequant path entirely. Quants follow their ggml block sizes. -uint fa_block_elems(uint ty) { - switch (ty) { - case FA_TYPE_F32: return 4u; - case FA_TYPE_F16: return 1u; - case FA_TYPE_Q4_0: return uint(QUANT_K_Q4_0); - case FA_TYPE_Q4_1: return uint(QUANT_K_Q4_1); - case FA_TYPE_Q5_0: return uint(QUANT_K_Q5_0); - case FA_TYPE_Q5_1: return uint(QUANT_K_Q5_1); - case FA_TYPE_Q8_0: return uint(QUANT_K_Q8_0); - case FA_TYPE_IQ4_NL: return uint(QUANT_K_IQ4_NL); - case FA_TYPE_BF16: return 1u; - default: return 1u; - } -} - -// QUANT_R_MMQ for FA-eligible K types. Q4_*/Q5_* store two nibbles per byte -// (R==2); Q8_0 stores one byte per element (R==1). Used to derive the number -// of int32s per 32-element block on the MMQ K path: ints_per_block == 8 / R. -uint fa_quant_r_mmq(uint ty) { - switch (ty) { - case FA_TYPE_Q4_0: return uint(QUANT_R_Q4_0); - case FA_TYPE_Q4_1: return uint(QUANT_R_Q4_1); - case FA_TYPE_Q5_0: return uint(QUANT_R_Q5_0); - case FA_TYPE_Q5_1: return uint(QUANT_R_Q5_1); - case FA_TYPE_Q8_0: return uint(QUANT_R_Q8_0); - default: return 1u; - } -} - -bool fa_type_needs_shmem(uint ty) { - switch (ty) { - case FA_TYPE_IQ4_NL: return true; - default: return false; - } -} - // These can't be `const` globals because GLSL forbids function calls in global // const initializers, even when the spec constants would let the driver fold // them. Macros expand at the use site and fold after specialization. diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer.comp b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer.comp new file mode 100644 index 000000000000..ba76ec72ca6c --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer.comp @@ -0,0 +1,151 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_KHR_shader_subgroup_basic : enable +#if USE_SUBGROUP_ADD +#extension GL_KHR_shader_subgroup_arithmetic : enable +#endif + +#define BINDING_IDX_K 0u + +#include "types.glsl" +#include "fa_types.glsl" +#define FaTypeV FA_TYPE_F32 + +layout(constant_id = 0) const uint FaTypeK = FA_TYPE_F32; +layout(constant_id = 1) const uint FaBlockBytesK = 4; +layout(constant_id = 2) const uint SUBGROUP_SIZE = 32; + +#include "flash_attn_dequant.glsl" + +// one workgroup computes one output element, one invocation per head element +#define HEAD_SIZE 128 + +layout(local_size_x = HEAD_SIZE, local_size_y = 1, local_size_z = 1) in; + +layout(binding = 0) readonly buffer QBuf { float q[]; }; +layout(binding = 1) readonly buffer KBufF16 { float16_t k_f16[]; }; +layout(binding = 1) readonly buffer KBufF32 { float k_f32[]; }; +layout(binding = 1) readonly buffer KBufBF16 { uint16_t k_bf16[]; }; +layout(binding = 2) readonly buffer WBuf { float weights[]; }; +layout(binding = 3) readonly buffer MBuf { float16_t mask[]; }; +layout(binding = 4) writeonly buffer DstBuf { float dst[]; }; + +layout(push_constant) uniform PushConstants { + uint n_kv; + uint n_heads; + uint n_tokens; + uint n_streams; + uint n_masks; + uint dispatch_x; + uint q_nb1; + uint q_nb2; + uint q_nb3; + uint k_nb2; + uint k_nb3; + uint w_nb1; + uint w_nb3; + uint m_nb1; + uint m_nb3; + uint d_nb1; + uint d_nb3; +}; + +shared float k_row[HEAD_SIZE]; + +#if USE_SUBGROUP_ADD +shared float sg_partials[HEAD_SIZE / SUBGROUP_SIZE]; +#else +shared float partials[HEAD_SIZE]; +#endif + +void main() { + const uint tid = gl_LocalInvocationID.x; + const uint output_idx = gl_WorkGroupID.y * dispatch_x + gl_WorkGroupID.x; + const uint n_outputs = n_kv * n_tokens * n_streams; + + if (fa_type_needs_shmem(FaTypeK)) { + init_iq_shmem(gl_WorkGroupSize); + } + + if (output_idx >= n_outputs) { + return; + } + + const uint ik = output_idx % n_kv; + const uint ts = output_idx / n_kv; + const uint t = ts % n_tokens; + const uint s = ts / n_tokens; + const uint k_offset = ik * k_nb2 + s * k_nb3; + + // k strides come in as bytes, so scale them down to the view being indexed + const uint k_block_elems = fa_block_elems(FaTypeK); + const uint k_elem_bytes = FaBlockBytesK / k_block_elems; + + if (FaTypeK == FA_TYPE_F16) { + k_row[tid] = float(k_f16[k_offset / k_elem_bytes + tid]); + } else if (FaTypeK == FA_TYPE_F32) { + k_row[tid] = k_f32[k_offset / k_elem_bytes + tid]; + } else if (FaTypeK == FA_TYPE_BF16) { + k_row[tid] = bf16_to_fp32(uint(k_bf16[k_offset / k_elem_bytes + tid])); + } else if (4 * tid < HEAD_SIZE) { + const uint coord = 4 * tid; + const uint ib = coord / k_block_elems; + const uint iqs = coord % k_block_elems; + const vec4 values = dequantize4(ib, iqs, k_offset / FaBlockBytesK, BINDING_IDX_K); + k_row[coord + 0] = values.x; + k_row[coord + 1] = values.y; + k_row[coord + 2] = values.z; + k_row[coord + 3] = values.w; + } + barrier(); + + const float k_val = k_row[tid]; + + float score = 0.0; + for (uint h = 0; h < n_heads; ++h) { + const float prod = q[h * q_nb1 + t * q_nb2 + s * q_nb3 + tid] * k_val; + +#if USE_SUBGROUP_ADD + const float sg_sum = subgroupAdd(prod); + if (gl_SubgroupInvocationID == 0) { + sg_partials[gl_SubgroupID] = sg_sum; + } + barrier(); + + if (tid == 0) { + float sum = 0.0; + [[unroll]] for (uint i = 0; i < HEAD_SIZE / SUBGROUP_SIZE; ++i) { + sum += sg_partials[i]; + } + score += max(sum, 0.0) * weights[h + t * w_nb1 + s * w_nb3]; + } + // the reads above must complete before the next iteration overwrites sg_partials + barrier(); +#else + partials[tid] = prod; + barrier(); + + [[unroll]] for (uint stride = HEAD_SIZE / 2; stride > 0; stride >>= 1) { + if (tid < stride) { + partials[tid] += partials[tid + stride]; + } + barrier(); + } + + if (tid == 0) { + score += max(partials[0], 0.0) * weights[h + t * w_nb1 + s * w_nb3]; + } + // the read of partials[0] above must complete before the next iteration + // overwrites partials[tid] + barrier(); +#endif + } + + if (tid == 0) { + const uint mask_offset = ik + t * m_nb1 + (s % n_masks) * m_nb3; + dst[ik + t * d_nb1 + s * d_nb3] = score + float(mask[mask_offset]); + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp index dbb99782cf73..0da943da9563 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -1069,6 +1069,12 @@ void process_shaders() { string_to_spv("gated_linear_attn_f32", "gla.comp", merge_maps(base_dict, {{"A_TYPE", "float"}})); + // Compile IQ4_NL support in so its shared LUT is available when K uses it. + // K quant type is selected at runtime via the FaTypeK spec constant. + std::map li_dict = {{"FLOAT_TYPE", "float"}, {"FLOAT_TYPEV4", "vec4"}, {"DATA_A_IQ4_NL", "1"}}; + string_to_spv("lightning_indexer_f32", "lightning_indexer.comp", li_dict); + string_to_spv("lightning_indexer_subgroup_f32", "lightning_indexer.comp", merge_maps(li_dict, {{"USE_SUBGROUP_ADD", "1"}})); + string_to_spv("rwkv_wkv7_f32", "wkv7.comp", merge_maps(base_dict, {{"A_TYPE", "float"}})); string_to_spv("gated_delta_net_f32", "gated_delta_net.comp", merge_maps(base_dict, {{"FLOAT_TYPE", "float"}, {"USE_SUBGROUP_ADD", "1"}, {"USE_SUBGROUP_CLUSTERED", "1"}})); diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index b35e9abc389b..fb123acfc2da 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -162,7 +162,12 @@ class LLM: TARGET_LAYERS = "{arch}.target_layers" TARGET_HIDDEN_SIZE = "{arch}.target_hidden_size" BLOCK_SIZE = "{arch}.block_size" + CONV_KERNEL_SIZE = "{arch}.conv_kernel_size" + CONV_GROUP_SIZE = "{arch}.conv_group_size" + SELECTOR_RANK = "{arch}.selector_rank" + SELECTOR_TOP_K = "{arch}.selector_top_k" SAMPLE_FROM_ANCHOR = "{arch}.sample_from_anchor" + HAS_CONFIDENCE_HEAD = "{arch}.has_confidence_head" NORM_BEFORE_RESIDUAL = "{arch}.norm_before_residual" NORM_BEFORE_FC = "{arch}.norm_before_fc" @@ -225,6 +230,19 @@ class HyperConnection: COUNT = "{arch}.hyper_connection.count" SINKHORN_ITERATIONS = "{arch}.hyper_connection.sinkhorn_iterations" EPSILON = "{arch}.hyper_connection.epsilon" + # absent means the mix projection is full rank (DeepSeek-V4 behaviour) + LOW_RANK = "{arch}.hyper_connection.low_rank" + + class PerLayerEmbedding: + LAYERS = "{arch}.ple.layers" + NGRAM_SIZE = "{arch}.ple.ngram_size" + HEADS_PER_NGRAM = "{arch}.ple.heads_per_ngram" + CONV_KERNEL = "{arch}.ple.conv_kernel" + LAYER_MULTIPLIERS = "{arch}.ple.layer_multipliers" + HEAD_OFFSETS = "{arch}.ple.head_offsets" + HEAD_VOCAB_SIZES = "{arch}.ple.head_vocab_sizes" + EOS_TOKEN_ID = "{arch}.ple.eos_token_id" + IMAGE_TOKEN_ID = "{arch}.ple.image_token_id" class Rope: DIMENSION_COUNT = "{arch}.rope.dimension_count" @@ -494,6 +512,7 @@ class MODEL_ARCH(IntEnum): QWEN3VLMOE = auto() QWEN35 = auto() QWEN35MOE = auto() + QWEN4EXP = auto() PHI2 = auto() PHI3 = auto() PHIMOE = auto() @@ -637,6 +656,9 @@ class MODEL_TENSOR(IntEnum): HC_HEAD_FN = auto() HC_HEAD_BASE = auto() HC_HEAD_SCALE = auto() + HC_HEAD_NORM = auto() # qwen4exp + HC_HEAD_DOWN = auto() # qwen4exp + HC_HEAD_UP = auto() # qwen4exp ROPE_FREQS = auto() ROPE_FACTORS_LONG = auto() ROPE_FACTORS_SHORT = auto() @@ -781,6 +803,20 @@ class MODEL_TENSOR(IntEnum): HC_FFN_FN = auto() HC_FFN_BASE = auto() HC_FFN_SCALE = auto() + HC_ATTN_NORM = auto() # qwen4exp + HC_ATTN_DOWN = auto() # qwen4exp + HC_ATTN_UP = auto() # qwen4exp + HC_ATTN_INJECT = auto() # qwen4exp + HC_FFN_NORM = auto() # qwen4exp + HC_FFN_DOWN = auto() # qwen4exp + HC_FFN_UP = auto() # qwen4exp + HC_FFN_INJECT = auto() # qwen4exp + PLE_KEY = auto() # qwen4exp + PLE_VALUE = auto() # qwen4exp + PLE_NORM_KEY = auto() # qwen4exp + PLE_NORM_QUERY = auto() # qwen4exp + PLE_NORM_CONV = auto() # qwen4exp + PLE_CONV1D = auto() # qwen4exp ATTN_COMPRESSOR_WKV = auto() ATTN_COMPRESSOR_WGATE = auto() ATTN_COMPRESSOR_APE = auto() @@ -1155,6 +1191,13 @@ class MODEL_TENSOR(IntEnum): DSPARK_MARKOV_W1 = auto() # markov head: prev-token embed DSPARK_MARKOV_W2 = auto() # markov head: bias projection DSPARK_CONF_PROJ = auto() # confidence head + DFLASH_ATTN_CONV_BASE = auto() + DFLASH_ATTN_CONV_PROJ = auto() + DFLASH_FFN_CONV_BASE = auto() + DFLASH_FFN_CONV_PROJ = auto() + DFLASH_SELECTOR_PREV = auto() + DFLASH_SELECTOR_NEXT = auto() + DFLASH_SELECTOR_HIDDEN = auto() # lfm2 audio A_ENC_NORM_CONV = auto() A_ENC_LINEAR_POS = auto() @@ -1226,6 +1269,7 @@ class MODEL_TENSOR(IntEnum): MODEL_ARCH.QWEN3VLMOE: "qwen3vlmoe", MODEL_ARCH.QWEN35: "qwen35", MODEL_ARCH.QWEN35MOE: "qwen35moe", + MODEL_ARCH.QWEN4EXP: "qwen4exp", MODEL_ARCH.PHI2: "phi2", MODEL_ARCH.PHI3: "phi3", MODEL_ARCH.PHIMOE: "phimoe", @@ -1368,6 +1412,9 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.HC_HEAD_FN: "output_hc_fn", MODEL_TENSOR.HC_HEAD_BASE: "output_hc_base", MODEL_TENSOR.HC_HEAD_SCALE: "output_hc_scale", + MODEL_TENSOR.HC_HEAD_NORM: "output_hc_norm", # qwen4exp + MODEL_TENSOR.HC_HEAD_DOWN: "output_hc_down", # qwen4exp + MODEL_TENSOR.HC_HEAD_UP: "output_hc_up", # qwen4exp MODEL_TENSOR.ROPE_FREQS: "rope_freqs", MODEL_TENSOR.ROPE_FACTORS_LONG: "rope_factors_long", MODEL_TENSOR.ROPE_FACTORS_SHORT: "rope_factors_short", @@ -1512,6 +1559,20 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.HC_FFN_FN: "blk.{bid}.hc_ffn_fn", MODEL_TENSOR.HC_FFN_BASE: "blk.{bid}.hc_ffn_base", MODEL_TENSOR.HC_FFN_SCALE: "blk.{bid}.hc_ffn_scale", + MODEL_TENSOR.HC_ATTN_NORM: "blk.{bid}.hc_attn_norm", # qwen4exp + MODEL_TENSOR.HC_ATTN_DOWN: "blk.{bid}.hc_attn_down", # qwen4exp + MODEL_TENSOR.HC_ATTN_UP: "blk.{bid}.hc_attn_up", # qwen4exp + MODEL_TENSOR.HC_ATTN_INJECT: "blk.{bid}.hc_attn_inject", # qwen4exp + MODEL_TENSOR.HC_FFN_NORM: "blk.{bid}.hc_ffn_norm", # qwen4exp + MODEL_TENSOR.HC_FFN_DOWN: "blk.{bid}.hc_ffn_down", # qwen4exp + MODEL_TENSOR.HC_FFN_UP: "blk.{bid}.hc_ffn_up", # qwen4exp + MODEL_TENSOR.HC_FFN_INJECT: "blk.{bid}.hc_ffn_inject", # qwen4exp + MODEL_TENSOR.PLE_KEY: "blk.{bid}.ple_key", # qwen4exp + MODEL_TENSOR.PLE_VALUE: "blk.{bid}.ple_value", # qwen4exp + MODEL_TENSOR.PLE_NORM_KEY: "blk.{bid}.ple_norm_key", # qwen4exp + MODEL_TENSOR.PLE_NORM_QUERY: "blk.{bid}.ple_norm_query", # qwen4exp + MODEL_TENSOR.PLE_NORM_CONV: "blk.{bid}.ple_norm_conv", # qwen4exp + MODEL_TENSOR.PLE_CONV1D: "blk.{bid}.ple_conv1d", # qwen4exp MODEL_TENSOR.ATTN_COMPRESSOR_WKV: "blk.{bid}.attn_compressor_kv", MODEL_TENSOR.ATTN_COMPRESSOR_WGATE: "blk.{bid}.attn_compressor_gate", MODEL_TENSOR.ATTN_COMPRESSOR_APE: "blk.{bid}.attn_compressor_ape", @@ -1912,6 +1973,13 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.DSPARK_MARKOV_W1: "markov_w1", MODEL_TENSOR.DSPARK_MARKOV_W2: "markov_w2", MODEL_TENSOR.DSPARK_CONF_PROJ: "conf_proj", + MODEL_TENSOR.DFLASH_ATTN_CONV_BASE: "blk.{bid}.attn_conv_base", + MODEL_TENSOR.DFLASH_ATTN_CONV_PROJ: "blk.{bid}.attn_conv_proj", + MODEL_TENSOR.DFLASH_FFN_CONV_BASE: "blk.{bid}.ffn_conv_base", + MODEL_TENSOR.DFLASH_FFN_CONV_PROJ: "blk.{bid}.ffn_conv_proj", + MODEL_TENSOR.DFLASH_SELECTOR_PREV: "selector_predecessor", + MODEL_TENSOR.DFLASH_SELECTOR_NEXT: "selector_successor", + MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN: "selector_hidden", MODEL_TENSOR.D2T: "d2t", } @@ -2812,6 +2880,58 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM, ], + MODEL_ARCH.QWEN4EXP: [ + MODEL_TENSOR.TOKEN_EMBD, + MODEL_TENSOR.OUTPUT, + # no OUTPUT_NORM / ATTN_NORM / ATTN_POST_NORM: hyper-connections replace every layer norm + MODEL_TENSOR.HC_HEAD_NORM, + MODEL_TENSOR.HC_HEAD_DOWN, + MODEL_TENSOR.HC_HEAD_UP, + MODEL_TENSOR.HC_ATTN_NORM, + MODEL_TENSOR.HC_ATTN_DOWN, + MODEL_TENSOR.HC_ATTN_UP, + MODEL_TENSOR.HC_ATTN_INJECT, + MODEL_TENSOR.HC_FFN_NORM, + MODEL_TENSOR.HC_FFN_DOWN, + MODEL_TENSOR.HC_FFN_UP, + MODEL_TENSOR.HC_FFN_INJECT, + # full attention layers: ATTN_Q holds [q|gate] interleaved per head + MODEL_TENSOR.ATTN_Q, + MODEL_TENSOR.ATTN_Q_NORM, + MODEL_TENSOR.ATTN_K, + MODEL_TENSOR.ATTN_K_NORM, + MODEL_TENSOR.ATTN_V, + MODEL_TENSOR.ATTN_OUT, + MODEL_TENSOR.INDEXER_Q_PROJ, + MODEL_TENSOR.INDEXER_K_PROJ, + MODEL_TENSOR.INDEXER_Q_NORM, + MODEL_TENSOR.INDEXER_K_NORM, + MODEL_TENSOR.ATTN_QKV, + MODEL_TENSOR.ATTN_GATE, + MODEL_TENSOR.SSM_A, + MODEL_TENSOR.SSM_CONV1D, + MODEL_TENSOR.SSM_DT, + MODEL_TENSOR.SSM_NORM, + MODEL_TENSOR.SSM_BETA, + MODEL_TENSOR.SSM_ALPHA, + MODEL_TENSOR.SSM_OUT, + MODEL_TENSOR.FFN_GATE_INP, + MODEL_TENSOR.FFN_GATE_INP_SHEXP, + MODEL_TENSOR.FFN_UP_SHEXP, + MODEL_TENSOR.FFN_DOWN_SHEXP, + MODEL_TENSOR.FFN_GATE_SHEXP, + MODEL_TENSOR.FFN_DOWN_EXP, + MODEL_TENSOR.FFN_UP_EXP, + MODEL_TENSOR.FFN_GATE_EXP, + MODEL_TENSOR.FFN_GATE_UP_EXP, + MODEL_TENSOR.PER_LAYER_TOKEN_EMBD, + MODEL_TENSOR.PLE_KEY, + MODEL_TENSOR.PLE_VALUE, + MODEL_TENSOR.PLE_NORM_KEY, + MODEL_TENSOR.PLE_NORM_QUERY, + MODEL_TENSOR.PLE_NORM_CONV, + MODEL_TENSOR.PLE_CONV1D, + ], MODEL_ARCH.PLAMO: [ MODEL_TENSOR.TOKEN_EMBD, MODEL_TENSOR.OUTPUT_NORM, @@ -5002,6 +5122,13 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.DSPARK_MARKOV_W1, MODEL_TENSOR.DSPARK_MARKOV_W2, MODEL_TENSOR.DSPARK_CONF_PROJ, + MODEL_TENSOR.DFLASH_ATTN_CONV_BASE, + MODEL_TENSOR.DFLASH_ATTN_CONV_PROJ, + MODEL_TENSOR.DFLASH_FFN_CONV_BASE, + MODEL_TENSOR.DFLASH_FFN_CONV_PROJ, + MODEL_TENSOR.DFLASH_SELECTOR_PREV, + MODEL_TENSOR.DFLASH_SELECTOR_NEXT, + MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN, ], MODEL_ARCH.MISTRAL4: [ MODEL_TENSOR.TOKEN_EMBD, diff --git a/gguf-py/gguf/gguf_writer.py b/gguf-py/gguf/gguf_writer.py index d8a96a27bdd5..1f309ad2eafd 100644 --- a/gguf-py/gguf/gguf_writer.py +++ b/gguf-py/gguf/gguf_writer.py @@ -993,9 +993,24 @@ def add_sliding_window(self, value: int) -> None: def add_block_size(self, value: int) -> None: self.add_uint32(Keys.LLM.BLOCK_SIZE.format(arch=self.arch), value) + def add_conv_kernel_size(self, value: int) -> None: + self.add_uint32(Keys.LLM.CONV_KERNEL_SIZE.format(arch=self.arch), value) + + def add_conv_group_size(self, value: int) -> None: + self.add_uint32(Keys.LLM.CONV_GROUP_SIZE.format(arch=self.arch), value) + + def add_selector_rank(self, value: int) -> None: + self.add_uint32(Keys.LLM.SELECTOR_RANK.format(arch=self.arch), value) + + def add_selector_top_k(self, value: int) -> None: + self.add_uint32(Keys.LLM.SELECTOR_TOP_K.format(arch=self.arch), value) + def add_sample_from_anchor(self, value: bool) -> None: self.add_bool(Keys.LLM.SAMPLE_FROM_ANCHOR.format(arch=self.arch), value) + def add_has_confidence_head(self, value: bool) -> None: + self.add_bool(Keys.LLM.HAS_CONFIDENCE_HEAD.format(arch=self.arch), value) + def add_target_layers(self, value: Sequence[int]) -> None: self.add_array(Keys.LLM.TARGET_LAYERS.format(arch=self.arch), value) @@ -1029,6 +1044,40 @@ def add_hyper_connection_sinkhorn_iterations(self, count: int) -> None: def add_hyper_connection_epsilon(self, value: float) -> None: self.add_float32(Keys.HyperConnection.EPSILON.format(arch=self.arch), value) + def add_hyper_connection_low_rank(self, value: int) -> None: + self.add_uint32(Keys.HyperConnection.LOW_RANK.format(arch=self.arch), value) + + def add_ple_layers(self, values: Sequence[int]) -> None: + self.add_array(Keys.PerLayerEmbedding.LAYERS.format(arch=self.arch), values) + + def add_ple_ngram_size(self, value: int) -> None: + self.add_uint32(Keys.PerLayerEmbedding.NGRAM_SIZE.format(arch=self.arch), value) + + def add_ple_heads_per_ngram(self, value: int) -> None: + self.add_uint32(Keys.PerLayerEmbedding.HEADS_PER_NGRAM.format(arch=self.arch), value) + + def add_ple_conv_kernel(self, value: int) -> None: + self.add_uint32(Keys.PerLayerEmbedding.CONV_KERNEL.format(arch=self.arch), value) + + # multipliers reach ~2.4e13; default INT32 inference would truncate them + def _add_u64_array(self, key: str, values: Sequence[int]) -> None: + self.add_key_value(key, list(values), GGUFValueType.ARRAY, GGUFValueType.UINT64) + + def add_ple_layer_multipliers(self, values: Sequence[int]) -> None: + self._add_u64_array(Keys.PerLayerEmbedding.LAYER_MULTIPLIERS.format(arch=self.arch), values) + + def add_ple_head_offsets(self, values: Sequence[int]) -> None: + self._add_u64_array(Keys.PerLayerEmbedding.HEAD_OFFSETS.format(arch=self.arch), values) + + def add_ple_head_vocab_sizes(self, values: Sequence[int]) -> None: + self._add_u64_array(Keys.PerLayerEmbedding.HEAD_VOCAB_SIZES.format(arch=self.arch), values) + + def add_ple_eos_token_id(self, value: int) -> None: + self.add_uint32(Keys.PerLayerEmbedding.EOS_TOKEN_ID.format(arch=self.arch), value) + + def add_ple_image_token_id(self, value: int) -> None: + self.add_uint32(Keys.PerLayerEmbedding.IMAGE_TOKEN_ID.format(arch=self.arch), value) + def add_attention_scale(self, value: float) -> None: self.add_float32(Keys.Attention.SCALE.format(arch=self.arch), value) diff --git a/gguf-py/gguf/lazy.py b/gguf-py/gguf/lazy.py index acbc79258a31..6a0aee881107 100644 --- a/gguf-py/gguf/lazy.py +++ b/gguf-py/gguf/lazy.py @@ -226,3 +226,64 @@ def tofile(self, *args, **kwargs): return eager.tofile(*args, **kwargs) # TODO: __array_function__ + + +# Tensor written to file one row-chunk at a time +class LazyChunkedTensor: + + def __init__( + self, chunks: list[Callable[[], np.ndarray]], shape: tuple[int, ...], dtype: DTypeLike, + qtype: Any = None, byteswap: bool = False, + ): + self._chunks = chunks + self._qtype = qtype + self._byteswap = byteswap + self.shape = tuple(shape) + self.dtype = np.dtype(dtype) + + @property + def nbytes(self) -> int: + n = self.dtype.itemsize + for d in self.shape: + n *= d + return n + + def numpy(self) -> LazyChunkedTensor: + return self + + def quantize(self, qtype: Any) -> LazyChunkedTensor: + from .constants import GGMLQuantizationType + from .quants import QuantError, quant_shape_to_byte_shape + + if qtype == GGMLQuantizationType.F32: + shape, dtype = self.shape, np.dtype(np.float32) + elif qtype == GGMLQuantizationType.F16: + shape, dtype = self.shape, np.dtype(np.float16) + else: + try: + shape, dtype = quant_shape_to_byte_shape(self.shape, qtype), np.dtype(np.uint8) + except ValueError as e: + # raised here and not per chunk, so callers can still fall back to F16 + raise QuantError(str(e)) from e + return LazyChunkedTensor(self._chunks, shape, dtype, qtype, self._byteswap) + + def byteswap(self, inplace: bool = False) -> LazyChunkedTensor: + if inplace: + raise NotImplementedError("a chunked tensor cannot be byteswapped in place") + return LazyChunkedTensor(self._chunks, self.shape, self.dtype, self._qtype, not self._byteswap) + + def tofile(self, *args, **kwargs) -> None: + from .quants import quantize + + written = 0 + for load_chunk in self._chunks: + chunk = load_chunk() + if self._qtype is not None: + # exact only because chunks split on rows, and blocks never cross one + chunk = quantize(chunk, self._qtype) + if self._byteswap: + chunk = chunk.byteswap(inplace=False) + chunk.tofile(*args, **kwargs) + written += chunk.nbytes + del chunk + assert written == self.nbytes, f"chunked tensor wrote {written} bytes, expected {self.nbytes}" diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index d83b25e967e1..103447d0d993 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -1357,6 +1357,34 @@ class TensorNameMap: "model.confidence_head.proj", # dspark ), + MODEL_TENSOR.DFLASH_ATTN_CONV_BASE: ( + "model.layers.{bid}.attention_conv.base_kernel", + ), + + MODEL_TENSOR.DFLASH_ATTN_CONV_PROJ: ( + "model.layers.{bid}.attention_conv.kernel_projection", + ), + + MODEL_TENSOR.DFLASH_FFN_CONV_BASE: ( + "model.layers.{bid}.mlp_conv.base_kernel", + ), + + MODEL_TENSOR.DFLASH_FFN_CONV_PROJ: ( + "model.layers.{bid}.mlp_conv.kernel_projection", + ), + + MODEL_TENSOR.DFLASH_SELECTOR_PREV: ( + "model.candidate_selector.predecessor_codebook", + ), + + MODEL_TENSOR.DFLASH_SELECTOR_NEXT: ( + "model.candidate_selector.successor_codebook", + ), + + MODEL_TENSOR.DFLASH_SELECTOR_HIDDEN: ( + "model.candidate_selector.hidden_projection", + ), + MODEL_TENSOR.CLS: ( "classifier", # jina "classifier.dense", # roberta @@ -2729,6 +2757,65 @@ class TensorNameMap: "model.layers.{bid}.post_attention_layernorm", ), }, + MODEL_ARCH.QWEN4EXP: { + MODEL_TENSOR.HC_ATTN_NORM: ( + "model.layers.{bid}.attn_hyper_connection.hc_norm", + ), + MODEL_TENSOR.HC_ATTN_DOWN: ( + "model.layers.{bid}.attn_hyper_connection.input_mix_weight_down", + ), + MODEL_TENSOR.HC_ATTN_UP: ( + "model.layers.{bid}.attn_hyper_connection.input_mix_weight_up", + ), + MODEL_TENSOR.HC_ATTN_INJECT: ( + "model.layers.{bid}.attn_hyper_connection.block_inject_weight", + ), + MODEL_TENSOR.HC_FFN_NORM: ( + "model.layers.{bid}.mlp_hyper_connection.hc_norm", + ), + MODEL_TENSOR.HC_FFN_DOWN: ( + "model.layers.{bid}.mlp_hyper_connection.input_mix_weight_down", + ), + MODEL_TENSOR.HC_FFN_UP: ( + "model.layers.{bid}.mlp_hyper_connection.input_mix_weight_up", + ), + MODEL_TENSOR.HC_FFN_INJECT: ( + "model.layers.{bid}.mlp_hyper_connection.block_inject_weight", + ), + MODEL_TENSOR.HC_HEAD_NORM: ( + "model.hyper_connection_mixer.hc_norm", + ), + MODEL_TENSOR.HC_HEAD_DOWN: ( + "model.hyper_connection_mixer.input_mix_weight_down", + ), + MODEL_TENSOR.HC_HEAD_UP: ( + "model.hyper_connection_mixer.input_mix_weight_up", + ), + MODEL_TENSOR.INDEXER_Q_NORM: ( + "model.layers.{bid}.self_attn.indexer.q_layernorm", + ), + MODEL_TENSOR.INDEXER_K_NORM: ( + "model.layers.{bid}.self_attn.indexer.k_layernorm", + ), + MODEL_TENSOR.PLE_KEY: ( + "model.layers.{bid}.ple.key_proj", + ), + MODEL_TENSOR.PLE_VALUE: ( + "model.layers.{bid}.ple.value_proj", + ), + MODEL_TENSOR.PLE_NORM_KEY: ( + "model.layers.{bid}.ple.norm_key", + ), + MODEL_TENSOR.PLE_NORM_QUERY: ( + "model.layers.{bid}.ple.norm_query", + ), + MODEL_TENSOR.PLE_NORM_CONV: ( + "model.layers.{bid}.ple.norm_conv", + ), + MODEL_TENSOR.PLE_CONV1D: ( + "model.layers.{bid}.ple.conv1d", + ), + }, } mapping: dict[str, tuple[MODEL_TENSOR, str]] diff --git a/include/llama.h b/include/llama.h index 84313cfd744a..49a758db2680 100644 --- a/include/llama.h +++ b/include/llama.h @@ -214,6 +214,12 @@ extern "C" { LLAMA_API const char * llama_load_mode_name(enum llama_load_mode load_mode); LLAMA_API enum llama_load_mode llama_load_mode_from_str(const char * str); + enum llama_tensor_read_lazy { + LLAMA_TENSOR_READ_LAZY_OFF = 0, // always read the whole tensor up front + LLAMA_TENSOR_READ_LAZY_AUTO = 1, // lazy only for marked tensors larger than 4 GiB (requires mmap) + LLAMA_TENSOR_READ_LAZY_ON = 2, // read the rows of tensors marked by the arch on demand (requires mmap) + }; + enum llama_context_type { LLAMA_CONTEXT_TYPE_DEFAULT = 0, LLAMA_CONTEXT_TYPE_MTP = 1, @@ -315,6 +321,8 @@ extern "C" { enum llama_split_mode split_mode; // how to split the model across multiple GPUs enum llama_load_mode load_mode; // how to load the model + enum llama_tensor_read_lazy tensor_read_lazy; // on-demand reading of tensors marked by the arch + // the GPU that is used for the entire model when split_mode is LLAMA_SPLIT_MODE_NONE int32_t main_gpu; @@ -437,6 +445,7 @@ extern "C" { const struct llama_model_kv_override * kv_overrides; // pointer to kv overrides const struct llama_model_tensor_override * tt_overrides; // pointer to tensor overrides const int32_t * prune_layers; // pointer to layer indices to prune + size_t max_buf_size; // max bytes of tensor rows kept in memory at once, 0 = default (8 GiB) } llama_model_quantize_params; typedef struct llama_logit_bias { diff --git a/scripts/pr2wt.sh b/scripts/pr2wt.sh index ae03a888cb64..ce0327eadb53 100755 --- a/scripts/pr2wt.sh +++ b/scripts/pr2wt.sh @@ -48,7 +48,11 @@ echo "org/repo: $org_repo" meta=$(curl -sSLf -H "Accept: application/vnd.github+json" "https://api.github.com/repos/$org_repo/pulls/$PR") -url_remote=$(echo "$meta" | jq -r '.head.repo.clone_url') +if [[ $url_origin =~ ^git@ ]]; then + url_remote=$(echo "$meta" | jq -r '.head.repo.ssh_url') +else + url_remote=$(echo "$meta" | jq -r '.head.repo.clone_url') +fi head_ref=$(echo "$meta" | jq -r '.head.ref') echo "url: $url_remote" diff --git a/scripts/snapdragon/adb/run-bench.sh b/scripts/snapdragon/adb/run-bench.sh deleted file mode 100755 index eaae80a77d6f..000000000000 --- a/scripts/snapdragon/adb/run-bench.sh +++ /dev/null @@ -1,49 +0,0 @@ -#!/bin/sh -# - -# Basedir on device -basedir=/data/local/tmp/llama.cpp - -branch=. -[ "$B" != "" ] && branch=$B - -adbserial= -[ "$S" != "" ] && adbserial="-s $S" - -adbhost= -[ "$H" != "" ] && adbhost="-H $H" - -model="Llama-3.2-3B-Instruct-Q4_0.gguf" -[ "$M" != "" ] && model="$M" - -device="HTP0" -[ "$D" != "" ] && device="$D" - -verbose= -[ "$V" != "" ] && verbose="GGML_HEXAGON_VERBOSE=$V" cli_opts="$cli_opts -v" - -profile= -[ "$PROF" != "" ] && profile="GGML_HEXAGON_PROFILE=$PROF" cli_opts="$cli_opts -v" - -opmask= -[ "$OPSTAGE" != "" ] && opmask="GGML_HEXAGON_OPSTAGE=$OPSTAGE" - -nhvx= -[ "$NHVX" != "" ] && nhvx="GGML_HEXAGON_NHVX=$NHVX" - -ndev= -[ "$NDEV" != "" ] && ndev="GGML_HEXAGON_NDEV=$NDEV" - -hb= -[ "$HB" != "" ] && hb="GGML_HEXAGON_HOSTBUF=$HB" - -set -x - -adb $adbserial $adbhost shell " \ - cd $basedir; \ - LD_LIBRARY_PATH=$basedir/$branch/lib \ - ADSP_LIBRARY_PATH=$basedir/$branch/lib \ - $ndev $nhvx $opmask $verbose $profile $hb ./$branch/bin/llama-bench --device $device --load-mode none -m $basedir/../gguf/$model \ - --poll 1000 -t 6 --cpu-mask 0xfc --cpu-strict 1 \ - --ubatch-size 1024 -fa 1 -ngl 99 $cli_opts $@ \ -" diff --git a/scripts/snapdragon/adb/run-cli.sh b/scripts/snapdragon/adb/run-cli.sh deleted file mode 100755 index 27a4a14195f0..000000000000 --- a/scripts/snapdragon/adb/run-cli.sh +++ /dev/null @@ -1,78 +0,0 @@ -#!/bin/sh -# - -# Basedir on device -basedir=/data/local/tmp/llama.cpp - -cli_opts= - -branch=. -[ "$B" != "" ] && branch=$B - -adbserial= -[ "$S" != "" ] && adbserial="-s $S" - -adbhost= -[ "$H" != "" ] && adbhost="-H $H" - -model="Llama-3.2-3B-Instruct-Q4_0.gguf" -[ "$M" != "" ] && model="$M" - -device="HTP0" -[ "$D" != "" ] && device="$D" - -verbose= -[ "$V" != "" ] && verbose="GGML_HEXAGON_VERBOSE=$V" cli_opts="$cli_opts -v" - -sched= -[ "$SCHED" != "" ] && sched="GGML_SCHED_DEBUG=2" cli_opts="$cli_opts -v" - -profile= -[ "$PROF" != "" ] && profile="GGML_HEXAGON_PROFILE=$PROF" cli_opts="$cli_opts -v" - -opmask= -[ "$OPSTAGE" != "" ] && opmask="GGML_HEXAGON_OPSTAGE=$OPSTAGE" - -nhvx= -[ "$NHVX" != "" ] && nhvx="GGML_HEXAGON_NHVX=$NHVX" - -hmx= -[ "$HMX" != "" ] && hmx="GGML_HEXAGON_USE_HMX=$HMX" - -ndev= -[ "$NDEV" != "" ] && ndev="GGML_HEXAGON_NDEV=$NDEV" - -hb= -[ "$HB" != "" ] && hb="GGML_HEXAGON_HOSTBUF=$HB" - -opbatch= -[ "$OB" != "" ] && opbatch="GGML_HEXAGON_OPBATCH=$OB" - -opqueue= -[ "$OQ" != "" ] && opqueue="GGML_HEXAGON_OPQUEUE=$OQ" - -opflt= -[ "$OF" != "" ] && opflt="GGML_HEXAGON_OPFILTER=$OF" - -vmem= -[ "$VM" != "" ] && opflt="GGML_HEXAGON_VMEM=$VM" - -mbuf= -[ "$MB" != "" ] && opflt="GGML_HEXAGON_MBUF=$MB" -vmem= -[ "$VM" != "" ] && vmem="GGML_HEXAGON_VMEM=$VM" - -mbuf= -[ "$MB" != "" ] && mbuf="GGML_HEXAGON_MBUF=$MB" -set -x - -adb $adbserial $adbhost shell " \ - cd $basedir; ulimit -c unlimited; \ - LD_LIBRARY_PATH=$basedir/$branch/lib \ - ADSP_LIBRARY_PATH=$basedir/$branch/lib \ - $verbose $sched $opmask $profile $nhvx $hmx $ndev $hb $opbatch $opqueue $opflt $vmem $mbuf \ - ./$branch/bin/llama-cli --load-mode none -m $basedir/../gguf/$model \ - --poll 1000 -t 6 --cpu-mask 0xfc --cpu-strict 1 \ - --ctx-size 8192 --ubatch-size 1024 -fa on \ - -ngl 99 --device $device $cli_opts $@ \ -" diff --git a/scripts/snapdragon/adb/run-completion.sh b/scripts/snapdragon/adb/run-completion.sh deleted file mode 100755 index 30893ed293a2..000000000000 --- a/scripts/snapdragon/adb/run-completion.sh +++ /dev/null @@ -1,86 +0,0 @@ -#!/bin/sh -# - -# Basedir on device -basedir=/data/local/tmp/llama.cpp - -cli_opts= - -branch=. -[ "$B" != "" ] && branch=$B - -adbserial= -[ "$S" != "" ] && adbserial="-s $S" - -adbhost= -[ "$H" != "" ] && adbhost="-H $H" - -model="Llama-3.2-3B-Instruct-Q4_0.gguf" -[ "$M" != "" ] && model="$M" - -device="HTP0" -[ "$D" != "" ] && device="$D" - -verbose= -[ "$V" != "" ] && verbose="GGML_HEXAGON_VERBOSE=$V" cli_opts="$cli_opts -v" - -sched= -[ "$SCHED" != "" ] && sched="GGML_SCHED_DEBUG=2" cli_opts="$cli_opts -v" - -profile= -[ "$PROF" != "" ] && profile="GGML_HEXAGON_PROFILE=$PROF" cli_opts="$cli_opts -v" - -opmask= -[ "$OPSTAGE" != "" ] && opmask="GGML_HEXAGON_OPSTAGE=$OPSTAGE" - -nhvx= -[ "$NHVX" != "" ] && nhvx="GGML_HEXAGON_NHVX=$NHVX" - -hmx= -[ "$HMX" != "" ] && hmx="GGML_HEXAGON_USE_HMX=$HMX" - -ndev= -[ "$NDEV" != "" ] && ndev="GGML_HEXAGON_NDEV=$NDEV" - -hb= -[ "$HB" != "" ] && hb="GGML_HEXAGON_HOSTBUF=$HB" - -opbatch= -[ "$OB" != "" ] && opbatch="GGML_HEXAGON_OPBATCH=$OB" - -opqueue= -[ "$OQ" != "" ] && opqueue="GGML_HEXAGON_OPQUEUE=$OQ" - -oppoll= -[ "$OP" != "" ] && oppoll="GGML_HEXAGON_OPPOLL=$OP" - -opflt= -[ "$OF" != "" ] && opflt="GGML_HEXAGON_OPFILTER=$OF" - -opfuse= -[ "$OC" != "" ] && opfuse="GGML_HEXAGON_OPFUSION=$OC" - -vmem= -[ "$VM" != "" ] && vmem="GGML_HEXAGON_VMEM=$VM" - -mbuf= -[ "$MB" != "" ] && mbuf="GGML_HEXAGON_MBUF=$MB" - -mmsel= -[ "$MM" != "" ] && mmsel="GGML_HEXAGON_MM_SELECT=$MM" - -fasel= -[ "$FA" != "" ] && fasel="GGML_HEXAGON_FA_SELECT=$FA" - -set -x - -adb $adbserial $adbhost shell " \ - cd $basedir; ulimit -c unlimited; \ - LD_LIBRARY_PATH=$basedir/$branch/lib \ - ADSP_LIBRARY_PATH=$basedir/$branch/lib \ - $verbose $sched $opmask $profile $nhvx $hmx $ndev $hb $opbatch $opqueue $oppoll $opflt $opfuse $vmem $mbuf $mmsel $fasel \ - ./$branch/bin/llama-completion --load-mode none -m $basedir/../gguf/$model \ - --poll 1000 -t 6 --cpu-mask 0xfc --cpu-strict 1 \ - --ctx-size 8192 --ubatch-size 1024 -fa on \ - -ngl 99 --device $device $cli_opts $@ \ -" diff --git a/scripts/snapdragon/adb/run-mtmd.sh b/scripts/snapdragon/adb/run-mtmd.sh deleted file mode 100755 index 65dd6ec59e51..000000000000 --- a/scripts/snapdragon/adb/run-mtmd.sh +++ /dev/null @@ -1,71 +0,0 @@ -#!/bin/sh -# - -# Basedir on device -basedir=/data/local/tmp/llama.cpp - -cli_opts= - -branch=. -[ "$B" != "" ] && branch=$B - -adbserial= -[ "$S" != "" ] && adbserial="-s $S" - -adbhost= -[ "$H" != "" ] && adbhost="-H $H" - -model="gemma-3-4b-it-Q4_0.gguf" -[ "$M" != "" ] && model="$M" - -mmproj="mmproj-F16.gguf" -[ "$MMPROJ" != "" ] && mmproj="$MMPROJ" - -image= -[ "$IMG" != "" ] && image="$IMG" - -device="HTP0" -[ "$D" != "" ] && device="$D" - -verbose= -[ "$V" != "" ] && verbose="GGML_HEXAGON_VERBOSE=$V" - -experimental="GGML_HEXAGON_EXPERIMENTAL=1" -[ "$E" != "" ] && experimental="GGML_HEXAGON_EXPERIMENTAL=$E" - -sched= -[ "$SCHED" != "" ] && sched="GGML_SCHED_DEBUG=2" cli_opts="$cli_opts -v" - -profile= -[ "$PROF" != "" ] && profile="GGML_HEXAGON_PROFILE=$PROF" - -opmask= -[ "$OPSTAGE" != "" ] && opmask="GGML_HEXAGON_OPSTAGE=$OPSTAGE" - -nhvx= -[ "$NHVX" != "" ] && nhvx="GGML_HEXAGON_NHVX=$NHVX" - -hmx= -[ "$HMX" != "" ] && hmx="GGML_HEXAGON_USE_HMX=$HMX" - -ndev= -[ "$NDEV" != "" ] && ndev="GGML_HEXAGON_NDEV=$NDEV" - -# MTMD backend device for vision model (defaults to CPU if not set) -mtmd_backend= -[ "$MTMD_DEVICE" != "" ] && mtmd_backend="MTMD_BACKEND_DEVICE=$MTMD_DEVICE" - -set -x - -adb $adbserial $adbhost shell " \ - cd $basedir; ulimit -c unlimited; \ - LD_LIBRARY_PATH=$basedir/$branch/lib \ - ADSP_LIBRARY_PATH=$basedir/$branch/lib \ - $verbose $experimental $sched $opmask $profile $hmx $nhvx $ndev $mtmd_backend \ - ./$branch/bin/llama-mtmd-cli --load-mode none -m $basedir/../gguf/$model \ - --mmproj $basedir/../gguf/$mmproj \ - --image $basedir/../gguf/$image \ - --poll 1000 -t 6 --cpu-mask 0xfc --cpu-strict 1 \ - --ctx-size 8192 --ubatch-size 1024 -fa on \ - -ngl 99 --device $device -v $cli_opts $@ \ -" diff --git a/scripts/snapdragon/adb/run-tool.sh b/scripts/snapdragon/adb/run-tool.sh deleted file mode 100755 index 8bf7ba3a5461..000000000000 --- a/scripts/snapdragon/adb/run-tool.sh +++ /dev/null @@ -1,72 +0,0 @@ -#!/bin/sh -# - -# Basedir on device -basedir=/data/local/tmp/llama.cpp - -cli_opts= - -branch=. -[ "$B" != "" ] && branch=$B - -adbserial= -[ "$S" != "" ] && adbserial="-s $S" - -adbhost= -[ "$H" != "" ] && adbhost="-H $H" - -device="HTP0" -[ "$D" != "" ] && device="$D" - -verbose= -[ "$V" != "" ] && verbose="GGML_HEXAGON_VERBOSE=$V" - -sched= -[ "$SCHED" != "" ] && sched="GGML_SCHED_DEBUG=2" cli_opts="$cli_opts -v" - -profile= -[ "$PROF" != "" ] && profile="GGML_HEXAGON_PROFILE=$PROF" - -opmask= -[ "$OPSTAGE" != "" ] && opmask="GGML_HEXAGON_OPSTAGE=$OPSTAGE" - -nhvx= -[ "$NHVX" != "" ] && nhvx="GGML_HEXAGON_NHVX=$NHVX" - -hmx= -[ "$HMX" != "" ] && hmx="GGML_HEXAGON_USE_HMX=$HMX" - -ndev= -[ "$NDEV" != "" ] && ndev="GGML_HEXAGON_NDEV=$NDEV" - -hb= -[ "$HB" != "" ] && hb="GGML_HEXAGON_HOSTBUF=$HB" - -opbatch= -[ "$OB" != "" ] && opbatch="GGML_HEXAGON_OPBATCH=$OB" - -opqueue= -[ "$OQ" != "" ] && opqueue="GGML_HEXAGON_OPQUEUE=$OQ" - -oppoll= -[ "$OP" != "" ] && oppoll="GGML_HEXAGON_OPPOLL=$OP" - -opfuse= -[ "$OC" != "" ] && opfuse="GGML_HEXAGON_OPFUSION=$OC" - -mmsel= -[ "$MM" != "" ] && mmsel="GGML_HEXAGON_MM_SELECT=$MM" - -fasel= -[ "$FA" != "" ] && fasel="GGML_HEXAGON_FA_SELECT=$FA" - -set -x - -tool=$1; shift - -adb $adbserial $adbhost shell " \ - cd $basedir; ulimit -c unlimited; \ - LD_LIBRARY_PATH=$basedir/$branch/lib \ - ADSP_LIBRARY_PATH=$basedir/$branch/lib \ - $verbose $sched $opmask $profile $nhvx $hmx $ndev $hb $opbatch $opqueue $oppoll $opfuse $mmsel $fasel ./$branch/bin/$tool $@ \ -" diff --git a/scripts/snapdragon/build.py b/scripts/snapdragon/build.py new file mode 100755 index 000000000000..02dcb930cb67 --- /dev/null +++ b/scripts/snapdragon/build.py @@ -0,0 +1,260 @@ +#!/usr/bin/env python3 +# +# Build llama.cpp for Snapdragon (via Docker or natively) and push to device. +# + +import sys +import os +import argparse +import subprocess +import platform +import shutil +import logging + +logger = logging.getLogger("build") + + +def parse_target(target_str): + if not target_str: + return None, None + if target_str.startswith("adb") or target_str.startswith("android"): + parts = target_str.split(":", 1) + serial = parts[1] if len(parts) > 1 else None + return "android", serial + elif target_str.startswith("lnx") or target_str.startswith("linux") or target_str.startswith("ubuntu"): + parts = target_str.split(":", 1) + host = parts[1] if len(parts) > 1 else None + return "linux", host + elif target_str in ("wos", "windows"): + return "windows", None + else: + return None, None + + +def get_uid_gid(): + if platform.system() != "Windows": + return [f"{os.getuid()}:{os.getgid()}"] + return [] + + +def main(): + logging.basicConfig(level=logging.INFO, format='%(message)s') + parser = argparse.ArgumentParser( + description="Build llama.cpp for Snapdragon using cross-compilation docker containers or natively." + ) + parser.add_argument("--target", default="android", help="Compilation target and deployment definition (e.g. android[:serial]/adb[:serial], linux:[user@]host/lnx:[user@]host/ubuntu:[user@]host, windows/wos) (default: android)") + parser.add_argument("--build-dir", help="Build directory name (defaults to build-TARGET[-dbg], e.g. build-android)") + parser.add_argument("--install-dir", help="Install directory name (defaults to pkg-TARGET[-dbg], e.g. pkg-android)") + parser.add_argument("--jobs", "-j", type=int, help="Number of build jobs (defaults to CPU thread count)") + parser.add_argument("--no-docker", action="store_true", help="Build natively on the host instead of in a docker container") + parser.add_argument("--preset", help="Override the CMake preset to use") + parser.add_argument("--debug", action="store_true", help="Build in debug mode (uses -debug presets instead of -release)") + + # Push options + parser.add_argument("--push", action="store_true", help="Push built package to the target device via ADB or SSH/SCP") + parser.add_argument("--target-dir", help="Target directory on the device (default: /data/local/tmp/llama.cpp for Android, ~/llama.cpp for Linux)") + + # Toolchain options + parser.add_argument("--toolchain-version", default="v0.7", help="Docker toolchain image version/tag (default: v0.7)") + parser.add_argument("--toolchain-url", default="ghcr.io/snapdragon-toolchain", help="Docker toolchain registry URL/namespace (default: ghcr.io/snapdragon-toolchain)") + + args = parser.parse_args() + + target_type, target_val = parse_target(args.target) + if not target_type: + logger.error(f"Error: Invalid target format '{args.target}'. Must be android[:serial]/adb[:serial], linux:[user@]host/lnx:[user@]host/ubuntu:[user@]host, or windows/wos.") + sys.exit(1) + + # Determine preset and check if it's debug + preset = args.preset + if preset: + is_debug = args.debug or ("debug" in preset.lower()) + else: + is_debug = args.debug + config_type = "debug" if is_debug else "release" + if args.no_docker: + if target_type == "windows" or platform.system() == "Windows": + preset = f"arm64-windows-snapdragon-{config_type}" + elif target_type == "linux": + preset = f"arm64-linux-snapdragon-{config_type}" + else: + preset = f"arm64-android-snapdragon-{config_type}" + else: + preset = f"arm64-linux-snapdragon-{config_type}" if target_type == "linux" else f"arm64-android-snapdragon-{config_type}" + + target_prefix = args.target.split(":", 1)[0] + suffix = "-dbg" if is_debug else "" + + build_dir = args.build_dir + if not build_dir: + build_dir = f"build-{target_prefix}{suffix}" + + install_dir = args.install_dir + if not install_dir: + install_dir = f"pkg-{target_prefix}{suffix}" + + repo_root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) + + # Ensure CMakeUserPresets.json is in the workspace root, update if docs version is newer + preset_src = os.path.join(repo_root, "docs", "backend", "snapdragon", "CMakeUserPresets.json") + preset_dst = os.path.join(repo_root, "CMakeUserPresets.json") + if os.path.exists(preset_src): + should_copy = False + if not os.path.exists(preset_dst): + should_copy = True + else: + # Check modification times + src_mtime = os.path.getmtime(preset_src) + dst_mtime = os.path.getmtime(preset_dst) + if src_mtime > dst_mtime: + preset_bak = preset_dst + ".bak" + logger.info(f"Docs CMakeUserPresets.json is newer. Backing up existing {preset_dst} to {preset_bak}") + shutil.copy2(preset_dst, preset_bak) + should_copy = True + + if should_copy: + logger.info(f"Copying CMakeUserPresets.json from {preset_src} to {preset_dst}") + shutil.copy2(preset_src, preset_dst) + else: + logger.warning("Warning: CMakeUserPresets.json not found in docs/backend/snapdragon/.") + + jobs = args.jobs if args.jobs else os.cpu_count() or 4 + + if target_type == "windows": + logger.info("Windows target selected. Forcing native compilation...") + args.no_docker = True + if platform.system() != "Windows": + logger.warning("Warning: Windows compilation is intended to run on Windows arm64 hosts.") + + if args.no_docker: + # Native/local host build + logger.info("Running native/local CMake build...") + install_prefix = os.path.join(repo_root, install_dir, "llama.cpp") + + # Configure + configure_cmd = ["cmake", f"--preset={preset}", "-B", build_dir] + logger.info(f"+ {' '.join(configure_cmd)}") + res = subprocess.run(configure_cmd, cwd=repo_root) + if res.returncode != 0: + logger.error("CMake configuration failed.") + sys.exit(res.returncode) + + # Build + build_cmd = ["cmake", "--build", build_dir, "-j", str(jobs)] + logger.info(f"+ {' '.join(build_cmd)}") + res = subprocess.run(build_cmd, cwd=repo_root) + if res.returncode != 0: + logger.error("CMake build failed.") + sys.exit(res.returncode) + + # Install + install_cmd = ["cmake", "--install", build_dir, "--prefix", install_prefix] + logger.info(f"+ {' '.join(install_cmd)}") + res = subprocess.run(install_cmd, cwd=repo_root) + if res.returncode != 0: + logger.error("CMake install failed.") + sys.exit(res.returncode) + else: + # Docker-based build + logger.info("Running Docker-based cross-compilation build...") + image_name = "arm64-linux" if target_type == "linux" else "arm64-android" + image = f"{args.toolchain_url}/{image_name}:{args.toolchain_version}" + + install_prefix_container = f"/workspace/{install_dir}/llama.cpp" + + build_sh_cmd = ( + f"cmake --preset {preset} -B /workspace/{build_dir} && " + f"cmake --build /workspace/{build_dir} -j {jobs} && " + f"cmake --install /workspace/{build_dir} --prefix {install_prefix_container}" + ) + + docker_cmd = [ + "docker", "run", "--rm", + "--volume", f"{repo_root}:/workspace", + "--workdir", "/workspace", + "--platform", "linux/amd64" + ] + uid_gid = get_uid_gid() + if uid_gid: + docker_cmd += ["-u", uid_gid[0]] + + docker_cmd += [image, "bash", "-c", build_sh_cmd] + + logger.info(f"+ {' '.join(docker_cmd)}") + res = subprocess.run(docker_cmd, cwd=repo_root) + if res.returncode != 0: + logger.error("Docker-based build failed.") + sys.exit(res.returncode) + + logger.info("\nBuild and installation completed successfully!") + + # Push/deploy if requested + if args.push: + src_path = os.path.join(repo_root, install_dir, "llama.cpp") + if not os.path.exists(src_path): + logger.error(f"Error: installation directory {src_path} does not exist. Cannot deploy.") + sys.exit(1) + + # Resolve target directory on device + target_dir = args.target_dir + if not target_dir: + target_dir = "/data/local/tmp/llama.cpp" if target_type == "android" else "~/llama.cpp" + target_dir = target_dir.rstrip("/") + + sub_items = [item for item in os.listdir(src_path) if not item.startswith(".")] + + if target_type == "android": + logger.info("\nPushing built artifacts to Android device via ADB...") + adb_cmd = ["adb"] + if target_val: # serial + adb_cmd += ["-s", target_val] + + # Clean stale package files on device + if sub_items: + clean_paths = " ".join(f"{target_dir}/{item}" for item in sub_items) + clean_cmd = adb_cmd + ["shell", f"rm -rf {clean_paths}"] + logger.info(f"+ {' '.join(clean_cmd)}") + subprocess.run(clean_cmd) + + # Android destination directory is target_dir + push_cmd = adb_cmd + ["push", os.path.join(src_path, "."), target_dir] + logger.info(f"+ {' '.join(push_cmd)}") + res = subprocess.run(push_cmd) + if res.returncode != 0: + logger.error("ADB push failed.") + sys.exit(res.returncode) + logger.info("ADB push completed successfully!") + + elif target_type == "linux": + ssh_host = target_val + if not ssh_host: + logger.error("Error: SSH host not specified in target (e.g. use linux:user@host, lnx:user@host, or ubuntu:user@host). Cannot deploy.") + sys.exit(1) + logger.info(f"\nDeploying built artifacts to Linux device {ssh_host} via SSH/SCP...") + + # Clean stale package files on remote host + if sub_items: + clean_paths = " ".join(f"{target_dir}/{item}" for item in sub_items) + clean_cmd = ["ssh", ssh_host, f"rm -rf {clean_paths}"] + logger.info(f"+ {' '.join(clean_cmd)}") + subprocess.run(clean_cmd) + + # Deploy to target_dir + deploy_cmd = ["scp", "-r", os.path.join(src_path, "."), f"{ssh_host}:{target_dir}"] + logger.info(f"+ {' '.join(deploy_cmd)}") + res = subprocess.run(deploy_cmd) + if res.returncode != 0: + logger.error("SSH/SCP deploy failed.") + sys.exit(res.returncode) + logger.info("SSH/SCP deploy completed successfully!") + + elif target_type == "windows": + logger.info("\nPush for Windows on Snapdragon (windows) target is currently a stub.") + + +if __name__ == "__main__": + try: + main() + except KeyboardInterrupt: + logger.info("\nInterrupted by user.") + sys.exit(130) diff --git a/scripts/snapdragon/ggml-hexagon-profile.py b/scripts/snapdragon/ggml-hexagon-profile.py index 97a3acd26c26..038d92fb5c47 100755 --- a/scripts/snapdragon/ggml-hexagon-profile.py +++ b/scripts/snapdragon/ggml-hexagon-profile.py @@ -34,6 +34,26 @@ r"trace-evt\s+(?P[A-Z_0-9\-]+):\s+thread\s+(?P\d+)\s+info\s+(?P\d+)\s+(?Pstart|stop)\s+(?P\d+)" ) +device_pattern = re.compile(r"\b(HTP\d+(?::\d+)?)\s+(?:profile-op|trace-evt)\b") + + +def extract_device(line): + m = device_pattern.search(line) + if m: + return m.group(1) + return "HTP0" + + +def device_matches(record_device, target_device): + targets = [t.strip() for t in target_device.split(',')] + for target in targets: + if record_device == target: + return True + if record_device.startswith(target + ":"): + return True + return False + + logger = logging.getLogger("ggml-hexagon-profile") @@ -72,7 +92,7 @@ def unwrap(self, raw): return raw + self.high_part -def parse_log(file_path, pmu_index=None): +def parse_log(file_path, pmu_index=None, limit=None, device_filter=None, op_filter_re=None): try: if file_path != "-": f = open(file_path, 'r', encoding='utf-8', errors='ignore') @@ -85,13 +105,22 @@ def parse_log(file_path, pmu_index=None): all_ops: List[Dict[str, Any]] = [] all_traces: List[Dict[str, Any]] = [] current_op: Optional[Dict[str, Any]] = None + ops_count_per_device = {} + if device_filter is not None: + for target in device_filter.split(','): + ops_count_per_device[target.strip()] = 0 + limit_reached = False - timestamp_pattern = re.compile(r"^(?P\d+)\.(?P\d+)\.(?P\d+)\.(?P\d+)\s+[A-Z]\s+") - unwrapper = None - trace_unwrapper = None + timestamp_pattern = re.compile(r"(?P\d+)\.(?P\d+)\.(?P\d+)\.(?P\d+)\s+[A-Z]\s+") + unwrappers = {} + last_batch_start = {} + trace_unwrappers = {} for line in f: - ts_match = timestamp_pattern.match(line) + if "profile-op" not in line and "trace-evt" not in line: + continue + + ts_match = timestamp_pattern.search(line) abs_usec = 0 if ts_match: abs_usec = ( @@ -100,8 +129,11 @@ def parse_log(file_path, pmu_index=None): + int(ts_match.group('us')) ) - if "|" in line and "profile-op" in line: - parts = [p.strip() for p in line.split("|")] + device = extract_device(line) + + idx = line.find("profile-op") + if idx != -1 and "|" in line[idx:]: + parts = [p.strip() for p in line[idx:].split("|")] prefix = parts[0] prefix_match = re.search(r"profile-op\s+(?P[A-Z_0-9+]+)", prefix) if not prefix_match: @@ -145,7 +177,6 @@ def parse_log(file_path, pmu_index=None): except (ValueError, IndexError): pmu_val = None - evt_val = None evt_val = None if types.startswith("evt-cnt "): try: @@ -158,14 +189,18 @@ def parse_log(file_path, pmu_index=None): if op_name == "OPBATCH": if cycles_start_raw: unwrapped_cycles_start = int(cycles_start_raw) - unwrapper = CycleUnwrapper(unwrapped_cycles_start) - trace_unwrapper = CycleUnwrapper(unwrapped_cycles_start) + unwrappers[device] = CycleUnwrapper(unwrapped_cycles_start) + last_batch_start[device] = unwrapped_cycles_start + for k in list(trace_unwrappers.keys()): + if k[0] == device: + del trace_unwrappers[k] else: - if cycles_start_raw and unwrapper is not None: - unwrapped_cycles_start = unwrapper.unwrap(int(cycles_start_raw)) + if cycles_start_raw: + device_unwrapper = unwrappers.get(device) + if device_unwrapper is not None: + unwrapped_cycles_start = device_unwrapper.unwrap(int(cycles_start_raw)) - idx = line.find("profile-op ") - op_text = line[idx + 11:].strip() if idx != -1 else line.strip() + op_text = re.sub(r"^profile-op\s+", "", line[idx:]).strip() if idx != -1 else line.strip() current_op = { 'name': op_name, @@ -180,24 +215,58 @@ def parse_log(file_path, pmu_index=None): 'pmu_val': pmu_val, 'evt_val': evt_val, 'abs_usec': abs_usec, - 'trace_events': [] + 'trace_events': [], + 'device': device } all_ops.append(current_op) + + # Check if matching early exit criteria + matched = False + matched_target = None + if device_filter is not None: + targets = [t.strip() for t in device_filter.split(',')] + for target in targets: + if device == target or device.startswith(target + ":"): + matched = True + matched_target = target + break + else: + matched = True + matched_target = device + + if op_filter_re is not None and not op_filter_re.search(op_text): + matched = False + + if matched: + if matched_target not in ops_count_per_device: + ops_count_per_device[matched_target] = 0 + ops_count_per_device[matched_target] += 1 + + if limit is not None and len(ops_count_per_device) > 0 and all(count >= limit for count in ops_count_per_device.values()): + limit_reached = True + + if limit_reached and op_name == "OPBATCH": + break continue trace_match = trace_pattern.search(line) if trace_match: + thread = int(trace_match.group('thread')) raw_cyc = int(trace_match.group('cycles')) unwrapped_cyc = None - if trace_unwrapper is not None: - unwrapped_cyc = trace_unwrapper.unwrap(raw_cyc) + th_key = (device, thread) + if th_key not in trace_unwrappers: + batch_start = last_batch_start.get(device) + trace_unwrappers[th_key] = CycleUnwrapper(batch_start) + unwrapped_cyc = trace_unwrappers[th_key].unwrap(raw_cyc) all_traces.append({ - 'thread': int(trace_match.group('thread')), + 'thread': thread, 'event': trace_match.group('event'), 'info': int(trace_match.group('info')), 'cycles': raw_cyc, 'unwrapped_cycles': unwrapped_cyc, - 'state': trace_match.group('state') + 'state': trace_match.group('state'), + 'device': device }) f.close() @@ -207,39 +276,45 @@ def parse_log(file_path, pmu_index=None): op['start_cycles'] = op['unwrapped_cycles_start'] op['end_cycles'] = op['start_cycles'] + op['cycles'] if op['start_cycles'] is not None else None - # Filter ops with valid start_cycles - valid_ops = [op for op in all_ops if op['start_cycles'] is not None and op['end_cycles'] is not None] + # Group ops by device + valid_ops_by_dev = defaultdict(list) + for op in all_ops: + if op['start_cycles'] is not None and op['end_cycles'] is not None: + valid_ops_by_dev[op['device']].append(op) + + # Group trace events by device + traces_by_dev = defaultdict(list) + for e in all_traces: + if e['unwrapped_cycles'] is not None: + traces_by_dev[e['device']].append(e) - # Separate OPBATCH ops from other ops - opbatch_ops = [op for op in valid_ops if op['name'] == "OPBATCH"] - other_ops = [op for op in valid_ops if op['name'] != "OPBATCH"] + for device, dev_ops in valid_ops_by_dev.items(): + opbatch_ops = [op for op in dev_ops if op['name'] == "OPBATCH"] + other_ops = [op for op in dev_ops if op['name'] != "OPBATCH"] - # Sort them by start_cycles to enable binary search - opbatch_ops.sort(key=lambda op: op['start_cycles']) - other_ops.sort(key=lambda op: op['start_cycles']) + opbatch_ops.sort(key=lambda op: op['start_cycles']) + other_ops.sort(key=lambda op: op['start_cycles']) - opbatch_starts = [op['start_cycles'] for op in opbatch_ops] - other_starts = [op['start_cycles'] for op in other_ops] + opbatch_starts = [op['start_cycles'] for op in opbatch_ops] + other_starts = [op['start_cycles'] for op in other_ops] - # Map trace events to any operator whose cycles contain them - for e in all_traces: - cyc = e['unwrapped_cycles'] - if cyc is None: - continue + dev_traces = traces_by_dev.get(device, []) + for e in dev_traces: + cyc = e['unwrapped_cycles'] - # Map to OPBATCH - idx = bisect.bisect_right(opbatch_starts, cyc) - 1 - if idx >= 0: - op = opbatch_ops[idx] - if op['start_cycles'] <= cyc <= op['end_cycles']: - op['trace_events'].append(e) + # Map to OPBATCH + idx = bisect.bisect_right(opbatch_starts, cyc) - 1 + if idx >= 0: + op = opbatch_ops[idx] + if op['start_cycles'] <= cyc <= op['end_cycles']: + op['trace_events'].append(e) - # Map to other ops - idx = bisect.bisect_right(other_starts, cyc) - 1 - if idx >= 0: - op = other_ops[idx] - if op['start_cycles'] <= cyc <= op['end_cycles']: - op['trace_events'].append(e) + # Map to other ops + idx = bisect.bisect_right(other_starts, cyc) - 1 + if idx >= 0: + op = other_ops[idx] + if op['start_cycles'] <= cyc <= op['end_cycles']: + op['trace_events'].append(e) return all_ops @@ -563,6 +638,7 @@ def main(): parser.add_argument("--timeline", type=str, nargs='?', const='summary', choices=["summary", "bubbles"], help="Output ASCII art event summary or thread idle bubble analysis (default: summary)") parser.add_argument("--filter", type=str, help="Regex filter matching against the original profile-op line") + parser.add_argument("--device", type=str, help="Device to filter by (e.g. HTP0, HTP0:0) or 'split' to generate separate reports per device") group = parser.add_mutually_exclusive_group() group.add_argument("--head", type=int, help="Limit to first N ops") @@ -586,29 +662,84 @@ def main(): logger.warning(f"Invalid width format '{w}'") final_pmu_name = (args.pmu_name or f"#{args.pmu_index}") if args.pmu_index is not None else None - ops = parse_log(args.logfile, pmu_index=args.pmu_index) + op_filter_re = None if args.filter: try: - filter_re = re.compile(args.filter) + op_filter_re = re.compile(args.filter) except re.error as e: logger.error(f"Invalid regex filter: {e}") sys.exit(1) - ops = [op for op in ops if filter_re.search(op['op_text'])] - - if args.head is not None: - ops = ops[:args.head] - elif args.tail is not None: - ops = ops[-args.tail:] - - if args.timeline: - for op in ops: - if args.timeline == "summary": - print_ascii_summary(op['name'], op['dims'], op['types'], op['usec'], op['cycles'], op['trace_events']) - elif args.timeline == "bubbles": - print_bubbles_timeline(op) + + limit = args.head if args.head is not None else None + device_filter = args.device if (args.device and args.device != "split") else None + ops = parse_log(args.logfile, pmu_index=args.pmu_index, limit=limit, device_filter=device_filter, op_filter_re=op_filter_re) + + if args.device and args.device != "split": + ops = [op for op in ops if device_matches(op['device'], args.device)] + + if args.device == "split": + unique_devices = sorted(list(set(op['device'] for op in ops))) + for dev in unique_devices: + dev_ops = [op for op in ops if device_matches(op['device'], dev)] + + if args.filter: + try: + filter_re = re.compile(args.filter) + except re.error as e: + logger.error(f"Invalid regex filter: {e}") + sys.exit(1) + dev_ops = [op for op in dev_ops if filter_re.search(op['op_text'])] + + if args.head is not None: + dev_ops = dev_ops[:args.head] + elif args.tail is not None: + dev_ops = dev_ops[-args.tail:] + + logger.info("\n=========================================") + logger.info(f" Device: {dev}") + logger.info("=========================================") + + if args.timeline: + for op in dev_ops: + if args.timeline == "summary": + print_ascii_summary(op['name'], op['dims'], op['types'], op['usec'], op['cycles'], op['trace_events']) + elif args.timeline == "bubbles": + print_bubbles_timeline(op) + else: + generate_report(dev_ops, args.top, overrides, args.sort, pmu_name=final_pmu_name) else: - generate_report(ops, args.top, overrides, args.sort, pmu_name=final_pmu_name) + if args.filter: + try: + filter_re = re.compile(args.filter) + except re.error as e: + logger.error(f"Invalid regex filter: {e}") + sys.exit(1) + ops = [op for op in ops if filter_re.search(op['op_text'])] + + if args.head is not None or args.tail is not None: + ops_by_dev = defaultdict(list) + for op in ops: + ops_by_dev[op['device']].append(op) + + filtered_ops = [] + for dev in sorted(ops_by_dev.keys()): + dev_ops = ops_by_dev[dev] + if args.head is not None: + dev_ops = dev_ops[:args.head] + elif args.tail is not None: + dev_ops = dev_ops[-args.tail:] + filtered_ops.extend(dev_ops) + ops = filtered_ops + + if args.timeline: + for op in ops: + if args.timeline == "summary": + print_ascii_summary(op['name'], op['dims'], op['types'], op['usec'], op['cycles'], op['trace_events']) + elif args.timeline == "bubbles": + print_bubbles_timeline(op) + else: + generate_report(ops, args.top, overrides, args.sort, pmu_name=final_pmu_name) if __name__ == "__main__": diff --git a/scripts/snapdragon/ggml-hexagon-trace.py b/scripts/snapdragon/ggml-hexagon-trace.py index 4755adfa1339..99bf771b85d8 100755 --- a/scripts/snapdragon/ggml-hexagon-trace.py +++ b/scripts/snapdragon/ggml-hexagon-trace.py @@ -20,6 +20,31 @@ r"trace-evt\s+(?P[A-Z_0-9\-]+):\s+thread\s+(?P\d+)\s+info\s+(?P\d+)\s+(?Pstart|stop)\s+(?P\d+)" ) +device_pattern = re.compile(r"\b(HTP\d+(?::\d+)?)\s+(?:profile-op|trace-evt)\b") + + +def extract_device(line): + m = device_pattern.search(line) + if m: + return m.group(1) + return "HTP0" + + +def device_matches(record_device, target_device): + targets = [t.strip() for t in target_device.split(',')] + for target in targets: + if record_device == target: + return True + if record_device.startswith(target + ":"): + return True + return False + + +def get_split_output_path(base_path, device_name): + safe_device = device_name.replace(':', '_') + root, ext = os.path.splitext(base_path) + return f"{root}-{safe_device}{ext}" + def normalize_event_name(evt_type, info=0): if evt_type == "HVX_COMP": @@ -54,7 +79,79 @@ def unwrap(self, raw): return raw + self.high_part -def parse_log(file_path): +class DeviceTimeMapper: + def __init__(self, dev, ops): + self.dev = dev + self.batches = [] + for op in ops: + if op.get('device') == dev and op.get('name') == 'OPBATCH' and op.get('unwrapped_cycles_start') is not None: + cycles = op.get('cycles', 0) + usec = op.get('usec', 0) + start_cyc = op['unwrapped_cycles_start'] + freq = (cycles / usec) if usec > 0 and cycles > 0 else 1000.0 + if freq <= 0: + freq = 1000.0 + self.batches.append({ + 'start_cycles': start_cyc, + 'cycles': cycles, + 'end_cycles': start_cyc + cycles, + 'usec': usec, + 'dur_ns': usec * 1000, + 'freq_mhz': freq, + }) + + self.batches.sort(key=lambda b: b['start_cycles']) + + for i, b in enumerate(self.batches): + if i == 0: + b['start_time_ns'] = 0 + else: + prev = self.batches[i - 1] + idle_cyc = max(0, b['start_cycles'] - prev['end_cycles']) + idle_ns = int(round((idle_cyc / prev['freq_mhz']) * 1000)) + b['start_time_ns'] = prev['start_time_ns'] + prev['dur_ns'] + idle_ns + + self.batch_starts = [b['start_cycles'] for b in self.batches] + + valid_starts = [op['unwrapped_cycles_start'] for op in ops if op.get('device') == dev and op.get('unwrapped_cycles_start') is not None] + self.min_cyc = min(valid_starts) if valid_starts else 0 + if self.batches: + self.default_freq = self.batches[0]['freq_mhz'] + else: + freqs = [op['cycles'] / op['usec'] for op in ops if op.get('device') == dev and op.get('usec', 0) > 0 and op.get('cycles', 0) > 0] + self.default_freq = statistics.mean(freqs) if freqs else 1000.0 + + def get_batch(self, cyc): + if not self.batches: + return None + idx = bisect.bisect_right(self.batch_starts, cyc) - 1 + if idx >= 0: + return self.batches[idx] + return self.batches[0] + + def get_freq(self, cyc=None): + if cyc is not None: + b = self.get_batch(cyc) + if b is not None: + return b['freq_mhz'] + return self.default_freq + + def cycle_to_ns(self, cyc): + if cyc is None: + return 0 + b = self.get_batch(cyc) + if b is not None: + return b['start_time_ns'] + int(round(((cyc - b['start_cycles']) / b['freq_mhz']) * 1000)) + return int(round(((cyc - self.min_cyc) / self.default_freq) * 1000)) + + def dur_cycles_to_ns(self, cyc_start, cyc_dur): + if cyc_dur is None: + return 0 + freq = self.get_freq(cyc_start) + return int(round((cyc_dur / freq) * 1000)) + + +def parse_log(file_path, limit=None, device_filter=None, op_filter_re=None): try: if file_path != "-": f = open(file_path, 'r', encoding='utf-8', errors='ignore') @@ -67,14 +164,25 @@ def parse_log(file_path): all_ops: List[Dict[str, Any]] = [] all_traces: List[Dict[str, Any]] = [] current_op: Optional[Dict[str, Any]] = None - unwrapper = None - trace_unwrapper = None + ops_count_per_device = {} + if device_filter is not None: + for target in device_filter.split(','): + ops_count_per_device[target.strip()] = 0 + limit_reached = False + unwrappers = {} + last_batch_start = {} + trace_unwrappers = {} line_idx = 0 for line in f: line_idx += 1 - if "|" in line and "profile-op" in line: - parts = [p.strip() for p in line.split("|")] + if "profile-op" not in line and "trace-evt" not in line: + continue + device = extract_device(line) + + idx = line.find("profile-op") + if idx != -1 and "|" in line[idx:]: + parts = [p.strip() for p in line[idx:].split("|")] prefix = parts[0] prefix_match = re.search(r"profile-op\s+(?P[A-Z_0-9+]+)", prefix) if not prefix_match: @@ -115,14 +223,18 @@ def parse_log(file_path): if op_name == "OPBATCH": if cycles_start_raw: unwrapped_cycles_start = int(cycles_start_raw) - unwrapper = CycleUnwrapper(unwrapped_cycles_start) - trace_unwrapper = CycleUnwrapper(unwrapped_cycles_start) + unwrappers[device] = CycleUnwrapper(unwrapped_cycles_start) + last_batch_start[device] = unwrapped_cycles_start + for k in list(trace_unwrappers.keys()): + if k[0] == device: + del trace_unwrappers[k] else: - if cycles_start_raw and unwrapper is not None: - unwrapped_cycles_start = unwrapper.unwrap(int(cycles_start_raw)) + if cycles_start_raw: + device_unwrapper = unwrappers.get(device) + if device_unwrapper is not None: + unwrapped_cycles_start = device_unwrapper.unwrap(int(cycles_start_raw)) - idx = line.find("profile-op ") - op_text = line[idx + 11:].strip() if idx != -1 else line.strip() + op_text = re.sub(r"^profile-op\s+", "", line[idx:]).strip() if idx != -1 else line.strip() evt_str = None if types.startswith("evt-cnt "): @@ -142,24 +254,59 @@ def parse_log(file_path): 'cycles_start': int(cycles_start_raw) if cycles_start_raw else None, 'unwrapped_cycles_start': unwrapped_cycles_start, 'trace_events': [], - 'line_num': line_idx + 'line_num': line_idx, + 'device': device } all_ops.append(current_op) + + # Check if matching early exit criteria + matched = False + matched_target = None + if device_filter is not None: + targets = [t.strip() for t in device_filter.split(',')] + for target in targets: + if device == target or device.startswith(target + ":"): + matched = True + matched_target = target + break + else: + matched = True + matched_target = device + + if op_filter_re is not None and not op_filter_re.search(op_text): + matched = False + + if matched: + if matched_target not in ops_count_per_device: + ops_count_per_device[matched_target] = 0 + ops_count_per_device[matched_target] += 1 + + if limit is not None and len(ops_count_per_device) > 0 and all(count >= limit for count in ops_count_per_device.values()): + limit_reached = True + + if limit_reached and op_name == "OPBATCH": + break continue trace_match = trace_pattern.search(line) if trace_match: + thread = int(trace_match.group('thread')) raw_cyc = int(trace_match.group('cycles')) unwrapped_cyc = None - if trace_unwrapper is not None: - unwrapped_cyc = trace_unwrapper.unwrap(raw_cyc) + th_key = (device, thread) + if th_key not in trace_unwrappers: + batch_start = last_batch_start.get(device) + trace_unwrappers[th_key] = CycleUnwrapper(batch_start) + unwrapped_cyc = trace_unwrappers[th_key].unwrap(raw_cyc) all_traces.append({ - 'thread': int(trace_match.group('thread')), + 'thread': thread, 'event': trace_match.group('event'), 'info': int(trace_match.group('info')), 'cycles': raw_cyc, 'unwrapped_cycles': unwrapped_cyc, - 'state': trace_match.group('state') + 'state': trace_match.group('state'), + 'line_num': line_idx, + 'device': device }) f.close() @@ -274,27 +421,24 @@ def generate_perfetto_trace(filtered_ops, trace_events, output_path): logger.warning("No operators found after filtering.") return - # Compute average frequency - frequencies = [] - for op in filtered_ops: - if op['usec'] > 0 and op['cycles'] > 0: - frequencies.append(op['cycles'] / op['usec']) - avg_freq_mhz = statistics.mean(frequencies) if frequencies else 1000.0 - if avg_freq_mhz <= 0: - avg_freq_mhz = 1000.0 - # Assign start and end cycles to each operator for op in filtered_ops: op['start_cycles'] = op['unwrapped_cycles_start'] - op['end_cycles'] = op['start_cycles'] + op['cycles'] + op['end_cycles'] = op['start_cycles'] + op['cycles'] if op['start_cycles'] is not None else None - global_min_cyc = min(op['start_cycles'] for op in filtered_ops if op['start_cycles'] is not None) + # Get list of unique devices present in the operations + unique_devices = sorted(list(set(op['device'] for op in filtered_ops))) + device_to_idx = {dev: idx for idx, dev in enumerate(unique_devices)} + time_mappers = {dev: DeviceTimeMapper(dev, filtered_ops) for dev in unique_devices} # Process events completed_events = [] if trace_events: trace_events = sorted(trace_events, key=lambda e: e['unwrapped_cycles']) - one_usec_cycles = max(avg_freq_mhz, 1.0) + + one_usec_cycles = {} + for dev in unique_devices: + one_usec_cycles[dev] = max(time_mappers[dev].get_freq(), 1.0) active_starts = {} for e in trace_events: @@ -303,31 +447,36 @@ def generate_perfetto_trace(filtered_ops, trace_events, output_path): info = e['info'] state = e['state'] cyc = e['unwrapped_cycles'] + dev = e['device'] - key = (t, evt, info) + key = (dev, t, evt, info) if state == 'start': # Handle missing stop (start followed by another start) if key in active_starts: - prev_start = active_starts[key] + prev_e = active_starts[key] completed_events.append({ 'thread': t, 'event': evt, 'info': info, - 'start_cyc': prev_start, - 'end_cyc': prev_start + one_usec_cycles, + 'start_cyc': prev_e['unwrapped_cycles'], + 'end_cyc': prev_e['unwrapped_cycles'] + one_usec_cycles.get(dev, 1000.0), + 'line_num': prev_e.get('line_num'), 'missing_stop': True, + 'device': dev }) - active_starts[key] = cyc + active_starts[key] = e elif state == 'stop': if key in active_starts: - start_cyc = active_starts[key] + prev_e = active_starts[key] del active_starts[key] completed_events.append({ 'thread': t, 'event': evt, 'info': info, - 'start_cyc': start_cyc, + 'start_cyc': prev_e['unwrapped_cycles'], 'end_cyc': cyc, + 'line_num': prev_e.get('line_num'), + 'device': dev }) else: # Handle missing start (stop without start) @@ -335,31 +484,36 @@ def generate_perfetto_trace(filtered_ops, trace_events, output_path): 'thread': t, 'event': evt, 'info': info, - 'start_cyc': cyc - one_usec_cycles, + 'start_cyc': cyc - one_usec_cycles.get(dev, 1000.0), 'end_cyc': cyc, + 'line_num': e.get('line_num'), 'missing_start': True, + 'device': dev }) # Clear remaining unmatched starts - for key, start_cyc in active_starts.items(): - t, evt, info = key + for key, prev_e in active_starts.items(): + dev, t, evt, info = key completed_events.append({ 'thread': t, 'event': evt, 'info': info, - 'start_cyc': start_cyc, - 'end_cyc': start_cyc + one_usec_cycles, + 'start_cyc': prev_e['unwrapped_cycles'], + 'end_cyc': prev_e['unwrapped_cycles'] + one_usec_cycles.get(dev, 1000.0), + 'line_num': prev_e.get('line_num'), 'missing_stop': True, + 'device': dev }) completed_events.sort(key=lambda e: e['start_cyc']) - # Convert event times to microseconds and apply clamp rounded to 1ns resolution (3 decimals) + # Convert event times to nanoseconds using per-device / per-batch time mapper for e in completed_events: - start_us = (e['start_cyc'] - global_min_cyc) / avg_freq_mhz - dur_us = (e['end_cyc'] - e['start_cyc']) / avg_freq_mhz - e['ts_ns'] = int(round(start_us * 1000)) - e['dur_ns'] = int(round(max(dur_us, 0.1) * 1000)) + dev = e['device'] + tm = time_mappers[dev] + e['ts_ns'] = tm.cycle_to_ns(e['start_cyc']) + dur_ns = tm.dur_cycles_to_ns(e['start_cyc'], e['end_cyc'] - e['start_cyc']) + e['dur_ns'] = max(dur_ns, 100) # Allocate slots (sub-tracks) to prevent overlaps on same virtual track active_slots = defaultdict(list) @@ -368,14 +522,15 @@ def generate_perfetto_trace(filtered_ops, trace_events, output_path): evt = e['event'] ts = e['ts_ns'] dur = e['dur_ns'] + dev = e['device'] norm_evt = normalize_event_name(evt, e['info']) if norm_evt == "DMA": - track_key = (t, "DMA") + track_key = (dev, t, "DMA") elif t == 10: - track_key = (t, "HMX") + track_key = (dev, t, "HMX") else: - track_key = (t, "HVX") + track_key = (dev, t, "HVX") slots = active_slots[track_key] allocated_slot = -1 @@ -395,6 +550,7 @@ def generate_perfetto_trace(filtered_ops, trace_events, output_path): t = e['thread'] evt = e['event'] slot = e['slot'] + dev = e['device'] norm_evt = normalize_event_name(evt, e['info']) if norm_evt == "DMA": @@ -408,56 +564,69 @@ def generate_perfetto_trace(filtered_ops, trace_events, output_path): evt_id = 2 t_sort = 1 if t == 10 else t + 2 + dev_idx = device_to_idx[dev] + # Unique UUID for each sub-track if t == 10: - uuid = 20 # HMX thread track UUID + uuid = dev_idx * 10000000 + 20 # HMX thread track UUID else: - uuid = int(t_sort * 1000000 + evt_id * 1000 + slot) + uuid = int(dev_idx * 10000000 + t_sort * 1000000 + evt_id * 1000 + slot) e['uuid'] = uuid - used_tracks[uuid] = (t, track_evt, slot) + used_tracks[uuid] = (dev, t, track_evt, slot) with open(output_path, "wb") as f: - # Define Process with EXPLICIT child sorting - proc_desc = make_process_descriptor(1, "HTP NPU") - proc_packet = make_trace_packet(0, track_descriptor=make_track_descriptor(1, process=proc_desc, child_ordering=3)) - write_trace_packet_to_file(f, proc_packet) - - # Define Operators Track (UUID = 2) as a thread track at rank 1, tid 8 - op_thread_desc = make_thread_descriptor(1, 8, "Ops", sort_index=1) - op_packet = make_trace_packet(0, track_descriptor=make_track_descriptor(2, parent_uuid=1, thread=op_thread_desc)) - write_trace_packet_to_file(f, op_packet) - - # Define HMX Thread Track (UUID = 20) at rank 2, tid 9 - hmx_thread_desc = make_thread_descriptor(1, 9, "HMX", sort_index=2) - hmx_packet = make_trace_packet(0, track_descriptor=make_track_descriptor(20, parent_uuid=1, thread=hmx_thread_desc)) - write_trace_packet_to_file(f, hmx_packet) - - # Define Thread Tracks (T0, T1, ..., T9) - unique_threads = sorted(list(set(t for (t, _, _) in used_tracks.values() if t != 10))) - for t in unique_threads: - thread_uuid = 10 + t - thread_name = f"T{t}" - # Sort order starts from index 3 (T0 -> 3, T1 -> 4, etc.) - sort_index = 3 + t - tid = 10 + t - thread_desc = make_thread_descriptor(1, tid, thread_name, sort_index=sort_index) - thread_packet = make_trace_packet(0, track_descriptor=make_track_descriptor( - thread_uuid, - parent_uuid=1, - thread=thread_desc, - sibling_order_rank=sort_index, - child_ordering=3 # Explicit child sorting for sub-tracks - )) - write_trace_packet_to_file(f, thread_packet) + for dev in unique_devices: + dev_idx = device_to_idx[dev] + pid = dev_idx + 1 + proc_uuid = dev_idx * 10000000 + 1 + + # Define Process with EXPLICIT child sorting + proc_name = dev + proc_desc = make_process_descriptor(pid, proc_name) + proc_packet = make_trace_packet(0, track_descriptor=make_track_descriptor(proc_uuid, process=proc_desc, child_ordering=3)) + write_trace_packet_to_file(f, proc_packet) + + # Define Operators Track as a thread track + op_track_uuid = dev_idx * 10000000 + 2 + op_tid = pid * 100 + 8 + op_thread_desc = make_thread_descriptor(pid, op_tid, "Ops", sort_index=1) + op_packet = make_trace_packet(0, track_descriptor=make_track_descriptor(op_track_uuid, parent_uuid=proc_uuid, thread=op_thread_desc)) + write_trace_packet_to_file(f, op_packet) + + # Define HMX Thread Track at rank 2 + hmx_track_uuid = dev_idx * 10000000 + 20 + hmx_tid = pid * 100 + 9 + hmx_thread_desc = make_thread_descriptor(pid, hmx_tid, "HMX", sort_index=2) + hmx_packet = make_trace_packet(0, track_descriptor=make_track_descriptor(hmx_track_uuid, parent_uuid=proc_uuid, thread=hmx_thread_desc)) + write_trace_packet_to_file(f, hmx_packet) + + # Define Thread Tracks (T0, T1, ..., T9) for this device + dev_used_tracks = {uuid: val for uuid, val in used_tracks.items() if val[0] == dev} + unique_threads = sorted(list(set(t for (_, t, _, _) in dev_used_tracks.values() if t != 10))) + for t in unique_threads: + thread_uuid = dev_idx * 10000000 + 10 + t + thread_name = f"T{t}" + sort_index = 3 + t + tid = pid * 100 + 10 + t + thread_desc = make_thread_descriptor(pid, tid, thread_name, sort_index=sort_index) + thread_packet = make_trace_packet(0, track_descriptor=make_track_descriptor( + thread_uuid, + parent_uuid=proc_uuid, + thread=thread_desc, + sibling_order_rank=sort_index, + child_ordering=3 # Explicit child sorting for sub-tracks + )) + write_trace_packet_to_file(f, thread_packet) # Define Track descriptors for sub-tracks parented to thread tracks for uuid in sorted(used_tracks.keys()): - if uuid == 20: + dev, t, evt, slot = used_tracks[uuid] + dev_idx = device_to_idx[dev] + if t == 10: continue - t, evt, slot = used_tracks[uuid] name = f"T{t} {evt}" rank = 0 if evt == "HVX" else 1 - parent_thread_uuid = 10 + t + parent_thread_uuid = dev_idx * 10000000 + 10 + t # Sibling merge behavior: 1 (SIBLING_MERGE_BEHAVIOR_BY_TRACK_NAME) track_desc = make_track_descriptor( uuid=uuid, @@ -470,15 +639,18 @@ def generate_perfetto_trace(filtered_ops, trace_events, output_path): write_trace_packet_to_file(f, track_packet) # Emit Operators - last_op_end_ns = 0 + last_op_end_ns = defaultdict(int) for op in filtered_ops: - op_start_ns = int(round(((op['start_cycles'] - global_min_cyc) / avg_freq_mhz) * 1000)) - op_dur_ns = int(round((op['cycles'] / avg_freq_mhz) * 1000)) + dev = op['device'] + dev_idx = device_to_idx[dev] + tm = time_mappers[dev] + op_start_ns = tm.cycle_to_ns(op['start_cycles']) + op_dur_ns = tm.dur_cycles_to_ns(op['start_cycles'], op['cycles']) if op['name'] != "OPBATCH": - if op_start_ns < last_op_end_ns: - op_start_ns = last_op_end_ns + if op_start_ns < last_op_end_ns[dev]: + op_start_ns = last_op_end_ns[dev] clamped_dur = max(op_dur_ns, 100) # Clamp to 100ns (0.1us) - last_op_end_ns = op_start_ns + clamped_dur + last_op_end_ns[dev] = op_start_ns + clamped_dur else: clamped_dur = max(op_dur_ns, 100) @@ -495,24 +667,41 @@ def generate_perfetto_trace(filtered_ops, trace_events, output_path): if 'evt' in op and op['evt']: debug_annots.append(make_debug_annotation("evt", string_val=op['evt'])) + op_track_uuid = dev_idx * 10000000 + 2 + # Slice Begin - evt_begin = make_track_event(1, 2, name=f"{op['name']} ({op['dims']})", category="operator", debug_annotations=debug_annots) + evt_begin = make_track_event(1, op_track_uuid, name=f"{op['name']} ({op['dims']})", category="operator", debug_annotations=debug_annots) packet_begin = make_trace_packet(op_start_ns, track_event=evt_begin) write_trace_packet_to_file(f, packet_begin) # Slice End - evt_end = make_track_event(2, 2) + evt_end = make_track_event(2, op_track_uuid) packet_end = make_trace_packet(op_start_ns + clamped_dur, track_event=evt_end) write_trace_packet_to_file(f, packet_end) # Emit Thread Trace Events for e in completed_events: norm_name = normalize_event_name(e['event'], e['info']) - name = f"DMA {e['info']}" if norm_name == "DMA" else norm_name + if norm_name == "DMA": + name = f"DMA {e['info']}" + elif norm_name == "FENCE": + name = f"FENCE {e['info']}" if e.get('info') is not None and e['info'] != 0 else "FENCE" + else: + name = norm_name + if e.get('missing_start') or e.get('missing_stop'): name += "!" debug_annots = [] + if 'line_num' in e and e['line_num'] is not None: + debug_annots.append(make_debug_annotation("line", int_val=e['line_num'])) + if norm_name == "FENCE" and e.get('info') is not None: + debug_annots.append(make_debug_annotation("seq", int_val=e['info'])) + elif norm_name == "DMA" and e.get('info') is not None: + debug_annots.append(make_debug_annotation("channel", int_val=e['info'])) + elif e.get('info') is not None and e['info'] != 0: + debug_annots.append(make_debug_annotation("info", int_val=e['info'])) + if e.get('missing_start'): debug_annots.append(make_debug_annotation("missing_start", string_val="true")) if e.get('missing_stop'): @@ -536,6 +725,7 @@ def main(): parser.add_argument("logfile", help="Path to hex-log profile file") parser.add_argument("-o", "--output", default="optrace.perfetto-trace", help="Output trace file path (default: optrace.perfetto-trace)") parser.add_argument("--filter", type=str, help="Regex filter matching against the original profile-op line") + parser.add_argument("--device", type=str, help="Device to filter by (e.g. HTP0, HTP0:0) or 'split' to generate separate files per device") group = parser.add_mutually_exclusive_group() group.add_argument("--head", type=int, help="Limit to first N ops") @@ -544,7 +734,21 @@ def main(): args = parser.parse_args() logging.basicConfig(level=logging.INFO, format='%(message)s') - ops, traces = parse_log(args.logfile) + op_filter_re = None + if args.filter: + try: + op_filter_re = re.compile(args.filter) + except re.error as e: + logger.error(f"Invalid regex filter: {e}") + sys.exit(1) + + limit = args.head if args.head is not None else None + device_filter = args.device if (args.device and args.device != "split") else None + ops, traces = parse_log(args.logfile, limit=limit, device_filter=device_filter, op_filter_re=op_filter_re) + + if args.device and args.device != "split": + ops = [op for op in ops if device_matches(op['device'], args.device)] + traces = [t for t in traces if device_matches(t['device'], args.device)] if args.filter: try: @@ -554,35 +758,60 @@ def main(): sys.exit(1) ops = [op for op in ops if filter_re.search(op['op_text'])] - if args.head is not None: - ops = ops[:args.head] - elif args.tail is not None: - ops = ops[-args.tail:] + if args.head is not None or args.tail is not None: + ops_by_dev = defaultdict(list) + for op in ops: + ops_by_dev[op['device']].append(op) + + filtered_ops = [] + for dev in sorted(ops_by_dev.keys()): + dev_ops = ops_by_dev[dev] + if args.head is not None: + dev_ops = dev_ops[:args.head] + elif args.tail is not None: + dev_ops = dev_ops[-args.tail:] + filtered_ops.extend(dev_ops) + ops = filtered_ops if args.filter or args.head is not None or args.tail is not None: - valid_ranges = [] + # Group valid ranges by device + valid_ranges_by_dev = defaultdict(list) for op in ops: start_cyc = op['unwrapped_cycles_start'] end_cyc = start_cyc + op['cycles'] if start_cyc is not None else None if start_cyc is not None and end_cyc is not None: - valid_ranges.append((start_cyc, end_cyc)) + valid_ranges_by_dev[op['device']].append((start_cyc, end_cyc)) - valid_ranges.sort(key=lambda r: r[0]) - range_starts = [r[0] for r in valid_ranges] + for dev in valid_ranges_by_dev: + valid_ranges_by_dev[dev].sort(key=lambda r: r[0]) + + range_starts_by_dev = {dev: [r[0] for r in ranges] for dev, ranges in valid_ranges_by_dev.items()} filtered_traces = [] for e in traces: cyc = e['unwrapped_cycles'] if cyc is None: continue + dev = e['device'] + range_starts = range_starts_by_dev.get(dev) + if not range_starts: + continue idx = bisect.bisect_right(range_starts, cyc) - 1 if idx >= 0: - start, end = valid_ranges[idx] + start, end = valid_ranges_by_dev[dev][idx] if start <= cyc <= end: filtered_traces.append(e) traces = filtered_traces - generate_perfetto_trace(ops, traces, args.output) + if args.device == "split": + unique_devices = sorted(list(set(op['device'] for op in ops))) + for dev in unique_devices: + dev_ops = [op for op in ops if device_matches(op['device'], dev)] + dev_traces = [t for t in traces if device_matches(t['device'], dev)] + out_path = get_split_output_path(args.output, dev) + generate_perfetto_trace(dev_ops, dev_traces, out_path) + else: + generate_perfetto_trace(ops, traces, args.output) if __name__ == "__main__": diff --git a/scripts/snapdragon/run.py b/scripts/snapdragon/run.py new file mode 100755 index 000000000000..d639b16b8250 --- /dev/null +++ b/scripts/snapdragon/run.py @@ -0,0 +1,405 @@ +#!/usr/bin/env python3 +# +# Run llama.cpp tools on Snapdragon devices (natively, via ADB, or SSH). +# + +import sys +import os +import argparse +import subprocess +import platform +import shlex +import logging + +logger = logging.getLogger("run") + + +def parse_target(target_str): + if not target_str: + return None, None + if target_str.startswith("adb") or target_str.startswith("android"): + parts = target_str.split(":", 1) + serial = parts[1] if len(parts) > 1 else None + return "android", serial + elif target_str.startswith("lnx") or target_str.startswith("linux") or target_str.startswith("ubuntu"): + parts = target_str.split(":", 1) + host = parts[1] if len(parts) > 1 else None + return "linux", host + elif target_str in ("wos", "windows"): + return "windows", None + else: + return None, None + + +def shlex_join(args_list): + if hasattr(shlex, 'join'): + return shlex.join(args_list) + import pipes + return " ".join(pipes.quote(x) for x in args_list) + + +def main(): + logging.basicConfig(level=logging.INFO, format='%(message)s') + # Split arguments at '--' + if '--' in sys.argv: + idx = sys.argv.index('--') + run_args = sys.argv[1:idx] + cmd_args = sys.argv[idx + 1:] + else: + run_args = sys.argv[1:] + cmd_args = [] + + parser = argparse.ArgumentParser( + description="Unified runner for llama.cpp tools on Snapdragon (natively, via ADB, or via SSH)." + ) + parser.add_argument("--target", help="Execution target (e.g. android[:serial]/adb[:serial], linux:[user@]host/lnx:[user@]host/ubuntu:[user@]host, windows/wos) (default: local run)") + parser.add_argument("--target-dir", help="Target directory on the device (default: /data/local/tmp/llama.cpp for Android, ~/llama.cpp for Linux)") + parser.add_argument("--install-dir", help="Install directory name (defaults to pkg-TARGET or pkg-TARGET-dbg prefix based on target)") + parser.add_argument("--debug", action="store_true", help="Use debug build (defaults to pkg-TARGET-dbg folder)") + parser.add_argument("--devices", "--device", "-d", help="Select execution devices (split into NPU and OpenCL GPUs automatically, default: HTP0)") + parser.add_argument("--verbose", help="Verbose level (enables both Hexagon and OpenCL kernel cache debugging)") + parser.add_argument("--profile", help="Profiling flag (enables Hexagon profiling and OpenCL autotuning)") + parser.add_argument("--sched-debug", action="store_true", help="Enable GGML/llama.cpp scheduler debug output (GGML_SCHED_DEBUG=2)") + parser.add_argument("--mtmd-device", help="Specify the backend device ID for Multi-Threaded Multi-Device setup (MTMD_BACKEND_DEVICE)") + + # Hexagon specific parameters + parser.add_argument("--hex-verbose", help="Enable verbose logging (GGML_HEXAGON_VERBOSE)") + parser.add_argument("--hex-profile", help="Enable NPU/Hexagon profiling and performance metrics print (GGML_HEXAGON_PROFILE)") + parser.add_argument("--hex-nhvx", help="Number of HVX units to use (GGML_HEXAGON_NHVX)") + parser.add_argument("--hex-nhmx", help="Number of HMX units to use. 0 disables HMX power-up (GGML_HEXAGON_NHMX)") + parser.add_argument("--hex-hostbuf", help="Enable host buffers (GGML_HEXAGON_HOSTBUF)") + parser.add_argument("--hex-opbatch", help="Maximum number of operations to batch into a single HTP execution (GGML_HEXAGON_OPBATCH)") + parser.add_argument("--hex-opqueue", help="Size of the asynchronous NPU operation queue (GGML_HEXAGON_OPQUEUE)") + parser.add_argument("--hex-oppoll", default="1", help="Enable (1) or Disable (0) polling for NPU opbatch completion (GGML_HEXAGON_OPPOLL) (default: 1)") + parser.add_argument("--hex-opfilter", help="Regex pattern to filter/select which operators are offloaded to NPU (GGML_HEXAGON_OPFILTER)") + parser.add_argument("--hex-opfusion", help="NPU graph node fusion optimization level (0: disabled, 1: enabled) (GGML_HEXAGON_OPFUSION)") + parser.add_argument("--hex-vmem", help="Maximum NPU VMEM size limit in MB to allocate (GGML_HEXAGON_VMEM)") + parser.add_argument("--hex-mbuf", help="Maximum host buffer size limit in MB to allocate (GGML_HEXAGON_MBUF)") + parser.add_argument("--hex-mm-select", help="Select MUL_MAT and MUL_MAT_ID kernel (GGML_HEXAGON_MM_SELECT) 3:HMX,2:HVX-tiled,1:HVX-flat,0:disable") + parser.add_argument("--hex-fa-select", help="Select Flash Attention kernel (GGML_HEXAGON_FA_SELECT) 2:HMX,1:HVX,0:disable") + parser.add_argument("--hex-ar-select", help="Select All-Reduce kernel (GGML_HEXAGON_AR_SELECT) 1:enable,0:disable") + parser.add_argument("--hex-etm", help="Enable Embedded Trace Macrocell hardware tracing / trace logging (GGML_HEXAGON_ETM)") + parser.add_argument("--hex-arch", help="Target Hexagon NPU architecture version override (v73, v75, v79, v81, etc.) (GGML_HEXAGON_ARCH)") + parser.add_argument("--hex-optrace", help="Trace buffer size in number of records (GGML_HEXAGON_OPTRACE)") + + # OpenCL specific parameters + parser.add_argument("--cl-platform", help="Select OpenCL platform name/regex (e.g. Qualified Qualcomm OpenCL platform) (GGML_OPENCL_PLATFORM)") + parser.add_argument("--cl-device", help="Select OpenCL device name/regex (e.g. Adreno GPU) (GGML_OPENCL_DEVICE)") + parser.add_argument("--cl-opfilter", help="Regex pattern to filter/select which operators are offloaded to OpenCL (GGML_OPENCL_OPFILTER)") + parser.add_argument("--cl-disable-fusion", action="store_true", help="Disable OpenCL kernel fusion optimizations (GGML_OPENCL_DISABLE_FUSION)") + parser.add_argument("--cl-cache-dir", help="Directory path to store compiled OpenCL program binaries (GGML_OPENCL_KERNEL_CACHE_DIR)") + parser.add_argument("--cl-cache-debug", help="Enable verbose debugging logs for the kernel caching system (GGML_OPENCL_KERNEL_CACHE_DEBUG)") + parser.add_argument("--cl-fa-tune", action="store_true", help="Enable automatic Flash Attention kernel autotuning (GGML_OPENCL_FA_TUNE)") + parser.add_argument("--cl-adreno-xmem", action="store_true", help="Enforce matmul using texture/image (xmem) memory paths on Adreno GPUs (GGML_OPENCL_ADRENO_XMEM_GEMM)") + parser.add_argument("--cl-adreno-large-buffer", action="store_true", help="Allow allocating larger buffer sizes on Adreno GPUs (GGML_OPENCL_ADRENO_USE_LARGE_BUFFER)") + + args = parser.parse_args(run_args) + + if not cmd_args: + parser.print_help() + logger.error("\nError: No command specified after '--'") + sys.exit(1) + + target_type = None + target_val = None + target_prefix = None + if args.target: + target_type, target_val = parse_target(args.target) + if not target_type: + logger.error(f"Error: Invalid target format '{args.target}'. Must be android[:serial]/adb[:serial], linux:[user@]host/lnx:[user@]host/ubuntu:[user@]host, or windows/wos.") + sys.exit(1) + target_prefix = args.target.split(":", 1)[0] + + # Resolve install directory + install_dir = args.install_dir + if not install_dir: + if target_prefix: + suffix = "-dbg" if args.debug else "" + install_dir = f"pkg-{target_prefix}{suffix}" + else: + # Smart branch folder detection for local run if default is not set + prefixes = ("wos", "windows", "lnx", "linux", "ubuntu", "adb", "android") + suffixes = ("-dbg", "") if args.debug else ("", "-dbg") + found = False + for suffix in suffixes: + for prefix in prefixes: + test_path = f"./pkg-{prefix}{suffix}/llama.cpp" + if os.path.exists(test_path): + install_dir = f"pkg-{prefix}{suffix}" + found = True + break + if found: + break + if not install_dir: + install_dir = "pkg-android" # Fallback default + + # Host side package path + package_path = os.path.join(install_dir, "llama.cpp") + + # Environment variables to map + env_vars = {} + + def set_env(env_name, opt_val): + if opt_val is not None: + env_vars[env_name] = str(opt_val) + elif env_name in os.environ: + env_vars[env_name] = os.environ[env_name] + + # Resolve and filter devices (HTP vs OpenCL) + devices_val = args.devices if args.devices is not None else "HTP0" + if devices_val.isdigit(): + hex_devices = devices_val + cl_device = "" + else: + parts = [p.strip() for p in devices_val.split(",")] + # Any device containing "htp" is Hexagon, rest is OpenCL + hex_parts = [p for p in parts if "htp" in p.lower()] + cl_parts = [p for p in parts if "htp" not in p.lower()] + hex_devices = ",".join(hex_parts) + cl_device = ",".join(cl_parts) + + # Set Hexagon devices + if hex_devices: + env_vars["GGML_HEXAGON_DEVICES"] = hex_devices + elif "GGML_HEXAGON_DEVICES" in os.environ: + env_vars["GGML_HEXAGON_DEVICES"] = os.environ["GGML_HEXAGON_DEVICES"] + + # Set OpenCL device (unless overridden by --cl-device) + final_cl_device = args.cl_device if args.cl_device is not None else cl_device + if final_cl_device: + env_vars["GGML_OPENCL_DEVICE"] = final_cl_device + elif "GGML_OPENCL_DEVICE" in os.environ: + env_vars["GGML_OPENCL_DEVICE"] = os.environ["GGML_OPENCL_DEVICE"] + + # Map shared & backend-specific parameters with correct overrides + + # Verbose logging mapping + hex_verbose_val = args.hex_verbose if args.hex_verbose is not None else args.verbose + set_env("GGML_HEXAGON_VERBOSE", hex_verbose_val) + + cl_cache_debug_val = args.cl_cache_debug if args.cl_cache_debug is not None else args.verbose + set_env("GGML_OPENCL_KERNEL_CACHE_DEBUG", cl_cache_debug_val) + + # Profiling mapping + hex_profile_val = args.hex_profile if args.hex_profile is not None else args.profile + set_env("GGML_HEXAGON_PROFILE", hex_profile_val) + + if args.cl_fa_tune or args.profile is not None: + env_vars["GGML_OPENCL_FA_TUNE"] = "1" + elif "GGML_OPENCL_FA_TUNE" in os.environ: + env_vars["GGML_OPENCL_FA_TUNE"] = os.environ["GGML_OPENCL_FA_TUNE"] + + # Other Hexagon environment variables + set_env("GGML_HEXAGON_NHVX", args.hex_nhvx) + set_env("GGML_HEXAGON_NHMX", args.hex_nhmx) + set_env("GGML_HEXAGON_HOSTBUF", args.hex_hostbuf) + set_env("GGML_HEXAGON_OPBATCH", args.hex_opbatch) + set_env("GGML_HEXAGON_OPQUEUE", args.hex_opqueue) + set_env("GGML_HEXAGON_OPPOLL", args.hex_oppoll) + set_env("GGML_HEXAGON_OPFILTER", args.hex_opfilter) + set_env("GGML_HEXAGON_OPFUSION", args.hex_opfusion) + set_env("GGML_HEXAGON_VMEM", args.hex_vmem) + set_env("GGML_HEXAGON_MBUF", args.hex_mbuf) + set_env("GGML_HEXAGON_MM_SELECT", args.hex_mm_select) + set_env("GGML_HEXAGON_FA_SELECT", args.hex_fa_select) + set_env("GGML_HEXAGON_AR_SELECT", args.hex_ar_select) + set_env("GGML_HEXAGON_ETM", args.hex_etm) + set_env("GGML_HEXAGON_ARCH", args.hex_arch) + set_env("GGML_HEXAGON_OPTRACE", args.hex_optrace) + set_env("MTMD_BACKEND_DEVICE", args.mtmd_device) + + # OpenCL environment variables + set_env("GGML_OPENCL_PLATFORM", args.cl_platform) + set_env("GGML_OPENCL_OPFILTER", args.cl_opfilter) + set_env("GGML_OPENCL_KERNEL_CACHE_DIR", args.cl_cache_dir) + + if args.cl_disable_fusion: + env_vars["GGML_OPENCL_DISABLE_FUSION"] = "1" + elif "GGML_OPENCL_DISABLE_FUSION" in os.environ: + env_vars["GGML_OPENCL_DISABLE_FUSION"] = os.environ["GGML_OPENCL_DISABLE_FUSION"] + + if args.cl_adreno_xmem: + env_vars["GGML_OPENCL_ADRENO_XMEM_GEMM"] = "1" + elif "GGML_OPENCL_ADRENO_XMEM_GEMM" in os.environ: + env_vars["GGML_OPENCL_ADRENO_XMEM_GEMM"] = os.environ["GGML_OPENCL_ADRENO_XMEM_GEMM"] + + if args.cl_adreno_large_buffer: + env_vars["GGML_OPENCL_ADRENO_USE_LARGE_BUFFER"] = "1" + elif "GGML_OPENCL_ADRENO_USE_LARGE_BUFFER" in os.environ: + env_vars["GGML_OPENCL_ADRENO_USE_LARGE_BUFFER"] = os.environ["GGML_OPENCL_ADRENO_USE_LARGE_BUFFER"] + + if args.sched_debug: + env_vars["GGML_SCHED_DEBUG"] = "2" + + # Resolve executable path + executable = cmd_args[0] + known_binaries = ["llama-cli", "llama-bench", "llama-completion", "llama-mtmd-cli", "test-backend-ops"] + if executable in known_binaries: + if target_type in ("android", "linux"): + resolved_exec = f"./bin/{executable}" + else: + if platform.system() == "Windows": + resolved_exec = os.path.normpath(os.path.join(package_path, "bin", f"{executable}.exe")) + else: + resolved_exec = os.path.normpath(os.path.join(package_path, "bin", executable)) + cmd_args[0] = resolved_exec + + # Infer device string to pass to the tool + basename = os.path.basename(executable) + if basename.endswith(".exe"): + basename = basename[:-4] + + device_val = None + if basename == "test-backend-ops": + for i in range(len(cmd_args)): + if cmd_args[i] in ("-p", "--params") and i + 1 < len(cmd_args): + val = cmd_args[i + 1] + new_val = "" + for j, char in enumerate(val): + if char in ('[', ']'): + if j > 0 and val[j - 1] == '\\': + new_val += char + else: + new_val += '\\' + char + else: + new_val += char + cmd_args[i + 1] = new_val + + has_b = any(arg == "-b" for arg in cmd_args) + if not has_b: + if args.devices: + if args.devices.isdigit(): + n = int(args.devices) + device_val = ",".join(f"HTP{i}" for i in range(n)) + else: + device_val = args.devices + elif "D" in os.environ: + device_val = os.environ["D"] + elif "DEVICE" in os.environ: + device_val = os.environ["DEVICE"] + else: + device_val = "HTP0" + if device_val: + cmd_args += ["-b", device_val] + else: + has_device = any(arg.startswith("--device") for arg in cmd_args) + if not has_device: + if args.devices: + if args.devices.isdigit(): + n = int(args.devices) + device_val = ",".join(f"HTP{i}" for i in range(n)) + else: + device_val = args.devices + elif "D" in os.environ: + device_val = os.environ["D"] + elif "DEVICE" in os.environ: + device_val = os.environ["DEVICE"] + else: + device_val = "HTP0" + if device_val: + cmd_args += ["--device", device_val] + + # Automatically add -v to known llama tools if sched-debug, verbose, or profile are set + verbose_trigger = ( + args.sched_debug + or args.verbose is not None + or args.profile is not None + or args.hex_verbose is not None + or args.hex_profile is not None + or args.hex_optrace is not None + ) + if verbose_trigger and basename in ("llama-cli", "llama-completion", "llama-bench", "llama-server", "llama-mtmd-cli"): + if "-v" not in cmd_args and "--verbose" not in cmd_args: + cmd_args.append("-v") + + # Inject defaults for llama-cli, llama-completion, and llama-server if not overridden by the user + if basename in ("llama-cli", "llama-completion", "llama-server"): + if "-ngl" not in cmd_args and "--n-gpu-layers" not in cmd_args: + cmd_args += ["-ngl", "99"] + if "--ubatch-size" not in cmd_args and "-ub" not in cmd_args: + cmd_args += ["--ubatch-size", "1024"] + if "-fa" not in cmd_args and "--flash-attn" not in cmd_args: + cmd_args += ["-fa", "on"] + + if basename in ("llama-cli", "llama-completion", "llama-server", "llama-bench"): + if "-t" not in cmd_args and "--threads" not in cmd_args: + cmd_args += ["-t", "6"] + + # Resolve target directory on device + target_dir = args.target_dir + if not target_dir: + target_dir = "/data/local/tmp/llama.cpp" if target_type == "android" else "~/llama.cpp" + + if target_type == "android": + # Run via ADB + adb_base = ["adb"] + if target_val: # serial + adb_base += ["-s", target_val] + + env_parts = [ + "LD_LIBRARY_PATH=./lib", + "ADSP_LIBRARY_PATH=./lib" + ] + for k, v in env_vars.items(): + env_parts.append(f"{k}={v}") + env_str = " ".join(env_parts) + + cmd_str = shlex_join(cmd_args) + adb_shell_cmd = f"cd {target_dir} && ulimit -c unlimited && {env_str} {cmd_str}" + full_cmd = adb_base + ["shell", adb_shell_cmd] + + logger.info(f"+ {' '.join(full_cmd)}") + res = subprocess.run(full_cmd) + sys.exit(res.returncode) + + elif target_type == "linux": + ssh_host = target_val + if not ssh_host: + logger.error("Error: SSH host not specified in target (e.g. use linux:user@host, lnx:user@host, or ubuntu:user@host). Cannot execute.") + sys.exit(1) + + # Linux remote run via SSH + env_parts = [ + "LD_LIBRARY_PATH=./lib", + "ADSP_LIBRARY_PATH=./lib" + ] + for k, v in env_vars.items(): + env_parts.append(f"{k}={v}") + env_str = " ".join(env_parts) + + cmd_str = shlex_join(cmd_args) + ssh_shell_cmd = f"cd {target_dir} && ulimit -c unlimited && {env_str} {cmd_str}" + full_cmd = ["ssh", ssh_host, ssh_shell_cmd] + + logger.info(f"+ {' '.join(full_cmd)}") + res = subprocess.run(full_cmd) + sys.exit(res.returncode) + + elif target_type == "windows": + logger.info("Windows target execution is currently a stub.") + sys.exit(0) + + else: + # Run locally + local_env = os.environ.copy() + lib_dir = os.path.normpath(os.path.join(package_path, "lib")) + local_env["ADSP_LIBRARY_PATH"] = lib_dir + if platform.system() == "Windows": + local_env["PATH"] = lib_dir + os.path.pathsep + local_env.get("PATH", "") + else: + local_env["LD_LIBRARY_PATH"] = lib_dir + os.path.pathsep + local_env.get("LD_LIBRARY_PATH", "") + + for k, v in env_vars.items(): + local_env[k] = v + + logger.info(f"+ {shlex_join(cmd_args)}") + res = subprocess.run(cmd_args, env=local_env) + sys.exit(res.returncode) + + +if __name__ == "__main__": + try: + main() + except KeyboardInterrupt: + logger.info("\nInterrupted by user.") + sys.exit(130) diff --git a/scripts/snapdragon/windows/run-bench.ps1 b/scripts/snapdragon/windows/run-bench.ps1 deleted file mode 100644 index 6eb656e66d32..000000000000 --- a/scripts/snapdragon/windows/run-bench.ps1 +++ /dev/null @@ -1,48 +0,0 @@ - -#!/usr/bin/env pwsh - -# Basedir on device -$basedir=".\pkg-snapdragon" - -$cli_opts=$args - -$model="Llama-3.2-3B-Instruct-Q4_0.gguf" -if ($null -ne $env:M) { - $model=$env:M -} - -$device="HTP0" -if ($null -ne $env:D) { - $device=$env:D -} - -if ($null -ne $env:V) { - $env:GGML_HEXAGON_VERBOSE=$env:V -} - -if ($null -ne $env:PROF) { - $env:GGML_HEXAGON_PROFILE=$env:PROF -} - -if ($null -ne $env:OPSTAGE) { - $env:GGML_HEXAGON_OPSTAGE=$env:OPSTAGE -} - -if ($null -ne $env:NHVX) { - $env:GGML_HEXAGON_NHVX=$env:NHVX -} - -if ($null -ne $env:NDEV) { - $env:GGML_HEXAGON_NDEV=$env:NDEV -} - -if ($null -ne $env:HB) { - $env:GGML_HEXAGON_HOSTBUF=$env:HB -} - -$env:ADSP_LIBRARY_PATH="$basedir\lib" - -& "$basedir\bin\llama-bench.exe" ` - --load-mode none -m $basedir\..\..\gguf\$model ` - --poll 1000 -t 6 --cpu-mask 0xfc --cpu-strict 1 ` - --ubatch-size 1024 -ngl 99 --device $device $cli_opts diff --git a/scripts/snapdragon/windows/run-cli.ps1 b/scripts/snapdragon/windows/run-cli.ps1 deleted file mode 100644 index 5da8bff33e34..000000000000 --- a/scripts/snapdragon/windows/run-cli.ps1 +++ /dev/null @@ -1,53 +0,0 @@ - -#!/usr/bin/env pwsh - -# Basedir on device -$basedir=".\pkg-snapdragon" - -$cli_opts=$args - -$model="Llama-3.2-3B-Instruct-Q4_0.gguf" -if ($null -ne $env:M) { - $model=$env:M -} - -$device="HTP0" -if ($null -ne $env:D) { - $device=$env:D -} - -if ($null -ne $env:V) { - $env:GGML_HEXAGON_VERBOSE=$env:V -} - -if ($null -ne $env:SCHED) { - $env:GGML_SCHED_DEBUG=$env:SCHED; $cli_opts="$cli_opts -v" -} - -if ($null -ne $env:PROF) { - $env:GGML_HEXAGON_PROFILE=$env:PROF -} - -if ($null -ne $env:OPSTAGE) { - $env:GGML_HEXAGON_OPSTAGE=$env:OPSTAGE -} - -if ($null -ne $env:NHVX) { - $env:GGML_HEXAGON_NHVX=$env:NHVX -} - -if ($null -ne $env:NDEV) { - $env:GGML_HEXAGON_NDEV=$env:NDEV -} - -if ($null -ne $env:HB) { - $env:GGML_HEXAGON_HOSTBUF=$env:HB -} - -$env:ADSP_LIBRARY_PATH="$basedir\lib" - -& "$basedir\bin\llama-cli.exe" ` - --load-mode none -m $basedir\..\..\gguf\$model ` - --poll 1000 -t 6 --cpu-mask 0xfc --cpu-strict 1 ` - --ctx-size 8192 --ubatch-size 1024 -fa on ` - -ngl 99 --device $device $cli_opts diff --git a/scripts/snapdragon/windows/run-completion.ps1 b/scripts/snapdragon/windows/run-completion.ps1 deleted file mode 100644 index 08ef139b7e2a..000000000000 --- a/scripts/snapdragon/windows/run-completion.ps1 +++ /dev/null @@ -1,53 +0,0 @@ - -#!/usr/bin/env pwsh - -# Basedir on device -$basedir=".\pkg-snapdragon" - -$cli_opts=$args - -$model="Llama-3.2-3B-Instruct-Q4_0.gguf" -if ($null -ne $env:M) { - $model=$env:M -} - -$device="HTP0" -if ($null -ne $env:D) { - $device=$env:D -} - -if ($null -ne $env:V) { - $env:GGML_HEXAGON_VERBOSE=$env:V -} - -if ($null -ne $env:SCHED) { - $env:GGML_SCHED_DEBUG=$env:SCHED; $cli_opts="$cli_opts -v" -} - -if ($null -ne $env:PROF) { - $env:GGML_HEXAGON_PROFILE=$env:PROF -} - -if ($null -ne $env:OPSTAGE) { - $env:GGML_HEXAGON_OPSTAGE=$env:OPSTAGE -} - -if ($null -ne $env:NHVX) { - $env:GGML_HEXAGON_NHVX=$env:NHVX -} - -if ($null -ne $env:NDEV) { - $env:GGML_HEXAGON_NDEV=$env:NDEV -} - -if ($null -ne $env:HB) { - $env:GGML_HEXAGON_HOSTBUF=$env:HB -} - -$env:ADSP_LIBRARY_PATH="$basedir\lib" - -& "$basedir\bin\llama-completion.exe" ` - --load-mode none -m $basedir\..\..\gguf\$model ` - --poll 1000 -t 6 --cpu-mask 0xfc --cpu-strict 1 ` - --ctx-size 8192 --ubatch-size 1024 -fa on ` - -ngl 99 -no-cnv --device $device $cli_opts diff --git a/scripts/snapdragon/windows/run-mtmd.ps1 b/scripts/snapdragon/windows/run-mtmd.ps1 deleted file mode 100644 index 6e270ec90b57..000000000000 --- a/scripts/snapdragon/windows/run-mtmd.ps1 +++ /dev/null @@ -1,68 +0,0 @@ -#!/usr/bin/env pwsh - -# Basedir on device -$basedir=".\pkg-snapdragon" - -$cli_opts=$args - -$model="gemma-3-4b-it-Q4_0.gguf" -if ($null -ne $env:M) { - $model=$env:M -} - -$mmproj="mmproj-F16.gguf" -if ($null -ne $env:MMPROJ) { - $mmproj=$env:MMPROJ -} - -$image="" -if ($null -ne $env:IMG) { - $image=$env:IMG -} - -$device="HTP0" -if ($null -ne $env:D) { - $device=$env:D -} - -if ($null -ne $env:V) { - $env:GGML_HEXAGON_VERBOSE=$env:V -} - -if ($null -ne $env:SCHED) { - $env:GGML_SCHED_DEBUG=$env:SCHED; $cli_opts="$cli_opts -v" -} - -if ($null -ne $env:PROF) { - $env:GGML_HEXAGON_PROFILE=$env:PROF -} - -if ($null -ne $env:OPSTAGE) { - $env:GGML_HEXAGON_OPSTAGE=$env:OPSTAGE -} - -if ($null -ne $env:NHVX) { - $env:GGML_HEXAGON_NHVX=$env:NHVX -} - -if ($null -ne $env:NDEV) { - $env:GGML_HEXAGON_NDEV=$env:NDEV -} - -if ($null -ne $env:HB) { - $env:GGML_HEXAGON_HOSTBUF=$env:HB -} - -if ($null -ne $env:MTMD_DEVICE) { - $env:MTMD_BACKEND_DEVICE=$env:MTMD_DEVICE -} - -$env:ADSP_LIBRARY_PATH="$basedir\lib" - -& "$basedir\bin\llama-mtmd-cli.exe" ` - --load-mode none -m $basedir\..\..\gguf\$model ` - --mmproj $basedir\..\..\gguf\$mmproj ` - --image $basedir\..\..\gguf\$image ` - --poll 1000 -t 6 --cpu-mask 0xfc --cpu-strict 1 ` - --ctx-size 8192 --ubatch-size 1024 -fa on ` - -ngl 99 --device $device -v $cli_opts diff --git a/scripts/snapdragon/windows/run-tool.ps1 b/scripts/snapdragon/windows/run-tool.ps1 deleted file mode 100644 index 15c880f2dbdf..000000000000 --- a/scripts/snapdragon/windows/run-tool.ps1 +++ /dev/null @@ -1,56 +0,0 @@ - -#!/usr/bin/env pwsh - -# Basedir on device -$basedir=".\pkg-snapdragon" - -if ($args.Count -eq 0) { - Write-Host "No arguments provided.Expected the tool and argument to run." - exit -1 -} - -$tool=$args[0] -$cli_opts=@() - -if ($args.Count -gt 1) { - $cli_opts=$args[1..($args.Count - 1)] - $remainingArgs = $args[1..($args.Count - 1)] -} - -$device="HTP0" -if ($null -ne $env:D) { - $device=$env:D -} - -if ($null -ne $env:V) { - $env:GGML_HEXAGON_VERBOSE=$env:V -} - -if ($null -ne $env:SCHED) { - $env:GGML_SCHED_DEBUG=$env:SCHED; $cli_opts="$cli_opts -v" -} - -if ($null -ne $env:PROF) { - $env:GGML_HEXAGON_PROFILE=$env:PROF -} - -if ($null -ne $env:OPSTAGE) { - $env:GGML_HEXAGON_OPSTAGE=$env:OPSTAGE -} - -if ($null -ne $env:NHVX) { - $env:GGML_HEXAGON_NHVX=$env:NHVX -} - -if ($null -ne $env:NDEV) { - $env:GGML_HEXAGON_NDEV=$env:NDEV -} - -if ($null -ne $env:HB) { - $env:GGML_HEXAGON_HOSTBUF=$env:HB -} - -$env:ADSP_LIBRARY_PATH="$basedir\lib" - -& "$basedir\bin\$tool" ` - $cli_opts diff --git a/scripts/snapdragon/windows/setup-build.ps1 b/scripts/snapdragon/windows/setup-build.ps1 deleted file mode 100644 index d8ef24d44132..000000000000 --- a/scripts/snapdragon/windows/setup-build.ps1 +++ /dev/null @@ -1,105 +0,0 @@ -# Requires Run as Administrator is NOT strictly necessary for User-scope env vars, -# but recommended for creating directories in C:\ root if permissions are restricted. - -$ErrorActionPreference = "Stop" - -# --- Configuration --- -$BaseDir = "C:\Qualcomm" - -# SDK 1: Hexagon -$HexagonUrl = "https://github.com/snapdragon-toolchain/hexagon-sdk/releases/download/v6.6.0.0/hexagon-sdk-v6.6.0.0-arm64-wos.tar.xz" -$HexagonParent = Join-Path $BaseDir "Hexagon_SDK" -$HexagonSdkVersion = "6.6.0.0" -$HexagonToolsVersion = "19.0.07" -$HexagonSdkTarget = Join-Path $HexagonParent $HexagonSdkVersion -$HexagonToolsTarget = Join-Path $HexagonSdkTarget "\tools\HEXAGON_Tools\$HexagonToolsVersion" - -# SDK 2: OpenCL -$OpenCLUrl = "https://github.com/snapdragon-toolchain/opencl-sdk/releases/download/v2.3.2/adreno-opencl-sdk-v2.3.2-arm64-wos.tar.xz" -$OpenCLParent = Join-Path $BaseDir "OpenCL_SDK" -$OpenCLVersion = "2.3.2" -$OpenCLTarget = Join-Path $OpenCLParent $OpenCLVersion - -# --- Helper Function --- -function Install-QualcommSDK { - param ( - [string]$Url, - [string]$ParentDir, - [string]$TargetDir, - [string]$Name - ) - - # 1. Create Parent Directory - if (-not (Test-Path -Path $ParentDir)) { - Write-Host "Creating directory: $ParentDir" -ForegroundColor Cyan - New-Item -Path $ParentDir -ItemType Directory -Force | Out-Null - } - - # 2. Check for Specific Version Directory - if (Test-Path -Path $TargetDir) { - Write-Host "$Name ($TargetDir) already exists. Skipping download." -ForegroundColor Green - } - else { - Write-Host "$Name not found. preparing to download..." -ForegroundColor Yellow - - # Create the target directory to extract into - New-Item -Path $TargetDir -ItemType Directory -Force | Out-Null - - # Define temporary archive path - $TempFile = Join-Path $ParentDir "temp_sdk.tar.xz" - - try { - # Download - Write-Host "Downloading from: $Url" - Invoke-WebRequest -Uri $Url -OutFile $TempFile - - # Untar - # Note: We assume Windows includes tar.exe (Win 10 build 17063+) - Write-Host "Extracting archive to $TargetDir..." - - # We use -C to extract contents INTO the target directory created above - tar -xJvf $TempFile -C $TargetDir\.. - - Write-Host "Extraction complete." -ForegroundColor Green - } - catch { - Write-Error "Failed to download or extract $Name. Error: $_" - # Cleanup target dir if failed so script tries again next time - Remove-Item -Path $TargetDir -Recurse -Force -ErrorAction SilentlyContinue - } - finally { - # Cleanup Archive - if (Test-Path $TempFile) { Remove-Item $TempFile -Force } - } - } -} - -# --- Execution --- - -# 1. Ensure Base C:\Qualcomm exists -if (-not (Test-Path $BaseDir)) { - New-Item -Path $BaseDir -ItemType Directory -Force | Out-Null -} - -# 2. Run Install Logic -Install-QualcommSDK -Url $HexagonUrl -ParentDir $HexagonParent -TargetDir $HexagonSdkTarget -Name "Hexagon SDK" -Install-QualcommSDK -Url $OpenCLUrl -ParentDir $OpenCLParent -TargetDir $OpenCLTarget -Name "OpenCL SDK" - -# --- Environment Variables --- - -Write-Host "`nSetting Environment Variables..." -ForegroundColor Cyan - -# Set OPENCL_SDK_ROOT -[System.Environment]::SetEnvironmentVariable('OPENCL_SDK_ROOT', $OpenCLTarget, [System.EnvironmentVariableTarget]::User) -$env:OPENCL_SDK_ROOT = $OpenCLTarget # Set for current session as well -Write-Host "OPENCL_SDK_ROOT set to: $OpenCLTarget" - -# Set HEXAGON_SDK_ROOT -[System.Environment]::SetEnvironmentVariable('HEXAGON_SDK_ROOT', $HexagonSdkTarget, [System.EnvironmentVariableTarget]::User) -$env:HEXAGON_SDK_ROOT = $HexagonSdkTarget # Set for current session as well -Write-Host "HEXAGON_SDK_ROOT set to: $HexagonSdkTarget" - -# Set HEXAGON_SDK_ROOT -[System.Environment]::SetEnvironmentVariable('HEXAGON_TOOLS_ROOT', $HexagonToolsTarget, [System.EnvironmentVariableTarget]::User) -$env:HEXAGON_TOOLS_ROOT = $HexagonToolsTarget # Set for current session as well -Write-Host "HEXAGON_TOOLS_ROOT set to: $HexagonToolsTarget" diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index c6df19f2ecf4..8922dc12adc7 100644 --- a/src/CMakeLists.txt +++ b/src/CMakeLists.txt @@ -31,6 +31,7 @@ add_library(llama llama-memory.cpp llama-memory-hybrid.cpp llama-memory-hybrid-iswa.cpp + llama-memory-hybrid-idx.cpp llama-memory-recurrent.cpp llama-mmap.cpp llama-model-loader.cpp diff --git a/src/llama-arch.cpp b/src/llama-arch.cpp index 2125a84cb776..8039ff30bf61 100644 --- a/src/llama-arch.cpp +++ b/src/llama-arch.cpp @@ -40,6 +40,7 @@ static const std::map LLM_ARCH_NAMES = { { LLM_ARCH_QWEN3VLMOE, "qwen3vlmoe" }, { LLM_ARCH_QWEN35, "qwen35" }, { LLM_ARCH_QWEN35MOE, "qwen35moe" }, + { LLM_ARCH_QWEN4EXP, "qwen4exp" }, { LLM_ARCH_PHI2, "phi2" }, { LLM_ARCH_PHI3, "phi3" }, { LLM_ARCH_PHIMOE, "phimoe" }, @@ -294,6 +295,17 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_HYPER_CONNECTION_COUNT, "%s.hyper_connection.count" }, { LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, "%s.hyper_connection.sinkhorn_iterations" }, { LLM_KV_HYPER_CONNECTION_EPSILON, "%s.hyper_connection.epsilon" }, + { LLM_KV_HYPER_CONNECTION_LOW_RANK, "%s.hyper_connection.low_rank" }, + + { LLM_KV_PLE_LAYERS, "%s.ple.layers" }, + { LLM_KV_PLE_NGRAM_SIZE, "%s.ple.ngram_size" }, + { LLM_KV_PLE_HEADS_PER_NGRAM, "%s.ple.heads_per_ngram" }, + { LLM_KV_PLE_CONV_KERNEL, "%s.ple.conv_kernel" }, + { LLM_KV_PLE_LAYER_MULTIPLIERS, "%s.ple.layer_multipliers" }, + { LLM_KV_PLE_HEAD_OFFSETS, "%s.ple.head_offsets" }, + { LLM_KV_PLE_HEAD_VOCAB_SIZES, "%s.ple.head_vocab_sizes" }, + { LLM_KV_PLE_EOS_TOKEN_ID, "%s.ple.eos_token_id" }, + { LLM_KV_PLE_IMAGE_TOKEN_ID, "%s.ple.image_token_id" }, { LLM_KV_HASH_LAYER_COUNT, "%s.hash_layer_count" }, @@ -345,6 +357,12 @@ static const std::map LLM_KV_NAMES = { { LLM_KV_NORM_BEFORE_RESIDUAL, "%s.norm_before_residual" }, { LLM_KV_NORM_BEFORE_FC, "%s.norm_before_fc" }, + { LLM_KV_DFLASH_BLOCK_SIZE, "%s.block_size" }, + { LLM_KV_DFLASH_CONV_KERNEL_SIZE, "%s.conv_kernel_size" }, + { LLM_KV_DFLASH_CONV_GROUP_SIZE, "%s.conv_group_size" }, + { LLM_KV_DFLASH_SELECTOR_RANK, "%s.selector_rank" }, + { LLM_KV_DFLASH_SELECTOR_TOP_K, "%s.selector_top_k" }, + { LLM_KV_SHORTCONV_L_CACHE, "%s.shortconv.l_cache" }, // inkling (private arch) @@ -513,12 +531,29 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_HC_HEAD_FN, "output_hc_fn" }, { LLM_TENSOR_HC_HEAD_BASE, "output_hc_base" }, { LLM_TENSOR_HC_HEAD_SCALE, "output_hc_scale" }, + { LLM_TENSOR_HC_HEAD_NORM, "output_hc_norm" }, + { LLM_TENSOR_HC_HEAD_DOWN, "output_hc_down" }, + { LLM_TENSOR_HC_HEAD_UP, "output_hc_up" }, { LLM_TENSOR_HC_ATTN_FN, "blk.%d.hc_attn_fn" }, { LLM_TENSOR_HC_ATTN_BASE, "blk.%d.hc_attn_base" }, { LLM_TENSOR_HC_ATTN_SCALE, "blk.%d.hc_attn_scale" }, { LLM_TENSOR_HC_FFN_FN, "blk.%d.hc_ffn_fn" }, { LLM_TENSOR_HC_FFN_BASE, "blk.%d.hc_ffn_base" }, { LLM_TENSOR_HC_FFN_SCALE, "blk.%d.hc_ffn_scale" }, + { LLM_TENSOR_HC_ATTN_NORM, "blk.%d.hc_attn_norm" }, + { LLM_TENSOR_HC_ATTN_DOWN, "blk.%d.hc_attn_down" }, + { LLM_TENSOR_HC_ATTN_UP, "blk.%d.hc_attn_up" }, + { LLM_TENSOR_HC_ATTN_INJECT, "blk.%d.hc_attn_inject" }, + { LLM_TENSOR_HC_FFN_NORM, "blk.%d.hc_ffn_norm" }, + { LLM_TENSOR_HC_FFN_DOWN, "blk.%d.hc_ffn_down" }, + { LLM_TENSOR_HC_FFN_UP, "blk.%d.hc_ffn_up" }, + { LLM_TENSOR_HC_FFN_INJECT, "blk.%d.hc_ffn_inject" }, + { LLM_TENSOR_PLE_KEY, "blk.%d.ple_key" }, + { LLM_TENSOR_PLE_VALUE, "blk.%d.ple_value" }, + { LLM_TENSOR_PLE_NORM_KEY, "blk.%d.ple_norm_key" }, + { LLM_TENSOR_PLE_NORM_QUERY, "blk.%d.ple_norm_query" }, + { LLM_TENSOR_PLE_NORM_CONV, "blk.%d.ple_norm_conv" }, + { LLM_TENSOR_PLE_CONV1D, "blk.%d.ple_conv1d" }, { LLM_TENSOR_ATTN_COMPRESSOR_WKV, "blk.%d.attn_compressor_kv" }, { LLM_TENSOR_ATTN_COMPRESSOR_WGATE, "blk.%d.attn_compressor_gate" }, { LLM_TENSOR_ATTN_COMPRESSOR_APE, "blk.%d.attn_compressor_ape" }, @@ -674,6 +709,13 @@ static const std::map LLM_TENSOR_NAMES = { { LLM_TENSOR_DSPARK_MARKOV_W1, "markov_w1" }, { LLM_TENSOR_DSPARK_MARKOV_W2, "markov_w2" }, { LLM_TENSOR_DSPARK_CONF_PROJ, "conf_proj" }, + { LLM_TENSOR_DFLASH_ATTN_CONV_BASE, "blk.%d.attn_conv_base" }, + { LLM_TENSOR_DFLASH_ATTN_CONV_PROJ, "blk.%d.attn_conv_proj" }, + { LLM_TENSOR_DFLASH_FFN_CONV_BASE, "blk.%d.ffn_conv_base" }, + { LLM_TENSOR_DFLASH_FFN_CONV_PROJ, "blk.%d.ffn_conv_proj" }, + { LLM_TENSOR_DFLASH_SELECTOR_PREV, "selector_predecessor" }, + { LLM_TENSOR_DFLASH_SELECTOR_NEXT, "selector_successor" }, + { LLM_TENSOR_DFLASH_SELECTOR_HIDDEN, "selector_hidden" }, }; // declare information about the model weight tensors: @@ -727,12 +769,29 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_HC_HEAD_FN, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, {LLM_TENSOR_HC_HEAD_BASE, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_ADD}}, {LLM_TENSOR_HC_HEAD_SCALE, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL}}, + {LLM_TENSOR_HC_HEAD_NORM, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL}}, + {LLM_TENSOR_HC_HEAD_DOWN, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_HC_HEAD_UP, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, {LLM_TENSOR_HC_ATTN_FN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_HC_ATTN_BASE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_ADD}}, {LLM_TENSOR_HC_ATTN_SCALE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, {LLM_TENSOR_HC_FFN_FN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_HC_FFN_BASE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_ADD}}, {LLM_TENSOR_HC_FFN_SCALE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_HC_ATTN_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_HC_ATTN_DOWN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_HC_ATTN_UP, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_HC_ATTN_INJECT, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_HC_FFN_NORM, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_HC_FFN_DOWN, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_HC_FFN_UP, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_HC_FFN_INJECT, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_PLE_KEY, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_PLE_VALUE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_PLE_NORM_KEY, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_PLE_NORM_QUERY, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_PLE_NORM_CONV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_PLE_CONV1D, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_SSM_CONV}}, {LLM_TENSOR_ATTN_COMPRESSOR_WKV, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_ATTN_COMPRESSOR_WGATE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, {LLM_TENSOR_ATTN_COMPRESSOR_APE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_GET_ROWS}}, @@ -949,6 +1008,13 @@ static const std::map LLM_TENSOR_INFOS = { {LLM_TENSOR_DSPARK_MARKOV_W1, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}}, {LLM_TENSOR_DSPARK_MARKOV_W2, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, {LLM_TENSOR_DSPARK_CONF_PROJ, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_DFLASH_ATTN_CONV_BASE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_DFLASH_ATTN_CONV_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_DFLASH_FFN_CONV_BASE, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL}}, + {LLM_TENSOR_DFLASH_FFN_CONV_PROJ, {LLM_TENSOR_LAYER_REPEATING, GGML_OP_MUL_MAT}}, + {LLM_TENSOR_DFLASH_SELECTOR_PREV, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}}, + {LLM_TENSOR_DFLASH_SELECTOR_NEXT, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_GET_ROWS}}, + {LLM_TENSOR_DFLASH_SELECTOR_HIDDEN, {LLM_TENSOR_LAYER_OUTPUT, GGML_OP_MUL_MAT}}, }; LLM_KV::LLM_KV(llm_arch arch, const char * suffix) : arch(arch), suffix(suffix) {} @@ -1042,6 +1108,7 @@ bool llm_arch_is_hybrid(const llm_arch & arch) { case LLM_ARCH_KIMI_K3: case LLM_ARCH_QWEN35: case LLM_ARCH_QWEN35MOE: + case LLM_ARCH_QWEN4EXP: case LLM_ARCH_DEEPSEEK4: case LLM_ARCH_MINIMAX_01: case LLM_ARCH_INKLING: diff --git a/src/llama-arch.h b/src/llama-arch.h index 490b1e0a4883..d6be8d40c196 100644 --- a/src/llama-arch.h +++ b/src/llama-arch.h @@ -45,6 +45,7 @@ enum llm_arch { LLM_ARCH_QWEN3VLMOE, LLM_ARCH_QWEN35, LLM_ARCH_QWEN35MOE, + LLM_ARCH_QWEN4EXP, LLM_ARCH_PHI2, LLM_ARCH_PHI3, LLM_ARCH_PHIMOE, @@ -299,6 +300,17 @@ enum llm_kv { LLM_KV_HYPER_CONNECTION_COUNT, LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, LLM_KV_HYPER_CONNECTION_EPSILON, + LLM_KV_HYPER_CONNECTION_LOW_RANK, + + LLM_KV_PLE_LAYERS, + LLM_KV_PLE_NGRAM_SIZE, + LLM_KV_PLE_HEADS_PER_NGRAM, + LLM_KV_PLE_CONV_KERNEL, + LLM_KV_PLE_LAYER_MULTIPLIERS, + LLM_KV_PLE_HEAD_OFFSETS, + LLM_KV_PLE_HEAD_VOCAB_SIZES, + LLM_KV_PLE_EOS_TOKEN_ID, + LLM_KV_PLE_IMAGE_TOKEN_ID, LLM_KV_HASH_LAYER_COUNT, @@ -388,6 +400,11 @@ enum llm_kv { LLM_KV_TARGET_LAYERS, LLM_KV_TARGET_HIDDEN_SIZE, + LLM_KV_DFLASH_BLOCK_SIZE, + LLM_KV_DFLASH_CONV_KERNEL_SIZE, + LLM_KV_DFLASH_CONV_GROUP_SIZE, + LLM_KV_DFLASH_SELECTOR_RANK, + LLM_KV_DFLASH_SELECTOR_TOP_K, LLM_KV_NORM_BEFORE_RESIDUAL, LLM_KV_NORM_BEFORE_FC, @@ -581,12 +598,29 @@ enum llm_tensor { LLM_TENSOR_HC_HEAD_FN, LLM_TENSOR_HC_HEAD_BASE, LLM_TENSOR_HC_HEAD_SCALE, + LLM_TENSOR_HC_HEAD_NORM, // qwen4exp + LLM_TENSOR_HC_HEAD_DOWN, // qwen4exp + LLM_TENSOR_HC_HEAD_UP, // qwen4exp LLM_TENSOR_HC_ATTN_FN, LLM_TENSOR_HC_ATTN_BASE, LLM_TENSOR_HC_ATTN_SCALE, LLM_TENSOR_HC_FFN_FN, LLM_TENSOR_HC_FFN_BASE, LLM_TENSOR_HC_FFN_SCALE, + LLM_TENSOR_HC_ATTN_NORM, // qwen4exp + LLM_TENSOR_HC_ATTN_DOWN, // qwen4exp + LLM_TENSOR_HC_ATTN_UP, // qwen4exp + LLM_TENSOR_HC_ATTN_INJECT, // qwen4exp + LLM_TENSOR_HC_FFN_NORM, // qwen4exp + LLM_TENSOR_HC_FFN_DOWN, // qwen4exp + LLM_TENSOR_HC_FFN_UP, // qwen4exp + LLM_TENSOR_HC_FFN_INJECT, // qwen4exp + LLM_TENSOR_PLE_KEY, // qwen4exp + LLM_TENSOR_PLE_VALUE, // qwen4exp + LLM_TENSOR_PLE_NORM_KEY, // qwen4exp + LLM_TENSOR_PLE_NORM_QUERY, // qwen4exp + LLM_TENSOR_PLE_NORM_CONV, // qwen4exp + LLM_TENSOR_PLE_CONV1D, // qwen4exp LLM_TENSOR_ATTN_COMPRESSOR_WKV, LLM_TENSOR_ATTN_COMPRESSOR_WGATE, LLM_TENSOR_ATTN_COMPRESSOR_APE, @@ -683,6 +717,13 @@ enum llm_tensor { LLM_TENSOR_DSPARK_MARKOV_W1, LLM_TENSOR_DSPARK_MARKOV_W2, LLM_TENSOR_DSPARK_CONF_PROJ, + LLM_TENSOR_DFLASH_ATTN_CONV_BASE, + LLM_TENSOR_DFLASH_ATTN_CONV_PROJ, + LLM_TENSOR_DFLASH_FFN_CONV_BASE, + LLM_TENSOR_DFLASH_FFN_CONV_PROJ, + LLM_TENSOR_DFLASH_SELECTOR_PREV, + LLM_TENSOR_DFLASH_SELECTOR_NEXT, + LLM_TENSOR_DFLASH_SELECTOR_HIDDEN, }; diff --git a/src/llama-context.cpp b/src/llama-context.cpp index 30e9d5d7febf..6d0dabe4bb58 100644 --- a/src/llama-context.cpp +++ b/src/llama-context.cpp @@ -2301,12 +2301,17 @@ uint32_t llama_context::graph_max_nodes(uint32_t n_tokens) const { model.arch == LLM_ARCH_BAILINGMOE3 || model.arch == LLM_ARCH_QWEN35 || model.arch == LLM_ARCH_QWEN35MOE || + model.arch == LLM_ARCH_QWEN4EXP || model.arch == LLM_ARCH_DEEPSEEK4 || (model.arch == LLM_ARCH_DFLASH && model.hparams.dsv4_hc_mult > 0) || model.arch == LLM_ARCH_NANBEIGE || model.arch == LLM_ARCH_MINIMAX_01 || model.arch == LLM_ARCH_MINIMAX_M3) { res = std::max(n_tokens * 40, 32u * model.n_tensors()); + } else if (model.arch == LLM_ARCH_DFLASH && model.hparams.dflash_selector_rank > 0) { + // DFlash2's convolutions and selector are shape work rather than matmuls, + // so they cost ~8.6 nodes per tensor against ~5.9 for a plain DFlash draft + res = std::max(1024u, 12u*model.n_tensors()); } else { res = std::max(1024u, 8u*model.n_tensors()); for (const auto & lora : model.loras) { diff --git a/src/llama-ext.h b/src/llama-ext.h index 35d6e58adfa8..92a759b7a0ae 100644 --- a/src/llama-ext.h +++ b/src/llama-ext.h @@ -120,6 +120,8 @@ LLAMA_API llama_context * llama_get_ctx_other(struct llama_context * ctx); // model/context data extraction // +LLAMA_API int32_t llama_model_dflash_selector_top_k(const struct llama_model * model); + // returns pointer to the target-model layer indices LLAMA_API const int32_t * llama_model_target_layer_ids (const struct llama_model * model); // returns the number of extracted layers from target model diff --git a/src/llama-hparams.cpp b/src/llama-hparams.cpp index 293619dbb4b1..9fbce9bdfa8a 100644 --- a/src/llama-hparams.cpp +++ b/src/llama-hparams.cpp @@ -206,7 +206,11 @@ uint32_t llama_hparams::n_embd_r() const { // TODO: maybe support other convolution strides than 1 // NOTE: since the first column of the conv_state is shifted out each time, it's not actually needed // Corresponds to Mamba's conv_states size - return (ssm_d_conv > 0 ? ssm_d_conv - 1 : 0) * (ssm_d_inner + 2*ssm_n_group*ssm_d_state); + const uint32_t n_conv = (ssm_d_conv > 0 ? ssm_d_conv - 1 : 0) * (ssm_d_inner + 2*ssm_n_group*ssm_d_state); + + // PLE conv history needs its own row: Meta splits cache_r_l by head, so a history packed behind the first is unaddressable + // it lives in cache_ple_r_l instead, mirrored like the rest of the PLE module + return n_conv; } uint32_t llama_hparams::n_embd_s() const { @@ -241,6 +245,23 @@ bool llama_hparams::is_recr(uint32_t il) const { GGML_ABORT("%s: il (%u) out of bounds (n_layer_all: %u)\n", __func__, il, n_layer_all); } +uint32_t llama_hparams::ple_conv_state() const { + if (ple_n_heads == 0 || ple_conv_kernel == 0) { + return 0; + } + + // dilation equals the n-gram size, matching the reference module + return (ple_conv_kernel - 1) * ple_ngram_size * dsv4_hc_mult * n_embd; +} + +bool llama_hparams::is_ple(uint32_t il) const { + if (il < n_layer_all) { + return is_ple_impl[il]; + } + + GGML_ABORT("%s: il (%u) out of bounds (n_layer_all: %u)\n", __func__, il, n_layer_all); +} + uint32_t llama_hparams::n_pos_per_embd() const { return rope_type == LLAMA_ROPE_TYPE_MROPE || rope_type == LLAMA_ROPE_TYPE_IMROPE ? 4 : 1; } diff --git a/src/llama-hparams.h b/src/llama-hparams.h index 6bd2721c7386..ebd038b4c2bf 100644 --- a/src/llama-hparams.h +++ b/src/llama-hparams.h @@ -3,12 +3,15 @@ #include "llama.h" #include +#include #include #include // bump if necessary #define LLAMA_MAX_LAYERS 512 #define LLAMA_MAX_EXPERTS 1024 // Kimi K3 +#define LLAMA_MAX_PLE_NGRAM 8 // qwen4exp +#define LLAMA_MAX_PLE_HEADS 64 // qwen4exp enum llama_expert_gating_func_type { LLAMA_EXPERT_GATING_FUNC_TYPE_NONE = 0, @@ -234,6 +237,12 @@ struct llama_hparams { // output embedding dimension (0 = use n_embd) uint32_t n_embd_out_impl = 0; + uint32_t dflash_block_size = 0; + uint32_t dflash_conv_kernel_size = 0; + uint32_t dflash_conv_group_size = 0; + uint32_t dflash_selector_rank = 0; + uint32_t dflash_selector_top_k = 0; + // llama4 smallthinker uint32_t n_moe_layer_step = 0; uint32_t n_no_rope_layer_step = 4; @@ -281,6 +290,30 @@ struct llama_hparams { float dsv4_hc_eps = 0.0f; std::array dsv4_compress_ratios; + // 0 = full rank (DeepSeek-V4) + uint32_t hc_low_rank = 0; + + uint32_t ple_ngram_size = 0; + uint32_t ple_heads_per_ngram = 0; + uint32_t ple_conv_kernel = 0; + uint32_t ple_n_heads = 0; // (ngram_size - 1) * heads_per_ngram + uint32_t ple_head_dim = 0; + uint32_t ple_eos_token_id = 0; + // the id the PLE hash stands in at image positions; 0 makes the loader fall back to EOS + uint32_t ple_image_token_id = 0; + // the file lists PLE layer indices, so this is never a per-layer gguf array and can hold one bit per layer + std::bitset is_ple_impl; + // the hash multipliers reach ~2e13 and have to stay 64-bit + std::array ple_layer_multipliers; + // head offsets and vocab sizes are token-space indices; the gather truncates them to int32 anyway + std::array ple_head_offsets; + std::array ple_head_vocab_sizes; + + bool is_ple(uint32_t il) const; + + // PLE conv history rows: (kernel - 1) * ngram_size; 0 without a PLE module + uint32_t ple_conv_state() const; + // qwen3vl deepstack // When parsed from GGUF, this implies the first N layers consume the first // N deepstack embeddings. Use deepstack_mapping_arr if you need a more diff --git a/src/llama-kv-cache.cpp b/src/llama-kv-cache.cpp index 2ddf16f7a6e6..a52018c1c293 100644 --- a/src/llama-kv-cache.cpp +++ b/src/llama-kv-cache.cpp @@ -6,6 +6,7 @@ #include "llama-context.h" #include +#include #include #include #include @@ -78,7 +79,8 @@ llama_kv_cache::llama_kv_cache( llama_memory_t mem_other, const layer_filter_cb & filter, const layer_reuse_cb & reuse, - const layer_share_cb & share) : + const layer_share_cb & share, + const char * name_tag) : model(model), hparams(hparams), v_trans(v_trans), n_seq_max(n_seq_max), n_stream(unified ? 1 : n_seq_max), n_pad(n_pad), n_swa(n_swa), swa_type(swa_type), other(static_cast(mem_other)), @@ -232,8 +234,8 @@ llama_kv_cache::llama_kv_cache( ggml_tensor * k = has_k ? ggml_new_tensor_3d(ctx, type_k, n_embd_k_gqa, kv_size, n_stream) : nullptr; ggml_tensor * v = has_v ? ggml_new_tensor_3d(ctx, type_v, n_embd_v_gqa, kv_size, n_stream) : nullptr; - has_k && ggml_format_name(k, "cache_k_l%d", il); - has_v && ggml_format_name(v, "cache_v_l%d", il); + has_k && ggml_format_name(k, "cache_%sk_l%d", name_tag, il); + has_v && ggml_format_name(v, "cache_%sv_l%d", name_tag, il); std::vector k_stream; std::vector v_stream; @@ -1129,7 +1131,7 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch & cells.pos_set(idx, ubatch.pos[i]); - if (ubatch.is_pos_2d() || ubatch.token) { + if (ubatch.is_pos_2d() || ubatch.token || hparams.ple_n_heads > 0) { llama_kv_cell_ext ext; if (ubatch.is_pos_2d()) { @@ -1139,6 +1141,12 @@ void llama_kv_cache::apply_ubatch(const slot_info & sinfo, const llama_ubatch & if (ubatch.token) { ext.tok = ubatch.token[i]; + } else if (hparams.ple_n_heads > 0) { + // embd batch (multimodal input) has no token ids, need to pad it with the correct ID for PLE layers + // TODO @ngxson : check if we can do the same as gemma 3n / gemma 4 + ext.tok = hparams.ple_image_token_id != 0 + ? (llama_token) hparams.ple_image_token_id + : (llama_token) hparams.ple_eos_token_id; } cells.ext_set(idx, ext); @@ -1892,7 +1900,8 @@ void llama_kv_cache::set_input_v_rot(ggml_tensor * dst) const { } bool llama_kv_cache::has_cell_ext() const { - return hparams.n_pos_per_embd() > 1; + // M-RoPE needs the 2D position, the PLE n-gram hash needs the token id + return hparams.n_pos_per_embd() > 1 || hparams.ple_n_heads > 0; } void llama_kv_cache::get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector & res) const { @@ -1921,6 +1930,8 @@ void llama_kv_cache::get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, st seqs.set(ubatch.seq_id_unq[s]); } + const llama_pos w0 = p_min - (llama_pos) n; + // (seq_id, pos) -> token, for every cell that could be a predecessor of a ubatch token std::unordered_map hist; @@ -1928,28 +1939,71 @@ void llama_kv_cache::get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, st return ((uint64_t) seq_id << 32) | (uint32_t) pos; }; + // handle M-RoPE gaps: multiple tokens share the same temporal pos + // TODO @ngxson : improve this in the future + std::array, LLAMA_MAX_SEQ> below; + below.fill({ -1, LLAMA_TOKEN_NULL }); + for (uint32_t s = 0; s < n_stream; ++s) { - v_cells[s].for_each_token_in(seqs, p_min - (llama_pos) n, p_max, + // p_max inclusive: an embd token looks up cells at its own (shared) position + v_cells[s].for_each_token_in(seqs, 0, p_max + 1, [&](llama_seq_id seq_id, llama_pos pos, llama_token tok) { - hist[key(seq_id, pos)] = tok; + if (pos >= w0) { + hist[key(seq_id, pos)] = tok; + } else if (pos > below[seq_id].first) { + below[seq_id] = { pos, tok }; + } }); } + // the token at pos p, or the nearest earlier one when p falls in an M-RoPE gap + const auto lookup = [&](llama_seq_id seq_id, llama_pos p) -> llama_token { + for (llama_pos q = p; q >= w0; --q) { + const auto it = hist.find(key(seq_id, q)); + if (it != hist.end()) { + return it->second; + } + } + return below[seq_id].second; + }; + + // an embd (multimodal) ubatch can repeat one position for a whole image, so positions + // do not encode the token order; resolve its predecessors by ubatch order instead + std::vector ord; // index among the ubatch tokens of the same seq + std::unordered_map> seq_idx; + + if (!ubatch.token) { + ord.resize(n_tokens); + for (uint32_t i = 0; i < n_tokens; ++i) { + auto & v = seq_idx[ubatch.seq_id[i][0]]; + ord[i] = v.size(); + v.push_back(i); + } + } + for (uint32_t i = 0; i < n_tokens; ++i) { // TODO: a token that belongs to more than one sequence has an ambiguous history. // the n-gram architectures have to reject such batches const llama_seq_id seq_id = ubatch.seq_id[i][0]; for (uint32_t j = 0; j < n; ++j) { - const llama_pos p = ubatch.pos[i] - (llama_pos) (n - j); + const llama_pos d = (llama_pos) (n - j); + + llama_pos p; + if (!ubatch.token) { + const auto & v = seq_idx[seq_id]; + const int64_t k = (int64_t) ord[i] - d; + // k >= 0: an earlier token of this very ubatch; k < 0: before the chunk + p = k >= 0 ? ubatch.pos[v[k]] : ubatch.pos[v[0]] + (llama_pos) k; + } else { + p = ubatch.pos[i] - d; + } + if (p < 0) { continue; } - const auto it = hist.find(key(seq_id, p)); - if (it != hist.end()) { - res[i*n + j] = it->second; - } + res[i*n + j] = lookup(seq_id, p); } } } @@ -2186,6 +2240,15 @@ void llama_kv_cache::state_write(llama_io_write_i & io, llama_seq_id seq_id, lla } void llama_kv_cache::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) { + state_read_sinfo(io, seq_id, flags, nullptr, nullptr); +} + +void llama_kv_cache::state_read_sinfo( + llama_io_read_i & io, + llama_seq_id seq_id, + llama_state_seq_flags flags, + slot_info_vec_t * sinfos_out, +const slot_info_vec_t * sinfos_in) { // TODO: refactor [TAG_KV_CACHE_SHARE_CELLS] if (other) { return; @@ -2196,17 +2259,35 @@ void llama_kv_cache::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama // TODO: fix incosistent handling of `seq_id < 0` and `seq_id == -1` in the codebase [TAG_LLAMA_SEQ_ID_NEG] GGML_ASSERT(seq_id == -1 || (seq_id >= 0 && (size_t) seq_id < seq_to_stream.size())); + if (sinfos_out) { + sinfos_out->assign(n_stream, slot_info{}); + } + + if (sinfos_in && sinfos_in->size() != n_stream) { + throw std::runtime_error("failed to restore kv cache: mirrored slot layout has the wrong stream count"); + } + uint32_t n_stream_cur; io.read(&n_stream_cur, sizeof(n_stream_cur)); if (n_stream_cur != n_stream) { throw std::runtime_error("n_stream mismatch"); } + // a whole-context restore replaces every stream, so the cache is emptied once here + // clear() resets all streams at once, so doing it per stream below would keep only the last one + if (seq_id == -1) { + clear(true); + } + for (uint32_t s = 0; s < n_stream; ++s) { uint32_t cell_count; io.read(&cell_count, sizeof(cell_count)); if (cell_count == 0) { + // a mirrored cache must be empty here as well, or the two no longer agree cell for cell + if (sinfos_in && !(*sinfos_in)[s].empty()) { + throw std::runtime_error("failed to restore kv cache: mirrored cache holds cells this one does not"); + } continue; } @@ -2215,7 +2296,7 @@ void llama_kv_cache::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama slot_info sinfo; bool res = true; - res = res && state_read_meta(io, strm, cell_count, sinfo, seq_id); + res = res && state_read_meta(io, strm, cell_count, sinfo, seq_id, sinfos_in ? &(*sinfos_in)[s] : nullptr); try { res = res && state_read_data(io, strm, cell_count, sinfo); @@ -2231,6 +2312,10 @@ void llama_kv_cache::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama } throw std::runtime_error("failed to restore kv cache"); } + + if (sinfos_out) { + (*sinfos_out)[s] = sinfo; + } } } @@ -2366,7 +2451,7 @@ void llama_kv_cache::state_write_data(llama_io_write_i & io, const cell_ranges_t } } -bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32_t cell_count, slot_info & sinfo, llama_seq_id dest_seq_id) { +bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32_t cell_count, slot_info & sinfo, llama_seq_id dest_seq_id, const slot_info * sinfo_in) { auto & cells = v_cells[strm]; auto & head = v_heads[strm]; @@ -2416,10 +2501,37 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32 ubatch.seq_id[i] = &dest_seq_id; } - sinfo = find_slot(ubatch, false); - if (sinfo.empty()) { - LLAMA_LOG_ERROR("%s: failed to find %d available cells in kv cache\n", __func__, cell_count); - return false; + if (sinfo_in) { + // this cache mirrors another one, so it takes that cache's layout instead of searching for its own cells + if (sinfo_in->empty() || sinfo_in->n_stream() != 1 || sinfo_in->idxs[0].size() != cell_count) { + LLAMA_LOG_ERROR("%s: mirrored slot layout holds %d cells, this cache restores %d\n", __func__, + sinfo_in->empty() ? 0 : (int) sinfo_in->idxs[0].size(), cell_count); + return false; + } + + sinfo = *sinfo_in; + + // the layout is cell indices, so it means the same in both caches only while their streams line up + sinfo.s0 = strm; + sinfo.s1 = strm; + sinfo.strm[0] = strm; + + // seq_rm above freed exactly the cells this sequence held + // anything else in the way is a cache that had already drifted, which this restore must not hide + for (uint32_t i = 0; i < cell_count; ++i) { + const uint32_t idx = sinfo.idxs[0][i]; + + if (idx >= cells.size() || !cells.is_empty(idx)) { + LLAMA_LOG_ERROR("%s: cell %u of the mirrored slot layout is not free\n", __func__, idx); + return false; + } + } + } else { + sinfo = find_slot(ubatch, false); + if (sinfo.empty()) { + LLAMA_LOG_ERROR("%s: failed to find %d available cells in kv cache\n", __func__, cell_count); + return false; + } } // note: apply_ubatch() rebuilds llama_kv_cell_ext from the ubatch @@ -2445,7 +2557,12 @@ bool llama_kv_cache::state_read_meta(llama_io_read_i & io, uint32_t strm, uint32 return false; } - clear(true); + // the cells go in from 0, so a mirrored cache lands on the same ones as long as it restores the same count. the layout itself carries no more information here + if (sinfo_in && (sinfo_in->empty() || sinfo_in->n_stream() != 1 || sinfo_in->idxs[0].size() != cell_count)) { + LLAMA_LOG_ERROR("%s: mirrored slot layout holds %d cells, this cache restores %d\n", __func__, + sinfo_in->empty() ? 0 : (int) sinfo_in->idxs[0].size(), cell_count); + return false; + } for (uint32_t i = 0; i < cell_count; ++i) { llama_pos pos; diff --git a/src/llama-kv-cache.h b/src/llama-kv-cache.h index 91cf7860d640..8c62300f8356 100644 --- a/src/llama-kv-cache.h +++ b/src/llama-kv-cache.h @@ -112,7 +112,9 @@ class llama_kv_cache : public llama_memory_i { llama_memory_t mem_other, const layer_filter_cb & filter, const layer_reuse_cb & reuse, - const layer_share_cb & share); + const layer_share_cb & share, + // a model can hold more than one cache, so the tensor names have to stay unique + const char * name_tag = ""); ~llama_kv_cache() = default; @@ -166,6 +168,17 @@ class llama_kv_cache : public llama_memory_i { const llama_kv_cells & get_cells(llama_seq_id seq_id) const; + // state_read, plus the cells the restored tokens were placed in + // a cache that mirrors another one (the qwen4exp indexer) must not search for its own cells: two searches agree only by luck + // sinfos_out: if set, filled with the layout used; a stream with no cells leaves an empty entry + // sinfos_in : if set, the layout to use instead of searching. one entry per stream, cell count must match the blob + void state_read_sinfo( + llama_io_read_i & io, + llama_seq_id seq_id, + llama_state_seq_flags flags, + slot_info_vec_t * sinfos_out, + const slot_info_vec_t * sinfos_in); + // // graph_build API // @@ -230,7 +243,10 @@ class llama_kv_cache : public llama_memory_i { bool has_cell_ext() const; // for every token of the ubatch, the ids of the n tokens that precede it in its sequence - // entries with no matching cell are set to LLAMA_TOKEN_NULL + // example for M-RoPE image case: tokens A B X X X C, where X is a 3-token image at pos 2 spanning positions 2..4: + // tok: A B X X X C + // pos: 0 1 2 2 2 5 + // prev, n=2: A -> [NULL, NULL], B -> [NULL, A], 3rd X -> [X, X], C -> [X, X] // note: used by n-gram input embeddings void get_prev_tokens(const llama_ubatch & ubatch, uint32_t n, std::vector & res) const; @@ -333,7 +349,8 @@ class llama_kv_cache : public llama_memory_i { void state_write_meta(llama_io_write_i & io, const cell_ranges_t & cr, llama_seq_id seq_id = -1) const; void state_write_data(llama_io_write_i & io, const cell_ranges_t & cr) const; - bool state_read_meta(llama_io_read_i & io, uint32_t strm, uint32_t cell_count, slot_info & sinfo, llama_seq_id dest_seq_id = -1); + // sinfo_in, when set, replaces the find_slot call: the cells are given by the caller + bool state_read_meta(llama_io_read_i & io, uint32_t strm, uint32_t cell_count, slot_info & sinfo, llama_seq_id dest_seq_id = -1, const slot_info * sinfo_in = nullptr); bool state_read_data(llama_io_read_i & io, uint32_t strm, uint32_t cell_count, const slot_info & sinfo); }; diff --git a/src/llama-memory-hybrid-idx.cpp b/src/llama-memory-hybrid-idx.cpp new file mode 100644 index 000000000000..d4e59d77e570 --- /dev/null +++ b/src/llama-memory-hybrid-idx.cpp @@ -0,0 +1,465 @@ +#include "llama-memory-hybrid-idx.h" + +#include "llama-impl.h" +#include "llama-batch.h" +#include "llama-io.h" +#include "llama-model.h" + +#include +#include +#include +#include +#include + +// +// llama_memory_hybrid_idx +// + +llama_memory_hybrid_idx::llama_memory_hybrid_idx( + const llama_model & model, + /* attn */ + ggml_type type_k, + ggml_type type_v, + bool v_trans, + uint32_t kv_size, + uint32_t n_pad, + uint32_t n_swa, + llama_swa_type swa_type, + /* recurrent */ + ggml_type type_r, + ggml_type type_s, + uint32_t rs_size, + /* common */ + uint32_t n_seq_max, + uint32_t n_rs_seq, + bool offload, + bool unified, + /* layer filters */ + const layer_filter_cb & filter_attn, + const layer_filter_cb & filter_recr, + const layer_filter_cb & filter_idx) : + llama_memory_hybrid( + model, + type_k, type_v, v_trans, kv_size, n_pad, n_swa, swa_type, + type_r, type_s, rs_size, + n_seq_max, n_rs_seq, offload, unified, + filter_attn, filter_recr), + hparams_idx(model.hparams), + mem_idx(filter_idx == nullptr ? nullptr : [&] { + // MQA with a single key head of indexer_head_size, as llama_kv_cache_dsa shapes its own + std::fill(hparams_idx.n_head_kv_arr.begin(), hparams_idx.n_head_kv_arr.end(), 1); + hparams_idx.n_embd_head_k_full = model.hparams.indexer_head_size; + + LLAMA_LOG_INFO("%s: creating indexer KV cache, size = %u cells\n", __func__, kv_size); + + return new llama_kv_cache( + model, hparams_idx, type_k, type_v, v_trans, offload, unified, + kv_size, n_seq_max, n_pad, n_swa, swa_type, + nullptr, filter_idx, nullptr, nullptr, "idx_"); + }()) {} + +llama_memory_context_ptr llama_memory_hybrid_idx::init_batch(llama_batch_allocr & balloc, uint32_t n_ubatch, bool embd_all) { + // note: repeats llama_memory_hybrid::init_batch, as the indexer needs the attention slot infos that the base context hides + do { + balloc.split_reset(); + + // follow the recurrent pattern for creating the ubatch splits + std::vector ubatches; + + while (true) { + llama_ubatch ubatch; + + if (embd_all) { + // if all tokens are output, split by sequence + ubatch = balloc.split_seq(n_ubatch); + } else { + // Use non-sequential split when KV cache is unified (needed for hellaswag/winogrande/multiple-choice) + const bool unified = (get_mem_attn()->get_n_stream() == 1); + + // [TAG_RECURRENT_ROLLBACK_SPLITS] + // the trailing (1 + n_rs_seq) tokens of each seq must stay in the same ubatch + // so that the rollback snapshots remain valid + const uint32_t n_rs_seq = get_mem_recr()->n_rs_seq; + + ubatch = balloc.split_equal(n_ubatch, !unified, n_rs_seq > 0 ? n_rs_seq + 1 : 0); + } + + if (ubatch.n_tokens == 0) { + break; + } + + ubatches.push_back(std::move(ubatch)); // NOLINT + } + + if (balloc.get_n_used() < balloc.get_n_tokens()) { + // failed to find a suitable split + break; + } + + // prepare the recurrent batches first + if (!get_mem_recr()->prepare(ubatches)) { + // TODO: will the recurrent cache be in an undefined context at this point? + LLAMA_LOG_ERROR("%s: failed to prepare recurrent ubatches\n", __func__); + return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); + } + + // prepare the attention cache + auto heads_attn = get_mem_attn()->prepare(ubatches); + if (heads_attn.empty()) { + LLAMA_LOG_ERROR("%s: failed to prepare attention ubatches\n", __func__); + return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); + } + + // the indexer uses the attention cache's slot layout; a separate one can drift from it + llama_kv_cache::slot_info_vec_t heads_idx; + if (mem_idx) { + heads_idx = heads_attn; + } + + return std::make_unique( + this, std::move(heads_attn), std::move(heads_idx), std::move(ubatches)); + } while(false); + + return std::make_unique(LLAMA_MEMORY_STATUS_FAILED_PREPARE); +} + +llama_memory_context_ptr llama_memory_hybrid_idx::init_full() { + return std::make_unique(this); +} + +llama_memory_context_ptr llama_memory_hybrid_idx::init_update(llama_context * lctx, bool optimize) { + return std::make_unique(this, lctx, optimize); +} + +void llama_memory_hybrid_idx::clear(bool data) { + llama_memory_hybrid::clear(data); + + if (mem_idx) { + mem_idx->clear(data); + } +} + +bool llama_memory_hybrid_idx::seq_rm(llama_seq_id seq_id, llama_pos p0, llama_pos p1) { + // same order as llama_memory_hybrid::seq_rm: the recurrent cache can refuse, so try it first + if (!get_mem_recr()->seq_rm(seq_id, p0, p1)) { + return false; + } + + if (mem_idx) { + mem_idx->seq_rm(seq_id, p0, p1); + } + + return get_mem_attn()->seq_rm(seq_id, p0, p1); +} + +void llama_memory_hybrid_idx::seq_cp(llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) { + llama_memory_hybrid::seq_cp(seq_id_src, seq_id_dst, p0, p1); + + if (mem_idx) { + mem_idx->seq_cp(seq_id_src, seq_id_dst, p0, p1); + } +} + +void llama_memory_hybrid_idx::seq_keep(llama_seq_id seq_id) { + llama_memory_hybrid::seq_keep(seq_id); + + if (mem_idx) { + mem_idx->seq_keep(seq_id); + } +} + +void llama_memory_hybrid_idx::seq_add(llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) { + llama_memory_hybrid::seq_add(seq_id, p0, p1, shift); + + if (mem_idx) { + mem_idx->seq_add(seq_id, p0, p1, shift); + } +} + +void llama_memory_hybrid_idx::seq_div(llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) { + llama_memory_hybrid::seq_div(seq_id, p0, p1, d); + + if (mem_idx) { + mem_idx->seq_div(seq_id, p0, p1, d); + } +} + +std::map llama_memory_hybrid_idx::memory_breakdown() const { + std::map mb = llama_memory_hybrid::memory_breakdown(); + + if (mem_idx) { + for (const auto & buft_size : mem_idx->memory_breakdown()) { + mb[buft_size.first] += buft_size.second; + } + } + + return mb; +} + +void llama_memory_hybrid_idx::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const { + llama_memory_hybrid::state_write(io, seq_id, flags); + + // [TAG_HYBRID_IDX_STATE] the indexer section goes last, so it is a pure suffix: an old reader stops early instead of misparsing it + // The indexer mirrors the attention cache, so it uses the same PARTIAL_ONLY gate. + if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) { + if (mem_idx) { + mem_idx->state_write(io, seq_id, flags); + } + } + +} + +void llama_memory_hybrid_idx::state_read(llama_io_read_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) { + // note: repeats llama_memory_hybrid::state_read + // the indexer needs the attention cache's cells, and a half-failed restore must leave all three caches alike + + // [TAG_HYBRID_IDX_SINFO] + // the indexer restore adopts the attention cache's layout instead of searching for cells of its own + // two find_slot calls agree only while both caches see the same occupancy, which a restore cannot promise + llama_kv_cache::slot_info_vec_t sinfos_attn; + + try { + if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) { + get_mem_attn()->state_read_sinfo(io, seq_id, flags, mem_idx ? &sinfos_attn : nullptr, nullptr); + } + + get_mem_recr()->state_read(io, seq_id, flags); + + // [TAG_HYBRID_IDX_STATE] must mirror the write order in state_write + if ((flags & LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY) == 0) { + if (mem_idx) { + mem_idx->state_read_sinfo(io, seq_id, flags, nullptr, &sinfos_attn); + } + } + + } catch (...) { + // a half-restored context is the one state the indexer cannot fix by itself: attention holds new cells, the indexer old ones + // drop what was being restored from all of them, which is a state they do agree on. + state_drop(seq_id); + + throw; + } +} + +void llama_memory_hybrid_idx::state_drop(llama_seq_id seq_id) { + // dropped directly, not via seq_rm: the recurrent cache may refuse it and then only the other two get cleared + if (seq_id < 0) { + clear(true); + + return; + } + + get_mem_attn()->seq_rm(seq_id, -1, -1); + get_mem_recr()->seq_rm(seq_id, -1, -1); + + if (mem_idx) { + mem_idx->seq_rm(seq_id, -1, -1); + } +} + +llama_kv_cache * llama_memory_hybrid_idx::get_mem_idx() const { + return mem_idx.get(); +} + +// +// llama_memory_hybrid_idx_context +// + +// streams in each ubatch's slot info, matching get_k/get_v's `ns` +static std::vector llama_memory_hybrid_idx_ns(const llama_kv_cache::slot_info_vec_t & sinfos) { + std::vector res; + res.reserve(sinfos.size()); + + for (const auto & sinfo : sinfos) { + res.push_back(sinfo.s1 - sinfo.s0 + 1); + } + + return res; +} + +llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context(llama_memory_status status) : + llama_memory_hybrid_context(status) {} + +llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context(llama_memory_hybrid_idx * mem) : + llama_memory_hybrid_context(mem), + mem(mem), + // graph reservation walks a full context, and qwen4exp builds the sparse attention only when this is set + // without it the reserved worst case is the dense graph, so ggml-alloc must grow the buffer on the first decode + ns_ubatch(mem->get_mem_idx() == nullptr ? + std::vector() : std::vector{ mem->get_mem_idx()->get_n_stream() }), + ctx_idx(mem->get_mem_idx() == nullptr ? nullptr : + new llama_kv_cache_context(mem->get_mem_idx())) {} + +llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context( + llama_memory_hybrid_idx * mem, + llama_context * lctx, + bool optimize) : + llama_memory_hybrid_context(mem, lctx, optimize), + mem(mem) {} + +llama_memory_hybrid_idx_context::llama_memory_hybrid_idx_context( + llama_memory_hybrid_idx * mem, + slot_info_vec_t sinfos_attn, + slot_info_vec_t sinfos_idx, + std::vector ubatches) : + // note: the base copies the ubatches; ctx_idx gets a copy of its own + llama_memory_hybrid_context(mem, std::move(sinfos_attn), ubatches), + mem(mem), + ns_ubatch(llama_memory_hybrid_idx_ns(sinfos_idx)), + ctx_idx(mem->get_mem_idx() == nullptr ? nullptr : + new llama_kv_cache_context(mem->get_mem_idx(), std::move(sinfos_idx), ubatches)) {} + +bool llama_memory_hybrid_idx_context::next() { + if (ctx_idx) { + ctx_idx->next(); + } + + ++i_cur; + + return llama_memory_hybrid_context::next(); +} + +bool llama_memory_hybrid_idx_context::apply() { + bool res = llama_memory_hybrid_context::apply(); + + if (ctx_idx) { + res = res & ctx_idx->apply(); + } + + return res; +} + +const llama_kv_cache_context * llama_memory_hybrid_idx_context::get_idx() const { + return static_cast(ctx_idx.get()); +} + +uint32_t llama_memory_hybrid_idx_context::get_n_stream() const { + GGML_ASSERT(i_cur < ns_ubatch.size()); + + return ns_ubatch[i_cur]; +} + +void llama_memory_hybrid_idx_context::set_input_qsa( + ggml_tensor * cell_blk, + ggml_tensor * blk_cells, + ggml_tensor * blk_pos, + ggml_tensor * bias, + const llama_ubatch * ubatch, + uint32_t ratio, + bool blk_bias) const { + GGML_ASSERT(ratio > 0); + GGML_ASSERT(mem != nullptr && mem->get_mem_idx() != nullptr); + + GGML_ASSERT(ggml_backend_buffer_is_host(cell_blk->buffer)); + + const int64_t n_kv = cell_blk->ne[0]; + const int64_t n_ns = cell_blk->ne[1]; // streams in this ubatch + const int64_t n_blocks = blk_pos->ne[0]/(4*n_ns); + const int64_t n_tokens = ubatch->n_tokens; + const int64_t r = ratio; + + GGML_ASSERT(n_tokens % n_ns == 0); + const int64_t n_tps = n_tokens/n_ns; // tokens per stream + + int32_t * dst_cell_blk = (int32_t *) cell_blk->data; + int32_t * dst_blk_cells = (int32_t *) blk_cells->data; + int32_t * dst_blk_pos = (int32_t *) blk_pos->data; + float * dst_bias = (float *) bias->data; + + // block b covers [b*ratio, (b+1)*ratio), so its first token is at b*ratio + // all mrope sections carry it: exact for text, approximate for images + for (int64_t sec = 0; sec < 4; ++sec) { + for (int64_t s = 0; s < n_ns; ++s) { + for (int64_t b = 0; b < n_blocks; ++b) { + dst_blk_pos[sec*(n_blocks*n_ns) + s*n_blocks + b] = (int32_t) (b*r); + } + } + } + + // one pass per stream: cell j is a different token in each, so no mapping is shared + std::vector blk_of(n_kv); + std::vector filled(n_blocks); + + for (int64_t s = 0; s < n_ns; ++s) { + // ubatch index s*n_tps belongs to this stream; ask which cells array it uses + const llama_seq_id seq_of_stream = ubatch->seq_id[s*n_tps][0]; + const auto & cells = mem->get_mem_idx()->get_cells(seq_of_stream); + + int32_t * cur_cell_blk = dst_cell_blk + s*n_kv; + int32_t * cur_blk_cells = dst_blk_cells + s*(r*n_blocks); + + // an incomplete block cannot be pooled; the bias below forces those tail cells in + // -1 means no usable block, and block 0 only keeps the gather in range + std::fill(blk_of.begin(), blk_of.end(), -1); + std::fill(filled.begin(), filled.end(), 0); + std::fill(cur_blk_cells, cur_blk_cells + r*n_blocks, 0); + + // a cell no block covers needs its own -inf, which a per-block bias cannot carry + // every cache path keeps the position below the cell window, so this stays false + bool oor = false; + + for (int64_t j = 0; j < n_kv; ++j) { + if (cells.is_empty(j)) { + continue; + } + + const llama_pos p = cells.pos_get(j); + const int64_t b = p/r; + + if (b >= n_blocks) { + oor = true; + continue; + } + + blk_of[j] = (int32_t) b; + cur_blk_cells[b*r + (p%r)] = (int32_t) j; + filled[b]++; + } + + GGML_ASSERT((!blk_bias || !oor) && "qsa: cell position runs past the cell window"); + + // per-block mode keeps an unpooled cell's real block, so the block's own -inf reaches it + // per-cell mode carries that -inf itself and only needs the gather in range + for (int64_t j = 0; j < n_kv; ++j) { + if (blk_of[j] >= 0 && filled[blk_of[j]] < r && !blk_bias) { + blk_of[j] = -1; + } + cur_cell_blk[j] = blk_of[j] < 0 ? 0 : blk_of[j]; + } + + for (int64_t ii = 0; ii < n_tps; ++ii) { + const int64_t i = s*n_tps + ii; + const llama_seq_id seq_id = ubatch->seq_id[i][0]; + const llama_pos q = ubatch->pos[i]; + + // the tail is an incomplete block and is always visible, as in the reference + const llama_pos tail_start = (q + 1)/r*r; + + if (blk_bias) { + // a block sits wholly inside or outside the tail, so one value covers it + // the caller adds the attention mask, which drops empty, foreign and future cells + float * cur_blk_bias = dst_bias + i*n_blocks; + + for (int64_t b = 0; b < n_blocks; ++b) { + // finite, so it can never meet a -inf and produce a nan + cur_blk_bias[b] = b*r >= tail_start ? 1e9f : (filled[b] < r ? -INFINITY : 0.0f); + } + + continue; + } + + float * cur_bias = dst_bias + i*n_kv; + + for (int64_t j = 0; j < n_kv; ++j) { + float v = -INFINITY; + + if (!cells.is_empty(j) && cells.seq_has(j, seq_id) && cells.pos_get(j) <= q) { + // finite, so it can never meet a -inf and produce a nan + v = cells.pos_get(j) >= tail_start ? 1e9f : (blk_of[j] < 0 ? -INFINITY : 0.0f); + } + + cur_bias[j] = v; + } + } + } +} diff --git a/src/llama-memory-hybrid-idx.h b/src/llama-memory-hybrid-idx.h new file mode 100644 index 000000000000..e3472646d0f6 --- /dev/null +++ b/src/llama-memory-hybrid-idx.h @@ -0,0 +1,156 @@ +#pragma once + +#include "llama-memory-hybrid.h" + +#include +#include + +// +// llama_memory_hybrid_idx +// + +// llama_memory_hybrid plus a third cache with one indexer key per token, for block-sparse attention (qwen4exp QSA) +// the indexer is a side buffer over the attention cells: same size, padding, streams and slots, so cell j is one token in both + +class llama_memory_hybrid_idx : public llama_memory_hybrid { +public: + llama_memory_hybrid_idx( + const llama_model & model, + /* attn */ + ggml_type type_k, + ggml_type type_v, + bool v_trans, + uint32_t kv_size, + uint32_t n_pad, + uint32_t n_swa, + llama_swa_type swa_type, + /* recurrent */ + ggml_type type_r, + ggml_type type_s, + uint32_t rs_size, + /* common */ + uint32_t n_seq_max, + uint32_t n_rs_seq, + bool offload, + bool unified, + /* layer filters */ + const layer_filter_cb & filter_attn, + const layer_filter_cb & filter_recr, + /* the indexer cache exists only if this is given */ + const layer_filter_cb & filter_idx); + + ~llama_memory_hybrid_idx() = default; + + // + // llama_memory_i + // + + llama_memory_context_ptr init_batch( + llama_batch_allocr & balloc, + uint32_t n_ubatch, + bool embd_all) override; + + llama_memory_context_ptr init_full() override; + + llama_memory_context_ptr init_update(llama_context * lctx, bool optimize) override; + + void clear(bool data) override; + + bool seq_rm (llama_seq_id seq_id, llama_pos p0, llama_pos p1) override; + void seq_cp (llama_seq_id seq_id_src, llama_seq_id seq_id_dst, llama_pos p0, llama_pos p1) override; + void seq_keep(llama_seq_id seq_id) override; + void seq_add (llama_seq_id seq_id, llama_pos p0, llama_pos p1, llama_pos shift) override; + void seq_div (llama_seq_id seq_id, llama_pos p0, llama_pos p1, int d) override; + + std::map memory_breakdown() const override; + + // state write/load + + void state_write(llama_io_write_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) const override; + void state_read (llama_io_read_i & io, llama_seq_id seq_id = -1, llama_state_seq_flags flags = 0) override; + + // + // llama_memory_hybrid_idx specific API + // + + llama_kv_cache * get_mem_idx() const; // nullptr when the model carries no indexer + +private: + // forget seq_id (all of it if seq_id < 0) in every cache at once, so a failed restore cannot leave the caches out of step + // seq_id < 0 drops the whole context, as the caches themselves do on a failed restore + void state_drop(llama_seq_id seq_id); + + // the indexer cache holds one key head per layer, so it needs its own hparams: + // llama_kv_cache keeps a reference to what it is given + llama_hparams hparams_idx; + + const std::unique_ptr mem_idx; +}; + +class llama_memory_hybrid_idx_context : public llama_memory_hybrid_context { +public: + using slot_info_vec_t = llama_kv_cache::slot_info_vec_t; + + // used for errors + explicit llama_memory_hybrid_idx_context(llama_memory_status status); + + // used to create a full-cache context + explicit llama_memory_hybrid_idx_context(llama_memory_hybrid_idx * mem); + + // used to create an update context + llama_memory_hybrid_idx_context( + llama_memory_hybrid_idx * mem, + llama_context * lctx, + bool optimize); + + // used to create a batch processing context from a batch + llama_memory_hybrid_idx_context( + llama_memory_hybrid_idx * mem, + slot_info_vec_t sinfos_attn, + slot_info_vec_t sinfos_idx, + std::vector ubatches); + + ~llama_memory_hybrid_idx_context() = default; + + // + // llama_memory_context_i + // + + bool next() override; + bool apply() override; + + // + // llama_memory_hybrid_idx_context specific API + // + + // nullptr with no indexer, and for the update context, which builds no sparse graph + const llama_kv_cache_context * get_idx() const; + + // streams in the current slot info, the `ns` of get_k/get_v; 1 if unified + uint32_t get_n_stream() const; + + // block-compressed sparse attention (qwen4exp QSA) over the cells of the indexer cache. + // Blocks cut the position line, not the cell array, so no caller assumes a contiguous layout: + // cell_blk I32 [n_kv, ns] block each cell belongs to + // blk_cells I32 [ratio*n_blocks, ns] cells making up each block + // blk_pos I32 [4*n_blocks*ns] mrope position rows of each block's first token + // bias F32 [n_kv, n_tokens/ns, ns] -inf where invisible, large where always visible + // blk_bias asks for the bias per block instead: [n_blocks, n_tokens/ns, ns] + // the caller then adds the attention mask, the only part of the bias that varies within a block + void set_input_qsa(ggml_tensor * cell_blk, ggml_tensor * blk_cells, ggml_tensor * blk_pos, + ggml_tensor * bias, const llama_ubatch * ubatch, uint32_t ratio, + bool blk_bias) const; + +private: + const llama_memory_hybrid_idx * mem = nullptr; + + // streams per ubatch, read from the slot infos before ctx_idx takes them + // declared first, so it is initialised while sinfos_idx is still intact + const std::vector ns_ubatch; + + // null unless the model has an indexer and this is a batch or full context + const llama_memory_context_ptr ctx_idx; + + // mirrors the base class's ubatch cursor, which is private there + size_t i_cur = 0; +}; diff --git a/src/llama-memory-recurrent.cpp b/src/llama-memory-recurrent.cpp index e2990972ef7e..57919accf095 100644 --- a/src/llama-memory-recurrent.cpp +++ b/src/llama-memory-recurrent.cpp @@ -51,7 +51,8 @@ llama_memory_recurrent::llama_memory_recurrent( auto it = ctx_map.find(buft); if (it == ctx_map.end()) { ggml_init_params params = { - /*.mem_size =*/ size_t(2u*n_layer*ggml_tensor_overhead()), + // r and s per layer, plus the separate PLE conv row where the model has one + /*.mem_size =*/ size_t((hparams.ple_conv_state() > 0 ? 3u : 2u)*n_layer*ggml_tensor_overhead()), /*.mem_buffer =*/ NULL, /*.no_alloc =*/ true, }; @@ -71,6 +72,7 @@ llama_memory_recurrent::llama_memory_recurrent( r_l.resize(n_layer); s_l.resize(n_layer); + p_l.resize(n_layer); for (int i = 0; i < n_layer; i++) { if (filter && !filter(i)) { @@ -103,6 +105,13 @@ llama_memory_recurrent::llama_memory_recurrent( ggml_format_name(s, "cache_s_l%d", i); r_l[i] = r; s_l[i] = s; + + // the PLE history needs its own row: Meta must mirror it while the delta-net conv state next door stays split + if (hparams.ple_conv_state() > 0 && hparams.is_ple(i)) { + ggml_tensor * p = ggml_new_tensor_2d(ctx, type_r, hparams.ple_conv_state(), n_rows); + ggml_format_name(p, "cache_ple_r_l%d", i); + p_l[i] = p; + } } // allocate tensors and initialize the buffers to avoid NaNs in the padding @@ -119,11 +128,13 @@ llama_memory_recurrent::llama_memory_recurrent( { const size_t memory_size_r = size_r_bytes(); const size_t memory_size_s = size_s_bytes(); + const size_t memory_size_p = size_p_bytes(); - LLAMA_LOG_INFO("%s: size = %7.2f MiB (%6u cells, %3d layers, %2u seqs %2u rs_seq), R (%s): %7.2f MiB, S (%s): %7.2f MiB\n", __func__, - (float)(memory_size_r + memory_size_s) / (1024.0f * 1024.0f), mem_size, n_layer, n_seq_max, n_rs_seq, + LLAMA_LOG_INFO("%s: size = %7.2f MiB (%6u cells, %3d layers, %2u seqs %2u rs_seq), R (%s): %7.2f MiB, S (%s): %7.2f MiB, P (%s): %7.2f MiB\n", __func__, + (float)(memory_size_r + memory_size_s + memory_size_p) / (1024.0f * 1024.0f), mem_size, n_layer, n_seq_max, n_rs_seq, ggml_type_name(type_r), (float)memory_size_r / (1024.0f * 1024.0f), - ggml_type_name(type_s), (float)memory_size_s / (1024.0f * 1024.0f)); + ggml_type_name(type_s), (float)memory_size_s / (1024.0f * 1024.0f), + ggml_type_name(type_r), (float)memory_size_p / (1024.0f * 1024.0f)); } } @@ -740,6 +751,18 @@ size_t llama_memory_recurrent::size_s_bytes() const { return size_s_bytes; } +size_t llama_memory_recurrent::size_p_bytes() const { + size_t size_p_bytes = 0; + + for (const auto & p : p_l) { + if (p != nullptr) { + size_p_bytes += ggml_nbytes(p); + } + } + + return size_p_bytes; +} + void llama_memory_recurrent::state_write(llama_io_write_i & io, llama_seq_id seq_id, llama_state_seq_flags flags) const { GGML_UNUSED(flags); @@ -899,6 +922,17 @@ void llama_memory_recurrent::state_write_data(llama_io_write_i & io, const std:: const size_t buf_size = range_size * r_size_row; io.write_tensor(r_l[il], range.first * r_size_row, buf_size); } + + // the PLE conv history is a second recurrent row, so it has to travel with the first + if (p_l[il] != nullptr) { + const uint64_t p_size_row = ggml_row_size(p_l[il]->type, hparams.ple_conv_state()); + io.write(&p_size_row, sizeof(p_size_row)); + + for (const auto & range : cell_ranges) { + const size_t range_size = range.second - range.first; + io.write_tensor(p_l[il], range.first * p_size_row, range_size * p_size_row); + } + } } if (!s_trans) { @@ -1097,6 +1131,20 @@ bool llama_memory_recurrent::state_read_data(llama_io_read_i & io, uint32_t cell // Read and set the keys for the whole cell range io.read_tensor(r_l[il], head * r_size_row, cell_count * r_size_row); } + + if (p_l[il] != nullptr) { + uint64_t p_size_row_ref; + io.read(&p_size_row_ref, sizeof(p_size_row_ref)); + const size_t p_size_row = ggml_row_size(p_l[il]->type, hparams.ple_conv_state()); + if (p_size_row != p_size_row_ref) { + LLAMA_LOG_ERROR("%s: mismatched ple row size (%zu != %zu, layer %d)\n", __func__, p_size_row, (size_t) p_size_row_ref, il); + return false; + } + + if (cell_count) { + io.read_tensor(p_l[il], head * p_size_row, cell_count * p_size_row); + } + } } if (!s_trans) { @@ -1251,6 +1299,10 @@ ggml_tensor * llama_memory_recurrent_context::get_s_l(int32_t il) const { return mem->s_l[il]; } +ggml_tensor * llama_memory_recurrent_context::get_p_l(int32_t il) const { + return mem->p_l[il]; +} + int32_t llama_memory_recurrent_context::s_copy(int i) const { const uint32_t cell_idx = i + mem->head; const int32_t src0 = mem->cells[cell_idx].src0; diff --git a/src/llama-memory-recurrent.h b/src/llama-memory-recurrent.h index b13b7b748f5e..4abb3f5cf5c0 100644 --- a/src/llama-memory-recurrent.h +++ b/src/llama-memory-recurrent.h @@ -111,6 +111,8 @@ class llama_memory_recurrent : public llama_memory_i { // per layer std::vector r_l; std::vector s_l; + // a second conv history that must stay replicated across devices, so it cannot share the r row + std::vector p_l; private: //const llama_model & model; @@ -125,6 +127,7 @@ class llama_memory_recurrent : public llama_memory_i { size_t size_r_bytes() const; size_t size_s_bytes() const; + size_t size_p_bytes() const; void state_write_meta(llama_io_write_i & io, const std::vector> & cell_ranges, llama_seq_id seq_id = -1) const; void state_write_data(llama_io_write_i & io, const std::vector> & cell_ranges) const; @@ -170,6 +173,7 @@ class llama_memory_recurrent_context : public llama_memory_context_i { ggml_tensor * get_r_l(int32_t il) const; ggml_tensor * get_s_l(int32_t il) const; + ggml_tensor * get_p_l(int32_t il) const; int32_t s_copy(int i) const; diff --git a/src/llama-mmap.cpp b/src/llama-mmap.cpp index ed572da7fb54..4d183cbc9c45 100644 --- a/src/llama-mmap.cpp +++ b/src/llama-mmap.cpp @@ -438,11 +438,34 @@ void llama_file::write_u32(uint32_t val) const { pimpl->write_u32(val); } // llama_mmap +#if defined(_POSIX_MAPPED_FILES) || defined(_WIN32) +// merge `ranges` and return their complement within [0, limit) +static llama_mmap::ranges ranges_complement(llama_mmap::ranges ranges, size_t limit) { + llama_mmap::ranges res; + std::sort(ranges.begin(), ranges.end()); + + size_t pos = 0; + for (const auto & range : ranges) { + const size_t beg = std::min(range.first, limit); + const size_t end = std::min(range.second, limit); + if (beg > pos) { + res.emplace_back(pos, beg); + } + pos = std::max(pos, end); + } + if (pos < limit) { + res.emplace_back(pos, limit); + } + + return res; +} +#endif + struct llama_mmap::impl { #ifdef _POSIX_MAPPED_FILES std::vector> mapped_fragments; - impl(struct llama_file * file, size_t prefetch, bool numa) { + impl(struct llama_file * file, size_t prefetch, bool numa, const llama_mmap::ranges & lazy_ranges) { size = file->size(); int fd = file->file_id(); int flags = MAP_SHARED; @@ -452,19 +475,35 @@ struct llama_mmap::impl { LLAMA_LOG_WARN("warning: posix_fadvise(.., POSIX_FADV_SEQUENTIAL) failed: %s\n", strerror(errno)); } - if (prefetch) { flags |= MAP_POPULATE; } + // MAP_POPULATE would fault in the lazy ranges too + if (prefetch && lazy_ranges.empty()) { flags |= MAP_POPULATE; } #endif addr = mmap(NULL, file->size(), PROT_READ, flags, fd, 0); if (addr == MAP_FAILED) { throw std::runtime_error(format("mmap failed: %s", strerror(errno))); } + // page-aligned madvise over [beg, end), clamped to the file + auto advise = [&](size_t beg, size_t end, int advice, const char * name) { + const size_t page_size = sysconf(_SC_PAGESIZE); + beg = beg & ~(page_size - 1); + end = std::min((end + page_size - 1) & ~(page_size - 1), file->size()); + if (beg >= end) { + return; + } + if (posix_madvise((char *) addr + beg, end - beg, advice)) { + LLAMA_LOG_WARN("warning: posix_madvise(.., %s) failed: %s\n", name, strerror(errno)); + } + }; + if (prefetch > 0) { - if (posix_madvise(addr, std::min(file->size(), prefetch), POSIX_MADV_WILLNEED)) { - LLAMA_LOG_WARN("warning: posix_madvise(.., POSIX_MADV_WILLNEED) failed: %s\n", - strerror(errno)); + for (const auto & range : ranges_complement(lazy_ranges, std::min(file->size(), prefetch))) { + advise(range.first, range.second, POSIX_MADV_WILLNEED, "POSIX_MADV_WILLNEED"); } } + for (const auto & range : lazy_ranges) { + advise(range.first, range.second, POSIX_MADV_RANDOM, "POSIX_MADV_RANDOM"); + } if (numa) { if (posix_madvise(addr, file->size(), POSIX_MADV_RANDOM)) { LLAMA_LOG_WARN("warning: posix_madvise(.., POSIX_MADV_RANDOM) failed: %s\n", @@ -533,7 +572,7 @@ struct llama_mmap::impl { #elif defined(_WIN32) HANDLE hMapping = nullptr; - impl(struct llama_file * file, size_t prefetch, bool numa) { + impl(struct llama_file * file, size_t prefetch, bool numa, const llama_mmap::ranges & lazy_ranges) { GGML_UNUSED(numa); size = file->size(); @@ -563,10 +602,15 @@ struct llama_mmap::impl { pPrefetchVirtualMemory = (decltype(pPrefetchVirtualMemory))(void *) GetProcAddress(hKernel32, "PrefetchVirtualMemory"); if (pPrefetchVirtualMemory) { - WIN32_MEMORY_RANGE_ENTRY range; - range.VirtualAddress = addr; - range.NumberOfBytes = (SIZE_T) std::min(size, prefetch); - if (!pPrefetchVirtualMemory(GetCurrentProcess(), 1, &range, 0)) { + std::vector entries; + for (const auto & range : ranges_complement(lazy_ranges, std::min(size, prefetch))) { + WIN32_MEMORY_RANGE_ENTRY entry; + entry.VirtualAddress = (char *) addr + range.first; + entry.NumberOfBytes = (SIZE_T) (range.second - range.first); + entries.push_back(entry); + } + if (!entries.empty() && + !pPrefetchVirtualMemory(GetCurrentProcess(), (ULONG_PTR) entries.size(), entries.data(), 0)) { LLAMA_LOG_WARN("warning: PrefetchVirtualMemory failed: %s\n", llama_format_win_err(GetLastError()).c_str()); } @@ -597,10 +641,11 @@ struct llama_mmap::impl { } } #else - impl(struct llama_file * file, size_t prefetch, bool numa) { + impl(struct llama_file * file, size_t prefetch, bool numa, const llama_mmap::ranges & lazy_ranges) { GGML_UNUSED(file); GGML_UNUSED(prefetch); GGML_UNUSED(numa); + GGML_UNUSED(lazy_ranges); throw std::runtime_error("mmap not supported"); } @@ -617,7 +662,8 @@ struct llama_mmap::impl { size_t size; }; -llama_mmap::llama_mmap(struct llama_file * file, size_t prefetch, bool numa) : pimpl(std::make_unique(file, prefetch, numa)) {} +llama_mmap::llama_mmap(struct llama_file * file, size_t prefetch, bool numa, + const ranges & lazy_ranges) : pimpl(std::make_unique(file, prefetch, numa, lazy_ranges)) {} llama_mmap::~llama_mmap() = default; size_t llama_mmap::size() const { return pimpl->size; } diff --git a/src/llama-mmap.h b/src/llama-mmap.h index b7d5c61e95ff..cc28c8a73fa5 100644 --- a/src/llama-mmap.h +++ b/src/llama-mmap.h @@ -2,6 +2,7 @@ #include #include +#include #include #include @@ -41,8 +42,12 @@ struct llama_file { }; struct llama_mmap { + // list of [first, last) byte ranges within a file + using ranges = std::vector>; + llama_mmap(const llama_mmap &) = delete; - llama_mmap(struct llama_file * file, size_t prefetch = (size_t) -1, bool numa = false); + llama_mmap(struct llama_file * file, size_t prefetch = (size_t) -1, bool numa = false, + const ranges & lazy_ranges = {}); ~llama_mmap(); size_t size() const; diff --git a/src/llama-model-loader.cpp b/src/llama-model-loader.cpp index 9b22cb05f29b..d9241022cf41 100644 --- a/src/llama-model-loader.cpp +++ b/src/llama-model-loader.cpp @@ -321,10 +321,11 @@ namespace GGUFMeta { case GGUF_TYPE_UINT32: case GGUF_TYPE_INT32: type_ok = (std::is_same::value) || (std::is_same::value); break; + case GGUF_TYPE_UINT64: type_ok = (std::is_same::value); break; case GGUF_TYPE_FLOAT32: type_ok = (std::is_same::value); break; case GGUF_TYPE_STRING: type_ok = (std::is_same::value); break; default: - throw std::runtime_error(format("%s is not a string/float32/uint32/int32 array", key.c_str())); + throw std::runtime_error(format("%s is not a string/float32/uint32/int32/uint64 array", key.c_str())); } if (!type_ok) { throw std::runtime_error(format("%s has wrong array element type %s", key.c_str(), gguf_type_name(arr_info.gt))); @@ -367,10 +368,11 @@ namespace GGUFMeta { case GGUF_TYPE_UINT32: case GGUF_TYPE_INT32: type_ok = (std::is_same::value) || (std::is_same::value); break; + case GGUF_TYPE_UINT64: type_ok = (std::is_same::value); break; case GGUF_TYPE_FLOAT32: type_ok = (std::is_same::value); break; case GGUF_TYPE_STRING: type_ok = (std::is_same::value); break; default: - throw std::runtime_error(format("%s is not a string/float32/uint32/int32 array", key.c_str())); + throw std::runtime_error(format("%s is not a string/float32/uint32/int32/uint64 array", key.c_str())); } if (!type_ok) { throw std::runtime_error(format("%s has wrong array element type %s", key.c_str(), gguf_type_name(arr_info.gt))); @@ -410,6 +412,9 @@ namespace GGUFMeta { template bool llama_model_loader::get_arr>(enum llm_kv kid, std::array & result, bool required); template bool llama_model_loader::get_arr>(enum llm_kv kid, std::vector & result, bool required); template bool llama_model_loader::get_arr>(enum llm_kv kid, std::array & result, bool required); + template bool llama_model_loader::get_arr>(enum llm_kv kid, std::vector & result, bool required); + template bool llama_model_loader::get_arr>(enum llm_kv kid, std::array & result, bool required); + template bool llama_model_loader::get_arr>(enum llm_kv kid, std::array & result, bool required); template bool llama_model_loader::get_key(const std::string & key, T & result, bool required) { @@ -1282,6 +1287,18 @@ struct ggml_tensor * llama_model_loader::create_tensor( return NULL; } + if ((flags & TENSOR_READ_LAZY) && use_mmap && tensor_read_lazy != LLAMA_TENSOR_READ_LAZY_OFF) { + // in auto mode, small tensors are cheap enough to keep resident + constexpr size_t auto_lazy_min_size = 4ull * 1024 * 1024 * 1024; + if (tensor_read_lazy == LLAMA_TENSOR_READ_LAZY_ON || ggml_nbytes(cur) > auto_lazy_min_size) { + const auto & w = require_weight(tn.str().c_str()); + lazy_tensor_ranges[w.idx].emplace_back(w.offs, w.offs + ggml_nbytes(cur)); + + LLAMA_LOG_INFO("%s: tensor %s (size = %zu MiB) lazy read enabled\n", + __func__, tn.str().c_str(), ggml_nbytes(cur)/1024/1024); + } + } + ggml_tensor t_meta = *cur; if (flags & TENSOR_ALLOW_RESHAPE) { for (size_t dim = 0; dim < GGML_MAX_DIMS; dim++) { @@ -1349,7 +1366,9 @@ void llama_model_loader::init_mappings(bool prefetch, llama_mlocks * mlock_mmaps if (use_mmap) { mappings.reserve(files.size()); mmaps_used.reserve(files.size()); - for (const auto & file : files) { + for (uint32_t idx = 0; idx < files.size(); idx++) { + const auto & file = files[idx]; + bool is_numa = false; auto * dev = ggml_backend_dev_by_type(GGML_BACKEND_DEVICE_TYPE_CPU); @@ -1361,7 +1380,11 @@ void llama_model_loader::init_mappings(bool prefetch, llama_mlocks * mlock_mmaps } } - std::unique_ptr mapping = std::make_unique(file.get(), prefetch ? -1 : 0, is_numa); + const auto it_lazy = lazy_tensor_ranges.find(idx); + static const llama_mmap::ranges no_lazy_ranges; + + std::unique_ptr mapping = std::make_unique(file.get(), prefetch ? -1 : 0, is_numa, + it_lazy != lazy_tensor_ranges.end() ? it_lazy->second : no_lazy_ranges); mmaps_used.emplace_back(mapping->size(), 0); if (mlock_mmaps) { std::unique_ptr mlock_mmap(new llama_mlock()); @@ -1400,27 +1423,26 @@ void llama_model_loader::unmap_weight(const llama_tensor_weight & w) const { mappings.at(w.idx)->unmap_fragment(w.offs, w.offs + ggml_nbytes(w.tensor)); } -void llama_model_loader::load_data_for(struct ggml_tensor * cur) const { - const auto & w = require_weight(ggml_get_name(cur)); +const void * llama_model_loader::load_data_range(const llama_tensor_weight & w, size_t offs, size_t size, void * buf) const { + GGML_ASSERT(offs + size <= ggml_nbytes(w.tensor)); + + const void * data = buf; if (use_mmap) { - const auto & mapping = mappings.at(w.idx); - if (cur->data == nullptr) { - cur->data = (uint8_t *)mapping->addr() + w.offs; - } else { - memcpy(cur->data, (uint8_t *)mapping->addr() + w.offs, ggml_nbytes(cur)); - } + data = (const uint8_t *) mappings.at(w.idx)->addr() + w.offs + offs; } else { - GGML_ASSERT(cur->data != nullptr); + GGML_ASSERT(buf != nullptr); GGML_ASSERT(w.idx < files.size()); const auto & file = files.at(w.idx); - file->seek(w.offs, SEEK_SET); - file->read_raw(cur->data, ggml_nbytes(cur)); + file->seek(w.offs + offs, SEEK_SET); + file->read_raw(buf, size); } - if (check_tensors && !ggml_validate_row_data(cur->type, cur->data, ggml_nbytes(cur))) { - throw std::runtime_error(format("tensor '%s' has invalid data", ggml_get_name(cur))); + if (check_tensors && !ggml_validate_row_data(w.tensor->type, data, size)) { + throw std::runtime_error(format("tensor '%s' has invalid data", ggml_get_name(w.tensor))); } + + return data; } bool llama_model_loader::load_all_data( diff --git a/src/llama-model-loader.h b/src/llama-model-loader.h index e9fe3592d421..407260e9907e 100644 --- a/src/llama-model-loader.h +++ b/src/llama-model-loader.h @@ -68,6 +68,7 @@ struct llama_model_loader { static const int TENSOR_SKIP = 1 << 2; static const int TENSOR_SKIP_IF_VIRTUAL = 1 << 3; static const int TENSOR_ALLOW_RESHAPE = 1 << 4; + static const int TENSOR_READ_LAZY = 1 << 5; // read rows on demand instead of loading whole tensor; requires mmap for now int n_kv = 0; int n_tensors = 0; @@ -82,12 +83,18 @@ struct llama_model_loader { bool no_alloc; bool load_mtp; + // set by the caller before the create_tensor() calls + enum llama_tensor_read_lazy tensor_read_lazy = LLAMA_TENSOR_READ_LAZY_OFF; + llama_files files; llama_ftype ftype; llama_fver fver; llama_mmaps mappings; + // byte ranges of TENSOR_READ_LAZY tensors, per file index + std::map lazy_tensor_ranges; + std::map weights_map; std::unordered_map kv_overrides; const llama_model_tensor_buft_override * tensor_buft_overrides; @@ -197,8 +204,9 @@ struct llama_model_loader { // release a weight's mmap pages void unmap_weight(const llama_tensor_weight & w) const; - // for backwards compatibility, does not support ggml-backend - void load_data_for(struct ggml_tensor * cur) const; + // read a byte range of a weight's data + // with mmap, returns a pointer into the mapping, otherwise reads into buf and returns buf + const void * load_data_range(const llama_tensor_weight & w, size_t offs, size_t size, void * buf) const; // Returns false if cancelled by progress_callback bool load_all_data( diff --git a/src/llama-model-saver.cpp b/src/llama-model-saver.cpp index 4c91b6f4ce45..fe0efd30e653 100644 --- a/src/llama-model-saver.cpp +++ b/src/llama-model-saver.cpp @@ -61,6 +61,10 @@ void llama_model_saver::add_kv(const enum llm_kv key, const int32_t value) { gguf_set_val_i32(gguf_ctx, llm_kv(key).c_str(), value); } +void llama_model_saver::add_kv(const enum llm_kv key, const uint64_t value) { + gguf_set_val_u64(gguf_ctx, llm_kv(key).c_str(), value); +} + void llama_model_saver::add_kv(const enum llm_kv key, const float value) { gguf_set_val_f32(gguf_ctx, llm_kv(key).c_str(), value); } @@ -114,6 +118,8 @@ void llama_model_saver::add_kv(const enum llm_kv key, const Container & value, c gguf_set_arr_data(gguf_ctx, llm_kv(key).c_str(), GGUF_TYPE_BOOL, value.data(), n_values); } else if (std::is_same::value) { gguf_set_arr_data(gguf_ctx, llm_kv(key).c_str(), GGUF_TYPE_INT32, value.data(), n_values); + } else if (std::is_same::value) { + gguf_set_arr_data(gguf_ctx, llm_kv(key).c_str(), GGUF_TYPE_UINT64, value.data(), n_values); } else if (std::is_same::value) { gguf_set_arr_data(gguf_ctx, llm_kv(key).c_str(), GGUF_TYPE_FLOAT32, value.data(), n_values); } else if (std::is_same::value) { @@ -125,6 +131,7 @@ void llama_model_saver::add_kv(const enum llm_kv key, const Container & value, c // instantiate for external usage: template void llama_model_saver::add_kv>(const enum llm_kv, const std::vector &, const bool); template void llama_model_saver::add_kv>(const enum llm_kv, const std::vector &, const bool); +template void llama_model_saver::add_kv>(const enum llm_kv, const std::vector &, const bool); void llama_model_saver::add_kv(const enum llm_kv key, const std::vector & value) { std::vector tmp(value.size()); @@ -309,6 +316,32 @@ void llama_model_saver::add_kv_from_model() { add_kv(LLM_KV_HYPER_CONNECTION_SINKHORN_ITERATIONS, hparams.dsv4_hc_sinkhorn_iters); add_kv(LLM_KV_HYPER_CONNECTION_EPSILON, hparams.dsv4_hc_eps); add_kv(LLM_KV_HASH_LAYER_COUNT, hparams.dsv4_hash_layer_count); + add_kv(LLM_KV_HYPER_CONNECTION_LOW_RANK, hparams.hc_low_rank); + + // the PLE group only means anything whole: write all of it or none + if (hparams.ple_n_heads > 0) { + std::vector ple_layers; + for (uint32_t il = 0; il < hparams.n_layer_all; ++il) { + if (hparams.is_ple_impl[il]) { + ple_layers.push_back(il); + } + } + add_kv(LLM_KV_PLE_LAYERS, ple_layers); + add_kv(LLM_KV_PLE_NGRAM_SIZE, hparams.ple_ngram_size); + add_kv(LLM_KV_PLE_HEADS_PER_NGRAM, hparams.ple_heads_per_ngram); + add_kv(LLM_KV_PLE_CONV_KERNEL, hparams.ple_conv_kernel); + add_kv(LLM_KV_PLE_EOS_TOKEN_ID, hparams.ple_eos_token_id); + add_kv(LLM_KV_EMBEDDING_LENGTH_PER_LAYER, hparams.ple_head_dim); + add_kv(LLM_KV_PLE_LAYER_MULTIPLIERS, std::vector( + hparams.ple_layer_multipliers.begin(), + hparams.ple_layer_multipliers.begin() + hparams.ple_ngram_size)); + add_kv(LLM_KV_PLE_HEAD_OFFSETS, std::vector( + hparams.ple_head_offsets.begin(), + hparams.ple_head_offsets.begin() + hparams.ple_n_heads)); + add_kv(LLM_KV_PLE_HEAD_VOCAB_SIZES, std::vector( + hparams.ple_head_vocab_sizes.begin(), + hparams.ple_head_vocab_sizes.begin() + hparams.ple_n_heads)); + } const float rope_scaling_factor = hparams.rope_freq_scale_train == 1.0f ? 0.0f : 1.0f/hparams.rope_freq_scale_train; @@ -443,6 +476,10 @@ void llama_model_saver::add_tensors_from_model() { add_tensor(model->hc_head_fn); add_tensor(model->hc_head_base); add_tensor(model->hc_head_scale); + add_tensor(model->per_layer_tok_embd); + add_tensor(model->hc_head_norm); + add_tensor(model->hc_head_down); + add_tensor(model->hc_head_up); for (const struct llama_layer & layer : model->layers) { for (size_t i = 0; i < sizeof(layer)/sizeof(struct ggml_tensor *); ++i) { diff --git a/src/llama-model-saver.h b/src/llama-model-saver.h index 36a715e2b6bf..95e19e666e7f 100644 --- a/src/llama-model-saver.h +++ b/src/llama-model-saver.h @@ -21,6 +21,7 @@ struct llama_model_saver { void add_kv(enum llm_kv key, uint32_t value); void add_kv(enum llm_kv key, int32_t value); + void add_kv(enum llm_kv key, uint64_t value); void add_kv(enum llm_kv key, float value); void add_kv(enum llm_kv key, bool value); void add_kv(enum llm_kv key, const char * value); diff --git a/src/llama-model.cpp b/src/llama-model.cpp index d2e929f9bd9b..dc6d39adef7b 100644 --- a/src/llama-model.cpp +++ b/src/llama-model.cpp @@ -16,6 +16,7 @@ #include "llama-kv-cache-dsv4.h" #include "llama-memory-hybrid.h" #include "llama-memory-hybrid-iswa.h" +#include "llama-memory-hybrid-idx.h" #include "llama-memory-recurrent.h" #include "llama.h" @@ -319,6 +320,8 @@ static llama_model * llama_model_mapping(llm_arch arch, const llama_model_params return new llama_model_qwen35(params); case LLM_ARCH_QWEN35MOE: return new llama_model_qwen35moe(params); + case LLM_ARCH_QWEN4EXP: + return new llama_model_qwen4exp(params); case LLM_ARCH_MISTRAL3: return new llama_model_mistral3(params); case LLM_ARCH_EAGLE3: @@ -378,6 +381,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str static const std::regex pattern_qkv_bias ("blk\\.\\d*\\.attn_qkv.bias"); static const std::regex pattern_qk_norm ("blk\\.\\d*\\.attn_(q|k)_norm\\.weight"); static const std::regex pattern_kv_cache ("cache_(k|v)_l\\d*"); + static const std::regex pattern_idx_cache ("cache_idx_(k|v)_l\\d*"); static const std::regex pattern_dsv4_state ("dsv4_(csa|hca|lid)_state_(kv|score)_l\\d*"); static const std::regex pattern_attn_sinks ("blk\\.\\d*\\.attn_sinks.weight"); static const std::regex pattern_attn_out_weight ("blk\\.\\d*\\.attn_output.weight"); @@ -393,6 +397,7 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str static const std::regex pattern_ssm_beta ("blk\\.\\d*\\.ssm_beta.weight"); static const std::regex pattern_ssm_beta_alpha ("blk\\.\\d*\\.ssm_ba.weight"); static const std::regex pattern_r_cache ("cache_r_l\\d*"); + static const std::regex pattern_ple_r_cache ("cache_ple_r_l\\d*"); static const std::regex pattern_s_cache ("cache_s_l\\d*"); static const std::regex pattern_ssm_conv1d ("blk\\.\\d*\\.ssm_conv1d.weight"); static const std::regex pattern_ssm_out_weight ("blk\\.\\d*\\.ssm_out.weight"); @@ -490,6 +495,16 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str } } + // the qsa indexer has one key head and its projections are mirrored, so its cache cannot be split + if (std::regex_match(tensor_name, pattern_idx_cache)) { + return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED); + } + + // the PLE table is model-level and its conv is mirrored, so every device runs the whole conv and needs the whole history + if (std::regex_match(tensor_name, pattern_ple_r_cache)) { + return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_MIRRORED); + } + // standard attention if (std::regex_match(tensor_name, pattern_q_weight) || std::regex_match(tensor_name, pattern_kv_weight)) { return get_tensor_config_impl(GGML_BACKEND_SPLIT_AXIS_1, "attn_output.weight", "ssm_out.weight"); @@ -578,7 +593,8 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str }; auto get_split_segments = [&](int axis, uint32_t il) -> std::vector> { - if (ud->model->arch == LLM_ARCH_QWEN3NEXT || ud->model->arch == LLM_ARCH_QWEN35 || ud->model->arch == LLM_ARCH_QWEN35MOE) { + if (ud->model->arch == LLM_ARCH_QWEN3NEXT || ud->model->arch == LLM_ARCH_QWEN35 || ud->model->arch == LLM_ARCH_QWEN35MOE || + ud->model->arch == LLM_ARCH_QWEN4EXP) { const int64_t head_k_dim = hparams.ssm_d_state; const int64_t head_v_dim = hparams.ssm_d_state; const int64_t n_k_heads = hparams.ssm_n_group; @@ -716,7 +732,8 @@ struct ggml_backend_meta_split_state llama_meta_device_get_split_state(const str if (std::regex_match(tensor_name, pattern_q_weight) || std::regex_match(tensor_name, pattern_q_bias)) { GGML_ASSERT(segments.size() == 1); // some models have Q gate tensors, for those cases the granularity needs to be doubled: - if (ud->model->arch == LLM_ARCH_QWEN3NEXT || ud->model->arch == LLM_ARCH_QWEN35 || ud->model->arch == LLM_ARCH_QWEN35MOE) { + if (ud->model->arch == LLM_ARCH_QWEN3NEXT || ud->model->arch == LLM_ARCH_QWEN35 || ud->model->arch == LLM_ARCH_QWEN35MOE || + ud->model->arch == LLM_ARCH_QWEN4EXP) { return {std::lcm(2*n_embd_q, blck_size_perf)}; } return {granularity_q}; @@ -929,6 +946,7 @@ const char * llm_type_name(llm_type type) { case LLM_TYPE_35B_A3B: return "35B.A3B"; case LLM_TYPE_48B_A3B: return "48B.A3B"; case LLM_TYPE_80B_A3B: return "80B.A3B"; + case LLM_TYPE_A3B: return "A3B"; case LLM_TYPE_100B_A6B: return "100B.A6B"; case LLM_TYPE_102B_A12B: return "102B.A12B"; case LLM_TYPE_106B_A12B: return "106B.A12B"; @@ -2433,6 +2451,10 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, // layer filters, so pick the right one here llama_memory_hybrid::layer_filter_cb filter_attn = nullptr; llama_memory_hybrid::layer_filter_cb filter_recr = nullptr; + // only the sparse-attention architectures use llama_memory_hybrid_idx + // a null filter_idx means the GGUF has no indexer tensors + llama_memory_hybrid::layer_filter_cb filter_idx = nullptr; + const bool needs_mem_idx = (arch == LLM_ARCH_QWEN4EXP); if (arch == LLM_ARCH_FALCON_H1 || arch == LLM_ARCH_INKLING) { // all layers have both an attention KV cache and a recurrent (conv) state filter_attn = [&](uint32_t) { return true; }; @@ -2444,13 +2466,20 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, filter_recr = [&](uint32_t il) { return hparams.is_recr(il) && hparams.n_ff(il) == 0; }; - } else if (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || arch == LLM_ARCH_MINIMAX_01) { + } else if (arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || arch == LLM_ARCH_QWEN4EXP || arch == LLM_ARCH_MINIMAX_01) { filter_attn = [&](uint32_t il) { return il < hparams.n_layer() && !hparams.is_recr(il); }; filter_recr = [&](uint32_t il) { return il < hparams.n_layer() && hparams.is_recr(il); }; + + if (arch == LLM_ARCH_QWEN4EXP && hparams.indexer_head_size > 0) { + // QSA runs on the dense-attention layers only + filter_idx = [&](uint32_t il) { + return il < hparams.n_layer() && !hparams.is_recr(il); + }; + } } if (hparams.swa_type != LLAMA_SWA_TYPE_NONE) { @@ -2473,6 +2502,27 @@ llama_memory_i * llama_model::create_memory(const llama_memory_params & params, /* unified */ cparams.kv_unified, /* filter_attn */ std::move(filter_attn), /* filter_recr */ std::move(filter_recr)); + } else if (needs_mem_idx) { + // sparse attention over a per-token indexer cache, in its own memory type + res = new llama_memory_hybrid_idx( + /* model */ *this, + /* attn_type_k */ params.type_k, + /* attn_type_v */ params.type_v, + /* attn_v_trans */ !cparams.flash_attn, + /* attn_kv_size */ cparams.n_ctx_seq, + /* attn_n_pad */ 1, + /* attn_n_swa */ hparams.n_swa, + /* attn_swa_type */ hparams.swa_type, + /* recurrent_type_k */ GGML_TYPE_F32, + /* recurrent_type_v */ GGML_TYPE_F32, + /* recurrent_kv_size */ std::max((uint32_t) 1, cparams.n_seq_max), + /* n_seq_max */ cparams.n_seq_max, + /* n_rs_seq */ cparams.n_rs_seq, + /* offload */ cparams.offload_kqv, + /* unified */ cparams.kv_unified, + /* filter_attn */ std::move(filter_attn), + /* filter_recr */ std::move(filter_recr), + /* filter_idx */ std::move(filter_idx)); } else { res = new llama_memory_hybrid( /* model */ *this, @@ -2634,6 +2684,7 @@ llama_model_params llama_model_default_params() { /*.n_gpu_layers =*/ -1, /*.split_mode =*/ LLAMA_SPLIT_MODE_LAYER, /*.load_mode =*/ LLAMA_LOAD_MODE_AUTO, + /*.tensor_read_lazy =*/ LLAMA_TENSOR_READ_LAZY_AUTO, /*.main_gpu =*/ 0, /*.tensor_split =*/ nullptr, /*.progress_callback =*/ nullptr, @@ -2686,6 +2737,10 @@ int32_t llama_model_n_layer_nextn(const llama_model * model) { return model->hparams.n_layer_nextn; } +int32_t llama_model_dflash_selector_top_k(const llama_model * model) { + return model->hparams.dflash_selector_top_k; +} + int32_t llama_model_n_head(const llama_model * model) { return model->hparams.n_head(); } @@ -2885,6 +2940,10 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { return LLAMA_ROPE_TYPE_NEOX; case LLM_ARCH_DFLASH: + // drafts for M-RoPE targets carry rope sections and follow the target's temporal dim + if (const auto & s = model->hparams.rope_sections; s[0] || s[1] || s[2] || s[3]) { + return LLAMA_ROPE_TYPE_MROPE; + } // DSV4 DSpark drafters use DeepSeek-V4's normal RoPE; legacy DFlash backbones are NeoX return model->hparams.dsv4_hc_mult > 0 ? LLAMA_ROPE_TYPE_NORM : LLAMA_ROPE_TYPE_NEOX; @@ -2895,6 +2954,7 @@ llama_rope_type llama_model_rope_type(const llama_model * model) { case LLM_ARCH_QWEN3VLMOE: case LLM_ARCH_QWEN35: case LLM_ARCH_QWEN35MOE: + case LLM_ARCH_QWEN4EXP: case LLM_ARCH_QWEN3TTS: return LLAMA_ROPE_TYPE_IMROPE; @@ -3071,7 +3131,8 @@ llama_model_base::llama_model_base(const struct llama_model_params & params) : l TENSOR_NOT_REQUIRED (llama_model_loader::TENSOR_NOT_REQUIRED), TENSOR_SKIP (llama_model_loader::TENSOR_SKIP), TENSOR_SKIP_IF_VIRTUAL(llama_model_loader::TENSOR_SKIP_IF_VIRTUAL), - TENSOR_ALLOW_RESHAPE (llama_model_loader::TENSOR_ALLOW_RESHAPE) {} + TENSOR_ALLOW_RESHAPE (llama_model_loader::TENSOR_ALLOW_RESHAPE), + TENSOR_READ_LAZY (llama_model_loader::TENSOR_READ_LAZY) {} ggml_tensor * llama_model_base::create_tensor(const LLM_TN_IMPL & tn, const std::initializer_list & ne, int flags) { GGML_ASSERT(ml != nullptr); diff --git a/src/llama-model.h b/src/llama-model.h index 6cb62c0caebf..eaa1435fdacd 100644 --- a/src/llama-model.h +++ b/src/llama-model.h @@ -129,6 +129,7 @@ enum llm_type { LLM_TYPE_35B_A3B, // Qwen3.5 LLM_TYPE_48B_A3B, // Kimi Linear LLM_TYPE_80B_A3B, // Qwen3 Next + LLM_TYPE_A3B, // Qwen3.8 Flash Next LLM_TYPE_100B_A6B, LLM_TYPE_102B_A12B, // Solar-Open LLM_TYPE_106B_A12B, // GLM-4.5-Air @@ -363,6 +364,11 @@ struct llama_layer { struct ggml_tensor * ffn_exp_probs_b = nullptr; struct ggml_tensor * ffn_gate_tid2eid = nullptr; + struct ggml_tensor * dflash_attn_conv_base = nullptr; + struct ggml_tensor * dflash_attn_conv_proj = nullptr; + struct ggml_tensor * dflash_ffn_conv_base = nullptr; + struct ggml_tensor * dflash_ffn_conv_proj = nullptr; + // mamba proj struct ggml_tensor * ssm_in = nullptr; struct ggml_tensor * ssm_x = nullptr; @@ -555,6 +561,22 @@ struct llama_layer { struct ggml_tensor * index_q_norm = nullptr; struct ggml_tensor * index_k_norm = nullptr; + struct ggml_tensor * hc_attn_norm = nullptr; + struct ggml_tensor * hc_attn_down = nullptr; + struct ggml_tensor * hc_attn_up = nullptr; + struct ggml_tensor * hc_attn_inject = nullptr; + struct ggml_tensor * hc_ffn_norm = nullptr; + struct ggml_tensor * hc_ffn_down = nullptr; + struct ggml_tensor * hc_ffn_up = nullptr; + struct ggml_tensor * hc_ffn_inject = nullptr; + + struct ggml_tensor * ple_key = nullptr; + struct ggml_tensor * ple_value = nullptr; + struct ggml_tensor * ple_norm_key = nullptr; + struct ggml_tensor * ple_norm_query = nullptr; + struct ggml_tensor * ple_norm_conv = nullptr; + struct ggml_tensor * ple_conv1d = nullptr; + // gemma4 layer output scale, reused for talkie embedding skip scale struct ggml_tensor * out_scale = nullptr; @@ -644,6 +666,10 @@ struct llama_model { struct ggml_tensor * altup_proj = nullptr; struct ggml_tensor * altup_unembd_proj = nullptr; struct ggml_tensor * per_layer_tok_embd = nullptr; + + struct ggml_tensor * hc_head_norm = nullptr; + struct ggml_tensor * hc_head_down = nullptr; + struct ggml_tensor * hc_head_up = nullptr; struct ggml_tensor * per_layer_model_proj = nullptr; struct ggml_tensor * per_layer_proj_norm = nullptr; @@ -655,9 +681,14 @@ struct llama_model { // dspark struct ggml_tensor * dspark_markov_w1 = nullptr; struct ggml_tensor * dspark_markov_w2 = nullptr; + struct ggml_tensor * dspark_markov_w2_s = nullptr; struct ggml_tensor * dspark_conf_proj = nullptr; struct ggml_tensor * dspark_conf_proj_b = nullptr; + struct ggml_tensor * dflash_selector_prev = nullptr; + struct ggml_tensor * dflash_selector_next = nullptr; + struct ggml_tensor * dflash_selector_hidden = nullptr; + // unified vector to store target-model extracted layer ids in eagle3, dflash, etc. std::vector target_layer_ids; @@ -765,6 +796,7 @@ struct llama_model_base : public llama_model { const int TENSOR_SKIP; const int TENSOR_SKIP_IF_VIRTUAL; const int TENSOR_ALLOW_RESHAPE; + const int TENSOR_READ_LAZY; explicit llama_model_base(const llama_model_params & params); virtual ~llama_model_base() = default; diff --git a/src/llama-quant.cpp b/src/llama-quant.cpp index a37e0d04de15..416556dc0fd1 100644 --- a/src/llama-quant.cpp +++ b/src/llama-quant.cpp @@ -38,6 +38,9 @@ enum class tensor_category { OTHER }; +// max amount of tensor data kept in memory while quantizing a single tensor +static const size_t LLAMA_QUANT_MAX_BUF_SIZE = 8ull*1024*1024*1024; + static void zeros(std::ofstream & file, size_t n) { char zero = 0; for (size_t i = 0; i < n; ++i) { @@ -211,31 +214,26 @@ struct tensor_metadata { // static void llama_tensor_dequantize_impl( - ggml_tensor * tensor, std::vector> & output, std::vector & workers, + ggml_type type, const void * data, float * f32_output, std::vector & workers, const size_t nelements, const int nthread ) { - if (output.size() < nelements) { - output.resize(nelements); - } - float * f32_output = (float *) output.data(); - - const ggml_type_traits * qtype = ggml_get_type_traits(tensor->type); - if (ggml_is_quantized(tensor->type)) { + const ggml_type_traits * qtype = ggml_get_type_traits(type); + if (ggml_is_quantized(type)) { if (qtype->to_float == NULL) { - throw std::runtime_error(format("type %s unsupported for integer quantization: no dequantization available", ggml_type_name(tensor->type))); + throw std::runtime_error(format("type %s unsupported for integer quantization: no dequantization available", ggml_type_name(type))); } - } else if (tensor->type != GGML_TYPE_F16 && - tensor->type != GGML_TYPE_BF16) { - throw std::runtime_error(format("cannot dequantize/convert tensor type %s", ggml_type_name(tensor->type))); + } else if (type != GGML_TYPE_F16 && + type != GGML_TYPE_BF16) { + throw std::runtime_error(format("cannot dequantize/convert tensor type %s", ggml_type_name(type))); } if (nthread < 2) { - if (tensor->type == GGML_TYPE_F16) { - ggml_fp16_to_fp32_row((ggml_fp16_t *)tensor->data, f32_output, nelements); - } else if (tensor->type == GGML_TYPE_BF16) { - ggml_bf16_to_fp32_row((ggml_bf16_t *)tensor->data, f32_output, nelements); - } else if (ggml_is_quantized(tensor->type)) { - qtype->to_float(tensor->data, f32_output, nelements); + if (type == GGML_TYPE_F16) { + ggml_fp16_to_fp32_row((const ggml_fp16_t *)data, f32_output, nelements); + } else if (type == GGML_TYPE_BF16) { + ggml_bf16_to_fp32_row((const ggml_bf16_t *)data, f32_output, nelements); + } else if (ggml_is_quantized(type)) { + qtype->to_float(data, f32_output, nelements); } else { GGML_ABORT("fatal error"); // unreachable } @@ -243,14 +241,14 @@ static void llama_tensor_dequantize_impl( } size_t block_size; - if (tensor->type == GGML_TYPE_F16 || - tensor->type == GGML_TYPE_BF16) { + if (type == GGML_TYPE_F16 || + type == GGML_TYPE_BF16) { block_size = 1; } else { - block_size = (size_t)ggml_blck_size(tensor->type); + block_size = (size_t)ggml_blck_size(type); } - size_t block_size_bytes = ggml_type_size(tensor->type); + size_t block_size_bytes = ggml_type_size(type); GGML_ASSERT(nelements % block_size == 0); size_t nblocks = nelements / block_size; @@ -265,16 +263,16 @@ static void llama_tensor_dequantize_impl( size_t thr_elems = thr_blocks * block_size; // number of elements for this thread size_t thr_block_bytes = thr_blocks * block_size_bytes; // number of input bytes for this thread - auto compute = [qtype] (ggml_type typ, uint8_t * inbuf, float * outbuf, int nels) { + auto compute = [qtype] (ggml_type typ, const uint8_t * inbuf, float * outbuf, int nels) { if (typ == GGML_TYPE_F16) { - ggml_fp16_to_fp32_row((ggml_fp16_t *)inbuf, outbuf, nels); + ggml_fp16_to_fp32_row((const ggml_fp16_t *)inbuf, outbuf, nels); } else if (typ == GGML_TYPE_BF16) { - ggml_bf16_to_fp32_row((ggml_bf16_t *)inbuf, outbuf, nels); + ggml_bf16_to_fp32_row((const ggml_bf16_t *)inbuf, outbuf, nels); } else { qtype->to_float(inbuf, outbuf, nels); } }; - workers.emplace_back(compute, tensor->type, (uint8_t *) tensor->data + in_buff_offs, f32_output + out_buff_offs, thr_elems); + workers.emplace_back(compute, type, (const uint8_t *) data + in_buff_offs, f32_output + out_buff_offs, thr_elems); in_buff_offs += thr_block_bytes; out_buff_offs += thr_elems; } @@ -410,6 +408,12 @@ static ggml_type tensor_type_fallback(quantize_state_impl & qs, const ggml_tenso case GGML_TYPE_Q5_K: return_type = GGML_TYPE_Q5_1; break; case GGML_TYPE_Q6_K: return_type = GGML_TYPE_Q8_0; break; default: + if (qk_k <= 32) { + // the target is already a 32-block type, so there is no smaller block to demote to + // the check below turns it into F16, as a 256-block type does when its fallback does not fit + return_type = target_type; + break; + } throw std::runtime_error(format("no tensor type fallback is defined for type %s", ggml_type_name(target_type))); } @@ -690,7 +694,21 @@ static ggml_type llama_tensor_get_type(quantize_state_impl & qs, const llama_mod return tensor->type; } if (params->token_embedding_type < GGML_TYPE_COUNT && tm.category == tensor_category::TOKEN_EMBD) { - return params->token_embedding_type; + // per_layer_token_embd follows --token-embedding-type by default, but it is a large + // separate table, so let an explicit --tensor-type name it + bool named = false; + if (std::strcmp(tensor->name, "per_layer_token_embd.weight") == 0) { + const std::string tensor_name(tensor->name); + for (const auto & [pattern, qtype] : qs.tensor_type_patterns) { + if (std::regex_search(tensor_name, pattern)) { + named = true; + break; + } + } + } + if (!named) { + return params->token_embedding_type; + } } if (params->output_tensor_type < GGML_TYPE_COUNT && tm.category == tensor_category::OUTPUT) { return params->output_tensor_type; @@ -1102,6 +1120,8 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: std::vector> work; std::vector> f32_conv_buf; + const size_t max_buf_size = params->max_buf_size ? params->max_buf_size : LLAMA_QUANT_MAX_BUF_SIZE; + int cur_split = -1; std::ofstream fout; auto close_ofstream = [&]() { @@ -1152,15 +1172,13 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: const size_t tensor_size = ggml_nbytes(tensor); - if (!params->dry_run) { - if (!ml.use_mmap) { - if (read_data.size() < tensor_size) { - read_data.resize(tensor_size); - } - tensor->data = read_data.data(); + // read a byte range of the current tensor + auto load_range = [&](size_t offs, size_t size) -> const void * { + if (!ml.use_mmap && read_data.size() < size) { + read_data.resize(size); } - ml.load_data_for(tensor); - } + return ml.load_data_range(weight, offs, size, read_data.data()); + }; LLAMA_LOG_INFO("[%4d/%4d] %-36s - [%s], type = %6s, ", ++idx, ml.n_tensors, @@ -1175,7 +1193,6 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: // in then there's nothing to do. bool quantize = cur_type != new_type; - void * new_data; size_t new_size; if (params->dry_run) { @@ -1199,12 +1216,18 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: } else { // no --dry-run, perform quantization if (!quantize) { - new_data = tensor->data; new_size = tensor_size; LLAMA_LOG_INFO("size = %8.3f MiB\n", tensor_size/1024.0/1024.0); - } else { - const int64_t nelements = ggml_nelements(tensor); + // copy in slabs of whole rows, so that each slab can be validated + const size_t row_size = ggml_row_size(tensor->type, tensor->ne[0]); + const size_t slab_size = std::max(row_size, (max_buf_size/row_size)*row_size); + + for (size_t offs = 0; offs < tensor_size; offs += slab_size) { + const size_t size = std::min(slab_size, tensor_size - offs); + fout.write((const char *) load_range(offs, size), size); + } + } else { const float * imatrix = nullptr; if (imatrix_data) { auto it = imatrix_data->find(tm.remapped_imatrix_name); @@ -1236,43 +1259,60 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: throw std::runtime_error(format("Missing importance matrix for tensor %s in a very low-bit quantization", tensor->name)); } - float * f32_data; - - if (tensor->type == GGML_TYPE_F32) { - f32_data = (float *) tensor->data; - } else if (ggml_is_quantized(tensor->type) && !params->allow_requantize) { + if (ggml_is_quantized(tensor->type) && !params->allow_requantize) { throw std::runtime_error(format("requantizing from type %s is disabled", ggml_type_name(tensor->type))); - } else { - llama_tensor_dequantize_impl(tensor, f32_conv_buf, workers, nelements, nthread); - f32_data = (float *) f32_conv_buf.data(); } LLAMA_LOG_INFO("converting to %s .. ", ggml_type_name(new_type)); fflush(stdout); - if (work.size() < (size_t)nelements * 4) { - work.resize(nelements * 4); // upper bound on size - } - new_data = work.data(); - const int64_t n_per_row = tensor->ne[0]; const int64_t nrows = tensor->ne[1]; + const size_t row_size_src = ggml_row_size(tensor->type, n_per_row); + const size_t row_size_dst = ggml_row_size(new_type, n_per_row); + + // process the rows in slabs, so that the buffers stay below max_buf_size + const size_t bytes_per_row = row_size_src + row_size_dst + (tensor->type == GGML_TYPE_F32 ? 0 : n_per_row*sizeof(float)); + const int64_t nrows_slab = std::max(1, std::min(nrows, max_buf_size/bytes_per_row)); + static const int64_t min_chunk_size = 32 * 512; const int64_t chunk_size = (n_per_row >= min_chunk_size ? n_per_row : n_per_row * ((min_chunk_size + n_per_row - 1)/n_per_row)); - const int64_t nelements_matrix = tensor->ne[0] * tensor->ne[1]; - const int64_t nchunk = (nelements_matrix + chunk_size - 1)/chunk_size; - const int64_t nthread_use = nthread > 1 ? std::max((int64_t)1, std::min((int64_t)nthread, nchunk)) : 1; - // quantize each expert separately since they have different importance matrices new_size = 0; for (int64_t i03 = 0; i03 < tensor->ne[2]; ++i03) { - const float * f32_data_03 = f32_data + i03 * nelements_matrix; - void * new_data_03 = (char *)new_data + ggml_row_size(new_type, n_per_row) * i03 * nrows; const float * imatrix_03 = imatrix ? imatrix + i03 * n_per_row : nullptr; - new_size += llama_tensor_quantize_impl(new_type, f32_data_03, new_data_03, chunk_size, nrows, n_per_row, imatrix_03, workers, nthread_use); + for (int64_t ir = 0; ir < nrows; ir += nrows_slab) { + const int64_t nrows_cur = std::min(nrows_slab, nrows - ir); + const int64_t nelements_cur = nrows_cur * n_per_row; + + const void * src = load_range((i03*nrows + ir)*row_size_src, nrows_cur*row_size_src); + + const float * f32_data; + if (tensor->type == GGML_TYPE_F32) { + f32_data = (const float *) src; + } else { + if (f32_conv_buf.size() < (size_t) nelements_cur) { + f32_conv_buf.resize(nelements_cur); + } + llama_tensor_dequantize_impl(tensor->type, src, (float *) f32_conv_buf.data(), workers, nelements_cur, nthread); + f32_data = (const float *) f32_conv_buf.data(); + } + + if (work.size() < nrows_cur*row_size_dst) { + work.resize(nrows_cur*row_size_dst); + } + + const int64_t nchunk = (nelements_cur + chunk_size - 1)/chunk_size; + const int64_t nthread_use = nthread > 1 ? std::max((int64_t)1, std::min((int64_t)nthread, nchunk)) : 1; + + const size_t size_cur = llama_tensor_quantize_impl(new_type, f32_data, work.data(), chunk_size, nrows_cur, n_per_row, imatrix_03, workers, nthread_use); + + fout.write((const char *) work.data(), size_cur); + new_size += size_cur; + } } LLAMA_LOG_INFO("size = %8.2f MiB -> %8.2f MiB\n", tensor_size/1024.0/1024.0, new_size/1024.0/1024.0); } @@ -1282,10 +1322,8 @@ static void llama_model_quantize_impl(const std::string & fname_inp, const std:: // update the gguf metadata as we go gguf_set_tensor_type(ctx_outs[cur_split].get(), metadata[i].name.c_str(), new_type); GGML_ASSERT(gguf_get_tensor_size(ctx_outs[cur_split].get(), gguf_find_tensor(ctx_outs[cur_split].get(), metadata[i].name.c_str())) == new_size); - gguf_set_tensor_data(ctx_outs[cur_split].get(), metadata[i].name.c_str(), new_data); - // write tensor data + padding - fout.write((const char *) new_data, new_size); + // tensor data is already written, add the padding zeros(fout, GGML_PAD(new_size, align) - new_size); // unmap the tensor to free memory @@ -1332,7 +1370,8 @@ llama_model_quantize_params llama_model_quantize_default_params() { /*.imatrix =*/ nullptr, /*.kv_overrides =*/ nullptr, /*.tensor_type =*/ nullptr, - /*.prune_layers =*/ nullptr + /*.prune_layers =*/ nullptr, + /*.max_buf_size =*/ LLAMA_QUANT_MAX_BUF_SIZE }; return result; diff --git a/src/llama.cpp b/src/llama.cpp index 1609fec88ddf..9c841ee35213 100644 --- a/src/llama.cpp +++ b/src/llama.cpp @@ -318,6 +318,8 @@ static std::pair llama_model_load(struct gguf_context * meta llama_model_loader ml(metadata, set_tensor_data, set_tensor_data_ud, fname, splits, file, params.load_mode, params.check_tensors, params.no_alloc, params.load_mtp, params.kv_overrides, params.tensor_buft_overrides); + ml.tensor_read_lazy = params.tensor_read_lazy; + ml.print_info(); std::unique_ptr model_ptr(llama_model_create(ml, params)); diff --git a/src/models/dflash.cpp b/src/models/dflash.cpp index ff40c16b22e4..f2c7d1d2462e 100644 --- a/src/models/dflash.cpp +++ b/src/models/dflash.cpp @@ -7,6 +7,18 @@ void llama_model_dflash::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, false); + hparams.f_final_logit_softcapping = 0.0f; + ml.get_key(LLM_KV_FINAL_LOGIT_SOFTCAPPING, hparams.f_final_logit_softcapping, false); + + // drafts for M-RoPE targets carry degenerate sections [n_rot/2, 0, 0, 0] + ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, false); + + ml.get_key(LLM_KV_DFLASH_BLOCK_SIZE, hparams.dflash_block_size, false); + ml.get_key(LLM_KV_DFLASH_CONV_KERNEL_SIZE, hparams.dflash_conv_kernel_size, false); + ml.get_key(LLM_KV_DFLASH_CONV_GROUP_SIZE, hparams.dflash_conv_group_size, false); + ml.get_key(LLM_KV_DFLASH_SELECTOR_RANK, hparams.dflash_selector_rank, false); + ml.get_key(LLM_KV_DFLASH_SELECTOR_TOP_K, hparams.dflash_selector_top_k, false); if (!ml.get_arr(LLM_KV_TARGET_LAYERS, target_layer_ids, false)) { throw std::runtime_error("DFlash model requires 'target_layers' in GGUF metadata"); @@ -103,15 +115,39 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { if (markov_meta) { const int64_t dspark_markov_rank = markov_meta->ne[0]; - dspark_markov_w1 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W1, "weight"), { dspark_markov_rank, n_vocab }, 0); - dspark_markov_w2 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "weight"), { dspark_markov_rank, n_vocab_draft }, 0); + dspark_markov_w1 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W1, "weight"), { dspark_markov_rank, n_vocab }, 0); + dspark_markov_w2 = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "weight"), { dspark_markov_rank, n_vocab_draft }, 0); + dspark_markov_w2_s = create_tensor(tn(LLM_TENSOR_DSPARK_MARKOV_W2, "scale"), { 1 }, TENSOR_NOT_REQUIRED); - dspark_conf_proj = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "weight"), { n_embd + dspark_markov_rank, 1 }, 0); + dspark_conf_proj = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "weight"), { n_embd + dspark_markov_rank, 1 }, TENSOR_NOT_REQUIRED); dspark_conf_proj_b = create_tensor(tn(LLM_TENSOR_DSPARK_CONF_PROJ, "bias"), { 1 }, TENSOR_NOT_REQUIRED); LLAMA_LOG_INFO("%s: DFlash with DSpark markov head (rank = %lld)\n", __func__, (long long) dspark_markov_rank); } + const struct ggml_tensor * selector_meta = ml->get_tensor_meta("selector_hidden.weight"); + if (selector_meta) { + const int64_t rank = hparams.dflash_selector_rank; + if (rank <= 0 || hparams.dflash_block_size <= 0 || hparams.dflash_selector_top_k <= 0 || + hparams.dflash_conv_kernel_size <= 0 || hparams.dflash_conv_group_size <= 0) { + throw std::runtime_error("DFlash2 model is missing conv/selector metadata"); + } + if (n_embd % hparams.dflash_conv_group_size != 0) { + throw std::runtime_error("DFlash2 hidden size must be divisible by conv_group_size"); + } + if (n_embd < hparams.dflash_selector_top_k * (hparams.dflash_selector_top_k + 1)) { + throw std::runtime_error("DFlash2 hidden size is too small for the selector lattice"); + } + + dflash_selector_prev = create_tensor(tn(LLM_TENSOR_DFLASH_SELECTOR_PREV, "weight"), { rank, n_vocab }, 0); + dflash_selector_next = create_tensor(tn(LLM_TENSOR_DFLASH_SELECTOR_NEXT, "weight"), { rank, n_vocab }, 0); + dflash_selector_hidden = create_tensor(tn(LLM_TENSOR_DFLASH_SELECTOR_HIDDEN, "weight"), { n_embd, rank }, 0); + + LLAMA_LOG_INFO("%s: DFlash2 conv kernel = %u, group = %u, selector rank = %u, top-k = %u\n", __func__, + hparams.dflash_conv_kernel_size, hparams.dflash_conv_group_size, + hparams.dflash_selector_rank, hparams.dflash_selector_top_k); + } + fc = create_tensor(tn(LLM_TENSOR_FC, "weight"), { n_embd_inp, n_embd }, 0); fc_s = create_tensor(tn(LLM_TENSOR_FC, "scale"), { 1 }, TENSOR_NOT_REQUIRED); output_norm_enc = create_tensor(tn(LLM_TENSOR_ENC_OUTPUT_NORM, "weight"), { n_embd }, 0); // encoder hidden_norm (after fc) @@ -184,10 +220,23 @@ void llama_model_dflash::load_arch_tensors(llama_model_loader &) { layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", i), { n_embd_head_k }, 0); layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", i), { n_embd_head_k }, 0); + // optional per-head attention sinks (e.g. Nemotron DSpark) + layer.attn_sinks = create_tensor(tn(LLM_TENSOR_ATTN_SINKS, "weight", i), { n_head }, TENSOR_NOT_REQUIRED); + layer.ffn_norm = create_tensor(tn(LLM_TENSOR_FFN_NORM, "weight", i), { n_embd }, 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); + + if (selector_meta) { + const int64_t kernel = hparams.dflash_conv_kernel_size; + const int64_t groups = n_embd / hparams.dflash_conv_group_size; + const int64_t projected = 2 * kernel * groups; + layer.dflash_attn_conv_base = create_tensor(tn(LLM_TENSOR_DFLASH_ATTN_CONV_BASE, i), { n_embd, kernel, 2 }, 0); + layer.dflash_attn_conv_proj = create_tensor(tn(LLM_TENSOR_DFLASH_ATTN_CONV_PROJ, "weight", i), { n_embd, projected }, 0); + layer.dflash_ffn_conv_base = create_tensor(tn(LLM_TENSOR_DFLASH_FFN_CONV_BASE, i), { n_embd, kernel, 2 }, 0); + layer.dflash_ffn_conv_proj = create_tensor(tn(LLM_TENSOR_DFLASH_FFN_CONV_PROJ, "weight", i), { n_embd, projected }, 0); + } } } @@ -245,7 +294,10 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model & ggml_tensor * w1 = model.dspark_markov_w1; ggml_tensor * w2 = model.dspark_markov_w2; - GGML_ASSERT(w1 && w2 && model.dspark_conf_proj && "DSpark markov/confidence weights not loaded"); + GGML_ASSERT(w1 && w2 && "DSpark markov weights not loaded"); + + // confidence head is optional + const bool has_conf = model.dspark_conf_proj != nullptr; ggml_tensor * base = res->t_logits; // [n_vocab, n_tokens] const int64_t n_vocab = base->ne[0]; @@ -276,23 +328,22 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model & ggml_tensor * prev = ggml_view_2d(ctx0, tokens, 1, n_blocks, token_stride, 0); prev = ggml_cont_1d(ctx0, prev, n_blocks); - // confidence head input: predicts per-position acceptance - ggml_tensor * conf_inp = res->t_embd; // [n_embd, n_tok] - ggml_tensor * cat = nullptr; ggml_tensor * cat_conf = nullptr; if (!sample_from_anchor) { // bonus anchor slot: pass the logits through unbiased, pad the (unread) confidence column - cat = ggml_cont(ctx0, ggml_view_2d(ctx0, base, n_vocab, n_blocks, base_stride, 0)); - cat_conf = ggml_sigmoid(ctx0, ggml_cont(ctx0, ggml_view_2d(ctx0, base, 1, n_blocks, base_stride, 0))); + cat = ggml_cont(ctx0, ggml_view_2d(ctx0, base, n_vocab, n_blocks, base_stride, 0)); + if (has_conf) { + cat_conf = ggml_sigmoid(ctx0, ggml_cont(ctx0, ggml_view_2d(ctx0, base, 1, n_blocks, base_stride, 0))); + } } // TODO: the in-graph chain is greedy (argmax); sampling params affect only the final // token pick, not the Markov conditioning path for (int64_t i = i_draft_beg; i < block_drafts; ++i) { - ggml_tensor * w1_prev = ggml_get_rows(ctx0, w1, prev); // [R, n_blocks] - ggml_tensor * bias = ggml_mul_mat(ctx0, w2, w1_prev); // [n_vocab_draft, n_blocks] + ggml_tensor * w1_prev = ggml_get_rows(ctx0, w1, prev); // [R, n_blocks] + ggml_tensor * bias = g.build_lora_mm(w2, w1_prev, model.dspark_markov_w2_s); // [n_vocab_draft, n_blocks] if (model.d2t) { // reduced draft vocab: scatter the bias to the target rows (base is -inf on the others) const int64_t n_draft_vocab = bias->ne[0]; @@ -309,17 +360,21 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model & cat = cat ? ggml_concat(ctx0, cat, col, 1) : col; - // conf(i) = sigmoid(conf_proj . [conf_inp(i); markov_w1[prev(i)]] + b) -- [1, n_blocks] - ggml_tensor * conf_inp_i = ggml_view_2d(ctx0, conf_inp, conf_inp->ne[0], n_blocks, - (size_t) block_drafts * conf_inp->nb[1], i*conf_inp->nb[1]); - ggml_tensor * feat = ggml_concat(ctx0, ggml_cont(ctx0, conf_inp_i), w1_prev, 0); - ggml_tensor * conf = ggml_mul_mat(ctx0, model.dspark_conf_proj, feat); - if (model.dspark_conf_proj_b) { - conf = ggml_add(ctx0, conf, model.dspark_conf_proj_b); - } - conf = ggml_sigmoid(ctx0, conf); + if (has_conf) { + // confidence head input: predicts per-position acceptance + ggml_tensor * conf_inp = res->t_embd; // [n_embd, n_tok] + // conf(i) = sigmoid(conf_proj . [conf_inp(i); markov_w1[prev(i)]] + b) -- [1, n_blocks] + ggml_tensor * conf_inp_i = ggml_view_2d(ctx0, conf_inp, conf_inp->ne[0], n_blocks, + (size_t) block_drafts * conf_inp->nb[1], i*conf_inp->nb[1]); + ggml_tensor * feat = ggml_concat(ctx0, ggml_cont(ctx0, conf_inp_i), w1_prev, 0); + ggml_tensor * conf = ggml_mul_mat(ctx0, model.dspark_conf_proj, feat); + if (model.dspark_conf_proj_b) { + conf = ggml_add(ctx0, conf, model.dspark_conf_proj_b); + } + conf = ggml_sigmoid(ctx0, conf); - cat_conf = cat_conf ? ggml_concat(ctx0, cat_conf, conf, 1) : conf; + cat_conf = cat_conf ? ggml_concat(ctx0, cat_conf, conf, 1) : conf; + } if (i + 1 < block_drafts) { prev = ggml_argmax(ctx0, col); @@ -331,7 +386,7 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model & out = ggml_cont(ctx0, ggml_permute(ctx0, out, 0, 2, 1, 3)); // [n_vocab, block_drafts, n_blocks] out = ggml_reshape_2d(ctx0, out, n_vocab, n_tok); - { + if (has_conf) { ggml_tensor * conf = ggml_reshape_3d(ctx0, cat_conf, 1, n_blocks, block_drafts); conf = ggml_cont(ctx0, ggml_permute(ctx0, conf, 0, 2, 1, 3)); conf = ggml_reshape_2d(ctx0, conf, 1, n_tok); @@ -346,6 +401,167 @@ static void build_dspark_markov_head(llm_graph_context & g, const llama_model & ggml_build_forward_expand(g.gf, out); } +static ggml_tensor * build_dflash2_conv( + llm_graph_context & g, + ggml_tensor * hidden, + ggml_tensor * dynamic, + ggml_tensor * base, + int side) { + const auto & hparams = g.hparams; + const int64_t hidden_size = hidden->ne[0]; + const int64_t n_tokens = hidden->ne[1]; + const int64_t n_blocks = g.ubatch.n_seqs_unq; + const int64_t kernel_size = hparams.dflash_conv_kernel_size; + const int64_t group_size = hparams.dflash_conv_group_size; + const int64_t n_groups = hidden_size / group_size; + + GGML_ASSERT(n_blocks > 0 && n_tokens % n_blocks == 0); + GGML_ASSERT(dynamic && base && side >= 0 && side < 2); + + const int64_t block_size = n_tokens / n_blocks; + ggml_context * ctx0 = g.ctx0; + // ggml_cont copies even when the tensor is already contiguous + if (!ggml_is_contiguous(hidden) || hidden->ne[1] != n_tokens) { + hidden = ggml_cont_2d(ctx0, hidden, hidden_size, n_tokens); + } + if (!ggml_is_contiguous(dynamic) || dynamic->ne[1] != n_tokens) { + dynamic = ggml_cont_2d(ctx0, dynamic, dynamic->ne[0], n_tokens); + } + ggml_tensor * blocks = ggml_reshape_3d(ctx0, hidden, hidden_size, block_size, n_blocks); + ggml_tensor * coeffs = ggml_reshape_4d(ctx0, dynamic, n_groups, kernel_size, 2, n_tokens); + ggml_tensor * coeffs_side = ggml_view_3d(ctx0, coeffs, n_groups, kernel_size, n_tokens, + coeffs->nb[1], coeffs->nb[3], side * coeffs->nb[2]); + + ggml_tensor * coeff_all = ggml_cont(ctx0, coeffs_side); + coeff_all = ggml_reshape_4d(ctx0, coeff_all, 1, n_groups, kernel_size, n_tokens); + coeff_all = ggml_repeat_4d(ctx0, coeff_all, group_size, n_groups, kernel_size, n_tokens); + + ggml_tensor * base_side = ggml_reshape_4d(ctx0, + ggml_view_1d(ctx0, base, hidden_size * kernel_size, side * base->nb[2]), + group_size, n_groups, kernel_size, 1); + + ggml_tensor * weight_all = ggml_add(ctx0, coeff_all, base_side); + + ggml_tensor * result = nullptr; + for (int64_t tap = 0; tap < kernel_size; ++tap) { + ggml_tensor * values = blocks; + if (tap > 0) { + ggml_tensor * zeros = ggml_fill(ctx0, + ggml_new_tensor_3d(ctx0, hidden->type, hidden_size, std::min(tap, block_size), n_blocks), 0.0f); + if (tap < block_size) { + ggml_tensor * previous = ggml_view_3d(ctx0, blocks, hidden_size, block_size - tap, n_blocks, + blocks->nb[1], blocks->nb[2], 0); + values = ggml_concat(ctx0, zeros, previous, 1); + } else { + values = zeros; + } + } + values = ggml_reshape_2d(ctx0, values, hidden_size, n_tokens); + + ggml_tensor * weight = ggml_reshape_2d(ctx0, + ggml_cont(ctx0, ggml_view_4d(ctx0, weight_all, group_size, n_groups, 1, n_tokens, + weight_all->nb[1], weight_all->nb[2], weight_all->nb[3], tap * weight_all->nb[2])), + hidden_size, n_tokens); + + ggml_tensor * term = ggml_mul(ctx0, weight, values); + result = result ? ggml_add(ctx0, result, term) : term; + } + return result; +} + +// DFlash2 selector: top-k candidates per block position plus the pairwise +// transition scores, packed into the nextn output slot for the CPU-side walk. +static void build_dflash2_selector(llm_graph_context & g, const llama_model & model, ggml_tensor * tokens) { + ggml_context * ctx0 = g.ctx0; + auto & res = g.res; + + const auto & hparams = g.hparams; + const int64_t n_tokens = g.n_tokens; + const int64_t n_embd = g.n_embd; + + const int64_t top_k = hparams.dflash_selector_top_k; + const int64_t rank = hparams.dflash_selector_rank; + const int64_t n_blocks = g.ubatch.n_seqs_unq; + GGML_ASSERT(n_blocks > 0 && n_tokens % n_blocks == 0); + GGML_ASSERT(res->t_logits->ne[1] == n_tokens); + if (!tokens) { + return; + } + + const int64_t tokens_per_block = n_tokens / n_blocks; + const int64_t block_size = std::min(tokens_per_block, hparams.dflash_block_size); + const int64_t row_used = top_k + top_k * top_k; + + ggml_tensor * candidates = ggml_top_k(ctx0, res->t_logits, top_k); + ggml_tensor * logits_rows = ggml_reshape_3d(ctx0, res->t_logits, 1, res->t_logits->ne[0], n_tokens); + ggml_tensor * unary = ggml_reshape_2d(ctx0, + ggml_get_rows(ctx0, logits_rows, candidates), top_k, n_tokens); + ggml_tensor * gate = g.build_lora_mm(model.dflash_selector_hidden, res->t_embd); + + // Everything below indexes [.., tokens_per_block, n_blocks]: the block + // position varies fastest, sequences are the outer dimension. + ggml_tensor * cand_blk = ggml_reshape_3d(ctx0, candidates, top_k, tokens_per_block, n_blocks); + ggml_tensor * unary_blk = ggml_reshape_3d(ctx0, unary, top_k, tokens_per_block, n_blocks); + ggml_tensor * gate_blk = ggml_reshape_3d(ctx0, gate, rank, tokens_per_block, n_blocks); + + // a position's score reads only the candidate sets at pos-1 and pos, so a run + // of positions has no internal dependency and scores in one batched matmul + auto score_run = [&](int64_t beg_pos, int64_t n_pos, ggml_tensor * pred_ids) { + ggml_tensor * cand_run = ggml_cont(ctx0, ggml_view_3d(ctx0, cand_blk, top_k, n_pos, n_blocks, + cand_blk->nb[1], cand_blk->nb[2], beg_pos * cand_blk->nb[1])); + ggml_tensor * unary_run = ggml_cont(ctx0, ggml_view_3d(ctx0, unary_blk, top_k, n_pos, n_blocks, + unary_blk->nb[1], unary_blk->nb[2], beg_pos * unary_blk->nb[1])); + ggml_tensor * gate_run = ggml_cont(ctx0, ggml_view_3d(ctx0, gate_blk, rank, n_pos, n_blocks, + gate_blk->nb[1], gate_blk->nb[2], beg_pos * gate_blk->nb[1])); + + const int64_t n_pred = pred_ids->ne[0] / (n_pos * n_blocks); + + ggml_tensor * successor = ggml_reshape_4d(ctx0, + ggml_get_rows(ctx0, model.dflash_selector_next, ggml_reshape_1d(ctx0, cand_run, top_k * n_pos * n_blocks)), + rank, top_k, n_pos, n_blocks); + ggml_tensor * predecessor = ggml_reshape_4d(ctx0, + ggml_get_rows(ctx0, model.dflash_selector_prev, pred_ids), + rank, n_pred, n_pos, n_blocks); + + ggml_tensor * gate_bcast = ggml_reshape_4d(ctx0, gate_run, rank, 1, n_pos, n_blocks); + ggml_tensor * cond = ggml_mul(ctx0, predecessor, ggml_repeat(ctx0, gate_bcast, predecessor)); + ggml_tensor * score = ggml_mul_mat(ctx0, successor, cond); + if (n_pred == 1) { + score = ggml_repeat_4d(ctx0, score, top_k, top_k, n_pos, n_blocks); + } + ggml_tensor * unary_bcast = ggml_reshape_4d(ctx0, unary_run, top_k, 1, n_pos, n_blocks); + score = ggml_add(ctx0, score, ggml_repeat(ctx0, unary_bcast, score)); + + ggml_tensor * row = ggml_concat(ctx0, + ggml_cast(ctx0, cand_run, GGML_TYPE_F32), + ggml_reshape_3d(ctx0, score, top_k * top_k, n_pos, n_blocks), 0); + return ggml_pad(ctx0, row, n_embd - row_used, 0, 0, 0); + }; + + ggml_tensor * packed = ggml_fill(ctx0, + ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, n_embd, 1, n_blocks), 0.0f); + + if (block_size > 1) { + // Position 1 alone: its predecessor is the anchor token, one id per + // sequence rather than a candidate set. + ggml_tensor * anchor_ids = ggml_cont_1d(ctx0, + ggml_view_2d(ctx0, tokens, 1, n_blocks, tokens_per_block * tokens->nb[0], 0), n_blocks); + packed = ggml_concat(ctx0, packed, score_run(1, 1, anchor_ids), 1); + } + if (block_size > 2) { + ggml_tensor * prev_ids = ggml_reshape_1d(ctx0, + ggml_cont(ctx0, ggml_view_3d(ctx0, cand_blk, top_k, block_size - 2, n_blocks, + cand_blk->nb[1], cand_blk->nb[2], cand_blk->nb[1])), + top_k * (block_size - 2) * n_blocks); + packed = ggml_concat(ctx0, packed, score_run(2, block_size - 2, prev_ids), 1); + } + + packed = ggml_reshape_2d(ctx0, packed, n_embd, block_size * n_blocks); + g.cb(packed, "dflash2_lattice", -1); + res->t_h_nextn = packed; + ggml_build_forward_expand(g.gf, packed); +} + // DFlash decoder, dual-mode by batch type: // * embd batch -> fused target features: project + inject K/V into the cache. // * token batch -> noise-block diffusion: attend over [committed, MASK...] to generate draft tokens @@ -370,6 +586,20 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra const float kq_scale = 1.0f/sqrtf(float(n_embd_head)); + // drafts for M-RoPE targets use degenerate sections (temporal dim only) + int sections[4]; + std::copy(std::begin(hparams.rope_sections), std::begin(hparams.rope_sections) + 4, sections); + + auto build_rope = [&](ggml_tensor * cur, ggml_tensor * pos) { + return rope_type == GGML_ROPE_TYPE_MROPE + ? ggml_rope_multi(ctx0, cur, pos, nullptr, + n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow) + : ggml_rope_ext(ctx0, cur, pos, nullptr, + n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + }; + // KV cache injection if (ubatch.embd) { auto inp = std::make_unique(n_embd); @@ -392,11 +622,7 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il); - Kcur = ggml_rope_ext( - ctx0, Kcur, inp_pos, nullptr, - n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow - ); + Kcur = build_rope(Kcur, inp_pos); cb(Kcur, "Kcur_injected", il); cb(Vcur, "Vcur_injected", il); @@ -450,6 +676,7 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra inp->tokens = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_tokens); ggml_set_input(inp->tokens); + res->t_inp_tokens = inp->tokens; ggml_tensor * inp_tokens = inp->tokens; @@ -464,6 +691,13 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra ggml_tensor * noise_norm = build_norm(inpL, layer.attn_norm, NULL, LLM_NORM_RMS, il); cb(noise_norm, "noise_norm", il); + ggml_tensor * attn_dynamic = nullptr; + if (layer.dflash_attn_conv_proj) { + attn_dynamic = build_lora_mm(layer.dflash_attn_conv_proj, noise_norm); + noise_norm = build_dflash2_conv(*this, noise_norm, attn_dynamic, layer.dflash_attn_conv_base, 0); + cb(noise_norm, "attn_conv_in", il); + } + ggml_tensor * Qcur = build_lora_mm(layer.wq, noise_norm); ggml_tensor * Kcur = build_lora_mm(layer.wk, noise_norm); ggml_tensor * Vcur = build_lora_mm(layer.wv, noise_norm); @@ -475,24 +709,21 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra Qcur = build_norm(Qcur, layer.attn_q_norm, NULL, LLM_NORM_RMS, il); Kcur = build_norm(Kcur, layer.attn_k_norm, NULL, LLM_NORM_RMS, il); - Qcur = ggml_rope_ext( - ctx0, Qcur, inp_pos, nullptr, - 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, nullptr, - n_rot, rope_type, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow - ); + Qcur = build_rope(Qcur, inp_pos); + Kcur = build_rope(Kcur, inp_pos); cb(Qcur, "Qcur", il); cb(Kcur, "Kcur", il); cb(Vcur, "Vcur", il); // cache-aware, non-causal attention ggml_tensor * cur = use_iswa - ? build_attn(inp_attn_iswa, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il) - : build_attn(inp_attn, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); + ? build_attn(inp_attn_iswa, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, layer.attn_sinks, nullptr, kq_scale, il) + : build_attn(inp_attn, layer.wo, NULL, NULL, Qcur, Kcur, Vcur, nullptr, layer.attn_sinks, nullptr, kq_scale, il); + + if (attn_dynamic) { + cur = build_dflash2_conv(*this, cur, attn_dynamic, layer.dflash_attn_conv_base, 1); + cb(cur, "attn_conv_out", il); + } ggml_tensor * ffn_inp = ggml_add(ctx0, cur, inpL); cb(ffn_inp, "ffn_inp", il); @@ -500,6 +731,13 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra cur = build_norm(ffn_inp, layer.ffn_norm, NULL, LLM_NORM_RMS, il); cb(cur, "ffn_norm", il); + ggml_tensor * ffn_dynamic = nullptr; + if (layer.dflash_ffn_conv_proj) { + ffn_dynamic = build_lora_mm(layer.dflash_ffn_conv_proj, cur); + cur = build_dflash2_conv(*this, cur, ffn_dynamic, layer.dflash_ffn_conv_base, 0); + cb(cur, "ffn_conv_in", il); + } + cur = build_ffn(cur, layer.ffn_up, NULL, layer.ffn_up_s, layer.ffn_gate, NULL, layer.ffn_gate_s, @@ -508,6 +746,11 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra LLM_FFN_SILU, LLM_FFN_PAR, il); cb(cur, "ffn_out", il); + if (ffn_dynamic) { + cur = build_dflash2_conv(*this, cur, ffn_dynamic, layer.dflash_ffn_conv_base, 1); + cb(cur, "ffn_conv_out", il); + } + cur = ggml_add(ctx0, cur, ffn_inp); cb(cur, "l_out", il); @@ -532,6 +775,19 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra cur = build_lora_mm(output, cur, output_s); + // DFlash2 feeds these logits to the selector, so they need the target's output + // transforms; DFlash1 and DSpark read them through the sampler instead + if (model.dflash_selector_hidden) { + if (hparams.f_logit_scale != 0.0f) { + cur = ggml_scale(ctx0, cur, hparams.f_logit_scale); + } + if (hparams.f_final_logit_softcapping > 0.0f) { + cur = ggml_scale(ctx0, cur, 1.0f / hparams.f_final_logit_softcapping); + cur = ggml_tanh(ctx0, cur); + cur = ggml_scale(ctx0, cur, hparams.f_final_logit_softcapping); + } + } + // reduced-draft-vocab exports: scatter the draft logits to the target vocabulary via d2t if (model.d2t) { const int64_t n_draft_vocab = cur->ne[0]; @@ -556,6 +812,10 @@ llama_model_dflash::graph::graph(const llama_model & model, const llm_gra if (model.dspark_markov_w1) { build_dspark_markov_head(*this, model, inp_tokens); } + + if (model.dflash_selector_hidden) { + build_dflash2_selector(*this, model, inp_tokens); + } } // DSV4 DSpark decoder, dual-mode by batch type (see the DFlash decoder above): diff --git a/src/models/gemma4.cpp b/src/models/gemma4.cpp index e44f423bdbc5..aa518c6df504 100644 --- a/src/models/gemma4.cpp +++ b/src/models/gemma4.cpp @@ -50,7 +50,7 @@ void llama_model_gemma4::load_arch_tensors(llama_model_loader &) { tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), {n_embd, n_vocab}, 0); if (n_embd_per_layer > 0) { - per_layer_tok_embd = create_tensor(tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight"), {n_embd_per_layer * n_layer, n_vocab}, 0); + per_layer_tok_embd = create_tensor(tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight"), {n_embd_per_layer * n_layer, n_vocab}, TENSOR_READ_LAZY); per_layer_model_proj = create_tensor(tn(LLM_TENSOR_PER_LAYER_MODEL_PROJ, "weight", 0), {n_embd, n_embd_per_layer * n_layer}, 0); per_layer_proj_norm = create_tensor(tn(LLM_TENSOR_PER_LAYER_PROJ_NORM, "weight", 0), {n_embd_per_layer}, 0); } diff --git a/src/models/minimax-01.cpp b/src/models/minimax-01.cpp index a6ccee1917eb..f14626b2c74f 100644 --- a/src/models/minimax-01.cpp +++ b/src/models/minimax-01.cpp @@ -174,11 +174,9 @@ class llm_graph_input_la : public llm_graph_input_i { bool can_reuse(const llm_graph_params & params) override { bool res = true; - if (params.ubatch.n_seq_tokens > 1) { - res &= ( inp_q_decay && inp_q_decay->ne[2] == params.ubatch.n_seq_tokens); - res &= ( inp_k_decay && inp_k_decay->ne[2] == params.ubatch.n_seq_tokens); - res &= (inp_diag_decay && inp_diag_decay->ne[1] == params.ubatch.n_seq_tokens); - } + res &= ( inp_q_decay && inp_q_decay->ne[2] == params.ubatch.n_seq_tokens); + res &= ( inp_k_decay && inp_k_decay->ne[2] == params.ubatch.n_seq_tokens); + res &= (inp_diag_decay && inp_diag_decay->ne[1] == params.ubatch.n_seq_tokens); return res; } @@ -223,19 +221,17 @@ llama_model_minimax_01::graph::graph(const llama_model & model, const llm_graph_ ggml_set_input(inp->inp_slopes); cb(inp->inp_slopes, "slopes", -1); - if (n_seq_tokens != 1) { - inp->inp_q_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, n_head, n_seq_tokens, n_seqs); - ggml_set_input(inp->inp_q_decay); - cb(inp->inp_q_decay, "q_decay_exp", -1); + inp->inp_q_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, n_head, n_seq_tokens, n_seqs); + ggml_set_input(inp->inp_q_decay); + cb(inp->inp_q_decay, "q_decay_exp", -1); - inp->inp_k_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, n_head, n_seq_tokens, n_seqs); - ggml_set_input(inp->inp_k_decay); - cb(inp->inp_k_decay, "k_decay_exp", -1); + inp->inp_k_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, n_head, n_seq_tokens, n_seqs); + ggml_set_input(inp->inp_k_decay); + cb(inp->inp_k_decay, "k_decay_exp", -1); - inp->inp_diag_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_seq_tokens, n_seq_tokens, n_head, n_seqs); - ggml_set_input(inp->inp_diag_decay); - cb(inp->inp_diag_decay, "diag_decay_exp", -1); - } + inp->inp_diag_decay = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, n_seq_tokens, n_seq_tokens, n_head, n_seqs); + ggml_set_input(inp->inp_diag_decay); + cb(inp->inp_diag_decay, "diag_decay_exp", -1); la = (llm_graph_input_la *) res->add_input(std::move(inp)); @@ -319,41 +315,8 @@ llama_model_minimax_01::graph::graph(const llama_model & model, const llm_graph_ ggml_tensor * qkv = nullptr; ggml_tensor * kv_new = nullptr; - - if (n_seq_tokens == 1) { - // lightning attention - optimized single token case for TG - - ggml_tensor * slopes_neg = ggml_scale(ctx0, slope_rate, -1.0); - cb(slopes_neg, "slopes_neg", il); - - ggml_tensor * ratio = ggml_exp(ctx0, slopes_neg); - cb(ratio, "ratio", il); - - ggml_tensor * ratio_3d = ggml_reshape_3d(ctx0, ratio, 1, 1, n_head); - cb(ratio_3d, "ratio3d", il); - - ggml_tensor * v_trans = ggml_cont(ctx0, ggml_permute(ctx0, Vcur, 1, 2, 0, 3)); - cb(v_trans, "v_trans", il); - - ggml_tensor * k_trans = ggml_cont(ctx0, ggml_permute(ctx0, Kcur, 1, 2, 0, 3)); - cb(k_trans, "k_trans", il); - - ggml_tensor * kv_cur = ggml_mul_mat(ctx0, k_trans, v_trans); - cb(kv_cur, "kv_cur", il); - - ggml_tensor * kv_old_s = ggml_mul(ctx0, kv_old, ratio_3d); - cb(kv_old_s, "kv_old_s", il); - - kv_new = ggml_add(ctx0, kv_old_s, kv_cur); - cb(kv_new, "kv_new", il); - - ggml_tensor * q_trans = ggml_permute(ctx0, Qcur, 0, 2, 1, 3); - cb(q_trans, "q_trans", il); - - qkv = ggml_mul_mat(ctx0, kv_new, q_trans); - cb(qkv, "qkv", il); - } else if(n_seq_tokens > 1) { - // lightning attention - general multi token case for PP + { + // lightning attention ggml_tensor * q_decay_exp = la->inp_q_decay; ggml_tensor * k_decay_exp = la->inp_k_decay; diff --git a/src/models/models.h b/src/models/models.h index 5a9496887099..bcf4745f3bf5 100644 --- a/src/models/models.h +++ b/src/models/models.h @@ -6,6 +6,9 @@ // note: almost all graphs require at least sqrtf, so include cmath globally #include +#include + +class llama_memory_hybrid_idx_context; // // base classes @@ -2285,6 +2288,108 @@ struct llama_model_qwen35 : public llama_model_base { }; +struct llama_model_qwen4exp : public llama_model_base { + llama_model_qwen4exp(const struct llama_model_params & params) : llama_model_base(params) {} + + class llm_graph_input_qsa; + + void load_arch_hparams(llama_model_loader & ml) override; + void load_arch_tensors(llama_model_loader & ml) override; + + struct graph : public llm_build_delta_net_base { + graph(const llama_model & model, const llm_graph_params & params); + private: + // HC replaces every layer norm: residual is [n_embd, hc, n_tokens] + ggml_tensor * build_hc_mix( + ggml_tensor * x, + ggml_tensor * w_norm, + ggml_tensor * w_down, + ggml_tensor * w_up, + ggml_tensor * w_inject, + ggml_tensor ** inject, + int il); + + ggml_tensor * build_hc_combine( + ggml_tensor * residual, + ggml_tensor * block_out, + ggml_tensor * inject, + int il); + + ggml_tensor * build_layer_attn( + llm_graph_input_attn_kv * inp_attn, + const llama_memory_hybrid_idx_context * mctx_hyb, + ggml_tensor * cur, + ggml_tensor * inp_pos, + int * sections, + int il); + + // dense self-attention restricted to the cells that top_k names + ggml_tensor * build_attn_qsa( + llm_graph_input_attn_kv * inp, + ggml_tensor * q_cur, + ggml_tensor * k_cur, + ggml_tensor * v_cur, + ggml_tensor * top_k, + float kq_scale, + int il); + + // the QSA cache layout inputs do not depend on the layer, only on its compress ratio, + // so the layers sharing a ratio share one input set + std::map qsa_inps; + + // QSA: token indices this layer's queries may attend to, or nullptr for dense + ggml_tensor * build_qsa_top_k( + const llama_memory_hybrid_idx_context * mctx_hyb, + ggml_tensor * cur, + ggml_tensor * inp_pos, + ggml_tensor * kq_mask, + int * sections, + int il); + + ggml_tensor * build_layer_attn_linear( + llm_graph_input_rs * inp, + ggml_tensor * cur, + int il); + + ggml_tensor * build_layer_ffn( + ggml_tensor * cur, + int il); + + ggml_tensor * build_norm_gated( + ggml_tensor * input, + ggml_tensor * weights, + ggml_tensor * gate, + int layer); + + // build_rs writes the state tensor in place, so one gather per cache tensor is reused + std::map rs_rows; + + // one conv history per cache tensor: delta-net and PLE each have their own + ggml_tensor * build_conv_state_at( + llm_graph_input_rs * inp, + ggml_tensor * conv_states_all, + ggml_tensor * x, + int64_t state_cols, + int64_t channels, + int il); + + ggml_tensor * build_ple( + llm_graph_input_rs * inp, + const llama_memory_hybrid_idx_context * mctx_hyb, + ggml_tensor * hidden, + int il); + + // returns pair of qkv, z + std::pair build_qkvz( + ggml_tensor * input, + int il); + + const llama_model & model; + }; + + std::unique_ptr build_arch_graph(const llm_graph_params & params) const override; +}; + struct llama_model_qwen35moe : public llama_model_base { llama_model_qwen35moe(const struct llama_model_params & params) : llama_model_base(params) {} void load_arch_hparams(llama_model_loader & ml) override; diff --git a/src/models/nanbeige.cpp b/src/models/nanbeige.cpp index 3a546600fa27..7d5a6bbdb6b4 100644 --- a/src/models/nanbeige.cpp +++ b/src/models/nanbeige.cpp @@ -103,6 +103,7 @@ llama_model_nanbeige::graph::graph(const llama_model & model, const llm_graph_pa ggml_tensor * inp_out_ids = build_inp_out_ids(); for (int il = 0; il < n_layer; ++il) { + res->t_layer_inp[il] = inpL; ggml_tensor * inpSA = inpL; cur = build_norm(inpL, model.layers[il].attn_norm, NULL, LLM_NORM_RMS, il); diff --git a/src/models/qwen4exp.cpp b/src/models/qwen4exp.cpp new file mode 100644 index 000000000000..acfdd5b50038 --- /dev/null +++ b/src/models/qwen4exp.cpp @@ -0,0 +1,1199 @@ +#include "models.h" +#include "llama-impl.h" +#include "llama-memory-hybrid-idx.h" +#include "llama-memory-recurrent.h" + +#include +#include + +void llama_model_qwen4exp::load_arch_hparams(llama_model_loader & ml) { + ml.get_key(LLM_KV_EXPERT_FEED_FORWARD_LENGTH, hparams.n_ff_exp, false); + ml.get_key(LLM_KV_EXPERT_SHARED_FEED_FORWARD_LENGTH, hparams.n_ff_shexp, false); + ml.get_key(LLM_KV_ATTENTION_LAYERNORM_RMS_EPS, hparams.f_norm_rms_eps); + + ml.get_key_or_arr(LLM_KV_ROPE_DIMENSION_SECTIONS, hparams.rope_sections, 4, true); + + ml.get_key(LLM_KV_SSM_CONV_KERNEL, hparams.ssm_d_conv); + ml.get_key(LLM_KV_SSM_INNER_SIZE, hparams.ssm_d_inner); + ml.get_key(LLM_KV_SSM_STATE_SIZE, hparams.ssm_d_state); + ml.get_key(LLM_KV_SSM_TIME_STEP_RANK, hparams.ssm_dt_rank); + ml.get_key(LLM_KV_SSM_GROUP_COUNT, hparams.ssm_n_group); + GGML_ASSERT(hparams.ssm_d_conv > 0 && hparams.ssm_d_inner > 0 && hparams.ssm_d_state > 0 && + hparams.ssm_dt_rank > 0 && hparams.ssm_n_group > 0); + + // HC; low_rank is qwen4exp-specific, DeepSeek-V4 leaves it absent (full rank) + ml.get_key(LLM_KV_HYPER_CONNECTION_COUNT, hparams.dsv4_hc_mult); + ml.get_key(LLM_KV_HYPER_CONNECTION_LOW_RANK, hparams.hc_low_rank); + GGML_ASSERT(hparams.dsv4_hc_mult > 0 && hparams.hc_low_rank > 0); + hparams.n_embd_out_impl = hparams.dsv4_hc_mult * hparams.n_embd; + + ml.get_key(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, hparams.indexer_n_head); + ml.get_key(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, hparams.indexer_head_size); + ml.get_key(LLM_KV_ATTENTION_INDEXER_TOP_K, hparams.indexer_top_k); + GGML_ASSERT(hparams.indexer_n_head > 0 + && hparams.indexer_head_size > 0 + && hparams.indexer_top_k > 0); + ml.get_key_or_arr(LLM_KV_ATTENTION_COMPRESS_RATIOS, hparams.dsv4_compress_ratios, hparams.n_layer_all, false); + + // PLE n-gram hash embeddings; if the key group is absent every field stays zero + hparams.is_ple_impl.reset(); + hparams.ple_n_heads = 0; + + uint32_t n_ple = 0; + ml.get_arr_n(LLM_KV_PLE_LAYERS, n_ple, false); + if (n_ple > 0) { + std::vector ple_layers; + ml.get_arr(LLM_KV_PLE_LAYERS, ple_layers); + GGML_ASSERT(n_ple == 1 && "qwen4exp supports only one PLE layer"); + for (uint32_t il : ple_layers) { + if (il >= hparams.n_layer_all) { + throw std::runtime_error(format("PLE layer %u is out of range", il)); + } + hparams.is_ple_impl.set(il); + } + + ml.get_key(LLM_KV_PLE_NGRAM_SIZE, hparams.ple_ngram_size); + ml.get_key(LLM_KV_PLE_HEADS_PER_NGRAM, hparams.ple_heads_per_ngram); + ml.get_key(LLM_KV_PLE_CONV_KERNEL, hparams.ple_conv_kernel); + ml.get_key(LLM_KV_PLE_EOS_TOKEN_ID, hparams.ple_eos_token_id); + // optional: files written before this key fall back to the EOS token + ml.get_key(LLM_KV_PLE_IMAGE_TOKEN_ID, hparams.ple_image_token_id, false); + ml.get_key(LLM_KV_EMBEDDING_LENGTH_PER_LAYER, hparams.n_embd_per_layer); + GGML_ASSERT(hparams.ple_conv_kernel > 0 && hparams.n_embd_per_layer > 0); + + hparams.ple_n_heads = (hparams.ple_ngram_size - 1) * hparams.ple_heads_per_ngram; + hparams.ple_head_dim = hparams.n_embd_per_layer; + if (hparams.ple_ngram_size < 2 || hparams.ple_ngram_size > LLAMA_MAX_PLE_NGRAM) { + throw std::runtime_error(format("PLE n-gram size %u is out of range", hparams.ple_ngram_size)); + } + if (hparams.ple_n_heads == 0 || hparams.ple_n_heads > LLAMA_MAX_PLE_HEADS) { + throw std::runtime_error(format("PLE head count %u is out of range", hparams.ple_n_heads)); + } + + ml.get_arr(LLM_KV_PLE_LAYER_MULTIPLIERS, hparams.ple_layer_multipliers); + + // the file stores the head ranges as uint64, so read at that width and narrow to the int32 the gather uses + std::array head_offsets = {}; + std::array head_vocab_sizes = {}; + ml.get_arr(LLM_KV_PLE_HEAD_OFFSETS, head_offsets); + ml.get_arr(LLM_KV_PLE_HEAD_VOCAB_SIZES, head_vocab_sizes); + for (uint32_t h = 0; h < hparams.ple_n_heads; ++h) { + if (head_vocab_sizes[h] == 0 || + head_offsets[h] > INT32_MAX || + head_vocab_sizes[h] > INT32_MAX || + head_offsets[h] + head_vocab_sizes[h] > INT32_MAX) { + throw std::runtime_error(format("PLE head %u range does not fit the int32 row index", h)); + } + hparams.ple_head_offsets[h] = (uint32_t) head_offsets[h]; + hparams.ple_head_vocab_sizes[h] = (uint32_t) head_vocab_sizes[h]; + } + } + + // linear attention everywhere except every full_attention_interval-th layer + if (!ml.get_key_or_arr(LLM_KV_ATTENTION_RECURRENT_LAYERS, hparams.is_recr_impl, hparams.n_layer_all, false)) { + uint32_t full_attn_interval = 4; + ml.get_key(LLM_KV_FULL_ATTENTION_INTERVAL, full_attn_interval, false); + GGML_ASSERT(full_attn_interval > 0); + for (uint32_t i = 0; i < hparams.n_layer_all; ++i) { + hparams.is_recr_impl[i] = (i < hparams.n_layer()) && ((i + 1) % full_attn_interval != 0); + } + } + + switch (hparams.n_layer()) { + case 48: type = LLM_TYPE_A3B; break; + default: type = LLM_TYPE_UNKNOWN; + } +} + +void llama_model_qwen4exp::load_arch_tensors(llama_model_loader & ml) { + LLAMA_LOAD_LOCALS; + + const int64_t hc = hparams.dsv4_hc_mult; + const int64_t hc_dim = hc * n_embd; + const int64_t hc_lr = hparams.hc_low_rank; + + tok_embd = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, 0); + + // there is no output_norm: the final hyper-connection mixer carries it + hc_head_norm = create_tensor(tn(LLM_TENSOR_HC_HEAD_NORM, "weight"), { hc_dim }, 0); + hc_head_down = create_tensor(tn(LLM_TENSOR_HC_HEAD_DOWN, "weight"), { hc_dim, hc_lr }, 0); + hc_head_up = create_tensor(tn(LLM_TENSOR_HC_HEAD_UP, "weight"), { hc_lr, hc_dim }, 0); + + output = create_tensor(tn(LLM_TENSOR_OUTPUT, "weight"), { n_embd, n_vocab }, TENSOR_NOT_REQUIRED); + if (output == NULL) { + output = create_tensor(tn(LLM_TENSOR_TOKEN_EMBD, "weight"), { n_embd, n_vocab }, TENSOR_DUPLICATED); + } + + // flat [ple_head_dim, n_rows] gather target; n_rows is padded, so read it back + if (hparams.ple_n_heads > 0) { + const std::string ple_name = tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight").str(); + const auto & ple_w = ml.require_weight(ple_name.c_str()); + const int64_t ple_rows = ple_w.tensor->ne[1]; + + // sanity check + for (uint32_t h = 0; h < hparams.ple_n_heads; ++h) { + if ((int64_t) hparams.ple_head_offsets[h] + hparams.ple_head_vocab_sizes[h] > ple_rows) { + throw std::runtime_error(format("PLE head %u range exceeds the %" PRId64 " table rows", h, ple_rows)); + } + } + per_layer_tok_embd = create_tensor(tn(LLM_TENSOR_PER_LAYER_TOKEN_EMBD, "weight"), + { hparams.ple_head_dim, ple_rows }, TENSOR_READ_LAZY); + } + + for (int il = 0; il < n_layer; ++il) { + auto & layer = layers[il]; + + const int64_t n_ff_exp = hparams.n_ff_exp ? hparams.n_ff_exp : n_ff / n_expert_used; + const int64_t n_ff_shexp = hparams.n_ff_shexp ? hparams.n_ff_shexp : n_ff; + + const int64_t head_k_dim = hparams.ssm_d_state; + const int64_t head_v_dim = hparams.ssm_d_state; + const int64_t n_k_heads = hparams.ssm_n_group; + const int64_t n_v_heads = hparams.ssm_dt_rank; + const int64_t key_dim = head_k_dim * n_k_heads; + const int64_t value_dim = head_v_dim * n_v_heads; + const int64_t conv_dim = key_dim * 2 + value_dim; + + // two HC modules per layer: before the token mixer, before the MoE + layer.hc_attn_norm = create_tensor(tn(LLM_TENSOR_HC_ATTN_NORM, "weight", il), { hc_dim }, 0); + layer.hc_attn_down = create_tensor(tn(LLM_TENSOR_HC_ATTN_DOWN, "weight", il), { hc_dim, hc_lr }, 0); + layer.hc_attn_up = create_tensor(tn(LLM_TENSOR_HC_ATTN_UP, "weight", il), { hc_lr, hc_dim }, 0); + layer.hc_attn_inject = create_tensor(tn(LLM_TENSOR_HC_ATTN_INJECT, "weight", il), { hc_dim, hc }, 0); + layer.hc_ffn_norm = create_tensor(tn(LLM_TENSOR_HC_FFN_NORM, "weight", il), { hc_dim }, 0); + layer.hc_ffn_down = create_tensor(tn(LLM_TENSOR_HC_FFN_DOWN, "weight", il), { hc_dim, hc_lr }, 0); + layer.hc_ffn_up = create_tensor(tn(LLM_TENSOR_HC_FFN_UP, "weight", il), { hc_lr, hc_dim }, 0); + layer.hc_ffn_inject = create_tensor(tn(LLM_TENSOR_HC_FFN_INJECT, "weight", il), { hc_dim, hc }, 0); + + if (!hparams.is_recr(il)) { + // full attention: wq holds [q|gate] interleaved per head + create_tensor_qkv(layer, il, n_embd, n_embd_head_k * n_head * 2, n_embd_k_gqa, n_embd_v_gqa, 0); + layer.wo = create_tensor(tn(LLM_TENSOR_ATTN_OUT, "weight", il), { n_embd_head_k * n_head, n_embd }, 0); + + layer.attn_q_norm = create_tensor(tn(LLM_TENSOR_ATTN_Q_NORM, "weight", il), { n_embd_head_k }, 0); + layer.attn_k_norm = create_tensor(tn(LLM_TENSOR_ATTN_K_NORM, "weight", il), { n_embd_head_k }, 0); + + const int64_t idx_dim = hparams.indexer_head_size; + layer.index_q_proj = create_tensor(tn(LLM_TENSOR_INDEXER_Q_PROJ, "weight", il), { n_embd, hparams.indexer_n_head * idx_dim }, 0); + layer.index_k_proj = create_tensor(tn(LLM_TENSOR_INDEXER_K_PROJ, "weight", il), { n_embd, idx_dim }, 0); + layer.index_q_norm = create_tensor(tn(LLM_TENSOR_INDEXER_Q_NORM, "weight", il), { idx_dim }, 0); + layer.index_k_norm = create_tensor(tn(LLM_TENSOR_INDEXER_K_NORM, "weight", il), { idx_dim }, 0); + } else { + layer.wqkv = create_tensor(tn(LLM_TENSOR_ATTN_QKV, "weight", il), { n_embd, key_dim * 2 + value_dim }, 0); + layer.wqkv_gate = create_tensor(tn(LLM_TENSOR_ATTN_GATE, "weight", il), { n_embd, value_dim }, 0); + layer.ssm_conv1d = create_tensor(tn(LLM_TENSOR_SSM_CONV1D, "weight", il), { hparams.ssm_d_conv, conv_dim }, 0); + layer.ssm_dt = create_tensor(tn(LLM_TENSOR_SSM_DT, "bias", il), { hparams.ssm_dt_rank }, 0); + layer.ssm_a = create_tensor(tn(LLM_TENSOR_SSM_A_NOSCAN, il), { hparams.ssm_dt_rank }, 0); + layer.ssm_beta = create_tensor(tn(LLM_TENSOR_SSM_BETA, "weight", il), { n_embd, n_v_heads }, 0); + layer.ssm_alpha = create_tensor(tn(LLM_TENSOR_SSM_ALPHA, "weight", il), { n_embd, n_v_heads }, 0); + layer.ssm_norm = create_tensor(tn(LLM_TENSOR_SSM_NORM, "weight", il), { head_v_dim }, 0); + layer.ssm_out = create_tensor(tn(LLM_TENSOR_SSM_OUT, "weight", il), { value_dim, n_embd }, 0); + } + + if (hparams.is_ple(il)) { + layer.ple_key = create_tensor(tn(LLM_TENSOR_PLE_KEY, "weight", il), { n_embd, hc_dim }, 0); + layer.ple_value = create_tensor(tn(LLM_TENSOR_PLE_VALUE, "weight", il), { n_embd, n_embd }, 0); + layer.ple_norm_key = create_tensor(tn(LLM_TENSOR_PLE_NORM_KEY, "weight", il), { hc_dim }, 0); + layer.ple_norm_query = create_tensor(tn(LLM_TENSOR_PLE_NORM_QUERY, "weight", il), { hc_dim }, 0); + layer.ple_norm_conv = create_tensor(tn(LLM_TENSOR_PLE_NORM_CONV, "weight", il), { hc_dim }, 0); + layer.ple_conv1d = create_tensor(tn(LLM_TENSOR_PLE_CONV1D, "weight", il), { hparams.ple_conv_kernel, hc_dim }, 0); + } + + layer.ffn_gate_inp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP, "weight", il), { n_embd, n_expert }, 0); + layer.ffn_down_exps = create_tensor(tn(LLM_TENSOR_FFN_DOWN_EXPS, "weight", il), { n_ff_exp, n_embd, n_expert }, 0); + create_tensor_gate_up_exps(layer, il, n_embd, n_ff_exp, n_expert, 0); + + layer.ffn_gate_inp_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_INP_SHEXP, "weight", il), { n_embd }, 0); + layer.ffn_gate_shexp = create_tensor(tn(LLM_TENSOR_FFN_GATE_SHEXP, "weight", il), { n_embd, n_ff_shexp }, 0); + layer.ffn_up_shexp = create_tensor(tn(LLM_TENSOR_FFN_UP_SHEXP, "weight", il), { n_embd, n_ff_shexp }, 0); + layer.ffn_down_shexp = create_tensor(tn(LLM_TENSOR_FFN_DOWN_SHEXP, "weight", il), { n_ff_shexp, n_embd }, 0); + } +} + +std::unique_ptr llama_model_qwen4exp::build_arch_graph(const llm_graph_params & params) const { + return std::make_unique(*this, params); +} + +// Hyper-connections keep hc parallel residual streams [n_embd, hc, T] in place of layer norms. +// Returns the mixed [n_embd, T] stream; `inject` gets the [hc, T] scatter weights. +ggml_tensor * llama_model_qwen4exp::graph::build_hc_mix( + ggml_tensor * x, + ggml_tensor * w_norm, + ggml_tensor * w_down, + ggml_tensor * w_up, + ggml_tensor * w_inject, + ggml_tensor ** inject, + int il) { + const int64_t hc = hparams.dsv4_hc_mult; + const int64_t hc_dim = hc * n_embd; + const int64_t nt = x->ne[2]; + + // grouped RMSNorm: reduce over one stream, then scale all streams with the [hc_dim] gamma + // the converter folded each gamma to (1 + w) + ggml_tensor * xn = ggml_rms_norm(ctx0, x, hparams.f_norm_rms_eps); + xn = ggml_reshape_2d(ctx0, xn, hc_dim, nt); + xn = ggml_mul(ctx0, xn, w_norm); + cb(xn, "hc_norm", il); + + ggml_tensor * lo = build_lora_mm(w_down, xn); + lo = ggml_silu(ctx0, ggml_scale(ctx0, lo, 1.0f / (float) hc)); + ggml_tensor * gate = ggml_sigmoid(ctx0, build_lora_mm(w_up, lo)); + cb(gate, "hc_gate", il); + + ggml_tensor * gated = ggml_mul(ctx0, xn, gate); + gated = ggml_reshape_3d(ctx0, gated, n_embd, hc, nt); + + // collapse the streams by their mean + ggml_tensor * mixed = ggml_view_2d(ctx0, gated, n_embd, nt, + ggml_row_size(gated->type, n_embd) * hc, 0); + mixed = ggml_cont(ctx0, mixed); + for (int64_t c = 1; c < hc; ++c) { + ggml_tensor * s = ggml_view_2d(ctx0, gated, n_embd, nt, + ggml_row_size(gated->type, n_embd) * hc, + ggml_row_size(gated->type, n_embd) * c); + mixed = ggml_add(ctx0, mixed, s); + } + mixed = ggml_scale(ctx0, mixed, 1.0f / (float) hc); + cb(mixed, "hc_mixed", il); + + if (inject) { + *inject = build_lora_mm(w_inject, xn); + cb(*inject, "hc_inject", il); + } + + return mixed; +} + +ggml_tensor * llama_model_qwen4exp::graph::build_hc_combine( + ggml_tensor * residual, + ggml_tensor * block_out, + ggml_tensor * inject, + int il) { + const int64_t hc = hparams.dsv4_hc_mult; + const int64_t nt = residual->ne[2]; + + // 2*sigmoid centres the scatter weights on 1, so a zero injection is a plain residual add + ggml_tensor * w = ggml_sigmoid(ctx0, ggml_scale(ctx0, inject, 1.0f / (float) hc)); + w = ggml_scale(ctx0, w, 2.0f); + w = ggml_reshape_3d(ctx0, w, 1, hc, nt); + + ggml_tensor * b = ggml_reshape_3d(ctx0, block_out, n_embd, 1, nt); + b = ggml_repeat_4d(ctx0, b, n_embd, hc, nt, 1); + + ggml_tensor * cur = ggml_add(ctx0, residual, ggml_mul(ctx0, b, w)); + cb(cur, "hc_combine", il); + + return cur; +} + +llama_model_qwen4exp::graph::graph(const llama_model & model, const llm_graph_params & params) : + llm_build_delta_net_base(params), model(model) { + const int64_t hc = hparams.dsv4_hc_mult; + + GGML_ASSERT(hparams.n_embd_head_v() == hparams.n_embd_head_k()); + + int sections[4]; + std::copy(std::begin(hparams.rope_sections), std::begin(hparams.rope_sections) + 4, sections); + + ggml_tensor * inpL = build_inp_embd(model.tok_embd); + cb(inpL, "model.input_embed", -1); + + auto * inp = build_inp_mem_hybrid(); + + // qwen4exp always builds llama_memory_hybrid_idx, so this downcast is safe + // the indexer cache inside it is absent when the GGUF has no indexer tensors + const auto * mctx_hyb = static_cast(inp->mctx); + + const llama_kv_cache_context * mctx_idx = mctx_hyb->get_idx(); + if (mctx_idx) { + GGML_ASSERT(mctx_idx->get_n_kv() == inp->mctx->get_attn()->get_n_kv() && + "the indexer cache must track the attention cache cell for cell"); + } + + ggml_tensor * inp_pos = build_inp_pos(); + ggml_tensor * inp_out_ids = build_inp_out_ids(); + + // the wide residual starts as hc identical copies of the embedding + ggml_tensor * res_hc = ggml_repeat_4d(ctx0, + ggml_reshape_3d(ctx0, inpL, n_embd, 1, n_tokens), + n_embd, hc, n_tokens, 1); + cb(res_hc, "hc_init", -1); + + for (int il = 0; il < n_layer; ++il) { + res->t_layer_inp[il] = res_hc; + + if (hparams.is_ple(il)) { + res_hc = build_ple(inp->get_recr(), mctx_hyb, res_hc, il); + } + + ggml_tensor * inject = nullptr; + ggml_tensor * cur = build_hc_mix(res_hc, + model.layers[il].hc_attn_norm, + model.layers[il].hc_attn_down, + model.layers[il].hc_attn_up, + model.layers[il].hc_attn_inject, + &inject, il); + + ggml_build_forward_expand(gf, cur); + + if (hparams.is_recr(il)) { + cur = build_layer_attn_linear(inp->get_recr(), cur, il); + } else { + cur = build_layer_attn(inp->get_attn(), mctx_hyb, cur, inp_pos, sections, il); + } + + if (il == n_layer - 1 && inp_out_ids) { + // everything below is per token, so drop the rows that produce no output + cur = ggml_get_rows(ctx0, cur, inp_out_ids); + inject = ggml_get_rows(ctx0, inject, inp_out_ids); + + res_hc = ggml_reshape_2d(ctx0, res_hc, n_embd*hc, res_hc->ne[2]); + res_hc = ggml_get_rows(ctx0, res_hc, inp_out_ids); + res_hc = ggml_reshape_3d(ctx0, res_hc, n_embd, hc, res_hc->ne[1]); + } + + res_hc = build_hc_combine(res_hc, cur, inject, il); + + cur = build_hc_mix(res_hc, + model.layers[il].hc_ffn_norm, + model.layers[il].hc_ffn_down, + model.layers[il].hc_ffn_up, + model.layers[il].hc_ffn_inject, + &inject, il); + + cur = build_layer_ffn(cur, il); + cb(cur, "ffn_out", il); + + res_hc = build_hc_combine(res_hc, cur, inject, il); + + // "l_last" is the layer output name that build_cvec and imatrix look for + cb(res_hc, "l_last", il); + } + + // the final mixer is the output norm: there is no separate one + ggml_tensor * cur = build_hc_mix(res_hc, + model.hc_head_norm, model.hc_head_down, model.hc_head_up, + nullptr, nullptr, -1); + + cb(cur, "result_norm", -1); + res->t_embd = cur; + + cur = build_lora_mm(model.output, cur, model.output_s); + cb(cur, "result_output", -1); + res->t_logits = cur; + + ggml_build_forward_expand(gf, cur); +} + +std::pair llama_model_qwen4exp::graph::build_qkvz( + ggml_tensor * input, + int il) { + const int64_t n_seqs = ubatch.n_seqs; + const int64_t n_seq_tokens = ubatch.n_seq_tokens; + + ggml_tensor * qkv_mixed = build_lora_mm(model.layers[il].wqkv, input, model.layers[il].wqkv_s); + qkv_mixed = ggml_reshape_3d(ctx0, qkv_mixed, qkv_mixed->ne[0], n_seq_tokens, n_seqs); + cb(qkv_mixed, "linear_attn_qkv_mixed", il); + + ggml_tensor * z = build_lora_mm(model.layers[il].wqkv_gate, input, model.layers[il].wqkv_gate_s); + cb(z, "z", il); + + return { qkv_mixed, z }; +} + +ggml_tensor * llama_model_qwen4exp::graph::build_norm_gated( + ggml_tensor * input, + ggml_tensor * weights, + ggml_tensor * gate, + int layer) { + // the one numerical difference from Qwen3.5's GDN: sigmoid output gate, not silu + ggml_tensor * normalized = build_norm(input, weights, nullptr, LLM_NORM_RMS, layer); + ggml_tensor * gated = ggml_sigmoid(ctx0, gate); + + return ggml_mul(ctx0, normalized, gated); +} + +// QSA attends to a budget of whole blocks of compress_ratio tokens, plus the incomplete tail +// one mean-pooled indexer key scores each block; set_input resolves the cache layout +class llama_model_qwen4exp::llm_graph_input_qsa : public llm_graph_input_i { +public: + llm_graph_input_qsa(const llama_memory_hybrid_idx_context * mctx, uint32_t ratio, bool blk_bias) : + mctx(mctx), ratio(ratio), blk_bias(blk_bias) {} + virtual ~llm_graph_input_qsa() = default; + + void set_input(const llama_ubatch * ubatch) override { + mctx->get_idx()->set_input_k_idxs(k_idxs, ubatch); + mctx->set_input_qsa(cell_blk, blk_cells, blk_pos, bias, ubatch, ratio, blk_bias); + } + + bool can_reuse(const llm_graph_params & params) override { + mctx = static_cast(params.mctx); + + const auto * idx = mctx->get_idx(); + if (idx == nullptr) { + return false; + } + + const int64_t n_kv = idx->get_n_kv(); + const int64_t n_stream = mctx->get_n_stream(); + const int64_t n_blocks = (n_kv + ratio - 1)/ratio; + + bool res = true; + + res &= params.ubatch.n_tokens % n_stream == 0; + + res &= k_idxs->ne[0] == params.ubatch.n_tokens; + res &= cell_blk->ne[0] == n_kv; + res &= cell_blk->ne[1] == n_stream; + res &= blk_cells->ne[0] == (int64_t) ratio*n_blocks; + res &= blk_pos->ne[0] == 4*n_blocks*n_stream; + res &= bias->ne[0] == (blk_bias ? n_blocks : n_kv); + res &= bias->ne[1] == params.ubatch.n_tokens/n_stream; + + return res; + } + + // per stream: a cell index names a different token in each stream + ggml_tensor * k_idxs = nullptr; // I32 [n_tokens] + ggml_tensor * cell_blk = nullptr; // I32 [n_kv, n_stream] + ggml_tensor * blk_cells = nullptr; // I32 [ratio*n_blocks, n_stream] + ggml_tensor * blk_pos = nullptr; // I32 [4*n_blocks*n_stream] + ggml_tensor * bias = nullptr; // F32 [n_blocks or n_kv, n_tokens/n_stream, n_stream] + + const llama_memory_hybrid_idx_context * mctx; + const uint32_t ratio; + + // the per-cell half of the bias is the attention mask, so only the per-block half is uploaded + const bool blk_bias; +}; + +ggml_tensor * llama_model_qwen4exp::graph::build_qsa_top_k( + const llama_memory_hybrid_idx_context * mctx_hyb, + ggml_tensor * cur, + ggml_tensor * inp_pos, + ggml_tensor * kq_mask, + int * sections, + int il) { + const llama_kv_cache_context * mctx_idx = mctx_hyb->get_idx(); + + const int64_t idx_dim = hparams.indexer_head_size; + const int64_t n_idx_h = hparams.indexer_n_head; + const int64_t r = hparams.dsv4_compress_ratios[il]; + const int64_t n_kv = mctx_idx->get_n_kv(); + + GGML_ASSERT(r > 0); + + const int64_t n_blocks = (n_kv + r - 1)/r; + + // build_attn_qsa and the KQ mask need the tokens to divide evenly across the streams + const int64_t n_stream = mctx_hyb->get_n_stream(); + GGML_ASSERT(n_tokens % n_stream == 0); + const int64_t n_tps = n_tokens/n_stream; + + // only the "which block is visible" half of the bias varies per block + // the rest is the visible/not test the attention mask already carries, so upload the per-block half only: 1/ratio of the cells + // alibi writes distances instead of a mask and non-causal keeps future cells, so both opt out + // the mask also holds an mrope rule for the query's own position, but only 2d image positions can differ there + const bool blk_bias = kq_mask != nullptr && + kq_mask->ne[0] == n_kv && kq_mask->ne[1] == n_tps && kq_mask->ne[3] == n_stream && + cparams.causal_attn && !hparams.use_alibi; + + // nothing above depends on the layer, so the layers sharing a ratio share one input set + llm_graph_input_qsa * inp = nullptr; + + const auto it = qsa_inps.find((uint32_t) r); + if (it != qsa_inps.end()) { + inp = it->second; + } else { + auto qsa = std::make_unique(mctx_hyb, (uint32_t) r, blk_bias); + + qsa->k_idxs = mctx_idx->build_input_k_idxs(ctx0, ubatch); + qsa->cell_blk = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, n_kv, n_stream); + qsa->blk_cells = ggml_new_tensor_2d(ctx0, GGML_TYPE_I32, r*n_blocks, n_stream); + qsa->blk_pos = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, 4*n_blocks*n_stream); + qsa->bias = ggml_new_tensor_3d(ctx0, GGML_TYPE_F32, blk_bias ? n_blocks : n_kv, n_tps, n_stream); + + ggml_set_input(qsa->cell_blk); + ggml_set_input(qsa->blk_cells); + ggml_set_input(qsa->blk_pos); + ggml_set_input(qsa->bias); + + inp = qsa.get(); + res->add_input(std::move(qsa)); + qsa_inps.emplace((uint32_t) r, inp); + } + + // cached indexer keys are raw: pooling precedes norm and rotation, so apply neither + ggml_tensor * k_raw = build_lora_mm(model.layers[il].index_k_proj, cur); + k_raw = ggml_reshape_3d(ctx0, k_raw, idx_dim, 1, n_tokens); + cb(k_raw, "indexer_k_raw", il); + + ggml_build_forward_expand(gf, mctx_idx->cpy_k(ctx0, k_raw, inp->k_idxs, il)); + + // one key head, so rows are contiguous. get_k gives [idx_dim, n_head_kv, n_kv, n_stream]. + ggml_tensor * k_all = mctx_idx->get_k(ctx0, il); + k_all = ggml_view_3d(ctx0, k_all, idx_dim, n_kv, n_stream, k_all->nb[2], k_all->nb[3], 0); + + // gathers per stream: blk_cells row s indexes stream s's own cells + ggml_tensor * members = ggml_get_rows(ctx0, k_all, inp->blk_cells); + members = ggml_reshape_4d(ctx0, members, idx_dim, r, n_blocks, n_stream); + + // mean over the block members; r is small, so summing slices beats a transpose plus sum_rows + ggml_tensor * pooled = nullptr; + for (int64_t i = 0; i < r; ++i) { + ggml_tensor * slice = ggml_cont(ctx0, + ggml_view_3d(ctx0, members, idx_dim, n_blocks, n_stream, + members->nb[2], members->nb[3], i*members->nb[1])); + pooled = pooled ? ggml_add(ctx0, pooled, slice) : slice; + } + pooled = ggml_scale(ctx0, pooled, 1.0f/(float) r); + cb(pooled, "indexer_k_pooled", il); + + // rope wants [n_dims, n_head, n_tokens]: lay every stream's blocks flat, split after. + pooled = ggml_reshape_3d(ctx0, pooled, idx_dim, 1, n_blocks*n_stream); + pooled = build_norm(pooled, model.layers[il].index_k_norm, nullptr, LLM_NORM_RMS, il); + pooled = ggml_rope_multi(ctx0, pooled, inp->blk_pos, nullptr, + n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + pooled = ggml_reshape_3d(ctx0, pooled, idx_dim, n_blocks, n_stream); + cb(pooled, "indexer_k", il); + + ggml_tensor * q = build_lora_mm(model.layers[il].index_q_proj, cur); + q = ggml_reshape_3d(ctx0, q, idx_dim, n_idx_h, n_tokens); + q = build_norm(q, model.layers[il].index_q_norm, nullptr, LLM_NORM_RMS, il); + q = ggml_rope_multi(ctx0, q, inp_pos, nullptr, + n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow); + cb(q, "indexer_q", il); + + // rectify each head dot product before the sum, as in the DeepSeek lightning indexer + // mul_mat matches ne[2], so the queries of stream s only meet the blocks of stream s + ggml_tensor * score = ggml_mul_mat(ctx0, pooled, + ggml_reshape_3d(ctx0, ggml_cont(ctx0, q), idx_dim, n_idx_h*n_tps, n_stream)); + score = ggml_reshape_4d(ctx0, score, n_blocks, n_idx_h, n_tps, n_stream); + score = ggml_relu(ctx0, score); + score = ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3)); + score = ggml_sum_rows(ctx0, score); + score = ggml_reshape_3d(ctx0, score, n_blocks, n_tps, n_stream); + cb(score, "indexer_score", il); + + // one value per block, so it is cheaper to bias here than after the cells are expanded + if (blk_bias) { + score = ggml_add(ctx0, score, inp->bias); + } + + // every token of a block gets the block score; the budget is whole blocks, so top-k cuts on a block boundary + ggml_tensor * expanded = ggml_get_rows(ctx0, + ggml_cont(ctx0, ggml_permute(ctx0, score, 1, 0, 2, 3)), inp->cell_blk); + expanded = ggml_cont(ctx0, ggml_permute(ctx0, expanded, 1, 0, 2, 3)); + + if (blk_bias) { + // flash attention keeps the mask in f16; the scores are f32 + ggml_tensor * mask = kq_mask->type == GGML_TYPE_F32 ? kq_mask : ggml_cast(ctx0, kq_mask, GGML_TYPE_F32); + expanded = ggml_add(ctx0, expanded, ggml_reshape_3d(ctx0, mask, n_kv, n_tps, n_stream)); + } else { + expanded = ggml_add(ctx0, expanded, inp->bias); + } + cb(expanded, "indexer_score_tokens", il); + + // the reference returns indexer_top_k + compress_ratio - 1: whole blocks plus the tail + const int64_t width = std::min(n_kv, (int64_t) hparams.indexer_top_k + r - 1); + + ggml_tensor * top_k = ggml_cont(ctx0, ggml_top_k(ctx0, expanded, width)); + + // build_attn_qsa reads [n_top_k, n_batch, 1, n_stream], matching the KQ mask. + top_k = ggml_reshape_4d(ctx0, top_k, width, n_tps, 1, n_stream); + cb(top_k, "indexer_top_k", il); + + return top_k; +} + +// Dense GQA self-attention restricted to the cells that top_k names. +// The mask build below copies the MLA sparse path in llm_graph_context::build_attn. +ggml_tensor * llama_model_qwen4exp::graph::build_attn_qsa( + llm_graph_input_attn_kv * inp, + ggml_tensor * q_cur, + ggml_tensor * k_cur, + ggml_tensor * v_cur, + ggml_tensor * top_k, + float kq_scale, + int il) { + // rotate q/k/v before they reach a quantized cache, as the dense path does. the indexer + // has already scored with its own query in build_qsa_top_k, so top_k is unaffected. + if (inp->self_k_rot) { + q_cur = llama_mul_mat_hadamard(ctx0, q_cur, inp->self_k_rot); + k_cur = llama_mul_mat_hadamard(ctx0, k_cur, inp->self_k_rot); + } + + if (inp->self_v_rot) { + v_cur = llama_mul_mat_hadamard(ctx0, v_cur, inp->self_v_rot); + } + + // these nodes are added to the graph together so that they are not reordered + // by doing so, the number of splits in the graph is reduced + // expand k later to enable rope fusion which directly writes into k-v cache + ggml_build_forward_expand(gf, q_cur); + ggml_build_forward_expand(gf, v_cur); + ggml_build_forward_expand(gf, k_cur); + + const auto * mctx_cur = inp->mctx; + + // store to KV cache + { + const auto & k_idxs = inp->get_k_idxs(); + const auto & v_idxs = inp->get_v_idxs(); + + ggml_build_forward_expand(gf, mctx_cur->cpy_k(ctx0, k_cur, k_idxs, il)); + ggml_build_forward_expand(gf, mctx_cur->cpy_v(ctx0, v_cur, v_idxs, il)); + } + + ggml_tensor * kq_mask = inp->get_kq_mask(); + + // prepare new kq mask - starts filled with -INFINITY + ggml_tensor * kq_mask_all = ggml_fill(ctx0, kq_mask, -INFINITY); + + // reshape KQ mask into tensor with rows of size 1: + // [n_kv, n_batch, 1, n_stream] -> [1, n_kv, n_batch, n_stream] + kq_mask_all = ggml_view_4d(ctx0, kq_mask_all, 1, kq_mask_all->ne[0], kq_mask_all->ne[1], kq_mask_all->ne[3], kq_mask_all->nb[0], kq_mask_all->nb[1], kq_mask_all->nb[2], 0); + + // reshape top_k indices: [n_top_k, n_batch, 1, n_stream] -> [n_top_k, n_batch, n_stream, 1] + ggml_tensor * top_k_3d = ggml_view_4d(ctx0, top_k, top_k->ne[0], top_k->ne[1], top_k->ne[3], 1, top_k->nb[1], top_k->nb[2], top_k->ne[3]*top_k->nb[3], 0); + + // prepare zero-filled tensor with rows of size 1: [1, n_top_k, n_batch, n_stream] + // this will be our source of zero values for unmasking top k mask elements + ggml_tensor * zeros = ggml_new_tensor_4d(ctx0, GGML_TYPE_F32, 1, top_k_3d->ne[0], top_k_3d->ne[1], top_k_3d->ne[2]); + zeros = ggml_fill(ctx0, zeros, 0.0f); + + // modify KQ mask by unmasking elements that are in top_k indices + // ggml_set_rows([1, n_kv, n_batch, n_stream], [1, n_top_k, n_batch, n_stream], [n_top_k, n_batch, n_stream, 1]) + ggml_tensor * kq_mask_top_k = ggml_set_rows(ctx0, kq_mask_all, zeros, top_k_3d); + + // reshape to restore the original shape of KQ mask: + // [1, n_kv, n_batch, n_stream] -> [n_kv, n_batch, 1, n_stream] + kq_mask_top_k = ggml_view_4d(ctx0, kq_mask_top_k, kq_mask_top_k->ne[1], kq_mask_top_k->ne[2], 1, kq_mask_top_k->ne[3], kq_mask_top_k->nb[2], kq_mask_top_k->nb[3], kq_mask_top_k->nb[3], 0); + + // combine with the original kq mask + kq_mask_top_k = ggml_add(ctx0, kq_mask_top_k, kq_mask); + + ggml_tensor * q = q_cur; + ggml_tensor * k = mctx_cur->get_k(ctx0, il); + ggml_tensor * v = mctx_cur->get_v(ctx0, il); + + ggml_tensor * cur = build_attn_mha(q, k, v, nullptr, kq_mask_top_k, nullptr, nullptr, kq_scale, il); + cb(cur, "kqv_out", il); + + // the rotation is its own inverse, so undo it on the value side of the output + if (inp->self_v_rot) { + cur = llama_mul_mat_hadamard(ctx0, cur, inp->self_v_rot); + } + + return cur; +} + +ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn( + llm_graph_input_attn_kv * inp, + const llama_memory_hybrid_idx_context * mctx_hyb, + ggml_tensor * cur, + ggml_tensor * inp_pos, + int * sections, + int il) { + const int64_t n_embd_head = hparams.n_embd_head_v(); + GGML_ASSERT(n_embd_head == hparams.n_embd_head_k()); + + // indexer reads the same block input as q/k/v; no cache or no ratio means dense + const bool qsa = mctx_hyb->get_idx() != nullptr && hparams.dsv4_compress_ratios[il] > 0; + + ggml_tensor * top_k = qsa ? build_qsa_top_k(mctx_hyb, cur, inp_pos, inp->get_kq_mask(), sections, il) : nullptr; + + // Qwen3Next uses a single Q projection that outputs query + gate + ggml_tensor * Qcur_full = build_lora_mm(model.layers[il].wq, cur, model.layers[il].wq_s); // [ (n_embd_head * 2) * n_head, n_tokens ] + cb(Qcur_full, "Qcur_full", il); + + ggml_tensor * Qcur = ggml_view_3d(ctx0, Qcur_full, n_embd_head, n_head, n_tokens, + ggml_element_size(Qcur_full) * n_embd_head * 2, + ggml_element_size(Qcur_full) * n_embd_head * 2 * n_head, 0); + cb(Qcur, "Qcur_reshaped", il); + + Qcur = build_norm(Qcur, model.layers[il].attn_q_norm, nullptr, LLM_NORM_RMS, il); + cb(Qcur, "Qcur_normed", il); + + ggml_tensor * Kcur = build_lora_mm(model.layers[il].wk, cur, model.layers[il].wk_s); + cb(Kcur, "Kcur", il); + + ggml_tensor * Vcur = build_lora_mm(model.layers[il].wv, cur, model.layers[il].wv_s); + cb(Vcur, "Vcur", il); + + Kcur = ggml_reshape_3d(ctx0, Kcur, n_embd_head, n_head_kv, n_tokens); + Kcur = build_norm(Kcur, model.layers[il].attn_k_norm, nullptr, LLM_NORM_RMS, il); + cb(Kcur, "Kcur_normed", il); + + ggml_tensor * gate = ggml_view_3d(ctx0, Qcur_full, n_embd_head, n_head, n_tokens, + ggml_element_size(Qcur_full) * n_embd_head * 2, + ggml_element_size(Qcur_full) * n_embd_head * 2 * n_head, + ggml_element_size(Qcur_full) * n_embd_head); + gate = ggml_cont_2d(ctx0, gate, n_embd_head * n_head, n_tokens); + cb(gate, "gate_reshaped", il); + + Vcur = ggml_reshape_3d(ctx0, Vcur, n_embd_head, n_head_kv, n_tokens); + + // Apply IMRoPE + Qcur = ggml_rope_multi( + ctx0, Qcur, inp_pos, nullptr, + n_rot, sections, rope_type, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow + ); + + Kcur = ggml_rope_multi( + ctx0, Kcur, inp_pos, nullptr, + n_rot, sections, 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; + + if (top_k) { + cur = build_attn_qsa(inp, Qcur, Kcur, Vcur, top_k, kq_scale, il); + } else { + cur = build_attn(inp, + nullptr, nullptr, nullptr, + Qcur, Kcur, Vcur, nullptr, nullptr, nullptr, kq_scale, il); + } + cb(cur, "attn_pregate", il); + + ggml_tensor * gate_sigmoid = ggml_sigmoid(ctx0, gate); + cb(gate_sigmoid, "gate_sigmoid", il); + + cur = ggml_mul(ctx0, cur, gate_sigmoid); + cb(cur, "attn_gated", il); + + cur = build_lora_mm(model.layers[il].wo, cur, model.layers[il].wo_s); + cb(cur, "attn_output", il); + + return cur; +} + +ggml_tensor * llama_model_qwen4exp::graph::build_layer_attn_linear( + llm_graph_input_rs * inp, + ggml_tensor * cur, + int il) { + const auto * mctx_cur = inp->mctx; + + const int64_t d_inner = hparams.ssm_d_inner; + const int64_t n_seqs = ubatch.n_seqs; + const int64_t head_k_dim = hparams.ssm_d_state; + const int64_t num_k_heads = hparams.ssm_n_group; + const int64_t num_v_heads = hparams.ssm_dt_rank; + const int64_t head_v_dim = hparams.ssm_d_state; + const int64_t n_seq_tokens = ubatch.n_seq_tokens; + + GGML_ASSERT(n_seqs != 0); + GGML_ASSERT(ubatch.equal_seqs()); + GGML_ASSERT(ubatch.n_tokens == n_seq_tokens * n_seqs); + GGML_ASSERT(head_v_dim * num_v_heads == d_inner); + + auto qkvz = build_qkvz(cur, il); + ggml_tensor * qkv_mixed = qkvz.first; + ggml_tensor * z = qkvz.second; + + ggml_tensor * beta = build_lora_mm(model.layers[il].ssm_beta, cur, model.layers[il].ssm_beta_s); + beta = ggml_reshape_4d(ctx0, beta, 1, num_v_heads, n_seq_tokens, n_seqs); + cb(beta, "beta", il); + + beta = ggml_sigmoid(ctx0, beta); + cb(beta, "beta_sigmoid", il); + + ggml_tensor * alpha = build_lora_mm(model.layers[il].ssm_alpha, cur, model.layers[il].ssm_alpha_s); + alpha = ggml_reshape_3d(ctx0, alpha, num_v_heads, n_seq_tokens, n_seqs); + cb(alpha, "alpha", il); + + ggml_tensor * alpha_biased = ggml_add(ctx0, alpha, model.layers[il].ssm_dt); + ggml_tensor * alpha_softplus = ggml_softplus(ctx0, alpha_biased); + cb(alpha_softplus, "a_softplus", il); + + ggml_tensor * gate = ggml_mul(ctx0, alpha_softplus, model.layers[il].ssm_a); // -A_log.exp() * softplus + cb(gate, "gate", il); + + gate = ggml_reshape_4d(ctx0, gate, 1, num_v_heads, n_seq_tokens, n_seqs); + + ggml_tensor * conv_states_all = mctx_cur->get_r_l(il); + ggml_tensor * ssm_states_all = mctx_cur->get_s_l(il); + + ggml_tensor * conv_kernel = model.layers[il].ssm_conv1d; + const int64_t conv_kernel_size = conv_kernel->ne[0]; + + // the channels must match how load_arch_tensors sizes wqkv, not ssm_d_inner + const int64_t conv_channels = head_k_dim * num_k_heads * 2 + head_v_dim * num_v_heads; + + ggml_tensor * conv_input = build_conv_state_at(inp, conv_states_all, qkv_mixed, + conv_kernel_size - 1, conv_channels, il); + + ggml_tensor * state = build_rs(inp, ssm_states_all, hparams.n_embd_s(), n_seqs); + state = ggml_reshape_4d(ctx0, state, head_v_dim, head_v_dim, num_v_heads, n_seqs); + cb(state, "state_predelta", il); + + ggml_tensor * conv_output_proper = ggml_ssm_conv(ctx0, conv_input, conv_kernel); + cb(conv_output_proper, "conv_output_raw", il); + + ggml_tensor * conv_output_silu = ggml_silu(ctx0, conv_output_proper); + cb(conv_output_silu, "conv_output_silu", il); + + ggml_tensor * conv_qkv_mix = conv_output_silu; + + int64_t nb1_qkv = ggml_row_size(conv_qkv_mix->type, conv_channels); + + // Extract the convolved Q, K, V from conv_output + ggml_tensor * q_conv = ggml_view_4d(ctx0, conv_qkv_mix, head_k_dim, num_k_heads, n_seq_tokens, n_seqs, + ggml_row_size(conv_qkv_mix->type, head_k_dim), + nb1_qkv, + nb1_qkv * n_seq_tokens, + 0); + + ggml_tensor * k_conv = ggml_view_4d(ctx0, conv_qkv_mix, head_k_dim, num_k_heads, n_seq_tokens, n_seqs, + ggml_row_size(conv_qkv_mix->type, head_k_dim), + nb1_qkv, + nb1_qkv * n_seq_tokens, + head_k_dim * num_k_heads * ggml_element_size(conv_qkv_mix)); + + ggml_tensor * v_conv = ggml_view_4d(ctx0, conv_qkv_mix, head_v_dim, num_v_heads, n_seq_tokens, n_seqs, + ggml_row_size(conv_qkv_mix->type, head_v_dim), + nb1_qkv, + nb1_qkv * n_seq_tokens, + ggml_row_size(conv_qkv_mix->type, 2 * head_k_dim * num_k_heads)); + + cb(q_conv, "q_conv", il); + cb(k_conv, "k_conv", il); + cb(v_conv, "v_conv", il); + + const float eps_norm = hparams.f_norm_rms_eps; + + q_conv = ggml_l2_norm(ctx0, q_conv, eps_norm); + k_conv = ggml_l2_norm(ctx0, k_conv, eps_norm); + + // repeat to match shapes when head keys != value keys; unneeded with the fused GDN + if (num_k_heads != num_v_heads && (!cparams.fused_gdn_ar || !cparams.fused_gdn_ch)) { + GGML_ASSERT(num_v_heads % num_k_heads == 0); + q_conv = ggml_repeat_4d(ctx0, q_conv, head_k_dim, num_v_heads, n_seq_tokens, n_seqs); + k_conv = ggml_repeat_4d(ctx0, k_conv, head_k_dim, num_v_heads, n_seq_tokens, n_seqs); + } + + cb(q_conv, "q_conv_predelta", il); + cb(k_conv, "k_conv_predelta", il); + cb(v_conv, "v_conv_predelta", il); + + ggml_tensor * output = build_recurrent_attn(inp, ssm_states_all, q_conv, k_conv, v_conv, gate, beta, state, il); + + ggml_tensor * z_2d = ggml_reshape_4d(ctx0, z, head_v_dim, num_v_heads, n_seq_tokens, n_seqs); + + // gated normalization, as self.norm(core_attn_out, z) in the reference + ggml_tensor * attn_out_norm = build_norm_gated(output, model.layers[il].ssm_norm, z_2d, il); + + ggml_tensor * final_output = ggml_reshape_3d(ctx0, attn_out_norm, head_v_dim * num_v_heads, n_seq_tokens, n_seqs); + cb(final_output, "final_output", il); + + cur = build_lora_mm(model.layers[il].ssm_out, final_output, model.layers[il].ssm_out_s); + cb(cur, "linear_attn_out", il); + + cur = ggml_reshape_2d(ctx0, cur, n_embd, n_seq_tokens * n_seqs); + + return cur; +} + +ggml_tensor * llama_model_qwen4exp::graph::build_layer_ffn(ggml_tensor * cur, const int il) { + GGML_ASSERT(model.layers[il].ffn_gate_inp != nullptr); + + 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, + model.layers[il].ffn_up_exps_s, + model.layers[il].ffn_gate_exps_s, + model.layers[il].ffn_down_exps_s); + cb(moe_out, "ffn_moe_out", il); + + // shared experts, as in the Qwen3Next reference + if (model.layers[il].ffn_up_shexp != nullptr) { + ggml_tensor * ffn_shexp = + build_ffn(cur, + model.layers[il].ffn_up_shexp, NULL, model.layers[il].ffn_up_shexp_s, + model.layers[il].ffn_gate_shexp, NULL, model.layers[il].ffn_gate_shexp_s, + model.layers[il].ffn_down_shexp, NULL, model.layers[il].ffn_down_shexp_s, + NULL, + LLM_FFN_SILU, LLM_FFN_PAR, il); + cb(ffn_shexp, "ffn_shexp", il); + + // shared expert has its own sigmoided gate (ffn_gate_inp_shexp, one value per token) + ggml_tensor * shared_gate = build_lora_mm(model.layers[il].ffn_gate_inp_shexp, cur); + cb(shared_gate, "shared_expert_gate", il); + + shared_gate = ggml_sigmoid(ctx0, shared_gate); + cb(shared_gate, "shared_expert_gate_sigmoid", il); + + ffn_shexp = ggml_mul(ctx0, ffn_shexp, shared_gate); + cb(ffn_shexp, "ffn_shexp_gated", il); + + cur = ggml_add(ctx0, moe_out, ffn_shexp); + cb(cur, "ffn_out", il); + } else { + cur = moe_out; + } + + return cur; +} + +// PLE n-gram hash embedding: each token gathers ple_n_heads rows of a shared table. +// mixed_n = (t[p]*m[0]) ^ ... ^ (t[p-n+1]*m[n-1]); row = mixed_n % vocab[h] + offset[h] +// The hash runs host-side because ggml has no int64 and no xor. EOS resets the window. + +class llm_graph_input_ple : public llm_graph_input_i { +public: + llm_graph_input_ple(const llama_model_qwen4exp & pmodel, + const llama_kv_cache_context * mctx) : pmodel(pmodel), mctx(mctx) {} + virtual ~llm_graph_input_ple() = default; + + void set_input(const llama_ubatch * ubatch) override; + + bool can_reuse(const llm_graph_params & params) override { + mctx = static_cast(params.mctx)->get_attn(); + return rows->ne[0] == (int64_t) pmodel.hparams.ple_n_heads * params.ubatch.n_tokens; + } + + ggml_tensor * rows = nullptr; // I32 [ple_n_heads * n_tokens] + + const llama_model_qwen4exp & pmodel; + + // the predecessor tokens live in the attention KV cells (ext.tok) + const llama_kv_cache_context * mctx; + + // scratch, reused across set_input() calls + std::vector prev; +}; + +void llm_graph_input_ple::set_input(const llama_ubatch * ubatch) { + const auto & hp = pmodel.hparams; + + // an image arrives as an embd batch, so ubatch->token is null, but every position still needs a row for ggml_get_rows + // stand in the image token id that the reference hashes, or EOS if the file has no such key + // gemma3n and gemma4 do the same with a hardcoded row 0 of per_layer_token_embd. + const llama_token img_tok = hp.ple_image_token_id != 0 + ? (llama_token) hp.ple_image_token_id + : (llama_token) hp.ple_eos_token_id; + auto tok_of = [&](int64_t k) -> llama_token { + return ubatch->token ? ubatch->token[k] : img_tok; + }; + + const int64_t n_tokens = ubatch->n_tokens; + const int64_t n_gram = hp.ple_ngram_size; + const int64_t n_heads = hp.ple_n_heads; + const int64_t per_gram = hp.ple_heads_per_ngram; + const int64_t eos = hp.ple_eos_token_id; + const int64_t n_prev = n_gram - 1; + + std::vector idx(n_heads * n_tokens); + + GGML_ASSERT(mctx != nullptr); + + for (int64_t i = 0; i < n_tokens; ++i) { + // the preceding tokens would be ambiguous, see get_prev_tokens() + GGML_ASSERT(ubatch->n_seq_id[i] == 1 && "PLE n-gram embeddings do not support tokens shared by multiple sequences"); + } + + // predecessors come from the KV cells (ext.tok); apply_ubatch() already stored this ubatch, so its own tokens count too + mctx->get_prev_tokens(*ubatch, n_prev, prev); + + for (int64_t i = 0; i < n_tokens; ++i) { + // an EOS in the window resets everything at or before it + // a missing predecessor (before the sequence start, or no cached cell) reads as EOS + // the EOS of the token itself does not cut its own context, as in the reference + std::vector ctx(n_gram); + ctx[0] = tok_of(i); + bool cut = false; + for (int64_t s = 1; s < n_gram; ++s) { + // predecessor s positions back; prev[] is oldest-first, missing entries are LLAMA_TOKEN_NULL + const llama_token t = cut ? LLAMA_TOKEN_NULL : prev[i*n_prev + (n_prev - s)]; + cut = cut || t < 0 || t == eos; + ctx[s] = cut ? eos : t; + } + + for (int64_t n = 2; n <= n_gram; ++n) { + uint64_t mixed = (uint64_t) ctx[0] * hp.ple_layer_multipliers[0]; + for (int64_t j = 1; j < n; ++j) { + mixed ^= (uint64_t) ctx[j] * hp.ple_layer_multipliers[j]; + } + const int64_t base = (n - 2) * per_gram; + for (int64_t g = 0; g < per_gram; ++g) { + const int64_t h_i = base + g; + idx[i * n_heads + h_i] = + (int32_t) (mixed % hp.ple_head_vocab_sizes[h_i] + hp.ple_head_offsets[h_i]); + } + } + } + + ggml_backend_tensor_set(rows, idx.data(), 0, idx.size()*ggml_element_size(rows)); +} + +// Read a conv history out of its own recurrent row and write the new tail back. +// The shared build_conv_state cannot do this: qwen4exp has two such rows per layer. +ggml_tensor * llama_model_qwen4exp::graph::build_conv_state_at( + llm_graph_input_rs * inp, + ggml_tensor * conv_states_all, + ggml_tensor * x, + int64_t state_cols, + int64_t channels, + int il) { + const auto * mctx_cur = inp->mctx; + + const auto kv_head = mctx_cur->get_head(); + + const int64_t n_seqs = ubatch.n_seqs; + const int64_t row_total = conv_states_all->ne[0]; + + // the row is exactly this convolution's state, so the gather is reused as a whole + GGML_ASSERT(state_cols * channels == row_total); + + auto it = rs_rows.find(conv_states_all); + if (it == rs_rows.end()) { + it = rs_rows.emplace(conv_states_all, build_rs(inp, conv_states_all, row_total, n_seqs)).first; + } + ggml_tensor * rows = it->second; + + ggml_tensor * state = ggml_reshape_3d(ctx0, rows, state_cols, channels, n_seqs); + cb(state, "conv_state_at", il); + + ggml_tensor * conv_input = ggml_concat(ctx0, state, ggml_transpose(ctx0, x), 0); + + // keep the last state_cols columns for the next ubatch + const size_t row_size = ggml_row_size(conv_states_all->type, row_total); + + ggml_tensor * tail = ggml_view_3d(ctx0, conv_input, + state_cols, channels, n_seqs, + conv_input->nb[1], conv_input->nb[2], + ggml_row_size(conv_input->type, conv_input->ne[0] - state_cols)); + + ggml_tensor * dst = ggml_view_2d(ctx0, conv_states_all, + state_cols * channels, n_seqs, + conv_states_all->nb[1], + kv_head * row_size); + + ggml_build_forward_expand(gf, ggml_cpy(ctx0, ggml_cont(ctx0, tail), dst)); + + return conv_input; +} + +ggml_tensor * llama_model_qwen4exp::graph::build_ple( + llm_graph_input_rs * inp, + const llama_memory_hybrid_idx_context * mctx_hyb, + ggml_tensor * hidden, + int il) { + const int64_t hc = hparams.dsv4_hc_mult; + const int64_t hc_dim = hc * n_embd; + const int64_t n_heads = hparams.ple_n_heads; + + // the attention cells see every ubatch regardless of the layer types + auto ple_inp = std::make_unique( + static_cast(model), mctx_hyb->get_attn()); + + ple_inp->rows = ggml_new_tensor_1d(ctx0, GGML_TYPE_I32, n_heads * n_tokens); + ggml_set_input(ple_inp->rows); + ggml_tensor * rows = ple_inp->rows; + res->add_input(std::move(ple_inp)); + + // gather then flatten the heads: get_rows lays the head dimension out slowest, as the reference does + ggml_tensor * emb = ggml_get_rows(ctx0, model.per_layer_tok_embd, rows); + emb = ggml_reshape_2d(ctx0, emb, hparams.ple_head_dim * n_heads, n_tokens); + cb(emb, "ple_embd", il); + + ggml_tensor * key = build_lora_mm(model.layers[il].ple_key, emb); + ggml_tensor * value = build_lora_mm(model.layers[il].ple_value, emb); + + // both norms group over one hc stream, with a weight over the whole hc*n_embd layout + auto grouped_norm = [&](ggml_tensor * x, ggml_tensor * w) { + ggml_tensor * t = ggml_reshape_3d(ctx0, x, n_embd, hc, n_tokens); + t = ggml_rms_norm(ctx0, t, hparams.f_norm_rms_eps); + t = ggml_reshape_2d(ctx0, t, hc_dim, n_tokens); + t = ggml_mul(ctx0, t, w); + return ggml_reshape_3d(ctx0, t, n_embd, hc, n_tokens); + }; + + key = grouped_norm(key, model.layers[il].ple_norm_key); + ggml_tensor * query = grouped_norm(hidden, model.layers[il].ple_norm_query); + + // per-stream dot product, then a signed square root before the sigmoid + ggml_tensor * s = ggml_sum_rows(ctx0, ggml_mul(ctx0, key, query)); + s = ggml_scale(ctx0, s, 1.0f / sqrtf((float) n_embd)); + + ggml_tensor * mag = ggml_sqrt(ctx0, ggml_clamp(ctx0, ggml_abs(ctx0, s), 1e-6f, INFINITY)); + ggml_tensor * gate = ggml_sigmoid(ctx0, ggml_mul(ctx0, ggml_sgn(ctx0, s), mag)); + cb(gate, "ple_gate", il); + + // [n_embd, 1, T] value broadcast across the hc streams, scaled by the gate + ggml_tensor * v3 = ggml_reshape_3d(ctx0, value, n_embd, 1, n_tokens); + v3 = ggml_repeat_4d(ctx0, v3, n_embd, hc, n_tokens, 1); + + ggml_tensor * gated = ggml_mul(ctx0, v3, gate); + cb(gated, "ple_gated_value", il); + + ggml_tensor * normalized = grouped_norm( + ggml_reshape_2d(ctx0, gated, hc_dim, n_tokens), + model.layers[il].ple_norm_conv); + normalized = ggml_reshape_2d(ctx0, normalized, hc_dim, n_tokens); + + // depthwise causal conv, dilated by the n-gram size, as a sum of shifted copies + // ggml_conv_1d_dw is documented as unreliable: + // out[c, t] = sum_k w[k, c] * x[c, t - (K-1-k)*dilation] + // The history of the earlier ubatches is prepended, so a chunked prefill matches a single-shot one. + const int64_t kern = hparams.ple_conv_kernel; + const int64_t dil = hparams.ple_ngram_size; + const int64_t hist = (kern - 1) * dil; + + // the conv history is per sequence, so the input carries the sequence axis too + const int64_t n_seqs = ubatch.n_seqs; + const int64_t n_seq_tokens = ubatch.n_seq_tokens; + + // [hist + n_seq_tokens, hc_dim, n_seqs], tokens on ne[0] + ggml_tensor * padded = build_conv_state_at(inp, inp->mctx->get_p_l(il), + ggml_reshape_3d(ctx0, normalized, hc_dim, n_seq_tokens, n_seqs), + hist, hc_dim, il); + + ggml_tensor * conv_out = nullptr; + for (int64_t k = 0; k < kern; ++k) { + // tap k reads (kern-1-k)*dilation positions back + const int64_t start = hist - (kern - 1 - k) * dil; + + ggml_tensor * shifted = ggml_cont(ctx0, + ggml_transpose(ctx0, + ggml_view_3d(ctx0, padded, n_seq_tokens, hc_dim, n_seqs, + padded->nb[1], padded->nb[2], + ggml_row_size(padded->type, start)))); + + // column k of the [kern, hc_dim] kernel is one weight per channel + ggml_tensor * wk = ggml_cont(ctx0, + ggml_view_2d(ctx0, model.layers[il].ple_conv1d, 1, hc_dim, + model.layers[il].ple_conv1d->nb[1], + k * model.layers[il].ple_conv1d->nb[0])); + // this kernel keeps the file type, so cast it before it multiplies an f32 activation + wk = ggml_reshape_1d(ctx0, wk, hc_dim); + if (wk->type != GGML_TYPE_F32) { + wk = ggml_cast(ctx0, wk, GGML_TYPE_F32); + } + + ggml_tensor * term = ggml_mul(ctx0, shifted, wk); + conv_out = conv_out ? ggml_add(ctx0, conv_out, term) : term; + } + + conv_out = ggml_silu(ctx0, conv_out); + conv_out = ggml_reshape_3d(ctx0, ggml_cont(ctx0, conv_out), n_embd, hc, n_tokens); + cb(conv_out, "ple_conv_out", il); + + return ggml_add(ctx0, hidden, ggml_add(ctx0, gated, conv_out)); +} diff --git a/tests/test-arg-parser.cpp b/tests/test-arg-parser.cpp index ba58f852eb4f..e0907631abd8 100644 --- a/tests/test-arg-parser.cpp +++ b/tests/test-arg-parser.cpp @@ -4,6 +4,7 @@ #include "llama.h" #include "speculative.h" +#include #include #include #include @@ -34,6 +35,62 @@ static void test(void) { std::numeric_limits::max(), std::numeric_limits::max()); + { + common_params_speculative spec; + spec.synth_len = 3.4; + + auto assert_invalid = [](const common_params_speculative & value, int32_t n_max) { + try { + common_speculative_synth_rates_resolve(&value, n_max); + assert(false); + } catch (const std::invalid_argument &) { + } + }; + + const auto rates = common_speculative_synth_rates_resolve(&spec, 4); + assert(rates.size() == 4); + assert(std::abs(rates[0] - 0.80581) < 1e-5); + assert(std::abs(rates[1] - 0.64933) < 1e-5); + assert(std::abs(rates[2] - 0.52323) < 1e-5); + assert(std::abs(rates[3] - 0.42163) < 1e-5); + assert(std::abs(1.0 + rates[0] + rates[1] + rates[2] + rates[3] - 3.4) < 1e-8); + + spec.synth_len = 1.0; + assert(common_speculative_synth_rates_resolve(&spec, 4) == std::vector({0.0, 0.0, 0.0, 0.0})); + + spec.synth_len = 5.0; + assert(common_speculative_synth_rates_resolve(&spec, 4) == std::vector({1.0, 1.0, 1.0, 1.0})); + + spec.synth_len = 5.1; + assert_invalid(spec, 4); + + spec.synth_len = std::numeric_limits::quiet_NaN(); + assert_invalid(spec, 4); + + spec.synth_len = 0.0; + assert_invalid(spec, 4); + + spec.synth_len = -1.0; + spec.synth_rates = {0.8, 0.6, 0.4}; + assert_invalid(spec, 4); + + spec.synth_rates = {0.8, 0.6, 0.4, 0.2}; + assert(common_speculative_synth_rates_resolve(&spec, 4) == spec.synth_rates); + + spec.synth_rates = {0.8, 0.9, 0.4, 0.2}; + assert_invalid(spec, 4); + + spec.synth_rates = {0.8, std::numeric_limits::quiet_NaN(), 0.4, 0.2}; + assert_invalid(spec, 4); + + spec.synth_rates = {0.8, 0.6, 0.4, -0.2}; + assert_invalid(spec, 4); + + spec.synth_rates = {0.8, 0.6, 0.4, 0.2}; + spec.synth_len = 3.0; + assert_invalid(spec, 4); + } + { common_params base; base.n_parallel = 4; @@ -197,6 +254,26 @@ static void test(void) { assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), params, LLAMA_EXAMPLE_SPECULATIVE)); assert(params.speculative.draft.n_max == 123); + { + common_params synth_params; + argv = {"binary_name", "--spec-synth-len", "3.4"}; + assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), synth_params, LLAMA_EXAMPLE_SERVER)); + assert(synth_params.speculative.synth_len == 3.4); + } + + { + common_params synth_params; + argv = {"binary_name", "--spec-synth-rates", "0.8,0.6,0.2"}; + assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), synth_params, LLAMA_EXAMPLE_SERVER)); + assert(synth_params.speculative.synth_rates == std::vector({0.8, 0.6, 0.2})); + } + + { + common_params synth_params; + argv = {"binary_name", "--spec-synth-len", "3.4x"}; + assert(false == common_params_parse(argv.size(), list_str_to_char(argv).data(), synth_params, LLAMA_EXAMPLE_SERVER)); + } + argv = {"binary_name", "-lm", "none"}; assert(true == common_params_parse(argv.size(), list_str_to_char(argv).data(), params, LLAMA_EXAMPLE_COMMON)); assert(params.load_mode == LLAMA_LOAD_MODE_NONE); diff --git a/tests/test-backend-ops.cpp b/tests/test-backend-ops.cpp index e4dc152def62..bf643918461a 100644 --- a/tests/test-backend-ops.cpp +++ b/tests/test-backend-ops.cpp @@ -2469,8 +2469,13 @@ struct test_set_rows : public test_case { // See dicussion here: https://github.com/ggml-org/llama.cpp/pull/23760#issuecomment-4566312209 double max_nmse_err(ggml_backend_t backend) override { ggml_backend_reg_t reg = ggml_backend_dev_backend_reg(ggml_backend_get_device(backend)); - if (type_dst == GGML_TYPE_Q8_0 && strcmp(ggml_backend_reg_name(reg), "WebGPU") == 0) { - return std::max(test_case::max_nmse_err(backend), 2e-7); + if (type_dst == GGML_TYPE_Q8_0) { + if (strcmp(ggml_backend_reg_name(reg), "WebGPU") == 0) { + return std::max(test_case::max_nmse_err(backend), 2e-7); + } + if (strcmp(ggml_backend_reg_name(reg), "HTP") == 0) { + return std::max(test_case::max_nmse_err(backend), 5e-6); + } } return test_case::max_nmse_err(backend); } @@ -9808,6 +9813,17 @@ static std::vector> make_test_cases_eval() { } } } + for (int k : {4, 8, 16, 32}) { + for (int nrows : {1, 8, 16}) { + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {202048, nrows, 1, 1}, k)); + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {151936, nrows, 1, 1}, k)); + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {8192, nrows, 1, 1}, k)); + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {8193, nrows, 1, 1}, k)); + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {8192, nrows, 1, 1}, k, true)); + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {202048, nrows, 1, 1}, k, true)); + } + } + for (int k : {1, 2, 3, 7, 15}) { test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {16, 10, 10, 10}, k)); test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {60, 10, 10, 10}, k)); @@ -10560,7 +10576,13 @@ static std::vector> make_test_cases_perf() { test_cases.emplace_back(new test_argsort(GGML_TYPE_F32, {200000, 16, 1, 1})); test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {2, 1, 1, 1}, 1)); - for (auto k : {1, 10, 40, 400}) { + // widths around the tiling threshold + for (auto cols : {4096, 8192, 12288, 16384, 24576, 32768, 65536, 131072}) { + for (auto nrows : {1, 16}) { + test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {cols, nrows, 1, 1}, 16)); + } + } + for (auto k : {1, 4, 8, 10, 16, 32, 40, 400}) { for (auto nrows : {1, 16}) { for (auto cols : {k, 1000, 65000, 200000}) { test_cases.emplace_back(new test_top_k(GGML_TYPE_F32, {cols, nrows, 1, 1}, k)); diff --git a/tests/test-llama-archs.cpp b/tests/test-llama-archs.cpp index 4e205c8e857c..5c7009c48fc5 100644 --- a/tests/test-llama-archs.cpp +++ b/tests/test-llama-archs.cpp @@ -249,8 +249,17 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { // MSA requires one indexer head per GQA (KV) head, unlike the DSA archs where the // indexer head count is independent of the main attention head count. + if (arch == LLM_ARCH_QWEN4EXP) { + ms.add_kv(LLM_KV_HYPER_CONNECTION_COUNT, uint32_t(4)); + ms.add_kv(LLM_KV_HYPER_CONNECTION_LOW_RANK, uint32_t(8)); + // without this the QSA layers fall back to dense and go uncovered + ms.add_kv(LLM_KV_ATTENTION_COMPRESS_RATIOS, std::vector(n_layer, 4)); + } + ms.add_kv(LLM_KV_ATTENTION_INDEXER_HEAD_COUNT, arch == LLM_ARCH_MINIMAX_M3 || arch == LLM_ARCH_DEEPSEEK4 ? n_head : uint32_t(1)); - ms.add_kv(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, uint32_t(64)); + // qwen4exp ropes indexer keys with the main rotary width, so its head can't be < n_rot + ms.add_kv(LLM_KV_ATTENTION_INDEXER_KEY_LENGTH, + arch == LLM_ARCH_QWEN4EXP ? n_embd_head : uint32_t(64)); ms.add_kv(LLM_KV_ATTENTION_INDEXER_TOP_K, uint32_t(8)); ms.add_kv(LLM_KV_ATTENTION_INDEXER_BLOCK_SIZE, uint32_t(4)); ms.add_kv(LLM_KV_ATTENTION_INDEXER_LOCAL_BLOCKS, uint32_t(1)); @@ -294,7 +303,7 @@ static gguf_context_ptr get_gguf_ctx(const llm_arch arch, const bool moe) { ms.add_kv(LLM_KV_XIELU_ALPHA_P, 1.0f); ms.add_kv(LLM_KV_XIELU_BETA, 1.0f); ms.add_kv(LLM_KV_XIELU_EPS, 1.0e-7f); - ms.add_kv(LLM_KV_SSM_INNER_SIZE, arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE ? 256 : 2*n_embd); + ms.add_kv(LLM_KV_SSM_INNER_SIZE, arch == LLM_ARCH_QWEN3NEXT || arch == LLM_ARCH_QWEN35 || arch == LLM_ARCH_QWEN35MOE || arch == LLM_ARCH_QWEN4EXP ? 256 : 2*n_embd); ms.add_kv(LLM_KV_SSM_CONV_KERNEL, uint32_t(4)); ms.add_kv(LLM_KV_SSM_STATE_SIZE, uint32_t(128)); ms.add_kv(LLM_KV_SSM_TIME_STEP_RANK, n_head); @@ -411,6 +420,7 @@ static bool moe_mandatory(const llm_arch arch) { case LLM_ARCH_QWEN3NEXT: case LLM_ARCH_QWEN3VLMOE: case LLM_ARCH_QWEN35MOE: + case LLM_ARCH_QWEN4EXP: case LLM_ARCH_PHIMOE: case LLM_ARCH_DBRX: case LLM_ARCH_OLMOE: @@ -510,7 +520,7 @@ static bool arch_supported(const llm_arch arch) { } // FIXME: these hit scheduler/view-backed-output issues with WebGPU on CI. #ifdef GGML_USE_WEBGPU - if (arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_DOTS3NOTE) { + if (arch == LLM_ARCH_DEEPSEEK32 || arch == LLM_ARCH_GLM_DSA || arch == LLM_ARCH_DOTS3NOTE || arch == LLM_ARCH_QWEN4EXP) { return false; } #endif // GGML_USE_WEBGPU diff --git a/tools/cli/README.md b/tools/cli/README.md index c9cbacafcd1d..163ee4fbaf0a 100644 --- a/tools/cli/README.md +++ b/tools/cli/README.md @@ -59,12 +59,14 @@ | `--mmap, --no-mmap` | DEPRECATED in favor of `--load-mode`: whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock)
(env: LLAMA_ARG_MMAP) | | `-dio, --direct-io, -ndio, --no-direct-io` | DEPRECATED in favor of `--load-mode`: use DirectIO if available
(env: LLAMA_ARG_DIO) | | `-lm, --load-mode MODE` | model loading mode (default: auto)
- auto: mmap, unless a device does not support it
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: force system to keep model in RAM rather than swapping or compressing
- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available

(env: LLAMA_ARG_LOAD_MODE) | +| `--tensor-read-lazy MODE` | on-demand reading of certain tensors, for example per-layer embeddings (default: auto)
- on: read the rows of such tensors from disk on demand instead of keeping them resident (requires mmap)
- auto: on, but only for tensors larger than 4 GiB
- off: always keep them resident
(env: LLAMA_ARG_TENSOR_READ_LAZY) | | `--numa TYPE` | attempt optimizations that help on some NUMA systems
- distribute: spread execution evenly over all nodes
- isolate: only spawn threads on CPUs on the node that execution started on
- numactl: use the CPU map provided by numactl
if run without this previously, it is recommended to drop the system page cache before using this
see https://github.com/ggml-org/llama.cpp/issues/1437
(env: LLAMA_ARG_NUMA) | | `-dev, --device ` | comma-separated list of devices to use for offloading (none = don't offload)
use --list-devices to see a list of available devices
(env: LLAMA_ARG_DEVICE) | | `--list-devices` | print list of available devices and exit | | `-ot, --override-tensor =,...` | override tensor buffer type
(env: LLAMA_ARG_OVERRIDE_TENSOR) | | `-cmoe, --cpu-moe` | keep all Mixture of Experts (MoE) weights in the CPU
(env: LLAMA_ARG_CPU_MOE) | | `-ncmoe, --n-cpu-moe N` | keep the Mixture of Experts (MoE) weights of the first N layers in the CPU
(env: LLAMA_ARG_N_CPU_MOE) | +| `-ncffn, --n-cpu-ffn N` | keep the dense FFN weights of the first N layers in the CPU
(dense models; for MoE expert weights use --n-cpu-moe)
(env: LLAMA_ARG_N_CPU_FFN) | | `-ngl, --gpu-layers, --n-gpu-layers N` | max. number of layers to store in VRAM, either an exact number, 'auto', or 'all' (default: auto)
(env: LLAMA_ARG_N_GPU_LAYERS) | | `-sm, --split-mode {none,layer,row,tensor}` | how to split the model across multiple GPUs, one of:
- none: use one GPU only
- layer (default): split layers and KV across GPUs (pipelined)
- row: split weight across GPUs by rows (parallelized)
- tensor: split weights and KV across GPUs (parallelized, EXPERIMENTAL)
(env: LLAMA_ARG_SPLIT_MODE) | | `-ts, --tensor-split N0,N1,N2,...` | fraction of the model to offload to each GPU, comma-separated list of proportions, e.g. 3,1
(env: LLAMA_ARG_TENSOR_SPLIT) | @@ -154,7 +156,6 @@ | `-sysf, --system-prompt-file FNAME` | a file containing the system prompt (default: none) | | `-r, --reverse-prompt PROMPT` | halt generation at PROMPT, return control in interactive mode | | `-sp, --special` | special tokens output enabled (default: false) | -| `-cnv, --conversation, -no-cnv, --no-conversation` | whether to run in conversation mode:
- does not print special tokens and suffix/prefix
- interactive mode is also enabled
(default: auto enabled if chat template is available) | | `-st, --single-turn` | run conversation for a single turn only, then exit when done
will not be interactive if first turn is predefined with --prompt
(default: false) | | `-mli, --multiline-input` | allows you to write or paste multiple lines without ending each in '\' | | `--warmup, --no-warmup` | whether to perform warmup with an empty run (default: enabled) | @@ -166,6 +167,9 @@ | `--image, --audio, --video FILE` | path to an image, audio, or video file. use with multimodal models, use comma-separated values for multiple files | | `--image-min-tokens N` | minimum number of tokens each image can take, only used by vision models with dynamic resolution (default: read from model)
(env: LLAMA_ARG_IMAGE_MIN_TOKENS) | | `--image-max-tokens N` | maximum number of tokens each image can take, only used by vision models with dynamic resolution (default: read from model)
(env: LLAMA_ARG_IMAGE_MAX_TOKENS) | +| `--video-fps N` | target video frame rate (default: 4.0)
(env: LLAMA_ARG_VIDEO_FPS) | +| `--video-timestamp-interval N` | interval in milliseconds between text timestamps (default: 5000)
(env: LLAMA_ARG_VIDEO_TIMESTAMP_INTERVAL) | +| `--video-ffmpeg-dir DIR` | path to the directory containing ffmpeg and ffprobe (default: search in PATH)
(env: LLAMA_ARG_VIDEO_FFMPEG_DIR) | | `-o, --output, --output-file FNAME` | output file (default: '') | | `--chat-template-kwargs STRING` | sets additional params for the json template parser, must be a valid json object string, e.g. '{"key1":"value1","key2":"value2"}'
(env: LLAMA_ARG_CHAT_TEMPLATE_KWARGS) | | `--jinja, --no-jinja` | whether to use jinja template engine for chat (default: enabled)
(env: LLAMA_ARG_JINJA) | @@ -197,6 +201,8 @@ | `--spec-draft-n-cpu-moe, --spec-draft-ncmoe, -ncmoed, --n-cpu-moe-draft N` | keep the Mixture of Experts (MoE) weights of the first N layers in the CPU for the draft model
(env: LLAMA_ARG_SPEC_DRAFT_N_CPU_MOE) | | `--spec-draft-n-max N` | number of tokens to draft for speculative decoding (default: 3)
(env: LLAMA_ARG_SPEC_DRAFT_N_MAX) | | `--spec-draft-n-min N` | minimum number of draft tokens to use for speculative decoding (default: 0)
(env: LLAMA_ARG_SPEC_DRAFT_N_MIN) | +| `--spec-synth-len L` | target mean synthetic acceptance length, including the target token (benchmarking only)
(env: LLAMA_ARG_SPEC_SYNTH_LEN) | +| `--spec-synth-rates P0,P1,...` | comma-separated unconditional per-position synthetic acceptance probabilities (benchmarking only)
(env: LLAMA_ARG_SPEC_SYNTH_RATES) | | `--spec-draft-p-split, --draft-p-split P` | speculative decoding split probability (default: 0.10)
(env: LLAMA_ARG_SPEC_DRAFT_P_SPLIT) | | `--spec-draft-p-min, --draft-p-min P` | minimum speculative decoding probability (greedy) (default: 0.00)
(env: LLAMA_ARG_SPEC_DRAFT_P_MIN) | | `--spec-draft-backend-sampling, --no-spec-draft-backend-sampling` | offload draft sampling to the backend (default: enabled)
(env: LLAMA_ARG_SPEC_DRAFT_BACKEND_SAMPLING) | diff --git a/tools/completion/README.md b/tools/completion/README.md index 833687dcad44..0cd86bac70fc 100644 --- a/tools/completion/README.md +++ b/tools/completion/README.md @@ -142,12 +142,14 @@ llama-completion.exe -m models\gemma-1.1-7b-it.Q4_K_M.gguf --ignore-eos -n -1 | `--mmap, --no-mmap` | DEPRECATED in favor of `--load-mode`: whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock)
(env: LLAMA_ARG_MMAP) | | `-dio, --direct-io, -ndio, --no-direct-io` | DEPRECATED in favor of `--load-mode`: use DirectIO if available
(env: LLAMA_ARG_DIO) | | `-lm, --load-mode MODE` | model loading mode (default: auto)
- auto: mmap, unless a device does not support it
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: force system to keep model in RAM rather than swapping or compressing
- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available

(env: LLAMA_ARG_LOAD_MODE) | +| `--tensor-read-lazy MODE` | on-demand reading of certain tensors, for example per-layer embeddings (default: auto)
- on: read the rows of such tensors from disk on demand instead of keeping them resident (requires mmap)
- auto: on, but only for tensors larger than 4 GiB
- off: always keep them resident
(env: LLAMA_ARG_TENSOR_READ_LAZY) | | `--numa TYPE` | attempt optimizations that help on some NUMA systems
- distribute: spread execution evenly over all nodes
- isolate: only spawn threads on CPUs on the node that execution started on
- numactl: use the CPU map provided by numactl
if run without this previously, it is recommended to drop the system page cache before using this
see https://github.com/ggml-org/llama.cpp/issues/1437
(env: LLAMA_ARG_NUMA) | | `-dev, --device ` | comma-separated list of devices to use for offloading (none = don't offload)
use --list-devices to see a list of available devices
(env: LLAMA_ARG_DEVICE) | | `--list-devices` | print list of available devices and exit | | `-ot, --override-tensor =,...` | override tensor buffer type
(env: LLAMA_ARG_OVERRIDE_TENSOR) | | `-cmoe, --cpu-moe` | keep all Mixture of Experts (MoE) weights in the CPU
(env: LLAMA_ARG_CPU_MOE) | | `-ncmoe, --n-cpu-moe N` | keep the Mixture of Experts (MoE) weights of the first N layers in the CPU
(env: LLAMA_ARG_N_CPU_MOE) | +| `-ncffn, --n-cpu-ffn N` | keep the dense FFN weights of the first N layers in the CPU
(dense models; for MoE expert weights use --n-cpu-moe)
(env: LLAMA_ARG_N_CPU_FFN) | | `-ngl, --gpu-layers, --n-gpu-layers N` | max. number of layers to store in VRAM, either an exact number, 'auto', or 'all' (default: auto)
(env: LLAMA_ARG_N_GPU_LAYERS) | | `-sm, --split-mode {none,layer,row,tensor}` | how to split the model across multiple GPUs, one of:
- none: use one GPU only
- layer (default): split layers and KV across GPUs (pipelined)
- row: split weight across GPUs by rows (parallelized)
- tensor: split weights and KV across GPUs (parallelized, EXPERIMENTAL)
(env: LLAMA_ARG_SPLIT_MODE) | | `-ts, --tensor-split N0,N1,N2,...` | fraction of the model to offload to each GPU, comma-separated list of proportions, e.g. 3,1
(env: LLAMA_ARG_TENSOR_SPLIT) | diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index a2da93b9a282..bc14d15c7968 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -1254,7 +1254,7 @@ struct cmd_params_instance { merged.reserve(merged.size() + (size_t) n_cpu_moe + 1); for (int i = 0; i < n_cpu_moe; ++i) { - patterns.push_back(llm_ffn_exps_block_regex(i)); + patterns.push_back(llm_ffn_block_regex(i, LLM_FFN_EXPS_REGEX)); merged.push_back({ patterns.back().c_str(), ggml_backend_cpu_buffer_type() }); } diff --git a/tools/mtmd/mtmd-cli.cpp b/tools/mtmd/mtmd-cli.cpp index f6c787fdb683..97678c6b2e59 100644 --- a/tools/mtmd/mtmd-cli.cpp +++ b/tools/mtmd/mtmd-cli.cpp @@ -87,6 +87,9 @@ struct mtmd_cli_context { mtmd::bitmaps bitmaps; std::vector videos; + mtmd_helper_init_opt init_opt = mtmd_helper_init_opt_default(); + std::string video_ffmpeg_bin_dir; + mtmd::batch_ptr mbatch; // chat template @@ -170,6 +173,12 @@ struct mtmd_cli_context { LOG_ERR("Failed to load vision model from %s\n", clip_path); exit(1); } + + video_ffmpeg_bin_dir = params.video_ffmpeg_bin_dir; + init_opt.video_params.fps_target = params.video_fps; + init_opt.video_params.timestamp_interval_ms = params.video_timestamp_interval_ms; + init_opt.video_params.ffmpeg_bin_dir = video_ffmpeg_bin_dir.empty() + ? nullptr : video_ffmpeg_bin_dir.c_str(); } bool check_antiprompt(const llama_tokens & generated_tokens) { @@ -184,7 +193,7 @@ struct mtmd_cli_context { } bool load_media(const std::string & fname) { - auto res = mtmd_helper_bitmap_init_from_file(ctx_vision.get(), fname.c_str(), false); + auto res = mtmd_helper_bitmap_init_from_file(ctx_vision.get(), fname.c_str(), false, init_opt); if (!res.bitmap) { return false; } diff --git a/tools/mtmd/mtmd-helper.cpp b/tools/mtmd/mtmd-helper.cpp index f1defb647728..77f9d58fe238 100644 --- a/tools/mtmd/mtmd-helper.cpp +++ b/tools/mtmd/mtmd-helper.cpp @@ -369,14 +369,18 @@ static bool is_webp_file(const unsigned char * buf, size_t len) { } #ifdef MTMD_VIDEO -static mtmd_bitmap * decode_webp_with_ffmpeg(mtmd_context * mctx, const unsigned char * buf, size_t len, bool placeholder); +static mtmd_bitmap * decode_webp_with_ffmpeg(mtmd_context * mctx, const unsigned char * buf, size_t len, bool placeholder, + const mtmd_helper_video_init_params & params); #endif -mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_buf(mtmd_context * ctx, const unsigned char * buf, size_t len, bool placeholder) { +mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_buf(mtmd_context * ctx, const unsigned char * buf, size_t len, bool placeholder, + mtmd_helper_init_opt opt) { // calculate the hash if needed std::string id; mtmd_bitmap * result = nullptr; + GGML_UNUSED(opt); // only used by video code paths + if (!placeholder) { // use sha256 to prevent cache poisoning id = hash_sha256_hex(buf, len); @@ -414,7 +418,7 @@ mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_buf(mtmd_context * ctx, #ifdef MTMD_VIDEO // stb_image does not support webp; decode it with ffmpeg as a single frame if (!result && is_webp_file(buf, len)) { - result = decode_webp_with_ffmpeg(ctx, buf, len, placeholder); + result = decode_webp_with_ffmpeg(ctx, buf, len, placeholder, opt.video_params); if (!result) { LOG_ERR("%s: failed to decode webp buffer\n", __func__); return {nullptr, nullptr}; @@ -427,8 +431,7 @@ mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_buf(mtmd_context * ctx, // last try: load as video #ifdef MTMD_VIDEO if (!result) { - auto params = mtmd_helper_video_init_params_default(); - auto video_ctx = mtmd_helper_video_init_from_buf(ctx, buf, len, params); + auto video_ctx = mtmd_helper_video_init_from_buf(ctx, buf, len, opt.video_params); if (!video_ctx) { LOG_ERR("%s: failed to decode buffer as either image/audio/video\n", __func__); return {nullptr, nullptr}; @@ -456,7 +459,8 @@ mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_buf(mtmd_context * ctx, return {nullptr, nullptr}; } -mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_file(mtmd_context * ctx, const char * fname, bool placeholder) { +mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_file(mtmd_context * ctx, const char * fname, bool placeholder, + mtmd_helper_init_opt opt) { #ifdef _WIN32 int wlen = MultiByteToWideChar(CP_UTF8, 0, fname, -1, NULL, 0); if (!wlen) { @@ -497,7 +501,7 @@ mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_file(mtmd_context * ctx, return {nullptr, nullptr}; } - return mtmd_helper_bitmap_init_from_buf(ctx, buf.data(), buf.size(), placeholder); + return mtmd_helper_bitmap_init_from_buf(ctx, buf.data(), buf.size(), placeholder, opt); } bool mtmd_helper_support_video(mtmd_context * ctx) { @@ -855,6 +859,12 @@ mtmd_helper_video_init_params mtmd_helper_video_init_params_default() { }; } +mtmd_helper_init_opt mtmd_helper_init_opt_default() { + return { + /* video_params */ mtmd_helper_video_init_params_default(), + }; +} + static std::string video_resolve_bin(const char * bin_dir, const char * name) { if (!bin_dir || bin_dir[0] == '\0') { return name; // rely on PATH @@ -876,8 +886,8 @@ static std::string video_resolve_bin(const char * bin_dir, const char * name) { } #ifdef MTMD_VIDEO -static mtmd_bitmap * decode_webp_with_ffmpeg(mtmd_context * mctx, const unsigned char * buf, size_t len, bool placeholder) { - auto params = mtmd_helper_video_init_params_default(); +static mtmd_bitmap * decode_webp_with_ffmpeg(mtmd_context * mctx, const unsigned char * buf, size_t len, bool placeholder, + const mtmd_helper_video_init_params & params) { mtmd_helper_video vctx; vctx.mctx = mctx; vctx.input_buf.assign(buf, buf + len); diff --git a/tools/mtmd/mtmd-helper.h b/tools/mtmd/mtmd-helper.h index 58dfb1525011..772e0f091bdc 100644 --- a/tools/mtmd/mtmd-helper.h +++ b/tools/mtmd/mtmd-helper.h @@ -23,6 +23,23 @@ extern "C" { struct mtmd_helper_video; typedef struct mtmd_helper_video mtmd_helper_video; +struct mtmd_helper_video_init_params { + float fps_target; // desired output fps; <= 0 means use the video's native fps, defaulted to 4.0f + const char * ffmpeg_bin_dir; // directory containing ffmpeg/ffprobe binaries; NULL means search PATH + int64_t timestamp_interval_ms; // interval for adding timestamp as text chunk (example: "[10m50.5s]"); <= 0 means no timestamp, defaulted to 5000ms + // TODO @ngxson : allow "placeholder" bitmap output for counting tokens +}; + +MTMD_API struct mtmd_helper_video_init_params mtmd_helper_video_init_params_default(void); + +// opt for mtmd_helper_bitmap_init_from_*() +struct mtmd_helper_init_opt { + struct mtmd_helper_video_init_params video_params; +}; +typedef struct mtmd_helper_init_opt mtmd_helper_init_opt; + +MTMD_API struct mtmd_helper_init_opt mtmd_helper_init_opt_default(void); + // Set callback for all future logging events. // If this is not called, or NULL is supplied, everything is output on stderr. // Note: this also call mtmd_log_set() internally @@ -40,7 +57,11 @@ struct mtmd_helper_bitmap_wrapper { // it calls mtmd_helper_bitmap_init_from_buf() internally // returns nullptr on failure // this function is thread-safe -MTMD_API struct mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_file(mtmd_context * ctx, const char * fname, bool placeholder); +MTMD_API struct mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_file( + mtmd_context * ctx, + const char * fname, + bool placeholder, + struct mtmd_helper_init_opt opt); // helper function to construct a mtmd_bitmap from a buffer containing a file // supported formats: @@ -53,7 +74,11 @@ MTMD_API struct mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_file(mtm // - output bitmap will have SHA-256 hash (hex string) as the ID // returns nullptr on failure // this function is thread-safe -MTMD_API struct mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_buf(mtmd_context * ctx, const unsigned char * buf, size_t len, bool placeholder); +MTMD_API struct mtmd_helper_bitmap_wrapper mtmd_helper_bitmap_init_from_buf( + mtmd_context * ctx, + const unsigned char * buf, size_t len, + bool placeholder, + struct mtmd_helper_init_opt opt); // helper to count the total number of tokens from a list of chunks, useful to keep track of KV cache MTMD_API size_t mtmd_helper_get_n_tokens(const mtmd_input_chunks * chunks); @@ -124,14 +149,7 @@ struct mtmd_helper_video_info { int32_t n_frames; // estimated total frames at effective fps (-1 if unknown) }; -struct mtmd_helper_video_init_params { - float fps_target; // desired output fps; <= 0 means use the video's native fps, defaulted to 4.0f - const char * ffmpeg_bin_dir; // directory containing ffmpeg/ffprobe binaries; NULL means search PATH - int64_t timestamp_interval_ms; // interval for adding timestamp as text chunk (example: "[10m50.5s]"); <= 0 means no timestamp, defaulted to 5000ms - // TODO @ngxson : allow "placeholder" bitmap output for counting tokens -}; - -MTMD_API struct mtmd_helper_video_init_params mtmd_helper_video_init_params_default(void); +// note: mtmd_helper_video_init_params is defined at the top, as it is part of mtmd_helper_init_opt // returns NULL on failure (ffprobe not found, file unreadable, etc.) MTMD_API mtmd_helper_video * mtmd_helper_video_init( diff --git a/tools/quantize/quantize.cpp b/tools/quantize/quantize.cpp index 8d03c8fcd427..38950036cd82 100644 --- a/tools/quantize/quantize.cpp +++ b/tools/quantize/quantize.cpp @@ -122,7 +122,7 @@ static bool try_parse_ftype(const std::string & ftype_str_in, llama_ftype & ftyp static void usage(const char * executable) { printf("usage: %s [--help] [--allow-requantize] [--leave-output-tensor] [--pure] [--imatrix] [--include-weights]\n", executable); printf(" [--exclude-weights] [--output-tensor-type] [--token-embedding-type] [--tensor-type] [--tensor-type-file]\n"); - printf(" [--prune-layers] [--keep-split] [--override-kv] [--dry-run]\n"); + printf(" [--prune-layers] [--keep-split] [--override-kv] [--dry-run] [--max-buffer-size]\n"); printf(" model-f32.gguf [model-quant.gguf] type [nthreads]\n\n"); printf(" --allow-requantize\n"); printf(" allow requantizing tensors that have already been quantized\n"); @@ -161,7 +161,10 @@ static void usage(const char * executable) { printf(" WARNING: this is an advanced option, use with care.\n"); printf(" --dry-run\n"); printf(" calculate and show the final quantization size without performing quantization\n"); - printf(" example: llama-quantize --dry-run model-f32.gguf Q4_K\n\n"); + printf(" example: llama-quantize --dry-run model-f32.gguf Q4_K\n"); + printf(" --max-buffer-size MiB\n"); + printf(" max amount of tensor rows kept in memory while quantizing one tensor (default: 8192)\n"); + printf(" lower it to quantize models with very large tensors on a machine with little RAM\n\n"); printf("note: --include-weights and --exclude-weights cannot be used together\n\n"); printf("-----------------------------------------------------------------------------\n"); printf(" allowed quantization types\n"); @@ -467,6 +470,16 @@ int llama_quantize(int argc, char ** argv) { } } else if (strcmp(argv[arg_idx], "--keep-split") == 0) { params.keep_split = true; + } else if (strcmp(argv[arg_idx], "--max-buffer-size") == 0) { + if (arg_idx == argc-1) { + usage(argv[0]); + } + const int mib = atoi(argv[++arg_idx]); + if (mib <= 0) { + fprintf(stderr, "%s: invalid --max-buffer-size '%s'\n", __func__, argv[arg_idx]); + return 1; + } + params.max_buf_size = (size_t) mib * 1024 * 1024; } else { usage(argv[0]); } diff --git a/tools/server/README.md b/tools/server/README.md index 93736c3edfa9..3c2228f34322 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -76,12 +76,14 @@ For the full list of features, please refer to [server's changelog](https://gith | `--mmap, --no-mmap` | DEPRECATED in favor of `--load-mode`: whether to memory-map model. (if mmap disabled, slower load but may reduce pageouts if not using mlock)
(env: LLAMA_ARG_MMAP) | | `-dio, --direct-io, -ndio, --no-direct-io` | DEPRECATED in favor of `--load-mode`: use DirectIO if available
(env: LLAMA_ARG_DIO) | | `-lm, --load-mode MODE` | model loading mode (default: auto)
- auto: mmap, unless a device does not support it
- none: no special loading mode
- mmap: memory-map model (if mmap disabled, slower load but may reduce pageouts if not using mlock)
- mlock: force system to keep model in RAM rather than swapping or compressing
- mmap+mlock: mmap + force system to keep model in RAM rather than swapping or compressing
- dio: use DirectIO if available

(env: LLAMA_ARG_LOAD_MODE) | +| `--tensor-read-lazy MODE` | on-demand reading of certain tensors, for example per-layer embeddings (default: auto)
- on: read the rows of such tensors from disk on demand instead of keeping them resident (requires mmap)
- auto: on, but only for tensors larger than 4 GiB
- off: always keep them resident
(env: LLAMA_ARG_TENSOR_READ_LAZY) | | `--numa TYPE` | attempt optimizations that help on some NUMA systems
- distribute: spread execution evenly over all nodes
- isolate: only spawn threads on CPUs on the node that execution started on
- numactl: use the CPU map provided by numactl
if run without this previously, it is recommended to drop the system page cache before using this
see https://github.com/ggml-org/llama.cpp/issues/1437
(env: LLAMA_ARG_NUMA) | | `-dev, --device ` | comma-separated list of devices to use for offloading (none = don't offload)
use --list-devices to see a list of available devices
(env: LLAMA_ARG_DEVICE) | | `--list-devices` | print list of available devices and exit | | `-ot, --override-tensor =,...` | override tensor buffer type
(env: LLAMA_ARG_OVERRIDE_TENSOR) | | `-cmoe, --cpu-moe` | keep all Mixture of Experts (MoE) weights in the CPU
(env: LLAMA_ARG_CPU_MOE) | | `-ncmoe, --n-cpu-moe N` | keep the Mixture of Experts (MoE) weights of the first N layers in the CPU
(env: LLAMA_ARG_N_CPU_MOE) | +| `-ncffn, --n-cpu-ffn N` | keep the dense FFN weights of the first N layers in the CPU
(dense models; for MoE expert weights use --n-cpu-moe)
(env: LLAMA_ARG_N_CPU_FFN) | | `-ngl, --gpu-layers, --n-gpu-layers N` | max. number of layers to store in VRAM, either an exact number, 'auto', or 'all' (default: auto)
(env: LLAMA_ARG_N_GPU_LAYERS) | | `-sm, --split-mode {none,layer,row,tensor}` | how to split the model across multiple GPUs, one of:
- none: use one GPU only
- layer (default): split layers and KV across GPUs (pipelined)
- row: split weight across GPUs by rows (parallelized)
- tensor: split weights and KV across GPUs (parallelized, EXPERIMENTAL)
(env: LLAMA_ARG_SPLIT_MODE) | | `-ts, --tensor-split N0,N1,N2,...` | fraction of the model to offload to each GPU, comma-separated list of proportions, e.g. 3,1
(env: LLAMA_ARG_TENSOR_SPLIT) | @@ -161,6 +163,7 @@ For the full list of features, please refer to [server's changelog](https://gith | -------- | ----------- | | `-lcs, --lookup-cache-static FNAME` | path to static lookup cache to use for lookup decoding (not updated by generation) | | `-lcd, --lookup-cache-dynamic FNAME` | path to dynamic lookup cache to use for lookup decoding (updated by generation) | +| `--kv-unified-per-slot N` | context limit per parallel slot (default: unset, behavior unchanged).
when set without -c/--ctx-size, the shared KV pool is sized to n_parallel*N
(env: LLAMA_ARG_KV_UNIFIED_PER_SLOT) | | `-ctxcp, --ctx-checkpoints, --swa-checkpoints N` | max number of context checkpoints to create per slot (default: 32)[(more info)](https://github.com/ggml-org/llama.cpp/pull/15293)
(env: LLAMA_ARG_CTX_CHECKPOINTS) | | `-cms, --checkpoint-min-step N` | minimum spacing between context checkpoints in tokens (default: 8192, 0 = no minimum)
(env: LLAMA_ARG_CHECKPOINT_MIN_SPACING_NT) | | `-cram, --cache-ram N` | set the maximum cache size in MiB (default: 8192, -1 - no limit, 0 - disable)[(more info)](https://github.com/ggml-org/llama.cpp/pull/16391)
(env: LLAMA_ARG_CACHE_RAM) | @@ -182,6 +185,9 @@ For the full list of features, please refer to [server's changelog](https://gith | `--image-min-tokens N` | minimum number of tokens each image can take, only used by vision models with dynamic resolution (default: read from model)
(env: LLAMA_ARG_IMAGE_MIN_TOKENS) | | `--image-max-tokens N` | maximum number of tokens each image can take, only used by vision models with dynamic resolution (default: read from model)
(env: LLAMA_ARG_IMAGE_MAX_TOKENS) | | `--mtmd-batch-max-tokens N` | maximum number of image tokens per batch when encoding images (default: 1024)
(env: LLAMA_ARG_MTMD_BATCH_MAX_TOKENS) | +| `--video-fps N` | target video frame rate (default: 4.0)
(env: LLAMA_ARG_VIDEO_FPS) | +| `--video-timestamp-interval N` | interval in milliseconds between text timestamps (default: 5000)
(env: LLAMA_ARG_VIDEO_TIMESTAMP_INTERVAL) | +| `--video-ffmpeg-dir DIR` | path to the directory containing ffmpeg and ffprobe (default: search in PATH)
(env: LLAMA_ARG_VIDEO_FFMPEG_DIR) | | `-a, --alias STRING` | set model name aliases, comma-separated (to be used by API)
(env: LLAMA_ARG_ALIAS) | | `--tags STRING` | set model tags, comma-separated (informational, not used for routing)
(env: LLAMA_ARG_TAGS) | | `--embd-normalize N` | normalisation for embeddings (default: 2) (-1=none, 0=max absolute int16, 1=taxicab, 2=euclidean, >2=p-norm) | @@ -256,6 +262,8 @@ For the full list of features, please refer to [server's changelog](https://gith | `--spec-draft-n-cpu-moe, --spec-draft-ncmoe, -ncmoed, --n-cpu-moe-draft N` | keep the Mixture of Experts (MoE) weights of the first N layers in the CPU for the draft model
(env: LLAMA_ARG_SPEC_DRAFT_N_CPU_MOE) | | `--spec-draft-n-max N` | number of tokens to draft for speculative decoding (default: 3)
(env: LLAMA_ARG_SPEC_DRAFT_N_MAX) | | `--spec-draft-n-min N` | minimum number of draft tokens to use for speculative decoding (default: 0)
(env: LLAMA_ARG_SPEC_DRAFT_N_MIN) | +| `--spec-synth-len L` | target mean synthetic acceptance length, including the target token (benchmarking only)
(env: LLAMA_ARG_SPEC_SYNTH_LEN) | +| `--spec-synth-rates P0,P1,...` | comma-separated unconditional per-position synthetic acceptance probabilities (benchmarking only)
(env: LLAMA_ARG_SPEC_SYNTH_RATES) | | `--spec-draft-p-split, --draft-p-split P` | speculative decoding split probability (default: 0.10)
(env: LLAMA_ARG_SPEC_DRAFT_P_SPLIT) | | `--spec-draft-p-min, --draft-p-min P` | minimum speculative decoding probability (greedy) (default: 0.00)
(env: LLAMA_ARG_SPEC_DRAFT_P_MIN) | | `--spec-draft-backend-sampling, --no-spec-draft-backend-sampling` | offload draft sampling to the backend (default: enabled)
(env: LLAMA_ARG_SPEC_DRAFT_BACKEND_SAMPLING) | diff --git a/tools/server/server-common.cpp b/tools/server/server-common.cpp index 7997d4016a68..c30955e89f03 100644 --- a/tools/server/server-common.cpp +++ b/tools/server/server-common.cpp @@ -910,12 +910,17 @@ size_t validate_utf8(const std::string& text) { return len; } -server_tokens process_mtmd_prompt(mtmd_context * mctx, const std::string & prompt, const std::vector & files, bool is_placeholder) { +server_tokens process_mtmd_prompt( + mtmd_context * mctx, + const std::string & prompt, + const std::vector & files, + const mtmd_helper_init_opt & init_opt, + bool is_placeholder) { // these will be freed upon going out of scope mtmd::bitmaps bitmaps; std::vector videos; for (auto & file : files) { - auto out = mtmd_helper_bitmap_init_from_buf(mctx, file.data(), file.size(), is_placeholder); + auto out = mtmd_helper_bitmap_init_from_buf(mctx, file.data(), file.size(), is_placeholder, init_opt); if (!out.bitmap) { throw std::runtime_error("Failed to load image or audio file"); } @@ -956,7 +961,7 @@ server_tokens process_mtmd_prompt(mtmd_context * mctx, const std::string & promp * - "prompt": [12, 34, "string", 56, 78] * - "prompt": { "prompt_string": "string", "multimodal_data": [ "base64" ] } */ -static server_tokens tokenize_input_subprompt(const llama_vocab * vocab, mtmd_context * mctx, const json & json_prompt, bool add_special, bool parse_special) { +static server_tokens tokenize_input_subprompt(const llama_vocab * vocab, mtmd_context * mctx, const json & json_prompt, bool add_special, bool parse_special, const mtmd_helper_init_opt & init_opt) { constexpr char JSON_STRING_PROMPT_KEY[] = "prompt_string"; constexpr char JSON_MTMD_DATA_KEY[] = "multimodal_data"; const bool has_mtmd = mctx != nullptr; @@ -979,7 +984,7 @@ static server_tokens tokenize_input_subprompt(const llama_vocab * vocab, mtmd_co for (const auto & entry : json_prompt.at(JSON_MTMD_DATA_KEY)) { files.push_back(base64_decode(entry)); } - return process_mtmd_prompt(mctx, json_prompt.at(JSON_STRING_PROMPT_KEY), files); + return process_mtmd_prompt(mctx, json_prompt.at(JSON_STRING_PROMPT_KEY), files, init_opt); } else { // Not multimodal, but contains a subobject. llama_tokens tmp = tokenize_mixed(vocab, json_prompt.at(JSON_STRING_PROMPT_KEY), add_special, parse_special); @@ -990,15 +995,15 @@ static server_tokens tokenize_input_subprompt(const llama_vocab * vocab, mtmd_co } } -std::vector tokenize_input_prompts(const llama_vocab * vocab, mtmd_context * mctx, const json & json_prompt, bool add_special, bool parse_special) { +std::vector tokenize_input_prompts(const llama_vocab * vocab, mtmd_context * mctx, const json & json_prompt, bool add_special, bool parse_special, const mtmd_helper_init_opt & init_opt) { std::vector result; if (json_prompt.is_array() && !json_is_array_and_contains_numbers(json_prompt)) { result.reserve(json_prompt.size()); for (const auto & p : json_prompt) { - result.push_back(tokenize_input_subprompt(vocab, mctx, p,add_special, parse_special)); + result.push_back(tokenize_input_subprompt(vocab, mctx, p, add_special, parse_special, init_opt)); } } else { - result.push_back(tokenize_input_subprompt(vocab, mctx, json_prompt, add_special, parse_special)); + result.push_back(tokenize_input_subprompt(vocab, mctx, json_prompt, add_special, parse_special, init_opt)); } if (result.empty()) { throw std::runtime_error("\"prompt\" must not be empty"); @@ -1787,7 +1792,8 @@ server_tokens format_prompt_rerank( const struct llama_vocab * vocab, mtmd_context * mctx, const std::string & query, - const std::string & doc) { + const std::string & doc, + const mtmd_helper_init_opt & init_opt) { server_tokens result = {}; const char * rerank_prompt = llama_model_chat_template(model, "rerank"); @@ -1796,12 +1802,12 @@ server_tokens format_prompt_rerank( std::string prompt = rerank_prompt; string_replace_all(prompt, "{query}" , query); string_replace_all(prompt, "{document}", doc ); - server_tokens tokens = tokenize_input_subprompt(vocab, mctx, prompt, false, true); + server_tokens tokens = tokenize_input_subprompt(vocab, mctx, prompt, false, true, init_opt); result.push_back(tokens); } else { // Get EOS token - use SEP token as fallback if EOS is not available - server_tokens query_tokens = tokenize_input_subprompt(vocab, mctx, query, false, false); - server_tokens doc_tokens = tokenize_input_subprompt(vocab, mctx, doc, false, false); + server_tokens query_tokens = tokenize_input_subprompt(vocab, mctx, query, false, false, init_opt); + server_tokens doc_tokens = tokenize_input_subprompt(vocab, mctx, doc, false, false, init_opt); llama_token eos_token = llama_vocab_eos(vocab); if (eos_token == LLAMA_TOKEN_NULL) { eos_token = llama_vocab_sep(vocab); diff --git a/tools/server/server-common.h b/tools/server/server-common.h index f8ea82ef4cf5..6c681a2cf56d 100644 --- a/tools/server/server-common.h +++ b/tools/server/server-common.h @@ -5,6 +5,7 @@ #include "llama.h" #include "chat.h" #include "mtmd.h" +#include "mtmd-helper.h" #include "json.h" @@ -269,7 +270,12 @@ size_t validate_utf8(const std::string& text); // process mtmd prompt, return the server_tokens containing both text tokens and media chunks // if is_placeholder is true, the media chunk will be treated as placeholder for counting tokens; the output tokens are not usable for actual inference (e.g. for submitting a task to server_queue) -server_tokens process_mtmd_prompt(mtmd_context * mctx, const std::string & prompt, const std::vector & files, bool is_placeholder = false); +server_tokens process_mtmd_prompt( + mtmd_context * mctx, + const std::string & prompt, + const std::vector & files, + const mtmd_helper_init_opt & init_opt, + bool is_placeholder = false); /** * break the input "prompt" object into multiple prompt if needed, then tokenize them @@ -289,7 +295,8 @@ std::vector tokenize_input_prompts( mtmd_context * mctx, const json & json_prompt, bool add_special, - bool parse_special); + bool parse_special, + const mtmd_helper_init_opt & init_opt); // // OAI utils @@ -538,7 +545,8 @@ server_tokens format_prompt_rerank( const struct llama_vocab * vocab, mtmd_context * mctx, const std::string & query, - const std::string & doc); + const std::string & doc, + const mtmd_helper_init_opt & init_opt); // simple implementation of a pipe // used for streaming data between threads diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index a9edbd7be8b4..f5477356d61d 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -23,6 +23,7 @@ #include #include #include +#include #include #include @@ -51,6 +52,50 @@ static common_speculative_output_limits server_output_limits(const common_params return result; } +// synthetic draft verification for benchmarking - accept draft tokens at random instead of by match with the target +// on replay the draft was already accepted before a context checkpoint restore, so repeat the same decisions +static std::vector server_sample_and_accept_synth( + common_sampler * smpl, + llama_context * ctx, + const std::vector & idxs, + const llama_tokens & draft, + const std::vector & synth_probs, + std::mt19937 & rng, + bool is_replay) { + GGML_ASSERT(idxs.size() == draft.size() + 1); + GGML_ASSERT(synth_probs.size() >= draft.size()); + + std::vector result; + result.reserve(idxs.size()); + + const llama_vocab * vocab = llama_model_get_vocab(llama_get_model(ctx)); + std::uniform_real_distribution dist(0.0, 1.0); + for (size_t i = 0; i < draft.size(); ++i) { + const llama_token id = common_sampler_sample(smpl, ctx, idxs[i]); + const bool accept = is_replay || dist(rng) < synth_probs[i]; + // do not accept a drafted EOG token - it would end the generation early + // on replay the last token is from the target and can be EOG, so skip this check + if (accept && (is_replay || !llama_vocab_is_eog(vocab, draft[i]))) { + // synthetic draft tokens do not advance grammar or reasoning state + // the last replay token is from the target and must advance both + const bool is_replay_target = is_replay && i + 1 == draft.size(); + common_sampler_accept(smpl, draft[i], is_replay_target); + result.push_back(draft[i]); + continue; + } + + common_sampler_accept(smpl, id, true); + result.push_back(id); + return result; + } + + const llama_token id = common_sampler_sample(smpl, ctx, idxs[draft.size()]); + common_sampler_accept(smpl, id, true); + result.push_back(id); + + return result; +} + // state diagram: https://github.com/ggml-org/llama.cpp/pull/9283 enum slot_state { SLOT_STATE_IDLE, @@ -211,6 +256,7 @@ struct server_slot { std::vector spec_i_batch; common_prompt_checkpoint spec_ckpt; bool spec_is_replay = false; + std::mt19937 spec_synth_rng; // TODO: move members that belong to the task (such as `generated_text`, `has_new_line`) to task_results_state // see https://github.com/ggml-org/llama.cpp/pull/18283#issuecomment-3710175837 @@ -794,6 +840,8 @@ struct server_context_impl { llama_model * model_tgt = nullptr; mtmd_context * mctx = nullptr; + // note: video_params.ffmpeg_bin_dir points into params_base, which outlives this struct + mtmd_helper_init_opt init_opt = mtmd_helper_init_opt_default(); const llama_vocab * vocab = nullptr; server_queue queue_tasks; @@ -1118,6 +1166,11 @@ struct server_context_impl { } SRV_INF("loaded multimodal model, '%s'\n", mmproj_path.c_str()); + init_opt.video_params.fps_target = params_base.video_fps; + init_opt.video_params.timestamp_interval_ms = params_base.video_timestamp_interval_ms; + init_opt.video_params.ffmpeg_bin_dir = params_base.video_ffmpeg_bin_dir.empty() + ? nullptr : params_base.video_ffmpeg_bin_dir.c_str(); + if (params_base.ctx_shift) { params_base.ctx_shift = false; SRV_WRN("%s\n", "ctx_shift is not supported by multimodal, it will be disabled"); @@ -1155,10 +1208,31 @@ struct server_context_impl { const int n_ctx_train = llama_model_n_ctx_train(model_tgt); - int n_ctx_slot = llama_n_ctx_seq(ctx_tgt); - if (n_ctx_slot > n_ctx_train) { - SRV_WRN("the slot context (%d) exceeds the training context of the model (%d) - capping\n", n_ctx_slot, n_ctx_train); - n_ctx_slot = n_ctx_train; + { + // note: the capping itself is done in n_ctx_slot(), here we only report it + const int n_ctx_seq = llama_n_ctx_seq(ctx_tgt); + + if (params_base.kv_unified_per_slot > 0) { + if (n_ctx_seq > params_base.kv_unified_per_slot) { + SRV_INF("capping per-slot context (%d) to --kv-unified-per-slot (%d)\n", + n_ctx_seq, params_base.kv_unified_per_slot); + } else if (params_base.kv_unified_per_slot > n_ctx_seq) { + // cap is above the per-slot pool capacity, so it can never bind + SRV_WRN( + "--kv-unified-per-slot (%d) exceeds the per-slot pool capacity (%d) - cap has no effect, " + "slots are limited to %d (raise the KV pool with -c, or unset -c to size it to " + "n_parallel * kv_unified_per_slot)\n", + params_base.kv_unified_per_slot, n_ctx_seq, n_ctx_seq); + } + } + + const int n_ctx_capped = params_base.kv_unified_per_slot > 0 ? + std::min(n_ctx_seq, params_base.kv_unified_per_slot) : n_ctx_seq; + + if (n_ctx_capped > n_ctx_train) { + SRV_WRN("the slot context (%d) exceeds the training context of the model (%d) - capping\n", + n_ctx_capped, n_ctx_train); + } } slots.clear(); @@ -1174,7 +1248,7 @@ struct server_context_impl { // setup slots SRV_INF("initializing, n_slots = %d, n_ctx_slot = %d, kv_unified = '%s'\n", - params_base.n_parallel, n_ctx_slot, params_base.kv_unified ? "true" : "false"); + params_base.n_parallel, n_ctx_slot(), params_base.kv_unified ? "true" : "false"); // initialize slots for (int i = 0; i < params_base.n_parallel; i++) { @@ -1187,6 +1261,9 @@ struct server_context_impl { spec.reset(common_speculative_init(params_base.speculative, params_base.n_parallel)); } catch (const std::exception & e) { SRV_ERR("failed to initialize speculative decoding context: %s\n", e.what()); + if (params_base.speculative.has_synth()) { + return false; + } } } @@ -1202,6 +1279,11 @@ struct server_context_impl { model_dft = nullptr; } + if (!spec && params_base.speculative.has_synth()) { + SRV_ERR("%s", "synthetic acceptance requires an initialized speculative decoding context\n"); + return false; + } + for (int i = 0; i < params_base.n_parallel; i++) { server_slot & slot = slots[i]; @@ -1210,7 +1292,7 @@ struct server_context_impl { slot.ctx_dft = ctx_dft; slot.mem.init(ctx_tgt, ctx_dft); slot.spec = spec.get(); - slot.n_ctx = n_ctx_slot; + slot.n_ctx = n_ctx_slot(); slot.mctx = mctx; slot.prompt.tokens.has_mtmd = mctx != nullptr; @@ -1710,6 +1792,13 @@ struct server_context_impl { SLT_TRC(slot, "sampler chain: %s\n", common_sampler_print(slot.smpl.get()).c_str()); SLT_TRC(slot, "sampler params: \n%s\n", task.params.sampling.print().c_str()); + + if (spec && !common_speculative_get_synth_probs(spec.get()).empty()) { + const uint32_t seed = task.params.sampling.seed == LLAMA_DEFAULT_SEED + ? std::random_device{}() + : task.params.sampling.seed; + slot.spec_synth_rng.seed(seed); + } } else { slot.smpl.reset(); } @@ -2134,9 +2223,9 @@ struct server_context_impl { try { auto & prompt = task.cli_prompt; if (mctx != nullptr) { - task.tokens = process_mtmd_prompt(mctx, prompt, task.cli_files); + task.tokens = process_mtmd_prompt(mctx, prompt, task.cli_files, init_opt); } else { - task.tokens = std::move(tokenize_input_prompts(vocab, mctx, prompt, true, true)[0]); + task.tokens = std::move(tokenize_input_prompts(vocab, mctx, prompt, true, true, init_opt)[0]); } task.cli_prompt.clear(); task.cli_files.clear(); @@ -3795,7 +3884,12 @@ struct server_context_impl { common_sampler_ptr smpl_save(common_sampler_clone(slot.smpl.get())); GGML_ASSERT(slot.spec_i_batch.size() == n_draft + 1); - auto accepted = common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft); + const auto & synth_probs = common_speculative_get_synth_probs(spec.get()); + auto accepted = synth_probs.empty() + ? common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft) + : server_sample_and_accept_synth( + slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft, + synth_probs, slot.spec_synth_rng, slot.spec_is_replay); slot.spec_i_batch.clear(); GGML_ASSERT(accepted.size() >= 1); @@ -3861,7 +3955,7 @@ struct server_context_impl { auto & n_accepted_per_pos = slot.n_accepted_per_pos; if (n_accepted_per_pos.empty()) { - n_accepted_per_pos.resize(common_speculative_n_max(¶ms_base.speculative), 0); + n_accepted_per_pos.resize(common_speculative_n_max(spec.get()), 0); } for (size_t i = 0; i < n_accepted && i < n_accepted_per_pos.size(); ++i) { n_accepted_per_pos[i]++; @@ -3902,8 +3996,15 @@ struct server_context_impl { }); } - int get_slot_n_ctx() { - return slots.back().n_ctx; + // context size of a single slot, capped by --kv-unified-per-slot and by the training context of the model + int n_ctx_slot() const { + int res = llama_n_ctx_seq(ctx_tgt); + + if (params_base.kv_unified_per_slot > 0) { + res = std::min(res, params_base.kv_unified_per_slot); + } + + return std::min(res, llama_model_n_ctx_train(model_tgt)); } server_response_reader get_response_reader() { @@ -4069,7 +4170,7 @@ server_context_meta server_context::get_meta() const { /* has_inp_audio */ impl->chat_params.allow_audio, /* has_inp_video */ impl->chat_params.allow_video, /* json_ui_settings */ impl->json_ui_settings, - /* slot_n_ctx */ impl->get_slot_n_ctx(), + /* slot_n_ctx */ impl->n_ctx_slot(), /* pooling_type */ llama_pooling_type(impl->ctx_tgt), /* chat_params */ impl->chat_params, @@ -4165,10 +4266,10 @@ std::unique_ptr server_routes::handle_completions_impl( if (res_type != TASK_RESPONSE_TYPE_NONE && ctx_server.mctx != nullptr) { // This is the case used by OAI compatible chat path with MTMD. TODO It can be moved to the path below. - inputs.push_back(process_mtmd_prompt(ctx_server.mctx, prompt.get(), files)); + inputs.push_back(process_mtmd_prompt(ctx_server.mctx, prompt.get(), files, ctx_server.init_opt)); } else { // Everything else, including multimodal completions. - inputs = tokenize_input_prompts(ctx_server.vocab, ctx_server.mctx, prompt, true, true); + inputs = tokenize_input_prompts(ctx_server.vocab, ctx_server.mctx, prompt, true, true, ctx_server.init_opt); } // tasks.reserve(inputs.size()); // TODO: this is inaccurate due to child tasks @@ -4752,7 +4853,7 @@ void server_routes::init_routes() { data["input_extra"] = input_extra; // default to empty array if it's not exist std::string prompt = json_value(data, "prompt", std::string()); - std::vector tokenized_prompts = tokenize_input_prompts(ctx_server.vocab, ctx_server.mctx, prompt, false, true); + std::vector tokenized_prompts = tokenize_input_prompts(ctx_server.vocab, ctx_server.mctx, prompt, false, true, ctx_server.init_opt); SRV_DBG("creating infill tasks, n_prompts = %d\n", (int) tokenized_prompts.size()); data["prompt"] = format_prompt_infill( ctx_server.vocab, @@ -4816,7 +4917,7 @@ void server_routes::init_routes() { }; this->post_chat_completions_tok = [this](const server_http_req & req) { - return handle_count_tokens(ctx_server.vocab, ctx_server.mctx, req, TASK_RESPONSE_TYPE_OAI_CHAT); + return handle_count_tokens(ctx_server.vocab, ctx_server.mctx, ctx_server.init_opt, req, TASK_RESPONSE_TYPE_OAI_CHAT); }; this->post_control = [this](const server_http_req & req) { @@ -4875,7 +4976,7 @@ void server_routes::init_routes() { }; this->post_responses_tok_oai = [this](const server_http_req & req) { - return handle_count_tokens(ctx_server.vocab, ctx_server.mctx, req, TASK_RESPONSE_TYPE_OAI_RESP); + return handle_count_tokens(ctx_server.vocab, ctx_server.mctx, ctx_server.init_opt, req, TASK_RESPONSE_TYPE_OAI_RESP); }; this->post_transcriptions_oai = [this](const server_http_req & req) { @@ -4925,7 +5026,7 @@ void server_routes::init_routes() { }; this->post_anthropic_count_tokens = [this](const server_http_req & req) { - return handle_count_tokens(ctx_server.vocab, ctx_server.mctx, req, TASK_RESPONSE_TYPE_ANTHROPIC); + return handle_count_tokens(ctx_server.vocab, ctx_server.mctx, ctx_server.init_opt, req, TASK_RESPONSE_TYPE_ANTHROPIC); }; // same with handle_chat_completions, but without inference part @@ -5058,7 +5159,7 @@ void server_routes::init_routes() { std::vector tasks; tasks.reserve(documents.size()); for (size_t i = 0; i < documents.size(); i++) { - auto tmp = format_prompt_rerank(ctx_server.model_tgt, ctx_server.vocab, ctx_server.mctx, query, documents[i]); + auto tmp = format_prompt_rerank(ctx_server.model_tgt, ctx_server.vocab, ctx_server.mctx, query, documents[i], ctx_server.init_opt); server_task task = server_task(SERVER_TASK_TYPE_RERANK); task.id = rd.get_new_id(); task.tokens = std::move(tmp); @@ -5296,7 +5397,7 @@ std::unique_ptr server_routes::handle_embeddings_impl(cons } } - auto tokenized_prompts = tokenize_input_prompts(ctx_server.vocab, ctx_server.mctx, prompt, true, true); + auto tokenized_prompts = tokenize_input_prompts(ctx_server.vocab, ctx_server.mctx, prompt, true, true, ctx_server.init_opt); for (const auto & tokens : tokenized_prompts) { // this check is necessary for models that do not add BOS token to the input if (tokens.empty()) { @@ -5357,7 +5458,7 @@ std::unique_ptr server_routes::handle_embeddings_impl(cons return res; } -std::unique_ptr server_routes::handle_count_tokens(const llama_vocab * vocab, mtmd_context * mctx, const server_http_req & req, task_response_type res_type) { +std::unique_ptr server_routes::handle_count_tokens(const llama_vocab * vocab, mtmd_context * mctx, const mtmd_helper_init_opt & init_opt, const server_http_req & req, task_response_type res_type) { auto res = create_response(); std::vector files; json body = json::parse(req.body); @@ -5395,7 +5496,7 @@ std::unique_ptr server_routes::handle_count_tokens(const l if (!prompt.is_string()) { throw std::runtime_error("for mtmd, input prompt must be a string."); } - n_tokens = process_mtmd_prompt(mctx, prompt.get(), files, true).size(); + n_tokens = process_mtmd_prompt(mctx, prompt.get(), files, init_opt, true).size(); } else { n_tokens = tokenize_mixed(vocab, prompt, true, true).size(); } diff --git a/tools/server/server-context.h b/tools/server/server-context.h index 5d464b8e8cb7..0acbbffa9e10 100644 --- a/tools/server/server-context.h +++ b/tools/server/server-context.h @@ -169,7 +169,7 @@ struct server_routes { std::unique_ptr handle_slots_restore(const server_http_req & req, int id_slot); std::unique_ptr handle_slots_erase(const server_http_req &, int id_slot); std::unique_ptr handle_embeddings_impl(const server_http_req & req, task_response_type res_type); - std::unique_ptr handle_count_tokens(const llama_vocab * vocab, mtmd_context * mctx, const server_http_req & req, task_response_type res_type); + std::unique_ptr handle_count_tokens(const llama_vocab * vocab, mtmd_context * mctx, const mtmd_helper_init_opt & init_opt, const server_http_req & req, task_response_type res_type); // using unique_ptr to allow late initialization of const std::unique_ptr meta; diff --git a/tools/server/server.cpp b/tools/server/server.cpp index 5fe2729ba1b2..22378b38c5ef 100644 --- a/tools/server/server.cpp +++ b/tools/server/server.cpp @@ -157,6 +157,18 @@ int llama_server(common_params & params, int argc, char ** argv) { } } + // size the KV pool from --kv-unified-per-slot, unless the user pinned it with -c + // or with -c 0 for max context + const bool ctx_pool_auto_sized = params.kv_unified_per_slot > 0 && + params.n_ctx == 0 && + (uint32_t) params.fit_params_min_ctx != UINT32_MAX; + + if (ctx_pool_auto_sized) { + params.n_ctx = params.n_parallel * params.kv_unified_per_slot; + SRV_INF("--kv-unified-per-slot: sizing KV pool to n_parallel * kv_unified_per_slot = %d * %d = %d\n", params.n_parallel, + params.kv_unified_per_slot, params.n_ctx); + } + // for consistency between server router mode and single-model mode, we set the same model name as alias auto model_name = params.model.get_name(); if (params.model_alias.empty() && !model_name.empty()) { diff --git a/tools/server/tests/unit/test_speculative.py b/tools/server/tests/unit/test_speculative.py index 5837195006bc..22b523954ec7 100644 --- a/tools/server/tests/unit/test_speculative.py +++ b/tools/server/tests/unit/test_speculative.py @@ -52,6 +52,18 @@ def test_with_and_without_draft(): assert tokens_no_draft == tokens_draft + server.stop() + create_server() + assert server.spec_draft_n_max is not None + server.spec_synth_rates = [0.0] * server.spec_draft_n_max + server.start() + res = server.make_request("POST", "/completion", data=request) + + assert res.status_code == 200 + assert res.body["timings"]["draft_n"] > 0 + assert res.body["timings"]["draft_n_accepted"] == 0 + assert res.body["tokens"] == tokens_no_draft + def test_different_draft_min_draft_max(): global server @@ -80,6 +92,66 @@ def test_different_draft_min_draft_max(): last_content = res.body["content"] +def test_synth_is_deterministic(): + global server + assert server.spec_draft_n_max is not None + server.spec_synth_rates = [0.75 ** (i + 1) for i in range(server.spec_draft_n_max)] + server.start() + + request = { + "prompt": "I believe the meaning of life is", + "temperature": 0.2, + "top_k": 5, + "seed": 4242, + "n_predict": 32, + } + responses = [server.make_request("POST", "/completion", data=request) for _ in range(2)] + + for res in responses: + assert res.status_code == 200 + assert res.body["timings"]["draft_n"] > 0 + assert responses[0].body["timings"]["draft_n"] == responses[1].body["timings"]["draft_n"] + assert responses[0].body["timings"]["draft_n_accepted"] == responses[1].body["timings"]["draft_n_accepted"] + + +def test_synth_ignores_target_tokens(): + global server + assert server.spec_draft_n_max is not None + server.spec_synth_rates = [1.0] * server.spec_draft_n_max + server.start() + + res = server.make_request("POST", "/completion", data={ + "prompt": "I believe the meaning of life is", + "temperature": 0.0, + "seed": 4242, + "n_predict": 32, + }) + + assert res.status_code == 200 + assert res.body["timings"]["draft_n"] > 0 + assert res.body["timings"]["draft_n_accepted"] == res.body["timings"]["draft_n"] + + res = server.make_request("POST", "/completion", data={ + "prompt": "I believe the meaning of life is", + "temperature": 0.0, + "seed": 4242, + "n_predict": 6, + "grammar": 'root ::= "a"{5,5}', + }) + assert res.status_code == 200, res.body + + res = server.make_request("POST", "/completion", data={ + "prompt": "Respond with only: OK", + "temperature": 0.0, + "seed": 4242, + "n_predict": 64, + "ignore_eos": True, + }) + assert res.status_code == 200, res.body + assert res.body["tokens_predicted"] == 64 + assert res.body["stop_type"] == "limit" + + def test_slot_ctx_not_exceeded(): global server server.n_ctx = 256 diff --git a/tools/server/tests/utils.py b/tools/server/tests/utils.py index a0d2dfa3c591..5a4f31a53e01 100644 --- a/tools/server/tests/utils.py +++ b/tools/server/tests/utils.py @@ -99,6 +99,8 @@ class ServerProcess: spec_type: str | None = None spec_draft_n_min: int | None = None spec_draft_n_max: int | None = None + spec_synth_len: float | None = None + spec_synth_rates: List[float] | None = None no_ui: bool | None = None jinja: bool | None = None reasoning_format: Literal['deepseek', 'none', 'nothink'] | None = None @@ -245,6 +247,11 @@ def start(self, timeout_seconds: int = DEFAULT_HTTP_TIMEOUT) -> None: server_args.extend(["--spec-draft-n-max", self.spec_draft_n_max]) if self.spec_draft_n_min: server_args.extend(["--spec-draft-n-min", self.spec_draft_n_min]) + if self.spec_synth_len is not None: + server_args.extend(["--spec-synth-len", self.spec_synth_len]) + if self.spec_synth_rates is not None: + rates = ",".join(str(rate) for rate in self.spec_synth_rates) + server_args.extend(["--spec-synth-rates", rates]) if self.no_ui: server_args.append("--no-ui") if self.no_models_autoload: diff --git a/tools/tts/tts.cpp b/tools/tts/tts.cpp index 368123baf539..6fd1936324ce 100644 --- a/tools/tts/tts.cpp +++ b/tools/tts/tts.cpp @@ -103,7 +103,7 @@ int main(int argc, char ** argv) { mtmd::bitmap_ptr speaker_bitmap; if (!params.tts_speaker_file.empty()) { - auto wrapper = mtmd_helper_bitmap_init_from_file(mctx.get(), params.tts_speaker_file.c_str(), false); + auto wrapper = mtmd_helper_bitmap_init_from_file(mctx.get(), params.tts_speaker_file.c_str(), false, mtmd_helper_init_opt_default()); if (!wrapper.bitmap) { LOG_ERR("failed to load speaker file %s\n", params.tts_speaker_file.c_str()); return 1; diff --git a/tools/ui/src/lib/components/app/badges/BadgesModality.svelte b/tools/ui/src/lib/components/app/badges/BadgesModality.svelte index 4eb3e7838d18..83b1b46affc6 100644 --- a/tools/ui/src/lib/components/app/badges/BadgesModality.svelte +++ b/tools/ui/src/lib/components/app/badges/BadgesModality.svelte @@ -1,5 +1,5 @@ -{#each modalities as modality (modality)} - {#if modality === ModelModality.VISION || modality === ModelModality.AUDIO || modality === ModelModality.VIDEO} - - {#if modality === ModelModality.VISION} - + const shownModalities = [ModelModality.VISION, ModelModality.AUDIO, ModelModality.VIDEO] as const; - Vision (Image) - {:else if modality === ModelModality.VIDEO} - {/each} 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 40819d6f1a17..893f8077dd6a 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatForm.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatForm.svelte @@ -23,7 +23,8 @@ FileExtensionText, KeyboardKey, MimeTypeText, - SpecialFileType + SpecialFileType, + ToolSource } from '$lib/enums'; import { useChatFormPickers } from '$lib/hooks/use-chat-form-pickers.svelte'; import { @@ -73,7 +74,6 @@ disabled?: boolean; isLoading?: boolean; placeholder?: string; - showMcpPromptButton?: boolean; showAddButton?: boolean; showModelSelector?: boolean; @@ -103,7 +103,6 @@ onValueChange, placeholder = 'Type a message...', showAddButton = true, - showMcpPromptButton = false, showModelSelector = true, uploadedFiles = $bindable([]), value = $bindable('') @@ -152,9 +151,18 @@ getServerHome: () => toolsStore.serverHome ?? null, getShowModelSelector: () => showModelSelector, getValue: () => value, - hasCwdTools: () => toolsStore.hasEnabledCwdTools, - hasPrompts: () => - mcpStore.hasPromptsCapability(conversationsStore.preferences.getAllMcpServerOverrides()), + hasCwdTools: () => conversationsStore.preferences.hasEnabledCwdTools(), + // policy-aware, same rule as the agentic flow: MCP category on and at + // least one globally-enabled server whose group key is not disabled + hasPrompts: () => { + const prefs = conversationsStore.preferences; + + if (!prefs.isCategoryEnabled(ToolSource.MCP)) return false; + + return mcpStore + .getServers() + .some((s) => s.enabled && prefs.isServerToolsEnabled(s.id) && s.url.trim()); + }, openModelSelector: () => chatFormActionsRef?.openModelSelector(), setCaretOffset: (offset) => inputRef?.setCaretOffset(offset), setValue: (v) => { @@ -620,8 +628,6 @@ isReasoning={chatStore.isReasoning} {isRecording} onFileUpload={handleFileUpload} - onMcpPromptClick={showMcpPromptButton ? () => pickers.openPromptPicker() : undefined} - onMcpResourcesClick={() => (isResourceDialogOpen = true)} onMcpSettingsClick={() => (isMcpServersDialogOpen = true)} onMicClick={handleMicClick} {onStop} @@ -635,7 +641,7 @@ - {#if toolsStore.hasEnabledCwdTools} + {#if conversationsStore.preferences.hasEnabledCwdTools()} - import { File, MessageSquare, Plus } from '@lucide/svelte'; + import { File, Image, MessageSquare, Mic, Plus, Video } from '@lucide/svelte'; import { ChatFormActionAddToolsSubmenu, McpLogo } from '$lib/components/app'; import { buttonVariants } from '$lib/components/ui/button'; import * as DropdownMenu from '$lib/components/ui/dropdown-menu'; @@ -8,10 +8,10 @@ import { ATTACHMENT_FILE_ITEMS, ATTACHMENT_TOOLTIP_TEXT, - ICON_CLASS_DEFAULT, - TOOLTIP_DELAY_DURATION + ICON_CLASS_DEFAULT } from '$lib/constants'; import { getChatFormActionsContext } from '$lib/contexts'; + import { AttachmentAction, AttachmentItemEnabledWhen } from '$lib/enums'; import { useAttachmentMenu } from '$lib/hooks/use-attachment-menu.svelte'; interface Props { @@ -30,21 +30,29 @@ const attachmentMenu = useAttachmentMenu( () => ({ hasAudioModality: chatFormActions.hasAudioModality, - hasMcpPromptsSupport: chatFormActions.hasMcpPromptsSupport, - hasMcpResourcesSupport: chatFormActions.hasMcpResourcesSupport, hasVideoModality: chatFormActions.hasVideoModality, hasVisionModality: chatFormActions.hasVisionModality }), () => ({ onFileUpload: chatFormActions.onFileUpload, - onMcpPromptClick: chatFormActions.onMcpPromptClick, - onMcpResourcesClick: chatFormActions.onMcpResourcesClick, onSystemPromptClick: chatFormActions.onSystemPromptClick }), () => { dropdownOpen = false; } ); + + const FILE_MODALITY_ICONS: Record = { + [AttachmentItemEnabledWhen.HAS_AUDIO_MODALITY]: { icon: Mic, label: 'Audio' }, + [AttachmentItemEnabledWhen.HAS_VIDEO_MODALITY]: { icon: Video, label: 'Video' }, + [AttachmentItemEnabledWhen.HAS_VISION_MODALITY]: { icon: Image, label: 'Vision' } + }; + + const supportedModalities = $derived.by(() => + ATTACHMENT_FILE_ITEMS.filter((item) => attachmentMenu.isItemEnabled(item.enabledWhen)) + .map((item) => FILE_MODALITY_ICONS[item.enabledWhen ?? '']) + .filter((modality) => modality !== undefined) + );
@@ -84,50 +92,32 @@ } }} > - - - + attachmentMenu.callbacks[AttachmentAction.FILE_UPLOAD]()} + > + + Add files - - - - {#each ATTACHMENT_FILE_ITEMS as item (item.id)} - {@const enabled = attachmentMenu.isItemEnabled(item.enabledWhen)} - {#if enabled} - attachmentMenu.callbacks[item.action]()} - > - - - {item.label} - - {:else if item.disabledTooltip} - - - {#snippet child({ props })} -
- - - - {item.label} - -
- {/snippet} - - - -

{item.disabledTooltip}

- - - {/if} - {/each} - - + + {#if supportedModalities.length > 0} + + {#each supportedModalities as modality (modality.label)} + + + + + + +

{modality.label}

+ + + {/each} + + {/if} + + - import { FolderOpen, Server, Zap } from '@lucide/svelte'; - import { McpLogo } from '$lib/components/app'; - import * as DropdownMenu from '$lib/components/ui/dropdown-menu'; - import { ICON_CLASS_DEFAULT } from '$lib/constants'; - import { getChatFormActionsContext } from '$lib/contexts'; - - const chatFormActions = getChatFormActionsContext(); - - function handleServersClick() { - chatFormActions.onMcpSettingsClick?.(); - } - - - - - - - MCP - - - - - - - Servers - - - {#if chatFormActions.hasMcpPromptsSupport} - - - - Prompts - - {/if} - - {#if chatFormActions.hasMcpResourcesSupport} - - - - Resources - - {/if} - - diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddReasoningSubmenu.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddReasoningSubmenu.svelte index 1b6fc4b02096..197d6c2e5903 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddReasoningSubmenu.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddReasoningSubmenu.svelte @@ -8,70 +8,68 @@ const reasoning = useReasoningMenu(); -{#if reasoning.modelSupportsThinking} - - - {#if reasoning.thinkingEnabled} - - {:else if reasoning.isOff} - - {:else} - - {/if} + + + {#if reasoning.isReasoningActive} + + {:else if reasoning.isOff} + + {:else} + + {/if} - - Reasoning + + Reasoning - - {reasoning.currentEffort} - + + {reasoning.currentEffort} - + + - - {#each reasoning.levels as level (level.value)} - {@const tokenLabel = reasoning.tokenLabel(level)} - reasoning.select(level)} - > - {#if reasoning.isSelected(level)} - - {:else} -
- {/if} + + {#each reasoning.levels as level (level.value)} + {@const tokenLabel = reasoning.tokenLabel(level)} + reasoning.select(level)} + > + {#if reasoning.isSelected(level)} + + {:else} +
+ {/if} - {level.label} + {level.label} - {#if tokenLabel} - - {tokenLabel} - - {/if} + {#if tokenLabel} + + {tokenLabel} + + {/if} - {#if level.hasInfo} - - - - + {#if level.hasInfo} + + + + - -

Maximum reasoning effort with extended context usage

- - - {/if} - - {/each} - - -{/if} + +

Maximum reasoning effort with extended context usage

+ + + {/if} + + {/each} + + diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddSheet.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddSheet.svelte index 2f69dc96de74..acd0f4d21161 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddSheet.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddSheet.svelte @@ -1,18 +1,18 @@
@@ -194,80 +186,15 @@ - (mcpExpanded = open)} open={mcpExpanded}> - - {#if mcpExpanded} - - {:else} - - {/if} - - - - MCP Servers - - - {mcpServers.length} server{mcpServers.length !== 1 ? 's' : ''} - - - - -
- {#each mcpServers as server (server.id)} - {@const healthState = mcpStore.getHealthCheckState(server.id)} - {@const hasError = healthState.status === HealthCheckStatus.ERROR} - {@const displayName = mcpStore.getServerLabel(server)} - {@const faviconUrl = mcpStore.getServerFavicon(server.id)} - {@const isEnabled = conversationsStore.preferences.isMcpServerEnabledForChat( - server.id - )} - - - {/each} +
- - + System Message + {#if toolsPanel.totalToolCount > 0} (toolsExpanded = open)} open={toolsExpanded}> @@ -289,40 +216,12 @@
- {#each toolsPanel.activeGroups as group (group.key)} - {@const checked = toolsPanel.isGroupChecked(group)} - {@const enabledCount = toolsPanel.getEnabledToolCount(group)} - {@const favicon = toolsPanel.getFavicon(group)} - - + {#each toolsPanel.mcpGroups as group (group.key)} + {@render sheetGroupRow(group)} {/each}
@@ -331,38 +230,55 @@ - - {#if chatFormActions.hasMcpPromptsSupport} - - {/if} - - {#if chatFormActions.hasMcpResourcesSupport} - - {/if}
+ +{#snippet sheetGroupRow(group: ToolGroup)} + {@const checkState = toolsPanel.getGroupCheckState(group)} + {@const enabledCount = toolsPanel.getEnabledToolCount(group)} + {@const favicon = toolsPanel.getFavicon(group)} + {@const groupDisabled = toolsPanel.isGroupDisabled(group)} + + +{/snippet} diff --git a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddToolsSubmenu.svelte b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddToolsSubmenu.svelte index 40fed27c70a6..f49544171480 100644 --- a/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddToolsSubmenu.svelte +++ b/tools/ui/src/lib/components/app/chat/ChatForm/ChatFormActions/ChatFormActionAdd/ChatFormActionAddToolsSubmenu.svelte @@ -7,6 +7,7 @@ import { CLI_FLAGS, ICON_CLASS_DEFAULT } from '$lib/constants'; import { useToolsPanel } from '$lib/hooks/use-tools-panel.svelte'; import { mcpStore, toolsStore } from '$lib/stores'; + import type { ToolGroup } from '$lib/types'; const toolsPanel = useToolsPanel(); const hasMcpServersAvailable = $derived(mcpStore.getServers().length > 0); @@ -62,95 +63,108 @@ {/if} {:else}
- {#each toolsPanel.activeGroups as group (group.key)} - {@const isExpanded = toolsPanel.expandedGroups.has(group.key)} - {@const checked = toolsPanel.isGroupChecked(group)} - {@const favicon = toolsPanel.getFavicon(group)} - - toolsPanel.toggleGroupExpanded(group.key)} - open={isExpanded} - > -
- - {#if isExpanded} - - {:else} - - {/if} - - - {#if favicon} - { - (e.currentTarget as HTMLImageElement).style.display = 'none'; - }} - src={favicon} - /> - {/if} - - {group.label} - - - - {toolsPanel.getEnabledToolCount(group)}/{group.tools.length} - - - - - - {#snippet child({ props })} - toolsPanel.toggleGroupByKey(group.key)} - /> - {/snippet} - - - -

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

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

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

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