From e688a2d0d4f68f3d7cbdba57250da9dea8a7e982 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Mon, 6 Jul 2026 18:36:53 +0200 Subject: [PATCH 1/7] feat(ggml): add ROCmFPX quant support for DS4 --- ggml/include/ggml.h | 27 +- ggml/rocmfp4/README.md | 1020 ++++++++++++++++++++++ ggml/rocmfp4/rocmfp4.c | 797 +++++++++++++++++ ggml/rocmfp4/rocmfp4.h | 58 ++ ggml/rocmfp4/rocmfp4_hip.cu | 85 ++ ggml/rocmfp4/rocmfp4_hip_codebook.cuh | 79 ++ ggml/rocmfp4/rocmfp4_hip_scale.cuh | 157 ++++ ggml/rocmfpx/README.md | 106 +++ ggml/rocmfpx/rocmfpx.c | 1140 +++++++++++++++++++++++++ ggml/rocmfpx/rocmfpx.h | 111 +++ ggml/rocmfpx/test_rocmfpx.c | 164 ++++ ggml/src/CMakeLists.txt | 4 + ggml/src/ggml-common.h | 6 + ggml/src/ggml-cpu/ops.cpp | 78 ++ ggml/src/ggml-cuda/common.cuh | 68 ++ ggml/src/ggml-cuda/convert.cu | 110 +++ ggml/src/ggml-cuda/cpy.cu | 131 +++ ggml/src/ggml-cuda/dequantize.cuh | 123 +++ ggml/src/ggml-cuda/fattn.cu | 136 ++- ggml/src/ggml-cuda/getrows.cu | 24 + ggml/src/ggml-cuda/ggml-cuda.cu | 25 +- ggml/src/ggml-cuda/mmvf.cu | 14 +- ggml/src/ggml-cuda/mmvq.cu | 73 +- ggml/src/ggml-cuda/unary.cu | 57 ++ ggml/src/ggml-cuda/unary.cuh | 9 + ggml/src/ggml-cuda/vecdotq.cuh | 372 ++++++++ ggml/src/ggml.c | 63 +- 27 files changed, 5023 insertions(+), 14 deletions(-) create mode 100644 ggml/rocmfp4/README.md create mode 100644 ggml/rocmfp4/rocmfp4.c create mode 100644 ggml/rocmfp4/rocmfp4.h create mode 100644 ggml/rocmfp4/rocmfp4_hip.cu create mode 100644 ggml/rocmfp4/rocmfp4_hip_codebook.cuh create mode 100644 ggml/rocmfp4/rocmfp4_hip_scale.cuh create mode 100644 ggml/rocmfpx/README.md create mode 100644 ggml/rocmfpx/rocmfpx.c create mode 100644 ggml/rocmfpx/rocmfpx.h create mode 100644 ggml/rocmfpx/test_rocmfpx.c diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index df742fbc5708..800854e7d462 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -430,7 +430,13 @@ extern "C" { GGML_TYPE_NVFP4 = 40, // NVFP4 (4 blocks, E4M3 scale) GGML_TYPE_Q1_0 = 41, GGML_TYPE_TQ3_0 = 42, // TurboQuant 3.5 bpv (3-bit Lloyd-Max + FWHT rotation) - GGML_TYPE_COUNT = 43, + GGML_TYPE_Q4_0_ROCMFP4 = 100, + GGML_TYPE_Q4_0_ROCMFP4_FAST = 101, + GGML_TYPE_Q6_0_ROCMFPX = 102, + GGML_TYPE_Q8_0_ROCMFPX = 103, + GGML_TYPE_Q3_0_ROCMFPX = 104, + GGML_TYPE_Q2_0_ROCMFP2 = 107, + GGML_TYPE_COUNT = 108, }; // precision @@ -468,6 +474,18 @@ extern "C" { GGML_FTYPE_MOSTLY_MXFP4 = 25, // except 1d tensors GGML_FTYPE_MOSTLY_NVFP4 = 26, // except 1d tensors GGML_FTYPE_MOSTLY_Q1_0 = 27, // except 1d tensors + GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4 = 100, + GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_LEAN = 101, + GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_COHERENT = 102, + GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_FAST = 103, + GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_FAST_COHERENT = 104, + GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_STRIX = 105, + GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_STRIX_LEAN = 106, + GGML_FTYPE_MOSTLY_Q6_0_ROCMFPX = 110, + GGML_FTYPE_MOSTLY_Q8_0_ROCMFPX = 111, + GGML_FTYPE_MOSTLY_Q3_0_ROCMFPX = 112, + GGML_FTYPE_MOSTLY_Q2_0_ROCMFP2 = 118, + GGML_FTYPE_MOSTLY_Q2_0_ROCMFP2_STRIX = 119, }; // available tensor operations: @@ -618,6 +636,7 @@ extern "C" { GGML_GLU_OP_GEGLU, GGML_GLU_OP_SWIGLU, GGML_GLU_OP_SWIGLU_OAI, + GGML_GLU_OP_SWIGLU_DS4, GGML_GLU_OP_GEGLU_ERF, GGML_GLU_OP_GEGLU_QUICK, @@ -1339,6 +1358,12 @@ extern "C" { struct ggml_tensor * a, struct ggml_tensor * b); + GGML_API struct ggml_tensor * ggml_swiglu_ds4_split( + struct ggml_context * ctx, + struct ggml_tensor * gate, + struct ggml_tensor * up, + float clamp); + GGML_API struct ggml_tensor * ggml_geglu_erf_split( struct ggml_context * ctx, struct ggml_tensor * a, diff --git a/ggml/rocmfp4/README.md b/ggml/rocmfp4/README.md new file mode 100644 index 000000000000..39bb79302cb2 --- /dev/null +++ b/ggml/rocmfp4/README.md @@ -0,0 +1,1020 @@ +# ROCmFP4 + +ROCmFP4 is the experimental `Q4_0_ROCMFP4` GGUF tensor type for Strix Halo. +The current AMD-tuned variant stores 32 weights per block as packed +E2M1-derived 4-bit values plus two unsigned E4M3 scale bytes, one per +16-weight half, for 18 bytes per block. The codebook stores half-scale signed +integer levels up to `10`, representing `5.0` raw-scale units after the scale +factor is applied. This keeps outlier pull lower than the original wider test +range while preserving a fast integer-dot backend shape. + +This directory owns the format-specific implementation. The rest of ggml only +registers and dispatches the type. + +Current status: +- The format runs on CPU, Vulkan, and ROCm/HIP in this custom tree. +- `Q4_0_ROCMFP4` is the pure 4.50 BPW dual-scale path. +- `Q4_0_ROCMFP4_LEAN` keeps ROCmFP4 for dense tensors but protects token + embeddings with `Q5_K`. On the Qwen3-4B Strix test pass this closed most of + the Q4_0 coherence gap while staying around normal Q4 size. +- `Q4_0_ROCMFP4_COHERENT` protects token embeddings with `Q6_K` and is the + quality-first ROCmFP4 preset. +- `Q4_0_ROCMFP4_FAST` is the 4.25 BPW single-scale speed path. It is the + smallest and fastest decode variant, but the pure version gives up too much + PPL to be the default quality target. +- `Q4_0_ROCMFP4_FAST_COHERENT` combines the fast 4.25 BPW transformer layout + with `Q6_K` token embeddings. On the Qwen3-4B Strix test pass it is the + current balanced AMD target: smaller than `Q4_0`, faster than `Q4_0` on + decode on both Vulkan and ROCm, and close to `Q4_0` PPL on the short + WikiText-2 check. +- `Q4_0_ROCMFP4_STRIX` is the current quality-biased Strix Halo preset. It + keeps most transformer tensors on `Q4_0_ROCMFP4_FAST`, protects token + embeddings with `Q6_K`, and uses the dual-scale `Q4_0_ROCMFP4` layout for + attention-K and attention-V tensors. On Qwen3-4B it improved the short + WikiText-2 PPL to `13.8865` at `4.49 BPW` while still beating the + same-flags stock `Q4_0` decode baselines on both Vulkan and ROCm. +- `Q4_0_ROCMFP4_STRIX_LEAN` is the compact Strix Halo preset. It keeps the + STRIX all-layer dual-scale attention-K/V protection, uses the FAST + single-scale transformer layout for the dense tensors, and protects token + embeddings/output with `Q5_K` instead of `Q6_K`. On the Qwen3-4B validation + pass it landed at `4.38 BPW`, improved short WikiText-2 PPL versus + `FAST_COHERENT`, and kept Vulkan decode in the `81 tok/s` band. +- A smaller first/last-layer-only K/V protection recipe was tested but not + promoted. It reached `4.48 BPW`, `80.13` Vulkan decode, and `69.85` ROCm + decode, but PPL regressed to `14.0167`, so the all-layer K/V STRIX preset + remains the quality target. +- The ROCm/HIP MMQ path for `Q4_0_ROCMFP4_FAST` uses one scale per 32-weight + block, matching the actual FAST layout instead of duplicating the scale into + two half-block slots. +- The ROCm/HIP vector-dot and MMQ loaders use a ROCmFP4-owned Codebook10 + expander backed by AMD `amdgcn_perm` constants. This avoids the generic + table-load helper on the hot ROCm path. +- The ROCm/HIP hot paths use a ROCmFP4-owned unaligned 32-bit quant-byte load + for packed nibble bytes. ROCmFP4 blocks are 17 or 18 bytes wide, so the older + byte-safe assembly path was conservative but expensive on Strix Halo HIP. + Direct unaligned dword loads improved ROCm `MUL_MAT` and FlashAttention + microbenchmarks while the Qwen3.6 27B MTP guard held `33.4 tok/s` short and + `27.7 tok/s` sustained. It is enabled by default and can be disabled with + `-DGGML_ROCMFP4_UNALIGNED_QS_DWORD_LOAD=0` for isolation. +- The ROCm/HIP FAST MMVQ/MMQ path now uses the same ROCmFP4-owned unaligned + quant-byte load instead of the generic byte-assembly helper. The full + promoted gate measured FAST ROCm `MUL_MAT` at `45.17`, `58.38`, `90.54`, + and `157.83` us for `n=1/2/4/8`, and the Qwen3.6 27B MTP guard improved to + `33.6 tok/s` short and `28.0 tok/s` sustained. The ROCm runtime guard now + tightens the FAST ceilings to protect this band. +- Qwen3.6 35B A3B MTP ROCmFP4 STRIX_LEAN was checked separately at native + `262144` context on ROCm0. In the reasoning-off discovery sweep, `n-max 1` + was best sustained at `72.2 tok/s`, while `n-max 5` was best burst at + `107.3 tok/s` but slower sustained at `64.9 tok/s`. With reasoning on, a + follow-up sweep found `n-max 2` best sustained at `92.6 tok/s` short and + `80.6 tok/s` sustained; `n-max 3` was close at `104.3` short and `80.1` + sustained, while `n-max 5` stayed burst-only at `98.6` short and `73.3` + sustained. KV-cache isolation then found that q8 main KV with q4 draft KV + improved the reasoning-on `n-max 2` profile to `93.7 tok/s` short and + `85.6 tok/s` sustained; full q8 main/draft KV reached `93.2` / `85.2`, + draft-only q8 stayed at `80.5` sustained, and K-only / V-only q8 main KV + regressed to `71.8` / `73.3` sustained. Rechecking draft depth under q8 main + KV moved the best sustained profile to `n-max 3` at `104.3 tok/s` short and + `89.3 tok/s` sustained, while `n-max 4` became the best burst-only profile + at `111.2 tok/s` short and `78.7 tok/s` sustained. The Pi sustained profile + therefore uses `n-max 3` with q8 main KV and q4 draft KV, while a separate + `n-max 4` burst profile is available for short-response experimentation. + The older `n-max 5` burst alias remains available for comparison. + Batch and CPU thread follow-ups did not beat the promoted runtime shape: + `-b 1024 -ub 512`, `-b 2048 -ub 512`, `-b 512 -ub 256`, + `-t 24 -tb 32`, and `-t 12 -tb 32` all measured in the `88.7`-`89.1 tok/s` + sustained band, so `-b 512 -ub 512 -t 16 -tb 32` stays promoted. + Current KV isolation rechecked the promoted shape at `104.3 tok/s` short and + `90.1 tok/s` sustained with q8 main KV, q4 draft KV, `n-max 3`, and + `p-min 0.25`; full q8 draft KV tied at `104.5` / `90.0`, while draft-only + q8 reached only `82.2` sustained. K-only and V-only q8 main KV regressed to + `70.5` and `74.6` sustained, so both accepted K and V need q8 and the draft + KV can stay q4. The updated default 35B guard then passed at `103.3 tok/s` + short and `90.1 tok/s` sustained. + Sampler-chain trims and backend sampling did not beat the promoted sustained + profile: `top_k;top_p;temperature` reached `104.2` / `87.5`, the milder + `penalties;top_k;top_p;min_p;temperature` chain reached `104.1` / `88.9`, + and `--backend-sampling` reached `104.3` / `89.2`, so the default sampler + path stays promoted. + A single-sequence MTP `process()` fast path for `-np 1` was also prototyped, + built, and rejected because the guard dropped to `104.1 tok/s` short and + `88.5 tok/s` sustained; the change was removed. + MoE `rows_per_block` compile-time variants were checked against the same 35B + guard and also rejected: `rows_per_block=4` measured `103.8` short / + `89.1 tok/s` sustained, `rows_per_block=3` measured `86.8 tok/s` sustained, + and `rows_per_block=1` measured `103.6` / `88.7 tok/s`, none beating the + promoted `104.3` / `89.3` band. + The Pi server profile was also started and stopped successfully with + `n_ctx = 262144`, `draft-mtp` initialized, built-in tools enabled, and + `thinking = 1`; ROCm reported no KFD PIDs after shutdown. The promoted + profile is now covered by the focused + `scripts/check-rocmfp4-qwen35-a3b-mtp-regression.sh` guard, which defaults + to `n-max 3`, q8 main KV, q4 draft KV, and `p-min 0.25` with `100.0` short + and `85.0 tok/s` sustained floors. + Follow-up `p-min 0.25/0.50/0.75`, `p-split 0.05/0.20`, and `n-min 1` + checks on the q4-KV baseline all tied the same `80.5`-`80.8 tok/s` sustained + band, while `p-min 0.25` on the q8-main/q4-draft path tied at `85.4 tok/s` + sustained, so the + conservative `n-min 0`, `p-min 0.0`, `p-split 0.10` defaults remained + promoted until the MTP internal sampler was changed from `top_k=1` to + `top_k=10`, making the `p-min` cutoff operate on a meaningful candidate + distribution while the draft loop still selects the top sorted candidate. + A pre-change p-min sweep on the actual promoted `n-max 3`, q8-main/q4-draft + profile measured `104.2` / `89.1` at `p-min 0.25`, `104.1` / `89.0` at + `0.50`, `104.0` / `89.3` at `0.75`, and `104.2` / `89.1` at `0.90`. + Follow-up p-split checks measured `103.7` / `88.9` at `0.05` and + `103.8` / `89.3` at `0.20`; n-min checks measured `103.8` / `88.9` for + `1` and `104.2` / `89.2` for `2`. None beat the promoted + sustained-plus-short profile. Combined `p-min`/`p-split` follow-ups also + tied rather than beat the default: `p-min 0.25` with `p-split 0.30` + repeated at `104.6` / `89.5`, `p-split 0.40` and `0.50` measured + `104.6` / `89.5` and `104.4` / `89.5`, and `p-split 0.70` / `0.90` + measured `104.8` / `89.5` and `104.7` / `89.5`; the same-session default + tied at `104.7` / `89.5`. With `top_k=10`, the 35B A3B profile now promotes + `--spec-draft-p-min 0.25`: default `p-min 0.0` measured `104.3` / `89.3`, + while `p-min 0.25` measured `103.9` / `90.0` after a first sustained run at + `89.8 tok/s`. The dense 27B profile should stay at `p-min 0.0`; the same + `p-min 0.25` filter regressed it to `24.6 tok/s` sustained. The full + all-regression harness can include the 35B guard with + `INCLUDE_QWEN35_A3B_GUARD=1`. + The MTP draft sampler now uses a probability-only top-10 helper for this + path. The draft loop only consumes the sorted top candidate and its + probability, so skipping the unused final RNG sampler selection preserved the + output path while improving the 35B A3B guard to `104.6 tok/s` short and + `90.2 tok/s` sustained. The dense 27B guard remained stable at `33.9` / + `28.1 tok/s`. + A fixed insertion top-10 replacement for `std::partial_sort` was tested and + rejected after the 35B A3B short guard dropped to `73.0 tok/s`, below the + `100.0 tok/s` floor. A `std::nth_element` plus top-k slice sort variant was + also rejected after the same guard dropped to `71.5 tok/s`. + Narrowing the MTP top-k probability accumulator from `double` to `float` was + also rejected: dense 27B still passed at `34.0` / `28.1 tok/s`, but the 35B + A3B short guard repeated below floor at `84.1` then `93.4 tok/s`; restoring + the `double` accumulator recovered the short guard to `104.4 tok/s`. + A `std::partial_sort_copy` top-k buffer variant built and passed the same + guard at `104.4 tok/s` short and `90.2 tok/s` sustained, but a same-session + promoted-build comparison measured `104.1` / `90.2`, so it was not promoted: + it added sampler buffer/clone complexity without a sustained decode gain. + Jackrong Qwopus3.6 27B v2 MTP BF16 was converted to STRIX_LEAN ROCmFP4 at + `4.34 BPW`. At native `262144` context with reasoning on, q4 main/draft KV, + and `n-max 4`, ROCm0 initially measured `34.9` / `29.6 tok/s`; increasing + only the batch to `-b 1024 -ub 512` moved sustained decode to `29.9 tok/s`. + Follow-up batch shapes `1280/512`, `1536/512`, and `1536/768` measured + `29.8`, `29.9`, and `29.8 tok/s` sustained respectively, so the smaller + `1024/512` profile remains promoted. Light acceptance filters `p-min 0.05` + and `0.10` also tied at `29.9 tok/s` sustained without beating the default. + Lower draft-depth checks with the same promoted batch rejected `n-max 1`, + `2`, and `3`, which measured only `19.9`, `26.6`, and `27.3 tok/s` + sustained. Qwopus therefore stays on `n-max 4`, unlike the 35B A3B profile + where `n-max 3` plus q8 main KV is best. + `--backend-sampling` tied sustained decode at `29.9 tok/s` but lowered prompt + throughput, so it is not promoted. Thread split checks at target/draft + `12/32` and `24/32` also tied `29.9 tok/s` sustained; the simpler default + `16/32` thread shape remains the recommended Qwopus profile. + KV isolation confirmed this is not like the 35B A3B q8-main profile: + draft-only q8 KV measured `35.1` / `29.8 tok/s`, and full q8 main plus q8 + draft KV measured `36.6` / `26.0 tok/s`. The full-q8 path improves burst + only and regresses sustained decode. Split accepted-KV checks also regressed: + q8 K only measured `35.7` / `22.3 tok/s`, and q8 V only measured `34.2` / + `24.7 tok/s`. Qwopus therefore keeps q4 main and q4 draft KV. + Vulkan0 measured `40.0` / `27.7 tok/s`, with Vulkan `n-max 3` and `n-max 5` + slower sustained; Vulkan `-b 1024 -ub 512` repeated the same `27.7 tok/s` + sustained band. q8 main KV regressed this model on both ROCm0 and Vulkan0. + A normal-path shortcut that only normalized `data[0].p` and filled the rest + of the top-k probabilities only for debug logging was rejected after the + 35B A3B short guard fell to `92.7 tok/s`, below the `100.0 tok/s` floor. + Disabling internal sampler timing with `sparams.no_perf = true` was also + tested and rejected after the same short guard dropped to `96.2 tok/s`. + Skipping the per-draft `common_sampler_reset()` in the MTP path was rejected + after the 35B A3B short guard dropped to `68.3 tok/s`; that reset remains + required to preserve the expected sampler/logit state for this helper. + A small MTP host-path cleanup now delays `llama_get_embeddings_pre_norm_ith()` + until the draft loop has confirmed that another draft token will be queued. + This avoids unused embedding-row pointer fetches on p-min rejects and final + `n-max` draft tokens. It passed the dense 27B guard at `34.0` / `28.1 tok/s`, + the 35B A3B guard at `104.4` / `90.1 tok/s`, and the Qwopus best ROCm0 + profile at `35.0` / `29.8 tok/s`. The default all-regression gate also + passed after this cleanup and ended with no KFD PIDs running. The 2026-05-25 + serial pass measured Qwen3.6 27B MTP at `33.9` / `27.9 tok/s`, ROCm runtime + FAST `45.66` / `57.81` / `88.27` / `155.05` us and dual-scale `49.16` / + `51.58` / `83.34` / `151.42` us for `n=1/2/4/8`, ROCm FlashAttention + `70.86` / `66.51` us for 64d dual-scale / FAST and `189.45` / `172.73` us + for Qwen-style 128d dual-scale / FAST, and ROCm CPY source-to-dual + `1106.89` / `1008.56` / `1006.60` us with source-to-FAST `1050.49` / + `958.98` / `950.50` us for F32/F16/BF16. + A second attempt to skip the final `n-max` `common_sampler_accept()` call was + rejected: the 35B A3B short check still reached `104.5 tok/s`, but sustained + decode fell to `81.1 tok/s`, below the `85.0 tok/s` floor. Reverting only + that sampler-accept change recovered the guard to `104.3` / `90.0 tok/s`. + Retesting the internal MTP sampler candidate count on the 35B A3B + reasoning-on profile rejected both directions around the promoted `top_k=10` + setting: `top_k=5` fell to `77.3 tok/s` sustained and `top_k=20` fell to + `69.6 tok/s` sustained. +- Reasoning-off checks on the final 35B q8-main/q4-draft profile are a + separate lower-throughput mode. `n-max 1/2/3/4` measured `77.7` / `73.9`, + `90.3` / `75.5`, `100.3` / `71.9`, and `85.7` / `66.1` short/sustained + tok/s respectively. If reasoning is disabled, `n-max 2` is currently the + best sustained profile in this bracket; the promoted fastest profile remains + reasoning-on `n-max 3`. +- Draft-thread-only checks on the promoted 35B profile also tied below the + promoted band. Keeping target threads at `16/32`, draft `8/16`, `16/16`, + and `24/32` all measured around `104.1`-`104.2 tok/s` short and + `89.1 tok/s` sustained, so the default matching draft thread counts remain. +- ROCm/HIP single-token MMVQ uses a full-block vector-dot ratio for the + dual-scale layout (`VDR_ROCMFP4_Q8_1_MMVQ=4`) while keeping the FAST layout + on the previous half-block ratio (`VDR_ROCMFP4_FAST_Q8_1_MMVQ=2`). This + lets dual-scale ROCmFP4 consume one full 32-value block per vector-dot call + without slowing the FAST dense-tensor path used by STRIX_LEAN. The focused + ROCm0 `MUL_MAT` guard improved dual-scale from the prior `78.81` us/run + serial pass to `54.89` us/run, and the Qwen3.6 27B MTP sustained guard + improved from `26.2` to `27.8 tok/s`. A broader version that also moved + FAST to `vdr=4` was rejected because sustained Qwen MTP dropped to + `24.2 tok/s`. A later FAST-only retest after the packed-byte load improved + the focused FAST ROCm guard to `41.37`, `49.29`, `80.91`, and `139.58` + us for `n=1/2/4/8`, but was still rejected because sustained Qwen MTP + dropped to `24.7 tok/s`. The remaining narrower FAST setting, + `GGML_ROCMFP4_FAST_Q8_1_MMVQ_VDR=1`, was also rejected after it failed the + focused FAST ROCm `n=1` guard at `60.18` us/run. The knob now rejects + invalid FAST MMVQ values at compile time; only `1`, `2`, and `4` are valid. +- A direct ROCmFP4 `vec_dot_q_cuda_dispatch` wrapper was tested in the + MMVQ kernels to bypass the generic constexpr function pointer call. It built + and passed the focused ROCm guard, measuring FAST `45.16` / `57.52` / + `89.44` / `156.49` us and dual-scale `50.83` / `51.11` / `84.51` / + `143.27` us for `n=1/2/4/8`, but Qwen3.6 27B MTP only reached `33.7 tok/s` + short and `27.9 tok/s` sustained. Because it did not beat the promoted + sustained band, the code change was removed. +- ROCm/HIP batched MMQ keeps the upstream-style `vdr=8` default for both + ROCmFP4 layouts, with ROCmFP4-owned compile-time test knobs + `GGML_ROCMFP4_Q8_1_MMQ_VDR` and `GGML_ROCMFP4_FAST_Q8_1_MMQ_VDR`. + FAST-only `vdr=4`, FAST-only `vdr=16`, and dual-scale `vdr=16` were tested + against the Qwen3.6 27B STRIX_LEAN ROCmFP4 bench and did not improve decode; + the tested runs all stayed at `13.56 tok/s` generation, while dual-scale + `vdr=16` slightly reduced prompt throughput. The default remains `8`. +- ROCm/HIP RDNA3.5 small-batch MMVQ now uses a ROCmFP4-specific two-warp + launch geometry for the ROCmFP4 layouts through `n=2`. Strix Halo previously + inherited the older RDNA2 one-warp table for this path. The promoted + `GGML_ROCMFP4_RDNA35_NWARPS=2` and + `GGML_ROCMFP4_RDNA35_NWARPS_MAX_NCOLS=2` defaults keep the single-token + microbench win and improve the guarded ROCm0 `n=2` shape: the latest serial + pass measured FAST/dual `66.56` / `58.40` us/run for `n=2`, versus the prior + `68.85` / `60.98` us/run band before extending the two-warp route. The + Qwen3.6 27B MTP guard held `33.5 tok/s` short and `27.7 tok/s` sustained, so + this is promoted as a backend micro-optimization with no sustained decode + regression. `GGML_ROCMFP4_RDNA35_NWARPS=4` was rejected because dual-scale + regressed to `57.87` us/run. The remaining 8-warp candidate was also rejected + because FAST `n=1` regressed to `59.33` us/run and failed the focused ROCm + runtime guard before multi-column checks. Extending the promoted two-warp + launch from `n=1` to `n=1..4` improved some focused multi-column ROCm + microbench rows (`dual n=2` reached `57.14` us/run), but Qwen3.6 27B MTP + sustained decode fell to `23.6 tok/s`, so the promoted upper bound stops at + `n=2`. A midpoint `GGML_ROCMFP4_RDNA35_NWARPS_MAX_NCOLS=3` build was also + checked on the Qwen3.6 35B A3B reasoning-on 262k profile; it reached + `87.5 tok/s` sustained versus `89.6 tok/s` for the same-session promoted + build, so `n=2` remains the default. +- ROCm/HIP RDNA3.5 wide-column rows-per-block now has ROCmFP4-owned compile-time + test knobs: `GGML_ROCMFP4_RDNA35_RPB_WIDE`, + `GGML_ROCMFP4_RDNA35_RPB_WIDE_DUAL`, and + `GGML_ROCMFP4_RDNA35_RPB_WIDE_FAST`. The defaults stay at `1`. A full + `RPB_WIDE=2` test improved FAST `n=8` `MUL_MAT` from `167.90` to + `131.72` us/run, but dual-scale collapsed from `148.11` to `1382.72` + us/run. A FAST-only `RPB_WIDE_FAST=2` build kept dual-scale safe + (`145.75` us/run) and improved FAST to `135.31` us/run, but Qwen3.6 + 27B MTP did not improve (`33.3 tok/s` short and `27.6 tok/s` sustained + at `n-max 4`; `n-max 5` remained burst-only at `45.5` short and `24.9` + sustained). Keep the knob off by default until a real decode guard benefits. +- ROCm/HIP RDNA3.5 `MUL_MAT_ID` routing now has a ROCmFP4-only compile-time + guard knob, `GGML_ROCMFP4_RDNA35_MMID_MAX_BATCH`, for testing whether MTP/MoE + batches should leave MMVQ earlier. The default keeps the accepted generic + RDNA3 behavior (`MMVQ_MAX_BATCH_SIZE`). A Strix test with the knob set to `1` + tied the promoted Qwen3.6 27B MTP path at `33.4 tok/s` short and `27.7 tok/s` + sustained. Follow-up threshold `2`, `3`, and `4` builds also tied at + `33.6` / `33.8` / `33.7 tok/s` short and `27.7` / `27.9` / `27.9 tok/s` + sustained. The same `3` threshold was then checked on the MoE-heavy + Qwen3.6 35B A3B MTP ROCmFP4 path and regressed to `95.8 tok/s` short and + `74.0 tok/s` sustained versus the promoted `104.4` / `89.3` band. Reusing + the older threshold `4` exploratory build on the same 35B guard measured + `104.2` / `89.2`, a near tie but still below the sustained promoted band. + No lower routing threshold is promoted. +- Target/draft "dual-stream" MTP overlap was inspected and is not a promoted + optimization. The current common MTP implementation verifies with the target + context, mirrors pre-norm embeddings into the MTP context, then drafts through + serial `llama_decode(ctx_dft, ...)` calls. ggml exposes async graph execution + and pipeline parallelism, but llama enables pipeline parallelism only for + multi-device layer-split offload cases. On single-device Strix Halo ROCm0, + target/draft overlap would require speculative scheduler changes and new + correctness guards rather than a runtime flag. The long-context optimization + focus remains KV-cache traffic, FlashAttention, memory bandwidth, and MTP + acceptance, while dual-scale ROCmFP4 mainly protects coherence. +- The MTP host loop now avoids two small sources of scheduler overhead. It + reuses a per-sequence drafting-state buffer instead of allocating a + `std::vector` on every draft call, and it tracks verify-batch sequence + bounds with one pass over token seq-ids instead of scanning every sequence for + every token. It also avoids copying all target verify rows when no previous + draft is pending, keeping only the carryover hidden row needed for the next + MTP step. The MTP path also uses a single-sequence batch append helper + instead of constructing a temporary sequence vector for each `common_batch_add` + call. It stores only non-final verify rows for partial-accept rollback, + copying the final row directly to `pending_h` for the next MTP step. It also + skips the draft-model decode call entirely when no sequence is drafting. + Debug candidate token formatting is skipped unless debug logging is enabled, + and the debug verbosity state is hoisted once per draft call instead of + rechecked inside the per-token loop. This is mainly a cleanup for + multi-sequence and long-running MTP sessions; the latest single-sequence + Qwen3.6 27B guard stayed in range at `33.7 tok/s` short and `27.9 tok/s` + sustained. +- The simple draft path also reuses an object-owned `uint8_t` drafting-state + buffer instead of allocating `std::vector` on every draft call. This + keeps both draft implementations on the same host-side allocation pattern; + the full serial gate stayed clean, with Qwen3.6 27B MTP at `33.8 tok/s` + short and `27.9 tok/s` sustained. +- The all-in-one regression harness now accepts a candidate `BUILD_DIR` and + derives separate `TEST_BACKEND_OPS_BIN` and `LLAMA_CLI_BIN` paths, so backend + microbench guards and Qwen CLI guards can run from the same candidate build. + DeepSeek remains opt-in only through `INCLUDE_DEEPSEEK_SMOKE=1`. +- ROCm/HIP vector FlashAttention now uses a ROCmFP4-only RDNA K/Q thread-group + default of `1` instead of the generic quantized-K RDNA default of `2`. This + affects only `Q4_0_ROCMFP4` and `Q4_0_ROCMFP4_FAST` K-cache cases; other + quantized FA paths keep the existing backend default. On the guarded ROCm0 + `FLASH_ATTN_EXT` shape, dual-scale improved from `122.33` to `113.03` + us/run, and FAST improved from `115.41` to `109.33` us/run. The Qwen3.6 + 27B MTP guard remained in the promoted band at `33.5 tok/s` short and + `27.6 tok/s` sustained. A wider `GGML_ROCMFP4_FATTN_KQ_NTHREADS=4` variant + was rejected because it regressed the same FA guard to `136.79` us/run + dual-scale and `124.37` us/run FAST. After the V-side default moved to `2`, + retesting `KQ_NTHREADS=2,V_NTHREADS=2` and `KQ_NTHREADS=4,V_NTHREADS=2` + still regressed the focused FA guard, so K/Q remains `1`. +- ROCm/HIP vector FlashAttention also uses a ROCmFP4-only V thread-group + default of `2`, down from the generic `D/4` path used by other quantized V + types. On the same guarded ROCm0 `FLASH_ATTN_EXT` shape, this moved + dual-scale from `113.03` to `85.63` us/run and FAST from `109.33` to + `80.74` us/run. The Qwen3.6 27B MTP guard stayed in range at `33.4 tok/s` + short and `27.6 tok/s` sustained, so this is promoted as the new default. + `GGML_ROCMFP4_FATTN_V_NTHREADS=4` also passed but was slower on the focused + FA guard; `1` was rejected because gfx1151 HIP compilation exceeded the 64 + KiB local-memory limit for ROCmFP4 FA instances. The FA guard now also + includes a Qwen3.6-style 128-head-dim, 8-KV-head, 12x-GQA ROCmFP4 shape + (`hsk=128,hsv=128,nh=8,nr23=[12,1],kv=7680,nb=1`). The accepted default + measured `246.91` us dual-scale and `219.23` us FAST on that shape. + Retesting `V_NTHREADS=4` on this wider shape produced a tiny dual-scale + improvement (`244.01` us) but regressed FAST (`224.21` us) and the existing + 64d guard, so it remains rejected. Retesting `KQ_NTHREADS=2,V_NTHREADS=2` + regressed the Qwen-style shape to `268.53` us dual-scale and `226.32` us + FAST, so K/Q remains `1`. +- ROCm/HIP vector FlashAttention now uses a ROCmFP4-owned single-half + Codebook10 expander in K/Q and V decode paths. These FA call sites already + know whether they need the low or high nibble stream, so they no longer pay + to expand both streams and discard one. The focused ROCm0 FA guard improved + from `86.13` / `81.23` us to `82.24` / `78.15` us for the 64d dual-scale / + FAST shapes, and from `247.06` / `221.19` us to `237.12` / `206.50` us for + the Qwen-style 128d shapes. The Qwen3.6 27B MTP guard held at `33.5 tok/s` + short and `27.7 tok/s` sustained in the full serial gate, so this is + promoted. +- ROCm/HIP vector FlashAttention also specializes the ROCmFP4 K/Q path when + `GGML_ROCMFP4_FATTN_KQ_NTHREADS=1`. Each thread owns the full head dot in + this promoted setting, so the kernel now expands both low/high packed + Codebook10 streams once per ROCmFP4 block and accumulates both half-blocks + together instead of loading the same packed bytes twice. The full serial gate + measured 64d FlashAttention at `81.62` / `78.13` us for dual-scale / FAST, + and the wider Qwen-style 128d guard improved from `237.12` / `206.50` us to + `228.58` / `199.32` us. Qwen3.6 27B MTP held `33.2 tok/s` short and + `27.7 tok/s` sustained, so this is promoted as a real Qwen-relevant FA + micro-optimization. + A post-specialization retest of `GGML_ROCMFP4_FATTN_V_NTHREADS=4` was + rejected. It passed the tightened FA guard and nudged the Qwen-style + dual-scale microbench to `228.33` us, but regressed 64d dual-scale to + `86.30` us and Qwen-style FAST to `201.81` us. Qwen3.6 27B MTP stayed at + `27.7 tok/s` sustained, so the promoted V-side default remains `2`. +- ROCm/HIP vector FlashAttention now uses a ROCmFP4-only V dequant rows-per- + thread default of `8`. The ROCmFP4 V helper supports 8-value chunks, and the + post-unaligned-load retest produced a real FA win: 64d dual-scale / FAST + `68.82` / `66.32` us and Qwen-style 128d dual-scale / FAST `201.13` / + `172.22` us. Qwen3.6 27B MTP held `33.5 tok/s` short and `27.7 tok/s` + sustained, so this is now promoted. It can be isolated with + `-DGGML_ROCMFP4_FATTN_V_ROWS_PER_THREAD=4` to return to the prior setting. + Retesting `GGML_ROCMFP4_FATTN_V_NTHREADS=4` on top of the 8-row default + improved only the Qwen-style dual-scale FA microbench (`187.72` us) while + regressing 64d dual-scale / FAST to `77.35` / `77.25` us and Qwen-style FAST + to `182.56` us. Qwen3.6 27B MTP tied at `33.5` / `27.7 tok/s`, so the + shipped V thread-group default remains `2`. + Retesting `GGML_ROCMFP4_FATTN_V_NTHREADS=8` on top of the 8-row default + regressed every focused FA row, including Qwen-style 128d dual-scale / FAST + at `208.94` / `200.45` us, so it was rejected without a longer MTP run. + Retesting `GGML_ROCMFP4_FATTN_KQ_NTHREADS=2` on top of the 8-row default + also regressed the focused FA guard: 64d dual-scale / FAST measured + `80.68` / `78.27` us and Qwen-style 128d measured `245.84` / `207.04` us. + Because the Qwen-style dual-scale row failed the guard, it was rejected + without a longer MTP run and K/Q remains `1`. +- ROCm/HIP vector FlashAttention now has a narrow dual-scale 128d V-thread + specialization: `GGML_ROCMFP4_FATTN_V_NTHREADS_D128_DUAL=4`. This keeps the + promoted 64d default and all FAST paths on `V_NTHREADS=2`, but uses the + previously promising 4-thread V grouping only for the Qwen-style dual-scale + 128d shape. The full serial gate measured FA at `70.37` / `66.31` us for + 64d dual-scale / FAST and `188.75` / `171.89` us for Qwen-style 128d + dual-scale / FAST, while Qwen3.6 27B MTP held `33.7 tok/s` short and + `27.9 tok/s` sustained with no KFD PIDs left running. + Retesting `GGML_ROCMFP4_FATTN_V_NTHREADS_D128_DUAL=8` was rejected. It + slightly improved the 64d dual-scale row (`69.40` us versus a same-session + `70.92` us promoted build), but regressed the Qwen-style 128d dual-scale row + from `194.40` to `211.10` us/run. Because long-context Qwen is the relevant + guard, the promoted D128 dual-scale V grouping remains `4`. + `GGML_ROCMFP4_FATTN_V_ROWS_PER_THREAD=16` was rejected during compilation: + the ROCmFP4 V dequantizer intentionally supports only `2`, `4`, and `8` + rows per thread, and the fixed-copy helper does not support the resulting + 32-byte move. The source now rejects unsupported values with a direct + ROCmFP4 compile-time error. + Retesting `GGML_ROCMFP4_FATTN_V_ROWS_PER_THREAD=2` on the promoted + D128-specialized build was rejected by the focused ROCm FlashAttention + guard. It regressed 64d dual-scale / FAST to `94.31` / `88.77` us and + Qwen-style 128d dual-scale / FAST to `270.00` / `223.88` us, so the + promoted rows-per-thread default remains `8`. +- Vulkan ROCmFP4 scale decode now uses a shared UE4M3 lookup table with + ROCmFP4's half-scale semantics. This moved the focused Vulkan dual-scale + `MUL_MAT` guard from `82.86`, `120.77`, and `181.28` us/run to `65.05`, + `83.07`, and `122.70` us/run for `n=1/2/4`. On Qwen3.6 27B MTP at 262k + context, the same change moved Vulkan sustained output from the older + `20.4 tok/s` at `--spec-draft-n-max 4` to `25.0 tok/s`; `n-max 3` now + reaches `25.3 tok/s` sustained. ROCm0 remains the promoted backend because + it still holds `33.5 tok/s` short and `27.7 tok/s` sustained. + A post-LUT Vulkan runtime sweep found a better fallback profile with q8 KV + and q8 draft KV: `--spec-draft-n-max 4` reached `34.8 tok/s` short and + `27.0 tok/s` sustained. q8 with `n-max 5` improved only the short burst + (`47.8 tok/s`) and regressed sustained output to `23.0 tok/s`; f16 KV + regressed sustained output to `22.5 tok/s`. + Follow-up isolation showed the Vulkan sustained gain comes from the main KV + cache rather than the draft KV cache: q4 main KV with q8 draft KV reached + only `34.6 tok/s` short and `25.0 tok/s` sustained, while q8 main KV with + q4 draft KV reached `34.7 tok/s` short and `26.9 tok/s` sustained. For a + Vulkan fallback, q8 main KV plus q4 draft KV is therefore the leaner near-tie + profile. Splitting q8 across only one accepted KV side was also rejected: + q8 main K with q4 main V reached `34.8 tok/s` short and `25.4 tok/s` + sustained, while q4 main K with q8 main V reached `34.6 tok/s` short and + `23.7 tok/s` sustained. The Vulkan fallback needs both accepted K and V at + q8. Adding mild acceptance filtering to the lean fallback + (`--spec-draft-p-min 0.25`) tied full q8/q8 at `27.0 tok/s` sustained while + keeping draft KV at q4; adding `--spec-draft-n-min 1`, stricter `p-min 0.75`, + and `p-split` checks stayed at `26.9 tok/s`. Retrying `n-max 5` with q8 main + KV, q4 draft KV, and + `p-min 0.25` still produced only a burst win (`47.8 tok/s` short) while + sustained output stayed at `23.0 tok/s`, so `n-max 4` remains the Vulkan + fallback setting. ROCm0 q4/q4 remains the overall promoted backend/profile. + Replacing the shared Codebook10 table with inline integer decode was tested + and rejected because it regressed the focused Vulkan FAST `n=1` guard to + `99.43` us/run. +- ROCm/HIP fallback dequant, copy, get-rows, GPU-side quantization scoring, + and standalone dequant helpers use ROCmFP4-owned HIP helpers for finite + scales and Codebook10 nibbles instead of relying on generic FP8 handling. + This keeps non-MMQ conversion paths aligned with the custom AMD format. +- A branchless HIP scalar Codebook10 nibble decoder was tested and rejected. + It passed the focused ROCm CPY, FlashAttention, and Qwen MTP guards, but + did not improve end-to-end decode. Qwen3.6 27B MTP measured `33.4 tok/s` + short and `27.7 tok/s` sustained, and the focused FlashAttention guard + slowed to `86.38` / `81.13` us for dual-scale / FAST, so the original + scalar decoder remains the default. +- CPU-side Codebook10 table decode in the quantizer/dequantizer was tested + and rejected. The full table variant failed the CPU quant guard with + dequantization at `49.04` cycles/32 for dual-scale and `84.23` cycles/32 for + FAST. The MSE-loop-only variant passed but slowed normal quantization to + `4183.33` / `4018.68` cycles/32 versus the restored arithmetic baseline at + `4034.33` / `3738.85` cycles/32, so the arithmetic CPU decoder remains the + default. +- UE4M3 scale decode in the ROCm/HIP software path uses a ROCmFP4-owned finite + scale decoder. It avoids `ldexpf`, builds normal FP32 values directly from + exponent/mantissa bits, and skips the generic FP8 NaN handling because + ROCmFP4 row validation already rejects non-finite scale bytes. +- A constant-memory HIP lookup table for the finite UE4M3 scale values was + tested and rejected. The latest isolated `GGML_ROCMFP4_USE_SCALE_LUT=1` + pass failed the focused ROCm runtime guard because FAST `n=1` regressed to + `69.48` us/run; prior checks also regressed ROCm CPY source-to-quant paths + and FlashAttention. The arithmetic finite-scale decoder remains the default. +- ROCm/HIP dequant conversion kernels use the same ROCmFP4 finite scale + decoder, keeping tensor conversion aligned with the hot MMQ/MMVQ backend + path instead of falling back to the generic FP8 helper. +- ROCm/HIP backend CPY now advertises and executes quantized + `Q4_0_ROCMFP4 -> F32` and `Q4_0_ROCMFP4_FAST -> F32` conversion paths. + This keeps diagnostic graph ops and fallback tensor conversion inside the + custom AMD decoder instead of being rejected by backend capability checks. + Contiguous q-to-f32 copies now use a ROCmFP4-specific packed-byte kernel: + one HIP thread reads one packed byte and writes the matching low/high + half-block output values. This gives coalesced output writes for normal + contiguous graph copies while the existing stride-aware block kernel remains + the fallback for views and non-contiguous tensors. +- ROCm/HIP backend CPY also supports `F16 -> Q4_0_ROCMFP4`, + `F16 -> Q4_0_ROCMFP4_FAST`, `BF16 -> Q4_0_ROCMFP4`, and + `BF16 -> Q4_0_ROCMFP4_FAST`. The kernels convert each 32-value half/bfloat + block to local FP32 and then run the same exhaustive ROCmFP4 scale search, + so runtime graph copies keep the coherence-first quantizer instead of + falling back to unsupported behavior. +- ROCm/HIP ROCmFP4 source-to-quant and quant-to-F32 CPY wrappers now use + normal multi-thread HIP launch geometry instead of launching one active + thread per quant block. This keeps F32/F16/BF16 runtime quantization and + ROCmFP4 dequantization on the same exact conversion math while removing a + launch-shape bottleneck in helper/fallback graph paths. +- ROCm/HIP ROCmFP4 quant-to-F32 CPY uses ROCmFP4-specific block dequant + helpers that decode each block scale once before unpacking all 32 values. + This avoids repeatedly decoding identical scale bytes through the generic + two-value dequant helper while preserving the same output values. For the + guarded contiguous `8192x512x2` ROCm0 CPY shape, the packed-byte path moved + dual-scale `Q4_0_ROCMFP4 -> F32` from the old `740` us/run band to + `184.48` us/run and FAST `Q4_0_ROCMFP4_FAST -> F32` to `169.99` us/run. + A split launch geometry is now used: source-to-ROCmFP4 quantization uses + 128-thread workgroups, while quant-to-F32 keeps the accepted 64-thread + packed-byte launch. A whole-path 128-thread launch was roughly tied overall + but slightly regressed quant-to-F32; isolating 128 threads to source + quantization preserves the dequant copy win while shaving the source paths. + A 256-thread launch regressed F16 source-to-quant paths and remains rejected. + After the FAST direct-value scoring win, FAST-only 256-thread and 64-thread + source launch splits were rechecked and rejected as well. The 256-thread + split regressed FAST F16 source-to-quant to `1047.24` us/run, and the + 64-thread split measured FAST F32/F16/BF16 `1055.32`, `955.35`, and + `954.78` us/run, so the shared 128-thread source launch remains promoted for + both ROCmFP4 layouts. +- ROCm/HIP contiguous quant-to-F32 CPY shared-scale staging was tested and + rejected. Decoding one block scale into shared memory and synchronizing each + 64-thread launch reduced duplicate scale decode work, but the synchronization + overhead regressed the guarded `8192x512x2` shape: dual-scale + `Q4_0_ROCMFP4 -> F32` moved from `181.84` to `188.63` us/run, and FAST + `Q4_0_ROCMFP4_FAST -> F32` moved from `170.65` to `180.03` us/run. The + direct packed-byte kernel remains promoted. +- ROCm/HIP backend CPY supports same-type packed-block copies for + `Q4_0_ROCMFP4 -> Q4_0_ROCMFP4` and + `Q4_0_ROCMFP4_FAST -> Q4_0_ROCMFP4_FAST`, including block-aligned views. + The kernel copies the packed 18-byte or 17-byte ROCmFP4 blocks directly, so + graph/view copies preserve exact bytes and avoid dequantize/requantize + fallback behavior. The launcher uses normal multi-thread HIP workgroups + rather than one-thread launches, so large packed-view copies scale with the + number of ROCmFP4 blocks. +- ROCm/HIP runtime quantization now finds the finite UE4M3 candidate nearest + `max_abs / 10` with a monotonic binary search, matching the CPU/Vulkan + reference tie behavior while avoiding the older 126-step linear nearest-scale + scan in the HIP copy utility. It also uses the same conservative lower-scale + pruning as the CPU quantizer, so runtime F32/F16/BF16-to-ROCmFP4 copies skip + smaller scales once their unavoidable clipped max-value error cannot beat the + current best exact scale. The HIP max scan uses the same plain absolute + compare as the CPU quantizer, preserving NaN/Inf handling while avoiding + `fmaxf` in the runtime quantization hot loop. +- ROCm/HIP runtime quantization specializes the exact scale search for the + only block shapes used by ROCmFP4 CPY: `0..15`, `16..31`, and `0..31`. + This keeps the same candidate order and tie behavior while letting HIP + compile the dual-scale and FAST source-to-quant paths with fixed offsets. +- ROCm/HIP dual-scale source-to-quant scoring uses a direct Codebook10 value + helper for 16-value half blocks. Final packed quantization still uses the + normal code-index helper, but scale scoring no longer has to map the selected + nibble back through the generic decode helper on the dual-scale path. +- ROCm/HIP FAST source-to-quant scoring now uses the same direct Codebook10 + value helper during scale search. This preserves the exact nearest-value + thresholds and final packed output, but avoids index-then-decode work in the + 32-value FAST scoring path. On the guarded `8192x512x2` ROCm0 CPY shape, + FAST source quantization improved from F32/F16/BF16 `1218.69`, `1138.32`, + and `1138.78` us/run to a repeat guard pass of `1047.21`, `950.93`, and + `951.00` us/run. +- ROCm/HIP `F32 -> ROCmFP4` runtime CPY stages each 32-value block into local + FP32 before running the exhaustive scale search, matching the existing + F16/BF16 source paths. This preserves exact F32 input values but avoids + rereading the global source pointer throughout scale scoring. On the guarded + `8192x512x2` ROCm0 CPY shape, dual-scale F32 source quantization improved + from roughly `9916` to `1117` us/run, and FAST improved from roughly + `10671` to `1231` us/run. +- ROCm/HIP contiguous-only source-to-quant CPY kernels for + `F32 -> ROCmFP4` were tested and rejected. They avoided the generic + multidimensional offset math, but on the guarded `8192x512x2` shape they + did not beat the existing multi-thread generic CPY path and slightly + regressed `F32 -> Q4_0_ROCMFP4_FAST`, so source-to-quant ROCmFP4 CPY keeps + the same guarded implementation for view and contiguous tensors. +- ROCm/HIP backend GET_ROWS supports both ROCmFP4 layouts. This gives pure + ROCmFP4 tensors the same direct row-gather coverage as stock small-block + quants on ROCm and keeps embedding-row access on the custom finite-scale + decoder. +- ROCm/HIP `MUL_MAT` support now covers `F16` activation tensors for both + ROCmFP4 layouts. The backend stages half activations to contiguous FP32 on + the GPU, including non-contiguous/views, then feeds the existing Q8 + activation quantizer and ROCmFP4 MMVQ/MMQ kernels. This keeps the forward + path on the AMD backend instead of rejecting the op and falling through to a + slower dequantized matrix path. The generic matmul runtime guard explicitly + allows this ROCmFP4 x F16 forward-inference case, so the support probe and + execution wrapper agree for batched activations. +- The standalone HIP dequant skeleton covers both the dual-scale and FAST + single-scale layouts, so future fused ROCm kernels can target the current + balanced FAST artifact without reintroducing the older scale path. +- The standalone HIP dequant launch now maps one thread to one packed + ROCmFP4 byte across a 256-thread global grid. This avoids the older + one-16-thread-block-per-32-values launch shape and gives future fused or + diagnostic dequant paths normal GPU occupancy scaling. +- Vulkan ROCmFP4 shaders also decode UE4M3 scales directly to the half-scale + value used by the codebook, matching CPU/HIP and avoiding repeated `* 0.5` + fixups at dequant and matmul call sites. +- Vulkan ROCmFP4 shaders keep a shared `kvalues_rocmfp4` Codebook10 table. + Arithmetic/direct Codebook10 decode variants compiled and preserved + coherence, but measured slower on Strix Halo Vulkan, so the table path + remains the active backend implementation. +- Vulkan `dequantize4()` for ROCmFP4 and ROCmFP4_FAST decodes each block's + UE4M3 scale once per 4-value vector instead of calling the 2-value helper + twice. This improves tested dual-scale ROCmFP4 small-batch matvec shapes + while keeping the promoted Vulkan runtime guard clean. +- Vulkan `Q4_0_ROCMFP4_FAST` matvec/MMQ kernels have a single-scale dot + specialization. They combine the two half-block dot sums and apply the one + FAST scale once, instead of taking the dual-scale path used by + `Q4_0_ROCMFP4`. +- Vulkan `Q4_0_ROCMFP4_FAST` MMQ stores its block scale as a scalar in + shared/register cache instead of duplicating it into a `vec2`. The focused + Strix Halo Vulkan `MUL_MAT` microbench for + `m=4096,n=1,k=14336,type_a=q4_0_rocmfp4_fast` improved from `62.18` to + `61.45` us/run. +- A Vulkan packed16 view for ROCmFP4 quant bytes was tested and rejected. It + compiled, but the focused FAST MMQ guard regressed to `83.03` us/run from the + accepted `61`-`62` us/run range, so ROCmFP4 Vulkan keeps the byte-view load + path. +- Vulkan backend CPY/SET_ROWS now has generated `F32 -> Q4_0_ROCMFP4`, + `F32 -> Q4_0_ROCMFP4_FAST`, `Q4_0_ROCMFP4 -> F32`, + `Q4_0_ROCMFP4_FAST -> F32`, and indexed SET_ROWS shaders. The SET_ROWS path + uses the same exhaustive finite UE4M3 scale search as the CPU reference, so + quantized K/V cache writes favor coherence over a cheap max-abs shortcut. +- Vulkan `F32 -> ROCmFP4` runtime quantization now uses the same exact ordered + UE4M3 scale search as CPU/HIP: find the scale nearest `max_abs / 10`, expand + outward, and stop a candidate once its partial error cannot beat the current + best scale. The candidate set is unchanged, so this avoids a slower linear + scan without falling back to a lower-quality shortcut. +- Vulkan source-to-ROCmFP4 runtime quantization now also prunes lower scale + candidates once clipping the block max alone cannot beat the current best + error. This mirrors the CPU/HIP scale search bound and keeps the candidate + set exact. The 2026-05-25 Vulkan CPY guard passed after this shader change: + F32/F16/BF16-to-dual measured `9525.39`, `2350.54`, and `2418.09` us/run, + dual-to-F32 `516.65` us/run, F32/F16/BF16-to-FAST `10111.85`, `2923.67`, + and `2949.42` us/run, and FAST-to-F32 `509.65` us/run. The full promoted + gate passed after rebuild with Qwen3.6 27B MTP at `33.9` / `28.0 tok/s` and + no KFD PIDs left running. +- Vulkan backend CPY also supports `F16 -> Q4_0_ROCMFP4`, + `F16 -> Q4_0_ROCMFP4_FAST`, `BF16 -> Q4_0_ROCMFP4`, and + `BF16 -> Q4_0_ROCMFP4_FAST`. The runtime quantization shader can load + `float`, `float16_t`, and BF16 source bits and then runs the same exact + ordered finite UE4M3 scale search. Backend tests keep a bounded NMSE + tolerance only for the half/bfloat runtime quantization cases because those + paths are inherently lossy around source-precision tie points; the + `F32 -> ROCmFP4` checks remain strict. +- Vulkan same-type CPY supports packed-block copies for + `Q4_0_ROCMFP4 -> Q4_0_ROCMFP4` and + `Q4_0_ROCMFP4_FAST -> Q4_0_ROCMFP4_FAST`, including block-aligned + non-contiguous/permuted/view copies. The non-contiguous path uses a + byte-addressed block shader and preserves exact 18-byte or 17-byte + ROCmFP4 blocks. Contiguous same-type dual-scale and FAST copies now both use + the direct byte-copy path, avoiding the generic halfword copy route for + ROCmFP4's custom 18-byte and 17-byte block layouts. +- Vulkan CPY is now covered by a dedicated regression guard so copy-path + changes cannot silently fall back or regress outside the ROCm-only CPY gate. + The guard runs same-type and source/dequant CPY correctness before measuring + the large copy performance shape, and now streams the performance phase + through `tee` so long Vulkan runs do not look idle to the command runner. + The ceilings were tightened after the lower-scale pruning gain, so the old + slow source-to-quant path no longer passes this guard. + On Strix Halo RADV Vulkan, the large guarded shape currently measures + F32/F16/BF16-to-dual at `9525.39`, `2350.54`, and `2418.09` us/run, + dual-to-F32 at `516.65` us/run, F32/F16/BF16-to-FAST at `10111.85`, + `2923.67`, and `2949.42` us/run, and FAST-to-F32 at `509.65` us/run. +- Vulkan scalar FlashAttention can now decode ROCmFP4 and ROCmFP4_FAST K/V + cache blocks. ROCmFP4 K/V is forced to the scalar FA path because the current + custom decode is not a coopmat/native matrix-core FP4 path. +- Vulkan scalar FlashAttention can use the integer-dot MMQ K path for both + ROCmFP4 K-cache layouts. The FAST layout expands each 4-bit Codebook10 value + into packed signed int8 lanes and uses its single UE4M3 scale as the K block + multiplier. The dual-scale layout also uses packed signed int8 lanes, but + splits the dot accumulation by half-block so each 16-value half uses its own + UE4M3 scale. This keeps the quality-biased STRIX K-cache path fast without + applying one scale to a two-scale block. +- Build and runtime verification generated the Vulkan SPIR-V entries for + ROCmFP4 copy/SET_ROWS shaders, linked `libggml-vulkan.so`, and passed + Vulkan ROCmFP4 CPY plus MUL_MAT smoke tests on Strix Halo. +- Row validation rejects invalid scale bytes outside finite unsigned UE4M3 + (`0x00` through `0x7e`) so corrupted custom GGUF tensors fail early. +- Quantization keeps the exhaustive 126-scale UE4M3 search for both normal and + imatrix paths. Candidate-window scale search was tested and improved GGUF + creation speed, but it regressed the Qwen3-4B short WikiText-2 PPL on the + compact FAST path, so it was rejected for coherence. +- The exhaustive scale search now visits the UE4M3 candidate nearest the + block's `max_abs / 10` first using a monotonic binary search, expands + outward, and exits a candidate scale once its partial error cannot beat the + current best scale. This remains exact because every finite scale is still + evaluated; on the Qwen3-4B FAST_COHERENT artifact it produced a + byte-identical GGUF while cutting FAST quantization cost sharply. +- CPU scale decode now uses a 127-entry finite UE4M3 half-scale table for the + exhaustive scale search and row dequantization. The table preserves the same + exact FP32 values as the former bit-construction decoder, but avoids + rebuilding them for every scale candidate. +- CPU scale search now has separate unweighted and imatrix-weighted MSE + scoring paths. The normal path avoids per-value weight branches, while the + imatrix path precomputes the row-energy calibration weight once per + 16/32-value block instead of recomputing the same `sqrtf` term for every + candidate scale. +- CPU normal quantization scans each block once for finite values and uses a + finite-only nearest-Codebook10 helper in the exhaustive MSE loop when the + block is clean. Non-finite blocks still use the guarded helper. This + promoted conversion-speed path measured focused dual-scale / FAST normal + quantization at `3844.38` / `3582.57` cycles per 32 values in the latest + guard. +- CPU imatrix-weighted scale-MSE scoring now uses the same finite-only + nearest-Codebook10 helper when the source block is finite. This preserves the + guarded path for non-finite input while improving FAST imatrix GGUF creation: + same-session pre-candidate FAST imatrix was `5258.07` cycles / 32 values, and + two guarded candidate passes measured `4448.73` and `4447.32`. Dual imatrix + stayed in the noisy guarded band, so the claim is limited to FAST imatrix. +- The weighted and unweighted exhaustive scale loops are split as well, so + the normal quantizer does not branch on imatrix state for every finite + UE4M3 candidate. This keeps the exact candidate order and tie behavior while + improving CPU GGUF creation speed for both ROCmFP4 layouts. +- The CPU quantizers fill every packed nibble byte in each ROCmFP4 block, so + they no longer clear `qs` before packing. This is an exact store cleanup: + correctness stayed clean, `q4_0_rocmfp4` quantize improved to `5244.46` + cycles / 32 values, `q4_0_rocmfp4_fast` improved to `4668.28` cycles / 32 + values. +- CPU scale selection scans block maxima with a plain absolute-value compare + instead of `fmaxf`. NaNs still do not raise the maximum, infinities are still + rejected before scale search, and finite output is unchanged. This moved + `q4_0_rocmfp4` quantize to `5228.64` cycles / 32 values and FAST to + `4516.45` cycles / 32 values. +- CPU dequantization and quantization-error scoring use the ROCmFP4 arithmetic + Codebook10 decoder, which avoids a table fetch per unpacked nibble. CPU + fallback vec-dot keeps the 16-entry codebook table because measured Strix + Halo fallback dot speed was better with the table there. This hybrid keeps + the dequant win without regressing CPU vec-dot. +- CPU fallback vec-dot now reuses each packed ROCmFP4 byte for both low and + high nibbles before table decode. This preserves exact output, moved focused + dual-scale vec-dot from `31.86` to the `29.77`-`29.82` cycles / 32 values + band, and kept FAST in the `27.04`-`27.06` cycles / 32 values band. +- Two additional CPU decode shortcuts were rechecked on 2026-05-24 and + rejected. Replacing arithmetic Codebook10 decode with table decode in + row dequantization and scale-MSE scoring slowed CPU dequantization from the + guarded `33`-cycle band to `51`-`84` cycles / 32 values. Returning a direct + decoded value during MSE scoring instead of index-then-decode also passed + correctness but regressed normal quantization in the measured guard, so the + original arithmetic decoder remains the CPU path. A narrower full-block + weighted-MSE-only retry was also rejected because it regressed the focused + FAST imatrix timing to `5007.73` cycles / 32 values. +- A direct finite-scale table helper for CPU quantizer scale-search candidates + was also rejected. It helped dual-scale normal quantization, but repeat + guard runs regressed FAST normal quantization to `4043.60` and + `4289.66` cycles / 32 values. The guarded scale helper remains in place + because the compact STRIX_LEAN preset depends heavily on FAST tensors. +- NaN-only and branchless finite-scan variants were also rejected for default. + The NaN-only scan was noisy and the same-session `isfinite()` scan measured + better dual-scale normal quantization (`3735.41` vs `3776.53` cycles / 32 + values), while the branchless boolean-and scan regressed FAST normal + quantization to `3773.44` cycles / 32 values. +- Per-value Codebook10 quantization uses exact nearest-neighbor thresholds + instead of a 16-entry scan. The hot quantizer path uses one reciprocal scale + per candidate/block and multiplies each value by that reciprocal instead of + dividing per value. On the Qwen3-4B FAST_COHERENT check this kept PPL tied + with the accepted artifact while cutting GGUF creation time further. +- CPU exhaustive scale search now prunes lower-scale candidates after the + clipped max-value error alone cannot beat the current best error. The + imatrix path uses the positive calibration weight of the max element for the + same conservative bound. This does not change accepted values because every + skipped smaller scale has a larger unavoidable clip error. The weighted path + also combines imatrix weight preparation, max-absolute scan, and max-value + pruning-weight selection into one pre-scan, while both exhaustive loops stop + once no valid lower or upper scale candidates remain. Latest 40-iteration + quant guard run: `q4_0_rocmfp4` normal `4049.28` cycles / 32 values, FAST + normal `3721.48`, imatrix dual-scale `5251.76`, imatrix FAST `4898.75`. +- Adding `GGML_RESTRICT` to the local CPU scale-scoring helper pointers was + tested and rejected. It compiled and passed the quant guard, but did not + improve normal quantization and made the imatrix guard noisier/slower in + the measured pass, so the existing helper signatures were kept. +- Regression guards: + - `scripts/check-rocmfp4-deepseek-regression.sh` is an optional compatibility + smoke guard for a second ROCmFP4-converted model family. It is not part of + the promoted-gain gate because no reproducible DeepSeek speedup has been + established yet. The default all-regression script skips it unless + `INCLUDE_DEEPSEEK_SMOKE=1` is set. + - `scripts/check-rocmfp4-qwen-mtp-regression.sh` verifies the promoted + Qwen3.6 27B MTP ROCmFP4 262k-context `draft-mtp` path. The stable guard + uses ROCm0 and `--spec-draft-n-max 4`. On ROCm0, `n-max 4` improved the + short guard prompt from `27.7` to `33.6 tok/s` versus `n-max 3`, and the + longer forced-output prompt from `25.6` to `26.3 tok/s`; the dual-scale + MMVQ VDR tune then moved the longer prompt to `27.8 tok/s`. `n-max 5` + and `n-max 6` improved only the bursty short prompt after the MMVQ tune + but regressed sustained output to `24.8` and `24.0 tok/s`, respectively. + Follow-up `n-max 5` checks with `--spec-draft-p-min 0.75`, + `--spec-draft-p-min 0.90`, and `--spec-draft-n-min 1` still held the + short prompt at `45.0`-`45.1 tok/s` but kept sustained output at only + `24.7 tok/s`, so `n-max 4` remains promoted. + The guard now checks both short and sustained prompts so high-acceptance + best-case gains do not hide sustained-output regressions. + Additional ROCm0 checks rejected larger batch sizes, smaller ubatch, + q8/f16 KV cache, `--swa-full`, `--no-host`, `--no-op-offload`, + `--no-repack`, polling changes, lower and higher CPU thread counts, + greedy sampling, backend sampling, p-split changes, sampler chain trimming, + and combined MTP+ngram speculation because none improved the sustained + prompt. After the MMVQ VDR tune, `-b 1024 -ub 512`, `-b 2048 -ub 512`, + and `-b 512 -ub 256` still tied the promoted `27.8 tok/s` sustained + result, while `-b 1024 -ub 1024` and + `--poll-batch 1 --spec-draft-poll-batch 1` stayed at `27.7 tok/s`. + q8 and f16 KV cache dropped sustained decode to `25.4` and `26.0 tok/s`. + The MTP draft loop now honors `--spec-draft-p-min`. With that cutoff active, + `n-max 4, p-min 0.75` tied the promoted path at `33.5 tok/s` short and + `27.6 tok/s` sustained, while `n-max 5` stayed burst-only at `45.0`-`45.1` + tok/s short and `24.6`-`24.7 tok/s` sustained for `p-min 0.75` and `0.90`. + Follow-up ROCm0 KV isolation confirmed this is not a draft-cache issue: + q8 main KV with q4 draft KV measured `33.5 tok/s` short and `25.3 tok/s` + sustained, q8 main K only measured `31.5` and `21.4 tok/s`, and q8 main V + only measured `31.1` and `21.9 tok/s`. The promoted ROCm0 path therefore + keeps both accepted and draft KV at q4. + `--spec-draft-p-split 0.05` and `0.20` both tied the `27.7 tok/s` + sustained guard, `-t 24 -tb 32 --spec-draft-threads 24 + --spec-draft-threads-batch 32` tied, and trimming the sampler chain to + `top_k;top_p;temperature` regressed sustained decode to `26.6 tok/s`. + Draft-only q8 KV (`SPEC_DRAFT_TYPE_K=q8_0`, + `SPEC_DRAFT_TYPE_V=q8_0`) with main KV left at q4 tied the promoted path + at `33.6 tok/s` short and `27.7 tok/s` sustained, so the leaner q4 draft + KV default remains promoted. + The heavier Qwen3.6 27B MTP `STRIX_MTP_Q6` ROCmFP4 profile was also + checked with the same ROCm0 guard. It reached `30.1 tok/s` on the short + prompt but only `21.3 tok/s` sustained, below the guard floor, so the + Q6/quality-biased profile is not promoted for throughput. + The guard script supports `SPEC_DRAFT_N_MAX`, `SPEC_DRAFT_N_MIN`, + `SPEC_DRAFT_P_MIN`, `SPEC_DRAFT_P_SPLIT`, `THREADS`, `THREADS_BATCH`, + `SPEC_DRAFT_THREADS`, `SPEC_DRAFT_THREADS_BATCH`, `SAMPLERS`, + `BATCH_SIZE`, `UBATCH_SIZE`, `CACHE_TYPE_K`, `CACHE_TYPE_V`, + `SPEC_DRAFT_TYPE_K`, `SPEC_DRAFT_TYPE_V`, `EXTRA_ARGS`, and + `SPEC_EXTRA_ARGS` environment overrides for + controlled sweeps while defaulting to the promoted `n-max 4`, `n-min 0`, + `p-min 0.0`, `p-split 0.10`, q4 KV, `-t 16`, `-tb 32`, `-b 512`, and + `-ub 512` settings. + - `scripts/check-rocmfp4-quant-regression.sh` runs quant correctness plus + CPU quantizer, dequantizer, and vec-dot cycle ceilings for both ROCmFP4 + block layouts. It checks normal and imatrix quantization, and now also + protects CPU fallback dequant/vec-dot paths so decode-helper experiments + cannot pass while slowing non-GPU fallbacks. Recent pass after keeping the + `isfinite()` finite-block scoring promotion and adding the weighted/imatrix + finite scorer: dual-scale quant `3844.38`, FAST quant `3582.57`, dual + dequant `33.59`, FAST dequant `33.13`, dual vec-dot `29.96`, FAST vec-dot + `27.03`, dual imatrix `5587.43`, and FAST imatrix `4447.32` cycles / 32 + values. + - `scripts/check-rocmfp4-vulkan-runtime-regression.sh` measures focused + Vulkan `MUL_MAT` runtime for ROCmFP4 FAST and dual-scale layouts, catching + shader regressions that can be hidden by end-to-end decode noise. The + guard covers `n=1`, `n=2`, `n=4`, and `n=8` for both layouts. Recent + tightened serial pass measured FAST `55.82`, `71.59`, `105.10`, and + `163.41` us/run for `n=1/2/4/8`, and dual-scale `64.87`, `83.24`, + `118.56`, and `194.27` us/run for `n=1/2/4/8`. The previous serial pass + after adding the shared ROCmFP4 UE4M3 + scale LUT was FAST `53.66`, `71.67`, `105.14` us/run and dual-scale + `65.05`, `83.07`, `122.70` us/run for the same shapes. + - `scripts/check-rocmfp4-vulkan-cpy-regression.sh` measures Vulkan0 CPY for + `F32/F16/BF16 -> Q4_0_ROCMFP4`, + `F32/F16/BF16 -> Q4_0_ROCMFP4_FAST`, and the matching quant-to-F32 + dequant copy paths on the same large shape used by the ROCm CPY guard. + It is included in the all-regression harness to catch accidental Vulkan + copy fallback or shader-routing regressions. + - `scripts/check-rocmfp4-rocm-runtime-regression.sh` measures the same + focused `MUL_MAT` shapes on ROCm0. The guard covers `n=1`, `n=2`, `n=4`, + and `n=8` for both FAST and dual-scale ROCmFP4 so MTP-style multi-column + regressions are caught instead of only protecting single-token matvec. + Recent tightened serial pass after extending the RDNA3.5 two-warp route + through `n=2` measured FAST `51.66`, `66.56`, `101.27`, and `168.84` + us/run for `n=1/2/4/8`, and dual-scale `53.66`, `58.40`, `87.32`, and + `148.26` us/run for `n=1/2/4/8`. The `n=8` guard was added after a + rejected wide-rows candidate regressed dual-scale `n=8` to `1382.72` + us/run. + - `scripts/check-rocmfp4-rocm-cpy-regression.sh` measures ROCm0 CPY for + `F32/F16/BF16 -> Q4_0_ROCMFP4`, + `F32/F16/BF16 -> Q4_0_ROCMFP4_FAST`, and the matching quant-to-F32 + dequant copy paths on a large `8192x512x2` shape. Recent serial-gate pass + after direct FAST value scoring: F32/F16/BF16-to-dual `1111.73`, + `1008.69`, `1006.93` us/run, dual-to-F32 `182.25` us/run, + F32/F16/BF16-to-FAST `1047.21`, `950.93`, `951.00` us/run, and + FAST-to-F32 `170.36` us/run. + - `scripts/check-rocmfp4-all-regression.sh` runs the promoted-gain gate: + quant, Vulkan runtime, Vulkan CPY, ROCm runtime, ROCm FlashAttention, + ROCm CPY, and Qwen MTP guards serially, then checks ROCm KFD PIDs. Runtime + microbenchmarks should use this serial path rather than being run in + parallel with other GPU/UMA workloads. DeepSeek is not part of this + promoted-gain gate because no reproducible DeepSeek speedup has been + established; run the separate compatibility smoke only when explicitly + needed with `INCLUDE_DEEPSEEK_SMOKE=1`. The K/Q block-pair FlashAttention + serial pass held Qwen3.6 27B MTP at `33.2 tok/s` short and `27.7 tok/s` + sustained, with no KFD PIDs left running. The same pass measured ROCm + runtime `MUL_MAT` at FAST `54.16` / `66.13` / `101.73` / `178.36` us and + dual-scale `51.65` / `57.98` / `88.33` / `148.48` us for `n=1/2/4/8`, + and ROCm FlashAttention at `81.62` / `78.13` us for 64d dual-scale / FAST + and `228.58` / `199.32` us for Qwen-style 128d dual-scale / FAST. + - The MTP `accept()` path now skips re-copying `pending_h` when acceptance + lands on the final verify row already staged by `process()`. The focused + Qwen MTP guard held `33.8 tok/s` short and `27.9 tok/s` sustained after + this host-side cleanup. + - The MTP `process()` path now stores only the non-final verify rows needed + for partial-accept rollback, while copying the final target hidden row + directly to `pending_h` for the next MTP step. The focused Qwen MTP guard + held `33.9 tok/s` short and `27.9 tok/s` sustained, and the full serial + gate held `33.7 tok/s` short and `27.9 tok/s` sustained after this cleanup. + - The simple draft and MTP draft paths now return before + `llama_decode(ctx_dft, batch)` when no sequence is actively drafting. The + focused Qwen MTP guard held `33.8 tok/s` short and `28.0 tok/s` sustained, + and the full serial gate held `33.7 tok/s` short and `27.9 tok/s` + sustained after this idle/no-draft cleanup. + - The simple draft path now uses a reusable `uint8_t` drafting-state buffer + instead of allocating `std::vector` per draft call. The focused Qwen + MTP guard held `33.7 tok/s` short and `28.0 tok/s` sustained, and the full + serial gate held `33.8 tok/s` short and `27.9 tok/s` sustained. + - The MTP draft loop now skips its debug-candidate loop entirely unless + debug logging is enabled, hoists the debug verbosity state once per draft + call, and uses direct vector indexing in the hot per-token path. The + focused Qwen MTP guard held `33.8 tok/s` short and `27.9 tok/s` sustained + after this cleanup, so it is kept as a low-risk host-side simplification + rather than claimed as a decode-speed gain. + - The simple and MTP draft paths now use a shared direct one-sequence batch + append helper instead of `common_batch_add(..., { seq_id }, ...)` in the + hot speculative path. This avoids a temporary sequence-id container without + changing multi-sequence batch behavior elsewhere. The focused Qwen MTP + guard held `33.9 tok/s` short and `27.9 tok/s` sustained, and the full + serial gate held `33.8 tok/s` short and `27.9 tok/s` sustained. + - The MTP verify-row buffer now reserves for the configured draft depth and + only grows when needed instead of shrinking/resizing on every verification + pass. The focused Qwen MTP guard held `33.8 tok/s` short and `27.9 tok/s` + sustained, and the full serial gate held `33.9 tok/s` short and + `27.9 tok/s` sustained with no KFD PIDs left running. + - The MTP `process()` path now copies retained verification hidden rows from + the target embedding buffer in one contiguous `memcpy()` instead of one + `llama_get_embeddings_pre_norm_ith()` plus `memcpy()` per row. Rollback + behavior is unchanged. The focused Qwen3.6 27B MTP guard held + `33.7 tok/s` short and `27.9 tok/s` sustained, and the 35B A3B guard held + `104.1 tok/s` short and `89.2 tok/s` sustained. The same build later + passed the full serial all-regression gate with `INCLUDE_QWEN35_A3B_GUARD=1`, + including Qwen3.6 27B MTP at `33.8` / `27.9 tok/s` and Qwen3.6 35B A3B MTP + at `104.1` / `89.3 tok/s`. + - A matching contiguous-pointer cleanup inside the MTP `draft()` loop was + tested and rejected. It replaced per-row + `llama_get_embeddings_pre_norm_ith(ctx_dft, i_batch)` calls with one + `llama_get_embeddings_pre_norm(ctx_dft)` pointer per draft decode + iteration. The 27B MTP guard held at `33.8` / `27.9 tok/s`, but the 35B A3B + repeat measured `104.3` / `88.7` and `104.3` / `89.2 tok/s`, below the + promoted `89.3 tok/s` sustained band, so the code change was removed. + - A single-sequence MTP `draft()` fast path was tested and rejected on + 2026-05-25. It removed the active-sequence bookkeeping loop for the common + `n_seq == 1` case and passed the dense 27B guard at `33.7` / `28.0 tok/s`, + but the 35B A3B sustained guard collapsed to `25.7 tok/s` despite a + passing `103.1 tok/s` short check. Reverting that path restored the 35B A3B + guard to `104.3` / `90.3 tok/s`, so the shared multi-sequence draft loop + remains the promoted implementation. + - A dual-scale-only finite-pack CPU quantizer shortcut was tested and + rejected on 2026-05-24. It passed correctness, but even after isolating + the shared scale chooser it regressed the protected FAST quant path + (`q4_0_rocmfp4_fast` normal quant rose into the `3882`-`4022` cycles / 32 + values band versus the clean `3623` cycles / 32 values baseline). The + final packing loops therefore remain on the guarded nearest-code helper. +- The Strix ROCmFP4 build script now builds the quant regression test + binaries (`test-quantize-fns`, `test-quantize-perf`, and + `test-backend-ops`) alongside `llama-cli`, `llama-quantize`, and + `llama-bench`, so quant and runtime guards work from a clean build + directory. +- `test-quantize-perf --imatrix` benchmarks the ROCmFP4 quality/coherence + path through `ggml_quantize_chunk(..., imatrix)` with synthetic importance + weights. The quant regression guard includes this mode to reject imatrix-side + CPU changes that would not show up in the normal non-imatrix timing. + +High-upside Strix-specific work: +- The next significant speed step is likely a fused ROCmFP4 decode-and-dot + ROCm matvec path. The current ROCmFP4 wins come from compact blocks and + avoiding generic decode/copy overhead; a fused path would keep Codebook10 + values and UE4M3 scales in registers/shared memory through the dot product + instead of materializing wider intermediates. +- A ROCmFP4-aware long-context attention path is the second major target. The + long-context limit is mostly KV traffic, FlashAttention shape efficiency, and + memory bandwidth. The current dual-scale layout protects coherence, but it is + not itself a large speed lever unless the attention kernels consume it + directly and efficiently. +- MTP target/draft overlap is a possible scheduler project, not a flag. The + current common MTP flow is serial on one Strix Halo ROCm device. Real overlap + would need scheduler changes plus acceptance/correctness guards, and it + should only be promoted if it beats the sustained Qwen 262k guard. +- Coherence should stay protected by tensor-aware ROCmFP4 profiles: dual-scale + for sensitive tensors and FAST only where guarded model tests show no + sustained decode or quality regression. Pure speed-only profiles are not + promoted in this tree without the serial regression gate. + +Hardware note: +- This is a special AMD-targeted ggml/llama.cpp quantization and backend + path. It includes custom Vulkan and ROCm/HIP handling for the new GGUF + types, but it is not yet a native rocWMMA FP4 tensor-core implementation. + Current speed gains come from the compact block layout and backend decode + paths; deeper rocWMMA/cooperative-matrix work is future optimization work. +- NVIDIA CUDA is disabled in the Strix-FP4 build (`-DGGML_CUDA=OFF`). Some + upstream llama.cpp HIP backend sources still live under + `ggml/src/ggml-cuda` and are compiled by HIP for AMD, but the ROCmFP4-owned + helper code and user-facing build/run path are ROCm/HIP/Vulkan targeted. + This tree also accepts `GGML_HIP_ENABLE_UNIFIED_MEMORY=1` as the AMD-named + alias for the upstream unified-memory switch. +- The bundled rocWMMA 7.1.0 headers expose gfx12 WMMA paths for FP8/BF8 and + integer 8-bit inputs, but no native FP4 input type or FP4 WMMA/MFMA builtin + is visible locally. A true matrix-core ROCmFP4 path therefore needs a + measured unpack/convert strategy first, such as ROCmFP4 Codebook10 to int8 + WMMA or FP8 WMMA tiles, before claiming native FP4 tensor-core execution. +- rocWMMA FlashAttention is intentionally opt-in via + `GGML_HIP_ROCWMMA_FATTN=ON scripts/build-strix-rocmfp4-mtp.sh`. + It currently compiles with the local rocWMMA headers in + `/path/to/third_party/rocWMMA`, but the Strix Halo benchmark + regressed sustained Qwen MTP decode (`23.3 tok/s` vs the promoted + `26.2 tok/s` default HIP FlashAttention path), so the default build keeps it + disabled. A follow-up on the 35B A3B ROCmFP4 MTP guard measured only + `99.7` / `76.1 tok/s` versus the promoted `104.4` / `89.3` band, so this + remains true for the MoE-heavy profile as well. +- TurboQuant and TriAttention are not present as runtime flags in this isolated + tree. Integrating them would require a source-level merge into `strix-fp4`; + do not claim support or promote them unless they beat the serial ROCmFP4 + regression gate. diff --git a/ggml/rocmfp4/rocmfp4.c b/ggml/rocmfp4/rocmfp4.c new file mode 100644 index 000000000000..bccefcec15da --- /dev/null +++ b/ggml/rocmfp4/rocmfp4.c @@ -0,0 +1,797 @@ +#define GGML_COMMON_DECL_C +#include "../src/ggml-common.h" + +#include "rocmfp4.h" + +#include +#include +#include + +// ggml-base is compiled architecture-neutral (no -mavx2), so SIMD for the hot +// CPU dot product is enabled per-function via a target attribute plus a runtime +// CPU check. This keeps the AVX2 path in one translation unit without exporting +// internals to ggml-cpu. Non-GNU/non-x86 builds fall back to the scalar loop. +#if defined(__GNUC__) && (defined(__x86_64__) || defined(__i386__)) && !defined(ROCMFP4_NO_AVX2) +#include +#define ROCMFP4_X86_AVX2_DISPATCH 1 +#endif + +// ROCmFP4 stores a signed integer FP4-like codebook at half-scale. It is +// E2M1-derived, but the largest magnitude is retuned from 12 to 10 after +// sampling Qwen3 dense tensors; this reduces outlier pull without changing the +// packed 4-bit layout or integer dot-product path. +static const int8_t rocmfp4_codebook[16] = { + 0, 1, 2, 3, 4, 6, 8, 10, + 0, -1, -2, -3, -4, -6, -8,-10, +}; + +static inline int8_t rocmfp4_decode(uint8_t q) { + q &= 0x0f; + const int mag3 = q & 0x07; + const int mag = mag3 <= 4 ? mag3 : 2*mag3 - 4; + return (q & 0x08) ? -mag : mag; +} + +static inline int8_t rocmfp4_decode_table(uint8_t q) { + return rocmfp4_codebook[q & 0x0f]; +} + +// Finite unsigned E4M3 scale bytes decoded to the half-scale values used by +// ROCmFP4. Keeping this as a table avoids rebuilding identical FP32 values for +// every candidate during exhaustive scale search. +#define ROCMFP4_SCALE_SUB(M) ((M) * 0x1p-10f) +#define ROCMFP4_SCALE_E1(M) ((8 + (M)) * 0x1p-10f) +#define ROCMFP4_SCALE_E2(M) ((8 + (M)) * 0x1p-9f) +#define ROCMFP4_SCALE_E3(M) ((8 + (M)) * 0x1p-8f) +#define ROCMFP4_SCALE_E4(M) ((8 + (M)) * 0x1p-7f) +#define ROCMFP4_SCALE_E5(M) ((8 + (M)) * 0x1p-6f) +#define ROCMFP4_SCALE_E6(M) ((8 + (M)) * 0x1p-5f) +#define ROCMFP4_SCALE_E7(M) ((8 + (M)) * 0x1p-4f) +#define ROCMFP4_SCALE_E8(M) ((8 + (M)) * 0x1p-3f) +#define ROCMFP4_SCALE_E9(M) ((8 + (M)) * 0x1p-2f) +#define ROCMFP4_SCALE_E10(M) ((8 + (M)) * 0x1p-1f) +#define ROCMFP4_SCALE_E11(M) ((8 + (M)) * 0x1p0f) +#define ROCMFP4_SCALE_E12(M) ((8 + (M)) * 0x1p1f) +#define ROCMFP4_SCALE_E13(M) ((8 + (M)) * 0x1p2f) +#define ROCMFP4_SCALE_E14(M) ((8 + (M)) * 0x1p3f) +#define ROCMFP4_SCALE_E15(M) ((8 + (M)) * 0x1p4f) + +static const float rocmfp4_scale_ue4m3_half[127] = { + ROCMFP4_SCALE_SUB(0), ROCMFP4_SCALE_SUB(1), ROCMFP4_SCALE_SUB(2), ROCMFP4_SCALE_SUB(3), + ROCMFP4_SCALE_SUB(4), ROCMFP4_SCALE_SUB(5), ROCMFP4_SCALE_SUB(6), ROCMFP4_SCALE_SUB(7), + ROCMFP4_SCALE_E1(0), ROCMFP4_SCALE_E1(1), ROCMFP4_SCALE_E1(2), ROCMFP4_SCALE_E1(3), + ROCMFP4_SCALE_E1(4), ROCMFP4_SCALE_E1(5), ROCMFP4_SCALE_E1(6), ROCMFP4_SCALE_E1(7), + ROCMFP4_SCALE_E2(0), ROCMFP4_SCALE_E2(1), ROCMFP4_SCALE_E2(2), ROCMFP4_SCALE_E2(3), + ROCMFP4_SCALE_E2(4), ROCMFP4_SCALE_E2(5), ROCMFP4_SCALE_E2(6), ROCMFP4_SCALE_E2(7), + ROCMFP4_SCALE_E3(0), ROCMFP4_SCALE_E3(1), ROCMFP4_SCALE_E3(2), ROCMFP4_SCALE_E3(3), + ROCMFP4_SCALE_E3(4), ROCMFP4_SCALE_E3(5), ROCMFP4_SCALE_E3(6), ROCMFP4_SCALE_E3(7), + ROCMFP4_SCALE_E4(0), ROCMFP4_SCALE_E4(1), ROCMFP4_SCALE_E4(2), ROCMFP4_SCALE_E4(3), + ROCMFP4_SCALE_E4(4), ROCMFP4_SCALE_E4(5), ROCMFP4_SCALE_E4(6), ROCMFP4_SCALE_E4(7), + ROCMFP4_SCALE_E5(0), ROCMFP4_SCALE_E5(1), ROCMFP4_SCALE_E5(2), ROCMFP4_SCALE_E5(3), + ROCMFP4_SCALE_E5(4), ROCMFP4_SCALE_E5(5), ROCMFP4_SCALE_E5(6), ROCMFP4_SCALE_E5(7), + ROCMFP4_SCALE_E6(0), ROCMFP4_SCALE_E6(1), ROCMFP4_SCALE_E6(2), ROCMFP4_SCALE_E6(3), + ROCMFP4_SCALE_E6(4), ROCMFP4_SCALE_E6(5), ROCMFP4_SCALE_E6(6), ROCMFP4_SCALE_E6(7), + ROCMFP4_SCALE_E7(0), ROCMFP4_SCALE_E7(1), ROCMFP4_SCALE_E7(2), ROCMFP4_SCALE_E7(3), + ROCMFP4_SCALE_E7(4), ROCMFP4_SCALE_E7(5), ROCMFP4_SCALE_E7(6), ROCMFP4_SCALE_E7(7), + ROCMFP4_SCALE_E8(0), ROCMFP4_SCALE_E8(1), ROCMFP4_SCALE_E8(2), ROCMFP4_SCALE_E8(3), + ROCMFP4_SCALE_E8(4), ROCMFP4_SCALE_E8(5), ROCMFP4_SCALE_E8(6), ROCMFP4_SCALE_E8(7), + ROCMFP4_SCALE_E9(0), ROCMFP4_SCALE_E9(1), ROCMFP4_SCALE_E9(2), ROCMFP4_SCALE_E9(3), + ROCMFP4_SCALE_E9(4), ROCMFP4_SCALE_E9(5), ROCMFP4_SCALE_E9(6), ROCMFP4_SCALE_E9(7), + ROCMFP4_SCALE_E10(0), ROCMFP4_SCALE_E10(1), ROCMFP4_SCALE_E10(2), ROCMFP4_SCALE_E10(3), + ROCMFP4_SCALE_E10(4), ROCMFP4_SCALE_E10(5), ROCMFP4_SCALE_E10(6), ROCMFP4_SCALE_E10(7), + ROCMFP4_SCALE_E11(0), ROCMFP4_SCALE_E11(1), ROCMFP4_SCALE_E11(2), ROCMFP4_SCALE_E11(3), + ROCMFP4_SCALE_E11(4), ROCMFP4_SCALE_E11(5), ROCMFP4_SCALE_E11(6), ROCMFP4_SCALE_E11(7), + ROCMFP4_SCALE_E12(0), ROCMFP4_SCALE_E12(1), ROCMFP4_SCALE_E12(2), ROCMFP4_SCALE_E12(3), + ROCMFP4_SCALE_E12(4), ROCMFP4_SCALE_E12(5), ROCMFP4_SCALE_E12(6), ROCMFP4_SCALE_E12(7), + ROCMFP4_SCALE_E13(0), ROCMFP4_SCALE_E13(1), ROCMFP4_SCALE_E13(2), ROCMFP4_SCALE_E13(3), + ROCMFP4_SCALE_E13(4), ROCMFP4_SCALE_E13(5), ROCMFP4_SCALE_E13(6), ROCMFP4_SCALE_E13(7), + ROCMFP4_SCALE_E14(0), ROCMFP4_SCALE_E14(1), ROCMFP4_SCALE_E14(2), ROCMFP4_SCALE_E14(3), + ROCMFP4_SCALE_E14(4), ROCMFP4_SCALE_E14(5), ROCMFP4_SCALE_E14(6), ROCMFP4_SCALE_E14(7), + ROCMFP4_SCALE_E15(0), ROCMFP4_SCALE_E15(1), ROCMFP4_SCALE_E15(2), ROCMFP4_SCALE_E15(3), + ROCMFP4_SCALE_E15(4), ROCMFP4_SCALE_E15(5), ROCMFP4_SCALE_E15(6), +}; + +#undef ROCMFP4_SCALE_SUB +#undef ROCMFP4_SCALE_E1 +#undef ROCMFP4_SCALE_E2 +#undef ROCMFP4_SCALE_E3 +#undef ROCMFP4_SCALE_E4 +#undef ROCMFP4_SCALE_E5 +#undef ROCMFP4_SCALE_E6 +#undef ROCMFP4_SCALE_E7 +#undef ROCMFP4_SCALE_E8 +#undef ROCMFP4_SCALE_E9 +#undef ROCMFP4_SCALE_E10 +#undef ROCMFP4_SCALE_E11 +#undef ROCMFP4_SCALE_E12 +#undef ROCMFP4_SCALE_E13 +#undef ROCMFP4_SCALE_E14 +#undef ROCMFP4_SCALE_E15 + +static inline float rocmfp4_ue4m3_to_fp32_half(uint8_t e) { + return e <= 0x7e ? rocmfp4_scale_ue4m3_half[e] : 0.0f; +} + +static inline uint8_t rocmfp4_best_index_scaled_finite(float x, float inv_scale_half) { + // Exact nearest-neighbor thresholds for Codebook10: + // 0, +/-1, +/-2, +/-3, +/-4, +/-6, +/-8, +/-10 + // Ties intentionally choose the lower-magnitude code, matching the former + // linear scan because the positive codes and zero appear first. + const float a = fabsf(x * inv_scale_half); + if (a <= 0.5f) { + return 0; + } + + const bool neg = x < 0.0f; + if (a <= 1.5f) { + return neg ? 9 : 1; + } + if (a <= 2.5f) { + return neg ? 10 : 2; + } + if (a <= 3.5f) { + return neg ? 11 : 3; + } + if (a <= 5.0f) { + return neg ? 12 : 4; + } + if (a <= 7.0f) { + return neg ? 13 : 5; + } + if (a <= 9.0f) { + return neg ? 14 : 6; + } + + return neg ? 15 : 7; +} + +static inline uint8_t rocmfp4_best_index_scaled(float x, float inv_scale_half) { + if (!isfinite(x)) { + return 0; + } + + return rocmfp4_best_index_scaled_finite(x, inv_scale_half); +} + +// Fused best-index + decode used only inside the exhaustive scale search. The +// scale search re-scans every block element for every candidate scale byte, so +// avoiding the code -> decode round-trip on the hottest quantize path matters. +// Returns the same signed Codebook10 magnitude that +// rocmfp4_decode(rocmfp4_best_index_scaled_finite(x, inv_scale_half)) produces, +// so quantized output is bit-identical to the previous path. +static inline float rocmfp4_decoded_mag_scaled_finite(float x, float inv_scale_half) { + const float a = fabsf(x * inv_scale_half); + + float mag; + if (a <= 0.5f) { + mag = 0.0f; + } else if (a <= 1.5f) { + mag = 1.0f; + } else if (a <= 2.5f) { + mag = 2.0f; + } else if (a <= 3.5f) { + mag = 3.0f; + } else if (a <= 5.0f) { + mag = 4.0f; + } else if (a <= 7.0f) { + mag = 6.0f; + } else if (a <= 9.0f) { + mag = 8.0f; + } else { + mag = 10.0f; + } + + return x < 0.0f ? -mag : mag; +} + +static inline float rocmfp4_decoded_mag_scaled(float x, float inv_scale_half) { + if (!isfinite(x)) { + return 0.0f; + } + + return rocmfp4_decoded_mag_scaled_finite(x, inv_scale_half); +} + +static inline bool rocmfp4_scale_is_valid(uint8_t e) { + // ROCmFP4 scale bytes are unsigned finite E4M3 values. 0x7f is NaN in the + // unsigned encoding and values with the sign bit set are not valid scales. + return e <= 0x7e; +} + +static float rocmfp4_block_mse_for_scale_unweighted( + const float * x, int n, int e, float best_err) { + const float scale_half = rocmfp4_ue4m3_to_fp32_half((uint8_t) e); + const float inv_scale_half = 1.0f / scale_half; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + const float y = rocmfp4_decoded_mag_scaled(x[i], inv_scale_half) * scale_half; + const float d = x[i] - y; + + err += d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +static float rocmfp4_block_mse_for_scale_unweighted_finite( + const float * x, int n, int e, float best_err) { + const float scale_half = rocmfp4_ue4m3_to_fp32_half((uint8_t) e); + const float inv_scale_half = 1.0f / scale_half; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + const float y = rocmfp4_decoded_mag_scaled_finite(x[i], inv_scale_half) * scale_half; + const float d = x[i] - y; + + err += d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +static float rocmfp4_block_mse_for_scale_weighted( + const float * x, int n, const float * mse_weights, int e, float best_err) { + const float scale_half = rocmfp4_ue4m3_to_fp32_half((uint8_t) e); + const float inv_scale_half = 1.0f / scale_half; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + const float y = rocmfp4_decoded_mag_scaled(x[i], inv_scale_half) * scale_half; + const float d = x[i] - y; + + err += mse_weights[i]*d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +static float rocmfp4_block_mse_for_scale_weighted_finite( + const float * x, int n, const float * mse_weights, int e, float best_err) { + const float scale_half = rocmfp4_ue4m3_to_fp32_half((uint8_t) e); + const float inv_scale_half = 1.0f / scale_half; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + const float y = rocmfp4_decoded_mag_scaled_finite(x[i], inv_scale_half) * scale_half; + const float d = x[i] - y; + + err += mse_weights[i]*d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +static void rocmfp4_prepare_mse_weights( + float * dst, const float * x, int n, const float * quant_weights, float sigma2, + float * max_abs, float * max_abs_weight, bool * all_finite) { + *max_abs = 0.0f; + *max_abs_weight = 0.0f; + *all_finite = true; + + for (int i = 0; i < n; ++i) { + const float qw = quant_weights[i]; + const float ax = fabsf(x[i]); + const float weight = isfinite(qw) && qw > 0.0f ? qw * sqrtf(sigma2 + x[i]*x[i]) : 0.0f; + *all_finite = *all_finite && isfinite(x[i]); + + if (ax > *max_abs) { + *max_abs = ax; + *max_abs_weight = weight; + } else if (ax == *max_abs && weight > *max_abs_weight) { + *max_abs_weight = weight; + } + + // Match llama.cpp's imatrix weighting style for Q4_0: calibration + // importance is scaled by row energy so large activations remain protected. + dst[i] = weight; + } +} + +static int rocmfp4_nearest_scale_ue4m3(float target_scale_half) { + if (!(target_scale_half > 0.0f) || !isfinite(target_scale_half)) { + return 1; + } + + int lo = 1; + int hi = 126; + while (lo < hi) { + const int mid = lo + (hi - lo) / 2; + if (rocmfp4_ue4m3_to_fp32_half((uint8_t) mid) < target_scale_half) { + lo = mid + 1; + } else { + hi = mid; + } + } + + if (lo == 1) { + return 1; + } + + const float hi_scale = rocmfp4_ue4m3_to_fp32_half((uint8_t) lo); + const float lo_scale = rocmfp4_ue4m3_to_fp32_half((uint8_t) (lo - 1)); + + // Match the former ascending nearest scan: exact midpoint ties keep the + // lower scale byte. + return (target_scale_half - lo_scale <= hi_scale - target_scale_half) ? lo - 1 : lo; +} + +static uint8_t rocmfp4_choose_scale_ue4m3_exhaustive_unweighted( + const float * x, int n, float max_abs, bool all_finite) { + const int start_e = rocmfp4_nearest_scale_ue4m3(max_abs / 10.0f); + + int best_e = 0; + float best_err = FLT_MAX; + bool lower_done = false; + + for (int delta = 0; delta <= 125; ++delta) { + const int e0 = start_e - delta; + if (!lower_done && e0 >= 1 && e0 <= 126) { + const float scale_half = rocmfp4_ue4m3_to_fp32_half((uint8_t) e0); + const float clip_delta = max_abs - 10.0f*scale_half; + if (clip_delta > 0.0f && clip_delta*clip_delta > best_err) { + lower_done = true; + } else { + const float err = all_finite ? + rocmfp4_block_mse_for_scale_unweighted_finite(x, n, e0, best_err) : + rocmfp4_block_mse_for_scale_unweighted(x, n, e0, best_err); + if (err < best_err || (err == best_err && e0 < best_e)) { + best_err = err; + best_e = e0; + } + } + } + + const int e1 = start_e + delta; + if (delta != 0 && e1 >= 1 && e1 <= 126) { + const float err = all_finite ? + rocmfp4_block_mse_for_scale_unweighted_finite(x, n, e1, best_err) : + rocmfp4_block_mse_for_scale_unweighted(x, n, e1, best_err); + if (err < best_err || (err == best_err && e1 < best_e)) { + best_err = err; + best_e = e1; + } + } + + if ((lower_done || e0 <= 1) && e1 >= 126) { + break; + } + } + + return (uint8_t) best_e; +} + +static uint8_t rocmfp4_choose_scale_ue4m3_exhaustive_weighted( + const float * x, int n, const float * mse_weights, float max_abs, float max_abs_weight, bool all_finite) { + const int start_e = rocmfp4_nearest_scale_ue4m3(max_abs / 10.0f); + + int best_e = 0; + float best_err = FLT_MAX; + bool lower_done = false; + + for (int delta = 0; delta <= 125; ++delta) { + const int e0 = start_e - delta; + if (!lower_done && e0 >= 1 && e0 <= 126) { + const float scale_half = rocmfp4_ue4m3_to_fp32_half((uint8_t) e0); + const float clip_delta = max_abs - 10.0f*scale_half; + if (max_abs_weight > 0.0f && clip_delta > 0.0f && max_abs_weight*clip_delta*clip_delta > best_err) { + lower_done = true; + } else { + const float err = all_finite ? + rocmfp4_block_mse_for_scale_weighted_finite(x, n, mse_weights, e0, best_err) : + rocmfp4_block_mse_for_scale_weighted(x, n, mse_weights, e0, best_err); + if (err < best_err || (err == best_err && e0 < best_e)) { + best_err = err; + best_e = e0; + } + } + } + + const int e1 = start_e + delta; + if (delta != 0 && e1 >= 1 && e1 <= 126) { + const float err = all_finite ? + rocmfp4_block_mse_for_scale_weighted_finite(x, n, mse_weights, e1, best_err) : + rocmfp4_block_mse_for_scale_weighted(x, n, mse_weights, e1, best_err); + if (err < best_err || (err == best_err && e1 < best_e)) { + best_err = err; + best_e = e1; + } + } + + if ((lower_done || e0 <= 1) && e1 >= 126) { + break; + } + } + + return (uint8_t) best_e; +} + +static uint8_t rocmfp4_choose_scale_ue4m3(const float * x, int n, const float * quant_weights, float sigma2) { + if (quant_weights) { + assert(n <= QK_ROCMFP4); + float mse_weights_buf[QK_ROCMFP4]; + float weighted_max_abs; + float max_abs_weight; + bool all_finite; + rocmfp4_prepare_mse_weights(mse_weights_buf, x, n, quant_weights, sigma2, &weighted_max_abs, &max_abs_weight, &all_finite); + if (!(weighted_max_abs > 0.0f) || !isfinite(weighted_max_abs)) { + return 0; + } + return rocmfp4_choose_scale_ue4m3_exhaustive_weighted(x, n, mse_weights_buf, weighted_max_abs, max_abs_weight, all_finite); + } + + float max_abs = 0.0f; + bool all_finite = true; + for (int i = 0; i < n; ++i) { + all_finite = all_finite && isfinite(x[i]); + const float ax = fabsf(x[i]); + if (ax > max_abs) { + max_abs = ax; + } + } + + if (!(max_abs > 0.0f) || !isfinite(max_abs)) { + return 0; + } + + return rocmfp4_choose_scale_ue4m3_exhaustive_unweighted(x, n, max_abs, all_finite); +} + +static void rocmfp4_quantize_row_q4_0_weighted( + const float * GGML_RESTRICT x, block_rocmfp4 * GGML_RESTRICT y, int64_t k, const float * GGML_RESTRICT quant_weights) { + assert(k % QK_ROCMFP4 == 0); + + float sum_x2 = 0.0f; + for (int64_t i = 0; i < k; ++i) { + sum_x2 += x[i]*x[i]; + } + const float sigma2 = sum_x2 / (float) k; + + const int64_t nb = k / QK_ROCMFP4; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib*QK_ROCMFP4; + const float * qw = quant_weights ? quant_weights + ib*QK_ROCMFP4 : NULL; + const uint8_t e0 = rocmfp4_choose_scale_ue4m3(xb, QK_ROCMFP4/2, qw, sigma2); + const uint8_t e1 = rocmfp4_choose_scale_ue4m3(xb + QK_ROCMFP4/2, QK_ROCMFP4/2, qw ? qw + QK_ROCMFP4/2 : NULL, sigma2); + const float scale_half0 = rocmfp4_ue4m3_to_fp32_half(e0); + const float scale_half1 = rocmfp4_ue4m3_to_fp32_half(e1); + const float inv_scale_half0 = scale_half0 > 0.0f ? 1.0f / scale_half0 : 0.0f; + const float inv_scale_half1 = scale_half1 > 0.0f ? 1.0f / scale_half1 : 0.0f; + + y[ib].e[0] = e0; + y[ib].e[1] = e1; + + for (int j = 0; j < QK_ROCMFP4/2; ++j) { + const uint8_t q0 = rocmfp4_best_index_scaled(xb[j], inv_scale_half0); + const uint8_t q1 = rocmfp4_best_index_scaled(xb[j + QK_ROCMFP4/2], inv_scale_half1); + y[ib].qs[j] = q0 | (q1 << 4); + } + } +} + +static void rocmfp4_quantize_row_q4_0_fast_weighted( + const float * GGML_RESTRICT x, block_rocmfp4_fast * GGML_RESTRICT y, int64_t k, const float * GGML_RESTRICT quant_weights) { + assert(k % QK_ROCMFP4 == 0); + + float sum_x2 = 0.0f; + for (int64_t i = 0; i < k; ++i) { + sum_x2 += x[i]*x[i]; + } + const float sigma2 = sum_x2 / (float) k; + + const int64_t nb = k / QK_ROCMFP4; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib*QK_ROCMFP4; + const float * qw = quant_weights ? quant_weights + ib*QK_ROCMFP4 : NULL; + const uint8_t e = rocmfp4_choose_scale_ue4m3(xb, QK_ROCMFP4, qw, sigma2); + const float scale_half = rocmfp4_ue4m3_to_fp32_half(e); + const float inv_scale_half = scale_half > 0.0f ? 1.0f / scale_half : 0.0f; + + y[ib].e = e; + + for (int j = 0; j < QK_ROCMFP4/2; ++j) { + const uint8_t q0 = rocmfp4_best_index_scaled(xb[j], inv_scale_half); + const uint8_t q1 = rocmfp4_best_index_scaled(xb[j + QK_ROCMFP4/2], inv_scale_half); + y[ib].qs[j] = q0 | (q1 << 4); + } + } +} + +void rocmfp4_quantize_row_q4_0_ref(const float * GGML_RESTRICT x, block_rocmfp4 * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP4 == 0); + + const int64_t nb = k / QK_ROCMFP4; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib*QK_ROCMFP4; + const uint8_t e0 = rocmfp4_choose_scale_ue4m3(xb, QK_ROCMFP4/2, NULL, 0.0f); + const uint8_t e1 = rocmfp4_choose_scale_ue4m3(xb + QK_ROCMFP4/2, QK_ROCMFP4/2, NULL, 0.0f); + const float scale_half0 = rocmfp4_ue4m3_to_fp32_half(e0); + const float scale_half1 = rocmfp4_ue4m3_to_fp32_half(e1); + const float inv_scale_half0 = scale_half0 > 0.0f ? 1.0f / scale_half0 : 0.0f; + const float inv_scale_half1 = scale_half1 > 0.0f ? 1.0f / scale_half1 : 0.0f; + + y[ib].e[0] = e0; + y[ib].e[1] = e1; + + for (int j = 0; j < QK_ROCMFP4/2; ++j) { + const uint8_t q0 = rocmfp4_best_index_scaled(xb[j], inv_scale_half0); + const uint8_t q1 = rocmfp4_best_index_scaled(xb[j + QK_ROCMFP4/2], inv_scale_half1); + y[ib].qs[j] = q0 | (q1 << 4); + } + } +} + +void rocmfp4_quantize_row_q4_0_fast_ref(const float * GGML_RESTRICT x, block_rocmfp4_fast * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP4 == 0); + + const int64_t nb = k / QK_ROCMFP4; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib*QK_ROCMFP4; + const uint8_t e = rocmfp4_choose_scale_ue4m3(xb, QK_ROCMFP4, NULL, 0.0f); + const float scale_half = rocmfp4_ue4m3_to_fp32_half(e); + const float inv_scale_half = scale_half > 0.0f ? 1.0f / scale_half : 0.0f; + + y[ib].e = e; + + for (int j = 0; j < QK_ROCMFP4/2; ++j) { + const uint8_t q0 = rocmfp4_best_index_scaled(xb[j], inv_scale_half); + const uint8_t q1 = rocmfp4_best_index_scaled(xb[j + QK_ROCMFP4/2], inv_scale_half); + y[ib].qs[j] = q0 | (q1 << 4); + } + } +} + +void rocmfp4_dequantize_row_q4_0(const block_rocmfp4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP4 == 0); + + const int64_t nb = k / QK_ROCMFP4; + for (int64_t ib = 0; ib < nb; ++ib) { + const float d0 = rocmfp4_ue4m3_to_fp32_half(x[ib].e[0]); + const float d1 = rocmfp4_ue4m3_to_fp32_half(x[ib].e[1]); + + for (int j = 0; j < QK_ROCMFP4/2; ++j) { + y[ib*QK_ROCMFP4 + j] = (float) rocmfp4_decode(x[ib].qs[j] & 0x0f) * d0; + y[ib*QK_ROCMFP4 + j + QK_ROCMFP4/2] = (float) rocmfp4_decode(x[ib].qs[j] >> 4) * d1; + } + } +} + +void rocmfp4_dequantize_row_q4_0_fast(const block_rocmfp4_fast * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP4 == 0); + + const int64_t nb = k / QK_ROCMFP4; + for (int64_t ib = 0; ib < nb; ++ib) { + const float d = rocmfp4_ue4m3_to_fp32_half(x[ib].e); + + for (int j = 0; j < QK_ROCMFP4/2; ++j) { + y[ib*QK_ROCMFP4 + j] = (float) rocmfp4_decode(x[ib].qs[j] & 0x0f) * d; + y[ib*QK_ROCMFP4 + j + QK_ROCMFP4/2] = (float) rocmfp4_decode(x[ib].qs[j] >> 4) * d; + } + } +} + +void rocmfp4_quantize_row_q4_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + rocmfp4_quantize_row_q4_0_ref(x, (block_rocmfp4 *) y, k); +} + +void rocmfp4_quantize_row_q4_0_fast(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + rocmfp4_quantize_row_q4_0_fast_ref(x, (block_rocmfp4_fast *) y, k); +} + +size_t rocmfp4_quantize_q4_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix) { + const size_t row_size = ggml_row_size(GGML_TYPE_Q4_0_ROCMFP4, n_per_row); + + if (!imatrix) { + rocmfp4_quantize_row_q4_0_ref(src, (block_rocmfp4 *) dst, nrows*n_per_row); + return nrows * row_size; + } + + char * qrow = (char *) dst; + for (int64_t row = 0; row < nrows; ++row) { + rocmfp4_quantize_row_q4_0_weighted(src, (block_rocmfp4 *) qrow, n_per_row, imatrix); + src += n_per_row; + qrow += row_size; + } + + return nrows * row_size; +} + +size_t rocmfp4_quantize_q4_0_fast(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix) { + const size_t row_size = ggml_row_size(GGML_TYPE_Q4_0_ROCMFP4_FAST, n_per_row); + + if (!imatrix) { + rocmfp4_quantize_row_q4_0_fast_ref(src, (block_rocmfp4_fast *) dst, nrows*n_per_row); + return nrows * row_size; + } + + char * qrow = (char *) dst; + for (int64_t row = 0; row < nrows; ++row) { + rocmfp4_quantize_row_q4_0_fast_weighted(src, (block_rocmfp4_fast *) qrow, n_per_row, imatrix); + src += n_per_row; + qrow += row_size; + } + + return nrows * row_size; +} + +bool rocmfp4_validate_row_data(const void * data, size_t nbytes) { + if (nbytes % sizeof(block_rocmfp4) != 0) { + return false; + } + + const block_rocmfp4 * blocks = (const block_rocmfp4 *) data; + const size_t nblocks = nbytes / sizeof(block_rocmfp4); + for (size_t i = 0; i < nblocks; ++i) { + if (!rocmfp4_scale_is_valid(blocks[i].e[0]) || !rocmfp4_scale_is_valid(blocks[i].e[1])) { + return false; + } + } + + return true; +} + +bool rocmfp4_validate_row_data_fast(const void * data, size_t nbytes) { + if (nbytes % sizeof(block_rocmfp4_fast) != 0) { + return false; + } + + const block_rocmfp4_fast * blocks = (const block_rocmfp4_fast *) data; + const size_t nblocks = nbytes / sizeof(block_rocmfp4_fast); + for (size_t i = 0; i < nblocks; ++i) { + if (!rocmfp4_scale_is_valid(blocks[i].e)) { + return false; + } + } + + return true; +} + +#ifdef ROCMFP4_X86_AVX2_DISPATCH +__attribute__((target("avx2"))) +static inline int rocmfp4_hsum_i32_8_avx2(__m256i v) { + __m128i s = _mm_add_epi32(_mm256_castsi256_si128(v), _mm256_extracti128_si256(v, 1)); + s = _mm_add_epi32(s, _mm_shuffle_epi32(s, _MM_SHUFFLE(1, 0, 3, 2))); + s = _mm_add_epi32(s, _mm_shuffle_epi32(s, _MM_SHUFFLE(2, 3, 0, 1))); + return _mm_cvtsi128_si32(s); +} + +// Decode one 32-weight block's low and high nibble streams through the +// Codebook10 table with a single PSHUFB, then integer-dot each against its half +// of the q8_0 block. Integer sums are order-independent, so sumi0/sumi1 match +// the scalar reference exactly and the float result is bit-identical. +__attribute__((target("avx2"))) +static inline void rocmfp4_block_isums_avx2( + const uint8_t * qs, const int8_t * q8, int * sumi0, int * sumi1) { + const __m128i tbl = _mm_loadu_si128((const __m128i *) rocmfp4_codebook); + const __m128i q = _mm_loadu_si128((const __m128i *) qs); + const __m128i lo = _mm_and_si128(q, _mm_set1_epi8(0x0F)); + const __m128i hi = _mm_and_si128(_mm_srli_epi16(q, 4), _mm_set1_epi8(0x0F)); + const __m128i dlo = _mm_shuffle_epi8(tbl, lo); + const __m128i dhi = _mm_shuffle_epi8(tbl, hi); + const __m128i ylo = _mm_loadu_si128((const __m128i *) q8); + const __m128i yhi = _mm_loadu_si128((const __m128i *) (q8 + QK_ROCMFP4/2)); + const __m256i pl = _mm256_madd_epi16(_mm256_cvtepi8_epi16(dlo), _mm256_cvtepi8_epi16(ylo)); + const __m256i ph = _mm256_madd_epi16(_mm256_cvtepi8_epi16(dhi), _mm256_cvtepi8_epi16(yhi)); + *sumi0 = rocmfp4_hsum_i32_8_avx2(pl); + *sumi1 = rocmfp4_hsum_i32_8_avx2(ph); +} + +__attribute__((target("avx2"))) +static void rocmfp4_vec_dot_q4_0_q8_0_avx2( + int nb, float * GGML_RESTRICT s, const block_rocmfp4 * GGML_RESTRICT x, const block_q8_0 * GGML_RESTRICT y) { + float sumf = 0.0f; + for (int ib = 0; ib < nb; ++ib) { + const float d0 = rocmfp4_ue4m3_to_fp32_half(x[ib].e[0]) * ggml_fp16_to_fp32(y[ib].d); + const float d1 = rocmfp4_ue4m3_to_fp32_half(x[ib].e[1]) * ggml_fp16_to_fp32(y[ib].d); + int sumi0, sumi1; + rocmfp4_block_isums_avx2(x[ib].qs, y[ib].qs, &sumi0, &sumi1); + sumf += d0 * (float) sumi0 + d1 * (float) sumi1; + } + *s = sumf; +} + +__attribute__((target("avx2"))) +static void rocmfp4_vec_dot_q4_0_fast_q8_0_avx2( + int nb, float * GGML_RESTRICT s, const block_rocmfp4_fast * GGML_RESTRICT x, const block_q8_0 * GGML_RESTRICT y) { + float sumf = 0.0f; + for (int ib = 0; ib < nb; ++ib) { + const float d = rocmfp4_ue4m3_to_fp32_half(x[ib].e) * ggml_fp16_to_fp32(y[ib].d); + int sumi0, sumi1; + rocmfp4_block_isums_avx2(x[ib].qs, y[ib].qs, &sumi0, &sumi1); + sumf += d * (float) (sumi0 + sumi1); + } + *s = sumf; +} +#endif // ROCMFP4_X86_AVX2_DISPATCH + +void rocmfp4_vec_dot_q4_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); + assert(nrc == 1); + GGML_UNUSED(nrc); + assert(n % QK_ROCMFP4 == 0); + assert(QK_ROCMFP4 == QK8_0); + + const block_rocmfp4 * GGML_RESTRICT x = (const block_rocmfp4 *) vx; + const block_q8_0 * GGML_RESTRICT y = (const block_q8_0 *) vy; + + const int nb = n / QK_ROCMFP4; + +#ifdef ROCMFP4_X86_AVX2_DISPATCH + if (__builtin_cpu_supports("avx2")) { + rocmfp4_vec_dot_q4_0_q8_0_avx2(nb, s, x, y); + return; + } +#endif + + float sumf = 0.0f; + + for (int ib = 0; ib < nb; ++ib) { + const float d0 = rocmfp4_ue4m3_to_fp32_half(x[ib].e[0]) * ggml_fp16_to_fp32(y[ib].d); + const float d1 = rocmfp4_ue4m3_to_fp32_half(x[ib].e[1]) * ggml_fp16_to_fp32(y[ib].d); + int sumi0 = 0; + int sumi1 = 0; + + for (int j = 0; j < QK_ROCMFP4/2; ++j) { + const uint8_t q = x[ib].qs[j]; + sumi0 += rocmfp4_decode_table(q) * y[ib].qs[j]; + sumi1 += rocmfp4_decode_table(q >> 4) * y[ib].qs[j + QK_ROCMFP4/2]; + } + + sumf += d0 * (float) sumi0 + d1 * (float) sumi1; + } + + *s = sumf; +} + +void rocmfp4_vec_dot_q4_0_fast_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); + assert(nrc == 1); + GGML_UNUSED(nrc); + assert(n % QK_ROCMFP4 == 0); + assert(QK_ROCMFP4 == QK8_0); + + const block_rocmfp4_fast * GGML_RESTRICT x = (const block_rocmfp4_fast *) vx; + const block_q8_0 * GGML_RESTRICT y = (const block_q8_0 *) vy; + + const int nb = n / QK_ROCMFP4; + +#ifdef ROCMFP4_X86_AVX2_DISPATCH + if (__builtin_cpu_supports("avx2")) { + rocmfp4_vec_dot_q4_0_fast_q8_0_avx2(nb, s, x, y); + return; + } +#endif + + float sumf = 0.0f; + + for (int ib = 0; ib < nb; ++ib) { + const float d = rocmfp4_ue4m3_to_fp32_half(x[ib].e) * ggml_fp16_to_fp32(y[ib].d); + int sumi = 0; + + for (int j = 0; j < QK_ROCMFP4/2; ++j) { + const uint8_t q = x[ib].qs[j]; + sumi += rocmfp4_decode_table(q) * y[ib].qs[j]; + sumi += rocmfp4_decode_table(q >> 4) * y[ib].qs[j + QK_ROCMFP4/2]; + } + + sumf += d * (float) sumi; + } + + *s = sumf; +} diff --git a/ggml/rocmfp4/rocmfp4.h b/ggml/rocmfp4/rocmfp4.h new file mode 100644 index 000000000000..9756f6ad4c3b --- /dev/null +++ b/ggml/rocmfp4/rocmfp4.h @@ -0,0 +1,58 @@ +#pragma once + +#include +#include +#include + +#include "ggml.h" + +#ifdef __cplusplus +extern "C" { +#endif + +#define QK_ROCMFP4 32 +#define QR_ROCMFP4 2 +#define QI_ROCMFP4 (QK_ROCMFP4 / (4 * QR_ROCMFP4)) +#define QS_ROCMFP4 32 + +// AMD-tuned compact layout: 16 bytes of packed E2M1-derived 4-bit codes, then +// one unsigned E4M3 scale byte per 16-weight half block. +typedef struct { + uint8_t qs[QK_ROCMFP4/2]; + uint8_t e[2]; +} block_rocmfp4; + +// Speed-focused layout: same 32 packed ROCmFP4 nibbles, but one UE4M3 scale +// for the whole block. This is a separate GGUF type so fast 4.25 BPW artifacts +// never alias the safer dual-scale format above. +typedef struct { + uint8_t qs[QK_ROCMFP4/2]; + uint8_t e; +} block_rocmfp4_fast; + +#if defined(__cplusplus) +static_assert(sizeof(block_rocmfp4) == QK_ROCMFP4/2 + 2*sizeof(uint8_t), "wrong rocmfp4 block size/padding"); +static_assert(sizeof(block_rocmfp4_fast) == QK_ROCMFP4/2 + sizeof(uint8_t), "wrong rocmfp4 fast block size/padding"); +#else +_Static_assert(sizeof(block_rocmfp4) == QK_ROCMFP4/2 + 2*sizeof(uint8_t), "wrong rocmfp4 block size/padding"); +_Static_assert(sizeof(block_rocmfp4_fast) == QK_ROCMFP4/2 + sizeof(uint8_t), "wrong rocmfp4 fast block size/padding"); +#endif + +GGML_API void rocmfp4_quantize_row_q4_0_ref(const float * GGML_RESTRICT x, block_rocmfp4 * GGML_RESTRICT y, int64_t k); +GGML_API void rocmfp4_dequantize_row_q4_0(const block_rocmfp4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void rocmfp4_quantize_row_q4_0_fast_ref(const float * GGML_RESTRICT x, block_rocmfp4_fast * GGML_RESTRICT y, int64_t k); +GGML_API void rocmfp4_dequantize_row_q4_0_fast(const block_rocmfp4_fast * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); + +GGML_API void rocmfp4_quantize_row_q4_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +GGML_API size_t rocmfp4_quantize_q4_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); +GGML_API void rocmfp4_quantize_row_q4_0_fast(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +GGML_API size_t rocmfp4_quantize_q4_0_fast(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); +GGML_API bool rocmfp4_validate_row_data(const void * data, size_t nbytes); +GGML_API bool rocmfp4_validate_row_data_fast(const void * data, size_t nbytes); + +GGML_API void rocmfp4_vec_dot_q4_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +GGML_API void rocmfp4_vec_dot_q4_0_fast_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +#ifdef __cplusplus +} +#endif diff --git a/ggml/rocmfp4/rocmfp4_hip.cu b/ggml/rocmfp4/rocmfp4_hip.cu new file mode 100644 index 000000000000..d2c9048c1876 --- /dev/null +++ b/ggml/rocmfp4/rocmfp4_hip.cu @@ -0,0 +1,85 @@ +#include "rocmfp4.h" + +#include + +#include "rocmfp4_hip_scale.cuh" + +// Standalone ROCm/HIP dequant kernel for integration tests and future fused +// paths. One lane owns one packed byte and writes the matching low/high +// half-block values, so each byte is read once. +extern "C" __global__ void rocmfp4_dequantize_q4_0_f32_kernel( + const block_rocmfp4 * __restrict__ x, + float * __restrict__ y, + int64_t k) { + const int64_t packed_idx = (int64_t) blockIdx.x*blockDim.x + threadIdx.x; + const int64_t nblocks = (k + QK_ROCMFP4 - 1) / QK_ROCMFP4; + const int64_t packed_count = nblocks * (QK_ROCMFP4/2); + + if (packed_idx >= packed_count) { + return; + } + + const int64_t ib = packed_idx / (QK_ROCMFP4/2); + const int tid = packed_idx - ib*(QK_ROCMFP4/2); + const int64_t base = ib*QK_ROCMFP4; + const uint8_t packed = x[ib].qs[tid]; + const float d0 = rocmfp4_ue4m3_to_fp32_half_finite(x[ib].e[0]); + const float d1 = rocmfp4_ue4m3_to_fp32_half_finite(x[ib].e[1]); + + if (base + tid < k) { + y[base + tid] = (float) rocmfp4_decode_i8(packed & 0x0f) * d0; + } + if (base + tid + QK_ROCMFP4/2 < k) { + y[base + tid + QK_ROCMFP4/2] = (float) rocmfp4_decode_i8(packed >> 4) * d1; + } +} + +extern "C" __global__ void rocmfp4_dequantize_q4_0_fast_f32_kernel( + const block_rocmfp4_fast * __restrict__ x, + float * __restrict__ y, + int64_t k) { + const int64_t packed_idx = (int64_t) blockIdx.x*blockDim.x + threadIdx.x; + const int64_t nblocks = (k + QK_ROCMFP4 - 1) / QK_ROCMFP4; + const int64_t packed_count = nblocks * (QK_ROCMFP4/2); + + if (packed_idx >= packed_count) { + return; + } + + const int64_t ib = packed_idx / (QK_ROCMFP4/2); + const int tid = packed_idx - ib*(QK_ROCMFP4/2); + const int64_t base = ib*QK_ROCMFP4; + const uint8_t packed = x[ib].qs[tid]; + const float d = rocmfp4_ue4m3_to_fp32_half_finite(x[ib].e); + + if (base + tid < k) { + y[base + tid] = (float) rocmfp4_decode_i8(packed & 0x0f) * d; + } + if (base + tid + QK_ROCMFP4/2 < k) { + y[base + tid + QK_ROCMFP4/2] = (float) rocmfp4_decode_i8(packed >> 4) * d; + } +} + +extern "C" void rocmfp4_hip_dequantize_q4_0_to_f32( + const void * src, + float * dst, + int64_t k, + hipStream_t stream) { + const int64_t nblocks = (k + QK_ROCMFP4 - 1) / QK_ROCMFP4; + const int64_t packed_count = nblocks * (QK_ROCMFP4/2); + const dim3 block(256); + const dim3 grid((unsigned int) ((packed_count + block.x - 1) / block.x)); + rocmfp4_dequantize_q4_0_f32_kernel<<>>((const block_rocmfp4 *) src, dst, k); +} + +extern "C" void rocmfp4_hip_dequantize_q4_0_fast_to_f32( + const void * src, + float * dst, + int64_t k, + hipStream_t stream) { + const int64_t nblocks = (k + QK_ROCMFP4 - 1) / QK_ROCMFP4; + const int64_t packed_count = nblocks * (QK_ROCMFP4/2); + const dim3 block(256); + const dim3 grid((unsigned int) ((packed_count + block.x - 1) / block.x)); + rocmfp4_dequantize_q4_0_fast_f32_kernel<<>>((const block_rocmfp4_fast *) src, dst, k); +} diff --git a/ggml/rocmfp4/rocmfp4_hip_codebook.cuh b/ggml/rocmfp4/rocmfp4_hip_codebook.cuh new file mode 100644 index 000000000000..1d91061129a3 --- /dev/null +++ b/ggml/rocmfp4/rocmfp4_hip_codebook.cuh @@ -0,0 +1,79 @@ +#pragma once + +#include "rocmfp4_hip_scale.cuh" + +#include +#include + +#ifndef GGML_ROCMFP4_UNALIGNED_QS_DWORD_LOAD +#define GGML_ROCMFP4_UNALIGNED_QS_DWORD_LOAD 1 +#endif + +static __device__ __forceinline__ int rocmfp4_get_qs_i32(const void * x, const int & i32) { +#if defined(GGML_USE_HIP) && GGML_ROCMFP4_UNALIGNED_QS_DWORD_LOAD + return *((const int *) ((const uint8_t *) x + 4*i32)); +#else + const uint8_t * x8 = (const uint8_t *) x; + + int x32 = x8[4*i32 + 0] << 0; + x32 |= x8[4*i32 + 1] << 8; + x32 |= x8[4*i32 + 2] << 16; + x32 |= x8[4*i32 + 3] << 24; + + return x32; +#endif +} + +// AMD-specific fast path for expanding eight packed ROCmFP4 nibbles into two +// int32 DP4A operands. This encodes the Codebook10 table directly as four +// 32-bit constants: +// [0, 1, 2, 3], [4, 6, 8, 10], [0, -1, -2, -3], [-4, -6, -8, -10] +// Avoiding the table pointer keeps the ROCm/HIP MMVQ/MMQ hot path fully local +// to this format. Non-HIP builds still use llama.cpp's generic table expander. +static __device__ __forceinline__ int2 rocmfp4_get_int_from_codebook_16(const int & q4, const int8_t * fallback_table) { +#if defined(GGML_USE_HIP) + constexpr uint32_t values0 = 0x03020100u; + constexpr uint32_t values1 = 0x0a080604u; + constexpr uint32_t values2 = 0xfdfeff00u; + constexpr uint32_t values3 = 0xf6f8fafcu; + + const uint32_t q_even = q4; + const uint32_t q_odd = q4 >> 4; + + const uint32_t v_even_low = __builtin_amdgcn_perm(values1, values0, q_even & 0x07070707u); + const uint32_t v_odd_low = __builtin_amdgcn_perm(values1, values0, q_odd & 0x07070707u); + const uint32_t v_even_high = __builtin_amdgcn_perm(values3, values2, q_even & 0x07070707u); + const uint32_t v_odd_high = __builtin_amdgcn_perm(values3, values2, q_odd & 0x07070707u); + + const uint32_t mask_even = 0x03020100u | ((q_even & 0x08080808u) >> 1); + const uint32_t mask_odd = 0x03020100u | ((q_odd & 0x08080808u) >> 1); + + return make_int2( + __builtin_amdgcn_perm(v_even_high, v_even_low, mask_even), + __builtin_amdgcn_perm(v_odd_high, v_odd_low, mask_odd)); +#else + return get_int_from_table_16(q4, fallback_table); +#endif +} + +// Variant for call sites that already selected either the low or high nibble +// stream and only need one DP4A operand. This avoids the extra odd/even table +// expansion work in ROCmFP4 FlashAttention K/V decode. +static __device__ __forceinline__ int rocmfp4_get_low_int_from_codebook_16(const int & q4, const int8_t * fallback_table) { +#if defined(GGML_USE_HIP) + constexpr uint32_t values0 = 0x03020100u; + constexpr uint32_t values1 = 0x0a080604u; + constexpr uint32_t values2 = 0xfdfeff00u; + constexpr uint32_t values3 = 0xf6f8fafcu; + + const uint32_t q = q4; + + const uint32_t v_low = __builtin_amdgcn_perm(values1, values0, q & 0x07070707u); + const uint32_t v_high = __builtin_amdgcn_perm(values3, values2, q & 0x07070707u); + const uint32_t mask = 0x03020100u | ((q & 0x08080808u) >> 1); + + return __builtin_amdgcn_perm(v_high, v_low, mask); +#else + return get_int_from_table_16(q4, fallback_table).x; +#endif +} diff --git a/ggml/rocmfp4/rocmfp4_hip_scale.cuh b/ggml/rocmfp4/rocmfp4_hip_scale.cuh new file mode 100644 index 000000000000..19525a4fb017 --- /dev/null +++ b/ggml/rocmfp4/rocmfp4_hip_scale.cuh @@ -0,0 +1,157 @@ +#pragma once + +#include +#include +#include + +#ifndef GGML_ROCMFP4_USE_SCALE_LUT +#define GGML_ROCMFP4_USE_SCALE_LUT 0 +#endif + +#if defined(GGML_USE_HIP) && GGML_ROCMFP4_USE_SCALE_LUT +#define ROCMFP4_SCALE_SUB(M) ((M) * 0x1p-10f) +#define ROCMFP4_SCALE_E1(M) ((8 + (M)) * 0x1p-10f) +#define ROCMFP4_SCALE_E2(M) ((8 + (M)) * 0x1p-9f) +#define ROCMFP4_SCALE_E3(M) ((8 + (M)) * 0x1p-8f) +#define ROCMFP4_SCALE_E4(M) ((8 + (M)) * 0x1p-7f) +#define ROCMFP4_SCALE_E5(M) ((8 + (M)) * 0x1p-6f) +#define ROCMFP4_SCALE_E6(M) ((8 + (M)) * 0x1p-5f) +#define ROCMFP4_SCALE_E7(M) ((8 + (M)) * 0x1p-4f) +#define ROCMFP4_SCALE_E8(M) ((8 + (M)) * 0x1p-3f) +#define ROCMFP4_SCALE_E9(M) ((8 + (M)) * 0x1p-2f) +#define ROCMFP4_SCALE_E10(M) ((8 + (M)) * 0x1p-1f) +#define ROCMFP4_SCALE_E11(M) ((8 + (M)) * 0x1p0f) +#define ROCMFP4_SCALE_E12(M) ((8 + (M)) * 0x1p1f) +#define ROCMFP4_SCALE_E13(M) ((8 + (M)) * 0x1p2f) +#define ROCMFP4_SCALE_E14(M) ((8 + (M)) * 0x1p3f) +#define ROCMFP4_SCALE_E15(M) ((8 + (M)) * 0x1p4f) + +static __device__ __constant__ const float rocmfp4_scale_ue4m3_half_lut[127] = { + ROCMFP4_SCALE_SUB(0), ROCMFP4_SCALE_SUB(1), ROCMFP4_SCALE_SUB(2), ROCMFP4_SCALE_SUB(3), + ROCMFP4_SCALE_SUB(4), ROCMFP4_SCALE_SUB(5), ROCMFP4_SCALE_SUB(6), ROCMFP4_SCALE_SUB(7), + ROCMFP4_SCALE_E1(0), ROCMFP4_SCALE_E1(1), ROCMFP4_SCALE_E1(2), ROCMFP4_SCALE_E1(3), + ROCMFP4_SCALE_E1(4), ROCMFP4_SCALE_E1(5), ROCMFP4_SCALE_E1(6), ROCMFP4_SCALE_E1(7), + ROCMFP4_SCALE_E2(0), ROCMFP4_SCALE_E2(1), ROCMFP4_SCALE_E2(2), ROCMFP4_SCALE_E2(3), + ROCMFP4_SCALE_E2(4), ROCMFP4_SCALE_E2(5), ROCMFP4_SCALE_E2(6), ROCMFP4_SCALE_E2(7), + ROCMFP4_SCALE_E3(0), ROCMFP4_SCALE_E3(1), ROCMFP4_SCALE_E3(2), ROCMFP4_SCALE_E3(3), + ROCMFP4_SCALE_E3(4), ROCMFP4_SCALE_E3(5), ROCMFP4_SCALE_E3(6), ROCMFP4_SCALE_E3(7), + ROCMFP4_SCALE_E4(0), ROCMFP4_SCALE_E4(1), ROCMFP4_SCALE_E4(2), ROCMFP4_SCALE_E4(3), + ROCMFP4_SCALE_E4(4), ROCMFP4_SCALE_E4(5), ROCMFP4_SCALE_E4(6), ROCMFP4_SCALE_E4(7), + ROCMFP4_SCALE_E5(0), ROCMFP4_SCALE_E5(1), ROCMFP4_SCALE_E5(2), ROCMFP4_SCALE_E5(3), + ROCMFP4_SCALE_E5(4), ROCMFP4_SCALE_E5(5), ROCMFP4_SCALE_E5(6), ROCMFP4_SCALE_E5(7), + ROCMFP4_SCALE_E6(0), ROCMFP4_SCALE_E6(1), ROCMFP4_SCALE_E6(2), ROCMFP4_SCALE_E6(3), + ROCMFP4_SCALE_E6(4), ROCMFP4_SCALE_E6(5), ROCMFP4_SCALE_E6(6), ROCMFP4_SCALE_E6(7), + ROCMFP4_SCALE_E7(0), ROCMFP4_SCALE_E7(1), ROCMFP4_SCALE_E7(2), ROCMFP4_SCALE_E7(3), + ROCMFP4_SCALE_E7(4), ROCMFP4_SCALE_E7(5), ROCMFP4_SCALE_E7(6), ROCMFP4_SCALE_E7(7), + ROCMFP4_SCALE_E8(0), ROCMFP4_SCALE_E8(1), ROCMFP4_SCALE_E8(2), ROCMFP4_SCALE_E8(3), + ROCMFP4_SCALE_E8(4), ROCMFP4_SCALE_E8(5), ROCMFP4_SCALE_E8(6), ROCMFP4_SCALE_E8(7), + ROCMFP4_SCALE_E9(0), ROCMFP4_SCALE_E9(1), ROCMFP4_SCALE_E9(2), ROCMFP4_SCALE_E9(3), + ROCMFP4_SCALE_E9(4), ROCMFP4_SCALE_E9(5), ROCMFP4_SCALE_E9(6), ROCMFP4_SCALE_E9(7), + ROCMFP4_SCALE_E10(0), ROCMFP4_SCALE_E10(1), ROCMFP4_SCALE_E10(2), ROCMFP4_SCALE_E10(3), + ROCMFP4_SCALE_E10(4), ROCMFP4_SCALE_E10(5), ROCMFP4_SCALE_E10(6), ROCMFP4_SCALE_E10(7), + ROCMFP4_SCALE_E11(0), ROCMFP4_SCALE_E11(1), ROCMFP4_SCALE_E11(2), ROCMFP4_SCALE_E11(3), + ROCMFP4_SCALE_E11(4), ROCMFP4_SCALE_E11(5), ROCMFP4_SCALE_E11(6), ROCMFP4_SCALE_E11(7), + ROCMFP4_SCALE_E12(0), ROCMFP4_SCALE_E12(1), ROCMFP4_SCALE_E12(2), ROCMFP4_SCALE_E12(3), + ROCMFP4_SCALE_E12(4), ROCMFP4_SCALE_E12(5), ROCMFP4_SCALE_E12(6), ROCMFP4_SCALE_E12(7), + ROCMFP4_SCALE_E13(0), ROCMFP4_SCALE_E13(1), ROCMFP4_SCALE_E13(2), ROCMFP4_SCALE_E13(3), + ROCMFP4_SCALE_E13(4), ROCMFP4_SCALE_E13(5), ROCMFP4_SCALE_E13(6), ROCMFP4_SCALE_E13(7), + ROCMFP4_SCALE_E14(0), ROCMFP4_SCALE_E14(1), ROCMFP4_SCALE_E14(2), ROCMFP4_SCALE_E14(3), + ROCMFP4_SCALE_E14(4), ROCMFP4_SCALE_E14(5), ROCMFP4_SCALE_E14(6), ROCMFP4_SCALE_E14(7), + ROCMFP4_SCALE_E15(0), ROCMFP4_SCALE_E15(1), ROCMFP4_SCALE_E15(2), ROCMFP4_SCALE_E15(3), + ROCMFP4_SCALE_E15(4), ROCMFP4_SCALE_E15(5), ROCMFP4_SCALE_E15(6), +}; + +#undef ROCMFP4_SCALE_SUB +#undef ROCMFP4_SCALE_E1 +#undef ROCMFP4_SCALE_E2 +#undef ROCMFP4_SCALE_E3 +#undef ROCMFP4_SCALE_E4 +#undef ROCMFP4_SCALE_E5 +#undef ROCMFP4_SCALE_E6 +#undef ROCMFP4_SCALE_E7 +#undef ROCMFP4_SCALE_E8 +#undef ROCMFP4_SCALE_E9 +#undef ROCMFP4_SCALE_E10 +#undef ROCMFP4_SCALE_E11 +#undef ROCMFP4_SCALE_E12 +#undef ROCMFP4_SCALE_E13 +#undef ROCMFP4_SCALE_E14 +#undef ROCMFP4_SCALE_E15 +#endif + +static __device__ __forceinline__ float rocmfp4_u32_as_f32(uint32_t bits) { +#if defined(GGML_USE_HIP) + return __uint_as_float(bits); +#else + float result; + memcpy(&result, &bits, sizeof(float)); + return result; +#endif +} + +// ROCmFP4 validates scale bytes before backend execution, so HIP/ROCm hot +// paths can decode finite unsigned E4M3 half-scales directly without the +// generic FP8 NaN handling used by other formats. +static __device__ __forceinline__ float rocmfp4_ue4m3_to_fp32_half_finite(uint8_t x) { +#if defined(GGML_USE_HIP) && GGML_ROCMFP4_USE_SCALE_LUT + return x <= 0x7e ? rocmfp4_scale_ue4m3_half_lut[x] : 0.0f; +#else + const int exp = (x >> 3) & 0xF; + const int man = x & 0x7; + + if (exp == 0) { + return (float) man * (1.0f / 1024.0f); + } + + const uint32_t bits = ((uint32_t) exp + 119u) << 23 | ((uint32_t) man << 20); + return rocmfp4_u32_as_f32(bits); +#endif +} + +static __device__ __forceinline__ float rocmfpx_ue4m3_to_fp32_finite(uint8_t x) { + if (x > 0x7e) { + return 0.0f; + } + + const int exp = (x >> 3) & 0xF; + const int man = x & 0x7; + + if (exp == 0) { + return (float) man * (1.0f / 1024.0f); + } + + const uint32_t bits = ((uint32_t) exp + 119u) << 23 | ((uint32_t) man << 20); + return rocmfp4_u32_as_f32(bits); +} + +static __device__ __forceinline__ uint8_t rocmfpx_nearest_scale_ue4m3_cuda(float target_scale) { + if (!(target_scale > 0.0f) || !isfinite(target_scale)) { + return 0; + } + + uint8_t lo = 1; + uint8_t hi = 0x7e; + while (lo < hi) { + const uint8_t mid = lo + (hi - lo) / 2; + if (rocmfpx_ue4m3_to_fp32_finite(mid) < target_scale) { + lo = mid + 1; + } else { + hi = mid; + } + } + + if (lo == 1) { + return 1; + } + + const float hi_scale = rocmfpx_ue4m3_to_fp32_finite(lo); + const float lo_scale = rocmfpx_ue4m3_to_fp32_finite((uint8_t) (lo - 1)); + return (target_scale - lo_scale <= hi_scale - target_scale) ? (uint8_t) (lo - 1) : lo; +} + +static __device__ __forceinline__ int8_t rocmfp4_decode_i8(uint8_t q) { + q &= 0x0f; + const int mag3 = q & 0x07; + const int mag = mag3 <= 4 ? mag3 : 2*mag3 - 4; + return (q & 0x08) ? -mag : mag; +} diff --git a/ggml/rocmfpx/README.md b/ggml/rocmfpx/README.md new file mode 100644 index 000000000000..a92eb46584b2 --- /dev/null +++ b/ggml/rocmfpx/README.md @@ -0,0 +1,106 @@ +# ROCmFPx Experimental Formats + +This folder contains the reference layer for the proposed ROCmFP3, ROCmFP6, and +ROCmFP8 quantization family. It is intentionally separate from `ggml/rocmfp4/` +so the promoted ROCmFP4 GGUF formats and kernels are not affected while the new +layouts are evaluated. + +## ROCmFP4 Instructions To Preserve + +ROCmFPX is a sibling model-weight quant family, not a new K/V-only compression +scheme. The promoted ROCmFP4 implementation remains the template for how the +family should behave in llama.cpp: + +- Keep 32-weight blocks so CPU, HIP, and Vulkan kernels can reuse the same + Q4/Q8-style reduction shape and GGUF row-size assumptions. +- Use finite unsigned UE4M3 scale bytes only. `0x7f` and sign-bit scale bytes + are invalid, matching the ROCmFP4 validation rule. +- Prefer reconstruction-MSE scale selection over plain max-abs scaling. ROCmFP4 + searches each 16-weight half-block; ROCmFP3 and ROCmFP6 follow the same + half-block policy, while ROCmFP8 currently uses one full-block scale. +- Preserve the ROCmFP4 kernel contract: CPU reference quant/dequant/dot first, + then HIP/Vulkan `CPY`, `GET_ROWS`, `SET_ROWS`, `MUL_MAT`, and `MUL_MAT_ID` + paths, with backend-op coverage before claiming runtime support. +- Keep dequant math explicit and deterministic: integer code times decoded + UE4M3 scale. ROCmFP4 uses the Codebook10 half-scale table; ROCmFPX formats + use their own integer code ranges but must retain the same finite-scale and + integer-dot discipline. + +The ROCmFP4 Codebook10 levels are not reused by FP3/FP6/FP8 directly: +`ROCmFP3` uses `0, +/-1, +/-2, +/-4`, `ROCmFP6` uses signed-magnitude levels up +to `31`, and `ROCmFP8` uses signed int8 levels clamped to `[-127, 127]`. +What is inherited is the block/scale/kernel/dequant contract. + +Current status (June 16, 2026): +- CPU reference quantize/dequantize exists for all three formats. +- `Q3_0_ROCMFPX`, `Q6_0_ROCMFPX`, and `Q8_0_ROCMFPX` are registered as + experimental GGUF tensor types. +- ROCm/HIP and Vulkan kernels support `CPY`, `GET_ROWS`, `SET_ROWS`, and + `MUL_MAT`/`MUL_MAT_ID` for all three formats. +- Qwen3-0.6B BF16 smoke tests pass on CPU, ROCm0, and Vulkan0. +- Default quant presets now include lean coherency routing: + - `Q3_0_ROCMFPX`: selective `Q5_K` on attention Q/O and early K/V, boosted + FFN-down at `Q5_K`, selective FFN-gate at `Q6_0_ROCMFPX`, bulk FFN-up on + `Q3_0_ROCMFPX`, embeddings/output at `Q4_0_ROCMFP4_FAST`. + - `Q6_0_ROCMFPX`: early attention and boosted FFN-down at `Q8_0_ROCMFPX`, + embeddings/output at `Q6_0_ROCMFPX`, bulk gate/up on `Q6_0_ROCMFPX`. + - `Q8_0_ROCMFPX`: pure FP8-family preset. +- Opt-in `*_AGENT` presets boost attention/FFN routing for tool-call / + Hermes / OpenClaw style workloads: + - `Q3_0_ROCMFPX_AGENT`, `Q6_0_ROCMFPX_AGENT`, `Q8_0_ROCMFPX_AGENT`. + - Routing is layered on top of LEAN; default presets are unchanged. +- FP3 and FP6 quantization use reconstruction-MSE scale selection per + 16-weight half-block. + +## Validation Script Index + +```text +scripts/check-rocmfpx-reference.sh # CPU reference math +scripts/check-rocmfpx-qwen-all.sh # core Qwen gates +scripts/check-rocmfpx-all.sh # qwen-all + optional smokes +scripts/check-rocmfpx-summary.sh # full JSON summary runner +scripts/sweep-rocmfpx-backend-ops.sh # test-backend-ops per backend +scripts/sweep-rocmfpx-agent-size-table.sh # LEAN vs AGENT MiB/BPW +scripts/sweep-rocmfpx-perplexity.sh # calibration PPL sweep +scripts/sweep-rocmfpx-decode-tune.sh # decode-tune matrix +scripts/build-rocmfpx-agent-fixtures.sh # proxy Hermes/OpenClaw AGENT GGUFs +``` + +## Layouts + +All formats use 32-weight blocks. + +| Format | Payload | Scale bytes | Block bytes | BPW | Purpose | +|---|---:|---:|---:|---:|---| +| `Q3_0_ROCMFPX` | 32 packed 3-bit codes | 2, one per 16 weights | 14 | 3.50 | Experimental low-bit candidate | +| `Q6_0_ROCMFPX` | 32 packed 6-bit codes | 2, one per 16 weights | 26 | 6.50 | Experimental quality candidate | +| `Q8_0_ROCMFPX` | 32 signed 8-bit codes | 1, one per 32 weights | 33 | 8.25 | Experimental high-quality reference | + +`ROCmFP3` uses a tiny signed codebook: `0, +/-1, +/-2, +/-4`. +`ROCmFP6` uses signed-magnitude integer levels up to `31`. +`ROCmFP8` uses signed int8 levels clamped to `[-127, 127]`. + +## Validation + +Reference math only: + +```bash +scripts/check-rocmfpx-reference.sh +``` + +Focused backend sweep from the experiment worktree: + +```bash +cmake --build build-strix-rocmfp4 --target test-backend-ops -j 8 +timeout 120 build-strix-rocmfp4/bin/test-backend-ops test -o MUL_MAT,GET_ROWS,CPY,SET_ROWS -b CPU +timeout 180 build-strix-rocmfp4/bin/test-backend-ops test -o MUL_MAT,GET_ROWS,CPY,SET_ROWS -b ROCm0 +timeout 180 build-strix-rocmfp4/bin/test-backend-ops test -o MUL_MAT,GET_ROWS,CPY,SET_ROWS -b Vulkan0 +``` + +Qwen3 BF16 coherency and decode-speed gates: + +```bash +MODEL=/home/caf/strix-fp4/models/rocmfpx-bf16-tests/Qwen3-0.6B-Q3_0_ROCMFPX_COHERENT-LEAN.gguf BACKEND=ROCm0 scripts/check-rocmfpx-qwen-coherency.sh +MODEL=/home/caf/strix-fp4/models/rocmfpx-bf16-tests/Qwen3-0.6B-Q3_0_ROCMFPX_COHERENT-LEAN.gguf BACKEND=ROCm0 scripts/check-rocmfpx-qwen-bench.sh +MODEL=/home/caf/strix-fp4/models/rocmfpx-bf16-tests/Qwen3-0.6B-Q3_0_ROCMFPX_COHERENT-LEAN.gguf BACKEND=ROCm0 scripts/check-rocmfpx-qwen-strict-json.sh +``` diff --git a/ggml/rocmfpx/rocmfpx.c b/ggml/rocmfpx/rocmfpx.c new file mode 100644 index 000000000000..65dbca72605a --- /dev/null +++ b/ggml/rocmfpx/rocmfpx.c @@ -0,0 +1,1140 @@ +#include "rocmfpx.h" + +#include +#include +#include + +// Finite unsigned E4M3 scale bytes decoded to FP32. Precomputed from the same +// exp/mant formula rocmfpx_ue4m3_to_fp32() used to evaluate with ldexpf(): +// exp == 0 -> mant * 2^-10 ; otherwise (8 + mant) * 2^(exp - 11). +// The scale search re-decodes candidate bytes for every block, and dequant +// decodes a scale for every element, so keeping this as a table (identical to +// the former per-call ldexpf result) removes the transcendental from both hot +// paths without changing any produced value. +#define ROCMFPX_SCALE_SUB(M) ((M) * 0x1p-10f) +#define ROCMFPX_SCALE_E(B, M) ((8 + (M)) * (B)) + +static const float rocmfpx_scale_ue4m3[127] = { + ROCMFPX_SCALE_SUB(0), ROCMFPX_SCALE_SUB(1), ROCMFPX_SCALE_SUB(2), ROCMFPX_SCALE_SUB(3), + ROCMFPX_SCALE_SUB(4), ROCMFPX_SCALE_SUB(5), ROCMFPX_SCALE_SUB(6), ROCMFPX_SCALE_SUB(7), + ROCMFPX_SCALE_E(0x1p-10f,0), ROCMFPX_SCALE_E(0x1p-10f,1), ROCMFPX_SCALE_E(0x1p-10f,2), ROCMFPX_SCALE_E(0x1p-10f,3), + ROCMFPX_SCALE_E(0x1p-10f,4), ROCMFPX_SCALE_E(0x1p-10f,5), ROCMFPX_SCALE_E(0x1p-10f,6), ROCMFPX_SCALE_E(0x1p-10f,7), + ROCMFPX_SCALE_E(0x1p-9f,0), ROCMFPX_SCALE_E(0x1p-9f,1), ROCMFPX_SCALE_E(0x1p-9f,2), ROCMFPX_SCALE_E(0x1p-9f,3), + ROCMFPX_SCALE_E(0x1p-9f,4), ROCMFPX_SCALE_E(0x1p-9f,5), ROCMFPX_SCALE_E(0x1p-9f,6), ROCMFPX_SCALE_E(0x1p-9f,7), + ROCMFPX_SCALE_E(0x1p-8f,0), ROCMFPX_SCALE_E(0x1p-8f,1), ROCMFPX_SCALE_E(0x1p-8f,2), ROCMFPX_SCALE_E(0x1p-8f,3), + ROCMFPX_SCALE_E(0x1p-8f,4), ROCMFPX_SCALE_E(0x1p-8f,5), ROCMFPX_SCALE_E(0x1p-8f,6), ROCMFPX_SCALE_E(0x1p-8f,7), + ROCMFPX_SCALE_E(0x1p-7f,0), ROCMFPX_SCALE_E(0x1p-7f,1), ROCMFPX_SCALE_E(0x1p-7f,2), ROCMFPX_SCALE_E(0x1p-7f,3), + ROCMFPX_SCALE_E(0x1p-7f,4), ROCMFPX_SCALE_E(0x1p-7f,5), ROCMFPX_SCALE_E(0x1p-7f,6), ROCMFPX_SCALE_E(0x1p-7f,7), + ROCMFPX_SCALE_E(0x1p-6f,0), ROCMFPX_SCALE_E(0x1p-6f,1), ROCMFPX_SCALE_E(0x1p-6f,2), ROCMFPX_SCALE_E(0x1p-6f,3), + ROCMFPX_SCALE_E(0x1p-6f,4), ROCMFPX_SCALE_E(0x1p-6f,5), ROCMFPX_SCALE_E(0x1p-6f,6), ROCMFPX_SCALE_E(0x1p-6f,7), + ROCMFPX_SCALE_E(0x1p-5f,0), ROCMFPX_SCALE_E(0x1p-5f,1), ROCMFPX_SCALE_E(0x1p-5f,2), ROCMFPX_SCALE_E(0x1p-5f,3), + ROCMFPX_SCALE_E(0x1p-5f,4), ROCMFPX_SCALE_E(0x1p-5f,5), ROCMFPX_SCALE_E(0x1p-5f,6), ROCMFPX_SCALE_E(0x1p-5f,7), + ROCMFPX_SCALE_E(0x1p-4f,0), ROCMFPX_SCALE_E(0x1p-4f,1), ROCMFPX_SCALE_E(0x1p-4f,2), ROCMFPX_SCALE_E(0x1p-4f,3), + ROCMFPX_SCALE_E(0x1p-4f,4), ROCMFPX_SCALE_E(0x1p-4f,5), ROCMFPX_SCALE_E(0x1p-4f,6), ROCMFPX_SCALE_E(0x1p-4f,7), + ROCMFPX_SCALE_E(0x1p-3f,0), ROCMFPX_SCALE_E(0x1p-3f,1), ROCMFPX_SCALE_E(0x1p-3f,2), ROCMFPX_SCALE_E(0x1p-3f,3), + ROCMFPX_SCALE_E(0x1p-3f,4), ROCMFPX_SCALE_E(0x1p-3f,5), ROCMFPX_SCALE_E(0x1p-3f,6), ROCMFPX_SCALE_E(0x1p-3f,7), + ROCMFPX_SCALE_E(0x1p-2f,0), ROCMFPX_SCALE_E(0x1p-2f,1), ROCMFPX_SCALE_E(0x1p-2f,2), ROCMFPX_SCALE_E(0x1p-2f,3), + ROCMFPX_SCALE_E(0x1p-2f,4), ROCMFPX_SCALE_E(0x1p-2f,5), ROCMFPX_SCALE_E(0x1p-2f,6), ROCMFPX_SCALE_E(0x1p-2f,7), + ROCMFPX_SCALE_E(0x1p-1f,0), ROCMFPX_SCALE_E(0x1p-1f,1), ROCMFPX_SCALE_E(0x1p-1f,2), ROCMFPX_SCALE_E(0x1p-1f,3), + ROCMFPX_SCALE_E(0x1p-1f,4), ROCMFPX_SCALE_E(0x1p-1f,5), ROCMFPX_SCALE_E(0x1p-1f,6), ROCMFPX_SCALE_E(0x1p-1f,7), + ROCMFPX_SCALE_E(0x1p0f,0), ROCMFPX_SCALE_E(0x1p0f,1), ROCMFPX_SCALE_E(0x1p0f,2), ROCMFPX_SCALE_E(0x1p0f,3), + ROCMFPX_SCALE_E(0x1p0f,4), ROCMFPX_SCALE_E(0x1p0f,5), ROCMFPX_SCALE_E(0x1p0f,6), ROCMFPX_SCALE_E(0x1p0f,7), + ROCMFPX_SCALE_E(0x1p1f,0), ROCMFPX_SCALE_E(0x1p1f,1), ROCMFPX_SCALE_E(0x1p1f,2), ROCMFPX_SCALE_E(0x1p1f,3), + ROCMFPX_SCALE_E(0x1p1f,4), ROCMFPX_SCALE_E(0x1p1f,5), ROCMFPX_SCALE_E(0x1p1f,6), ROCMFPX_SCALE_E(0x1p1f,7), + ROCMFPX_SCALE_E(0x1p2f,0), ROCMFPX_SCALE_E(0x1p2f,1), ROCMFPX_SCALE_E(0x1p2f,2), ROCMFPX_SCALE_E(0x1p2f,3), + ROCMFPX_SCALE_E(0x1p2f,4), ROCMFPX_SCALE_E(0x1p2f,5), ROCMFPX_SCALE_E(0x1p2f,6), ROCMFPX_SCALE_E(0x1p2f,7), + ROCMFPX_SCALE_E(0x1p3f,0), ROCMFPX_SCALE_E(0x1p3f,1), ROCMFPX_SCALE_E(0x1p3f,2), ROCMFPX_SCALE_E(0x1p3f,3), + ROCMFPX_SCALE_E(0x1p3f,4), ROCMFPX_SCALE_E(0x1p3f,5), ROCMFPX_SCALE_E(0x1p3f,6), ROCMFPX_SCALE_E(0x1p3f,7), + ROCMFPX_SCALE_E(0x1p4f,0), ROCMFPX_SCALE_E(0x1p4f,1), ROCMFPX_SCALE_E(0x1p4f,2), ROCMFPX_SCALE_E(0x1p4f,3), + ROCMFPX_SCALE_E(0x1p4f,4), ROCMFPX_SCALE_E(0x1p4f,5), ROCMFPX_SCALE_E(0x1p4f,6), +}; + +#undef ROCMFPX_SCALE_SUB +#undef ROCMFPX_SCALE_E + +float rocmfpx_ue4m3_to_fp32(uint8_t e) { + return rocmfpx_scale_is_valid(e) ? rocmfpx_scale_ue4m3[e] : 0.0f; +} + +bool rocmfpx_scale_is_valid(uint8_t e) { + return e <= 0x7e; +} + +size_t rocmfpx_row_size_fp2(int64_t k) { + assert(k % QK_ROCMFP2 == 0); + return (size_t) (k / QK_ROCMFP2) * sizeof(block_rocmfp2); +} + +size_t rocmfpx_row_size_fp3(int64_t k) { + assert(k % QK_ROCMFP3 == 0); + return (size_t) (k / QK_ROCMFP3) * sizeof(block_rocmfp3); +} + +size_t rocmfpx_row_size_fp6(int64_t k) { + assert(k % QK_ROCMFP6 == 0); + return (size_t) (k / QK_ROCMFP6) * sizeof(block_rocmfp6); +} + +size_t rocmfpx_row_size_fp8(int64_t k) { + assert(k % QK_ROCMFP8 == 0); + return (size_t) (k / QK_ROCMFP8) * sizeof(block_rocmfp8); +} + +static uint8_t rocmfpx_nearest_scale_ue4m3(float target) { + if (!(target > 0.0f) || !isfinite(target)) { + return 0; + } + + uint8_t best_e = 1; + float best_err = fabsf(rocmfpx_ue4m3_to_fp32(best_e) - target); + + for (int e = 2; e <= 0x7e; ++e) { + const float err = fabsf(rocmfpx_ue4m3_to_fp32((uint8_t) e) - target); + if (err < best_err) { + best_err = err; + best_e = (uint8_t) e; + } + } + + return best_e; +} + +static float rocmfpx_max_abs(const float * x, int n) { + float max_abs = 0.0f; + + for (int i = 0; i < n; ++i) { + if (!isfinite(x[i])) { + continue; + } + + const float ax = fabsf(x[i]); + if (ax > max_abs) { + max_abs = ax; + } + } + + return max_abs; +} + +static void rocmfpx_prepare_mse_weights( + float * dst, const float * x, int n, const float * quant_weights, float sigma2, + float * max_abs, float * max_abs_weight) { + *max_abs = 0.0f; + *max_abs_weight = 0.0f; + + for (int i = 0; i < n; ++i) { + const float ax = fabsf(x[i]); + const float qw = quant_weights[i]; + const float weight = isfinite(qw) && qw > 0.0f && isfinite(x[i]) ? qw * sqrtf(sigma2 + x[i]*x[i]) : 0.0f; + + if (isfinite(x[i])) { + if (ax > *max_abs) { + *max_abs = ax; + *max_abs_weight = weight; + } else if (ax == *max_abs && weight > *max_abs_weight) { + *max_abs_weight = weight; + } + } + + // Match llama.cpp imatrix weighting style: calibration importance is + // scaled by row energy so large activations stay protected. + dst[i] = weight; + } +} + +static void rocmfpx_set_bits(uint8_t * dst, int bit_pos, int nbits, uint32_t code) { + for (int bit = 0; bit < nbits; ++bit) { + const int absolute_bit = bit_pos + bit; + const int byte_index = absolute_bit >> 3; + const int bit_index = absolute_bit & 7; + + if ((code >> bit) & 1u) { + dst[byte_index] |= (uint8_t) (1u << bit_index); + } + } +} + +static uint32_t rocmfpx_get_bits(const uint8_t * src, int bit_pos, int nbits) { + uint32_t code = 0; + + for (int bit = 0; bit < nbits; ++bit) { + const int absolute_bit = bit_pos + bit; + const int byte_index = absolute_bit >> 3; + const int bit_index = absolute_bit & 7; + + code |= (uint32_t) ((src[byte_index] >> bit_index) & 1u) << bit; + } + + return code; +} + +// Starting 2-bit ROCmFP2 codebook. Keep this single definition easy to tune. +// TODO: affine scale+min would likely improve quality, but it would break the +// ROCmFPx family's unsigned-UE4M3 scale contract; revisit for a v2 layout. +static const float kvalues_rocmfp2[4] = ROCMFP2_KVALUES_INIT; + +static uint8_t rocmfpx_quantize_fp2_code(float x, float inv_scale) { + if (!isfinite(x) || inv_scale <= 0.0f) { + return 1; + } + + const float q = x * inv_scale; + uint8_t best_code = 0; + float best_err = fabsf(q - kvalues_rocmfp2[0]); + + for (uint8_t code = 1; code < 4; ++code) { + const float err = fabsf(q - kvalues_rocmfp2[code]); + if (err < best_err) { + best_err = err; + best_code = code; + } + } + + return best_code; +} + +static inline float rocmfpx_fp2_decoded_mag(float x, float inv_scale) { + return kvalues_rocmfp2[rocmfpx_quantize_fp2_code(x, inv_scale)]; +} + +static float rocmfpx_fp2_block_mse_for_scale(const float * x, int n, uint8_t e, float best_err) { + const float scale = rocmfpx_ue4m3_to_fp32(e); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + if (!isfinite(x[i])) { + continue; + } + + const float y = rocmfpx_fp2_decoded_mag(x[i], inv_scale) * scale; + const float d = x[i] - y; + + err += d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +static float rocmfpx_fp2_block_weighted_mse_for_scale(const float * x, int n, const float * mse_weights, uint8_t e, float best_err) { + const float scale = rocmfpx_ue4m3_to_fp32(e); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + if (!isfinite(x[i])) { + continue; + } + + const float y = rocmfpx_fp2_decoded_mag(x[i], inv_scale) * scale; + const float d = x[i] - y; + + err += mse_weights[i]*d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +static uint8_t rocmfpx_choose_scale_fp2_mse_impl(const float * x, int n, const float * mse_weights, float max_abs, float max_abs_weight) { + const uint8_t start_e = rocmfpx_nearest_scale_ue4m3(max_abs / 2.0f); + uint8_t best_e = start_e; + float best_err = INFINITY; + bool lower_done = false; + + for (int delta = 0; delta <= 125; ++delta) { + const int e0 = (int) start_e - delta; + if (!lower_done && e0 >= 1 && e0 <= 126) { + const float scale = rocmfpx_ue4m3_to_fp32((uint8_t) e0); + const float clip_delta = max_abs - 2.0f*scale; + const float clip_err = mse_weights ? max_abs_weight*clip_delta*clip_delta : clip_delta*clip_delta; + if (clip_delta > 0.0f && clip_err > best_err) { + lower_done = true; + } else { + const float err = mse_weights ? + rocmfpx_fp2_block_weighted_mse_for_scale(x, n, mse_weights, (uint8_t) e0, best_err) : + rocmfpx_fp2_block_mse_for_scale(x, n, (uint8_t) e0, best_err); + if (err < best_err || (err == best_err && e0 < best_e)) { + best_err = err; + best_e = (uint8_t) e0; + } + } + } + + const int e1 = (int) start_e + delta; + if (delta != 0 && e1 >= 1 && e1 <= 126) { + const float err = mse_weights ? + rocmfpx_fp2_block_weighted_mse_for_scale(x, n, mse_weights, (uint8_t) e1, best_err) : + rocmfpx_fp2_block_mse_for_scale(x, n, (uint8_t) e1, best_err); + if (err < best_err || (err == best_err && e1 < best_e)) { + best_err = err; + best_e = (uint8_t) e1; + } + } + + if ((lower_done || e0 <= 1) && e1 >= 126) { + break; + } + } + + return best_e; +} + +static uint8_t rocmfpx_choose_scale_fp2_mse(const float * x, int n) { + const float max_abs = rocmfpx_max_abs(x, n); + if (!(max_abs > 0.0f) || !isfinite(max_abs)) { + return 0; + } + + return rocmfpx_choose_scale_fp2_mse_impl(x, n, NULL, max_abs, 0.0f); +} + +static uint8_t rocmfpx_choose_scale_fp2_weighted_mse(const float * x, int n, const float * quant_weights, float sigma2) { + assert(n <= QK_ROCMFP2); + float mse_weights[QK_ROCMFP2]; + float max_abs; + float max_abs_weight; + rocmfpx_prepare_mse_weights(mse_weights, x, n, quant_weights, sigma2, &max_abs, &max_abs_weight); + if (!(max_abs > 0.0f) || !isfinite(max_abs)) { + return 0; + } + + return rocmfpx_choose_scale_fp2_mse_impl(x, n, mse_weights, max_abs, max_abs_weight); +} + +static int rocmfpx_decode_fp3_code(uint8_t code) { + static const int mag[4] = { 0, 1, 2, 4 }; + const int value = mag[code & 3u]; + return (code & 4u) ? -value : value; +} + +static uint8_t rocmfpx_quantize_fp3_code(float x, float inv_scale) { + if (!isfinite(x) || inv_scale <= 0.0f) { + return 0; + } + + const float ax = fabsf(x * inv_scale); + uint8_t mag; + + if (ax <= 0.5f) { + mag = 0; + } else if (ax <= 1.5f) { + mag = 1; + } else if (ax <= 3.0f) { + mag = 2; + } else { + mag = 3; + } + + return mag == 0 ? 0 : (uint8_t) ((x < 0.0f ? 4u : 0u) | mag); +} + +// Fused threshold + decode used only inside the exhaustive scale search, which +// re-scans every element for every candidate scale byte. Returns the same +// signed decoded magnitude that +// rocmfpx_decode_fp3_code(rocmfpx_quantize_fp3_code(x, inv_scale)) produces +// (fp3 magnitudes {0,1,2,4}), so quantized output stays bit-identical. +static inline float rocmfpx_fp3_decoded_mag(float x, float inv_scale) { + const float a = fabsf(x * inv_scale); + float mag; + if (a <= 0.5f) { + return 0.0f; + } else if (a <= 1.5f) { + mag = 1.0f; + } else if (a <= 3.0f) { + mag = 2.0f; + } else { + mag = 4.0f; + } + return x < 0.0f ? -mag : mag; +} + +static float rocmfpx_fp3_block_mse_for_scale(const float * x, int n, uint8_t e, float best_err) { + const float scale = rocmfpx_ue4m3_to_fp32(e); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + if (!isfinite(x[i])) { + continue; + } + + const float y = rocmfpx_fp3_decoded_mag(x[i], inv_scale) * scale; + const float d = x[i] - y; + + err += d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +static float rocmfpx_fp3_block_weighted_mse_for_scale(const float * x, int n, const float * mse_weights, uint8_t e, float best_err) { + const float scale = rocmfpx_ue4m3_to_fp32(e); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + if (!isfinite(x[i])) { + continue; + } + + const float y = rocmfpx_fp3_decoded_mag(x[i], inv_scale) * scale; + const float d = x[i] - y; + + err += mse_weights[i]*d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +static uint8_t rocmfpx_choose_scale_fp3_mse_impl(const float * x, int n, const float * mse_weights, float max_abs, float max_abs_weight) { + const uint8_t start_e = rocmfpx_nearest_scale_ue4m3(max_abs / 4.0f); + uint8_t best_e = start_e; + float best_err = INFINITY; + bool lower_done = false; + + for (int delta = 0; delta <= 125; ++delta) { + const int e0 = (int) start_e - delta; + if (!lower_done && e0 >= 1 && e0 <= 126) { + const float scale = rocmfpx_ue4m3_to_fp32((uint8_t) e0); + const float clip_delta = max_abs - 4.0f*scale; + const float clip_err = mse_weights ? max_abs_weight*clip_delta*clip_delta : clip_delta*clip_delta; + if (clip_delta > 0.0f && clip_err > best_err) { + lower_done = true; + } else { + const float err = mse_weights ? + rocmfpx_fp3_block_weighted_mse_for_scale(x, n, mse_weights, (uint8_t) e0, best_err) : + rocmfpx_fp3_block_mse_for_scale(x, n, (uint8_t) e0, best_err); + if (err < best_err || (err == best_err && e0 < best_e)) { + best_err = err; + best_e = (uint8_t) e0; + } + } + } + + const int e1 = (int) start_e + delta; + if (delta != 0 && e1 >= 1 && e1 <= 126) { + const float err = mse_weights ? + rocmfpx_fp3_block_weighted_mse_for_scale(x, n, mse_weights, (uint8_t) e1, best_err) : + rocmfpx_fp3_block_mse_for_scale(x, n, (uint8_t) e1, best_err); + if (err < best_err || (err == best_err && e1 < best_e)) { + best_err = err; + best_e = (uint8_t) e1; + } + } + + if ((lower_done || e0 <= 1) && e1 >= 126) { + break; + } + } + + return best_e; +} + +static uint8_t rocmfpx_choose_scale_fp3_mse(const float * x, int n) { + const float max_abs = rocmfpx_max_abs(x, n); + if (!(max_abs > 0.0f) || !isfinite(max_abs)) { + return 0; + } + + return rocmfpx_choose_scale_fp3_mse_impl(x, n, NULL, max_abs, 0.0f); +} + +static uint8_t rocmfpx_choose_scale_fp3_weighted_mse(const float * x, int n, const float * quant_weights, float sigma2) { + assert(n <= QK_ROCMFP3); + float mse_weights[QK_ROCMFP3]; + float max_abs; + float max_abs_weight; + rocmfpx_prepare_mse_weights(mse_weights, x, n, quant_weights, sigma2, &max_abs, &max_abs_weight); + if (!(max_abs > 0.0f) || !isfinite(max_abs)) { + return 0; + } + + return rocmfpx_choose_scale_fp3_mse_impl(x, n, mse_weights, max_abs, max_abs_weight); +} + +static int rocmfpx_decode_fp6_code(uint8_t code) { + const int value = code & 31u; + return (code & 32u) ? -value : value; +} + +static uint8_t rocmfpx_quantize_fp6_code(float x, float inv_scale) { + if (!isfinite(x) || inv_scale <= 0.0f) { + return 0; + } + + int mag = (int) lroundf(fabsf(x * inv_scale)); + if (mag > 31) { + mag = 31; + } + + return mag == 0 ? 0 : (uint8_t) ((x < 0.0f ? 32u : 0u) | (uint8_t) mag); +} + +// Fused round + clamp + decode for the fp6 scale search. Returns the same signed +// decoded magnitude as rocmfpx_decode_fp6_code(rocmfpx_quantize_fp6_code(...)) +// (nearest integer in [0,31], signed), keeping quantized output bit-identical. +static inline float rocmfpx_fp6_decoded_mag(float x, float inv_scale) { + int mag = (int) lroundf(fabsf(x * inv_scale)); + if (mag > 31) { + mag = 31; + } + if (mag == 0) { + return 0.0f; + } + return x < 0.0f ? -(float) mag : (float) mag; +} + +static float rocmfpx_fp6_block_mse_for_scale(const float * x, int n, uint8_t e, float best_err) { + const float scale = rocmfpx_ue4m3_to_fp32(e); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + if (!isfinite(x[i])) { + continue; + } + const float y = rocmfpx_fp6_decoded_mag(x[i], inv_scale) * scale; + const float d = x[i] - y; + err += d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +static float rocmfpx_fp6_block_weighted_mse_for_scale(const float * x, int n, const float * mse_weights, uint8_t e, float best_err) { + const float scale = rocmfpx_ue4m3_to_fp32(e); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + if (!isfinite(x[i])) { + continue; + } + const float y = rocmfpx_fp6_decoded_mag(x[i], inv_scale) * scale; + const float d = x[i] - y; + err += mse_weights[i]*d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +static uint8_t rocmfpx_choose_scale_fp6_mse_impl(const float * x, int n, const float * mse_weights, float max_abs, float max_abs_weight) { + const uint8_t start_e = rocmfpx_nearest_scale_ue4m3(max_abs / 31.0f); + uint8_t best_e = start_e; + float best_err = INFINITY; + bool lower_done = false; + + for (int delta = 0; delta <= 125; ++delta) { + const int e0 = (int) start_e - delta; + if (!lower_done && e0 >= 1 && e0 <= 126) { + const float scale = rocmfpx_ue4m3_to_fp32((uint8_t) e0); + const float clip_delta = max_abs - 31.0f*scale; + const float clip_err = mse_weights ? max_abs_weight*clip_delta*clip_delta : clip_delta*clip_delta; + if (clip_delta > 0.0f && clip_err > best_err) { + lower_done = true; + } else { + const float err = mse_weights ? + rocmfpx_fp6_block_weighted_mse_for_scale(x, n, mse_weights, (uint8_t) e0, best_err) : + rocmfpx_fp6_block_mse_for_scale(x, n, (uint8_t) e0, best_err); + if (err < best_err || (err == best_err && e0 < best_e)) { + best_err = err; + best_e = (uint8_t) e0; + } + } + } + + const int e1 = (int) start_e + delta; + if (delta != 0 && e1 >= 1 && e1 <= 126) { + const float err = mse_weights ? + rocmfpx_fp6_block_weighted_mse_for_scale(x, n, mse_weights, (uint8_t) e1, best_err) : + rocmfpx_fp6_block_mse_for_scale(x, n, (uint8_t) e1, best_err); + if (err < best_err || (err == best_err && e1 < best_e)) { + best_err = err; + best_e = (uint8_t) e1; + } + } + + if ((lower_done || e0 <= 1) && e1 >= 126) { + break; + } + } + + return best_e; +} + +static uint8_t rocmfpx_choose_scale_fp6_mse(const float * x, int n) { + const float max_abs = rocmfpx_max_abs(x, n); + if (!(max_abs > 0.0f) || !isfinite(max_abs)) { + return 0; + } + + return rocmfpx_choose_scale_fp6_mse_impl(x, n, NULL, max_abs, 0.0f); +} + +static uint8_t rocmfpx_choose_scale_fp6_weighted_mse(const float * x, int n, const float * quant_weights, float sigma2) { + assert(n <= QK_ROCMFP6); + float mse_weights[QK_ROCMFP6]; + float max_abs; + float max_abs_weight; + rocmfpx_prepare_mse_weights(mse_weights, x, n, quant_weights, sigma2, &max_abs, &max_abs_weight); + if (!(max_abs > 0.0f) || !isfinite(max_abs)) { + return 0; + } + + return rocmfpx_choose_scale_fp6_mse_impl(x, n, mse_weights, max_abs, max_abs_weight); +} + +static int8_t rocmfpx_quantize_fp8_code(float x, float inv_scale) { + if (!isfinite(x) || inv_scale <= 0.0f) { + return 0; + } + + int q = (int) lroundf(x * inv_scale); + if (q > 127) { + q = 127; + } else if (q < -127) { + q = -127; + } + + return (int8_t) q; +} + +static float rocmfpx_fp8_block_weighted_mse_for_scale(const float * x, int n, const float * mse_weights, uint8_t e, float best_err) { + const float scale = rocmfpx_ue4m3_to_fp32(e); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + if (!isfinite(x[i])) { + continue; + } + + const int8_t code = rocmfpx_quantize_fp8_code(x[i], inv_scale); + const float y = (float) code * scale; + const float d = x[i] - y; + + err += mse_weights[i]*d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +static uint8_t rocmfpx_choose_scale_fp8_weighted_mse(const float * x, int n, const float * quant_weights, float sigma2) { + assert(n <= QK_ROCMFP8); + float mse_weights[QK_ROCMFP8]; + float max_abs; + float max_abs_weight; + rocmfpx_prepare_mse_weights(mse_weights, x, n, quant_weights, sigma2, &max_abs, &max_abs_weight); + if (!(max_abs > 0.0f) || !isfinite(max_abs)) { + return 0; + } + + const uint8_t start_e = rocmfpx_nearest_scale_ue4m3(max_abs / 127.0f); + uint8_t best_e = start_e; + float best_err = INFINITY; + bool lower_done = false; + + for (int delta = 0; delta <= 125; ++delta) { + const int e0 = (int) start_e - delta; + if (!lower_done && e0 >= 1 && e0 <= 126) { + const float scale = rocmfpx_ue4m3_to_fp32((uint8_t) e0); + const float clip_delta = max_abs - 127.0f*scale; + if (clip_delta > 0.0f && max_abs_weight*clip_delta*clip_delta > best_err) { + lower_done = true; + } else { + const float err = rocmfpx_fp8_block_weighted_mse_for_scale(x, n, mse_weights, (uint8_t) e0, best_err); + if (err < best_err || (err == best_err && e0 < best_e)) { + best_err = err; + best_e = (uint8_t) e0; + } + } + } + + const int e1 = (int) start_e + delta; + if (delta != 0 && e1 >= 1 && e1 <= 126) { + const float err = rocmfpx_fp8_block_weighted_mse_for_scale(x, n, mse_weights, (uint8_t) e1, best_err); + if (err < best_err || (err == best_err && e1 < best_e)) { + best_err = err; + best_e = (uint8_t) e1; + } + } + + if ((lower_done || e0 <= 1) && e1 >= 126) { + break; + } + } + + return best_e; +} + +void rocmfpx_quantize_row_fp2_ref(const float * GGML_RESTRICT x, block_rocmfp2 * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP2 == 0); + + const int64_t nb = k / QK_ROCMFP2; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib*QK_ROCMFP2; + block_rocmfp2 * yb = y + ib; + + memset(yb->qs, 0, sizeof(yb->qs)); + + for (int half = 0; half < 2; ++half) { + const float * xh = xb + half*(QK_ROCMFP2/2); + yb->e[half] = rocmfpx_choose_scale_fp2_mse(xh, QK_ROCMFP2/2); + + const float scale = rocmfpx_ue4m3_to_fp32(yb->e[half]); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + + for (int j = 0; j < QK_ROCMFP2/2; ++j) { + const int i = half*(QK_ROCMFP2/2) + j; + const uint8_t code = rocmfpx_quantize_fp2_code(xb[i], inv_scale); + yb->qs[i >> 2] |= (uint8_t) (code << (2*(i & 3))); + } + } + } +} + +static void rocmfpx_quantize_row_fp2_weighted( + const float * GGML_RESTRICT x, block_rocmfp2 * GGML_RESTRICT y, int64_t k, const float * GGML_RESTRICT quant_weights) { + assert(k % QK_ROCMFP2 == 0); + + float sum_x2 = 0.0f; + for (int64_t i = 0; i < k; ++i) { + sum_x2 += isfinite(x[i]) ? x[i]*x[i] : 0.0f; + } + const float sigma2 = sum_x2 / (float) k; + + const int64_t nb = k / QK_ROCMFP2; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib*QK_ROCMFP2; + const float * qw = quant_weights ? quant_weights + ib*QK_ROCMFP2 : NULL; + block_rocmfp2 * yb = y + ib; + + memset(yb->qs, 0, sizeof(yb->qs)); + + for (int half = 0; half < 2; ++half) { + const int half_off = half*(QK_ROCMFP2/2); + const float * xh = xb + half_off; + const float * qh = qw ? qw + half_off : NULL; + yb->e[half] = qh ? + rocmfpx_choose_scale_fp2_weighted_mse(xh, QK_ROCMFP2/2, qh, sigma2) : + rocmfpx_choose_scale_fp2_mse(xh, QK_ROCMFP2/2); + + const float scale = rocmfpx_ue4m3_to_fp32(yb->e[half]); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + + for (int j = 0; j < QK_ROCMFP2/2; ++j) { + const int i = half_off + j; + const uint8_t code = rocmfpx_quantize_fp2_code(xb[i], inv_scale); + yb->qs[i >> 2] |= (uint8_t) (code << (2*(i & 3))); + } + } + } +} + +void rocmfpx_dequantize_row_fp2(const block_rocmfp2 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP2 == 0); + + const int64_t nb = k / QK_ROCMFP2; + for (int64_t ib = 0; ib < nb; ++ib) { + const block_rocmfp2 * xb = x + ib; + float * yb = y + ib*QK_ROCMFP2; + + for (int i = 0; i < QK_ROCMFP2; ++i) { + const float scale = rocmfpx_ue4m3_to_fp32(xb->e[i >= QK_ROCMFP2/2]); + const uint8_t code = (uint8_t) ((xb->qs[i >> 2] >> (2*(i & 3))) & 3u); + yb[i] = kvalues_rocmfp2[code] * scale; + } + } +} + +void rocmfpx_quantize_row_fp2(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + rocmfpx_quantize_row_fp2_ref(x, (block_rocmfp2 *) y, k); +} + +size_t rocmfpx_quantize_fp2(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix) { + const size_t row_size = rocmfpx_row_size_fp2(n_per_row); + char * qrow = (char *) dst; + + for (int64_t row = 0; row < nrows; ++row) { + if (imatrix) { + rocmfpx_quantize_row_fp2_weighted(src + row*n_per_row, (block_rocmfp2 *) qrow, n_per_row, imatrix); + } else { + rocmfpx_quantize_row_fp2_ref(src + row*n_per_row, (block_rocmfp2 *) qrow, n_per_row); + } + qrow += row_size; + } + + return (size_t) nrows * row_size; +} + +void rocmfpx_quantize_row_fp3_ref(const float * GGML_RESTRICT x, block_rocmfp3 * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP3 == 0); + + const int64_t nb = k / QK_ROCMFP3; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib*QK_ROCMFP3; + block_rocmfp3 * yb = y + ib; + + memset(yb->qs, 0, sizeof(yb->qs)); + + for (int half = 0; half < 2; ++half) { + const float * xh = xb + half*(QK_ROCMFP3/2); + yb->e[half] = rocmfpx_choose_scale_fp3_mse(xh, QK_ROCMFP3/2); + + const float scale = rocmfpx_ue4m3_to_fp32(yb->e[half]); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + + for (int j = 0; j < QK_ROCMFP3/2; ++j) { + const int i = half*(QK_ROCMFP3/2) + j; + const uint8_t code = rocmfpx_quantize_fp3_code(xb[i], inv_scale); + rocmfpx_set_bits(yb->qs, i*3, 3, code); + } + } + } +} + +static void rocmfpx_quantize_row_fp3_weighted( + const float * GGML_RESTRICT x, block_rocmfp3 * GGML_RESTRICT y, int64_t k, const float * GGML_RESTRICT quant_weights) { + assert(k % QK_ROCMFP3 == 0); + + float sum_x2 = 0.0f; + for (int64_t i = 0; i < k; ++i) { + sum_x2 += isfinite(x[i]) ? x[i]*x[i] : 0.0f; + } + const float sigma2 = sum_x2 / (float) k; + + const int64_t nb = k / QK_ROCMFP3; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib*QK_ROCMFP3; + const float * qw = quant_weights ? quant_weights + ib*QK_ROCMFP3 : NULL; + block_rocmfp3 * yb = y + ib; + + memset(yb->qs, 0, sizeof(yb->qs)); + + for (int half = 0; half < 2; ++half) { + const int half_off = half*(QK_ROCMFP3/2); + const float * xh = xb + half_off; + const float * qh = qw ? qw + half_off : NULL; + yb->e[half] = qh ? + rocmfpx_choose_scale_fp3_weighted_mse(xh, QK_ROCMFP3/2, qh, sigma2) : + rocmfpx_choose_scale_fp3_mse(xh, QK_ROCMFP3/2); + + const float scale = rocmfpx_ue4m3_to_fp32(yb->e[half]); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + + for (int j = 0; j < QK_ROCMFP3/2; ++j) { + const int i = half_off + j; + const uint8_t code = rocmfpx_quantize_fp3_code(xb[i], inv_scale); + rocmfpx_set_bits(yb->qs, i*3, 3, code); + } + } + } +} + +void rocmfpx_dequantize_row_fp3(const block_rocmfp3 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP3 == 0); + + const int64_t nb = k / QK_ROCMFP3; + for (int64_t ib = 0; ib < nb; ++ib) { + const block_rocmfp3 * xb = x + ib; + float * yb = y + ib*QK_ROCMFP3; + + for (int i = 0; i < QK_ROCMFP3; ++i) { + const float scale = rocmfpx_ue4m3_to_fp32(xb->e[i >= QK_ROCMFP3/2]); + const uint8_t code = (uint8_t) rocmfpx_get_bits(xb->qs, i*3, 3); + yb[i] = (float) rocmfpx_decode_fp3_code(code) * scale; + } + } +} + +void rocmfpx_quantize_row_fp3(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + rocmfpx_quantize_row_fp3_ref(x, (block_rocmfp3 *) y, k); +} + +size_t rocmfpx_quantize_fp3(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix) { + const size_t row_size = rocmfpx_row_size_fp3(n_per_row); + char * qrow = (char *) dst; + + for (int64_t row = 0; row < nrows; ++row) { + if (imatrix) { + rocmfpx_quantize_row_fp3_weighted(src + row*n_per_row, (block_rocmfp3 *) qrow, n_per_row, imatrix); + } else { + rocmfpx_quantize_row_fp3_ref(src + row*n_per_row, (block_rocmfp3 *) qrow, n_per_row); + } + qrow += row_size; + } + + return (size_t) nrows * row_size; +} + +static void rocmfpx_quantize_row_fp6_weighted( + const float * GGML_RESTRICT x, block_rocmfp6 * GGML_RESTRICT y, int64_t k, const float * GGML_RESTRICT quant_weights) { + assert(k % QK_ROCMFP6 == 0); + + float sum_x2 = 0.0f; + for (int64_t i = 0; i < k; ++i) { + sum_x2 += isfinite(x[i]) ? x[i]*x[i] : 0.0f; + } + const float sigma2 = sum_x2 / (float) k; + + const int64_t nb = k / QK_ROCMFP6; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib*QK_ROCMFP6; + const float * qw = quant_weights ? quant_weights + ib*QK_ROCMFP6 : NULL; + block_rocmfp6 * yb = y + ib; + + memset(yb->qs, 0, sizeof(yb->qs)); + + for (int half = 0; half < 2; ++half) { + const int half_off = half*(QK_ROCMFP6/2); + const float * xh = xb + half_off; + const float * qh = qw ? qw + half_off : NULL; + yb->e[half] = qh ? + rocmfpx_choose_scale_fp6_weighted_mse(xh, QK_ROCMFP6/2, qh, sigma2) : + rocmfpx_choose_scale_fp6_mse(xh, QK_ROCMFP6/2); + + const float scale = rocmfpx_ue4m3_to_fp32(yb->e[half]); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + + for (int j = 0; j < QK_ROCMFP6/2; ++j) { + const int i = half_off + j; + const uint8_t code = rocmfpx_quantize_fp6_code(xb[i], inv_scale); + rocmfpx_set_bits(yb->qs, i*6, 6, code); + } + } + } +} + +void rocmfpx_quantize_row_fp6_ref(const float * GGML_RESTRICT x, block_rocmfp6 * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP6 == 0); + + const int64_t nb = k / QK_ROCMFP6; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib*QK_ROCMFP6; + block_rocmfp6 * yb = y + ib; + + memset(yb->qs, 0, sizeof(yb->qs)); + + for (int half = 0; half < 2; ++half) { + const float * xh = xb + half*(QK_ROCMFP6/2); + yb->e[half] = rocmfpx_choose_scale_fp6_mse(xh, QK_ROCMFP6/2); + + const float scale = rocmfpx_ue4m3_to_fp32(yb->e[half]); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + + for (int j = 0; j < QK_ROCMFP6/2; ++j) { + const int i = half*(QK_ROCMFP6/2) + j; + const uint8_t code = rocmfpx_quantize_fp6_code(xb[i], inv_scale); + rocmfpx_set_bits(yb->qs, i*6, 6, code); + } + } + } +} + +void rocmfpx_dequantize_row_fp6(const block_rocmfp6 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP6 == 0); + + const int64_t nb = k / QK_ROCMFP6; + for (int64_t ib = 0; ib < nb; ++ib) { + const block_rocmfp6 * xb = x + ib; + float * yb = y + ib*QK_ROCMFP6; + + for (int i = 0; i < QK_ROCMFP6; ++i) { + const float scale = rocmfpx_ue4m3_to_fp32(xb->e[i >= QK_ROCMFP6/2]); + const uint8_t code = (uint8_t) rocmfpx_get_bits(xb->qs, i*6, 6); + yb[i] = (float) rocmfpx_decode_fp6_code(code) * scale; + } + } +} + +void rocmfpx_quantize_row_fp6(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + rocmfpx_quantize_row_fp6_ref(x, (block_rocmfp6 *) y, k); +} + +size_t rocmfpx_quantize_fp6(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix) { + const size_t row_size = rocmfpx_row_size_fp6(n_per_row); + char * qrow = (char *) dst; + + for (int64_t row = 0; row < nrows; ++row) { + if (imatrix) { + rocmfpx_quantize_row_fp6_weighted(src + row*n_per_row, (block_rocmfp6 *) qrow, n_per_row, imatrix); + } else { + rocmfpx_quantize_row_fp6_ref(src + row*n_per_row, (block_rocmfp6 *) qrow, n_per_row); + } + qrow += row_size; + } + + return (size_t) nrows * row_size; +} + +static void rocmfpx_quantize_row_fp8_weighted( + const float * GGML_RESTRICT x, block_rocmfp8 * GGML_RESTRICT y, int64_t k, const float * GGML_RESTRICT quant_weights) { + assert(k % QK_ROCMFP8 == 0); + + float sum_x2 = 0.0f; + for (int64_t i = 0; i < k; ++i) { + sum_x2 += isfinite(x[i]) ? x[i]*x[i] : 0.0f; + } + const float sigma2 = sum_x2 / (float) k; + + const int64_t nb = k / QK_ROCMFP8; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib*QK_ROCMFP8; + const float * qw = quant_weights ? quant_weights + ib*QK_ROCMFP8 : NULL; + block_rocmfp8 * yb = y + ib; + + yb->e = qw ? rocmfpx_choose_scale_fp8_weighted_mse(xb, QK_ROCMFP8, qw, sigma2) : + rocmfpx_nearest_scale_ue4m3(rocmfpx_max_abs(xb, QK_ROCMFP8) / 127.0f); + + const float scale = rocmfpx_ue4m3_to_fp32(yb->e); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + + for (int i = 0; i < QK_ROCMFP8; ++i) { + yb->qs[i] = rocmfpx_quantize_fp8_code(xb[i], inv_scale); + } + } +} + +void rocmfpx_quantize_row_fp8_ref(const float * GGML_RESTRICT x, block_rocmfp8 * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP8 == 0); + + const int64_t nb = k / QK_ROCMFP8; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib*QK_ROCMFP8; + block_rocmfp8 * yb = y + ib; + + const float max_abs = rocmfpx_max_abs(xb, QK_ROCMFP8); + yb->e = rocmfpx_nearest_scale_ue4m3(max_abs / 127.0f); + + const float scale = rocmfpx_ue4m3_to_fp32(yb->e); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + + for (int i = 0; i < QK_ROCMFP8; ++i) { + yb->qs[i] = rocmfpx_quantize_fp8_code(xb[i], inv_scale); + } + } +} + +void rocmfpx_dequantize_row_fp8(const block_rocmfp8 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP8 == 0); + + const int64_t nb = k / QK_ROCMFP8; + for (int64_t ib = 0; ib < nb; ++ib) { + const block_rocmfp8 * xb = x + ib; + float * yb = y + ib*QK_ROCMFP8; + + const float scale = rocmfpx_ue4m3_to_fp32(xb->e); + for (int i = 0; i < QK_ROCMFP8; ++i) { + yb[i] = (float) xb->qs[i] * scale; + } + } +} + +void rocmfpx_quantize_row_fp8(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + rocmfpx_quantize_row_fp8_ref(x, (block_rocmfp8 *) y, k); +} + +size_t rocmfpx_quantize_fp8(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix) { + const size_t row_size = rocmfpx_row_size_fp8(n_per_row); + char * qrow = (char *) dst; + + for (int64_t row = 0; row < nrows; ++row) { + if (imatrix) { + rocmfpx_quantize_row_fp8_weighted(src + row*n_per_row, (block_rocmfp8 *) qrow, n_per_row, imatrix); + } else { + rocmfpx_quantize_row_fp8_ref(src + row*n_per_row, (block_rocmfp8 *) qrow, n_per_row); + } + qrow += row_size; + } + + return (size_t) nrows * row_size; +} + +bool rocmfpx_validate_row_data_fp2(const void * data, size_t nbytes) { + if (nbytes % sizeof(block_rocmfp2) != 0) { + return false; + } + + const block_rocmfp2 * blocks = (const block_rocmfp2 *) data; + const size_t nb = nbytes / sizeof(block_rocmfp2); + + for (size_t i = 0; i < nb; ++i) { + if (!rocmfpx_scale_is_valid(blocks[i].e[0]) || !rocmfpx_scale_is_valid(blocks[i].e[1])) { + return false; + } + } + + return true; +} + +bool rocmfpx_validate_row_data_fp3(const void * data, size_t nbytes) { + if (nbytes % sizeof(block_rocmfp3) != 0) { + return false; + } + + const block_rocmfp3 * blocks = (const block_rocmfp3 *) data; + const size_t nb = nbytes / sizeof(block_rocmfp3); + + for (size_t i = 0; i < nb; ++i) { + if (!rocmfpx_scale_is_valid(blocks[i].e[0]) || !rocmfpx_scale_is_valid(blocks[i].e[1])) { + return false; + } + } + + return true; +} + +bool rocmfpx_validate_row_data_fp6(const void * data, size_t nbytes) { + if (nbytes % sizeof(block_rocmfp6) != 0) { + return false; + } + + const block_rocmfp6 * blocks = (const block_rocmfp6 *) data; + const size_t nb = nbytes / sizeof(block_rocmfp6); + + for (size_t i = 0; i < nb; ++i) { + if (!rocmfpx_scale_is_valid(blocks[i].e[0]) || !rocmfpx_scale_is_valid(blocks[i].e[1])) { + return false; + } + } + + return true; +} + +bool rocmfpx_validate_row_data_fp8(const void * data, size_t nbytes) { + if (nbytes % sizeof(block_rocmfp8) != 0) { + return false; + } + + const block_rocmfp8 * blocks = (const block_rocmfp8 *) data; + const size_t nb = nbytes / sizeof(block_rocmfp8); + + for (size_t i = 0; i < nb; ++i) { + if (!rocmfpx_scale_is_valid(blocks[i].e)) { + return false; + } + } + + return true; +} diff --git a/ggml/rocmfpx/rocmfpx.h b/ggml/rocmfpx/rocmfpx.h new file mode 100644 index 000000000000..f9d6cb896ab6 --- /dev/null +++ b/ggml/rocmfpx/rocmfpx.h @@ -0,0 +1,111 @@ +#pragma once + +#include +#include +#include + +#include "ggml.h" + +#ifdef __cplusplus +extern "C" { +#endif + +#define QK_ROCMFPX 32 + +#define QK_ROCMFP2 QK_ROCMFPX +#define QK_ROCMFP3 QK_ROCMFPX +#define QK_ROCMFP6 QK_ROCMFPX +#define QK_ROCMFP8 QK_ROCMFPX + +#define QS_ROCMFP2 ((QK_ROCMFP2 * 2) / 8) +#define QS_ROCMFP3 ((QK_ROCMFP3 * 3) / 8) +#define QS_ROCMFP6 ((QK_ROCMFP6 * 6) / 8) +#define QS_ROCMFP8 QK_ROCMFP8 + +#define QR_ROCMFP2 1 +#define QI_ROCMFP2 (QS_ROCMFP2 / 4) // 8B qs = 2 int32; decoupled from QR_ROCMFP2 (QR stays 1 for convert.cu dequant) + +#define ROCMFP2_KVALUE_0_I8 (-1) +#define ROCMFP2_KVALUE_1_I8 0 +#define ROCMFP2_KVALUE_2_I8 1 +#define ROCMFP2_KVALUE_3_I8 2 +#define ROCMFP2_KVALUES_INIT { (float) ROCMFP2_KVALUE_0_I8, (float) ROCMFP2_KVALUE_1_I8, (float) ROCMFP2_KVALUE_2_I8, (float) ROCMFP2_KVALUE_3_I8 } + +#define QR_ROCMFP3 1 +#define QI_ROCMFP3 (QK_ROCMFP3 / (4 * QR_ROCMFP3)) + +#define QR_ROCMFP6 1 +#define QI_ROCMFP6 (QK_ROCMFP6 / (4 * QR_ROCMFP6)) + +#define QR_ROCMFP8 1 +#define QI_ROCMFP8 (QK_ROCMFP8 / (4 * QR_ROCMFP8)) + +// AMD-native experimental family layouts. The GGUF types are registered, but +// the layouts stay isolated from the promoted ROCmFP4 formats while evaluated. +typedef struct { + uint8_t qs[QS_ROCMFP2]; + uint8_t e[2]; +} block_rocmfp2; + +typedef struct { + uint8_t qs[QS_ROCMFP3]; + uint8_t e[2]; +} block_rocmfp3; + +typedef struct { + uint8_t qs[QS_ROCMFP6]; + uint8_t e[2]; +} block_rocmfp6; + +typedef struct { + int8_t qs[QS_ROCMFP8]; + uint8_t e; +} block_rocmfp8; + +#if defined(__cplusplus) +static_assert(sizeof(block_rocmfp2) == QS_ROCMFP2 + 2*sizeof(uint8_t), "wrong rocmfp2 block size/padding"); +static_assert(sizeof(block_rocmfp3) == QS_ROCMFP3 + 2*sizeof(uint8_t), "wrong rocmfp3 block size/padding"); +static_assert(sizeof(block_rocmfp6) == QS_ROCMFP6 + 2*sizeof(uint8_t), "wrong rocmfp6 block size/padding"); +static_assert(sizeof(block_rocmfp8) == QS_ROCMFP8 + sizeof(uint8_t), "wrong rocmfp8 block size/padding"); +#else +_Static_assert(sizeof(block_rocmfp2) == QS_ROCMFP2 + 2*sizeof(uint8_t), "wrong rocmfp2 block size/padding"); +_Static_assert(sizeof(block_rocmfp3) == QS_ROCMFP3 + 2*sizeof(uint8_t), "wrong rocmfp3 block size/padding"); +_Static_assert(sizeof(block_rocmfp6) == QS_ROCMFP6 + 2*sizeof(uint8_t), "wrong rocmfp6 block size/padding"); +_Static_assert(sizeof(block_rocmfp8) == QS_ROCMFP8 + sizeof(uint8_t), "wrong rocmfp8 block size/padding"); +#endif + +GGML_API float rocmfpx_ue4m3_to_fp32(uint8_t e); +GGML_API bool rocmfpx_scale_is_valid(uint8_t e); +GGML_API size_t rocmfpx_row_size_fp2(int64_t k); +GGML_API size_t rocmfpx_row_size_fp3(int64_t k); +GGML_API size_t rocmfpx_row_size_fp6(int64_t k); +GGML_API size_t rocmfpx_row_size_fp8(int64_t k); + +GGML_API void rocmfpx_quantize_row_fp2_ref(const float * GGML_RESTRICT x, block_rocmfp2 * GGML_RESTRICT y, int64_t k); +GGML_API void rocmfpx_dequantize_row_fp2(const block_rocmfp2 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void rocmfpx_quantize_row_fp2(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +GGML_API size_t rocmfpx_quantize_fp2(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); + +GGML_API void rocmfpx_quantize_row_fp3_ref(const float * GGML_RESTRICT x, block_rocmfp3 * GGML_RESTRICT y, int64_t k); +GGML_API void rocmfpx_dequantize_row_fp3(const block_rocmfp3 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void rocmfpx_quantize_row_fp3(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +GGML_API size_t rocmfpx_quantize_fp3(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); + +GGML_API void rocmfpx_quantize_row_fp6_ref(const float * GGML_RESTRICT x, block_rocmfp6 * GGML_RESTRICT y, int64_t k); +GGML_API void rocmfpx_dequantize_row_fp6(const block_rocmfp6 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void rocmfpx_quantize_row_fp6(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +GGML_API size_t rocmfpx_quantize_fp6(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); + +GGML_API void rocmfpx_quantize_row_fp8_ref(const float * GGML_RESTRICT x, block_rocmfp8 * GGML_RESTRICT y, int64_t k); +GGML_API void rocmfpx_dequantize_row_fp8(const block_rocmfp8 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void rocmfpx_quantize_row_fp8(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +GGML_API size_t rocmfpx_quantize_fp8(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); + +GGML_API bool rocmfpx_validate_row_data_fp2(const void * data, size_t nbytes); +GGML_API bool rocmfpx_validate_row_data_fp3(const void * data, size_t nbytes); +GGML_API bool rocmfpx_validate_row_data_fp6(const void * data, size_t nbytes); +GGML_API bool rocmfpx_validate_row_data_fp8(const void * data, size_t nbytes); + +#ifdef __cplusplus +} +#endif diff --git a/ggml/rocmfpx/test_rocmfpx.c b/ggml/rocmfpx/test_rocmfpx.c new file mode 100644 index 000000000000..5bae72bd7f65 --- /dev/null +++ b/ggml/rocmfpx/test_rocmfpx.c @@ -0,0 +1,164 @@ +#include "rocmfpx.h" + +#include +#include +#include + +static void fill_row(float * x, int n) { + for (int i = 0; i < n; ++i) { + const float wave = 0.75f*sinf((float) i * 0.37f) + 0.25f*cosf((float) i * 0.13f); + const float ramp = ((float) (i % 11) - 5.0f) * 0.035f; + x[i] = wave + ramp; + } + + x[7] = 3.25f; + x[19] = -2.75f; + x[43] = 1.875f; +} + +static float mse(const float * a, const float * b, int n) { + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + const float d = a[i] - b[i]; + err += d*d; + } + + return err / (float) n; +} + +static float weighted_mse(const float * a, const float * b, const float * w, int n) { + float err = 0.0f; + float sum_w = 0.0f; + + for (int i = 0; i < n; ++i) { + const float d = a[i] - b[i]; + err += w[i]*d*d; + sum_w += w[i]; + } + + return sum_w > 0.0f ? err / sum_w : 0.0f; +} + +static void check_weighted_imatrix_fp3(void) { + enum { N = QK_ROCMFP3 }; + + float src[N]; + float imatrix[N]; + float plain[N]; + float weighted[N]; + block_rocmfp3 q_plain[N / QK_ROCMFP3]; + block_rocmfp3 q_weighted[N / QK_ROCMFP3]; + + for (int i = 0; i < N; ++i) { + src[i] = (i % 2) ? 0.21f : -0.21f; + imatrix[i] = 100.0f; + } + + src[0] = 9.0f; + imatrix[0] = 0.0f; + + rocmfpx_quantize_fp3(src, q_plain, 1, N, NULL); + rocmfpx_quantize_fp3(src, q_weighted, 1, N, imatrix); + rocmfpx_dequantize_row_fp3(q_plain, plain, N); + rocmfpx_dequantize_row_fp3(q_weighted, weighted, N); + + const float plain_err = weighted_mse(src, plain, imatrix, N); + const float weighted_err = weighted_mse(src, weighted, imatrix, N); + + printf("ROCmFP3 imatrix weighted_mse: plain=%g weighted=%g\n", plain_err, weighted_err); + assert(weighted_err < plain_err); +} + +static void check_weighted_imatrix_fp2(void) { + enum { N = QK_ROCMFP2 }; + + float src[N]; + float imatrix[N]; + float plain[N]; + float weighted[N]; + block_rocmfp2 q_plain[N / QK_ROCMFP2]; + block_rocmfp2 q_weighted[N / QK_ROCMFP2]; + + for (int i = 0; i < N; ++i) { + src[i] = (i % 2) ? 0.21f : -0.21f; + imatrix[i] = 100.0f; + } + + src[0] = 9.0f; + imatrix[0] = 0.0f; + + rocmfpx_quantize_fp2(src, q_plain, 1, N, NULL); + rocmfpx_quantize_fp2(src, q_weighted, 1, N, imatrix); + rocmfpx_dequantize_row_fp2(q_plain, plain, N); + rocmfpx_dequantize_row_fp2(q_weighted, weighted, N); + + const float plain_err = weighted_mse(src, plain, imatrix, N); + const float weighted_err = weighted_mse(src, weighted, imatrix, N); + + printf("ROCmFP2 imatrix weighted_mse: plain=%g weighted=%g\n", plain_err, weighted_err); + assert(weighted_err < plain_err); +} + +int main(void) { + enum { N = 64 }; + + float src[N]; + float fp2[N]; + float fp3[N]; + float fp6[N]; + float fp8[N]; + + block_rocmfp2 q2[N / QK_ROCMFP2]; + block_rocmfp3 q3[N / QK_ROCMFP3]; + block_rocmfp6 q6[N / QK_ROCMFP6]; + block_rocmfp8 q8[N / QK_ROCMFP8]; + + fill_row(src, N); + + rocmfpx_quantize_row_fp2_ref(src, q2, N); + rocmfpx_quantize_row_fp3_ref(src, q3, N); + rocmfpx_quantize_row_fp6_ref(src, q6, N); + rocmfpx_quantize_row_fp8_ref(src, q8, N); + + assert(rocmfpx_validate_row_data_fp2(q2, sizeof(q2))); + assert(rocmfpx_validate_row_data_fp3(q3, sizeof(q3))); + assert(rocmfpx_validate_row_data_fp6(q6, sizeof(q6))); + assert(rocmfpx_validate_row_data_fp8(q8, sizeof(q8))); + + rocmfpx_dequantize_row_fp2(q2, fp2, N); + rocmfpx_dequantize_row_fp3(q3, fp3, N); + rocmfpx_dequantize_row_fp6(q6, fp6, N); + rocmfpx_dequantize_row_fp8(q8, fp8, N); + + const float mse2 = mse(src, fp2, N); + const float mse3 = mse(src, fp3, N); + const float mse6 = mse(src, fp6, N); + const float mse8 = mse(src, fp8, N); + + printf("ROCmFP2: block=%zu row=%zu bpw=%.2f mse=%g\n", + sizeof(block_rocmfp2), rocmfpx_row_size_fp2(N), + 8.0f*(float) sizeof(block_rocmfp2)/(float) QK_ROCMFP2, mse2); + printf("ROCmFP3: block=%zu row=%zu bpw=%.2f mse=%g\n", + sizeof(block_rocmfp3), rocmfpx_row_size_fp3(N), + 8.0f*(float) sizeof(block_rocmfp3)/(float) QK_ROCMFP3, mse3); + printf("ROCmFP6: block=%zu row=%zu bpw=%.2f mse=%g\n", + sizeof(block_rocmfp6), rocmfpx_row_size_fp6(N), + 8.0f*(float) sizeof(block_rocmfp6)/(float) QK_ROCMFP6, mse6); + printf("ROCmFP8: block=%zu row=%zu bpw=%.2f mse=%g\n", + sizeof(block_rocmfp8), rocmfpx_row_size_fp8(N), + 8.0f*(float) sizeof(block_rocmfp8)/(float) QK_ROCMFP8, mse8); + + assert(sizeof(block_rocmfp2) == 10); + assert(isfinite(mse2)); + assert(isfinite(mse3)); + assert(isfinite(mse6)); + assert(isfinite(mse8)); + assert(mse8 < mse6); + assert(mse6 < mse3); + + check_weighted_imatrix_fp2(); + check_weighted_imatrix_fp3(); + + return 0; +} diff --git a/ggml/src/CMakeLists.txt b/ggml/src/CMakeLists.txt index 48fbe208d900..315eafe2edbb 100644 --- a/ggml/src/CMakeLists.txt +++ b/ggml/src/CMakeLists.txt @@ -205,6 +205,10 @@ add_library(ggml-base ggml-threading.cpp ggml-threading.h ggml-quants.c + ../rocmfp4/rocmfp4.c + ../rocmfp4/rocmfp4.h + ../rocmfpx/rocmfpx.c + ../rocmfpx/rocmfpx.h ggml-quants.h gguf.cpp) diff --git a/ggml/src/ggml-common.h b/ggml/src/ggml-common.h index d47c976968f1..e52a2005dda9 100644 --- a/ggml/src/ggml-common.h +++ b/ggml/src/ggml-common.h @@ -1132,6 +1132,12 @@ GGML_TABLE_BEGIN(int8_t, kvalues_mxfp4, 16) 0, 1, 2, 3, 4, 6, 8, 12, 0, -1, -2, -3, -4, -6, -8, -12, GGML_TABLE_END() +// ROCmFP4 uses an E2M1-derived value set with the largest level retuned from +// 12 to 10, plus dual half-block UE4M3 scales. +GGML_TABLE_BEGIN(int8_t, kvalues_rocmfp4, 16) + 0, 1, 2, 3, 4, 6, 8, 10, 0, -1, -2, -3, -4, -6, -8, -10, +GGML_TABLE_END() + #define NGRID_IQ1S 2048 #define IQ1S_DELTA 0.125f #define IQ1M_DELTA 0.125f diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index c8f811406ccc..6dff193386ea 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -3271,6 +3271,80 @@ static void ggml_compute_forward_swiglu( } } +// ggml_compute_forward_swiglu_ds4 + +static void ggml_compute_forward_swiglu_ds4_f32( + const ggml_compute_params * params, + ggml_tensor * dst) { + + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + char * src0_d = (char *) src0->data; + char * src1_d = (char *) (src1 ? src1->data : src0->data); + const size_t src0_o = src0->nb[1]; + const size_t src1_o = src1 ? src1->nb[1] : src0->nb[1]; + + GGML_ASSERT(ggml_is_contiguous_1(src0)); + GGML_ASSERT(ggml_is_contiguous_1(dst)); + GGML_ASSERT(src1); + GGML_ASSERT(ggml_is_contiguous_1(src1)); + GGML_ASSERT(src0->type == src1->type); + + const int ith = params->ith; + const int nth = params->nth; + + const int nc = src0->ne[0]; + const int nr = ggml_nrows(src0); + + GGML_ASSERT(dst->ne[0] == nc); + GGML_ASSERT(ggml_nrows(dst) == nr); + + const float clamp = ggml_get_op_params_f32(dst, 2); + + const int dr = (nr + nth - 1)/nth; + const int ir0 = dr*ith; + const int ir1 = MIN(ir0 + dr, nr); + + for (int i1 = ir0; i1 < ir1; i1++) { + float * gate_p = (float *) (src0_d + i1*src0_o); + float * up_p = (float *) (src1_d + i1*src1_o); + float * dst_p = (float *) ((char *) dst->data + i1*(dst->nb[1])); + + for (int k = 0; k < nc; k++) { + const float gate = std::min(gate_p[k], clamp); + const float up = std::clamp(up_p[k], -clamp, clamp); + dst_p[k] = up * gate / (1.f + expf(-gate)); + } + +#ifndef NDEBUG + for (int k = 0; k < nc; k++) { + const float x = dst_p[k]; + GGML_UNUSED(x); + assert(!isnan(x)); + assert(!isinf(x)); + } +#endif // NDEBUG + } +} + +static void ggml_compute_forward_swiglu_ds4( + const ggml_compute_params * params, + ggml_tensor * dst) { + + const ggml_tensor * src0 = dst->src[0]; + + switch (src0->type) { + case GGML_TYPE_F32: + { + ggml_compute_forward_swiglu_ds4_f32(params, dst); + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} + // ggml_compute_forward_swiglu_oai static void ggml_compute_forward_swiglu_oai_f32( @@ -9801,6 +9875,10 @@ void ggml_compute_forward_glu( { ggml_compute_forward_swiglu_oai(params, dst); } break; + case GGML_GLU_OP_SWIGLU_DS4: + { + ggml_compute_forward_swiglu_ds4(params, dst); + } break; case GGML_GLU_OP_GEGLU_ERF: { ggml_compute_forward_geglu_erf(params, dst); diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index 9df1461a5749..eb7d125a1e46 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -19,6 +19,8 @@ #endif #endif #include "ggml-common.h" +#include "../../rocmfp4/rocmfp4.h" +#include "../../rocmfpx/rocmfpx.h" #include #include @@ -29,6 +31,26 @@ #include #include +#ifndef GGML_ROCMFP6_EXPANDED_DEVICE +#define GGML_ROCMFP6_EXPANDED_DEVICE 0 +#endif + +// Optional device-only ROCmFP6 layout. GGUF/CPU storage remains the packed +// block_rocmfp6 layout; experimental ROCm builds may expand qs to signed +// bytes to avoid bit unpacking in hot matmul/FA kernels. +struct block_rocmfp6_expanded { + int8_t qs[QK_ROCMFP6]; + uint8_t e[2]; +}; + +static_assert(sizeof(block_rocmfp6_expanded) == QK_ROCMFP6 + 2*sizeof(uint8_t), "wrong expanded rocmfp6 block size/padding"); + +#if GGML_ROCMFP6_EXPANDED_DEVICE +using block_rocmfp6_device = block_rocmfp6_expanded; +#else +using block_rocmfp6_device = block_rocmfp6; +#endif + #if defined(GGML_USE_HIP) #include "vendors/hip.h" #elif defined(GGML_USE_MUSA) @@ -959,6 +981,48 @@ struct ggml_cuda_type_traits { static constexpr int qi = QI8_0; }; +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK_ROCMFP4; + static constexpr int qr = QR_ROCMFP4; + static constexpr int qi = QI_ROCMFP4; +}; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK_ROCMFP4; + static constexpr int qr = QR_ROCMFP4; + static constexpr int qi = QI_ROCMFP4; +}; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK_ROCMFP2; + static constexpr int qr = QR_ROCMFP2; + static constexpr int qi = QI_ROCMFP2; +}; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK_ROCMFP3; + static constexpr int qr = QR_ROCMFP3; + static constexpr int qi = QI_ROCMFP3; +}; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK_ROCMFP6; + static constexpr int qr = QR_ROCMFP6; + static constexpr int qi = QI_ROCMFP6; +}; + +template<> +struct ggml_cuda_type_traits { + static constexpr int qk = QK_ROCMFP8; + static constexpr int qr = QR_ROCMFP8; + static constexpr int qi = QI_ROCMFP8; +}; + template<> struct ggml_cuda_type_traits { static constexpr int qk = QK_MXFP4; @@ -1466,10 +1530,14 @@ struct ggml_cuda_mm_fusion_args_host { const ggml_tensor * gate = nullptr; const ggml_tensor * gate_bias = nullptr; ggml_glu_op glu_op; + float glu_param0 = 0.0f; + float glu_param1 = 0.0f; }; struct ggml_cuda_mm_fusion_args_device { const void * x_bias = nullptr; const void * gate = nullptr; const void * gate_bias = nullptr; ggml_glu_op glu_op; + float glu_param0 = 0.0f; + float glu_param1 = 0.0f; }; diff --git a/ggml/src/ggml-cuda/convert.cu b/ggml/src/ggml-cuda/convert.cu index c675034af7c0..28d042aa01ab 100644 --- a/ggml/src/ggml-cuda/convert.cu +++ b/ggml/src/ggml-cuda/convert.cu @@ -1,6 +1,7 @@ #include "convert.cuh" #include "dequantize.cuh" #include "tq3-quant.cuh" +#include "../../rocmfp4/rocmfp4_hip_scale.cuh" #include @@ -487,6 +488,43 @@ static __global__ void dequantize_block_mxfp4(const void * __restrict__ vx, dst_ } } +template +static __global__ void dequantize_block_rocmfp4(const void * __restrict__ vx, dst_t * __restrict__ yy) { + + const int64_t i = blockIdx.x; + const block_rocmfp4 * x = (const block_rocmfp4 *) vx + i*(QK_K/QK_ROCMFP4); + + const int64_t tid = threadIdx.x; + const int64_t il = tid/8; // 0...3 + const int64_t ib = tid%8; // 0...7 + dst_t * y = yy + i*QK_K + 32*ib + 4*il; + const uint8_t * q4 = x[ib].qs + 4*il; + const float d0 = rocmfp4_ue4m3_to_fp32_half_finite(x[ib].e[0]); + const float d1 = rocmfp4_ue4m3_to_fp32_half_finite(x[ib].e[1]); + for (int j = 0; j < 4; ++j) { + y[j+ 0] = d0 * rocmfp4_decode_i8(q4[j]); + y[j+16] = d1 * rocmfp4_decode_i8(q4[j] >> 4); + } +} + +template +static __global__ void dequantize_block_rocmfp4_fast(const void * __restrict__ vx, dst_t * __restrict__ yy) { + + const int64_t i = blockIdx.x; + const block_rocmfp4_fast * x = (const block_rocmfp4_fast *) vx + i*(QK_K/QK_ROCMFP4); + + const int64_t tid = threadIdx.x; + const int64_t il = tid/8; // 0...3 + const int64_t ib = tid%8; // 0...7 + dst_t * y = yy + i*QK_K + 32*ib + 4*il; + const uint8_t * q4 = x[ib].qs + 4*il; + const float d = rocmfp4_ue4m3_to_fp32_half_finite(x[ib].e); + for (int j = 0; j < 4; ++j) { + y[j+ 0] = d * rocmfp4_decode_i8(q4[j]); + y[j+16] = d * rocmfp4_decode_i8(q4[j] >> 4); + } +} + template static void dequantize_block_cuda(const void * vx, dst_t * y, const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t ne03, @@ -645,6 +683,18 @@ static void dequantize_row_mxfp4_cuda(const void * vx, dst_t * y, const int64_t dequantize_block_mxfp4<<>>(vx, y); } +template +static void dequantize_row_rocmfp4_hip(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) { + const int nb = (k + QK_K - 1) / QK_K; + dequantize_block_rocmfp4<<>>(vx, y); +} + +template +static void dequantize_row_rocmfp4_fast_hip(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) { + const int nb = (k + QK_K - 1) / QK_K; + dequantize_block_rocmfp4_fast<<>>(vx, y); +} + template static __global__ void dequantize_block_nvfp4( const void * __restrict__ vx, @@ -786,6 +836,18 @@ to_fp16_cuda_t ggml_get_to_fp16_cuda(ggml_type type) { return dequantize_row_mxfp4_cuda; case GGML_TYPE_NVFP4: return dequantize_row_nvfp4_cuda; + case GGML_TYPE_Q4_0_ROCMFP4: + return dequantize_row_rocmfp4_hip; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + return dequantize_row_rocmfp4_fast_hip; + case GGML_TYPE_Q2_0_ROCMFP2: + return dequantize_block_cont_cuda; + case GGML_TYPE_Q3_0_ROCMFPX: + return dequantize_block_cont_cuda; + case GGML_TYPE_Q6_0_ROCMFPX: + return dequantize_block_cont_cuda; + case GGML_TYPE_Q8_0_ROCMFPX: + return dequantize_block_cont_cuda; case GGML_TYPE_F32: return convert_unary_cont_cuda; case GGML_TYPE_BF16: @@ -841,6 +903,18 @@ to_fp32_cuda_t ggml_get_to_fp32_cuda(ggml_type type) { return dequantize_row_mxfp4_cuda; case GGML_TYPE_NVFP4: return dequantize_row_nvfp4_cuda; + case GGML_TYPE_Q4_0_ROCMFP4: + return dequantize_row_rocmfp4_hip; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + return dequantize_row_rocmfp4_fast_hip; + case GGML_TYPE_Q2_0_ROCMFP2: + return dequantize_block_cont_cuda; + case GGML_TYPE_Q3_0_ROCMFPX: + return dequantize_block_cont_cuda; + case GGML_TYPE_Q6_0_ROCMFPX: + return dequantize_block_cont_cuda; + case GGML_TYPE_Q8_0_ROCMFPX: + return dequantize_block_cont_cuda; case GGML_TYPE_F16: return convert_unary_cont_cuda; case GGML_TYPE_BF16: @@ -864,6 +938,18 @@ to_fp16_nc_cuda_t ggml_get_to_fp16_nc_cuda(ggml_type type) { return dequantize_block_cuda; case GGML_TYPE_Q8_0: return dequantize_block_cuda; + case GGML_TYPE_Q4_0_ROCMFP4: + return dequantize_block_cuda; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + return dequantize_block_cuda; + case GGML_TYPE_Q2_0_ROCMFP2: + return dequantize_block_cuda; + case GGML_TYPE_Q3_0_ROCMFPX: + return dequantize_block_cuda; + case GGML_TYPE_Q6_0_ROCMFPX: + return dequantize_block_cuda; + case GGML_TYPE_Q8_0_ROCMFPX: + return dequantize_block_cuda; case GGML_TYPE_BF16: return convert_unary_cuda; default: @@ -885,6 +971,18 @@ to_bf16_nc_cuda_t ggml_get_to_bf16_nc_cuda(ggml_type type) { return dequantize_block_cuda; case GGML_TYPE_Q8_0: return dequantize_block_cuda; + case GGML_TYPE_Q4_0_ROCMFP4: + return dequantize_block_cuda; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + return dequantize_block_cuda; + case GGML_TYPE_Q2_0_ROCMFP2: + return dequantize_block_cuda; + case GGML_TYPE_Q3_0_ROCMFPX: + return dequantize_block_cuda; + case GGML_TYPE_Q6_0_ROCMFPX: + return dequantize_block_cuda; + case GGML_TYPE_Q8_0_ROCMFPX: + return dequantize_block_cuda; case GGML_TYPE_F16: return convert_unary_cuda; default: @@ -906,6 +1004,18 @@ to_fp32_nc_cuda_t ggml_get_to_fp32_nc_cuda(ggml_type type) { return dequantize_block_cuda; case GGML_TYPE_Q8_0: return dequantize_block_cuda; + case GGML_TYPE_Q4_0_ROCMFP4: + return dequantize_block_cuda; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + return dequantize_block_cuda; + case GGML_TYPE_Q2_0_ROCMFP2: + return dequantize_block_cuda; + case GGML_TYPE_Q3_0_ROCMFPX: + return dequantize_block_cuda; + case GGML_TYPE_Q6_0_ROCMFPX: + return dequantize_block_cuda; + case GGML_TYPE_Q8_0_ROCMFPX: + return dequantize_block_cuda; case GGML_TYPE_BF16: return convert_unary_cuda; default: diff --git a/ggml/src/ggml-cuda/cpy.cu b/ggml/src/ggml-cuda/cpy.cu index 372fc98b1af8..3d7a698b9b53 100644 --- a/ggml/src/ggml-cuda/cpy.cu +++ b/ggml/src/ggml-cuda/cpy.cu @@ -1,6 +1,7 @@ #include "cpy.cuh" #include "dequantize.cuh" #include "cpy-utils.cuh" +#include "../../rocmfp4/rocmfp4_hip_scale.cuh" #if defined(GGML_USE_MUSA) && defined(GGML_MUSA_MUDNN_COPY) #include "ggml-musa/mudnn.cuh" #endif // GGML_USE_MUSA && GGML_MUSA_MUDNN_COPY @@ -119,6 +120,72 @@ static __device__ void cpy_blck_q_f32(const char * cxi, char * cdsti) { } } +static __device__ void cpy_blck_rocmfp4_f32(const char * cxi, char * cdsti) { + const block_rocmfp4 * x = (const block_rocmfp4 *) cxi; + float * cdstf = (float *) cdsti; + + const float d0 = rocmfp4_ue4m3_to_fp32_half_finite(x->e[0]); + const float d1 = rocmfp4_ue4m3_to_fp32_half_finite(x->e[1]); + +#pragma unroll + for (int j = 0; j < QK_ROCMFP4/2; ++j) { + const uint8_t q = x->qs[j]; + cdstf[j] = d0 * (float) rocmfp4_decode_i8(q); + cdstf[j + QK_ROCMFP4/2] = d1 * (float) rocmfp4_decode_i8(q >> 4); + } +} + +static __device__ void cpy_blck_rocmfp4_fast_f32(const char * cxi, char * cdsti) { + const block_rocmfp4_fast * x = (const block_rocmfp4_fast *) cxi; + float * cdstf = (float *) cdsti; + + const float d = rocmfp4_ue4m3_to_fp32_half_finite(x->e); + +#pragma unroll + for (int j = 0; j < QK_ROCMFP4/2; ++j) { + const uint8_t q = x->qs[j]; + cdstf[j] = d * (float) rocmfp4_decode_i8(q); + cdstf[j + QK_ROCMFP4/2] = d * (float) rocmfp4_decode_i8(q >> 4); + } +} + +static __global__ void cpy_rocmfp4_f32_contiguous(const block_rocmfp4 * cx, float * cdst, const int64_t ne) { + const int64_t packed_idx = (int64_t)blockDim.x * blockIdx.x + threadIdx.x; + const int64_t packed_count = (ne / QK_ROCMFP4) * (QK_ROCMFP4/2); + + if (packed_idx >= packed_count) { + return; + } + + const int64_t ib = packed_idx >> 4; + const int j = packed_idx & 0x0f; + const int64_t base = ib * QK_ROCMFP4; + const uint8_t q = cx[ib].qs[j]; + const float d0 = rocmfp4_ue4m3_to_fp32_half_finite(cx[ib].e[0]); + const float d1 = rocmfp4_ue4m3_to_fp32_half_finite(cx[ib].e[1]); + + cdst[base + j] = d0 * (float) rocmfp4_decode_i8(q); + cdst[base + j + QK_ROCMFP4/2] = d1 * (float) rocmfp4_decode_i8(q >> 4); +} + +static __global__ void cpy_rocmfp4_fast_f32_contiguous(const block_rocmfp4_fast * cx, float * cdst, const int64_t ne) { + const int64_t packed_idx = (int64_t)blockDim.x * blockIdx.x + threadIdx.x; + const int64_t packed_count = (ne / QK_ROCMFP4) * (QK_ROCMFP4/2); + + if (packed_idx >= packed_count) { + return; + } + + const int64_t ib = packed_idx >> 4; + const int j = packed_idx & 0x0f; + const int64_t base = ib * QK_ROCMFP4; + const uint8_t q = cx[ib].qs[j]; + const float d = rocmfp4_ue4m3_to_fp32_half_finite(cx[ib].e); + + cdst[base + j] = d * (float) rocmfp4_decode_i8(q); + cdst[base + j + QK_ROCMFP4/2] = d * (float) rocmfp4_decode_i8(q >> 4); +} + template static __global__ void cpy_f32_q(const char * cx, char * cdst, const int64_t ne, const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t nb00, const int64_t nb01, const int64_t nb02, @@ -308,6 +375,56 @@ static void ggml_cpy_q4_1_f32_cuda( ne10, ne11, ne12, nb10, nb11, nb12, nb13); } +static void ggml_cpy_rocmfp4_f32_hip( + const char * cx, char * cdst, const int64_t ne, + const int64_t ne00, const int64_t ne01, const int64_t ne02, + const int64_t nb00, const int64_t nb01, const int64_t nb02, + const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12, + const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13, + cudaStream_t stream) { + GGML_ASSERT(ne % QK_ROCMFP4 == 0); + const int64_t num_blocks = ne / QK_ROCMFP4; + GGML_ASSERT(num_blocks < UINT_MAX); + cpy_q_f32<<>>( + cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, + ne10, ne11, ne12, nb10, nb11, nb12, nb13); +} + +static void ggml_cpy_rocmfp4_f32_contiguous_hip( + const char * cx, char * cdst, const int64_t ne, cudaStream_t stream) { + GGML_ASSERT(ne % QK_ROCMFP4 == 0); + const int64_t packed_count = (ne / QK_ROCMFP4) * (QK_ROCMFP4/2); + const int64_t num_blocks = (packed_count + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE; + GGML_ASSERT(num_blocks < UINT_MAX); + cpy_rocmfp4_f32_contiguous<<>>( + (const block_rocmfp4 *) cx, (float *) cdst, ne); +} + +static void ggml_cpy_rocmfp4_fast_f32_hip( + const char * cx, char * cdst, const int64_t ne, + const int64_t ne00, const int64_t ne01, const int64_t ne02, + const int64_t nb00, const int64_t nb01, const int64_t nb02, + const int64_t nb03, const int64_t ne10, const int64_t ne11, const int64_t ne12, + const int64_t nb10, const int64_t nb11, const int64_t nb12, const int64_t nb13, + cudaStream_t stream) { + GGML_ASSERT(ne % QK_ROCMFP4 == 0); + const int64_t num_blocks = ne / QK_ROCMFP4; + GGML_ASSERT(num_blocks < UINT_MAX); + cpy_q_f32<<>>( + cx, cdst, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, + ne10, ne11, ne12, nb10, nb11, nb12, nb13); +} + +static void ggml_cpy_rocmfp4_fast_f32_contiguous_hip( + const char * cx, char * cdst, const int64_t ne, cudaStream_t stream) { + GGML_ASSERT(ne % QK_ROCMFP4 == 0); + const int64_t packed_count = (ne / QK_ROCMFP4) * (QK_ROCMFP4/2); + const int64_t num_blocks = (packed_count + CUDA_CPY_BLOCK_SIZE - 1) / CUDA_CPY_BLOCK_SIZE; + GGML_ASSERT(num_blocks < UINT_MAX); + cpy_rocmfp4_fast_f32_contiguous<<>>( + (const block_rocmfp4_fast *) cx, (float *) cdst, ne); +} + static void ggml_cpy_f32_q5_0_cuda( const char * cx, char * cdst, const int64_t ne, const int64_t ne00, const int64_t ne01, const int64_t ne02, const int64_t nb00, const int64_t nb01, const int64_t nb02, @@ -556,6 +673,20 @@ void ggml_cuda_cpy(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, gg } else if (src0->type == GGML_TYPE_Q4_0 && src1->type == GGML_TYPE_F32) { ggml_cpy_q4_0_f32_cuda (src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream); + } else if (src0->type == GGML_TYPE_Q4_0_ROCMFP4 && src1->type == GGML_TYPE_F32) { + if (contiguous_srcs) { + ggml_cpy_rocmfp4_f32_contiguous_hip(src0_ddc, src1_ddc, ne, main_stream); + } else { + ggml_cpy_rocmfp4_f32_hip + (src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream); + } + } else if (src0->type == GGML_TYPE_Q4_0_ROCMFP4_FAST && src1->type == GGML_TYPE_F32) { + if (contiguous_srcs) { + ggml_cpy_rocmfp4_fast_f32_contiguous_hip(src0_ddc, src1_ddc, ne, main_stream); + } else { + ggml_cpy_rocmfp4_fast_f32_hip + (src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream); + } } else if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_Q4_1) { ggml_cpy_f32_q4_1_cuda (src0_ddc, src1_ddc, ne, ne00, ne01, ne02, nb00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb13, main_stream); diff --git a/ggml/src/ggml-cuda/dequantize.cuh b/ggml/src/ggml-cuda/dequantize.cuh index 1801bd0489bc..dc660e1aabf5 100644 --- a/ggml/src/ggml-cuda/dequantize.cuh +++ b/ggml/src/ggml-cuda/dequantize.cuh @@ -1,5 +1,6 @@ #include "common.cuh" #include "tq3-quant.cuh" +#include "../../rocmfp4/rocmfp4_hip_scale.cuh" static __device__ __forceinline__ void dequantize_q4_0(const void * vx, const int64_t ib, const int iqs, float2 & v){ const block_q4_0 * x = (const block_q4_0 *) vx; @@ -76,3 +77,125 @@ static __device__ __forceinline__ void dequantize_q8_0(const void * vx, const in v.x *= d; v.y *= d; } + +static __device__ __forceinline__ void dequantize_rocmfp4(const void * vx, const int64_t ib, const int iqs, float2 & v) { + const block_rocmfp4 * x = (const block_rocmfp4 *) vx; + + const int q = x[ib].qs[iqs]; + const float d0 = rocmfp4_ue4m3_to_fp32_half_finite(x[ib].e[0]); + const float d1 = rocmfp4_ue4m3_to_fp32_half_finite(x[ib].e[1]); + + v.x = d0 * rocmfp4_decode_i8(q); + v.y = d1 * rocmfp4_decode_i8(q >> 4); +} + +static __device__ __forceinline__ void dequantize_rocmfp4_fast(const void * vx, const int64_t ib, const int iqs, float2 & v) { + const block_rocmfp4_fast * x = (const block_rocmfp4_fast *) vx; + + const int q = x[ib].qs[iqs]; + const float d = rocmfp4_ue4m3_to_fp32_half_finite(x[ib].e); + + v.x = d * rocmfp4_decode_i8(q); + v.y = d * rocmfp4_decode_i8(q >> 4); +} + +template +static __device__ __forceinline__ uint32_t rocmfpx_load_qs_window_cuda(const uint8_t * src, const int byte_pos) { + uint32_t v = (uint32_t) src[byte_pos + 0]; + + if (byte_pos + 1 < qs) { + v |= (uint32_t) src[byte_pos + 1] << 8; + } + if (byte_pos + 2 < qs) { + v |= (uint32_t) src[byte_pos + 2] << 16; + } + + return v; +} + +static __device__ __forceinline__ uint32_t rocmfpx_get_fp3_code_cuda(const uint8_t * src, const int i) { + const int bit_pos = i * 3; + const int byte_pos = bit_pos >> 3; + const int shift = bit_pos & 7; + return (rocmfpx_load_qs_window_cuda(src, byte_pos) >> shift) & 7u; +} + +static __device__ __forceinline__ uint32_t rocmfpx_get_fp2_code_cuda(const uint8_t * src, const int i) { + return (src[i >> 2] >> (2*(i & 3))) & 3u; +} + +static __device__ __forceinline__ uint32_t rocmfpx_get_fp6_code_cuda(const uint8_t * src, const int i) { + const int bit_pos = i * 6; + const int byte_pos = bit_pos >> 3; + const int shift = bit_pos & 7; + return (rocmfpx_load_qs_window_cuda(src, byte_pos) >> shift) & 63u; +} + +static __device__ __forceinline__ int rocmfpx_decode_fp3_code_cuda(const uint32_t code) { + const uint32_t mag_code = code & 3u; + const int mag = mag_code == 3u ? 4 : (int) mag_code; + return (code & 4u) ? -mag : mag; +} + +static __device__ __forceinline__ float rocmfpx_decode_fp2_code_cuda(const uint32_t code) { + switch (code & 3u) { + case 0: return (float) ROCMFP2_KVALUE_0_I8; + case 1: return (float) ROCMFP2_KVALUE_1_I8; + case 2: return (float) ROCMFP2_KVALUE_2_I8; + default: return (float) ROCMFP2_KVALUE_3_I8; + } +} + +static __device__ __forceinline__ int rocmfpx_decode_fp6_code_cuda(const uint32_t code) { + const int mag = (int) (code & 31u); + return (code & 32u) ? -mag : mag; +} + +static __device__ __forceinline__ void dequantize_rocmfpx_fp3(const void * vx, const int64_t ib, const int iqs, float2 & v) { + const block_rocmfp3 * x = (const block_rocmfp3 *) vx; + + const int i0 = iqs + 0; + const int i1 = iqs + 1; + const float d0 = rocmfpx_ue4m3_to_fp32_finite(x[ib].e[i0 >= QK_ROCMFP3/2]); + const float d1 = rocmfpx_ue4m3_to_fp32_finite(x[ib].e[i1 >= QK_ROCMFP3/2]); + + v.x = d0 * (float) rocmfpx_decode_fp3_code_cuda(rocmfpx_get_fp3_code_cuda(x[ib].qs, i0)); + v.y = d1 * (float) rocmfpx_decode_fp3_code_cuda(rocmfpx_get_fp3_code_cuda(x[ib].qs, i1)); +} + +static __device__ __forceinline__ void dequantize_rocmfpx_fp2(const void * vx, const int64_t ib, const int iqs, float2 & v) { + const block_rocmfp2 * x = (const block_rocmfp2 *) vx; + + const int i0 = iqs + 0; + const int i1 = iqs + 1; + const float d0 = rocmfpx_ue4m3_to_fp32_finite(x[ib].e[i0 >= QK_ROCMFP2/2]); + const float d1 = rocmfpx_ue4m3_to_fp32_finite(x[ib].e[i1 >= QK_ROCMFP2/2]); + + v.x = d0 * rocmfpx_decode_fp2_code_cuda(rocmfpx_get_fp2_code_cuda(x[ib].qs, i0)); + v.y = d1 * rocmfpx_decode_fp2_code_cuda(rocmfpx_get_fp2_code_cuda(x[ib].qs, i1)); +} + +static __device__ __forceinline__ void dequantize_rocmfpx_fp6(const void * vx, const int64_t ib, const int iqs, float2 & v) { + const block_rocmfp6_device * x = (const block_rocmfp6_device *) vx; + + const int i0 = iqs + 0; + const int i1 = iqs + 1; + const float d0 = rocmfpx_ue4m3_to_fp32_finite(x[ib].e[i0 >= QK_ROCMFP6/2]); + const float d1 = rocmfpx_ue4m3_to_fp32_finite(x[ib].e[i1 >= QK_ROCMFP6/2]); + +#if GGML_ROCMFP6_EXPANDED_DEVICE + v.x = d0 * (float) x[ib].qs[i0]; + v.y = d1 * (float) x[ib].qs[i1]; +#else + v.x = d0 * (float) rocmfpx_decode_fp6_code_cuda(rocmfpx_get_fp6_code_cuda(x[ib].qs, i0)); + v.y = d1 * (float) rocmfpx_decode_fp6_code_cuda(rocmfpx_get_fp6_code_cuda(x[ib].qs, i1)); +#endif +} + +static __device__ __forceinline__ void dequantize_rocmfpx_fp8(const void * vx, const int64_t ib, const int iqs, float2 & v) { + const block_rocmfp8 * x = (const block_rocmfp8 *) vx; + + const float d = rocmfpx_ue4m3_to_fp32_finite(x[ib].e); + v.x = d * (float) x[ib].qs[iqs + 0]; + v.y = d * (float) x[ib].qs[iqs + 1]; +} diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index 0eff4f1b7c6a..e991bf4a30de 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -7,6 +7,133 @@ #include "fattn-chunked.cuh" #include "fattn.cuh" +__device__ static float ds4_fa_block_sum(float v) { + __shared__ float smem[256]; + const int tid = threadIdx.x; + smem[tid] = v; + __syncthreads(); + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (tid < stride) smem[tid] += smem[tid + stride]; + __syncthreads(); + } + return smem[0]; +} + +__device__ static float ds4_fa_block_max(float v) { + __shared__ float smem[256]; + const int tid = threadIdx.x; + smem[tid] = v; + __syncthreads(); + for (int stride = blockDim.x / 2; stride > 0; stride >>= 1) { + if (tid < stride) smem[tid] = fmaxf(smem[tid], smem[tid + stride]); + __syncthreads(); + } + return smem[0]; +} + +__global__ static void ds4_flash_attn_d512_f32_shared_kv_kernel( + float * dst, + const float * q, + const float * k, + const float * v, + const float * sinks, + int n_tokens, + int n_heads, + int n_kv) { + constexpr int D = 512; + const int t = (int) blockIdx.x; + const int h = (int) blockIdx.y; + const int tid = (int) threadIdx.x; + if (t >= n_tokens || h >= n_heads) return; + + extern __shared__ float scores[]; + const float * qh = q + ((size_t) h * (size_t) n_tokens + (size_t) t) * D; + const float scale = rsqrtf((float) D); + + float local_max = sinks ? sinks[h] : -3.402823466e38f; + for (int r = tid; r < n_kv; r += blockDim.x) { + const float * kr = k + (size_t) r * D; + float dot = 0.0f; +#pragma unroll + for (int d = 0; d < D; ++d) { + dot += qh[d] * kr[d]; + } + const float s = dot * scale; + scores[r] = s; + local_max = fmaxf(local_max, s); + } + const float max_score = ds4_fa_block_max(local_max); + + float local_sum = 0.0f; + for (int r = tid; r < n_kv; r += blockDim.x) { + const float w = expf(scores[r] - max_score); + scores[r] = w; + local_sum += w; + } + if (tid == 0 && sinks) { + local_sum += expf(sinks[h] - max_score); + } + const float denom = ds4_fa_block_sum(local_sum); + const float inv_denom = 1.0f / denom; + + for (int d = tid; d < D; d += blockDim.x) { + float acc = 0.0f; + for (int r = 0; r < n_kv; ++r) { + acc += scores[r] * v[(size_t) r * D + d]; + } + dst[((size_t) t * (size_t) n_heads + (size_t) h) * D + d] = acc * inv_denom; + } +} + +static bool ggml_cuda_ds4_flash_attn_d512_f32(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const char * e = getenv("DFLASH_DS4_FLASH_ATTN"); + if (!e || !e[0] || e[0] == (char)48) { + return false; + } + const ggml_tensor * Q = dst->src[0]; + const ggml_tensor * K = dst->src[1]; + const ggml_tensor * V = dst->src[2]; + const ggml_tensor * mask = dst->src[3]; + const ggml_tensor * sinks = dst->src[4]; + if (!Q || !K || !V || mask || + Q->type != GGML_TYPE_F32 || K->type != GGML_TYPE_F32 || V->type != GGML_TYPE_F32 || + dst->type != GGML_TYPE_F32 || + Q->ne[0] != 512 || K->ne[0] != 512 || V->ne[0] != 512 || + K->ne[2] != 1 || V->ne[2] != 1 || + Q->ne[3] != 1 || K->ne[3] != 1 || V->ne[3] != 1 || + dst->ne[0] != 512 || dst->ne[1] != Q->ne[2] || dst->ne[2] != Q->ne[1] || + Q->nb[0] != (int64_t) sizeof(float) || + K->nb[0] != (int64_t) sizeof(float) || + V->nb[0] != (int64_t) sizeof(float) || + dst->nb[0] != (int64_t) sizeof(float)) { + return false; + } + if (sinks && (sinks->type != GGML_TYPE_F32 || sinks->ne[0] != Q->ne[2])) { + return false; + } + + const int n_tokens = (int) Q->ne[1]; + const int n_heads = (int) Q->ne[2]; + const int n_kv = (int) K->ne[1]; + if (n_tokens <= 0 || n_heads <= 0 || n_kv <= 0) { + return false; + } + + cudaStream_t stream = ctx.stream(); + dim3 grid((unsigned) n_tokens, (unsigned) n_heads, 1); + const size_t shmem = (size_t) n_kv * sizeof(float); + ds4_flash_attn_d512_f32_shared_kv_kernel<<>>( + (float *) dst->data, + (const float *) Q->data, + (const float *) K->data, + (const float *) V->data, + sinks ? (const float *) sinks->data : nullptr, + n_tokens, + n_heads, + n_kv); + return true; +} + template static void ggml_cuda_flash_attn_ext_mma_f16_switch_ncols1(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; @@ -354,6 +481,10 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const } } } + static const bool ds4_allow_d512_nomask = [] { + const char * e = getenv("DFLASH_DS4_FLASH_ATTN"); + return e && e[0] && e[0] != (char)48; + }(); const int cc = ggml_cuda_info().devices[device].cc; @@ -374,7 +505,7 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const if (V->ne[0] != K->ne[0]) { return BEST_FATTN_KERNEL_NONE; } - if (!gqa_opt_applies) { + if (!gqa_opt_applies && !ds4_allow_d512_nomask) { return BEST_FATTN_KERNEL_NONE; } break; @@ -570,6 +701,9 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ggml_cuda_set_device(ctx.device); + if (ggml_cuda_ds4_flash_attn_d512_f32(ctx, dst)) { + return; + } switch (ggml_cuda_get_best_fattn_kernel(ggml_cuda_get_device(), dst)) { case BEST_FATTN_KERNEL_NONE: GGML_ABORT("fatal error"); diff --git a/ggml/src/ggml-cuda/getrows.cu b/ggml/src/ggml-cuda/getrows.cu index 1d7c6d17a525..73a28e5298e9 100644 --- a/ggml/src/ggml-cuda/getrows.cu +++ b/ggml/src/ggml-cuda/getrows.cu @@ -199,6 +199,30 @@ static void ggml_cuda_get_rows_switch_src0_type( get_rows_cuda_q(src0_d, src1_d, dst_d, ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); break; + case GGML_TYPE_Q4_0_ROCMFP4: + get_rows_cuda_q(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + get_rows_cuda_q(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_Q2_0_ROCMFP2: + get_rows_cuda_q(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_Q3_0_ROCMFPX: + get_rows_cuda_q(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_Q6_0_ROCMFPX: + get_rows_cuda_q(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; + case GGML_TYPE_Q8_0_ROCMFPX: + get_rows_cuda_q(src0_d, src1_d, dst_d, + ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); + break; case GGML_TYPE_TQ3_0: get_rows_cuda_q(src0_d, src1_d, dst_d, ne00, nb01, nb02, nb03, ne10, ne11, ne12, nb10, nb11, nb12, nb1, nb2, nb3, stream); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index c09e3d7c8fd1..9d1c870879a4 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -2361,7 +2361,12 @@ static bool ggml_cuda_should_fuse_mul_mat(const ggml_tensor * ffn_up, return false; } - static constexpr std::array valid_glu_ops = { GGML_GLU_OP_SWIGLU, GGML_GLU_OP_GEGLU, GGML_GLU_OP_SWIGLU_OAI }; + static constexpr std::array valid_glu_ops = { + GGML_GLU_OP_SWIGLU, + GGML_GLU_OP_GEGLU, + GGML_GLU_OP_SWIGLU_OAI, + GGML_GLU_OP_SWIGLU_DS4, + }; if (std::find(valid_glu_ops.begin(), valid_glu_ops.end(), ggml_get_glu_op(glu)) == valid_glu_ops.end()) { return false; @@ -2417,6 +2422,12 @@ static bool ggml_cuda_should_fuse_mul_mat_vec_f(const ggml_tensor * tensor) { return use_mul_mat_vec_f; } +static inline void ggml_cuda_set_fusion_glu_params(ggml_cuda_mm_fusion_args_host & fusion_data, const ggml_tensor * glu) { + fusion_data.glu_op = ggml_get_glu_op(glu); + fusion_data.glu_param0 = ggml_get_op_params_f32(glu, 2); + fusion_data.glu_param1 = ggml_get_op_params_f32(glu, 3); +} + static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) { ggml_tensor * src0 = tensor->src[0]; ggml_tensor * src1 = tensor->src[1]; @@ -2846,6 +2857,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg case GGML_GLU_OP_SWIGLU_OAI: ggml_cuda_op_swiglu_oai(ctx, dst); break; + case GGML_GLU_OP_SWIGLU_DS4: + ggml_cuda_op_swiglu_ds4(ctx, dst); + break; case GGML_GLU_OP_GEGLU_ERF: ggml_cuda_op_geglu_erf(ctx, dst); break; @@ -4033,7 +4047,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud fusion_data.gate = gate_n->src[0]; fusion_data.x_bias = up_bias_tensor; fusion_data.gate_bias = gate_bias_tensor; - fusion_data.glu_op = ggml_get_glu_op(glu); + ggml_cuda_set_fusion_glu_params(fusion_data, glu); ggml_cuda_mul_mat_vec_f(*cuda_ctx, src0, src1, ids, glu, &fusion_data); fused_mul_mat_vec = true; @@ -4046,7 +4060,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud fusion_data.gate = gate_n->src[0]; fusion_data.x_bias = up_bias_tensor; fusion_data.gate_bias = gate_bias_tensor; - fusion_data.glu_op = ggml_get_glu_op(glu); + ggml_cuda_set_fusion_glu_params(fusion_data, glu); ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, glu, &fusion_data); fused_mul_mat_vec = true; @@ -4070,7 +4084,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud if (ggml_cuda_should_fuse_mul_mat_vec_f(up)) { ggml_cuda_mm_fusion_args_host fusion_data{}; fusion_data.gate = gate->src[0]; - fusion_data.glu_op = ggml_get_glu_op(glu); + ggml_cuda_set_fusion_glu_params(fusion_data, glu); ggml_cuda_mul_mat_vec_f(*cuda_ctx, src0, src1, ids, glu, &fusion_data); fused_mul_mat_vec = true; @@ -4081,7 +4095,7 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud if (ggml_cuda_should_fuse_mul_mat_vec_q(up)) { ggml_cuda_mm_fusion_args_host fusion_data{}; fusion_data.gate = gate->src[0]; - fusion_data.glu_op = ggml_get_glu_op(glu); + ggml_cuda_set_fusion_glu_params(fusion_data, glu); ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, glu, &fusion_data); fused_mul_mat_vec = true; @@ -4936,6 +4950,7 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g case GGML_GLU_OP_GEGLU: case GGML_GLU_OP_SWIGLU: case GGML_GLU_OP_SWIGLU_OAI: + case GGML_GLU_OP_SWIGLU_DS4: case GGML_GLU_OP_GEGLU_ERF: case GGML_GLU_OP_GEGLU_QUICK: return ggml_is_contiguous_1(op->src[0]); diff --git a/ggml/src/ggml-cuda/mmvf.cu b/ggml/src/ggml-cuda/mmvf.cu index d91472024296..6e4db8e8bbcf 100644 --- a/ggml/src/ggml-cuda/mmvf.cu +++ b/ggml/src/ggml-cuda/mmvf.cu @@ -51,6 +51,8 @@ static __global__ void mul_mat_vec_f( bool use_bias = false; bool use_gate_bias = false; ggml_glu_op glu_op = ggml_glu_op::GGML_GLU_OP_SWIGLU; + float glu_param0 = 0.0f; + float glu_param1 = 0.0f; const T * gate_x = nullptr; const float * x_bias = nullptr; const float * gate_bias = nullptr; @@ -60,6 +62,8 @@ static __global__ void mul_mat_vec_f( use_bias = fusion.x_bias != nullptr; use_gate_bias = fusion.gate_bias != nullptr; glu_op = fusion.glu_op; + glu_param0 = fusion.glu_param0; + glu_param1 = fusion.glu_param1; if (use_gate) { gate_x = static_cast(fusion.gate); @@ -357,7 +361,11 @@ static __global__ void mul_mat_vec_f( value *= ggml_cuda_op_gelu_single(gate_value); break; case GGML_GLU_OP_SWIGLU_OAI: { - value = ggml_cuda_op_swiglu_oai_single(gate_value, value); + value = ggml_cuda_op_swiglu_oai_single(gate_value, value, glu_param0, glu_param1); + break; + } + case GGML_GLU_OP_SWIGLU_DS4: { + value = ggml_cuda_op_swiglu_ds4_single(gate_value, value, glu_param0); break; } default: @@ -369,7 +377,7 @@ static __global__ void mul_mat_vec_f( dst[tid*stride_col_dst + row] = value; if constexpr (!has_fusion) { - GGML_UNUSED_VARS(use_gate, use_bias, use_gate_bias, glu_op, gate_x, x_bias, gate_bias, sumf_gate); + GGML_UNUSED_VARS(use_gate, use_bias, use_gate_bias, glu_op, glu_param0, glu_param1, gate_x, x_bias, gate_bias, sumf_gate); } } @@ -668,6 +676,8 @@ void ggml_cuda_mul_mat_vec_f(ggml_backend_cuda_context & ctx, const ggml_tensor fusion_local.gate_bias = fusion->gate_bias->data; } fusion_local.glu_op = fusion->glu_op; + fusion_local.glu_param0 = fusion->glu_param0; + fusion_local.glu_param1 = fusion->glu_param1; } const int64_t s01 = src0->nb[1] / ts_src0; diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index b674727178fa..558cae933516 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -17,6 +17,12 @@ static constexpr __device__ vec_dot_q_cuda_t get_vec_dot_q_cuda(ggml_type type) case GGML_TYPE_Q8_0: return vec_dot_q8_0_q8_1; case GGML_TYPE_MXFP4: return vec_dot_mxfp4_q8_1; case GGML_TYPE_NVFP4: return vec_dot_nvfp4_q8_1; + case GGML_TYPE_Q4_0_ROCMFP4: return vec_dot_rocmfp4_q8_1; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: return vec_dot_rocmfp4_fast_q8_1; + case GGML_TYPE_Q2_0_ROCMFP2: return vec_dot_rocmfpx_fp2_q8_1; + case GGML_TYPE_Q3_0_ROCMFPX: return vec_dot_rocmfpx_fp3_q8_1; + case GGML_TYPE_Q6_0_ROCMFPX: return vec_dot_rocmfpx_fp6_q8_1; + case GGML_TYPE_Q8_0_ROCMFPX: return vec_dot_rocmfpx_fp8_q8_1; case GGML_TYPE_Q2_K: return vec_dot_q2_K_q8_1; case GGML_TYPE_Q3_K: return vec_dot_q3_K_q8_1; case GGML_TYPE_Q4_K: return vec_dot_q4_K_q8_1; @@ -44,6 +50,12 @@ static constexpr __host__ __device__ int get_vdr_mmvq(ggml_type type) { case GGML_TYPE_Q8_0: return VDR_Q8_0_Q8_1_MMVQ; case GGML_TYPE_MXFP4: return VDR_MXFP4_Q8_1_MMVQ; case GGML_TYPE_NVFP4: return VDR_NVFP4_Q8_1_MMVQ; + case GGML_TYPE_Q4_0_ROCMFP4: return VDR_ROCMFP4_Q8_1_MMVQ; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: return VDR_ROCMFP4_FAST_Q8_1_MMVQ; + case GGML_TYPE_Q2_0_ROCMFP2: return VDR_ROCMFP2_Q8_1_MMVQ; + case GGML_TYPE_Q3_0_ROCMFPX: return VDR_ROCMFP3_Q8_1_MMVQ; + case GGML_TYPE_Q6_0_ROCMFPX: return VDR_ROCMFP6_Q8_1_MMVQ; + case GGML_TYPE_Q8_0_ROCMFPX: return VDR_ROCMFP8_Q8_1_MMVQ; case GGML_TYPE_Q2_K: return VDR_Q2_K_Q8_1_MMVQ; case GGML_TYPE_Q3_K: return VDR_Q3_K_Q8_1_MMVQ; case GGML_TYPE_Q4_K: return VDR_Q4_K_Q8_1_MMVQ; @@ -438,6 +450,8 @@ static __global__ void mul_mat_vec_q( const float * x_bias = nullptr; const float * gate_bias = nullptr; ggml_glu_op active_glu; + float glu_param0 = 0.0f; + float glu_param1 = 0.0f; if constexpr (has_fusion) { use_gate = fusion.gate != nullptr; @@ -447,6 +461,8 @@ static __global__ void mul_mat_vec_q( x_bias = (const float *) fusion.x_bias; gate_bias = (const float *) fusion.gate_bias; active_glu = fusion.glu_op; + glu_param0 = fusion.glu_param0; + glu_param1 = fusion.glu_param1; } @@ -580,7 +596,11 @@ static __global__ void mul_mat_vec_q( result *= ggml_cuda_op_gelu_single(gate_value); break; case GGML_GLU_OP_SWIGLU_OAI: { - result = ggml_cuda_op_swiglu_oai_single(gate_value, result); + result = ggml_cuda_op_swiglu_oai_single(gate_value, result, glu_param0, glu_param1); + break; + } + case GGML_GLU_OP_SWIGLU_DS4: { + result = ggml_cuda_op_swiglu_ds4_single(gate_value, result, glu_param0); break; } default: @@ -594,7 +614,7 @@ static __global__ void mul_mat_vec_q( } if constexpr (!has_fusion) { - GGML_UNUSED_VARS(use_gate, use_bias, use_gate_bias, active_glu, gate_bias, x_bias, tmp_gate); + GGML_UNUSED_VARS(use_gate, use_bias, use_gate_bias, active_glu, glu_param0, glu_param1, gate_bias, x_bias, tmp_gate); } } @@ -644,6 +664,8 @@ static __global__ void mul_mat_vec_q_moe( const float * x_bias = nullptr; const float * gate_bias = nullptr; ggml_glu_op active_glu; + float glu_param0 = 0.0f; + float glu_param1 = 0.0f; if constexpr (has_fusion) { use_gate = fusion.gate != nullptr; @@ -653,6 +675,8 @@ static __global__ void mul_mat_vec_q_moe( x_bias = (const float *) fusion.x_bias; gate_bias = (const float *) fusion.gate_bias; active_glu = fusion.glu_op; + glu_param0 = fusion.glu_param0; + glu_param1 = fusion.glu_param1; } // partial sum for each thread @@ -705,7 +729,10 @@ static __global__ void mul_mat_vec_q_moe( result *= ggml_cuda_op_gelu_single(gate_value); break; case GGML_GLU_OP_SWIGLU_OAI: - result = ggml_cuda_op_swiglu_oai_single(gate_value, result); + result = ggml_cuda_op_swiglu_oai_single(gate_value, result, glu_param0, glu_param1); + break; + case GGML_GLU_OP_SWIGLU_DS4: + result = ggml_cuda_op_swiglu_ds4_single(gate_value, result, glu_param0); break; default: result = result * gate_value; @@ -717,7 +744,7 @@ static __global__ void mul_mat_vec_q_moe( } if constexpr (!has_fusion) { - GGML_UNUSED_VARS(use_gate, use_bias, use_gate_bias, vgate, x_bias, gate_bias, active_glu, tmp_gate); + GGML_UNUSED_VARS(use_gate, use_bias, use_gate_bias, vgate, x_bias, gate_bias, active_glu, glu_param0, glu_param1, tmp_gate); } } @@ -1066,6 +1093,42 @@ static void mul_mat_vec_q_switch_type( nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); break; + case GGML_TYPE_Q4_0_ROCMFP4: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; + case GGML_TYPE_Q2_0_ROCMFP2: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; + case GGML_TYPE_Q3_0_ROCMFPX: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; + case GGML_TYPE_Q6_0_ROCMFPX: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; + case GGML_TYPE_Q8_0_ROCMFPX: + mul_mat_vec_q_switch_ncols_dst + (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + break; case GGML_TYPE_Q2_K: mul_mat_vec_q_switch_ncols_dst (vx, vy, ids, fusion, dst, ncols_x, nrows_x, ncols_dst, stride_row_x, stride_col_y, stride_col_dst, @@ -1203,6 +1266,8 @@ void ggml_cuda_mul_mat_vec_q( fusion_local.gate_bias = fusion->gate_bias->data; } fusion_local.glu_op = fusion->glu_op; + fusion_local.glu_param0 = fusion->glu_param0; + fusion_local.glu_param1 = fusion->glu_param1; } // If src0 is a temporary compute buffer, clear any potential padding. diff --git a/ggml/src/ggml-cuda/unary.cu b/ggml/src/ggml-cuda/unary.cu index 4ad30fa1f353..ed7f3d02a6bb 100644 --- a/ggml/src/ggml-cuda/unary.cu +++ b/ggml/src/ggml-cuda/unary.cu @@ -338,6 +338,63 @@ void ggml_cuda_op_swiglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ggml_cuda_op_unary_gated(ctx, dst); } +// swiglu_ds4 + +template +static __global__ void swiglu_ds4_kernel(const T * gate, const T * up, T * dst, const int64_t k, const int64_t n, const int64_t o0, const int64_t o1, float limit) { + const int64_t i = int64_t(blockDim.x)*blockIdx.x + threadIdx.x; + + if (i >= k) { + return; + } + + const int64_t j0 = (i / n) * o0 + (i % n); + const int64_t j1 = o0 == o1 ? j0 : (i / n) * o1 + (i % n); + + const float gate_v = gate[j0]; + const float up_v = up[j1]; + + dst[i] = ggml_cuda_op_swiglu_ds4_single(gate_v, up_v, limit); +} + +template +static void swiglu_ds4_cuda(const T * gate, const T * up, T * dst, const int64_t k, const int64_t n, const int64_t o0, const int64_t o1, const float limit, cudaStream_t stream) { + const int64_t num_blocks = (k + CUDA_GLU_BLOCK_SIZE - 1) / CUDA_GLU_BLOCK_SIZE; + swiglu_ds4_kernel<<>>(gate, up, dst, k, n, o0, o1, limit); +} + +void ggml_cuda_op_swiglu_ds4(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + void * src0_d = src0->data; + void * src1_d = src1 ? src1->data : src0->data; + const int64_t src0_o = src0->nb[1]; + const int64_t src1_o = src1 ? src1->nb[1] : src0->nb[1]; + void * dst_d = dst->data; + const int64_t nc = src1 ? src0->ne[0] : src0->ne[0] / 2; + cudaStream_t stream = ctx.stream(); + + GGML_ASSERT(ggml_is_contiguous_1(src0)); + GGML_ASSERT(src0->nb[0] == ggml_element_size(src0)); + GGML_ASSERT(ggml_is_contiguous(dst)); + + GGML_ASSERT(src0->type == GGML_TYPE_F32); + GGML_ASSERT( dst->type == GGML_TYPE_F32); + GGML_ASSERT(src0->type == dst->type); + GGML_ASSERT(dst->ne[0] == nc); + GGML_ASSERT(ggml_nrows(dst) == ggml_nrows(src0)); + + GGML_ASSERT(src1); + GGML_ASSERT(ggml_is_contiguous_1(src1)); + GGML_ASSERT(src1->nb[0] == ggml_element_size(src1)); + GGML_ASSERT(src1->ne[0] == nc); + GGML_ASSERT(src0->type == src1->type); + + const float limit = ggml_get_op_params_f32(dst, 2); + + swiglu_ds4_cuda((float *) src0_d, (float *) src1_d, (float *)dst_d, ggml_nelements(dst), nc, src0_o / sizeof(float), src1_o / sizeof(float), limit, stream); +} + void ggml_cuda_op_geglu_erf(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ggml_cuda_op_unary_gated(ctx, dst); } diff --git a/ggml/src/ggml-cuda/unary.cuh b/ggml/src/ggml-cuda/unary.cuh index f1dd2183a6c7..854b27e35097 100644 --- a/ggml/src/ggml-cuda/unary.cuh +++ b/ggml/src/ggml-cuda/unary.cuh @@ -81,6 +81,8 @@ void ggml_cuda_op_geglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst); void ggml_cuda_op_swiglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst); +void ggml_cuda_op_swiglu_ds4(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + void ggml_cuda_op_swiglu_oai(ggml_backend_cuda_context & ctx, ggml_tensor * dst); void ggml_cuda_op_geglu_erf(ggml_backend_cuda_context & ctx, ggml_tensor * dst); @@ -110,3 +112,10 @@ __device__ __forceinline__ float ggml_cuda_op_swiglu_oai_single(float x, float g out_glu = out_glu * (1.0f + g); return out_glu; } + +__device__ __forceinline__ float ggml_cuda_op_swiglu_ds4_single(float gate, float up, float limit) { + gate = fminf(gate, limit); + up = fmaxf(fminf(up, limit), -limit); + + return up * gate / (1.0f + expf(-gate)); +} diff --git a/ggml/src/ggml-cuda/vecdotq.cuh b/ggml/src/ggml-cuda/vecdotq.cuh index 40b2b41e7e82..296615392019 100644 --- a/ggml/src/ggml-cuda/vecdotq.cuh +++ b/ggml/src/ggml-cuda/vecdotq.cuh @@ -1,6 +1,7 @@ #pragma once #include "common.cuh" +#include "../../rocmfp4/rocmfp4_hip_codebook.cuh" #include @@ -322,6 +323,377 @@ static __device__ __forceinline__ float vec_dot_mxfp4_q8_1( return d * sumi; } +// === ROCMFP vec_dot device functions (ported from ROCmFPX) === +#ifndef GGML_ROCMFP4_Q8_1_MMQ_VDR +#define GGML_ROCMFP4_Q8_1_MMQ_VDR 8 +#endif + +#ifndef GGML_ROCMFP4_FAST_Q8_1_MMQ_VDR +#define GGML_ROCMFP4_FAST_Q8_1_MMQ_VDR GGML_ROCMFP4_Q8_1_MMQ_VDR +#endif + +#ifndef GGML_ROCMFP4_FAST_Q8_1_MMVQ_VDR +#define GGML_ROCMFP4_FAST_Q8_1_MMVQ_VDR 2 +#endif + +#if GGML_ROCMFP4_FAST_Q8_1_MMVQ_VDR != 1 && \ + GGML_ROCMFP4_FAST_Q8_1_MMVQ_VDR != 2 && \ + GGML_ROCMFP4_FAST_Q8_1_MMVQ_VDR != 4 +#error "GGML_ROCMFP4_FAST_Q8_1_MMVQ_VDR must be 1, 2, or 4" +#endif + +#define VDR_ROCMFP4_Q8_1_MMVQ 4 +#define VDR_ROCMFP4_Q8_1_MMQ GGML_ROCMFP4_Q8_1_MMQ_VDR +#define VDR_ROCMFP4_FAST_Q8_1_MMVQ GGML_ROCMFP4_FAST_Q8_1_MMVQ_VDR +#define VDR_ROCMFP4_FAST_Q8_1_MMQ GGML_ROCMFP4_FAST_Q8_1_MMQ_VDR +#define VDR_ROCMFP2_Q8_1_MMVQ 1 +#define VDR_ROCMFP3_Q8_1_MMVQ 2 +#ifndef VDR_ROCMFP6_Q8_1_MMVQ +#define VDR_ROCMFP6_Q8_1_MMVQ 4 +#endif +#ifndef GGML_ROCMFP6_FAST_SIGNMAG_PACK +#define GGML_ROCMFP6_FAST_SIGNMAG_PACK 0 +#endif +#ifndef GGML_ROCMFP6_MMVQ_HALF_BLOCK_SPLIT +// Enabled by default: hoisting the half-block scale selection out of the FP6 +// MMVQ dot loop measured ~1.8% faster tg on gfx1151 (Qwen3-0.6B Q6_0_ROCMFPX, +// 211.8 -> 215.7 t/s) with bit-identical results. Set to 0 to restore the +// per-element branch. +#define GGML_ROCMFP6_MMVQ_HALF_BLOCK_SPLIT 1 +#endif +#define VDR_ROCMFP8_Q8_1_MMVQ 2 + +#define VDR_ROCMFP2_Q8_1_MMQ 4 +#define VDR_ROCMFP3_Q8_1_MMQ 4 +#ifndef VDR_ROCMFP6_Q8_1_MMQ +#define VDR_ROCMFP6_Q8_1_MMQ 4 +#endif +#define VDR_ROCMFP8_Q8_1_MMQ 8 + +static __device__ __forceinline__ uint32_t rocmfpx_get_bits_vec_cuda(const uint8_t * src, const int bit_pos, const int nbits) { + uint32_t code = 0; + +#pragma unroll + for (int bit = 0; bit < nbits; ++bit) { + const int src_bit = bit_pos + bit; + code |= ((uint32_t) ((src[src_bit >> 3] >> (src_bit & 7)) & 1u)) << bit; + } + + return code; +} + +static __device__ __forceinline__ int rocmfpx_decode_fp3_code_vec_cuda(const uint32_t code) { + const uint32_t mag_code = code & 3u; + const int mag = mag_code == 3u ? 4 : (int) mag_code; + return (code & 4u) ? -mag : mag; +} + +static __device__ __forceinline__ int rocmfpx_decode_fp2_code_vec_cuda(const uint32_t code) { + switch (code & 3u) { + case 0: return ROCMFP2_KVALUE_0_I8; + case 1: return ROCMFP2_KVALUE_1_I8; + case 2: return ROCMFP2_KVALUE_2_I8; + default: return ROCMFP2_KVALUE_3_I8; + } +} + +static __device__ __forceinline__ int rocmfpx_decode_fp6_code_vec_cuda(const uint32_t code) { +#if GGML_ROCMFP6_FAST_SIGNMAG_PACK + const int mag = (int) (code & 31u); + const int sign = -((int) ((code >> 5) & 1u)); + return (mag ^ sign) - sign; +#else + const int mag = (int) (code & 31u); + return (code & 32u) ? -mag : mag; +#endif +} + +static __device__ __forceinline__ int rocmfpx_pack4_fp6_bits24_vec_cuda(const uint32_t bits24) { + const char4 v = make_char4( + (int8_t) rocmfpx_decode_fp6_code_vec_cuda(bits24 & 63u), + (int8_t) rocmfpx_decode_fp6_code_vec_cuda((bits24 >> 6) & 63u), + (int8_t) rocmfpx_decode_fp6_code_vec_cuda((bits24 >> 12) & 63u), + (int8_t) rocmfpx_decode_fp6_code_vec_cuda((bits24 >> 18) & 63u)); + return *((const int *) &v); +} + +static __device__ __forceinline__ int rocmfpx_pack4_fp3_vec_cuda(const uint8_t * qs, const int base) { + const char4 v = make_char4( + (int8_t) rocmfpx_decode_fp3_code_vec_cuda(rocmfpx_get_bits_vec_cuda(qs, (base + 0)*3, 3)), + (int8_t) rocmfpx_decode_fp3_code_vec_cuda(rocmfpx_get_bits_vec_cuda(qs, (base + 1)*3, 3)), + (int8_t) rocmfpx_decode_fp3_code_vec_cuda(rocmfpx_get_bits_vec_cuda(qs, (base + 2)*3, 3)), + (int8_t) rocmfpx_decode_fp3_code_vec_cuda(rocmfpx_get_bits_vec_cuda(qs, (base + 3)*3, 3))); + return *((const int *) &v); +} + +static __device__ __forceinline__ int rocmfpx_pack4_fp2_bits8_vec_cuda(const uint32_t bits8) { +#if defined(GGML_USE_HIP) + const uint32_t values = + ((uint32_t) (uint8_t) (int8_t) ROCMFP2_KVALUE_0_I8) | + ((uint32_t) (uint8_t) (int8_t) ROCMFP2_KVALUE_1_I8 << 8) | + ((uint32_t) (uint8_t) (int8_t) ROCMFP2_KVALUE_2_I8 << 16) | + ((uint32_t) (uint8_t) (int8_t) ROCMFP2_KVALUE_3_I8 << 24); + const uint32_t selectors = + ((bits8 >> 0) & 3u) | + (((bits8 >> 2) & 3u) << 8) | + (((bits8 >> 4) & 3u) << 16) | + (((bits8 >> 6) & 3u) << 24); + return (int) __builtin_amdgcn_perm(0, values, selectors); +#else + const char4 v = make_char4( + (int8_t) rocmfpx_decode_fp2_code_vec_cuda((bits8 >> 0) & 3u), + (int8_t) rocmfpx_decode_fp2_code_vec_cuda((bits8 >> 2) & 3u), + (int8_t) rocmfpx_decode_fp2_code_vec_cuda((bits8 >> 4) & 3u), + (int8_t) rocmfpx_decode_fp2_code_vec_cuda((bits8 >> 6) & 3u)); + return *((const int *) &v); +#endif +} + +static __device__ __forceinline__ int rocmfpx_pack4_fp2_vec_cuda(const uint8_t * qs, const int base) { + return rocmfpx_pack4_fp2_bits8_vec_cuda(qs[base >> 2]); +} + +static __device__ __forceinline__ int rocmfpx_pack4_fp6_vec_cuda(const uint8_t * qs, const int base) { +#if GGML_ROCMFP6_FAST_SIGNMAG_PACK + uint32_t qs0, qs1, qs2, qs3, qs4, qs5; + memcpy(&qs0, qs + 0, 4); + memcpy(&qs1, qs + 4, 4); + memcpy(&qs2, qs + 8, 4); + memcpy(&qs3, qs + 12, 4); + memcpy(&qs4, qs + 16, 4); + memcpy(&qs5, qs + 20, 4); + + const uint32_t words[7] = { qs0, qs1, qs2, qs3, qs4, qs5, 0 }; + const int start_bit = 6 * base; + const int reg_idx = start_bit >> 5; + const int reg_shift = start_bit & 31; + const uint32_t val_low = words[reg_idx]; + const uint32_t val_high = words[reg_idx + 1]; + const uint32_t bits24 = (reg_shift == 0) ? (val_low & 0xFFFFFFu) : + (((val_low >> reg_shift) | (val_high << (32 - reg_shift))) & 0xFFFFFFu); + + return rocmfpx_pack4_fp6_bits24_vec_cuda(bits24); +#else + const char4 v = make_char4( + (int8_t) rocmfpx_decode_fp6_code_vec_cuda(rocmfpx_get_bits_vec_cuda(qs, (base + 0)*6, 6)), + (int8_t) rocmfpx_decode_fp6_code_vec_cuda(rocmfpx_get_bits_vec_cuda(qs, (base + 1)*6, 6)), + (int8_t) rocmfpx_decode_fp6_code_vec_cuda(rocmfpx_get_bits_vec_cuda(qs, (base + 2)*6, 6)), + (int8_t) rocmfpx_decode_fp6_code_vec_cuda(rocmfpx_get_bits_vec_cuda(qs, (base + 3)*6, 6))); + return *((const int *) &v); +#endif +} + +static __device__ __forceinline__ int rocmfpx_pack4_fp6_expanded_vec_cuda(const int8_t * qs, const int base) { + const char4 v = make_char4(qs[base + 0], qs[base + 1], qs[base + 2], qs[base + 3]); + return *((const int *) &v); +} + +static __device__ __forceinline__ int rocmfpx_pack4_fp6_device_vec_cuda(const block_rocmfp6_device * bq6, const int base) { +#if GGML_ROCMFP6_EXPANDED_DEVICE + return rocmfpx_pack4_fp6_expanded_vec_cuda(bq6->qs, base); +#else + return rocmfpx_pack4_fp6_vec_cuda(bq6->qs, base); +#endif +} + +static __device__ __forceinline__ float vec_dot_rocmfp4_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_rocmfp4 * bq4 = (const block_rocmfp4 *) vbq + kbx; + + const int * q8 = (const int *) bq8_1->qs + iqs; + + int sumi0 = 0; + int sumi1 = 0; +#pragma unroll + for (int l = 0; l < VDR_ROCMFP4_Q8_1_MMVQ; ++l) { + const int aux_q4 = rocmfp4_get_qs_i32(bq4->qs, iqs + l); + const int2 v = rocmfp4_get_int_from_codebook_16(aux_q4, kvalues_rocmfp4); + + sumi0 = ggml_cuda_dp4a(v.x, q8[l + 0], sumi0); + sumi1 = ggml_cuda_dp4a(v.y, q8[l + 4], sumi1); + } + + const float db = __low2float(bq8_1->ds); + return db * (rocmfp4_ue4m3_to_fp32_half_finite(bq4->e[0]) * sumi0 + rocmfp4_ue4m3_to_fp32_half_finite(bq4->e[1]) * sumi1); +} + +static __device__ __forceinline__ float vec_dot_rocmfp4_fast_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_rocmfp4_fast * bq4 = (const block_rocmfp4_fast *) vbq + kbx; + + const int * q8 = (const int *) bq8_1->qs + iqs; + + int sumi = 0; +#pragma unroll + for (int l = 0; l < VDR_ROCMFP4_FAST_Q8_1_MMVQ; ++l) { + const int aux_q4 = rocmfp4_get_qs_i32(bq4->qs, iqs + l); + const int2 v = rocmfp4_get_int_from_codebook_16(aux_q4, kvalues_rocmfp4); + + sumi = ggml_cuda_dp4a(v.x, q8[l + 0], sumi); + sumi = ggml_cuda_dp4a(v.y, q8[l + 4], sumi); + } + + return __low2float(bq8_1->ds) * rocmfp4_ue4m3_to_fp32_half_finite(bq4->e) * sumi; +} + +static __device__ __forceinline__ float vec_dot_rocmfpx_fp2_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + // int32-consistent MMVQ layout: QI_ROCMFP2=2 (8B qs = 2 int32), VDR=1. + // iqs in {0,1} selects one 16-weight half-block: qs bytes [4*iqs .. 4*iqs+3] + scale e[iqs]. + // qs is a 10-byte-strided, byte-aligned array -> read byte-wise (uint8_t index); NEVER cast qs to int*/int2*. + const block_rocmfp2 * bq2 = (const block_rocmfp2 *) vbq + kbx; + + int sumi = 0; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const int val_packed = rocmfpx_pack4_fp2_bits8_vec_cuda((uint32_t) bq2->qs[4*iqs + j]); + const int u = get_int_b4(bq8_1->qs, 4*iqs + j); + sumi = ggml_cuda_dp4a(val_packed, u, sumi); + } + + const float db = __low2float(bq8_1->ds); + return db * rocmfpx_ue4m3_to_fp32_finite(bq2->e[iqs]) * sumi; +} + +static __device__ __forceinline__ float vec_dot_rocmfpx_fp3_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_rocmfp3 * bq3 = (const block_rocmfp3 *) vbq + kbx; + + uint32_t qs0, qs1, qs2; + memcpy(&qs0, bq3->qs + 0, 4); + memcpy(&qs1, bq3->qs + 4, 4); + memcpy(&qs2, bq3->qs + 8, 4); + + const uint32_t qs[4] = { qs0, qs1, qs2, 0 }; + + int sumi0 = 0; + int sumi1 = 0; + + // The two half-block scales (e[0]/e[1]) split at element QK_ROCMFP3/2. base + // < QK_ROCMFP3/2 is equivalent to (iqs+i) < QK_ROCMFP3/8, so for a VDR + // window that lies entirely in one half the accumulator choice is loop + // invariant and can be hoisted out of the unrolled loop. A straddling window + // (only possible for VDRs that cross the midpoint) still uses the exact + // per-element branch, so results are bit-identical either way. + const bool fp3_first_half = iqs + VDR_ROCMFP3_Q8_1_MMVQ <= QK_ROCMFP3/8; + const bool fp3_second_half = iqs >= QK_ROCMFP3/8; + +#pragma unroll + for (int i = 0; i < VDR_ROCMFP3_Q8_1_MMVQ; ++i) { + const int base = 4 * (iqs + i); + const int start_bit = 12 * (iqs + i); + const int reg_idx = start_bit >> 5; + const int reg_shift = start_bit & 31; + const uint32_t val_low = qs[reg_idx]; + const uint32_t val_high = qs[reg_idx + 1]; + const uint32_t bits12 = (reg_shift == 0) ? (val_low & 0xFFFu) : (((val_low >> reg_shift) | (val_high << (32 - reg_shift))) & 0xFFFu); + + const char4 v = make_char4( + (int8_t) rocmfpx_decode_fp3_code_vec_cuda(bits12 & 7u), + (int8_t) rocmfpx_decode_fp3_code_vec_cuda((bits12 >> 3) & 7u), + (int8_t) rocmfpx_decode_fp3_code_vec_cuda((bits12 >> 6) & 7u), + (int8_t) rocmfpx_decode_fp3_code_vec_cuda((bits12 >> 9) & 7u)); + const int val_packed = *((const int *) &v); + + const int u = get_int_b4(bq8_1->qs, iqs + i); + + if (fp3_first_half) { + sumi0 = ggml_cuda_dp4a(val_packed, u, sumi0); + } else if (fp3_second_half) { + sumi1 = ggml_cuda_dp4a(val_packed, u, sumi1); + } else if (base < QK_ROCMFP3/2) { + sumi0 = ggml_cuda_dp4a(val_packed, u, sumi0); + } else { + sumi1 = ggml_cuda_dp4a(val_packed, u, sumi1); + } + } + + const float db = __low2float(bq8_1->ds); + return db * (rocmfpx_ue4m3_to_fp32_finite(bq3->e[0]) * sumi0 + rocmfpx_ue4m3_to_fp32_finite(bq3->e[1]) * sumi1); +} + +static __device__ __forceinline__ float vec_dot_rocmfpx_fp6_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_rocmfp6_device * bq6 = (const block_rocmfp6_device *) vbq + kbx; + +#if !GGML_ROCMFP6_EXPANDED_DEVICE + uint32_t qs0, qs1, qs2, qs3, qs4, qs5; + memcpy(&qs0, bq6->qs + 0, 4); + memcpy(&qs1, bq6->qs + 4, 4); + memcpy(&qs2, bq6->qs + 8, 4); + memcpy(&qs3, bq6->qs + 12, 4); + memcpy(&qs4, bq6->qs + 16, 4); + memcpy(&qs5, bq6->qs + 20, 4); + + const uint32_t qs[6] = { qs0, qs1, qs2, qs3, qs4, qs5 }; +#endif + + int sumi0 = 0; + int sumi1 = 0; + +#if GGML_ROCMFP6_MMVQ_HALF_BLOCK_SPLIT + const bool fp6_first_half = iqs + VDR_ROCMFP6_Q8_1_MMVQ <= QK_ROCMFP6/8; + const bool fp6_second_half = iqs >= QK_ROCMFP6/8; +#endif + +#pragma unroll + for (int i = 0; i < VDR_ROCMFP6_Q8_1_MMVQ; ++i) { + const int base = 4 * (iqs + i); +#if GGML_ROCMFP6_EXPANDED_DEVICE + const int val_packed = rocmfpx_pack4_fp6_device_vec_cuda(bq6, base); +#else + const int start_bit = 6 * base; + const int reg_idx = start_bit >> 5; + const int reg_shift = start_bit & 31; + const uint32_t val_low = qs[reg_idx]; + const uint32_t val_high = qs[reg_idx + 1]; + const uint32_t bits24 = (reg_shift == 0) ? (val_low & 0xFFFFFFu) : + (((val_low >> reg_shift) | (val_high << (32 - reg_shift))) & 0xFFFFFFu); + + const int val_packed = rocmfpx_pack4_fp6_bits24_vec_cuda(bits24); +#endif + const int u = get_int_b4(bq8_1->qs, iqs + i); + +#if GGML_ROCMFP6_MMVQ_HALF_BLOCK_SPLIT + if (fp6_first_half) { + sumi0 = ggml_cuda_dp4a(val_packed, u, sumi0); + } else if (fp6_second_half) { + sumi1 = ggml_cuda_dp4a(val_packed, u, sumi1); + } else +#endif + if (base < QK_ROCMFP6/2) { + sumi0 = ggml_cuda_dp4a(val_packed, u, sumi0); + } else { + sumi1 = ggml_cuda_dp4a(val_packed, u, sumi1); + } + } + + const float db = __low2float(bq8_1->ds); + return db * (rocmfpx_ue4m3_to_fp32_finite(bq6->e[0]) * sumi0 + rocmfpx_ue4m3_to_fp32_finite(bq6->e[1]) * sumi1); +} + +static __device__ __forceinline__ float vec_dot_rocmfpx_fp8_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_rocmfp8 * bq8 = (const block_rocmfp8 *) vbq + kbx; + + int v[VDR_ROCMFP8_Q8_1_MMVQ]; + int u[VDR_ROCMFP8_Q8_1_MMVQ]; + +#pragma unroll + for (int i = 0; i < VDR_ROCMFP8_Q8_1_MMVQ; ++i) { + v[i] = get_int_b1(bq8->qs, iqs + i); + u[i] = get_int_b4(bq8_1->qs, iqs + i); + } + + return vec_dot_q8_0_q8_1_impl( + v, u, rocmfpx_ue4m3_to_fp32_finite(bq8->e), __low2half(bq8_1->ds)); +} + #define VDR_NVFP4_Q8_1_MMVQ 4 #define VDR_NVFP4_Q8_1_MMQ 8 diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index c1ffccafb7e3..3b5c6bcafa23 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -9,6 +9,8 @@ // FIXME: required here for quantization functions #include "ggml-quants.h" +#include "../rocmfp4/rocmfp4.h" +#include "../rocmfpx/rocmfpx.h" #ifdef GGML_USE_CPU_HBM #include @@ -719,6 +721,54 @@ static const struct ggml_type_traits type_traits[GGML_TYPE_COUNT] = { .to_float = (ggml_to_float_t) dequantize_row_q8_0, .from_float_ref = (ggml_from_float_t) quantize_row_q8_0_ref, }, + [GGML_TYPE_Q4_0_ROCMFP4] = { + .type_name = "q4_0_rocmfp4", + .blck_size = QK_ROCMFP4, + .type_size = sizeof(block_rocmfp4), + .is_quantized = true, + .to_float = (ggml_to_float_t) rocmfp4_dequantize_row_q4_0, + .from_float_ref = (ggml_from_float_t) rocmfp4_quantize_row_q4_0_ref, + }, + [GGML_TYPE_Q4_0_ROCMFP4_FAST] = { + .type_name = "q4_0_rocmfp4_fast", + .blck_size = QK_ROCMFP4, + .type_size = sizeof(block_rocmfp4_fast), + .is_quantized = true, + .to_float = (ggml_to_float_t) rocmfp4_dequantize_row_q4_0_fast, + .from_float_ref = (ggml_from_float_t) rocmfp4_quantize_row_q4_0_fast_ref, + }, + [GGML_TYPE_Q3_0_ROCMFPX] = { + .type_name = "q3_0_rocmfpx", + .blck_size = QK_ROCMFP3, + .type_size = sizeof(block_rocmfp3), + .is_quantized = true, + .to_float = (ggml_to_float_t) rocmfpx_dequantize_row_fp3, + .from_float_ref = (ggml_from_float_t) rocmfpx_quantize_row_fp3_ref, + }, + [GGML_TYPE_Q2_0_ROCMFP2] = { + .type_name = "q2_0_rocmfp2", + .blck_size = QK_ROCMFP2, + .type_size = sizeof(block_rocmfp2), + .is_quantized = true, + .to_float = (ggml_to_float_t) rocmfpx_dequantize_row_fp2, + .from_float_ref = (ggml_from_float_t) rocmfpx_quantize_row_fp2_ref, + }, + [GGML_TYPE_Q6_0_ROCMFPX] = { + .type_name = "q6_0_rocmfpx", + .blck_size = QK_ROCMFP6, + .type_size = sizeof(block_rocmfp6), + .is_quantized = true, + .to_float = (ggml_to_float_t) rocmfpx_dequantize_row_fp6, + .from_float_ref = (ggml_from_float_t) rocmfpx_quantize_row_fp6_ref, + }, + [GGML_TYPE_Q8_0_ROCMFPX] = { + .type_name = "q8_0_rocmfpx", + .blck_size = QK_ROCMFP8, + .type_size = sizeof(block_rocmfp8), + .is_quantized = true, + .to_float = (ggml_to_float_t) rocmfpx_dequantize_row_fp8, + .from_float_ref = (ggml_from_float_t) rocmfpx_quantize_row_fp8_ref, + }, [GGML_TYPE_Q8_1] = { .type_name = "q8_1", .blck_size = QK8_1, @@ -1225,11 +1275,12 @@ static const char * GGML_GLU_OP_NAME[GGML_GLU_OP_COUNT] = { "GEGLU", "SWIGLU", "SWIGLU_OAI", + "SWIGLU_DS4", "GEGLU_ERF", "GEGLU_QUICK", }; -static_assert(GGML_GLU_OP_COUNT == 6, "GGML_GLU_OP_COUNT != 6"); +static_assert(GGML_GLU_OP_COUNT == 7, "GGML_GLU_OP_COUNT != 7"); static_assert(sizeof(struct ggml_object)%GGML_MEM_ALIGN == 0, "ggml_object size must be a multiple of GGML_MEM_ALIGN"); @@ -3025,6 +3076,16 @@ struct ggml_tensor * ggml_swiglu_split( return ggml_glu_impl(ctx, a, b, GGML_GLU_OP_SWIGLU, false); } +struct ggml_tensor * ggml_swiglu_ds4_split( + struct ggml_context * ctx, + struct ggml_tensor * gate, + struct ggml_tensor * up, + float clamp) { + struct ggml_tensor * result = ggml_glu_impl(ctx, gate, up, GGML_GLU_OP_SWIGLU_DS4, false); + ggml_set_op_params_f32(result, 2, clamp); + return result; +} + // ggml_geglu_erf struct ggml_tensor * ggml_geglu_erf( From a971ec1bf7756edf01c19e47c44358cb18041176 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Tue, 7 Jul 2026 01:07:07 +0200 Subject: [PATCH 2/7] fix(hip): advertise ROCmFPX matmul support --- ggml/src/ggml-cuda/ggml-cuda.cu | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 9d1c870879a4..d45913164daf 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -5009,6 +5009,12 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g case GGML_TYPE_Q8_0: case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: + case GGML_TYPE_Q4_0_ROCMFP4: + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + case GGML_TYPE_Q2_0_ROCMFP2: + case GGML_TYPE_Q3_0_ROCMFPX: + case GGML_TYPE_Q6_0_ROCMFPX: + case GGML_TYPE_Q8_0_ROCMFPX: case GGML_TYPE_Q2_K: case GGML_TYPE_Q3_K: case GGML_TYPE_Q4_K: @@ -5044,6 +5050,12 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g case GGML_TYPE_Q5_0: case GGML_TYPE_Q5_1: case GGML_TYPE_Q8_0: + case GGML_TYPE_Q4_0_ROCMFP4: + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + case GGML_TYPE_Q2_0_ROCMFP2: + case GGML_TYPE_Q3_0_ROCMFPX: + case GGML_TYPE_Q6_0_ROCMFPX: + case GGML_TYPE_Q8_0_ROCMFPX: case GGML_TYPE_TQ3_0: return true; default: From 7f54096110f01cb6301391fd31f89ba48058ab64 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Tue, 7 Jul 2026 02:39:42 +0200 Subject: [PATCH 3/7] fix(hip): preserve DS4 SwiGLU fusion rounding order --- ggml/src/ggml-cuda/unary.cuh | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/ggml/src/ggml-cuda/unary.cuh b/ggml/src/ggml-cuda/unary.cuh index 854b27e35097..1864ce2dc432 100644 --- a/ggml/src/ggml-cuda/unary.cuh +++ b/ggml/src/ggml-cuda/unary.cuh @@ -117,5 +117,6 @@ __device__ __forceinline__ float ggml_cuda_op_swiglu_ds4_single(float gate, floa gate = fminf(gate, limit); up = fmaxf(fminf(up, limit), -limit); - return up * gate / (1.0f + expf(-gate)); + const float silu = gate / (1.0f + expf(-gate)); + return silu * up; } From 113d8812025243ebf6d3370cf4ce22cc4646faa0 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Tue, 7 Jul 2026 10:44:48 +0200 Subject: [PATCH 4/7] fix(hip): support f16 concat kernels --- ggml/src/ggml-cuda/concat.cu | 97 ++++++++++++++++++++------------- ggml/src/ggml-cuda/ggml-cuda.cu | 5 +- 2 files changed, 63 insertions(+), 39 deletions(-) diff --git a/ggml/src/ggml-cuda/concat.cu b/ggml/src/ggml-cuda/concat.cu index e9ffd274b996..51a3efa53694 100644 --- a/ggml/src/ggml-cuda/concat.cu +++ b/ggml/src/ggml-cuda/concat.cu @@ -1,7 +1,8 @@ #include "concat.cuh" // contiguous kernels -static __global__ void concat_f32_dim0(const float * x, const float * y, float * dst, const int ne0, const int ne00) { +template +static __global__ void concat_dim0(const T * x, const T * y, T * dst, const int ne0, const int ne00) { int nidx = threadIdx.x + blockIdx.x * blockDim.x; if (nidx >= ne0) { return; @@ -27,7 +28,8 @@ static __global__ void concat_f32_dim0(const float * x, const float * y, float * } } -static __global__ void concat_f32_dim1(const float * x, const float * y, float * dst, const int ne0, const int ne01) { +template +static __global__ void concat_dim1(const T * x, const T * y, T * dst, const int ne0, const int ne01) { int nidx = threadIdx.x + blockIdx.x * blockDim.x; if (nidx >= ne0) { return; @@ -53,7 +55,8 @@ static __global__ void concat_f32_dim1(const float * x, const float * y, float * } } -static __global__ void concat_f32_dim2(const float * x, const float * y, float * dst, const int ne0, const int ne02) { +template +static __global__ void concat_dim2(const T * x, const T * y, T * dst, const int ne0, const int ne02) { int nidx = threadIdx.x + blockIdx.x * blockDim.x; if (nidx >= ne0) { return; @@ -79,24 +82,25 @@ static __global__ void concat_f32_dim2(const float * x, const float * y, float * } } -static void concat_f32_cuda(const float * x, const float * y, float * dst, int ne00, int ne01, int ne02, int ne0, int ne1, int ne2, int dim, cudaStream_t stream) { +template +static void concat_cuda(const T * x, const T * y, T * dst, int ne00, int ne01, int ne02, int ne0, int ne1, int ne2, int dim, cudaStream_t stream) { int num_blocks = (ne0 + CUDA_CONCAT_BLOCK_SIZE - 1) / CUDA_CONCAT_BLOCK_SIZE; dim3 gridDim(num_blocks, ne1, ne2); if (dim == 0) { - concat_f32_dim0<<>>(x, y, dst, ne0, ne00); + concat_dim0<<>>(x, y, dst, ne0, ne00); return; } if (dim == 1) { - concat_f32_dim1<<>>(x, y, dst, ne0, ne01); + concat_dim1<<>>(x, y, dst, ne0, ne01); return; } - concat_f32_dim2<<>>(x, y, dst, ne0, ne02); + concat_dim2<<>>(x, y, dst, ne0, ne02); } // non-contiguous kernel (slow) -template +template static __global__ void __launch_bounds__(CUDA_CONCAT_BLOCK_SIZE) - concat_f32_non_cont( + concat_non_cont( const char * src0, const char * src1, char * dst, @@ -130,54 +134,44 @@ static __global__ void __launch_bounds__(CUDA_CONCAT_BLOCK_SIZE) const int64_t i2 = blockIdx.y; const int64_t i1 = blockIdx.x; - const float * x; + const T * x; for (int64_t i0 = threadIdx.x; i0 < ne0; i0 += blockDim.x) { if (i0 < ne00 && i1 < ne01 && i2 < ne02 && i3 < ne03) { - x = (const float *)(src0 + (i3 )*nb03 + (i2 )*nb02 + (i1 )*nb01 + (i0 )*nb00); + x = (const T *)(src0 + (i3 )*nb03 + (i2 )*nb02 + (i1 )*nb01 + (i0 )*nb00); } else { if constexpr (dim == 0) { - x = (const float *) (src1 + i3 * nb13 + i2 * nb12 + i1 * nb11 + (i0 - ne00) * nb10); + x = (const T *) (src1 + i3 * nb13 + i2 * nb12 + i1 * nb11 + (i0 - ne00) * nb10); } else if constexpr (dim == 1) { - x = (const float *) (src1 + i3 * nb13 + i2 * nb12 + (i1 - ne01) * nb11 + i0 * nb10); + x = (const T *) (src1 + i3 * nb13 + i2 * nb12 + (i1 - ne01) * nb11 + i0 * nb10); } else if constexpr (dim == 2) { - x = (const float *) (src1 + i3 * nb13 + (i2 - ne02) * nb12 + i1 * nb11 + i0 * nb10); + x = (const T *) (src1 + i3 * nb13 + (i2 - ne02) * nb12 + i1 * nb11 + i0 * nb10); } else if constexpr (dim == 3) { - x = (const float *) (src1 + (i3 - ne03) * nb13 + i2 * nb12 + i1 * nb11 + i0 * nb10); + x = (const T *) (src1 + (i3 - ne03) * nb13 + i2 * nb12 + i1 * nb11 + i0 * nb10); } } - float * y = (float *)(dst + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0); + T * y = (T *)(dst + i3*nb3 + i2*nb2 + i1*nb1 + i0*nb0); *y = *x; } } - -void ggml_cuda_op_concat(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { - const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; - +template +static void concat_cuda_typed(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, int dim) { cudaStream_t stream = ctx.stream(); - const int32_t dim = ((int32_t *) dst->op_params)[0]; - - GGML_ASSERT(src0->type == GGML_TYPE_F32); - GGML_ASSERT(src1->type == GGML_TYPE_F32); - GGML_ASSERT(dst->type == GGML_TYPE_F32); - if (ggml_is_contiguous(src0) && ggml_is_contiguous(src1)) { - const float * src0_d = (const float *)src0->data; - const float * src1_d = (const float *)src1->data; - - float * dst_d = (float *)dst->data; + const T * src0_d = (const T *)src0->data; + const T * src1_d = (const T *)src1->data; + T * dst_d = (T *)dst->data; if (dim != 3) { for (int i3 = 0; i3 < dst->ne[3]; i3++) { - concat_f32_cuda( - src0_d + i3 * (src0->nb[3] / 4), - src1_d + i3 * (src1->nb[3] / 4), - dst_d + i3 * ( dst->nb[3] / 4), + concat_cuda( + src0_d + i3 * (src0->nb[3] / sizeof(T)), + src1_d + i3 * (src1->nb[3] / sizeof(T)), + dst_d + i3 * ( dst->nb[3] / sizeof(T)), src0->ne[0], src0->ne[1], src0->ne[2], dst->ne[0], dst->ne[1], dst->ne[2], dim, stream); } @@ -185,13 +179,13 @@ void ggml_cuda_op_concat(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const size_t size0 = ggml_nbytes(src0); const size_t size1 = ggml_nbytes(src1); - CUDA_CHECK(cudaMemcpyAsync(dst_d, src0_d, size0, cudaMemcpyDeviceToDevice, stream)); - CUDA_CHECK(cudaMemcpyAsync(dst_d + size0/4, src1_d, size1, cudaMemcpyDeviceToDevice, stream)); + CUDA_CHECK(cudaMemcpyAsync(dst_d, src0_d, size0, cudaMemcpyDeviceToDevice, stream)); + CUDA_CHECK(cudaMemcpyAsync((char *)dst_d + size0, src1_d, size1, cudaMemcpyDeviceToDevice, stream)); } } else { dim3 grid_dim(dst->ne[1], dst->ne[2], dst->ne[3]); auto launch_kernel = [&](auto dim) { - concat_f32_non_cont<<>>( + concat_non_cont<<>>( (const char *) src0->data, (const char *) src1->data, (char *) dst->data, src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3], src0->nb[0], src0->nb[1], src0->nb[2], src0->nb[3], @@ -219,3 +213,30 @@ void ggml_cuda_op_concat(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { } } } + +void ggml_cuda_op_concat(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + + const int32_t dim = ((int32_t *) dst->op_params)[0]; + + GGML_ASSERT(src0->type == src1->type); + GGML_ASSERT(src0->type == dst->type); + + switch (src0->type) { + case GGML_TYPE_F32: + concat_cuda_typed(ctx, src0, src1, dst, dim); + break; + case GGML_TYPE_F16: + concat_cuda_typed(ctx, src0, src1, dst, dim); + break; + case GGML_TYPE_BF16: + concat_cuda_typed(ctx, src0, src1, dst, dim); + break; + case GGML_TYPE_I8: + concat_cuda_typed(ctx, src0, src1, dst, dim); + break; + default: + GGML_ABORT("unsupported concat type %s", ggml_type_name(src0->type)); + } +} diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index d45913164daf..066dd2680d89 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -5162,7 +5162,10 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g case GGML_OP_CONCAT: { ggml_type src0_type = op->src[0]->type; - return src0_type != GGML_TYPE_I32 && src0_type != GGML_TYPE_I16; + return src0_type == GGML_TYPE_F32 || + src0_type == GGML_TYPE_F16 || + src0_type == GGML_TYPE_BF16 || + src0_type == GGML_TYPE_I8; } break; case GGML_OP_CONV_TRANSPOSE_1D: { From 7c3a6cbacdce06c958bf29a151ea36cbb2558d44 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Wed, 8 Jul 2026 00:27:21 +0200 Subject: [PATCH 5/7] feat(cuda): add fused DS4 hyper-connection op (GGML_OP_DS4_HC) Adds ds4-hc kernels (pre/post/out modes) used by the opt-in DFLASH_DS4_FUSED_DECODE single-graph decode path on DeepSeek V4 Flash. The fn/fn_out matvecs stay on ggml mul_mat; the op fuses sigmoid gates, sinkhorn normalization and stream mixing. CPU fallback included; the op is inert unless emitted by the server. --- ggml/include/ggml.h | 33 + ggml/include/ggml.h.ds4bak | 2903 +++++++++ ggml/src/ggml-cpu/ggml-cpu.c | 4 + ggml/src/ggml-cpu/ggml-cpu.c.ds4bak | 3782 +++++++++++ ggml/src/ggml-cuda/ds4-hc.cu | 219 + ggml/src/ggml-cuda/ds4-hc.cuh | 10 + ggml/src/ggml-cuda/ggml-cuda.cu | 6 + ggml/src/ggml-cuda/ggml-cuda.cu.ds4bak | 5559 ++++++++++++++++ ggml/src/ggml.c | 102 +- ggml/src/ggml.c.ds4bak | 8040 ++++++++++++++++++++++++ 10 files changed, 20656 insertions(+), 2 deletions(-) create mode 100644 ggml/include/ggml.h.ds4bak create mode 100644 ggml/src/ggml-cpu/ggml-cpu.c.ds4bak create mode 100644 ggml/src/ggml-cuda/ds4-hc.cu create mode 100644 ggml/src/ggml-cuda/ds4-hc.cuh create mode 100644 ggml/src/ggml-cuda/ggml-cuda.cu.ds4bak create mode 100644 ggml/src/ggml.c.ds4bak diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 800854e7d462..2fcc45ba44f5 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -601,6 +601,8 @@ extern "C" { GGML_OP_MOE_FUSED, // Fused MoE FFN: gate+up+swiglu+down+weighted_sum+shared_expert + GGML_OP_DS4_HC, // Fused DeepSeek4 hyper-connection pre/post/out mixing + GGML_OP_COUNT, }; @@ -2424,6 +2426,37 @@ extern "C" { struct ggml_tensor * experts, struct ggml_tensor * expert_weights); + // Fused DeepSeek4 hyper-connection helpers (decode, n_tokens == 1). + // ggml_ds4_hc_pre: mix[2*n_hc+n_hc^2] + base + hc_state[n_embd*n_hc] -> + // dst[n_embd + 2*n_hc + n_hc^2] = { working, split(pre,post,comb) } + GGML_API struct ggml_tensor * ggml_ds4_hc_pre( + struct ggml_context * ctx, + struct ggml_tensor * mix, + struct ggml_tensor * base, + struct ggml_tensor * hc_state, + int n_hc, + int sinkhorn_iters, + float pre_scale, + float post_scale, + float comb_scale); + + // ggml_ds4_hc_post: residual hc_state + block_out + split -> new hc_state + GGML_API struct ggml_tensor * ggml_ds4_hc_post( + struct ggml_context * ctx, + struct ggml_tensor * residual_hc, + struct ggml_tensor * block_out, + struct ggml_tensor * split, + int n_hc); + + // ggml_ds4_hc_out: output-stage merge of hc streams into one embedding + GGML_API struct ggml_tensor * ggml_ds4_hc_out( + struct ggml_context * ctx, + struct ggml_tensor * mix, + struct ggml_tensor * base, + struct ggml_tensor * hc_state, + int n_hc, + float pre_scale); + // TODO: needs to be adapted to ggml_flash_attn_ext GGML_API struct ggml_tensor * ggml_flash_attn_back( struct ggml_context * ctx, diff --git a/ggml/include/ggml.h.ds4bak b/ggml/include/ggml.h.ds4bak new file mode 100644 index 000000000000..800854e7d462 --- /dev/null +++ b/ggml/include/ggml.h.ds4bak @@ -0,0 +1,2903 @@ +#pragma once + +// +// GGML Tensor Library +// +// This documentation is still a work in progress. +// If you wish some specific topics to be covered, feel free to drop a comment: +// +// https://github.com/ggml-org/whisper.cpp/issues/40 +// +// ## Overview +// +// This library implements: +// +// - a set of tensor operations +// - automatic differentiation +// - basic optimization algorithms +// +// The aim of this library is to provide a minimalistic approach for various machine learning tasks. This includes, +// but is not limited to, the following: +// +// - linear regression +// - support vector machines +// - neural networks +// +// The library allows the user to define a certain function using the available tensor operations. This function +// definition is represented internally via a computation graph. Each tensor operation in the function definition +// corresponds to a node in the graph. Having the computation graph defined, the user can choose to compute the +// function's value and/or its gradient with respect to the input variables. Optionally, the function can be optimized +// using one of the available optimization algorithms. +// +// For example, here we define the function: f(x) = a*x^2 + b +// +// { +// struct ggml_init_params params = { +// .mem_size = 16*1024*1024, +// .mem_buffer = NULL, +// }; +// +// // memory allocation happens here +// struct ggml_context * ctx = ggml_init(params); +// +// struct ggml_tensor * x = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); +// +// ggml_set_param(ctx, x); // x is an input variable +// +// struct ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); +// struct ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); +// struct ggml_tensor * x2 = ggml_mul(ctx, x, x); +// struct ggml_tensor * f = ggml_add(ctx, ggml_mul(ctx, a, x2), b); +// +// ... +// } +// +// Notice that the function definition above does not involve any actual computation. The computation is performed only +// when the user explicitly requests it. For example, to compute the function's value at x = 2.0: +// +// { +// ... +// +// struct ggml_cgraph * gf = ggml_new_graph(ctx); +// ggml_build_forward_expand(gf, f); +// +// // set the input variable and parameter values +// ggml_set_f32(x, 2.0f); +// ggml_set_f32(a, 3.0f); +// ggml_set_f32(b, 4.0f); +// +// ggml_graph_compute_with_ctx(ctx, &gf, n_threads); +// +// printf("f = %f\n", ggml_get_f32_1d(f, 0)); +// +// ... +// } +// +// The actual computation is performed in the ggml_graph_compute() function. +// +// The ggml_new_tensor_...() functions create new tensors. They are allocated in the memory buffer provided to the +// ggml_init() function. You have to be careful not to exceed the memory buffer size. Therefore, you have to know +// in advance how much memory you need for your computation. Alternatively, you can allocate a large enough memory +// and after defining the computation graph, call the ggml_used_mem() function to find out how much memory was +// actually needed. +// +// The ggml_set_param() function marks a tensor as an input variable. This is used by the automatic +// differentiation and optimization algorithms. +// +// The described approach allows to define the function graph once and then compute its forward or backward graphs +// multiple times. All computations will use the same memory buffer allocated in the ggml_init() function. This way +// the user can avoid the memory allocation overhead at runtime. +// +// The library supports multi-dimensional tensors - up to 4 dimensions. The FP16 and FP32 data types are first class +// citizens, but in theory the library can be extended to support FP8 and integer data types. +// +// Each tensor operation produces a new tensor. Initially the library was envisioned to support only the use of unary +// and binary operations. Most of the available operations fall into one of these two categories. With time, it became +// clear that the library needs to support more complex operations. The way to support these operations is not clear +// yet, but a few examples are demonstrated in the following operations: +// +// - ggml_permute() +// - ggml_conv_1d_1s() +// - ggml_conv_1d_2s() +// +// For each tensor operator, the library implements a forward and backward computation function. The forward function +// computes the output tensor value given the input tensor values. The backward function computes the adjoint of the +// input tensors given the adjoint of the output tensor. For a detailed explanation of what this means, take a +// calculus class, or watch the following video: +// +// What is Automatic Differentiation? +// https://www.youtube.com/watch?v=wG_nF1awSSY +// +// +// ## Tensor data (struct ggml_tensor) +// +// The tensors are stored in memory via the ggml_tensor struct. The structure provides information about the size of +// the tensor, the data type, and the memory buffer where the tensor data is stored. Additionally, it contains +// pointers to the "source" tensors - i.e. the tensors that were used to compute the current tensor. For example: +// +// { +// struct ggml_tensor * c = ggml_add(ctx, a, b); +// +// assert(c->src[0] == a); +// assert(c->src[1] == b); +// } +// +// The multi-dimensional tensors are stored in row-major order. The ggml_tensor struct contains fields for the +// number of elements in each dimension ("ne") as well as the number of bytes ("nb", a.k.a. stride). This allows +// to store tensors that are not contiguous in memory, which is useful for operations such as transposition and +// permutation. All tensor operations have to take the stride into account and not assume that the tensor is +// contiguous in memory. +// +// The data of the tensor is accessed via the "data" pointer. For example: +// +// { +// const int nx = 2; +// const int ny = 3; +// +// struct ggml_tensor * a = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, nx, ny); +// +// for (int y = 0; y < ny; y++) { +// for (int x = 0; x < nx; x++) { +// *(float *) ((char *) a->data + y*a->nb[1] + x*a->nb[0]) = x + y; +// } +// } +// +// ... +// } +// +// Alternatively, there are helper functions, such as ggml_get_f32_1d() and ggml_set_f32_1d() that can be used. +// +// ## The matrix multiplication operator (ggml_mul_mat) +// +// TODO +// +// +// ## Multi-threading +// +// TODO +// +// +// ## Overview of ggml.c +// +// TODO +// +// +// ## SIMD optimizations +// +// TODO +// +// +// ## Debugging ggml +// +// TODO +// +// + +#ifdef GGML_SHARED +# if defined(_WIN32) && !defined(__MINGW32__) +# ifdef GGML_BUILD +# define GGML_API __declspec(dllexport) extern +# else +# define GGML_API __declspec(dllimport) extern +# endif +# else +# define GGML_API __attribute__ ((visibility ("default"))) extern +# endif +#else +# define GGML_API extern +#endif + +// TODO: support for clang +#ifdef __GNUC__ +# define GGML_DEPRECATED(func, hint) func __attribute__((deprecated(hint))) +#elif defined(_MSC_VER) +# define GGML_DEPRECATED(func, hint) __declspec(deprecated(hint)) func +#else +# define GGML_DEPRECATED(func, hint) func +#endif + +#ifndef __GNUC__ +# define GGML_ATTRIBUTE_FORMAT(...) +#elif defined(__MINGW32__) && !defined(__clang__) +# define GGML_ATTRIBUTE_FORMAT(...) __attribute__((format(gnu_printf, __VA_ARGS__))) +#else +# define GGML_ATTRIBUTE_FORMAT(...) __attribute__((format(printf, __VA_ARGS__))) +#endif + +#if defined(_WIN32) && !defined(_WIN32_WINNT) +# define _WIN32_WINNT 0x0A00 +#endif + +#include +#include +#include +#include + +#define GGML_FILE_MAGIC 0x67676d6c // "ggml" +#define GGML_FILE_VERSION 2 + +#define GGML_QNT_VERSION 2 // bump this on quantization format changes +#define GGML_QNT_VERSION_FACTOR 1000 // do not change this + +#define GGML_MAX_DIMS 4 +#define GGML_MAX_PARAMS 2048 +#define GGML_MAX_SRC 10 +#define GGML_MAX_N_THREADS 512 +#define GGML_MAX_OP_PARAMS 64 + +#ifndef GGML_MAX_NAME +# define GGML_MAX_NAME 64 +#endif + +#define GGML_DEFAULT_N_THREADS 4 +#define GGML_DEFAULT_GRAPH_SIZE 2048 + +#if UINTPTR_MAX == 0xFFFFFFFF + #define GGML_MEM_ALIGN 4 +#elif defined(__EMSCRIPTEN__) +// emscripten uses max_align_t == 8, so we need GGML_MEM_ALIGN == 8 for 64-bit wasm. +// (for 32-bit wasm, the first conditional is true and GGML_MEM_ALIGN stays 4.) +// ref: https://github.com/ggml-org/llama.cpp/pull/18628 + #define GGML_MEM_ALIGN 8 +#else + #define GGML_MEM_ALIGN 16 +#endif + +#define GGML_EXIT_SUCCESS 0 +#define GGML_EXIT_ABORTED 1 + +// TODO: convert to enum https://github.com/ggml-org/llama.cpp/pull/16187#discussion_r2388538726 +#define GGML_ROPE_TYPE_NORMAL 0 +#define GGML_ROPE_TYPE_NEOX 2 +#define GGML_ROPE_TYPE_MROPE 8 +#define GGML_ROPE_TYPE_VISION 24 +#define GGML_ROPE_TYPE_IMROPE 40 // binary: 101000 + +#define GGML_MROPE_SECTIONS 4 + +#define GGML_UNUSED(x) (void)(x) +#ifdef __CUDACC__ +template +__host__ __device__ constexpr inline void ggml_unused_vars_impl(Args&&...) noexcept {} +#define GGML_UNUSED_VARS(...) ggml_unused_vars_impl(__VA_ARGS__) +#else +#define GGML_UNUSED_VARS(...) do { (void)sizeof((__VA_ARGS__, 0)); } while(0) +#endif // __CUDACC__ + +#define GGML_PAD(x, n) (((x) + (n) - 1) & ~((n) - 1)) + +#ifndef NDEBUG +# define GGML_UNREACHABLE() do { fprintf(stderr, "statement should be unreachable\n"); abort(); } while(0) +#elif defined(__GNUC__) +# define GGML_UNREACHABLE() __builtin_unreachable() +#elif defined(_MSC_VER) +# define GGML_UNREACHABLE() __assume(0) +#else +# define GGML_UNREACHABLE() ((void) 0) +#endif + +#ifdef __cplusplus +# define GGML_NORETURN [[noreturn]] +#elif defined(_MSC_VER) +# define GGML_NORETURN __declspec(noreturn) +#else +# define GGML_NORETURN _Noreturn +#endif + +#define GGML_ABORT(...) ggml_abort(__FILE__, __LINE__, __VA_ARGS__) +#define GGML_ASSERT(x) if (!(x)) GGML_ABORT("GGML_ASSERT(%s) failed", #x) + +// used to copy the number of elements and stride in bytes of tensors into local variables. +// main purpose is to reduce code duplication and improve readability. +// +// example: +// +// GGML_TENSOR_LOCALS(int64_t, ne1, src1, ne); +// GGML_TENSOR_LOCALS(size_t, nb1, src1, nb); +// +#define GGML_TENSOR_LOCALS_1(type, prefix, pointer, array) \ + const type prefix##0 = (pointer) ? (pointer)->array[0] : 0; \ + GGML_UNUSED(prefix##0); +#define GGML_TENSOR_LOCALS_2(type, prefix, pointer, array) \ + GGML_TENSOR_LOCALS_1 (type, prefix, pointer, array) \ + const type prefix##1 = (pointer) ? (pointer)->array[1] : 0; \ + GGML_UNUSED(prefix##1); +#define GGML_TENSOR_LOCALS_3(type, prefix, pointer, array) \ + GGML_TENSOR_LOCALS_2 (type, prefix, pointer, array) \ + const type prefix##2 = (pointer) ? (pointer)->array[2] : 0; \ + GGML_UNUSED(prefix##2); +#define GGML_TENSOR_LOCALS(type, prefix, pointer, array) \ + GGML_TENSOR_LOCALS_3 (type, prefix, pointer, array) \ + const type prefix##3 = (pointer) ? (pointer)->array[3] : 0; \ + GGML_UNUSED(prefix##3); + +#define GGML_TENSOR_UNARY_OP_LOCALS \ + GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) \ + GGML_TENSOR_LOCALS(size_t, nb0, src0, nb) \ + GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) \ + GGML_TENSOR_LOCALS(size_t, nb, dst, nb) + +#define GGML_TENSOR_BINARY_OP_LOCALS \ + GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) \ + GGML_TENSOR_LOCALS(size_t, nb0, src0, nb) \ + GGML_TENSOR_LOCALS(int64_t, ne1, src1, ne) \ + GGML_TENSOR_LOCALS(size_t, nb1, src1, nb) \ + GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) \ + GGML_TENSOR_LOCALS(size_t, nb, dst, nb) + +#define GGML_TENSOR_TERNARY_OP_LOCALS \ + GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) \ + GGML_TENSOR_LOCALS(size_t, nb0, src0, nb) \ + GGML_TENSOR_LOCALS(int64_t, ne1, src1, ne) \ + GGML_TENSOR_LOCALS(size_t, nb1, src1, nb) \ + GGML_TENSOR_LOCALS(int64_t, ne2, src2, ne) \ + GGML_TENSOR_LOCALS(size_t, nb2, src2, nb) \ + GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) \ + GGML_TENSOR_LOCALS(size_t, nb, dst, nb) + +#define GGML_TENSOR_BINARY_OP_LOCALS01 \ + GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) \ + GGML_TENSOR_LOCALS(size_t, nb0, src0, nb) \ + GGML_TENSOR_LOCALS(int64_t, ne1, src1, ne) \ + GGML_TENSOR_LOCALS(size_t, nb1, src1, nb) + +#ifdef __cplusplus +extern "C" { +#endif + + // Function type used in fatal error callbacks + typedef void (*ggml_abort_callback_t)(const char * error_message); + + // Set the abort callback (passing null will restore original abort functionality: printing a message to stdout) + // Returns the old callback for chaining + GGML_API ggml_abort_callback_t ggml_set_abort_callback(ggml_abort_callback_t callback); + + GGML_NORETURN GGML_ATTRIBUTE_FORMAT(3, 4) + GGML_API void ggml_abort(const char * file, int line, const char * fmt, ...); + + enum ggml_status { + GGML_STATUS_ALLOC_FAILED = -2, + GGML_STATUS_FAILED = -1, + GGML_STATUS_SUCCESS = 0, + GGML_STATUS_ABORTED = 1, + }; + + // get ggml_status name string + GGML_API const char * ggml_status_to_string(enum ggml_status status); + + // ieee 754-2008 half-precision float16 + // todo: make this not an integral type + typedef uint16_t ggml_fp16_t; + GGML_API float ggml_fp16_to_fp32(ggml_fp16_t); + GGML_API ggml_fp16_t ggml_fp32_to_fp16(float); + GGML_API void ggml_fp16_to_fp32_row(const ggml_fp16_t *, float *, int64_t); + GGML_API void ggml_fp32_to_fp16_row(const float *, ggml_fp16_t *, int64_t); + + // google brain half-precision bfloat16 + typedef struct { uint16_t bits; } ggml_bf16_t; + GGML_API ggml_bf16_t ggml_fp32_to_bf16(float); + GGML_API float ggml_bf16_to_fp32(ggml_bf16_t); // consider just doing << 16 + GGML_API void ggml_bf16_to_fp32_row(const ggml_bf16_t *, float *, int64_t); + GGML_API void ggml_fp32_to_bf16_row_ref(const float *, ggml_bf16_t *, int64_t); + GGML_API void ggml_fp32_to_bf16_row(const float *, ggml_bf16_t *, int64_t); + + struct ggml_object; + struct ggml_context; + struct ggml_cgraph; + + // NOTE: always add types at the end of the enum to keep backward compatibility + enum ggml_type { + GGML_TYPE_F32 = 0, + GGML_TYPE_F16 = 1, + GGML_TYPE_Q4_0 = 2, + GGML_TYPE_Q4_1 = 3, + // GGML_TYPE_Q4_2 = 4, support has been removed + // GGML_TYPE_Q4_3 = 5, support has been removed + GGML_TYPE_Q5_0 = 6, + GGML_TYPE_Q5_1 = 7, + GGML_TYPE_Q8_0 = 8, + GGML_TYPE_Q8_1 = 9, + GGML_TYPE_Q2_K = 10, + GGML_TYPE_Q3_K = 11, + GGML_TYPE_Q4_K = 12, + GGML_TYPE_Q5_K = 13, + GGML_TYPE_Q6_K = 14, + GGML_TYPE_Q8_K = 15, + GGML_TYPE_IQ2_XXS = 16, + GGML_TYPE_IQ2_XS = 17, + GGML_TYPE_IQ3_XXS = 18, + GGML_TYPE_IQ1_S = 19, + GGML_TYPE_IQ4_NL = 20, + GGML_TYPE_IQ3_S = 21, + GGML_TYPE_IQ2_S = 22, + GGML_TYPE_IQ4_XS = 23, + GGML_TYPE_I8 = 24, + GGML_TYPE_I16 = 25, + GGML_TYPE_I32 = 26, + GGML_TYPE_I64 = 27, + GGML_TYPE_F64 = 28, + GGML_TYPE_IQ1_M = 29, + GGML_TYPE_BF16 = 30, + // GGML_TYPE_Q4_0_4_4 = 31, support has been removed from gguf files + // GGML_TYPE_Q4_0_4_8 = 32, + // GGML_TYPE_Q4_0_8_8 = 33, + GGML_TYPE_TQ1_0 = 34, + GGML_TYPE_TQ2_0 = 35, + // GGML_TYPE_IQ4_NL_4_4 = 36, + // GGML_TYPE_IQ4_NL_4_8 = 37, + // GGML_TYPE_IQ4_NL_8_8 = 38, + GGML_TYPE_MXFP4 = 39, // MXFP4 (1 block) + GGML_TYPE_NVFP4 = 40, // NVFP4 (4 blocks, E4M3 scale) + GGML_TYPE_Q1_0 = 41, + GGML_TYPE_TQ3_0 = 42, // TurboQuant 3.5 bpv (3-bit Lloyd-Max + FWHT rotation) + GGML_TYPE_Q4_0_ROCMFP4 = 100, + GGML_TYPE_Q4_0_ROCMFP4_FAST = 101, + GGML_TYPE_Q6_0_ROCMFPX = 102, + GGML_TYPE_Q8_0_ROCMFPX = 103, + GGML_TYPE_Q3_0_ROCMFPX = 104, + GGML_TYPE_Q2_0_ROCMFP2 = 107, + GGML_TYPE_COUNT = 108, + }; + + // precision + enum ggml_prec { + GGML_PREC_DEFAULT = 0, // stored as ggml_tensor.op_params, 0 by default + GGML_PREC_F32 = 10, + }; + + // model file types + enum ggml_ftype { + GGML_FTYPE_UNKNOWN = -1, + GGML_FTYPE_ALL_F32 = 0, + GGML_FTYPE_MOSTLY_F16 = 1, // except 1d tensors + GGML_FTYPE_MOSTLY_Q4_0 = 2, // except 1d tensors + GGML_FTYPE_MOSTLY_Q4_1 = 3, // except 1d tensors + GGML_FTYPE_MOSTLY_Q4_1_SOME_F16 = 4, // tok_embeddings.weight and output.weight are F16 + GGML_FTYPE_MOSTLY_Q8_0 = 7, // except 1d tensors + GGML_FTYPE_MOSTLY_Q5_0 = 8, // except 1d tensors + GGML_FTYPE_MOSTLY_Q5_1 = 9, // except 1d tensors + GGML_FTYPE_MOSTLY_Q2_K = 10, // except 1d tensors + GGML_FTYPE_MOSTLY_Q3_K = 11, // except 1d tensors + GGML_FTYPE_MOSTLY_Q4_K = 12, // except 1d tensors + GGML_FTYPE_MOSTLY_Q5_K = 13, // except 1d tensors + GGML_FTYPE_MOSTLY_Q6_K = 14, // except 1d tensors + GGML_FTYPE_MOSTLY_IQ2_XXS = 15, // except 1d tensors + GGML_FTYPE_MOSTLY_IQ2_XS = 16, // except 1d tensors + GGML_FTYPE_MOSTLY_IQ3_XXS = 17, // except 1d tensors + GGML_FTYPE_MOSTLY_IQ1_S = 18, // except 1d tensors + GGML_FTYPE_MOSTLY_IQ4_NL = 19, // except 1d tensors + GGML_FTYPE_MOSTLY_IQ3_S = 20, // except 1d tensors + GGML_FTYPE_MOSTLY_IQ2_S = 21, // except 1d tensors + GGML_FTYPE_MOSTLY_IQ4_XS = 22, // except 1d tensors + GGML_FTYPE_MOSTLY_IQ1_M = 23, // except 1d tensors + GGML_FTYPE_MOSTLY_BF16 = 24, // except 1d tensors + GGML_FTYPE_MOSTLY_MXFP4 = 25, // except 1d tensors + GGML_FTYPE_MOSTLY_NVFP4 = 26, // except 1d tensors + GGML_FTYPE_MOSTLY_Q1_0 = 27, // except 1d tensors + GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4 = 100, + GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_LEAN = 101, + GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_COHERENT = 102, + GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_FAST = 103, + GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_FAST_COHERENT = 104, + GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_STRIX = 105, + GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_STRIX_LEAN = 106, + GGML_FTYPE_MOSTLY_Q6_0_ROCMFPX = 110, + GGML_FTYPE_MOSTLY_Q8_0_ROCMFPX = 111, + GGML_FTYPE_MOSTLY_Q3_0_ROCMFPX = 112, + GGML_FTYPE_MOSTLY_Q2_0_ROCMFP2 = 118, + GGML_FTYPE_MOSTLY_Q2_0_ROCMFP2_STRIX = 119, + }; + + // available tensor operations: + enum ggml_op { + GGML_OP_NONE = 0, + + GGML_OP_DUP, + GGML_OP_ADD, + GGML_OP_ADD_ID, + GGML_OP_ADD1, + GGML_OP_ACC, + GGML_OP_SUB, + GGML_OP_MUL, + GGML_OP_DIV, + GGML_OP_SQR, + GGML_OP_SQRT, + GGML_OP_LOG, + GGML_OP_SIN, + GGML_OP_COS, + GGML_OP_SUM, + GGML_OP_SUM_ROWS, + GGML_OP_CUMSUM, + GGML_OP_MEAN, + GGML_OP_ARGMAX, + GGML_OP_COUNT_EQUAL, + GGML_OP_REPEAT, + GGML_OP_REPEAT_BACK, + GGML_OP_CONCAT, + GGML_OP_SILU_BACK, + GGML_OP_NORM, // normalize + GGML_OP_RMS_NORM, + GGML_OP_RMS_NORM_BACK, + GGML_OP_GROUP_NORM, + GGML_OP_L2_NORM, + + GGML_OP_MUL_MAT, + GGML_OP_MUL_MAT_ID, + GGML_OP_OUT_PROD, + + GGML_OP_SCALE, + GGML_OP_SET, + GGML_OP_CPY, + GGML_OP_CONT, + GGML_OP_RESHAPE, + GGML_OP_VIEW, + GGML_OP_PERMUTE, + GGML_OP_TRANSPOSE, + GGML_OP_GET_ROWS, + GGML_OP_GET_ROWS_BACK, + GGML_OP_SET_ROWS, + GGML_OP_DIAG, + GGML_OP_DIAG_MASK_INF, + GGML_OP_DIAG_MASK_ZERO, + GGML_OP_SOFT_MAX, + GGML_OP_SOFT_MAX_BACK, + GGML_OP_ROPE, + GGML_OP_ROPE_BACK, + GGML_OP_CLAMP, + GGML_OP_CONV_TRANSPOSE_1D, + GGML_OP_IM2COL, + GGML_OP_IM2COL_BACK, + GGML_OP_IM2COL_3D, + GGML_OP_CONV_2D, + GGML_OP_CONV_3D, + GGML_OP_CONV_2D_DW, + GGML_OP_CONV_TRANSPOSE_2D, + GGML_OP_POOL_1D, + GGML_OP_POOL_2D, + GGML_OP_POOL_2D_BACK, + GGML_OP_UPSCALE, + GGML_OP_PAD, + GGML_OP_PAD_REFLECT_1D, + GGML_OP_ROLL, + GGML_OP_ARANGE, + GGML_OP_TIMESTEP_EMBEDDING, + GGML_OP_ARGSORT, + GGML_OP_TOP_K, + GGML_OP_LEAKY_RELU, + GGML_OP_TRI, + GGML_OP_FILL, + + GGML_OP_FLASH_ATTN_EXT, + GGML_OP_FLASH_ATTN_SPARSE, + GGML_OP_FLASH_ATTN_BACK, + GGML_OP_SSM_CONV, + GGML_OP_SSM_SCAN, + GGML_OP_WIN_PART, + GGML_OP_WIN_UNPART, + GGML_OP_GET_REL_POS, + GGML_OP_ADD_REL_POS, + GGML_OP_RWKV_WKV6, + GGML_OP_GATED_LINEAR_ATTN, + GGML_OP_RWKV_WKV7, + GGML_OP_SOLVE_TRI, + GGML_OP_GATED_DELTA_NET, + + GGML_OP_UNARY, + + GGML_OP_MAP_CUSTOM1, + GGML_OP_MAP_CUSTOM2, + GGML_OP_MAP_CUSTOM3, + + GGML_OP_CUSTOM, + + GGML_OP_CROSS_ENTROPY_LOSS, + GGML_OP_CROSS_ENTROPY_LOSS_BACK, + GGML_OP_OPT_STEP_ADAMW, + GGML_OP_OPT_STEP_SGD, + + GGML_OP_GLU, + + GGML_OP_TURBO_WHT, // FWHT rotation for TurboQuant KV cache + + GGML_OP_MOE_FUSED, // Fused MoE FFN: gate+up+swiglu+down+weighted_sum+shared_expert + + GGML_OP_COUNT, + }; + + enum ggml_unary_op { + GGML_UNARY_OP_ABS, + GGML_UNARY_OP_SGN, + GGML_UNARY_OP_NEG, + GGML_UNARY_OP_STEP, + GGML_UNARY_OP_TANH, + GGML_UNARY_OP_ELU, + GGML_UNARY_OP_RELU, + GGML_UNARY_OP_SIGMOID, + GGML_UNARY_OP_GELU, + GGML_UNARY_OP_GELU_QUICK, + GGML_UNARY_OP_SILU, + GGML_UNARY_OP_HARDSWISH, + GGML_UNARY_OP_HARDSIGMOID, + GGML_UNARY_OP_EXP, + GGML_UNARY_OP_EXPM1, + GGML_UNARY_OP_SOFTPLUS, + GGML_UNARY_OP_GELU_ERF, + GGML_UNARY_OP_XIELU, + GGML_UNARY_OP_FLOOR, + GGML_UNARY_OP_CEIL, + GGML_UNARY_OP_ROUND, + GGML_UNARY_OP_TRUNC, + + GGML_UNARY_OP_COUNT, + }; + + enum ggml_glu_op { + GGML_GLU_OP_REGLU, + GGML_GLU_OP_GEGLU, + GGML_GLU_OP_SWIGLU, + GGML_GLU_OP_SWIGLU_OAI, + GGML_GLU_OP_SWIGLU_DS4, + GGML_GLU_OP_GEGLU_ERF, + GGML_GLU_OP_GEGLU_QUICK, + + GGML_GLU_OP_COUNT, + }; + + enum ggml_object_type { + GGML_OBJECT_TYPE_TENSOR, + GGML_OBJECT_TYPE_GRAPH, + GGML_OBJECT_TYPE_WORK_BUFFER + }; + + enum ggml_log_level { + GGML_LOG_LEVEL_NONE = 0, + GGML_LOG_LEVEL_DEBUG = 1, + GGML_LOG_LEVEL_INFO = 2, + GGML_LOG_LEVEL_WARN = 3, + GGML_LOG_LEVEL_ERROR = 4, + GGML_LOG_LEVEL_CONT = 5, // continue previous log + }; + + // this tensor... + enum ggml_tensor_flag { + GGML_TENSOR_FLAG_INPUT = 1, // ...is an input for the GGML compute graph + GGML_TENSOR_FLAG_OUTPUT = 2, // ...is an output for the GGML compute graph + GGML_TENSOR_FLAG_PARAM = 4, // ...contains trainable parameters + GGML_TENSOR_FLAG_LOSS = 8, // ...defines loss for numerical optimization (multiple loss tensors add up) + GGML_TENSOR_FLAG_COMPUTE = 16, // ...must be computed + }; + + enum ggml_tri_type { + GGML_TRI_TYPE_UPPER_DIAG = 0, + GGML_TRI_TYPE_UPPER = 1, + GGML_TRI_TYPE_LOWER_DIAG = 2, + GGML_TRI_TYPE_LOWER = 3 + }; + + struct ggml_init_params { + // memory pool + size_t mem_size; // bytes + void * mem_buffer; // if NULL, memory will be allocated internally + bool no_alloc; // don't allocate memory for the tensor data + }; + + // n-dimensional tensor + struct ggml_tensor { + enum ggml_type type; + + struct ggml_backend_buffer * buffer; + + int64_t ne[GGML_MAX_DIMS]; // number of elements + size_t nb[GGML_MAX_DIMS]; // stride in bytes: + // nb[0] = ggml_type_size(type) + // nb[1] = nb[0] * (ne[0] / ggml_blck_size(type)) + padding + // nb[i] = nb[i-1] * ne[i-1] + + // compute data + enum ggml_op op; + + // op params - allocated as int32_t for alignment + int32_t op_params[GGML_MAX_OP_PARAMS / sizeof(int32_t)]; + + int32_t flags; + + struct ggml_tensor * src[GGML_MAX_SRC]; + + // source tensor and offset for views + struct ggml_tensor * view_src; + size_t view_offs; + + void * data; + + char name[GGML_MAX_NAME]; + + void * extra; // extra things e.g. for ggml-cuda.cu + + char padding[8]; + }; + + static const size_t GGML_TENSOR_SIZE = sizeof(struct ggml_tensor); + + // Abort callback + // If not NULL, called before ggml computation + // If it returns true, the computation is aborted + typedef bool (*ggml_abort_callback)(void * data); + + + // + // GUID + // + + // GUID types + typedef uint8_t ggml_guid[16]; + typedef ggml_guid * ggml_guid_t; + + GGML_API bool ggml_guid_matches(ggml_guid_t guid_a, ggml_guid_t guid_b); + + // misc + + GGML_API const char * ggml_version(void); + GGML_API const char * ggml_commit(void); + + GGML_API void ggml_time_init(void); // call this once at the beginning of the program + GGML_API int64_t ggml_time_ms(void); + GGML_API int64_t ggml_time_us(void); + GGML_API int64_t ggml_cycles(void); + GGML_API int64_t ggml_cycles_per_ms(void); + + // accepts a UTF-8 path, even on Windows + GGML_API FILE * ggml_fopen(const char * fname, const char * mode); + + GGML_API void ggml_print_object (const struct ggml_object * obj); + GGML_API void ggml_print_objects(const struct ggml_context * ctx); + + GGML_API int64_t ggml_nelements (const struct ggml_tensor * tensor); + GGML_API int64_t ggml_nrows (const struct ggml_tensor * tensor); + GGML_API size_t ggml_nbytes (const struct ggml_tensor * tensor); + GGML_API size_t ggml_nbytes_pad(const struct ggml_tensor * tensor); // same as ggml_nbytes() but padded to GGML_MEM_ALIGN + + GGML_API int64_t ggml_blck_size(enum ggml_type type); + GGML_API size_t ggml_type_size(enum ggml_type type); // size in bytes for all elements in a block + GGML_API size_t ggml_row_size (enum ggml_type type, int64_t ne); // size in bytes for all elements in a row + + GGML_DEPRECATED( + GGML_API double ggml_type_sizef(enum ggml_type type), // ggml_type_size()/ggml_blck_size() as float + "use ggml_row_size() instead"); + + GGML_API const char * ggml_type_name(enum ggml_type type); + GGML_API const char * ggml_op_name (enum ggml_op op); + GGML_API const char * ggml_op_symbol(enum ggml_op op); + + GGML_API const char * ggml_unary_op_name(enum ggml_unary_op op); + GGML_API const char * ggml_glu_op_name(enum ggml_glu_op op); + GGML_API const char * ggml_op_desc(const struct ggml_tensor * t); // unary or op name + + GGML_API size_t ggml_element_size(const struct ggml_tensor * tensor); + + GGML_API bool ggml_is_quantized(enum ggml_type type); + + // TODO: temporary until model loading of ggml examples is refactored + GGML_API enum ggml_type ggml_ftype_to_ggml_type(enum ggml_ftype ftype); + + GGML_API bool ggml_is_transposed(const struct ggml_tensor * tensor); + GGML_API bool ggml_is_permuted (const struct ggml_tensor * tensor); + GGML_API bool ggml_is_empty (const struct ggml_tensor * tensor); + GGML_API bool ggml_is_view (const struct ggml_tensor * tensor); + GGML_API bool ggml_is_scalar (const struct ggml_tensor * tensor); + GGML_API bool ggml_is_vector (const struct ggml_tensor * tensor); + GGML_API bool ggml_is_matrix (const struct ggml_tensor * tensor); + GGML_API bool ggml_is_3d (const struct ggml_tensor * tensor); + GGML_API int ggml_n_dims (const struct ggml_tensor * tensor); // returns 1 for scalars + + // returns whether the tensor elements can be iterated over with a flattened index (no gaps, no permutation) + GGML_API bool ggml_is_contiguous (const struct ggml_tensor * tensor); + GGML_API bool ggml_is_contiguous_0(const struct ggml_tensor * tensor); // same as ggml_is_contiguous() + GGML_API bool ggml_is_contiguous_1(const struct ggml_tensor * tensor); // contiguous for dims >= 1 + GGML_API bool ggml_is_contiguous_2(const struct ggml_tensor * tensor); // contiguous for dims >= 2 + + // returns whether the tensor elements are allocated as one contiguous block of memory (no gaps, but permutation ok) + GGML_API bool ggml_is_contiguously_allocated(const struct ggml_tensor * tensor); + + // true for tensor that is stored in memory as CxWxHxN and has been permuted to WxHxCxN + GGML_API bool ggml_is_contiguous_channels(const struct ggml_tensor * tensor); + + // true if the elements in dimension 0 are contiguous, or there is just 1 block of elements + GGML_API bool ggml_is_contiguous_rows(const struct ggml_tensor * tensor); + + GGML_API bool ggml_are_same_shape (const struct ggml_tensor * t0, const struct ggml_tensor * t1); + GGML_API bool ggml_are_same_stride(const struct ggml_tensor * t0, const struct ggml_tensor * t1); + + GGML_API bool ggml_can_repeat(const struct ggml_tensor * t0, const struct ggml_tensor * t1); + + // use this to compute the memory overhead of a tensor + GGML_API size_t ggml_tensor_overhead(void); + + GGML_API bool ggml_validate_row_data(enum ggml_type type, const void * data, size_t nbytes); + + // main + + GGML_API struct ggml_context * ggml_init (struct ggml_init_params params); + GGML_API void ggml_reset(struct ggml_context * ctx); + GGML_API void ggml_free (struct ggml_context * ctx); + + GGML_API size_t ggml_used_mem(const struct ggml_context * ctx); + + GGML_API bool ggml_get_no_alloc(struct ggml_context * ctx); + GGML_API void ggml_set_no_alloc(struct ggml_context * ctx, bool no_alloc); + + GGML_API void * ggml_get_mem_buffer (const struct ggml_context * ctx); + GGML_API size_t ggml_get_mem_size (const struct ggml_context * ctx); + GGML_API size_t ggml_get_max_tensor_size(const struct ggml_context * ctx); + + GGML_API struct ggml_tensor * ggml_new_tensor( + struct ggml_context * ctx, + enum ggml_type type, + int n_dims, + const int64_t *ne); + + GGML_API struct ggml_tensor * ggml_new_tensor_1d( + struct ggml_context * ctx, + enum ggml_type type, + int64_t ne0); + + GGML_API struct ggml_tensor * ggml_new_tensor_2d( + struct ggml_context * ctx, + enum ggml_type type, + int64_t ne0, + int64_t ne1); + + GGML_API struct ggml_tensor * ggml_new_tensor_3d( + struct ggml_context * ctx, + enum ggml_type type, + int64_t ne0, + int64_t ne1, + int64_t ne2); + + GGML_API struct ggml_tensor * ggml_new_tensor_4d( + struct ggml_context * ctx, + enum ggml_type type, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3); + + GGML_API void * ggml_new_buffer(struct ggml_context * ctx, size_t nbytes); + + GGML_API struct ggml_tensor * ggml_dup_tensor (struct ggml_context * ctx, const struct ggml_tensor * src); + GGML_API struct ggml_tensor * ggml_view_tensor(struct ggml_context * ctx, struct ggml_tensor * src); + + // Context tensor enumeration and lookup + GGML_API struct ggml_tensor * ggml_get_first_tensor(const struct ggml_context * ctx); + GGML_API struct ggml_tensor * ggml_get_next_tensor (const struct ggml_context * ctx, struct ggml_tensor * tensor); + GGML_API struct ggml_tensor * ggml_get_tensor(struct ggml_context * ctx, const char * name); + + // Converts a flat index into coordinates + GGML_API void ggml_unravel_index(const struct ggml_tensor * tensor, int64_t i, int64_t * i0, int64_t * i1, int64_t * i2, int64_t * i3); + + GGML_API enum ggml_unary_op ggml_get_unary_op(const struct ggml_tensor * tensor); + GGML_API enum ggml_glu_op ggml_get_glu_op(const struct ggml_tensor * tensor); + + GGML_API void * ggml_get_data (const struct ggml_tensor * tensor); + GGML_API float * ggml_get_data_f32(const struct ggml_tensor * tensor); + + GGML_API const char * ggml_get_name (const struct ggml_tensor * tensor); + GGML_API struct ggml_tensor * ggml_set_name ( struct ggml_tensor * tensor, const char * name); + GGML_ATTRIBUTE_FORMAT(2, 3) + GGML_API struct ggml_tensor * ggml_format_name( struct ggml_tensor * tensor, const char * fmt, ...); + + // Tensor flags + GGML_API void ggml_set_input(struct ggml_tensor * tensor); + GGML_API void ggml_set_output(struct ggml_tensor * tensor); + GGML_API void ggml_set_param(struct ggml_tensor * tensor); + GGML_API void ggml_set_loss(struct ggml_tensor * tensor); + + // + // operations on tensors with backpropagation + // + + GGML_API struct ggml_tensor * ggml_dup( + struct ggml_context * ctx, + struct ggml_tensor * a); + + // in-place, returns view(a) + GGML_API struct ggml_tensor * ggml_dup_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_add( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + GGML_API struct ggml_tensor * ggml_add_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + GGML_API struct ggml_tensor * ggml_add_cast( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + enum ggml_type type); + + // dst[i0, i1, i2] = a[i0, i1, i2] + b[i0, ids[i1, i2]] + GGML_API struct ggml_tensor * ggml_add_id( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * ids); + + GGML_DEPRECATED(GGML_API struct ggml_tensor * ggml_add1( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b), + "use ggml_add instead"); + + GGML_DEPRECATED(GGML_API struct ggml_tensor * ggml_add1_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b), + "use ggml_add_inplace instead"); + + // dst = a + // view(dst, nb1, nb2, nb3, offset) += b + // return dst + GGML_API struct ggml_tensor * ggml_acc( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t nb1, + size_t nb2, + size_t nb3, + size_t offset); + + GGML_API struct ggml_tensor * ggml_acc_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t nb1, + size_t nb2, + size_t nb3, + size_t offset); + + GGML_API struct ggml_tensor * ggml_sub( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + GGML_API struct ggml_tensor * ggml_sub_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + GGML_API struct ggml_tensor * ggml_mul( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + GGML_API struct ggml_tensor * ggml_mul_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + GGML_API struct ggml_tensor * ggml_div( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + GGML_API struct ggml_tensor * ggml_div_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + GGML_API struct ggml_tensor * ggml_sqr( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_sqr_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_sqrt( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_sqrt_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_log( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_log_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_expm1( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_expm1_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_softplus( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_softplus_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_sin( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_sin_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_cos( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_cos_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + // return scalar + GGML_API struct ggml_tensor * ggml_sum( + struct ggml_context * ctx, + struct ggml_tensor * a); + + // sums along rows, with input shape [a,b,c,d] return shape [1,b,c,d] + GGML_API struct ggml_tensor * ggml_sum_rows( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_cumsum( + struct ggml_context * ctx, + struct ggml_tensor * a); + + // mean along rows + GGML_API struct ggml_tensor * ggml_mean( + struct ggml_context * ctx, + struct ggml_tensor * a); + + // argmax along rows + GGML_API struct ggml_tensor * ggml_argmax( + struct ggml_context * ctx, + struct ggml_tensor * a); + + // count number of equal elements in a and b + GGML_API struct ggml_tensor * ggml_count_equal( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + // if a is the same shape as b, and a is not parameter, return a + // otherwise, return a new tensor: repeat(a) to fit in b + GGML_API struct ggml_tensor * ggml_repeat( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + // repeat a to the specified shape + GGML_API struct ggml_tensor * ggml_repeat_4d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3); + + // sums repetitions in a into shape of b + GGML_API struct ggml_tensor * ggml_repeat_back( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); // sum up values that are adjacent in dims > 0 instead of repeated with same stride + + // concat a and b along dim + // used in stable-diffusion + GGML_API struct ggml_tensor * ggml_concat( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int dim); + + GGML_API struct ggml_tensor * ggml_abs( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_abs_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_sgn( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_sgn_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_neg( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_neg_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_step( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_step_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_tanh( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_tanh_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_elu( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_elu_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_relu( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_leaky_relu( + struct ggml_context * ctx, + struct ggml_tensor * a, float negative_slope, bool inplace); + + GGML_API struct ggml_tensor * ggml_relu_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_sigmoid( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_sigmoid_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_gelu( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_gelu_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + // GELU using erf (error function) when possible + // some backends may fallback to approximation based on Abramowitz and Stegun formula + GGML_API struct ggml_tensor * ggml_gelu_erf( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_gelu_erf_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_gelu_quick( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_gelu_quick_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_silu( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_silu_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + // a - x + // b - dy + GGML_API struct ggml_tensor * ggml_silu_back( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + // hardswish(x) = x * relu6(x + 3) / 6 + GGML_API struct ggml_tensor * ggml_hardswish( + struct ggml_context * ctx, + struct ggml_tensor * a); + + // hardsigmoid(x) = relu6(x + 3) / 6 + GGML_API struct ggml_tensor * ggml_hardsigmoid( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_exp( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_exp_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_floor( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_floor_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_ceil( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_ceil_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_round( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_round_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + /** + * Truncates the fractional part of each element in the tensor (towards zero). + * For example: trunc(3.7) = 3.0, trunc(-2.9) = -2.0 + * Similar to std::trunc in C/C++. + */ + + GGML_API struct ggml_tensor * ggml_trunc( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_trunc_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + + + // xIELU activation function + // x = x * (c_a(alpha_n) + c_b(alpha_p, beta) * sigmoid(beta * x)) + eps * (x > 0) + // where c_a = softplus and c_b(a, b) = softplus(a) + b are constraining functions + // that constrain the positive and negative source alpha values respectively + GGML_API struct ggml_tensor * ggml_xielu( + struct ggml_context * ctx, + struct ggml_tensor * a, + float alpha_n, + float alpha_p, + float beta, + float eps); + + // gated linear unit ops + // A: n columns, r rows, + // result is n / 2 columns, r rows, + // expects gate in second half of row, unless swapped is true + GGML_API struct ggml_tensor * ggml_glu( + struct ggml_context * ctx, + struct ggml_tensor * a, + enum ggml_glu_op op, + bool swapped); + + GGML_API struct ggml_tensor * ggml_reglu( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_reglu_swapped( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_geglu( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_geglu_swapped( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_swiglu( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_swiglu_swapped( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_geglu_erf( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_geglu_erf_swapped( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_geglu_quick( + struct ggml_context * ctx, + struct ggml_tensor * a); + + GGML_API struct ggml_tensor * ggml_geglu_quick_swapped( + struct ggml_context * ctx, + struct ggml_tensor * a); + + // A: n columns, r rows, + // B: n columns, r rows, + GGML_API struct ggml_tensor * ggml_glu_split( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + enum ggml_glu_op op); + + GGML_API struct ggml_tensor * ggml_reglu_split( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + GGML_API struct ggml_tensor * ggml_geglu_split( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + GGML_API struct ggml_tensor * ggml_swiglu_split( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + GGML_API struct ggml_tensor * ggml_swiglu_ds4_split( + struct ggml_context * ctx, + struct ggml_tensor * gate, + struct ggml_tensor * up, + float clamp); + + GGML_API struct ggml_tensor * ggml_geglu_erf_split( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + GGML_API struct ggml_tensor * ggml_geglu_quick_split( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + GGML_API struct ggml_tensor * ggml_swiglu_oai( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + float alpha, + float limit); + + // normalize along rows + GGML_API struct ggml_tensor * ggml_norm( + struct ggml_context * ctx, + struct ggml_tensor * a, + float eps); + + GGML_API struct ggml_tensor * ggml_norm_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + float eps); + + GGML_API struct ggml_tensor * ggml_rms_norm( + struct ggml_context * ctx, + struct ggml_tensor * a, + float eps); + + GGML_API struct ggml_tensor * ggml_rms_norm_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + float eps); + + // group normalize along ne0*ne1*n_groups + // used in stable-diffusion + GGML_API struct ggml_tensor * ggml_group_norm( + struct ggml_context * ctx, + struct ggml_tensor * a, + int n_groups, + float eps); + + GGML_API struct ggml_tensor * ggml_group_norm_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + int n_groups, + float eps); + + // l2 normalize along rows + // used in rwkv v7 + GGML_API struct ggml_tensor * ggml_l2_norm( + struct ggml_context * ctx, + struct ggml_tensor * a, + float eps); + + GGML_API struct ggml_tensor * ggml_l2_norm_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + float eps); + + // a - x + // b - dy + GGML_API struct ggml_tensor * ggml_rms_norm_back( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + float eps); + + // A: k columns, n rows => [ne03, ne02, n, k] + // B: k columns, m rows (i.e. we transpose it internally) => [ne03 * x, ne02 * y, m, k] + // result is n columns, m rows => [ne03 * x, ne02 * y, m, n] + GGML_API struct ggml_tensor * ggml_mul_mat( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + // change the precision of a matrix multiplication + // set to GGML_PREC_F32 for higher precision (useful for phi-2) + GGML_API void ggml_mul_mat_set_prec( + struct ggml_tensor * a, + enum ggml_prec prec); + + // indirect matrix multiplication + GGML_API struct ggml_tensor * ggml_mul_mat_id( + struct ggml_context * ctx, + struct ggml_tensor * as, + struct ggml_tensor * b, + struct ggml_tensor * ids); + + // A: m columns, n rows, + // B: p columns, n rows, + // result is m columns, p rows + GGML_API struct ggml_tensor * ggml_out_prod( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + // + // operations on tensors without backpropagation + // + + GGML_API struct ggml_tensor * ggml_scale( + struct ggml_context * ctx, + struct ggml_tensor * a, + float s); + + // in-place, returns view(a) + GGML_API struct ggml_tensor * ggml_scale_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + float s); + + // x = s * a + b + GGML_API struct ggml_tensor * ggml_scale_bias( + struct ggml_context * ctx, + struct ggml_tensor * a, + float s, + float b); + + GGML_API struct ggml_tensor * ggml_scale_bias_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + float s, + float b); + + // b -> view(a,offset,nb1,nb2,3), return modified a + GGML_API struct ggml_tensor * ggml_set( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t nb1, + size_t nb2, + size_t nb3, + size_t offset); // in bytes + + // b -> view(a,offset,nb1,nb2,3), return view(a) + GGML_API struct ggml_tensor * ggml_set_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t nb1, + size_t nb2, + size_t nb3, + size_t offset); // in bytes + + GGML_API struct ggml_tensor * ggml_set_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t offset); // in bytes + + GGML_API struct ggml_tensor * ggml_set_1d_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t offset); // in bytes + + // b -> view(a,offset,nb1,nb2,3), return modified a + GGML_API struct ggml_tensor * ggml_set_2d( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t nb1, + size_t offset); // in bytes + + // b -> view(a,offset,nb1,nb2,3), return view(a) + GGML_API struct ggml_tensor * ggml_set_2d_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t nb1, + size_t offset); // in bytes + + // a -> b, return view(b) + GGML_API struct ggml_tensor * ggml_cpy( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + // note: casting from f32 to i32 will discard the fractional part + GGML_API struct ggml_tensor * ggml_cast( + struct ggml_context * ctx, + struct ggml_tensor * a, + enum ggml_type type); + + // make contiguous + GGML_API struct ggml_tensor * ggml_cont( + struct ggml_context * ctx, + struct ggml_tensor * a); + + // make contiguous, with new shape + GGML_API struct ggml_tensor * ggml_cont_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0); + + GGML_API struct ggml_tensor * ggml_cont_2d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1); + + GGML_API struct ggml_tensor * ggml_cont_3d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + int64_t ne2); + + GGML_API struct ggml_tensor * ggml_cont_4d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3); + + // return view(a), b specifies the new shape + // TODO: when we start computing gradient, make a copy instead of view + GGML_API struct ggml_tensor * ggml_reshape( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + // return view(a) + // TODO: when we start computing gradient, make a copy instead of view + GGML_API struct ggml_tensor * ggml_reshape_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0); + + GGML_API struct ggml_tensor * ggml_reshape_2d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1); + + // return view(a) + // TODO: when we start computing gradient, make a copy instead of view + GGML_API struct ggml_tensor * ggml_reshape_3d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + int64_t ne2); + + GGML_API struct ggml_tensor * ggml_reshape_4d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3); + + // offset in bytes + GGML_API struct ggml_tensor * ggml_view_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + size_t offset); + + GGML_API struct ggml_tensor * ggml_view_2d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + size_t nb1, // row stride in bytes + size_t offset); + + GGML_API struct ggml_tensor * ggml_view_3d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + int64_t ne2, + size_t nb1, // row stride in bytes + size_t nb2, // slice stride in bytes + size_t offset); + + GGML_API struct ggml_tensor * ggml_view_4d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3, + size_t nb1, // row stride in bytes + size_t nb2, // slice stride in bytes + size_t nb3, + size_t offset); + + GGML_API struct ggml_tensor * ggml_permute( + struct ggml_context * ctx, + struct ggml_tensor * a, + int axis0, + int axis1, + int axis2, + int axis3); + + // alias for ggml_permute(ctx, a, 1, 0, 2, 3) + GGML_API struct ggml_tensor * ggml_transpose( + struct ggml_context * ctx, + struct ggml_tensor * a); + + // supports 4D a: + // a [n_embd, ne1, ne2, ne3] + // b I32 [n_rows, ne2, ne3, 1] + // + // return [n_embd, n_rows, ne2, ne3] + GGML_API struct ggml_tensor * ggml_get_rows( + struct ggml_context * ctx, + struct ggml_tensor * a, // data + struct ggml_tensor * b); // row indices + + GGML_API struct ggml_tensor * ggml_get_rows_back( + struct ggml_context * ctx, + struct ggml_tensor * a, // gradients of ggml_get_rows result + struct ggml_tensor * b, // row indices + struct ggml_tensor * c); // data for ggml_get_rows, only used for its shape + + // a TD [n_embd, ne1, ne2, ne3] + // b TS [n_embd, n_rows, ne02, ne03] | ne02 == ne2, ne03 == ne3 + // c I64 [n_rows, ne11, ne12, 1] | c[i] in [0, ne1) + // + // undefined behavior if destination rows overlap + // + // broadcast: + // ne2 % ne11 == 0 + // ne3 % ne12 == 0 + // + // return view(a) + GGML_API struct ggml_tensor * ggml_set_rows( + struct ggml_context * ctx, + struct ggml_tensor * a, // destination + struct ggml_tensor * b, // source + struct ggml_tensor * c); // row indices + + GGML_API struct ggml_tensor * ggml_diag( + struct ggml_context * ctx, + struct ggml_tensor * a); + + // set elements above the diagonal to -INF + GGML_API struct ggml_tensor * ggml_diag_mask_inf( + struct ggml_context * ctx, + struct ggml_tensor * a, + int n_past); + + // in-place, returns view(a) + GGML_API struct ggml_tensor * ggml_diag_mask_inf_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + int n_past); + + // set elements above the diagonal to 0 + GGML_API struct ggml_tensor * ggml_diag_mask_zero( + struct ggml_context * ctx, + struct ggml_tensor * a, + int n_past); + + // in-place, returns view(a) + GGML_API struct ggml_tensor * ggml_diag_mask_zero_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + int n_past); + + GGML_API struct ggml_tensor * ggml_soft_max( + struct ggml_context * ctx, + struct ggml_tensor * a); + + // in-place, returns view(a) + GGML_API struct ggml_tensor * ggml_soft_max_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a); + + // a [ne0, ne01, ne02, ne03] + // mask [ne0, ne11, ne12, ne13] | ne11 >= ne01, F16 or F32, optional + // + // broadcast: + // ne02 % ne12 == 0 + // ne03 % ne13 == 0 + // + // fused soft_max(a*scale + mask*(ALiBi slope)) + // max_bias = 0.0f for no ALiBi + GGML_API struct ggml_tensor * ggml_soft_max_ext( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * mask, + float scale, + float max_bias); + + GGML_API struct ggml_tensor * ggml_soft_max_ext_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * mask, + float scale, + float max_bias); + + GGML_API void ggml_soft_max_add_sinks( + struct ggml_tensor * a, + struct ggml_tensor * sinks); + + GGML_API struct ggml_tensor * ggml_soft_max_ext_back( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + float scale, + float max_bias); + + // in-place, returns view(a) + GGML_API struct ggml_tensor * ggml_soft_max_ext_back_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + float scale, + float max_bias); + + // rotary position embedding + // if (mode & 1) - skip n_past elements (NOT SUPPORTED) + // if (mode & GGML_ROPE_TYPE_NEOX) - GPT-NeoX style + // + // b is an int32 vector with size a->ne[2], it contains the positions + GGML_API struct ggml_tensor * ggml_rope( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int n_dims, + int mode); + + // in-place, returns view(a) + GGML_API struct ggml_tensor * ggml_rope_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int n_dims, + int mode); + + // custom RoPE + // c is freq factors (e.g. phi3-128k), (optional) + GGML_API struct ggml_tensor * ggml_rope_ext( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c, + int n_dims, + int mode, + int n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow); + + GGML_API struct ggml_tensor * ggml_rope_multi( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c, + int n_dims, + int sections[GGML_MROPE_SECTIONS], + int mode, + int n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow); + + // in-place, returns view(a) + GGML_API struct ggml_tensor * ggml_rope_ext_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c, + int n_dims, + int mode, + int n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow); + + GGML_API struct ggml_tensor * ggml_rope_multi_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c, + int n_dims, + int sections[GGML_MROPE_SECTIONS], + int mode, + int n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow); + + GGML_DEPRECATED(GGML_API struct ggml_tensor * ggml_rope_custom( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int n_dims, + int mode, + int n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow), + "use ggml_rope_ext instead"); + + GGML_DEPRECATED(GGML_API struct ggml_tensor * ggml_rope_custom_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int n_dims, + int mode, + int n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow), + "use ggml_rope_ext_inplace instead"); + + // compute correction dims for YaRN RoPE scaling + GGML_API void ggml_rope_yarn_corr_dims( + int n_dims, int n_ctx_orig, float freq_base, float beta_fast, float beta_slow, float dims[2]); + + // rotary position embedding backward, i.e compute dx from dy + // a - dy + GGML_API struct ggml_tensor * ggml_rope_ext_back( + struct ggml_context * ctx, + struct ggml_tensor * a, // gradients of ggml_rope result + struct ggml_tensor * b, // positions + struct ggml_tensor * c, // freq factors + int n_dims, + int mode, + int n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow); + + GGML_API struct ggml_tensor * ggml_rope_multi_back( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c, + int n_dims, + int sections[4], + int mode, + int n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow); + + + // clamp + // in-place, returns view(a) + GGML_API struct ggml_tensor * ggml_clamp( + struct ggml_context * ctx, + struct ggml_tensor * a, + float min, + float max); + + // im2col + // converts data into a format that effectively results in a convolution when combined with matrix multiplication + GGML_API struct ggml_tensor * ggml_im2col( + struct ggml_context * ctx, + struct ggml_tensor * a, // convolution kernel + struct ggml_tensor * b, // data + int s0, // stride dimension 0 + int s1, // stride dimension 1 + int p0, // padding dimension 0 + int p1, // padding dimension 1 + int d0, // dilation dimension 0 + int d1, // dilation dimension 1 + bool is_2D, + enum ggml_type dst_type); + + GGML_API struct ggml_tensor * ggml_im2col_back( + struct ggml_context * ctx, + struct ggml_tensor * a, // convolution kernel + struct ggml_tensor * b, // gradient of im2col output + int64_t * ne, // shape of im2col input + int s0, // stride dimension 0 + int s1, // stride dimension 1 + int p0, // padding dimension 0 + int p1, // padding dimension 1 + int d0, // dilation dimension 0 + int d1, // dilation dimension 1 + bool is_2D); + + GGML_API struct ggml_tensor * ggml_conv_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, // convolution kernel + struct ggml_tensor * b, // data + int s0, // stride + int p0, // padding + int d0); // dilation + + // conv_1d with padding = half + // alias for ggml_conv_1d(a, b, s, a->ne[0]/2, d) + GGML_API struct ggml_tensor* ggml_conv_1d_ph( + struct ggml_context * ctx, + struct ggml_tensor * a, // convolution kernel + struct ggml_tensor * b, // data + int s, // stride + int d); // dilation + + // depthwise + // TODO: this is very likely wrong for some cases! - needs more testing + GGML_API struct ggml_tensor * ggml_conv_1d_dw( + struct ggml_context * ctx, + struct ggml_tensor * a, // convolution kernel + struct ggml_tensor * b, // data + int s0, // stride + int p0, // padding + int d0); // dilation + + GGML_API struct ggml_tensor * ggml_conv_1d_dw_ph( + struct ggml_context * ctx, + struct ggml_tensor * a, // convolution kernel + struct ggml_tensor * b, // data + int s0, // stride + int d0); // dilation + + GGML_API struct ggml_tensor * ggml_conv_transpose_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, // convolution kernel + struct ggml_tensor * b, // data + int s0, // stride + int p0, // padding + int d0); // dilation + + GGML_API struct ggml_tensor * ggml_conv_2d( + struct ggml_context * ctx, + struct ggml_tensor * a, // convolution kernel + struct ggml_tensor * b, // data + int s0, // stride dimension 0 + int s1, // stride dimension 1 + int p0, // padding dimension 0 + int p1, // padding dimension 1 + int d0, // dilation dimension 0 + int d1); // dilation dimension 1 + + GGML_API struct ggml_tensor * ggml_im2col_3d( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int64_t IC, + int s0, // stride width + int s1, // stride height + int s2, // stride depth + int p0, // padding width + int p1, // padding height + int p2, // padding depth + int d0, // dilation width + int d1, // dilation height + int d2, // dilation depth + enum ggml_type dst_type); + + // a: [OC*IC, KD, KH, KW] + // b: [N*IC, ID, IH, IW] + // result: [N*OC, OD, OH, OW] + GGML_API struct ggml_tensor * ggml_conv_3d( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int64_t IC, + int s0, // stride width + int s1, // stride height + int s2, // stride depth + int p0, // padding width + int p1, // padding height + int p2, // padding depth + int d0, // dilation width + int d1, // dilation height + int d2 // dilation depth + ); + + // kernel size is a->ne[0] x a->ne[1] + // stride is equal to kernel size + // padding is zero + // example: + // a: 16 16 3 768 + // b: 1024 1024 3 1 + // res: 64 64 768 1 + // used in sam + GGML_API struct ggml_tensor * ggml_conv_2d_sk_p0( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + // kernel size is a->ne[0] x a->ne[1] + // stride is 1 + // padding is half + // example: + // a: 3 3 256 256 + // b: 64 64 256 1 + // res: 64 64 256 1 + // used in sam + GGML_API struct ggml_tensor * ggml_conv_2d_s1_ph( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b); + + // depthwise (via im2col and mul_mat) + GGML_API struct ggml_tensor * ggml_conv_2d_dw( + struct ggml_context * ctx, + struct ggml_tensor * a, // convolution kernel + struct ggml_tensor * b, // data + int s0, // stride dimension 0 + int s1, // stride dimension 1 + int p0, // padding dimension 0 + int p1, // padding dimension 1 + int d0, // dilation dimension 0 + int d1); // dilation dimension 1 + + // Depthwise 2D convolution + // may be faster than ggml_conv_2d_dw, but not available in all backends + // a: KW KH 1 C convolution kernel + // b: W H C N input data + // res: W_out H_out C N + GGML_API struct ggml_tensor * ggml_conv_2d_dw_direct( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int stride0, + int stride1, + int pad0, + int pad1, + int dilation0, + int dilation1); + + GGML_API struct ggml_tensor * ggml_conv_transpose_2d_p0( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int stride); + + GGML_API struct ggml_tensor * ggml_conv_2d_direct( + struct ggml_context * ctx, + struct ggml_tensor * a, // convolution kernel [KW, KH, IC, OC] + struct ggml_tensor * b, // input data [W, H, C, N] + int s0, // stride dimension 0 + int s1, // stride dimension 1 + int p0, // padding dimension 0 + int p1, // padding dimension 1 + int d0, // dilation dimension 0 + int d1); // dilation dimension 1 + + GGML_API struct ggml_tensor * ggml_conv_3d_direct( + struct ggml_context * ctx, + struct ggml_tensor * a, // kernel [KW, KH, KD, IC * OC] + struct ggml_tensor * b, // input [W, H, D, C * N] + int s0, // stride + int s1, + int s2, + int p0, // padding + int p1, + int p2, + int d0, // dilation + int d1, + int d2, + int n_channels, + int n_batch, + int n_channels_out); + + enum ggml_op_pool { + GGML_OP_POOL_MAX, + GGML_OP_POOL_AVG, + GGML_OP_POOL_COUNT, + }; + + GGML_API struct ggml_tensor * ggml_pool_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, + enum ggml_op_pool op, + int k0, // kernel size + int s0, // stride + int p0); // padding + + // the result will have 2*p0 padding for the first dimension + // and 2*p1 padding for the second dimension + GGML_API struct ggml_tensor * ggml_pool_2d( + struct ggml_context * ctx, + struct ggml_tensor * a, + enum ggml_op_pool op, + int k0, + int k1, + int s0, + int s1, + float p0, + float p1); + + GGML_API struct ggml_tensor * ggml_pool_2d_back( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * af, // "a"/input used in forward pass + enum ggml_op_pool op, + int k0, + int k1, + int s0, + int s1, + float p0, + float p1); + + enum ggml_scale_mode { + GGML_SCALE_MODE_NEAREST = 0, + GGML_SCALE_MODE_BILINEAR = 1, + GGML_SCALE_MODE_BICUBIC = 2, + + GGML_SCALE_MODE_COUNT + }; + + enum ggml_scale_flag { + GGML_SCALE_FLAG_ALIGN_CORNERS = (1 << 8), + GGML_SCALE_FLAG_ANTIALIAS = (1 << 9), + }; + + // interpolate + // multiplies ne0 and ne1 by scale factor + GGML_API struct ggml_tensor * ggml_upscale( + struct ggml_context * ctx, + struct ggml_tensor * a, + int scale_factor, + enum ggml_scale_mode mode); + + // interpolate + // interpolate scale to specified dimensions + GGML_DEPRECATED(GGML_API struct ggml_tensor * ggml_upscale_ext( + struct ggml_context * ctx, + struct ggml_tensor * a, + int ne0, + int ne1, + int ne2, + int ne3, + enum ggml_scale_mode mode), + "use ggml_interpolate instead"); + + // Up- or downsamples the input to the specified size. + // 2D scale modes (eg. bilinear) are applied to the first two dimensions. + GGML_API struct ggml_tensor * ggml_interpolate( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3, + uint32_t mode); // ggml_scale_mode [ | ggml_scale_flag...] + + // pad each dimension with zeros: [x, ..., x] -> [x, ..., x, 0, ..., 0] + GGML_API struct ggml_tensor * ggml_pad( + struct ggml_context * ctx, + struct ggml_tensor * a, + int p0, + int p1, + int p2, + int p3); + + // pad each dimension with values on the other side of the torus (looping around) + GGML_API struct ggml_tensor * ggml_pad_circular( + struct ggml_context * ctx, + struct ggml_tensor * a, + int p0, + int p1, + int p2, + int p3); + + GGML_API struct ggml_tensor * ggml_pad_ext( + struct ggml_context * ctx, + struct ggml_tensor * a, + int lp0, + int rp0, + int lp1, + int rp1, + int lp2, + int rp2, + int lp3, + int rp3 + ); + + // pad each dimension with values on the other side of the torus (looping around) + GGML_API struct ggml_tensor * ggml_pad_ext_circular( + struct ggml_context * ctx, + struct ggml_tensor * a, + int lp0, + int rp0, + int lp1, + int rp1, + int lp2, + int rp2, + int lp3, + int rp3); + + // pad each dimension with reflection: [a, b, c, d] -> [b, a, b, c, d, c] + GGML_API struct ggml_tensor * ggml_pad_reflect_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int p0, + int p1); + + // Move tensor elements by an offset given for each dimension. Elements that + // are shifted beyond the last position are wrapped around to the beginning. + GGML_API struct ggml_tensor * ggml_roll( + struct ggml_context * ctx, + struct ggml_tensor * a, + int shift0, + int shift1, + int shift2, + int shift3); + + // Convert matrix into a triangular one (upper, strict upper, lower or strict lower) by writing + // zeroes everywhere outside the masked area + GGML_API struct ggml_tensor * ggml_tri( + struct ggml_context * ctx, + struct ggml_tensor * a, + enum ggml_tri_type type); + + // Fill tensor a with constant c + GGML_API struct ggml_tensor * ggml_fill( + struct ggml_context * ctx, + struct ggml_tensor * a, + float c); + + GGML_API struct ggml_tensor * ggml_fill_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + float c); + + // Ref: https://github.com/CompVis/stable-diffusion/blob/main/ldm/modules/diffusionmodules/util.py#L151 + // timesteps: [N,] + // return: [N, dim] + GGML_API struct ggml_tensor * ggml_timestep_embedding( + struct ggml_context * ctx, + struct ggml_tensor * timesteps, + int dim, + int max_period); + + // sort rows + enum ggml_sort_order { + GGML_SORT_ORDER_ASC, + GGML_SORT_ORDER_DESC, + }; + + GGML_API struct ggml_tensor * ggml_argsort( + struct ggml_context * ctx, + struct ggml_tensor * a, + enum ggml_sort_order order); + + // similar to ggml_top_k but implemented as `argsort` + `view` + GGML_API struct ggml_tensor * ggml_argsort_top_k( + struct ggml_context * ctx, + struct ggml_tensor * a, + int k); + + // top k elements per row + // note: the resulting top k indices are in no particular order + GGML_API struct ggml_tensor * ggml_top_k( + struct ggml_context * ctx, + struct ggml_tensor * a, + int k); + + GGML_API struct ggml_tensor * ggml_arange( + struct ggml_context * ctx, + float start, + float stop, + float step); + + // q: [n_embd_k, n_batch, n_head, ne3 ] + // k: [n_embd_k, n_kv, n_head_kv, ne3 ] + // v: [n_embd_v, n_kv, n_head_kv, ne3 ] !! not transposed !! + // mask: [n_kv, n_batch, ne32, ne33] + // res: [n_embd_v, n_head, n_batch, ne3 ] !! permuted !! + // + // broadcast: + // n_head % n_head_kv == 0 + // n_head % ne32 == 0 + // ne3 % ne33 == 0 + // + GGML_API struct ggml_tensor * ggml_flash_attn_ext( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * mask, + float scale, + float max_bias, + float logit_softcap); + + GGML_API void ggml_flash_attn_ext_set_prec( + struct ggml_tensor * a, + enum ggml_prec prec); + + GGML_API enum ggml_prec ggml_flash_attn_ext_get_prec( + const struct ggml_tensor * a); + + GGML_API void ggml_flash_attn_ext_add_sinks( + struct ggml_tensor * a, + struct ggml_tensor * sinks); + + // Block-sparse flash attention (pFlash). + // Selects top-scoring K-blocks dynamically based on alpha threshold. + // alpha=1.0 selects all blocks (equivalent to dense attention). + // Q: [D, S, H, B], K: [D, S, Hk, B], V: [D, S, Hk, B] (ggml layout) + // Returns: [D, H, S, B] F32 (same shape as ggml_flash_attn_ext output) + GGML_API struct ggml_tensor * ggml_flash_attn_sparse( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + float scale, + float alpha); + + // TurboQuant FWHT rotation. direction: 0 = forward, 1 = inverse. + // Applies signs1 -> FWHT -> signs2 (forward) or signs2 -> FWHT -> signs1 (inverse). + // Used for KV cache rotation in TurboQuant quantization types (TQ3_0). + GGML_API struct ggml_tensor * ggml_turbo_wht( + struct ggml_context * ctx, + struct ggml_tensor * a, + int direction); + + GGML_API struct ggml_tensor * ggml_moe_fused( + struct ggml_context * ctx, + struct ggml_tensor * input, + struct ggml_tensor * gate_w, + struct ggml_tensor * up_w, + struct ggml_tensor * down_w, + struct ggml_tensor * expert_ids, + struct ggml_tensor * expert_weights, + struct ggml_tensor * sh_gate_w, + struct ggml_tensor * sh_up_w, + struct ggml_tensor * sh_down_w, + struct ggml_tensor * sh_gate_inp_w, + int64_t n_embd, + int64_t ff_dim, + int64_t n_expert_used); + + GGML_API struct ggml_tensor * ggml_laguna_moe_combine( + struct ggml_context * ctx, + struct ggml_tensor * experts, + struct ggml_tensor * expert_weights); + + // TODO: needs to be adapted to ggml_flash_attn_ext + GGML_API struct ggml_tensor * ggml_flash_attn_back( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * d, + bool masked); + + GGML_API struct ggml_tensor * ggml_ssm_conv( + struct ggml_context * ctx, + struct ggml_tensor * sx, + struct ggml_tensor * c); + + // dflash extension: tree-mode ssm_conv for DDTree-style spec + // decoding. parent_ids is an int32 tensor of shape [n_tokens, n_seqs] + // where entry [t, s] is the new-token index of t's tree parent, or -1 if + // t's parent is "before the block" (walks into the old conv state region). + // The CUDA kernel walks the parent chain K-1 times per token to build the + // correct conv window, so siblings get their own parent lineage instead + // of sharing the DFS-neighbour window. Mirrors SGLang's + // causal_conv1d_triton HAS_EAGLE_TREE_CUSTOM_ATTN_MASK path. + GGML_API struct ggml_tensor * ggml_ssm_conv_tree( + struct ggml_context * ctx, + struct ggml_tensor * sx, + struct ggml_tensor * c, + struct ggml_tensor * parent_ids); + + GGML_API struct ggml_tensor * ggml_ssm_scan( + struct ggml_context * ctx, + struct ggml_tensor * s, + struct ggml_tensor * x, + struct ggml_tensor * dt, + struct ggml_tensor * A, + struct ggml_tensor * B, + struct ggml_tensor * C, + struct ggml_tensor * ids); + + // partition into non-overlapping windows with padding if needed + // example: + // a: 768 64 64 1 + // w: 14 + // res: 768 14 14 25 + // used in sam + GGML_API struct ggml_tensor * ggml_win_part( + struct ggml_context * ctx, + struct ggml_tensor * a, + int w); + + // reverse of ggml_win_part + // used in sam + GGML_API struct ggml_tensor * ggml_win_unpart( + struct ggml_context * ctx, + struct ggml_tensor * a, + int w0, + int h0, + int w); + + GGML_API struct ggml_tensor * ggml_unary( + struct ggml_context * ctx, + struct ggml_tensor * a, + enum ggml_unary_op op); + + GGML_API struct ggml_tensor * ggml_unary_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + enum ggml_unary_op op); + + // used in sam + GGML_API struct ggml_tensor * ggml_get_rel_pos( + struct ggml_context * ctx, + struct ggml_tensor * a, + int qh, + int kh); + + // used in sam + GGML_API struct ggml_tensor * ggml_add_rel_pos( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * pw, + struct ggml_tensor * ph); + + GGML_API struct ggml_tensor * ggml_add_rel_pos_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * pw, + struct ggml_tensor * ph); + + GGML_API struct ggml_tensor * ggml_rwkv_wkv6( + struct ggml_context * ctx, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * r, + struct ggml_tensor * tf, + struct ggml_tensor * td, + struct ggml_tensor * state); + + GGML_API struct ggml_tensor * ggml_gated_linear_attn( + struct ggml_context * ctx, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * q, + struct ggml_tensor * g, + struct ggml_tensor * state, + float scale); + + GGML_API struct ggml_tensor * ggml_rwkv_wkv7( + struct ggml_context * ctx, + struct ggml_tensor * r, + struct ggml_tensor * w, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * state); + + /* Solves a specific equation of the form Ax=B, where A is a triangular matrix + * without zeroes on the diagonal (i.e. invertible). + * B can have any number of columns, but must have the same number of rows as A + * If A is [n, n] and B is [n, m], then the result will be [n, m] as well + * Has O(n^3) complexity (unlike most matrix ops out there), so use on cases + * where n > 100 sparingly, pre-chunk if necessary. + * + * If left = false, solves xA=B instead + * If lower = false, assumes upper triangular instead + * If uni = true, assumes diagonal of A to be all ones (will override actual values) + * + * TODO: currently only lower, right, non-unitriangular variant is implemented + */ + GGML_API struct ggml_tensor * ggml_solve_tri( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + bool left, + bool lower, + bool uni); + + // TODO: add ggml_gated_delta_net_set_bcast() to be able to configure Q, K broadcast type: tiled vs interleaved [TAG_GGML_GDN_BCAST] + // ref: https://github.com/ggml-org/llama.cpp/pull/19468#discussion_r2786394306 + GGML_API struct ggml_tensor * ggml_gated_delta_net( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * g, + struct ggml_tensor * beta, + struct ggml_tensor * state); + + GGML_API void ggml_gated_delta_net_set_skip_intermediate( + struct ggml_tensor * tensor, + bool skip_intermediate); + + // dflash extension: tree-mode gated delta net for DDTree-style + // speculative decoding verify. `parent_ids` is an int32 tensor of shape + // [n_tokens, n_seqs] where entry [t, s] is the index within sequence s of + // the parent token in the DFS-flattened tree (or -1 for a root-level + // node). At each token step t > 0 in the recurrence, if parent_ids[t] is + // not (t - 1), the kernel reloads the recurrent state from the + // intermediate-states region at parent_ids[t] instead of continuing + // sequentially. This lets one verify forward pass correctly process + // multiple tree branches (siblings) without cross-contamination. + GGML_API struct ggml_tensor * ggml_gated_delta_net_tree( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * g, + struct ggml_tensor * beta, + struct ggml_tensor * state, + struct ggml_tensor * parent_ids); + + // dflash extension: tree-mode with direct intermediate-state writes to a + // persistent external buffer. Identical to ggml_gated_delta_net_tree but + // the per-token intermediate states are written to `persist_inter->data` + // (f32 or f16, [S_v, S_v, H, n_tokens, n_seqs], contiguous) instead of + // the default internal region of the result tensor. Eliminates a + // downstream ggml_cpy into the persistent cache buffer, saving ~5-10 ms + // per verify step on a 27B hybrid target. + GGML_API struct ggml_tensor * ggml_gated_delta_net_tree_persist( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * g, + struct ggml_tensor * beta, + struct ggml_tensor * state, + struct ggml_tensor * parent_ids, + struct ggml_tensor * persist_inter); + + // custom operators + + typedef void (*ggml_custom1_op_t)(struct ggml_tensor * dst , const struct ggml_tensor * a, int ith, int nth, void * userdata); + typedef void (*ggml_custom2_op_t)(struct ggml_tensor * dst , const struct ggml_tensor * a, const struct ggml_tensor * b, int ith, int nth, void * userdata); + typedef void (*ggml_custom3_op_t)(struct ggml_tensor * dst , const struct ggml_tensor * a, const struct ggml_tensor * b, const struct ggml_tensor * c, int ith, int nth, void * userdata); + +#define GGML_N_TASKS_MAX (-1) + // n_tasks == GGML_N_TASKS_MAX means to use max number of tasks + + GGML_API struct ggml_tensor * ggml_map_custom1( + struct ggml_context * ctx, + struct ggml_tensor * a, + ggml_custom1_op_t fun, + int n_tasks, + void * userdata); + + GGML_API struct ggml_tensor * ggml_map_custom1_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + ggml_custom1_op_t fun, + int n_tasks, + void * userdata); + + GGML_API struct ggml_tensor * ggml_map_custom2( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + ggml_custom2_op_t fun, + int n_tasks, + void * userdata); + + GGML_API struct ggml_tensor * ggml_map_custom2_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + ggml_custom2_op_t fun, + int n_tasks, + void * userdata); + + GGML_API struct ggml_tensor * ggml_map_custom3( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c, + ggml_custom3_op_t fun, + int n_tasks, + void * userdata); + + GGML_API struct ggml_tensor * ggml_map_custom3_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c, + ggml_custom3_op_t fun, + int n_tasks, + void * userdata); + + typedef void (*ggml_custom_op_t)(struct ggml_tensor * dst , int ith, int nth, void * userdata); + + GGML_API struct ggml_tensor * ggml_custom_4d( + struct ggml_context * ctx, + enum ggml_type type, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3, + struct ggml_tensor ** args, + int n_args, + ggml_custom_op_t fun, + int n_tasks, + void * userdata); + + GGML_API struct ggml_tensor * ggml_custom_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor ** args, + int n_args, + ggml_custom_op_t fun, + int n_tasks, + void * userdata); + + // loss function + + GGML_API struct ggml_tensor * ggml_cross_entropy_loss( + struct ggml_context * ctx, + struct ggml_tensor * a, // logits + struct ggml_tensor * b); // labels + + GGML_API struct ggml_tensor * ggml_cross_entropy_loss_back( + struct ggml_context * ctx, + struct ggml_tensor * a, // logits + struct ggml_tensor * b, // labels + struct ggml_tensor * c); // gradients of cross_entropy_loss result + + // AdamW optimizer step + // Paper: https://arxiv.org/pdf/1711.05101v3.pdf + // PyTorch: https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html + GGML_API struct ggml_tensor * ggml_opt_step_adamw( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * grad, + struct ggml_tensor * m, + struct ggml_tensor * v, + struct ggml_tensor * adamw_params); // parameters such as the learning rate + + // stochastic gradient descent step (with weight decay) + GGML_API struct ggml_tensor * ggml_opt_step_sgd( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * grad, + struct ggml_tensor * sgd_params); // alpha, weight decay + + // build forward multiple tensors and select one of them for computing + // this is useful for creating graphs that have constant topology but compute different things based on the input + // ref: https://github.com/ggml-org/llama.cpp/pull/18550 + // + // nodes: + // | - build forward into the graph but do not compute + // c - build forward into the graph and compute + // + // | | ... c ... | + // | | ... c ... | + // | | ... c ... | + // [0 1 ... idx ... n-1] <-- ggml_build_forward_select(..., n, idx) + // c + // c + // + // example: + // struct ggml_tensor * curs[3]; + // + // curs[0] = compute0(...); + // curs[1] = compute1(...); + // curs[2] = compute2(...); + // + // int idx = select_branch(some_input); + // + // struct ggml_tensor * out = ggml_build_forward_select(cgraph, curs, 3, idx); + // + GGML_API struct ggml_tensor * ggml_build_forward_select( + struct ggml_cgraph * cgraph, + struct ggml_tensor ** tensors, + int n_tensors, + int idx); + + GGML_API void ggml_build_forward_expand( + struct ggml_cgraph * cgraph, + struct ggml_tensor * tensor); + + GGML_API void ggml_build_backward_expand( + struct ggml_context * ctx, // context for gradient computation + struct ggml_cgraph * cgraph, + struct ggml_tensor ** grad_accs); + + // graph allocation in a context + GGML_API struct ggml_cgraph * ggml_new_graph (struct ggml_context * ctx); // size = GGML_DEFAULT_GRAPH_SIZE, grads = false + GGML_API struct ggml_cgraph * ggml_new_graph_custom(struct ggml_context * ctx, size_t size, bool grads); + GGML_API struct ggml_cgraph * ggml_graph_dup (struct ggml_context * ctx, struct ggml_cgraph * cgraph, bool force_grads); + GGML_API void ggml_graph_cpy (struct ggml_cgraph * src, struct ggml_cgraph * dst); + GGML_API void ggml_graph_reset (struct ggml_cgraph * cgraph); // set regular grads + optimizer momenta to 0, set loss grad to 1 + GGML_API void ggml_graph_clear (struct ggml_cgraph * cgraph); + + GGML_API int ggml_graph_size (struct ggml_cgraph * cgraph); + GGML_API struct ggml_tensor * ggml_graph_node (struct ggml_cgraph * cgraph, int i); // if i < 0, returns nodes[n_nodes + i] + GGML_API struct ggml_tensor ** ggml_graph_nodes (struct ggml_cgraph * cgraph); + GGML_API int ggml_graph_n_nodes(struct ggml_cgraph * cgraph); + + GGML_API void ggml_graph_add_node(struct ggml_cgraph * cgraph, struct ggml_tensor * tensor); + + GGML_API size_t ggml_graph_overhead(void); + GGML_API size_t ggml_graph_overhead_custom(size_t size, bool grads); + + GGML_API struct ggml_tensor * ggml_graph_get_tensor (const struct ggml_cgraph * cgraph, const char * name); + GGML_API struct ggml_tensor * ggml_graph_get_grad (const struct ggml_cgraph * cgraph, const struct ggml_tensor * node); + GGML_API struct ggml_tensor * ggml_graph_get_grad_acc(const struct ggml_cgraph * cgraph, const struct ggml_tensor * node); + + // print info and performance information for the graph + GGML_API void ggml_graph_print(const struct ggml_cgraph * cgraph); + + // dump the graph into a file using the dot format + GGML_API void ggml_graph_dump_dot(const struct ggml_cgraph * gb, const struct ggml_cgraph * cgraph, const char * filename); + + // TODO these functions were sandwiched in the old optimization interface, is there a better place for them? + typedef void (*ggml_log_callback)(enum ggml_log_level level, const char * text, void * user_data); + + // Set callback for all future logging events. + // If this is not called, or NULL is supplied, everything is output on stderr. + GGML_API void ggml_log_get(ggml_log_callback * log_callback, void ** user_data); + GGML_API void ggml_log_set(ggml_log_callback log_callback, void * user_data); + + GGML_API struct ggml_tensor * ggml_set_zero(struct ggml_tensor * tensor); + + // + // quantization + // + + // - ggml_quantize_init can be called multiple times with the same type + // it will only initialize the quantization tables for the first call or after ggml_quantize_free + // automatically called by ggml_quantize_chunk for convenience + // + // - ggml_quantize_free will free any memory allocated by ggml_quantize_init + // call this at the end of the program to avoid memory leaks + // + // note: these are thread-safe + // + GGML_API void ggml_quantize_init(enum ggml_type type); + GGML_API void ggml_quantize_free(void); + + // some quantization type cannot be used without an importance matrix + GGML_API bool ggml_quantize_requires_imatrix(enum ggml_type type); + + // calls ggml_quantize_init internally (i.e. can allocate memory) + GGML_API size_t ggml_quantize_chunk( + enum ggml_type type, + const float * src, + void * dst, + int64_t start, + int64_t nrows, + int64_t n_per_row, + const float * imatrix); + +#ifdef __cplusplus + // restrict not standard in C++ +# if defined(__GNUC__) +# define GGML_RESTRICT __restrict__ +# elif defined(__clang__) +# define GGML_RESTRICT __restrict +# elif defined(_MSC_VER) +# define GGML_RESTRICT __restrict +# else +# define GGML_RESTRICT +# endif +#else +# if defined (_MSC_VER) && (__STDC_VERSION__ < 201112L) +# define GGML_RESTRICT __restrict +# else +# define GGML_RESTRICT restrict +# endif +#endif + typedef void (*ggml_to_float_t) (const void * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); + typedef void (*ggml_from_float_t)(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); + + struct ggml_type_traits { + const char * type_name; + int64_t blck_size; + int64_t blck_size_interleave; // interleave elements in blocks + size_t type_size; + bool is_quantized; + ggml_to_float_t to_float; + ggml_from_float_t from_float_ref; + }; + + GGML_API const struct ggml_type_traits * ggml_get_type_traits(enum ggml_type type); + + // ggml threadpool + // TODO: currently, only a few functions are in the base ggml API, while the rest are in the CPU backend + // the goal should be to create an API that other backends can use move everything to the ggml base + + // scheduling priorities + enum ggml_sched_priority { + GGML_SCHED_PRIO_LOW = -1, + GGML_SCHED_PRIO_NORMAL, + GGML_SCHED_PRIO_MEDIUM, + GGML_SCHED_PRIO_HIGH, + GGML_SCHED_PRIO_REALTIME + }; + + // threadpool params + // Use ggml_threadpool_params_default() or ggml_threadpool_params_init() to populate the defaults + struct ggml_threadpool_params { + bool cpumask[GGML_MAX_N_THREADS]; // mask of cpu cores (all-zeros means use default affinity settings) + int n_threads; // number of threads + enum ggml_sched_priority prio; // thread priority + uint32_t poll; // polling level (0 - no polling, 100 - aggressive polling) + bool strict_cpu; // strict cpu placement + bool paused; // start in paused state + }; + + struct ggml_threadpool; // forward declaration, see ggml.c + + typedef struct ggml_threadpool * ggml_threadpool_t; + + GGML_API struct ggml_threadpool_params ggml_threadpool_params_default(int n_threads); + GGML_API void ggml_threadpool_params_init (struct ggml_threadpool_params * p, int n_threads); + GGML_API bool ggml_threadpool_params_match (const struct ggml_threadpool_params * p0, const struct ggml_threadpool_params * p1); + +#ifdef __cplusplus +} +#endif diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index f0e14a00673c..170f1e4a942d 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -1826,6 +1826,10 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm { GGML_ABORT("GGML_OP_MOE_FUSED is only implemented for CUDA"); } + case GGML_OP_DS4_HC: + { + GGML_ABORT("GGML_OP_DS4_HC is only implemented for CUDA"); + } case GGML_OP_OUT_PROD: { ggml_compute_forward_out_prod(params, tensor); diff --git a/ggml/src/ggml-cpu/ggml-cpu.c.ds4bak b/ggml/src/ggml-cpu/ggml-cpu.c.ds4bak new file mode 100644 index 000000000000..f0e14a00673c --- /dev/null +++ b/ggml/src/ggml-cpu/ggml-cpu.c.ds4bak @@ -0,0 +1,3782 @@ +#define _CRT_SECURE_NO_DEPRECATE // Disables "unsafe" warnings on Windows +#define _USE_MATH_DEFINES // For M_PI on MSVC + +#include "ggml-backend-impl.h" +#include "ggml-backend.h" +#include "traits.h" +#include "ggml-cpu-impl.h" +#include "ggml-impl.h" +#include "quants.h" +#include "ggml-threading.h" +#include "unary-ops.h" +#include "binary-ops.h" +#include "vec.h" +#include "ops.h" +#include "ggml.h" +#include "common.h" + +#if defined(_MSC_VER) || defined(__MINGW32__) +#include // using malloc.h with MSC/MINGW +#elif !defined(__FreeBSD__) && !defined(__NetBSD__) && !defined(__OpenBSD__) +#include +#endif + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#if defined(__gnu_linux__) +#include +#endif + +#ifdef GGML_USE_OPENMP +#include +#endif + +#if defined(__ARM_FEATURE_SVE) || defined(__ARM_FEATURE_MATMUL_INT8) +#undef GGML_USE_LLAMAFILE +#endif + +#ifdef GGML_USE_LLAMAFILE +#include "llamafile/sgemm.h" +#endif + +// Note: once we move threading into a separate C++ file +// will use std::hardware_destructive_interference_size instead of hardcoding it here +// and we'll use C++ attribute syntax. +#define GGML_CACHE_LINE 64 + +#if defined(__clang__) || defined(__GNUC__) +#define GGML_CACHE_ALIGN __attribute__((aligned(GGML_CACHE_LINE))) +#endif + +#if defined(__has_feature) +#if __has_feature(thread_sanitizer) +#define GGML_TSAN_ENABLED 1 +#endif +#else // __has_feature +#if defined(__SANITIZE_THREAD__) +#define GGML_TSAN_ENABLED 1 +#endif +#endif // __has_feature + +#define UNUSED GGML_UNUSED +#define SWAP(x, y, T) do { T SWAP = x; (x) = y; (y) = SWAP; } while (0) + +// precomputed f32 table for f16 (256 KB) (simd-mappings.h) +float ggml_table_f32_f16[1 << 16]; + +// precomputed f32 table for e8m0 half (1 KB) (simd-mappings.h) +float ggml_table_f32_e8m0_half[1 << 8]; + +#if defined(__ARM_ARCH) +struct ggml_arm_arch_features_type { + int sve_cnt; +} ggml_arm_arch_features = { 0 }; +#endif + +#if defined(__riscv) +struct ggml_riscv_arch_features_type { + int rvv_vlen; +} ggml_riscv_arch_features = { 0 }; +#endif + +#if defined(_WIN32) + +#define WIN32_LEAN_AND_MEAN +#ifndef NOMINMAX + #define NOMINMAX +#endif +#include + +#if defined(_MSC_VER) && !defined(__clang__) +#define GGML_CACHE_ALIGN __declspec(align(GGML_CACHE_LINE)) + +typedef volatile LONG atomic_int; +typedef atomic_int atomic_bool; +typedef atomic_int atomic_flag; + +#define ATOMIC_FLAG_INIT 0 + +typedef enum { + memory_order_relaxed, + memory_order_consume, + memory_order_acquire, + memory_order_release, + memory_order_acq_rel, + memory_order_seq_cst +} memory_order; + +static void atomic_store(atomic_int * ptr, LONG val) { + InterlockedExchange(ptr, val); +} +static void atomic_store_explicit(atomic_int * ptr, LONG val, memory_order mo) { + // TODO: add support for explicit memory order + InterlockedExchange(ptr, val); +} +static LONG atomic_load(atomic_int * ptr) { + return InterlockedCompareExchange(ptr, 0, 0); +} +static LONG atomic_load_explicit(atomic_int * ptr, memory_order mo) { + // TODO: add support for explicit memory order + return InterlockedCompareExchange(ptr, 0, 0); +} +static LONG atomic_fetch_add(atomic_int * ptr, LONG inc) { + return InterlockedExchangeAdd(ptr, inc); +} +static LONG atomic_fetch_add_explicit(atomic_int * ptr, LONG inc, memory_order mo) { + // TODO: add support for explicit memory order + return InterlockedExchangeAdd(ptr, inc); +} +static atomic_bool atomic_flag_test_and_set(atomic_flag * ptr) { + return InterlockedExchange(ptr, 1); +} +static void atomic_flag_clear(atomic_flag * ptr) { + InterlockedExchange(ptr, 0); +} +static void atomic_thread_fence(memory_order mo) { + MemoryBarrier(); +} +#else // clang +#include +#endif + +typedef HANDLE pthread_t; + +typedef DWORD thread_ret_t; +static int pthread_create(pthread_t * out, void * unused, thread_ret_t(*func)(void *), void * arg) { + (void) unused; + HANDLE handle = CreateThread(NULL, 0, (LPTHREAD_START_ROUTINE) func, arg, 0, NULL); + if (handle == NULL) + { + return EAGAIN; + } + + *out = handle; + return 0; +} + +static int pthread_join(pthread_t thread, void * unused) { + (void) unused; + int ret = (int) WaitForSingleObject(thread, INFINITE); + CloseHandle(thread); + return ret; +} + +static int sched_yield (void) { + Sleep (0); + return 0; +} +#else + +#include +#include +#include +#if defined(__FreeBSD__) +#include +#endif + +typedef void * thread_ret_t; + +#include +#include +#include + +#endif + +typedef pthread_t ggml_thread_t; + +#define GGML_THREADPOOL_N_THREADS_MASK (0xffffU) +#define GGML_THREADPOOL_N_THREADS_BITS (16) + +#if defined(__APPLE__) +#include +#include +#include +#endif + +static const struct ggml_type_traits_cpu type_traits_cpu[GGML_TYPE_COUNT] = { + [GGML_TYPE_F32] = { + .from_float = (ggml_from_float_t) ggml_cpu_fp32_to_fp32, + .vec_dot = (ggml_vec_dot_t) ggml_vec_dot_f32, + .vec_dot_type = GGML_TYPE_F32, + .nrows = 1, + }, + [GGML_TYPE_F16] = { + .from_float = (ggml_from_float_t) ggml_cpu_fp32_to_fp16, + .vec_dot = (ggml_vec_dot_t) ggml_vec_dot_f16, + .vec_dot_type = GGML_TYPE_F16, + .nrows = 1, + }, + [GGML_TYPE_Q1_0] = { + .from_float = quantize_row_q1_0, + .vec_dot = ggml_vec_dot_q1_0_q8_0, + .vec_dot_type = GGML_TYPE_Q8_0, + .nrows = 1, + }, + [GGML_TYPE_Q4_0] = { + .from_float = quantize_row_q4_0, + .vec_dot = ggml_vec_dot_q4_0_q8_0, + .vec_dot_type = GGML_TYPE_Q8_0, +#if defined (__ARM_FEATURE_MATMUL_INT8) + .nrows = 2, +#else + .nrows = 1, +#endif + }, + [GGML_TYPE_Q4_1] = { + .from_float = quantize_row_q4_1, + .vec_dot = ggml_vec_dot_q4_1_q8_1, + .vec_dot_type = GGML_TYPE_Q8_1, +#if defined (__ARM_FEATURE_MATMUL_INT8) + .nrows = 2, +#else + .nrows = 1, +#endif + }, + [GGML_TYPE_Q5_0] = { + .from_float = quantize_row_q5_0, + .vec_dot = ggml_vec_dot_q5_0_q8_0, + .vec_dot_type = GGML_TYPE_Q8_0, + .nrows = 1, + }, + [GGML_TYPE_Q5_1] = { + .from_float = quantize_row_q5_1, + .vec_dot = ggml_vec_dot_q5_1_q8_1, + .vec_dot_type = GGML_TYPE_Q8_1, + .nrows = 1, + }, + [GGML_TYPE_Q8_0] = { + .from_float = quantize_row_q8_0, + .vec_dot = ggml_vec_dot_q8_0_q8_0, + .vec_dot_type = GGML_TYPE_Q8_0, +#if defined (__ARM_FEATURE_MATMUL_INT8) + .nrows = 2, +#else + .nrows = 1, +#endif + }, + [GGML_TYPE_Q8_1] = { + .from_float = quantize_row_q8_1, + .vec_dot_type = GGML_TYPE_Q8_1, + .nrows = 1, + }, + [GGML_TYPE_MXFP4] = { + .from_float = quantize_row_mxfp4, + .vec_dot = ggml_vec_dot_mxfp4_q8_0, + .vec_dot_type = GGML_TYPE_Q8_0, + .nrows = 1, + }, + [GGML_TYPE_NVFP4] = { + .from_float = quantize_row_nvfp4, + .vec_dot = ggml_vec_dot_nvfp4_q8_0, + .vec_dot_type = GGML_TYPE_Q8_0, + .nrows = 1, + }, + [GGML_TYPE_Q2_K] = { + .from_float = quantize_row_q2_K, + .vec_dot = ggml_vec_dot_q2_K_q8_K, + .vec_dot_type = GGML_TYPE_Q8_K, + .nrows = 1, + }, + [GGML_TYPE_Q3_K] = { + .from_float = quantize_row_q3_K, + .vec_dot = ggml_vec_dot_q3_K_q8_K, + .vec_dot_type = GGML_TYPE_Q8_K, + .nrows = 1, + }, + [GGML_TYPE_Q4_K] = { + .from_float = quantize_row_q4_K, + .vec_dot = ggml_vec_dot_q4_K_q8_K, + .vec_dot_type = GGML_TYPE_Q8_K, +#if defined (__ARM_FEATURE_MATMUL_INT8) + .nrows = 2, +#else + .nrows = 1, +#endif + }, + [GGML_TYPE_Q5_K] = { + .from_float = quantize_row_q5_K, + .vec_dot = ggml_vec_dot_q5_K_q8_K, + .vec_dot_type = GGML_TYPE_Q8_K, + .nrows = 1, + }, + [GGML_TYPE_Q6_K] = { + .from_float = quantize_row_q6_K, + .vec_dot = ggml_vec_dot_q6_K_q8_K, + .vec_dot_type = GGML_TYPE_Q8_K, +#if defined (__ARM_FEATURE_MATMUL_INT8) + .nrows = 2, +#else + .nrows = 1, +#endif + }, + [GGML_TYPE_IQ2_XXS] = { + .from_float = NULL, + .vec_dot = ggml_vec_dot_iq2_xxs_q8_K, + .vec_dot_type = GGML_TYPE_Q8_K, + .nrows = 1, + }, + [GGML_TYPE_IQ2_XS] = { + .from_float = NULL, + .vec_dot = ggml_vec_dot_iq2_xs_q8_K, + .vec_dot_type = GGML_TYPE_Q8_K, + .nrows = 1, + }, + [GGML_TYPE_IQ3_XXS] = { + // NOTE: from_float for iq3 and iq2_s was removed because these quants require initialization in ggml_quantize_init + //.from_float = quantize_row_iq3_xxs, + .vec_dot = ggml_vec_dot_iq3_xxs_q8_K, + .vec_dot_type = GGML_TYPE_Q8_K, + .nrows = 1, + }, + [GGML_TYPE_IQ3_S] = { + //.from_float = quantize_row_iq3_s, + .vec_dot = ggml_vec_dot_iq3_s_q8_K, + .vec_dot_type = GGML_TYPE_Q8_K, + .nrows = 1, + }, + [GGML_TYPE_IQ2_S] = { + //.from_float = quantize_row_iq2_s, + .vec_dot = ggml_vec_dot_iq2_s_q8_K, + .vec_dot_type = GGML_TYPE_Q8_K, + .nrows = 1, + }, + [GGML_TYPE_IQ1_S] = { + .from_float = NULL, + .vec_dot = ggml_vec_dot_iq1_s_q8_K, + .vec_dot_type = GGML_TYPE_Q8_K, + .nrows = 1, + }, + [GGML_TYPE_IQ1_M] = { + .from_float = NULL, + .vec_dot = ggml_vec_dot_iq1_m_q8_K, + .vec_dot_type = GGML_TYPE_Q8_K, + .nrows = 1, + }, + [GGML_TYPE_IQ4_NL] = { + .from_float = quantize_row_iq4_nl, + .vec_dot = ggml_vec_dot_iq4_nl_q8_0, + .vec_dot_type = GGML_TYPE_Q8_0, + .nrows = 1, + }, + [GGML_TYPE_IQ4_XS] = { + .from_float = quantize_row_iq4_xs, + .vec_dot = ggml_vec_dot_iq4_xs_q8_K, + .vec_dot_type = GGML_TYPE_Q8_K, + .nrows = 1, + }, + [GGML_TYPE_Q8_K] = { + .from_float = quantize_row_q8_K, + }, + [GGML_TYPE_BF16] = { + .from_float = (ggml_from_float_t) ggml_cpu_fp32_to_bf16, + .vec_dot = (ggml_vec_dot_t) ggml_vec_dot_bf16, + .vec_dot_type = GGML_TYPE_BF16, + .nrows = 1, + }, + [GGML_TYPE_TQ1_0] = { + .from_float = quantize_row_tq1_0, + .vec_dot = ggml_vec_dot_tq1_0_q8_K, + .vec_dot_type = GGML_TYPE_Q8_K, + .nrows = 1, + }, + [GGML_TYPE_TQ2_0] = { + .from_float = quantize_row_tq2_0, + .vec_dot = ggml_vec_dot_tq2_0_q8_K, + .vec_dot_type = GGML_TYPE_Q8_K, + .nrows = 1, + }, + [GGML_TYPE_I32] = { + .from_float = (ggml_from_float_t) ggml_cpu_fp32_to_i32, + }, +}; + +const struct ggml_type_traits_cpu * ggml_get_type_traits_cpu(enum ggml_type type) { + return &type_traits_cpu[type]; +} + +// +// Threading defs +// + +typedef pthread_t ggml_thread_t; + +#if defined(_WIN32) + +typedef CONDITION_VARIABLE ggml_cond_t; +typedef SRWLOCK ggml_mutex_t; + +#define ggml_mutex_init(m) InitializeSRWLock(m) +#define ggml_mutex_destroy(m) +#define ggml_mutex_lock(m) AcquireSRWLockExclusive(m) +#define ggml_mutex_unlock(m) ReleaseSRWLockExclusive(m) +#define ggml_mutex_lock_shared(m) AcquireSRWLockShared(m) +#define ggml_mutex_unlock_shared(m) ReleaseSRWLockShared(m) + +#define ggml_cond_init(c) InitializeConditionVariable(c) +#define ggml_cond_destroy(c) +#define ggml_cond_wait(c, m) SleepConditionVariableSRW(c, m, INFINITE, CONDITION_VARIABLE_LOCKMODE_SHARED) +#define ggml_cond_broadcast(c) WakeAllConditionVariable(c) + +#define ggml_thread_create pthread_create +#define ggml_thread_join pthread_join + +#else + +typedef pthread_cond_t ggml_cond_t; +typedef pthread_mutex_t ggml_mutex_t; + +#define ggml_mutex_init(m) pthread_mutex_init(m, NULL) +#define ggml_mutex_destroy(m) pthread_mutex_destroy(m) +#define ggml_mutex_lock(m) pthread_mutex_lock(m) +#define ggml_mutex_unlock(m) pthread_mutex_unlock(m) +#define ggml_mutex_lock_shared(m) pthread_mutex_lock(m) +#define ggml_mutex_unlock_shared(m) pthread_mutex_unlock(m) + +#define ggml_lock_init(x) UNUSED(x) +#define ggml_lock_destroy(x) UNUSED(x) +#if defined(__x86_64__) || (defined(_MSC_VER) && defined(_M_AMD64)) +#define ggml_lock_lock(x) _mm_pause() +#else +#define ggml_lock_lock(x) UNUSED(x) +#endif +#define ggml_lock_unlock(x) UNUSED(x) + +#define GGML_LOCK_INITIALIZER 0 +#define ggml_cond_init(c) pthread_cond_init(c, NULL) +#define ggml_cond_destroy(c) pthread_cond_destroy(c) +#define ggml_cond_wait(c, m) pthread_cond_wait(c, m) +#define ggml_cond_broadcast(c) pthread_cond_broadcast(c) + +#define ggml_thread_create pthread_create +#define ggml_thread_join pthread_join + +#endif + +// Threadpool def +struct ggml_threadpool { + ggml_mutex_t mutex; // mutex for cond.var + ggml_cond_t cond; // cond.var for waiting for new work + + struct ggml_cgraph * cgraph; + struct ggml_cplan * cplan; + + // synchronization primitives + atomic_int n_graph; // updated when there is work to be done (i.e each graph) holds graph and active thread counts. + atomic_int GGML_CACHE_ALIGN n_barrier; + atomic_int GGML_CACHE_ALIGN n_barrier_passed; + atomic_int GGML_CACHE_ALIGN current_chunk; // currently processing chunk during Mat_Mul, shared between all the threads. + + // these are atomic as an annotation for thread-sanitizer + atomic_bool stop; // Used for stopping the threadpool altogether + atomic_bool pause; // Used for pausing the threadpool or individual threads + atomic_int abort; // Used for aborting processing of a graph + + struct ggml_compute_state * workers; // per thread state + int n_threads; // Number of threads in the pool + int32_t prio; // Scheduling priority + uint32_t poll; // Polling level (0 - no polling) + + enum ggml_status ec; +}; + +// Per-thread state +struct ggml_compute_state { +#ifndef GGML_USE_OPENMP + ggml_thread_t thrd; + int last_graph; + bool pending; +#endif + bool cpumask[GGML_MAX_N_THREADS]; + struct ggml_threadpool * threadpool; + int ith; +}; + +// Helpers for polling loops +#if defined(__aarch64__) && ( defined(__clang__) || defined(__GNUC__) ) +static inline void ggml_thread_cpu_relax(void) { + __asm__ volatile("yield" ::: "memory"); +} +#elif defined(__x86_64__) +static inline void ggml_thread_cpu_relax(void) { + _mm_pause(); +} +#elif defined(__riscv) +static inline void ggml_thread_cpu_relax(void) { + #ifdef __riscv_zihintpause + __asm__ __volatile__ ("pause"); + #else + /* Encoding of the pause instruction */ + __asm__ __volatile__ (".4byte 0x100000F"); + #endif +} +#else +static inline void ggml_thread_cpu_relax(void) {;} +#endif + +// +// NUMA support +// + +#define GGML_NUMA_MAX_NODES 8 +#define GGML_NUMA_MAX_CPUS 512 + +struct ggml_numa_node { + uint32_t cpus[GGML_NUMA_MAX_CPUS]; // hardware threads on this node + uint32_t n_cpus; +}; + +struct ggml_numa_nodes { + enum ggml_numa_strategy numa_strategy; + struct ggml_numa_node nodes[GGML_NUMA_MAX_NODES]; + uint32_t n_nodes; + uint32_t total_cpus; // hardware threads on system + uint32_t current_node; // node on which main process is execting +#if defined(__gnu_linux__) + cpu_set_t cpuset; // cpuset from numactl +#else + uint32_t cpuset; // no NUMA support outside of Linux at this time. Use a portable datatype +#endif +}; + +// +// ggml state +// + +struct ggml_state { + struct ggml_numa_nodes numa; +}; + +static struct ggml_state g_state = {0}; + +void ggml_barrier(struct ggml_threadpool * tp) { + int n_threads = atomic_load_explicit(&tp->n_graph, memory_order_relaxed) & GGML_THREADPOOL_N_THREADS_MASK; + if (n_threads == 1) { + return; + } + +#ifdef GGML_USE_OPENMP + #pragma omp barrier +#else + int n_passed = atomic_load_explicit(&tp->n_barrier_passed, memory_order_relaxed); + + // enter barrier (full seq-cst fence) + int n_barrier = atomic_fetch_add_explicit(&tp->n_barrier, 1, memory_order_seq_cst); + + if (n_barrier == (n_threads - 1)) { + // last thread + atomic_store_explicit(&tp->n_barrier, 0, memory_order_relaxed); + + // exit barrier (full seq-cst fence) + atomic_fetch_add_explicit(&tp->n_barrier_passed, 1, memory_order_seq_cst); + return; + } + + // wait for other threads + while (atomic_load_explicit(&tp->n_barrier_passed, memory_order_relaxed) == n_passed) { + ggml_thread_cpu_relax(); + } + + // exit barrier (full seq-cst fence) + // TSAN doesn't support standalone fence yet, we use a dummy read-modify-write instead + #ifdef GGML_TSAN_ENABLED + atomic_fetch_add_explicit(&tp->n_barrier_passed, 0, memory_order_seq_cst); + #else + atomic_thread_fence(memory_order_seq_cst); + #endif +#endif +} + +void ggml_threadpool_chunk_set(struct ggml_threadpool * tp, int value) { + atomic_store_explicit(&tp->current_chunk, value, memory_order_relaxed); +} + +int ggml_threadpool_chunk_add(struct ggml_threadpool * tp, int value) { + return atomic_fetch_add_explicit(&tp->current_chunk, value, memory_order_relaxed); +} + +#if defined(__gnu_linux__) +static cpu_set_t ggml_get_numa_affinity(void) { + cpu_set_t cpuset; + pthread_t thread; + thread = pthread_self(); + CPU_ZERO(&cpuset); + pthread_getaffinity_np(thread, sizeof(cpu_set_t), &cpuset); + return cpuset; +} +#else +static uint32_t ggml_get_numa_affinity(void) { + return 0; // no NUMA support +} +#endif + +void ggml_numa_init(enum ggml_numa_strategy numa_flag) { + if (g_state.numa.n_nodes > 0) { + fprintf(stderr, "ggml_numa_init: NUMA already initialized\n"); + + return; + } + +#if defined(__gnu_linux__) + struct stat st; + char path[256]; + int rv; + + // set numa scheme + g_state.numa.numa_strategy = numa_flag; + + GGML_PRINT_DEBUG("numa strategy %u\n",g_state.numa.numa_strategy); + + g_state.numa.cpuset = ggml_get_numa_affinity(); + + // enumerate nodes + while (g_state.numa.n_nodes < GGML_NUMA_MAX_NODES) { + rv = snprintf(path, sizeof(path), "/sys/devices/system/node/node%u", g_state.numa.n_nodes); + GGML_ASSERT(rv > 0 && (unsigned)rv < sizeof(path)); + if (stat(path, &st) != 0) { break; } + ++g_state.numa.n_nodes; + } + + // enumerate CPUs + while (g_state.numa.total_cpus < GGML_NUMA_MAX_CPUS) { + rv = snprintf(path, sizeof(path), "/sys/devices/system/cpu/cpu%u", g_state.numa.total_cpus); + GGML_ASSERT(rv > 0 && (unsigned)rv < sizeof(path)); + if (stat(path, &st) != 0) { break; } + ++g_state.numa.total_cpus; + } + + GGML_PRINT_DEBUG("found %u numa nodes, %u CPUs\n", g_state.numa.n_nodes, g_state.numa.total_cpus); + + // figure out which node we're on + uint current_cpu; + int getcpu_ret = 0; +#if __GLIBC__ > 2 || (__GLIBC__ == 2 && __GLIBC_MINOR__ > 33) || defined(__COSMOPOLITAN__) + getcpu_ret = getcpu(¤t_cpu, &g_state.numa.current_node); +#else + // old glibc doesn't have a wrapper for this call. Fall back on direct syscall +# if !defined(SYS_getcpu) && defined(SYS_get_cpu) +# define SYS_getcpu SYS_get_cpu // some older glibc versions use this name +# endif + getcpu_ret = syscall(SYS_getcpu, ¤t_cpu, &g_state.numa.current_node); +#endif + + if (g_state.numa.n_nodes < 1 || g_state.numa.total_cpus < 1 || getcpu_ret != 0) { + g_state.numa.n_nodes = 0; + return; + } + + GGML_PRINT_DEBUG("found our process on numa node %u, CPU %u\n", g_state.numa.current_node, current_cpu); + + for (uint32_t n = 0; n < g_state.numa.n_nodes; ++n) { + struct ggml_numa_node * node = &g_state.numa.nodes[n]; + GGML_PRINT_DEBUG("CPUs on node %u:", n); + node->n_cpus = 0; + for (uint32_t c = 0; c < g_state.numa.total_cpus; ++c) { + rv = snprintf(path, sizeof(path), "/sys/devices/system/node/node%u/cpu%u", n, c); + GGML_ASSERT(rv > 0 && (unsigned)rv < sizeof(path)); + if (stat(path, &st) == 0) { + node->cpus[node->n_cpus++] = c; + GGML_PRINT_DEBUG(" %u", c); + } + } + GGML_PRINT_DEBUG("\n"); + } + + if (ggml_is_numa()) { + FILE *fptr = fopen("/proc/sys/kernel/numa_balancing", "r"); + if (fptr != NULL) { + char buf[42]; + if (fgets(buf, sizeof(buf), fptr) && strncmp(buf, "0\n", sizeof(buf)) != 0) { + GGML_LOG_WARN("/proc/sys/kernel/numa_balancing is enabled, this has been observed to impair performance\n"); + } + fclose(fptr); + } + } +#else + UNUSED(numa_flag); + // TODO +#endif +} + +bool ggml_is_numa(void) { + return g_state.numa.n_nodes > 1; +} + +#if defined(__ARM_ARCH) +#if defined(__aarch64__) && defined(__ARM_FEATURE_SVE) +#include +static void ggml_init_arm_arch_features(void) { + ggml_arm_arch_features.sve_cnt = svcntb(); +} +#else +static void ggml_init_arm_arch_features(void) {} +#endif +#endif // __ARM_ARCH + +#if defined(__riscv) && defined(__riscv_v_intrinsic) +#include +static void ggml_init_riscv_arch_features(void) { + ggml_riscv_arch_features.rvv_vlen = __riscv_vlenb(); +} +#else +static void ggml_init_riscv_arch_features(void) {} +#endif + +struct ggml_tensor * ggml_new_i32(struct ggml_context * ctx, int32_t value) { + GGML_ASSERT(!ggml_get_no_alloc(ctx)); + + struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); + + ggml_set_i32(result, value); + + return result; +} + +struct ggml_tensor * ggml_new_f32(struct ggml_context * ctx, float value) { + GGML_ASSERT(!ggml_get_no_alloc(ctx)); + + struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); + + ggml_set_f32(result, value); + + return result; +} + +struct ggml_tensor * ggml_set_i32 (struct ggml_tensor * tensor, int32_t value) { + const int n = ggml_nrows(tensor); + const int nc = tensor->ne[0]; + const size_t n1 = tensor->nb[1]; + + char * const data = tensor->data; + + switch (tensor->type) { + case GGML_TYPE_I8: + { + assert(tensor->nb[0] == sizeof(int8_t)); + for (int i = 0; i < n; i++) { + ggml_vec_set_i8(nc, (int8_t *)(data + i*n1), value); + } + } break; + case GGML_TYPE_I16: + { + assert(tensor->nb[0] == sizeof(int16_t)); + for (int i = 0; i < n; i++) { + ggml_vec_set_i16(nc, (int16_t *)(data + i*n1), value); + } + } break; + case GGML_TYPE_I32: + { + assert(tensor->nb[0] == sizeof(int32_t)); + for (int i = 0; i < n; i++) { + ggml_vec_set_i32(nc, (int32_t *)(data + i*n1), value); + } + } break; + case GGML_TYPE_F16: + { + assert(tensor->nb[0] == sizeof(ggml_fp16_t)); + for (int i = 0; i < n; i++) { + ggml_vec_set_f16(nc, (ggml_fp16_t *)(data + i*n1), GGML_CPU_FP32_TO_FP16(value)); + } + } break; + case GGML_TYPE_BF16: + { + assert(tensor->nb[0] == sizeof(ggml_fp16_t)); + for (int i = 0; i < n; i++) { + ggml_vec_set_bf16(nc, (ggml_bf16_t *)(data + i*n1), GGML_FP32_TO_BF16(value)); + } + } break; + case GGML_TYPE_F32: + { + assert(tensor->nb[0] == sizeof(float)); + for (int i = 0; i < n; i++) { + ggml_vec_set_f32(nc, (float *)(data + i*n1), value); + } + } break; + default: + { + GGML_ABORT("fatal error"); + } + } + + return tensor; +} + +struct ggml_tensor * ggml_set_f32(struct ggml_tensor * tensor, float value) { + const int n = ggml_nrows(tensor); + const int nc = tensor->ne[0]; + const size_t n1 = tensor->nb[1]; + + char * const data = tensor->data; + + switch (tensor->type) { + case GGML_TYPE_I8: + { + assert(tensor->nb[0] == sizeof(int8_t)); + for (int i = 0; i < n; i++) { + ggml_vec_set_i8(nc, (int8_t *)(data + i*n1), value); + } + } break; + case GGML_TYPE_I16: + { + assert(tensor->nb[0] == sizeof(int16_t)); + for (int i = 0; i < n; i++) { + ggml_vec_set_i16(nc, (int16_t *)(data + i*n1), value); + } + } break; + case GGML_TYPE_I32: + { + assert(tensor->nb[0] == sizeof(int32_t)); + for (int i = 0; i < n; i++) { + ggml_vec_set_i32(nc, (int32_t *)(data + i*n1), value); + } + } break; + case GGML_TYPE_F16: + { + assert(tensor->nb[0] == sizeof(ggml_fp16_t)); + for (int i = 0; i < n; i++) { + ggml_vec_set_f16(nc, (ggml_fp16_t *)(data + i*n1), GGML_CPU_FP32_TO_FP16(value)); + } + } break; + case GGML_TYPE_BF16: + { + assert(tensor->nb[0] == sizeof(ggml_bf16_t)); + for (int i = 0; i < n; i++) { + ggml_vec_set_bf16(nc, (ggml_bf16_t *)(data + i*n1), GGML_FP32_TO_BF16(value)); + } + } break; + case GGML_TYPE_F32: + { + assert(tensor->nb[0] == sizeof(float)); + for (int i = 0; i < n; i++) { + ggml_vec_set_f32(nc, (float *)(data + i*n1), value); + } + } break; + default: + { + GGML_ABORT("fatal error"); + } + } + + return tensor; +} + +int32_t ggml_get_i32_1d(const struct ggml_tensor * tensor, int i) { + if (!ggml_is_contiguous(tensor)) { + int64_t id[4] = { 0, 0, 0, 0 }; + ggml_unravel_index(tensor, i, &id[0], &id[1], &id[2], &id[3]); + return ggml_get_i32_nd(tensor, id[0], id[1], id[2], id[3]); + } + switch (tensor->type) { + case GGML_TYPE_I8: + { + GGML_ASSERT(tensor->nb[0] == sizeof(int8_t)); + return ((int8_t *)(tensor->data))[i]; + } + case GGML_TYPE_I16: + { + GGML_ASSERT(tensor->nb[0] == sizeof(int16_t)); + return ((int16_t *)(tensor->data))[i]; + } + case GGML_TYPE_I32: + { + GGML_ASSERT(tensor->nb[0] == sizeof(int32_t)); + return ((int32_t *)(tensor->data))[i]; + } + case GGML_TYPE_F16: + { + GGML_ASSERT(tensor->nb[0] == sizeof(ggml_fp16_t)); + return GGML_CPU_FP16_TO_FP32(((ggml_fp16_t *)(tensor->data))[i]); + } + case GGML_TYPE_BF16: + { + GGML_ASSERT(tensor->nb[0] == sizeof(ggml_bf16_t)); + return GGML_BF16_TO_FP32(((ggml_bf16_t *)(tensor->data))[i]); + } + case GGML_TYPE_F32: + { + GGML_ASSERT(tensor->nb[0] == sizeof(float)); + return ((float *)(tensor->data))[i]; + } + default: + { + GGML_ABORT("fatal error"); + } + } +} + +void ggml_set_i32_1d(const struct ggml_tensor * tensor, int i, int32_t value) { + if (!ggml_is_contiguous(tensor)) { + int64_t id[4] = { 0, 0, 0, 0 }; + ggml_unravel_index(tensor, i, &id[0], &id[1], &id[2], &id[3]); + ggml_set_i32_nd(tensor, id[0], id[1], id[2], id[3], value); + return; + } + switch (tensor->type) { + case GGML_TYPE_I8: + { + GGML_ASSERT(tensor->nb[0] == sizeof(int8_t)); + ((int8_t *)(tensor->data))[i] = value; + } break; + case GGML_TYPE_I16: + { + GGML_ASSERT(tensor->nb[0] == sizeof(int16_t)); + ((int16_t *)(tensor->data))[i] = value; + } break; + case GGML_TYPE_I32: + { + GGML_ASSERT(tensor->nb[0] == sizeof(int32_t)); + ((int32_t *)(tensor->data))[i] = value; + } break; + case GGML_TYPE_F16: + { + GGML_ASSERT(tensor->nb[0] == sizeof(ggml_fp16_t)); + ((ggml_fp16_t *)(tensor->data))[i] = GGML_CPU_FP32_TO_FP16(value); + } break; + case GGML_TYPE_BF16: + { + GGML_ASSERT(tensor->nb[0] == sizeof(ggml_bf16_t)); + ((ggml_bf16_t *)(tensor->data))[i] = GGML_FP32_TO_BF16(value); + } break; + case GGML_TYPE_F32: + { + GGML_ASSERT(tensor->nb[0] == sizeof(float)); + ((float *)(tensor->data))[i] = value; + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} + +int32_t ggml_get_i32_nd(const struct ggml_tensor * tensor, int i0, int i1, int i2, int i3) { + void * data = (char *) tensor->data + i0*tensor->nb[0] + i1*tensor->nb[1] + i2*tensor->nb[2] + i3*tensor->nb[3]; + switch (tensor->type) { + case GGML_TYPE_I8: + return ((int8_t *) data)[0]; + case GGML_TYPE_I16: + return ((int16_t *) data)[0]; + case GGML_TYPE_I32: + return ((int32_t *) data)[0]; + case GGML_TYPE_F16: + return GGML_CPU_FP16_TO_FP32(((ggml_fp16_t *) data)[0]); + case GGML_TYPE_BF16: + return GGML_BF16_TO_FP32(((ggml_bf16_t *) data)[0]); + case GGML_TYPE_F32: + return ((float *) data)[0]; + default: + GGML_ABORT("fatal error"); + } +} + +void ggml_set_i32_nd(const struct ggml_tensor * tensor, int i0, int i1, int i2, int i3, int32_t value) { + void * data = (char *) tensor->data + i0*tensor->nb[0] + i1*tensor->nb[1] + i2*tensor->nb[2] + i3*tensor->nb[3]; + switch (tensor->type) { + case GGML_TYPE_I8: + { + ((int8_t *)(data))[0] = value; + } break; + case GGML_TYPE_I16: + { + ((int16_t *)(data))[0] = value; + } break; + case GGML_TYPE_I32: + { + ((int32_t *)(data))[0] = value; + } break; + case GGML_TYPE_F16: + { + ((ggml_fp16_t *)(data))[0] = GGML_CPU_FP32_TO_FP16(value); + } break; + case GGML_TYPE_BF16: + { + ((ggml_bf16_t *)(data))[0] = GGML_FP32_TO_BF16(value); + } break; + case GGML_TYPE_F32: + { + ((float *)(data))[0] = value; + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} + +float ggml_get_f32_1d(const struct ggml_tensor * tensor, int i) { + if (!ggml_is_contiguous(tensor)) { + int64_t id[4] = { 0, 0, 0, 0 }; + ggml_unravel_index(tensor, i, &id[0], &id[1], &id[2], &id[3]); + return ggml_get_f32_nd(tensor, id[0], id[1], id[2], id[3]); + } + switch (tensor->type) { + case GGML_TYPE_I8: + { + return ((int8_t *)(tensor->data))[i]; + } + case GGML_TYPE_I16: + { + return ((int16_t *)(tensor->data))[i]; + } + case GGML_TYPE_I32: + { + return ((int32_t *)(tensor->data))[i]; + } + case GGML_TYPE_F16: + { + return GGML_CPU_FP16_TO_FP32(((ggml_fp16_t *)(tensor->data))[i]); + } + case GGML_TYPE_BF16: + { + return GGML_BF16_TO_FP32(((ggml_bf16_t *)(tensor->data))[i]); + } + case GGML_TYPE_F32: + { + return ((float *)(tensor->data))[i]; + } + default: + { + GGML_ABORT("fatal error"); + } + } +} + +void ggml_set_f32_1d(const struct ggml_tensor * tensor, int i, float value) { + if (!ggml_is_contiguous(tensor)) { + int64_t id[4] = { 0, 0, 0, 0 }; + ggml_unravel_index(tensor, i, &id[0], &id[1], &id[2], &id[3]); + ggml_set_f32_nd(tensor, id[0], id[1], id[2], id[3], value); + return; + } + switch (tensor->type) { + case GGML_TYPE_I8: + { + ((int8_t *)(tensor->data))[i] = value; + } break; + case GGML_TYPE_I16: + { + ((int16_t *)(tensor->data))[i] = value; + } break; + case GGML_TYPE_I32: + { + ((int32_t *)(tensor->data))[i] = value; + } break; + case GGML_TYPE_F16: + { + ((ggml_fp16_t *)(tensor->data))[i] = GGML_CPU_FP32_TO_FP16(value); + } break; + case GGML_TYPE_BF16: + { + ((ggml_bf16_t *)(tensor->data))[i] = GGML_FP32_TO_BF16(value); + } break; + case GGML_TYPE_F32: + { + ((float *)(tensor->data))[i] = value; + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} + +float ggml_get_f32_nd(const struct ggml_tensor * tensor, int i0, int i1, int i2, int i3) { + void * data = (char *) tensor->data + i0*tensor->nb[0] + i1*tensor->nb[1] + i2*tensor->nb[2] + i3*tensor->nb[3]; + switch (tensor->type) { + case GGML_TYPE_I8: + return ((int8_t *) data)[0]; + case GGML_TYPE_I16: + return ((int16_t *) data)[0]; + case GGML_TYPE_I32: + return ((int32_t *) data)[0]; + case GGML_TYPE_F16: + return GGML_CPU_FP16_TO_FP32(((ggml_fp16_t *) data)[0]); + case GGML_TYPE_BF16: + return GGML_BF16_TO_FP32(((ggml_bf16_t *) data)[0]); + case GGML_TYPE_F32: + return ((float *) data)[0]; + default: + GGML_ABORT("fatal error"); + } +} + +void ggml_set_f32_nd(const struct ggml_tensor * tensor, int i0, int i1, int i2, int i3, float value) { + void * data = (char *) tensor->data + i0*tensor->nb[0] + i1*tensor->nb[1] + i2*tensor->nb[2] + i3*tensor->nb[3]; + switch (tensor->type) { + case GGML_TYPE_I8: + { + ((int8_t *)(data))[0] = value; + } break; + case GGML_TYPE_I16: + { + ((int16_t *)(data))[0] = value; + } break; + case GGML_TYPE_I32: + { + ((int32_t *)(data))[0] = value; + } break; + case GGML_TYPE_F16: + { + ((ggml_fp16_t *)(data))[0] = GGML_CPU_FP32_TO_FP16(value); + } break; + case GGML_TYPE_BF16: + { + ((ggml_bf16_t *)(data))[0] = GGML_FP32_TO_BF16(value); + } break; + case GGML_TYPE_F32: + { + ((float *)(data))[0] = value; + } break; + default: + { + GGML_ABORT("fatal error"); + } + } +} + +//////////////////////////////////////////////////////////////////////////////// + +// ggml_compute_forward_mul_mat + +static void ggml_compute_forward_mul_mat_one_chunk( + const struct ggml_compute_params * params, + struct ggml_tensor * dst, + const enum ggml_type type, + const int64_t num_rows_per_vec_dot, + const int64_t ir0_start, + const int64_t ir0_end, + const int64_t ir1_start, + const int64_t ir1_end) { + + const struct ggml_tensor * src0 = dst->src[0]; + const struct ggml_tensor * src1 = dst->src[1]; + + GGML_TENSOR_BINARY_OP_LOCALS + + const bool src1_cont = ggml_is_contiguous(src1); + + ggml_vec_dot_t const vec_dot = type_traits_cpu[type].vec_dot; + enum ggml_type const vec_dot_type = type_traits_cpu[type].vec_dot_type; + + // broadcast factors + const int64_t r2 = ne12 / ne02; + const int64_t r3 = ne13 / ne03; + + //printf("ir0_start = %6lld, ir0_end = %6lld, ir1_start = %6lld, ir1_end = %6lld\n", ir0_start, ir0_end, ir1_start, ir1_end); + + // threads with no work simply yield (not sure if it helps) + if (ir0_start >= ir0_end || ir1_start >= ir1_end) { + return; + } + + const void * wdata = (src1->type == vec_dot_type) ? src1->data : params->wdata; + const size_t row_size = ggml_row_size(vec_dot_type, ne10); + + assert(ne12 % ne02 == 0); + assert(ne13 % ne03 == 0); + + // block-tiling attempt + const int64_t blck_0 = 16; + const int64_t blck_1 = 16; + + const size_t src1_col_stride = src1_cont || src1->type != vec_dot_type ? row_size : nb11; + + // attempt to reduce false-sharing (does not seem to make a difference) + // 16 * 2, accounting for mmla kernels + float tmp[32]; + + for (int64_t iir1 = ir1_start; iir1 < ir1_end; iir1 += blck_1) { + for (int64_t iir0 = ir0_start; iir0 < ir0_end; iir0 += blck_0) { + for (int64_t ir1 = iir1; ir1 < iir1 + blck_1 && ir1 < ir1_end; ir1 += num_rows_per_vec_dot) { + const int64_t i13 = (ir1 / (ne12 * ne1)); + const int64_t i12 = (ir1 - i13 * ne12 * ne1) / ne1; + const int64_t i11 = (ir1 - i13 * ne12 * ne1 - i12 * ne1); + + // broadcast src0 into src1 + const int64_t i03 = i13 / r3; + const int64_t i02 = i12 / r2; + + const int64_t i1 = i11; + const int64_t i2 = i12; + const int64_t i3 = i13; + + const char * src0_row = (const char*)src0->data + (0 + i02 * nb02 + i03 * nb03); + + // desc: when src1 is not a contiguous memory block we have to calculate the offset using the strides + // if it is, then we have either copied the data to params->wdata and made it contiguous or we are using + // the original src1 data pointer, so we should index using the indices directly + // TODO: this is a bit of a hack, we should probably have a better way to handle this + const char * src1_col = (const char*)wdata + + (src1_cont || src1->type != vec_dot_type + ? (i11 + i12 * ne11 + i13 * ne12 * ne11) * row_size + : (i11 * nb11 + i12 * nb12 + i13 * nb13)); + float * dst_col = (float*)((char*)dst->data + (i1 * nb1 + i2 * nb2 + i3 * nb3)); + + //for (int64_t ir0 = iir0; ir0 < iir0 + blck_0 && ir0 < ir0_end; ++ir0) { + // vec_dot(ne00, &dst_col[ir0], src0_row + ir0*nb01, src1_col); + //} + + for (int64_t ir0 = iir0; ir0 < iir0 + blck_0 && ir0 < ir0_end; ir0 += num_rows_per_vec_dot) { + vec_dot(ne00, &tmp[ir0 - iir0], (num_rows_per_vec_dot > 1 ? 16 : 0), src0_row + ir0 * nb01, (num_rows_per_vec_dot > 1 ? nb01 : 0), src1_col, (num_rows_per_vec_dot > 1 ? src1_col_stride : 0), num_rows_per_vec_dot); + } + + for (int cn = 0; cn < num_rows_per_vec_dot; ++cn) { + memcpy(&dst_col[iir0 + cn * nb1 / nb0], tmp + (cn * 16), (MIN(iir0 + blck_0, ir0_end) - iir0) * sizeof(float)); + } + } + } + } +} + +void ggml_compute_forward_mul_mat( + const struct ggml_compute_params * params, + struct ggml_tensor * dst) { + + const struct ggml_tensor * src0 = dst->src[0]; + const struct ggml_tensor * src1 = dst->src[1]; + + GGML_TENSOR_BINARY_OP_LOCALS + + const int ith = params->ith; + const int nth = params->nth; + + enum ggml_type const vec_dot_type = type_traits_cpu[src0->type].vec_dot_type; + ggml_from_float_t const from_float = type_traits_cpu[vec_dot_type].from_float; + int64_t const vec_dot_num_rows = type_traits_cpu[src0->type].nrows; + + GGML_ASSERT(ne0 == ne01); + GGML_ASSERT(ne1 == ne11); + GGML_ASSERT(ne2 == ne12); + GGML_ASSERT(ne3 == ne13); + + // we don't support permuted src0 or src1 + GGML_ASSERT(nb00 == ggml_type_size(src0->type)); + GGML_ASSERT(nb10 == ggml_type_size(src1->type)); + + // dst cannot be transposed or permuted + GGML_ASSERT(nb0 == sizeof(float)); + GGML_ASSERT(nb0 <= nb1); + GGML_ASSERT(nb1 <= nb2); + GGML_ASSERT(nb2 <= nb3); + + // nb01 >= nb00 - src0 is not transposed + // compute by src0 rows + + // TODO: extract to "extra_op" +#if GGML_USE_LLAMAFILE + // broadcast factors + const int64_t r2 = ne12 / ne02; + const int64_t r3 = ne13 / ne03; + + const bool src1_cont = ggml_is_contiguous(src1); + + if (src1_cont) { + for (int64_t i13 = 0; i13 < ne13; i13++) + for (int64_t i12 = 0; i12 < ne12; i12++) + if (!llamafile_sgemm(params, + ne01, ne11, ne00/ggml_blck_size(src0->type), + (const char *)src0->data + i12/r2*nb02 + i13/r3*nb03, + nb01/ggml_type_size(src0->type), + (const char *)src1->data + i12*nb12 + i13*nb13, + nb11/ggml_type_size(src1->type), + (char *)dst->data + i12*nb2 + i13*nb3, + nb1/ggml_type_size(dst->type), + src0->type, + src1->type, + dst->type)) + goto UseGgmlGemm1; + return; + } +UseGgmlGemm1:; +#endif + + if (src1->type != vec_dot_type) { + char * wdata = params->wdata; + + const size_t nbw0 = ggml_type_size(vec_dot_type); + const size_t nbw1 = ggml_row_size(vec_dot_type, ne10); + const size_t nbw2 = nbw1*ne11; + const size_t nbw3 = nbw2*ne12; + + assert(params->wsize >= ne13*nbw3); + GGML_ASSERT(src1->type == GGML_TYPE_F32); + + #if 0 + for (int64_t i13 = 0; i13 < ne13; ++i13) { + for (int64_t i12 = 0; i12 < ne12; ++i12) { + for (int64_t i11 = ith; i11 < ne11; i11 += nth) { + from_float((float *)((char *) src1->data + i13*nb13 + i12*nb12 + i11*nb11), + (void *) (wdata + i13*nbw3 + i12*nbw2 + i11*nbw1), + ne10); + } + } + } + #else + for (int64_t i13 = 0; i13 < ne13; ++i13) { + for (int64_t i12 = 0; i12 < ne12; ++i12) { + for (int64_t i11 = 0; i11 < ne11; ++i11) { + size_t bs = ggml_blck_size(vec_dot_type); + int64_t ne10_block_start = (ith * ne10/bs) / nth; + int64_t ne10_block_end = ((ith + 1) * ne10/bs) / nth; + from_float((float *)((char *) src1->data + i13*nb13 + i12*nb12 + i11*nb11 + ne10_block_start*bs*nb10), + (void *) (wdata + i13*nbw3 + i12*nbw2 + i11*nbw1 + ne10_block_start*nbw0), + (ne10_block_end - ne10_block_start) * bs); + } + } + } + #endif + } + + if (ith == 0) { + // Every thread starts at ith, so the first unprocessed chunk is nth. This save a bit of coordination right at the start. + atomic_store_explicit(¶ms->threadpool->current_chunk, nth, memory_order_relaxed); + } + + ggml_barrier(params->threadpool); + +#if GGML_USE_LLAMAFILE + if (src1->type != vec_dot_type) { + const void* wdata = (src1->type == vec_dot_type) ? src1->data : params->wdata; + const size_t row_size = ggml_row_size(vec_dot_type, ne10); + + for (int64_t i13 = 0; i13 < ne13; i13++) + for (int64_t i12 = 0; i12 < ne12; i12++) + if (!llamafile_sgemm(params, + ne01, ne11, ne00/ggml_blck_size(src0->type), + (const char *)src0->data + i12/r2*nb02 + i13/r3*nb03, + nb01/ggml_type_size(src0->type), + (const char *)wdata + (i12*ne11 + i13*ne12*ne11)*row_size, + row_size/ggml_type_size(vec_dot_type), + (char *)dst->data + i12*nb2 + i13*nb3, + nb1/ggml_type_size(dst->type), + src0->type, + vec_dot_type, + dst->type)) + goto UseGgmlGemm2; + return; + } +UseGgmlGemm2:; +#endif + + // This is the size of the first dimension of the result, so we can iterate that way. (see the ASSERT above, these are the same numbers) + const int64_t nr0 = ne0; + + // This is the size of the rest of the dimensions of the result + const int64_t nr1 = ne1 * ne2 * ne3; + + // Now select a reasonable chunk size. + int chunk_size = 16; + + // We need to step up the size if it's small + if (nr0 == 1 || nr1 == 1) { + chunk_size = 64; + } + + // distribute the work across the inner or outer loop based on which one is larger + // The number of chunks in the 0/1 dim. + // CEIL(nr0/chunk_size) + int64_t nchunk0 = (nr0 + chunk_size - 1) / chunk_size; + int64_t nchunk1 = (nr1 + chunk_size - 1) / chunk_size; + + // If the chunking is poor for the number of threads on this setup, scrap the whole plan. Re-chunk it by thread. + // Also, chunking by thread was measured to have perform better on NUMA systems. See https://github.com/ggml-org/llama.cpp/pull/6915 + // In theory, chunking should be just as useful on NUMA and non NUMA systems, but testing disagreed with that. + if (nchunk0 * nchunk1 < nth * 4 || ggml_is_numa()) { + // distribute the thread work across the inner or outer loop based on which one is larger + nchunk0 = nr0 > nr1 ? nth : 1; // parallelize by src0 rows + nchunk1 = nr0 > nr1 ? 1 : nth; // parallelize by src1 rows + } + + // The number of elements in each chunk + const int64_t dr0 = (nr0 + nchunk0 - 1) / nchunk0; + const int64_t dr1 = (nr1 + nchunk1 - 1) / nchunk1; + + // The first chunk comes from our thread_id, the rest will get auto-assigned. + int current_chunk = ith; + + while (current_chunk < nchunk0 * nchunk1) { + const int64_t ith0 = current_chunk % nchunk0; + const int64_t ith1 = current_chunk / nchunk0; + + const int64_t ir0_start = dr0 * ith0; + const int64_t ir0_end = MIN(ir0_start + dr0, nr0); + + const int64_t ir1_start = dr1 * ith1; + const int64_t ir1_end = MIN(ir1_start + dr1, nr1); + + // dot kernels can handle 1 row and col at a time, but mmla kernels can process 2 rows and cols + int64_t num_rows_per_vec_dot = vec_dot_num_rows; + + // these checks are needed to avoid crossing dim1 boundaries + // can be optimized, but the logic would become more complicated, so keeping it like this for simplicity + if ((nr0 % 2 != 0) || (ne11 % 2 != 0) || ((ir0_end - ir0_start) % 2 != 0) || ((ir1_end - ir1_start) % 2 != 0)) { + num_rows_per_vec_dot = 1; + } + ggml_compute_forward_mul_mat_one_chunk(params, dst, src0->type, num_rows_per_vec_dot, ir0_start, ir0_end, ir1_start, ir1_end); + + if (nth >= nchunk0 * nchunk1) { + break; + } + + current_chunk = atomic_fetch_add_explicit(¶ms->threadpool->current_chunk, 1, memory_order_relaxed); + } +} + +// ggml_compute_forward_mul_mat_id + +#define MMID_MATRIX_ROW(row_id, i1) matrix_rows[(row_id)*ids->ne[0]*ids->ne[1] + (i1)] + +struct mmid_row_mapping { + int32_t i1; + int32_t i2; +}; + +static void ggml_compute_forward_mul_mat_id_one_chunk( + struct ggml_tensor * dst, + const struct ggml_tensor * src0, + const struct ggml_tensor * src1, + const struct ggml_tensor * ids, + const int64_t cur_a, + const int64_t ir0_start, + const int64_t ir0_end, + const int64_t ir1_start, + const int64_t ir1_end, + const char * src0_cur, + const struct mmid_row_mapping * matrix_rows, + const size_t row_size, + const bool src1_cont, + const void * wdata) { + + GGML_TENSOR_BINARY_OP_LOCALS + + const enum ggml_type type = src0->type; + + ggml_vec_dot_t const vec_dot = type_traits_cpu[type].vec_dot; + enum ggml_type const vec_dot_type = type_traits_cpu[type].vec_dot_type; + + const int64_t blck_0 = 16; + const int64_t blck_1 = 16; + + float tmp[16]; + + for (int64_t iir1 = ir1_start; iir1 < ir1_end; iir1 += blck_1) { + for (int64_t iir0 = ir0_start; iir0 < ir0_end; iir0 += blck_0) { + for (int64_t ir1 = iir1; ir1 < iir1 + blck_1 && ir1 < ir1_end; ++ir1) { + const int64_t _i12 = ir1; // logical row index for this expert + + struct mmid_row_mapping row_mapping = MMID_MATRIX_ROW(cur_a, _i12); + const int id = row_mapping.i1; // selected expert index + + const int64_t i11 = id % ne11; + const int64_t i12 = row_mapping.i2; // row index in src1 + + const int64_t i1 = id; // selected expert index + const int64_t i2 = i12; // row + + // desc: when src1 is not a contiguous memory block we have to calculate the offset using the strides + // if it is, then we have either copied the data to params->wdata and made it contiguous or we are using + // the original src1 data pointer, so we should index using the indices directly + // TODO: this is a bit of a hack, we should probably have a better way to handle this + const char * src1_col = (const char *) wdata + + (src1_cont || src1->type != vec_dot_type + ? (i11 + i12*ne11)*row_size + : (i11*nb11 + i12*nb12)); + + float * dst_col = (float *) ((char *) dst->data + (i1*nb1 + i2*nb2)); + + for (int64_t ir0 = iir0; ir0 < iir0 + blck_0 && ir0 < ir0_end; ++ir0) { + vec_dot(ne00, &tmp[ir0 - iir0], 0, src0_cur + ir0*nb01, 0, src1_col, 0, 1); + } + + memcpy(&dst_col[iir0], tmp, (MIN(iir0 + blck_0, ir0_end) - iir0)*sizeof(float)); + } + } + } +} + +static void * incr_ptr_aligned(void ** p, size_t size, size_t align) { + + void * ptr = *p; + ptr = (void *) GGML_PAD((uintptr_t) ptr, align); + *p = (void *) ((char *) ptr + size); + return ptr; +} + +static void ggml_compute_forward_mul_mat_id( + const struct ggml_compute_params * params, + struct ggml_tensor * dst) { + + const struct ggml_tensor * src0 = dst->src[0]; + const struct ggml_tensor * src1 = dst->src[1]; + const struct ggml_tensor * ids = dst->src[2]; + + GGML_TENSOR_BINARY_OP_LOCALS + + const int ith = params->ith; + const int nth = params->nth; + + const enum ggml_type type = src0->type; + + const bool src1_cont = ggml_is_contiguous(src1); + + enum ggml_type const vec_dot_type = type_traits_cpu[type].vec_dot_type; + ggml_from_float_t const from_float = type_traits_cpu[vec_dot_type].from_float; + + // we don't support permuted src0 or src1 + GGML_ASSERT(nb00 == ggml_type_size(type)); + GGML_ASSERT(nb10 == ggml_type_size(src1->type)); + + // dst cannot be transposed or permuted + GGML_ASSERT(nb0 == sizeof(float)); + GGML_ASSERT(nb0 <= nb1); + GGML_ASSERT(nb1 <= nb2); + GGML_ASSERT(nb2 <= nb3); + + // row groups + const int n_ids = ids->ne[0]; // n_expert_used + const int n_as = ne02; // n_expert + + void * wdata_cur = params->wdata; + + if (src1->type != vec_dot_type) { + incr_ptr_aligned(&wdata_cur, ggml_row_size(vec_dot_type, ggml_nelements(src1)), sizeof(int64_t)); + } + + int64_t * matrix_row_counts = // [n_as] + incr_ptr_aligned(&wdata_cur, n_as*sizeof(int64_t), sizeof(int64_t)); + + struct mmid_row_mapping * matrix_rows = // [n_as][ids->ne[0]*ids->ne[1]] + incr_ptr_aligned(&wdata_cur, n_as*ids->ne[0]*ids->ne[1]*sizeof(struct mmid_row_mapping), sizeof(int64_t)); + + char (*atomic_current_chunk)[CACHE_LINE_SIZE] = // [n_as] + incr_ptr_aligned(&wdata_cur, CACHE_LINE_SIZE * n_as, CACHE_LINE_SIZE); + + GGML_ASSERT(params->wsize >= (size_t)((char *) wdata_cur - (char *) params->wdata)); + + if (src1->type != vec_dot_type) { + char * wdata = params->wdata; + + const size_t nbw0 = ggml_type_size(vec_dot_type); + const size_t nbw1 = ggml_row_size(vec_dot_type, ne10); + const size_t nbw2 = nbw1*ne11; + const size_t nbw3 = nbw2*ne12; + + assert(params->wsize >= ne13*nbw3); + GGML_ASSERT(src1->type == GGML_TYPE_F32); + +#if 0 + for (int64_t i13 = 0; i13 < ne13; ++i13) { + for (int64_t i12 = ith; i12 < ne12; i12 += nth) { + for (int64_t i11 = 0; i11 < ne11; ++i11) { + from_float((float *)((char *) src1->data + i13*nb13 + i12*nb12 + i11*nb11), + (void *) (wdata + i13*nbw3 + i12*nbw2 + i11*nbw1), + ne10); + } + } + } +#else + for (int64_t i13 = 0; i13 < ne13; ++i13) { + for (int64_t i12 = 0; i12 < ne12; ++i12) { + for (int64_t i11 = 0; i11 < ne11; ++i11) { + size_t bs = ggml_blck_size(vec_dot_type); + int64_t ne10_block_start = (ith * ne10/bs) / nth; + int64_t ne10_block_end = ((ith + 1) * ne10/bs) / nth; + from_float((float *)((char *) src1->data + i13*nb13 + i12*nb12 + i11*nb11 + ne10_block_start*bs*nb10), + (void *) (wdata + i13*nbw3 + i12*nbw2 + i11*nbw1 + ne10_block_start*nbw0), + (ne10_block_end - ne10_block_start) * bs); + } + } + } +#endif + } + + if (ith == 0) { + // initialize matrix_row_counts + memset(matrix_row_counts, 0, n_as*sizeof(int64_t)); + + // group rows by src0 matrix + for (int64_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) { + for (int id = 0; id < n_ids; ++id) { + const int32_t i02 = *(const int32_t *) ((const char *) ids->data + iid1*ids->nb[1] + id*ids->nb[0]); + + assert(i02 >= 0 && i02 < n_as); + + MMID_MATRIX_ROW(i02, matrix_row_counts[i02]) = (struct mmid_row_mapping) {id, iid1}; + matrix_row_counts[i02] += 1; + } + } + } + + // reset current_chunk + for (int cur_a = ith; cur_a < n_as; cur_a += nth) { + atomic_int * current_chunk_ctr = (atomic_int *)(atomic_current_chunk + cur_a); + *current_chunk_ctr = nth; + } + + ggml_barrier(params->threadpool); + + for (int cur_a = 0; cur_a < n_as; ++cur_a) { + const int64_t cne1 = matrix_row_counts[cur_a]; + + if (cne1 == 0) { + continue; + } + + const char * src0_cur = (const char *) src0->data + cur_a * nb02; + const void * wdata = (src1->type == vec_dot_type) ? src1->data : params->wdata; + const size_t row_size = ggml_row_size(vec_dot_type, ne10); + + const int64_t nr0 = ne01; + const int64_t nr1 = cne1; + + int chunk_size = 16; + if (nr0 == 1 || nr1 == 1) { + chunk_size = 64; + } + + // disable for NUMA + const bool disable_chunking = ggml_is_numa(); + + int64_t nchunk0 = (nr0 + chunk_size - 1) / chunk_size; + int64_t nchunk1 = (nr1 + chunk_size - 1) / chunk_size; + + if (nchunk0 * nchunk1 < nth * 4 || disable_chunking) { + nchunk0 = nr0 > nr1 ? nth : 1; + nchunk1 = nr0 > nr1 ? 1 : nth; + } + + const int64_t dr0 = (nr0 + nchunk0 - 1) / nchunk0; + const int64_t dr1 = (nr1 + nchunk1 - 1) / nchunk1; + + int current_chunk = ith; + + atomic_int * current_chunk_ctr = (atomic_int *)(atomic_current_chunk + cur_a); + + while (current_chunk < nchunk0 * nchunk1) { + const int64_t ith0 = current_chunk % nchunk0; + const int64_t ith1 = current_chunk / nchunk0; + + const int64_t ir0_start = dr0 * ith0; + const int64_t ir0_end = MIN(ir0_start + dr0, nr0); + + const int64_t ir1_start = dr1 * ith1; + const int64_t ir1_end = MIN(ir1_start + dr1, nr1); + + ggml_compute_forward_mul_mat_id_one_chunk( + dst, src0, src1, ids, cur_a, + ir0_start, ir0_end, ir1_start, ir1_end, + src0_cur, matrix_rows, row_size, src1_cont, wdata + ); + + if (nth >= nchunk0 * nchunk1) { + break; + } + + current_chunk = atomic_fetch_add_explicit(current_chunk_ctr, 1, memory_order_relaxed); + } + } +} + +///////////////////////////////// + +static void ggml_compute_forward(struct ggml_compute_params * params, struct ggml_tensor * tensor) { + GGML_ASSERT(params); + + if (tensor->op == GGML_OP_NONE || ggml_is_empty(tensor)) { + return; + } + + // extra_buffer op? + if (ggml_cpu_extra_compute_forward(params, tensor)) { + return; + } + + switch (tensor->op) { + case GGML_OP_DUP: + { + ggml_compute_forward_dup(params, tensor); + } break; + case GGML_OP_ADD: + { + ggml_compute_forward_add(params, tensor); + } break; + case GGML_OP_ADD_ID: + { + ggml_compute_forward_add_id(params, tensor); + } break; + case GGML_OP_ADD1: + { + ggml_compute_forward_add1(params, tensor); + } break; + case GGML_OP_ACC: + { + ggml_compute_forward_acc(params, tensor); + } break; + case GGML_OP_SUB: + { + ggml_compute_forward_sub(params, tensor); + } break; + case GGML_OP_MUL: + { + ggml_compute_forward_mul(params, tensor); + } break; + case GGML_OP_DIV: + { + ggml_compute_forward_div(params, tensor); + } break; + case GGML_OP_SQR: + { + ggml_compute_forward_sqr(params, tensor); + } break; + case GGML_OP_SQRT: + { + ggml_compute_forward_sqrt(params, tensor); + } break; + case GGML_OP_LOG: + { + ggml_compute_forward_log(params, tensor); + } break; + case GGML_OP_SIN: + { + ggml_compute_forward_sin(params, tensor); + } break; + case GGML_OP_COS: + { + ggml_compute_forward_cos(params, tensor); + } break; + case GGML_OP_SUM: + { + ggml_compute_forward_sum(params, tensor); + } break; + case GGML_OP_SUM_ROWS: + { + ggml_compute_forward_sum_rows(params, tensor); + } break; + case GGML_OP_CUMSUM: + { + ggml_compute_forward_cumsum(params, tensor); + } break; + case GGML_OP_MEAN: + { + ggml_compute_forward_mean(params, tensor); + } break; + case GGML_OP_ARGMAX: + { + ggml_compute_forward_argmax(params, tensor); + } break; + case GGML_OP_COUNT_EQUAL: + { + ggml_compute_forward_count_equal(params, tensor); + } break; + case GGML_OP_REPEAT: + { + ggml_compute_forward_repeat(params, tensor); + } break; + case GGML_OP_REPEAT_BACK: + { + ggml_compute_forward_repeat_back(params, tensor); + } break; + case GGML_OP_CONCAT: + { + ggml_compute_forward_concat(params, tensor); + } break; + case GGML_OP_SILU_BACK: + { + ggml_compute_forward_silu_back(params, tensor); + } break; + case GGML_OP_NORM: + { + ggml_compute_forward_norm(params, tensor); + } break; + case GGML_OP_RMS_NORM: + { + ggml_compute_forward_rms_norm(params, tensor); + } break; + case GGML_OP_RMS_NORM_BACK: + { + ggml_compute_forward_rms_norm_back(params, tensor); + } break; + case GGML_OP_GROUP_NORM: + { + ggml_compute_forward_group_norm(params, tensor); + } break; + case GGML_OP_L2_NORM: + { + ggml_compute_forward_l2_norm(params, tensor); + } break; + case GGML_OP_MUL_MAT: + { + ggml_compute_forward_mul_mat(params, tensor); + } break; + case GGML_OP_MUL_MAT_ID: + { + ggml_compute_forward_mul_mat_id(params, tensor); + } break; + case GGML_OP_MOE_FUSED: + { + GGML_ABORT("GGML_OP_MOE_FUSED is only implemented for CUDA"); + } + case GGML_OP_OUT_PROD: + { + ggml_compute_forward_out_prod(params, tensor); + } break; + case GGML_OP_SCALE: + { + ggml_compute_forward_scale(params, tensor); + } break; + case GGML_OP_SET: + { + ggml_compute_forward_set(params, tensor); + } break; + case GGML_OP_CPY: + { + ggml_compute_forward_cpy(params, tensor); + } break; + case GGML_OP_CONT: + { + ggml_compute_forward_cont(params, tensor); + } break; + case GGML_OP_GET_ROWS: + { + ggml_compute_forward_get_rows(params, tensor); + } break; + case GGML_OP_GET_ROWS_BACK: + { + ggml_compute_forward_get_rows_back(params, tensor); + } break; + case GGML_OP_SET_ROWS: + { + ggml_compute_forward_set_rows(params, tensor); + } break; + case GGML_OP_DIAG: + { + ggml_compute_forward_diag(params, tensor); + } break; + case GGML_OP_DIAG_MASK_INF: + { + ggml_compute_forward_diag_mask_inf(params, tensor); + } break; + case GGML_OP_DIAG_MASK_ZERO: + { + ggml_compute_forward_diag_mask_zero(params, tensor); + } break; + case GGML_OP_SOFT_MAX: + { + ggml_compute_forward_soft_max(params, tensor); + } break; + case GGML_OP_SOFT_MAX_BACK: + { + ggml_compute_forward_soft_max_ext_back(params, tensor); + } break; + case GGML_OP_ROPE: + { + ggml_compute_forward_rope(params, tensor); + } break; + case GGML_OP_ROPE_BACK: + { + ggml_compute_forward_rope_back(params, tensor); + } break; + case GGML_OP_CLAMP: + { + ggml_compute_forward_clamp(params, tensor); + } break; + case GGML_OP_CONV_TRANSPOSE_1D: + { + ggml_compute_forward_conv_transpose_1d(params, tensor); + } break; + case GGML_OP_IM2COL: + { + ggml_compute_forward_im2col(params, tensor); + } break; + case GGML_OP_IM2COL_BACK: + { + ggml_compute_forward_im2col_back_f32(params, tensor); + } break; + case GGML_OP_IM2COL_3D: + { + ggml_compute_forward_im2col_3d(params, tensor); + } break; + case GGML_OP_CONV_2D: + { + ggml_compute_forward_conv_2d(params, tensor); + } break; + case GGML_OP_CONV_3D: + { + ggml_compute_forward_conv_3d(params, tensor); + } break; + case GGML_OP_CONV_2D_DW: + { + ggml_compute_forward_conv_2d_dw(params, tensor); + } break; + case GGML_OP_CONV_TRANSPOSE_2D: + { + ggml_compute_forward_conv_transpose_2d(params, tensor); + } break; + case GGML_OP_POOL_1D: + { + ggml_compute_forward_pool_1d(params, tensor); + } break; + case GGML_OP_POOL_2D: + { + ggml_compute_forward_pool_2d(params, tensor); + } break; + case GGML_OP_POOL_2D_BACK: + { + ggml_compute_forward_pool_2d_back(params, tensor); + } break; + case GGML_OP_UPSCALE: + { + ggml_compute_forward_upscale(params, tensor); + } break; + case GGML_OP_PAD: + { + ggml_compute_forward_pad(params, tensor); + } break; + case GGML_OP_PAD_REFLECT_1D: + { + ggml_compute_forward_pad_reflect_1d(params, tensor); + } break; + case GGML_OP_ROLL: + { + ggml_compute_forward_roll(params, tensor); + } break; + case GGML_OP_ARANGE: + { + ggml_compute_forward_arange(params, tensor); + } break; + case GGML_OP_TIMESTEP_EMBEDDING: + { + ggml_compute_forward_timestep_embedding(params, tensor); + } break; + case GGML_OP_ARGSORT: + { + ggml_compute_forward_argsort(params, tensor); + } break; + case GGML_OP_TOP_K: + { + ggml_compute_forward_top_k(params, tensor); + } break; + case GGML_OP_LEAKY_RELU: + { + ggml_compute_forward_leaky_relu(params, tensor); + } break; + case GGML_OP_TRI: + { + ggml_compute_forward_tri(params, tensor); + } break; + case GGML_OP_FILL: + { + ggml_compute_forward_fill(params, tensor); + } break; + case GGML_OP_FLASH_ATTN_EXT: + { + ggml_compute_forward_flash_attn_ext(params, tensor); + } break; + case GGML_OP_FLASH_ATTN_SPARSE: + { + GGML_ABORT("GGML_OP_FLASH_ATTN_SPARSE is only supported on the CUDA backend"); + } + case GGML_OP_FLASH_ATTN_BACK: + { + int32_t t = ggml_get_op_params_i32(tensor, 0); + GGML_ASSERT(t == 0 || t == 1); + bool masked = t != 0; + ggml_compute_forward_flash_attn_back(params, masked, tensor); + } break; + case GGML_OP_SSM_CONV: + { + ggml_compute_forward_ssm_conv(params, tensor); + } break; + case GGML_OP_SSM_SCAN: + { + ggml_compute_forward_ssm_scan(params, tensor); + } break; + case GGML_OP_WIN_PART: + { + ggml_compute_forward_win_part(params, tensor); + } break; + case GGML_OP_WIN_UNPART: + { + ggml_compute_forward_win_unpart(params, tensor); + } break; + case GGML_OP_UNARY: + { + ggml_compute_forward_unary(params, tensor); + } break; + case GGML_OP_GLU: + { + ggml_compute_forward_glu(params, tensor); + } break; + case GGML_OP_TURBO_WHT: + { + // CUDA-only op; CPU fallback is identity (no rotation) + GGML_ASSERT(false && "TURBO_WHT is CUDA-only, use CUDA backend"); + } + case GGML_OP_GET_REL_POS: + { + ggml_compute_forward_get_rel_pos(params, tensor); + } break; + case GGML_OP_ADD_REL_POS: + { + ggml_compute_forward_add_rel_pos(params, tensor); + } break; + case GGML_OP_RWKV_WKV6: + { + ggml_compute_forward_rwkv_wkv6(params, tensor); + } break; + case GGML_OP_GATED_LINEAR_ATTN: + { + ggml_compute_forward_gla(params, tensor); + } break; + case GGML_OP_RWKV_WKV7: + { + ggml_compute_forward_rwkv_wkv7(params, tensor); + } break; + case GGML_OP_SOLVE_TRI: + { + ggml_compute_forward_solve_tri(params, tensor); + } break; + case GGML_OP_GATED_DELTA_NET: + { + ggml_compute_forward_gated_delta_net(params, tensor); + } break; + case GGML_OP_MAP_CUSTOM1: + { + ggml_compute_forward_map_custom1(params, tensor); + } + break; + case GGML_OP_MAP_CUSTOM2: + { + ggml_compute_forward_map_custom2(params, tensor); + } + break; + case GGML_OP_MAP_CUSTOM3: + { + ggml_compute_forward_map_custom3(params, tensor); + } + break; + case GGML_OP_CUSTOM: + { + ggml_compute_forward_custom(params, tensor); + } + break; + case GGML_OP_CROSS_ENTROPY_LOSS: + { + ggml_compute_forward_cross_entropy_loss(params, tensor); + } + break; + case GGML_OP_CROSS_ENTROPY_LOSS_BACK: + { + ggml_compute_forward_cross_entropy_loss_back(params, tensor); + } + break; + case GGML_OP_OPT_STEP_ADAMW: + { + ggml_compute_forward_opt_step_adamw(params, tensor); + } + break; + case GGML_OP_OPT_STEP_SGD: + { + ggml_compute_forward_opt_step_sgd(params, tensor); + } + break; + case GGML_OP_NONE: + { + // nop + } break; + case GGML_OP_RESHAPE: + { + // nop + } break; + case GGML_OP_PERMUTE: + { + // nop + } break; + case GGML_OP_VIEW: + { + // nop + } break; + case GGML_OP_TRANSPOSE: + { + // nop + } break; + case GGML_OP_COUNT: + { + GGML_ABORT("fatal error"); + } + } +} + +// Android's libc implementation "bionic" does not support setting affinity +#if defined(__gnu_linux__) +static void set_numa_thread_affinity(int thread_n) { + if (!ggml_is_numa()) { + return; + } + + int node_num; + int rv; + size_t setsize = CPU_ALLOC_SIZE(g_state.numa.total_cpus); + + switch(g_state.numa.numa_strategy) { + case GGML_NUMA_STRATEGY_DISTRIBUTE: + // run thread on node_num thread_n / (threads per node) + node_num = thread_n % g_state.numa.n_nodes; + break; + case GGML_NUMA_STRATEGY_ISOLATE: + // run thread on current_node + node_num = g_state.numa.current_node; + break; + case GGML_NUMA_STRATEGY_NUMACTL: + // use the cpuset that numactl gave us + rv = pthread_setaffinity_np(pthread_self(), setsize, &g_state.numa.cpuset); + if (rv) { + fprintf(stderr, "warning: pthread_setaffinity_np() failed: %s\n",strerror(rv)); + } + return; + default: + return; + } + + struct ggml_numa_node * node = &g_state.numa.nodes[node_num]; + + cpu_set_t * cpus = CPU_ALLOC(g_state.numa.total_cpus); + CPU_ZERO_S(setsize, cpus); + for (size_t i = 0; i < node->n_cpus; ++i) { + CPU_SET_S(node->cpus[i], setsize, cpus); + } + + rv = pthread_setaffinity_np(pthread_self(), setsize, cpus); + if (rv) { + fprintf(stderr, "warning: pthread_setaffinity_np() failed: %s\n", strerror(rv)); + } + + CPU_FREE(cpus); +} + +static void clear_numa_thread_affinity(void) { + if (!ggml_is_numa()) { + return; + } + + size_t setsize = CPU_ALLOC_SIZE(g_state.numa.total_cpus); + + cpu_set_t * cpus = CPU_ALLOC(g_state.numa.total_cpus); + CPU_ZERO_S(setsize, cpus); + for (unsigned i = 0; i < g_state.numa.total_cpus; ++i) { + CPU_SET_S(i, setsize, cpus); + } + + int rv = pthread_setaffinity_np(pthread_self(), setsize, cpus); + if (rv) { + fprintf(stderr, "warning: pthread_setaffinity_np() failed: %s\n", strerror(rv)); + } + + CPU_FREE(cpus); +} +#else +// TODO: Windows etc. +// (the linux implementation may also work on BSD, someone should test) +static void set_numa_thread_affinity(int thread_n) { UNUSED(thread_n); } +static void clear_numa_thread_affinity(void) {} +#endif + +static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { + int n_tasks = 0; + + if (ggml_is_empty(node)) { + // no need to multi-thread a no-op + n_tasks = 1; + return n_tasks; + } + + switch (node->op) { + case GGML_OP_CPY: + case GGML_OP_DUP: + case GGML_OP_CONT: + case GGML_OP_ADD: + case GGML_OP_ADD_ID: + case GGML_OP_ADD1: + case GGML_OP_ACC: + case GGML_OP_CUMSUM: + case GGML_OP_TRI: + case GGML_OP_FILL: + { + n_tasks = n_threads; + } break; + case GGML_OP_SUB: + case GGML_OP_SQR: + case GGML_OP_SQRT: + case GGML_OP_LOG: + case GGML_OP_SIN: + case GGML_OP_COS: + case GGML_OP_SUM: + case GGML_OP_SUM_ROWS: + case GGML_OP_MEAN: + case GGML_OP_ARGMAX: + { + n_tasks = 1; + } break; + case GGML_OP_COUNT_EQUAL: + case GGML_OP_SOLVE_TRI: + case GGML_OP_GATED_DELTA_NET: + { + n_tasks = n_threads; + } break; + case GGML_OP_REPEAT: + case GGML_OP_REPEAT_BACK: + case GGML_OP_LEAKY_RELU: + { + n_tasks = 1; + } break; + case GGML_OP_UNARY: + switch (ggml_get_unary_op(node)) { + case GGML_UNARY_OP_ABS: + case GGML_UNARY_OP_SGN: + case GGML_UNARY_OP_NEG: + case GGML_UNARY_OP_STEP: + case GGML_UNARY_OP_TANH: + case GGML_UNARY_OP_ELU: + case GGML_UNARY_OP_RELU: + case GGML_UNARY_OP_SIGMOID: + case GGML_UNARY_OP_HARDSWISH: + case GGML_UNARY_OP_HARDSIGMOID: + case GGML_UNARY_OP_EXP: + case GGML_UNARY_OP_SOFTPLUS: + case GGML_UNARY_OP_EXPM1: + case GGML_UNARY_OP_FLOOR: + case GGML_UNARY_OP_CEIL: + case GGML_UNARY_OP_ROUND: + case GGML_UNARY_OP_TRUNC: + { + n_tasks = 1; + } break; + + case GGML_UNARY_OP_GELU: + case GGML_UNARY_OP_GELU_ERF: + case GGML_UNARY_OP_GELU_QUICK: + case GGML_UNARY_OP_SILU: + case GGML_UNARY_OP_XIELU: + { + n_tasks = n_threads; + } break; + default: + GGML_ABORT("fatal error"); + } + break; + case GGML_OP_GLU: + switch (ggml_get_glu_op(node)) { + case GGML_GLU_OP_REGLU: + case GGML_GLU_OP_GEGLU: + case GGML_GLU_OP_SWIGLU: + case GGML_GLU_OP_SWIGLU_OAI: + case GGML_GLU_OP_GEGLU_ERF: + case GGML_GLU_OP_GEGLU_QUICK: + { + n_tasks = n_threads; + } break; + default: + break; + } + break; + case GGML_OP_TURBO_WHT: + { + n_tasks = n_threads; + } break; + case GGML_OP_SILU_BACK: + case GGML_OP_MUL: + case GGML_OP_DIV: + case GGML_OP_NORM: + case GGML_OP_RMS_NORM: + case GGML_OP_RMS_NORM_BACK: + case GGML_OP_L2_NORM: + case GGML_OP_GROUP_NORM: + case GGML_OP_CONCAT: + case GGML_OP_MUL_MAT: + case GGML_OP_MUL_MAT_ID: + case GGML_OP_OUT_PROD: + { + n_tasks = n_threads; + } break; + case GGML_OP_GET_ROWS: + case GGML_OP_SET_ROWS: + { + // FIXME: get_rows can use additional threads, but the cost of launching additional threads + // decreases performance with GPU offloading + //n_tasks = n_threads; + n_tasks = 1; + } break; + case GGML_OP_SCALE: + case GGML_OP_SET: + case GGML_OP_RESHAPE: + case GGML_OP_VIEW: + case GGML_OP_PERMUTE: + case GGML_OP_TRANSPOSE: + case GGML_OP_GET_ROWS_BACK: + case GGML_OP_DIAG: + { + n_tasks = 1; + } break; + case GGML_OP_DIAG_MASK_ZERO: + case GGML_OP_DIAG_MASK_INF: + case GGML_OP_SOFT_MAX_BACK: + case GGML_OP_ROPE: + case GGML_OP_ROPE_BACK: + case GGML_OP_ADD_REL_POS: + { + n_tasks = n_threads; + } break; + case GGML_OP_CLAMP: + { + n_tasks = 1; //TODO + } break; + case GGML_OP_SOFT_MAX: + { + n_tasks = MIN(n_threads, ggml_nrows(node->src[0])); + } break; + case GGML_OP_IM2COL: + case GGML_OP_IM2COL_BACK: + case GGML_OP_IM2COL_3D: + case GGML_OP_CONV_2D: + case GGML_OP_CONV_3D: + case GGML_OP_CONV_2D_DW: + case GGML_OP_CONV_TRANSPOSE_1D: + case GGML_OP_CONV_TRANSPOSE_2D: + { + n_tasks = n_threads; + } break; + case GGML_OP_POOL_1D: + case GGML_OP_POOL_2D: + case GGML_OP_POOL_2D_BACK: + { + n_tasks = 1; + } break; + case GGML_OP_UPSCALE: + case GGML_OP_PAD: + case GGML_OP_PAD_REFLECT_1D: + case GGML_OP_ROLL: + case GGML_OP_ARANGE: + case GGML_OP_TIMESTEP_EMBEDDING: + case GGML_OP_ARGSORT: + case GGML_OP_TOP_K: + case GGML_OP_FLASH_ATTN_EXT: + case GGML_OP_FLASH_ATTN_SPARSE: + case GGML_OP_FLASH_ATTN_BACK: + case GGML_OP_SSM_CONV: + case GGML_OP_SSM_SCAN: + { + n_tasks = n_threads; + } break; + case GGML_OP_RWKV_WKV6: + case GGML_OP_GATED_LINEAR_ATTN: + case GGML_OP_RWKV_WKV7: + { + const int64_t n_heads = node->src[1]->ne[1]; + n_tasks = MIN(n_threads, n_heads); + } break; + case GGML_OP_WIN_PART: + case GGML_OP_WIN_UNPART: + case GGML_OP_GET_REL_POS: + { + n_tasks = 1; + } break; + case GGML_OP_MAP_CUSTOM1: + { + struct ggml_map_custom1_op_params p; + memcpy(&p, node->op_params, sizeof(p)); + if (p.n_tasks == GGML_N_TASKS_MAX) { + n_tasks = n_threads; + } else { + n_tasks = MIN(p.n_tasks, n_threads); + } + } break; + case GGML_OP_MAP_CUSTOM2: + { + struct ggml_map_custom2_op_params p; + memcpy(&p, node->op_params, sizeof(p)); + if (p.n_tasks == GGML_N_TASKS_MAX) { + n_tasks = n_threads; + } else { + n_tasks = MIN(p.n_tasks, n_threads); + } + } break; + case GGML_OP_MAP_CUSTOM3: + { + struct ggml_map_custom3_op_params p; + memcpy(&p, node->op_params, sizeof(p)); + if (p.n_tasks == GGML_N_TASKS_MAX) { + n_tasks = n_threads; + } else { + n_tasks = MIN(p.n_tasks, n_threads); + } + } break; + case GGML_OP_CUSTOM: + { + struct ggml_custom_op_params p; + memcpy(&p, node->op_params, sizeof(p)); + if (p.n_tasks == GGML_N_TASKS_MAX) { + n_tasks = n_threads; + } else { + n_tasks = MIN(p.n_tasks, n_threads); + } + } break; + case GGML_OP_CROSS_ENTROPY_LOSS: + case GGML_OP_CROSS_ENTROPY_LOSS_BACK: + case GGML_OP_OPT_STEP_ADAMW: + case GGML_OP_OPT_STEP_SGD: + { + n_tasks = n_threads; + } break; + case GGML_OP_NONE: + { + n_tasks = 1; + } break; + case GGML_OP_COUNT: + { + GGML_ABORT("fatal error"); + } + default: + { + fprintf(stderr, "%s: op not implemented: ", __func__); + if (node->op < GGML_OP_COUNT) { + fprintf(stderr, "%s\n", ggml_op_name(node->op)); + } else { + fprintf(stderr, "%d\n", node->op); + } + GGML_ABORT("fatal error"); + } + } + + assert(n_tasks > 0); + + return n_tasks; +} + +static thread_ret_t ggml_graph_compute_secondary_thread(void* data); + +#if defined(_WIN32) +#include "windows.h" + +// TODO: support > 64 CPUs +static bool ggml_thread_apply_affinity(bool * mask) { + HANDLE h = GetCurrentThread(); + uint64_t bitmask = 0ULL; + + assert(GGML_MAX_N_THREADS >= 64); + + for (int32_t i = 0; i < 8; i++) { + int32_t idx = i * 8; + uint8_t val = 0; + val |= mask[idx + 0] << 0; + val |= mask[idx + 1] << 1; + val |= mask[idx + 2] << 2; + val |= mask[idx + 3] << 3; + val |= mask[idx + 4] << 4; + val |= mask[idx + 5] << 5; + val |= mask[idx + 6] << 6; + val |= mask[idx + 7] << 7; + bitmask |= (uint64_t)val << idx; + } + + for (int32_t i = 64; i < GGML_MAX_N_THREADS; i++) { + if (mask[i]) { + fprintf(stderr, "warn: setting thread-affinity for > 64 CPUs isn't supported on windows!\n"); + break; + } + } + + DWORD_PTR m = (DWORD_PTR)bitmask; + + m = SetThreadAffinityMask(h, m); + + return m != 0; +} + +static bool ggml_thread_apply_priority(int32_t prio) { + // Note that on Windows the Process Priority Class must be updated in order to set Thread priority. + // This is up to the applications. + DWORD p = THREAD_PRIORITY_NORMAL; + switch (prio) { + case GGML_SCHED_PRIO_LOW: p = THREAD_PRIORITY_BELOW_NORMAL; break; + case GGML_SCHED_PRIO_NORMAL: p = THREAD_PRIORITY_NORMAL; break; + case GGML_SCHED_PRIO_MEDIUM: p = THREAD_PRIORITY_ABOVE_NORMAL; break; + case GGML_SCHED_PRIO_HIGH: p = THREAD_PRIORITY_HIGHEST; break; + case GGML_SCHED_PRIO_REALTIME: p = THREAD_PRIORITY_TIME_CRITICAL; break; + } + + if (prio != GGML_SCHED_PRIO_LOW) { + // Tell Windows that this thread should not be throttled (needs its own CPU core). + // Newer Windows 11 versions aggressively park (offline) CPU cores and often place + // all our threads onto the first 4 cores which results in terrible performance with + // n_threads > 4 + #if _WIN32_WINNT >= 0x0602 + THREAD_POWER_THROTTLING_STATE t; + ZeroMemory(&t, sizeof(t)); + t.Version = THREAD_POWER_THROTTLING_CURRENT_VERSION; + t.ControlMask = THREAD_POWER_THROTTLING_EXECUTION_SPEED; + t.StateMask = 0; + + if (!SetThreadInformation(GetCurrentThread(), ThreadPowerThrottling, &t, sizeof(t))) { + GGML_LOG_DEBUG("failed to disable thread power throttling %d : (%d)\n", prio, (int) GetLastError()); + return false; + } + #endif + } + + if (prio == GGML_SCHED_PRIO_NORMAL) { + // Keep inherited policy/priority + return true; + } + + if (!SetThreadPriority(GetCurrentThread(), p)) { + fprintf(stderr, "warn: failed to set thread priority %d : (%d)\n", prio, (int) GetLastError()); + return false; + } + + return true; +} + +#elif defined(__APPLE__) +#include +#include + +static bool ggml_thread_apply_affinity(const bool * mask) { + // Not supported on Apple platforms + UNUSED(mask); + return true; +} + +static bool ggml_thread_apply_priority(int32_t prio) { + struct sched_param p; + int32_t policy = SCHED_OTHER; + switch (prio) { + // TODO: there seems to be no way to set lower prio on Apple platforms + case GGML_SCHED_PRIO_LOW: policy = SCHED_OTHER; p.sched_priority = 0; break; + case GGML_SCHED_PRIO_NORMAL: policy = SCHED_OTHER; p.sched_priority = 0; break; + case GGML_SCHED_PRIO_MEDIUM: policy = SCHED_FIFO; p.sched_priority = 40; break; + case GGML_SCHED_PRIO_HIGH: policy = SCHED_FIFO; p.sched_priority = 80; break; + case GGML_SCHED_PRIO_REALTIME: policy = SCHED_FIFO; p.sched_priority = 90; break; + } + + if (prio == GGML_SCHED_PRIO_NORMAL) { + // Keep inherited policy/priority + return true; + } + + int32_t err = pthread_setschedparam(pthread_self(), policy, &p); + if (err != 0) { + fprintf(stderr, "warn: failed to set thread priority %d : %s (%d)\n", prio, strerror(err), err); + return false; + } + + return true; +} + +#elif defined(__gnu_linux__) +// TODO: this may not work on BSD, to be verified + +static bool ggml_thread_apply_affinity(const bool * mask) { + cpu_set_t cpuset; + int err; + + CPU_ZERO(&cpuset); + + for (uint32_t i = 0; i < GGML_MAX_N_THREADS; i++) { + if (mask[i]) { + GGML_PRINT_DEBUG("Thread %lx: adding %d to cpuset\n", pthread_self(), i); + CPU_SET(i, &cpuset); + } + } + +#ifdef __ANDROID__ + err = sched_setaffinity(0, sizeof(cpuset), &cpuset); + if (err < 0) { + err = errno; + } +#else + err = pthread_setaffinity_np(pthread_self(), sizeof(cpuset), &cpuset); +#endif + if (err != 0) { + fprintf(stderr, "warn: failed to set affinity mask 0x%llx : %s (%d)\n", (unsigned long long)mask, strerror(err), err); + return false; + } + + return true; +} + +static bool ggml_thread_apply_priority(int32_t prio) { + struct sched_param p; + int32_t policy = SCHED_OTHER; + switch (prio) { + case GGML_SCHED_PRIO_LOW: policy = SCHED_BATCH; p.sched_priority = 0; break; + case GGML_SCHED_PRIO_NORMAL: policy = SCHED_OTHER; p.sched_priority = 0; break; + case GGML_SCHED_PRIO_MEDIUM: policy = SCHED_FIFO; p.sched_priority = 40; break; + case GGML_SCHED_PRIO_HIGH: policy = SCHED_FIFO; p.sched_priority = 80; break; + case GGML_SCHED_PRIO_REALTIME: policy = SCHED_FIFO; p.sched_priority = 90; break; + } + + if (prio == GGML_SCHED_PRIO_NORMAL) { + // Keep inherited policy/priority + return true; + } + + int32_t err = pthread_setschedparam(pthread_self(), policy, &p); + if (err != 0) { + fprintf(stderr, "warn: failed to set thread priority %d : %s (%d)\n", prio, strerror(err), err); + return false; + } + + return true; +} + +#else // unsupported platforms + +static bool ggml_thread_apply_affinity(const bool * mask) { + UNUSED(mask); + return true; +} + +static bool ggml_thread_apply_priority(int32_t prio) { + UNUSED(prio); + return true; +} + +#endif + +static bool ggml_thread_cpumask_is_valid(const bool * mask) { + for (int i = 0; i < GGML_MAX_N_THREADS; i++) { + if (mask[i]) { return true; } + } + return false; +} + +static void ggml_thread_cpumask_next(const bool * global_mask, bool * local_mask, bool strict, int32_t* iter) { + if (!strict) { + memcpy(local_mask, global_mask, GGML_MAX_N_THREADS); + return; + } else { + memset(local_mask, 0, GGML_MAX_N_THREADS); + int32_t base_idx = *iter; + for (int32_t i = 0; i < GGML_MAX_N_THREADS; i++) { + int32_t idx = base_idx + i; + if (idx >= GGML_MAX_N_THREADS) { + // Just a cheaper modulo + idx -= GGML_MAX_N_THREADS; + } + if (global_mask[idx]) { + local_mask[idx] = 1; + *iter = idx + 1; + return; + } + } + } +} + +void ggml_threadpool_free(struct ggml_threadpool* threadpool) { + if (!threadpool) return; + + const int n_threads = threadpool->n_threads; + +#ifndef GGML_USE_OPENMP + struct ggml_compute_state* workers = threadpool->workers; + + ggml_mutex_lock(&threadpool->mutex); + + threadpool->stop = true; + threadpool->pause = false; + + ggml_cond_broadcast(&threadpool->cond); + ggml_mutex_unlock(&threadpool->mutex); + + for (int j = 1; j < n_threads; j++) { + int32_t rc = ggml_thread_join(workers[j].thrd, NULL); + GGML_ASSERT(rc == GGML_EXIT_SUCCESS || rc == GGML_EXIT_ABORTED); + UNUSED(rc); + } + + ggml_mutex_destroy(&threadpool->mutex); + ggml_cond_destroy(&threadpool->cond); +#endif // GGML_USE_OPENMP + + const size_t workers_size = sizeof(struct ggml_compute_state) * n_threads; + ggml_aligned_free(threadpool->workers, workers_size); + ggml_aligned_free(threadpool, sizeof(struct ggml_threadpool)); +} + +#ifndef GGML_USE_OPENMP +// pause/resume must be called under mutex +static void ggml_threadpool_pause_locked(struct ggml_threadpool * threadpool) { + GGML_PRINT_DEBUG("Pausing threadpool\n"); + threadpool->pause = true; + ggml_cond_broadcast(&threadpool->cond); +} + +static void ggml_threadpool_resume_locked(struct ggml_threadpool * threadpool) { + GGML_PRINT_DEBUG("Resuming threadpool\n"); + threadpool->pause = false; + ggml_cond_broadcast(&threadpool->cond); +} +#endif + +void ggml_threadpool_pause(struct ggml_threadpool * threadpool) { +#ifndef GGML_USE_OPENMP + ggml_mutex_lock(&threadpool->mutex); + if (!threadpool->pause) { + ggml_threadpool_pause_locked(threadpool); + } + ggml_mutex_unlock(&threadpool->mutex); +#else + UNUSED(threadpool); +#endif +} + +void ggml_threadpool_resume(struct ggml_threadpool * threadpool) { +#ifndef GGML_USE_OPENMP + ggml_mutex_lock(&threadpool->mutex); + if (threadpool->pause) { + ggml_threadpool_resume_locked(threadpool); + } + ggml_mutex_unlock(&threadpool->mutex); +#else + UNUSED(threadpool); +#endif +} + +struct ggml_cplan ggml_graph_plan( + const struct ggml_cgraph * cgraph, + int n_threads, + struct ggml_threadpool * threadpool) { + + if (threadpool == NULL) { + //GGML_PRINT_DEBUG("Threadpool is not specified. Will create a disposable threadpool : n_threads %d\n", n_threads); + } + if (n_threads <= 0) { + n_threads = threadpool ? threadpool->n_threads : GGML_DEFAULT_N_THREADS; + } + +#if defined(__EMSCRIPTEN__) && !defined(__EMSCRIPTEN_PTHREADS__) + // Emscripten without pthreads support can only use a single thread + n_threads = 1; +#endif + + size_t work_size = 0; + + struct ggml_cplan cplan; + memset(&cplan, 0, sizeof(struct ggml_cplan)); + + int max_tasks = 1; + + // thread scheduling for the different operations + work buffer size estimation + for (int i = 0; i < cgraph->n_nodes; i++) { + struct ggml_tensor * node = cgraph->nodes[i]; + + const int n_tasks = ggml_get_n_tasks(node, n_threads); + + max_tasks = MAX(max_tasks, n_tasks); + + size_t cur = 0; + + if (!ggml_cpu_extra_work_size(n_threads, node, &cur)) { + switch (node->op) { + case GGML_OP_CPY: + case GGML_OP_DUP: + { + if (ggml_is_quantized(node->type) || + // F16 -> BF16 and BF16 -> F16 copies go through intermediate F32 + (node->src[0]->type == GGML_TYPE_F16 && node->src[1] && node->src[1]->type == GGML_TYPE_BF16) || + (node->src[0]->type == GGML_TYPE_BF16 && node->src[1] && node->src[1]->type == GGML_TYPE_F16) || + // conversion between F32 and I32 + (node->src[0]->type == GGML_TYPE_F32 && node->src[1] && node->src[1]->type == GGML_TYPE_I32) || + (node->src[0]->type == GGML_TYPE_I32 && node->src[1] && node->src[1]->type == GGML_TYPE_F32)) { + cur = ggml_type_size(GGML_TYPE_F32) * node->ne[0] * n_tasks; + } + } break; + case GGML_OP_ADD: + case GGML_OP_ADD_ID: + case GGML_OP_ADD1: + { + if (ggml_is_quantized(node->src[0]->type)) { + cur = ggml_type_size(GGML_TYPE_F32) * node->src[0]->ne[0] * n_tasks; + } + } break; + case GGML_OP_ACC: + { + if (ggml_is_quantized(node->src[0]->type)) { + cur = ggml_type_size(GGML_TYPE_F32) * node->src[1]->ne[0] * n_tasks; + } + } break; + case GGML_OP_COUNT_EQUAL: + { + cur = ggml_type_size(node->type)*n_tasks; + } break; + case GGML_OP_MUL_MAT: + { + const enum ggml_type vec_dot_type = type_traits_cpu[node->src[0]->type].vec_dot_type; + + if (node->src[1]->type != vec_dot_type) { + cur = ggml_row_size(vec_dot_type, ggml_nelements(node->src[1])); + } + } break; + case GGML_OP_MUL_MAT_ID: + { + cur = 0; + const struct ggml_tensor * src0 = node->src[0]; + const struct ggml_tensor * src1 = node->src[1]; + const struct ggml_tensor * ids = node->src[2]; + const enum ggml_type vec_dot_type = type_traits_cpu[src0->type].vec_dot_type; + const int n_as = src0->ne[2]; + // src1 + if (src1->type != vec_dot_type) { + cur += ggml_row_size(vec_dot_type, ggml_nelements(src1)) + sizeof(int64_t); + } + // matrix_row_counts + cur += n_as * sizeof(int64_t) + sizeof(int64_t); + // matrix_rows + cur += n_as*ids->ne[0]*ids->ne[1]*sizeof(struct mmid_row_mapping) + sizeof(int64_t); + // atomic_current_chunk + cur += CACHE_LINE_SIZE*n_as + CACHE_LINE_SIZE; + } break; + case GGML_OP_OUT_PROD: + { + if (ggml_is_quantized(node->src[0]->type)) { + cur = ggml_type_size(GGML_TYPE_F32) * node->src[0]->ne[0] * n_tasks; + } + } break; + case GGML_OP_SOFT_MAX: + case GGML_OP_ROPE: + case GGML_OP_ROPE_BACK: + { + cur = ggml_type_size(GGML_TYPE_F32) * node->ne[0] * n_tasks; + } break; + case GGML_OP_CONV_TRANSPOSE_1D: + { + GGML_ASSERT(node->src[0]->ne[3] == 1); + GGML_ASSERT(node->src[1]->ne[2] == 1); + GGML_ASSERT(node->src[1]->ne[3] == 1); + + const int64_t ne00 = node->src[0]->ne[0]; // K + const int64_t ne01 = node->src[0]->ne[1]; // Cout + const int64_t ne02 = node->src[0]->ne[2]; // Cin + const int64_t ne10 = node->src[1]->ne[0]; // L + const int64_t ne11 = node->src[1]->ne[1]; // Cin + + if ((node->src[0]->type == GGML_TYPE_F16 || + node->src[0]->type == GGML_TYPE_BF16) && + node->src[1]->type == GGML_TYPE_F32) { + cur += sizeof(ggml_fp16_t)*ne00*ne01*ne02; + cur += sizeof(ggml_fp16_t)*ne10*ne11; + } else if (node->src[0]->type == GGML_TYPE_F32 && + node->src[1]->type == GGML_TYPE_F32) { + cur += sizeof(float)*ne00*ne01*ne02; + cur += sizeof(float)*ne10*ne11; + } else { + GGML_ABORT("fatal error"); + } + } break; + case GGML_OP_CONV_2D: + case GGML_OP_CONV_3D: + { + cur = GGML_IM2COL_WORK_SIZE; + } break; + case GGML_OP_CONV_TRANSPOSE_2D: + { + const int64_t ne00 = node->src[0]->ne[0]; // W + const int64_t ne01 = node->src[0]->ne[1]; // H + const int64_t ne02 = node->src[0]->ne[2]; // Channels Out + const int64_t ne03 = node->src[0]->ne[3]; // Channels In + + const int64_t ne10 = node->src[1]->ne[0]; // W + const int64_t ne11 = node->src[1]->ne[1]; // H + const int64_t ne12 = node->src[1]->ne[2]; // Channels In + + GGML_ASSERT(node->src[0]->type == GGML_TYPE_F16 || node->src[0]->type == GGML_TYPE_F32); + GGML_ASSERT(node->src[1]->type == GGML_TYPE_F32); + + cur += ggml_type_size(node->src[0]->type) * ne00 * ne01 * ne02 * ne03; + cur += ggml_type_size(node->src[0]->type) * ne10 * ne11 * ne12; + + } break; + case GGML_OP_TOP_K: + { + cur += sizeof(int32_t)*node->src[0]->ne[0]*n_tasks; + } break; + case GGML_OP_FLASH_ATTN_EXT: + { + const int64_t neq2 = node->src[0]->ne[2]; // number of query heads + const int64_t DK = node->src[1]->ne[0]; + const int64_t DV = node->src[2]->ne[0]; + + // Tiled flash attention scratch (tile sizes defined in common.h) + // Per-thread: Q_q + KQ + mask + VKQ32 + V32 + K_f32 + padding + size_t prefill = sizeof(float)*(GGML_FA_TILE_Q*DK + 2*GGML_FA_TILE_Q*GGML_FA_TILE_KV + GGML_FA_TILE_Q*DV + GGML_FA_TILE_KV*DV + GGML_FA_TILE_KV*DK)*n_tasks; + + // Decode path: n_kv_chunks = n_tasks (one chunk per thread) + // Per-thread: VKQ accmulator (DV), partial M, partial S + intra-thread scratch for V, Q and VKQ + size_t n_chunks = n_tasks; + size_t decode = sizeof(float)*(neq2*n_chunks*(2+DV) + n_tasks*(DK + 2*DV)); + + cur += MAX(prefill, decode); + } break; + case GGML_OP_FLASH_ATTN_BACK: + { + const int64_t D = node->src[0]->ne[0]; + const int64_t ne11 = ggml_up(node->src[1]->ne[1], GGML_SOFT_MAX_UNROLL); + const int64_t mxDn = MAX(D, ne11) * 2; // *2 because of S and SM in ggml_compute_forward_flash_attn_back + if (node->src[1]->type == GGML_TYPE_F32) { + cur = sizeof(float)*mxDn*n_tasks; // TODO: this can become (n_tasks-1) + cur += sizeof(float)*mxDn*n_tasks; // this is overestimated by x2 + } else if (node->src[1]->type == GGML_TYPE_F16) { + cur = sizeof(float)*mxDn*n_tasks; // TODO: this can become (n_tasks-1) + cur += sizeof(float)*mxDn*n_tasks; // this is overestimated by x2 + } else if (node->src[1]->type == GGML_TYPE_BF16) { + cur = sizeof(float)*mxDn*n_tasks; // TODO: this can become (n_tasks-1) + cur += sizeof(float)*mxDn*n_tasks; // this is overestimated by x2 + } + } break; + + case GGML_OP_CROSS_ENTROPY_LOSS: + { + cur = ggml_type_size(node->type)*(n_tasks + node->src[0]->ne[0]*n_tasks); + } break; + case GGML_OP_GATED_DELTA_NET: + { + const int64_t S_v = node->src[2]->ne[0]; + cur = S_v * sizeof(float) * n_tasks; + } break; + case GGML_OP_COUNT: + { + GGML_ABORT("fatal error"); + } + default: + break; + } + } + + work_size = MAX(work_size, cur); + } + + if (work_size > 0) { + work_size += CACHE_LINE_SIZE*(n_threads); + } + + cplan.threadpool = threadpool; + cplan.n_threads = MIN(max_tasks, n_threads); + cplan.work_size = work_size; + cplan.work_data = NULL; + + return cplan; +} + +static thread_ret_t ggml_graph_compute_thread(void * data) { + struct ggml_compute_state * state = (struct ggml_compute_state *) data; + struct ggml_threadpool * tp = state->threadpool; + + const struct ggml_cgraph * cgraph = tp->cgraph; + const struct ggml_cplan * cplan = tp->cplan; + + set_numa_thread_affinity(state->ith); + + struct ggml_compute_params params = { + /*.ith =*/ state->ith, + /*.nth =*/ atomic_load_explicit(&tp->n_graph, memory_order_relaxed) & GGML_THREADPOOL_N_THREADS_MASK, + /*.wsize =*/ cplan->work_size, + /*.wdata =*/ cplan->work_data, + /*.threadpool =*/ tp, + /*.use_ref =*/ cplan->use_ref, + }; + +#ifdef GGML_USE_OPENMP + GGML_PRINT_DEBUG("thread #%d compute-start cplan %p\n", state->ith, (const void *)cplan); +#else + GGML_PRINT_DEBUG("thread #%d compute-start cplan %p last-graph %d\n", state->ith, (const void *)cplan, state->last_graph); +#endif + + for (int node_n = 0; node_n < cgraph->n_nodes && atomic_load_explicit(&tp->abort, memory_order_relaxed) != node_n; node_n++) { + struct ggml_tensor * node = cgraph->nodes[node_n]; + + if (ggml_op_is_empty(node->op)) { + // skip NOPs + continue; + } + + if ((node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { + continue; + } + + ggml_compute_forward(¶ms, node); + + if (state->ith == 0 && cplan->abort_callback && + cplan->abort_callback(cplan->abort_callback_data)) { + atomic_store_explicit(&tp->abort, node_n + 1, memory_order_relaxed); + tp->ec = GGML_STATUS_ABORTED; + } + + if (node_n + 1 < cgraph->n_nodes) { + ggml_barrier(state->threadpool); + } + } + +#ifdef GGML_USE_OPENMP + GGML_PRINT_DEBUG("thread #%d compute-done cplan %p\n", state->ith, (const void *)cplan); +#else + GGML_PRINT_DEBUG("thread #%d compute-done cplan %p last-graph %d\n", state->ith, (const void *)cplan, state->last_graph); +#endif + + ggml_barrier(state->threadpool); + + return 0; +} + +#ifndef GGML_USE_OPENMP + +// check if thread is ready to proceed (exit from polling or sleeping) +// returns true if loops should exit, sets state->pending to indicate new work +static inline bool ggml_graph_compute_thread_ready(struct ggml_compute_state * state) { + struct ggml_threadpool * threadpool = state->threadpool; + + if (state->pending || threadpool->stop || threadpool->pause) { return true; } + + // check for new graph/work + int n_graph = atomic_load_explicit(&threadpool->n_graph, memory_order_relaxed); + int n_threads = n_graph & GGML_THREADPOOL_N_THREADS_MASK; + if (n_graph != state->last_graph) { + state->pending = (state->ith < n_threads); + state->last_graph = n_graph; + return true; + } + + return false; +} + +// sync thread state after polling +static inline void ggml_graph_compute_thread_sync(struct ggml_compute_state * state) { + // TSAN doesn't support standalone fence yet, we use a dummy read-modify-write instead + #ifdef GGML_TSAN_ENABLED + atomic_fetch_add_explicit(&state->threadpool->n_graph, 0, memory_order_seq_cst); + #else + atomic_thread_fence(memory_order_seq_cst); + #endif + UNUSED(state); +} + +static inline bool ggml_graph_compute_poll_for_work(struct ggml_compute_state * state) { + struct ggml_threadpool * threadpool = state->threadpool; + + // This seems to make 0 ... 100 a decent range for polling level across modern processors. + // Perhaps, we can adjust it dynamically based on load and things. + const uint64_t n_rounds = 1024UL * 128 * threadpool->poll; + + for (uint64_t i=0; !ggml_graph_compute_thread_ready(state) && i < n_rounds; i++) { + // No new work. Keep polling. + ggml_thread_cpu_relax(); + } + + return state->pending; +} + +static inline bool ggml_graph_compute_check_for_work(struct ggml_compute_state * state) { + struct ggml_threadpool * threadpool = state->threadpool; + + if (ggml_graph_compute_poll_for_work(state)) { + ggml_graph_compute_thread_sync(state); + return state->pending; + } + + ggml_mutex_lock_shared(&threadpool->mutex); + while (!ggml_graph_compute_thread_ready(state)) { + // No new work. Wait for the signal. + GGML_PRINT_DEBUG("thread #%d waiting for work (sleeping)\n", state->ith); + ggml_cond_wait(&threadpool->cond, &threadpool->mutex); + } + ggml_mutex_unlock_shared(&threadpool->mutex); + + return state->pending; +} + +static thread_ret_t ggml_graph_compute_secondary_thread(void* data) { + struct ggml_compute_state * state = (struct ggml_compute_state *) data; + struct ggml_threadpool * threadpool = state->threadpool; + + ggml_thread_apply_priority(threadpool->prio); + if (ggml_thread_cpumask_is_valid(state->cpumask)) { + ggml_thread_apply_affinity(state->cpumask); + } + + while (true) { + // Check if we need to sleep + while (threadpool->pause) { + GGML_PRINT_DEBUG("thread #%d inside pause loop\n", state->ith); + ggml_mutex_lock_shared(&threadpool->mutex); + if (threadpool->pause) { + ggml_cond_wait(&threadpool->cond, &threadpool->mutex); + } + GGML_PRINT_DEBUG("thread #%d resuming after wait\n", state->ith); + ggml_mutex_unlock_shared(&threadpool->mutex); + } + + // This needs to be checked for after the cond_wait + if (threadpool->stop) break; + + // Check if there is new work + // The main thread is the only one that can dispatch new work + + ggml_graph_compute_check_for_work(state); + if (state->pending) { + state->pending = false; + ggml_graph_compute_thread(state); + } + } + + return (thread_ret_t) 0; +} + +// Start processing new graph +static void ggml_graph_compute_kickoff(struct ggml_threadpool * threadpool, int n_threads) +{ + // Always take the mutex here because the worker threads are doing hybrid poll/wait + + ggml_mutex_lock(&threadpool->mutex); + + // Update the number of active threads and the graph count + int n_graph = atomic_load_explicit(&threadpool->n_graph, memory_order_relaxed) >> GGML_THREADPOOL_N_THREADS_BITS; + n_graph = ((n_graph + 1) << GGML_THREADPOOL_N_THREADS_BITS) | (n_threads & GGML_THREADPOOL_N_THREADS_MASK); + + GGML_PRINT_DEBUG("compute-kickoff: n_threads %d n_graph %d\n", n_threads, n_graph); + + // Indicate the graph is ready to be processed + // We need the full seq-cst fence here because of the polling threads (used in thread_sync) + atomic_store_explicit(&threadpool->n_graph, n_graph, memory_order_seq_cst); + + if (threadpool->pause) { + // Update main thread prio and affinity to match the threadpool settings + ggml_thread_apply_priority(threadpool->prio); + if (ggml_thread_cpumask_is_valid(threadpool->workers[0].cpumask)) { + ggml_thread_apply_affinity(threadpool->workers[0].cpumask); + } + + // resume does cond broadcast + ggml_threadpool_resume_locked(threadpool); + } else { + ggml_cond_broadcast(&threadpool->cond); + } + + ggml_mutex_unlock(&threadpool->mutex); +} + +#endif // GGML_USE_OPENMP + +static struct ggml_threadpool * ggml_threadpool_new_impl( + struct ggml_threadpool_params * tpp, + struct ggml_cgraph * cgraph, + struct ggml_cplan * cplan) { + + struct ggml_threadpool * threadpool = + ggml_aligned_malloc(sizeof(struct ggml_threadpool)); + { + threadpool->cgraph = cgraph; + threadpool->cplan = cplan; + threadpool->n_graph = 0; + threadpool->n_barrier = 0; + threadpool->n_barrier_passed = 0; + threadpool->current_chunk = 0; + threadpool->stop = false; + threadpool->pause = tpp->paused; + threadpool->abort = -1; + threadpool->workers = NULL; + threadpool->n_threads = tpp->n_threads; + threadpool->poll = tpp->poll; + threadpool->prio = tpp->prio; + threadpool->ec = GGML_STATUS_SUCCESS; + } + + // Allocate and init workers state + const size_t workers_size = sizeof(struct ggml_compute_state) * tpp->n_threads; + struct ggml_compute_state * workers = ggml_aligned_malloc(workers_size); + + memset(workers, 0, workers_size); + for (int j = 0; j < tpp->n_threads; j++) { + workers[j].threadpool = threadpool; + workers[j].ith = j; + } + + threadpool->workers = workers; + +#ifdef GGML_USE_OPENMP + int32_t cpumask_iter = 0; + + // Compute CPU masks for each thread + for (int j = 0; j < tpp->n_threads; j++) { + ggml_thread_cpumask_next(tpp->cpumask, workers[j].cpumask, tpp->strict_cpu, &cpumask_iter); + } +#else // GGML_USE_OPENMP + ggml_mutex_init(&threadpool->mutex); + ggml_cond_init(&threadpool->cond); + + // Spin the threads for all workers, and update CPU placements. + // Place the main thread last (towards the higher numbered CPU cores). + + int32_t cpumask_iter = 0; + + for (int j = 1; j < tpp->n_threads; j++) { + ggml_thread_cpumask_next(tpp->cpumask, workers[j].cpumask, tpp->strict_cpu, &cpumask_iter); + + int32_t rc = ggml_thread_create(&workers[j].thrd, NULL, ggml_graph_compute_secondary_thread, &workers[j]); + GGML_ASSERT(rc == 0); + } + + ggml_thread_cpumask_next(tpp->cpumask, workers[0].cpumask, tpp->strict_cpu, &cpumask_iter); + + if (!threadpool->pause) { + // Update main thread prio and affinity at the start, otherwise we'll do it in resume + ggml_thread_apply_priority(threadpool->prio); + if (ggml_thread_cpumask_is_valid(threadpool->workers[0].cpumask)) { + ggml_thread_apply_affinity(threadpool->workers[0].cpumask); + } + } +#endif // GGML_USE_OPENMP + + return threadpool; +} + +struct ggml_threadpool * ggml_threadpool_new(struct ggml_threadpool_params * tpp) { + return ggml_threadpool_new_impl(tpp, NULL, NULL); +} + +enum ggml_status ggml_graph_compute(struct ggml_cgraph * cgraph, struct ggml_cplan * cplan) { + ggml_cpu_init(); + + GGML_ASSERT(cplan); + GGML_ASSERT(cplan->n_threads > 0); + GGML_ASSERT(cplan->work_size == 0 || cplan->work_data != NULL); + + int n_threads = cplan->n_threads; + struct ggml_threadpool * threadpool = cplan->threadpool; + + bool disposable_threadpool = false; + + if (threadpool == NULL) { + //GGML_PRINT_DEBUG("Threadpool is not specified. Will create a disposable threadpool : n_threads %d\n", n_threads); + disposable_threadpool = true; + + struct ggml_threadpool_params ttp = ggml_threadpool_params_default(n_threads); + threadpool = ggml_threadpool_new_impl(&ttp, cgraph, cplan); + } else { + // Reset some of the parameters that need resetting + // No worker threads should be accessing the parameters below at this stage + threadpool->cgraph = cgraph; + threadpool->cplan = cplan; + threadpool->current_chunk = 0; + threadpool->abort = -1; + threadpool->ec = GGML_STATUS_SUCCESS; + } + +#ifdef GGML_USE_OPENMP + if (n_threads > 1) { + #pragma omp parallel num_threads(n_threads) + { + #pragma omp single + { + // update the number of threads from the actual number of threads that we got from OpenMP + n_threads = omp_get_num_threads(); + atomic_store_explicit(&threadpool->n_graph, n_threads, memory_order_relaxed); + } + + // Apply thread CPU mask and priority + int ith = omp_get_thread_num(); + + ggml_thread_apply_priority(threadpool->prio); + if (ggml_thread_cpumask_is_valid(threadpool->workers[ith].cpumask)) { + ggml_thread_apply_affinity(threadpool->workers[ith].cpumask); + } + ggml_graph_compute_thread(&threadpool->workers[ith]); + } + } else { + atomic_store_explicit(&threadpool->n_graph, 1, memory_order_relaxed); + ggml_graph_compute_thread(&threadpool->workers[0]); + } +#else + if (n_threads > threadpool->n_threads) { + GGML_LOG_WARN("cplan requested more threads (%d) than available (%d)\n", n_threads, threadpool->n_threads); + n_threads = threadpool->n_threads; + } + + // Kick all threads to start the new graph + ggml_graph_compute_kickoff(threadpool, n_threads); + + // This is a work thread too + ggml_graph_compute_thread(&threadpool->workers[0]); +#endif + + // don't leave affinity set on the main thread + clear_numa_thread_affinity(); + + enum ggml_status ret = threadpool->ec; + + if (disposable_threadpool) { + ggml_threadpool_free(threadpool); + } + + return ret; +} + +enum ggml_status ggml_graph_compute_with_ctx(struct ggml_context * ctx, struct ggml_cgraph * cgraph, int n_threads) { + struct ggml_cplan cplan = ggml_graph_plan(cgraph, n_threads, NULL); + + cplan.work_data = (uint8_t *)ggml_new_buffer(ctx, cplan.work_size); + + return ggml_graph_compute(cgraph, &cplan); +} + +void ggml_cpu_fp32_to_fp32(const float * x, float * y, int64_t n) { + memcpy(y, x, n * sizeof(float)); +} + +void ggml_cpu_fp32_to_fp16(const float * x, ggml_fp16_t * y, int64_t n) { + int64_t i = 0; +#if defined(__F16C__) +#if defined(__AVX512F__) + for (; i + 15 < n; i += 16) { + __m512 x_vec = _mm512_loadu_ps(x + i); + __m256i y_vec = _mm512_cvtps_ph(x_vec, _MM_FROUND_TO_NEAREST_INT); + _mm256_storeu_si256((__m256i *)(y + i), y_vec); + } +#endif + for (; i + 7 < n; i += 8) { + __m256 x_vec = _mm256_loadu_ps(x + i); + __m128i y_vec = _mm256_cvtps_ph(x_vec, _MM_FROUND_TO_NEAREST_INT); + _mm_storeu_si128((__m128i *)(y + i), y_vec); + } + for (; i + 3 < n; i += 4) { + __m128 x_vec = _mm_loadu_ps(x + i); + __m128i y_vec = _mm_cvtps_ph(x_vec, _MM_FROUND_TO_NEAREST_INT); + _mm_storel_epi64((__m128i *)(y + i), y_vec); + } +#elif defined(__riscv_zvfh) + for (int vl; i < n; i += vl) { + vl = __riscv_vsetvl_e32m2(n - i); + vfloat32m2_t vx = __riscv_vle32_v_f32m2(&x[i], vl); + vfloat16m1_t vy = __riscv_vfncvt_f_f_w_f16m1(vx, vl); + __riscv_vse16_v_f16m1((_Float16 *)&y[i], vy, vl); + } +#endif + for (; i < n; ++i) { + y[i] = GGML_CPU_FP32_TO_FP16(x[i]); + } +} + +void ggml_cpu_fp16_to_fp32(const ggml_fp16_t * x, float * y, int64_t n) { + int64_t i = 0; +#if defined(__F16C__) +#if defined(__AVX512F__) + for (; i + 15 < n; i += 16) { + __m256i x_vec = _mm256_loadu_si256((const __m256i *)(x + i)); + __m512 y_vec = _mm512_cvtph_ps(x_vec); + _mm512_storeu_ps(y + i, y_vec); + } +#endif + for (; i + 7 < n; i += 8) { + __m128i x_vec = _mm_loadu_si128((const __m128i *)(x + i)); + __m256 y_vec = _mm256_cvtph_ps(x_vec); + _mm256_storeu_ps(y + i, y_vec); + } + for (; i + 3 < n; i += 4) { + __m128i x_vec = _mm_loadl_epi64((const __m128i *)(x + i)); + __m128 y_vec = _mm_cvtph_ps(x_vec); + _mm_storeu_ps(y + i, y_vec); + } + +#elif defined(__riscv_v_intrinsic) && defined(__riscv_zvfhmin) + // calculate step size + const int epr = __riscv_vsetvlmax_e16m2(); + const int step = epr * 2; + const int np = (n & ~(step - 1)); + + // unroll by 2 + for (; i < np; i += step) { + vfloat16m2_t ax0 = __riscv_vle16_v_f16m2((const _Float16*)x + i, epr); + vfloat32m4_t ay0 = __riscv_vfwcvt_f_f_v_f32m4(ax0, epr); + __riscv_vse32_v_f32m4(y + i, ay0, epr); + + vfloat16m2_t ax1 = __riscv_vle16_v_f16m2((const _Float16*)x + i + epr, epr); + vfloat32m4_t ay1 = __riscv_vfwcvt_f_f_v_f32m4(ax1, epr); + __riscv_vse32_v_f32m4(y + i + epr, ay1, epr); + } + + // leftovers + int vl; + for (i = np; i < n; i += vl) { + vl = __riscv_vsetvl_e16m2(n - i); + vfloat16m2_t ax0 = __riscv_vle16_v_f16m2((const _Float16*)x + i, vl); + vfloat32m4_t ay0 = __riscv_vfwcvt_f_f_v_f32m4(ax0, vl); + __riscv_vse32_v_f32m4(y + i, ay0, vl); + } + +#endif + + for (; i < n; ++i) { + y[i] = GGML_CPU_FP16_TO_FP32(x[i]); + } +} + +void ggml_cpu_fp32_to_bf16(const float * x, ggml_bf16_t * y, int64_t n) { + int64_t i = 0; + for (; i < n; ++i) { + y[i] = GGML_FP32_TO_BF16(x[i]); + } +} + +void ggml_cpu_fp32_to_i32(const float * x, int32_t * y, int64_t n) { + int64_t i = 0; + for (; i < n; ++i) { + y[i] = x[i]; + } +} + +void ggml_cpu_bf16_to_fp32(const ggml_bf16_t * x, float * y, int64_t n) { + int64_t i = 0; +#if defined(__AVX2__) +#if defined(__AVX512F__) + for (; i + 15 < n; i += 16) { + _mm512_storeu_ps(y + i, + _mm512_castsi512_ps( + _mm512_slli_epi32( + _mm512_cvtepu16_epi32( + _mm256_loadu_si256( + (const __m256i *)(x + i))), + 16))); + } +#endif + for (; i + 7 < n; i += 8) { + _mm256_storeu_ps(y + i, + _mm256_castsi256_ps( + _mm256_slli_epi32( + _mm256_cvtepu16_epi32( + _mm_loadu_si128( + (const __m128i *)(x + i))), + 16))); + } +#elif defined(__riscv_v_intrinsic) && defined(__riscv_zvfbfmin) + // calculate step size + const int epr = __riscv_vsetvlmax_e16m2(); + const int step = epr * 2; + const int np = (n & ~(step - 1)); + + // unroll by 2 + for (; i < np; i += step) { + vbfloat16m2_t ax0 = __riscv_vle16_v_bf16m2((const __bf16*)x + i, epr); + vfloat32m4_t ay0 = __riscv_vfwcvtbf16_f_f_v_f32m4(ax0, epr); + __riscv_vse32_v_f32m4(y + i, ay0, epr); + + vbfloat16m2_t ax1 = __riscv_vle16_v_bf16m2((const __bf16*)x + i + epr, epr); + vfloat32m4_t ay1 = __riscv_vfwcvtbf16_f_f_v_f32m4(ax1, epr); + __riscv_vse32_v_f32m4(y + i + epr, ay1, epr); + } + + // leftovers + int vl; + for (i = np; i < n; i += vl) { + vl = __riscv_vsetvl_e16m2(n - i); + vbfloat16m2_t ax0 = __riscv_vle16_v_bf16m2((const __bf16*)x + i, vl); + vfloat32m4_t ay0 = __riscv_vfwcvtbf16_f_f_v_f32m4(ax0, vl); + __riscv_vse32_v_f32m4(y + i, ay0, vl); + } +#endif + for (; i < n; i++) { + y[i] = GGML_BF16_TO_FP32(x[i]); + } +} + +int ggml_cpu_has_avx(void) { +#if defined(__AVX__) + return 1; +#else + return 0; +#endif +} + +int ggml_cpu_has_avx_vnni(void) { +#if defined(__AVXVNNI__) + return 1; +#else + return 0; +#endif +} + +int ggml_cpu_has_avx2(void) { +#if defined(__AVX2__) + return 1; +#else + return 0; +#endif +} + +int ggml_cpu_has_avx512(void) { +#if defined(__AVX512F__) + return 1; +#else + return 0; +#endif +} + +int ggml_cpu_has_avx512_vbmi(void) { +#if defined(__AVX512VBMI__) + return 1; +#else + return 0; +#endif +} + +int ggml_cpu_has_avx512_vnni(void) { +#if defined(__AVX512VNNI__) + return 1; +#else + return 0; +#endif +} + +int ggml_cpu_has_avx512_bf16(void) { +#if defined(__AVX512BF16__) + return 1; +#else + return 0; +#endif +} + +int ggml_cpu_has_amx_int8(void) { +#if defined(__AMX_INT8__) + return 1; +#else + return 0; +#endif +} + +int ggml_cpu_has_bmi2(void) { +#if defined(__BMI2__) + return 1; +#else + return 0; +#endif +} + +int ggml_cpu_has_fma(void) { +#if defined(__FMA__) + return 1; +#else + return 0; +#endif +} + +int ggml_cpu_has_arm_fma(void) { +#if defined(__ARM_FEATURE_FMA) + return 1; +#else + return 0; +#endif +} + +int ggml_cpu_has_riscv_v(void) { +#if defined(__riscv_v_intrinsic) + return 1; +#else + return 0; +#endif +} + +int ggml_cpu_get_rvv_vlen(void) { +#if defined(__riscv) && defined(__riscv_v_intrinsic) + return ggml_riscv_arch_features.rvv_vlen; +#else + return 0; +#endif +} + +int ggml_cpu_has_f16c(void) { +#if defined(__F16C__) + return 1; +#else + return 0; +#endif +} + +int ggml_cpu_has_fp16_va(void) { +#if defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC) + return 1; +#else + return 0; +#endif +} + +int ggml_cpu_has_wasm_simd(void) { +#if defined(__wasm_simd128__) + return 1; +#else + return 0; +#endif +} + +int ggml_cpu_has_llamafile(void) { +#if defined(GGML_USE_LLAMAFILE) + return 1; +#else + return 0; +#endif +} + +int ggml_cpu_has_sse3(void) { +#if defined(__SSE3__) + return 1; +#else + return 0; +#endif +} + +int ggml_cpu_has_ssse3(void) { +#if defined(__SSSE3__) + return 1; +#else + return 0; +#endif +} + +int ggml_cpu_has_vsx(void) { +#if defined(__POWER9_VECTOR__) + return 1; +#else + return 0; +#endif +} + +int ggml_cpu_has_vxe(void) { +#if defined(__VXE__) || defined(__VXE2__) + return 1; +#else + return 0; +#endif +} + +int ggml_cpu_has_neon(void) { +#if defined(__ARM_ARCH) && defined(__ARM_NEON) + return 1; +#else + return 0; +#endif +} + +int ggml_cpu_has_dotprod(void) { +#if defined(__ARM_ARCH) && defined(__ARM_FEATURE_DOTPROD) + return 1; +#else + return 0; +#endif +} + +int ggml_cpu_has_sve(void) { +#if defined(__ARM_ARCH) && defined(__ARM_FEATURE_SVE) + return 1; +#else + return 0; +#endif +} + +int ggml_cpu_has_matmul_int8(void) { +#if defined(__ARM_ARCH) && defined(__ARM_FEATURE_MATMUL_INT8) + return 1; +#else + return 0; +#endif +} + +int ggml_cpu_get_sve_cnt(void) { +#if defined(__ARM_ARCH) && defined(__ARM_FEATURE_SVE) + return ggml_arm_arch_features.sve_cnt; +#else + return 0; +#endif +} + +int ggml_cpu_has_sme(void) { +#if defined(__ARM_ARCH) && defined(__ARM_FEATURE_SME) + return 1; +#else + return 0; +#endif +} + +void ggml_cpu_init(void) { + // needed to initialize ggml_time + { + struct ggml_init_params params = { 0, NULL, false }; + struct ggml_context * ctx = ggml_init(params); + ggml_free(ctx); + } + + ggml_critical_section_start(); + + static bool is_first_call = true; + + if (is_first_call) { + // initialize GELU, Quick GELU, SILU and EXP F32 tables + { + const uint64_t t_start = ggml_time_us(); UNUSED(t_start); + + for (int i = 0; i < (1 << 16); ++i) { + union { + uint16_t u16; + ggml_fp16_t fp16; + } u = {i}; + float f = GGML_COMPUTE_FP16_TO_FP32(u.fp16); + ggml_table_f32_f16[i] = f; + ggml_table_gelu_f16[i] = GGML_CPU_FP32_TO_FP16(ggml_gelu_f32(f)); + ggml_table_gelu_quick_f16[i] = GGML_CPU_FP32_TO_FP16(ggml_gelu_quick_f32(f)); + } + + // initialize E8M0 half table (256 entries) + for (int i = 0; i < (1 << 8); ++i) { + ggml_table_f32_e8m0_half[i] = GGML_E8M0_TO_FP32_HALF(i); + } + + const uint64_t t_end = ggml_time_us(); UNUSED(t_end); + + GGML_PRINT_DEBUG("%s: GELU, Quick GELU, SILU and EXP tables initialized in %f ms\n", __func__, (t_end - t_start)/1000.0); + +#ifdef GGML_USE_OPENMP + //if (!getenv("OMP_WAIT_POLICY")) { + // // set the wait policy to active, so that OpenMP threads don't sleep + // setenv("OMP_WAIT_POLICY", "active", 0) + //} + + if (!getenv("KMP_BLOCKTIME")) { + // set the time to wait before sleeping a thread + // this is less aggressive than setting the wait policy to active, but should achieve similar results in most cases +#ifdef _WIN32 + _putenv_s("KMP_BLOCKTIME", "200"); // 200ms +#else + setenv("KMP_BLOCKTIME", "200", 0); // 200ms +#endif + } +#endif + } + +#if defined(__ARM_ARCH) + ggml_init_arm_arch_features(); +#endif + +#if defined(__riscv) + ggml_init_riscv_arch_features(); +#endif + + is_first_call = false; + } + + ggml_critical_section_end(); +} diff --git a/ggml/src/ggml-cuda/ds4-hc.cu b/ggml/src/ggml-cuda/ds4-hc.cu new file mode 100644 index 000000000000..1a331bbed485 --- /dev/null +++ b/ggml/src/ggml-cuda/ds4-hc.cu @@ -0,0 +1,219 @@ +#include "ds4-hc.cuh" + +// Fused DeepSeek4 hyper-connection ops. +// +// mode 0 (pre): src0 = mix [mix_dim] (f32, from fn @ rms_norm(hc_state)) +// src1 = base [mix_dim] (f32) +// src2 = hc_state [n_embd*n_hc] (f32, raw residual streams) +// dst = [n_embd + mix_dim]: +// dst[0..n_embd) = working vector (pre-mixed input) +// dst[n_embd..n_embd+mix) = split = {pre[n_hc], post[n_hc], comb[n_hc*n_hc]} +// Math matches cpu_hc_sinkhorn + finish_hc_pre_from_mix_into in +// deepseek4_graph.cpp (sigmoid gates + Sinkhorn-normalized combine). +// +// mode 1 (post): src0 = residual hc_state [n_embd*n_hc] +// src1 = block_out [n_embd] +// src2 = split [mix_dim] (view of a mode-0 dst tail) +// dst = new hc_state [n_embd*n_hc]: +// dst[h*n_embd+d] = post[h]*block_out[d] +// + sum_src comb[h + src*n_hc] * residual[src*n_embd+d] +// +// mode 2 (out): src0 = mix [n_hc] +// src1 = base [n_hc] +// src2 = hc_state [n_embd*n_hc] +// dst = [n_embd]: weights[h] = sigmoid(mix[h]*s0+base[h]) + 1e-6; +// dst[d] = sum_h weights[h]*hc_state[h*n_embd+d] + +#define DS4_HC_SINKHORN_EPS 1.0e-6f +#define DS4_HC_MAX_HC 8 +#define DS4_HC_MAX_MIX (2*DS4_HC_MAX_HC + DS4_HC_MAX_HC*DS4_HC_MAX_HC) + +static __device__ __forceinline__ float ds4_hc_sigmoid(float x) { + return 1.0f / (1.0f + expf(-x)); +} + +static __device__ void ds4_hc_sinkhorn_split( + const float * mix, + const float * base, + float pre_scale, + float post_scale, + float comb_scale, + int n_hc, + int iters, + float * split) { + for (int i = 0; i < n_hc; ++i) { + split[i] = ds4_hc_sigmoid(mix[i] * pre_scale + base[i]) + DS4_HC_SINKHORN_EPS; + } + for (int i = 0; i < n_hc; ++i) { + split[n_hc + i] = 2.0f * ds4_hc_sigmoid(mix[n_hc + i] * post_scale + base[n_hc + i]); + } + + float c[DS4_HC_MAX_HC * DS4_HC_MAX_HC]; + for (int dst_i = 0; dst_i < n_hc; ++dst_i) { + float row_max = -1.0e30f; + for (int src_i = 0; src_i < n_hc; ++src_i) { + const int idx = src_i + dst_i * n_hc; + const float v = mix[2 * n_hc + idx] * comb_scale + base[2 * n_hc + idx]; + c[idx] = v; + row_max = v > row_max ? v : row_max; + } + float row_sum = 0.0f; + for (int src_i = 0; src_i < n_hc; ++src_i) { + const int idx = src_i + dst_i * n_hc; + c[idx] = expf(c[idx] - row_max); + row_sum += c[idx]; + } + const float inv = 1.0f / row_sum; + for (int src_i = 0; src_i < n_hc; ++src_i) { + c[src_i + dst_i * n_hc] = c[src_i + dst_i * n_hc] * inv + DS4_HC_SINKHORN_EPS; + } + } + for (int src_i = 0; src_i < n_hc; ++src_i) { + float sum = 0.0f; + for (int dst_i = 0; dst_i < n_hc; ++dst_i) sum += c[src_i + dst_i * n_hc]; + const float inv = 1.0f / (sum + DS4_HC_SINKHORN_EPS); + for (int dst_i = 0; dst_i < n_hc; ++dst_i) c[src_i + dst_i * n_hc] *= inv; + } + for (int iter = 1; iter < iters; ++iter) { + for (int dst_i = 0; dst_i < n_hc; ++dst_i) { + float sum = 0.0f; + for (int src_i = 0; src_i < n_hc; ++src_i) sum += c[src_i + dst_i * n_hc]; + const float inv = 1.0f / (sum + DS4_HC_SINKHORN_EPS); + for (int src_i = 0; src_i < n_hc; ++src_i) c[src_i + dst_i * n_hc] *= inv; + } + for (int src_i = 0; src_i < n_hc; ++src_i) { + float sum = 0.0f; + for (int dst_i = 0; dst_i < n_hc; ++dst_i) sum += c[src_i + dst_i * n_hc]; + const float inv = 1.0f / (sum + DS4_HC_SINKHORN_EPS); + for (int dst_i = 0; dst_i < n_hc; ++dst_i) c[src_i + dst_i * n_hc] *= inv; + } + } + for (int i = 0; i < n_hc * n_hc; ++i) { + split[2 * n_hc + i] = c[i]; + } +} + +static __global__ void ds4_hc_pre_kernel( + const float * __restrict__ mix, + const float * __restrict__ base, + const float * __restrict__ hc_state, + float * __restrict__ dst, + int n_embd, + int n_hc, + int iters, + float pre_scale, + float post_scale, + float comb_scale) { + __shared__ float split[DS4_HC_MAX_MIX]; + const int mix_dim = 2 * n_hc + n_hc * n_hc; + const int tid = threadIdx.x; + + if (tid == 0) { + ds4_hc_sinkhorn_split(mix, base, pre_scale, post_scale, comb_scale, n_hc, iters, split); + for (int i = 0; i < mix_dim; ++i) { + dst[n_embd + i] = split[i]; + } + } + __syncthreads(); + + for (int d = tid; d < n_embd; d += blockDim.x) { + float acc = 0.0f; + for (int h = 0; h < n_hc; ++h) { + acc += split[h] * hc_state[(size_t) h * n_embd + d]; + } + dst[d] = acc; + } +} + +static __global__ void ds4_hc_post_kernel( + const float * __restrict__ residual, + const float * __restrict__ block_out, + const float * __restrict__ split, + float * __restrict__ dst, + int n_embd, + int n_hc) { + const int i = blockIdx.x * blockDim.x + threadIdx.x; + const int total = n_embd * n_hc; + if (i >= total) { + return; + } + const int h = i / n_embd; + const int d = i - h * n_embd; + const float * post = split + n_hc; + const float * comb = split + 2 * n_hc; + float acc = block_out[d] * post[h]; + for (int src = 0; src < n_hc; ++src) { + acc += comb[h + src * n_hc] * residual[(size_t) src * n_embd + d]; + } + dst[i] = acc; +} + +static __global__ void ds4_hc_out_kernel( + const float * __restrict__ mix, + const float * __restrict__ base, + const float * __restrict__ hc_state, + float * __restrict__ dst, + int n_embd, + int n_hc, + float pre_scale) { + const int d = blockIdx.x * blockDim.x + threadIdx.x; + if (d >= n_embd) { + return; + } + float acc = 0.0f; + for (int h = 0; h < n_hc; ++h) { + const float wgt = ds4_hc_sigmoid(mix[h] * pre_scale + base[h]) + DS4_HC_SINKHORN_EPS; + acc += wgt * hc_state[(size_t) h * n_embd + d]; + } + dst[d] = acc; +} + +void ggml_cuda_op_ds4_hc(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + const ggml_tensor * src2 = dst->src[2]; + + GGML_ASSERT(src0 && src0->type == GGML_TYPE_F32); + GGML_ASSERT(src1 && src1->type == GGML_TYPE_F32); + GGML_ASSERT(src2 && src2->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + + const int mode = ggml_get_op_params_i32(dst, 0); + const int n_embd = ggml_get_op_params_i32(dst, 1); + const int n_hc = ggml_get_op_params_i32(dst, 2); + + GGML_ASSERT(n_hc > 0 && n_hc <= DS4_HC_MAX_HC); + + cudaStream_t stream = ctx.stream(); + + switch (mode) { + case 0: { + const int iters = ggml_get_op_params_i32(dst, 3); + const float pre_scale = ggml_get_op_params_f32(dst, 4); + const float post_scale = ggml_get_op_params_f32(dst, 5); + const float comb_scale = ggml_get_op_params_f32(dst, 6); + ds4_hc_pre_kernel<<<1, 256, 0, stream>>>( + (const float *) src0->data, (const float *) src1->data, + (const float *) src2->data, (float *) dst->data, + n_embd, n_hc, iters, pre_scale, post_scale, comb_scale); + } break; + case 1: { + const int total = n_embd * n_hc; + const int blocks = (total + 255) / 256; + ds4_hc_post_kernel<<>>( + (const float *) src0->data, (const float *) src1->data, + (const float *) src2->data, (float *) dst->data, + n_embd, n_hc); + } break; + case 2: { + const float pre_scale = ggml_get_op_params_f32(dst, 4); + const int blocks = (n_embd + 255) / 256; + ds4_hc_out_kernel<<>>( + (const float *) src0->data, (const float *) src1->data, + (const float *) src2->data, (float *) dst->data, + n_embd, n_hc, pre_scale); + } break; + default: + GGML_ABORT("ds4_hc: unknown mode"); + } +} diff --git a/ggml/src/ggml-cuda/ds4-hc.cuh b/ggml/src/ggml-cuda/ds4-hc.cuh new file mode 100644 index 000000000000..1a2cdf5eb5b5 --- /dev/null +++ b/ggml/src/ggml-cuda/ds4-hc.cuh @@ -0,0 +1,10 @@ +#pragma once + +#include "ggml-cuda/common.cuh" + +// Fused DeepSeek4 hyper-connection (HC) decode ops. See ds4-hc.cu for the +// per-mode contract (pre / post / out). Used by the opt-in +// DFLASH_DS4_FUSED_DECODE single-graph decode path; output is deterministic +// but not bit-identical to the CPU HC reference (expf ULP differences +// amplified by the sinkhorn iterations). +void ggml_cuda_op_ds4_hc(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 066dd2680d89..4e5a9242e67b 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -64,6 +64,7 @@ #include "ggml-cuda/cumsum.cuh" #include "ggml-cuda/fill.cuh" #include "ggml-cuda/moe-fused.cuh" +#include "ggml-cuda/ds4-hc.cuh" #include "ggml.h" #include @@ -2879,6 +2880,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg case GGML_OP_MOE_FUSED: ggml_cuda_op_moe_fused(ctx, dst); break; + case GGML_OP_DS4_HC: + ggml_cuda_op_ds4_hc(ctx, dst); + break; case GGML_OP_GROUP_NORM: ggml_cuda_op_group_norm(ctx, dst); break; @@ -4964,6 +4968,8 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g return op->src[0]->nb[0] == ggml_type_size(op->src[0]->type); case GGML_OP_MOE_FUSED: return true; + case GGML_OP_DS4_HC: + return true; case GGML_OP_MUL_MAT: case GGML_OP_MUL_MAT_ID: { diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu.ds4bak b/ggml/src/ggml-cuda/ggml-cuda.cu.ds4bak new file mode 100644 index 000000000000..066dd2680d89 --- /dev/null +++ b/ggml/src/ggml-cuda/ggml-cuda.cu.ds4bak @@ -0,0 +1,5559 @@ +#include "ggml-cuda.h" +#include "ggml-impl.h" +#include "ggml-backend-impl.h" + +#include "ggml-cuda/common.cuh" +#include "ggml-cuda/acc.cuh" +#include "ggml-cuda/add-id.cuh" +#include "ggml-cuda/arange.cuh" +#include "ggml-cuda/argmax.cuh" +#include "ggml-cuda/argsort.cuh" +#include "ggml-cuda/binbcast.cuh" +#include "ggml-cuda/clamp.cuh" +#include "ggml-cuda/concat.cuh" +#include "ggml-cuda/conv-transpose-1d.cuh" +#include "ggml-cuda/conv2d.cuh" +#include "ggml-cuda/conv2d-dw.cuh" +#include "ggml-cuda/conv2d-transpose.cuh" +#include "ggml-cuda/convert.cuh" +#include "ggml-cuda/count-equal.cuh" +#include "ggml-cuda/cpy.cuh" +#include "ggml-cuda/cross-entropy-loss.cuh" +#include "ggml-cuda/cumsum.cuh" +#include "ggml-cuda/diagmask.cuh" +#include "ggml-cuda/diag.cuh" +#include "ggml-cuda/fattn.cuh" +#include "ggml-cuda/fattn-sparse.cuh" +#include "ggml-cuda/getrows.cuh" +#include "ggml-cuda/turbo-wht.cuh" +#include "ggml-cuda/im2col.cuh" +#include "ggml-cuda/mmf.cuh" +#include "ggml-cuda/mmq.cuh" +#include "ggml-cuda/mmvf.cuh" +#include "ggml-cuda/mmvq.cuh" +#include "ggml-cuda/norm.cuh" +#include "ggml-cuda/opt-step-adamw.cuh" +#include "ggml-cuda/opt-step-sgd.cuh" +#include "ggml-cuda/out-prod.cuh" +#include "ggml-cuda/pad.cuh" +#include "ggml-cuda/pool2d.cuh" +#include "ggml-cuda/quantize.cuh" +#include "ggml-cuda/rope.cuh" +#include "ggml-cuda/roll.cuh" +#include "ggml-cuda/scale.cuh" +#include "ggml-cuda/softcap.cuh" +#include "ggml-cuda/softmax.cuh" +#include "ggml-cuda/ssm-conv.cuh" +#include "ggml-cuda/ssm-scan.cuh" +#include "ggml-cuda/sum.cuh" +#include "ggml-cuda/sumrows.cuh" +#include "ggml-cuda/top-k.cuh" +#include "ggml-cuda/mean.cuh" +#include "ggml-cuda/tsembd.cuh" +#include "ggml-cuda/topk-moe.cuh" +#include "ggml-cuda/unary.cuh" +#include "ggml-cuda/upscale.cuh" +#include "ggml-cuda/wkv.cuh" +#include "ggml-cuda/gla.cuh" +#include "ggml-cuda/gated_delta_net.cuh" +#include "ggml-cuda/set.cuh" +#include "ggml-cuda/set-rows.cuh" +#include "ggml-cuda/pad_reflect_1d.cuh" +#include "ggml-cuda/solve_tri.cuh" +#include "ggml-cuda/tri.cuh" +#include "ggml-cuda/cumsum.cuh" +#include "ggml-cuda/fill.cuh" +#include "ggml-cuda/moe-fused.cuh" +#include "ggml.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +static_assert(sizeof(half) == sizeof(ggml_fp16_t), "wrong fp16 size"); + +[[noreturn]] +void ggml_cuda_error(const char * stmt, const char * func, const char * file, int line, const char * msg) { + int id = -1; // in case cudaGetDevice fails + (void)cudaGetDevice(&id); + + GGML_LOG_ERROR(GGML_CUDA_NAME " error: %s\n", msg); + GGML_LOG_ERROR(" current device: %d, in function %s at %s:%d\n", id, func, file, line); + GGML_LOG_ERROR(" %s\n", stmt); + // abort with GGML_ABORT to get a stack trace + GGML_ABORT(GGML_CUDA_NAME " error"); +} + +// this is faster on Windows +// probably because the Windows CUDA libraries forget to make this check before invoking the drivers +void ggml_cuda_set_device(int device) { + int current_device; + CUDA_CHECK(cudaGetDevice(¤t_device)); + + if (device == current_device) { + return; + } + + CUDA_CHECK(cudaSetDevice(device)); +} + +int ggml_cuda_get_device() { + int id; + CUDA_CHECK(cudaGetDevice(&id)); + return id; +} + +// dflash: auto-enable unified (managed) memory on integrated GPUs (APU shared RAM) +// ONLY for allocations large enough that a plain device buffer plus the mmap source +// page cache (~2x size) would not fit in RAM and would thrash (e.g. an ~86GB DeepSeek-V4 +// model on a 125GB box). Managed memory has measurable alloc + access overhead on +// this APU, so small models that fit stay on the faster hipMalloc path (verified: +// 16GB model loads 3s / decodes 11.8 tok/s on hipMalloc vs 15s / 10.9 on managed). +// The cached ggml_cuda_info().devices[].integrated is hard-forced false (#15034 +// dodge), so probe a FRESH cudaDeviceProp. Opt out: DFLASH_HIP_NO_AUTO_UMA=1. +// Force-all: GGML_CUDA_ENABLE_UNIFIED_MEMORY. Tune gate: DFLASH_HIP_UMA_MIN_FRAC. +static size_t ggml_cuda_total_ram_bytes() { + static const size_t cached = []() -> size_t { + size_t bytes = 0; + FILE * f = fopen("/proc/meminfo", "r"); + if (f) { + char line[256]; + unsigned long kb = 0; + while (fgets(line, sizeof(line), f)) { + if (sscanf(line, "MemTotal: %lu kB", &kb) == 1) { bytes = (size_t) kb * 1024; break; } + } + fclose(f); + } + return bytes; + }(); + return cached; +} + +static bool ggml_cuda_device_use_uma(int device, size_t size) { + if (getenv("GGML_CUDA_ENABLE_UNIFIED_MEMORY") != nullptr) { + return true; + } + if (getenv("DFLASH_HIP_NO_AUTO_UMA") != nullptr) { + return false; + } + static const std::array integrated = []() { + std::array flags{}; + int n = 0; + if (cudaGetDeviceCount(&n) != cudaSuccess) { n = 0; } + for (int i = 0; i < n && i < GGML_CUDA_MAX_DEVICES; i++) { + cudaDeviceProp prop; + flags[i] = (cudaGetDeviceProperties(&prop, i) == cudaSuccess && prop.integrated); + } + return flags; + }(); + if (device < 0 || device >= GGML_CUDA_MAX_DEVICES || !integrated[device]) { + return false; + } + size_t total_ram = ggml_cuda_total_ram_bytes(); + if (total_ram == 0) { + return false; // unknown RAM: stay conservative on the legacy path + } + double frac = 0.45; + const char * fenv = getenv("DFLASH_HIP_UMA_MIN_FRAC"); + if (fenv != nullptr) { + double v = atof(fenv); + if (v > 0.0 && v < 1.0) { frac = v; } + } + bool use = (double) size > frac * (double) total_ram; + if (use) { + GGML_LOG_INFO("ggml_cuda: device %d integrated, alloc %.1f GiB (> %.0f%% of %.1f GiB RAM) " + "-> unified (managed) memory; small models stay on hipMalloc " + "(DFLASH_HIP_NO_AUTO_UMA=1 to disable)\n", + device, (double) size / 1073741824.0, frac * 100.0, + (double) total_ram / 1073741824.0); + } + return use; +} + +static cudaError_t ggml_cuda_device_malloc(void ** ptr, size_t size, int device, bool * out_managed = nullptr) { + ggml_cuda_set_device(device); + cudaError_t err; + bool managed = false; + if (ggml_cuda_device_use_uma(device, size)) { + err = cudaMallocManaged(ptr, size); + if (err == cudaSuccess) { + managed = true; + } +#if defined(GGML_USE_HIP) + if (err == hipSuccess) { + // hipMemAdviseSetCoarseGrain is an optional performance hint; + // ignore errors (e.g. hipErrorInvalidValue on some APU/iGPU configs). + (void)cudaMemAdvise(*ptr, size, hipMemAdviseSetCoarseGrain, device); + (void)hipGetLastError(); // clear any error + } + + // fall back to cudaMalloc if not supported (e.g. on Windows) + if (err == hipErrorNotSupported) { + static bool warned_unsupported = false; + if (!warned_unsupported) { + GGML_LOG_WARN("hipMallocManaged unsupported, falling back to hipMalloc.\n"); + warned_unsupported = true; + } + + err = cudaMalloc(ptr, size); + managed = false; + } +#endif // defined(GGML_USE_HIP) + } else { + err = cudaMalloc(ptr, size); + } + if (out_managed != nullptr) { + *out_managed = managed && (err == cudaSuccess); + } + return err; +} + +#if defined(GGML_USE_HIP) +static int ggml_cuda_parse_id(char devName[]) { + // A list of possible Target IDs can be found under the rocclr/clr repo in device.cpp + // these values are not stable so this is susceptible to breakage + // https://github.com/ROCm/clr/blob/amd-staging/rocclr/device/device.cpp + int archMajor = 0x0; + int archMinor = 0x0; + int archNum = GGML_CUDA_CC_OFFSET_AMD; + int archLen = strlen(devName); + char archName[archLen + 1]; + + // strip leading 'gfx' while copying into our buffer + if (archLen > 3) { + strcpy(archName, &devName[3]); + archLen -= 3; + } + + // trim trailing :xnack- or :sramecc- statuses + archLen = strcspn(archName, ":"); + archName[archLen] = '\0'; + + // tease out the version information + if (archLen > 8) { + // versions labeled generic use '-' as delimiter + // strip the trailing "-generic" then iterate through what remains + if ((strstr(archName, "-generic"))) { + archName[archLen - 8] = '\0'; + char * pch; + if ((pch = strtok(archName, "-"))) { + archMajor = (int)strtoul(pch, 0, 16); + if ((pch = strtok(NULL, "-"))) { + archMinor = 0x10 * (int)strtoul(pch, 0, 16); + } + } + } + } else if (archLen >= 3) { + // last two digits should be the minor * 0x10 + stepping + archMinor = (int)strtoul(&archName[archLen - 2], 0, 16); + archName[archLen - 2] = '\0'; + + // only the major version remains + archMajor = (int)strtoul(archName, 0, 16); + } + archNum += archMajor * 0x100; + archNum += archMinor; + return archNum; +} +#endif // defined(GGML_USE_HIP) + +static ggml_cuda_device_info ggml_cuda_init() { + ggml_cuda_device_info info = {}; + + cudaError_t err = cudaGetDeviceCount(&info.device_count); + if (err != cudaSuccess) { + GGML_LOG_ERROR("%s: failed to initialize " GGML_CUDA_NAME ": %s\n", __func__, cudaGetErrorString(err)); + return info; + } + + GGML_ASSERT(info.device_count <= GGML_CUDA_MAX_DEVICES); + + int64_t total_vram = 0; + for (int id = 0; id < info.device_count; ++id) { + cudaDeviceProp prop; + CUDA_CHECK(cudaGetDeviceProperties(&prop, id)); + total_vram += prop.totalGlobalMem; + } + GGML_LOG_INFO("%s: found %d " GGML_CUDA_NAME " devices (Total VRAM: %zu MiB):\n", + __func__, info.device_count, (size_t)(total_vram / (1024 * 1024))); + total_vram = 0; + + std::vector> turing_devices_without_mma; + for (int id = 0; id < info.device_count; ++id) { + int device_vmm = 0; + +#if defined(GGML_USE_VMM) + CUdevice device; + CU_CHECK(cuDeviceGet(&device, id)); + CU_CHECK(cuDeviceGetAttribute(&device_vmm, CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED, device)); + + if (device_vmm) { + CUmemAllocationProp alloc_prop = {}; + alloc_prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; + alloc_prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; + alloc_prop.location.id = id; + CU_CHECK(cuMemGetAllocationGranularity(&info.devices[id].vmm_granularity, &alloc_prop, CU_MEM_ALLOC_GRANULARITY_RECOMMENDED)); + } +#endif // defined(GGML_USE_VMM) + info.devices[id].vmm = !!device_vmm; + + cudaDeviceProp prop; + CUDA_CHECK(cudaGetDeviceProperties(&prop, id)); + + info.default_tensor_split[id] = total_vram; + total_vram += prop.totalGlobalMem; + info.devices[id].integrated = false; // Temporarily disabled due to issues with corrupted output (e.g. #15034) + info.devices[id].nsm = prop.multiProcessorCount; + info.devices[id].smpb = prop.sharedMemPerBlock; + info.devices[id].warp_size = prop.warpSize; + +#ifndef GGML_USE_MUSA + int supports_coop_launch = 0; + CUDA_CHECK(cudaDeviceGetAttribute(&supports_coop_launch, cudaDevAttrCooperativeLaunch, id)); + info.devices[id].supports_cooperative_launch = !!supports_coop_launch; +#else + info.devices[id].supports_cooperative_launch = false; +#endif // !(GGML_USE_MUSA) + +#if defined(GGML_USE_HIP) + info.devices[id].smpbo = prop.sharedMemPerBlock; + + info.devices[id].cc = ggml_cuda_parse_id(prop.gcnArchName); + if ((info.devices[id].cc & 0xff00) == 0x0) { + GGML_LOG_WARN("invalid architecture ID received for device %d %s: %s cc %d.%d\n", + id, prop.name, prop.gcnArchName, prop.major, prop.minor); + + // Fallback to prop.major and prop.minor + if (prop.major > 0) { + info.devices[id].cc = GGML_CUDA_CC_OFFSET_AMD + prop.major * 0x100; + info.devices[id].cc += prop.minor * 0x10; + } + } + GGML_LOG_INFO(" Device %d: %s, %s (0x%x), VMM: %s, Wave Size: %d, VRAM: %zu MiB\n", + id, prop.name, prop.gcnArchName, info.devices[id].cc & 0xffff, + device_vmm ? "yes" : "no", prop.warpSize, + (size_t)(prop.totalGlobalMem / (1024 * 1024))); +#elif defined(GGML_USE_MUSA) + // FIXME: Ensure compatibility with varying warp sizes across different MUSA archs. + info.devices[id].warp_size = 32; + info.devices[id].smpbo = prop.sharedMemPerBlockOptin; + info.devices[id].cc = GGML_CUDA_CC_OFFSET_MTHREADS + prop.major * 0x100; + info.devices[id].cc += prop.minor * 0x10; + GGML_LOG_INFO(" Device %d: %s, compute capability %d.%d, VMM: %s, VRAM: %zu MiB\n", + id, prop.name, prop.major, prop.minor, device_vmm ? "yes" : "no", + (size_t)(prop.totalGlobalMem / (1024 * 1024))); +#else + info.devices[id].smpbo = prop.sharedMemPerBlockOptin; + info.devices[id].cc = 100*prop.major + 10*prop.minor; + GGML_LOG_INFO(" Device %d: %s, compute capability %d.%d, VMM: %s, VRAM: %zu MiB\n", + id, prop.name, prop.major, prop.minor, device_vmm ? "yes" : "no", + (size_t)(prop.totalGlobalMem / (1024 * 1024))); + std::string device_name(prop.name); + if (device_name == "NVIDIA GeForce MX450") { + turing_devices_without_mma.push_back({ id, device_name }); + } else if (device_name == "NVIDIA GeForce MX550") { + turing_devices_without_mma.push_back({ id, device_name }); + } else if (device_name.substr(0, 21) == "NVIDIA GeForce GTX 16") { + turing_devices_without_mma.push_back({ id, device_name }); + } + + // Temporary performance fix: + // Setting device scheduling strategy for iGPUs with cc121 to "spinning" to avoid delays in cuda synchronize calls. + // TODO: Check for future drivers the default scheduling strategy and + // remove this call again when cudaDeviceScheduleSpin is default. + if (prop.major == 12 && prop.minor == 1) { + CUDA_CHECK(cudaSetDevice(id)); + CUDA_CHECK(cudaSetDeviceFlags(cudaDeviceScheduleSpin)); + } + +#endif // defined(GGML_USE_HIP) + } + + if (ggml_cuda_highest_compiled_arch(GGML_CUDA_CC_TURING) >= GGML_CUDA_CC_TURING && !turing_devices_without_mma.empty()) { + GGML_LOG_INFO("The following devices will have suboptimal performance due to a lack of tensor cores:\n"); + for (size_t device_pos = 0; device_pos < turing_devices_without_mma.size(); device_pos++) { + GGML_LOG_INFO( + " Device %d: %s\n", turing_devices_without_mma[device_pos].first, turing_devices_without_mma[device_pos].second.c_str()); + } + GGML_LOG_INFO( + "Consider compiling with CMAKE_CUDA_ARCHITECTURES=61-virtual;80-virtual and DGGML_CUDA_FORCE_MMQ to force the use of the Pascal code for Turing.\n"); + } + + for (int id = 0; id < info.device_count; ++id) { + info.default_tensor_split[id] /= total_vram; + } + + // configure logging to stdout + // CUBLAS_CHECK(cublasLoggerConfigure(1, 1, 0, nullptr)); + + for (int id = 0; id < info.device_count; ++id) { + ggml_cuda_set_device(id); + for (int id_other = 0; id_other < info.device_count; ++id_other) { + if (id == id_other) { + continue; + } + int can_access_peer; + CUDA_CHECK(cudaDeviceCanAccessPeer(&can_access_peer, id, id_other)); + if (can_access_peer) { + CUDA_CHECK(cudaDeviceEnablePeerAccess(id_other, 0)); + } + } + } + +#ifdef GGML_USE_NCCL + int dev_ids[GGML_CUDA_MAX_DEVICES]; + for (int id = 0; id < info.device_count; ++id) { + dev_ids[id] = id; + } + NCCL_CHECK(ncclCommInitAll(info.comms, info.device_count, dev_ids)); +#endif // GGML_USE_NCCL + + return info; +} + +const ggml_cuda_device_info & ggml_cuda_info() { + static ggml_cuda_device_info info = ggml_cuda_init(); + return info; +} + +// #define DEBUG_CUDA_MALLOC + +// buffer pool for cuda (legacy) +struct ggml_cuda_pool_leg : public ggml_cuda_pool { + static const int MAX_BUFFERS = 256; + + int device; + struct ggml_cuda_buffer { + void * ptr = nullptr; + size_t size = 0; + }; + + ggml_cuda_buffer buffer_pool[MAX_BUFFERS] = {}; + size_t pool_size = 0; + + explicit ggml_cuda_pool_leg(int device) : + device(device) { + } + + ~ggml_cuda_pool_leg() { + ggml_cuda_set_device(device); + for (int i = 0; i < MAX_BUFFERS; ++i) { + ggml_cuda_buffer & b = buffer_pool[i]; + if (b.ptr != nullptr) { + CUDA_CHECK(cudaFree(b.ptr)); + pool_size -= b.size; + } + } + GGML_ASSERT(pool_size == 0); + } + + void * alloc(size_t size, size_t * actual_size) override { +#ifdef DEBUG_CUDA_MALLOC + int nnz = 0; + size_t max_size = 0; +#endif + size_t best_diff = 1ull << 36; + int ibest = -1; + for (int i = 0; i < MAX_BUFFERS; ++i) { + ggml_cuda_buffer& b = buffer_pool[i]; + if (b.ptr != nullptr) { +#ifdef DEBUG_CUDA_MALLOC + ++nnz; + if (b.size > max_size) max_size = b.size; +#endif + if (b.size >= size) { + size_t diff = b.size - size; + if (diff < best_diff) { + best_diff = diff; + ibest = i; + if (!best_diff) { + void * ptr = b.ptr; + *actual_size = b.size; + b.ptr = nullptr; + b.size = 0; + return ptr; + } + } + } + } + } + if (ibest >= 0) { + ggml_cuda_buffer& b = buffer_pool[ibest]; + void * ptr = b.ptr; + *actual_size = b.size; + b.ptr = nullptr; + b.size = 0; + return ptr; + } + void * ptr; + size_t look_ahead_size = (size_t) (1.05 * size); + look_ahead_size = 256 * ((look_ahead_size + 255)/256); + ggml_cuda_set_device(device); + CUDA_CHECK(ggml_cuda_device_malloc(&ptr, look_ahead_size, device)); + *actual_size = look_ahead_size; + pool_size += look_ahead_size; +#ifdef DEBUG_CUDA_MALLOC + GGML_LOG_INFO("%s[%d]: %d buffers, max_size = %u MB, pool_size = %u MB, requested %u MB\n", __func__, device, nnz, + (uint32_t)(max_size / 1024 / 1024), (uint32_t)(pool_size / 1024 / 1024), (uint32_t)(size / 1024 / 1024)); +#endif + return ptr; + } + + void free(void * ptr, size_t size) override { + for (int i = 0; i < MAX_BUFFERS; ++i) { + ggml_cuda_buffer& b = buffer_pool[i]; + if (b.ptr == nullptr) { + b.ptr = ptr; + b.size = size; + return; + } + } + GGML_LOG_DEBUG(GGML_CUDA_NAME " buffer pool full, increase MAX_CUDA_BUFFERS\n"); + ggml_cuda_set_device(device); + CUDA_CHECK(cudaFree(ptr)); + pool_size -= size; + } +}; + +// pool with virtual memory +#if defined(GGML_USE_VMM) +struct ggml_cuda_pool_vmm : public ggml_cuda_pool { + static const size_t CUDA_POOL_VMM_MAX_SIZE = 1ull << 35; // 32 GB + + int device; + CUdeviceptr pool_addr = 0; + size_t pool_used = 0; + size_t pool_size = 0; + size_t granularity; +#if defined(GGML_USE_HIP) + std::vector> mappings; +#endif + + explicit ggml_cuda_pool_vmm(int device) : + device(device), + granularity(ggml_cuda_info().devices[device].vmm_granularity) { + } + + ~ggml_cuda_pool_vmm() { + if (pool_addr != 0) { +#if defined(GGML_USE_HIP) + // Workaround for https://github.com/ROCm/ROCR-Runtime/issues/285 + for (std::pair & mapping : mappings) { + CU_CHECK(cuMemUnmap(mapping.first, mapping.second)); + } +#else + CU_CHECK(cuMemUnmap(pool_addr, pool_size)); +#endif + CU_CHECK(cuMemAddressFree(pool_addr, CUDA_POOL_VMM_MAX_SIZE)); + } + } + + void * alloc(size_t size, size_t * actual_size) override { + // round up the allocation size to the alignment to ensure that all allocations are aligned for all data types + const size_t alignment = 128; + size = alignment * ((size + alignment - 1) / alignment); + + size_t avail = pool_size - pool_used; + + if (size > avail) { + // round up to the next multiple of the granularity + size_t reserve_size = size - avail; + reserve_size = granularity * ((reserve_size + granularity - 1) / granularity); + + GGML_ASSERT(pool_size + reserve_size <= CUDA_POOL_VMM_MAX_SIZE); + + // Pool extension via cuMem* must run with the device idle. Without + // this sync, cuMemSetAccess below can fire while previously-queued + // async kernels (compute, tensor copies) are still in flight on + // the default stream and return CUDA_ERROR_NOT_READY. The CU_CHECK + // macro hits GGML_ABORT on that error and the resulting + // signal/longjmp ends up corrupting the just-mapped region: the + // mapping is registered but access permissions are never set, so + // every subsequent read/write to it silently misbehaves. Easiest + // visible failure: prefix-cache snapshots stored into freshly + // extended pool memory return zeroed/garbled state on restore. + // Repro: any workload that queues compute then immediately + // alloc_ctx_tensors on the same backend. + CUDA_CHECK(cudaDeviceSynchronize()); + + // allocate more physical memory + CUmemAllocationProp prop = {}; + prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; + prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; + prop.location.id = device; + CUmemGenericAllocationHandle handle; + CU_CHECK(cuMemCreate(&handle, reserve_size, &prop, 0)); + + // reserve virtual address space (if not already reserved) + if (pool_addr == 0) { + CU_CHECK(cuMemAddressReserve(&pool_addr, CUDA_POOL_VMM_MAX_SIZE, 0, 0, 0)); + } + + // map at the end of the pool + CUdeviceptr start_ptr = (CUdeviceptr)((char *)(pool_addr) + pool_size); + CU_CHECK(cuMemMap(start_ptr, reserve_size, 0, handle, 0)); +#if defined(GGML_USE_HIP) + mappings.push_back({start_ptr, reserve_size}); +#endif + + // the memory allocation handle is no longer needed after mapping + CU_CHECK(cuMemRelease(handle)); + + // set access + CUmemAccessDesc access = {}; + access.location.type = CU_MEM_LOCATION_TYPE_DEVICE; + access.location.id = device; + access.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; + // The sync above is not enough under concurrency: other threads + // (snapshot D2H, a second client, the draft backend) can queue + // async work between it and this remap, and cuMemSetAccess then + // returns CUDA_ERROR_NOT_READY (observed on WSL2 at >19GB load). + // NOT_READY is transient: re-sync and retry instead of aborting. + CUresult set_err = cuMemSetAccess(start_ptr, reserve_size, &access, 1); + for (int retry = 0; set_err == CUDA_ERROR_NOT_READY && retry < 100; ++retry) { + CUDA_CHECK(cudaDeviceSynchronize()); + set_err = cuMemSetAccess(start_ptr, reserve_size, &access, 1); + } + CU_CHECK(set_err); + + // add to the pool + pool_size += reserve_size; + + //printf("cuda pool[%d]: size increased to %llu MB (reserved %llu MB)\n", + // device, (unsigned long long) (pool_size/1024/1024), + // (unsigned long long) (reserve_size/1024/1024)); + } + + GGML_ASSERT(pool_addr != 0); + + void * ptr = (void *) ((CUdeviceptr)((char *)(pool_addr) + pool_used)); + *actual_size = size; + pool_used += size; + +#ifdef DEBUG_CUDA_MALLOC + printf("cuda pool[%d]: allocated %llu bytes at %llx\n", device, (unsigned long long) size, ptr); +#endif + + return ptr; + } + + void free(void * ptr, size_t size) override { +#ifdef DEBUG_CUDA_MALLOC + printf("cuda pool[%d]: freed %llu bytes at %llx\n", device, (unsigned long long) size, ptr); +#endif + + pool_used -= size; + + // all deallocations must be in reverse order of the allocations + GGML_ASSERT(ptr == (void *) ((char *)(pool_addr) + pool_used)); + } +}; +#endif // defined(GGML_USE_VMM) + +std::unique_ptr ggml_backend_cuda_context::new_pool_for_device(int device, + [[maybe_unused]] int stream_no) { +#if defined(GGML_USE_VMM) + if (ggml_cuda_info().devices[device].vmm) { + return std::unique_ptr(new ggml_cuda_pool_vmm(device)); + } +#endif // defined(GGML_USE_VMM) + return std::unique_ptr(new ggml_cuda_pool_leg(device)); +} + +// destroying a cuBLAS handle while a graph is being captured in a different thread can result in a CUDA error +// this lock is used to ensure that no cuBLAS handle is destroyed while a graph is being captured + +static std::mutex ggml_cuda_lock; +static std::condition_variable ggml_cuda_lock_cv; +static std::atomic ggml_cuda_lock_counter; + +ggml_backend_cuda_context::~ggml_backend_cuda_context() { + std::unique_lock lock(ggml_cuda_lock); + ggml_cuda_lock_cv.wait(lock, []{ return ggml_cuda_lock_counter.load(std::memory_order_relaxed) == 0; }); + + if (copy_event != nullptr) { + CUDA_CHECK(cudaEventDestroy(copy_event)); + } + for (int i = 0; i < GGML_CUDA_MAX_DEVICES; ++i) { + for (int j = 0; j < GGML_CUDA_MAX_STREAMS; ++j) { + if (streams[i][j] != nullptr) { + CUDA_CHECK(cudaStreamDestroy(streams[i][j])); + } + } + if (cublas_handles[i] != nullptr) { + CUBLAS_CHECK(cublasDestroy(cublas_handles[i])); + } + } +} + + +// cuda buffer + +struct ggml_backend_cuda_buffer_context { + int device; + void * dev_ptr = nullptr; + bool is_managed = false; // dflash: dev_ptr is unified/managed (CPU-addressable) + std::string name; + + ggml_backend_cuda_buffer_context(int device, void * dev_ptr) : + device(device), dev_ptr(dev_ptr), + name(GGML_CUDA_NAME + std::to_string(device)) { + } + + ~ggml_backend_cuda_buffer_context() { + CUDA_CHECK(cudaFree(dev_ptr)); + } +}; + +static void ggml_backend_cuda_buffer_free_buffer(ggml_backend_buffer_t buffer) { + ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *)buffer->context; + delete ctx; +} + +static bool ggml_backend_buffer_is_cuda(ggml_backend_buffer_t buffer) { + return buffer->iface.free_buffer == ggml_backend_cuda_buffer_free_buffer; +} + +static void * ggml_backend_cuda_buffer_get_base(ggml_backend_buffer_t buffer) { + ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *)buffer->context; + return ctx->dev_ptr; +} + +// dflash: query whether a buffer is unified/managed (CPU-addressable) so loaders can +// stream weights straight into it with explicit reads instead of a host->device copy. +extern "C" bool ggml_backend_cuda_buffer_is_managed(ggml_backend_buffer_t buffer) { + if (buffer == nullptr || !ggml_backend_buffer_is_cuda(buffer)) { + return false; + } + return ((ggml_backend_cuda_buffer_context *) buffer->context)->is_managed; +} + +static enum ggml_status ggml_backend_cuda_buffer_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) { + ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *)buffer->context; + + if (tensor->view_src != NULL) { + assert(tensor->view_src->buffer->buft == buffer->buft); + return GGML_STATUS_SUCCESS; + } + + if (ggml_is_quantized(tensor->type) && tensor->view_src == nullptr && ggml_backend_buffer_get_usage(buffer) != GGML_BACKEND_BUFFER_USAGE_COMPUTE) { + // initialize padding to 0 to avoid possible NaN values + const size_t original_size = ggml_nbytes(tensor); + const size_t padded_size = ggml_backend_buft_get_alloc_size(buffer->buft, tensor); + + if (padded_size > original_size) { + ggml_cuda_set_device(ctx->device); + CUDA_CHECK(cudaMemset((char *)tensor->data + original_size, 0, padded_size - original_size)); + } + } + return GGML_STATUS_SUCCESS; +} + +static void ggml_backend_cuda_buffer_memset_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, uint8_t value, size_t offset, size_t size) { + ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *) buffer->context; + + ggml_cuda_set_device(ctx->device); + CUDA_CHECK(cudaMemsetAsync((char *) tensor->data + offset, value, size, cudaStreamPerThread)); + CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread)); +} + +static void ggml_backend_cuda_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { + ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *) buffer->context; + + ggml_cuda_set_device(ctx->device); + if (ctx->is_managed) { + // Unified/managed memory is CPU-addressable: copy weights straight in with a + // host memcpy, skipping the pageable host->device staged copy + per-tensor + // synchronize that dominate large-model load on shared-memory (integrated) + // GPUs. GPU visibility is established by the backend synchronize before the + // first graph compute. + char * dst = (char *) tensor->data + offset; + const char * src = (const char *) data; + const size_t par_threshold = (size_t) 16 * 1024 * 1024; + unsigned nth = std::thread::hardware_concurrency(); + if (nth == 0) nth = 4; + if (nth > 16) nth = 16; + if (size >= par_threshold && nth > 1) { + // Split large copies across threads so cold-cache page-ins and the + // RAM->RAM copy run in parallel (disjoint byte ranges, no sharing). + std::vector workers; + workers.reserve(nth); + const size_t chunk = (size + nth - 1) / nth; + for (unsigned t = 0; t < nth; t++) { + const size_t s0 = (size_t) t * chunk; + if (s0 >= size) break; + const size_t n = (chunk < size - s0) ? chunk : (size - s0); + workers.emplace_back([dst, src, s0, n]() { + memcpy(dst + s0, src + s0, n); + }); + } + for (auto & w : workers) { + w.join(); + } + } else { + memcpy(dst, src, size); + } + return; + } + CUDA_CHECK(cudaMemcpyAsync((char *) tensor->data + offset, data, size, cudaMemcpyHostToDevice, cudaStreamPerThread)); + CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread)); +} + +static void ggml_backend_cuda_buffer_get_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * tensor, void * data, size_t offset, size_t size) { + ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *) buffer->context; + + ggml_cuda_set_device(ctx->device); + CUDA_CHECK(cudaMemcpyAsync(data, (const char *) tensor->data + offset, size, cudaMemcpyDeviceToHost, cudaStreamPerThread)); + CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread)); +} + +static void ggml_backend_cuda_buffer_set_tensor_2d(ggml_backend_buffer_t buffer, 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) { + ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *) buffer->context; + + ggml_cuda_set_device(ctx->device); + CUDA_CHECK(cudaMemcpy2DAsync( + (char *) tensor->data + offset, stride_tensor, data, stride_data, size, n_copies, cudaMemcpyHostToDevice, cudaStreamPerThread)); + CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread)); +} + +static void ggml_backend_cuda_buffer_get_tensor_2d(ggml_backend_buffer_t buffer, 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) { + ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *)buffer->context; + + ggml_cuda_set_device(ctx->device); + CUDA_CHECK(cudaMemcpy2DAsync( + data, stride_data, (const char *) tensor->data + offset, stride_tensor, size, n_copies, cudaMemcpyDeviceToHost, cudaStreamPerThread)); + CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread)); +} + +static bool ggml_backend_cuda_buffer_cpy_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * src, ggml_tensor * dst) { + if (ggml_backend_buffer_is_cuda(src->buffer)) { + ggml_backend_cuda_buffer_context * src_ctx = (ggml_backend_cuda_buffer_context *)src->buffer->context; + ggml_backend_cuda_buffer_context * dst_ctx = (ggml_backend_cuda_buffer_context *)dst->buffer->context; + if (src_ctx->device == dst_ctx->device) { + CUDA_CHECK(cudaMemcpyAsync(dst->data, src->data, ggml_nbytes(src), cudaMemcpyDeviceToDevice, cudaStreamPerThread)); + } else { +#ifdef GGML_CUDA_NO_PEER_COPY + return false; +#else + CUDA_CHECK(cudaMemcpyPeerAsync(dst->data, dst_ctx->device, src->data, src_ctx->device, ggml_nbytes(src), cudaStreamPerThread)); +#endif + } + CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread)); + return true; + } + return false; + + GGML_UNUSED(buffer); +} + +static void ggml_backend_cuda_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { + ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *)buffer->context; + + ggml_cuda_set_device(ctx->device); + CUDA_CHECK(cudaMemsetAsync(ctx->dev_ptr, value, buffer->size, cudaStreamPerThread)); + CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread)); +} + +static const ggml_backend_buffer_i ggml_backend_cuda_buffer_interface = { + /* .free_buffer = */ ggml_backend_cuda_buffer_free_buffer, + /* .get_base = */ ggml_backend_cuda_buffer_get_base, + /* .init_tensor = */ ggml_backend_cuda_buffer_init_tensor, + /* .memset_tensor = */ ggml_backend_cuda_buffer_memset_tensor, + /* .set_tensor = */ ggml_backend_cuda_buffer_set_tensor, + /* .get_tensor = */ ggml_backend_cuda_buffer_get_tensor, + /* .set_tensor_2d = */ ggml_backend_cuda_buffer_set_tensor_2d, + /* .get_tensor_2d = */ ggml_backend_cuda_buffer_get_tensor_2d, + /* .cpy_tensor = */ ggml_backend_cuda_buffer_cpy_tensor, + /* .clear = */ ggml_backend_cuda_buffer_clear, + /* .reset = */ NULL, +}; + +// cuda buffer type +struct ggml_backend_cuda_buffer_type_context { + int device; + std::string name; +}; + +static const char * ggml_backend_cuda_buffer_type_get_name(ggml_backend_buffer_type_t buft) { + ggml_backend_cuda_buffer_type_context * ctx = (ggml_backend_cuda_buffer_type_context *)buft->context; + + return ctx->name.c_str(); +} + +static bool ggml_backend_buft_is_cuda(ggml_backend_buffer_type_t buft) { + return buft->iface.get_name == ggml_backend_cuda_buffer_type_get_name; +} + +static ggml_backend_buffer_t ggml_backend_cuda_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { + ggml_backend_cuda_buffer_type_context * buft_ctx = (ggml_backend_cuda_buffer_type_context *)buft->context; + + ggml_cuda_set_device(buft_ctx->device); + + void * dev_ptr; + bool is_managed = false; + cudaError_t err = ggml_cuda_device_malloc(&dev_ptr, size, buft_ctx->device, &is_managed); + if (err != cudaSuccess) { + // clear the error + (void)cudaGetLastError(); + GGML_LOG_ERROR("%s: allocating %.2f MiB on device %d: cudaMalloc failed: %s\n", __func__, size / 1024.0 / 1024.0, buft_ctx->device, cudaGetErrorString(err)); + return nullptr; + } + + ggml_backend_cuda_buffer_context * ctx = new ggml_backend_cuda_buffer_context(buft_ctx->device, dev_ptr); + ctx->is_managed = is_managed; + + return ggml_backend_buffer_init(buft, ggml_backend_cuda_buffer_interface, ctx, size); +} + +static size_t ggml_backend_cuda_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) { + return 128; + + GGML_UNUSED(buft); +} + +static size_t ggml_backend_cuda_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) { + size_t size = ggml_nbytes(tensor); + int64_t ne0 = tensor->ne[0]; + + if (ggml_is_quantized(tensor->type)) { + if (ne0 % MATRIX_ROW_PADDING != 0) { + GGML_ASSERT(tensor->nb[0] == ggml_element_size(tensor)); + size += ggml_row_size(tensor->type, MATRIX_ROW_PADDING - ne0 % MATRIX_ROW_PADDING); + } + } + + return size; + + GGML_UNUSED(buft); +} + +static const ggml_backend_buffer_type_i ggml_backend_cuda_buffer_type_interface = { + /* .get_name = */ ggml_backend_cuda_buffer_type_get_name, + /* .alloc_buffer = */ ggml_backend_cuda_buffer_type_alloc_buffer, + /* .get_alignment = */ ggml_backend_cuda_buffer_type_get_alignment, + /* .get_max_size = */ NULL, // defaults to SIZE_MAX + /* .get_alloc_size = */ ggml_backend_cuda_buffer_type_get_alloc_size, + /* .is_host = */ NULL, +}; + +ggml_backend_buffer_type_t ggml_backend_cuda_buffer_type(int device) { + static std::mutex mutex; + std::lock_guard lock(mutex); + + if (device >= ggml_backend_cuda_get_device_count()) { + return nullptr; + } + + static ggml_backend_buffer_type ggml_backend_cuda_buffer_types[GGML_CUDA_MAX_DEVICES]; + + static bool ggml_backend_cuda_buffer_type_initialized = false; + + if (!ggml_backend_cuda_buffer_type_initialized) { + for (int i = 0; i < ggml_backend_cuda_get_device_count(); i++) { + ggml_backend_cuda_buffer_types[i] = { + /* .iface = */ ggml_backend_cuda_buffer_type_interface, + /* .device = */ ggml_backend_reg_dev_get(ggml_backend_cuda_reg(), i), + /* .context = */ new ggml_backend_cuda_buffer_type_context{i, GGML_CUDA_NAME + std::to_string(i)}, + }; + } + ggml_backend_cuda_buffer_type_initialized = true; + } + + return &ggml_backend_cuda_buffer_types[device]; +} + +// cuda split buffer + +static int64_t get_row_rounding(const std::array & tensor_split) { + int64_t row_rounding = 0; + for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { + if (tensor_split[id] >= (id + 1 < ggml_backend_cuda_get_device_count() ? tensor_split[id + 1] : 1.0f)) { + continue; + } + + const int cc = ggml_cuda_info().devices[id].cc; + row_rounding = std::max(row_rounding, (int64_t)get_mmq_y_host(cc)); + } + return row_rounding; +} + +static void get_row_split(int64_t * row_low, int64_t * row_high, const ggml_tensor * tensor, const std::array & tensor_split, int id) { + const int64_t nrows = ggml_nrows(tensor); + const int64_t rounding = get_row_rounding(tensor_split); + + *row_low = id == 0 ? 0 : nrows*tensor_split[id]; + *row_low -= *row_low % rounding; + + if (id == ggml_backend_cuda_get_device_count() - 1) { + *row_high = nrows; + } else { + *row_high = nrows*tensor_split[id + 1]; + *row_high -= *row_high % rounding; + } +} + +static size_t ggml_nbytes_split(const struct ggml_tensor * tensor, int nrows_split) { + static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); + + return nrows_split*ggml_row_size(tensor->type, tensor->ne[0]); +} + +struct ggml_backend_cuda_split_buffer_type_context { + int main_device; + std::array tensor_split; + std::string name; +}; + +struct ggml_backend_cuda_split_buffer_context { + ~ggml_backend_cuda_split_buffer_context() { + for (ggml_tensor_extra_gpu * extra : tensor_extras) { + for (int id = 0; id < GGML_CUDA_MAX_DEVICES; ++id) { + for (int64_t is = 0; is < GGML_CUDA_MAX_STREAMS; ++is) { + if (extra->events[id][is] != nullptr) { + CUDA_CHECK(cudaEventDestroy(extra->events[id][is])); + } + } + if (extra->data_device[id] != nullptr) { + CUDA_CHECK(cudaFree(extra->data_device[id])); + } + } + delete extra; + } + } + + std::vector tensor_extras; +}; + + +static void ggml_backend_cuda_split_buffer_free_buffer(ggml_backend_buffer_t buffer) { + ggml_backend_cuda_split_buffer_context * ctx = (ggml_backend_cuda_split_buffer_context *)buffer->context; + delete ctx; +} + +static void * ggml_backend_cuda_split_buffer_get_base(ggml_backend_buffer_t buffer) { + // the pointers are stored in the tensor extras, this is just a dummy address and never dereferenced + return (void *)0x1000; + + GGML_UNUSED(buffer); +} + +static enum ggml_status ggml_backend_cuda_split_buffer_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) { + GGML_ASSERT(tensor->view_src == nullptr); // views of split tensors are not supported + GGML_ASSERT(ggml_is_contiguous(tensor) && "split buffers only supported for contiguous tensors"); + + ggml_backend_cuda_split_buffer_context * ctx = (ggml_backend_cuda_split_buffer_context *)buffer->context; + ggml_backend_cuda_split_buffer_type_context * buft_ctx = (ggml_backend_cuda_split_buffer_type_context *)buffer->buft->context; + + const int64_t ne0 = tensor->ne[0]; + + ggml_tensor_extra_gpu * extra = new ggml_tensor_extra_gpu{}; + ctx->tensor_extras.push_back(extra); + + for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { + int64_t row_low, row_high; + get_row_split(&row_low, &row_high, tensor, buft_ctx->tensor_split, id); + + int64_t nrows_split = row_high - row_low; + if (nrows_split == 0) { + continue; + } + + size_t size = ggml_nbytes_split(tensor, nrows_split); + const size_t original_size = size; + + // pad last row to a multiple of 512 elements to avoid out-of-bounds memory accesses + if (ne0 % MATRIX_ROW_PADDING != 0) { + size += ggml_row_size(tensor->type, MATRIX_ROW_PADDING - ne0 % MATRIX_ROW_PADDING); + } + + // FIXME: do not crash if cudaMalloc fails + // currently, init_tensor cannot fail, it needs to be fixed in ggml-backend first + ggml_cuda_set_device(id); + char * buf; + CUDA_CHECK(ggml_cuda_device_malloc((void**)&buf, size, id)); + + // set padding to 0 to avoid possible NaN values + if (size > original_size) { + CUDA_CHECK(cudaMemset(buf + original_size, 0, size - original_size)); + } + + extra->data_device[id] = buf; + + for (int64_t is = 0; is < GGML_CUDA_MAX_STREAMS; ++is) { + CUDA_CHECK(cudaEventCreateWithFlags(&extra->events[id][is], cudaEventDisableTiming)); + } + } + tensor->extra = extra; + return GGML_STATUS_SUCCESS; +} + +static void ggml_backend_cuda_split_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { + // split tensors must always be set in their entirety at once + GGML_ASSERT(offset == 0); + GGML_ASSERT(size == ggml_nbytes(tensor)); + GGML_ASSERT(ggml_is_contiguous(tensor) && "split buffers only supported for contiguous tensors"); + + ggml_backend_cuda_split_buffer_type_context * buft_ctx = (ggml_backend_cuda_split_buffer_type_context *)buffer->buft->context; + + const int64_t ne0 = tensor->ne[0]; + const size_t nb1 = tensor->nb[1]; + ggml_tensor_extra_gpu * extra = (ggml_tensor_extra_gpu *)tensor->extra; + + for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { + int64_t row_low, row_high; + get_row_split(&row_low, &row_high, tensor, buft_ctx->tensor_split, id); + + int64_t nrows_split = row_high - row_low; + if (nrows_split == 0) { + continue; + } + + const size_t offset_split = row_low*nb1; + size_t size = ggml_nbytes_split(tensor, nrows_split); + const size_t original_size = size; + + // pad last row to a multiple of 512 elements to avoid out-of-bounds memory accesses + if (ne0 % MATRIX_ROW_PADDING != 0) { + size += ggml_row_size(tensor->type, MATRIX_ROW_PADDING - ne0 % MATRIX_ROW_PADDING); + } + + const char * buf_host = (const char *)data + offset_split; + CUDA_CHECK(cudaMemcpyAsync(extra->data_device[id], buf_host, original_size, cudaMemcpyHostToDevice, cudaStreamPerThread)); + } + + for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { + CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread)); + } +} + +static void ggml_backend_cuda_split_buffer_get_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * tensor, void * data, size_t offset, size_t size) { + // split tensors must always be set in their entirety at once + GGML_ASSERT(offset == 0); + GGML_ASSERT(size == ggml_nbytes(tensor)); + GGML_ASSERT(ggml_is_contiguous(tensor) && "split buffers only supported for contiguous tensors"); + + ggml_backend_cuda_split_buffer_type_context * buft_ctx = (ggml_backend_cuda_split_buffer_type_context *)buffer->buft->context; + + const int64_t ne0 = tensor->ne[0]; + const size_t nb1 = tensor->nb[1]; + ggml_tensor_extra_gpu * extra = (ggml_tensor_extra_gpu *)tensor->extra; + + for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { + int64_t row_low, row_high; + get_row_split(&row_low, &row_high, tensor, buft_ctx->tensor_split, id); + + int64_t nrows_split = row_high - row_low; + if (nrows_split == 0) { + continue; + } + + const size_t offset_split = row_low*nb1; + size_t size = ggml_nbytes_split(tensor, nrows_split); + const size_t original_size = size; + + // pad last row to a multiple of 512 elements to avoid out-of-bounds memory accesses + if (ne0 % MATRIX_ROW_PADDING != 0) { + size += ggml_row_size(tensor->type, MATRIX_ROW_PADDING - ne0 % MATRIX_ROW_PADDING); + } + + char * buf_host = (char *)data + offset_split; + CUDA_CHECK(cudaMemcpyAsync(buf_host, extra->data_device[id], original_size, cudaMemcpyDeviceToHost, cudaStreamPerThread)); + } + + for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { + CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread)); + } +} + +static void ggml_backend_cuda_split_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { + GGML_UNUSED(buffer); + GGML_UNUSED(value); +} + +static const ggml_backend_buffer_i ggml_backend_cuda_split_buffer_interface = { + /* .free_buffer = */ ggml_backend_cuda_split_buffer_free_buffer, + /* .get_base = */ ggml_backend_cuda_split_buffer_get_base, + /* .init_tensor = */ ggml_backend_cuda_split_buffer_init_tensor, + /* .memset_tensor = */ NULL, + /* .set_tensor = */ ggml_backend_cuda_split_buffer_set_tensor, + /* .get_tensor = */ ggml_backend_cuda_split_buffer_get_tensor, + /* .set_tensor_2d = */ NULL, + /* .get_tensor_2d = */ NULL, + /* .cpy_tensor = */ NULL, + /* .clear = */ ggml_backend_cuda_split_buffer_clear, + /* .reset = */ NULL, +}; + +// cuda split buffer type + +static const char * ggml_backend_cuda_split_buffer_type_get_name(ggml_backend_buffer_type_t buft) { + ggml_backend_cuda_split_buffer_type_context * ctx = (ggml_backend_cuda_split_buffer_type_context *)buft->context; + + return ctx->name.c_str(); +} + +static bool ggml_backend_buft_is_cuda_split(ggml_backend_buffer_type_t buft) { + return buft->iface.get_name == ggml_backend_cuda_split_buffer_type_get_name; +} + +static ggml_backend_buffer_t ggml_backend_cuda_split_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { + // since we don't know the exact split after rounding, we cannot allocate the device buffers at this point + // instead, we allocate them for each tensor separately in init_tensor + // however, the size still represents the maximum cumulative size of all the device buffers after the tensors are allocated, + // as returned by get_alloc_size. this limit is enforced during tensor allocation by ggml-alloc, so it must be correct. + ggml_backend_cuda_split_buffer_context * ctx = new ggml_backend_cuda_split_buffer_context(); + + return ggml_backend_buffer_init(buft, ggml_backend_cuda_split_buffer_interface, ctx, size); +} + +static size_t ggml_backend_cuda_split_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) { + return 128; + + GGML_UNUSED(buft); +} + +static size_t ggml_backend_cuda_split_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) { + ggml_backend_cuda_split_buffer_type_context * ctx = (ggml_backend_cuda_split_buffer_type_context *)buft->context; + GGML_ASSERT(ggml_is_contiguous(tensor) && "split buffers only supported for contiguous tensors"); + + size_t total_size = 0; + + const int64_t ne0 = tensor->ne[0]; + + for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { + int64_t row_low, row_high; + get_row_split(&row_low, &row_high, tensor, ctx->tensor_split, id); + + int64_t nrows_split = row_high - row_low; + if (nrows_split == 0) { + continue; + } + + total_size += ggml_nbytes_split(tensor, nrows_split); + + // pad last row to a multiple of 512 elements to avoid out-of-bounds memory accesses + if (ne0 % MATRIX_ROW_PADDING != 0) { + total_size += ggml_row_size(tensor->type, MATRIX_ROW_PADDING - ne0 % MATRIX_ROW_PADDING); + } + } + + return total_size; +} + +static bool ggml_backend_cuda_split_buffer_type_is_host(ggml_backend_buffer_type_t buft) { + return false; + + GGML_UNUSED(buft); +} + +static const ggml_backend_buffer_type_i ggml_backend_cuda_split_buffer_type_interface = { + /* .get_name = */ ggml_backend_cuda_split_buffer_type_get_name, + /* .alloc_buffer = */ ggml_backend_cuda_split_buffer_type_alloc_buffer, + /* .get_alignment = */ ggml_backend_cuda_split_buffer_type_get_alignment, + /* .get_max_size = */ NULL, // defaults to SIZE_MAX + /* .get_alloc_size = */ ggml_backend_cuda_split_buffer_type_get_alloc_size, + /* .is_host = */ ggml_backend_cuda_split_buffer_type_is_host, +}; + +bool ggml_backend_cuda_allreduce_tensor(ggml_backend_t * backends, struct ggml_tensor ** tensors, size_t n_backends) { +#ifdef GGML_USE_NCCL + const int64_t ne = ggml_nelements(tensors[0]); + // FIXME the input of llm_graph_context::build_in_out_ids can produce a tensor with 0 elements if n_outputs == 0 + // This then causes a crash in this function + if (ne == 0) { + return true; + } + for (size_t i = 0; i < n_backends; ++i) { + GGML_ASSERT(tensors[i] != nullptr); + GGML_ASSERT(ggml_nelements(tensors[i]) == ne); + GGML_ASSERT(ggml_is_contiguously_allocated(tensors[i])); + } + + const ggml_cuda_device_info info = ggml_cuda_info(); + + // For small tensors, simply reduce them as FP32. + // The following heuristic for how "small" a tensor should be is based on RTX 4090s connected via 16x PCIe 4.0. + if ((n_backends <= 2 && ne < 32768) || (n_backends == 3 && ne < 131072) || (n_backends >= 4 && ne < 262144)) { + NCCL_CHECK(ncclGroupStart()); + for (size_t i = 0; i < n_backends; ++i) { + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backends[i]->context; + NCCL_CHECK(ncclAllReduce(tensors[i]->data, tensors[i]->data, ne, ncclFloat, ncclSum, info.comms[cuda_ctx->device], cuda_ctx->stream())); + } + NCCL_CHECK(ncclGroupEnd()); + + return true; + } + + // For large tensors it's faster to compress them to BF16 for the reduction: + to_bf16_cuda_t to_bf16 = ggml_get_to_bf16_cuda(GGML_TYPE_F32); + to_fp32_cuda_t to_fp32 = ggml_get_to_fp32_cuda(GGML_TYPE_BF16); + + ggml_cuda_pool_alloc tmp[GGML_CUDA_MAX_DEVICES]; + for (size_t i = 0; i < n_backends; ++i) { + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backends[i]->context; + tmp[i].pool = &cuda_ctx->pool(); + tmp[i].alloc(ne); + + ggml_cuda_set_device(i); + to_bf16(tensors[i]->data, tmp[i].get(), ne, cuda_ctx->stream()); + CUDA_CHECK(cudaGetLastError()); + } + + NCCL_CHECK(ncclGroupStart()); + for (size_t i = 0; i < n_backends; ++i) { + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backends[i]->context; + NCCL_CHECK(ncclAllReduce(tmp[i].get(), tmp[i].get(), ne, ncclBfloat16, ncclSum, info.comms[cuda_ctx->device], cuda_ctx->stream())); + } + NCCL_CHECK(ncclGroupEnd()); + + for (size_t i = 0; i < n_backends; ++i) { + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backends[i]->context; + + ggml_cuda_set_device(i); + to_fp32(tmp[i].get(), (float *) tensors[i]->data, ne, cuda_ctx->stream()); + CUDA_CHECK(cudaGetLastError()); + } + + return true; +#else + // If NCCL is installed it is used by default for optimal performance. + // However, NVIDIA does not distribute NCCL with CUDA so users may be unwittingly missing this package. + // RCCL is disabled by default, users are explicitly opting in. + // Therefore print no warning for RCCL. +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + static bool warning_printed = false; + if (!warning_printed) { + GGML_LOG_WARN("%s: NVIDIA Collective Communications Library (NCCL) is unavailable, multi GPU performance will be suboptimal\n", __func__); + warning_printed = true; + } +#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + GGML_UNUSED_VARS(backends, tensors, n_backends); + return false; +#endif // GGML_USE_NCCL +} + +ggml_backend_buffer_type_t ggml_backend_cuda_split_buffer_type(int main_device, const float * tensor_split) { + static std::mutex mutex; + std::lock_guard lock(mutex); + + static std::map>, struct ggml_backend_buffer_type> buft_map; + + std::array tensor_split_arr = {}; + + bool all_zero = tensor_split == nullptr || std::all_of(tensor_split, tensor_split + GGML_CUDA_MAX_DEVICES, [](float x) { return x == 0.0f; }); + if (all_zero) { + tensor_split_arr = ggml_cuda_info().default_tensor_split; + } else { + float split_sum = 0.0f; + for (int i = 0; i < ggml_backend_cuda_get_device_count(); ++i) { + tensor_split_arr[i] = split_sum; + split_sum += tensor_split[i]; + } + for (int i = 0; i < ggml_backend_cuda_get_device_count(); ++i) { + tensor_split_arr[i] /= split_sum; + } + } + + auto it = buft_map.find({main_device, tensor_split_arr}); + if (it != buft_map.end()) { + return &it->second; + } + auto * ctx = new ggml_backend_cuda_split_buffer_type_context{ + main_device, + tensor_split_arr, + GGML_CUDA_NAME + std::to_string(main_device) + "_Split", + }; + + struct ggml_backend_buffer_type buft { + /* .iface = */ ggml_backend_cuda_split_buffer_type_interface, + /* .device = */ ggml_backend_reg_dev_get(ggml_backend_cuda_reg(), main_device), + /* .context = */ ctx, + }; + + auto result = buft_map.emplace(std::make_pair(main_device, tensor_split_arr), buft); + return &result.first->second; +} + +// host buffer type + +static const char * ggml_backend_cuda_host_buffer_type_name(ggml_backend_buffer_type_t buft) { + return GGML_CUDA_NAME "_Host"; + + GGML_UNUSED(buft); +} + +static bool ggml_backend_buft_is_cuda_host(ggml_backend_buffer_type_t buft) { + return buft->iface.get_name == ggml_backend_cuda_host_buffer_type_name; +} + +static void ggml_backend_cuda_host_buffer_free_buffer(ggml_backend_buffer_t buffer) { + CUDA_CHECK(cudaFreeHost(buffer->context)); +} + +static void * ggml_cuda_host_malloc(size_t size) { + if (getenv("GGML_CUDA_NO_PINNED") != nullptr) { + return nullptr; + } + + void * ptr = nullptr; + cudaError_t err = cudaMallocHost((void **) &ptr, size); + if (err != cudaSuccess) { + // clear the error + (void)cudaGetLastError(); + GGML_LOG_DEBUG("%s: failed to allocate %.2f MiB of pinned memory: %s\n", __func__, + size / 1024.0 / 1024.0, cudaGetErrorString(err)); + return nullptr; + } + + return ptr; +} + +static ggml_backend_buffer_t ggml_backend_cuda_host_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { + void * ptr = ggml_cuda_host_malloc(size); + + if (ptr == nullptr) { + // fallback to cpu buffer + return ggml_backend_buft_alloc_buffer(ggml_backend_cpu_buffer_type(), size); + } + + ggml_backend_buffer_t buffer = ggml_backend_cpu_buffer_from_ptr(ptr, size); + buffer->buft = buft; + buffer->iface.free_buffer = ggml_backend_cuda_host_buffer_free_buffer; + + return buffer; +} + +ggml_backend_buffer_type_t ggml_backend_cuda_host_buffer_type() { + static struct ggml_backend_buffer_type ggml_backend_cuda_buffer_type_host = { + /* .iface = */ { + /* .get_name = */ ggml_backend_cuda_host_buffer_type_name, + /* .alloc_buffer = */ ggml_backend_cuda_host_buffer_type_alloc_buffer, + /* .get_alignment = */ ggml_backend_cpu_buffer_type()->iface.get_alignment, + /* .get_max_size = */ NULL, // defaults to SIZE_MAX + /* .get_alloc_size = */ ggml_backend_cpu_buffer_type()->iface.get_alloc_size, + /* .is_host = */ ggml_backend_cpu_buffer_type()->iface.is_host, + }, + /* .device = */ ggml_backend_reg_dev_get(ggml_backend_cuda_reg(), 0), + /* .context = */ nullptr, + }; + + return &ggml_backend_cuda_buffer_type_host; +} + +//static bool ggml_backend_buffer_is_cuda_host(ggml_backend_buffer_t buffer) { +// return buffer->buft->iface.get_name == ggml_backend_cuda_host_buffer_type_name; +//} + +/// kernels + +typedef void (*ggml_cuda_op_mul_mat_t)( + ggml_backend_cuda_context & ctx, + const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const char * src0_dd_i, const float * src1_ddf_i, + const char * src1_ddq_i, float * dst_dd_i, const int64_t row_low, const int64_t row_high, const int64_t src1_ncols, + const int64_t src1_padded_row_size, cudaStream_t stream); + +#ifndef GGML_CUDA_PEER_MAX_BATCH_SIZE +#define GGML_CUDA_PEER_MAX_BATCH_SIZE 128 +#endif // GGML_CUDA_PEER_MAX_BATCH_SIZE + +#define MUL_MAT_SRC1_COL_STRIDE 128 + +static cudaError_t ggml_cuda_cpy_tensor_2d( + void * dst, const struct ggml_tensor * src, int64_t i3, int64_t i2, int64_t i1_low, int64_t i1_high, cudaStream_t stream) { + + const char * src_ptr = (const char *) src->data; + char * dst_ptr = (char *) dst; + + const int64_t ne0 = src->ne[0]; + const int64_t nb0 = src->nb[0]; + const int64_t nb1 = src->nb[1]; + const int64_t nb2 = src->nb[2]; + const int64_t nb3 = src->nb[3]; + const enum ggml_type type = src->type; + const int64_t ts = ggml_type_size(type); + const int64_t bs = ggml_blck_size(type); + const int64_t i1_diff = i1_high - i1_low; + + const char * x = src_ptr + i1_low*nb1 + i2*nb2 + i3*nb3; + if (nb0 == ts && nb1 == ts*ne0/bs) { + return cudaMemcpyAsync(dst_ptr, x, i1_diff*nb1, cudaMemcpyDeviceToDevice, stream); + } else if (nb0 == ts) { + return cudaMemcpy2DAsync(dst_ptr, ts*ne0/bs, x, nb1, ts*ne0/bs, i1_diff, cudaMemcpyDeviceToDevice, stream); + } else { + for (int64_t i1 = 0; i1 < i1_diff; i1++) { + const void * rx = (const void *) ((const char *) x + i1*nb1); + void * rd = (void *) (dst_ptr + i1*ts*ne0/bs); + // pretend the row is a matrix with cols=1 + cudaError_t r = cudaMemcpy2DAsync(rd, ts/bs, rx, nb0, ts/bs, ne0, cudaMemcpyDeviceToDevice, stream); + if (r != cudaSuccess) { + return r; + } + } + return cudaSuccess; + } +} + +struct cublas_force_compute_type { + bool fp32 = false; + bool fp16 = false; +}; + +static const cublas_force_compute_type & ggml_cuda_cublas_get_force_compute_type() { + static const cublas_force_compute_type compute_type = [] { + cublas_force_compute_type result; + + const bool ggml_cuda_force_cublas_compute_32f_env = getenv("GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F") != nullptr; + const bool ggml_cuda_force_cublas_compute_16f_env = getenv("GGML_CUDA_FORCE_CUBLAS_COMPUTE_16F") != nullptr; + + GGML_ASSERT(ggml_cuda_force_cublas_compute_16f_env == false || ggml_cuda_force_cublas_compute_32f_env == false); + + if (ggml_cuda_force_cublas_compute_32f_env) { + GGML_LOG_INFO("Detected GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F\n"); + result.fp32 = true; + } else if (ggml_cuda_force_cublas_compute_16f_env) { + GGML_LOG_INFO("Detected GGML_CUDA_FORCE_CUBLAS_COMPUTE_16F\n"); + result.fp16 = true; + } + + return result; + }(); + + return compute_type; +} + +static void ggml_cuda_op_mul_mat_cublas( + ggml_backend_cuda_context & ctx, + const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const char * src0_dd_i, const float * src1_ddf_i, + const char * src1_ddq_i, float * dst_dd_i, const int64_t row_low, const int64_t row_high, const int64_t src1_ncols, + const int64_t src1_padded_row_size, cudaStream_t stream) { + + GGML_ASSERT(src0_dd_i != nullptr); + GGML_ASSERT(src1_ddf_i != nullptr); + GGML_ASSERT(dst_dd_i != nullptr); + + const int64_t ne00 = src0->ne[0]; + const int64_t ne10 = src1->ne[0]; + + const int64_t ne0 = dst->ne[0]; + + const int64_t row_diff = row_high - row_low; + + int id = ggml_cuda_get_device(); + + // the main device has a larger memory buffer to hold the results from all GPUs + // ldc == nrows of the matrix that cuBLAS writes into + int64_t ldc = id == ctx.device ? ne0 : row_diff; + + const int cc = ggml_cuda_info().devices[id].cc; + + const bool supports_bf16 = GGML_CUDA_CC_IS_NVIDIA(cc) || GGML_CUDA_CC_IS_AMD(cc) || + (GGML_CUDA_CC_IS_MTHREADS(cc) && cc >= GGML_CUDA_CC_QY2); + + const bool use_fp16 = + src0->type != GGML_TYPE_NVFP4 && + (src0->type == GGML_TYPE_F16 || ggml_is_quantized(src0->type)) && + ggml_is_contiguous(src0) && + row_diff == src0->ne[1] && + dst->op_params[0] == GGML_PREC_DEFAULT; + + if (supports_bf16 && src0->type == GGML_TYPE_BF16 && ggml_is_contiguous(src0) && row_diff == src0->ne[1]) { + ggml_cuda_pool_alloc src1_as_bf16(ctx.pool(id)); + if (src1->type != GGML_TYPE_BF16) { + const to_bf16_cuda_t to_bf16_cuda = ggml_get_to_bf16_cuda(src1->type); + GGML_ASSERT(to_bf16_cuda != nullptr); + size_t ne = src1_ncols*ne10; + src1_as_bf16.alloc(ne); + to_bf16_cuda(src1_ddf_i, src1_as_bf16.get(), ne, stream); + } + const nv_bfloat16 * src1_ptr = src1->type == GGML_TYPE_BF16 ? (const nv_bfloat16 *) src1_ddf_i : src1_as_bf16.get(); + const nv_bfloat16 * src0_ptr = (const nv_bfloat16 *)src0_dd_i; + ggml_cuda_pool_alloc dst_bf16(ctx.pool(id), row_diff*src1_ncols); + + const float alpha_f32 = 1.0f; + const float beta_f32 = 0.0f; + + CUBLAS_CHECK(cublasSetStream(ctx.cublas_handle(id), stream)); + CUBLAS_CHECK( + cublasGemmEx(ctx.cublas_handle(id), CUBLAS_OP_T, CUBLAS_OP_N, + row_diff, src1_ncols, ne10, + &alpha_f32, src0_ptr, CUDA_R_16BF, ne00, + src1_ptr, CUDA_R_16BF, ne10, + &beta_f32, dst_bf16.get(), CUDA_R_16BF, ldc, + CUBLAS_COMPUTE_32F, + CUBLAS_GEMM_DEFAULT_TENSOR_OP)); + + const to_fp32_cuda_t to_fp32_cuda = ggml_get_to_fp32_cuda(GGML_TYPE_BF16); + to_fp32_cuda(dst_bf16.get(), dst_dd_i, row_diff*src1_ncols, stream); + } else if (fast_fp16_hardware_available(cc) && use_fp16) { + // convert src0 and src1 to fp16, multiply as fp16, convert dst to fp32 + ggml_cuda_pool_alloc src0_as_f16(ctx.pool(id)); + if (src0->type != GGML_TYPE_F16) { + const to_fp16_cuda_t to_fp16_cuda = ggml_get_to_fp16_cuda(src0->type); + GGML_ASSERT(to_fp16_cuda != nullptr); + size_t ne = row_diff*ne00; + src0_as_f16.alloc(ne); + to_fp16_cuda(src0_dd_i, src0_as_f16.get(), ne, stream); + } + const half * src0_ptr = src0->type == GGML_TYPE_F16 ? (const half *) src0_dd_i : src0_as_f16.get(); + + ggml_cuda_pool_alloc src1_as_f16(ctx.pool(id)); + if (src1->type != GGML_TYPE_F16) { + const to_fp16_cuda_t to_fp16_cuda = ggml_get_to_fp16_cuda(src1->type); + GGML_ASSERT(to_fp16_cuda != nullptr); + size_t ne = src1_ncols*ne10; + src1_as_f16.alloc(ne); + to_fp16_cuda(src1_ddf_i, src1_as_f16.get(), ne, stream); + } + const half * src1_ptr = src1->type == GGML_TYPE_F16 ? (const half *) src1_ddf_i : src1_as_f16.get(); + + CUBLAS_CHECK(cublasSetStream(ctx.cublas_handle(id), stream)); + + const auto & force_compute_type = ggml_cuda_cublas_get_force_compute_type(); + + if (!force_compute_type.fp16 && (GGML_CUDA_CC_IS_CDNA(cc) + || GGML_CUDA_CC_IS_RDNA4(cc) + || cc == GGML_CUDA_CC_VOLTA + || force_compute_type.fp32)) + { + const float alpha = 1.0f; + const float beta = 0.0f; + CUBLAS_CHECK( + cublasGemmEx(ctx.cublas_handle(id), CUBLAS_OP_T, CUBLAS_OP_N, + row_diff, src1_ncols, ne10, + &alpha, src0_ptr, CUDA_R_16F, ne00, + src1_ptr, CUDA_R_16F, ne10, + &beta, dst_dd_i, CUDA_R_32F, ldc, + CUBLAS_COMPUTE_32F, + CUBLAS_GEMM_DEFAULT_TENSOR_OP)); + } else { + ggml_cuda_pool_alloc dst_f16(ctx.pool(id), row_diff*src1_ncols); + + const half alpha_f16 = 1.0f; + const half beta_f16 = 0.0f; + + CUBLAS_CHECK( + cublasGemmEx(ctx.cublas_handle(id), CUBLAS_OP_T, CUBLAS_OP_N, + row_diff, src1_ncols, ne10, + &alpha_f16, src0_ptr, CUDA_R_16F, ne00, + src1_ptr, CUDA_R_16F, ne10, + &beta_f16, dst_f16.get(), CUDA_R_16F, ldc, + CUBLAS_COMPUTE_16F, + CUBLAS_GEMM_DEFAULT_TENSOR_OP)); + + const to_fp32_cuda_t to_fp32_cuda = ggml_get_to_fp32_cuda(GGML_TYPE_F16); + to_fp32_cuda(dst_f16.get(), dst_dd_i, row_diff*src1_ncols, stream); + } + } else { + ggml_cuda_pool_alloc src0_ddq_as_f32(ctx.pool(id)); + ggml_cuda_pool_alloc src1_ddq_as_f32(ctx.pool(id)); + + if (src0->type != GGML_TYPE_F32) { + const to_fp32_cuda_t to_fp32_cuda = ggml_get_to_fp32_cuda(src0->type); + GGML_ASSERT(to_fp32_cuda != nullptr); + src0_ddq_as_f32.alloc(row_diff*ne00); + to_fp32_cuda(src0_dd_i, src0_ddq_as_f32.get(), row_diff*ne00, stream); + } + if (src1->type != GGML_TYPE_F32) { + const to_fp32_cuda_t to_fp32_cuda = ggml_get_to_fp32_cuda(src1->type); + GGML_ASSERT(to_fp32_cuda != nullptr); + src1_ddq_as_f32.alloc(src1_ncols*ne10); + to_fp32_cuda(src1_ddf_i, src1_ddq_as_f32.get(), src1_ncols*ne10, stream); + } + + const float * src0_ddf_i = src0->type == GGML_TYPE_F32 ? (const float *) src0_dd_i : src0_ddq_as_f32.get(); + const float * src1_ddf1_i = src1->type == GGML_TYPE_F32 ? (const float *) src1_ddf_i : src1_ddq_as_f32.get(); + + const float alpha = 1.0f; + const float beta = 0.0f; + + CUBLAS_CHECK(cublasSetStream(ctx.cublas_handle(id), stream)); + CUBLAS_CHECK( + cublasSgemm(ctx.cublas_handle(id), CUBLAS_OP_T, CUBLAS_OP_N, + row_diff, src1_ncols, ne10, + &alpha, src0_ddf_i, ne00, + src1_ddf1_i, ne10, + &beta, dst_dd_i, ldc)); + } + + GGML_UNUSED_VARS(dst, src1_ddq_i, src1_padded_row_size); +} + +static cudaError_t ggml_cuda_Memcpy2DPeerAsync( + void * dst, int dstDevice, size_t dpitch, void * src, int srcDevice, size_t spitch, size_t width, size_t height, cudaStream_t stream) { + +#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) + // cudaMemcpy2DAsync may fail with copies between vmm pools of different devices + cudaMemcpy3DPeerParms p = {}; + p.dstDevice = dstDevice; + p.dstPtr = make_cudaPitchedPtr(dst, dpitch, dpitch, height); + p.srcDevice = srcDevice; + p.srcPtr = make_cudaPitchedPtr(src, spitch, spitch, height); + p.extent = make_cudaExtent(width, height, 1); + return cudaMemcpy3DPeerAsync(&p, stream); +#else + // HIP does not support cudaMemcpy3DPeerAsync or vmm pools + GGML_UNUSED(dstDevice); + GGML_UNUSED(srcDevice); + return cudaMemcpy2DAsync(dst, dpitch, src, spitch, width, height, cudaMemcpyDeviceToDevice, stream); +#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) +} + +static void ggml_cuda_op_mul_mat( + ggml_backend_cuda_context & ctx, + const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, ggml_cuda_op_mul_mat_t op, + quantize_cuda_t quantize_src1) { + + const int64_t ne00 = src0->ne[0]; + const int64_t ne01 = src0->ne[1]; + const int64_t ne02 = src0->ne[2]; + const int64_t ne03 = src0->ne[3]; + + const int64_t ne10 = src1->ne[0]; + const int64_t ne11 = src1->ne[1]; + const int64_t ne12 = src1->ne[2]; + const int64_t ne13 = src1->ne[3]; + const int64_t nrows1 = ggml_nrows(src1); + + const int64_t ne0 = dst->ne[0]; + const int64_t ne1 = dst->ne[1]; + + // const int64_t nb10 = src1->nb[0]; + const int64_t nb11 = src1->nb[1]; + const int64_t nb12 = src1->nb[2]; + const int64_t nb13 = src1->nb[3]; + + const int64_t nb2 = dst->nb[2]; + const int64_t nb3 = dst->nb[3]; + + ggml_backend_cuda_buffer_context * src1_ctx = (ggml_backend_cuda_buffer_context *) src1->buffer->context; + ggml_backend_cuda_buffer_context * dst_ctx = (ggml_backend_cuda_buffer_context *) dst->buffer->context; + + GGML_ASSERT(src1->type == GGML_TYPE_F32 || (src1->ne[2] == 1 && src1->ne[3] == 1)); + + GGML_ASSERT(ne12 % ne02 == 0); + GGML_ASSERT(ne13 % ne03 == 0); + + const int64_t i02_divisor = ne12 / ne02; + const int64_t i03_divisor = ne13 / ne03; + + const size_t src0_ts = ggml_type_size(src0->type); + const size_t src0_bs = ggml_blck_size(src0->type); + const size_t q8_1_ts = sizeof(block_q8_1); + const size_t q8_1_bs = QK8_1; + + const bool src0_is_contiguous = ggml_is_contiguous(src0); + const bool src1_is_contiguous = ggml_is_contiguous(src1); + + const int64_t src1_padded_col_size = GGML_PAD(ne10, MATRIX_ROW_PADDING); + + const bool split = ggml_backend_buft_is_cuda_split(src0->buffer->buft); + GGML_ASSERT(!(split && ne02 > 1)); + GGML_ASSERT(!(split && ne03 > 1)); + GGML_ASSERT(!(split && ne02 < ne12)); + GGML_ASSERT(!(split && ne03 < ne13)); + + ggml_tensor_extra_gpu * src0_extra = split ? (ggml_tensor_extra_gpu *) src0->extra : nullptr; + + + std::array tensor_split; + if (split) { + ggml_backend_cuda_split_buffer_type_context * buft_ctx = (ggml_backend_cuda_split_buffer_type_context *) src0->buffer->buft->context; + tensor_split = buft_ctx->tensor_split; + } + + struct dev_data { + int cc; + + ggml_cuda_pool_alloc src0_dd_alloc; + ggml_cuda_pool_alloc src1_ddf_alloc; + ggml_cuda_pool_alloc src1_ddq_alloc; + ggml_cuda_pool_alloc dst_dd_alloc; + + char * src0_dd = nullptr; + float * src1_ddf = nullptr; // float + char * src1_ddq = nullptr; // q8_1 + float * dst_dd = nullptr; + + int64_t row_low; + int64_t row_high; + }; + + dev_data dev[GGML_CUDA_MAX_DEVICES]; + + int used_devices = 0; + + for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { + dev[id].cc = ggml_cuda_info().devices[id].cc; + + // by default, use all rows + dev[id].row_low = 0; + dev[id].row_high = ne01; + + // for multi GPU, get the row boundaries from tensor split + // and round to mul_mat_q tile sizes + if (split) { + const int64_t rounding = get_row_rounding(tensor_split); + + if (id != 0) { + dev[id].row_low = ne01*tensor_split[id]; + if (dev[id].row_low < ne01) { + dev[id].row_low -= dev[id].row_low % rounding; + } + } + + if (id != ggml_backend_cuda_get_device_count() - 1) { + dev[id].row_high = ne01*tensor_split[id + 1]; + if (dev[id].row_high < ne01) { + dev[id].row_high -= dev[id].row_high % rounding; + } + } + } + } + + for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { + if ((!split && id != ctx.device) || dev[id].row_low == dev[id].row_high) { + continue; + } + + used_devices++; + + const bool src1_on_device = id == src1_ctx->device; + const bool dst_on_device = id == dst_ctx->device; + + ggml_cuda_set_device(id); + cudaStream_t stream = ctx.stream(id, 0); + + if (src0_is_contiguous) { + dev[id].src0_dd = split ? (char *) src0_extra->data_device[id] : (char *) src0->data; + } else { + // If src0 is not contiguous it will be copied to a temporary buffer. + // This buffer needs to be cleared entirely because multiple regions will function as padding. + const size_t nbytes_data = ggml_nbytes(src0); + const size_t nbytes_padding = ggml_row_size(src0->type, MATRIX_ROW_PADDING - ne00 % MATRIX_ROW_PADDING); + dev[id].src0_dd = dev[id].src0_dd_alloc.alloc(ctx.pool(id), nbytes_data + nbytes_padding); + CUDA_CHECK(cudaMemsetAsync(dev[id].src0_dd, 0, nbytes_data + nbytes_padding, stream)); + } + + // If src0 is on a temporary compute buffer (partial offloading) there may be some padding that needs to be cleared: + if (ne00 % MATRIX_ROW_PADDING != 0 && ggml_is_quantized(src0->type) && ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE && src0->view_src == nullptr) { + GGML_ASSERT(ggml_is_contiguously_allocated(src0)); + GGML_ASSERT(!src0->view_src); + const size_t nbytes_data = ggml_row_size(src0->type, (dev[id].row_high - dev[id].row_low)*ne00); + const size_t nbytes_padding = ggml_row_size(src0->type, MATRIX_ROW_PADDING - ne00 % MATRIX_ROW_PADDING); + CUDA_CHECK(cudaMemsetAsync(dev[id].src0_dd + nbytes_data, 0, nbytes_padding, stream)); + } + + if (src1_on_device && src1_is_contiguous) { + dev[id].src1_ddf = (float *) src1->data; + } else { + dev[id].src1_ddf = dev[id].src1_ddf_alloc.alloc(ctx.pool(id), ggml_nelements(src1)); + } + + if (quantize_src1) { + size_t src_1_ddq_size = nrows1*src1_padded_col_size*q8_1_ts/q8_1_bs; + if (quantize_src1 == quantize_mmq_q8_1_cuda) { + src_1_ddq_size += get_mmq_x_max_host(dev[id].cc)*sizeof(block_q8_1_mmq); + } + dev[id].src1_ddq = dev[id].src1_ddq_alloc.alloc(ctx.pool(id), src_1_ddq_size); + + if (src1_on_device && src1_is_contiguous) { + quantize_src1( + dev[id].src1_ddf, nullptr, dev[id].src1_ddq, src0->type, ne10, + nb11/sizeof(float), nb12/sizeof(float), nb13/sizeof(float), + src1_padded_col_size, ne11, ne12, ne13, stream); + CUDA_CHECK(cudaGetLastError()); + } + } + + if (dst_on_device) { + dev[id].dst_dd = (float *) dst->data; + } else { + const size_t size_dst_ddf = split ? (dev[id].row_high - dev[id].row_low)*ne1 : ggml_nelements(dst); + dev[id].dst_dd = dev[id].dst_dd_alloc.alloc(ctx.pool(id), size_dst_ddf); + } + } + + // if multiple devices are used they need to wait for the main device + // here an event is recorded that signals that the main device has finished calculating the input data + if (split && used_devices > 1) { + ggml_cuda_set_device(ctx.device); + CUDA_CHECK(cudaEventRecord(src0_extra->events[ctx.device][0], ctx.stream())); + } + + const int64_t src1_col_stride = split && used_devices > 1 ? MUL_MAT_SRC1_COL_STRIDE : ne11; + for (int64_t src1_col_0 = 0; src1_col_0 < ne11; src1_col_0 += src1_col_stride) { + const int64_t is = split ? (src1_col_0/src1_col_stride) % GGML_CUDA_MAX_STREAMS : 0; + const int64_t src1_ncols = src1_col_0 + src1_col_stride > ne11 ? ne11 - src1_col_0 : src1_col_stride; + + for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { + if ((!split && id != ctx.device) || dev[id].row_low == dev[id].row_high) { + continue; + } + + const bool src1_on_device = id == src1_ctx->device; + const bool dst_on_device = id == dst_ctx->device; + const int64_t row_diff = dev[id].row_high - dev[id].row_low; + + ggml_cuda_set_device(id); + cudaStream_t stream = ctx.stream(id, is); + + // wait for main GPU data if necessary + if (split && (id != ctx.device || is != 0)) { + CUDA_CHECK(cudaStreamWaitEvent(stream, src0_extra->events[ctx.device][0], 0)); + } + + for (int64_t i0 = 0; i0 < ne13*ne12; ++i0) { + const int64_t i03 = i0 / ne12; + const int64_t i02 = i0 % ne12; + + size_t src1_ddq_i_offset = i0*ne11 * src1_padded_col_size*q8_1_ts/q8_1_bs; + if (quantize_src1 == quantize_mmq_q8_1_cuda) { + src1_ddq_i_offset += src1_col_0 * sizeof(block_q8_1_mmq); + } else { + src1_ddq_i_offset += src1_col_0 * src1_padded_col_size*q8_1_ts/q8_1_bs; + } + + // for split tensors the data begins at i0 == i0_offset_low + const size_t nbytes_src0_matrix = ne01*ne00*src0_ts / src0_bs; + char * src0_dd_i = dev[id].src0_dd + ((i03/i03_divisor)*ne02 + (i02/i02_divisor)) * nbytes_src0_matrix; + float * src1_ddf_i = dev[id].src1_ddf + (i0*ne11 + src1_col_0) * ne10; + char * src1_ddq_i = dev[id].src1_ddq + src1_ddq_i_offset; + float * dst_dd_i = dev[id].dst_dd + (i0*ne1 + src1_col_0) * (dst_on_device ? ne0 : row_diff); + + // the main device memory buffer can be on VRAM scratch, with space for all partial results + // in that case an offset on dst_ddf_i is needed + if (id == ctx.device) { + dst_dd_i += dev[id].row_low; // offset is 0 if no tensor split + } + + // copy src0, src1 to device if necessary + if (src1_is_contiguous) { + if (id != ctx.device) { + if (quantize_src1) { + char * src1_ddq_i_source = dev[ctx.device].src1_ddq + src1_ddq_i_offset; + if (quantize_src1 == quantize_mmq_q8_1_cuda) { + const size_t pitch = ne11*sizeof(block_q8_1_mmq); + const size_t width = src1_ncols*sizeof(block_q8_1_mmq); + const size_t height = src1_padded_col_size/(4*QK8_1); + CUDA_CHECK(ggml_cuda_Memcpy2DPeerAsync(src1_ddq_i, id, pitch, src1_ddq_i_source, ctx.device, pitch, width, height, stream)); + } else { + CUDA_CHECK(cudaMemcpyPeerAsync( + src1_ddq_i, id, src1_ddq_i_source, ctx.device, src1_ncols*src1_padded_col_size*q8_1_ts/q8_1_bs, stream)); + } + } else { + float * src1_ddf_i_source = (float *) src1->data; + src1_ddf_i_source += (i0*ne11 + src1_col_0) * ne10; + CUDA_CHECK(cudaMemcpyPeerAsync(src1_ddf_i, id, src1_ddf_i_source, ctx.device, + src1_ncols*ne10*sizeof(float), stream)); + } + } + } else if (src1_on_device && !src1_is_contiguous) { + CUDA_CHECK(ggml_cuda_cpy_tensor_2d( + src1_ddf_i, src1, i03, i02, src1_col_0, src1_col_0+src1_ncols, stream)); + } else { + GGML_ABORT("fatal error"); + } + + if (quantize_src1 && !src1_is_contiguous) { + quantize_src1( + src1_ddf_i, nullptr, src1_ddq_i, src0->type, ne10, ne10, ne11*ne10, ne12*ne11*ne10, + src1_padded_col_size, src1_ncols, 1, 1, stream); + CUDA_CHECK(cudaGetLastError()); + } + + if (src1_col_0 == 0 && !src0_is_contiguous && i03 % i03_divisor == 0 && i02 % i02_divisor == 0) { + CUDA_CHECK(ggml_cuda_cpy_tensor_2d( + src0_dd_i, src0, i03/i03_divisor, i02/i02_divisor, dev[id].row_low, dev[id].row_high, stream)); + } + + // do the computation + op(ctx, src0, src1, dst, src0_dd_i, src1_ddf_i, src1_ddq_i, dst_dd_i, + dev[id].row_low, dev[id].row_high, src1_ncols, src1_padded_col_size, stream); + CUDA_CHECK(cudaGetLastError()); + + // copy dst to host or other device if necessary + if (!dst_on_device) { + void * dst_off_device = dst->data; + if (split) { + // src0 = weight matrix is saved as a transposed matrix for better memory layout. + // dst is NOT transposed. + // The outputs of matrix matrix multiplications can therefore NOT simply be concatenated for >1 GPU. + // Instead they need to be copied to the correct slice in ne0 = dst row index. + // If dst is a vector with ne0 == 1 then you don't have to do this but it still produces correct results. + float * dhf_dst_i = (float *) ((char *) dst_off_device + i02*nb2 + i03*nb3); + GGML_ASSERT(dst->nb[1] == ne0*sizeof(float)); + dhf_dst_i += src1_col_0*ne0 + dev[id].row_low; + CUDA_CHECK(ggml_cuda_Memcpy2DPeerAsync( + dhf_dst_i, ctx.device, ne0*sizeof(float), dst_dd_i, id, row_diff*sizeof(float), row_diff*sizeof(float), src1_ncols, stream)); + } else { + float * dhf_dst_i = (float *) ((char *) dst_off_device + i02*nb2 + i03*nb3); + GGML_ASSERT(dst->nb[1] == ne0*sizeof(float)); + dhf_dst_i += src1_col_0*ne0; + CUDA_CHECK(cudaMemcpyAsync(dhf_dst_i, dst_dd_i, src1_ncols*ne0*sizeof(float), cudaMemcpyDeviceToDevice, stream)); + } + } + + // add event for the main device to wait on until other device is done + if (split && (id != ctx.device || is != 0)) { + CUDA_CHECK(cudaEventRecord(src0_extra->events[id][is], stream)); + } + } + } + } + + // main device waits for all other devices to be finished + if (split && ggml_backend_cuda_get_device_count() > 1) { + int64_t is_max = (ne11 + MUL_MAT_SRC1_COL_STRIDE - 1) / MUL_MAT_SRC1_COL_STRIDE; + is_max = is_max <= GGML_CUDA_MAX_STREAMS ? is_max : GGML_CUDA_MAX_STREAMS; + + ggml_cuda_set_device(ctx.device); + for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { + if (dev[id].row_low == dev[id].row_high) { + continue; + } + for (int64_t is = 0; is < is_max; ++is) { + CUDA_CHECK(cudaStreamWaitEvent(ctx.stream(), src0_extra->events[id][is], 0)); + } + } + } +} + +static __global__ void k_compute_batched_ptrs( + const void * src0_as_f16, const void * src1_as_f16, char * dst, + const void ** ptrs_src, void ** ptrs_dst, + int64_t ne12, int64_t ne13, + int64_t ne23, + size_t nb02, size_t nb03, + size_t nb12, size_t nb13, + size_t nbd2, size_t nbd3, + int64_t r2, int64_t r3) { + const int64_t i13 = blockIdx.x * blockDim.x + threadIdx.x; + const int64_t i12 = blockIdx.y * blockDim.y + threadIdx.y; + + if (i13 >= ne13 || i12 >= ne12) { + return; + } + + const int64_t i03 = i13 / r3; + const int64_t i02 = i12 / r2; + + ptrs_src[0*ne23 + i12 + i13*ne12] = (const char *) src0_as_f16 + i02*nb02 + i03*nb03; + ptrs_src[1*ne23 + i12 + i13*ne12] = (const char *) src1_as_f16 + i12*nb12 + i13*nb13; + ptrs_dst[0*ne23 + i12 + i13*ne12] = ( char *) dst + i12*nbd2 + i13*nbd3; +} + +// Type traits for mapping ggml types to CUDA/cuBLAS types +template +struct batched_mul_mat_traits; + +template<> +struct batched_mul_mat_traits { + using cuda_type = float; + static inline const cublasComputeType_t compute_type = CUBLAS_COMPUTE_32F; + static inline const cudaDataType_t data_type = CUDA_R_32F; + static inline const ggml_type ggml_type_val = GGML_TYPE_F32; + static inline const float alpha = 1.0f; + static inline const float beta = 0.0f; + static inline const void* get_alpha() { static const float val = alpha; return &val; } + static inline const void* get_beta() { static const float val = beta; return &val; } + static inline auto get_nc_converter(ggml_type src_type) { return ggml_get_to_fp32_nc_cuda(src_type); } +}; + +template<> +struct batched_mul_mat_traits { + using cuda_type = nv_bfloat16; + static inline const cublasComputeType_t compute_type = CUBLAS_COMPUTE_32F; + static inline const cudaDataType_t data_type = CUDA_R_16BF; + static inline const ggml_type ggml_type_val = GGML_TYPE_BF16; + static inline const float alpha = 1.0f; + static inline const float beta = 0.0f; + static inline const void* get_alpha() { static const float val = alpha; return &val; } + static inline const void* get_beta() { static const float val = beta; return &val; } + static inline auto get_nc_converter(ggml_type src_type) { return ggml_get_to_bf16_nc_cuda(src_type); } +}; + +template<> +struct batched_mul_mat_traits { + using cuda_type = half; + static inline const cublasComputeType_t compute_type = CUBLAS_COMPUTE_16F; + static inline const cudaDataType_t data_type = CUDA_R_16F; + static inline const ggml_type ggml_type_val = GGML_TYPE_F16; + static inline const half alpha = 1.0; + static inline const half beta = 0.0; + static inline const void* get_alpha() { static const half val = alpha; return &val; } + static inline const void* get_beta() { static const half val = beta; return &val; } + static inline auto get_nc_converter(ggml_type src_type) { return ggml_get_to_fp16_nc_cuda(src_type); } +}; + +template +static void ggml_cuda_mul_mat_batched_cublas_impl(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { + using traits = batched_mul_mat_traits; + using cuda_t = typename traits::cuda_type; + + GGML_ASSERT(!ggml_is_transposed(src0)); + GGML_ASSERT(!ggml_is_transposed(src1)); + GGML_ASSERT(!ggml_backend_buft_is_cuda_split(src0->buffer->buft)); + GGML_ASSERT(src0->type == src0_type); + GGML_ASSERT(ggml_is_contiguous(dst)); + + // Byte offsets and tensor dimensions are currently used in an inconsistent way for dst. + // As long as dst is contiguous this does not matter though. + + GGML_TENSOR_BINARY_OP_LOCALS + + const int64_t ne_dst = ggml_nelements(dst); + cudaStream_t main_stream = ctx.stream(); + CUBLAS_CHECK(cublasSetStream(ctx.cublas_handle(), main_stream)); + + float * dst_ddf = (float *) dst->data; + const size_t ts_src1 = ggml_type_size(src1->type); + GGML_ASSERT(nb10 == ts_src1); + int64_t s11 = nb11 / ts_src1; + int64_t s12 = nb12 / ts_src1; + int64_t s13 = nb13 / ts_src1; + + const cuda_t * src0_ptr = nullptr; + const cuda_t * src1_ptr = nullptr; + + ggml_cuda_pool_alloc src0_alloc(ctx.pool()); + ggml_cuda_pool_alloc src1_alloc(ctx.pool()); + + bool is_src0_cont_2 = ggml_is_contiguous_2(src0); + bool is_src1_cont_2 = ggml_is_contiguous_2(src1); + + // Handle src0 + src0_ptr = (const cuda_t *) src0->data; + + // Handle src1 - convert if necessary + if (src1->type == src0_type) { + src1_ptr = (const cuda_t *) src1->data; + } else { + // Convert src1 to target type using traits conversion functions + const int64_t ne_src1 = ggml_nelements(src1); + src1_alloc.alloc(ne_src1); + + const auto convert_func = traits::get_nc_converter(src1->type); + GGML_ASSERT(convert_func != nullptr); + convert_func(src1->data, src1_alloc.get(), ne10, ne11, ne12, ne13, s11, s12, s13, main_stream); + src1_ptr = src1_alloc.get(); + s11 = ne10; + s12 = ne11*s11; + s13 = ne12*s12; + + is_src1_cont_2 = true; + } + + // Setup destination buffer + ggml_cuda_pool_alloc dst_temp(ctx.pool()); + char * dst_t; + size_t nbd2 = dst->nb[2]; + size_t nbd3 = dst->nb[3]; + + cublasComputeType_t cu_compute_type = traits::compute_type; + cudaDataType_t cu_data_type = traits::data_type; + cudaDataType_t cu_data_type_a = traits::data_type; + cudaDataType_t cu_data_type_b = traits::data_type; + const void * alpha = traits::get_alpha(); + const void * beta = traits::get_beta(); + + const auto & force_compute_type = ggml_cuda_cublas_get_force_compute_type(); + + int id = ggml_cuda_get_device(); + const int cc = ggml_cuda_info().devices[id].cc; + static constexpr bool is_src0_type_f16 = src0_type == GGML_TYPE_F16; + + // bf16 and fp32 are already being computed in fp32 (ensure it using static_assert), + // so checking necessity of forced fp32 only for fp16 src0_type + static_assert(is_src0_type_f16 || traits::compute_type == CUBLAS_COMPUTE_32F); + + const bool need_compute_32f = is_src0_type_f16 && !force_compute_type.fp16 && (GGML_CUDA_CC_IS_CDNA(cc) + || GGML_CUDA_CC_IS_RDNA4(cc) + || cc == GGML_CUDA_CC_VOLTA + || force_compute_type.fp32); + + if (dst->op_params[0] == GGML_PREC_DEFAULT && !need_compute_32f) { + if constexpr (src0_type == GGML_TYPE_F32) { + dst_t = (char *) dst_ddf; // Direct F32 output + } else { + dst_t = (char *) dst_temp.alloc(ne_dst); + nbd2 /= sizeof(float) / sizeof(cuda_t); + nbd3 /= sizeof(float) / sizeof(cuda_t); + } + } else { + dst_t = (char *) dst_ddf; + cu_compute_type = batched_mul_mat_traits::compute_type; + cu_data_type = batched_mul_mat_traits::data_type; + alpha = batched_mul_mat_traits::get_alpha(); + beta = batched_mul_mat_traits::get_beta(); + } + + GGML_ASSERT(ne12 % ne02 == 0); + GGML_ASSERT(ne13 % ne03 == 0); + + // broadcast factors + const int64_t r2 = ne12/ne02; + const int64_t r3 = ne13/ne03; + + if (r2 == 1 && r3 == 1 && is_src0_cont_2 && is_src1_cont_2) { + // with a [0, 2, 1, 3] perm. and ne02==1 the matrix strides need to be determined from dim 3: + const int64_t sma = ne02 == 1 ? nb03/nb00 : nb02/nb00; + const int64_t smb = ne12 == 1 ? s13 : s12; + + // there is no broadcast and src0, src1 are contiguous across dims 2, 3 + // use cublasGemmStridedBatchedEx + CUBLAS_CHECK( + cublasGemmStridedBatchedEx(ctx.cublas_handle(), CUBLAS_OP_T, CUBLAS_OP_N, + ne01, ne11, ne10, + alpha, src0_ptr, cu_data_type_a, nb01/nb00, sma, // strideA + src1_ptr, cu_data_type_b, s11, smb, // strideB + beta, dst_t, cu_data_type, ne0, ne1*ne0, // strideC + ne12*ne13, + cu_compute_type, + CUBLAS_GEMM_DEFAULT_TENSOR_OP)); + } else { + // use cublasGemmBatchedEx + const int64_t ne23 = ne12*ne13; + + ggml_cuda_pool_alloc ptrs_src(ctx.pool(), 2*ne23); + ggml_cuda_pool_alloc< void *> ptrs_dst(ctx.pool(), 1*ne23); + + size_t src1_stride_size = sizeof(cuda_t); + + const int threads_x = 16; + const int threads_y = 16; + dim3 block_dims(threads_x, threads_y); + + dim3 grid_dims( + (ne13 + threads_x - 1) / threads_x, + (ne12 + threads_y - 1) / threads_y + ); + k_compute_batched_ptrs<<>>( + src0_ptr, src1_ptr, dst_t, + ptrs_src.get(), ptrs_dst.get(), + ne12, ne13, + ne23, + nb02, nb03, + (src1->type == src0_type) ? nb12 : s12*src1_stride_size, + (src1->type == src0_type) ? nb13 : s13*src1_stride_size, + nbd2, nbd3, + r2, r3); + + CUDA_CHECK(cudaGetLastError()); + + CUBLAS_CHECK( + cublasGemmBatchedEx(ctx.cublas_handle(), CUBLAS_OP_T, CUBLAS_OP_N, + ne01, ne11, ne10, + alpha, (const void **) (ptrs_src.get() + 0*ne23), cu_data_type_a, nb01/nb00, + (const void **) (ptrs_src.get() + 1*ne23), cu_data_type_b, s11, + beta, ( void **) (ptrs_dst.get() + 0*ne23), cu_data_type, ne0, + ne23, + cu_compute_type, + CUBLAS_GEMM_DEFAULT_TENSOR_OP)); + } + + // Convert output back to F32 if needed + if (dst->op_params[0] == GGML_PREC_DEFAULT && cu_data_type != CUDA_R_32F) { + const to_fp32_cuda_t to_fp32_cuda = ggml_get_to_fp32_cuda(traits::ggml_type_val); + to_fp32_cuda(dst_temp.get(), dst_ddf, ne_dst, main_stream); + } +} + +static void ggml_cuda_mul_mat_batched_cublas(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { + GGML_ASSERT(src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_BF16 || src0->type == GGML_TYPE_F32); + + switch (src0->type) { + case GGML_TYPE_F32: + ggml_cuda_mul_mat_batched_cublas_impl(ctx, src0, src1, dst); + break; + case GGML_TYPE_BF16: + ggml_cuda_mul_mat_batched_cublas_impl(ctx, src0, src1, dst); + break; + case GGML_TYPE_F16: + ggml_cuda_mul_mat_batched_cublas_impl(ctx, src0, src1, dst); + break; + default: + GGML_ABORT("Unsupported type"); + } +} + +static bool ggml_cuda_should_fuse_mul_mat(const ggml_tensor * ffn_up, + const ggml_tensor * ffn_gate, + const ggml_tensor * glu, + const ggml_tensor * ffn_up_bias = nullptr, + const ggml_tensor * ffn_gate_bias = nullptr) { + const bool has_bias = ffn_up_bias != nullptr || ffn_gate_bias != nullptr; + + if (has_bias && (!ffn_up_bias || !ffn_gate_bias)) { + return false; + } + + const bool is_mul_mat = ffn_up->op == GGML_OP_MUL_MAT && ffn_gate->op == GGML_OP_MUL_MAT && glu->op == GGML_OP_GLU; + const bool is_mul_mat_id = ffn_up->op == GGML_OP_MUL_MAT_ID && ffn_gate->op == GGML_OP_MUL_MAT_ID && glu->op == GGML_OP_GLU; + + GGML_ASSERT(ffn_up && ffn_gate && glu); + + if (!is_mul_mat && !is_mul_mat_id) { + return false; + } + + const ggml_op expected_bias_op = is_mul_mat ? GGML_OP_ADD : GGML_OP_ADD_ID; + + if (has_bias) { + if (ffn_up_bias->op != expected_bias_op || ffn_gate_bias->op != expected_bias_op) { + return false; + } + + if (glu->src[0] != ffn_gate_bias || glu->src[1] != ffn_up_bias) { + return false; + } + + if (expected_bias_op == GGML_OP_ADD) { + const bool up_has_mul = ffn_up_bias->src[0] == ffn_up || ffn_up_bias->src[1] == ffn_up; + const bool gate_has_mul = ffn_gate_bias->src[0] == ffn_gate || ffn_gate_bias->src[1] == ffn_gate; + if (!up_has_mul || !gate_has_mul) { + return false; + } + } else { // GGML_OP_ADD_ID + if (ffn_up_bias->src[0] != ffn_up || ffn_gate_bias->src[0] != ffn_gate) { + return false; + } + if (ffn_up_bias->src[2] != ffn_up->src[2] || ffn_gate_bias->src[2] != ffn_gate->src[2]) { + return false; + } + } + } else { + if (glu->src[0] != ffn_gate && glu->src[1] != ffn_up) { + return false; + } + } + + if (ffn_up->src[0]->type != ffn_gate->src[0]->type || !ggml_are_same_shape(ffn_up->src[0], ffn_gate->src[0]) || + !ggml_are_same_stride(ffn_up->src[0], ffn_gate->src[0])) { + return false; + } + + if (ffn_up->src[1] != ffn_gate->src[1]) { + return false; + } + + if (ffn_up->src[2] && (ffn_up->src[2] != ffn_gate->src[2])) { + return false; + } + + static constexpr std::array valid_glu_ops = { + GGML_GLU_OP_SWIGLU, + GGML_GLU_OP_GEGLU, + GGML_GLU_OP_SWIGLU_OAI, + GGML_GLU_OP_SWIGLU_DS4, + }; + + if (std::find(valid_glu_ops.begin(), valid_glu_ops.end(), ggml_get_glu_op(glu)) == valid_glu_ops.end()) { + return false; + } + + if (const bool swapped = ggml_get_op_params_i32(glu, 1); swapped) { + return false; + } + + const bool split = ggml_backend_buft_is_cuda_split(ffn_up->src[0]->buffer->buft) || + ggml_backend_buft_is_cuda_split(ffn_gate->src[0]->buffer->buft); + + //TODO: add support for fusion for split buffers + if (split) { + return false; + } + + return true; +} + +static bool ggml_cuda_should_fuse_mul_mat_vec_f(const ggml_tensor * tensor) { + ggml_tensor * src0 = tensor->src[0]; + ggml_tensor * src1 = tensor->src[1]; + const ggml_tensor * dst = tensor; + + const bool is_mul_mat_id = tensor->op == GGML_OP_MUL_MAT_ID; + + bool use_mul_mat_vec_f = + (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_BF16) && + src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32; + + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + use_mul_mat_vec_f = use_mul_mat_vec_f && ggml_cuda_should_use_mmvf(src0->type, cc, src0->ne, src0->nb, is_mul_mat_id ? src1->ne[2] : src1->ne[1]); + + const bool split = ggml_backend_buft_is_cuda_split(src0->buffer->buft) || + ggml_backend_buft_is_cuda_split(src1->buffer->buft); + + //TODO: add support for fusion for split buffers + if (split) { + return false; + } + + //we only support fusion for ncols_dst = 1 + if (tensor->op == GGML_OP_MUL_MAT && dst->ne[1] != 1) { + return false; + } + + if (tensor->op == GGML_OP_MUL_MAT_ID && dst->ne[2] != 1) { + return false; + } + + + return use_mul_mat_vec_f; +} + +static inline void ggml_cuda_set_fusion_glu_params(ggml_cuda_mm_fusion_args_host & fusion_data, const ggml_tensor * glu) { + fusion_data.glu_op = ggml_get_glu_op(glu); + fusion_data.glu_param0 = ggml_get_op_params_f32(glu, 2); + fusion_data.glu_param1 = ggml_get_op_params_f32(glu, 3); +} + +static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) { + ggml_tensor * src0 = tensor->src[0]; + ggml_tensor * src1 = tensor->src[1]; + const ggml_tensor * dst = tensor; + const bool is_mul_mat_id = tensor->op == GGML_OP_MUL_MAT_ID; + + const bool bad_padding_clear = ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE && + ggml_nbytes(src0) != ggml_backend_buffer_get_alloc_size(src0->buffer, src0) && + src0->view_src; + + const int64_t ncols_dst = is_mul_mat_id ? dst->ne[2] : src1->ne[1]; + bool use_mul_mat_vec_q = ggml_is_quantized(src0->type) && !bad_padding_clear && src1->type == GGML_TYPE_F32 && + dst->type == GGML_TYPE_F32 && ncols_dst <= MMVQ_MAX_BATCH_SIZE; + + // fusion is not universally faster on Pascal + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + if (cc <= GGML_CUDA_CC_PASCAL) { + return false; + } + if (is_mul_mat_id) { + const int mmvq_mmid_max = get_mmvq_mmid_max_batch(src0->type, cc); + if (dst->ne[2] > mmvq_mmid_max) { + return false; + } + } + if (tensor->op == GGML_OP_MUL_MAT && dst->ne[1] != 1) { + return false; + } + const bool split = ggml_backend_buft_is_cuda_split(src0->buffer->buft) || + ggml_backend_buft_is_cuda_split(src1->buffer->buft); + + //TODO: add support for fusion for split buffers + if (split) { + return false; + } + + return use_mul_mat_vec_q; +} + +static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { + const bool split = ggml_backend_buft_is_cuda_split(src0->buffer->buft); + + // If src0 is a temporary compute buffer it may have some padding that needs to be cleared for mul_mat_vec_q or mul_mat_q. + // But if src0 is also a view of another tensor then this cannot be done safely because it may overwrite valid tensor data. + // Therefore, in such cases use cuBLAS. + const bool bad_padding_clear = ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE + && ggml_nbytes(src0) != ggml_backend_buffer_get_alloc_size(src0->buffer, src0) && src0->view_src; + + bool use_mul_mat_vec_f = (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_BF16) + && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32; + bool use_mul_mat_f = !ggml_is_quantized(src0->type) + && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32; + // LUCE_MMVQ_MAX_NCOLS: lower the MMVQ ncols ceiling for plain mul_mat so + // that MMQ takes small multi-token batches (spec-decode verify widths). + static const int luce_mmvq_max_ncols = []() { + const char * e = getenv("LUCE_MMVQ_MAX_NCOLS"); + const int v = e ? atoi(e) : MMVQ_MAX_BATCH_SIZE; + return v > 0 ? v : MMVQ_MAX_BATCH_SIZE; + }(); + bool use_mul_mat_vec_q = ggml_is_quantized(src0->type) && !bad_padding_clear + && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32 + && src1->ne[1] <= luce_mmvq_max_ncols; + bool use_mul_mat_q = ggml_is_quantized(src0->type) && !bad_padding_clear + && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32; + + bool any_gpus_with_slow_fp16 = false; + + if (split) { + ggml_backend_cuda_split_buffer_type_context * buft_ctx = (ggml_backend_cuda_split_buffer_type_context *) src0->buffer->buft->context; + auto & tensor_split = buft_ctx->tensor_split; + for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { + // skip devices that are not going to do any work: + if (tensor_split[id] >= (id + 1 < ggml_backend_cuda_get_device_count() ? tensor_split[id + 1] : 1.0f)) { + continue; + } + + const int cc = ggml_cuda_info().devices[id].cc; + const int warp_size = ggml_cuda_info().devices[id].warp_size; + use_mul_mat_q = use_mul_mat_q && ggml_cuda_should_use_mmq(src0->type, cc, src1->ne[1], /*n_experts=*/0); + use_mul_mat_f = use_mul_mat_f && ggml_cuda_should_use_mmf(src0->type, cc, warp_size, src0->ne, src0->nb, src1->ne[1], /*mul_mat_id=*/false); + use_mul_mat_vec_f = use_mul_mat_vec_f && ggml_cuda_should_use_mmvf(src0->type, cc, src0->ne, src0->nb, src1->ne[1]); + any_gpus_with_slow_fp16 = any_gpus_with_slow_fp16 || !fast_fp16_hardware_available(cc); + } + } else { + const int cc = ggml_cuda_info().devices[ctx.device].cc; + const int warp_size = ggml_cuda_info().devices[ctx.device].warp_size; + use_mul_mat_q = use_mul_mat_q && ggml_cuda_should_use_mmq(src0->type, cc, src1->ne[1], /*n_experts=*/0); + use_mul_mat_f = use_mul_mat_f && ggml_cuda_should_use_mmf(src0->type, cc, warp_size, src0->ne, src0->nb, src1->ne[1], /*mul_mat_id=*/false); + use_mul_mat_vec_f = use_mul_mat_vec_f && ggml_cuda_should_use_mmvf(src0->type, cc, src0->ne, src0->nb, src1->ne[1]); + any_gpus_with_slow_fp16 = any_gpus_with_slow_fp16 || !fast_fp16_hardware_available(cc); + } + + // debug helpers + //printf("src0: %8d %8d %8d %8d\n", src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3]); + //printf(" %8d %8d %8d %8d\n", src0->nb[0], src0->nb[1], src0->nb[2], src0->nb[3]); + //printf("src1: %8d %8d %8d %8d\n", src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3]); + //printf(" %8d %8d %8d %8d\n", src1->nb[0], src1->nb[1], src1->nb[2], src1->nb[3]); + //printf("src0 is contiguous %d, transposed %d, type = %s, name = %s\n", ggml_is_contiguous(src0), ggml_is_transposed(src0), ggml_type_name(src0->type), src0->name); + //printf("src1 is contiguous %d, transposed %d, type = %s, name = %s\n", ggml_is_contiguous(src1), ggml_is_transposed(src1), ggml_type_name(src1->type), src1->name); + + //TODO update for generic tensor parallelism + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + bool use_batched_cublas_f16 = src0->type == GGML_TYPE_F16 && (src1->type == GGML_TYPE_F16 || !any_gpus_with_slow_fp16); + bool use_batched_cublas_bf16 = src0->type == GGML_TYPE_BF16 && bf16_mma_hardware_available(cc); + bool use_batched_cublas_f32 = src0->type == GGML_TYPE_F32; + + if (!split && use_mul_mat_vec_f) { + // the custom F16 vector kernel can be used over batched cuBLAS GEMM + // but this is only faster for GPUs without tensor cores or with a thin src0 matrix (particularly KQV in attention) + ggml_cuda_mul_mat_vec_f(ctx, src0, src1, nullptr, dst); + } else if (!split && use_mul_mat_f) { + ggml_cuda_mul_mat_f(ctx, src0, src1, nullptr, dst); + } else if (!split && use_mul_mat_vec_q) { + ggml_cuda_mul_mat_vec_q(ctx, src0, src1, nullptr, dst); + } else if (!split && use_mul_mat_q) { + ggml_cuda_mul_mat_q(ctx, src0, src1, nullptr, dst); + } else if (!split && (use_batched_cublas_f16 || use_batched_cublas_bf16 || use_batched_cublas_f32) + && !ggml_is_transposed(src0) && !ggml_is_transposed(src1) && src1->ne[2]*src1->ne[3] > 1) { + // general KQ + KQV multi-batch without FlashAttention + ggml_cuda_mul_mat_batched_cublas(ctx, src0, src1, dst); + } else if (use_mul_mat_vec_f) { + ggml_cuda_op_mul_mat(ctx, src0, src1, dst, ggml_cuda_op_mul_mat_vec_f, nullptr); + } else if (use_mul_mat_vec_q) { + ggml_cuda_op_mul_mat(ctx, src0, src1, dst, ggml_cuda_op_mul_mat_vec_q, quantize_row_q8_1_cuda); + } else if (use_mul_mat_q) { + ggml_cuda_op_mul_mat(ctx, src0, src1, dst, ggml_cuda_op_mul_mat_q, quantize_mmq_q8_1_cuda); + } else { + ggml_cuda_op_mul_mat(ctx, src0, src1, dst, ggml_cuda_op_mul_mat_cublas, nullptr); + } +} + +static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + const ggml_tensor * ids = dst->src[2]; + + GGML_ASSERT(src1->type == GGML_TYPE_F32); + GGML_ASSERT(dst->type == GGML_TYPE_F32); + GGML_ASSERT(!ggml_backend_buft_is_cuda_split(src0->buffer->buft) && "mul_mat_id does not support split buffers"); + + GGML_TENSOR_BINARY_OP_LOCALS + + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + + // [TAG_MUL_MAT_ID_CUDA_GRAPHS] + if (src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { + static_assert(MMVQ_MAX_BATCH_SIZE == MMVF_MAX_BATCH_SIZE); + if (ne2 <= MMVQ_MAX_BATCH_SIZE) { + if (ggml_is_quantized(src0->type)) { + const int mmvq_mmid_max = get_mmvq_mmid_max_batch(src0->type, cc); + if (ne2 <= mmvq_mmid_max) { + ggml_cuda_mul_mat_vec_q(ctx, src0, src1, ids, dst); + return; + } + } else { + if (GGML_CUDA_CC_IS_AMD(cc)) { + ggml_cuda_mul_mat_vec_f(ctx, src0, src1, ids, dst); + return; + } + } + } + + if (ggml_cuda_should_use_mmq(src0->type, cc, ne12, /*n_experts=*/ne02)) { + ggml_cuda_mul_mat_q(ctx, src0, src1, ids, dst); + return; + } + + if (ggml_cuda_should_use_mmf(src0->type, cc, WARP_SIZE, src0->ne, src0->nb, src1->ne[2], /*mul_mat_id=*/true)) { + ggml_cuda_mul_mat_f(ctx, src0, src1, ids, dst); + return; + } + } + + // note: this path should not be reached when recording CUDA graphs, because it requires stream synchronization + // TODO: add asserts to verify this. should work with CUDA, HIP, etc. + cudaStream_t stream = ctx.stream(); + + GGML_ASSERT(nb12 % nb11 == 0); + GGML_ASSERT(nb2 % nb1 == 0); + + const ggml_type type_src1_sorted = (src0->type == GGML_TYPE_F16 && !fast_fp16_hardware_available(cc)) + || ggml_is_quantized(src0->type) ? GGML_TYPE_F32 : src0->type; + const ggml_type type_dst_sorted = GGML_TYPE_F32; + const size_t ts_src1_sorted = ggml_type_size(type_src1_sorted); + const size_t ts_dst_sorted = ggml_type_size(type_dst_sorted); + + const int64_t n_expert_used = ids->ne[0]; + const int64_t ne_get_rows = ne12 * n_expert_used; + + std::vector ids_to_sorted_host; + ids_to_sorted_host.reserve(2*ne_get_rows); + std::vector ids_from_sorted_host(ne_get_rows); + + ggml_cuda_pool_alloc ids_buf_dev(ctx.pool(), 2*ne_get_rows); + + std::vector tokens_per_expert(ne02); + + ggml_cuda_pool_alloc src1_sorted(ctx.pool(), ne12*n_expert_used*ne10*ts_src1_sorted); + ggml_cuda_pool_alloc dst_sorted(ctx.pool(), ne2 *n_expert_used* ne0*ts_dst_sorted); + + std::vector ids_host(ggml_nbytes(ids)); + CUDA_CHECK(cudaMemcpyAsync(ids_host.data(), ids->data, ggml_nbytes(ids), cudaMemcpyDeviceToHost, stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + + for (int64_t i02 = 0; i02 < ne02; ++i02) { // expert matrices + for (int64_t i12 = 0; i12 < ne12; ++i12) { // tokens + for (int64_t iex = 0; iex < n_expert_used; ++iex) { + const int32_t expert_to_use = *(const int32_t *)(ids_host.data() + i12*ids->nb[1] + iex*ids->nb[0]); + assert(expert_to_use >= 0 && expert_to_use < ne02); + if (expert_to_use == i02) { + ids_from_sorted_host[i12*n_expert_used + iex] = ids_to_sorted_host.size(); + ids_to_sorted_host.push_back(i12*ne11 + iex % ne11); + tokens_per_expert[i02]++; + break; + } + } + } + } + GGML_ASSERT(ids_to_sorted_host.size() == size_t(ne_get_rows)); + + ids_to_sorted_host.insert(ids_to_sorted_host.end(), ids_from_sorted_host.begin(), ids_from_sorted_host.end()); + + CUDA_CHECK(cudaMemcpyAsync(ids_buf_dev.ptr, ids_to_sorted_host.data(), 2*ne_get_rows*sizeof(int32_t), cudaMemcpyHostToDevice, stream)); + CUDA_CHECK(cudaStreamSynchronize(stream)); + + const int32_t * ids_to_sorted = ids_buf_dev.ptr + 0*ne_get_rows; + const int32_t * ids_from_sorted = ids_buf_dev.ptr + 1*ne_get_rows; + + get_rows_cuda(src1->data, src1->type, ids_to_sorted, src1_sorted.ptr, type_src1_sorted, + ne10, nb11, nb12, nb13, + ne_get_rows, 1, 1, sizeof(int32_t), ne_get_rows*sizeof(int32_t), ne_get_rows*sizeof(int32_t), + ne10*ts_src1_sorted, ne_get_rows*ne10*ts_src1_sorted, ne_get_rows*ne10*ts_src1_sorted, stream); + CUDA_CHECK(cudaGetLastError()); + + char * src1_data_cur = (char *) src1_sorted.ptr; + char * dst_data_cur = (char *) dst_sorted.ptr; + for (int64_t i02 = 0; i02 < ne02; ++i02) { + if (tokens_per_expert[i02] == 0) { + continue; + } + + ggml_tensor src0_slice = *src0; + src0_slice.ne[2] = 1; + src0_slice.nb[3] = src0_slice.nb[2]; + src0_slice.op = GGML_OP_VIEW; + src0_slice.view_src = dst->src[0]; // non-const pointer to src0 + src0_slice.data = (char *) src0->data + i02*nb02; + + ggml_tensor src1_slice; + memset(&src1_slice, 0, sizeof(src1_slice)); + src1_slice.buffer = src1->buffer; + src1_slice.type = type_src1_sorted; + src1_slice.ne[0] = ne10; + src1_slice.ne[1] = tokens_per_expert[i02]; + src1_slice.ne[2] = 1; + src1_slice.ne[3] = 1; + src1_slice.nb[0] = ts_src1_sorted; + src1_slice.nb[1] = src1_slice.ne[0] * src1_slice.nb[0]; + src1_slice.nb[2] = src1_slice.ne[1] * src1_slice.nb[1]; + src1_slice.nb[3] = src1_slice.ne[2] * src1_slice.nb[2]; + src1_slice.data = src1_data_cur; + + ggml_tensor dst_slice; + memset(&dst_slice, 0, sizeof(dst_slice)); + dst_slice.buffer = dst->buffer; + dst_slice.type = type_dst_sorted; + dst_slice.ne[0] = ne0; + dst_slice.ne[1] = tokens_per_expert[i02]; + dst_slice.ne[2] = 1; + dst_slice.ne[3] = 1; + dst_slice.nb[0] = ts_dst_sorted; + dst_slice.nb[1] = dst_slice.ne[0] * dst_slice.nb[0]; + dst_slice.nb[2] = dst_slice.ne[1] * dst_slice.nb[1]; + dst_slice.nb[3] = dst_slice.ne[2] * dst_slice.nb[2]; + dst_slice.data = dst_data_cur; + + ggml_cuda_mul_mat(ctx, &src0_slice, &src1_slice, &dst_slice); + CUDA_CHECK(cudaGetLastError()); + + src1_data_cur += src1_slice.nb[2]; + dst_data_cur += dst_slice.nb[2]; + } + + get_rows_cuda(dst_sorted.ptr, type_dst_sorted, ids_from_sorted, dst->data, dst->type, + ne0, ne0*ts_dst_sorted, ne_get_rows*ne0*ts_dst_sorted, ne_get_rows*ne0*ts_dst_sorted, + ne_get_rows, 1, 1, sizeof(int32_t), ne_get_rows*sizeof(int32_t), ne_get_rows*sizeof(int32_t), + nb1, nb2, nb3, stream); +} + +static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct ggml_tensor * dst) { + switch (dst->op) { + case GGML_OP_ARGMAX: + ggml_cuda_argmax(ctx, dst); + break; + case GGML_OP_COUNT_EQUAL: + ggml_cuda_count_equal(ctx, dst); + break; + case GGML_OP_REPEAT: + ggml_cuda_op_repeat(ctx, dst); + break; + case GGML_OP_REPEAT_BACK: + ggml_cuda_op_repeat_back(ctx, dst); + break; + case GGML_OP_GET_ROWS: + ggml_cuda_op_get_rows(ctx, dst); + break; + case GGML_OP_GET_ROWS_BACK: + ggml_cuda_op_get_rows_back(ctx, dst); + break; + case GGML_OP_SET_ROWS: + ggml_cuda_op_set_rows(ctx, dst); + break; + case GGML_OP_SET: + ggml_cuda_op_set(ctx, dst); + break; + case GGML_OP_DUP: + ggml_cuda_dup(ctx, dst); + break; + case GGML_OP_CPY: + ggml_cuda_cpy(ctx, dst->src[0], dst->src[1]); + break; + case GGML_OP_CONT: + ggml_cuda_dup(ctx, dst); + break; + case GGML_OP_ADD: + case GGML_OP_ADD1: // TODO: more efficient implementation + ggml_cuda_op_add(ctx, dst); + break; + case GGML_OP_ADD_ID: + ggml_cuda_op_add_id(ctx, dst); + break; + case GGML_OP_SUB: + ggml_cuda_op_sub(ctx, dst); + break; + case GGML_OP_ACC: + ggml_cuda_op_acc(ctx, dst); + break; + case GGML_OP_MUL: + ggml_cuda_op_mul(ctx, dst); + break; + case GGML_OP_DIV: + ggml_cuda_op_div(ctx, dst); + break; + case GGML_OP_UNARY: + switch (ggml_get_unary_op(dst)) { + case GGML_UNARY_OP_ABS: + ggml_cuda_op_abs(ctx, dst); + break; + case GGML_UNARY_OP_SGN: + ggml_cuda_op_sgn(ctx, dst); + break; + case GGML_UNARY_OP_NEG: + ggml_cuda_op_neg(ctx, dst); + break; + case GGML_UNARY_OP_STEP: + ggml_cuda_op_step(ctx, dst); + break; + case GGML_UNARY_OP_GELU: + ggml_cuda_op_gelu(ctx, dst); + break; + case GGML_UNARY_OP_SILU: + ggml_cuda_op_silu(ctx, dst); + break; + case GGML_UNARY_OP_GELU_ERF: + ggml_cuda_op_gelu_erf(ctx, dst); + break; + case GGML_UNARY_OP_GELU_QUICK: + ggml_cuda_op_gelu_quick(ctx, dst); + break; + case GGML_UNARY_OP_TANH: + ggml_cuda_op_tanh(ctx, dst); + break; + case GGML_UNARY_OP_RELU: + ggml_cuda_op_relu(ctx, dst); + break; + case GGML_UNARY_OP_SIGMOID: + ggml_cuda_op_sigmoid(ctx, dst); + break; + case GGML_UNARY_OP_HARDSIGMOID: + ggml_cuda_op_hardsigmoid(ctx, dst); + break; + case GGML_UNARY_OP_HARDSWISH: + ggml_cuda_op_hardswish(ctx, dst); + break; + case GGML_UNARY_OP_EXP: + ggml_cuda_op_exp(ctx, dst); + break; + case GGML_UNARY_OP_ELU: + ggml_cuda_op_elu(ctx, dst); + break; + case GGML_UNARY_OP_XIELU: + ggml_cuda_op_xielu(ctx, dst); + break; + case GGML_UNARY_OP_FLOOR: + ggml_cuda_op_floor(ctx, dst); + break; + case GGML_UNARY_OP_CEIL: + ggml_cuda_op_ceil(ctx, dst); + break; + case GGML_UNARY_OP_ROUND: + ggml_cuda_op_round(ctx, dst); + break; + case GGML_UNARY_OP_TRUNC: + ggml_cuda_op_trunc(ctx, dst); + break; + case GGML_UNARY_OP_EXPM1: + ggml_cuda_op_expm1(ctx, dst); + break; + case GGML_UNARY_OP_SOFTPLUS: + ggml_cuda_op_softplus(ctx, dst); + break; + default: + return false; + } + break; + case GGML_OP_GLU: + switch (ggml_get_glu_op(dst)) { + case GGML_GLU_OP_REGLU: + ggml_cuda_op_reglu(ctx, dst); + break; + case GGML_GLU_OP_GEGLU: + ggml_cuda_op_geglu(ctx, dst); + break; + case GGML_GLU_OP_SWIGLU: + ggml_cuda_op_swiglu(ctx, dst); + break; + case GGML_GLU_OP_SWIGLU_OAI: + ggml_cuda_op_swiglu_oai(ctx, dst); + break; + case GGML_GLU_OP_SWIGLU_DS4: + ggml_cuda_op_swiglu_ds4(ctx, dst); + break; + case GGML_GLU_OP_GEGLU_ERF: + ggml_cuda_op_geglu_erf(ctx, dst); + break; + case GGML_GLU_OP_GEGLU_QUICK: + ggml_cuda_op_geglu_quick(ctx, dst); + break; + default: + return false; + } + break; + case GGML_OP_NORM: + ggml_cuda_op_norm(ctx, dst); + break; + case GGML_OP_TURBO_WHT: + ggml_cuda_op_turbo_wht(ctx, dst); + break; + case GGML_OP_MOE_FUSED: + ggml_cuda_op_moe_fused(ctx, dst); + break; + case GGML_OP_GROUP_NORM: + ggml_cuda_op_group_norm(ctx, dst); + break; + case GGML_OP_L2_NORM: + ggml_cuda_op_l2_norm(ctx, dst); + break; + case GGML_OP_CONCAT: + ggml_cuda_op_concat(ctx, dst); + break; + case GGML_OP_UPSCALE: + ggml_cuda_op_upscale(ctx, dst); + break; + case GGML_OP_PAD: + ggml_cuda_op_pad(ctx, dst); + break; + case GGML_OP_PAD_REFLECT_1D: + ggml_cuda_op_pad_reflect_1d(ctx, dst); + break; + case GGML_OP_ARANGE: + ggml_cuda_op_arange(ctx, dst); + break; + case GGML_OP_TIMESTEP_EMBEDDING: + ggml_cuda_op_timestep_embedding(ctx, dst); + break; + case GGML_OP_LEAKY_RELU: + ggml_cuda_op_leaky_relu(ctx, dst); + break; + case GGML_OP_SILU_BACK: + ggml_cuda_op_silu_back(ctx, dst); + break; + case GGML_OP_RMS_NORM: + ggml_cuda_op_rms_norm(ctx, dst); + break; + case GGML_OP_RMS_NORM_BACK: + ggml_cuda_op_rms_norm_back(ctx, dst); + break; + case GGML_OP_MUL_MAT: + ggml_cuda_mul_mat(ctx, dst->src[0], dst->src[1], dst); + break; + case GGML_OP_MUL_MAT_ID: + ggml_cuda_mul_mat_id(ctx, dst); + break; + case GGML_OP_OUT_PROD: + ggml_cuda_out_prod(ctx, dst); + break; + case GGML_OP_SCALE: + ggml_cuda_op_scale(ctx, dst); + break; + case GGML_OP_SQR: + ggml_cuda_op_sqr(ctx, dst); + break; + case GGML_OP_SQRT: + ggml_cuda_op_sqrt(ctx, dst); + break; + case GGML_OP_SIN: + ggml_cuda_op_sin(ctx, dst); + break; + case GGML_OP_COS: + ggml_cuda_op_cos(ctx, dst); + break; + case GGML_OP_CLAMP: + ggml_cuda_op_clamp(ctx, dst); + break; + case GGML_OP_LOG: + ggml_cuda_op_log(ctx, dst); + break; + case GGML_OP_NONE: + case GGML_OP_RESHAPE: + case GGML_OP_VIEW: + case GGML_OP_PERMUTE: + case GGML_OP_TRANSPOSE: + break; + case GGML_OP_DIAG: + ggml_cuda_op_diag(ctx, dst); + break; + case GGML_OP_DIAG_MASK_INF: + ggml_cuda_op_diag_mask_inf(ctx, dst); + break; + case GGML_OP_SOFT_MAX: + ggml_cuda_op_soft_max(ctx, dst); + break; + case GGML_OP_SOFT_MAX_BACK: + ggml_cuda_op_soft_max_back(ctx, dst); + break; + case GGML_OP_ROPE: + ggml_cuda_op_rope(ctx, dst); + break; + case GGML_OP_ROPE_BACK: + ggml_cuda_op_rope_back(ctx, dst); + break; + case GGML_OP_ROLL: + ggml_cuda_op_roll(ctx, dst); + break; + case GGML_OP_IM2COL: + ggml_cuda_op_im2col(ctx, dst); + break; + case GGML_OP_IM2COL_3D: + ggml_cuda_op_im2col_3d(ctx, dst); + break; + case GGML_OP_CONV_2D: + ggml_cuda_op_conv2d(ctx, dst); + break; + case GGML_OP_CONV_2D_DW: + ggml_cuda_op_conv2d_dw(ctx, dst); + break; + case GGML_OP_CONV_TRANSPOSE_2D: + ggml_cuda_conv_2d_transpose_p0(ctx, dst); + break; + case GGML_OP_CONV_TRANSPOSE_1D: + ggml_cuda_op_conv_transpose_1d(ctx,dst); + break; + case GGML_OP_POOL_2D: + ggml_cuda_op_pool2d(ctx, dst); + break; + case GGML_OP_SUM: + ggml_cuda_op_sum(ctx, dst); + break; + case GGML_OP_CUMSUM: + ggml_cuda_op_cumsum(ctx, dst); + break; + case GGML_OP_SUM_ROWS: + ggml_cuda_op_sum_rows(ctx, dst); + break; + case GGML_OP_MEAN: + ggml_cuda_op_mean(ctx, dst); + break; + case GGML_OP_SSM_CONV: + ggml_cuda_op_ssm_conv(ctx, dst); + break; + case GGML_OP_SSM_SCAN: + ggml_cuda_op_ssm_scan(ctx, dst); + break; + case GGML_OP_TOP_K: + ggml_cuda_op_top_k(ctx, dst); + break; + case GGML_OP_ARGSORT: + ggml_cuda_op_argsort(ctx, dst); + break; + case GGML_OP_FLASH_ATTN_EXT: + ggml_cuda_flash_attn_ext(ctx, dst); + break; + case GGML_OP_FLASH_ATTN_SPARSE: + ggml_cuda_flash_attn_sparse(ctx, dst); + break; + case GGML_OP_CROSS_ENTROPY_LOSS: + ggml_cuda_cross_entropy_loss(ctx, dst); + break; + case GGML_OP_TRI: + ggml_cuda_op_tri(ctx, dst); + break; + case GGML_OP_RWKV_WKV6: + ggml_cuda_op_rwkv_wkv6(ctx, dst); + break; + case GGML_OP_GATED_LINEAR_ATTN: + ggml_cuda_op_gated_linear_attn(ctx, dst); + break; + case GGML_OP_GATED_DELTA_NET: + ggml_cuda_op_gated_delta_net(ctx, dst); + break; + case GGML_OP_RWKV_WKV7: + ggml_cuda_op_rwkv_wkv7(ctx, dst); + break; + case GGML_OP_CROSS_ENTROPY_LOSS_BACK: + ggml_cuda_cross_entropy_loss_back(ctx, dst); + break; + case GGML_OP_OPT_STEP_ADAMW: + ggml_cuda_opt_step_adamw(ctx, dst); + break; + case GGML_OP_OPT_STEP_SGD: + ggml_cuda_opt_step_sgd(ctx, dst); + break; + case GGML_OP_SOLVE_TRI: + ggml_cuda_op_solve_tri(ctx, dst); + break; + case GGML_OP_FILL: + ggml_cuda_op_fill(ctx, dst); + break; + default: + return false; + } + + cudaError_t err = cudaGetLastError(); + if (err != cudaSuccess) { + GGML_LOG_ERROR("%s: %s failed: node='%s' [%ld,%ld] src0='%s' t=%d [%ld,%ld] src1='%s' [%ld,%ld]\n", + __func__, ggml_op_desc(dst), dst->name, (long)dst->ne[0], (long)dst->ne[1], + dst->src[0] ? dst->src[0]->name : "-", dst->src[0] ? (int)dst->src[0]->type : -1, + dst->src[0] ? (long)dst->src[0]->ne[0] : 0, dst->src[0] ? (long)dst->src[0]->ne[1] : 0, + dst->src[1] ? dst->src[1]->name : "-", + dst->src[1] ? (long)dst->src[1]->ne[0] : 0, dst->src[1] ? (long)dst->src[1]->ne[1] : 0); + CUDA_CHECK(err); + } + + return true; +} + +//////////////////////////////////////////////////////////////////////////////// + +// backend + +static const char * ggml_backend_cuda_get_name(ggml_backend_t backend) { + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context; + + return cuda_ctx->name.c_str(); +} + +static void ggml_backend_cuda_free(ggml_backend_t backend) { + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context; + + delete cuda_ctx; + delete backend; +} + +static void ggml_backend_cuda_set_tensor_async(ggml_backend_t backend, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context; + ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; + + GGML_ASSERT(buf->buft == ggml_backend_cuda_buffer_type(cuda_ctx->device) && "unsupported buffer type"); + + CUDA_CHECK(cudaMemcpyAsync((char *) tensor->data + offset, data, size, cudaMemcpyHostToDevice, cuda_ctx->stream())); +} + +static void ggml_backend_cuda_get_tensor_async(ggml_backend_t backend, const ggml_tensor * tensor, void * data, size_t offset, size_t size) { + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context; + ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; + + GGML_ASSERT(buf->buft == ggml_backend_cuda_buffer_type(cuda_ctx->device) && "unsupported buffer type"); + + CUDA_CHECK(cudaMemcpyAsync(data, (const char *) tensor->data + offset, size, cudaMemcpyDeviceToHost, cuda_ctx->stream())); +} + +static void ggml_backend_cuda_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) { + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context; + ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; + + GGML_ASSERT(buf->buft == ggml_backend_cuda_buffer_type(cuda_ctx->device) && "unsupported buffer type"); + + CUDA_CHECK(cudaMemcpy2DAsync( + (char *) tensor->data + offset, stride_tensor, data, stride_data, size, n_copies, cudaMemcpyHostToDevice, cuda_ctx->stream())); +} + +static void ggml_backend_cuda_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) { + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context; + ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; + + GGML_ASSERT(buf->buft == ggml_backend_cuda_buffer_type(cuda_ctx->device) && "unsupported buffer type"); + + CUDA_CHECK(cudaMemcpy2DAsync( + data, stride_data, (const char *) tensor->data + offset, stride_tensor, size, n_copies, cudaMemcpyDeviceToHost, cuda_ctx->stream())); +} + +static bool ggml_backend_cuda_cpy_tensor_async(ggml_backend_t backend_src, ggml_backend_t backend_dst, const ggml_tensor * src, ggml_tensor * dst) { + ggml_backend_buffer_t buf_src = src->view_src ? src->view_src->buffer : src->buffer; + ggml_backend_buffer_t buf_dst = dst->view_src ? dst->view_src->buffer : dst->buffer; + + if (!ggml_backend_is_cuda(backend_src) || !ggml_backend_is_cuda(backend_dst)) { + return false; + } + + if (!ggml_backend_buffer_is_cuda(buf_src) || !ggml_backend_buffer_is_cuda(buf_dst)) { + return false; + } + + // device -> device copy + ggml_backend_cuda_context * cuda_ctx_src = (ggml_backend_cuda_context *) backend_src->context; + ggml_backend_cuda_context * cuda_ctx_dst = (ggml_backend_cuda_context *) backend_dst->context; + + ggml_backend_cuda_buffer_context * buf_ctx_src = (ggml_backend_cuda_buffer_context *) buf_src->context; + ggml_backend_cuda_buffer_context * buf_ctx_dst = (ggml_backend_cuda_buffer_context *) buf_dst->context; + + if (cuda_ctx_src->device != buf_ctx_src->device || cuda_ctx_dst->device != buf_ctx_dst->device) { +#ifndef NDEBUG + GGML_LOG_DEBUG("%s: backend and buffer devices do not match\n", __func__); +#endif // NDEBUG + return false; + } + + if (backend_src != backend_dst) { + // copy on src stream + if (cuda_ctx_src->device == cuda_ctx_dst->device) { + CUDA_CHECK(cudaMemcpyAsync(dst->data, src->data, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, cuda_ctx_src->stream())); + } else { +#ifdef GGML_CUDA_NO_PEER_COPY + return false; +#else + CUDA_CHECK(cudaMemcpyPeerAsync(dst->data, cuda_ctx_dst->device, src->data, cuda_ctx_src->device, ggml_nbytes(dst), cuda_ctx_src->stream())); +#endif // GGML_CUDA_NO_PEER_COPY + } + + // record event on src stream after the copy + if (!cuda_ctx_src->copy_event) { + ggml_cuda_set_device(cuda_ctx_src->device); + CUDA_CHECK(cudaEventCreateWithFlags(&cuda_ctx_src->copy_event, cudaEventDisableTiming)); + } + + CUDA_CHECK(cudaEventRecord(cuda_ctx_src->copy_event, cuda_ctx_src->stream())); + + // wait on dst stream for the copy to complete + CUDA_CHECK(cudaStreamWaitEvent(cuda_ctx_dst->stream(), cuda_ctx_src->copy_event, 0)); + } else { + // src and dst are on the same backend + CUDA_CHECK(cudaMemcpyAsync(dst->data, src->data, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, cuda_ctx_src->stream())); + } + return true; +} + +static void ggml_backend_cuda_synchronize(ggml_backend_t backend) { + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context; + + CUDA_CHECK(cudaStreamSynchronize(cuda_ctx->stream())); + + GGML_UNUSED(backend); +} + +#ifdef USE_CUDA_GRAPH +static bool ggml_cuda_graph_check_compability(ggml_cgraph * cgraph) { + + bool use_cuda_graph = true; + // Loop over nodes in GGML graph to obtain info needed for CUDA graph + + for (int i = 0; i < cgraph->n_nodes; i++) { + ggml_tensor * node = cgraph->nodes[i]; + + if (ggml_is_empty(node) || node->op == GGML_OP_RESHAPE || node->op == GGML_OP_TRANSPOSE || node->op == GGML_OP_VIEW || node->op == GGML_OP_PERMUTE || node->op == GGML_OP_NONE) { + continue; + } + + if (node->src[0] && node->src[0]->buffer && ggml_backend_buft_is_cuda_split(node->src[0]->buffer->buft)) { + use_cuda_graph = false; // Split buffers are not supported by CUDA graph capture +#ifndef NDEBUG + GGML_LOG_DEBUG("%s: disabling CUDA graphs due to split buffer\n", __func__); +#endif + } + + // [TAG_MUL_MAT_ID_CUDA_GRAPHS] + if (node->op == GGML_OP_MUL_MAT_ID) { + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + const int mmvq_mmid_max = get_mmvq_mmid_max_batch(node->src[0]->type, cc); + if (!ggml_is_quantized(node->src[0]->type) || node->ne[2] > mmvq_mmid_max) { + // under these conditions, the mul_mat_id operation will need to synchronize the stream, so we cannot use CUDA graphs + // TODO: figure out a way to enable for larger batch sizes, without hurting performance + // ref: https://github.com/ggml-org/llama.cpp/pull/18958 + use_cuda_graph = false; +#ifndef NDEBUG + GGML_LOG_DEBUG("%s: disabling CUDA graphs due to unsupported node type\n", __func__); +#endif + } + } + + if (!use_cuda_graph) { + break; + } + } + + return use_cuda_graph; +} + +static const void * ggml_cuda_graph_get_key(ggml_cgraph * cgraph) { + return cgraph->nodes[0]; +} + +static bool ggml_cuda_graph_update_required(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph) { + bool res = false; + + const void * graph_key = ggml_cuda_graph_get_key(cgraph); + ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); + + // Check if the graph size has changed + if ((int)graph->node_props.size() != cgraph->n_nodes) { + res = true; + graph->node_props.resize(cgraph->n_nodes); + } + + for (int i = 0; i < cgraph->n_nodes; i++) { + ggml_cuda_graph::node_properties prop = {}; + memcpy(&prop.node, cgraph->nodes[i], sizeof(ggml_tensor)); + + for (int j = 0; j < GGML_MAX_SRC; ++j) { + if (cgraph->nodes[i]->src[j]) { + prop.node_src_data_ptrs[j] = cgraph->nodes[i]->src[j]->data; + memcpy(prop.node_src_ne[j], cgraph->nodes[i]->src[j]->ne, sizeof(prop.node_src_ne[j])); + memcpy(prop.node_src_nb[j], cgraph->nodes[i]->src[j]->nb, sizeof(prop.node_src_nb[j])); + } + } + + if (res || memcmp(&graph->node_props[i], &prop, sizeof(prop)) != 0) { + graph->node_props[i] = prop; + res = true; + } + } + + return res; +} + +static void ggml_cuda_graph_update_executable(ggml_backend_cuda_context * cuda_ctx, const void * graph_key) { + ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); + +#if CUDART_VERSION >= 12000 + cudaGraphExecUpdateResultInfo result_info; + cudaError_t stat = cudaGraphExecUpdate(graph->instance, graph->graph, &result_info); +#else + cudaGraphNode_t errorNode; + cudaGraphExecUpdateResult result_info; + cudaError_t stat = cudaGraphExecUpdate(graph->instance, graph->graph, &errorNode, &result_info); +#endif // CUDART_VERSION >= 12000 + + if (stat == cudaErrorGraphExecUpdateFailure) { +#ifndef NDEBUG + GGML_LOG_DEBUG("%s: CUDA graph update failed\n", __func__); +#endif + + // The pre-existing graph exec cannot be updated due to violated constraints + // so instead clear error and re-instantiate + (void)cudaGetLastError(); + CUDA_CHECK(cudaGraphExecDestroy(graph->instance)); + graph->instance = nullptr; + CUDA_CHECK(cudaGraphInstantiate(&graph->instance, graph->graph, NULL, NULL, 0)); + } else { + GGML_ASSERT(stat == cudaSuccess); + } +} +#endif // USE_CUDA_GRAPH + +static bool ggml_cuda_should_fuse_rope_set_rows(const ggml_tensor * rope, + const ggml_tensor * view, + const ggml_tensor * set_rows) { + + if (rope->op != GGML_OP_ROPE || view->op != GGML_OP_VIEW || set_rows->op != GGML_OP_SET_ROWS) { + return false; + } + // ne3 not tested + if (rope->src[0]->ne[3] != 1) { + return false; + } + + if (set_rows->type != GGML_TYPE_F32 && set_rows->type != GGML_TYPE_F16) { + return false; + } + + if (set_rows->src[1]->type != GGML_TYPE_I64) { + return false; + } + + // The view should flatten two dims of rope into one dim + if (!ggml_is_contiguous(view) || view->ne[0] != rope->ne[0] * rope->ne[1]) { + return false; + } + + // Only norm/neox shaders have the fusion code + const int mode = ((const int32_t *) rope->op_params)[2]; + if (mode != GGML_ROPE_TYPE_NORMAL && mode != GGML_ROPE_TYPE_NEOX) { + return false; + } + + return true; +} + +static bool ggml_cuda_topk_moe_fusion(const struct ggml_cgraph * cgraph, int node_idx, ggml_cuda_topk_moe_args & args) { + args.sigmoid = false; + args.softmax = false; + args.delayed_softmax = false; + args.prob_bias = false; + args.norm = false; + + const int n_nodes = cgraph->n_nodes; + ggml_tensor ** nodes = cgraph->nodes; + + if (nodes[node_idx]->op == GGML_OP_SOFT_MAX) { + args.softmax = true; + } + + if (nodes[node_idx]->op == GGML_OP_UNARY) { + if (ggml_get_unary_op(nodes[node_idx]) != GGML_UNARY_OP_SIGMOID) { + return false; + } + args.sigmoid = true; + } + + if (nodes[node_idx]->op == GGML_OP_ARGSORT) { + args.delayed_softmax = true; + } + + node_idx++; + + if (args.sigmoid || args.softmax) { + // SOFTMAX -> RESHAPE + if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_RESHAPE || + nodes[node_idx]->src[0] != nodes[node_idx - 1]) { + return false; + } + ggml_tensor * probs_reshaped = nodes[node_idx]; + node_idx++; + + if (node_idx >= n_nodes) { + return false; + } + + // src of bias add is the unreshaped probs (-2 instead of -1) + if (nodes[node_idx]->op == GGML_OP_ADD && nodes[node_idx]->src[0] == nodes[node_idx - 2]) { + args.prob_bias = true; + node_idx++; + } + // RESHAPE/ADD -> ARGSORT + if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_ARGSORT) { + return false; + } + + if (args.prob_bias && nodes[node_idx]->src[0] != nodes[node_idx - 1]) { + return false; + } else if (!args.prob_bias && nodes[node_idx]->src[0] != nodes[node_idx - 2]) { + return false; + } + + node_idx++; + + // ARGSORT-> VIEW + if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_VIEW || + nodes[node_idx]->src[0] != nodes[node_idx - 1]) { + return false; + } + node_idx++; + + if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_GET_ROWS) { + return false; + } + + // GET_ROWS + if (nodes[node_idx]->src[0] != probs_reshaped || nodes[node_idx]->src[1] != nodes[node_idx - 1]) { + return false; + } + node_idx++; + } else if (args.delayed_softmax) { + if (node_idx - 2 < 0) { + return false; + } + ggml_tensor * probs_reshaped = nodes[node_idx - 2]; + + // VIEW->ARGSORT + if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_VIEW || + nodes[node_idx]->src[0] != nodes[node_idx - 1]) { + return false; + } + node_idx++; + + // GET_ROWS + if (node_idx >= n_nodes || nodes[node_idx]->src[1] != nodes[node_idx - 1] || + nodes[node_idx]->src[0] != probs_reshaped) { + return false; + } + node_idx++; + + static const std::vector remaining_ops = { GGML_OP_RESHAPE, GGML_OP_SOFT_MAX, GGML_OP_RESHAPE }; + + for (const ggml_op op : remaining_ops) { + if (node_idx >= n_nodes || nodes[node_idx]->op != op || nodes[node_idx]->src[0] != nodes[node_idx - 1]) { + return false; + } + node_idx++; + } + } + + // At this point we can check for norm + scale. Everything is now at least valid till the norm + if (node_idx >= n_nodes) { + return true; + } + + if (nodes[node_idx]->op == GGML_OP_RESHAPE) { + //check RESHAPE->SUM_ROWS->CLAMP->DIV->RESHAPE + static const std::vector norm_ops = { GGML_OP_RESHAPE, GGML_OP_SUM_ROWS, GGML_OP_CLAMP }; + + args.norm = true; + for (const ggml_op op : norm_ops) { + if (nodes[node_idx]->op == op && nodes[node_idx]->src[0] == nodes[node_idx - 1]) { + node_idx++; + } else { + args.norm = false; + return true; + } + } + + // DIV <- CLAMP, RESHAPE + if (nodes[node_idx]->op != GGML_OP_DIV || nodes[node_idx]->src[1] != nodes[node_idx - 1] || + nodes[node_idx]->src[0] != nodes[node_idx - 3]) { + args.norm = false; + return true; + } + node_idx++; + + if (nodes[node_idx]->op != GGML_OP_RESHAPE || nodes[node_idx]->src[0] != nodes[node_idx - 1]) { + args.norm = false; + return true; + } + + node_idx++; + } + + if (nodes[node_idx]->op == GGML_OP_SCALE && nodes[node_idx]->src[0] == nodes[node_idx - 1]) { + args.scale = true; + } + + return true; +} + +// returns whether the write (out) nodes overwrite the read nodes in operation +static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph, + const int node_idx, + const int node_count, + const int * out_nodes, + const int out_count, + const bool is_topk_moe = false) { + auto nodes_overlap = [&](const ggml_tensor * a, const ggml_tensor * b) { + const int64_t a_start = (int64_t) a->data; + const int64_t a_end = a_start + ggml_backend_buft_get_alloc_size(a->buffer->buft, a); + + const int64_t b_start = (int64_t) b->data; + const int64_t b_end = b_start + ggml_backend_buft_get_alloc_size(b->buffer->buft, b); + + if ((b_start <= a_start && a_start < b_end) || (a_start <= b_start && b_start < a_end)) { + return true; + } + + return false; + }; + + bool is_ok = true; + // exception for topk-moe, as each row is read entirely before writing + if (ggml_nrows(cgraph->nodes[node_idx]) == 1 && is_topk_moe) { + return true; + } + + for (int i = 0; i < out_count; ++i) { + const ggml_tensor * dst = cgraph->nodes[out_nodes[i]]; + + for (int j = node_idx; j < node_idx + node_count; ++j) { + // Loop over all srcs of all nodes in the fusion. If the src overlaps + // the destination and the src is not an intermediate node that's being + // elided, then disable fusion. + + for (int src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) { + const ggml_tensor * src = cgraph->nodes[j]->src[src_idx]; + + if (!src || src->op == GGML_OP_NONE) { + continue; + } + + if (nodes_overlap(dst, src)) { + bool found = false; + + for (int k = node_idx; k < j; ++k) { + if (cgraph->nodes[k] == src) { + found = true; + break; + } + } + + if (!found) { + is_ok = false; + break; + } + } + } + } + } + + return is_ok; +} + + +static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, + int node_idx, + std::initializer_list ops, + std::initializer_list unary_ops) { +#ifndef NDEBUG + const size_t num_unary = std::count(ops.begin(), ops.end(), GGML_OP_UNARY); + GGML_ASSERT(unary_ops.size() == num_unary); +#endif + + const auto is_equal = [](const std::initializer_list & list1, + const std::initializer_list & list2) { + return std::equal(list1.begin(), list1.end(), list2.begin(), list2.end()); + }; + + std::initializer_list mul_mat_bias_glu_ops = { GGML_OP_MUL_MAT, GGML_OP_ADD, GGML_OP_MUL_MAT, GGML_OP_ADD, GGML_OP_GLU }; + std::initializer_list mul_mat_id_bias_glu_ops = { GGML_OP_MUL_MAT_ID, GGML_OP_ADD_ID, GGML_OP_MUL_MAT_ID, GGML_OP_ADD_ID, GGML_OP_GLU }; + + std::initializer_list mul_mat_id_glu_ops = { GGML_OP_MUL_MAT_ID, GGML_OP_MUL_MAT_ID, GGML_OP_GLU }; + std::initializer_list mul_mat_glu_ops = { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT, GGML_OP_GLU }; + + if ((is_equal(mul_mat_bias_glu_ops, ops) || is_equal(mul_mat_id_bias_glu_ops, ops)) && + ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 4 })) { + const ggml_tensor * ffn_gate = cgraph->nodes[node_idx]; + const ggml_tensor * ffn_gate_bias = cgraph->nodes[node_idx + 1]; + const ggml_tensor * ffn_up = cgraph->nodes[node_idx + 2]; + const ggml_tensor * ffn_up_bias = cgraph->nodes[node_idx + 3]; + const ggml_tensor * glu = cgraph->nodes[node_idx + 4]; + + if (ggml_cuda_should_fuse_mul_mat(ffn_up, ffn_gate, glu, ffn_up_bias, ffn_gate_bias)) { + int out_nodes[] = { node_idx + 4 }; + return ggml_cuda_check_fusion_memory_ranges(cgraph, node_idx, (int)ops.size(), out_nodes, 1); + } + } + + if ((is_equal(mul_mat_id_glu_ops, ops) || is_equal(mul_mat_glu_ops, ops)) && + ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 2 })) { + const ggml_tensor * ffn_gate = cgraph->nodes[node_idx]; + const ggml_tensor * ffn_up = cgraph->nodes[node_idx + 1]; + const ggml_tensor * glu = cgraph->nodes[node_idx + 2]; + + if (ggml_cuda_should_fuse_mul_mat(ffn_up, ffn_gate, glu)) { + int out_nodes[] = { node_idx + 2 }; + return ggml_cuda_check_fusion_memory_ranges(cgraph, node_idx, (int)ops.size(), out_nodes, 1); + } + } + + std::initializer_list rope_set_rows_ops = { GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS }; + + if (is_equal(rope_set_rows_ops, ops) && ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 2 })) { + const ggml_tensor * rope = cgraph->nodes[node_idx]; + const ggml_tensor * view = cgraph->nodes[node_idx + 1]; + const ggml_tensor * set_rows = cgraph->nodes[node_idx + 2]; + + if (ggml_cuda_should_fuse_rope_set_rows(rope, view, set_rows)) { + return true; + } + } + + if (!ggml_can_fuse(cgraph, node_idx, ops)) { + return false; + } + + if ((ops.size() == 2 || ops.size() == 3) && ops.begin()[0] == GGML_OP_RMS_NORM && ops.begin()[1] == GGML_OP_MUL) { + const ggml_tensor *rms_norm = cgraph->nodes[node_idx]; + const ggml_tensor *mul = cgraph->nodes[node_idx+1]; + const ggml_tensor *add = nullptr; + + if (ops.size() == 3 && ops.begin()[2] == GGML_OP_ADD) { + add = cgraph->nodes[node_idx+2]; + } + + GGML_ASSERT(rms_norm->src[0]->type == GGML_TYPE_F32); + GGML_ASSERT(rms_norm->type == GGML_TYPE_F32); + + //rms norm only supports F32 + if (mul->src[0]->type != GGML_TYPE_F32 || + mul->src[1]->type != GGML_TYPE_F32 || + mul->type != GGML_TYPE_F32) { + return false; + } + + if (add && (add->src[0]->type != GGML_TYPE_F32 || + add->src[1]->type != GGML_TYPE_F32 || + add->type != GGML_TYPE_F32) ) { + return false; + } + + //if rms norm is the B operand, then we don't handle broadcast + if (rms_norm == mul->src[1] && !ggml_are_same_shape(mul->src[0], rms_norm)) { + return false; + } + + //rms_norm kernel assumes contiguous rows + if (!ggml_is_contiguous_rows(mul->src[0]) || !ggml_is_contiguous_rows(mul->src[1])) { + return false; + } + + if (add && (!ggml_is_contiguous(add->src[0]) || !ggml_is_contiguous_rows(add->src[1]))) { + return false; + } + + return true; + } + + if (ops.size() == 2 && ops.begin()[0] == GGML_OP_SSM_CONV && ops.begin()[1] == GGML_OP_UNARY + && unary_ops.size() == 1 && unary_ops.begin()[0] == GGML_UNARY_OP_SILU) { + const ggml_tensor * ssm_conv = cgraph->nodes[node_idx]; + const ggml_tensor * silu = cgraph->nodes[node_idx+1]; + + if (ssm_conv->type != GGML_TYPE_F32 || silu->type != GGML_TYPE_F32) { + return false; + } + + return true; + } + + if (ops.size() == 2 && ops.begin()[0] == GGML_OP_UNARY && ops.begin()[1] == GGML_OP_MUL + && unary_ops.size() == 1 && (unary_ops.begin()[0] == GGML_UNARY_OP_SILU || unary_ops.begin()[0] == GGML_UNARY_OP_SIGMOID || unary_ops.begin()[0] == GGML_UNARY_OP_SOFTPLUS)) { + const ggml_tensor * unary = cgraph->nodes[node_idx]; + const ggml_tensor * mul = cgraph->nodes[node_idx+1]; + + if (ggml_get_unary_op(unary) != unary_ops.begin()[0]) { + return false; + } + + if (unary->type != GGML_TYPE_F32 && unary->type != GGML_TYPE_F16) { + return false; + } + + if (unary->type != mul->type) { + return false; + } + + const ggml_tensor * other = (mul->src[0] == unary) ? mul->src[1] : mul->src[0]; + if (other->type != unary->type) { + return false; + } + if (!ggml_is_contiguous_1(other) || !ggml_is_contiguous_1(unary->src[0]) || !ggml_are_same_shape(other, unary)) { + return false; + } + + return true; + } + + if (ops.size() == 3 && ops.begin()[0] == GGML_OP_SCALE && ops.begin()[1] == GGML_OP_UNARY && ops.begin()[2] == GGML_OP_SCALE + && unary_ops.size() == 1 && unary_ops.begin()[0] == GGML_UNARY_OP_TANH) { + const ggml_tensor *scale = cgraph->nodes[node_idx]; + const ggml_tensor *tanh = cgraph->nodes[node_idx+1]; + const ggml_tensor *scale2 = cgraph->nodes[node_idx+2]; + + GGML_ASSERT(scale->src[0]->type == GGML_TYPE_F32); + GGML_ASSERT(scale->type == GGML_TYPE_F32); + + if (ggml_get_unary_op(tanh) != GGML_UNARY_OP_TANH) { + return false; + } + + // Check for bias + if (ggml_get_op_params_f32(scale, 1) != 0.0f || ggml_get_op_params_f32(scale2, 1) != 0.0f) { + return false; + } + + return true; + } + + return false; +} + +static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, const bool use_cuda_graph, const bool cuda_graph_update_required, const void * graph_key) { + bool graph_evaluated_or_captured = false; + + // flag used to determine whether it is an integrated_gpu + const bool integrated = ggml_cuda_info().devices[cuda_ctx->device].integrated; + + ggml_cuda_stream_context & stream_ctx = cuda_ctx->stream_context(); + bool is_concurrent_event_active = false; + ggml_cuda_concurrent_event * concurrent_event = nullptr; + bool should_launch_concurrent_events = false; + + const auto try_launch_concurrent_event = [&](const ggml_tensor * node) { + if (stream_ctx.concurrent_events.find(node) != stream_ctx.concurrent_events.end()) { + concurrent_event = &stream_ctx.concurrent_events[node]; + + is_concurrent_event_active = true; + + GGML_LOG_DEBUG("Launching %d streams at %s\n", concurrent_event->n_streams, node->name); + + cudaStream_t main_stream = cuda_ctx->stream(); // this should be stream 0 + GGML_ASSERT(cuda_ctx->curr_stream_no == 0); + CUDA_CHECK(cudaEventRecord(concurrent_event->fork_event, main_stream)); + + for (int i = 1; i <= concurrent_event->n_streams; ++i) { + cudaStream_t stream = cuda_ctx->stream(cuda_ctx->device, i); + CUDA_CHECK(cudaStreamWaitEvent(stream, concurrent_event->fork_event)); + } + } + }; + + while (!graph_evaluated_or_captured) { + // Only perform the graph execution if CUDA graphs are not enabled, or we are capturing the graph. + // With the use of CUDA graphs, the execution will be performed by the graph launch. + if (!use_cuda_graph || cuda_graph_update_required) { + [[maybe_unused]] int prev_i = 0; + + if (stream_ctx.concurrent_events.size() > 0) { + should_launch_concurrent_events = true; + for (const auto & [tensor, event] : stream_ctx.concurrent_events) { + should_launch_concurrent_events = should_launch_concurrent_events && event.is_valid(); + } + } + + if (should_launch_concurrent_events) { + // Restore original node order within each concurrent region to enable fusion within streams + + std::unordered_map node_to_idx; + node_to_idx.reserve(cgraph->n_nodes); + for (int i = 0; i < cgraph->n_nodes; ++i) { + node_to_idx[cgraph->nodes[i]] = i; + } + + for (auto & [fork_node, event] : stream_ctx.concurrent_events) { + // Find positions of all nodes from this event in the current graph + std::vector positions; + positions.reserve(event.original_order.size()); + + bool all_found = true; + for (const ggml_tensor * orig_node : event.original_order) { + auto it = node_to_idx.find(orig_node); + if (it != node_to_idx.end()) { + positions.push_back(it->second); + } else { + all_found = false; + break; + } + } + + if (!all_found || positions.size() != event.original_order.size()) { + continue; + } + + // Sort positions to get contiguous range + std::vector sorted_positions = positions; + std::sort(sorted_positions.begin(), sorted_positions.end()); + + bool is_contiguous = true; + for (size_t i = 1; i < sorted_positions.size(); ++i) { + if (sorted_positions[i] != sorted_positions[i-1] + 1) { + is_contiguous = false; + break; + } + } + + if (!is_contiguous) { + continue; + } + + // Restore original order at the sorted positions + int start_pos = sorted_positions[0]; + for (size_t i = 0; i < event.original_order.size(); ++i) { + cgraph->nodes[start_pos + i] = const_cast(event.original_order[i]); + } + } + } else { + stream_ctx.concurrent_events.clear(); + } + + for (int i = 0; i < cgraph->n_nodes; i++) { + ggml_tensor * node = cgraph->nodes[i]; + if (is_concurrent_event_active) { + GGML_ASSERT(concurrent_event); + + if (node == concurrent_event->join_node) { + cuda_ctx->curr_stream_no = 0; + for (int i = 1; i <= concurrent_event->n_streams; ++i) { + // Wait on join events of forked streams in the main stream + CUDA_CHECK(cudaEventRecord(concurrent_event->join_events[i - 1], + cuda_ctx->stream(cuda_ctx->device, i))); + CUDA_CHECK(cudaStreamWaitEvent(cuda_ctx->stream(), concurrent_event->join_events[i - 1])); + } + + is_concurrent_event_active = false; + concurrent_event = nullptr; + } else { + GGML_ASSERT (concurrent_event->stream_mapping.find(node) != concurrent_event->stream_mapping.end()); + cuda_ctx->curr_stream_no = concurrent_event->stream_mapping[node]; + GGML_LOG_DEBUG("Setting stream no to %d for node %s\n", cuda_ctx->curr_stream_no, node->name); + } + } else if (i - prev_i > 1) { + //the previous node was fused + const ggml_tensor * prev_node = cgraph->nodes[i - 1]; + try_launch_concurrent_event(prev_node); + + if (is_concurrent_event_active) { + cuda_ctx->curr_stream_no = concurrent_event->stream_mapping[node]; + GGML_LOG_DEBUG("Setting stream no to %d for node %s\n", cuda_ctx->curr_stream_no, node->name); + } + } + +#ifdef GGML_CUDA_DEBUG + const int nodes_fused = i - prev_i - 1; + if (nodes_fused > 0) { + GGML_LOG_INFO("nodes_fused: %d\n", nodes_fused); + } +#endif + prev_i = i; + + if (ggml_is_empty(node) || node->op == GGML_OP_RESHAPE || node->op == GGML_OP_TRANSPOSE || node->op == GGML_OP_VIEW || node->op == GGML_OP_PERMUTE || node->op == GGML_OP_NONE) { + continue; + } + + if ((node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { + continue; + } + + // start of fusion operations + static bool disable_fusion = (getenv("GGML_CUDA_DISABLE_FUSION") != nullptr); + if (!disable_fusion) { + ggml_cuda_topk_moe_args args; + + if (cgraph->nodes[i]->op == GGML_OP_UNARY || cgraph->nodes[i]->op == GGML_OP_SOFT_MAX || + cgraph->nodes[i]->op == GGML_OP_ARGSORT) { + const bool can_fuse = ggml_cuda_topk_moe_fusion(cgraph, i, args); + + std::vector ops; + + if (can_fuse) { + const ggml_tensor * logits = node->src[0]; + ggml_tensor * weights = nullptr; + ggml_tensor * ids = nullptr; + const ggml_tensor * bias = nullptr; + const ggml_tensor * clamp = nullptr; + const ggml_tensor * scale = nullptr; + + if (!args.delayed_softmax) { + ggml_op gating_op = args.sigmoid ? GGML_OP_UNARY : GGML_OP_SOFT_MAX; + int out_nodes[2]; // nodes which can't be elided + + if (args.prob_bias) { + bias = cgraph->nodes[i + 2]->src[1]; + ops.insert(ops.end(), { gating_op, GGML_OP_RESHAPE, GGML_OP_ADD, GGML_OP_ARGSORT, + GGML_OP_VIEW, GGML_OP_GET_ROWS }); + out_nodes[0] = i + 4; + ids = cgraph->nodes[i + 4]; + } else { + ops.insert(ops.end(), { gating_op, GGML_OP_RESHAPE, GGML_OP_ARGSORT, GGML_OP_VIEW, + GGML_OP_GET_ROWS }); + out_nodes[0] = i + 3; + ids = cgraph->nodes[i + 3]; + } + + if (args.norm) { + ops.insert(ops.end(), { GGML_OP_RESHAPE, GGML_OP_SUM_ROWS, GGML_OP_CLAMP, + GGML_OP_DIV, GGML_OP_RESHAPE }); + clamp = cgraph->nodes[i + ops.size() - 3]; + } + if (args.scale) { + ops.insert(ops.end(), { GGML_OP_SCALE }); + scale = cgraph->nodes[i + ops.size() - 1]; + } + + weights = cgraph->nodes[i + ops.size() - 1]; + out_nodes[1] = i + ops.size() - 1; + + if (ggml_can_fuse_subgraph(cgraph, i, ops.size(), ops.data(), out_nodes, 2) && + ggml_cuda_should_use_topk_moe(node, logits, weights, ids) && + ggml_cuda_check_fusion_memory_ranges(cgraph, i, ops.size(), out_nodes, 2, /*is_topk_moe=*/ true)) { + ggml_cuda_op_topk_moe(*cuda_ctx, logits, weights, ids, clamp, scale, bias, args); + i += ops.size() - 1; + continue; + } + } else if (!args.norm && !args.prob_bias) { + //special case gpt-oss, no norm, no bias. + ops.insert(ops.end(), { GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS, + GGML_OP_RESHAPE, GGML_OP_SOFT_MAX, GGML_OP_RESHAPE }); + weights = cgraph->nodes[i + 5]; + ids = cgraph->nodes[i + 1]; + const ggml_tensor * softmax = cgraph->nodes[i + 4]; + + int out_nodes[2] = { i + 1, i + 5 }; + if (ggml_can_fuse_subgraph(cgraph, i, ops.size(), ops.data(), out_nodes, 2) && + ggml_cuda_should_use_topk_moe(softmax, logits, weights, ids) && + ggml_cuda_check_fusion_memory_ranges(cgraph, i, ops.size(), out_nodes, 2, /*is_topk_moe=*/ true)) { + ggml_cuda_op_topk_moe(*cuda_ctx, logits, weights, ids, clamp, scale, bias, args); + i += ops.size() - 1; + continue; + } + } + } + } + + if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS }, {})) { + ggml_tensor * rope = cgraph->nodes[i]; + ggml_tensor * set_rows = cgraph->nodes[i + 2]; + + ggml_cuda_op_rope_fused(*cuda_ctx, rope, set_rows); + i += 2; + continue; + } + + if (node->op == GGML_OP_ADD || node->op == GGML_OP_MUL) { + int n_fuse = 0; + ggml_op ops[8]; + std::fill(ops, ops + 8, node->op); + + for (; n_fuse <= 6; ++n_fuse){ + if (!ggml_can_fuse(cgraph, i + n_fuse, ops + n_fuse, 2)) { + break; + } + if (cgraph->nodes[i + n_fuse] != cgraph->nodes[i + n_fuse + 1]->src[0]) { + break; + } + if (!ggml_are_same_layout(cgraph->nodes[i + n_fuse]->src[1], cgraph->nodes[i + n_fuse + 1]->src[1])) { + break; + } + } + + n_fuse++; + + if (n_fuse > 1) { + ggml_tensor fused_node; + memcpy(&fused_node, node, sizeof(ggml_tensor)); + for (int j = 0; j < n_fuse - 1; ++j) { + fused_node.src[j + 2] = cgraph->nodes[i + j + 1]->src[1]; + } + fused_node.data = cgraph->nodes[i + n_fuse - 1]->data; + if (node->op == GGML_OP_ADD) { + ggml_cuda_op_fused_add(*cuda_ctx, &fused_node, n_fuse); + } else { + ggml_cuda_op_fused_mul(*cuda_ctx, &fused_node, n_fuse); + } + i += n_fuse - 1; + + continue; + } + } + + bool fused_mul_mat_vec = false; + int fused_node_count = 0; + + for (ggml_op op : { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT_ID }) { + const ggml_op bias_op = op == GGML_OP_MUL_MAT ? GGML_OP_ADD : GGML_OP_ADD_ID; + + if (ggml_cuda_can_fuse(cgraph, i, { op, bias_op, op, bias_op, GGML_OP_GLU }, {})) { + ggml_tensor * glu = cgraph->nodes[i + 4]; + ggml_tensor * gate_bias_n = glu->src[0]; + ggml_tensor * up_bias_n = glu->src[1]; + + //we don't assume the order for {gate, up}. Instead infer it from the bias tensor + ggml_tensor * gate_n = nullptr; + ggml_tensor * up_n = nullptr; + + if (gate_bias_n->src[0] == cgraph->nodes[i] || gate_bias_n->src[1] == cgraph->nodes[i]) { + gate_n = cgraph->nodes[i]; + up_n = cgraph->nodes[i + 2]; + } else if (gate_bias_n->src[0] == cgraph->nodes[i + 2] || gate_bias_n->src[1] == cgraph->nodes[i + 2]) { + gate_n = cgraph->nodes[i + 2]; + up_n = cgraph->nodes[i]; + } else { + continue; + } + + auto get_bias_tensor = [](const ggml_tensor * bias_node, const ggml_tensor * mul_node, ggml_op op_bias) { + if (op_bias == GGML_OP_ADD) { + if (bias_node->src[0] == mul_node) { + return bias_node->src[1]; + } + if (bias_node->src[1] == mul_node) { + return bias_node->src[0]; + } + return (ggml_tensor *) nullptr; + } + GGML_ASSERT(op_bias == GGML_OP_ADD_ID); + GGML_ASSERT(bias_node->src[0] == mul_node); + return bias_node->src[1]; + }; + + ggml_tensor * up_bias_tensor = get_bias_tensor(up_bias_n, up_n, bias_op); + ggml_tensor * gate_bias_tensor = get_bias_tensor(gate_bias_n, gate_n, bias_op); + + if (!up_bias_tensor || !gate_bias_tensor) { + continue; + } + + // we don't support repeating adds + if (bias_op == GGML_OP_ADD && + (!ggml_are_same_shape(gate_bias_n->src[0], gate_bias_n->src[1]) || + !ggml_are_same_shape(up_bias_n->src[0], up_bias_n->src[1]))) { + continue; + } + + const ggml_tensor * src0 = up_n->src[0]; + const ggml_tensor * src1 = up_n->src[1]; + const ggml_tensor * ids = up_n->src[2]; + + if (ggml_cuda_should_fuse_mul_mat_vec_f(up_n)) { + ggml_cuda_mm_fusion_args_host fusion_data{}; + fusion_data.gate = gate_n->src[0]; + fusion_data.x_bias = up_bias_tensor; + fusion_data.gate_bias = gate_bias_tensor; + ggml_cuda_set_fusion_glu_params(fusion_data, glu); + + ggml_cuda_mul_mat_vec_f(*cuda_ctx, src0, src1, ids, glu, &fusion_data); + fused_mul_mat_vec = true; + fused_node_count = 5; + break; + } + + if (ggml_cuda_should_fuse_mul_mat_vec_q(up_n)) { + ggml_cuda_mm_fusion_args_host fusion_data{}; + fusion_data.gate = gate_n->src[0]; + fusion_data.x_bias = up_bias_tensor; + fusion_data.gate_bias = gate_bias_tensor; + ggml_cuda_set_fusion_glu_params(fusion_data, glu); + + ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, glu, &fusion_data); + fused_mul_mat_vec = true; + fused_node_count = 5; + break; + } + } else if (ggml_cuda_can_fuse(cgraph, i, { op, op, GGML_OP_GLU }, {})) { + ggml_tensor * glu = cgraph->nodes[i + 2]; + ggml_tensor * gate = glu->src[0]; + ggml_tensor * up = glu->src[1]; + + bool ok = (gate == cgraph->nodes[i] && up == cgraph->nodes[i + 1]) + || (gate == cgraph->nodes[i + 1] && up == cgraph->nodes[i]); + + if (!ok) continue; + + const ggml_tensor * src0 = up->src[0]; + const ggml_tensor * src1 = up->src[1]; + const ggml_tensor * ids = up->src[2]; + + if (ggml_cuda_should_fuse_mul_mat_vec_f(up)) { + ggml_cuda_mm_fusion_args_host fusion_data{}; + fusion_data.gate = gate->src[0]; + ggml_cuda_set_fusion_glu_params(fusion_data, glu); + + ggml_cuda_mul_mat_vec_f(*cuda_ctx, src0, src1, ids, glu, &fusion_data); + fused_mul_mat_vec = true; + fused_node_count = 3; + break; + } + + if (ggml_cuda_should_fuse_mul_mat_vec_q(up)) { + ggml_cuda_mm_fusion_args_host fusion_data{}; + fusion_data.gate = gate->src[0]; + ggml_cuda_set_fusion_glu_params(fusion_data, glu); + + ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, glu, &fusion_data); + fused_mul_mat_vec = true; + fused_node_count = 3; + break; + } + } + } + + if (fused_mul_mat_vec) { + i += fused_node_count - 1; + continue; + } + + fused_mul_mat_vec = false; + fused_node_count = 0; + + for (ggml_op op : { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT_ID }) { + const ggml_op bias_op = op == GGML_OP_MUL_MAT ? GGML_OP_ADD : GGML_OP_ADD_ID; + + if (!ggml_can_fuse(cgraph, i, { op, bias_op })) { + continue; + } + + ggml_tensor * mm_node = cgraph->nodes[i]; + ggml_tensor * bias_node = cgraph->nodes[i + 1]; + + ggml_tensor * bias_tensor = nullptr; + if (bias_op == GGML_OP_ADD) { + if (bias_node->src[0] == mm_node) { + bias_tensor = bias_node->src[1]; + } else if (bias_node->src[1] == mm_node) { + bias_tensor = bias_node->src[0]; + } else { + continue; + } + } else { + if (bias_node->src[0] != mm_node) { + continue; + } + bias_tensor = bias_node->src[1]; + } + + const ggml_tensor * src0 = mm_node->src[0]; + const ggml_tensor * src1 = mm_node->src[1]; + const ggml_tensor * ids = mm_node->src[2]; + + if (bias_op == GGML_OP_ADD_ID && bias_node->src[2] != ids) { + continue; + } + + if (bias_op == GGML_OP_ADD && !ggml_are_same_shape(bias_node->src[0], bias_node->src[1])) { + continue; + } + + ggml_cuda_mm_fusion_args_host fusion_data{}; + fusion_data.x_bias = bias_tensor; + + if (ggml_cuda_should_fuse_mul_mat_vec_f(mm_node)) { + ggml_cuda_mul_mat_vec_f(*cuda_ctx, src0, src1, ids, bias_node, &fusion_data); + fused_mul_mat_vec = true; + fused_node_count = 2; + break; + } + + if (ggml_cuda_should_fuse_mul_mat_vec_q(mm_node)) { + ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, bias_node, &fusion_data); + fused_mul_mat_vec = true; + fused_node_count = 2; + break; + } + } + + if (fused_mul_mat_vec) { + i += fused_node_count - 1; + continue; + } + + if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD}, {})) { + ggml_cuda_op_rms_norm_fused_add(*cuda_ctx, node, cgraph->nodes[i+1], cgraph->nodes[i+2]); + i += 2; + continue; + } + + if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL}, {})) { + ggml_cuda_op_rms_norm_fused(*cuda_ctx, node, cgraph->nodes[i+1]); + i++; + continue; + } + + if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_SSM_CONV, GGML_OP_UNARY }, { GGML_UNARY_OP_SILU })) { + ggml_cuda_op_ssm_conv(*cuda_ctx, node, cgraph->nodes[i+1]); + i++; + continue; + } + + if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_UNARY, GGML_OP_MUL }, { GGML_UNARY_OP_SILU }) || + ggml_cuda_can_fuse(cgraph, i, { GGML_OP_UNARY, GGML_OP_MUL }, { GGML_UNARY_OP_SIGMOID }) || + ggml_cuda_can_fuse(cgraph, i, { GGML_OP_UNARY, GGML_OP_MUL }, { GGML_UNARY_OP_SOFTPLUS })) { + ggml_cuda_op_unary_mul(*cuda_ctx, node, cgraph->nodes[i+1]); + i++; + continue; + } + + if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_SCALE, GGML_OP_UNARY, GGML_OP_SCALE }, { GGML_UNARY_OP_TANH })) { + i += 2; + ggml_cuda_op_softcap(*cuda_ctx, cgraph->nodes[i], node); + continue; + } + } +#ifndef NDEBUG + assert(node->buffer->buft == ggml_backend_cuda_buffer_type(cuda_ctx->device)); + for (int j = 0; j < GGML_MAX_SRC; j++) { + if (node->src[j] != nullptr) { + assert(node->src[j]->buffer); + assert(node->src[j]->buffer->buft == ggml_backend_cuda_buffer_type(cuda_ctx->device) || + ggml_backend_buft_is_cuda_split(node->src[j]->buffer->buft) || (integrated && ggml_backend_buft_is_cuda_host(node->src[j]->buffer->buft))); + } + } +#else + GGML_UNUSED(integrated); +#endif // NDEBUG + + bool ok = ggml_cuda_compute_forward(*cuda_ctx, node); + if (!ok) { + GGML_LOG_ERROR("%s: op not supported %s (%s)\n", __func__, node->name, ggml_op_name(node->op)); + } + GGML_ASSERT(ok); + + if (!is_concurrent_event_active) { + try_launch_concurrent_event(node); + } + } + } + +#ifdef USE_CUDA_GRAPH + ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); + if (use_cuda_graph && cuda_graph_update_required) { // End CUDA graph capture + if (graph->graph != nullptr) { + CUDA_CHECK(cudaGraphDestroy(graph->graph)); + graph->graph = nullptr; + } + + CUDA_CHECK(cudaStreamEndCapture(cuda_ctx->stream(), &graph->graph)); + graph_evaluated_or_captured = true; // CUDA graph has been captured + + std::lock_guard lock(ggml_cuda_lock); + if (ggml_cuda_lock_counter.fetch_sub(1, std::memory_order_relaxed) == 1) { + ggml_cuda_lock_cv.notify_all(); + } + } else { + graph_evaluated_or_captured = true; // ggml graph has been directly evaluated + } + } + + if (use_cuda_graph) { + ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); + if (graph->instance == nullptr) { // Create executable graph from captured graph. + CUDA_CHECK(cudaGraphInstantiate(&graph->instance, graph->graph, NULL, NULL, 0)); + } + if (cuda_graph_update_required) { // Update graph executable + ggml_cuda_graph_update_executable(cuda_ctx, graph_key); + } + // Launch graph + CUDA_CHECK(cudaGraphLaunch(graph->instance, cuda_ctx->stream())); +#else + GGML_UNUSED(graph_key); + graph_evaluated_or_captured = true; +#endif // USE_CUDA_GRAPH + } +} + +#ifdef USE_CUDA_GRAPH +static bool ggml_cuda_graph_set_enabled(ggml_backend_cuda_context * cuda_ctx, const void * graph_key) { + ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); + + if (graph->graph == nullptr) { + if (ggml_cuda_info().devices[cuda_ctx->device].cc < GGML_CUDA_CC_AMPERE) { + if (!graph->disable_due_to_gpu_arch) { + GGML_LOG_DEBUG("%s: disabling CUDA graphs due to GPU architecture\n", __func__); + } + graph->disable_due_to_gpu_arch = true; + } + } + + return graph->is_enabled(); +} +#endif // USE_CUDA_GRAPH + +static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, ggml_cgraph * cgraph) { + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context; + + ggml_cuda_set_device(cuda_ctx->device); + + // LIFO-free last evaluation's memoized q8_1 activations. + while (!cuda_ctx->luce_q8_memo.empty()) { + cuda_ctx->luce_q8_memo.pop_back(); + } + + bool use_cuda_graph = false; + bool cuda_graph_update_required = false; + const void * graph_key = nullptr; + +#ifdef USE_CUDA_GRAPH + graph_key = ggml_cuda_graph_get_key(cgraph); + + ggml_cuda_graph_set_enabled(cuda_ctx, graph_key); + + ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); + if (graph->is_enabled()) { + const bool graph_compatible = ggml_cuda_graph_check_compability(cgraph); + if (graph_compatible) { + const bool properties_changed = ggml_cuda_graph_update_required(cuda_ctx, cgraph); + + if (!graph->warmup_complete) { + // Warmup: need at least 2 calls with no property change on the 2nd call + if (!properties_changed) { + graph->warmup_complete = true; + GGML_LOG_DEBUG("%s: CUDA graph warmup complete\n", __func__); + use_cuda_graph = true; + cuda_graph_update_required = true; + } + // else: properties changed or first call - execute directly (use_cuda_graph stays false) + } else { + // Post-warmup: normal CUDA graph operation + if (properties_changed) { + // Properties changed - reset warmup, execute directly until stable again + graph->warmup_complete = false; + GGML_LOG_DEBUG("%s: CUDA graph warmup reset\n", __func__); + } else { + use_cuda_graph = true; + cuda_graph_update_required = graph->instance == nullptr; + } + } + } + } +#endif // USE_CUDA_GRAPH + +#ifdef USE_CUDA_GRAPH + { + static const bool luce_graph_stats = getenv("GGML_CUDA_GRAPH_STATS") != nullptr; + if (luce_graph_stats && graph_key) { + ggml_cuda_graph * g = cuda_ctx->cuda_graph(graph_key); + g->stat_total++; + if (use_cuda_graph && !cuda_graph_update_required) g->stat_replay++; + else if (use_cuda_graph) g->stat_capture++; + else g->stat_eager++; + if (g->stat_total % 200 == 0) { + GGML_LOG_INFO("[cuda-graph-stats] key=%p n_nodes=%d total=%llu replay=%llu capture=%llu eager=%llu enabled=%d\n", + graph_key, cgraph->n_nodes, + (unsigned long long)g->stat_total, (unsigned long long)g->stat_replay, + (unsigned long long)g->stat_capture, (unsigned long long)g->stat_eager, + (int)g->is_enabled()); + } + } + } +#endif // USE_CUDA_GRAPH + + if (use_cuda_graph && cuda_graph_update_required) { + // Start CUDA graph capture + { + std::lock_guard lock(ggml_cuda_lock); + ggml_cuda_lock_counter.fetch_add(1, std::memory_order_relaxed); + } + + CUDA_CHECK(cudaStreamBeginCapture(cuda_ctx->stream(), cudaStreamCaptureModeRelaxed)); + } + + ggml_cuda_graph_evaluate_and_capture(cuda_ctx, cgraph, use_cuda_graph, cuda_graph_update_required, graph_key); + + return GGML_STATUS_SUCCESS; +} + +static void ggml_backend_cuda_event_record(ggml_backend_t backend, ggml_backend_event_t event) { + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context; + + CUDA_CHECK(cudaEventRecord((cudaEvent_t)event->context, cuda_ctx->stream())); +} + +static void ggml_backend_cuda_event_wait(ggml_backend_t backend, ggml_backend_event_t event) { + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context; + + if (ggml_backend_is_cuda(backend)) { + CUDA_CHECK(cudaStreamWaitEvent(cuda_ctx->stream(), (cudaEvent_t)event->context, 0)); + } else { +#if 0 + // untested + auto wait_fn = [](void * user_data) { + ggml_backend_event_t event = (ggml_backend_event_t)user_data; + ggml_backend_event_synchronize(event); + }; + + CUDA_CHECK(cudaLaunchHostFunc(cuda_ctx->stream(), wait_fn, event)); +#endif + GGML_ABORT("fatal error"); + } +} + +static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph * cgraph) { + ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context; + +#ifdef USE_CUDA_GRAPH + const void * graph_key = ggml_cuda_graph_get_key(cgraph); + const bool use_cuda_graph = ggml_cuda_graph_set_enabled(cuda_ctx, graph_key); +#else + const bool use_cuda_graph = false; + GGML_UNUSED(cuda_ctx); + GGML_UNUSED(cgraph); +#endif + + static bool enable_graph_optimization = [] { + const char * env = getenv("GGML_CUDA_GRAPH_OPT"); + return env != nullptr && atoi(env) == 1; + }(); + + if (!enable_graph_optimization) { + return; + } + + ggml_cuda_stream_context & stream_context = cuda_ctx->stream_context(); + stream_context.reset(); + + if (!use_cuda_graph || ggml_backend_cuda_get_device_count() != 1) { + return; + } + + // number of out-degrees for a particular node + std::unordered_map fan_out; + // reverse mapping of node to index in the cgraph + std::unordered_map node_indices; + + const auto & is_noop = [](const ggml_tensor * node) -> bool { + return ggml_is_empty(node) || node->op == GGML_OP_NONE || node->op == GGML_OP_RESHAPE || + node->op == GGML_OP_TRANSPOSE || node->op == GGML_OP_VIEW || node->op == GGML_OP_PERMUTE; + }; + + const auto & depends_on = [](const ggml_tensor * dst, const ggml_tensor * src) -> bool { + for (uint32_t s = 0; s < GGML_MAX_SRC; ++s) { + if (dst->src[s] == src) { + return true; + } + } + // implicit dependency if they view the same tensor + const ggml_tensor * dst2 = dst->view_src ? dst->view_src : dst; + const ggml_tensor * src2 = src->view_src ? src->view_src : src; + if (dst2 == src2) { + return true; + } + return false; + }; + + for (int node_idx = 0; node_idx < cgraph->n_nodes; node_idx++) { + const ggml_tensor * node = cgraph->nodes[node_idx]; + node_indices[node] = node_idx; + + if (is_noop(node)) { + continue; + } + for (int src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) { + const ggml_tensor * src = cgraph->nodes[node_idx]->src[src_idx]; + //TODO: check why nrows > 1 fails + if (node && !is_noop(node) && ggml_nrows(node) <= 1) { + fan_out[src] += 1; + } + } + } + + // Target Q, K, V for concurrency + // this is a more general way to find nodes which can be candidates for concurrency (although it has not been tested for anything else): + // 1. find fan-out (fork) nodes where the same input is used at least N times (in QKV, it would be "attn-norm") + // 2. find the join node, where 2 or more of the outputs are required (in QKV, this would "KQ" or "flash-attn") + // 3. account for all branches from the fork to the join + // 4. To extend lifetimes of the tensors, we interleave the branches (see below for more details) + // 5. save the original cgraph and restore it in graph_compute, to enable fusion within streams + // See discussion: https://github.com/ggml-org/llama.cpp/pull/16991#issuecomment-3522620030 + + const int min_fan_out = 3; + const int max_fan_out = 3; + + // store {fork_idx, join_idx} + std::vector> concurrent_node_ranges; + + for (const auto & [root_node, count] : fan_out) { + if (count >= min_fan_out && count <= max_fan_out) { + const int root_node_idx = node_indices[root_node]; + + // only optimize for attn_norm + // TODO: make this more generic + if (!strstr(root_node->name, "attn_norm")) { + continue; + } + + bool is_part_of_event = false; + for (const auto & [start, end] : concurrent_node_ranges) { + if (root_node_idx >= start && root_node_idx <= end) { + is_part_of_event = true; + } + } + + if (is_part_of_event) { + continue; + } + + std::vector> nodes_per_branch; + for (int i = root_node_idx + 1; i < cgraph->n_nodes; ++i) { + const ggml_tensor * node = cgraph->nodes[i]; + if (!is_noop(node) && depends_on(node, root_node)) { + nodes_per_branch.push_back({ node }); + } + } + + GGML_ASSERT(nodes_per_branch.size() == (size_t) count); + + //find the join point + const ggml_tensor * join_node = nullptr; + + const auto & belongs_to_branch = [&](const ggml_tensor * node, + const std::vector & branch) -> bool { + for (const ggml_tensor * n : branch) { + if (depends_on(node, n)) { + return true; + } + } + return false; + }; + + for (int i = root_node_idx + 1; i < cgraph->n_nodes; ++i) { + const ggml_tensor * curr_node = cgraph->nodes[i]; + + int num_joins = 0; + for (size_t branch_idx = 0; branch_idx < nodes_per_branch.size(); branch_idx++) { + if (belongs_to_branch(curr_node, nodes_per_branch[branch_idx])) { + num_joins++; + } + } + + if (num_joins >= 2) { + join_node = curr_node; + break; + } + + bool found_branch = false; + for (size_t branch_idx = 0; branch_idx < nodes_per_branch.size(); branch_idx++) { + std::vector & branch_vec = nodes_per_branch[branch_idx]; + if (belongs_to_branch(curr_node, branch_vec)) { + //continue accumulating + if (std::find(branch_vec.begin(), branch_vec.end(), curr_node) == branch_vec.end()) { + branch_vec.push_back(curr_node); + } + found_branch = true; + } + } + + if (!found_branch && is_noop(curr_node)) { + // we can put it in any branch because it will be ignored + nodes_per_branch[0].push_back({ curr_node }); + } + } + + if (join_node) { + //Create ggml_cuda_concurrent_event + ggml_cuda_concurrent_event concurrent_event(nodes_per_branch.size()); + concurrent_event.join_node = join_node; + + for (size_t branch_idx = 0; branch_idx < nodes_per_branch.size(); branch_idx++) { + for (const ggml_tensor * n : nodes_per_branch[branch_idx]) { + concurrent_event.stream_mapping[n] = branch_idx + 1; + } + } + + int fork_node_idx = node_indices[root_node]; + int join_node_idx = node_indices[join_node]; + + int current_branch_idx = 0; + int current_node_idx = fork_node_idx + 1; + const int n_branches = nodes_per_branch.size(); + + int total_branch_nodes = 0; + for (std::vector branch_nodes : nodes_per_branch) { + total_branch_nodes += branch_nodes.size(); + } + + // there are other nodes in the middle which are unaccounted for + // usually (cpy) nodes, then ignore this fork + if (join_node_idx - fork_node_idx - 1 != total_branch_nodes) { + GGML_LOG_DEBUG( + "Skipping %s because the number of nodes in the middle is not equal to the total number of " + "branch nodes %d != %d\n", + root_node->name, join_node_idx - fork_node_idx - 1, total_branch_nodes); + continue; + } + + // Save the original order of nodes in this region before interleaving + // This is used later to restore grouping for fusion within streams + concurrent_event.original_order.reserve(total_branch_nodes); + for (int i = fork_node_idx + 1; i < join_node_idx; ++i) { + concurrent_event.original_order.push_back(cgraph->nodes[i]); + } + + std::unordered_map & concurrent_events = cuda_ctx->stream_context().concurrent_events; + GGML_ASSERT(concurrent_events.find(root_node) == concurrent_events.end()); + concurrent_events.emplace(root_node, std::move(concurrent_event)); + GGML_LOG_DEBUG("Adding stream at node %s %p\n", root_node->name, root_node); + concurrent_node_ranges.emplace_back(fork_node_idx, join_node_idx); + + // interleave tensors to extend lifetimes so that ggml graph doesn't recycle them + // example transformation: + // [attn-norm, QMul, QNorm, QRope, KMul, KNorm, KRope, VMul, attn] -> + // [attn-norm, QMul, KMul, VMul, QNorm, VNorm, QRope, KRope, attn] + while (current_node_idx < join_node_idx) { + std::vector & branch_nodes = nodes_per_branch[current_branch_idx]; + + bool has_node = false; + for (std::vector branch_node : nodes_per_branch) { + has_node |= branch_node.size() > 0; + } + + GGML_ASSERT(has_node); + + if (branch_nodes.empty()) { + current_branch_idx = (current_branch_idx + 1) % n_branches; + continue; + } + + cgraph->nodes[current_node_idx] = const_cast(branch_nodes.front()); + current_node_idx++; + branch_nodes.erase(branch_nodes.begin()); + + // append all empty nodes + while (!branch_nodes.empty() && is_noop(branch_nodes.front())) { + cgraph->nodes[current_node_idx] = const_cast(branch_nodes.front()); + current_node_idx++; + branch_nodes.erase(branch_nodes.begin()); + } + + current_branch_idx = (current_branch_idx + 1) % n_branches; + } + } + } + } +} + +static const ggml_backend_i ggml_backend_cuda_interface = { + /* .get_name = */ ggml_backend_cuda_get_name, + /* .free = */ ggml_backend_cuda_free, + /* .set_tensor_async = */ ggml_backend_cuda_set_tensor_async, + /* .get_tensor_async = */ ggml_backend_cuda_get_tensor_async, + /* .get_tensor_2d_async = */ ggml_backend_cuda_set_tensor_2d_async, + /* .set_tensor_2d_async = */ ggml_backend_cuda_get_tensor_2d_async, + /* .cpy_tensor_async = */ ggml_backend_cuda_cpy_tensor_async, + /* .synchronize = */ ggml_backend_cuda_synchronize, + /* .graph_plan_create = */ NULL, + /* .graph_plan_free = */ NULL, + /* .graph_plan_update = */ NULL, + /* .graph_plan_compute = */ NULL, + /* .graph_compute = */ ggml_backend_cuda_graph_compute, + /* .event_record = */ ggml_backend_cuda_event_record, + /* .event_wait = */ ggml_backend_cuda_event_wait, + /* .graph_optimize = */ ggml_backend_cuda_graph_optimize, +}; + +static ggml_guid_t ggml_backend_cuda_guid() { + static ggml_guid guid = { 0x2c, 0xdd, 0xe8, 0x1c, 0x65, 0xb3, 0x65, 0x73, 0x6a, 0x12, 0x88, 0x61, 0x1c, 0xc9, 0xdc, 0x25 }; + return &guid; +} + +bool ggml_backend_is_cuda(ggml_backend_t backend) { + return backend != NULL && ggml_guid_matches(backend->guid, ggml_backend_cuda_guid()); +} + +int ggml_backend_cuda_get_device_count() { + return ggml_cuda_info().device_count; +} + +void ggml_backend_cuda_get_device_description(int device, char * description, size_t description_size) { + cudaDeviceProp prop; + CUDA_CHECK(cudaGetDeviceProperties(&prop, device)); + snprintf(description, description_size, "%s", prop.name); +} + +void ggml_backend_cuda_get_device_memory(int device, size_t * free, size_t * total) { + ggml_cuda_set_device(device); + + CUDA_CHECK(cudaMemGetInfo(free, total)); +} + +bool ggml_backend_cuda_register_host_buffer(void * buffer, size_t size) { + if (getenv("GGML_CUDA_REGISTER_HOST") == nullptr) { + return false; + } + +#if CUDART_VERSION >= 11010 || defined(GGML_USE_MUSA) || defined(GGML_USE_HIP) + cudaError_t err = cudaHostRegister(buffer, size, cudaHostRegisterPortable | cudaHostRegisterReadOnly); + if (err != cudaSuccess) { + // clear the error + (void)cudaGetLastError(); + + GGML_LOG_DEBUG("%s: failed to register %.2f MiB of pinned memory: %s\n", __func__, + size / 1024.0 / 1024.0, cudaGetErrorString(err)); + return false; + } + return true; +#else + GGML_UNUSED(buffer); + GGML_UNUSED(size); + return false; +#endif // CUDART_VERSION >= 11010 || defined(GGML_USE_MUSA) +} + +void ggml_backend_cuda_unregister_host_buffer(void * buffer) { + if (getenv("GGML_CUDA_REGISTER_HOST") == nullptr) { + return; + } + + cudaError_t err = cudaHostUnregister(buffer); + if (err != cudaSuccess) { + // clear the error + (void)cudaGetLastError(); + } +} + + +// backend device + +struct ggml_backend_cuda_device_context { + int device; + std::string name; + std::string description; + std::string pci_bus_id; + int op_offload_min_batch_size; +}; + +static const char * ggml_backend_cuda_device_get_name(ggml_backend_dev_t dev) { + ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *)dev->context; + return ctx->name.c_str(); +} + +static const char * ggml_backend_cuda_device_get_description(ggml_backend_dev_t dev) { + ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *)dev->context; + return ctx->description.c_str(); +} + +#if defined(__linux__) +// Helper function to get available memory from /proc/meminfo for UMA systems +static bool ggml_backend_cuda_get_available_uma_memory(long * available_memory_kb, long * free_swap_kb) { + FILE * meminfo_file = nullptr; + // 2KB buffer for reading /proc/meminfo since it does not report size info, should be enough + const size_t BUFFER_SIZE = 2048; + auto file_buffer = std::make_unique(BUFFER_SIZE); + size_t bytes_read = 0; + long huge_tlb_total_pages = -1; + long huge_tlb_free_pages = -1; + long huge_tlb_page_size = -1; + + if (available_memory_kb == nullptr || free_swap_kb == nullptr) { + return false; + } + + meminfo_file = fopen("/proc/meminfo", "r"); + if (meminfo_file == nullptr) { + GGML_LOG_ERROR("%s: failed to open /proc/meminfo\n", __func__); + return false; + } + + // Read file into buffer + bytes_read = fread(file_buffer.get(), 1, BUFFER_SIZE - 1, meminfo_file); + fclose(meminfo_file); + + if (bytes_read == 0) { + GGML_LOG_ERROR("%s: failed to read from /proc/meminfo\n", __func__); + return false; + } + file_buffer[bytes_read] = '\0'; + + *available_memory_kb = -1; + *free_swap_kb = -1; + + // Parse the file buffer line by line + char * line = file_buffer.get(); + char * line_next; + while (line < file_buffer.get() + bytes_read) { + // Find the end of the current line + line_next = strchr(line, '\n'); + if (line_next != nullptr) { + *line_next = '\0'; + line_next++; + } else { + line_next = file_buffer.get() + bytes_read; + } + + long value; + if (sscanf(line, "MemAvailable: %ld kB", &value) == 1) { + *available_memory_kb = value; + } else if (sscanf(line, "SwapFree: %ld kB", &value) == 1) { + *free_swap_kb = value; + } else if (sscanf(line, "HugePages_Total: %ld", &value) == 1) { + huge_tlb_total_pages = value; + } else if (sscanf(line, "HugePages_Free: %ld", &value) == 1) { + huge_tlb_free_pages = value; + } else if (sscanf(line, "Hugepagesize: %ld kB", &value) == 1) { + huge_tlb_page_size = value; + } + + line = line_next; + } + + if (huge_tlb_total_pages != 0 && huge_tlb_total_pages != -1) { + *available_memory_kb = huge_tlb_free_pages * huge_tlb_page_size; + + // Hugetlbfs pages are not swappable. + *free_swap_kb = 0; + } + + GGML_LOG_DEBUG("%s: final available_memory_kb: %ld\n", __func__, *available_memory_kb); + return true; +} +#endif // defined(__linux__) + +static void ggml_backend_cuda_device_get_memory(ggml_backend_dev_t dev, size_t * free, size_t * total) { + ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *)dev->context; + ggml_cuda_set_device(ctx->device); + CUDA_CHECK(cudaMemGetInfo(free, total)); + +// ref: https://github.com/ggml-org/llama.cpp/pull/17368 +#if defined(__linux__) + // Check if this is a UMA (Unified Memory Architecture) system + cudaDeviceProp prop; + CUDA_CHECK(cudaGetDeviceProperties(&prop, ctx->device)); + + // Check if UMA is explicitly enabled via environment variable + bool uma_env = getenv("GGML_CUDA_ENABLE_UNIFIED_MEMORY") != nullptr; + bool is_uma = prop.integrated > 0 || uma_env; + + if (is_uma) { + // For UMA systems (like DGX Spark), use system memory info + long available_memory_kb = 0; + long free_swap_kb = 0; + + if (ggml_backend_cuda_get_available_uma_memory(&available_memory_kb, &free_swap_kb) && available_memory_kb > 0) { + *free = (size_t)available_memory_kb * 1024; + } else { + GGML_LOG_ERROR("%s: /proc/meminfo reading failed, using cudaMemGetInfo\n", __func__); + } + } +#endif // defined(__linux__) + +} + +static enum ggml_backend_dev_type ggml_backend_cuda_device_get_type(ggml_backend_dev_t dev) { + GGML_UNUSED(dev); + return GGML_BACKEND_DEVICE_TYPE_GPU; +} + +static void ggml_backend_cuda_device_get_props(ggml_backend_dev_t dev, ggml_backend_dev_props * props) { + ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *)dev->context; + + props->name = ggml_backend_cuda_device_get_name(dev); + props->description = ggml_backend_cuda_device_get_description(dev); + props->type = ggml_backend_cuda_device_get_type(dev); + props->device_id = ctx->pci_bus_id.empty() ? nullptr : ctx->pci_bus_id.c_str(); + ggml_backend_cuda_device_get_memory(dev, &props->memory_free, &props->memory_total); + + bool host_buffer = getenv("GGML_CUDA_NO_PINNED") == nullptr; +#ifdef GGML_CUDA_NO_PEER_COPY + bool events = false; +#else + bool events = true; +#endif + + props->caps = { + /* .async = */ true, + /* .host_buffer = */ host_buffer, + /* .buffer_from_host_ptr = */ false, + /* .events = */ events, + }; +} + +static ggml_backend_t ggml_backend_cuda_device_init_backend(ggml_backend_dev_t dev, const char * params) { + GGML_UNUSED(params); + ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *)dev->context; + return ggml_backend_cuda_init(ctx->device); +} + +static ggml_backend_buffer_type_t ggml_backend_cuda_device_get_buffer_type(ggml_backend_dev_t dev) { + ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *)dev->context; + return ggml_backend_cuda_buffer_type(ctx->device); +} + +static ggml_backend_buffer_type_t ggml_backend_cuda_device_get_host_buffer_type(ggml_backend_dev_t dev) { + GGML_UNUSED(dev); + return ggml_backend_cuda_host_buffer_type(); +} + +// TODO: move these functions here +static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) { + ggml_backend_cuda_device_context * dev_ctx = (ggml_backend_cuda_device_context *) dev->context; + + // split buffers can only be used with GGML_OP_MUL_MAT + if (op->op != GGML_OP_MUL_MAT) { + for (int i = 0; i < GGML_MAX_SRC; i++) { + if (op->src[i] && op->src[i]->buffer && ggml_backend_buft_is_cuda_split(op->src[i]->buffer->buft)) { + return false; + } + } + } + + // check if all the sources are allocated on this device + for (int i = 0; i < GGML_MAX_SRC; i++) { + if (op->src[i] && op->src[i]->buffer && ggml_backend_buft_is_cuda(op->src[i]->buffer->buft)) { + ggml_backend_cuda_buffer_type_context * buft_ctx = (ggml_backend_cuda_buffer_type_context *)op->src[i]->buffer->buft->context; + if (buft_ctx->device != dev_ctx->device) { + return false; + } + } + } + + switch (op->op) { + case GGML_OP_UNARY: + switch (ggml_get_unary_op(op)) { + case GGML_UNARY_OP_ABS: + case GGML_UNARY_OP_SGN: + case GGML_UNARY_OP_NEG: + case GGML_UNARY_OP_STEP: + case GGML_UNARY_OP_GELU: + case GGML_UNARY_OP_SILU: + case GGML_UNARY_OP_RELU: + case GGML_UNARY_OP_SIGMOID: + case GGML_UNARY_OP_HARDSIGMOID: + case GGML_UNARY_OP_HARDSWISH: + case GGML_UNARY_OP_GELU_ERF: + case GGML_UNARY_OP_GELU_QUICK: + case GGML_UNARY_OP_TANH: + case GGML_UNARY_OP_EXP: + case GGML_UNARY_OP_EXPM1: + case GGML_UNARY_OP_SOFTPLUS: + case GGML_UNARY_OP_ELU: + case GGML_UNARY_OP_XIELU: + case GGML_UNARY_OP_FLOOR: + case GGML_UNARY_OP_CEIL: + case GGML_UNARY_OP_ROUND: + case GGML_UNARY_OP_TRUNC: + // TODO: should become: + //return ggml_is_contiguous_rows(op->src[0]); + return ggml_is_contiguous(op->src[0]); + default: + return false; + } + break; + case GGML_OP_GLU: + switch (ggml_get_glu_op(op)) { + case GGML_GLU_OP_REGLU: + case GGML_GLU_OP_GEGLU: + case GGML_GLU_OP_SWIGLU: + case GGML_GLU_OP_SWIGLU_OAI: + case GGML_GLU_OP_SWIGLU_DS4: + case GGML_GLU_OP_GEGLU_ERF: + case GGML_GLU_OP_GEGLU_QUICK: + return ggml_is_contiguous_1(op->src[0]); + default: + return false; + } + break; + case GGML_OP_TURBO_WHT: + // Only requires dim0 contiguous (nb[0] == sizeof(float)); + // the kernel handles strided dim1/dim2 via separate src/dst strides. + return op->src[0]->nb[0] == ggml_type_size(op->src[0]->type); + case GGML_OP_MOE_FUSED: + return true; + case GGML_OP_MUL_MAT: + case GGML_OP_MUL_MAT_ID: + { + struct ggml_tensor * a = op->src[0]; + struct ggml_tensor * b = op->src[1]; + if (a->buffer && ggml_backend_buft_is_cuda_split(a->buffer->buft)) { + if (a->ne[2] > 1 || a->ne[3] > 1) { + return false; + } + // for small weight matrices the active device can end up without any rows, don't use row split in those cases + // this avoids some edge cases (and the performance would not be good anyways) + ggml_backend_cuda_split_buffer_type_context * buft_ctx = (ggml_backend_cuda_split_buffer_type_context *) a->buffer->buft->context; + int64_t row_low; + int64_t row_high; + get_row_split(&row_low, &row_high, a, buft_ctx->tensor_split, dev_ctx->device); + if (row_low == row_high) { + return false; + } + } + if (b->type == GGML_TYPE_F16 && a->type != GGML_TYPE_F16) { + return false; + } +#ifdef GGML_USE_MUSA + const int cc = ggml_cuda_info().devices[dev_ctx->device].cc; + if (b->ne[2]*b->ne[3] > 1 && !ggml_is_transposed(a) && !ggml_is_transposed(b)) { + if (GGML_CUDA_CC_IS_QY1(cc) && op->op == GGML_OP_MUL_MAT && + a->type == GGML_TYPE_F16 && b->type == GGML_TYPE_F16) { + return false; + } + if (GGML_CUDA_CC_IS_QY2(cc) && op->op == GGML_OP_MUL_MAT_ID && + a->type == GGML_TYPE_Q2_K && b->type == GGML_TYPE_F32) { + return false; + } + } +#endif // GGML_USE_MUSA + switch (a->type) { + case GGML_TYPE_F32: + case GGML_TYPE_F16: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_MXFP4: + case GGML_TYPE_NVFP4: + case GGML_TYPE_Q4_0_ROCMFP4: + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + case GGML_TYPE_Q2_0_ROCMFP2: + case GGML_TYPE_Q3_0_ROCMFPX: + case GGML_TYPE_Q6_0_ROCMFPX: + case GGML_TYPE_Q8_0_ROCMFPX: + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_Q8_K: + case GGML_TYPE_IQ1_M: + case GGML_TYPE_IQ1_S: + case GGML_TYPE_IQ2_S: + case GGML_TYPE_IQ2_XS: + case GGML_TYPE_IQ2_XXS: + case GGML_TYPE_IQ3_S: + case GGML_TYPE_IQ3_XXS: + case GGML_TYPE_IQ4_NL: + case GGML_TYPE_IQ4_XS: + case GGML_TYPE_BF16: + return true; + default: + return false; + } + } break; + case GGML_OP_OUT_PROD: + return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32; + case GGML_OP_GET_ROWS: + { + switch (op->src[0]->type) { + case GGML_TYPE_F16: + case GGML_TYPE_F32: + case GGML_TYPE_BF16: + case GGML_TYPE_I32: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q4_0_ROCMFP4: + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + case GGML_TYPE_Q2_0_ROCMFP2: + case GGML_TYPE_Q3_0_ROCMFPX: + case GGML_TYPE_Q6_0_ROCMFPX: + case GGML_TYPE_Q8_0_ROCMFPX: + case GGML_TYPE_TQ3_0: + return true; + default: + return false; + } + } break; + case GGML_OP_GET_ROWS_BACK: + { + return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && op->ne[2] == 1 && op->ne[3] == 1; + } break; + case GGML_OP_SET_ROWS: + { + return (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_BF16 || + op->type == GGML_TYPE_Q4_0 || op->type == GGML_TYPE_Q4_1 || op->type == GGML_TYPE_Q5_0 || + op->type == GGML_TYPE_Q5_1 || op->type == GGML_TYPE_Q8_0 || op->type == GGML_TYPE_IQ4_NL || + op->type == GGML_TYPE_TQ3_0) && + op->src[0]->type == GGML_TYPE_F32 && + (op->src[1]->type == GGML_TYPE_I64 || op->src[1]->type == GGML_TYPE_I32); + } break; + case GGML_OP_SET: + { + const ggml_type t = op->type; + return (t == GGML_TYPE_F32 || t == GGML_TYPE_I32) && + t == op->src[0]->type && + t == op->src[1]->type; + } break; + case GGML_OP_CPY: + { + ggml_type src0_type = op->src[0]->type; + ggml_type src1_type = op->src[1]->type; + if ((src0_type == GGML_TYPE_F32 || src0_type == GGML_TYPE_BF16 || src0_type == GGML_TYPE_F16) && + (src1_type == GGML_TYPE_F32 || src1_type == GGML_TYPE_BF16 || src1_type == GGML_TYPE_F16) + ) { + return true; + } + if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_Q8_0) { + return true; + } + if (src0_type == GGML_TYPE_Q8_0 && src1_type == GGML_TYPE_F32) { + return true; + } + if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_Q4_0) { + return true; + } + if (src0_type == GGML_TYPE_Q4_0 && src1_type == GGML_TYPE_F32) { + return true; + } + if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_Q4_1) { + return true; + } + if (src0_type == GGML_TYPE_Q4_1 && src1_type == GGML_TYPE_F32) { + return true; + } + if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_Q5_0) { + return true; + } + if (src0_type == GGML_TYPE_Q5_0 && src1_type == GGML_TYPE_F32) { + return true; + } + if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_Q5_1) { + return true; + } + if (src0_type == GGML_TYPE_Q5_1 && src1_type == GGML_TYPE_F32) { + return true; + } + if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_IQ4_NL) { + return true; + } + if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_TQ3_0) { + return true; + } + if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_I32) { + return true; + } + if (src0_type == GGML_TYPE_I32 && src1_type == GGML_TYPE_F32) { + return true; + } + if (src0_type == GGML_TYPE_I32 && src1_type == GGML_TYPE_I32) { + return true; + } + if (src0_type == src1_type && ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1])) { + return true; + } + return false; + } break; + case GGML_OP_DUP: + { + ggml_type src0_type = op->src[0]->type; + return src0_type != GGML_TYPE_I32 && src0_type != GGML_TYPE_I16; + } break; + case GGML_OP_ARGMAX: + case GGML_OP_COUNT_EQUAL: + { + return true; + } break; + case GGML_OP_REPEAT: + { + // the CUDA REPEAT path only implements F32/F16; other types assert at runtime + ggml_type src0_type = op->src[0]->type; + return src0_type == GGML_TYPE_F32 || src0_type == GGML_TYPE_F16; + } break; + case GGML_OP_REPEAT_BACK: + return op->type == GGML_TYPE_F32 && (op->src[0]->ne[2]*op->src[0]->ne[3]) <= (1 << 15); + case GGML_OP_CONCAT: + { + ggml_type src0_type = op->src[0]->type; + return src0_type == GGML_TYPE_F32 || + src0_type == GGML_TYPE_F16 || + src0_type == GGML_TYPE_BF16 || + src0_type == GGML_TYPE_I8; + } break; + case GGML_OP_CONV_TRANSPOSE_1D: + { + ggml_type src0_type = op->src[0]->type; + ggml_type src1_type = op->src[1]->type; + if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_F32) { + return true; + } + return false; + } break; + case GGML_OP_SILU_BACK: + return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32; + break; + case GGML_OP_NORM: + case GGML_OP_RMS_NORM: + case GGML_OP_L2_NORM: + return true; + case GGML_OP_RMS_NORM_BACK: + return ggml_is_contiguous(op->src[0]); + break; + case GGML_OP_NONE: + case GGML_OP_RESHAPE: + case GGML_OP_VIEW: + case GGML_OP_PERMUTE: + case GGML_OP_TRANSPOSE: + case GGML_OP_ADD: + case GGML_OP_ADD_ID: + case GGML_OP_ADD1: + case GGML_OP_SUB: + case GGML_OP_MUL: + case GGML_OP_DIV: + case GGML_OP_SCALE: + case GGML_OP_SQR: + case GGML_OP_SQRT: + case GGML_OP_SIN: + case GGML_OP_COS: + case GGML_OP_CLAMP: + case GGML_OP_LOG: + return true; + case GGML_OP_SSM_SCAN: { + if (op->src[3]->ne[0] == 1) { + // Mamba2 + // (kernel only supports (d_state == 128 || d_state == 256) && d_head % 16 == 0) + return (op->src[0]->ne[0] == 128 || op->src[0]->ne[0] == 256) && op->src[0]->ne[1] % 16 == 0; + } else { + // Mamba + // (kernel only supports d_state == 16, d_head == 1, n_head % 128 == 0, n_group == 1) + return op->src[0]->ne[0] == 16 && op->src[0]->ne[1] == 1 && op->src[0]->ne[2] % 128 == 0 && op->src[4]->ne[1] == 1; + } + } + case GGML_OP_SSM_CONV: { + // assumes d_inner % threads == 0 + return op->src[0]->ne[1] % 128 == 0; + } + case GGML_OP_CONT: + return true; + case GGML_OP_DIAG_MASK_INF: + return true; + case GGML_OP_SOFT_MAX: + return true; + case GGML_OP_SOFT_MAX_BACK: { + float max_bias = 0.0f; + memcpy(&max_bias, (const float *) op->op_params + 1, sizeof(float)); + return max_bias == 0.0f; + } + case GGML_OP_ROLL: + if(op->src[0]->type == GGML_TYPE_F32) { + return true; + } + return false; + case GGML_OP_ROPE: + case GGML_OP_ROPE_BACK: { + return op->src[0]->nb[0] == ggml_type_size(op->src[0]->type) && ggml_is_contiguous_2(op->src[0]); + } + case GGML_OP_IM2COL: + case GGML_OP_IM2COL_3D: + case GGML_OP_CONV_2D: + case GGML_OP_CONV_2D_DW: + case GGML_OP_CONV_TRANSPOSE_2D: + case GGML_OP_POOL_2D: + return true; + case GGML_OP_ACC: + // TODO: extend support like so: + //return ggml_is_contiguous_rows(op->src[0]) && ggml_is_contiguous_rows(op->src[1]); + return ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1]); + case GGML_OP_SUM: + return ggml_is_contiguous_rows(op->src[0]); + case GGML_OP_TOP_K: + case GGML_OP_ARGSORT: +#ifndef GGML_CUDA_USE_CUB + return op->src[0]->ne[0] <= 1024; +#else + return true; +#endif + case GGML_OP_SUM_ROWS: + case GGML_OP_MEAN: + case GGML_OP_GROUP_NORM: + return ggml_is_contiguous(op->src[0]); + case GGML_OP_PAD: + return true; + case GGML_OP_UPSCALE: + case GGML_OP_PAD_REFLECT_1D: + case GGML_OP_ARANGE: + case GGML_OP_TIMESTEP_EMBEDDING: + case GGML_OP_LEAKY_RELU: + case GGML_OP_RWKV_WKV6: + case GGML_OP_GATED_LINEAR_ATTN: + case GGML_OP_RWKV_WKV7: + return true; + case GGML_OP_GATED_DELTA_NET: + //TODO: enable once MUSA compiler is solved https://github.com/ggml-org/llama.cpp/pull/19504#issuecomment-4018634327 +#ifdef GGML_USE_MUSA + return false; +#else + return true; +#endif // GGML_USE_MUSA + case GGML_OP_FLASH_ATTN_EXT: + return ggml_cuda_flash_attn_ext_supported(dev_ctx->device, op); + case GGML_OP_FLASH_ATTN_SPARSE: + return true; // Always supported on CUDA + case GGML_OP_CROSS_ENTROPY_LOSS: + case GGML_OP_CROSS_ENTROPY_LOSS_BACK: + case GGML_OP_OPT_STEP_ADAMW: + case GGML_OP_OPT_STEP_SGD: + case GGML_OP_FILL: + case GGML_OP_CUMSUM: + case GGML_OP_TRI: + case GGML_OP_DIAG: + case GGML_OP_SOLVE_TRI: + return true; + + default: + return false; + } +} + +static bool ggml_backend_cuda_device_supports_buft(ggml_backend_dev_t dev, ggml_backend_buffer_type_t buft) { + ggml_backend_cuda_device_context * dev_ctx = (ggml_backend_cuda_device_context *) dev->context; + const bool integrated = ggml_cuda_info().devices[dev_ctx->device].integrated; + return (((ggml_backend_buft_is_cuda(buft) || ggml_backend_buft_is_cuda_split(buft)) && buft->device == dev) || (integrated && ggml_backend_buft_is_cuda_host(buft))); +} + +static int64_t get_op_batch_size(const ggml_tensor * op) { + switch (op->op) { + case GGML_OP_GET_ROWS: + return 0; + case GGML_OP_MUL_MAT: + return op->ne[1]; + case GGML_OP_MUL_MAT_ID: + case GGML_OP_ROPE: + case GGML_OP_ROPE_BACK: + return op->ne[2]; + default: + return ggml_nrows(op); + } +} + +static bool ggml_backend_cuda_device_offload_op(ggml_backend_dev_t dev, const ggml_tensor * op) { + ggml_backend_cuda_device_context * dev_ctx = (ggml_backend_cuda_device_context *) dev->context; + + return get_op_batch_size(op) >= dev_ctx->op_offload_min_batch_size; +} + +static ggml_backend_event_t ggml_backend_cuda_device_event_new(ggml_backend_dev_t dev) { +#ifdef GGML_CUDA_NO_PEER_COPY + return nullptr; +#else + ggml_backend_cuda_device_context * dev_ctx = (ggml_backend_cuda_device_context *)dev->context; + + ggml_cuda_set_device(dev_ctx->device); + + cudaEvent_t event; + CUDA_CHECK(cudaEventCreateWithFlags(&event, cudaEventDisableTiming)); + + return new ggml_backend_event { + /* .device = */ dev, + /* .context = */ event, + }; +#endif +} + +static void ggml_backend_cuda_device_event_free(ggml_backend_dev_t dev, ggml_backend_event_t event) { + GGML_UNUSED(dev); + + CUDA_CHECK(cudaEventDestroy((cudaEvent_t)event->context)); + delete event; +} + +static void ggml_backend_cuda_device_event_synchronize(ggml_backend_dev_t dev, ggml_backend_event_t event) { + GGML_UNUSED(dev); + CUDA_CHECK(cudaEventSynchronize((cudaEvent_t)event->context)); +} + +static const ggml_backend_device_i ggml_backend_cuda_device_interface = { + /* .get_name = */ ggml_backend_cuda_device_get_name, + /* .get_description = */ ggml_backend_cuda_device_get_description, + /* .get_memory = */ ggml_backend_cuda_device_get_memory, + /* .get_type = */ ggml_backend_cuda_device_get_type, + /* .get_props = */ ggml_backend_cuda_device_get_props, + /* .init_backend = */ ggml_backend_cuda_device_init_backend, + /* .get_buffer_type = */ ggml_backend_cuda_device_get_buffer_type, + /* .get_host_buffer_type = */ ggml_backend_cuda_device_get_host_buffer_type, + /* .buffer_from_host_ptr = */ NULL, + /* .supports_op = */ ggml_backend_cuda_device_supports_op, + /* .supports_buft = */ ggml_backend_cuda_device_supports_buft, + /* .offload_op = */ ggml_backend_cuda_device_offload_op, + /* .event_new = */ ggml_backend_cuda_device_event_new, + /* .event_free = */ ggml_backend_cuda_device_event_free, + /* .event_synchronize = */ ggml_backend_cuda_device_event_synchronize, +}; + +// backend reg + +struct ggml_backend_cuda_reg_context { + std::vector devices; +}; + +static const char * ggml_backend_cuda_reg_get_name(ggml_backend_reg_t reg) { + GGML_UNUSED(reg); + return GGML_CUDA_NAME; +} + +static size_t ggml_backend_cuda_reg_get_device_count(ggml_backend_reg_t reg) { + ggml_backend_cuda_reg_context * ctx = (ggml_backend_cuda_reg_context *)reg->context; + return ctx->devices.size(); +} + +static ggml_backend_dev_t ggml_backend_cuda_reg_get_device(ggml_backend_reg_t reg, size_t index) { + ggml_backend_cuda_reg_context * ctx = (ggml_backend_cuda_reg_context *)reg->context; + GGML_ASSERT(index < ctx->devices.size()); + return ctx->devices[index]; +} + +static ggml_backend_feature * ggml_backend_cuda_get_features(ggml_backend_reg_t reg) { + static std::vector features = []() { + std::vector features; + #define _STRINGIFY(...) #__VA_ARGS__ + #define STRINGIFY(...) _STRINGIFY(__VA_ARGS__) + + #ifdef __CUDA_ARCH_LIST__ + features.push_back({ "ARCHS", STRINGIFY(__CUDA_ARCH_LIST__) }); + #endif + + #ifdef GGML_CUDA_FORCE_MMQ + features.push_back({ "FORCE_MMQ", "1" }); + #endif + + #ifdef GGML_CUDA_FORCE_CUBLAS + features.push_back({ "FORCE_CUBLAS", "1" }); + #endif + + #ifndef GGML_USE_VMM + features.push_back({ "NO_VMM", "1" }); + #endif + + #ifdef GGML_CUDA_NO_PEER_COPY + features.push_back({ "NO_PEER_COPY", "1" }); + #endif + + #ifdef GGML_CUDA_USE_GRAPHS + features.push_back({ "USE_GRAPHS", "1" }); + #endif + + #ifdef GGML_CUDA_PEER_MAX_BATCH_SIZE + features.push_back({ "PEER_MAX_BATCH_SIZE", STRINGIFY(GGML_CUDA_PEER_MAX_BATCH_SIZE) }); + #endif + + #ifdef GGML_CUDA_FA_ALL_QUANTS + features.push_back({ "FA_ALL_QUANTS", "1" }); + #endif + + { + const auto & info = ggml_cuda_info(); + for (int id = 0; id < info.device_count; ++id) { + if (blackwell_mma_available(info.devices[id].cc)) { + features.push_back({ "BLACKWELL_NATIVE_FP4", "1"}); + break; + } + } + } + + #undef _STRINGIFY + #undef STRINGIFY + + features.push_back({ nullptr, nullptr }); + + return features; + }(); + + return features.data(); + + GGML_UNUSED(reg); +} + +static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, const char * name) { + GGML_UNUSED(reg); + if (strcmp(name, "ggml_backend_allreduce_tensor") == 0) { + return (void *)ggml_backend_cuda_allreduce_tensor; + } + if (strcmp(name, "ggml_backend_split_buffer_type") == 0) { + return (void *)ggml_backend_cuda_split_buffer_type; + } + if (strcmp(name, "ggml_backend_register_host_buffer") == 0) { + return (void *)ggml_backend_cuda_register_host_buffer; + } + if (strcmp(name, "ggml_backend_unregister_host_buffer") == 0) { + return (void *)ggml_backend_cuda_unregister_host_buffer; + } + if (strcmp(name, "ggml_backend_get_features") == 0) { + return (void *)ggml_backend_cuda_get_features; + } + return nullptr; +} + +static const ggml_backend_reg_i ggml_backend_cuda_reg_interface = { + /* .get_name = */ ggml_backend_cuda_reg_get_name, + /* .get_device_count = */ ggml_backend_cuda_reg_get_device_count, + /* .get_device = */ ggml_backend_cuda_reg_get_device, + /* .get_proc_address = */ ggml_backend_cuda_reg_get_proc_address, +}; + +// backend registry +ggml_backend_reg_t ggml_backend_cuda_reg() { + static ggml_backend_reg reg; + static bool initialized = false; + + { + static std::mutex mutex; + std::lock_guard lock(mutex); + if (!initialized) { + ggml_backend_cuda_reg_context * ctx = new ggml_backend_cuda_reg_context; + const int min_batch_size = getenv("GGML_OP_OFFLOAD_MIN_BATCH") ? atoi(getenv("GGML_OP_OFFLOAD_MIN_BATCH")) : 32; + + for (int i = 0; i < ggml_cuda_info().device_count; i++) { + ggml_backend_cuda_device_context * dev_ctx = new ggml_backend_cuda_device_context; + dev_ctx->device = i; + dev_ctx->name = GGML_CUDA_NAME + std::to_string(i); + + cudaDeviceProp prop; + CUDA_CHECK(cudaGetDeviceProperties(&prop, i)); + dev_ctx->description = prop.name; + + char pci_bus_id[16] = {}; + snprintf(pci_bus_id, sizeof(pci_bus_id), "%04x:%02x:%02x.0", prop.pciDomainID, prop.pciBusID, prop.pciDeviceID); + dev_ctx->pci_bus_id = pci_bus_id; + dev_ctx->op_offload_min_batch_size = min_batch_size; + + ggml_backend_dev_t dev = new ggml_backend_device { + /* .iface = */ ggml_backend_cuda_device_interface, + /* .reg = */ ®, + /* .context = */ dev_ctx + }; + ctx->devices.push_back(dev); + } + + reg = ggml_backend_reg { + /* .api_version = */ GGML_BACKEND_API_VERSION, + /* .iface = */ ggml_backend_cuda_reg_interface, + /* .context = */ ctx + }; + } + + initialized = true; + } + + return ® +} + +ggml_backend_t ggml_backend_cuda_init(int device) { + if (device < 0 || device >= ggml_backend_cuda_get_device_count()) { + GGML_LOG_ERROR("%s: invalid device %d\n", __func__, device); + return nullptr; + } + + ggml_backend_cuda_context * ctx = new ggml_backend_cuda_context(device); + if (ctx == nullptr) { + GGML_LOG_ERROR("%s: failed to allocate context\n", __func__); + return nullptr; + } + + ggml_backend_t cuda_backend = new ggml_backend { + /* .guid = */ ggml_backend_cuda_guid(), + /* .iface = */ ggml_backend_cuda_interface, + /* .device = */ ggml_backend_reg_dev_get(ggml_backend_cuda_reg(), device), + /* .context = */ ctx, + }; + + return cuda_backend; +} + +GGML_BACKEND_DL_IMPL(ggml_backend_cuda_reg) diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 3b5c6bcafa23..ef22633b3552 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -1124,9 +1124,13 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = { "GLU", "TURBO_WHT", + + "MOE_FUSED", + + "DS4_HC", }; -static_assert(GGML_OP_COUNT == 99, "GGML_OP_COUNT != 99"); +static_assert(GGML_OP_COUNT == 100, "GGML_OP_COUNT != 100"); static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "none", @@ -1237,9 +1241,13 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { "glu(x)", "turbo_wht(a)", + + "moe_fused(x)", + + "ds4_hc(x)", }; -static_assert(GGML_OP_COUNT == 99, "GGML_OP_COUNT != 99"); +static_assert(GGML_OP_COUNT == 100, "GGML_OP_COUNT != 100"); static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); @@ -8038,3 +8046,93 @@ struct ggml_tensor * ggml_laguna_moe_combine( return result; } + +struct ggml_tensor * ggml_ds4_hc_pre( + struct ggml_context * ctx, + struct ggml_tensor * mix, + struct ggml_tensor * base, + struct ggml_tensor * hc_state, + int n_hc, + int sinkhorn_iters, + float pre_scale, + float post_scale, + float comb_scale) { + GGML_ASSERT(mix->type == GGML_TYPE_F32); + GGML_ASSERT(base->type == GGML_TYPE_F32); + GGML_ASSERT(hc_state->type == GGML_TYPE_F32); + GGML_ASSERT(n_hc > 0 && n_hc <= 8); + const int64_t mix_dim = 2*(int64_t)n_hc + (int64_t)n_hc*n_hc; + GGML_ASSERT(ggml_nelements(mix) == mix_dim); + GGML_ASSERT(ggml_nelements(base) >= mix_dim); + GGML_ASSERT(ggml_nelements(hc_state) % n_hc == 0); + const int64_t n_embd = ggml_nelements(hc_state) / n_hc; + + struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embd + mix_dim); + result->op = GGML_OP_DS4_HC; + result->src[0] = mix; + result->src[1] = base; + result->src[2] = hc_state; + ggml_set_op_params_i32(result, 0, 0); + ggml_set_op_params_i32(result, 1, (int32_t) n_embd); + ggml_set_op_params_i32(result, 2, (int32_t) n_hc); + ggml_set_op_params_i32(result, 3, (int32_t) sinkhorn_iters); + ggml_set_op_params_f32(result, 4, pre_scale); + ggml_set_op_params_f32(result, 5, post_scale); + ggml_set_op_params_f32(result, 6, comb_scale); + return result; +} + +struct ggml_tensor * ggml_ds4_hc_post( + struct ggml_context * ctx, + struct ggml_tensor * residual_hc, + struct ggml_tensor * block_out, + struct ggml_tensor * split, + int n_hc) { + GGML_ASSERT(residual_hc->type == GGML_TYPE_F32); + GGML_ASSERT(block_out->type == GGML_TYPE_F32); + GGML_ASSERT(split->type == GGML_TYPE_F32); + GGML_ASSERT(n_hc > 0 && n_hc <= 8); + const int64_t mix_dim = 2*(int64_t)n_hc + (int64_t)n_hc*n_hc; + GGML_ASSERT(ggml_nelements(split) == mix_dim); + GGML_ASSERT(ggml_nelements(residual_hc) % n_hc == 0); + const int64_t n_embd = ggml_nelements(residual_hc) / n_hc; + GGML_ASSERT(ggml_nelements(block_out) == n_embd); + + struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, (int64_t) n_embd * n_hc); + result->op = GGML_OP_DS4_HC; + result->src[0] = residual_hc; + result->src[1] = block_out; + result->src[2] = split; + ggml_set_op_params_i32(result, 0, 1); + ggml_set_op_params_i32(result, 1, (int32_t) n_embd); + ggml_set_op_params_i32(result, 2, (int32_t) n_hc); + return result; +} + +struct ggml_tensor * ggml_ds4_hc_out( + struct ggml_context * ctx, + struct ggml_tensor * mix, + struct ggml_tensor * base, + struct ggml_tensor * hc_state, + int n_hc, + float pre_scale) { + GGML_ASSERT(mix->type == GGML_TYPE_F32); + GGML_ASSERT(base->type == GGML_TYPE_F32); + GGML_ASSERT(hc_state->type == GGML_TYPE_F32); + GGML_ASSERT(n_hc > 0 && n_hc <= 8); + GGML_ASSERT(ggml_nelements(mix) >= n_hc); + GGML_ASSERT(ggml_nelements(base) >= n_hc); + GGML_ASSERT(ggml_nelements(hc_state) % n_hc == 0); + const int64_t n_embd = ggml_nelements(hc_state) / n_hc; + + struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, n_embd); + result->op = GGML_OP_DS4_HC; + result->src[0] = mix; + result->src[1] = base; + result->src[2] = hc_state; + ggml_set_op_params_i32(result, 0, 2); + ggml_set_op_params_i32(result, 1, (int32_t) n_embd); + ggml_set_op_params_i32(result, 2, (int32_t) n_hc); + ggml_set_op_params_f32(result, 4, pre_scale); + return result; +} diff --git a/ggml/src/ggml.c.ds4bak b/ggml/src/ggml.c.ds4bak new file mode 100644 index 000000000000..3b5c6bcafa23 --- /dev/null +++ b/ggml/src/ggml.c.ds4bak @@ -0,0 +1,8040 @@ +#define _CRT_SECURE_NO_DEPRECATE // Disables "unsafe" warnings on Windows +#define _USE_MATH_DEFINES // For M_PI on MSVC + +#include "ggml-backend.h" +#include "ggml-impl.h" +#include "ggml-threading.h" +#include "ggml-cpu.h" +#include "ggml.h" + +// FIXME: required here for quantization functions +#include "ggml-quants.h" +#include "../rocmfp4/rocmfp4.h" +#include "../rocmfpx/rocmfpx.h" + +#ifdef GGML_USE_CPU_HBM +#include +#endif + +#if defined(_MSC_VER) || defined(__MINGW32__) +#include // using malloc.h with MSC/MINGW +#elif !defined(__FreeBSD__) && !defined(__NetBSD__) && !defined(__OpenBSD__) +#include +#endif + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#if defined(__gnu_linux__) +#include +#endif + +#if defined(__APPLE__) +#include +#include +#include +#endif + +#if defined(_WIN32) +#define WIN32_LEAN_AND_MEAN +#ifndef NOMINMAX + #define NOMINMAX +#endif +#include +#endif + +#define UNUSED GGML_UNUSED + +// Needed for ggml_fp32_to_bf16_row() +#if defined(__AVX512BF16__) +#if defined(_MSC_VER) +#define m512i(p) p +#else +#include +#define m512i(p) (__m512i)(p) +#endif // defined(_MSC_VER) +#endif // defined(__AVX512BF16__) + +#if defined(__linux__) || \ + defined(__FreeBSD__) || defined(__NetBSD__) || defined(__OpenBSD__) || \ + (defined(__APPLE__) && !TARGET_OS_TV && !TARGET_OS_WATCH) + +#include +#include +#include +#include +#if defined(__linux__) +#include +#endif + +#if defined(__ANDROID__) +#include +#include +#include + +struct backtrace_state { + void ** current; + void ** end; +}; + +static _Unwind_Reason_Code unwind_callback(struct _Unwind_Context* context, void* arg) { + struct backtrace_state * state = (struct backtrace_state *)arg; + uintptr_t pc = _Unwind_GetIP(context); + if (pc) { + if (state->current == state->end) { + return _URC_END_OF_STACK; + } else { + *state->current++ = (void*)pc; + } + } + return _URC_NO_REASON; +} + +static void ggml_print_backtrace_symbols(void) { + const int max = 100; + void* buffer[max]; + + struct backtrace_state state = {buffer, buffer + max}; + _Unwind_Backtrace(unwind_callback, &state); + + int count = state.current - buffer; + + for (int idx = 0; idx < count; ++idx) { + const void * addr = buffer[idx]; + const char * symbol = ""; + + Dl_info info; + if (dladdr(addr, &info) && info.dli_sname) { + symbol = info.dli_sname; + } + + fprintf(stderr, "%d: %p %s\n", idx, addr, symbol); + } +} +#elif defined(__linux__) && defined(__GLIBC__) +#include +static void ggml_print_backtrace_symbols(void) { + void * trace[100]; + int nptrs = backtrace(trace, sizeof(trace)/sizeof(trace[0])); + backtrace_symbols_fd(trace, nptrs, STDERR_FILENO); +} +#elif defined(__APPLE__) +#include +static void ggml_print_backtrace_symbols(void) { + void * trace[100]; + int nptrs = backtrace(trace, sizeof(trace)/sizeof(trace[0])); + backtrace_symbols_fd(trace, nptrs, STDERR_FILENO); +} +#else +static void ggml_print_backtrace_symbols(void) { + // platform not supported +} +#endif + +void ggml_print_backtrace(void) { + const char * GGML_NO_BACKTRACE = getenv("GGML_NO_BACKTRACE"); + if (GGML_NO_BACKTRACE) { + return; + } +#if defined(__APPLE__) + // On macOS, fork+debugger attachment is problematic due to: + // 1. libdispatch "poisons" forked child processes + // 2. lldb has issues attaching to parent from forked child + // Use simple backtrace() instead to avoid Terminal.app crashes + const char * GGML_BACKTRACE_LLDB = getenv("GGML_BACKTRACE_LLDB"); + if (!GGML_BACKTRACE_LLDB) { + fprintf(stderr, "WARNING: Using native backtrace. Set GGML_BACKTRACE_LLDB for more info.\n"); + fprintf(stderr, "WARNING: GGML_BACKTRACE_LLDB may cause native MacOS Terminal.app to crash.\n"); + fprintf(stderr, "See: https://github.com/ggml-org/llama.cpp/pull/17869\n"); + ggml_print_backtrace_symbols(); + return; + } +#endif +#if defined(__linux__) + FILE * f = fopen("/proc/self/status", "r"); + size_t size = 0; + char * line = NULL; + ssize_t length = 0; + while ((length = getline(&line, &size, f)) > 0) { + if (!strncmp(line, "TracerPid:", sizeof("TracerPid:") - 1) && + (length != sizeof("TracerPid:\t0\n") - 1 || line[length - 2] != '0')) { + // Already being debugged, and the breakpoint is the later abort() + free(line); + fclose(f); + return; + } + } + free(line); + fclose(f); + int lock[2] = { -1, -1 }; + (void) !pipe(lock); // Don't start gdb until after PR_SET_PTRACER +#endif + const int parent_pid = getpid(); + const int child_pid = fork(); + if (child_pid < 0) { // error +#if defined(__linux__) + close(lock[1]); + close(lock[0]); +#endif + return; + } else if (child_pid == 0) { // child + char attach[32]; + snprintf(attach, sizeof(attach), "attach %d", parent_pid); +#if defined(__linux__) + close(lock[1]); + (void) !read(lock[0], lock, 1); + close(lock[0]); +#endif + // try gdb + execlp("gdb", "gdb", "--batch", + "-ex", "set style enabled on", + "-ex", attach, + "-ex", "bt -frame-info source-and-location", + "-ex", "detach", + "-ex", "quit", + (char *) NULL); + // try lldb + execlp("lldb", "lldb", "--batch", + "-o", "bt", + "-o", "quit", + "-p", &attach[sizeof("attach ") - 1], + (char *) NULL); + // gdb failed, fallback to backtrace_symbols + ggml_print_backtrace_symbols(); + _Exit(0); + } else { // parent +#if defined(__linux__) + prctl(PR_SET_PTRACER, child_pid); + close(lock[1]); + close(lock[0]); +#endif + waitpid(child_pid, NULL, 0); + } +} +#else +void ggml_print_backtrace(void) { + // platform not supported +} +#endif + +static ggml_abort_callback_t g_abort_callback = NULL; + +// Set the abort callback (passing null will restore original abort functionality: printing a message to stdout) +GGML_API ggml_abort_callback_t ggml_set_abort_callback(ggml_abort_callback_t callback) { + ggml_abort_callback_t ret_val = g_abort_callback; + g_abort_callback = callback; + return ret_val; +} + +void ggml_abort(const char * file, int line, const char * fmt, ...) { + fflush(stdout); + + char message[2048]; + int offset = snprintf(message, sizeof(message), "%s:%d: ", file, line); + + va_list args; + va_start(args, fmt); + vsnprintf(message + offset, sizeof(message) - offset, fmt, args); + va_end(args); + + if (g_abort_callback) { + g_abort_callback(message); + } else { + // default: print error and backtrace to stderr + fprintf(stderr, "%s\n", message); + ggml_print_backtrace(); + } + + abort(); +} + +// ggml_print_backtrace is registered with std::set_terminate by ggml.cpp + +// +// logging +// + +struct ggml_logger_state { + ggml_log_callback log_callback; + void * log_callback_user_data; +}; +static struct ggml_logger_state g_logger_state = {ggml_log_callback_default, NULL}; + +static void ggml_log_internal_v(enum ggml_log_level level, const char * format, va_list args) { + if (format == NULL) { + return; + } + va_list args_copy; + va_copy(args_copy, args); + char buffer[128]; + int len = vsnprintf(buffer, 128, format, args); + if (len < 128) { + g_logger_state.log_callback(level, buffer, g_logger_state.log_callback_user_data); + } else { + char * buffer2 = (char *) calloc(len + 1, sizeof(char)); + vsnprintf(buffer2, len + 1, format, args_copy); + buffer2[len] = 0; + g_logger_state.log_callback(level, buffer2, g_logger_state.log_callback_user_data); + free(buffer2); + } + va_end(args_copy); +} + +void ggml_log_internal(enum ggml_log_level level, const char * format, ...) { + va_list args; + va_start(args, format); + ggml_log_internal_v(level, format, args); + va_end(args); +} + +void ggml_log_callback_default(enum ggml_log_level level, const char * text, void * user_data) { + (void) level; + (void) user_data; + fputs(text, stderr); + fflush(stderr); +} + +// +// end of logging block +// + +#ifdef GGML_USE_ACCELERATE +// uncomment to use vDSP for soft max computation +// note: not sure if it is actually faster +//#define GGML_SOFT_MAX_ACCELERATE +#endif + + +void * ggml_aligned_malloc(size_t size) { +#if defined(__s390x__) + const int alignment = 256; +#else + const int alignment = 64; +#endif + +#if defined(_MSC_VER) || defined(__MINGW32__) + return _aligned_malloc(size, alignment); +#else + if (size == 0) { + GGML_LOG_WARN("Behavior may be unexpected when allocating 0 bytes for ggml_aligned_malloc!\n"); + return NULL; + } + void * aligned_memory = NULL; + #ifdef GGML_USE_CPU_HBM + int result = hbw_posix_memalign(&aligned_memory, alignment, size); + #elif TARGET_OS_OSX + GGML_UNUSED(alignment); + kern_return_t alloc_status = vm_allocate((vm_map_t) mach_task_self(), (vm_address_t *) &aligned_memory, size, VM_FLAGS_ANYWHERE); + int result = EFAULT; + switch (alloc_status) { + case KERN_SUCCESS: + result = 0; + break; + case KERN_INVALID_ADDRESS: + result = EINVAL; + break; + case KERN_NO_SPACE: + result = ENOMEM; + break; + default: + result = EFAULT; + break; + } + #else + int result = posix_memalign(&aligned_memory, alignment, size); + #endif + if (result != 0) { + // Handle allocation failure + const char *error_desc = "unknown allocation error"; + switch (result) { + case EINVAL: + error_desc = "invalid alignment value"; + break; + case ENOMEM: + error_desc = "insufficient memory"; + break; + } + GGML_LOG_ERROR("%s: %s (attempted to allocate %6.2f MB)\n", __func__, error_desc, size/(1024.0*1024.0)); + return NULL; + } + return aligned_memory; +#endif +} + +void ggml_aligned_free(void * ptr, size_t size) { + GGML_UNUSED(size); +#if defined(_MSC_VER) || defined(__MINGW32__) + _aligned_free(ptr); +#elif GGML_USE_CPU_HBM + if (ptr != NULL) { + hbw_free(ptr); + } +#elif TARGET_OS_OSX + if (ptr != NULL) { + vm_deallocate((vm_map_t)mach_task_self(), (vm_address_t)ptr, size); + } +#else + free(ptr); +#endif +} + + +inline static void * ggml_malloc(size_t size) { + if (size == 0) { + GGML_LOG_WARN("Behavior may be unexpected when allocating 0 bytes for ggml_malloc!\n"); + return NULL; + } + void * result = malloc(size); + if (result == NULL) { + GGML_LOG_ERROR("%s: failed to allocate %6.2f MB\n", __func__, size/(1024.0*1024.0)); + GGML_ABORT("fatal error"); + } + return result; +} + +// calloc +inline static void * ggml_calloc(size_t num, size_t size) { + if (num == 0 || size == 0) { + GGML_LOG_WARN("Behavior may be unexpected when allocating 0 bytes for ggml_calloc!\n"); + return NULL; + } + void * result = calloc(num, size); + if (result == NULL) { + GGML_LOG_ERROR("%s: failed to allocate %6.2f MB\n", __func__, size/(1024.0*1024.0)); + GGML_ABORT("fatal error"); + } + return result; +} + +#define GGML_MALLOC(size) ggml_malloc(size) +#define GGML_CALLOC(num, size) ggml_calloc(num, size) + +#define GGML_FREE(ptr) free(ptr) + +const char * ggml_status_to_string(enum ggml_status status) { + switch (status) { + case GGML_STATUS_ALLOC_FAILED: return "GGML status: error (failed to allocate memory)"; + case GGML_STATUS_FAILED: return "GGML status: error (operation failed)"; + case GGML_STATUS_SUCCESS: return "GGML status: success"; + case GGML_STATUS_ABORTED: return "GGML status: warning (operation aborted)"; + } + + return "GGML status: unknown"; +} + +float ggml_fp16_to_fp32(ggml_fp16_t x) { +#define ggml_fp16_to_fp32 do_not_use__ggml_fp16_to_fp32__in_ggml + return GGML_FP16_TO_FP32(x); +} + +ggml_fp16_t ggml_fp32_to_fp16(float x) { +#define ggml_fp32_to_fp16 do_not_use__ggml_fp32_to_fp16__in_ggml + return GGML_FP32_TO_FP16(x); +} + +float ggml_bf16_to_fp32(ggml_bf16_t x) { +#define ggml_bf16_to_fp32 do_not_use__ggml_bf16_to_fp32__in_ggml + return GGML_BF16_TO_FP32(x); // it just left shifts +} + +ggml_bf16_t ggml_fp32_to_bf16(float x) { +#define ggml_fp32_to_bf16 do_not_use__ggml_fp32_to_bf16__in_ggml + return GGML_FP32_TO_BF16(x); +} + +void ggml_fp16_to_fp32_row(const ggml_fp16_t * x, float * y, int64_t n) { + for (int64_t i = 0; i < n; i++) { + y[i] = GGML_FP16_TO_FP32(x[i]); + } +} + +void ggml_fp32_to_fp16_row(const float * x, ggml_fp16_t * y, int64_t n) { + int i = 0; + for (; i < n; ++i) { + y[i] = GGML_FP32_TO_FP16(x[i]); + } +} + +void ggml_bf16_to_fp32_row(const ggml_bf16_t * x, float * y, int64_t n) { + int i = 0; + for (; i < n; ++i) { + y[i] = GGML_BF16_TO_FP32(x[i]); + } +} + +void ggml_fp32_to_bf16_row_ref(const float * x, ggml_bf16_t * y, int64_t n) { + for (int i = 0; i < n; i++) { + y[i] = ggml_compute_fp32_to_bf16(x[i]); + } +} + +void ggml_fp32_to_bf16_row(const float * x, ggml_bf16_t * y, int64_t n) { + int i = 0; +#if defined(__AVX512BF16__) + // subnormals are flushed to zero on this platform + for (; i + 32 <= n; i += 32) { + _mm512_storeu_si512( + (__m512i *)(y + i), + m512i(_mm512_cvtne2ps_pbh(_mm512_loadu_ps(x + i + 16), + _mm512_loadu_ps(x + i)))); + } +#endif + for (; i < n; i++) { + y[i] = GGML_FP32_TO_BF16(x[i]); + } +} + +bool ggml_guid_matches(ggml_guid_t guid_a, ggml_guid_t guid_b) { + return memcmp(guid_a, guid_b, sizeof(ggml_guid)) == 0; +} + +const char * ggml_version(void) { + return GGML_VERSION; +} + +const char * ggml_commit(void) { + return GGML_COMMIT; +} + +// +// timing +// + +#if defined(_MSC_VER) || defined(__MINGW32__) +static int64_t timer_freq, timer_start; +void ggml_time_init(void) { + LARGE_INTEGER t; + QueryPerformanceFrequency(&t); + timer_freq = t.QuadPart; + + // The multiplication by 1000 or 1000000 below can cause an overflow if timer_freq + // and the uptime is high enough. + // We subtract the program start time to reduce the likelihood of that happening. + QueryPerformanceCounter(&t); + timer_start = t.QuadPart; +} +int64_t ggml_time_ms(void) { + LARGE_INTEGER t; + QueryPerformanceCounter(&t); + return ((t.QuadPart-timer_start) * 1000) / timer_freq; +} +int64_t ggml_time_us(void) { + LARGE_INTEGER t; + QueryPerformanceCounter(&t); + return ((t.QuadPart-timer_start) * 1000000) / timer_freq; +} +#else +void ggml_time_init(void) {} +int64_t ggml_time_ms(void) { + struct timespec ts; + clock_gettime(CLOCK_MONOTONIC, &ts); + return (int64_t)ts.tv_sec*1000 + (int64_t)ts.tv_nsec/1000000; +} + +int64_t ggml_time_us(void) { + struct timespec ts; + clock_gettime(CLOCK_MONOTONIC, &ts); + return (int64_t)ts.tv_sec*1000000 + (int64_t)ts.tv_nsec/1000; +} +#endif + +int64_t ggml_cycles(void) { + return clock(); +} + +int64_t ggml_cycles_per_ms(void) { + return CLOCKS_PER_SEC/1000; +} + +// +// cross-platform UTF-8 file paths +// + +#ifdef _WIN32 +static wchar_t * ggml_mbstowcs(const char * mbs) { + int wlen = MultiByteToWideChar(CP_UTF8, 0, mbs, -1, NULL, 0); + if (!wlen) { + errno = EINVAL; + return NULL; + } + + wchar_t * wbuf = GGML_MALLOC(wlen * sizeof(wchar_t)); + wlen = MultiByteToWideChar(CP_UTF8, 0, mbs, -1, wbuf, wlen); + if (!wlen) { + GGML_FREE(wbuf); + errno = EINVAL; + return NULL; + } + + return wbuf; +} +#endif + +FILE * ggml_fopen(const char * fname, const char * mode) { +#ifdef _WIN32 + FILE * file = NULL; + + // convert fname (UTF-8) + wchar_t * wfname = ggml_mbstowcs(fname); + if (wfname) { + // convert mode (ANSI) + wchar_t * wmode = GGML_MALLOC((strlen(mode) + 1) * sizeof(wchar_t)); + wchar_t * wmode_p = wmode; + do { + *wmode_p++ = (wchar_t)*mode; + } while (*mode++); + + // open file + file = _wfopen(wfname, wmode); + + GGML_FREE(wfname); + GGML_FREE(wmode); + } + + return file; +#else + return fopen(fname, mode); +#endif + +} + +static const struct ggml_type_traits type_traits[GGML_TYPE_COUNT] = { + [GGML_TYPE_I8] = { + .type_name = "i8", + .blck_size = 1, + .type_size = sizeof(int8_t), + .is_quantized = false, + }, + [GGML_TYPE_I16] = { + .type_name = "i16", + .blck_size = 1, + .type_size = sizeof(int16_t), + .is_quantized = false, + }, + [GGML_TYPE_I32] = { + .type_name = "i32", + .blck_size = 1, + .type_size = sizeof(int32_t), + .is_quantized = false, + }, + [GGML_TYPE_I64] = { + .type_name = "i64", + .blck_size = 1, + .type_size = sizeof(int64_t), + .is_quantized = false, + }, + [GGML_TYPE_F64] = { + .type_name = "f64", + .blck_size = 1, + .type_size = sizeof(double), + .is_quantized = false, + }, + [GGML_TYPE_F32] = { + .type_name = "f32", + .blck_size = 1, + .type_size = sizeof(float), + .is_quantized = false, + }, + [GGML_TYPE_F16] = { + .type_name = "f16", + .blck_size = 1, + .type_size = sizeof(ggml_fp16_t), + .is_quantized = false, + .to_float = (ggml_to_float_t) ggml_fp16_to_fp32_row, + .from_float_ref = (ggml_from_float_t) ggml_fp32_to_fp16_row, + }, + [GGML_TYPE_Q1_0] = { + .type_name = "q1_0", + .blck_size = QK1_0, + .type_size = sizeof(block_q1_0), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_q1_0, + .from_float_ref = (ggml_from_float_t) quantize_row_q1_0_ref, + }, + [GGML_TYPE_TQ3_0] = { + .type_name = "tq3_0", + .blck_size = QK_TQ3_0, + .type_size = sizeof(block_tq3_0), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_tq3_0, + .from_float_ref = (ggml_from_float_t) quantize_row_tq3_0_ref, + }, + [GGML_TYPE_Q4_0] = { + .type_name = "q4_0", + .blck_size = QK4_0, + .type_size = sizeof(block_q4_0), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_q4_0, + .from_float_ref = (ggml_from_float_t) quantize_row_q4_0_ref, + }, + [GGML_TYPE_Q4_1] = { + .type_name = "q4_1", + .blck_size = QK4_1, + .type_size = sizeof(block_q4_1), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_q4_1, + .from_float_ref = (ggml_from_float_t) quantize_row_q4_1_ref, + }, + [4] = { // GGML_TYPE_Q4_2 + .type_name = "DEPRECATED", + .blck_size = 0, + .type_size = 0, + .is_quantized = false, + }, + [5] = { // GGML_TYPE_Q4_3 + .type_name = "DEPRECATED", + .blck_size = 0, + .type_size = 0, + .is_quantized = false, + }, + [GGML_TYPE_Q5_0] = { + .type_name = "q5_0", + .blck_size = QK5_0, + .type_size = sizeof(block_q5_0), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_q5_0, + .from_float_ref = (ggml_from_float_t) quantize_row_q5_0_ref, + }, + [GGML_TYPE_Q5_1] = { + .type_name = "q5_1", + .blck_size = QK5_1, + .type_size = sizeof(block_q5_1), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_q5_1, + .from_float_ref = (ggml_from_float_t) quantize_row_q5_1_ref, + }, + [GGML_TYPE_Q8_0] = { + .type_name = "q8_0", + .blck_size = QK8_0, + .type_size = sizeof(block_q8_0), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_q8_0, + .from_float_ref = (ggml_from_float_t) quantize_row_q8_0_ref, + }, + [GGML_TYPE_Q4_0_ROCMFP4] = { + .type_name = "q4_0_rocmfp4", + .blck_size = QK_ROCMFP4, + .type_size = sizeof(block_rocmfp4), + .is_quantized = true, + .to_float = (ggml_to_float_t) rocmfp4_dequantize_row_q4_0, + .from_float_ref = (ggml_from_float_t) rocmfp4_quantize_row_q4_0_ref, + }, + [GGML_TYPE_Q4_0_ROCMFP4_FAST] = { + .type_name = "q4_0_rocmfp4_fast", + .blck_size = QK_ROCMFP4, + .type_size = sizeof(block_rocmfp4_fast), + .is_quantized = true, + .to_float = (ggml_to_float_t) rocmfp4_dequantize_row_q4_0_fast, + .from_float_ref = (ggml_from_float_t) rocmfp4_quantize_row_q4_0_fast_ref, + }, + [GGML_TYPE_Q3_0_ROCMFPX] = { + .type_name = "q3_0_rocmfpx", + .blck_size = QK_ROCMFP3, + .type_size = sizeof(block_rocmfp3), + .is_quantized = true, + .to_float = (ggml_to_float_t) rocmfpx_dequantize_row_fp3, + .from_float_ref = (ggml_from_float_t) rocmfpx_quantize_row_fp3_ref, + }, + [GGML_TYPE_Q2_0_ROCMFP2] = { + .type_name = "q2_0_rocmfp2", + .blck_size = QK_ROCMFP2, + .type_size = sizeof(block_rocmfp2), + .is_quantized = true, + .to_float = (ggml_to_float_t) rocmfpx_dequantize_row_fp2, + .from_float_ref = (ggml_from_float_t) rocmfpx_quantize_row_fp2_ref, + }, + [GGML_TYPE_Q6_0_ROCMFPX] = { + .type_name = "q6_0_rocmfpx", + .blck_size = QK_ROCMFP6, + .type_size = sizeof(block_rocmfp6), + .is_quantized = true, + .to_float = (ggml_to_float_t) rocmfpx_dequantize_row_fp6, + .from_float_ref = (ggml_from_float_t) rocmfpx_quantize_row_fp6_ref, + }, + [GGML_TYPE_Q8_0_ROCMFPX] = { + .type_name = "q8_0_rocmfpx", + .blck_size = QK_ROCMFP8, + .type_size = sizeof(block_rocmfp8), + .is_quantized = true, + .to_float = (ggml_to_float_t) rocmfpx_dequantize_row_fp8, + .from_float_ref = (ggml_from_float_t) rocmfpx_quantize_row_fp8_ref, + }, + [GGML_TYPE_Q8_1] = { + .type_name = "q8_1", + .blck_size = QK8_1, + .type_size = sizeof(block_q8_1), + .is_quantized = true, + .from_float_ref = (ggml_from_float_t) quantize_row_q8_1_ref, + }, + [GGML_TYPE_MXFP4] = { + .type_name = "mxfp4", + .blck_size = QK_MXFP4, + .type_size = sizeof(block_mxfp4), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_mxfp4, + .from_float_ref = (ggml_from_float_t)quantize_row_mxfp4_ref, + }, + [GGML_TYPE_NVFP4] = { + .type_name = "nvfp4", + .blck_size = QK_NVFP4, + .type_size = sizeof(block_nvfp4), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_nvfp4, + .from_float_ref = (ggml_from_float_t)quantize_row_nvfp4_ref, + }, + [GGML_TYPE_Q2_K] = { + .type_name = "q2_K", + .blck_size = QK_K, + .type_size = sizeof(block_q2_K), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_q2_K, + .from_float_ref = (ggml_from_float_t) quantize_row_q2_K_ref, + }, + [GGML_TYPE_Q3_K] = { + .type_name = "q3_K", + .blck_size = QK_K, + .type_size = sizeof(block_q3_K), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_q3_K, + .from_float_ref = (ggml_from_float_t) quantize_row_q3_K_ref, + }, + [GGML_TYPE_Q4_K] = { + .type_name = "q4_K", + .blck_size = QK_K, + .type_size = sizeof(block_q4_K), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_q4_K, + .from_float_ref = (ggml_from_float_t) quantize_row_q4_K_ref, + }, + [GGML_TYPE_Q5_K] = { + .type_name = "q5_K", + .blck_size = QK_K, + .type_size = sizeof(block_q5_K), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_q5_K, + .from_float_ref = (ggml_from_float_t) quantize_row_q5_K_ref, + }, + [GGML_TYPE_Q6_K] = { + .type_name = "q6_K", + .blck_size = QK_K, + .type_size = sizeof(block_q6_K), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_q6_K, + .from_float_ref = (ggml_from_float_t) quantize_row_q6_K_ref, + }, + [GGML_TYPE_IQ2_XXS] = { + .type_name = "iq2_xxs", + .blck_size = QK_K, + .type_size = sizeof(block_iq2_xxs), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_iq2_xxs, + .from_float_ref = NULL, + }, + [GGML_TYPE_IQ2_XS] = { + .type_name = "iq2_xs", + .blck_size = QK_K, + .type_size = sizeof(block_iq2_xs), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_iq2_xs, + .from_float_ref = NULL, + }, + [GGML_TYPE_IQ3_XXS] = { + .type_name = "iq3_xxs", + .blck_size = QK_K, + .type_size = sizeof(block_iq3_xxs), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_iq3_xxs, + .from_float_ref = (ggml_from_float_t)quantize_row_iq3_xxs_ref, + }, + [GGML_TYPE_IQ3_S] = { + .type_name = "iq3_s", + .blck_size = QK_K, + .type_size = sizeof(block_iq3_s), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_iq3_s, + .from_float_ref = (ggml_from_float_t)quantize_row_iq3_s_ref, + }, + [GGML_TYPE_IQ2_S] = { + .type_name = "iq2_s", + .blck_size = QK_K, + .type_size = sizeof(block_iq2_s), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_iq2_s, + .from_float_ref = (ggml_from_float_t)quantize_row_iq2_s_ref, + }, + [GGML_TYPE_IQ1_S] = { + .type_name = "iq1_s", + .blck_size = QK_K, + .type_size = sizeof(block_iq1_s), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_iq1_s, + .from_float_ref = NULL, + }, + [GGML_TYPE_IQ1_M] = { + .type_name = "iq1_m", + .blck_size = QK_K, + .type_size = sizeof(block_iq1_m), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_iq1_m, + .from_float_ref = NULL, + }, + [GGML_TYPE_IQ4_NL] = { + .type_name = "iq4_nl", + .blck_size = QK4_NL, + .type_size = sizeof(block_iq4_nl), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_iq4_nl, + .from_float_ref = (ggml_from_float_t)quantize_row_iq4_nl_ref, + }, + [GGML_TYPE_IQ4_XS] = { + .type_name = "iq4_xs", + .blck_size = QK_K, + .type_size = sizeof(block_iq4_xs), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_iq4_xs, + .from_float_ref = (ggml_from_float_t)quantize_row_iq4_xs_ref, + }, + [GGML_TYPE_Q8_K] = { + .type_name = "q8_K", + .blck_size = QK_K, + .type_size = sizeof(block_q8_K), + .is_quantized = true, + }, + [GGML_TYPE_BF16] = { + .type_name = "bf16", + .blck_size = 1, + .type_size = sizeof(ggml_bf16_t), + .is_quantized = false, + .to_float = (ggml_to_float_t) ggml_bf16_to_fp32_row, + .from_float_ref = (ggml_from_float_t) ggml_fp32_to_bf16_row_ref, + }, + [31] = { // GGML_TYPE_Q4_0_4_4 + .type_name = "TYPE_Q4_0_4_4 REMOVED, use Q4_0 with runtime repacking", + .blck_size = 0, + .type_size = 0, + .is_quantized = false, + }, + [32] = { // GGML_TYPE_Q4_0_4_8 + .type_name = "TYPE_Q4_0_4_8 REMOVED, use Q4_0 with runtime repacking", + .blck_size = 0, + .type_size = 0, + .is_quantized = false, + }, + [33] = { // GGML_TYPE_Q4_0_8_8 + .type_name = "TYPE_Q4_0_8_8 REMOVED, use Q4_0 with runtime repacking", + .blck_size = 0, + .type_size = 0, + .is_quantized = false, + }, + [GGML_TYPE_TQ1_0] = { + .type_name = "tq1_0", + .blck_size = QK_K, + .type_size = sizeof(block_tq1_0), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_tq1_0, + .from_float_ref = (ggml_from_float_t) quantize_row_tq1_0_ref, + }, + [GGML_TYPE_TQ2_0] = { + .type_name = "tq2_0", + .blck_size = QK_K, + .type_size = sizeof(block_tq2_0), + .is_quantized = true, + .to_float = (ggml_to_float_t) dequantize_row_tq2_0, + .from_float_ref = (ggml_from_float_t) quantize_row_tq2_0_ref, + }, + [36] = { // GGML_TYPE_IQ4_NL_4_4 + .type_name = "TYPE_IQ4_NL_4_4 REMOVED, use IQ4_NL with runtime repacking", + .blck_size = 0, + .type_size = 0, + .is_quantized = false, + }, + [37] = { // GGML_TYPE_IQ4_NL_4_8 + .type_name = "TYPE_IQ4_NL_4_8 REMOVED, use IQ4_NL with runtime repacking", + .blck_size = 0, + .type_size = 0, + .is_quantized = false, + }, + [38] = { // GGML_TYPE_IQ4_NL_8_8 + .type_name = "TYPE_IQ4_NL_8_8 REMOVED, use IQ4_NL with runtime repacking", + .blck_size = 0, + .type_size = 0, + .is_quantized = false, + }, +}; + +const struct ggml_type_traits * ggml_get_type_traits(enum ggml_type type) { + assert(type >= 0); + assert(type < GGML_TYPE_COUNT); + return &type_traits[type]; +} + +// +// ggml object +// + +struct ggml_object { + size_t offs; + size_t size; + + struct ggml_object * next; + + enum ggml_object_type type; + + char padding[4]; +}; + +static const size_t GGML_OBJECT_SIZE = sizeof(struct ggml_object); + +// +// ggml context +// + +struct ggml_context { + size_t mem_size; + void * mem_buffer; + bool mem_buffer_owned; + bool no_alloc; + + int n_objects; + + struct ggml_object * objects_begin; + struct ggml_object * objects_end; +}; + +// +// data types +// + +static const char * GGML_OP_NAME[GGML_OP_COUNT] = { + "NONE", + + "DUP", + "ADD", + "ADD_ID", + "ADD1", + "ACC", + "SUB", + "MUL", + "DIV", + "SQR", + "SQRT", + "LOG", + "SIN", + "COS", + "SUM", + "SUM_ROWS", + "CUMSUM", + "MEAN", + "ARGMAX", + "COUNT_EQUAL", + "REPEAT", + "REPEAT_BACK", + "CONCAT", + "SILU_BACK", + "NORM", + "RMS_NORM", + "RMS_NORM_BACK", + "GROUP_NORM", + "L2_NORM", + + "MUL_MAT", + "MUL_MAT_ID", + "OUT_PROD", + + "SCALE", + "SET", + "CPY", + "CONT", + "RESHAPE", + "VIEW", + "PERMUTE", + "TRANSPOSE", + "GET_ROWS", + "GET_ROWS_BACK", + "SET_ROWS", + "DIAG", + "DIAG_MASK_INF", + "DIAG_MASK_ZERO", + "SOFT_MAX", + "SOFT_MAX_BACK", + "ROPE", + "ROPE_BACK", + "CLAMP", + "CONV_TRANSPOSE_1D", + "IM2COL", + "IM2COL_BACK", + "IM2COL_3D", + "CONV_2D", + "CONV_3D", + "CONV_2D_DW", + "CONV_TRANSPOSE_2D", + "POOL_1D", + "POOL_2D", + "POOL_2D_BACK", + "UPSCALE", + "PAD", + "PAD_REFLECT_1D", + "ROLL", + "ARANGE", + "TIMESTEP_EMBEDDING", + "ARGSORT", + "TOP_K", + "LEAKY_RELU", + "TRI", + "FILL", + + "FLASH_ATTN_EXT", + "FLASH_ATTN_SPARSE", + "FLASH_ATTN_BACK", + "SSM_CONV", + "SSM_SCAN", + "WIN_PART", + "WIN_UNPART", + "GET_REL_POS", + "ADD_REL_POS", + "RWKV_WKV6", + "GATED_LINEAR_ATTN", + "RWKV_WKV7", + "SOLVE_TRI", + "GATED_DELTA_NET", + + "UNARY", + + "MAP_CUSTOM1", + "MAP_CUSTOM2", + "MAP_CUSTOM3", + + "CUSTOM", + + "CROSS_ENTROPY_LOSS", + "CROSS_ENTROPY_LOSS_BACK", + "OPT_STEP_ADAMW", + "OPT_STEP_SGD", + + "GLU", + + "TURBO_WHT", +}; + +static_assert(GGML_OP_COUNT == 99, "GGML_OP_COUNT != 99"); + +static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { + "none", + + "x", + "x+y", + "x[i]+y", + "x+y", + "view(x,nb,offset)+=y->x", + "x-y", + "x*y", + "x/y", + "x^2", + "√x", + "log(x)", + "sin(x)", + "cos(x)", + "Σx", + "Σx_k", + "cumsum(x)", + "Σx/n", + "argmax(x)", + "count_equal(x)", + "repeat(x)", + "repeat_back(x)", + "concat(x, y)", + "silu_back(x)", + "norm(x)", + "rms_norm(x)", + "rms_norm_back(x)", + "group_norm(x)", + "l2_norm(x)", + + "X*Y", + "X[i]*Y", + "X*Y", + + "x*v", + "y-\\>view(x)", + "x-\\>y", + "cont(x)", + "reshape(x)", + "view(x)", + "permute(x)", + "transpose(x)", + "get_rows(x)", + "get_rows_back(x)", + "set_rows(x)", + "diag(x)", + "diag_mask_inf(x)", + "diag_mask_zero(x)", + "soft_max(x)", + "soft_max_back(x)", + "rope(x)", + "rope_back(x)", + "clamp(x)", + "conv_transpose_1d(x)", + "im2col(x)", + "im2col_back(x)", + "im2col_3d(x)", + "conv_2d(x)", + "conv_3d(x)", + "conv_2d_dw(x)", + "conv_transpose_2d(x)", + "pool_1d(x)", + "pool_2d(x)", + "pool_2d_back(x)", + "upscale(x)", + "pad(x)", + "pad_reflect_1d(x)", + "roll(x)", + "arange(start, stop, step)", + "timestep_embedding(timesteps, dim, max_period)", + "argsort(x)", + "top_k(x)", + "leaky_relu(x)", + "tri(x)", + "fill(x, c)", + + "flash_attn_ext(x)", + "sparse_fa(x)", + "flash_attn_back(x)", + "ssm_conv(x)", + "ssm_scan(x)", + "win_part(x)", + "win_unpart(x)", + "get_rel_pos(x)", + "add_rel_pos(x)", + "rwkv_wkv6(k, v, r, tf, td, s)", + "gated_linear_attn(k, v, q, gate, s)", + "rwkv_wkv7(r, w, k, v, a, b, s)", + "A X = B, A triangular, solve X", + "gated_delta_net(q, k, v, g, beta, s)", + + "unary(x)", + + "map_custom(x)", + "map_custom(x,y)", + "map_custom(x,y,z)", + + "custom(x)", + + "cross_entropy_loss(x,y)", + "cross_entropy_loss_back(x,y)", + "adamw(x)", + "sgd(x)", + + "glu(x)", + + "turbo_wht(a)", +}; + +static_assert(GGML_OP_COUNT == 99, "GGML_OP_COUNT != 99"); + +static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); + +static const char * GGML_UNARY_OP_NAME[GGML_UNARY_OP_COUNT] = { + "ABS", + "SGN", + "NEG", + "STEP", + "TANH", + "ELU", + "RELU", + "SIGMOID", + "GELU", + "GELU_QUICK", + "SILU", + "HARDSWISH", + "HARDSIGMOID", + "EXP", + "EXPM1", + "SOFTPLUS", + "GELU_ERF", + "XIELU", + "FLOOR", + "CEIL", + "ROUND", + "TRUNC", +}; + +static_assert(GGML_UNARY_OP_COUNT == 22, "GGML_UNARY_OP_COUNT != 22"); + +static const char * GGML_GLU_OP_NAME[GGML_GLU_OP_COUNT] = { + "REGLU", + "GEGLU", + "SWIGLU", + "SWIGLU_OAI", + "SWIGLU_DS4", + "GEGLU_ERF", + "GEGLU_QUICK", +}; + +static_assert(GGML_GLU_OP_COUNT == 7, "GGML_GLU_OP_COUNT != 7"); + + +static_assert(sizeof(struct ggml_object)%GGML_MEM_ALIGN == 0, "ggml_object size must be a multiple of GGML_MEM_ALIGN"); +static_assert(sizeof(struct ggml_tensor)%GGML_MEM_ALIGN == 0, "ggml_tensor size must be a multiple of GGML_MEM_ALIGN"); + + +//////////////////////////////////////////////////////////////////////////////// + +void ggml_print_object(const struct ggml_object * obj) { + GGML_LOG_INFO(" - ggml_object: type = %d, offset = %zu, size = %zu, next = %p\n", + obj->type, obj->offs, obj->size, (const void *) obj->next); +} + +void ggml_print_objects(const struct ggml_context * ctx) { + struct ggml_object * obj = ctx->objects_begin; + + GGML_LOG_INFO("%s: objects in context %p:\n", __func__, (const void *) ctx); + + while (obj != NULL) { + ggml_print_object(obj); + obj = obj->next; + } + + GGML_LOG_INFO("%s: --- end ---\n", __func__); +} + +int64_t ggml_nelements(const struct ggml_tensor * tensor) { + static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); + + return tensor->ne[0]*tensor->ne[1]*tensor->ne[2]*tensor->ne[3]; +} + +int64_t ggml_nrows(const struct ggml_tensor * tensor) { + static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); + + return tensor->ne[1]*tensor->ne[2]*tensor->ne[3]; +} + +size_t ggml_nbytes(const struct ggml_tensor * tensor) { + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (tensor->ne[i] <= 0) { + return 0; + } + } + + size_t nbytes; + const size_t blck_size = ggml_blck_size(tensor->type); + if (blck_size == 1) { + nbytes = ggml_type_size(tensor->type); + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + nbytes += (tensor->ne[i] - 1)*tensor->nb[i]; + } + } + else { + nbytes = tensor->ne[0]*tensor->nb[0]/blck_size; + for (int i = 1; i < GGML_MAX_DIMS; ++i) { + nbytes += (tensor->ne[i] - 1)*tensor->nb[i]; + } + } + + return nbytes; +} + +size_t ggml_nbytes_pad(const struct ggml_tensor * tensor) { + return GGML_PAD(ggml_nbytes(tensor), GGML_MEM_ALIGN); +} + +int64_t ggml_blck_size(enum ggml_type type) { + assert(type >= 0); + assert(type < GGML_TYPE_COUNT); + return type_traits[type].blck_size; +} + +size_t ggml_type_size(enum ggml_type type) { + assert(type >= 0); + assert(type < GGML_TYPE_COUNT); + return type_traits[type].type_size; +} + +size_t ggml_row_size(enum ggml_type type, int64_t ne) { + assert(type >= 0); + assert(type < GGML_TYPE_COUNT); + assert(ne % ggml_blck_size(type) == 0); + return ggml_type_size(type)*ne/ggml_blck_size(type); +} + +double ggml_type_sizef(enum ggml_type type) { + assert(type >= 0); + assert(type < GGML_TYPE_COUNT); + return ((double)(type_traits[type].type_size))/type_traits[type].blck_size; +} + +const char * ggml_type_name(enum ggml_type type) { + assert(type >= 0); + assert(type < GGML_TYPE_COUNT); + return type_traits[type].type_name; +} + +bool ggml_is_quantized(enum ggml_type type) { + assert(type >= 0); + assert(type < GGML_TYPE_COUNT); + return type_traits[type].is_quantized; +} + +const char * ggml_op_name(enum ggml_op op) { + return GGML_OP_NAME[op]; +} + +const char * ggml_op_symbol(enum ggml_op op) { + return GGML_OP_SYMBOL[op]; +} + +const char * ggml_unary_op_name(enum ggml_unary_op op) { + return GGML_UNARY_OP_NAME[op]; +} + +const char * ggml_glu_op_name(enum ggml_glu_op op) { + return GGML_GLU_OP_NAME[op]; +} + +const char * ggml_op_desc(const struct ggml_tensor * t) { + if (t->op == GGML_OP_UNARY) { + enum ggml_unary_op uop = ggml_get_unary_op(t); + return ggml_unary_op_name(uop); + } + if (t->op == GGML_OP_GLU) { + enum ggml_glu_op gop = ggml_get_glu_op(t); + return ggml_glu_op_name(gop); + } + return ggml_op_name(t->op); +} + +size_t ggml_element_size(const struct ggml_tensor * tensor) { + return ggml_type_size(tensor->type); +} + +bool ggml_is_scalar(const struct ggml_tensor * tensor) { + static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); + + return tensor->ne[0] == 1 && tensor->ne[1] == 1 && tensor->ne[2] == 1 && tensor->ne[3] == 1; +} + +bool ggml_is_vector(const struct ggml_tensor * tensor) { + static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); + + return tensor->ne[1] == 1 && tensor->ne[2] == 1 && tensor->ne[3] == 1; +} + +bool ggml_is_matrix(const struct ggml_tensor * tensor) { + static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); + + return tensor->ne[2] == 1 && tensor->ne[3] == 1; +} + +bool ggml_is_3d(const struct ggml_tensor * tensor) { + return tensor->ne[3] == 1; +} + +int ggml_n_dims(const struct ggml_tensor * tensor) { + for (int i = GGML_MAX_DIMS - 1; i >= 1; --i) { + if (tensor->ne[i] > 1) { + return i + 1; + } + } + return 1; +} + +enum ggml_type ggml_ftype_to_ggml_type(enum ggml_ftype ftype) { + enum ggml_type wtype = GGML_TYPE_COUNT; + + switch (ftype) { + case GGML_FTYPE_ALL_F32: wtype = GGML_TYPE_F32; break; + case GGML_FTYPE_MOSTLY_F16: wtype = GGML_TYPE_F16; break; + case GGML_FTYPE_MOSTLY_BF16: wtype = GGML_TYPE_BF16; break; + case GGML_FTYPE_MOSTLY_Q4_0: wtype = GGML_TYPE_Q4_0; break; + case GGML_FTYPE_MOSTLY_Q4_1: wtype = GGML_TYPE_Q4_1; break; + case GGML_FTYPE_MOSTLY_Q1_0: wtype = GGML_TYPE_Q1_0; break; + case GGML_FTYPE_MOSTLY_Q5_0: wtype = GGML_TYPE_Q5_0; break; + case GGML_FTYPE_MOSTLY_Q5_1: wtype = GGML_TYPE_Q5_1; break; + case GGML_FTYPE_MOSTLY_Q8_0: wtype = GGML_TYPE_Q8_0; break; + case GGML_FTYPE_MOSTLY_MXFP4: wtype = GGML_TYPE_MXFP4; break; + case GGML_FTYPE_MOSTLY_NVFP4: wtype = GGML_TYPE_NVFP4; break; + case GGML_FTYPE_MOSTLY_Q2_K: wtype = GGML_TYPE_Q2_K; break; + case GGML_FTYPE_MOSTLY_Q3_K: wtype = GGML_TYPE_Q3_K; break; + case GGML_FTYPE_MOSTLY_Q4_K: wtype = GGML_TYPE_Q4_K; break; + case GGML_FTYPE_MOSTLY_Q5_K: wtype = GGML_TYPE_Q5_K; break; + case GGML_FTYPE_MOSTLY_Q6_K: wtype = GGML_TYPE_Q6_K; break; + case GGML_FTYPE_MOSTLY_IQ2_XXS: wtype = GGML_TYPE_IQ2_XXS; break; + case GGML_FTYPE_MOSTLY_IQ2_XS: wtype = GGML_TYPE_IQ2_XS; break; + case GGML_FTYPE_MOSTLY_IQ3_XXS: wtype = GGML_TYPE_IQ3_XXS; break; + case GGML_FTYPE_MOSTLY_IQ1_S: wtype = GGML_TYPE_IQ1_S; break; + case GGML_FTYPE_MOSTLY_IQ1_M: wtype = GGML_TYPE_IQ1_M; break; + case GGML_FTYPE_MOSTLY_IQ4_NL: wtype = GGML_TYPE_IQ4_NL; break; + case GGML_FTYPE_MOSTLY_IQ4_XS: wtype = GGML_TYPE_IQ4_XS; break; + case GGML_FTYPE_MOSTLY_IQ3_S: wtype = GGML_TYPE_IQ3_S; break; + case GGML_FTYPE_MOSTLY_IQ2_S: wtype = GGML_TYPE_IQ2_S; break; + case GGML_FTYPE_UNKNOWN: wtype = GGML_TYPE_COUNT; break; + case GGML_FTYPE_MOSTLY_Q4_1_SOME_F16: wtype = GGML_TYPE_COUNT; break; + } + + GGML_ASSERT(wtype != GGML_TYPE_COUNT); + + return wtype; +} + +size_t ggml_tensor_overhead(void) { + return GGML_OBJECT_SIZE + GGML_TENSOR_SIZE; +} + +bool ggml_is_transposed(const struct ggml_tensor * tensor) { + return tensor->nb[0] > tensor->nb[1]; +} + +static bool ggml_is_contiguous_n(const struct ggml_tensor * tensor, int n) { + size_t next_nb = ggml_type_size(tensor->type); + if (tensor->ne[0] != ggml_blck_size(tensor->type) && tensor->nb[0] != next_nb) { + return false; + } + next_nb *= tensor->ne[0]/ggml_blck_size(tensor->type); + for (int i = 1; i < GGML_MAX_DIMS; i++) { + if (i > n) { + if (tensor->ne[i] != 1 && tensor->nb[i] != next_nb) { + return false; + } + next_nb *= tensor->ne[i]; + } else { + // this dimension does not need to be contiguous + next_nb = tensor->ne[i]*tensor->nb[i]; + } + } + return true; +} + +bool ggml_is_contiguous(const struct ggml_tensor * tensor) { + return ggml_is_contiguous_0(tensor); +} + +bool ggml_is_contiguous_0(const struct ggml_tensor * tensor) { + return ggml_is_contiguous_n(tensor, 0); +} + +bool ggml_is_contiguous_1(const struct ggml_tensor * tensor) { + return ggml_is_contiguous_n(tensor, 1); +} + +bool ggml_is_contiguous_2(const struct ggml_tensor * tensor) { + return ggml_is_contiguous_n(tensor, 2); +} + +bool ggml_is_contiguously_allocated(const struct ggml_tensor * tensor) { + return ggml_nbytes(tensor) == ggml_nelements(tensor) * ggml_type_size(tensor->type)/ggml_blck_size(tensor->type); +} + +bool ggml_is_permuted(const struct ggml_tensor * tensor) { + static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); + + return tensor->nb[0] > tensor->nb[1] || tensor->nb[1] > tensor->nb[2] || tensor->nb[2] > tensor->nb[3]; +} + +bool ggml_is_contiguous_channels(const struct ggml_tensor * tensor) { + return + tensor->nb[0] > tensor->nb[2] && + tensor->nb[1] > tensor->nb[0] && + tensor->nb[2] == ggml_type_size(tensor->type); +} + +bool ggml_is_contiguous_rows(const struct ggml_tensor * tensor) { + return + tensor->ne[0] == ggml_blck_size(tensor->type) || + tensor->nb[0] == ggml_type_size(tensor->type); +} + +static inline bool ggml_is_padded_1d(const struct ggml_tensor * tensor) { + static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); + + return + tensor->nb[0] == ggml_type_size(tensor->type) && + tensor->nb[2] == tensor->nb[1]*tensor->ne[1] && + tensor->nb[3] == tensor->nb[2]*tensor->ne[2]; +} + +bool ggml_is_empty(const struct ggml_tensor * tensor) { + for (int i = 0; i < GGML_MAX_DIMS; ++i) { + if (tensor->ne[i] == 0) { + // empty if any dimension has no elements + return true; + } + } + return false; +} + +bool ggml_are_same_shape(const struct ggml_tensor * t0, const struct ggml_tensor * t1) { + static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); + + return + (t0->ne[0] == t1->ne[0]) && + (t0->ne[1] == t1->ne[1]) && + (t0->ne[2] == t1->ne[2]) && + (t0->ne[3] == t1->ne[3]); +} + +bool ggml_are_same_stride(const struct ggml_tensor * t0, const struct ggml_tensor * t1) { + static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); + + return + (t0->nb[0] == t1->nb[0]) && + (t0->nb[1] == t1->nb[1]) && + (t0->nb[2] == t1->nb[2]) && + (t0->nb[3] == t1->nb[3]); +} + +bool ggml_is_view(const struct ggml_tensor * t) { + return ggml_impl_is_view(t); +} + +// check if t1 can be represented as a repetition of t0 +bool ggml_can_repeat(const struct ggml_tensor * t0, const struct ggml_tensor * t1) { + static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); + + return ggml_is_empty(t0) ? ggml_is_empty(t1) : + (t1->ne[0]%t0->ne[0] == 0) && + (t1->ne[1]%t0->ne[1] == 0) && + (t1->ne[2]%t0->ne[2] == 0) && + (t1->ne[3]%t0->ne[3] == 0); +} + +static inline bool ggml_can_repeat_rows(const struct ggml_tensor * t0, const struct ggml_tensor * t1) { + static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); + + return (t0->ne[0] == t1->ne[0]) && ggml_can_repeat(t0, t1); +} + +// assert that pointer is aligned to GGML_MEM_ALIGN +#define GGML_ASSERT_ALIGNED(ptr) \ + GGML_ASSERT(((uintptr_t) (ptr))%GGML_MEM_ALIGN == 0) + +//////////////////////////////////////////////////////////////////////////////// + +struct ggml_context * ggml_init(struct ggml_init_params params) { + static bool is_first_call = true; + + ggml_critical_section_start(); + + if (is_first_call) { + // initialize time system (required on Windows) + ggml_time_init(); + + is_first_call = false; + } + + ggml_critical_section_end(); + + struct ggml_context * ctx = GGML_MALLOC(sizeof(struct ggml_context)); + + // allow to call ggml_init with 0 size + if (params.mem_size == 0) { + params.mem_size = GGML_MEM_ALIGN; + } + + const size_t mem_size = params.mem_buffer ? params.mem_size : GGML_PAD(params.mem_size, GGML_MEM_ALIGN); + + *ctx = (struct ggml_context) { + /*.mem_size =*/ mem_size, + /*.mem_buffer =*/ params.mem_buffer ? params.mem_buffer : ggml_aligned_malloc(mem_size), + /*.mem_buffer_owned =*/ params.mem_buffer ? false : true, + /*.no_alloc =*/ params.no_alloc, + /*.n_objects =*/ 0, + /*.objects_begin =*/ NULL, + /*.objects_end =*/ NULL, + }; + + GGML_ASSERT(ctx->mem_buffer != NULL); + + GGML_ASSERT_ALIGNED(ctx->mem_buffer); + + GGML_PRINT_DEBUG("%s: context initialized\n", __func__); + + return ctx; +} + +void ggml_reset(struct ggml_context * ctx) { + if (ctx == NULL) { + return; + } + + ctx->n_objects = 0; + ctx->objects_begin = NULL; + ctx->objects_end = NULL; +} + +void ggml_free(struct ggml_context * ctx) { + if (ctx == NULL) { + return; + } + + if (ctx->mem_buffer_owned) { + ggml_aligned_free(ctx->mem_buffer, ctx->mem_size); + } + + GGML_FREE(ctx); +} + +size_t ggml_used_mem(const struct ggml_context * ctx) { + return ctx->objects_end == NULL ? 0 : ctx->objects_end->offs + ctx->objects_end->size; +} + +bool ggml_get_no_alloc(struct ggml_context * ctx) { + return ctx->no_alloc; +} + +void ggml_set_no_alloc(struct ggml_context * ctx, bool no_alloc) { + ctx->no_alloc = no_alloc; +} + +void * ggml_get_mem_buffer(const struct ggml_context * ctx) { + return ctx->mem_buffer; +} + +size_t ggml_get_mem_size(const struct ggml_context * ctx) { + return ctx->mem_size; +} + +size_t ggml_get_max_tensor_size(const struct ggml_context * ctx) { + size_t max_size = 0; + + for (struct ggml_tensor * tensor = ggml_get_first_tensor(ctx); tensor != NULL; tensor = ggml_get_next_tensor(ctx, tensor)) { + size_t bytes = ggml_nbytes(tensor); + max_size = MAX(max_size, bytes); + } + + return max_size; +} + +//////////////////////////////////////////////////////////////////////////////// + +static struct ggml_object * ggml_new_object(struct ggml_context * ctx, enum ggml_object_type type, size_t size) { + // always insert objects at the end of the context's memory pool + struct ggml_object * obj_cur = ctx->objects_end; + + const size_t cur_offs = obj_cur == NULL ? 0 : obj_cur->offs; + const size_t cur_size = obj_cur == NULL ? 0 : obj_cur->size; + const size_t cur_end = cur_offs + cur_size; + + // align to GGML_MEM_ALIGN + GGML_ASSERT(size <= SIZE_MAX - (GGML_MEM_ALIGN - 1)); + size_t size_needed = GGML_PAD(size, GGML_MEM_ALIGN); + + char * const mem_buffer = ctx->mem_buffer; + struct ggml_object * const obj_new = (struct ggml_object *)(mem_buffer + cur_end); + + // integer overflow checks + if (cur_end > SIZE_MAX - size_needed) { + GGML_LOG_WARN("%s: overflow detected in cur_end (%zu) + size_needed (%zu)\n", __func__, cur_end, size_needed); + return NULL; + } + if (cur_end + size_needed > SIZE_MAX - GGML_OBJECT_SIZE) { + GGML_LOG_WARN("%s: overflow detected in cur_end (%zu) + size_needed (%zu) + GGML_OBJECT_SIZE (%zu)\n", __func__, + cur_end, size_needed, (size_t) GGML_OBJECT_SIZE); + return NULL; + } + + if (cur_end + size_needed + GGML_OBJECT_SIZE > ctx->mem_size) { + GGML_LOG_WARN("%s: not enough space in the context's memory pool (needed %zu, available %zu)\n", + __func__, cur_end + size_needed + GGML_OBJECT_SIZE, ctx->mem_size); +#ifndef NDEBUG + GGML_ABORT("not enough space in the context's memory pool"); +#endif + return NULL; + } + + *obj_new = (struct ggml_object) { + .offs = cur_end + GGML_OBJECT_SIZE, + .size = size_needed, + .next = NULL, + .type = type, + }; + + GGML_ASSERT_ALIGNED(mem_buffer + obj_new->offs); + + if (obj_cur != NULL) { + obj_cur->next = obj_new; + } else { + // this is the first object in this context + ctx->objects_begin = obj_new; + } + + ctx->objects_end = obj_new; + + //printf("%s: inserted new object at %zu, size = %zu\n", __func__, cur_end, obj_new->size); + + return obj_new; +} + +static struct ggml_tensor * ggml_new_tensor_impl( + struct ggml_context * ctx, + enum ggml_type type, + int n_dims, + const int64_t * ne, + struct ggml_tensor * view_src, + size_t view_offs) { + + GGML_ASSERT(type >= 0 && type < GGML_TYPE_COUNT); + GGML_ASSERT(n_dims >= 1 && n_dims <= GGML_MAX_DIMS); + + // find the base tensor and absolute offset + if (view_src != NULL && view_src->view_src != NULL) { + view_offs += view_src->view_offs; + view_src = view_src->view_src; + } + + size_t data_size = ggml_row_size(type, ne[0]); + for (int i = 1; i < n_dims; i++) { + data_size *= ne[i]; + } + + GGML_ASSERT(view_src == NULL || data_size == 0 || data_size + view_offs <= ggml_nbytes(view_src)); + + void * data = view_src != NULL ? view_src->data : NULL; + if (data != NULL) { + data = (char *) data + view_offs; + } + + size_t obj_alloc_size = 0; + + if (view_src == NULL && !ctx->no_alloc) { + // allocate tensor data in the context's memory pool + obj_alloc_size = data_size; + } + + GGML_ASSERT(GGML_TENSOR_SIZE <= SIZE_MAX - obj_alloc_size); + + struct ggml_object * const obj_new = ggml_new_object(ctx, GGML_OBJECT_TYPE_TENSOR, GGML_TENSOR_SIZE + obj_alloc_size); + GGML_ASSERT(obj_new); + + struct ggml_tensor * const result = (struct ggml_tensor *)((char *)ctx->mem_buffer + obj_new->offs); + + *result = (struct ggml_tensor) { + /*.type =*/ type, + /*.buffer =*/ NULL, + /*.ne =*/ { 1, 1, 1, 1 }, + /*.nb =*/ { 0, 0, 0, 0 }, + /*.op =*/ GGML_OP_NONE, + /*.op_params =*/ { 0 }, + /*.flags =*/ 0, + /*.src =*/ { NULL }, + /*.view_src =*/ view_src, + /*.view_offs =*/ view_offs, + /*.data =*/ obj_alloc_size > 0 ? (void *)(result + 1) : data, + /*.name =*/ { 0 }, + /*.extra =*/ NULL, + /*.padding =*/ { 0 }, + }; + + // TODO: this should not be needed as long as we don't rely on aligned SIMD loads + //GGML_ASSERT_ALIGNED(result->data); + + for (int i = 0; i < n_dims; i++) { + result->ne[i] = ne[i]; + } + + result->nb[0] = ggml_type_size(type); + result->nb[1] = result->nb[0]*(result->ne[0]/ggml_blck_size(type)); + for (int i = 2; i < GGML_MAX_DIMS; i++) { + result->nb[i] = result->nb[i - 1]*result->ne[i - 1]; + } + + ctx->n_objects++; + + return result; +} + +struct ggml_tensor * ggml_new_tensor( + struct ggml_context * ctx, + enum ggml_type type, + int n_dims, + const int64_t * ne) { + return ggml_new_tensor_impl(ctx, type, n_dims, ne, NULL, 0); +} + +struct ggml_tensor * ggml_new_tensor_1d( + struct ggml_context * ctx, + enum ggml_type type, + int64_t ne0) { + return ggml_new_tensor(ctx, type, 1, &ne0); +} + +struct ggml_tensor * ggml_new_tensor_2d( + struct ggml_context * ctx, + enum ggml_type type, + int64_t ne0, + int64_t ne1) { + const int64_t ne[2] = { ne0, ne1 }; + return ggml_new_tensor(ctx, type, 2, ne); +} + +struct ggml_tensor * ggml_new_tensor_3d( + struct ggml_context * ctx, + enum ggml_type type, + int64_t ne0, + int64_t ne1, + int64_t ne2) { + const int64_t ne[3] = { ne0, ne1, ne2 }; + return ggml_new_tensor(ctx, type, 3, ne); +} + +struct ggml_tensor * ggml_new_tensor_4d( + struct ggml_context * ctx, + enum ggml_type type, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3) { + const int64_t ne[4] = { ne0, ne1, ne2, ne3 }; + return ggml_new_tensor(ctx, type, 4, ne); +} + +void * ggml_new_buffer(struct ggml_context * ctx, size_t nbytes) { + struct ggml_object * obj = ggml_new_object(ctx, GGML_OBJECT_TYPE_WORK_BUFFER, nbytes); + + return (uint8_t *)ctx->mem_buffer + obj->offs; +} + +struct ggml_tensor * ggml_dup_tensor(struct ggml_context * ctx, const struct ggml_tensor * src) { + return ggml_new_tensor(ctx, src->type, GGML_MAX_DIMS, src->ne); +} + +void ggml_unravel_index(const struct ggml_tensor * tensor, int64_t i, int64_t * i0, int64_t * i1, int64_t * i2, int64_t * i3) { + const int64_t ne2 = tensor->ne[2]; + const int64_t ne1 = tensor->ne[1]; + const int64_t ne0 = tensor->ne[0]; + + const int64_t i3_ = (i/(ne2*ne1*ne0)); + const int64_t i2_ = (i - i3_*ne2*ne1*ne0)/(ne1*ne0); + const int64_t i1_ = (i - i3_*ne2*ne1*ne0 - i2_*ne1*ne0)/ne0; + const int64_t i0_ = (i - i3_*ne2*ne1*ne0 - i2_*ne1*ne0 - i1_*ne0); + + if (i0) { + * i0 = i0_; + } + if (i1) { + * i1 = i1_; + } + if (i2) { + * i2 = i2_; + } + if (i3) { + * i3 = i3_; + } +} + +void * ggml_get_data(const struct ggml_tensor * tensor) { + return tensor->data; +} + +float * ggml_get_data_f32(const struct ggml_tensor * tensor) { + assert(tensor->type == GGML_TYPE_F32); + return (float *)(tensor->data); +} + +enum ggml_unary_op ggml_get_unary_op(const struct ggml_tensor * tensor) { + GGML_ASSERT(tensor->op == GGML_OP_UNARY); + return (enum ggml_unary_op) ggml_get_op_params_i32(tensor, 0); +} + +enum ggml_glu_op ggml_get_glu_op(const struct ggml_tensor * tensor) { + GGML_ASSERT(tensor->op == GGML_OP_GLU); + return (enum ggml_glu_op) ggml_get_op_params_i32(tensor, 0); +} + +const char * ggml_get_name(const struct ggml_tensor * tensor) { + return tensor->name; +} + +struct ggml_tensor * ggml_set_name(struct ggml_tensor * tensor, const char * name) { + size_t i; + for (i = 0; i < sizeof(tensor->name) - 1 && name[i] != '\0'; i++) { + tensor->name[i] = name[i]; + } + tensor->name[i] = '\0'; + return tensor; +} + +struct ggml_tensor * ggml_format_name(struct ggml_tensor * tensor, const char * fmt, ...) { + va_list args; + va_start(args, fmt); + vsnprintf(tensor->name, sizeof(tensor->name), fmt, args); + va_end(args); + return tensor; +} + +struct ggml_tensor * ggml_view_tensor( + struct ggml_context * ctx, + struct ggml_tensor * src) { + struct ggml_tensor * result = ggml_new_tensor_impl(ctx, src->type, GGML_MAX_DIMS, src->ne, src, 0); + ggml_format_name(result, "%s (view)", src->name); + + for (int i = 0; i < GGML_MAX_DIMS; i++) { + result->nb[i] = src->nb[i]; + } + + return result; +} + +struct ggml_tensor * ggml_get_first_tensor(const struct ggml_context * ctx) { + struct ggml_object * obj = ctx->objects_begin; + + char * const mem_buffer = ctx->mem_buffer; + + while (obj != NULL) { + if (obj->type == GGML_OBJECT_TYPE_TENSOR) { + return (struct ggml_tensor *)(mem_buffer + obj->offs); + } + + obj = obj->next; + } + + return NULL; +} + +struct ggml_tensor * ggml_get_next_tensor(const struct ggml_context * ctx, struct ggml_tensor * tensor) { + struct ggml_object * obj = (struct ggml_object *) ((char *)tensor - GGML_OBJECT_SIZE); + obj = obj->next; + + char * const mem_buffer = ctx->mem_buffer; + + while (obj != NULL) { + if (obj->type == GGML_OBJECT_TYPE_TENSOR) { + return (struct ggml_tensor *)(mem_buffer + obj->offs); + } + + obj = obj->next; + } + + return NULL; +} + +struct ggml_tensor * ggml_get_tensor(struct ggml_context * ctx, const char * name) { + struct ggml_object * obj = ctx->objects_begin; + + char * const mem_buffer = ctx->mem_buffer; + + while (obj != NULL) { + if (obj->type == GGML_OBJECT_TYPE_TENSOR) { + struct ggml_tensor * cur = (struct ggml_tensor *)(mem_buffer + obj->offs); + if (strcmp(cur->name, name) == 0) { + return cur; + } + } + + obj = obj->next; + } + + return NULL; +} + +//////////////////////////////////////////////////////////////////////////////// + +// ggml_dup + +static struct ggml_tensor * ggml_dup_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + bool inplace) { + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + result->op = GGML_OP_DUP; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_dup( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_dup_impl(ctx, a, false); +} + +struct ggml_tensor * ggml_dup_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_dup_impl(ctx, a, true); +} + +// ggml_add + +static struct ggml_tensor * ggml_add_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + bool inplace) { + GGML_ASSERT(ggml_can_repeat(b, a)); + + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + result->op = GGML_OP_ADD; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +struct ggml_tensor * ggml_add( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_add_impl(ctx, a, b, false); +} + +struct ggml_tensor * ggml_add_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_add_impl(ctx, a, b, true); +} + +// ggml_add_cast + +static struct ggml_tensor * ggml_add_cast_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + enum ggml_type type) { + // TODO: support less-strict constraint + // GGML_ASSERT(ggml_can_repeat(b, a)); + GGML_ASSERT(ggml_can_repeat_rows(b, a)); + + // currently only supported for quantized input and f16 + GGML_ASSERT(ggml_is_quantized(a->type) || + a->type == GGML_TYPE_F16 || + a->type == GGML_TYPE_BF16); + + struct ggml_tensor * result = ggml_new_tensor(ctx, type, GGML_MAX_DIMS, a->ne); + + result->op = GGML_OP_ADD; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +struct ggml_tensor * ggml_add_cast( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + enum ggml_type type) { + return ggml_add_cast_impl(ctx, a, b, type); +} + +struct ggml_tensor * ggml_add_id( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * ids) { + + GGML_ASSERT(a->ne[0] == b->ne[0]); + GGML_ASSERT(a->ne[1] == ids->ne[0]); + GGML_ASSERT(a->ne[2] == ids->ne[1]); + GGML_ASSERT(ids->type == GGML_TYPE_I32); + + struct ggml_tensor * result = ggml_dup_tensor(ctx, a); + + result->op = GGML_OP_ADD_ID; + result->src[0] = a; + result->src[1] = b; + result->src[2] = ids; + + return result; +} + +// ggml_add1 + +static struct ggml_tensor * ggml_add1_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + bool inplace) { + GGML_ASSERT(ggml_is_scalar(b)); + GGML_ASSERT(ggml_is_padded_1d(a)); + + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + result->op = GGML_OP_ADD1; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +struct ggml_tensor * ggml_add1( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_add1_impl(ctx, a, b, false); +} + +struct ggml_tensor * ggml_add1_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_add1_impl(ctx, a, b, true); +} + +// ggml_acc + +static struct ggml_tensor * ggml_acc_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t nb1, + size_t nb2, + size_t nb3, + size_t offset, + bool inplace) { + GGML_ASSERT(ggml_nelements(b) <= ggml_nelements(a)); + GGML_ASSERT(ggml_is_contiguous(a)); + GGML_ASSERT(a->type == GGML_TYPE_F32); + GGML_ASSERT(b->type == GGML_TYPE_F32); + + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + int32_t params[] = { nb1, nb2, nb3, offset, inplace ? 1 : 0 }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_ACC; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +struct ggml_tensor * ggml_acc( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t nb1, + size_t nb2, + size_t nb3, + size_t offset) { + return ggml_acc_impl(ctx, a, b, nb1, nb2, nb3, offset, false); +} + +struct ggml_tensor * ggml_acc_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t nb1, + size_t nb2, + size_t nb3, + size_t offset) { + return ggml_acc_impl(ctx, a, b, nb1, nb2, nb3, offset, true); +} + +// ggml_sub + +static struct ggml_tensor * ggml_sub_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + bool inplace) { + GGML_ASSERT(ggml_can_repeat(b, a)); + + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + result->op = GGML_OP_SUB; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +struct ggml_tensor * ggml_sub( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_sub_impl(ctx, a, b, false); +} + +struct ggml_tensor * ggml_sub_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_sub_impl(ctx, a, b, true); +} + +// ggml_mul + +static struct ggml_tensor * ggml_mul_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + bool inplace) { + GGML_ASSERT(ggml_can_repeat(b, a)); + + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + result->op = GGML_OP_MUL; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +struct ggml_tensor * ggml_mul( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_mul_impl(ctx, a, b, false); +} + +struct ggml_tensor * ggml_mul_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_mul_impl(ctx, a, b, true); +} + +// ggml_div + +static struct ggml_tensor * ggml_div_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + bool inplace) { + GGML_ASSERT(ggml_can_repeat(b, a)); + + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + result->op = GGML_OP_DIV; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +struct ggml_tensor * ggml_div( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_div_impl(ctx, a, b, false); +} + +struct ggml_tensor * ggml_div_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_div_impl(ctx, a, b, true); +} + +// ggml_sqr + +static struct ggml_tensor * ggml_sqr_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + bool inplace) { + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + result->op = GGML_OP_SQR; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_sqr( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_sqr_impl(ctx, a, false); +} + +struct ggml_tensor * ggml_sqr_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_sqr_impl(ctx, a, true); +} + +// ggml_sqrt + +static struct ggml_tensor * ggml_sqrt_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + bool inplace) { + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + result->op = GGML_OP_SQRT; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_sqrt( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_sqrt_impl(ctx, a, false); +} + +struct ggml_tensor * ggml_sqrt_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_sqrt_impl(ctx, a, true); +} + +// ggml_log + +static struct ggml_tensor * ggml_log_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + bool inplace) { + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + result->op = GGML_OP_LOG; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_log( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_log_impl(ctx, a, false); +} + +struct ggml_tensor * ggml_log_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_log_impl(ctx, a, true); +} + +struct ggml_tensor * ggml_expm1( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_EXPM1); +} + +struct ggml_tensor * ggml_expm1_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_EXPM1); +} + +struct ggml_tensor * ggml_softplus( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_SOFTPLUS); +} + +struct ggml_tensor * ggml_softplus_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_SOFTPLUS); +} + +// ggml_sin + +static struct ggml_tensor * ggml_sin_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + bool inplace) { + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + result->op = GGML_OP_SIN; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_sin( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_sin_impl(ctx, a, false); +} + +struct ggml_tensor * ggml_sin_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_sin_impl(ctx, a, true); +} + +// ggml_cos + +static struct ggml_tensor * ggml_cos_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + bool inplace) { + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + result->op = GGML_OP_COS; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_cos( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_cos_impl(ctx, a, false); +} + +struct ggml_tensor * ggml_cos_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_cos_impl(ctx, a, true); +} + +// ggml_sum + +struct ggml_tensor * ggml_sum( + struct ggml_context * ctx, + struct ggml_tensor * a) { + struct ggml_tensor * result = ggml_new_tensor_1d(ctx, a->type, 1); + + result->op = GGML_OP_SUM; + result->src[0] = a; + + return result; +} + +// ggml_sum_rows + +struct ggml_tensor * ggml_sum_rows( + struct ggml_context * ctx, + struct ggml_tensor * a) { + int64_t ne[GGML_MAX_DIMS] = { 1 }; + for (int i = 1; i < GGML_MAX_DIMS; ++i) { + ne[i] = a->ne[i]; + } + + struct ggml_tensor * result = ggml_new_tensor(ctx, a->type, GGML_MAX_DIMS, ne); + + result->op = GGML_OP_SUM_ROWS; + result->src[0] = a; + + return result; +} + +// ggml_cumsum + +struct ggml_tensor * ggml_cumsum( + struct ggml_context * ctx, + struct ggml_tensor * a) { + GGML_ASSERT(a->type == GGML_TYPE_F32); + + struct ggml_tensor * result = ggml_dup_tensor(ctx, a); + + result->op = GGML_OP_CUMSUM; + result->src[0] = a; + + return result; +} + +// ggml_mean + +struct ggml_tensor * ggml_mean( + struct ggml_context * ctx, + struct ggml_tensor * a) { + int64_t ne[4] = { 1, a->ne[1], a->ne[2], a->ne[3] }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + result->op = GGML_OP_MEAN; + result->src[0] = a; + + return result; +} + +// ggml_argmax + +struct ggml_tensor * ggml_argmax( + struct ggml_context * ctx, + struct ggml_tensor * a) { + GGML_ASSERT(ggml_is_matrix(a)); + GGML_ASSERT(a->ne[0] <= INT32_MAX); + + struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, a->ne[1]); + + result->op = GGML_OP_ARGMAX; + result->src[0] = a; + + return result; +} + +// ggml_count_equal + +struct ggml_tensor * ggml_count_equal( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + GGML_ASSERT(ggml_are_same_shape(a, b)); + + struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, 1); + + result->op = GGML_OP_COUNT_EQUAL; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +// ggml_repeat + +struct ggml_tensor * ggml_repeat( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + GGML_ASSERT(ggml_can_repeat(a, b)); + + struct ggml_tensor * result = ggml_new_tensor(ctx, a->type, GGML_MAX_DIMS, b->ne); + + result->op = GGML_OP_REPEAT; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_repeat_4d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3) { + const bool can_repeat = ggml_is_empty(a) || ( + (ne0 % a->ne[0] == 0) && + (ne1 % a->ne[1] == 0) && + (ne2 % a->ne[2] == 0) && + (ne3 % a->ne[3] == 0) + ); + GGML_ASSERT(can_repeat); + + struct ggml_tensor * result = ggml_new_tensor_4d(ctx, a->type, ne0, ne1, ne2, ne3); + + result->op = GGML_OP_REPEAT; + result->src[0] = a; + + return result; +} + +// ggml_repeat_back + +struct ggml_tensor * ggml_repeat_back( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + GGML_ASSERT(ggml_can_repeat(b, a)); + + struct ggml_tensor * result = ggml_new_tensor(ctx, a->type, GGML_MAX_DIMS, b->ne); + + result->op = GGML_OP_REPEAT_BACK; + result->src[0] = a; + + return result; +} + +// ggml_concat + +struct ggml_tensor * ggml_concat( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int dim) { + GGML_ASSERT(dim >= 0 && dim < GGML_MAX_DIMS); + GGML_ASSERT(a->type == b->type); + + int64_t ne[GGML_MAX_DIMS]; + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + if (d == dim) { + ne[d] = a->ne[d] + b->ne[d]; + continue; + } + GGML_ASSERT(a->ne[d] == b->ne[d]); + ne[d] = a->ne[d]; + } + + struct ggml_tensor * result = ggml_new_tensor(ctx, a->type, GGML_MAX_DIMS, ne); + + ggml_set_op_params_i32(result, 0, dim); + + result->op = GGML_OP_CONCAT; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +// ggml_abs + +struct ggml_tensor * ggml_abs( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_ABS); +} + +struct ggml_tensor * ggml_abs_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_ABS); +} + +// ggml_sgn + +struct ggml_tensor * ggml_sgn( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_SGN); +} + +struct ggml_tensor * ggml_sgn_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_SGN); +} + +// ggml_neg + +struct ggml_tensor * ggml_neg( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_NEG); +} + +struct ggml_tensor * ggml_neg_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_NEG); +} + +// ggml_step + +struct ggml_tensor * ggml_step( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_STEP); +} + +struct ggml_tensor * ggml_step_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_STEP); +} + +// ggml_tanh + +struct ggml_tensor * ggml_tanh( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_TANH); +} + +struct ggml_tensor * ggml_tanh_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_TANH); +} + +// ggml_elu + +struct ggml_tensor * ggml_elu( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_ELU); +} + +struct ggml_tensor * ggml_elu_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_ELU); +} + +// ggml_relu + +struct ggml_tensor * ggml_relu( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_RELU); +} + +struct ggml_tensor * ggml_relu_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_RELU); +} + +// ggml_leaky_relu + +struct ggml_tensor * ggml_leaky_relu( + struct ggml_context * ctx, + struct ggml_tensor * a, + float negative_slope, + bool inplace) { + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + ggml_set_op_params(result, &negative_slope, sizeof(negative_slope)); + + result->op = GGML_OP_LEAKY_RELU; + result->src[0] = a; + + return result; +} + +// ggml_sigmoid + +struct ggml_tensor * ggml_sigmoid( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_SIGMOID); +} + +struct ggml_tensor * ggml_sigmoid_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_SIGMOID); +} + +// ggml_gelu + +struct ggml_tensor * ggml_gelu( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_GELU); +} + +struct ggml_tensor * ggml_gelu_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_GELU); +} + +// ggml_gelu_erf + +struct ggml_tensor * ggml_gelu_erf( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_GELU_ERF); +} + +struct ggml_tensor * ggml_gelu_erf_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_GELU_ERF); +} + +// ggml_gelu_quick + +struct ggml_tensor * ggml_gelu_quick( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_GELU_QUICK); +} + +struct ggml_tensor * ggml_gelu_quick_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_GELU_QUICK); +} + +// ggml_silu + +struct ggml_tensor * ggml_silu( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_SILU); +} + +struct ggml_tensor * ggml_silu_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_SILU); +} + +// ggml_xielu + +struct ggml_tensor * ggml_xielu( + struct ggml_context * ctx, + struct ggml_tensor * a, + float alpha_n, + float alpha_p, + float beta, + float eps) { + struct ggml_tensor * result = ggml_dup_tensor(ctx, a); + + ggml_set_op_params_i32(result, 0, (int32_t) GGML_UNARY_OP_XIELU); + ggml_set_op_params_f32(result, 1, beta + ggml_compute_softplus_f32(alpha_n)); + ggml_set_op_params_f32(result, 2, ggml_compute_softplus_f32(alpha_p)); + ggml_set_op_params_f32(result, 3, beta); + ggml_set_op_params_f32(result, 4, eps); + + result->op = GGML_OP_UNARY; + result->src[0] = a; + + return result; +} + +// ggml_silu_back + +struct ggml_tensor * ggml_silu_back( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + struct ggml_tensor * result = ggml_dup_tensor(ctx, a); + + result->op = GGML_OP_SILU_BACK; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +// ggml hardswish + +struct ggml_tensor * ggml_hardswish( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_HARDSWISH); +} + +// ggml hardsigmoid + +struct ggml_tensor * ggml_hardsigmoid( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_HARDSIGMOID); +} + +// ggml exp + +struct ggml_tensor * ggml_exp( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_EXP); +} + +struct ggml_tensor * ggml_exp_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_EXP); +} + +// ggml_glu + +static struct ggml_tensor * ggml_glu_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + enum ggml_glu_op op, + bool swapped) { + GGML_ASSERT(ggml_is_contiguous_1(a)); + + if (b) { + GGML_ASSERT(ggml_is_contiguous_1(b)); + GGML_ASSERT(ggml_are_same_shape(a, b)); + GGML_ASSERT(a->type == b->type); + } + + int64_t ne[GGML_MAX_DIMS] = { a->ne[0] / 2 }; for (int i = 1; i < GGML_MAX_DIMS; i++) ne[i] = a->ne[i]; + struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, GGML_MAX_DIMS, b ? a->ne : ne, NULL, 0); + + ggml_set_op_params_i32(result, 0, (int32_t) op); + ggml_set_op_params_i32(result, 1, (int32_t) swapped); + + result->op = GGML_OP_GLU; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +// ggml_floor + +struct ggml_tensor * ggml_floor( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_FLOOR); +} + +struct ggml_tensor * ggml_floor_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_FLOOR); +} + +// ggml_ceil + +struct ggml_tensor * ggml_ceil( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_CEIL); +} + +struct ggml_tensor * ggml_ceil_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_CEIL); +} + +//ggml_round + +struct ggml_tensor * ggml_round( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_ROUND); +} + +struct ggml_tensor * ggml_round_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_ROUND); +} + +//ggml_trunc + +struct ggml_tensor * ggml_trunc( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary(ctx, a, GGML_UNARY_OP_TRUNC); +} + +struct ggml_tensor * ggml_trunc_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_TRUNC); +} + +struct ggml_tensor * ggml_glu( + struct ggml_context * ctx, + struct ggml_tensor * a, + enum ggml_glu_op op, + bool swapped) { + return ggml_glu_impl(ctx, a, NULL, op, swapped); +} + +struct ggml_tensor * ggml_glu_split( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + enum ggml_glu_op op) { + return ggml_glu_impl(ctx, a, b, op, false); +} + +// ggml_reglu + +struct ggml_tensor * ggml_reglu( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_REGLU, false); +} + +struct ggml_tensor * ggml_reglu_swapped( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_REGLU, true); +} + +struct ggml_tensor * ggml_reglu_split( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_glu_impl(ctx, a, b, GGML_GLU_OP_REGLU, false); +} + +// ggml_geglu + +struct ggml_tensor * ggml_geglu( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_GEGLU, false); +} + +struct ggml_tensor * ggml_geglu_swapped( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_GEGLU, true); +} + +struct ggml_tensor * ggml_geglu_split( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_glu_impl(ctx, a, b, GGML_GLU_OP_GEGLU, false); +} + +// ggml_swiglu + +struct ggml_tensor * ggml_swiglu( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_SWIGLU, false); +} + +struct ggml_tensor * ggml_swiglu_swapped( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_SWIGLU, true); +} + +struct ggml_tensor * ggml_swiglu_split( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_glu_impl(ctx, a, b, GGML_GLU_OP_SWIGLU, false); +} + +struct ggml_tensor * ggml_swiglu_ds4_split( + struct ggml_context * ctx, + struct ggml_tensor * gate, + struct ggml_tensor * up, + float clamp) { + struct ggml_tensor * result = ggml_glu_impl(ctx, gate, up, GGML_GLU_OP_SWIGLU_DS4, false); + ggml_set_op_params_f32(result, 2, clamp); + return result; +} + +// ggml_geglu_erf + +struct ggml_tensor * ggml_geglu_erf( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_GEGLU_ERF, false); +} + +struct ggml_tensor * ggml_geglu_erf_swapped( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_GEGLU_ERF, true); +} + +struct ggml_tensor * ggml_geglu_erf_split( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_glu_impl(ctx, a, b, GGML_GLU_OP_GEGLU_ERF, false); +} + +// ggml_geglu_quick + +struct ggml_tensor * ggml_geglu_quick( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_GEGLU_QUICK, false); +} + +struct ggml_tensor * ggml_geglu_quick_swapped( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_GEGLU_QUICK, true); +} + +struct ggml_tensor * ggml_geglu_quick_split( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_glu_impl(ctx, a, b, GGML_GLU_OP_GEGLU_QUICK, false); +} + +struct ggml_tensor * ggml_swiglu_oai( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + float alpha, + float limit) { + struct ggml_tensor * result = ggml_glu_impl(ctx, a, b, GGML_GLU_OP_SWIGLU_OAI, false); + ggml_set_op_params_f32(result, 2, alpha); + ggml_set_op_params_f32(result, 3, limit); + + return result; +} + +// ggml_norm + +static struct ggml_tensor * ggml_norm_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + float eps, + bool inplace) { + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + ggml_set_op_params(result, &eps, sizeof(eps)); + + result->op = GGML_OP_NORM; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_norm( + struct ggml_context * ctx, + struct ggml_tensor * a, + float eps) { + return ggml_norm_impl(ctx, a, eps, false); +} + +struct ggml_tensor * ggml_norm_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + float eps) { + return ggml_norm_impl(ctx, a, eps, true); +} + +// ggml_rms_norm + +static struct ggml_tensor * ggml_rms_norm_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + float eps, + bool inplace) { + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + ggml_set_op_params(result, &eps, sizeof(eps)); + + result->op = GGML_OP_RMS_NORM; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_rms_norm( + struct ggml_context * ctx, + struct ggml_tensor * a, + float eps) { + return ggml_rms_norm_impl(ctx, a, eps, false); +} + +struct ggml_tensor * ggml_rms_norm_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + float eps) { + return ggml_rms_norm_impl(ctx, a, eps, true); +} + +// ggml_rms_norm_back + +struct ggml_tensor * ggml_rms_norm_back( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + float eps) { + struct ggml_tensor * result = ggml_dup_tensor(ctx, a); + + ggml_set_op_params(result, &eps, sizeof(eps)); + + result->op = GGML_OP_RMS_NORM_BACK; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +// ggml_group_norm + +static struct ggml_tensor * ggml_group_norm_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + int n_groups, + float eps, + bool inplace) { + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + ggml_set_op_params_i32(result, 0, n_groups); + ggml_set_op_params_f32(result, 1, eps); + + result->op = GGML_OP_GROUP_NORM; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_group_norm( + struct ggml_context * ctx, + struct ggml_tensor * a, + int n_groups, + float eps) { + return ggml_group_norm_impl(ctx, a, n_groups, eps, false); +} + +struct ggml_tensor * ggml_group_norm_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + int n_groups, + float eps) { + return ggml_group_norm_impl(ctx, a, n_groups, eps, true); +} + +// ggml_l2_norm + +static struct ggml_tensor * ggml_l2_norm_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + float eps, + bool inplace) { + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + ggml_set_op_params_f32(result, 0, eps); + + result->op = GGML_OP_L2_NORM; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_l2_norm( + struct ggml_context * ctx, + struct ggml_tensor * a, + float eps) { + return ggml_l2_norm_impl(ctx, a, eps, false); +} + +struct ggml_tensor * ggml_l2_norm_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + float eps) { + return ggml_l2_norm_impl(ctx, a, eps, true); +} + +// ggml_mul_mat + +static inline bool ggml_can_mul_mat(const struct ggml_tensor * t0, const struct ggml_tensor * t1) { + static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); + + return (t0->ne[0] == t1->ne[0]) && + (t1->ne[2]%t0->ne[2] == 0) && // verify t0 is broadcastable + (t1->ne[3]%t0->ne[3] == 0); +} + +struct ggml_tensor * ggml_mul_mat( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + GGML_ASSERT(ggml_can_mul_mat(a, b)); + GGML_ASSERT(!ggml_is_transposed(a)); + + const int64_t ne[4] = { a->ne[1], b->ne[1], b->ne[2], b->ne[3] }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + result->op = GGML_OP_MUL_MAT; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +void ggml_mul_mat_set_prec( + struct ggml_tensor * a, + enum ggml_prec prec) { + GGML_ASSERT(a->op == GGML_OP_MUL_MAT); + + const int32_t prec_i32 = (int32_t) prec; + + ggml_set_op_params_i32(a, 0, prec_i32); +} + +// ggml_mul_mat_id + +/* + c = ggml_mul_mat_id(ctx, as, b, ids); + + as -> [cols, rows, n_expert] + b -> [cols, n_expert_used, n_tokens] + ids -> [n_expert_used, n_tokens] (i32) + c -> [rows, n_expert_used, n_tokens] + + in b, n_expert_used can be broadcasted to match the n_expert_used of ids + + c ~= as[:,:,i] @ b[:,i%r,t], i = ids[e,t] for all e,t in ids +*/ +struct ggml_tensor * ggml_mul_mat_id( + struct ggml_context * ctx, + struct ggml_tensor * as, + struct ggml_tensor * b, + struct ggml_tensor * ids) { + GGML_ASSERT(!ggml_is_transposed(as)); + GGML_ASSERT(ids->type == GGML_TYPE_I32); + + GGML_ASSERT(as->ne[3] == 1); // as is 3d (one matrix per expert) + GGML_ASSERT(b->ne[3] == 1); // b is 3d + GGML_ASSERT(ids->ne[2] == 1 && ids->ne[3] == 1); // ids is 2d + GGML_ASSERT(ids->ne[1] == b->ne[2]); // must have an expert list per b row + GGML_ASSERT(as->ne[0] == b->ne[0]); // can_mul_mat + GGML_ASSERT(ids->ne[0] % b->ne[1] == 0); // can broadcast + + const int64_t ne[4] = { as->ne[1], ids->ne[0], b->ne[2], 1 }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + result->op = GGML_OP_MUL_MAT_ID; + result->src[0] = as; + result->src[1] = b; + result->src[2] = ids; + + return result; +} + +// ggml_out_prod + +static inline bool ggml_can_out_prod(const struct ggml_tensor * t0, const struct ggml_tensor * t1) { + static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); + + return (t0->ne[1] == t1->ne[1]) && + (t1->ne[2]%t0->ne[2] == 0) && // verify t0 is broadcastable + (t1->ne[3]%t0->ne[3] == 0); +} + +struct ggml_tensor * ggml_out_prod( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + GGML_ASSERT(ggml_can_out_prod(a, b)); + GGML_ASSERT(!ggml_is_transposed(a)); + + // a is broadcastable to b for ne[2] and ne[3] -> use b->ne[2] and b->ne[3] + const int64_t ne[4] = { a->ne[0], b->ne[0], b->ne[2], b->ne[3] }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + result->op = GGML_OP_OUT_PROD; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +// ggml_scale + +static struct ggml_tensor * ggml_scale_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + float s, + float b, + bool inplace) { + GGML_ASSERT(ggml_is_padded_1d(a)); + + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + float params[2] = { s, b }; + ggml_set_op_params(result, ¶ms, sizeof(params)); + + result->op = GGML_OP_SCALE; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_scale( + struct ggml_context * ctx, + struct ggml_tensor * a, + float s) { + return ggml_scale_impl(ctx, a, s, 0.0, false); +} + +struct ggml_tensor * ggml_scale_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + float s) { + return ggml_scale_impl(ctx, a, s, 0.0, true); +} + +struct ggml_tensor * ggml_scale_bias( + struct ggml_context * ctx, + struct ggml_tensor * a, + float s, + float b) { + return ggml_scale_impl(ctx, a, s, b, false); +} + +struct ggml_tensor * ggml_scale_bias_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + float s, + float b) { + return ggml_scale_impl(ctx, a, s, b, true); +} + +// ggml_set + +static struct ggml_tensor * ggml_set_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t nb1, + size_t nb2, + size_t nb3, + size_t offset, + bool inplace) { + GGML_ASSERT(ggml_nelements(a) >= ggml_nelements(b)); + + // make a view of the destination + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + GGML_ASSERT(offset < (size_t)(1 << 30)); + int32_t params[] = { nb1, nb2, nb3, offset, inplace ? 1 : 0 }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_SET; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +struct ggml_tensor * ggml_set( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t nb1, + size_t nb2, + size_t nb3, + size_t offset) { + return ggml_set_impl(ctx, a, b, nb1, nb2, nb3, offset, false); +} + +struct ggml_tensor * ggml_set_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t nb1, + size_t nb2, + size_t nb3, + size_t offset) { + return ggml_set_impl(ctx, a, b, nb1, nb2, nb3, offset, true); +} + +struct ggml_tensor * ggml_set_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t offset) { + return ggml_set_impl(ctx, a, b, a->nb[1], a->nb[2], a->nb[3], offset, false); +} + +struct ggml_tensor * ggml_set_1d_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t offset) { + return ggml_set_impl(ctx, a, b, a->nb[1], a->nb[2], a->nb[3], offset, true); +} + +struct ggml_tensor * ggml_set_2d( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t nb1, + size_t offset) { + return ggml_set_impl(ctx, a, b, nb1, a->nb[2], a->nb[3], offset, false); +} + +struct ggml_tensor * ggml_set_2d_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + size_t nb1, + size_t offset) { + return ggml_set_impl(ctx, a, b, nb1, a->nb[2], a->nb[3], offset, true); +} + +// ggml_cpy + +static struct ggml_tensor * ggml_cpy_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + GGML_ASSERT(ggml_nelements(a) == ggml_nelements(b)); + + // make a view of the destination + struct ggml_tensor * result = ggml_view_tensor(ctx, b); + if (strlen(b->name) > 0) { + ggml_format_name(result, "%s (copy of %s)", b->name, a->name); + } else { + ggml_format_name(result, "%s (copy)", a->name); + } + + result->op = GGML_OP_CPY; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +struct ggml_tensor * ggml_cpy( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_cpy_impl(ctx, a, b); +} + +struct ggml_tensor * ggml_cast( + struct ggml_context * ctx, + struct ggml_tensor * a, + enum ggml_type type) { + struct ggml_tensor * result = ggml_new_tensor(ctx, type, GGML_MAX_DIMS, a->ne); + ggml_format_name(result, "%s (copy)", a->name); + + result->op = GGML_OP_CPY; + result->src[0] = a; + result->src[1] = result; // note: this self-reference might seem redundant, but it's actually needed by some + // backends for consistency with ggml_cpy_impl() above + + return result; +} + +// ggml_cont + +static struct ggml_tensor * ggml_cont_impl( + struct ggml_context * ctx, + struct ggml_tensor * a) { + struct ggml_tensor * result = ggml_dup_tensor(ctx, a); + ggml_format_name(result, "%s (cont)", a->name); + + result->op = GGML_OP_CONT; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_cont( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_cont_impl(ctx, a); +} + +// make contiguous, with new shape +GGML_API struct ggml_tensor * ggml_cont_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0) { + return ggml_cont_4d(ctx, a, ne0, 1, 1, 1); +} + +GGML_API struct ggml_tensor * ggml_cont_2d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1) { + return ggml_cont_4d(ctx, a, ne0, ne1, 1, 1); +} + +GGML_API struct ggml_tensor * ggml_cont_3d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + int64_t ne2) { + return ggml_cont_4d(ctx, a, ne0, ne1, ne2, 1); +} + +struct ggml_tensor * ggml_cont_4d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3) { + GGML_ASSERT(ggml_nelements(a) == (ne0*ne1*ne2*ne3)); + + struct ggml_tensor * result = ggml_new_tensor_4d(ctx, a->type, ne0, ne1, ne2, ne3); + ggml_format_name(result, "%s (cont)", a->name); + + result->op = GGML_OP_CONT; + result->src[0] = a; + + return result; +} + +// ggml_reshape + +struct ggml_tensor * ggml_reshape( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + GGML_ASSERT(ggml_is_contiguous(a)); + // as only the shape of b is relevant, and not its memory layout, b is allowed to be non contiguous. + GGML_ASSERT(ggml_nelements(a) == ggml_nelements(b)); + + struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, GGML_MAX_DIMS, b->ne, a, 0); + ggml_format_name(result, "%s (reshaped)", a->name); + + result->op = GGML_OP_RESHAPE; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_reshape_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0) { + GGML_ASSERT(ggml_is_contiguous(a)); + GGML_ASSERT(ggml_nelements(a) == ne0); + + const int64_t ne[1] = { ne0 }; + struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, 1, ne, a, 0); + ggml_format_name(result, "%s (reshaped)", a->name); + + result->op = GGML_OP_RESHAPE; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_reshape_2d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1) { + GGML_ASSERT(ggml_is_contiguous(a)); + GGML_ASSERT(ggml_nelements(a) == ne0*ne1); + + const int64_t ne[2] = { ne0, ne1 }; + struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, 2, ne, a, 0); + ggml_format_name(result, "%s (reshaped)", a->name); + + result->op = GGML_OP_RESHAPE; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_reshape_3d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + int64_t ne2) { + GGML_ASSERT(ggml_is_contiguous(a)); + GGML_ASSERT(ggml_nelements(a) == ne0*ne1*ne2); + + const int64_t ne[3] = { ne0, ne1, ne2 }; + struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, 3, ne, a, 0); + ggml_format_name(result, "%s (reshaped)", a->name); + + result->op = GGML_OP_RESHAPE; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_reshape_4d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3) { + GGML_ASSERT(ggml_is_contiguous(a)); + GGML_ASSERT(ggml_nelements(a) == ne0*ne1*ne2*ne3); + + const int64_t ne[4] = { ne0, ne1, ne2, ne3 }; + struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, 4, ne, a, 0); + ggml_format_name(result, "%s (reshaped)", a->name); + + result->op = GGML_OP_RESHAPE; + result->src[0] = a; + + return result; +} + +static struct ggml_tensor * ggml_view_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + int n_dims, + const int64_t * ne, + size_t offset) { + struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, n_dims, ne, a, offset); + ggml_format_name(result, "%s (view)", a->name); + + ggml_set_op_params(result, &offset, sizeof(offset)); + + result->op = GGML_OP_VIEW; + result->src[0] = a; + + return result; +} + +// ggml_view_1d + +struct ggml_tensor * ggml_view_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + size_t offset) { + struct ggml_tensor * result = ggml_view_impl(ctx, a, 1, &ne0, offset); + + return result; +} + +// ggml_view_2d + +struct ggml_tensor * ggml_view_2d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + size_t nb1, + size_t offset) { + const int64_t ne[2] = { ne0, ne1 }; + + struct ggml_tensor * result = ggml_view_impl(ctx, a, 2, ne, offset); + + result->nb[1] = nb1; + result->nb[2] = result->nb[1]*ne1; + result->nb[3] = result->nb[2]; + + return result; +} + +// ggml_view_3d + +struct ggml_tensor * ggml_view_3d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + int64_t ne2, + size_t nb1, + size_t nb2, + size_t offset) { + const int64_t ne[3] = { ne0, ne1, ne2 }; + + struct ggml_tensor * result = ggml_view_impl(ctx, a, 3, ne, offset); + + result->nb[1] = nb1; + result->nb[2] = nb2; + result->nb[3] = result->nb[2]*ne2; + + return result; +} + +// ggml_view_4d + +struct ggml_tensor * ggml_view_4d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3, + size_t nb1, + size_t nb2, + size_t nb3, + size_t offset) { + const int64_t ne[4] = { ne0, ne1, ne2, ne3 }; + + struct ggml_tensor * result = ggml_view_impl(ctx, a, 4, ne, offset); + + result->nb[1] = nb1; + result->nb[2] = nb2; + result->nb[3] = nb3; + + return result; +} + +// ggml_permute + +struct ggml_tensor * ggml_permute( + struct ggml_context * ctx, + struct ggml_tensor * a, + int axis0, + int axis1, + int axis2, + int axis3) { + GGML_ASSERT(axis0 >= 0 && axis0 < GGML_MAX_DIMS); + GGML_ASSERT(axis1 >= 0 && axis1 < GGML_MAX_DIMS); + GGML_ASSERT(axis2 >= 0 && axis2 < GGML_MAX_DIMS); + GGML_ASSERT(axis3 >= 0 && axis3 < GGML_MAX_DIMS); + + GGML_ASSERT(axis0 != axis1); + GGML_ASSERT(axis0 != axis2); + GGML_ASSERT(axis0 != axis3); + GGML_ASSERT(axis1 != axis2); + GGML_ASSERT(axis1 != axis3); + GGML_ASSERT(axis2 != axis3); + + struct ggml_tensor * result = ggml_view_tensor(ctx, a); + ggml_format_name(result, "%s (permuted)", a->name); + + int ne[GGML_MAX_DIMS]; + int nb[GGML_MAX_DIMS]; + + ne[axis0] = a->ne[0]; + ne[axis1] = a->ne[1]; + ne[axis2] = a->ne[2]; + ne[axis3] = a->ne[3]; + + nb[axis0] = a->nb[0]; + nb[axis1] = a->nb[1]; + nb[axis2] = a->nb[2]; + nb[axis3] = a->nb[3]; + + result->ne[0] = ne[0]; + result->ne[1] = ne[1]; + result->ne[2] = ne[2]; + result->ne[3] = ne[3]; + + result->nb[0] = nb[0]; + result->nb[1] = nb[1]; + result->nb[2] = nb[2]; + result->nb[3] = nb[3]; + + result->op = GGML_OP_PERMUTE; + result->src[0] = a; + + int32_t params[] = { axis0, axis1, axis2, axis3 }; + ggml_set_op_params(result, params, sizeof(params)); + + return result; +} + +// ggml_transpose + +struct ggml_tensor * ggml_transpose( + struct ggml_context * ctx, + struct ggml_tensor * a) { + struct ggml_tensor * result = ggml_view_tensor(ctx, a); + ggml_format_name(result, "%s (transposed)", a->name); + + result->ne[0] = a->ne[1]; + result->ne[1] = a->ne[0]; + + result->nb[0] = a->nb[1]; + result->nb[1] = a->nb[0]; + + result->op = GGML_OP_TRANSPOSE; + result->src[0] = a; + + return result; +} + +// ggml_get_rows + +struct ggml_tensor * ggml_get_rows( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + GGML_ASSERT(a->ne[2] == b->ne[1]); + GGML_ASSERT(a->ne[3] == b->ne[2]); + GGML_ASSERT(b->ne[3] == 1); + GGML_ASSERT(b->type == GGML_TYPE_I32); + + // TODO: implement non F32 return + enum ggml_type type = GGML_TYPE_F32; + if (a->type == GGML_TYPE_I32) { + type = a->type; + } + struct ggml_tensor * result = ggml_new_tensor_4d(ctx, type, a->ne[0], b->ne[0], b->ne[1], b->ne[2]); + + result->op = GGML_OP_GET_ROWS; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +// ggml_get_rows_back + +struct ggml_tensor * ggml_get_rows_back( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c) { + GGML_ASSERT(ggml_is_matrix(a) && ggml_is_vector(b) && b->type == GGML_TYPE_I32); + GGML_ASSERT(ggml_is_matrix(c) && (a->ne[0] == c->ne[0])); + + // TODO: implement non F32 return + //struct ggml_tensor * result = ggml_new_tensor_2d(ctx, a->type, a->ne[0], b->ne[0]); + struct ggml_tensor * result = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, c->ne[0], c->ne[1]); + + result->op = GGML_OP_GET_ROWS_BACK; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +// ggml_set_rows + +struct ggml_tensor * ggml_set_rows( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c) { + GGML_ASSERT(a->ne[0] == b->ne[0]); + GGML_ASSERT(a->ne[2] == b->ne[2]); + GGML_ASSERT(a->ne[3] == b->ne[3]); + GGML_ASSERT(b->ne[1] == c->ne[0]); + GGML_ASSERT(b->ne[2] % c->ne[1] == 0); + GGML_ASSERT(b->ne[3] % c->ne[2] == 0); + GGML_ASSERT(c->ne[3] == 1); + GGML_ASSERT(b->type == GGML_TYPE_F32); + GGML_ASSERT(c->type == GGML_TYPE_I64 || c->type == GGML_TYPE_I32); + + GGML_ASSERT(ggml_is_contiguous_rows(a)); + GGML_ASSERT(ggml_is_contiguous_rows(b)); + + struct ggml_tensor * result = ggml_view_tensor(ctx, a); + + result->op = GGML_OP_SET_ROWS; + result->src[0] = b; + result->src[1] = c; + result->src[2] = a; // note: order is weird due to legacy reasons (https://github.com/ggml-org/llama.cpp/pull/16063#discussion_r2385795931) + + return result; +} + +// ggml_diag + +struct ggml_tensor * ggml_diag( + struct ggml_context * ctx, + struct ggml_tensor * a) { + GGML_ASSERT(a->ne[1] == 1); + + const int64_t ne[4] = { a->ne[0], a->ne[0], a->ne[2], a->ne[3] }; + struct ggml_tensor * result = ggml_new_tensor(ctx, a->type, 4, ne); + + result->op = GGML_OP_DIAG; + result->src[0] = a; + + return result; +} + +// ggml_diag_mask_inf + +static struct ggml_tensor * ggml_diag_mask_inf_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + int n_past, + bool inplace) { + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + int32_t params[] = { n_past }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_DIAG_MASK_INF; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_diag_mask_inf( + struct ggml_context * ctx, + struct ggml_tensor * a, + int n_past) { + return ggml_diag_mask_inf_impl(ctx, a, n_past, false); +} + +struct ggml_tensor * ggml_diag_mask_inf_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + int n_past) { + return ggml_diag_mask_inf_impl(ctx, a, n_past, true); +} + +// ggml_diag_mask_zero + +static struct ggml_tensor * ggml_diag_mask_zero_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + int n_past, + bool inplace) { + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + int32_t params[] = { n_past }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_DIAG_MASK_ZERO; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_diag_mask_zero( + struct ggml_context * ctx, + struct ggml_tensor * a, + int n_past) { + return ggml_diag_mask_zero_impl(ctx, a, n_past, false); +} + +struct ggml_tensor * ggml_diag_mask_zero_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + int n_past) { + return ggml_diag_mask_zero_impl(ctx, a, n_past, true); +} + +// ggml_soft_max + +static struct ggml_tensor * ggml_soft_max_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * mask, + float scale, + float max_bias, + bool inplace) { + GGML_ASSERT(ggml_is_contiguous(a)); + + if (mask) { + GGML_ASSERT(mask->type == GGML_TYPE_F16 || mask->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(mask)); + GGML_ASSERT(mask->ne[0] == a->ne[0]); + GGML_ASSERT(mask->ne[1] >= a->ne[1]); + GGML_ASSERT(a->ne[2]%mask->ne[2] == 0); + GGML_ASSERT(a->ne[3]%mask->ne[3] == 0); + } + + if (max_bias > 0.0f) { + GGML_ASSERT(mask); + } + + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + float params[] = { scale, max_bias }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_SOFT_MAX; + result->src[0] = a; + result->src[1] = mask; + + return result; +} + +struct ggml_tensor * ggml_soft_max( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_soft_max_impl(ctx, a, NULL, 1.0f, 0.0f, false); +} + +struct ggml_tensor * ggml_soft_max_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a) { + return ggml_soft_max_impl(ctx, a, NULL, 1.0f, 0.0f, true); +} + +struct ggml_tensor * ggml_soft_max_ext( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * mask, + float scale, + float max_bias) { + return ggml_soft_max_impl(ctx, a, mask, scale, max_bias, false); +} + +struct ggml_tensor * ggml_soft_max_ext_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * mask, + float scale, + float max_bias) { + return ggml_soft_max_impl(ctx, a, mask, scale, max_bias, true); +} + +void ggml_soft_max_add_sinks( + struct ggml_tensor * a, + struct ggml_tensor * sinks) { + if (!sinks) { + a->src[2] = NULL; + return; + } + + GGML_ASSERT(a->op == GGML_OP_SOFT_MAX); + GGML_ASSERT(a->src[2] == NULL); + GGML_ASSERT(a->src[0]->ne[2] == sinks->ne[0]); + GGML_ASSERT(sinks->type == GGML_TYPE_F32); + + a->src[2] = sinks; +} + +// ggml_soft_max_ext_back + +static struct ggml_tensor * ggml_soft_max_ext_back_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + float scale, + float max_bias, + bool inplace) { + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + result->op = GGML_OP_SOFT_MAX_BACK; + result->src[0] = a; + result->src[1] = b; + + memcpy((float *) result->op_params + 0, &scale, sizeof(float)); + memcpy((float *) result->op_params + 1, &max_bias, sizeof(float)); + + return result; +} + +struct ggml_tensor * ggml_soft_max_ext_back( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + float scale, + float max_bias) { + return ggml_soft_max_ext_back_impl(ctx, a, b, scale, max_bias, false); +} + +struct ggml_tensor * ggml_soft_max_ext_back_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + float scale, + float max_bias) { + return ggml_soft_max_ext_back_impl(ctx, a, b, scale, max_bias, true); +} + +// ggml_rope + +static struct ggml_tensor * ggml_rope_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c, + int n_dims, + int sections[GGML_MROPE_SECTIONS], + int mode, + int n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow, + bool inplace) { + GGML_ASSERT((mode & 1) == 0 && "mode & 1 == 1 is no longer supported"); + + GGML_ASSERT(ggml_is_vector(b)); + GGML_ASSERT(b->type == GGML_TYPE_I32); + + bool mrope_used = mode & GGML_ROPE_TYPE_MROPE; + if (mrope_used) { + GGML_ASSERT(a->ne[2] * 4 == b->ne[0]); // mrope expecting 4 position ids per token + } else { + GGML_ASSERT(a->ne[2] == b->ne[0]); + } + + if (c) { + GGML_ASSERT(c->type == GGML_TYPE_F32); + GGML_ASSERT(c->ne[0] >= n_dims / 2); + } + + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + int32_t params[15] = { /*n_past*/ 0, n_dims, mode, /*n_ctx*/ 0, n_ctx_orig }; + memcpy(params + 5, &freq_base, sizeof(float)); + memcpy(params + 6, &freq_scale, sizeof(float)); + memcpy(params + 7, &ext_factor, sizeof(float)); + memcpy(params + 8, &attn_factor, sizeof(float)); + memcpy(params + 9, &beta_fast, sizeof(float)); + memcpy(params + 10, &beta_slow, sizeof(float)); + if (mrope_used && sections) { + memcpy(params + 11, sections, sizeof(int32_t) * GGML_MROPE_SECTIONS); + } else { + memset(params + 11, 0, sizeof(int32_t) * GGML_MROPE_SECTIONS); + } + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_ROPE; + result->src[0] = a; + result->src[1] = b; + result->src[2] = c; + + return result; +} + +struct ggml_tensor * ggml_rope( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int n_dims, + int mode) { + return ggml_rope_impl( + ctx, a, b, NULL, n_dims, NULL, mode, 0, 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f, false + ); +} + +struct ggml_tensor * ggml_rope_multi( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c, + int n_dims, + int sections[GGML_MROPE_SECTIONS], + int mode, + int n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + return ggml_rope_impl( + ctx, a, b, c, n_dims, sections, mode, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow, false + ); +} + +struct ggml_tensor * ggml_rope_multi_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c, + int n_dims, + int sections[GGML_MROPE_SECTIONS], + int mode, + int n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + return ggml_rope_impl( + ctx, a, b, c, n_dims, sections, mode, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow, true + ); +} + +struct ggml_tensor * ggml_rope_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int n_dims, + int mode) { + return ggml_rope_impl( + ctx, a, b, NULL, n_dims, NULL, mode, 0, 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f, true + ); +} + +struct ggml_tensor * ggml_rope_ext( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c, + int n_dims, + int mode, + int n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + return ggml_rope_impl( + ctx, a, b, c, n_dims, NULL, mode, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow, false + ); +} + +struct ggml_tensor * ggml_rope_ext_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c, + int n_dims, + int mode, + int n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + return ggml_rope_impl( + ctx, a, b, c, n_dims, NULL, mode, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow, true + ); +} + +struct ggml_tensor * ggml_rope_custom( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int n_dims, + int mode, + int n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + return ggml_rope_impl( + ctx, a, b, NULL, n_dims, NULL, mode, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow, false + ); +} + +struct ggml_tensor * ggml_rope_custom_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int n_dims, + int mode, + int n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + return ggml_rope_impl( + ctx, a, b, NULL, n_dims, NULL, mode, n_ctx_orig, freq_base, freq_scale, + ext_factor, attn_factor, beta_fast, beta_slow, true + ); +} + +// Apparently solving `n_rot = 2pi * x * base^((2 * max_pos_emb) / n_dims)` for x, we get +// `corr_dim(n_rot) = n_dims * log(max_pos_emb / (n_rot * 2pi)) / (2 * log(base))` +static float ggml_rope_yarn_corr_dim(int n_dims, int n_ctx_orig, float n_rot, float base) { + return n_dims * logf(n_ctx_orig / (n_rot * 2 * (float)M_PI)) / (2 * logf(base)); +} + +void ggml_rope_yarn_corr_dims( + int n_dims, int n_ctx_orig, float freq_base, float beta_fast, float beta_slow, float dims[2] +) { + // start and end correction dims + float start = floorf(ggml_rope_yarn_corr_dim(n_dims, n_ctx_orig, beta_fast, freq_base)); + float end = ceilf(ggml_rope_yarn_corr_dim(n_dims, n_ctx_orig, beta_slow, freq_base)); + dims[0] = MAX(0, start); + dims[1] = MIN(n_dims - 1, end); +} + +// ggml_rope_back + +struct ggml_tensor * ggml_rope_ext_back( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c, + int n_dims, + int mode, + int n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + struct ggml_tensor * result = ggml_rope_ext( + ctx, a, b, c, n_dims, mode, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); + result->op = GGML_OP_ROPE_BACK; + return result; +} + +struct ggml_tensor * ggml_rope_multi_back( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c, + int n_dims, + int sections[4], + int mode, + int n_ctx_orig, + float freq_base, + float freq_scale, + float ext_factor, + float attn_factor, + float beta_fast, + float beta_slow) { + struct ggml_tensor * result = ggml_rope_multi( + ctx, a, b, c, n_dims, sections, mode, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); + result->op = GGML_OP_ROPE_BACK; + return result; +} +// ggml_clamp + +struct ggml_tensor * ggml_clamp( + struct ggml_context * ctx, + struct ggml_tensor * a, + float min, + float max) { + // TODO: when implement backward, fix this: + struct ggml_tensor * result = ggml_view_tensor(ctx, a); + + float params[] = { min, max }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_CLAMP; + result->src[0] = a; + + return result; +} + +static int64_t ggml_calc_conv_output_size(int64_t ins, int64_t ks, int s, int p, int d) { + return (ins + 2 * p - d * (ks - 1) - 1) / s + 1; +} + +// im2col: [N, IC, IH, IW] => [N, OH, OW, IC*KH*KW] +// a: [OC，IC, KH, KW] +// b: [N, IC, IH, IW] +// result: [N, OH, OW, IC*KH*KW] +struct ggml_tensor * ggml_im2col( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int s0, + int s1, + int p0, + int p1, + int d0, + int d1, + bool is_2D, + enum ggml_type dst_type) { + if (is_2D) { + GGML_ASSERT(a->ne[2] == b->ne[2]); + } else { + //GGML_ASSERT(b->ne[1] % a->ne[1] == 0); + GGML_ASSERT(b->ne[1] == a->ne[1]); + GGML_ASSERT(b->ne[3] == 1); + } + + const int64_t OH = is_2D ? ggml_calc_conv_output_size(b->ne[1], a->ne[1], s1, p1, d1) : 0; + const int64_t OW = ggml_calc_conv_output_size(b->ne[0], a->ne[0], s0, p0, d0); + + GGML_ASSERT((!is_2D || OH > 0) && "b too small compared to a"); + GGML_ASSERT((OW > 0) && "b too small compared to a"); + + const int64_t ne[4] = { + is_2D ? (a->ne[2] * a->ne[1] * a->ne[0]) : a->ne[1] * a->ne[0], + OW, + is_2D ? OH : b->ne[2], + is_2D ? b->ne[3] : 1, + }; + + struct ggml_tensor * result = ggml_new_tensor(ctx, dst_type, 4, ne); + int32_t params[] = { s0, s1, p0, p1, d0, d1, (is_2D ? 1 : 0) }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_IM2COL; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +struct ggml_tensor * ggml_im2col_back( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int64_t * ne, + int s0, + int s1, + int p0, + int p1, + int d0, + int d1, + bool is_2D) { + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + int32_t params[] = { s0, s1, p0, p1, d0, d1, (is_2D ? 1 : 0) }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_IM2COL_BACK; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +// ggml_conv_1d + +struct ggml_tensor * ggml_conv_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int s0, + int p0, + int d0) { + struct ggml_tensor * im2col = ggml_im2col(ctx, a, b, s0, 0, p0, 0, d0, 0, false, GGML_TYPE_F16); // [N, OL, IC * K] + + struct ggml_tensor * result = + ggml_mul_mat(ctx, + ggml_reshape_2d(ctx, im2col, im2col->ne[0], (im2col->ne[2] * im2col->ne[1])), // [N, OL, IC * K] => [N*OL, IC * K] + ggml_reshape_2d(ctx, a, (a->ne[0] * a->ne[1]), a->ne[2])); // [OC，IC, K] => [OC, IC * K] + + result = ggml_reshape_3d(ctx, result, im2col->ne[1], a->ne[2], im2col->ne[2]); // [N, OC, OL] + + return result; +} + +// ggml_conv_1d_ph + +struct ggml_tensor* ggml_conv_1d_ph( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int s, + int d) { + return ggml_conv_1d(ctx, a, b, s, a->ne[0] / 2, d); +} + +// ggml_conv_1d_dw + +struct ggml_tensor * ggml_conv_1d_dw( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int s0, + int p0, + int d0) { + struct ggml_tensor * new_b = ggml_reshape_4d(ctx, b, b->ne[0], 1, b->ne[1], b->ne[2]); + + struct ggml_tensor * im2col = ggml_im2col(ctx, a, new_b, s0, 0, p0, 0, d0, 0, false, GGML_TYPE_F16); + + struct ggml_tensor * result = ggml_mul_mat(ctx, im2col, a); + + result = ggml_reshape_3d(ctx, result, result->ne[0], result->ne[2], 1); + + return result; +} + +// ggml_conv_1d_dw_ph + +struct ggml_tensor * ggml_conv_1d_dw_ph( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int s0, + int d0) { + return ggml_conv_1d_dw(ctx, a, b, s0, a->ne[0] / 2, d0); +} + +// ggml_conv_transpose_1d + +static int64_t ggml_calc_conv_transpose_1d_output_size(int64_t ins, int64_t ks, int s, int p, int d) { + return (ins - 1) * s - 2 * p + d * (ks - 1) + 1; +} + +GGML_API struct ggml_tensor * ggml_conv_transpose_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int s0, + int p0, + int d0) { + GGML_ASSERT(ggml_is_matrix(b)); + GGML_ASSERT(a->ne[2] == b->ne[1]); + GGML_ASSERT(a->ne[3] == 1); + + GGML_ASSERT(p0 == 0); + GGML_ASSERT(d0 == 1); + + const int64_t ne[4] = { + ggml_calc_conv_transpose_1d_output_size(b->ne[0], a->ne[0], s0, 0 /*p0*/, 1 /*d0*/), + a->ne[1], b->ne[2], 1, + }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + int32_t params[] = { s0, p0, d0 }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_CONV_TRANSPOSE_1D; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +// ggml_conv_2d + +// a: [OC，IC, KH, KW] +// b: [N, IC, IH, IW] +// result: [N, OC, OH, OW] +struct ggml_tensor * ggml_conv_2d( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int s0, + int s1, + int p0, + int p1, + int d0, + int d1) { + struct ggml_tensor * im2col = ggml_im2col(ctx, a, b, s0, s1, p0, p1, d0, d1, true, a->type); // [N, OH, OW, IC * KH * KW] + + struct ggml_tensor * result = + ggml_mul_mat(ctx, + ggml_reshape_2d(ctx, im2col, im2col->ne[0], im2col->ne[3] * im2col->ne[2] * im2col->ne[1]), // [N, OH, OW, IC * KH * KW] => [N*OH*OW, IC * KH * KW] + ggml_reshape_2d(ctx, a, (a->ne[0] * a->ne[1] * a->ne[2]), a->ne[3])); // [OC，IC, KH, KW] => [OC, IC * KH * KW] + + result = ggml_reshape_4d(ctx, result, im2col->ne[1], im2col->ne[2], im2col->ne[3], a->ne[3]); // [OC, N, OH, OW] + result = ggml_cont(ctx, ggml_permute(ctx, result, 0, 1, 3, 2)); // [N, OC, OH, OW] + + + return result; +} + +// a: [OC*IC, KD, KH, KW] +// b: [N*IC, ID, IH, IW] +// result: [N*OD, OH, OW, IC * KD * KH * KW] +struct ggml_tensor * ggml_im2col_3d( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int64_t IC, + int s0, // stride width + int s1, // stride height + int s2, // stride depth + int p0, // padding width + int p1, // padding height + int p2, // padding depth + int d0, // dilation width + int d1, // dilation height + int d2, // dilation depth + enum ggml_type dst_type) { + const int64_t N = b->ne[3] / IC; + const int64_t ID = b->ne[2]; + const int64_t IH = b->ne[1]; + const int64_t IW = b->ne[0]; + + const int64_t OC = a->ne[3] / IC; + UNUSED(OC); + const int64_t KD = a->ne[2]; + const int64_t KH = a->ne[1]; + const int64_t KW = a->ne[0]; + const int64_t OD = ggml_calc_conv_output_size(ID, KD, s2, p2, d2); + const int64_t OH = ggml_calc_conv_output_size(IH, KH, s1, p1, d1); + const int64_t OW = ggml_calc_conv_output_size(IW, KW, s0, p0, d0); + + GGML_ASSERT((OD > 0) && "b too small compared to a"); + GGML_ASSERT((OH > 0) && "b too small compared to a"); + GGML_ASSERT((OW > 0) && "b too small compared to a"); + + + const int64_t ne[4] = {KW*KH*KD*IC, OW, OH, OD*N}; + + struct ggml_tensor * result = ggml_new_tensor(ctx, dst_type, 4, ne); + int32_t params[] = { s0, s1, s2, p0, p1, p2, d0, d1, d2, (int32_t)IC}; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_IM2COL_3D; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +// a: [OC*IC, KD, KH, KW] +// b: [N*IC, ID, IH, IW] +// result: [N*OC, OD, OH, OW] +struct ggml_tensor * ggml_conv_3d( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int64_t IC, + int s0, // stride width + int s1, // stride height + int s2, // stride depth + int p0, // padding width + int p1, // padding height + int p2, // padding depth + int d0, // dilation width + int d1, // dilation height + int d2 // dilation depth + ) { + struct ggml_tensor * im2col = ggml_im2col_3d(ctx, a, b, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2, a->type); // [N*OD, OH, OW, IC * KD * KH * KW] + + int64_t OC = a->ne[3] / IC; + int64_t N = b->ne[3] / IC; + struct ggml_tensor * result = + ggml_mul_mat(ctx, + ggml_reshape_2d(ctx, im2col, im2col->ne[0], im2col->ne[3] * im2col->ne[2] * im2col->ne[1]), // [N*OD, OH, OW, IC * KD * KH * KW] => [N*OD*OH*OW, IC * KD * KH * KW] + ggml_reshape_2d(ctx, a, (a->ne[0] * a->ne[1] * a->ne[2] * IC), OC)); // [OC*IC, KD, KH, KW] => [OC, IC * KD * KH * KW] + + int64_t OD = im2col->ne[3] / N; + result = ggml_reshape_4d(ctx, result, im2col->ne[1]*im2col->ne[2], OD, N, OC); // [OC, N*OD*OH*OW] => [OC, N, OD, OH*OW] + result = ggml_cont(ctx, ggml_permute(ctx, result, 0, 1, 3, 2)); // [N, OC, OD, OH*OW] + result = ggml_reshape_4d(ctx, result, im2col->ne[1], im2col->ne[2], OD, OC * N); // [N*OC, OD, OH, OW] + + return result; +} + +// ggml_conv_2d_sk_p0 + +struct ggml_tensor * ggml_conv_2d_sk_p0( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_conv_2d(ctx, a, b, a->ne[0], a->ne[1], 0, 0, 1, 1); +} + +// ggml_conv_2d_s1_ph + +struct ggml_tensor * ggml_conv_2d_s1_ph( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + return ggml_conv_2d(ctx, a, b, 1, 1, a->ne[0] / 2, a->ne[1] / 2, 1, 1); +} + +// ggml_conv_2d_dw + +struct ggml_tensor * ggml_conv_2d_dw( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int s0, + int s1, + int p0, + int p1, + int d0, + int d1) { + struct ggml_tensor * new_a = ggml_reshape_4d(ctx, a, a->ne[0], a->ne[1], 1, a->ne[2] * a->ne[3]); + struct ggml_tensor * im2col = ggml_im2col(ctx, new_a, + ggml_reshape_4d(ctx, b, b->ne[0], b->ne[1], 1, b->ne[2] * b->ne[3]), + s0, s1, p0, p1, d0, d1, true, GGML_TYPE_F16); // [N * IC, OH, OW, KH * KW] + struct ggml_tensor * new_b = ggml_reshape_4d(ctx, im2col, im2col->ne[0], im2col->ne[2] * im2col->ne[1], b->ne[2], b->ne[3]); // [N * IC, OH, OW, KH * KW] => [N, IC, OH * OW, KH * KW] + + new_a = ggml_reshape_4d(ctx, new_a, (new_a->ne[0] * new_a->ne[1]), new_a->ne[2], new_a->ne[3], 1); // [OC，1, KH, KW] => [1, OC, 1, KH * KW] + struct ggml_tensor * result = ggml_mul_mat(ctx, new_a, new_b); + result = ggml_reshape_4d(ctx, result, im2col->ne[1], im2col->ne[2], b->ne[2], b->ne[3]); // [N, OC, OH, OW] + + return result; +} + +// ggml_conv_2d_dw_direct + +struct ggml_tensor * ggml_conv_2d_dw_direct( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int stride0, + int stride1, + int pad0, + int pad1, + int dilation0, + int dilation1) { + GGML_ASSERT(a->ne[2] == 1); + GGML_ASSERT(a->ne[3] == b->ne[2]); + int64_t ne[4]; + ne[0] = ggml_calc_conv_output_size(b->ne[0], a->ne[0], stride0, pad0, dilation0); + ne[1] = ggml_calc_conv_output_size(b->ne[1], a->ne[1], stride1, pad1, dilation1); + ne[2] = b->ne[2]; + ne[3] = b->ne[3]; + + struct ggml_tensor * result = ggml_new_tensor(ctx, b->type, 4, ne); + + if (ggml_is_contiguous_channels(b)) { + // Result will be permuted the same way as input (CWHN order) + const int64_t type_size = ggml_type_size(result->type); + GGML_ASSERT(ggml_blck_size(result->type) == 1); + result->nb[0] = result->ne[2] * type_size; + result->nb[1] = result->ne[0] * result->nb[0]; + result->nb[2] = type_size; + } + + int32_t params[] = { stride0, stride1, pad0, pad1, dilation0, dilation1 }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_CONV_2D_DW; + result->src[0] = a; + result->src[1] = b; + return result; +} + +// ggml_conv_2d_direct + +struct ggml_tensor * ggml_conv_2d_direct( + struct ggml_context * ctx, + struct ggml_tensor * a, // convolution kernel [KW, KH, IC, OC] + struct ggml_tensor * b, // input data [W, H, C, N] + int s0, // stride dimension 0 + int s1, // stride dimension 1 + int p0, // padding dimension 0 + int p1, // padding dimension 1 + int d0, // dilation dimension 0 + int d1) {// dilation dimension 1 + + GGML_ASSERT(a->ne[2] == b->ne[2]); + //GGML_ASSERT(a->type == b->type); + + int64_t ne[4]; + ne[0] = ggml_calc_conv_output_size(b->ne[0], a->ne[0], s0, p0, d0); + ne[1] = ggml_calc_conv_output_size(b->ne[1], a->ne[1], s1, p1, d1); + ne[2] = a->ne[3]; + ne[3] = b->ne[3]; + + struct ggml_tensor * result = ggml_new_tensor(ctx, b->type, 4, ne); + + ggml_set_op_params_i32(result, 0, s0); + ggml_set_op_params_i32(result, 1, s1); + ggml_set_op_params_i32(result, 2, p0); + ggml_set_op_params_i32(result, 3, p1); + ggml_set_op_params_i32(result, 4, d0); + ggml_set_op_params_i32(result, 5, d1); + + result->op = GGML_OP_CONV_2D; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +// ggml_conv_3d_direct + +struct ggml_tensor * ggml_conv_3d_direct( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int s0, + int s1, + int s2, + int p0, + int p1, + int p2, + int d0, + int d1, + int d2, + int c, + int n, + int oc) { + + GGML_ASSERT(a->ne[3] == (int64_t) c * oc); + GGML_ASSERT(b->ne[3] == (int64_t) c * n); + + int64_t ne[4]; + ne[0] = ggml_calc_conv_output_size(b->ne[0], a->ne[0], s0, p0, d0); + ne[1] = ggml_calc_conv_output_size(b->ne[1], a->ne[1], s1, p1, d1); + ne[2] = ggml_calc_conv_output_size(b->ne[2], a->ne[2], s2, p2, d2); + ne[3] = (int64_t) oc * n; + + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + ggml_set_op_params_i32(result, 0, s0); + ggml_set_op_params_i32(result, 1, s1); + ggml_set_op_params_i32(result, 2, s2); + ggml_set_op_params_i32(result, 3, p0); + ggml_set_op_params_i32(result, 4, p1); + ggml_set_op_params_i32(result, 5, p2); + ggml_set_op_params_i32(result, 6, d0); + ggml_set_op_params_i32(result, 7, d1); + ggml_set_op_params_i32(result, 8, d2); + ggml_set_op_params_i32(result, 9, c); + ggml_set_op_params_i32(result, 10, n); + ggml_set_op_params_i32(result, 11, oc); + + result->op = GGML_OP_CONV_3D; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +// ggml_conv_transpose_2d_p0 + +static int64_t ggml_calc_conv_transpose_output_size(int64_t ins, int64_t ks, int s, int p) { + return (ins - 1) * s - 2 * p + ks; +} + +struct ggml_tensor * ggml_conv_transpose_2d_p0( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + int stride) { + GGML_ASSERT(a->ne[3] == b->ne[2]); + + const int64_t ne[4] = { + ggml_calc_conv_transpose_output_size(b->ne[0], a->ne[0], stride, 0 /*p0*/), + ggml_calc_conv_transpose_output_size(b->ne[1], a->ne[1], stride, 0 /*p1*/), + a->ne[2], b->ne[3], + }; + + struct ggml_tensor* result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + ggml_set_op_params_i32(result, 0, stride); + + result->op = GGML_OP_CONV_TRANSPOSE_2D; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +// ggml_pool_* + +static int64_t ggml_calc_pool_output_size(int64_t ins, int ks, int s, float p) { + return (ins + 2 * p - ks) / s + 1; +} + +// ggml_pool_1d + +struct ggml_tensor * ggml_pool_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, + enum ggml_op_pool op, + int k0, + int s0, + int p0) { + const int64_t ne[4] = { + ggml_calc_pool_output_size(a->ne[0], k0, s0, p0), + a->ne[1], + a->ne[2], + a->ne[3], + }; + GGML_ASSERT(ne[0] > 0); + + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + int32_t params[] = { op, k0, s0, p0 }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_POOL_1D; + result->src[0] = a; + + return result; +} + +// ggml_pool_2d + +struct ggml_tensor * ggml_pool_2d( + struct ggml_context * ctx, + struct ggml_tensor * a, + enum ggml_op_pool op, + int k0, + int k1, + int s0, + int s1, + float p0, + float p1) { + struct ggml_tensor * result; + const int64_t ne[4] = { + ggml_calc_pool_output_size(a->ne[0], k0, s0, p0), + ggml_calc_pool_output_size(a->ne[1], k1, s1, p1), + a->ne[2], + a->ne[3], + }; + GGML_ASSERT(ne[0] > 0); + GGML_ASSERT(ne[1] > 0); + + result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + int32_t params[] = { op, k0, k1, s0, s1, p0, p1 }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_POOL_2D; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_pool_2d_back( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * af, + enum ggml_op_pool op, + int k0, + int k1, + int s0, + int s1, + float p0, + float p1) { + struct ggml_tensor * result; + result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, af->ne); + + int32_t params[] = { op, k0, k1, s0, s1, p0, p1 }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_POOL_2D_BACK; + result->src[0] = a; + result->src[1] = af; + + return result; +} + +// ggml_upscale / ggml_interpolate + +static struct ggml_tensor * ggml_interpolate_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3, + uint32_t mode) { + GGML_ASSERT((mode & 0xFF) < GGML_SCALE_MODE_COUNT); + // TODO: implement antialias for modes other than bilinear + GGML_ASSERT(!(mode & GGML_SCALE_FLAG_ANTIALIAS) || (mode & 0xFF) == GGML_SCALE_MODE_BILINEAR); + GGML_ASSERT(a->type == GGML_TYPE_F32); + + struct ggml_tensor * result = ggml_new_tensor_4d(ctx, a->type, ne0, ne1, ne2, ne3); + + ggml_set_op_params_i32(result, 0, (int32_t)mode); + + result->op = GGML_OP_UPSCALE; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_upscale( + struct ggml_context * ctx, + struct ggml_tensor * a, + int scale_factor, + enum ggml_scale_mode mode) { + GGML_ASSERT(scale_factor > 1); + return ggml_interpolate_impl(ctx, a, a->ne[0] * scale_factor, a->ne[1] * scale_factor, a->ne[2], a->ne[3], mode); +} + +struct ggml_tensor * ggml_upscale_ext( + struct ggml_context * ctx, + struct ggml_tensor * a, + int ne0, + int ne1, + int ne2, + int ne3, + enum ggml_scale_mode mode) { + return ggml_interpolate_impl(ctx, a, ne0, ne1, ne2, ne3, mode); +} + +struct ggml_tensor * ggml_interpolate( + struct ggml_context * ctx, + struct ggml_tensor * a, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3, + uint32_t mode) { + return ggml_interpolate_impl(ctx, a, ne0, ne1, ne2, ne3, mode); +} + +// ggml_pad + +struct ggml_tensor * ggml_pad( + struct ggml_context * ctx, + struct ggml_tensor * a, + int p0, + int p1, + int p2, + int p3) { + return ggml_pad_ext(ctx, a, 0, p0, 0, p1, 0, p2, 0, p3); +} + +// ggml_pad_circular + +struct ggml_tensor * ggml_pad_circular( + struct ggml_context * ctx, + struct ggml_tensor * a, + int p0, + int p1, + int p2, + int p3) { + return ggml_pad_ext_circular(ctx, a, 0, p0, 0, p1, 0, p2, 0, p3); +} + +struct ggml_tensor * ggml_pad_ext( + struct ggml_context * ctx, + struct ggml_tensor * a, + int lp0, + int rp0, + int lp1, + int rp1, + int lp2, + int rp2, + int lp3, + int rp3 + ) { + struct ggml_tensor * result = ggml_new_tensor_4d(ctx, a->type, + a->ne[0] + lp0 + rp0, + a->ne[1] + lp1 + rp1, + a->ne[2] + lp2 + rp2, + a->ne[3] + lp3 + rp3); + + ggml_set_op_params_i32(result, 0, lp0); + ggml_set_op_params_i32(result, 1, rp0); + ggml_set_op_params_i32(result, 2, lp1); + ggml_set_op_params_i32(result, 3, rp1); + ggml_set_op_params_i32(result, 4, lp2); + ggml_set_op_params_i32(result, 5, rp2); + ggml_set_op_params_i32(result, 6, lp3); + ggml_set_op_params_i32(result, 7, rp3); + ggml_set_op_params_i32(result, 8, 0); // not circular by default + + + result->op = GGML_OP_PAD; + result->src[0] = a; + + return result; +} + +// ggml_pad_ext_circular + +struct ggml_tensor * ggml_pad_ext_circular( + struct ggml_context * ctx, + struct ggml_tensor * a, + int lp0, + int rp0, + int lp1, + int rp1, + int lp2, + int rp2, + int lp3, + int rp3 + ) { + struct ggml_tensor * result = ggml_pad_ext(ctx, a, lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3); + ggml_set_op_params_i32(result, 8, 1); // circular + return result; +} + +// ggml_pad_reflect_1d + +struct ggml_tensor * ggml_pad_reflect_1d( + struct ggml_context * ctx, + struct ggml_tensor * a, + int p0, + int p1) { + GGML_ASSERT(p0 >= 0); + GGML_ASSERT(p1 >= 0); + + GGML_ASSERT(p0 < a->ne[0]); // padding length on each size must be less than the + GGML_ASSERT(p1 < a->ne[0]); // existing length of the dimension being padded + + GGML_ASSERT(ggml_is_contiguous(a)); + GGML_ASSERT(a->type == GGML_TYPE_F32); + + struct ggml_tensor * result = ggml_new_tensor_4d(ctx, a->type, + a->ne[0] + p0 + p1, + a->ne[1], + a->ne[2], + a->ne[3]); + + int32_t params[] = { p0, p1 }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_PAD_REFLECT_1D; + result->src[0] = a; + + return result; +} + +// ggml_roll + +struct ggml_tensor * ggml_roll( + struct ggml_context * ctx, + struct ggml_tensor * a, + int shift0, + int shift1, + int shift2, + int shift3) { + GGML_ASSERT(a->nb[0] == ggml_type_size(a->type)); + GGML_ASSERT(abs(shift0) < a->ne[0]); + GGML_ASSERT(abs(shift1) < a->ne[1]); + GGML_ASSERT(abs(shift2) < a->ne[2]); + GGML_ASSERT(abs(shift3) < a->ne[3]); + + struct ggml_tensor * result = ggml_dup_tensor(ctx, a); + + ggml_set_op_params_i32(result, 0, shift0); + ggml_set_op_params_i32(result, 1, shift1); + ggml_set_op_params_i32(result, 2, shift2); + ggml_set_op_params_i32(result, 3, shift3); + + result->op = GGML_OP_ROLL; + result->src[0] = a; + + return result; +} + +// ggml_timestep_embedding + +struct ggml_tensor * ggml_timestep_embedding( + struct ggml_context * ctx, + struct ggml_tensor * timesteps, + int dim, + int max_period) { + + struct ggml_tensor * result = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, dim, timesteps->ne[0]); + + ggml_set_op_params_i32(result, 0, dim); + ggml_set_op_params_i32(result, 1, max_period); + + result->op = GGML_OP_TIMESTEP_EMBEDDING; + result->src[0] = timesteps; + + return result; +} + +// ggml_tri + +struct ggml_tensor * ggml_tri( + struct ggml_context * ctx, + struct ggml_tensor * a, + enum ggml_tri_type type) { + GGML_ASSERT(a->type == GGML_TYPE_F32); + + GGML_ASSERT(ggml_is_contiguous(a)); + GGML_ASSERT(a->ne[0] == a->ne[1]); + + struct ggml_tensor * result = ggml_dup_tensor(ctx, a); + + ggml_set_op_params_i32(result, 0, type); + + result->op = GGML_OP_TRI; + result->src[0] = a; + + return result; +} + +// ggml_fill + +static struct ggml_tensor * ggml_fill_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + float c, + bool inplace) { + GGML_ASSERT(a->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(a)); + + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + ggml_set_op_params_f32(result, 0, c); + + result->op = GGML_OP_FILL; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_fill( + struct ggml_context * ctx, + struct ggml_tensor * a, + float c) { + return ggml_fill_impl(ctx, a, c, false); +} + +struct ggml_tensor * ggml_fill_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + float c) { + return ggml_fill_impl(ctx, a, c, true); +} + +// ggml_argsort + +struct ggml_tensor * ggml_argsort( + struct ggml_context * ctx, + struct ggml_tensor * a, + enum ggml_sort_order order) { + GGML_ASSERT(a->ne[0] <= INT32_MAX); + + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_I32, GGML_MAX_DIMS, a->ne); + + ggml_set_op_params_i32(result, 0, (int32_t) order); + + result->op = GGML_OP_ARGSORT; + result->src[0] = a; + + return result; +} + +// ggml_argsort_top_k + +struct ggml_tensor * ggml_argsort_top_k( + struct ggml_context * ctx, + struct ggml_tensor * a, + int k) { + GGML_ASSERT(a->ne[0] >= k); + + struct ggml_tensor * result = ggml_argsort(ctx, a, GGML_SORT_ORDER_DESC); + + result = ggml_view_4d(ctx, result, + k, result->ne[1], result->ne[2], result->ne[3], + result->nb[1], result->nb[2], result->nb[3], + 0); + + return result; +} + +// ggml_top_k + +struct ggml_tensor * ggml_top_k( + struct ggml_context * ctx, + struct ggml_tensor * a, + int k) { + GGML_ASSERT(a->ne[0] >= k); + + struct ggml_tensor * result = ggml_new_tensor_4d(ctx, GGML_TYPE_I32, k, a->ne[1], a->ne[2], a->ne[3]); + + result->op = GGML_OP_TOP_K; + result->src[0] = a; + + return result; +} + +// ggml_arange + +struct ggml_tensor * ggml_arange( + struct ggml_context * ctx, + float start, + float stop, + float step) { + GGML_ASSERT(stop > start); + + const int64_t steps = (int64_t) ceilf((stop - start) / step); + + struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, steps); + + ggml_set_op_params_f32(result, 0, start); + ggml_set_op_params_f32(result, 1, stop); + ggml_set_op_params_f32(result, 2, step); + + result->op = GGML_OP_ARANGE; + + return result; +} + +// ggml_flash_attn_ext + +struct ggml_tensor * ggml_flash_attn_ext( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * mask, + float scale, + float max_bias, + float logit_softcap) { + GGML_ASSERT(ggml_can_mul_mat(k, q)); + // TODO: check if vT can be multiplied by (k*qT) + + GGML_ASSERT(q->ne[3] == k->ne[3]); + GGML_ASSERT(q->ne[3] == v->ne[3]); + + if (mask) { + GGML_ASSERT(mask->type == GGML_TYPE_F16); + GGML_ASSERT(ggml_is_contiguous(mask)); + //GGML_ASSERT(ggml_can_repeat_rows(mask, qk)); + + GGML_ASSERT(q->ne[2] % mask->ne[2] == 0); + GGML_ASSERT(q->ne[3] % mask->ne[3] == 0); + } + + if (max_bias > 0.0f) { + GGML_ASSERT(mask); + } + + // permute(0, 2, 1, 3) + int64_t ne[4] = { v->ne[0], q->ne[2], q->ne[1], q->ne[3] }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + float params[] = { scale, max_bias, logit_softcap }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_FLASH_ATTN_EXT; + result->src[0] = q; + result->src[1] = k; + result->src[2] = v; + result->src[3] = mask; + + return result; +} + +void ggml_flash_attn_ext_set_prec( + struct ggml_tensor * a, + enum ggml_prec prec) { + GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); + + const int32_t prec_i32 = (int32_t) prec; + + ggml_set_op_params_i32(a, 3, prec_i32); // scale is on first pos, max_bias on second +} + +enum ggml_prec ggml_flash_attn_ext_get_prec( + const struct ggml_tensor * a) { + GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); + + const int32_t prec_i32 = ggml_get_op_params_i32(a, 3); + + return (enum ggml_prec) prec_i32; +} + +void ggml_flash_attn_ext_add_sinks( + struct ggml_tensor * a, + struct ggml_tensor * sinks) { + if (!sinks) { + a->src[4] = NULL; + return; + } + + GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); + GGML_ASSERT(a->src[4] == NULL); + GGML_ASSERT(a->src[0]->ne[2] == sinks->ne[0]); + GGML_ASSERT(sinks->type == GGML_TYPE_F32); + + a->src[4] = sinks; +} + +// ggml_flash_attn_sparse + +struct ggml_tensor * ggml_flash_attn_sparse( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + float scale, + float alpha) { + + GGML_ASSERT(ggml_can_mul_mat(k, q)); + + // Output shape matches ggml_flash_attn_ext: [v->ne[0], q->ne[2], q->ne[1], q->ne[3]] + int64_t ne[4] = { v->ne[0], q->ne[2], q->ne[1], q->ne[3] }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + float params[] = { scale, alpha }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_FLASH_ATTN_SPARSE; + result->src[0] = q; + result->src[1] = k; + result->src[2] = v; + + return result; +} + +// ggml_flash_attn_back + +struct ggml_tensor * ggml_flash_attn_back( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * d, + bool masked) { + GGML_ABORT("TODO: adapt to ggml_flash_attn_ext() changes"); + + GGML_ASSERT(ggml_can_mul_mat(k, q)); + // TODO: check if vT can be multiplied by (k*qT) + + // d shape [D,N,ne2,ne3] + // q shape [D,N,ne2,ne3] + // k shape [D,M,kvne2,ne3] + // v shape [M,D,kvne2,ne3] + + const int64_t D = q->ne[0]; + const int64_t N = q->ne[1]; + const int64_t M = k->ne[1]; + const int64_t ne2 = q->ne[2]; + const int64_t ne3 = q->ne[3]; + const int64_t kvne2 = k->ne[2]; + + GGML_ASSERT(k->ne[0] == D); + GGML_ASSERT(v->ne[0] == M); + GGML_ASSERT(v->ne[1] == D); + GGML_ASSERT(d->ne[0] == D); + GGML_ASSERT(d->ne[1] == N); + GGML_ASSERT(k->ne[2] == kvne2); + GGML_ASSERT(k->ne[3] == ne3); + GGML_ASSERT(v->ne[2] == kvne2); + GGML_ASSERT(v->ne[3] == ne3); + GGML_ASSERT(d->ne[2] == ne2); + GGML_ASSERT(d->ne[3] == ne3); + + GGML_ASSERT(ne2 % kvne2 == 0); + + // store gradients of q, k and v as continuous tensors concatenated in result. + // note: v and gradv are actually transposed, i.e. v->ne[0] != D. + const int64_t elem_q = ggml_nelements(q); + const int64_t elem_k = ggml_nelements(k); + const int64_t elem_v = ggml_nelements(v); + + enum ggml_type result_type = GGML_TYPE_F32; + GGML_ASSERT(ggml_blck_size(result_type) == 1); + const size_t tsize = ggml_type_size(result_type); + + const size_t offs_q = 0; + const size_t offs_k = offs_q + GGML_PAD(elem_q * tsize, GGML_MEM_ALIGN); + const size_t offs_v = offs_k + GGML_PAD(elem_k * tsize, GGML_MEM_ALIGN); + const size_t end = offs_v + GGML_PAD(elem_v * tsize, GGML_MEM_ALIGN); + + const size_t nelements = (end + tsize - 1)/tsize; + + struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, nelements); + + int32_t masked_i = masked ? 1 : 0; + ggml_set_op_params(result, &masked_i, sizeof(masked_i)); + + result->op = GGML_OP_FLASH_ATTN_BACK; + result->src[0] = q; + result->src[1] = k; + result->src[2] = v; + result->src[3] = d; + + return result; +} + +// ggml_ssm_conv + +struct ggml_tensor * ggml_ssm_conv( + struct ggml_context * ctx, + struct ggml_tensor * sx, + struct ggml_tensor * c) { + GGML_ASSERT(ggml_is_3d(sx)); + GGML_ASSERT(ggml_is_matrix(c)); + + const int64_t d_conv = c->ne[0]; + const int64_t d_inner = c->ne[1]; + const int64_t n_t = sx->ne[0] - d_conv + 1; // tokens per sequence + const int64_t n_s = sx->ne[2]; + + // TODO: maybe support other strides than 1? + GGML_ASSERT(sx->ne[0] == d_conv - 1 + n_t); + GGML_ASSERT(sx->ne[1] == d_inner); + GGML_ASSERT(n_t >= 0); + + struct ggml_tensor * result = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, d_inner, n_t, n_s); + + result->op = GGML_OP_SSM_CONV; + result->src[0] = sx; + result->src[1] = c; + + return result; +} + +// dflash: tree-mode variant. Same op, with parent_ids plumbed into +// src[2] so the CUDA kernel gathers each token's window along its tree +// parent chain instead of the DFS-neighbour window. +struct ggml_tensor * ggml_ssm_conv_tree( + struct ggml_context * ctx, + struct ggml_tensor * sx, + struct ggml_tensor * c, + struct ggml_tensor * parent_ids) { + struct ggml_tensor * result = ggml_ssm_conv(ctx, sx, c); + + GGML_ASSERT(parent_ids != NULL); + GGML_ASSERT(parent_ids->type == GGML_TYPE_I32); + GGML_ASSERT(ggml_is_contiguous(parent_ids)); + + const int64_t n_t = sx->ne[0] - c->ne[0] + 1; + const int64_t n_s = sx->ne[2]; + GGML_ASSERT(ggml_nelements(parent_ids) == n_t * n_s); + + result->src[2] = parent_ids; + + return result; +} + +// ggml_ssm_scan + +struct ggml_tensor * ggml_ssm_scan( + struct ggml_context * ctx, + struct ggml_tensor * s, + struct ggml_tensor * x, + struct ggml_tensor * dt, + struct ggml_tensor * A, + struct ggml_tensor * B, + struct ggml_tensor * C, + struct ggml_tensor * ids) { + GGML_ASSERT(ggml_is_contiguous(s)); + GGML_ASSERT(ggml_is_contiguous(dt)); + GGML_ASSERT(ggml_is_contiguous(A)); + GGML_ASSERT(x->nb[0] == ggml_type_size(x->type)); + GGML_ASSERT(B->nb[0] == ggml_type_size(B->type)); + GGML_ASSERT(C->nb[0] == ggml_type_size(C->type)); + GGML_ASSERT(x->nb[1] == x->ne[0]*x->nb[0]); + GGML_ASSERT(B->nb[1] == B->ne[0]*B->nb[0]); + GGML_ASSERT(C->nb[1] == C->ne[0]*C->nb[0]); + GGML_ASSERT(ggml_are_same_shape(B, C)); + GGML_ASSERT(ids->type == GGML_TYPE_I32); + + { + const int64_t d_state = s->ne[0]; + const int64_t head_dim = x->ne[0]; + const int64_t n_head = x->ne[1]; + const int64_t n_seq_tokens = x->ne[2]; + const int64_t n_seqs = x->ne[3]; + + GGML_ASSERT(dt->ne[0] == n_head); + GGML_ASSERT(dt->ne[1] == n_seq_tokens); + GGML_ASSERT(dt->ne[2] == n_seqs); + GGML_ASSERT(ggml_is_3d(dt)); + GGML_ASSERT(s->ne[1] == head_dim); + GGML_ASSERT(s->ne[2] == n_head); + GGML_ASSERT(B->ne[0] == d_state); + GGML_ASSERT(B->ne[2] == n_seq_tokens); + GGML_ASSERT(B->ne[3] == n_seqs); + GGML_ASSERT(ids->ne[0] == n_seqs); + GGML_ASSERT(ggml_is_vector(ids)); + GGML_ASSERT(A->ne[1] == n_head); + GGML_ASSERT(ggml_is_matrix(A)); + + if (A->ne[0] != 1) { + // Mamba-1 has more granular decay factors + GGML_ASSERT(A->ne[0] == d_state); + } + } + + // concatenated y + ssm_states + struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ggml_nelements(x) + s->ne[0]*s->ne[1]*s->ne[2]*ids->ne[0]); + + result->op = GGML_OP_SSM_SCAN; + result->src[0] = s; + result->src[1] = x; + result->src[2] = dt; + result->src[3] = A; + result->src[4] = B; + result->src[5] = C; + result->src[6] = ids; + + return result; +} + +// ggml_win_part + +struct ggml_tensor * ggml_win_part( + struct ggml_context * ctx, + struct ggml_tensor * a, + int w) { + GGML_ASSERT(a->ne[3] == 1); + GGML_ASSERT(a->type == GGML_TYPE_F32); + + // padding + const int px = (w - a->ne[1]%w)%w; + const int py = (w - a->ne[2]%w)%w; + + const int npx = (px + a->ne[1])/w; + const int npy = (py + a->ne[2])/w; + const int np = npx*npy; + + const int64_t ne[4] = { a->ne[0], w, w, np, }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + int32_t params[] = { npx, npy, w }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_WIN_PART; + result->src[0] = a; + + return result; +} + +// ggml_win_unpart + +struct ggml_tensor * ggml_win_unpart( + struct ggml_context * ctx, + struct ggml_tensor * a, + int w0, + int h0, + int w) { + GGML_ASSERT(a->type == GGML_TYPE_F32); + + const int64_t ne[4] = { a->ne[0], w0, h0, 1, }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 3, ne); + + int32_t params[] = { w }; + ggml_set_op_params(result, params, sizeof(params)); + + result->op = GGML_OP_WIN_UNPART; + result->src[0] = a; + + return result; +} + +// ggml_get_rel_pos + +struct ggml_tensor * ggml_get_rel_pos( + struct ggml_context * ctx, + struct ggml_tensor * a, + int qh, + int kh) { + GGML_ASSERT(qh == kh); + GGML_ASSERT(2*MAX(qh, kh) - 1 == a->ne[1]); + + const int64_t ne[4] = { a->ne[0], kh, qh, 1, }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F16, 3, ne); + + result->op = GGML_OP_GET_REL_POS; + result->src[0] = a; + + return result; +} + +// ggml_add_rel_pos + +static struct ggml_tensor * ggml_add_rel_pos_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * pw, + struct ggml_tensor * ph, + bool inplace) { + GGML_ASSERT(ggml_are_same_shape(pw, ph)); + GGML_ASSERT(ggml_is_contiguous(a)); + GGML_ASSERT(ggml_is_contiguous(pw)); + GGML_ASSERT(ggml_is_contiguous(ph)); + GGML_ASSERT(ph->type == GGML_TYPE_F32); + GGML_ASSERT(pw->type == GGML_TYPE_F32); + GGML_ASSERT(pw->ne[3] == a->ne[2]); + GGML_ASSERT(pw->ne[0]*pw->ne[0] == a->ne[0]); + GGML_ASSERT(pw->ne[1]*pw->ne[2] == a->ne[1]); + + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + ggml_set_op_params_i32(result, 0, inplace ? 1 : 0); + + result->op = GGML_OP_ADD_REL_POS; + result->src[0] = a; + result->src[1] = pw; + result->src[2] = ph; + + return result; +} + +struct ggml_tensor * ggml_add_rel_pos( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * pw, + struct ggml_tensor * ph) { + return ggml_add_rel_pos_impl(ctx, a, pw, ph, false); +} + +struct ggml_tensor * ggml_add_rel_pos_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * pw, + struct ggml_tensor * ph) { + return ggml_add_rel_pos_impl(ctx, a, pw, ph, true); +} + +// ggml_rwkv_wkv6 + +struct ggml_tensor * ggml_rwkv_wkv6( + struct ggml_context * ctx, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * r, + struct ggml_tensor * tf, + struct ggml_tensor * td, + struct ggml_tensor * state) { + GGML_ASSERT(ggml_is_contiguous(k)); + GGML_ASSERT(ggml_is_contiguous(v)); + GGML_ASSERT(ggml_is_contiguous(r)); + GGML_ASSERT(ggml_is_contiguous(tf)); + GGML_ASSERT(ggml_is_contiguous(td)); + GGML_ASSERT(ggml_is_contiguous(state)); + + const int64_t S = k->ne[0]; + const int64_t H = k->ne[1]; + const int64_t n_tokens = k->ne[2]; + const int64_t n_seqs = state->ne[1]; + { + GGML_ASSERT(v->ne[0] == S && v->ne[1] == H && v->ne[2] == n_tokens); + GGML_ASSERT(r->ne[0] == S && r->ne[1] == H && r->ne[2] == n_tokens); + GGML_ASSERT(td->ne[0] == S && td->ne[1] == H && td->ne[2] == n_tokens); + GGML_ASSERT(ggml_nelements(state) == S * S * H * n_seqs); + } + + // concat output and new_state + const int64_t ne[4] = { S * H, n_tokens + S * n_seqs, 1, 1 }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + result->op = GGML_OP_RWKV_WKV6; + result->src[0] = k; + result->src[1] = v; + result->src[2] = r; + result->src[3] = tf; + result->src[4] = td; + result->src[5] = state; + + return result; +} + +// ggml_gated_linear_attn + +struct ggml_tensor * ggml_gated_linear_attn( + struct ggml_context * ctx, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * q, + struct ggml_tensor * g, + struct ggml_tensor * state, + float scale) { + GGML_ASSERT(ggml_is_contiguous(k)); + GGML_ASSERT(ggml_is_contiguous(v)); + GGML_ASSERT(ggml_is_contiguous(q)); + GGML_ASSERT(ggml_is_contiguous(g)); + GGML_ASSERT(ggml_is_contiguous(state)); + + const int64_t S = k->ne[0]; + const int64_t H = k->ne[1]; + const int64_t n_tokens = k->ne[2]; + const int64_t n_seqs = state->ne[1]; + { + GGML_ASSERT(v->ne[0] == S && v->ne[1] == H && v->ne[2] == n_tokens); + GGML_ASSERT(q->ne[0] == S && q->ne[1] == H && q->ne[2] == n_tokens); + GGML_ASSERT(g->ne[0] == S && g->ne[1] == H && g->ne[2] == n_tokens); + GGML_ASSERT(ggml_nelements(state) == S * S * H * n_seqs); + } + + // concat output and new_state + const int64_t ne[4] = { S * H, n_tokens + S * n_seqs, 1, 1 }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + ggml_set_op_params_f32(result, 0, scale); + + result->op = GGML_OP_GATED_LINEAR_ATTN; + result->src[0] = k; + result->src[1] = v; + result->src[2] = q; + result->src[3] = g; + result->src[4] = state; + + return result; +} + +// ggml_rwkv_wkv7 + +struct ggml_tensor * ggml_rwkv_wkv7( + struct ggml_context * ctx, + struct ggml_tensor * r, + struct ggml_tensor * w, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * state) { + GGML_ASSERT(ggml_is_contiguous(r)); + GGML_ASSERT(ggml_is_contiguous(w)); + GGML_ASSERT(ggml_is_contiguous(k)); + GGML_ASSERT(ggml_is_contiguous(v)); + GGML_ASSERT(ggml_is_contiguous(a)); + GGML_ASSERT(ggml_is_contiguous(b)); + GGML_ASSERT(ggml_is_contiguous(state)); + + const int64_t S = k->ne[0]; + const int64_t H = k->ne[1]; + const int64_t n_tokens = k->ne[2]; + const int64_t n_seqs = state->ne[1]; + { + GGML_ASSERT(w->ne[0] == S && w->ne[1] == H && w->ne[2] == n_tokens); + GGML_ASSERT(k->ne[0] == S && k->ne[1] == H && k->ne[2] == n_tokens); + GGML_ASSERT(v->ne[0] == S && v->ne[1] == H && v->ne[2] == n_tokens); + GGML_ASSERT(a->ne[0] == S && a->ne[1] == H && a->ne[2] == n_tokens); + GGML_ASSERT(b->ne[0] == S && b->ne[1] == H && b->ne[2] == n_tokens); + GGML_ASSERT(ggml_nelements(state) == S * S * H * n_seqs); + } + + // concat output and new_state + const int64_t ne[4] = { S * H, n_tokens + S * n_seqs, 1, 1 }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + result->op = GGML_OP_RWKV_WKV7; + result->src[0] = r; + result->src[1] = w; + result->src[2] = k; + result->src[3] = v; + result->src[4] = a; + result->src[5] = b; + result->src[6] = state; + + return result; +} + +// ggml_unary + +static struct ggml_tensor * ggml_unary_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + enum ggml_unary_op op, + bool inplace) { + GGML_ASSERT(ggml_is_contiguous_rows(a)); + + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + ggml_set_op_params_i32(result, 0, (int32_t) op); + + result->op = GGML_OP_UNARY; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_unary( + struct ggml_context * ctx, + struct ggml_tensor * a, + enum ggml_unary_op op) { + return ggml_unary_impl(ctx, a, op, false); +} + +struct ggml_tensor * ggml_unary_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + enum ggml_unary_op op) { + return ggml_unary_impl(ctx, a, op, true); +} + +// ggml_map_custom1 + +static struct ggml_tensor * ggml_map_custom1_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + const ggml_custom1_op_t fun, + int n_tasks, + void * userdata, + bool inplace) { + GGML_ASSERT(n_tasks == GGML_N_TASKS_MAX || n_tasks > 0); + + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + struct ggml_map_custom1_op_params params = { + /*.fun =*/ fun, + /*.n_tasks =*/ n_tasks, + /*.userdata =*/ userdata + }; + ggml_set_op_params(result, ¶ms, sizeof(params)); + + result->op = GGML_OP_MAP_CUSTOM1; + result->src[0] = a; + + return result; +} + +struct ggml_tensor * ggml_map_custom1( + struct ggml_context * ctx, + struct ggml_tensor * a, + const ggml_custom1_op_t fun, + int n_tasks, + void * userdata) { + return ggml_map_custom1_impl(ctx, a, fun, n_tasks, userdata, false); +} + +struct ggml_tensor * ggml_map_custom1_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + const ggml_custom1_op_t fun, + int n_tasks, + void * userdata) { + return ggml_map_custom1_impl(ctx, a, fun, n_tasks, userdata, true); +} + +// ggml_map_custom2 + +static struct ggml_tensor * ggml_map_custom2_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + const ggml_custom2_op_t fun, + int n_tasks, + void * userdata, + bool inplace) { + GGML_ASSERT(n_tasks == GGML_N_TASKS_MAX || n_tasks > 0); + + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + struct ggml_map_custom2_op_params params = { + /*.fun =*/ fun, + /*.n_tasks =*/ n_tasks, + /*.userdata =*/ userdata + }; + ggml_set_op_params(result, ¶ms, sizeof(params)); + + result->op = GGML_OP_MAP_CUSTOM2; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +struct ggml_tensor * ggml_map_custom2( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + const ggml_custom2_op_t fun, + int n_tasks, + void * userdata) { + return ggml_map_custom2_impl(ctx, a, b, fun, n_tasks, userdata, false); +} + +struct ggml_tensor * ggml_map_custom2_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + const ggml_custom2_op_t fun, + int n_tasks, + void * userdata) { + return ggml_map_custom2_impl(ctx, a, b, fun, n_tasks, userdata, true); +} + +// ggml_map_custom3 + +static struct ggml_tensor * ggml_map_custom3_impl( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c, + const ggml_custom3_op_t fun, + int n_tasks, + void * userdata, + bool inplace) { + GGML_ASSERT(n_tasks == GGML_N_TASKS_MAX || n_tasks > 0); + + struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); + + struct ggml_map_custom3_op_params params = { + /*.fun =*/ fun, + /*.n_tasks =*/ n_tasks, + /*.userdata =*/ userdata + }; + ggml_set_op_params(result, ¶ms, sizeof(params)); + + result->op = GGML_OP_MAP_CUSTOM3; + result->src[0] = a; + result->src[1] = b; + result->src[2] = c; + + return result; +} + +struct ggml_tensor * ggml_map_custom3( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c, + const ggml_custom3_op_t fun, + int n_tasks, + void * userdata) { + return ggml_map_custom3_impl(ctx, a, b, c, fun, n_tasks, userdata, false); +} + +struct ggml_tensor * ggml_map_custom3_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c, + const ggml_custom3_op_t fun, + int n_tasks, + void * userdata) { + return ggml_map_custom3_impl(ctx, a, b, c, fun, n_tasks, userdata, true); +} + +struct ggml_tensor * ggml_custom_4d( + struct ggml_context * ctx, + enum ggml_type type, + int64_t ne0, + int64_t ne1, + int64_t ne2, + int64_t ne3, + struct ggml_tensor ** args, + int n_args, + ggml_custom_op_t fun, + int n_tasks, + void * userdata) { + + GGML_ASSERT(n_args < GGML_MAX_SRC); + + struct ggml_tensor * result = ggml_new_tensor_4d(ctx, type, ne0, ne1, ne2, ne3); + + struct ggml_custom_op_params params = { + /*.fun =*/ fun, + /*.n_tasks =*/ n_tasks, + /*.userdata =*/ userdata + }; + ggml_set_op_params(result, ¶ms, sizeof(params)); + + result->op = GGML_OP_CUSTOM; + for (int i = 0; i < n_args; i++) { + result->src[i] = args[i]; + } + + return result; +} + +struct ggml_tensor * ggml_custom_inplace( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor ** args, + int n_args, + ggml_custom_op_t fun, + int n_tasks, + void * userdata) { + + GGML_ASSERT(n_args < GGML_MAX_SRC - 1); + + struct ggml_tensor * result = ggml_view_tensor(ctx, a); + + struct ggml_custom_op_params params = { + /*.fun =*/ fun, + /*.n_tasks =*/ n_tasks, + /*.userdata =*/ userdata + }; + ggml_set_op_params(result, ¶ms, sizeof(params)); + + result->op = GGML_OP_CUSTOM; + result->src[0] = a; + for (int i = 0; i < n_args; i++) { + result->src[i + 1] = args[i]; + } + + return result; +} +// ggml_cross_entropy_loss + +struct ggml_tensor * ggml_cross_entropy_loss( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b) { + GGML_ASSERT(ggml_are_same_shape(a, b)); + + struct ggml_tensor * result = ggml_new_tensor_1d(ctx, a->type, 1); + + result->op = GGML_OP_CROSS_ENTROPY_LOSS; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +// ggml_cross_entropy_loss_back + +struct ggml_tensor * ggml_cross_entropy_loss_back( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + struct ggml_tensor * c) { + GGML_ASSERT(ggml_is_scalar(a)); + GGML_ASSERT(ggml_are_same_shape(b, c)); + + struct ggml_tensor * result = ggml_dup_tensor(ctx, b); + + result->op = GGML_OP_CROSS_ENTROPY_LOSS_BACK; + result->src[0] = a; + result->src[1] = b; + result->src[2] = c; + + return result; +} + +// opt_step_adamw + +struct ggml_tensor * ggml_opt_step_adamw( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * grad, + struct ggml_tensor * m, + struct ggml_tensor * v, + struct ggml_tensor * adamw_params) { + GGML_ASSERT(a->flags & GGML_TENSOR_FLAG_PARAM); + GGML_ASSERT(ggml_are_same_shape(a, grad)); + GGML_ASSERT(ggml_are_same_shape(a, m)); + GGML_ASSERT(ggml_are_same_shape(a, v)); + GGML_ASSERT(adamw_params->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_nelements(adamw_params) == 7); + + struct ggml_tensor * result = ggml_view_tensor(ctx, a); + + result->op = GGML_OP_OPT_STEP_ADAMW; + result->src[0] = a; + result->src[1] = grad; + result->src[2] = m; + result->src[3] = v; + result->src[4] = adamw_params; + + return result; +} + +// opt_step_sgd + +struct ggml_tensor * ggml_opt_step_sgd( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * grad, + struct ggml_tensor * params) { + GGML_ASSERT(a->flags & GGML_TENSOR_FLAG_PARAM); + GGML_ASSERT(ggml_are_same_shape(a, grad)); + GGML_ASSERT(params->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_nelements(params) == 2); + + struct ggml_tensor * result = ggml_view_tensor(ctx, a); + + result->op = GGML_OP_OPT_STEP_SGD; + result->src[0] = a; + result->src[1] = grad; + result->src[2] = params; + + return result; +} + +// solve_tri + +struct ggml_tensor * ggml_solve_tri( + struct ggml_context * ctx, + struct ggml_tensor * a, + struct ggml_tensor * b, + bool left, + bool lower, + bool uni) { + GGML_ASSERT(a->type == GGML_TYPE_F32); + GGML_ASSERT(b->type == GGML_TYPE_F32); + + // A must be square and lower diagonal + GGML_ASSERT(a->ne[0] == a->ne[1]); + // B must have same outer dimension as A + GGML_ASSERT(a->ne[1] == b->ne[1]); + + // batch dimensions must be equal + GGML_ASSERT(a->ne[2] == b->ne[2]); + GGML_ASSERT(a->ne[3] == b->ne[3]); + + GGML_ASSERT(ggml_is_contiguous(a)); + GGML_ASSERT(ggml_is_contiguous(b)); + + GGML_ASSERT(lower && left && !uni); // TODO: support other variants + + struct ggml_tensor * result = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, b->ne[0], b->ne[1], b->ne[2], b->ne[3]); + + result->op = GGML_OP_SOLVE_TRI; + result->src[0] = a; + result->src[1] = b; + + return result; +} + +// ggml_gated_delta_net + +struct ggml_tensor * ggml_gated_delta_net( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * g, + struct ggml_tensor * beta, + struct ggml_tensor * state) { + GGML_ASSERT(ggml_is_contiguous_rows(q)); + GGML_ASSERT(ggml_is_contiguous_rows(k)); + GGML_ASSERT(ggml_is_contiguous_rows(v)); + GGML_ASSERT(ggml_is_contiguous(g)); + GGML_ASSERT(ggml_is_contiguous(beta)); + GGML_ASSERT(ggml_is_contiguous(state)); + + GGML_ASSERT(q->type == GGML_TYPE_F32); + GGML_ASSERT(k->type == GGML_TYPE_F32); + GGML_ASSERT(v->type == GGML_TYPE_F32); + GGML_ASSERT(g->type == GGML_TYPE_F32); + GGML_ASSERT(beta->type == GGML_TYPE_F32); + GGML_ASSERT(state->type == GGML_TYPE_F32); + + const int64_t S_v = v->ne[0]; + const int64_t H = v->ne[1]; + const int64_t n_tokens = v->ne[2]; + const int64_t n_seqs = v->ne[3]; + + // gate: scalar [1, H, T, B] or vector [S_v, H, T, B] (KDA) + GGML_ASSERT(g->ne[0] == 1 || g->ne[0] == S_v); + GGML_ASSERT(beta->ne[0] == 1); + + GGML_ASSERT(ggml_nelements(state) == S_v * S_v * H * n_seqs); + + // Pack output, final new_state, and per-step intermediate states into one tensor. + // Layout (in units of `S_v * H`-wide rows): + // [ attn_output: n_tokens*n_seqs | final_state: S_v*n_seqs | intermediate_states: S_v*n_tokens*n_seqs ] + // + // The final_state slot is kept for backward compatibility with stock llama.cpp + // callers that read state at offset S_v*H*n_tokens*n_seqs. The intermediate_states + // region is a dflash extension: for each token t in [0, n_tokens), it holds the + // recurrent state after processing token t. Used by the spec decoding loop to + // roll back SSM state to the accepted prefix without a full replay forward pass. + const int64_t ne[4] = { S_v * H, n_tokens * n_seqs + S_v * n_seqs + S_v * n_tokens * n_seqs, 1, 1 }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + result->op = GGML_OP_GATED_DELTA_NET; + result->src[0] = q; + result->src[1] = k; + result->src[2] = v; + result->src[3] = g; + result->src[4] = beta; + result->src[5] = state; + + return result; +} + +void ggml_gated_delta_net_set_skip_intermediate( + struct ggml_tensor * tensor, + bool skip_intermediate) { + GGML_ASSERT(tensor != NULL); + GGML_ASSERT(tensor->op == GGML_OP_GATED_DELTA_NET); + ggml_set_op_params_i32(tensor, 0, skip_intermediate ? 1 : 0); + + const struct ggml_tensor * v = tensor->src[2]; + GGML_ASSERT(v != NULL); + + const int64_t S_v = v->ne[0]; + const int64_t n_tokens = v->ne[2]; + const int64_t n_seqs = v->ne[3]; + + // Compact only the plain chain path. Tree/persistent variants need + // intermediate states for branch reloads or explicit capture storage. + const bool can_compact = tensor->src[6] == NULL && tensor->src[7] == NULL; + tensor->ne[1] = n_tokens*n_seqs + S_v*n_seqs; + if (!skip_intermediate || !can_compact) { + tensor->ne[1] += S_v*n_tokens*n_seqs; + } + tensor->nb[2] = tensor->nb[1]*tensor->ne[1]; + tensor->nb[3] = tensor->nb[2]*tensor->ne[2]; +} + +// dflash: tree-mode variant. Same op, with parent_ids plumbed into +// src[6] so the CUDA kernel can branch-reload state at DFS transitions. +struct ggml_tensor * ggml_gated_delta_net_tree( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * g, + struct ggml_tensor * beta, + struct ggml_tensor * state, + struct ggml_tensor * parent_ids) { + struct ggml_tensor * result = ggml_gated_delta_net(ctx, q, k, v, g, beta, state); + + GGML_ASSERT(parent_ids != NULL); + GGML_ASSERT(parent_ids->type == GGML_TYPE_I32); + GGML_ASSERT(ggml_is_contiguous(parent_ids)); + + const int64_t n_tokens = v->ne[2]; + const int64_t n_seqs = v->ne[3]; + GGML_ASSERT(ggml_nelements(parent_ids) == n_tokens * n_seqs); + + result->src[6] = parent_ids; + + return result; +} + +// dflash: tree-mode + external persistent intermediate buffer. The +// kernel writes per-token intermediate states DIRECTLY into persist_inter's +// memory, skipping the cost of copying them out of the result tensor's +// internal region after graph_compute. +struct ggml_tensor * ggml_gated_delta_net_tree_persist( + struct ggml_context * ctx, + struct ggml_tensor * q, + struct ggml_tensor * k, + struct ggml_tensor * v, + struct ggml_tensor * g, + struct ggml_tensor * beta, + struct ggml_tensor * state, + struct ggml_tensor * parent_ids, + struct ggml_tensor * persist_inter) { + struct ggml_tensor * result = ggml_gated_delta_net_tree( + ctx, q, k, v, g, beta, state, parent_ids); + + GGML_ASSERT(persist_inter != NULL); + GGML_ASSERT(persist_inter->type == GGML_TYPE_F32 || + persist_inter->type == GGML_TYPE_F16); + GGML_ASSERT(ggml_is_contiguous(persist_inter)); + + const int64_t S_v = v->ne[0]; + const int64_t H = v->ne[1]; + const int64_t n_tokens = v->ne[2]; + const int64_t n_seqs = v->ne[3]; + GGML_ASSERT(ggml_nelements(persist_inter) >= S_v * S_v * H * n_tokens * n_seqs); + + result->src[7] = persist_inter; + + return result; +} + +//////////////////////////////////////////////////////////////////////////////// + +struct ggml_hash_set ggml_hash_set_new(size_t size) { + size = ggml_hash_size(size); + struct ggml_hash_set result; + result.size = size; + result.keys = GGML_MALLOC(sizeof(struct ggml_tensor *) * size); + result.used = GGML_CALLOC(ggml_bitset_size(size), sizeof(ggml_bitset_t)); + return result; +} + +void ggml_hash_set_reset(struct ggml_hash_set * hash_set) { + memset(hash_set->used, 0, sizeof(ggml_bitset_t) * ggml_bitset_size(hash_set->size)); +} + +void ggml_hash_set_free(struct ggml_hash_set * hash_set) { + GGML_FREE(hash_set->used); + GGML_FREE(hash_set->keys); +} + +size_t ggml_hash_size(size_t min_sz) { + // next primes after powers of two + static const size_t primes[] = { + 2, 3, 5, 11, 17, 37, 67, 131, 257, 521, 1031, + 2053, 4099, 8209, 16411, 32771, 65537, 131101, + 262147, 524309, 1048583, 2097169, 4194319, 8388617, + 16777259, 33554467, 67108879, 134217757, 268435459, + 536870923, 1073741827, 2147483659 + }; + static const size_t n_primes = sizeof(primes)/sizeof(primes[0]); + + // find the smallest prime that is larger or equal than min_sz + size_t l = 0; + size_t r = n_primes; + while (l < r) { + size_t m = (l + r)/2; + if (primes[m] < min_sz) { + l = m + 1; + } else { + r = m; + } + } + size_t sz = l < n_primes ? primes[l] : min_sz | 1; + return sz; +} + +struct hash_map { + struct ggml_hash_set set; + struct ggml_tensor ** vals; +}; + +static struct hash_map * ggml_new_hash_map(size_t size) { + struct hash_map * result = GGML_MALLOC(sizeof(struct hash_map)); + result->set = ggml_hash_set_new(size); + result->vals = GGML_CALLOC(result->set.size, sizeof(struct ggml_tensor *)); + return result; +} + +static void ggml_hash_map_free(struct hash_map * map) { + ggml_hash_set_free(&map->set); + GGML_FREE(map->vals); + GGML_FREE(map); +} + +// utility functions to change gradients +// isrc is the index of tensor in cgraph->visited_has_set.keys +// the corresponding gradient (accumulators) are also at position isrc +// if tensor has a gradient accumulator, modify that accumulator in-place +// else if there is no gradient for tensor, set the corresponding value +// else, just add/subtract/etc. the gradients + +static void ggml_add_or_set( + struct ggml_context * ctx, + struct ggml_cgraph * cgraph, + size_t isrc, + struct ggml_tensor * tensor) { + struct ggml_tensor * src = cgraph->visited_hash_set.keys[isrc]; + GGML_ASSERT(src); + if (cgraph->grads[isrc]) { + cgraph->grads[isrc] = ggml_add_impl(ctx, cgraph->grads[isrc], tensor, /*inplace =*/ cgraph->grad_accs[isrc]); + } else { + cgraph->grads[isrc] = tensor; + } + ggml_format_name(cgraph->grads[isrc], "grad for %s", src->name); + ggml_build_forward_expand(cgraph, cgraph->grads[isrc]); +} + +static void ggml_acc_or_set( + struct ggml_context * ctx, + struct ggml_cgraph * cgraph, + size_t isrc, + struct ggml_tensor * tensor, + const size_t nb1, + const size_t nb2, + const size_t nb3, + const size_t offset) { + struct ggml_tensor * src = cgraph->visited_hash_set.keys[isrc]; + GGML_ASSERT(src); + if (cgraph->grads[isrc]) { + cgraph->grads[isrc] = ggml_acc_impl(ctx, cgraph->grads[isrc], tensor, nb1, nb2, nb3, offset, cgraph->grad_accs[isrc]); + } else { + struct ggml_tensor * a_zero = ggml_scale(ctx, src, 0.0f); // FIXME this is going to produce NaN if a contains inf/NaN + cgraph->grads[isrc] = ggml_acc_impl(ctx, a_zero, tensor, nb1, nb2, nb3, offset, false); + } + ggml_format_name(cgraph->grads[isrc], "grad for %s", cgraph->visited_hash_set.keys[isrc]->name); + ggml_build_forward_expand(cgraph, cgraph->grads[isrc]); +} + +static void ggml_add1_or_set( + struct ggml_context * ctx, + struct ggml_cgraph * cgraph, + size_t isrc, + struct ggml_tensor * tensor) { + struct ggml_tensor * src = cgraph->visited_hash_set.keys[isrc]; + GGML_ASSERT(src); + if (cgraph->grads[isrc]) { + cgraph->grads[isrc] = ggml_add1_impl(ctx, cgraph->grads[isrc], tensor, cgraph->grad_accs[isrc]); + } else { + cgraph->grads[isrc] = ggml_repeat(ctx, tensor, src); + } + ggml_format_name(cgraph->grads[isrc], "grad for %s", src->name); + ggml_build_forward_expand(cgraph, cgraph->grads[isrc]); +} + +static void ggml_sub_or_set( + struct ggml_context * ctx, + struct ggml_cgraph * cgraph, + size_t isrc, + struct ggml_tensor * tensor) { + struct ggml_tensor * src = cgraph->visited_hash_set.keys[isrc]; + GGML_ASSERT(src); + if (cgraph->grads[isrc]) { + cgraph->grads[isrc] = ggml_sub_impl(ctx, cgraph->grads[isrc], tensor, cgraph->grad_accs[isrc]); + } else { + cgraph->grads[isrc] = ggml_neg(ctx, tensor); + } + ggml_format_name(cgraph->grads[isrc], "grad for %s", src->name); + ggml_build_forward_expand(cgraph, cgraph->grads[isrc]); +} + +static void ggml_compute_backward( + struct ggml_context * ctx, struct ggml_cgraph * cgraph, int i, const bool * grads_needed) { + struct ggml_tensor * tensor = cgraph->nodes[i]; + struct ggml_tensor * grad = ggml_graph_get_grad(cgraph, tensor); + + if (!grad) { + return; + } + + struct ggml_tensor * src0 = tensor->src[0]; + struct ggml_tensor * src1 = tensor->src[1]; + struct ggml_tensor * src2 = tensor->src[2]; + struct ggml_hash_set * hash_set = &cgraph->visited_hash_set; + const size_t isrc0 = src0 ? ggml_hash_find(hash_set, src0) : (size_t) -1; + const size_t isrc1 = src1 ? ggml_hash_find(hash_set, src1) : (size_t) -1; + const size_t isrc2 = src2 ? ggml_hash_find(hash_set, src2) : (size_t) -1; + const bool src0_needs_grads = src0 && isrc0 != GGML_HASHSET_FULL && ggml_bitset_get(hash_set->used, isrc0) && grads_needed[isrc0]; + const bool src1_needs_grads = src1 && isrc1 != GGML_HASHSET_FULL && ggml_bitset_get(hash_set->used, isrc1) && grads_needed[isrc1]; + const bool src2_needs_grads = src2 && isrc2 != GGML_HASHSET_FULL && ggml_bitset_get(hash_set->used, isrc2) && grads_needed[isrc2]; + + switch (tensor->op) { + case GGML_OP_DUP: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, grad); + } + } break; + case GGML_OP_ADD: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, grad); + } + if (src1_needs_grads) { + struct ggml_tensor * tmp = grad; + if (!ggml_are_same_shape(src0, src1)) { + tmp = ggml_repeat_back(ctx, tmp, src1); + } + ggml_add_or_set(ctx, cgraph, isrc1, tmp); + } + } break; + case GGML_OP_ADD1: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, grad); + } + if (src1_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc1, ggml_mean(ctx, grad)); // TODO: should probably be sum instead of mean + } + } break; + case GGML_OP_ACC: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, grad); + } + if (src1_needs_grads) { + const size_t nb1 = ((int32_t *) tensor->op_params)[0]; + const size_t nb2 = ((int32_t *) tensor->op_params)[1]; + const size_t nb3 = ((int32_t *) tensor->op_params)[2]; + const size_t offset = ((int32_t *) tensor->op_params)[3]; + + struct ggml_tensor * tensor_grad_view = ggml_view_4d(ctx, + grad, src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3], + nb1, nb2, nb3, offset); + + ggml_add_or_set(ctx, cgraph, isrc1, ggml_reshape(ctx, ggml_cont(ctx, tensor_grad_view), src1)); + } + } break; + case GGML_OP_SUB: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, grad); + } + if (src1_needs_grads) { + ggml_sub_or_set(ctx, cgraph, isrc1, grad); + } + } break; + case GGML_OP_MUL: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_mul(ctx, grad, src1)); + } + if (src1_needs_grads) { + struct ggml_tensor * tmp = ggml_mul(ctx, src0, grad); + if (!ggml_are_same_shape(src0, src1)) { + tmp = ggml_repeat_back(ctx, tmp, src1); + } + ggml_add_or_set(ctx, cgraph, isrc1, tmp); + } + } break; + case GGML_OP_DIV: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_div(ctx, grad, src1)); + } + if (src1_needs_grads) { + ggml_sub_or_set(ctx, cgraph, isrc1, ggml_mul(ctx, grad, ggml_div(ctx, tensor, src1))); + } + } break; + case GGML_OP_SQR: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_scale(ctx, ggml_mul(ctx, src0, grad), 2.0f)); + } + } break; + case GGML_OP_SQRT: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_scale(ctx, ggml_div(ctx, grad, tensor), 0.5f)); + } + } break; + case GGML_OP_LOG: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_div(ctx, grad, src0)); + } + } break; + case GGML_OP_SIN: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_mul(ctx, grad, ggml_cos(ctx, src0))); + } + } break; + case GGML_OP_COS: { + if (src0_needs_grads) { + ggml_sub_or_set(ctx, cgraph, isrc0, ggml_mul(ctx, grad, ggml_sin(ctx, src0))); + } + } break; + case GGML_OP_SUM: { + if (src0_needs_grads) { + ggml_add1_or_set(ctx, cgraph, isrc0, grad); + } + } break; + case GGML_OP_SUM_ROWS: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_repeat(ctx, grad, src0)); + } + } break; + case GGML_OP_MEAN: { + if (src0_needs_grads) { + ggml_add1_or_set(ctx, cgraph, isrc0, ggml_scale_impl(ctx, grad, 1.0f/src0->ne[0], 0.0, false)); + } + } break; + case GGML_OP_REPEAT: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_repeat_back(ctx, grad, src0)); + } + } break; + case GGML_OP_REPEAT_BACK: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_repeat(ctx, grad, src0)); + } + } break; + case GGML_OP_RMS_NORM: { + if (src0_needs_grads) { + float eps; + memcpy(&eps, tensor->op_params, sizeof(float)); + ggml_add_or_set(ctx, cgraph, isrc0, ggml_rms_norm_back(ctx, grad, src0, eps)); + } + } break; + case GGML_OP_MUL_MAT: { + // https://cs231n.github.io/optimization-2/#staged + // # forward pass + // s0 = np.random.randn(5, 10) + // s1 = np.random.randn(10, 3) + // t = s0.dot(s1) + + // # now suppose we had the gradient on t from above in the circuit + // dt = np.random.randn(*t.shape) # same shape as t + // ds0 = dt.dot(s1.T) #.T gives the transpose of the matrix + // ds1 = t.T.dot(dt) + + // tensor.shape [m,p,qq,rr] + // src0.shape [n,m,q1,r1] + // src1.shape [n,p,qq,rr] + + if (src0_needs_grads) { + GGML_ASSERT(grad->ne[2] == src1->ne[2]); + GGML_ASSERT(grad->ne[3] == src1->ne[3]); + struct ggml_tensor * tmp = + ggml_out_prod(ctx, // [n,m,qq,rr] + src1, // [n,p,qq,rr] + grad); // [m,p,qq,rr] + if (!ggml_are_same_shape(tmp, src0)) { + GGML_ASSERT(tmp->ne[0] == src0->ne[0]); + GGML_ASSERT(tmp->ne[1] == src0->ne[1]); + GGML_ASSERT(tmp->ne[3] == 1); + + const int64_t nr2 = tmp->ne[2] / src0->ne[2]; + const size_t nb2 = tmp->nb[2] * nr2; + const size_t nb3 = tmp->nb[2]; + + tmp = ggml_view_4d(ctx, tmp, src0->ne[0], src0->ne[1], src0->ne[2], nr2, tmp->nb[1], nb2, nb3, 0); + tmp = ggml_repeat_back(ctx, tmp, src0); + } + ggml_add_or_set(ctx, cgraph, isrc0, tmp); + } + if (src1_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc1, + // ggml_mul_mat(ctx, // [n,p,qq,rr] + // ggml_cont(ctx, // [m,n,q1,r1] + // ggml_transpose(ctx, src0)), // [m,n,q1,r1] + // grad), // [m,p,qq,rr] + + // when src0 is bigger than tensor->grad (this is mostly the case in llama), + // avoid transpose of src0, rather transpose smaller tensor->grad + // and then use ggml_out_prod + ggml_out_prod(ctx, // [n,p,qq,rr] + src0, // [n,m,q1,r1] + ggml_transpose(ctx, // [p,m,qq,rr] + grad))); // [m,p,qq,rr] + } + } break; + case GGML_OP_SCALE: { + if (src0_needs_grads) { + float s; + memcpy(&s, tensor->op_params, sizeof(float)); + ggml_add_or_set(ctx, cgraph, isrc0, ggml_scale_impl(ctx, grad, s, 0.0, false)); + } + } break; + case GGML_OP_SET: { + const size_t nb1 = ((const int32_t *) tensor->op_params)[0]; + const size_t nb2 = ((const int32_t *) tensor->op_params)[1]; + const size_t nb3 = ((const int32_t *) tensor->op_params)[2]; + const size_t offset = ((const int32_t *) tensor->op_params)[3]; + + struct ggml_tensor * tensor_grad_view = NULL; + + if (src0_needs_grads || src1_needs_grads) { + GGML_ASSERT(src0->type == tensor->type); + GGML_ASSERT(!cgraph->grads[isrc0] || cgraph->grads[isrc0]->type == grad->type); + GGML_ASSERT(!cgraph->grads[isrc1] || !src1_needs_grads || cgraph->grads[isrc1]->type == grad->type); + + tensor_grad_view = ggml_view_4d(ctx, + grad, src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3], + nb1, nb2, nb3, offset); + } + + if (src0_needs_grads) { + struct ggml_tensor * tmp = ggml_neg(ctx, tensor_grad_view); + ggml_add_or_set(ctx, cgraph, isrc0, ggml_acc_impl(ctx, grad, tmp, nb1, nb2, nb3, offset, false)); + } + + if (src1_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc1, ggml_reshape(ctx, ggml_cont(ctx, tensor_grad_view), src1)); + } + } break; + case GGML_OP_CPY: { + // cpy overwrites value of src1 by src0 and returns view(src1) + // the overwriting is mathematically equivalent to: + // tensor = src0 * 1 + src1 * 0 + if (src0_needs_grads) { + // dsrc0 = dtensor * 1 + ggml_add_or_set(ctx, cgraph, isrc0, ggml_reshape(ctx, grad, src0)); + } + if (src1_needs_grads) { + // dsrc1 = dtensor * 0 -> noop + } + } break; + case GGML_OP_CONT: { + // same as cpy + if (src0_needs_grads) { + GGML_ASSERT(!cgraph->grads[isrc0] || ggml_is_contiguous(cgraph->grads[isrc0])); + GGML_ASSERT(ggml_is_contiguous(grad)); + GGML_ASSERT(ggml_nelements(tensor) == ggml_nelements(src0)); + ggml_add_or_set(ctx, cgraph, isrc0, + ggml_are_same_shape(tensor, src0) ? grad : ggml_reshape(ctx, grad, src0)); + } + } break; + case GGML_OP_RESHAPE: { + if (src0_needs_grads) { + struct ggml_tensor * grad_cont = ggml_is_contiguous(grad) ? grad : ggml_cont(ctx, grad); + ggml_add_or_set(ctx, cgraph, isrc0, ggml_reshape(ctx, grad_cont, src0)); + } + } break; + case GGML_OP_VIEW: { + if (src0_needs_grads) { + size_t offset; + + memcpy(&offset, tensor->op_params, sizeof(offset)); + + size_t nb1 = tensor->nb[1]; + size_t nb2 = tensor->nb[2]; + size_t nb3 = tensor->nb[3]; + + if (cgraph->grads[isrc0] && src0->type != cgraph->grads[isrc0]->type) { + // gradient is typically F32, but src0 could be other type + size_t ng = ggml_element_size(cgraph->grads[isrc0]); + size_t n0 = ggml_element_size(src0); + GGML_ASSERT(offset % n0 == 0); + GGML_ASSERT(nb1 % n0 == 0); + GGML_ASSERT(nb2 % n0 == 0); + GGML_ASSERT(nb3 % n0 == 0); + offset = (offset / n0) * ng; + nb1 = (nb1 / n0) * ng; + nb2 = (nb2 / n0) * ng; + nb3 = (nb3 / n0) * ng; + } + + ggml_acc_or_set(ctx, cgraph, isrc0, grad, nb1, nb2, nb3, offset); + } + } break; + case GGML_OP_PERMUTE: { + if (src0_needs_grads) { + const int32_t * axes = (const int32_t *) tensor->op_params; + const int axis0 = axes[0] & 0x3; + const int axis1 = axes[1] & 0x3; + const int axis2 = axes[2] & 0x3; + const int axis3 = axes[3] & 0x3; + int axb[4] = {0,0,0,0}; // axes backward + axb[axis0] = 0; + axb[axis1] = 1; + axb[axis2] = 2; + axb[axis3] = 3; + ggml_add_or_set(ctx, cgraph, isrc0, ggml_permute(ctx, grad, axb[0], axb[1], axb[2], axb[3])); + } + } break; + case GGML_OP_TRANSPOSE: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_transpose(ctx, grad)); + } + } break; + case GGML_OP_GET_ROWS: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_get_rows_back(ctx, grad, src1, src0)); + } + if (src1_needs_grads) { + // noop + } + } break; + case GGML_OP_DIAG_MASK_INF: { + if (src0_needs_grads) { + /* ggml_diag_mask_inf_impl() shouldn't be here */ + /* ref: https://github.com/ggml-org/llama.cpp/pull/4203#discussion_r1412377992 */ + const int n_past = ((const int32_t *) tensor->op_params)[0]; + ggml_add_or_set(ctx, cgraph, isrc0, ggml_diag_mask_zero_impl(ctx, grad, n_past, false)); + } + } break; + case GGML_OP_DIAG_MASK_ZERO: { + if (src0_needs_grads) { + const int n_past = ((const int32_t *) tensor->op_params)[0]; + ggml_add_or_set(ctx, cgraph, isrc0, ggml_diag_mask_zero_impl(ctx, grad, n_past, false)); + } + } break; + case GGML_OP_SOFT_MAX: { + if (src0_needs_grads) { + float scale = 1.0f; + float max_bias = 0.0f; + + memcpy(&scale, (const float *) tensor->op_params + 0, sizeof(float)); + memcpy(&max_bias, (const float *) tensor->op_params + 1, sizeof(float)); + + ggml_add_or_set(ctx, cgraph, isrc0, ggml_soft_max_ext_back(ctx, grad, tensor, scale, max_bias)); + } + GGML_ASSERT((!src1 || !src1_needs_grads) && "backward pass for softmax mask not implemented"); + } break; + case GGML_OP_ROPE: { + if (src0_needs_grads) { + //const int n_past = ((int32_t *) tensor->op_params)[0]; + const int n_dims = ((const int32_t *) tensor->op_params)[1]; + const int mode = ((const int32_t *) tensor->op_params)[2]; + //const int n_ctx = ((int32_t *) tensor->op_params)[3]; + const int n_ctx_orig = ((const int32_t *) tensor->op_params)[4]; + float freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow; + int sections[4] = {0, 0, 0, 0}; + + memcpy(&freq_base, (const float *) tensor->op_params + 5, sizeof(float)); + memcpy(&freq_scale, (const float *) tensor->op_params + 6, sizeof(float)); + memcpy(&ext_factor, (const float *) tensor->op_params + 7, sizeof(float)); + memcpy(&attn_factor, (const float *) tensor->op_params + 8, sizeof(float)); + memcpy(&beta_fast, (const float *) tensor->op_params + 9, sizeof(float)); + memcpy(&beta_slow, (const float *) tensor->op_params + 10, sizeof(float)); + memcpy(§ions, tensor->op_params + 11, sizeof(sections)); + + struct ggml_tensor * rope_back = grad->ne[2] == src1->ne[0] ? + ggml_rope_ext_back(ctx, grad, src1, src2, n_dims, + mode, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow) : + ggml_rope_multi_back(ctx, grad, src1, src2, n_dims, sections, + mode, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); + ggml_add_or_set(ctx, cgraph, isrc0, rope_back); + } + GGML_ASSERT((!src2 || !src2_needs_grads) && "gradients for freq factors not implemented"); + } break; + case GGML_OP_IM2COL: { + if (src1_needs_grads) { + const int32_t s0 = ggml_get_op_params_i32(tensor, 0); + const int32_t s1 = ggml_get_op_params_i32(tensor, 1); + const int32_t p0 = ggml_get_op_params_i32(tensor, 2); + const int32_t p1 = ggml_get_op_params_i32(tensor, 3); + const int32_t d0 = ggml_get_op_params_i32(tensor, 4); + const int32_t d1 = ggml_get_op_params_i32(tensor, 5); + const bool is_2D = ggml_get_op_params_i32(tensor, 6) == 1; + + ggml_add_or_set(ctx, cgraph, isrc1, ggml_im2col_back(ctx, grad, src0, src1->ne, s0, s1, p0, p1, d0, d1, is_2D)); + } + } break; + case GGML_OP_POOL_2D: { + if (src0_needs_grads) { + const enum ggml_op_pool op = ggml_get_op_params_i32(tensor, 0); + const int32_t k0 = ggml_get_op_params_i32(tensor, 1); + const int32_t k1 = ggml_get_op_params_i32(tensor, 2); + const int32_t s0 = ggml_get_op_params_i32(tensor, 3); + const int32_t s1 = ggml_get_op_params_i32(tensor, 4); + const int32_t p0 = ggml_get_op_params_i32(tensor, 5); + const int32_t p1 = ggml_get_op_params_i32(tensor, 6); + + ggml_add_or_set(ctx, cgraph, isrc0, ggml_pool_2d_back(ctx, grad, src0, op, k0, k1, s0, s1, p0, p1)); + } + } break; + case GGML_OP_WIN_PART: + case GGML_OP_WIN_UNPART: + case GGML_OP_UNARY: { + switch (ggml_get_unary_op(tensor)) { + case GGML_UNARY_OP_ABS: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_mul(ctx, ggml_sgn(ctx, src0), grad)); + } + } break; + case GGML_UNARY_OP_SGN: { + // noop + } break; + case GGML_UNARY_OP_NEG: { + if (src0_needs_grads) { + ggml_sub_or_set(ctx, cgraph, isrc0, grad); + } + } break; + case GGML_UNARY_OP_STEP: { + // noop + } break; + case GGML_UNARY_OP_RELU: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_mul(ctx, ggml_step(ctx, src0), grad)); + } + } break; + case GGML_UNARY_OP_SILU: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_silu_back(ctx, grad, src0)); + } + } break; + case GGML_UNARY_OP_EXP: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_mul(ctx, tensor, grad)); + } + } break; + case GGML_UNARY_OP_EXPM1: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_mul(ctx, grad, ggml_exp(ctx, src0))); + } + } break; + case GGML_UNARY_OP_SOFTPLUS: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_mul(ctx, grad, ggml_sigmoid(ctx, src0))); + } + } break; + default: { + fprintf(stderr, "%s: unsupported unary op for backward pass: %s\n", + __func__, ggml_unary_op_name(ggml_get_unary_op(tensor))); + GGML_ABORT("fatal error"); + } //break; + } + } break; + case GGML_OP_CROSS_ENTROPY_LOSS: { + if (src0_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc0, ggml_cross_entropy_loss_back(ctx, grad, src0, src1)); + } + GGML_ASSERT(!src1_needs_grads && "backward pass for labels not implemented"); + } break; + case GGML_OP_GLU: { + switch (ggml_get_glu_op(tensor)) { + case GGML_GLU_OP_SWIGLU: { + if (src0_needs_grads) { + GGML_ASSERT(src1 && "backward pass only implemented for split swiglu"); + ggml_add_or_set(ctx, cgraph, isrc0, ggml_silu_back(ctx, ggml_mul(ctx, grad, src1), src0)); + } + if (src1_needs_grads) { + ggml_add_or_set(ctx, cgraph, isrc1, ggml_mul(ctx, ggml_silu(ctx, src0), grad)); + } + } break; + default: { + GGML_ABORT("unsupported glu op for backward pass: %s", ggml_glu_op_name(ggml_get_glu_op(tensor))); + } //break; + } + } break; + case GGML_OP_NONE: { + // noop + } break; + case GGML_OP_COUNT: + default: { + GGML_ABORT("%s: unsupported ggml op for backward pass: %s\n", __func__, ggml_op_name(tensor->op)); + } //break; + } + + GGML_ASSERT(!src0_needs_grads || ggml_are_same_shape(src0, cgraph->grads[isrc0])); + GGML_ASSERT(!src1_needs_grads || ggml_are_same_shape(src1, cgraph->grads[isrc1])); + GGML_ASSERT(!src2_needs_grads || ggml_are_same_shape(src2, cgraph->grads[isrc2])); +} + +static size_t ggml_visit_parents_graph(struct ggml_cgraph * cgraph, struct ggml_tensor * node, bool compute) { + if (node->op != GGML_OP_NONE && compute) { + node->flags |= GGML_TENSOR_FLAG_COMPUTE; + } + + const size_t node_hash_pos = ggml_hash_find(&cgraph->visited_hash_set, node); + GGML_ASSERT(node_hash_pos != GGML_HASHSET_FULL); + + if (ggml_bitset_get(cgraph->visited_hash_set.used, node_hash_pos)) { + // already visited + + if (compute) { + // update the compute flag regardless + for (int i = 0; i < GGML_MAX_SRC; ++i) { + struct ggml_tensor * src = node->src[i]; + if (src && ((src->flags & GGML_TENSOR_FLAG_COMPUTE) == 0)) { + ggml_visit_parents_graph(cgraph, src, true); + } + } + } + + return node_hash_pos; + } + + // This is the first time we see this node in the current graph. + cgraph->visited_hash_set.keys[node_hash_pos] = node; + ggml_bitset_set(cgraph->visited_hash_set.used, node_hash_pos); + cgraph->use_counts[node_hash_pos] = 0; + + for (int i = 0; i < GGML_MAX_SRC; ++i) { + const int k = + (cgraph->order == GGML_CGRAPH_EVAL_ORDER_LEFT_TO_RIGHT) ? i : + (cgraph->order == GGML_CGRAPH_EVAL_ORDER_RIGHT_TO_LEFT) ? (GGML_MAX_SRC-1-i) : + /* unknown order, just fall back to using i */ i; + + struct ggml_tensor * src = node->src[k]; + if (src) { + const size_t src_hash_pos = ggml_visit_parents_graph(cgraph, src, compute); + + // Update the use count for this operand. + cgraph->use_counts[src_hash_pos]++; + } + } + + if (node->op == GGML_OP_NONE && !(node->flags & GGML_TENSOR_FLAG_PARAM)) { + // reached a leaf node, not part of the gradient graph (e.g. a constant) + GGML_ASSERT(cgraph->n_leafs < cgraph->size); + + if (strlen(node->name) == 0) { + ggml_format_name(node, "leaf_%d", cgraph->n_leafs); + } + + cgraph->leafs[cgraph->n_leafs] = node; + cgraph->n_leafs++; + } else { + GGML_ASSERT(cgraph->n_nodes < cgraph->size); + + if (strlen(node->name) == 0) { + ggml_format_name(node, "node_%d", cgraph->n_nodes); + } + + cgraph->nodes[cgraph->n_nodes] = node; + cgraph->n_nodes++; + } + + return node_hash_pos; +} + +static void ggml_build_forward_impl(struct ggml_cgraph * cgraph, struct ggml_tensor * tensor, bool expand, bool compute) { + if (!expand) { + // TODO: this branch isn't accessible anymore, maybe move this to ggml_build_forward_expand + ggml_graph_clear(cgraph); + } + + const int n_old = cgraph->n_nodes; + + ggml_visit_parents_graph(cgraph, tensor, compute); + + const int n_new = cgraph->n_nodes - n_old; + GGML_PRINT_DEBUG("%s: visited %d new nodes\n", __func__, n_new); + + if (n_new > 0) { + // the last added node should always be starting point + GGML_ASSERT(cgraph->nodes[cgraph->n_nodes - 1] == tensor); + } +} + +struct ggml_tensor * ggml_build_forward_select( + struct ggml_cgraph * cgraph, + struct ggml_tensor ** tensors, + int n_tensors, + int idx) { + GGML_ASSERT(idx >= 0 && idx < n_tensors); + + for (int i = 0; i < n_tensors; i++) { + ggml_build_forward_impl(cgraph, tensors[i], true, i == idx ? true : false); + } + + return tensors[idx]; +} + +void ggml_build_forward_expand(struct ggml_cgraph * cgraph, struct ggml_tensor * tensor) { + ggml_build_forward_impl(cgraph, tensor, true, true); +} + +void ggml_build_backward_expand( + struct ggml_context * ctx, + struct ggml_cgraph * cgraph, + struct ggml_tensor ** grad_accs) { + GGML_ASSERT(cgraph->n_nodes > 0); + GGML_ASSERT(cgraph->grads); + GGML_ASSERT(cgraph->grad_accs); + + const int n_nodes_f = cgraph->n_nodes; + + memset(cgraph->grads, 0, cgraph->visited_hash_set.size*sizeof(struct ggml_tensor *)); + memset(cgraph->grad_accs, 0, cgraph->visited_hash_set.size*sizeof(struct ggml_tensor *)); + bool * grads_needed = calloc(cgraph->visited_hash_set.size, sizeof(bool)); + + { + bool any_params = false; + bool any_loss = false; + for (int i = 0; i < n_nodes_f; ++i) { + struct ggml_tensor * node = cgraph->nodes[i]; + any_params = any_params || (node->flags & GGML_TENSOR_FLAG_PARAM); + any_loss = any_loss || (node->flags & GGML_TENSOR_FLAG_LOSS); + } + GGML_ASSERT(any_params && "no trainable parameters found, did you forget to call ggml_set_param?"); + GGML_ASSERT(any_loss && "no training loss found, did you forget to call ggml_set_loss?"); + } + + for (int i = 0; i < n_nodes_f; ++i) { + struct ggml_tensor * node = cgraph->nodes[i]; + + if (node->type == GGML_TYPE_I32) { + continue; + } + + bool node_needs_grad = (node->flags & GGML_TENSOR_FLAG_PARAM) || (node->flags & GGML_TENSOR_FLAG_LOSS); + bool ignore_src[GGML_MAX_SRC] = {false}; + switch (node->op) { + // gradients in node->src[0] for one reason or another have no effect on output gradients + case GGML_OP_IM2COL: // only used for its shape + case GGML_OP_IM2COL_BACK: // same as IM2COL + ignore_src[0] = true; + break; + case GGML_OP_UNARY: { + const enum ggml_unary_op uop = ggml_get_unary_op(node); + // SGN and STEP unary ops are piecewise constant + if (uop == GGML_UNARY_OP_SGN || uop == GGML_UNARY_OP_STEP) { + ignore_src[0] = true; + } + } break; + + // gradients in node->src[1] for one reason or another have no effect on output gradients + case GGML_OP_CPY: // gradients in CPY target are irrelevant + case GGML_OP_GET_ROWS: // row indices not differentiable + case GGML_OP_GET_ROWS_BACK: // same as for GET_ROWS + case GGML_OP_ROPE: // positions not differentiable + ignore_src[1] = true; + break; + + default: + break; + } + for (int j = 0; j < GGML_MAX_SRC; ++j) { + if (!node->src[j] || ignore_src[j] || !grads_needed[ggml_hash_find(&cgraph->visited_hash_set, node->src[j])]) { + continue; + } + GGML_ASSERT(node->src[j]->type == GGML_TYPE_F32 || node->src[j]->type == GGML_TYPE_F16); + node_needs_grad = true; + break; + } + if (!node_needs_grad) { + continue; + } + + // inplace operations are currently not supported + GGML_ASSERT(!node->view_src || node->op == GGML_OP_CPY || node->op == GGML_OP_VIEW || + node->op == GGML_OP_RESHAPE || node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE); + + const size_t ihash = ggml_hash_find(&cgraph->visited_hash_set, node); + GGML_ASSERT(ihash != GGML_HASHSET_FULL); + GGML_ASSERT(ggml_bitset_get(cgraph->visited_hash_set.used, ihash)); + if (grad_accs && grad_accs[i]) { + cgraph->grad_accs[ihash] = grad_accs[i]; + cgraph->grads[ihash] = cgraph->grad_accs[ihash]; + } else if (node->flags & GGML_TENSOR_FLAG_LOSS) { + // loss tensors always need a gradient accumulator + cgraph->grad_accs[ihash] = ggml_new_tensor(ctx, GGML_TYPE_F32, GGML_MAX_DIMS, node->ne); + cgraph->grads[ihash] = cgraph->grad_accs[ihash]; + } + grads_needed[ihash] = true; + } + + for (int i = n_nodes_f - 1; i >= 0; --i) { + // inplace operations to add gradients are not created by ggml_compute_backward except for gradient accumulation + // use allocator to automatically make inplace operations + ggml_compute_backward(ctx, cgraph, i, grads_needed); + } + + free(grads_needed); +} + +static void * incr_ptr_aligned(void ** p, size_t size, size_t align) { + void * ptr = *p; + ptr = (void *) GGML_PAD((uintptr_t) ptr, align); + *p = (void *) ((char *) ptr + size); + return ptr; +} + +static size_t ggml_graph_nbytes(size_t size, bool grads) { + size_t hash_size = ggml_hash_size(size * 2); + void * p = 0; + incr_ptr_aligned(&p, sizeof(struct ggml_cgraph), 1); + incr_ptr_aligned(&p, size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); // nodes + incr_ptr_aligned(&p, size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); // leafs + incr_ptr_aligned(&p, hash_size * sizeof(int32_t), sizeof(int32_t)); // use_counts + incr_ptr_aligned(&p, hash_size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); // hash keys + if (grads) { + incr_ptr_aligned(&p, hash_size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); // grads + incr_ptr_aligned(&p, hash_size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); // grad_accs + } + incr_ptr_aligned(&p, ggml_bitset_size(hash_size) * sizeof(ggml_bitset_t), sizeof(ggml_bitset_t)); + + size_t nbytes = (size_t) p; + return nbytes; +} + +size_t ggml_graph_overhead_custom(size_t size, bool grads) { + return GGML_OBJECT_SIZE + GGML_PAD(ggml_graph_nbytes(size, grads), GGML_MEM_ALIGN); +} + +size_t ggml_graph_overhead(void) { + return ggml_graph_overhead_custom(GGML_DEFAULT_GRAPH_SIZE, false); +} + +struct ggml_cgraph * ggml_new_graph_custom(struct ggml_context * ctx, size_t size, bool grads) { + const size_t obj_size = ggml_graph_nbytes(size, grads); + struct ggml_object * obj = ggml_new_object(ctx, GGML_OBJECT_TYPE_GRAPH, obj_size); + struct ggml_cgraph * cgraph = (struct ggml_cgraph *) ((char *) ctx->mem_buffer + obj->offs); + + // the size of the hash table is doubled since it needs to hold both nodes and leafs + size_t hash_size = ggml_hash_size(size * 2); + + void * p = cgraph + 1; + + struct ggml_tensor ** nodes_ptr = incr_ptr_aligned(&p, size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); + struct ggml_tensor ** leafs_ptr = incr_ptr_aligned(&p, size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); + int32_t * use_counts_ptr = incr_ptr_aligned(&p, hash_size * sizeof(int32_t), sizeof(int32_t)); + struct ggml_tensor ** hash_keys_ptr = incr_ptr_aligned(&p, hash_size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); + struct ggml_tensor ** grads_ptr = grads ? incr_ptr_aligned(&p, hash_size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)) : NULL; + struct ggml_tensor ** grad_accs_ptr = grads ? incr_ptr_aligned(&p, hash_size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)) : NULL; + + ggml_bitset_t * hash_used = incr_ptr_aligned(&p, ggml_bitset_size(hash_size) * sizeof(ggml_bitset_t), sizeof(ggml_bitset_t)); + + // check that we allocated the correct amount of memory + assert(obj_size == (size_t)((char *)p - (char *)cgraph)); + + *cgraph = (struct ggml_cgraph) { + /*.size =*/ size, + /*.n_nodes =*/ 0, + /*.n_leafs =*/ 0, + /*.nodes =*/ nodes_ptr, + /*.grads =*/ grads_ptr, + /*.grad_accs =*/ grad_accs_ptr, + /*.leafs =*/ leafs_ptr, + /*.use_counts =*/ use_counts_ptr, + /*.hash_table =*/ { hash_size, hash_used, hash_keys_ptr }, + /*.order =*/ GGML_CGRAPH_EVAL_ORDER_LEFT_TO_RIGHT, + }; + + ggml_hash_set_reset(&cgraph->visited_hash_set); + if (grads) { + memset(cgraph->grads, 0, hash_size*sizeof(struct ggml_tensor *)); + memset(cgraph->grad_accs, 0, hash_size*sizeof(struct ggml_tensor *)); + } + + return cgraph; +} + +struct ggml_cgraph * ggml_new_graph(struct ggml_context * ctx) { + return ggml_new_graph_custom(ctx, GGML_DEFAULT_GRAPH_SIZE, false); +} + +struct ggml_cgraph ggml_graph_view(struct ggml_cgraph * cgraph0, int i0, int i1) { + struct ggml_cgraph cgraph = { + /*.size =*/ 0, + /*.n_nodes =*/ i1 - i0, + /*.n_leafs =*/ 0, + /*.nodes =*/ cgraph0->nodes + i0, + /*.grads =*/ NULL, // gradients would need visited_hash_set + /*.grad_accs =*/ NULL, + /*.leafs =*/ NULL, + /*.use_counts =*/ cgraph0->use_counts, + /*.visited_hash_set =*/ cgraph0->visited_hash_set, + /*.order =*/ cgraph0->order, + }; + + return cgraph; +} + +void ggml_graph_cpy(struct ggml_cgraph * src, struct ggml_cgraph * dst) { + GGML_ASSERT(dst->size >= src->n_leafs); + GGML_ASSERT(dst->size >= src->n_nodes); + GGML_ASSERT(dst->visited_hash_set.size >= src->visited_hash_set.size); + + dst->n_leafs = src->n_leafs; + dst->n_nodes = src->n_nodes; + dst->order = src->order; + + for (int i = 0; i < src->n_leafs; ++i) { + dst->leafs[i] = src->leafs[i]; + } + + for (int i = 0; i < src->n_nodes; ++i) { + dst->nodes[i] = src->nodes[i]; + } + + for (size_t i = 0; i < src->visited_hash_set.size; ++i) { + // copy all hashset keys (tensors) that are in use + if (ggml_bitset_get(src->visited_hash_set.used, i)) { + size_t new_hash_pos = ggml_hash_insert(&dst->visited_hash_set, src->visited_hash_set.keys[i]); + dst->use_counts[new_hash_pos] = src->use_counts[i]; + } + } + + if (dst->grads) { + memset(dst->grads, 0, dst->visited_hash_set.size*sizeof(struct ggml_tensor *)); + memset(dst->grad_accs, 0, dst->visited_hash_set.size*sizeof(struct ggml_tensor *)); + } + if (src->grads) { + GGML_ASSERT(dst->grads != NULL); + GGML_ASSERT(dst->grad_accs != NULL); + for (int i = 0; i < src->n_nodes; ++i) { + const size_t igrad_src = ggml_hash_find(&src->visited_hash_set, src->nodes[i]); + const size_t igrad_dst = ggml_hash_find(&dst->visited_hash_set, dst->nodes[i]); + + GGML_ASSERT(igrad_src != GGML_HASHSET_FULL); + GGML_ASSERT(ggml_bitset_get(src->visited_hash_set.used, igrad_src)); + GGML_ASSERT(igrad_dst != GGML_HASHSET_FULL); + GGML_ASSERT(ggml_bitset_get(dst->visited_hash_set.used, igrad_dst)); + + dst->grads[igrad_dst] = src->grads[igrad_src]; + dst->grad_accs[igrad_dst] = src->grad_accs[igrad_src]; + } + } +} + +struct ggml_cgraph * ggml_graph_dup(struct ggml_context * ctx, struct ggml_cgraph * cgraph, bool force_grads) { + struct ggml_cgraph * result = ggml_new_graph_custom(ctx, cgraph->size, cgraph->grads || force_grads); + ggml_graph_cpy(cgraph, result); + return result; +} + +struct ggml_tensor * ggml_set_zero(struct ggml_tensor * tensor) { + if (ggml_is_empty(tensor)) { + return tensor; + } + if (tensor->buffer) { + ggml_backend_tensor_memset(tensor, 0, 0, ggml_nbytes(tensor)); + } else { + GGML_ASSERT(tensor->data); + memset(tensor->data, 0, ggml_nbytes(tensor)); + } + return tensor; +} + +void ggml_graph_reset(struct ggml_cgraph * cgraph) { + if (!cgraph) { + return; + } + GGML_ASSERT(cgraph->grads != NULL); + + for (int i = 0; i < cgraph->n_nodes; i++) { + struct ggml_tensor * node = cgraph->nodes[i]; + struct ggml_tensor * grad_acc = ggml_graph_get_grad_acc(cgraph, node); + + if (node->op == GGML_OP_OPT_STEP_ADAMW) { + // clear momenta + ggml_set_zero(node->src[2]); + ggml_set_zero(node->src[3]); + } + + // initial gradients of loss should be 1, 0 otherwise + if (grad_acc) { + if (node->flags & GGML_TENSOR_FLAG_LOSS) { + GGML_ASSERT(grad_acc->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_scalar(grad_acc)); + + const float onef = 1.0f; + if (grad_acc->buffer) { + ggml_backend_tensor_set(grad_acc, &onef, 0, sizeof(float)); + } else { + GGML_ASSERT(grad_acc->data); + *((float *) grad_acc->data) = onef; + } + } else { + ggml_set_zero(grad_acc); + } + } + } +} + +void ggml_graph_clear(struct ggml_cgraph * cgraph) { + cgraph->n_leafs = 0; + cgraph->n_nodes = 0; + ggml_hash_set_reset(&cgraph->visited_hash_set); +} + +int ggml_graph_size(struct ggml_cgraph * cgraph) { + return cgraph->size; +} + +struct ggml_tensor * ggml_graph_node(struct ggml_cgraph * cgraph, int i) { + if (i < 0) { + GGML_ASSERT(cgraph->n_nodes + i >= 0); + return cgraph->nodes[cgraph->n_nodes + i]; + } + + GGML_ASSERT(i < cgraph->n_nodes); + return cgraph->nodes[i]; +} + +struct ggml_tensor ** ggml_graph_nodes(struct ggml_cgraph * cgraph) { + return cgraph->nodes; +} + +int ggml_graph_n_nodes(struct ggml_cgraph * cgraph) { + return cgraph->n_nodes; +} + +void ggml_graph_add_node(struct ggml_cgraph * cgraph, struct ggml_tensor * tensor) { + GGML_ASSERT(cgraph->size > cgraph->n_nodes); + cgraph->nodes[cgraph->n_nodes] = tensor; + cgraph->n_nodes++; +} + +struct ggml_tensor * ggml_graph_get_tensor(const struct ggml_cgraph * cgraph, const char * name) { + for (int i = 0; i < cgraph->n_leafs; i++) { + struct ggml_tensor * leaf = cgraph->leafs[i]; + + if (strcmp(leaf->name, name) == 0) { + return leaf; + } + } + + for (int i = 0; i < cgraph->n_nodes; i++) { + struct ggml_tensor * node = cgraph->nodes[i]; + + if (strcmp(node->name, name) == 0) { + return node; + } + } + + return NULL; +} + +struct ggml_tensor * ggml_graph_get_grad(const struct ggml_cgraph * cgraph, const struct ggml_tensor * node) { + const size_t igrad = ggml_hash_find(&cgraph->visited_hash_set, node); + return igrad != GGML_HASHSET_FULL && ggml_bitset_get(cgraph->visited_hash_set.used, igrad) && cgraph->grads ? cgraph->grads[igrad] : NULL; +} + +struct ggml_tensor * ggml_graph_get_grad_acc(const struct ggml_cgraph * cgraph, const struct ggml_tensor * node) { + const size_t igrad = ggml_hash_find(&cgraph->visited_hash_set, node); + return igrad != GGML_HASHSET_FULL && ggml_bitset_get(cgraph->visited_hash_set.used, igrad) && cgraph->grad_accs ? cgraph->grad_accs[igrad] : NULL; +} + +void ggml_graph_print(const struct ggml_cgraph * cgraph) { + GGML_LOG_INFO("=== GRAPH ===\n"); + + GGML_LOG_INFO("n_nodes = %d\n", cgraph->n_nodes); + for (int i = 0; i < cgraph->n_nodes; i++) { + struct ggml_tensor * node = cgraph->nodes[i]; + + GGML_LOG_INFO(" - %3d: [ %5" PRId64 ", %5" PRId64 ", %5" PRId64 "] %16s %s\n", + i, + node->ne[0], node->ne[1], node->ne[2], + ggml_op_name(node->op), (node->flags & GGML_TENSOR_FLAG_PARAM) ? "x" : + ggml_graph_get_grad(cgraph, node) ? "g" : " "); + } + + GGML_LOG_INFO("n_leafs = %d\n", cgraph->n_leafs); + for (int i = 0; i < cgraph->n_leafs; i++) { + struct ggml_tensor * node = cgraph->leafs[i]; + + GGML_LOG_INFO(" - %3d: [ %5" PRId64 ", %5" PRId64 "] %8s %16s\n", + i, + node->ne[0], node->ne[1], + ggml_op_name(node->op), + ggml_get_name(node)); + } + + GGML_LOG_INFO("========================================\n"); +} + +static int ggml_node_list_find_tensor(const struct ggml_cgraph * cgraph, + const int * idxs, + int count, + const struct ggml_tensor * tensor) { + GGML_ASSERT(cgraph && idxs); + for (int i = 0; i < count; ++i) { + const int node_idx = idxs[i]; + + if (node_idx >= cgraph->n_nodes) { + return -1; + } + if (cgraph->nodes[node_idx] == tensor) { + return i; + } + } + return -1; +} + +bool ggml_can_fuse_subgraph_ext(const struct ggml_cgraph * cgraph, + const int * node_idxs, + int count, + const enum ggml_op * ops, + const int * outputs, + int num_outputs) { + GGML_ASSERT(outputs && num_outputs > 0); + + for (int i = 0; i < count; ++i) { + if (node_idxs[i] >= cgraph->n_nodes) { + return false; + } + + const struct ggml_tensor * node = cgraph->nodes[node_idxs[i]]; + + if (node->op != ops[i]) { + return false; + } + + if ((node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { + return false; + } + + if (ggml_node_list_find_tensor(cgraph, outputs, num_outputs, node) != -1) { + continue; + } + + if (node->flags & GGML_TENSOR_FLAG_OUTPUT) { + return false; + } + + int subgraph_uses = 0; + for (int j = i + 1; j < count; ++j) { + const struct ggml_tensor * other_node = cgraph->nodes[node_idxs[j]]; + for (int src_idx = 0; src_idx < GGML_MAX_SRC; src_idx++) { + if (other_node->src[src_idx] == node) { + subgraph_uses++; + } + } + } + + if (subgraph_uses != ggml_node_get_use_count(cgraph, node_idxs[i])) { + return false; + } + + // if node is a view, check if the view_src and all it's parent view_srcs are within the subgraph + struct ggml_tensor * view_src = node->view_src; + while (view_src) { + if (ggml_node_list_find_tensor(cgraph, node_idxs, count, view_src) == -1) { + return false; + } + view_src = view_src->view_src; + } + } + + return true; +} + +// check if node is part of the graph +static bool ggml_graph_find(const struct ggml_cgraph * cgraph, const struct ggml_tensor * node) { + if (cgraph == NULL) { + return true; + } + + for (int i = 0; i < cgraph->n_nodes; i++) { + if (cgraph->nodes[i] == node) { + return true; + } + } + + return false; +} + +static struct ggml_tensor * ggml_graph_get_parent(const struct ggml_cgraph * cgraph, const struct ggml_tensor * node) { + for (int i = 0; i < cgraph->n_nodes; i++) { + struct ggml_tensor * parent = cgraph->nodes[i]; + struct ggml_tensor * grad = ggml_graph_get_grad(cgraph, parent); + + if (grad == node) { + return parent; + } + } + + return NULL; +} + +static void ggml_graph_dump_dot_node_edge(FILE * fp, const struct ggml_cgraph * gb, struct ggml_tensor * node, struct ggml_tensor * parent, const char * label) { + struct ggml_tensor * gparent = ggml_graph_get_parent(gb, node); + struct ggml_tensor * gparent0 = ggml_graph_get_parent(gb, parent); + fprintf(fp, " \"%p\" -> \"%p\" [ arrowhead = %s; style = %s; label = \"%s\"; ]\n", + gparent0 ? (void *) gparent0 : (void *) parent, + gparent ? (void *) gparent : (void *) node, + gparent ? "empty" : "vee", + gparent ? "dashed" : "solid", + label); +} + +static void ggml_graph_dump_dot_leaf_edge(FILE * fp, struct ggml_tensor * node, struct ggml_tensor * parent, const char * label) { + fprintf(fp, " \"%p\" -> \"%p\" [ label = \"%s\"; ]\n", + (void *) parent, + (void *) node, + label); +} + +void ggml_graph_dump_dot(const struct ggml_cgraph * gb, const struct ggml_cgraph * cgraph, const char * filename) { + char color[16]; + + FILE * fp = ggml_fopen(filename, "w"); + GGML_ASSERT(fp); + + fprintf(fp, "digraph G {\n"); + fprintf(fp, " newrank = true;\n"); + fprintf(fp, " rankdir = TB;\n"); + + for (int i = 0; i < gb->n_nodes; i++) { + struct ggml_tensor * node = gb->nodes[i]; + struct ggml_tensor * grad = ggml_graph_get_grad(gb, node); + + if (ggml_graph_get_parent(gb, node) != NULL) { + continue; + } + + if (node->flags & GGML_TENSOR_FLAG_PARAM) { + snprintf(color, sizeof(color), "yellow"); + } else if (grad) { + if (ggml_graph_find(cgraph, node)) { + snprintf(color, sizeof(color), "green"); + } else { + snprintf(color, sizeof(color), "lightblue"); + } + } else { + snprintf(color, sizeof(color), "white"); + } + + fprintf(fp, " \"%p\" [ " + "style = filled; fillcolor = %s; shape = record; " + "label=\"", + (void *) node, color); + + if (strlen(node->name) > 0) { + fprintf(fp, "%s (%s)|", node->name, ggml_type_name(node->type)); + } else { + fprintf(fp, "(%s)|", ggml_type_name(node->type)); + } + + if (ggml_is_matrix(node)) { + fprintf(fp, "%d [%" PRId64 ", %" PRId64 "] | %s", i, node->ne[0], node->ne[1], ggml_op_symbol(node->op)); + } else { + fprintf(fp, "%d [%" PRId64 ", %" PRId64 ", %" PRId64 "] | %s", i, node->ne[0], node->ne[1], node->ne[2], ggml_op_symbol(node->op)); + } + + if (grad) { + fprintf(fp, " | %s\"; ]\n", ggml_op_symbol(grad->op)); + } else { + fprintf(fp, "\"; ]\n"); + } + } + + for (int i = 0; i < gb->n_leafs; i++) { + struct ggml_tensor * node = gb->leafs[i]; + + snprintf(color, sizeof(color), "pink"); + + fprintf(fp, " \"%p\" [ " + "style = filled; fillcolor = %s; shape = record; " + "label=\"", + (void *) node, color); + + if (strlen(node->name) > 0) { + fprintf(fp, "%s (%s)|", node->name, ggml_type_name(node->type)); + } else { + fprintf(fp, "(%s)|", ggml_type_name(node->type)); + } + + fprintf(fp, "CONST %d [%" PRId64 ", %" PRId64 "]", i, node->ne[0], node->ne[1]); + if (ggml_nelements(node) < 5 && node->data != NULL) { + fprintf(fp, " | ("); + for (int j = 0; j < ggml_nelements(node); j++) { + // FIXME: use ggml-backend to obtain the tensor data + //if (node->type == GGML_TYPE_I8 || node->type == GGML_TYPE_I16 || node->type == GGML_TYPE_I32) { + // fprintf(fp, "%d", ggml_get_i32_1d(node, j)); + //} + //else if (node->type == GGML_TYPE_F32 || + // node->type == GGML_TYPE_F16 || + // node->type == GGML_TYPE_BF16) { + // fprintf(fp, "%.1e", (double)ggml_get_f32_1d(node, j)); + //} + //else + { + fprintf(fp, "#"); + } + if (j < ggml_nelements(node) - 1) { + fprintf(fp, ", "); + } + } + fprintf(fp, ")"); + } + fprintf(fp, "\"; ]\n"); + } + + for (int i = 0; i < gb->n_nodes; i++) { + struct ggml_tensor * node = gb->nodes[i]; + + for (int j = 0; j < GGML_MAX_SRC; j++) { + if (node->src[j]) { + char label[16]; + snprintf(label, sizeof(label), "src %d", j); + ggml_graph_dump_dot_node_edge(fp, gb, node, node->src[j], label); + } + } + } + + for (int i = 0; i < gb->n_leafs; i++) { + struct ggml_tensor * node = gb->leafs[i]; + + for (int j = 0; j < GGML_MAX_SRC; j++) { + if (node->src[j]) { + char label[16]; + snprintf(label, sizeof(label), "src %d", j); + ggml_graph_dump_dot_leaf_edge(fp, node, node->src[j], label); + } + } + } + + fprintf(fp, "}\n"); + + fclose(fp); + + GGML_LOG_INFO("%s: dot -Tpng %s -o %s.png && open %s.png\n", __func__, filename, filename, filename); +} + +//////////////////////////////////////////////////////////////////////////////// + +void ggml_set_input(struct ggml_tensor * tensor) { + tensor->flags |= GGML_TENSOR_FLAG_INPUT; +} + +void ggml_set_output(struct ggml_tensor * tensor) { + tensor->flags |= GGML_TENSOR_FLAG_OUTPUT; +} + +void ggml_set_param(struct ggml_tensor * tensor) { + GGML_ASSERT(tensor->op == GGML_OP_NONE); + tensor->flags |= GGML_TENSOR_FLAG_PARAM; +} + +void ggml_set_loss(struct ggml_tensor * tensor) { + GGML_ASSERT(ggml_is_scalar(tensor)); + GGML_ASSERT(tensor->type == GGML_TYPE_F32); + tensor->flags |= GGML_TENSOR_FLAG_LOSS; +} + +//////////////////////////////////////////////////////////////////////////////// + +void ggml_quantize_init(enum ggml_type type) { + ggml_critical_section_start(); + + switch (type) { + case GGML_TYPE_IQ2_XXS: + case GGML_TYPE_IQ2_XS: + case GGML_TYPE_IQ2_S: + case GGML_TYPE_IQ1_S: + case GGML_TYPE_IQ1_M: iq2xs_init_impl(type); break; + case GGML_TYPE_IQ3_XXS: iq3xs_init_impl(256); break; + case GGML_TYPE_IQ3_S: iq3xs_init_impl(512); break; + default: // nothing + break; + } + + ggml_critical_section_end(); +} + +void ggml_quantize_free(void) { + ggml_critical_section_start(); + + iq2xs_free_impl(GGML_TYPE_IQ2_XXS); + iq2xs_free_impl(GGML_TYPE_IQ2_XS); + iq2xs_free_impl(GGML_TYPE_IQ2_S); + iq2xs_free_impl(GGML_TYPE_IQ1_S); + iq2xs_free_impl(GGML_TYPE_IQ1_M); + iq3xs_free_impl(256); + iq3xs_free_impl(512); + + ggml_critical_section_end(); +} + +bool ggml_quantize_requires_imatrix(enum ggml_type type) { + return + type == GGML_TYPE_IQ2_XXS || + type == GGML_TYPE_IQ2_XS || + type == GGML_TYPE_IQ1_S;// || + //type == GGML_TYPE_IQ1_M; +} + +size_t ggml_quantize_chunk( + enum ggml_type type, + const float * src, + void * dst, + int64_t start, + int64_t nrows, + int64_t n_per_row, + const float * imatrix) { + const int64_t n = (int64_t) nrows * n_per_row; + + if (ggml_quantize_requires_imatrix(type)) { + GGML_ASSERT(imatrix != NULL); + } + + GGML_ASSERT(start % type_traits[type].blck_size == 0); + GGML_ASSERT(start % n_per_row == 0); + + ggml_quantize_init(type); // this is noop if already initialized + + const size_t start_row = start / n_per_row; + const size_t row_size = ggml_row_size(type, n_per_row); + + size_t result = 0; + + switch (type) { + case GGML_TYPE_Q1_0: result = quantize_q1_0(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_Q4_0: result = quantize_q4_0(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_Q4_1: result = quantize_q4_1(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_Q5_0: result = quantize_q5_0(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_Q5_1: result = quantize_q5_1(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_Q8_0: result = quantize_q8_0(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_MXFP4: result = quantize_mxfp4(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_NVFP4: result = quantize_nvfp4(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_Q2_K: result = quantize_q2_K(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_Q3_K: result = quantize_q3_K(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_Q4_K: result = quantize_q4_K(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_Q5_K: result = quantize_q5_K(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_Q6_K: result = quantize_q6_K(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_TQ1_0: result = quantize_tq1_0(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_TQ2_0: result = quantize_tq2_0(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_IQ2_XXS: result = quantize_iq2_xxs(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_IQ2_XS: result = quantize_iq2_xs (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_IQ3_XXS: result = quantize_iq3_xxs(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_IQ3_S: result = quantize_iq3_s (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_IQ2_S: result = quantize_iq2_s (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_IQ1_S: result = quantize_iq1_s (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_IQ1_M: result = quantize_iq1_m (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_IQ4_NL: result = quantize_iq4_nl (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_IQ4_XS: result = quantize_iq4_xs (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_F16: + { + size_t elemsize = sizeof(ggml_fp16_t); + ggml_fp32_to_fp16_row(src + start, (ggml_fp16_t *)dst + start, n); + result = n * elemsize; + } break; + case GGML_TYPE_BF16: + { + size_t elemsize = sizeof(ggml_bf16_t); + ggml_fp32_to_bf16_row_ref(src + start, (ggml_bf16_t *)dst + start, n); + result = n * elemsize; + } break; + case GGML_TYPE_F32: + { + size_t elemsize = sizeof(float); + result = n * elemsize; + memcpy((uint8_t *)dst + start * elemsize, src + start, result); + } break; + default: + assert(false); + } + + GGML_ASSERT(result == nrows * row_size); + + return result; +} + +//////////////////////////////////////////////////////////////////////////////// + +void ggml_log_get(ggml_log_callback * log_callback, void ** user_data) { + *log_callback = g_logger_state.log_callback; + *user_data = g_logger_state.log_callback_user_data; +} + +void ggml_log_set(ggml_log_callback log_callback, void * user_data) { + g_logger_state.log_callback = log_callback ? log_callback : ggml_log_callback_default; + g_logger_state.log_callback_user_data = user_data; +} + +void ggml_threadpool_params_init(struct ggml_threadpool_params * p, int n_threads) { + p->n_threads = n_threads; + p->prio = 0; // default priority (usually means normal or inherited) + p->poll = 50; // hybrid-polling enabled + p->strict_cpu = false; // no strict placement (all threads share same cpumask) + p->paused = false; // threads are ready to go + memset(p->cpumask, 0, GGML_MAX_N_THREADS); // all-zero means use the default affinity (usually inherited) +} + +struct ggml_threadpool_params ggml_threadpool_params_default(int n_threads) { + struct ggml_threadpool_params p; + ggml_threadpool_params_init(&p, n_threads); + return p; +} + +bool ggml_threadpool_params_match(const struct ggml_threadpool_params * p0, const struct ggml_threadpool_params * p1) { + if (p0->n_threads != p1->n_threads ) return false; + if (p0->prio != p1->prio ) return false; + if (p0->poll != p1->poll ) return false; + if (p0->strict_cpu != p1->strict_cpu ) return false; + return memcmp(p0->cpumask, p1->cpumask, GGML_MAX_N_THREADS) == 0; +} + +struct ggml_tensor * ggml_turbo_wht( + struct ggml_context * ctx, + struct ggml_tensor * a, + int direction) { + GGML_ASSERT(ggml_n_dims(a) >= 2); + struct ggml_tensor * result = ggml_new_tensor(ctx, a->type, GGML_MAX_DIMS, a->ne); + result->op = GGML_OP_TURBO_WHT; + result->src[0] = a; + ggml_set_op_params_i32(result, 0, direction); + return result; +} + +struct ggml_tensor * ggml_moe_fused( + struct ggml_context * ctx, + struct ggml_tensor * input, + struct ggml_tensor * gate_w, + struct ggml_tensor * up_w, + struct ggml_tensor * down_w, + struct ggml_tensor * expert_ids, + struct ggml_tensor * expert_weights, + struct ggml_tensor * sh_gate_w, + struct ggml_tensor * sh_up_w, + struct ggml_tensor * sh_down_w, + struct ggml_tensor * sh_gate_inp_w, + int64_t n_embd, + int64_t ff_dim, + int64_t n_expert_used) { + GGML_ASSERT(ggml_is_contiguous(input)); + GGML_ASSERT(input->type == GGML_TYPE_F32); + + const int64_t ne[4] = { n_embd, 1, 1, 1 }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); + + result->op = GGML_OP_MOE_FUSED; + result->src[0] = input; + result->src[1] = gate_w; + result->src[2] = up_w; + result->src[3] = down_w; + result->src[4] = expert_ids; + result->src[5] = expert_weights; + result->src[6] = sh_gate_w; + result->src[7] = sh_up_w; + result->src[8] = sh_down_w; + result->src[9] = sh_gate_inp_w; + + ggml_set_op_params_i32(result, 0, (int32_t) n_embd); + ggml_set_op_params_i32(result, 1, (int32_t) ff_dim); + ggml_set_op_params_i32(result, 2, (int32_t) n_expert_used); + + return result; +} + +struct ggml_tensor * ggml_laguna_moe_combine( + struct ggml_context * ctx, + struct ggml_tensor * experts, + struct ggml_tensor * expert_weights) { + GGML_ASSERT(experts->type == GGML_TYPE_F32); + GGML_ASSERT(expert_weights->type == GGML_TYPE_F32); + GGML_ASSERT(experts->ne[1] == expert_weights->ne[0]); + GGML_ASSERT(experts->ne[2] == expert_weights->ne[1]); + + const int64_t ne[4] = { experts->ne[0], experts->ne[2], 1, 1 }; + struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 2, ne); + + result->op = GGML_OP_MOE_FUSED; + result->src[0] = experts; + result->src[1] = expert_weights; + + ggml_set_op_params_i32(result, 0, -1); + ggml_set_op_params_i32(result, 1, (int32_t) experts->ne[0]); + ggml_set_op_params_i32(result, 2, (int32_t) experts->ne[1]); + ggml_set_op_params_i32(result, 3, (int32_t) experts->ne[2]); + + return result; +} From 53a91585f2406ff2329c8969bba4d85f40148a12 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Wed, 8 Jul 2026 11:15:40 +0200 Subject: [PATCH 6/7] chore: remove stray editor backup files --- ggml/include/ggml.h.ds4bak | 2903 --------- ggml/src/ggml-cpu/ggml-cpu.c.ds4bak | 3782 ----------- ggml/src/ggml-cuda/ggml-cuda.cu.ds4bak | 5559 ---------------- ggml/src/ggml.c.ds4bak | 8040 ------------------------ 4 files changed, 20284 deletions(-) delete mode 100644 ggml/include/ggml.h.ds4bak delete mode 100644 ggml/src/ggml-cpu/ggml-cpu.c.ds4bak delete mode 100644 ggml/src/ggml-cuda/ggml-cuda.cu.ds4bak delete mode 100644 ggml/src/ggml.c.ds4bak diff --git a/ggml/include/ggml.h.ds4bak b/ggml/include/ggml.h.ds4bak deleted file mode 100644 index 800854e7d462..000000000000 --- a/ggml/include/ggml.h.ds4bak +++ /dev/null @@ -1,2903 +0,0 @@ -#pragma once - -// -// GGML Tensor Library -// -// This documentation is still a work in progress. -// If you wish some specific topics to be covered, feel free to drop a comment: -// -// https://github.com/ggml-org/whisper.cpp/issues/40 -// -// ## Overview -// -// This library implements: -// -// - a set of tensor operations -// - automatic differentiation -// - basic optimization algorithms -// -// The aim of this library is to provide a minimalistic approach for various machine learning tasks. This includes, -// but is not limited to, the following: -// -// - linear regression -// - support vector machines -// - neural networks -// -// The library allows the user to define a certain function using the available tensor operations. This function -// definition is represented internally via a computation graph. Each tensor operation in the function definition -// corresponds to a node in the graph. Having the computation graph defined, the user can choose to compute the -// function's value and/or its gradient with respect to the input variables. Optionally, the function can be optimized -// using one of the available optimization algorithms. -// -// For example, here we define the function: f(x) = a*x^2 + b -// -// { -// struct ggml_init_params params = { -// .mem_size = 16*1024*1024, -// .mem_buffer = NULL, -// }; -// -// // memory allocation happens here -// struct ggml_context * ctx = ggml_init(params); -// -// struct ggml_tensor * x = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); -// -// ggml_set_param(ctx, x); // x is an input variable -// -// struct ggml_tensor * a = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); -// struct ggml_tensor * b = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); -// struct ggml_tensor * x2 = ggml_mul(ctx, x, x); -// struct ggml_tensor * f = ggml_add(ctx, ggml_mul(ctx, a, x2), b); -// -// ... -// } -// -// Notice that the function definition above does not involve any actual computation. The computation is performed only -// when the user explicitly requests it. For example, to compute the function's value at x = 2.0: -// -// { -// ... -// -// struct ggml_cgraph * gf = ggml_new_graph(ctx); -// ggml_build_forward_expand(gf, f); -// -// // set the input variable and parameter values -// ggml_set_f32(x, 2.0f); -// ggml_set_f32(a, 3.0f); -// ggml_set_f32(b, 4.0f); -// -// ggml_graph_compute_with_ctx(ctx, &gf, n_threads); -// -// printf("f = %f\n", ggml_get_f32_1d(f, 0)); -// -// ... -// } -// -// The actual computation is performed in the ggml_graph_compute() function. -// -// The ggml_new_tensor_...() functions create new tensors. They are allocated in the memory buffer provided to the -// ggml_init() function. You have to be careful not to exceed the memory buffer size. Therefore, you have to know -// in advance how much memory you need for your computation. Alternatively, you can allocate a large enough memory -// and after defining the computation graph, call the ggml_used_mem() function to find out how much memory was -// actually needed. -// -// The ggml_set_param() function marks a tensor as an input variable. This is used by the automatic -// differentiation and optimization algorithms. -// -// The described approach allows to define the function graph once and then compute its forward or backward graphs -// multiple times. All computations will use the same memory buffer allocated in the ggml_init() function. This way -// the user can avoid the memory allocation overhead at runtime. -// -// The library supports multi-dimensional tensors - up to 4 dimensions. The FP16 and FP32 data types are first class -// citizens, but in theory the library can be extended to support FP8 and integer data types. -// -// Each tensor operation produces a new tensor. Initially the library was envisioned to support only the use of unary -// and binary operations. Most of the available operations fall into one of these two categories. With time, it became -// clear that the library needs to support more complex operations. The way to support these operations is not clear -// yet, but a few examples are demonstrated in the following operations: -// -// - ggml_permute() -// - ggml_conv_1d_1s() -// - ggml_conv_1d_2s() -// -// For each tensor operator, the library implements a forward and backward computation function. The forward function -// computes the output tensor value given the input tensor values. The backward function computes the adjoint of the -// input tensors given the adjoint of the output tensor. For a detailed explanation of what this means, take a -// calculus class, or watch the following video: -// -// What is Automatic Differentiation? -// https://www.youtube.com/watch?v=wG_nF1awSSY -// -// -// ## Tensor data (struct ggml_tensor) -// -// The tensors are stored in memory via the ggml_tensor struct. The structure provides information about the size of -// the tensor, the data type, and the memory buffer where the tensor data is stored. Additionally, it contains -// pointers to the "source" tensors - i.e. the tensors that were used to compute the current tensor. For example: -// -// { -// struct ggml_tensor * c = ggml_add(ctx, a, b); -// -// assert(c->src[0] == a); -// assert(c->src[1] == b); -// } -// -// The multi-dimensional tensors are stored in row-major order. The ggml_tensor struct contains fields for the -// number of elements in each dimension ("ne") as well as the number of bytes ("nb", a.k.a. stride). This allows -// to store tensors that are not contiguous in memory, which is useful for operations such as transposition and -// permutation. All tensor operations have to take the stride into account and not assume that the tensor is -// contiguous in memory. -// -// The data of the tensor is accessed via the "data" pointer. For example: -// -// { -// const int nx = 2; -// const int ny = 3; -// -// struct ggml_tensor * a = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, nx, ny); -// -// for (int y = 0; y < ny; y++) { -// for (int x = 0; x < nx; x++) { -// *(float *) ((char *) a->data + y*a->nb[1] + x*a->nb[0]) = x + y; -// } -// } -// -// ... -// } -// -// Alternatively, there are helper functions, such as ggml_get_f32_1d() and ggml_set_f32_1d() that can be used. -// -// ## The matrix multiplication operator (ggml_mul_mat) -// -// TODO -// -// -// ## Multi-threading -// -// TODO -// -// -// ## Overview of ggml.c -// -// TODO -// -// -// ## SIMD optimizations -// -// TODO -// -// -// ## Debugging ggml -// -// TODO -// -// - -#ifdef GGML_SHARED -# if defined(_WIN32) && !defined(__MINGW32__) -# ifdef GGML_BUILD -# define GGML_API __declspec(dllexport) extern -# else -# define GGML_API __declspec(dllimport) extern -# endif -# else -# define GGML_API __attribute__ ((visibility ("default"))) extern -# endif -#else -# define GGML_API extern -#endif - -// TODO: support for clang -#ifdef __GNUC__ -# define GGML_DEPRECATED(func, hint) func __attribute__((deprecated(hint))) -#elif defined(_MSC_VER) -# define GGML_DEPRECATED(func, hint) __declspec(deprecated(hint)) func -#else -# define GGML_DEPRECATED(func, hint) func -#endif - -#ifndef __GNUC__ -# define GGML_ATTRIBUTE_FORMAT(...) -#elif defined(__MINGW32__) && !defined(__clang__) -# define GGML_ATTRIBUTE_FORMAT(...) __attribute__((format(gnu_printf, __VA_ARGS__))) -#else -# define GGML_ATTRIBUTE_FORMAT(...) __attribute__((format(printf, __VA_ARGS__))) -#endif - -#if defined(_WIN32) && !defined(_WIN32_WINNT) -# define _WIN32_WINNT 0x0A00 -#endif - -#include -#include -#include -#include - -#define GGML_FILE_MAGIC 0x67676d6c // "ggml" -#define GGML_FILE_VERSION 2 - -#define GGML_QNT_VERSION 2 // bump this on quantization format changes -#define GGML_QNT_VERSION_FACTOR 1000 // do not change this - -#define GGML_MAX_DIMS 4 -#define GGML_MAX_PARAMS 2048 -#define GGML_MAX_SRC 10 -#define GGML_MAX_N_THREADS 512 -#define GGML_MAX_OP_PARAMS 64 - -#ifndef GGML_MAX_NAME -# define GGML_MAX_NAME 64 -#endif - -#define GGML_DEFAULT_N_THREADS 4 -#define GGML_DEFAULT_GRAPH_SIZE 2048 - -#if UINTPTR_MAX == 0xFFFFFFFF - #define GGML_MEM_ALIGN 4 -#elif defined(__EMSCRIPTEN__) -// emscripten uses max_align_t == 8, so we need GGML_MEM_ALIGN == 8 for 64-bit wasm. -// (for 32-bit wasm, the first conditional is true and GGML_MEM_ALIGN stays 4.) -// ref: https://github.com/ggml-org/llama.cpp/pull/18628 - #define GGML_MEM_ALIGN 8 -#else - #define GGML_MEM_ALIGN 16 -#endif - -#define GGML_EXIT_SUCCESS 0 -#define GGML_EXIT_ABORTED 1 - -// TODO: convert to enum https://github.com/ggml-org/llama.cpp/pull/16187#discussion_r2388538726 -#define GGML_ROPE_TYPE_NORMAL 0 -#define GGML_ROPE_TYPE_NEOX 2 -#define GGML_ROPE_TYPE_MROPE 8 -#define GGML_ROPE_TYPE_VISION 24 -#define GGML_ROPE_TYPE_IMROPE 40 // binary: 101000 - -#define GGML_MROPE_SECTIONS 4 - -#define GGML_UNUSED(x) (void)(x) -#ifdef __CUDACC__ -template -__host__ __device__ constexpr inline void ggml_unused_vars_impl(Args&&...) noexcept {} -#define GGML_UNUSED_VARS(...) ggml_unused_vars_impl(__VA_ARGS__) -#else -#define GGML_UNUSED_VARS(...) do { (void)sizeof((__VA_ARGS__, 0)); } while(0) -#endif // __CUDACC__ - -#define GGML_PAD(x, n) (((x) + (n) - 1) & ~((n) - 1)) - -#ifndef NDEBUG -# define GGML_UNREACHABLE() do { fprintf(stderr, "statement should be unreachable\n"); abort(); } while(0) -#elif defined(__GNUC__) -# define GGML_UNREACHABLE() __builtin_unreachable() -#elif defined(_MSC_VER) -# define GGML_UNREACHABLE() __assume(0) -#else -# define GGML_UNREACHABLE() ((void) 0) -#endif - -#ifdef __cplusplus -# define GGML_NORETURN [[noreturn]] -#elif defined(_MSC_VER) -# define GGML_NORETURN __declspec(noreturn) -#else -# define GGML_NORETURN _Noreturn -#endif - -#define GGML_ABORT(...) ggml_abort(__FILE__, __LINE__, __VA_ARGS__) -#define GGML_ASSERT(x) if (!(x)) GGML_ABORT("GGML_ASSERT(%s) failed", #x) - -// used to copy the number of elements and stride in bytes of tensors into local variables. -// main purpose is to reduce code duplication and improve readability. -// -// example: -// -// GGML_TENSOR_LOCALS(int64_t, ne1, src1, ne); -// GGML_TENSOR_LOCALS(size_t, nb1, src1, nb); -// -#define GGML_TENSOR_LOCALS_1(type, prefix, pointer, array) \ - const type prefix##0 = (pointer) ? (pointer)->array[0] : 0; \ - GGML_UNUSED(prefix##0); -#define GGML_TENSOR_LOCALS_2(type, prefix, pointer, array) \ - GGML_TENSOR_LOCALS_1 (type, prefix, pointer, array) \ - const type prefix##1 = (pointer) ? (pointer)->array[1] : 0; \ - GGML_UNUSED(prefix##1); -#define GGML_TENSOR_LOCALS_3(type, prefix, pointer, array) \ - GGML_TENSOR_LOCALS_2 (type, prefix, pointer, array) \ - const type prefix##2 = (pointer) ? (pointer)->array[2] : 0; \ - GGML_UNUSED(prefix##2); -#define GGML_TENSOR_LOCALS(type, prefix, pointer, array) \ - GGML_TENSOR_LOCALS_3 (type, prefix, pointer, array) \ - const type prefix##3 = (pointer) ? (pointer)->array[3] : 0; \ - GGML_UNUSED(prefix##3); - -#define GGML_TENSOR_UNARY_OP_LOCALS \ - GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) \ - GGML_TENSOR_LOCALS(size_t, nb0, src0, nb) \ - GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) \ - GGML_TENSOR_LOCALS(size_t, nb, dst, nb) - -#define GGML_TENSOR_BINARY_OP_LOCALS \ - GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) \ - GGML_TENSOR_LOCALS(size_t, nb0, src0, nb) \ - GGML_TENSOR_LOCALS(int64_t, ne1, src1, ne) \ - GGML_TENSOR_LOCALS(size_t, nb1, src1, nb) \ - GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) \ - GGML_TENSOR_LOCALS(size_t, nb, dst, nb) - -#define GGML_TENSOR_TERNARY_OP_LOCALS \ - GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) \ - GGML_TENSOR_LOCALS(size_t, nb0, src0, nb) \ - GGML_TENSOR_LOCALS(int64_t, ne1, src1, ne) \ - GGML_TENSOR_LOCALS(size_t, nb1, src1, nb) \ - GGML_TENSOR_LOCALS(int64_t, ne2, src2, ne) \ - GGML_TENSOR_LOCALS(size_t, nb2, src2, nb) \ - GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) \ - GGML_TENSOR_LOCALS(size_t, nb, dst, nb) - -#define GGML_TENSOR_BINARY_OP_LOCALS01 \ - GGML_TENSOR_LOCALS(int64_t, ne0, src0, ne) \ - GGML_TENSOR_LOCALS(size_t, nb0, src0, nb) \ - GGML_TENSOR_LOCALS(int64_t, ne1, src1, ne) \ - GGML_TENSOR_LOCALS(size_t, nb1, src1, nb) - -#ifdef __cplusplus -extern "C" { -#endif - - // Function type used in fatal error callbacks - typedef void (*ggml_abort_callback_t)(const char * error_message); - - // Set the abort callback (passing null will restore original abort functionality: printing a message to stdout) - // Returns the old callback for chaining - GGML_API ggml_abort_callback_t ggml_set_abort_callback(ggml_abort_callback_t callback); - - GGML_NORETURN GGML_ATTRIBUTE_FORMAT(3, 4) - GGML_API void ggml_abort(const char * file, int line, const char * fmt, ...); - - enum ggml_status { - GGML_STATUS_ALLOC_FAILED = -2, - GGML_STATUS_FAILED = -1, - GGML_STATUS_SUCCESS = 0, - GGML_STATUS_ABORTED = 1, - }; - - // get ggml_status name string - GGML_API const char * ggml_status_to_string(enum ggml_status status); - - // ieee 754-2008 half-precision float16 - // todo: make this not an integral type - typedef uint16_t ggml_fp16_t; - GGML_API float ggml_fp16_to_fp32(ggml_fp16_t); - GGML_API ggml_fp16_t ggml_fp32_to_fp16(float); - GGML_API void ggml_fp16_to_fp32_row(const ggml_fp16_t *, float *, int64_t); - GGML_API void ggml_fp32_to_fp16_row(const float *, ggml_fp16_t *, int64_t); - - // google brain half-precision bfloat16 - typedef struct { uint16_t bits; } ggml_bf16_t; - GGML_API ggml_bf16_t ggml_fp32_to_bf16(float); - GGML_API float ggml_bf16_to_fp32(ggml_bf16_t); // consider just doing << 16 - GGML_API void ggml_bf16_to_fp32_row(const ggml_bf16_t *, float *, int64_t); - GGML_API void ggml_fp32_to_bf16_row_ref(const float *, ggml_bf16_t *, int64_t); - GGML_API void ggml_fp32_to_bf16_row(const float *, ggml_bf16_t *, int64_t); - - struct ggml_object; - struct ggml_context; - struct ggml_cgraph; - - // NOTE: always add types at the end of the enum to keep backward compatibility - enum ggml_type { - GGML_TYPE_F32 = 0, - GGML_TYPE_F16 = 1, - GGML_TYPE_Q4_0 = 2, - GGML_TYPE_Q4_1 = 3, - // GGML_TYPE_Q4_2 = 4, support has been removed - // GGML_TYPE_Q4_3 = 5, support has been removed - GGML_TYPE_Q5_0 = 6, - GGML_TYPE_Q5_1 = 7, - GGML_TYPE_Q8_0 = 8, - GGML_TYPE_Q8_1 = 9, - GGML_TYPE_Q2_K = 10, - GGML_TYPE_Q3_K = 11, - GGML_TYPE_Q4_K = 12, - GGML_TYPE_Q5_K = 13, - GGML_TYPE_Q6_K = 14, - GGML_TYPE_Q8_K = 15, - GGML_TYPE_IQ2_XXS = 16, - GGML_TYPE_IQ2_XS = 17, - GGML_TYPE_IQ3_XXS = 18, - GGML_TYPE_IQ1_S = 19, - GGML_TYPE_IQ4_NL = 20, - GGML_TYPE_IQ3_S = 21, - GGML_TYPE_IQ2_S = 22, - GGML_TYPE_IQ4_XS = 23, - GGML_TYPE_I8 = 24, - GGML_TYPE_I16 = 25, - GGML_TYPE_I32 = 26, - GGML_TYPE_I64 = 27, - GGML_TYPE_F64 = 28, - GGML_TYPE_IQ1_M = 29, - GGML_TYPE_BF16 = 30, - // GGML_TYPE_Q4_0_4_4 = 31, support has been removed from gguf files - // GGML_TYPE_Q4_0_4_8 = 32, - // GGML_TYPE_Q4_0_8_8 = 33, - GGML_TYPE_TQ1_0 = 34, - GGML_TYPE_TQ2_0 = 35, - // GGML_TYPE_IQ4_NL_4_4 = 36, - // GGML_TYPE_IQ4_NL_4_8 = 37, - // GGML_TYPE_IQ4_NL_8_8 = 38, - GGML_TYPE_MXFP4 = 39, // MXFP4 (1 block) - GGML_TYPE_NVFP4 = 40, // NVFP4 (4 blocks, E4M3 scale) - GGML_TYPE_Q1_0 = 41, - GGML_TYPE_TQ3_0 = 42, // TurboQuant 3.5 bpv (3-bit Lloyd-Max + FWHT rotation) - GGML_TYPE_Q4_0_ROCMFP4 = 100, - GGML_TYPE_Q4_0_ROCMFP4_FAST = 101, - GGML_TYPE_Q6_0_ROCMFPX = 102, - GGML_TYPE_Q8_0_ROCMFPX = 103, - GGML_TYPE_Q3_0_ROCMFPX = 104, - GGML_TYPE_Q2_0_ROCMFP2 = 107, - GGML_TYPE_COUNT = 108, - }; - - // precision - enum ggml_prec { - GGML_PREC_DEFAULT = 0, // stored as ggml_tensor.op_params, 0 by default - GGML_PREC_F32 = 10, - }; - - // model file types - enum ggml_ftype { - GGML_FTYPE_UNKNOWN = -1, - GGML_FTYPE_ALL_F32 = 0, - GGML_FTYPE_MOSTLY_F16 = 1, // except 1d tensors - GGML_FTYPE_MOSTLY_Q4_0 = 2, // except 1d tensors - GGML_FTYPE_MOSTLY_Q4_1 = 3, // except 1d tensors - GGML_FTYPE_MOSTLY_Q4_1_SOME_F16 = 4, // tok_embeddings.weight and output.weight are F16 - GGML_FTYPE_MOSTLY_Q8_0 = 7, // except 1d tensors - GGML_FTYPE_MOSTLY_Q5_0 = 8, // except 1d tensors - GGML_FTYPE_MOSTLY_Q5_1 = 9, // except 1d tensors - GGML_FTYPE_MOSTLY_Q2_K = 10, // except 1d tensors - GGML_FTYPE_MOSTLY_Q3_K = 11, // except 1d tensors - GGML_FTYPE_MOSTLY_Q4_K = 12, // except 1d tensors - GGML_FTYPE_MOSTLY_Q5_K = 13, // except 1d tensors - GGML_FTYPE_MOSTLY_Q6_K = 14, // except 1d tensors - GGML_FTYPE_MOSTLY_IQ2_XXS = 15, // except 1d tensors - GGML_FTYPE_MOSTLY_IQ2_XS = 16, // except 1d tensors - GGML_FTYPE_MOSTLY_IQ3_XXS = 17, // except 1d tensors - GGML_FTYPE_MOSTLY_IQ1_S = 18, // except 1d tensors - GGML_FTYPE_MOSTLY_IQ4_NL = 19, // except 1d tensors - GGML_FTYPE_MOSTLY_IQ3_S = 20, // except 1d tensors - GGML_FTYPE_MOSTLY_IQ2_S = 21, // except 1d tensors - GGML_FTYPE_MOSTLY_IQ4_XS = 22, // except 1d tensors - GGML_FTYPE_MOSTLY_IQ1_M = 23, // except 1d tensors - GGML_FTYPE_MOSTLY_BF16 = 24, // except 1d tensors - GGML_FTYPE_MOSTLY_MXFP4 = 25, // except 1d tensors - GGML_FTYPE_MOSTLY_NVFP4 = 26, // except 1d tensors - GGML_FTYPE_MOSTLY_Q1_0 = 27, // except 1d tensors - GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4 = 100, - GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_LEAN = 101, - GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_COHERENT = 102, - GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_FAST = 103, - GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_FAST_COHERENT = 104, - GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_STRIX = 105, - GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_STRIX_LEAN = 106, - GGML_FTYPE_MOSTLY_Q6_0_ROCMFPX = 110, - GGML_FTYPE_MOSTLY_Q8_0_ROCMFPX = 111, - GGML_FTYPE_MOSTLY_Q3_0_ROCMFPX = 112, - GGML_FTYPE_MOSTLY_Q2_0_ROCMFP2 = 118, - GGML_FTYPE_MOSTLY_Q2_0_ROCMFP2_STRIX = 119, - }; - - // available tensor operations: - enum ggml_op { - GGML_OP_NONE = 0, - - GGML_OP_DUP, - GGML_OP_ADD, - GGML_OP_ADD_ID, - GGML_OP_ADD1, - GGML_OP_ACC, - GGML_OP_SUB, - GGML_OP_MUL, - GGML_OP_DIV, - GGML_OP_SQR, - GGML_OP_SQRT, - GGML_OP_LOG, - GGML_OP_SIN, - GGML_OP_COS, - GGML_OP_SUM, - GGML_OP_SUM_ROWS, - GGML_OP_CUMSUM, - GGML_OP_MEAN, - GGML_OP_ARGMAX, - GGML_OP_COUNT_EQUAL, - GGML_OP_REPEAT, - GGML_OP_REPEAT_BACK, - GGML_OP_CONCAT, - GGML_OP_SILU_BACK, - GGML_OP_NORM, // normalize - GGML_OP_RMS_NORM, - GGML_OP_RMS_NORM_BACK, - GGML_OP_GROUP_NORM, - GGML_OP_L2_NORM, - - GGML_OP_MUL_MAT, - GGML_OP_MUL_MAT_ID, - GGML_OP_OUT_PROD, - - GGML_OP_SCALE, - GGML_OP_SET, - GGML_OP_CPY, - GGML_OP_CONT, - GGML_OP_RESHAPE, - GGML_OP_VIEW, - GGML_OP_PERMUTE, - GGML_OP_TRANSPOSE, - GGML_OP_GET_ROWS, - GGML_OP_GET_ROWS_BACK, - GGML_OP_SET_ROWS, - GGML_OP_DIAG, - GGML_OP_DIAG_MASK_INF, - GGML_OP_DIAG_MASK_ZERO, - GGML_OP_SOFT_MAX, - GGML_OP_SOFT_MAX_BACK, - GGML_OP_ROPE, - GGML_OP_ROPE_BACK, - GGML_OP_CLAMP, - GGML_OP_CONV_TRANSPOSE_1D, - GGML_OP_IM2COL, - GGML_OP_IM2COL_BACK, - GGML_OP_IM2COL_3D, - GGML_OP_CONV_2D, - GGML_OP_CONV_3D, - GGML_OP_CONV_2D_DW, - GGML_OP_CONV_TRANSPOSE_2D, - GGML_OP_POOL_1D, - GGML_OP_POOL_2D, - GGML_OP_POOL_2D_BACK, - GGML_OP_UPSCALE, - GGML_OP_PAD, - GGML_OP_PAD_REFLECT_1D, - GGML_OP_ROLL, - GGML_OP_ARANGE, - GGML_OP_TIMESTEP_EMBEDDING, - GGML_OP_ARGSORT, - GGML_OP_TOP_K, - GGML_OP_LEAKY_RELU, - GGML_OP_TRI, - GGML_OP_FILL, - - GGML_OP_FLASH_ATTN_EXT, - GGML_OP_FLASH_ATTN_SPARSE, - GGML_OP_FLASH_ATTN_BACK, - GGML_OP_SSM_CONV, - GGML_OP_SSM_SCAN, - GGML_OP_WIN_PART, - GGML_OP_WIN_UNPART, - GGML_OP_GET_REL_POS, - GGML_OP_ADD_REL_POS, - GGML_OP_RWKV_WKV6, - GGML_OP_GATED_LINEAR_ATTN, - GGML_OP_RWKV_WKV7, - GGML_OP_SOLVE_TRI, - GGML_OP_GATED_DELTA_NET, - - GGML_OP_UNARY, - - GGML_OP_MAP_CUSTOM1, - GGML_OP_MAP_CUSTOM2, - GGML_OP_MAP_CUSTOM3, - - GGML_OP_CUSTOM, - - GGML_OP_CROSS_ENTROPY_LOSS, - GGML_OP_CROSS_ENTROPY_LOSS_BACK, - GGML_OP_OPT_STEP_ADAMW, - GGML_OP_OPT_STEP_SGD, - - GGML_OP_GLU, - - GGML_OP_TURBO_WHT, // FWHT rotation for TurboQuant KV cache - - GGML_OP_MOE_FUSED, // Fused MoE FFN: gate+up+swiglu+down+weighted_sum+shared_expert - - GGML_OP_COUNT, - }; - - enum ggml_unary_op { - GGML_UNARY_OP_ABS, - GGML_UNARY_OP_SGN, - GGML_UNARY_OP_NEG, - GGML_UNARY_OP_STEP, - GGML_UNARY_OP_TANH, - GGML_UNARY_OP_ELU, - GGML_UNARY_OP_RELU, - GGML_UNARY_OP_SIGMOID, - GGML_UNARY_OP_GELU, - GGML_UNARY_OP_GELU_QUICK, - GGML_UNARY_OP_SILU, - GGML_UNARY_OP_HARDSWISH, - GGML_UNARY_OP_HARDSIGMOID, - GGML_UNARY_OP_EXP, - GGML_UNARY_OP_EXPM1, - GGML_UNARY_OP_SOFTPLUS, - GGML_UNARY_OP_GELU_ERF, - GGML_UNARY_OP_XIELU, - GGML_UNARY_OP_FLOOR, - GGML_UNARY_OP_CEIL, - GGML_UNARY_OP_ROUND, - GGML_UNARY_OP_TRUNC, - - GGML_UNARY_OP_COUNT, - }; - - enum ggml_glu_op { - GGML_GLU_OP_REGLU, - GGML_GLU_OP_GEGLU, - GGML_GLU_OP_SWIGLU, - GGML_GLU_OP_SWIGLU_OAI, - GGML_GLU_OP_SWIGLU_DS4, - GGML_GLU_OP_GEGLU_ERF, - GGML_GLU_OP_GEGLU_QUICK, - - GGML_GLU_OP_COUNT, - }; - - enum ggml_object_type { - GGML_OBJECT_TYPE_TENSOR, - GGML_OBJECT_TYPE_GRAPH, - GGML_OBJECT_TYPE_WORK_BUFFER - }; - - enum ggml_log_level { - GGML_LOG_LEVEL_NONE = 0, - GGML_LOG_LEVEL_DEBUG = 1, - GGML_LOG_LEVEL_INFO = 2, - GGML_LOG_LEVEL_WARN = 3, - GGML_LOG_LEVEL_ERROR = 4, - GGML_LOG_LEVEL_CONT = 5, // continue previous log - }; - - // this tensor... - enum ggml_tensor_flag { - GGML_TENSOR_FLAG_INPUT = 1, // ...is an input for the GGML compute graph - GGML_TENSOR_FLAG_OUTPUT = 2, // ...is an output for the GGML compute graph - GGML_TENSOR_FLAG_PARAM = 4, // ...contains trainable parameters - GGML_TENSOR_FLAG_LOSS = 8, // ...defines loss for numerical optimization (multiple loss tensors add up) - GGML_TENSOR_FLAG_COMPUTE = 16, // ...must be computed - }; - - enum ggml_tri_type { - GGML_TRI_TYPE_UPPER_DIAG = 0, - GGML_TRI_TYPE_UPPER = 1, - GGML_TRI_TYPE_LOWER_DIAG = 2, - GGML_TRI_TYPE_LOWER = 3 - }; - - struct ggml_init_params { - // memory pool - size_t mem_size; // bytes - void * mem_buffer; // if NULL, memory will be allocated internally - bool no_alloc; // don't allocate memory for the tensor data - }; - - // n-dimensional tensor - struct ggml_tensor { - enum ggml_type type; - - struct ggml_backend_buffer * buffer; - - int64_t ne[GGML_MAX_DIMS]; // number of elements - size_t nb[GGML_MAX_DIMS]; // stride in bytes: - // nb[0] = ggml_type_size(type) - // nb[1] = nb[0] * (ne[0] / ggml_blck_size(type)) + padding - // nb[i] = nb[i-1] * ne[i-1] - - // compute data - enum ggml_op op; - - // op params - allocated as int32_t for alignment - int32_t op_params[GGML_MAX_OP_PARAMS / sizeof(int32_t)]; - - int32_t flags; - - struct ggml_tensor * src[GGML_MAX_SRC]; - - // source tensor and offset for views - struct ggml_tensor * view_src; - size_t view_offs; - - void * data; - - char name[GGML_MAX_NAME]; - - void * extra; // extra things e.g. for ggml-cuda.cu - - char padding[8]; - }; - - static const size_t GGML_TENSOR_SIZE = sizeof(struct ggml_tensor); - - // Abort callback - // If not NULL, called before ggml computation - // If it returns true, the computation is aborted - typedef bool (*ggml_abort_callback)(void * data); - - - // - // GUID - // - - // GUID types - typedef uint8_t ggml_guid[16]; - typedef ggml_guid * ggml_guid_t; - - GGML_API bool ggml_guid_matches(ggml_guid_t guid_a, ggml_guid_t guid_b); - - // misc - - GGML_API const char * ggml_version(void); - GGML_API const char * ggml_commit(void); - - GGML_API void ggml_time_init(void); // call this once at the beginning of the program - GGML_API int64_t ggml_time_ms(void); - GGML_API int64_t ggml_time_us(void); - GGML_API int64_t ggml_cycles(void); - GGML_API int64_t ggml_cycles_per_ms(void); - - // accepts a UTF-8 path, even on Windows - GGML_API FILE * ggml_fopen(const char * fname, const char * mode); - - GGML_API void ggml_print_object (const struct ggml_object * obj); - GGML_API void ggml_print_objects(const struct ggml_context * ctx); - - GGML_API int64_t ggml_nelements (const struct ggml_tensor * tensor); - GGML_API int64_t ggml_nrows (const struct ggml_tensor * tensor); - GGML_API size_t ggml_nbytes (const struct ggml_tensor * tensor); - GGML_API size_t ggml_nbytes_pad(const struct ggml_tensor * tensor); // same as ggml_nbytes() but padded to GGML_MEM_ALIGN - - GGML_API int64_t ggml_blck_size(enum ggml_type type); - GGML_API size_t ggml_type_size(enum ggml_type type); // size in bytes for all elements in a block - GGML_API size_t ggml_row_size (enum ggml_type type, int64_t ne); // size in bytes for all elements in a row - - GGML_DEPRECATED( - GGML_API double ggml_type_sizef(enum ggml_type type), // ggml_type_size()/ggml_blck_size() as float - "use ggml_row_size() instead"); - - GGML_API const char * ggml_type_name(enum ggml_type type); - GGML_API const char * ggml_op_name (enum ggml_op op); - GGML_API const char * ggml_op_symbol(enum ggml_op op); - - GGML_API const char * ggml_unary_op_name(enum ggml_unary_op op); - GGML_API const char * ggml_glu_op_name(enum ggml_glu_op op); - GGML_API const char * ggml_op_desc(const struct ggml_tensor * t); // unary or op name - - GGML_API size_t ggml_element_size(const struct ggml_tensor * tensor); - - GGML_API bool ggml_is_quantized(enum ggml_type type); - - // TODO: temporary until model loading of ggml examples is refactored - GGML_API enum ggml_type ggml_ftype_to_ggml_type(enum ggml_ftype ftype); - - GGML_API bool ggml_is_transposed(const struct ggml_tensor * tensor); - GGML_API bool ggml_is_permuted (const struct ggml_tensor * tensor); - GGML_API bool ggml_is_empty (const struct ggml_tensor * tensor); - GGML_API bool ggml_is_view (const struct ggml_tensor * tensor); - GGML_API bool ggml_is_scalar (const struct ggml_tensor * tensor); - GGML_API bool ggml_is_vector (const struct ggml_tensor * tensor); - GGML_API bool ggml_is_matrix (const struct ggml_tensor * tensor); - GGML_API bool ggml_is_3d (const struct ggml_tensor * tensor); - GGML_API int ggml_n_dims (const struct ggml_tensor * tensor); // returns 1 for scalars - - // returns whether the tensor elements can be iterated over with a flattened index (no gaps, no permutation) - GGML_API bool ggml_is_contiguous (const struct ggml_tensor * tensor); - GGML_API bool ggml_is_contiguous_0(const struct ggml_tensor * tensor); // same as ggml_is_contiguous() - GGML_API bool ggml_is_contiguous_1(const struct ggml_tensor * tensor); // contiguous for dims >= 1 - GGML_API bool ggml_is_contiguous_2(const struct ggml_tensor * tensor); // contiguous for dims >= 2 - - // returns whether the tensor elements are allocated as one contiguous block of memory (no gaps, but permutation ok) - GGML_API bool ggml_is_contiguously_allocated(const struct ggml_tensor * tensor); - - // true for tensor that is stored in memory as CxWxHxN and has been permuted to WxHxCxN - GGML_API bool ggml_is_contiguous_channels(const struct ggml_tensor * tensor); - - // true if the elements in dimension 0 are contiguous, or there is just 1 block of elements - GGML_API bool ggml_is_contiguous_rows(const struct ggml_tensor * tensor); - - GGML_API bool ggml_are_same_shape (const struct ggml_tensor * t0, const struct ggml_tensor * t1); - GGML_API bool ggml_are_same_stride(const struct ggml_tensor * t0, const struct ggml_tensor * t1); - - GGML_API bool ggml_can_repeat(const struct ggml_tensor * t0, const struct ggml_tensor * t1); - - // use this to compute the memory overhead of a tensor - GGML_API size_t ggml_tensor_overhead(void); - - GGML_API bool ggml_validate_row_data(enum ggml_type type, const void * data, size_t nbytes); - - // main - - GGML_API struct ggml_context * ggml_init (struct ggml_init_params params); - GGML_API void ggml_reset(struct ggml_context * ctx); - GGML_API void ggml_free (struct ggml_context * ctx); - - GGML_API size_t ggml_used_mem(const struct ggml_context * ctx); - - GGML_API bool ggml_get_no_alloc(struct ggml_context * ctx); - GGML_API void ggml_set_no_alloc(struct ggml_context * ctx, bool no_alloc); - - GGML_API void * ggml_get_mem_buffer (const struct ggml_context * ctx); - GGML_API size_t ggml_get_mem_size (const struct ggml_context * ctx); - GGML_API size_t ggml_get_max_tensor_size(const struct ggml_context * ctx); - - GGML_API struct ggml_tensor * ggml_new_tensor( - struct ggml_context * ctx, - enum ggml_type type, - int n_dims, - const int64_t *ne); - - GGML_API struct ggml_tensor * ggml_new_tensor_1d( - struct ggml_context * ctx, - enum ggml_type type, - int64_t ne0); - - GGML_API struct ggml_tensor * ggml_new_tensor_2d( - struct ggml_context * ctx, - enum ggml_type type, - int64_t ne0, - int64_t ne1); - - GGML_API struct ggml_tensor * ggml_new_tensor_3d( - struct ggml_context * ctx, - enum ggml_type type, - int64_t ne0, - int64_t ne1, - int64_t ne2); - - GGML_API struct ggml_tensor * ggml_new_tensor_4d( - struct ggml_context * ctx, - enum ggml_type type, - int64_t ne0, - int64_t ne1, - int64_t ne2, - int64_t ne3); - - GGML_API void * ggml_new_buffer(struct ggml_context * ctx, size_t nbytes); - - GGML_API struct ggml_tensor * ggml_dup_tensor (struct ggml_context * ctx, const struct ggml_tensor * src); - GGML_API struct ggml_tensor * ggml_view_tensor(struct ggml_context * ctx, struct ggml_tensor * src); - - // Context tensor enumeration and lookup - GGML_API struct ggml_tensor * ggml_get_first_tensor(const struct ggml_context * ctx); - GGML_API struct ggml_tensor * ggml_get_next_tensor (const struct ggml_context * ctx, struct ggml_tensor * tensor); - GGML_API struct ggml_tensor * ggml_get_tensor(struct ggml_context * ctx, const char * name); - - // Converts a flat index into coordinates - GGML_API void ggml_unravel_index(const struct ggml_tensor * tensor, int64_t i, int64_t * i0, int64_t * i1, int64_t * i2, int64_t * i3); - - GGML_API enum ggml_unary_op ggml_get_unary_op(const struct ggml_tensor * tensor); - GGML_API enum ggml_glu_op ggml_get_glu_op(const struct ggml_tensor * tensor); - - GGML_API void * ggml_get_data (const struct ggml_tensor * tensor); - GGML_API float * ggml_get_data_f32(const struct ggml_tensor * tensor); - - GGML_API const char * ggml_get_name (const struct ggml_tensor * tensor); - GGML_API struct ggml_tensor * ggml_set_name ( struct ggml_tensor * tensor, const char * name); - GGML_ATTRIBUTE_FORMAT(2, 3) - GGML_API struct ggml_tensor * ggml_format_name( struct ggml_tensor * tensor, const char * fmt, ...); - - // Tensor flags - GGML_API void ggml_set_input(struct ggml_tensor * tensor); - GGML_API void ggml_set_output(struct ggml_tensor * tensor); - GGML_API void ggml_set_param(struct ggml_tensor * tensor); - GGML_API void ggml_set_loss(struct ggml_tensor * tensor); - - // - // operations on tensors with backpropagation - // - - GGML_API struct ggml_tensor * ggml_dup( - struct ggml_context * ctx, - struct ggml_tensor * a); - - // in-place, returns view(a) - GGML_API struct ggml_tensor * ggml_dup_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_add( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b); - - GGML_API struct ggml_tensor * ggml_add_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b); - - GGML_API struct ggml_tensor * ggml_add_cast( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - enum ggml_type type); - - // dst[i0, i1, i2] = a[i0, i1, i2] + b[i0, ids[i1, i2]] - GGML_API struct ggml_tensor * ggml_add_id( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * ids); - - GGML_DEPRECATED(GGML_API struct ggml_tensor * ggml_add1( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b), - "use ggml_add instead"); - - GGML_DEPRECATED(GGML_API struct ggml_tensor * ggml_add1_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b), - "use ggml_add_inplace instead"); - - // dst = a - // view(dst, nb1, nb2, nb3, offset) += b - // return dst - GGML_API struct ggml_tensor * ggml_acc( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - size_t nb1, - size_t nb2, - size_t nb3, - size_t offset); - - GGML_API struct ggml_tensor * ggml_acc_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - size_t nb1, - size_t nb2, - size_t nb3, - size_t offset); - - GGML_API struct ggml_tensor * ggml_sub( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b); - - GGML_API struct ggml_tensor * ggml_sub_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b); - - GGML_API struct ggml_tensor * ggml_mul( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b); - - GGML_API struct ggml_tensor * ggml_mul_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b); - - GGML_API struct ggml_tensor * ggml_div( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b); - - GGML_API struct ggml_tensor * ggml_div_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b); - - GGML_API struct ggml_tensor * ggml_sqr( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_sqr_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_sqrt( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_sqrt_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_log( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_log_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_expm1( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_expm1_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_softplus( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_softplus_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_sin( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_sin_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_cos( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_cos_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a); - - // return scalar - GGML_API struct ggml_tensor * ggml_sum( - struct ggml_context * ctx, - struct ggml_tensor * a); - - // sums along rows, with input shape [a,b,c,d] return shape [1,b,c,d] - GGML_API struct ggml_tensor * ggml_sum_rows( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_cumsum( - struct ggml_context * ctx, - struct ggml_tensor * a); - - // mean along rows - GGML_API struct ggml_tensor * ggml_mean( - struct ggml_context * ctx, - struct ggml_tensor * a); - - // argmax along rows - GGML_API struct ggml_tensor * ggml_argmax( - struct ggml_context * ctx, - struct ggml_tensor * a); - - // count number of equal elements in a and b - GGML_API struct ggml_tensor * ggml_count_equal( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b); - - // if a is the same shape as b, and a is not parameter, return a - // otherwise, return a new tensor: repeat(a) to fit in b - GGML_API struct ggml_tensor * ggml_repeat( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b); - - // repeat a to the specified shape - GGML_API struct ggml_tensor * ggml_repeat_4d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - int64_t ne1, - int64_t ne2, - int64_t ne3); - - // sums repetitions in a into shape of b - GGML_API struct ggml_tensor * ggml_repeat_back( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b); // sum up values that are adjacent in dims > 0 instead of repeated with same stride - - // concat a and b along dim - // used in stable-diffusion - GGML_API struct ggml_tensor * ggml_concat( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int dim); - - GGML_API struct ggml_tensor * ggml_abs( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_abs_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_sgn( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_sgn_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_neg( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_neg_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_step( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_step_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_tanh( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_tanh_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_elu( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_elu_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_relu( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_leaky_relu( - struct ggml_context * ctx, - struct ggml_tensor * a, float negative_slope, bool inplace); - - GGML_API struct ggml_tensor * ggml_relu_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_sigmoid( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_sigmoid_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_gelu( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_gelu_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a); - - // GELU using erf (error function) when possible - // some backends may fallback to approximation based on Abramowitz and Stegun formula - GGML_API struct ggml_tensor * ggml_gelu_erf( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_gelu_erf_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_gelu_quick( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_gelu_quick_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_silu( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_silu_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a); - - // a - x - // b - dy - GGML_API struct ggml_tensor * ggml_silu_back( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b); - - // hardswish(x) = x * relu6(x + 3) / 6 - GGML_API struct ggml_tensor * ggml_hardswish( - struct ggml_context * ctx, - struct ggml_tensor * a); - - // hardsigmoid(x) = relu6(x + 3) / 6 - GGML_API struct ggml_tensor * ggml_hardsigmoid( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_exp( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_exp_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_floor( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_floor_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_ceil( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_ceil_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_round( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_round_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a); - - /** - * Truncates the fractional part of each element in the tensor (towards zero). - * For example: trunc(3.7) = 3.0, trunc(-2.9) = -2.0 - * Similar to std::trunc in C/C++. - */ - - GGML_API struct ggml_tensor * ggml_trunc( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_trunc_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a); - - - - // xIELU activation function - // x = x * (c_a(alpha_n) + c_b(alpha_p, beta) * sigmoid(beta * x)) + eps * (x > 0) - // where c_a = softplus and c_b(a, b) = softplus(a) + b are constraining functions - // that constrain the positive and negative source alpha values respectively - GGML_API struct ggml_tensor * ggml_xielu( - struct ggml_context * ctx, - struct ggml_tensor * a, - float alpha_n, - float alpha_p, - float beta, - float eps); - - // gated linear unit ops - // A: n columns, r rows, - // result is n / 2 columns, r rows, - // expects gate in second half of row, unless swapped is true - GGML_API struct ggml_tensor * ggml_glu( - struct ggml_context * ctx, - struct ggml_tensor * a, - enum ggml_glu_op op, - bool swapped); - - GGML_API struct ggml_tensor * ggml_reglu( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_reglu_swapped( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_geglu( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_geglu_swapped( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_swiglu( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_swiglu_swapped( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_geglu_erf( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_geglu_erf_swapped( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_geglu_quick( - struct ggml_context * ctx, - struct ggml_tensor * a); - - GGML_API struct ggml_tensor * ggml_geglu_quick_swapped( - struct ggml_context * ctx, - struct ggml_tensor * a); - - // A: n columns, r rows, - // B: n columns, r rows, - GGML_API struct ggml_tensor * ggml_glu_split( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - enum ggml_glu_op op); - - GGML_API struct ggml_tensor * ggml_reglu_split( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b); - - GGML_API struct ggml_tensor * ggml_geglu_split( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b); - - GGML_API struct ggml_tensor * ggml_swiglu_split( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b); - - GGML_API struct ggml_tensor * ggml_swiglu_ds4_split( - struct ggml_context * ctx, - struct ggml_tensor * gate, - struct ggml_tensor * up, - float clamp); - - GGML_API struct ggml_tensor * ggml_geglu_erf_split( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b); - - GGML_API struct ggml_tensor * ggml_geglu_quick_split( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b); - - GGML_API struct ggml_tensor * ggml_swiglu_oai( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - float alpha, - float limit); - - // normalize along rows - GGML_API struct ggml_tensor * ggml_norm( - struct ggml_context * ctx, - struct ggml_tensor * a, - float eps); - - GGML_API struct ggml_tensor * ggml_norm_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - float eps); - - GGML_API struct ggml_tensor * ggml_rms_norm( - struct ggml_context * ctx, - struct ggml_tensor * a, - float eps); - - GGML_API struct ggml_tensor * ggml_rms_norm_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - float eps); - - // group normalize along ne0*ne1*n_groups - // used in stable-diffusion - GGML_API struct ggml_tensor * ggml_group_norm( - struct ggml_context * ctx, - struct ggml_tensor * a, - int n_groups, - float eps); - - GGML_API struct ggml_tensor * ggml_group_norm_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - int n_groups, - float eps); - - // l2 normalize along rows - // used in rwkv v7 - GGML_API struct ggml_tensor * ggml_l2_norm( - struct ggml_context * ctx, - struct ggml_tensor * a, - float eps); - - GGML_API struct ggml_tensor * ggml_l2_norm_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - float eps); - - // a - x - // b - dy - GGML_API struct ggml_tensor * ggml_rms_norm_back( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - float eps); - - // A: k columns, n rows => [ne03, ne02, n, k] - // B: k columns, m rows (i.e. we transpose it internally) => [ne03 * x, ne02 * y, m, k] - // result is n columns, m rows => [ne03 * x, ne02 * y, m, n] - GGML_API struct ggml_tensor * ggml_mul_mat( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b); - - // change the precision of a matrix multiplication - // set to GGML_PREC_F32 for higher precision (useful for phi-2) - GGML_API void ggml_mul_mat_set_prec( - struct ggml_tensor * a, - enum ggml_prec prec); - - // indirect matrix multiplication - GGML_API struct ggml_tensor * ggml_mul_mat_id( - struct ggml_context * ctx, - struct ggml_tensor * as, - struct ggml_tensor * b, - struct ggml_tensor * ids); - - // A: m columns, n rows, - // B: p columns, n rows, - // result is m columns, p rows - GGML_API struct ggml_tensor * ggml_out_prod( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b); - - // - // operations on tensors without backpropagation - // - - GGML_API struct ggml_tensor * ggml_scale( - struct ggml_context * ctx, - struct ggml_tensor * a, - float s); - - // in-place, returns view(a) - GGML_API struct ggml_tensor * ggml_scale_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - float s); - - // x = s * a + b - GGML_API struct ggml_tensor * ggml_scale_bias( - struct ggml_context * ctx, - struct ggml_tensor * a, - float s, - float b); - - GGML_API struct ggml_tensor * ggml_scale_bias_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - float s, - float b); - - // b -> view(a,offset,nb1,nb2,3), return modified a - GGML_API struct ggml_tensor * ggml_set( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - size_t nb1, - size_t nb2, - size_t nb3, - size_t offset); // in bytes - - // b -> view(a,offset,nb1,nb2,3), return view(a) - GGML_API struct ggml_tensor * ggml_set_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - size_t nb1, - size_t nb2, - size_t nb3, - size_t offset); // in bytes - - GGML_API struct ggml_tensor * ggml_set_1d( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - size_t offset); // in bytes - - GGML_API struct ggml_tensor * ggml_set_1d_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - size_t offset); // in bytes - - // b -> view(a,offset,nb1,nb2,3), return modified a - GGML_API struct ggml_tensor * ggml_set_2d( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - size_t nb1, - size_t offset); // in bytes - - // b -> view(a,offset,nb1,nb2,3), return view(a) - GGML_API struct ggml_tensor * ggml_set_2d_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - size_t nb1, - size_t offset); // in bytes - - // a -> b, return view(b) - GGML_API struct ggml_tensor * ggml_cpy( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b); - - // note: casting from f32 to i32 will discard the fractional part - GGML_API struct ggml_tensor * ggml_cast( - struct ggml_context * ctx, - struct ggml_tensor * a, - enum ggml_type type); - - // make contiguous - GGML_API struct ggml_tensor * ggml_cont( - struct ggml_context * ctx, - struct ggml_tensor * a); - - // make contiguous, with new shape - GGML_API struct ggml_tensor * ggml_cont_1d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0); - - GGML_API struct ggml_tensor * ggml_cont_2d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - int64_t ne1); - - GGML_API struct ggml_tensor * ggml_cont_3d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - int64_t ne1, - int64_t ne2); - - GGML_API struct ggml_tensor * ggml_cont_4d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - int64_t ne1, - int64_t ne2, - int64_t ne3); - - // return view(a), b specifies the new shape - // TODO: when we start computing gradient, make a copy instead of view - GGML_API struct ggml_tensor * ggml_reshape( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b); - - // return view(a) - // TODO: when we start computing gradient, make a copy instead of view - GGML_API struct ggml_tensor * ggml_reshape_1d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0); - - GGML_API struct ggml_tensor * ggml_reshape_2d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - int64_t ne1); - - // return view(a) - // TODO: when we start computing gradient, make a copy instead of view - GGML_API struct ggml_tensor * ggml_reshape_3d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - int64_t ne1, - int64_t ne2); - - GGML_API struct ggml_tensor * ggml_reshape_4d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - int64_t ne1, - int64_t ne2, - int64_t ne3); - - // offset in bytes - GGML_API struct ggml_tensor * ggml_view_1d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - size_t offset); - - GGML_API struct ggml_tensor * ggml_view_2d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - int64_t ne1, - size_t nb1, // row stride in bytes - size_t offset); - - GGML_API struct ggml_tensor * ggml_view_3d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - int64_t ne1, - int64_t ne2, - size_t nb1, // row stride in bytes - size_t nb2, // slice stride in bytes - size_t offset); - - GGML_API struct ggml_tensor * ggml_view_4d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - int64_t ne1, - int64_t ne2, - int64_t ne3, - size_t nb1, // row stride in bytes - size_t nb2, // slice stride in bytes - size_t nb3, - size_t offset); - - GGML_API struct ggml_tensor * ggml_permute( - struct ggml_context * ctx, - struct ggml_tensor * a, - int axis0, - int axis1, - int axis2, - int axis3); - - // alias for ggml_permute(ctx, a, 1, 0, 2, 3) - GGML_API struct ggml_tensor * ggml_transpose( - struct ggml_context * ctx, - struct ggml_tensor * a); - - // supports 4D a: - // a [n_embd, ne1, ne2, ne3] - // b I32 [n_rows, ne2, ne3, 1] - // - // return [n_embd, n_rows, ne2, ne3] - GGML_API struct ggml_tensor * ggml_get_rows( - struct ggml_context * ctx, - struct ggml_tensor * a, // data - struct ggml_tensor * b); // row indices - - GGML_API struct ggml_tensor * ggml_get_rows_back( - struct ggml_context * ctx, - struct ggml_tensor * a, // gradients of ggml_get_rows result - struct ggml_tensor * b, // row indices - struct ggml_tensor * c); // data for ggml_get_rows, only used for its shape - - // a TD [n_embd, ne1, ne2, ne3] - // b TS [n_embd, n_rows, ne02, ne03] | ne02 == ne2, ne03 == ne3 - // c I64 [n_rows, ne11, ne12, 1] | c[i] in [0, ne1) - // - // undefined behavior if destination rows overlap - // - // broadcast: - // ne2 % ne11 == 0 - // ne3 % ne12 == 0 - // - // return view(a) - GGML_API struct ggml_tensor * ggml_set_rows( - struct ggml_context * ctx, - struct ggml_tensor * a, // destination - struct ggml_tensor * b, // source - struct ggml_tensor * c); // row indices - - GGML_API struct ggml_tensor * ggml_diag( - struct ggml_context * ctx, - struct ggml_tensor * a); - - // set elements above the diagonal to -INF - GGML_API struct ggml_tensor * ggml_diag_mask_inf( - struct ggml_context * ctx, - struct ggml_tensor * a, - int n_past); - - // in-place, returns view(a) - GGML_API struct ggml_tensor * ggml_diag_mask_inf_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - int n_past); - - // set elements above the diagonal to 0 - GGML_API struct ggml_tensor * ggml_diag_mask_zero( - struct ggml_context * ctx, - struct ggml_tensor * a, - int n_past); - - // in-place, returns view(a) - GGML_API struct ggml_tensor * ggml_diag_mask_zero_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - int n_past); - - GGML_API struct ggml_tensor * ggml_soft_max( - struct ggml_context * ctx, - struct ggml_tensor * a); - - // in-place, returns view(a) - GGML_API struct ggml_tensor * ggml_soft_max_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a); - - // a [ne0, ne01, ne02, ne03] - // mask [ne0, ne11, ne12, ne13] | ne11 >= ne01, F16 or F32, optional - // - // broadcast: - // ne02 % ne12 == 0 - // ne03 % ne13 == 0 - // - // fused soft_max(a*scale + mask*(ALiBi slope)) - // max_bias = 0.0f for no ALiBi - GGML_API struct ggml_tensor * ggml_soft_max_ext( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * mask, - float scale, - float max_bias); - - GGML_API struct ggml_tensor * ggml_soft_max_ext_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * mask, - float scale, - float max_bias); - - GGML_API void ggml_soft_max_add_sinks( - struct ggml_tensor * a, - struct ggml_tensor * sinks); - - GGML_API struct ggml_tensor * ggml_soft_max_ext_back( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - float scale, - float max_bias); - - // in-place, returns view(a) - GGML_API struct ggml_tensor * ggml_soft_max_ext_back_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - float scale, - float max_bias); - - // rotary position embedding - // if (mode & 1) - skip n_past elements (NOT SUPPORTED) - // if (mode & GGML_ROPE_TYPE_NEOX) - GPT-NeoX style - // - // b is an int32 vector with size a->ne[2], it contains the positions - GGML_API struct ggml_tensor * ggml_rope( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int n_dims, - int mode); - - // in-place, returns view(a) - GGML_API struct ggml_tensor * ggml_rope_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int n_dims, - int mode); - - // custom RoPE - // c is freq factors (e.g. phi3-128k), (optional) - GGML_API struct ggml_tensor * ggml_rope_ext( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * c, - int n_dims, - int mode, - int n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow); - - GGML_API struct ggml_tensor * ggml_rope_multi( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * c, - int n_dims, - int sections[GGML_MROPE_SECTIONS], - int mode, - int n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow); - - // in-place, returns view(a) - GGML_API struct ggml_tensor * ggml_rope_ext_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * c, - int n_dims, - int mode, - int n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow); - - GGML_API struct ggml_tensor * ggml_rope_multi_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * c, - int n_dims, - int sections[GGML_MROPE_SECTIONS], - int mode, - int n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow); - - GGML_DEPRECATED(GGML_API struct ggml_tensor * ggml_rope_custom( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int n_dims, - int mode, - int n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow), - "use ggml_rope_ext instead"); - - GGML_DEPRECATED(GGML_API struct ggml_tensor * ggml_rope_custom_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int n_dims, - int mode, - int n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow), - "use ggml_rope_ext_inplace instead"); - - // compute correction dims for YaRN RoPE scaling - GGML_API void ggml_rope_yarn_corr_dims( - int n_dims, int n_ctx_orig, float freq_base, float beta_fast, float beta_slow, float dims[2]); - - // rotary position embedding backward, i.e compute dx from dy - // a - dy - GGML_API struct ggml_tensor * ggml_rope_ext_back( - struct ggml_context * ctx, - struct ggml_tensor * a, // gradients of ggml_rope result - struct ggml_tensor * b, // positions - struct ggml_tensor * c, // freq factors - int n_dims, - int mode, - int n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow); - - GGML_API struct ggml_tensor * ggml_rope_multi_back( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * c, - int n_dims, - int sections[4], - int mode, - int n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow); - - - // clamp - // in-place, returns view(a) - GGML_API struct ggml_tensor * ggml_clamp( - struct ggml_context * ctx, - struct ggml_tensor * a, - float min, - float max); - - // im2col - // converts data into a format that effectively results in a convolution when combined with matrix multiplication - GGML_API struct ggml_tensor * ggml_im2col( - struct ggml_context * ctx, - struct ggml_tensor * a, // convolution kernel - struct ggml_tensor * b, // data - int s0, // stride dimension 0 - int s1, // stride dimension 1 - int p0, // padding dimension 0 - int p1, // padding dimension 1 - int d0, // dilation dimension 0 - int d1, // dilation dimension 1 - bool is_2D, - enum ggml_type dst_type); - - GGML_API struct ggml_tensor * ggml_im2col_back( - struct ggml_context * ctx, - struct ggml_tensor * a, // convolution kernel - struct ggml_tensor * b, // gradient of im2col output - int64_t * ne, // shape of im2col input - int s0, // stride dimension 0 - int s1, // stride dimension 1 - int p0, // padding dimension 0 - int p1, // padding dimension 1 - int d0, // dilation dimension 0 - int d1, // dilation dimension 1 - bool is_2D); - - GGML_API struct ggml_tensor * ggml_conv_1d( - struct ggml_context * ctx, - struct ggml_tensor * a, // convolution kernel - struct ggml_tensor * b, // data - int s0, // stride - int p0, // padding - int d0); // dilation - - // conv_1d with padding = half - // alias for ggml_conv_1d(a, b, s, a->ne[0]/2, d) - GGML_API struct ggml_tensor* ggml_conv_1d_ph( - struct ggml_context * ctx, - struct ggml_tensor * a, // convolution kernel - struct ggml_tensor * b, // data - int s, // stride - int d); // dilation - - // depthwise - // TODO: this is very likely wrong for some cases! - needs more testing - GGML_API struct ggml_tensor * ggml_conv_1d_dw( - struct ggml_context * ctx, - struct ggml_tensor * a, // convolution kernel - struct ggml_tensor * b, // data - int s0, // stride - int p0, // padding - int d0); // dilation - - GGML_API struct ggml_tensor * ggml_conv_1d_dw_ph( - struct ggml_context * ctx, - struct ggml_tensor * a, // convolution kernel - struct ggml_tensor * b, // data - int s0, // stride - int d0); // dilation - - GGML_API struct ggml_tensor * ggml_conv_transpose_1d( - struct ggml_context * ctx, - struct ggml_tensor * a, // convolution kernel - struct ggml_tensor * b, // data - int s0, // stride - int p0, // padding - int d0); // dilation - - GGML_API struct ggml_tensor * ggml_conv_2d( - struct ggml_context * ctx, - struct ggml_tensor * a, // convolution kernel - struct ggml_tensor * b, // data - int s0, // stride dimension 0 - int s1, // stride dimension 1 - int p0, // padding dimension 0 - int p1, // padding dimension 1 - int d0, // dilation dimension 0 - int d1); // dilation dimension 1 - - GGML_API struct ggml_tensor * ggml_im2col_3d( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int64_t IC, - int s0, // stride width - int s1, // stride height - int s2, // stride depth - int p0, // padding width - int p1, // padding height - int p2, // padding depth - int d0, // dilation width - int d1, // dilation height - int d2, // dilation depth - enum ggml_type dst_type); - - // a: [OC*IC, KD, KH, KW] - // b: [N*IC, ID, IH, IW] - // result: [N*OC, OD, OH, OW] - GGML_API struct ggml_tensor * ggml_conv_3d( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int64_t IC, - int s0, // stride width - int s1, // stride height - int s2, // stride depth - int p0, // padding width - int p1, // padding height - int p2, // padding depth - int d0, // dilation width - int d1, // dilation height - int d2 // dilation depth - ); - - // kernel size is a->ne[0] x a->ne[1] - // stride is equal to kernel size - // padding is zero - // example: - // a: 16 16 3 768 - // b: 1024 1024 3 1 - // res: 64 64 768 1 - // used in sam - GGML_API struct ggml_tensor * ggml_conv_2d_sk_p0( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b); - - // kernel size is a->ne[0] x a->ne[1] - // stride is 1 - // padding is half - // example: - // a: 3 3 256 256 - // b: 64 64 256 1 - // res: 64 64 256 1 - // used in sam - GGML_API struct ggml_tensor * ggml_conv_2d_s1_ph( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b); - - // depthwise (via im2col and mul_mat) - GGML_API struct ggml_tensor * ggml_conv_2d_dw( - struct ggml_context * ctx, - struct ggml_tensor * a, // convolution kernel - struct ggml_tensor * b, // data - int s0, // stride dimension 0 - int s1, // stride dimension 1 - int p0, // padding dimension 0 - int p1, // padding dimension 1 - int d0, // dilation dimension 0 - int d1); // dilation dimension 1 - - // Depthwise 2D convolution - // may be faster than ggml_conv_2d_dw, but not available in all backends - // a: KW KH 1 C convolution kernel - // b: W H C N input data - // res: W_out H_out C N - GGML_API struct ggml_tensor * ggml_conv_2d_dw_direct( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int stride0, - int stride1, - int pad0, - int pad1, - int dilation0, - int dilation1); - - GGML_API struct ggml_tensor * ggml_conv_transpose_2d_p0( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int stride); - - GGML_API struct ggml_tensor * ggml_conv_2d_direct( - struct ggml_context * ctx, - struct ggml_tensor * a, // convolution kernel [KW, KH, IC, OC] - struct ggml_tensor * b, // input data [W, H, C, N] - int s0, // stride dimension 0 - int s1, // stride dimension 1 - int p0, // padding dimension 0 - int p1, // padding dimension 1 - int d0, // dilation dimension 0 - int d1); // dilation dimension 1 - - GGML_API struct ggml_tensor * ggml_conv_3d_direct( - struct ggml_context * ctx, - struct ggml_tensor * a, // kernel [KW, KH, KD, IC * OC] - struct ggml_tensor * b, // input [W, H, D, C * N] - int s0, // stride - int s1, - int s2, - int p0, // padding - int p1, - int p2, - int d0, // dilation - int d1, - int d2, - int n_channels, - int n_batch, - int n_channels_out); - - enum ggml_op_pool { - GGML_OP_POOL_MAX, - GGML_OP_POOL_AVG, - GGML_OP_POOL_COUNT, - }; - - GGML_API struct ggml_tensor * ggml_pool_1d( - struct ggml_context * ctx, - struct ggml_tensor * a, - enum ggml_op_pool op, - int k0, // kernel size - int s0, // stride - int p0); // padding - - // the result will have 2*p0 padding for the first dimension - // and 2*p1 padding for the second dimension - GGML_API struct ggml_tensor * ggml_pool_2d( - struct ggml_context * ctx, - struct ggml_tensor * a, - enum ggml_op_pool op, - int k0, - int k1, - int s0, - int s1, - float p0, - float p1); - - GGML_API struct ggml_tensor * ggml_pool_2d_back( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * af, // "a"/input used in forward pass - enum ggml_op_pool op, - int k0, - int k1, - int s0, - int s1, - float p0, - float p1); - - enum ggml_scale_mode { - GGML_SCALE_MODE_NEAREST = 0, - GGML_SCALE_MODE_BILINEAR = 1, - GGML_SCALE_MODE_BICUBIC = 2, - - GGML_SCALE_MODE_COUNT - }; - - enum ggml_scale_flag { - GGML_SCALE_FLAG_ALIGN_CORNERS = (1 << 8), - GGML_SCALE_FLAG_ANTIALIAS = (1 << 9), - }; - - // interpolate - // multiplies ne0 and ne1 by scale factor - GGML_API struct ggml_tensor * ggml_upscale( - struct ggml_context * ctx, - struct ggml_tensor * a, - int scale_factor, - enum ggml_scale_mode mode); - - // interpolate - // interpolate scale to specified dimensions - GGML_DEPRECATED(GGML_API struct ggml_tensor * ggml_upscale_ext( - struct ggml_context * ctx, - struct ggml_tensor * a, - int ne0, - int ne1, - int ne2, - int ne3, - enum ggml_scale_mode mode), - "use ggml_interpolate instead"); - - // Up- or downsamples the input to the specified size. - // 2D scale modes (eg. bilinear) are applied to the first two dimensions. - GGML_API struct ggml_tensor * ggml_interpolate( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - int64_t ne1, - int64_t ne2, - int64_t ne3, - uint32_t mode); // ggml_scale_mode [ | ggml_scale_flag...] - - // pad each dimension with zeros: [x, ..., x] -> [x, ..., x, 0, ..., 0] - GGML_API struct ggml_tensor * ggml_pad( - struct ggml_context * ctx, - struct ggml_tensor * a, - int p0, - int p1, - int p2, - int p3); - - // pad each dimension with values on the other side of the torus (looping around) - GGML_API struct ggml_tensor * ggml_pad_circular( - struct ggml_context * ctx, - struct ggml_tensor * a, - int p0, - int p1, - int p2, - int p3); - - GGML_API struct ggml_tensor * ggml_pad_ext( - struct ggml_context * ctx, - struct ggml_tensor * a, - int lp0, - int rp0, - int lp1, - int rp1, - int lp2, - int rp2, - int lp3, - int rp3 - ); - - // pad each dimension with values on the other side of the torus (looping around) - GGML_API struct ggml_tensor * ggml_pad_ext_circular( - struct ggml_context * ctx, - struct ggml_tensor * a, - int lp0, - int rp0, - int lp1, - int rp1, - int lp2, - int rp2, - int lp3, - int rp3); - - // pad each dimension with reflection: [a, b, c, d] -> [b, a, b, c, d, c] - GGML_API struct ggml_tensor * ggml_pad_reflect_1d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int p0, - int p1); - - // Move tensor elements by an offset given for each dimension. Elements that - // are shifted beyond the last position are wrapped around to the beginning. - GGML_API struct ggml_tensor * ggml_roll( - struct ggml_context * ctx, - struct ggml_tensor * a, - int shift0, - int shift1, - int shift2, - int shift3); - - // Convert matrix into a triangular one (upper, strict upper, lower or strict lower) by writing - // zeroes everywhere outside the masked area - GGML_API struct ggml_tensor * ggml_tri( - struct ggml_context * ctx, - struct ggml_tensor * a, - enum ggml_tri_type type); - - // Fill tensor a with constant c - GGML_API struct ggml_tensor * ggml_fill( - struct ggml_context * ctx, - struct ggml_tensor * a, - float c); - - GGML_API struct ggml_tensor * ggml_fill_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - float c); - - // Ref: https://github.com/CompVis/stable-diffusion/blob/main/ldm/modules/diffusionmodules/util.py#L151 - // timesteps: [N,] - // return: [N, dim] - GGML_API struct ggml_tensor * ggml_timestep_embedding( - struct ggml_context * ctx, - struct ggml_tensor * timesteps, - int dim, - int max_period); - - // sort rows - enum ggml_sort_order { - GGML_SORT_ORDER_ASC, - GGML_SORT_ORDER_DESC, - }; - - GGML_API struct ggml_tensor * ggml_argsort( - struct ggml_context * ctx, - struct ggml_tensor * a, - enum ggml_sort_order order); - - // similar to ggml_top_k but implemented as `argsort` + `view` - GGML_API struct ggml_tensor * ggml_argsort_top_k( - struct ggml_context * ctx, - struct ggml_tensor * a, - int k); - - // top k elements per row - // note: the resulting top k indices are in no particular order - GGML_API struct ggml_tensor * ggml_top_k( - struct ggml_context * ctx, - struct ggml_tensor * a, - int k); - - GGML_API struct ggml_tensor * ggml_arange( - struct ggml_context * ctx, - float start, - float stop, - float step); - - // q: [n_embd_k, n_batch, n_head, ne3 ] - // k: [n_embd_k, n_kv, n_head_kv, ne3 ] - // v: [n_embd_v, n_kv, n_head_kv, ne3 ] !! not transposed !! - // mask: [n_kv, n_batch, ne32, ne33] - // res: [n_embd_v, n_head, n_batch, ne3 ] !! permuted !! - // - // broadcast: - // n_head % n_head_kv == 0 - // n_head % ne32 == 0 - // ne3 % ne33 == 0 - // - GGML_API struct ggml_tensor * ggml_flash_attn_ext( - struct ggml_context * ctx, - struct ggml_tensor * q, - struct ggml_tensor * k, - struct ggml_tensor * v, - struct ggml_tensor * mask, - float scale, - float max_bias, - float logit_softcap); - - GGML_API void ggml_flash_attn_ext_set_prec( - struct ggml_tensor * a, - enum ggml_prec prec); - - GGML_API enum ggml_prec ggml_flash_attn_ext_get_prec( - const struct ggml_tensor * a); - - GGML_API void ggml_flash_attn_ext_add_sinks( - struct ggml_tensor * a, - struct ggml_tensor * sinks); - - // Block-sparse flash attention (pFlash). - // Selects top-scoring K-blocks dynamically based on alpha threshold. - // alpha=1.0 selects all blocks (equivalent to dense attention). - // Q: [D, S, H, B], K: [D, S, Hk, B], V: [D, S, Hk, B] (ggml layout) - // Returns: [D, H, S, B] F32 (same shape as ggml_flash_attn_ext output) - GGML_API struct ggml_tensor * ggml_flash_attn_sparse( - struct ggml_context * ctx, - struct ggml_tensor * q, - struct ggml_tensor * k, - struct ggml_tensor * v, - float scale, - float alpha); - - // TurboQuant FWHT rotation. direction: 0 = forward, 1 = inverse. - // Applies signs1 -> FWHT -> signs2 (forward) or signs2 -> FWHT -> signs1 (inverse). - // Used for KV cache rotation in TurboQuant quantization types (TQ3_0). - GGML_API struct ggml_tensor * ggml_turbo_wht( - struct ggml_context * ctx, - struct ggml_tensor * a, - int direction); - - GGML_API struct ggml_tensor * ggml_moe_fused( - struct ggml_context * ctx, - struct ggml_tensor * input, - struct ggml_tensor * gate_w, - struct ggml_tensor * up_w, - struct ggml_tensor * down_w, - struct ggml_tensor * expert_ids, - struct ggml_tensor * expert_weights, - struct ggml_tensor * sh_gate_w, - struct ggml_tensor * sh_up_w, - struct ggml_tensor * sh_down_w, - struct ggml_tensor * sh_gate_inp_w, - int64_t n_embd, - int64_t ff_dim, - int64_t n_expert_used); - - GGML_API struct ggml_tensor * ggml_laguna_moe_combine( - struct ggml_context * ctx, - struct ggml_tensor * experts, - struct ggml_tensor * expert_weights); - - // TODO: needs to be adapted to ggml_flash_attn_ext - GGML_API struct ggml_tensor * ggml_flash_attn_back( - struct ggml_context * ctx, - struct ggml_tensor * q, - struct ggml_tensor * k, - struct ggml_tensor * v, - struct ggml_tensor * d, - bool masked); - - GGML_API struct ggml_tensor * ggml_ssm_conv( - struct ggml_context * ctx, - struct ggml_tensor * sx, - struct ggml_tensor * c); - - // dflash extension: tree-mode ssm_conv for DDTree-style spec - // decoding. parent_ids is an int32 tensor of shape [n_tokens, n_seqs] - // where entry [t, s] is the new-token index of t's tree parent, or -1 if - // t's parent is "before the block" (walks into the old conv state region). - // The CUDA kernel walks the parent chain K-1 times per token to build the - // correct conv window, so siblings get their own parent lineage instead - // of sharing the DFS-neighbour window. Mirrors SGLang's - // causal_conv1d_triton HAS_EAGLE_TREE_CUSTOM_ATTN_MASK path. - GGML_API struct ggml_tensor * ggml_ssm_conv_tree( - struct ggml_context * ctx, - struct ggml_tensor * sx, - struct ggml_tensor * c, - struct ggml_tensor * parent_ids); - - GGML_API struct ggml_tensor * ggml_ssm_scan( - struct ggml_context * ctx, - struct ggml_tensor * s, - struct ggml_tensor * x, - struct ggml_tensor * dt, - struct ggml_tensor * A, - struct ggml_tensor * B, - struct ggml_tensor * C, - struct ggml_tensor * ids); - - // partition into non-overlapping windows with padding if needed - // example: - // a: 768 64 64 1 - // w: 14 - // res: 768 14 14 25 - // used in sam - GGML_API struct ggml_tensor * ggml_win_part( - struct ggml_context * ctx, - struct ggml_tensor * a, - int w); - - // reverse of ggml_win_part - // used in sam - GGML_API struct ggml_tensor * ggml_win_unpart( - struct ggml_context * ctx, - struct ggml_tensor * a, - int w0, - int h0, - int w); - - GGML_API struct ggml_tensor * ggml_unary( - struct ggml_context * ctx, - struct ggml_tensor * a, - enum ggml_unary_op op); - - GGML_API struct ggml_tensor * ggml_unary_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - enum ggml_unary_op op); - - // used in sam - GGML_API struct ggml_tensor * ggml_get_rel_pos( - struct ggml_context * ctx, - struct ggml_tensor * a, - int qh, - int kh); - - // used in sam - GGML_API struct ggml_tensor * ggml_add_rel_pos( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * pw, - struct ggml_tensor * ph); - - GGML_API struct ggml_tensor * ggml_add_rel_pos_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * pw, - struct ggml_tensor * ph); - - GGML_API struct ggml_tensor * ggml_rwkv_wkv6( - struct ggml_context * ctx, - struct ggml_tensor * k, - struct ggml_tensor * v, - struct ggml_tensor * r, - struct ggml_tensor * tf, - struct ggml_tensor * td, - struct ggml_tensor * state); - - GGML_API struct ggml_tensor * ggml_gated_linear_attn( - struct ggml_context * ctx, - struct ggml_tensor * k, - struct ggml_tensor * v, - struct ggml_tensor * q, - struct ggml_tensor * g, - struct ggml_tensor * state, - float scale); - - GGML_API struct ggml_tensor * ggml_rwkv_wkv7( - struct ggml_context * ctx, - struct ggml_tensor * r, - struct ggml_tensor * w, - struct ggml_tensor * k, - struct ggml_tensor * v, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * state); - - /* Solves a specific equation of the form Ax=B, where A is a triangular matrix - * without zeroes on the diagonal (i.e. invertible). - * B can have any number of columns, but must have the same number of rows as A - * If A is [n, n] and B is [n, m], then the result will be [n, m] as well - * Has O(n^3) complexity (unlike most matrix ops out there), so use on cases - * where n > 100 sparingly, pre-chunk if necessary. - * - * If left = false, solves xA=B instead - * If lower = false, assumes upper triangular instead - * If uni = true, assumes diagonal of A to be all ones (will override actual values) - * - * TODO: currently only lower, right, non-unitriangular variant is implemented - */ - GGML_API struct ggml_tensor * ggml_solve_tri( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - bool left, - bool lower, - bool uni); - - // TODO: add ggml_gated_delta_net_set_bcast() to be able to configure Q, K broadcast type: tiled vs interleaved [TAG_GGML_GDN_BCAST] - // ref: https://github.com/ggml-org/llama.cpp/pull/19468#discussion_r2786394306 - GGML_API struct ggml_tensor * ggml_gated_delta_net( - struct ggml_context * ctx, - struct ggml_tensor * q, - struct ggml_tensor * k, - struct ggml_tensor * v, - struct ggml_tensor * g, - struct ggml_tensor * beta, - struct ggml_tensor * state); - - GGML_API void ggml_gated_delta_net_set_skip_intermediate( - struct ggml_tensor * tensor, - bool skip_intermediate); - - // dflash extension: tree-mode gated delta net for DDTree-style - // speculative decoding verify. `parent_ids` is an int32 tensor of shape - // [n_tokens, n_seqs] where entry [t, s] is the index within sequence s of - // the parent token in the DFS-flattened tree (or -1 for a root-level - // node). At each token step t > 0 in the recurrence, if parent_ids[t] is - // not (t - 1), the kernel reloads the recurrent state from the - // intermediate-states region at parent_ids[t] instead of continuing - // sequentially. This lets one verify forward pass correctly process - // multiple tree branches (siblings) without cross-contamination. - GGML_API struct ggml_tensor * ggml_gated_delta_net_tree( - struct ggml_context * ctx, - struct ggml_tensor * q, - struct ggml_tensor * k, - struct ggml_tensor * v, - struct ggml_tensor * g, - struct ggml_tensor * beta, - struct ggml_tensor * state, - struct ggml_tensor * parent_ids); - - // dflash extension: tree-mode with direct intermediate-state writes to a - // persistent external buffer. Identical to ggml_gated_delta_net_tree but - // the per-token intermediate states are written to `persist_inter->data` - // (f32 or f16, [S_v, S_v, H, n_tokens, n_seqs], contiguous) instead of - // the default internal region of the result tensor. Eliminates a - // downstream ggml_cpy into the persistent cache buffer, saving ~5-10 ms - // per verify step on a 27B hybrid target. - GGML_API struct ggml_tensor * ggml_gated_delta_net_tree_persist( - struct ggml_context * ctx, - struct ggml_tensor * q, - struct ggml_tensor * k, - struct ggml_tensor * v, - struct ggml_tensor * g, - struct ggml_tensor * beta, - struct ggml_tensor * state, - struct ggml_tensor * parent_ids, - struct ggml_tensor * persist_inter); - - // custom operators - - typedef void (*ggml_custom1_op_t)(struct ggml_tensor * dst , const struct ggml_tensor * a, int ith, int nth, void * userdata); - typedef void (*ggml_custom2_op_t)(struct ggml_tensor * dst , const struct ggml_tensor * a, const struct ggml_tensor * b, int ith, int nth, void * userdata); - typedef void (*ggml_custom3_op_t)(struct ggml_tensor * dst , const struct ggml_tensor * a, const struct ggml_tensor * b, const struct ggml_tensor * c, int ith, int nth, void * userdata); - -#define GGML_N_TASKS_MAX (-1) - // n_tasks == GGML_N_TASKS_MAX means to use max number of tasks - - GGML_API struct ggml_tensor * ggml_map_custom1( - struct ggml_context * ctx, - struct ggml_tensor * a, - ggml_custom1_op_t fun, - int n_tasks, - void * userdata); - - GGML_API struct ggml_tensor * ggml_map_custom1_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - ggml_custom1_op_t fun, - int n_tasks, - void * userdata); - - GGML_API struct ggml_tensor * ggml_map_custom2( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - ggml_custom2_op_t fun, - int n_tasks, - void * userdata); - - GGML_API struct ggml_tensor * ggml_map_custom2_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - ggml_custom2_op_t fun, - int n_tasks, - void * userdata); - - GGML_API struct ggml_tensor * ggml_map_custom3( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * c, - ggml_custom3_op_t fun, - int n_tasks, - void * userdata); - - GGML_API struct ggml_tensor * ggml_map_custom3_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * c, - ggml_custom3_op_t fun, - int n_tasks, - void * userdata); - - typedef void (*ggml_custom_op_t)(struct ggml_tensor * dst , int ith, int nth, void * userdata); - - GGML_API struct ggml_tensor * ggml_custom_4d( - struct ggml_context * ctx, - enum ggml_type type, - int64_t ne0, - int64_t ne1, - int64_t ne2, - int64_t ne3, - struct ggml_tensor ** args, - int n_args, - ggml_custom_op_t fun, - int n_tasks, - void * userdata); - - GGML_API struct ggml_tensor * ggml_custom_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor ** args, - int n_args, - ggml_custom_op_t fun, - int n_tasks, - void * userdata); - - // loss function - - GGML_API struct ggml_tensor * ggml_cross_entropy_loss( - struct ggml_context * ctx, - struct ggml_tensor * a, // logits - struct ggml_tensor * b); // labels - - GGML_API struct ggml_tensor * ggml_cross_entropy_loss_back( - struct ggml_context * ctx, - struct ggml_tensor * a, // logits - struct ggml_tensor * b, // labels - struct ggml_tensor * c); // gradients of cross_entropy_loss result - - // AdamW optimizer step - // Paper: https://arxiv.org/pdf/1711.05101v3.pdf - // PyTorch: https://pytorch.org/docs/stable/generated/torch.optim.AdamW.html - GGML_API struct ggml_tensor * ggml_opt_step_adamw( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * grad, - struct ggml_tensor * m, - struct ggml_tensor * v, - struct ggml_tensor * adamw_params); // parameters such as the learning rate - - // stochastic gradient descent step (with weight decay) - GGML_API struct ggml_tensor * ggml_opt_step_sgd( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * grad, - struct ggml_tensor * sgd_params); // alpha, weight decay - - // build forward multiple tensors and select one of them for computing - // this is useful for creating graphs that have constant topology but compute different things based on the input - // ref: https://github.com/ggml-org/llama.cpp/pull/18550 - // - // nodes: - // | - build forward into the graph but do not compute - // c - build forward into the graph and compute - // - // | | ... c ... | - // | | ... c ... | - // | | ... c ... | - // [0 1 ... idx ... n-1] <-- ggml_build_forward_select(..., n, idx) - // c - // c - // - // example: - // struct ggml_tensor * curs[3]; - // - // curs[0] = compute0(...); - // curs[1] = compute1(...); - // curs[2] = compute2(...); - // - // int idx = select_branch(some_input); - // - // struct ggml_tensor * out = ggml_build_forward_select(cgraph, curs, 3, idx); - // - GGML_API struct ggml_tensor * ggml_build_forward_select( - struct ggml_cgraph * cgraph, - struct ggml_tensor ** tensors, - int n_tensors, - int idx); - - GGML_API void ggml_build_forward_expand( - struct ggml_cgraph * cgraph, - struct ggml_tensor * tensor); - - GGML_API void ggml_build_backward_expand( - struct ggml_context * ctx, // context for gradient computation - struct ggml_cgraph * cgraph, - struct ggml_tensor ** grad_accs); - - // graph allocation in a context - GGML_API struct ggml_cgraph * ggml_new_graph (struct ggml_context * ctx); // size = GGML_DEFAULT_GRAPH_SIZE, grads = false - GGML_API struct ggml_cgraph * ggml_new_graph_custom(struct ggml_context * ctx, size_t size, bool grads); - GGML_API struct ggml_cgraph * ggml_graph_dup (struct ggml_context * ctx, struct ggml_cgraph * cgraph, bool force_grads); - GGML_API void ggml_graph_cpy (struct ggml_cgraph * src, struct ggml_cgraph * dst); - GGML_API void ggml_graph_reset (struct ggml_cgraph * cgraph); // set regular grads + optimizer momenta to 0, set loss grad to 1 - GGML_API void ggml_graph_clear (struct ggml_cgraph * cgraph); - - GGML_API int ggml_graph_size (struct ggml_cgraph * cgraph); - GGML_API struct ggml_tensor * ggml_graph_node (struct ggml_cgraph * cgraph, int i); // if i < 0, returns nodes[n_nodes + i] - GGML_API struct ggml_tensor ** ggml_graph_nodes (struct ggml_cgraph * cgraph); - GGML_API int ggml_graph_n_nodes(struct ggml_cgraph * cgraph); - - GGML_API void ggml_graph_add_node(struct ggml_cgraph * cgraph, struct ggml_tensor * tensor); - - GGML_API size_t ggml_graph_overhead(void); - GGML_API size_t ggml_graph_overhead_custom(size_t size, bool grads); - - GGML_API struct ggml_tensor * ggml_graph_get_tensor (const struct ggml_cgraph * cgraph, const char * name); - GGML_API struct ggml_tensor * ggml_graph_get_grad (const struct ggml_cgraph * cgraph, const struct ggml_tensor * node); - GGML_API struct ggml_tensor * ggml_graph_get_grad_acc(const struct ggml_cgraph * cgraph, const struct ggml_tensor * node); - - // print info and performance information for the graph - GGML_API void ggml_graph_print(const struct ggml_cgraph * cgraph); - - // dump the graph into a file using the dot format - GGML_API void ggml_graph_dump_dot(const struct ggml_cgraph * gb, const struct ggml_cgraph * cgraph, const char * filename); - - // TODO these functions were sandwiched in the old optimization interface, is there a better place for them? - typedef void (*ggml_log_callback)(enum ggml_log_level level, const char * text, void * user_data); - - // Set callback for all future logging events. - // If this is not called, or NULL is supplied, everything is output on stderr. - GGML_API void ggml_log_get(ggml_log_callback * log_callback, void ** user_data); - GGML_API void ggml_log_set(ggml_log_callback log_callback, void * user_data); - - GGML_API struct ggml_tensor * ggml_set_zero(struct ggml_tensor * tensor); - - // - // quantization - // - - // - ggml_quantize_init can be called multiple times with the same type - // it will only initialize the quantization tables for the first call or after ggml_quantize_free - // automatically called by ggml_quantize_chunk for convenience - // - // - ggml_quantize_free will free any memory allocated by ggml_quantize_init - // call this at the end of the program to avoid memory leaks - // - // note: these are thread-safe - // - GGML_API void ggml_quantize_init(enum ggml_type type); - GGML_API void ggml_quantize_free(void); - - // some quantization type cannot be used without an importance matrix - GGML_API bool ggml_quantize_requires_imatrix(enum ggml_type type); - - // calls ggml_quantize_init internally (i.e. can allocate memory) - GGML_API size_t ggml_quantize_chunk( - enum ggml_type type, - const float * src, - void * dst, - int64_t start, - int64_t nrows, - int64_t n_per_row, - const float * imatrix); - -#ifdef __cplusplus - // restrict not standard in C++ -# if defined(__GNUC__) -# define GGML_RESTRICT __restrict__ -# elif defined(__clang__) -# define GGML_RESTRICT __restrict -# elif defined(_MSC_VER) -# define GGML_RESTRICT __restrict -# else -# define GGML_RESTRICT -# endif -#else -# if defined (_MSC_VER) && (__STDC_VERSION__ < 201112L) -# define GGML_RESTRICT __restrict -# else -# define GGML_RESTRICT restrict -# endif -#endif - typedef void (*ggml_to_float_t) (const void * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); - typedef void (*ggml_from_float_t)(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); - - struct ggml_type_traits { - const char * type_name; - int64_t blck_size; - int64_t blck_size_interleave; // interleave elements in blocks - size_t type_size; - bool is_quantized; - ggml_to_float_t to_float; - ggml_from_float_t from_float_ref; - }; - - GGML_API const struct ggml_type_traits * ggml_get_type_traits(enum ggml_type type); - - // ggml threadpool - // TODO: currently, only a few functions are in the base ggml API, while the rest are in the CPU backend - // the goal should be to create an API that other backends can use move everything to the ggml base - - // scheduling priorities - enum ggml_sched_priority { - GGML_SCHED_PRIO_LOW = -1, - GGML_SCHED_PRIO_NORMAL, - GGML_SCHED_PRIO_MEDIUM, - GGML_SCHED_PRIO_HIGH, - GGML_SCHED_PRIO_REALTIME - }; - - // threadpool params - // Use ggml_threadpool_params_default() or ggml_threadpool_params_init() to populate the defaults - struct ggml_threadpool_params { - bool cpumask[GGML_MAX_N_THREADS]; // mask of cpu cores (all-zeros means use default affinity settings) - int n_threads; // number of threads - enum ggml_sched_priority prio; // thread priority - uint32_t poll; // polling level (0 - no polling, 100 - aggressive polling) - bool strict_cpu; // strict cpu placement - bool paused; // start in paused state - }; - - struct ggml_threadpool; // forward declaration, see ggml.c - - typedef struct ggml_threadpool * ggml_threadpool_t; - - GGML_API struct ggml_threadpool_params ggml_threadpool_params_default(int n_threads); - GGML_API void ggml_threadpool_params_init (struct ggml_threadpool_params * p, int n_threads); - GGML_API bool ggml_threadpool_params_match (const struct ggml_threadpool_params * p0, const struct ggml_threadpool_params * p1); - -#ifdef __cplusplus -} -#endif diff --git a/ggml/src/ggml-cpu/ggml-cpu.c.ds4bak b/ggml/src/ggml-cpu/ggml-cpu.c.ds4bak deleted file mode 100644 index f0e14a00673c..000000000000 --- a/ggml/src/ggml-cpu/ggml-cpu.c.ds4bak +++ /dev/null @@ -1,3782 +0,0 @@ -#define _CRT_SECURE_NO_DEPRECATE // Disables "unsafe" warnings on Windows -#define _USE_MATH_DEFINES // For M_PI on MSVC - -#include "ggml-backend-impl.h" -#include "ggml-backend.h" -#include "traits.h" -#include "ggml-cpu-impl.h" -#include "ggml-impl.h" -#include "quants.h" -#include "ggml-threading.h" -#include "unary-ops.h" -#include "binary-ops.h" -#include "vec.h" -#include "ops.h" -#include "ggml.h" -#include "common.h" - -#if defined(_MSC_VER) || defined(__MINGW32__) -#include // using malloc.h with MSC/MINGW -#elif !defined(__FreeBSD__) && !defined(__NetBSD__) && !defined(__OpenBSD__) -#include -#endif - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#if defined(__gnu_linux__) -#include -#endif - -#ifdef GGML_USE_OPENMP -#include -#endif - -#if defined(__ARM_FEATURE_SVE) || defined(__ARM_FEATURE_MATMUL_INT8) -#undef GGML_USE_LLAMAFILE -#endif - -#ifdef GGML_USE_LLAMAFILE -#include "llamafile/sgemm.h" -#endif - -// Note: once we move threading into a separate C++ file -// will use std::hardware_destructive_interference_size instead of hardcoding it here -// and we'll use C++ attribute syntax. -#define GGML_CACHE_LINE 64 - -#if defined(__clang__) || defined(__GNUC__) -#define GGML_CACHE_ALIGN __attribute__((aligned(GGML_CACHE_LINE))) -#endif - -#if defined(__has_feature) -#if __has_feature(thread_sanitizer) -#define GGML_TSAN_ENABLED 1 -#endif -#else // __has_feature -#if defined(__SANITIZE_THREAD__) -#define GGML_TSAN_ENABLED 1 -#endif -#endif // __has_feature - -#define UNUSED GGML_UNUSED -#define SWAP(x, y, T) do { T SWAP = x; (x) = y; (y) = SWAP; } while (0) - -// precomputed f32 table for f16 (256 KB) (simd-mappings.h) -float ggml_table_f32_f16[1 << 16]; - -// precomputed f32 table for e8m0 half (1 KB) (simd-mappings.h) -float ggml_table_f32_e8m0_half[1 << 8]; - -#if defined(__ARM_ARCH) -struct ggml_arm_arch_features_type { - int sve_cnt; -} ggml_arm_arch_features = { 0 }; -#endif - -#if defined(__riscv) -struct ggml_riscv_arch_features_type { - int rvv_vlen; -} ggml_riscv_arch_features = { 0 }; -#endif - -#if defined(_WIN32) - -#define WIN32_LEAN_AND_MEAN -#ifndef NOMINMAX - #define NOMINMAX -#endif -#include - -#if defined(_MSC_VER) && !defined(__clang__) -#define GGML_CACHE_ALIGN __declspec(align(GGML_CACHE_LINE)) - -typedef volatile LONG atomic_int; -typedef atomic_int atomic_bool; -typedef atomic_int atomic_flag; - -#define ATOMIC_FLAG_INIT 0 - -typedef enum { - memory_order_relaxed, - memory_order_consume, - memory_order_acquire, - memory_order_release, - memory_order_acq_rel, - memory_order_seq_cst -} memory_order; - -static void atomic_store(atomic_int * ptr, LONG val) { - InterlockedExchange(ptr, val); -} -static void atomic_store_explicit(atomic_int * ptr, LONG val, memory_order mo) { - // TODO: add support for explicit memory order - InterlockedExchange(ptr, val); -} -static LONG atomic_load(atomic_int * ptr) { - return InterlockedCompareExchange(ptr, 0, 0); -} -static LONG atomic_load_explicit(atomic_int * ptr, memory_order mo) { - // TODO: add support for explicit memory order - return InterlockedCompareExchange(ptr, 0, 0); -} -static LONG atomic_fetch_add(atomic_int * ptr, LONG inc) { - return InterlockedExchangeAdd(ptr, inc); -} -static LONG atomic_fetch_add_explicit(atomic_int * ptr, LONG inc, memory_order mo) { - // TODO: add support for explicit memory order - return InterlockedExchangeAdd(ptr, inc); -} -static atomic_bool atomic_flag_test_and_set(atomic_flag * ptr) { - return InterlockedExchange(ptr, 1); -} -static void atomic_flag_clear(atomic_flag * ptr) { - InterlockedExchange(ptr, 0); -} -static void atomic_thread_fence(memory_order mo) { - MemoryBarrier(); -} -#else // clang -#include -#endif - -typedef HANDLE pthread_t; - -typedef DWORD thread_ret_t; -static int pthread_create(pthread_t * out, void * unused, thread_ret_t(*func)(void *), void * arg) { - (void) unused; - HANDLE handle = CreateThread(NULL, 0, (LPTHREAD_START_ROUTINE) func, arg, 0, NULL); - if (handle == NULL) - { - return EAGAIN; - } - - *out = handle; - return 0; -} - -static int pthread_join(pthread_t thread, void * unused) { - (void) unused; - int ret = (int) WaitForSingleObject(thread, INFINITE); - CloseHandle(thread); - return ret; -} - -static int sched_yield (void) { - Sleep (0); - return 0; -} -#else - -#include -#include -#include -#if defined(__FreeBSD__) -#include -#endif - -typedef void * thread_ret_t; - -#include -#include -#include - -#endif - -typedef pthread_t ggml_thread_t; - -#define GGML_THREADPOOL_N_THREADS_MASK (0xffffU) -#define GGML_THREADPOOL_N_THREADS_BITS (16) - -#if defined(__APPLE__) -#include -#include -#include -#endif - -static const struct ggml_type_traits_cpu type_traits_cpu[GGML_TYPE_COUNT] = { - [GGML_TYPE_F32] = { - .from_float = (ggml_from_float_t) ggml_cpu_fp32_to_fp32, - .vec_dot = (ggml_vec_dot_t) ggml_vec_dot_f32, - .vec_dot_type = GGML_TYPE_F32, - .nrows = 1, - }, - [GGML_TYPE_F16] = { - .from_float = (ggml_from_float_t) ggml_cpu_fp32_to_fp16, - .vec_dot = (ggml_vec_dot_t) ggml_vec_dot_f16, - .vec_dot_type = GGML_TYPE_F16, - .nrows = 1, - }, - [GGML_TYPE_Q1_0] = { - .from_float = quantize_row_q1_0, - .vec_dot = ggml_vec_dot_q1_0_q8_0, - .vec_dot_type = GGML_TYPE_Q8_0, - .nrows = 1, - }, - [GGML_TYPE_Q4_0] = { - .from_float = quantize_row_q4_0, - .vec_dot = ggml_vec_dot_q4_0_q8_0, - .vec_dot_type = GGML_TYPE_Q8_0, -#if defined (__ARM_FEATURE_MATMUL_INT8) - .nrows = 2, -#else - .nrows = 1, -#endif - }, - [GGML_TYPE_Q4_1] = { - .from_float = quantize_row_q4_1, - .vec_dot = ggml_vec_dot_q4_1_q8_1, - .vec_dot_type = GGML_TYPE_Q8_1, -#if defined (__ARM_FEATURE_MATMUL_INT8) - .nrows = 2, -#else - .nrows = 1, -#endif - }, - [GGML_TYPE_Q5_0] = { - .from_float = quantize_row_q5_0, - .vec_dot = ggml_vec_dot_q5_0_q8_0, - .vec_dot_type = GGML_TYPE_Q8_0, - .nrows = 1, - }, - [GGML_TYPE_Q5_1] = { - .from_float = quantize_row_q5_1, - .vec_dot = ggml_vec_dot_q5_1_q8_1, - .vec_dot_type = GGML_TYPE_Q8_1, - .nrows = 1, - }, - [GGML_TYPE_Q8_0] = { - .from_float = quantize_row_q8_0, - .vec_dot = ggml_vec_dot_q8_0_q8_0, - .vec_dot_type = GGML_TYPE_Q8_0, -#if defined (__ARM_FEATURE_MATMUL_INT8) - .nrows = 2, -#else - .nrows = 1, -#endif - }, - [GGML_TYPE_Q8_1] = { - .from_float = quantize_row_q8_1, - .vec_dot_type = GGML_TYPE_Q8_1, - .nrows = 1, - }, - [GGML_TYPE_MXFP4] = { - .from_float = quantize_row_mxfp4, - .vec_dot = ggml_vec_dot_mxfp4_q8_0, - .vec_dot_type = GGML_TYPE_Q8_0, - .nrows = 1, - }, - [GGML_TYPE_NVFP4] = { - .from_float = quantize_row_nvfp4, - .vec_dot = ggml_vec_dot_nvfp4_q8_0, - .vec_dot_type = GGML_TYPE_Q8_0, - .nrows = 1, - }, - [GGML_TYPE_Q2_K] = { - .from_float = quantize_row_q2_K, - .vec_dot = ggml_vec_dot_q2_K_q8_K, - .vec_dot_type = GGML_TYPE_Q8_K, - .nrows = 1, - }, - [GGML_TYPE_Q3_K] = { - .from_float = quantize_row_q3_K, - .vec_dot = ggml_vec_dot_q3_K_q8_K, - .vec_dot_type = GGML_TYPE_Q8_K, - .nrows = 1, - }, - [GGML_TYPE_Q4_K] = { - .from_float = quantize_row_q4_K, - .vec_dot = ggml_vec_dot_q4_K_q8_K, - .vec_dot_type = GGML_TYPE_Q8_K, -#if defined (__ARM_FEATURE_MATMUL_INT8) - .nrows = 2, -#else - .nrows = 1, -#endif - }, - [GGML_TYPE_Q5_K] = { - .from_float = quantize_row_q5_K, - .vec_dot = ggml_vec_dot_q5_K_q8_K, - .vec_dot_type = GGML_TYPE_Q8_K, - .nrows = 1, - }, - [GGML_TYPE_Q6_K] = { - .from_float = quantize_row_q6_K, - .vec_dot = ggml_vec_dot_q6_K_q8_K, - .vec_dot_type = GGML_TYPE_Q8_K, -#if defined (__ARM_FEATURE_MATMUL_INT8) - .nrows = 2, -#else - .nrows = 1, -#endif - }, - [GGML_TYPE_IQ2_XXS] = { - .from_float = NULL, - .vec_dot = ggml_vec_dot_iq2_xxs_q8_K, - .vec_dot_type = GGML_TYPE_Q8_K, - .nrows = 1, - }, - [GGML_TYPE_IQ2_XS] = { - .from_float = NULL, - .vec_dot = ggml_vec_dot_iq2_xs_q8_K, - .vec_dot_type = GGML_TYPE_Q8_K, - .nrows = 1, - }, - [GGML_TYPE_IQ3_XXS] = { - // NOTE: from_float for iq3 and iq2_s was removed because these quants require initialization in ggml_quantize_init - //.from_float = quantize_row_iq3_xxs, - .vec_dot = ggml_vec_dot_iq3_xxs_q8_K, - .vec_dot_type = GGML_TYPE_Q8_K, - .nrows = 1, - }, - [GGML_TYPE_IQ3_S] = { - //.from_float = quantize_row_iq3_s, - .vec_dot = ggml_vec_dot_iq3_s_q8_K, - .vec_dot_type = GGML_TYPE_Q8_K, - .nrows = 1, - }, - [GGML_TYPE_IQ2_S] = { - //.from_float = quantize_row_iq2_s, - .vec_dot = ggml_vec_dot_iq2_s_q8_K, - .vec_dot_type = GGML_TYPE_Q8_K, - .nrows = 1, - }, - [GGML_TYPE_IQ1_S] = { - .from_float = NULL, - .vec_dot = ggml_vec_dot_iq1_s_q8_K, - .vec_dot_type = GGML_TYPE_Q8_K, - .nrows = 1, - }, - [GGML_TYPE_IQ1_M] = { - .from_float = NULL, - .vec_dot = ggml_vec_dot_iq1_m_q8_K, - .vec_dot_type = GGML_TYPE_Q8_K, - .nrows = 1, - }, - [GGML_TYPE_IQ4_NL] = { - .from_float = quantize_row_iq4_nl, - .vec_dot = ggml_vec_dot_iq4_nl_q8_0, - .vec_dot_type = GGML_TYPE_Q8_0, - .nrows = 1, - }, - [GGML_TYPE_IQ4_XS] = { - .from_float = quantize_row_iq4_xs, - .vec_dot = ggml_vec_dot_iq4_xs_q8_K, - .vec_dot_type = GGML_TYPE_Q8_K, - .nrows = 1, - }, - [GGML_TYPE_Q8_K] = { - .from_float = quantize_row_q8_K, - }, - [GGML_TYPE_BF16] = { - .from_float = (ggml_from_float_t) ggml_cpu_fp32_to_bf16, - .vec_dot = (ggml_vec_dot_t) ggml_vec_dot_bf16, - .vec_dot_type = GGML_TYPE_BF16, - .nrows = 1, - }, - [GGML_TYPE_TQ1_0] = { - .from_float = quantize_row_tq1_0, - .vec_dot = ggml_vec_dot_tq1_0_q8_K, - .vec_dot_type = GGML_TYPE_Q8_K, - .nrows = 1, - }, - [GGML_TYPE_TQ2_0] = { - .from_float = quantize_row_tq2_0, - .vec_dot = ggml_vec_dot_tq2_0_q8_K, - .vec_dot_type = GGML_TYPE_Q8_K, - .nrows = 1, - }, - [GGML_TYPE_I32] = { - .from_float = (ggml_from_float_t) ggml_cpu_fp32_to_i32, - }, -}; - -const struct ggml_type_traits_cpu * ggml_get_type_traits_cpu(enum ggml_type type) { - return &type_traits_cpu[type]; -} - -// -// Threading defs -// - -typedef pthread_t ggml_thread_t; - -#if defined(_WIN32) - -typedef CONDITION_VARIABLE ggml_cond_t; -typedef SRWLOCK ggml_mutex_t; - -#define ggml_mutex_init(m) InitializeSRWLock(m) -#define ggml_mutex_destroy(m) -#define ggml_mutex_lock(m) AcquireSRWLockExclusive(m) -#define ggml_mutex_unlock(m) ReleaseSRWLockExclusive(m) -#define ggml_mutex_lock_shared(m) AcquireSRWLockShared(m) -#define ggml_mutex_unlock_shared(m) ReleaseSRWLockShared(m) - -#define ggml_cond_init(c) InitializeConditionVariable(c) -#define ggml_cond_destroy(c) -#define ggml_cond_wait(c, m) SleepConditionVariableSRW(c, m, INFINITE, CONDITION_VARIABLE_LOCKMODE_SHARED) -#define ggml_cond_broadcast(c) WakeAllConditionVariable(c) - -#define ggml_thread_create pthread_create -#define ggml_thread_join pthread_join - -#else - -typedef pthread_cond_t ggml_cond_t; -typedef pthread_mutex_t ggml_mutex_t; - -#define ggml_mutex_init(m) pthread_mutex_init(m, NULL) -#define ggml_mutex_destroy(m) pthread_mutex_destroy(m) -#define ggml_mutex_lock(m) pthread_mutex_lock(m) -#define ggml_mutex_unlock(m) pthread_mutex_unlock(m) -#define ggml_mutex_lock_shared(m) pthread_mutex_lock(m) -#define ggml_mutex_unlock_shared(m) pthread_mutex_unlock(m) - -#define ggml_lock_init(x) UNUSED(x) -#define ggml_lock_destroy(x) UNUSED(x) -#if defined(__x86_64__) || (defined(_MSC_VER) && defined(_M_AMD64)) -#define ggml_lock_lock(x) _mm_pause() -#else -#define ggml_lock_lock(x) UNUSED(x) -#endif -#define ggml_lock_unlock(x) UNUSED(x) - -#define GGML_LOCK_INITIALIZER 0 -#define ggml_cond_init(c) pthread_cond_init(c, NULL) -#define ggml_cond_destroy(c) pthread_cond_destroy(c) -#define ggml_cond_wait(c, m) pthread_cond_wait(c, m) -#define ggml_cond_broadcast(c) pthread_cond_broadcast(c) - -#define ggml_thread_create pthread_create -#define ggml_thread_join pthread_join - -#endif - -// Threadpool def -struct ggml_threadpool { - ggml_mutex_t mutex; // mutex for cond.var - ggml_cond_t cond; // cond.var for waiting for new work - - struct ggml_cgraph * cgraph; - struct ggml_cplan * cplan; - - // synchronization primitives - atomic_int n_graph; // updated when there is work to be done (i.e each graph) holds graph and active thread counts. - atomic_int GGML_CACHE_ALIGN n_barrier; - atomic_int GGML_CACHE_ALIGN n_barrier_passed; - atomic_int GGML_CACHE_ALIGN current_chunk; // currently processing chunk during Mat_Mul, shared between all the threads. - - // these are atomic as an annotation for thread-sanitizer - atomic_bool stop; // Used for stopping the threadpool altogether - atomic_bool pause; // Used for pausing the threadpool or individual threads - atomic_int abort; // Used for aborting processing of a graph - - struct ggml_compute_state * workers; // per thread state - int n_threads; // Number of threads in the pool - int32_t prio; // Scheduling priority - uint32_t poll; // Polling level (0 - no polling) - - enum ggml_status ec; -}; - -// Per-thread state -struct ggml_compute_state { -#ifndef GGML_USE_OPENMP - ggml_thread_t thrd; - int last_graph; - bool pending; -#endif - bool cpumask[GGML_MAX_N_THREADS]; - struct ggml_threadpool * threadpool; - int ith; -}; - -// Helpers for polling loops -#if defined(__aarch64__) && ( defined(__clang__) || defined(__GNUC__) ) -static inline void ggml_thread_cpu_relax(void) { - __asm__ volatile("yield" ::: "memory"); -} -#elif defined(__x86_64__) -static inline void ggml_thread_cpu_relax(void) { - _mm_pause(); -} -#elif defined(__riscv) -static inline void ggml_thread_cpu_relax(void) { - #ifdef __riscv_zihintpause - __asm__ __volatile__ ("pause"); - #else - /* Encoding of the pause instruction */ - __asm__ __volatile__ (".4byte 0x100000F"); - #endif -} -#else -static inline void ggml_thread_cpu_relax(void) {;} -#endif - -// -// NUMA support -// - -#define GGML_NUMA_MAX_NODES 8 -#define GGML_NUMA_MAX_CPUS 512 - -struct ggml_numa_node { - uint32_t cpus[GGML_NUMA_MAX_CPUS]; // hardware threads on this node - uint32_t n_cpus; -}; - -struct ggml_numa_nodes { - enum ggml_numa_strategy numa_strategy; - struct ggml_numa_node nodes[GGML_NUMA_MAX_NODES]; - uint32_t n_nodes; - uint32_t total_cpus; // hardware threads on system - uint32_t current_node; // node on which main process is execting -#if defined(__gnu_linux__) - cpu_set_t cpuset; // cpuset from numactl -#else - uint32_t cpuset; // no NUMA support outside of Linux at this time. Use a portable datatype -#endif -}; - -// -// ggml state -// - -struct ggml_state { - struct ggml_numa_nodes numa; -}; - -static struct ggml_state g_state = {0}; - -void ggml_barrier(struct ggml_threadpool * tp) { - int n_threads = atomic_load_explicit(&tp->n_graph, memory_order_relaxed) & GGML_THREADPOOL_N_THREADS_MASK; - if (n_threads == 1) { - return; - } - -#ifdef GGML_USE_OPENMP - #pragma omp barrier -#else - int n_passed = atomic_load_explicit(&tp->n_barrier_passed, memory_order_relaxed); - - // enter barrier (full seq-cst fence) - int n_barrier = atomic_fetch_add_explicit(&tp->n_barrier, 1, memory_order_seq_cst); - - if (n_barrier == (n_threads - 1)) { - // last thread - atomic_store_explicit(&tp->n_barrier, 0, memory_order_relaxed); - - // exit barrier (full seq-cst fence) - atomic_fetch_add_explicit(&tp->n_barrier_passed, 1, memory_order_seq_cst); - return; - } - - // wait for other threads - while (atomic_load_explicit(&tp->n_barrier_passed, memory_order_relaxed) == n_passed) { - ggml_thread_cpu_relax(); - } - - // exit barrier (full seq-cst fence) - // TSAN doesn't support standalone fence yet, we use a dummy read-modify-write instead - #ifdef GGML_TSAN_ENABLED - atomic_fetch_add_explicit(&tp->n_barrier_passed, 0, memory_order_seq_cst); - #else - atomic_thread_fence(memory_order_seq_cst); - #endif -#endif -} - -void ggml_threadpool_chunk_set(struct ggml_threadpool * tp, int value) { - atomic_store_explicit(&tp->current_chunk, value, memory_order_relaxed); -} - -int ggml_threadpool_chunk_add(struct ggml_threadpool * tp, int value) { - return atomic_fetch_add_explicit(&tp->current_chunk, value, memory_order_relaxed); -} - -#if defined(__gnu_linux__) -static cpu_set_t ggml_get_numa_affinity(void) { - cpu_set_t cpuset; - pthread_t thread; - thread = pthread_self(); - CPU_ZERO(&cpuset); - pthread_getaffinity_np(thread, sizeof(cpu_set_t), &cpuset); - return cpuset; -} -#else -static uint32_t ggml_get_numa_affinity(void) { - return 0; // no NUMA support -} -#endif - -void ggml_numa_init(enum ggml_numa_strategy numa_flag) { - if (g_state.numa.n_nodes > 0) { - fprintf(stderr, "ggml_numa_init: NUMA already initialized\n"); - - return; - } - -#if defined(__gnu_linux__) - struct stat st; - char path[256]; - int rv; - - // set numa scheme - g_state.numa.numa_strategy = numa_flag; - - GGML_PRINT_DEBUG("numa strategy %u\n",g_state.numa.numa_strategy); - - g_state.numa.cpuset = ggml_get_numa_affinity(); - - // enumerate nodes - while (g_state.numa.n_nodes < GGML_NUMA_MAX_NODES) { - rv = snprintf(path, sizeof(path), "/sys/devices/system/node/node%u", g_state.numa.n_nodes); - GGML_ASSERT(rv > 0 && (unsigned)rv < sizeof(path)); - if (stat(path, &st) != 0) { break; } - ++g_state.numa.n_nodes; - } - - // enumerate CPUs - while (g_state.numa.total_cpus < GGML_NUMA_MAX_CPUS) { - rv = snprintf(path, sizeof(path), "/sys/devices/system/cpu/cpu%u", g_state.numa.total_cpus); - GGML_ASSERT(rv > 0 && (unsigned)rv < sizeof(path)); - if (stat(path, &st) != 0) { break; } - ++g_state.numa.total_cpus; - } - - GGML_PRINT_DEBUG("found %u numa nodes, %u CPUs\n", g_state.numa.n_nodes, g_state.numa.total_cpus); - - // figure out which node we're on - uint current_cpu; - int getcpu_ret = 0; -#if __GLIBC__ > 2 || (__GLIBC__ == 2 && __GLIBC_MINOR__ > 33) || defined(__COSMOPOLITAN__) - getcpu_ret = getcpu(¤t_cpu, &g_state.numa.current_node); -#else - // old glibc doesn't have a wrapper for this call. Fall back on direct syscall -# if !defined(SYS_getcpu) && defined(SYS_get_cpu) -# define SYS_getcpu SYS_get_cpu // some older glibc versions use this name -# endif - getcpu_ret = syscall(SYS_getcpu, ¤t_cpu, &g_state.numa.current_node); -#endif - - if (g_state.numa.n_nodes < 1 || g_state.numa.total_cpus < 1 || getcpu_ret != 0) { - g_state.numa.n_nodes = 0; - return; - } - - GGML_PRINT_DEBUG("found our process on numa node %u, CPU %u\n", g_state.numa.current_node, current_cpu); - - for (uint32_t n = 0; n < g_state.numa.n_nodes; ++n) { - struct ggml_numa_node * node = &g_state.numa.nodes[n]; - GGML_PRINT_DEBUG("CPUs on node %u:", n); - node->n_cpus = 0; - for (uint32_t c = 0; c < g_state.numa.total_cpus; ++c) { - rv = snprintf(path, sizeof(path), "/sys/devices/system/node/node%u/cpu%u", n, c); - GGML_ASSERT(rv > 0 && (unsigned)rv < sizeof(path)); - if (stat(path, &st) == 0) { - node->cpus[node->n_cpus++] = c; - GGML_PRINT_DEBUG(" %u", c); - } - } - GGML_PRINT_DEBUG("\n"); - } - - if (ggml_is_numa()) { - FILE *fptr = fopen("/proc/sys/kernel/numa_balancing", "r"); - if (fptr != NULL) { - char buf[42]; - if (fgets(buf, sizeof(buf), fptr) && strncmp(buf, "0\n", sizeof(buf)) != 0) { - GGML_LOG_WARN("/proc/sys/kernel/numa_balancing is enabled, this has been observed to impair performance\n"); - } - fclose(fptr); - } - } -#else - UNUSED(numa_flag); - // TODO -#endif -} - -bool ggml_is_numa(void) { - return g_state.numa.n_nodes > 1; -} - -#if defined(__ARM_ARCH) -#if defined(__aarch64__) && defined(__ARM_FEATURE_SVE) -#include -static void ggml_init_arm_arch_features(void) { - ggml_arm_arch_features.sve_cnt = svcntb(); -} -#else -static void ggml_init_arm_arch_features(void) {} -#endif -#endif // __ARM_ARCH - -#if defined(__riscv) && defined(__riscv_v_intrinsic) -#include -static void ggml_init_riscv_arch_features(void) { - ggml_riscv_arch_features.rvv_vlen = __riscv_vlenb(); -} -#else -static void ggml_init_riscv_arch_features(void) {} -#endif - -struct ggml_tensor * ggml_new_i32(struct ggml_context * ctx, int32_t value) { - GGML_ASSERT(!ggml_get_no_alloc(ctx)); - - struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, 1); - - ggml_set_i32(result, value); - - return result; -} - -struct ggml_tensor * ggml_new_f32(struct ggml_context * ctx, float value) { - GGML_ASSERT(!ggml_get_no_alloc(ctx)); - - struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, 1); - - ggml_set_f32(result, value); - - return result; -} - -struct ggml_tensor * ggml_set_i32 (struct ggml_tensor * tensor, int32_t value) { - const int n = ggml_nrows(tensor); - const int nc = tensor->ne[0]; - const size_t n1 = tensor->nb[1]; - - char * const data = tensor->data; - - switch (tensor->type) { - case GGML_TYPE_I8: - { - assert(tensor->nb[0] == sizeof(int8_t)); - for (int i = 0; i < n; i++) { - ggml_vec_set_i8(nc, (int8_t *)(data + i*n1), value); - } - } break; - case GGML_TYPE_I16: - { - assert(tensor->nb[0] == sizeof(int16_t)); - for (int i = 0; i < n; i++) { - ggml_vec_set_i16(nc, (int16_t *)(data + i*n1), value); - } - } break; - case GGML_TYPE_I32: - { - assert(tensor->nb[0] == sizeof(int32_t)); - for (int i = 0; i < n; i++) { - ggml_vec_set_i32(nc, (int32_t *)(data + i*n1), value); - } - } break; - case GGML_TYPE_F16: - { - assert(tensor->nb[0] == sizeof(ggml_fp16_t)); - for (int i = 0; i < n; i++) { - ggml_vec_set_f16(nc, (ggml_fp16_t *)(data + i*n1), GGML_CPU_FP32_TO_FP16(value)); - } - } break; - case GGML_TYPE_BF16: - { - assert(tensor->nb[0] == sizeof(ggml_fp16_t)); - for (int i = 0; i < n; i++) { - ggml_vec_set_bf16(nc, (ggml_bf16_t *)(data + i*n1), GGML_FP32_TO_BF16(value)); - } - } break; - case GGML_TYPE_F32: - { - assert(tensor->nb[0] == sizeof(float)); - for (int i = 0; i < n; i++) { - ggml_vec_set_f32(nc, (float *)(data + i*n1), value); - } - } break; - default: - { - GGML_ABORT("fatal error"); - } - } - - return tensor; -} - -struct ggml_tensor * ggml_set_f32(struct ggml_tensor * tensor, float value) { - const int n = ggml_nrows(tensor); - const int nc = tensor->ne[0]; - const size_t n1 = tensor->nb[1]; - - char * const data = tensor->data; - - switch (tensor->type) { - case GGML_TYPE_I8: - { - assert(tensor->nb[0] == sizeof(int8_t)); - for (int i = 0; i < n; i++) { - ggml_vec_set_i8(nc, (int8_t *)(data + i*n1), value); - } - } break; - case GGML_TYPE_I16: - { - assert(tensor->nb[0] == sizeof(int16_t)); - for (int i = 0; i < n; i++) { - ggml_vec_set_i16(nc, (int16_t *)(data + i*n1), value); - } - } break; - case GGML_TYPE_I32: - { - assert(tensor->nb[0] == sizeof(int32_t)); - for (int i = 0; i < n; i++) { - ggml_vec_set_i32(nc, (int32_t *)(data + i*n1), value); - } - } break; - case GGML_TYPE_F16: - { - assert(tensor->nb[0] == sizeof(ggml_fp16_t)); - for (int i = 0; i < n; i++) { - ggml_vec_set_f16(nc, (ggml_fp16_t *)(data + i*n1), GGML_CPU_FP32_TO_FP16(value)); - } - } break; - case GGML_TYPE_BF16: - { - assert(tensor->nb[0] == sizeof(ggml_bf16_t)); - for (int i = 0; i < n; i++) { - ggml_vec_set_bf16(nc, (ggml_bf16_t *)(data + i*n1), GGML_FP32_TO_BF16(value)); - } - } break; - case GGML_TYPE_F32: - { - assert(tensor->nb[0] == sizeof(float)); - for (int i = 0; i < n; i++) { - ggml_vec_set_f32(nc, (float *)(data + i*n1), value); - } - } break; - default: - { - GGML_ABORT("fatal error"); - } - } - - return tensor; -} - -int32_t ggml_get_i32_1d(const struct ggml_tensor * tensor, int i) { - if (!ggml_is_contiguous(tensor)) { - int64_t id[4] = { 0, 0, 0, 0 }; - ggml_unravel_index(tensor, i, &id[0], &id[1], &id[2], &id[3]); - return ggml_get_i32_nd(tensor, id[0], id[1], id[2], id[3]); - } - switch (tensor->type) { - case GGML_TYPE_I8: - { - GGML_ASSERT(tensor->nb[0] == sizeof(int8_t)); - return ((int8_t *)(tensor->data))[i]; - } - case GGML_TYPE_I16: - { - GGML_ASSERT(tensor->nb[0] == sizeof(int16_t)); - return ((int16_t *)(tensor->data))[i]; - } - case GGML_TYPE_I32: - { - GGML_ASSERT(tensor->nb[0] == sizeof(int32_t)); - return ((int32_t *)(tensor->data))[i]; - } - case GGML_TYPE_F16: - { - GGML_ASSERT(tensor->nb[0] == sizeof(ggml_fp16_t)); - return GGML_CPU_FP16_TO_FP32(((ggml_fp16_t *)(tensor->data))[i]); - } - case GGML_TYPE_BF16: - { - GGML_ASSERT(tensor->nb[0] == sizeof(ggml_bf16_t)); - return GGML_BF16_TO_FP32(((ggml_bf16_t *)(tensor->data))[i]); - } - case GGML_TYPE_F32: - { - GGML_ASSERT(tensor->nb[0] == sizeof(float)); - return ((float *)(tensor->data))[i]; - } - default: - { - GGML_ABORT("fatal error"); - } - } -} - -void ggml_set_i32_1d(const struct ggml_tensor * tensor, int i, int32_t value) { - if (!ggml_is_contiguous(tensor)) { - int64_t id[4] = { 0, 0, 0, 0 }; - ggml_unravel_index(tensor, i, &id[0], &id[1], &id[2], &id[3]); - ggml_set_i32_nd(tensor, id[0], id[1], id[2], id[3], value); - return; - } - switch (tensor->type) { - case GGML_TYPE_I8: - { - GGML_ASSERT(tensor->nb[0] == sizeof(int8_t)); - ((int8_t *)(tensor->data))[i] = value; - } break; - case GGML_TYPE_I16: - { - GGML_ASSERT(tensor->nb[0] == sizeof(int16_t)); - ((int16_t *)(tensor->data))[i] = value; - } break; - case GGML_TYPE_I32: - { - GGML_ASSERT(tensor->nb[0] == sizeof(int32_t)); - ((int32_t *)(tensor->data))[i] = value; - } break; - case GGML_TYPE_F16: - { - GGML_ASSERT(tensor->nb[0] == sizeof(ggml_fp16_t)); - ((ggml_fp16_t *)(tensor->data))[i] = GGML_CPU_FP32_TO_FP16(value); - } break; - case GGML_TYPE_BF16: - { - GGML_ASSERT(tensor->nb[0] == sizeof(ggml_bf16_t)); - ((ggml_bf16_t *)(tensor->data))[i] = GGML_FP32_TO_BF16(value); - } break; - case GGML_TYPE_F32: - { - GGML_ASSERT(tensor->nb[0] == sizeof(float)); - ((float *)(tensor->data))[i] = value; - } break; - default: - { - GGML_ABORT("fatal error"); - } - } -} - -int32_t ggml_get_i32_nd(const struct ggml_tensor * tensor, int i0, int i1, int i2, int i3) { - void * data = (char *) tensor->data + i0*tensor->nb[0] + i1*tensor->nb[1] + i2*tensor->nb[2] + i3*tensor->nb[3]; - switch (tensor->type) { - case GGML_TYPE_I8: - return ((int8_t *) data)[0]; - case GGML_TYPE_I16: - return ((int16_t *) data)[0]; - case GGML_TYPE_I32: - return ((int32_t *) data)[0]; - case GGML_TYPE_F16: - return GGML_CPU_FP16_TO_FP32(((ggml_fp16_t *) data)[0]); - case GGML_TYPE_BF16: - return GGML_BF16_TO_FP32(((ggml_bf16_t *) data)[0]); - case GGML_TYPE_F32: - return ((float *) data)[0]; - default: - GGML_ABORT("fatal error"); - } -} - -void ggml_set_i32_nd(const struct ggml_tensor * tensor, int i0, int i1, int i2, int i3, int32_t value) { - void * data = (char *) tensor->data + i0*tensor->nb[0] + i1*tensor->nb[1] + i2*tensor->nb[2] + i3*tensor->nb[3]; - switch (tensor->type) { - case GGML_TYPE_I8: - { - ((int8_t *)(data))[0] = value; - } break; - case GGML_TYPE_I16: - { - ((int16_t *)(data))[0] = value; - } break; - case GGML_TYPE_I32: - { - ((int32_t *)(data))[0] = value; - } break; - case GGML_TYPE_F16: - { - ((ggml_fp16_t *)(data))[0] = GGML_CPU_FP32_TO_FP16(value); - } break; - case GGML_TYPE_BF16: - { - ((ggml_bf16_t *)(data))[0] = GGML_FP32_TO_BF16(value); - } break; - case GGML_TYPE_F32: - { - ((float *)(data))[0] = value; - } break; - default: - { - GGML_ABORT("fatal error"); - } - } -} - -float ggml_get_f32_1d(const struct ggml_tensor * tensor, int i) { - if (!ggml_is_contiguous(tensor)) { - int64_t id[4] = { 0, 0, 0, 0 }; - ggml_unravel_index(tensor, i, &id[0], &id[1], &id[2], &id[3]); - return ggml_get_f32_nd(tensor, id[0], id[1], id[2], id[3]); - } - switch (tensor->type) { - case GGML_TYPE_I8: - { - return ((int8_t *)(tensor->data))[i]; - } - case GGML_TYPE_I16: - { - return ((int16_t *)(tensor->data))[i]; - } - case GGML_TYPE_I32: - { - return ((int32_t *)(tensor->data))[i]; - } - case GGML_TYPE_F16: - { - return GGML_CPU_FP16_TO_FP32(((ggml_fp16_t *)(tensor->data))[i]); - } - case GGML_TYPE_BF16: - { - return GGML_BF16_TO_FP32(((ggml_bf16_t *)(tensor->data))[i]); - } - case GGML_TYPE_F32: - { - return ((float *)(tensor->data))[i]; - } - default: - { - GGML_ABORT("fatal error"); - } - } -} - -void ggml_set_f32_1d(const struct ggml_tensor * tensor, int i, float value) { - if (!ggml_is_contiguous(tensor)) { - int64_t id[4] = { 0, 0, 0, 0 }; - ggml_unravel_index(tensor, i, &id[0], &id[1], &id[2], &id[3]); - ggml_set_f32_nd(tensor, id[0], id[1], id[2], id[3], value); - return; - } - switch (tensor->type) { - case GGML_TYPE_I8: - { - ((int8_t *)(tensor->data))[i] = value; - } break; - case GGML_TYPE_I16: - { - ((int16_t *)(tensor->data))[i] = value; - } break; - case GGML_TYPE_I32: - { - ((int32_t *)(tensor->data))[i] = value; - } break; - case GGML_TYPE_F16: - { - ((ggml_fp16_t *)(tensor->data))[i] = GGML_CPU_FP32_TO_FP16(value); - } break; - case GGML_TYPE_BF16: - { - ((ggml_bf16_t *)(tensor->data))[i] = GGML_FP32_TO_BF16(value); - } break; - case GGML_TYPE_F32: - { - ((float *)(tensor->data))[i] = value; - } break; - default: - { - GGML_ABORT("fatal error"); - } - } -} - -float ggml_get_f32_nd(const struct ggml_tensor * tensor, int i0, int i1, int i2, int i3) { - void * data = (char *) tensor->data + i0*tensor->nb[0] + i1*tensor->nb[1] + i2*tensor->nb[2] + i3*tensor->nb[3]; - switch (tensor->type) { - case GGML_TYPE_I8: - return ((int8_t *) data)[0]; - case GGML_TYPE_I16: - return ((int16_t *) data)[0]; - case GGML_TYPE_I32: - return ((int32_t *) data)[0]; - case GGML_TYPE_F16: - return GGML_CPU_FP16_TO_FP32(((ggml_fp16_t *) data)[0]); - case GGML_TYPE_BF16: - return GGML_BF16_TO_FP32(((ggml_bf16_t *) data)[0]); - case GGML_TYPE_F32: - return ((float *) data)[0]; - default: - GGML_ABORT("fatal error"); - } -} - -void ggml_set_f32_nd(const struct ggml_tensor * tensor, int i0, int i1, int i2, int i3, float value) { - void * data = (char *) tensor->data + i0*tensor->nb[0] + i1*tensor->nb[1] + i2*tensor->nb[2] + i3*tensor->nb[3]; - switch (tensor->type) { - case GGML_TYPE_I8: - { - ((int8_t *)(data))[0] = value; - } break; - case GGML_TYPE_I16: - { - ((int16_t *)(data))[0] = value; - } break; - case GGML_TYPE_I32: - { - ((int32_t *)(data))[0] = value; - } break; - case GGML_TYPE_F16: - { - ((ggml_fp16_t *)(data))[0] = GGML_CPU_FP32_TO_FP16(value); - } break; - case GGML_TYPE_BF16: - { - ((ggml_bf16_t *)(data))[0] = GGML_FP32_TO_BF16(value); - } break; - case GGML_TYPE_F32: - { - ((float *)(data))[0] = value; - } break; - default: - { - GGML_ABORT("fatal error"); - } - } -} - -//////////////////////////////////////////////////////////////////////////////// - -// ggml_compute_forward_mul_mat - -static void ggml_compute_forward_mul_mat_one_chunk( - const struct ggml_compute_params * params, - struct ggml_tensor * dst, - const enum ggml_type type, - const int64_t num_rows_per_vec_dot, - const int64_t ir0_start, - const int64_t ir0_end, - const int64_t ir1_start, - const int64_t ir1_end) { - - const struct ggml_tensor * src0 = dst->src[0]; - const struct ggml_tensor * src1 = dst->src[1]; - - GGML_TENSOR_BINARY_OP_LOCALS - - const bool src1_cont = ggml_is_contiguous(src1); - - ggml_vec_dot_t const vec_dot = type_traits_cpu[type].vec_dot; - enum ggml_type const vec_dot_type = type_traits_cpu[type].vec_dot_type; - - // broadcast factors - const int64_t r2 = ne12 / ne02; - const int64_t r3 = ne13 / ne03; - - //printf("ir0_start = %6lld, ir0_end = %6lld, ir1_start = %6lld, ir1_end = %6lld\n", ir0_start, ir0_end, ir1_start, ir1_end); - - // threads with no work simply yield (not sure if it helps) - if (ir0_start >= ir0_end || ir1_start >= ir1_end) { - return; - } - - const void * wdata = (src1->type == vec_dot_type) ? src1->data : params->wdata; - const size_t row_size = ggml_row_size(vec_dot_type, ne10); - - assert(ne12 % ne02 == 0); - assert(ne13 % ne03 == 0); - - // block-tiling attempt - const int64_t blck_0 = 16; - const int64_t blck_1 = 16; - - const size_t src1_col_stride = src1_cont || src1->type != vec_dot_type ? row_size : nb11; - - // attempt to reduce false-sharing (does not seem to make a difference) - // 16 * 2, accounting for mmla kernels - float tmp[32]; - - for (int64_t iir1 = ir1_start; iir1 < ir1_end; iir1 += blck_1) { - for (int64_t iir0 = ir0_start; iir0 < ir0_end; iir0 += blck_0) { - for (int64_t ir1 = iir1; ir1 < iir1 + blck_1 && ir1 < ir1_end; ir1 += num_rows_per_vec_dot) { - const int64_t i13 = (ir1 / (ne12 * ne1)); - const int64_t i12 = (ir1 - i13 * ne12 * ne1) / ne1; - const int64_t i11 = (ir1 - i13 * ne12 * ne1 - i12 * ne1); - - // broadcast src0 into src1 - const int64_t i03 = i13 / r3; - const int64_t i02 = i12 / r2; - - const int64_t i1 = i11; - const int64_t i2 = i12; - const int64_t i3 = i13; - - const char * src0_row = (const char*)src0->data + (0 + i02 * nb02 + i03 * nb03); - - // desc: when src1 is not a contiguous memory block we have to calculate the offset using the strides - // if it is, then we have either copied the data to params->wdata and made it contiguous or we are using - // the original src1 data pointer, so we should index using the indices directly - // TODO: this is a bit of a hack, we should probably have a better way to handle this - const char * src1_col = (const char*)wdata + - (src1_cont || src1->type != vec_dot_type - ? (i11 + i12 * ne11 + i13 * ne12 * ne11) * row_size - : (i11 * nb11 + i12 * nb12 + i13 * nb13)); - float * dst_col = (float*)((char*)dst->data + (i1 * nb1 + i2 * nb2 + i3 * nb3)); - - //for (int64_t ir0 = iir0; ir0 < iir0 + blck_0 && ir0 < ir0_end; ++ir0) { - // vec_dot(ne00, &dst_col[ir0], src0_row + ir0*nb01, src1_col); - //} - - for (int64_t ir0 = iir0; ir0 < iir0 + blck_0 && ir0 < ir0_end; ir0 += num_rows_per_vec_dot) { - vec_dot(ne00, &tmp[ir0 - iir0], (num_rows_per_vec_dot > 1 ? 16 : 0), src0_row + ir0 * nb01, (num_rows_per_vec_dot > 1 ? nb01 : 0), src1_col, (num_rows_per_vec_dot > 1 ? src1_col_stride : 0), num_rows_per_vec_dot); - } - - for (int cn = 0; cn < num_rows_per_vec_dot; ++cn) { - memcpy(&dst_col[iir0 + cn * nb1 / nb0], tmp + (cn * 16), (MIN(iir0 + blck_0, ir0_end) - iir0) * sizeof(float)); - } - } - } - } -} - -void ggml_compute_forward_mul_mat( - const struct ggml_compute_params * params, - struct ggml_tensor * dst) { - - const struct ggml_tensor * src0 = dst->src[0]; - const struct ggml_tensor * src1 = dst->src[1]; - - GGML_TENSOR_BINARY_OP_LOCALS - - const int ith = params->ith; - const int nth = params->nth; - - enum ggml_type const vec_dot_type = type_traits_cpu[src0->type].vec_dot_type; - ggml_from_float_t const from_float = type_traits_cpu[vec_dot_type].from_float; - int64_t const vec_dot_num_rows = type_traits_cpu[src0->type].nrows; - - GGML_ASSERT(ne0 == ne01); - GGML_ASSERT(ne1 == ne11); - GGML_ASSERT(ne2 == ne12); - GGML_ASSERT(ne3 == ne13); - - // we don't support permuted src0 or src1 - GGML_ASSERT(nb00 == ggml_type_size(src0->type)); - GGML_ASSERT(nb10 == ggml_type_size(src1->type)); - - // dst cannot be transposed or permuted - GGML_ASSERT(nb0 == sizeof(float)); - GGML_ASSERT(nb0 <= nb1); - GGML_ASSERT(nb1 <= nb2); - GGML_ASSERT(nb2 <= nb3); - - // nb01 >= nb00 - src0 is not transposed - // compute by src0 rows - - // TODO: extract to "extra_op" -#if GGML_USE_LLAMAFILE - // broadcast factors - const int64_t r2 = ne12 / ne02; - const int64_t r3 = ne13 / ne03; - - const bool src1_cont = ggml_is_contiguous(src1); - - if (src1_cont) { - for (int64_t i13 = 0; i13 < ne13; i13++) - for (int64_t i12 = 0; i12 < ne12; i12++) - if (!llamafile_sgemm(params, - ne01, ne11, ne00/ggml_blck_size(src0->type), - (const char *)src0->data + i12/r2*nb02 + i13/r3*nb03, - nb01/ggml_type_size(src0->type), - (const char *)src1->data + i12*nb12 + i13*nb13, - nb11/ggml_type_size(src1->type), - (char *)dst->data + i12*nb2 + i13*nb3, - nb1/ggml_type_size(dst->type), - src0->type, - src1->type, - dst->type)) - goto UseGgmlGemm1; - return; - } -UseGgmlGemm1:; -#endif - - if (src1->type != vec_dot_type) { - char * wdata = params->wdata; - - const size_t nbw0 = ggml_type_size(vec_dot_type); - const size_t nbw1 = ggml_row_size(vec_dot_type, ne10); - const size_t nbw2 = nbw1*ne11; - const size_t nbw3 = nbw2*ne12; - - assert(params->wsize >= ne13*nbw3); - GGML_ASSERT(src1->type == GGML_TYPE_F32); - - #if 0 - for (int64_t i13 = 0; i13 < ne13; ++i13) { - for (int64_t i12 = 0; i12 < ne12; ++i12) { - for (int64_t i11 = ith; i11 < ne11; i11 += nth) { - from_float((float *)((char *) src1->data + i13*nb13 + i12*nb12 + i11*nb11), - (void *) (wdata + i13*nbw3 + i12*nbw2 + i11*nbw1), - ne10); - } - } - } - #else - for (int64_t i13 = 0; i13 < ne13; ++i13) { - for (int64_t i12 = 0; i12 < ne12; ++i12) { - for (int64_t i11 = 0; i11 < ne11; ++i11) { - size_t bs = ggml_blck_size(vec_dot_type); - int64_t ne10_block_start = (ith * ne10/bs) / nth; - int64_t ne10_block_end = ((ith + 1) * ne10/bs) / nth; - from_float((float *)((char *) src1->data + i13*nb13 + i12*nb12 + i11*nb11 + ne10_block_start*bs*nb10), - (void *) (wdata + i13*nbw3 + i12*nbw2 + i11*nbw1 + ne10_block_start*nbw0), - (ne10_block_end - ne10_block_start) * bs); - } - } - } - #endif - } - - if (ith == 0) { - // Every thread starts at ith, so the first unprocessed chunk is nth. This save a bit of coordination right at the start. - atomic_store_explicit(¶ms->threadpool->current_chunk, nth, memory_order_relaxed); - } - - ggml_barrier(params->threadpool); - -#if GGML_USE_LLAMAFILE - if (src1->type != vec_dot_type) { - const void* wdata = (src1->type == vec_dot_type) ? src1->data : params->wdata; - const size_t row_size = ggml_row_size(vec_dot_type, ne10); - - for (int64_t i13 = 0; i13 < ne13; i13++) - for (int64_t i12 = 0; i12 < ne12; i12++) - if (!llamafile_sgemm(params, - ne01, ne11, ne00/ggml_blck_size(src0->type), - (const char *)src0->data + i12/r2*nb02 + i13/r3*nb03, - nb01/ggml_type_size(src0->type), - (const char *)wdata + (i12*ne11 + i13*ne12*ne11)*row_size, - row_size/ggml_type_size(vec_dot_type), - (char *)dst->data + i12*nb2 + i13*nb3, - nb1/ggml_type_size(dst->type), - src0->type, - vec_dot_type, - dst->type)) - goto UseGgmlGemm2; - return; - } -UseGgmlGemm2:; -#endif - - // This is the size of the first dimension of the result, so we can iterate that way. (see the ASSERT above, these are the same numbers) - const int64_t nr0 = ne0; - - // This is the size of the rest of the dimensions of the result - const int64_t nr1 = ne1 * ne2 * ne3; - - // Now select a reasonable chunk size. - int chunk_size = 16; - - // We need to step up the size if it's small - if (nr0 == 1 || nr1 == 1) { - chunk_size = 64; - } - - // distribute the work across the inner or outer loop based on which one is larger - // The number of chunks in the 0/1 dim. - // CEIL(nr0/chunk_size) - int64_t nchunk0 = (nr0 + chunk_size - 1) / chunk_size; - int64_t nchunk1 = (nr1 + chunk_size - 1) / chunk_size; - - // If the chunking is poor for the number of threads on this setup, scrap the whole plan. Re-chunk it by thread. - // Also, chunking by thread was measured to have perform better on NUMA systems. See https://github.com/ggml-org/llama.cpp/pull/6915 - // In theory, chunking should be just as useful on NUMA and non NUMA systems, but testing disagreed with that. - if (nchunk0 * nchunk1 < nth * 4 || ggml_is_numa()) { - // distribute the thread work across the inner or outer loop based on which one is larger - nchunk0 = nr0 > nr1 ? nth : 1; // parallelize by src0 rows - nchunk1 = nr0 > nr1 ? 1 : nth; // parallelize by src1 rows - } - - // The number of elements in each chunk - const int64_t dr0 = (nr0 + nchunk0 - 1) / nchunk0; - const int64_t dr1 = (nr1 + nchunk1 - 1) / nchunk1; - - // The first chunk comes from our thread_id, the rest will get auto-assigned. - int current_chunk = ith; - - while (current_chunk < nchunk0 * nchunk1) { - const int64_t ith0 = current_chunk % nchunk0; - const int64_t ith1 = current_chunk / nchunk0; - - const int64_t ir0_start = dr0 * ith0; - const int64_t ir0_end = MIN(ir0_start + dr0, nr0); - - const int64_t ir1_start = dr1 * ith1; - const int64_t ir1_end = MIN(ir1_start + dr1, nr1); - - // dot kernels can handle 1 row and col at a time, but mmla kernels can process 2 rows and cols - int64_t num_rows_per_vec_dot = vec_dot_num_rows; - - // these checks are needed to avoid crossing dim1 boundaries - // can be optimized, but the logic would become more complicated, so keeping it like this for simplicity - if ((nr0 % 2 != 0) || (ne11 % 2 != 0) || ((ir0_end - ir0_start) % 2 != 0) || ((ir1_end - ir1_start) % 2 != 0)) { - num_rows_per_vec_dot = 1; - } - ggml_compute_forward_mul_mat_one_chunk(params, dst, src0->type, num_rows_per_vec_dot, ir0_start, ir0_end, ir1_start, ir1_end); - - if (nth >= nchunk0 * nchunk1) { - break; - } - - current_chunk = atomic_fetch_add_explicit(¶ms->threadpool->current_chunk, 1, memory_order_relaxed); - } -} - -// ggml_compute_forward_mul_mat_id - -#define MMID_MATRIX_ROW(row_id, i1) matrix_rows[(row_id)*ids->ne[0]*ids->ne[1] + (i1)] - -struct mmid_row_mapping { - int32_t i1; - int32_t i2; -}; - -static void ggml_compute_forward_mul_mat_id_one_chunk( - struct ggml_tensor * dst, - const struct ggml_tensor * src0, - const struct ggml_tensor * src1, - const struct ggml_tensor * ids, - const int64_t cur_a, - const int64_t ir0_start, - const int64_t ir0_end, - const int64_t ir1_start, - const int64_t ir1_end, - const char * src0_cur, - const struct mmid_row_mapping * matrix_rows, - const size_t row_size, - const bool src1_cont, - const void * wdata) { - - GGML_TENSOR_BINARY_OP_LOCALS - - const enum ggml_type type = src0->type; - - ggml_vec_dot_t const vec_dot = type_traits_cpu[type].vec_dot; - enum ggml_type const vec_dot_type = type_traits_cpu[type].vec_dot_type; - - const int64_t blck_0 = 16; - const int64_t blck_1 = 16; - - float tmp[16]; - - for (int64_t iir1 = ir1_start; iir1 < ir1_end; iir1 += blck_1) { - for (int64_t iir0 = ir0_start; iir0 < ir0_end; iir0 += blck_0) { - for (int64_t ir1 = iir1; ir1 < iir1 + blck_1 && ir1 < ir1_end; ++ir1) { - const int64_t _i12 = ir1; // logical row index for this expert - - struct mmid_row_mapping row_mapping = MMID_MATRIX_ROW(cur_a, _i12); - const int id = row_mapping.i1; // selected expert index - - const int64_t i11 = id % ne11; - const int64_t i12 = row_mapping.i2; // row index in src1 - - const int64_t i1 = id; // selected expert index - const int64_t i2 = i12; // row - - // desc: when src1 is not a contiguous memory block we have to calculate the offset using the strides - // if it is, then we have either copied the data to params->wdata and made it contiguous or we are using - // the original src1 data pointer, so we should index using the indices directly - // TODO: this is a bit of a hack, we should probably have a better way to handle this - const char * src1_col = (const char *) wdata + - (src1_cont || src1->type != vec_dot_type - ? (i11 + i12*ne11)*row_size - : (i11*nb11 + i12*nb12)); - - float * dst_col = (float *) ((char *) dst->data + (i1*nb1 + i2*nb2)); - - for (int64_t ir0 = iir0; ir0 < iir0 + blck_0 && ir0 < ir0_end; ++ir0) { - vec_dot(ne00, &tmp[ir0 - iir0], 0, src0_cur + ir0*nb01, 0, src1_col, 0, 1); - } - - memcpy(&dst_col[iir0], tmp, (MIN(iir0 + blck_0, ir0_end) - iir0)*sizeof(float)); - } - } - } -} - -static void * incr_ptr_aligned(void ** p, size_t size, size_t align) { - - void * ptr = *p; - ptr = (void *) GGML_PAD((uintptr_t) ptr, align); - *p = (void *) ((char *) ptr + size); - return ptr; -} - -static void ggml_compute_forward_mul_mat_id( - const struct ggml_compute_params * params, - struct ggml_tensor * dst) { - - const struct ggml_tensor * src0 = dst->src[0]; - const struct ggml_tensor * src1 = dst->src[1]; - const struct ggml_tensor * ids = dst->src[2]; - - GGML_TENSOR_BINARY_OP_LOCALS - - const int ith = params->ith; - const int nth = params->nth; - - const enum ggml_type type = src0->type; - - const bool src1_cont = ggml_is_contiguous(src1); - - enum ggml_type const vec_dot_type = type_traits_cpu[type].vec_dot_type; - ggml_from_float_t const from_float = type_traits_cpu[vec_dot_type].from_float; - - // we don't support permuted src0 or src1 - GGML_ASSERT(nb00 == ggml_type_size(type)); - GGML_ASSERT(nb10 == ggml_type_size(src1->type)); - - // dst cannot be transposed or permuted - GGML_ASSERT(nb0 == sizeof(float)); - GGML_ASSERT(nb0 <= nb1); - GGML_ASSERT(nb1 <= nb2); - GGML_ASSERT(nb2 <= nb3); - - // row groups - const int n_ids = ids->ne[0]; // n_expert_used - const int n_as = ne02; // n_expert - - void * wdata_cur = params->wdata; - - if (src1->type != vec_dot_type) { - incr_ptr_aligned(&wdata_cur, ggml_row_size(vec_dot_type, ggml_nelements(src1)), sizeof(int64_t)); - } - - int64_t * matrix_row_counts = // [n_as] - incr_ptr_aligned(&wdata_cur, n_as*sizeof(int64_t), sizeof(int64_t)); - - struct mmid_row_mapping * matrix_rows = // [n_as][ids->ne[0]*ids->ne[1]] - incr_ptr_aligned(&wdata_cur, n_as*ids->ne[0]*ids->ne[1]*sizeof(struct mmid_row_mapping), sizeof(int64_t)); - - char (*atomic_current_chunk)[CACHE_LINE_SIZE] = // [n_as] - incr_ptr_aligned(&wdata_cur, CACHE_LINE_SIZE * n_as, CACHE_LINE_SIZE); - - GGML_ASSERT(params->wsize >= (size_t)((char *) wdata_cur - (char *) params->wdata)); - - if (src1->type != vec_dot_type) { - char * wdata = params->wdata; - - const size_t nbw0 = ggml_type_size(vec_dot_type); - const size_t nbw1 = ggml_row_size(vec_dot_type, ne10); - const size_t nbw2 = nbw1*ne11; - const size_t nbw3 = nbw2*ne12; - - assert(params->wsize >= ne13*nbw3); - GGML_ASSERT(src1->type == GGML_TYPE_F32); - -#if 0 - for (int64_t i13 = 0; i13 < ne13; ++i13) { - for (int64_t i12 = ith; i12 < ne12; i12 += nth) { - for (int64_t i11 = 0; i11 < ne11; ++i11) { - from_float((float *)((char *) src1->data + i13*nb13 + i12*nb12 + i11*nb11), - (void *) (wdata + i13*nbw3 + i12*nbw2 + i11*nbw1), - ne10); - } - } - } -#else - for (int64_t i13 = 0; i13 < ne13; ++i13) { - for (int64_t i12 = 0; i12 < ne12; ++i12) { - for (int64_t i11 = 0; i11 < ne11; ++i11) { - size_t bs = ggml_blck_size(vec_dot_type); - int64_t ne10_block_start = (ith * ne10/bs) / nth; - int64_t ne10_block_end = ((ith + 1) * ne10/bs) / nth; - from_float((float *)((char *) src1->data + i13*nb13 + i12*nb12 + i11*nb11 + ne10_block_start*bs*nb10), - (void *) (wdata + i13*nbw3 + i12*nbw2 + i11*nbw1 + ne10_block_start*nbw0), - (ne10_block_end - ne10_block_start) * bs); - } - } - } -#endif - } - - if (ith == 0) { - // initialize matrix_row_counts - memset(matrix_row_counts, 0, n_as*sizeof(int64_t)); - - // group rows by src0 matrix - for (int64_t iid1 = 0; iid1 < ids->ne[1]; ++iid1) { - for (int id = 0; id < n_ids; ++id) { - const int32_t i02 = *(const int32_t *) ((const char *) ids->data + iid1*ids->nb[1] + id*ids->nb[0]); - - assert(i02 >= 0 && i02 < n_as); - - MMID_MATRIX_ROW(i02, matrix_row_counts[i02]) = (struct mmid_row_mapping) {id, iid1}; - matrix_row_counts[i02] += 1; - } - } - } - - // reset current_chunk - for (int cur_a = ith; cur_a < n_as; cur_a += nth) { - atomic_int * current_chunk_ctr = (atomic_int *)(atomic_current_chunk + cur_a); - *current_chunk_ctr = nth; - } - - ggml_barrier(params->threadpool); - - for (int cur_a = 0; cur_a < n_as; ++cur_a) { - const int64_t cne1 = matrix_row_counts[cur_a]; - - if (cne1 == 0) { - continue; - } - - const char * src0_cur = (const char *) src0->data + cur_a * nb02; - const void * wdata = (src1->type == vec_dot_type) ? src1->data : params->wdata; - const size_t row_size = ggml_row_size(vec_dot_type, ne10); - - const int64_t nr0 = ne01; - const int64_t nr1 = cne1; - - int chunk_size = 16; - if (nr0 == 1 || nr1 == 1) { - chunk_size = 64; - } - - // disable for NUMA - const bool disable_chunking = ggml_is_numa(); - - int64_t nchunk0 = (nr0 + chunk_size - 1) / chunk_size; - int64_t nchunk1 = (nr1 + chunk_size - 1) / chunk_size; - - if (nchunk0 * nchunk1 < nth * 4 || disable_chunking) { - nchunk0 = nr0 > nr1 ? nth : 1; - nchunk1 = nr0 > nr1 ? 1 : nth; - } - - const int64_t dr0 = (nr0 + nchunk0 - 1) / nchunk0; - const int64_t dr1 = (nr1 + nchunk1 - 1) / nchunk1; - - int current_chunk = ith; - - atomic_int * current_chunk_ctr = (atomic_int *)(atomic_current_chunk + cur_a); - - while (current_chunk < nchunk0 * nchunk1) { - const int64_t ith0 = current_chunk % nchunk0; - const int64_t ith1 = current_chunk / nchunk0; - - const int64_t ir0_start = dr0 * ith0; - const int64_t ir0_end = MIN(ir0_start + dr0, nr0); - - const int64_t ir1_start = dr1 * ith1; - const int64_t ir1_end = MIN(ir1_start + dr1, nr1); - - ggml_compute_forward_mul_mat_id_one_chunk( - dst, src0, src1, ids, cur_a, - ir0_start, ir0_end, ir1_start, ir1_end, - src0_cur, matrix_rows, row_size, src1_cont, wdata - ); - - if (nth >= nchunk0 * nchunk1) { - break; - } - - current_chunk = atomic_fetch_add_explicit(current_chunk_ctr, 1, memory_order_relaxed); - } - } -} - -///////////////////////////////// - -static void ggml_compute_forward(struct ggml_compute_params * params, struct ggml_tensor * tensor) { - GGML_ASSERT(params); - - if (tensor->op == GGML_OP_NONE || ggml_is_empty(tensor)) { - return; - } - - // extra_buffer op? - if (ggml_cpu_extra_compute_forward(params, tensor)) { - return; - } - - switch (tensor->op) { - case GGML_OP_DUP: - { - ggml_compute_forward_dup(params, tensor); - } break; - case GGML_OP_ADD: - { - ggml_compute_forward_add(params, tensor); - } break; - case GGML_OP_ADD_ID: - { - ggml_compute_forward_add_id(params, tensor); - } break; - case GGML_OP_ADD1: - { - ggml_compute_forward_add1(params, tensor); - } break; - case GGML_OP_ACC: - { - ggml_compute_forward_acc(params, tensor); - } break; - case GGML_OP_SUB: - { - ggml_compute_forward_sub(params, tensor); - } break; - case GGML_OP_MUL: - { - ggml_compute_forward_mul(params, tensor); - } break; - case GGML_OP_DIV: - { - ggml_compute_forward_div(params, tensor); - } break; - case GGML_OP_SQR: - { - ggml_compute_forward_sqr(params, tensor); - } break; - case GGML_OP_SQRT: - { - ggml_compute_forward_sqrt(params, tensor); - } break; - case GGML_OP_LOG: - { - ggml_compute_forward_log(params, tensor); - } break; - case GGML_OP_SIN: - { - ggml_compute_forward_sin(params, tensor); - } break; - case GGML_OP_COS: - { - ggml_compute_forward_cos(params, tensor); - } break; - case GGML_OP_SUM: - { - ggml_compute_forward_sum(params, tensor); - } break; - case GGML_OP_SUM_ROWS: - { - ggml_compute_forward_sum_rows(params, tensor); - } break; - case GGML_OP_CUMSUM: - { - ggml_compute_forward_cumsum(params, tensor); - } break; - case GGML_OP_MEAN: - { - ggml_compute_forward_mean(params, tensor); - } break; - case GGML_OP_ARGMAX: - { - ggml_compute_forward_argmax(params, tensor); - } break; - case GGML_OP_COUNT_EQUAL: - { - ggml_compute_forward_count_equal(params, tensor); - } break; - case GGML_OP_REPEAT: - { - ggml_compute_forward_repeat(params, tensor); - } break; - case GGML_OP_REPEAT_BACK: - { - ggml_compute_forward_repeat_back(params, tensor); - } break; - case GGML_OP_CONCAT: - { - ggml_compute_forward_concat(params, tensor); - } break; - case GGML_OP_SILU_BACK: - { - ggml_compute_forward_silu_back(params, tensor); - } break; - case GGML_OP_NORM: - { - ggml_compute_forward_norm(params, tensor); - } break; - case GGML_OP_RMS_NORM: - { - ggml_compute_forward_rms_norm(params, tensor); - } break; - case GGML_OP_RMS_NORM_BACK: - { - ggml_compute_forward_rms_norm_back(params, tensor); - } break; - case GGML_OP_GROUP_NORM: - { - ggml_compute_forward_group_norm(params, tensor); - } break; - case GGML_OP_L2_NORM: - { - ggml_compute_forward_l2_norm(params, tensor); - } break; - case GGML_OP_MUL_MAT: - { - ggml_compute_forward_mul_mat(params, tensor); - } break; - case GGML_OP_MUL_MAT_ID: - { - ggml_compute_forward_mul_mat_id(params, tensor); - } break; - case GGML_OP_MOE_FUSED: - { - GGML_ABORT("GGML_OP_MOE_FUSED is only implemented for CUDA"); - } - case GGML_OP_OUT_PROD: - { - ggml_compute_forward_out_prod(params, tensor); - } break; - case GGML_OP_SCALE: - { - ggml_compute_forward_scale(params, tensor); - } break; - case GGML_OP_SET: - { - ggml_compute_forward_set(params, tensor); - } break; - case GGML_OP_CPY: - { - ggml_compute_forward_cpy(params, tensor); - } break; - case GGML_OP_CONT: - { - ggml_compute_forward_cont(params, tensor); - } break; - case GGML_OP_GET_ROWS: - { - ggml_compute_forward_get_rows(params, tensor); - } break; - case GGML_OP_GET_ROWS_BACK: - { - ggml_compute_forward_get_rows_back(params, tensor); - } break; - case GGML_OP_SET_ROWS: - { - ggml_compute_forward_set_rows(params, tensor); - } break; - case GGML_OP_DIAG: - { - ggml_compute_forward_diag(params, tensor); - } break; - case GGML_OP_DIAG_MASK_INF: - { - ggml_compute_forward_diag_mask_inf(params, tensor); - } break; - case GGML_OP_DIAG_MASK_ZERO: - { - ggml_compute_forward_diag_mask_zero(params, tensor); - } break; - case GGML_OP_SOFT_MAX: - { - ggml_compute_forward_soft_max(params, tensor); - } break; - case GGML_OP_SOFT_MAX_BACK: - { - ggml_compute_forward_soft_max_ext_back(params, tensor); - } break; - case GGML_OP_ROPE: - { - ggml_compute_forward_rope(params, tensor); - } break; - case GGML_OP_ROPE_BACK: - { - ggml_compute_forward_rope_back(params, tensor); - } break; - case GGML_OP_CLAMP: - { - ggml_compute_forward_clamp(params, tensor); - } break; - case GGML_OP_CONV_TRANSPOSE_1D: - { - ggml_compute_forward_conv_transpose_1d(params, tensor); - } break; - case GGML_OP_IM2COL: - { - ggml_compute_forward_im2col(params, tensor); - } break; - case GGML_OP_IM2COL_BACK: - { - ggml_compute_forward_im2col_back_f32(params, tensor); - } break; - case GGML_OP_IM2COL_3D: - { - ggml_compute_forward_im2col_3d(params, tensor); - } break; - case GGML_OP_CONV_2D: - { - ggml_compute_forward_conv_2d(params, tensor); - } break; - case GGML_OP_CONV_3D: - { - ggml_compute_forward_conv_3d(params, tensor); - } break; - case GGML_OP_CONV_2D_DW: - { - ggml_compute_forward_conv_2d_dw(params, tensor); - } break; - case GGML_OP_CONV_TRANSPOSE_2D: - { - ggml_compute_forward_conv_transpose_2d(params, tensor); - } break; - case GGML_OP_POOL_1D: - { - ggml_compute_forward_pool_1d(params, tensor); - } break; - case GGML_OP_POOL_2D: - { - ggml_compute_forward_pool_2d(params, tensor); - } break; - case GGML_OP_POOL_2D_BACK: - { - ggml_compute_forward_pool_2d_back(params, tensor); - } break; - case GGML_OP_UPSCALE: - { - ggml_compute_forward_upscale(params, tensor); - } break; - case GGML_OP_PAD: - { - ggml_compute_forward_pad(params, tensor); - } break; - case GGML_OP_PAD_REFLECT_1D: - { - ggml_compute_forward_pad_reflect_1d(params, tensor); - } break; - case GGML_OP_ROLL: - { - ggml_compute_forward_roll(params, tensor); - } break; - case GGML_OP_ARANGE: - { - ggml_compute_forward_arange(params, tensor); - } break; - case GGML_OP_TIMESTEP_EMBEDDING: - { - ggml_compute_forward_timestep_embedding(params, tensor); - } break; - case GGML_OP_ARGSORT: - { - ggml_compute_forward_argsort(params, tensor); - } break; - case GGML_OP_TOP_K: - { - ggml_compute_forward_top_k(params, tensor); - } break; - case GGML_OP_LEAKY_RELU: - { - ggml_compute_forward_leaky_relu(params, tensor); - } break; - case GGML_OP_TRI: - { - ggml_compute_forward_tri(params, tensor); - } break; - case GGML_OP_FILL: - { - ggml_compute_forward_fill(params, tensor); - } break; - case GGML_OP_FLASH_ATTN_EXT: - { - ggml_compute_forward_flash_attn_ext(params, tensor); - } break; - case GGML_OP_FLASH_ATTN_SPARSE: - { - GGML_ABORT("GGML_OP_FLASH_ATTN_SPARSE is only supported on the CUDA backend"); - } - case GGML_OP_FLASH_ATTN_BACK: - { - int32_t t = ggml_get_op_params_i32(tensor, 0); - GGML_ASSERT(t == 0 || t == 1); - bool masked = t != 0; - ggml_compute_forward_flash_attn_back(params, masked, tensor); - } break; - case GGML_OP_SSM_CONV: - { - ggml_compute_forward_ssm_conv(params, tensor); - } break; - case GGML_OP_SSM_SCAN: - { - ggml_compute_forward_ssm_scan(params, tensor); - } break; - case GGML_OP_WIN_PART: - { - ggml_compute_forward_win_part(params, tensor); - } break; - case GGML_OP_WIN_UNPART: - { - ggml_compute_forward_win_unpart(params, tensor); - } break; - case GGML_OP_UNARY: - { - ggml_compute_forward_unary(params, tensor); - } break; - case GGML_OP_GLU: - { - ggml_compute_forward_glu(params, tensor); - } break; - case GGML_OP_TURBO_WHT: - { - // CUDA-only op; CPU fallback is identity (no rotation) - GGML_ASSERT(false && "TURBO_WHT is CUDA-only, use CUDA backend"); - } - case GGML_OP_GET_REL_POS: - { - ggml_compute_forward_get_rel_pos(params, tensor); - } break; - case GGML_OP_ADD_REL_POS: - { - ggml_compute_forward_add_rel_pos(params, tensor); - } break; - case GGML_OP_RWKV_WKV6: - { - ggml_compute_forward_rwkv_wkv6(params, tensor); - } break; - case GGML_OP_GATED_LINEAR_ATTN: - { - ggml_compute_forward_gla(params, tensor); - } break; - case GGML_OP_RWKV_WKV7: - { - ggml_compute_forward_rwkv_wkv7(params, tensor); - } break; - case GGML_OP_SOLVE_TRI: - { - ggml_compute_forward_solve_tri(params, tensor); - } break; - case GGML_OP_GATED_DELTA_NET: - { - ggml_compute_forward_gated_delta_net(params, tensor); - } break; - case GGML_OP_MAP_CUSTOM1: - { - ggml_compute_forward_map_custom1(params, tensor); - } - break; - case GGML_OP_MAP_CUSTOM2: - { - ggml_compute_forward_map_custom2(params, tensor); - } - break; - case GGML_OP_MAP_CUSTOM3: - { - ggml_compute_forward_map_custom3(params, tensor); - } - break; - case GGML_OP_CUSTOM: - { - ggml_compute_forward_custom(params, tensor); - } - break; - case GGML_OP_CROSS_ENTROPY_LOSS: - { - ggml_compute_forward_cross_entropy_loss(params, tensor); - } - break; - case GGML_OP_CROSS_ENTROPY_LOSS_BACK: - { - ggml_compute_forward_cross_entropy_loss_back(params, tensor); - } - break; - case GGML_OP_OPT_STEP_ADAMW: - { - ggml_compute_forward_opt_step_adamw(params, tensor); - } - break; - case GGML_OP_OPT_STEP_SGD: - { - ggml_compute_forward_opt_step_sgd(params, tensor); - } - break; - case GGML_OP_NONE: - { - // nop - } break; - case GGML_OP_RESHAPE: - { - // nop - } break; - case GGML_OP_PERMUTE: - { - // nop - } break; - case GGML_OP_VIEW: - { - // nop - } break; - case GGML_OP_TRANSPOSE: - { - // nop - } break; - case GGML_OP_COUNT: - { - GGML_ABORT("fatal error"); - } - } -} - -// Android's libc implementation "bionic" does not support setting affinity -#if defined(__gnu_linux__) -static void set_numa_thread_affinity(int thread_n) { - if (!ggml_is_numa()) { - return; - } - - int node_num; - int rv; - size_t setsize = CPU_ALLOC_SIZE(g_state.numa.total_cpus); - - switch(g_state.numa.numa_strategy) { - case GGML_NUMA_STRATEGY_DISTRIBUTE: - // run thread on node_num thread_n / (threads per node) - node_num = thread_n % g_state.numa.n_nodes; - break; - case GGML_NUMA_STRATEGY_ISOLATE: - // run thread on current_node - node_num = g_state.numa.current_node; - break; - case GGML_NUMA_STRATEGY_NUMACTL: - // use the cpuset that numactl gave us - rv = pthread_setaffinity_np(pthread_self(), setsize, &g_state.numa.cpuset); - if (rv) { - fprintf(stderr, "warning: pthread_setaffinity_np() failed: %s\n",strerror(rv)); - } - return; - default: - return; - } - - struct ggml_numa_node * node = &g_state.numa.nodes[node_num]; - - cpu_set_t * cpus = CPU_ALLOC(g_state.numa.total_cpus); - CPU_ZERO_S(setsize, cpus); - for (size_t i = 0; i < node->n_cpus; ++i) { - CPU_SET_S(node->cpus[i], setsize, cpus); - } - - rv = pthread_setaffinity_np(pthread_self(), setsize, cpus); - if (rv) { - fprintf(stderr, "warning: pthread_setaffinity_np() failed: %s\n", strerror(rv)); - } - - CPU_FREE(cpus); -} - -static void clear_numa_thread_affinity(void) { - if (!ggml_is_numa()) { - return; - } - - size_t setsize = CPU_ALLOC_SIZE(g_state.numa.total_cpus); - - cpu_set_t * cpus = CPU_ALLOC(g_state.numa.total_cpus); - CPU_ZERO_S(setsize, cpus); - for (unsigned i = 0; i < g_state.numa.total_cpus; ++i) { - CPU_SET_S(i, setsize, cpus); - } - - int rv = pthread_setaffinity_np(pthread_self(), setsize, cpus); - if (rv) { - fprintf(stderr, "warning: pthread_setaffinity_np() failed: %s\n", strerror(rv)); - } - - CPU_FREE(cpus); -} -#else -// TODO: Windows etc. -// (the linux implementation may also work on BSD, someone should test) -static void set_numa_thread_affinity(int thread_n) { UNUSED(thread_n); } -static void clear_numa_thread_affinity(void) {} -#endif - -static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) { - int n_tasks = 0; - - if (ggml_is_empty(node)) { - // no need to multi-thread a no-op - n_tasks = 1; - return n_tasks; - } - - switch (node->op) { - case GGML_OP_CPY: - case GGML_OP_DUP: - case GGML_OP_CONT: - case GGML_OP_ADD: - case GGML_OP_ADD_ID: - case GGML_OP_ADD1: - case GGML_OP_ACC: - case GGML_OP_CUMSUM: - case GGML_OP_TRI: - case GGML_OP_FILL: - { - n_tasks = n_threads; - } break; - case GGML_OP_SUB: - case GGML_OP_SQR: - case GGML_OP_SQRT: - case GGML_OP_LOG: - case GGML_OP_SIN: - case GGML_OP_COS: - case GGML_OP_SUM: - case GGML_OP_SUM_ROWS: - case GGML_OP_MEAN: - case GGML_OP_ARGMAX: - { - n_tasks = 1; - } break; - case GGML_OP_COUNT_EQUAL: - case GGML_OP_SOLVE_TRI: - case GGML_OP_GATED_DELTA_NET: - { - n_tasks = n_threads; - } break; - case GGML_OP_REPEAT: - case GGML_OP_REPEAT_BACK: - case GGML_OP_LEAKY_RELU: - { - n_tasks = 1; - } break; - case GGML_OP_UNARY: - switch (ggml_get_unary_op(node)) { - case GGML_UNARY_OP_ABS: - case GGML_UNARY_OP_SGN: - case GGML_UNARY_OP_NEG: - case GGML_UNARY_OP_STEP: - case GGML_UNARY_OP_TANH: - case GGML_UNARY_OP_ELU: - case GGML_UNARY_OP_RELU: - case GGML_UNARY_OP_SIGMOID: - case GGML_UNARY_OP_HARDSWISH: - case GGML_UNARY_OP_HARDSIGMOID: - case GGML_UNARY_OP_EXP: - case GGML_UNARY_OP_SOFTPLUS: - case GGML_UNARY_OP_EXPM1: - case GGML_UNARY_OP_FLOOR: - case GGML_UNARY_OP_CEIL: - case GGML_UNARY_OP_ROUND: - case GGML_UNARY_OP_TRUNC: - { - n_tasks = 1; - } break; - - case GGML_UNARY_OP_GELU: - case GGML_UNARY_OP_GELU_ERF: - case GGML_UNARY_OP_GELU_QUICK: - case GGML_UNARY_OP_SILU: - case GGML_UNARY_OP_XIELU: - { - n_tasks = n_threads; - } break; - default: - GGML_ABORT("fatal error"); - } - break; - case GGML_OP_GLU: - switch (ggml_get_glu_op(node)) { - case GGML_GLU_OP_REGLU: - case GGML_GLU_OP_GEGLU: - case GGML_GLU_OP_SWIGLU: - case GGML_GLU_OP_SWIGLU_OAI: - case GGML_GLU_OP_GEGLU_ERF: - case GGML_GLU_OP_GEGLU_QUICK: - { - n_tasks = n_threads; - } break; - default: - break; - } - break; - case GGML_OP_TURBO_WHT: - { - n_tasks = n_threads; - } break; - case GGML_OP_SILU_BACK: - case GGML_OP_MUL: - case GGML_OP_DIV: - case GGML_OP_NORM: - case GGML_OP_RMS_NORM: - case GGML_OP_RMS_NORM_BACK: - case GGML_OP_L2_NORM: - case GGML_OP_GROUP_NORM: - case GGML_OP_CONCAT: - case GGML_OP_MUL_MAT: - case GGML_OP_MUL_MAT_ID: - case GGML_OP_OUT_PROD: - { - n_tasks = n_threads; - } break; - case GGML_OP_GET_ROWS: - case GGML_OP_SET_ROWS: - { - // FIXME: get_rows can use additional threads, but the cost of launching additional threads - // decreases performance with GPU offloading - //n_tasks = n_threads; - n_tasks = 1; - } break; - case GGML_OP_SCALE: - case GGML_OP_SET: - case GGML_OP_RESHAPE: - case GGML_OP_VIEW: - case GGML_OP_PERMUTE: - case GGML_OP_TRANSPOSE: - case GGML_OP_GET_ROWS_BACK: - case GGML_OP_DIAG: - { - n_tasks = 1; - } break; - case GGML_OP_DIAG_MASK_ZERO: - case GGML_OP_DIAG_MASK_INF: - case GGML_OP_SOFT_MAX_BACK: - case GGML_OP_ROPE: - case GGML_OP_ROPE_BACK: - case GGML_OP_ADD_REL_POS: - { - n_tasks = n_threads; - } break; - case GGML_OP_CLAMP: - { - n_tasks = 1; //TODO - } break; - case GGML_OP_SOFT_MAX: - { - n_tasks = MIN(n_threads, ggml_nrows(node->src[0])); - } break; - case GGML_OP_IM2COL: - case GGML_OP_IM2COL_BACK: - case GGML_OP_IM2COL_3D: - case GGML_OP_CONV_2D: - case GGML_OP_CONV_3D: - case GGML_OP_CONV_2D_DW: - case GGML_OP_CONV_TRANSPOSE_1D: - case GGML_OP_CONV_TRANSPOSE_2D: - { - n_tasks = n_threads; - } break; - case GGML_OP_POOL_1D: - case GGML_OP_POOL_2D: - case GGML_OP_POOL_2D_BACK: - { - n_tasks = 1; - } break; - case GGML_OP_UPSCALE: - case GGML_OP_PAD: - case GGML_OP_PAD_REFLECT_1D: - case GGML_OP_ROLL: - case GGML_OP_ARANGE: - case GGML_OP_TIMESTEP_EMBEDDING: - case GGML_OP_ARGSORT: - case GGML_OP_TOP_K: - case GGML_OP_FLASH_ATTN_EXT: - case GGML_OP_FLASH_ATTN_SPARSE: - case GGML_OP_FLASH_ATTN_BACK: - case GGML_OP_SSM_CONV: - case GGML_OP_SSM_SCAN: - { - n_tasks = n_threads; - } break; - case GGML_OP_RWKV_WKV6: - case GGML_OP_GATED_LINEAR_ATTN: - case GGML_OP_RWKV_WKV7: - { - const int64_t n_heads = node->src[1]->ne[1]; - n_tasks = MIN(n_threads, n_heads); - } break; - case GGML_OP_WIN_PART: - case GGML_OP_WIN_UNPART: - case GGML_OP_GET_REL_POS: - { - n_tasks = 1; - } break; - case GGML_OP_MAP_CUSTOM1: - { - struct ggml_map_custom1_op_params p; - memcpy(&p, node->op_params, sizeof(p)); - if (p.n_tasks == GGML_N_TASKS_MAX) { - n_tasks = n_threads; - } else { - n_tasks = MIN(p.n_tasks, n_threads); - } - } break; - case GGML_OP_MAP_CUSTOM2: - { - struct ggml_map_custom2_op_params p; - memcpy(&p, node->op_params, sizeof(p)); - if (p.n_tasks == GGML_N_TASKS_MAX) { - n_tasks = n_threads; - } else { - n_tasks = MIN(p.n_tasks, n_threads); - } - } break; - case GGML_OP_MAP_CUSTOM3: - { - struct ggml_map_custom3_op_params p; - memcpy(&p, node->op_params, sizeof(p)); - if (p.n_tasks == GGML_N_TASKS_MAX) { - n_tasks = n_threads; - } else { - n_tasks = MIN(p.n_tasks, n_threads); - } - } break; - case GGML_OP_CUSTOM: - { - struct ggml_custom_op_params p; - memcpy(&p, node->op_params, sizeof(p)); - if (p.n_tasks == GGML_N_TASKS_MAX) { - n_tasks = n_threads; - } else { - n_tasks = MIN(p.n_tasks, n_threads); - } - } break; - case GGML_OP_CROSS_ENTROPY_LOSS: - case GGML_OP_CROSS_ENTROPY_LOSS_BACK: - case GGML_OP_OPT_STEP_ADAMW: - case GGML_OP_OPT_STEP_SGD: - { - n_tasks = n_threads; - } break; - case GGML_OP_NONE: - { - n_tasks = 1; - } break; - case GGML_OP_COUNT: - { - GGML_ABORT("fatal error"); - } - default: - { - fprintf(stderr, "%s: op not implemented: ", __func__); - if (node->op < GGML_OP_COUNT) { - fprintf(stderr, "%s\n", ggml_op_name(node->op)); - } else { - fprintf(stderr, "%d\n", node->op); - } - GGML_ABORT("fatal error"); - } - } - - assert(n_tasks > 0); - - return n_tasks; -} - -static thread_ret_t ggml_graph_compute_secondary_thread(void* data); - -#if defined(_WIN32) -#include "windows.h" - -// TODO: support > 64 CPUs -static bool ggml_thread_apply_affinity(bool * mask) { - HANDLE h = GetCurrentThread(); - uint64_t bitmask = 0ULL; - - assert(GGML_MAX_N_THREADS >= 64); - - for (int32_t i = 0; i < 8; i++) { - int32_t idx = i * 8; - uint8_t val = 0; - val |= mask[idx + 0] << 0; - val |= mask[idx + 1] << 1; - val |= mask[idx + 2] << 2; - val |= mask[idx + 3] << 3; - val |= mask[idx + 4] << 4; - val |= mask[idx + 5] << 5; - val |= mask[idx + 6] << 6; - val |= mask[idx + 7] << 7; - bitmask |= (uint64_t)val << idx; - } - - for (int32_t i = 64; i < GGML_MAX_N_THREADS; i++) { - if (mask[i]) { - fprintf(stderr, "warn: setting thread-affinity for > 64 CPUs isn't supported on windows!\n"); - break; - } - } - - DWORD_PTR m = (DWORD_PTR)bitmask; - - m = SetThreadAffinityMask(h, m); - - return m != 0; -} - -static bool ggml_thread_apply_priority(int32_t prio) { - // Note that on Windows the Process Priority Class must be updated in order to set Thread priority. - // This is up to the applications. - DWORD p = THREAD_PRIORITY_NORMAL; - switch (prio) { - case GGML_SCHED_PRIO_LOW: p = THREAD_PRIORITY_BELOW_NORMAL; break; - case GGML_SCHED_PRIO_NORMAL: p = THREAD_PRIORITY_NORMAL; break; - case GGML_SCHED_PRIO_MEDIUM: p = THREAD_PRIORITY_ABOVE_NORMAL; break; - case GGML_SCHED_PRIO_HIGH: p = THREAD_PRIORITY_HIGHEST; break; - case GGML_SCHED_PRIO_REALTIME: p = THREAD_PRIORITY_TIME_CRITICAL; break; - } - - if (prio != GGML_SCHED_PRIO_LOW) { - // Tell Windows that this thread should not be throttled (needs its own CPU core). - // Newer Windows 11 versions aggressively park (offline) CPU cores and often place - // all our threads onto the first 4 cores which results in terrible performance with - // n_threads > 4 - #if _WIN32_WINNT >= 0x0602 - THREAD_POWER_THROTTLING_STATE t; - ZeroMemory(&t, sizeof(t)); - t.Version = THREAD_POWER_THROTTLING_CURRENT_VERSION; - t.ControlMask = THREAD_POWER_THROTTLING_EXECUTION_SPEED; - t.StateMask = 0; - - if (!SetThreadInformation(GetCurrentThread(), ThreadPowerThrottling, &t, sizeof(t))) { - GGML_LOG_DEBUG("failed to disable thread power throttling %d : (%d)\n", prio, (int) GetLastError()); - return false; - } - #endif - } - - if (prio == GGML_SCHED_PRIO_NORMAL) { - // Keep inherited policy/priority - return true; - } - - if (!SetThreadPriority(GetCurrentThread(), p)) { - fprintf(stderr, "warn: failed to set thread priority %d : (%d)\n", prio, (int) GetLastError()); - return false; - } - - return true; -} - -#elif defined(__APPLE__) -#include -#include - -static bool ggml_thread_apply_affinity(const bool * mask) { - // Not supported on Apple platforms - UNUSED(mask); - return true; -} - -static bool ggml_thread_apply_priority(int32_t prio) { - struct sched_param p; - int32_t policy = SCHED_OTHER; - switch (prio) { - // TODO: there seems to be no way to set lower prio on Apple platforms - case GGML_SCHED_PRIO_LOW: policy = SCHED_OTHER; p.sched_priority = 0; break; - case GGML_SCHED_PRIO_NORMAL: policy = SCHED_OTHER; p.sched_priority = 0; break; - case GGML_SCHED_PRIO_MEDIUM: policy = SCHED_FIFO; p.sched_priority = 40; break; - case GGML_SCHED_PRIO_HIGH: policy = SCHED_FIFO; p.sched_priority = 80; break; - case GGML_SCHED_PRIO_REALTIME: policy = SCHED_FIFO; p.sched_priority = 90; break; - } - - if (prio == GGML_SCHED_PRIO_NORMAL) { - // Keep inherited policy/priority - return true; - } - - int32_t err = pthread_setschedparam(pthread_self(), policy, &p); - if (err != 0) { - fprintf(stderr, "warn: failed to set thread priority %d : %s (%d)\n", prio, strerror(err), err); - return false; - } - - return true; -} - -#elif defined(__gnu_linux__) -// TODO: this may not work on BSD, to be verified - -static bool ggml_thread_apply_affinity(const bool * mask) { - cpu_set_t cpuset; - int err; - - CPU_ZERO(&cpuset); - - for (uint32_t i = 0; i < GGML_MAX_N_THREADS; i++) { - if (mask[i]) { - GGML_PRINT_DEBUG("Thread %lx: adding %d to cpuset\n", pthread_self(), i); - CPU_SET(i, &cpuset); - } - } - -#ifdef __ANDROID__ - err = sched_setaffinity(0, sizeof(cpuset), &cpuset); - if (err < 0) { - err = errno; - } -#else - err = pthread_setaffinity_np(pthread_self(), sizeof(cpuset), &cpuset); -#endif - if (err != 0) { - fprintf(stderr, "warn: failed to set affinity mask 0x%llx : %s (%d)\n", (unsigned long long)mask, strerror(err), err); - return false; - } - - return true; -} - -static bool ggml_thread_apply_priority(int32_t prio) { - struct sched_param p; - int32_t policy = SCHED_OTHER; - switch (prio) { - case GGML_SCHED_PRIO_LOW: policy = SCHED_BATCH; p.sched_priority = 0; break; - case GGML_SCHED_PRIO_NORMAL: policy = SCHED_OTHER; p.sched_priority = 0; break; - case GGML_SCHED_PRIO_MEDIUM: policy = SCHED_FIFO; p.sched_priority = 40; break; - case GGML_SCHED_PRIO_HIGH: policy = SCHED_FIFO; p.sched_priority = 80; break; - case GGML_SCHED_PRIO_REALTIME: policy = SCHED_FIFO; p.sched_priority = 90; break; - } - - if (prio == GGML_SCHED_PRIO_NORMAL) { - // Keep inherited policy/priority - return true; - } - - int32_t err = pthread_setschedparam(pthread_self(), policy, &p); - if (err != 0) { - fprintf(stderr, "warn: failed to set thread priority %d : %s (%d)\n", prio, strerror(err), err); - return false; - } - - return true; -} - -#else // unsupported platforms - -static bool ggml_thread_apply_affinity(const bool * mask) { - UNUSED(mask); - return true; -} - -static bool ggml_thread_apply_priority(int32_t prio) { - UNUSED(prio); - return true; -} - -#endif - -static bool ggml_thread_cpumask_is_valid(const bool * mask) { - for (int i = 0; i < GGML_MAX_N_THREADS; i++) { - if (mask[i]) { return true; } - } - return false; -} - -static void ggml_thread_cpumask_next(const bool * global_mask, bool * local_mask, bool strict, int32_t* iter) { - if (!strict) { - memcpy(local_mask, global_mask, GGML_MAX_N_THREADS); - return; - } else { - memset(local_mask, 0, GGML_MAX_N_THREADS); - int32_t base_idx = *iter; - for (int32_t i = 0; i < GGML_MAX_N_THREADS; i++) { - int32_t idx = base_idx + i; - if (idx >= GGML_MAX_N_THREADS) { - // Just a cheaper modulo - idx -= GGML_MAX_N_THREADS; - } - if (global_mask[idx]) { - local_mask[idx] = 1; - *iter = idx + 1; - return; - } - } - } -} - -void ggml_threadpool_free(struct ggml_threadpool* threadpool) { - if (!threadpool) return; - - const int n_threads = threadpool->n_threads; - -#ifndef GGML_USE_OPENMP - struct ggml_compute_state* workers = threadpool->workers; - - ggml_mutex_lock(&threadpool->mutex); - - threadpool->stop = true; - threadpool->pause = false; - - ggml_cond_broadcast(&threadpool->cond); - ggml_mutex_unlock(&threadpool->mutex); - - for (int j = 1; j < n_threads; j++) { - int32_t rc = ggml_thread_join(workers[j].thrd, NULL); - GGML_ASSERT(rc == GGML_EXIT_SUCCESS || rc == GGML_EXIT_ABORTED); - UNUSED(rc); - } - - ggml_mutex_destroy(&threadpool->mutex); - ggml_cond_destroy(&threadpool->cond); -#endif // GGML_USE_OPENMP - - const size_t workers_size = sizeof(struct ggml_compute_state) * n_threads; - ggml_aligned_free(threadpool->workers, workers_size); - ggml_aligned_free(threadpool, sizeof(struct ggml_threadpool)); -} - -#ifndef GGML_USE_OPENMP -// pause/resume must be called under mutex -static void ggml_threadpool_pause_locked(struct ggml_threadpool * threadpool) { - GGML_PRINT_DEBUG("Pausing threadpool\n"); - threadpool->pause = true; - ggml_cond_broadcast(&threadpool->cond); -} - -static void ggml_threadpool_resume_locked(struct ggml_threadpool * threadpool) { - GGML_PRINT_DEBUG("Resuming threadpool\n"); - threadpool->pause = false; - ggml_cond_broadcast(&threadpool->cond); -} -#endif - -void ggml_threadpool_pause(struct ggml_threadpool * threadpool) { -#ifndef GGML_USE_OPENMP - ggml_mutex_lock(&threadpool->mutex); - if (!threadpool->pause) { - ggml_threadpool_pause_locked(threadpool); - } - ggml_mutex_unlock(&threadpool->mutex); -#else - UNUSED(threadpool); -#endif -} - -void ggml_threadpool_resume(struct ggml_threadpool * threadpool) { -#ifndef GGML_USE_OPENMP - ggml_mutex_lock(&threadpool->mutex); - if (threadpool->pause) { - ggml_threadpool_resume_locked(threadpool); - } - ggml_mutex_unlock(&threadpool->mutex); -#else - UNUSED(threadpool); -#endif -} - -struct ggml_cplan ggml_graph_plan( - const struct ggml_cgraph * cgraph, - int n_threads, - struct ggml_threadpool * threadpool) { - - if (threadpool == NULL) { - //GGML_PRINT_DEBUG("Threadpool is not specified. Will create a disposable threadpool : n_threads %d\n", n_threads); - } - if (n_threads <= 0) { - n_threads = threadpool ? threadpool->n_threads : GGML_DEFAULT_N_THREADS; - } - -#if defined(__EMSCRIPTEN__) && !defined(__EMSCRIPTEN_PTHREADS__) - // Emscripten without pthreads support can only use a single thread - n_threads = 1; -#endif - - size_t work_size = 0; - - struct ggml_cplan cplan; - memset(&cplan, 0, sizeof(struct ggml_cplan)); - - int max_tasks = 1; - - // thread scheduling for the different operations + work buffer size estimation - for (int i = 0; i < cgraph->n_nodes; i++) { - struct ggml_tensor * node = cgraph->nodes[i]; - - const int n_tasks = ggml_get_n_tasks(node, n_threads); - - max_tasks = MAX(max_tasks, n_tasks); - - size_t cur = 0; - - if (!ggml_cpu_extra_work_size(n_threads, node, &cur)) { - switch (node->op) { - case GGML_OP_CPY: - case GGML_OP_DUP: - { - if (ggml_is_quantized(node->type) || - // F16 -> BF16 and BF16 -> F16 copies go through intermediate F32 - (node->src[0]->type == GGML_TYPE_F16 && node->src[1] && node->src[1]->type == GGML_TYPE_BF16) || - (node->src[0]->type == GGML_TYPE_BF16 && node->src[1] && node->src[1]->type == GGML_TYPE_F16) || - // conversion between F32 and I32 - (node->src[0]->type == GGML_TYPE_F32 && node->src[1] && node->src[1]->type == GGML_TYPE_I32) || - (node->src[0]->type == GGML_TYPE_I32 && node->src[1] && node->src[1]->type == GGML_TYPE_F32)) { - cur = ggml_type_size(GGML_TYPE_F32) * node->ne[0] * n_tasks; - } - } break; - case GGML_OP_ADD: - case GGML_OP_ADD_ID: - case GGML_OP_ADD1: - { - if (ggml_is_quantized(node->src[0]->type)) { - cur = ggml_type_size(GGML_TYPE_F32) * node->src[0]->ne[0] * n_tasks; - } - } break; - case GGML_OP_ACC: - { - if (ggml_is_quantized(node->src[0]->type)) { - cur = ggml_type_size(GGML_TYPE_F32) * node->src[1]->ne[0] * n_tasks; - } - } break; - case GGML_OP_COUNT_EQUAL: - { - cur = ggml_type_size(node->type)*n_tasks; - } break; - case GGML_OP_MUL_MAT: - { - const enum ggml_type vec_dot_type = type_traits_cpu[node->src[0]->type].vec_dot_type; - - if (node->src[1]->type != vec_dot_type) { - cur = ggml_row_size(vec_dot_type, ggml_nelements(node->src[1])); - } - } break; - case GGML_OP_MUL_MAT_ID: - { - cur = 0; - const struct ggml_tensor * src0 = node->src[0]; - const struct ggml_tensor * src1 = node->src[1]; - const struct ggml_tensor * ids = node->src[2]; - const enum ggml_type vec_dot_type = type_traits_cpu[src0->type].vec_dot_type; - const int n_as = src0->ne[2]; - // src1 - if (src1->type != vec_dot_type) { - cur += ggml_row_size(vec_dot_type, ggml_nelements(src1)) + sizeof(int64_t); - } - // matrix_row_counts - cur += n_as * sizeof(int64_t) + sizeof(int64_t); - // matrix_rows - cur += n_as*ids->ne[0]*ids->ne[1]*sizeof(struct mmid_row_mapping) + sizeof(int64_t); - // atomic_current_chunk - cur += CACHE_LINE_SIZE*n_as + CACHE_LINE_SIZE; - } break; - case GGML_OP_OUT_PROD: - { - if (ggml_is_quantized(node->src[0]->type)) { - cur = ggml_type_size(GGML_TYPE_F32) * node->src[0]->ne[0] * n_tasks; - } - } break; - case GGML_OP_SOFT_MAX: - case GGML_OP_ROPE: - case GGML_OP_ROPE_BACK: - { - cur = ggml_type_size(GGML_TYPE_F32) * node->ne[0] * n_tasks; - } break; - case GGML_OP_CONV_TRANSPOSE_1D: - { - GGML_ASSERT(node->src[0]->ne[3] == 1); - GGML_ASSERT(node->src[1]->ne[2] == 1); - GGML_ASSERT(node->src[1]->ne[3] == 1); - - const int64_t ne00 = node->src[0]->ne[0]; // K - const int64_t ne01 = node->src[0]->ne[1]; // Cout - const int64_t ne02 = node->src[0]->ne[2]; // Cin - const int64_t ne10 = node->src[1]->ne[0]; // L - const int64_t ne11 = node->src[1]->ne[1]; // Cin - - if ((node->src[0]->type == GGML_TYPE_F16 || - node->src[0]->type == GGML_TYPE_BF16) && - node->src[1]->type == GGML_TYPE_F32) { - cur += sizeof(ggml_fp16_t)*ne00*ne01*ne02; - cur += sizeof(ggml_fp16_t)*ne10*ne11; - } else if (node->src[0]->type == GGML_TYPE_F32 && - node->src[1]->type == GGML_TYPE_F32) { - cur += sizeof(float)*ne00*ne01*ne02; - cur += sizeof(float)*ne10*ne11; - } else { - GGML_ABORT("fatal error"); - } - } break; - case GGML_OP_CONV_2D: - case GGML_OP_CONV_3D: - { - cur = GGML_IM2COL_WORK_SIZE; - } break; - case GGML_OP_CONV_TRANSPOSE_2D: - { - const int64_t ne00 = node->src[0]->ne[0]; // W - const int64_t ne01 = node->src[0]->ne[1]; // H - const int64_t ne02 = node->src[0]->ne[2]; // Channels Out - const int64_t ne03 = node->src[0]->ne[3]; // Channels In - - const int64_t ne10 = node->src[1]->ne[0]; // W - const int64_t ne11 = node->src[1]->ne[1]; // H - const int64_t ne12 = node->src[1]->ne[2]; // Channels In - - GGML_ASSERT(node->src[0]->type == GGML_TYPE_F16 || node->src[0]->type == GGML_TYPE_F32); - GGML_ASSERT(node->src[1]->type == GGML_TYPE_F32); - - cur += ggml_type_size(node->src[0]->type) * ne00 * ne01 * ne02 * ne03; - cur += ggml_type_size(node->src[0]->type) * ne10 * ne11 * ne12; - - } break; - case GGML_OP_TOP_K: - { - cur += sizeof(int32_t)*node->src[0]->ne[0]*n_tasks; - } break; - case GGML_OP_FLASH_ATTN_EXT: - { - const int64_t neq2 = node->src[0]->ne[2]; // number of query heads - const int64_t DK = node->src[1]->ne[0]; - const int64_t DV = node->src[2]->ne[0]; - - // Tiled flash attention scratch (tile sizes defined in common.h) - // Per-thread: Q_q + KQ + mask + VKQ32 + V32 + K_f32 + padding - size_t prefill = sizeof(float)*(GGML_FA_TILE_Q*DK + 2*GGML_FA_TILE_Q*GGML_FA_TILE_KV + GGML_FA_TILE_Q*DV + GGML_FA_TILE_KV*DV + GGML_FA_TILE_KV*DK)*n_tasks; - - // Decode path: n_kv_chunks = n_tasks (one chunk per thread) - // Per-thread: VKQ accmulator (DV), partial M, partial S + intra-thread scratch for V, Q and VKQ - size_t n_chunks = n_tasks; - size_t decode = sizeof(float)*(neq2*n_chunks*(2+DV) + n_tasks*(DK + 2*DV)); - - cur += MAX(prefill, decode); - } break; - case GGML_OP_FLASH_ATTN_BACK: - { - const int64_t D = node->src[0]->ne[0]; - const int64_t ne11 = ggml_up(node->src[1]->ne[1], GGML_SOFT_MAX_UNROLL); - const int64_t mxDn = MAX(D, ne11) * 2; // *2 because of S and SM in ggml_compute_forward_flash_attn_back - if (node->src[1]->type == GGML_TYPE_F32) { - cur = sizeof(float)*mxDn*n_tasks; // TODO: this can become (n_tasks-1) - cur += sizeof(float)*mxDn*n_tasks; // this is overestimated by x2 - } else if (node->src[1]->type == GGML_TYPE_F16) { - cur = sizeof(float)*mxDn*n_tasks; // TODO: this can become (n_tasks-1) - cur += sizeof(float)*mxDn*n_tasks; // this is overestimated by x2 - } else if (node->src[1]->type == GGML_TYPE_BF16) { - cur = sizeof(float)*mxDn*n_tasks; // TODO: this can become (n_tasks-1) - cur += sizeof(float)*mxDn*n_tasks; // this is overestimated by x2 - } - } break; - - case GGML_OP_CROSS_ENTROPY_LOSS: - { - cur = ggml_type_size(node->type)*(n_tasks + node->src[0]->ne[0]*n_tasks); - } break; - case GGML_OP_GATED_DELTA_NET: - { - const int64_t S_v = node->src[2]->ne[0]; - cur = S_v * sizeof(float) * n_tasks; - } break; - case GGML_OP_COUNT: - { - GGML_ABORT("fatal error"); - } - default: - break; - } - } - - work_size = MAX(work_size, cur); - } - - if (work_size > 0) { - work_size += CACHE_LINE_SIZE*(n_threads); - } - - cplan.threadpool = threadpool; - cplan.n_threads = MIN(max_tasks, n_threads); - cplan.work_size = work_size; - cplan.work_data = NULL; - - return cplan; -} - -static thread_ret_t ggml_graph_compute_thread(void * data) { - struct ggml_compute_state * state = (struct ggml_compute_state *) data; - struct ggml_threadpool * tp = state->threadpool; - - const struct ggml_cgraph * cgraph = tp->cgraph; - const struct ggml_cplan * cplan = tp->cplan; - - set_numa_thread_affinity(state->ith); - - struct ggml_compute_params params = { - /*.ith =*/ state->ith, - /*.nth =*/ atomic_load_explicit(&tp->n_graph, memory_order_relaxed) & GGML_THREADPOOL_N_THREADS_MASK, - /*.wsize =*/ cplan->work_size, - /*.wdata =*/ cplan->work_data, - /*.threadpool =*/ tp, - /*.use_ref =*/ cplan->use_ref, - }; - -#ifdef GGML_USE_OPENMP - GGML_PRINT_DEBUG("thread #%d compute-start cplan %p\n", state->ith, (const void *)cplan); -#else - GGML_PRINT_DEBUG("thread #%d compute-start cplan %p last-graph %d\n", state->ith, (const void *)cplan, state->last_graph); -#endif - - for (int node_n = 0; node_n < cgraph->n_nodes && atomic_load_explicit(&tp->abort, memory_order_relaxed) != node_n; node_n++) { - struct ggml_tensor * node = cgraph->nodes[node_n]; - - if (ggml_op_is_empty(node->op)) { - // skip NOPs - continue; - } - - if ((node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { - continue; - } - - ggml_compute_forward(¶ms, node); - - if (state->ith == 0 && cplan->abort_callback && - cplan->abort_callback(cplan->abort_callback_data)) { - atomic_store_explicit(&tp->abort, node_n + 1, memory_order_relaxed); - tp->ec = GGML_STATUS_ABORTED; - } - - if (node_n + 1 < cgraph->n_nodes) { - ggml_barrier(state->threadpool); - } - } - -#ifdef GGML_USE_OPENMP - GGML_PRINT_DEBUG("thread #%d compute-done cplan %p\n", state->ith, (const void *)cplan); -#else - GGML_PRINT_DEBUG("thread #%d compute-done cplan %p last-graph %d\n", state->ith, (const void *)cplan, state->last_graph); -#endif - - ggml_barrier(state->threadpool); - - return 0; -} - -#ifndef GGML_USE_OPENMP - -// check if thread is ready to proceed (exit from polling or sleeping) -// returns true if loops should exit, sets state->pending to indicate new work -static inline bool ggml_graph_compute_thread_ready(struct ggml_compute_state * state) { - struct ggml_threadpool * threadpool = state->threadpool; - - if (state->pending || threadpool->stop || threadpool->pause) { return true; } - - // check for new graph/work - int n_graph = atomic_load_explicit(&threadpool->n_graph, memory_order_relaxed); - int n_threads = n_graph & GGML_THREADPOOL_N_THREADS_MASK; - if (n_graph != state->last_graph) { - state->pending = (state->ith < n_threads); - state->last_graph = n_graph; - return true; - } - - return false; -} - -// sync thread state after polling -static inline void ggml_graph_compute_thread_sync(struct ggml_compute_state * state) { - // TSAN doesn't support standalone fence yet, we use a dummy read-modify-write instead - #ifdef GGML_TSAN_ENABLED - atomic_fetch_add_explicit(&state->threadpool->n_graph, 0, memory_order_seq_cst); - #else - atomic_thread_fence(memory_order_seq_cst); - #endif - UNUSED(state); -} - -static inline bool ggml_graph_compute_poll_for_work(struct ggml_compute_state * state) { - struct ggml_threadpool * threadpool = state->threadpool; - - // This seems to make 0 ... 100 a decent range for polling level across modern processors. - // Perhaps, we can adjust it dynamically based on load and things. - const uint64_t n_rounds = 1024UL * 128 * threadpool->poll; - - for (uint64_t i=0; !ggml_graph_compute_thread_ready(state) && i < n_rounds; i++) { - // No new work. Keep polling. - ggml_thread_cpu_relax(); - } - - return state->pending; -} - -static inline bool ggml_graph_compute_check_for_work(struct ggml_compute_state * state) { - struct ggml_threadpool * threadpool = state->threadpool; - - if (ggml_graph_compute_poll_for_work(state)) { - ggml_graph_compute_thread_sync(state); - return state->pending; - } - - ggml_mutex_lock_shared(&threadpool->mutex); - while (!ggml_graph_compute_thread_ready(state)) { - // No new work. Wait for the signal. - GGML_PRINT_DEBUG("thread #%d waiting for work (sleeping)\n", state->ith); - ggml_cond_wait(&threadpool->cond, &threadpool->mutex); - } - ggml_mutex_unlock_shared(&threadpool->mutex); - - return state->pending; -} - -static thread_ret_t ggml_graph_compute_secondary_thread(void* data) { - struct ggml_compute_state * state = (struct ggml_compute_state *) data; - struct ggml_threadpool * threadpool = state->threadpool; - - ggml_thread_apply_priority(threadpool->prio); - if (ggml_thread_cpumask_is_valid(state->cpumask)) { - ggml_thread_apply_affinity(state->cpumask); - } - - while (true) { - // Check if we need to sleep - while (threadpool->pause) { - GGML_PRINT_DEBUG("thread #%d inside pause loop\n", state->ith); - ggml_mutex_lock_shared(&threadpool->mutex); - if (threadpool->pause) { - ggml_cond_wait(&threadpool->cond, &threadpool->mutex); - } - GGML_PRINT_DEBUG("thread #%d resuming after wait\n", state->ith); - ggml_mutex_unlock_shared(&threadpool->mutex); - } - - // This needs to be checked for after the cond_wait - if (threadpool->stop) break; - - // Check if there is new work - // The main thread is the only one that can dispatch new work - - ggml_graph_compute_check_for_work(state); - if (state->pending) { - state->pending = false; - ggml_graph_compute_thread(state); - } - } - - return (thread_ret_t) 0; -} - -// Start processing new graph -static void ggml_graph_compute_kickoff(struct ggml_threadpool * threadpool, int n_threads) -{ - // Always take the mutex here because the worker threads are doing hybrid poll/wait - - ggml_mutex_lock(&threadpool->mutex); - - // Update the number of active threads and the graph count - int n_graph = atomic_load_explicit(&threadpool->n_graph, memory_order_relaxed) >> GGML_THREADPOOL_N_THREADS_BITS; - n_graph = ((n_graph + 1) << GGML_THREADPOOL_N_THREADS_BITS) | (n_threads & GGML_THREADPOOL_N_THREADS_MASK); - - GGML_PRINT_DEBUG("compute-kickoff: n_threads %d n_graph %d\n", n_threads, n_graph); - - // Indicate the graph is ready to be processed - // We need the full seq-cst fence here because of the polling threads (used in thread_sync) - atomic_store_explicit(&threadpool->n_graph, n_graph, memory_order_seq_cst); - - if (threadpool->pause) { - // Update main thread prio and affinity to match the threadpool settings - ggml_thread_apply_priority(threadpool->prio); - if (ggml_thread_cpumask_is_valid(threadpool->workers[0].cpumask)) { - ggml_thread_apply_affinity(threadpool->workers[0].cpumask); - } - - // resume does cond broadcast - ggml_threadpool_resume_locked(threadpool); - } else { - ggml_cond_broadcast(&threadpool->cond); - } - - ggml_mutex_unlock(&threadpool->mutex); -} - -#endif // GGML_USE_OPENMP - -static struct ggml_threadpool * ggml_threadpool_new_impl( - struct ggml_threadpool_params * tpp, - struct ggml_cgraph * cgraph, - struct ggml_cplan * cplan) { - - struct ggml_threadpool * threadpool = - ggml_aligned_malloc(sizeof(struct ggml_threadpool)); - { - threadpool->cgraph = cgraph; - threadpool->cplan = cplan; - threadpool->n_graph = 0; - threadpool->n_barrier = 0; - threadpool->n_barrier_passed = 0; - threadpool->current_chunk = 0; - threadpool->stop = false; - threadpool->pause = tpp->paused; - threadpool->abort = -1; - threadpool->workers = NULL; - threadpool->n_threads = tpp->n_threads; - threadpool->poll = tpp->poll; - threadpool->prio = tpp->prio; - threadpool->ec = GGML_STATUS_SUCCESS; - } - - // Allocate and init workers state - const size_t workers_size = sizeof(struct ggml_compute_state) * tpp->n_threads; - struct ggml_compute_state * workers = ggml_aligned_malloc(workers_size); - - memset(workers, 0, workers_size); - for (int j = 0; j < tpp->n_threads; j++) { - workers[j].threadpool = threadpool; - workers[j].ith = j; - } - - threadpool->workers = workers; - -#ifdef GGML_USE_OPENMP - int32_t cpumask_iter = 0; - - // Compute CPU masks for each thread - for (int j = 0; j < tpp->n_threads; j++) { - ggml_thread_cpumask_next(tpp->cpumask, workers[j].cpumask, tpp->strict_cpu, &cpumask_iter); - } -#else // GGML_USE_OPENMP - ggml_mutex_init(&threadpool->mutex); - ggml_cond_init(&threadpool->cond); - - // Spin the threads for all workers, and update CPU placements. - // Place the main thread last (towards the higher numbered CPU cores). - - int32_t cpumask_iter = 0; - - for (int j = 1; j < tpp->n_threads; j++) { - ggml_thread_cpumask_next(tpp->cpumask, workers[j].cpumask, tpp->strict_cpu, &cpumask_iter); - - int32_t rc = ggml_thread_create(&workers[j].thrd, NULL, ggml_graph_compute_secondary_thread, &workers[j]); - GGML_ASSERT(rc == 0); - } - - ggml_thread_cpumask_next(tpp->cpumask, workers[0].cpumask, tpp->strict_cpu, &cpumask_iter); - - if (!threadpool->pause) { - // Update main thread prio and affinity at the start, otherwise we'll do it in resume - ggml_thread_apply_priority(threadpool->prio); - if (ggml_thread_cpumask_is_valid(threadpool->workers[0].cpumask)) { - ggml_thread_apply_affinity(threadpool->workers[0].cpumask); - } - } -#endif // GGML_USE_OPENMP - - return threadpool; -} - -struct ggml_threadpool * ggml_threadpool_new(struct ggml_threadpool_params * tpp) { - return ggml_threadpool_new_impl(tpp, NULL, NULL); -} - -enum ggml_status ggml_graph_compute(struct ggml_cgraph * cgraph, struct ggml_cplan * cplan) { - ggml_cpu_init(); - - GGML_ASSERT(cplan); - GGML_ASSERT(cplan->n_threads > 0); - GGML_ASSERT(cplan->work_size == 0 || cplan->work_data != NULL); - - int n_threads = cplan->n_threads; - struct ggml_threadpool * threadpool = cplan->threadpool; - - bool disposable_threadpool = false; - - if (threadpool == NULL) { - //GGML_PRINT_DEBUG("Threadpool is not specified. Will create a disposable threadpool : n_threads %d\n", n_threads); - disposable_threadpool = true; - - struct ggml_threadpool_params ttp = ggml_threadpool_params_default(n_threads); - threadpool = ggml_threadpool_new_impl(&ttp, cgraph, cplan); - } else { - // Reset some of the parameters that need resetting - // No worker threads should be accessing the parameters below at this stage - threadpool->cgraph = cgraph; - threadpool->cplan = cplan; - threadpool->current_chunk = 0; - threadpool->abort = -1; - threadpool->ec = GGML_STATUS_SUCCESS; - } - -#ifdef GGML_USE_OPENMP - if (n_threads > 1) { - #pragma omp parallel num_threads(n_threads) - { - #pragma omp single - { - // update the number of threads from the actual number of threads that we got from OpenMP - n_threads = omp_get_num_threads(); - atomic_store_explicit(&threadpool->n_graph, n_threads, memory_order_relaxed); - } - - // Apply thread CPU mask and priority - int ith = omp_get_thread_num(); - - ggml_thread_apply_priority(threadpool->prio); - if (ggml_thread_cpumask_is_valid(threadpool->workers[ith].cpumask)) { - ggml_thread_apply_affinity(threadpool->workers[ith].cpumask); - } - ggml_graph_compute_thread(&threadpool->workers[ith]); - } - } else { - atomic_store_explicit(&threadpool->n_graph, 1, memory_order_relaxed); - ggml_graph_compute_thread(&threadpool->workers[0]); - } -#else - if (n_threads > threadpool->n_threads) { - GGML_LOG_WARN("cplan requested more threads (%d) than available (%d)\n", n_threads, threadpool->n_threads); - n_threads = threadpool->n_threads; - } - - // Kick all threads to start the new graph - ggml_graph_compute_kickoff(threadpool, n_threads); - - // This is a work thread too - ggml_graph_compute_thread(&threadpool->workers[0]); -#endif - - // don't leave affinity set on the main thread - clear_numa_thread_affinity(); - - enum ggml_status ret = threadpool->ec; - - if (disposable_threadpool) { - ggml_threadpool_free(threadpool); - } - - return ret; -} - -enum ggml_status ggml_graph_compute_with_ctx(struct ggml_context * ctx, struct ggml_cgraph * cgraph, int n_threads) { - struct ggml_cplan cplan = ggml_graph_plan(cgraph, n_threads, NULL); - - cplan.work_data = (uint8_t *)ggml_new_buffer(ctx, cplan.work_size); - - return ggml_graph_compute(cgraph, &cplan); -} - -void ggml_cpu_fp32_to_fp32(const float * x, float * y, int64_t n) { - memcpy(y, x, n * sizeof(float)); -} - -void ggml_cpu_fp32_to_fp16(const float * x, ggml_fp16_t * y, int64_t n) { - int64_t i = 0; -#if defined(__F16C__) -#if defined(__AVX512F__) - for (; i + 15 < n; i += 16) { - __m512 x_vec = _mm512_loadu_ps(x + i); - __m256i y_vec = _mm512_cvtps_ph(x_vec, _MM_FROUND_TO_NEAREST_INT); - _mm256_storeu_si256((__m256i *)(y + i), y_vec); - } -#endif - for (; i + 7 < n; i += 8) { - __m256 x_vec = _mm256_loadu_ps(x + i); - __m128i y_vec = _mm256_cvtps_ph(x_vec, _MM_FROUND_TO_NEAREST_INT); - _mm_storeu_si128((__m128i *)(y + i), y_vec); - } - for (; i + 3 < n; i += 4) { - __m128 x_vec = _mm_loadu_ps(x + i); - __m128i y_vec = _mm_cvtps_ph(x_vec, _MM_FROUND_TO_NEAREST_INT); - _mm_storel_epi64((__m128i *)(y + i), y_vec); - } -#elif defined(__riscv_zvfh) - for (int vl; i < n; i += vl) { - vl = __riscv_vsetvl_e32m2(n - i); - vfloat32m2_t vx = __riscv_vle32_v_f32m2(&x[i], vl); - vfloat16m1_t vy = __riscv_vfncvt_f_f_w_f16m1(vx, vl); - __riscv_vse16_v_f16m1((_Float16 *)&y[i], vy, vl); - } -#endif - for (; i < n; ++i) { - y[i] = GGML_CPU_FP32_TO_FP16(x[i]); - } -} - -void ggml_cpu_fp16_to_fp32(const ggml_fp16_t * x, float * y, int64_t n) { - int64_t i = 0; -#if defined(__F16C__) -#if defined(__AVX512F__) - for (; i + 15 < n; i += 16) { - __m256i x_vec = _mm256_loadu_si256((const __m256i *)(x + i)); - __m512 y_vec = _mm512_cvtph_ps(x_vec); - _mm512_storeu_ps(y + i, y_vec); - } -#endif - for (; i + 7 < n; i += 8) { - __m128i x_vec = _mm_loadu_si128((const __m128i *)(x + i)); - __m256 y_vec = _mm256_cvtph_ps(x_vec); - _mm256_storeu_ps(y + i, y_vec); - } - for (; i + 3 < n; i += 4) { - __m128i x_vec = _mm_loadl_epi64((const __m128i *)(x + i)); - __m128 y_vec = _mm_cvtph_ps(x_vec); - _mm_storeu_ps(y + i, y_vec); - } - -#elif defined(__riscv_v_intrinsic) && defined(__riscv_zvfhmin) - // calculate step size - const int epr = __riscv_vsetvlmax_e16m2(); - const int step = epr * 2; - const int np = (n & ~(step - 1)); - - // unroll by 2 - for (; i < np; i += step) { - vfloat16m2_t ax0 = __riscv_vle16_v_f16m2((const _Float16*)x + i, epr); - vfloat32m4_t ay0 = __riscv_vfwcvt_f_f_v_f32m4(ax0, epr); - __riscv_vse32_v_f32m4(y + i, ay0, epr); - - vfloat16m2_t ax1 = __riscv_vle16_v_f16m2((const _Float16*)x + i + epr, epr); - vfloat32m4_t ay1 = __riscv_vfwcvt_f_f_v_f32m4(ax1, epr); - __riscv_vse32_v_f32m4(y + i + epr, ay1, epr); - } - - // leftovers - int vl; - for (i = np; i < n; i += vl) { - vl = __riscv_vsetvl_e16m2(n - i); - vfloat16m2_t ax0 = __riscv_vle16_v_f16m2((const _Float16*)x + i, vl); - vfloat32m4_t ay0 = __riscv_vfwcvt_f_f_v_f32m4(ax0, vl); - __riscv_vse32_v_f32m4(y + i, ay0, vl); - } - -#endif - - for (; i < n; ++i) { - y[i] = GGML_CPU_FP16_TO_FP32(x[i]); - } -} - -void ggml_cpu_fp32_to_bf16(const float * x, ggml_bf16_t * y, int64_t n) { - int64_t i = 0; - for (; i < n; ++i) { - y[i] = GGML_FP32_TO_BF16(x[i]); - } -} - -void ggml_cpu_fp32_to_i32(const float * x, int32_t * y, int64_t n) { - int64_t i = 0; - for (; i < n; ++i) { - y[i] = x[i]; - } -} - -void ggml_cpu_bf16_to_fp32(const ggml_bf16_t * x, float * y, int64_t n) { - int64_t i = 0; -#if defined(__AVX2__) -#if defined(__AVX512F__) - for (; i + 15 < n; i += 16) { - _mm512_storeu_ps(y + i, - _mm512_castsi512_ps( - _mm512_slli_epi32( - _mm512_cvtepu16_epi32( - _mm256_loadu_si256( - (const __m256i *)(x + i))), - 16))); - } -#endif - for (; i + 7 < n; i += 8) { - _mm256_storeu_ps(y + i, - _mm256_castsi256_ps( - _mm256_slli_epi32( - _mm256_cvtepu16_epi32( - _mm_loadu_si128( - (const __m128i *)(x + i))), - 16))); - } -#elif defined(__riscv_v_intrinsic) && defined(__riscv_zvfbfmin) - // calculate step size - const int epr = __riscv_vsetvlmax_e16m2(); - const int step = epr * 2; - const int np = (n & ~(step - 1)); - - // unroll by 2 - for (; i < np; i += step) { - vbfloat16m2_t ax0 = __riscv_vle16_v_bf16m2((const __bf16*)x + i, epr); - vfloat32m4_t ay0 = __riscv_vfwcvtbf16_f_f_v_f32m4(ax0, epr); - __riscv_vse32_v_f32m4(y + i, ay0, epr); - - vbfloat16m2_t ax1 = __riscv_vle16_v_bf16m2((const __bf16*)x + i + epr, epr); - vfloat32m4_t ay1 = __riscv_vfwcvtbf16_f_f_v_f32m4(ax1, epr); - __riscv_vse32_v_f32m4(y + i + epr, ay1, epr); - } - - // leftovers - int vl; - for (i = np; i < n; i += vl) { - vl = __riscv_vsetvl_e16m2(n - i); - vbfloat16m2_t ax0 = __riscv_vle16_v_bf16m2((const __bf16*)x + i, vl); - vfloat32m4_t ay0 = __riscv_vfwcvtbf16_f_f_v_f32m4(ax0, vl); - __riscv_vse32_v_f32m4(y + i, ay0, vl); - } -#endif - for (; i < n; i++) { - y[i] = GGML_BF16_TO_FP32(x[i]); - } -} - -int ggml_cpu_has_avx(void) { -#if defined(__AVX__) - return 1; -#else - return 0; -#endif -} - -int ggml_cpu_has_avx_vnni(void) { -#if defined(__AVXVNNI__) - return 1; -#else - return 0; -#endif -} - -int ggml_cpu_has_avx2(void) { -#if defined(__AVX2__) - return 1; -#else - return 0; -#endif -} - -int ggml_cpu_has_avx512(void) { -#if defined(__AVX512F__) - return 1; -#else - return 0; -#endif -} - -int ggml_cpu_has_avx512_vbmi(void) { -#if defined(__AVX512VBMI__) - return 1; -#else - return 0; -#endif -} - -int ggml_cpu_has_avx512_vnni(void) { -#if defined(__AVX512VNNI__) - return 1; -#else - return 0; -#endif -} - -int ggml_cpu_has_avx512_bf16(void) { -#if defined(__AVX512BF16__) - return 1; -#else - return 0; -#endif -} - -int ggml_cpu_has_amx_int8(void) { -#if defined(__AMX_INT8__) - return 1; -#else - return 0; -#endif -} - -int ggml_cpu_has_bmi2(void) { -#if defined(__BMI2__) - return 1; -#else - return 0; -#endif -} - -int ggml_cpu_has_fma(void) { -#if defined(__FMA__) - return 1; -#else - return 0; -#endif -} - -int ggml_cpu_has_arm_fma(void) { -#if defined(__ARM_FEATURE_FMA) - return 1; -#else - return 0; -#endif -} - -int ggml_cpu_has_riscv_v(void) { -#if defined(__riscv_v_intrinsic) - return 1; -#else - return 0; -#endif -} - -int ggml_cpu_get_rvv_vlen(void) { -#if defined(__riscv) && defined(__riscv_v_intrinsic) - return ggml_riscv_arch_features.rvv_vlen; -#else - return 0; -#endif -} - -int ggml_cpu_has_f16c(void) { -#if defined(__F16C__) - return 1; -#else - return 0; -#endif -} - -int ggml_cpu_has_fp16_va(void) { -#if defined(__ARM_FEATURE_FP16_VECTOR_ARITHMETIC) - return 1; -#else - return 0; -#endif -} - -int ggml_cpu_has_wasm_simd(void) { -#if defined(__wasm_simd128__) - return 1; -#else - return 0; -#endif -} - -int ggml_cpu_has_llamafile(void) { -#if defined(GGML_USE_LLAMAFILE) - return 1; -#else - return 0; -#endif -} - -int ggml_cpu_has_sse3(void) { -#if defined(__SSE3__) - return 1; -#else - return 0; -#endif -} - -int ggml_cpu_has_ssse3(void) { -#if defined(__SSSE3__) - return 1; -#else - return 0; -#endif -} - -int ggml_cpu_has_vsx(void) { -#if defined(__POWER9_VECTOR__) - return 1; -#else - return 0; -#endif -} - -int ggml_cpu_has_vxe(void) { -#if defined(__VXE__) || defined(__VXE2__) - return 1; -#else - return 0; -#endif -} - -int ggml_cpu_has_neon(void) { -#if defined(__ARM_ARCH) && defined(__ARM_NEON) - return 1; -#else - return 0; -#endif -} - -int ggml_cpu_has_dotprod(void) { -#if defined(__ARM_ARCH) && defined(__ARM_FEATURE_DOTPROD) - return 1; -#else - return 0; -#endif -} - -int ggml_cpu_has_sve(void) { -#if defined(__ARM_ARCH) && defined(__ARM_FEATURE_SVE) - return 1; -#else - return 0; -#endif -} - -int ggml_cpu_has_matmul_int8(void) { -#if defined(__ARM_ARCH) && defined(__ARM_FEATURE_MATMUL_INT8) - return 1; -#else - return 0; -#endif -} - -int ggml_cpu_get_sve_cnt(void) { -#if defined(__ARM_ARCH) && defined(__ARM_FEATURE_SVE) - return ggml_arm_arch_features.sve_cnt; -#else - return 0; -#endif -} - -int ggml_cpu_has_sme(void) { -#if defined(__ARM_ARCH) && defined(__ARM_FEATURE_SME) - return 1; -#else - return 0; -#endif -} - -void ggml_cpu_init(void) { - // needed to initialize ggml_time - { - struct ggml_init_params params = { 0, NULL, false }; - struct ggml_context * ctx = ggml_init(params); - ggml_free(ctx); - } - - ggml_critical_section_start(); - - static bool is_first_call = true; - - if (is_first_call) { - // initialize GELU, Quick GELU, SILU and EXP F32 tables - { - const uint64_t t_start = ggml_time_us(); UNUSED(t_start); - - for (int i = 0; i < (1 << 16); ++i) { - union { - uint16_t u16; - ggml_fp16_t fp16; - } u = {i}; - float f = GGML_COMPUTE_FP16_TO_FP32(u.fp16); - ggml_table_f32_f16[i] = f; - ggml_table_gelu_f16[i] = GGML_CPU_FP32_TO_FP16(ggml_gelu_f32(f)); - ggml_table_gelu_quick_f16[i] = GGML_CPU_FP32_TO_FP16(ggml_gelu_quick_f32(f)); - } - - // initialize E8M0 half table (256 entries) - for (int i = 0; i < (1 << 8); ++i) { - ggml_table_f32_e8m0_half[i] = GGML_E8M0_TO_FP32_HALF(i); - } - - const uint64_t t_end = ggml_time_us(); UNUSED(t_end); - - GGML_PRINT_DEBUG("%s: GELU, Quick GELU, SILU and EXP tables initialized in %f ms\n", __func__, (t_end - t_start)/1000.0); - -#ifdef GGML_USE_OPENMP - //if (!getenv("OMP_WAIT_POLICY")) { - // // set the wait policy to active, so that OpenMP threads don't sleep - // setenv("OMP_WAIT_POLICY", "active", 0) - //} - - if (!getenv("KMP_BLOCKTIME")) { - // set the time to wait before sleeping a thread - // this is less aggressive than setting the wait policy to active, but should achieve similar results in most cases -#ifdef _WIN32 - _putenv_s("KMP_BLOCKTIME", "200"); // 200ms -#else - setenv("KMP_BLOCKTIME", "200", 0); // 200ms -#endif - } -#endif - } - -#if defined(__ARM_ARCH) - ggml_init_arm_arch_features(); -#endif - -#if defined(__riscv) - ggml_init_riscv_arch_features(); -#endif - - is_first_call = false; - } - - ggml_critical_section_end(); -} diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu.ds4bak b/ggml/src/ggml-cuda/ggml-cuda.cu.ds4bak deleted file mode 100644 index 066dd2680d89..000000000000 --- a/ggml/src/ggml-cuda/ggml-cuda.cu.ds4bak +++ /dev/null @@ -1,5559 +0,0 @@ -#include "ggml-cuda.h" -#include "ggml-impl.h" -#include "ggml-backend-impl.h" - -#include "ggml-cuda/common.cuh" -#include "ggml-cuda/acc.cuh" -#include "ggml-cuda/add-id.cuh" -#include "ggml-cuda/arange.cuh" -#include "ggml-cuda/argmax.cuh" -#include "ggml-cuda/argsort.cuh" -#include "ggml-cuda/binbcast.cuh" -#include "ggml-cuda/clamp.cuh" -#include "ggml-cuda/concat.cuh" -#include "ggml-cuda/conv-transpose-1d.cuh" -#include "ggml-cuda/conv2d.cuh" -#include "ggml-cuda/conv2d-dw.cuh" -#include "ggml-cuda/conv2d-transpose.cuh" -#include "ggml-cuda/convert.cuh" -#include "ggml-cuda/count-equal.cuh" -#include "ggml-cuda/cpy.cuh" -#include "ggml-cuda/cross-entropy-loss.cuh" -#include "ggml-cuda/cumsum.cuh" -#include "ggml-cuda/diagmask.cuh" -#include "ggml-cuda/diag.cuh" -#include "ggml-cuda/fattn.cuh" -#include "ggml-cuda/fattn-sparse.cuh" -#include "ggml-cuda/getrows.cuh" -#include "ggml-cuda/turbo-wht.cuh" -#include "ggml-cuda/im2col.cuh" -#include "ggml-cuda/mmf.cuh" -#include "ggml-cuda/mmq.cuh" -#include "ggml-cuda/mmvf.cuh" -#include "ggml-cuda/mmvq.cuh" -#include "ggml-cuda/norm.cuh" -#include "ggml-cuda/opt-step-adamw.cuh" -#include "ggml-cuda/opt-step-sgd.cuh" -#include "ggml-cuda/out-prod.cuh" -#include "ggml-cuda/pad.cuh" -#include "ggml-cuda/pool2d.cuh" -#include "ggml-cuda/quantize.cuh" -#include "ggml-cuda/rope.cuh" -#include "ggml-cuda/roll.cuh" -#include "ggml-cuda/scale.cuh" -#include "ggml-cuda/softcap.cuh" -#include "ggml-cuda/softmax.cuh" -#include "ggml-cuda/ssm-conv.cuh" -#include "ggml-cuda/ssm-scan.cuh" -#include "ggml-cuda/sum.cuh" -#include "ggml-cuda/sumrows.cuh" -#include "ggml-cuda/top-k.cuh" -#include "ggml-cuda/mean.cuh" -#include "ggml-cuda/tsembd.cuh" -#include "ggml-cuda/topk-moe.cuh" -#include "ggml-cuda/unary.cuh" -#include "ggml-cuda/upscale.cuh" -#include "ggml-cuda/wkv.cuh" -#include "ggml-cuda/gla.cuh" -#include "ggml-cuda/gated_delta_net.cuh" -#include "ggml-cuda/set.cuh" -#include "ggml-cuda/set-rows.cuh" -#include "ggml-cuda/pad_reflect_1d.cuh" -#include "ggml-cuda/solve_tri.cuh" -#include "ggml-cuda/tri.cuh" -#include "ggml-cuda/cumsum.cuh" -#include "ggml-cuda/fill.cuh" -#include "ggml-cuda/moe-fused.cuh" -#include "ggml.h" - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include - -static_assert(sizeof(half) == sizeof(ggml_fp16_t), "wrong fp16 size"); - -[[noreturn]] -void ggml_cuda_error(const char * stmt, const char * func, const char * file, int line, const char * msg) { - int id = -1; // in case cudaGetDevice fails - (void)cudaGetDevice(&id); - - GGML_LOG_ERROR(GGML_CUDA_NAME " error: %s\n", msg); - GGML_LOG_ERROR(" current device: %d, in function %s at %s:%d\n", id, func, file, line); - GGML_LOG_ERROR(" %s\n", stmt); - // abort with GGML_ABORT to get a stack trace - GGML_ABORT(GGML_CUDA_NAME " error"); -} - -// this is faster on Windows -// probably because the Windows CUDA libraries forget to make this check before invoking the drivers -void ggml_cuda_set_device(int device) { - int current_device; - CUDA_CHECK(cudaGetDevice(¤t_device)); - - if (device == current_device) { - return; - } - - CUDA_CHECK(cudaSetDevice(device)); -} - -int ggml_cuda_get_device() { - int id; - CUDA_CHECK(cudaGetDevice(&id)); - return id; -} - -// dflash: auto-enable unified (managed) memory on integrated GPUs (APU shared RAM) -// ONLY for allocations large enough that a plain device buffer plus the mmap source -// page cache (~2x size) would not fit in RAM and would thrash (e.g. an ~86GB DeepSeek-V4 -// model on a 125GB box). Managed memory has measurable alloc + access overhead on -// this APU, so small models that fit stay on the faster hipMalloc path (verified: -// 16GB model loads 3s / decodes 11.8 tok/s on hipMalloc vs 15s / 10.9 on managed). -// The cached ggml_cuda_info().devices[].integrated is hard-forced false (#15034 -// dodge), so probe a FRESH cudaDeviceProp. Opt out: DFLASH_HIP_NO_AUTO_UMA=1. -// Force-all: GGML_CUDA_ENABLE_UNIFIED_MEMORY. Tune gate: DFLASH_HIP_UMA_MIN_FRAC. -static size_t ggml_cuda_total_ram_bytes() { - static const size_t cached = []() -> size_t { - size_t bytes = 0; - FILE * f = fopen("/proc/meminfo", "r"); - if (f) { - char line[256]; - unsigned long kb = 0; - while (fgets(line, sizeof(line), f)) { - if (sscanf(line, "MemTotal: %lu kB", &kb) == 1) { bytes = (size_t) kb * 1024; break; } - } - fclose(f); - } - return bytes; - }(); - return cached; -} - -static bool ggml_cuda_device_use_uma(int device, size_t size) { - if (getenv("GGML_CUDA_ENABLE_UNIFIED_MEMORY") != nullptr) { - return true; - } - if (getenv("DFLASH_HIP_NO_AUTO_UMA") != nullptr) { - return false; - } - static const std::array integrated = []() { - std::array flags{}; - int n = 0; - if (cudaGetDeviceCount(&n) != cudaSuccess) { n = 0; } - for (int i = 0; i < n && i < GGML_CUDA_MAX_DEVICES; i++) { - cudaDeviceProp prop; - flags[i] = (cudaGetDeviceProperties(&prop, i) == cudaSuccess && prop.integrated); - } - return flags; - }(); - if (device < 0 || device >= GGML_CUDA_MAX_DEVICES || !integrated[device]) { - return false; - } - size_t total_ram = ggml_cuda_total_ram_bytes(); - if (total_ram == 0) { - return false; // unknown RAM: stay conservative on the legacy path - } - double frac = 0.45; - const char * fenv = getenv("DFLASH_HIP_UMA_MIN_FRAC"); - if (fenv != nullptr) { - double v = atof(fenv); - if (v > 0.0 && v < 1.0) { frac = v; } - } - bool use = (double) size > frac * (double) total_ram; - if (use) { - GGML_LOG_INFO("ggml_cuda: device %d integrated, alloc %.1f GiB (> %.0f%% of %.1f GiB RAM) " - "-> unified (managed) memory; small models stay on hipMalloc " - "(DFLASH_HIP_NO_AUTO_UMA=1 to disable)\n", - device, (double) size / 1073741824.0, frac * 100.0, - (double) total_ram / 1073741824.0); - } - return use; -} - -static cudaError_t ggml_cuda_device_malloc(void ** ptr, size_t size, int device, bool * out_managed = nullptr) { - ggml_cuda_set_device(device); - cudaError_t err; - bool managed = false; - if (ggml_cuda_device_use_uma(device, size)) { - err = cudaMallocManaged(ptr, size); - if (err == cudaSuccess) { - managed = true; - } -#if defined(GGML_USE_HIP) - if (err == hipSuccess) { - // hipMemAdviseSetCoarseGrain is an optional performance hint; - // ignore errors (e.g. hipErrorInvalidValue on some APU/iGPU configs). - (void)cudaMemAdvise(*ptr, size, hipMemAdviseSetCoarseGrain, device); - (void)hipGetLastError(); // clear any error - } - - // fall back to cudaMalloc if not supported (e.g. on Windows) - if (err == hipErrorNotSupported) { - static bool warned_unsupported = false; - if (!warned_unsupported) { - GGML_LOG_WARN("hipMallocManaged unsupported, falling back to hipMalloc.\n"); - warned_unsupported = true; - } - - err = cudaMalloc(ptr, size); - managed = false; - } -#endif // defined(GGML_USE_HIP) - } else { - err = cudaMalloc(ptr, size); - } - if (out_managed != nullptr) { - *out_managed = managed && (err == cudaSuccess); - } - return err; -} - -#if defined(GGML_USE_HIP) -static int ggml_cuda_parse_id(char devName[]) { - // A list of possible Target IDs can be found under the rocclr/clr repo in device.cpp - // these values are not stable so this is susceptible to breakage - // https://github.com/ROCm/clr/blob/amd-staging/rocclr/device/device.cpp - int archMajor = 0x0; - int archMinor = 0x0; - int archNum = GGML_CUDA_CC_OFFSET_AMD; - int archLen = strlen(devName); - char archName[archLen + 1]; - - // strip leading 'gfx' while copying into our buffer - if (archLen > 3) { - strcpy(archName, &devName[3]); - archLen -= 3; - } - - // trim trailing :xnack- or :sramecc- statuses - archLen = strcspn(archName, ":"); - archName[archLen] = '\0'; - - // tease out the version information - if (archLen > 8) { - // versions labeled generic use '-' as delimiter - // strip the trailing "-generic" then iterate through what remains - if ((strstr(archName, "-generic"))) { - archName[archLen - 8] = '\0'; - char * pch; - if ((pch = strtok(archName, "-"))) { - archMajor = (int)strtoul(pch, 0, 16); - if ((pch = strtok(NULL, "-"))) { - archMinor = 0x10 * (int)strtoul(pch, 0, 16); - } - } - } - } else if (archLen >= 3) { - // last two digits should be the minor * 0x10 + stepping - archMinor = (int)strtoul(&archName[archLen - 2], 0, 16); - archName[archLen - 2] = '\0'; - - // only the major version remains - archMajor = (int)strtoul(archName, 0, 16); - } - archNum += archMajor * 0x100; - archNum += archMinor; - return archNum; -} -#endif // defined(GGML_USE_HIP) - -static ggml_cuda_device_info ggml_cuda_init() { - ggml_cuda_device_info info = {}; - - cudaError_t err = cudaGetDeviceCount(&info.device_count); - if (err != cudaSuccess) { - GGML_LOG_ERROR("%s: failed to initialize " GGML_CUDA_NAME ": %s\n", __func__, cudaGetErrorString(err)); - return info; - } - - GGML_ASSERT(info.device_count <= GGML_CUDA_MAX_DEVICES); - - int64_t total_vram = 0; - for (int id = 0; id < info.device_count; ++id) { - cudaDeviceProp prop; - CUDA_CHECK(cudaGetDeviceProperties(&prop, id)); - total_vram += prop.totalGlobalMem; - } - GGML_LOG_INFO("%s: found %d " GGML_CUDA_NAME " devices (Total VRAM: %zu MiB):\n", - __func__, info.device_count, (size_t)(total_vram / (1024 * 1024))); - total_vram = 0; - - std::vector> turing_devices_without_mma; - for (int id = 0; id < info.device_count; ++id) { - int device_vmm = 0; - -#if defined(GGML_USE_VMM) - CUdevice device; - CU_CHECK(cuDeviceGet(&device, id)); - CU_CHECK(cuDeviceGetAttribute(&device_vmm, CU_DEVICE_ATTRIBUTE_VIRTUAL_MEMORY_MANAGEMENT_SUPPORTED, device)); - - if (device_vmm) { - CUmemAllocationProp alloc_prop = {}; - alloc_prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; - alloc_prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; - alloc_prop.location.id = id; - CU_CHECK(cuMemGetAllocationGranularity(&info.devices[id].vmm_granularity, &alloc_prop, CU_MEM_ALLOC_GRANULARITY_RECOMMENDED)); - } -#endif // defined(GGML_USE_VMM) - info.devices[id].vmm = !!device_vmm; - - cudaDeviceProp prop; - CUDA_CHECK(cudaGetDeviceProperties(&prop, id)); - - info.default_tensor_split[id] = total_vram; - total_vram += prop.totalGlobalMem; - info.devices[id].integrated = false; // Temporarily disabled due to issues with corrupted output (e.g. #15034) - info.devices[id].nsm = prop.multiProcessorCount; - info.devices[id].smpb = prop.sharedMemPerBlock; - info.devices[id].warp_size = prop.warpSize; - -#ifndef GGML_USE_MUSA - int supports_coop_launch = 0; - CUDA_CHECK(cudaDeviceGetAttribute(&supports_coop_launch, cudaDevAttrCooperativeLaunch, id)); - info.devices[id].supports_cooperative_launch = !!supports_coop_launch; -#else - info.devices[id].supports_cooperative_launch = false; -#endif // !(GGML_USE_MUSA) - -#if defined(GGML_USE_HIP) - info.devices[id].smpbo = prop.sharedMemPerBlock; - - info.devices[id].cc = ggml_cuda_parse_id(prop.gcnArchName); - if ((info.devices[id].cc & 0xff00) == 0x0) { - GGML_LOG_WARN("invalid architecture ID received for device %d %s: %s cc %d.%d\n", - id, prop.name, prop.gcnArchName, prop.major, prop.minor); - - // Fallback to prop.major and prop.minor - if (prop.major > 0) { - info.devices[id].cc = GGML_CUDA_CC_OFFSET_AMD + prop.major * 0x100; - info.devices[id].cc += prop.minor * 0x10; - } - } - GGML_LOG_INFO(" Device %d: %s, %s (0x%x), VMM: %s, Wave Size: %d, VRAM: %zu MiB\n", - id, prop.name, prop.gcnArchName, info.devices[id].cc & 0xffff, - device_vmm ? "yes" : "no", prop.warpSize, - (size_t)(prop.totalGlobalMem / (1024 * 1024))); -#elif defined(GGML_USE_MUSA) - // FIXME: Ensure compatibility with varying warp sizes across different MUSA archs. - info.devices[id].warp_size = 32; - info.devices[id].smpbo = prop.sharedMemPerBlockOptin; - info.devices[id].cc = GGML_CUDA_CC_OFFSET_MTHREADS + prop.major * 0x100; - info.devices[id].cc += prop.minor * 0x10; - GGML_LOG_INFO(" Device %d: %s, compute capability %d.%d, VMM: %s, VRAM: %zu MiB\n", - id, prop.name, prop.major, prop.minor, device_vmm ? "yes" : "no", - (size_t)(prop.totalGlobalMem / (1024 * 1024))); -#else - info.devices[id].smpbo = prop.sharedMemPerBlockOptin; - info.devices[id].cc = 100*prop.major + 10*prop.minor; - GGML_LOG_INFO(" Device %d: %s, compute capability %d.%d, VMM: %s, VRAM: %zu MiB\n", - id, prop.name, prop.major, prop.minor, device_vmm ? "yes" : "no", - (size_t)(prop.totalGlobalMem / (1024 * 1024))); - std::string device_name(prop.name); - if (device_name == "NVIDIA GeForce MX450") { - turing_devices_without_mma.push_back({ id, device_name }); - } else if (device_name == "NVIDIA GeForce MX550") { - turing_devices_without_mma.push_back({ id, device_name }); - } else if (device_name.substr(0, 21) == "NVIDIA GeForce GTX 16") { - turing_devices_without_mma.push_back({ id, device_name }); - } - - // Temporary performance fix: - // Setting device scheduling strategy for iGPUs with cc121 to "spinning" to avoid delays in cuda synchronize calls. - // TODO: Check for future drivers the default scheduling strategy and - // remove this call again when cudaDeviceScheduleSpin is default. - if (prop.major == 12 && prop.minor == 1) { - CUDA_CHECK(cudaSetDevice(id)); - CUDA_CHECK(cudaSetDeviceFlags(cudaDeviceScheduleSpin)); - } - -#endif // defined(GGML_USE_HIP) - } - - if (ggml_cuda_highest_compiled_arch(GGML_CUDA_CC_TURING) >= GGML_CUDA_CC_TURING && !turing_devices_without_mma.empty()) { - GGML_LOG_INFO("The following devices will have suboptimal performance due to a lack of tensor cores:\n"); - for (size_t device_pos = 0; device_pos < turing_devices_without_mma.size(); device_pos++) { - GGML_LOG_INFO( - " Device %d: %s\n", turing_devices_without_mma[device_pos].first, turing_devices_without_mma[device_pos].second.c_str()); - } - GGML_LOG_INFO( - "Consider compiling with CMAKE_CUDA_ARCHITECTURES=61-virtual;80-virtual and DGGML_CUDA_FORCE_MMQ to force the use of the Pascal code for Turing.\n"); - } - - for (int id = 0; id < info.device_count; ++id) { - info.default_tensor_split[id] /= total_vram; - } - - // configure logging to stdout - // CUBLAS_CHECK(cublasLoggerConfigure(1, 1, 0, nullptr)); - - for (int id = 0; id < info.device_count; ++id) { - ggml_cuda_set_device(id); - for (int id_other = 0; id_other < info.device_count; ++id_other) { - if (id == id_other) { - continue; - } - int can_access_peer; - CUDA_CHECK(cudaDeviceCanAccessPeer(&can_access_peer, id, id_other)); - if (can_access_peer) { - CUDA_CHECK(cudaDeviceEnablePeerAccess(id_other, 0)); - } - } - } - -#ifdef GGML_USE_NCCL - int dev_ids[GGML_CUDA_MAX_DEVICES]; - for (int id = 0; id < info.device_count; ++id) { - dev_ids[id] = id; - } - NCCL_CHECK(ncclCommInitAll(info.comms, info.device_count, dev_ids)); -#endif // GGML_USE_NCCL - - return info; -} - -const ggml_cuda_device_info & ggml_cuda_info() { - static ggml_cuda_device_info info = ggml_cuda_init(); - return info; -} - -// #define DEBUG_CUDA_MALLOC - -// buffer pool for cuda (legacy) -struct ggml_cuda_pool_leg : public ggml_cuda_pool { - static const int MAX_BUFFERS = 256; - - int device; - struct ggml_cuda_buffer { - void * ptr = nullptr; - size_t size = 0; - }; - - ggml_cuda_buffer buffer_pool[MAX_BUFFERS] = {}; - size_t pool_size = 0; - - explicit ggml_cuda_pool_leg(int device) : - device(device) { - } - - ~ggml_cuda_pool_leg() { - ggml_cuda_set_device(device); - for (int i = 0; i < MAX_BUFFERS; ++i) { - ggml_cuda_buffer & b = buffer_pool[i]; - if (b.ptr != nullptr) { - CUDA_CHECK(cudaFree(b.ptr)); - pool_size -= b.size; - } - } - GGML_ASSERT(pool_size == 0); - } - - void * alloc(size_t size, size_t * actual_size) override { -#ifdef DEBUG_CUDA_MALLOC - int nnz = 0; - size_t max_size = 0; -#endif - size_t best_diff = 1ull << 36; - int ibest = -1; - for (int i = 0; i < MAX_BUFFERS; ++i) { - ggml_cuda_buffer& b = buffer_pool[i]; - if (b.ptr != nullptr) { -#ifdef DEBUG_CUDA_MALLOC - ++nnz; - if (b.size > max_size) max_size = b.size; -#endif - if (b.size >= size) { - size_t diff = b.size - size; - if (diff < best_diff) { - best_diff = diff; - ibest = i; - if (!best_diff) { - void * ptr = b.ptr; - *actual_size = b.size; - b.ptr = nullptr; - b.size = 0; - return ptr; - } - } - } - } - } - if (ibest >= 0) { - ggml_cuda_buffer& b = buffer_pool[ibest]; - void * ptr = b.ptr; - *actual_size = b.size; - b.ptr = nullptr; - b.size = 0; - return ptr; - } - void * ptr; - size_t look_ahead_size = (size_t) (1.05 * size); - look_ahead_size = 256 * ((look_ahead_size + 255)/256); - ggml_cuda_set_device(device); - CUDA_CHECK(ggml_cuda_device_malloc(&ptr, look_ahead_size, device)); - *actual_size = look_ahead_size; - pool_size += look_ahead_size; -#ifdef DEBUG_CUDA_MALLOC - GGML_LOG_INFO("%s[%d]: %d buffers, max_size = %u MB, pool_size = %u MB, requested %u MB\n", __func__, device, nnz, - (uint32_t)(max_size / 1024 / 1024), (uint32_t)(pool_size / 1024 / 1024), (uint32_t)(size / 1024 / 1024)); -#endif - return ptr; - } - - void free(void * ptr, size_t size) override { - for (int i = 0; i < MAX_BUFFERS; ++i) { - ggml_cuda_buffer& b = buffer_pool[i]; - if (b.ptr == nullptr) { - b.ptr = ptr; - b.size = size; - return; - } - } - GGML_LOG_DEBUG(GGML_CUDA_NAME " buffer pool full, increase MAX_CUDA_BUFFERS\n"); - ggml_cuda_set_device(device); - CUDA_CHECK(cudaFree(ptr)); - pool_size -= size; - } -}; - -// pool with virtual memory -#if defined(GGML_USE_VMM) -struct ggml_cuda_pool_vmm : public ggml_cuda_pool { - static const size_t CUDA_POOL_VMM_MAX_SIZE = 1ull << 35; // 32 GB - - int device; - CUdeviceptr pool_addr = 0; - size_t pool_used = 0; - size_t pool_size = 0; - size_t granularity; -#if defined(GGML_USE_HIP) - std::vector> mappings; -#endif - - explicit ggml_cuda_pool_vmm(int device) : - device(device), - granularity(ggml_cuda_info().devices[device].vmm_granularity) { - } - - ~ggml_cuda_pool_vmm() { - if (pool_addr != 0) { -#if defined(GGML_USE_HIP) - // Workaround for https://github.com/ROCm/ROCR-Runtime/issues/285 - for (std::pair & mapping : mappings) { - CU_CHECK(cuMemUnmap(mapping.first, mapping.second)); - } -#else - CU_CHECK(cuMemUnmap(pool_addr, pool_size)); -#endif - CU_CHECK(cuMemAddressFree(pool_addr, CUDA_POOL_VMM_MAX_SIZE)); - } - } - - void * alloc(size_t size, size_t * actual_size) override { - // round up the allocation size to the alignment to ensure that all allocations are aligned for all data types - const size_t alignment = 128; - size = alignment * ((size + alignment - 1) / alignment); - - size_t avail = pool_size - pool_used; - - if (size > avail) { - // round up to the next multiple of the granularity - size_t reserve_size = size - avail; - reserve_size = granularity * ((reserve_size + granularity - 1) / granularity); - - GGML_ASSERT(pool_size + reserve_size <= CUDA_POOL_VMM_MAX_SIZE); - - // Pool extension via cuMem* must run with the device idle. Without - // this sync, cuMemSetAccess below can fire while previously-queued - // async kernels (compute, tensor copies) are still in flight on - // the default stream and return CUDA_ERROR_NOT_READY. The CU_CHECK - // macro hits GGML_ABORT on that error and the resulting - // signal/longjmp ends up corrupting the just-mapped region: the - // mapping is registered but access permissions are never set, so - // every subsequent read/write to it silently misbehaves. Easiest - // visible failure: prefix-cache snapshots stored into freshly - // extended pool memory return zeroed/garbled state on restore. - // Repro: any workload that queues compute then immediately - // alloc_ctx_tensors on the same backend. - CUDA_CHECK(cudaDeviceSynchronize()); - - // allocate more physical memory - CUmemAllocationProp prop = {}; - prop.type = CU_MEM_ALLOCATION_TYPE_PINNED; - prop.location.type = CU_MEM_LOCATION_TYPE_DEVICE; - prop.location.id = device; - CUmemGenericAllocationHandle handle; - CU_CHECK(cuMemCreate(&handle, reserve_size, &prop, 0)); - - // reserve virtual address space (if not already reserved) - if (pool_addr == 0) { - CU_CHECK(cuMemAddressReserve(&pool_addr, CUDA_POOL_VMM_MAX_SIZE, 0, 0, 0)); - } - - // map at the end of the pool - CUdeviceptr start_ptr = (CUdeviceptr)((char *)(pool_addr) + pool_size); - CU_CHECK(cuMemMap(start_ptr, reserve_size, 0, handle, 0)); -#if defined(GGML_USE_HIP) - mappings.push_back({start_ptr, reserve_size}); -#endif - - // the memory allocation handle is no longer needed after mapping - CU_CHECK(cuMemRelease(handle)); - - // set access - CUmemAccessDesc access = {}; - access.location.type = CU_MEM_LOCATION_TYPE_DEVICE; - access.location.id = device; - access.flags = CU_MEM_ACCESS_FLAGS_PROT_READWRITE; - // The sync above is not enough under concurrency: other threads - // (snapshot D2H, a second client, the draft backend) can queue - // async work between it and this remap, and cuMemSetAccess then - // returns CUDA_ERROR_NOT_READY (observed on WSL2 at >19GB load). - // NOT_READY is transient: re-sync and retry instead of aborting. - CUresult set_err = cuMemSetAccess(start_ptr, reserve_size, &access, 1); - for (int retry = 0; set_err == CUDA_ERROR_NOT_READY && retry < 100; ++retry) { - CUDA_CHECK(cudaDeviceSynchronize()); - set_err = cuMemSetAccess(start_ptr, reserve_size, &access, 1); - } - CU_CHECK(set_err); - - // add to the pool - pool_size += reserve_size; - - //printf("cuda pool[%d]: size increased to %llu MB (reserved %llu MB)\n", - // device, (unsigned long long) (pool_size/1024/1024), - // (unsigned long long) (reserve_size/1024/1024)); - } - - GGML_ASSERT(pool_addr != 0); - - void * ptr = (void *) ((CUdeviceptr)((char *)(pool_addr) + pool_used)); - *actual_size = size; - pool_used += size; - -#ifdef DEBUG_CUDA_MALLOC - printf("cuda pool[%d]: allocated %llu bytes at %llx\n", device, (unsigned long long) size, ptr); -#endif - - return ptr; - } - - void free(void * ptr, size_t size) override { -#ifdef DEBUG_CUDA_MALLOC - printf("cuda pool[%d]: freed %llu bytes at %llx\n", device, (unsigned long long) size, ptr); -#endif - - pool_used -= size; - - // all deallocations must be in reverse order of the allocations - GGML_ASSERT(ptr == (void *) ((char *)(pool_addr) + pool_used)); - } -}; -#endif // defined(GGML_USE_VMM) - -std::unique_ptr ggml_backend_cuda_context::new_pool_for_device(int device, - [[maybe_unused]] int stream_no) { -#if defined(GGML_USE_VMM) - if (ggml_cuda_info().devices[device].vmm) { - return std::unique_ptr(new ggml_cuda_pool_vmm(device)); - } -#endif // defined(GGML_USE_VMM) - return std::unique_ptr(new ggml_cuda_pool_leg(device)); -} - -// destroying a cuBLAS handle while a graph is being captured in a different thread can result in a CUDA error -// this lock is used to ensure that no cuBLAS handle is destroyed while a graph is being captured - -static std::mutex ggml_cuda_lock; -static std::condition_variable ggml_cuda_lock_cv; -static std::atomic ggml_cuda_lock_counter; - -ggml_backend_cuda_context::~ggml_backend_cuda_context() { - std::unique_lock lock(ggml_cuda_lock); - ggml_cuda_lock_cv.wait(lock, []{ return ggml_cuda_lock_counter.load(std::memory_order_relaxed) == 0; }); - - if (copy_event != nullptr) { - CUDA_CHECK(cudaEventDestroy(copy_event)); - } - for (int i = 0; i < GGML_CUDA_MAX_DEVICES; ++i) { - for (int j = 0; j < GGML_CUDA_MAX_STREAMS; ++j) { - if (streams[i][j] != nullptr) { - CUDA_CHECK(cudaStreamDestroy(streams[i][j])); - } - } - if (cublas_handles[i] != nullptr) { - CUBLAS_CHECK(cublasDestroy(cublas_handles[i])); - } - } -} - - -// cuda buffer - -struct ggml_backend_cuda_buffer_context { - int device; - void * dev_ptr = nullptr; - bool is_managed = false; // dflash: dev_ptr is unified/managed (CPU-addressable) - std::string name; - - ggml_backend_cuda_buffer_context(int device, void * dev_ptr) : - device(device), dev_ptr(dev_ptr), - name(GGML_CUDA_NAME + std::to_string(device)) { - } - - ~ggml_backend_cuda_buffer_context() { - CUDA_CHECK(cudaFree(dev_ptr)); - } -}; - -static void ggml_backend_cuda_buffer_free_buffer(ggml_backend_buffer_t buffer) { - ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *)buffer->context; - delete ctx; -} - -static bool ggml_backend_buffer_is_cuda(ggml_backend_buffer_t buffer) { - return buffer->iface.free_buffer == ggml_backend_cuda_buffer_free_buffer; -} - -static void * ggml_backend_cuda_buffer_get_base(ggml_backend_buffer_t buffer) { - ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *)buffer->context; - return ctx->dev_ptr; -} - -// dflash: query whether a buffer is unified/managed (CPU-addressable) so loaders can -// stream weights straight into it with explicit reads instead of a host->device copy. -extern "C" bool ggml_backend_cuda_buffer_is_managed(ggml_backend_buffer_t buffer) { - if (buffer == nullptr || !ggml_backend_buffer_is_cuda(buffer)) { - return false; - } - return ((ggml_backend_cuda_buffer_context *) buffer->context)->is_managed; -} - -static enum ggml_status ggml_backend_cuda_buffer_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) { - ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *)buffer->context; - - if (tensor->view_src != NULL) { - assert(tensor->view_src->buffer->buft == buffer->buft); - return GGML_STATUS_SUCCESS; - } - - if (ggml_is_quantized(tensor->type) && tensor->view_src == nullptr && ggml_backend_buffer_get_usage(buffer) != GGML_BACKEND_BUFFER_USAGE_COMPUTE) { - // initialize padding to 0 to avoid possible NaN values - const size_t original_size = ggml_nbytes(tensor); - const size_t padded_size = ggml_backend_buft_get_alloc_size(buffer->buft, tensor); - - if (padded_size > original_size) { - ggml_cuda_set_device(ctx->device); - CUDA_CHECK(cudaMemset((char *)tensor->data + original_size, 0, padded_size - original_size)); - } - } - return GGML_STATUS_SUCCESS; -} - -static void ggml_backend_cuda_buffer_memset_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, uint8_t value, size_t offset, size_t size) { - ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *) buffer->context; - - ggml_cuda_set_device(ctx->device); - CUDA_CHECK(cudaMemsetAsync((char *) tensor->data + offset, value, size, cudaStreamPerThread)); - CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread)); -} - -static void ggml_backend_cuda_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { - ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *) buffer->context; - - ggml_cuda_set_device(ctx->device); - if (ctx->is_managed) { - // Unified/managed memory is CPU-addressable: copy weights straight in with a - // host memcpy, skipping the pageable host->device staged copy + per-tensor - // synchronize that dominate large-model load on shared-memory (integrated) - // GPUs. GPU visibility is established by the backend synchronize before the - // first graph compute. - char * dst = (char *) tensor->data + offset; - const char * src = (const char *) data; - const size_t par_threshold = (size_t) 16 * 1024 * 1024; - unsigned nth = std::thread::hardware_concurrency(); - if (nth == 0) nth = 4; - if (nth > 16) nth = 16; - if (size >= par_threshold && nth > 1) { - // Split large copies across threads so cold-cache page-ins and the - // RAM->RAM copy run in parallel (disjoint byte ranges, no sharing). - std::vector workers; - workers.reserve(nth); - const size_t chunk = (size + nth - 1) / nth; - for (unsigned t = 0; t < nth; t++) { - const size_t s0 = (size_t) t * chunk; - if (s0 >= size) break; - const size_t n = (chunk < size - s0) ? chunk : (size - s0); - workers.emplace_back([dst, src, s0, n]() { - memcpy(dst + s0, src + s0, n); - }); - } - for (auto & w : workers) { - w.join(); - } - } else { - memcpy(dst, src, size); - } - return; - } - CUDA_CHECK(cudaMemcpyAsync((char *) tensor->data + offset, data, size, cudaMemcpyHostToDevice, cudaStreamPerThread)); - CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread)); -} - -static void ggml_backend_cuda_buffer_get_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * tensor, void * data, size_t offset, size_t size) { - ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *) buffer->context; - - ggml_cuda_set_device(ctx->device); - CUDA_CHECK(cudaMemcpyAsync(data, (const char *) tensor->data + offset, size, cudaMemcpyDeviceToHost, cudaStreamPerThread)); - CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread)); -} - -static void ggml_backend_cuda_buffer_set_tensor_2d(ggml_backend_buffer_t buffer, 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) { - ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *) buffer->context; - - ggml_cuda_set_device(ctx->device); - CUDA_CHECK(cudaMemcpy2DAsync( - (char *) tensor->data + offset, stride_tensor, data, stride_data, size, n_copies, cudaMemcpyHostToDevice, cudaStreamPerThread)); - CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread)); -} - -static void ggml_backend_cuda_buffer_get_tensor_2d(ggml_backend_buffer_t buffer, 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) { - ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *)buffer->context; - - ggml_cuda_set_device(ctx->device); - CUDA_CHECK(cudaMemcpy2DAsync( - data, stride_data, (const char *) tensor->data + offset, stride_tensor, size, n_copies, cudaMemcpyDeviceToHost, cudaStreamPerThread)); - CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread)); -} - -static bool ggml_backend_cuda_buffer_cpy_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * src, ggml_tensor * dst) { - if (ggml_backend_buffer_is_cuda(src->buffer)) { - ggml_backend_cuda_buffer_context * src_ctx = (ggml_backend_cuda_buffer_context *)src->buffer->context; - ggml_backend_cuda_buffer_context * dst_ctx = (ggml_backend_cuda_buffer_context *)dst->buffer->context; - if (src_ctx->device == dst_ctx->device) { - CUDA_CHECK(cudaMemcpyAsync(dst->data, src->data, ggml_nbytes(src), cudaMemcpyDeviceToDevice, cudaStreamPerThread)); - } else { -#ifdef GGML_CUDA_NO_PEER_COPY - return false; -#else - CUDA_CHECK(cudaMemcpyPeerAsync(dst->data, dst_ctx->device, src->data, src_ctx->device, ggml_nbytes(src), cudaStreamPerThread)); -#endif - } - CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread)); - return true; - } - return false; - - GGML_UNUSED(buffer); -} - -static void ggml_backend_cuda_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { - ggml_backend_cuda_buffer_context * ctx = (ggml_backend_cuda_buffer_context *)buffer->context; - - ggml_cuda_set_device(ctx->device); - CUDA_CHECK(cudaMemsetAsync(ctx->dev_ptr, value, buffer->size, cudaStreamPerThread)); - CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread)); -} - -static const ggml_backend_buffer_i ggml_backend_cuda_buffer_interface = { - /* .free_buffer = */ ggml_backend_cuda_buffer_free_buffer, - /* .get_base = */ ggml_backend_cuda_buffer_get_base, - /* .init_tensor = */ ggml_backend_cuda_buffer_init_tensor, - /* .memset_tensor = */ ggml_backend_cuda_buffer_memset_tensor, - /* .set_tensor = */ ggml_backend_cuda_buffer_set_tensor, - /* .get_tensor = */ ggml_backend_cuda_buffer_get_tensor, - /* .set_tensor_2d = */ ggml_backend_cuda_buffer_set_tensor_2d, - /* .get_tensor_2d = */ ggml_backend_cuda_buffer_get_tensor_2d, - /* .cpy_tensor = */ ggml_backend_cuda_buffer_cpy_tensor, - /* .clear = */ ggml_backend_cuda_buffer_clear, - /* .reset = */ NULL, -}; - -// cuda buffer type -struct ggml_backend_cuda_buffer_type_context { - int device; - std::string name; -}; - -static const char * ggml_backend_cuda_buffer_type_get_name(ggml_backend_buffer_type_t buft) { - ggml_backend_cuda_buffer_type_context * ctx = (ggml_backend_cuda_buffer_type_context *)buft->context; - - return ctx->name.c_str(); -} - -static bool ggml_backend_buft_is_cuda(ggml_backend_buffer_type_t buft) { - return buft->iface.get_name == ggml_backend_cuda_buffer_type_get_name; -} - -static ggml_backend_buffer_t ggml_backend_cuda_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { - ggml_backend_cuda_buffer_type_context * buft_ctx = (ggml_backend_cuda_buffer_type_context *)buft->context; - - ggml_cuda_set_device(buft_ctx->device); - - void * dev_ptr; - bool is_managed = false; - cudaError_t err = ggml_cuda_device_malloc(&dev_ptr, size, buft_ctx->device, &is_managed); - if (err != cudaSuccess) { - // clear the error - (void)cudaGetLastError(); - GGML_LOG_ERROR("%s: allocating %.2f MiB on device %d: cudaMalloc failed: %s\n", __func__, size / 1024.0 / 1024.0, buft_ctx->device, cudaGetErrorString(err)); - return nullptr; - } - - ggml_backend_cuda_buffer_context * ctx = new ggml_backend_cuda_buffer_context(buft_ctx->device, dev_ptr); - ctx->is_managed = is_managed; - - return ggml_backend_buffer_init(buft, ggml_backend_cuda_buffer_interface, ctx, size); -} - -static size_t ggml_backend_cuda_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) { - return 128; - - GGML_UNUSED(buft); -} - -static size_t ggml_backend_cuda_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) { - size_t size = ggml_nbytes(tensor); - int64_t ne0 = tensor->ne[0]; - - if (ggml_is_quantized(tensor->type)) { - if (ne0 % MATRIX_ROW_PADDING != 0) { - GGML_ASSERT(tensor->nb[0] == ggml_element_size(tensor)); - size += ggml_row_size(tensor->type, MATRIX_ROW_PADDING - ne0 % MATRIX_ROW_PADDING); - } - } - - return size; - - GGML_UNUSED(buft); -} - -static const ggml_backend_buffer_type_i ggml_backend_cuda_buffer_type_interface = { - /* .get_name = */ ggml_backend_cuda_buffer_type_get_name, - /* .alloc_buffer = */ ggml_backend_cuda_buffer_type_alloc_buffer, - /* .get_alignment = */ ggml_backend_cuda_buffer_type_get_alignment, - /* .get_max_size = */ NULL, // defaults to SIZE_MAX - /* .get_alloc_size = */ ggml_backend_cuda_buffer_type_get_alloc_size, - /* .is_host = */ NULL, -}; - -ggml_backend_buffer_type_t ggml_backend_cuda_buffer_type(int device) { - static std::mutex mutex; - std::lock_guard lock(mutex); - - if (device >= ggml_backend_cuda_get_device_count()) { - return nullptr; - } - - static ggml_backend_buffer_type ggml_backend_cuda_buffer_types[GGML_CUDA_MAX_DEVICES]; - - static bool ggml_backend_cuda_buffer_type_initialized = false; - - if (!ggml_backend_cuda_buffer_type_initialized) { - for (int i = 0; i < ggml_backend_cuda_get_device_count(); i++) { - ggml_backend_cuda_buffer_types[i] = { - /* .iface = */ ggml_backend_cuda_buffer_type_interface, - /* .device = */ ggml_backend_reg_dev_get(ggml_backend_cuda_reg(), i), - /* .context = */ new ggml_backend_cuda_buffer_type_context{i, GGML_CUDA_NAME + std::to_string(i)}, - }; - } - ggml_backend_cuda_buffer_type_initialized = true; - } - - return &ggml_backend_cuda_buffer_types[device]; -} - -// cuda split buffer - -static int64_t get_row_rounding(const std::array & tensor_split) { - int64_t row_rounding = 0; - for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { - if (tensor_split[id] >= (id + 1 < ggml_backend_cuda_get_device_count() ? tensor_split[id + 1] : 1.0f)) { - continue; - } - - const int cc = ggml_cuda_info().devices[id].cc; - row_rounding = std::max(row_rounding, (int64_t)get_mmq_y_host(cc)); - } - return row_rounding; -} - -static void get_row_split(int64_t * row_low, int64_t * row_high, const ggml_tensor * tensor, const std::array & tensor_split, int id) { - const int64_t nrows = ggml_nrows(tensor); - const int64_t rounding = get_row_rounding(tensor_split); - - *row_low = id == 0 ? 0 : nrows*tensor_split[id]; - *row_low -= *row_low % rounding; - - if (id == ggml_backend_cuda_get_device_count() - 1) { - *row_high = nrows; - } else { - *row_high = nrows*tensor_split[id + 1]; - *row_high -= *row_high % rounding; - } -} - -static size_t ggml_nbytes_split(const struct ggml_tensor * tensor, int nrows_split) { - static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); - - return nrows_split*ggml_row_size(tensor->type, tensor->ne[0]); -} - -struct ggml_backend_cuda_split_buffer_type_context { - int main_device; - std::array tensor_split; - std::string name; -}; - -struct ggml_backend_cuda_split_buffer_context { - ~ggml_backend_cuda_split_buffer_context() { - for (ggml_tensor_extra_gpu * extra : tensor_extras) { - for (int id = 0; id < GGML_CUDA_MAX_DEVICES; ++id) { - for (int64_t is = 0; is < GGML_CUDA_MAX_STREAMS; ++is) { - if (extra->events[id][is] != nullptr) { - CUDA_CHECK(cudaEventDestroy(extra->events[id][is])); - } - } - if (extra->data_device[id] != nullptr) { - CUDA_CHECK(cudaFree(extra->data_device[id])); - } - } - delete extra; - } - } - - std::vector tensor_extras; -}; - - -static void ggml_backend_cuda_split_buffer_free_buffer(ggml_backend_buffer_t buffer) { - ggml_backend_cuda_split_buffer_context * ctx = (ggml_backend_cuda_split_buffer_context *)buffer->context; - delete ctx; -} - -static void * ggml_backend_cuda_split_buffer_get_base(ggml_backend_buffer_t buffer) { - // the pointers are stored in the tensor extras, this is just a dummy address and never dereferenced - return (void *)0x1000; - - GGML_UNUSED(buffer); -} - -static enum ggml_status ggml_backend_cuda_split_buffer_init_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor) { - GGML_ASSERT(tensor->view_src == nullptr); // views of split tensors are not supported - GGML_ASSERT(ggml_is_contiguous(tensor) && "split buffers only supported for contiguous tensors"); - - ggml_backend_cuda_split_buffer_context * ctx = (ggml_backend_cuda_split_buffer_context *)buffer->context; - ggml_backend_cuda_split_buffer_type_context * buft_ctx = (ggml_backend_cuda_split_buffer_type_context *)buffer->buft->context; - - const int64_t ne0 = tensor->ne[0]; - - ggml_tensor_extra_gpu * extra = new ggml_tensor_extra_gpu{}; - ctx->tensor_extras.push_back(extra); - - for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { - int64_t row_low, row_high; - get_row_split(&row_low, &row_high, tensor, buft_ctx->tensor_split, id); - - int64_t nrows_split = row_high - row_low; - if (nrows_split == 0) { - continue; - } - - size_t size = ggml_nbytes_split(tensor, nrows_split); - const size_t original_size = size; - - // pad last row to a multiple of 512 elements to avoid out-of-bounds memory accesses - if (ne0 % MATRIX_ROW_PADDING != 0) { - size += ggml_row_size(tensor->type, MATRIX_ROW_PADDING - ne0 % MATRIX_ROW_PADDING); - } - - // FIXME: do not crash if cudaMalloc fails - // currently, init_tensor cannot fail, it needs to be fixed in ggml-backend first - ggml_cuda_set_device(id); - char * buf; - CUDA_CHECK(ggml_cuda_device_malloc((void**)&buf, size, id)); - - // set padding to 0 to avoid possible NaN values - if (size > original_size) { - CUDA_CHECK(cudaMemset(buf + original_size, 0, size - original_size)); - } - - extra->data_device[id] = buf; - - for (int64_t is = 0; is < GGML_CUDA_MAX_STREAMS; ++is) { - CUDA_CHECK(cudaEventCreateWithFlags(&extra->events[id][is], cudaEventDisableTiming)); - } - } - tensor->extra = extra; - return GGML_STATUS_SUCCESS; -} - -static void ggml_backend_cuda_split_buffer_set_tensor(ggml_backend_buffer_t buffer, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { - // split tensors must always be set in their entirety at once - GGML_ASSERT(offset == 0); - GGML_ASSERT(size == ggml_nbytes(tensor)); - GGML_ASSERT(ggml_is_contiguous(tensor) && "split buffers only supported for contiguous tensors"); - - ggml_backend_cuda_split_buffer_type_context * buft_ctx = (ggml_backend_cuda_split_buffer_type_context *)buffer->buft->context; - - const int64_t ne0 = tensor->ne[0]; - const size_t nb1 = tensor->nb[1]; - ggml_tensor_extra_gpu * extra = (ggml_tensor_extra_gpu *)tensor->extra; - - for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { - int64_t row_low, row_high; - get_row_split(&row_low, &row_high, tensor, buft_ctx->tensor_split, id); - - int64_t nrows_split = row_high - row_low; - if (nrows_split == 0) { - continue; - } - - const size_t offset_split = row_low*nb1; - size_t size = ggml_nbytes_split(tensor, nrows_split); - const size_t original_size = size; - - // pad last row to a multiple of 512 elements to avoid out-of-bounds memory accesses - if (ne0 % MATRIX_ROW_PADDING != 0) { - size += ggml_row_size(tensor->type, MATRIX_ROW_PADDING - ne0 % MATRIX_ROW_PADDING); - } - - const char * buf_host = (const char *)data + offset_split; - CUDA_CHECK(cudaMemcpyAsync(extra->data_device[id], buf_host, original_size, cudaMemcpyHostToDevice, cudaStreamPerThread)); - } - - for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { - CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread)); - } -} - -static void ggml_backend_cuda_split_buffer_get_tensor(ggml_backend_buffer_t buffer, const ggml_tensor * tensor, void * data, size_t offset, size_t size) { - // split tensors must always be set in their entirety at once - GGML_ASSERT(offset == 0); - GGML_ASSERT(size == ggml_nbytes(tensor)); - GGML_ASSERT(ggml_is_contiguous(tensor) && "split buffers only supported for contiguous tensors"); - - ggml_backend_cuda_split_buffer_type_context * buft_ctx = (ggml_backend_cuda_split_buffer_type_context *)buffer->buft->context; - - const int64_t ne0 = tensor->ne[0]; - const size_t nb1 = tensor->nb[1]; - ggml_tensor_extra_gpu * extra = (ggml_tensor_extra_gpu *)tensor->extra; - - for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { - int64_t row_low, row_high; - get_row_split(&row_low, &row_high, tensor, buft_ctx->tensor_split, id); - - int64_t nrows_split = row_high - row_low; - if (nrows_split == 0) { - continue; - } - - const size_t offset_split = row_low*nb1; - size_t size = ggml_nbytes_split(tensor, nrows_split); - const size_t original_size = size; - - // pad last row to a multiple of 512 elements to avoid out-of-bounds memory accesses - if (ne0 % MATRIX_ROW_PADDING != 0) { - size += ggml_row_size(tensor->type, MATRIX_ROW_PADDING - ne0 % MATRIX_ROW_PADDING); - } - - char * buf_host = (char *)data + offset_split; - CUDA_CHECK(cudaMemcpyAsync(buf_host, extra->data_device[id], original_size, cudaMemcpyDeviceToHost, cudaStreamPerThread)); - } - - for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { - CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread)); - } -} - -static void ggml_backend_cuda_split_buffer_clear(ggml_backend_buffer_t buffer, uint8_t value) { - GGML_UNUSED(buffer); - GGML_UNUSED(value); -} - -static const ggml_backend_buffer_i ggml_backend_cuda_split_buffer_interface = { - /* .free_buffer = */ ggml_backend_cuda_split_buffer_free_buffer, - /* .get_base = */ ggml_backend_cuda_split_buffer_get_base, - /* .init_tensor = */ ggml_backend_cuda_split_buffer_init_tensor, - /* .memset_tensor = */ NULL, - /* .set_tensor = */ ggml_backend_cuda_split_buffer_set_tensor, - /* .get_tensor = */ ggml_backend_cuda_split_buffer_get_tensor, - /* .set_tensor_2d = */ NULL, - /* .get_tensor_2d = */ NULL, - /* .cpy_tensor = */ NULL, - /* .clear = */ ggml_backend_cuda_split_buffer_clear, - /* .reset = */ NULL, -}; - -// cuda split buffer type - -static const char * ggml_backend_cuda_split_buffer_type_get_name(ggml_backend_buffer_type_t buft) { - ggml_backend_cuda_split_buffer_type_context * ctx = (ggml_backend_cuda_split_buffer_type_context *)buft->context; - - return ctx->name.c_str(); -} - -static bool ggml_backend_buft_is_cuda_split(ggml_backend_buffer_type_t buft) { - return buft->iface.get_name == ggml_backend_cuda_split_buffer_type_get_name; -} - -static ggml_backend_buffer_t ggml_backend_cuda_split_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { - // since we don't know the exact split after rounding, we cannot allocate the device buffers at this point - // instead, we allocate them for each tensor separately in init_tensor - // however, the size still represents the maximum cumulative size of all the device buffers after the tensors are allocated, - // as returned by get_alloc_size. this limit is enforced during tensor allocation by ggml-alloc, so it must be correct. - ggml_backend_cuda_split_buffer_context * ctx = new ggml_backend_cuda_split_buffer_context(); - - return ggml_backend_buffer_init(buft, ggml_backend_cuda_split_buffer_interface, ctx, size); -} - -static size_t ggml_backend_cuda_split_buffer_type_get_alignment(ggml_backend_buffer_type_t buft) { - return 128; - - GGML_UNUSED(buft); -} - -static size_t ggml_backend_cuda_split_buffer_type_get_alloc_size(ggml_backend_buffer_type_t buft, const ggml_tensor * tensor) { - ggml_backend_cuda_split_buffer_type_context * ctx = (ggml_backend_cuda_split_buffer_type_context *)buft->context; - GGML_ASSERT(ggml_is_contiguous(tensor) && "split buffers only supported for contiguous tensors"); - - size_t total_size = 0; - - const int64_t ne0 = tensor->ne[0]; - - for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { - int64_t row_low, row_high; - get_row_split(&row_low, &row_high, tensor, ctx->tensor_split, id); - - int64_t nrows_split = row_high - row_low; - if (nrows_split == 0) { - continue; - } - - total_size += ggml_nbytes_split(tensor, nrows_split); - - // pad last row to a multiple of 512 elements to avoid out-of-bounds memory accesses - if (ne0 % MATRIX_ROW_PADDING != 0) { - total_size += ggml_row_size(tensor->type, MATRIX_ROW_PADDING - ne0 % MATRIX_ROW_PADDING); - } - } - - return total_size; -} - -static bool ggml_backend_cuda_split_buffer_type_is_host(ggml_backend_buffer_type_t buft) { - return false; - - GGML_UNUSED(buft); -} - -static const ggml_backend_buffer_type_i ggml_backend_cuda_split_buffer_type_interface = { - /* .get_name = */ ggml_backend_cuda_split_buffer_type_get_name, - /* .alloc_buffer = */ ggml_backend_cuda_split_buffer_type_alloc_buffer, - /* .get_alignment = */ ggml_backend_cuda_split_buffer_type_get_alignment, - /* .get_max_size = */ NULL, // defaults to SIZE_MAX - /* .get_alloc_size = */ ggml_backend_cuda_split_buffer_type_get_alloc_size, - /* .is_host = */ ggml_backend_cuda_split_buffer_type_is_host, -}; - -bool ggml_backend_cuda_allreduce_tensor(ggml_backend_t * backends, struct ggml_tensor ** tensors, size_t n_backends) { -#ifdef GGML_USE_NCCL - const int64_t ne = ggml_nelements(tensors[0]); - // FIXME the input of llm_graph_context::build_in_out_ids can produce a tensor with 0 elements if n_outputs == 0 - // This then causes a crash in this function - if (ne == 0) { - return true; - } - for (size_t i = 0; i < n_backends; ++i) { - GGML_ASSERT(tensors[i] != nullptr); - GGML_ASSERT(ggml_nelements(tensors[i]) == ne); - GGML_ASSERT(ggml_is_contiguously_allocated(tensors[i])); - } - - const ggml_cuda_device_info info = ggml_cuda_info(); - - // For small tensors, simply reduce them as FP32. - // The following heuristic for how "small" a tensor should be is based on RTX 4090s connected via 16x PCIe 4.0. - if ((n_backends <= 2 && ne < 32768) || (n_backends == 3 && ne < 131072) || (n_backends >= 4 && ne < 262144)) { - NCCL_CHECK(ncclGroupStart()); - for (size_t i = 0; i < n_backends; ++i) { - ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backends[i]->context; - NCCL_CHECK(ncclAllReduce(tensors[i]->data, tensors[i]->data, ne, ncclFloat, ncclSum, info.comms[cuda_ctx->device], cuda_ctx->stream())); - } - NCCL_CHECK(ncclGroupEnd()); - - return true; - } - - // For large tensors it's faster to compress them to BF16 for the reduction: - to_bf16_cuda_t to_bf16 = ggml_get_to_bf16_cuda(GGML_TYPE_F32); - to_fp32_cuda_t to_fp32 = ggml_get_to_fp32_cuda(GGML_TYPE_BF16); - - ggml_cuda_pool_alloc tmp[GGML_CUDA_MAX_DEVICES]; - for (size_t i = 0; i < n_backends; ++i) { - ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backends[i]->context; - tmp[i].pool = &cuda_ctx->pool(); - tmp[i].alloc(ne); - - ggml_cuda_set_device(i); - to_bf16(tensors[i]->data, tmp[i].get(), ne, cuda_ctx->stream()); - CUDA_CHECK(cudaGetLastError()); - } - - NCCL_CHECK(ncclGroupStart()); - for (size_t i = 0; i < n_backends; ++i) { - ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backends[i]->context; - NCCL_CHECK(ncclAllReduce(tmp[i].get(), tmp[i].get(), ne, ncclBfloat16, ncclSum, info.comms[cuda_ctx->device], cuda_ctx->stream())); - } - NCCL_CHECK(ncclGroupEnd()); - - for (size_t i = 0; i < n_backends; ++i) { - ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backends[i]->context; - - ggml_cuda_set_device(i); - to_fp32(tmp[i].get(), (float *) tensors[i]->data, ne, cuda_ctx->stream()); - CUDA_CHECK(cudaGetLastError()); - } - - return true; -#else - // If NCCL is installed it is used by default for optimal performance. - // However, NVIDIA does not distribute NCCL with CUDA so users may be unwittingly missing this package. - // RCCL is disabled by default, users are explicitly opting in. - // Therefore print no warning for RCCL. -#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) - static bool warning_printed = false; - if (!warning_printed) { - GGML_LOG_WARN("%s: NVIDIA Collective Communications Library (NCCL) is unavailable, multi GPU performance will be suboptimal\n", __func__); - warning_printed = true; - } -#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) - GGML_UNUSED_VARS(backends, tensors, n_backends); - return false; -#endif // GGML_USE_NCCL -} - -ggml_backend_buffer_type_t ggml_backend_cuda_split_buffer_type(int main_device, const float * tensor_split) { - static std::mutex mutex; - std::lock_guard lock(mutex); - - static std::map>, struct ggml_backend_buffer_type> buft_map; - - std::array tensor_split_arr = {}; - - bool all_zero = tensor_split == nullptr || std::all_of(tensor_split, tensor_split + GGML_CUDA_MAX_DEVICES, [](float x) { return x == 0.0f; }); - if (all_zero) { - tensor_split_arr = ggml_cuda_info().default_tensor_split; - } else { - float split_sum = 0.0f; - for (int i = 0; i < ggml_backend_cuda_get_device_count(); ++i) { - tensor_split_arr[i] = split_sum; - split_sum += tensor_split[i]; - } - for (int i = 0; i < ggml_backend_cuda_get_device_count(); ++i) { - tensor_split_arr[i] /= split_sum; - } - } - - auto it = buft_map.find({main_device, tensor_split_arr}); - if (it != buft_map.end()) { - return &it->second; - } - auto * ctx = new ggml_backend_cuda_split_buffer_type_context{ - main_device, - tensor_split_arr, - GGML_CUDA_NAME + std::to_string(main_device) + "_Split", - }; - - struct ggml_backend_buffer_type buft { - /* .iface = */ ggml_backend_cuda_split_buffer_type_interface, - /* .device = */ ggml_backend_reg_dev_get(ggml_backend_cuda_reg(), main_device), - /* .context = */ ctx, - }; - - auto result = buft_map.emplace(std::make_pair(main_device, tensor_split_arr), buft); - return &result.first->second; -} - -// host buffer type - -static const char * ggml_backend_cuda_host_buffer_type_name(ggml_backend_buffer_type_t buft) { - return GGML_CUDA_NAME "_Host"; - - GGML_UNUSED(buft); -} - -static bool ggml_backend_buft_is_cuda_host(ggml_backend_buffer_type_t buft) { - return buft->iface.get_name == ggml_backend_cuda_host_buffer_type_name; -} - -static void ggml_backend_cuda_host_buffer_free_buffer(ggml_backend_buffer_t buffer) { - CUDA_CHECK(cudaFreeHost(buffer->context)); -} - -static void * ggml_cuda_host_malloc(size_t size) { - if (getenv("GGML_CUDA_NO_PINNED") != nullptr) { - return nullptr; - } - - void * ptr = nullptr; - cudaError_t err = cudaMallocHost((void **) &ptr, size); - if (err != cudaSuccess) { - // clear the error - (void)cudaGetLastError(); - GGML_LOG_DEBUG("%s: failed to allocate %.2f MiB of pinned memory: %s\n", __func__, - size / 1024.0 / 1024.0, cudaGetErrorString(err)); - return nullptr; - } - - return ptr; -} - -static ggml_backend_buffer_t ggml_backend_cuda_host_buffer_type_alloc_buffer(ggml_backend_buffer_type_t buft, size_t size) { - void * ptr = ggml_cuda_host_malloc(size); - - if (ptr == nullptr) { - // fallback to cpu buffer - return ggml_backend_buft_alloc_buffer(ggml_backend_cpu_buffer_type(), size); - } - - ggml_backend_buffer_t buffer = ggml_backend_cpu_buffer_from_ptr(ptr, size); - buffer->buft = buft; - buffer->iface.free_buffer = ggml_backend_cuda_host_buffer_free_buffer; - - return buffer; -} - -ggml_backend_buffer_type_t ggml_backend_cuda_host_buffer_type() { - static struct ggml_backend_buffer_type ggml_backend_cuda_buffer_type_host = { - /* .iface = */ { - /* .get_name = */ ggml_backend_cuda_host_buffer_type_name, - /* .alloc_buffer = */ ggml_backend_cuda_host_buffer_type_alloc_buffer, - /* .get_alignment = */ ggml_backend_cpu_buffer_type()->iface.get_alignment, - /* .get_max_size = */ NULL, // defaults to SIZE_MAX - /* .get_alloc_size = */ ggml_backend_cpu_buffer_type()->iface.get_alloc_size, - /* .is_host = */ ggml_backend_cpu_buffer_type()->iface.is_host, - }, - /* .device = */ ggml_backend_reg_dev_get(ggml_backend_cuda_reg(), 0), - /* .context = */ nullptr, - }; - - return &ggml_backend_cuda_buffer_type_host; -} - -//static bool ggml_backend_buffer_is_cuda_host(ggml_backend_buffer_t buffer) { -// return buffer->buft->iface.get_name == ggml_backend_cuda_host_buffer_type_name; -//} - -/// kernels - -typedef void (*ggml_cuda_op_mul_mat_t)( - ggml_backend_cuda_context & ctx, - const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const char * src0_dd_i, const float * src1_ddf_i, - const char * src1_ddq_i, float * dst_dd_i, const int64_t row_low, const int64_t row_high, const int64_t src1_ncols, - const int64_t src1_padded_row_size, cudaStream_t stream); - -#ifndef GGML_CUDA_PEER_MAX_BATCH_SIZE -#define GGML_CUDA_PEER_MAX_BATCH_SIZE 128 -#endif // GGML_CUDA_PEER_MAX_BATCH_SIZE - -#define MUL_MAT_SRC1_COL_STRIDE 128 - -static cudaError_t ggml_cuda_cpy_tensor_2d( - void * dst, const struct ggml_tensor * src, int64_t i3, int64_t i2, int64_t i1_low, int64_t i1_high, cudaStream_t stream) { - - const char * src_ptr = (const char *) src->data; - char * dst_ptr = (char *) dst; - - const int64_t ne0 = src->ne[0]; - const int64_t nb0 = src->nb[0]; - const int64_t nb1 = src->nb[1]; - const int64_t nb2 = src->nb[2]; - const int64_t nb3 = src->nb[3]; - const enum ggml_type type = src->type; - const int64_t ts = ggml_type_size(type); - const int64_t bs = ggml_blck_size(type); - const int64_t i1_diff = i1_high - i1_low; - - const char * x = src_ptr + i1_low*nb1 + i2*nb2 + i3*nb3; - if (nb0 == ts && nb1 == ts*ne0/bs) { - return cudaMemcpyAsync(dst_ptr, x, i1_diff*nb1, cudaMemcpyDeviceToDevice, stream); - } else if (nb0 == ts) { - return cudaMemcpy2DAsync(dst_ptr, ts*ne0/bs, x, nb1, ts*ne0/bs, i1_diff, cudaMemcpyDeviceToDevice, stream); - } else { - for (int64_t i1 = 0; i1 < i1_diff; i1++) { - const void * rx = (const void *) ((const char *) x + i1*nb1); - void * rd = (void *) (dst_ptr + i1*ts*ne0/bs); - // pretend the row is a matrix with cols=1 - cudaError_t r = cudaMemcpy2DAsync(rd, ts/bs, rx, nb0, ts/bs, ne0, cudaMemcpyDeviceToDevice, stream); - if (r != cudaSuccess) { - return r; - } - } - return cudaSuccess; - } -} - -struct cublas_force_compute_type { - bool fp32 = false; - bool fp16 = false; -}; - -static const cublas_force_compute_type & ggml_cuda_cublas_get_force_compute_type() { - static const cublas_force_compute_type compute_type = [] { - cublas_force_compute_type result; - - const bool ggml_cuda_force_cublas_compute_32f_env = getenv("GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F") != nullptr; - const bool ggml_cuda_force_cublas_compute_16f_env = getenv("GGML_CUDA_FORCE_CUBLAS_COMPUTE_16F") != nullptr; - - GGML_ASSERT(ggml_cuda_force_cublas_compute_16f_env == false || ggml_cuda_force_cublas_compute_32f_env == false); - - if (ggml_cuda_force_cublas_compute_32f_env) { - GGML_LOG_INFO("Detected GGML_CUDA_FORCE_CUBLAS_COMPUTE_32F\n"); - result.fp32 = true; - } else if (ggml_cuda_force_cublas_compute_16f_env) { - GGML_LOG_INFO("Detected GGML_CUDA_FORCE_CUBLAS_COMPUTE_16F\n"); - result.fp16 = true; - } - - return result; - }(); - - return compute_type; -} - -static void ggml_cuda_op_mul_mat_cublas( - ggml_backend_cuda_context & ctx, - const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const char * src0_dd_i, const float * src1_ddf_i, - const char * src1_ddq_i, float * dst_dd_i, const int64_t row_low, const int64_t row_high, const int64_t src1_ncols, - const int64_t src1_padded_row_size, cudaStream_t stream) { - - GGML_ASSERT(src0_dd_i != nullptr); - GGML_ASSERT(src1_ddf_i != nullptr); - GGML_ASSERT(dst_dd_i != nullptr); - - const int64_t ne00 = src0->ne[0]; - const int64_t ne10 = src1->ne[0]; - - const int64_t ne0 = dst->ne[0]; - - const int64_t row_diff = row_high - row_low; - - int id = ggml_cuda_get_device(); - - // the main device has a larger memory buffer to hold the results from all GPUs - // ldc == nrows of the matrix that cuBLAS writes into - int64_t ldc = id == ctx.device ? ne0 : row_diff; - - const int cc = ggml_cuda_info().devices[id].cc; - - const bool supports_bf16 = GGML_CUDA_CC_IS_NVIDIA(cc) || GGML_CUDA_CC_IS_AMD(cc) || - (GGML_CUDA_CC_IS_MTHREADS(cc) && cc >= GGML_CUDA_CC_QY2); - - const bool use_fp16 = - src0->type != GGML_TYPE_NVFP4 && - (src0->type == GGML_TYPE_F16 || ggml_is_quantized(src0->type)) && - ggml_is_contiguous(src0) && - row_diff == src0->ne[1] && - dst->op_params[0] == GGML_PREC_DEFAULT; - - if (supports_bf16 && src0->type == GGML_TYPE_BF16 && ggml_is_contiguous(src0) && row_diff == src0->ne[1]) { - ggml_cuda_pool_alloc src1_as_bf16(ctx.pool(id)); - if (src1->type != GGML_TYPE_BF16) { - const to_bf16_cuda_t to_bf16_cuda = ggml_get_to_bf16_cuda(src1->type); - GGML_ASSERT(to_bf16_cuda != nullptr); - size_t ne = src1_ncols*ne10; - src1_as_bf16.alloc(ne); - to_bf16_cuda(src1_ddf_i, src1_as_bf16.get(), ne, stream); - } - const nv_bfloat16 * src1_ptr = src1->type == GGML_TYPE_BF16 ? (const nv_bfloat16 *) src1_ddf_i : src1_as_bf16.get(); - const nv_bfloat16 * src0_ptr = (const nv_bfloat16 *)src0_dd_i; - ggml_cuda_pool_alloc dst_bf16(ctx.pool(id), row_diff*src1_ncols); - - const float alpha_f32 = 1.0f; - const float beta_f32 = 0.0f; - - CUBLAS_CHECK(cublasSetStream(ctx.cublas_handle(id), stream)); - CUBLAS_CHECK( - cublasGemmEx(ctx.cublas_handle(id), CUBLAS_OP_T, CUBLAS_OP_N, - row_diff, src1_ncols, ne10, - &alpha_f32, src0_ptr, CUDA_R_16BF, ne00, - src1_ptr, CUDA_R_16BF, ne10, - &beta_f32, dst_bf16.get(), CUDA_R_16BF, ldc, - CUBLAS_COMPUTE_32F, - CUBLAS_GEMM_DEFAULT_TENSOR_OP)); - - const to_fp32_cuda_t to_fp32_cuda = ggml_get_to_fp32_cuda(GGML_TYPE_BF16); - to_fp32_cuda(dst_bf16.get(), dst_dd_i, row_diff*src1_ncols, stream); - } else if (fast_fp16_hardware_available(cc) && use_fp16) { - // convert src0 and src1 to fp16, multiply as fp16, convert dst to fp32 - ggml_cuda_pool_alloc src0_as_f16(ctx.pool(id)); - if (src0->type != GGML_TYPE_F16) { - const to_fp16_cuda_t to_fp16_cuda = ggml_get_to_fp16_cuda(src0->type); - GGML_ASSERT(to_fp16_cuda != nullptr); - size_t ne = row_diff*ne00; - src0_as_f16.alloc(ne); - to_fp16_cuda(src0_dd_i, src0_as_f16.get(), ne, stream); - } - const half * src0_ptr = src0->type == GGML_TYPE_F16 ? (const half *) src0_dd_i : src0_as_f16.get(); - - ggml_cuda_pool_alloc src1_as_f16(ctx.pool(id)); - if (src1->type != GGML_TYPE_F16) { - const to_fp16_cuda_t to_fp16_cuda = ggml_get_to_fp16_cuda(src1->type); - GGML_ASSERT(to_fp16_cuda != nullptr); - size_t ne = src1_ncols*ne10; - src1_as_f16.alloc(ne); - to_fp16_cuda(src1_ddf_i, src1_as_f16.get(), ne, stream); - } - const half * src1_ptr = src1->type == GGML_TYPE_F16 ? (const half *) src1_ddf_i : src1_as_f16.get(); - - CUBLAS_CHECK(cublasSetStream(ctx.cublas_handle(id), stream)); - - const auto & force_compute_type = ggml_cuda_cublas_get_force_compute_type(); - - if (!force_compute_type.fp16 && (GGML_CUDA_CC_IS_CDNA(cc) - || GGML_CUDA_CC_IS_RDNA4(cc) - || cc == GGML_CUDA_CC_VOLTA - || force_compute_type.fp32)) - { - const float alpha = 1.0f; - const float beta = 0.0f; - CUBLAS_CHECK( - cublasGemmEx(ctx.cublas_handle(id), CUBLAS_OP_T, CUBLAS_OP_N, - row_diff, src1_ncols, ne10, - &alpha, src0_ptr, CUDA_R_16F, ne00, - src1_ptr, CUDA_R_16F, ne10, - &beta, dst_dd_i, CUDA_R_32F, ldc, - CUBLAS_COMPUTE_32F, - CUBLAS_GEMM_DEFAULT_TENSOR_OP)); - } else { - ggml_cuda_pool_alloc dst_f16(ctx.pool(id), row_diff*src1_ncols); - - const half alpha_f16 = 1.0f; - const half beta_f16 = 0.0f; - - CUBLAS_CHECK( - cublasGemmEx(ctx.cublas_handle(id), CUBLAS_OP_T, CUBLAS_OP_N, - row_diff, src1_ncols, ne10, - &alpha_f16, src0_ptr, CUDA_R_16F, ne00, - src1_ptr, CUDA_R_16F, ne10, - &beta_f16, dst_f16.get(), CUDA_R_16F, ldc, - CUBLAS_COMPUTE_16F, - CUBLAS_GEMM_DEFAULT_TENSOR_OP)); - - const to_fp32_cuda_t to_fp32_cuda = ggml_get_to_fp32_cuda(GGML_TYPE_F16); - to_fp32_cuda(dst_f16.get(), dst_dd_i, row_diff*src1_ncols, stream); - } - } else { - ggml_cuda_pool_alloc src0_ddq_as_f32(ctx.pool(id)); - ggml_cuda_pool_alloc src1_ddq_as_f32(ctx.pool(id)); - - if (src0->type != GGML_TYPE_F32) { - const to_fp32_cuda_t to_fp32_cuda = ggml_get_to_fp32_cuda(src0->type); - GGML_ASSERT(to_fp32_cuda != nullptr); - src0_ddq_as_f32.alloc(row_diff*ne00); - to_fp32_cuda(src0_dd_i, src0_ddq_as_f32.get(), row_diff*ne00, stream); - } - if (src1->type != GGML_TYPE_F32) { - const to_fp32_cuda_t to_fp32_cuda = ggml_get_to_fp32_cuda(src1->type); - GGML_ASSERT(to_fp32_cuda != nullptr); - src1_ddq_as_f32.alloc(src1_ncols*ne10); - to_fp32_cuda(src1_ddf_i, src1_ddq_as_f32.get(), src1_ncols*ne10, stream); - } - - const float * src0_ddf_i = src0->type == GGML_TYPE_F32 ? (const float *) src0_dd_i : src0_ddq_as_f32.get(); - const float * src1_ddf1_i = src1->type == GGML_TYPE_F32 ? (const float *) src1_ddf_i : src1_ddq_as_f32.get(); - - const float alpha = 1.0f; - const float beta = 0.0f; - - CUBLAS_CHECK(cublasSetStream(ctx.cublas_handle(id), stream)); - CUBLAS_CHECK( - cublasSgemm(ctx.cublas_handle(id), CUBLAS_OP_T, CUBLAS_OP_N, - row_diff, src1_ncols, ne10, - &alpha, src0_ddf_i, ne00, - src1_ddf1_i, ne10, - &beta, dst_dd_i, ldc)); - } - - GGML_UNUSED_VARS(dst, src1_ddq_i, src1_padded_row_size); -} - -static cudaError_t ggml_cuda_Memcpy2DPeerAsync( - void * dst, int dstDevice, size_t dpitch, void * src, int srcDevice, size_t spitch, size_t width, size_t height, cudaStream_t stream) { - -#if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) - // cudaMemcpy2DAsync may fail with copies between vmm pools of different devices - cudaMemcpy3DPeerParms p = {}; - p.dstDevice = dstDevice; - p.dstPtr = make_cudaPitchedPtr(dst, dpitch, dpitch, height); - p.srcDevice = srcDevice; - p.srcPtr = make_cudaPitchedPtr(src, spitch, spitch, height); - p.extent = make_cudaExtent(width, height, 1); - return cudaMemcpy3DPeerAsync(&p, stream); -#else - // HIP does not support cudaMemcpy3DPeerAsync or vmm pools - GGML_UNUSED(dstDevice); - GGML_UNUSED(srcDevice); - return cudaMemcpy2DAsync(dst, dpitch, src, spitch, width, height, cudaMemcpyDeviceToDevice, stream); -#endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) -} - -static void ggml_cuda_op_mul_mat( - ggml_backend_cuda_context & ctx, - const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, ggml_cuda_op_mul_mat_t op, - quantize_cuda_t quantize_src1) { - - const int64_t ne00 = src0->ne[0]; - const int64_t ne01 = src0->ne[1]; - const int64_t ne02 = src0->ne[2]; - const int64_t ne03 = src0->ne[3]; - - const int64_t ne10 = src1->ne[0]; - const int64_t ne11 = src1->ne[1]; - const int64_t ne12 = src1->ne[2]; - const int64_t ne13 = src1->ne[3]; - const int64_t nrows1 = ggml_nrows(src1); - - const int64_t ne0 = dst->ne[0]; - const int64_t ne1 = dst->ne[1]; - - // const int64_t nb10 = src1->nb[0]; - const int64_t nb11 = src1->nb[1]; - const int64_t nb12 = src1->nb[2]; - const int64_t nb13 = src1->nb[3]; - - const int64_t nb2 = dst->nb[2]; - const int64_t nb3 = dst->nb[3]; - - ggml_backend_cuda_buffer_context * src1_ctx = (ggml_backend_cuda_buffer_context *) src1->buffer->context; - ggml_backend_cuda_buffer_context * dst_ctx = (ggml_backend_cuda_buffer_context *) dst->buffer->context; - - GGML_ASSERT(src1->type == GGML_TYPE_F32 || (src1->ne[2] == 1 && src1->ne[3] == 1)); - - GGML_ASSERT(ne12 % ne02 == 0); - GGML_ASSERT(ne13 % ne03 == 0); - - const int64_t i02_divisor = ne12 / ne02; - const int64_t i03_divisor = ne13 / ne03; - - const size_t src0_ts = ggml_type_size(src0->type); - const size_t src0_bs = ggml_blck_size(src0->type); - const size_t q8_1_ts = sizeof(block_q8_1); - const size_t q8_1_bs = QK8_1; - - const bool src0_is_contiguous = ggml_is_contiguous(src0); - const bool src1_is_contiguous = ggml_is_contiguous(src1); - - const int64_t src1_padded_col_size = GGML_PAD(ne10, MATRIX_ROW_PADDING); - - const bool split = ggml_backend_buft_is_cuda_split(src0->buffer->buft); - GGML_ASSERT(!(split && ne02 > 1)); - GGML_ASSERT(!(split && ne03 > 1)); - GGML_ASSERT(!(split && ne02 < ne12)); - GGML_ASSERT(!(split && ne03 < ne13)); - - ggml_tensor_extra_gpu * src0_extra = split ? (ggml_tensor_extra_gpu *) src0->extra : nullptr; - - - std::array tensor_split; - if (split) { - ggml_backend_cuda_split_buffer_type_context * buft_ctx = (ggml_backend_cuda_split_buffer_type_context *) src0->buffer->buft->context; - tensor_split = buft_ctx->tensor_split; - } - - struct dev_data { - int cc; - - ggml_cuda_pool_alloc src0_dd_alloc; - ggml_cuda_pool_alloc src1_ddf_alloc; - ggml_cuda_pool_alloc src1_ddq_alloc; - ggml_cuda_pool_alloc dst_dd_alloc; - - char * src0_dd = nullptr; - float * src1_ddf = nullptr; // float - char * src1_ddq = nullptr; // q8_1 - float * dst_dd = nullptr; - - int64_t row_low; - int64_t row_high; - }; - - dev_data dev[GGML_CUDA_MAX_DEVICES]; - - int used_devices = 0; - - for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { - dev[id].cc = ggml_cuda_info().devices[id].cc; - - // by default, use all rows - dev[id].row_low = 0; - dev[id].row_high = ne01; - - // for multi GPU, get the row boundaries from tensor split - // and round to mul_mat_q tile sizes - if (split) { - const int64_t rounding = get_row_rounding(tensor_split); - - if (id != 0) { - dev[id].row_low = ne01*tensor_split[id]; - if (dev[id].row_low < ne01) { - dev[id].row_low -= dev[id].row_low % rounding; - } - } - - if (id != ggml_backend_cuda_get_device_count() - 1) { - dev[id].row_high = ne01*tensor_split[id + 1]; - if (dev[id].row_high < ne01) { - dev[id].row_high -= dev[id].row_high % rounding; - } - } - } - } - - for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { - if ((!split && id != ctx.device) || dev[id].row_low == dev[id].row_high) { - continue; - } - - used_devices++; - - const bool src1_on_device = id == src1_ctx->device; - const bool dst_on_device = id == dst_ctx->device; - - ggml_cuda_set_device(id); - cudaStream_t stream = ctx.stream(id, 0); - - if (src0_is_contiguous) { - dev[id].src0_dd = split ? (char *) src0_extra->data_device[id] : (char *) src0->data; - } else { - // If src0 is not contiguous it will be copied to a temporary buffer. - // This buffer needs to be cleared entirely because multiple regions will function as padding. - const size_t nbytes_data = ggml_nbytes(src0); - const size_t nbytes_padding = ggml_row_size(src0->type, MATRIX_ROW_PADDING - ne00 % MATRIX_ROW_PADDING); - dev[id].src0_dd = dev[id].src0_dd_alloc.alloc(ctx.pool(id), nbytes_data + nbytes_padding); - CUDA_CHECK(cudaMemsetAsync(dev[id].src0_dd, 0, nbytes_data + nbytes_padding, stream)); - } - - // If src0 is on a temporary compute buffer (partial offloading) there may be some padding that needs to be cleared: - if (ne00 % MATRIX_ROW_PADDING != 0 && ggml_is_quantized(src0->type) && ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE && src0->view_src == nullptr) { - GGML_ASSERT(ggml_is_contiguously_allocated(src0)); - GGML_ASSERT(!src0->view_src); - const size_t nbytes_data = ggml_row_size(src0->type, (dev[id].row_high - dev[id].row_low)*ne00); - const size_t nbytes_padding = ggml_row_size(src0->type, MATRIX_ROW_PADDING - ne00 % MATRIX_ROW_PADDING); - CUDA_CHECK(cudaMemsetAsync(dev[id].src0_dd + nbytes_data, 0, nbytes_padding, stream)); - } - - if (src1_on_device && src1_is_contiguous) { - dev[id].src1_ddf = (float *) src1->data; - } else { - dev[id].src1_ddf = dev[id].src1_ddf_alloc.alloc(ctx.pool(id), ggml_nelements(src1)); - } - - if (quantize_src1) { - size_t src_1_ddq_size = nrows1*src1_padded_col_size*q8_1_ts/q8_1_bs; - if (quantize_src1 == quantize_mmq_q8_1_cuda) { - src_1_ddq_size += get_mmq_x_max_host(dev[id].cc)*sizeof(block_q8_1_mmq); - } - dev[id].src1_ddq = dev[id].src1_ddq_alloc.alloc(ctx.pool(id), src_1_ddq_size); - - if (src1_on_device && src1_is_contiguous) { - quantize_src1( - dev[id].src1_ddf, nullptr, dev[id].src1_ddq, src0->type, ne10, - nb11/sizeof(float), nb12/sizeof(float), nb13/sizeof(float), - src1_padded_col_size, ne11, ne12, ne13, stream); - CUDA_CHECK(cudaGetLastError()); - } - } - - if (dst_on_device) { - dev[id].dst_dd = (float *) dst->data; - } else { - const size_t size_dst_ddf = split ? (dev[id].row_high - dev[id].row_low)*ne1 : ggml_nelements(dst); - dev[id].dst_dd = dev[id].dst_dd_alloc.alloc(ctx.pool(id), size_dst_ddf); - } - } - - // if multiple devices are used they need to wait for the main device - // here an event is recorded that signals that the main device has finished calculating the input data - if (split && used_devices > 1) { - ggml_cuda_set_device(ctx.device); - CUDA_CHECK(cudaEventRecord(src0_extra->events[ctx.device][0], ctx.stream())); - } - - const int64_t src1_col_stride = split && used_devices > 1 ? MUL_MAT_SRC1_COL_STRIDE : ne11; - for (int64_t src1_col_0 = 0; src1_col_0 < ne11; src1_col_0 += src1_col_stride) { - const int64_t is = split ? (src1_col_0/src1_col_stride) % GGML_CUDA_MAX_STREAMS : 0; - const int64_t src1_ncols = src1_col_0 + src1_col_stride > ne11 ? ne11 - src1_col_0 : src1_col_stride; - - for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { - if ((!split && id != ctx.device) || dev[id].row_low == dev[id].row_high) { - continue; - } - - const bool src1_on_device = id == src1_ctx->device; - const bool dst_on_device = id == dst_ctx->device; - const int64_t row_diff = dev[id].row_high - dev[id].row_low; - - ggml_cuda_set_device(id); - cudaStream_t stream = ctx.stream(id, is); - - // wait for main GPU data if necessary - if (split && (id != ctx.device || is != 0)) { - CUDA_CHECK(cudaStreamWaitEvent(stream, src0_extra->events[ctx.device][0], 0)); - } - - for (int64_t i0 = 0; i0 < ne13*ne12; ++i0) { - const int64_t i03 = i0 / ne12; - const int64_t i02 = i0 % ne12; - - size_t src1_ddq_i_offset = i0*ne11 * src1_padded_col_size*q8_1_ts/q8_1_bs; - if (quantize_src1 == quantize_mmq_q8_1_cuda) { - src1_ddq_i_offset += src1_col_0 * sizeof(block_q8_1_mmq); - } else { - src1_ddq_i_offset += src1_col_0 * src1_padded_col_size*q8_1_ts/q8_1_bs; - } - - // for split tensors the data begins at i0 == i0_offset_low - const size_t nbytes_src0_matrix = ne01*ne00*src0_ts / src0_bs; - char * src0_dd_i = dev[id].src0_dd + ((i03/i03_divisor)*ne02 + (i02/i02_divisor)) * nbytes_src0_matrix; - float * src1_ddf_i = dev[id].src1_ddf + (i0*ne11 + src1_col_0) * ne10; - char * src1_ddq_i = dev[id].src1_ddq + src1_ddq_i_offset; - float * dst_dd_i = dev[id].dst_dd + (i0*ne1 + src1_col_0) * (dst_on_device ? ne0 : row_diff); - - // the main device memory buffer can be on VRAM scratch, with space for all partial results - // in that case an offset on dst_ddf_i is needed - if (id == ctx.device) { - dst_dd_i += dev[id].row_low; // offset is 0 if no tensor split - } - - // copy src0, src1 to device if necessary - if (src1_is_contiguous) { - if (id != ctx.device) { - if (quantize_src1) { - char * src1_ddq_i_source = dev[ctx.device].src1_ddq + src1_ddq_i_offset; - if (quantize_src1 == quantize_mmq_q8_1_cuda) { - const size_t pitch = ne11*sizeof(block_q8_1_mmq); - const size_t width = src1_ncols*sizeof(block_q8_1_mmq); - const size_t height = src1_padded_col_size/(4*QK8_1); - CUDA_CHECK(ggml_cuda_Memcpy2DPeerAsync(src1_ddq_i, id, pitch, src1_ddq_i_source, ctx.device, pitch, width, height, stream)); - } else { - CUDA_CHECK(cudaMemcpyPeerAsync( - src1_ddq_i, id, src1_ddq_i_source, ctx.device, src1_ncols*src1_padded_col_size*q8_1_ts/q8_1_bs, stream)); - } - } else { - float * src1_ddf_i_source = (float *) src1->data; - src1_ddf_i_source += (i0*ne11 + src1_col_0) * ne10; - CUDA_CHECK(cudaMemcpyPeerAsync(src1_ddf_i, id, src1_ddf_i_source, ctx.device, - src1_ncols*ne10*sizeof(float), stream)); - } - } - } else if (src1_on_device && !src1_is_contiguous) { - CUDA_CHECK(ggml_cuda_cpy_tensor_2d( - src1_ddf_i, src1, i03, i02, src1_col_0, src1_col_0+src1_ncols, stream)); - } else { - GGML_ABORT("fatal error"); - } - - if (quantize_src1 && !src1_is_contiguous) { - quantize_src1( - src1_ddf_i, nullptr, src1_ddq_i, src0->type, ne10, ne10, ne11*ne10, ne12*ne11*ne10, - src1_padded_col_size, src1_ncols, 1, 1, stream); - CUDA_CHECK(cudaGetLastError()); - } - - if (src1_col_0 == 0 && !src0_is_contiguous && i03 % i03_divisor == 0 && i02 % i02_divisor == 0) { - CUDA_CHECK(ggml_cuda_cpy_tensor_2d( - src0_dd_i, src0, i03/i03_divisor, i02/i02_divisor, dev[id].row_low, dev[id].row_high, stream)); - } - - // do the computation - op(ctx, src0, src1, dst, src0_dd_i, src1_ddf_i, src1_ddq_i, dst_dd_i, - dev[id].row_low, dev[id].row_high, src1_ncols, src1_padded_col_size, stream); - CUDA_CHECK(cudaGetLastError()); - - // copy dst to host or other device if necessary - if (!dst_on_device) { - void * dst_off_device = dst->data; - if (split) { - // src0 = weight matrix is saved as a transposed matrix for better memory layout. - // dst is NOT transposed. - // The outputs of matrix matrix multiplications can therefore NOT simply be concatenated for >1 GPU. - // Instead they need to be copied to the correct slice in ne0 = dst row index. - // If dst is a vector with ne0 == 1 then you don't have to do this but it still produces correct results. - float * dhf_dst_i = (float *) ((char *) dst_off_device + i02*nb2 + i03*nb3); - GGML_ASSERT(dst->nb[1] == ne0*sizeof(float)); - dhf_dst_i += src1_col_0*ne0 + dev[id].row_low; - CUDA_CHECK(ggml_cuda_Memcpy2DPeerAsync( - dhf_dst_i, ctx.device, ne0*sizeof(float), dst_dd_i, id, row_diff*sizeof(float), row_diff*sizeof(float), src1_ncols, stream)); - } else { - float * dhf_dst_i = (float *) ((char *) dst_off_device + i02*nb2 + i03*nb3); - GGML_ASSERT(dst->nb[1] == ne0*sizeof(float)); - dhf_dst_i += src1_col_0*ne0; - CUDA_CHECK(cudaMemcpyAsync(dhf_dst_i, dst_dd_i, src1_ncols*ne0*sizeof(float), cudaMemcpyDeviceToDevice, stream)); - } - } - - // add event for the main device to wait on until other device is done - if (split && (id != ctx.device || is != 0)) { - CUDA_CHECK(cudaEventRecord(src0_extra->events[id][is], stream)); - } - } - } - } - - // main device waits for all other devices to be finished - if (split && ggml_backend_cuda_get_device_count() > 1) { - int64_t is_max = (ne11 + MUL_MAT_SRC1_COL_STRIDE - 1) / MUL_MAT_SRC1_COL_STRIDE; - is_max = is_max <= GGML_CUDA_MAX_STREAMS ? is_max : GGML_CUDA_MAX_STREAMS; - - ggml_cuda_set_device(ctx.device); - for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { - if (dev[id].row_low == dev[id].row_high) { - continue; - } - for (int64_t is = 0; is < is_max; ++is) { - CUDA_CHECK(cudaStreamWaitEvent(ctx.stream(), src0_extra->events[id][is], 0)); - } - } - } -} - -static __global__ void k_compute_batched_ptrs( - const void * src0_as_f16, const void * src1_as_f16, char * dst, - const void ** ptrs_src, void ** ptrs_dst, - int64_t ne12, int64_t ne13, - int64_t ne23, - size_t nb02, size_t nb03, - size_t nb12, size_t nb13, - size_t nbd2, size_t nbd3, - int64_t r2, int64_t r3) { - const int64_t i13 = blockIdx.x * blockDim.x + threadIdx.x; - const int64_t i12 = blockIdx.y * blockDim.y + threadIdx.y; - - if (i13 >= ne13 || i12 >= ne12) { - return; - } - - const int64_t i03 = i13 / r3; - const int64_t i02 = i12 / r2; - - ptrs_src[0*ne23 + i12 + i13*ne12] = (const char *) src0_as_f16 + i02*nb02 + i03*nb03; - ptrs_src[1*ne23 + i12 + i13*ne12] = (const char *) src1_as_f16 + i12*nb12 + i13*nb13; - ptrs_dst[0*ne23 + i12 + i13*ne12] = ( char *) dst + i12*nbd2 + i13*nbd3; -} - -// Type traits for mapping ggml types to CUDA/cuBLAS types -template -struct batched_mul_mat_traits; - -template<> -struct batched_mul_mat_traits { - using cuda_type = float; - static inline const cublasComputeType_t compute_type = CUBLAS_COMPUTE_32F; - static inline const cudaDataType_t data_type = CUDA_R_32F; - static inline const ggml_type ggml_type_val = GGML_TYPE_F32; - static inline const float alpha = 1.0f; - static inline const float beta = 0.0f; - static inline const void* get_alpha() { static const float val = alpha; return &val; } - static inline const void* get_beta() { static const float val = beta; return &val; } - static inline auto get_nc_converter(ggml_type src_type) { return ggml_get_to_fp32_nc_cuda(src_type); } -}; - -template<> -struct batched_mul_mat_traits { - using cuda_type = nv_bfloat16; - static inline const cublasComputeType_t compute_type = CUBLAS_COMPUTE_32F; - static inline const cudaDataType_t data_type = CUDA_R_16BF; - static inline const ggml_type ggml_type_val = GGML_TYPE_BF16; - static inline const float alpha = 1.0f; - static inline const float beta = 0.0f; - static inline const void* get_alpha() { static const float val = alpha; return &val; } - static inline const void* get_beta() { static const float val = beta; return &val; } - static inline auto get_nc_converter(ggml_type src_type) { return ggml_get_to_bf16_nc_cuda(src_type); } -}; - -template<> -struct batched_mul_mat_traits { - using cuda_type = half; - static inline const cublasComputeType_t compute_type = CUBLAS_COMPUTE_16F; - static inline const cudaDataType_t data_type = CUDA_R_16F; - static inline const ggml_type ggml_type_val = GGML_TYPE_F16; - static inline const half alpha = 1.0; - static inline const half beta = 0.0; - static inline const void* get_alpha() { static const half val = alpha; return &val; } - static inline const void* get_beta() { static const half val = beta; return &val; } - static inline auto get_nc_converter(ggml_type src_type) { return ggml_get_to_fp16_nc_cuda(src_type); } -}; - -template -static void ggml_cuda_mul_mat_batched_cublas_impl(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { - using traits = batched_mul_mat_traits; - using cuda_t = typename traits::cuda_type; - - GGML_ASSERT(!ggml_is_transposed(src0)); - GGML_ASSERT(!ggml_is_transposed(src1)); - GGML_ASSERT(!ggml_backend_buft_is_cuda_split(src0->buffer->buft)); - GGML_ASSERT(src0->type == src0_type); - GGML_ASSERT(ggml_is_contiguous(dst)); - - // Byte offsets and tensor dimensions are currently used in an inconsistent way for dst. - // As long as dst is contiguous this does not matter though. - - GGML_TENSOR_BINARY_OP_LOCALS - - const int64_t ne_dst = ggml_nelements(dst); - cudaStream_t main_stream = ctx.stream(); - CUBLAS_CHECK(cublasSetStream(ctx.cublas_handle(), main_stream)); - - float * dst_ddf = (float *) dst->data; - const size_t ts_src1 = ggml_type_size(src1->type); - GGML_ASSERT(nb10 == ts_src1); - int64_t s11 = nb11 / ts_src1; - int64_t s12 = nb12 / ts_src1; - int64_t s13 = nb13 / ts_src1; - - const cuda_t * src0_ptr = nullptr; - const cuda_t * src1_ptr = nullptr; - - ggml_cuda_pool_alloc src0_alloc(ctx.pool()); - ggml_cuda_pool_alloc src1_alloc(ctx.pool()); - - bool is_src0_cont_2 = ggml_is_contiguous_2(src0); - bool is_src1_cont_2 = ggml_is_contiguous_2(src1); - - // Handle src0 - src0_ptr = (const cuda_t *) src0->data; - - // Handle src1 - convert if necessary - if (src1->type == src0_type) { - src1_ptr = (const cuda_t *) src1->data; - } else { - // Convert src1 to target type using traits conversion functions - const int64_t ne_src1 = ggml_nelements(src1); - src1_alloc.alloc(ne_src1); - - const auto convert_func = traits::get_nc_converter(src1->type); - GGML_ASSERT(convert_func != nullptr); - convert_func(src1->data, src1_alloc.get(), ne10, ne11, ne12, ne13, s11, s12, s13, main_stream); - src1_ptr = src1_alloc.get(); - s11 = ne10; - s12 = ne11*s11; - s13 = ne12*s12; - - is_src1_cont_2 = true; - } - - // Setup destination buffer - ggml_cuda_pool_alloc dst_temp(ctx.pool()); - char * dst_t; - size_t nbd2 = dst->nb[2]; - size_t nbd3 = dst->nb[3]; - - cublasComputeType_t cu_compute_type = traits::compute_type; - cudaDataType_t cu_data_type = traits::data_type; - cudaDataType_t cu_data_type_a = traits::data_type; - cudaDataType_t cu_data_type_b = traits::data_type; - const void * alpha = traits::get_alpha(); - const void * beta = traits::get_beta(); - - const auto & force_compute_type = ggml_cuda_cublas_get_force_compute_type(); - - int id = ggml_cuda_get_device(); - const int cc = ggml_cuda_info().devices[id].cc; - static constexpr bool is_src0_type_f16 = src0_type == GGML_TYPE_F16; - - // bf16 and fp32 are already being computed in fp32 (ensure it using static_assert), - // so checking necessity of forced fp32 only for fp16 src0_type - static_assert(is_src0_type_f16 || traits::compute_type == CUBLAS_COMPUTE_32F); - - const bool need_compute_32f = is_src0_type_f16 && !force_compute_type.fp16 && (GGML_CUDA_CC_IS_CDNA(cc) - || GGML_CUDA_CC_IS_RDNA4(cc) - || cc == GGML_CUDA_CC_VOLTA - || force_compute_type.fp32); - - if (dst->op_params[0] == GGML_PREC_DEFAULT && !need_compute_32f) { - if constexpr (src0_type == GGML_TYPE_F32) { - dst_t = (char *) dst_ddf; // Direct F32 output - } else { - dst_t = (char *) dst_temp.alloc(ne_dst); - nbd2 /= sizeof(float) / sizeof(cuda_t); - nbd3 /= sizeof(float) / sizeof(cuda_t); - } - } else { - dst_t = (char *) dst_ddf; - cu_compute_type = batched_mul_mat_traits::compute_type; - cu_data_type = batched_mul_mat_traits::data_type; - alpha = batched_mul_mat_traits::get_alpha(); - beta = batched_mul_mat_traits::get_beta(); - } - - GGML_ASSERT(ne12 % ne02 == 0); - GGML_ASSERT(ne13 % ne03 == 0); - - // broadcast factors - const int64_t r2 = ne12/ne02; - const int64_t r3 = ne13/ne03; - - if (r2 == 1 && r3 == 1 && is_src0_cont_2 && is_src1_cont_2) { - // with a [0, 2, 1, 3] perm. and ne02==1 the matrix strides need to be determined from dim 3: - const int64_t sma = ne02 == 1 ? nb03/nb00 : nb02/nb00; - const int64_t smb = ne12 == 1 ? s13 : s12; - - // there is no broadcast and src0, src1 are contiguous across dims 2, 3 - // use cublasGemmStridedBatchedEx - CUBLAS_CHECK( - cublasGemmStridedBatchedEx(ctx.cublas_handle(), CUBLAS_OP_T, CUBLAS_OP_N, - ne01, ne11, ne10, - alpha, src0_ptr, cu_data_type_a, nb01/nb00, sma, // strideA - src1_ptr, cu_data_type_b, s11, smb, // strideB - beta, dst_t, cu_data_type, ne0, ne1*ne0, // strideC - ne12*ne13, - cu_compute_type, - CUBLAS_GEMM_DEFAULT_TENSOR_OP)); - } else { - // use cublasGemmBatchedEx - const int64_t ne23 = ne12*ne13; - - ggml_cuda_pool_alloc ptrs_src(ctx.pool(), 2*ne23); - ggml_cuda_pool_alloc< void *> ptrs_dst(ctx.pool(), 1*ne23); - - size_t src1_stride_size = sizeof(cuda_t); - - const int threads_x = 16; - const int threads_y = 16; - dim3 block_dims(threads_x, threads_y); - - dim3 grid_dims( - (ne13 + threads_x - 1) / threads_x, - (ne12 + threads_y - 1) / threads_y - ); - k_compute_batched_ptrs<<>>( - src0_ptr, src1_ptr, dst_t, - ptrs_src.get(), ptrs_dst.get(), - ne12, ne13, - ne23, - nb02, nb03, - (src1->type == src0_type) ? nb12 : s12*src1_stride_size, - (src1->type == src0_type) ? nb13 : s13*src1_stride_size, - nbd2, nbd3, - r2, r3); - - CUDA_CHECK(cudaGetLastError()); - - CUBLAS_CHECK( - cublasGemmBatchedEx(ctx.cublas_handle(), CUBLAS_OP_T, CUBLAS_OP_N, - ne01, ne11, ne10, - alpha, (const void **) (ptrs_src.get() + 0*ne23), cu_data_type_a, nb01/nb00, - (const void **) (ptrs_src.get() + 1*ne23), cu_data_type_b, s11, - beta, ( void **) (ptrs_dst.get() + 0*ne23), cu_data_type, ne0, - ne23, - cu_compute_type, - CUBLAS_GEMM_DEFAULT_TENSOR_OP)); - } - - // Convert output back to F32 if needed - if (dst->op_params[0] == GGML_PREC_DEFAULT && cu_data_type != CUDA_R_32F) { - const to_fp32_cuda_t to_fp32_cuda = ggml_get_to_fp32_cuda(traits::ggml_type_val); - to_fp32_cuda(dst_temp.get(), dst_ddf, ne_dst, main_stream); - } -} - -static void ggml_cuda_mul_mat_batched_cublas(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { - GGML_ASSERT(src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_BF16 || src0->type == GGML_TYPE_F32); - - switch (src0->type) { - case GGML_TYPE_F32: - ggml_cuda_mul_mat_batched_cublas_impl(ctx, src0, src1, dst); - break; - case GGML_TYPE_BF16: - ggml_cuda_mul_mat_batched_cublas_impl(ctx, src0, src1, dst); - break; - case GGML_TYPE_F16: - ggml_cuda_mul_mat_batched_cublas_impl(ctx, src0, src1, dst); - break; - default: - GGML_ABORT("Unsupported type"); - } -} - -static bool ggml_cuda_should_fuse_mul_mat(const ggml_tensor * ffn_up, - const ggml_tensor * ffn_gate, - const ggml_tensor * glu, - const ggml_tensor * ffn_up_bias = nullptr, - const ggml_tensor * ffn_gate_bias = nullptr) { - const bool has_bias = ffn_up_bias != nullptr || ffn_gate_bias != nullptr; - - if (has_bias && (!ffn_up_bias || !ffn_gate_bias)) { - return false; - } - - const bool is_mul_mat = ffn_up->op == GGML_OP_MUL_MAT && ffn_gate->op == GGML_OP_MUL_MAT && glu->op == GGML_OP_GLU; - const bool is_mul_mat_id = ffn_up->op == GGML_OP_MUL_MAT_ID && ffn_gate->op == GGML_OP_MUL_MAT_ID && glu->op == GGML_OP_GLU; - - GGML_ASSERT(ffn_up && ffn_gate && glu); - - if (!is_mul_mat && !is_mul_mat_id) { - return false; - } - - const ggml_op expected_bias_op = is_mul_mat ? GGML_OP_ADD : GGML_OP_ADD_ID; - - if (has_bias) { - if (ffn_up_bias->op != expected_bias_op || ffn_gate_bias->op != expected_bias_op) { - return false; - } - - if (glu->src[0] != ffn_gate_bias || glu->src[1] != ffn_up_bias) { - return false; - } - - if (expected_bias_op == GGML_OP_ADD) { - const bool up_has_mul = ffn_up_bias->src[0] == ffn_up || ffn_up_bias->src[1] == ffn_up; - const bool gate_has_mul = ffn_gate_bias->src[0] == ffn_gate || ffn_gate_bias->src[1] == ffn_gate; - if (!up_has_mul || !gate_has_mul) { - return false; - } - } else { // GGML_OP_ADD_ID - if (ffn_up_bias->src[0] != ffn_up || ffn_gate_bias->src[0] != ffn_gate) { - return false; - } - if (ffn_up_bias->src[2] != ffn_up->src[2] || ffn_gate_bias->src[2] != ffn_gate->src[2]) { - return false; - } - } - } else { - if (glu->src[0] != ffn_gate && glu->src[1] != ffn_up) { - return false; - } - } - - if (ffn_up->src[0]->type != ffn_gate->src[0]->type || !ggml_are_same_shape(ffn_up->src[0], ffn_gate->src[0]) || - !ggml_are_same_stride(ffn_up->src[0], ffn_gate->src[0])) { - return false; - } - - if (ffn_up->src[1] != ffn_gate->src[1]) { - return false; - } - - if (ffn_up->src[2] && (ffn_up->src[2] != ffn_gate->src[2])) { - return false; - } - - static constexpr std::array valid_glu_ops = { - GGML_GLU_OP_SWIGLU, - GGML_GLU_OP_GEGLU, - GGML_GLU_OP_SWIGLU_OAI, - GGML_GLU_OP_SWIGLU_DS4, - }; - - if (std::find(valid_glu_ops.begin(), valid_glu_ops.end(), ggml_get_glu_op(glu)) == valid_glu_ops.end()) { - return false; - } - - if (const bool swapped = ggml_get_op_params_i32(glu, 1); swapped) { - return false; - } - - const bool split = ggml_backend_buft_is_cuda_split(ffn_up->src[0]->buffer->buft) || - ggml_backend_buft_is_cuda_split(ffn_gate->src[0]->buffer->buft); - - //TODO: add support for fusion for split buffers - if (split) { - return false; - } - - return true; -} - -static bool ggml_cuda_should_fuse_mul_mat_vec_f(const ggml_tensor * tensor) { - ggml_tensor * src0 = tensor->src[0]; - ggml_tensor * src1 = tensor->src[1]; - const ggml_tensor * dst = tensor; - - const bool is_mul_mat_id = tensor->op == GGML_OP_MUL_MAT_ID; - - bool use_mul_mat_vec_f = - (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_BF16) && - src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32; - - const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; - use_mul_mat_vec_f = use_mul_mat_vec_f && ggml_cuda_should_use_mmvf(src0->type, cc, src0->ne, src0->nb, is_mul_mat_id ? src1->ne[2] : src1->ne[1]); - - const bool split = ggml_backend_buft_is_cuda_split(src0->buffer->buft) || - ggml_backend_buft_is_cuda_split(src1->buffer->buft); - - //TODO: add support for fusion for split buffers - if (split) { - return false; - } - - //we only support fusion for ncols_dst = 1 - if (tensor->op == GGML_OP_MUL_MAT && dst->ne[1] != 1) { - return false; - } - - if (tensor->op == GGML_OP_MUL_MAT_ID && dst->ne[2] != 1) { - return false; - } - - - return use_mul_mat_vec_f; -} - -static inline void ggml_cuda_set_fusion_glu_params(ggml_cuda_mm_fusion_args_host & fusion_data, const ggml_tensor * glu) { - fusion_data.glu_op = ggml_get_glu_op(glu); - fusion_data.glu_param0 = ggml_get_op_params_f32(glu, 2); - fusion_data.glu_param1 = ggml_get_op_params_f32(glu, 3); -} - -static bool ggml_cuda_should_fuse_mul_mat_vec_q(const ggml_tensor * tensor) { - ggml_tensor * src0 = tensor->src[0]; - ggml_tensor * src1 = tensor->src[1]; - const ggml_tensor * dst = tensor; - const bool is_mul_mat_id = tensor->op == GGML_OP_MUL_MAT_ID; - - const bool bad_padding_clear = ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE && - ggml_nbytes(src0) != ggml_backend_buffer_get_alloc_size(src0->buffer, src0) && - src0->view_src; - - const int64_t ncols_dst = is_mul_mat_id ? dst->ne[2] : src1->ne[1]; - bool use_mul_mat_vec_q = ggml_is_quantized(src0->type) && !bad_padding_clear && src1->type == GGML_TYPE_F32 && - dst->type == GGML_TYPE_F32 && ncols_dst <= MMVQ_MAX_BATCH_SIZE; - - // fusion is not universally faster on Pascal - const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; - if (cc <= GGML_CUDA_CC_PASCAL) { - return false; - } - if (is_mul_mat_id) { - const int mmvq_mmid_max = get_mmvq_mmid_max_batch(src0->type, cc); - if (dst->ne[2] > mmvq_mmid_max) { - return false; - } - } - if (tensor->op == GGML_OP_MUL_MAT && dst->ne[1] != 1) { - return false; - } - const bool split = ggml_backend_buft_is_cuda_split(src0->buffer->buft) || - ggml_backend_buft_is_cuda_split(src1->buffer->buft); - - //TODO: add support for fusion for split buffers - if (split) { - return false; - } - - return use_mul_mat_vec_q; -} - -static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { - const bool split = ggml_backend_buft_is_cuda_split(src0->buffer->buft); - - // If src0 is a temporary compute buffer it may have some padding that needs to be cleared for mul_mat_vec_q or mul_mat_q. - // But if src0 is also a view of another tensor then this cannot be done safely because it may overwrite valid tensor data. - // Therefore, in such cases use cuBLAS. - const bool bad_padding_clear = ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE - && ggml_nbytes(src0) != ggml_backend_buffer_get_alloc_size(src0->buffer, src0) && src0->view_src; - - bool use_mul_mat_vec_f = (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_BF16) - && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32; - bool use_mul_mat_f = !ggml_is_quantized(src0->type) - && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32; - // LUCE_MMVQ_MAX_NCOLS: lower the MMVQ ncols ceiling for plain mul_mat so - // that MMQ takes small multi-token batches (spec-decode verify widths). - static const int luce_mmvq_max_ncols = []() { - const char * e = getenv("LUCE_MMVQ_MAX_NCOLS"); - const int v = e ? atoi(e) : MMVQ_MAX_BATCH_SIZE; - return v > 0 ? v : MMVQ_MAX_BATCH_SIZE; - }(); - bool use_mul_mat_vec_q = ggml_is_quantized(src0->type) && !bad_padding_clear - && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32 - && src1->ne[1] <= luce_mmvq_max_ncols; - bool use_mul_mat_q = ggml_is_quantized(src0->type) && !bad_padding_clear - && src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32; - - bool any_gpus_with_slow_fp16 = false; - - if (split) { - ggml_backend_cuda_split_buffer_type_context * buft_ctx = (ggml_backend_cuda_split_buffer_type_context *) src0->buffer->buft->context; - auto & tensor_split = buft_ctx->tensor_split; - for (int id = 0; id < ggml_backend_cuda_get_device_count(); ++id) { - // skip devices that are not going to do any work: - if (tensor_split[id] >= (id + 1 < ggml_backend_cuda_get_device_count() ? tensor_split[id + 1] : 1.0f)) { - continue; - } - - const int cc = ggml_cuda_info().devices[id].cc; - const int warp_size = ggml_cuda_info().devices[id].warp_size; - use_mul_mat_q = use_mul_mat_q && ggml_cuda_should_use_mmq(src0->type, cc, src1->ne[1], /*n_experts=*/0); - use_mul_mat_f = use_mul_mat_f && ggml_cuda_should_use_mmf(src0->type, cc, warp_size, src0->ne, src0->nb, src1->ne[1], /*mul_mat_id=*/false); - use_mul_mat_vec_f = use_mul_mat_vec_f && ggml_cuda_should_use_mmvf(src0->type, cc, src0->ne, src0->nb, src1->ne[1]); - any_gpus_with_slow_fp16 = any_gpus_with_slow_fp16 || !fast_fp16_hardware_available(cc); - } - } else { - const int cc = ggml_cuda_info().devices[ctx.device].cc; - const int warp_size = ggml_cuda_info().devices[ctx.device].warp_size; - use_mul_mat_q = use_mul_mat_q && ggml_cuda_should_use_mmq(src0->type, cc, src1->ne[1], /*n_experts=*/0); - use_mul_mat_f = use_mul_mat_f && ggml_cuda_should_use_mmf(src0->type, cc, warp_size, src0->ne, src0->nb, src1->ne[1], /*mul_mat_id=*/false); - use_mul_mat_vec_f = use_mul_mat_vec_f && ggml_cuda_should_use_mmvf(src0->type, cc, src0->ne, src0->nb, src1->ne[1]); - any_gpus_with_slow_fp16 = any_gpus_with_slow_fp16 || !fast_fp16_hardware_available(cc); - } - - // debug helpers - //printf("src0: %8d %8d %8d %8d\n", src0->ne[0], src0->ne[1], src0->ne[2], src0->ne[3]); - //printf(" %8d %8d %8d %8d\n", src0->nb[0], src0->nb[1], src0->nb[2], src0->nb[3]); - //printf("src1: %8d %8d %8d %8d\n", src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3]); - //printf(" %8d %8d %8d %8d\n", src1->nb[0], src1->nb[1], src1->nb[2], src1->nb[3]); - //printf("src0 is contiguous %d, transposed %d, type = %s, name = %s\n", ggml_is_contiguous(src0), ggml_is_transposed(src0), ggml_type_name(src0->type), src0->name); - //printf("src1 is contiguous %d, transposed %d, type = %s, name = %s\n", ggml_is_contiguous(src1), ggml_is_transposed(src1), ggml_type_name(src1->type), src1->name); - - //TODO update for generic tensor parallelism - const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; - bool use_batched_cublas_f16 = src0->type == GGML_TYPE_F16 && (src1->type == GGML_TYPE_F16 || !any_gpus_with_slow_fp16); - bool use_batched_cublas_bf16 = src0->type == GGML_TYPE_BF16 && bf16_mma_hardware_available(cc); - bool use_batched_cublas_f32 = src0->type == GGML_TYPE_F32; - - if (!split && use_mul_mat_vec_f) { - // the custom F16 vector kernel can be used over batched cuBLAS GEMM - // but this is only faster for GPUs without tensor cores or with a thin src0 matrix (particularly KQV in attention) - ggml_cuda_mul_mat_vec_f(ctx, src0, src1, nullptr, dst); - } else if (!split && use_mul_mat_f) { - ggml_cuda_mul_mat_f(ctx, src0, src1, nullptr, dst); - } else if (!split && use_mul_mat_vec_q) { - ggml_cuda_mul_mat_vec_q(ctx, src0, src1, nullptr, dst); - } else if (!split && use_mul_mat_q) { - ggml_cuda_mul_mat_q(ctx, src0, src1, nullptr, dst); - } else if (!split && (use_batched_cublas_f16 || use_batched_cublas_bf16 || use_batched_cublas_f32) - && !ggml_is_transposed(src0) && !ggml_is_transposed(src1) && src1->ne[2]*src1->ne[3] > 1) { - // general KQ + KQV multi-batch without FlashAttention - ggml_cuda_mul_mat_batched_cublas(ctx, src0, src1, dst); - } else if (use_mul_mat_vec_f) { - ggml_cuda_op_mul_mat(ctx, src0, src1, dst, ggml_cuda_op_mul_mat_vec_f, nullptr); - } else if (use_mul_mat_vec_q) { - ggml_cuda_op_mul_mat(ctx, src0, src1, dst, ggml_cuda_op_mul_mat_vec_q, quantize_row_q8_1_cuda); - } else if (use_mul_mat_q) { - ggml_cuda_op_mul_mat(ctx, src0, src1, dst, ggml_cuda_op_mul_mat_q, quantize_mmq_q8_1_cuda); - } else { - ggml_cuda_op_mul_mat(ctx, src0, src1, dst, ggml_cuda_op_mul_mat_cublas, nullptr); - } -} - -static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { - const ggml_tensor * src0 = dst->src[0]; - const ggml_tensor * src1 = dst->src[1]; - const ggml_tensor * ids = dst->src[2]; - - GGML_ASSERT(src1->type == GGML_TYPE_F32); - GGML_ASSERT(dst->type == GGML_TYPE_F32); - GGML_ASSERT(!ggml_backend_buft_is_cuda_split(src0->buffer->buft) && "mul_mat_id does not support split buffers"); - - GGML_TENSOR_BINARY_OP_LOCALS - - const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; - - // [TAG_MUL_MAT_ID_CUDA_GRAPHS] - if (src1->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32) { - static_assert(MMVQ_MAX_BATCH_SIZE == MMVF_MAX_BATCH_SIZE); - if (ne2 <= MMVQ_MAX_BATCH_SIZE) { - if (ggml_is_quantized(src0->type)) { - const int mmvq_mmid_max = get_mmvq_mmid_max_batch(src0->type, cc); - if (ne2 <= mmvq_mmid_max) { - ggml_cuda_mul_mat_vec_q(ctx, src0, src1, ids, dst); - return; - } - } else { - if (GGML_CUDA_CC_IS_AMD(cc)) { - ggml_cuda_mul_mat_vec_f(ctx, src0, src1, ids, dst); - return; - } - } - } - - if (ggml_cuda_should_use_mmq(src0->type, cc, ne12, /*n_experts=*/ne02)) { - ggml_cuda_mul_mat_q(ctx, src0, src1, ids, dst); - return; - } - - if (ggml_cuda_should_use_mmf(src0->type, cc, WARP_SIZE, src0->ne, src0->nb, src1->ne[2], /*mul_mat_id=*/true)) { - ggml_cuda_mul_mat_f(ctx, src0, src1, ids, dst); - return; - } - } - - // note: this path should not be reached when recording CUDA graphs, because it requires stream synchronization - // TODO: add asserts to verify this. should work with CUDA, HIP, etc. - cudaStream_t stream = ctx.stream(); - - GGML_ASSERT(nb12 % nb11 == 0); - GGML_ASSERT(nb2 % nb1 == 0); - - const ggml_type type_src1_sorted = (src0->type == GGML_TYPE_F16 && !fast_fp16_hardware_available(cc)) - || ggml_is_quantized(src0->type) ? GGML_TYPE_F32 : src0->type; - const ggml_type type_dst_sorted = GGML_TYPE_F32; - const size_t ts_src1_sorted = ggml_type_size(type_src1_sorted); - const size_t ts_dst_sorted = ggml_type_size(type_dst_sorted); - - const int64_t n_expert_used = ids->ne[0]; - const int64_t ne_get_rows = ne12 * n_expert_used; - - std::vector ids_to_sorted_host; - ids_to_sorted_host.reserve(2*ne_get_rows); - std::vector ids_from_sorted_host(ne_get_rows); - - ggml_cuda_pool_alloc ids_buf_dev(ctx.pool(), 2*ne_get_rows); - - std::vector tokens_per_expert(ne02); - - ggml_cuda_pool_alloc src1_sorted(ctx.pool(), ne12*n_expert_used*ne10*ts_src1_sorted); - ggml_cuda_pool_alloc dst_sorted(ctx.pool(), ne2 *n_expert_used* ne0*ts_dst_sorted); - - std::vector ids_host(ggml_nbytes(ids)); - CUDA_CHECK(cudaMemcpyAsync(ids_host.data(), ids->data, ggml_nbytes(ids), cudaMemcpyDeviceToHost, stream)); - CUDA_CHECK(cudaStreamSynchronize(stream)); - - for (int64_t i02 = 0; i02 < ne02; ++i02) { // expert matrices - for (int64_t i12 = 0; i12 < ne12; ++i12) { // tokens - for (int64_t iex = 0; iex < n_expert_used; ++iex) { - const int32_t expert_to_use = *(const int32_t *)(ids_host.data() + i12*ids->nb[1] + iex*ids->nb[0]); - assert(expert_to_use >= 0 && expert_to_use < ne02); - if (expert_to_use == i02) { - ids_from_sorted_host[i12*n_expert_used + iex] = ids_to_sorted_host.size(); - ids_to_sorted_host.push_back(i12*ne11 + iex % ne11); - tokens_per_expert[i02]++; - break; - } - } - } - } - GGML_ASSERT(ids_to_sorted_host.size() == size_t(ne_get_rows)); - - ids_to_sorted_host.insert(ids_to_sorted_host.end(), ids_from_sorted_host.begin(), ids_from_sorted_host.end()); - - CUDA_CHECK(cudaMemcpyAsync(ids_buf_dev.ptr, ids_to_sorted_host.data(), 2*ne_get_rows*sizeof(int32_t), cudaMemcpyHostToDevice, stream)); - CUDA_CHECK(cudaStreamSynchronize(stream)); - - const int32_t * ids_to_sorted = ids_buf_dev.ptr + 0*ne_get_rows; - const int32_t * ids_from_sorted = ids_buf_dev.ptr + 1*ne_get_rows; - - get_rows_cuda(src1->data, src1->type, ids_to_sorted, src1_sorted.ptr, type_src1_sorted, - ne10, nb11, nb12, nb13, - ne_get_rows, 1, 1, sizeof(int32_t), ne_get_rows*sizeof(int32_t), ne_get_rows*sizeof(int32_t), - ne10*ts_src1_sorted, ne_get_rows*ne10*ts_src1_sorted, ne_get_rows*ne10*ts_src1_sorted, stream); - CUDA_CHECK(cudaGetLastError()); - - char * src1_data_cur = (char *) src1_sorted.ptr; - char * dst_data_cur = (char *) dst_sorted.ptr; - for (int64_t i02 = 0; i02 < ne02; ++i02) { - if (tokens_per_expert[i02] == 0) { - continue; - } - - ggml_tensor src0_slice = *src0; - src0_slice.ne[2] = 1; - src0_slice.nb[3] = src0_slice.nb[2]; - src0_slice.op = GGML_OP_VIEW; - src0_slice.view_src = dst->src[0]; // non-const pointer to src0 - src0_slice.data = (char *) src0->data + i02*nb02; - - ggml_tensor src1_slice; - memset(&src1_slice, 0, sizeof(src1_slice)); - src1_slice.buffer = src1->buffer; - src1_slice.type = type_src1_sorted; - src1_slice.ne[0] = ne10; - src1_slice.ne[1] = tokens_per_expert[i02]; - src1_slice.ne[2] = 1; - src1_slice.ne[3] = 1; - src1_slice.nb[0] = ts_src1_sorted; - src1_slice.nb[1] = src1_slice.ne[0] * src1_slice.nb[0]; - src1_slice.nb[2] = src1_slice.ne[1] * src1_slice.nb[1]; - src1_slice.nb[3] = src1_slice.ne[2] * src1_slice.nb[2]; - src1_slice.data = src1_data_cur; - - ggml_tensor dst_slice; - memset(&dst_slice, 0, sizeof(dst_slice)); - dst_slice.buffer = dst->buffer; - dst_slice.type = type_dst_sorted; - dst_slice.ne[0] = ne0; - dst_slice.ne[1] = tokens_per_expert[i02]; - dst_slice.ne[2] = 1; - dst_slice.ne[3] = 1; - dst_slice.nb[0] = ts_dst_sorted; - dst_slice.nb[1] = dst_slice.ne[0] * dst_slice.nb[0]; - dst_slice.nb[2] = dst_slice.ne[1] * dst_slice.nb[1]; - dst_slice.nb[3] = dst_slice.ne[2] * dst_slice.nb[2]; - dst_slice.data = dst_data_cur; - - ggml_cuda_mul_mat(ctx, &src0_slice, &src1_slice, &dst_slice); - CUDA_CHECK(cudaGetLastError()); - - src1_data_cur += src1_slice.nb[2]; - dst_data_cur += dst_slice.nb[2]; - } - - get_rows_cuda(dst_sorted.ptr, type_dst_sorted, ids_from_sorted, dst->data, dst->type, - ne0, ne0*ts_dst_sorted, ne_get_rows*ne0*ts_dst_sorted, ne_get_rows*ne0*ts_dst_sorted, - ne_get_rows, 1, 1, sizeof(int32_t), ne_get_rows*sizeof(int32_t), ne_get_rows*sizeof(int32_t), - nb1, nb2, nb3, stream); -} - -static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct ggml_tensor * dst) { - switch (dst->op) { - case GGML_OP_ARGMAX: - ggml_cuda_argmax(ctx, dst); - break; - case GGML_OP_COUNT_EQUAL: - ggml_cuda_count_equal(ctx, dst); - break; - case GGML_OP_REPEAT: - ggml_cuda_op_repeat(ctx, dst); - break; - case GGML_OP_REPEAT_BACK: - ggml_cuda_op_repeat_back(ctx, dst); - break; - case GGML_OP_GET_ROWS: - ggml_cuda_op_get_rows(ctx, dst); - break; - case GGML_OP_GET_ROWS_BACK: - ggml_cuda_op_get_rows_back(ctx, dst); - break; - case GGML_OP_SET_ROWS: - ggml_cuda_op_set_rows(ctx, dst); - break; - case GGML_OP_SET: - ggml_cuda_op_set(ctx, dst); - break; - case GGML_OP_DUP: - ggml_cuda_dup(ctx, dst); - break; - case GGML_OP_CPY: - ggml_cuda_cpy(ctx, dst->src[0], dst->src[1]); - break; - case GGML_OP_CONT: - ggml_cuda_dup(ctx, dst); - break; - case GGML_OP_ADD: - case GGML_OP_ADD1: // TODO: more efficient implementation - ggml_cuda_op_add(ctx, dst); - break; - case GGML_OP_ADD_ID: - ggml_cuda_op_add_id(ctx, dst); - break; - case GGML_OP_SUB: - ggml_cuda_op_sub(ctx, dst); - break; - case GGML_OP_ACC: - ggml_cuda_op_acc(ctx, dst); - break; - case GGML_OP_MUL: - ggml_cuda_op_mul(ctx, dst); - break; - case GGML_OP_DIV: - ggml_cuda_op_div(ctx, dst); - break; - case GGML_OP_UNARY: - switch (ggml_get_unary_op(dst)) { - case GGML_UNARY_OP_ABS: - ggml_cuda_op_abs(ctx, dst); - break; - case GGML_UNARY_OP_SGN: - ggml_cuda_op_sgn(ctx, dst); - break; - case GGML_UNARY_OP_NEG: - ggml_cuda_op_neg(ctx, dst); - break; - case GGML_UNARY_OP_STEP: - ggml_cuda_op_step(ctx, dst); - break; - case GGML_UNARY_OP_GELU: - ggml_cuda_op_gelu(ctx, dst); - break; - case GGML_UNARY_OP_SILU: - ggml_cuda_op_silu(ctx, dst); - break; - case GGML_UNARY_OP_GELU_ERF: - ggml_cuda_op_gelu_erf(ctx, dst); - break; - case GGML_UNARY_OP_GELU_QUICK: - ggml_cuda_op_gelu_quick(ctx, dst); - break; - case GGML_UNARY_OP_TANH: - ggml_cuda_op_tanh(ctx, dst); - break; - case GGML_UNARY_OP_RELU: - ggml_cuda_op_relu(ctx, dst); - break; - case GGML_UNARY_OP_SIGMOID: - ggml_cuda_op_sigmoid(ctx, dst); - break; - case GGML_UNARY_OP_HARDSIGMOID: - ggml_cuda_op_hardsigmoid(ctx, dst); - break; - case GGML_UNARY_OP_HARDSWISH: - ggml_cuda_op_hardswish(ctx, dst); - break; - case GGML_UNARY_OP_EXP: - ggml_cuda_op_exp(ctx, dst); - break; - case GGML_UNARY_OP_ELU: - ggml_cuda_op_elu(ctx, dst); - break; - case GGML_UNARY_OP_XIELU: - ggml_cuda_op_xielu(ctx, dst); - break; - case GGML_UNARY_OP_FLOOR: - ggml_cuda_op_floor(ctx, dst); - break; - case GGML_UNARY_OP_CEIL: - ggml_cuda_op_ceil(ctx, dst); - break; - case GGML_UNARY_OP_ROUND: - ggml_cuda_op_round(ctx, dst); - break; - case GGML_UNARY_OP_TRUNC: - ggml_cuda_op_trunc(ctx, dst); - break; - case GGML_UNARY_OP_EXPM1: - ggml_cuda_op_expm1(ctx, dst); - break; - case GGML_UNARY_OP_SOFTPLUS: - ggml_cuda_op_softplus(ctx, dst); - break; - default: - return false; - } - break; - case GGML_OP_GLU: - switch (ggml_get_glu_op(dst)) { - case GGML_GLU_OP_REGLU: - ggml_cuda_op_reglu(ctx, dst); - break; - case GGML_GLU_OP_GEGLU: - ggml_cuda_op_geglu(ctx, dst); - break; - case GGML_GLU_OP_SWIGLU: - ggml_cuda_op_swiglu(ctx, dst); - break; - case GGML_GLU_OP_SWIGLU_OAI: - ggml_cuda_op_swiglu_oai(ctx, dst); - break; - case GGML_GLU_OP_SWIGLU_DS4: - ggml_cuda_op_swiglu_ds4(ctx, dst); - break; - case GGML_GLU_OP_GEGLU_ERF: - ggml_cuda_op_geglu_erf(ctx, dst); - break; - case GGML_GLU_OP_GEGLU_QUICK: - ggml_cuda_op_geglu_quick(ctx, dst); - break; - default: - return false; - } - break; - case GGML_OP_NORM: - ggml_cuda_op_norm(ctx, dst); - break; - case GGML_OP_TURBO_WHT: - ggml_cuda_op_turbo_wht(ctx, dst); - break; - case GGML_OP_MOE_FUSED: - ggml_cuda_op_moe_fused(ctx, dst); - break; - case GGML_OP_GROUP_NORM: - ggml_cuda_op_group_norm(ctx, dst); - break; - case GGML_OP_L2_NORM: - ggml_cuda_op_l2_norm(ctx, dst); - break; - case GGML_OP_CONCAT: - ggml_cuda_op_concat(ctx, dst); - break; - case GGML_OP_UPSCALE: - ggml_cuda_op_upscale(ctx, dst); - break; - case GGML_OP_PAD: - ggml_cuda_op_pad(ctx, dst); - break; - case GGML_OP_PAD_REFLECT_1D: - ggml_cuda_op_pad_reflect_1d(ctx, dst); - break; - case GGML_OP_ARANGE: - ggml_cuda_op_arange(ctx, dst); - break; - case GGML_OP_TIMESTEP_EMBEDDING: - ggml_cuda_op_timestep_embedding(ctx, dst); - break; - case GGML_OP_LEAKY_RELU: - ggml_cuda_op_leaky_relu(ctx, dst); - break; - case GGML_OP_SILU_BACK: - ggml_cuda_op_silu_back(ctx, dst); - break; - case GGML_OP_RMS_NORM: - ggml_cuda_op_rms_norm(ctx, dst); - break; - case GGML_OP_RMS_NORM_BACK: - ggml_cuda_op_rms_norm_back(ctx, dst); - break; - case GGML_OP_MUL_MAT: - ggml_cuda_mul_mat(ctx, dst->src[0], dst->src[1], dst); - break; - case GGML_OP_MUL_MAT_ID: - ggml_cuda_mul_mat_id(ctx, dst); - break; - case GGML_OP_OUT_PROD: - ggml_cuda_out_prod(ctx, dst); - break; - case GGML_OP_SCALE: - ggml_cuda_op_scale(ctx, dst); - break; - case GGML_OP_SQR: - ggml_cuda_op_sqr(ctx, dst); - break; - case GGML_OP_SQRT: - ggml_cuda_op_sqrt(ctx, dst); - break; - case GGML_OP_SIN: - ggml_cuda_op_sin(ctx, dst); - break; - case GGML_OP_COS: - ggml_cuda_op_cos(ctx, dst); - break; - case GGML_OP_CLAMP: - ggml_cuda_op_clamp(ctx, dst); - break; - case GGML_OP_LOG: - ggml_cuda_op_log(ctx, dst); - break; - case GGML_OP_NONE: - case GGML_OP_RESHAPE: - case GGML_OP_VIEW: - case GGML_OP_PERMUTE: - case GGML_OP_TRANSPOSE: - break; - case GGML_OP_DIAG: - ggml_cuda_op_diag(ctx, dst); - break; - case GGML_OP_DIAG_MASK_INF: - ggml_cuda_op_diag_mask_inf(ctx, dst); - break; - case GGML_OP_SOFT_MAX: - ggml_cuda_op_soft_max(ctx, dst); - break; - case GGML_OP_SOFT_MAX_BACK: - ggml_cuda_op_soft_max_back(ctx, dst); - break; - case GGML_OP_ROPE: - ggml_cuda_op_rope(ctx, dst); - break; - case GGML_OP_ROPE_BACK: - ggml_cuda_op_rope_back(ctx, dst); - break; - case GGML_OP_ROLL: - ggml_cuda_op_roll(ctx, dst); - break; - case GGML_OP_IM2COL: - ggml_cuda_op_im2col(ctx, dst); - break; - case GGML_OP_IM2COL_3D: - ggml_cuda_op_im2col_3d(ctx, dst); - break; - case GGML_OP_CONV_2D: - ggml_cuda_op_conv2d(ctx, dst); - break; - case GGML_OP_CONV_2D_DW: - ggml_cuda_op_conv2d_dw(ctx, dst); - break; - case GGML_OP_CONV_TRANSPOSE_2D: - ggml_cuda_conv_2d_transpose_p0(ctx, dst); - break; - case GGML_OP_CONV_TRANSPOSE_1D: - ggml_cuda_op_conv_transpose_1d(ctx,dst); - break; - case GGML_OP_POOL_2D: - ggml_cuda_op_pool2d(ctx, dst); - break; - case GGML_OP_SUM: - ggml_cuda_op_sum(ctx, dst); - break; - case GGML_OP_CUMSUM: - ggml_cuda_op_cumsum(ctx, dst); - break; - case GGML_OP_SUM_ROWS: - ggml_cuda_op_sum_rows(ctx, dst); - break; - case GGML_OP_MEAN: - ggml_cuda_op_mean(ctx, dst); - break; - case GGML_OP_SSM_CONV: - ggml_cuda_op_ssm_conv(ctx, dst); - break; - case GGML_OP_SSM_SCAN: - ggml_cuda_op_ssm_scan(ctx, dst); - break; - case GGML_OP_TOP_K: - ggml_cuda_op_top_k(ctx, dst); - break; - case GGML_OP_ARGSORT: - ggml_cuda_op_argsort(ctx, dst); - break; - case GGML_OP_FLASH_ATTN_EXT: - ggml_cuda_flash_attn_ext(ctx, dst); - break; - case GGML_OP_FLASH_ATTN_SPARSE: - ggml_cuda_flash_attn_sparse(ctx, dst); - break; - case GGML_OP_CROSS_ENTROPY_LOSS: - ggml_cuda_cross_entropy_loss(ctx, dst); - break; - case GGML_OP_TRI: - ggml_cuda_op_tri(ctx, dst); - break; - case GGML_OP_RWKV_WKV6: - ggml_cuda_op_rwkv_wkv6(ctx, dst); - break; - case GGML_OP_GATED_LINEAR_ATTN: - ggml_cuda_op_gated_linear_attn(ctx, dst); - break; - case GGML_OP_GATED_DELTA_NET: - ggml_cuda_op_gated_delta_net(ctx, dst); - break; - case GGML_OP_RWKV_WKV7: - ggml_cuda_op_rwkv_wkv7(ctx, dst); - break; - case GGML_OP_CROSS_ENTROPY_LOSS_BACK: - ggml_cuda_cross_entropy_loss_back(ctx, dst); - break; - case GGML_OP_OPT_STEP_ADAMW: - ggml_cuda_opt_step_adamw(ctx, dst); - break; - case GGML_OP_OPT_STEP_SGD: - ggml_cuda_opt_step_sgd(ctx, dst); - break; - case GGML_OP_SOLVE_TRI: - ggml_cuda_op_solve_tri(ctx, dst); - break; - case GGML_OP_FILL: - ggml_cuda_op_fill(ctx, dst); - break; - default: - return false; - } - - cudaError_t err = cudaGetLastError(); - if (err != cudaSuccess) { - GGML_LOG_ERROR("%s: %s failed: node='%s' [%ld,%ld] src0='%s' t=%d [%ld,%ld] src1='%s' [%ld,%ld]\n", - __func__, ggml_op_desc(dst), dst->name, (long)dst->ne[0], (long)dst->ne[1], - dst->src[0] ? dst->src[0]->name : "-", dst->src[0] ? (int)dst->src[0]->type : -1, - dst->src[0] ? (long)dst->src[0]->ne[0] : 0, dst->src[0] ? (long)dst->src[0]->ne[1] : 0, - dst->src[1] ? dst->src[1]->name : "-", - dst->src[1] ? (long)dst->src[1]->ne[0] : 0, dst->src[1] ? (long)dst->src[1]->ne[1] : 0); - CUDA_CHECK(err); - } - - return true; -} - -//////////////////////////////////////////////////////////////////////////////// - -// backend - -static const char * ggml_backend_cuda_get_name(ggml_backend_t backend) { - ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context; - - return cuda_ctx->name.c_str(); -} - -static void ggml_backend_cuda_free(ggml_backend_t backend) { - ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context; - - delete cuda_ctx; - delete backend; -} - -static void ggml_backend_cuda_set_tensor_async(ggml_backend_t backend, ggml_tensor * tensor, const void * data, size_t offset, size_t size) { - ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context; - ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; - - GGML_ASSERT(buf->buft == ggml_backend_cuda_buffer_type(cuda_ctx->device) && "unsupported buffer type"); - - CUDA_CHECK(cudaMemcpyAsync((char *) tensor->data + offset, data, size, cudaMemcpyHostToDevice, cuda_ctx->stream())); -} - -static void ggml_backend_cuda_get_tensor_async(ggml_backend_t backend, const ggml_tensor * tensor, void * data, size_t offset, size_t size) { - ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context; - ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; - - GGML_ASSERT(buf->buft == ggml_backend_cuda_buffer_type(cuda_ctx->device) && "unsupported buffer type"); - - CUDA_CHECK(cudaMemcpyAsync(data, (const char *) tensor->data + offset, size, cudaMemcpyDeviceToHost, cuda_ctx->stream())); -} - -static void ggml_backend_cuda_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) { - ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context; - ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; - - GGML_ASSERT(buf->buft == ggml_backend_cuda_buffer_type(cuda_ctx->device) && "unsupported buffer type"); - - CUDA_CHECK(cudaMemcpy2DAsync( - (char *) tensor->data + offset, stride_tensor, data, stride_data, size, n_copies, cudaMemcpyHostToDevice, cuda_ctx->stream())); -} - -static void ggml_backend_cuda_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) { - ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context; - ggml_backend_buffer_t buf = tensor->view_src ? tensor->view_src->buffer : tensor->buffer; - - GGML_ASSERT(buf->buft == ggml_backend_cuda_buffer_type(cuda_ctx->device) && "unsupported buffer type"); - - CUDA_CHECK(cudaMemcpy2DAsync( - data, stride_data, (const char *) tensor->data + offset, stride_tensor, size, n_copies, cudaMemcpyDeviceToHost, cuda_ctx->stream())); -} - -static bool ggml_backend_cuda_cpy_tensor_async(ggml_backend_t backend_src, ggml_backend_t backend_dst, const ggml_tensor * src, ggml_tensor * dst) { - ggml_backend_buffer_t buf_src = src->view_src ? src->view_src->buffer : src->buffer; - ggml_backend_buffer_t buf_dst = dst->view_src ? dst->view_src->buffer : dst->buffer; - - if (!ggml_backend_is_cuda(backend_src) || !ggml_backend_is_cuda(backend_dst)) { - return false; - } - - if (!ggml_backend_buffer_is_cuda(buf_src) || !ggml_backend_buffer_is_cuda(buf_dst)) { - return false; - } - - // device -> device copy - ggml_backend_cuda_context * cuda_ctx_src = (ggml_backend_cuda_context *) backend_src->context; - ggml_backend_cuda_context * cuda_ctx_dst = (ggml_backend_cuda_context *) backend_dst->context; - - ggml_backend_cuda_buffer_context * buf_ctx_src = (ggml_backend_cuda_buffer_context *) buf_src->context; - ggml_backend_cuda_buffer_context * buf_ctx_dst = (ggml_backend_cuda_buffer_context *) buf_dst->context; - - if (cuda_ctx_src->device != buf_ctx_src->device || cuda_ctx_dst->device != buf_ctx_dst->device) { -#ifndef NDEBUG - GGML_LOG_DEBUG("%s: backend and buffer devices do not match\n", __func__); -#endif // NDEBUG - return false; - } - - if (backend_src != backend_dst) { - // copy on src stream - if (cuda_ctx_src->device == cuda_ctx_dst->device) { - CUDA_CHECK(cudaMemcpyAsync(dst->data, src->data, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, cuda_ctx_src->stream())); - } else { -#ifdef GGML_CUDA_NO_PEER_COPY - return false; -#else - CUDA_CHECK(cudaMemcpyPeerAsync(dst->data, cuda_ctx_dst->device, src->data, cuda_ctx_src->device, ggml_nbytes(dst), cuda_ctx_src->stream())); -#endif // GGML_CUDA_NO_PEER_COPY - } - - // record event on src stream after the copy - if (!cuda_ctx_src->copy_event) { - ggml_cuda_set_device(cuda_ctx_src->device); - CUDA_CHECK(cudaEventCreateWithFlags(&cuda_ctx_src->copy_event, cudaEventDisableTiming)); - } - - CUDA_CHECK(cudaEventRecord(cuda_ctx_src->copy_event, cuda_ctx_src->stream())); - - // wait on dst stream for the copy to complete - CUDA_CHECK(cudaStreamWaitEvent(cuda_ctx_dst->stream(), cuda_ctx_src->copy_event, 0)); - } else { - // src and dst are on the same backend - CUDA_CHECK(cudaMemcpyAsync(dst->data, src->data, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, cuda_ctx_src->stream())); - } - return true; -} - -static void ggml_backend_cuda_synchronize(ggml_backend_t backend) { - ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context; - - CUDA_CHECK(cudaStreamSynchronize(cuda_ctx->stream())); - - GGML_UNUSED(backend); -} - -#ifdef USE_CUDA_GRAPH -static bool ggml_cuda_graph_check_compability(ggml_cgraph * cgraph) { - - bool use_cuda_graph = true; - // Loop over nodes in GGML graph to obtain info needed for CUDA graph - - for (int i = 0; i < cgraph->n_nodes; i++) { - ggml_tensor * node = cgraph->nodes[i]; - - if (ggml_is_empty(node) || node->op == GGML_OP_RESHAPE || node->op == GGML_OP_TRANSPOSE || node->op == GGML_OP_VIEW || node->op == GGML_OP_PERMUTE || node->op == GGML_OP_NONE) { - continue; - } - - if (node->src[0] && node->src[0]->buffer && ggml_backend_buft_is_cuda_split(node->src[0]->buffer->buft)) { - use_cuda_graph = false; // Split buffers are not supported by CUDA graph capture -#ifndef NDEBUG - GGML_LOG_DEBUG("%s: disabling CUDA graphs due to split buffer\n", __func__); -#endif - } - - // [TAG_MUL_MAT_ID_CUDA_GRAPHS] - if (node->op == GGML_OP_MUL_MAT_ID) { - const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; - const int mmvq_mmid_max = get_mmvq_mmid_max_batch(node->src[0]->type, cc); - if (!ggml_is_quantized(node->src[0]->type) || node->ne[2] > mmvq_mmid_max) { - // under these conditions, the mul_mat_id operation will need to synchronize the stream, so we cannot use CUDA graphs - // TODO: figure out a way to enable for larger batch sizes, without hurting performance - // ref: https://github.com/ggml-org/llama.cpp/pull/18958 - use_cuda_graph = false; -#ifndef NDEBUG - GGML_LOG_DEBUG("%s: disabling CUDA graphs due to unsupported node type\n", __func__); -#endif - } - } - - if (!use_cuda_graph) { - break; - } - } - - return use_cuda_graph; -} - -static const void * ggml_cuda_graph_get_key(ggml_cgraph * cgraph) { - return cgraph->nodes[0]; -} - -static bool ggml_cuda_graph_update_required(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph) { - bool res = false; - - const void * graph_key = ggml_cuda_graph_get_key(cgraph); - ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); - - // Check if the graph size has changed - if ((int)graph->node_props.size() != cgraph->n_nodes) { - res = true; - graph->node_props.resize(cgraph->n_nodes); - } - - for (int i = 0; i < cgraph->n_nodes; i++) { - ggml_cuda_graph::node_properties prop = {}; - memcpy(&prop.node, cgraph->nodes[i], sizeof(ggml_tensor)); - - for (int j = 0; j < GGML_MAX_SRC; ++j) { - if (cgraph->nodes[i]->src[j]) { - prop.node_src_data_ptrs[j] = cgraph->nodes[i]->src[j]->data; - memcpy(prop.node_src_ne[j], cgraph->nodes[i]->src[j]->ne, sizeof(prop.node_src_ne[j])); - memcpy(prop.node_src_nb[j], cgraph->nodes[i]->src[j]->nb, sizeof(prop.node_src_nb[j])); - } - } - - if (res || memcmp(&graph->node_props[i], &prop, sizeof(prop)) != 0) { - graph->node_props[i] = prop; - res = true; - } - } - - return res; -} - -static void ggml_cuda_graph_update_executable(ggml_backend_cuda_context * cuda_ctx, const void * graph_key) { - ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); - -#if CUDART_VERSION >= 12000 - cudaGraphExecUpdateResultInfo result_info; - cudaError_t stat = cudaGraphExecUpdate(graph->instance, graph->graph, &result_info); -#else - cudaGraphNode_t errorNode; - cudaGraphExecUpdateResult result_info; - cudaError_t stat = cudaGraphExecUpdate(graph->instance, graph->graph, &errorNode, &result_info); -#endif // CUDART_VERSION >= 12000 - - if (stat == cudaErrorGraphExecUpdateFailure) { -#ifndef NDEBUG - GGML_LOG_DEBUG("%s: CUDA graph update failed\n", __func__); -#endif - - // The pre-existing graph exec cannot be updated due to violated constraints - // so instead clear error and re-instantiate - (void)cudaGetLastError(); - CUDA_CHECK(cudaGraphExecDestroy(graph->instance)); - graph->instance = nullptr; - CUDA_CHECK(cudaGraphInstantiate(&graph->instance, graph->graph, NULL, NULL, 0)); - } else { - GGML_ASSERT(stat == cudaSuccess); - } -} -#endif // USE_CUDA_GRAPH - -static bool ggml_cuda_should_fuse_rope_set_rows(const ggml_tensor * rope, - const ggml_tensor * view, - const ggml_tensor * set_rows) { - - if (rope->op != GGML_OP_ROPE || view->op != GGML_OP_VIEW || set_rows->op != GGML_OP_SET_ROWS) { - return false; - } - // ne3 not tested - if (rope->src[0]->ne[3] != 1) { - return false; - } - - if (set_rows->type != GGML_TYPE_F32 && set_rows->type != GGML_TYPE_F16) { - return false; - } - - if (set_rows->src[1]->type != GGML_TYPE_I64) { - return false; - } - - // The view should flatten two dims of rope into one dim - if (!ggml_is_contiguous(view) || view->ne[0] != rope->ne[0] * rope->ne[1]) { - return false; - } - - // Only norm/neox shaders have the fusion code - const int mode = ((const int32_t *) rope->op_params)[2]; - if (mode != GGML_ROPE_TYPE_NORMAL && mode != GGML_ROPE_TYPE_NEOX) { - return false; - } - - return true; -} - -static bool ggml_cuda_topk_moe_fusion(const struct ggml_cgraph * cgraph, int node_idx, ggml_cuda_topk_moe_args & args) { - args.sigmoid = false; - args.softmax = false; - args.delayed_softmax = false; - args.prob_bias = false; - args.norm = false; - - const int n_nodes = cgraph->n_nodes; - ggml_tensor ** nodes = cgraph->nodes; - - if (nodes[node_idx]->op == GGML_OP_SOFT_MAX) { - args.softmax = true; - } - - if (nodes[node_idx]->op == GGML_OP_UNARY) { - if (ggml_get_unary_op(nodes[node_idx]) != GGML_UNARY_OP_SIGMOID) { - return false; - } - args.sigmoid = true; - } - - if (nodes[node_idx]->op == GGML_OP_ARGSORT) { - args.delayed_softmax = true; - } - - node_idx++; - - if (args.sigmoid || args.softmax) { - // SOFTMAX -> RESHAPE - if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_RESHAPE || - nodes[node_idx]->src[0] != nodes[node_idx - 1]) { - return false; - } - ggml_tensor * probs_reshaped = nodes[node_idx]; - node_idx++; - - if (node_idx >= n_nodes) { - return false; - } - - // src of bias add is the unreshaped probs (-2 instead of -1) - if (nodes[node_idx]->op == GGML_OP_ADD && nodes[node_idx]->src[0] == nodes[node_idx - 2]) { - args.prob_bias = true; - node_idx++; - } - // RESHAPE/ADD -> ARGSORT - if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_ARGSORT) { - return false; - } - - if (args.prob_bias && nodes[node_idx]->src[0] != nodes[node_idx - 1]) { - return false; - } else if (!args.prob_bias && nodes[node_idx]->src[0] != nodes[node_idx - 2]) { - return false; - } - - node_idx++; - - // ARGSORT-> VIEW - if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_VIEW || - nodes[node_idx]->src[0] != nodes[node_idx - 1]) { - return false; - } - node_idx++; - - if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_GET_ROWS) { - return false; - } - - // GET_ROWS - if (nodes[node_idx]->src[0] != probs_reshaped || nodes[node_idx]->src[1] != nodes[node_idx - 1]) { - return false; - } - node_idx++; - } else if (args.delayed_softmax) { - if (node_idx - 2 < 0) { - return false; - } - ggml_tensor * probs_reshaped = nodes[node_idx - 2]; - - // VIEW->ARGSORT - if (node_idx >= n_nodes || nodes[node_idx]->op != GGML_OP_VIEW || - nodes[node_idx]->src[0] != nodes[node_idx - 1]) { - return false; - } - node_idx++; - - // GET_ROWS - if (node_idx >= n_nodes || nodes[node_idx]->src[1] != nodes[node_idx - 1] || - nodes[node_idx]->src[0] != probs_reshaped) { - return false; - } - node_idx++; - - static const std::vector remaining_ops = { GGML_OP_RESHAPE, GGML_OP_SOFT_MAX, GGML_OP_RESHAPE }; - - for (const ggml_op op : remaining_ops) { - if (node_idx >= n_nodes || nodes[node_idx]->op != op || nodes[node_idx]->src[0] != nodes[node_idx - 1]) { - return false; - } - node_idx++; - } - } - - // At this point we can check for norm + scale. Everything is now at least valid till the norm - if (node_idx >= n_nodes) { - return true; - } - - if (nodes[node_idx]->op == GGML_OP_RESHAPE) { - //check RESHAPE->SUM_ROWS->CLAMP->DIV->RESHAPE - static const std::vector norm_ops = { GGML_OP_RESHAPE, GGML_OP_SUM_ROWS, GGML_OP_CLAMP }; - - args.norm = true; - for (const ggml_op op : norm_ops) { - if (nodes[node_idx]->op == op && nodes[node_idx]->src[0] == nodes[node_idx - 1]) { - node_idx++; - } else { - args.norm = false; - return true; - } - } - - // DIV <- CLAMP, RESHAPE - if (nodes[node_idx]->op != GGML_OP_DIV || nodes[node_idx]->src[1] != nodes[node_idx - 1] || - nodes[node_idx]->src[0] != nodes[node_idx - 3]) { - args.norm = false; - return true; - } - node_idx++; - - if (nodes[node_idx]->op != GGML_OP_RESHAPE || nodes[node_idx]->src[0] != nodes[node_idx - 1]) { - args.norm = false; - return true; - } - - node_idx++; - } - - if (nodes[node_idx]->op == GGML_OP_SCALE && nodes[node_idx]->src[0] == nodes[node_idx - 1]) { - args.scale = true; - } - - return true; -} - -// returns whether the write (out) nodes overwrite the read nodes in operation -static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph, - const int node_idx, - const int node_count, - const int * out_nodes, - const int out_count, - const bool is_topk_moe = false) { - auto nodes_overlap = [&](const ggml_tensor * a, const ggml_tensor * b) { - const int64_t a_start = (int64_t) a->data; - const int64_t a_end = a_start + ggml_backend_buft_get_alloc_size(a->buffer->buft, a); - - const int64_t b_start = (int64_t) b->data; - const int64_t b_end = b_start + ggml_backend_buft_get_alloc_size(b->buffer->buft, b); - - if ((b_start <= a_start && a_start < b_end) || (a_start <= b_start && b_start < a_end)) { - return true; - } - - return false; - }; - - bool is_ok = true; - // exception for topk-moe, as each row is read entirely before writing - if (ggml_nrows(cgraph->nodes[node_idx]) == 1 && is_topk_moe) { - return true; - } - - for (int i = 0; i < out_count; ++i) { - const ggml_tensor * dst = cgraph->nodes[out_nodes[i]]; - - for (int j = node_idx; j < node_idx + node_count; ++j) { - // Loop over all srcs of all nodes in the fusion. If the src overlaps - // the destination and the src is not an intermediate node that's being - // elided, then disable fusion. - - for (int src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) { - const ggml_tensor * src = cgraph->nodes[j]->src[src_idx]; - - if (!src || src->op == GGML_OP_NONE) { - continue; - } - - if (nodes_overlap(dst, src)) { - bool found = false; - - for (int k = node_idx; k < j; ++k) { - if (cgraph->nodes[k] == src) { - found = true; - break; - } - } - - if (!found) { - is_ok = false; - break; - } - } - } - } - } - - return is_ok; -} - - -static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, - int node_idx, - std::initializer_list ops, - std::initializer_list unary_ops) { -#ifndef NDEBUG - const size_t num_unary = std::count(ops.begin(), ops.end(), GGML_OP_UNARY); - GGML_ASSERT(unary_ops.size() == num_unary); -#endif - - const auto is_equal = [](const std::initializer_list & list1, - const std::initializer_list & list2) { - return std::equal(list1.begin(), list1.end(), list2.begin(), list2.end()); - }; - - std::initializer_list mul_mat_bias_glu_ops = { GGML_OP_MUL_MAT, GGML_OP_ADD, GGML_OP_MUL_MAT, GGML_OP_ADD, GGML_OP_GLU }; - std::initializer_list mul_mat_id_bias_glu_ops = { GGML_OP_MUL_MAT_ID, GGML_OP_ADD_ID, GGML_OP_MUL_MAT_ID, GGML_OP_ADD_ID, GGML_OP_GLU }; - - std::initializer_list mul_mat_id_glu_ops = { GGML_OP_MUL_MAT_ID, GGML_OP_MUL_MAT_ID, GGML_OP_GLU }; - std::initializer_list mul_mat_glu_ops = { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT, GGML_OP_GLU }; - - if ((is_equal(mul_mat_bias_glu_ops, ops) || is_equal(mul_mat_id_bias_glu_ops, ops)) && - ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 4 })) { - const ggml_tensor * ffn_gate = cgraph->nodes[node_idx]; - const ggml_tensor * ffn_gate_bias = cgraph->nodes[node_idx + 1]; - const ggml_tensor * ffn_up = cgraph->nodes[node_idx + 2]; - const ggml_tensor * ffn_up_bias = cgraph->nodes[node_idx + 3]; - const ggml_tensor * glu = cgraph->nodes[node_idx + 4]; - - if (ggml_cuda_should_fuse_mul_mat(ffn_up, ffn_gate, glu, ffn_up_bias, ffn_gate_bias)) { - int out_nodes[] = { node_idx + 4 }; - return ggml_cuda_check_fusion_memory_ranges(cgraph, node_idx, (int)ops.size(), out_nodes, 1); - } - } - - if ((is_equal(mul_mat_id_glu_ops, ops) || is_equal(mul_mat_glu_ops, ops)) && - ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 2 })) { - const ggml_tensor * ffn_gate = cgraph->nodes[node_idx]; - const ggml_tensor * ffn_up = cgraph->nodes[node_idx + 1]; - const ggml_tensor * glu = cgraph->nodes[node_idx + 2]; - - if (ggml_cuda_should_fuse_mul_mat(ffn_up, ffn_gate, glu)) { - int out_nodes[] = { node_idx + 2 }; - return ggml_cuda_check_fusion_memory_ranges(cgraph, node_idx, (int)ops.size(), out_nodes, 1); - } - } - - std::initializer_list rope_set_rows_ops = { GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS }; - - if (is_equal(rope_set_rows_ops, ops) && ggml_can_fuse_subgraph(cgraph, node_idx, ops, { node_idx + 2 })) { - const ggml_tensor * rope = cgraph->nodes[node_idx]; - const ggml_tensor * view = cgraph->nodes[node_idx + 1]; - const ggml_tensor * set_rows = cgraph->nodes[node_idx + 2]; - - if (ggml_cuda_should_fuse_rope_set_rows(rope, view, set_rows)) { - return true; - } - } - - if (!ggml_can_fuse(cgraph, node_idx, ops)) { - return false; - } - - if ((ops.size() == 2 || ops.size() == 3) && ops.begin()[0] == GGML_OP_RMS_NORM && ops.begin()[1] == GGML_OP_MUL) { - const ggml_tensor *rms_norm = cgraph->nodes[node_idx]; - const ggml_tensor *mul = cgraph->nodes[node_idx+1]; - const ggml_tensor *add = nullptr; - - if (ops.size() == 3 && ops.begin()[2] == GGML_OP_ADD) { - add = cgraph->nodes[node_idx+2]; - } - - GGML_ASSERT(rms_norm->src[0]->type == GGML_TYPE_F32); - GGML_ASSERT(rms_norm->type == GGML_TYPE_F32); - - //rms norm only supports F32 - if (mul->src[0]->type != GGML_TYPE_F32 || - mul->src[1]->type != GGML_TYPE_F32 || - mul->type != GGML_TYPE_F32) { - return false; - } - - if (add && (add->src[0]->type != GGML_TYPE_F32 || - add->src[1]->type != GGML_TYPE_F32 || - add->type != GGML_TYPE_F32) ) { - return false; - } - - //if rms norm is the B operand, then we don't handle broadcast - if (rms_norm == mul->src[1] && !ggml_are_same_shape(mul->src[0], rms_norm)) { - return false; - } - - //rms_norm kernel assumes contiguous rows - if (!ggml_is_contiguous_rows(mul->src[0]) || !ggml_is_contiguous_rows(mul->src[1])) { - return false; - } - - if (add && (!ggml_is_contiguous(add->src[0]) || !ggml_is_contiguous_rows(add->src[1]))) { - return false; - } - - return true; - } - - if (ops.size() == 2 && ops.begin()[0] == GGML_OP_SSM_CONV && ops.begin()[1] == GGML_OP_UNARY - && unary_ops.size() == 1 && unary_ops.begin()[0] == GGML_UNARY_OP_SILU) { - const ggml_tensor * ssm_conv = cgraph->nodes[node_idx]; - const ggml_tensor * silu = cgraph->nodes[node_idx+1]; - - if (ssm_conv->type != GGML_TYPE_F32 || silu->type != GGML_TYPE_F32) { - return false; - } - - return true; - } - - if (ops.size() == 2 && ops.begin()[0] == GGML_OP_UNARY && ops.begin()[1] == GGML_OP_MUL - && unary_ops.size() == 1 && (unary_ops.begin()[0] == GGML_UNARY_OP_SILU || unary_ops.begin()[0] == GGML_UNARY_OP_SIGMOID || unary_ops.begin()[0] == GGML_UNARY_OP_SOFTPLUS)) { - const ggml_tensor * unary = cgraph->nodes[node_idx]; - const ggml_tensor * mul = cgraph->nodes[node_idx+1]; - - if (ggml_get_unary_op(unary) != unary_ops.begin()[0]) { - return false; - } - - if (unary->type != GGML_TYPE_F32 && unary->type != GGML_TYPE_F16) { - return false; - } - - if (unary->type != mul->type) { - return false; - } - - const ggml_tensor * other = (mul->src[0] == unary) ? mul->src[1] : mul->src[0]; - if (other->type != unary->type) { - return false; - } - if (!ggml_is_contiguous_1(other) || !ggml_is_contiguous_1(unary->src[0]) || !ggml_are_same_shape(other, unary)) { - return false; - } - - return true; - } - - if (ops.size() == 3 && ops.begin()[0] == GGML_OP_SCALE && ops.begin()[1] == GGML_OP_UNARY && ops.begin()[2] == GGML_OP_SCALE - && unary_ops.size() == 1 && unary_ops.begin()[0] == GGML_UNARY_OP_TANH) { - const ggml_tensor *scale = cgraph->nodes[node_idx]; - const ggml_tensor *tanh = cgraph->nodes[node_idx+1]; - const ggml_tensor *scale2 = cgraph->nodes[node_idx+2]; - - GGML_ASSERT(scale->src[0]->type == GGML_TYPE_F32); - GGML_ASSERT(scale->type == GGML_TYPE_F32); - - if (ggml_get_unary_op(tanh) != GGML_UNARY_OP_TANH) { - return false; - } - - // Check for bias - if (ggml_get_op_params_f32(scale, 1) != 0.0f || ggml_get_op_params_f32(scale2, 1) != 0.0f) { - return false; - } - - return true; - } - - return false; -} - -static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, const bool use_cuda_graph, const bool cuda_graph_update_required, const void * graph_key) { - bool graph_evaluated_or_captured = false; - - // flag used to determine whether it is an integrated_gpu - const bool integrated = ggml_cuda_info().devices[cuda_ctx->device].integrated; - - ggml_cuda_stream_context & stream_ctx = cuda_ctx->stream_context(); - bool is_concurrent_event_active = false; - ggml_cuda_concurrent_event * concurrent_event = nullptr; - bool should_launch_concurrent_events = false; - - const auto try_launch_concurrent_event = [&](const ggml_tensor * node) { - if (stream_ctx.concurrent_events.find(node) != stream_ctx.concurrent_events.end()) { - concurrent_event = &stream_ctx.concurrent_events[node]; - - is_concurrent_event_active = true; - - GGML_LOG_DEBUG("Launching %d streams at %s\n", concurrent_event->n_streams, node->name); - - cudaStream_t main_stream = cuda_ctx->stream(); // this should be stream 0 - GGML_ASSERT(cuda_ctx->curr_stream_no == 0); - CUDA_CHECK(cudaEventRecord(concurrent_event->fork_event, main_stream)); - - for (int i = 1; i <= concurrent_event->n_streams; ++i) { - cudaStream_t stream = cuda_ctx->stream(cuda_ctx->device, i); - CUDA_CHECK(cudaStreamWaitEvent(stream, concurrent_event->fork_event)); - } - } - }; - - while (!graph_evaluated_or_captured) { - // Only perform the graph execution if CUDA graphs are not enabled, or we are capturing the graph. - // With the use of CUDA graphs, the execution will be performed by the graph launch. - if (!use_cuda_graph || cuda_graph_update_required) { - [[maybe_unused]] int prev_i = 0; - - if (stream_ctx.concurrent_events.size() > 0) { - should_launch_concurrent_events = true; - for (const auto & [tensor, event] : stream_ctx.concurrent_events) { - should_launch_concurrent_events = should_launch_concurrent_events && event.is_valid(); - } - } - - if (should_launch_concurrent_events) { - // Restore original node order within each concurrent region to enable fusion within streams - - std::unordered_map node_to_idx; - node_to_idx.reserve(cgraph->n_nodes); - for (int i = 0; i < cgraph->n_nodes; ++i) { - node_to_idx[cgraph->nodes[i]] = i; - } - - for (auto & [fork_node, event] : stream_ctx.concurrent_events) { - // Find positions of all nodes from this event in the current graph - std::vector positions; - positions.reserve(event.original_order.size()); - - bool all_found = true; - for (const ggml_tensor * orig_node : event.original_order) { - auto it = node_to_idx.find(orig_node); - if (it != node_to_idx.end()) { - positions.push_back(it->second); - } else { - all_found = false; - break; - } - } - - if (!all_found || positions.size() != event.original_order.size()) { - continue; - } - - // Sort positions to get contiguous range - std::vector sorted_positions = positions; - std::sort(sorted_positions.begin(), sorted_positions.end()); - - bool is_contiguous = true; - for (size_t i = 1; i < sorted_positions.size(); ++i) { - if (sorted_positions[i] != sorted_positions[i-1] + 1) { - is_contiguous = false; - break; - } - } - - if (!is_contiguous) { - continue; - } - - // Restore original order at the sorted positions - int start_pos = sorted_positions[0]; - for (size_t i = 0; i < event.original_order.size(); ++i) { - cgraph->nodes[start_pos + i] = const_cast(event.original_order[i]); - } - } - } else { - stream_ctx.concurrent_events.clear(); - } - - for (int i = 0; i < cgraph->n_nodes; i++) { - ggml_tensor * node = cgraph->nodes[i]; - if (is_concurrent_event_active) { - GGML_ASSERT(concurrent_event); - - if (node == concurrent_event->join_node) { - cuda_ctx->curr_stream_no = 0; - for (int i = 1; i <= concurrent_event->n_streams; ++i) { - // Wait on join events of forked streams in the main stream - CUDA_CHECK(cudaEventRecord(concurrent_event->join_events[i - 1], - cuda_ctx->stream(cuda_ctx->device, i))); - CUDA_CHECK(cudaStreamWaitEvent(cuda_ctx->stream(), concurrent_event->join_events[i - 1])); - } - - is_concurrent_event_active = false; - concurrent_event = nullptr; - } else { - GGML_ASSERT (concurrent_event->stream_mapping.find(node) != concurrent_event->stream_mapping.end()); - cuda_ctx->curr_stream_no = concurrent_event->stream_mapping[node]; - GGML_LOG_DEBUG("Setting stream no to %d for node %s\n", cuda_ctx->curr_stream_no, node->name); - } - } else if (i - prev_i > 1) { - //the previous node was fused - const ggml_tensor * prev_node = cgraph->nodes[i - 1]; - try_launch_concurrent_event(prev_node); - - if (is_concurrent_event_active) { - cuda_ctx->curr_stream_no = concurrent_event->stream_mapping[node]; - GGML_LOG_DEBUG("Setting stream no to %d for node %s\n", cuda_ctx->curr_stream_no, node->name); - } - } - -#ifdef GGML_CUDA_DEBUG - const int nodes_fused = i - prev_i - 1; - if (nodes_fused > 0) { - GGML_LOG_INFO("nodes_fused: %d\n", nodes_fused); - } -#endif - prev_i = i; - - if (ggml_is_empty(node) || node->op == GGML_OP_RESHAPE || node->op == GGML_OP_TRANSPOSE || node->op == GGML_OP_VIEW || node->op == GGML_OP_PERMUTE || node->op == GGML_OP_NONE) { - continue; - } - - if ((node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { - continue; - } - - // start of fusion operations - static bool disable_fusion = (getenv("GGML_CUDA_DISABLE_FUSION") != nullptr); - if (!disable_fusion) { - ggml_cuda_topk_moe_args args; - - if (cgraph->nodes[i]->op == GGML_OP_UNARY || cgraph->nodes[i]->op == GGML_OP_SOFT_MAX || - cgraph->nodes[i]->op == GGML_OP_ARGSORT) { - const bool can_fuse = ggml_cuda_topk_moe_fusion(cgraph, i, args); - - std::vector ops; - - if (can_fuse) { - const ggml_tensor * logits = node->src[0]; - ggml_tensor * weights = nullptr; - ggml_tensor * ids = nullptr; - const ggml_tensor * bias = nullptr; - const ggml_tensor * clamp = nullptr; - const ggml_tensor * scale = nullptr; - - if (!args.delayed_softmax) { - ggml_op gating_op = args.sigmoid ? GGML_OP_UNARY : GGML_OP_SOFT_MAX; - int out_nodes[2]; // nodes which can't be elided - - if (args.prob_bias) { - bias = cgraph->nodes[i + 2]->src[1]; - ops.insert(ops.end(), { gating_op, GGML_OP_RESHAPE, GGML_OP_ADD, GGML_OP_ARGSORT, - GGML_OP_VIEW, GGML_OP_GET_ROWS }); - out_nodes[0] = i + 4; - ids = cgraph->nodes[i + 4]; - } else { - ops.insert(ops.end(), { gating_op, GGML_OP_RESHAPE, GGML_OP_ARGSORT, GGML_OP_VIEW, - GGML_OP_GET_ROWS }); - out_nodes[0] = i + 3; - ids = cgraph->nodes[i + 3]; - } - - if (args.norm) { - ops.insert(ops.end(), { GGML_OP_RESHAPE, GGML_OP_SUM_ROWS, GGML_OP_CLAMP, - GGML_OP_DIV, GGML_OP_RESHAPE }); - clamp = cgraph->nodes[i + ops.size() - 3]; - } - if (args.scale) { - ops.insert(ops.end(), { GGML_OP_SCALE }); - scale = cgraph->nodes[i + ops.size() - 1]; - } - - weights = cgraph->nodes[i + ops.size() - 1]; - out_nodes[1] = i + ops.size() - 1; - - if (ggml_can_fuse_subgraph(cgraph, i, ops.size(), ops.data(), out_nodes, 2) && - ggml_cuda_should_use_topk_moe(node, logits, weights, ids) && - ggml_cuda_check_fusion_memory_ranges(cgraph, i, ops.size(), out_nodes, 2, /*is_topk_moe=*/ true)) { - ggml_cuda_op_topk_moe(*cuda_ctx, logits, weights, ids, clamp, scale, bias, args); - i += ops.size() - 1; - continue; - } - } else if (!args.norm && !args.prob_bias) { - //special case gpt-oss, no norm, no bias. - ops.insert(ops.end(), { GGML_OP_ARGSORT, GGML_OP_VIEW, GGML_OP_GET_ROWS, - GGML_OP_RESHAPE, GGML_OP_SOFT_MAX, GGML_OP_RESHAPE }); - weights = cgraph->nodes[i + 5]; - ids = cgraph->nodes[i + 1]; - const ggml_tensor * softmax = cgraph->nodes[i + 4]; - - int out_nodes[2] = { i + 1, i + 5 }; - if (ggml_can_fuse_subgraph(cgraph, i, ops.size(), ops.data(), out_nodes, 2) && - ggml_cuda_should_use_topk_moe(softmax, logits, weights, ids) && - ggml_cuda_check_fusion_memory_ranges(cgraph, i, ops.size(), out_nodes, 2, /*is_topk_moe=*/ true)) { - ggml_cuda_op_topk_moe(*cuda_ctx, logits, weights, ids, clamp, scale, bias, args); - i += ops.size() - 1; - continue; - } - } - } - } - - if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_ROPE, GGML_OP_VIEW, GGML_OP_SET_ROWS }, {})) { - ggml_tensor * rope = cgraph->nodes[i]; - ggml_tensor * set_rows = cgraph->nodes[i + 2]; - - ggml_cuda_op_rope_fused(*cuda_ctx, rope, set_rows); - i += 2; - continue; - } - - if (node->op == GGML_OP_ADD || node->op == GGML_OP_MUL) { - int n_fuse = 0; - ggml_op ops[8]; - std::fill(ops, ops + 8, node->op); - - for (; n_fuse <= 6; ++n_fuse){ - if (!ggml_can_fuse(cgraph, i + n_fuse, ops + n_fuse, 2)) { - break; - } - if (cgraph->nodes[i + n_fuse] != cgraph->nodes[i + n_fuse + 1]->src[0]) { - break; - } - if (!ggml_are_same_layout(cgraph->nodes[i + n_fuse]->src[1], cgraph->nodes[i + n_fuse + 1]->src[1])) { - break; - } - } - - n_fuse++; - - if (n_fuse > 1) { - ggml_tensor fused_node; - memcpy(&fused_node, node, sizeof(ggml_tensor)); - for (int j = 0; j < n_fuse - 1; ++j) { - fused_node.src[j + 2] = cgraph->nodes[i + j + 1]->src[1]; - } - fused_node.data = cgraph->nodes[i + n_fuse - 1]->data; - if (node->op == GGML_OP_ADD) { - ggml_cuda_op_fused_add(*cuda_ctx, &fused_node, n_fuse); - } else { - ggml_cuda_op_fused_mul(*cuda_ctx, &fused_node, n_fuse); - } - i += n_fuse - 1; - - continue; - } - } - - bool fused_mul_mat_vec = false; - int fused_node_count = 0; - - for (ggml_op op : { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT_ID }) { - const ggml_op bias_op = op == GGML_OP_MUL_MAT ? GGML_OP_ADD : GGML_OP_ADD_ID; - - if (ggml_cuda_can_fuse(cgraph, i, { op, bias_op, op, bias_op, GGML_OP_GLU }, {})) { - ggml_tensor * glu = cgraph->nodes[i + 4]; - ggml_tensor * gate_bias_n = glu->src[0]; - ggml_tensor * up_bias_n = glu->src[1]; - - //we don't assume the order for {gate, up}. Instead infer it from the bias tensor - ggml_tensor * gate_n = nullptr; - ggml_tensor * up_n = nullptr; - - if (gate_bias_n->src[0] == cgraph->nodes[i] || gate_bias_n->src[1] == cgraph->nodes[i]) { - gate_n = cgraph->nodes[i]; - up_n = cgraph->nodes[i + 2]; - } else if (gate_bias_n->src[0] == cgraph->nodes[i + 2] || gate_bias_n->src[1] == cgraph->nodes[i + 2]) { - gate_n = cgraph->nodes[i + 2]; - up_n = cgraph->nodes[i]; - } else { - continue; - } - - auto get_bias_tensor = [](const ggml_tensor * bias_node, const ggml_tensor * mul_node, ggml_op op_bias) { - if (op_bias == GGML_OP_ADD) { - if (bias_node->src[0] == mul_node) { - return bias_node->src[1]; - } - if (bias_node->src[1] == mul_node) { - return bias_node->src[0]; - } - return (ggml_tensor *) nullptr; - } - GGML_ASSERT(op_bias == GGML_OP_ADD_ID); - GGML_ASSERT(bias_node->src[0] == mul_node); - return bias_node->src[1]; - }; - - ggml_tensor * up_bias_tensor = get_bias_tensor(up_bias_n, up_n, bias_op); - ggml_tensor * gate_bias_tensor = get_bias_tensor(gate_bias_n, gate_n, bias_op); - - if (!up_bias_tensor || !gate_bias_tensor) { - continue; - } - - // we don't support repeating adds - if (bias_op == GGML_OP_ADD && - (!ggml_are_same_shape(gate_bias_n->src[0], gate_bias_n->src[1]) || - !ggml_are_same_shape(up_bias_n->src[0], up_bias_n->src[1]))) { - continue; - } - - const ggml_tensor * src0 = up_n->src[0]; - const ggml_tensor * src1 = up_n->src[1]; - const ggml_tensor * ids = up_n->src[2]; - - if (ggml_cuda_should_fuse_mul_mat_vec_f(up_n)) { - ggml_cuda_mm_fusion_args_host fusion_data{}; - fusion_data.gate = gate_n->src[0]; - fusion_data.x_bias = up_bias_tensor; - fusion_data.gate_bias = gate_bias_tensor; - ggml_cuda_set_fusion_glu_params(fusion_data, glu); - - ggml_cuda_mul_mat_vec_f(*cuda_ctx, src0, src1, ids, glu, &fusion_data); - fused_mul_mat_vec = true; - fused_node_count = 5; - break; - } - - if (ggml_cuda_should_fuse_mul_mat_vec_q(up_n)) { - ggml_cuda_mm_fusion_args_host fusion_data{}; - fusion_data.gate = gate_n->src[0]; - fusion_data.x_bias = up_bias_tensor; - fusion_data.gate_bias = gate_bias_tensor; - ggml_cuda_set_fusion_glu_params(fusion_data, glu); - - ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, glu, &fusion_data); - fused_mul_mat_vec = true; - fused_node_count = 5; - break; - } - } else if (ggml_cuda_can_fuse(cgraph, i, { op, op, GGML_OP_GLU }, {})) { - ggml_tensor * glu = cgraph->nodes[i + 2]; - ggml_tensor * gate = glu->src[0]; - ggml_tensor * up = glu->src[1]; - - bool ok = (gate == cgraph->nodes[i] && up == cgraph->nodes[i + 1]) - || (gate == cgraph->nodes[i + 1] && up == cgraph->nodes[i]); - - if (!ok) continue; - - const ggml_tensor * src0 = up->src[0]; - const ggml_tensor * src1 = up->src[1]; - const ggml_tensor * ids = up->src[2]; - - if (ggml_cuda_should_fuse_mul_mat_vec_f(up)) { - ggml_cuda_mm_fusion_args_host fusion_data{}; - fusion_data.gate = gate->src[0]; - ggml_cuda_set_fusion_glu_params(fusion_data, glu); - - ggml_cuda_mul_mat_vec_f(*cuda_ctx, src0, src1, ids, glu, &fusion_data); - fused_mul_mat_vec = true; - fused_node_count = 3; - break; - } - - if (ggml_cuda_should_fuse_mul_mat_vec_q(up)) { - ggml_cuda_mm_fusion_args_host fusion_data{}; - fusion_data.gate = gate->src[0]; - ggml_cuda_set_fusion_glu_params(fusion_data, glu); - - ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, glu, &fusion_data); - fused_mul_mat_vec = true; - fused_node_count = 3; - break; - } - } - } - - if (fused_mul_mat_vec) { - i += fused_node_count - 1; - continue; - } - - fused_mul_mat_vec = false; - fused_node_count = 0; - - for (ggml_op op : { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT_ID }) { - const ggml_op bias_op = op == GGML_OP_MUL_MAT ? GGML_OP_ADD : GGML_OP_ADD_ID; - - if (!ggml_can_fuse(cgraph, i, { op, bias_op })) { - continue; - } - - ggml_tensor * mm_node = cgraph->nodes[i]; - ggml_tensor * bias_node = cgraph->nodes[i + 1]; - - ggml_tensor * bias_tensor = nullptr; - if (bias_op == GGML_OP_ADD) { - if (bias_node->src[0] == mm_node) { - bias_tensor = bias_node->src[1]; - } else if (bias_node->src[1] == mm_node) { - bias_tensor = bias_node->src[0]; - } else { - continue; - } - } else { - if (bias_node->src[0] != mm_node) { - continue; - } - bias_tensor = bias_node->src[1]; - } - - const ggml_tensor * src0 = mm_node->src[0]; - const ggml_tensor * src1 = mm_node->src[1]; - const ggml_tensor * ids = mm_node->src[2]; - - if (bias_op == GGML_OP_ADD_ID && bias_node->src[2] != ids) { - continue; - } - - if (bias_op == GGML_OP_ADD && !ggml_are_same_shape(bias_node->src[0], bias_node->src[1])) { - continue; - } - - ggml_cuda_mm_fusion_args_host fusion_data{}; - fusion_data.x_bias = bias_tensor; - - if (ggml_cuda_should_fuse_mul_mat_vec_f(mm_node)) { - ggml_cuda_mul_mat_vec_f(*cuda_ctx, src0, src1, ids, bias_node, &fusion_data); - fused_mul_mat_vec = true; - fused_node_count = 2; - break; - } - - if (ggml_cuda_should_fuse_mul_mat_vec_q(mm_node)) { - ggml_cuda_mul_mat_vec_q(*cuda_ctx, src0, src1, ids, bias_node, &fusion_data); - fused_mul_mat_vec = true; - fused_node_count = 2; - break; - } - } - - if (fused_mul_mat_vec) { - i += fused_node_count - 1; - continue; - } - - if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL, GGML_OP_ADD}, {})) { - ggml_cuda_op_rms_norm_fused_add(*cuda_ctx, node, cgraph->nodes[i+1], cgraph->nodes[i+2]); - i += 2; - continue; - } - - if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL}, {})) { - ggml_cuda_op_rms_norm_fused(*cuda_ctx, node, cgraph->nodes[i+1]); - i++; - continue; - } - - if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_SSM_CONV, GGML_OP_UNARY }, { GGML_UNARY_OP_SILU })) { - ggml_cuda_op_ssm_conv(*cuda_ctx, node, cgraph->nodes[i+1]); - i++; - continue; - } - - if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_UNARY, GGML_OP_MUL }, { GGML_UNARY_OP_SILU }) || - ggml_cuda_can_fuse(cgraph, i, { GGML_OP_UNARY, GGML_OP_MUL }, { GGML_UNARY_OP_SIGMOID }) || - ggml_cuda_can_fuse(cgraph, i, { GGML_OP_UNARY, GGML_OP_MUL }, { GGML_UNARY_OP_SOFTPLUS })) { - ggml_cuda_op_unary_mul(*cuda_ctx, node, cgraph->nodes[i+1]); - i++; - continue; - } - - if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_SCALE, GGML_OP_UNARY, GGML_OP_SCALE }, { GGML_UNARY_OP_TANH })) { - i += 2; - ggml_cuda_op_softcap(*cuda_ctx, cgraph->nodes[i], node); - continue; - } - } -#ifndef NDEBUG - assert(node->buffer->buft == ggml_backend_cuda_buffer_type(cuda_ctx->device)); - for (int j = 0; j < GGML_MAX_SRC; j++) { - if (node->src[j] != nullptr) { - assert(node->src[j]->buffer); - assert(node->src[j]->buffer->buft == ggml_backend_cuda_buffer_type(cuda_ctx->device) || - ggml_backend_buft_is_cuda_split(node->src[j]->buffer->buft) || (integrated && ggml_backend_buft_is_cuda_host(node->src[j]->buffer->buft))); - } - } -#else - GGML_UNUSED(integrated); -#endif // NDEBUG - - bool ok = ggml_cuda_compute_forward(*cuda_ctx, node); - if (!ok) { - GGML_LOG_ERROR("%s: op not supported %s (%s)\n", __func__, node->name, ggml_op_name(node->op)); - } - GGML_ASSERT(ok); - - if (!is_concurrent_event_active) { - try_launch_concurrent_event(node); - } - } - } - -#ifdef USE_CUDA_GRAPH - ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); - if (use_cuda_graph && cuda_graph_update_required) { // End CUDA graph capture - if (graph->graph != nullptr) { - CUDA_CHECK(cudaGraphDestroy(graph->graph)); - graph->graph = nullptr; - } - - CUDA_CHECK(cudaStreamEndCapture(cuda_ctx->stream(), &graph->graph)); - graph_evaluated_or_captured = true; // CUDA graph has been captured - - std::lock_guard lock(ggml_cuda_lock); - if (ggml_cuda_lock_counter.fetch_sub(1, std::memory_order_relaxed) == 1) { - ggml_cuda_lock_cv.notify_all(); - } - } else { - graph_evaluated_or_captured = true; // ggml graph has been directly evaluated - } - } - - if (use_cuda_graph) { - ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); - if (graph->instance == nullptr) { // Create executable graph from captured graph. - CUDA_CHECK(cudaGraphInstantiate(&graph->instance, graph->graph, NULL, NULL, 0)); - } - if (cuda_graph_update_required) { // Update graph executable - ggml_cuda_graph_update_executable(cuda_ctx, graph_key); - } - // Launch graph - CUDA_CHECK(cudaGraphLaunch(graph->instance, cuda_ctx->stream())); -#else - GGML_UNUSED(graph_key); - graph_evaluated_or_captured = true; -#endif // USE_CUDA_GRAPH - } -} - -#ifdef USE_CUDA_GRAPH -static bool ggml_cuda_graph_set_enabled(ggml_backend_cuda_context * cuda_ctx, const void * graph_key) { - ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); - - if (graph->graph == nullptr) { - if (ggml_cuda_info().devices[cuda_ctx->device].cc < GGML_CUDA_CC_AMPERE) { - if (!graph->disable_due_to_gpu_arch) { - GGML_LOG_DEBUG("%s: disabling CUDA graphs due to GPU architecture\n", __func__); - } - graph->disable_due_to_gpu_arch = true; - } - } - - return graph->is_enabled(); -} -#endif // USE_CUDA_GRAPH - -static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, ggml_cgraph * cgraph) { - ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context; - - ggml_cuda_set_device(cuda_ctx->device); - - // LIFO-free last evaluation's memoized q8_1 activations. - while (!cuda_ctx->luce_q8_memo.empty()) { - cuda_ctx->luce_q8_memo.pop_back(); - } - - bool use_cuda_graph = false; - bool cuda_graph_update_required = false; - const void * graph_key = nullptr; - -#ifdef USE_CUDA_GRAPH - graph_key = ggml_cuda_graph_get_key(cgraph); - - ggml_cuda_graph_set_enabled(cuda_ctx, graph_key); - - ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); - if (graph->is_enabled()) { - const bool graph_compatible = ggml_cuda_graph_check_compability(cgraph); - if (graph_compatible) { - const bool properties_changed = ggml_cuda_graph_update_required(cuda_ctx, cgraph); - - if (!graph->warmup_complete) { - // Warmup: need at least 2 calls with no property change on the 2nd call - if (!properties_changed) { - graph->warmup_complete = true; - GGML_LOG_DEBUG("%s: CUDA graph warmup complete\n", __func__); - use_cuda_graph = true; - cuda_graph_update_required = true; - } - // else: properties changed or first call - execute directly (use_cuda_graph stays false) - } else { - // Post-warmup: normal CUDA graph operation - if (properties_changed) { - // Properties changed - reset warmup, execute directly until stable again - graph->warmup_complete = false; - GGML_LOG_DEBUG("%s: CUDA graph warmup reset\n", __func__); - } else { - use_cuda_graph = true; - cuda_graph_update_required = graph->instance == nullptr; - } - } - } - } -#endif // USE_CUDA_GRAPH - -#ifdef USE_CUDA_GRAPH - { - static const bool luce_graph_stats = getenv("GGML_CUDA_GRAPH_STATS") != nullptr; - if (luce_graph_stats && graph_key) { - ggml_cuda_graph * g = cuda_ctx->cuda_graph(graph_key); - g->stat_total++; - if (use_cuda_graph && !cuda_graph_update_required) g->stat_replay++; - else if (use_cuda_graph) g->stat_capture++; - else g->stat_eager++; - if (g->stat_total % 200 == 0) { - GGML_LOG_INFO("[cuda-graph-stats] key=%p n_nodes=%d total=%llu replay=%llu capture=%llu eager=%llu enabled=%d\n", - graph_key, cgraph->n_nodes, - (unsigned long long)g->stat_total, (unsigned long long)g->stat_replay, - (unsigned long long)g->stat_capture, (unsigned long long)g->stat_eager, - (int)g->is_enabled()); - } - } - } -#endif // USE_CUDA_GRAPH - - if (use_cuda_graph && cuda_graph_update_required) { - // Start CUDA graph capture - { - std::lock_guard lock(ggml_cuda_lock); - ggml_cuda_lock_counter.fetch_add(1, std::memory_order_relaxed); - } - - CUDA_CHECK(cudaStreamBeginCapture(cuda_ctx->stream(), cudaStreamCaptureModeRelaxed)); - } - - ggml_cuda_graph_evaluate_and_capture(cuda_ctx, cgraph, use_cuda_graph, cuda_graph_update_required, graph_key); - - return GGML_STATUS_SUCCESS; -} - -static void ggml_backend_cuda_event_record(ggml_backend_t backend, ggml_backend_event_t event) { - ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context; - - CUDA_CHECK(cudaEventRecord((cudaEvent_t)event->context, cuda_ctx->stream())); -} - -static void ggml_backend_cuda_event_wait(ggml_backend_t backend, ggml_backend_event_t event) { - ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *)backend->context; - - if (ggml_backend_is_cuda(backend)) { - CUDA_CHECK(cudaStreamWaitEvent(cuda_ctx->stream(), (cudaEvent_t)event->context, 0)); - } else { -#if 0 - // untested - auto wait_fn = [](void * user_data) { - ggml_backend_event_t event = (ggml_backend_event_t)user_data; - ggml_backend_event_synchronize(event); - }; - - CUDA_CHECK(cudaLaunchHostFunc(cuda_ctx->stream(), wait_fn, event)); -#endif - GGML_ABORT("fatal error"); - } -} - -static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph * cgraph) { - ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context; - -#ifdef USE_CUDA_GRAPH - const void * graph_key = ggml_cuda_graph_get_key(cgraph); - const bool use_cuda_graph = ggml_cuda_graph_set_enabled(cuda_ctx, graph_key); -#else - const bool use_cuda_graph = false; - GGML_UNUSED(cuda_ctx); - GGML_UNUSED(cgraph); -#endif - - static bool enable_graph_optimization = [] { - const char * env = getenv("GGML_CUDA_GRAPH_OPT"); - return env != nullptr && atoi(env) == 1; - }(); - - if (!enable_graph_optimization) { - return; - } - - ggml_cuda_stream_context & stream_context = cuda_ctx->stream_context(); - stream_context.reset(); - - if (!use_cuda_graph || ggml_backend_cuda_get_device_count() != 1) { - return; - } - - // number of out-degrees for a particular node - std::unordered_map fan_out; - // reverse mapping of node to index in the cgraph - std::unordered_map node_indices; - - const auto & is_noop = [](const ggml_tensor * node) -> bool { - return ggml_is_empty(node) || node->op == GGML_OP_NONE || node->op == GGML_OP_RESHAPE || - node->op == GGML_OP_TRANSPOSE || node->op == GGML_OP_VIEW || node->op == GGML_OP_PERMUTE; - }; - - const auto & depends_on = [](const ggml_tensor * dst, const ggml_tensor * src) -> bool { - for (uint32_t s = 0; s < GGML_MAX_SRC; ++s) { - if (dst->src[s] == src) { - return true; - } - } - // implicit dependency if they view the same tensor - const ggml_tensor * dst2 = dst->view_src ? dst->view_src : dst; - const ggml_tensor * src2 = src->view_src ? src->view_src : src; - if (dst2 == src2) { - return true; - } - return false; - }; - - for (int node_idx = 0; node_idx < cgraph->n_nodes; node_idx++) { - const ggml_tensor * node = cgraph->nodes[node_idx]; - node_indices[node] = node_idx; - - if (is_noop(node)) { - continue; - } - for (int src_idx = 0; src_idx < GGML_MAX_SRC; ++src_idx) { - const ggml_tensor * src = cgraph->nodes[node_idx]->src[src_idx]; - //TODO: check why nrows > 1 fails - if (node && !is_noop(node) && ggml_nrows(node) <= 1) { - fan_out[src] += 1; - } - } - } - - // Target Q, K, V for concurrency - // this is a more general way to find nodes which can be candidates for concurrency (although it has not been tested for anything else): - // 1. find fan-out (fork) nodes where the same input is used at least N times (in QKV, it would be "attn-norm") - // 2. find the join node, where 2 or more of the outputs are required (in QKV, this would "KQ" or "flash-attn") - // 3. account for all branches from the fork to the join - // 4. To extend lifetimes of the tensors, we interleave the branches (see below for more details) - // 5. save the original cgraph and restore it in graph_compute, to enable fusion within streams - // See discussion: https://github.com/ggml-org/llama.cpp/pull/16991#issuecomment-3522620030 - - const int min_fan_out = 3; - const int max_fan_out = 3; - - // store {fork_idx, join_idx} - std::vector> concurrent_node_ranges; - - for (const auto & [root_node, count] : fan_out) { - if (count >= min_fan_out && count <= max_fan_out) { - const int root_node_idx = node_indices[root_node]; - - // only optimize for attn_norm - // TODO: make this more generic - if (!strstr(root_node->name, "attn_norm")) { - continue; - } - - bool is_part_of_event = false; - for (const auto & [start, end] : concurrent_node_ranges) { - if (root_node_idx >= start && root_node_idx <= end) { - is_part_of_event = true; - } - } - - if (is_part_of_event) { - continue; - } - - std::vector> nodes_per_branch; - for (int i = root_node_idx + 1; i < cgraph->n_nodes; ++i) { - const ggml_tensor * node = cgraph->nodes[i]; - if (!is_noop(node) && depends_on(node, root_node)) { - nodes_per_branch.push_back({ node }); - } - } - - GGML_ASSERT(nodes_per_branch.size() == (size_t) count); - - //find the join point - const ggml_tensor * join_node = nullptr; - - const auto & belongs_to_branch = [&](const ggml_tensor * node, - const std::vector & branch) -> bool { - for (const ggml_tensor * n : branch) { - if (depends_on(node, n)) { - return true; - } - } - return false; - }; - - for (int i = root_node_idx + 1; i < cgraph->n_nodes; ++i) { - const ggml_tensor * curr_node = cgraph->nodes[i]; - - int num_joins = 0; - for (size_t branch_idx = 0; branch_idx < nodes_per_branch.size(); branch_idx++) { - if (belongs_to_branch(curr_node, nodes_per_branch[branch_idx])) { - num_joins++; - } - } - - if (num_joins >= 2) { - join_node = curr_node; - break; - } - - bool found_branch = false; - for (size_t branch_idx = 0; branch_idx < nodes_per_branch.size(); branch_idx++) { - std::vector & branch_vec = nodes_per_branch[branch_idx]; - if (belongs_to_branch(curr_node, branch_vec)) { - //continue accumulating - if (std::find(branch_vec.begin(), branch_vec.end(), curr_node) == branch_vec.end()) { - branch_vec.push_back(curr_node); - } - found_branch = true; - } - } - - if (!found_branch && is_noop(curr_node)) { - // we can put it in any branch because it will be ignored - nodes_per_branch[0].push_back({ curr_node }); - } - } - - if (join_node) { - //Create ggml_cuda_concurrent_event - ggml_cuda_concurrent_event concurrent_event(nodes_per_branch.size()); - concurrent_event.join_node = join_node; - - for (size_t branch_idx = 0; branch_idx < nodes_per_branch.size(); branch_idx++) { - for (const ggml_tensor * n : nodes_per_branch[branch_idx]) { - concurrent_event.stream_mapping[n] = branch_idx + 1; - } - } - - int fork_node_idx = node_indices[root_node]; - int join_node_idx = node_indices[join_node]; - - int current_branch_idx = 0; - int current_node_idx = fork_node_idx + 1; - const int n_branches = nodes_per_branch.size(); - - int total_branch_nodes = 0; - for (std::vector branch_nodes : nodes_per_branch) { - total_branch_nodes += branch_nodes.size(); - } - - // there are other nodes in the middle which are unaccounted for - // usually (cpy) nodes, then ignore this fork - if (join_node_idx - fork_node_idx - 1 != total_branch_nodes) { - GGML_LOG_DEBUG( - "Skipping %s because the number of nodes in the middle is not equal to the total number of " - "branch nodes %d != %d\n", - root_node->name, join_node_idx - fork_node_idx - 1, total_branch_nodes); - continue; - } - - // Save the original order of nodes in this region before interleaving - // This is used later to restore grouping for fusion within streams - concurrent_event.original_order.reserve(total_branch_nodes); - for (int i = fork_node_idx + 1; i < join_node_idx; ++i) { - concurrent_event.original_order.push_back(cgraph->nodes[i]); - } - - std::unordered_map & concurrent_events = cuda_ctx->stream_context().concurrent_events; - GGML_ASSERT(concurrent_events.find(root_node) == concurrent_events.end()); - concurrent_events.emplace(root_node, std::move(concurrent_event)); - GGML_LOG_DEBUG("Adding stream at node %s %p\n", root_node->name, root_node); - concurrent_node_ranges.emplace_back(fork_node_idx, join_node_idx); - - // interleave tensors to extend lifetimes so that ggml graph doesn't recycle them - // example transformation: - // [attn-norm, QMul, QNorm, QRope, KMul, KNorm, KRope, VMul, attn] -> - // [attn-norm, QMul, KMul, VMul, QNorm, VNorm, QRope, KRope, attn] - while (current_node_idx < join_node_idx) { - std::vector & branch_nodes = nodes_per_branch[current_branch_idx]; - - bool has_node = false; - for (std::vector branch_node : nodes_per_branch) { - has_node |= branch_node.size() > 0; - } - - GGML_ASSERT(has_node); - - if (branch_nodes.empty()) { - current_branch_idx = (current_branch_idx + 1) % n_branches; - continue; - } - - cgraph->nodes[current_node_idx] = const_cast(branch_nodes.front()); - current_node_idx++; - branch_nodes.erase(branch_nodes.begin()); - - // append all empty nodes - while (!branch_nodes.empty() && is_noop(branch_nodes.front())) { - cgraph->nodes[current_node_idx] = const_cast(branch_nodes.front()); - current_node_idx++; - branch_nodes.erase(branch_nodes.begin()); - } - - current_branch_idx = (current_branch_idx + 1) % n_branches; - } - } - } - } -} - -static const ggml_backend_i ggml_backend_cuda_interface = { - /* .get_name = */ ggml_backend_cuda_get_name, - /* .free = */ ggml_backend_cuda_free, - /* .set_tensor_async = */ ggml_backend_cuda_set_tensor_async, - /* .get_tensor_async = */ ggml_backend_cuda_get_tensor_async, - /* .get_tensor_2d_async = */ ggml_backend_cuda_set_tensor_2d_async, - /* .set_tensor_2d_async = */ ggml_backend_cuda_get_tensor_2d_async, - /* .cpy_tensor_async = */ ggml_backend_cuda_cpy_tensor_async, - /* .synchronize = */ ggml_backend_cuda_synchronize, - /* .graph_plan_create = */ NULL, - /* .graph_plan_free = */ NULL, - /* .graph_plan_update = */ NULL, - /* .graph_plan_compute = */ NULL, - /* .graph_compute = */ ggml_backend_cuda_graph_compute, - /* .event_record = */ ggml_backend_cuda_event_record, - /* .event_wait = */ ggml_backend_cuda_event_wait, - /* .graph_optimize = */ ggml_backend_cuda_graph_optimize, -}; - -static ggml_guid_t ggml_backend_cuda_guid() { - static ggml_guid guid = { 0x2c, 0xdd, 0xe8, 0x1c, 0x65, 0xb3, 0x65, 0x73, 0x6a, 0x12, 0x88, 0x61, 0x1c, 0xc9, 0xdc, 0x25 }; - return &guid; -} - -bool ggml_backend_is_cuda(ggml_backend_t backend) { - return backend != NULL && ggml_guid_matches(backend->guid, ggml_backend_cuda_guid()); -} - -int ggml_backend_cuda_get_device_count() { - return ggml_cuda_info().device_count; -} - -void ggml_backend_cuda_get_device_description(int device, char * description, size_t description_size) { - cudaDeviceProp prop; - CUDA_CHECK(cudaGetDeviceProperties(&prop, device)); - snprintf(description, description_size, "%s", prop.name); -} - -void ggml_backend_cuda_get_device_memory(int device, size_t * free, size_t * total) { - ggml_cuda_set_device(device); - - CUDA_CHECK(cudaMemGetInfo(free, total)); -} - -bool ggml_backend_cuda_register_host_buffer(void * buffer, size_t size) { - if (getenv("GGML_CUDA_REGISTER_HOST") == nullptr) { - return false; - } - -#if CUDART_VERSION >= 11010 || defined(GGML_USE_MUSA) || defined(GGML_USE_HIP) - cudaError_t err = cudaHostRegister(buffer, size, cudaHostRegisterPortable | cudaHostRegisterReadOnly); - if (err != cudaSuccess) { - // clear the error - (void)cudaGetLastError(); - - GGML_LOG_DEBUG("%s: failed to register %.2f MiB of pinned memory: %s\n", __func__, - size / 1024.0 / 1024.0, cudaGetErrorString(err)); - return false; - } - return true; -#else - GGML_UNUSED(buffer); - GGML_UNUSED(size); - return false; -#endif // CUDART_VERSION >= 11010 || defined(GGML_USE_MUSA) -} - -void ggml_backend_cuda_unregister_host_buffer(void * buffer) { - if (getenv("GGML_CUDA_REGISTER_HOST") == nullptr) { - return; - } - - cudaError_t err = cudaHostUnregister(buffer); - if (err != cudaSuccess) { - // clear the error - (void)cudaGetLastError(); - } -} - - -// backend device - -struct ggml_backend_cuda_device_context { - int device; - std::string name; - std::string description; - std::string pci_bus_id; - int op_offload_min_batch_size; -}; - -static const char * ggml_backend_cuda_device_get_name(ggml_backend_dev_t dev) { - ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *)dev->context; - return ctx->name.c_str(); -} - -static const char * ggml_backend_cuda_device_get_description(ggml_backend_dev_t dev) { - ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *)dev->context; - return ctx->description.c_str(); -} - -#if defined(__linux__) -// Helper function to get available memory from /proc/meminfo for UMA systems -static bool ggml_backend_cuda_get_available_uma_memory(long * available_memory_kb, long * free_swap_kb) { - FILE * meminfo_file = nullptr; - // 2KB buffer for reading /proc/meminfo since it does not report size info, should be enough - const size_t BUFFER_SIZE = 2048; - auto file_buffer = std::make_unique(BUFFER_SIZE); - size_t bytes_read = 0; - long huge_tlb_total_pages = -1; - long huge_tlb_free_pages = -1; - long huge_tlb_page_size = -1; - - if (available_memory_kb == nullptr || free_swap_kb == nullptr) { - return false; - } - - meminfo_file = fopen("/proc/meminfo", "r"); - if (meminfo_file == nullptr) { - GGML_LOG_ERROR("%s: failed to open /proc/meminfo\n", __func__); - return false; - } - - // Read file into buffer - bytes_read = fread(file_buffer.get(), 1, BUFFER_SIZE - 1, meminfo_file); - fclose(meminfo_file); - - if (bytes_read == 0) { - GGML_LOG_ERROR("%s: failed to read from /proc/meminfo\n", __func__); - return false; - } - file_buffer[bytes_read] = '\0'; - - *available_memory_kb = -1; - *free_swap_kb = -1; - - // Parse the file buffer line by line - char * line = file_buffer.get(); - char * line_next; - while (line < file_buffer.get() + bytes_read) { - // Find the end of the current line - line_next = strchr(line, '\n'); - if (line_next != nullptr) { - *line_next = '\0'; - line_next++; - } else { - line_next = file_buffer.get() + bytes_read; - } - - long value; - if (sscanf(line, "MemAvailable: %ld kB", &value) == 1) { - *available_memory_kb = value; - } else if (sscanf(line, "SwapFree: %ld kB", &value) == 1) { - *free_swap_kb = value; - } else if (sscanf(line, "HugePages_Total: %ld", &value) == 1) { - huge_tlb_total_pages = value; - } else if (sscanf(line, "HugePages_Free: %ld", &value) == 1) { - huge_tlb_free_pages = value; - } else if (sscanf(line, "Hugepagesize: %ld kB", &value) == 1) { - huge_tlb_page_size = value; - } - - line = line_next; - } - - if (huge_tlb_total_pages != 0 && huge_tlb_total_pages != -1) { - *available_memory_kb = huge_tlb_free_pages * huge_tlb_page_size; - - // Hugetlbfs pages are not swappable. - *free_swap_kb = 0; - } - - GGML_LOG_DEBUG("%s: final available_memory_kb: %ld\n", __func__, *available_memory_kb); - return true; -} -#endif // defined(__linux__) - -static void ggml_backend_cuda_device_get_memory(ggml_backend_dev_t dev, size_t * free, size_t * total) { - ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *)dev->context; - ggml_cuda_set_device(ctx->device); - CUDA_CHECK(cudaMemGetInfo(free, total)); - -// ref: https://github.com/ggml-org/llama.cpp/pull/17368 -#if defined(__linux__) - // Check if this is a UMA (Unified Memory Architecture) system - cudaDeviceProp prop; - CUDA_CHECK(cudaGetDeviceProperties(&prop, ctx->device)); - - // Check if UMA is explicitly enabled via environment variable - bool uma_env = getenv("GGML_CUDA_ENABLE_UNIFIED_MEMORY") != nullptr; - bool is_uma = prop.integrated > 0 || uma_env; - - if (is_uma) { - // For UMA systems (like DGX Spark), use system memory info - long available_memory_kb = 0; - long free_swap_kb = 0; - - if (ggml_backend_cuda_get_available_uma_memory(&available_memory_kb, &free_swap_kb) && available_memory_kb > 0) { - *free = (size_t)available_memory_kb * 1024; - } else { - GGML_LOG_ERROR("%s: /proc/meminfo reading failed, using cudaMemGetInfo\n", __func__); - } - } -#endif // defined(__linux__) - -} - -static enum ggml_backend_dev_type ggml_backend_cuda_device_get_type(ggml_backend_dev_t dev) { - GGML_UNUSED(dev); - return GGML_BACKEND_DEVICE_TYPE_GPU; -} - -static void ggml_backend_cuda_device_get_props(ggml_backend_dev_t dev, ggml_backend_dev_props * props) { - ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *)dev->context; - - props->name = ggml_backend_cuda_device_get_name(dev); - props->description = ggml_backend_cuda_device_get_description(dev); - props->type = ggml_backend_cuda_device_get_type(dev); - props->device_id = ctx->pci_bus_id.empty() ? nullptr : ctx->pci_bus_id.c_str(); - ggml_backend_cuda_device_get_memory(dev, &props->memory_free, &props->memory_total); - - bool host_buffer = getenv("GGML_CUDA_NO_PINNED") == nullptr; -#ifdef GGML_CUDA_NO_PEER_COPY - bool events = false; -#else - bool events = true; -#endif - - props->caps = { - /* .async = */ true, - /* .host_buffer = */ host_buffer, - /* .buffer_from_host_ptr = */ false, - /* .events = */ events, - }; -} - -static ggml_backend_t ggml_backend_cuda_device_init_backend(ggml_backend_dev_t dev, const char * params) { - GGML_UNUSED(params); - ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *)dev->context; - return ggml_backend_cuda_init(ctx->device); -} - -static ggml_backend_buffer_type_t ggml_backend_cuda_device_get_buffer_type(ggml_backend_dev_t dev) { - ggml_backend_cuda_device_context * ctx = (ggml_backend_cuda_device_context *)dev->context; - return ggml_backend_cuda_buffer_type(ctx->device); -} - -static ggml_backend_buffer_type_t ggml_backend_cuda_device_get_host_buffer_type(ggml_backend_dev_t dev) { - GGML_UNUSED(dev); - return ggml_backend_cuda_host_buffer_type(); -} - -// TODO: move these functions here -static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const ggml_tensor * op) { - ggml_backend_cuda_device_context * dev_ctx = (ggml_backend_cuda_device_context *) dev->context; - - // split buffers can only be used with GGML_OP_MUL_MAT - if (op->op != GGML_OP_MUL_MAT) { - for (int i = 0; i < GGML_MAX_SRC; i++) { - if (op->src[i] && op->src[i]->buffer && ggml_backend_buft_is_cuda_split(op->src[i]->buffer->buft)) { - return false; - } - } - } - - // check if all the sources are allocated on this device - for (int i = 0; i < GGML_MAX_SRC; i++) { - if (op->src[i] && op->src[i]->buffer && ggml_backend_buft_is_cuda(op->src[i]->buffer->buft)) { - ggml_backend_cuda_buffer_type_context * buft_ctx = (ggml_backend_cuda_buffer_type_context *)op->src[i]->buffer->buft->context; - if (buft_ctx->device != dev_ctx->device) { - return false; - } - } - } - - switch (op->op) { - case GGML_OP_UNARY: - switch (ggml_get_unary_op(op)) { - case GGML_UNARY_OP_ABS: - case GGML_UNARY_OP_SGN: - case GGML_UNARY_OP_NEG: - case GGML_UNARY_OP_STEP: - case GGML_UNARY_OP_GELU: - case GGML_UNARY_OP_SILU: - case GGML_UNARY_OP_RELU: - case GGML_UNARY_OP_SIGMOID: - case GGML_UNARY_OP_HARDSIGMOID: - case GGML_UNARY_OP_HARDSWISH: - case GGML_UNARY_OP_GELU_ERF: - case GGML_UNARY_OP_GELU_QUICK: - case GGML_UNARY_OP_TANH: - case GGML_UNARY_OP_EXP: - case GGML_UNARY_OP_EXPM1: - case GGML_UNARY_OP_SOFTPLUS: - case GGML_UNARY_OP_ELU: - case GGML_UNARY_OP_XIELU: - case GGML_UNARY_OP_FLOOR: - case GGML_UNARY_OP_CEIL: - case GGML_UNARY_OP_ROUND: - case GGML_UNARY_OP_TRUNC: - // TODO: should become: - //return ggml_is_contiguous_rows(op->src[0]); - return ggml_is_contiguous(op->src[0]); - default: - return false; - } - break; - case GGML_OP_GLU: - switch (ggml_get_glu_op(op)) { - case GGML_GLU_OP_REGLU: - case GGML_GLU_OP_GEGLU: - case GGML_GLU_OP_SWIGLU: - case GGML_GLU_OP_SWIGLU_OAI: - case GGML_GLU_OP_SWIGLU_DS4: - case GGML_GLU_OP_GEGLU_ERF: - case GGML_GLU_OP_GEGLU_QUICK: - return ggml_is_contiguous_1(op->src[0]); - default: - return false; - } - break; - case GGML_OP_TURBO_WHT: - // Only requires dim0 contiguous (nb[0] == sizeof(float)); - // the kernel handles strided dim1/dim2 via separate src/dst strides. - return op->src[0]->nb[0] == ggml_type_size(op->src[0]->type); - case GGML_OP_MOE_FUSED: - return true; - case GGML_OP_MUL_MAT: - case GGML_OP_MUL_MAT_ID: - { - struct ggml_tensor * a = op->src[0]; - struct ggml_tensor * b = op->src[1]; - if (a->buffer && ggml_backend_buft_is_cuda_split(a->buffer->buft)) { - if (a->ne[2] > 1 || a->ne[3] > 1) { - return false; - } - // for small weight matrices the active device can end up without any rows, don't use row split in those cases - // this avoids some edge cases (and the performance would not be good anyways) - ggml_backend_cuda_split_buffer_type_context * buft_ctx = (ggml_backend_cuda_split_buffer_type_context *) a->buffer->buft->context; - int64_t row_low; - int64_t row_high; - get_row_split(&row_low, &row_high, a, buft_ctx->tensor_split, dev_ctx->device); - if (row_low == row_high) { - return false; - } - } - if (b->type == GGML_TYPE_F16 && a->type != GGML_TYPE_F16) { - return false; - } -#ifdef GGML_USE_MUSA - const int cc = ggml_cuda_info().devices[dev_ctx->device].cc; - if (b->ne[2]*b->ne[3] > 1 && !ggml_is_transposed(a) && !ggml_is_transposed(b)) { - if (GGML_CUDA_CC_IS_QY1(cc) && op->op == GGML_OP_MUL_MAT && - a->type == GGML_TYPE_F16 && b->type == GGML_TYPE_F16) { - return false; - } - if (GGML_CUDA_CC_IS_QY2(cc) && op->op == GGML_OP_MUL_MAT_ID && - a->type == GGML_TYPE_Q2_K && b->type == GGML_TYPE_F32) { - return false; - } - } -#endif // GGML_USE_MUSA - switch (a->type) { - case GGML_TYPE_F32: - case GGML_TYPE_F16: - case GGML_TYPE_Q4_0: - case GGML_TYPE_Q4_1: - case GGML_TYPE_Q5_0: - case GGML_TYPE_Q5_1: - case GGML_TYPE_Q8_0: - case GGML_TYPE_MXFP4: - case GGML_TYPE_NVFP4: - case GGML_TYPE_Q4_0_ROCMFP4: - case GGML_TYPE_Q4_0_ROCMFP4_FAST: - case GGML_TYPE_Q2_0_ROCMFP2: - case GGML_TYPE_Q3_0_ROCMFPX: - case GGML_TYPE_Q6_0_ROCMFPX: - case GGML_TYPE_Q8_0_ROCMFPX: - case GGML_TYPE_Q2_K: - case GGML_TYPE_Q3_K: - case GGML_TYPE_Q4_K: - case GGML_TYPE_Q5_K: - case GGML_TYPE_Q6_K: - case GGML_TYPE_Q8_K: - case GGML_TYPE_IQ1_M: - case GGML_TYPE_IQ1_S: - case GGML_TYPE_IQ2_S: - case GGML_TYPE_IQ2_XS: - case GGML_TYPE_IQ2_XXS: - case GGML_TYPE_IQ3_S: - case GGML_TYPE_IQ3_XXS: - case GGML_TYPE_IQ4_NL: - case GGML_TYPE_IQ4_XS: - case GGML_TYPE_BF16: - return true; - default: - return false; - } - } break; - case GGML_OP_OUT_PROD: - return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32; - case GGML_OP_GET_ROWS: - { - switch (op->src[0]->type) { - case GGML_TYPE_F16: - case GGML_TYPE_F32: - case GGML_TYPE_BF16: - case GGML_TYPE_I32: - case GGML_TYPE_Q4_0: - case GGML_TYPE_Q4_1: - case GGML_TYPE_Q5_0: - case GGML_TYPE_Q5_1: - case GGML_TYPE_Q8_0: - case GGML_TYPE_Q4_0_ROCMFP4: - case GGML_TYPE_Q4_0_ROCMFP4_FAST: - case GGML_TYPE_Q2_0_ROCMFP2: - case GGML_TYPE_Q3_0_ROCMFPX: - case GGML_TYPE_Q6_0_ROCMFPX: - case GGML_TYPE_Q8_0_ROCMFPX: - case GGML_TYPE_TQ3_0: - return true; - default: - return false; - } - } break; - case GGML_OP_GET_ROWS_BACK: - { - return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && op->ne[2] == 1 && op->ne[3] == 1; - } break; - case GGML_OP_SET_ROWS: - { - return (op->type == GGML_TYPE_F32 || op->type == GGML_TYPE_F16 || op->type == GGML_TYPE_BF16 || - op->type == GGML_TYPE_Q4_0 || op->type == GGML_TYPE_Q4_1 || op->type == GGML_TYPE_Q5_0 || - op->type == GGML_TYPE_Q5_1 || op->type == GGML_TYPE_Q8_0 || op->type == GGML_TYPE_IQ4_NL || - op->type == GGML_TYPE_TQ3_0) && - op->src[0]->type == GGML_TYPE_F32 && - (op->src[1]->type == GGML_TYPE_I64 || op->src[1]->type == GGML_TYPE_I32); - } break; - case GGML_OP_SET: - { - const ggml_type t = op->type; - return (t == GGML_TYPE_F32 || t == GGML_TYPE_I32) && - t == op->src[0]->type && - t == op->src[1]->type; - } break; - case GGML_OP_CPY: - { - ggml_type src0_type = op->src[0]->type; - ggml_type src1_type = op->src[1]->type; - if ((src0_type == GGML_TYPE_F32 || src0_type == GGML_TYPE_BF16 || src0_type == GGML_TYPE_F16) && - (src1_type == GGML_TYPE_F32 || src1_type == GGML_TYPE_BF16 || src1_type == GGML_TYPE_F16) - ) { - return true; - } - if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_Q8_0) { - return true; - } - if (src0_type == GGML_TYPE_Q8_0 && src1_type == GGML_TYPE_F32) { - return true; - } - if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_Q4_0) { - return true; - } - if (src0_type == GGML_TYPE_Q4_0 && src1_type == GGML_TYPE_F32) { - return true; - } - if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_Q4_1) { - return true; - } - if (src0_type == GGML_TYPE_Q4_1 && src1_type == GGML_TYPE_F32) { - return true; - } - if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_Q5_0) { - return true; - } - if (src0_type == GGML_TYPE_Q5_0 && src1_type == GGML_TYPE_F32) { - return true; - } - if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_Q5_1) { - return true; - } - if (src0_type == GGML_TYPE_Q5_1 && src1_type == GGML_TYPE_F32) { - return true; - } - if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_IQ4_NL) { - return true; - } - if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_TQ3_0) { - return true; - } - if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_I32) { - return true; - } - if (src0_type == GGML_TYPE_I32 && src1_type == GGML_TYPE_F32) { - return true; - } - if (src0_type == GGML_TYPE_I32 && src1_type == GGML_TYPE_I32) { - return true; - } - if (src0_type == src1_type && ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1])) { - return true; - } - return false; - } break; - case GGML_OP_DUP: - { - ggml_type src0_type = op->src[0]->type; - return src0_type != GGML_TYPE_I32 && src0_type != GGML_TYPE_I16; - } break; - case GGML_OP_ARGMAX: - case GGML_OP_COUNT_EQUAL: - { - return true; - } break; - case GGML_OP_REPEAT: - { - // the CUDA REPEAT path only implements F32/F16; other types assert at runtime - ggml_type src0_type = op->src[0]->type; - return src0_type == GGML_TYPE_F32 || src0_type == GGML_TYPE_F16; - } break; - case GGML_OP_REPEAT_BACK: - return op->type == GGML_TYPE_F32 && (op->src[0]->ne[2]*op->src[0]->ne[3]) <= (1 << 15); - case GGML_OP_CONCAT: - { - ggml_type src0_type = op->src[0]->type; - return src0_type == GGML_TYPE_F32 || - src0_type == GGML_TYPE_F16 || - src0_type == GGML_TYPE_BF16 || - src0_type == GGML_TYPE_I8; - } break; - case GGML_OP_CONV_TRANSPOSE_1D: - { - ggml_type src0_type = op->src[0]->type; - ggml_type src1_type = op->src[1]->type; - if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_F32) { - return true; - } - return false; - } break; - case GGML_OP_SILU_BACK: - return ggml_is_contiguous(op->src[0]) && op->src[0]->type == GGML_TYPE_F32; - break; - case GGML_OP_NORM: - case GGML_OP_RMS_NORM: - case GGML_OP_L2_NORM: - return true; - case GGML_OP_RMS_NORM_BACK: - return ggml_is_contiguous(op->src[0]); - break; - case GGML_OP_NONE: - case GGML_OP_RESHAPE: - case GGML_OP_VIEW: - case GGML_OP_PERMUTE: - case GGML_OP_TRANSPOSE: - case GGML_OP_ADD: - case GGML_OP_ADD_ID: - case GGML_OP_ADD1: - case GGML_OP_SUB: - case GGML_OP_MUL: - case GGML_OP_DIV: - case GGML_OP_SCALE: - case GGML_OP_SQR: - case GGML_OP_SQRT: - case GGML_OP_SIN: - case GGML_OP_COS: - case GGML_OP_CLAMP: - case GGML_OP_LOG: - return true; - case GGML_OP_SSM_SCAN: { - if (op->src[3]->ne[0] == 1) { - // Mamba2 - // (kernel only supports (d_state == 128 || d_state == 256) && d_head % 16 == 0) - return (op->src[0]->ne[0] == 128 || op->src[0]->ne[0] == 256) && op->src[0]->ne[1] % 16 == 0; - } else { - // Mamba - // (kernel only supports d_state == 16, d_head == 1, n_head % 128 == 0, n_group == 1) - return op->src[0]->ne[0] == 16 && op->src[0]->ne[1] == 1 && op->src[0]->ne[2] % 128 == 0 && op->src[4]->ne[1] == 1; - } - } - case GGML_OP_SSM_CONV: { - // assumes d_inner % threads == 0 - return op->src[0]->ne[1] % 128 == 0; - } - case GGML_OP_CONT: - return true; - case GGML_OP_DIAG_MASK_INF: - return true; - case GGML_OP_SOFT_MAX: - return true; - case GGML_OP_SOFT_MAX_BACK: { - float max_bias = 0.0f; - memcpy(&max_bias, (const float *) op->op_params + 1, sizeof(float)); - return max_bias == 0.0f; - } - case GGML_OP_ROLL: - if(op->src[0]->type == GGML_TYPE_F32) { - return true; - } - return false; - case GGML_OP_ROPE: - case GGML_OP_ROPE_BACK: { - return op->src[0]->nb[0] == ggml_type_size(op->src[0]->type) && ggml_is_contiguous_2(op->src[0]); - } - case GGML_OP_IM2COL: - case GGML_OP_IM2COL_3D: - case GGML_OP_CONV_2D: - case GGML_OP_CONV_2D_DW: - case GGML_OP_CONV_TRANSPOSE_2D: - case GGML_OP_POOL_2D: - return true; - case GGML_OP_ACC: - // TODO: extend support like so: - //return ggml_is_contiguous_rows(op->src[0]) && ggml_is_contiguous_rows(op->src[1]); - return ggml_is_contiguous(op->src[0]) && ggml_is_contiguous(op->src[1]); - case GGML_OP_SUM: - return ggml_is_contiguous_rows(op->src[0]); - case GGML_OP_TOP_K: - case GGML_OP_ARGSORT: -#ifndef GGML_CUDA_USE_CUB - return op->src[0]->ne[0] <= 1024; -#else - return true; -#endif - case GGML_OP_SUM_ROWS: - case GGML_OP_MEAN: - case GGML_OP_GROUP_NORM: - return ggml_is_contiguous(op->src[0]); - case GGML_OP_PAD: - return true; - case GGML_OP_UPSCALE: - case GGML_OP_PAD_REFLECT_1D: - case GGML_OP_ARANGE: - case GGML_OP_TIMESTEP_EMBEDDING: - case GGML_OP_LEAKY_RELU: - case GGML_OP_RWKV_WKV6: - case GGML_OP_GATED_LINEAR_ATTN: - case GGML_OP_RWKV_WKV7: - return true; - case GGML_OP_GATED_DELTA_NET: - //TODO: enable once MUSA compiler is solved https://github.com/ggml-org/llama.cpp/pull/19504#issuecomment-4018634327 -#ifdef GGML_USE_MUSA - return false; -#else - return true; -#endif // GGML_USE_MUSA - case GGML_OP_FLASH_ATTN_EXT: - return ggml_cuda_flash_attn_ext_supported(dev_ctx->device, op); - case GGML_OP_FLASH_ATTN_SPARSE: - return true; // Always supported on CUDA - case GGML_OP_CROSS_ENTROPY_LOSS: - case GGML_OP_CROSS_ENTROPY_LOSS_BACK: - case GGML_OP_OPT_STEP_ADAMW: - case GGML_OP_OPT_STEP_SGD: - case GGML_OP_FILL: - case GGML_OP_CUMSUM: - case GGML_OP_TRI: - case GGML_OP_DIAG: - case GGML_OP_SOLVE_TRI: - return true; - - default: - return false; - } -} - -static bool ggml_backend_cuda_device_supports_buft(ggml_backend_dev_t dev, ggml_backend_buffer_type_t buft) { - ggml_backend_cuda_device_context * dev_ctx = (ggml_backend_cuda_device_context *) dev->context; - const bool integrated = ggml_cuda_info().devices[dev_ctx->device].integrated; - return (((ggml_backend_buft_is_cuda(buft) || ggml_backend_buft_is_cuda_split(buft)) && buft->device == dev) || (integrated && ggml_backend_buft_is_cuda_host(buft))); -} - -static int64_t get_op_batch_size(const ggml_tensor * op) { - switch (op->op) { - case GGML_OP_GET_ROWS: - return 0; - case GGML_OP_MUL_MAT: - return op->ne[1]; - case GGML_OP_MUL_MAT_ID: - case GGML_OP_ROPE: - case GGML_OP_ROPE_BACK: - return op->ne[2]; - default: - return ggml_nrows(op); - } -} - -static bool ggml_backend_cuda_device_offload_op(ggml_backend_dev_t dev, const ggml_tensor * op) { - ggml_backend_cuda_device_context * dev_ctx = (ggml_backend_cuda_device_context *) dev->context; - - return get_op_batch_size(op) >= dev_ctx->op_offload_min_batch_size; -} - -static ggml_backend_event_t ggml_backend_cuda_device_event_new(ggml_backend_dev_t dev) { -#ifdef GGML_CUDA_NO_PEER_COPY - return nullptr; -#else - ggml_backend_cuda_device_context * dev_ctx = (ggml_backend_cuda_device_context *)dev->context; - - ggml_cuda_set_device(dev_ctx->device); - - cudaEvent_t event; - CUDA_CHECK(cudaEventCreateWithFlags(&event, cudaEventDisableTiming)); - - return new ggml_backend_event { - /* .device = */ dev, - /* .context = */ event, - }; -#endif -} - -static void ggml_backend_cuda_device_event_free(ggml_backend_dev_t dev, ggml_backend_event_t event) { - GGML_UNUSED(dev); - - CUDA_CHECK(cudaEventDestroy((cudaEvent_t)event->context)); - delete event; -} - -static void ggml_backend_cuda_device_event_synchronize(ggml_backend_dev_t dev, ggml_backend_event_t event) { - GGML_UNUSED(dev); - CUDA_CHECK(cudaEventSynchronize((cudaEvent_t)event->context)); -} - -static const ggml_backend_device_i ggml_backend_cuda_device_interface = { - /* .get_name = */ ggml_backend_cuda_device_get_name, - /* .get_description = */ ggml_backend_cuda_device_get_description, - /* .get_memory = */ ggml_backend_cuda_device_get_memory, - /* .get_type = */ ggml_backend_cuda_device_get_type, - /* .get_props = */ ggml_backend_cuda_device_get_props, - /* .init_backend = */ ggml_backend_cuda_device_init_backend, - /* .get_buffer_type = */ ggml_backend_cuda_device_get_buffer_type, - /* .get_host_buffer_type = */ ggml_backend_cuda_device_get_host_buffer_type, - /* .buffer_from_host_ptr = */ NULL, - /* .supports_op = */ ggml_backend_cuda_device_supports_op, - /* .supports_buft = */ ggml_backend_cuda_device_supports_buft, - /* .offload_op = */ ggml_backend_cuda_device_offload_op, - /* .event_new = */ ggml_backend_cuda_device_event_new, - /* .event_free = */ ggml_backend_cuda_device_event_free, - /* .event_synchronize = */ ggml_backend_cuda_device_event_synchronize, -}; - -// backend reg - -struct ggml_backend_cuda_reg_context { - std::vector devices; -}; - -static const char * ggml_backend_cuda_reg_get_name(ggml_backend_reg_t reg) { - GGML_UNUSED(reg); - return GGML_CUDA_NAME; -} - -static size_t ggml_backend_cuda_reg_get_device_count(ggml_backend_reg_t reg) { - ggml_backend_cuda_reg_context * ctx = (ggml_backend_cuda_reg_context *)reg->context; - return ctx->devices.size(); -} - -static ggml_backend_dev_t ggml_backend_cuda_reg_get_device(ggml_backend_reg_t reg, size_t index) { - ggml_backend_cuda_reg_context * ctx = (ggml_backend_cuda_reg_context *)reg->context; - GGML_ASSERT(index < ctx->devices.size()); - return ctx->devices[index]; -} - -static ggml_backend_feature * ggml_backend_cuda_get_features(ggml_backend_reg_t reg) { - static std::vector features = []() { - std::vector features; - #define _STRINGIFY(...) #__VA_ARGS__ - #define STRINGIFY(...) _STRINGIFY(__VA_ARGS__) - - #ifdef __CUDA_ARCH_LIST__ - features.push_back({ "ARCHS", STRINGIFY(__CUDA_ARCH_LIST__) }); - #endif - - #ifdef GGML_CUDA_FORCE_MMQ - features.push_back({ "FORCE_MMQ", "1" }); - #endif - - #ifdef GGML_CUDA_FORCE_CUBLAS - features.push_back({ "FORCE_CUBLAS", "1" }); - #endif - - #ifndef GGML_USE_VMM - features.push_back({ "NO_VMM", "1" }); - #endif - - #ifdef GGML_CUDA_NO_PEER_COPY - features.push_back({ "NO_PEER_COPY", "1" }); - #endif - - #ifdef GGML_CUDA_USE_GRAPHS - features.push_back({ "USE_GRAPHS", "1" }); - #endif - - #ifdef GGML_CUDA_PEER_MAX_BATCH_SIZE - features.push_back({ "PEER_MAX_BATCH_SIZE", STRINGIFY(GGML_CUDA_PEER_MAX_BATCH_SIZE) }); - #endif - - #ifdef GGML_CUDA_FA_ALL_QUANTS - features.push_back({ "FA_ALL_QUANTS", "1" }); - #endif - - { - const auto & info = ggml_cuda_info(); - for (int id = 0; id < info.device_count; ++id) { - if (blackwell_mma_available(info.devices[id].cc)) { - features.push_back({ "BLACKWELL_NATIVE_FP4", "1"}); - break; - } - } - } - - #undef _STRINGIFY - #undef STRINGIFY - - features.push_back({ nullptr, nullptr }); - - return features; - }(); - - return features.data(); - - GGML_UNUSED(reg); -} - -static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, const char * name) { - GGML_UNUSED(reg); - if (strcmp(name, "ggml_backend_allreduce_tensor") == 0) { - return (void *)ggml_backend_cuda_allreduce_tensor; - } - if (strcmp(name, "ggml_backend_split_buffer_type") == 0) { - return (void *)ggml_backend_cuda_split_buffer_type; - } - if (strcmp(name, "ggml_backend_register_host_buffer") == 0) { - return (void *)ggml_backend_cuda_register_host_buffer; - } - if (strcmp(name, "ggml_backend_unregister_host_buffer") == 0) { - return (void *)ggml_backend_cuda_unregister_host_buffer; - } - if (strcmp(name, "ggml_backend_get_features") == 0) { - return (void *)ggml_backend_cuda_get_features; - } - return nullptr; -} - -static const ggml_backend_reg_i ggml_backend_cuda_reg_interface = { - /* .get_name = */ ggml_backend_cuda_reg_get_name, - /* .get_device_count = */ ggml_backend_cuda_reg_get_device_count, - /* .get_device = */ ggml_backend_cuda_reg_get_device, - /* .get_proc_address = */ ggml_backend_cuda_reg_get_proc_address, -}; - -// backend registry -ggml_backend_reg_t ggml_backend_cuda_reg() { - static ggml_backend_reg reg; - static bool initialized = false; - - { - static std::mutex mutex; - std::lock_guard lock(mutex); - if (!initialized) { - ggml_backend_cuda_reg_context * ctx = new ggml_backend_cuda_reg_context; - const int min_batch_size = getenv("GGML_OP_OFFLOAD_MIN_BATCH") ? atoi(getenv("GGML_OP_OFFLOAD_MIN_BATCH")) : 32; - - for (int i = 0; i < ggml_cuda_info().device_count; i++) { - ggml_backend_cuda_device_context * dev_ctx = new ggml_backend_cuda_device_context; - dev_ctx->device = i; - dev_ctx->name = GGML_CUDA_NAME + std::to_string(i); - - cudaDeviceProp prop; - CUDA_CHECK(cudaGetDeviceProperties(&prop, i)); - dev_ctx->description = prop.name; - - char pci_bus_id[16] = {}; - snprintf(pci_bus_id, sizeof(pci_bus_id), "%04x:%02x:%02x.0", prop.pciDomainID, prop.pciBusID, prop.pciDeviceID); - dev_ctx->pci_bus_id = pci_bus_id; - dev_ctx->op_offload_min_batch_size = min_batch_size; - - ggml_backend_dev_t dev = new ggml_backend_device { - /* .iface = */ ggml_backend_cuda_device_interface, - /* .reg = */ ®, - /* .context = */ dev_ctx - }; - ctx->devices.push_back(dev); - } - - reg = ggml_backend_reg { - /* .api_version = */ GGML_BACKEND_API_VERSION, - /* .iface = */ ggml_backend_cuda_reg_interface, - /* .context = */ ctx - }; - } - - initialized = true; - } - - return ® -} - -ggml_backend_t ggml_backend_cuda_init(int device) { - if (device < 0 || device >= ggml_backend_cuda_get_device_count()) { - GGML_LOG_ERROR("%s: invalid device %d\n", __func__, device); - return nullptr; - } - - ggml_backend_cuda_context * ctx = new ggml_backend_cuda_context(device); - if (ctx == nullptr) { - GGML_LOG_ERROR("%s: failed to allocate context\n", __func__); - return nullptr; - } - - ggml_backend_t cuda_backend = new ggml_backend { - /* .guid = */ ggml_backend_cuda_guid(), - /* .iface = */ ggml_backend_cuda_interface, - /* .device = */ ggml_backend_reg_dev_get(ggml_backend_cuda_reg(), device), - /* .context = */ ctx, - }; - - return cuda_backend; -} - -GGML_BACKEND_DL_IMPL(ggml_backend_cuda_reg) diff --git a/ggml/src/ggml.c.ds4bak b/ggml/src/ggml.c.ds4bak deleted file mode 100644 index 3b5c6bcafa23..000000000000 --- a/ggml/src/ggml.c.ds4bak +++ /dev/null @@ -1,8040 +0,0 @@ -#define _CRT_SECURE_NO_DEPRECATE // Disables "unsafe" warnings on Windows -#define _USE_MATH_DEFINES // For M_PI on MSVC - -#include "ggml-backend.h" -#include "ggml-impl.h" -#include "ggml-threading.h" -#include "ggml-cpu.h" -#include "ggml.h" - -// FIXME: required here for quantization functions -#include "ggml-quants.h" -#include "../rocmfp4/rocmfp4.h" -#include "../rocmfpx/rocmfpx.h" - -#ifdef GGML_USE_CPU_HBM -#include -#endif - -#if defined(_MSC_VER) || defined(__MINGW32__) -#include // using malloc.h with MSC/MINGW -#elif !defined(__FreeBSD__) && !defined(__NetBSD__) && !defined(__OpenBSD__) -#include -#endif - -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#include -#if defined(__gnu_linux__) -#include -#endif - -#if defined(__APPLE__) -#include -#include -#include -#endif - -#if defined(_WIN32) -#define WIN32_LEAN_AND_MEAN -#ifndef NOMINMAX - #define NOMINMAX -#endif -#include -#endif - -#define UNUSED GGML_UNUSED - -// Needed for ggml_fp32_to_bf16_row() -#if defined(__AVX512BF16__) -#if defined(_MSC_VER) -#define m512i(p) p -#else -#include -#define m512i(p) (__m512i)(p) -#endif // defined(_MSC_VER) -#endif // defined(__AVX512BF16__) - -#if defined(__linux__) || \ - defined(__FreeBSD__) || defined(__NetBSD__) || defined(__OpenBSD__) || \ - (defined(__APPLE__) && !TARGET_OS_TV && !TARGET_OS_WATCH) - -#include -#include -#include -#include -#if defined(__linux__) -#include -#endif - -#if defined(__ANDROID__) -#include -#include -#include - -struct backtrace_state { - void ** current; - void ** end; -}; - -static _Unwind_Reason_Code unwind_callback(struct _Unwind_Context* context, void* arg) { - struct backtrace_state * state = (struct backtrace_state *)arg; - uintptr_t pc = _Unwind_GetIP(context); - if (pc) { - if (state->current == state->end) { - return _URC_END_OF_STACK; - } else { - *state->current++ = (void*)pc; - } - } - return _URC_NO_REASON; -} - -static void ggml_print_backtrace_symbols(void) { - const int max = 100; - void* buffer[max]; - - struct backtrace_state state = {buffer, buffer + max}; - _Unwind_Backtrace(unwind_callback, &state); - - int count = state.current - buffer; - - for (int idx = 0; idx < count; ++idx) { - const void * addr = buffer[idx]; - const char * symbol = ""; - - Dl_info info; - if (dladdr(addr, &info) && info.dli_sname) { - symbol = info.dli_sname; - } - - fprintf(stderr, "%d: %p %s\n", idx, addr, symbol); - } -} -#elif defined(__linux__) && defined(__GLIBC__) -#include -static void ggml_print_backtrace_symbols(void) { - void * trace[100]; - int nptrs = backtrace(trace, sizeof(trace)/sizeof(trace[0])); - backtrace_symbols_fd(trace, nptrs, STDERR_FILENO); -} -#elif defined(__APPLE__) -#include -static void ggml_print_backtrace_symbols(void) { - void * trace[100]; - int nptrs = backtrace(trace, sizeof(trace)/sizeof(trace[0])); - backtrace_symbols_fd(trace, nptrs, STDERR_FILENO); -} -#else -static void ggml_print_backtrace_symbols(void) { - // platform not supported -} -#endif - -void ggml_print_backtrace(void) { - const char * GGML_NO_BACKTRACE = getenv("GGML_NO_BACKTRACE"); - if (GGML_NO_BACKTRACE) { - return; - } -#if defined(__APPLE__) - // On macOS, fork+debugger attachment is problematic due to: - // 1. libdispatch "poisons" forked child processes - // 2. lldb has issues attaching to parent from forked child - // Use simple backtrace() instead to avoid Terminal.app crashes - const char * GGML_BACKTRACE_LLDB = getenv("GGML_BACKTRACE_LLDB"); - if (!GGML_BACKTRACE_LLDB) { - fprintf(stderr, "WARNING: Using native backtrace. Set GGML_BACKTRACE_LLDB for more info.\n"); - fprintf(stderr, "WARNING: GGML_BACKTRACE_LLDB may cause native MacOS Terminal.app to crash.\n"); - fprintf(stderr, "See: https://github.com/ggml-org/llama.cpp/pull/17869\n"); - ggml_print_backtrace_symbols(); - return; - } -#endif -#if defined(__linux__) - FILE * f = fopen("/proc/self/status", "r"); - size_t size = 0; - char * line = NULL; - ssize_t length = 0; - while ((length = getline(&line, &size, f)) > 0) { - if (!strncmp(line, "TracerPid:", sizeof("TracerPid:") - 1) && - (length != sizeof("TracerPid:\t0\n") - 1 || line[length - 2] != '0')) { - // Already being debugged, and the breakpoint is the later abort() - free(line); - fclose(f); - return; - } - } - free(line); - fclose(f); - int lock[2] = { -1, -1 }; - (void) !pipe(lock); // Don't start gdb until after PR_SET_PTRACER -#endif - const int parent_pid = getpid(); - const int child_pid = fork(); - if (child_pid < 0) { // error -#if defined(__linux__) - close(lock[1]); - close(lock[0]); -#endif - return; - } else if (child_pid == 0) { // child - char attach[32]; - snprintf(attach, sizeof(attach), "attach %d", parent_pid); -#if defined(__linux__) - close(lock[1]); - (void) !read(lock[0], lock, 1); - close(lock[0]); -#endif - // try gdb - execlp("gdb", "gdb", "--batch", - "-ex", "set style enabled on", - "-ex", attach, - "-ex", "bt -frame-info source-and-location", - "-ex", "detach", - "-ex", "quit", - (char *) NULL); - // try lldb - execlp("lldb", "lldb", "--batch", - "-o", "bt", - "-o", "quit", - "-p", &attach[sizeof("attach ") - 1], - (char *) NULL); - // gdb failed, fallback to backtrace_symbols - ggml_print_backtrace_symbols(); - _Exit(0); - } else { // parent -#if defined(__linux__) - prctl(PR_SET_PTRACER, child_pid); - close(lock[1]); - close(lock[0]); -#endif - waitpid(child_pid, NULL, 0); - } -} -#else -void ggml_print_backtrace(void) { - // platform not supported -} -#endif - -static ggml_abort_callback_t g_abort_callback = NULL; - -// Set the abort callback (passing null will restore original abort functionality: printing a message to stdout) -GGML_API ggml_abort_callback_t ggml_set_abort_callback(ggml_abort_callback_t callback) { - ggml_abort_callback_t ret_val = g_abort_callback; - g_abort_callback = callback; - return ret_val; -} - -void ggml_abort(const char * file, int line, const char * fmt, ...) { - fflush(stdout); - - char message[2048]; - int offset = snprintf(message, sizeof(message), "%s:%d: ", file, line); - - va_list args; - va_start(args, fmt); - vsnprintf(message + offset, sizeof(message) - offset, fmt, args); - va_end(args); - - if (g_abort_callback) { - g_abort_callback(message); - } else { - // default: print error and backtrace to stderr - fprintf(stderr, "%s\n", message); - ggml_print_backtrace(); - } - - abort(); -} - -// ggml_print_backtrace is registered with std::set_terminate by ggml.cpp - -// -// logging -// - -struct ggml_logger_state { - ggml_log_callback log_callback; - void * log_callback_user_data; -}; -static struct ggml_logger_state g_logger_state = {ggml_log_callback_default, NULL}; - -static void ggml_log_internal_v(enum ggml_log_level level, const char * format, va_list args) { - if (format == NULL) { - return; - } - va_list args_copy; - va_copy(args_copy, args); - char buffer[128]; - int len = vsnprintf(buffer, 128, format, args); - if (len < 128) { - g_logger_state.log_callback(level, buffer, g_logger_state.log_callback_user_data); - } else { - char * buffer2 = (char *) calloc(len + 1, sizeof(char)); - vsnprintf(buffer2, len + 1, format, args_copy); - buffer2[len] = 0; - g_logger_state.log_callback(level, buffer2, g_logger_state.log_callback_user_data); - free(buffer2); - } - va_end(args_copy); -} - -void ggml_log_internal(enum ggml_log_level level, const char * format, ...) { - va_list args; - va_start(args, format); - ggml_log_internal_v(level, format, args); - va_end(args); -} - -void ggml_log_callback_default(enum ggml_log_level level, const char * text, void * user_data) { - (void) level; - (void) user_data; - fputs(text, stderr); - fflush(stderr); -} - -// -// end of logging block -// - -#ifdef GGML_USE_ACCELERATE -// uncomment to use vDSP for soft max computation -// note: not sure if it is actually faster -//#define GGML_SOFT_MAX_ACCELERATE -#endif - - -void * ggml_aligned_malloc(size_t size) { -#if defined(__s390x__) - const int alignment = 256; -#else - const int alignment = 64; -#endif - -#if defined(_MSC_VER) || defined(__MINGW32__) - return _aligned_malloc(size, alignment); -#else - if (size == 0) { - GGML_LOG_WARN("Behavior may be unexpected when allocating 0 bytes for ggml_aligned_malloc!\n"); - return NULL; - } - void * aligned_memory = NULL; - #ifdef GGML_USE_CPU_HBM - int result = hbw_posix_memalign(&aligned_memory, alignment, size); - #elif TARGET_OS_OSX - GGML_UNUSED(alignment); - kern_return_t alloc_status = vm_allocate((vm_map_t) mach_task_self(), (vm_address_t *) &aligned_memory, size, VM_FLAGS_ANYWHERE); - int result = EFAULT; - switch (alloc_status) { - case KERN_SUCCESS: - result = 0; - break; - case KERN_INVALID_ADDRESS: - result = EINVAL; - break; - case KERN_NO_SPACE: - result = ENOMEM; - break; - default: - result = EFAULT; - break; - } - #else - int result = posix_memalign(&aligned_memory, alignment, size); - #endif - if (result != 0) { - // Handle allocation failure - const char *error_desc = "unknown allocation error"; - switch (result) { - case EINVAL: - error_desc = "invalid alignment value"; - break; - case ENOMEM: - error_desc = "insufficient memory"; - break; - } - GGML_LOG_ERROR("%s: %s (attempted to allocate %6.2f MB)\n", __func__, error_desc, size/(1024.0*1024.0)); - return NULL; - } - return aligned_memory; -#endif -} - -void ggml_aligned_free(void * ptr, size_t size) { - GGML_UNUSED(size); -#if defined(_MSC_VER) || defined(__MINGW32__) - _aligned_free(ptr); -#elif GGML_USE_CPU_HBM - if (ptr != NULL) { - hbw_free(ptr); - } -#elif TARGET_OS_OSX - if (ptr != NULL) { - vm_deallocate((vm_map_t)mach_task_self(), (vm_address_t)ptr, size); - } -#else - free(ptr); -#endif -} - - -inline static void * ggml_malloc(size_t size) { - if (size == 0) { - GGML_LOG_WARN("Behavior may be unexpected when allocating 0 bytes for ggml_malloc!\n"); - return NULL; - } - void * result = malloc(size); - if (result == NULL) { - GGML_LOG_ERROR("%s: failed to allocate %6.2f MB\n", __func__, size/(1024.0*1024.0)); - GGML_ABORT("fatal error"); - } - return result; -} - -// calloc -inline static void * ggml_calloc(size_t num, size_t size) { - if (num == 0 || size == 0) { - GGML_LOG_WARN("Behavior may be unexpected when allocating 0 bytes for ggml_calloc!\n"); - return NULL; - } - void * result = calloc(num, size); - if (result == NULL) { - GGML_LOG_ERROR("%s: failed to allocate %6.2f MB\n", __func__, size/(1024.0*1024.0)); - GGML_ABORT("fatal error"); - } - return result; -} - -#define GGML_MALLOC(size) ggml_malloc(size) -#define GGML_CALLOC(num, size) ggml_calloc(num, size) - -#define GGML_FREE(ptr) free(ptr) - -const char * ggml_status_to_string(enum ggml_status status) { - switch (status) { - case GGML_STATUS_ALLOC_FAILED: return "GGML status: error (failed to allocate memory)"; - case GGML_STATUS_FAILED: return "GGML status: error (operation failed)"; - case GGML_STATUS_SUCCESS: return "GGML status: success"; - case GGML_STATUS_ABORTED: return "GGML status: warning (operation aborted)"; - } - - return "GGML status: unknown"; -} - -float ggml_fp16_to_fp32(ggml_fp16_t x) { -#define ggml_fp16_to_fp32 do_not_use__ggml_fp16_to_fp32__in_ggml - return GGML_FP16_TO_FP32(x); -} - -ggml_fp16_t ggml_fp32_to_fp16(float x) { -#define ggml_fp32_to_fp16 do_not_use__ggml_fp32_to_fp16__in_ggml - return GGML_FP32_TO_FP16(x); -} - -float ggml_bf16_to_fp32(ggml_bf16_t x) { -#define ggml_bf16_to_fp32 do_not_use__ggml_bf16_to_fp32__in_ggml - return GGML_BF16_TO_FP32(x); // it just left shifts -} - -ggml_bf16_t ggml_fp32_to_bf16(float x) { -#define ggml_fp32_to_bf16 do_not_use__ggml_fp32_to_bf16__in_ggml - return GGML_FP32_TO_BF16(x); -} - -void ggml_fp16_to_fp32_row(const ggml_fp16_t * x, float * y, int64_t n) { - for (int64_t i = 0; i < n; i++) { - y[i] = GGML_FP16_TO_FP32(x[i]); - } -} - -void ggml_fp32_to_fp16_row(const float * x, ggml_fp16_t * y, int64_t n) { - int i = 0; - for (; i < n; ++i) { - y[i] = GGML_FP32_TO_FP16(x[i]); - } -} - -void ggml_bf16_to_fp32_row(const ggml_bf16_t * x, float * y, int64_t n) { - int i = 0; - for (; i < n; ++i) { - y[i] = GGML_BF16_TO_FP32(x[i]); - } -} - -void ggml_fp32_to_bf16_row_ref(const float * x, ggml_bf16_t * y, int64_t n) { - for (int i = 0; i < n; i++) { - y[i] = ggml_compute_fp32_to_bf16(x[i]); - } -} - -void ggml_fp32_to_bf16_row(const float * x, ggml_bf16_t * y, int64_t n) { - int i = 0; -#if defined(__AVX512BF16__) - // subnormals are flushed to zero on this platform - for (; i + 32 <= n; i += 32) { - _mm512_storeu_si512( - (__m512i *)(y + i), - m512i(_mm512_cvtne2ps_pbh(_mm512_loadu_ps(x + i + 16), - _mm512_loadu_ps(x + i)))); - } -#endif - for (; i < n; i++) { - y[i] = GGML_FP32_TO_BF16(x[i]); - } -} - -bool ggml_guid_matches(ggml_guid_t guid_a, ggml_guid_t guid_b) { - return memcmp(guid_a, guid_b, sizeof(ggml_guid)) == 0; -} - -const char * ggml_version(void) { - return GGML_VERSION; -} - -const char * ggml_commit(void) { - return GGML_COMMIT; -} - -// -// timing -// - -#if defined(_MSC_VER) || defined(__MINGW32__) -static int64_t timer_freq, timer_start; -void ggml_time_init(void) { - LARGE_INTEGER t; - QueryPerformanceFrequency(&t); - timer_freq = t.QuadPart; - - // The multiplication by 1000 or 1000000 below can cause an overflow if timer_freq - // and the uptime is high enough. - // We subtract the program start time to reduce the likelihood of that happening. - QueryPerformanceCounter(&t); - timer_start = t.QuadPart; -} -int64_t ggml_time_ms(void) { - LARGE_INTEGER t; - QueryPerformanceCounter(&t); - return ((t.QuadPart-timer_start) * 1000) / timer_freq; -} -int64_t ggml_time_us(void) { - LARGE_INTEGER t; - QueryPerformanceCounter(&t); - return ((t.QuadPart-timer_start) * 1000000) / timer_freq; -} -#else -void ggml_time_init(void) {} -int64_t ggml_time_ms(void) { - struct timespec ts; - clock_gettime(CLOCK_MONOTONIC, &ts); - return (int64_t)ts.tv_sec*1000 + (int64_t)ts.tv_nsec/1000000; -} - -int64_t ggml_time_us(void) { - struct timespec ts; - clock_gettime(CLOCK_MONOTONIC, &ts); - return (int64_t)ts.tv_sec*1000000 + (int64_t)ts.tv_nsec/1000; -} -#endif - -int64_t ggml_cycles(void) { - return clock(); -} - -int64_t ggml_cycles_per_ms(void) { - return CLOCKS_PER_SEC/1000; -} - -// -// cross-platform UTF-8 file paths -// - -#ifdef _WIN32 -static wchar_t * ggml_mbstowcs(const char * mbs) { - int wlen = MultiByteToWideChar(CP_UTF8, 0, mbs, -1, NULL, 0); - if (!wlen) { - errno = EINVAL; - return NULL; - } - - wchar_t * wbuf = GGML_MALLOC(wlen * sizeof(wchar_t)); - wlen = MultiByteToWideChar(CP_UTF8, 0, mbs, -1, wbuf, wlen); - if (!wlen) { - GGML_FREE(wbuf); - errno = EINVAL; - return NULL; - } - - return wbuf; -} -#endif - -FILE * ggml_fopen(const char * fname, const char * mode) { -#ifdef _WIN32 - FILE * file = NULL; - - // convert fname (UTF-8) - wchar_t * wfname = ggml_mbstowcs(fname); - if (wfname) { - // convert mode (ANSI) - wchar_t * wmode = GGML_MALLOC((strlen(mode) + 1) * sizeof(wchar_t)); - wchar_t * wmode_p = wmode; - do { - *wmode_p++ = (wchar_t)*mode; - } while (*mode++); - - // open file - file = _wfopen(wfname, wmode); - - GGML_FREE(wfname); - GGML_FREE(wmode); - } - - return file; -#else - return fopen(fname, mode); -#endif - -} - -static const struct ggml_type_traits type_traits[GGML_TYPE_COUNT] = { - [GGML_TYPE_I8] = { - .type_name = "i8", - .blck_size = 1, - .type_size = sizeof(int8_t), - .is_quantized = false, - }, - [GGML_TYPE_I16] = { - .type_name = "i16", - .blck_size = 1, - .type_size = sizeof(int16_t), - .is_quantized = false, - }, - [GGML_TYPE_I32] = { - .type_name = "i32", - .blck_size = 1, - .type_size = sizeof(int32_t), - .is_quantized = false, - }, - [GGML_TYPE_I64] = { - .type_name = "i64", - .blck_size = 1, - .type_size = sizeof(int64_t), - .is_quantized = false, - }, - [GGML_TYPE_F64] = { - .type_name = "f64", - .blck_size = 1, - .type_size = sizeof(double), - .is_quantized = false, - }, - [GGML_TYPE_F32] = { - .type_name = "f32", - .blck_size = 1, - .type_size = sizeof(float), - .is_quantized = false, - }, - [GGML_TYPE_F16] = { - .type_name = "f16", - .blck_size = 1, - .type_size = sizeof(ggml_fp16_t), - .is_quantized = false, - .to_float = (ggml_to_float_t) ggml_fp16_to_fp32_row, - .from_float_ref = (ggml_from_float_t) ggml_fp32_to_fp16_row, - }, - [GGML_TYPE_Q1_0] = { - .type_name = "q1_0", - .blck_size = QK1_0, - .type_size = sizeof(block_q1_0), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_q1_0, - .from_float_ref = (ggml_from_float_t) quantize_row_q1_0_ref, - }, - [GGML_TYPE_TQ3_0] = { - .type_name = "tq3_0", - .blck_size = QK_TQ3_0, - .type_size = sizeof(block_tq3_0), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_tq3_0, - .from_float_ref = (ggml_from_float_t) quantize_row_tq3_0_ref, - }, - [GGML_TYPE_Q4_0] = { - .type_name = "q4_0", - .blck_size = QK4_0, - .type_size = sizeof(block_q4_0), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_q4_0, - .from_float_ref = (ggml_from_float_t) quantize_row_q4_0_ref, - }, - [GGML_TYPE_Q4_1] = { - .type_name = "q4_1", - .blck_size = QK4_1, - .type_size = sizeof(block_q4_1), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_q4_1, - .from_float_ref = (ggml_from_float_t) quantize_row_q4_1_ref, - }, - [4] = { // GGML_TYPE_Q4_2 - .type_name = "DEPRECATED", - .blck_size = 0, - .type_size = 0, - .is_quantized = false, - }, - [5] = { // GGML_TYPE_Q4_3 - .type_name = "DEPRECATED", - .blck_size = 0, - .type_size = 0, - .is_quantized = false, - }, - [GGML_TYPE_Q5_0] = { - .type_name = "q5_0", - .blck_size = QK5_0, - .type_size = sizeof(block_q5_0), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_q5_0, - .from_float_ref = (ggml_from_float_t) quantize_row_q5_0_ref, - }, - [GGML_TYPE_Q5_1] = { - .type_name = "q5_1", - .blck_size = QK5_1, - .type_size = sizeof(block_q5_1), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_q5_1, - .from_float_ref = (ggml_from_float_t) quantize_row_q5_1_ref, - }, - [GGML_TYPE_Q8_0] = { - .type_name = "q8_0", - .blck_size = QK8_0, - .type_size = sizeof(block_q8_0), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_q8_0, - .from_float_ref = (ggml_from_float_t) quantize_row_q8_0_ref, - }, - [GGML_TYPE_Q4_0_ROCMFP4] = { - .type_name = "q4_0_rocmfp4", - .blck_size = QK_ROCMFP4, - .type_size = sizeof(block_rocmfp4), - .is_quantized = true, - .to_float = (ggml_to_float_t) rocmfp4_dequantize_row_q4_0, - .from_float_ref = (ggml_from_float_t) rocmfp4_quantize_row_q4_0_ref, - }, - [GGML_TYPE_Q4_0_ROCMFP4_FAST] = { - .type_name = "q4_0_rocmfp4_fast", - .blck_size = QK_ROCMFP4, - .type_size = sizeof(block_rocmfp4_fast), - .is_quantized = true, - .to_float = (ggml_to_float_t) rocmfp4_dequantize_row_q4_0_fast, - .from_float_ref = (ggml_from_float_t) rocmfp4_quantize_row_q4_0_fast_ref, - }, - [GGML_TYPE_Q3_0_ROCMFPX] = { - .type_name = "q3_0_rocmfpx", - .blck_size = QK_ROCMFP3, - .type_size = sizeof(block_rocmfp3), - .is_quantized = true, - .to_float = (ggml_to_float_t) rocmfpx_dequantize_row_fp3, - .from_float_ref = (ggml_from_float_t) rocmfpx_quantize_row_fp3_ref, - }, - [GGML_TYPE_Q2_0_ROCMFP2] = { - .type_name = "q2_0_rocmfp2", - .blck_size = QK_ROCMFP2, - .type_size = sizeof(block_rocmfp2), - .is_quantized = true, - .to_float = (ggml_to_float_t) rocmfpx_dequantize_row_fp2, - .from_float_ref = (ggml_from_float_t) rocmfpx_quantize_row_fp2_ref, - }, - [GGML_TYPE_Q6_0_ROCMFPX] = { - .type_name = "q6_0_rocmfpx", - .blck_size = QK_ROCMFP6, - .type_size = sizeof(block_rocmfp6), - .is_quantized = true, - .to_float = (ggml_to_float_t) rocmfpx_dequantize_row_fp6, - .from_float_ref = (ggml_from_float_t) rocmfpx_quantize_row_fp6_ref, - }, - [GGML_TYPE_Q8_0_ROCMFPX] = { - .type_name = "q8_0_rocmfpx", - .blck_size = QK_ROCMFP8, - .type_size = sizeof(block_rocmfp8), - .is_quantized = true, - .to_float = (ggml_to_float_t) rocmfpx_dequantize_row_fp8, - .from_float_ref = (ggml_from_float_t) rocmfpx_quantize_row_fp8_ref, - }, - [GGML_TYPE_Q8_1] = { - .type_name = "q8_1", - .blck_size = QK8_1, - .type_size = sizeof(block_q8_1), - .is_quantized = true, - .from_float_ref = (ggml_from_float_t) quantize_row_q8_1_ref, - }, - [GGML_TYPE_MXFP4] = { - .type_name = "mxfp4", - .blck_size = QK_MXFP4, - .type_size = sizeof(block_mxfp4), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_mxfp4, - .from_float_ref = (ggml_from_float_t)quantize_row_mxfp4_ref, - }, - [GGML_TYPE_NVFP4] = { - .type_name = "nvfp4", - .blck_size = QK_NVFP4, - .type_size = sizeof(block_nvfp4), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_nvfp4, - .from_float_ref = (ggml_from_float_t)quantize_row_nvfp4_ref, - }, - [GGML_TYPE_Q2_K] = { - .type_name = "q2_K", - .blck_size = QK_K, - .type_size = sizeof(block_q2_K), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_q2_K, - .from_float_ref = (ggml_from_float_t) quantize_row_q2_K_ref, - }, - [GGML_TYPE_Q3_K] = { - .type_name = "q3_K", - .blck_size = QK_K, - .type_size = sizeof(block_q3_K), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_q3_K, - .from_float_ref = (ggml_from_float_t) quantize_row_q3_K_ref, - }, - [GGML_TYPE_Q4_K] = { - .type_name = "q4_K", - .blck_size = QK_K, - .type_size = sizeof(block_q4_K), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_q4_K, - .from_float_ref = (ggml_from_float_t) quantize_row_q4_K_ref, - }, - [GGML_TYPE_Q5_K] = { - .type_name = "q5_K", - .blck_size = QK_K, - .type_size = sizeof(block_q5_K), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_q5_K, - .from_float_ref = (ggml_from_float_t) quantize_row_q5_K_ref, - }, - [GGML_TYPE_Q6_K] = { - .type_name = "q6_K", - .blck_size = QK_K, - .type_size = sizeof(block_q6_K), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_q6_K, - .from_float_ref = (ggml_from_float_t) quantize_row_q6_K_ref, - }, - [GGML_TYPE_IQ2_XXS] = { - .type_name = "iq2_xxs", - .blck_size = QK_K, - .type_size = sizeof(block_iq2_xxs), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_iq2_xxs, - .from_float_ref = NULL, - }, - [GGML_TYPE_IQ2_XS] = { - .type_name = "iq2_xs", - .blck_size = QK_K, - .type_size = sizeof(block_iq2_xs), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_iq2_xs, - .from_float_ref = NULL, - }, - [GGML_TYPE_IQ3_XXS] = { - .type_name = "iq3_xxs", - .blck_size = QK_K, - .type_size = sizeof(block_iq3_xxs), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_iq3_xxs, - .from_float_ref = (ggml_from_float_t)quantize_row_iq3_xxs_ref, - }, - [GGML_TYPE_IQ3_S] = { - .type_name = "iq3_s", - .blck_size = QK_K, - .type_size = sizeof(block_iq3_s), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_iq3_s, - .from_float_ref = (ggml_from_float_t)quantize_row_iq3_s_ref, - }, - [GGML_TYPE_IQ2_S] = { - .type_name = "iq2_s", - .blck_size = QK_K, - .type_size = sizeof(block_iq2_s), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_iq2_s, - .from_float_ref = (ggml_from_float_t)quantize_row_iq2_s_ref, - }, - [GGML_TYPE_IQ1_S] = { - .type_name = "iq1_s", - .blck_size = QK_K, - .type_size = sizeof(block_iq1_s), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_iq1_s, - .from_float_ref = NULL, - }, - [GGML_TYPE_IQ1_M] = { - .type_name = "iq1_m", - .blck_size = QK_K, - .type_size = sizeof(block_iq1_m), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_iq1_m, - .from_float_ref = NULL, - }, - [GGML_TYPE_IQ4_NL] = { - .type_name = "iq4_nl", - .blck_size = QK4_NL, - .type_size = sizeof(block_iq4_nl), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_iq4_nl, - .from_float_ref = (ggml_from_float_t)quantize_row_iq4_nl_ref, - }, - [GGML_TYPE_IQ4_XS] = { - .type_name = "iq4_xs", - .blck_size = QK_K, - .type_size = sizeof(block_iq4_xs), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_iq4_xs, - .from_float_ref = (ggml_from_float_t)quantize_row_iq4_xs_ref, - }, - [GGML_TYPE_Q8_K] = { - .type_name = "q8_K", - .blck_size = QK_K, - .type_size = sizeof(block_q8_K), - .is_quantized = true, - }, - [GGML_TYPE_BF16] = { - .type_name = "bf16", - .blck_size = 1, - .type_size = sizeof(ggml_bf16_t), - .is_quantized = false, - .to_float = (ggml_to_float_t) ggml_bf16_to_fp32_row, - .from_float_ref = (ggml_from_float_t) ggml_fp32_to_bf16_row_ref, - }, - [31] = { // GGML_TYPE_Q4_0_4_4 - .type_name = "TYPE_Q4_0_4_4 REMOVED, use Q4_0 with runtime repacking", - .blck_size = 0, - .type_size = 0, - .is_quantized = false, - }, - [32] = { // GGML_TYPE_Q4_0_4_8 - .type_name = "TYPE_Q4_0_4_8 REMOVED, use Q4_0 with runtime repacking", - .blck_size = 0, - .type_size = 0, - .is_quantized = false, - }, - [33] = { // GGML_TYPE_Q4_0_8_8 - .type_name = "TYPE_Q4_0_8_8 REMOVED, use Q4_0 with runtime repacking", - .blck_size = 0, - .type_size = 0, - .is_quantized = false, - }, - [GGML_TYPE_TQ1_0] = { - .type_name = "tq1_0", - .blck_size = QK_K, - .type_size = sizeof(block_tq1_0), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_tq1_0, - .from_float_ref = (ggml_from_float_t) quantize_row_tq1_0_ref, - }, - [GGML_TYPE_TQ2_0] = { - .type_name = "tq2_0", - .blck_size = QK_K, - .type_size = sizeof(block_tq2_0), - .is_quantized = true, - .to_float = (ggml_to_float_t) dequantize_row_tq2_0, - .from_float_ref = (ggml_from_float_t) quantize_row_tq2_0_ref, - }, - [36] = { // GGML_TYPE_IQ4_NL_4_4 - .type_name = "TYPE_IQ4_NL_4_4 REMOVED, use IQ4_NL with runtime repacking", - .blck_size = 0, - .type_size = 0, - .is_quantized = false, - }, - [37] = { // GGML_TYPE_IQ4_NL_4_8 - .type_name = "TYPE_IQ4_NL_4_8 REMOVED, use IQ4_NL with runtime repacking", - .blck_size = 0, - .type_size = 0, - .is_quantized = false, - }, - [38] = { // GGML_TYPE_IQ4_NL_8_8 - .type_name = "TYPE_IQ4_NL_8_8 REMOVED, use IQ4_NL with runtime repacking", - .blck_size = 0, - .type_size = 0, - .is_quantized = false, - }, -}; - -const struct ggml_type_traits * ggml_get_type_traits(enum ggml_type type) { - assert(type >= 0); - assert(type < GGML_TYPE_COUNT); - return &type_traits[type]; -} - -// -// ggml object -// - -struct ggml_object { - size_t offs; - size_t size; - - struct ggml_object * next; - - enum ggml_object_type type; - - char padding[4]; -}; - -static const size_t GGML_OBJECT_SIZE = sizeof(struct ggml_object); - -// -// ggml context -// - -struct ggml_context { - size_t mem_size; - void * mem_buffer; - bool mem_buffer_owned; - bool no_alloc; - - int n_objects; - - struct ggml_object * objects_begin; - struct ggml_object * objects_end; -}; - -// -// data types -// - -static const char * GGML_OP_NAME[GGML_OP_COUNT] = { - "NONE", - - "DUP", - "ADD", - "ADD_ID", - "ADD1", - "ACC", - "SUB", - "MUL", - "DIV", - "SQR", - "SQRT", - "LOG", - "SIN", - "COS", - "SUM", - "SUM_ROWS", - "CUMSUM", - "MEAN", - "ARGMAX", - "COUNT_EQUAL", - "REPEAT", - "REPEAT_BACK", - "CONCAT", - "SILU_BACK", - "NORM", - "RMS_NORM", - "RMS_NORM_BACK", - "GROUP_NORM", - "L2_NORM", - - "MUL_MAT", - "MUL_MAT_ID", - "OUT_PROD", - - "SCALE", - "SET", - "CPY", - "CONT", - "RESHAPE", - "VIEW", - "PERMUTE", - "TRANSPOSE", - "GET_ROWS", - "GET_ROWS_BACK", - "SET_ROWS", - "DIAG", - "DIAG_MASK_INF", - "DIAG_MASK_ZERO", - "SOFT_MAX", - "SOFT_MAX_BACK", - "ROPE", - "ROPE_BACK", - "CLAMP", - "CONV_TRANSPOSE_1D", - "IM2COL", - "IM2COL_BACK", - "IM2COL_3D", - "CONV_2D", - "CONV_3D", - "CONV_2D_DW", - "CONV_TRANSPOSE_2D", - "POOL_1D", - "POOL_2D", - "POOL_2D_BACK", - "UPSCALE", - "PAD", - "PAD_REFLECT_1D", - "ROLL", - "ARANGE", - "TIMESTEP_EMBEDDING", - "ARGSORT", - "TOP_K", - "LEAKY_RELU", - "TRI", - "FILL", - - "FLASH_ATTN_EXT", - "FLASH_ATTN_SPARSE", - "FLASH_ATTN_BACK", - "SSM_CONV", - "SSM_SCAN", - "WIN_PART", - "WIN_UNPART", - "GET_REL_POS", - "ADD_REL_POS", - "RWKV_WKV6", - "GATED_LINEAR_ATTN", - "RWKV_WKV7", - "SOLVE_TRI", - "GATED_DELTA_NET", - - "UNARY", - - "MAP_CUSTOM1", - "MAP_CUSTOM2", - "MAP_CUSTOM3", - - "CUSTOM", - - "CROSS_ENTROPY_LOSS", - "CROSS_ENTROPY_LOSS_BACK", - "OPT_STEP_ADAMW", - "OPT_STEP_SGD", - - "GLU", - - "TURBO_WHT", -}; - -static_assert(GGML_OP_COUNT == 99, "GGML_OP_COUNT != 99"); - -static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = { - "none", - - "x", - "x+y", - "x[i]+y", - "x+y", - "view(x,nb,offset)+=y->x", - "x-y", - "x*y", - "x/y", - "x^2", - "√x", - "log(x)", - "sin(x)", - "cos(x)", - "Σx", - "Σx_k", - "cumsum(x)", - "Σx/n", - "argmax(x)", - "count_equal(x)", - "repeat(x)", - "repeat_back(x)", - "concat(x, y)", - "silu_back(x)", - "norm(x)", - "rms_norm(x)", - "rms_norm_back(x)", - "group_norm(x)", - "l2_norm(x)", - - "X*Y", - "X[i]*Y", - "X*Y", - - "x*v", - "y-\\>view(x)", - "x-\\>y", - "cont(x)", - "reshape(x)", - "view(x)", - "permute(x)", - "transpose(x)", - "get_rows(x)", - "get_rows_back(x)", - "set_rows(x)", - "diag(x)", - "diag_mask_inf(x)", - "diag_mask_zero(x)", - "soft_max(x)", - "soft_max_back(x)", - "rope(x)", - "rope_back(x)", - "clamp(x)", - "conv_transpose_1d(x)", - "im2col(x)", - "im2col_back(x)", - "im2col_3d(x)", - "conv_2d(x)", - "conv_3d(x)", - "conv_2d_dw(x)", - "conv_transpose_2d(x)", - "pool_1d(x)", - "pool_2d(x)", - "pool_2d_back(x)", - "upscale(x)", - "pad(x)", - "pad_reflect_1d(x)", - "roll(x)", - "arange(start, stop, step)", - "timestep_embedding(timesteps, dim, max_period)", - "argsort(x)", - "top_k(x)", - "leaky_relu(x)", - "tri(x)", - "fill(x, c)", - - "flash_attn_ext(x)", - "sparse_fa(x)", - "flash_attn_back(x)", - "ssm_conv(x)", - "ssm_scan(x)", - "win_part(x)", - "win_unpart(x)", - "get_rel_pos(x)", - "add_rel_pos(x)", - "rwkv_wkv6(k, v, r, tf, td, s)", - "gated_linear_attn(k, v, q, gate, s)", - "rwkv_wkv7(r, w, k, v, a, b, s)", - "A X = B, A triangular, solve X", - "gated_delta_net(q, k, v, g, beta, s)", - - "unary(x)", - - "map_custom(x)", - "map_custom(x,y)", - "map_custom(x,y,z)", - - "custom(x)", - - "cross_entropy_loss(x,y)", - "cross_entropy_loss_back(x,y)", - "adamw(x)", - "sgd(x)", - - "glu(x)", - - "turbo_wht(a)", -}; - -static_assert(GGML_OP_COUNT == 99, "GGML_OP_COUNT != 99"); - -static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2"); - -static const char * GGML_UNARY_OP_NAME[GGML_UNARY_OP_COUNT] = { - "ABS", - "SGN", - "NEG", - "STEP", - "TANH", - "ELU", - "RELU", - "SIGMOID", - "GELU", - "GELU_QUICK", - "SILU", - "HARDSWISH", - "HARDSIGMOID", - "EXP", - "EXPM1", - "SOFTPLUS", - "GELU_ERF", - "XIELU", - "FLOOR", - "CEIL", - "ROUND", - "TRUNC", -}; - -static_assert(GGML_UNARY_OP_COUNT == 22, "GGML_UNARY_OP_COUNT != 22"); - -static const char * GGML_GLU_OP_NAME[GGML_GLU_OP_COUNT] = { - "REGLU", - "GEGLU", - "SWIGLU", - "SWIGLU_OAI", - "SWIGLU_DS4", - "GEGLU_ERF", - "GEGLU_QUICK", -}; - -static_assert(GGML_GLU_OP_COUNT == 7, "GGML_GLU_OP_COUNT != 7"); - - -static_assert(sizeof(struct ggml_object)%GGML_MEM_ALIGN == 0, "ggml_object size must be a multiple of GGML_MEM_ALIGN"); -static_assert(sizeof(struct ggml_tensor)%GGML_MEM_ALIGN == 0, "ggml_tensor size must be a multiple of GGML_MEM_ALIGN"); - - -//////////////////////////////////////////////////////////////////////////////// - -void ggml_print_object(const struct ggml_object * obj) { - GGML_LOG_INFO(" - ggml_object: type = %d, offset = %zu, size = %zu, next = %p\n", - obj->type, obj->offs, obj->size, (const void *) obj->next); -} - -void ggml_print_objects(const struct ggml_context * ctx) { - struct ggml_object * obj = ctx->objects_begin; - - GGML_LOG_INFO("%s: objects in context %p:\n", __func__, (const void *) ctx); - - while (obj != NULL) { - ggml_print_object(obj); - obj = obj->next; - } - - GGML_LOG_INFO("%s: --- end ---\n", __func__); -} - -int64_t ggml_nelements(const struct ggml_tensor * tensor) { - static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); - - return tensor->ne[0]*tensor->ne[1]*tensor->ne[2]*tensor->ne[3]; -} - -int64_t ggml_nrows(const struct ggml_tensor * tensor) { - static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); - - return tensor->ne[1]*tensor->ne[2]*tensor->ne[3]; -} - -size_t ggml_nbytes(const struct ggml_tensor * tensor) { - for (int i = 0; i < GGML_MAX_DIMS; ++i) { - if (tensor->ne[i] <= 0) { - return 0; - } - } - - size_t nbytes; - const size_t blck_size = ggml_blck_size(tensor->type); - if (blck_size == 1) { - nbytes = ggml_type_size(tensor->type); - for (int i = 0; i < GGML_MAX_DIMS; ++i) { - nbytes += (tensor->ne[i] - 1)*tensor->nb[i]; - } - } - else { - nbytes = tensor->ne[0]*tensor->nb[0]/blck_size; - for (int i = 1; i < GGML_MAX_DIMS; ++i) { - nbytes += (tensor->ne[i] - 1)*tensor->nb[i]; - } - } - - return nbytes; -} - -size_t ggml_nbytes_pad(const struct ggml_tensor * tensor) { - return GGML_PAD(ggml_nbytes(tensor), GGML_MEM_ALIGN); -} - -int64_t ggml_blck_size(enum ggml_type type) { - assert(type >= 0); - assert(type < GGML_TYPE_COUNT); - return type_traits[type].blck_size; -} - -size_t ggml_type_size(enum ggml_type type) { - assert(type >= 0); - assert(type < GGML_TYPE_COUNT); - return type_traits[type].type_size; -} - -size_t ggml_row_size(enum ggml_type type, int64_t ne) { - assert(type >= 0); - assert(type < GGML_TYPE_COUNT); - assert(ne % ggml_blck_size(type) == 0); - return ggml_type_size(type)*ne/ggml_blck_size(type); -} - -double ggml_type_sizef(enum ggml_type type) { - assert(type >= 0); - assert(type < GGML_TYPE_COUNT); - return ((double)(type_traits[type].type_size))/type_traits[type].blck_size; -} - -const char * ggml_type_name(enum ggml_type type) { - assert(type >= 0); - assert(type < GGML_TYPE_COUNT); - return type_traits[type].type_name; -} - -bool ggml_is_quantized(enum ggml_type type) { - assert(type >= 0); - assert(type < GGML_TYPE_COUNT); - return type_traits[type].is_quantized; -} - -const char * ggml_op_name(enum ggml_op op) { - return GGML_OP_NAME[op]; -} - -const char * ggml_op_symbol(enum ggml_op op) { - return GGML_OP_SYMBOL[op]; -} - -const char * ggml_unary_op_name(enum ggml_unary_op op) { - return GGML_UNARY_OP_NAME[op]; -} - -const char * ggml_glu_op_name(enum ggml_glu_op op) { - return GGML_GLU_OP_NAME[op]; -} - -const char * ggml_op_desc(const struct ggml_tensor * t) { - if (t->op == GGML_OP_UNARY) { - enum ggml_unary_op uop = ggml_get_unary_op(t); - return ggml_unary_op_name(uop); - } - if (t->op == GGML_OP_GLU) { - enum ggml_glu_op gop = ggml_get_glu_op(t); - return ggml_glu_op_name(gop); - } - return ggml_op_name(t->op); -} - -size_t ggml_element_size(const struct ggml_tensor * tensor) { - return ggml_type_size(tensor->type); -} - -bool ggml_is_scalar(const struct ggml_tensor * tensor) { - static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); - - return tensor->ne[0] == 1 && tensor->ne[1] == 1 && tensor->ne[2] == 1 && tensor->ne[3] == 1; -} - -bool ggml_is_vector(const struct ggml_tensor * tensor) { - static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); - - return tensor->ne[1] == 1 && tensor->ne[2] == 1 && tensor->ne[3] == 1; -} - -bool ggml_is_matrix(const struct ggml_tensor * tensor) { - static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); - - return tensor->ne[2] == 1 && tensor->ne[3] == 1; -} - -bool ggml_is_3d(const struct ggml_tensor * tensor) { - return tensor->ne[3] == 1; -} - -int ggml_n_dims(const struct ggml_tensor * tensor) { - for (int i = GGML_MAX_DIMS - 1; i >= 1; --i) { - if (tensor->ne[i] > 1) { - return i + 1; - } - } - return 1; -} - -enum ggml_type ggml_ftype_to_ggml_type(enum ggml_ftype ftype) { - enum ggml_type wtype = GGML_TYPE_COUNT; - - switch (ftype) { - case GGML_FTYPE_ALL_F32: wtype = GGML_TYPE_F32; break; - case GGML_FTYPE_MOSTLY_F16: wtype = GGML_TYPE_F16; break; - case GGML_FTYPE_MOSTLY_BF16: wtype = GGML_TYPE_BF16; break; - case GGML_FTYPE_MOSTLY_Q4_0: wtype = GGML_TYPE_Q4_0; break; - case GGML_FTYPE_MOSTLY_Q4_1: wtype = GGML_TYPE_Q4_1; break; - case GGML_FTYPE_MOSTLY_Q1_0: wtype = GGML_TYPE_Q1_0; break; - case GGML_FTYPE_MOSTLY_Q5_0: wtype = GGML_TYPE_Q5_0; break; - case GGML_FTYPE_MOSTLY_Q5_1: wtype = GGML_TYPE_Q5_1; break; - case GGML_FTYPE_MOSTLY_Q8_0: wtype = GGML_TYPE_Q8_0; break; - case GGML_FTYPE_MOSTLY_MXFP4: wtype = GGML_TYPE_MXFP4; break; - case GGML_FTYPE_MOSTLY_NVFP4: wtype = GGML_TYPE_NVFP4; break; - case GGML_FTYPE_MOSTLY_Q2_K: wtype = GGML_TYPE_Q2_K; break; - case GGML_FTYPE_MOSTLY_Q3_K: wtype = GGML_TYPE_Q3_K; break; - case GGML_FTYPE_MOSTLY_Q4_K: wtype = GGML_TYPE_Q4_K; break; - case GGML_FTYPE_MOSTLY_Q5_K: wtype = GGML_TYPE_Q5_K; break; - case GGML_FTYPE_MOSTLY_Q6_K: wtype = GGML_TYPE_Q6_K; break; - case GGML_FTYPE_MOSTLY_IQ2_XXS: wtype = GGML_TYPE_IQ2_XXS; break; - case GGML_FTYPE_MOSTLY_IQ2_XS: wtype = GGML_TYPE_IQ2_XS; break; - case GGML_FTYPE_MOSTLY_IQ3_XXS: wtype = GGML_TYPE_IQ3_XXS; break; - case GGML_FTYPE_MOSTLY_IQ1_S: wtype = GGML_TYPE_IQ1_S; break; - case GGML_FTYPE_MOSTLY_IQ1_M: wtype = GGML_TYPE_IQ1_M; break; - case GGML_FTYPE_MOSTLY_IQ4_NL: wtype = GGML_TYPE_IQ4_NL; break; - case GGML_FTYPE_MOSTLY_IQ4_XS: wtype = GGML_TYPE_IQ4_XS; break; - case GGML_FTYPE_MOSTLY_IQ3_S: wtype = GGML_TYPE_IQ3_S; break; - case GGML_FTYPE_MOSTLY_IQ2_S: wtype = GGML_TYPE_IQ2_S; break; - case GGML_FTYPE_UNKNOWN: wtype = GGML_TYPE_COUNT; break; - case GGML_FTYPE_MOSTLY_Q4_1_SOME_F16: wtype = GGML_TYPE_COUNT; break; - } - - GGML_ASSERT(wtype != GGML_TYPE_COUNT); - - return wtype; -} - -size_t ggml_tensor_overhead(void) { - return GGML_OBJECT_SIZE + GGML_TENSOR_SIZE; -} - -bool ggml_is_transposed(const struct ggml_tensor * tensor) { - return tensor->nb[0] > tensor->nb[1]; -} - -static bool ggml_is_contiguous_n(const struct ggml_tensor * tensor, int n) { - size_t next_nb = ggml_type_size(tensor->type); - if (tensor->ne[0] != ggml_blck_size(tensor->type) && tensor->nb[0] != next_nb) { - return false; - } - next_nb *= tensor->ne[0]/ggml_blck_size(tensor->type); - for (int i = 1; i < GGML_MAX_DIMS; i++) { - if (i > n) { - if (tensor->ne[i] != 1 && tensor->nb[i] != next_nb) { - return false; - } - next_nb *= tensor->ne[i]; - } else { - // this dimension does not need to be contiguous - next_nb = tensor->ne[i]*tensor->nb[i]; - } - } - return true; -} - -bool ggml_is_contiguous(const struct ggml_tensor * tensor) { - return ggml_is_contiguous_0(tensor); -} - -bool ggml_is_contiguous_0(const struct ggml_tensor * tensor) { - return ggml_is_contiguous_n(tensor, 0); -} - -bool ggml_is_contiguous_1(const struct ggml_tensor * tensor) { - return ggml_is_contiguous_n(tensor, 1); -} - -bool ggml_is_contiguous_2(const struct ggml_tensor * tensor) { - return ggml_is_contiguous_n(tensor, 2); -} - -bool ggml_is_contiguously_allocated(const struct ggml_tensor * tensor) { - return ggml_nbytes(tensor) == ggml_nelements(tensor) * ggml_type_size(tensor->type)/ggml_blck_size(tensor->type); -} - -bool ggml_is_permuted(const struct ggml_tensor * tensor) { - static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); - - return tensor->nb[0] > tensor->nb[1] || tensor->nb[1] > tensor->nb[2] || tensor->nb[2] > tensor->nb[3]; -} - -bool ggml_is_contiguous_channels(const struct ggml_tensor * tensor) { - return - tensor->nb[0] > tensor->nb[2] && - tensor->nb[1] > tensor->nb[0] && - tensor->nb[2] == ggml_type_size(tensor->type); -} - -bool ggml_is_contiguous_rows(const struct ggml_tensor * tensor) { - return - tensor->ne[0] == ggml_blck_size(tensor->type) || - tensor->nb[0] == ggml_type_size(tensor->type); -} - -static inline bool ggml_is_padded_1d(const struct ggml_tensor * tensor) { - static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); - - return - tensor->nb[0] == ggml_type_size(tensor->type) && - tensor->nb[2] == tensor->nb[1]*tensor->ne[1] && - tensor->nb[3] == tensor->nb[2]*tensor->ne[2]; -} - -bool ggml_is_empty(const struct ggml_tensor * tensor) { - for (int i = 0; i < GGML_MAX_DIMS; ++i) { - if (tensor->ne[i] == 0) { - // empty if any dimension has no elements - return true; - } - } - return false; -} - -bool ggml_are_same_shape(const struct ggml_tensor * t0, const struct ggml_tensor * t1) { - static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); - - return - (t0->ne[0] == t1->ne[0]) && - (t0->ne[1] == t1->ne[1]) && - (t0->ne[2] == t1->ne[2]) && - (t0->ne[3] == t1->ne[3]); -} - -bool ggml_are_same_stride(const struct ggml_tensor * t0, const struct ggml_tensor * t1) { - static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); - - return - (t0->nb[0] == t1->nb[0]) && - (t0->nb[1] == t1->nb[1]) && - (t0->nb[2] == t1->nb[2]) && - (t0->nb[3] == t1->nb[3]); -} - -bool ggml_is_view(const struct ggml_tensor * t) { - return ggml_impl_is_view(t); -} - -// check if t1 can be represented as a repetition of t0 -bool ggml_can_repeat(const struct ggml_tensor * t0, const struct ggml_tensor * t1) { - static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); - - return ggml_is_empty(t0) ? ggml_is_empty(t1) : - (t1->ne[0]%t0->ne[0] == 0) && - (t1->ne[1]%t0->ne[1] == 0) && - (t1->ne[2]%t0->ne[2] == 0) && - (t1->ne[3]%t0->ne[3] == 0); -} - -static inline bool ggml_can_repeat_rows(const struct ggml_tensor * t0, const struct ggml_tensor * t1) { - static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); - - return (t0->ne[0] == t1->ne[0]) && ggml_can_repeat(t0, t1); -} - -// assert that pointer is aligned to GGML_MEM_ALIGN -#define GGML_ASSERT_ALIGNED(ptr) \ - GGML_ASSERT(((uintptr_t) (ptr))%GGML_MEM_ALIGN == 0) - -//////////////////////////////////////////////////////////////////////////////// - -struct ggml_context * ggml_init(struct ggml_init_params params) { - static bool is_first_call = true; - - ggml_critical_section_start(); - - if (is_first_call) { - // initialize time system (required on Windows) - ggml_time_init(); - - is_first_call = false; - } - - ggml_critical_section_end(); - - struct ggml_context * ctx = GGML_MALLOC(sizeof(struct ggml_context)); - - // allow to call ggml_init with 0 size - if (params.mem_size == 0) { - params.mem_size = GGML_MEM_ALIGN; - } - - const size_t mem_size = params.mem_buffer ? params.mem_size : GGML_PAD(params.mem_size, GGML_MEM_ALIGN); - - *ctx = (struct ggml_context) { - /*.mem_size =*/ mem_size, - /*.mem_buffer =*/ params.mem_buffer ? params.mem_buffer : ggml_aligned_malloc(mem_size), - /*.mem_buffer_owned =*/ params.mem_buffer ? false : true, - /*.no_alloc =*/ params.no_alloc, - /*.n_objects =*/ 0, - /*.objects_begin =*/ NULL, - /*.objects_end =*/ NULL, - }; - - GGML_ASSERT(ctx->mem_buffer != NULL); - - GGML_ASSERT_ALIGNED(ctx->mem_buffer); - - GGML_PRINT_DEBUG("%s: context initialized\n", __func__); - - return ctx; -} - -void ggml_reset(struct ggml_context * ctx) { - if (ctx == NULL) { - return; - } - - ctx->n_objects = 0; - ctx->objects_begin = NULL; - ctx->objects_end = NULL; -} - -void ggml_free(struct ggml_context * ctx) { - if (ctx == NULL) { - return; - } - - if (ctx->mem_buffer_owned) { - ggml_aligned_free(ctx->mem_buffer, ctx->mem_size); - } - - GGML_FREE(ctx); -} - -size_t ggml_used_mem(const struct ggml_context * ctx) { - return ctx->objects_end == NULL ? 0 : ctx->objects_end->offs + ctx->objects_end->size; -} - -bool ggml_get_no_alloc(struct ggml_context * ctx) { - return ctx->no_alloc; -} - -void ggml_set_no_alloc(struct ggml_context * ctx, bool no_alloc) { - ctx->no_alloc = no_alloc; -} - -void * ggml_get_mem_buffer(const struct ggml_context * ctx) { - return ctx->mem_buffer; -} - -size_t ggml_get_mem_size(const struct ggml_context * ctx) { - return ctx->mem_size; -} - -size_t ggml_get_max_tensor_size(const struct ggml_context * ctx) { - size_t max_size = 0; - - for (struct ggml_tensor * tensor = ggml_get_first_tensor(ctx); tensor != NULL; tensor = ggml_get_next_tensor(ctx, tensor)) { - size_t bytes = ggml_nbytes(tensor); - max_size = MAX(max_size, bytes); - } - - return max_size; -} - -//////////////////////////////////////////////////////////////////////////////// - -static struct ggml_object * ggml_new_object(struct ggml_context * ctx, enum ggml_object_type type, size_t size) { - // always insert objects at the end of the context's memory pool - struct ggml_object * obj_cur = ctx->objects_end; - - const size_t cur_offs = obj_cur == NULL ? 0 : obj_cur->offs; - const size_t cur_size = obj_cur == NULL ? 0 : obj_cur->size; - const size_t cur_end = cur_offs + cur_size; - - // align to GGML_MEM_ALIGN - GGML_ASSERT(size <= SIZE_MAX - (GGML_MEM_ALIGN - 1)); - size_t size_needed = GGML_PAD(size, GGML_MEM_ALIGN); - - char * const mem_buffer = ctx->mem_buffer; - struct ggml_object * const obj_new = (struct ggml_object *)(mem_buffer + cur_end); - - // integer overflow checks - if (cur_end > SIZE_MAX - size_needed) { - GGML_LOG_WARN("%s: overflow detected in cur_end (%zu) + size_needed (%zu)\n", __func__, cur_end, size_needed); - return NULL; - } - if (cur_end + size_needed > SIZE_MAX - GGML_OBJECT_SIZE) { - GGML_LOG_WARN("%s: overflow detected in cur_end (%zu) + size_needed (%zu) + GGML_OBJECT_SIZE (%zu)\n", __func__, - cur_end, size_needed, (size_t) GGML_OBJECT_SIZE); - return NULL; - } - - if (cur_end + size_needed + GGML_OBJECT_SIZE > ctx->mem_size) { - GGML_LOG_WARN("%s: not enough space in the context's memory pool (needed %zu, available %zu)\n", - __func__, cur_end + size_needed + GGML_OBJECT_SIZE, ctx->mem_size); -#ifndef NDEBUG - GGML_ABORT("not enough space in the context's memory pool"); -#endif - return NULL; - } - - *obj_new = (struct ggml_object) { - .offs = cur_end + GGML_OBJECT_SIZE, - .size = size_needed, - .next = NULL, - .type = type, - }; - - GGML_ASSERT_ALIGNED(mem_buffer + obj_new->offs); - - if (obj_cur != NULL) { - obj_cur->next = obj_new; - } else { - // this is the first object in this context - ctx->objects_begin = obj_new; - } - - ctx->objects_end = obj_new; - - //printf("%s: inserted new object at %zu, size = %zu\n", __func__, cur_end, obj_new->size); - - return obj_new; -} - -static struct ggml_tensor * ggml_new_tensor_impl( - struct ggml_context * ctx, - enum ggml_type type, - int n_dims, - const int64_t * ne, - struct ggml_tensor * view_src, - size_t view_offs) { - - GGML_ASSERT(type >= 0 && type < GGML_TYPE_COUNT); - GGML_ASSERT(n_dims >= 1 && n_dims <= GGML_MAX_DIMS); - - // find the base tensor and absolute offset - if (view_src != NULL && view_src->view_src != NULL) { - view_offs += view_src->view_offs; - view_src = view_src->view_src; - } - - size_t data_size = ggml_row_size(type, ne[0]); - for (int i = 1; i < n_dims; i++) { - data_size *= ne[i]; - } - - GGML_ASSERT(view_src == NULL || data_size == 0 || data_size + view_offs <= ggml_nbytes(view_src)); - - void * data = view_src != NULL ? view_src->data : NULL; - if (data != NULL) { - data = (char *) data + view_offs; - } - - size_t obj_alloc_size = 0; - - if (view_src == NULL && !ctx->no_alloc) { - // allocate tensor data in the context's memory pool - obj_alloc_size = data_size; - } - - GGML_ASSERT(GGML_TENSOR_SIZE <= SIZE_MAX - obj_alloc_size); - - struct ggml_object * const obj_new = ggml_new_object(ctx, GGML_OBJECT_TYPE_TENSOR, GGML_TENSOR_SIZE + obj_alloc_size); - GGML_ASSERT(obj_new); - - struct ggml_tensor * const result = (struct ggml_tensor *)((char *)ctx->mem_buffer + obj_new->offs); - - *result = (struct ggml_tensor) { - /*.type =*/ type, - /*.buffer =*/ NULL, - /*.ne =*/ { 1, 1, 1, 1 }, - /*.nb =*/ { 0, 0, 0, 0 }, - /*.op =*/ GGML_OP_NONE, - /*.op_params =*/ { 0 }, - /*.flags =*/ 0, - /*.src =*/ { NULL }, - /*.view_src =*/ view_src, - /*.view_offs =*/ view_offs, - /*.data =*/ obj_alloc_size > 0 ? (void *)(result + 1) : data, - /*.name =*/ { 0 }, - /*.extra =*/ NULL, - /*.padding =*/ { 0 }, - }; - - // TODO: this should not be needed as long as we don't rely on aligned SIMD loads - //GGML_ASSERT_ALIGNED(result->data); - - for (int i = 0; i < n_dims; i++) { - result->ne[i] = ne[i]; - } - - result->nb[0] = ggml_type_size(type); - result->nb[1] = result->nb[0]*(result->ne[0]/ggml_blck_size(type)); - for (int i = 2; i < GGML_MAX_DIMS; i++) { - result->nb[i] = result->nb[i - 1]*result->ne[i - 1]; - } - - ctx->n_objects++; - - return result; -} - -struct ggml_tensor * ggml_new_tensor( - struct ggml_context * ctx, - enum ggml_type type, - int n_dims, - const int64_t * ne) { - return ggml_new_tensor_impl(ctx, type, n_dims, ne, NULL, 0); -} - -struct ggml_tensor * ggml_new_tensor_1d( - struct ggml_context * ctx, - enum ggml_type type, - int64_t ne0) { - return ggml_new_tensor(ctx, type, 1, &ne0); -} - -struct ggml_tensor * ggml_new_tensor_2d( - struct ggml_context * ctx, - enum ggml_type type, - int64_t ne0, - int64_t ne1) { - const int64_t ne[2] = { ne0, ne1 }; - return ggml_new_tensor(ctx, type, 2, ne); -} - -struct ggml_tensor * ggml_new_tensor_3d( - struct ggml_context * ctx, - enum ggml_type type, - int64_t ne0, - int64_t ne1, - int64_t ne2) { - const int64_t ne[3] = { ne0, ne1, ne2 }; - return ggml_new_tensor(ctx, type, 3, ne); -} - -struct ggml_tensor * ggml_new_tensor_4d( - struct ggml_context * ctx, - enum ggml_type type, - int64_t ne0, - int64_t ne1, - int64_t ne2, - int64_t ne3) { - const int64_t ne[4] = { ne0, ne1, ne2, ne3 }; - return ggml_new_tensor(ctx, type, 4, ne); -} - -void * ggml_new_buffer(struct ggml_context * ctx, size_t nbytes) { - struct ggml_object * obj = ggml_new_object(ctx, GGML_OBJECT_TYPE_WORK_BUFFER, nbytes); - - return (uint8_t *)ctx->mem_buffer + obj->offs; -} - -struct ggml_tensor * ggml_dup_tensor(struct ggml_context * ctx, const struct ggml_tensor * src) { - return ggml_new_tensor(ctx, src->type, GGML_MAX_DIMS, src->ne); -} - -void ggml_unravel_index(const struct ggml_tensor * tensor, int64_t i, int64_t * i0, int64_t * i1, int64_t * i2, int64_t * i3) { - const int64_t ne2 = tensor->ne[2]; - const int64_t ne1 = tensor->ne[1]; - const int64_t ne0 = tensor->ne[0]; - - const int64_t i3_ = (i/(ne2*ne1*ne0)); - const int64_t i2_ = (i - i3_*ne2*ne1*ne0)/(ne1*ne0); - const int64_t i1_ = (i - i3_*ne2*ne1*ne0 - i2_*ne1*ne0)/ne0; - const int64_t i0_ = (i - i3_*ne2*ne1*ne0 - i2_*ne1*ne0 - i1_*ne0); - - if (i0) { - * i0 = i0_; - } - if (i1) { - * i1 = i1_; - } - if (i2) { - * i2 = i2_; - } - if (i3) { - * i3 = i3_; - } -} - -void * ggml_get_data(const struct ggml_tensor * tensor) { - return tensor->data; -} - -float * ggml_get_data_f32(const struct ggml_tensor * tensor) { - assert(tensor->type == GGML_TYPE_F32); - return (float *)(tensor->data); -} - -enum ggml_unary_op ggml_get_unary_op(const struct ggml_tensor * tensor) { - GGML_ASSERT(tensor->op == GGML_OP_UNARY); - return (enum ggml_unary_op) ggml_get_op_params_i32(tensor, 0); -} - -enum ggml_glu_op ggml_get_glu_op(const struct ggml_tensor * tensor) { - GGML_ASSERT(tensor->op == GGML_OP_GLU); - return (enum ggml_glu_op) ggml_get_op_params_i32(tensor, 0); -} - -const char * ggml_get_name(const struct ggml_tensor * tensor) { - return tensor->name; -} - -struct ggml_tensor * ggml_set_name(struct ggml_tensor * tensor, const char * name) { - size_t i; - for (i = 0; i < sizeof(tensor->name) - 1 && name[i] != '\0'; i++) { - tensor->name[i] = name[i]; - } - tensor->name[i] = '\0'; - return tensor; -} - -struct ggml_tensor * ggml_format_name(struct ggml_tensor * tensor, const char * fmt, ...) { - va_list args; - va_start(args, fmt); - vsnprintf(tensor->name, sizeof(tensor->name), fmt, args); - va_end(args); - return tensor; -} - -struct ggml_tensor * ggml_view_tensor( - struct ggml_context * ctx, - struct ggml_tensor * src) { - struct ggml_tensor * result = ggml_new_tensor_impl(ctx, src->type, GGML_MAX_DIMS, src->ne, src, 0); - ggml_format_name(result, "%s (view)", src->name); - - for (int i = 0; i < GGML_MAX_DIMS; i++) { - result->nb[i] = src->nb[i]; - } - - return result; -} - -struct ggml_tensor * ggml_get_first_tensor(const struct ggml_context * ctx) { - struct ggml_object * obj = ctx->objects_begin; - - char * const mem_buffer = ctx->mem_buffer; - - while (obj != NULL) { - if (obj->type == GGML_OBJECT_TYPE_TENSOR) { - return (struct ggml_tensor *)(mem_buffer + obj->offs); - } - - obj = obj->next; - } - - return NULL; -} - -struct ggml_tensor * ggml_get_next_tensor(const struct ggml_context * ctx, struct ggml_tensor * tensor) { - struct ggml_object * obj = (struct ggml_object *) ((char *)tensor - GGML_OBJECT_SIZE); - obj = obj->next; - - char * const mem_buffer = ctx->mem_buffer; - - while (obj != NULL) { - if (obj->type == GGML_OBJECT_TYPE_TENSOR) { - return (struct ggml_tensor *)(mem_buffer + obj->offs); - } - - obj = obj->next; - } - - return NULL; -} - -struct ggml_tensor * ggml_get_tensor(struct ggml_context * ctx, const char * name) { - struct ggml_object * obj = ctx->objects_begin; - - char * const mem_buffer = ctx->mem_buffer; - - while (obj != NULL) { - if (obj->type == GGML_OBJECT_TYPE_TENSOR) { - struct ggml_tensor * cur = (struct ggml_tensor *)(mem_buffer + obj->offs); - if (strcmp(cur->name, name) == 0) { - return cur; - } - } - - obj = obj->next; - } - - return NULL; -} - -//////////////////////////////////////////////////////////////////////////////// - -// ggml_dup - -static struct ggml_tensor * ggml_dup_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - bool inplace) { - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - result->op = GGML_OP_DUP; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_dup( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_dup_impl(ctx, a, false); -} - -struct ggml_tensor * ggml_dup_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_dup_impl(ctx, a, true); -} - -// ggml_add - -static struct ggml_tensor * ggml_add_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - bool inplace) { - GGML_ASSERT(ggml_can_repeat(b, a)); - - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - result->op = GGML_OP_ADD; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -struct ggml_tensor * ggml_add( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_add_impl(ctx, a, b, false); -} - -struct ggml_tensor * ggml_add_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_add_impl(ctx, a, b, true); -} - -// ggml_add_cast - -static struct ggml_tensor * ggml_add_cast_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - enum ggml_type type) { - // TODO: support less-strict constraint - // GGML_ASSERT(ggml_can_repeat(b, a)); - GGML_ASSERT(ggml_can_repeat_rows(b, a)); - - // currently only supported for quantized input and f16 - GGML_ASSERT(ggml_is_quantized(a->type) || - a->type == GGML_TYPE_F16 || - a->type == GGML_TYPE_BF16); - - struct ggml_tensor * result = ggml_new_tensor(ctx, type, GGML_MAX_DIMS, a->ne); - - result->op = GGML_OP_ADD; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -struct ggml_tensor * ggml_add_cast( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - enum ggml_type type) { - return ggml_add_cast_impl(ctx, a, b, type); -} - -struct ggml_tensor * ggml_add_id( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * ids) { - - GGML_ASSERT(a->ne[0] == b->ne[0]); - GGML_ASSERT(a->ne[1] == ids->ne[0]); - GGML_ASSERT(a->ne[2] == ids->ne[1]); - GGML_ASSERT(ids->type == GGML_TYPE_I32); - - struct ggml_tensor * result = ggml_dup_tensor(ctx, a); - - result->op = GGML_OP_ADD_ID; - result->src[0] = a; - result->src[1] = b; - result->src[2] = ids; - - return result; -} - -// ggml_add1 - -static struct ggml_tensor * ggml_add1_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - bool inplace) { - GGML_ASSERT(ggml_is_scalar(b)); - GGML_ASSERT(ggml_is_padded_1d(a)); - - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - result->op = GGML_OP_ADD1; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -struct ggml_tensor * ggml_add1( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_add1_impl(ctx, a, b, false); -} - -struct ggml_tensor * ggml_add1_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_add1_impl(ctx, a, b, true); -} - -// ggml_acc - -static struct ggml_tensor * ggml_acc_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - size_t nb1, - size_t nb2, - size_t nb3, - size_t offset, - bool inplace) { - GGML_ASSERT(ggml_nelements(b) <= ggml_nelements(a)); - GGML_ASSERT(ggml_is_contiguous(a)); - GGML_ASSERT(a->type == GGML_TYPE_F32); - GGML_ASSERT(b->type == GGML_TYPE_F32); - - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - int32_t params[] = { nb1, nb2, nb3, offset, inplace ? 1 : 0 }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_ACC; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -struct ggml_tensor * ggml_acc( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - size_t nb1, - size_t nb2, - size_t nb3, - size_t offset) { - return ggml_acc_impl(ctx, a, b, nb1, nb2, nb3, offset, false); -} - -struct ggml_tensor * ggml_acc_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - size_t nb1, - size_t nb2, - size_t nb3, - size_t offset) { - return ggml_acc_impl(ctx, a, b, nb1, nb2, nb3, offset, true); -} - -// ggml_sub - -static struct ggml_tensor * ggml_sub_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - bool inplace) { - GGML_ASSERT(ggml_can_repeat(b, a)); - - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - result->op = GGML_OP_SUB; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -struct ggml_tensor * ggml_sub( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_sub_impl(ctx, a, b, false); -} - -struct ggml_tensor * ggml_sub_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_sub_impl(ctx, a, b, true); -} - -// ggml_mul - -static struct ggml_tensor * ggml_mul_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - bool inplace) { - GGML_ASSERT(ggml_can_repeat(b, a)); - - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - result->op = GGML_OP_MUL; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -struct ggml_tensor * ggml_mul( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_mul_impl(ctx, a, b, false); -} - -struct ggml_tensor * ggml_mul_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_mul_impl(ctx, a, b, true); -} - -// ggml_div - -static struct ggml_tensor * ggml_div_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - bool inplace) { - GGML_ASSERT(ggml_can_repeat(b, a)); - - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - result->op = GGML_OP_DIV; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -struct ggml_tensor * ggml_div( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_div_impl(ctx, a, b, false); -} - -struct ggml_tensor * ggml_div_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_div_impl(ctx, a, b, true); -} - -// ggml_sqr - -static struct ggml_tensor * ggml_sqr_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - bool inplace) { - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - result->op = GGML_OP_SQR; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_sqr( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_sqr_impl(ctx, a, false); -} - -struct ggml_tensor * ggml_sqr_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_sqr_impl(ctx, a, true); -} - -// ggml_sqrt - -static struct ggml_tensor * ggml_sqrt_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - bool inplace) { - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - result->op = GGML_OP_SQRT; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_sqrt( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_sqrt_impl(ctx, a, false); -} - -struct ggml_tensor * ggml_sqrt_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_sqrt_impl(ctx, a, true); -} - -// ggml_log - -static struct ggml_tensor * ggml_log_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - bool inplace) { - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - result->op = GGML_OP_LOG; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_log( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_log_impl(ctx, a, false); -} - -struct ggml_tensor * ggml_log_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_log_impl(ctx, a, true); -} - -struct ggml_tensor * ggml_expm1( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_EXPM1); -} - -struct ggml_tensor * ggml_expm1_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_EXPM1); -} - -struct ggml_tensor * ggml_softplus( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_SOFTPLUS); -} - -struct ggml_tensor * ggml_softplus_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_SOFTPLUS); -} - -// ggml_sin - -static struct ggml_tensor * ggml_sin_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - bool inplace) { - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - result->op = GGML_OP_SIN; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_sin( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_sin_impl(ctx, a, false); -} - -struct ggml_tensor * ggml_sin_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_sin_impl(ctx, a, true); -} - -// ggml_cos - -static struct ggml_tensor * ggml_cos_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - bool inplace) { - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - result->op = GGML_OP_COS; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_cos( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_cos_impl(ctx, a, false); -} - -struct ggml_tensor * ggml_cos_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_cos_impl(ctx, a, true); -} - -// ggml_sum - -struct ggml_tensor * ggml_sum( - struct ggml_context * ctx, - struct ggml_tensor * a) { - struct ggml_tensor * result = ggml_new_tensor_1d(ctx, a->type, 1); - - result->op = GGML_OP_SUM; - result->src[0] = a; - - return result; -} - -// ggml_sum_rows - -struct ggml_tensor * ggml_sum_rows( - struct ggml_context * ctx, - struct ggml_tensor * a) { - int64_t ne[GGML_MAX_DIMS] = { 1 }; - for (int i = 1; i < GGML_MAX_DIMS; ++i) { - ne[i] = a->ne[i]; - } - - struct ggml_tensor * result = ggml_new_tensor(ctx, a->type, GGML_MAX_DIMS, ne); - - result->op = GGML_OP_SUM_ROWS; - result->src[0] = a; - - return result; -} - -// ggml_cumsum - -struct ggml_tensor * ggml_cumsum( - struct ggml_context * ctx, - struct ggml_tensor * a) { - GGML_ASSERT(a->type == GGML_TYPE_F32); - - struct ggml_tensor * result = ggml_dup_tensor(ctx, a); - - result->op = GGML_OP_CUMSUM; - result->src[0] = a; - - return result; -} - -// ggml_mean - -struct ggml_tensor * ggml_mean( - struct ggml_context * ctx, - struct ggml_tensor * a) { - int64_t ne[4] = { 1, a->ne[1], a->ne[2], a->ne[3] }; - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - - result->op = GGML_OP_MEAN; - result->src[0] = a; - - return result; -} - -// ggml_argmax - -struct ggml_tensor * ggml_argmax( - struct ggml_context * ctx, - struct ggml_tensor * a) { - GGML_ASSERT(ggml_is_matrix(a)); - GGML_ASSERT(a->ne[0] <= INT32_MAX); - - struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_I32, a->ne[1]); - - result->op = GGML_OP_ARGMAX; - result->src[0] = a; - - return result; -} - -// ggml_count_equal - -struct ggml_tensor * ggml_count_equal( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - GGML_ASSERT(ggml_are_same_shape(a, b)); - - struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_I64, 1); - - result->op = GGML_OP_COUNT_EQUAL; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -// ggml_repeat - -struct ggml_tensor * ggml_repeat( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - GGML_ASSERT(ggml_can_repeat(a, b)); - - struct ggml_tensor * result = ggml_new_tensor(ctx, a->type, GGML_MAX_DIMS, b->ne); - - result->op = GGML_OP_REPEAT; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_repeat_4d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3) { - const bool can_repeat = ggml_is_empty(a) || ( - (ne0 % a->ne[0] == 0) && - (ne1 % a->ne[1] == 0) && - (ne2 % a->ne[2] == 0) && - (ne3 % a->ne[3] == 0) - ); - GGML_ASSERT(can_repeat); - - struct ggml_tensor * result = ggml_new_tensor_4d(ctx, a->type, ne0, ne1, ne2, ne3); - - result->op = GGML_OP_REPEAT; - result->src[0] = a; - - return result; -} - -// ggml_repeat_back - -struct ggml_tensor * ggml_repeat_back( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - GGML_ASSERT(ggml_can_repeat(b, a)); - - struct ggml_tensor * result = ggml_new_tensor(ctx, a->type, GGML_MAX_DIMS, b->ne); - - result->op = GGML_OP_REPEAT_BACK; - result->src[0] = a; - - return result; -} - -// ggml_concat - -struct ggml_tensor * ggml_concat( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int dim) { - GGML_ASSERT(dim >= 0 && dim < GGML_MAX_DIMS); - GGML_ASSERT(a->type == b->type); - - int64_t ne[GGML_MAX_DIMS]; - for (int d = 0; d < GGML_MAX_DIMS; ++d) { - if (d == dim) { - ne[d] = a->ne[d] + b->ne[d]; - continue; - } - GGML_ASSERT(a->ne[d] == b->ne[d]); - ne[d] = a->ne[d]; - } - - struct ggml_tensor * result = ggml_new_tensor(ctx, a->type, GGML_MAX_DIMS, ne); - - ggml_set_op_params_i32(result, 0, dim); - - result->op = GGML_OP_CONCAT; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -// ggml_abs - -struct ggml_tensor * ggml_abs( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_ABS); -} - -struct ggml_tensor * ggml_abs_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_ABS); -} - -// ggml_sgn - -struct ggml_tensor * ggml_sgn( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_SGN); -} - -struct ggml_tensor * ggml_sgn_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_SGN); -} - -// ggml_neg - -struct ggml_tensor * ggml_neg( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_NEG); -} - -struct ggml_tensor * ggml_neg_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_NEG); -} - -// ggml_step - -struct ggml_tensor * ggml_step( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_STEP); -} - -struct ggml_tensor * ggml_step_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_STEP); -} - -// ggml_tanh - -struct ggml_tensor * ggml_tanh( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_TANH); -} - -struct ggml_tensor * ggml_tanh_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_TANH); -} - -// ggml_elu - -struct ggml_tensor * ggml_elu( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_ELU); -} - -struct ggml_tensor * ggml_elu_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_ELU); -} - -// ggml_relu - -struct ggml_tensor * ggml_relu( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_RELU); -} - -struct ggml_tensor * ggml_relu_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_RELU); -} - -// ggml_leaky_relu - -struct ggml_tensor * ggml_leaky_relu( - struct ggml_context * ctx, - struct ggml_tensor * a, - float negative_slope, - bool inplace) { - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - ggml_set_op_params(result, &negative_slope, sizeof(negative_slope)); - - result->op = GGML_OP_LEAKY_RELU; - result->src[0] = a; - - return result; -} - -// ggml_sigmoid - -struct ggml_tensor * ggml_sigmoid( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_SIGMOID); -} - -struct ggml_tensor * ggml_sigmoid_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_SIGMOID); -} - -// ggml_gelu - -struct ggml_tensor * ggml_gelu( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_GELU); -} - -struct ggml_tensor * ggml_gelu_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_GELU); -} - -// ggml_gelu_erf - -struct ggml_tensor * ggml_gelu_erf( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_GELU_ERF); -} - -struct ggml_tensor * ggml_gelu_erf_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_GELU_ERF); -} - -// ggml_gelu_quick - -struct ggml_tensor * ggml_gelu_quick( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_GELU_QUICK); -} - -struct ggml_tensor * ggml_gelu_quick_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_GELU_QUICK); -} - -// ggml_silu - -struct ggml_tensor * ggml_silu( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_SILU); -} - -struct ggml_tensor * ggml_silu_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_SILU); -} - -// ggml_xielu - -struct ggml_tensor * ggml_xielu( - struct ggml_context * ctx, - struct ggml_tensor * a, - float alpha_n, - float alpha_p, - float beta, - float eps) { - struct ggml_tensor * result = ggml_dup_tensor(ctx, a); - - ggml_set_op_params_i32(result, 0, (int32_t) GGML_UNARY_OP_XIELU); - ggml_set_op_params_f32(result, 1, beta + ggml_compute_softplus_f32(alpha_n)); - ggml_set_op_params_f32(result, 2, ggml_compute_softplus_f32(alpha_p)); - ggml_set_op_params_f32(result, 3, beta); - ggml_set_op_params_f32(result, 4, eps); - - result->op = GGML_OP_UNARY; - result->src[0] = a; - - return result; -} - -// ggml_silu_back - -struct ggml_tensor * ggml_silu_back( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - struct ggml_tensor * result = ggml_dup_tensor(ctx, a); - - result->op = GGML_OP_SILU_BACK; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -// ggml hardswish - -struct ggml_tensor * ggml_hardswish( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_HARDSWISH); -} - -// ggml hardsigmoid - -struct ggml_tensor * ggml_hardsigmoid( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_HARDSIGMOID); -} - -// ggml exp - -struct ggml_tensor * ggml_exp( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_EXP); -} - -struct ggml_tensor * ggml_exp_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_EXP); -} - -// ggml_glu - -static struct ggml_tensor * ggml_glu_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - enum ggml_glu_op op, - bool swapped) { - GGML_ASSERT(ggml_is_contiguous_1(a)); - - if (b) { - GGML_ASSERT(ggml_is_contiguous_1(b)); - GGML_ASSERT(ggml_are_same_shape(a, b)); - GGML_ASSERT(a->type == b->type); - } - - int64_t ne[GGML_MAX_DIMS] = { a->ne[0] / 2 }; for (int i = 1; i < GGML_MAX_DIMS; i++) ne[i] = a->ne[i]; - struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, GGML_MAX_DIMS, b ? a->ne : ne, NULL, 0); - - ggml_set_op_params_i32(result, 0, (int32_t) op); - ggml_set_op_params_i32(result, 1, (int32_t) swapped); - - result->op = GGML_OP_GLU; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -// ggml_floor - -struct ggml_tensor * ggml_floor( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_FLOOR); -} - -struct ggml_tensor * ggml_floor_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_FLOOR); -} - -// ggml_ceil - -struct ggml_tensor * ggml_ceil( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_CEIL); -} - -struct ggml_tensor * ggml_ceil_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_CEIL); -} - -//ggml_round - -struct ggml_tensor * ggml_round( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_ROUND); -} - -struct ggml_tensor * ggml_round_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_ROUND); -} - -//ggml_trunc - -struct ggml_tensor * ggml_trunc( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary(ctx, a, GGML_UNARY_OP_TRUNC); -} - -struct ggml_tensor * ggml_trunc_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_unary_inplace(ctx, a, GGML_UNARY_OP_TRUNC); -} - -struct ggml_tensor * ggml_glu( - struct ggml_context * ctx, - struct ggml_tensor * a, - enum ggml_glu_op op, - bool swapped) { - return ggml_glu_impl(ctx, a, NULL, op, swapped); -} - -struct ggml_tensor * ggml_glu_split( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - enum ggml_glu_op op) { - return ggml_glu_impl(ctx, a, b, op, false); -} - -// ggml_reglu - -struct ggml_tensor * ggml_reglu( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_REGLU, false); -} - -struct ggml_tensor * ggml_reglu_swapped( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_REGLU, true); -} - -struct ggml_tensor * ggml_reglu_split( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_glu_impl(ctx, a, b, GGML_GLU_OP_REGLU, false); -} - -// ggml_geglu - -struct ggml_tensor * ggml_geglu( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_GEGLU, false); -} - -struct ggml_tensor * ggml_geglu_swapped( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_GEGLU, true); -} - -struct ggml_tensor * ggml_geglu_split( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_glu_impl(ctx, a, b, GGML_GLU_OP_GEGLU, false); -} - -// ggml_swiglu - -struct ggml_tensor * ggml_swiglu( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_SWIGLU, false); -} - -struct ggml_tensor * ggml_swiglu_swapped( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_SWIGLU, true); -} - -struct ggml_tensor * ggml_swiglu_split( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_glu_impl(ctx, a, b, GGML_GLU_OP_SWIGLU, false); -} - -struct ggml_tensor * ggml_swiglu_ds4_split( - struct ggml_context * ctx, - struct ggml_tensor * gate, - struct ggml_tensor * up, - float clamp) { - struct ggml_tensor * result = ggml_glu_impl(ctx, gate, up, GGML_GLU_OP_SWIGLU_DS4, false); - ggml_set_op_params_f32(result, 2, clamp); - return result; -} - -// ggml_geglu_erf - -struct ggml_tensor * ggml_geglu_erf( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_GEGLU_ERF, false); -} - -struct ggml_tensor * ggml_geglu_erf_swapped( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_GEGLU_ERF, true); -} - -struct ggml_tensor * ggml_geglu_erf_split( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_glu_impl(ctx, a, b, GGML_GLU_OP_GEGLU_ERF, false); -} - -// ggml_geglu_quick - -struct ggml_tensor * ggml_geglu_quick( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_GEGLU_QUICK, false); -} - -struct ggml_tensor * ggml_geglu_quick_swapped( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_glu_impl(ctx, a, NULL, GGML_GLU_OP_GEGLU_QUICK, true); -} - -struct ggml_tensor * ggml_geglu_quick_split( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_glu_impl(ctx, a, b, GGML_GLU_OP_GEGLU_QUICK, false); -} - -struct ggml_tensor * ggml_swiglu_oai( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - float alpha, - float limit) { - struct ggml_tensor * result = ggml_glu_impl(ctx, a, b, GGML_GLU_OP_SWIGLU_OAI, false); - ggml_set_op_params_f32(result, 2, alpha); - ggml_set_op_params_f32(result, 3, limit); - - return result; -} - -// ggml_norm - -static struct ggml_tensor * ggml_norm_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - float eps, - bool inplace) { - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - ggml_set_op_params(result, &eps, sizeof(eps)); - - result->op = GGML_OP_NORM; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_norm( - struct ggml_context * ctx, - struct ggml_tensor * a, - float eps) { - return ggml_norm_impl(ctx, a, eps, false); -} - -struct ggml_tensor * ggml_norm_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - float eps) { - return ggml_norm_impl(ctx, a, eps, true); -} - -// ggml_rms_norm - -static struct ggml_tensor * ggml_rms_norm_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - float eps, - bool inplace) { - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - ggml_set_op_params(result, &eps, sizeof(eps)); - - result->op = GGML_OP_RMS_NORM; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_rms_norm( - struct ggml_context * ctx, - struct ggml_tensor * a, - float eps) { - return ggml_rms_norm_impl(ctx, a, eps, false); -} - -struct ggml_tensor * ggml_rms_norm_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - float eps) { - return ggml_rms_norm_impl(ctx, a, eps, true); -} - -// ggml_rms_norm_back - -struct ggml_tensor * ggml_rms_norm_back( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - float eps) { - struct ggml_tensor * result = ggml_dup_tensor(ctx, a); - - ggml_set_op_params(result, &eps, sizeof(eps)); - - result->op = GGML_OP_RMS_NORM_BACK; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -// ggml_group_norm - -static struct ggml_tensor * ggml_group_norm_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - int n_groups, - float eps, - bool inplace) { - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - ggml_set_op_params_i32(result, 0, n_groups); - ggml_set_op_params_f32(result, 1, eps); - - result->op = GGML_OP_GROUP_NORM; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_group_norm( - struct ggml_context * ctx, - struct ggml_tensor * a, - int n_groups, - float eps) { - return ggml_group_norm_impl(ctx, a, n_groups, eps, false); -} - -struct ggml_tensor * ggml_group_norm_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - int n_groups, - float eps) { - return ggml_group_norm_impl(ctx, a, n_groups, eps, true); -} - -// ggml_l2_norm - -static struct ggml_tensor * ggml_l2_norm_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - float eps, - bool inplace) { - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - ggml_set_op_params_f32(result, 0, eps); - - result->op = GGML_OP_L2_NORM; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_l2_norm( - struct ggml_context * ctx, - struct ggml_tensor * a, - float eps) { - return ggml_l2_norm_impl(ctx, a, eps, false); -} - -struct ggml_tensor * ggml_l2_norm_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - float eps) { - return ggml_l2_norm_impl(ctx, a, eps, true); -} - -// ggml_mul_mat - -static inline bool ggml_can_mul_mat(const struct ggml_tensor * t0, const struct ggml_tensor * t1) { - static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); - - return (t0->ne[0] == t1->ne[0]) && - (t1->ne[2]%t0->ne[2] == 0) && // verify t0 is broadcastable - (t1->ne[3]%t0->ne[3] == 0); -} - -struct ggml_tensor * ggml_mul_mat( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - GGML_ASSERT(ggml_can_mul_mat(a, b)); - GGML_ASSERT(!ggml_is_transposed(a)); - - const int64_t ne[4] = { a->ne[1], b->ne[1], b->ne[2], b->ne[3] }; - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - - result->op = GGML_OP_MUL_MAT; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -void ggml_mul_mat_set_prec( - struct ggml_tensor * a, - enum ggml_prec prec) { - GGML_ASSERT(a->op == GGML_OP_MUL_MAT); - - const int32_t prec_i32 = (int32_t) prec; - - ggml_set_op_params_i32(a, 0, prec_i32); -} - -// ggml_mul_mat_id - -/* - c = ggml_mul_mat_id(ctx, as, b, ids); - - as -> [cols, rows, n_expert] - b -> [cols, n_expert_used, n_tokens] - ids -> [n_expert_used, n_tokens] (i32) - c -> [rows, n_expert_used, n_tokens] - - in b, n_expert_used can be broadcasted to match the n_expert_used of ids - - c ~= as[:,:,i] @ b[:,i%r,t], i = ids[e,t] for all e,t in ids -*/ -struct ggml_tensor * ggml_mul_mat_id( - struct ggml_context * ctx, - struct ggml_tensor * as, - struct ggml_tensor * b, - struct ggml_tensor * ids) { - GGML_ASSERT(!ggml_is_transposed(as)); - GGML_ASSERT(ids->type == GGML_TYPE_I32); - - GGML_ASSERT(as->ne[3] == 1); // as is 3d (one matrix per expert) - GGML_ASSERT(b->ne[3] == 1); // b is 3d - GGML_ASSERT(ids->ne[2] == 1 && ids->ne[3] == 1); // ids is 2d - GGML_ASSERT(ids->ne[1] == b->ne[2]); // must have an expert list per b row - GGML_ASSERT(as->ne[0] == b->ne[0]); // can_mul_mat - GGML_ASSERT(ids->ne[0] % b->ne[1] == 0); // can broadcast - - const int64_t ne[4] = { as->ne[1], ids->ne[0], b->ne[2], 1 }; - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - - result->op = GGML_OP_MUL_MAT_ID; - result->src[0] = as; - result->src[1] = b; - result->src[2] = ids; - - return result; -} - -// ggml_out_prod - -static inline bool ggml_can_out_prod(const struct ggml_tensor * t0, const struct ggml_tensor * t1) { - static_assert(GGML_MAX_DIMS == 4, "GGML_MAX_DIMS is not 4 - update this function"); - - return (t0->ne[1] == t1->ne[1]) && - (t1->ne[2]%t0->ne[2] == 0) && // verify t0 is broadcastable - (t1->ne[3]%t0->ne[3] == 0); -} - -struct ggml_tensor * ggml_out_prod( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - GGML_ASSERT(ggml_can_out_prod(a, b)); - GGML_ASSERT(!ggml_is_transposed(a)); - - // a is broadcastable to b for ne[2] and ne[3] -> use b->ne[2] and b->ne[3] - const int64_t ne[4] = { a->ne[0], b->ne[0], b->ne[2], b->ne[3] }; - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - - result->op = GGML_OP_OUT_PROD; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -// ggml_scale - -static struct ggml_tensor * ggml_scale_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - float s, - float b, - bool inplace) { - GGML_ASSERT(ggml_is_padded_1d(a)); - - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - float params[2] = { s, b }; - ggml_set_op_params(result, ¶ms, sizeof(params)); - - result->op = GGML_OP_SCALE; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_scale( - struct ggml_context * ctx, - struct ggml_tensor * a, - float s) { - return ggml_scale_impl(ctx, a, s, 0.0, false); -} - -struct ggml_tensor * ggml_scale_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - float s) { - return ggml_scale_impl(ctx, a, s, 0.0, true); -} - -struct ggml_tensor * ggml_scale_bias( - struct ggml_context * ctx, - struct ggml_tensor * a, - float s, - float b) { - return ggml_scale_impl(ctx, a, s, b, false); -} - -struct ggml_tensor * ggml_scale_bias_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - float s, - float b) { - return ggml_scale_impl(ctx, a, s, b, true); -} - -// ggml_set - -static struct ggml_tensor * ggml_set_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - size_t nb1, - size_t nb2, - size_t nb3, - size_t offset, - bool inplace) { - GGML_ASSERT(ggml_nelements(a) >= ggml_nelements(b)); - - // make a view of the destination - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - GGML_ASSERT(offset < (size_t)(1 << 30)); - int32_t params[] = { nb1, nb2, nb3, offset, inplace ? 1 : 0 }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_SET; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -struct ggml_tensor * ggml_set( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - size_t nb1, - size_t nb2, - size_t nb3, - size_t offset) { - return ggml_set_impl(ctx, a, b, nb1, nb2, nb3, offset, false); -} - -struct ggml_tensor * ggml_set_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - size_t nb1, - size_t nb2, - size_t nb3, - size_t offset) { - return ggml_set_impl(ctx, a, b, nb1, nb2, nb3, offset, true); -} - -struct ggml_tensor * ggml_set_1d( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - size_t offset) { - return ggml_set_impl(ctx, a, b, a->nb[1], a->nb[2], a->nb[3], offset, false); -} - -struct ggml_tensor * ggml_set_1d_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - size_t offset) { - return ggml_set_impl(ctx, a, b, a->nb[1], a->nb[2], a->nb[3], offset, true); -} - -struct ggml_tensor * ggml_set_2d( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - size_t nb1, - size_t offset) { - return ggml_set_impl(ctx, a, b, nb1, a->nb[2], a->nb[3], offset, false); -} - -struct ggml_tensor * ggml_set_2d_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - size_t nb1, - size_t offset) { - return ggml_set_impl(ctx, a, b, nb1, a->nb[2], a->nb[3], offset, true); -} - -// ggml_cpy - -static struct ggml_tensor * ggml_cpy_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - GGML_ASSERT(ggml_nelements(a) == ggml_nelements(b)); - - // make a view of the destination - struct ggml_tensor * result = ggml_view_tensor(ctx, b); - if (strlen(b->name) > 0) { - ggml_format_name(result, "%s (copy of %s)", b->name, a->name); - } else { - ggml_format_name(result, "%s (copy)", a->name); - } - - result->op = GGML_OP_CPY; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -struct ggml_tensor * ggml_cpy( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_cpy_impl(ctx, a, b); -} - -struct ggml_tensor * ggml_cast( - struct ggml_context * ctx, - struct ggml_tensor * a, - enum ggml_type type) { - struct ggml_tensor * result = ggml_new_tensor(ctx, type, GGML_MAX_DIMS, a->ne); - ggml_format_name(result, "%s (copy)", a->name); - - result->op = GGML_OP_CPY; - result->src[0] = a; - result->src[1] = result; // note: this self-reference might seem redundant, but it's actually needed by some - // backends for consistency with ggml_cpy_impl() above - - return result; -} - -// ggml_cont - -static struct ggml_tensor * ggml_cont_impl( - struct ggml_context * ctx, - struct ggml_tensor * a) { - struct ggml_tensor * result = ggml_dup_tensor(ctx, a); - ggml_format_name(result, "%s (cont)", a->name); - - result->op = GGML_OP_CONT; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_cont( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_cont_impl(ctx, a); -} - -// make contiguous, with new shape -GGML_API struct ggml_tensor * ggml_cont_1d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0) { - return ggml_cont_4d(ctx, a, ne0, 1, 1, 1); -} - -GGML_API struct ggml_tensor * ggml_cont_2d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - int64_t ne1) { - return ggml_cont_4d(ctx, a, ne0, ne1, 1, 1); -} - -GGML_API struct ggml_tensor * ggml_cont_3d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - int64_t ne1, - int64_t ne2) { - return ggml_cont_4d(ctx, a, ne0, ne1, ne2, 1); -} - -struct ggml_tensor * ggml_cont_4d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - int64_t ne1, - int64_t ne2, - int64_t ne3) { - GGML_ASSERT(ggml_nelements(a) == (ne0*ne1*ne2*ne3)); - - struct ggml_tensor * result = ggml_new_tensor_4d(ctx, a->type, ne0, ne1, ne2, ne3); - ggml_format_name(result, "%s (cont)", a->name); - - result->op = GGML_OP_CONT; - result->src[0] = a; - - return result; -} - -// ggml_reshape - -struct ggml_tensor * ggml_reshape( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - GGML_ASSERT(ggml_is_contiguous(a)); - // as only the shape of b is relevant, and not its memory layout, b is allowed to be non contiguous. - GGML_ASSERT(ggml_nelements(a) == ggml_nelements(b)); - - struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, GGML_MAX_DIMS, b->ne, a, 0); - ggml_format_name(result, "%s (reshaped)", a->name); - - result->op = GGML_OP_RESHAPE; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_reshape_1d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0) { - GGML_ASSERT(ggml_is_contiguous(a)); - GGML_ASSERT(ggml_nelements(a) == ne0); - - const int64_t ne[1] = { ne0 }; - struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, 1, ne, a, 0); - ggml_format_name(result, "%s (reshaped)", a->name); - - result->op = GGML_OP_RESHAPE; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_reshape_2d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - int64_t ne1) { - GGML_ASSERT(ggml_is_contiguous(a)); - GGML_ASSERT(ggml_nelements(a) == ne0*ne1); - - const int64_t ne[2] = { ne0, ne1 }; - struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, 2, ne, a, 0); - ggml_format_name(result, "%s (reshaped)", a->name); - - result->op = GGML_OP_RESHAPE; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_reshape_3d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - int64_t ne1, - int64_t ne2) { - GGML_ASSERT(ggml_is_contiguous(a)); - GGML_ASSERT(ggml_nelements(a) == ne0*ne1*ne2); - - const int64_t ne[3] = { ne0, ne1, ne2 }; - struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, 3, ne, a, 0); - ggml_format_name(result, "%s (reshaped)", a->name); - - result->op = GGML_OP_RESHAPE; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_reshape_4d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - int64_t ne1, - int64_t ne2, - int64_t ne3) { - GGML_ASSERT(ggml_is_contiguous(a)); - GGML_ASSERT(ggml_nelements(a) == ne0*ne1*ne2*ne3); - - const int64_t ne[4] = { ne0, ne1, ne2, ne3 }; - struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, 4, ne, a, 0); - ggml_format_name(result, "%s (reshaped)", a->name); - - result->op = GGML_OP_RESHAPE; - result->src[0] = a; - - return result; -} - -static struct ggml_tensor * ggml_view_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - int n_dims, - const int64_t * ne, - size_t offset) { - struct ggml_tensor * result = ggml_new_tensor_impl(ctx, a->type, n_dims, ne, a, offset); - ggml_format_name(result, "%s (view)", a->name); - - ggml_set_op_params(result, &offset, sizeof(offset)); - - result->op = GGML_OP_VIEW; - result->src[0] = a; - - return result; -} - -// ggml_view_1d - -struct ggml_tensor * ggml_view_1d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - size_t offset) { - struct ggml_tensor * result = ggml_view_impl(ctx, a, 1, &ne0, offset); - - return result; -} - -// ggml_view_2d - -struct ggml_tensor * ggml_view_2d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - int64_t ne1, - size_t nb1, - size_t offset) { - const int64_t ne[2] = { ne0, ne1 }; - - struct ggml_tensor * result = ggml_view_impl(ctx, a, 2, ne, offset); - - result->nb[1] = nb1; - result->nb[2] = result->nb[1]*ne1; - result->nb[3] = result->nb[2]; - - return result; -} - -// ggml_view_3d - -struct ggml_tensor * ggml_view_3d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - int64_t ne1, - int64_t ne2, - size_t nb1, - size_t nb2, - size_t offset) { - const int64_t ne[3] = { ne0, ne1, ne2 }; - - struct ggml_tensor * result = ggml_view_impl(ctx, a, 3, ne, offset); - - result->nb[1] = nb1; - result->nb[2] = nb2; - result->nb[3] = result->nb[2]*ne2; - - return result; -} - -// ggml_view_4d - -struct ggml_tensor * ggml_view_4d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - int64_t ne1, - int64_t ne2, - int64_t ne3, - size_t nb1, - size_t nb2, - size_t nb3, - size_t offset) { - const int64_t ne[4] = { ne0, ne1, ne2, ne3 }; - - struct ggml_tensor * result = ggml_view_impl(ctx, a, 4, ne, offset); - - result->nb[1] = nb1; - result->nb[2] = nb2; - result->nb[3] = nb3; - - return result; -} - -// ggml_permute - -struct ggml_tensor * ggml_permute( - struct ggml_context * ctx, - struct ggml_tensor * a, - int axis0, - int axis1, - int axis2, - int axis3) { - GGML_ASSERT(axis0 >= 0 && axis0 < GGML_MAX_DIMS); - GGML_ASSERT(axis1 >= 0 && axis1 < GGML_MAX_DIMS); - GGML_ASSERT(axis2 >= 0 && axis2 < GGML_MAX_DIMS); - GGML_ASSERT(axis3 >= 0 && axis3 < GGML_MAX_DIMS); - - GGML_ASSERT(axis0 != axis1); - GGML_ASSERT(axis0 != axis2); - GGML_ASSERT(axis0 != axis3); - GGML_ASSERT(axis1 != axis2); - GGML_ASSERT(axis1 != axis3); - GGML_ASSERT(axis2 != axis3); - - struct ggml_tensor * result = ggml_view_tensor(ctx, a); - ggml_format_name(result, "%s (permuted)", a->name); - - int ne[GGML_MAX_DIMS]; - int nb[GGML_MAX_DIMS]; - - ne[axis0] = a->ne[0]; - ne[axis1] = a->ne[1]; - ne[axis2] = a->ne[2]; - ne[axis3] = a->ne[3]; - - nb[axis0] = a->nb[0]; - nb[axis1] = a->nb[1]; - nb[axis2] = a->nb[2]; - nb[axis3] = a->nb[3]; - - result->ne[0] = ne[0]; - result->ne[1] = ne[1]; - result->ne[2] = ne[2]; - result->ne[3] = ne[3]; - - result->nb[0] = nb[0]; - result->nb[1] = nb[1]; - result->nb[2] = nb[2]; - result->nb[3] = nb[3]; - - result->op = GGML_OP_PERMUTE; - result->src[0] = a; - - int32_t params[] = { axis0, axis1, axis2, axis3 }; - ggml_set_op_params(result, params, sizeof(params)); - - return result; -} - -// ggml_transpose - -struct ggml_tensor * ggml_transpose( - struct ggml_context * ctx, - struct ggml_tensor * a) { - struct ggml_tensor * result = ggml_view_tensor(ctx, a); - ggml_format_name(result, "%s (transposed)", a->name); - - result->ne[0] = a->ne[1]; - result->ne[1] = a->ne[0]; - - result->nb[0] = a->nb[1]; - result->nb[1] = a->nb[0]; - - result->op = GGML_OP_TRANSPOSE; - result->src[0] = a; - - return result; -} - -// ggml_get_rows - -struct ggml_tensor * ggml_get_rows( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - GGML_ASSERT(a->ne[2] == b->ne[1]); - GGML_ASSERT(a->ne[3] == b->ne[2]); - GGML_ASSERT(b->ne[3] == 1); - GGML_ASSERT(b->type == GGML_TYPE_I32); - - // TODO: implement non F32 return - enum ggml_type type = GGML_TYPE_F32; - if (a->type == GGML_TYPE_I32) { - type = a->type; - } - struct ggml_tensor * result = ggml_new_tensor_4d(ctx, type, a->ne[0], b->ne[0], b->ne[1], b->ne[2]); - - result->op = GGML_OP_GET_ROWS; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -// ggml_get_rows_back - -struct ggml_tensor * ggml_get_rows_back( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * c) { - GGML_ASSERT(ggml_is_matrix(a) && ggml_is_vector(b) && b->type == GGML_TYPE_I32); - GGML_ASSERT(ggml_is_matrix(c) && (a->ne[0] == c->ne[0])); - - // TODO: implement non F32 return - //struct ggml_tensor * result = ggml_new_tensor_2d(ctx, a->type, a->ne[0], b->ne[0]); - struct ggml_tensor * result = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, c->ne[0], c->ne[1]); - - result->op = GGML_OP_GET_ROWS_BACK; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -// ggml_set_rows - -struct ggml_tensor * ggml_set_rows( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * c) { - GGML_ASSERT(a->ne[0] == b->ne[0]); - GGML_ASSERT(a->ne[2] == b->ne[2]); - GGML_ASSERT(a->ne[3] == b->ne[3]); - GGML_ASSERT(b->ne[1] == c->ne[0]); - GGML_ASSERT(b->ne[2] % c->ne[1] == 0); - GGML_ASSERT(b->ne[3] % c->ne[2] == 0); - GGML_ASSERT(c->ne[3] == 1); - GGML_ASSERT(b->type == GGML_TYPE_F32); - GGML_ASSERT(c->type == GGML_TYPE_I64 || c->type == GGML_TYPE_I32); - - GGML_ASSERT(ggml_is_contiguous_rows(a)); - GGML_ASSERT(ggml_is_contiguous_rows(b)); - - struct ggml_tensor * result = ggml_view_tensor(ctx, a); - - result->op = GGML_OP_SET_ROWS; - result->src[0] = b; - result->src[1] = c; - result->src[2] = a; // note: order is weird due to legacy reasons (https://github.com/ggml-org/llama.cpp/pull/16063#discussion_r2385795931) - - return result; -} - -// ggml_diag - -struct ggml_tensor * ggml_diag( - struct ggml_context * ctx, - struct ggml_tensor * a) { - GGML_ASSERT(a->ne[1] == 1); - - const int64_t ne[4] = { a->ne[0], a->ne[0], a->ne[2], a->ne[3] }; - struct ggml_tensor * result = ggml_new_tensor(ctx, a->type, 4, ne); - - result->op = GGML_OP_DIAG; - result->src[0] = a; - - return result; -} - -// ggml_diag_mask_inf - -static struct ggml_tensor * ggml_diag_mask_inf_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - int n_past, - bool inplace) { - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - int32_t params[] = { n_past }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_DIAG_MASK_INF; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_diag_mask_inf( - struct ggml_context * ctx, - struct ggml_tensor * a, - int n_past) { - return ggml_diag_mask_inf_impl(ctx, a, n_past, false); -} - -struct ggml_tensor * ggml_diag_mask_inf_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - int n_past) { - return ggml_diag_mask_inf_impl(ctx, a, n_past, true); -} - -// ggml_diag_mask_zero - -static struct ggml_tensor * ggml_diag_mask_zero_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - int n_past, - bool inplace) { - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - int32_t params[] = { n_past }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_DIAG_MASK_ZERO; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_diag_mask_zero( - struct ggml_context * ctx, - struct ggml_tensor * a, - int n_past) { - return ggml_diag_mask_zero_impl(ctx, a, n_past, false); -} - -struct ggml_tensor * ggml_diag_mask_zero_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - int n_past) { - return ggml_diag_mask_zero_impl(ctx, a, n_past, true); -} - -// ggml_soft_max - -static struct ggml_tensor * ggml_soft_max_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * mask, - float scale, - float max_bias, - bool inplace) { - GGML_ASSERT(ggml_is_contiguous(a)); - - if (mask) { - GGML_ASSERT(mask->type == GGML_TYPE_F16 || mask->type == GGML_TYPE_F32); - GGML_ASSERT(ggml_is_contiguous(mask)); - GGML_ASSERT(mask->ne[0] == a->ne[0]); - GGML_ASSERT(mask->ne[1] >= a->ne[1]); - GGML_ASSERT(a->ne[2]%mask->ne[2] == 0); - GGML_ASSERT(a->ne[3]%mask->ne[3] == 0); - } - - if (max_bias > 0.0f) { - GGML_ASSERT(mask); - } - - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - float params[] = { scale, max_bias }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_SOFT_MAX; - result->src[0] = a; - result->src[1] = mask; - - return result; -} - -struct ggml_tensor * ggml_soft_max( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_soft_max_impl(ctx, a, NULL, 1.0f, 0.0f, false); -} - -struct ggml_tensor * ggml_soft_max_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a) { - return ggml_soft_max_impl(ctx, a, NULL, 1.0f, 0.0f, true); -} - -struct ggml_tensor * ggml_soft_max_ext( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * mask, - float scale, - float max_bias) { - return ggml_soft_max_impl(ctx, a, mask, scale, max_bias, false); -} - -struct ggml_tensor * ggml_soft_max_ext_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * mask, - float scale, - float max_bias) { - return ggml_soft_max_impl(ctx, a, mask, scale, max_bias, true); -} - -void ggml_soft_max_add_sinks( - struct ggml_tensor * a, - struct ggml_tensor * sinks) { - if (!sinks) { - a->src[2] = NULL; - return; - } - - GGML_ASSERT(a->op == GGML_OP_SOFT_MAX); - GGML_ASSERT(a->src[2] == NULL); - GGML_ASSERT(a->src[0]->ne[2] == sinks->ne[0]); - GGML_ASSERT(sinks->type == GGML_TYPE_F32); - - a->src[2] = sinks; -} - -// ggml_soft_max_ext_back - -static struct ggml_tensor * ggml_soft_max_ext_back_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - float scale, - float max_bias, - bool inplace) { - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - result->op = GGML_OP_SOFT_MAX_BACK; - result->src[0] = a; - result->src[1] = b; - - memcpy((float *) result->op_params + 0, &scale, sizeof(float)); - memcpy((float *) result->op_params + 1, &max_bias, sizeof(float)); - - return result; -} - -struct ggml_tensor * ggml_soft_max_ext_back( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - float scale, - float max_bias) { - return ggml_soft_max_ext_back_impl(ctx, a, b, scale, max_bias, false); -} - -struct ggml_tensor * ggml_soft_max_ext_back_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - float scale, - float max_bias) { - return ggml_soft_max_ext_back_impl(ctx, a, b, scale, max_bias, true); -} - -// ggml_rope - -static struct ggml_tensor * ggml_rope_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * c, - int n_dims, - int sections[GGML_MROPE_SECTIONS], - int mode, - int n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow, - bool inplace) { - GGML_ASSERT((mode & 1) == 0 && "mode & 1 == 1 is no longer supported"); - - GGML_ASSERT(ggml_is_vector(b)); - GGML_ASSERT(b->type == GGML_TYPE_I32); - - bool mrope_used = mode & GGML_ROPE_TYPE_MROPE; - if (mrope_used) { - GGML_ASSERT(a->ne[2] * 4 == b->ne[0]); // mrope expecting 4 position ids per token - } else { - GGML_ASSERT(a->ne[2] == b->ne[0]); - } - - if (c) { - GGML_ASSERT(c->type == GGML_TYPE_F32); - GGML_ASSERT(c->ne[0] >= n_dims / 2); - } - - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - int32_t params[15] = { /*n_past*/ 0, n_dims, mode, /*n_ctx*/ 0, n_ctx_orig }; - memcpy(params + 5, &freq_base, sizeof(float)); - memcpy(params + 6, &freq_scale, sizeof(float)); - memcpy(params + 7, &ext_factor, sizeof(float)); - memcpy(params + 8, &attn_factor, sizeof(float)); - memcpy(params + 9, &beta_fast, sizeof(float)); - memcpy(params + 10, &beta_slow, sizeof(float)); - if (mrope_used && sections) { - memcpy(params + 11, sections, sizeof(int32_t) * GGML_MROPE_SECTIONS); - } else { - memset(params + 11, 0, sizeof(int32_t) * GGML_MROPE_SECTIONS); - } - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_ROPE; - result->src[0] = a; - result->src[1] = b; - result->src[2] = c; - - return result; -} - -struct ggml_tensor * ggml_rope( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int n_dims, - int mode) { - return ggml_rope_impl( - ctx, a, b, NULL, n_dims, NULL, mode, 0, 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f, false - ); -} - -struct ggml_tensor * ggml_rope_multi( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * c, - int n_dims, - int sections[GGML_MROPE_SECTIONS], - int mode, - int n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow) { - return ggml_rope_impl( - ctx, a, b, c, n_dims, sections, mode, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow, false - ); -} - -struct ggml_tensor * ggml_rope_multi_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * c, - int n_dims, - int sections[GGML_MROPE_SECTIONS], - int mode, - int n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow) { - return ggml_rope_impl( - ctx, a, b, c, n_dims, sections, mode, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow, true - ); -} - -struct ggml_tensor * ggml_rope_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int n_dims, - int mode) { - return ggml_rope_impl( - ctx, a, b, NULL, n_dims, NULL, mode, 0, 10000.0f, 1.0f, 0.0f, 1.0f, 0.0f, 0.0f, true - ); -} - -struct ggml_tensor * ggml_rope_ext( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * c, - int n_dims, - int mode, - int n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow) { - return ggml_rope_impl( - ctx, a, b, c, n_dims, NULL, mode, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow, false - ); -} - -struct ggml_tensor * ggml_rope_ext_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * c, - int n_dims, - int mode, - int n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow) { - return ggml_rope_impl( - ctx, a, b, c, n_dims, NULL, mode, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow, true - ); -} - -struct ggml_tensor * ggml_rope_custom( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int n_dims, - int mode, - int n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow) { - return ggml_rope_impl( - ctx, a, b, NULL, n_dims, NULL, mode, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow, false - ); -} - -struct ggml_tensor * ggml_rope_custom_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int n_dims, - int mode, - int n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow) { - return ggml_rope_impl( - ctx, a, b, NULL, n_dims, NULL, mode, n_ctx_orig, freq_base, freq_scale, - ext_factor, attn_factor, beta_fast, beta_slow, true - ); -} - -// Apparently solving `n_rot = 2pi * x * base^((2 * max_pos_emb) / n_dims)` for x, we get -// `corr_dim(n_rot) = n_dims * log(max_pos_emb / (n_rot * 2pi)) / (2 * log(base))` -static float ggml_rope_yarn_corr_dim(int n_dims, int n_ctx_orig, float n_rot, float base) { - return n_dims * logf(n_ctx_orig / (n_rot * 2 * (float)M_PI)) / (2 * logf(base)); -} - -void ggml_rope_yarn_corr_dims( - int n_dims, int n_ctx_orig, float freq_base, float beta_fast, float beta_slow, float dims[2] -) { - // start and end correction dims - float start = floorf(ggml_rope_yarn_corr_dim(n_dims, n_ctx_orig, beta_fast, freq_base)); - float end = ceilf(ggml_rope_yarn_corr_dim(n_dims, n_ctx_orig, beta_slow, freq_base)); - dims[0] = MAX(0, start); - dims[1] = MIN(n_dims - 1, end); -} - -// ggml_rope_back - -struct ggml_tensor * ggml_rope_ext_back( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * c, - int n_dims, - int mode, - int n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow) { - struct ggml_tensor * result = ggml_rope_ext( - ctx, a, b, c, n_dims, mode, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); - result->op = GGML_OP_ROPE_BACK; - return result; -} - -struct ggml_tensor * ggml_rope_multi_back( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * c, - int n_dims, - int sections[4], - int mode, - int n_ctx_orig, - float freq_base, - float freq_scale, - float ext_factor, - float attn_factor, - float beta_fast, - float beta_slow) { - struct ggml_tensor * result = ggml_rope_multi( - ctx, a, b, c, n_dims, sections, mode, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); - result->op = GGML_OP_ROPE_BACK; - return result; -} -// ggml_clamp - -struct ggml_tensor * ggml_clamp( - struct ggml_context * ctx, - struct ggml_tensor * a, - float min, - float max) { - // TODO: when implement backward, fix this: - struct ggml_tensor * result = ggml_view_tensor(ctx, a); - - float params[] = { min, max }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_CLAMP; - result->src[0] = a; - - return result; -} - -static int64_t ggml_calc_conv_output_size(int64_t ins, int64_t ks, int s, int p, int d) { - return (ins + 2 * p - d * (ks - 1) - 1) / s + 1; -} - -// im2col: [N, IC, IH, IW] => [N, OH, OW, IC*KH*KW] -// a: [OC，IC, KH, KW] -// b: [N, IC, IH, IW] -// result: [N, OH, OW, IC*KH*KW] -struct ggml_tensor * ggml_im2col( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int s0, - int s1, - int p0, - int p1, - int d0, - int d1, - bool is_2D, - enum ggml_type dst_type) { - if (is_2D) { - GGML_ASSERT(a->ne[2] == b->ne[2]); - } else { - //GGML_ASSERT(b->ne[1] % a->ne[1] == 0); - GGML_ASSERT(b->ne[1] == a->ne[1]); - GGML_ASSERT(b->ne[3] == 1); - } - - const int64_t OH = is_2D ? ggml_calc_conv_output_size(b->ne[1], a->ne[1], s1, p1, d1) : 0; - const int64_t OW = ggml_calc_conv_output_size(b->ne[0], a->ne[0], s0, p0, d0); - - GGML_ASSERT((!is_2D || OH > 0) && "b too small compared to a"); - GGML_ASSERT((OW > 0) && "b too small compared to a"); - - const int64_t ne[4] = { - is_2D ? (a->ne[2] * a->ne[1] * a->ne[0]) : a->ne[1] * a->ne[0], - OW, - is_2D ? OH : b->ne[2], - is_2D ? b->ne[3] : 1, - }; - - struct ggml_tensor * result = ggml_new_tensor(ctx, dst_type, 4, ne); - int32_t params[] = { s0, s1, p0, p1, d0, d1, (is_2D ? 1 : 0) }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_IM2COL; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -struct ggml_tensor * ggml_im2col_back( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int64_t * ne, - int s0, - int s1, - int p0, - int p1, - int d0, - int d1, - bool is_2D) { - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - int32_t params[] = { s0, s1, p0, p1, d0, d1, (is_2D ? 1 : 0) }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_IM2COL_BACK; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -// ggml_conv_1d - -struct ggml_tensor * ggml_conv_1d( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int s0, - int p0, - int d0) { - struct ggml_tensor * im2col = ggml_im2col(ctx, a, b, s0, 0, p0, 0, d0, 0, false, GGML_TYPE_F16); // [N, OL, IC * K] - - struct ggml_tensor * result = - ggml_mul_mat(ctx, - ggml_reshape_2d(ctx, im2col, im2col->ne[0], (im2col->ne[2] * im2col->ne[1])), // [N, OL, IC * K] => [N*OL, IC * K] - ggml_reshape_2d(ctx, a, (a->ne[0] * a->ne[1]), a->ne[2])); // [OC，IC, K] => [OC, IC * K] - - result = ggml_reshape_3d(ctx, result, im2col->ne[1], a->ne[2], im2col->ne[2]); // [N, OC, OL] - - return result; -} - -// ggml_conv_1d_ph - -struct ggml_tensor* ggml_conv_1d_ph( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int s, - int d) { - return ggml_conv_1d(ctx, a, b, s, a->ne[0] / 2, d); -} - -// ggml_conv_1d_dw - -struct ggml_tensor * ggml_conv_1d_dw( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int s0, - int p0, - int d0) { - struct ggml_tensor * new_b = ggml_reshape_4d(ctx, b, b->ne[0], 1, b->ne[1], b->ne[2]); - - struct ggml_tensor * im2col = ggml_im2col(ctx, a, new_b, s0, 0, p0, 0, d0, 0, false, GGML_TYPE_F16); - - struct ggml_tensor * result = ggml_mul_mat(ctx, im2col, a); - - result = ggml_reshape_3d(ctx, result, result->ne[0], result->ne[2], 1); - - return result; -} - -// ggml_conv_1d_dw_ph - -struct ggml_tensor * ggml_conv_1d_dw_ph( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int s0, - int d0) { - return ggml_conv_1d_dw(ctx, a, b, s0, a->ne[0] / 2, d0); -} - -// ggml_conv_transpose_1d - -static int64_t ggml_calc_conv_transpose_1d_output_size(int64_t ins, int64_t ks, int s, int p, int d) { - return (ins - 1) * s - 2 * p + d * (ks - 1) + 1; -} - -GGML_API struct ggml_tensor * ggml_conv_transpose_1d( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int s0, - int p0, - int d0) { - GGML_ASSERT(ggml_is_matrix(b)); - GGML_ASSERT(a->ne[2] == b->ne[1]); - GGML_ASSERT(a->ne[3] == 1); - - GGML_ASSERT(p0 == 0); - GGML_ASSERT(d0 == 1); - - const int64_t ne[4] = { - ggml_calc_conv_transpose_1d_output_size(b->ne[0], a->ne[0], s0, 0 /*p0*/, 1 /*d0*/), - a->ne[1], b->ne[2], 1, - }; - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - - int32_t params[] = { s0, p0, d0 }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_CONV_TRANSPOSE_1D; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -// ggml_conv_2d - -// a: [OC，IC, KH, KW] -// b: [N, IC, IH, IW] -// result: [N, OC, OH, OW] -struct ggml_tensor * ggml_conv_2d( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int s0, - int s1, - int p0, - int p1, - int d0, - int d1) { - struct ggml_tensor * im2col = ggml_im2col(ctx, a, b, s0, s1, p0, p1, d0, d1, true, a->type); // [N, OH, OW, IC * KH * KW] - - struct ggml_tensor * result = - ggml_mul_mat(ctx, - ggml_reshape_2d(ctx, im2col, im2col->ne[0], im2col->ne[3] * im2col->ne[2] * im2col->ne[1]), // [N, OH, OW, IC * KH * KW] => [N*OH*OW, IC * KH * KW] - ggml_reshape_2d(ctx, a, (a->ne[0] * a->ne[1] * a->ne[2]), a->ne[3])); // [OC，IC, KH, KW] => [OC, IC * KH * KW] - - result = ggml_reshape_4d(ctx, result, im2col->ne[1], im2col->ne[2], im2col->ne[3], a->ne[3]); // [OC, N, OH, OW] - result = ggml_cont(ctx, ggml_permute(ctx, result, 0, 1, 3, 2)); // [N, OC, OH, OW] - - - return result; -} - -// a: [OC*IC, KD, KH, KW] -// b: [N*IC, ID, IH, IW] -// result: [N*OD, OH, OW, IC * KD * KH * KW] -struct ggml_tensor * ggml_im2col_3d( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int64_t IC, - int s0, // stride width - int s1, // stride height - int s2, // stride depth - int p0, // padding width - int p1, // padding height - int p2, // padding depth - int d0, // dilation width - int d1, // dilation height - int d2, // dilation depth - enum ggml_type dst_type) { - const int64_t N = b->ne[3] / IC; - const int64_t ID = b->ne[2]; - const int64_t IH = b->ne[1]; - const int64_t IW = b->ne[0]; - - const int64_t OC = a->ne[3] / IC; - UNUSED(OC); - const int64_t KD = a->ne[2]; - const int64_t KH = a->ne[1]; - const int64_t KW = a->ne[0]; - const int64_t OD = ggml_calc_conv_output_size(ID, KD, s2, p2, d2); - const int64_t OH = ggml_calc_conv_output_size(IH, KH, s1, p1, d1); - const int64_t OW = ggml_calc_conv_output_size(IW, KW, s0, p0, d0); - - GGML_ASSERT((OD > 0) && "b too small compared to a"); - GGML_ASSERT((OH > 0) && "b too small compared to a"); - GGML_ASSERT((OW > 0) && "b too small compared to a"); - - - const int64_t ne[4] = {KW*KH*KD*IC, OW, OH, OD*N}; - - struct ggml_tensor * result = ggml_new_tensor(ctx, dst_type, 4, ne); - int32_t params[] = { s0, s1, s2, p0, p1, p2, d0, d1, d2, (int32_t)IC}; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_IM2COL_3D; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -// a: [OC*IC, KD, KH, KW] -// b: [N*IC, ID, IH, IW] -// result: [N*OC, OD, OH, OW] -struct ggml_tensor * ggml_conv_3d( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int64_t IC, - int s0, // stride width - int s1, // stride height - int s2, // stride depth - int p0, // padding width - int p1, // padding height - int p2, // padding depth - int d0, // dilation width - int d1, // dilation height - int d2 // dilation depth - ) { - struct ggml_tensor * im2col = ggml_im2col_3d(ctx, a, b, IC, s0, s1, s2, p0, p1, p2, d0, d1, d2, a->type); // [N*OD, OH, OW, IC * KD * KH * KW] - - int64_t OC = a->ne[3] / IC; - int64_t N = b->ne[3] / IC; - struct ggml_tensor * result = - ggml_mul_mat(ctx, - ggml_reshape_2d(ctx, im2col, im2col->ne[0], im2col->ne[3] * im2col->ne[2] * im2col->ne[1]), // [N*OD, OH, OW, IC * KD * KH * KW] => [N*OD*OH*OW, IC * KD * KH * KW] - ggml_reshape_2d(ctx, a, (a->ne[0] * a->ne[1] * a->ne[2] * IC), OC)); // [OC*IC, KD, KH, KW] => [OC, IC * KD * KH * KW] - - int64_t OD = im2col->ne[3] / N; - result = ggml_reshape_4d(ctx, result, im2col->ne[1]*im2col->ne[2], OD, N, OC); // [OC, N*OD*OH*OW] => [OC, N, OD, OH*OW] - result = ggml_cont(ctx, ggml_permute(ctx, result, 0, 1, 3, 2)); // [N, OC, OD, OH*OW] - result = ggml_reshape_4d(ctx, result, im2col->ne[1], im2col->ne[2], OD, OC * N); // [N*OC, OD, OH, OW] - - return result; -} - -// ggml_conv_2d_sk_p0 - -struct ggml_tensor * ggml_conv_2d_sk_p0( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_conv_2d(ctx, a, b, a->ne[0], a->ne[1], 0, 0, 1, 1); -} - -// ggml_conv_2d_s1_ph - -struct ggml_tensor * ggml_conv_2d_s1_ph( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - return ggml_conv_2d(ctx, a, b, 1, 1, a->ne[0] / 2, a->ne[1] / 2, 1, 1); -} - -// ggml_conv_2d_dw - -struct ggml_tensor * ggml_conv_2d_dw( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int s0, - int s1, - int p0, - int p1, - int d0, - int d1) { - struct ggml_tensor * new_a = ggml_reshape_4d(ctx, a, a->ne[0], a->ne[1], 1, a->ne[2] * a->ne[3]); - struct ggml_tensor * im2col = ggml_im2col(ctx, new_a, - ggml_reshape_4d(ctx, b, b->ne[0], b->ne[1], 1, b->ne[2] * b->ne[3]), - s0, s1, p0, p1, d0, d1, true, GGML_TYPE_F16); // [N * IC, OH, OW, KH * KW] - struct ggml_tensor * new_b = ggml_reshape_4d(ctx, im2col, im2col->ne[0], im2col->ne[2] * im2col->ne[1], b->ne[2], b->ne[3]); // [N * IC, OH, OW, KH * KW] => [N, IC, OH * OW, KH * KW] - - new_a = ggml_reshape_4d(ctx, new_a, (new_a->ne[0] * new_a->ne[1]), new_a->ne[2], new_a->ne[3], 1); // [OC，1, KH, KW] => [1, OC, 1, KH * KW] - struct ggml_tensor * result = ggml_mul_mat(ctx, new_a, new_b); - result = ggml_reshape_4d(ctx, result, im2col->ne[1], im2col->ne[2], b->ne[2], b->ne[3]); // [N, OC, OH, OW] - - return result; -} - -// ggml_conv_2d_dw_direct - -struct ggml_tensor * ggml_conv_2d_dw_direct( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int stride0, - int stride1, - int pad0, - int pad1, - int dilation0, - int dilation1) { - GGML_ASSERT(a->ne[2] == 1); - GGML_ASSERT(a->ne[3] == b->ne[2]); - int64_t ne[4]; - ne[0] = ggml_calc_conv_output_size(b->ne[0], a->ne[0], stride0, pad0, dilation0); - ne[1] = ggml_calc_conv_output_size(b->ne[1], a->ne[1], stride1, pad1, dilation1); - ne[2] = b->ne[2]; - ne[3] = b->ne[3]; - - struct ggml_tensor * result = ggml_new_tensor(ctx, b->type, 4, ne); - - if (ggml_is_contiguous_channels(b)) { - // Result will be permuted the same way as input (CWHN order) - const int64_t type_size = ggml_type_size(result->type); - GGML_ASSERT(ggml_blck_size(result->type) == 1); - result->nb[0] = result->ne[2] * type_size; - result->nb[1] = result->ne[0] * result->nb[0]; - result->nb[2] = type_size; - } - - int32_t params[] = { stride0, stride1, pad0, pad1, dilation0, dilation1 }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_CONV_2D_DW; - result->src[0] = a; - result->src[1] = b; - return result; -} - -// ggml_conv_2d_direct - -struct ggml_tensor * ggml_conv_2d_direct( - struct ggml_context * ctx, - struct ggml_tensor * a, // convolution kernel [KW, KH, IC, OC] - struct ggml_tensor * b, // input data [W, H, C, N] - int s0, // stride dimension 0 - int s1, // stride dimension 1 - int p0, // padding dimension 0 - int p1, // padding dimension 1 - int d0, // dilation dimension 0 - int d1) {// dilation dimension 1 - - GGML_ASSERT(a->ne[2] == b->ne[2]); - //GGML_ASSERT(a->type == b->type); - - int64_t ne[4]; - ne[0] = ggml_calc_conv_output_size(b->ne[0], a->ne[0], s0, p0, d0); - ne[1] = ggml_calc_conv_output_size(b->ne[1], a->ne[1], s1, p1, d1); - ne[2] = a->ne[3]; - ne[3] = b->ne[3]; - - struct ggml_tensor * result = ggml_new_tensor(ctx, b->type, 4, ne); - - ggml_set_op_params_i32(result, 0, s0); - ggml_set_op_params_i32(result, 1, s1); - ggml_set_op_params_i32(result, 2, p0); - ggml_set_op_params_i32(result, 3, p1); - ggml_set_op_params_i32(result, 4, d0); - ggml_set_op_params_i32(result, 5, d1); - - result->op = GGML_OP_CONV_2D; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -// ggml_conv_3d_direct - -struct ggml_tensor * ggml_conv_3d_direct( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int s0, - int s1, - int s2, - int p0, - int p1, - int p2, - int d0, - int d1, - int d2, - int c, - int n, - int oc) { - - GGML_ASSERT(a->ne[3] == (int64_t) c * oc); - GGML_ASSERT(b->ne[3] == (int64_t) c * n); - - int64_t ne[4]; - ne[0] = ggml_calc_conv_output_size(b->ne[0], a->ne[0], s0, p0, d0); - ne[1] = ggml_calc_conv_output_size(b->ne[1], a->ne[1], s1, p1, d1); - ne[2] = ggml_calc_conv_output_size(b->ne[2], a->ne[2], s2, p2, d2); - ne[3] = (int64_t) oc * n; - - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - - ggml_set_op_params_i32(result, 0, s0); - ggml_set_op_params_i32(result, 1, s1); - ggml_set_op_params_i32(result, 2, s2); - ggml_set_op_params_i32(result, 3, p0); - ggml_set_op_params_i32(result, 4, p1); - ggml_set_op_params_i32(result, 5, p2); - ggml_set_op_params_i32(result, 6, d0); - ggml_set_op_params_i32(result, 7, d1); - ggml_set_op_params_i32(result, 8, d2); - ggml_set_op_params_i32(result, 9, c); - ggml_set_op_params_i32(result, 10, n); - ggml_set_op_params_i32(result, 11, oc); - - result->op = GGML_OP_CONV_3D; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -// ggml_conv_transpose_2d_p0 - -static int64_t ggml_calc_conv_transpose_output_size(int64_t ins, int64_t ks, int s, int p) { - return (ins - 1) * s - 2 * p + ks; -} - -struct ggml_tensor * ggml_conv_transpose_2d_p0( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - int stride) { - GGML_ASSERT(a->ne[3] == b->ne[2]); - - const int64_t ne[4] = { - ggml_calc_conv_transpose_output_size(b->ne[0], a->ne[0], stride, 0 /*p0*/), - ggml_calc_conv_transpose_output_size(b->ne[1], a->ne[1], stride, 0 /*p1*/), - a->ne[2], b->ne[3], - }; - - struct ggml_tensor* result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - - ggml_set_op_params_i32(result, 0, stride); - - result->op = GGML_OP_CONV_TRANSPOSE_2D; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -// ggml_pool_* - -static int64_t ggml_calc_pool_output_size(int64_t ins, int ks, int s, float p) { - return (ins + 2 * p - ks) / s + 1; -} - -// ggml_pool_1d - -struct ggml_tensor * ggml_pool_1d( - struct ggml_context * ctx, - struct ggml_tensor * a, - enum ggml_op_pool op, - int k0, - int s0, - int p0) { - const int64_t ne[4] = { - ggml_calc_pool_output_size(a->ne[0], k0, s0, p0), - a->ne[1], - a->ne[2], - a->ne[3], - }; - GGML_ASSERT(ne[0] > 0); - - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - - int32_t params[] = { op, k0, s0, p0 }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_POOL_1D; - result->src[0] = a; - - return result; -} - -// ggml_pool_2d - -struct ggml_tensor * ggml_pool_2d( - struct ggml_context * ctx, - struct ggml_tensor * a, - enum ggml_op_pool op, - int k0, - int k1, - int s0, - int s1, - float p0, - float p1) { - struct ggml_tensor * result; - const int64_t ne[4] = { - ggml_calc_pool_output_size(a->ne[0], k0, s0, p0), - ggml_calc_pool_output_size(a->ne[1], k1, s1, p1), - a->ne[2], - a->ne[3], - }; - GGML_ASSERT(ne[0] > 0); - GGML_ASSERT(ne[1] > 0); - - result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - - int32_t params[] = { op, k0, k1, s0, s1, p0, p1 }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_POOL_2D; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_pool_2d_back( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * af, - enum ggml_op_pool op, - int k0, - int k1, - int s0, - int s1, - float p0, - float p1) { - struct ggml_tensor * result; - result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, af->ne); - - int32_t params[] = { op, k0, k1, s0, s1, p0, p1 }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_POOL_2D_BACK; - result->src[0] = a; - result->src[1] = af; - - return result; -} - -// ggml_upscale / ggml_interpolate - -static struct ggml_tensor * ggml_interpolate_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - int64_t ne1, - int64_t ne2, - int64_t ne3, - uint32_t mode) { - GGML_ASSERT((mode & 0xFF) < GGML_SCALE_MODE_COUNT); - // TODO: implement antialias for modes other than bilinear - GGML_ASSERT(!(mode & GGML_SCALE_FLAG_ANTIALIAS) || (mode & 0xFF) == GGML_SCALE_MODE_BILINEAR); - GGML_ASSERT(a->type == GGML_TYPE_F32); - - struct ggml_tensor * result = ggml_new_tensor_4d(ctx, a->type, ne0, ne1, ne2, ne3); - - ggml_set_op_params_i32(result, 0, (int32_t)mode); - - result->op = GGML_OP_UPSCALE; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_upscale( - struct ggml_context * ctx, - struct ggml_tensor * a, - int scale_factor, - enum ggml_scale_mode mode) { - GGML_ASSERT(scale_factor > 1); - return ggml_interpolate_impl(ctx, a, a->ne[0] * scale_factor, a->ne[1] * scale_factor, a->ne[2], a->ne[3], mode); -} - -struct ggml_tensor * ggml_upscale_ext( - struct ggml_context * ctx, - struct ggml_tensor * a, - int ne0, - int ne1, - int ne2, - int ne3, - enum ggml_scale_mode mode) { - return ggml_interpolate_impl(ctx, a, ne0, ne1, ne2, ne3, mode); -} - -struct ggml_tensor * ggml_interpolate( - struct ggml_context * ctx, - struct ggml_tensor * a, - int64_t ne0, - int64_t ne1, - int64_t ne2, - int64_t ne3, - uint32_t mode) { - return ggml_interpolate_impl(ctx, a, ne0, ne1, ne2, ne3, mode); -} - -// ggml_pad - -struct ggml_tensor * ggml_pad( - struct ggml_context * ctx, - struct ggml_tensor * a, - int p0, - int p1, - int p2, - int p3) { - return ggml_pad_ext(ctx, a, 0, p0, 0, p1, 0, p2, 0, p3); -} - -// ggml_pad_circular - -struct ggml_tensor * ggml_pad_circular( - struct ggml_context * ctx, - struct ggml_tensor * a, - int p0, - int p1, - int p2, - int p3) { - return ggml_pad_ext_circular(ctx, a, 0, p0, 0, p1, 0, p2, 0, p3); -} - -struct ggml_tensor * ggml_pad_ext( - struct ggml_context * ctx, - struct ggml_tensor * a, - int lp0, - int rp0, - int lp1, - int rp1, - int lp2, - int rp2, - int lp3, - int rp3 - ) { - struct ggml_tensor * result = ggml_new_tensor_4d(ctx, a->type, - a->ne[0] + lp0 + rp0, - a->ne[1] + lp1 + rp1, - a->ne[2] + lp2 + rp2, - a->ne[3] + lp3 + rp3); - - ggml_set_op_params_i32(result, 0, lp0); - ggml_set_op_params_i32(result, 1, rp0); - ggml_set_op_params_i32(result, 2, lp1); - ggml_set_op_params_i32(result, 3, rp1); - ggml_set_op_params_i32(result, 4, lp2); - ggml_set_op_params_i32(result, 5, rp2); - ggml_set_op_params_i32(result, 6, lp3); - ggml_set_op_params_i32(result, 7, rp3); - ggml_set_op_params_i32(result, 8, 0); // not circular by default - - - result->op = GGML_OP_PAD; - result->src[0] = a; - - return result; -} - -// ggml_pad_ext_circular - -struct ggml_tensor * ggml_pad_ext_circular( - struct ggml_context * ctx, - struct ggml_tensor * a, - int lp0, - int rp0, - int lp1, - int rp1, - int lp2, - int rp2, - int lp3, - int rp3 - ) { - struct ggml_tensor * result = ggml_pad_ext(ctx, a, lp0, rp0, lp1, rp1, lp2, rp2, lp3, rp3); - ggml_set_op_params_i32(result, 8, 1); // circular - return result; -} - -// ggml_pad_reflect_1d - -struct ggml_tensor * ggml_pad_reflect_1d( - struct ggml_context * ctx, - struct ggml_tensor * a, - int p0, - int p1) { - GGML_ASSERT(p0 >= 0); - GGML_ASSERT(p1 >= 0); - - GGML_ASSERT(p0 < a->ne[0]); // padding length on each size must be less than the - GGML_ASSERT(p1 < a->ne[0]); // existing length of the dimension being padded - - GGML_ASSERT(ggml_is_contiguous(a)); - GGML_ASSERT(a->type == GGML_TYPE_F32); - - struct ggml_tensor * result = ggml_new_tensor_4d(ctx, a->type, - a->ne[0] + p0 + p1, - a->ne[1], - a->ne[2], - a->ne[3]); - - int32_t params[] = { p0, p1 }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_PAD_REFLECT_1D; - result->src[0] = a; - - return result; -} - -// ggml_roll - -struct ggml_tensor * ggml_roll( - struct ggml_context * ctx, - struct ggml_tensor * a, - int shift0, - int shift1, - int shift2, - int shift3) { - GGML_ASSERT(a->nb[0] == ggml_type_size(a->type)); - GGML_ASSERT(abs(shift0) < a->ne[0]); - GGML_ASSERT(abs(shift1) < a->ne[1]); - GGML_ASSERT(abs(shift2) < a->ne[2]); - GGML_ASSERT(abs(shift3) < a->ne[3]); - - struct ggml_tensor * result = ggml_dup_tensor(ctx, a); - - ggml_set_op_params_i32(result, 0, shift0); - ggml_set_op_params_i32(result, 1, shift1); - ggml_set_op_params_i32(result, 2, shift2); - ggml_set_op_params_i32(result, 3, shift3); - - result->op = GGML_OP_ROLL; - result->src[0] = a; - - return result; -} - -// ggml_timestep_embedding - -struct ggml_tensor * ggml_timestep_embedding( - struct ggml_context * ctx, - struct ggml_tensor * timesteps, - int dim, - int max_period) { - - struct ggml_tensor * result = ggml_new_tensor_2d(ctx, GGML_TYPE_F32, dim, timesteps->ne[0]); - - ggml_set_op_params_i32(result, 0, dim); - ggml_set_op_params_i32(result, 1, max_period); - - result->op = GGML_OP_TIMESTEP_EMBEDDING; - result->src[0] = timesteps; - - return result; -} - -// ggml_tri - -struct ggml_tensor * ggml_tri( - struct ggml_context * ctx, - struct ggml_tensor * a, - enum ggml_tri_type type) { - GGML_ASSERT(a->type == GGML_TYPE_F32); - - GGML_ASSERT(ggml_is_contiguous(a)); - GGML_ASSERT(a->ne[0] == a->ne[1]); - - struct ggml_tensor * result = ggml_dup_tensor(ctx, a); - - ggml_set_op_params_i32(result, 0, type); - - result->op = GGML_OP_TRI; - result->src[0] = a; - - return result; -} - -// ggml_fill - -static struct ggml_tensor * ggml_fill_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - float c, - bool inplace) { - GGML_ASSERT(a->type == GGML_TYPE_F32); - GGML_ASSERT(ggml_is_contiguous(a)); - - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - ggml_set_op_params_f32(result, 0, c); - - result->op = GGML_OP_FILL; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_fill( - struct ggml_context * ctx, - struct ggml_tensor * a, - float c) { - return ggml_fill_impl(ctx, a, c, false); -} - -struct ggml_tensor * ggml_fill_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - float c) { - return ggml_fill_impl(ctx, a, c, true); -} - -// ggml_argsort - -struct ggml_tensor * ggml_argsort( - struct ggml_context * ctx, - struct ggml_tensor * a, - enum ggml_sort_order order) { - GGML_ASSERT(a->ne[0] <= INT32_MAX); - - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_I32, GGML_MAX_DIMS, a->ne); - - ggml_set_op_params_i32(result, 0, (int32_t) order); - - result->op = GGML_OP_ARGSORT; - result->src[0] = a; - - return result; -} - -// ggml_argsort_top_k - -struct ggml_tensor * ggml_argsort_top_k( - struct ggml_context * ctx, - struct ggml_tensor * a, - int k) { - GGML_ASSERT(a->ne[0] >= k); - - struct ggml_tensor * result = ggml_argsort(ctx, a, GGML_SORT_ORDER_DESC); - - result = ggml_view_4d(ctx, result, - k, result->ne[1], result->ne[2], result->ne[3], - result->nb[1], result->nb[2], result->nb[3], - 0); - - return result; -} - -// ggml_top_k - -struct ggml_tensor * ggml_top_k( - struct ggml_context * ctx, - struct ggml_tensor * a, - int k) { - GGML_ASSERT(a->ne[0] >= k); - - struct ggml_tensor * result = ggml_new_tensor_4d(ctx, GGML_TYPE_I32, k, a->ne[1], a->ne[2], a->ne[3]); - - result->op = GGML_OP_TOP_K; - result->src[0] = a; - - return result; -} - -// ggml_arange - -struct ggml_tensor * ggml_arange( - struct ggml_context * ctx, - float start, - float stop, - float step) { - GGML_ASSERT(stop > start); - - const int64_t steps = (int64_t) ceilf((stop - start) / step); - - struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, steps); - - ggml_set_op_params_f32(result, 0, start); - ggml_set_op_params_f32(result, 1, stop); - ggml_set_op_params_f32(result, 2, step); - - result->op = GGML_OP_ARANGE; - - return result; -} - -// ggml_flash_attn_ext - -struct ggml_tensor * ggml_flash_attn_ext( - struct ggml_context * ctx, - struct ggml_tensor * q, - struct ggml_tensor * k, - struct ggml_tensor * v, - struct ggml_tensor * mask, - float scale, - float max_bias, - float logit_softcap) { - GGML_ASSERT(ggml_can_mul_mat(k, q)); - // TODO: check if vT can be multiplied by (k*qT) - - GGML_ASSERT(q->ne[3] == k->ne[3]); - GGML_ASSERT(q->ne[3] == v->ne[3]); - - if (mask) { - GGML_ASSERT(mask->type == GGML_TYPE_F16); - GGML_ASSERT(ggml_is_contiguous(mask)); - //GGML_ASSERT(ggml_can_repeat_rows(mask, qk)); - - GGML_ASSERT(q->ne[2] % mask->ne[2] == 0); - GGML_ASSERT(q->ne[3] % mask->ne[3] == 0); - } - - if (max_bias > 0.0f) { - GGML_ASSERT(mask); - } - - // permute(0, 2, 1, 3) - int64_t ne[4] = { v->ne[0], q->ne[2], q->ne[1], q->ne[3] }; - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - - float params[] = { scale, max_bias, logit_softcap }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_FLASH_ATTN_EXT; - result->src[0] = q; - result->src[1] = k; - result->src[2] = v; - result->src[3] = mask; - - return result; -} - -void ggml_flash_attn_ext_set_prec( - struct ggml_tensor * a, - enum ggml_prec prec) { - GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); - - const int32_t prec_i32 = (int32_t) prec; - - ggml_set_op_params_i32(a, 3, prec_i32); // scale is on first pos, max_bias on second -} - -enum ggml_prec ggml_flash_attn_ext_get_prec( - const struct ggml_tensor * a) { - GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); - - const int32_t prec_i32 = ggml_get_op_params_i32(a, 3); - - return (enum ggml_prec) prec_i32; -} - -void ggml_flash_attn_ext_add_sinks( - struct ggml_tensor * a, - struct ggml_tensor * sinks) { - if (!sinks) { - a->src[4] = NULL; - return; - } - - GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); - GGML_ASSERT(a->src[4] == NULL); - GGML_ASSERT(a->src[0]->ne[2] == sinks->ne[0]); - GGML_ASSERT(sinks->type == GGML_TYPE_F32); - - a->src[4] = sinks; -} - -// ggml_flash_attn_sparse - -struct ggml_tensor * ggml_flash_attn_sparse( - struct ggml_context * ctx, - struct ggml_tensor * q, - struct ggml_tensor * k, - struct ggml_tensor * v, - float scale, - float alpha) { - - GGML_ASSERT(ggml_can_mul_mat(k, q)); - - // Output shape matches ggml_flash_attn_ext: [v->ne[0], q->ne[2], q->ne[1], q->ne[3]] - int64_t ne[4] = { v->ne[0], q->ne[2], q->ne[1], q->ne[3] }; - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - - float params[] = { scale, alpha }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_FLASH_ATTN_SPARSE; - result->src[0] = q; - result->src[1] = k; - result->src[2] = v; - - return result; -} - -// ggml_flash_attn_back - -struct ggml_tensor * ggml_flash_attn_back( - struct ggml_context * ctx, - struct ggml_tensor * q, - struct ggml_tensor * k, - struct ggml_tensor * v, - struct ggml_tensor * d, - bool masked) { - GGML_ABORT("TODO: adapt to ggml_flash_attn_ext() changes"); - - GGML_ASSERT(ggml_can_mul_mat(k, q)); - // TODO: check if vT can be multiplied by (k*qT) - - // d shape [D,N,ne2,ne3] - // q shape [D,N,ne2,ne3] - // k shape [D,M,kvne2,ne3] - // v shape [M,D,kvne2,ne3] - - const int64_t D = q->ne[0]; - const int64_t N = q->ne[1]; - const int64_t M = k->ne[1]; - const int64_t ne2 = q->ne[2]; - const int64_t ne3 = q->ne[3]; - const int64_t kvne2 = k->ne[2]; - - GGML_ASSERT(k->ne[0] == D); - GGML_ASSERT(v->ne[0] == M); - GGML_ASSERT(v->ne[1] == D); - GGML_ASSERT(d->ne[0] == D); - GGML_ASSERT(d->ne[1] == N); - GGML_ASSERT(k->ne[2] == kvne2); - GGML_ASSERT(k->ne[3] == ne3); - GGML_ASSERT(v->ne[2] == kvne2); - GGML_ASSERT(v->ne[3] == ne3); - GGML_ASSERT(d->ne[2] == ne2); - GGML_ASSERT(d->ne[3] == ne3); - - GGML_ASSERT(ne2 % kvne2 == 0); - - // store gradients of q, k and v as continuous tensors concatenated in result. - // note: v and gradv are actually transposed, i.e. v->ne[0] != D. - const int64_t elem_q = ggml_nelements(q); - const int64_t elem_k = ggml_nelements(k); - const int64_t elem_v = ggml_nelements(v); - - enum ggml_type result_type = GGML_TYPE_F32; - GGML_ASSERT(ggml_blck_size(result_type) == 1); - const size_t tsize = ggml_type_size(result_type); - - const size_t offs_q = 0; - const size_t offs_k = offs_q + GGML_PAD(elem_q * tsize, GGML_MEM_ALIGN); - const size_t offs_v = offs_k + GGML_PAD(elem_k * tsize, GGML_MEM_ALIGN); - const size_t end = offs_v + GGML_PAD(elem_v * tsize, GGML_MEM_ALIGN); - - const size_t nelements = (end + tsize - 1)/tsize; - - struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, nelements); - - int32_t masked_i = masked ? 1 : 0; - ggml_set_op_params(result, &masked_i, sizeof(masked_i)); - - result->op = GGML_OP_FLASH_ATTN_BACK; - result->src[0] = q; - result->src[1] = k; - result->src[2] = v; - result->src[3] = d; - - return result; -} - -// ggml_ssm_conv - -struct ggml_tensor * ggml_ssm_conv( - struct ggml_context * ctx, - struct ggml_tensor * sx, - struct ggml_tensor * c) { - GGML_ASSERT(ggml_is_3d(sx)); - GGML_ASSERT(ggml_is_matrix(c)); - - const int64_t d_conv = c->ne[0]; - const int64_t d_inner = c->ne[1]; - const int64_t n_t = sx->ne[0] - d_conv + 1; // tokens per sequence - const int64_t n_s = sx->ne[2]; - - // TODO: maybe support other strides than 1? - GGML_ASSERT(sx->ne[0] == d_conv - 1 + n_t); - GGML_ASSERT(sx->ne[1] == d_inner); - GGML_ASSERT(n_t >= 0); - - struct ggml_tensor * result = ggml_new_tensor_3d(ctx, GGML_TYPE_F32, d_inner, n_t, n_s); - - result->op = GGML_OP_SSM_CONV; - result->src[0] = sx; - result->src[1] = c; - - return result; -} - -// dflash: tree-mode variant. Same op, with parent_ids plumbed into -// src[2] so the CUDA kernel gathers each token's window along its tree -// parent chain instead of the DFS-neighbour window. -struct ggml_tensor * ggml_ssm_conv_tree( - struct ggml_context * ctx, - struct ggml_tensor * sx, - struct ggml_tensor * c, - struct ggml_tensor * parent_ids) { - struct ggml_tensor * result = ggml_ssm_conv(ctx, sx, c); - - GGML_ASSERT(parent_ids != NULL); - GGML_ASSERT(parent_ids->type == GGML_TYPE_I32); - GGML_ASSERT(ggml_is_contiguous(parent_ids)); - - const int64_t n_t = sx->ne[0] - c->ne[0] + 1; - const int64_t n_s = sx->ne[2]; - GGML_ASSERT(ggml_nelements(parent_ids) == n_t * n_s); - - result->src[2] = parent_ids; - - return result; -} - -// ggml_ssm_scan - -struct ggml_tensor * ggml_ssm_scan( - struct ggml_context * ctx, - struct ggml_tensor * s, - struct ggml_tensor * x, - struct ggml_tensor * dt, - struct ggml_tensor * A, - struct ggml_tensor * B, - struct ggml_tensor * C, - struct ggml_tensor * ids) { - GGML_ASSERT(ggml_is_contiguous(s)); - GGML_ASSERT(ggml_is_contiguous(dt)); - GGML_ASSERT(ggml_is_contiguous(A)); - GGML_ASSERT(x->nb[0] == ggml_type_size(x->type)); - GGML_ASSERT(B->nb[0] == ggml_type_size(B->type)); - GGML_ASSERT(C->nb[0] == ggml_type_size(C->type)); - GGML_ASSERT(x->nb[1] == x->ne[0]*x->nb[0]); - GGML_ASSERT(B->nb[1] == B->ne[0]*B->nb[0]); - GGML_ASSERT(C->nb[1] == C->ne[0]*C->nb[0]); - GGML_ASSERT(ggml_are_same_shape(B, C)); - GGML_ASSERT(ids->type == GGML_TYPE_I32); - - { - const int64_t d_state = s->ne[0]; - const int64_t head_dim = x->ne[0]; - const int64_t n_head = x->ne[1]; - const int64_t n_seq_tokens = x->ne[2]; - const int64_t n_seqs = x->ne[3]; - - GGML_ASSERT(dt->ne[0] == n_head); - GGML_ASSERT(dt->ne[1] == n_seq_tokens); - GGML_ASSERT(dt->ne[2] == n_seqs); - GGML_ASSERT(ggml_is_3d(dt)); - GGML_ASSERT(s->ne[1] == head_dim); - GGML_ASSERT(s->ne[2] == n_head); - GGML_ASSERT(B->ne[0] == d_state); - GGML_ASSERT(B->ne[2] == n_seq_tokens); - GGML_ASSERT(B->ne[3] == n_seqs); - GGML_ASSERT(ids->ne[0] == n_seqs); - GGML_ASSERT(ggml_is_vector(ids)); - GGML_ASSERT(A->ne[1] == n_head); - GGML_ASSERT(ggml_is_matrix(A)); - - if (A->ne[0] != 1) { - // Mamba-1 has more granular decay factors - GGML_ASSERT(A->ne[0] == d_state); - } - } - - // concatenated y + ssm_states - struct ggml_tensor * result = ggml_new_tensor_1d(ctx, GGML_TYPE_F32, ggml_nelements(x) + s->ne[0]*s->ne[1]*s->ne[2]*ids->ne[0]); - - result->op = GGML_OP_SSM_SCAN; - result->src[0] = s; - result->src[1] = x; - result->src[2] = dt; - result->src[3] = A; - result->src[4] = B; - result->src[5] = C; - result->src[6] = ids; - - return result; -} - -// ggml_win_part - -struct ggml_tensor * ggml_win_part( - struct ggml_context * ctx, - struct ggml_tensor * a, - int w) { - GGML_ASSERT(a->ne[3] == 1); - GGML_ASSERT(a->type == GGML_TYPE_F32); - - // padding - const int px = (w - a->ne[1]%w)%w; - const int py = (w - a->ne[2]%w)%w; - - const int npx = (px + a->ne[1])/w; - const int npy = (py + a->ne[2])/w; - const int np = npx*npy; - - const int64_t ne[4] = { a->ne[0], w, w, np, }; - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - - int32_t params[] = { npx, npy, w }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_WIN_PART; - result->src[0] = a; - - return result; -} - -// ggml_win_unpart - -struct ggml_tensor * ggml_win_unpart( - struct ggml_context * ctx, - struct ggml_tensor * a, - int w0, - int h0, - int w) { - GGML_ASSERT(a->type == GGML_TYPE_F32); - - const int64_t ne[4] = { a->ne[0], w0, h0, 1, }; - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 3, ne); - - int32_t params[] = { w }; - ggml_set_op_params(result, params, sizeof(params)); - - result->op = GGML_OP_WIN_UNPART; - result->src[0] = a; - - return result; -} - -// ggml_get_rel_pos - -struct ggml_tensor * ggml_get_rel_pos( - struct ggml_context * ctx, - struct ggml_tensor * a, - int qh, - int kh) { - GGML_ASSERT(qh == kh); - GGML_ASSERT(2*MAX(qh, kh) - 1 == a->ne[1]); - - const int64_t ne[4] = { a->ne[0], kh, qh, 1, }; - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F16, 3, ne); - - result->op = GGML_OP_GET_REL_POS; - result->src[0] = a; - - return result; -} - -// ggml_add_rel_pos - -static struct ggml_tensor * ggml_add_rel_pos_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * pw, - struct ggml_tensor * ph, - bool inplace) { - GGML_ASSERT(ggml_are_same_shape(pw, ph)); - GGML_ASSERT(ggml_is_contiguous(a)); - GGML_ASSERT(ggml_is_contiguous(pw)); - GGML_ASSERT(ggml_is_contiguous(ph)); - GGML_ASSERT(ph->type == GGML_TYPE_F32); - GGML_ASSERT(pw->type == GGML_TYPE_F32); - GGML_ASSERT(pw->ne[3] == a->ne[2]); - GGML_ASSERT(pw->ne[0]*pw->ne[0] == a->ne[0]); - GGML_ASSERT(pw->ne[1]*pw->ne[2] == a->ne[1]); - - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - ggml_set_op_params_i32(result, 0, inplace ? 1 : 0); - - result->op = GGML_OP_ADD_REL_POS; - result->src[0] = a; - result->src[1] = pw; - result->src[2] = ph; - - return result; -} - -struct ggml_tensor * ggml_add_rel_pos( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * pw, - struct ggml_tensor * ph) { - return ggml_add_rel_pos_impl(ctx, a, pw, ph, false); -} - -struct ggml_tensor * ggml_add_rel_pos_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * pw, - struct ggml_tensor * ph) { - return ggml_add_rel_pos_impl(ctx, a, pw, ph, true); -} - -// ggml_rwkv_wkv6 - -struct ggml_tensor * ggml_rwkv_wkv6( - struct ggml_context * ctx, - struct ggml_tensor * k, - struct ggml_tensor * v, - struct ggml_tensor * r, - struct ggml_tensor * tf, - struct ggml_tensor * td, - struct ggml_tensor * state) { - GGML_ASSERT(ggml_is_contiguous(k)); - GGML_ASSERT(ggml_is_contiguous(v)); - GGML_ASSERT(ggml_is_contiguous(r)); - GGML_ASSERT(ggml_is_contiguous(tf)); - GGML_ASSERT(ggml_is_contiguous(td)); - GGML_ASSERT(ggml_is_contiguous(state)); - - const int64_t S = k->ne[0]; - const int64_t H = k->ne[1]; - const int64_t n_tokens = k->ne[2]; - const int64_t n_seqs = state->ne[1]; - { - GGML_ASSERT(v->ne[0] == S && v->ne[1] == H && v->ne[2] == n_tokens); - GGML_ASSERT(r->ne[0] == S && r->ne[1] == H && r->ne[2] == n_tokens); - GGML_ASSERT(td->ne[0] == S && td->ne[1] == H && td->ne[2] == n_tokens); - GGML_ASSERT(ggml_nelements(state) == S * S * H * n_seqs); - } - - // concat output and new_state - const int64_t ne[4] = { S * H, n_tokens + S * n_seqs, 1, 1 }; - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - - result->op = GGML_OP_RWKV_WKV6; - result->src[0] = k; - result->src[1] = v; - result->src[2] = r; - result->src[3] = tf; - result->src[4] = td; - result->src[5] = state; - - return result; -} - -// ggml_gated_linear_attn - -struct ggml_tensor * ggml_gated_linear_attn( - struct ggml_context * ctx, - struct ggml_tensor * k, - struct ggml_tensor * v, - struct ggml_tensor * q, - struct ggml_tensor * g, - struct ggml_tensor * state, - float scale) { - GGML_ASSERT(ggml_is_contiguous(k)); - GGML_ASSERT(ggml_is_contiguous(v)); - GGML_ASSERT(ggml_is_contiguous(q)); - GGML_ASSERT(ggml_is_contiguous(g)); - GGML_ASSERT(ggml_is_contiguous(state)); - - const int64_t S = k->ne[0]; - const int64_t H = k->ne[1]; - const int64_t n_tokens = k->ne[2]; - const int64_t n_seqs = state->ne[1]; - { - GGML_ASSERT(v->ne[0] == S && v->ne[1] == H && v->ne[2] == n_tokens); - GGML_ASSERT(q->ne[0] == S && q->ne[1] == H && q->ne[2] == n_tokens); - GGML_ASSERT(g->ne[0] == S && g->ne[1] == H && g->ne[2] == n_tokens); - GGML_ASSERT(ggml_nelements(state) == S * S * H * n_seqs); - } - - // concat output and new_state - const int64_t ne[4] = { S * H, n_tokens + S * n_seqs, 1, 1 }; - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - - ggml_set_op_params_f32(result, 0, scale); - - result->op = GGML_OP_GATED_LINEAR_ATTN; - result->src[0] = k; - result->src[1] = v; - result->src[2] = q; - result->src[3] = g; - result->src[4] = state; - - return result; -} - -// ggml_rwkv_wkv7 - -struct ggml_tensor * ggml_rwkv_wkv7( - struct ggml_context * ctx, - struct ggml_tensor * r, - struct ggml_tensor * w, - struct ggml_tensor * k, - struct ggml_tensor * v, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * state) { - GGML_ASSERT(ggml_is_contiguous(r)); - GGML_ASSERT(ggml_is_contiguous(w)); - GGML_ASSERT(ggml_is_contiguous(k)); - GGML_ASSERT(ggml_is_contiguous(v)); - GGML_ASSERT(ggml_is_contiguous(a)); - GGML_ASSERT(ggml_is_contiguous(b)); - GGML_ASSERT(ggml_is_contiguous(state)); - - const int64_t S = k->ne[0]; - const int64_t H = k->ne[1]; - const int64_t n_tokens = k->ne[2]; - const int64_t n_seqs = state->ne[1]; - { - GGML_ASSERT(w->ne[0] == S && w->ne[1] == H && w->ne[2] == n_tokens); - GGML_ASSERT(k->ne[0] == S && k->ne[1] == H && k->ne[2] == n_tokens); - GGML_ASSERT(v->ne[0] == S && v->ne[1] == H && v->ne[2] == n_tokens); - GGML_ASSERT(a->ne[0] == S && a->ne[1] == H && a->ne[2] == n_tokens); - GGML_ASSERT(b->ne[0] == S && b->ne[1] == H && b->ne[2] == n_tokens); - GGML_ASSERT(ggml_nelements(state) == S * S * H * n_seqs); - } - - // concat output and new_state - const int64_t ne[4] = { S * H, n_tokens + S * n_seqs, 1, 1 }; - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - - result->op = GGML_OP_RWKV_WKV7; - result->src[0] = r; - result->src[1] = w; - result->src[2] = k; - result->src[3] = v; - result->src[4] = a; - result->src[5] = b; - result->src[6] = state; - - return result; -} - -// ggml_unary - -static struct ggml_tensor * ggml_unary_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - enum ggml_unary_op op, - bool inplace) { - GGML_ASSERT(ggml_is_contiguous_rows(a)); - - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - ggml_set_op_params_i32(result, 0, (int32_t) op); - - result->op = GGML_OP_UNARY; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_unary( - struct ggml_context * ctx, - struct ggml_tensor * a, - enum ggml_unary_op op) { - return ggml_unary_impl(ctx, a, op, false); -} - -struct ggml_tensor * ggml_unary_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - enum ggml_unary_op op) { - return ggml_unary_impl(ctx, a, op, true); -} - -// ggml_map_custom1 - -static struct ggml_tensor * ggml_map_custom1_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - const ggml_custom1_op_t fun, - int n_tasks, - void * userdata, - bool inplace) { - GGML_ASSERT(n_tasks == GGML_N_TASKS_MAX || n_tasks > 0); - - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - struct ggml_map_custom1_op_params params = { - /*.fun =*/ fun, - /*.n_tasks =*/ n_tasks, - /*.userdata =*/ userdata - }; - ggml_set_op_params(result, ¶ms, sizeof(params)); - - result->op = GGML_OP_MAP_CUSTOM1; - result->src[0] = a; - - return result; -} - -struct ggml_tensor * ggml_map_custom1( - struct ggml_context * ctx, - struct ggml_tensor * a, - const ggml_custom1_op_t fun, - int n_tasks, - void * userdata) { - return ggml_map_custom1_impl(ctx, a, fun, n_tasks, userdata, false); -} - -struct ggml_tensor * ggml_map_custom1_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - const ggml_custom1_op_t fun, - int n_tasks, - void * userdata) { - return ggml_map_custom1_impl(ctx, a, fun, n_tasks, userdata, true); -} - -// ggml_map_custom2 - -static struct ggml_tensor * ggml_map_custom2_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - const ggml_custom2_op_t fun, - int n_tasks, - void * userdata, - bool inplace) { - GGML_ASSERT(n_tasks == GGML_N_TASKS_MAX || n_tasks > 0); - - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - struct ggml_map_custom2_op_params params = { - /*.fun =*/ fun, - /*.n_tasks =*/ n_tasks, - /*.userdata =*/ userdata - }; - ggml_set_op_params(result, ¶ms, sizeof(params)); - - result->op = GGML_OP_MAP_CUSTOM2; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -struct ggml_tensor * ggml_map_custom2( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - const ggml_custom2_op_t fun, - int n_tasks, - void * userdata) { - return ggml_map_custom2_impl(ctx, a, b, fun, n_tasks, userdata, false); -} - -struct ggml_tensor * ggml_map_custom2_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - const ggml_custom2_op_t fun, - int n_tasks, - void * userdata) { - return ggml_map_custom2_impl(ctx, a, b, fun, n_tasks, userdata, true); -} - -// ggml_map_custom3 - -static struct ggml_tensor * ggml_map_custom3_impl( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * c, - const ggml_custom3_op_t fun, - int n_tasks, - void * userdata, - bool inplace) { - GGML_ASSERT(n_tasks == GGML_N_TASKS_MAX || n_tasks > 0); - - struct ggml_tensor * result = inplace ? ggml_view_tensor(ctx, a) : ggml_dup_tensor(ctx, a); - - struct ggml_map_custom3_op_params params = { - /*.fun =*/ fun, - /*.n_tasks =*/ n_tasks, - /*.userdata =*/ userdata - }; - ggml_set_op_params(result, ¶ms, sizeof(params)); - - result->op = GGML_OP_MAP_CUSTOM3; - result->src[0] = a; - result->src[1] = b; - result->src[2] = c; - - return result; -} - -struct ggml_tensor * ggml_map_custom3( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * c, - const ggml_custom3_op_t fun, - int n_tasks, - void * userdata) { - return ggml_map_custom3_impl(ctx, a, b, c, fun, n_tasks, userdata, false); -} - -struct ggml_tensor * ggml_map_custom3_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * c, - const ggml_custom3_op_t fun, - int n_tasks, - void * userdata) { - return ggml_map_custom3_impl(ctx, a, b, c, fun, n_tasks, userdata, true); -} - -struct ggml_tensor * ggml_custom_4d( - struct ggml_context * ctx, - enum ggml_type type, - int64_t ne0, - int64_t ne1, - int64_t ne2, - int64_t ne3, - struct ggml_tensor ** args, - int n_args, - ggml_custom_op_t fun, - int n_tasks, - void * userdata) { - - GGML_ASSERT(n_args < GGML_MAX_SRC); - - struct ggml_tensor * result = ggml_new_tensor_4d(ctx, type, ne0, ne1, ne2, ne3); - - struct ggml_custom_op_params params = { - /*.fun =*/ fun, - /*.n_tasks =*/ n_tasks, - /*.userdata =*/ userdata - }; - ggml_set_op_params(result, ¶ms, sizeof(params)); - - result->op = GGML_OP_CUSTOM; - for (int i = 0; i < n_args; i++) { - result->src[i] = args[i]; - } - - return result; -} - -struct ggml_tensor * ggml_custom_inplace( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor ** args, - int n_args, - ggml_custom_op_t fun, - int n_tasks, - void * userdata) { - - GGML_ASSERT(n_args < GGML_MAX_SRC - 1); - - struct ggml_tensor * result = ggml_view_tensor(ctx, a); - - struct ggml_custom_op_params params = { - /*.fun =*/ fun, - /*.n_tasks =*/ n_tasks, - /*.userdata =*/ userdata - }; - ggml_set_op_params(result, ¶ms, sizeof(params)); - - result->op = GGML_OP_CUSTOM; - result->src[0] = a; - for (int i = 0; i < n_args; i++) { - result->src[i + 1] = args[i]; - } - - return result; -} -// ggml_cross_entropy_loss - -struct ggml_tensor * ggml_cross_entropy_loss( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b) { - GGML_ASSERT(ggml_are_same_shape(a, b)); - - struct ggml_tensor * result = ggml_new_tensor_1d(ctx, a->type, 1); - - result->op = GGML_OP_CROSS_ENTROPY_LOSS; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -// ggml_cross_entropy_loss_back - -struct ggml_tensor * ggml_cross_entropy_loss_back( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - struct ggml_tensor * c) { - GGML_ASSERT(ggml_is_scalar(a)); - GGML_ASSERT(ggml_are_same_shape(b, c)); - - struct ggml_tensor * result = ggml_dup_tensor(ctx, b); - - result->op = GGML_OP_CROSS_ENTROPY_LOSS_BACK; - result->src[0] = a; - result->src[1] = b; - result->src[2] = c; - - return result; -} - -// opt_step_adamw - -struct ggml_tensor * ggml_opt_step_adamw( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * grad, - struct ggml_tensor * m, - struct ggml_tensor * v, - struct ggml_tensor * adamw_params) { - GGML_ASSERT(a->flags & GGML_TENSOR_FLAG_PARAM); - GGML_ASSERT(ggml_are_same_shape(a, grad)); - GGML_ASSERT(ggml_are_same_shape(a, m)); - GGML_ASSERT(ggml_are_same_shape(a, v)); - GGML_ASSERT(adamw_params->type == GGML_TYPE_F32); - GGML_ASSERT(ggml_nelements(adamw_params) == 7); - - struct ggml_tensor * result = ggml_view_tensor(ctx, a); - - result->op = GGML_OP_OPT_STEP_ADAMW; - result->src[0] = a; - result->src[1] = grad; - result->src[2] = m; - result->src[3] = v; - result->src[4] = adamw_params; - - return result; -} - -// opt_step_sgd - -struct ggml_tensor * ggml_opt_step_sgd( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * grad, - struct ggml_tensor * params) { - GGML_ASSERT(a->flags & GGML_TENSOR_FLAG_PARAM); - GGML_ASSERT(ggml_are_same_shape(a, grad)); - GGML_ASSERT(params->type == GGML_TYPE_F32); - GGML_ASSERT(ggml_nelements(params) == 2); - - struct ggml_tensor * result = ggml_view_tensor(ctx, a); - - result->op = GGML_OP_OPT_STEP_SGD; - result->src[0] = a; - result->src[1] = grad; - result->src[2] = params; - - return result; -} - -// solve_tri - -struct ggml_tensor * ggml_solve_tri( - struct ggml_context * ctx, - struct ggml_tensor * a, - struct ggml_tensor * b, - bool left, - bool lower, - bool uni) { - GGML_ASSERT(a->type == GGML_TYPE_F32); - GGML_ASSERT(b->type == GGML_TYPE_F32); - - // A must be square and lower diagonal - GGML_ASSERT(a->ne[0] == a->ne[1]); - // B must have same outer dimension as A - GGML_ASSERT(a->ne[1] == b->ne[1]); - - // batch dimensions must be equal - GGML_ASSERT(a->ne[2] == b->ne[2]); - GGML_ASSERT(a->ne[3] == b->ne[3]); - - GGML_ASSERT(ggml_is_contiguous(a)); - GGML_ASSERT(ggml_is_contiguous(b)); - - GGML_ASSERT(lower && left && !uni); // TODO: support other variants - - struct ggml_tensor * result = ggml_new_tensor_4d(ctx, GGML_TYPE_F32, b->ne[0], b->ne[1], b->ne[2], b->ne[3]); - - result->op = GGML_OP_SOLVE_TRI; - result->src[0] = a; - result->src[1] = b; - - return result; -} - -// ggml_gated_delta_net - -struct ggml_tensor * ggml_gated_delta_net( - struct ggml_context * ctx, - struct ggml_tensor * q, - struct ggml_tensor * k, - struct ggml_tensor * v, - struct ggml_tensor * g, - struct ggml_tensor * beta, - struct ggml_tensor * state) { - GGML_ASSERT(ggml_is_contiguous_rows(q)); - GGML_ASSERT(ggml_is_contiguous_rows(k)); - GGML_ASSERT(ggml_is_contiguous_rows(v)); - GGML_ASSERT(ggml_is_contiguous(g)); - GGML_ASSERT(ggml_is_contiguous(beta)); - GGML_ASSERT(ggml_is_contiguous(state)); - - GGML_ASSERT(q->type == GGML_TYPE_F32); - GGML_ASSERT(k->type == GGML_TYPE_F32); - GGML_ASSERT(v->type == GGML_TYPE_F32); - GGML_ASSERT(g->type == GGML_TYPE_F32); - GGML_ASSERT(beta->type == GGML_TYPE_F32); - GGML_ASSERT(state->type == GGML_TYPE_F32); - - const int64_t S_v = v->ne[0]; - const int64_t H = v->ne[1]; - const int64_t n_tokens = v->ne[2]; - const int64_t n_seqs = v->ne[3]; - - // gate: scalar [1, H, T, B] or vector [S_v, H, T, B] (KDA) - GGML_ASSERT(g->ne[0] == 1 || g->ne[0] == S_v); - GGML_ASSERT(beta->ne[0] == 1); - - GGML_ASSERT(ggml_nelements(state) == S_v * S_v * H * n_seqs); - - // Pack output, final new_state, and per-step intermediate states into one tensor. - // Layout (in units of `S_v * H`-wide rows): - // [ attn_output: n_tokens*n_seqs | final_state: S_v*n_seqs | intermediate_states: S_v*n_tokens*n_seqs ] - // - // The final_state slot is kept for backward compatibility with stock llama.cpp - // callers that read state at offset S_v*H*n_tokens*n_seqs. The intermediate_states - // region is a dflash extension: for each token t in [0, n_tokens), it holds the - // recurrent state after processing token t. Used by the spec decoding loop to - // roll back SSM state to the accepted prefix without a full replay forward pass. - const int64_t ne[4] = { S_v * H, n_tokens * n_seqs + S_v * n_seqs + S_v * n_tokens * n_seqs, 1, 1 }; - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - - result->op = GGML_OP_GATED_DELTA_NET; - result->src[0] = q; - result->src[1] = k; - result->src[2] = v; - result->src[3] = g; - result->src[4] = beta; - result->src[5] = state; - - return result; -} - -void ggml_gated_delta_net_set_skip_intermediate( - struct ggml_tensor * tensor, - bool skip_intermediate) { - GGML_ASSERT(tensor != NULL); - GGML_ASSERT(tensor->op == GGML_OP_GATED_DELTA_NET); - ggml_set_op_params_i32(tensor, 0, skip_intermediate ? 1 : 0); - - const struct ggml_tensor * v = tensor->src[2]; - GGML_ASSERT(v != NULL); - - const int64_t S_v = v->ne[0]; - const int64_t n_tokens = v->ne[2]; - const int64_t n_seqs = v->ne[3]; - - // Compact only the plain chain path. Tree/persistent variants need - // intermediate states for branch reloads or explicit capture storage. - const bool can_compact = tensor->src[6] == NULL && tensor->src[7] == NULL; - tensor->ne[1] = n_tokens*n_seqs + S_v*n_seqs; - if (!skip_intermediate || !can_compact) { - tensor->ne[1] += S_v*n_tokens*n_seqs; - } - tensor->nb[2] = tensor->nb[1]*tensor->ne[1]; - tensor->nb[3] = tensor->nb[2]*tensor->ne[2]; -} - -// dflash: tree-mode variant. Same op, with parent_ids plumbed into -// src[6] so the CUDA kernel can branch-reload state at DFS transitions. -struct ggml_tensor * ggml_gated_delta_net_tree( - struct ggml_context * ctx, - struct ggml_tensor * q, - struct ggml_tensor * k, - struct ggml_tensor * v, - struct ggml_tensor * g, - struct ggml_tensor * beta, - struct ggml_tensor * state, - struct ggml_tensor * parent_ids) { - struct ggml_tensor * result = ggml_gated_delta_net(ctx, q, k, v, g, beta, state); - - GGML_ASSERT(parent_ids != NULL); - GGML_ASSERT(parent_ids->type == GGML_TYPE_I32); - GGML_ASSERT(ggml_is_contiguous(parent_ids)); - - const int64_t n_tokens = v->ne[2]; - const int64_t n_seqs = v->ne[3]; - GGML_ASSERT(ggml_nelements(parent_ids) == n_tokens * n_seqs); - - result->src[6] = parent_ids; - - return result; -} - -// dflash: tree-mode + external persistent intermediate buffer. The -// kernel writes per-token intermediate states DIRECTLY into persist_inter's -// memory, skipping the cost of copying them out of the result tensor's -// internal region after graph_compute. -struct ggml_tensor * ggml_gated_delta_net_tree_persist( - struct ggml_context * ctx, - struct ggml_tensor * q, - struct ggml_tensor * k, - struct ggml_tensor * v, - struct ggml_tensor * g, - struct ggml_tensor * beta, - struct ggml_tensor * state, - struct ggml_tensor * parent_ids, - struct ggml_tensor * persist_inter) { - struct ggml_tensor * result = ggml_gated_delta_net_tree( - ctx, q, k, v, g, beta, state, parent_ids); - - GGML_ASSERT(persist_inter != NULL); - GGML_ASSERT(persist_inter->type == GGML_TYPE_F32 || - persist_inter->type == GGML_TYPE_F16); - GGML_ASSERT(ggml_is_contiguous(persist_inter)); - - const int64_t S_v = v->ne[0]; - const int64_t H = v->ne[1]; - const int64_t n_tokens = v->ne[2]; - const int64_t n_seqs = v->ne[3]; - GGML_ASSERT(ggml_nelements(persist_inter) >= S_v * S_v * H * n_tokens * n_seqs); - - result->src[7] = persist_inter; - - return result; -} - -//////////////////////////////////////////////////////////////////////////////// - -struct ggml_hash_set ggml_hash_set_new(size_t size) { - size = ggml_hash_size(size); - struct ggml_hash_set result; - result.size = size; - result.keys = GGML_MALLOC(sizeof(struct ggml_tensor *) * size); - result.used = GGML_CALLOC(ggml_bitset_size(size), sizeof(ggml_bitset_t)); - return result; -} - -void ggml_hash_set_reset(struct ggml_hash_set * hash_set) { - memset(hash_set->used, 0, sizeof(ggml_bitset_t) * ggml_bitset_size(hash_set->size)); -} - -void ggml_hash_set_free(struct ggml_hash_set * hash_set) { - GGML_FREE(hash_set->used); - GGML_FREE(hash_set->keys); -} - -size_t ggml_hash_size(size_t min_sz) { - // next primes after powers of two - static const size_t primes[] = { - 2, 3, 5, 11, 17, 37, 67, 131, 257, 521, 1031, - 2053, 4099, 8209, 16411, 32771, 65537, 131101, - 262147, 524309, 1048583, 2097169, 4194319, 8388617, - 16777259, 33554467, 67108879, 134217757, 268435459, - 536870923, 1073741827, 2147483659 - }; - static const size_t n_primes = sizeof(primes)/sizeof(primes[0]); - - // find the smallest prime that is larger or equal than min_sz - size_t l = 0; - size_t r = n_primes; - while (l < r) { - size_t m = (l + r)/2; - if (primes[m] < min_sz) { - l = m + 1; - } else { - r = m; - } - } - size_t sz = l < n_primes ? primes[l] : min_sz | 1; - return sz; -} - -struct hash_map { - struct ggml_hash_set set; - struct ggml_tensor ** vals; -}; - -static struct hash_map * ggml_new_hash_map(size_t size) { - struct hash_map * result = GGML_MALLOC(sizeof(struct hash_map)); - result->set = ggml_hash_set_new(size); - result->vals = GGML_CALLOC(result->set.size, sizeof(struct ggml_tensor *)); - return result; -} - -static void ggml_hash_map_free(struct hash_map * map) { - ggml_hash_set_free(&map->set); - GGML_FREE(map->vals); - GGML_FREE(map); -} - -// utility functions to change gradients -// isrc is the index of tensor in cgraph->visited_has_set.keys -// the corresponding gradient (accumulators) are also at position isrc -// if tensor has a gradient accumulator, modify that accumulator in-place -// else if there is no gradient for tensor, set the corresponding value -// else, just add/subtract/etc. the gradients - -static void ggml_add_or_set( - struct ggml_context * ctx, - struct ggml_cgraph * cgraph, - size_t isrc, - struct ggml_tensor * tensor) { - struct ggml_tensor * src = cgraph->visited_hash_set.keys[isrc]; - GGML_ASSERT(src); - if (cgraph->grads[isrc]) { - cgraph->grads[isrc] = ggml_add_impl(ctx, cgraph->grads[isrc], tensor, /*inplace =*/ cgraph->grad_accs[isrc]); - } else { - cgraph->grads[isrc] = tensor; - } - ggml_format_name(cgraph->grads[isrc], "grad for %s", src->name); - ggml_build_forward_expand(cgraph, cgraph->grads[isrc]); -} - -static void ggml_acc_or_set( - struct ggml_context * ctx, - struct ggml_cgraph * cgraph, - size_t isrc, - struct ggml_tensor * tensor, - const size_t nb1, - const size_t nb2, - const size_t nb3, - const size_t offset) { - struct ggml_tensor * src = cgraph->visited_hash_set.keys[isrc]; - GGML_ASSERT(src); - if (cgraph->grads[isrc]) { - cgraph->grads[isrc] = ggml_acc_impl(ctx, cgraph->grads[isrc], tensor, nb1, nb2, nb3, offset, cgraph->grad_accs[isrc]); - } else { - struct ggml_tensor * a_zero = ggml_scale(ctx, src, 0.0f); // FIXME this is going to produce NaN if a contains inf/NaN - cgraph->grads[isrc] = ggml_acc_impl(ctx, a_zero, tensor, nb1, nb2, nb3, offset, false); - } - ggml_format_name(cgraph->grads[isrc], "grad for %s", cgraph->visited_hash_set.keys[isrc]->name); - ggml_build_forward_expand(cgraph, cgraph->grads[isrc]); -} - -static void ggml_add1_or_set( - struct ggml_context * ctx, - struct ggml_cgraph * cgraph, - size_t isrc, - struct ggml_tensor * tensor) { - struct ggml_tensor * src = cgraph->visited_hash_set.keys[isrc]; - GGML_ASSERT(src); - if (cgraph->grads[isrc]) { - cgraph->grads[isrc] = ggml_add1_impl(ctx, cgraph->grads[isrc], tensor, cgraph->grad_accs[isrc]); - } else { - cgraph->grads[isrc] = ggml_repeat(ctx, tensor, src); - } - ggml_format_name(cgraph->grads[isrc], "grad for %s", src->name); - ggml_build_forward_expand(cgraph, cgraph->grads[isrc]); -} - -static void ggml_sub_or_set( - struct ggml_context * ctx, - struct ggml_cgraph * cgraph, - size_t isrc, - struct ggml_tensor * tensor) { - struct ggml_tensor * src = cgraph->visited_hash_set.keys[isrc]; - GGML_ASSERT(src); - if (cgraph->grads[isrc]) { - cgraph->grads[isrc] = ggml_sub_impl(ctx, cgraph->grads[isrc], tensor, cgraph->grad_accs[isrc]); - } else { - cgraph->grads[isrc] = ggml_neg(ctx, tensor); - } - ggml_format_name(cgraph->grads[isrc], "grad for %s", src->name); - ggml_build_forward_expand(cgraph, cgraph->grads[isrc]); -} - -static void ggml_compute_backward( - struct ggml_context * ctx, struct ggml_cgraph * cgraph, int i, const bool * grads_needed) { - struct ggml_tensor * tensor = cgraph->nodes[i]; - struct ggml_tensor * grad = ggml_graph_get_grad(cgraph, tensor); - - if (!grad) { - return; - } - - struct ggml_tensor * src0 = tensor->src[0]; - struct ggml_tensor * src1 = tensor->src[1]; - struct ggml_tensor * src2 = tensor->src[2]; - struct ggml_hash_set * hash_set = &cgraph->visited_hash_set; - const size_t isrc0 = src0 ? ggml_hash_find(hash_set, src0) : (size_t) -1; - const size_t isrc1 = src1 ? ggml_hash_find(hash_set, src1) : (size_t) -1; - const size_t isrc2 = src2 ? ggml_hash_find(hash_set, src2) : (size_t) -1; - const bool src0_needs_grads = src0 && isrc0 != GGML_HASHSET_FULL && ggml_bitset_get(hash_set->used, isrc0) && grads_needed[isrc0]; - const bool src1_needs_grads = src1 && isrc1 != GGML_HASHSET_FULL && ggml_bitset_get(hash_set->used, isrc1) && grads_needed[isrc1]; - const bool src2_needs_grads = src2 && isrc2 != GGML_HASHSET_FULL && ggml_bitset_get(hash_set->used, isrc2) && grads_needed[isrc2]; - - switch (tensor->op) { - case GGML_OP_DUP: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, grad); - } - } break; - case GGML_OP_ADD: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, grad); - } - if (src1_needs_grads) { - struct ggml_tensor * tmp = grad; - if (!ggml_are_same_shape(src0, src1)) { - tmp = ggml_repeat_back(ctx, tmp, src1); - } - ggml_add_or_set(ctx, cgraph, isrc1, tmp); - } - } break; - case GGML_OP_ADD1: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, grad); - } - if (src1_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc1, ggml_mean(ctx, grad)); // TODO: should probably be sum instead of mean - } - } break; - case GGML_OP_ACC: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, grad); - } - if (src1_needs_grads) { - const size_t nb1 = ((int32_t *) tensor->op_params)[0]; - const size_t nb2 = ((int32_t *) tensor->op_params)[1]; - const size_t nb3 = ((int32_t *) tensor->op_params)[2]; - const size_t offset = ((int32_t *) tensor->op_params)[3]; - - struct ggml_tensor * tensor_grad_view = ggml_view_4d(ctx, - grad, src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3], - nb1, nb2, nb3, offset); - - ggml_add_or_set(ctx, cgraph, isrc1, ggml_reshape(ctx, ggml_cont(ctx, tensor_grad_view), src1)); - } - } break; - case GGML_OP_SUB: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, grad); - } - if (src1_needs_grads) { - ggml_sub_or_set(ctx, cgraph, isrc1, grad); - } - } break; - case GGML_OP_MUL: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_mul(ctx, grad, src1)); - } - if (src1_needs_grads) { - struct ggml_tensor * tmp = ggml_mul(ctx, src0, grad); - if (!ggml_are_same_shape(src0, src1)) { - tmp = ggml_repeat_back(ctx, tmp, src1); - } - ggml_add_or_set(ctx, cgraph, isrc1, tmp); - } - } break; - case GGML_OP_DIV: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_div(ctx, grad, src1)); - } - if (src1_needs_grads) { - ggml_sub_or_set(ctx, cgraph, isrc1, ggml_mul(ctx, grad, ggml_div(ctx, tensor, src1))); - } - } break; - case GGML_OP_SQR: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_scale(ctx, ggml_mul(ctx, src0, grad), 2.0f)); - } - } break; - case GGML_OP_SQRT: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_scale(ctx, ggml_div(ctx, grad, tensor), 0.5f)); - } - } break; - case GGML_OP_LOG: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_div(ctx, grad, src0)); - } - } break; - case GGML_OP_SIN: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_mul(ctx, grad, ggml_cos(ctx, src0))); - } - } break; - case GGML_OP_COS: { - if (src0_needs_grads) { - ggml_sub_or_set(ctx, cgraph, isrc0, ggml_mul(ctx, grad, ggml_sin(ctx, src0))); - } - } break; - case GGML_OP_SUM: { - if (src0_needs_grads) { - ggml_add1_or_set(ctx, cgraph, isrc0, grad); - } - } break; - case GGML_OP_SUM_ROWS: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_repeat(ctx, grad, src0)); - } - } break; - case GGML_OP_MEAN: { - if (src0_needs_grads) { - ggml_add1_or_set(ctx, cgraph, isrc0, ggml_scale_impl(ctx, grad, 1.0f/src0->ne[0], 0.0, false)); - } - } break; - case GGML_OP_REPEAT: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_repeat_back(ctx, grad, src0)); - } - } break; - case GGML_OP_REPEAT_BACK: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_repeat(ctx, grad, src0)); - } - } break; - case GGML_OP_RMS_NORM: { - if (src0_needs_grads) { - float eps; - memcpy(&eps, tensor->op_params, sizeof(float)); - ggml_add_or_set(ctx, cgraph, isrc0, ggml_rms_norm_back(ctx, grad, src0, eps)); - } - } break; - case GGML_OP_MUL_MAT: { - // https://cs231n.github.io/optimization-2/#staged - // # forward pass - // s0 = np.random.randn(5, 10) - // s1 = np.random.randn(10, 3) - // t = s0.dot(s1) - - // # now suppose we had the gradient on t from above in the circuit - // dt = np.random.randn(*t.shape) # same shape as t - // ds0 = dt.dot(s1.T) #.T gives the transpose of the matrix - // ds1 = t.T.dot(dt) - - // tensor.shape [m,p,qq,rr] - // src0.shape [n,m,q1,r1] - // src1.shape [n,p,qq,rr] - - if (src0_needs_grads) { - GGML_ASSERT(grad->ne[2] == src1->ne[2]); - GGML_ASSERT(grad->ne[3] == src1->ne[3]); - struct ggml_tensor * tmp = - ggml_out_prod(ctx, // [n,m,qq,rr] - src1, // [n,p,qq,rr] - grad); // [m,p,qq,rr] - if (!ggml_are_same_shape(tmp, src0)) { - GGML_ASSERT(tmp->ne[0] == src0->ne[0]); - GGML_ASSERT(tmp->ne[1] == src0->ne[1]); - GGML_ASSERT(tmp->ne[3] == 1); - - const int64_t nr2 = tmp->ne[2] / src0->ne[2]; - const size_t nb2 = tmp->nb[2] * nr2; - const size_t nb3 = tmp->nb[2]; - - tmp = ggml_view_4d(ctx, tmp, src0->ne[0], src0->ne[1], src0->ne[2], nr2, tmp->nb[1], nb2, nb3, 0); - tmp = ggml_repeat_back(ctx, tmp, src0); - } - ggml_add_or_set(ctx, cgraph, isrc0, tmp); - } - if (src1_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc1, - // ggml_mul_mat(ctx, // [n,p,qq,rr] - // ggml_cont(ctx, // [m,n,q1,r1] - // ggml_transpose(ctx, src0)), // [m,n,q1,r1] - // grad), // [m,p,qq,rr] - - // when src0 is bigger than tensor->grad (this is mostly the case in llama), - // avoid transpose of src0, rather transpose smaller tensor->grad - // and then use ggml_out_prod - ggml_out_prod(ctx, // [n,p,qq,rr] - src0, // [n,m,q1,r1] - ggml_transpose(ctx, // [p,m,qq,rr] - grad))); // [m,p,qq,rr] - } - } break; - case GGML_OP_SCALE: { - if (src0_needs_grads) { - float s; - memcpy(&s, tensor->op_params, sizeof(float)); - ggml_add_or_set(ctx, cgraph, isrc0, ggml_scale_impl(ctx, grad, s, 0.0, false)); - } - } break; - case GGML_OP_SET: { - const size_t nb1 = ((const int32_t *) tensor->op_params)[0]; - const size_t nb2 = ((const int32_t *) tensor->op_params)[1]; - const size_t nb3 = ((const int32_t *) tensor->op_params)[2]; - const size_t offset = ((const int32_t *) tensor->op_params)[3]; - - struct ggml_tensor * tensor_grad_view = NULL; - - if (src0_needs_grads || src1_needs_grads) { - GGML_ASSERT(src0->type == tensor->type); - GGML_ASSERT(!cgraph->grads[isrc0] || cgraph->grads[isrc0]->type == grad->type); - GGML_ASSERT(!cgraph->grads[isrc1] || !src1_needs_grads || cgraph->grads[isrc1]->type == grad->type); - - tensor_grad_view = ggml_view_4d(ctx, - grad, src1->ne[0], src1->ne[1], src1->ne[2], src1->ne[3], - nb1, nb2, nb3, offset); - } - - if (src0_needs_grads) { - struct ggml_tensor * tmp = ggml_neg(ctx, tensor_grad_view); - ggml_add_or_set(ctx, cgraph, isrc0, ggml_acc_impl(ctx, grad, tmp, nb1, nb2, nb3, offset, false)); - } - - if (src1_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc1, ggml_reshape(ctx, ggml_cont(ctx, tensor_grad_view), src1)); - } - } break; - case GGML_OP_CPY: { - // cpy overwrites value of src1 by src0 and returns view(src1) - // the overwriting is mathematically equivalent to: - // tensor = src0 * 1 + src1 * 0 - if (src0_needs_grads) { - // dsrc0 = dtensor * 1 - ggml_add_or_set(ctx, cgraph, isrc0, ggml_reshape(ctx, grad, src0)); - } - if (src1_needs_grads) { - // dsrc1 = dtensor * 0 -> noop - } - } break; - case GGML_OP_CONT: { - // same as cpy - if (src0_needs_grads) { - GGML_ASSERT(!cgraph->grads[isrc0] || ggml_is_contiguous(cgraph->grads[isrc0])); - GGML_ASSERT(ggml_is_contiguous(grad)); - GGML_ASSERT(ggml_nelements(tensor) == ggml_nelements(src0)); - ggml_add_or_set(ctx, cgraph, isrc0, - ggml_are_same_shape(tensor, src0) ? grad : ggml_reshape(ctx, grad, src0)); - } - } break; - case GGML_OP_RESHAPE: { - if (src0_needs_grads) { - struct ggml_tensor * grad_cont = ggml_is_contiguous(grad) ? grad : ggml_cont(ctx, grad); - ggml_add_or_set(ctx, cgraph, isrc0, ggml_reshape(ctx, grad_cont, src0)); - } - } break; - case GGML_OP_VIEW: { - if (src0_needs_grads) { - size_t offset; - - memcpy(&offset, tensor->op_params, sizeof(offset)); - - size_t nb1 = tensor->nb[1]; - size_t nb2 = tensor->nb[2]; - size_t nb3 = tensor->nb[3]; - - if (cgraph->grads[isrc0] && src0->type != cgraph->grads[isrc0]->type) { - // gradient is typically F32, but src0 could be other type - size_t ng = ggml_element_size(cgraph->grads[isrc0]); - size_t n0 = ggml_element_size(src0); - GGML_ASSERT(offset % n0 == 0); - GGML_ASSERT(nb1 % n0 == 0); - GGML_ASSERT(nb2 % n0 == 0); - GGML_ASSERT(nb3 % n0 == 0); - offset = (offset / n0) * ng; - nb1 = (nb1 / n0) * ng; - nb2 = (nb2 / n0) * ng; - nb3 = (nb3 / n0) * ng; - } - - ggml_acc_or_set(ctx, cgraph, isrc0, grad, nb1, nb2, nb3, offset); - } - } break; - case GGML_OP_PERMUTE: { - if (src0_needs_grads) { - const int32_t * axes = (const int32_t *) tensor->op_params; - const int axis0 = axes[0] & 0x3; - const int axis1 = axes[1] & 0x3; - const int axis2 = axes[2] & 0x3; - const int axis3 = axes[3] & 0x3; - int axb[4] = {0,0,0,0}; // axes backward - axb[axis0] = 0; - axb[axis1] = 1; - axb[axis2] = 2; - axb[axis3] = 3; - ggml_add_or_set(ctx, cgraph, isrc0, ggml_permute(ctx, grad, axb[0], axb[1], axb[2], axb[3])); - } - } break; - case GGML_OP_TRANSPOSE: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_transpose(ctx, grad)); - } - } break; - case GGML_OP_GET_ROWS: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_get_rows_back(ctx, grad, src1, src0)); - } - if (src1_needs_grads) { - // noop - } - } break; - case GGML_OP_DIAG_MASK_INF: { - if (src0_needs_grads) { - /* ggml_diag_mask_inf_impl() shouldn't be here */ - /* ref: https://github.com/ggml-org/llama.cpp/pull/4203#discussion_r1412377992 */ - const int n_past = ((const int32_t *) tensor->op_params)[0]; - ggml_add_or_set(ctx, cgraph, isrc0, ggml_diag_mask_zero_impl(ctx, grad, n_past, false)); - } - } break; - case GGML_OP_DIAG_MASK_ZERO: { - if (src0_needs_grads) { - const int n_past = ((const int32_t *) tensor->op_params)[0]; - ggml_add_or_set(ctx, cgraph, isrc0, ggml_diag_mask_zero_impl(ctx, grad, n_past, false)); - } - } break; - case GGML_OP_SOFT_MAX: { - if (src0_needs_grads) { - float scale = 1.0f; - float max_bias = 0.0f; - - memcpy(&scale, (const float *) tensor->op_params + 0, sizeof(float)); - memcpy(&max_bias, (const float *) tensor->op_params + 1, sizeof(float)); - - ggml_add_or_set(ctx, cgraph, isrc0, ggml_soft_max_ext_back(ctx, grad, tensor, scale, max_bias)); - } - GGML_ASSERT((!src1 || !src1_needs_grads) && "backward pass for softmax mask not implemented"); - } break; - case GGML_OP_ROPE: { - if (src0_needs_grads) { - //const int n_past = ((int32_t *) tensor->op_params)[0]; - const int n_dims = ((const int32_t *) tensor->op_params)[1]; - const int mode = ((const int32_t *) tensor->op_params)[2]; - //const int n_ctx = ((int32_t *) tensor->op_params)[3]; - const int n_ctx_orig = ((const int32_t *) tensor->op_params)[4]; - float freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow; - int sections[4] = {0, 0, 0, 0}; - - memcpy(&freq_base, (const float *) tensor->op_params + 5, sizeof(float)); - memcpy(&freq_scale, (const float *) tensor->op_params + 6, sizeof(float)); - memcpy(&ext_factor, (const float *) tensor->op_params + 7, sizeof(float)); - memcpy(&attn_factor, (const float *) tensor->op_params + 8, sizeof(float)); - memcpy(&beta_fast, (const float *) tensor->op_params + 9, sizeof(float)); - memcpy(&beta_slow, (const float *) tensor->op_params + 10, sizeof(float)); - memcpy(§ions, tensor->op_params + 11, sizeof(sections)); - - struct ggml_tensor * rope_back = grad->ne[2] == src1->ne[0] ? - ggml_rope_ext_back(ctx, grad, src1, src2, n_dims, - mode, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow) : - ggml_rope_multi_back(ctx, grad, src1, src2, n_dims, sections, - mode, n_ctx_orig, freq_base, freq_scale, ext_factor, attn_factor, beta_fast, beta_slow); - ggml_add_or_set(ctx, cgraph, isrc0, rope_back); - } - GGML_ASSERT((!src2 || !src2_needs_grads) && "gradients for freq factors not implemented"); - } break; - case GGML_OP_IM2COL: { - if (src1_needs_grads) { - const int32_t s0 = ggml_get_op_params_i32(tensor, 0); - const int32_t s1 = ggml_get_op_params_i32(tensor, 1); - const int32_t p0 = ggml_get_op_params_i32(tensor, 2); - const int32_t p1 = ggml_get_op_params_i32(tensor, 3); - const int32_t d0 = ggml_get_op_params_i32(tensor, 4); - const int32_t d1 = ggml_get_op_params_i32(tensor, 5); - const bool is_2D = ggml_get_op_params_i32(tensor, 6) == 1; - - ggml_add_or_set(ctx, cgraph, isrc1, ggml_im2col_back(ctx, grad, src0, src1->ne, s0, s1, p0, p1, d0, d1, is_2D)); - } - } break; - case GGML_OP_POOL_2D: { - if (src0_needs_grads) { - const enum ggml_op_pool op = ggml_get_op_params_i32(tensor, 0); - const int32_t k0 = ggml_get_op_params_i32(tensor, 1); - const int32_t k1 = ggml_get_op_params_i32(tensor, 2); - const int32_t s0 = ggml_get_op_params_i32(tensor, 3); - const int32_t s1 = ggml_get_op_params_i32(tensor, 4); - const int32_t p0 = ggml_get_op_params_i32(tensor, 5); - const int32_t p1 = ggml_get_op_params_i32(tensor, 6); - - ggml_add_or_set(ctx, cgraph, isrc0, ggml_pool_2d_back(ctx, grad, src0, op, k0, k1, s0, s1, p0, p1)); - } - } break; - case GGML_OP_WIN_PART: - case GGML_OP_WIN_UNPART: - case GGML_OP_UNARY: { - switch (ggml_get_unary_op(tensor)) { - case GGML_UNARY_OP_ABS: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_mul(ctx, ggml_sgn(ctx, src0), grad)); - } - } break; - case GGML_UNARY_OP_SGN: { - // noop - } break; - case GGML_UNARY_OP_NEG: { - if (src0_needs_grads) { - ggml_sub_or_set(ctx, cgraph, isrc0, grad); - } - } break; - case GGML_UNARY_OP_STEP: { - // noop - } break; - case GGML_UNARY_OP_RELU: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_mul(ctx, ggml_step(ctx, src0), grad)); - } - } break; - case GGML_UNARY_OP_SILU: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_silu_back(ctx, grad, src0)); - } - } break; - case GGML_UNARY_OP_EXP: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_mul(ctx, tensor, grad)); - } - } break; - case GGML_UNARY_OP_EXPM1: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_mul(ctx, grad, ggml_exp(ctx, src0))); - } - } break; - case GGML_UNARY_OP_SOFTPLUS: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_mul(ctx, grad, ggml_sigmoid(ctx, src0))); - } - } break; - default: { - fprintf(stderr, "%s: unsupported unary op for backward pass: %s\n", - __func__, ggml_unary_op_name(ggml_get_unary_op(tensor))); - GGML_ABORT("fatal error"); - } //break; - } - } break; - case GGML_OP_CROSS_ENTROPY_LOSS: { - if (src0_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc0, ggml_cross_entropy_loss_back(ctx, grad, src0, src1)); - } - GGML_ASSERT(!src1_needs_grads && "backward pass for labels not implemented"); - } break; - case GGML_OP_GLU: { - switch (ggml_get_glu_op(tensor)) { - case GGML_GLU_OP_SWIGLU: { - if (src0_needs_grads) { - GGML_ASSERT(src1 && "backward pass only implemented for split swiglu"); - ggml_add_or_set(ctx, cgraph, isrc0, ggml_silu_back(ctx, ggml_mul(ctx, grad, src1), src0)); - } - if (src1_needs_grads) { - ggml_add_or_set(ctx, cgraph, isrc1, ggml_mul(ctx, ggml_silu(ctx, src0), grad)); - } - } break; - default: { - GGML_ABORT("unsupported glu op for backward pass: %s", ggml_glu_op_name(ggml_get_glu_op(tensor))); - } //break; - } - } break; - case GGML_OP_NONE: { - // noop - } break; - case GGML_OP_COUNT: - default: { - GGML_ABORT("%s: unsupported ggml op for backward pass: %s\n", __func__, ggml_op_name(tensor->op)); - } //break; - } - - GGML_ASSERT(!src0_needs_grads || ggml_are_same_shape(src0, cgraph->grads[isrc0])); - GGML_ASSERT(!src1_needs_grads || ggml_are_same_shape(src1, cgraph->grads[isrc1])); - GGML_ASSERT(!src2_needs_grads || ggml_are_same_shape(src2, cgraph->grads[isrc2])); -} - -static size_t ggml_visit_parents_graph(struct ggml_cgraph * cgraph, struct ggml_tensor * node, bool compute) { - if (node->op != GGML_OP_NONE && compute) { - node->flags |= GGML_TENSOR_FLAG_COMPUTE; - } - - const size_t node_hash_pos = ggml_hash_find(&cgraph->visited_hash_set, node); - GGML_ASSERT(node_hash_pos != GGML_HASHSET_FULL); - - if (ggml_bitset_get(cgraph->visited_hash_set.used, node_hash_pos)) { - // already visited - - if (compute) { - // update the compute flag regardless - for (int i = 0; i < GGML_MAX_SRC; ++i) { - struct ggml_tensor * src = node->src[i]; - if (src && ((src->flags & GGML_TENSOR_FLAG_COMPUTE) == 0)) { - ggml_visit_parents_graph(cgraph, src, true); - } - } - } - - return node_hash_pos; - } - - // This is the first time we see this node in the current graph. - cgraph->visited_hash_set.keys[node_hash_pos] = node; - ggml_bitset_set(cgraph->visited_hash_set.used, node_hash_pos); - cgraph->use_counts[node_hash_pos] = 0; - - for (int i = 0; i < GGML_MAX_SRC; ++i) { - const int k = - (cgraph->order == GGML_CGRAPH_EVAL_ORDER_LEFT_TO_RIGHT) ? i : - (cgraph->order == GGML_CGRAPH_EVAL_ORDER_RIGHT_TO_LEFT) ? (GGML_MAX_SRC-1-i) : - /* unknown order, just fall back to using i */ i; - - struct ggml_tensor * src = node->src[k]; - if (src) { - const size_t src_hash_pos = ggml_visit_parents_graph(cgraph, src, compute); - - // Update the use count for this operand. - cgraph->use_counts[src_hash_pos]++; - } - } - - if (node->op == GGML_OP_NONE && !(node->flags & GGML_TENSOR_FLAG_PARAM)) { - // reached a leaf node, not part of the gradient graph (e.g. a constant) - GGML_ASSERT(cgraph->n_leafs < cgraph->size); - - if (strlen(node->name) == 0) { - ggml_format_name(node, "leaf_%d", cgraph->n_leafs); - } - - cgraph->leafs[cgraph->n_leafs] = node; - cgraph->n_leafs++; - } else { - GGML_ASSERT(cgraph->n_nodes < cgraph->size); - - if (strlen(node->name) == 0) { - ggml_format_name(node, "node_%d", cgraph->n_nodes); - } - - cgraph->nodes[cgraph->n_nodes] = node; - cgraph->n_nodes++; - } - - return node_hash_pos; -} - -static void ggml_build_forward_impl(struct ggml_cgraph * cgraph, struct ggml_tensor * tensor, bool expand, bool compute) { - if (!expand) { - // TODO: this branch isn't accessible anymore, maybe move this to ggml_build_forward_expand - ggml_graph_clear(cgraph); - } - - const int n_old = cgraph->n_nodes; - - ggml_visit_parents_graph(cgraph, tensor, compute); - - const int n_new = cgraph->n_nodes - n_old; - GGML_PRINT_DEBUG("%s: visited %d new nodes\n", __func__, n_new); - - if (n_new > 0) { - // the last added node should always be starting point - GGML_ASSERT(cgraph->nodes[cgraph->n_nodes - 1] == tensor); - } -} - -struct ggml_tensor * ggml_build_forward_select( - struct ggml_cgraph * cgraph, - struct ggml_tensor ** tensors, - int n_tensors, - int idx) { - GGML_ASSERT(idx >= 0 && idx < n_tensors); - - for (int i = 0; i < n_tensors; i++) { - ggml_build_forward_impl(cgraph, tensors[i], true, i == idx ? true : false); - } - - return tensors[idx]; -} - -void ggml_build_forward_expand(struct ggml_cgraph * cgraph, struct ggml_tensor * tensor) { - ggml_build_forward_impl(cgraph, tensor, true, true); -} - -void ggml_build_backward_expand( - struct ggml_context * ctx, - struct ggml_cgraph * cgraph, - struct ggml_tensor ** grad_accs) { - GGML_ASSERT(cgraph->n_nodes > 0); - GGML_ASSERT(cgraph->grads); - GGML_ASSERT(cgraph->grad_accs); - - const int n_nodes_f = cgraph->n_nodes; - - memset(cgraph->grads, 0, cgraph->visited_hash_set.size*sizeof(struct ggml_tensor *)); - memset(cgraph->grad_accs, 0, cgraph->visited_hash_set.size*sizeof(struct ggml_tensor *)); - bool * grads_needed = calloc(cgraph->visited_hash_set.size, sizeof(bool)); - - { - bool any_params = false; - bool any_loss = false; - for (int i = 0; i < n_nodes_f; ++i) { - struct ggml_tensor * node = cgraph->nodes[i]; - any_params = any_params || (node->flags & GGML_TENSOR_FLAG_PARAM); - any_loss = any_loss || (node->flags & GGML_TENSOR_FLAG_LOSS); - } - GGML_ASSERT(any_params && "no trainable parameters found, did you forget to call ggml_set_param?"); - GGML_ASSERT(any_loss && "no training loss found, did you forget to call ggml_set_loss?"); - } - - for (int i = 0; i < n_nodes_f; ++i) { - struct ggml_tensor * node = cgraph->nodes[i]; - - if (node->type == GGML_TYPE_I32) { - continue; - } - - bool node_needs_grad = (node->flags & GGML_TENSOR_FLAG_PARAM) || (node->flags & GGML_TENSOR_FLAG_LOSS); - bool ignore_src[GGML_MAX_SRC] = {false}; - switch (node->op) { - // gradients in node->src[0] for one reason or another have no effect on output gradients - case GGML_OP_IM2COL: // only used for its shape - case GGML_OP_IM2COL_BACK: // same as IM2COL - ignore_src[0] = true; - break; - case GGML_OP_UNARY: { - const enum ggml_unary_op uop = ggml_get_unary_op(node); - // SGN and STEP unary ops are piecewise constant - if (uop == GGML_UNARY_OP_SGN || uop == GGML_UNARY_OP_STEP) { - ignore_src[0] = true; - } - } break; - - // gradients in node->src[1] for one reason or another have no effect on output gradients - case GGML_OP_CPY: // gradients in CPY target are irrelevant - case GGML_OP_GET_ROWS: // row indices not differentiable - case GGML_OP_GET_ROWS_BACK: // same as for GET_ROWS - case GGML_OP_ROPE: // positions not differentiable - ignore_src[1] = true; - break; - - default: - break; - } - for (int j = 0; j < GGML_MAX_SRC; ++j) { - if (!node->src[j] || ignore_src[j] || !grads_needed[ggml_hash_find(&cgraph->visited_hash_set, node->src[j])]) { - continue; - } - GGML_ASSERT(node->src[j]->type == GGML_TYPE_F32 || node->src[j]->type == GGML_TYPE_F16); - node_needs_grad = true; - break; - } - if (!node_needs_grad) { - continue; - } - - // inplace operations are currently not supported - GGML_ASSERT(!node->view_src || node->op == GGML_OP_CPY || node->op == GGML_OP_VIEW || - node->op == GGML_OP_RESHAPE || node->op == GGML_OP_PERMUTE || node->op == GGML_OP_TRANSPOSE); - - const size_t ihash = ggml_hash_find(&cgraph->visited_hash_set, node); - GGML_ASSERT(ihash != GGML_HASHSET_FULL); - GGML_ASSERT(ggml_bitset_get(cgraph->visited_hash_set.used, ihash)); - if (grad_accs && grad_accs[i]) { - cgraph->grad_accs[ihash] = grad_accs[i]; - cgraph->grads[ihash] = cgraph->grad_accs[ihash]; - } else if (node->flags & GGML_TENSOR_FLAG_LOSS) { - // loss tensors always need a gradient accumulator - cgraph->grad_accs[ihash] = ggml_new_tensor(ctx, GGML_TYPE_F32, GGML_MAX_DIMS, node->ne); - cgraph->grads[ihash] = cgraph->grad_accs[ihash]; - } - grads_needed[ihash] = true; - } - - for (int i = n_nodes_f - 1; i >= 0; --i) { - // inplace operations to add gradients are not created by ggml_compute_backward except for gradient accumulation - // use allocator to automatically make inplace operations - ggml_compute_backward(ctx, cgraph, i, grads_needed); - } - - free(grads_needed); -} - -static void * incr_ptr_aligned(void ** p, size_t size, size_t align) { - void * ptr = *p; - ptr = (void *) GGML_PAD((uintptr_t) ptr, align); - *p = (void *) ((char *) ptr + size); - return ptr; -} - -static size_t ggml_graph_nbytes(size_t size, bool grads) { - size_t hash_size = ggml_hash_size(size * 2); - void * p = 0; - incr_ptr_aligned(&p, sizeof(struct ggml_cgraph), 1); - incr_ptr_aligned(&p, size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); // nodes - incr_ptr_aligned(&p, size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); // leafs - incr_ptr_aligned(&p, hash_size * sizeof(int32_t), sizeof(int32_t)); // use_counts - incr_ptr_aligned(&p, hash_size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); // hash keys - if (grads) { - incr_ptr_aligned(&p, hash_size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); // grads - incr_ptr_aligned(&p, hash_size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); // grad_accs - } - incr_ptr_aligned(&p, ggml_bitset_size(hash_size) * sizeof(ggml_bitset_t), sizeof(ggml_bitset_t)); - - size_t nbytes = (size_t) p; - return nbytes; -} - -size_t ggml_graph_overhead_custom(size_t size, bool grads) { - return GGML_OBJECT_SIZE + GGML_PAD(ggml_graph_nbytes(size, grads), GGML_MEM_ALIGN); -} - -size_t ggml_graph_overhead(void) { - return ggml_graph_overhead_custom(GGML_DEFAULT_GRAPH_SIZE, false); -} - -struct ggml_cgraph * ggml_new_graph_custom(struct ggml_context * ctx, size_t size, bool grads) { - const size_t obj_size = ggml_graph_nbytes(size, grads); - struct ggml_object * obj = ggml_new_object(ctx, GGML_OBJECT_TYPE_GRAPH, obj_size); - struct ggml_cgraph * cgraph = (struct ggml_cgraph *) ((char *) ctx->mem_buffer + obj->offs); - - // the size of the hash table is doubled since it needs to hold both nodes and leafs - size_t hash_size = ggml_hash_size(size * 2); - - void * p = cgraph + 1; - - struct ggml_tensor ** nodes_ptr = incr_ptr_aligned(&p, size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); - struct ggml_tensor ** leafs_ptr = incr_ptr_aligned(&p, size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); - int32_t * use_counts_ptr = incr_ptr_aligned(&p, hash_size * sizeof(int32_t), sizeof(int32_t)); - struct ggml_tensor ** hash_keys_ptr = incr_ptr_aligned(&p, hash_size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)); - struct ggml_tensor ** grads_ptr = grads ? incr_ptr_aligned(&p, hash_size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)) : NULL; - struct ggml_tensor ** grad_accs_ptr = grads ? incr_ptr_aligned(&p, hash_size * sizeof(struct ggml_tensor *), sizeof(struct ggml_tensor *)) : NULL; - - ggml_bitset_t * hash_used = incr_ptr_aligned(&p, ggml_bitset_size(hash_size) * sizeof(ggml_bitset_t), sizeof(ggml_bitset_t)); - - // check that we allocated the correct amount of memory - assert(obj_size == (size_t)((char *)p - (char *)cgraph)); - - *cgraph = (struct ggml_cgraph) { - /*.size =*/ size, - /*.n_nodes =*/ 0, - /*.n_leafs =*/ 0, - /*.nodes =*/ nodes_ptr, - /*.grads =*/ grads_ptr, - /*.grad_accs =*/ grad_accs_ptr, - /*.leafs =*/ leafs_ptr, - /*.use_counts =*/ use_counts_ptr, - /*.hash_table =*/ { hash_size, hash_used, hash_keys_ptr }, - /*.order =*/ GGML_CGRAPH_EVAL_ORDER_LEFT_TO_RIGHT, - }; - - ggml_hash_set_reset(&cgraph->visited_hash_set); - if (grads) { - memset(cgraph->grads, 0, hash_size*sizeof(struct ggml_tensor *)); - memset(cgraph->grad_accs, 0, hash_size*sizeof(struct ggml_tensor *)); - } - - return cgraph; -} - -struct ggml_cgraph * ggml_new_graph(struct ggml_context * ctx) { - return ggml_new_graph_custom(ctx, GGML_DEFAULT_GRAPH_SIZE, false); -} - -struct ggml_cgraph ggml_graph_view(struct ggml_cgraph * cgraph0, int i0, int i1) { - struct ggml_cgraph cgraph = { - /*.size =*/ 0, - /*.n_nodes =*/ i1 - i0, - /*.n_leafs =*/ 0, - /*.nodes =*/ cgraph0->nodes + i0, - /*.grads =*/ NULL, // gradients would need visited_hash_set - /*.grad_accs =*/ NULL, - /*.leafs =*/ NULL, - /*.use_counts =*/ cgraph0->use_counts, - /*.visited_hash_set =*/ cgraph0->visited_hash_set, - /*.order =*/ cgraph0->order, - }; - - return cgraph; -} - -void ggml_graph_cpy(struct ggml_cgraph * src, struct ggml_cgraph * dst) { - GGML_ASSERT(dst->size >= src->n_leafs); - GGML_ASSERT(dst->size >= src->n_nodes); - GGML_ASSERT(dst->visited_hash_set.size >= src->visited_hash_set.size); - - dst->n_leafs = src->n_leafs; - dst->n_nodes = src->n_nodes; - dst->order = src->order; - - for (int i = 0; i < src->n_leafs; ++i) { - dst->leafs[i] = src->leafs[i]; - } - - for (int i = 0; i < src->n_nodes; ++i) { - dst->nodes[i] = src->nodes[i]; - } - - for (size_t i = 0; i < src->visited_hash_set.size; ++i) { - // copy all hashset keys (tensors) that are in use - if (ggml_bitset_get(src->visited_hash_set.used, i)) { - size_t new_hash_pos = ggml_hash_insert(&dst->visited_hash_set, src->visited_hash_set.keys[i]); - dst->use_counts[new_hash_pos] = src->use_counts[i]; - } - } - - if (dst->grads) { - memset(dst->grads, 0, dst->visited_hash_set.size*sizeof(struct ggml_tensor *)); - memset(dst->grad_accs, 0, dst->visited_hash_set.size*sizeof(struct ggml_tensor *)); - } - if (src->grads) { - GGML_ASSERT(dst->grads != NULL); - GGML_ASSERT(dst->grad_accs != NULL); - for (int i = 0; i < src->n_nodes; ++i) { - const size_t igrad_src = ggml_hash_find(&src->visited_hash_set, src->nodes[i]); - const size_t igrad_dst = ggml_hash_find(&dst->visited_hash_set, dst->nodes[i]); - - GGML_ASSERT(igrad_src != GGML_HASHSET_FULL); - GGML_ASSERT(ggml_bitset_get(src->visited_hash_set.used, igrad_src)); - GGML_ASSERT(igrad_dst != GGML_HASHSET_FULL); - GGML_ASSERT(ggml_bitset_get(dst->visited_hash_set.used, igrad_dst)); - - dst->grads[igrad_dst] = src->grads[igrad_src]; - dst->grad_accs[igrad_dst] = src->grad_accs[igrad_src]; - } - } -} - -struct ggml_cgraph * ggml_graph_dup(struct ggml_context * ctx, struct ggml_cgraph * cgraph, bool force_grads) { - struct ggml_cgraph * result = ggml_new_graph_custom(ctx, cgraph->size, cgraph->grads || force_grads); - ggml_graph_cpy(cgraph, result); - return result; -} - -struct ggml_tensor * ggml_set_zero(struct ggml_tensor * tensor) { - if (ggml_is_empty(tensor)) { - return tensor; - } - if (tensor->buffer) { - ggml_backend_tensor_memset(tensor, 0, 0, ggml_nbytes(tensor)); - } else { - GGML_ASSERT(tensor->data); - memset(tensor->data, 0, ggml_nbytes(tensor)); - } - return tensor; -} - -void ggml_graph_reset(struct ggml_cgraph * cgraph) { - if (!cgraph) { - return; - } - GGML_ASSERT(cgraph->grads != NULL); - - for (int i = 0; i < cgraph->n_nodes; i++) { - struct ggml_tensor * node = cgraph->nodes[i]; - struct ggml_tensor * grad_acc = ggml_graph_get_grad_acc(cgraph, node); - - if (node->op == GGML_OP_OPT_STEP_ADAMW) { - // clear momenta - ggml_set_zero(node->src[2]); - ggml_set_zero(node->src[3]); - } - - // initial gradients of loss should be 1, 0 otherwise - if (grad_acc) { - if (node->flags & GGML_TENSOR_FLAG_LOSS) { - GGML_ASSERT(grad_acc->type == GGML_TYPE_F32); - GGML_ASSERT(ggml_is_scalar(grad_acc)); - - const float onef = 1.0f; - if (grad_acc->buffer) { - ggml_backend_tensor_set(grad_acc, &onef, 0, sizeof(float)); - } else { - GGML_ASSERT(grad_acc->data); - *((float *) grad_acc->data) = onef; - } - } else { - ggml_set_zero(grad_acc); - } - } - } -} - -void ggml_graph_clear(struct ggml_cgraph * cgraph) { - cgraph->n_leafs = 0; - cgraph->n_nodes = 0; - ggml_hash_set_reset(&cgraph->visited_hash_set); -} - -int ggml_graph_size(struct ggml_cgraph * cgraph) { - return cgraph->size; -} - -struct ggml_tensor * ggml_graph_node(struct ggml_cgraph * cgraph, int i) { - if (i < 0) { - GGML_ASSERT(cgraph->n_nodes + i >= 0); - return cgraph->nodes[cgraph->n_nodes + i]; - } - - GGML_ASSERT(i < cgraph->n_nodes); - return cgraph->nodes[i]; -} - -struct ggml_tensor ** ggml_graph_nodes(struct ggml_cgraph * cgraph) { - return cgraph->nodes; -} - -int ggml_graph_n_nodes(struct ggml_cgraph * cgraph) { - return cgraph->n_nodes; -} - -void ggml_graph_add_node(struct ggml_cgraph * cgraph, struct ggml_tensor * tensor) { - GGML_ASSERT(cgraph->size > cgraph->n_nodes); - cgraph->nodes[cgraph->n_nodes] = tensor; - cgraph->n_nodes++; -} - -struct ggml_tensor * ggml_graph_get_tensor(const struct ggml_cgraph * cgraph, const char * name) { - for (int i = 0; i < cgraph->n_leafs; i++) { - struct ggml_tensor * leaf = cgraph->leafs[i]; - - if (strcmp(leaf->name, name) == 0) { - return leaf; - } - } - - for (int i = 0; i < cgraph->n_nodes; i++) { - struct ggml_tensor * node = cgraph->nodes[i]; - - if (strcmp(node->name, name) == 0) { - return node; - } - } - - return NULL; -} - -struct ggml_tensor * ggml_graph_get_grad(const struct ggml_cgraph * cgraph, const struct ggml_tensor * node) { - const size_t igrad = ggml_hash_find(&cgraph->visited_hash_set, node); - return igrad != GGML_HASHSET_FULL && ggml_bitset_get(cgraph->visited_hash_set.used, igrad) && cgraph->grads ? cgraph->grads[igrad] : NULL; -} - -struct ggml_tensor * ggml_graph_get_grad_acc(const struct ggml_cgraph * cgraph, const struct ggml_tensor * node) { - const size_t igrad = ggml_hash_find(&cgraph->visited_hash_set, node); - return igrad != GGML_HASHSET_FULL && ggml_bitset_get(cgraph->visited_hash_set.used, igrad) && cgraph->grad_accs ? cgraph->grad_accs[igrad] : NULL; -} - -void ggml_graph_print(const struct ggml_cgraph * cgraph) { - GGML_LOG_INFO("=== GRAPH ===\n"); - - GGML_LOG_INFO("n_nodes = %d\n", cgraph->n_nodes); - for (int i = 0; i < cgraph->n_nodes; i++) { - struct ggml_tensor * node = cgraph->nodes[i]; - - GGML_LOG_INFO(" - %3d: [ %5" PRId64 ", %5" PRId64 ", %5" PRId64 "] %16s %s\n", - i, - node->ne[0], node->ne[1], node->ne[2], - ggml_op_name(node->op), (node->flags & GGML_TENSOR_FLAG_PARAM) ? "x" : - ggml_graph_get_grad(cgraph, node) ? "g" : " "); - } - - GGML_LOG_INFO("n_leafs = %d\n", cgraph->n_leafs); - for (int i = 0; i < cgraph->n_leafs; i++) { - struct ggml_tensor * node = cgraph->leafs[i]; - - GGML_LOG_INFO(" - %3d: [ %5" PRId64 ", %5" PRId64 "] %8s %16s\n", - i, - node->ne[0], node->ne[1], - ggml_op_name(node->op), - ggml_get_name(node)); - } - - GGML_LOG_INFO("========================================\n"); -} - -static int ggml_node_list_find_tensor(const struct ggml_cgraph * cgraph, - const int * idxs, - int count, - const struct ggml_tensor * tensor) { - GGML_ASSERT(cgraph && idxs); - for (int i = 0; i < count; ++i) { - const int node_idx = idxs[i]; - - if (node_idx >= cgraph->n_nodes) { - return -1; - } - if (cgraph->nodes[node_idx] == tensor) { - return i; - } - } - return -1; -} - -bool ggml_can_fuse_subgraph_ext(const struct ggml_cgraph * cgraph, - const int * node_idxs, - int count, - const enum ggml_op * ops, - const int * outputs, - int num_outputs) { - GGML_ASSERT(outputs && num_outputs > 0); - - for (int i = 0; i < count; ++i) { - if (node_idxs[i] >= cgraph->n_nodes) { - return false; - } - - const struct ggml_tensor * node = cgraph->nodes[node_idxs[i]]; - - if (node->op != ops[i]) { - return false; - } - - if ((node->flags & GGML_TENSOR_FLAG_COMPUTE) == 0) { - return false; - } - - if (ggml_node_list_find_tensor(cgraph, outputs, num_outputs, node) != -1) { - continue; - } - - if (node->flags & GGML_TENSOR_FLAG_OUTPUT) { - return false; - } - - int subgraph_uses = 0; - for (int j = i + 1; j < count; ++j) { - const struct ggml_tensor * other_node = cgraph->nodes[node_idxs[j]]; - for (int src_idx = 0; src_idx < GGML_MAX_SRC; src_idx++) { - if (other_node->src[src_idx] == node) { - subgraph_uses++; - } - } - } - - if (subgraph_uses != ggml_node_get_use_count(cgraph, node_idxs[i])) { - return false; - } - - // if node is a view, check if the view_src and all it's parent view_srcs are within the subgraph - struct ggml_tensor * view_src = node->view_src; - while (view_src) { - if (ggml_node_list_find_tensor(cgraph, node_idxs, count, view_src) == -1) { - return false; - } - view_src = view_src->view_src; - } - } - - return true; -} - -// check if node is part of the graph -static bool ggml_graph_find(const struct ggml_cgraph * cgraph, const struct ggml_tensor * node) { - if (cgraph == NULL) { - return true; - } - - for (int i = 0; i < cgraph->n_nodes; i++) { - if (cgraph->nodes[i] == node) { - return true; - } - } - - return false; -} - -static struct ggml_tensor * ggml_graph_get_parent(const struct ggml_cgraph * cgraph, const struct ggml_tensor * node) { - for (int i = 0; i < cgraph->n_nodes; i++) { - struct ggml_tensor * parent = cgraph->nodes[i]; - struct ggml_tensor * grad = ggml_graph_get_grad(cgraph, parent); - - if (grad == node) { - return parent; - } - } - - return NULL; -} - -static void ggml_graph_dump_dot_node_edge(FILE * fp, const struct ggml_cgraph * gb, struct ggml_tensor * node, struct ggml_tensor * parent, const char * label) { - struct ggml_tensor * gparent = ggml_graph_get_parent(gb, node); - struct ggml_tensor * gparent0 = ggml_graph_get_parent(gb, parent); - fprintf(fp, " \"%p\" -> \"%p\" [ arrowhead = %s; style = %s; label = \"%s\"; ]\n", - gparent0 ? (void *) gparent0 : (void *) parent, - gparent ? (void *) gparent : (void *) node, - gparent ? "empty" : "vee", - gparent ? "dashed" : "solid", - label); -} - -static void ggml_graph_dump_dot_leaf_edge(FILE * fp, struct ggml_tensor * node, struct ggml_tensor * parent, const char * label) { - fprintf(fp, " \"%p\" -> \"%p\" [ label = \"%s\"; ]\n", - (void *) parent, - (void *) node, - label); -} - -void ggml_graph_dump_dot(const struct ggml_cgraph * gb, const struct ggml_cgraph * cgraph, const char * filename) { - char color[16]; - - FILE * fp = ggml_fopen(filename, "w"); - GGML_ASSERT(fp); - - fprintf(fp, "digraph G {\n"); - fprintf(fp, " newrank = true;\n"); - fprintf(fp, " rankdir = TB;\n"); - - for (int i = 0; i < gb->n_nodes; i++) { - struct ggml_tensor * node = gb->nodes[i]; - struct ggml_tensor * grad = ggml_graph_get_grad(gb, node); - - if (ggml_graph_get_parent(gb, node) != NULL) { - continue; - } - - if (node->flags & GGML_TENSOR_FLAG_PARAM) { - snprintf(color, sizeof(color), "yellow"); - } else if (grad) { - if (ggml_graph_find(cgraph, node)) { - snprintf(color, sizeof(color), "green"); - } else { - snprintf(color, sizeof(color), "lightblue"); - } - } else { - snprintf(color, sizeof(color), "white"); - } - - fprintf(fp, " \"%p\" [ " - "style = filled; fillcolor = %s; shape = record; " - "label=\"", - (void *) node, color); - - if (strlen(node->name) > 0) { - fprintf(fp, "%s (%s)|", node->name, ggml_type_name(node->type)); - } else { - fprintf(fp, "(%s)|", ggml_type_name(node->type)); - } - - if (ggml_is_matrix(node)) { - fprintf(fp, "%d [%" PRId64 ", %" PRId64 "] | %s", i, node->ne[0], node->ne[1], ggml_op_symbol(node->op)); - } else { - fprintf(fp, "%d [%" PRId64 ", %" PRId64 ", %" PRId64 "] | %s", i, node->ne[0], node->ne[1], node->ne[2], ggml_op_symbol(node->op)); - } - - if (grad) { - fprintf(fp, " | %s\"; ]\n", ggml_op_symbol(grad->op)); - } else { - fprintf(fp, "\"; ]\n"); - } - } - - for (int i = 0; i < gb->n_leafs; i++) { - struct ggml_tensor * node = gb->leafs[i]; - - snprintf(color, sizeof(color), "pink"); - - fprintf(fp, " \"%p\" [ " - "style = filled; fillcolor = %s; shape = record; " - "label=\"", - (void *) node, color); - - if (strlen(node->name) > 0) { - fprintf(fp, "%s (%s)|", node->name, ggml_type_name(node->type)); - } else { - fprintf(fp, "(%s)|", ggml_type_name(node->type)); - } - - fprintf(fp, "CONST %d [%" PRId64 ", %" PRId64 "]", i, node->ne[0], node->ne[1]); - if (ggml_nelements(node) < 5 && node->data != NULL) { - fprintf(fp, " | ("); - for (int j = 0; j < ggml_nelements(node); j++) { - // FIXME: use ggml-backend to obtain the tensor data - //if (node->type == GGML_TYPE_I8 || node->type == GGML_TYPE_I16 || node->type == GGML_TYPE_I32) { - // fprintf(fp, "%d", ggml_get_i32_1d(node, j)); - //} - //else if (node->type == GGML_TYPE_F32 || - // node->type == GGML_TYPE_F16 || - // node->type == GGML_TYPE_BF16) { - // fprintf(fp, "%.1e", (double)ggml_get_f32_1d(node, j)); - //} - //else - { - fprintf(fp, "#"); - } - if (j < ggml_nelements(node) - 1) { - fprintf(fp, ", "); - } - } - fprintf(fp, ")"); - } - fprintf(fp, "\"; ]\n"); - } - - for (int i = 0; i < gb->n_nodes; i++) { - struct ggml_tensor * node = gb->nodes[i]; - - for (int j = 0; j < GGML_MAX_SRC; j++) { - if (node->src[j]) { - char label[16]; - snprintf(label, sizeof(label), "src %d", j); - ggml_graph_dump_dot_node_edge(fp, gb, node, node->src[j], label); - } - } - } - - for (int i = 0; i < gb->n_leafs; i++) { - struct ggml_tensor * node = gb->leafs[i]; - - for (int j = 0; j < GGML_MAX_SRC; j++) { - if (node->src[j]) { - char label[16]; - snprintf(label, sizeof(label), "src %d", j); - ggml_graph_dump_dot_leaf_edge(fp, node, node->src[j], label); - } - } - } - - fprintf(fp, "}\n"); - - fclose(fp); - - GGML_LOG_INFO("%s: dot -Tpng %s -o %s.png && open %s.png\n", __func__, filename, filename, filename); -} - -//////////////////////////////////////////////////////////////////////////////// - -void ggml_set_input(struct ggml_tensor * tensor) { - tensor->flags |= GGML_TENSOR_FLAG_INPUT; -} - -void ggml_set_output(struct ggml_tensor * tensor) { - tensor->flags |= GGML_TENSOR_FLAG_OUTPUT; -} - -void ggml_set_param(struct ggml_tensor * tensor) { - GGML_ASSERT(tensor->op == GGML_OP_NONE); - tensor->flags |= GGML_TENSOR_FLAG_PARAM; -} - -void ggml_set_loss(struct ggml_tensor * tensor) { - GGML_ASSERT(ggml_is_scalar(tensor)); - GGML_ASSERT(tensor->type == GGML_TYPE_F32); - tensor->flags |= GGML_TENSOR_FLAG_LOSS; -} - -//////////////////////////////////////////////////////////////////////////////// - -void ggml_quantize_init(enum ggml_type type) { - ggml_critical_section_start(); - - switch (type) { - case GGML_TYPE_IQ2_XXS: - case GGML_TYPE_IQ2_XS: - case GGML_TYPE_IQ2_S: - case GGML_TYPE_IQ1_S: - case GGML_TYPE_IQ1_M: iq2xs_init_impl(type); break; - case GGML_TYPE_IQ3_XXS: iq3xs_init_impl(256); break; - case GGML_TYPE_IQ3_S: iq3xs_init_impl(512); break; - default: // nothing - break; - } - - ggml_critical_section_end(); -} - -void ggml_quantize_free(void) { - ggml_critical_section_start(); - - iq2xs_free_impl(GGML_TYPE_IQ2_XXS); - iq2xs_free_impl(GGML_TYPE_IQ2_XS); - iq2xs_free_impl(GGML_TYPE_IQ2_S); - iq2xs_free_impl(GGML_TYPE_IQ1_S); - iq2xs_free_impl(GGML_TYPE_IQ1_M); - iq3xs_free_impl(256); - iq3xs_free_impl(512); - - ggml_critical_section_end(); -} - -bool ggml_quantize_requires_imatrix(enum ggml_type type) { - return - type == GGML_TYPE_IQ2_XXS || - type == GGML_TYPE_IQ2_XS || - type == GGML_TYPE_IQ1_S;// || - //type == GGML_TYPE_IQ1_M; -} - -size_t ggml_quantize_chunk( - enum ggml_type type, - const float * src, - void * dst, - int64_t start, - int64_t nrows, - int64_t n_per_row, - const float * imatrix) { - const int64_t n = (int64_t) nrows * n_per_row; - - if (ggml_quantize_requires_imatrix(type)) { - GGML_ASSERT(imatrix != NULL); - } - - GGML_ASSERT(start % type_traits[type].blck_size == 0); - GGML_ASSERT(start % n_per_row == 0); - - ggml_quantize_init(type); // this is noop if already initialized - - const size_t start_row = start / n_per_row; - const size_t row_size = ggml_row_size(type, n_per_row); - - size_t result = 0; - - switch (type) { - case GGML_TYPE_Q1_0: result = quantize_q1_0(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_Q4_0: result = quantize_q4_0(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_Q4_1: result = quantize_q4_1(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_Q5_0: result = quantize_q5_0(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_Q5_1: result = quantize_q5_1(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_Q8_0: result = quantize_q8_0(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_MXFP4: result = quantize_mxfp4(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_NVFP4: result = quantize_nvfp4(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_Q2_K: result = quantize_q2_K(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_Q3_K: result = quantize_q3_K(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_Q4_K: result = quantize_q4_K(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_Q5_K: result = quantize_q5_K(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_Q6_K: result = quantize_q6_K(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_TQ1_0: result = quantize_tq1_0(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_TQ2_0: result = quantize_tq2_0(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_IQ2_XXS: result = quantize_iq2_xxs(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_IQ2_XS: result = quantize_iq2_xs (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_IQ3_XXS: result = quantize_iq3_xxs(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_IQ3_S: result = quantize_iq3_s (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_IQ2_S: result = quantize_iq2_s (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_IQ1_S: result = quantize_iq1_s (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_IQ1_M: result = quantize_iq1_m (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_IQ4_NL: result = quantize_iq4_nl (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_IQ4_XS: result = quantize_iq4_xs (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; - case GGML_TYPE_F16: - { - size_t elemsize = sizeof(ggml_fp16_t); - ggml_fp32_to_fp16_row(src + start, (ggml_fp16_t *)dst + start, n); - result = n * elemsize; - } break; - case GGML_TYPE_BF16: - { - size_t elemsize = sizeof(ggml_bf16_t); - ggml_fp32_to_bf16_row_ref(src + start, (ggml_bf16_t *)dst + start, n); - result = n * elemsize; - } break; - case GGML_TYPE_F32: - { - size_t elemsize = sizeof(float); - result = n * elemsize; - memcpy((uint8_t *)dst + start * elemsize, src + start, result); - } break; - default: - assert(false); - } - - GGML_ASSERT(result == nrows * row_size); - - return result; -} - -//////////////////////////////////////////////////////////////////////////////// - -void ggml_log_get(ggml_log_callback * log_callback, void ** user_data) { - *log_callback = g_logger_state.log_callback; - *user_data = g_logger_state.log_callback_user_data; -} - -void ggml_log_set(ggml_log_callback log_callback, void * user_data) { - g_logger_state.log_callback = log_callback ? log_callback : ggml_log_callback_default; - g_logger_state.log_callback_user_data = user_data; -} - -void ggml_threadpool_params_init(struct ggml_threadpool_params * p, int n_threads) { - p->n_threads = n_threads; - p->prio = 0; // default priority (usually means normal or inherited) - p->poll = 50; // hybrid-polling enabled - p->strict_cpu = false; // no strict placement (all threads share same cpumask) - p->paused = false; // threads are ready to go - memset(p->cpumask, 0, GGML_MAX_N_THREADS); // all-zero means use the default affinity (usually inherited) -} - -struct ggml_threadpool_params ggml_threadpool_params_default(int n_threads) { - struct ggml_threadpool_params p; - ggml_threadpool_params_init(&p, n_threads); - return p; -} - -bool ggml_threadpool_params_match(const struct ggml_threadpool_params * p0, const struct ggml_threadpool_params * p1) { - if (p0->n_threads != p1->n_threads ) return false; - if (p0->prio != p1->prio ) return false; - if (p0->poll != p1->poll ) return false; - if (p0->strict_cpu != p1->strict_cpu ) return false; - return memcmp(p0->cpumask, p1->cpumask, GGML_MAX_N_THREADS) == 0; -} - -struct ggml_tensor * ggml_turbo_wht( - struct ggml_context * ctx, - struct ggml_tensor * a, - int direction) { - GGML_ASSERT(ggml_n_dims(a) >= 2); - struct ggml_tensor * result = ggml_new_tensor(ctx, a->type, GGML_MAX_DIMS, a->ne); - result->op = GGML_OP_TURBO_WHT; - result->src[0] = a; - ggml_set_op_params_i32(result, 0, direction); - return result; -} - -struct ggml_tensor * ggml_moe_fused( - struct ggml_context * ctx, - struct ggml_tensor * input, - struct ggml_tensor * gate_w, - struct ggml_tensor * up_w, - struct ggml_tensor * down_w, - struct ggml_tensor * expert_ids, - struct ggml_tensor * expert_weights, - struct ggml_tensor * sh_gate_w, - struct ggml_tensor * sh_up_w, - struct ggml_tensor * sh_down_w, - struct ggml_tensor * sh_gate_inp_w, - int64_t n_embd, - int64_t ff_dim, - int64_t n_expert_used) { - GGML_ASSERT(ggml_is_contiguous(input)); - GGML_ASSERT(input->type == GGML_TYPE_F32); - - const int64_t ne[4] = { n_embd, 1, 1, 1 }; - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne); - - result->op = GGML_OP_MOE_FUSED; - result->src[0] = input; - result->src[1] = gate_w; - result->src[2] = up_w; - result->src[3] = down_w; - result->src[4] = expert_ids; - result->src[5] = expert_weights; - result->src[6] = sh_gate_w; - result->src[7] = sh_up_w; - result->src[8] = sh_down_w; - result->src[9] = sh_gate_inp_w; - - ggml_set_op_params_i32(result, 0, (int32_t) n_embd); - ggml_set_op_params_i32(result, 1, (int32_t) ff_dim); - ggml_set_op_params_i32(result, 2, (int32_t) n_expert_used); - - return result; -} - -struct ggml_tensor * ggml_laguna_moe_combine( - struct ggml_context * ctx, - struct ggml_tensor * experts, - struct ggml_tensor * expert_weights) { - GGML_ASSERT(experts->type == GGML_TYPE_F32); - GGML_ASSERT(expert_weights->type == GGML_TYPE_F32); - GGML_ASSERT(experts->ne[1] == expert_weights->ne[0]); - GGML_ASSERT(experts->ne[2] == expert_weights->ne[1]); - - const int64_t ne[4] = { experts->ne[0], experts->ne[2], 1, 1 }; - struct ggml_tensor * result = ggml_new_tensor(ctx, GGML_TYPE_F32, 2, ne); - - result->op = GGML_OP_MOE_FUSED; - result->src[0] = experts; - result->src[1] = expert_weights; - - ggml_set_op_params_i32(result, 0, -1); - ggml_set_op_params_i32(result, 1, (int32_t) experts->ne[0]); - ggml_set_op_params_i32(result, 2, (int32_t) experts->ne[1]); - ggml_set_op_params_i32(result, 3, (int32_t) experts->ne[2]); - - return result; -} From 27597a9d7a10660fe2682a4eecb125105a202587 Mon Sep 17 00:00:00 2001 From: mrciffa <49000955+davide221@users.noreply.github.com> Date: Wed, 8 Jul 2026 13:00:37 +0200 Subject: [PATCH 7/7] perf(cuda): keep DS4_HC sinkhorn in registers (20x kernel speedup) The generic sinkhorn's runtime-bound loops forced its c[] matrix into scratch (private, VRAM-backed) memory, so the serial thread-0 sinkhorn cost ~97us per launch (measured on gfx1151; 86 launches per decoded token ~ 8ms/token). A fully-unrolled template variant keeps c[] in registers: 5.3us per launch. Also grid over n_embd instead of a single block and stage mix/base through shared memory. DS4-Flash fused decode on Strix gfx1151: 19.1 -> 23.0 tok/s. --- ggml/src/ggml-cuda/ds4-hc.cu | 164 +++++++++++++++++++++++++++++++++-- 1 file changed, 156 insertions(+), 8 deletions(-) diff --git a/ggml/src/ggml-cuda/ds4-hc.cu b/ggml/src/ggml-cuda/ds4-hc.cu index 1a331bbed485..1c6b8697287c 100644 --- a/ggml/src/ggml-cuda/ds4-hc.cu +++ b/ggml/src/ggml-cuda/ds4-hc.cu @@ -93,6 +93,131 @@ static __device__ void ds4_hc_sinkhorn_split( } } +// Fully-unrolled variant: with compile-time NHC the c[] matrix lives in +// registers. The generic version's runtime-bound loops force c[] into scratch +// (private, VRAM-backed) memory, making the serial sinkhorn ~20x slower +// (97us vs 5us measured on gfx1151). +template +static __device__ void ds4_hc_sinkhorn_split_t( + const float * mix, + const float * base, + float pre_scale, + float post_scale, + float comb_scale, + int iters, + float * split) { + #pragma unroll + for (int i = 0; i < NHC; ++i) { + split[i] = ds4_hc_sigmoid(mix[i] * pre_scale + base[i]) + DS4_HC_SINKHORN_EPS; + } + #pragma unroll + for (int i = 0; i < NHC; ++i) { + split[NHC + i] = 2.0f * ds4_hc_sigmoid(mix[NHC + i] * post_scale + base[NHC + i]); + } + float c[NHC * NHC]; + #pragma unroll + for (int dst_i = 0; dst_i < NHC; ++dst_i) { + float row_max = -1.0e30f; + #pragma unroll + for (int src_i = 0; src_i < NHC; ++src_i) { + const int idx = src_i + dst_i * NHC; + const float v = mix[2 * NHC + idx] * comb_scale + base[2 * NHC + idx]; + c[idx] = v; + row_max = v > row_max ? v : row_max; + } + float row_sum = 0.0f; + #pragma unroll + for (int src_i = 0; src_i < NHC; ++src_i) { + const int idx = src_i + dst_i * NHC; + c[idx] = expf(c[idx] - row_max); + row_sum += c[idx]; + } + const float inv = 1.0f / row_sum; + #pragma unroll + for (int src_i = 0; src_i < NHC; ++src_i) { + c[src_i + dst_i * NHC] = c[src_i + dst_i * NHC] * inv + DS4_HC_SINKHORN_EPS; + } + } + #pragma unroll + for (int src_i = 0; src_i < NHC; ++src_i) { + float sum = 0.0f; + #pragma unroll + for (int dst_i = 0; dst_i < NHC; ++dst_i) sum += c[src_i + dst_i * NHC]; + const float inv = 1.0f / (sum + DS4_HC_SINKHORN_EPS); + #pragma unroll + for (int dst_i = 0; dst_i < NHC; ++dst_i) c[src_i + dst_i * NHC] *= inv; + } + for (int iter = 1; iter < iters; ++iter) { + #pragma unroll + for (int dst_i = 0; dst_i < NHC; ++dst_i) { + float sum = 0.0f; + #pragma unroll + for (int src_i = 0; src_i < NHC; ++src_i) sum += c[src_i + dst_i * NHC]; + const float inv = 1.0f / (sum + DS4_HC_SINKHORN_EPS); + #pragma unroll + for (int src_i = 0; src_i < NHC; ++src_i) c[src_i + dst_i * NHC] *= inv; + } + #pragma unroll + for (int src_i = 0; src_i < NHC; ++src_i) { + float sum = 0.0f; + #pragma unroll + for (int dst_i = 0; dst_i < NHC; ++dst_i) sum += c[src_i + dst_i * NHC]; + const float inv = 1.0f / (sum + DS4_HC_SINKHORN_EPS); + #pragma unroll + for (int dst_i = 0; dst_i < NHC; ++dst_i) c[src_i + dst_i * NHC] *= inv; + } + } + #pragma unroll + for (int i = 0; i < NHC * NHC; ++i) { + split[2 * NHC + i] = c[i]; + } +} + +template +static __global__ void ds4_hc_pre_kernel_t( + const float * __restrict__ mix, + const float * __restrict__ base, + const float * __restrict__ hc_state, + float * __restrict__ dst, + int n_embd, + int iters, + float pre_scale, + float post_scale, + float comb_scale) { + __shared__ float split[DS4_HC_MAX_MIX]; + __shared__ float s_mix[DS4_HC_MAX_MIX]; + __shared__ float s_base[DS4_HC_MAX_MIX]; + constexpr int mix_dim = 2 * NHC + NHC * NHC; + const int tid = threadIdx.x; + + if (tid < mix_dim) { + s_mix[tid] = mix[tid]; + s_base[tid] = base[tid]; + } + __syncthreads(); + + if (tid == 0) { + ds4_hc_sinkhorn_split_t(s_mix, s_base, pre_scale, post_scale, comb_scale, iters, split); + if (blockIdx.x == 0) { + #pragma unroll + for (int i = 0; i < mix_dim; ++i) { + dst[n_embd + i] = split[i]; + } + } + } + __syncthreads(); + + const int d = (int) blockIdx.x * blockDim.x + tid; + if (d < n_embd) { + float acc = 0.0f; + #pragma unroll + for (int h = 0; h < NHC; ++h) { + acc += split[h] * hc_state[(size_t) h * n_embd + d]; + } + dst[d] = acc; + } +} + static __global__ void ds4_hc_pre_kernel( const float * __restrict__ mix, const float * __restrict__ base, @@ -105,18 +230,33 @@ static __global__ void ds4_hc_pre_kernel( float post_scale, float comb_scale) { __shared__ float split[DS4_HC_MAX_MIX]; + __shared__ float s_mix[DS4_HC_MAX_MIX]; + __shared__ float s_base[DS4_HC_MAX_MIX]; const int mix_dim = 2 * n_hc + n_hc * n_hc; const int tid = threadIdx.x; + // Stage mix/base cooperatively: base lives in managed (UMA) memory where + // serial scalar loads cost ~2us each; one parallel coalesced load instead. + if (tid < mix_dim) { + s_mix[tid] = mix[tid]; + s_base[tid] = base[tid]; + } + __syncthreads(); + + // Each block redoes the (tiny) sinkhorn into shared memory so the mix + // loop below can spread across the whole GPU instead of one CU. if (tid == 0) { - ds4_hc_sinkhorn_split(mix, base, pre_scale, post_scale, comb_scale, n_hc, iters, split); - for (int i = 0; i < mix_dim; ++i) { - dst[n_embd + i] = split[i]; + ds4_hc_sinkhorn_split(s_mix, s_base, pre_scale, post_scale, comb_scale, n_hc, iters, split); + if (blockIdx.x == 0) { + for (int i = 0; i < mix_dim; ++i) { + dst[n_embd + i] = split[i]; + } } } __syncthreads(); - for (int d = tid; d < n_embd; d += blockDim.x) { + const int d = (int) blockIdx.x * blockDim.x + tid; + if (d < n_embd) { float acc = 0.0f; for (int h = 0; h < n_hc; ++h) { acc += split[h] * hc_state[(size_t) h * n_embd + d]; @@ -192,10 +332,18 @@ void ggml_cuda_op_ds4_hc(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const float pre_scale = ggml_get_op_params_f32(dst, 4); const float post_scale = ggml_get_op_params_f32(dst, 5); const float comb_scale = ggml_get_op_params_f32(dst, 6); - ds4_hc_pre_kernel<<<1, 256, 0, stream>>>( - (const float *) src0->data, (const float *) src1->data, - (const float *) src2->data, (float *) dst->data, - n_embd, n_hc, iters, pre_scale, post_scale, comb_scale); + const int pre_blocks = (n_embd + 255) / 256; + if (n_hc == 4) { + ds4_hc_pre_kernel_t<4><<>>( + (const float *) src0->data, (const float *) src1->data, + (const float *) src2->data, (float *) dst->data, + n_embd, iters, pre_scale, post_scale, comb_scale); + } else { + ds4_hc_pre_kernel<<>>( + (const float *) src0->data, (const float *) src1->data, + (const float *) src2->data, (float *) dst->data, + n_embd, n_hc, iters, pre_scale, post_scale, comb_scale); + } } break; case 1: { const int total = n_embd * n_hc;