From cec8151711d88ed9e6d2f2a0ce25a8f4614e1cac Mon Sep 17 00:00:00 2001 From: BREW <81044362+b-re-w@users.noreply.github.com> Date: Fri, 16 Oct 2026 03:21:21 +0900 Subject: [PATCH 1/2] vendor candle-kernels with a toolkit-version guard on __hmax_nan/__hmin_nan (#86) CUDA 12.6's cuda_fp16.hpp defines __hmax_nan/__hmin_nan for every arch, and candle-kernels 0.11.0 compatibility.cuh defined them again under __CUDA_ARCH__ < 800, so build-cuda-wheel.yml has been red on every branch since 2026-10-04. candle-kernels is now vendored like the other two forks (cuda12-candle-kernels-0.11.0.patch, vendor_candle.sh produces and --checks three trees, [patch.crates-io] in torch_c), and the block is guarded on the toolkit: < 800 and CUDA < 12.2. The 12.2 threshold follows upstream candle's later fix and is an assumption; a wrong threshold fails loudly. - build-cuda-wheel.yml now triggers on the vendored trees and the lock. - test_int8.py: patch wiring, lock, the exact guard, --check drift refusal. Nullified: stripping the version clause turns three vendoring tests red. Local: cargo build --release --locked passes (no nvcc here; CI builds it). --- .github/workflows/build-cuda-wheel.yml | 4 + scripts/vendor/vendor_candle.sh | 18 +- tests/numerics/test_int8.py | 63 +- torchnative/rust/torch_c/Cargo.lock | 2 - torchnative/rust/torch_c/Cargo.toml | 11 + .../candle-kernels/.cargo_vcs_info.json | 6 + .../rust/vendor/candle-kernels/Cargo.lock | 441 ++ .../rust/vendor/candle-kernels/Cargo.toml | 40 + .../vendor/candle-kernels/Cargo.toml.orig | 15 + .../rust/vendor/candle-kernels/README.md | 4 + .../rust/vendor/candle-kernels/build.rs | 75 + .../rust/vendor/candle-kernels/src/affine.cu | 51 + .../rust/vendor/candle-kernels/src/binary.cu | 114 + .../candle-kernels/src/binary_op_macros.cuh | 69 + .../rust/vendor/candle-kernels/src/cast.cu | 227 + .../candle-kernels/src/compatibility.cuh | 167 + .../rust/vendor/candle-kernels/src/conv.cu | 896 +++ .../vendor/candle-kernels/src/cuda_utils.cuh | 224 + .../rust/vendor/candle-kernels/src/ffi.rs | 631 +++ .../rust/vendor/candle-kernels/src/fill.cu | 87 + .../vendor/candle-kernels/src/indexing.cu | 457 ++ .../rust/vendor/candle-kernels/src/lib.rs | 82 + .../src/mmq_gguf/mmq_common.cuh | 534 ++ .../candle-kernels/src/mmq_gguf/mmq_gguf.cuh | 3957 ++++++++++++++ .../src/mmq_gguf/mmq_instance_q2_k.cu | 136 + .../src/mmq_gguf/mmq_instance_q3_k.cu | 136 + .../src/mmq_gguf/mmq_instance_q4_0.cu | 136 + .../src/mmq_gguf/mmq_instance_q4_1.cu | 136 + .../src/mmq_gguf/mmq_instance_q4_k.cu | 136 + .../src/mmq_gguf/mmq_instance_q5_0.cu | 136 + .../src/mmq_gguf/mmq_instance_q5_1.cu | 136 + .../src/mmq_gguf/mmq_instance_q5_k.cu | 136 + .../src/mmq_gguf/mmq_instance_q6_k.cu | 136 + .../src/mmq_gguf/mmq_instance_q8_0.cu | 136 + .../candle-kernels/src/mmq_gguf/mmq_mma.cuh | 1409 +++++ .../src/mmq_gguf/mmq_quantize.cu | 145 + .../src/mmq_gguf/mmq_vecdotq.cuh | 1269 +++++ .../vendor/candle-kernels/src/mmvq_gguf.cu | 1049 ++++ .../vendor/candle-kernels/src/moe/gguf.cuh | 1438 +++++ .../vendor/candle-kernels/src/moe/moe_gguf.cu | 216 + .../candle-kernels/src/moe/moe_utils.cuh | 188 + .../vendor/candle-kernels/src/moe/moe_wmma.cu | 287 + .../candle-kernels/src/moe/moe_wmma_gguf.cu | 422 ++ .../rust/vendor/candle-kernels/src/ptx.rs | 11 + .../vendor/candle-kernels/src/quantized.cu | 4845 +++++++++++++++++ .../rust/vendor/candle-kernels/src/reduce.cu | 692 +++ .../rust/vendor/candle-kernels/src/sort.cu | 90 + .../rust/vendor/candle-kernels/src/ternary.cu | 71 + .../rust/vendor/candle-kernels/src/unary.cu | 233 + .../vendor/cuda12-candle-kernels-0.11.0.patch | 17 + 50 files changed, 21910 insertions(+), 7 deletions(-) create mode 100644 torchnative/rust/vendor/candle-kernels/.cargo_vcs_info.json create mode 100644 torchnative/rust/vendor/candle-kernels/Cargo.lock create mode 100644 torchnative/rust/vendor/candle-kernels/Cargo.toml create mode 100644 torchnative/rust/vendor/candle-kernels/Cargo.toml.orig create mode 100644 torchnative/rust/vendor/candle-kernels/README.md create mode 100644 torchnative/rust/vendor/candle-kernels/build.rs create mode 100644 torchnative/rust/vendor/candle-kernels/src/affine.cu create mode 100644 torchnative/rust/vendor/candle-kernels/src/binary.cu create mode 100644 torchnative/rust/vendor/candle-kernels/src/binary_op_macros.cuh create mode 100644 torchnative/rust/vendor/candle-kernels/src/cast.cu create mode 100644 torchnative/rust/vendor/candle-kernels/src/compatibility.cuh create mode 100644 torchnative/rust/vendor/candle-kernels/src/conv.cu create mode 100644 torchnative/rust/vendor/candle-kernels/src/cuda_utils.cuh create mode 100644 torchnative/rust/vendor/candle-kernels/src/ffi.rs create mode 100644 torchnative/rust/vendor/candle-kernels/src/fill.cu create mode 100644 torchnative/rust/vendor/candle-kernels/src/indexing.cu create mode 100644 torchnative/rust/vendor/candle-kernels/src/lib.rs create mode 100644 torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_common.cuh create mode 100644 torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_gguf.cuh create mode 100644 torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q2_k.cu create mode 100644 torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q3_k.cu create mode 100644 torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q4_0.cu create mode 100644 torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q4_1.cu create mode 100644 torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q4_k.cu create mode 100644 torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q5_0.cu create mode 100644 torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q5_1.cu create mode 100644 torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q5_k.cu create mode 100644 torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q6_k.cu create mode 100644 torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q8_0.cu create mode 100644 torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_mma.cuh create mode 100644 torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_quantize.cu create mode 100644 torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_vecdotq.cuh create mode 100644 torchnative/rust/vendor/candle-kernels/src/mmvq_gguf.cu create mode 100644 torchnative/rust/vendor/candle-kernels/src/moe/gguf.cuh create mode 100644 torchnative/rust/vendor/candle-kernels/src/moe/moe_gguf.cu create mode 100644 torchnative/rust/vendor/candle-kernels/src/moe/moe_utils.cuh create mode 100644 torchnative/rust/vendor/candle-kernels/src/moe/moe_wmma.cu create mode 100644 torchnative/rust/vendor/candle-kernels/src/moe/moe_wmma_gguf.cu create mode 100644 torchnative/rust/vendor/candle-kernels/src/ptx.rs create mode 100644 torchnative/rust/vendor/candle-kernels/src/quantized.cu create mode 100644 torchnative/rust/vendor/candle-kernels/src/reduce.cu create mode 100644 torchnative/rust/vendor/candle-kernels/src/sort.cu create mode 100644 torchnative/rust/vendor/candle-kernels/src/ternary.cu create mode 100644 torchnative/rust/vendor/candle-kernels/src/unary.cu create mode 100644 torchnative/rust/vendor/cuda12-candle-kernels-0.11.0.patch diff --git a/.github/workflows/build-cuda-wheel.yml b/.github/workflows/build-cuda-wheel.yml index 24f5b9db..694a8975 100644 --- a/.github/workflows/build-cuda-wheel.yml +++ b/.github/workflows/build-cuda-wheel.yml @@ -49,6 +49,10 @@ on: - "torchnative/rust/torch_c/Cargo.toml" - "torchnative/rust/torch_c/build.rs" - "torchnative/rust/torch_c/src/device.rs" + - "torchnative/rust/torch_c/Cargo.lock" + - "torchnative/rust/vendor/candle-kernels/**" + - "torchnative/rust/vendor/cuda12-candle-kernels-0.11.0.patch" + - "torchnative/rust/vendor/candle-core/**" permissions: contents: read diff --git a/scripts/vendor/vendor_candle.sh b/scripts/vendor/vendor_candle.sh index b903162a..2a5162b0 100644 --- a/scripts/vendor/vendor_candle.sh +++ b/scripts/vendor/vendor_candle.sh @@ -3,8 +3,9 @@ # the committed copies are exactly those forks. # # sh scripts/vendor/vendor_candle.sh write torchnative/rust/vendor/candle-core/ -# and torchnative/rust/vendor/candle-metal-kernels/ -# sh scripts/vendor/vendor_candle.sh --check rebuild both in a temp dir and diff +# torchnative/rust/vendor/candle-metal-kernels/ +# and torchnative/rust/vendor/candle-kernels/ +# sh scripts/vendor/vendor_candle.sh --check rebuild all three in a temp dir and diff # # TWO crates, because `DType::I8` is two halves that are useless apart. The # `candle-core` half can build an int8 buffer on Metal; the @@ -29,6 +30,12 @@ # candle-metal-kernels 0.11.0, pinned by the sha256 the lock already # recorded, plus torchnative/rust/vendor/int8-candle-metal-kernels-0.11.0.patch. # +# candle-kernels 0.11.0 (the CUDA kernels), pinned by the sha256 the lock +# recorded, plus torchnative/rust/vendor/cuda12-candle-kernels-0.11.0.patch. +# This third fork has nothing to do with I8: the published compatibility.cuh +# redefines __hmax_nan/__hmin_nan that CUDA 12.2+ headers already define, so +# no CUDA 12.6 build compiles (issue #86). The patch guards that one block. +# # The *published* crate rather than a clone of candle's repository, on purpose: # `cargo package` normalises `Cargo.toml`, so the published crate stands alone. # The repository's `candle-core/Cargo.toml` inherits `version.workspace = true` @@ -60,7 +67,7 @@ esac # and `TORCHNATIVE_CANDLE_CRATE` stay scoped to candle-core, which is what # test_int8.py's wrong-crate refusal test drives; the second crate has its own # pair so a test can drive either without disturbing the other. -crates="candle-core candle-metal-kernels" +crates="candle-core candle-metal-kernels candle-kernels" sha256_for() { case "$1" in @@ -68,6 +75,8 @@ sha256_for() { echo 5ecb245093b0f791b89d3420c3df9c6d49c60ab63ba54db896bf8a3baf486706 ;; candle-metal-kernels) echo 242e83c6acf639bb273c929d73c67a882bb4dd08a140f121096e19ba2f213d3e ;; + candle-kernels) + echo 67450168a281bbb195a14cc85cf164c7a12023a54a03409a3324f476346e0789 ;; *) echo "sha256_for: unknown crate $1" >&2; exit 2 ;; esac } @@ -76,6 +85,7 @@ patch_for() { case "$1" in candle-core) echo "$repo/torchnative/rust/vendor/int8-candle-$version-cpu.patch" ;; candle-metal-kernels) echo "$repo/torchnative/rust/vendor/int8-candle-metal-kernels-$version.patch" ;; + candle-kernels) echo "$repo/torchnative/rust/vendor/cuda12-candle-kernels-$version.patch" ;; *) echo "patch_for: unknown crate $1" >&2; exit 2 ;; esac } @@ -84,6 +94,7 @@ dest_for() { case "$1" in candle-core) echo "${TORCHNATIVE_CANDLE_DIR:-$repo/torchnative/rust/vendor/candle-core}" ;; candle-metal-kernels) echo "${TORCHNATIVE_CANDLE_KERNELS_DIR:-$repo/torchnative/rust/vendor/candle-metal-kernels}" ;; + candle-kernels) echo "${TORCHNATIVE_CANDLE_CUDA_KERNELS_DIR:-$repo/torchnative/rust/vendor/candle-kernels}" ;; esac } @@ -91,6 +102,7 @@ override_for() { case "$1" in candle-core) echo "${TORCHNATIVE_CANDLE_CRATE:-}" ;; candle-metal-kernels) echo "${TORCHNATIVE_CANDLE_KERNELS_CRATE:-}" ;; + candle-kernels) echo "${TORCHNATIVE_CANDLE_CUDA_KERNELS_CRATE:-}" ;; esac } diff --git a/tests/numerics/test_int8.py b/tests/numerics/test_int8.py index 346b3e11..a36b3d02 100644 --- a/tests/numerics/test_int8.py +++ b/tests/numerics/test_int8.py @@ -41,6 +41,7 @@ _SCRIPT = os.path.join(_REPO_ROOT, "scripts", "vendor", "vendor_candle.sh") _FORK = os.path.join(_REPO_ROOT, "torchnative", "rust", "vendor", "candle-core") _KERNELS_FORK = os.path.join(_REPO_ROOT, "torchnative", "rust", "vendor", "candle-metal-kernels") +_CUDA_KERNELS_FORK = os.path.join(_REPO_ROOT, "torchnative", "rust", "vendor", "candle-kernels") # -------------------------------------------------------------------------- @@ -152,8 +153,8 @@ def test_the_committed_metal_kernels_fork_is_the_pinned_crate_plus_the_patch(): assert proc.returncode == 0, "exit %d\n%s\n%s" % (proc.returncode, proc.stdout, proc.stderr) assert ("sha256 242e83c6acf639bb273c929d73c67a882bb4dd08a140f121096e19ba2f213d3e verified" in proc.stdout), proc.stdout - assert proc.stdout.count("byte for byte") == 2, ( - "--check reported %d trees, not 2. Both forks have to be covered, or " + assert proc.stdout.count("byte for byte") == 3, ( + "--check reported %d trees, not 3. Every fork has to be covered, or " "one of them can drift silently.\n%s" % (proc.stdout.count("byte for byte"), proc.stdout)) @@ -196,6 +197,64 @@ def test_the_metal_kernels_script_refuses_a_crate_that_is_not_the_pinned_one(): shutil.rmtree(tmp) +# -------------------------------------------------------------------------- +# the THIRD fork -- `candle-kernels` (CUDA), issue #86 +# +# Not an I8 half: the published `compatibility.cuh` redefines `__hmax_nan` and +# `__hmin_nan`, which CUDA 12.2+ headers define for every architecture, so no +# CUDA 12.6 build compiled. Same discipline as the two forks above. This Mac +# has no nvcc, so these tests can only prove the fork is wired in and carries +# the guard; build-cuda-wheel.yml is the compiler that proves it builds. +# -------------------------------------------------------------------------- + + +def test_the_cuda_kernels_patch_path_is_relative_and_is_the_tree_the_script_writes(): + with open(os.path.join(_CRATE_DIR, "Cargo.toml"), "rb") as f: + manifest = tomllib.load(f) + path = manifest["patch"]["crates-io"]["candle-kernels"]["path"] + assert not os.path.isabs(path), path + resolved = os.path.normpath(os.path.join(_CRATE_DIR, path)) + assert resolved == _CUDA_KERNELS_FORK, (resolved, _CUDA_KERNELS_FORK) + with open(os.path.join(resolved, "Cargo.toml"), "rb") as f: + package = tomllib.load(f)["package"] + assert (package["name"], package["version"]) == ("candle-kernels", "0.11.0"), package + + +def test_the_lock_resolves_cuda_kernels_from_the_fork_not_the_registry(): + with open(os.path.join(_CRATE_DIR, "Cargo.lock"), "rb") as f: + lock = tomllib.load(f) + entries = [p for p in lock["package"] if p["name"] == "candle-kernels"] + assert len(entries) == 1, entries + assert entries[0]["version"] == "0.11.0", entries[0] + assert "source" not in entries[0], ( + "Cargo.lock resolves candle-kernels from %s -- the CUDA build uses the " + "unpatched compatibility.cuh and fails on CUDA 12.2+" % entries[0].get("source")) + + +def test_the_cuda_kernels_fork_guards_hmax_nan_on_the_toolkit_version(): + with open(os.path.join(_CUDA_KERNELS_FORK, "src", "compatibility.cuh"), encoding="utf-8") as f: + text = f.read() + guard = ("#if __CUDA_ARCH__ < 800 && (__CUDACC_VER_MAJOR__ < 12 || " + "(__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ < 2))\n" + "__device__ __forceinline__ __half __hmax_nan(") + assert guard in text, "compatibility.cuh defines __hmax_nan without the CUDA 12.2 guard" + assert "#if __CUDA_ARCH__ < 800\n" not in text, "the unguarded block is still present" + + +def test_the_cuda_kernels_check_refuses_a_tree_that_drifted_from_the_patch(): + tmp = tempfile.mkdtemp(prefix="cuda-k-drift-") + try: + copy = os.path.join(tmp, "candle-kernels") + shutil.copytree(_CUDA_KERNELS_FORK, copy) + with open(os.path.join(copy, "src", "compatibility.cuh"), "a", encoding="utf-8") as f: + f.write("\n// an edit made to the fork and not to the patch\n") + proc = _run_script("--check", env_extra={"TORCHNATIVE_CANDLE_CUDA_KERNELS_DIR": copy}) + assert proc.returncode != 0, proc.stdout + assert "compatibility.cuh" in proc.stderr, proc.stderr + finally: + shutil.rmtree(tmp) + + def test_the_fork_script_refuses_a_crate_that_is_not_the_pinned_one(): tmp = tempfile.mkdtemp(prefix="int8-crate-") try: diff --git a/torchnative/rust/torch_c/Cargo.lock b/torchnative/rust/torch_c/Cargo.lock index 61749b75..fd3f1289 100644 --- a/torchnative/rust/torch_c/Cargo.lock +++ b/torchnative/rust/torch_c/Cargo.lock @@ -160,8 +160,6 @@ dependencies = [ [[package]] name = "candle-kernels" version = "0.11.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "67450168a281bbb195a14cc85cf164c7a12023a54a03409a3324f476346e0789" dependencies = [ "cudaforge", ] diff --git a/torchnative/rust/torch_c/Cargo.toml b/torchnative/rust/torch_c/Cargo.toml index 1cc5732b..51234fd0 100644 --- a/torchnative/rust/torch_c/Cargo.toml +++ b/torchnative/rust/torch_c/Cargo.toml @@ -301,3 +301,14 @@ candle-core = { path = "../vendor/candle-core" } # entry for a crate absent from a target's graph costs that target nothing, # which is why Android, Linux and wasm builds are unaffected. candle-metal-kernels = { path = "../vendor/candle-metal-kernels" } + +# The CUDA kernels, forked for a build fix and not for I8 (issue #86). +# candle-kernels 0.11.0's `src/compatibility.cuh` defines `__hmax_nan` and +# `__hmin_nan` whenever `__CUDA_ARCH__ < 800`, and CUDA 12.2+ `cuda_fp16.hpp` +# defines both for every architecture, so every `.cu` file failed with "has +# already been defined" on the CUDA 12.6 toolkit `build-cuda-wheel.yml` uses. +# `torchnative/rust/vendor/cuda12-candle-kernels-0.11.0.patch` adds a toolkit +# version guard to that one block; vendored, pinned and `--check`ed exactly as +# the two forks above. Reached only through `candle-core`'s `cuda` feature, so +# a target without it pays nothing for the entry. +candle-kernels = { path = "../vendor/candle-kernels" } diff --git a/torchnative/rust/vendor/candle-kernels/.cargo_vcs_info.json b/torchnative/rust/vendor/candle-kernels/.cargo_vcs_info.json new file mode 100644 index 00000000..d2e1804f --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/.cargo_vcs_info.json @@ -0,0 +1,6 @@ +{ + "git": { + "sha1": "31f35b147389700ed2a178ee66a91c3cc25cc80d" + }, + "path_in_vcs": "candle-kernels" +} \ No newline at end of file diff --git a/torchnative/rust/vendor/candle-kernels/Cargo.lock b/torchnative/rust/vendor/candle-kernels/Cargo.lock new file mode 100644 index 00000000..efdca174 --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/Cargo.lock @@ -0,0 +1,441 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "anyhow" +version = "1.0.103" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2a4385e2e34eb35d6b3efe798b9eb88096925d87726c0798709bf56d9ed84af3" + +[[package]] +name = "bitflags" +version = "2.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b4388bee8683e3d04af747c73422af53102d2bd24d9eadb6cbc100baef4b43f8" + +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + +[[package]] +name = "candle-kernels" +version = "0.11.0" +dependencies = [ + "cudaforge", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + +[[package]] +name = "crossbeam-deque" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9dd111b7b7f7d55b72c0a6ae361660ee5853c9af73f70c3c2ef6858b950e2e51" +dependencies = [ + "crossbeam-epoch", + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-epoch" +version = "0.9.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5b82ac4a3c2ca9c3460964f020e1402edd5753411d7737aa39c3714ad1b5420e" +dependencies = [ + "crossbeam-utils", +] + +[[package]] +name = "crossbeam-utils" +version = "0.8.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d0a5c400df2834b80a4c3327b3aad3a4c4cd4de0629063962b03235697506a28" + +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "cudaforge" +version = "0.1.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8d475ff95b9e4096878f47fe0ca5dbad0e23849a79fc225305ea06c2f25771cd" +dependencies = [ + "anyhow", + "fs2", + "glob", + "num_cpus", + "rayon", + "serde", + "serde_json", + "sha2", + "thiserror", + "walkdir", + "which", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + +[[package]] +name = "either" +version = "1.16.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" + +[[package]] +name = "env_home" +version = "0.1.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7f84e12ccf0a7ddc17a6c41c93326024c42920d7ee630d04950e6926645c0fe" + +[[package]] +name = "errno" +version = "0.3.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "39cab71617ae0d63f51a36d69f866391735b51691dbda63cf6f96d042b63efeb" +dependencies = [ + "libc", + "windows-sys", +] + +[[package]] +name = "fs2" +version = "0.4.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9564fc758e15025b46aa6643b1b77d047d1a56a1aea6e01002ac0c7026876213" +dependencies = [ + "libc", + "winapi", +] + +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + +[[package]] +name = "glob" +version = "0.3.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" + +[[package]] +name = "hermit-abi" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc0fef456e4baa96da950455cd02c081ca953b141298e41db3fc7e36b1da849c" + +[[package]] +name = "itoa" +version = "1.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8f42a60cbdf9a97f5d2305f08a87dc4e09308d1276d28c869c684d7777685682" + +[[package]] +name = "libc" +version = "0.2.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + +[[package]] +name = "memchr" +version = "2.8.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88904434abc2901f197fe8cc55f0445e7ded921dba5911dad2e2b39b48e663c4" + +[[package]] +name = "num_cpus" +version = "1.17.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "91df4bbde75afed763b708b7eee1e8e7651e02d97f6d5dd763e89367e957b23b" +dependencies = [ + "hermit-abi", + "libc", +] + +[[package]] +name = "proc-macro2" +version = "1.0.106" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8fd00f0bb2e90d81d1044c2b32617f68fcb9fa3bb7640c23e9c748e53fb30934" +dependencies = [ + "unicode-ident", +] + +[[package]] +name = "quote" +version = "1.0.46" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dfbc457d0c7a0759a614551b11a6409e5951f6c7537be1f1b7682b9ae9230368" +dependencies = [ + "proc-macro2", +] + +[[package]] +name = "rayon" +version = "1.12.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fb39b166781f92d482534ef4b4b1b2568f42613b53e5b6c160e24cfbfa30926d" +dependencies = [ + "either", + "rayon-core", +] + +[[package]] +name = "rayon-core" +version = "1.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "22e18b0f0062d30d4230b2e85ff77fdfe4326feb054b9783a3460d8435c8ab91" +dependencies = [ + "crossbeam-deque", + "crossbeam-utils", +] + +[[package]] +name = "rustix" +version = "1.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6fe4565b9518b83ef4f91bb47ce29620ca828bd32cb7e408f0062e9930ba190" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys", +] + +[[package]] +name = "same-file" +version = "1.0.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "93fc1dc3aaa9bfed95e02e6eadabb4baf7e3078b0bd1b4d7b6b0b68378900502" +dependencies = [ + "winapi-util", +] + +[[package]] +name = "serde" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9a8e94ea7f378bd32cbbd37198a4a91436180c5bb472411e48b5ec2e2124ae9e" +dependencies = [ + "serde_core", + "serde_derive", +] + +[[package]] +name = "serde_core" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "41d385c7d4ca58e59fc732af25c3983b67ac852c1a25000afe1175de458b67ad" +dependencies = [ + "serde_derive", +] + +[[package]] +name = "serde_derive" +version = "1.0.228" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d540f220d3187173da220f885ab66608367b6574e925011a9353e4badda91d79" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "serde_json" +version = "1.0.150" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e8014e44b4736ed0538adeecded0fce2a272f22dc9578a7eb6b2d9993c74cfb9" +dependencies = [ + "itoa", + "memchr", + "serde", + "serde_core", + "zmij", +] + +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + +[[package]] +name = "syn" +version = "2.0.118" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1b9ae57f904213ebb649ce6895b8a66c66f0203b9319718f69a5612a065b1422" +dependencies = [ + "proc-macro2", + "quote", + "unicode-ident", +] + +[[package]] +name = "thiserror" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4288b5bcbc7920c07a1149a35cf9590a2aa808e0bc1eafaade0b80947865fbc4" +dependencies = [ + "thiserror-impl", +] + +[[package]] +name = "thiserror-impl" +version = "2.0.18" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ebc4ee7f67670e9b64d05fa4253e753e016c6c95ff35b89b7941d6b856dec1d5" +dependencies = [ + "proc-macro2", + "quote", + "syn", +] + +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + +[[package]] +name = "unicode-ident" +version = "1.0.24" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6e4313cd5fcd3dad5cafa179702e2b244f760991f45397d14d4ebf38247da75" + +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + +[[package]] +name = "walkdir" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "29790946404f91d9c5d06f9874efddea1dc06c5efe94541a7d6863108e3a5e4b" +dependencies = [ + "same-file", + "winapi-util", +] + +[[package]] +name = "which" +version = "7.0.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "24d643ce3fd3e5b54854602a080f34fb10ab75e0b813ee32d00ca2b44fa74762" +dependencies = [ + "either", + "env_home", + "rustix", + "winsafe", +] + +[[package]] +name = "winapi" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419" +dependencies = [ + "winapi-i686-pc-windows-gnu", + "winapi-x86_64-pc-windows-gnu", +] + +[[package]] +name = "winapi-i686-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6" + +[[package]] +name = "winapi-util" +version = "0.1.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2a7b1c03c876122aa43f3020e6c3c3ee5c05081c9a00739faf7503aeba10d22" +dependencies = [ + "windows-sys", +] + +[[package]] +name = "winapi-x86_64-pc-windows-gnu" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f" + +[[package]] +name = "windows-link" +version = "0.2.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0805222e57f7521d6a62e36fa9163bc891acd422f971defe97d64e70d0a4fe5" + +[[package]] +name = "windows-sys" +version = "0.61.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ae137229bcbd6cdf0f7b80a31df61766145077ddf49416a728b02cb3921ff3fc" +dependencies = [ + "windows-link", +] + +[[package]] +name = "winsafe" +version = "0.0.19" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d135d17ab770252ad95e9a872d365cf3090e3be864a34ab46f48555993efc904" + +[[package]] +name = "zmij" +version = "1.0.21" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8848ee67ecc8aedbaf3e4122217aff892639231befc6a1b58d29fff4c2cabaa" diff --git a/torchnative/rust/vendor/candle-kernels/Cargo.toml b/torchnative/rust/vendor/candle-kernels/Cargo.toml new file mode 100644 index 00000000..60f8a78a --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/Cargo.toml @@ -0,0 +1,40 @@ +# THIS FILE IS AUTOMATICALLY GENERATED BY CARGO +# +# When uploading crates to the registry Cargo will automatically +# "normalize" Cargo.toml files for maximal compatibility +# with all versions of Cargo and also rewrite `path` dependencies +# to registry (e.g., crates.io) dependencies. +# +# If you are reading this file be aware that the original Cargo.toml +# will likely look very different (and much more reasonable). +# See Cargo.toml.orig for the original contents. + +[package] +edition = "2021" +name = "candle-kernels" +version = "0.11.0" +build = "build.rs" +autolib = false +autobins = false +autoexamples = false +autotests = false +autobenches = false +description = "CUDA kernels for Candle" +readme = "README.md" +keywords = [ + "blas", + "tensor", + "machine-learning", +] +categories = ["science"] +license = "MIT OR Apache-2.0" +repository = "https://github.com/huggingface/candle" + +[lib] +name = "candle_kernels" +path = "src/lib.rs" + +[dependencies] + +[build-dependencies.cudaforge] +version = "0.1.2" diff --git a/torchnative/rust/vendor/candle-kernels/Cargo.toml.orig b/torchnative/rust/vendor/candle-kernels/Cargo.toml.orig new file mode 100644 index 00000000..9e5efd44 --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/Cargo.toml.orig @@ -0,0 +1,15 @@ +[package] +name = "candle-kernels" +version = "0.11.0" +edition = "2021" + +description = "CUDA kernels for Candle" +repository = "https://github.com/huggingface/candle" +keywords = ["blas", "tensor", "machine-learning"] +categories = ["science"] +license = "MIT OR Apache-2.0" + +[dependencies] + +[build-dependencies] +cudaforge = "0.1.2" diff --git a/torchnative/rust/vendor/candle-kernels/README.md b/torchnative/rust/vendor/candle-kernels/README.md new file mode 100644 index 00000000..1043f31f --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/README.md @@ -0,0 +1,4 @@ +# candle-kernels + +This crate contains CUDA kernels used from candle. Some of these implementations +come from the [dfdx crate](https://github.com/coreylowman/dfdx). diff --git a/torchnative/rust/vendor/candle-kernels/build.rs b/torchnative/rust/vendor/candle-kernels/build.rs new file mode 100644 index 00000000..925345c6 --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/build.rs @@ -0,0 +1,75 @@ +use cudaforge::{KernelBuilder, Result}; +use std::env; +use std::path::PathBuf; + +fn main() -> Result<()> { + println!("cargo::rerun-if-changed=build.rs"); + println!("cargo::rerun-if-changed=src/compatibility.cuh"); + println!("cargo::rerun-if-changed=src/cuda_utils.cuh"); + println!("cargo::rerun-if-changed=src/binary_op_macros.cuh"); + + // Build for PTX + let out_dir = PathBuf::from(env::var("OUT_DIR").unwrap()); + let ptx_path = out_dir.join("ptx.rs"); + let bindings = KernelBuilder::new() + .source_dir("src") // Scan src/ for .cu files + .exclude(&["moe_*.cu", "mmvq_gguf.cu", "mmq_*.cu"]) // Exclude statically compiled kernels from ptx build + .arg("--expt-relaxed-constexpr") + .arg("-std=c++17") + .arg("-O3") + .build_ptx()?; + + bindings.write(&ptx_path)?; + + let mut moe_builder = KernelBuilder::default() + .source_files(vec![ + "src/moe/moe_gguf.cu", + "src/moe/moe_wmma.cu", + "src/moe/moe_wmma_gguf.cu", + "src/mmvq_gguf.cu", + "src/mmq_gguf/mmq_quantize.cu", + "src/mmq_gguf/mmq_instance_q4_0.cu", + "src/mmq_gguf/mmq_instance_q4_1.cu", + "src/mmq_gguf/mmq_instance_q5_0.cu", + "src/mmq_gguf/mmq_instance_q5_1.cu", + "src/mmq_gguf/mmq_instance_q8_0.cu", + "src/mmq_gguf/mmq_instance_q2_k.cu", + "src/mmq_gguf/mmq_instance_q3_k.cu", + "src/mmq_gguf/mmq_instance_q4_k.cu", + "src/mmq_gguf/mmq_instance_q5_k.cu", + "src/mmq_gguf/mmq_instance_q6_k.cu", + ]) + .arg("--expt-relaxed-constexpr") + .arg("-std=c++17") + .arg("-O3"); + + // Disable bf16 WMMA kernels on GPUs older than sm_80 (Ampere). + // bf16 WMMA fragments require compute capability >= 8.0. + let compute_cap = cudaforge::detect_compute_cap() + .map(|arch| arch.base()) + .unwrap_or(80); + if compute_cap < 80 { + moe_builder = moe_builder.arg("-DNO_BF16_KERNEL"); + } + + let mut is_target_msvc = false; + if let Ok(target) = std::env::var("TARGET") { + if target.contains("msvc") { + is_target_msvc = true; + moe_builder = moe_builder.arg("-D_USE_MATH_DEFINES"); + } + } + + if !is_target_msvc { + moe_builder = moe_builder.arg("-Xcompiler").arg("-fPIC"); + } + + moe_builder.build_lib(out_dir.join("libmoe.a"))?; + println!("cargo:rustc-link-search={}", out_dir.display()); + println!("cargo:rustc-link-lib=moe"); + println!("cargo:rustc-link-lib=dylib=cudart"); + if !is_target_msvc { + println!("cargo:rustc-link-lib=stdc++"); + } + Ok(()) +} diff --git a/torchnative/rust/vendor/candle-kernels/src/affine.cu b/torchnative/rust/vendor/candle-kernels/src/affine.cu new file mode 100644 index 00000000..5f5cc158 --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/affine.cu @@ -0,0 +1,51 @@ +#include "cuda_utils.cuh" +#include + +#define AFFINE_OP(TYPENAME, FN_NAME, AFFINE) \ +extern "C" __global__ void FN_NAME( \ + const size_t numel, \ + const size_t num_dims, \ + const size_t *info, \ + const TYPENAME *inp, \ + TYPENAME *out, \ + const TYPENAME mul, \ + const TYPENAME add \ +) { \ + const size_t *dims = info; \ + const size_t *strides = info + num_dims; \ + if (info == nullptr || is_contiguous(num_dims, dims, strides)) { \ + for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { \ + TYPENAME x = inp ? inp[i] : out[i]; \ + out[i] = AFFINE; \ + } \ + } \ + else { \ + for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { \ + unsigned strided_i = get_strided_index(i, num_dims, dims, strides); \ + TYPENAME x = inp ? inp[strided_i] : out[i]; \ + out[i] = AFFINE; \ + } \ + } \ +} \ + +#if __CUDA_ARCH__ >= 800 +AFFINE_OP(__nv_bfloat16, affine_bf16, x * mul + add) +#endif + +#if __CUDA_ARCH__ >= 890 +#define F8E4M3_TO_FLOAT(x) __half2float(__nv_cvt_fp8_to_halfraw(x.__x, __NV_E4M3)) + +AFFINE_OP(__nv_fp8_e4m3, affine_f8_e4m3, __nv_fp8_e4m3(F8E4M3_TO_FLOAT(x) * F8E4M3_TO_FLOAT(mul) + F8E4M3_TO_FLOAT(add))) +#endif + +#if __CUDA_ARCH__ >= 530 +AFFINE_OP(__half, affine_f16, x * mul + add) +#endif + +AFFINE_OP(float, affine_f32, x * mul + add) +AFFINE_OP(double, affine_f64, x * mul + add) +AFFINE_OP(uint8_t, affine_u8, x * mul + add) +AFFINE_OP(uint32_t, affine_u32, x * mul + add) +AFFINE_OP(int16_t, affine_i16, x * mul + add) +AFFINE_OP(int32_t, affine_i32, x * mul + add) +AFFINE_OP(int64_t, affine_i64, x * mul + add) diff --git a/torchnative/rust/vendor/candle-kernels/src/binary.cu b/torchnative/rust/vendor/candle-kernels/src/binary.cu new file mode 100644 index 00000000..971a2c43 --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/binary.cu @@ -0,0 +1,114 @@ +#include "binary_op_macros.cuh" +#include + +#if __CUDA_ARCH__ >= 800 +BINARY_OP(__nv_bfloat16, badd_bf16, x + y) +BINARY_OP(__nv_bfloat16, bdiv_bf16, x / y) +BINARY_OP(__nv_bfloat16, bmul_bf16, x * y) +BINARY_OP(__nv_bfloat16, bsub_bf16, x - y) +BINARY_OP(__nv_bfloat16, bmaximum_bf16, maxg(x, y)) +BINARY_OP(__nv_bfloat16, bminimum_bf16, ming(x, y)) +BINARY_OP_OUT(__nv_bfloat16, uint8_t, eq_bf16, x == y) +BINARY_OP_OUT(__nv_bfloat16, uint8_t, ne_bf16, x != y) +BINARY_OP_OUT(__nv_bfloat16, uint8_t, lt_bf16, x < y) +BINARY_OP_OUT(__nv_bfloat16, uint8_t, le_bf16, x <= y) +BINARY_OP_OUT(__nv_bfloat16, uint8_t, gt_bf16, x > y) +BINARY_OP_OUT(__nv_bfloat16, uint8_t, ge_bf16, x >= y) + +#define F8E4M3_TO_FLOAT(x) __half2float(__nv_cvt_fp8_to_halfraw(x.__x, __NV_E4M3)) + +BINARY_OP(__nv_fp8_e4m3, badd_f8_e4m3, __nv_fp8_e4m3(F8E4M3_TO_FLOAT(x) + F8E4M3_TO_FLOAT(y))) +BINARY_OP(__nv_fp8_e4m3, bdiv_f8_e4m3, __nv_fp8_e4m3(F8E4M3_TO_FLOAT(x) / F8E4M3_TO_FLOAT(y))) +BINARY_OP(__nv_fp8_e4m3, bmul_f8_e4m3, __nv_fp8_e4m3(F8E4M3_TO_FLOAT(x) * F8E4M3_TO_FLOAT(y))) +BINARY_OP(__nv_fp8_e4m3, bsub_f8_e4m3, __nv_fp8_e4m3(F8E4M3_TO_FLOAT(x) - F8E4M3_TO_FLOAT(y))) +BINARY_OP(__nv_fp8_e4m3, bmaximum_f8_e4m3, maxg(x, y)) +BINARY_OP(__nv_fp8_e4m3, bminimum_f8_e4m3, ming(x, y)) +BINARY_OP_OUT(__nv_fp8_e4m3, uint8_t, eq_f8_e4m3, F8E4M3_TO_FLOAT(x) == F8E4M3_TO_FLOAT(y)) +BINARY_OP_OUT(__nv_fp8_e4m3, uint8_t, ne_f8_e4m3, F8E4M3_TO_FLOAT(x) != F8E4M3_TO_FLOAT(y)) +BINARY_OP_OUT(__nv_fp8_e4m3, uint8_t, lt_f8_e4m3, F8E4M3_TO_FLOAT(x) < F8E4M3_TO_FLOAT(y)) +BINARY_OP_OUT(__nv_fp8_e4m3, uint8_t, le_f8_e4m3, F8E4M3_TO_FLOAT(x) <= F8E4M3_TO_FLOAT(y)) +BINARY_OP_OUT(__nv_fp8_e4m3, uint8_t, gt_f8_e4m3, F8E4M3_TO_FLOAT(x) > F8E4M3_TO_FLOAT(y)) +BINARY_OP_OUT(__nv_fp8_e4m3, uint8_t, ge_f8_e4m3, F8E4M3_TO_FLOAT(x) >= F8E4M3_TO_FLOAT(y)) +#endif + +#if __CUDA_ARCH__ >= 530 +BINARY_OP(__half, badd_f16, x + y) +BINARY_OP(__half, bdiv_f16, x / y) +BINARY_OP(__half, bmul_f16, x * y) +BINARY_OP(__half, bsub_f16, x - y) +BINARY_OP(__half, bmaximum_f16, maxg(x, y)) +BINARY_OP(__half, bminimum_f16, ming(x, y)) +BINARY_OP_OUT(__half, uint8_t, eq_f16, x == y) +BINARY_OP_OUT(__half, uint8_t, ne_f16, x != y) +BINARY_OP_OUT(__half, uint8_t, lt_f16, x < y) +BINARY_OP_OUT(__half, uint8_t, le_f16, x <= y) +BINARY_OP_OUT(__half, uint8_t, gt_f16, x > y) +BINARY_OP_OUT(__half, uint8_t, ge_f16, x >= y) +#endif + +BINARY_OP(float, badd_f32, x + y) +BINARY_OP(double, badd_f64, x + y); +BINARY_OP(uint8_t, badd_u8, x + y); +BINARY_OP(uint32_t, badd_u32, x + y); +BINARY_OP(int64_t, badd_i64, x + y); +BINARY_OP(float, bdiv_f32, x / y) +BINARY_OP(double, bdiv_f64, x / y); +BINARY_OP(uint8_t, bdiv_u8, x / y); +BINARY_OP(uint32_t, bdiv_u32, x / y); +BINARY_OP(int64_t, bdiv_i64, x / y); +BINARY_OP(float, bmul_f32, x * y) +BINARY_OP(double, bmul_f64, x * y); +BINARY_OP(uint8_t, bmul_u8, x * y); +BINARY_OP(uint32_t, bmul_u32, x * y); +BINARY_OP(int64_t, bmul_i64, x * y); +BINARY_OP(float, bsub_f32, x - y) +BINARY_OP(double, bsub_f64, x - y); +BINARY_OP(uint8_t, bsub_u8, x - y); +BINARY_OP(uint32_t, bsub_u32, x - y); +BINARY_OP(int64_t, bsub_i64, x - y); +BINARY_OP(float, bminimum_f32, ming(x, y)); +BINARY_OP(double, bminimum_f64, ming(x, y)); +BINARY_OP(uint8_t, bminimum_u8, ming(x, y)); +BINARY_OP(uint32_t, bminimum_u32, ming(x, y)); +BINARY_OP(int64_t, bminimum_i64, ming(x, y)); +BINARY_OP(float, bmaximum_f32, maxg(x, y)); +BINARY_OP(double, bmaximum_f64, maxg(x, y)); +BINARY_OP(uint8_t, bmaximum_u8, maxg(x, y)); +BINARY_OP(uint32_t, bmaximum_u32, maxg(x, y)); +BINARY_OP(int64_t, bmaximum_i64, maxg(x, y)); + +BINARY_OP_OUT(float, uint8_t, eq_f32, x == y) +BINARY_OP_OUT(double, uint8_t, eq_f64, x == y) +BINARY_OP_OUT(uint8_t, uint8_t, eq_u8, x == y) +BINARY_OP_OUT(uint32_t, uint8_t, eq_u32, x == y) +BINARY_OP_OUT(int64_t, uint8_t, eq_i64, x == y) + +BINARY_OP_OUT(float, uint8_t, ne_f32, x != y) +BINARY_OP_OUT(double, uint8_t, ne_f64, x != y) +BINARY_OP_OUT(uint8_t, uint8_t, ne_u8, x != y) +BINARY_OP_OUT(uint32_t, uint8_t, ne_u32, x != y) +BINARY_OP_OUT(int64_t, uint8_t, ne_i64, x != y) + +BINARY_OP_OUT(float, uint8_t, lt_f32, x < y) +BINARY_OP_OUT(double, uint8_t, lt_f64, x < y) +BINARY_OP_OUT(uint8_t, uint8_t, lt_u8, x < y) +BINARY_OP_OUT(uint32_t, uint8_t, lt_u32, x < y) +BINARY_OP_OUT(int64_t, uint8_t, lt_i64, x < y) + +BINARY_OP_OUT(float, uint8_t, le_f32, x <= y) +BINARY_OP_OUT(double, uint8_t, le_f64, x <= y) +BINARY_OP_OUT(uint8_t, uint8_t, le_u8, x <= y) +BINARY_OP_OUT(uint32_t, uint8_t, le_u32, x <= y) +BINARY_OP_OUT(int64_t, uint8_t, le_i64, x <= y) + +BINARY_OP_OUT(float, uint8_t, gt_f32, x > y) +BINARY_OP_OUT(double, uint8_t, gt_f64, x > y) +BINARY_OP_OUT(uint8_t, uint8_t, gt_u8, x > y) +BINARY_OP_OUT(uint32_t, uint8_t, gt_u32, x > y) +BINARY_OP_OUT(int64_t, uint8_t, gt_i64, x > y) + +BINARY_OP_OUT(float, uint8_t, ge_f32, x >= y) +BINARY_OP_OUT(double, uint8_t, ge_f64, x >= y) +BINARY_OP_OUT(uint8_t, uint8_t, ge_u8, x >= y) +BINARY_OP_OUT(uint32_t, uint8_t, ge_u32, x >= y) +BINARY_OP_OUT(int64_t, uint8_t, ge_i64, x >= y) diff --git a/torchnative/rust/vendor/candle-kernels/src/binary_op_macros.cuh b/torchnative/rust/vendor/candle-kernels/src/binary_op_macros.cuh new file mode 100644 index 00000000..9cb00874 --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/binary_op_macros.cuh @@ -0,0 +1,69 @@ +#include "cuda_utils.cuh" + +#define BINARY_OP_OUT(TYPENAME, OUT_TYPENAME, FN_NAME, FUNC) \ +extern "C" __global__ void FN_NAME( \ + const size_t numel, \ + const size_t num_dims, \ + const size_t *dims_and_strides, \ + const TYPENAME *lhs, \ + const TYPENAME *rhs, \ + OUT_TYPENAME *out \ +) { \ + const size_t *dims = dims_and_strides; \ + const size_t *lhs_strides = dims_and_strides + 1 * num_dims; \ + const size_t *rhs_strides = dims_and_strides + 2 * num_dims; \ + bool lhs_cont = dims_and_strides == nullptr || is_contiguous(num_dims, dims, lhs_strides); \ + bool rhs_cont = dims_and_strides == nullptr || is_contiguous(num_dims, dims, rhs_strides); \ + if (lhs_cont && rhs_cont) { \ + for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { \ + TYPENAME x = lhs[i]; \ + TYPENAME y = rhs[i]; \ + out[i] = FUNC; \ + } \ + } else if (lhs_cont) { \ + for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { \ + unsigned int tmp_i = i; \ + unsigned int rhs_i = 0; \ + for (int d = num_dims - 1; d >= 0; d--) { \ + unsigned int i_dim = tmp_i % dims[d]; \ + rhs_i += i_dim * rhs_strides[d]; \ + tmp_i /= dims[d]; \ + } \ + TYPENAME x = lhs[i]; \ + TYPENAME y = rhs[rhs_i]; \ + out[i] = FUNC; \ + } \ + } else if (rhs_cont) { \ + for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { \ + unsigned int tmp_i = i; \ + unsigned int lhs_i = 0; \ + for (int d = num_dims - 1; d >= 0; d--) { \ + unsigned int i_dim = tmp_i % dims[d]; \ + lhs_i += i_dim * lhs_strides[d]; \ + tmp_i /= dims[d]; \ + } \ + TYPENAME x = lhs[lhs_i]; \ + TYPENAME y = rhs[i]; \ + out[i] = FUNC; \ + } \ + } else { \ + for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { \ + unsigned int tmp_i = i; \ + unsigned int lhs_i = 0; \ + unsigned int rhs_i = 0; \ + for (int d = num_dims - 1; d >= 0; d--) { \ + unsigned int i_dim = tmp_i % dims[d]; \ + lhs_i += i_dim * lhs_strides[d]; \ + rhs_i += i_dim * rhs_strides[d]; \ + tmp_i /= dims[d]; \ + } \ + TYPENAME x = lhs[lhs_i]; \ + TYPENAME y = rhs[rhs_i]; \ + out[i] = FUNC; \ + } \ + } \ +} \ + + +#define BINARY_OP(TYPENAME, FN_NAME, FUNC) \ + BINARY_OP_OUT(TYPENAME, TYPENAME, FN_NAME, FUNC) diff --git a/torchnative/rust/vendor/candle-kernels/src/cast.cu b/torchnative/rust/vendor/candle-kernels/src/cast.cu new file mode 100644 index 00000000..1b38f58e --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/cast.cu @@ -0,0 +1,227 @@ +#include "cuda_utils.cuh" +#include + +template +__device__ void cast_( + const size_t numel, + const size_t num_dims, + const size_t *info, + const S *inp, + T *out +) { + const size_t *dims = info; + const size_t *strides = info + num_dims; + if (info == nullptr || is_contiguous(num_dims, dims, strides)) { + for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { + out[i] = inp[i]; + } + } + else { + for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { + unsigned strided_i = get_strided_index(i, num_dims, dims, strides); + out[i] = inp[strided_i]; + } + } +} + +#define F8E4M3_TO_FLOAT(x) __half2float(__nv_cvt_fp8_to_halfraw(x.__x, __NV_E4M3)) + +template +__device__ void cast_fp8_( + const size_t numel, + const size_t num_dims, + const size_t *info, + const __nv_fp8_e4m3 *inp, + T *out +) { + const size_t *dims = info; + const size_t *strides = info + num_dims; + if (info == nullptr || is_contiguous(num_dims, dims, strides)) { + for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { + out[i] = F8E4M3_TO_FLOAT(inp[i]); + } + } + else { + for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { + unsigned strided_i = get_strided_index(i, num_dims, dims, strides); + out[i] = F8E4M3_TO_FLOAT(inp[strided_i]); + } + } +} +template +__device__ void cast_fp8_into_( + const size_t numel, + const size_t num_dims, + const size_t *info, + const S *inp, + __nv_fp8_e4m3 *out +) { + const size_t *dims = info; + const size_t *strides = info + num_dims; + if (info == nullptr || is_contiguous(num_dims, dims, strides)) { + for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { + out[i] = __nv_fp8_e4m3((float)inp[i]); + } + } + else { + for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { + unsigned strided_i = get_strided_index(i, num_dims, dims, strides); + out[i] = __nv_fp8_e4m3((float)inp[strided_i]); + } + } +} + +template +__device__ void cast_through( + const size_t numel, + const size_t num_dims, + const size_t *info, + const S *inp, + T *out +) { + const size_t *dims = info; + const size_t *strides = info + num_dims; + if (info == nullptr || is_contiguous(num_dims, dims, strides)) { + for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { + out[i] = static_cast(static_cast(inp[i])); + } + } + else { + for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { + unsigned strided_i = get_strided_index(i, num_dims, dims, strides); + out[i] = static_cast(static_cast(inp[strided_i])); + } + } +} + + +#define CAST_OP(SRC_TYPENAME, DST_TYPENAME, FN_NAME) \ +extern "C" __global__ void FN_NAME( \ + const size_t numel, \ + const size_t num_dims, \ + const size_t *info, \ + const SRC_TYPENAME *inp, \ + DST_TYPENAME *out \ +) { \ + cast_(numel, num_dims, info, inp, out); \ +} \ + + +#define CAST_OP_FP8(SRC_TYPENAME, DST_TYPENAME, FN_NAME) \ +extern "C" __global__ void FN_NAME( \ + const size_t numel, \ + const size_t num_dims, \ + const size_t *info, \ + const SRC_TYPENAME *inp, \ + DST_TYPENAME *out \ +) { \ + cast_fp8_(numel, num_dims, info, inp, out); \ +} \ + + +#define CAST_OP_FP8_INTO(SRC_TYPENAME, DST_TYPENAME, FN_NAME) \ +extern "C" __global__ void FN_NAME( \ + const size_t numel, \ + const size_t num_dims, \ + const size_t *info, \ + const SRC_TYPENAME *inp, \ + DST_TYPENAME *out \ +) { \ + cast_fp8_into_(numel, num_dims, info, inp, out); \ +} \ + +#define CAST_THROUGH_OP(SRC_TYPENAME, DST_TYPENAME, INT_TYPENAME, FN_NAME) \ +extern "C" __global__ void FN_NAME( \ + const size_t numel, \ + const size_t num_dims, \ + const size_t *info, \ + const SRC_TYPENAME *inp, \ + DST_TYPENAME *out \ +) { \ + cast_through(numel, num_dims, info, inp, out); \ +} \ + +#if __CUDA_ARCH__ >= 800 +CAST_OP(__nv_bfloat16, __nv_bfloat16, cast_bf16_bf16) +CAST_OP(__nv_fp8_e4m3, __nv_fp8_e4m3, cast_f8_e4m3_f8_e4m3) + +CAST_OP(__nv_bfloat16, uint32_t, cast_bf16_u32) +CAST_OP(__nv_bfloat16, float, cast_bf16_f32) +CAST_OP(__nv_bfloat16, double, cast_bf16_f64) +CAST_OP(uint8_t, __nv_bfloat16, cast_u8_bf16) +CAST_OP(uint32_t, __nv_bfloat16, cast_u32_bf16) +CAST_OP(float, __nv_bfloat16, cast_f32_bf16) +CAST_OP(double, __nv_bfloat16, cast_f64_bf16) +CAST_THROUGH_OP(__nv_bfloat16, uint8_t, float, cast_bf16_u8) +CAST_THROUGH_OP(__nv_bfloat16, __half, float, cast_bf16_f16) +CAST_THROUGH_OP(__half, __nv_bfloat16, float, cast_f16_bf16) + +CAST_OP_FP8(__nv_fp8_e4m3, float, cast_f8_e4m3_f32) +CAST_OP_FP8_INTO(float, __nv_fp8_e4m3, cast_f32_f8_e4m3) +CAST_OP_FP8(__nv_fp8_e4m3, uint8_t, cast_f8_e4m3_u8) +CAST_OP_FP8(__nv_fp8_e4m3, __half, cast_f8_e4m3_f16) +CAST_OP_FP8(__nv_fp8_e4m3, double, cast_f8_e4m3_f64) +CAST_OP_FP8_INTO(__half, __nv_fp8_e4m3, cast_f16_f8_e4m3) +CAST_OP_FP8_INTO(double, __nv_fp8_e4m3, cast_f64_f8_e4m3) +CAST_OP_FP8_INTO(uint8_t, __nv_fp8_e4m3, cast_u8_f8_e4m3) +CAST_OP_FP8_INTO(int32_t, __nv_fp8_e4m3, cast_i32_f8_e4m3) +CAST_OP_FP8(__nv_fp8_e4m3, int32_t, cast_f8_e4m3_i32) +CAST_OP_FP8(__nv_fp8_e4m3, __nv_bfloat16, cast_f8_e4m3_bf16) +CAST_OP_FP8_INTO(__nv_bfloat16, __nv_fp8_e4m3, cast_bf16_f8_e4m3) +#else +#include +#if CUDA_VERSION >= 11000 +CAST_OP(__nv_bfloat16, float, cast_bf16_f32) +CAST_OP(float, __nv_bfloat16, cast_f32_bf16) +CAST_THROUGH_OP(__nv_bfloat16, uint8_t, float, cast_bf16_u8) +CAST_THROUGH_OP(__nv_bfloat16, __half, float, cast_bf16_f16) +CAST_THROUGH_OP(__nv_bfloat16, double, float, cast_bf16_f64) +CAST_THROUGH_OP(__half, __nv_bfloat16, float, cast_f16_bf16) +CAST_THROUGH_OP(double, __nv_bfloat16, float, cast_f64_bf16) +CAST_THROUGH_OP(uint8_t, __nv_bfloat16, float, cast_u8_bf16) +CAST_THROUGH_OP(__nv_bfloat16, __nv_fp8_e4m3, float, cast_bf16_f8_e4m3) +#endif +#endif + +#if __CUDA_ARCH__ >= 530 +CAST_OP(__half, __half, cast_f16_f16) + +CAST_THROUGH_OP(__half, uint8_t, float, cast_f16_u8) +CAST_OP(__half, uint32_t, cast_f16_u32) +CAST_OP(__half, float, cast_f16_f32) +CAST_OP(__half, double, cast_f16_f64) +CAST_OP(uint8_t, __half, cast_u8_f16 ) +CAST_OP(uint32_t, __half, cast_u32_f16) +CAST_OP(float, __half, cast_f32_f16) +CAST_OP(double, __half, cast_f64_f16) +#endif + +CAST_OP(uint32_t, uint32_t, cast_u32_u32) +CAST_OP(uint32_t, uint8_t, cast_u32_u8 ) +CAST_OP(uint32_t, int64_t, cast_u32_i64 ) +CAST_OP(uint32_t, float, cast_u32_f32) +CAST_OP(uint32_t, double, cast_u32_f64) + +CAST_OP(uint8_t, uint32_t, cast_u8_u32) +CAST_OP(uint8_t, uint8_t, cast_u8_u8 ) +CAST_OP(uint8_t, int64_t, cast_u8_i64 ) +CAST_OP(uint8_t, float, cast_u8_f32) +CAST_OP(uint8_t, double, cast_u8_f64) + +CAST_OP(int64_t, uint32_t, cast_i64_u32) +CAST_OP(int64_t, uint8_t, cast_i64_u8 ) +CAST_OP(int64_t, int64_t, cast_i64_i64 ) +CAST_OP(int64_t, float, cast_i64_f32) +CAST_OP(int64_t, double, cast_i64_f64) + +CAST_OP(float, uint8_t, cast_f32_u8 ) +CAST_OP(float, uint32_t, cast_f32_u32) +CAST_OP(float, int64_t, cast_f32_i64 ) +CAST_OP(float, float, cast_f32_f32) +CAST_OP(float, double, cast_f32_f64) + +CAST_OP(double, uint8_t, cast_f64_u8 ) +CAST_OP(double, uint32_t, cast_f64_u32) +CAST_OP(double, int64_t, cast_f64_i64 ) +CAST_OP(double, float, cast_f64_f32) +CAST_OP(double, double, cast_f64_f64) diff --git a/torchnative/rust/vendor/candle-kernels/src/compatibility.cuh b/torchnative/rust/vendor/candle-kernels/src/compatibility.cuh new file mode 100644 index 00000000..71a30dcd --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/compatibility.cuh @@ -0,0 +1,167 @@ +#include "cuda_fp16.h" +#include "cuda_bf16.h" +#include "cuda_fp8.h" + +// Table showing which features are supported on which compute capability +// https://docs.nvidia.com/cuda/cuda-c-programming-guide/#features-and-technical-specifications + +// FIXME: the minimum compute capabilities are just guesses since the table is not specific enough + +// torchnative (issue #86): CUDA 12.2 made cuda_fp16.hpp define __hmax_nan and +// __hmin_nan for every architecture, so defining them again here is a +// redefinition error on any toolkit from 12.2 on (measured on 12.6 with +// CUDA_COMPUTE_CAP=75). Define them only for an older toolkit on an older arch. +#if __CUDA_ARCH__ < 800 && (__CUDACC_VER_MAJOR__ < 12 || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ < 2)) +__device__ __forceinline__ __half __hmax_nan(__half a, __half b) { + return __hisnan(a) ? a : (__hisnan(b) ? b : __hmax(a, b)); +} +__device__ __forceinline__ __half __hmin_nan(__half a, __half b) { + return __hisnan(a) ? a : (__hisnan(b) ? b : __hmin(a, b)); +} +#endif + +#if __CUDA_ARCH__ < 600 +// Copied from https://docs.nvidia.com/cuda/cuda-c-programming-guide/#atomic-functions +__device__ double atomicAdd(double* address, double val) { + unsigned long long int* address_as_ull = (unsigned long long int*)address; + unsigned long long int old = *address_as_ull, assumed; + + do { + assumed = old; + old = atomicCAS(address_as_ull, assumed, + __double_as_longlong(val + + __longlong_as_double(assumed))); + + // Note: uses integer comparison to avoid hang in case of NaN (since NaN != NaN) + } while (assumed != old); + + return __longlong_as_double(old); +} +#endif + +#if __CUDA_ARCH__ < 700 +// https://docs.nvidia.com/cuda/cuda-c-programming-guide/index.html#atomicadd +// The 16-bit __half floating-point version of atomicAdd() is only supported by devices of compute capability 7.x and higher. +// Solution adapted from https://github.com/torch/cutorch/blob/master/lib/THC/THCAtomics.cuh#L96-L119 +//__device__ __half atomicAdd(__half *address, __half val) { + // unsigned int *address_as_ui = (unsigned int *) ((char *)address - ((size_t)address & 2)); + // unsigned int old = *address_as_ui; + // unsigned int assumed; + // bool unaligned = (size_t) address & 2; + // do { + // assumed = old; + // unsigned int hsum; + // hsum = unaligned ? (old >> 16) : (old & 0xffff); + // hsum = __half_as_ushort(__ushort_as_half(hsum) + val); + // old = atomicCAS(address_as_ui, assumed, + // unaligned ? (old & 0xffff) | (hsum << 16) : (old & 0xffff0000) | hsum + // ); + + // } while (assumed != old); + // return __ushort_as_half(unaligned ? (old >> 16) : (old & 0xffff)); +//} +#endif + + +__device__ __forceinline__ __half atomicMaxf(__half* address, __half val) { +#if __CUDA_ARCH__ < 700 + // On older GPUs we do not have access to atomicCAS for shorts, so we have to do some trickery. + // Solution adapted from https://github.com/torch/cutorch/blob/master/lib/THC/THCAtomics.cuh#L96-L119 + unsigned int *address_as_ui = (unsigned int *) ((char *)address - ((size_t)address & 2)); + unsigned int old = *address_as_ui; + unsigned int assumed; + bool unaligned = (size_t) address & 2; + do { + assumed = old; + unsigned int hmax; + hmax = unaligned ? (old >> 16) : (old & 0xffff); + hmax = __half_as_ushort(__hmax_nan(val, __ushort_as_half(hmax))); + old = atomicCAS(address_as_ui, assumed, + unaligned ? (old & 0xffff) | (hmax << 16) : (old & 0xffff0000) | hmax + ); + + } while (assumed != old); + return __ushort_as_half(unaligned ? (old >> 16) : (old & 0xffff)); +#else + // Based on https://docs.nvidia.com/cuda/cuda-c-programming-guide/#atomic-functions + unsigned short int* casted_address = (unsigned short int*)address; + unsigned short int old = *casted_address; + unsigned short int assumed; + do { + assumed = old; + old = atomicCAS(casted_address, assumed, __half_as_ushort(__hmax_nan(val, __ushort_as_half(assumed)))); + // Note: uses integer comparison to avoid hang in case of NaN (since NaN != NaN) + } while (assumed != old); + return __ushort_as_half(old); +#endif +} + +// atomicMax is not implemented for floats, +// solution copied https://stackoverflow.com/questions/17399119/how-do-i-use-atomicmax-on-floating-point-values-in-cuda +__device__ __forceinline__ float atomicMaxf(float * addr, float value) { + if (signbit(value)) { + return __uint_as_float(atomicMin((unsigned int *)addr, __float_as_uint(value))); + } else { + return __int_as_float(atomicMax((int *)addr, __float_as_int(value))); + } +} + +__device__ __forceinline__ double atomicMaxf(double * addr, double value) { + if (signbit(value)) { + return __longlong_as_double(atomicMin((unsigned long long int *)addr, __double_as_longlong(value))); + } else { + return __longlong_as_double(atomicMax((long long int *)addr, __double_as_longlong(value))); + } +} + + +__device__ __forceinline__ __half atomicMinf(__half* address, __half val) { +#if __CUDA_ARCH__ < 700 + // On older GPUs we do not have access to atomicCAS for shorts, so we have to do some trickery. + // Solution adapted from https://github.com/torch/cutorch/blob/master/lib/THC/THCAtomics.cuh#L96-L119 + unsigned int *address_as_ui = (unsigned int *) ((char *)address - ((size_t)address & 2)); + unsigned int old = *address_as_ui; + unsigned int assumed; + bool unaligned = (size_t) address & 2; + do { + assumed = old; + unsigned int hmin; + hmin = unaligned ? (old >> 16) : (old & 0xffff); + hmin = __half_as_ushort(__hmin_nan(val, __ushort_as_half(hmin))); + old = atomicCAS(address_as_ui, assumed, + unaligned ? (old & 0xffff) | (hmin << 16) : (old & 0xffff0000) | hmin + ); + + } while (assumed != old); + return __ushort_as_half(unaligned ? (old >> 16) : (old & 0xffff)); +#else + // Based on https://docs.nvidia.com/cuda/cuda-c-programming-guide/#atomic-functions + unsigned short int* casted_address = (unsigned short int*)address; + unsigned short int old = *casted_address; + unsigned short int assumed; + do { + assumed = old; + old = atomicCAS(casted_address, assumed, __half_as_ushort(__hmin_nan(val, __ushort_as_half(assumed)))); + // Note: uses integer comparison to avoid hang in case of NaN (since NaN != NaN) + } while (assumed != old); + return __ushort_as_half(old); +#endif +} + +// atomicMin is not implemented for floats, +// solution copied https://stackoverflow.com/questions/17399119/how-do-i-use-atomicmax-on-floating-point-values-in-cuda +__device__ __forceinline__ float atomicMinf(float * addr, float value) { + if (signbit(value)) { + return __uint_as_float(atomicMax((unsigned int *)addr, __float_as_uint(value))); + } else { + return __int_as_float(atomicMin((int *)addr, __float_as_int(value))); + } +} + +__device__ __forceinline__ double atomicMinf(double * addr, double value) { + if (signbit(value)) { + return __longlong_as_double(atomicMax((unsigned long long int *)addr, __double_as_longlong(value))); + } else { + return __longlong_as_double(atomicMin((long long int *)addr, __double_as_longlong(value))); + } +} diff --git a/torchnative/rust/vendor/candle-kernels/src/conv.cu b/torchnative/rust/vendor/candle-kernels/src/conv.cu new file mode 100644 index 00000000..a901c35e --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/conv.cu @@ -0,0 +1,896 @@ +#include "cuda_utils.cuh" +#include + +// Naive implementation of conv1d. +template +__device__ void conv1d( + const size_t src_numel, + const size_t l_out, + const size_t stride, + const size_t padding, + const size_t dilation, + const size_t *info, + const T *src, + const T *kernel, + T *dst +) { + // src: (b_size, c_in, l_in) + // k: (c_out, c_in, k_size) + const size_t *src_dims = info; + const size_t *src_s = info + 3; + const size_t *k_dims = info + 6; + const size_t *k_s = info + 9; + const size_t dst_i = blockIdx.x * blockDim.x + threadIdx.x; + const size_t k_size = k_dims[2]; + const size_t c_out = k_dims[0]; + const size_t c_in = src_dims[1]; + const size_t l_in = src_dims[2]; + if (dst_i >= src_dims[0] * c_out * l_out) { + return; + } + + // TODO + const size_t b_idx = dst_i / (l_out * c_out); + const size_t dst_c_idx = (dst_i / l_out) % c_out; + const size_t dst_l = dst_i % l_out; + + const size_t src_idx0 = b_idx * src_s[0]; + A d = 0; + for (size_t offset = 0; offset < k_size; ++offset) { + size_t src_l = (stride * dst_l + offset) * dilation; + if (src_l < padding || src_l >= padding + l_in) { + continue; + } + src_l -= padding; + for (size_t src_c_idx = 0; src_c_idx < c_in; ++src_c_idx) { + const size_t src_idx = src_idx0 + src_c_idx * src_s[1] + src_l * src_s[2]; + const size_t k_idx = dst_c_idx * k_s[0] + src_c_idx * k_s[1] + offset * k_s[2]; + d += static_cast(src[src_idx]) * static_cast(kernel[k_idx]); + } + } + dst[dst_i] = static_cast(d); +} + +template +__device__ void im2col1d( + const size_t numel, + const size_t l_out, + const size_t l_k, + const size_t stride, + const size_t padding, + const size_t dilation, + const size_t *info, + const T *src, + T *dst +) { + const size_t thread_i = blockIdx.x * blockDim.x + threadIdx.x; + // dst: (b_size, l_out, c_in, l_k) + // src: (b_size, c_in, l_in) + if (thread_i >= numel) { + return; + } + const size_t *src_dims = info; + const size_t *src_s = info + 3; + const size_t c_in = src_dims[1]; + const size_t l_in = src_dims[2]; + + const size_t dst_s1 = c_in; + const size_t dst_s0 = l_out * dst_s1; + + size_t tmp_dst_i = thread_i; + const size_t b_idx = tmp_dst_i / dst_s0; + tmp_dst_i -= b_idx * dst_s0; + const size_t l_idx = tmp_dst_i / dst_s1; + tmp_dst_i -= l_idx * dst_s1; + const size_t c_idx = tmp_dst_i; + for (size_t l_k_idx = 0; l_k_idx < l_k; ++l_k_idx) { + size_t src_l_idx = l_idx * stride + l_k_idx * dilation; + size_t dst_i = thread_i * l_k + l_k_idx; + if (src_l_idx < padding || src_l_idx >= l_in + padding) { + dst[dst_i] = static_cast(0); + } + else { + src_l_idx -= padding; + const size_t src_i = b_idx * src_s[0] + c_idx * src_s[1] + src_l_idx * src_s[2]; + dst[dst_i] = src[src_i]; + } + } +} + +template +__device__ void col2im1d( + const size_t dst_el, + const size_t l_out, + const size_t l_in, + const size_t c_out, + const size_t k_size, + const size_t stride, + const T *src, + T *dst +) { + const size_t dst_i = blockIdx.x * blockDim.x + threadIdx.x; + // src: (b_size, l_in, c_out, l_k) + // dst: (b_size, c_out, l_out) + if (dst_i >= dst_el) { + return; + } + + const size_t dst_s0 = c_out * l_out; + const size_t dst_s1 = l_out; + const size_t src_s0 = c_out * k_size * l_in; + const size_t src_s1 = c_out * k_size; + const size_t src_s2 = k_size; + + size_t tmp_dst_i = dst_i; + const size_t b_idx = tmp_dst_i / dst_s0; + tmp_dst_i -= b_idx * dst_s0; + const size_t c_idx = tmp_dst_i / dst_s1; + tmp_dst_i -= c_idx * dst_s1; + const int l_out_idx = tmp_dst_i; + + dst[dst_i] = static_cast(0); + + int l_in_idx = l_out_idx / stride; + int k0 = l_out_idx - l_in_idx * stride; + // l_out_idx = l_in_idx * stride + k0 + for (; k0 < k_size && l_in_idx >= 0; k0 += stride, --l_in_idx) { + if (l_in_idx < l_in) { + const size_t src_i = b_idx * src_s0 + l_in_idx * src_s1 + c_idx * src_s2 + k0; + dst[dst_i] += src[src_i]; + } + } +} + +template +__device__ void im2col( + const size_t dst_numel, + const size_t h_out, + const size_t w_out, + const size_t h_k, + const size_t w_k, + const size_t stride, + const size_t padding, + const size_t dilation, + const size_t *info, + const T *src, + T *dst +) { + const size_t dst_i = blockIdx.x * blockDim.x + threadIdx.x; + // dst: (b_size, h_out, w_out, c_in, h_k, w_k) + // src: (b_size, c_in, h_in, w_in) + if (dst_i >= dst_numel) { + return; + } + const size_t *src_dims = info; + const size_t *src_s = info + 4; + const size_t c_in = src_dims[1]; + const size_t h_in = src_dims[2]; + const size_t w_in = src_dims[3]; + + const size_t dst_s4 = w_k; + const size_t dst_s3 = h_k * dst_s4; + const size_t dst_s2 = c_in * dst_s3; + const size_t dst_s1 = w_out * dst_s2; + const size_t dst_s0 = h_out * dst_s1; + + size_t tmp_dst_i = dst_i; + const size_t b_idx = tmp_dst_i / dst_s0; + tmp_dst_i -= b_idx * dst_s0; + const size_t h_idx = tmp_dst_i / dst_s1; + tmp_dst_i -= h_idx * dst_s1; + const size_t w_idx = tmp_dst_i / dst_s2; + tmp_dst_i -= w_idx * dst_s2; + const size_t c_idx = tmp_dst_i / dst_s3; + tmp_dst_i -= c_idx * dst_s3; + const size_t h_k_idx = tmp_dst_i / dst_s4; + tmp_dst_i -= h_k_idx * dst_s4; + const size_t w_k_idx = tmp_dst_i; + size_t src_h_idx = h_idx * stride + h_k_idx * dilation; + size_t src_w_idx = w_idx * stride + w_k_idx * dilation; + if (src_h_idx < padding || src_h_idx >= h_in + padding) { + dst[dst_i] = static_cast(0); + } + else if (src_w_idx < padding || src_w_idx >= w_in + padding) { + dst[dst_i] = static_cast(0); + } + else { + src_h_idx -= padding; + src_w_idx -= padding; + const size_t src_i = + b_idx * src_s[0] + + c_idx * src_s[1] + + src_h_idx * src_s[2] + + src_w_idx * src_s[3]; + dst[dst_i] = src[src_i]; + } +} + +// Naive implementation of conv2d. +template +__device__ void conv2d( + const size_t src_numel, + const size_t w_out, + const size_t h_out, + const size_t stride, + const size_t padding, + const size_t dilation, + const size_t *info, + const T *src, + const T *kernel, + T *dst +) { + const size_t dst_i = blockIdx.x * blockDim.x + threadIdx.x; + // src: (b_size, c_in, h_in, w_in) + // k: (c_out, c_in, h_k, w_k) + const size_t *src_dims = info; + const size_t *src_s = info + 4; + const size_t *k_dims = info + 8; + const size_t *k_s = info + 12; + const size_t h_k = k_dims[2]; + const size_t w_k = k_dims[3]; + const size_t c_out = k_dims[0]; + const size_t c_in = src_dims[1]; + const size_t h_in = src_dims[2]; + const size_t w_in = src_dims[3]; + if (dst_i >= src_dims[0] * c_out * w_out * h_out) { + return; + } + + // TODO + const size_t b_idx = dst_i / (w_out * h_out * c_out); + const size_t dst_c_idx = (dst_i / (w_out * h_out)) % c_out; + // NCHW layout. + const size_t dst_h = (dst_i / w_out) % h_out; + const size_t dst_w = dst_i % w_out; + + const size_t src_idx0 = b_idx * src_s[0]; + A d = 0; + for (size_t w_offset = 0; w_offset < w_k; ++w_offset) { + size_t src_w = stride * dst_w + w_offset * dilation; + if (src_w < padding || src_w >= w_in + padding) { + continue; + } + src_w -= padding; + for (size_t h_offset = 0; h_offset < h_k; ++h_offset) { + size_t src_h = stride * dst_h + h_offset * dilation; + if (src_h < padding || src_h >= h_in + padding) { + continue; + } + src_h -= padding; + for (size_t src_c_idx = 0; src_c_idx < c_in; ++src_c_idx) { + const size_t src_idx = src_idx0 + src_c_idx * src_s[1] + src_h * src_s[2] + src_w * src_s[3]; + const size_t k_idx = dst_c_idx * k_s[0] + src_c_idx * k_s[1] + h_offset * k_s[2] + w_offset * k_s[3]; + d += static_cast(src[src_idx]) * static_cast(kernel[k_idx]); + } + } + } + dst[dst_i] = static_cast(d); +} + +// Naive implementation of conv_transpose1d. +template +__device__ void conv_transpose1d( + const size_t src_numel, + const size_t l_out, + const size_t stride, + const size_t padding, + const size_t out_padding, + const size_t dilation, + const size_t *info, + const T *src, + const T *kernel, + T *dst +) { + const size_t dst_i = blockIdx.x * blockDim.x + threadIdx.x; + // src: (b_size, c_in, l_in) + // k: (c_in, c_out, l_k) + const size_t *src_dims = info; + const size_t *src_s = info + 3; + const size_t *k_dims = info + 6; + const size_t *k_s = info + 9; + const size_t l_k = k_dims[2]; + const size_t c_out = k_dims[1]; + const size_t c_in = src_dims[1]; + const size_t l_in = src_dims[2]; + if (dst_i >= src_dims[0] * c_out * l_out) { + return; + } + + // TODO + const size_t b_idx = dst_i / (l_out * c_out); + const size_t dst_c_idx = (dst_i / l_out) % c_out; + // NCL layout. + const size_t out_x = dst_i % l_out; + + const size_t src_idx0 = b_idx * src_s[0]; + A d = 0; + for (int k_x = 0; k_x < (int)l_k; ++k_x) { + // let out_x = inp_x * p.stride + k_x * p.dilation - p.padding; + int inp_x_stride = (int)(out_x + padding) - k_x * dilation; + if (inp_x_stride < 0 || inp_x_stride % stride) { + continue; + } + int inp_x = inp_x_stride / stride; + if (inp_x >= l_in) continue; + for (size_t src_c_idx = 0; src_c_idx < c_in; ++src_c_idx) { + const size_t src_idx = src_idx0 + src_c_idx * src_s[1] + inp_x * src_s[2]; + const size_t k_idx = src_c_idx * k_s[0] + dst_c_idx * k_s[1] + k_x * k_s[2]; + d += static_cast(src[src_idx]) * static_cast(kernel[k_idx]); + } + } + dst[dst_i] = static_cast(d); +} + +// Naive implementation of conv_transpose2d. +template +__device__ void conv_transpose2d( + const size_t src_numel, + const size_t w_out, + const size_t h_out, + const size_t stride, + const size_t padding, + const size_t out_padding, + const size_t dilation, + const size_t *info, + const T *src, + const T *kernel, + T *dst +) { + const size_t dst_i = blockIdx.x * blockDim.x + threadIdx.x; + // src: (b_size, c_in, h_in, w_in) + // k: (c_in, c_out, h_k, w_k) + const size_t *src_dims = info; + const size_t *src_s = info + 4; + const size_t *k_dims = info + 8; + const size_t *k_s = info + 12; + const size_t h_k = k_dims[2]; + const size_t w_k = k_dims[3]; + const size_t c_out = k_dims[1]; + const size_t c_in = src_dims[1]; + const size_t h_in = src_dims[2]; + const size_t w_in = src_dims[3]; + if (dst_i >= src_dims[0] * c_out * w_out * h_out) { + return; + } + + // TODO + const size_t b_idx = dst_i / (w_out * h_out * c_out); + const size_t dst_c_idx = (dst_i / (w_out * h_out)) % c_out; + // NCHW layout. + const size_t out_y = (dst_i / w_out) % h_out; + const size_t out_x = dst_i % w_out; + + const size_t src_idx0 = b_idx * src_s[0]; + A d = 0; + for (int k_x = 0; k_x < (int)w_k; ++k_x) { + // let out_x = inp_x * p.stride + k_x * p.dilation - p.padding; + int inp_x_stride = (int)(out_x + padding) - k_x * dilation; + if (inp_x_stride < 0 || inp_x_stride % stride) { + continue; + } + int inp_x = inp_x_stride / stride; + if (inp_x >= w_in) continue; + for (int k_y = 0; k_y < (int)h_k; ++k_y) { + int inp_y_stride = (int)(out_y + padding) - k_y * dilation; + if (inp_y_stride < 0 || inp_y_stride % stride) { + continue; + } + int inp_y = inp_y_stride / stride; + if (inp_y >= h_in) continue; + for (size_t src_c_idx = 0; src_c_idx < c_in; ++src_c_idx) { + const size_t src_idx = src_idx0 + src_c_idx * src_s[1] + inp_y * src_s[2] + inp_x * src_s[3]; + const size_t k_idx = src_c_idx * k_s[0] + dst_c_idx * k_s[1] + k_y * k_s[2] + k_x * k_s[3]; + d += static_cast(src[src_idx]) * static_cast(kernel[k_idx]); + } + } + } + dst[dst_i] = static_cast(d); +} + +template +__device__ void avg_pool2d( + const size_t src_numel, + const size_t w_k, + const size_t h_k, + const size_t w_stride, + const size_t h_stride, + const size_t *info, + const T *src, + T *dst +) { + const size_t dst_i = blockIdx.x * blockDim.x + threadIdx.x; + // src: (b_size, c_in, w_in, h_in) + const size_t *src_dims = info; + const size_t *src_s = info + 4; + + const size_t c = src_dims[1]; + const size_t w_in = src_dims[2]; + const size_t h_in = src_dims[3]; + + const size_t w_out = (w_in - w_k) / w_stride + 1; + const size_t h_out = (h_in - h_k) / h_stride + 1; + if (dst_i >= src_dims[0] * c * w_out * h_out) { + return; + } + + // TODO: Improve this. + const size_t b_idx = dst_i / (w_out * h_out * c); + const size_t c_idx = (dst_i / (w_out * h_out)) % c; + const size_t dst_w = (dst_i / h_out) % w_out; + const size_t dst_h = dst_i % h_out; + + const size_t src_idx0 = b_idx * src_s[0]; + const float scale = 1.0 / (w_k * h_k); + A d = 0; + for (size_t w_offset = 0; w_offset < w_k; ++w_offset) { + size_t src_w = w_stride * dst_w + w_offset; + if (src_w >= w_in) { + continue; + } + for (size_t h_offset = 0; h_offset < h_k; ++h_offset) { + size_t src_h = h_stride * dst_h + h_offset; + if (src_h >= h_in) { + continue; + } + const size_t src_idx = src_idx0 + c_idx * src_s[1] + src_w * src_s[2] + src_h * src_s[3]; + d += static_cast(src[src_idx]); + } + } + dst[dst_i] = static_cast(d * scale); +} + +template +__device__ void max_pool2d( + const size_t src_numel, + const size_t w_k, + const size_t h_k, + const size_t w_stride, + const size_t h_stride, + const size_t *info, + const T *src, + T *dst +) { + const size_t dst_i = blockIdx.x * blockDim.x + threadIdx.x; + // src: (b_size, c_in, w_in, h_in) + const size_t *src_dims = info; + const size_t *src_s = info + 4; + + const size_t c = src_dims[1]; + const size_t w_in = src_dims[2]; + const size_t h_in = src_dims[3]; + + const size_t w_out = (w_in - w_k) / w_stride + 1; + const size_t h_out = (h_in - h_k) / h_stride + 1; + if (dst_i >= src_dims[0] * c * w_out * h_out) { + return; + } + + // TODO: Improve this. + const size_t b_idx = dst_i / (w_out * h_out * c); + const size_t c_idx = (dst_i / (w_out * h_out)) % c; + const size_t dst_w = (dst_i / h_out) % w_out; + const size_t dst_h = dst_i % h_out; + + const size_t src_idx0 = b_idx * src_s[0]; + T d = 0; + bool set = false; + for (size_t w_offset = 0; w_offset < w_k; ++w_offset) { + size_t src_w = w_stride * dst_w + w_offset; + if (src_w >= w_in) { + continue; + } + for (size_t h_offset = 0; h_offset < h_k; ++h_offset) { + size_t src_h = h_stride * dst_h + h_offset; + if (src_h >= h_in) { + continue; + } + const size_t src_idx = src_idx0 + c_idx * src_s[1] + src_w * src_s[2] + src_h * src_s[3]; + if (set) { + d = maxg(d, src[src_idx]); + } + else { + d = src[src_idx]; + set = true; + } + } + } + dst[dst_i] = d; +} + +template +__device__ void upsample_nearest2d( + const size_t w_out, + const size_t h_out, + const double w_scale, + const double h_scale, + const size_t *info, + const T *src, + T *dst +) { + const size_t dst_i = blockIdx.x * blockDim.x + threadIdx.x; + // src: (b_size, c_in, w_in, h_in) + const size_t *src_dims = info; + const size_t *src_s = info + 4; + + const size_t c = src_dims[1]; + const size_t w_in = src_dims[2]; + const size_t h_in = src_dims[3]; + + if (dst_i >= src_dims[0] * c * w_out * h_out) { + return; + } + + // TODO: Improve this. + const size_t b_idx = dst_i / (w_out * h_out * c); + const size_t c_idx = (dst_i / (w_out * h_out)) % c; + const size_t dst_w = (dst_i / h_out) % w_out; + const size_t dst_h = dst_i % h_out; + + size_t src_w = static_cast(dst_w * w_scale); + size_t src_h = static_cast(dst_h * h_scale); + if (src_w >= w_in) { + src_w = w_in - 1; + } + if (src_h >= h_in) { + src_h = h_in - 1; + } + + const size_t src_i = b_idx * src_s[0] + c_idx * src_s[1] + src_w * src_s[2] + src_h * src_s[3]; + dst[dst_i] = src[src_i]; +} + +template +__device__ void upsample_bilinear2d( + const size_t w_out, + const size_t h_out, + const bool align_corners, + const bool has_scale_h, + const double scale_h_factor, + const bool has_scale_w, + const double scale_w_factor, + const size_t *info, + const scalar_t *src, + scalar_t *dst +) { + const size_t dst_i = blockIdx.x * blockDim.x + threadIdx.x; + + // src: (b_size, c_in, h_in, w_in) // Standard NCHW layout + const size_t *src_dims = info; + const size_t *src_s = info + 4; + + const size_t c = src_dims[1]; + const size_t h_in = src_dims[2]; // dims[2] = height + const size_t w_in = src_dims[3]; // dims[3] = width + + if (dst_i >= src_dims[0] * c * h_out * w_out) { + return; + } + + // Compute output position (NCHW layout) + const size_t b_idx = dst_i / (h_out * w_out * c); + const size_t c_idx = (dst_i / (h_out * w_out)) % c; + const size_t dst_h = (dst_i / w_out) % h_out; + const size_t dst_w = dst_i % w_out; + + // Calculate scale factors following PyTorch's area_pixel_compute_scale logic + double h_scale, w_scale; + if (align_corners) { + h_scale = (h_out > 1) ? static_cast(h_in - 1) / (h_out - 1) : 0.0; + w_scale = (w_out > 1) ? static_cast(w_in - 1) / (w_out - 1) : 0.0; + } else { + // PyTorch's compute_scales_value logic + h_scale = has_scale_h ? (1.0 / scale_h_factor) : (static_cast(h_in) / h_out); + w_scale = has_scale_w ? (1.0 / scale_w_factor) : (static_cast(w_in) / w_out); + } + + // Compute source position (floating point) + double src_h_fp, src_w_fp; + if (align_corners) { + src_h_fp = h_scale * dst_h; + src_w_fp = w_scale * dst_w; + } else { + src_h_fp = h_scale * (dst_h + 0.5) - 0.5; + src_w_fp = w_scale * (dst_w + 0.5) - 0.5; + } + + // Clamp to valid range + src_h_fp = fmax(0.0, src_h_fp); + src_w_fp = fmax(0.0, src_w_fp); + + // Get integer indices + size_t h0 = static_cast(floor(src_h_fp)); + size_t w0 = static_cast(floor(src_w_fp)); + size_t h1 = min(h0 + 1, h_in - 1); + size_t w1 = min(w0 + 1, w_in - 1); + + // Compute interpolation weights + double weight_h = src_h_fp - h0; + double weight_w = src_w_fp - w0; + weight_h = fmin(fmax(weight_h, 0.0), 1.0); + weight_w = fmin(fmax(weight_w, 0.0), 1.0); + + // Get base index + const size_t base = b_idx * src_s[0] + c_idx * src_s[1]; + + // Read four neighboring pixels + const scalar_t v00 = src[base + h0 * src_s[2] + w0 * src_s[3]]; + const scalar_t v10 = src[base + h0 * src_s[2] + w1 * src_s[3]]; + const scalar_t v01 = src[base + h1 * src_s[2] + w0 * src_s[3]]; + const scalar_t v11 = src[base + h1 * src_s[2] + w1 * src_s[3]]; + + // Bilinear interpolation + // Convert to double for computation to avoid type issues with __half and __nv_bfloat16 + const double v00_d = static_cast(v00); + const double v10_d = static_cast(v10); + const double v01_d = static_cast(v01); + const double v11_d = static_cast(v11); + + const double v_top = v00_d * (1.0 - weight_w) + v10_d * weight_w; + const double v_bottom = v01_d * (1.0 - weight_w) + v11_d * weight_w; + const double value = v_top * (1.0 - weight_h) + v_bottom * weight_h; + + dst[dst_i] = static_cast(value); +} + + +#define CONV1D_OP(TYPENAME, TYPEACC, FN_NAME) \ +extern "C" __global__ void FN_NAME( \ + const size_t src_numel, \ + const size_t num_dims, \ + const size_t stride, \ + const size_t padding, \ + const size_t dilation, \ + const size_t *info, \ + const TYPENAME *src, \ + const TYPENAME *kernel, \ + TYPENAME *dst \ +) { \ + conv1d(src_numel, num_dims, stride, padding, dilation, info, src, kernel, dst); \ +} \ + +#define CONV2D_OP(TYPENAME, TYPEACC, FN_NAME) \ +extern "C" __global__ void FN_NAME( \ + const size_t src_numel, \ + const size_t w_out, \ + const size_t h_out, \ + const size_t stride, \ + const size_t padding, \ + const size_t dilation, \ + const size_t *info, \ + const TYPENAME *src, \ + const TYPENAME *kernel, \ + TYPENAME *dst \ +) { \ + conv2d(src_numel, w_out, h_out, stride, padding, dilation, info, src, kernel, dst); \ +} \ + +#define IM2COL1D_OP(TYPENAME, FN_NAME) \ +extern "C" __global__ void FN_NAME( \ + const size_t dst_numel, \ + const size_t l_out, \ + const size_t l_k, \ + const size_t stride, \ + const size_t padding, \ + const size_t dilation, \ + const size_t *info, \ + const TYPENAME *src, \ + TYPENAME *dst \ +) { \ + im2col1d(dst_numel, l_out, l_k, stride, padding, dilation, info, src, dst); \ +} \ + +#define COL2IM1D_OP(TYPENAME, FN_NAME) \ +extern "C" __global__ void FN_NAME( \ + const size_t dst_el, \ + const size_t l_out, \ + const size_t l_in, \ + const size_t c_out, \ + const size_t k_size, \ + const size_t stride, \ + const TYPENAME *src, \ + TYPENAME *dst \ +) { \ + col2im1d(dst_el, l_out, l_in, c_out, k_size, stride, src, dst); \ +} \ + +#define IM2COL_OP(TYPENAME, FN_NAME) \ +extern "C" __global__ void FN_NAME( \ + const size_t dst_numel, \ + const size_t h_out, \ + const size_t w_out, \ + const size_t h_k, \ + const size_t w_k, \ + const size_t stride, \ + const size_t padding, \ + const size_t dilation, \ + const size_t *info, \ + const TYPENAME *src, \ + TYPENAME *dst \ +) { \ + im2col(dst_numel, h_out, w_out, h_k, w_k, stride, padding, dilation, info, src, dst); \ +} \ + +#define CONVT1D_OP(TYPENAME, TYPEACC, FN_NAME) \ +extern "C" __global__ void FN_NAME( \ + const size_t src_numel, \ + const size_t l_out, \ + const size_t stride, \ + const size_t padding, \ + const size_t out_padding, \ + const size_t dilation, \ + const size_t *info, \ + const TYPENAME *src, \ + const TYPENAME *kernel, \ + TYPENAME *dst \ +) { \ + conv_transpose1d(src_numel, l_out, stride, padding, out_padding, dilation, info, src, kernel, dst); \ +} \ + +#define CONVT2D_OP(TYPENAME, TYPEACC, FN_NAME) \ +extern "C" __global__ void FN_NAME( \ + const size_t src_numel, \ + const size_t w_out, \ + const size_t h_out, \ + const size_t stride, \ + const size_t padding, \ + const size_t out_padding, \ + const size_t dilation, \ + const size_t *info, \ + const TYPENAME *src, \ + const TYPENAME *kernel, \ + TYPENAME *dst \ +) { \ + conv_transpose2d(src_numel, w_out, h_out, stride, padding, out_padding, dilation, info, src, kernel, dst); \ +} \ + +#define AVG_POOL2D_OP(TYPENAME, TYPEACC, FN_NAME) \ +extern "C" __global__ void FN_NAME( \ + const size_t src_numel, \ + const size_t w_k, \ + const size_t h_k, \ + const size_t w_stride, \ + const size_t h_stride, \ + const size_t *info, \ + const TYPENAME *src, \ + TYPENAME *dst \ +) { \ + avg_pool2d(src_numel, w_k, h_k, w_stride, h_stride, info, src, dst); \ +} \ + +#define MAX_POOL2D_OP(TYPENAME, FN_NAME) \ +extern "C" __global__ void FN_NAME( \ + const size_t src_numel, \ + const size_t w_k, \ + const size_t h_k, \ + const size_t w_stride, \ + const size_t h_stride, \ + const size_t *info, \ + const TYPENAME *src, \ + TYPENAME *dst \ +) { \ + max_pool2d(src_numel, w_k, h_k, w_stride, h_stride, info, src, dst); \ +} \ + +#define UPSAMPLE_NEAREST2D_OP(TYPENAME, FN_NAME) \ +extern "C" __global__ void FN_NAME( \ + const size_t w_out, \ + const size_t h_out, \ + const double w_scale, \ + const double h_scale, \ + const size_t *info, \ + const TYPENAME *src, \ + TYPENAME *dst \ +) { \ + upsample_nearest2d(w_out, h_out, w_scale, h_scale, info, src, dst); \ +} \ + +#define UPSAMPLE_BILINEAR2D_OP(TYPENAME, FN_NAME) \ +extern "C" __global__ void FN_NAME( \ + const size_t w_out, \ + const size_t h_out, \ + const bool align_corners, \ + const bool has_scale_h, \ + const double scale_h_factor, \ + const bool has_scale_w, \ + const double scale_w_factor, \ + const size_t *info, \ + const TYPENAME *src, \ + TYPENAME *dst \ +) { \ + upsample_bilinear2d(w_out, h_out, align_corners, has_scale_h, scale_h_factor, has_scale_w, scale_w_factor, info, src, dst); \ +} \ + +#if __CUDA_ARCH__ >= 800 +CONV1D_OP(__nv_bfloat16, float, conv1d_bf16) +CONV2D_OP(__nv_bfloat16, float, conv2d_bf16) +CONVT1D_OP(__nv_bfloat16, float, conv_transpose1d_bf16) +CONVT2D_OP(__nv_bfloat16, float, conv_transpose2d_bf16) +AVG_POOL2D_OP(__nv_bfloat16, float, avg_pool2d_bf16) +MAX_POOL2D_OP(__nv_bfloat16, max_pool2d_bf16) +UPSAMPLE_NEAREST2D_OP(__nv_bfloat16, upsample_nearest2d_bf16) +UPSAMPLE_BILINEAR2D_OP(__nv_bfloat16, upsample_bilinear2d_bf16) +IM2COL_OP(__nv_bfloat16, im2col_bf16) +IM2COL1D_OP(__nv_bfloat16, im2col1d_bf16) +COL2IM1D_OP(__nv_bfloat16, col2im1d_bf16) + +// NOTE: No conv ops for f8 +// CONV1D_OP(__nv_bfloat16, float, conv1d_f8_e5m) +// CONV2D_OP(__nv_fp8_e4m3, float, conv2d_f8_e5m) +// CONVT1D_OP(__nv_fp8_e4m3, float, conv_transpose1d_f8_e5m) +// CONVT2D_OP(__nv_fp8_e4m3, float, conv_transpose2d_f8_e5m) +// AVG_POOL2D_OP(__nv_fp8_e4m3, float, avg_pool2d_f8_e5m) +// MAX_POOL2D_OP(__nv_fp8_e4m3, max_pool2d_f8_e5m) +// UPSAMPLE_NEAREST2D_OP(__nv_fp8_e4m3, upsample_nearest2d_f8_e5m) +// IM2COL_OP(__nv_fp8_e4m3, im2col_f8_e5m) +// IM2COL1D_OP(__nv_fp8_e4m3, im2col1d_f8_e5m) +// COL2IM1D_OP(__nv_fp8_e4m3, col2im1d_f8_e5m) +#endif + +#if __CUDA_ARCH__ >= 530 +CONV1D_OP(__half, float, conv1d_f16) +CONV2D_OP(__half, float, conv2d_f16) +CONVT1D_OP(__half, float, conv_transpose1d_f16) +CONVT2D_OP(__half, float, conv_transpose2d_f16) +AVG_POOL2D_OP(__half, float, avg_pool2d_f16) +MAX_POOL2D_OP(__half, max_pool2d_f16) +UPSAMPLE_NEAREST2D_OP(__half, upsample_nearest2d_f16) +UPSAMPLE_BILINEAR2D_OP(__half, upsample_bilinear2d_f16) +IM2COL_OP(__half, im2col_f16) +IM2COL1D_OP(__half, im2col1d_f16) +COL2IM1D_OP(__half, col2im1d_f16) +#endif + +CONV1D_OP(float, float, conv1d_f32) +CONV1D_OP(double, double, conv1d_f64) +CONV1D_OP(uint8_t, uint8_t, conv1d_u8) +CONV1D_OP(uint32_t, uint32_t, conv1d_u32) + +CONV2D_OP(float, float, conv2d_f32) +CONV2D_OP(double, double, conv2d_f64) +CONV2D_OP(uint8_t, uint8_t, conv2d_u8) +CONV2D_OP(uint32_t, uint32_t, conv2d_u32) + +CONVT1D_OP(float, float, conv_transpose1d_f32) +CONVT1D_OP(double, double, conv_transpose1d_f64) +CONVT1D_OP(uint8_t, uint8_t, conv_transpose1d_u8) +CONVT1D_OP(uint32_t, uint32_t, conv_transpose1d_u32) + +CONVT2D_OP(float, float, conv_transpose2d_f32) +CONVT2D_OP(double, double, conv_transpose2d_f64) +CONVT2D_OP(uint8_t, uint8_t, conv_transpose2d_u8) +CONVT2D_OP(uint32_t, uint32_t, conv_transpose2d_u32) + +AVG_POOL2D_OP(float, float, avg_pool2d_f32) +AVG_POOL2D_OP(double, double, avg_pool2d_f64) +AVG_POOL2D_OP(uint8_t, uint8_t, avg_pool2d_u8) +AVG_POOL2D_OP(uint32_t, uint32_t, avg_pool2d_u32) + +MAX_POOL2D_OP(float, max_pool2d_f32) +MAX_POOL2D_OP(double, max_pool2d_f64) +MAX_POOL2D_OP(uint8_t, max_pool2d_u8) +MAX_POOL2D_OP(uint32_t, max_pool2d_u32) + +UPSAMPLE_NEAREST2D_OP(float, upsample_nearest2d_f32) +UPSAMPLE_NEAREST2D_OP(double, upsample_nearest2d_f64) +UPSAMPLE_NEAREST2D_OP(uint8_t, upsample_nearest2d_u8) +UPSAMPLE_NEAREST2D_OP(uint32_t, upsample_nearest2d_u32) + +UPSAMPLE_BILINEAR2D_OP(float, upsample_bilinear2d_f32) +UPSAMPLE_BILINEAR2D_OP(double, upsample_bilinear2d_f64) +UPSAMPLE_BILINEAR2D_OP(uint8_t, upsample_bilinear2d_u8) +UPSAMPLE_BILINEAR2D_OP(uint32_t, upsample_bilinear2d_u32) + +IM2COL_OP(float, im2col_f32) +IM2COL_OP(double, im2col_f64) +IM2COL_OP(uint8_t, im2col_u8) +IM2COL_OP(uint32_t, im2col_u32) + +IM2COL1D_OP(float, im2col1d_f32) +IM2COL1D_OP(double, im2col1d_f64) +IM2COL1D_OP(uint8_t, im2col1d_u8) +IM2COL1D_OP(uint32_t, im2col1d_u32) + +COL2IM1D_OP(float, col2im1d_f32) +COL2IM1D_OP(double, col2im1d_f64) +COL2IM1D_OP(uint8_t, col2im1d_u8) +COL2IM1D_OP(uint32_t, col2im1d_u32) diff --git a/torchnative/rust/vendor/candle-kernels/src/cuda_utils.cuh b/torchnative/rust/vendor/candle-kernels/src/cuda_utils.cuh new file mode 100644 index 00000000..eb1400b4 --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/cuda_utils.cuh @@ -0,0 +1,224 @@ +#include "compatibility.cuh" +#include +#include + +// TODO: This is often used to check that the data is contiguous so that +// kernels can be easily mapped. However this only returns true for row +// major, if all the inputs are column major, we could apply the fast path +// too (but we wouldn't if some of them are row major and some column major). +__device__ bool is_contiguous( + const size_t num_dims, + const size_t *dims, + const size_t *strides +) { + size_t acc = 1; + for (unsigned int d = 0; d < num_dims; d++) { + unsigned int dim_idx = num_dims - 1 - d; + if (dims[dim_idx] > 1 && acc != strides[dim_idx]) { + return false; + } + acc *= dims[dim_idx]; + } + return true; +} + +__device__ unsigned int get_strided_index( + unsigned int idx, + const size_t num_dims, + const size_t *dims, + const size_t *strides +) { + unsigned int strided_i = 0; + for (unsigned int d = 0; d < num_dims; d++) { + unsigned int dim_idx = num_dims - 1 - d; + strided_i += (idx % dims[dim_idx]) * strides[dim_idx]; + idx /= dims[dim_idx]; + } + return strided_i; +} + +__device__ unsigned int restrided( + const unsigned int strided_i, + const size_t num_dims, + const size_t *dims, + const size_t *strides, + const size_t *new_strides +) { + unsigned int idx = 0; + for (int d = 0; d < num_dims; d++) { + idx += (strides[d] == 0 ? 0 : (strided_i / strides[d]) % dims[d]) * new_strides[d]; + } + return idx; +} + +// Sourced from https://graphics.stanford.edu/~seander/bithacks.html#RoundUpPowerOf2 +// Input must be less than or equal to 2 ^ 16 +// used in reductions +__device__ __forceinline__ unsigned int next_power_of_two(unsigned int v) { + v--; + v |= v >> 1; + v |= v >> 2; + v |= v >> 4; + v |= v >> 8; + v++; + return v; +} + +// Efficiently computes the sum of each chunk in "data" of size chunk_len, and +// stores the sums in out[i / chunk_len] +template +__device__ void chunk_sum( + const size_t chunk_len, + const T data, + T* out +) { + __shared__ T buf[1024]; + + // assumes that threads where i >= numel have already exited + unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; + unsigned int block_i = threadIdx.x; + + // Fall back to atomicAdd if chunk_len is small to reduce overhead + if (chunk_len <= 2) { + atomicAdd(out + i / chunk_len, data); + return; + } + buf[block_i] = data; + + unsigned int chunk_i = i % chunk_len; + unsigned int chunk_start = max((int)(block_i - chunk_i), 0); + unsigned int chunk_end = min((unsigned int)(block_i + chunk_len - chunk_i), blockDim.x); + + chunk_i = block_i - chunk_start; + + size_t max_chunk_len = min(chunk_end - chunk_start, blockDim.x); + size_t incr = next_power_of_two(max_chunk_len) >> 1; + + __syncthreads(); + + // Uses sequential addressing as discussed in + // https://developer.download.nvidia.com/assets/cuda/files/reduction.pdf + for (; incr > 0; incr >>= 1) { + unsigned int block_i_2 = block_i + incr; + + if (block_i_2 < chunk_end && chunk_i < incr) { + // This is sound because __syncthreads and the conditions above + // ensure that no data races occur + buf[block_i] += buf[block_i_2]; + } + + __syncthreads(); + } + + if (block_i == chunk_start) { + atomicAdd(out + i / chunk_len, buf[block_i]); + } +} + +__device__ __forceinline__ bool isnang(float a) { return isnan(a); } +__device__ __forceinline__ bool isnang(double a) { return isnan(a); } +__device__ __forceinline__ float recipg(float a) { return 1.0 / a; } +__device__ __forceinline__ double recipg(double a) { return 1.0 / a; } +__device__ __forceinline__ float cosg(float a) { return cosf(a); } +__device__ __forceinline__ double cosg(double a) { return cos(a); } +__device__ __forceinline__ float sing(float a) { return sinf(a); } +__device__ __forceinline__ double sing(double a) { return sin(a); } +__device__ __forceinline__ float sqrtg(float a) { return sqrtf(a); } +__device__ __forceinline__ double sqrtg(double a) { return sqrt(a); } +__device__ __forceinline__ float powg(float a, float b) { return powf(a, b); } +__device__ __forceinline__ double powg(double a, double b) { return pow(a, b); } +__device__ __forceinline__ float tanhg(float a) { return tanhf(a); } +__device__ __forceinline__ double tanhg(double a) { return tanh(a); } +__device__ __forceinline__ float erfg(float a) { return erff(a); } +__device__ __forceinline__ double erfg(double a) { return erf(a); } +__device__ __forceinline__ float ceilg(float a) { return ceilf(a); } +__device__ __forceinline__ double ceilg(double a) { return ceil(a); } +__device__ __forceinline__ float floorg(float a) { return floorf(a); } +__device__ __forceinline__ double floorg(double a) { return floor(a); } +__device__ __forceinline__ float roundg(float a) { return roundf(a); } +__device__ __forceinline__ double roundg(double a) { return round(a); } +__device__ __forceinline__ float normcdfg(float a) { return normcdff(a); } +__device__ __forceinline__ double normcdfg(double a) { return normcdf(a); } +__device__ __forceinline__ float maxg(float a, float b) { return fmaxf(a, b); } +__device__ __forceinline__ double maxg(double a, double b) { return fmax(a, b); } +__device__ __forceinline__ float ming(float a, float b) { return fminf(a, b); } +__device__ __forceinline__ double ming(double a, double b) { return fmin(a, b); } +__device__ __forceinline__ float logg(float a) { return logf(a); } +__device__ __forceinline__ double logg(double a) { return log(a); } +__device__ __forceinline__ float expg(float a) { return expf(a); } +__device__ __forceinline__ double expg(double a) { return exp(a); } +__device__ __forceinline__ float absg(float a) { return fabsf(a); } +__device__ __forceinline__ double absg(double a) { return fabs(a); } +__device__ __forceinline__ float copysigng(float a, float b) { return copysignf(a, b); } +__device__ __forceinline__ double copysigng(double a, double b) { return copysign(a, b); } + +__device__ __forceinline__ int64_t ming(int64_t a, int64_t b) { return min(a, b); } +__device__ __forceinline__ int64_t maxg(int64_t a, int64_t b) { return max(a, b); } +__device__ __forceinline__ uint32_t ming(uint32_t a, uint32_t b) { return min(a, b); } +__device__ __forceinline__ uint32_t maxg(uint32_t a, uint32_t b) { return max(a, b); } +__device__ __forceinline__ uint8_t ming(uint8_t a, uint8_t b) { return min(a, b); } +__device__ __forceinline__ uint8_t maxg(uint8_t a, uint8_t b) { return max(a, b); } +#if __CUDA_ARCH__ >= 530 +__device__ __forceinline__ __half powg(__half a, __half b) { return __float2half(powf(__half2float(a), __half2float(b))); } +__device__ __forceinline__ bool isnang(__half a) { return __hisnan(a); } +__device__ __forceinline__ __half sqrtg(__half a) { return hsqrt(a); } +__device__ __forceinline__ __half cosg(__half a) { return hcos(a); } +__device__ __forceinline__ __half sing(__half a) { return hsin(a); } +__device__ __forceinline__ __half recipg(__half a) { __half one = 1.0; return one / a; } +__device__ __forceinline__ __half maxg(__half a, __half b) { return __hmax_nan(a, b); } +__device__ __forceinline__ __half tanhg(__half a) { return __float2half(tanhf(__half2float(a))); } +__device__ __forceinline__ __half erfg(__half a) { return __float2half(erff(__half2float(a))); } +__device__ __forceinline__ __half ceilg(__half a) { return __float2half(ceilf(__half2float(a))); } +__device__ __forceinline__ __half floorg(__half a) { return __float2half(floorf(__half2float(a))); } +__device__ __forceinline__ __half roundg(__half a) { return __float2half(roundf(__half2float(a))); } +__device__ __forceinline__ __half normcdfg(__half a) { return __float2half(normcdff(__half2float(a))); } +__device__ __forceinline__ __half ming(__half a, __half b) { return __hmin_nan(a, b); } +__device__ __forceinline__ __half logg(__half a) { return hlog(a); } +__device__ __forceinline__ __half expg(__half a) { return hexp(a); } +__device__ __forceinline__ __half absg(__half a) { return __habs(a); } +__device__ __forceinline__ __half copysigng(__half a, __half b) { return __float2half(copysignf(__half2float(a), __half2float(b))); } +#endif + +#if __CUDA_ARCH__ >= 800 +__device__ __forceinline__ __nv_bfloat16 powg(__nv_bfloat16 a, __nv_bfloat16 b) { return __float2bfloat16(powf(__bfloat162float(a), __bfloat162float(b))); } +__device__ __forceinline__ bool isnang(__nv_bfloat16 a) { return __hisnan(a); } +__device__ __forceinline__ __nv_bfloat16 sqrtg(__nv_bfloat16 a) { return hsqrt(a); } +__device__ __forceinline__ __nv_bfloat16 cosg(__nv_bfloat16 a) { return hcos(a); } +__device__ __forceinline__ __nv_bfloat16 sing(__nv_bfloat16 a) { return hsin(a); } +__device__ __forceinline__ __nv_bfloat16 recipg(__nv_bfloat16 a) { __nv_bfloat16 one = 1.0; return one / a; } +__device__ __forceinline__ __nv_bfloat16 maxg(__nv_bfloat16 a, __nv_bfloat16 b) { return __hmax_nan(a, b); } +__device__ __forceinline__ __nv_bfloat16 tanhg(__nv_bfloat16 a) { return __float2bfloat16(tanhf(__bfloat162float(a))); } +__device__ __forceinline__ __nv_bfloat16 erfg(__nv_bfloat16 a) { return __float2bfloat16(erff(__bfloat162float(a))); } +__device__ __forceinline__ __nv_bfloat16 ceilg(__nv_bfloat16 a) { return __float2bfloat16(ceilf(__bfloat162float(a))); } +__device__ __forceinline__ __nv_bfloat16 floorg(__nv_bfloat16 a) { return __float2bfloat16(floorf(__bfloat162float(a))); } +__device__ __forceinline__ __nv_bfloat16 roundg(__nv_bfloat16 a) { return __float2bfloat16(roundf(__bfloat162float(a))); } +__device__ __forceinline__ __nv_bfloat16 normcdfg(__nv_bfloat16 a) { return __float2bfloat16(normcdff(__bfloat162float(a))); } +__device__ __forceinline__ __nv_bfloat16 ming(__nv_bfloat16 a, __nv_bfloat16 b) { return __hmin_nan(a, b); } +__device__ __forceinline__ __nv_bfloat16 logg(__nv_bfloat16 a) { return hlog(a); } +__device__ __forceinline__ __nv_bfloat16 expg(__nv_bfloat16 a) { return hexp(a); } +__device__ __forceinline__ __nv_bfloat16 absg(__nv_bfloat16 a) { return __habs(a); } +__device__ __forceinline__ __nv_bfloat16 copysigng(__nv_bfloat16 a, __nv_bfloat16 b) { return __float2bfloat16(copysignf(__bfloat162float(a), __bfloat162float(b))); } + +#define F8E4M3_TO_FLOAT(x) __half2float(__nv_cvt_fp8_to_halfraw(x.__x, __NV_E4M3)) + +__device__ __forceinline__ __nv_fp8_e4m3 powg(__nv_fp8_e4m3 a, __nv_fp8_e4m3 b) { return __nv_fp8_e4m3(powf(F8E4M3_TO_FLOAT(a), F8E4M3_TO_FLOAT(b))); } +__device__ __forceinline__ bool isnang(__nv_fp8_e4m3 a) { return isnan(F8E4M3_TO_FLOAT(a)); } +__device__ __forceinline__ __nv_fp8_e4m3 sqrtg(__nv_fp8_e4m3 a) { return __nv_fp8_e4m3(sqrtf(F8E4M3_TO_FLOAT(a))); } +__device__ __forceinline__ __nv_fp8_e4m3 cosg(__nv_fp8_e4m3 a) { return __nv_fp8_e4m3(cosf(F8E4M3_TO_FLOAT(a))); } +__device__ __forceinline__ __nv_fp8_e4m3 sing(__nv_fp8_e4m3 a) { return __nv_fp8_e4m3(sinf(F8E4M3_TO_FLOAT(a))); } +__device__ __forceinline__ __nv_fp8_e4m3 recipg(__nv_fp8_e4m3 a) { return __nv_fp8_e4m3(1. / F8E4M3_TO_FLOAT(a)); } +__device__ __forceinline__ __nv_fp8_e4m3 maxg(__nv_fp8_e4m3 a, __nv_fp8_e4m3 b) { return __nv_fp8_e4m3(fmaxf(F8E4M3_TO_FLOAT(a), F8E4M3_TO_FLOAT(b))); } +__device__ __forceinline__ __nv_fp8_e4m3 tanhg(__nv_fp8_e4m3 a) { return __nv_fp8_e4m3(tanhf(F8E4M3_TO_FLOAT(a))); } +__device__ __forceinline__ __nv_fp8_e4m3 erfg(__nv_fp8_e4m3 a) { return __nv_fp8_e4m3(erff(F8E4M3_TO_FLOAT(a))); } +__device__ __forceinline__ __nv_fp8_e4m3 ceilg(__nv_fp8_e4m3 a) { return __nv_fp8_e4m3(ceilf(F8E4M3_TO_FLOAT(a))); } +__device__ __forceinline__ __nv_fp8_e4m3 floorg(__nv_fp8_e4m3 a) { return __nv_fp8_e4m3(floorf(F8E4M3_TO_FLOAT(a))); } +__device__ __forceinline__ __nv_fp8_e4m3 roundg(__nv_fp8_e4m3 a) { return __nv_fp8_e4m3(roundf(F8E4M3_TO_FLOAT(a))); } +__device__ __forceinline__ __nv_fp8_e4m3 normcdfg(__nv_fp8_e4m3 a) { return __nv_fp8_e4m3(normcdff(F8E4M3_TO_FLOAT(a))); } +__device__ __forceinline__ __nv_fp8_e4m3 ming(__nv_fp8_e4m3 a, __nv_fp8_e4m3 b) { return __nv_fp8_e4m3(fminf(F8E4M3_TO_FLOAT(a), F8E4M3_TO_FLOAT(b))); } +__device__ __forceinline__ __nv_fp8_e4m3 logg(__nv_fp8_e4m3 a) { return __nv_fp8_e4m3(logf(F8E4M3_TO_FLOAT(a))); } +__device__ __forceinline__ __nv_fp8_e4m3 expg(__nv_fp8_e4m3 a) { return __nv_fp8_e4m3(expf(F8E4M3_TO_FLOAT(a))); } +__device__ __forceinline__ __nv_fp8_e4m3 absg(__nv_fp8_e4m3 a) { return __nv_fp8_e4m3(fabsf(F8E4M3_TO_FLOAT(a))); } +__device__ __forceinline__ __nv_fp8_e4m3 copysigng(__nv_fp8_e4m3 a, __nv_fp8_e4m3 b) { return __nv_fp8_e4m3(copysignf(F8E4M3_TO_FLOAT(a), F8E4M3_TO_FLOAT(b))); } + + +#endif diff --git a/torchnative/rust/vendor/candle-kernels/src/ffi.rs b/torchnative/rust/vendor/candle-kernels/src/ffi.rs new file mode 100644 index 00000000..11dcbacd --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/ffi.rs @@ -0,0 +1,631 @@ +use core::ffi::c_void; +#[allow(dead_code)] +#[allow(improper_ctypes)] +extern "C" { + // for unquntized models + pub fn moe_gemm_wmma( + input: *const c_void, // device pointer [size_m, size_k] + weights: *const c_void, // device pointer [num_experts, size_n, size_k] + sorted_token_ids: *const i32, // device pointer [size_m] + expert_ids: *const i32, // host array [size_m] (expert id per sorted token) + topk_weights: *const f32, + output: *mut c_void, // device pointer [size_m, size_n] + expert_counts: *mut i32, // pre-allocated buffer [num_experts] + expert_offsets: *mut i32, // pre-allocated buffer [num_experts + 1] + num_experts: i32, + topk: i32, + size_m: i32, + size_n: i32, + size_k: i32, + dtype: i32, // 0=float16, 1=bf16 (for input/output) + is_prefill: bool, + stream: i64, + ); + + pub fn moe_gemm_gguf( + input: *const f32, // input [size_m, size_k] + weights: *const c_void, // weights [num_experts, size_n, size_k] + sorted_token_ids: *const i32, + expert_ids: *const i32, + topk_weights: *const f32, // device ptr or nullptr + output: *mut c_void, // float output [size_m, size_n] + num_experts: i32, + topk: i32, + size_m: i32, + size_n: i32, + size_k: i32, + gguf_dtype: i32, // Q8_0: 0, Q4K: 1, Q2K: 2, Q3k: 3, Q5K: 4, Q6K: 5 (for weights) + stream: i64, + ); + + pub fn moe_gemm_gguf_prefill( + input: *const c_void, // input [size_m, size_k] + weights: *const u8, // weights [num_experts, size_n, size_k] + sorted_token_ids: *const i32, + expert_ids: *const i32, //must be host ptr + topk_weights: *const f32, // device ptr or nullptr + output: *mut c_void, // float output [size_m, size_n] + num_experts: i32, + topk: i32, + size_m: i32, + size_n: i32, + size_k: i32, + input_dtype: i32, // 0=f16, 1=bf16 (for inputs) + gguf_dtype: i32, //Q8_0: 0, Q4K: 1, Q2K: 2, Q3k: 3, Q5K: 4, Q6K: 5 (for weights) + stream: i64, + ); + + // ============== Dense GGUF MMVQ launchers (from mmvq_gguf.cu) ============== + + // BF16 output launchers + pub fn launch_mmvq_gguf_q4_0_bf16_plain( + vx: *const c_void, + vy: *const c_void, + dst: *mut c_void, + ncols_x: i32, + nrows_x: i32, + stride_col_y: i32, + stride_col_dst: i32, + b_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmvq_gguf_q4_1_bf16_plain( + vx: *const c_void, + vy: *const c_void, + dst: *mut c_void, + ncols_x: i32, + nrows_x: i32, + stride_col_y: i32, + stride_col_dst: i32, + b_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmvq_gguf_q5_0_bf16_plain( + vx: *const c_void, + vy: *const c_void, + dst: *mut c_void, + ncols_x: i32, + nrows_x: i32, + stride_col_y: i32, + stride_col_dst: i32, + b_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmvq_gguf_q5_1_bf16_plain( + vx: *const c_void, + vy: *const c_void, + dst: *mut c_void, + ncols_x: i32, + nrows_x: i32, + stride_col_y: i32, + stride_col_dst: i32, + b_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmvq_gguf_q8_0_bf16_plain( + vx: *const c_void, + vy: *const c_void, + dst: *mut c_void, + ncols_x: i32, + nrows_x: i32, + stride_col_y: i32, + stride_col_dst: i32, + b_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmvq_gguf_q2_k_bf16_plain( + vx: *const c_void, + vy: *const c_void, + dst: *mut c_void, + ncols_x: i32, + nrows_x: i32, + stride_col_y: i32, + stride_col_dst: i32, + b_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmvq_gguf_q3_k_bf16_plain( + vx: *const c_void, + vy: *const c_void, + dst: *mut c_void, + ncols_x: i32, + nrows_x: i32, + stride_col_y: i32, + stride_col_dst: i32, + b_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmvq_gguf_q4_k_bf16_plain( + vx: *const c_void, + vy: *const c_void, + dst: *mut c_void, + ncols_x: i32, + nrows_x: i32, + stride_col_y: i32, + stride_col_dst: i32, + b_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmvq_gguf_q5_k_bf16_plain( + vx: *const c_void, + vy: *const c_void, + dst: *mut c_void, + ncols_x: i32, + nrows_x: i32, + stride_col_y: i32, + stride_col_dst: i32, + b_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmvq_gguf_q6_k_bf16_plain( + vx: *const c_void, + vy: *const c_void, + dst: *mut c_void, + ncols_x: i32, + nrows_x: i32, + stride_col_y: i32, + stride_col_dst: i32, + b_size: i32, + stream: *mut c_void, + ); + + // F32 output launchers + pub fn launch_mmvq_gguf_q4_0_f32_plain( + vx: *const c_void, + vy: *const c_void, + dst: *mut c_void, + ncols_x: i32, + nrows_x: i32, + stride_col_y: i32, + stride_col_dst: i32, + b_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmvq_gguf_q4_1_f32_plain( + vx: *const c_void, + vy: *const c_void, + dst: *mut c_void, + ncols_x: i32, + nrows_x: i32, + stride_col_y: i32, + stride_col_dst: i32, + b_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmvq_gguf_q5_0_f32_plain( + vx: *const c_void, + vy: *const c_void, + dst: *mut c_void, + ncols_x: i32, + nrows_x: i32, + stride_col_y: i32, + stride_col_dst: i32, + b_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmvq_gguf_q5_1_f32_plain( + vx: *const c_void, + vy: *const c_void, + dst: *mut c_void, + ncols_x: i32, + nrows_x: i32, + stride_col_y: i32, + stride_col_dst: i32, + b_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmvq_gguf_q8_0_f32_plain( + vx: *const c_void, + vy: *const c_void, + dst: *mut c_void, + ncols_x: i32, + nrows_x: i32, + stride_col_y: i32, + stride_col_dst: i32, + b_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmvq_gguf_q2_k_f32_plain( + vx: *const c_void, + vy: *const c_void, + dst: *mut c_void, + ncols_x: i32, + nrows_x: i32, + stride_col_y: i32, + stride_col_dst: i32, + b_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmvq_gguf_q3_k_f32_plain( + vx: *const c_void, + vy: *const c_void, + dst: *mut c_void, + ncols_x: i32, + nrows_x: i32, + stride_col_y: i32, + stride_col_dst: i32, + b_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmvq_gguf_q4_k_f32_plain( + vx: *const c_void, + vy: *const c_void, + dst: *mut c_void, + ncols_x: i32, + nrows_x: i32, + stride_col_y: i32, + stride_col_dst: i32, + b_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmvq_gguf_q5_k_f32_plain( + vx: *const c_void, + vy: *const c_void, + dst: *mut c_void, + ncols_x: i32, + nrows_x: i32, + stride_col_y: i32, + stride_col_dst: i32, + b_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmvq_gguf_q6_k_f32_plain( + vx: *const c_void, + vy: *const c_void, + dst: *mut c_void, + ncols_x: i32, + nrows_x: i32, + stride_col_y: i32, + stride_col_dst: i32, + b_size: i32, + stream: *mut c_void, + ); + + pub fn launch_mmvq_gguf_q4_0_f16_plain( + vx: *const c_void, + vy: *const c_void, + dst: *mut c_void, + ncols_x: i32, + nrows_x: i32, + stride_col_y: i32, + stride_col_dst: i32, + b_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmvq_gguf_q4_1_f16_plain( + vx: *const c_void, + vy: *const c_void, + dst: *mut c_void, + ncols_x: i32, + nrows_x: i32, + stride_col_y: i32, + stride_col_dst: i32, + b_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmvq_gguf_q5_0_f16_plain( + vx: *const c_void, + vy: *const c_void, + dst: *mut c_void, + ncols_x: i32, + nrows_x: i32, + stride_col_y: i32, + stride_col_dst: i32, + b_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmvq_gguf_q5_1_f16_plain( + vx: *const c_void, + vy: *const c_void, + dst: *mut c_void, + ncols_x: i32, + nrows_x: i32, + stride_col_y: i32, + stride_col_dst: i32, + b_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmvq_gguf_q8_0_f16_plain( + vx: *const c_void, + vy: *const c_void, + dst: *mut c_void, + ncols_x: i32, + nrows_x: i32, + stride_col_y: i32, + stride_col_dst: i32, + b_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmvq_gguf_q2_k_f16_plain( + vx: *const c_void, + vy: *const c_void, + dst: *mut c_void, + ncols_x: i32, + nrows_x: i32, + stride_col_y: i32, + stride_col_dst: i32, + b_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmvq_gguf_q3_k_f16_plain( + vx: *const c_void, + vy: *const c_void, + dst: *mut c_void, + ncols_x: i32, + nrows_x: i32, + stride_col_y: i32, + stride_col_dst: i32, + b_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmvq_gguf_q4_k_f16_plain( + vx: *const c_void, + vy: *const c_void, + dst: *mut c_void, + ncols_x: i32, + nrows_x: i32, + stride_col_y: i32, + stride_col_dst: i32, + b_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmvq_gguf_q5_k_f16_plain( + vx: *const c_void, + vy: *const c_void, + dst: *mut c_void, + ncols_x: i32, + nrows_x: i32, + stride_col_y: i32, + stride_col_dst: i32, + b_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmvq_gguf_q6_k_f16_plain( + vx: *const c_void, + vy: *const c_void, + dst: *mut c_void, + ncols_x: i32, + nrows_x: i32, + stride_col_y: i32, + stride_col_dst: i32, + b_size: i32, + stream: *mut c_void, + ); + + // Quantize launchers (activation → Q8_1) + pub fn launch_mmvq_gguf_quantize_q8_1_bf16( + x: *const c_void, + vy: *mut c_void, + kx: i32, + kx_padded: i32, + num_rows: i32, + stream: *mut c_void, + ); + pub fn launch_mmvq_gguf_quantize_q8_1_f16( + x: *const c_void, + vy: *mut c_void, + kx: i32, + kx_padded: i32, + num_rows: i32, + stream: *mut c_void, + ); + pub fn launch_mmvq_gguf_quantize_q8_1_f32( + x: *const c_void, + vy: *mut c_void, + kx: i32, + kx_padded: i32, + num_rows: i32, + stream: *mut c_void, + ); + + // ============== Dense GGUF MMQ launchers (from mmq_gguf/) ============== + + // MMQ quantize launchers (f32 -> block_q8_1_mmq, 3 scale layouts) + pub fn launch_mmq_quantize_q8_1_D4( + x: *const c_void, + ids: *const i32, + vy: *mut c_void, + type_x: i32, + ne00: i64, + s01: i64, + s02: i64, + s03: i64, + ne0: i64, + ne1: i64, + ne2: i64, + ne3: i64, + stream: *mut c_void, + ); + pub fn launch_mmq_quantize_q8_1_DS4( + x: *const c_void, + ids: *const i32, + vy: *mut c_void, + type_x: i32, + ne00: i64, + s01: i64, + s02: i64, + s03: i64, + ne0: i64, + ne1: i64, + ne2: i64, + ne3: i64, + stream: *mut c_void, + ); + pub fn launch_mmq_quantize_q8_1_D2S6( + x: *const c_void, + ids: *const i32, + vy: *mut c_void, + type_x: i32, + ne00: i64, + s01: i64, + s02: i64, + s03: i64, + ne0: i64, + ne1: i64, + ne2: i64, + ne3: i64, + stream: *mut c_void, + ); + + // MMQ matmul launchers (one per quant type) + pub fn launch_mmq_gguf_q4_0( + tmp_fixup: *mut c_void, + x: *const c_void, + y: *const c_void, + dst: *mut c_void, + ncols_x: i64, + nrows_x: i64, + ncols_y: i64, + stride_row_x: i64, + stride_col_dst: i64, + cc: i32, + nsm: i32, + smpbo: i64, + warp_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmq_gguf_q4_1( + tmp_fixup: *mut c_void, + x: *const c_void, + y: *const c_void, + dst: *mut c_void, + ncols_x: i64, + nrows_x: i64, + ncols_y: i64, + stride_row_x: i64, + stride_col_dst: i64, + cc: i32, + nsm: i32, + smpbo: i64, + warp_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmq_gguf_q5_0( + tmp_fixup: *mut c_void, + x: *const c_void, + y: *const c_void, + dst: *mut c_void, + ncols_x: i64, + nrows_x: i64, + ncols_y: i64, + stride_row_x: i64, + stride_col_dst: i64, + cc: i32, + nsm: i32, + smpbo: i64, + warp_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmq_gguf_q5_1( + tmp_fixup: *mut c_void, + x: *const c_void, + y: *const c_void, + dst: *mut c_void, + ncols_x: i64, + nrows_x: i64, + ncols_y: i64, + stride_row_x: i64, + stride_col_dst: i64, + cc: i32, + nsm: i32, + smpbo: i64, + warp_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmq_gguf_q8_0( + tmp_fixup: *mut c_void, + x: *const c_void, + y: *const c_void, + dst: *mut c_void, + ncols_x: i64, + nrows_x: i64, + ncols_y: i64, + stride_row_x: i64, + stride_col_dst: i64, + cc: i32, + nsm: i32, + smpbo: i64, + warp_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmq_gguf_q2_k( + tmp_fixup: *mut c_void, + x: *const c_void, + y: *const c_void, + dst: *mut c_void, + ncols_x: i64, + nrows_x: i64, + ncols_y: i64, + stride_row_x: i64, + stride_col_dst: i64, + cc: i32, + nsm: i32, + smpbo: i64, + warp_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmq_gguf_q3_k( + tmp_fixup: *mut c_void, + x: *const c_void, + y: *const c_void, + dst: *mut c_void, + ncols_x: i64, + nrows_x: i64, + ncols_y: i64, + stride_row_x: i64, + stride_col_dst: i64, + cc: i32, + nsm: i32, + smpbo: i64, + warp_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmq_gguf_q4_k( + tmp_fixup: *mut c_void, + x: *const c_void, + y: *const c_void, + dst: *mut c_void, + ncols_x: i64, + nrows_x: i64, + ncols_y: i64, + stride_row_x: i64, + stride_col_dst: i64, + cc: i32, + nsm: i32, + smpbo: i64, + warp_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmq_gguf_q5_k( + tmp_fixup: *mut c_void, + x: *const c_void, + y: *const c_void, + dst: *mut c_void, + ncols_x: i64, + nrows_x: i64, + ncols_y: i64, + stride_row_x: i64, + stride_col_dst: i64, + cc: i32, + nsm: i32, + smpbo: i64, + warp_size: i32, + stream: *mut c_void, + ); + pub fn launch_mmq_gguf_q6_k( + tmp_fixup: *mut c_void, + x: *const c_void, + y: *const c_void, + dst: *mut c_void, + ncols_x: i64, + nrows_x: i64, + ncols_y: i64, + stride_row_x: i64, + stride_col_dst: i64, + cc: i32, + nsm: i32, + smpbo: i64, + warp_size: i32, + stream: *mut c_void, + ); +} diff --git a/torchnative/rust/vendor/candle-kernels/src/fill.cu b/torchnative/rust/vendor/candle-kernels/src/fill.cu new file mode 100644 index 00000000..5e2d7ffc --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/fill.cu @@ -0,0 +1,87 @@ +#include +#include "cuda_fp16.h" +#include "cuda_utils.cuh" + +template +__device__ void fill_with(T *buf, T value, const size_t numel) { + for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { + buf[i] = value; + } +} +extern "C" __global__ void fill_u8(uint8_t *buf, uint8_t value, const size_t numel) { fill_with(buf, value, numel); } +extern "C" __global__ void fill_u32(uint32_t *buf, uint32_t value, const size_t numel) { fill_with(buf, value, numel); } +extern "C" __global__ void fill_i64(int64_t *buf, int64_t value, const size_t numel) { fill_with(buf, value, numel); } +extern "C" __global__ void fill_f32(float *buf, float value, const size_t numel) { fill_with(buf, value, numel); } +extern "C" __global__ void fill_f64(double *buf, double value, const size_t numel) { fill_with(buf, value, numel); } + +template +__device__ void copy2d(const T *src, T *dst, uint32_t d1, uint32_t d2, uint32_t src_s, uint32_t dst_s) { + uint32_t idx = blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= d1 * d2) { + return; + } + uint32_t idx1 = idx / d2; + uint32_t idx2 = idx - d2 * idx1; + dst[idx1 * dst_s + idx2] = src[idx1 * src_s + idx2]; +} + +#define COPY2D_OP(TYPENAME, FNNAME) \ +extern "C" __global__ \ +void FNNAME(const TYPENAME *src, TYPENAME *dst, uint32_t d1, uint32_t d2, uint32_t src_s, uint32_t dst_s) { \ + copy2d(src, dst, d1, d2, src_s, dst_s); \ +} \ + +COPY2D_OP(float, copy2d_f32) +COPY2D_OP(double, copy2d_f64) +COPY2D_OP(uint8_t, copy2d_u8) +COPY2D_OP(uint32_t, copy2d_u32) +COPY2D_OP(int64_t, copy2d_i64) + +#define CONST_SET_OP(TYPENAME, FN_NAME) \ +extern "C" __global__ void FN_NAME( \ + const size_t numel, \ + const size_t num_dims, \ + const size_t *info, \ + const TYPENAME inp, \ + TYPENAME *out \ +) { \ + const size_t *dims = info; \ + const size_t *strides = info + num_dims; \ + if (info == nullptr || is_contiguous(num_dims, dims, strides)) { \ + for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { \ + out[i] = inp; \ + } \ + } \ + else { \ + for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { \ + unsigned strided_i = get_strided_index(i, num_dims, dims, strides); \ + out[strided_i] = inp; \ + } \ + } \ +} \ + +CONST_SET_OP(float, const_set_f32) +CONST_SET_OP(double, const_set_f64) +CONST_SET_OP(uint8_t, const_set_u8) +CONST_SET_OP(uint32_t, const_set_u32) +CONST_SET_OP(int64_t, const_set_i64) + + +#if __CUDA_ARCH__ >= 530 +extern "C" __global__ void fill_f16(__half *buf, __half value, const size_t numel) { fill_with(buf, value, numel); } +COPY2D_OP(__half, copy2d_f16) +CONST_SET_OP(__half, const_set_f16) +#endif + +#if __CUDA_ARCH__ >= 800 +#include +#include + +extern "C" __global__ void fill_bf16(__nv_bfloat16 *buf, __nv_bfloat16 value, const size_t numel) { fill_with(buf, value, numel); } +COPY2D_OP(__nv_bfloat16, copy2d_bf16) +CONST_SET_OP(__nv_bfloat16, const_set_bf16) + +extern "C" __global__ void fill_f8_e4m3(__nv_fp8_e4m3 *buf, __nv_fp8_e4m3 value, const size_t numel) { fill_with(buf, value, numel); } +COPY2D_OP(__nv_fp8_e4m3, copy2d_f8_e4m3) +CONST_SET_OP(__nv_fp8_e4m3, const_set_f8_e4m3) +#endif diff --git a/torchnative/rust/vendor/candle-kernels/src/indexing.cu b/torchnative/rust/vendor/candle-kernels/src/indexing.cu new file mode 100644 index 00000000..d0eb718c --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/indexing.cu @@ -0,0 +1,457 @@ +// WARNING: THIS IS ONLY VALID ASSUMING THAT inp IS CONTIGUOUS! +// TODO: proper error reporting when ids are larger than v_size. +#include "cuda_utils.cuh" +#include + +template +__host__ __device__ +constexpr T max_value(); + +template <> +__host__ __device__ +constexpr int64_t max_value() { + return 0x7FFFFFFFFFFFFFFFLL; +} + +template <> +__host__ __device__ +constexpr uint32_t max_value() { + return 0xFFFFFFFFu; +} + +template <> +__host__ __device__ +constexpr uint8_t max_value() { + return 0xFFu; +} + +template <> +__host__ __device__ +constexpr int32_t max_value() { + return 0x7FFFFFFF; +} + +template <> +__host__ __device__ +constexpr int16_t max_value() { + return 0x7FFF; +} + +template +__device__ void index_select( + const size_t numel, + const size_t num_dims, + const size_t *info, + const I *ids, + const T *inp, + T *out, + const size_t left_size, + const size_t src_dim_size, + const size_t ids_dim_size, + const size_t right_size +) { + const size_t *dims = info; + const size_t *strides = info + num_dims; + bool b = is_contiguous(num_dims, dims, strides); + for (unsigned int dst_i = blockIdx.x * blockDim.x + threadIdx.x; dst_i < numel; dst_i += blockDim.x * gridDim.x) { + unsigned int left_i = dst_i / (ids_dim_size * right_size); + unsigned int id_i = dst_i / right_size % ids_dim_size; + unsigned int right_i = dst_i % right_size; + if (ids[id_i] == max_value()) { + out[dst_i] = static_cast(0); + } else { + assert(ids[id_i] < src_dim_size); + unsigned int src_i = left_i * (src_dim_size * right_size) + ids[id_i] * right_size + right_i; + unsigned strided_i = b ? src_i : get_strided_index(src_i, num_dims, dims, strides); + out[dst_i] = inp[strided_i]; + } + } +} + +#define IS_OP(TYPENAME, INDEX_TYPENAME, FN_NAME) \ +extern "C" __global__ void FN_NAME( \ + const size_t numel, \ + const size_t num_dims, \ + const size_t *info, \ + const INDEX_TYPENAME *ids, \ + const TYPENAME *inp, \ + TYPENAME *out, \ + const size_t left_size, \ + const size_t src_dim_size, \ + const size_t ids_dim_size, \ + const size_t right_size \ +) { index_select(numel, num_dims, info, ids, inp, out, left_size, src_dim_size, ids_dim_size, right_size); } \ + +template +__device__ void gather( + const size_t numel, + const I *ids, + const T *inp, + T *out, + const size_t left_size, + const size_t src_dim_size, + const size_t ids_dim_size, + const size_t right_size +) { + for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { + size_t post = i % right_size; + const I idx = ids[i]; + if (ids[i] == max_value()) { + out[i] = static_cast(0); + } else { + assert(idx < src_dim_size); + size_t pre = i / (right_size * ids_dim_size); + size_t src_i = (pre * src_dim_size + idx) * right_size + post; + out[i] = inp[src_i]; + } + } +} + +#define GATHER_OP(TYPENAME, INDEX_TYPENAME, FN_NAME) \ +extern "C" __global__ void FN_NAME( \ + const size_t numel, \ + const INDEX_TYPENAME *ids, \ + const TYPENAME *inp, \ + TYPENAME *out, \ + const size_t left_size, \ + const size_t src_dim_size, \ + const size_t ids_dim_size, \ + const size_t right_size \ +) { gather(numel, ids, inp, out, left_size, src_dim_size, ids_dim_size, right_size); } \ + +template +__device__ void index_add( + const I *ids, + const size_t ids_dim_size, + const T *inp, + T *out, + const size_t left_size, + const size_t src_dim_size, + const size_t dst_dim_size, + const size_t right_size +) { + const size_t numel = left_size * right_size; + for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { + const size_t pre = i / right_size; + const size_t post = i % right_size; + for (unsigned int j = 0; j < ids_dim_size; ++j) { + const I idx = ids[j]; + const size_t src_i = (pre * ids_dim_size + j) * right_size + post; + if (idx < max_value()) { + assert(idx < dst_dim_size); + const size_t dst_i = (pre * dst_dim_size + idx) * right_size + post; + out[dst_i] += inp[src_i]; + } + } + } +} + +#if __CUDA_ARCH__ >= 890 +#define F8E4M3_TO_FLOAT(x) __half2float(__nv_cvt_fp8_to_halfraw(x.__x, __NV_E4M3)) + +template +__device__ void scatter_add_f8( + const I *ids, + const __nv_fp8_e4m3 *inp, + __nv_fp8_e4m3 *out, + const size_t left_size, + const size_t src_dim_size, + const size_t dst_dim_size, + const size_t right_size +) { + const size_t numel = left_size * right_size; + for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { + const size_t pre = i / right_size; + const size_t post = i % right_size; + for (unsigned int j = 0; j < src_dim_size; ++j) { + const size_t src_i = (pre * src_dim_size + j) * right_size + post; + const size_t idx = ids[src_i]; + const size_t dst_i = (pre * dst_dim_size + idx) * right_size + post; + out[dst_i] = __nv_fp8_e4m3(F8E4M3_TO_FLOAT(out[dst_i]) + F8E4M3_TO_FLOAT(inp[src_i])); + } + } +} + +template +__device__ void index_add_f8( + const I *ids, + const size_t ids_dim_size, + const __nv_fp8_e4m3 *inp, + __nv_fp8_e4m3 *out, + const size_t left_size, + const size_t src_dim_size, + const size_t dst_dim_size, + const size_t right_size +) { + const size_t numel = left_size * right_size; + for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { + const size_t pre = i / right_size; + const size_t post = i % right_size; + for (unsigned int j = 0; j < ids_dim_size; ++j) { + const size_t idx = ids[j]; + const size_t src_i = (pre * ids_dim_size + j) * right_size + post; + const size_t dst_i = (pre * dst_dim_size + idx) * right_size + post; + out[dst_i] = __nv_fp8_e4m3(F8E4M3_TO_FLOAT(out[dst_i]) + F8E4M3_TO_FLOAT(inp[src_i])); + } + } +} +#endif + +#define IA_OP(TYPENAME, INDEX_TYPENAME, FN_NAME) \ +extern "C" __global__ void FN_NAME( \ + const INDEX_TYPENAME *ids, \ + const size_t ids_dim_size, \ + const TYPENAME *inp, \ + TYPENAME *out, \ + const size_t left_size, \ + const size_t src_dim_size, \ + const size_t dst_dim_size, \ + const size_t right_size \ +) { index_add(ids, ids_dim_size, inp, out, left_size, src_dim_size, dst_dim_size, right_size); } \ + +#define IA_OP_F8(TYPENAME, INDEX_TYPENAME, FN_NAME) \ +extern "C" __global__ void FN_NAME( \ + const INDEX_TYPENAME *ids, \ + const size_t ids_dim_size, \ + const TYPENAME *inp, \ + TYPENAME *out, \ + const size_t left_size, \ + const size_t src_dim_size, \ + const size_t dst_dim_size, \ + const size_t right_size \ +) { index_add_f8(ids, ids_dim_size, inp, out, left_size, src_dim_size, dst_dim_size, right_size); } \ + +template +__device__ void scatter( + const I *ids, + const T *inp, + T *out, + const size_t left_size, + const size_t src_dim_size, + const size_t dst_dim_size, + const size_t right_size +) { + const size_t numel = left_size * right_size; + for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { + const size_t pre = i / right_size; + const size_t post = i % right_size; + for (unsigned int j = 0; j < src_dim_size; ++j) { + const size_t src_i = (pre * src_dim_size + j) * right_size + post; + const I idx = ids[src_i]; + if (idx < max_value()) { + assert(idx < dst_dim_size); + const size_t dst_i = (pre * dst_dim_size + idx) * right_size + post; + out[dst_i] = inp[src_i]; + } + } + } +} + +template +__device__ void scatter_add( + const I *ids, + const T *inp, + T *out, + const size_t left_size, + const size_t src_dim_size, + const size_t dst_dim_size, + const size_t right_size +) { + const size_t numel = left_size * right_size; + for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { + const size_t pre = i / right_size; + const size_t post = i % right_size; + for (unsigned int j = 0; j < src_dim_size; ++j) { + const size_t src_i = (pre * src_dim_size + j) * right_size + post; + const I idx = ids[src_i]; + if (idx < max_value()) { + assert(idx < dst_dim_size); + const size_t dst_i = (pre * dst_dim_size + idx) * right_size + post; + out[dst_i] += inp[src_i]; + } + } + } +} + +#define S_OP(TYPENAME, INDEX_TYPENAME, FN_NAME) \ +extern "C" __global__ void FN_NAME( \ + const INDEX_TYPENAME *ids, \ + const TYPENAME *inp, \ + TYPENAME *out, \ + const size_t left_size, \ + const size_t src_dim_size, \ + const size_t dst_dim_size, \ + const size_t right_size \ +) { scatter(ids, inp, out, left_size, src_dim_size, dst_dim_size, right_size); } \ + +#define SA_OP(TYPENAME, INDEX_TYPENAME, FN_NAME) \ +extern "C" __global__ void FN_NAME( \ + const INDEX_TYPENAME *ids, \ + const TYPENAME *inp, \ + TYPENAME *out, \ + const size_t left_size, \ + const size_t src_dim_size, \ + const size_t dst_dim_size, \ + const size_t right_size \ +) { scatter_add(ids, inp, out, left_size, src_dim_size, dst_dim_size, right_size); } \ + +#define SA_OP_F8(TYPENAME, INDEX_TYPENAME, FN_NAME) \ +extern "C" __global__ void FN_NAME( \ + const INDEX_TYPENAME *ids, \ + const TYPENAME *inp, \ + TYPENAME *out, \ + const size_t left_size, \ + const size_t src_dim_size, \ + const size_t dst_dim_size, \ + const size_t right_size \ +) { scatter_add_f8(ids, inp, out, left_size, src_dim_size, dst_dim_size, right_size); } \ + + +#if __CUDA_ARCH__ >= 800 +IS_OP(__nv_bfloat16, int64_t, is_i64_bf16) +IS_OP(__nv_bfloat16, uint32_t, is_u32_bf16) +IS_OP(__nv_bfloat16, uint8_t, is_u8_bf16) +GATHER_OP(__nv_bfloat16, int64_t, gather_i64_bf16) +GATHER_OP(__nv_bfloat16, uint32_t, gather_u32_bf16) +GATHER_OP(__nv_bfloat16, uint8_t, gather_u8_bf16) +IA_OP(__nv_bfloat16, int64_t, ia_i64_bf16) +IA_OP(__nv_bfloat16, uint32_t, ia_u32_bf16) +IA_OP(__nv_bfloat16, uint8_t, ia_u8_bf16) +SA_OP(__nv_bfloat16, int64_t, sa_i64_bf16) +SA_OP(__nv_bfloat16, uint32_t, sa_u32_bf16) +SA_OP(__nv_bfloat16, uint8_t, sa_u8_bf16) +S_OP(__nv_bfloat16, int64_t, s_i64_bf16) +S_OP(__nv_bfloat16, uint32_t, s_u32_bf16) +S_OP(__nv_bfloat16, uint8_t, s_u8_bf16) +#endif + +#if __CUDA_ARCH__ >= 890 +IS_OP(__nv_fp8_e4m3, int16_t, is_i16_f8_e4m3) +IS_OP(__nv_fp8_e4m3, int32_t, is_i32_f8_e4m3) +IS_OP(__nv_fp8_e4m3, int64_t, is_i64_f8_e4m3) +IS_OP(__nv_fp8_e4m3, uint32_t, is_u32_f8_e4m3) +IS_OP(__nv_fp8_e4m3, uint8_t, is_u8_f8_e4m3) +GATHER_OP(__nv_fp8_e4m3, int16_t, gather_i16_f8_e4m3) +GATHER_OP(__nv_fp8_e4m3, int32_t, gather_i32_f8_e4m3) +GATHER_OP(__nv_fp8_e4m3, int64_t, gather_i64_f8_e4m3) +GATHER_OP(__nv_fp8_e4m3, uint32_t, gather_u32_f8_e4m3) +GATHER_OP(__nv_fp8_e4m3, uint8_t, gather_u8_f8_e4m3) +IA_OP_F8(__nv_fp8_e4m3, int16_t, ia_i16_f8_e4m3) +IA_OP_F8(__nv_fp8_e4m3, int32_t, ia_i32_f8_e4m3) +IA_OP_F8(__nv_fp8_e4m3, int64_t, ia_i64_f8_e4m3) +IA_OP_F8(__nv_fp8_e4m3, uint32_t, ia_u32_f8_e4m3) +IA_OP_F8(__nv_fp8_e4m3, uint8_t, ia_u8_f8_e4m3) +SA_OP_F8(__nv_fp8_e4m3, int16_t, sa_i16_f8_e4m3) +SA_OP_F8(__nv_fp8_e4m3, int32_t, sa_i32_f8_e4m3) +SA_OP_F8(__nv_fp8_e4m3, int64_t, sa_i64_f8_e4m3) +SA_OP_F8(__nv_fp8_e4m3, uint32_t, sa_u32_f8_e4m3) +SA_OP_F8(__nv_fp8_e4m3, uint8_t, sa_u8_f8_e4m3) +#endif + +#if __CUDA_ARCH__ >= 530 +IS_OP(__half, int64_t, is_i64_f16) +IS_OP(__half, uint32_t, is_u32_f16) +IS_OP(__half, uint8_t, is_u8_f16) +GATHER_OP(__half, int64_t, gather_i64_f16) +GATHER_OP(__half, uint32_t, gather_u32_f16) +GATHER_OP(__half, uint8_t, gather_u8_f16) +IA_OP(__half, int64_t, ia_i64_f16) +IA_OP(__half, uint32_t, ia_u32_f16) +IA_OP(__half, uint8_t, ia_u8_f16) +SA_OP(__half, int64_t, sa_i64_f16) +SA_OP(__half, uint32_t, sa_u32_f16) +SA_OP(__half, uint8_t, sa_u8_f16) +S_OP(__half, int64_t, s_i64_f16) +S_OP(__half, uint32_t, s_u32_f16) +S_OP(__half, uint8_t, s_u8_f16) +#endif + +IS_OP(float, int64_t, is_i64_f32) +IS_OP(double, int64_t, is_i64_f64) +IS_OP(uint8_t, int64_t, is_i64_u8) +IS_OP(uint32_t, int64_t, is_i64_u32) +IS_OP(int64_t, int64_t, is_i64_i64) + +IS_OP(float, uint32_t, is_u32_f32) +IS_OP(double, uint32_t, is_u32_f64) +IS_OP(uint8_t, uint32_t, is_u32_u8) +IS_OP(int64_t, uint32_t, is_u32_i64) +IS_OP(uint32_t, uint32_t, is_u32_u32) + +IS_OP(float, uint8_t, is_u8_f32) +IS_OP(double, uint8_t, is_u8_f64) +IS_OP(uint8_t, uint8_t, is_u8_u8) +IS_OP(uint32_t, uint8_t, is_u8_u32) +IS_OP(int64_t, uint8_t, is_u8_i64) + +GATHER_OP(float, int64_t, gather_i64_f32) +GATHER_OP(double, int64_t, gather_i64_f64) +GATHER_OP(uint8_t, int64_t, gather_i64_u8) +GATHER_OP(uint32_t, int64_t, gather_i64_u32) +GATHER_OP(int64_t, int64_t, gather_i64_i64) + +GATHER_OP(float, uint32_t, gather_u32_f32) +GATHER_OP(double, uint32_t, gather_u32_f64) +GATHER_OP(uint8_t, uint32_t, gather_u32_u8) +GATHER_OP(int64_t, uint32_t, gather_u32_i64) +GATHER_OP(uint32_t, uint32_t, gather_u32_u32) + +GATHER_OP(float, uint8_t, gather_u8_f32) +GATHER_OP(double, uint8_t, gather_u8_f64) +GATHER_OP(uint8_t, uint8_t, gather_u8_u8) +GATHER_OP(uint32_t, uint8_t, gather_u8_u32) +GATHER_OP(int64_t, uint8_t, gather_u8_i64) + +IA_OP(float, int64_t, ia_i64_f32) +IA_OP(double, int64_t, ia_i64_f64) +IA_OP(uint8_t, int64_t, ia_i64_u8) +IA_OP(int64_t, int64_t, ia_i64_i64) +IA_OP(uint32_t, int64_t, ia_i64_u32) + +IA_OP(float, uint32_t, ia_u32_f32) +IA_OP(double, uint32_t, ia_u32_f64) +IA_OP(uint8_t, uint32_t, ia_u32_u8) +IA_OP(int64_t, uint32_t, ia_u32_i64) +IA_OP(uint32_t, uint32_t, ia_u32_u32) + +IA_OP(float, uint8_t, ia_u8_f32) +IA_OP(double, uint8_t, ia_u8_f64) +IA_OP(uint8_t, uint8_t, ia_u8_u8) +IA_OP(uint32_t, uint8_t, ia_u8_u32) +IA_OP(int64_t, uint8_t, ia_u8_i64) + +SA_OP(float, int64_t, sa_i64_f32) +SA_OP(double, int64_t, sa_i64_f64) +SA_OP(uint8_t, int64_t, sa_i64_u8) +SA_OP(int64_t, int64_t, sa_i64_i64) +SA_OP(uint32_t, int64_t, sa_i64_u32) + +SA_OP(float, uint32_t, sa_u32_f32) +SA_OP(double, uint32_t, sa_u32_f64) +SA_OP(uint8_t, uint32_t, sa_u32_u8) +SA_OP(int64_t, uint32_t, sa_u32_i64) +SA_OP(uint32_t, uint32_t, sa_u32_u32) + +SA_OP(float, uint8_t, sa_u8_f32) +SA_OP(double, uint8_t, sa_u8_f64) +SA_OP(uint8_t, uint8_t, sa_u8_u8) +SA_OP(uint32_t, uint8_t, sa_u8_u32) +SA_OP(int64_t, uint8_t, sa_u8_i64) + +S_OP(float, int64_t, s_i64_f32) +S_OP(double, int64_t, s_i64_f64) +S_OP(uint8_t, int64_t, s_i64_u8) +S_OP(int64_t, int64_t, s_i64_i64) +S_OP(uint32_t, int64_t, s_i64_u32) + +S_OP(float, uint32_t, s_u32_f32) +S_OP(double, uint32_t, s_u32_f64) +S_OP(uint8_t, uint32_t, s_u32_u8) +S_OP(int64_t, uint32_t, s_u32_i64) +S_OP(uint32_t, uint32_t, s_u32_u32) + +S_OP(float, uint8_t, s_u8_f32) +S_OP(double, uint8_t, s_u8_f64) +S_OP(uint8_t, uint8_t, s_u8_u8) +S_OP(uint32_t, uint8_t, s_u8_u32) +S_OP(int64_t, uint8_t, s_u8_i64) diff --git a/torchnative/rust/vendor/candle-kernels/src/lib.rs b/torchnative/rust/vendor/candle-kernels/src/lib.rs new file mode 100644 index 00000000..cfc57326 --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/lib.rs @@ -0,0 +1,82 @@ +mod ptx { + include!(concat!(env!("OUT_DIR"), "/ptx.rs")); +} + +#[repr(u32)] +#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)] +pub enum Id { + Affine, + Binary, + Cast, + Conv, + Fill, + Indexing, + Quantized, + Reduce, + Sort, + Ternary, + Unary, +} + +pub const ALL_IDS: [Id; 11] = [ + Id::Affine, + Id::Binary, + Id::Cast, + Id::Conv, + Id::Fill, + Id::Indexing, + Id::Quantized, + Id::Reduce, + Id::Sort, + Id::Ternary, + Id::Unary, +]; + +pub struct Module { + index: usize, + ptx: &'static str, +} + +impl Module { + pub fn index(&self) -> usize { + self.index + } + + pub fn ptx(&self) -> &'static str { + self.ptx + } +} + +const fn module_index(id: Id) -> usize { + let mut i = 0; + while i < ALL_IDS.len() { + if ALL_IDS[i] as u32 == id as u32 { + return i; + } + i += 1; + } + panic!("id not found") +} + +macro_rules! mdl { + ($cst:ident, $id:ident) => { + pub const $cst: Module = Module { + index: module_index(Id::$id), + ptx: ptx::$cst, + }; + }; +} + +mdl!(AFFINE, Affine); +mdl!(BINARY, Binary); +mdl!(CAST, Cast); +mdl!(CONV, Conv); +mdl!(FILL, Fill); +mdl!(INDEXING, Indexing); +mdl!(QUANTIZED, Quantized); +mdl!(REDUCE, Reduce); +mdl!(SORT, Sort); +mdl!(TERNARY, Ternary); +mdl!(UNARY, Unary); + +pub mod ffi; diff --git a/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_common.cuh b/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_common.cuh new file mode 100644 index 00000000..bb4fc864 --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_common.cuh @@ -0,0 +1,534 @@ +// Self-contained common header for MMQ GGUF kernels. +// Replaces llama.cpp's common.cuh + ggml.h + ggml-common.h for standalone compilation. +#pragma once + +#include +#include +#include + +#include "cuda_fp16.h" +#include "cuda_bf16.h" + +// ============================================================ +// Basic macros +// ============================================================ + +#define WARP_SIZE 32 +#define MATRIX_ROW_PADDING 512 +#define GGML_PAD(x, n) (((x) + (n) - 1) & ~((n) - 1)) +#define GGML_UNUSED(x) (void)(x) +#define GGML_CUDA_MAX_DEVICES 16 + +#define STRINGIZE_IMPL(...) #__VA_ARGS__ +#define STRINGIZE(...) STRINGIZE_IMPL(__VA_ARGS__) + +// ============================================================ +// ggml_type enum (matching llama.cpp values) +// ============================================================ + +enum ggml_type { + GGML_TYPE_F32 = 0, + GGML_TYPE_F16 = 1, + GGML_TYPE_Q4_0 = 2, + GGML_TYPE_Q4_1 = 3, + 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_IQ1_M = 29, + GGML_TYPE_BF16 = 30, + GGML_TYPE_MXFP4 = 39, + GGML_TYPE_NVFP4 = 40, +}; + +// ============================================================ +// Quantization constants +// ============================================================ + +#define QK_K 256 +#define K_SCALE_SIZE 12 + +#define QK4_0 32 +#define QR4_0 2 +#define QI4_0 (QK4_0 / (4 * QR4_0)) + +#define QK4_1 32 +#define QR4_1 2 +#define QI4_1 (QK4_1 / (4 * QR4_1)) + +#define QK_MXFP4 32 +#define QR_MXFP4 2 +#define QI_MXFP4 (QK_MXFP4 / (4 * QR_MXFP4)) + +#define QK_NVFP4 64 +#define QK_NVFP4_SUB 16 +#define QR_NVFP4 2 +#define QI_NVFP4 (QK_NVFP4 / (4 * QR_NVFP4)) + +#define QK5_0 32 +#define QR5_0 2 +#define QI5_0 (QK5_0 / (4 * QR5_0)) + +#define QK5_1 32 +#define QR5_1 2 +#define QI5_1 (QK5_1 / (4 * QR5_1)) + +#define QK8_0 32 +#define QR8_0 1 +#define QI8_0 (QK8_0 / (4 * QR8_0)) + +#define QK8_1 32 +#define QR8_1 1 +#define QI8_1 (QK8_1 / (4 * QR8_1)) + +#define QR2_K 4 +#define QI2_K (QK_K / (4 * QR2_K)) + +#define QR3_K 4 +#define QI3_K (QK_K / (4 * QR3_K)) + +#define QR4_K 2 +#define QI4_K (QK_K / (4 * QR4_K)) + +#define QR5_K 2 +#define QI5_K (QK_K / (4 * QR5_K)) + +#define QR6_K 2 +#define QI6_K (QK_K / (4 * QR6_K)) + +// IQ constants (needed for template compilation even if not instantiated) +#define QR2_XXS 4 +#define QI2_XXS (QK_K / (4 * QR2_XXS)) +#define QR2_XS 4 +#define QI2_XS (QK_K / (4 * QR2_XS)) +#define QR2_S 4 +#define QI2_S (QK_K / (4 * QR2_S)) +#define QR3_XXS 4 +#define QI3_XXS (QK_K / (4 * QR3_XXS)) +#define QR3_S 4 +#define QI3_S (QK_K / (4 * QR3_S)) +#define QR1_S 8 +#define QI1_S (QK_K / (4 * QR1_S)) +#define QR1_M 8 +#define QI1_M (QK_K / (4 * QR1_M)) +#define QK4_NL 32 +#define QR4_NL 2 +#define QI4_NL (QK4_NL / (4 * QR4_NL)) +#define QR4_XS 2 +#define QI4_XS (QK_K / (4 * QR4_XS)) +#define QR3_XS 4 +#define QI3_XS (QK_K / (4 * QR3_XS)) + +// ============================================================ +// Block type definitions (CUDA half/half2) +// ============================================================ + +typedef struct { half d; uint8_t qs[QK4_0 / 2]; } block_q4_0; +typedef struct { half2 dm; uint8_t qs[QK4_1 / 2]; } block_q4_1; +typedef struct { uint8_t e; uint8_t qs[QK_MXFP4/2]; } block_mxfp4; +typedef struct { uint8_t d[QK_NVFP4/QK_NVFP4_SUB]; uint8_t qs[QK_NVFP4/2]; } block_nvfp4; +typedef struct { half d; uint8_t qh[4]; uint8_t qs[QK5_0 / 2]; } block_q5_0; +typedef struct { half2 dm; uint8_t qh[4]; uint8_t qs[QK5_1 / 2]; } block_q5_1; +typedef struct { half d; int8_t qs[QK8_0]; } block_q8_0; +typedef struct { half2 ds; int8_t qs[QK8_1]; } block_q8_1; + +typedef struct { + uint8_t scales[QK_K/16]; + uint8_t qs[QK_K/4]; + half2 dm; +} block_q2_K; + +typedef struct { + uint8_t hmask[QK_K/8]; + uint8_t qs[QK_K/4]; + uint8_t scales[12]; + half d; +} block_q3_K; + +typedef struct { + half2 dm; + uint8_t scales[K_SCALE_SIZE]; + uint8_t qs[QK_K/2]; +} block_q4_K; + +typedef struct { + half2 dm; + uint8_t scales[K_SCALE_SIZE]; + uint8_t qh[QK_K/8]; + uint8_t qs[QK_K/2]; +} block_q5_K; + +typedef struct { + uint8_t ql[QK_K/2]; + uint8_t qh[QK_K/4]; + int8_t scales[QK_K/16]; + half d; +} block_q6_K; + +// IQ block types (needed for template compilation) +typedef struct { half d; uint16_t qs[QK_K/8]; } block_iq2_xxs; +typedef struct { half d; uint16_t qs[QK_K/8]; uint8_t scales[QK_K/32]; } block_iq2_xs; +typedef struct { half d; uint8_t qs[QK_K/4]; uint8_t qh[QK_K/32]; uint8_t scales[QK_K/32]; } block_iq2_s; +typedef struct { half d; uint8_t qs[3*QK_K/8]; } block_iq3_xxs; +#define IQ3S_N_SCALE QK_K/64 +typedef struct { half d; uint8_t qs[QK_K/4]; uint8_t qh[QK_K/32]; uint8_t signs[QK_K/8]; uint8_t scales[IQ3S_N_SCALE]; } block_iq3_s; +typedef struct { half d; uint8_t qs[QK_K/8]; uint16_t qh[QK_K/32]; } block_iq1_s; +typedef struct { uint8_t qs[QK_K/8]; uint8_t qh[QK_K/16]; uint8_t scales[QK_K/32]; } block_iq1_m; +typedef struct { half d; uint8_t qs[QK4_NL/2]; } block_iq4_nl; +typedef struct { half d; uint16_t scales_h; uint8_t scales_l[QK_K/64]; uint8_t qs[QK_K/2]; } block_iq4_xs; + +// ============================================================ +// Architecture detection +// ============================================================ + +#define GGML_CUDA_CC_PASCAL 600 +#define GGML_CUDA_CC_DP4A 610 +#define GGML_CUDA_CC_VOLTA 700 +#define GGML_CUDA_CC_TURING 750 +#define GGML_CUDA_CC_AMPERE 800 +#define GGML_CUDA_CC_ADA_LOVELACE 890 +#define GGML_CUDA_CC_BLACKWELL 1200 +#define GGML_CUDA_CC_DGX_SPARK 1210 +#define GGML_CUDA_CC_RUBIN 1300 + +#define GGML_CUDA_CC_OFFSET_AMD 0x1000000 +#define GGML_CUDA_CC_OFFSET_MTHREADS 0x0100000 +#define GGML_CUDA_CC_IS_NVIDIA(cc) (cc < GGML_CUDA_CC_OFFSET_MTHREADS) +#define GGML_CUDA_CC_IS_AMD(cc) (cc >= GGML_CUDA_CC_OFFSET_AMD) + +// AMD CC constants (needed for compile-time checks even though we target NVIDIA) +#define GGML_CUDA_CC_CDNA1 (GGML_CUDA_CC_OFFSET_AMD + 0x908) +#define GGML_CUDA_CC_RDNA1 (GGML_CUDA_CC_OFFSET_AMD + 0x1010) +#define GGML_CUDA_CC_RDNA2 (GGML_CUDA_CC_OFFSET_AMD + 0x1030) +#define GGML_CUDA_CC_RDNA3 (GGML_CUDA_CC_OFFSET_AMD + 0x1100) +#define GGML_CUDA_CC_RDNA3_5 (GGML_CUDA_CC_OFFSET_AMD + 0x1150) +#define GGML_CUDA_CC_RDNA4 (GGML_CUDA_CC_OFFSET_AMD + 0x1200) +#define GGML_CUDA_CC_CDNA3 (GGML_CUDA_CC_OFFSET_AMD + 0x942) + +#define GGML_CUDA_CC_IS_RDNA(cc) (cc >= GGML_CUDA_CC_RDNA1) +#define GGML_CUDA_CC_IS_RDNA1(cc) (cc >= GGML_CUDA_CC_RDNA1 && cc < GGML_CUDA_CC_RDNA2) +#define GGML_CUDA_CC_IS_RDNA3_0(cc) (cc >= GGML_CUDA_CC_RDNA3 && cc < GGML_CUDA_CC_RDNA3_5) +#define GGML_CUDA_CC_IS_RDNA3_5(cc) (cc >= GGML_CUDA_CC_RDNA3_5 && cc < GGML_CUDA_CC_RDNA4) +#define GGML_CUDA_CC_IS_RDNA3(cc) (GGML_CUDA_CC_IS_RDNA3_0(cc) || GGML_CUDA_CC_IS_RDNA3_5(cc)) +#define GGML_CUDA_CC_IS_RDNA4(cc) (cc >= GGML_CUDA_CC_RDNA4) +#define GGML_CUDA_CC_IS_CDNA(cc) (cc >= GGML_CUDA_CC_CDNA1 && cc < GGML_CUDA_CC_RDNA1) +#define GGML_CUDA_CC_IS_CDNA3(cc) (cc >= GGML_CUDA_CC_CDNA3 && cc < GGML_CUDA_CC_RDNA1) + +// Compile-time architecture detection +#ifdef __CUDA_ARCH_LIST__ +constexpr bool ggml_cuda_has_arch_impl(int) { return false; } + +template +constexpr bool ggml_cuda_has_arch_impl(const int arch, const int first, Archs... rest) { + return arch == first || ggml_cuda_has_arch_impl(arch, rest...); +} + +constexpr bool ggml_cuda_has_arch(const int arch) { + return ggml_cuda_has_arch_impl(arch, __CUDA_ARCH_LIST__); +} + +constexpr int ggml_cuda_highest_compiled_arch_impl(const int /*arch*/, const int cur) { + if (cur == 0) return -1; + return cur; +} + +template +constexpr int ggml_cuda_highest_compiled_arch_impl(const int arch, const int cur, const int first, Archs... rest) { + if (first <= arch && first > cur) { + return ggml_cuda_highest_compiled_arch_impl(arch, first, rest...); + } else { + return ggml_cuda_highest_compiled_arch_impl(arch, cur, rest...); + } +} + +constexpr int ggml_cuda_highest_compiled_arch(const int arch) { + return ggml_cuda_highest_compiled_arch_impl(arch, 0, __CUDA_ARCH_LIST__); +} +#else +static int ggml_cuda_highest_compiled_arch(const int arch) { + return arch; +} +#endif // __CUDA_ARCH_LIST__ + +// FP16 availability +#if __CUDA_ARCH__ >= GGML_CUDA_CC_PASCAL +#define FP16_AVAILABLE +#endif + +#if defined(FP16_AVAILABLE) && __CUDA_ARCH__ != 610 +#define FAST_FP16_AVAILABLE +#endif + +// MMA (tensor core) availability +#if __CUDA_ARCH__ == GGML_CUDA_CC_VOLTA +#define VOLTA_MMA_AVAILABLE +#endif + +#if __CUDA_ARCH__ >= GGML_CUDA_CC_TURING +#define TURING_MMA_AVAILABLE +#endif + +#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE +#define AMPERE_MMA_AVAILABLE +#endif + +#if __CUDA_ARCH__ >= GGML_CUDA_CC_BLACKWELL && __CUDA_ARCH__ < GGML_CUDA_CC_RUBIN +#define BLACKWELL_MMA_AVAILABLE +#endif + +#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE +#define CP_ASYNC_AVAILABLE +#endif + +#if defined(TURING_MMA_AVAILABLE) +#define LDMATRIX_TRANS_AVAILABLE +#endif + +// Host-side architecture query functions +static bool fp16_mma_hardware_available(const int cc) { + return GGML_CUDA_CC_IS_NVIDIA(cc) && cc >= GGML_CUDA_CC_VOLTA; +} + +static bool amd_mfma_available(const int /*cc*/) { return false; } // NVIDIA only +static bool amd_wmma_available(const int /*cc*/) { return false; } // NVIDIA only + +static bool turing_mma_available(const int cc) { + return GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_TURING; +} + +static bool blackwell_mma_available(const int cc) { + return GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_BLACKWELL && + ggml_cuda_highest_compiled_arch(cc) < GGML_CUDA_CC_RUBIN; +} + +// ============================================================ +// Device helpers +// ============================================================ + +static constexpr __device__ int ggml_cuda_get_physical_warp_size() { + return 32; // NVIDIA only +} + +// NO_DEVICE_CODE: called from unused template paths to satisfy compiler +[[noreturn]] +static __device__ void no_device_code( + const char * file_name, const int line, const char * function_name, const int arch, const char * arch_list) { + printf("%s:%d: ERROR: CUDA kernel %s has no device code for arch %d. Compiled for: %s\n", + file_name, line, function_name, arch, arch_list); + __trap(); + GGML_UNUSED(no_device_code); +} + +#ifdef __CUDA_ARCH__ +#define NO_DEVICE_CODE no_device_code(__FILE__, __LINE__, __FUNCTION__, __CUDA_ARCH__, STRINGIZE(__CUDA_ARCH_LIST__)) +#else +#define NO_DEVICE_CODE +#endif + +#ifdef __CUDA_ARCH__ +#define GGML_ABORT(msg) do { printf("GGML_ABORT: %s\n", msg); __trap(); } while(0) +#define GGML_ASSERT(x) do { if (!(x)) { printf("GGML_ASSERT failed: %s\n", #x); __trap(); } } while(0) +#else +#define GGML_ABORT(msg) do { fprintf(stderr, "GGML_ABORT: %s\n", msg); abort(); } while(0) +#define GGML_ASSERT(x) do { if (!(x)) { fprintf(stderr, "GGML_ASSERT failed: %s\n", #x); abort(); } } while(0) +#endif + +// dp4a intrinsic +static __device__ __forceinline__ int ggml_cuda_dp4a(const int a, const int b, int c) { +#if __CUDA_ARCH__ >= GGML_CUDA_CC_DP4A + return __dp4a(a, b, c); +#else + const int8_t * a8 = (const int8_t *) &a; + const int8_t * b8 = (const int8_t *) &b; + return c + a8[0]*b8[0] + a8[1]*b8[1] + a8[2]*b8[2] + a8[3]*b8[3]; +#endif +} + +// Warp reductions +template +static __device__ __forceinline__ int warp_reduce_sum(int x) { +#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE + return __reduce_add_sync(0xffffffff, x); +#else +#pragma unroll + for (int offset = width/2; offset > 0; offset >>= 1) { + x += __shfl_xor_sync(0xffffffff, x, offset, width); + } + return x; +#endif +} + +template +static __device__ __forceinline__ float warp_reduce_sum(float x) { +#pragma unroll + for (int offset = width/2; offset > 0; offset >>= 1) { + x += __shfl_xor_sync(0xffffffff, x, offset, width); + } + return x; +} + +template +static __device__ __forceinline__ int warp_reduce_max(int x) { +#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE + return __reduce_max_sync(0xffffffff, x); +#else +#pragma unroll + for (int offset = width/2; offset > 0; offset >>= 1) { + x = max(x, __shfl_xor_sync(0xffffffff, x, offset, width)); + } + return x; +#endif +} + +// CUDA_SET_SHARED_MEMORY_LIMIT +#define CUDA_SET_SHARED_MEMORY_LIMIT(kernel, nbytes) \ + do { \ + static bool raised[GGML_CUDA_MAX_DEVICES] = {false}; \ + int dev; cudaGetDevice(&dev); \ + if (!raised[dev]) { \ + cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, nbytes); \ + raised[dev] = true; \ + } \ + } while(0) + +// ============================================================ +// ggml_cuda_type_traits +// ============================================================ + +template +struct ggml_cuda_type_traits; + +template<> struct ggml_cuda_type_traits { static constexpr int qk = 1; static constexpr int qr = 1; }; +template<> struct ggml_cuda_type_traits { static constexpr int qk = QK4_0; static constexpr int qr = QR4_0; static constexpr int qi = QI4_0; }; +template<> struct ggml_cuda_type_traits { static constexpr int qk = QK4_1; static constexpr int qr = QR4_1; static constexpr int qi = QI4_1; }; +template<> struct ggml_cuda_type_traits { static constexpr int qk = QK5_0; static constexpr int qr = QR5_0; static constexpr int qi = QI5_0; }; +template<> struct ggml_cuda_type_traits { static constexpr int qk = QK5_1; static constexpr int qr = QR5_1; static constexpr int qi = QI5_1; }; +template<> struct ggml_cuda_type_traits { static constexpr int qk = QK8_0; static constexpr int qr = QR8_0; static constexpr int qi = QI8_0; }; +template<> struct ggml_cuda_type_traits { static constexpr int qk = QK_K; static constexpr int qr = QR2_K; static constexpr int qi = QI2_K; }; +template<> struct ggml_cuda_type_traits { static constexpr int qk = QK_K; static constexpr int qr = QR3_K; static constexpr int qi = QI3_K; }; +template<> struct ggml_cuda_type_traits { static constexpr int qk = QK_K; static constexpr int qr = QR4_K; static constexpr int qi = QI4_K; }; +template<> struct ggml_cuda_type_traits { static constexpr int qk = QK_K; static constexpr int qr = QR5_K; static constexpr int qi = QI5_K; }; +template<> struct ggml_cuda_type_traits { static constexpr int qk = QK_K; static constexpr int qr = QR6_K; static constexpr int qi = QI6_K; }; +template<> struct ggml_cuda_type_traits { static constexpr int qk = QK_MXFP4; static constexpr int qr = QR_MXFP4; static constexpr int qi = QI_MXFP4; }; +template<> struct ggml_cuda_type_traits { static constexpr int qk = QK_NVFP4; static constexpr int qr = QR_NVFP4; static constexpr int qi = QI_NVFP4; }; +template<> struct ggml_cuda_type_traits { static constexpr int qk = QK_K; static constexpr int qr = QR2_XXS; static constexpr int qi = QI2_XXS; }; +template<> struct ggml_cuda_type_traits { static constexpr int qk = QK_K; static constexpr int qr = QR2_XS; static constexpr int qi = QI2_XS; }; +template<> struct ggml_cuda_type_traits { static constexpr int qk = QK_K; static constexpr int qr = QR2_S; static constexpr int qi = QI2_S; }; +template<> struct ggml_cuda_type_traits { static constexpr int qk = QK_K; static constexpr int qr = QR3_XXS; static constexpr int qi = QI3_XXS; }; +template<> struct ggml_cuda_type_traits { static constexpr int qk = QK_K; static constexpr int qr = QR3_S; static constexpr int qi = QI3_S; }; +template<> struct ggml_cuda_type_traits { static constexpr int qk = QK_K; static constexpr int qr = QR1_S; static constexpr int qi = QI1_S; }; +template<> struct ggml_cuda_type_traits { static constexpr int qk = QK4_NL; static constexpr int qr = QR4_NL; static constexpr int qi = QI4_NL; }; +template<> struct ggml_cuda_type_traits { static constexpr int qk = QK_K; static constexpr int qr = QR4_XS; static constexpr int qi = QI4_XS; }; + +// ============================================================ +// Additional macros and helpers +// ============================================================ + +template +__host__ __device__ constexpr inline void ggml_unused_vars_impl(Args&&...) noexcept {} +#define GGML_UNUSED_VARS(...) ggml_unused_vars_impl(__VA_ARGS__) + +// Maximum number of bytes that can be copied in a single instruction. +static constexpr __device__ int ggml_cuda_get_max_cpy_bytes() { +#if __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA + return 16; +#else + return 8; +#endif +} + +// Device memcpy helper for register<->SRAM copies +template +static __device__ __forceinline__ void ggml_cuda_memcpy_1(void * __restrict__ dst, const void * __restrict__ src) { + static_assert( + nbytes <= ggml_cuda_get_max_cpy_bytes() || alignment == 0, + "Alignment misuse in ggml_cuda_memcpy_1"); + if constexpr (alignment != 0) { + static_assert(nbytes % alignment == 0, "bad alignment"); + } + constexpr int nb_per_cpy = alignment == 0 ? nbytes : alignment; +#pragma unroll + for (int i = 0; i < nbytes/nb_per_cpy; ++i) { + if constexpr (nb_per_cpy == 1) { + ((char *) dst)[i] = ((const char *) src)[i]; + } else if constexpr (nb_per_cpy == 2) { + ((short *) dst)[i] = ((const short *) src)[i]; + } else if constexpr (nb_per_cpy == 4) { + ((int *) dst)[i] = ((const int *) src)[i]; + } else if constexpr (nb_per_cpy == 8) { + ((int2 *) dst)[i] = ((const int2 *) src)[i]; + } else if constexpr (nb_per_cpy == 16) { + ((int4 *) dst)[i] = ((const int4 *) src)[i]; + } else { + static_assert(nbytes == 0 && nbytes == -1, "bad nbytes"); + } + } +} + +// E8M0/UE4M3 float conversion helpers (for MXFP4/NVFP4) +static __device__ __forceinline__ float ggml_cuda_e8m0_to_fp32(uint8_t x) { + uint32_t bits; + if (x == 0) { bits = 0x00400000; } else { bits = (uint32_t) x << 23; } + float result; + memcpy(&result, &bits, sizeof(float)); + return result; +} + +static __device__ __forceinline__ float ggml_cuda_ue4m3_to_fp32(uint8_t x) { + if (x == 0 || (x == 0x7F && x != 0xFF)) { return 0.0f; } + const int exp = (x >> 3) & 0xF; + const int man = x & 0x7; + float raw; + if (exp == 0) { raw = ldexpf((float) man, -9); } else { raw = ldexpf(1.0f + (float) man / 8.0f, exp - 7); } + return static_cast(raw / 2); +} + +// IQ/MXFP4 lookup table stubs (needed for compilation even though we only instantiate standard quant types) +// These are device constants from ggml-common.h. We provide minimal stubs. +// The functions referencing them are only called for IQ/MXFP4 types which we never instantiate. +static const __device__ int8_t kvalues_mxfp4[16] = {0, 2, 4, 6, 8, 10, 12, 14, 1, 3, 5, 7, 9, 11, 13, 15}; +static const __device__ int8_t kvalues_iq4nl[16] = {-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113}; +static const __device__ uint64_t iq2xxs_grid[256] = {0}; +static const __device__ uint64_t iq2xs_grid[512] = {0}; +static const __device__ uint64_t iq2s_grid[1024] = {0}; +static const __device__ uint32_t iq3xxs_grid[256] = {0}; +static const __device__ uint32_t iq3s_grid[512] = {0}; +static const __device__ uint32_t iq1s_grid_gpu[512] = {0}; +#define IQ1S_DELTA 0.125f +#define IQ1M_DELTA 0.125f +typedef union { half f16; uint16_t u16; } iq1m_scale_t; + +// ============================================================ +// ggml_cuda_unroll helper (used by some kernels) +// ============================================================ + +template +struct ggml_cuda_unroll { + template + __device__ void operator()(const Func & f, Args... args) const { + f(n - 1, args...); + ggml_cuda_unroll{}(f, args...); + } +}; + +template <> +struct ggml_cuda_unroll<1> { + template + __device__ void operator()(const Func & f, Args... args) const { + f(0, args...); + } +}; diff --git a/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_gguf.cuh b/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_gguf.cuh new file mode 100644 index 00000000..04c5f79b --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_gguf.cuh @@ -0,0 +1,3957 @@ +#pragma once + +#include "mmq_common.cuh" +#include "mmq_vecdotq.cuh" +#include "mmq_mma.cuh" + +#include +#include + +using namespace ggml_cuda_mma; + +#define MMQ_DP4A_MAX_BATCH_SIZE 64 // Max. batch size to use for dp4a MMQ kernels when FP16 tensor cores are available. +#define MMQ_ITER_K 256 +#define MMQ_ITER_K_MXFP4_FP4 512 +#define MMQ_NWARPS 8 + +typedef void (*load_tiles_mmq_t)(const char * __restrict__ x, int * x_tile, const int kbx0, const int i_max, const int stride); +typedef void (*vec_dot_mmq_t)(const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00); +typedef void (*mmq_write_back_t)(const float * __restrict__ sum, const int32_t * __restrict__ get_rows_to_sorted, + float * __restrict__ dst, const int stride, const int i_max, const int j_max); + +enum mmq_q8_1_ds_layout { + MMQ_Q8_1_DS_LAYOUT_D4, + MMQ_Q8_1_DS_LAYOUT_DS4, + MMQ_Q8_1_DS_LAYOUT_D2S6, +}; + +struct block_q8_1_mmq { + // The y float data is converted to a data layout that can simply be copied to shared memory as a contiguous block. + // The y float data is first grouped as blocks of 128 values. + // These blocks are then treated as individual data values and transposed. + // + // To avoid shared memory bank conflicts each block is padded with 16 bytes. + // This padding is also used to store block scales/partial sums. + // The scales multiplied with the quantized data are equal to the unquantized values. + // The partial sums are obtained by summing up a subgroup of the contained values (prior to quantization) + // and are only needed for performance reasons. + // + // The exact data stored depends on the x data type. + union { + float d4[4]; // 1 32 bit scale per 32 values, stored as d0,d1,d2,d3 + half2 ds4[4]; // 1 16 bit scale + 1 16 bit partial sum per 32 values, stored as d0,s0,d1,s1,d2,s2,d3,s3 + half d2s6[8]; // 1 16 bit scale per 64 values + 1 16 bit partial sum per 16 values for the first 96 values, + // stored as d0,d1,s1,s2,s3,s4,s5 + }; + int8_t qs[4*QK8_1]; // 128 values quantized to 8 bit each +}; + +struct block_fp4_mmq { + uint32_t d4[4]; // 8 E8M0 scales (1 per 32 values), 2 packed per uint32: d4[0]={s0,s1}, d4[1]={s2,s3}, etc. + int8_t qs[4 * 32]; // 256 FP4 values packed as 4-bit pairs (2 per byte), 8 blocks of 32 values +}; + +static_assert(sizeof(block_q8_1_mmq) == 4*QK8_1 + 4*sizeof(half2), "Unexpected block_q8_1_mmq size"); +static_assert(sizeof(block_q8_1_mmq) == 4*sizeof(block_q8_1), "Unexpected block_q8_1_mmq size"); +static_assert(sizeof(block_fp4_mmq) == sizeof(block_q8_1_mmq), "Unexpected block_fp4_mmq size"); + +static mmq_q8_1_ds_layout mmq_get_q8_1_ds_layout(const ggml_type type_x) { + switch (type_x) { + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + return MMQ_Q8_1_DS_LAYOUT_DS4; + case GGML_TYPE_Q5_0: + return MMQ_Q8_1_DS_LAYOUT_D4; + case GGML_TYPE_Q5_1: + return MMQ_Q8_1_DS_LAYOUT_DS4; + case GGML_TYPE_Q8_0: + return MMQ_Q8_1_DS_LAYOUT_D4; + case GGML_TYPE_MXFP4: + return MMQ_Q8_1_DS_LAYOUT_D4; + case GGML_TYPE_NVFP4: + return MMQ_Q8_1_DS_LAYOUT_D4; + case GGML_TYPE_Q2_K: + return MMQ_Q8_1_DS_LAYOUT_D2S6; + case GGML_TYPE_Q3_K: + return MMQ_Q8_1_DS_LAYOUT_D4; + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + return MMQ_Q8_1_DS_LAYOUT_DS4; + case GGML_TYPE_Q6_K: + case GGML_TYPE_IQ2_XXS: + case GGML_TYPE_IQ2_XS: + case GGML_TYPE_IQ2_S: + case GGML_TYPE_IQ3_XXS: + case GGML_TYPE_IQ3_S: + return MMQ_Q8_1_DS_LAYOUT_D4; + case GGML_TYPE_IQ1_S: + return MMQ_Q8_1_DS_LAYOUT_DS4; + case GGML_TYPE_IQ4_XS: + case GGML_TYPE_IQ4_NL: + return MMQ_Q8_1_DS_LAYOUT_D4; + default: + GGML_ABORT("fatal error"); + break; + } +} + +struct tile_x_sizes { + int qs; + int dm; + int sc; +}; + +static int get_mmq_x_max_host(const int cc) { + return (amd_mfma_available(cc) || turing_mma_available(cc) || amd_wmma_available(cc)) ? 128 : + GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA ? +#ifdef GGML_CUDA_FORCE_MMQ + 128 : 64; +#else + MMQ_DP4A_MAX_BATCH_SIZE : 64; +#endif // GGML_CUDA_FORCE_MMQ +} + +static constexpr __device__ int get_mmq_x_max_device() { +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + return 128; +#else // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) + +#if defined(GGML_USE_HIP) + return 64; +#else // defined(GGML_USE_HIP) + +#if __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA +#ifdef GGML_CUDA_FORCE_MMQ + return 128; +#else // GGML_CUDA_FORCE_MMQ + return MMQ_DP4A_MAX_BATCH_SIZE; +#endif // GGML_CUDA_FORCE_MMQ +#else // __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA + return 64; +#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA + +#endif // defined(GGML_USE_HIP) +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) +} + +static int get_mmq_y_host(const int cc) { + return GGML_CUDA_CC_IS_AMD(cc) ? (GGML_CUDA_CC_IS_RDNA1(cc) ? 64 : 128) : + ((GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 128 : 64); +} + +static constexpr __device__ int get_iter_k([[maybe_unused]] const ggml_type type) { +#if defined(BLACKWELL_MMA_AVAILABLE) + return type == GGML_TYPE_MXFP4 ? MMQ_ITER_K_MXFP4_FP4 : MMQ_ITER_K; +#else + return MMQ_ITER_K; +#endif // defined(BLACKWELL_MMA_AVAILABLE) +} + +static constexpr __device__ int get_mmq_y_device() { +#if defined(GGML_USE_HIP) +#if defined(RDNA1) + return 64; +#else + return 128; +#endif // defined RDNA1 +#else +#if __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA + return 128; +#else + return 64; +#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA +#endif // defined(GGML_USE_HIP) +} + +// Decouple shared memory tile sizes from WARP_SIZE to allow for different warp sizes. +// The K dimension of the tiles has either, +// 1*MMQ_TILE_NE_K==32 (always for TILE_Y_K) or 2*MMQ_TILE_NE_K==64 (typically for TILE_X_K), +// 32 bit elements for the quantized data (does not include scales). +// In other words, the size of the quantized data in the K dimension is a multiple of MMQ_TILE_NE_K. +// The final tile size in K direction is padded to avoid shared memory bank conflicts, +// in terms of 32 bit elements that means K % 2 == 1 for dp4a or K % 8 == 4 for mma. +#define MMQ_TILE_NE_K 32 + +#define MMQ_DP4A_TXS_Q4_0 tile_x_sizes{mmq_y*MMQ_TILE_NE_K + mmq_y, mmq_y*MMQ_TILE_NE_K/QI4_0 + mmq_y/QI4_0, 0} +#define MMQ_DP4A_TXS_Q4_1 tile_x_sizes{mmq_y*MMQ_TILE_NE_K + mmq_y, mmq_y*MMQ_TILE_NE_K/QI4_1 + mmq_y/QI4_1, 0} +#define MMQ_DP4A_TXS_Q8_0 tile_x_sizes{mmq_y*MMQ_TILE_NE_K*2 + mmq_y, mmq_y*MMQ_TILE_NE_K*2/QI8_0 + mmq_y/(QI8_0/2), 0} +#define MMQ_DP4A_TXS_Q8_0_16 tile_x_sizes{mmq_y*MMQ_TILE_NE_K*2 + mmq_y, mmq_y*MMQ_TILE_NE_K*4/QI8_0 + mmq_y/(QI8_0/4), 0} +#define MMQ_DP4A_TXS_Q8_1 tile_x_sizes{mmq_y*MMQ_TILE_NE_K*2 + mmq_y, mmq_y*MMQ_TILE_NE_K*2/QI8_1 + mmq_y/(QI8_1/2), 0} +#define MMQ_DP4A_TXS_Q2_K tile_x_sizes{mmq_y*MMQ_TILE_NE_K*2 + mmq_y, mmq_y*MMQ_TILE_NE_K + mmq_y, 0} +#define MMQ_DP4A_TXS_Q3_K tile_x_sizes{mmq_y*MMQ_TILE_NE_K*2 + mmq_y, mmq_y, mmq_y*MMQ_TILE_NE_K/8 + mmq_y/8} +#define MMQ_DP4A_TXS_Q4_K tile_x_sizes{mmq_y*MMQ_TILE_NE_K + mmq_y, mmq_y*MMQ_TILE_NE_K/QI4_K, mmq_y*MMQ_TILE_NE_K/8 + mmq_y/8} +#define MMQ_DP4A_TXS_Q5_K tile_x_sizes{mmq_y*MMQ_TILE_NE_K*2 + mmq_y, mmq_y*MMQ_TILE_NE_K/QI5_K + mmq_y/QI5_K, mmq_y*MMQ_TILE_NE_K/8 + mmq_y/8} +#define MMQ_DP4A_TXS_Q6_K tile_x_sizes{mmq_y*MMQ_TILE_NE_K*2 + mmq_y, mmq_y*MMQ_TILE_NE_K/QI6_K + mmq_y/QI6_K, mmq_y*MMQ_TILE_NE_K/8 + mmq_y/8} + +static constexpr __host__ __device__ tile_x_sizes mmq_get_dp4a_tile_x_sizes(ggml_type type, int mmq_y) { + switch (type) { + case GGML_TYPE_Q4_0: return MMQ_DP4A_TXS_Q4_0; + case GGML_TYPE_Q4_1: return MMQ_DP4A_TXS_Q4_1; + case GGML_TYPE_Q5_0: return MMQ_DP4A_TXS_Q8_0; + case GGML_TYPE_Q5_1: return MMQ_DP4A_TXS_Q8_1; + case GGML_TYPE_Q8_0: return MMQ_DP4A_TXS_Q8_0; + case GGML_TYPE_MXFP4: return MMQ_DP4A_TXS_Q8_1; + case GGML_TYPE_NVFP4: return MMQ_DP4A_TXS_Q8_0_16; + case GGML_TYPE_Q2_K: return MMQ_DP4A_TXS_Q2_K; + case GGML_TYPE_Q3_K: return MMQ_DP4A_TXS_Q3_K; + case GGML_TYPE_Q4_K: return MMQ_DP4A_TXS_Q4_K; + case GGML_TYPE_Q5_K: return MMQ_DP4A_TXS_Q5_K; + case GGML_TYPE_Q6_K: return MMQ_DP4A_TXS_Q6_K; + case GGML_TYPE_IQ2_XXS: return MMQ_DP4A_TXS_Q8_0; + case GGML_TYPE_IQ2_XS: return MMQ_DP4A_TXS_Q8_0_16; + case GGML_TYPE_IQ2_S: return MMQ_DP4A_TXS_Q8_0_16; + case GGML_TYPE_IQ3_XXS: return MMQ_DP4A_TXS_Q8_0; + case GGML_TYPE_IQ3_S: return MMQ_DP4A_TXS_Q8_0; + case GGML_TYPE_IQ1_S: return MMQ_DP4A_TXS_Q8_0; + case GGML_TYPE_IQ4_XS: return MMQ_DP4A_TXS_Q8_0; + case GGML_TYPE_IQ4_NL: return MMQ_DP4A_TXS_Q8_0; + default: return tile_x_sizes{0, 0, 0}; + } +} + +#define MMQ_MMA_TILE_X_K_Q8_0 (2*MMQ_TILE_NE_K + 2*MMQ_TILE_NE_K/QI8_0 + 4) +#define MMQ_MMA_TILE_X_K_FP4 (2*MMQ_TILE_NE_K + 8 + 4) // MXFP4 +#define MMQ_MMA_TILE_X_K_NVFP4 (2*MMQ_TILE_NE_K + MMQ_TILE_NE_K/2 + 4) // NVFP4 +#define MMQ_MMA_TILE_X_K_Q8_1 (2*MMQ_TILE_NE_K + 2*MMQ_TILE_NE_K/QI8_0 + 4) +#define MMQ_MMA_TILE_X_K_Q2_K (2*MMQ_TILE_NE_K + MMQ_TILE_NE_K + 4) +#define MMQ_MMA_TILE_X_K_Q3_K (2*MMQ_TILE_NE_K + MMQ_TILE_NE_K/2 + 4) +#define MMQ_MMA_TILE_X_K_Q6_K (2*MMQ_TILE_NE_K + MMQ_TILE_NE_K/QI6_K + MMQ_TILE_NE_K/8 + 7) + +static_assert(MMQ_MMA_TILE_X_K_Q8_0 % 8 == 4, "Wrong padding."); +static_assert(MMQ_MMA_TILE_X_K_Q8_1 % 8 == 4, "Wrong padding."); +static_assert(MMQ_MMA_TILE_X_K_Q2_K % 8 == 4, "Wrong padding."); +static_assert(MMQ_MMA_TILE_X_K_Q3_K % 8 == 4, "Wrong padding."); +static_assert(MMQ_MMA_TILE_X_K_Q6_K % 8 == 4, "Wrong padding."); +static_assert(MMQ_MMA_TILE_X_K_FP4 % 8 == 4, "Wrong padding."); +static_assert(MMQ_MMA_TILE_X_K_FP4 == MMQ_MMA_TILE_X_K_Q8_1, "Wrong tile size for MXFP4"); +static_assert(MMQ_MMA_TILE_X_K_NVFP4 % 8 == 4, "Wrong padding."); + + +static constexpr __host__ __device__ int mmq_get_mma_tile_x_k(ggml_type type) { + switch (type) { + case GGML_TYPE_Q4_0: return MMQ_MMA_TILE_X_K_Q8_0; + case GGML_TYPE_Q4_1: return MMQ_MMA_TILE_X_K_Q8_1; + case GGML_TYPE_Q5_0: return MMQ_MMA_TILE_X_K_Q8_0; + case GGML_TYPE_Q5_1: return MMQ_MMA_TILE_X_K_Q8_1; + case GGML_TYPE_Q8_0: return MMQ_MMA_TILE_X_K_Q8_0; + // tile sizes are the same for Q8_1 and FP4 for blackwell + case GGML_TYPE_MXFP4: return MMQ_MMA_TILE_X_K_Q8_1; + case GGML_TYPE_NVFP4: return MMQ_MMA_TILE_X_K_NVFP4; + case GGML_TYPE_Q2_K: return MMQ_MMA_TILE_X_K_Q2_K; + case GGML_TYPE_Q3_K: return MMQ_MMA_TILE_X_K_Q3_K; + case GGML_TYPE_Q4_K: return MMQ_MMA_TILE_X_K_Q8_1; + case GGML_TYPE_Q5_K: return MMQ_MMA_TILE_X_K_Q8_1; + case GGML_TYPE_Q6_K: return MMQ_MMA_TILE_X_K_Q6_K; + case GGML_TYPE_IQ2_XXS: return MMQ_MMA_TILE_X_K_Q8_0; + case GGML_TYPE_IQ2_XS: return MMQ_MMA_TILE_X_K_Q3_K; + case GGML_TYPE_IQ2_S: return MMQ_MMA_TILE_X_K_Q3_K; + case GGML_TYPE_IQ3_XXS: return MMQ_MMA_TILE_X_K_Q8_0; + case GGML_TYPE_IQ3_S: return MMQ_MMA_TILE_X_K_Q8_0; + case GGML_TYPE_IQ1_S: return MMQ_MMA_TILE_X_K_Q8_0; + case GGML_TYPE_IQ4_XS: return MMQ_MMA_TILE_X_K_Q8_0; + case GGML_TYPE_IQ4_NL: return MMQ_MMA_TILE_X_K_Q8_0; + default: return 0; + } +} + +// block_q8_1_mmq has (128 8-bit ints == 32 32-bit ints + 4 32-bit scales) +#define MMQ_TILE_Y_K (MMQ_TILE_NE_K + MMQ_TILE_NE_K / QI8_1) +#define MMQ_TILE_Y_FP4_K MMQ_TILE_Y_K + +static int mmq_get_granularity_host(const int mmq_x, const int cc) { + if (amd_mfma_available(cc) || amd_wmma_available(cc)) { + return mmq_x >= 128 ? 32 : 16; + } else if (turing_mma_available(cc) && mmq_x >= 48) { + return 16; + } else { + return 8; + } +} + +#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) +static constexpr __device__ int mmq_get_granularity_device(const int mmq_x) { + return mmq_x >= 128 ? 32 : 16; +} +#elif defined(TURING_MMA_AVAILABLE) +static constexpr __device__ int mmq_get_granularity_device(const int mmq_x) { + return mmq_x >= 48 ? 16 : 8; +} +#else +static constexpr __device__ int mmq_get_granularity_device(const int /*mmq_x*/) { + return 8; +} +#endif // AMD_MFMA_AVAILABLE + +#if defined(GGML_USE_HIP) +static int mmq_get_nwarps_host(const int cc, const int warp_size) { + return amd_mfma_available(cc) ? 8 : 256/warp_size; +} +#else +static int mmq_get_nwarps_host(const int /*cc*/, const int warp_size) { + return 256/warp_size; +} +#endif // (GGML_USE_HIP) + +static constexpr __device__ int mmq_get_nwarps_device() { +#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + return 8; +#else + return 256/ggml_cuda_get_physical_warp_size(); +#endif // AMD_MFMA_AVAILABLE +} + +// ------------------------------------------------------------ + +template static __device__ __forceinline__ void load_tiles_q4_0( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + 2*MMQ_TILE_NE_K); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_0, mmq_y); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = MMQ_ITER_K / (4 * QR4_0); + constexpr int nrows = warp_size / threads_per_row; + const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x; + const int kbx = txi / QI4_0; + const int kqsx = txi % QI4_0; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nrows*nwarps) { + int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row); + + if (need_check) { + i = min(i, i_max); + } + + const block_q4_0 * bxi = (const block_q4_0 *) x + kbx0 + i*stride + kbx; + const int qs0 = get_int_b2(bxi->qs, kqsx); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + kbx*(2*QI4_0) + kqsx + 0] = __vsubss4((qs0 >> 0) & 0x0F0F0F0F, 0x08080808); + x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + kbx*(2*QI4_0) + kqsx + QI4_0] = __vsubss4((qs0 >> 4) & 0x0F0F0F0F, 0x08080808); +#else + x_qs[i*(MMQ_TILE_NE_K + 1) + txi] = qs0; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) + } + + constexpr int blocks_per_tile_x_row = MMQ_TILE_NE_K / QI4_0; + constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row; + const int kbxd = threadIdx.x % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * rows_per_warp) { + int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row; + + if (need_check) { + i = min(i, i_max); + } + + const block_q4_0 * bxi = (const block_q4_0 *) x + kbx0 + i*stride + kbxd; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_df[i*MMQ_MMA_TILE_X_K_Q8_0 + kbxd] = bxi->d; +#else + x_df[i*(MMQ_TILE_NE_K/QI4_0) + i/QI4_0 + kbxd] = bxi->d; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +template +static __device__ __forceinline__ void vec_dot_q4_0_q8_1_dp4a( + const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_0, mmq_y); + const int * x_qs = (const int *) x; + const float * x_df = (const float *) x_qs + txs.qs; + const int * y_qs = (const int *) y + 4; + const half2 * y_ds = (const half2 *) y; + +// #pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR4_0*VDR_Q4_0_Q8_1_MMQ) { + const int k0 = k00 + k01; + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += nwarps) { + const int j = j0 + threadIdx.y; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += warp_size) { + const int i = i0 + threadIdx.x; + + const int kyqs = QI8_1 * ((k01/2) / (QI8_1/2)) + (k01/2) % (QI8_1/2); + + int u[2*VDR_Q4_0_Q8_1_MMQ]; + +#pragma unroll + for (int l = 0; l < VDR_Q4_0_Q8_1_MMQ; ++l) { + u[2*l+0] = y_qs[j*MMQ_TILE_Y_K + kyqs + l]; + u[2*l+1] = y_qs[j*MMQ_TILE_Y_K + kyqs + (l + QI4_0)]; + } + + sum[j0/nwarps*mmq_y/warp_size + i0/warp_size] += vec_dot_q4_0_q8_1_impl + (&x_qs[i*(MMQ_TILE_NE_K + 1) + k0/QR4_0], u, + x_df[i*(MMQ_TILE_NE_K/QI4_0) + i/QI4_0 + k0/(QR4_0*QI4_0)], y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]); + } + } + } +} + +template static __device__ __forceinline__ void load_tiles_q4_1( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + int * x_qs = (int *) x_tile; + half2 * x_dm = (half2 *) (x_qs + 2*MMQ_TILE_NE_K); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_1, mmq_y); + int * x_qs = (int *) x_tile; + half2 * x_dm = (half2 *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = MMQ_ITER_K / (4 * QR4_1); + constexpr int nrows = warp_size / threads_per_row; + const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x; + const int kbx = txi / QI4_1; + const int kqsx = txi % QI4_1; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nrows*nwarps) { + int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row); + + if (need_check) { + i = min(i, i_max); + } + + const block_q4_1 * bxi = (const block_q4_1 *) x + kbx0 + i*stride + kbx; + const int qs0 = get_int_b4(bxi->qs, kqsx); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*MMQ_MMA_TILE_X_K_Q8_1 + kbx*(2*QI4_1) + kqsx + 0] = (qs0 >> 0) & 0x0F0F0F0F; + x_qs[i*MMQ_MMA_TILE_X_K_Q8_1 + kbx*(2*QI4_1) + kqsx + QI4_1] = (qs0 >> 4) & 0x0F0F0F0F; +#else + x_qs[i*(MMQ_TILE_NE_K + 1) + txi] = qs0; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + + constexpr int blocks_per_tile_x_row = MMQ_TILE_NE_K / QI4_1; + constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row; + const int kbxd = threadIdx.x % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * rows_per_warp) { + int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row; + + if (need_check) { + i = min(i, i_max); + } + + const block_q4_1 * bxi = (const block_q4_1 *) x + kbx0 + i*stride + kbxd; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_dm[i*MMQ_MMA_TILE_X_K_Q8_1 + kbxd] = bxi->dm; +#else + x_dm[i*(MMQ_TILE_NE_K/QI4_1) + i/QI4_1 + kbxd] = bxi->dm; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +template +static __device__ __forceinline__ void vec_dot_q4_1_q8_1_dp4a( + const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_1, mmq_y); + const int * x_qs = (const int *) x; + const half2 * x_dm = (const half2 *) x_qs + txs.qs; + const int * y_qs = (const int *) y + 4; + const half2 * y_ds = (const half2 *) y; + +// #pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR4_1*VDR_Q4_1_Q8_1_MMQ) { + const int k0 = k00 + k01; + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += nwarps) { + const int j = j0 + threadIdx.y; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += warp_size) { + const int i = i0 + threadIdx.x; + + const int kyqs = QI8_1 * ((k01/2) / (QI8_1/2)) + (k01/2) % (QI8_1/2); + + int u[2*VDR_Q4_1_Q8_1_MMQ]; + +#pragma unroll + for (int l = 0; l < VDR_Q4_1_Q8_1_MMQ; ++l) { + u[2*l+0] = y_qs[j*MMQ_TILE_Y_K + kyqs + l]; + u[2*l+1] = y_qs[j*MMQ_TILE_Y_K + kyqs + (l + QI4_1)]; + } + + sum[j0/nwarps*mmq_y/warp_size + i0/warp_size] += vec_dot_q4_1_q8_1_impl + (&x_qs[i*(MMQ_TILE_NE_K + 1) + k0/QR4_1], u, + x_dm[i*(MMQ_TILE_NE_K/QI4_1) + i/QI4_1 + k0/(QR4_1*QI4_1)], y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]); + } + } + } +} + +template static __device__ __forceinline__ void load_tiles_q5_0( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q5_0, mmq_y); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = MMQ_ITER_K / (4 * QR5_0); + constexpr int nrows = warp_size / threads_per_row; + const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x; + const int kbx = txi / QI5_0; + const int kqsx = txi % QI5_0; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nrows*nwarps) { + int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row); + + if (need_check) { + i = min(i, i_max); + } + + const block_q5_0 * bxi = (const block_q5_0 *) x + kbx0 + i*stride + kbx; + + const int ql = get_int_b2(bxi->qs, kqsx); + const int qh = get_int_b2(bxi->qh, 0) >> (4 * kqsx); + + int qs0 = (ql >> 0) & 0x0F0F0F0F; + qs0 |= (qh << 4) & 0x00000010; // 0 -> 4 + qs0 |= (qh << 11) & 0x00001000; // 1 -> 12 + qs0 |= (qh << 18) & 0x00100000; // 2 -> 20 + qs0 |= (qh << 25) & 0x10000000; // 3 -> 28 + qs0 = __vsubss4(qs0, 0x10101010); // subtract 16 + + int qs1 = (ql >> 4) & 0x0F0F0F0F; + qs1 |= (qh >> 12) & 0x00000010; // 16 -> 4 + qs1 |= (qh >> 5) & 0x00001000; // 17 -> 12 + qs1 |= (qh << 2) & 0x00100000; // 18 -> 20 + qs1 |= (qh << 9) & 0x10000000; // 19 -> 28 + qs1 = __vsubss4(qs1, 0x10101010); // subtract 16 + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + kbx*(2*QI5_0) + kqsx + 0] = qs0; + x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + kbx*(2*QI5_0) + kqsx + QI5_0] = qs1; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + kbx*(2*QI5_0) + kqsx + 0] = qs0; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + kbx*(2*QI5_0) + kqsx + QI5_0] = qs1; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + + constexpr int blocks_per_tile_x_row = MMQ_TILE_NE_K / QI5_0; + constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row; + const int kbxd = threadIdx.x % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * rows_per_warp) { + int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row; + + if (need_check) { + i = min(i, i_max); + } + + const block_q5_0 * bxi = (const block_q5_0 *) x + kbx0 + i*stride + kbxd; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_df[i*MMQ_MMA_TILE_X_K_Q8_0 + kbxd] = bxi->d; +#else + x_df[i*(MMQ_TILE_NE_K/QI5_0) + i/QI5_0 + kbxd] = bxi->d; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +template static __device__ __forceinline__ void load_tiles_q5_1( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + int * x_qs = (int *) x_tile; + half2 * x_dm = (half2 *) (x_qs + 2*MMQ_TILE_NE_K); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q5_1, mmq_y); + int * x_qs = (int *) x_tile; + half2 * x_dm = (half2 *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = MMQ_ITER_K / (4 * QR5_1); + constexpr int nrows = warp_size / threads_per_row; + const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x; + const int kbx = txi / QI5_1; + const int kqsx = txi % QI5_1; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nrows*nwarps) { + int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row); + + if (need_check) { + i = min(i, i_max); + } + + const block_q5_1 * bxi = (const block_q5_1 *) x + kbx0 + i*stride + kbx; + + const int ql = get_int_b4(bxi->qs, kqsx); + const int qh = get_int_b4(bxi->qh, 0) >> (4 * kqsx); + + int qs0 = (ql >> 0) & 0x0F0F0F0F; + qs0 |= (qh << 4) & 0x00000010; // 0 -> 4 + qs0 |= (qh << 11) & 0x00001000; // 1 -> 12 + qs0 |= (qh << 18) & 0x00100000; // 2 -> 20 + qs0 |= (qh << 25) & 0x10000000; // 3 -> 28 + + int qs1 = (ql >> 4) & 0x0F0F0F0F; + qs1 |= (qh >> 12) & 0x00000010; // 16 -> 4 + qs1 |= (qh >> 5) & 0x00001000; // 17 -> 12 + qs1 |= (qh << 2) & 0x00100000; // 18 -> 20 + qs1 |= (qh << 9) & 0x10000000; // 19 -> 28 + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*MMQ_MMA_TILE_X_K_Q8_1 + kbx*(2*QI5_1) + kqsx + 0] = qs0; + x_qs[i*MMQ_MMA_TILE_X_K_Q8_1 + kbx*(2*QI5_1) + kqsx + QI5_1] = qs1; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + kbx*(2*QI5_1) + kqsx + 0] = qs0; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + kbx*(2*QI5_1) + kqsx + QI5_1] = qs1; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + + constexpr int blocks_per_tile_x_row = MMQ_TILE_NE_K / QI5_1; + constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row; + const int kbxd = threadIdx.x % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * rows_per_warp) { + int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row; + + if (need_check) { + i = min(i, i_max); + } + + const block_q5_1 * bxi = (const block_q5_1 *) x + kbx0 + i*stride + kbxd; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_dm[i*MMQ_MMA_TILE_X_K_Q8_1 + kbxd] = bxi->dm; +#else + x_dm[i*(MMQ_TILE_NE_K/QI5_1) + i/QI5_1 + kbxd] = bxi->dm; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +template static __device__ __forceinline__ void load_tiles_q8_0( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_tile + 2*MMQ_TILE_NE_K); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q8_0, mmq_y); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + // MMQ_ITER_K / (4 * QR8_0) == 64 required. but NV has only 32 threads per warp + constexpr int threads_per_row = 32; + constexpr int nrows = warp_size / threads_per_row; + const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x; + const int kbx = txi / QI8_0; + const int kqsx = txi % QI8_0; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nrows*nwarps) { + int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row); + + if (need_check) { + i = min(i, i_max); + } + + const block_q8_0 * bxi = (const block_q8_0 *) x + kbx0 + i*stride + kbx; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + 0 + txi] = get_int_b2(bxi[0].qs, kqsx); + x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + MMQ_TILE_NE_K + txi] = get_int_b2(bxi[MMQ_TILE_NE_K/QI8_0].qs, kqsx); +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + 0 + txi] = get_int_b2(bxi[0].qs, kqsx); + x_qs[i*(2*MMQ_TILE_NE_K + 1) + MMQ_TILE_NE_K + txi] = get_int_b2(bxi[MMQ_TILE_NE_K/QI8_0].qs, kqsx); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + + constexpr int blocks_per_tile_x_row = 2*MMQ_TILE_NE_K / QI8_0; + constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row; + const int kbxd = threadIdx.x % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * rows_per_warp) { + int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row; + + if (need_check) { + i = min(i, i_max); + } + + const block_q8_0 * bxi = (const block_q8_0 *) x + kbx0 + i*stride + kbxd; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_df[i*MMQ_MMA_TILE_X_K_Q8_0 + kbxd] = bxi->d; +#else + x_df[i*(2*MMQ_TILE_NE_K/QI8_0) + i/(QI8_0/2) + kbxd] = bxi->d; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +template static __device__ __forceinline__ void load_tiles_mxfp4( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_MXFP4, mmq_y); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = MMQ_ITER_K / (4 * QR_MXFP4); + constexpr int nrows = warp_size / threads_per_row; + const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x; + const int kbx = txi / QI_MXFP4; + const int kqsx = txi % QI_MXFP4; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nrows*nwarps) { + int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row); + + if (need_check) { + i = min(i, i_max); + } + + const block_mxfp4 * bxi = (const block_mxfp4 *) x + kbx0 + i*stride + kbx; + + const int aux_q4 = get_int_b1(bxi->qs, kqsx); + const int2 v = get_int_from_table_16(aux_q4, kvalues_mxfp4); + const int k0 = kbx * (2 * QI_MXFP4) + kqsx; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*MMQ_MMA_TILE_X_K_Q8_1 + k0 + 0] = v.x; + x_qs[i*MMQ_MMA_TILE_X_K_Q8_1 + k0 + QI_MXFP4] = v.y; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + 0] = v.x; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + QI_MXFP4] = v.y; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + + constexpr int blocks_per_tile_x_row = MMQ_TILE_NE_K / QI_MXFP4; + constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row; + const int kbxd = threadIdx.x % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * rows_per_warp) { + int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row; + + if (need_check) { + i = min(i, i_max); + } + + const block_mxfp4 * bxi = (const block_mxfp4 *) x + kbx0 + i*stride + kbxd; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_df[i*MMQ_MMA_TILE_X_K_Q8_1 + kbxd] = ggml_cuda_e8m0_to_fp32(bxi->e)*0.5f; +#else + x_df[i*(MMQ_TILE_NE_K/QI_MXFP4) + i/QI_MXFP4 + kbxd] = ggml_cuda_e8m0_to_fp32(bxi->e)*0.5f; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +template +static __device__ __forceinline__ void load_tiles_mxfp4_fp4(const char * __restrict__ x, + int * __restrict__ x_tile, + const int kbx0, + const int i_max, + const int stride) { + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + + int * x_qs = (int *) x_tile; + uint32_t * x_sc = (uint32_t *) (x_qs + 2 * MMQ_TILE_NE_K); + + const int txi = threadIdx.x; + + constexpr int iter_k = get_iter_k(GGML_TYPE_MXFP4); + + constexpr int threads_per_row = iter_k / QK_MXFP4; // each thread processes 1 block + constexpr int rows_per_warp = warp_size / threads_per_row; + const int kbx = txi % threads_per_row; + const int row_in_warp = txi / threads_per_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += rows_per_warp * nwarps) { + int i = i0 + threadIdx.y * rows_per_warp + row_in_warp; + + if constexpr (need_check) { + i = min(i, i_max); + } + + const block_mxfp4 * bxi = (const block_mxfp4 *) x + kbx0 + i * stride + kbx; + + // quantize_mxfp4_mmq permutes nibbles to match the quantized format + const int k0 = kbx * 4; + memcpy(x_qs + i * MMQ_MMA_TILE_X_K_FP4 + k0, bxi->qs, 16); + + // Load E8M0 scales: pack 2 consecutive scales into one uint32 + if (kbx % 2 == 0) { + uint32_t e = bxi->e; + e |= ((bxi + 1)->e << 8); + x_sc[i * MMQ_MMA_TILE_X_K_FP4 + kbx / 2] = e; + } + } +} + + +template +static __device__ __forceinline__ void load_tiles_nvfp4(const char * __restrict__ x, + int * __restrict__ x_tile, + const int kb0, + const int i_max, + const int stride) { + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_NVFP4, mmq_y); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = MMQ_ITER_K / QK_NVFP4; + constexpr int rows_per_warp = warp_size / threads_per_row; + const int kbx = threadIdx.x % threads_per_row; + const int row_in_warp = threadIdx.x / threads_per_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += rows_per_warp * nwarps) { + int i = i0 + threadIdx.y * rows_per_warp + row_in_warp; + + if constexpr (need_check) { + i = min(i, i_max); + } + + const block_nvfp4 * bxi = (const block_nvfp4 *) x + kb0 + i * stride + kbx; + const uint32_t * __restrict__ src_qs = reinterpret_cast(bxi->qs); + const int kqs = 16 * kbx; + const int ksc = 4 * kbx; + +#pragma unroll + for (int sub = 0; sub < QK_NVFP4 / QK_NVFP4_SUB; ++sub) { + const int2 q0 = get_int_from_table_16(src_qs[2 * sub + 0], kvalues_mxfp4); + const int2 q1 = get_int_from_table_16(src_qs[2 * sub + 1], kvalues_mxfp4); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i * MMQ_MMA_TILE_X_K_NVFP4 + kqs + 4 * sub + 0] = q0.x; + x_qs[i * MMQ_MMA_TILE_X_K_NVFP4 + kqs + 4 * sub + 1] = q1.x; + x_qs[i * MMQ_MMA_TILE_X_K_NVFP4 + kqs + 4 * sub + 2] = q0.y; + x_qs[i * MMQ_MMA_TILE_X_K_NVFP4 + kqs + 4 * sub + 3] = q1.y; + x_df[i * MMQ_MMA_TILE_X_K_NVFP4 + ksc + sub] = ggml_cuda_ue4m3_to_fp32(bxi->d[sub]); +#else + x_qs[i * (2 * MMQ_TILE_NE_K + 1) + kqs + 4 * sub + 0] = q0.x; + x_qs[i * (2 * MMQ_TILE_NE_K + 1) + kqs + 4 * sub + 1] = q1.x; + x_qs[i * (2 * MMQ_TILE_NE_K + 1) + kqs + 4 * sub + 2] = q0.y; + x_qs[i * (2 * MMQ_TILE_NE_K + 1) + kqs + 4 * sub + 3] = q1.y; + x_df[i * (2 * MMQ_TILE_NE_K * 2 / QI_NVFP4) + i / (QK_NVFP4_SUB / QI_NVFP4) + ksc + sub] = ggml_cuda_ue4m3_to_fp32(bxi->d[sub]); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + } +} + +template +static __device__ __forceinline__ void vec_dot_q8_0_q8_1_dp4a( + const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q8_0, mmq_y); + const int * x_qs = (const int *) x; + const float * x_df = (const float *) x_qs + txs.qs; + const int * y_qs = (const int *) y + 4; + const float * y_df = (const float *) y; + +// #pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += VDR_Q8_0_Q8_1_MMQ) { + const int k0 = k00 + k01; + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += nwarps) { + const int j = j0 + threadIdx.y; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += warp_size) { + const int i = i0 + threadIdx.x; + + sum[j0/nwarps*mmq_y/warp_size + i0/warp_size] += vec_dot_q8_0_q8_1_impl + (&x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k0 % MMQ_TILE_NE_K], + x_df[i*(2*MMQ_TILE_NE_K/QI8_0) + i/(QI8_0/2) + k0/QI8_0], y_df[j*MMQ_TILE_Y_K + (k0/QI8_1) % (MMQ_TILE_NE_K/QI8_1)]); + } + } + } +} + +template +static __device__ __forceinline__ void vec_dot_q8_0_q8_1_mma( + const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { +#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr data_layout input_layout = get_input_data_layout(); + typedef tile<16, 8, int, input_layout> tile_A; + typedef tile<16, 8, int, input_layout> tile_B; + typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C; + + constexpr int granularity = mmq_get_granularity_device(mmq_x); + constexpr int rows_per_warp = granularity; + constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. + + y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); + + const int * x_qs = (const int *) x; + const float * x_df = (const float *) x_qs + 2*MMQ_TILE_NE_K; + const int * y_qs = (const int *) y + 4; + const float * y_df = (const float *) y; + const half2 * y_ds = (const half2 *) y; + + const int i0 = (threadIdx.y / ntx) * rows_per_warp; + + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) { + const int k0 = k00 + k01; + + tile_A A[ntx]; +#pragma unroll + for (int n = 0; n < ntx; ++n) { + load_generic(A[n], x_qs + (i0 + n*tile_A::I)*MMQ_MMA_TILE_X_K_Q8_0 + k0, MMQ_MMA_TILE_X_K_Q8_0); + } + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += ntx*tile_C::J) { + tile_B B; + load_generic(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); + + float dB; + const int j = j0 + tile_C::get_j(0); + if (ds_layout == MMQ_Q8_1_DS_LAYOUT_D4) { + dB = y_df[j*MMQ_TILE_Y_K + k01/QI8_1]; + } else { + dB = __low2float(y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]); + } + +#pragma unroll + for (int n = 0; n < ntx; ++n) { + tile_C C; + mma(C, A[n], B); + +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + const int i = i0 + n*tile_A::I + tile_C::get_i(l); + const float dA = x_df[i*MMQ_MMA_TILE_X_K_Q8_0 + k0/QI8_0]; + sum[(j0/tile_C::J + n)*tile_C::ne + l] += C.x[l]*dA*dB; + } + } + } + } +#else + typedef tile<16, 8, int> tile_A; + typedef tile< 8, 8, int> tile_B; + typedef tile<16, 8, int> tile_C; + + constexpr int granularity = mmq_get_granularity_device(mmq_x); + constexpr int rows_per_warp = 2 * granularity; + constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. + + y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); + + const int * x_qs = (const int *) x; + const float * x_df = (const float *) x_qs + 2*MMQ_TILE_NE_K; + const int * y_qs = (const int *) y + 4; + const float * y_df = (const float *) y; + const half2 * y_ds = (const half2 *) y; + + tile_A A[ntx][MMQ_TILE_NE_K/QI8_0]; + float dA[ntx][tile_C::ne/2][MMQ_TILE_NE_K/QI8_0]; + + const int i0 = (threadIdx.y/ntx)*rows_per_warp; + +#pragma unroll + for (int n = 0; n < ntx; ++n) { +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) { + const int k0 = k00 + k01; + + load_ldmatrix(A[n][k01/QI8_0], x_qs + (i0 + n*tile_A::I)*MMQ_MMA_TILE_X_K_Q8_0 + k0, MMQ_MMA_TILE_X_K_Q8_0); + } + +#pragma unroll + for (int l = 0; l < tile_C::ne/2; ++l) { + const int i = i0 + n*tile_A::I + tile_C::get_i(2*l); + +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) { + const int k0 = k00 + k01; + + dA[n][l][k01/QI8_0] = x_df[i*MMQ_MMA_TILE_X_K_Q8_0 + k0/QI8_0]; + } + } + } + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += ntx*tile_C::J) { +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) { + tile_B B; + float dB[tile_C::ne/2]; + + load_generic(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); // faster than load_ldmatrix + +#pragma unroll + for (int l = 0; l < tile_C::ne/2; ++l) { + const int j = j0 + tile_C::get_j(l); + + if (ds_layout == MMQ_Q8_1_DS_LAYOUT_D4) { + dB[l] = y_df[j*MMQ_TILE_Y_K + k01/QI8_1]; + } else { + dB[l] = __low2float(y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]); + } + } + +#pragma unroll + for (int n = 0; n < ntx; ++n) { + tile_C C; + mma(C, A[n][k01/QI8_0], B); + +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + sum[(j0/tile_C::J + n)*tile_C::ne + l] += C.x[l]*dA[n][l/2][k01/QI8_0]*dB[l%2]; + } + } + } + } +#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) +} + +template +static __device__ __forceinline__ void vec_dot_mxfp4_mxfp4_mma(const int * __restrict__ x, + const int * __restrict__ y, + float * __restrict__ sum, + const int k00) { + typedef tile<16, 8, int> tile_A; + typedef tile<8, 8, int> tile_B; + typedef tile<16, 8, float> tile_C; // Output is float for native scaled MMA + + constexpr int granularity = mmq_get_granularity_device(mmq_x); + constexpr int rows_per_warp = 2 * granularity; + constexpr int ntx = rows_per_warp / tile_C::I; // Number of x minitiles per warp. + + y += (threadIdx.y % ntx) * (tile_C::J * MMQ_TILE_Y_FP4_K); + + // Match layout from load_tiles_mxfp4_fp4 + const int * x_qs = (const int *) x; + const uint32_t * x_sc = (const uint32_t *) (x_qs + 2 * MMQ_TILE_NE_K); + const int * y_qs = (const int *) y + 4; + const uint32_t * y_sc = (const uint32_t *) y; + + // tile_A has a length of 64 logical values vs. 32 values in block_mxfp4 + tile_A A[ntx][MMQ_TILE_NE_K / (2 * QI_MXFP4)]; + uint32_t scaleA[ntx][MMQ_TILE_NE_K / (2 * QI_MXFP4)]; + + // Block scale + // Each thread has to point to a 4 byte scale value + // https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-block-scaling + + const int i0 = (threadIdx.y / ntx) * rows_per_warp; + +#pragma unroll + for (int n = 0; n < ntx; ++n) { +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 2 * QI_MXFP4) { + const int k0 = k00 + k01; + + load_ldmatrix(A[n][k01 / (2 * QI_MXFP4)], x_qs + (i0 + n * tile_A::I) * MMQ_MMA_TILE_X_K_FP4 + k0, + MMQ_MMA_TILE_X_K_FP4); + + // based on block-scaling document, 2 threads in each quad need to supply to the scale value + const int tidx = threadIdx.x / 4 + (threadIdx.x % 2) * 8; + scaleA[n][k01 / (2 * QI_MXFP4)] = + *(x_sc + (i0 + n * tile_A::I + tidx) * MMQ_MMA_TILE_X_K_FP4 + k0 / (2 * QI_MXFP4)); + } + } + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += ntx * tile_C::J) { +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 2 * QI_MXFP4) { + tile_B B; + uint32_t scaleB; // 2xN scales + + load_generic(B, y_qs + j0 * MMQ_TILE_Y_FP4_K + k01, MMQ_TILE_Y_FP4_K); + + scaleB = y_sc[(j0 + threadIdx.x / 4) * MMQ_TILE_Y_FP4_K + k01 / (2 * QI_MXFP4)]; + +#pragma unroll + for (int n = 0; n < ntx; ++n) { + tile_C C; + + mma_block_scaled(C, A[n][k01 / (2 * QI_MXFP4)], B, scaleA[n][k01 / (2 * QI_MXFP4)], scaleB); +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + sum[(j0 / tile_C::J + n) * tile_C::ne + l] += C.x[l]; + } + } + } + } +} + +template +static __device__ __forceinline__ void vec_dot_q8_1_q8_1_dp4a( + const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q5_1, mmq_y); + const int * x_qs = (const int *) x; + const half2 * x_dm = (const half2 *) x_qs + txs.qs; + const int * y_qs = (const int *) y + 4; + const half2 * y_ds = (const half2 *) y; + +// #pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += VDR_Q8_0_Q8_1_MMQ) { + const int k0 = k00 + k01; + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += nwarps) { + const int j = j0 + threadIdx.y; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += warp_size) { + const int i = i0 + threadIdx.x; + + sum[j0/nwarps*mmq_y/warp_size + i0/warp_size] += vec_dot_q8_1_q8_1_impl + (&x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01], + x_dm[i*(MMQ_TILE_NE_K/QI5_1) + i/QI5_1 + k0/QI8_1], y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]); + } + } + } +} + +template +static __device__ __forceinline__ void vec_dot_q8_1_q8_1_mma( + const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { +#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr data_layout input_layout = get_input_data_layout(); + typedef tile<16, 8, int, input_layout> tile_A; + typedef tile<16, 8, int, input_layout> tile_B; + typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C; + + constexpr int granularity = mmq_get_granularity_device(mmq_x); + constexpr int rows_per_warp = granularity; + constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. + + y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); + + const int * x_qs = (const int *) x; + const half2 * x_dm = (const half2 *) x_qs + 2*MMQ_TILE_NE_K; + const int * y_qs = (const int *) y + 4; + const half2 * y_dm = (const half2 *) y; + + const int i0 = (threadIdx.y / ntx) * rows_per_warp; + + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) { + const int k0 = k00 + k01; + + tile_A A[ntx]; +#pragma unroll + for (int n = 0; n < ntx; ++n) { + load_generic(A[n], x_qs + (i0 + n*tile_A::I)*MMQ_MMA_TILE_X_K_Q8_1 + k0, MMQ_MMA_TILE_X_K_Q8_1); + } + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += ntx*tile_C::J) { + tile_B B; + load_generic(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); + + const int j = j0 + tile_C::get_j(0); + const float2 dsB = __half22float2(y_dm[j*MMQ_TILE_Y_K + k01/QI8_1]); + +#pragma unroll + for (int n = 0; n < ntx; ++n) { + tile_C C; + mma(C, A[n], B); + +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + const int i = i0 + n*tile_A::I + tile_C::get_i(l); + float2 dmA = __half22float2(x_dm[i*MMQ_MMA_TILE_X_K_Q8_1 + k0/QI8_1]); + sum[(j0/tile_C::J + n)*tile_C::ne + l] += dmA.x*dsB.x*C.x[l]; + sum[(j0/tile_C::J + n)*tile_C::ne + l] += dmA.y*dsB.y; + } + } + } + } +#else + typedef tile<16, 8, int> tile_A; + typedef tile< 8, 8, int> tile_B; + typedef tile<16, 8, int> tile_C; + + constexpr int granularity = mmq_get_granularity_device(mmq_x); + constexpr int rows_per_warp = 2 * granularity; + constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. + + y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); + + const int * x_qs = (const int *) x; + const half2 * x_dm = (const half2 *) x_qs + 2*MMQ_TILE_NE_K; + const int * y_qs = (const int *) y + 4; + const half2 * y_dm = (const half2 *) y; + + tile_A A[ntx][MMQ_TILE_NE_K/QI8_1]; + float2 dmA[ntx][tile_C::ne/2][MMQ_TILE_NE_K/QI8_1]; + + const int i0 = (threadIdx.y/ntx)*rows_per_warp; + +#pragma unroll + for (int n = 0; n < ntx; ++n) { +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) { + const int k0 = k00 + k01; + + load_ldmatrix(A[n][k01/QI8_1], x_qs + (i0 + n*tile_A::I)*MMQ_MMA_TILE_X_K_Q8_1 + k0, MMQ_MMA_TILE_X_K_Q8_1); + } + +#pragma unroll + for (int l = 0; l < tile_C::ne/2; ++l) { + const int i = i0 + n*tile_A::I + tile_C::get_i(2*l); + +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) { + const int k0 = k00 + k01; + + dmA[n][l][k01/QI8_1] = __half22float2(x_dm[i*MMQ_MMA_TILE_X_K_Q8_1 + k0/QI8_1]); + } + } + } + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += ntx*tile_C::J) { +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) { + tile_B B; + float2 dsB[tile_C::ne/2]; + + load_generic(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); // faster than load_ldmatrix + +#pragma unroll + for (int l = 0; l < tile_C::ne/2; ++l) { + const int j = j0 + tile_C::get_j(l); + + dsB[l] = __half22float2(y_dm[j*MMQ_TILE_Y_K + k01/QI8_1]); + } + +#pragma unroll + for (int n = 0; n < ntx; ++n) { + tile_C C; + mma(C, A[n][k01/QI8_1], B); + +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + sum[(j0/tile_C::J + n)*tile_C::ne + l] += dmA[n][l/2][k01/QI8_1].x*dsB[l%2].x*C.x[l]; + sum[(j0/tile_C::J + n)*tile_C::ne + l] += dmA[n][l/2][k01/QI8_1].y*dsB[l%2].y; + } + } + } + } +#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) +} + +// Used for NVFP4, Q3_K, IQ2_S, and IQ2_XS +template +static __device__ __forceinline__ void vec_dot_q8_0_16_q8_1_dp4a( + const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + + constexpr tile_x_sizes txs = MMQ_DP4A_TXS_Q8_0_16; + const int * x_qs = (const int *) x; + const float * x_df = (const float *) x_qs + txs.qs; + const int * y_qs = (const int *) y + 4; + const float * y_df = (const float *) y; + +// #pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) { + const int k0 = k00 + k01; + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += nwarps) { + const int j = j0 + threadIdx.y; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += warp_size) { + const int i = i0 + threadIdx.x; + + sum[j0/nwarps*mmq_y/warp_size + i0/warp_size] += vec_dot_q8_0_16_q8_1_impl( + &x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0], + &y_qs[j*MMQ_TILE_Y_K + k01], + &x_df[i*(2*MMQ_TILE_NE_K*2/QI8_0) + i/(QI8_0/4) + k0/(QI8_0/2)], + y_df[j*MMQ_TILE_Y_K + k01/QI8_1]); + } + } + } +} + +// Used for Q3_K, IQ2_S, and IQ2_XS: +template +static __device__ __forceinline__ void vec_dot_q8_0_16_q8_1_mma( + const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { +#if defined(AMD_MFMA_AVAILABLE) + constexpr data_layout input_layout = get_input_data_layout(); + typedef tile<16, 8, int, input_layout> tile_A; + typedef tile<16, 8, int, input_layout> tile_B; + typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C; + typedef tile<64, 2, int, input_layout> tile_load; + + constexpr int granularity = mmq_get_granularity_device(mmq_x); + constexpr int rows_per_warp = granularity; + constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. + + y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); + + const int * x_qs = (const int *) x; + const float * x_df = (const float *) x_qs + MMQ_TILE_NE_K*2; + const int * y_qs = (const int *) y + 4; + const float * y_df = (const float *) y; + + const int i0 = (threadIdx.y / ntx) * rows_per_warp; + + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 4) { + const int k0 = k00 + k01; + + tile_A A[ntx]; +#pragma unroll + for (int n = 0; n < ntx; ++n) { + load_generic(((tile_load *) A)[n], x_qs + (i0 + n*tile_A::I)*MMQ_MMA_TILE_X_K_Q3_K + k0, MMQ_MMA_TILE_X_K_Q3_K); + } + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += ntx*tile_C::J) { + tile_B B[1]; + load_generic(((tile_load *) B)[0], y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); + + const int j = j0 + tile_C::get_j(0); + const float dB = y_df[j*MMQ_TILE_Y_K + k01/QI8_1] / 2; + +#pragma unroll + for (int n = 0; n < ntx; ++n) { + tile_C C; + mma(C, A[n], B[0]); + +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + const int i = i0 + n*tile_C::I + tile_C::get_i(l); + sum[(j0/tile_C::J + n)*tile_C::ne + l] += C.x[l] * x_df[i*MMQ_MMA_TILE_X_K_Q3_K + k0/4] * dB; + } + } + } + } +#elif defined(AMD_WMMA_AVAILABLE) //wmma instructions can handle 16x4 tiles, does not require loading 64x2 tiles + constexpr data_layout input_layout = get_input_data_layout(); + typedef tile<16, 4, int, input_layout> tile_A; + typedef tile<16, 4, int, input_layout> tile_B; + typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C; + + constexpr int granularity = mmq_get_granularity_device(mmq_x); + constexpr int rows_per_warp = granularity; + constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. + + y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); + + const int * x_qs = (const int *) x; + const float * x_df = (const float *) x_qs + MMQ_TILE_NE_K*2; + const int * y_qs = (const int *) y + 4; + const float * y_df = (const float *) y; + + const int i0 = (threadIdx.y / ntx) * rows_per_warp; + + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 4) { + const int k0 = k00 + k01; + + tile_A A[ntx]; +#pragma unroll + for (int n = 0; n < ntx; ++n) { + load_generic(A[n], x_qs + (i0 + n*tile_A::I)*MMQ_MMA_TILE_X_K_Q3_K + k0, MMQ_MMA_TILE_X_K_Q3_K); + } + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += ntx*tile_C::J) { + tile_B B; + load_generic(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); + + const int j = j0 + tile_C::get_j(0); + const float dB = y_df[j*MMQ_TILE_Y_K + k01/QI8_1]; + +#pragma unroll + for (int n = 0; n < ntx; ++n) { + tile_C C; + mma(C, A[n], B); + +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + const int i = i0 + n*tile_C::I + tile_C::get_i(l); + sum[(j0/tile_C::J + n)*tile_C::ne + l] += C.x[l] * x_df[i*MMQ_MMA_TILE_X_K_Q3_K + k0/4] * dB; + } + } + } + } +#elif defined(TURING_MMA_AVAILABLE) + + typedef tile<16, 4, int> tile_A; + typedef tile<16, 8, int> tile_A_8; + typedef tile< 8, 4, int> tile_B; + typedef tile<16, 8, int> tile_C; + + constexpr int granularity = mmq_get_granularity_device(mmq_x); + constexpr int rows_per_warp = 2 * granularity; + constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. + + y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); + + const int * x_qs = (const int *) x; + const float * x_df = (const float *) x_qs + MMQ_TILE_NE_K*2; + const int * y_qs = (const int *) y + 4; + const float * y_df = (const float *) y; + + const int i0 = (threadIdx.y / ntx) * (ntx*tile_A::I); + + tile_A A[ntx][8]; + float dA[ntx][tile_C::ne/2][8]; + +#pragma unroll + for (int n = 0; n < ntx; ++n) { +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 8) { + const int k0 = k00 + k01; + + load_ldmatrix(((tile_A_8 *) A[n])[k01/8], x_qs + (i0 + n*tile_A::I)*MMQ_MMA_TILE_X_K_Q3_K + k0, MMQ_MMA_TILE_X_K_Q3_K); + } + +#pragma unroll + for (int l = 0; l < tile_C::ne/2; ++l) { + const int i = i0 + n*tile_C::I + tile_C::get_i(2*l); + +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 4) { + const int k0 = k00 + k01; + + dA[n][l][k01/4] = x_df[i*MMQ_MMA_TILE_X_K_Q3_K + k0/4]; + } + } + } + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += ntx*tile_C::J) { +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR3_K*VDR_Q3_K_Q8_1_MMQ) { + tile_B B[2]; + float dB[tile_C::ne/2]; + + // Here load_generic is faster than load_ldmatrix. + load_generic(B[0], y_qs + j0*MMQ_TILE_Y_K + (k01 + 0), MMQ_TILE_Y_K); + load_generic(B[1], y_qs + j0*MMQ_TILE_Y_K + (k01 + tile_B::J), MMQ_TILE_Y_K); + +#pragma unroll + for (int l = 0; l < tile_C::ne/2; ++l) { + const int j = j0 + tile_C::get_j(l); + + dB[l] = y_df[j*MMQ_TILE_Y_K + k01/QI8_1]; + } + +#pragma unroll + for (int n = 0; n < ntx; ++n) { + tile_C C[2]; + mma(C[0], A[n][k01/4 + 0], B[0]); + mma(C[1], A[n][k01/4 + 1], B[1]); + +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + sum[(j0/tile_C::J + n)*tile_C::ne + l] += dB[l%2]*(C[0].x[l]*dA[n][l/2][k01/4 + 0] + C[1].x[l]*dA[n][l/2][k01/4 + 1]); + } + } + } + } +#else + GGML_UNUSED_VARS(x, y, sum, k00); + NO_DEVICE_CODE; +#endif // AMD_MFMA_AVAILABLE || AMD_WMMA_AVAILABLE +} + +template static __device__ __forceinline__ void load_tiles_q2_K( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int nwarps = mmq_get_nwarps_device(); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + int * x_qs = (int *) x_tile; + half2 * x_dm = (half2 *) (x_qs + 2*MMQ_TILE_NE_K); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q2_K, mmq_y); + int * x_qs = (int *) x_tile; + half2 * x_dm = (half2 *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = MMQ_ITER_K / (4 * QR2_K); + constexpr int nrows = ggml_cuda_get_physical_warp_size() / threads_per_row; + const int kqsx = threadIdx.x % threads_per_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nrows*nwarps) { + int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row; + + if (need_check) { + i = min(i, i_max); + } + + const block_q2_K * bxi = (const block_q2_K *) x + kbx0 + i*stride; + + const int x_ql_0 = get_int_b2(bxi->qs, kqsx); + +#pragma unroll + for (int l = 0; l < QR2_K; ++l) { + const int k = (kqsx/8)*32 + l*8 + kqsx % 8; + + const int x_qs_k = (x_ql_0 >> (2*l)) & 0x03030303; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*MMQ_MMA_TILE_X_K_Q2_K + k] = x_qs_k; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + k] = x_qs_k; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + + const int sc_m = bxi->scales[kqsx]; +#ifdef FAST_FP16_AVAILABLE + const half2 x_dm_ik = __hmul2(bxi->dm, make_half2(sc_m & 0x0F, sc_m >> 4)); +#else + const float2 bxi_dmf = __half22float2(bxi->dm); + const half2 x_dm_ik = make_half2(bxi_dmf.x*(sc_m & 0x0F), bxi_dmf.y*(sc_m >> 4)); +#endif // FAST_FP16_AVAILABLE + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_dm[i*MMQ_MMA_TILE_X_K_Q2_K + kqsx] = x_dm_ik; +#else + x_dm[i*(MMQ_TILE_NE_K + 1) + kqsx] = x_dm_ik; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +template +static __device__ __forceinline__ void vec_dot_q2_K_q8_1_dp4a( + const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q2_K, mmq_y); + const int * x_qs = (const int *) x; + const half2 * x_dm = (const half2 *) x_qs + txs.qs; + const int * y_qs = (const int *) y + 4; + const half2 * y_ds = (const half2 *) y; + + float2 y_df[mmq_x/nwarps]; +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += nwarps) { + const int j = j0 + threadIdx.y; + + y_df[j0/nwarps] = __half22float2(y_ds[j*MMQ_TILE_Y_K]); + } + +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K/2; k01 += QR2_K*VDR_Q2_K_Q8_1_MMQ) { + const int k0 = k00 + k01; + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += nwarps) { + const int j = j0 + threadIdx.y; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += warp_size) { + const int i = i0 + threadIdx.x; + + constexpr int ns = 2; + sum[j0/nwarps*mmq_y/warp_size + i0/warp_size] += vec_dot_q2_K_q8_1_impl_mmq( + &x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01], + &x_dm[i*(MMQ_TILE_NE_K + 1) + k0/4], k01 < MMQ_TILE_NE_K/2 ? y_df[j0/nwarps].x : y_df[j0/nwarps].y, + &y_ds[j*MMQ_TILE_Y_K + (1 + k01/QI8_1)]); + } + } + } + + // Some compilers fail to unroll the loop over k01 if there is a conditional statement for ns in the inner loop. + // As a workaround 2 separate loops are used instead. +#pragma unroll + for (int k01 = MMQ_TILE_NE_K/2; k01 < MMQ_TILE_NE_K; k01 += QR2_K*VDR_Q2_K_Q8_1_MMQ) { + const int k0 = k00 + k01; + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += nwarps) { + const int j = j0 + threadIdx.y; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += warp_size) { + const int i = i0 + threadIdx.x; + + constexpr int ns = 1; + sum[j0/nwarps*mmq_y/warp_size + i0/warp_size] += vec_dot_q2_K_q8_1_impl_mmq( + &x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01], + &x_dm[i*(MMQ_TILE_NE_K + 1) + k0/4], k01 < MMQ_TILE_NE_K/2 ? y_df[j0/nwarps].x : y_df[j0/nwarps].y, + &y_ds[j*MMQ_TILE_Y_K + (1 + k01/QI8_1)]); + } + } + } +} + +template +static __device__ __forceinline__ void vec_dot_q2_K_q8_1_mma( + const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { +#if defined(AMD_MFMA_AVAILABLE) + constexpr data_layout input_layout = get_input_data_layout(); + typedef tile<16, 8, int, input_layout> tile_A; + typedef tile<16, 8, int, input_layout> tile_B; + typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C; + typedef tile<64, 2, int, input_layout> tile_load; + + constexpr int granularity = mmq_get_granularity_device(mmq_x); + constexpr int rows_per_warp = granularity; + constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. + + y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); + + const int * x_qs = (const int *) x; + const half2 * x_dm = (const half2 *) x_qs + MMQ_TILE_NE_K*2; + const int * y_qs = (const int *) y + 4; + const half2 * y_ds = (const half2 *) y; + + const int i0 = (threadIdx.y / ntx) * rows_per_warp; + + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 4) { + const int k0 = k00 + k01; + + tile_A A[ntx]; +#pragma unroll + for (int n = 0; n < ntx; ++n) { + load_generic(((tile_load *) A)[n], x_qs + (i0 + n*tile_A::I)*MMQ_MMA_TILE_X_K_Q2_K + k0, MMQ_MMA_TILE_X_K_Q2_K); + } + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += ntx*tile_C::J) { + tile_B B[1]; + load_generic(((tile_load *) B)[0], y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); + + const int j = j0 + tile_C::get_j(0); + const float dB = (k01 < MMQ_TILE_NE_K/2) ? __half22float2(y_ds[j*MMQ_TILE_Y_K]).x/2 : __half22float2(y_ds[j*MMQ_TILE_Y_K]).y/2; + const float sB = (k01 >= MMQ_TILE_NE_K * 3/4) ? 0 + : (((k01/4)%2) ? __half22float2(y_ds[j*MMQ_TILE_Y_K + (1 + k01/QI8_1)]).y + : __half22float2(y_ds[j*MMQ_TILE_Y_K + (1 + k01/QI8_1)]).x); + + tile_C Cm; + if (k01 >= MMQ_TILE_NE_K * 3/4) { + tile_A A1; + A1.x[0] = 0x01010101; + A1.x[1] = 0x01010101; + mma(Cm, A1, B[0]); + } + +#pragma unroll + for (int n = 0; n < ntx; ++n) { + tile_C Cd; + mma(Cd, A[n], B[0]); + +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + const int i = i0 + n*tile_C::I + tile_C::get_i(l); + const float2 dm = __half22float2(x_dm[i*MMQ_MMA_TILE_X_K_Q2_K + k0/4]); + float tmp = Cd.x[l]*dm.x; + if (k01 >= MMQ_TILE_NE_K * 3/4) { + tmp -= Cm.x[l]*dm.y; + } + sum[(j0/tile_C::J + n)*tile_C::ne + l] += tmp*dB; + sum[(j0/tile_C::J + n)*tile_C::ne + l] -= dm.y*sB; + } + } + } + } +#elif defined(AMD_WMMA_AVAILABLE) //wmma instructions can handle 16x4 tiles, does not require loading 64x2 tiles + constexpr data_layout input_layout = get_input_data_layout(); + typedef tile<16, 4, int, input_layout> tile_A; + typedef tile<16, 4, int, input_layout> tile_B; + typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C; + + constexpr int granularity = mmq_get_granularity_device(mmq_x); + constexpr int rows_per_warp = granularity; + constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. + + y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); + + const int * x_qs = (const int *) x; + const half2 * x_dm = (const half2 *) x_qs + MMQ_TILE_NE_K*2; + const int * y_qs = (const int *) y + 4; + const half2 * y_ds = (const half2 *) y; + + const int i0 = (threadIdx.y / ntx) * rows_per_warp; + + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 4) { + const int k0 = k00 + k01; + + tile_A A[ntx]; +#pragma unroll + for (int n = 0; n < ntx; ++n) { + load_generic(A[n], x_qs + (i0 + n*tile_A::I)*MMQ_MMA_TILE_X_K_Q2_K + k0, MMQ_MMA_TILE_X_K_Q2_K); + } + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += ntx*tile_C::J) { + tile_B B; + load_generic(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); + + const int j = j0 + tile_C::get_j(0); + const float dB = (k01 < MMQ_TILE_NE_K/2) ? __half22float2(y_ds[j*MMQ_TILE_Y_K]).x : __half22float2(y_ds[j*MMQ_TILE_Y_K]).y; + const float sB = (k01 >= MMQ_TILE_NE_K * 3/4) ? 0 + : (((k01/4)%2) ? __half22float2(y_ds[j*MMQ_TILE_Y_K + (1 + k01/QI8_1)]).y + : __half22float2(y_ds[j*MMQ_TILE_Y_K + (1 + k01/QI8_1)]).x); + + tile_C Cm; + if (k01 >= MMQ_TILE_NE_K * 3/4) { + tile_A A1; +#pragma unroll + for (int l = 0; l < tile_A::ne; ++l) { + A1.x[l] = 0x01010101; + } + mma(Cm, A1, B); + } + +#pragma unroll + for (int n = 0; n < ntx; ++n) { + tile_C Cd; + mma(Cd, A[n], B); + +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + const int i = i0 + n*tile_C::I + tile_C::get_i(l); + const float2 dm = __half22float2(x_dm[i*MMQ_MMA_TILE_X_K_Q2_K + k0/4]); + float tmp = Cd.x[l]*dm.x; + if (k01 >= MMQ_TILE_NE_K * 3/4) { + tmp -= Cm.x[l]*dm.y; + } + sum[(j0/tile_C::J + n)*tile_C::ne + l] += tmp*dB; + sum[(j0/tile_C::J + n)*tile_C::ne + l] -= dm.y*sB; + } + } + } + } +#elif defined(TURING_MMA_AVAILABLE) + + typedef tile<16, 4, int> tile_A; + typedef tile<16, 8, int> tile_A_8; + typedef tile< 8, 4, int> tile_B; + typedef tile<16, 8, int> tile_C; + + constexpr int granularity = mmq_get_granularity_device(mmq_x); + constexpr int rows_per_warp = 2 * granularity; + constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. + + y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); + + const int * x_qs = (const int *) x; + const half2 * x_dm = (const half2 *) x_qs + MMQ_TILE_NE_K*2; + const int * y_qs = (const int *) y + 4; + const half2 * y_ds = (const half2 *) y; + + const int i0 = (threadIdx.y / ntx) * (ntx*tile_A::I); + + tile_A A[ntx][8]; + float dA[ntx][tile_C::ne/2][8]; + float mA[ntx][tile_C::ne/2][8]; + +#pragma unroll + for (int n = 0; n < ntx; ++n) { +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) { + const int k0 = k00 + k01; + + load_ldmatrix(((tile_A_8 *) A[n])[k01/QI8_1], x_qs + (i0 + n*tile_A::I)*MMQ_MMA_TILE_X_K_Q2_K + k0, MMQ_MMA_TILE_X_K_Q2_K); + } + } + +#pragma unroll + for (int n = 0; n < ntx; ++n) { +#pragma unroll + for (int l = 0; l < tile_C::ne/2; ++l) { + const int i = i0 + n*tile_C::I + tile_C::get_i(2*l); + +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1/2) { + const int k0 = k00 + k01; + + const float2 dm = __half22float2(x_dm[i*MMQ_MMA_TILE_X_K_Q2_K + k0/(QI8_1/2)]); + + dA[n][l][k01/(QI8_1/2)] = dm.x; + mA[n][l][k01/(QI8_1/2)] = dm.y; + } + } + } + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += ntx*tile_C::J) { + float2 dB[tile_C::ne/2]; + +#pragma unroll + for (int l = 0; l < tile_C::ne/2; ++l) { + const int j = j0 + tile_C::get_j(l); + + dB[l] = __half22float2(y_ds[j*MMQ_TILE_Y_K]); + } + +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_1) { + tile_B B[2]; + + // Here load_generic is faster than load_ldmatrix. + load_generic(B[0], y_qs + j0*MMQ_TILE_Y_K + (k01 + 0), MMQ_TILE_Y_K); + load_generic(B[1], y_qs + j0*MMQ_TILE_Y_K + (k01 + tile_B::J), MMQ_TILE_Y_K); + + tile_C Cm[2]; + if (k01 >= MMQ_TILE_NE_K * 3/4) { + tile_A A1; + A1.x[0] = 0x01010101; + A1.x[1] = 0x01010101; + mma(Cm[0], A1, B[0]); + mma(Cm[1], A1, B[1]); + } + +#pragma unroll + for (int n = 0; n < ntx; ++n) { + tile_C Cd[2]; + + mma(Cd[0], A[n][k01/4 + 0], B[0]); + mma(Cd[1], A[n][k01/4 + 1], B[1]); + +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + float tmp = Cd[0].x[l]*dA[n][l/2][k01/4 + 0] + Cd[1].x[l]*dA[n][l/2][k01/4 + 1]; + if (k01 >= MMQ_TILE_NE_K * 3/4) { + tmp -= Cm[0].x[l]*mA[n][l/2][k01/4 + 0] + Cm[1].x[l]*mA[n][l/2][k01/4 + 1]; + } + sum[(j0/tile_C::J + n)*tile_C::ne + l] += tmp*(k01 < MMQ_TILE_NE_K/2 ? dB[l%2].x : dB[l%2].y); + } + } + } + +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K * 3/4; k01 += QI8_1) { + float2 sB[tile_C::ne/2]; + +#pragma unroll + for (int l = 0; l < tile_C::ne/2; ++l) { + const int j = j0 + tile_C::get_j(l); + + sB[l] = __half22float2(y_ds[j*MMQ_TILE_Y_K + (1 + k01/QI8_1)]); + } + +#pragma unroll + for (int n = 0; n < ntx; ++n) { +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + sum[(j0/tile_C::J + n)*tile_C::ne + l] -= mA[n][l/2][k01/4 + 0]*sB[l%2].x; + sum[(j0/tile_C::J + n)*tile_C::ne + l] -= mA[n][l/2][k01/4 + 1]*sB[l%2].y; + } + } + } + } +#else + GGML_UNUSED_VARS(x, y, sum, k00); + NO_DEVICE_CODE; +#endif // AMD_MFMA_AVAILABLE || AMD_WMMA_AVAILABLE +} + +template static __device__ __forceinline__ void load_tiles_q3_K( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q3_K, mmq_y); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); + int * x_sc = (int *) (x_df + txs.dm); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) + + constexpr int threads_per_row = MMQ_ITER_K / (4 * QR3_K); + constexpr int nrows = warp_size / threads_per_row; + const int kqsx = threadIdx.x % threads_per_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nrows*nwarps) { + int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row; + + if (need_check) { + i = min(i, i_max); + } + + const block_q3_K * bxi = (const block_q3_K *) x + kbx0 + i*stride; + + const int x_ql_0 = get_int_b2(bxi->qs, kqsx); + const int x_qh_0 = get_int_b2(bxi->hmask, kqsx % (QI3_K/2)) >> (4 * (kqsx / (QI3_K/2))); + +#pragma unroll + for (int l = 0; l < QR3_K; ++l) { + const int k = (kqsx/8)*32 + l*8 + kqsx % 8; + + const int x_ql_k = (x_ql_0 >> (2*l)) & 0x03030303; + const int x_qh_k = ((x_qh_0 >> l) << 2) & 0x04040404; + + const int x_qs_k = __vsubss4(x_ql_k | x_qh_k, 0x04040404); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*MMQ_MMA_TILE_X_K_Q3_K + k] = x_qs_k; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + k] = x_qs_k; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + } + + constexpr int rows_per_warp = warp_size / 4; +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps*rows_per_warp) { + int i = i0 + threadIdx.y*rows_per_warp + threadIdx.x/4; + + if (need_check) { + i = min(i, i_max); + } + + const block_q3_K * bxi = (const block_q3_K *) x + kbx0 + i*stride; + + const int ksc = threadIdx.x % 4; + + const int ksc_low = ksc % (QI3_K/8); + const int shift_low = 4 * (ksc / (QI3_K/8)); + const int sc_low = (get_int_b2(bxi->scales, ksc_low) >> shift_low) & 0x0F0F0F0F; + + const int ksc_high = QI3_K/8; + const int shift_high = 2 * ksc; + const int sc_high = ((get_int_b2(bxi->scales, ksc_high) >> shift_high) << 4) & 0x30303030; + + const int sc = __vsubss4(sc_low | sc_high, 0x20202020); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + const int8_t * sc8 = (const int8_t *) ≻ + const float d = bxi->d; + +#pragma unroll + for (int l = 0; l < int(sizeof(int)); ++l) { + x_df[i*MMQ_MMA_TILE_X_K_Q3_K + sizeof(int)*ksc + l] = d*sc8[l]; + } +#else + x_sc[i*(MMQ_TILE_NE_K/8) + i/8 + ksc] = sc; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + +#if !(defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE)) +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps*warp_size) { + int i = (i0 + threadIdx.y*warp_size + threadIdx.x) % mmq_y; + + if (need_check) { + i = min(i, i_max); + } + + const block_q3_K * bxi = (const block_q3_K *) x + kbx0 + i*stride; + + x_df[i] = bxi->d; + } +#endif // !(defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE)) || defined(AMD_WMMA_AVAILABLE) +} + +template +static __device__ __forceinline__ void vec_dot_q3_K_q8_1_dp4a( + const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q3_K, mmq_y); + const int * x_qs = (const int *) x; + const float * x_df = (const float *) x_qs + txs.qs; + const int * x_sc = (const int *) x_df + txs.dm; + const int * y_qs = (const int *) y + 4; + const float * y_df = (const float *) y; + +// #pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR3_K*VDR_Q3_K_Q8_1_MMQ) { + const int k0 = k00 + k01; + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += nwarps) { + const int j = j0 + threadIdx.y; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += warp_size) { + const int i = i0 + threadIdx.x; + + const int8_t * scales = ((const int8_t *) (x_sc + i*(MMQ_TILE_NE_K/8) + i/8)) + k0/4; + + sum[j0/nwarps*mmq_y/warp_size + i0/warp_size] += vec_dot_q3_K_q8_1_impl_mmq( + &x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01], scales, + x_df[i], y_df[j*MMQ_TILE_Y_K + k01/QI8_1]); + } + } + } +} + +static __device__ __forceinline__ int unpack_scales_q45_K(const int * scales, const int ksc) { + // scale arrangement after the following two lines: + // - ksc == 0: sc0, sc1, sc2, sc3 + // - ksc == 1: sc4, sc5, sc6, sc7 + // - ksc == 2: m0, m1, m2, m3 + // - ksc == 3: m4, m5, m6, m7 + return ((scales[(ksc%2) + (ksc!=0)] >> (4 * (ksc & (ksc/2)))) & 0x0F0F0F0F) | // lower 4 bits + ((scales[ksc/2] >> (2 * (ksc % 2))) & 0x30303030); // upper 2 bits +} + +template static __device__ __forceinline__ void load_tiles_q4_K( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + int * x_qs = (int *) x_tile; + half2 * x_dm = (half2 *) (x_qs + 2*MMQ_TILE_NE_K); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_K, mmq_y); + int * x_qs = (int *) x_tile; + half2 * x_dm = (half2 *) (x_qs + txs.qs); + int * x_sc = (int *) (x_dm + txs.dm); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = MMQ_ITER_K / (4 * QR4_K); + constexpr int nrows = warp_size / threads_per_row; + const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nrows*nwarps) { + int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row); + + if (need_check) { + i = min(i, i_max); + } + + const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride; + const int qs0 = get_int_b4(bxi->qs, txi); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*MMQ_MMA_TILE_X_K_Q8_1 + 16*(txi/8) + txi % 8 + 0] = (qs0 >> 0) & 0x0F0F0F0F; + x_qs[i*MMQ_MMA_TILE_X_K_Q8_1 + 16*(txi/8) + txi % 8 + 8] = (qs0 >> 4) & 0x0F0F0F0F; +#else + x_qs[i*(MMQ_TILE_NE_K + 1) + txi] = qs0; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr int rows_per_warp = warp_size / 2; +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps*rows_per_warp) { +#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + // Need if on AMD instead of % because warp_size == 64 + // This causes double work and throughput loss (MI300X) + // H100 loses about 100 t/s with 'if' condition over '%' + int i = i0 + threadIdx.y*rows_per_warp + threadIdx.x/2; + if (i < mmq_y) { +#else + int i = (i0 + threadIdx.y*rows_per_warp + threadIdx.x/2) % mmq_y; + { +#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + if (need_check) { + i = min(i, i_max); + } + + const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride; + + const int * scales = (const int *) bxi->scales; + const int ksc = threadIdx.x % 2; + + const int sc32 = unpack_scales_q45_K(scales, ksc + 0); + const int m32 = unpack_scales_q45_K(scales, ksc + 2); + + const uint8_t * sc8 = (const uint8_t *) &sc32; + const uint8_t * m8 = (const uint8_t *) &m32; + + const half2 dm = bxi->dm * make_half2(1.0f, -1.0f); + + #pragma unroll + for (int l = 0; l < sizeof(int); ++l) { + x_dm[i*MMQ_MMA_TILE_X_K_Q8_1 + sizeof(int)*ksc + l] = dm*make_half2(sc8[l], m8[l]); + } + } + } +#else +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps*warp_size) { + int i = (i0 + threadIdx.y*warp_size + threadIdx.x) % mmq_y; + + if (need_check) { + i = min(i, i_max); + } + + const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride; + + x_dm[i] = bxi->dm; + } + constexpr int rows_per_warp = warp_size / 4; +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps*rows_per_warp) { + int i = (i0 + threadIdx.y*rows_per_warp + threadIdx.x/(MMQ_TILE_NE_K/8)) % mmq_y; + + if (need_check) { + i = min(i, i_max); + } + + const block_q4_K * bxi = (const block_q4_K *) x + kbx0 + i*stride + (threadIdx.x % (MMQ_TILE_NE_K/8)) / (QI4_K/8); + + const int * scales = (const int *) bxi->scales; + + const int ksc = threadIdx.x % (MMQ_TILE_NE_K/8); + const int scales8 = unpack_scales_q45_K(scales, ksc); + + x_sc[i*(MMQ_TILE_NE_K/8) + i/8 + ksc] = scales8; + } +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) +} + +template +static __device__ __forceinline__ void vec_dot_q4_K_q8_1_dp4a( + const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q4_K, mmq_y); + const int * x_qs = (const int *) x; + const half2 * x_dm = (const half2 *) x_qs + txs.qs; + const int * x_sc = (const int *) x_dm + txs.dm; + const int * y_qs = (const int *) y + 4; + const half2 * y_ds = (const half2 *) y; + +// #pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR4_K*VDR_Q4_K_Q8_1_MMQ) { + const int k0 = k00 + k01; + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += nwarps) { + const int j = j0 + threadIdx.y; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += warp_size) { + const int i = i0 + threadIdx.x; + + const uint8_t * sc = (const uint8_t *) &x_sc[i * (MMQ_TILE_NE_K/8) + i/8 + k0/32] + 2*(k01/16); + + sum[j0/nwarps*mmq_y/warp_size + i0/warp_size] += vec_dot_q4_K_q8_1_impl_mmq( + &x_qs[i*(MMQ_TILE_NE_K + 1) + k0/2], &y_qs[j*MMQ_TILE_Y_K + k01], sc, sc+8, + x_dm[i], &y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]); + } + } + } +} + +template static __device__ __forceinline__ void load_tiles_q5_K( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + int * x_qs = (int *) x_tile; + half2 * x_dm = (half2 *) (x_qs + MMQ_TILE_NE_K*2); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q5_K, mmq_y); + int * x_qs = (int *) x_tile; + half2 * x_dm = (half2 *) (x_qs + txs.qs); + int * x_sc = (int *) (x_dm + txs.dm); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) + + constexpr int threads_per_row = MMQ_ITER_K / (4 * QR5_K); + constexpr int nrows = warp_size / threads_per_row; + const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nrows*nwarps) { + int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row); + + if (need_check) { + i = min(i, i_max); + } + + const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride; + const int ky = QR5_K*txi; + + const int ql = get_int_b4(bxi->qs, txi); + const int ql0 = (ql >> 0) & 0x0F0F0F0F; + const int ql1 = (ql >> 4) & 0x0F0F0F0F; + + const int qh = get_int_b4(bxi->qh, txi % (QI5_K/4)); + const int qh0 = ((qh >> (2 * (txi / (QI5_K/4)) + 0)) << 4) & 0x10101010; + const int qh1 = ((qh >> (2 * (txi / (QI5_K/4)) + 1)) << 4) & 0x10101010; + + const int kq0 = ky - ky % (QI5_K/2) + txi % (QI5_K/4) + 0; + const int kq1 = ky - ky % (QI5_K/2) + txi % (QI5_K/4) + QI5_K/4; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*MMQ_MMA_TILE_X_K_Q8_1 + kq0] = ql0 | qh0; + x_qs[i*MMQ_MMA_TILE_X_K_Q8_1 + kq1] = ql1 | qh1; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + kq0] = ql0 | qh0; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + kq1] = ql1 | qh1; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr int rows_per_warp = warp_size / 2; +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps*rows_per_warp) { +#if defined(AMD_MFMA_AVAILABLE) + // Need if on AMD instead of % because warp_size == 64 + // This causes double work and throughput loss (MI300X) + // H100 loses about 100 t/s with 'if' condition over '%' + int i = i0 + threadIdx.y*rows_per_warp + threadIdx.x/2; + if (i < mmq_y) { +#else + int i = (i0 + threadIdx.y*rows_per_warp + threadIdx.x/2) % mmq_y; + { +#endif // defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + if (need_check) { + i = min(i, i_max); + } + + const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride; + + const int * scales = (const int *) bxi->scales; + const int ksc = threadIdx.x % 2; + + const int sc32 = unpack_scales_q45_K(scales, ksc + 0); + const int m32 = unpack_scales_q45_K(scales, ksc + 2); + + const uint8_t * sc8 = (const uint8_t *) &sc32; + const uint8_t * m8 = (const uint8_t *) &m32; + + const half2 dm = bxi->dm * make_half2(1.0f, -1.0f); + +#pragma unroll + for (int l = 0; l < int(sizeof(int)); ++l) { + x_dm[i*MMQ_MMA_TILE_X_K_Q8_1 + sizeof(int)*ksc + l] = dm*make_half2(sc8[l], m8[l]); + } + } + } +#else +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps*warp_size) { + int i = (i0 + threadIdx.y*warp_size + threadIdx.x) % mmq_y; + + if (need_check) { + i = min(i, i_max); + } + + const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride; + + x_dm[i] = bxi->dm; + } + + constexpr int rows_per_warp = warp_size / 4; +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps*rows_per_warp) { + int i = (i0 + threadIdx.y*rows_per_warp + threadIdx.x/(MMQ_TILE_NE_K/8)) % mmq_y; + + if (need_check) { + i = min(i, i_max); + } + + const block_q5_K * bxi = (const block_q5_K *) x + kbx0 + i*stride; + + const int * scales = (const int *) bxi->scales; + + const int ksc = threadIdx.x % (MMQ_TILE_NE_K/8); + const int scales8 = unpack_scales_q45_K(scales, ksc); + + x_sc[i*(MMQ_TILE_NE_K/8) + i/8 + ksc] = scales8; + } +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) +} + +template +static __device__ __forceinline__ void vec_dot_q5_K_q8_1_dp4a( + const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q5_K, mmq_y); + const int * x_qs = (const int *) x; + const half2 * x_dm = (const half2 *) x_qs + txs.qs; + const int * x_sc = (const int *) x_dm + txs.dm; + const int * y_qs = (const int *) y + 4; + const half2 * y_ds = (const half2 *) y; + +// #pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR5_K*VDR_Q5_K_Q8_1_MMQ) { + const int k0 = k00 + k01; + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += nwarps) { + const int j = j0 + threadIdx.y; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += warp_size) { + const int i = i0 + threadIdx.x; + + const uint8_t * sc = ((const uint8_t *) &x_sc[i * (MMQ_TILE_NE_K/8) + i/8 + k00/32]) + 2*(k01/16); + + sum[j0/nwarps*mmq_y/warp_size + i0/warp_size] += vec_dot_q5_K_q8_1_impl_mmq( + &x_qs[i*(QR5_K*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01], sc, sc+8, + x_dm[i], &y_ds[j*MMQ_TILE_Y_K + k01/QI8_1]); + } + } + } +} + +template static __device__ __forceinline__ void load_tiles_q6_K( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2); + int * x_sc = (int *) (x_df + MMQ_TILE_NE_K/QI6_K); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q6_K, mmq_y); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); + int * x_sc = (int *) (x_df + txs.dm); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = MMQ_ITER_K / (4 * QR6_K); + constexpr int nrows = warp_size / threads_per_row; + const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nrows*nwarps) { + int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row); + + if (need_check) { + i = min(i, i_max); + } + + const block_q6_K * bxi = (const block_q6_K *) x + kbx0 + i*stride; + + const int ql = get_int_b2(bxi->ql, txi); + const int ql0 = (ql >> 0) & 0x0F0F0F0F; + const int ql1 = (ql >> 4) & 0x0F0F0F0F; + + const int qh = get_int_b2(bxi->qh, (QI6_K/4) * (txi / (QI6_K/2)) + txi % (QI6_K/4)); + const int qh0 = ((qh >> ((txi & 0x08) >> 2)) << 4) & 0x30303030; + const int qh1 = (qh >> ((txi & 0x08) >> 2)) & 0x30303030; + + const int kq0 = 2*txi - txi % (QI6_K/2) + 0; + const int kq1 = 2*txi - txi % (QI6_K/2) + QI6_K/2; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*MMQ_MMA_TILE_X_K_Q6_K + kq0] = __vsubss4(ql0 | qh0, 0x20202020); + x_qs[i*MMQ_MMA_TILE_X_K_Q6_K + kq1] = __vsubss4(ql1 | qh1, 0x20202020); +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + kq0] = __vsubss4(ql0 | qh0, 0x20202020); + x_qs[i*(2*MMQ_TILE_NE_K + 1) + kq1] = __vsubss4(ql1 | qh1, 0x20202020); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps*warp_size) { + int i = (i0 + threadIdx.y*warp_size + threadIdx.x) % mmq_y; + + if (need_check) { + i = min(i, i_max); + } + + const block_q6_K * bxi = (const block_q6_K *) x + kbx0 + i*stride; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_df[i*MMQ_MMA_TILE_X_K_Q6_K] = bxi->d; +#else + x_df[i*(MMQ_TILE_NE_K/QI6_K) + i/QI6_K] = bxi->d; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + + constexpr int rows_per_warp = warp_size / 4; +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps*rows_per_warp) { + int i = (i0 + threadIdx.y*rows_per_warp + threadIdx.x/(MMQ_TILE_NE_K/8)) % mmq_y; + + if (need_check) { + i = min(i, i_max); + } + + const block_q6_K * bxi = (const block_q6_K *) x + kbx0 + i*stride + (threadIdx.x % (MMQ_TILE_NE_K/8)) / 4; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_sc[i*MMQ_MMA_TILE_X_K_Q6_K + threadIdx.x%4] = get_int_b2(bxi->scales, threadIdx.x % (MMQ_TILE_NE_K/8)); +#else + x_sc[i*(MMQ_TILE_NE_K/8) + i/8 + threadIdx.x%(MMQ_TILE_NE_K/8)] = get_int_b2(bxi->scales, threadIdx.x%(QI6_K/8)); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +template +static __device__ __forceinline__ void vec_dot_q6_K_q8_1_dp4a( + const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_Q6_K, mmq_y); + const int * x_qs = (const int *) x; + const float * x_df = (const float *) x_qs + txs.qs; + const int * x_sc = (const int *) x_df + txs.dm; + const int * y_qs = (const int *) y + 4; + const float * y_df = (const float *) y; + +// #pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QR6_K*VDR_Q6_K_Q8_1_MMQ) { + const int k0 = k00 + k01; + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += nwarps) { + const int j = j0 + threadIdx.y; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += warp_size) { + const int i = i0 + threadIdx.x; + + const int8_t * sc = ((const int8_t *) &x_sc[i * (MMQ_TILE_NE_K/8) + i/8 + k0/16]); + + sum[j0/nwarps*mmq_y/warp_size + i0/warp_size] += vec_dot_q6_K_q8_1_impl_mmq( + &x_qs[i*(QR6_K*MMQ_TILE_NE_K + 1) + k0], &y_qs[j*MMQ_TILE_Y_K + k01], sc, + x_df[i*(MMQ_TILE_NE_K/QI6_K) + i/QI6_K], &y_df[j*MMQ_TILE_Y_K + k01/QI8_1]); + } + } + } +} + +template +static __device__ __forceinline__ void vec_dot_q6_K_q8_1_mma( + const int * __restrict__ x, const int * __restrict__ y, float * __restrict__ sum, const int k00) { +#if defined(AMD_MFMA_AVAILABLE) + constexpr data_layout input_layout = get_input_data_layout(); + typedef tile<16, 8, int, input_layout> tile_A; + typedef tile<16, 8, int, input_layout> tile_B; + typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C; + typedef tile<64, 2, int, input_layout> tile_load; + + constexpr int granularity = mmq_get_granularity_device(mmq_x); + constexpr int rows_per_warp = granularity; + constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. + + y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); + + const int * x_qs = (const int *) x; + const float * x_df = (const float *) x_qs + MMQ_TILE_NE_K*2; + const int * x_sc = (const int *) x_df + MMQ_TILE_NE_K/QI6_K; + const int * y_qs = (const int *) y + 4; + const float * y_df = (const float *) y; + + const int i0 = (threadIdx.y / ntx) * rows_per_warp; + + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 4) { + const int k0 = k00 + k01; + + tile_A A[ntx]; +#pragma unroll + for (int n = 0; n < ntx; ++n) { + load_generic(((tile_load *) A)[n], x_qs + (i0 + n*tile_A::I)*MMQ_MMA_TILE_X_K_Q6_K + k0, MMQ_MMA_TILE_X_K_Q6_K); + } + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += ntx*tile_C::J) { + tile_B B[1]; + load_generic(((tile_load *) B)[0], y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); + + const int j = j0 + tile_C::get_j(0); + const float dB = y_df[j*MMQ_TILE_Y_K + k01/QI8_1] / 2; + +#pragma unroll + for (int n = 0; n < ntx; ++n) { + tile_C C; + mma(C, A[n], B[0]); + +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + const int i = i0 + n*tile_C::I + tile_C::get_i(l); + const int8_t * sc = (const int8_t *) (x_sc + i*MMQ_MMA_TILE_X_K_Q6_K + k00/16); + sum[(j0/tile_C::J + n)*tile_C::ne + l] += C.x[l] * sc[k01/4] * x_df[i*MMQ_MMA_TILE_X_K_Q6_K] * dB; + } + } + } + } +#elif defined(AMD_WMMA_AVAILABLE) //wmma instructions can handle 16x4 tiles, does not require loading 64x2 tiles + constexpr data_layout input_layout = get_input_data_layout(); + typedef tile<16, 4, int, input_layout> tile_A; + typedef tile<16, 4, int, input_layout> tile_B; + typedef tile<16, 16, int, DATA_LAYOUT_J_MAJOR> tile_C; + + constexpr int granularity = mmq_get_granularity_device(mmq_x); + constexpr int rows_per_warp = granularity; + constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. + + y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); + + const int * x_qs = (const int *) x; + const float * x_df = (const float *) x_qs + MMQ_TILE_NE_K*2; + const int * x_sc = (const int *) x_df + MMQ_TILE_NE_K/QI6_K; + const int * y_qs = (const int *) y + 4; + const float * y_df = (const float *) y; + + const int i0 = (threadIdx.y / ntx) * rows_per_warp; + + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 4) { + const int k0 = k00 + k01; + + tile_A A[ntx]; +#pragma unroll + for (int n = 0; n < ntx; ++n) { + load_generic(A[n], x_qs + (i0 + n*tile_A::I)*MMQ_MMA_TILE_X_K_Q6_K + k0, MMQ_MMA_TILE_X_K_Q6_K); + } + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += ntx*tile_C::J) { + tile_B B; + load_generic(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); + + const int j = j0 + tile_C::get_j(0); + const float dB = y_df[j*MMQ_TILE_Y_K + k01/QI8_1]; + +#pragma unroll + for (int n = 0; n < ntx; ++n) { + tile_C C; + mma(C, A[n], B); + +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + const int i = i0 + n*tile_C::I + tile_C::get_i(l); + const int8_t * sc = (const int8_t *) (x_sc + i*MMQ_MMA_TILE_X_K_Q6_K + k00/16); + sum[(j0/tile_C::J + n)*tile_C::ne + l] += C.x[l] * sc[k01/4] * x_df[i*MMQ_MMA_TILE_X_K_Q6_K] * dB; + } + } + } + } +#elif defined(TURING_MMA_AVAILABLE) + + typedef tile<16, 4, int> tile_A; + typedef tile< 8, 4, int> tile_B; + typedef tile<16, 8, int> tile_C; + + constexpr int granularity = mmq_get_granularity_device(mmq_x); + constexpr int rows_per_warp = 2 * granularity; + constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. + + y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); + + const int * x_qs = (const int *) x; + const float * x_df = (const float *) x_qs + MMQ_TILE_NE_K*2; + const int * x_sc = (const int *) x_df + MMQ_TILE_NE_K/QI6_K; + const int * y_qs = (const int *) y + 4; + const float * y_df = (const float *) y; + + const int i0 = (threadIdx.y / ntx) * (ntx*tile_A::I); + + tile_A A[ntx][8]; + int scA[ntx][tile_C::ne/2][8]; + float dA[ntx][tile_C::ne/2]; + +#pragma unroll + for (int n = 0; n < ntx; ++n) { +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 8) { + const int k0 = k00 + k01; + + load_ldmatrix(A[n][k01/4 + 0], x_qs + (i0 + n*tile_A::I)*MMQ_MMA_TILE_X_K_Q6_K + (k0 + 0), MMQ_MMA_TILE_X_K_Q6_K); + load_ldmatrix(A[n][k01/4 + 1], x_qs + (i0 + n*tile_A::I)*MMQ_MMA_TILE_X_K_Q6_K + (k0 + tile_A::J), MMQ_MMA_TILE_X_K_Q6_K); + } + +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 16) { + const int k0 = k00 + k01; + +#pragma unroll + for (int l = 0; l < tile_C::ne/2; ++l) { + const int i = i0 + n*tile_C::I + tile_C::get_i(2*l); + + const int sc_packed = x_sc[i*MMQ_MMA_TILE_X_K_Q6_K + k0/16]; + const int8_t * sc = (const int8_t *) &sc_packed; + +#pragma unroll + for (int ksc = 0; ksc < sizeof(int); ++ksc) { + scA[n][l][k01/4 + ksc] = sc[ksc]; + } + } + } + +#pragma unroll + for (int l = 0; l < tile_C::ne/2; ++l) { + const int i = i0 + n*tile_C::I + tile_C::get_i(2*l); + + dA[n][l] = x_df[i*MMQ_MMA_TILE_X_K_Q6_K]; + } + } + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += ntx*tile_C::J) { + float tmp[ntx][tile_C::ne] = {{0.0f}}; + +#pragma unroll + for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += 8) { + tile_B B[2]; + float dB[tile_C::ne/2]; + + // Here load_generic is faster than load_ldmatrix. + load_generic(B[0], y_qs + j0*MMQ_TILE_Y_K + 0 + k01, MMQ_TILE_Y_K); + load_generic(B[1], y_qs + j0*MMQ_TILE_Y_K + tile_B::J + k01, MMQ_TILE_Y_K); + +#pragma unroll + for (int l = 0; l < tile_C::ne/2; ++l) { + const int j = j0 + tile_C::get_j(l); + + dB[l] = y_df[j*MMQ_TILE_Y_K + k01/QI8_1]; + } + +#pragma unroll + for (int n = 0; n < ntx; ++n) { + tile_C C[2]; + mma(C[0], A[n][k01/4 + 0], B[0]); + mma(C[1], A[n][k01/4 + 1], B[1]); + +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + tmp[n][l] += (C[0].x[l]*scA[n][l/2][k01/4 + 0] + C[1].x[l]*scA[n][l/2][k01/4 + 1])*dB[l%2]; + } + } + } + +#pragma unroll + for (int n = 0; n < ntx; ++n) { +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + sum[(j0/tile_C::J + n)*tile_C::ne + l] += tmp[n][l]*dA[n][l/2]; + } + } + } +#else + GGML_UNUSED_VARS(x, y, sum, k00); + NO_DEVICE_CODE; +#endif // AMD_MFMA_AVAILABLE || AMD_WMMA_AVAILABLE +} + +template static __device__ __forceinline__ void load_tiles_iq4_nl( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ4_NL, mmq_y); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = MMQ_ITER_K / (4 * QR4_NL); + constexpr int nrows = warp_size / threads_per_row; + const int txi = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x; + const int kbx = txi / QI4_NL; + const int kqsx = txi % QI4_NL; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nrows*nwarps) { + int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row); + + if (need_check) { + i = min(i, i_max); + } + + const block_iq4_nl * bxi = (const block_iq4_nl *) x + kbx0 + i*stride + kbx; + + const int aux_q4 = get_int_b2(bxi->qs, kqsx); + const int2 v = get_int_from_table_16(aux_q4, kvalues_iq4nl); + const int k0 = kbx * (2 * QI4_NL) + kqsx; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + k0 + 0] = v.x; + x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + k0 + QI4_NL] = v.y; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + 0] = v.x; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + QI4_NL] = v.y; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + + constexpr int blocks_per_tile_x_row = MMQ_TILE_NE_K / QI4_NL; + constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row; + const int kbxd = threadIdx.x % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * rows_per_warp) { + int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / blocks_per_tile_x_row; + + if (need_check) { + i = min(i, i_max); + } + + const block_iq4_nl * bxi = (const block_iq4_nl *) x + kbx0 + i*stride + kbxd; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_df[i*MMQ_MMA_TILE_X_K_Q8_0 + kbxd] = __half2float(bxi->d); +#else + x_df[i*(MMQ_TILE_NE_K/QI4_NL) + i/QI4_NL + kbxd] = __half2float(bxi->d); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +template static __device__ __forceinline__ void load_tiles_iq2_xxs( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ2_XXS, mmq_y); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = (MMQ_ITER_K / (4 * QR2_XXS)) / 2; + constexpr int nrows = warp_size / threads_per_row; + const int kqsx = warp_size > threads_per_row ? threadIdx.x % threads_per_row : threadIdx.x; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * nrows) { + int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row; + + if (need_check) { + i = min(i, i_max); + } + + const block_iq2_xxs * bxi = (const block_iq2_xxs *) x + kbx0 + i*stride; + + const int q2 = get_int_b2(bxi->qs, 2*kqsx+0); + const uint8_t * aux8 = (const uint8_t *) &q2; + const uint32_t aux32 = get_int_b2(bxi->qs, 2*kqsx+1); + +#pragma unroll + for (int l = 0; l < QR2_XXS; ++l) { + const uint2 grid_pos = ((const uint2*)iq2xxs_grid)[aux8[l]]; + const uint32_t signs = unpack_ksigns(aux32 >> (7 * l)); + + const int signs0 = __vcmpne4(signs & 0x08040201, 0); + const int grid0 = __vsub4(grid_pos.x ^ signs0, signs0); + + const int signs1 = __vcmpne4(signs & 0x80402010, 0); + const int grid1 = __vsub4(grid_pos.y ^ signs1, signs1); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + 8*kqsx + (2*l + 0)] = grid0; + x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + 8*kqsx + (2*l + 1)] = grid1; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 0)] = grid0; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 1)] = grid1; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + + const int ls = aux32 >> 27 | 1; // (scale * 2 + 1) + const float d = bxi->d; +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_df[i*MMQ_MMA_TILE_X_K_Q8_0 + kqsx] = d * ls / 8; // (d * scale + d / 2) / 4 +#else + x_df[i*(MMQ_TILE_NE_K/4) + i/4 + kqsx] = d * ls / 8; // (d * scale + d / 2) / 4 +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +template static __device__ __forceinline__ void load_tiles_iq2_xs( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2); +#else + constexpr tile_x_sizes txs = MMQ_DP4A_TXS_Q8_0_16; + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = (MMQ_ITER_K / (4 * QR2_XS)) / 2; + constexpr int nrows = warp_size / threads_per_row; + const int kqsx = threadIdx.x % threads_per_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * nrows) { + int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row; + + if (need_check) { + i = min(i, i_max); + } + + const block_iq2_xs * bxi = (const block_iq2_xs *) x + kbx0 + i*stride; + + const int2 q2_packed = make_int2(get_int_b2(bxi->qs, 2*kqsx+0), get_int_b2(bxi->qs, 2*kqsx+1)); + const uint16_t * q2 = (const uint16_t *) &q2_packed; + + #pragma unroll + for (int l = 0; l < QR2_XS; ++l) { + const uint2 grid_pos = ((const uint2*)iq2xs_grid)[q2[l] & 0x1FF]; + const uint32_t signs = unpack_ksigns(q2[l] >> 9); + + const int signs0 = __vcmpne4(signs & 0x08040201, 0); + const int grid_l = __vsub4(grid_pos.x ^ signs0, signs0); + + const int signs1 = __vcmpne4(signs & 0x80402010, 0); + const int grid_h = __vsub4(grid_pos.y ^ signs1, signs1); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*MMQ_MMA_TILE_X_K_Q3_K + 8*kqsx + (2*l + 0)] = grid_l; + x_qs[i*MMQ_MMA_TILE_X_K_Q3_K + 8*kqsx + (2*l + 1)] = grid_h; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 0)] = grid_l; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 1)] = grid_h; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + + const int ls = bxi->scales[kqsx]; + const float d = bxi->d; +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_df[i*MMQ_MMA_TILE_X_K_Q3_K + 2*kqsx+0] = ((ls & 0x0F)*d + d/2)/4; + x_df[i*MMQ_MMA_TILE_X_K_Q3_K + 2*kqsx+1] = ((ls >> 4)*d + d/2)/4; +#else + x_df[i*(2*MMQ_TILE_NE_K*2/QI8_0) + i/(QI8_0/4) + 2*kqsx+0] = ((ls & 0x0F)*d + d/2)/4; + x_df[i*(2*MMQ_TILE_NE_K*2/QI8_0) + i/(QI8_0/4) + 2*kqsx+1] = ((ls >> 4)*d + d/2)/4; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +template static __device__ __forceinline__ void load_tiles_iq2_s( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ2_S, mmq_y); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr int threads_per_row = (MMQ_ITER_K / (4 * QR2_S)) / 2; + constexpr int nrows = warp_size / threads_per_row; + const int kqsx = threadIdx.x % threads_per_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * nrows) { + int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row; + + if (need_check) { + i = min(i, i_max); + } + + const block_iq2_s * bxi = (const block_iq2_s *) x + kbx0 + i*stride; + + const int qs_packed = get_int_b2(bxi->qs, kqsx); + const uint8_t * qs = (const uint8_t *) &qs_packed; + + const int qh = bxi->qh[kqsx]; + + const int signs_packed_32 = get_int_b2(bxi->qs, QK_K/32 + kqsx); + const uint8_t * signs_packed_8 = (const uint8_t *) &signs_packed_32; + +#pragma unroll + for (int l = 0; l < QR2_S; ++l) { + const int * grid_pos = (const int *)(iq2s_grid + (qs[l] | ((qh << (8-2*l)) & 0x300))); + + const int signs0 = __vcmpne4(((signs_packed_8[l] & 0x03) << 7) | ((signs_packed_8[l] & 0x0C) << 21), 0x00000000); + const int signs1 = __vcmpne4(((signs_packed_8[l] & 0x30) << 3) | ((signs_packed_8[l] & 0xC0) << 17), 0x00000000); + + const int grid_l = __vsub4(grid_pos[0] ^ signs0, signs0); + const int grid_h = __vsub4(grid_pos[1] ^ signs1, signs1); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*MMQ_MMA_TILE_X_K_Q3_K + 8*kqsx + (2*l + 0)] = grid_l; + x_qs[i*MMQ_MMA_TILE_X_K_Q3_K + 8*kqsx + (2*l + 1)] = grid_h; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 0)] = grid_l; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 1)] = grid_h; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + + const int ls = bxi->scales[kqsx]; + const float d = bxi->d; +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_df[i*MMQ_MMA_TILE_X_K_Q3_K + 2*kqsx+0] = ((ls & 0x0F)*d + d/2)/4; + x_df[i*MMQ_MMA_TILE_X_K_Q3_K + 2*kqsx+1] = ((ls >> 4)*d + d/2)/4; +#else + x_df[i*(2*MMQ_TILE_NE_K*2/QI8_0) + i/(QI8_0/4) + 2*kqsx+0] = ((ls & 0x0F)*d + d/2)/4; + x_df[i*(2*MMQ_TILE_NE_K*2/QI8_0) + i/(QI8_0/4) + 2*kqsx+1] = ((ls >> 4)*d + d/2)/4; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +template static __device__ __forceinline__ void load_tiles_iq3_xxs( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ3_XXS, mmq_y); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = (MMQ_ITER_K / (4 * QR3_XXS)) / 2; + constexpr int nrows = warp_size / threads_per_row; + const int kqsx = threadIdx.x % threads_per_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * nrows) { + int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row; + + if (need_check) { + i = min(i, i_max); + } + + const block_iq3_xxs * bxi = (const block_iq3_xxs *) x + kbx0 + i*stride; + + const int2 q3_packed = make_int2(get_int_b2(bxi->qs, 2*kqsx+0), get_int_b2(bxi->qs, 2*kqsx+1)); + const uint8_t * q3 = (const uint8_t *) &q3_packed; + const uint32_t aux32 = get_int_b2(bxi->qs, QK_K/16 + kqsx); + +#pragma unroll + for (int l = 0; l < QR3_XXS; ++l) { + const int2 grid_pos = make_int2(iq3xxs_grid[q3[2*l+0]], iq3xxs_grid[q3[2*l+1]]); + const uint32_t signs = unpack_ksigns(aux32 >> (7*l)); + + const int signs0 = __vcmpne4(signs & 0x08040201, 0); + const int grid_l = __vsub4(grid_pos.x ^ signs0, signs0); + + const int signs1 = __vcmpne4(signs & 0x80402010, 0); + const int grid_h = __vsub4(grid_pos.y ^ signs1, signs1); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + 8*kqsx + (2*l + 0)] = grid_l; + x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + 8*kqsx + (2*l + 1)] = grid_h; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 0)] = grid_l; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l + 1)] = grid_h; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + + const int ls = aux32 >> 28; + const float d = bxi->d; +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_df[i*MMQ_MMA_TILE_X_K_Q8_0 + kqsx] = (ls*d + d/2)/2; +#else + x_df[i*(MMQ_TILE_NE_K/4) + i/4 + kqsx] = (ls*d + d/2)/2; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +template static __device__ __forceinline__ void load_tiles_iq3_s( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ3_S, mmq_y); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = (MMQ_ITER_K / (4 * QR3_S)) / 2; + constexpr int nrows = warp_size / threads_per_row; + const int kqsx = threadIdx.x % threads_per_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * nrows) { + int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row; + + if (need_check) { + i = min(i, i_max); + } + + const block_iq3_s * bxi = (const block_iq3_s *) x + kbx0 + i*stride; + + const int2 qs_packed = make_int2(get_int_b2(bxi->qs, 2*kqsx+0), get_int_b2(bxi->qs, 2*kqsx+1)); + const uint8_t * qs = (const uint8_t *) &qs_packed; + + const int qh = bxi->qh[kqsx]; + + const int signs_packed_32 = get_int_b2(bxi->signs, kqsx); + const uint8_t * signs_packed_8 = (const uint8_t *) &signs_packed_32; + +#pragma unroll + for (int l = 0; l < QR3_S; ++l) { + const int2 grid_pos = make_int2( + iq3s_grid[qs[2*l+0] | ((qh << (8 - 2*l)) & 0x100)], + iq3s_grid[qs[2*l+1] | ((qh << (7 - 2*l)) & 0x100)]); + + const int signs0 = __vcmpne4(((signs_packed_8[l] & 0x03) << 7) | ((signs_packed_8[l] & 0x0C) << 21), 0x00000000); + const int signs1 = __vcmpne4(((signs_packed_8[l] & 0x30) << 3) | ((signs_packed_8[l] & 0xC0) << 17), 0x00000000); + + const int grid_l = __vsub4(grid_pos.x ^ signs0, signs0); + const int grid_h = __vsub4(grid_pos.y ^ signs1, signs1); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + 8*kqsx + (2*l+0)] = grid_l; + x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + 8*kqsx + (2*l+1)] = grid_h; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l+0)] = grid_l; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l+1)] = grid_h; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + + const int ls = 1 + 2*((bxi->scales[kqsx/2] >> (((2*kqsx) << 1) & 0x04)) & 0x0F); + const float d = bxi->d; +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_df[i*MMQ_MMA_TILE_X_K_Q8_0 + kqsx] = ls*d; +#else + x_df[i*(MMQ_TILE_NE_K/4) + i/4 + kqsx] = ls*d; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +template static __device__ __forceinline__ void load_tiles_iq1_s( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + int * x_qs = (int *) x_tile; + half2 * x_ds = (half2 *) (x_qs + MMQ_TILE_NE_K*2); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ3_S, mmq_y); + int * x_qs = (int *) x_tile; + half2 * x_ds = (half2 *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = MMQ_ITER_K / (4 * QR1_S); + constexpr int nrows = warp_size / threads_per_row; + const int kqsx = threadIdx.x % threads_per_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * nrows) { + int i = i0 + threadIdx.y*nrows + threadIdx.x/threads_per_row; + + if (need_check) { + i = min(i, i_max); + } + + const block_iq1_s * bxi = (const block_iq1_s *) x + kbx0 + i*stride; + + const int qs_packed = get_int_b2(bxi->qs, kqsx); + const uint8_t * qs = (const uint8_t *) &qs_packed; + + const int qh = bxi->qh[kqsx]; + + #pragma unroll + for (int l = 0; l < QR1_S/2; ++l) { + const int grid = iq1s_grid_gpu[qs[l] | (((qh >> (3*l)) & 0x07) << 8)]; + + const int grid0 = (grid >> 0) & 0x0F0F0F0F; + const int grid1 = (grid >> 4) & 0x0F0F0F0F; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*MMQ_MMA_TILE_X_K_Q8_1 + 8*kqsx + (2*l+0)] = grid0; + x_qs[i*MMQ_MMA_TILE_X_K_Q8_1 + 8*kqsx + (2*l+1)] = grid1; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l+0)] = grid0; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + 8*kqsx + (2*l+1)] = grid1; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + + const float d1q = __half2float(bxi->d) * (((qh >> 11) & 0x0E) + 1); + const float delta = -1.0f + IQ1S_DELTA - (qh & 0x8000) * (2.0f*IQ1S_DELTA/0x8000); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_ds[i*MMQ_MMA_TILE_X_K_Q8_1 + kqsx] = make_half2(d1q, d1q*delta); +#else + x_ds[i*(MMQ_TILE_NE_K/4) + i/4 + kqsx] = make_half2(d1q, d1q*delta); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +template static __device__ __forceinline__ void load_tiles_iq4_xs( + const char * __restrict__ x, int * __restrict__ x_tile, const int kbx0, const int i_max, const int stride) { + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + MMQ_TILE_NE_K*2); +#else + constexpr tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(GGML_TYPE_IQ4_XS, mmq_y); + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_qs + txs.qs); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + + constexpr int threads_per_row = MMQ_ITER_K / (4 * QR4_XS); + constexpr int nrows = warp_size / threads_per_row; + const int kqsx = threadIdx.x % threads_per_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nrows*nwarps) { + int i = i0 + (nrows == 1 ? threadIdx.y : threadIdx.y*nrows + threadIdx.x/threads_per_row); + + if (need_check) { + i = min(i, i_max); + } + + const block_iq4_xs * bxi = (const block_iq4_xs *) x + kbx0 + i*stride; + + const int aux_q4 = get_int_b4(bxi->qs, kqsx); + const int2 v = get_int_from_table_16(aux_q4, kvalues_iq4nl); + const int k0 = 8 * (kqsx / 4) + kqsx % 4; + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + k0 + 0] = v.x; + x_qs[i*MMQ_MMA_TILE_X_K_Q8_0 + k0 + 4] = v.y; +#else + x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + 0] = v.x; + x_qs[i*(2*MMQ_TILE_NE_K + 1) + k0 + 4] = v.y; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } + + constexpr int rows_per_warp = warp_size / 8; +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * rows_per_warp) { + int i = i0 + threadIdx.y * rows_per_warp + threadIdx.x / (MMQ_TILE_NE_K/4); + + if (need_check) { + i = min(i, i_max); + } + + const block_iq4_xs * bxi = (const block_iq4_xs *) x + kbx0 + i*stride; + + const float d = __half2float(bxi->d); + + const int ls = ((bxi->scales_l[(threadIdx.x % 8)/2] >> (4*(threadIdx.x % 2))) & 0x0F) + | (((bxi->scales_h >> (2*(threadIdx.x % 8))) & 0x03) << 4); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + x_df[i*MMQ_MMA_TILE_X_K_Q8_0 + threadIdx.x % 8] = d * (ls - 32); +#else + x_df[i*(MMQ_TILE_NE_K/4) + i/4 + threadIdx.x % 8] = d * (ls - 32); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + } +} + +template +static __device__ __forceinline__ void mmq_write_back_dp4a( + const float * __restrict__ sum, const int32_t * __restrict__ ids_dst, float * __restrict__ dst, + const int stride, const int i_max, const int j_max) { + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += nwarps) { + const int j = j0 + threadIdx.y; + + if (j > j_max) { + return; + } + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += warp_size) { + const int i = i0 + threadIdx.x; + + if (need_check && i > i_max) { + continue; + } + + dst[ids_dst[j]*stride + i] = sum[(j0/nwarps) * (mmq_y/warp_size) + i0/warp_size]; + } + } +} + +template +static __device__ __forceinline__ void mmq_write_back_mma( + const float * __restrict__ sum, const int * __restrict__ ids_dst, float * __restrict__ dst, + const int stride, const int i_max, const int j_max) { + + constexpr int granularity = mmq_get_granularity_device(mmq_x); + constexpr int nwarps = mmq_get_nwarps_device(); + +#if defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr int tileC_IJ = mmq_get_granularity_device(0); + typedef tile tile_C; + constexpr int rows_per_warp = granularity; +#else + typedef tile<16, 8, int> tile_C; + constexpr int rows_per_warp = 2 * granularity; +#endif // defined(AMD_MFMA_AVAILABLE) + constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. + + const int i0 = (threadIdx.y / ntx) * (ntx*tile_C::I); +#if defined(TURING_MMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + static_assert(nwarps*tile_C::I == mmq_y, "nwarps*tile_C::I != mmq_y"); +#else + GGML_UNUSED(nwarps); +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += ntx*tile_C::J) { +#pragma unroll + for (int n = 0; n < ntx; ++n) { +#pragma unroll + for (int l = 0; l < tile_C::ne; ++l) { + const int j = j0 + (threadIdx.y % ntx) * tile_C::J + tile_C::get_j(l); + + if (j > j_max) { + continue; + } + + const int i = i0 + n*tile_C::I + tile_C::get_i(l); + + if (need_check && i > i_max) { + continue; + } + + dst[ids_dst[j]*stride + i] = sum[(j0/tile_C::J + n)*tile_C::ne + l]; + } + } + } +} + +// ------------------------------------------------------------------------------------------------------------------------------------- + +template +struct mmq_type_traits; + +template +struct mmq_type_traits { + static constexpr int vdr = VDR_Q4_0_Q8_1_MMQ; + static constexpr load_tiles_mmq_t load_tiles = load_tiles_q4_0; + static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_q8_1_mma; + static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q4_0_q8_1_dp4a; +}; + +template +struct mmq_type_traits { + static constexpr int vdr = VDR_Q4_1_Q8_1_MMQ; + static constexpr load_tiles_mmq_t load_tiles = load_tiles_q4_1; + static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_1_q8_1_mma; + static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q4_1_q8_1_dp4a; +}; + +template +struct mmq_type_traits { + static constexpr int vdr = VDR_Q5_0_Q8_1_MMQ; + static constexpr load_tiles_mmq_t load_tiles = load_tiles_q5_0; + static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_q8_1_mma; + static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_q8_1_dp4a; +}; + +template +struct mmq_type_traits { + static constexpr int vdr = VDR_Q5_1_Q8_1_MMQ; + static constexpr load_tiles_mmq_t load_tiles = load_tiles_q5_1; + static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_1_q8_1_mma; + static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_1_q8_1_dp4a; +}; + +template +struct mmq_type_traits { + static constexpr int vdr = VDR_Q8_0_Q8_1_MMQ; + static constexpr load_tiles_mmq_t load_tiles = load_tiles_q8_0; + static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_q8_1_mma; + static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_q8_1_dp4a; +}; + +template +struct mmq_type_traits { + static constexpr int vdr = VDR_MXFP4_Q8_1_MMQ; +#ifdef BLACKWELL_MMA_AVAILABLE + static constexpr load_tiles_mmq_t load_tiles = load_tiles_mxfp4_fp4; + static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_mxfp4_mxfp4_mma; +#else + static constexpr load_tiles_mmq_t load_tiles = load_tiles_mxfp4; + static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_q8_1_mma; +#endif // BLACKWELL_MMA_AVAILABLE + static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_q8_1_dp4a; +}; + +template +struct mmq_type_traits { + static constexpr int vdr = VDR_NVFP4_Q8_1_MMQ; + static constexpr load_tiles_mmq_t load_tiles = load_tiles_nvfp4; + static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_16_q8_1_mma; + static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_16_q8_1_dp4a; +}; + +template +struct mmq_type_traits { + static constexpr int vdr = VDR_Q2_K_Q8_1_MMQ; + static constexpr load_tiles_mmq_t load_tiles = load_tiles_q2_K; + static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q2_K_q8_1_mma; + static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q2_K_q8_1_dp4a; +}; + +template +struct mmq_type_traits { + static constexpr int vdr = VDR_Q3_K_Q8_1_MMQ; + static constexpr load_tiles_mmq_t load_tiles = load_tiles_q3_K; + static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_16_q8_1_mma; + static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q3_K_q8_1_dp4a; +}; + +template +struct mmq_type_traits { + static constexpr int vdr = VDR_Q4_K_Q8_1_MMQ; + static constexpr load_tiles_mmq_t load_tiles = load_tiles_q4_K; + static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_1_q8_1_mma; + static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q4_K_q8_1_dp4a; +}; + +template +struct mmq_type_traits { + static constexpr int vdr = VDR_Q5_K_Q8_1_MMQ; + static constexpr load_tiles_mmq_t load_tiles = load_tiles_q5_K; + static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_1_q8_1_mma; + static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q5_K_q8_1_dp4a; +}; + +template +struct mmq_type_traits { + static constexpr int vdr = VDR_Q6_K_Q8_1_MMQ; + static constexpr load_tiles_mmq_t load_tiles = load_tiles_q6_K; + static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q6_K_q8_1_mma; + static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q6_K_q8_1_dp4a; +}; + +template +struct mmq_type_traits { + static constexpr int vdr = VDR_IQ2_XXS_Q8_1_MMQ; + static constexpr load_tiles_mmq_t load_tiles = load_tiles_iq2_xxs; + static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_q8_1_mma; + static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_q8_1_dp4a; +}; + +template +struct mmq_type_traits { + static constexpr int vdr = VDR_IQ2_XS_Q8_1_MMQ; + static constexpr load_tiles_mmq_t load_tiles = load_tiles_iq2_xs; + static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_16_q8_1_mma; + static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_16_q8_1_dp4a; +}; + +template +struct mmq_type_traits { + static constexpr int vdr = VDR_IQ2_S_Q8_1_MMQ; + static constexpr load_tiles_mmq_t load_tiles = load_tiles_iq2_s; + static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_16_q8_1_mma; + static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_16_q8_1_dp4a; +}; + +template +struct mmq_type_traits { + static constexpr int vdr = VDR_IQ3_XXS_Q8_1_MMQ; + static constexpr load_tiles_mmq_t load_tiles = load_tiles_iq3_xxs; + static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_q8_1_mma; + static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_q8_1_dp4a; +}; + +template +struct mmq_type_traits { + static constexpr int vdr = VDR_IQ3_S_Q8_1_MMQ; + static constexpr load_tiles_mmq_t load_tiles = load_tiles_iq3_s; + static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_q8_1_mma; + static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_q8_1_dp4a; +}; + +template +struct mmq_type_traits { + static constexpr int vdr = VDR_IQ1_S_Q8_1_MMQ; + static constexpr load_tiles_mmq_t load_tiles = load_tiles_iq1_s; + static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_1_q8_1_mma; + static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_1_q8_1_dp4a; +}; + +template +struct mmq_type_traits { + static constexpr int vdr = VDR_IQ4_NL_Q8_1_MMQ; + static constexpr load_tiles_mmq_t load_tiles = load_tiles_iq4_nl; + static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_q8_1_mma; + static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_q8_1_dp4a; +}; + +template +struct mmq_type_traits { + static constexpr int vdr = VDR_IQ4_XS_Q8_1_MMQ; + static constexpr load_tiles_mmq_t load_tiles = load_tiles_iq4_xs; + static constexpr vec_dot_mmq_t vec_dot_mma = vec_dot_q8_0_q8_1_mma; + static constexpr vec_dot_mmq_t vec_dot_dp4a = vec_dot_q8_0_q8_1_dp4a; +}; + +template +static __device__ __forceinline__ void mul_mat_q_process_tile( + const char * __restrict__ x, const int offset_x, const int * __restrict__ y, + const int * __restrict__ ids_dst, float * __restrict__ dst, float * __restrict__ tmp_fixup, + const int stride_row_x, const int ncols_y, const int stride_col_dst, + const int tile_x_max_i, const int tile_y_max_j, const int kb0_start, const int kb0_stop) { + + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int qk = ggml_cuda_type_traits::qk; + constexpr int mmq_y = get_mmq_y_device(); + constexpr load_tiles_mmq_t load_tiles = mmq_type_traits::load_tiles; + + extern __shared__ int data_mul_mat_q[]; + int * tile_y = data_mul_mat_q + mmq_x; + int * tile_x = tile_y + GGML_PAD(mmq_x*MMQ_TILE_Y_K, nwarps*warp_size); + +#if defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + constexpr vec_dot_mmq_t vec_dot = mmq_type_traits::vec_dot_mma; + constexpr mmq_write_back_t write_back = mmq_write_back_mma; +#else + constexpr vec_dot_mmq_t vec_dot = mmq_type_traits::vec_dot_dp4a; + constexpr mmq_write_back_t write_back = mmq_write_back_dp4a; +#endif // defined(AMD_MFMA_AVAILABLE) || defined(TURING_MMA_AVAILABLE) || defined(AMD_WMMA_AVAILABLE) + +#if defined(BLACKWELL_MMA_AVAILABLE) + // FP4 tile stores 8 blocks + constexpr int ne_block = (type == GGML_TYPE_MXFP4) ? 8 * QK_MXFP4 : 4 * QK8_1; +#else + constexpr int ne_block = 4 * QK8_1; +#endif // defined(BLACKWELL_MMA_AVAILABLE) + + constexpr int ITER_K = get_iter_k(type); + constexpr int blocks_per_iter = ITER_K / qk; + + float sum[mmq_x*mmq_y / (nwarps*warp_size)] = {0.0f}; + + constexpr int sz = sizeof(block_q8_1_mmq) / sizeof(int); + + for (int kb0 = kb0_start; kb0 < kb0_stop; kb0 += blocks_per_iter) { + load_tiles(x, tile_x, offset_x + kb0, tile_x_max_i, stride_row_x); + { + const int * by0 = y + ncols_y * (kb0 * qk / ne_block) * sz; +#pragma unroll + for (int l0 = 0; l0 < mmq_x * MMQ_TILE_Y_K; l0 += nwarps * warp_size) { + int l = l0 + threadIdx.y*warp_size + threadIdx.x; + + tile_y[l] = by0[l]; + } + } + + __syncthreads(); + + vec_dot(tile_x, tile_y, sum, 0); + + __syncthreads(); + + { + const int * by0 = y + ncols_y * ((kb0 * qk / ne_block) * sz + sz); +#pragma unroll + for (int l0 = 0; l0 < mmq_x * MMQ_TILE_Y_K; l0 += nwarps * warp_size) { + int l = l0 + threadIdx.y*warp_size + threadIdx.x; + + tile_y[l] = by0[l]; + } + } + + __syncthreads(); + + vec_dot(tile_x, tile_y, sum, MMQ_TILE_NE_K); + + __syncthreads(); + } + + if (fixup) { + write_back(sum, ids_dst, tmp_fixup + blockIdx.x*(mmq_x*mmq_y), mmq_y, mmq_y, mmq_x); + } else { + write_back(sum, ids_dst, dst, stride_col_dst, tile_x_max_i, tile_y_max_j); + } +} + + +// The mul_mat_q kernel implements "stream-k" work partitioning as described in https://arxiv.org/abs/2301.03598 + +template +#if defined(GGML_USE_HIP) +#if defined(RDNA4) || defined(RDNA3) || defined(RDNA2) || defined(CDNA) || defined(GCN) + __launch_bounds__(ggml_cuda_get_physical_warp_size()*mmq_get_nwarps_device(), 2) +#endif // defined(RDNA4) || defined(RDNA3) || defined(RDNA2) || defined(CDNA) || defined(GCN) +#else +#if __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA + __launch_bounds__(ggml_cuda_get_physical_warp_size()*mmq_get_nwarps_device(), 1) +#else + __launch_bounds__(ggml_cuda_get_physical_warp_size()*mmq_get_nwarps_device(), 2) +#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA +#endif // defined(GGML_USE_HIP) +static __global__ void mul_mat_q( + const char * __restrict__ x, const int * __restrict__ y, const int32_t * __restrict__ ids_dst, + const int32_t * __restrict__ expert_bounds, float * __restrict__ dst, float * __restrict__ tmp_fixup, + const int ncols_x, const int nrows_x, const int ncols_dst, const int stride_row_x, const int ncols_y, const int stride_col_dst, + const int channel_ratio, const int nchannels_y, const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst, + const int sample_ratio, const int nsamples_y, const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst, + const int ncols_max) { + + // Skip unused template specializations for faster compilation: + if (mmq_x > get_mmq_x_max_device() || mmq_x % mmq_get_granularity_device(mmq_x) != 0) { + NO_DEVICE_CODE; + return; + } + + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + + constexpr int qk = ggml_cuda_type_traits::qk; + constexpr int mmq_y = get_mmq_y_device(); + + const int ntx = (ncols_max + mmq_x - 1) / mmq_x; // Number of tiles x + const int nty = (nrows_x + mmq_y - 1) / mmq_y; // Number of tiles y + + // Initialize the ids for writing back data with just the index. + // For regular matrix multiplications this is never changed. + // For MoE the correct indices are loaded from ids_dst. + extern __shared__ int ids_dst_shared[]; // Stored at beginning of shared memory. +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += nwarps*warp_size) { + const int j = j0 + threadIdx.y*warp_size + threadIdx.x; + + if (j0 + nwarps*warp_size > mmq_x && j >= mmq_x) { + break; + } + + ids_dst_shared[j] = j; + } + __syncthreads(); + + // On non-CDNA AMD or old CUDA the performance with stream-k was worse, use conventional tiling instead: +#if (defined(GGML_USE_HIP) && !defined(CDNA)) || __CUDA_ARCH__ < GGML_CUDA_CC_VOLTA + { + const int wt = blockIdx.z / nchannels_y; + const int zt = blockIdx.z - wt*nchannels_y; + const int jt = blockIdx.y; + const int it = blockIdx.x; + + // Defaults for regular matrix multiplication: + int col_low = 0; + int col_high = ncols_dst; + int col_diff = ncols_dst; + int offset_y = wt*stride_sample_y + zt*stride_channel_y; + int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*mmq_x*stride_col_dst; + + if (ids_dst) { + col_low = expert_bounds[zt + 0]; + col_high = expert_bounds[zt + 1]; + col_diff = col_high - col_low; + + offset_y = 0; + offset_dst = 0; + + if (jt*mmq_x >= col_diff) { + return; + } + + // __syncthreads(); // There is no previous tile that could cause a race condition. +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += nwarps*warp_size) { + const int j = j0 + threadIdx.y*warp_size + threadIdx.x; + + if (j0 + nwarps*warp_size > mmq_x && j >= mmq_x) { + break; + } + + ids_dst_shared[j] = ids_dst[col_low + jt*mmq_x + j]; + } + __syncthreads(); + } + + offset_y += (col_low + jt*mmq_x)*(sizeof(block_q8_1_mmq)/sizeof(int)); + offset_dst += it*mmq_y; + + const int tile_x_max_i = nrows_x - it*mmq_y - 1; + const int tile_y_max_j = col_diff - jt*mmq_x - 1; + + const int offset_x = (wt/sample_ratio)*stride_sample_x + (zt/channel_ratio)*stride_channel_x + it*mmq_y*stride_row_x; + + constexpr bool fixup = false; + mul_mat_q_process_tile + (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, stride_row_x, ncols_y, stride_col_dst, + tile_x_max_i, tile_y_max_j, 0, ncols_x/qk); + return; + } +#endif // (defined(GGML_USE_HIP) && !defined(CDNA3)) || __CUDA_ARCH__ < GGML_CUDA_CC_VOLTA + + constexpr int ITER_K = get_iter_k(type); + + const int64_t blocks_per_ne00 = ncols_x / qk; + constexpr int blocks_per_iter = ITER_K / qk; + + // kbc == k block continuous, current index in continuous ijk space. + int64_t kbc = (int64_t) blockIdx.x *nsamples_y*nchannels_y*ntx*nty*blocks_per_ne00 / gridDim.x; + int64_t kbc_stop = (int64_t)(blockIdx.x + 1)*nsamples_y*nchannels_y*ntx*nty*blocks_per_ne00 / gridDim.x; + + kbc -= (kbc % blocks_per_ne00) % blocks_per_iter; + kbc_stop -= (kbc_stop % blocks_per_ne00) % blocks_per_iter; + + // kb0 == k index when doing the matrix multiplication for an output tile. + int kb0_start = kbc % blocks_per_ne00; + int kb0_stop = min(blocks_per_ne00, kb0_start + kbc_stop - kbc); + while (kbc < kbc_stop && kb0_stop == blocks_per_ne00) { + int tmp = kbc; + const int it = tmp / (nsamples_y*nchannels_y*ntx*blocks_per_ne00); + tmp -= it * (nsamples_y*nchannels_y*ntx*blocks_per_ne00); + const int wt = tmp / (nchannels_y*ntx*blocks_per_ne00); + tmp -= wt * (nchannels_y*ntx*blocks_per_ne00); + const int zt = tmp / (ntx*blocks_per_ne00); + tmp -= zt * (ntx*blocks_per_ne00); + const int jt = tmp / blocks_per_ne00; + + // Defaults for regular matrix multiplication: + int col_low = 0; + int col_high = ncols_dst; + int col_diff = ncols_dst; + int offset_y = wt*stride_sample_y + zt*stride_channel_y; + int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*mmq_x*stride_col_dst; + + if (ids_dst) { + col_low = expert_bounds[zt + 0]; + col_high = expert_bounds[zt + 1]; + col_diff = col_high - col_low; + + offset_y = 0; + offset_dst = 0; + + if (jt*mmq_x >= col_diff) { + kbc += blocks_per_ne00; + kbc -= kbc % blocks_per_ne00; + + kb0_start = 0; + kb0_stop = min(blocks_per_ne00, kbc_stop - kbc); + + continue; + } + + __syncthreads(); +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += nwarps*warp_size) { + const int j = j0 + threadIdx.y*warp_size + threadIdx.x; + + if (j0 + nwarps*warp_size > mmq_x && j >= mmq_x) { + break; + } + + ids_dst_shared[j] = ids_dst[col_low + jt*mmq_x + j]; + } + __syncthreads(); + } + + offset_y += (col_low + jt * mmq_x) * (sizeof(block_q8_1_mmq) / sizeof(int)); + offset_dst += it*mmq_y; + + const int tile_x_max_i = nrows_x - it*mmq_y - 1; + const int tile_y_max_j = col_diff - jt*mmq_x - 1; + + const int offset_x = (wt/sample_ratio)*stride_sample_x + (zt/channel_ratio)*stride_channel_x + it*mmq_y*stride_row_x; + + constexpr bool fixup = false; // All but (potentially) the last iterations write their data to dst rather than the fixup buffer. + mul_mat_q_process_tile + (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, stride_row_x, ncols_y, stride_col_dst, + tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop); + + kbc += blocks_per_ne00; + kbc -= kbc % blocks_per_ne00; + + kb0_start = 0; + kb0_stop = min(blocks_per_ne00, kbc_stop - kbc); + } + + if (kbc >= kbc_stop) { + return; + } + + int tmp = kbc; + const int it = tmp / (nsamples_y*nchannels_y*ntx*blocks_per_ne00); + tmp -= it * (nsamples_y*nchannels_y*ntx*blocks_per_ne00); + const int wt = tmp / (nchannels_y*ntx*blocks_per_ne00); + tmp -= wt * (nchannels_y*ntx*blocks_per_ne00); + const int zt = tmp / (ntx*blocks_per_ne00); + tmp -= zt * (ntx*blocks_per_ne00); + const int jt = tmp / blocks_per_ne00; + + // Defaults for regular matrix multiplication: + int col_low = 0; + int col_high = ncols_dst; + int col_diff = ncols_dst; + int offset_y = wt*stride_sample_y + zt*stride_channel_y; + int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*mmq_x*stride_col_dst; + + if (ids_dst) { + col_low = expert_bounds[zt + 0]; + col_high = expert_bounds[zt + 1]; + col_diff = col_high - col_low; + + offset_y = 0; + offset_dst = 0; + + if (jt*mmq_x >= col_diff) { + return; + } + + // The memory layout for the fixup buffer is always contiguous, therefore reset ids: + __syncthreads(); +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += nwarps*warp_size) { + const int j = j0 + threadIdx.y*warp_size + threadIdx.x; + + if (j0 + nwarps*warp_size > mmq_x && j >= mmq_x) { + break; + } + + ids_dst_shared[j] = j; + } + __syncthreads(); + } + + offset_y += (col_low + jt * mmq_x) * (sizeof(block_q8_1_mmq) / sizeof(int)); + offset_dst += it*mmq_y; + + const int tile_x_max_i = nrows_x - it*mmq_y - 1; + const int tile_y_max_j = col_diff - jt*mmq_x - 1; + + const int offset_x = (wt/sample_ratio)*stride_sample_x + (zt/channel_ratio)*stride_channel_x + it*mmq_y*stride_row_x; + + constexpr bool fixup = true; // Last index writes its data to fixup buffer to avoid data races with other blocks. + mul_mat_q_process_tile + (x, offset_x, y + offset_y, ids_dst_shared, dst + offset_dst, tmp_fixup, stride_row_x, ncols_y, stride_col_dst, + tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop); +} + +template +static __global__ void mul_mat_q_stream_k_fixup(const int32_t * ids_dst, + const int32_t * expert_bounds, + float * __restrict__ dst, + const float * __restrict__ tmp_last_tile, + const int ncols_x, + const int nrows_x, + const int ncols_dst, + const size_t stride_col_dst, + const int nchannels_y, + const size_t stride_channel_dst, + const int nsamples_y, + const size_t stride_sample_dst, + const int ncols_max) { + constexpr int mmq_y = get_mmq_y_device(); + constexpr int qk = ggml_cuda_type_traits::qk; + constexpr int ITER_K = get_iter_k(type); + + constexpr int blocks_per_iter = ITER_K / qk; + const int64_t blocks_per_ne00 = ncols_x / qk; + + constexpr int nwarps = mmq_get_nwarps_device(); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + + float sum[mmq_x*mmq_y / (nwarps*warp_size)] = {0.0f}; + + const int ntx = (ncols_max + mmq_x - 1) / mmq_x; + const int nty = (nrows_x + mmq_y - 1) / mmq_y; + + const int bidx0 = blockIdx.x; + + // kbc == k block continuous, current index in continuous ijk space. + int64_t kbc0 = (int64_t) bidx0 *nsamples_y*nchannels_y*ntx*nty*blocks_per_ne00 / gridDim.x; + int64_t kbc0_stop = (int64_t)(bidx0 + 1)*nsamples_y*nchannels_y*ntx*nty*blocks_per_ne00 / gridDim.x; + + kbc0 -= (kbc0 % blocks_per_ne00) % blocks_per_iter; + kbc0_stop -= (kbc0_stop % blocks_per_ne00) % blocks_per_iter; + + const bool did_not_have_any_data = kbc0 == kbc0_stop; + const bool wrote_beginning_of_tile = kbc0 % blocks_per_ne00 == 0; + const bool did_not_write_last = kbc0/blocks_per_ne00 == kbc0_stop/blocks_per_ne00 && kbc0_stop % blocks_per_ne00 != 0; + if (did_not_have_any_data || wrote_beginning_of_tile || did_not_write_last) { + return; + } + + bool any_fixup = false; + + // Iterate over previous blocks and sum up partial sums written to fixup buffer. + // All CUDA blocks that get here must have a previous block that needs a fixup. + int64_t bidx = bidx0 - 1; + int64_t kbc_stop = kbc0; + while(true) { + int64_t kbc = bidx*nsamples_y*nchannels_y*ntx*nty*blocks_per_ne00 / gridDim.x; + kbc -= (kbc % blocks_per_ne00) % blocks_per_iter; + + if (kbc == kbc_stop) { // Did not have any data. + bidx--; + kbc_stop = kbc; + continue; + } + + any_fixup = true; + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += nwarps) { + const int j = j0 + threadIdx.y; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += warp_size) { + const int i = i0 + threadIdx.x; + + sum[(j0/nwarps) * (mmq_y/warp_size) + i0/warp_size] += tmp_last_tile[bidx*(mmq_x*mmq_y) + j*mmq_y + i]; + } + } + + // If this block started in a previous tile we are done and don't need to combine additional partial results. + if (kbc % blocks_per_ne00 == 0 || kbc/blocks_per_ne00 < kbc0/blocks_per_ne00) { + break; + } + bidx--; + kbc_stop = kbc; + } + + if (!any_fixup) { + return; + } + + int tmp = kbc0; + const int it = tmp / (nsamples_y*nchannels_y*ntx*blocks_per_ne00); + tmp -= it * (nsamples_y*nchannels_y*ntx*blocks_per_ne00); + const int wt = tmp / (nchannels_y*ntx*blocks_per_ne00); + tmp -= wt * (nchannels_y*ntx*blocks_per_ne00); + const int zt = tmp / (ntx*blocks_per_ne00); + tmp -= zt * (ntx*blocks_per_ne00); + const int jt = tmp / blocks_per_ne00; + + if (!ids_dst) { + const int offset_dst = wt*stride_sample_dst + zt*stride_channel_dst + jt*mmq_x*stride_col_dst + it*mmq_y; + dst += offset_dst; + + const int i_max = nrows_x - it*mmq_y - 1; + const int j_max = ncols_dst - jt*mmq_x - 1; + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += nwarps) { + const int j = j0 + threadIdx.y; + + if (j > j_max) { + return; + } + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += warp_size) { + const int i = i0 + threadIdx.x; + + if (need_check && i > i_max) { + continue; + } + + dst[j*stride_col_dst + i] += sum[(j0/nwarps) * (mmq_y/warp_size) + i0/warp_size]; + } + } + return; + } + + __shared__ int ids_dst_shared[mmq_x]; + const int col_low = expert_bounds[zt + 0]; + const int col_high = expert_bounds[zt + 1]; + const int col_diff = col_high - col_low; + + for (int j = threadIdx.y*warp_size + threadIdx.x; j < mmq_x; j += nwarps*warp_size) { + ids_dst_shared[j] = ids_dst[col_low + jt*mmq_x + j]; + } + __syncthreads(); + + const int offset_dst = it*mmq_y; + dst += offset_dst; + + const int i_max = nrows_x - it*mmq_y - 1; + const int j_max = col_diff - jt*mmq_x - 1; + +#pragma unroll + for (int j0 = 0; j0 < mmq_x; j0 += nwarps) { + const int j = j0 + threadIdx.y; + + if (j > j_max) { + return; + } + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += warp_size) { + const int i = i0 + threadIdx.x; + + if (need_check && i > i_max) { + continue; + } + + dst[ids_dst_shared[j]*stride_col_dst + i] += sum[(j0/nwarps) * (mmq_y/warp_size) + i0/warp_size]; + } + } +} + +struct mmq_args { + const char * x; ggml_type type_x; const int * y; const int32_t * ids_dst; const int32_t * expert_bounds; float * dst; + int64_t ncols_x; int64_t nrows_x; int64_t ncols_dst; int64_t stride_row_x; int64_t ncols_y; int64_t nrows_dst; + int64_t nchannels_x; int64_t nchannels_y; int64_t stride_channel_x; int64_t stride_channel_y; int64_t stride_channel_dst; + int64_t nsamples_x; int64_t nsamples_y; int64_t stride_sample_x; int64_t stride_sample_y; int64_t stride_sample_dst; + bool use_stream_k; int64_t ncols_max; +}; + +template +static size_t mmq_get_nbytes_shared(const int mmq_x, const int mmq_y, const int cc, const int warp_size, const int nwarps) { + const tile_x_sizes txs = mmq_get_dp4a_tile_x_sizes(type, mmq_y); + const int mmq_tile_x_k = mmq_get_mma_tile_x_k(type); + const size_t nbs_ids = mmq_x*sizeof(int); + const size_t nbs_x = (turing_mma_available(cc) || amd_mfma_available(cc) || amd_wmma_available(cc)) ? mmq_y*mmq_tile_x_k*sizeof(int) : txs.qs*sizeof(int) + txs.dm*sizeof(half2) + txs.sc*sizeof(int); + const size_t nbs_y = mmq_x * (sizeof(block_q8_1_mmq)); + return nbs_ids + nbs_x + GGML_PAD(nbs_y, nwarps*warp_size*sizeof(int)); +} diff --git a/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q2_k.cu b/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q2_k.cu new file mode 100644 index 00000000..0cacee3e --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q2_k.cu @@ -0,0 +1,136 @@ +#include "mmq_common.cuh" +#include "mmq_gguf.cuh" + +// Force instantiation of all mmq_x variants for GGML_TYPE_Q2_K +template +static void instantiate_mmq_q2_k(float * tmp_fixup, + const mmq_args & args, cudaStream_t stream, + int cc, int nsm, size_t smpbo, int warp_size_host) { + + const int mmq_y = (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 128 : 64; + const int nwarps = 256 / warp_size_host; + const int nbytes_shared = mmq_get_nbytes_shared(mmq_x, mmq_y, cc, warp_size_host, nwarps); + const int nty = (args.nrows_x + mmq_y - 1) / mmq_y; + const int ntx = (args.ncols_max + mmq_x - 1) / mmq_x; + const int ntzw = args.nchannels_y * args.nsamples_y; + const dim3 block_dims(warp_size_host, nwarps, 1); + const int channel_ratio = args.nchannels_y / args.nchannels_x; + const int sample_ratio = args.nsamples_y / args.nsamples_x; + + CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); + CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); + + if (!args.use_stream_k) { + const dim3 grid(nty, ntx, ntzw); + if (args.nrows_x % mmq_y == 0) { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + } else { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + } + return; + } + + // Stream-k + const dim3 grid_sk(nsm, 1, 1); + const bool fixup_needed = ntx * nty * ntzw % nsm != 0; + + + if (args.nrows_x % mmq_y == 0) { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + if (fixup_needed) { + mul_mat_q_stream_k_fixup<<>>( + args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, args.ncols_x, args.nrows_x, args.ncols_dst, + args.nrows_dst, args.nchannels_y, args.stride_channel_dst, args.nsamples_y, args.stride_sample_dst, + args.ncols_max); + } + } else { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + if (fixup_needed) { + mul_mat_q_stream_k_fixup<<>>( + args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, args.ncols_x, args.nrows_x, args.ncols_dst, + args.nrows_dst, args.nchannels_y, args.stride_channel_dst, args.nsamples_y, args.stride_sample_dst, + args.ncols_max); + } + } +} + +static void launch_mmq_case_q2_k(float * tmp_fixup, const mmq_args & args, cudaStream_t stream, + int cc, int nsm, size_t smpbo, int warp_size_host) { + + const int mmq_x_max = (turing_mma_available(cc)) ? 128 : + (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 64 : 64; + const int mmq_y = (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 128 : 64; + const int nwarps = 256 / warp_size_host; + + int mmq_x_best = 0; + int ntiles_x_best = INT_MAX; + for (int mmq_x = 8; mmq_x <= mmq_x_max && ntiles_x_best > 1; mmq_x += 8) { + const int granularity = (turing_mma_available(cc) && mmq_x >= 48) ? 16 : 8; + if (mmq_x % granularity != 0) continue; + const size_t nbs = mmq_get_nbytes_shared(mmq_x, mmq_y, cc, warp_size_host, nwarps); + if (nbs > smpbo) continue; + const int ntiles_x = (args.ncols_max + mmq_x - 1) / mmq_x; + if (ntiles_x < ntiles_x_best) { mmq_x_best = mmq_x; ntiles_x_best = ntiles_x; } + } + + switch (mmq_x_best) { + case 8: instantiate_mmq_q2_k< 8>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 16: instantiate_mmq_q2_k< 16>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 24: instantiate_mmq_q2_k< 24>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 32: instantiate_mmq_q2_k< 32>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 40: instantiate_mmq_q2_k< 40>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 48: instantiate_mmq_q2_k< 48>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 56: instantiate_mmq_q2_k< 56>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 64: instantiate_mmq_q2_k< 64>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 72: instantiate_mmq_q2_k< 72>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 80: instantiate_mmq_q2_k< 80>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 88: instantiate_mmq_q2_k< 88>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 96: instantiate_mmq_q2_k< 96>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 104: instantiate_mmq_q2_k<104>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 112: instantiate_mmq_q2_k<112>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 120: instantiate_mmq_q2_k<120>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 128: instantiate_mmq_q2_k<128>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + default: break; + } +} + +extern "C" void launch_mmq_gguf_q2_k( + void *tmp_fixup_ptr, + const void *x, const void *y_q8_1_mmq, void *dst, + int64_t ncols_x, int64_t nrows_x, int64_t ncols_y, + int64_t stride_row_x, int64_t stride_col_dst, + int cc, int nsm, int64_t smpbo, int warp_size_host, + void *stream) { + + const bool use_stream_k = (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA); + + const mmq_args args = { + (const char *)x, GGML_TYPE_Q2_K, (const int *)y_q8_1_mmq, nullptr, nullptr, (float *)dst, + ncols_x, nrows_x, ncols_y, stride_row_x, ncols_y, nrows_x, + 1, 1, 0, 0, 0, + 1, 1, 0, 0, 0, + use_stream_k, ncols_y + }; + + launch_mmq_case_q2_k((float *)tmp_fixup_ptr, args, (cudaStream_t)stream, cc, nsm, smpbo, warp_size_host); +} diff --git a/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q3_k.cu b/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q3_k.cu new file mode 100644 index 00000000..46783e72 --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q3_k.cu @@ -0,0 +1,136 @@ +#include "mmq_common.cuh" +#include "mmq_gguf.cuh" + +// Force instantiation of all mmq_x variants for GGML_TYPE_Q3_K +template +static void instantiate_mmq_q3_k(float * tmp_fixup, + const mmq_args & args, cudaStream_t stream, + int cc, int nsm, size_t smpbo, int warp_size_host) { + + const int mmq_y = (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 128 : 64; + const int nwarps = 256 / warp_size_host; + const int nbytes_shared = mmq_get_nbytes_shared(mmq_x, mmq_y, cc, warp_size_host, nwarps); + const int nty = (args.nrows_x + mmq_y - 1) / mmq_y; + const int ntx = (args.ncols_max + mmq_x - 1) / mmq_x; + const int ntzw = args.nchannels_y * args.nsamples_y; + const dim3 block_dims(warp_size_host, nwarps, 1); + const int channel_ratio = args.nchannels_y / args.nchannels_x; + const int sample_ratio = args.nsamples_y / args.nsamples_x; + + CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); + CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); + + if (!args.use_stream_k) { + const dim3 grid(nty, ntx, ntzw); + if (args.nrows_x % mmq_y == 0) { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + } else { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + } + return; + } + + // Stream-k + const dim3 grid_sk(nsm, 1, 1); + const bool fixup_needed = ntx * nty * ntzw % nsm != 0; + + + if (args.nrows_x % mmq_y == 0) { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + if (fixup_needed) { + mul_mat_q_stream_k_fixup<<>>( + args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, args.ncols_x, args.nrows_x, args.ncols_dst, + args.nrows_dst, args.nchannels_y, args.stride_channel_dst, args.nsamples_y, args.stride_sample_dst, + args.ncols_max); + } + } else { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + if (fixup_needed) { + mul_mat_q_stream_k_fixup<<>>( + args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, args.ncols_x, args.nrows_x, args.ncols_dst, + args.nrows_dst, args.nchannels_y, args.stride_channel_dst, args.nsamples_y, args.stride_sample_dst, + args.ncols_max); + } + } +} + +static void launch_mmq_case_q3_k(float * tmp_fixup, const mmq_args & args, cudaStream_t stream, + int cc, int nsm, size_t smpbo, int warp_size_host) { + + const int mmq_x_max = (turing_mma_available(cc)) ? 128 : + (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 64 : 64; + const int mmq_y = (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 128 : 64; + const int nwarps = 256 / warp_size_host; + + int mmq_x_best = 0; + int ntiles_x_best = INT_MAX; + for (int mmq_x = 8; mmq_x <= mmq_x_max && ntiles_x_best > 1; mmq_x += 8) { + const int granularity = (turing_mma_available(cc) && mmq_x >= 48) ? 16 : 8; + if (mmq_x % granularity != 0) continue; + const size_t nbs = mmq_get_nbytes_shared(mmq_x, mmq_y, cc, warp_size_host, nwarps); + if (nbs > smpbo) continue; + const int ntiles_x = (args.ncols_max + mmq_x - 1) / mmq_x; + if (ntiles_x < ntiles_x_best) { mmq_x_best = mmq_x; ntiles_x_best = ntiles_x; } + } + + switch (mmq_x_best) { + case 8: instantiate_mmq_q3_k< 8>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 16: instantiate_mmq_q3_k< 16>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 24: instantiate_mmq_q3_k< 24>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 32: instantiate_mmq_q3_k< 32>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 40: instantiate_mmq_q3_k< 40>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 48: instantiate_mmq_q3_k< 48>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 56: instantiate_mmq_q3_k< 56>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 64: instantiate_mmq_q3_k< 64>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 72: instantiate_mmq_q3_k< 72>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 80: instantiate_mmq_q3_k< 80>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 88: instantiate_mmq_q3_k< 88>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 96: instantiate_mmq_q3_k< 96>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 104: instantiate_mmq_q3_k<104>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 112: instantiate_mmq_q3_k<112>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 120: instantiate_mmq_q3_k<120>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 128: instantiate_mmq_q3_k<128>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + default: break; + } +} + +extern "C" void launch_mmq_gguf_q3_k( + void *tmp_fixup_ptr, + const void *x, const void *y_q8_1_mmq, void *dst, + int64_t ncols_x, int64_t nrows_x, int64_t ncols_y, + int64_t stride_row_x, int64_t stride_col_dst, + int cc, int nsm, int64_t smpbo, int warp_size_host, + void *stream) { + + const bool use_stream_k = (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA); + + const mmq_args args = { + (const char *)x, GGML_TYPE_Q3_K, (const int *)y_q8_1_mmq, nullptr, nullptr, (float *)dst, + ncols_x, nrows_x, ncols_y, stride_row_x, ncols_y, nrows_x, + 1, 1, 0, 0, 0, + 1, 1, 0, 0, 0, + use_stream_k, ncols_y + }; + + launch_mmq_case_q3_k((float *)tmp_fixup_ptr, args, (cudaStream_t)stream, cc, nsm, smpbo, warp_size_host); +} diff --git a/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q4_0.cu b/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q4_0.cu new file mode 100644 index 00000000..db453659 --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q4_0.cu @@ -0,0 +1,136 @@ +#include "mmq_common.cuh" +#include "mmq_gguf.cuh" + +// Force instantiation of all mmq_x variants for GGML_TYPE_Q4_0 +template +static void instantiate_mmq_q4_0(float * tmp_fixup, + const mmq_args & args, cudaStream_t stream, + int cc, int nsm, size_t smpbo, int warp_size_host) { + + const int mmq_y = (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 128 : 64; + const int nwarps = 256 / warp_size_host; + const int nbytes_shared = mmq_get_nbytes_shared(mmq_x, mmq_y, cc, warp_size_host, nwarps); + const int nty = (args.nrows_x + mmq_y - 1) / mmq_y; + const int ntx = (args.ncols_max + mmq_x - 1) / mmq_x; + const int ntzw = args.nchannels_y * args.nsamples_y; + const dim3 block_dims(warp_size_host, nwarps, 1); + const int channel_ratio = args.nchannels_y / args.nchannels_x; + const int sample_ratio = args.nsamples_y / args.nsamples_x; + + CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); + CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); + + if (!args.use_stream_k) { + const dim3 grid(nty, ntx, ntzw); + if (args.nrows_x % mmq_y == 0) { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + } else { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + } + return; + } + + // Stream-k + const dim3 grid_sk(nsm, 1, 1); + const bool fixup_needed = ntx * nty * ntzw % nsm != 0; + + + if (args.nrows_x % mmq_y == 0) { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + if (fixup_needed) { + mul_mat_q_stream_k_fixup<<>>( + args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, args.ncols_x, args.nrows_x, args.ncols_dst, + args.nrows_dst, args.nchannels_y, args.stride_channel_dst, args.nsamples_y, args.stride_sample_dst, + args.ncols_max); + } + } else { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + if (fixup_needed) { + mul_mat_q_stream_k_fixup<<>>( + args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, args.ncols_x, args.nrows_x, args.ncols_dst, + args.nrows_dst, args.nchannels_y, args.stride_channel_dst, args.nsamples_y, args.stride_sample_dst, + args.ncols_max); + } + } +} + +static void launch_mmq_case_q4_0(float * tmp_fixup, const mmq_args & args, cudaStream_t stream, + int cc, int nsm, size_t smpbo, int warp_size_host) { + + const int mmq_x_max = (turing_mma_available(cc)) ? 128 : + (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 64 : 64; + const int mmq_y = (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 128 : 64; + const int nwarps = 256 / warp_size_host; + + int mmq_x_best = 0; + int ntiles_x_best = INT_MAX; + for (int mmq_x = 8; mmq_x <= mmq_x_max && ntiles_x_best > 1; mmq_x += 8) { + const int granularity = (turing_mma_available(cc) && mmq_x >= 48) ? 16 : 8; + if (mmq_x % granularity != 0) continue; + const size_t nbs = mmq_get_nbytes_shared(mmq_x, mmq_y, cc, warp_size_host, nwarps); + if (nbs > smpbo) continue; + const int ntiles_x = (args.ncols_max + mmq_x - 1) / mmq_x; + if (ntiles_x < ntiles_x_best) { mmq_x_best = mmq_x; ntiles_x_best = ntiles_x; } + } + + switch (mmq_x_best) { + case 8: instantiate_mmq_q4_0< 8>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 16: instantiate_mmq_q4_0< 16>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 24: instantiate_mmq_q4_0< 24>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 32: instantiate_mmq_q4_0< 32>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 40: instantiate_mmq_q4_0< 40>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 48: instantiate_mmq_q4_0< 48>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 56: instantiate_mmq_q4_0< 56>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 64: instantiate_mmq_q4_0< 64>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 72: instantiate_mmq_q4_0< 72>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 80: instantiate_mmq_q4_0< 80>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 88: instantiate_mmq_q4_0< 88>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 96: instantiate_mmq_q4_0< 96>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 104: instantiate_mmq_q4_0<104>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 112: instantiate_mmq_q4_0<112>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 120: instantiate_mmq_q4_0<120>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 128: instantiate_mmq_q4_0<128>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + default: break; + } +} + +extern "C" void launch_mmq_gguf_q4_0( + void *tmp_fixup_ptr, + const void *x, const void *y_q8_1_mmq, void *dst, + int64_t ncols_x, int64_t nrows_x, int64_t ncols_y, + int64_t stride_row_x, int64_t stride_col_dst, + int cc, int nsm, int64_t smpbo, int warp_size_host, + void *stream) { + + const bool use_stream_k = (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA); + + const mmq_args args = { + (const char *)x, GGML_TYPE_Q4_0, (const int *)y_q8_1_mmq, nullptr, nullptr, (float *)dst, + ncols_x, nrows_x, ncols_y, stride_row_x, ncols_y, nrows_x, + 1, 1, 0, 0, 0, + 1, 1, 0, 0, 0, + use_stream_k, ncols_y + }; + + launch_mmq_case_q4_0((float *)tmp_fixup_ptr, args, (cudaStream_t)stream, cc, nsm, smpbo, warp_size_host); +} diff --git a/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q4_1.cu b/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q4_1.cu new file mode 100644 index 00000000..1ea0d110 --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q4_1.cu @@ -0,0 +1,136 @@ +#include "mmq_common.cuh" +#include "mmq_gguf.cuh" + +// Force instantiation of all mmq_x variants for GGML_TYPE_Q4_1 +template +static void instantiate_mmq_q4_1(float * tmp_fixup, + const mmq_args & args, cudaStream_t stream, + int cc, int nsm, size_t smpbo, int warp_size_host) { + + const int mmq_y = (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 128 : 64; + const int nwarps = 256 / warp_size_host; + const int nbytes_shared = mmq_get_nbytes_shared(mmq_x, mmq_y, cc, warp_size_host, nwarps); + const int nty = (args.nrows_x + mmq_y - 1) / mmq_y; + const int ntx = (args.ncols_max + mmq_x - 1) / mmq_x; + const int ntzw = args.nchannels_y * args.nsamples_y; + const dim3 block_dims(warp_size_host, nwarps, 1); + const int channel_ratio = args.nchannels_y / args.nchannels_x; + const int sample_ratio = args.nsamples_y / args.nsamples_x; + + CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); + CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); + + if (!args.use_stream_k) { + const dim3 grid(nty, ntx, ntzw); + if (args.nrows_x % mmq_y == 0) { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + } else { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + } + return; + } + + // Stream-k + const dim3 grid_sk(nsm, 1, 1); + const bool fixup_needed = ntx * nty * ntzw % nsm != 0; + + + if (args.nrows_x % mmq_y == 0) { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + if (fixup_needed) { + mul_mat_q_stream_k_fixup<<>>( + args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, args.ncols_x, args.nrows_x, args.ncols_dst, + args.nrows_dst, args.nchannels_y, args.stride_channel_dst, args.nsamples_y, args.stride_sample_dst, + args.ncols_max); + } + } else { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + if (fixup_needed) { + mul_mat_q_stream_k_fixup<<>>( + args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, args.ncols_x, args.nrows_x, args.ncols_dst, + args.nrows_dst, args.nchannels_y, args.stride_channel_dst, args.nsamples_y, args.stride_sample_dst, + args.ncols_max); + } + } +} + +static void launch_mmq_case_q4_1(float * tmp_fixup, const mmq_args & args, cudaStream_t stream, + int cc, int nsm, size_t smpbo, int warp_size_host) { + + const int mmq_x_max = (turing_mma_available(cc)) ? 128 : + (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 64 : 64; + const int mmq_y = (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 128 : 64; + const int nwarps = 256 / warp_size_host; + + int mmq_x_best = 0; + int ntiles_x_best = INT_MAX; + for (int mmq_x = 8; mmq_x <= mmq_x_max && ntiles_x_best > 1; mmq_x += 8) { + const int granularity = (turing_mma_available(cc) && mmq_x >= 48) ? 16 : 8; + if (mmq_x % granularity != 0) continue; + const size_t nbs = mmq_get_nbytes_shared(mmq_x, mmq_y, cc, warp_size_host, nwarps); + if (nbs > smpbo) continue; + const int ntiles_x = (args.ncols_max + mmq_x - 1) / mmq_x; + if (ntiles_x < ntiles_x_best) { mmq_x_best = mmq_x; ntiles_x_best = ntiles_x; } + } + + switch (mmq_x_best) { + case 8: instantiate_mmq_q4_1< 8>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 16: instantiate_mmq_q4_1< 16>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 24: instantiate_mmq_q4_1< 24>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 32: instantiate_mmq_q4_1< 32>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 40: instantiate_mmq_q4_1< 40>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 48: instantiate_mmq_q4_1< 48>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 56: instantiate_mmq_q4_1< 56>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 64: instantiate_mmq_q4_1< 64>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 72: instantiate_mmq_q4_1< 72>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 80: instantiate_mmq_q4_1< 80>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 88: instantiate_mmq_q4_1< 88>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 96: instantiate_mmq_q4_1< 96>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 104: instantiate_mmq_q4_1<104>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 112: instantiate_mmq_q4_1<112>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 120: instantiate_mmq_q4_1<120>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 128: instantiate_mmq_q4_1<128>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + default: break; + } +} + +extern "C" void launch_mmq_gguf_q4_1( + void *tmp_fixup_ptr, + const void *x, const void *y_q8_1_mmq, void *dst, + int64_t ncols_x, int64_t nrows_x, int64_t ncols_y, + int64_t stride_row_x, int64_t stride_col_dst, + int cc, int nsm, int64_t smpbo, int warp_size_host, + void *stream) { + + const bool use_stream_k = (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA); + + const mmq_args args = { + (const char *)x, GGML_TYPE_Q4_1, (const int *)y_q8_1_mmq, nullptr, nullptr, (float *)dst, + ncols_x, nrows_x, ncols_y, stride_row_x, ncols_y, nrows_x, + 1, 1, 0, 0, 0, + 1, 1, 0, 0, 0, + use_stream_k, ncols_y + }; + + launch_mmq_case_q4_1((float *)tmp_fixup_ptr, args, (cudaStream_t)stream, cc, nsm, smpbo, warp_size_host); +} diff --git a/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q4_k.cu b/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q4_k.cu new file mode 100644 index 00000000..f3029578 --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q4_k.cu @@ -0,0 +1,136 @@ +#include "mmq_common.cuh" +#include "mmq_gguf.cuh" + +// Force instantiation of all mmq_x variants for GGML_TYPE_Q4_K +template +static void instantiate_mmq_q4_k(float * tmp_fixup, + const mmq_args & args, cudaStream_t stream, + int cc, int nsm, size_t smpbo, int warp_size_host) { + + const int mmq_y = (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 128 : 64; + const int nwarps = 256 / warp_size_host; + const int nbytes_shared = mmq_get_nbytes_shared(mmq_x, mmq_y, cc, warp_size_host, nwarps); + const int nty = (args.nrows_x + mmq_y - 1) / mmq_y; + const int ntx = (args.ncols_max + mmq_x - 1) / mmq_x; + const int ntzw = args.nchannels_y * args.nsamples_y; + const dim3 block_dims(warp_size_host, nwarps, 1); + const int channel_ratio = args.nchannels_y / args.nchannels_x; + const int sample_ratio = args.nsamples_y / args.nsamples_x; + + CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); + CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); + + if (!args.use_stream_k) { + const dim3 grid(nty, ntx, ntzw); + if (args.nrows_x % mmq_y == 0) { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + } else { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + } + return; + } + + // Stream-k + const dim3 grid_sk(nsm, 1, 1); + const bool fixup_needed = ntx * nty * ntzw % nsm != 0; + + + if (args.nrows_x % mmq_y == 0) { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + if (fixup_needed) { + mul_mat_q_stream_k_fixup<<>>( + args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, args.ncols_x, args.nrows_x, args.ncols_dst, + args.nrows_dst, args.nchannels_y, args.stride_channel_dst, args.nsamples_y, args.stride_sample_dst, + args.ncols_max); + } + } else { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + if (fixup_needed) { + mul_mat_q_stream_k_fixup<<>>( + args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, args.ncols_x, args.nrows_x, args.ncols_dst, + args.nrows_dst, args.nchannels_y, args.stride_channel_dst, args.nsamples_y, args.stride_sample_dst, + args.ncols_max); + } + } +} + +static void launch_mmq_case_q4_k(float * tmp_fixup, const mmq_args & args, cudaStream_t stream, + int cc, int nsm, size_t smpbo, int warp_size_host) { + + const int mmq_x_max = (turing_mma_available(cc)) ? 128 : + (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 64 : 64; + const int mmq_y = (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 128 : 64; + const int nwarps = 256 / warp_size_host; + + int mmq_x_best = 0; + int ntiles_x_best = INT_MAX; + for (int mmq_x = 8; mmq_x <= mmq_x_max && ntiles_x_best > 1; mmq_x += 8) { + const int granularity = (turing_mma_available(cc) && mmq_x >= 48) ? 16 : 8; + if (mmq_x % granularity != 0) continue; + const size_t nbs = mmq_get_nbytes_shared(mmq_x, mmq_y, cc, warp_size_host, nwarps); + if (nbs > smpbo) continue; + const int ntiles_x = (args.ncols_max + mmq_x - 1) / mmq_x; + if (ntiles_x < ntiles_x_best) { mmq_x_best = mmq_x; ntiles_x_best = ntiles_x; } + } + + switch (mmq_x_best) { + case 8: instantiate_mmq_q4_k< 8>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 16: instantiate_mmq_q4_k< 16>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 24: instantiate_mmq_q4_k< 24>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 32: instantiate_mmq_q4_k< 32>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 40: instantiate_mmq_q4_k< 40>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 48: instantiate_mmq_q4_k< 48>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 56: instantiate_mmq_q4_k< 56>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 64: instantiate_mmq_q4_k< 64>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 72: instantiate_mmq_q4_k< 72>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 80: instantiate_mmq_q4_k< 80>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 88: instantiate_mmq_q4_k< 88>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 96: instantiate_mmq_q4_k< 96>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 104: instantiate_mmq_q4_k<104>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 112: instantiate_mmq_q4_k<112>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 120: instantiate_mmq_q4_k<120>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 128: instantiate_mmq_q4_k<128>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + default: break; + } +} + +extern "C" void launch_mmq_gguf_q4_k( + void *tmp_fixup_ptr, + const void *x, const void *y_q8_1_mmq, void *dst, + int64_t ncols_x, int64_t nrows_x, int64_t ncols_y, + int64_t stride_row_x, int64_t stride_col_dst, + int cc, int nsm, int64_t smpbo, int warp_size_host, + void *stream) { + + const bool use_stream_k = (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA); + + const mmq_args args = { + (const char *)x, GGML_TYPE_Q4_K, (const int *)y_q8_1_mmq, nullptr, nullptr, (float *)dst, + ncols_x, nrows_x, ncols_y, stride_row_x, ncols_y, nrows_x, + 1, 1, 0, 0, 0, + 1, 1, 0, 0, 0, + use_stream_k, ncols_y + }; + + launch_mmq_case_q4_k((float *)tmp_fixup_ptr, args, (cudaStream_t)stream, cc, nsm, smpbo, warp_size_host); +} diff --git a/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q5_0.cu b/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q5_0.cu new file mode 100644 index 00000000..acc46a65 --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q5_0.cu @@ -0,0 +1,136 @@ +#include "mmq_common.cuh" +#include "mmq_gguf.cuh" + +// Force instantiation of all mmq_x variants for GGML_TYPE_Q5_0 +template +static void instantiate_mmq_q5_0(float * tmp_fixup, + const mmq_args & args, cudaStream_t stream, + int cc, int nsm, size_t smpbo, int warp_size_host) { + + const int mmq_y = (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 128 : 64; + const int nwarps = 256 / warp_size_host; + const int nbytes_shared = mmq_get_nbytes_shared(mmq_x, mmq_y, cc, warp_size_host, nwarps); + const int nty = (args.nrows_x + mmq_y - 1) / mmq_y; + const int ntx = (args.ncols_max + mmq_x - 1) / mmq_x; + const int ntzw = args.nchannels_y * args.nsamples_y; + const dim3 block_dims(warp_size_host, nwarps, 1); + const int channel_ratio = args.nchannels_y / args.nchannels_x; + const int sample_ratio = args.nsamples_y / args.nsamples_x; + + CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); + CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); + + if (!args.use_stream_k) { + const dim3 grid(nty, ntx, ntzw); + if (args.nrows_x % mmq_y == 0) { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + } else { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + } + return; + } + + // Stream-k + const dim3 grid_sk(nsm, 1, 1); + const bool fixup_needed = ntx * nty * ntzw % nsm != 0; + + + if (args.nrows_x % mmq_y == 0) { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + if (fixup_needed) { + mul_mat_q_stream_k_fixup<<>>( + args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, args.ncols_x, args.nrows_x, args.ncols_dst, + args.nrows_dst, args.nchannels_y, args.stride_channel_dst, args.nsamples_y, args.stride_sample_dst, + args.ncols_max); + } + } else { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + if (fixup_needed) { + mul_mat_q_stream_k_fixup<<>>( + args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, args.ncols_x, args.nrows_x, args.ncols_dst, + args.nrows_dst, args.nchannels_y, args.stride_channel_dst, args.nsamples_y, args.stride_sample_dst, + args.ncols_max); + } + } +} + +static void launch_mmq_case_q5_0(float * tmp_fixup, const mmq_args & args, cudaStream_t stream, + int cc, int nsm, size_t smpbo, int warp_size_host) { + + const int mmq_x_max = (turing_mma_available(cc)) ? 128 : + (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 64 : 64; + const int mmq_y = (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 128 : 64; + const int nwarps = 256 / warp_size_host; + + int mmq_x_best = 0; + int ntiles_x_best = INT_MAX; + for (int mmq_x = 8; mmq_x <= mmq_x_max && ntiles_x_best > 1; mmq_x += 8) { + const int granularity = (turing_mma_available(cc) && mmq_x >= 48) ? 16 : 8; + if (mmq_x % granularity != 0) continue; + const size_t nbs = mmq_get_nbytes_shared(mmq_x, mmq_y, cc, warp_size_host, nwarps); + if (nbs > smpbo) continue; + const int ntiles_x = (args.ncols_max + mmq_x - 1) / mmq_x; + if (ntiles_x < ntiles_x_best) { mmq_x_best = mmq_x; ntiles_x_best = ntiles_x; } + } + + switch (mmq_x_best) { + case 8: instantiate_mmq_q5_0< 8>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 16: instantiate_mmq_q5_0< 16>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 24: instantiate_mmq_q5_0< 24>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 32: instantiate_mmq_q5_0< 32>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 40: instantiate_mmq_q5_0< 40>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 48: instantiate_mmq_q5_0< 48>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 56: instantiate_mmq_q5_0< 56>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 64: instantiate_mmq_q5_0< 64>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 72: instantiate_mmq_q5_0< 72>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 80: instantiate_mmq_q5_0< 80>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 88: instantiate_mmq_q5_0< 88>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 96: instantiate_mmq_q5_0< 96>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 104: instantiate_mmq_q5_0<104>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 112: instantiate_mmq_q5_0<112>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 120: instantiate_mmq_q5_0<120>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 128: instantiate_mmq_q5_0<128>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + default: break; + } +} + +extern "C" void launch_mmq_gguf_q5_0( + void *tmp_fixup_ptr, + const void *x, const void *y_q8_1_mmq, void *dst, + int64_t ncols_x, int64_t nrows_x, int64_t ncols_y, + int64_t stride_row_x, int64_t stride_col_dst, + int cc, int nsm, int64_t smpbo, int warp_size_host, + void *stream) { + + const bool use_stream_k = (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA); + + const mmq_args args = { + (const char *)x, GGML_TYPE_Q5_0, (const int *)y_q8_1_mmq, nullptr, nullptr, (float *)dst, + ncols_x, nrows_x, ncols_y, stride_row_x, ncols_y, nrows_x, + 1, 1, 0, 0, 0, + 1, 1, 0, 0, 0, + use_stream_k, ncols_y + }; + + launch_mmq_case_q5_0((float *)tmp_fixup_ptr, args, (cudaStream_t)stream, cc, nsm, smpbo, warp_size_host); +} diff --git a/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q5_1.cu b/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q5_1.cu new file mode 100644 index 00000000..bcb968cd --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q5_1.cu @@ -0,0 +1,136 @@ +#include "mmq_common.cuh" +#include "mmq_gguf.cuh" + +// Force instantiation of all mmq_x variants for GGML_TYPE_Q5_1 +template +static void instantiate_mmq_q5_1(float * tmp_fixup, + const mmq_args & args, cudaStream_t stream, + int cc, int nsm, size_t smpbo, int warp_size_host) { + + const int mmq_y = (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 128 : 64; + const int nwarps = 256 / warp_size_host; + const int nbytes_shared = mmq_get_nbytes_shared(mmq_x, mmq_y, cc, warp_size_host, nwarps); + const int nty = (args.nrows_x + mmq_y - 1) / mmq_y; + const int ntx = (args.ncols_max + mmq_x - 1) / mmq_x; + const int ntzw = args.nchannels_y * args.nsamples_y; + const dim3 block_dims(warp_size_host, nwarps, 1); + const int channel_ratio = args.nchannels_y / args.nchannels_x; + const int sample_ratio = args.nsamples_y / args.nsamples_x; + + CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); + CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); + + if (!args.use_stream_k) { + const dim3 grid(nty, ntx, ntzw); + if (args.nrows_x % mmq_y == 0) { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + } else { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + } + return; + } + + // Stream-k + const dim3 grid_sk(nsm, 1, 1); + const bool fixup_needed = ntx * nty * ntzw % nsm != 0; + + + if (args.nrows_x % mmq_y == 0) { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + if (fixup_needed) { + mul_mat_q_stream_k_fixup<<>>( + args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, args.ncols_x, args.nrows_x, args.ncols_dst, + args.nrows_dst, args.nchannels_y, args.stride_channel_dst, args.nsamples_y, args.stride_sample_dst, + args.ncols_max); + } + } else { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + if (fixup_needed) { + mul_mat_q_stream_k_fixup<<>>( + args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, args.ncols_x, args.nrows_x, args.ncols_dst, + args.nrows_dst, args.nchannels_y, args.stride_channel_dst, args.nsamples_y, args.stride_sample_dst, + args.ncols_max); + } + } +} + +static void launch_mmq_case_q5_1(float * tmp_fixup, const mmq_args & args, cudaStream_t stream, + int cc, int nsm, size_t smpbo, int warp_size_host) { + + const int mmq_x_max = (turing_mma_available(cc)) ? 128 : + (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 64 : 64; + const int mmq_y = (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 128 : 64; + const int nwarps = 256 / warp_size_host; + + int mmq_x_best = 0; + int ntiles_x_best = INT_MAX; + for (int mmq_x = 8; mmq_x <= mmq_x_max && ntiles_x_best > 1; mmq_x += 8) { + const int granularity = (turing_mma_available(cc) && mmq_x >= 48) ? 16 : 8; + if (mmq_x % granularity != 0) continue; + const size_t nbs = mmq_get_nbytes_shared(mmq_x, mmq_y, cc, warp_size_host, nwarps); + if (nbs > smpbo) continue; + const int ntiles_x = (args.ncols_max + mmq_x - 1) / mmq_x; + if (ntiles_x < ntiles_x_best) { mmq_x_best = mmq_x; ntiles_x_best = ntiles_x; } + } + + switch (mmq_x_best) { + case 8: instantiate_mmq_q5_1< 8>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 16: instantiate_mmq_q5_1< 16>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 24: instantiate_mmq_q5_1< 24>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 32: instantiate_mmq_q5_1< 32>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 40: instantiate_mmq_q5_1< 40>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 48: instantiate_mmq_q5_1< 48>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 56: instantiate_mmq_q5_1< 56>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 64: instantiate_mmq_q5_1< 64>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 72: instantiate_mmq_q5_1< 72>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 80: instantiate_mmq_q5_1< 80>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 88: instantiate_mmq_q5_1< 88>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 96: instantiate_mmq_q5_1< 96>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 104: instantiate_mmq_q5_1<104>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 112: instantiate_mmq_q5_1<112>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 120: instantiate_mmq_q5_1<120>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 128: instantiate_mmq_q5_1<128>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + default: break; + } +} + +extern "C" void launch_mmq_gguf_q5_1( + void *tmp_fixup_ptr, + const void *x, const void *y_q8_1_mmq, void *dst, + int64_t ncols_x, int64_t nrows_x, int64_t ncols_y, + int64_t stride_row_x, int64_t stride_col_dst, + int cc, int nsm, int64_t smpbo, int warp_size_host, + void *stream) { + + const bool use_stream_k = (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA); + + const mmq_args args = { + (const char *)x, GGML_TYPE_Q5_1, (const int *)y_q8_1_mmq, nullptr, nullptr, (float *)dst, + ncols_x, nrows_x, ncols_y, stride_row_x, ncols_y, nrows_x, + 1, 1, 0, 0, 0, + 1, 1, 0, 0, 0, + use_stream_k, ncols_y + }; + + launch_mmq_case_q5_1((float *)tmp_fixup_ptr, args, (cudaStream_t)stream, cc, nsm, smpbo, warp_size_host); +} diff --git a/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q5_k.cu b/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q5_k.cu new file mode 100644 index 00000000..f9d668d4 --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q5_k.cu @@ -0,0 +1,136 @@ +#include "mmq_common.cuh" +#include "mmq_gguf.cuh" + +// Force instantiation of all mmq_x variants for GGML_TYPE_Q5_K +template +static void instantiate_mmq_q5_k(float * tmp_fixup, + const mmq_args & args, cudaStream_t stream, + int cc, int nsm, size_t smpbo, int warp_size_host) { + + const int mmq_y = (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 128 : 64; + const int nwarps = 256 / warp_size_host; + const int nbytes_shared = mmq_get_nbytes_shared(mmq_x, mmq_y, cc, warp_size_host, nwarps); + const int nty = (args.nrows_x + mmq_y - 1) / mmq_y; + const int ntx = (args.ncols_max + mmq_x - 1) / mmq_x; + const int ntzw = args.nchannels_y * args.nsamples_y; + const dim3 block_dims(warp_size_host, nwarps, 1); + const int channel_ratio = args.nchannels_y / args.nchannels_x; + const int sample_ratio = args.nsamples_y / args.nsamples_x; + + CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); + CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); + + if (!args.use_stream_k) { + const dim3 grid(nty, ntx, ntzw); + if (args.nrows_x % mmq_y == 0) { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + } else { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + } + return; + } + + // Stream-k + const dim3 grid_sk(nsm, 1, 1); + const bool fixup_needed = ntx * nty * ntzw % nsm != 0; + + + if (args.nrows_x % mmq_y == 0) { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + if (fixup_needed) { + mul_mat_q_stream_k_fixup<<>>( + args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, args.ncols_x, args.nrows_x, args.ncols_dst, + args.nrows_dst, args.nchannels_y, args.stride_channel_dst, args.nsamples_y, args.stride_sample_dst, + args.ncols_max); + } + } else { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + if (fixup_needed) { + mul_mat_q_stream_k_fixup<<>>( + args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, args.ncols_x, args.nrows_x, args.ncols_dst, + args.nrows_dst, args.nchannels_y, args.stride_channel_dst, args.nsamples_y, args.stride_sample_dst, + args.ncols_max); + } + } +} + +static void launch_mmq_case_q5_k(float * tmp_fixup, const mmq_args & args, cudaStream_t stream, + int cc, int nsm, size_t smpbo, int warp_size_host) { + + const int mmq_x_max = (turing_mma_available(cc)) ? 128 : + (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 64 : 64; + const int mmq_y = (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 128 : 64; + const int nwarps = 256 / warp_size_host; + + int mmq_x_best = 0; + int ntiles_x_best = INT_MAX; + for (int mmq_x = 8; mmq_x <= mmq_x_max && ntiles_x_best > 1; mmq_x += 8) { + const int granularity = (turing_mma_available(cc) && mmq_x >= 48) ? 16 : 8; + if (mmq_x % granularity != 0) continue; + const size_t nbs = mmq_get_nbytes_shared(mmq_x, mmq_y, cc, warp_size_host, nwarps); + if (nbs > smpbo) continue; + const int ntiles_x = (args.ncols_max + mmq_x - 1) / mmq_x; + if (ntiles_x < ntiles_x_best) { mmq_x_best = mmq_x; ntiles_x_best = ntiles_x; } + } + + switch (mmq_x_best) { + case 8: instantiate_mmq_q5_k< 8>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 16: instantiate_mmq_q5_k< 16>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 24: instantiate_mmq_q5_k< 24>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 32: instantiate_mmq_q5_k< 32>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 40: instantiate_mmq_q5_k< 40>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 48: instantiate_mmq_q5_k< 48>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 56: instantiate_mmq_q5_k< 56>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 64: instantiate_mmq_q5_k< 64>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 72: instantiate_mmq_q5_k< 72>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 80: instantiate_mmq_q5_k< 80>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 88: instantiate_mmq_q5_k< 88>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 96: instantiate_mmq_q5_k< 96>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 104: instantiate_mmq_q5_k<104>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 112: instantiate_mmq_q5_k<112>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 120: instantiate_mmq_q5_k<120>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 128: instantiate_mmq_q5_k<128>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + default: break; + } +} + +extern "C" void launch_mmq_gguf_q5_k( + void *tmp_fixup_ptr, + const void *x, const void *y_q8_1_mmq, void *dst, + int64_t ncols_x, int64_t nrows_x, int64_t ncols_y, + int64_t stride_row_x, int64_t stride_col_dst, + int cc, int nsm, int64_t smpbo, int warp_size_host, + void *stream) { + + const bool use_stream_k = (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA); + + const mmq_args args = { + (const char *)x, GGML_TYPE_Q5_K, (const int *)y_q8_1_mmq, nullptr, nullptr, (float *)dst, + ncols_x, nrows_x, ncols_y, stride_row_x, ncols_y, nrows_x, + 1, 1, 0, 0, 0, + 1, 1, 0, 0, 0, + use_stream_k, ncols_y + }; + + launch_mmq_case_q5_k((float *)tmp_fixup_ptr, args, (cudaStream_t)stream, cc, nsm, smpbo, warp_size_host); +} diff --git a/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q6_k.cu b/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q6_k.cu new file mode 100644 index 00000000..ee8bdcbc --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q6_k.cu @@ -0,0 +1,136 @@ +#include "mmq_common.cuh" +#include "mmq_gguf.cuh" + +// Force instantiation of all mmq_x variants for GGML_TYPE_Q6_K +template +static void instantiate_mmq_q6_k(float * tmp_fixup, + const mmq_args & args, cudaStream_t stream, + int cc, int nsm, size_t smpbo, int warp_size_host) { + + const int mmq_y = (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 128 : 64; + const int nwarps = 256 / warp_size_host; + const int nbytes_shared = mmq_get_nbytes_shared(mmq_x, mmq_y, cc, warp_size_host, nwarps); + const int nty = (args.nrows_x + mmq_y - 1) / mmq_y; + const int ntx = (args.ncols_max + mmq_x - 1) / mmq_x; + const int ntzw = args.nchannels_y * args.nsamples_y; + const dim3 block_dims(warp_size_host, nwarps, 1); + const int channel_ratio = args.nchannels_y / args.nchannels_x; + const int sample_ratio = args.nsamples_y / args.nsamples_x; + + CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); + CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); + + if (!args.use_stream_k) { + const dim3 grid(nty, ntx, ntzw); + if (args.nrows_x % mmq_y == 0) { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + } else { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + } + return; + } + + // Stream-k + const dim3 grid_sk(nsm, 1, 1); + const bool fixup_needed = ntx * nty * ntzw % nsm != 0; + + + if (args.nrows_x % mmq_y == 0) { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + if (fixup_needed) { + mul_mat_q_stream_k_fixup<<>>( + args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, args.ncols_x, args.nrows_x, args.ncols_dst, + args.nrows_dst, args.nchannels_y, args.stride_channel_dst, args.nsamples_y, args.stride_sample_dst, + args.ncols_max); + } + } else { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + if (fixup_needed) { + mul_mat_q_stream_k_fixup<<>>( + args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, args.ncols_x, args.nrows_x, args.ncols_dst, + args.nrows_dst, args.nchannels_y, args.stride_channel_dst, args.nsamples_y, args.stride_sample_dst, + args.ncols_max); + } + } +} + +static void launch_mmq_case_q6_k(float * tmp_fixup, const mmq_args & args, cudaStream_t stream, + int cc, int nsm, size_t smpbo, int warp_size_host) { + + const int mmq_x_max = (turing_mma_available(cc)) ? 128 : + (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 64 : 64; + const int mmq_y = (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 128 : 64; + const int nwarps = 256 / warp_size_host; + + int mmq_x_best = 0; + int ntiles_x_best = INT_MAX; + for (int mmq_x = 8; mmq_x <= mmq_x_max && ntiles_x_best > 1; mmq_x += 8) { + const int granularity = (turing_mma_available(cc) && mmq_x >= 48) ? 16 : 8; + if (mmq_x % granularity != 0) continue; + const size_t nbs = mmq_get_nbytes_shared(mmq_x, mmq_y, cc, warp_size_host, nwarps); + if (nbs > smpbo) continue; + const int ntiles_x = (args.ncols_max + mmq_x - 1) / mmq_x; + if (ntiles_x < ntiles_x_best) { mmq_x_best = mmq_x; ntiles_x_best = ntiles_x; } + } + + switch (mmq_x_best) { + case 8: instantiate_mmq_q6_k< 8>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 16: instantiate_mmq_q6_k< 16>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 24: instantiate_mmq_q6_k< 24>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 32: instantiate_mmq_q6_k< 32>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 40: instantiate_mmq_q6_k< 40>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 48: instantiate_mmq_q6_k< 48>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 56: instantiate_mmq_q6_k< 56>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 64: instantiate_mmq_q6_k< 64>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 72: instantiate_mmq_q6_k< 72>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 80: instantiate_mmq_q6_k< 80>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 88: instantiate_mmq_q6_k< 88>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 96: instantiate_mmq_q6_k< 96>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 104: instantiate_mmq_q6_k<104>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 112: instantiate_mmq_q6_k<112>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 120: instantiate_mmq_q6_k<120>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 128: instantiate_mmq_q6_k<128>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + default: break; + } +} + +extern "C" void launch_mmq_gguf_q6_k( + void *tmp_fixup_ptr, + const void *x, const void *y_q8_1_mmq, void *dst, + int64_t ncols_x, int64_t nrows_x, int64_t ncols_y, + int64_t stride_row_x, int64_t stride_col_dst, + int cc, int nsm, int64_t smpbo, int warp_size_host, + void *stream) { + + const bool use_stream_k = (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA); + + const mmq_args args = { + (const char *)x, GGML_TYPE_Q6_K, (const int *)y_q8_1_mmq, nullptr, nullptr, (float *)dst, + ncols_x, nrows_x, ncols_y, stride_row_x, ncols_y, nrows_x, + 1, 1, 0, 0, 0, + 1, 1, 0, 0, 0, + use_stream_k, ncols_y + }; + + launch_mmq_case_q6_k((float *)tmp_fixup_ptr, args, (cudaStream_t)stream, cc, nsm, smpbo, warp_size_host); +} diff --git a/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q8_0.cu b/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q8_0.cu new file mode 100644 index 00000000..65c31026 --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_instance_q8_0.cu @@ -0,0 +1,136 @@ +#include "mmq_common.cuh" +#include "mmq_gguf.cuh" + +// Force instantiation of all mmq_x variants for GGML_TYPE_Q8_0 +template +static void instantiate_mmq_q8_0(float * tmp_fixup, + const mmq_args & args, cudaStream_t stream, + int cc, int nsm, size_t smpbo, int warp_size_host) { + + const int mmq_y = (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 128 : 64; + const int nwarps = 256 / warp_size_host; + const int nbytes_shared = mmq_get_nbytes_shared(mmq_x, mmq_y, cc, warp_size_host, nwarps); + const int nty = (args.nrows_x + mmq_y - 1) / mmq_y; + const int ntx = (args.ncols_max + mmq_x - 1) / mmq_x; + const int ntzw = args.nchannels_y * args.nsamples_y; + const dim3 block_dims(warp_size_host, nwarps, 1); + const int channel_ratio = args.nchannels_y / args.nchannels_x; + const int sample_ratio = args.nsamples_y / args.nsamples_x; + + CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); + CUDA_SET_SHARED_MEMORY_LIMIT((mul_mat_q), nbytes_shared); + + if (!args.use_stream_k) { + const dim3 grid(nty, ntx, ntzw); + if (args.nrows_x % mmq_y == 0) { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + } else { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + } + return; + } + + // Stream-k + const dim3 grid_sk(nsm, 1, 1); + const bool fixup_needed = ntx * nty * ntzw % nsm != 0; + + + if (args.nrows_x % mmq_y == 0) { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + if (fixup_needed) { + mul_mat_q_stream_k_fixup<<>>( + args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, args.ncols_x, args.nrows_x, args.ncols_dst, + args.nrows_dst, args.nchannels_y, args.stride_channel_dst, args.nsamples_y, args.stride_sample_dst, + args.ncols_max); + } + } else { + mul_mat_q<<>>( + args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, + args.ncols_x, args.nrows_x, args.ncols_dst, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio, args.nchannels_y, args.stride_channel_x, args.stride_channel_y, args.stride_channel_dst, + sample_ratio, args.nsamples_y, args.stride_sample_x, args.stride_sample_y, args.stride_sample_dst, + args.ncols_max); + if (fixup_needed) { + mul_mat_q_stream_k_fixup<<>>( + args.ids_dst, args.expert_bounds, args.dst, tmp_fixup, args.ncols_x, args.nrows_x, args.ncols_dst, + args.nrows_dst, args.nchannels_y, args.stride_channel_dst, args.nsamples_y, args.stride_sample_dst, + args.ncols_max); + } + } +} + +static void launch_mmq_case_q8_0(float * tmp_fixup, const mmq_args & args, cudaStream_t stream, + int cc, int nsm, size_t smpbo, int warp_size_host) { + + const int mmq_x_max = (turing_mma_available(cc)) ? 128 : + (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 64 : 64; + const int mmq_y = (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA) ? 128 : 64; + const int nwarps = 256 / warp_size_host; + + int mmq_x_best = 0; + int ntiles_x_best = INT_MAX; + for (int mmq_x = 8; mmq_x <= mmq_x_max && ntiles_x_best > 1; mmq_x += 8) { + const int granularity = (turing_mma_available(cc) && mmq_x >= 48) ? 16 : 8; + if (mmq_x % granularity != 0) continue; + const size_t nbs = mmq_get_nbytes_shared(mmq_x, mmq_y, cc, warp_size_host, nwarps); + if (nbs > smpbo) continue; + const int ntiles_x = (args.ncols_max + mmq_x - 1) / mmq_x; + if (ntiles_x < ntiles_x_best) { mmq_x_best = mmq_x; ntiles_x_best = ntiles_x; } + } + + switch (mmq_x_best) { + case 8: instantiate_mmq_q8_0< 8>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 16: instantiate_mmq_q8_0< 16>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 24: instantiate_mmq_q8_0< 24>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 32: instantiate_mmq_q8_0< 32>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 40: instantiate_mmq_q8_0< 40>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 48: instantiate_mmq_q8_0< 48>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 56: instantiate_mmq_q8_0< 56>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 64: instantiate_mmq_q8_0< 64>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 72: instantiate_mmq_q8_0< 72>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 80: instantiate_mmq_q8_0< 80>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 88: instantiate_mmq_q8_0< 88>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 96: instantiate_mmq_q8_0< 96>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 104: instantiate_mmq_q8_0<104>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 112: instantiate_mmq_q8_0<112>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 120: instantiate_mmq_q8_0<120>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + case 128: instantiate_mmq_q8_0<128>(tmp_fixup, args, stream, cc, nsm, smpbo, warp_size_host); break; + default: break; + } +} + +extern "C" void launch_mmq_gguf_q8_0( + void *tmp_fixup_ptr, + const void *x, const void *y_q8_1_mmq, void *dst, + int64_t ncols_x, int64_t nrows_x, int64_t ncols_y, + int64_t stride_row_x, int64_t stride_col_dst, + int cc, int nsm, int64_t smpbo, int warp_size_host, + void *stream) { + + const bool use_stream_k = (GGML_CUDA_CC_IS_NVIDIA(cc) && ggml_cuda_highest_compiled_arch(cc) >= GGML_CUDA_CC_VOLTA); + + const mmq_args args = { + (const char *)x, GGML_TYPE_Q8_0, (const int *)y_q8_1_mmq, nullptr, nullptr, (float *)dst, + ncols_x, nrows_x, ncols_y, stride_row_x, ncols_y, nrows_x, + 1, 1, 0, 0, 0, + 1, 1, 0, 0, 0, + use_stream_k, ncols_y + }; + + launch_mmq_case_q8_0((float *)tmp_fixup_ptr, args, (cudaStream_t)stream, cc, nsm, smpbo, warp_size_host); +} diff --git a/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_mma.cuh b/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_mma.cuh new file mode 100644 index 00000000..1c441905 --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_mma.cuh @@ -0,0 +1,1409 @@ +#pragma once +// This file contains primitives that expose the tensor core PTX instructions for CUDA code. +// The primitives can be used in a similar way as the nvcuda::wmma interface but with a well-defined memory layout. +// The documentation for the PTX instructions can be found under: +// https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#matrix-multiply-accumulate-operation-using-mma-instruction +// +// Like with nvcuda::wmma there are three types of matrix tiles: A, B, and C with A @ B = C. +// A is a row-major matrix with shape M x K. +// B is a column-major matrix with shape K x N. +// C is a column-major matrix with shape M x N. +// A, B, and C are represented using the same fundamental data type: a row-major matrix with I rows and J columns. +// Note that J is measured in physical 32 bit elements instead of logical elements. +// The methods get_i and get_j can be used to get the physical 32 bit index of the lth element of a thread within a tile. +// All matrix tiles have ne physical 32 bit elements per warp. +// +// As described in the PTX documentation, all pointers for load_ldmatrix must be to shared memory and aligned to 16 bytes. +// The API in this file also assumes that the pointers for load_generic are aligned to 16 bytes, unaligned pointers are considered undefined behavior. + +#include "mmq_common.cuh" + +// On Volta each warp is doing 4 8x8 mma operations in parallel. +// The basic memory layout for a 32x8 output tile is to stack 4 input tiles in I direction and to mirror the B tile. +// However, the i indices in this file are by default permuted to simplify the index calculations. +// #define GGML_CUDA_MMA_NO_VOLTA_PERM + +#if CUDART_VERSION >= 11080 + +static __device__ __forceinline__ int ggml_cuda_movmatrix(const int x) { + int ret = 0; + +#ifdef TURING_MMA_AVAILABLE + asm("movmatrix.sync.aligned.m8n8.trans.b16 %0, %1;" + : "=r"(ret) : "r"(x)); +#else + GGML_UNUSED(x); + NO_DEVICE_CODE; +#endif // defined(TURING_MMA_AVAILABLE) + return ret; +} + +#else + +static __device__ __forceinline__ int ggml_cuda_movmatrix(const int x) { + // Imagine transposing row-major matrix to column-major matrix. + const int src_i_low = 2 * (threadIdx.x % 4); + const int src_i_high = src_i_low + 1; + const int src_j = threadIdx.x / 4; + + const int src_laneid_low = src_i_low * 4 + src_j / 2; + const int src_laneid_high = src_i_high * 4 + src_j / 2; + + const int shift_low = ((src_j + 0) % 2) * 16; + const int shift_high = ((src_j + 1) % 2) * 16; + + const int ret_low = (__shfl_sync(0xFFFFFFFF, x, src_laneid_low, WARP_SIZE) >> shift_low) & 0x0000FFFF; + const int ret_high = (__shfl_sync(0xFFFFFFFF, x, src_laneid_high, WARP_SIZE) << shift_high) & 0xFFFF0000; + + return ret_low | ret_high; +} + +#endif // CUDART_VERSION >= 11080 + +static __device__ __forceinline__ half2 ggml_cuda_movmatrix(const half2 x) { + half2 ret; + *((int *) &ret) = ggml_cuda_movmatrix(*((const int *) &x)); + return ret; +} + +namespace ggml_cuda_mma { + + // Some architectures like Volta or CDNA3 perform multiple matrix multiplications per warp in parallel, + // effectively the warp is being split into subgroups of threads that each perform a single mma instruction. + // In those cases the data can be split in different ways across the warp. + enum data_layout { + // By default the data uses the I direction as its major dimension and the J direction as its minor dimension. + // For the A/C matrices this means I major == row major, J major == column major. + // For the B matrix this means I major == column major, J major == row major. + // MIRRORED == Each data value is held exactly once per thread subgroup. + DATA_LAYOUT_I_MAJOR = 0, // Always used for Turing, Ampere, Ada Lovelace, consumer Blackwell, matrix A&B for RDNA4 and CDNA. + DATA_LAYOUT_J_MAJOR = 10, // Matrix C for CDNA and RDNA4, int and float matrix C for RDNA3. + DATA_LAYOUT_I_MAJOR_MIRRORED = 20, // Volta, matrix A&B for RDNA3. + DATA_LAYOUT_J_MAJOR_MIRRORED = 30, + }; + // Implemented mma combinations are: + // - (I_MAJOR, I_MAJOR) -> I_MAJOR + // - (I_MAJOR, I_MAJOR_MIRRORED) -> I_MAJOR + // - (I_MAJOR, J_MAJOR_MIRRORED) -> I_MAJOR + + static constexpr bool is_i_major(const data_layout dl) { + return dl == DATA_LAYOUT_I_MAJOR || + dl == DATA_LAYOUT_I_MAJOR_MIRRORED; + } + + static constexpr __device__ data_layout get_input_data_layout() { +#if defined(RDNA3) || __CUDA_ARCH__ == GGML_CUDA_CC_VOLTA + return DATA_LAYOUT_I_MAJOR_MIRRORED; +#else + return DATA_LAYOUT_I_MAJOR; +#endif // defined(RDNA3) || __CUDA_ARCH__ == GGML_CUDA_CC_VOLTA + } + + template + struct tile {}; + + template + struct tile { + static constexpr int I = I_; + static constexpr int J = J_; + static constexpr data_layout dl = DATA_LAYOUT_I_MAJOR; + +#if defined(AMD_MFMA_AVAILABLE) + static constexpr int ne = I * J / 64; + T x[ne] = {0}; + + static constexpr __device__ bool supported() { + if (I == 64 && J == 2) return true; + if (I == 16 && J == 8) return true; + if (I == 32 && J == 4) return true; + if (I == 16 && J == 16) return true; + if (I == 32 && J == 32) return true; + return false; + } + + static __device__ __forceinline__ int get_i(const int l) { + if constexpr (I == 64 && J == 2) { // Special tile size to load <16, 4> as <16, 8> + return threadIdx.x % 16; + } else if constexpr (I == 16 && J == 8) { + return threadIdx.x % 16; + } else if constexpr (I == 32 && J == 4) { + return threadIdx.x % 32; + } else if constexpr (I == 16 && J == 16) { + return threadIdx.x % 16; + } else if constexpr (I == 32 && J == 32) { + return threadIdx.x % 32; + } else { + NO_DEVICE_CODE; + return -1; + } + } + + static __device__ __forceinline__ int get_j(const int l) { + if constexpr (I == 64 && J == 2) { // Special tile size to load <16, 4> as <16, 8> + return (2 * ((threadIdx.x / 16) % 2) + l); + } else if constexpr (I == 16 && J == 8) { + return 2 * (threadIdx.x / 16) + l; + } else if constexpr (I == 32 && J == 4) { + return 2 * (threadIdx.x / 32) + l; + } else if constexpr (I == 16 && J == 16) { + return 4 * (threadIdx.x / 16) + l; + } else if constexpr (I == 32 && J == 32) { + return 4 * (threadIdx.x / 32) + 8 * (l / 4) + (l % 4); + } else { + NO_DEVICE_CODE; + return -1; + } + } +#elif __CUDA_ARCH__ == GGML_CUDA_CC_VOLTA + static constexpr int ne = I * J / 32; + T x[ne] = {0}; + + static constexpr __device__ bool supported() { + if (I == 32 && J == 8) return true; + return false; + } + + static __device__ __forceinline__ int get_i(const int l) { + if constexpr (I == 32 && J == 8) { +#ifdef GGML_CUDA_MMA_NO_VOLTA_PERM + return (((threadIdx.x % 16) / 4) * 8) + ((threadIdx.x / 16) * 4) + (l & 2) + (threadIdx.x % 2); +#else + return (l & 2) + (threadIdx.x & ~2); +#endif // GGML_CUDA_MMA_NO_VOLTA_PERM + } else { + NO_DEVICE_CODE; + return -1; + } + } + + static __device__ __forceinline__ int get_j(const int l) { + if constexpr (I == 32 && J == 8) { + return (threadIdx.x & 2) + (l & (4 + 1)); + } else { + NO_DEVICE_CODE; + return -1; + } + } +#elif defined(AMD_WMMA_AVAILABLE) + static constexpr int ne = I * J / 32; + T x[ne] = {0}; + + static constexpr __device__ bool supported() { + if (I == 16 && J == 16) return true; + if (I == 16 && J == 8) return true; + if (I == 16 && J == 4) return true; + return false; + } + + static __device__ __forceinline__ int get_i(const int l) { + if constexpr (supported()) { + return threadIdx.x % 16; + } else { + NO_DEVICE_CODE; + return -1; + } + } + + static __device__ __forceinline__ int get_j(const int l) { + if constexpr (I == 16 && J == 16) { +#if defined(RDNA3) + if constexpr (std::is_same_v || std::is_same_v) { + // matrix C + return 2 * l + (threadIdx.x / 16); + } else { + // matrix A&B + return l; + } +#else + // matrix C is the transposed matrix A&B on RDNA4 + return ne * (threadIdx.x / 16) + l; +#endif // defined(RDNA3) + } else if constexpr (I == 16 && J == 8) { + // mmq input for RDNA4 + return ne * (threadIdx.x / 16) + l; + } else if constexpr (I == 16 && J == 4) { + return ne * (threadIdx.x / 16) + l; + } else { + NO_DEVICE_CODE; + return -1; + } + } +#else + static constexpr int ne = I * J / 32; + T x[ne] = {0}; + + static constexpr __device__ bool supported() { + if (I == 8 && J == 4) return true; + if (I == 8 && J == 8) return true; + if (I == 16 && J == 8) return true; + if (I == 16 && J == 16) return true; + if (I == 32 && J == 8) return true; + return false; + } + + static __device__ __forceinline__ int get_i(const int l) { + if constexpr (I == 8 && J == 4) { + return threadIdx.x / 4; + } else if constexpr (I == 8 && J == 8) { + return threadIdx.x / 4; + } else if constexpr (I == 16 && J == 8) { + return ((l / 2) * 8) + (threadIdx.x / 4); + } else if constexpr (I == 16 && J == 16) { + return (((l / 2) % 2) * 8) + (threadIdx.x / 4); + } else if constexpr (I == 32 && J == 8) { + return tile<16, 8, T>::get_i(l); // Memory layout simply repeated with same pattern in i direction. + } else { + NO_DEVICE_CODE; + return -1; + } + } + + static __device__ __forceinline__ int get_j(const int l) { + if constexpr (I == 8 && J == 4) { + return threadIdx.x % 4; + } else if constexpr (I == 8 && J == 8) { + return (l * 4) + (threadIdx.x % 4); + } else if constexpr (I == 16 && J == 8) { + return ((threadIdx.x % 4) * 2) + (l % 2); + } else if constexpr (I == 16 && J == 16) { + return ((l / 4) * 8) + ((threadIdx.x % 4) * 2) + (l % 2); + } else if constexpr (I == 32 && J == 8) { + return tile<16, 8, T>::get_j(l); // Memory layout simply repeated with same pattern in i direction. + } else { + NO_DEVICE_CODE; + return -1; + } + } +#endif // defined(GGML_USE_HIP) + }; + + template + struct tile { + static constexpr int I = I_; + static constexpr int J = J_; + static constexpr data_layout dl = DATA_LAYOUT_I_MAJOR; + +#if __CUDA_ARCH__ == GGML_CUDA_CC_VOLTA + static constexpr int ne = I * J / WARP_SIZE; + half2 x[ne] = {{0.0f, 0.0f}}; + + static constexpr __device__ bool supported() { + if (I == 32 && J == 4) return true; + return false; + } + + static __device__ __forceinline__ int get_i(const int l) { + if constexpr (I == 32 && J == 4) { +#ifdef GGML_CUDA_MMA_NO_VOLTA_PERM + return (((threadIdx.x % 16) / 4) * 8) + ((threadIdx.x / 16) * 4) + (threadIdx.x % 4); +#else + return threadIdx.x; +#endif // GGML_CUDA_MMA_NO_VOLTA_PERM + } else { + NO_DEVICE_CODE; + return -1; + } + } + + static __device__ __forceinline__ int get_j(const int l) { + if constexpr (I == 32 && J == 4) { + return l; + } else { + NO_DEVICE_CODE; + return -1; + } + } +#elif defined(AMD_WMMA_AVAILABLE) + static constexpr int ne = I * J / 32; + half2 x[ne] = {{0.0f, 0.0f}}; + + static constexpr __device__ bool supported() { + if (I == 16 && J == 8) return true; + return false; + } + + static __device__ __forceinline__ int get_i(const int l) { + if constexpr (I == 16 && J == 8) { + return threadIdx.x % 16; + } else { + NO_DEVICE_CODE; + return -1; + } + } + + static __device__ __forceinline__ int get_j(const int l) { + if constexpr (I == 16 && J == 8) { + return ne * (threadIdx.x / 16) + l; + } else { + NO_DEVICE_CODE; + return -1; + } + } +#elif defined(AMD_MFMA_AVAILABLE) + static constexpr int ne = I * J / 64; + half2 x[ne] = {{0.0f, 0.0f}}; + + static constexpr __device__ bool supported() { + if (I == 16 && J == 8) return true; + return false; + } + + static __device__ __forceinline__ int get_i(const int l) { + if constexpr (I == 16 && J == 8) { + return threadIdx.x % 16; + } else { + NO_DEVICE_CODE; + return -1; + } + } + + static __device__ __forceinline__ int get_j(const int l) { + if constexpr (I == 16 && J == 8) { + return ne * (threadIdx.x / 16) + l; + } else { + NO_DEVICE_CODE; + return -1; + } + } +#else + static constexpr int ne = I * J / WARP_SIZE; + half2 x[ne] = {{0.0f, 0.0f}}; + + static constexpr __device__ bool supported() { + if (I == 8 && J == 4) return true; + if (I == 8 && J == 8) return true; + if (I == 16 && J == 8) return true; + if (I == 16 && J == 16) return true; + if (I == 32 && J == 8) return true; + return false; + } + + static __device__ __forceinline__ int get_i(const int l) { + if constexpr (I == 8 && J == 8) { + return threadIdx.x / 4; + } else if constexpr (I == 16 && J == 4) { + return (l * 8) + (threadIdx.x / 4); + } else if constexpr (I == 16 && J == 8) { + return ((l % 2) * 8) + (threadIdx.x / 4); + } else if constexpr (I == 32 && J == 8) { + return ((l / 4) * 16) + ((l % 2) * 8) + (threadIdx.x / 4); + } else { + NO_DEVICE_CODE; + return -1; + } + } + + static __device__ __forceinline__ int get_j(const int l) { + if constexpr (I == 8 && J == 8) { + return (l * 4) + (threadIdx.x % 4); + } else if constexpr (I == 16 && J == 4) { + return threadIdx.x % 4; + } else if constexpr (I == 16 && J == 8) { + return ((l / 2) * 4) + (threadIdx.x % 4); + } else if constexpr (I == 32 && J == 8) { + return ((l & 2) * 2) + (threadIdx.x % 4); + } else { + NO_DEVICE_CODE; + return -1; + } + } +#endif // __CUDA_ARCH__ == GGML_CUDA_CC_VOLTA + }; + + template + struct tile { + static constexpr int I = I_; + static constexpr int J = J_; + static constexpr data_layout dl = DATA_LAYOUT_I_MAJOR; + +#if defined(AMD_WMMA_AVAILABLE) + static constexpr int ne = tile::ne; + nv_bfloat162 x[ne] = {{0.0f, 0.0f}}; + + static constexpr __device__ bool supported() { + return tile::supported(); + } + + static __device__ __forceinline__ int get_i(const int l) { + return tile::get_i(l); + } + + static __device__ __forceinline__ int get_j(const int l) { + return tile::get_j(l); + } +#elif defined(AMD_MFMA_AVAILABLE) + static constexpr int ne = tile::ne; + nv_bfloat162 x[ne] = {{0.0f, 0.0f}}; + + static constexpr __device__ bool supported() { + return tile::supported(); + } + + static __device__ __forceinline__ int get_i(const int l) { + return tile::get_i(l); + } + + static __device__ __forceinline__ int get_j(const int l) { + return tile::get_j(l); + } +#else + static constexpr int ne = I * J / WARP_SIZE; + nv_bfloat162 x[ne] = {{0.0f, 0.0f}}; + + static constexpr __device__ bool supported() { + if (I == 8 && J == 8) return true; + if (I == 16 && J == 4) return true; + if (I == 16 && J == 8) return true; + return false; + } + + static __device__ __forceinline__ int get_i(const int l) { + if constexpr (I == 8 && J == 8) { + return threadIdx.x / 4; + } else if constexpr (I == 16 && J == 4) { + return (l * 8) + (threadIdx.x / 4); + } else if constexpr (I == 16 && J == 8) { + return ((l % 2) * 8) + (threadIdx.x / 4); + } else { + NO_DEVICE_CODE; + return -1; + } + } + + static __device__ __forceinline__ int get_j(const int l) { + if constexpr (I == 8 && J == 8) { + return (l * 4) + (threadIdx.x % 4); + } else if constexpr (I == 16 && J == 4) { + return threadIdx.x % 4; + } else if constexpr (I == 16 && J == 8) { + return ((l / 2) * 4) + (threadIdx.x % 4); + } else { + NO_DEVICE_CODE; + return -1; + } + } +#endif // defined(AMD_WMMA_AVAILABLE) + }; + + template + struct tile { + static constexpr int I = I_; + static constexpr int J = J_; + static constexpr data_layout dl = DATA_LAYOUT_J_MAJOR; + + static constexpr int ne = tile::ne; + T x[ne] = {0}; + + static constexpr __device__ bool supported() { + return tile::supported(); + } + + static __device__ __forceinline__ int get_i(const int l) { + return tile::get_j(l); + } + + static __device__ __forceinline__ int get_j(const int l) { + return tile::get_i(l); + } + }; + + template + struct tile { + static constexpr int I = I_; + static constexpr int J = J_; + static constexpr data_layout dl = DATA_LAYOUT_I_MAJOR_MIRRORED; + + // RDNA3 + static constexpr int ne = I * J / 32 * 2; + + T x[ne] = {0}; + + static constexpr __device__ bool supported() { + if (I == 16 && J == 16) return true; + if (I == 16 && J == 8) return true; + if (I == 16 && J == 4) return true; + return false; + } + + static __device__ __forceinline__ int get_i(const int /*l*/) { + if constexpr (supported()) { + return threadIdx.x % 16; + } else { + NO_DEVICE_CODE; + return -1; + } + } + + static __device__ __forceinline__ int get_j(const int l) { + if constexpr (supported()) { + return l; + } else { + NO_DEVICE_CODE; + return -1; + } + } + }; + + template + struct tile { + static constexpr int I = I_; + static constexpr int J = J_; + static constexpr data_layout dl = DATA_LAYOUT_I_MAJOR_MIRRORED; +#if defined(RDNA3) + static constexpr int ne = tile::ne; + + half2 x[ne] = {{0.0f, 0.0f}}; + + static constexpr __device__ bool supported() { + return tile::supported(); + } + + static __device__ __forceinline__ int get_i(const int l) { + return tile::get_i(l); + } + + static __device__ __forceinline__ int get_j(const int l) { + return tile::get_j(l); + } +#else // Volta + static constexpr int ne = I * J / (WARP_SIZE/4); + + half2 x[ne] = {{0.0f, 0.0f}}; + + static constexpr __device__ bool supported() { + if (I == 8 && J == 4) return true; + return false; + } + + static __device__ __forceinline__ int get_i(const int /*l*/) { + if constexpr (I == 8 && J == 4) { + return ((threadIdx.x / 16) * 4) + (threadIdx.x % 4); + } else { + NO_DEVICE_CODE; + return -1; + } + } + + static __device__ __forceinline__ int get_j(const int l) { + if constexpr (I == 8 && J == 4) { + return l; + } else { + NO_DEVICE_CODE; + return -1; + } + } +#endif // defined(RDNA3) + }; + + template + struct tile { + static constexpr int I = I_; + static constexpr int J = J_; + static constexpr data_layout dl = DATA_LAYOUT_I_MAJOR_MIRRORED; + static constexpr int ne = tile::ne; + + nv_bfloat162 x[ne] = {{0.0f, 0.0f}}; + + static constexpr __device__ bool supported() { + return tile::supported(); + } + + static __device__ __forceinline__ int get_i(const int l) { + return tile::get_i(l); + } + + static __device__ __forceinline__ int get_j(const int l) { + return tile::get_j(l); + } + }; + + template + struct tile { + static constexpr int I = I_; + static constexpr int J = J_; + static constexpr data_layout dl = DATA_LAYOUT_J_MAJOR_MIRRORED; + static constexpr int ne = I * J / (WARP_SIZE/4); + + half2 x[ne] = {{0.0f, 0.0f}}; + + static constexpr __device__ bool supported() { + if (I == 8 && J == 4) return true; + return false; + } + + static __device__ __forceinline__ int get_i(const int l) { + if constexpr (I == 8 && J == 4) { + return ((l / 2) * 4) + (threadIdx.x % 4); + } else { + NO_DEVICE_CODE; + return -1; + } + } + + static __device__ __forceinline__ int get_j(const int l) { + if constexpr (I == 8 && J == 4) { + return ((threadIdx.x / 16) * 2) + (l % 2); + } else { + NO_DEVICE_CODE; + return -1; + } + } + }; + +#if defined(TURING_MMA_AVAILABLE) + template + static __device__ __forceinline__ tile get_half2(const tile & tile_float) { + tile ret; +#pragma unroll + for (int l0 = 0; l0 < tile_float.ne; l0 += 2) { + ret.x[l0/2] = make_half2(tile_float.x[l0 + 0], tile_float.x[l0 + 1]); + } + return ret; + } + + static __device__ __forceinline__ tile<8, 8, half2> get_transposed(const tile<16, 4, half2> & t) { + tile<8, 8, half2> ret; + ret.x[0] = ggml_cuda_movmatrix(t.x[0]); + ret.x[1] = ggml_cuda_movmatrix(t.x[1]); + + return ret; + } +#elif defined(AMD_WMMA_AVAILABLE) || defined(AMD_MFMA_AVAILABLE) + template + static __device__ __forceinline__ tile get_half2(const tile & tile_float) { + tile ret; +#pragma unroll + for (int l0 = 0; l0 < tile_float.ne; l0 += 2) { + ret.x[l0/2] = make_half2(tile_float.x[l0 + 0], tile_float.x[l0 + 1]); + } + return ret; + } + + static __device__ __forceinline__ tile<8, 8, half2> get_transposed(const tile<16, 4, half2> & t) { + NO_DEVICE_CODE; + return tile<8, 8, half2>{}; + } +#else // Volta + template + static __device__ __forceinline__ tile get_half2(const tile & tile_float) { + tile ret; +#pragma unroll + for (int l0 = 0; l0 < tile_float.ne; l0 += 4) { + ret.x[l0/2 + 0] = make_half2(tile_float.x[l0 + 0], tile_float.x[l0 + 1]); + ret.x[l0/2 + 1] = make_half2(tile_float.x[l0 + 2], tile_float.x[l0 + 3]); + + // On Volta FP16 and FP32 tiles have a different memory layout, + // for the conversion threads with an offset of 2 need to exchange half their values: + ret.x[l0/2 + (((threadIdx.x % 4) / 2) ^ 1)] = __shfl_xor_sync( + 0xFFFFFFFF, ret.x[l0/2 + (((threadIdx.x % 4) / 2) ^ 1)], 2, WARP_SIZE); + } + return ret; + } +#endif // defined(TURING_MMA_AVAILABLE) + + static __device__ __forceinline__ void make_identity_mat(tile<16, 8, half2> & t) { +#if defined(RDNA4) + const int row = t.get_i(0); + const int left_right = t.get_j(0) / 4; + const int up_down = row / 8; + const int idx = row % 8; + reinterpret_cast(t.x)[idx] = left_right == up_down ? 1.0f : 0.0f; +#else + GGML_UNUSED_VARS(t); + NO_DEVICE_CODE; +#endif // defined(RDNA4) + } + + template + static __device__ __forceinline__ void load_generic(tile & t, const T * __restrict__ xs0, const int stride) { +#if defined(AMD_MFMA_AVAILABLE) + if constexpr (I == 64 && J == 2) { // Special tile size to load <16, 4> as <16, 8> +#pragma unroll + for (int l = 0; l < t.ne; ++l) { + t.x[l] = xs0[t.get_i(l)*stride + t.get_j(l)]; + } + } else { + ggml_cuda_memcpy_1(t.x, xs0 + t.get_i(0) * stride + t.get_j(0)); + } +#elif defined(AMD_WMMA_AVAILABLE) + // All wmma layout has contiguous data when i-major. + if constexpr (is_i_major(dl)) { + // the data must be aligned to 16 bytes when bigger than ggml_cuda_get_max_cpy_bytes() + constexpr int aligned_copy_bytes = ggml_cuda_get_max_cpy_bytes(); + if constexpr (sizeof(t.x) > aligned_copy_bytes) { + static_assert(sizeof(t.x) % aligned_copy_bytes == 0, "bad type size"); + constexpr int aligned_copy_count = sizeof(t.x)/aligned_copy_bytes; +#pragma unroll + for (int i = 0; i < aligned_copy_count; ++i) { + ggml_cuda_memcpy_1(t.x + t.ne/aligned_copy_count*i, xs0 + t.get_i(0) * stride + t.get_j(t.ne/aligned_copy_count*i)); + } + } else { + ggml_cuda_memcpy_1(t.x, xs0 + t.get_i(0) * stride + t.get_j(0)); + } + } else { +#pragma unroll + for (int l = 0; l < t.ne; ++l) { + t.x[l] = xs0[t.get_i(l)*stride + t.get_j(l)]; + } + } +#else +#pragma unroll + for (int l = 0; l < t.ne; ++l) { + t.x[l] = xs0[t.get_i(l)*stride + t.get_j(l)]; + } +#endif // defined(AMD_MFMA_AVAILABLE) + } + + template + static __device__ __forceinline__ void load_ldmatrix( + tile<8, 8, T> & t, const T * __restrict__ xs0, const int stride) { +#ifdef TURING_MMA_AVAILABLE + int * xi = (int *) t.x; + const int * xs = (const int *) xs0 + (threadIdx.x % t.I) * stride + ((threadIdx.x / t.I) * (t.J / 2)) % t.J; + asm volatile("ldmatrix.sync.aligned.m8n8.x2.b16 {%0, %1}, [%2];" + : "=r"(xi[0]), "=r"(xi[1]) + : "l"(xs)); +#else + load_generic(t, xs0, stride); +#endif // TURING_MMA_AVAILABLE + } + + template + static __device__ __forceinline__ void load_ldmatrix( + tile<16, 4, T> & t, const T * __restrict__ xs0, const int stride) { +#ifdef TURING_MMA_AVAILABLE + int * xi = (int *) t.x; + const int * xs = (const int *) xs0 + (threadIdx.x % t.I) * stride; + asm volatile("ldmatrix.sync.aligned.m8n8.x2.b16 {%0, %1}, [%2];" + : "=r"(xi[0]), "=r"(xi[1]) + : "l"(xs)); +#else +#if __CUDA_ARCH__ == GGML_CUDA_CC_VOLTA + GGML_UNUSED_VARS(t, xs0, stride); + NO_DEVICE_CODE; +#else + load_generic(t, xs0, stride); +#endif // __CUDA_ARCH__ == GGML_CUDA_CC_VOLTA +#endif // TURING_MMA_AVAILABLE + } + + template + static __device__ __forceinline__ void load_ldmatrix( + tile<16, 8, T, dl> & t, const T * __restrict__ xs0, const int stride) { +#if defined(TURING_MMA_AVAILABLE) + int * xi = (int * ) t.x; + const int * xs = (const int *) xs0 + (threadIdx.x % t.I) * stride + (threadIdx.x / t.I) * (t.J / 2); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.b16 {%0, %1, %2, %3}, [%4];" + : "=r"(xi[0]), "=r"(xi[1]), "=r"(xi[2]), "=r"(xi[3]) + : "l"(xs)); +#else +#if __CUDA_ARCH__ == GGML_CUDA_CC_VOLTA +#if 1 + // TODO: more generic handling + static_assert(sizeof(T) == 4, "bad type size"); + ggml_cuda_memcpy_1<4*sizeof(T)>(t.x + 0, xs0 + t.get_i(0)*stride + 0); + ggml_cuda_memcpy_1<4*sizeof(T)>(t.x + 4, xs0 + t.get_i(4)*stride + 4); +#else + load_generic(t, xs0, stride); +#endif // 1 +#else + load_generic(t, xs0, stride); +#endif // __CUDA_ARCH__ == GGML_CUDA_CC_VOLTA +#endif // TURING_MMA_AVAILABLE + } + + static __device__ __forceinline__ void load_ldmatrix( + tile<8, 4, half2, DATA_LAYOUT_I_MAJOR_MIRRORED> & t, const half2 * __restrict__ xs0, const int stride) { + ggml_cuda_memcpy_1<4*sizeof(half2)>(t.x, xs0 + t.get_i(0)*stride); + } + + static __device__ __forceinline__ void load_ldmatrix( + tile<8, 4, half2, DATA_LAYOUT_J_MAJOR_MIRRORED> & t, const half2 * __restrict__ xs0, const int stride) { +#pragma unroll + for (int l0 = 0; l0 < t.ne; l0 += 2) { + ggml_cuda_memcpy_1<2*sizeof(half2)>(t.x + l0, xs0 + t.get_i(l0)*stride + t.get_j(l0)); + } + } + + static __device__ __forceinline__ void load_ldmatrix( + tile<32, 4, half2> & t, const half2 * __restrict__ xs0, const int stride) { +#if __CUDA_ARCH__ == GGML_CUDA_CC_VOLTA + ggml_cuda_memcpy_1<4*sizeof(half2)>(t.x, xs0 + t.get_i(0)*stride); +#else + GGML_UNUSED_VARS(t, xs0, stride); + NO_DEVICE_CODE; +#endif // __CUDA_ARCH__ == GGML_CUDA_CC_VOLTA + } + + template + static __device__ __forceinline__ void load_ldmatrix_trans( + tile<16, 8, T> & t, const T * __restrict__ xs0, const int stride) { +#ifdef TURING_MMA_AVAILABLE + int * xi = (int * ) t.x; + const int * xs = (const int *) xs0 + (threadIdx.x % t.I) * stride + (threadIdx.x / t.I) * (t.J / 2); + asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.b16 {%0, %1, %2, %3}, [%4];" + : "=r"(xi[0]), "=r"(xi[2]), "=r"(xi[1]), "=r"(xi[3]) + : "l"(xs)); +#else + GGML_UNUSED_VARS(t, xs0, stride); + NO_DEVICE_CODE; +#endif // TURING_MMA_AVAILABLE + } + + static __device__ __forceinline__ void mma( + tile<16, 8, int> & D, const tile<16, 4, int> & A, const tile<8, 4, int> & B) { +#ifdef TURING_MMA_AVAILABLE +#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE + asm("mma.sync.aligned.m16n8k16.row.col.s32.s8.s8.s32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {%0, %1, %2, %3};" + : "+r"(D.x[0]), "+r"(D.x[1]), "+r"(D.x[2]), "+r"(D.x[3]) + : "r"(A.x[0]), "r"(A.x[1]), "r"(B.x[0])); +#else + // On Turing m16n8k16 mma is not available, use 2x m8n8k16 mma instead: + asm("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 {%0, %1}, {%2}, {%3}, {%0, %1};" + : "+r"(D.x[0]), "+r"(D.x[1]) + : "r"(A.x[0]), "r"(B.x[0])); + asm("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 {%0, %1}, {%2}, {%3}, {%0, %1};" + : "+r"(D.x[2]), "+r"(D.x[3]) + : "r"(A.x[1]), "r"(B.x[0])); +#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // TURING_MMA_AVAILABLE + } + + static __device__ __forceinline__ void mma( + tile<16, 8, int> & D, const tile<16, 8, int> & A, const tile<8, 8, int> & B) { +#ifdef TURING_MMA_AVAILABLE +#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE + asm("mma.sync.aligned.m16n8k32.row.col.s32.s8.s8.s32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};" + : "+r"(D.x[0]), "+r"(D.x[1]), "+r"(D.x[2]), "+r"(D.x[3]) + : "r"(A.x[0]), "r"(A.x[1]), "r"(A.x[2]), "r"(A.x[3]), "r"(B.x[0]), "r"(B.x[1])); +#else + // On Turing m16n8k32 mma is not available, use 4x m8n8k16 mma instead: + asm("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 {%0, %1}, {%2}, {%3}, {%0, %1};" + : "+r"(D.x[0]), "+r"(D.x[1]) + : "r"(A.x[0]), "r"(B.x[0])); + asm("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 {%0, %1}, {%2}, {%3}, {%0, %1};" + : "+r"(D.x[2]), "+r"(D.x[3]) + : "r"(A.x[1]), "r"(B.x[0])); + asm("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 {%0, %1}, {%2}, {%3}, {%0, %1};" + : "+r"(D.x[0]), "+r"(D.x[1]) + : "r"(A.x[2]), "r"(B.x[1])); + asm("mma.sync.aligned.m8n8k16.row.col.s32.s8.s8.s32 {%0, %1}, {%2}, {%3}, {%0, %1};" + : "+r"(D.x[2]), "+r"(D.x[3]) + : "r"(A.x[3]), "r"(B.x[1])); +#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // TURING_MMA_AVAILABLE + } + + static __device__ __forceinline__ void mma( + tile<16, 4, half2> & D, const tile<16, 8, half2> & A, const tile<8, 8, half2> & B) { +#ifdef TURING_MMA_AVAILABLE + const int * Axi = (const int *) A.x; + const int * Bxi = (const int *) B.x; + int * Dxi = (int *) D.x; +#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE + asm("mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3, %4, %5}, {%6, %7}, {%0, %1};" + : "+r"(Dxi[0]), "+r"(Dxi[1]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[0]), "r"(Bxi[1])); +#else + // On Turing m16n8k16 mma is not available, use 2x m8n8k8 mma instead: + asm("mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3}, {%4}, {%0, %1};" + : "+r"(Dxi[0]), "+r"(Dxi[1]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Bxi[0])); + asm("mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3}, {%4}, {%0, %1};" + : "+r"(Dxi[0]), "+r"(Dxi[1]) + : "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[1])); +#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // TURING_MMA_AVAILABLE + } + + static __device__ __forceinline__ void mma( + tile<16, 8, half2> & D, const tile<16, 8, half2> & A, const tile<16, 8, half2> & B) { +#ifdef TURING_MMA_AVAILABLE + const int * Axi = (const int *) A.x; + const int * Bxi = (const int *) B.x; + int * Dxi = (int *) D.x; +#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE + asm("mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3, %4, %5}, {%6, %7}, {%0, %1};" + : "+r"(Dxi[0]), "+r"(Dxi[1]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[0]), "r"(Bxi[2])); + asm("mma.sync.aligned.m16n8k16.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3, %4, %5}, {%6, %7}, {%0, %1};" + : "+r"(Dxi[2]), "+r"(Dxi[3]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[1]), "r"(Bxi[3])); +#else + // On Turing m16n8k16 mma is not available, use 4x m8n8k8 mma instead: + asm("mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3}, {%4}, {%0, %1};" + : "+r"(Dxi[0]), "+r"(Dxi[1]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Bxi[0])); + asm("mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3}, {%4}, {%0, %1};" + : "+r"(Dxi[0]), "+r"(Dxi[1]) + : "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[2])); + asm("mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3}, {%4}, {%0, %1};" + : "+r"(Dxi[2]), "+r"(Dxi[3]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Bxi[1])); + asm("mma.sync.aligned.m16n8k8.row.col.f16.f16.f16.f16 {%0, %1}, {%2, %3}, {%4}, {%0, %1};" + : "+r"(Dxi[2]), "+r"(Dxi[3]) + : "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[3])); +#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE +#elif defined(AMD_WMMA_AVAILABLE) +#if defined(RDNA4) + using halfx8_t = __attribute__((ext_vector_type(8))) _Float16; + halfx8_t& acc_frag = reinterpret_cast(D.x[0]); + const halfx8_t& a_frag = reinterpret_cast(A.x[0]); + const halfx8_t& b_frag = reinterpret_cast(B.x[0]); + acc_frag = __builtin_amdgcn_wmma_f16_16x16x16_f16_w32_gfx12(a_frag, b_frag, acc_frag); +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // defined(RDNA4) +#elif defined(AMD_MFMA_AVAILABLE) + // MFMA: FP16 input, FP32 accumulate, convert back to half2. + using halfx4_t = __attribute__((ext_vector_type(4))) _Float16; + using floatx4_t = __attribute__((ext_vector_type(4))) float; + + // Convert existing half2 accumulator to float for MFMA: + floatx4_t acc_f32; + { + const halfx4_t acc_h = reinterpret_cast(D.x[0]); +#pragma unroll + for (int i = 0; i < 4; ++i) { + acc_f32[i] = (float)acc_h[i]; + } + } + + const halfx4_t& a_frag = reinterpret_cast(A.x[0]); + const halfx4_t& b_frag = reinterpret_cast(B.x[0]); + acc_f32 = __builtin_amdgcn_mfma_f32_16x16x16f16(a_frag, b_frag, acc_f32, 0, 0, 0); + + // Convert back to half2: + { + halfx4_t result_h; +#pragma unroll + for (int i = 0; i < 4; ++i) { + result_h[i] = (_Float16)acc_f32[i]; + } + reinterpret_cast(D.x[0]) = result_h; + } +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // TURING_MMA_AVAILABLE + } + + template + static __device__ __forceinline__ void mma( + tile<16, 8, float, dl_d> & D, const tile<16, 8, float, dl_ab> & A, const tile<8, 8, float, dl_ab> & B) { +#ifdef AMPERE_MMA_AVAILABLE + const int * Axi = (const int *) A.x; + const int * Bxi = (const int *) B.x; + int * Dxi = (int *) D.x; + asm("mma.sync.aligned.m16n8k8.row.col.f32.tf32.tf32.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};" + : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[0]), "r"(Bxi[1])); +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // AMPERE_MMA_AVAILABLE + } + + template + static __device__ __forceinline__ void mma( + tile<16, 16, float, dl_d> & D, const tile<16, 8, float, dl_ab> & A, const tile<16, 8, float, dl_ab> & B) { +#ifdef AMD_MFMA_AVAILABLE + using floatx4_t = __attribute__((ext_vector_type(4))) float; + floatx4_t& acc_frag = reinterpret_cast(D.x[0]); +#if defined(CDNA3) + using floatx2_t = __attribute__((ext_vector_type(2))) float; + const floatx2_t& a_frag = reinterpret_cast(A.x[0]); + const floatx2_t& b_frag = reinterpret_cast(B.x[0]); + acc_frag = __builtin_amdgcn_mfma_f32_16x16x8_xf32(a_frag, b_frag, acc_frag, 0, 0, 0); +#elif defined(CDNA2) || defined(CDNA1) +#pragma unroll + for (int i = 0; i < 2; ++i) { + acc_frag = __builtin_amdgcn_mfma_f32_16x16x4f32(A.x[i], B.x[i], acc_frag, 0, 0, 0); + } +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // defined(CDNA3) +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // AMD_MFMA_AVAILABLE + } + + static __device__ __forceinline__ void mma_block_scaled(tile<16, 8, float> & D, + const tile<16, 8, int> & A, + const tile<8, 8, int> & B, + uint32_t a_scale, + uint32_t b_scale) { +#ifdef BLACKWELL_MMA_AVAILABLE + const int * Axi = (const int *) A.x; + const int * Bxi = (const int *) B.x; + float * Dxi = (float *) D.x; + + asm volatile( + "mma.sync.aligned.kind::mxf4.block_scale.scale_vec::2X.m16n8k64.row.col.f32.e2m1.e2m1.f32.ue8m0 " + "{%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3}, " + "%10, {0, 0}, %11, {0, 0};" + : "+f"(Dxi[0]), "+f"(Dxi[1]), "+f"(Dxi[2]), "+f"(Dxi[3]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[0]), "r"(Bxi[1]), "r"(a_scale), "r"(b_scale)); +#else + GGML_UNUSED_VARS(D, A, B, a_scale, b_scale); +#endif // BLACKWELL_MMA_AVAILABLE + } + + static __device__ __forceinline__ void mma( + tile<16, 8, float> & D, const tile<16, 8, half2> & A, const tile<8, 8, half2> & B) { +#ifdef TURING_MMA_AVAILABLE + const int * Axi = (const int *) A.x; + const int * Bxi = (const int *) B.x; + int * Dxi = (int *) D.x; +#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE + asm("mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};" + : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[0]), "r"(Bxi[1])); +#else + // On Turing m16n8k16 mma is not available, use 2x m8n8k8 mma instead: + asm("mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {%0, %1, %2, %3};" + : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Bxi[0])); + asm("mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {%0, %1, %2, %3};" + : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]) + : "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[1])); +#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // TURING_MMA_AVAILABLE + } + + static __device__ __forceinline__ void mma( + tile<16, 8, float> & D, const tile<16, 8, nv_bfloat162> & A, const tile<8, 8, nv_bfloat162> & B) { +#ifdef AMPERE_MMA_AVAILABLE + const int * Axi = (const int *) A.x; + const int * Bxi = (const int *) B.x; + int * Dxi = (int *) D.x; + asm("mma.sync.aligned.m16n8k16.row.col.f32.bf16.bf16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};" + : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[0]), "r"(Bxi[1])); +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // AMPERE_MMA_AVAILABLE + } + + template + static __device__ __forceinline__ void mma( + tile<16, 16, float, dl_d> & D, const tile<16, 8, half2, dl_ab> & A, const tile<16, 8, half2, dl_ab> & B) { +#ifdef TURING_MMA_AVAILABLE + const int * Axi = (const int *) A.x; + const int * Bxi = (const int *) B.x; + int * Dxi = (int *) D.x; +#if __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE + asm("mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};" + : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[0]), "r"(Bxi[2])); + asm("mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5, %6, %7}, {%8, %9}, {%0, %1, %2, %3};" + : "+r"(Dxi[4]), "+r"(Dxi[5]), "+r"(Dxi[6]), "+r"(Dxi[7]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[1]), "r"(Bxi[3])); +#else + // On Turing m16n8k16 mma is not available, use 4x m8n8k8 mma instead: + asm("mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {%0, %1, %2, %3};" + : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Bxi[0])); + asm("mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {%0, %1, %2, %3};" + : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]) + : "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[2])); + asm("mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {%0, %1, %2, %3};" + : "+r"(Dxi[4]), "+r"(Dxi[5]), "+r"(Dxi[6]), "+r"(Dxi[7]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Bxi[1])); + asm("mma.sync.aligned.m16n8k8.row.col.f32.f16.f16.f32 {%0, %1, %2, %3}, {%4, %5}, {%6}, {%0, %1, %2, %3};" + : "+r"(Dxi[4]), "+r"(Dxi[5]), "+r"(Dxi[6]), "+r"(Dxi[7]) + : "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[3])); +#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_AMPERE +#elif defined(AMD_WMMA_AVAILABLE) +#if defined(RDNA4) + using halfx8_t = __attribute__((ext_vector_type(8))) _Float16; + using floatx8_t = __attribute__((ext_vector_type(8))) float; + floatx8_t& acc_frag = reinterpret_cast(D.x[0]); + const halfx8_t& a_frag = reinterpret_cast(A.x[0]); + const halfx8_t& b_frag = reinterpret_cast(B.x[0]); + acc_frag = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32_gfx12(a_frag, b_frag, acc_frag); +#elif defined(RDNA3) + using halfx16_t = __attribute__((ext_vector_type(16))) _Float16; + using floatx8_t = __attribute__((ext_vector_type(8))) float; + floatx8_t& acc_frag = reinterpret_cast(D.x[0]); + const halfx16_t& a_frag = reinterpret_cast(A.x[0]); + const halfx16_t& b_frag = reinterpret_cast(B.x[0]); + acc_frag = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(a_frag, b_frag, acc_frag); +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // RDNA4 +#elif defined(AMD_MFMA_AVAILABLE) + using halfx4_t = __attribute__((ext_vector_type(4))) _Float16; + using floatx4_t = __attribute__((ext_vector_type(4))) float; + floatx4_t& acc_frag = reinterpret_cast(D.x[0]); + const halfx4_t& a_frag = reinterpret_cast(A.x[0]); + const halfx4_t& b_frag = reinterpret_cast(B.x[0]); + acc_frag = __builtin_amdgcn_mfma_f32_16x16x16f16(a_frag, b_frag, acc_frag, 0, 0, 0); +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // TURING_MMA_AVAILABLE + } + + template + static __device__ __forceinline__ void mma( + tile<16, 16, float, dl_d> & D, const tile<16, 8, nv_bfloat162, dl_ab> & A, const tile<16, 8, nv_bfloat162, dl_ab> & B) { +#if defined(AMD_WMMA_AVAILABLE) +#if defined(RDNA4) + using bf16x8_t = __attribute__((ext_vector_type(8))) __bf16; + using floatx8_t = __attribute__((ext_vector_type(8))) float; + floatx8_t& acc_frag = reinterpret_cast(D.x[0]); + const bf16x8_t& a_frag = reinterpret_cast(A.x[0]); + const bf16x8_t& b_frag = reinterpret_cast(B.x[0]); + acc_frag = __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32_gfx12(a_frag, b_frag, acc_frag); +#elif defined(RDNA3) + using bf16x16_t = __attribute__((ext_vector_type(16))) __bf16; + using floatx8_t = __attribute__((ext_vector_type(8))) float; + floatx8_t& acc_frag = reinterpret_cast(D.x[0]); + const bf16x16_t& a_frag = reinterpret_cast(A.x[0]); + const bf16x16_t& b_frag = reinterpret_cast(B.x[0]); + acc_frag = __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32(a_frag, b_frag, acc_frag); +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // defined(RDNA4) +#elif defined(AMD_MFMA_AVAILABLE) + using floatx4_t = __attribute__((ext_vector_type(4))) float; + floatx4_t& acc_frag = reinterpret_cast(D.x[0]); +#if defined(CDNA3) || defined(CDNA2) + using bf16x4_t = __attribute__((ext_vector_type(4))) __bf16; + const bf16x4_t& a_frag = reinterpret_cast(A.x[0]); + const bf16x4_t& b_frag = reinterpret_cast(B.x[0]); + acc_frag = __builtin_amdgcn_mfma_f32_16x16x16bf16_1k(a_frag, b_frag, acc_frag, 0, 0, 0); +#elif defined(CDNA1) +#pragma unroll + for (int i = 0; i < 2; ++i) { + using bf16x2_t = __attribute__((ext_vector_type(2))) __bf16; + const bf16x2_t& a_frag = reinterpret_cast(A.x[i]); + const bf16x2_t& b_frag = reinterpret_cast(B.x[i]); + acc_frag = __builtin_amdgcn_mfma_f32_16x16x8bf16(a_frag, b_frag, acc_frag, 0, 0, 0); + } +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // defined(CDNA3) || defined(CDNA2) +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // defined(AMD_WMMA_AVAILABLE) + } + + template + static __device__ __forceinline__ void mma( + tile<16, 16, int, dl_d> & D, const tile<16, 8, int, dl_ab> & A, const tile<16, 8, int, dl_ab> & B) { +#if defined(AMD_MFMA_AVAILABLE) + using int32x4_t = __attribute__((__vector_size__(4 * sizeof(int)))) int; + int32x4_t * acc = (int32x4_t *) D.x; +#if defined(CDNA3) + acc[0] = __builtin_amdgcn_mfma_i32_16x16x32_i8(((int64_t *) A.x)[0], + ((int64_t *) B.x)[0], + acc[0], + 0, 0, 0); +#elif defined(CDNA2) || defined(CDNA) + acc[0] = __builtin_amdgcn_mfma_i32_16x16x16i8(A.x[0], + B.x[0], + acc[0], + 0, 0, 0); + acc[0] = __builtin_amdgcn_mfma_i32_16x16x16i8(A.x[1], + B.x[1], + acc[0], + 0, 0, 0); +#endif // defined(CDNA3) + +#elif defined(AMD_WMMA_AVAILABLE) + + using int32x8_t = __attribute__((__vector_size__(8 * sizeof(int)))) int; + int32x8_t * acc = (int32x8_t *) D.x; + +#if defined(RDNA4) + using int32x2_t = __attribute__((__vector_size__(2 * sizeof(int)))) int; + int32x2_t * a_vec = (int32x2_t *) A.x; + int32x2_t * b_vec = (int32x2_t *) B.x; + + acc[0] = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32_gfx12( + true, + a_vec[0], + true, + b_vec[0], + acc[0], + true + ); + + acc[0] = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32_gfx12( + true, + a_vec[1], + true, + b_vec[1], + acc[0], + true + ); + +#elif defined(RDNA3) + using int32x4_t = __attribute__((__vector_size__(4 * sizeof(int)))) int; + int32x4_t * a_vec = (int32x4_t *) A.x; + int32x4_t * b_vec = (int32x4_t *) B.x; + + acc[0] = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32( + true, + a_vec[0], + true, + b_vec[0], + acc[0], + true + ); + + acc[0] = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32( + true, + a_vec[1], + true, + b_vec[1], + acc[0], + true + ); +#endif // RDNA4 + +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // AMD_MFMA_AVAILABLE + } + + static __device__ __forceinline__ void mma( + tile<32, 32, int> & D, const tile<32, 4, int> & A, const tile<32, 4, int> & B) { +#if defined(AMD_MFMA_AVAILABLE) + using int32x16_t = __attribute__((__vector_size__(16 * sizeof(int)))) int; + int32x16_t * acc = (int32x16_t *) D.x; +#if defined(CDNA3) + acc[0] = __builtin_amdgcn_mfma_i32_32x32x16_i8(((int64_t *) A.x)[0], + ((int64_t *) B.x)[0], + acc[0], + 0, 0, 0); +#elif defined(CDNA2) || defined(CDNA) + acc[0] = __builtin_amdgcn_mfma_i32_32x32x8i8(A.x[0], + B.x[0], + acc[0], + 0, 0, 0); + acc[0] = __builtin_amdgcn_mfma_i32_32x32x8i8(A.x[1], + B.x[1], + acc[0], + 0, 0, 0); +#endif // defined(CDNA3) + +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // AMD_MFMA_AVAILABLE + } + + template + static __device__ __forceinline__ void mma( + tile<32, J, T1> & D, const tile<32, K, T2> & A, const tile & B) { + tile <16, J, T1> * D16 = reinterpret_cast< tile<16, J, T1> *>(&D); + const tile<16, K, T2> * A16 = reinterpret_cast *>(&A); + mma(D16[0], A16[0], B); + mma(D16[1], A16[1], B); + } + + static __device__ __forceinline__ void mma( + tile<32, 8, float> & D, const tile<32, 4, half2> & A, const tile<8, 4, half2, DATA_LAYOUT_I_MAJOR_MIRRORED> & B) { +#if __CUDA_ARCH__ == GGML_CUDA_CC_VOLTA + const int * Axi = (const int *) A.x; + const int * Bxi = (const int *) B.x; + int * Dxi = (int *) D.x; + asm("mma.sync.aligned.m8n8k4.row.col.f32.f16.f16.f32 " + "{%0, %1, %2, %3, %4, %5, %6, %7}, {%8, %9}, {%10, %11}, {%0, %1, %2, %3, %4, %5, %6, %7};" + : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]), "+r"(Dxi[4]), "+r"(Dxi[5]), "+r"(Dxi[6]), "+r"(Dxi[7]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Bxi[0]), "r"(Bxi[1])); + asm("mma.sync.aligned.m8n8k4.row.col.f32.f16.f16.f32 " + "{%0, %1, %2, %3, %4, %5, %6, %7}, {%8, %9}, {%10, %11}, {%0, %1, %2, %3, %4, %5, %6, %7};" + : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]), "+r"(Dxi[4]), "+r"(Dxi[5]), "+r"(Dxi[6]), "+r"(Dxi[7]) + : "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[2]), "r"(Bxi[3])); +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA + } + + static __device__ __forceinline__ void mma( + tile<32, 4, half2> & D, const tile<32, 4, half2> & A, const tile<8, 4, half2, DATA_LAYOUT_J_MAJOR_MIRRORED> & B) { +#if __CUDA_ARCH__ == GGML_CUDA_CC_VOLTA + const int * Axi = (const int *) A.x; + const int * Bxi = (const int *) B.x; + int * Dxi = (int *) D.x; + asm("mma.sync.aligned.m8n8k4.row.row.f16.f16.f16.f16 " + "{%0, %1, %2, %3}, {%4, %5}, {%6, %7}, {%0, %1, %2, %3};" + : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]) + : "r"(Axi[0]), "r"(Axi[1]), "r"(Bxi[0]), "r"(Bxi[1])); + asm("mma.sync.aligned.m8n8k4.row.row.f16.f16.f16.f16 " + "{%0, %1, %2, %3}, {%4, %5}, {%6, %7}, {%0, %1, %2, %3};" + : "+r"(Dxi[0]), "+r"(Dxi[1]), "+r"(Dxi[2]), "+r"(Dxi[3]) + : "r"(Axi[2]), "r"(Axi[3]), "r"(Bxi[2]), "r"(Bxi[3])); +#else + GGML_UNUSED_VARS(D, A, B); + NO_DEVICE_CODE; +#endif // __CUDA_ARCH__ >= GGML_CUDA_CC_VOLTA + } + + template + static __device__ __forceinline__ void mma( + tile<16, 16, int, dl_d> & D, const tile<16, 4, int, dl_ab> & A, const tile<16, 4, int, dl_ab> & B) { +#if defined(AMD_WMMA_AVAILABLE) + using int32x8_t = __attribute__((__vector_size__(8 * sizeof(int)))) int; + int32x8_t * acc = (int32x8_t *) D.x; +#if defined(RDNA4) + using int32x2_t = __attribute__((__vector_size__(2 * sizeof(int)))) int; + int32x2_t * a_vec = (int32x2_t *) A.x; + int32x2_t * b_vec = (int32x2_t *) B.x; + + acc[0] = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32_gfx12( + true, + a_vec[0], + true, + b_vec[0], + acc[0], + false + ); +#elif defined(RDNA3) + using int32x4_t = __attribute__((__vector_size__(4 * sizeof(int)))) int; + int32x4_t * a_vec = (int32x4_t *) A.x; + int32x4_t * b_vec = (int32x4_t *) B.x; + + acc[0] = __builtin_amdgcn_wmma_i32_16x16x16_iu8_w32( + true, + a_vec[0], + true, + b_vec[0], + acc[0], + false + ); +#endif // RDNA4 +#else + GGML_UNUSED(D); + GGML_UNUSED(A); + GGML_UNUSED(B); + NO_DEVICE_CODE; +#endif // AMD_WMMA_AVAILABLE + } +} diff --git a/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_quantize.cu b/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_quantize.cu new file mode 100644 index 00000000..5c4a032e --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_quantize.cu @@ -0,0 +1,145 @@ +// MMQ quantize kernel: f32 -> block_q8_1_mmq +// Adapted from llama.cpp quantize.cu + +#include "mmq_common.cuh" +#include "mmq_gguf.cuh" + +#define CUDA_QUANTIZE_BLOCK_SIZE_MMQ 128 + +static_assert(MATRIX_ROW_PADDING % (4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ) == 0, "Risk of out-of-bounds access."); + +template +static __global__ void quantize_mmq_q8_1( + const float * __restrict__ x, const int32_t * __restrict__ ids, void * __restrict__ vy, + const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03, + const int64_t ne0, const int ne1, const int ne2) { + + constexpr int vals_per_scale = ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 64 : 32; + constexpr int vals_per_sum = ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6 ? 16 : 32; + + const int64_t i0 = ((int64_t)blockDim.x*blockIdx.y + threadIdx.x)*4; + + if (i0 >= ne0) { + return; + } + + const int64_t i1 = blockIdx.x; + const int64_t i2 = blockIdx.z % ne2; + const int64_t i3 = blockIdx.z / ne2; + + const int64_t i00 = i0; + const int64_t i01 = ids ? ids[i1] : i1; + const int64_t i02 = i2; + const int64_t i03 = i3; + + const float4 * x4 = (const float4 *) x; + + block_q8_1_mmq * y = (block_q8_1_mmq *) vy; + + const int64_t ib0 = blockIdx.z*((int64_t)gridDim.x*gridDim.y*blockDim.x/QK8_1); // first block of channel + const int64_t ib = ib0 + (i0 / (4*QK8_1))*ne1 + blockIdx.x; // block index in channel + const int64_t iqs = i0 % (4*QK8_1); // quant index in block + + // Load 4 floats per thread and calculate max. abs. value between them: + const float4 xi = i0 < ne00 ? x4[(i03*s03 + i02*s02 + i01*s01 + i00)/4] : make_float4(0.0f, 0.0f, 0.0f, 0.0f); + float amax = fabsf(xi.x); + amax = fmaxf(amax, fabsf(xi.y)); + amax = fmaxf(amax, fabsf(xi.z)); + amax = fmaxf(amax, fabsf(xi.w)); + + // Exchange max. abs. value between vals_per_scale/4 threads. +#pragma unroll + for (int offset = vals_per_scale/8; offset > 0; offset >>= 1) { + amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFF, amax, offset, WARP_SIZE)); + } + + float sum; + if (ds_layout != MMQ_Q8_1_DS_LAYOUT_D4) { + sum = xi.x + xi.y + xi.z + xi.w; + + // Calculate sums across vals_per_sum/4 threads. +#pragma unroll + for (int offset = vals_per_sum/8; offset > 0; offset >>= 1) { + sum += __shfl_xor_sync(0xFFFFFFFF, sum, offset, WARP_SIZE); + } + } + + const float d_inv = 127.0f / amax; + char4 q; + q.x = roundf(xi.x*d_inv); + q.y = roundf(xi.y*d_inv); + q.z = roundf(xi.z*d_inv); + q.w = roundf(xi.w*d_inv); + + // Write back 4 int8 values as a single 32 bit value for better memory bandwidth: + char4 * yqs4 = (char4 *) y[ib].qs; + yqs4[iqs/4] = q; + + if (ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6) { + if (iqs % 16 != 0 || iqs >= 96) { + return; + } + + y[ib].d2s6[2 + iqs/16] = sum; + + if (iqs % 64 != 0) { + return; + } + + const float d = 1.0f / d_inv; + + y[ib].d2s6[iqs/64] = d; + + return; + } + + if (iqs % 32 != 0) { + return; + } + + const float d = 1.0f / d_inv; + + if (ds_layout == MMQ_Q8_1_DS_LAYOUT_DS4) { + y[ib].ds4[iqs/32] = make_half2(d, sum); + } else { + y[ib].d4[iqs/32] = d; + } +} + +// C-linkage quantize launchers + +extern "C" void launch_mmq_quantize_q8_1_D4( + const void *x_f32, const int32_t *ids, void *vy, int type_x, + int64_t ne00, int64_t s01, int64_t s02, int64_t s03, + int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, void *stream) { + const int64_t block_num_y = (ne0 + 4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ - 1) / (4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ); + const dim3 num_blocks(ne1, block_num_y, ne2*ne3); + const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1); + quantize_mmq_q8_1 + <<>>( + (const float *)x_f32, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2); +} + +extern "C" void launch_mmq_quantize_q8_1_DS4( + const void *x_f32, const int32_t *ids, void *vy, int type_x, + int64_t ne00, int64_t s01, int64_t s02, int64_t s03, + int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, void *stream) { + const int64_t block_num_y = (ne0 + 4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ - 1) / (4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ); + const dim3 num_blocks(ne1, block_num_y, ne2*ne3); + const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1); + quantize_mmq_q8_1 + <<>>( + (const float *)x_f32, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2); +} + +extern "C" void launch_mmq_quantize_q8_1_D2S6( + const void *x_f32, const int32_t *ids, void *vy, int type_x, + int64_t ne00, int64_t s01, int64_t s02, int64_t s03, + int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, void *stream) { + const int64_t block_num_y = (ne0 + 4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ - 1) / (4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ); + const dim3 num_blocks(ne1, block_num_y, ne2*ne3); + const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1); + quantize_mmq_q8_1 + <<>>( + (const float *)x_f32, ids, vy, ne00, s01, s02, s03, ne0, ne1, ne2); +} diff --git a/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_vecdotq.cuh b/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_vecdotq.cuh new file mode 100644 index 00000000..b3fdf4b4 --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/mmq_gguf/mmq_vecdotq.cuh @@ -0,0 +1,1269 @@ +#pragma once + +#include "mmq_common.cuh" + +#include + +static __device__ __forceinline__ int get_int_b1(const void * x, const int & i32) { + 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; +} + +static __device__ __forceinline__ int get_int_b2(const void * x, const int & i32) { + const uint16_t * x16 = (const uint16_t *) x; // assume at least 2 byte alignment + + int x32 = x16[2*i32 + 0] << 0; + x32 |= x16[2*i32 + 1] << 16; + + return x32; +} + +static __device__ __forceinline__ int get_int_b4(const void * x, const int & i32) { + return ((const int *) x)[i32]; // assume at least 4 byte alignment +} + +// q4 contains 8 indices with 4 bit each. +// This function selects those bytes from table that are at those indices and returns them as int2. +// The first int contains the bytes with even indices in q4, the second int contains the bytes with odd indices in q4. +static __device__ __forceinline__ int2 get_int_from_table_16(const int & q4, const int8_t * table) { +#if defined(GGML_USE_HIP) + // Load the 16-byte table into four 32-bit unsigned integers. + const uint32_t *values = (const uint32_t *)table; + + const uint32_t q_even = q4; + const uint32_t q_odd = (q4 >> 4); + + // Perform lookups in the lower half of the table (indices 0-7). + uint32_t v_even_low = __builtin_amdgcn_perm(values[1], values[0], q_even & 0x07070707); + uint32_t v_odd_low = __builtin_amdgcn_perm(values[1], values[0], q_odd & 0x07070707); + + // Perform lookups in the upper half of the table (indices 8-15). + uint32_t v_even_high = __builtin_amdgcn_perm(values[3], values[2], q_even & 0x07070707); + uint32_t v_odd_high = __builtin_amdgcn_perm(values[3], values[2], q_odd & 0x07070707); + + // Select between the low and high results based on the MSB of each index nibble. + uint32_t mask_even = 0x03020100 | ((q_even & 0x08080808) >> 1); + uint32_t res_x = __builtin_amdgcn_perm(v_even_high, v_even_low, mask_even); + uint32_t mask_odd = 0x03020100 | ((q_odd & 0x08080808) >> 1); + uint32_t res_y = __builtin_amdgcn_perm(v_odd_high, v_odd_low, mask_odd); + + return make_int2(res_x, res_y); +#elif !defined(GGML_USE_MUSA) + // CUDA does not have an instruction for selecting bytes with 4 bit indices. + // However, __byte_perm is an instruction that selects bytes with 3 bit indices that can be used instead. + const uint32_t * table32 = (const uint32_t *) table; + + // __byte_perm selects bytes based on the lower 16 bits in its third argument. + // Therefore, do 2 iterations over the 32 bits in q4 with 0 and 16 shift. + // To handle the fourth bit, first call _byte_perm both for the low and the high 64 bit of table, using the low 3 bits. + // Then, call __byte_perm again to select from the low and high bytes based on the fourth bit. + uint32_t tmp[2]; + const uint32_t low_high_selection_indices = (0x32103210 | ((q4 & 0x88888888) >> 1)); +#pragma unroll + for (uint32_t i = 0; i < 2; ++i) { + const uint32_t shift = 16 * i; + + const uint32_t low = __byte_perm(table32[0], table32[1], q4 >> shift); + const uint32_t high = __byte_perm(table32[2], table32[3], q4 >> shift); + tmp[i] = __byte_perm(low, high, low_high_selection_indices >> shift); + } + + // tmp contains the bytes from tyble in the same order as the 4 bit indices in q4. + // However, for the result we need ints with all even/odd 4 bit indices in q4. + // Therefore, 2 more calls to __byte_perm to put the bytes in the correct order. + return make_int2(__byte_perm(tmp[0], tmp[1], 0x6420), __byte_perm(tmp[0], tmp[1], 0x7531)); +#else + // Generic implementation. + const int q0_32 = (q4 >> 0) & 0x0F0F0F0F; + const int8_t * q0_8 = (const int8_t *) &q0_32; + const char4 val0_8 = make_char4( + table[q0_8[0]], table[q0_8[1]], table[q0_8[2]], table[q0_8[3]]); + + const int q1_32 = (q4 >> 4) & 0x0F0F0F0F; + const int8_t * q1_8 = (const int8_t *) &q1_32; + const char4 val1_8 = make_char4( + table[q1_8[0]], table[q1_8[1]], table[q1_8[2]], table[q1_8[3]]); + + return make_int2(*((const int *) &val0_8), *((const int *) &val1_8)); +#endif +} + +static __device__ __forceinline__ uint32_t unpack_ksigns(const uint8_t v) { + // v is a 7 bit int, with the 8th sign being encodable as popcnt + // with xor we can "correct" the bit instead of having to mask + const uint32_t p = __popc(v) & 1; + const uint32_t s = v ^ p << 7; + // broadcast over uint to allow for 0x08040201 / 0x80402010 as selectors + return s * 0x01010101; +} + +// VDR = vec dot ratio, how many contiguous integers each thread processes when the vec dot kernel is called +// MMVQ = mul_mat_vec_q, MMQ = mul_mat_q + +#define VDR_Q4_0_Q8_1_MMVQ 2 +#define VDR_Q4_0_Q8_1_MMQ 4 + +template static __device__ __forceinline__ float vec_dot_q4_0_q8_1_impl( + const int * v, const int * u, const float & d4, const half2 & ds8) { + + int sumi = 0; + +#pragma unroll + for (int i = 0; i < vdr; ++i) { + const int vi0 = (v[i] >> 0) & 0x0F0F0F0F; + const int vi1 = (v[i] >> 4) & 0x0F0F0F0F; + + // SIMD dot product of quantized values + sumi = ggml_cuda_dp4a(vi0, u[2*i+0], sumi); + sumi = ggml_cuda_dp4a(vi1, u[2*i+1], sumi); + } + + const float2 ds8f = __half22float2(ds8); + + // second part effectively subtracts 8 from each quant value + return d4 * (sumi * ds8f.x - (8*vdr/QI4_0) * ds8f.y); +} + +#define VDR_Q4_1_Q8_1_MMVQ 2 +#define VDR_Q4_1_Q8_1_MMQ 4 + +template static __device__ __forceinline__ float vec_dot_q4_1_q8_1_impl( + const int * v, const int * u, const half2 & dm4, const half2 & ds8) { + + int sumi = 0; + +#pragma unroll + for (int i = 0; i < vdr; ++i) { + const int vi0 = (v[i] >> 0) & 0x0F0F0F0F; + const int vi1 = (v[i] >> 4) & 0x0F0F0F0F; + + // SIMD dot product of quantized values + sumi = ggml_cuda_dp4a(vi0, u[2*i+0], sumi); + sumi = ggml_cuda_dp4a(vi1, u[2*i+1], sumi); + } + +#ifdef FAST_FP16_AVAILABLE + const float2 tmp = __half22float2(__hmul2(dm4, ds8)); + const float d4d8 = tmp.x; + const float m4s8 = tmp.y; +#else + const float2 dm4f = __half22float2(dm4); + const float2 ds8f = __half22float2(ds8); + const float d4d8 = dm4f.x * ds8f.x; + const float m4s8 = dm4f.y * ds8f.y; +#endif // FAST_FP16_AVAILABLE + + // scale second part of sum by QI8_1/(vdr * QR4_1) to compensate for multiple threads adding it + return sumi * d4d8 + m4s8 / (QI8_1 / (vdr * QR4_1)); +} + +#define VDR_Q5_0_Q8_1_MMVQ 2 +#define VDR_Q5_0_Q8_1_MMQ 4 + +template static __device__ __forceinline__ float vec_dot_q5_0_q8_1_impl( + const int * vl, const int * vh, const int * u, const float & d5, const half2 & ds8) { + + int sumi = 0; + +#pragma unroll + for (int i = 0; i < vdr; ++i) { + int vi0 = (vl[i] >> 0) & 0x0F0F0F0F; // lower 4 qs bits, still need qh as 5th bits + vi0 |= (vh[i] << 4) & 0x00000010; // 0 -> 4 + vi0 |= (vh[i] << 11) & 0x00001000; // 1 -> 12 + vi0 |= (vh[i] << 18) & 0x00100000; // 2 -> 20 + vi0 |= (vh[i] << 25) & 0x10000000; // 3 -> 28 + sumi = ggml_cuda_dp4a(vi0, u[2*i+0], sumi); // SIMD dot product of quantized values + + int vi1 = (vl[i] >> 4) & 0x0F0F0F0F; // upper 4 qs bits, still need qh as 5th bits + vi1 |= (vh[i] >> 12) & 0x00000010; // 16 -> 4 + vi1 |= (vh[i] >> 5) & 0x00001000; // 17 -> 12 + vi1 |= (vh[i] << 2) & 0x00100000; // 18 -> 20 + vi1 |= (vh[i] << 9) & 0x10000000; // 19 -> 28 + sumi = ggml_cuda_dp4a(vi1, u[2*i+1], sumi); // SIMD dot product of quantized values + } + + const float2 ds8f = __half22float2(ds8); + + // second part effectively subtracts 16 from each quant value + return d5 * (sumi * ds8f.x - (16*vdr/QI5_0) * ds8f.y); +} + +#define VDR_Q5_1_Q8_1_MMVQ 2 +#define VDR_Q5_1_Q8_1_MMQ 4 + +template static __device__ __forceinline__ float vec_dot_q5_1_q8_1_impl( + const int * vl, const int * vh, const int * u, const half2 & dm5, const half2 & ds8) { + + int sumi = 0; + +#pragma unroll + for (int i = 0; i < vdr; ++i) { + int vi0 = (vl[i] >> 0) & 0x0F0F0F0F; // lower 4 qs bits, still need qh as 5th bits + vi0 |= (vh[i] << 4) & 0x00000010; // 0 -> 4 + vi0 |= (vh[i] << 11) & 0x00001000; // 1 -> 12 + vi0 |= (vh[i] << 18) & 0x00100000; // 2 -> 20 + vi0 |= (vh[i] << 25) & 0x10000000; // 3 -> 28 + sumi = ggml_cuda_dp4a(vi0, u[2*i+0], sumi); // SIMD dot product of quantized values + + int vi1 = (vl[i] >> 4) & 0x0F0F0F0F; // upper 4 qs bits, still need qh as 5th bits + vi1 |= (vh[i] >> 12) & 0x00000010; // 16 -> 4 + vi1 |= (vh[i] >> 5) & 0x00001000; // 17 -> 12 + vi1 |= (vh[i] << 2) & 0x00100000; // 18 -> 20 + vi1 |= (vh[i] << 9) & 0x10000000; // 19 -> 28 + sumi = ggml_cuda_dp4a(vi1, u[2*i+1], sumi); // SIMD dot product of quantized values + } + +#ifdef FAST_FP16_AVAILABLE + const float2 tmp = __half22float2(__hmul2(dm5, ds8)); + const float d5d8 = tmp.x; + const float m5s8 = tmp.y; +#else + const float2 dm5f = __half22float2(dm5); + const float2 ds8f = __half22float2(ds8); + const float d5d8 = dm5f.x * ds8f.x; + const float m5s8 = dm5f.y * ds8f.y; +#endif // FAST_FP16_AVAILABLE + + // scale second part of sum by QI5_1 / vdr to compensate for multiple threads adding it + return sumi*d5d8 + m5s8 / (QI5_1 / vdr); +} + +#define VDR_Q8_0_Q8_1_MMVQ 2 +#define VDR_Q8_0_Q8_1_MMQ 8 + +template static __device__ __forceinline__ T vec_dot_q8_0_q8_1_impl( + const int * v, const int * u, const T & d8_0, const T & d8_1) { + + int sumi = 0; + +#pragma unroll + for (int i = 0; i < vdr; ++i) { + // SIMD dot product of quantized values + sumi = ggml_cuda_dp4a(v[i], u[i], sumi); + } + + return d8_0*d8_1 * ((T) sumi); +} + +template static __device__ __forceinline__ float vec_dot_q8_1_q8_1_impl( + const int * v, const int * u, const half2 & dm8, const half2 & ds8) { + + int sumi = 0; + +#pragma unroll + for (int i = 0; i < vdr; ++i) { + // SIMD dot product of quantized values + sumi = ggml_cuda_dp4a(v[i], u[i], sumi); + } + +#ifdef FAST_FP16_AVAILABLE + const float2 tmp = __half22float2(__hmul2(dm8, ds8)); + const float d8d8 = tmp.x; + const float m8s8 = tmp.y; +#else + const float2 dm8f = __half22float2(dm8); + const float2 ds8f = __half22float2(ds8); + const float d8d8 = dm8f.x * ds8f.x; + const float m8s8 = dm8f.y * ds8f.y; +#endif // FAST_FP16_AVAILABLE + + // scale second part of sum by QI8_1/ vdr to compensate for multiple threads adding it + return sumi*d8d8 + m8s8 / (QI8_1 / vdr); +} + +template static __device__ __forceinline__ float vec_dot_q8_0_16_q8_1_impl( + const int * v, const int * u, const float * d8_0, const float & d8_1) { + + float sumf = 0.0f; + +#pragma unroll + for (int i0 = 0; i0 < vdr; i0 += QI8_0/2) { + int sumi = 0; + +#pragma unroll + for (int i = i0; i < i0 + QI8_0/2; ++i) { + // SIMD dot product of quantized values + sumi = ggml_cuda_dp4a(v[i], u[i], sumi); + } + + sumf += d8_0[i0/(QI8_0/2)]*sumi; + } + + return d8_1*sumf; +} + +#define VDR_MXFP4_Q8_1_MMVQ 2 +#define VDR_MXFP4_Q8_1_MMQ 4 + +static __device__ __forceinline__ float vec_dot_mxfp4_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_mxfp4 * bq4 = (const block_mxfp4 *) vbq + kbx; + + const int * q8 = (const int *) bq8_1->qs + iqs; + + int sumi = 0; +#pragma unroll + for (int l = 0; l < VDR_MXFP4_Q8_1_MMVQ; ++l) { + const int aux_q4 = get_int_b1(bq4->qs, iqs + l); + const int2 v = get_int_from_table_16(aux_q4, kvalues_mxfp4); + + sumi = ggml_cuda_dp4a(v.x, q8[l + 0], sumi); + sumi = ggml_cuda_dp4a(v.y, q8[l + 4], sumi); + } + + const float d = ggml_cuda_e8m0_to_fp32(bq4->e) * 0.5f * __low2float(bq8_1->ds); + return d * sumi; +} + +#define VDR_NVFP4_Q8_1_MMVQ 4 +#define VDR_NVFP4_Q8_1_MMQ 8 + +static __device__ __forceinline__ float vec_dot_nvfp4_q8_1( + const void * __restrict__ vbq, + const block_q8_1 * __restrict__ bq8_1, + const int32_t & kbx, + const int32_t & iqs) { + + const block_nvfp4 * bq4 = (const block_nvfp4 *) vbq + kbx; + float sum = 0.0f; +#pragma unroll + for (int i = 0; i < VDR_NVFP4_Q8_1_MMVQ/2; i++) { + const int32_t iqs0 = iqs + 2*i; + const int32_t iqs1 = iqs0 + 1; + const int32_t is = iqs0 >> 1; + const int2 v0 = get_int_from_table_16(get_int_b4(bq4->qs, iqs0), kvalues_mxfp4); + const int2 v1 = get_int_from_table_16(get_int_b4(bq4->qs, iqs1), kvalues_mxfp4); + const block_q8_1 * bq8 = bq8_1 + (is >> 1); + const int32_t i8 = ((is & 1) << 2); + + int sumi = ggml_cuda_dp4a(v0.x, get_int_b4(bq8->qs, i8 + 0), 0); + sumi = ggml_cuda_dp4a(v0.y, get_int_b4(bq8->qs, i8 + 2), sumi); + sumi = ggml_cuda_dp4a(v1.x, get_int_b4(bq8->qs, i8 + 1), sumi); + sumi = ggml_cuda_dp4a(v1.y, get_int_b4(bq8->qs, i8 + 3), sumi); + + const float d = ggml_cuda_ue4m3_to_fp32(bq4->d[is]) * __low2float(bq8->ds); + sum += d * float(sumi); + } + + return sum; +} +#define VDR_Q2_K_Q8_1_MMVQ 1 +#define VDR_Q2_K_Q8_1_MMQ 4 + +// contiguous v/x values +static __device__ __forceinline__ float vec_dot_q2_K_q8_1_impl_mmvq( + const int & v, const int * __restrict__ u, const uint8_t * __restrict__ scales, + const half2 & dm2, const float * __restrict__ d8) { + + float sumf_d = 0.0f; + float sumf_m = 0.0f; + +#pragma unroll + for (int i = 0; i < QR2_K; ++i) { + const int sc = scales[2*i]; + + const int vi = (v >> (2*i)) & 0x03030303; + + sumf_d += d8[i] * (ggml_cuda_dp4a(vi, u[i], 0) * (sc & 0xF)); // SIMD dot product + + // fill int with 4x m + int m = sc >> 4; + m |= m << 8; + m |= m << 16; + sumf_m += d8[i] * ggml_cuda_dp4a(m, u[i], 0); // multiply constant q2_K part with sum of q8_1 values + } + + const float2 dm2f = __half22float2(dm2); + + return dm2f.x*sumf_d - dm2f.y*sumf_m; +} + +// contiguous v/x + u/y values +template +static __device__ __forceinline__ float vec_dot_q2_K_q8_1_impl_mmq( + const int * __restrict__ v, const int * __restrict__ u, const half2 * dm2, const float & d8, const half2 * s8) { + + float sumf = 0.0f; + float sumf_d8 = 0.0f; + +#pragma unroll + for (int i0 = 0; i0 < QR2_K*VDR_Q2_K_Q8_1_MMQ; i0 += QI8_1) { + const float2 dm2f0 = __half22float2(dm2[i0/(QI8_1/2) + 0]); + int sumi_d0 = 0; + + const float2 dm2f1 = __half22float2(dm2[i0/(QI8_1/2) + 1]); + int sumi_d1 = 0; + +#pragma unroll + for (int i = i0; i < i0 + QI8_1/2; ++i) { + sumi_d0 = ggml_cuda_dp4a(v[i], u[i], sumi_d0); + } + sumf_d8 += dm2f0.x * sumi_d0; + +#pragma unroll + for (int i = i0 + QI8_1/2; i < i0 + QI8_1; ++i) { + sumi_d1 = ggml_cuda_dp4a(v[i], u[i], sumi_d1); + } + sumf_d8 += dm2f1.x * sumi_d1; + + if (i0/QI8_1 < ns8) { + const float2 s8f = __half22float2(s8[i0/QI8_1]); + sumf -= dm2f0.y*s8f.x; + sumf -= dm2f1.y*s8f.y; + } else { + int sumi_m0 = 0; +#pragma unroll + for (int i = i0; i < i0 + QI8_1/2; ++i) { + sumi_m0 = ggml_cuda_dp4a(0x01010101, u[i], sumi_m0); + } + sumf_d8 -= dm2f0.y * sumi_m0; + + int sumi_m1 = 0; +#pragma unroll + for (int i = i0 + QI8_1/2; i < i0 + QI8_1; ++i) { + sumi_m1 = ggml_cuda_dp4a(0x01010101, u[i], sumi_m1); + } + sumf_d8 -= dm2f1.y * sumi_m1; + } + } + + return sumf + d8*sumf_d8; +} + +#define VDR_Q3_K_Q8_1_MMVQ 1 +#define VDR_Q3_K_Q8_1_MMQ 2 + +// contiguous v/x values +static __device__ __forceinline__ float vec_dot_q3_K_q8_1_impl_mmvq( + const int & vl, const int & vh, const int * __restrict__ u, const uint8_t * __restrict__ scales, + const int & scale_offset, const float & d3, const float * __restrict__ d8) { + + float sumf = 0.0f; + +#pragma unroll + for (int i = 0; i < QR3_K; ++i) { + const int isc = scale_offset + 2*i; + + const int isc_low = isc % (QK_K/32); + const int sc_shift_low = 4 * (isc / (QK_K/32)); + const int sc_low = (scales[isc_low] >> sc_shift_low) & 0xF; + + const int isc_high = isc % (QK_K/64); + const int sc_shift_high = 2 * (isc / (QK_K/64)); + const int sc_high = ((scales[(QK_K/32) + isc_high] >> sc_shift_high) & 3) << 4; + + const int sc = (sc_low | sc_high) - 32; + + const int vil = (vl >> (2*i)) & 0x03030303; + + const int vih = ((vh >> i) << 2) & 0x04040404; + + const int vi = __vsubss4(vil, vih); + + sumf += d8[i] * (ggml_cuda_dp4a(vi, u[i], 0) * sc); // SIMD dot product + } + + return d3 * sumf; +} + +// contiguous v/x + u/y values +static __device__ __forceinline__ float vec_dot_q3_K_q8_1_impl_mmq( + const int * __restrict__ v, const int * __restrict__ u, const int8_t * __restrict__ scales, + const float & d3, const float & d8) { + + int sumi = 0; + +#pragma unroll + for (int i0 = 0; i0 < QR3_K*VDR_Q3_K_Q8_1_MMQ; i0 += QI8_1/2) { + int sumi_sc = 0; + +#pragma unroll + for (int i = i0; i < i0 + QI8_1/2; ++i) { + sumi_sc = ggml_cuda_dp4a(v[i], u[i], sumi_sc); // SIMD dot product + } + + sumi += sumi_sc * scales[i0 / (QI8_1/2)]; + } + + return d3*d8 * sumi; +} + +#define VDR_Q4_K_Q8_1_MMVQ 2 +#define VDR_Q4_K_Q8_1_MMQ 8 + +// contiguous v/x values +static __device__ __forceinline__ float vec_dot_q4_K_q8_1_impl_vmmq( + const int * __restrict__ v, const int * __restrict__ u, const uint8_t * __restrict__ sc, + const uint8_t * __restrict__ m, const half2 & dm4, const float * __restrict__ d8) { + + float sumf_d = 0.0f; + float sumf_m = 0.0f; + +#pragma unroll + for (int i = 0; i < QR4_K; ++i) { + const int v0i = (v[0] >> (4*i)) & 0x0F0F0F0F; + const int v1i = (v[1] >> (4*i)) & 0x0F0F0F0F; + + const int dot1 = ggml_cuda_dp4a(v1i, u[2*i+1], ggml_cuda_dp4a(v0i, u[2*i+0], 0)); // SIMD dot product + const int dot2 = ggml_cuda_dp4a(0x01010101, u[2*i+1], ggml_cuda_dp4a(0x01010101, u[2*i+0], 0)); // sum of u + + sumf_d += d8[i] * (dot1 * sc[i]); + sumf_m += d8[i] * (dot2 * m[i]); // multiply constant part of q4_K with sum of q8_1 values + } + + const float2 dm4f = __half22float2(dm4); + + return dm4f.x*sumf_d - dm4f.y*sumf_m; +} + +// contiguous v/x + u/y values +static __device__ __forceinline__ float vec_dot_q4_K_q8_1_impl_mmq( + const int * __restrict__ v, const int * __restrict__ u, const uint8_t * __restrict__ sc, + const uint8_t * __restrict__ m, const half2 & dm4, const half2 * __restrict__ ds8) { + + float sumf_d = 0.0f; + float sumf_m = 0.0f; + +#pragma unroll + for (int i = 0; i < QR4_K*VDR_Q4_K_Q8_1_MMQ/QI8_1; ++i) { + int sumi_d = 0; + +#pragma unroll + for (int j = 0; j < QI8_1; ++j) { + sumi_d = ggml_cuda_dp4a((v[j] >> (4*i)) & 0x0F0F0F0F, u[i*QI8_1 + j], sumi_d); // SIMD dot product + } + + const float2 ds8f = __half22float2(ds8[i]); + + sumf_d += ds8f.x * (sc[i] * sumi_d); + sumf_m += ds8f.y * m[i]; // sum of q8_1 block * q4_K min val + } + + const float2 dm4f = __half22float2(dm4); + + return dm4f.x*sumf_d - dm4f.y*sumf_m; +} + +#define VDR_Q5_K_Q8_1_MMVQ 2 +#define VDR_Q5_K_Q8_1_MMQ 8 + +// contiguous v/x values +static __device__ __forceinline__ float vec_dot_q5_K_q8_1_impl_vmmq( + const int * __restrict__ vl, const int * __restrict__ vh, const int * __restrict__ u, const uint8_t * __restrict__ sc, + const uint8_t * __restrict__ m, const half2 & dm5, const float * __restrict__ d8) { + + float sumf_d = 0.0f; + float sumf_m = 0.0f; + +#pragma unroll + for (int i = 0; i < QR5_K; ++i) { + const int vl0i = (vl[0] >> (4*i)) & 0x0F0F0F0F; + const int vl1i = (vl[1] >> (4*i)) & 0x0F0F0F0F; + + const int vh0i = ((vh[0] >> i) << 4) & 0x10101010; + const int vh1i = ((vh[1] >> i) << 4) & 0x10101010; + + const int v0i = vl0i | vh0i; + const int v1i = vl1i | vh1i; + + const int dot1 = ggml_cuda_dp4a(v0i, u[2*i+0], ggml_cuda_dp4a(v1i, u[2*i+1], 0)); // SIMD dot product + const int dot2 = ggml_cuda_dp4a(0x01010101, u[2*i+0], ggml_cuda_dp4a(0x01010101, u[2*i+1], 0)); // sum of u + + sumf_d += d8[i] * (dot1 * sc[i]); + sumf_m += d8[i] * (dot2 * m[i]); + + } + + const float2 dm5f = __half22float2(dm5); + + return dm5f.x*sumf_d - dm5f.y*sumf_m; +} + +// contiguous v/x + u/y values +static __device__ __forceinline__ float vec_dot_q5_K_q8_1_impl_mmq( + const int * __restrict__ v, const int * __restrict__ u, const uint8_t * __restrict__ sc, + const uint8_t * __restrict__ m, const half2 & dm4, const half2 * __restrict__ ds8) { + + float sumf_d = 0.0f; + float sumf_m = 0.0f; + +#pragma unroll + for (int i = 0; i < QR5_K*VDR_Q5_K_Q8_1_MMQ/QI8_1; ++i) { + int sumi_d = 0; + +#pragma unroll + for (int j = 0; j < QI8_1; ++j) { + sumi_d = ggml_cuda_dp4a(v[i*QI8_1 + j], u[i*QI8_1 + j], sumi_d); // SIMD dot product + } + + const float2 ds8f = __half22float2(ds8[i]); + + sumf_d += ds8f.x * (sc[i] * sumi_d); + sumf_m += ds8f.y * m[i]; // sum of q8_1 block * q4_K min val + } + + const float2 dm4f = __half22float2(dm4); + + return dm4f.x*sumf_d - dm4f.y*sumf_m; +} + +#define VDR_Q6_K_Q8_1_MMVQ 1 +#define VDR_Q6_K_Q8_1_MMQ 8 + +// contiguous v/x values +static __device__ __forceinline__ float vec_dot_q6_K_q8_1_impl_mmvq( + const int & vl, const int & vh, const int * __restrict__ u, const int8_t * __restrict__ scales, + const float & d, const float * __restrict__ d8) { + + float sumf = 0.0f; + +#pragma unroll + for (int i = 0; i < QR6_K; ++i) { + const int sc = scales[4*i]; + + const int vil = (vl >> (4*i)) & 0x0F0F0F0F; + + const int vih = ((vh >> (4*i)) << 4) & 0x30303030; + + const int vi = __vsubss4((vil | vih), 0x20202020); // vi = (vil | vih) - 32 + + sumf += d8[i] * (ggml_cuda_dp4a(vi, u[i], 0) * sc); // SIMD dot product + } + + return d*sumf; +} + +// contiguous v/x + u/y values +static __device__ __forceinline__ float vec_dot_q6_K_q8_1_impl_mmq( + const int * __restrict__ v, const int * __restrict__ u, const int8_t * __restrict__ sc, + const float & d6, const float * __restrict__ d8) { + + float sumf_d = 0.0f; + + const int sc_packed = get_int_b4(sc, 0); + const int8_t * sc_reg = (const int8_t *) &sc_packed; + +#pragma unroll + for (int i0 = 0; i0 < VDR_Q6_K_Q8_1_MMQ; i0 += 4) { + int2 sumi_d = {0, 0}; // 2 q6_K scales per q8_1 scale + +#pragma unroll + for (int i = i0; i < i0 + 2; ++i) { + sumi_d.x = ggml_cuda_dp4a(v[2*i+0], u[2*i+0], sumi_d.x); // SIMD dot product + sumi_d.x = ggml_cuda_dp4a(v[2*i+1], u[2*i+1], sumi_d.x); // SIMD dot product + + sumi_d.y = ggml_cuda_dp4a(v[2*i+4], u[2*i+4], sumi_d.y); // SIMD dot product + sumi_d.y = ggml_cuda_dp4a(v[2*i+5], u[2*i+5], sumi_d.y); // SIMD dot product + } + + sumf_d += d8[i0/4] * (sc_reg[i0/2+0]*sumi_d.x + sc_reg[i0/2+1]*sumi_d.y); + } + + return d6 * sumf_d; +} + +static __device__ __forceinline__ float vec_dot_q4_0_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_q4_0 * bq4_0 = (const block_q4_0 *) vbq + kbx; + + int v[VDR_Q4_0_Q8_1_MMVQ]; + int u[2*VDR_Q4_0_Q8_1_MMVQ]; + +#pragma unroll + for (int i = 0; i < VDR_Q4_0_Q8_1_MMVQ; ++i) { + v[i] = get_int_b2(bq4_0->qs, iqs + i); + u[2*i+0] = get_int_b4(bq8_1->qs, iqs + i); + u[2*i+1] = get_int_b4(bq8_1->qs, iqs + i + QI4_0); + } + + return vec_dot_q4_0_q8_1_impl(v, u, bq4_0->d, bq8_1->ds); +} + + +static __device__ __forceinline__ float vec_dot_q4_1_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_q4_1 * bq4_1 = (const block_q4_1 *) vbq + kbx; + + int v[VDR_Q4_1_Q8_1_MMVQ]; + int u[2*VDR_Q4_1_Q8_1_MMVQ]; + +#pragma unroll + for (int i = 0; i < VDR_Q4_1_Q8_1_MMVQ; ++i) { + v[i] = get_int_b4(bq4_1->qs, iqs + i); + u[2*i+0] = get_int_b4(bq8_1->qs, iqs + i); + u[2*i+1] = get_int_b4(bq8_1->qs, iqs + i + QI4_1); + } + + return vec_dot_q4_1_q8_1_impl(v, u, bq4_1->dm, bq8_1->ds); +} + +static __device__ __forceinline__ float vec_dot_q5_0_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_q5_0 * bq5_0 = (const block_q5_0 *) vbq + kbx; + + int vl[VDR_Q5_0_Q8_1_MMVQ]; + int vh[VDR_Q5_0_Q8_1_MMVQ]; + int u[2*VDR_Q5_0_Q8_1_MMVQ]; + +#pragma unroll + for (int i = 0; i < VDR_Q5_0_Q8_1_MMVQ; ++i) { + vl[i] = get_int_b2(bq5_0->qs, iqs + i); + vh[i] = get_int_b2(bq5_0->qh, 0) >> (4 * (iqs + i)); + u[2*i+0] = get_int_b4(bq8_1->qs, iqs + i); + u[2*i+1] = get_int_b4(bq8_1->qs, iqs + i + QI5_0); + } + + return vec_dot_q5_0_q8_1_impl(vl, vh, u, bq5_0->d, bq8_1->ds); +} + +static __device__ __forceinline__ float vec_dot_q5_1_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_q5_1 * bq5_1 = (const block_q5_1 *) vbq + kbx; + + int vl[VDR_Q5_1_Q8_1_MMVQ]; + int vh[VDR_Q5_1_Q8_1_MMVQ]; + int u[2*VDR_Q5_1_Q8_1_MMVQ]; + +#pragma unroll + for (int i = 0; i < VDR_Q5_1_Q8_1_MMVQ; ++i) { + vl[i] = get_int_b4(bq5_1->qs, iqs + i); + vh[i] = get_int_b4(bq5_1->qh, 0) >> (4 * (iqs + i)); + u[2*i+0] = get_int_b4(bq8_1->qs, iqs + i); + u[2*i+1] = get_int_b4(bq8_1->qs, iqs + i + QI5_1); + } + + return vec_dot_q5_1_q8_1_impl(vl, vh, u, bq5_1->dm, bq8_1->ds); +} + +static __device__ __forceinline__ float vec_dot_q8_0_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_q8_0 * bq8_0 = (const block_q8_0 *) vbq + kbx; + + int v[VDR_Q8_0_Q8_1_MMVQ]; + int u[VDR_Q8_0_Q8_1_MMVQ]; + +#pragma unroll + for (int i = 0; i < VDR_Q8_0_Q8_1_MMVQ; ++i) { + v[i] = get_int_b2(bq8_0->qs, iqs + i); + u[i] = get_int_b4(bq8_1->qs, iqs + i); + } + + return vec_dot_q8_0_q8_1_impl(v, u, bq8_0->d, __low2half(bq8_1->ds)); +} + +static __device__ __forceinline__ float vec_dot_q2_K_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_q2_K * bq2_K = (const block_q2_K *) vbq + kbx; + + const int bq8_offset = QR2_K * (iqs / QI8_1); + const int scale_offset = iqs - iqs % QI8_1 + (iqs % QI8_1) / (QI8_1/2); + + const uint8_t * scales = bq2_K->scales + scale_offset; + + const int v = get_int_b4(bq2_K->qs, iqs); + int u[QR2_K]; + float d8[QR2_K]; + +#pragma unroll + for (int i = 0; i < QR2_K; ++ i) { + u[i] = get_int_b4(bq8_1[bq8_offset + i].qs, iqs % QI8_1); + d8[i] = __low2float(bq8_1[bq8_offset + i].ds); + } + + return vec_dot_q2_K_q8_1_impl_mmvq(v, u, scales, bq2_K->dm, d8); +} + +static __device__ __forceinline__ float vec_dot_q3_K_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_q3_K * bq3_K = (const block_q3_K *) vbq + kbx; + + const int bq8_offset = QR3_K * (iqs / (QI3_K/2)); + const int scale_offset = iqs - iqs % QI8_1 + (iqs % QI8_1) / (QI8_1/2); + + const float d = bq3_K->d; + + const int vl = get_int_b2(bq3_K->qs, iqs); + + // invert the mask with ~ so that a 0/1 results in 4/0 being subtracted + const int vh = ~get_int_b2(bq3_K->hmask, iqs % (QI3_K/2)) >> bq8_offset; + + int u[QR3_K]; + float d8[QR3_K]; + +#pragma unroll + for (int i = 0; i < QR3_K; ++i) { + u[i] = get_int_b4(bq8_1[bq8_offset + i].qs, iqs % QI8_1); + d8[i] = __low2float(bq8_1[bq8_offset + i].ds); + } + + return vec_dot_q3_K_q8_1_impl_mmvq(vl, vh, u, bq3_K->scales, scale_offset, d, d8); +} + +static __device__ __forceinline__ float vec_dot_q4_K_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_q4_K * bq4_K = (const block_q4_K *) vbq + kbx; + + int v[2]; + int u[2*QR4_K]; + float d8[QR4_K]; + + // iqs is in 0,2..30. bq8_offset = iqs/4 -> bq8_offset = 0, 2, 4, 6 + const int bq8_offset = QR4_K * ((iqs/2) / (QI8_1/2)); + + // iqs = 0....3 -> bq8_offset = 0, want q4_offset = 0, 4, 8, 12 + // iqs = 4....7 -> bq8_offset = 2, want q4_offset = 32, 36, 40, 44 + // iqs = 8...11 -> bq8_offset = 4, want q4_offset = 64, 68, 72, 76 + // iqs = 12..15 -> bq8_offset = 6, want q4_offset = 96, 100, 104, 108 + + const int * q4 = (const int *)(bq4_K->qs + 16 * bq8_offset + 4 * ((iqs/2)%4)); + v[0] = q4[0]; + v[1] = q4[4]; + + const uint16_t * scales = (const uint16_t *)bq4_K->scales; + uint16_t aux[2]; + const int j = bq8_offset/2; + if (j < 2) { + aux[0] = scales[j+0] & 0x3f3f; + aux[1] = scales[j+2] & 0x3f3f; + } else { + aux[0] = ((scales[j+2] >> 0) & 0x0f0f) | ((scales[j-2] & 0xc0c0) >> 2); + aux[1] = ((scales[j+2] >> 4) & 0x0f0f) | ((scales[j-0] & 0xc0c0) >> 2); + } + const uint8_t * sc = (const uint8_t *)aux; + const uint8_t * m = sc + 2; + + for (int i = 0; i < QR4_K; ++i) { + const block_q8_1 * bq8i = bq8_1 + bq8_offset + i; + d8[i] = __low2float(bq8i->ds); + + const int * q8 = (const int *)bq8i->qs + ((iqs/2)%4); + u[2*i+0] = q8[0]; + u[2*i+1] = q8[4]; + } + + return vec_dot_q4_K_q8_1_impl_vmmq(v, u, sc, m, bq4_K->dm, d8); +} + +static __device__ __forceinline__ float vec_dot_q5_K_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_q5_K * bq5_K = (const block_q5_K *) vbq + kbx; + + int vl[2]; + int vh[2]; + int u[2*QR5_K]; + float d8[QR5_K]; + + const int bq8_offset = QR5_K * ((iqs/2) / (QI8_1/2)); + const int * ql = (const int *)(bq5_K->qs + 16 * bq8_offset + 4 * ((iqs/2)%4)); + const int * qh = (const int *)(bq5_K->qh + 4 * ((iqs/2)%4)); + + vl[0] = ql[0]; + vl[1] = ql[4]; + + vh[0] = qh[0] >> bq8_offset; + vh[1] = qh[4] >> bq8_offset; + + const uint16_t * scales = (const uint16_t *)bq5_K->scales; + uint16_t aux[2]; + const int j = bq8_offset/2; + if (j < 2) { + aux[0] = scales[j+0] & 0x3f3f; + aux[1] = scales[j+2] & 0x3f3f; + } else { + aux[0] = ((scales[j+2] >> 0) & 0x0f0f) | ((scales[j-2] & 0xc0c0) >> 2); + aux[1] = ((scales[j+2] >> 4) & 0x0f0f) | ((scales[j-0] & 0xc0c0) >> 2); + } + const uint8_t * sc = (const uint8_t *)aux; + const uint8_t * m = sc + 2; + +#pragma unroll + for (int i = 0; i < QR5_K; ++i) { + const block_q8_1 * bq8i = bq8_1 + bq8_offset + i; + d8[i] = __low2float(bq8i->ds); + + const int * q8 = (const int *)bq8i->qs + ((iqs/2)%4); + u[2*i+0] = q8[0]; + u[2*i+1] = q8[4]; + } + + return vec_dot_q5_K_q8_1_impl_vmmq(vl, vh, u, sc, m, bq5_K->dm, d8); +} + +static __device__ __forceinline__ float vec_dot_q6_K_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_q6_K * bq6_K = (const block_q6_K *) vbq + kbx; + + const int bq8_offset = 2 * QR6_K * (iqs / (QI6_K/2)) + (iqs % (QI6_K/2)) / (QI6_K/4); + const int scale_offset = (QI6_K/4) * (iqs / (QI6_K/2)) + (iqs % (QI6_K/2)) / (QI6_K/8); + const int vh_shift = 2 * ((iqs % (QI6_K/2)) / (QI6_K/4)); + + const int vl = get_int_b2(bq6_K->ql, iqs); + const int vh = get_int_b2(bq6_K->qh, (QI6_K/4) * (iqs / (QI6_K/2)) + iqs % (QI6_K/4)) >> vh_shift; + + const int8_t * scales = bq6_K->scales + scale_offset; + + int u[QR6_K]; + float d8[QR6_K]; + +#pragma unroll + for (int i = 0; i < QR6_K; ++i) { + u[i] = get_int_b4(bq8_1[bq8_offset + 2*i].qs, iqs % QI8_1); + d8[i] = __low2float(bq8_1[bq8_offset + 2*i].ds); + } + + return vec_dot_q6_K_q8_1_impl_mmvq(vl, vh, u, scales, bq6_K->d, d8); +} + +#define VDR_IQ2_XXS_Q8_1_MMVQ 2 +#define VDR_IQ2_XXS_Q8_1_MMQ 2 + +static __device__ __forceinline__ float vec_dot_iq2_xxs_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_iq2_xxs * bq2 = (const block_iq2_xxs *) vbq + kbx; + + const int q2 = get_int_b2(bq2->qs, iqs); + const uint8_t * aux8 = (const uint8_t *) &q2; + const uint32_t aux32 = get_int_b2(bq2->qs, iqs + 1); + + int sumi = 0; +#pragma unroll + for (int k0 = 0; k0 < 8; k0 += 2) { + const uint2 grid_pos = ((const uint2*)iq2xxs_grid)[aux8[k0/2]]; + const uint32_t signs = unpack_ksigns(aux32 >> (7 * k0 / 2)); + + const int signs0 = __vcmpne4(signs & 0x08040201, 0); + const int grid0 = __vsub4(grid_pos.x ^ signs0, signs0); + const int u0 = get_int_b4(bq8_1[iqs/2].qs, k0 + 0); + sumi = ggml_cuda_dp4a(grid0, u0, sumi); + + const int signs1 = __vcmpne4(signs & 0x80402010, 0); + const int grid1 = __vsub4(grid_pos.y ^ signs1, signs1); + const int u1 = get_int_b4(bq8_1[iqs/2].qs, k0 + 1); + sumi = ggml_cuda_dp4a(grid1, u1, sumi); + } + + const int ls = aux32 >> 27 | 1; // (scale * 2 + 1) + sumi = sumi * ls / 8; // (sumi * scale + sumi / 2) / 4 + const float d = __half2float(bq2->d) * __low2float(bq8_1[iqs/2].ds); + return d * sumi; +} + +#define VDR_IQ2_XS_Q8_1_MMVQ 2 +#define VDR_IQ2_XS_Q8_1_MMQ 2 + +static __device__ __forceinline__ float vec_dot_iq2_xs_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_iq2_xs * bq2 = (const block_iq2_xs *) vbq + kbx; + + const int2 q2_packed = make_int2(get_int_b2(bq2->qs, iqs + 0), get_int_b2(bq2->qs, iqs + 1)); + const uint16_t * q2 = (const uint16_t *) &q2_packed; + const int ls0 = bq2->scales[iqs/2] & 0x0F; + const int ls1 = bq2->scales[iqs/2] >> 4; + + int sumi0 = 0; + int sumi1 = 0; +#pragma unroll + for (int l0 = 0; l0 < 8; l0 += 2) { + const uint2 grid_pos = ((const uint2*)iq2xs_grid)[q2[l0/2] & 0x1FF]; + const uint32_t signs = unpack_ksigns(q2[l0/2] >> 9); + + const int signs0 = __vcmpne4(signs & 0x08040201, 0); + const int grid_l = __vsub4(grid_pos.x ^ signs0, signs0); + const int u0 = get_int_b4(bq8_1[iqs/2].qs, l0 + 0); + + const int signs1 = __vcmpne4(signs & 0x80402010, 0); + const int grid_h = __vsub4(grid_pos.y ^ signs1, signs1); + const int u1 = get_int_b4(bq8_1[iqs/2].qs, l0 + 1); + + if (l0 < 4) { + sumi0 = ggml_cuda_dp4a(grid_l, u0, sumi0); + sumi0 = ggml_cuda_dp4a(grid_h, u1, sumi0); + } else { + sumi1 = ggml_cuda_dp4a(grid_l, u0, sumi1); + sumi1 = ggml_cuda_dp4a(grid_h, u1, sumi1); + } + } + const int sumi = (sumi0*ls0 + sumi1*ls1 + (sumi0 + sumi1)/2)/4; + const float d = __half2float(bq2->d) * __low2float(bq8_1[iqs/2].ds); + return d * sumi; +} + +#define VDR_IQ2_S_Q8_1_MMVQ 2 +#define VDR_IQ2_S_Q8_1_MMQ 2 + +static __device__ __forceinline__ float vec_dot_iq2_s_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_iq2_s * bq2 = (const block_iq2_s *) vbq + kbx; + + const int qs_packed = get_int_b2(bq2->qs, iqs/2); + const uint8_t * qs = (const uint8_t *) &qs_packed; + + const int qh = bq2->qh[iqs/2]; + + const int signs_packed_32 = get_int_b2(bq2->qs, QK_K/32 + iqs/2); + const uint8_t * signs_packed_8 = (const uint8_t *) &signs_packed_32; + + const int ls0 = bq2->scales[iqs/2] & 0x0F; + const int ls1 = bq2->scales[iqs/2] >> 4; + + int sumi0 = 0; + int sumi1 = 0; +#pragma unroll + for (int l0 = 0; l0 < 8; l0 += 2) { + const int * grid_pos = (const int *)(iq2s_grid + (qs[l0/2] | ((qh << (8-l0)) & 0x300))); + + const int signs0 = __vcmpne4(((signs_packed_8[l0/2] & 0x03) << 7) | ((signs_packed_8[l0/2] & 0x0C) << 21), 0x00000000); + const int signs1 = __vcmpne4(((signs_packed_8[l0/2] & 0x30) << 3) | ((signs_packed_8[l0/2] & 0xC0) << 17), 0x00000000); + + const int grid_l = __vsub4(grid_pos[0] ^ signs0, signs0); + const int grid_h = __vsub4(grid_pos[1] ^ signs1, signs1); + + const int u0 = get_int_b4(bq8_1[iqs/2].qs, l0 + 0); + const int u1 = get_int_b4(bq8_1[iqs/2].qs, l0 + 1); + + if (l0 < 4) { + sumi0 = ggml_cuda_dp4a(grid_l, u0, sumi0); + sumi0 = ggml_cuda_dp4a(grid_h, u1, sumi0); + } else { + sumi1 = ggml_cuda_dp4a(grid_l, u0, sumi1); + sumi1 = ggml_cuda_dp4a(grid_h, u1, sumi1); + } + } + const int sumi = (sumi0*ls0 + sumi1*ls1 + (sumi0 + sumi1)/2)/4; + + const float d = __half2float(bq2->d) * __low2float(bq8_1[iqs/2].ds); + return d * sumi; +} + +#define VDR_IQ3_XXS_Q8_1_MMVQ 2 +#define VDR_IQ3_XXS_Q8_1_MMQ 2 + +static __device__ __forceinline__ float vec_dot_iq3_xxs_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_iq3_xxs * bq3 = (const block_iq3_xxs *) vbq + kbx; + + const int2 q3_packed = make_int2(get_int_b2(bq3->qs, iqs), get_int_b2(bq3->qs, iqs+1)); + const uint8_t * q3 = (const uint8_t *) &q3_packed; + const uint32_t aux32 = get_int_b2(bq3->qs, QK_K/16 + iqs/2); + + int sumi = 0; +#pragma unroll + for (int l0 = 0; l0 < 8; l0 += 2) { + const int2 grid_pos = make_int2(iq3xxs_grid[q3[l0 + 0]], iq3xxs_grid[q3[l0 + 1]]); + const uint32_t signs = unpack_ksigns(aux32 >> (7*l0/2)); + + const int signs0 = __vcmpne4(signs & 0x08040201, 0); + const int grid_l = __vsub4(grid_pos.x ^ signs0, signs0); + + const int u0 = get_int_b4(bq8_1[iqs/2].qs, l0 + 0); + + const int signs1 = __vcmpne4(signs & 0x80402010, 0); + const int grid_h = __vsub4(grid_pos.y ^ signs1, signs1); + + const int u1 = get_int_b4(bq8_1[iqs/2].qs, l0 + 1); + + sumi = ggml_cuda_dp4a(grid_l, u0, sumi); + sumi = ggml_cuda_dp4a(grid_h, u1, sumi); + } + + const int ls = aux32 >> 28; + sumi = (ls*sumi + sumi/2)/2; + const float d = __half2float(bq3->d) * __low2float(bq8_1[iqs/2].ds); + return d * sumi; +} + +#define VDR_IQ3_S_Q8_1_MMVQ 2 +#define VDR_IQ3_S_Q8_1_MMQ 2 + +// TODO: don't use lookup table for signs +static __device__ __forceinline__ float vec_dot_iq3_s_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_iq3_s * bq3 = (const block_iq3_s *) vbq + kbx; + + const int2 qs_packed = make_int2(get_int_b2(bq3->qs, iqs + 0), get_int_b2(bq3->qs, iqs + 1)); + const uint8_t * qs = (const uint8_t *) &qs_packed; + + const int qh = bq3->qh[iqs/2]; + + const int signs_packed_32 = get_int_b2(bq3->signs, iqs/2); + const uint8_t * signs_packed_8 = (const uint8_t *) &signs_packed_32; + + int sumi = 0; +#pragma unroll + for (int l0 = 0; l0 < 8; l0 += 2) { + const int2 grid_pos = make_int2( + iq3s_grid[qs[l0 + 0] | ((qh << (8 - l0)) & 0x100)], + iq3s_grid[qs[l0 + 1] | ((qh << (7 - l0)) & 0x100)]); + + const int signs0 = __vcmpne4(((signs_packed_8[l0/2] & 0x03) << 7) | ((signs_packed_8[l0/2] & 0x0C) << 21), 0x00000000); + const int signs1 = __vcmpne4(((signs_packed_8[l0/2] & 0x30) << 3) | ((signs_packed_8[l0/2] & 0xC0) << 17), 0x00000000); + + const int grid_l = __vsub4(grid_pos.x ^ signs0, signs0); + const int grid_h = __vsub4(grid_pos.y ^ signs1, signs1); + + const int u0 = get_int_b4(bq8_1[iqs/2].qs, l0 + 0); + const int u1 = get_int_b4(bq8_1[iqs/2].qs, l0 + 1); + + sumi = ggml_cuda_dp4a(grid_l, u0, sumi); + sumi = ggml_cuda_dp4a(grid_h, u1, sumi); + } + + sumi *= 1 + 2*((bq3->scales[iqs/4] >> ((iqs << 1) & 0x04)) & 0x0F); + + const float d = __half2float(bq3->d) * __low2float(bq8_1[iqs/2].ds); + return d * sumi; +} + +#define VDR_IQ1_S_Q8_1_MMVQ 1 +#define VDR_IQ1_S_Q8_1_MMQ 1 + +static __device__ __forceinline__ float vec_dot_iq1_s_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + const block_iq1_s * bq1 = (const block_iq1_s *) vbq + kbx; + + const int qs_packed = get_int_b2(bq1->qs, iqs); + const uint8_t * qs = (const uint8_t *) &qs_packed; + + const int qh = bq1->qh[iqs]; + + int sumi = 0; +#pragma unroll + for (int l0 = 0; l0 < 8; l0 += 2) { + const int grid = iq1s_grid_gpu[qs[l0/2] | (((qh >> 3*(l0/2)) & 0x07) << 8)]; + + const int grid0 = (grid >> 0) & 0x0F0F0F0F; + const int grid1 = (grid >> 4) & 0x0F0F0F0F; + + const int u0 = get_int_b4(bq8_1[iqs].qs, l0 + 0); + const int u1 = get_int_b4(bq8_1[iqs].qs, l0 + 1); + + sumi = ggml_cuda_dp4a(grid0, u0, sumi); + sumi = ggml_cuda_dp4a(grid1, u1, sumi); + } + + const float d1q = __half2float(bq1->d) * (((qh >> 11) & 0x0E) + 1); + const float delta = -1.0f + IQ1S_DELTA - (qh & 0x8000) * (2.0f*IQ1S_DELTA/0x8000); + const float2 ds = __half22float2(bq8_1[iqs].ds); + return d1q * (ds.x*sumi + ds.y*delta); +} + +#define VDR_IQ1_M_Q8_1_MMVQ 1 +#define VDR_IQ1_M_Q8_1_MMQ 1 + +static __device__ __forceinline__ float vec_dot_iq1_m_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_iq1_m * bq1 = (const block_iq1_m *) vbq + kbx; + + const int qs_packed = get_int_b4(bq1->qs, iqs); + const uint8_t * qs = (const uint8_t *) &qs_packed; + + int sumi[2] = {0}; + float sumf[2] = {0.0f}; +#pragma unroll + for (int l0 = 0; l0 < 8; l0 += 2) { + const int qhl = bq1->qh[2*iqs + l0/4] >> (4 * ((l0/2) % 2)); + + const int grid = iq1s_grid_gpu[qs[l0/2] | ((qhl & 0x07) << 8)]; + + const int grid0 = (grid >> 0) & 0x0F0F0F0F; + const int grid1 = (grid >> 4) & 0x0F0F0F0F; + + const int u0 = get_int_b4(bq8_1[iqs].qs, l0 + 0); + const int u1 = get_int_b4(bq8_1[iqs].qs, l0 + 1); + + sumi[l0/4] = ggml_cuda_dp4a(grid0, u0, sumi[l0/4]); + sumi[l0/4] = ggml_cuda_dp4a(grid1, u1, sumi[l0/4]); + + const float delta = -1.0f + IQ1M_DELTA - (qhl & 0x08) * (2.0f*IQ1M_DELTA/0x08); + int sumy = 0; + sumy = ggml_cuda_dp4a(u0, 0x01010101, sumy); + sumy = ggml_cuda_dp4a(u1, 0x01010101, sumy); + sumf[l0/4] += delta*sumy; + } + + const uint16_t * sc = (const uint16_t *) bq1->scales; + + iq1m_scale_t scale; + scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00F0) | ((sc[2] >> 4) & 0x0F00) | (sc[3] & 0xF000); + const float d = __half2float(scale.f16) * __low2float(bq8_1[iqs].ds); + + const int tmp = sc[iqs/2] >> (6*(iqs%2)); + const int sc0 = 2*((tmp >> 0) & 0x07) + 1; + const int sc1 = 2*((tmp >> 3) & 0x07) + 1; + return d * ((sumi[0] + sumf[0]) * sc0 + (sumi[1] + sumf[1]) * sc1); +} + +#define VDR_IQ4_NL_Q8_1_MMVQ 2 +#define VDR_IQ4_NL_Q8_1_MMQ 4 + +static __device__ __forceinline__ float vec_dot_iq4_nl_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_iq4_nl * bq4 = (const block_iq4_nl *) vbq + kbx; + + const int * q8 = (const int *) bq8_1->qs + iqs; + + int sumi = 0; +#pragma unroll + for (int l = 0; l < VDR_Q4_0_Q8_1_MMVQ; ++l) { + const int aux_q4 = get_int_b2(bq4->qs, iqs + l); + const int2 v = get_int_from_table_16(aux_q4, kvalues_iq4nl); + + sumi = ggml_cuda_dp4a(v.x, q8[l + 0], sumi); + sumi = ggml_cuda_dp4a(v.y, q8[l + 4], sumi); + } + + const float d = __half2float(bq4->d) * __low2float(bq8_1->ds); + return d * sumi; +} + +#define VDR_IQ4_XS_Q8_1_MMVQ 4 +#define VDR_IQ4_XS_Q8_1_MMQ 4 + +static __device__ __forceinline__ float vec_dot_iq4_xs_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + + const block_iq4_xs * bq4 = (const block_iq4_xs *) vbq + kbx; + + int sumi = 0; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const int aux_q4 = get_int_b4(bq4->qs, iqs + j); + const int2 v = get_int_from_table_16(aux_q4, kvalues_iq4nl); + + const int u0 = get_int_b4(bq8_1[iqs/4].qs, j + 0); + const int u1 = get_int_b4(bq8_1[iqs/4].qs, j + 4); + + sumi = ggml_cuda_dp4a(v.x, u0, sumi); + sumi = ggml_cuda_dp4a(v.y, u1, sumi); + } + + const int ls = ((bq4->scales_l[iqs/8] >> (iqs & 0x04)) & 0x0F) | (((bq4->scales_h >> (iqs/2)) & 0x03) << 4); + sumi *= ls - 32; + + const float d = __half2float(bq4->d) * __low2float(bq8_1[iqs/4].ds); + return d * sumi; +} diff --git a/torchnative/rust/vendor/candle-kernels/src/mmvq_gguf.cu b/torchnative/rust/vendor/candle-kernels/src/mmvq_gguf.cu new file mode 100644 index 00000000..7b89ad61 --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/mmvq_gguf.cu @@ -0,0 +1,1049 @@ +// GGUF matvec kernels with Q8_1-quantized activations. +// Adapted from llama.cpp's CUDA mmvq path for the GGUF types used here. + +#include "cuda_bf16.h" +#include "cuda_fp16.h" +#include + +// Constants, types, and helpers shared with the indexed MoE kernels. + +#define WARP_SIZE 32 +#define CUDA_QUANTIZE_BLOCK_SIZE 256 +#define K_QUANTS_PER_ITERATION 2 +#define QK_K 256 +#define K_SCALE_SIZE 12 + +// Matches candle's MATRIX_ROW_PADDING. +#define MATRIX_ROW_PADDING 512 + +typedef uint16_t ggml_fp16_t; + +static __device__ __forceinline__ float warp_reduce_sum_f32(float x) { +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + x += __shfl_xor_sync(0xffffffff, x, mask, WARP_SIZE); + } + return x; +} + +static __device__ __forceinline__ float warp_reduce_max_f32(float x) { +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + x = fmaxf(x, __shfl_xor_sync(0xffffffff, x, mask, WARP_SIZE)); + } + return x; +} + +static __device__ __forceinline__ int get_int_from_int8(const int8_t *x8, + const int &i32) { + const uint16_t *x16 = (const uint16_t *)(x8 + sizeof(int) * i32); + int x32 = 0; + x32 |= x16[0] << 0; + x32 |= x16[1] << 16; + return x32; +} + +static __device__ __forceinline__ int get_int_from_uint8(const uint8_t *x8, + const int &i32) { + const uint16_t *x16 = (const uint16_t *)(x8 + sizeof(int) * i32); + int x32 = 0; + x32 |= x16[0] << 0; + x32 |= x16[1] << 16; + return x32; +} + +static __device__ __forceinline__ int +get_int_from_int8_aligned(const int8_t *x8, const int &i32) { + return *((const int *)(x8 + sizeof(int) * i32)); +} + +static __device__ __forceinline__ int +get_int_from_uint8_aligned(const uint8_t *x8, const int &i32) { + return *((const int *)(x8 + sizeof(int) * i32)); +} + +#define MIN_CC_DP4A 610 + +static __device__ __forceinline__ int ggml_cuda_dp4a(const int a, const int b, + int c) { +#if __CUDA_ARCH__ >= MIN_CC_DP4A + return __dp4a(a, b, c); +#else + const int8_t *a8 = (const int8_t *)&a; + const int8_t *b8 = (const int8_t *)&b; + return c + a8[0] * b8[0] + a8[1] * b8[1] + a8[2] * b8[2] + a8[3] * b8[3]; +#endif +} + +// --------------------------------------------------------------------------- +// Block structs for each supported quantization type +// --------------------------------------------------------------------------- + +#define QK8_0 32 +#define QR8_0 1 +#define QI8_0 (QK8_0 / (4 * QR8_0)) +typedef struct { + half d; + int8_t qs[QK8_0]; +} block_q8_0; + +#define QK8_1 32 +#define QR8_1 1 +#define QI8_1 (QK8_1 / (4 * QR8_1)) +typedef struct { + half2 ds; + int8_t qs[QK8_0]; +} block_q8_1; + +#define QR2_K 4 +#define QI2_K (QK_K / (4 * QR2_K)) +typedef struct { + uint8_t scales[QK_K / 16]; + uint8_t qs[QK_K / 4]; + half2 dm; +} block_q2_K; + +#define QR3_K 4 +#define QI3_K (QK_K / (4 * QR3_K)) +typedef struct { + uint8_t hmask[QK_K / 8]; + uint8_t qs[QK_K / 4]; + uint8_t scales[K_SCALE_SIZE]; + half d; +} block_q3_K; + +#define QR4_K 2 +#define QI4_K (QK_K / (4 * QR4_K)) +typedef struct { + half2 dm; + uint8_t scales[3 * QK_K / 64]; + uint8_t qs[QK_K / 2]; +} block_q4_K; + +#define QR5_K 2 +#define QI5_K (QK_K / (4 * QR5_K)) +typedef struct { + half2 dm; + uint8_t scales[K_SCALE_SIZE]; + uint8_t qh[QK_K / 8]; + uint8_t qs[QK_K / 2]; +} block_q5_K; + +#define QR6_K 2 +#define QI6_K (QK_K / (4 * QR6_K)) +typedef struct { + uint8_t ql[QK_K / 2]; + uint8_t qh[QK_K / 4]; + int8_t scales[QK_K / 16]; + half d; +} block_q6_K; + +#define QK4_0 32 +#define QR4_0 2 +#define QI4_0 (QK4_0 / (4 * QR4_0)) +typedef struct { + half d; + uint8_t qs[QK4_0 / 2]; +} block_q4_0; + +#define QK4_1 32 +#define QR4_1 2 +#define QI4_1 (QK4_1 / (4 * QR4_1)) +typedef struct { + half2 dm; + uint8_t qs[QK4_1 / 2]; +} block_q4_1; + +#define QK5_0 32 +#define QR5_0 2 +#define QI5_0 (QK5_0 / (4 * QR5_0)) +typedef struct { + half d; + uint8_t qh[4]; + uint8_t qs[QK5_0 / 2]; +} block_q5_0; + +#define QK5_1 32 +#define QR5_1 2 +#define QI5_1 (QK5_1 / (4 * QR5_1)) +typedef struct { + half2 dm; + uint8_t qh[4]; + uint8_t qs[QK5_1 / 2]; +} block_q5_1; + +// VDR = vec-dot unroll factor per type. +#define VDR_Q4_0_Q8_1_MMVQ 2 +#define VDR_Q4_1_Q8_1_MMVQ 2 +#define VDR_Q5_0_Q8_1_MMVQ 2 +#define VDR_Q5_1_Q8_1_MMVQ 2 +#define VDR_Q8_0_Q8_1_MMVQ 2 +#define VDR_Q8_1_Q8_1_MMVQ 2 +#define VDR_Q2_K_Q8_1_MMVQ 1 +#define VDR_Q3_K_Q8_1_MMVQ 1 +#define VDR_Q4_K_Q8_1_MMVQ 2 +#define VDR_Q5_K_Q8_1_MMVQ 2 +#define VDR_Q6_K_Q8_1_MMVQ 1 + +// --------------------------------------------------------------------------- +// vec_dot impl helpers (per-quant-type) +// --------------------------------------------------------------------------- + +template +static __device__ __forceinline__ float +vec_dot_q4_0_q8_1_impl(const int *v, const int *u, const float &d4, + const half2 &ds8) { + int sumi = 0; +#pragma unroll + for (int i = 0; i < vdr; ++i) { + const int vi0 = (v[i] >> 0) & 0x0F0F0F0F; + const int vi1 = (v[i] >> 4) & 0x0F0F0F0F; + sumi = ggml_cuda_dp4a(vi0, u[2 * i + 0], sumi); + sumi = ggml_cuda_dp4a(vi1, u[2 * i + 1], sumi); + } + const float2 ds8f = __half22float2(ds8); + return d4 * (sumi * ds8f.x - (8 * vdr / QI4_0) * ds8f.y); +} + +template +static __device__ __forceinline__ float +vec_dot_q4_1_q8_1_impl(const int *v, const int *u, const half2 &dm4, + const half2 &ds8) { + int sumi = 0; +#pragma unroll + for (int i = 0; i < vdr; ++i) { + const int vi0 = (v[i] >> 0) & 0x0F0F0F0F; + const int vi1 = (v[i] >> 4) & 0x0F0F0F0F; + sumi = ggml_cuda_dp4a(vi0, u[2 * i + 0], sumi); + sumi = ggml_cuda_dp4a(vi1, u[2 * i + 1], sumi); + } + const float2 dm4f = __half22float2(dm4); + const float2 ds8f = __half22float2(ds8); + const float d4d8 = dm4f.x * ds8f.x; + const float m4s8 = dm4f.y * ds8f.y; + return sumi * d4d8 + m4s8 / (QI8_1 / (vdr * QR4_1)); +} + +template +static __device__ __forceinline__ float +vec_dot_q5_0_q8_1_impl(const int *vl, const int *vh, const int *u, + const float &d5, const half2 &ds8) { + int sumi = 0; +#pragma unroll + for (int i = 0; i < vdr; ++i) { + int vi0 = (vl[i] >> 0) & 0x0F0F0F0F; + vi0 |= (vh[i] << 4) & 0x00000010; + vi0 |= (vh[i] << 11) & 0x00001000; + vi0 |= (vh[i] << 18) & 0x00100000; + vi0 |= (vh[i] << 25) & 0x10000000; + sumi = ggml_cuda_dp4a(vi0, u[2 * i + 0], sumi); + + int vi1 = (vl[i] >> 4) & 0x0F0F0F0F; + vi1 |= (vh[i] >> 12) & 0x00000010; + vi1 |= (vh[i] >> 5) & 0x00001000; + vi1 |= (vh[i] << 2) & 0x00100000; + vi1 |= (vh[i] << 9) & 0x10000000; + sumi = ggml_cuda_dp4a(vi1, u[2 * i + 1], sumi); + } + const float2 ds8f = __half22float2(ds8); + return d5 * (sumi * ds8f.x - (16 * vdr / QI5_0) * ds8f.y); +} + +template +static __device__ __forceinline__ float +vec_dot_q5_1_q8_1_impl(const int *vl, const int *vh, const int *u, + const half2 &dm5, const half2 &ds8) { + int sumi = 0; +#pragma unroll + for (int i = 0; i < vdr; ++i) { + int vi0 = (vl[i] >> 0) & 0x0F0F0F0F; + vi0 |= (vh[i] << 4) & 0x00000010; + vi0 |= (vh[i] << 11) & 0x00001000; + vi0 |= (vh[i] << 18) & 0x00100000; + vi0 |= (vh[i] << 25) & 0x10000000; + sumi = ggml_cuda_dp4a(vi0, u[2 * i + 0], sumi); + + int vi1 = (vl[i] >> 4) & 0x0F0F0F0F; + vi1 |= (vh[i] >> 12) & 0x00000010; + vi1 |= (vh[i] >> 5) & 0x00001000; + vi1 |= (vh[i] << 2) & 0x00100000; + vi1 |= (vh[i] << 9) & 0x10000000; + sumi = ggml_cuda_dp4a(vi1, u[2 * i + 1], sumi); + } + const float2 dm5f = __half22float2(dm5); + const float2 ds8f = __half22float2(ds8); + const float d5d8 = dm5f.x * ds8f.x; + const float m5s8 = dm5f.y * ds8f.y; + return sumi * d5d8 + m5s8 / (QI5_1 / vdr); +} + +template +static __device__ __forceinline__ float +vec_dot_q8_0_q8_1_impl(const int *v, const int *u, const half &d8_0, + const half &d8_1) { + int sumi = 0; +#pragma unroll + for (int i = 0; i < vdr; ++i) { + sumi = ggml_cuda_dp4a(v[i], u[i], sumi); + } + return sumi * __half2float(d8_0) * __half2float(d8_1); +} + +static __device__ __forceinline__ float +vec_dot_q2_K_q8_1_impl_mmvq(const int &v, const int *__restrict__ u, + const uint8_t *__restrict__ scales, + const half2 &dm2, const float *__restrict__ d8) { + float sumf_d = 0.0f; + float sumf_m = 0.0f; +#pragma unroll + for (int i = 0; i < QR2_K; ++i) { + const int sc = scales[2 * i]; + const int vi = (v >> (2 * i)) & 0x03030303; + sumf_d += d8[i] * (ggml_cuda_dp4a(vi, u[i], 0) * (sc & 0xF)); + int m = sc >> 4; + m |= m << 8; + m |= m << 16; + sumf_m += d8[i] * ggml_cuda_dp4a(m, u[i], 0); + } + const float2 dm2f = __half22float2(dm2); + return dm2f.x * sumf_d - dm2f.y * sumf_m; +} + +static __device__ __forceinline__ float vec_dot_q3_K_q8_1_impl_mmvq( + const int &vl, const int &vh, const int *__restrict__ u, + const uint8_t *__restrict__ scales, const int &scale_offset, + const float &d3, const float *__restrict__ d8) { + float sumf = 0.0f; +#pragma unroll + for (int i = 0; i < QR3_K; ++i) { + const int isc = scale_offset + 2 * i; + const int isc_low = isc % (QK_K / 32); + const int sc_shift_low = 4 * (isc / (QK_K / 32)); + const int sc_low = (scales[isc_low] >> sc_shift_low) & 0xF; + const int isc_high = isc % (QK_K / 64); + const int sc_shift_high = 2 * (isc / (QK_K / 64)); + const int sc_high = ((scales[(QK_K / 32) + isc_high] >> sc_shift_high) & 3) + << 4; + const int sc = (sc_low | sc_high) - 32; + const int vil = (vl >> (2 * i)) & 0x03030303; + const int vih = ((vh >> i) << 2) & 0x04040404; + const int vi = __vsubss4(vil, vih); + sumf += d8[i] * (ggml_cuda_dp4a(vi, u[i], 0) * sc); + } + return d3 * sumf; +} + +static __device__ __forceinline__ float vec_dot_q4_K_q8_1_impl_vmmq( + const int *__restrict__ v, const int *__restrict__ u, + const uint8_t *__restrict__ sc, const uint8_t *__restrict__ m, + const half2 &dm4, const float *__restrict__ d8) { + float sumf_d = 0.0f; + float sumf_m = 0.0f; +#pragma unroll + for (int i = 0; i < QR4_K; ++i) { + const int v0i = (v[0] >> (4 * i)) & 0x0F0F0F0F; + const int v1i = (v[1] >> (4 * i)) & 0x0F0F0F0F; + const int dot1 = + ggml_cuda_dp4a(v1i, u[2 * i + 1], ggml_cuda_dp4a(v0i, u[2 * i + 0], 0)); + const int dot2 = ggml_cuda_dp4a( + 0x01010101, u[2 * i + 1], ggml_cuda_dp4a(0x01010101, u[2 * i + 0], 0)); + sumf_d += d8[i] * (dot1 * sc[i]); + sumf_m += d8[i] * (dot2 * m[i]); + } + const float2 dm4f = __half22float2(dm4); + return dm4f.x * sumf_d - dm4f.y * sumf_m; +} + +static __device__ __forceinline__ float vec_dot_q5_K_q8_1_impl_vmmq( + const int *__restrict__ vl, const int *__restrict__ vh, + const int *__restrict__ u, const uint8_t *__restrict__ sc, + const uint8_t *__restrict__ m, const half2 &dm5, + const float *__restrict__ d8) { + float sumf_d = 0.0f; + float sumf_m = 0.0f; +#pragma unroll + for (int i = 0; i < QR5_K; ++i) { + const int vl0i = (vl[0] >> (4 * i)) & 0x0F0F0F0F; + const int vl1i = (vl[1] >> (4 * i)) & 0x0F0F0F0F; + const int vh0i = ((vh[0] >> i) << 4) & 0x10101010; + const int vh1i = ((vh[1] >> i) << 4) & 0x10101010; + const int v0i = vl0i | vh0i; + const int v1i = vl1i | vh1i; + const int dot1 = + ggml_cuda_dp4a(v0i, u[2 * i + 0], ggml_cuda_dp4a(v1i, u[2 * i + 1], 0)); + const int dot2 = ggml_cuda_dp4a( + 0x01010101, u[2 * i + 0], ggml_cuda_dp4a(0x01010101, u[2 * i + 1], 0)); + sumf_d += d8[i] * (dot1 * sc[i]); + sumf_m += d8[i] * (dot2 * m[i]); + } + const float2 dm5f = __half22float2(dm5); + return dm5f.x * sumf_d - dm5f.y * sumf_m; +} + +static __device__ __forceinline__ float +vec_dot_q6_K_q8_1_impl_mmvq(const int &vl, const int &vh, + const int *__restrict__ u, + const int8_t *__restrict__ scales, const float &d, + const float *__restrict__ d8) { + float sumf = 0.0f; +#pragma unroll + for (int i = 0; i < QR6_K; ++i) { + const int sc = scales[4 * i]; + const int vil = (vl >> (4 * i)) & 0x0F0F0F0F; + const int vih = ((vh >> (4 * i)) << 4) & 0x30303030; + const int vi = __vsubss4((vil | vih), 0x20202020); + sumf += d8[i] * (ggml_cuda_dp4a(vi, u[i], 0) * sc); + } + return d * sumf; +} + +// vec_dot wrappers for each quant type. + +typedef float (*vec_dot_q_cuda_t)(const void *__restrict__ vbq, + const block_q8_1 *__restrict__ bq8_1, + const int &kbx, const int &iqs); + +static __device__ __forceinline__ float +vec_dot_q4_0_q8_1(const void *__restrict__ vbq, + const block_q8_1 *__restrict__ bq8_1, const int &kbx, + const int &iqs) { + const block_q4_0 *bq4_0 = (const block_q4_0 *)vbq + kbx; + int v[VDR_Q4_0_Q8_1_MMVQ]; + int u[2 * VDR_Q4_0_Q8_1_MMVQ]; +#pragma unroll + for (int i = 0; i < VDR_Q4_0_Q8_1_MMVQ; ++i) { + v[i] = get_int_from_uint8(bq4_0->qs, iqs + i); + u[2 * i + 0] = get_int_from_int8_aligned(bq8_1->qs, iqs + i); + u[2 * i + 1] = get_int_from_int8_aligned(bq8_1->qs, iqs + i + QI4_0); + } + return vec_dot_q4_0_q8_1_impl(v, u, bq4_0->d, bq8_1->ds); +} + +static __device__ __forceinline__ float +vec_dot_q4_1_q8_1(const void *__restrict__ vbq, + const block_q8_1 *__restrict__ bq8_1, const int &kbx, + const int &iqs) { + const block_q4_1 *bq4_1 = (const block_q4_1 *)vbq + kbx; + int v[VDR_Q4_1_Q8_1_MMVQ]; + int u[2 * VDR_Q4_1_Q8_1_MMVQ]; +#pragma unroll + for (int i = 0; i < VDR_Q4_1_Q8_1_MMVQ; ++i) { + v[i] = get_int_from_uint8_aligned(bq4_1->qs, iqs + i); + u[2 * i + 0] = get_int_from_int8_aligned(bq8_1->qs, iqs + i); + u[2 * i + 1] = get_int_from_int8_aligned(bq8_1->qs, iqs + i + QI4_1); + } + return vec_dot_q4_1_q8_1_impl(v, u, bq4_1->dm, bq8_1->ds); +} + +static __device__ __forceinline__ float +vec_dot_q5_0_q8_1(const void *__restrict__ vbq, + const block_q8_1 *__restrict__ bq8_1, const int &kbx, + const int &iqs) { + const block_q5_0 *bq5_0 = (const block_q5_0 *)vbq + kbx; + int vl[VDR_Q5_0_Q8_1_MMVQ]; + int vh[VDR_Q5_0_Q8_1_MMVQ]; + int u[2 * VDR_Q5_0_Q8_1_MMVQ]; +#pragma unroll + for (int i = 0; i < VDR_Q5_0_Q8_1_MMVQ; ++i) { + vl[i] = get_int_from_uint8(bq5_0->qs, iqs + i); + vh[i] = get_int_from_uint8(bq5_0->qh, 0) >> (4 * (iqs + i)); + u[2 * i + 0] = get_int_from_int8_aligned(bq8_1->qs, iqs + i); + u[2 * i + 1] = get_int_from_int8_aligned(bq8_1->qs, iqs + i + QI5_0); + } + return vec_dot_q5_0_q8_1_impl(vl, vh, u, bq5_0->d, + bq8_1->ds); +} + +static __device__ __forceinline__ float +vec_dot_q5_1_q8_1(const void *__restrict__ vbq, + const block_q8_1 *__restrict__ bq8_1, const int &kbx, + const int &iqs) { + const block_q5_1 *bq5_1 = (const block_q5_1 *)vbq + kbx; + int vl[VDR_Q5_1_Q8_1_MMVQ]; + int vh[VDR_Q5_1_Q8_1_MMVQ]; + int u[2 * VDR_Q5_1_Q8_1_MMVQ]; +#pragma unroll + for (int i = 0; i < VDR_Q5_1_Q8_1_MMVQ; ++i) { + vl[i] = get_int_from_uint8_aligned(bq5_1->qs, iqs + i); + vh[i] = get_int_from_uint8_aligned(bq5_1->qh, 0) >> (4 * (iqs + i)); + u[2 * i + 0] = get_int_from_int8_aligned(bq8_1->qs, iqs + i); + u[2 * i + 1] = get_int_from_int8_aligned(bq8_1->qs, iqs + i + QI5_1); + } + return vec_dot_q5_1_q8_1_impl(vl, vh, u, bq5_1->dm, + bq8_1->ds); +} + +static __device__ __forceinline__ float +vec_dot_q8_0_q8_1(const void *__restrict__ vbq, + const block_q8_1 *__restrict__ bq8_1, const int &kbx, + const int &iqs) { + const block_q8_0 *bq8_0 = (const block_q8_0 *)vbq + kbx; + int v[VDR_Q8_0_Q8_1_MMVQ]; + int u[VDR_Q8_0_Q8_1_MMVQ]; +#pragma unroll + for (int i = 0; i < VDR_Q8_0_Q8_1_MMVQ; ++i) { + v[i] = get_int_from_int8(bq8_0->qs, iqs + i); + u[i] = get_int_from_int8_aligned(bq8_1->qs, iqs + i); + } + return vec_dot_q8_0_q8_1_impl(v, u, bq8_0->d, + __low2half(bq8_1->ds)); +} + +static __device__ __forceinline__ float +vec_dot_q2_K_q8_1(const void *__restrict__ vbq, + const block_q8_1 *__restrict__ bq8_1, const int &kbx, + const int &iqs) { + const block_q2_K *bq2_K = (const block_q2_K *)vbq + kbx; + const int bq8_offset = QR2_K * (iqs / QI8_1); + const int scale_offset = iqs - iqs % QI8_1 + (iqs % QI8_1) / (QI8_1 / 2); + const uint8_t *scales = bq2_K->scales + scale_offset; + const int v = get_int_from_uint8_aligned(bq2_K->qs, iqs); + int u[QR2_K]; + float d8[QR2_K]; +#pragma unroll + for (int i = 0; i < QR2_K; ++i) { + u[i] = get_int_from_int8_aligned(bq8_1[bq8_offset + i].qs, iqs % QI8_1); + d8[i] = __low2float(bq8_1[bq8_offset + i].ds); + } + return vec_dot_q2_K_q8_1_impl_mmvq(v, u, scales, bq2_K->dm, d8); +} + +static __device__ __forceinline__ float +vec_dot_q3_K_q8_1(const void *__restrict__ vbq, + const block_q8_1 *__restrict__ bq8_1, const int &kbx, + const int &iqs) { + const block_q3_K *bq3_K = (const block_q3_K *)vbq + kbx; + const int bq8_offset = QR3_K * (iqs / (QI3_K / 2)); + const int scale_offset = iqs - iqs % QI8_1 + (iqs % QI8_1) / (QI8_1 / 2); + const float d = bq3_K->d; + const int vl = get_int_from_uint8(bq3_K->qs, iqs); + const int vh = + ~get_int_from_uint8(bq3_K->hmask, iqs % (QI3_K / 2)) >> bq8_offset; + int u[QR3_K]; + float d8[QR3_K]; +#pragma unroll + for (int i = 0; i < QR3_K; ++i) { + u[i] = get_int_from_int8_aligned(bq8_1[bq8_offset + i].qs, iqs % QI8_1); + d8[i] = __low2float(bq8_1[bq8_offset + i].ds); + } + return vec_dot_q3_K_q8_1_impl_mmvq(vl, vh, u, bq3_K->scales, scale_offset, d, + d8); +} + +static __device__ __forceinline__ float +vec_dot_q4_K_q8_1(const void *__restrict__ vbq, + const block_q8_1 *__restrict__ bq8_1, const int &kbx, + const int &iqs) { + const block_q4_K *bq4_K = (const block_q4_K *)vbq + kbx; + int v[2]; + int u[2 * QR4_K]; + float d8[QR4_K]; + const int bq8_offset = QR4_K * ((iqs / 2) / (QI8_1 / 2)); + const int *q4 = + (const int *)(bq4_K->qs + 16 * bq8_offset + 4 * ((iqs / 2) % 4)); + v[0] = q4[0]; + v[1] = q4[4]; + const uint16_t *scales = (const uint16_t *)bq4_K->scales; + uint16_t aux[2]; + const int j = bq8_offset / 2; + if (j < 2) { + aux[0] = scales[j + 0] & 0x3f3f; + aux[1] = scales[j + 2] & 0x3f3f; + } else { + aux[0] = ((scales[j + 2] >> 0) & 0x0f0f) | ((scales[j - 2] & 0xc0c0) >> 2); + aux[1] = ((scales[j + 2] >> 4) & 0x0f0f) | ((scales[j - 0] & 0xc0c0) >> 2); + } + const uint8_t *sc = (const uint8_t *)aux; + const uint8_t *m = sc + 2; + for (int i = 0; i < QR4_K; ++i) { + const block_q8_1 *bq8i = bq8_1 + bq8_offset + i; + d8[i] = __low2float(bq8i->ds); + const int *q8 = (const int *)bq8i->qs + ((iqs / 2) % 4); + u[2 * i + 0] = q8[0]; + u[2 * i + 1] = q8[4]; + } + return vec_dot_q4_K_q8_1_impl_vmmq(v, u, sc, m, bq4_K->dm, d8); +} + +static __device__ __forceinline__ float +vec_dot_q5_K_q8_1(const void *__restrict__ vbq, + const block_q8_1 *__restrict__ bq8_1, const int &kbx, + const int &iqs) { + const block_q5_K *bq5_K = (const block_q5_K *)vbq + kbx; + int vl[2]; + int vh[2]; + int u[2 * QR5_K]; + float d8[QR5_K]; + const int bq8_offset = QR5_K * ((iqs / 2) / (QI8_1 / 2)); + const int *ql = + (const int *)(bq5_K->qs + 16 * bq8_offset + 4 * ((iqs / 2) % 4)); + const int *qh = (const int *)(bq5_K->qh + 4 * ((iqs / 2) % 4)); + vl[0] = ql[0]; + vl[1] = ql[4]; + vh[0] = qh[0] >> bq8_offset; + vh[1] = qh[4] >> bq8_offset; + const uint16_t *scales = (const uint16_t *)bq5_K->scales; + uint16_t aux[2]; + const int j = bq8_offset / 2; + if (j < 2) { + aux[0] = scales[j + 0] & 0x3f3f; + aux[1] = scales[j + 2] & 0x3f3f; + } else { + aux[0] = ((scales[j + 2] >> 0) & 0x0f0f) | ((scales[j - 2] & 0xc0c0) >> 2); + aux[1] = ((scales[j + 2] >> 4) & 0x0f0f) | ((scales[j - 0] & 0xc0c0) >> 2); + } + const uint8_t *sc = (const uint8_t *)aux; + const uint8_t *m = sc + 2; +#pragma unroll + for (int i = 0; i < QR5_K; ++i) { + const block_q8_1 *bq8i = bq8_1 + bq8_offset + i; + d8[i] = __low2float(bq8i->ds); + const int *q8 = (const int *)bq8i->qs + ((iqs / 2) % 4); + u[2 * i + 0] = q8[0]; + u[2 * i + 1] = q8[4]; + } + return vec_dot_q5_K_q8_1_impl_vmmq(vl, vh, u, sc, m, bq5_K->dm, d8); +} + +static __device__ __forceinline__ float +vec_dot_q6_K_q8_1(const void *__restrict__ vbq, + const block_q8_1 *__restrict__ bq8_1, const int &kbx, + const int &iqs) { + const block_q6_K *bq6_K = (const block_q6_K *)vbq + kbx; + const int bq8_offset = + 2 * QR6_K * (iqs / (QI6_K / 2)) + (iqs % (QI6_K / 2)) / (QI6_K / 4); + const int scale_offset = + (QI6_K / 4) * (iqs / (QI6_K / 2)) + (iqs % (QI6_K / 2)) / (QI6_K / 8); + const int vh_shift = 2 * ((iqs % (QI6_K / 2)) / (QI6_K / 4)); + const int vl = get_int_from_uint8(bq6_K->ql, iqs); + const int vh = + get_int_from_uint8(bq6_K->qh, (QI6_K / 4) * (iqs / (QI6_K / 2)) + + iqs % (QI6_K / 4)) >> + vh_shift; + const int8_t *scales = bq6_K->scales + scale_offset; + int u[QR6_K]; + float d8[QR6_K]; +#pragma unroll + for (int i = 0; i < QR6_K; ++i) { + u[i] = get_int_from_int8_aligned(bq8_1[bq8_offset + 2 * i].qs, iqs % QI8_1); + d8[i] = __low2float(bq8_1[bq8_offset + 2 * i].ds); + } + return vec_dot_q6_K_q8_1_impl_mmvq(vl, vh, u, scales, bq6_K->d, d8); +} + +// Core mat-vec-q template. + +static constexpr __device__ int mmvq_nwarps_for(int ncols_dst) { + return (ncols_dst <= 4) ? 4 : 2; +} + +static constexpr __device__ int mmvq_rows_per_cuda_block_for(int ncols_dst) { + return (ncols_dst == 1) ? 1 : 2; +} + +template +static __device__ void mmvq_core_impl( + const void *__restrict__ vx, + const block_q8_1 *__restrict__ y, + dst_t *__restrict__ dst, + const int ncols_x, const int nrows_x, + const int stride_col_y, const int stride_col_dst) { + + constexpr int nwarps = mmvq_nwarps_for(ncols_dst); + constexpr int rows_per_cuda_block = mmvq_rows_per_cuda_block_for(ncols_dst); + + const int tid = WARP_SIZE * threadIdx.y + threadIdx.x; + const int row0 = rows_per_cuda_block * blockIdx.x; + const int blocks_per_row_x = ncols_x / qk; + constexpr int blocks_per_iter = vdr * nwarps * WARP_SIZE / qi; + + // Partial sums. + float tmp[ncols_dst][rows_per_cuda_block] = {{0.0f}}; + + for (int kbx = tid / (qi / vdr); kbx < blocks_per_row_x; + kbx += blocks_per_iter) { + const int kby = kbx * (qk / QK8_1); + const int kqs = vdr * (tid % (qi / vdr)); + +#pragma unroll + for (int j = 0; j < ncols_dst; ++j) { +#pragma unroll + for (int i = 0; i < rows_per_cuda_block; ++i) { + const int row = row0 + i; + const int weight_kbx = row * blocks_per_row_x + kbx; + tmp[j][i] += + vec_dot_q_cuda(vx, &y[j * stride_col_y + kby], weight_kbx, kqs); + } + } + } + + __shared__ float tmp_shared[nwarps - 1 > 0 ? nwarps - 1 : 1][ncols_dst] + [rows_per_cuda_block][WARP_SIZE]; + + if (threadIdx.y > 0) { +#pragma unroll + for (int j = 0; j < ncols_dst; ++j) { +#pragma unroll + for (int i = 0; i < rows_per_cuda_block; ++i) { + tmp_shared[threadIdx.y - 1][j][i][threadIdx.x] = tmp[j][i]; + } + } + } + __syncthreads(); + if (threadIdx.y > 0) { + return; + } + +#pragma unroll + for (int j = 0; j < ncols_dst; ++j) { +#pragma unroll + for (int i = 0; i < rows_per_cuda_block; ++i) { +#pragma unroll + for (int l = 0; l < nwarps - 1; ++l) { + tmp[j][i] += tmp_shared[l][j][i][threadIdx.x]; + } + tmp[j][i] = warp_reduce_sum_f32(tmp[j][i]); + } + if (threadIdx.x < rows_per_cuda_block && + (rows_per_cuda_block == 1 || + uint32_t(row0 + threadIdx.x) < (uint32_t)nrows_x)) { + dst[j * stride_col_dst + row0 + threadIdx.x] = + (dst_t)tmp[j][threadIdx.x]; + } + } +} + +// --------------------------------------------------------------------------- +// Extern-C kernel entry points +// +// Macro expands `MMVQ_PLAIN_ENTRY(tag, block_q_t, qk, qi, vdr, vec_dot, +// dst_tag, dst_c_type, ncols)` into one `__global__` function for batch +// size `ncols`. The Rust launcher switches on batch size 1..=8. +// --------------------------------------------------------------------------- + +#define MMVQ_PLAIN_ENTRY(tag, block_q_t, qk_val, qi_val, vdr_val, vec_dot, \ + dst_tag, dst_c_type, ncols) \ + extern "C" __global__ void \ + mmvq_gguf_##tag##_##dst_tag##_plain_cuda##ncols( \ + const void *__restrict__ vx, const void *__restrict__ vy, \ + dst_c_type *__restrict__ dst, const int ncols_x, const int nrows_x, \ + const int stride_col_y, const int stride_col_dst) { \ + mmvq_core_impl(vx, (const block_q8_1 *)vy, dst, ncols_x, nrows_x, \ + stride_col_y, stride_col_dst); \ + } + +// -- plain entries for all 10 supported quant types, batch sizes 1..8, bf16 + f16 + f32 -- +#define MMVQ_PLAIN_BATCH_SET(tag, block_q_t, qk_val, qi_val, vdr_val, vec_dot) \ + MMVQ_PLAIN_ENTRY(tag, block_q_t, qk_val, qi_val, vdr_val, vec_dot, bf16, \ + __nv_bfloat16, 1) \ + MMVQ_PLAIN_ENTRY(tag, block_q_t, qk_val, qi_val, vdr_val, vec_dot, bf16, \ + __nv_bfloat16, 2) \ + MMVQ_PLAIN_ENTRY(tag, block_q_t, qk_val, qi_val, vdr_val, vec_dot, bf16, \ + __nv_bfloat16, 3) \ + MMVQ_PLAIN_ENTRY(tag, block_q_t, qk_val, qi_val, vdr_val, vec_dot, bf16, \ + __nv_bfloat16, 4) \ + MMVQ_PLAIN_ENTRY(tag, block_q_t, qk_val, qi_val, vdr_val, vec_dot, bf16, \ + __nv_bfloat16, 5) \ + MMVQ_PLAIN_ENTRY(tag, block_q_t, qk_val, qi_val, vdr_val, vec_dot, bf16, \ + __nv_bfloat16, 6) \ + MMVQ_PLAIN_ENTRY(tag, block_q_t, qk_val, qi_val, vdr_val, vec_dot, bf16, \ + __nv_bfloat16, 7) \ + MMVQ_PLAIN_ENTRY(tag, block_q_t, qk_val, qi_val, vdr_val, vec_dot, bf16, \ + __nv_bfloat16, 8) \ + MMVQ_PLAIN_ENTRY(tag, block_q_t, qk_val, qi_val, vdr_val, vec_dot, f16, \ + half, 1) \ + MMVQ_PLAIN_ENTRY(tag, block_q_t, qk_val, qi_val, vdr_val, vec_dot, f16, \ + half, 2) \ + MMVQ_PLAIN_ENTRY(tag, block_q_t, qk_val, qi_val, vdr_val, vec_dot, f16, \ + half, 3) \ + MMVQ_PLAIN_ENTRY(tag, block_q_t, qk_val, qi_val, vdr_val, vec_dot, f16, \ + half, 4) \ + MMVQ_PLAIN_ENTRY(tag, block_q_t, qk_val, qi_val, vdr_val, vec_dot, f16, \ + half, 5) \ + MMVQ_PLAIN_ENTRY(tag, block_q_t, qk_val, qi_val, vdr_val, vec_dot, f16, \ + half, 6) \ + MMVQ_PLAIN_ENTRY(tag, block_q_t, qk_val, qi_val, vdr_val, vec_dot, f16, \ + half, 7) \ + MMVQ_PLAIN_ENTRY(tag, block_q_t, qk_val, qi_val, vdr_val, vec_dot, f16, \ + half, 8) \ + MMVQ_PLAIN_ENTRY(tag, block_q_t, qk_val, qi_val, vdr_val, vec_dot, f32, \ + float, 1) \ + MMVQ_PLAIN_ENTRY(tag, block_q_t, qk_val, qi_val, vdr_val, vec_dot, f32, \ + float, 2) \ + MMVQ_PLAIN_ENTRY(tag, block_q_t, qk_val, qi_val, vdr_val, vec_dot, f32, \ + float, 3) \ + MMVQ_PLAIN_ENTRY(tag, block_q_t, qk_val, qi_val, vdr_val, vec_dot, f32, \ + float, 4) \ + MMVQ_PLAIN_ENTRY(tag, block_q_t, qk_val, qi_val, vdr_val, vec_dot, f32, \ + float, 5) \ + MMVQ_PLAIN_ENTRY(tag, block_q_t, qk_val, qi_val, vdr_val, vec_dot, f32, \ + float, 6) \ + MMVQ_PLAIN_ENTRY(tag, block_q_t, qk_val, qi_val, vdr_val, vec_dot, f32, \ + float, 7) \ + MMVQ_PLAIN_ENTRY(tag, block_q_t, qk_val, qi_val, vdr_val, vec_dot, f32, \ + float, 8) + +MMVQ_PLAIN_BATCH_SET(q4_0, block_q4_0, QK4_0, QI4_0, VDR_Q4_0_Q8_1_MMVQ, + vec_dot_q4_0_q8_1) +MMVQ_PLAIN_BATCH_SET(q4_1, block_q4_1, QK4_1, QI4_1, VDR_Q4_1_Q8_1_MMVQ, + vec_dot_q4_1_q8_1) +MMVQ_PLAIN_BATCH_SET(q5_0, block_q5_0, QK5_0, QI5_0, VDR_Q5_0_Q8_1_MMVQ, + vec_dot_q5_0_q8_1) +MMVQ_PLAIN_BATCH_SET(q5_1, block_q5_1, QK5_1, QI5_1, VDR_Q5_1_Q8_1_MMVQ, + vec_dot_q5_1_q8_1) +MMVQ_PLAIN_BATCH_SET(q8_0, block_q8_0, QK8_0, QI8_0, VDR_Q8_0_Q8_1_MMVQ, + vec_dot_q8_0_q8_1) +MMVQ_PLAIN_BATCH_SET(q2_k, block_q2_K, QK_K, QI2_K, VDR_Q2_K_Q8_1_MMVQ, + vec_dot_q2_K_q8_1) +MMVQ_PLAIN_BATCH_SET(q3_k, block_q3_K, QK_K, QI3_K, VDR_Q3_K_Q8_1_MMVQ, + vec_dot_q3_K_q8_1) +MMVQ_PLAIN_BATCH_SET(q4_k, block_q4_K, QK_K, QI4_K, VDR_Q4_K_Q8_1_MMVQ, + vec_dot_q4_K_q8_1) +MMVQ_PLAIN_BATCH_SET(q5_k, block_q5_K, QK_K, QI5_K, VDR_Q5_K_Q8_1_MMVQ, + vec_dot_q5_K_q8_1) +MMVQ_PLAIN_BATCH_SET(q6_k, block_q6_K, QK_K, QI6_K, VDR_Q6_K_Q8_1_MMVQ, + vec_dot_q6_K_q8_1) + +// Padding-aware BF16/F16/F32 -> Q8_1 quantize kernels. + +extern "C" __global__ void +mmvq_gguf_quantize_q8_1_bf16(const __nv_bfloat16 *__restrict__ x, + void *__restrict__ vy, const int kx, + const int kx_padded) { + const int ix = blockDim.x * blockIdx.x + threadIdx.x; + if (ix >= kx_padded) { + return; + } + const int iy = blockDim.y * blockIdx.y + threadIdx.y; + const int i_padded = iy * kx_padded + ix; + + block_q8_1 *y = (block_q8_1 *)vy; + const int ib = i_padded / QK8_1; + const int iqs = i_padded % QK8_1; + + const float xi = (ix < kx) ? __bfloat162float(x[iy * kx + ix]) : 0.0f; + float amax = fabsf(xi); + float sum = xi; + + amax = warp_reduce_max_f32(amax); + sum = warp_reduce_sum_f32(sum); + + const float d = amax / 127.0f; + const int8_t q = (amax == 0.0f) ? 0 : (int8_t)roundf(xi / d); + + y[ib].qs[iqs] = q; + + if (iqs > 0) { + return; + } + reinterpret_cast(y[ib].ds.x) = (half)d; + reinterpret_cast(y[ib].ds.y) = (half)sum; +} + +extern "C" __global__ void +mmvq_gguf_quantize_q8_1_f16(const half *__restrict__ x, + void *__restrict__ vy, const int kx, + const int kx_padded) { + const int ix = blockDim.x * blockIdx.x + threadIdx.x; + if (ix >= kx_padded) { + return; + } + const int iy = blockDim.y * blockIdx.y + threadIdx.y; + const int i_padded = iy * kx_padded + ix; + + block_q8_1 *y = (block_q8_1 *)vy; + const int ib = i_padded / QK8_1; + const int iqs = i_padded % QK8_1; + + const float xi = (ix < kx) ? __half2float(x[iy * kx + ix]) : 0.0f; + float amax = fabsf(xi); + float sum = xi; + + amax = warp_reduce_max_f32(amax); + sum = warp_reduce_sum_f32(sum); + + const float d = amax / 127.0f; + const int8_t q = (amax == 0.0f) ? 0 : (int8_t)roundf(xi / d); + + y[ib].qs[iqs] = q; + + if (iqs > 0) { + return; + } + reinterpret_cast(y[ib].ds.x) = (half)d; + reinterpret_cast(y[ib].ds.y) = (half)sum; +} + +extern "C" __global__ void +mmvq_gguf_quantize_q8_1_f32(const float *__restrict__ x, + void *__restrict__ vy, const int kx, + const int kx_padded) { + const int ix = blockDim.x * blockIdx.x + threadIdx.x; + if (ix >= kx_padded) { + return; + } + const int iy = blockDim.y * blockIdx.y + threadIdx.y; + const int i_padded = iy * kx_padded + ix; + + block_q8_1 *y = (block_q8_1 *)vy; + const int ib = i_padded / QK8_1; + const int iqs = i_padded % QK8_1; + + const float xi = (ix < kx) ? x[iy * kx + ix] : 0.0f; + float amax = fabsf(xi); + float sum = xi; + + amax = warp_reduce_max_f32(amax); + sum = warp_reduce_sum_f32(sum); + + const float d = amax / 127.0f; + const int8_t q = (amax == 0.0f) ? 0 : (int8_t)roundf(xi / d); + + y[ib].qs[iqs] = q; + + if (iqs > 0) { + return; + } + reinterpret_cast(y[ib].ds.x) = (half)d; + reinterpret_cast(y[ib].ds.y) = (half)sum; +} + +// Host-side launchers used by `candle-kernels/src/ffi.rs`. + +#define MMVQ_LAUNCHER_PLAIN(tag, dst_tag, dst_c_type) \ + extern "C" void launch_mmvq_gguf_##tag##_##dst_tag##_plain( \ + const void *vx, const void *vy, void *dst, int ncols_x, int nrows_x, \ + int stride_col_y, int stride_col_dst, int b_size, void *stream) { \ + const unsigned int rows_per_block = (b_size <= 1) ? 1 : 2; \ + const unsigned int nblocks = \ + (unsigned int)((nrows_x + rows_per_block - 1) / rows_per_block); \ + unsigned int nwarps; \ + if (b_size <= 4) { \ + nwarps = 4; \ + } else { \ + nwarps = 2; \ + } \ + dim3 grid(nblocks, 1, 1); \ + dim3 block(WARP_SIZE, nwarps, 1); \ + cudaStream_t s = static_cast(stream); \ + switch (b_size) { \ + case 1: \ + mmvq_gguf_##tag##_##dst_tag##_plain_cuda1<<>>( \ + vx, vy, (dst_c_type *)dst, ncols_x, nrows_x, stride_col_y, \ + stride_col_dst); \ + break; \ + case 2: \ + mmvq_gguf_##tag##_##dst_tag##_plain_cuda2<<>>( \ + vx, vy, (dst_c_type *)dst, ncols_x, nrows_x, stride_col_y, \ + stride_col_dst); \ + break; \ + case 3: \ + mmvq_gguf_##tag##_##dst_tag##_plain_cuda3<<>>( \ + vx, vy, (dst_c_type *)dst, ncols_x, nrows_x, stride_col_y, \ + stride_col_dst); \ + break; \ + case 4: \ + mmvq_gguf_##tag##_##dst_tag##_plain_cuda4<<>>( \ + vx, vy, (dst_c_type *)dst, ncols_x, nrows_x, stride_col_y, \ + stride_col_dst); \ + break; \ + case 5: \ + mmvq_gguf_##tag##_##dst_tag##_plain_cuda5<<>>( \ + vx, vy, (dst_c_type *)dst, ncols_x, nrows_x, stride_col_y, \ + stride_col_dst); \ + break; \ + case 6: \ + mmvq_gguf_##tag##_##dst_tag##_plain_cuda6<<>>( \ + vx, vy, (dst_c_type *)dst, ncols_x, nrows_x, stride_col_y, \ + stride_col_dst); \ + break; \ + case 7: \ + mmvq_gguf_##tag##_##dst_tag##_plain_cuda7<<>>( \ + vx, vy, (dst_c_type *)dst, ncols_x, nrows_x, stride_col_y, \ + stride_col_dst); \ + break; \ + case 8: \ + mmvq_gguf_##tag##_##dst_tag##_plain_cuda8<<>>( \ + vx, vy, (dst_c_type *)dst, ncols_x, nrows_x, stride_col_y, \ + stride_col_dst); \ + break; \ + default: \ + break; \ + } \ + } + +MMVQ_LAUNCHER_PLAIN(q4_0, bf16, __nv_bfloat16) +MMVQ_LAUNCHER_PLAIN(q4_1, bf16, __nv_bfloat16) +MMVQ_LAUNCHER_PLAIN(q5_0, bf16, __nv_bfloat16) +MMVQ_LAUNCHER_PLAIN(q5_1, bf16, __nv_bfloat16) +MMVQ_LAUNCHER_PLAIN(q8_0, bf16, __nv_bfloat16) +MMVQ_LAUNCHER_PLAIN(q2_k, bf16, __nv_bfloat16) +MMVQ_LAUNCHER_PLAIN(q3_k, bf16, __nv_bfloat16) +MMVQ_LAUNCHER_PLAIN(q4_k, bf16, __nv_bfloat16) +MMVQ_LAUNCHER_PLAIN(q5_k, bf16, __nv_bfloat16) +MMVQ_LAUNCHER_PLAIN(q6_k, bf16, __nv_bfloat16) + +MMVQ_LAUNCHER_PLAIN(q4_0, f16, half) +MMVQ_LAUNCHER_PLAIN(q4_1, f16, half) +MMVQ_LAUNCHER_PLAIN(q5_0, f16, half) +MMVQ_LAUNCHER_PLAIN(q5_1, f16, half) +MMVQ_LAUNCHER_PLAIN(q8_0, f16, half) +MMVQ_LAUNCHER_PLAIN(q2_k, f16, half) +MMVQ_LAUNCHER_PLAIN(q3_k, f16, half) +MMVQ_LAUNCHER_PLAIN(q4_k, f16, half) +MMVQ_LAUNCHER_PLAIN(q5_k, f16, half) +MMVQ_LAUNCHER_PLAIN(q6_k, f16, half) + +MMVQ_LAUNCHER_PLAIN(q4_0, f32, float) +MMVQ_LAUNCHER_PLAIN(q4_1, f32, float) +MMVQ_LAUNCHER_PLAIN(q5_0, f32, float) +MMVQ_LAUNCHER_PLAIN(q5_1, f32, float) +MMVQ_LAUNCHER_PLAIN(q8_0, f32, float) +MMVQ_LAUNCHER_PLAIN(q2_k, f32, float) +MMVQ_LAUNCHER_PLAIN(q3_k, f32, float) +MMVQ_LAUNCHER_PLAIN(q4_k, f32, float) +MMVQ_LAUNCHER_PLAIN(q5_k, f32, float) +MMVQ_LAUNCHER_PLAIN(q6_k, f32, float) + +// Quantize launchers. + +extern "C" void launch_mmvq_gguf_quantize_q8_1_bf16(const void *x, void *vy, + int kx, int kx_padded, + int num_rows, + void *stream) { + const int num_blocks_x = + (kx_padded + CUDA_QUANTIZE_BLOCK_SIZE - 1) / CUDA_QUANTIZE_BLOCK_SIZE; + dim3 grid(num_blocks_x, num_rows, 1); + dim3 block(CUDA_QUANTIZE_BLOCK_SIZE, 1, 1); + cudaStream_t s = static_cast(stream); + mmvq_gguf_quantize_q8_1_bf16<<>>( + (const __nv_bfloat16 *)x, vy, kx, kx_padded); +} + +extern "C" void launch_mmvq_gguf_quantize_q8_1_f16(const void *x, void *vy, + int kx, int kx_padded, + int num_rows, + void *stream) { + const int num_blocks_x = + (kx_padded + CUDA_QUANTIZE_BLOCK_SIZE - 1) / CUDA_QUANTIZE_BLOCK_SIZE; + dim3 grid(num_blocks_x, num_rows, 1); + dim3 block(CUDA_QUANTIZE_BLOCK_SIZE, 1, 1); + cudaStream_t s = static_cast(stream); + mmvq_gguf_quantize_q8_1_f16<<>>( + (const half *)x, vy, kx, kx_padded); +} + +extern "C" void launch_mmvq_gguf_quantize_q8_1_f32(const void *x, void *vy, + int kx, int kx_padded, + int num_rows, + void *stream) { + const int num_blocks_x = + (kx_padded + CUDA_QUANTIZE_BLOCK_SIZE - 1) / CUDA_QUANTIZE_BLOCK_SIZE; + dim3 grid(num_blocks_x, num_rows, 1); + dim3 block(CUDA_QUANTIZE_BLOCK_SIZE, 1, 1); + cudaStream_t s = static_cast(stream); + mmvq_gguf_quantize_q8_1_f32<<>>( + (const float *)x, vy, kx, kx_padded); +} diff --git a/torchnative/rust/vendor/candle-kernels/src/moe/gguf.cuh b/torchnative/rust/vendor/candle-kernels/src/moe/gguf.cuh new file mode 100644 index 00000000..7e325969 --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/moe/gguf.cuh @@ -0,0 +1,1438 @@ +// Kernels adapted from llama.cpp ggml-cuda.cu +// https://github.com/ggerganov/llama.cpp/blob/master/ggml-cuda.cu +#include "cuda_fp16.h" +#include "cuda_bf16.h" +#include + +#define GGML_UNUSED(x) (void)(x) +#define GGML_CUDA_ASSUME(x) + +#ifdef GGML_QKK_64 +#define QK_K 64 +#define K_SCALE_SIZE 4 +#else +#define QK_K 256 +#define K_SCALE_SIZE 12 +#endif + +#undef GGML_CUDA_F16 +#define GGML_CUDA_DMMV_X 32 +#define CUDA_QUANTIZE_BLOCK_SIZE 256 +#define CUDA_DEQUANTIZE_BLOCK_SIZE 256 +#define K_QUANTS_PER_ITERATION 2 + +typedef uint16_t ggml_fp16_t; +typedef float dfloat; // dequantize float +typedef float2 dfloat2; +typedef void (*dequantize_kernel_t)(const void * vx, const int ib, const int iqs, dfloat2 & v); + +static __device__ __forceinline__ float warp_reduce_sum(float x) { +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + x += __shfl_xor_sync(0xffffffff, x, mask, 32); + } + return x; +} + +static __device__ __forceinline__ float warp_reduce_max(float x) { +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + x = fmaxf(x, __shfl_xor_sync(0xffffffff, x, mask, 32)); + } + return x; +} + +static __device__ __forceinline__ int get_int_from_int8(const int8_t * x8, const int & i32) { + const uint16_t * x16 = (const uint16_t *) (x8 + sizeof(int) * i32); // assume at least 2 byte alignment + + int x32 = 0; + x32 |= x16[0] << 0; + x32 |= x16[1] << 16; + + return x32; +} + +static __device__ __forceinline__ int get_int_from_uint8(const uint8_t * x8, const int & i32) { + const uint16_t * x16 = (const uint16_t *) (x8 + sizeof(int) * i32); // assume at least 2 byte alignment + + int x32 = 0; + x32 |= x16[0] << 0; + x32 |= x16[1] << 16; + + return x32; +} + +static __device__ __forceinline__ int get_int_from_int8_aligned(const int8_t * x8, const int & i32) { + return *((const int *) (x8 + sizeof(int) * i32)); // assume at least 4 byte alignment +} + +static __device__ __forceinline__ int get_int_from_uint8_aligned(const uint8_t * x8, const int & i32) { + return *((const int *) (x8 + sizeof(int) * i32)); // assume at least 4 byte alignment +} + + +#define WARP_SIZE 32 +#define CUDART_HMAX 11070 // CUDA 11.7, min. ver. for which __hmax and __hmax2 are known to work (may be higher than needed) + +#define CUDA_CC_PASCAL 600 +#define MIN_CC_DP4A 610 // minimum compute capability for __dp4a, an intrinsic for byte-wise dot products +#define CUDA_CC_VOLTA 700 +#define CC_OFFSET_AMD 1000000 +#define CC_RDNA1 (CC_OFFSET_AMD + 1010) +#define CC_RDNA2 (CC_OFFSET_AMD + 1030) +#define CC_RDNA3 (CC_OFFSET_AMD + 1100) + +static __device__ __forceinline__ int ggml_cuda_dp4a(const int a, const int b, int c) { +#if __CUDA_ARCH__ >= MIN_CC_DP4A + return __dp4a(a, b, c); +#else // __CUDA_ARCH__ >= MIN_CC_DP4A + const int8_t * a8 = (const int8_t *) &a; + const int8_t * b8 = (const int8_t *) &b; + return c + a8[0]*b8[0] + a8[1]*b8[1] + a8[2]*b8[2] + a8[3]*b8[3]; +#endif // __CUDA_ARCH__ >= MIN_CC_DP4A +} + + +#define MMQ_X_Q4_0_RDNA2 64 +#define MMQ_Y_Q4_0_RDNA2 128 +#define NWARPS_Q4_0_RDNA2 8 +#define MMQ_X_Q4_0_RDNA1 64 +#define MMQ_Y_Q4_0_RDNA1 64 +#define NWARPS_Q4_0_RDNA1 8 +#if defined(CUDA_USE_TENSOR_CORES) +#define MMQ_X_Q4_0_AMPERE 4 +#define MMQ_Y_Q4_0_AMPERE 32 +#define NWARPS_Q4_0_AMPERE 4 +#else +#define MMQ_X_Q4_0_AMPERE 64 +#define MMQ_Y_Q4_0_AMPERE 128 +#define NWARPS_Q4_0_AMPERE 4 +#endif +#define MMQ_X_Q4_0_PASCAL 64 +#define MMQ_Y_Q4_0_PASCAL 64 +#define NWARPS_Q4_0_PASCAL 8 + +#define MMQ_X_Q4_1_RDNA2 64 +#define MMQ_Y_Q4_1_RDNA2 128 +#define NWARPS_Q4_1_RDNA2 8 +#define MMQ_X_Q4_1_RDNA1 64 +#define MMQ_Y_Q4_1_RDNA1 64 +#define NWARPS_Q4_1_RDNA1 8 +#if defined(CUDA_USE_TENSOR_CORES) +#define MMQ_X_Q4_1_AMPERE 4 +#define MMQ_Y_Q4_1_AMPERE 32 +#define NWARPS_Q4_1_AMPERE 4 +#else +#define MMQ_X_Q4_1_AMPERE 64 +#define MMQ_Y_Q4_1_AMPERE 128 +#define NWARPS_Q4_1_AMPERE 4 +#endif +#define MMQ_X_Q4_1_PASCAL 64 +#define MMQ_Y_Q4_1_PASCAL 64 +#define NWARPS_Q4_1_PASCAL 8 + +#define MMQ_X_Q5_0_RDNA2 64 +#define MMQ_Y_Q5_0_RDNA2 128 +#define NWARPS_Q5_0_RDNA2 8 +#define MMQ_X_Q5_0_RDNA1 64 +#define MMQ_Y_Q5_0_RDNA1 64 +#define NWARPS_Q5_0_RDNA1 8 +#if defined(CUDA_USE_TENSOR_CORES) +#define MMQ_X_Q5_0_AMPERE 4 +#define MMQ_Y_Q5_0_AMPERE 32 +#define NWARPS_Q5_0_AMPERE 4 +#else +#define MMQ_X_Q5_0_AMPERE 128 +#define MMQ_Y_Q5_0_AMPERE 64 +#define NWARPS_Q5_0_AMPERE 4 +#endif +#define MMQ_X_Q5_0_PASCAL 64 +#define MMQ_Y_Q5_0_PASCAL 64 +#define NWARPS_Q5_0_PASCAL 8 + +#define MMQ_X_Q5_1_RDNA2 64 +#define MMQ_Y_Q5_1_RDNA2 128 +#define NWARPS_Q5_1_RDNA2 8 +#define MMQ_X_Q5_1_RDNA1 64 +#define MMQ_Y_Q5_1_RDNA1 64 +#define NWARPS_Q5_1_RDNA1 8 +#if defined(CUDA_USE_TENSOR_CORES) +#define MMQ_X_Q5_1_AMPERE 4 +#define MMQ_Y_Q5_1_AMPERE 32 +#define NWARPS_Q5_1_AMPERE 4 +#else +#define MMQ_X_Q5_1_AMPERE 128 +#define MMQ_Y_Q5_1_AMPERE 64 +#define NWARPS_Q5_1_AMPERE 4 +#endif +#define MMQ_X_Q5_1_PASCAL 64 +#define MMQ_Y_Q5_1_PASCAL 64 +#define NWARPS_Q5_1_PASCAL 8 + +#define MMQ_X_Q8_0_RDNA2 64 +#define MMQ_Y_Q8_0_RDNA2 128 +#define NWARPS_Q8_0_RDNA2 8 +#define MMQ_X_Q8_0_RDNA1 64 +#define MMQ_Y_Q8_0_RDNA1 64 +#define NWARPS_Q8_0_RDNA1 8 +#if defined(CUDA_USE_TENSOR_CORES) +#define MMQ_X_Q8_0_AMPERE 4 +#define MMQ_Y_Q8_0_AMPERE 32 +#define NWARPS_Q8_0_AMPERE 4 +#else +#define MMQ_X_Q8_0_AMPERE 128 +#define MMQ_Y_Q8_0_AMPERE 64 +#define NWARPS_Q8_0_AMPERE 4 +#endif +#define MMQ_X_Q8_0_PASCAL 64 +#define MMQ_Y_Q8_0_PASCAL 64 +#define NWARPS_Q8_0_PASCAL 8 + +#define MMQ_X_Q2_K_RDNA2 64 +#define MMQ_Y_Q2_K_RDNA2 128 +#define NWARPS_Q2_K_RDNA2 8 +#define MMQ_X_Q2_K_RDNA1 128 +#define MMQ_Y_Q2_K_RDNA1 32 +#define NWARPS_Q2_K_RDNA1 8 +#if defined(CUDA_USE_TENSOR_CORES) +#define MMQ_X_Q2_K_AMPERE 4 +#define MMQ_Y_Q2_K_AMPERE 32 +#define NWARPS_Q2_K_AMPERE 4 +#else +#define MMQ_X_Q2_K_AMPERE 64 +#define MMQ_Y_Q2_K_AMPERE 128 +#define NWARPS_Q2_K_AMPERE 4 +#endif +#define MMQ_X_Q2_K_PASCAL 64 +#define MMQ_Y_Q2_K_PASCAL 64 +#define NWARPS_Q2_K_PASCAL 8 + +#define MMQ_X_Q3_K_RDNA2 128 +#define MMQ_Y_Q3_K_RDNA2 64 +#define NWARPS_Q3_K_RDNA2 8 +#define MMQ_X_Q3_K_RDNA1 32 +#define MMQ_Y_Q3_K_RDNA1 128 +#define NWARPS_Q3_K_RDNA1 8 +#if defined(CUDA_USE_TENSOR_CORES) +#define MMQ_X_Q3_K_AMPERE 4 +#define MMQ_Y_Q3_K_AMPERE 32 +#define NWARPS_Q3_K_AMPERE 4 +#else +#define MMQ_X_Q3_K_AMPERE 128 +#define MMQ_Y_Q3_K_AMPERE 128 +#define NWARPS_Q3_K_AMPERE 4 +#endif +#define MMQ_X_Q3_K_PASCAL 64 +#define MMQ_Y_Q3_K_PASCAL 64 +#define NWARPS_Q3_K_PASCAL 8 + +#define MMQ_X_Q4_K_RDNA2 64 +#define MMQ_Y_Q4_K_RDNA2 128 +#define NWARPS_Q4_K_RDNA2 8 +#define MMQ_X_Q4_K_RDNA1 32 +#define MMQ_Y_Q4_K_RDNA1 64 +#define NWARPS_Q4_K_RDNA1 8 +#if defined(CUDA_USE_TENSOR_CORES) +#define MMQ_X_Q4_K_AMPERE 4 +#define MMQ_Y_Q4_K_AMPERE 32 +#define NWARPS_Q4_K_AMPERE 4 +#else +#define MMQ_X_Q4_K_AMPERE 64 +#define MMQ_Y_Q4_K_AMPERE 128 +#define NWARPS_Q4_K_AMPERE 4 +#endif +#define MMQ_X_Q4_K_PASCAL 64 +#define MMQ_Y_Q4_K_PASCAL 64 +#define NWARPS_Q4_K_PASCAL 8 + +#define MMQ_X_Q5_K_RDNA2 64 +#define MMQ_Y_Q5_K_RDNA2 128 +#define NWARPS_Q5_K_RDNA2 8 +#define MMQ_X_Q5_K_RDNA1 32 +#define MMQ_Y_Q5_K_RDNA1 64 +#define NWARPS_Q5_K_RDNA1 8 +#if defined(CUDA_USE_TENSOR_CORES) +#define MMQ_X_Q5_K_AMPERE 4 +#define MMQ_Y_Q5_K_AMPERE 32 +#define NWARPS_Q5_K_AMPERE 4 +#else +#define MMQ_X_Q5_K_AMPERE 64 +#define MMQ_Y_Q5_K_AMPERE 128 +#define NWARPS_Q5_K_AMPERE 4 +#endif +#define MMQ_X_Q5_K_PASCAL 64 +#define MMQ_Y_Q5_K_PASCAL 64 +#define NWARPS_Q5_K_PASCAL 8 + +#define MMQ_X_Q6_K_RDNA2 64 +#define MMQ_Y_Q6_K_RDNA2 128 +#define NWARPS_Q6_K_RDNA2 8 +#define MMQ_X_Q6_K_RDNA1 32 +#define MMQ_Y_Q6_K_RDNA1 64 +#define NWARPS_Q6_K_RDNA1 8 +#if defined(CUDA_USE_TENSOR_CORES) +#define MMQ_X_Q6_K_AMPERE 4 +#define MMQ_Y_Q6_K_AMPERE 32 +#define NWARPS_Q6_K_AMPERE 4 +#else +#define MMQ_X_Q6_K_AMPERE 64 +#define MMQ_Y_Q6_K_AMPERE 64 +#define NWARPS_Q6_K_AMPERE 4 +#endif +#define MMQ_X_Q6_K_PASCAL 64 +#define MMQ_Y_Q6_K_PASCAL 64 +#define NWARPS_Q6_K_PASCAL 8 + + +// QK = number of values after dequantization +// QR = QK / number of values before dequantization +// QI = number of 32 bit integers before dequantization + +#define QK4_0 32 +#define QR4_0 2 +#define QI4_0 (QK4_0 / (4 * QR4_0)) +typedef struct { + half d; // delta + uint8_t qs[QK4_0 / 2]; // nibbles / quants +} block_q4_0; +static_assert(sizeof(block_q4_0) == sizeof(ggml_fp16_t) + QK4_0 / 2, "wrong q4_0 block size/padding"); + +#define QK4_1 32 +#define QR4_1 2 +#define QI4_1 (QK4_1 / (4 * QR4_1)) +typedef struct { + half2 dm; // dm.x = delta, dm.y = min + uint8_t qs[QK4_1 / 2]; // nibbles / quants +} block_q4_1; +static_assert(sizeof(block_q4_1) == sizeof(ggml_fp16_t) * 2 + QK4_1 / 2, "wrong q4_1 block size/padding"); + +#define QK5_0 32 +#define QR5_0 2 +#define QI5_0 (QK5_0 / (4 * QR5_0)) +typedef struct { + half d; // delta + uint8_t qh[4]; // 5-th bit of quants + uint8_t qs[QK5_0 / 2]; // nibbles / quants +} block_q5_0; +static_assert(sizeof(block_q5_0) == sizeof(ggml_fp16_t) + sizeof(uint32_t) + QK5_0 / 2, "wrong q5_0 block size/padding"); + +#define QK5_1 32 +#define QR5_1 2 +#define QI5_1 (QK5_1 / (4 * QR5_1)) +typedef struct { + half2 dm; // dm.x = delta, dm.y = min + uint8_t qh[4]; // 5-th bit of quants + uint8_t qs[QK5_1 / 2]; // nibbles / quants +} block_q5_1; +static_assert(sizeof(block_q5_1) == 2 * sizeof(ggml_fp16_t) + sizeof(uint32_t) + QK5_1 / 2, "wrong q5_1 block size/padding"); + +#define QK8_0 32 +#define QR8_0 1 +#define QI8_0 (QK8_0 / (4 * QR8_0)) +typedef struct { + half d; // delta + int8_t qs[QK8_0]; // quants +} block_q8_0; +static_assert(sizeof(block_q8_0) == sizeof(ggml_fp16_t) + QK8_0, "wrong q8_0 block size/padding"); + +#define QK8_1 32 +#define QR8_1 1 +#define QI8_1 (QK8_1 / (4 * QR8_1)) +typedef struct { + half2 ds; // ds.x = delta, ds.y = sum + int8_t qs[QK8_0]; // quants +} block_q8_1; +static_assert(sizeof(block_q8_1) == 2*sizeof(ggml_fp16_t) + QK8_0, "wrong q8_1 block size/padding"); + +typedef float (*vec_dot_q_cuda_t)(const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs); + +#define QR2_K 4 +#define QI2_K (QK_K / (4*QR2_K)) +typedef struct { + uint8_t scales[QK_K/16]; // scales and mins, quantized with 4 bits + uint8_t qs[QK_K/4]; // quants + half2 dm; // super-block scale for quantized scales/mins +} block_q2_K; +static_assert(sizeof(block_q2_K) == 2*sizeof(ggml_fp16_t) + QK_K/16 + QK_K/4, "wrong q2_K block size/padding"); + +#define QR3_K 4 +#define QI3_K (QK_K / (4*QR3_K)) +typedef struct { + uint8_t hmask[QK_K/8]; // quants - high bit + uint8_t qs[QK_K/4]; // quants - low 2 bits +#ifdef GGML_QKK_64 + uint8_t scales[2]; // scales, quantized with 8 bits +#else + uint8_t scales[K_SCALE_SIZE]; // scales, quantized with 6 bits +#endif + half d; // super-block scale +} block_q3_K; +//static_assert(sizeof(block_q3_K) == sizeof(ggml_fp16_t) + QK_K / 4 + QK_K / 8 + K_SCALE_SIZE, "wrong q3_K block size/padding"); + +#define QR4_K 2 +#define QI4_K (QK_K / (4*QR4_K)) +#ifdef GGML_QKK_64 +typedef struct { + half dm[2]; // super-block scales/mins + uint8_t scales[2]; // 4-bit block scales/mins + uint8_t qs[QK_K/2]; // 4--bit quants +} block_q4_K; +static_assert(sizeof(block_q4_K) == sizeof(half2) + QK_K/2 + 2, "wrong q4_K block size/padding"); +#else +typedef struct { + half2 dm; // super-block scale for quantized scales/mins + uint8_t scales[3*QK_K/64]; // scales, quantized with 6 bits + uint8_t qs[QK_K/2]; // 4--bit quants +} block_q4_K; +static_assert(sizeof(block_q4_K) == 2*sizeof(ggml_fp16_t) + 3*QK_K/64 + QK_K/2, "wrong q4_K block size/padding"); +#endif + +#define QR5_K 2 +#define QI5_K (QK_K / (4*QR5_K)) +#ifdef GGML_QKK_64 +typedef struct { + half d; // super-block scale + int8_t scales[QK_K/16]; // block scales + uint8_t qh[QK_K/8]; // quants, high bit + uint8_t qs[QK_K/2]; // quants, low 4 bits +} block_q5_K; +static_assert(sizeof(block_q5_K) == sizeof(ggml_fp16_t) + QK_K/2 + QK_K/8 + QK_K/16, "wrong q5_K block size/padding"); +#else +typedef struct { + half2 dm; // super-block scale for quantized scales/mins + uint8_t scales[K_SCALE_SIZE]; // scales and mins, quantized with 6 bits + uint8_t qh[QK_K/8]; // quants, high bit + uint8_t qs[QK_K/2]; // quants, low 4 bits +} block_q5_K; +static_assert(sizeof(block_q5_K) == 2*sizeof(ggml_fp16_t) + K_SCALE_SIZE + QK_K/2 + QK_K/8, "wrong q5_K block size/padding"); +#endif + +#define QR6_K 2 +#define QI6_K (QK_K / (4*QR6_K)) +typedef struct { + uint8_t ql[QK_K/2]; // quants, lower 4 bits + uint8_t qh[QK_K/4]; // quants, upper 2 bits + int8_t scales[QK_K/16]; // scales + half d; // delta +} block_q6_K; +static_assert(sizeof(block_q6_K) == sizeof(ggml_fp16_t) + 13*QK_K/16, "wrong q6_K block size/padding"); + +// In llama.cpp this is only used for intermediate quantization and dot products +typedef struct { + float d; // delta + int8_t qs[QK_K]; // quants + int16_t bsums[QK_K/16]; // sum of quants in groups of 16 +} block_q8_K; +static_assert(sizeof(block_q8_K) == sizeof(float) + QK_K + QK_K/16*sizeof(int16_t), "wrong q8_K block size/padding"); + + +// VDR = vec dot ratio, how many contiguous integers each thread processes when the vec dot kernel is called +// MMVQ = mul_mat_vec_q, MMQ = mul_mat_q + +#define VDR_Q4_0_Q8_1_MMVQ 2 +#define VDR_Q4_0_Q8_1_MMQ 4 + +template static __device__ __forceinline__ float vec_dot_q4_0_q8_1_impl( + const int * v, const int * u, const float & d4, const half2 & ds8) { + + int sumi = 0; + +#pragma unroll + for (int i = 0; i < vdr; ++i) { + const int vi0 = (v[i] >> 0) & 0x0F0F0F0F; + const int vi1 = (v[i] >> 4) & 0x0F0F0F0F; + + // SIMD dot product of quantized values + sumi = ggml_cuda_dp4a(vi0, u[2*i+0], sumi); + sumi = ggml_cuda_dp4a(vi1, u[2*i+1], sumi); + } + + const float2 ds8f = __half22float2(ds8); + + // second part effectively subtracts 8 from each quant value + return d4 * (sumi * ds8f.x - (8*vdr/QI4_0) * ds8f.y); +} + +#define VDR_Q4_1_Q8_1_MMVQ 2 +#define VDR_Q4_1_Q8_1_MMQ 4 + +template static __device__ __forceinline__ float vec_dot_q4_1_q8_1_impl( + const int * v, const int * u, const half2 & dm4, const half2 & ds8) { + int sumi = 0; + +#pragma unroll + for (int i = 0; i < vdr; ++i) { + const int vi0 = (v[i] >> 0) & 0x0F0F0F0F; + const int vi1 = (v[i] >> 4) & 0x0F0F0F0F; + + // SIMD dot product of quantized values + sumi = ggml_cuda_dp4a(vi0, u[2*i+0], sumi); + sumi = ggml_cuda_dp4a(vi1, u[2*i+1], sumi); + } + +#ifdef GGML_CUDA_F16 + const float2 tmp = __half22float2(__hmul2(dm4, ds8)); + const float d4d8 = tmp.x; + const float m4s8 = tmp.y; +#else + const float2 dm4f = __half22float2(dm4); + const float2 ds8f = __half22float2(ds8); + const float d4d8 = dm4f.x * ds8f.x; + const float m4s8 = dm4f.y * ds8f.y; +#endif // GGML_CUDA_F16 + + // scale second part of sum by QI8_1/(vdr * QR4_1) to compensate for multiple threads adding it + return sumi * d4d8 + m4s8 / (QI8_1 / (vdr * QR4_1)); +} + +#define VDR_Q5_0_Q8_1_MMVQ 2 +#define VDR_Q5_0_Q8_1_MMQ 4 + +template static __device__ __forceinline__ float vec_dot_q5_0_q8_1_impl( + const int * vl, const int * vh, const int * u, const float & d5, const half2 & ds8) { + + int sumi = 0; + +#pragma unroll + for (int i = 0; i < vdr; ++i) { + int vi0 = (vl[i] >> 0) & 0x0F0F0F0F; // lower 4 qs bits, still need qh as 5th bits + vi0 |= (vh[i] << 4) & 0x00000010; // 0 -> 4 + vi0 |= (vh[i] << 11) & 0x00001000; // 1 -> 12 + vi0 |= (vh[i] << 18) & 0x00100000; // 2 -> 20 + vi0 |= (vh[i] << 25) & 0x10000000; // 3 -> 28 + sumi = ggml_cuda_dp4a(vi0, u[2*i+0], sumi); // SIMD dot product of quantized values + + int vi1 = (vl[i] >> 4) & 0x0F0F0F0F; // upper 4 qs bits, still need qh as 5th bits + vi1 |= (vh[i] >> 12) & 0x00000010; // 16 -> 4 + vi1 |= (vh[i] >> 5) & 0x00001000; // 17 -> 12 + vi1 |= (vh[i] << 2) & 0x00100000; // 18 -> 20 + vi1 |= (vh[i] << 9) & 0x10000000; // 19 -> 28 + sumi = ggml_cuda_dp4a(vi1, u[2*i+1], sumi); // SIMD dot product of quantized values + } + + const float2 ds8f = __half22float2(ds8); + + // second part effectively subtracts 16 from each quant value + return d5 * (sumi * ds8f.x - (16*vdr/QI5_0) * ds8f.y); +} + +#define VDR_Q5_1_Q8_1_MMVQ 2 +#define VDR_Q5_1_Q8_1_MMQ 4 + +template static __device__ __forceinline__ float vec_dot_q5_1_q8_1_impl( + const int * vl, const int * vh, const int * u, const half2 & dm5, const half2 & ds8) { + + int sumi = 0; + +#pragma unroll + for (int i = 0; i < vdr; ++i) { + int vi0 = (vl[i] >> 0) & 0x0F0F0F0F; // lower 4 qs bits, still need qh as 5th bits + vi0 |= (vh[i] << 4) & 0x00000010; // 0 -> 4 + vi0 |= (vh[i] << 11) & 0x00001000; // 1 -> 12 + vi0 |= (vh[i] << 18) & 0x00100000; // 2 -> 20 + vi0 |= (vh[i] << 25) & 0x10000000; // 3 -> 28 + sumi = ggml_cuda_dp4a(vi0, u[2*i+0], sumi); // SIMD dot product of quantized values + + int vi1 = (vl[i] >> 4) & 0x0F0F0F0F; // upper 4 qs bits, still need qh as 5th bits + vi1 |= (vh[i] >> 12) & 0x00000010; // 16 -> 4 + vi1 |= (vh[i] >> 5) & 0x00001000; // 17 -> 12 + vi1 |= (vh[i] << 2) & 0x00100000; // 18 -> 20 + vi1 |= (vh[i] << 9) & 0x10000000; // 19 -> 28 + sumi = ggml_cuda_dp4a(vi1, u[2*i+1], sumi); // SIMD dot product of quantized values + } + +#ifdef GGML_CUDA_F16 + const float2 tmp = __half22float2(__hmul2(dm5, ds8)); + const float d5d8 = tmp.x; + const float m5s8 = tmp.y; +#else + const float2 dm5f = __half22float2(dm5); + const float2 ds8f = __half22float2(ds8); + const float d5d8 = dm5f.x * ds8f.x; + const float m5s8 = dm5f.y * ds8f.y; +#endif // GGML_CUDA_F16 + + // scale second part of sum by QI5_1 / vdr to compensate for multiple threads adding it + return sumi*d5d8 + m5s8 / (QI5_1 / vdr); +} + +#define VDR_Q8_0_Q8_1_MMVQ 2 +#define VDR_Q8_0_Q8_1_MMQ 8 + +template static __device__ __forceinline__ float vec_dot_q8_0_q8_1_impl( + const int * v, const int * u, const float & d8_0, const float & d8_1) { + + int sumi = 0; + +#pragma unroll + for (int i = 0; i < vdr; ++i) { + // SIMD dot product of quantized values + sumi = ggml_cuda_dp4a(v[i], u[i], sumi); + } + + return d8_0*d8_1 * sumi; +} + +template static __device__ __forceinline__ float vec_dot_q8_1_q8_1_impl( + const int * v, const int * u, const half2 & dm8, const half2 & ds8) { + + int sumi = 0; + +#pragma unroll + for (int i = 0; i < vdr; ++i) { + // SIMD dot product of quantized values + sumi = ggml_cuda_dp4a(v[i], u[i], sumi); + } + +#ifdef GGML_CUDA_F16 + const float2 tmp = __half22float2(__hmul2(dm8, ds8)); + const float d8d8 = tmp.x; + const float m8s8 = tmp.y; +#else + const float2 dm8f = __half22float2(dm8); + const float2 ds8f = __half22float2(ds8); + const float d8d8 = dm8f.x * ds8f.x; + const float m8s8 = dm8f.y * ds8f.y; +#endif // GGML_CUDA_F16 + + // scale second part of sum by QI8_1/ vdr to compensate for multiple threads adding it + return sumi*d8d8 + m8s8 / (QI8_1 / vdr); +} + +#define VDR_Q2_K_Q8_1_MMVQ 1 +#define VDR_Q2_K_Q8_1_MMQ 2 + +// contiguous v/x values +static __device__ __forceinline__ float vec_dot_q2_K_q8_1_impl_mmvq( + const int & v, const int * __restrict__ u, const uint8_t * __restrict__ scales, + const half2 & dm2, const float * __restrict__ d8) { + + float sumf_d = 0.0f; + float sumf_m = 0.0f; + +#pragma unroll + for (int i = 0; i < QR2_K; ++i) { + const int sc = scales[2*i]; + + const int vi = (v >> (2*i)) & 0x03030303; + + sumf_d += d8[i] * (ggml_cuda_dp4a(vi, u[i], 0) * (sc & 0xF)); // SIMD dot product + + // fill int with 4x m + int m = sc >> 4; + m |= m << 8; + m |= m << 16; + sumf_m += d8[i] * ggml_cuda_dp4a(m, u[i], 0); // multiply constant q2_K part with sum of q8_1 values + } + + const float2 dm2f = __half22float2(dm2); + + return dm2f.x*sumf_d - dm2f.y*sumf_m; +} + +// contiguous u/y values +static __device__ __forceinline__ float vec_dot_q2_K_q8_1_impl_mmq( + const int * __restrict__ v, const int * __restrict__ u, const uint8_t * __restrict__ scales, + const half2 & dm2, const float & d8) { + + int sumi_d = 0; + int sumi_m = 0; + +#pragma unroll + for (int i0 = 0; i0 < QI8_1; i0 += QI8_1/2) { + int sumi_d_sc = 0; + + const int sc = scales[i0 / (QI8_1/2)]; + + // fill int with 4x m + int m = sc >> 4; + m |= m << 8; + m |= m << 16; + +#pragma unroll + for (int i = i0; i < i0 + QI8_1/2; ++i) { + sumi_d_sc = ggml_cuda_dp4a(v[i], u[i], sumi_d_sc); // SIMD dot product + sumi_m = ggml_cuda_dp4a(m, u[i], sumi_m); // multiply sum of q8_1 values with m + } + + sumi_d += sumi_d_sc * (sc & 0xF); + } + + const float2 dm2f = __half22float2(dm2); + + return d8 * (dm2f.x*sumi_d - dm2f.y*sumi_m); +} + +#define VDR_Q3_K_Q8_1_MMVQ 1 +#define VDR_Q3_K_Q8_1_MMQ 2 + +// contiguous v/x values +static __device__ __forceinline__ float vec_dot_q3_K_q8_1_impl_mmvq( + const int & vl, const int & vh, const int * __restrict__ u, const uint8_t * __restrict__ scales, + const int & scale_offset, const float & d3, const float * __restrict__ d8) { + + float sumf = 0.0f; + +#pragma unroll + for (int i = 0; i < QR3_K; ++i) { + const int isc = scale_offset + 2*i; + + const int isc_low = isc % (QK_K/32); + const int sc_shift_low = 4 * (isc / (QK_K/32)); + const int sc_low = (scales[isc_low] >> sc_shift_low) & 0xF; + + const int isc_high = isc % (QK_K/64); + const int sc_shift_high = 2 * (isc / (QK_K/64)); + const int sc_high = ((scales[(QK_K/32) + isc_high] >> sc_shift_high) & 3) << 4; + + const int sc = (sc_low | sc_high) - 32; + + const int vil = (vl >> (2*i)) & 0x03030303; + + const int vih = ((vh >> i) << 2) & 0x04040404; + + const int vi = __vsubss4(vil, vih); + + sumf += d8[i] * (ggml_cuda_dp4a(vi, u[i], 0) * sc); // SIMD dot product + } + + return d3 * sumf; +} + +// contiguous u/y values +static __device__ __forceinline__ float vec_dot_q3_K_q8_1_impl_mmq( + const int * __restrict__ v, const int * __restrict__ u, const int8_t * __restrict__ scales, + const float & d3, const float & d8) { + + int sumi = 0; + +#pragma unroll + for (int i0 = 0; i0 < QR3_K*VDR_Q3_K_Q8_1_MMQ; i0 += QI8_1/2) { + int sumi_sc = 0; + + for (int i = i0; i < i0 + QI8_1/2; ++i) { + sumi_sc = ggml_cuda_dp4a(v[i], u[i], sumi_sc); // SIMD dot product + } + + sumi += sumi_sc * scales[i0 / (QI8_1/2)]; + } + + return d3*d8 * sumi; +} + +#define VDR_Q4_K_Q8_1_MMVQ 2 +#define VDR_Q4_K_Q8_1_MMQ 8 + +// contiguous v/x values +static __device__ __forceinline__ float vec_dot_q4_K_q8_1_impl_vmmq( + const int * __restrict__ v, const int * __restrict__ u, const uint8_t * __restrict__ sc, + const uint8_t * __restrict__ m, const half2 & dm4, const float * __restrict__ d8) { + + float sumf_d = 0.0f; + float sumf_m = 0.0f; + +#pragma unroll + for (int i = 0; i < QR4_K; ++i) { + const int v0i = (v[0] >> (4*i)) & 0x0F0F0F0F; + const int v1i = (v[1] >> (4*i)) & 0x0F0F0F0F; + + const int dot1 = ggml_cuda_dp4a(v1i, u[2*i+1], ggml_cuda_dp4a(v0i, u[2*i+0], 0)); // SIMD dot product + const int dot2 = ggml_cuda_dp4a(0x01010101, u[2*i+1], ggml_cuda_dp4a(0x01010101, u[2*i+0], 0)); // sum of u + + sumf_d += d8[i] * (dot1 * sc[i]); + sumf_m += d8[i] * (dot2 * m[i]); // multiply constant part of q4_K with sum of q8_1 values + } + + const float2 dm4f = __half22float2(dm4); + + return dm4f.x*sumf_d - dm4f.y*sumf_m; +} + +// contiguous u/y values +static __device__ __forceinline__ float vec_dot_q4_K_q8_1_impl_mmq( + const int * __restrict__ v, const int * __restrict__ u, const uint8_t * __restrict__ sc, + const uint8_t * __restrict__ m, const half2 & dm4, const half2 * __restrict__ ds8) { + + float sumf_d = 0.0f; + float sumf_m = 0.0f; + +#pragma unroll + for (int i = 0; i < QR4_K*VDR_Q4_K_Q8_1_MMQ/QI8_1; ++i) { + int sumi_d = 0; + +#pragma unroll + for (int j = 0; j < QI8_1; ++j) { + sumi_d = ggml_cuda_dp4a((v[j] >> (4*i)) & 0x0F0F0F0F, u[i*QI8_1 + j], sumi_d); // SIMD dot product + } + + const float2 ds8f = __half22float2(ds8[i]); + + sumf_d += ds8f.x * (sc[i] * sumi_d); + sumf_m += ds8f.y * m[i]; // sum of q8_1 block * q4_K min val + } + + const float2 dm4f = __half22float2(dm4); + + return dm4f.x*sumf_d - dm4f.y*sumf_m; +} + +#define VDR_Q5_K_Q8_1_MMVQ 2 +#define VDR_Q5_K_Q8_1_MMQ 8 + +// contiguous v/x values +static __device__ __forceinline__ float vec_dot_q5_K_q8_1_impl_vmmq( + const int * __restrict__ vl, const int * __restrict__ vh, const int * __restrict__ u, const uint8_t * __restrict__ sc, + const uint8_t * __restrict__ m, const half2 & dm5, const float * __restrict__ d8) { + + float sumf_d = 0.0f; + float sumf_m = 0.0f; + +#pragma unroll + for (int i = 0; i < QR5_K; ++i) { + const int vl0i = (vl[0] >> (4*i)) & 0x0F0F0F0F; + const int vl1i = (vl[1] >> (4*i)) & 0x0F0F0F0F; + + const int vh0i = ((vh[0] >> i) << 4) & 0x10101010; + const int vh1i = ((vh[1] >> i) << 4) & 0x10101010; + + const int v0i = vl0i | vh0i; + const int v1i = vl1i | vh1i; + + const int dot1 = ggml_cuda_dp4a(v0i, u[2*i+0], ggml_cuda_dp4a(v1i, u[2*i+1], 0)); // SIMD dot product + const int dot2 = ggml_cuda_dp4a(0x01010101, u[2*i+0], ggml_cuda_dp4a(0x01010101, u[2*i+1], 0)); // sum of u + + sumf_d += d8[i] * (dot1 * sc[i]); + sumf_m += d8[i] * (dot2 * m[i]); + + } + + const float2 dm5f = __half22float2(dm5); + + return dm5f.x*sumf_d - dm5f.y*sumf_m; +} + +// contiguous u/y values +static __device__ __forceinline__ float vec_dot_q5_K_q8_1_impl_mmq( + const int * __restrict__ v, const int * __restrict__ u, const uint8_t * __restrict__ sc, + const uint8_t * __restrict__ m, const half2 & dm4, const half2 * __restrict__ ds8) { + + float sumf_d = 0.0f; + float sumf_m = 0.0f; + +#pragma unroll + for (int i = 0; i < QR5_K*VDR_Q5_K_Q8_1_MMQ/QI8_1; ++i) { + int sumi_d = 0; + +#pragma unroll + for (int j = 0; j < QI8_1; ++j) { + sumi_d = ggml_cuda_dp4a(v[i*QI8_1 + j], u[i*QI8_1 + j], sumi_d); // SIMD dot product + } + + const float2 ds8f = __half22float2(ds8[i]); + + sumf_d += ds8f.x * (sc[i] * sumi_d); + sumf_m += ds8f.y * m[i]; // sum of q8_1 block * q4_K min val + } + + const float2 dm4f = __half22float2(dm4); + + return dm4f.x*sumf_d - dm4f.y*sumf_m; +} + +#define VDR_Q6_K_Q8_1_MMVQ 1 +#define VDR_Q6_K_Q8_1_MMQ 8 + +// contiguous v/x values +static __device__ __forceinline__ float vec_dot_q6_K_q8_1_impl_mmvq( + const int & vl, const int & vh, const int * __restrict__ u, const int8_t * __restrict__ scales, + const float & d, const float * __restrict__ d8) { + + float sumf = 0.0f; + +#pragma unroll + for (int i = 0; i < QR6_K; ++i) { + const int sc = scales[4*i]; + + const int vil = (vl >> (4*i)) & 0x0F0F0F0F; + + const int vih = ((vh >> (4*i)) << 4) & 0x30303030; + + const int vi = __vsubss4((vil | vih), 0x20202020); // vi = (vil | vih) - 32 + + sumf += d8[i] * (ggml_cuda_dp4a(vi, u[i], 0) * sc); // SIMD dot product + } + + return d*sumf; +} + +// contiguous u/y values +static __device__ __forceinline__ float vec_dot_q6_K_q8_1_impl_mmq( + const int * __restrict__ v, const int * __restrict__ u, const int8_t * __restrict__ sc, + const float & d6, const float * __restrict__ d8) { + + float sumf_d = 0.0f; + +#pragma unroll + for (int i0 = 0; i0 < VDR_Q6_K_Q8_1_MMQ; i0 += 4) { + int2 sumi_d = {0, 0}; // 2 q6_K scales per q8_1 scale + +#pragma unroll + for (int i = i0; i < i0 + 2; ++i) { + sumi_d.x = ggml_cuda_dp4a(v[2*i+0], u[2*i+0], sumi_d.x); // SIMD dot product + sumi_d.x = ggml_cuda_dp4a(v[2*i+1], u[2*i+1], sumi_d.x); // SIMD dot product + + sumi_d.y = ggml_cuda_dp4a(v[2*i+4], u[2*i+4], sumi_d.y); // SIMD dot product + sumi_d.y = ggml_cuda_dp4a(v[2*i+5], u[2*i+5], sumi_d.y); // SIMD dot product + } + + sumf_d += d8[i0/4] * (sc[i0/2+0]*sumi_d.x + sc[i0/2+1]*sumi_d.y); + } + + return d6 * sumf_d; +} + +static __device__ __forceinline__ float vec_dot_q4_0_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { + + const block_q4_0 * bq4_0 = (const block_q4_0 *) vbq; + + int v[VDR_Q4_0_Q8_1_MMVQ]; + int u[2*VDR_Q4_0_Q8_1_MMVQ]; + +#pragma unroll + for (int i = 0; i < VDR_Q4_0_Q8_1_MMVQ; ++i) { + v[i] = get_int_from_uint8(bq4_0->qs, iqs + i); + u[2*i+0] = get_int_from_int8_aligned(bq8_1->qs, iqs + i); + u[2*i+1] = get_int_from_int8_aligned(bq8_1->qs, iqs + i + QI4_0); + } + + return vec_dot_q4_0_q8_1_impl(v, u, bq4_0->d, bq8_1->ds); +} + + +static __device__ __forceinline__ float vec_dot_q4_1_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { + + const block_q4_1 * bq4_1 = (const block_q4_1 *) vbq; + + int v[VDR_Q4_1_Q8_1_MMVQ]; + int u[2*VDR_Q4_1_Q8_1_MMVQ]; + +#pragma unroll + for (int i = 0; i < VDR_Q4_1_Q8_1_MMVQ; ++i) { + v[i] = get_int_from_uint8_aligned(bq4_1->qs, iqs + i); + u[2*i+0] = get_int_from_int8_aligned(bq8_1->qs, iqs + i); + u[2*i+1] = get_int_from_int8_aligned(bq8_1->qs, iqs + i + QI4_1); + } + + return vec_dot_q4_1_q8_1_impl(v, u, bq4_1->dm, bq8_1->ds); +} + +static __device__ __forceinline__ float vec_dot_q5_0_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { + + const block_q5_0 * bq5_0 = (const block_q5_0 *) vbq; + + int vl[VDR_Q5_0_Q8_1_MMVQ]; + int vh[VDR_Q5_0_Q8_1_MMVQ]; + int u[2*VDR_Q5_0_Q8_1_MMVQ]; + +#pragma unroll + for (int i = 0; i < VDR_Q5_0_Q8_1_MMVQ; ++i) { + vl[i] = get_int_from_uint8(bq5_0->qs, iqs + i); + vh[i] = get_int_from_uint8(bq5_0->qh, 0) >> (4 * (iqs + i)); + u[2*i+0] = get_int_from_int8_aligned(bq8_1->qs, iqs + i); + u[2*i+1] = get_int_from_int8_aligned(bq8_1->qs, iqs + i + QI5_0); + } + + return vec_dot_q5_0_q8_1_impl(vl, vh, u, bq5_0->d, bq8_1->ds); +} + +static __device__ __forceinline__ float vec_dot_q5_1_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { + + const block_q5_1 * bq5_1 = (const block_q5_1 *) vbq; + + int vl[VDR_Q5_1_Q8_1_MMVQ]; + int vh[VDR_Q5_1_Q8_1_MMVQ]; + int u[2*VDR_Q5_1_Q8_1_MMVQ]; + +#pragma unroll + for (int i = 0; i < VDR_Q5_1_Q8_1_MMVQ; ++i) { + vl[i] = get_int_from_uint8_aligned(bq5_1->qs, iqs + i); + vh[i] = get_int_from_uint8_aligned(bq5_1->qh, 0) >> (4 * (iqs + i)); + u[2*i+0] = get_int_from_int8_aligned(bq8_1->qs, iqs + i); + u[2*i+1] = get_int_from_int8_aligned(bq8_1->qs, iqs + i + QI5_1); + } + + return vec_dot_q5_1_q8_1_impl(vl, vh, u, bq5_1->dm, bq8_1->ds); +} + +static __device__ __forceinline__ float vec_dot_q8_0_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { + + const block_q8_0 * bq8_0 = (const block_q8_0 *) vbq; + + int v[VDR_Q8_0_Q8_1_MMVQ]; + int u[VDR_Q8_0_Q8_1_MMVQ]; + +#pragma unroll + for (int i = 0; i < VDR_Q8_0_Q8_1_MMVQ; ++i) { + v[i] = get_int_from_int8(bq8_0->qs, iqs + i); + u[i] = get_int_from_int8_aligned(bq8_1->qs, iqs + i); + } + + return vec_dot_q8_0_q8_1_impl(v, u, bq8_0->d, __low2half(bq8_1->ds)); +} + +static __device__ __forceinline__ float vec_dot_q2_K_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { + + const block_q2_K * bq2_K = (const block_q2_K *) vbq; + + const int bq8_offset = QR2_K * (iqs / QI8_1); + const int scale_offset = iqs - iqs % QI8_1 + (iqs % QI8_1) / (QI8_1/2); + + const uint8_t * scales = bq2_K->scales + scale_offset; + + const int v = get_int_from_uint8_aligned(bq2_K->qs, iqs); + int u[QR2_K]; + float d8[QR2_K]; + +#pragma unroll + for (int i = 0; i < QR2_K; ++ i) { + u[i] = get_int_from_int8_aligned(bq8_1[bq8_offset + i].qs, iqs % QI8_1); + d8[i] = __low2float(bq8_1[bq8_offset + i].ds); + } + + return vec_dot_q2_K_q8_1_impl_mmvq(v, u, scales, bq2_K->dm, d8); +} + +static __device__ __forceinline__ float vec_dot_q3_K_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { + + const block_q3_K * bq3_K = (const block_q3_K *) vbq; + + const int bq8_offset = QR3_K * (iqs / (QI3_K/2)); + const int scale_offset = iqs - iqs % QI8_1 + (iqs % QI8_1) / (QI8_1/2); + + const float d = bq3_K->d; + + const int vl = get_int_from_uint8(bq3_K->qs, iqs); + + // invert the mask with ~ so that a 0/1 results in 4/0 being subtracted + const int vh = ~get_int_from_uint8(bq3_K->hmask, iqs % (QI3_K/2)) >> bq8_offset; + + int u[QR3_K]; + float d8[QR3_K]; + +#pragma unroll + for (int i = 0; i < QR3_K; ++i) { + u[i] = get_int_from_int8_aligned(bq8_1[bq8_offset + i].qs, iqs % QI8_1); + d8[i] = __low2float(bq8_1[bq8_offset + i].ds); + } + + return vec_dot_q3_K_q8_1_impl_mmvq(vl, vh, u, bq3_K->scales, scale_offset, d, d8); +} + +static __device__ __forceinline__ float vec_dot_q4_K_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { + +#ifndef GGML_QKK_64 + const block_q4_K * bq4_K = (const block_q4_K *) vbq; + + int v[2]; + int u[2*QR4_K]; + float d8[QR4_K]; + + // iqs is in 0,2..30. bq8_offset = iqs/4 -> bq8_offset = 0, 2, 4, 6 + const int bq8_offset = QR4_K * ((iqs/2) / (QI8_1/2)); + + // iqs = 0....3 -> bq8_offset = 0, want q4_offset = 0, 4, 8, 12 + // iqs = 4....7 -> bq8_offset = 2, want q4_offset = 32, 36, 40, 44 + // iqs = 8...11 -> bq8_offset = 4, want q4_offset = 64, 68, 72, 76 + // iqs = 12..15 -> bq8_offset = 6, want q4_offset = 96, 100, 104, 108 + + const int * q4 = (const int *)(bq4_K->qs + 16 * bq8_offset + 4 * ((iqs/2)%4)); + v[0] = q4[0]; + v[1] = q4[4]; + + const uint16_t * scales = (const uint16_t *)bq4_K->scales; + uint16_t aux[2]; + const int j = bq8_offset/2; + if (j < 2) { + aux[0] = scales[j+0] & 0x3f3f; + aux[1] = scales[j+2] & 0x3f3f; + } else { + aux[0] = ((scales[j+2] >> 0) & 0x0f0f) | ((scales[j-2] & 0xc0c0) >> 2); + aux[1] = ((scales[j+2] >> 4) & 0x0f0f) | ((scales[j-0] & 0xc0c0) >> 2); + } + const uint8_t * sc = (const uint8_t *)aux; + const uint8_t * m = sc + 2; + + for (int i = 0; i < QR4_K; ++i) { + const block_q8_1 * bq8i = bq8_1 + bq8_offset + i; + d8[i] = __low2float(bq8i->ds); + + const int * q8 = (const int *)bq8i->qs + ((iqs/2)%4); + u[2*i+0] = q8[0]; + u[2*i+1] = q8[4]; + } + + return vec_dot_q4_K_q8_1_impl_vmmq(v, u, sc, m, bq4_K->dm, d8); + +#else + + const block_q4_K * bq4_K = (const block_q4_K *) vbq; + + float sumf_d = 0.0f; + float sumf_m = 0.0f; + + uint16_t aux16[2]; + const uint8_t * s = (const uint8_t *)aux16; + + const uint16_t * a = (const uint16_t *)bq4_K->scales; + aux16[0] = a[0] & 0x0f0f; + aux16[1] = (a[0] >> 4) & 0x0f0f; + + const float dall = bq4_K->dm[0]; + const float dmin = bq4_K->dm[1]; + + const float d8_1 = __low2float(bq8_1[0].ds); + const float d8_2 = __low2float(bq8_1[1].ds); + + const int ui1 = *((const int *)bq8_1[0].qs + (iqs/2)); + const int ui2 = *((const int *)bq8_1[0].qs + (iqs/2) + 4); + const int ui3 = *((const int *)bq8_1[1].qs + (iqs/2)); + const int ui4 = *((const int *)bq8_1[1].qs + (iqs/2) + 4); + + const int * q4 = (const int *)bq4_K->qs + (iqs/2); + const int v1 = q4[0]; + const int v2 = q4[4]; + + const int dot1 = ggml_cuda_dp4a(ui2, v2 & 0x0f0f0f0f, ggml_cuda_dp4a(ui1, v1 & 0x0f0f0f0f, 0)); + const int dot2 = ggml_cuda_dp4a(ui4, (v2 >> 4) & 0x0f0f0f0f, ggml_cuda_dp4a(ui3, (v1 >> 4) & 0x0f0f0f0f, 0)); + const int dot3 = ggml_cuda_dp4a(0x01010101, ui2, ggml_cuda_dp4a(0x01010101, ui1, 0)); + const int dot4 = ggml_cuda_dp4a(0x01010101, ui4, ggml_cuda_dp4a(0x01010101, ui3, 0)); + + sumf_d += d8_1 * (dot1 * s[0]) + d8_2 * (dot2 * s[1]); + sumf_m += d8_1 * (dot3 * s[2]) + d8_2 * (dot4 * s[3]); + + return dall * sumf_d - dmin * sumf_m; +#endif +} + +static __device__ __forceinline__ float vec_dot_q5_K_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { + +#ifndef GGML_QKK_64 + const block_q5_K * bq5_K = (const block_q5_K *) vbq; + + int vl[2]; + int vh[2]; + int u[2*QR5_K]; + float d8[QR5_K]; + + const int bq8_offset = QR5_K * ((iqs/2) / (QI8_1/2)); + const int * ql = (const int *)(bq5_K->qs + 16 * bq8_offset + 4 * ((iqs/2)%4)); + const int * qh = (const int *)(bq5_K->qh + 4 * ((iqs/2)%4)); + + vl[0] = ql[0]; + vl[1] = ql[4]; + + vh[0] = qh[0] >> bq8_offset; + vh[1] = qh[4] >> bq8_offset; + + const uint16_t * scales = (const uint16_t *)bq5_K->scales; + uint16_t aux[2]; + const int j = bq8_offset/2; + if (j < 2) { + aux[0] = scales[j+0] & 0x3f3f; + aux[1] = scales[j+2] & 0x3f3f; + } else { + aux[0] = ((scales[j+2] >> 0) & 0x0f0f) | ((scales[j-2] & 0xc0c0) >> 2); + aux[1] = ((scales[j+2] >> 4) & 0x0f0f) | ((scales[j-0] & 0xc0c0) >> 2); + } + const uint8_t * sc = (const uint8_t *)aux; + const uint8_t * m = sc + 2; + +#pragma unroll + for (int i = 0; i < QR5_K; ++i) { + const block_q8_1 * bq8i = bq8_1 + bq8_offset + i; + d8[i] = __low2float(bq8i->ds); + + const int * q8 = (const int *)bq8i->qs + ((iqs/2)%4); + u[2*i+0] = q8[0]; + u[2*i+1] = q8[4]; + } + + return vec_dot_q5_K_q8_1_impl_vmmq(vl, vh, u, sc, m, bq5_K->dm, d8); + +#else + + const block_q5_K * bq5_K = (const block_q5_K *) vbq; + + const int8_t * s = bq5_K->scales; + + const float d = bq5_K->d; + + const float d8_1 = __low2half(bq8_1[0].ds); + const float d8_2 = __low2half(bq8_1[1].ds); + + const int ui1 = *((const int *)bq8_1[0].qs + (iqs/2)); + const int ui2 = *((const int *)bq8_1[0].qs + (iqs/2) + 4); + const int ui3 = *((const int *)bq8_1[1].qs + (iqs/2)); + const int ui4 = *((const int *)bq8_1[1].qs + (iqs/2) + 4); + + const int * ql = (const int *)bq5_K->qs + (iqs/2); + const int vl1 = ql[0]; + const int vl2 = ql[4]; + + const int step = 4 * (iqs/2); // 0, 4, 8, 12 + const int im = step/8; // = 0 for iqs = 0, 2, = 1 for iqs = 4, 6 + const int in = step%8; // 0, 4, 0, 4 + const int vh = (*((const int *)(bq5_K->qh + in))) >> im; + + const int v1 = (((vh << 4) & 0x10101010) ^ 0x10101010) | ((vl1 >> 0) & 0x0f0f0f0f); + const int v2 = (((vh << 2) & 0x10101010) ^ 0x10101010) | ((vl2 >> 0) & 0x0f0f0f0f); + const int v3 = (((vh >> 0) & 0x10101010) ^ 0x10101010) | ((vl1 >> 4) & 0x0f0f0f0f); + const int v4 = (((vh >> 2) & 0x10101010) ^ 0x10101010) | ((vl2 >> 4) & 0x0f0f0f0f); + + const float sumf_d = d8_1 * (ggml_cuda_dp4a(ui1, v1, 0) * s[0] + ggml_cuda_dp4a(ui2, v2, 0) * s[1]) + + d8_2 * (ggml_cuda_dp4a(ui3, v3, 0) * s[2] + ggml_cuda_dp4a(ui4, v4, 0) * s[3]); + + return d * sumf_d; +#endif +} + +static __device__ __forceinline__ float vec_dot_q6_K_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { + + const block_q6_K * bq6_K = (const block_q6_K *) vbq; + + const int bq8_offset = 2 * QR6_K * (iqs / (QI6_K/2)) + (iqs % (QI6_K/2)) / (QI6_K/4); + const int scale_offset = (QI6_K/4) * (iqs / (QI6_K/2)) + (iqs % (QI6_K/2)) / (QI6_K/8); + const int vh_shift = 2 * ((iqs % (QI6_K/2)) / (QI6_K/4)); + + const int vl = get_int_from_uint8(bq6_K->ql, iqs); + const int vh = get_int_from_uint8(bq6_K->qh, (QI6_K/4) * (iqs / (QI6_K/2)) + iqs % (QI6_K/4)) >> vh_shift; + + const int8_t * scales = bq6_K->scales + scale_offset; + + int u[QR6_K]; + float d8[QR6_K]; + +#pragma unroll + for (int i = 0; i < QR6_K; ++i) { + u[i] = get_int_from_int8_aligned(bq8_1[bq8_offset + 2*i].qs, iqs % QI8_1); + d8[i] = __low2float(bq8_1[bq8_offset + 2*i].ds); + } + + return vec_dot_q6_K_q8_1_impl_mmvq(vl, vh, u, scales, bq6_K->d, d8); +} + +static __global__ void quantize_q8_1(const float * __restrict__ x, void * __restrict__ vy, const int kx, const int kx_padded) { + const int ix = blockDim.x*blockIdx.x + threadIdx.x; + if (ix >= kx_padded) { + return; + } + const int iy = blockDim.y*blockIdx.y + threadIdx.y; + const int i_padded = iy*kx_padded + ix; + block_q8_1 * y = (block_q8_1 *) vy; + + const int ib = i_padded / QK8_1; // block index + const int iqs = i_padded % QK8_1; // quant index + + const float xi = ix < kx ? x[iy*kx + ix] : 0.0f; + float amax = fabsf(xi); + float sum = xi; + + amax = warp_reduce_max(amax); + sum = warp_reduce_sum(sum); + + const float d = amax / 127; + const int8_t q = amax == 0.0f ? 0 : roundf(xi / d); + + y[ib].qs[iqs] = q; + if (iqs > 0) { + return; + } + reinterpret_cast(y[ib].ds.x) = d; + reinterpret_cast(y[ib].ds.y) = sum; +} + +template +static __device__ __forceinline__ dst_t convert_from_half(half val) { + return val; +} + +template<> +__device__ __forceinline__ nv_bfloat16 convert_from_half(half val) { +#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 + return __float2bfloat16(__half2float(val)); +#else + return __half2float(val); +#endif // defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800 +} + +template<> +__device__ __forceinline__ float convert_from_half(half val) { + return __half2float(val); +} + +template +inline __device__ void dequantize_block_q2_K(const void * __restrict__ vx, dst_t * __restrict__ yy) { + + const auto i = 0; //we only need dequant one block in each call + const block_q2_K * x = (const block_q2_K *) vx; + + const auto tid = threadIdx.x; + const int n = tid/32; + const int l = tid - 32*n; + const int is = 8*n + l/16; + + const uint8_t q = x[i].qs[32*n + l]; + dst_t * y = yy + i*QK_K + 128*n; + + half dall = __low2half(x[i].dm); + half dmin = __high2half(x[i].dm); + y[l+ 0] = convert_from_half(__hsub(__hmul(dall, __int2half_rn((x[i].scales[is+0] & 0xF) * ((q >> 0) & 3))), __hmul(dmin, __int2half_rn(x[i].scales[is+0] >> 4)))); + y[l+32] = convert_from_half(__hsub(__hmul(dall, __int2half_rn((x[i].scales[is+2] & 0xF) * ((q >> 2) & 3))), __hmul(dmin, __int2half_rn(x[i].scales[is+2] >> 4)))); + y[l+64] = convert_from_half(__hsub(__hmul(dall, __int2half_rn((x[i].scales[is+4] & 0xF) * ((q >> 4) & 3))), __hmul(dmin, __int2half_rn(x[i].scales[is+4] >> 4)))); + y[l+96] = convert_from_half(__hsub(__hmul(dall, __int2half_rn((x[i].scales[is+6] & 0xF) * ((q >> 6) & 3))), __hmul(dmin, __int2half_rn(x[i].scales[is+6] >> 4)))); +} + +template +inline __device__ void dequantize_block_q3_K(const void * __restrict__ vx, dst_t * __restrict__ yy) { + + const auto i = 0; + const block_q3_K * x = (const block_q3_K *) vx; + + const auto r = threadIdx.x/4; + const int tid = r/2; + const int is0 = r%2; + const int l0 = 16*is0 + 4*(threadIdx.x%4); + const int n = tid / 4; + const int j = tid - 4*n; + + uint8_t m = 1 << (4*n + j); + int is = 8*n + 2*j + is0; + int shift = 2*j; + + int8_t us = is < 4 ? (x[i].scales[is-0] & 0xF) | (((x[i].scales[is+8] >> 0) & 3) << 4) : + is < 8 ? (x[i].scales[is-0] & 0xF) | (((x[i].scales[is+4] >> 2) & 3) << 4) : + is < 12 ? (x[i].scales[is-8] >> 4) | (((x[i].scales[is+0] >> 4) & 3) << 4) : + (x[i].scales[is-8] >> 4) | (((x[i].scales[is-4] >> 6) & 3) << 4); + half d_all = x[i].d; + half dl = __hmul(d_all, __int2half_rn(us - 32)); + + dst_t * y = yy + i*QK_K + 128*n + 32*j; + const uint8_t * q = x[i].qs + 32*n; + const uint8_t * hm = x[i].hmask; + + for (int l = l0; l < l0+4; ++l) { + y[l] = convert_from_half(__hmul(dl, __int2half_rn((int8_t)((q[l] >> shift) & 3) - ((hm[l] & m) ? 0 : 4)))); + } +} + +static inline __device__ void get_scale_min_k4(int j, const uint8_t * q, uint8_t & d, uint8_t & m) { + if (j < 4) { + d = q[j] & 63; m = q[j + 4] & 63; + } else { + d = (q[j+4] & 0xF) | ((q[j-4] >> 6) << 4); + m = (q[j+4] >> 4) | ((q[j-0] >> 6) << 4); + } +} + +template +inline __device__ void dequantize_block_q4_K(const void * __restrict__ vx, dst_t * __restrict__ yy) { + const block_q4_K * x = (const block_q4_K *) vx; + + const auto i = 0; + + // assume 32 threads + const auto tid = threadIdx.x; + const int il = tid/8; + const int ir = tid%8; + const int is = 2*il; + const int n = 4; + + dst_t * y = yy + i*QK_K + 64*il + n*ir; + + const half dall = __low2half(x[i].dm); + const half dmin = __high2half(x[i].dm); + + const uint8_t * q = x[i].qs + 32*il + n*ir; + + uint8_t sc, m; + get_scale_min_k4(is + 0, x[i].scales, sc, m); + const half d1 = __hmul(dall, __int2half_rn(sc)); + const half m1 = __hmul(dmin, __int2half_rn(m)); + get_scale_min_k4(is + 1, x[i].scales, sc, m); + const half d2 = __hmul(dall, __int2half_rn(sc)); + const half m2 = __hmul(dmin, __int2half_rn(m)); + for (int l = 0; l < n; ++l) { + y[l + 0] = convert_from_half(__hsub(__hmul(d1, __int2half_rn(q[l] & 0xF)), m1)); + y[l +32] = convert_from_half(__hsub(__hmul(d2, __int2half_rn(q[l] >> 4)), m2)); + } +} + +template +inline __device__ void dequantize_block_q5_K(const void * __restrict__ vx, dst_t * __restrict__ yy) { + const block_q5_K * x = (const block_q5_K *) vx; + + const auto i = 0; + + // assume 64 threads - this is very slightly better than the one below + const auto tid = threadIdx.x; + const int il = tid/16; // il is in 0...3 + const int ir = tid%16; // ir is in 0...15 + const int is = 2*il; // is is in 0...6 + + dst_t * y = yy + i*QK_K + 64*il + 2*ir; + + const half dall = __low2half(x[i].dm); + const half dmin = __high2half(x[i].dm); + + const uint8_t * ql = x[i].qs + 32*il + 2*ir; + const uint8_t * qh = x[i].qh + 2*ir; + + uint8_t sc, m; + get_scale_min_k4(is + 0, x[i].scales, sc, m); + const half d1 = __hmul(dall, __int2half_rn(sc)); const half m1 = __hmul(dmin, __int2half_rn(m)); + get_scale_min_k4(is + 1, x[i].scales, sc, m); + const half d2 = __hmul(dall, __int2half_rn(sc)); const half m2 = __hmul(dmin, __int2half_rn(m)); + + uint8_t hm = 1 << (2*il); + y[ 0] = convert_from_half(__hsub(__hmul(d1, __int2half_rn((ql[0] & 0xF) + (qh[0] & hm ? 16 : 0))), m1)); + y[ 1] = convert_from_half(__hsub(__hmul(d1, __int2half_rn((ql[1] & 0xF) + (qh[1] & hm ? 16 : 0))), m1)); + hm <<= 1; + y[32] = convert_from_half(__hsub(__hmul(d2, __int2half_rn((ql[0] >> 4) + (qh[0] & hm ? 16 : 0))), m2)); + y[33] = convert_from_half(__hsub(__hmul(d2, __int2half_rn((ql[1] >> 4) + (qh[1] & hm ? 16 : 0))), m2)); +} + +template +inline __device__ void dequantize_block_q6_K(const void * __restrict__ vx, dst_t * __restrict__ yy) { + const block_q6_K * x = (const block_q6_K *) vx; + + const auto i = 0; + + // assume 64 threads - this is very slightly better than the one below + const auto tid = threadIdx.x; + const int ip = tid/32; // ip is 0 or 1 + const int il = tid - 32*ip; // 0...32 + const int is = 8*ip + il/16; + + dst_t * y = yy + i*QK_K + 128*ip + il; + + const half d = x[i].d; + + const uint8_t * ql = x[i].ql + 64*ip + il; + const uint8_t qh = x[i].qh[32*ip + il]; + const int8_t * sc = x[i].scales + is; + + y[ 0] = convert_from_half(__hmul(d, __int2half_rn(sc[0] * ((int8_t)((ql[ 0] & 0xF) | (((qh >> 0) & 3) << 4)) - 32)))); + y[32] = convert_from_half(__hmul(d, __int2half_rn(sc[2] * ((int8_t)((ql[32] & 0xF) | (((qh >> 2) & 3) << 4)) - 32)))); + y[64] = convert_from_half(__hmul(d, __int2half_rn(sc[4] * ((int8_t)((ql[ 0] >> 4) | (((qh >> 4) & 3) << 4)) - 32)))); + y[96] = convert_from_half(__hmul(d, __int2half_rn(sc[6] * ((int8_t)((ql[32] >> 4) | (((qh >> 6) & 3) << 4)) - 32)))); +} \ No newline at end of file diff --git a/torchnative/rust/vendor/candle-kernels/src/moe/moe_gguf.cu b/torchnative/rust/vendor/candle-kernels/src/moe/moe_gguf.cu new file mode 100644 index 00000000..92704e6a --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/moe/moe_gguf.cu @@ -0,0 +1,216 @@ +/** + * @brief CUDA kernel for Mixture-of-Experts (MoE) GEMM using GGUF quantized weights. + * + * This kernel performs a dot-product between quantized input tokens and + * quantized expert weight matrices, accumulating into float outputs. + * It supports per-token top-k weighting and tiling along the K dimension + * for efficient vectorized execution. + * + * Adapted from: https://github.com/guoqingbao/attention.rs/tree/main/src/kernels/src/moe_gemm_gguf.cu + */ +#include "gguf.cuh" +#include +#include +#include +#include +#include +#include +constexpr int MATRIX_ROW_PADDING = 512; + +constexpr int pad(int size, int padding) { + if (padding == 0) return size; // avoid divide-by-zero + return ((size + padding - 1) / padding) * padding; +} + +// Optional helper if you want ceil division explicitly +constexpr int ceil_div(int a, int b) { + return (a + b - 1) / b; +} + +namespace vllm_rs { + +/* +* Template Parameters: + * @tparam T Type of output elements (float, half, etc.) + * @tparam qk Quantization block size for weights (e.g., 32) + * @tparam qi Quantization block size for inputs (e.g., 32) + * @tparam block_q_t Type of quantized weight block (e.g., block_q8_0) + * @tparam vdr Vectorization factor (number of elements per lane) + * @tparam vec_dot_q_cuda Function for computing vectorized dot-product between quantized blocks + * + * Kernel Parameters: + * @param all_weights Pointer to all expert weight matrices, [num_experts, N, K] (quantized) + * @param all_inputs Pointer to all input tokens, [M_total, K] (quantized) + * @param sorted_token_ids Sorted token indices for batch processing + * @param expert_ids Expert ID for each token + * @param topk_weights Optional top-k MoE weight per token + * @param all_outputs Output buffer [M_total, N] (float) + * @param num_experts Number of experts + * @param topk Top-k experts selected per token + * @param size_m Number of tokens processed (M dimension) + * @param size_n Output feature dimension (N dimension) + * @param size_k Input feature dimension (K dimension) + * @param k_padded Padded K dimension for GGUF stride +*/ +template +__global__ void moe_gemm_gguf_kernel( + const void * __restrict__ all_weights, // [num_experts, N, K] (quantized) + const void * __restrict__ all_inputs, // [M_total, K] (quantized, M_total is total tokens) + const int32_t* __restrict__ sorted_token_ids,// [M] (M = num tokens processed) + const int32_t* __restrict__ expert_ids, // [M] + const float* __restrict__ topk_weights, // [M] + float * __restrict__ all_outputs, // [M_total, N] (float) + int num_experts, + int topk, + int size_m, int size_n, int size_k, // M, N, K are the logical dims + int k_padded // Padded K-dim for GGUF stride +) { + const int laneId = threadIdx.x; + const int wrapId = threadIdx.y; + const int nWraps = blockDim.y; + const int row = blockIdx.x * nWraps + wrapId; // This is the 'n' dimension (output row) + const int m_idx = blockIdx.y; // This is the 'm' dimension (token index) + + // This block computes the dot product for `output[token_id][n_row]` + + if (row >= size_n || m_idx >= size_m) { + return; + } + + // strides + const size_t weight_expert_stride_bytes = (size_t)(size_n * size_k) / qk * sizeof(block_q_t); + const size_t input_task_stride_bytes = (size_t)k_padded / QK8_1 * sizeof(block_q8_1); + const size_t output_task_stride_elems = (size_t)size_n; + + const int token_id = sorted_token_ids[m_idx]; // The *actual* row in input/output tensors + const int expert = expert_ids[m_idx]; + + // If expert is invalid, this token does not participate. + if (expert < 0 || expert >= num_experts) return; + + // Get the scaling factor for this token/expert pair + const float scale = (topk_weights) ? topk_weights[token_id] : 1.0f; + + const block_q_t * __restrict__ w_expert = + (const block_q_t *)((const char *)all_weights + (size_t)expert * weight_expert_stride_bytes); + + const int input_index = topk_weights ? token_id : (token_id / topk); + const block_q8_1 * __restrict__ y_ptr = + (const block_q8_1 *)((const char *)all_inputs + (size_t)input_index * input_task_stride_bytes); + + // dot-product tiling along k + const int blocks_per_row_x = size_k / qk; + const int blocks_per_iter = vdr * WARP_SIZE / qi; // no nwarps factor: one warp per batch item + + extern __shared__ int8_t shared_bytes[]; + block_q_t* w_shared_row = reinterpret_cast(shared_bytes); + for (int i = laneId; i < blocks_per_row_x; i += WARP_SIZE) { + w_shared_row[wrapId * blocks_per_row_x + i] = w_expert[row * blocks_per_row_x + i]; + } + __syncthreads(); + + // accumulators for rows_per_block rows (usually 1) + float acc = 0.0f; + + #pragma unroll + for (int kbx = laneId / (qi / vdr); kbx < blocks_per_row_x; kbx += blocks_per_iter) { + const int kby = kbx * (qk / QK8_1); + const int kqs = vdr * (laneId % (qi / vdr)); + acc += vec_dot_q_cuda( + // &w_expert[kbx + row * blocks_per_row_x], + &w_shared_row[wrapId * blocks_per_row_x + kbx], + &y_ptr[kby], + kqs); + } + + float v = warp_reduce_sum(acc) * scale; + if (laneId == 0) { + float * __restrict__ out_ptr = + all_outputs + ((size_t)token_id) * output_task_stride_elems; + out_ptr[row] = v; + } +} + +} + +#define LAUNCH_MOE_GGUF(qk, qi, block_q_t, vdr, vec_dot_q_cuda) \ + const int shared_bytes = size_k / qk * sizeof(block_q_t) * nWraps + 1024;\ + vllm_rs::moe_gemm_gguf_kernel \ + <<>>(\ + weights, y_q8_1,\ + sorted_token_ids, expert_ids, topk_weights,\ + outputs,\ + num_experts, topk,\ + size_m, size_n, size_k,\ + kx_padded\ + );\ + + +extern "C" void moe_gemm_gguf( + const float* inputs, //must be float + const void* weights, + const int32_t* sorted_token_ids, + const int32_t* expert_ids, + const float* topk_weights, + float* outputs, + int num_experts, + int topk, + int size_m, // M (num tokens to process) + int size_n, // N (output dim) + int size_k, // K (input dim) + int quant_type, // Q8_0: 0, Q4K: 1, Q2K: 2, Q3k: 3, Q5K: 4, Q6K: 5, + cudaStream_t stream +) { + const int QUANTIZE_BLOCK_SIZE = CUDA_QUANTIZE_BLOCK_SIZE; + const int kx_padded = pad(size_k, MATRIX_ROW_PADDING); + const int num_blocks = ceil_div(kx_padded, QUANTIZE_BLOCK_SIZE); + int m = topk_weights ? size_m : size_m / topk; + dim3 grid_dim_quant(num_blocks, m, 1); + dim3 block_dim_quant(QUANTIZE_BLOCK_SIZE, 1, 1); + int y_size_in_bytes = + m * (kx_padded / QK8_1 * sizeof(block_q8_1)); + void* y_q8_1 = nullptr; + cudaMallocAsync(&y_q8_1, y_size_in_bytes, stream); + quantize_q8_1<<>>(inputs, y_q8_1, size_k, kx_padded); + + const int nWraps = 4; + dim3 grid_dim(ceil_div(size_n, nWraps), size_m, 1); + dim3 block_dim(WARP_SIZE, nWraps, 1); + + //Q8_0: 0, Q4K: 1, Q2K: 2, Q3k: 3, Q5K: 4, Q6K: 5, + switch (quant_type) { + case 0: // Q8_0 + { + LAUNCH_MOE_GGUF(QK8_0, QI8_0, block_q8_0, VDR_Q8_0_Q8_1_MMVQ, vec_dot_q8_0_q8_1); + break; + } + case 1: // Q4K + { + LAUNCH_MOE_GGUF(QK_K, QI4_K, block_q4_K, VDR_Q4_K_Q8_1_MMVQ, vec_dot_q4_K_q8_1); + break; + } + case 2: // Q2_K + { + LAUNCH_MOE_GGUF(QK_K, QI2_K, block_q2_K, VDR_Q2_K_Q8_1_MMVQ, vec_dot_q2_K_q8_1); + break; + } + case 3: // Q3_K + { + LAUNCH_MOE_GGUF(QK_K, QI3_K, block_q3_K, VDR_Q3_K_Q8_1_MMVQ, vec_dot_q3_K_q8_1); + break; + } + case 4: // Q5_K + { + LAUNCH_MOE_GGUF(QK_K, QI5_K, block_q5_K, VDR_Q5_K_Q8_1_MMVQ, vec_dot_q5_K_q8_1); + break; + } + case 5: // Q6K + { + LAUNCH_MOE_GGUF(QK_K, QI6_K, block_q6_K, VDR_Q6_K_Q8_1_MMVQ, vec_dot_q6_K_q8_1); + break; + } + default: + break; + } + cudaFreeAsync(y_q8_1, stream); +} \ No newline at end of file diff --git a/torchnative/rust/vendor/candle-kernels/src/moe/moe_utils.cuh b/torchnative/rust/vendor/candle-kernels/src/moe/moe_utils.cuh new file mode 100644 index 00000000..59643408 --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/moe/moe_utils.cuh @@ -0,0 +1,188 @@ +#undef __CUDA_FP8_TYPES_EXIST__ +#include +#include +#include +#include +#include + +/** + * @brief Counts the number of tokens assigned to each expert. + * + * @param expert_ids Device pointer to the sorted expert IDs [size_m]. + * @param expert_counts Device pointer to the output counts [num_experts] + * (must be pre-initialized to zero). + * @param size_m Total number of tokens. + */ +static __global__ void count_tokens_per_expert_kernel( + const int32_t* expert_ids, + int32_t* expert_counts, + int size_m) +{ + int i = blockIdx.x * blockDim.x + threadIdx.x; + if (i < size_m) { + int32_t expert_id = expert_ids[i]; + // expert_id is from a sorted list, so we assume it's valid + // (i.e., 0 <= expert_id < num_experts) + atomicAdd(&expert_counts[expert_id], 1); + } +} + +/** + * @brief Calculates expert offsets array on the GPU. + * + * @param d_expert_ids Device pointer to sorted expert IDs [size_m]. + * @param size_m Total number of tokens. + * @param d_expert_offsets Device pointer for output offsets [num_experts + 1]. + * @param num_experts Number of experts. + * @param stream CUDA stream. + */ +static void calculate_expert_offsets( + const int32_t* d_expert_ids, + int size_m, + int32_t* d_expert_counts, + int32_t* d_expert_offsets, + int num_experts, + cudaStream_t stream +) { + // 1. Zero-initialize the counts buffer + cudaMemsetAsync(d_expert_counts, 0, num_experts * sizeof(int32_t), stream); + + // 2. Launch kernel to count tokens per expert + int threads = 256; + int blocks = (size_m + threads - 1) / threads; + count_tokens_per_expert_kernel<<>>( + d_expert_ids, d_expert_counts, size_m + ); + + // 3. Perform prefix sum (scan) + // We will use inclusive_scan on [counts] and store results in [offsets + 1] + // This is a common and efficient pattern. + + // Wrap raw pointers for Thrust + thrust::device_ptr d_counts_ptr(d_expert_counts); + thrust::device_ptr d_offsets_ptr(d_expert_offsets); + + // Run inclusive scan. + // Input: [c0, c1, c2, ...] (size num_experts) + // Output: [c0, c0+c1, c0+c1+c2, ...] (stored at offsets[1]) + thrust::inclusive_scan( + thrust::cuda::par.on(stream), // Execute on the specified stream + d_counts_ptr, // Input start + d_counts_ptr + num_experts, // Input end + d_offsets_ptr + 1 // Output start (shifted by 1) + ); + + // 4. Set the first offset (offsets[0]) to 0 + // This completes the exclusive scan. + cudaMemsetAsync(d_expert_offsets, 0, sizeof(int32_t), stream); +} + + +// This performs an EXCLUSIVE scan: [c0, c1] -> [0, c0, c0+c1] +// Assumptions: num_experts <= 1024 (fits in one block) +static __global__ void expert_prefix_sum_kernel( + const int32_t* __restrict__ counts, + int32_t* __restrict__ offsets, + int num_experts +) { + // Use shared memory for fast scanning + // Size needs to be enough for num_experts + extern __shared__ int32_t temp_storage[]; + + int tid = threadIdx.x; + + // We pad with 0 if tid >= num_experts + int val = (tid < num_experts) ? counts[tid] : 0; + temp_storage[tid] = val; + + __syncthreads(); + + // Hillis-Steele Parallel Scan (Inclusive in shared mem) + for (int offset = 1; offset < blockDim.x; offset <<= 1) { + int temp_val = 0; + if (tid >= offset) { + temp_val = temp_storage[tid - offset]; + } + __syncthreads(); + if (tid >= offset) { + temp_storage[tid] += temp_val; + } + __syncthreads(); + } + + // The result at temp_storage[i] is the inclusive sum of counts[0..i] + // We want offsets[i] = inclusive_sum[i-1] + // We want offsets[0] = 0 + + if (tid < num_experts) { + // Shift right: Offset[i+1] gets the inclusive sum up to i + offsets[tid + 1] = temp_storage[tid]; + + // Handle the first element separately + if (tid == 0) { + offsets[0] = 0; + } + } +} + +static void calculate_expert_offsets_light( + const int32_t* d_expert_ids, + int size_m, + int32_t* d_expert_counts, + int32_t* d_expert_offsets, + int num_experts, + cudaStream_t stream +) { + cudaMemsetAsync(d_expert_counts, 0, num_experts * sizeof(int32_t), stream); + + int threads = 256; + int blocks = (size_m + threads - 1) / threads; + count_tokens_per_expert_kernel<<>>( + d_expert_ids, d_expert_counts, size_m + ); + + // We launch exactly one block with 'num_experts' threads (or next power of 2) + // We need shared memory size = threads * sizeof(int32_t) + int scan_threads = num_experts; + + // Round up scan_threads to next power of 2 if needed, + // or just use a fixed size like 1024 if num_experts is small enough. + if (scan_threads < 32) scan_threads = 32; + else if (scan_threads > 1024) { + // Error: This custom kernel only supports up to 1024 experts + // Handle error or assert here + } + + size_t smem_size = scan_threads * sizeof(int32_t); + + expert_prefix_sum_kernel<<<1, scan_threads, smem_size, stream>>>( + d_expert_counts, + d_expert_offsets, + num_experts + ); +} + +namespace vllm_rs { + +inline __device__ uint16_t float_to_half(float f) { + union { + uint32_t u32; + uint16_t u16[2]; + } tmp; +#ifndef USE_ROCM + asm volatile("cvt.rn.f16.f32 %0, %1;\n" : "=h"(tmp.u16[0]) : "f"(f)); +#else + asm volatile("v_cvt_f16_f32 %0, %1;\n" : "=v"(tmp.u32) : "v"(f)); +#endif + return tmp.u16[0]; +} + +inline __device__ void from_float(half& dst, float src) { + dst = static_cast(float_to_half(src)); +} + +inline __device__ void from_float(__nv_bfloat16& dst, float src) { + dst = __float2bfloat16(src); +} + +} \ No newline at end of file diff --git a/torchnative/rust/vendor/candle-kernels/src/moe/moe_wmma.cu b/torchnative/rust/vendor/candle-kernels/src/moe/moe_wmma.cu new file mode 100644 index 00000000..b95d77d7 --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/moe/moe_wmma.cu @@ -0,0 +1,287 @@ +/** + * @brief WMMA-based grouped MoE GEMM kernel. + * + * Each block computes a tile of the output corresponding to: + * - One expert segment (group of tokens routed to the same expert) + * - One N-dimension tile (a sub-block of the expert's output features) + * + * The kernel loads input activations and expert weights in tiles using shared memory, + * performs matrix multiplication using Tensor Cores (WMMA), and accumulates results + * into a shared C tile. The final results are written atomically into the global + * output buffer to support multi-expert (top-k > 1) routing where tokens appear in + * multiple experts’ outputs. + * + * Adapted from https://github.com/guoqingbao/attention.rs/tree/main/src/kernels/src/moe_gemm_wmma.cu + */ + +#include +#include +#include +#include +#include +#include +#include +#include +#include "moe_utils.cuh" +using namespace nvcuda::wmma; + +namespace vllm_rs { + +#define CEILDIV(x,y) (((x) + (y) - 1) / (y)) + +constexpr int WMMA_K = 16; +using VecT = float4; + +// Vectorized load size (float4 = 128 bits = 8 half/bfloat16 values) +constexpr int VEC_SIZE = 8; +constexpr int NUM_VECS = 32; + +// We use 4 Warps (128 threads) per block +constexpr int WARPS_PER_BLOCK = 4; // 4 warps +constexpr int BLOCK_THREADS = 128; // 128 threads + +constexpr int M_BLK = 32; +constexpr int N_BLK = 32; +constexpr int K_BLK = WMMA_K; // 16 + + +/** + * @brief WMMA-based grouped MoE GEMM kernel. + * + * @tparam T Data type: half or nv_bfloat16 + * + * @param input [size_m or size_m/topk, size_k] + * @param weights [num_experts, size_n, size_k] compacted expert weights + * @param sorted_token_ids [size_m] mapping of per-token row indices (sorted by expert) + * @param expert_offsets [num_experts] array of {start, len} tokens indices for each expert + * @param topk_weights [size_m] optional per-token scaling weights (nullptr if unused) + * @param output [size_m, size_n] global output buffer (must be zero-initialized) + * @param num_experts Total number of experts + * @param topk Number of experts each token is routed to + * @param size_m Number of tokens + * @param size_n Output hidden dimension (per expert) + * @param size_k Input hidden dimension +*/ +template +__global__ void moe_gemm_grouped_kernel( + const T* __restrict__ input, // [size_m, size_k] + const T* __restrict__ weights, // [num_experts, size_n, size_k] + const int32_t* __restrict__ sorted_token_ids, // [size_m] + const int32_t* __restrict__ expert_offsets, // [num_experts] + const float* __restrict__ topk_weights, // [size_m] + T* __restrict__ output, // [size_m, size_n] (Zero-initialized) + const int num_experts, const int topk, + const int32_t size_m, + const int32_t size_n, + const int32_t size_k +) { + // Get Segment and N-Tile for this Block + const int expert_id = blockIdx.x; + const int n_tile_idx = blockIdx.y; + if (expert_id < 0 || expert_id >= num_experts) return; + const int segment_start = expert_offsets[expert_id]; + const int segment_end = expert_offsets[expert_id + 1]; + const int num_rows_in_segment = segment_end - segment_start; + + if (num_rows_in_segment == 0) return; + + const int n_base = n_tile_idx * N_BLK; + if (n_base >= size_n) return; + + const T* expert_w = weights + (size_t)expert_id * (size_t)size_n * (size_t)size_k; + + extern __shared__ uint8_t smem_bytes[]; + + // A tile: [M_BLK, K_BLK] (row-major) + T* A_sh = reinterpret_cast(smem_bytes); + // B tile: [N_BLK, K_BLK] (row-major) + T* B_sh = reinterpret_cast(A_sh + M_BLK * K_BLK); + uint8_t* C_ptr = reinterpret_cast(B_sh + N_BLK * K_BLK); + + // align next pointer to float alignment + size_t offset = reinterpret_cast(C_ptr) % alignof(float); + if (offset != 0) { + C_ptr += (alignof(float) - offset); + } + float* C_sh = reinterpret_cast(C_ptr); // shared scratch for final per-block tile writes + + const int threadId = threadIdx.x; + const int warpId = threadId / 32; + const int laneId = threadId % 32; + const int warp_m_idx = warpId / WARPS_N; + const int warp_n_idx = warpId % WARPS_N; + + const int B_ELEMS_PER_BLOCK = N_BLK * K_BLK; + const int VEC_ELEMS_B = B_ELEMS_PER_BLOCK / VEC_SIZE; // 512 / 8 = 64 + const int A_ELEMS_PER_BLOCK = M_BLK * K_BLK; + const int VEC_ELEMS_A = A_ELEMS_PER_BLOCK / VEC_SIZE; // 512 / 8 = 64 + VecT zero_vec; + zero_vec.x = zero_vec.y = zero_vec.z = zero_vec.w = 0.0f; + + for (int m_base = 0; m_base < num_rows_in_segment; m_base += M_BLK) { + // We'll accumulate full-K results in per-warp fragments (initialized here) + fragment c_frag; + fill_fragment(c_frag, 0.0f); + + // For every k_block we will load B_sh and A_sh for this m_base subsequently + for (int k_base = 0; k_base < size_k; k_base += K_BLK) { + // Load B Tile (Weights) into B_sh + for (int i = threadId; i < VEC_ELEMS_B; i += BLOCK_THREADS) { + int idx = i * VEC_SIZE; // element index (0..511) + int n_local = idx / K_BLK; + int k_local = idx % K_BLK; + + int n_global = n_base + n_local; + int k_global = k_base + k_local; + + // this should be always satisfied since k dim aligned to 8 + if (n_global < size_n && k_global < size_k) { + *reinterpret_cast(&B_sh[n_local * K_BLK + k_local]) = *reinterpret_cast( + &expert_w[(size_t)n_global * size_k + k_global] + ); + } else { + *reinterpret_cast(&B_sh[n_local * K_BLK + k_local]) = zero_vec; + } + } + + // Load A Tile (Inputs) into A_sh for this m_base and this k_base + for (int i = threadId; i < VEC_ELEMS_A; i += BLOCK_THREADS) { + int idx = i * VEC_SIZE; // element index + int m_local = idx / K_BLK; + int k_local = idx % K_BLK; + + int m_seg = m_base + m_local; // row index within segment + int k_global = k_base + k_local; + + if (m_seg < num_rows_in_segment && k_global < size_k) { + int token_pair_index = segment_start + m_seg; + int token_index = sorted_token_ids[token_pair_index]; + int input_index = token_index / (topk_weights? 1: topk); + *reinterpret_cast(&A_sh[m_local * K_BLK + k_local]) = *reinterpret_cast( + &input[(size_t)input_index * size_k + k_global] + ); + } else { + // in case m dim in this segment not aligned to 8 + *reinterpret_cast(&A_sh[m_local * K_BLK + k_local]) = zero_vec; + } + } + + __syncthreads(); + + // Compute (Warp-level) : update c_frag for this k_block + fragment a_frag; + fragment b_frag; + + // Point this warp to its tile in shared memory + const T* A_sh_ptr = A_sh + (warp_m_idx * WMMA_M * K_BLK); + const T* B_sh_ptr = B_sh + (warp_n_idx * WMMA_N * K_BLK); + + load_matrix_sync(a_frag, A_sh_ptr, K_BLK); + load_matrix_sync(b_frag, B_sh_ptr, K_BLK); + + // Accumulate into c_frag (which persists across k_base iterations) + mma_sync(c_frag, a_frag, b_frag, c_frag); + __syncthreads(); // Fix shared memory mismatch on V100 + } // end k_base loop (we have a fully-accumulated c_frag for this m_base tile) + + // Store the accumulated c_frag to C_sh (shared) once per warp + // Point this warp to its 16x16 tile *within* the 32x32 C_sh + float* C_sh_ptr = C_sh + (warp_m_idx * WMMA_M * N_BLK) + (warp_n_idx * WMMA_N); + // store the full accumulated 16x16 tile (note ld = N_BLK, result in row-major in C_sh) + store_matrix_sync(C_sh_ptr, c_frag, N_BLK, mem_row_major); + + __syncthreads(); + + // Cooperative Store from C_sh to Global + // 128 threads write [M_BLK, N_BLK] = [32, 32] = 1024 elements + const int C_ELEMS_PER_BLOCK = M_BLK * N_BLK; + for (int i = threadId; i < C_ELEMS_PER_BLOCK; i += BLOCK_THREADS) { + int m_local_c = i / N_BLK; // row in C_sh (0..31) + int n_local_c = i % N_BLK; // col in C_sh (0..31) + + int m_seg = m_base + m_local_c; // row index within segment + int n_global = n_base + n_local_c; // col index in output + + if (m_seg < num_rows_in_segment && n_global < size_n) { + int token_pair_index = segment_start + m_seg; + if (token_pair_index < size_m) { + int token_index = sorted_token_ids[token_pair_index]; + float val = C_sh[m_local_c * N_BLK + n_local_c]; + if (topk_weights) { + val *= topk_weights[token_index]; + } + from_float(output[(size_t)token_index * size_n + n_global], val); + } + } + } + } // end m_base loop +} + +} + +#define LAUNCH_MOE_WMMA(DTYPE, WMMA_M, WMMA_N, WARPS_N)\ + vllm_rs::moe_gemm_grouped_kernel<<>>(\ + reinterpret_cast(input),\ + reinterpret_cast(weights),\ + sorted_token_ids,\ + expert_offsets,\ + topk_weights,\ + reinterpret_cast(output),\ + num_experts, topk,\ + size_m, size_n, size_k \ + );\ + +extern "C" void moe_gemm_wmma( + const void* input, // [size_m, size_k] + const void* weights, // [num_experts, size_n, size_k] + const int32_t* sorted_token_ids, // [size_m] (Device) + const int32_t* expert_ids, // [size_m * topk] + const float* topk_weights, // [size_m] (Device, can be nullptr) + void* output, // [size_m, size_n] + int32_t* expert_counts, // prealloc [num_experts] + int32_t* expert_offsets, // prealloc [num_experts + 1] + int num_experts, + int topk, + int size_m, + int size_n, + int size_k, + int data_type, // 0 = half, 1 = bfloat16 + bool is_prefill, + cudaStream_t stream +) { + if (is_prefill) { + calculate_expert_offsets(expert_ids, size_m, expert_counts, expert_offsets, num_experts, stream); + } else { + calculate_expert_offsets_light(expert_ids, size_m, expert_counts, expert_offsets, num_experts, stream); + } + + int grid_n = CEILDIV(size_n, vllm_rs::N_BLK); + dim3 grid(num_experts, grid_n, 1); + dim3 block(vllm_rs::BLOCK_THREADS, 1, 1); + + // Shared memory: A_sh[M_BLK, K_BLK] + B_sh[N_BLK, K_BLK] + size_t A_sh_bytes = vllm_rs::M_BLK * vllm_rs::K_BLK * 2; // (32*16 * 2) = 1024 + size_t B_sh_bytes = vllm_rs::N_BLK * vllm_rs::K_BLK * 2; // (32*16 * 2) = 1024 + size_t C_sh_bytes = vllm_rs::M_BLK * vllm_rs::N_BLK * sizeof(float); + size_t AB_bytes = A_sh_bytes + B_sh_bytes; + size_t pad = (16 - (AB_bytes % 16)) % 16; + size_t smem_bytes = AB_bytes + pad + C_sh_bytes; // ~6KB total needed + + if (data_type == 0) { // half + if (is_prefill) { + LAUNCH_MOE_WMMA(half, 16, 16, 2) + } else { + // we use smaller M_tile and larger N_tile for decoding + LAUNCH_MOE_WMMA(half, 8, 32, 1) + } + } +#ifndef NO_BF16_KERNEL + else if (data_type == 1) { // bfloat16 + if (is_prefill) { + LAUNCH_MOE_WMMA(nv_bfloat16, 16, 16, 2) + } else { + LAUNCH_MOE_WMMA(nv_bfloat16, 8, 32, 1) + } + } +#endif +} \ No newline at end of file diff --git a/torchnative/rust/vendor/candle-kernels/src/moe/moe_wmma_gguf.cu b/torchnative/rust/vendor/candle-kernels/src/moe/moe_wmma_gguf.cu new file mode 100644 index 00000000..0d3701ee --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/moe/moe_wmma_gguf.cu @@ -0,0 +1,422 @@ +/** + * @brief CUDA kernel for Mixture-of-Experts (MoE) GEMM with GGUF quantized weights and Tensor Core. + * + * This kernel performs batched GEMM where the weight matrix is stored in GGUF + * quantized format (uint8_t blocks). It supports top-k expert selection and + * segmented expert layouts. Uses shared memory tiles and WMMA (tensor cores) + * for efficient computation. + * + * Adapted from: https://github.com/guoqingbao/attention.rs/tree/main/src/kernels/src/moe_wmma_gguf.cu + */ +#include "gguf.cuh" +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include "moe_utils.cuh" +using namespace nvcuda::wmma; + +// Constants from original kernel +constexpr int WMMA_M = 16; +constexpr int WMMA_N = 16; +constexpr int WMMA_K = 16; // This is fixed by the hardware instruction +using VecT = float4; + +constexpr int VEC_SIZE = 8; +constexpr int WARPS_M = 2; +constexpr int WARPS_N = 2; +constexpr int WARPS_PER_BLOCK = WARPS_M * WARPS_N; // 4 warps + +constexpr int M_BLK = WARPS_M * WMMA_M; // 32 +constexpr int N_BLK = WARPS_N * WMMA_N; // 32 + +// Helper for ceiling division +#define CEILDIV(A, B) (((A) + (B)-1) / (B)) + +// --- GGUF Dequantization Function (Warp-level) --- +/** + * @brief Dequantizes a single GGUF block using one warp (32 threads). + * + * @tparam T Output type (half or nv_bfloat16) + * @param dequant_out Pointer to output in shared mem [qk] + * @param quant_in Pointer to input GGUF block in shared mem + * @param type GGUF type + * @param qk Quantization group size (32 or 256) + * @param laneId threadIdx.x % 32 + */ +template +__forceinline__ __device__ void dequantize_block_warp( + T* dequant_out, + const uint8_t* quant_in, + int gguf_dtype //Q8_0: 0, Q4K: 1, Q2K: 2, Q3k: 3, Q5K: 4, Q6K: 5, +) { + using namespace nvcuda; + switch (gguf_dtype) { + case 0: { // qk = 32, q8_0 + // Block: half d (2B), int8_t qs[32] (32B) + int laneId = threadIdx.x; + const half* d_ptr = (const half*)quant_in; + const int8_t* qs = (const int8_t*)(quant_in + 2); + + // Lane 0 loads scale and broadcasts to all other lanes + half d_val = (laneId == 0) ? *d_ptr : (half)0.0f; + d_val = __shfl_sync(0xFFFFFFFF, d_val, 0); + float d_f = __half2float(d_val); + + // 32 lanes dequantize 32 values + if (laneId < QK8_0) { // qk should be 32 + dequant_out[laneId] = T( (float)qs[laneId] * d_f ); + } + break; + } + case 1: { // q4k, 32 lanes + dequantize_block_q4_K(quant_in, dequant_out); + break; + } + case 2: { // q2k, 64 lanes + dequantize_block_q2_K(quant_in, dequant_out); + break; + } + case 3: { // q3k, 64 lanes + dequantize_block_q3_K(quant_in, dequant_out); + break; + } + case 4: { // q5k, 64 lanes + dequantize_block_q5_K(quant_in, dequant_out); + break; + } + case 5: { // q6k, 64 lanes + dequantize_block_q6_K(quant_in, dequant_out); + break; + } + default: + break; + } +} + +/* +* Template Parameters: + * @tparam T Type of input/output (float, half, etc.) + * @tparam qk Quantization block size (e.g., 32) + * @tparam block_q_t Type representing a single GGUF block (e.g., block_q8_0) + * @tparam wrap_size Warp size used for thread tiling (usually 32) + * + * Kernel Parameters: + * @param input Input matrix [size_m, size_k] + * @param weights GGUF quantized weights buffer (uint8_t blocks) + * @param sorted_token_ids Array of sorted token indices for MoE routing + * @param expert_offsets [num_experts] array of {start, len} tokens indices for each expert + * @param topk_weights Top-k MoE weights per token (optional) + * @param output Output matrix [size_m, size_n] + * @param num_experts Number of experts in the MoE + * @param topk Number of top experts selected per token + * @param size_m Number of input rows / tokens + * @param size_n Output feature dimension + * @param size_k Input feature dimension + * @param gguf_dtype GGUF quantization type ID (e.g., Q8_0) +*/ +template +__global__ void moe_gemm_gguf_prefill_kernel( + const T* __restrict__ input, + const uint8_t* __restrict__ weights, // Now uint8_t* + const int32_t* __restrict__ sorted_token_ids, + const int32_t* __restrict__ expert_offsets, + const float* __restrict__ topk_weights, + float* __restrict__ output, + const int num_experts, const int topk, + const int32_t size_m, + const int32_t size_n, + const int32_t size_k, + const int gguf_dtype +) { + const int expert_id = blockIdx.x; + const int n_tile_idx = blockIdx.y; + + if (expert_id < 0 || expert_id >= num_experts) return; + const int segment_start = expert_offsets[expert_id]; + const int segment_end = expert_offsets[expert_id + 1]; + const int num_rows_in_segment = segment_end - segment_start; + + if (num_rows_in_segment == 0) return; + constexpr int BLOCK_THREADS = WARPS_PER_BLOCK * wrap_size; // 128 threads + + const int n_base = n_tile_idx * N_BLK; + if (n_base >= size_n) return; + + const size_t block_size_bytes = sizeof(block_q_t); + const size_t expert_w_row_stride_bytes = (size_k / qk) * block_size_bytes; + const uint8_t* expert_w = weights + (size_t)expert_id * size_n * expert_w_row_stride_bytes; + + extern __shared__ uint8_t smem_bytes[]; + + // 1. A tile: [M_BLK, qk] (dequantized) + T* A_sh = reinterpret_cast(smem_bytes); + size_t A_sh_bytes = (size_t)M_BLK * qk * sizeof(T); + + // 2. B tile: [N_BLK, qk] (dequantized) + uint8_t* B_sh_ptr = smem_bytes + A_sh_bytes; + size_t B_sh_bytes = (size_t)N_BLK * qk * sizeof(T); + + // 3. B quantized tile: [N_BLK * block_size_bytes] (raw GGUF) + uint8_t* B_quant_sh_ptr = B_sh_ptr + B_sh_bytes; + size_t B_quant_sh_bytes = (size_t)N_BLK * block_size_bytes; + + // 4. C tile: [M_BLK, N_BLK] (float accumulator) + uint8_t* C_sh_ptr = B_quant_sh_ptr + B_quant_sh_bytes; + size_t C_sh_offset = reinterpret_cast(C_sh_ptr) % alignof(float); + if (C_sh_offset != 0) C_sh_ptr += (alignof(float) - C_sh_offset); + + // Final aligned shared memory pointers + T* B_sh = reinterpret_cast(B_sh_ptr); + uint8_t* B_quant_sh = reinterpret_cast(B_quant_sh_ptr); + float* C_sh = reinterpret_cast(C_sh_ptr); + + const int laneId = threadIdx.x; + const int warpId = threadIdx.y; + const int threadId = warpId * wrap_size + laneId; + const int warp_m_idx = warpId / WARPS_N; + const int warp_n_idx = warpId % WARPS_N; + + const size_t A_ELEMS_PER_BLOCK = (size_t)M_BLK * qk; + const size_t VEC_ELEMS_A = A_ELEMS_PER_BLOCK / VEC_SIZE; + VecT zero_vec; + zero_vec.x = zero_vec.y = zero_vec.z = zero_vec.w = 0.0f; + + for (int m_base = 0; m_base < num_rows_in_segment; m_base += M_BLK) { + + // Per-warp accumulator fragment + fragment c_frag; + fill_fragment(c_frag, 0.0f); + + // K-Loop: Strides by GGUF block size `qk` + for (int k_base = 0; k_base < size_k; k_base += qk) { + + // Load A Tile (Inputs) into A_sh + #pragma unroll + for (size_t i = threadId; i < VEC_ELEMS_A; i += BLOCK_THREADS) { + size_t idx = i * VEC_SIZE; // element index + size_t m_local = idx / qk; + size_t k_local = idx % qk; + + int m_seg = m_base + m_local; + int k_global = k_base + k_local; + + if (m_seg < num_rows_in_segment && k_global < size_k) { + int token_pair_index = segment_start + m_seg; + int token_index = sorted_token_ids[token_pair_index]; + int input_index = token_index / (topk_weights? 1: topk); + *reinterpret_cast(&A_sh[m_local * qk + k_local]) = *reinterpret_cast( + &input[(size_t)input_index * size_k + k_global] + ); + } else { + *reinterpret_cast(&A_sh[m_local * qk + k_local]) = zero_vec; + } + } + + // Load B Tile (Quantized) into B_quant_sh + const size_t k_base_offset_bytes = (k_base / qk) * block_size_bytes; + constexpr int ROWS_PER_WARP = N_BLK / WARPS_PER_BLOCK; + + #pragma unroll + for (int row = 0; row < ROWS_PER_WARP; ++row) { + int n_local = warpId * ROWS_PER_WARP + row; + int n_global = n_base + n_local; + if (n_local < N_BLK && n_global < size_n) { + block_q_t* dest_ptr = reinterpret_cast(B_quant_sh + n_local * block_size_bytes); + const block_q_t* src_ptr = reinterpret_cast(expert_w + (size_t)n_global * expert_w_row_stride_bytes + k_base_offset_bytes); + *dest_ptr = *src_ptr; + } + } + + __syncthreads(); + + // Dequantize B from B_quant_sh to B_sh + #pragma unroll + for (int row = 0; row < ROWS_PER_WARP; ++row) { + int n_local = warpId * ROWS_PER_WARP + row; + int n_global = n_base + n_local; + if (n_local < N_BLK && n_global < size_n) { + const uint8_t* quant_ptr = B_quant_sh + n_local * block_size_bytes; + T* dequant_ptr = B_sh + n_local * qk; // Stride by qk + // Dequantize one block using this warp + dequantize_block_warp(dequant_ptr, quant_ptr, gguf_dtype); + } + } + + __syncthreads(); + + // Inner WMMA Loop + // A_sh and B_sh are now dequantized and in shared mem + // We loop over the K-dim (now `qk`) using the hardware `WMMA_K` + #pragma unroll + for (int k_tile = 0; k_tile < qk; k_tile += WMMA_K) { + fragment a_frag; + fragment b_frag; + + // Point to the correct 16x16 tile inside the [M_BLK, qk] / [N_BLK, qk] buffers + const T* A_sh_ptr = A_sh + (warp_m_idx * WMMA_M * qk) + k_tile; + const T* B_sh_ptr = B_sh + (warp_n_idx * WMMA_N * qk) + k_tile; + + load_matrix_sync(a_frag, A_sh_ptr, qk); // Stride is qk + load_matrix_sync(b_frag, B_sh_ptr, qk); // Stride is qk + + mma_sync(c_frag, a_frag, b_frag, c_frag); + } + } // end k_base loop + + // Store C_frag to C_sh + float* C_sh_ptr_warp = C_sh + (warp_m_idx * WMMA_M * N_BLK) + (warp_n_idx * WMMA_N); + store_matrix_sync(C_sh_ptr_warp, c_frag, N_BLK, mem_row_major); + __syncthreads(); + + // Cooperative Store to Global + const int C_ELEMS_PER_BLOCK = M_BLK * N_BLK; + #pragma unroll + for (int i = threadId; i < C_ELEMS_PER_BLOCK; i += BLOCK_THREADS) { + int m_local_c = i / N_BLK; + int n_local_c = i % N_BLK; + int m_seg = m_base + m_local_c; + int n_global = n_base + n_local_c; + + if (m_seg < num_rows_in_segment && n_global < size_n) { + int token_pair_index = segment_start + m_seg; + if (token_pair_index < size_m) { + int token_index = sorted_token_ids[token_pair_index]; + float val = C_sh[m_local_c * N_BLK + n_local_c]; + if (topk_weights) { + val *= topk_weights[token_index]; + } + output[(size_t)token_index * size_n + n_global] = val; + } + } + } + } // end m_base loop +} + +#define LAUNCH_MOE_GGUF_PREFILL(DTYPE) \ + if (gguf_type == 0) {\ + dim3 block(32, WARPS_PER_BLOCK, 1);\ + moe_gemm_gguf_prefill_kernel<<>>(\ + reinterpret_cast(input),\ + reinterpret_cast(weights),\ + sorted_token_ids, expert_offsets, topk_weights,\ + output, num_experts, topk, size_m, size_n, size_k, gguf_type\ + );\ + } else if (gguf_type == 1) {\ + dim3 block(32, WARPS_PER_BLOCK, 1);\ + moe_gemm_gguf_prefill_kernel<<>>(\ + reinterpret_cast(input),\ + reinterpret_cast(weights),\ + sorted_token_ids, expert_offsets, topk_weights,\ + output, num_experts, topk, size_m, size_n, size_k, gguf_type\ + );\ + } else if (gguf_type == 2) {\ + dim3 block(64, WARPS_PER_BLOCK, 1);\ + moe_gemm_gguf_prefill_kernel<<>>(\ + reinterpret_cast(input),\ + reinterpret_cast(weights),\ + sorted_token_ids, expert_offsets, topk_weights,\ + output, num_experts, topk, size_m, size_n, size_k, gguf_type\ + );\ + } else if (gguf_type == 3) {\ + dim3 block(64, WARPS_PER_BLOCK, 1);\ + moe_gemm_gguf_prefill_kernel<<>>(\ + reinterpret_cast(input),\ + reinterpret_cast(weights),\ + sorted_token_ids, expert_offsets, topk_weights,\ + output, num_experts, topk, size_m, size_n, size_k, gguf_type\ + );\ + } else if (gguf_type == 4) { \ + dim3 block(64, WARPS_PER_BLOCK, 1);\ + moe_gemm_gguf_prefill_kernel<<>>(\ + reinterpret_cast(input),\ + reinterpret_cast(weights),\ + sorted_token_ids, expert_offsets, topk_weights,\ + output, num_experts, topk, size_m, size_n, size_k, gguf_type\ + );\ + } else if (gguf_type == 5) { \ + dim3 block(64, WARPS_PER_BLOCK, 1);\ + moe_gemm_gguf_prefill_kernel<<>>(\ + reinterpret_cast(input),\ + reinterpret_cast(weights),\ + sorted_token_ids, expert_offsets, topk_weights,\ + output, num_experts, topk, size_m, size_n, size_k, gguf_type\ + );\ + } + + +extern "C" void moe_gemm_gguf_prefill( + const void* input, + const uint8_t* weights, + const int32_t* sorted_token_ids, + const int32_t* expert_ids, + const float* topk_weights, + float* output, + int num_experts, + int topk, + int size_m, + int size_n, + int size_k, + int input_dtype, // 0 = half, 1 = bfloat16 + int gguf_type, //Q8_0: 0, Q4K: 1, Q2K: 2, Q3k: 3, Q5K: 4, Q6K: 5, + cudaStream_t stream +) { + int32_t* expert_counts; + cudaMallocAsync(&expert_counts, num_experts * sizeof(int32_t), stream); + + int32_t* expert_offsets; + cudaMallocAsync(&expert_offsets, (num_experts + 1) * sizeof(int32_t), stream); + calculate_expert_offsets(expert_ids, size_m, expert_counts, expert_offsets, num_experts, stream); + + int grid_n = CEILDIV(size_n, N_BLK); + dim3 grid(num_experts, grid_n, 1); + + size_t qk = QK_K; + size_t block_size_bytes = sizeof(block_q6_K); + if (gguf_type == 0) { //Q8_0: 0, + block_size_bytes = sizeof(block_q8_0); + qk = QK8_0; + } else if (gguf_type == 1) {// Q4K: 1, + block_size_bytes = sizeof(block_q4_K); + } else if (gguf_type == 2) {// Q2K: 2, + block_size_bytes = sizeof(block_q2_K); + } else if (gguf_type == 3) {//Q3K: 3, + block_size_bytes = sizeof(block_q3_K); + } else if (gguf_type == 4) {//Q5K: 4, + block_size_bytes = sizeof(block_q5_K); + } + + // 1. A tile: [M_BLK, qk] (dequantized) + size_t A_sh_bytes = (size_t)M_BLK * qk * 2; // 2 for half/bfloat16 + + // 2. B tile: [N_BLK, qk] (dequantized) + size_t B_sh_bytes = (size_t)N_BLK * qk * 2; + + // 3. B quantized tile: [N_BLK * block_size_bytes] + size_t B_quant_sh_bytes = (size_t)N_BLK * block_size_bytes; + + // 4. C tile: [M_BLK, N_BLK] (float accumulator) + size_t C_sh_bytes = (size_t)M_BLK * N_BLK * sizeof(float); + + // Add up, with padding for C + size_t smem_bytes = A_sh_bytes + B_sh_bytes + B_quant_sh_bytes; + size_t C_sh_offset = smem_bytes % alignof(float); + if (C_sh_offset != 0) smem_bytes += (alignof(float) - C_sh_offset); + smem_bytes += C_sh_bytes; + + if (input_dtype == 0) { + LAUNCH_MOE_GGUF_PREFILL(half); + } else { +#ifndef NO_BF16_KERNEL + LAUNCH_MOE_GGUF_PREFILL(nv_bfloat16); +#endif + } + cudaFreeAsync(expert_counts, stream); + cudaFreeAsync(expert_offsets, stream); +} diff --git a/torchnative/rust/vendor/candle-kernels/src/ptx.rs b/torchnative/rust/vendor/candle-kernels/src/ptx.rs new file mode 100644 index 00000000..1c73d6b7 --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/ptx.rs @@ -0,0 +1,11 @@ +pub const AFFINE: &str = include_str!(concat!(env!("OUT_DIR"), "/affine.ptx")); +pub const BINARY: &str = include_str!(concat!(env!("OUT_DIR"), "/binary.ptx")); +pub const CAST: &str = include_str!(concat!(env!("OUT_DIR"), "/cast.ptx")); +pub const CONV: &str = include_str!(concat!(env!("OUT_DIR"), "/conv.ptx")); +pub const FILL: &str = include_str!(concat!(env!("OUT_DIR"), "/fill.ptx")); +pub const INDEXING: &str = include_str!(concat!(env!("OUT_DIR"), "/indexing.ptx")); +pub const QUANTIZED: &str = include_str!(concat!(env!("OUT_DIR"), "/quantized.ptx")); +pub const REDUCE: &str = include_str!(concat!(env!("OUT_DIR"), "/reduce.ptx")); +pub const SORT: &str = include_str!(concat!(env!("OUT_DIR"), "/sort.ptx")); +pub const TERNARY: &str = include_str!(concat!(env!("OUT_DIR"), "/ternary.ptx")); +pub const UNARY: &str = include_str!(concat!(env!("OUT_DIR"), "/unary.ptx")); diff --git a/torchnative/rust/vendor/candle-kernels/src/quantized.cu b/torchnative/rust/vendor/candle-kernels/src/quantized.cu new file mode 100644 index 00000000..a9f24900 --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/quantized.cu @@ -0,0 +1,4845 @@ +// Kernels adapted from llama.cpp ggml-cuda.cu and ggml-cuda/getrows.cu +// https://github.com/ggerganov/llama.cpp/blob/master/ggml-cuda.cu +#include "cuda_fp16.h" +#include "cuda_bf16.h" +#include + +#define GGML_UNUSED(x) (void)(x) +#define GGML_CUDA_ASSUME(x) + +#ifdef GGML_QKK_64 +#define QK_K 64 +#define K_SCALE_SIZE 4 +#else +#define QK_K 256 +#define K_SCALE_SIZE 12 +#endif + +#undef GGML_CUDA_F16 +#define GGML_CUDA_DMMV_X 32 +#define CUDA_QUANTIZE_BLOCK_SIZE 256 +#define CUDA_DEQUANTIZE_BLOCK_SIZE 256 +#define CUDA_GET_ROWS_BLOCK_SIZE 256 +#define K_QUANTS_PER_ITERATION 2 + +typedef uint16_t ggml_fp16_t; +typedef float dfloat; // dequantize float +typedef float2 dfloat2; +typedef void (*dequantize_kernel_t)(const void * vx, const int ib, const int iqs, dfloat2 & v); + +static __device__ __forceinline__ float warp_reduce_sum(float x) { +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + x += __shfl_xor_sync(0xffffffff, x, mask, 32); + } + return x; +} + +static __device__ __forceinline__ float warp_reduce_max(float x) { +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + x = fmaxf(x, __shfl_xor_sync(0xffffffff, x, mask, 32)); + } + return x; +} + +static __device__ __forceinline__ int get_int_from_int8(const int8_t * x8, const int & i32) { + const uint16_t * x16 = (const uint16_t *) (x8 + sizeof(int) * i32); // assume at least 2 byte alignment + + int x32 = 0; + x32 |= x16[0] << 0; + x32 |= x16[1] << 16; + + return x32; +} + +static __device__ __forceinline__ int get_int_from_uint8(const uint8_t * x8, const int & i32) { + const uint16_t * x16 = (const uint16_t *) (x8 + sizeof(int) * i32); // assume at least 2 byte alignment + + int x32 = 0; + x32 |= x16[0] << 0; + x32 |= x16[1] << 16; + + return x32; +} + +static __device__ __forceinline__ int get_int_from_int8_aligned(const int8_t * x8, const int & i32) { + return *((const int *) (x8 + sizeof(int) * i32)); // assume at least 4 byte alignment +} + +static __device__ __forceinline__ int get_int_from_uint8_aligned(const uint8_t * x8, const int & i32) { + return *((const int *) (x8 + sizeof(int) * i32)); // assume at least 4 byte alignment +} + + +#define WARP_SIZE 32 +#define CUDART_HMAX 11070 // CUDA 11.7, min. ver. for which __hmax and __hmax2 are known to work (may be higher than needed) + +#define CUDA_CC_PASCAL 600 +#define MIN_CC_DP4A 610 // minimum compute capability for __dp4a, an intrinsic for byte-wise dot products +#define CUDA_CC_VOLTA 700 +#define CC_OFFSET_AMD 1000000 +#define CC_RDNA1 (CC_OFFSET_AMD + 1010) +#define CC_RDNA2 (CC_OFFSET_AMD + 1030) +#define CC_RDNA3 (CC_OFFSET_AMD + 1100) + +static __device__ __forceinline__ int ggml_cuda_dp4a(const int a, const int b, int c) { +#if __CUDA_ARCH__ >= MIN_CC_DP4A + return __dp4a(a, b, c); +#else // __CUDA_ARCH__ >= MIN_CC_DP4A + const int8_t * a8 = (const int8_t *) &a; + const int8_t * b8 = (const int8_t *) &b; + return c + a8[0]*b8[0] + a8[1]*b8[1] + a8[2]*b8[2] + a8[3]*b8[3]; +#endif // __CUDA_ARCH__ >= MIN_CC_DP4A +} + + +#define MMQ_X_Q4_0_RDNA2 64 +#define MMQ_Y_Q4_0_RDNA2 128 +#define NWARPS_Q4_0_RDNA2 8 +#define MMQ_X_Q4_0_RDNA1 64 +#define MMQ_Y_Q4_0_RDNA1 64 +#define NWARPS_Q4_0_RDNA1 8 +#if defined(CUDA_USE_TENSOR_CORES) +#define MMQ_X_Q4_0_AMPERE 4 +#define MMQ_Y_Q4_0_AMPERE 32 +#define NWARPS_Q4_0_AMPERE 4 +#else +#define MMQ_X_Q4_0_AMPERE 64 +#define MMQ_Y_Q4_0_AMPERE 128 +#define NWARPS_Q4_0_AMPERE 4 +#endif +#define MMQ_X_Q4_0_PASCAL 64 +#define MMQ_Y_Q4_0_PASCAL 64 +#define NWARPS_Q4_0_PASCAL 8 + +#define MMQ_X_Q4_1_RDNA2 64 +#define MMQ_Y_Q4_1_RDNA2 128 +#define NWARPS_Q4_1_RDNA2 8 +#define MMQ_X_Q4_1_RDNA1 64 +#define MMQ_Y_Q4_1_RDNA1 64 +#define NWARPS_Q4_1_RDNA1 8 +#if defined(CUDA_USE_TENSOR_CORES) +#define MMQ_X_Q4_1_AMPERE 4 +#define MMQ_Y_Q4_1_AMPERE 32 +#define NWARPS_Q4_1_AMPERE 4 +#else +#define MMQ_X_Q4_1_AMPERE 64 +#define MMQ_Y_Q4_1_AMPERE 128 +#define NWARPS_Q4_1_AMPERE 4 +#endif +#define MMQ_X_Q4_1_PASCAL 64 +#define MMQ_Y_Q4_1_PASCAL 64 +#define NWARPS_Q4_1_PASCAL 8 + +#define MMQ_X_Q5_0_RDNA2 64 +#define MMQ_Y_Q5_0_RDNA2 128 +#define NWARPS_Q5_0_RDNA2 8 +#define MMQ_X_Q5_0_RDNA1 64 +#define MMQ_Y_Q5_0_RDNA1 64 +#define NWARPS_Q5_0_RDNA1 8 +#if defined(CUDA_USE_TENSOR_CORES) +#define MMQ_X_Q5_0_AMPERE 4 +#define MMQ_Y_Q5_0_AMPERE 32 +#define NWARPS_Q5_0_AMPERE 4 +#else +#define MMQ_X_Q5_0_AMPERE 128 +#define MMQ_Y_Q5_0_AMPERE 64 +#define NWARPS_Q5_0_AMPERE 4 +#endif +#define MMQ_X_Q5_0_PASCAL 64 +#define MMQ_Y_Q5_0_PASCAL 64 +#define NWARPS_Q5_0_PASCAL 8 + +#define MMQ_X_Q5_1_RDNA2 64 +#define MMQ_Y_Q5_1_RDNA2 128 +#define NWARPS_Q5_1_RDNA2 8 +#define MMQ_X_Q5_1_RDNA1 64 +#define MMQ_Y_Q5_1_RDNA1 64 +#define NWARPS_Q5_1_RDNA1 8 +#if defined(CUDA_USE_TENSOR_CORES) +#define MMQ_X_Q5_1_AMPERE 4 +#define MMQ_Y_Q5_1_AMPERE 32 +#define NWARPS_Q5_1_AMPERE 4 +#else +#define MMQ_X_Q5_1_AMPERE 128 +#define MMQ_Y_Q5_1_AMPERE 64 +#define NWARPS_Q5_1_AMPERE 4 +#endif +#define MMQ_X_Q5_1_PASCAL 64 +#define MMQ_Y_Q5_1_PASCAL 64 +#define NWARPS_Q5_1_PASCAL 8 + +#define MMQ_X_Q8_0_RDNA2 64 +#define MMQ_Y_Q8_0_RDNA2 128 +#define NWARPS_Q8_0_RDNA2 8 +#define MMQ_X_Q8_0_RDNA1 64 +#define MMQ_Y_Q8_0_RDNA1 64 +#define NWARPS_Q8_0_RDNA1 8 +#if defined(CUDA_USE_TENSOR_CORES) +#define MMQ_X_Q8_0_AMPERE 4 +#define MMQ_Y_Q8_0_AMPERE 32 +#define NWARPS_Q8_0_AMPERE 4 +#else +#define MMQ_X_Q8_0_AMPERE 128 +#define MMQ_Y_Q8_0_AMPERE 64 +#define NWARPS_Q8_0_AMPERE 4 +#endif +#define MMQ_X_Q8_0_PASCAL 64 +#define MMQ_Y_Q8_0_PASCAL 64 +#define NWARPS_Q8_0_PASCAL 8 + +#define MMQ_X_Q2_K_RDNA2 64 +#define MMQ_Y_Q2_K_RDNA2 128 +#define NWARPS_Q2_K_RDNA2 8 +#define MMQ_X_Q2_K_RDNA1 128 +#define MMQ_Y_Q2_K_RDNA1 32 +#define NWARPS_Q2_K_RDNA1 8 +#if defined(CUDA_USE_TENSOR_CORES) +#define MMQ_X_Q2_K_AMPERE 4 +#define MMQ_Y_Q2_K_AMPERE 32 +#define NWARPS_Q2_K_AMPERE 4 +#else +#define MMQ_X_Q2_K_AMPERE 64 +#define MMQ_Y_Q2_K_AMPERE 128 +#define NWARPS_Q2_K_AMPERE 4 +#endif +#define MMQ_X_Q2_K_PASCAL 64 +#define MMQ_Y_Q2_K_PASCAL 64 +#define NWARPS_Q2_K_PASCAL 8 + +#define MMQ_X_Q3_K_RDNA2 128 +#define MMQ_Y_Q3_K_RDNA2 64 +#define NWARPS_Q3_K_RDNA2 8 +#define MMQ_X_Q3_K_RDNA1 32 +#define MMQ_Y_Q3_K_RDNA1 128 +#define NWARPS_Q3_K_RDNA1 8 +#if defined(CUDA_USE_TENSOR_CORES) +#define MMQ_X_Q3_K_AMPERE 4 +#define MMQ_Y_Q3_K_AMPERE 32 +#define NWARPS_Q3_K_AMPERE 4 +#else +#define MMQ_X_Q3_K_AMPERE 128 +#define MMQ_Y_Q3_K_AMPERE 128 +#define NWARPS_Q3_K_AMPERE 4 +#endif +#define MMQ_X_Q3_K_PASCAL 64 +#define MMQ_Y_Q3_K_PASCAL 64 +#define NWARPS_Q3_K_PASCAL 8 + +#define MMQ_X_Q4_K_RDNA2 64 +#define MMQ_Y_Q4_K_RDNA2 128 +#define NWARPS_Q4_K_RDNA2 8 +#define MMQ_X_Q4_K_RDNA1 32 +#define MMQ_Y_Q4_K_RDNA1 64 +#define NWARPS_Q4_K_RDNA1 8 +#if defined(CUDA_USE_TENSOR_CORES) +#define MMQ_X_Q4_K_AMPERE 4 +#define MMQ_Y_Q4_K_AMPERE 32 +#define NWARPS_Q4_K_AMPERE 4 +#else +#define MMQ_X_Q4_K_AMPERE 64 +#define MMQ_Y_Q4_K_AMPERE 128 +#define NWARPS_Q4_K_AMPERE 4 +#endif +#define MMQ_X_Q4_K_PASCAL 64 +#define MMQ_Y_Q4_K_PASCAL 64 +#define NWARPS_Q4_K_PASCAL 8 + +#define MMQ_X_Q5_K_RDNA2 64 +#define MMQ_Y_Q5_K_RDNA2 128 +#define NWARPS_Q5_K_RDNA2 8 +#define MMQ_X_Q5_K_RDNA1 32 +#define MMQ_Y_Q5_K_RDNA1 64 +#define NWARPS_Q5_K_RDNA1 8 +#if defined(CUDA_USE_TENSOR_CORES) +#define MMQ_X_Q5_K_AMPERE 4 +#define MMQ_Y_Q5_K_AMPERE 32 +#define NWARPS_Q5_K_AMPERE 4 +#else +#define MMQ_X_Q5_K_AMPERE 64 +#define MMQ_Y_Q5_K_AMPERE 128 +#define NWARPS_Q5_K_AMPERE 4 +#endif +#define MMQ_X_Q5_K_PASCAL 64 +#define MMQ_Y_Q5_K_PASCAL 64 +#define NWARPS_Q5_K_PASCAL 8 + +#define MMQ_X_Q6_K_RDNA2 64 +#define MMQ_Y_Q6_K_RDNA2 128 +#define NWARPS_Q6_K_RDNA2 8 +#define MMQ_X_Q6_K_RDNA1 32 +#define MMQ_Y_Q6_K_RDNA1 64 +#define NWARPS_Q6_K_RDNA1 8 +#if defined(CUDA_USE_TENSOR_CORES) +#define MMQ_X_Q6_K_AMPERE 4 +#define MMQ_Y_Q6_K_AMPERE 32 +#define NWARPS_Q6_K_AMPERE 4 +#else +#define MMQ_X_Q6_K_AMPERE 64 +#define MMQ_Y_Q6_K_AMPERE 64 +#define NWARPS_Q6_K_AMPERE 4 +#endif +#define MMQ_X_Q6_K_PASCAL 64 +#define MMQ_Y_Q6_K_PASCAL 64 +#define NWARPS_Q6_K_PASCAL 8 + + +// QK = number of values after dequantization +// QR = QK / number of values before dequantization +// QI = number of 32 bit integers before dequantization + +#define QK4_0 32 +#define QR4_0 2 +#define QI4_0 (QK4_0 / (4 * QR4_0)) +typedef struct { + half d; // delta + uint8_t qs[QK4_0 / 2]; // nibbles / quants +} block_q4_0; +static_assert(sizeof(block_q4_0) == sizeof(ggml_fp16_t) + QK4_0 / 2, "wrong q4_0 block size/padding"); + +#define QK4_1 32 +#define QR4_1 2 +#define QI4_1 (QK4_1 / (4 * QR4_1)) +typedef struct { + half2 dm; // dm.x = delta, dm.y = min + uint8_t qs[QK4_1 / 2]; // nibbles / quants +} block_q4_1; +static_assert(sizeof(block_q4_1) == sizeof(ggml_fp16_t) * 2 + QK4_1 / 2, "wrong q4_1 block size/padding"); + +#define QK5_0 32 +#define QR5_0 2 +#define QI5_0 (QK5_0 / (4 * QR5_0)) +typedef struct { + half d; // delta + uint8_t qh[4]; // 5-th bit of quants + uint8_t qs[QK5_0 / 2]; // nibbles / quants +} block_q5_0; +static_assert(sizeof(block_q5_0) == sizeof(ggml_fp16_t) + sizeof(uint32_t) + QK5_0 / 2, "wrong q5_0 block size/padding"); + +#define QK5_1 32 +#define QR5_1 2 +#define QI5_1 (QK5_1 / (4 * QR5_1)) +typedef struct { + half2 dm; // dm.x = delta, dm.y = min + uint8_t qh[4]; // 5-th bit of quants + uint8_t qs[QK5_1 / 2]; // nibbles / quants +} block_q5_1; +static_assert(sizeof(block_q5_1) == 2 * sizeof(ggml_fp16_t) + sizeof(uint32_t) + QK5_1 / 2, "wrong q5_1 block size/padding"); + +#define QK8_0 32 +#define QR8_0 1 +#define QI8_0 (QK8_0 / (4 * QR8_0)) +typedef struct { + half d; // delta + int8_t qs[QK8_0]; // quants +} block_q8_0; +static_assert(sizeof(block_q8_0) == sizeof(ggml_fp16_t) + QK8_0, "wrong q8_0 block size/padding"); + +#define QK8_1 32 +#define QR8_1 1 +#define QI8_1 (QK8_1 / (4 * QR8_1)) +typedef struct { + half2 ds; // ds.x = delta, ds.y = sum + int8_t qs[QK8_0]; // quants +} block_q8_1; +static_assert(sizeof(block_q8_1) == 2*sizeof(ggml_fp16_t) + QK8_0, "wrong q8_1 block size/padding"); + +typedef float (*vec_dot_q_cuda_t)(const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs); +typedef void (*allocate_tiles_cuda_t)(int ** x_ql, half2 ** x_dm, int ** x_qh, int ** x_sc); +typedef void (*load_tiles_cuda_t)( + const void * __restrict__ vx, int * __restrict__ x_ql, half2 * __restrict__ x_dm, int * __restrict__ x_qh, + int * __restrict__ x_sc, const int & i_offset, const int & i_max, const int & k, const int & blocks_per_row); +typedef float (*vec_dot_q_mul_mat_cuda_t)( + const int * __restrict__ x_ql, const half2 * __restrict__ x_dm, const int * __restrict__ x_qh, const int * __restrict__ x_sc, + const int * __restrict__ y_qs, const half2 * __restrict__ y_ms, const int & i, const int & j, const int & k); + +#define QR2_K 4 +#define QI2_K (QK_K / (4*QR2_K)) +typedef struct { + uint8_t scales[QK_K/16]; // scales and mins, quantized with 4 bits + uint8_t qs[QK_K/4]; // quants + half2 dm; // super-block scale for quantized scales/mins +} block_q2_K; +static_assert(sizeof(block_q2_K) == 2*sizeof(ggml_fp16_t) + QK_K/16 + QK_K/4, "wrong q2_K block size/padding"); + +#define QR3_K 4 +#define QI3_K (QK_K / (4*QR3_K)) +typedef struct { + uint8_t hmask[QK_K/8]; // quants - high bit + uint8_t qs[QK_K/4]; // quants - low 2 bits +#ifdef GGML_QKK_64 + uint8_t scales[2]; // scales, quantized with 8 bits +#else + uint8_t scales[K_SCALE_SIZE]; // scales, quantized with 6 bits +#endif + half d; // super-block scale +} block_q3_K; +//static_assert(sizeof(block_q3_K) == sizeof(ggml_fp16_t) + QK_K / 4 + QK_K / 8 + K_SCALE_SIZE, "wrong q3_K block size/padding"); + +#define QR4_K 2 +#define QI4_K (QK_K / (4*QR4_K)) +#ifdef GGML_QKK_64 +typedef struct { + half dm[2]; // super-block scales/mins + uint8_t scales[2]; // 4-bit block scales/mins + uint8_t qs[QK_K/2]; // 4--bit quants +} block_q4_K; +static_assert(sizeof(block_q4_K) == sizeof(half2) + QK_K/2 + 2, "wrong q4_K block size/padding"); +#else +typedef struct { + half2 dm; // super-block scale for quantized scales/mins + uint8_t scales[3*QK_K/64]; // scales, quantized with 6 bits + uint8_t qs[QK_K/2]; // 4--bit quants +} block_q4_K; +static_assert(sizeof(block_q4_K) == 2*sizeof(ggml_fp16_t) + 3*QK_K/64 + QK_K/2, "wrong q4_K block size/padding"); +#endif + +#define QR5_K 2 +#define QI5_K (QK_K / (4*QR5_K)) +#ifdef GGML_QKK_64 +typedef struct { + half d; // super-block scale + int8_t scales[QK_K/16]; // block scales + uint8_t qh[QK_K/8]; // quants, high bit + uint8_t qs[QK_K/2]; // quants, low 4 bits +} block_q5_K; +static_assert(sizeof(block_q5_K) == sizeof(ggml_fp16_t) + QK_K/2 + QK_K/8 + QK_K/16, "wrong q5_K block size/padding"); +#else +typedef struct { + half2 dm; // super-block scale for quantized scales/mins + uint8_t scales[K_SCALE_SIZE]; // scales and mins, quantized with 6 bits + uint8_t qh[QK_K/8]; // quants, high bit + uint8_t qs[QK_K/2]; // quants, low 4 bits +} block_q5_K; +static_assert(sizeof(block_q5_K) == 2*sizeof(ggml_fp16_t) + K_SCALE_SIZE + QK_K/2 + QK_K/8, "wrong q5_K block size/padding"); +#endif + +#define QR6_K 2 +#define QI6_K (QK_K / (4*QR6_K)) +typedef struct { + uint8_t ql[QK_K/2]; // quants, lower 4 bits + uint8_t qh[QK_K/4]; // quants, upper 2 bits + int8_t scales[QK_K/16]; // scales + half d; // delta +} block_q6_K; +static_assert(sizeof(block_q6_K) == sizeof(ggml_fp16_t) + 13*QK_K/16, "wrong q6_K block size/padding"); + +// In llama.cpp this is only used for intermediate quantization and dot products +typedef struct { + float d; // delta + int8_t qs[QK_K]; // quants + int16_t bsums[QK_K/16]; // sum of quants in groups of 16 +} block_q8_K; +static_assert(sizeof(block_q8_K) == sizeof(float) + QK_K + QK_K/16*sizeof(int16_t), "wrong q8_K block size/padding"); + + +template +static __device__ __forceinline__ void mul_mat_q( + const void * __restrict__ vx, const void * __restrict__ vy, float * __restrict__ dst, + const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y, const int nrows_dst) { + + const block_q_t * x = (const block_q_t *) vx; + const block_q8_1 * y = (const block_q8_1 *) vy; + + const int blocks_per_row_x = ncols_x / qk; + const int blocks_per_col_y = nrows_y / QK8_1; + const int blocks_per_warp = WARP_SIZE / qi; + + const int & ncols_dst = ncols_y; + + const int row_dst_0 = blockIdx.x*mmq_y; + const int & row_x_0 = row_dst_0; + + const int col_dst_0 = blockIdx.y*mmq_x; + const int & col_y_0 = col_dst_0; + + int * tile_x_ql = nullptr; + half2 * tile_x_dm = nullptr; + int * tile_x_qh = nullptr; + int * tile_x_sc = nullptr; + + allocate_tiles(&tile_x_ql, &tile_x_dm, &tile_x_qh, &tile_x_sc); + + __shared__ int tile_y_qs[mmq_x * WARP_SIZE]; + __shared__ half2 tile_y_ds[mmq_x * WARP_SIZE/QI8_1]; + + float sum[mmq_y/WARP_SIZE][mmq_x/nwarps] = {{0.0f}}; + + for (int ib0 = 0; ib0 < blocks_per_row_x; ib0 += blocks_per_warp) { + + load_tiles(x + row_x_0*blocks_per_row_x + ib0, tile_x_ql, tile_x_dm, tile_x_qh, tile_x_sc, + threadIdx.y, nrows_x-row_x_0-1, threadIdx.x, blocks_per_row_x); + +#pragma unroll + for (int ir = 0; ir < qr; ++ir) { + const int kqs = ir*WARP_SIZE + threadIdx.x; + const int kbxd = kqs / QI8_1; + +#pragma unroll + for (int i = 0; i < mmq_x; i += nwarps) { + const int col_y_eff = min(col_y_0 + threadIdx.y + i, ncols_y-1); // to prevent out-of-bounds memory accesses + + const block_q8_1 * by0 = &y[col_y_eff*blocks_per_col_y + ib0 * (qk/QK8_1) + kbxd]; + + const int index_y = (threadIdx.y + i) * WARP_SIZE + kqs % WARP_SIZE; + tile_y_qs[index_y] = get_int_from_int8_aligned(by0->qs, threadIdx.x % QI8_1); + } + +#pragma unroll + for (int ids0 = 0; ids0 < mmq_x; ids0 += nwarps * QI8_1) { + const int ids = (ids0 + threadIdx.y * QI8_1 + threadIdx.x / (WARP_SIZE/QI8_1)) % mmq_x; + const int kby = threadIdx.x % (WARP_SIZE/QI8_1); + const int col_y_eff = min(col_y_0 + ids, ncols_y-1); + + // if the sum is not needed it's faster to transform the scale to f32 ahead of time + const half2 * dsi_src = &y[col_y_eff*blocks_per_col_y + ib0 * (qk/QK8_1) + ir*(WARP_SIZE/QI8_1) + kby].ds; + half2 * dsi_dst = &tile_y_ds[ids * (WARP_SIZE/QI8_1) + kby]; + if (need_sum) { + *dsi_dst = *dsi_src; + } else { + float * dfi_dst = (float *) dsi_dst; + *dfi_dst = __low2half(*dsi_src); + } + } + + __syncthreads(); + +// #pragma unroll // unrolling this loop causes too much register pressure + for (int k = ir*WARP_SIZE/qr; k < (ir+1)*WARP_SIZE/qr; k += vdr) { +#pragma unroll + for (int j = 0; j < mmq_x; j += nwarps) { +#pragma unroll + for (int i = 0; i < mmq_y; i += WARP_SIZE) { + sum[i/WARP_SIZE][j/nwarps] += vec_dot( + tile_x_ql, tile_x_dm, tile_x_qh, tile_x_sc, tile_y_qs, tile_y_ds, + threadIdx.x + i, threadIdx.y + j, k); + } + } + } + + __syncthreads(); + } + } + +#pragma unroll + for (int j = 0; j < mmq_x; j += nwarps) { + const int col_dst = col_dst_0 + j + threadIdx.y; + + if (col_dst >= ncols_dst) { + return; + } + +#pragma unroll + for (int i = 0; i < mmq_y; i += WARP_SIZE) { + const int row_dst = row_dst_0 + threadIdx.x + i; + + if (row_dst >= nrows_dst) { + continue; + } + + dst[col_dst*nrows_dst + row_dst] = sum[i/WARP_SIZE][j/nwarps]; + } + } +} + +template static __device__ __forceinline__ void load_tiles_q4_0( + const void * __restrict__ vx, int * __restrict__ x_ql, half2 * __restrict__ x_dm, int * __restrict__ x_qh, + int * __restrict__ x_sc, const int & i_offset, const int & i_max, const int & k, const int & blocks_per_row) { + (void)x_qh; (void)x_sc; + + const int kbx = k / QI4_0; + const int kqsx = k % QI4_0; + + const block_q4_0 * bx0 = (const block_q4_0 *) vx; + + float * x_dmf = (float *) x_dm; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { + int i = i0 + i_offset; + + if (need_check) { + i = min(i, i_max); + } + + const block_q4_0 * bxi = bx0 + i*blocks_per_row + kbx; + + x_ql[i * (WARP_SIZE + 1) + k] = get_int_from_uint8(bxi->qs, kqsx); + // x_dmf[i * (WARP_SIZE/QI4_0) + i / QI4_0 + kbx] = bxi->d; + } + + const int blocks_per_tile_x_row = WARP_SIZE / QI4_0; + const int kbxd = k % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI4_0) { + int i = i0 + i_offset * QI4_0 + k / blocks_per_tile_x_row; + + if (need_check) { + i = min(i, i_max); + } + + const block_q4_0 * bxi = bx0 + i*blocks_per_row + kbxd; + + x_dmf[i * (WARP_SIZE/QI4_0) + i / QI4_0 + kbxd] = bxi->d; + } +} + +template static __device__ __forceinline__ void load_tiles_q4_1( + const void * __restrict__ vx, int * __restrict__ x_ql, half2 * __restrict__ x_dm, int * __restrict__ x_qh, + int * __restrict__ x_sc, const int & i_offset, const int & i_max, const int & k, const int & blocks_per_row) { + GGML_UNUSED(x_qh); GGML_UNUSED(x_sc); + + GGML_CUDA_ASSUME(i_offset >= 0); + GGML_CUDA_ASSUME(i_offset < nwarps); + GGML_CUDA_ASSUME(k >= 0); + GGML_CUDA_ASSUME(k < WARP_SIZE); + + const int kbx = k / QI4_1; + const int kqsx = k % QI4_1; + + const block_q4_1 * bx0 = (const block_q4_1 *) vx; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { + int i = i0 + i_offset; + + if (need_check) { + i = min(i, i_max); + } + + const block_q4_1 * bxi = bx0 + i*blocks_per_row + kbx; + + x_ql[i * (WARP_SIZE + 1) + k] = get_int_from_uint8_aligned(bxi->qs, kqsx); + } + + const int blocks_per_tile_x_row = WARP_SIZE / QI4_1; + const int kbxd = k % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI4_1) { + int i = i0 + i_offset * QI4_1 + k / blocks_per_tile_x_row; + + if (need_check) { + i = min(i, i_max); + } + + const block_q4_1 * bxi = bx0 + i*blocks_per_row + kbxd; + + x_dm[i * (WARP_SIZE/QI4_1) + i / QI4_1 + kbxd] = bxi->dm; + } +} + + +template static __device__ __forceinline__ void allocate_tiles_q4_0(int ** x_ql, half2 ** x_dm, int ** x_qh, int ** x_sc) { + (void)x_qh; (void)x_sc; + + __shared__ int tile_x_qs[mmq_y * (WARP_SIZE) + mmq_y]; + __shared__ float tile_x_d[mmq_y * (WARP_SIZE/QI4_0) + mmq_y/QI4_0]; + + *x_ql = tile_x_qs; + *x_dm = (half2 *) tile_x_d; +} + +template static __device__ __forceinline__ void allocate_tiles_q4_1(int ** x_ql, half2 ** x_dm, int ** x_qh, int ** x_sc) { + GGML_UNUSED(x_qh); GGML_UNUSED(x_sc); + + __shared__ int tile_x_qs[mmq_y * (WARP_SIZE) + + mmq_y]; + __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE/QI4_1) + mmq_y/QI4_1]; + + *x_ql = tile_x_qs; + *x_dm = tile_x_dm; +} + +static __device__ __forceinline__ void dequantize_q4_0(const void * vx, const int ib, const int iqs, dfloat2 & v){ + const block_q4_0 * x = (const block_q4_0 *) vx; + + const dfloat d = x[ib].d; + + const int vui = x[ib].qs[iqs]; + + v.x = vui & 0xF; + v.y = vui >> 4; + +#ifdef GGML_CUDA_F16 + v = __hsub2(v, {8.0f, 8.0f}); + v = __hmul2(v, {d, d}); +#else + v.x = (v.x - 8.0f) * d; + v.y = (v.y - 8.0f) * d; +#endif // GGML_CUDA_F16 +} + +static __device__ __forceinline__ void dequantize_q4_1(const void * vx, const int ib, const int iqs, dfloat2 & v){ + const block_q4_1 * x = (const block_q4_1 *) vx; + + const dfloat d = __low2half(x[ib].dm); + const dfloat m = __high2half(x[ib].dm); + + const int vui = x[ib].qs[iqs]; + + v.x = vui & 0xF; + v.y = vui >> 4; + +#ifdef GGML_CUDA_F16 + v = __hmul2(v, {d, d}); + v = __hadd2(v, {m, m}); +#else + v.x = (v.x * d) + m; + v.y = (v.y * d) + m; +#endif // GGML_CUDA_F16 +} + +static __device__ __forceinline__ void dequantize_q5_0(const void * vx, const int ib, const int iqs, dfloat2 & v){ + const block_q5_0 * x = (const block_q5_0 *) vx; + + const dfloat d = x[ib].d; + + uint32_t qh; + memcpy(&qh, x[ib].qh, sizeof(qh)); + + const int xh_0 = ((qh >> (iqs + 0)) << 4) & 0x10; + const int xh_1 = ((qh >> (iqs + 12)) ) & 0x10; + + v.x = ((x[ib].qs[iqs] & 0xf) | xh_0); + v.y = ((x[ib].qs[iqs] >> 4) | xh_1); + +#ifdef GGML_CUDA_F16 + v = __hsub2(v, {16.0f, 16.0f}); + v = __hmul2(v, {d, d}); +#else + v.x = (v.x - 16.0f) * d; + v.y = (v.y - 16.0f) * d; +#endif // GGML_CUDA_F16 +} + +static __device__ __forceinline__ void dequantize_q5_1(const void * vx, const int ib, const int iqs, dfloat2 & v){ + const block_q5_1 * x = (const block_q5_1 *) vx; + + const dfloat d = __low2half(x[ib].dm); + const dfloat m = __high2half(x[ib].dm); + + uint32_t qh; + memcpy(&qh, x[ib].qh, sizeof(qh)); + + const int xh_0 = ((qh >> (iqs + 0)) << 4) & 0x10; + const int xh_1 = ((qh >> (iqs + 12)) ) & 0x10; + + v.x = ((x[ib].qs[iqs] & 0xf) | xh_0); + v.y = ((x[ib].qs[iqs] >> 4) | xh_1); + +#ifdef GGML_CUDA_F16 + v = __hmul2(v, {d, d}); + v = __hadd2(v, {m, m}); +#else + v.x = (v.x * d) + m; + v.y = (v.y * d) + m; +#endif // GGML_CUDA_F16 +} + +static __device__ __forceinline__ void dequantize_q8_0(const void * vx, const int ib, const int iqs, dfloat2 & v){ + const block_q8_0 * x = (const block_q8_0 *) vx; + + const dfloat d = x[ib].d; + + v.x = x[ib].qs[iqs + 0]; + v.y = x[ib].qs[iqs + 1]; + +#ifdef GGML_CUDA_F16 + v = __hmul2(v, {d, d}); +#else + v.x *= d; + v.y *= d; +#endif // GGML_CUDA_F16 +} + + +template +static __device__ void dequantize_block(const void * __restrict__ vx, dst_t * __restrict__ y, const int k) { + const int i = 2*(blockDim.x*blockIdx.x + threadIdx.x); + + if (i >= k) { + return; + } + + const int ib = i/qk; // block index + const int iqs = (i%qk)/qr; // quant index + const int iybs = i - i%qk; // y block start index + const int y_offset = qr == 1 ? 1 : qk/2; + + // dequantize + dfloat2 v; + dequantize_kernel(vx, ib, iqs, v); + + y[iybs + iqs + 0] = v.x; + y[iybs + iqs + y_offset] = v.y; +} + +template +static __device__ void dequantize_block_q4_0(const void * __restrict__ vx, dst_t * __restrict__ yy, int nb32) { + + const int64_t i = blockIdx.x; + + // assume 32 threads + const int tid = threadIdx.x; + const int il = tid/8; + const int ir = tid%8; + const int64_t ib = 8*i + ir; + if (ib >= nb32) { + return; + } + + dst_t * y = yy + 256*i + 32*ir + 4*il; + + const block_q4_0 * x = (const block_q4_0 *)vx + ib; + const float d = __half2float(x->d); + const float dm = -8*d; + + const uint8_t * q = x->qs + 4*il; + + for (int l = 0; l < 4; ++l) { + y[l+ 0] = d * (q[l] & 0xF) + dm; + y[l+16] = d * (q[l] >> 4) + dm; + } +} + +template +static __device__ void dequantize_block_q4_1(const void * __restrict__ vx, dst_t * __restrict__ yy, int nb32) { + + const int64_t i = blockIdx.x; + + // assume 32 threads + const int tid = threadIdx.x; + const int il = tid/8; + const int ir = tid%8; + const int64_t ib = 8*i + ir; + if (ib >= nb32) { + return; + } + + dst_t * y = yy + 256*i + 32*ir + 4*il; + + const block_q4_1 * x = (const block_q4_1 *)vx + ib; + const float2 d = __half22float2(x->dm); + + const uint8_t * q = x->qs + 4*il; + + for (int l = 0; l < 4; ++l) { + y[l+ 0] = d.x * (q[l] & 0xF) + d.y; + y[l+16] = d.x * (q[l] >> 4) + d.y; + } +} + +//================================== k-quants + +template +static __device__ void dequantize_block_q2_K(const void * __restrict__ vx, dst_t * __restrict__ yy) { + + const int i = blockIdx.x; + const block_q2_K * x = (const block_q2_K *) vx; + + const int tid = threadIdx.x; +#if QK_K == 256 + const int n = tid/32; + const int l = tid - 32*n; + const int is = 8*n + l/16; + + const uint8_t q = x[i].qs[32*n + l]; + dst_t * y = yy + i*QK_K + 128*n; + + float dall = __low2half(x[i].dm); + float dmin = __high2half(x[i].dm); + y[l+ 0] = dall * (x[i].scales[is+0] & 0xF) * ((q >> 0) & 3) - dmin * (x[i].scales[is+0] >> 4); + y[l+32] = dall * (x[i].scales[is+2] & 0xF) * ((q >> 2) & 3) - dmin * (x[i].scales[is+2] >> 4); + y[l+64] = dall * (x[i].scales[is+4] & 0xF) * ((q >> 4) & 3) - dmin * (x[i].scales[is+4] >> 4); + y[l+96] = dall * (x[i].scales[is+6] & 0xF) * ((q >> 6) & 3) - dmin * (x[i].scales[is+6] >> 4); +#else + const int is = tid/16; // 0 or 1 + const int il = tid%16; // 0...15 + const uint8_t q = x[i].qs[il] >> (2*is); + dst_t * y = yy + i*QK_K + 16*is + il; + float dall = __low2half(x[i].dm); + float dmin = __high2half(x[i].dm); + y[ 0] = dall * (x[i].scales[is+0] & 0xF) * ((q >> 0) & 3) - dmin * (x[i].scales[is+0] >> 4); + y[32] = dall * (x[i].scales[is+2] & 0xF) * ((q >> 4) & 3) - dmin * (x[i].scales[is+2] >> 4); +#endif + +} + +template +static __device__ void dequantize_block_q3_K(const void * __restrict__ vx, dst_t * __restrict__ yy) { + + const int i = blockIdx.x; + const block_q3_K * x = (const block_q3_K *) vx; + +#if QK_K == 256 + const int r = threadIdx.x/4; + const int tid = r/2; + const int is0 = r%2; + const int l0 = 16*is0 + 4*(threadIdx.x%4); + const int n = tid / 4; + const int j = tid - 4*n; + + uint8_t m = 1 << (4*n + j); + int is = 8*n + 2*j + is0; + int shift = 2*j; + + int8_t us = is < 4 ? (x[i].scales[is-0] & 0xF) | (((x[i].scales[is+8] >> 0) & 3) << 4) : + is < 8 ? (x[i].scales[is-0] & 0xF) | (((x[i].scales[is+4] >> 2) & 3) << 4) : + is < 12 ? (x[i].scales[is-8] >> 4) | (((x[i].scales[is+0] >> 4) & 3) << 4) : + (x[i].scales[is-8] >> 4) | (((x[i].scales[is-4] >> 6) & 3) << 4); + float d_all = x[i].d; + float dl = d_all * (us - 32); + + dst_t * y = yy + i*QK_K + 128*n + 32*j; + const uint8_t * q = x[i].qs + 32*n; + const uint8_t * hm = x[i].hmask; + + for (int l = l0; l < l0+4; ++l) y[l] = dl * ((int8_t)((q[l] >> shift) & 3) - ((hm[l] & m) ? 0 : 4)); +#else + const int tid = threadIdx.x; + const int is = tid/16; // 0 or 1 + const int il = tid%16; // 0...15 + const int im = il/8; // 0...1 + const int in = il%8; // 0...7 + + dst_t * y = yy + i*QK_K + 16*is + il; + + const uint8_t q = x[i].qs[il] >> (2*is); + const uint8_t h = x[i].hmask[in] >> (2*is + im); + const float d = (float)x[i].d; + + if (is == 0) { + y[ 0] = d * ((x[i].scales[0] & 0xF) - 8) * ((int8_t)((q >> 0) & 3) - ((h >> 0) & 1 ? 0 : 4)); + y[32] = d * ((x[i].scales[1] & 0xF) - 8) * ((int8_t)((q >> 4) & 3) - ((h >> 4) & 1 ? 0 : 4)); + } else { + y[ 0] = d * ((x[i].scales[0] >> 4) - 8) * ((int8_t)((q >> 0) & 3) - ((h >> 0) & 1 ? 0 : 4)); + y[32] = d * ((x[i].scales[1] >> 4) - 8) * ((int8_t)((q >> 4) & 3) - ((h >> 4) & 1 ? 0 : 4)); + } +#endif + +} + +#if QK_K == 256 +static inline __device__ void get_scale_min_k4(int j, const uint8_t * q, uint8_t & d, uint8_t & m) { + if (j < 4) { + d = q[j] & 63; m = q[j + 4] & 63; + } else { + d = (q[j+4] & 0xF) | ((q[j-4] >> 6) << 4); + m = (q[j+4] >> 4) | ((q[j-0] >> 6) << 4); + } +} +#endif + +template +static __device__ void dequantize_block_q4_K(const void * __restrict__ vx, dst_t * __restrict__ yy) { + const block_q4_K * x = (const block_q4_K *) vx; + + const int i = blockIdx.x; + +#if QK_K == 256 + // assume 32 threads + const int tid = threadIdx.x; + const int il = tid/8; + const int ir = tid%8; + const int is = 2*il; + const int n = 4; + + dst_t * y = yy + i*QK_K + 64*il + n*ir; + + const float dall = __low2half(x[i].dm); + const float dmin = __high2half(x[i].dm); + + const uint8_t * q = x[i].qs + 32*il + n*ir; + + uint8_t sc, m; + get_scale_min_k4(is + 0, x[i].scales, sc, m); + const float d1 = dall * sc; const float m1 = dmin * m; + get_scale_min_k4(is + 1, x[i].scales, sc, m); + const float d2 = dall * sc; const float m2 = dmin * m; + for (int l = 0; l < n; ++l) { + y[l + 0] = d1 * (q[l] & 0xF) - m1; + y[l +32] = d2 * (q[l] >> 4) - m2; + } +#else + const int tid = threadIdx.x; + const uint8_t * q = x[i].qs; + dst_t * y = yy + i*QK_K; + const float d = (float)x[i].dm[0]; + const float m = (float)x[i].dm[1]; + y[tid+ 0] = d * (x[i].scales[0] & 0xF) * (q[tid] & 0xF) - m * (x[i].scales[0] >> 4); + y[tid+32] = d * (x[i].scales[1] & 0xF) * (q[tid] >> 4) - m * (x[i].scales[1] >> 4); +#endif +} + +template +static __device__ void dequantize_block_q5_K(const void * __restrict__ vx, dst_t * __restrict__ yy) { + const block_q5_K * x = (const block_q5_K *) vx; + + const int i = blockIdx.x; + +#if QK_K == 256 + // assume 64 threads - this is very slightly better than the one below + const int tid = threadIdx.x; + const int il = tid/16; // il is in 0...3 + const int ir = tid%16; // ir is in 0...15 + const int is = 2*il; // is is in 0...6 + + dst_t * y = yy + i*QK_K + 64*il + 2*ir; + + const float dall = __low2half(x[i].dm); + const float dmin = __high2half(x[i].dm); + + const uint8_t * ql = x[i].qs + 32*il + 2*ir; + const uint8_t * qh = x[i].qh + 2*ir; + + uint8_t sc, m; + get_scale_min_k4(is + 0, x[i].scales, sc, m); + const float d1 = dall * sc; const float m1 = dmin * m; + get_scale_min_k4(is + 1, x[i].scales, sc, m); + const float d2 = dall * sc; const float m2 = dmin * m; + + uint8_t hm = 1 << (2*il); + y[ 0] = d1 * ((ql[ 0] & 0xF) + (qh[ 0] & hm ? 16 : 0)) - m1; + y[ 1] = d1 * ((ql[ 1] & 0xF) + (qh[ 1] & hm ? 16 : 0)) - m1; + hm <<= 1; + y[32] = d2 * ((ql[ 0] >> 4) + (qh[ 0] & hm ? 16 : 0)) - m2; + y[33] = d2 * ((ql[ 1] >> 4) + (qh[ 1] & hm ? 16 : 0)) - m2; +#else + const int tid = threadIdx.x; + const uint8_t q = x[i].qs[tid]; + const int im = tid/8; // 0...3 + const int in = tid%8; // 0...7 + const int is = tid/16; // 0 or 1 + const uint8_t h = x[i].qh[in] >> im; + const float d = x[i].d; + dst_t * y = yy + i*QK_K + tid; + y[ 0] = d * x[i].scales[is+0] * ((q & 0xF) - ((h >> 0) & 1 ? 0 : 16)); + y[32] = d * x[i].scales[is+2] * ((q >> 4) - ((h >> 4) & 1 ? 0 : 16)); +#endif +} + +template +static __device__ void dequantize_block_q6_K(const void * __restrict__ vx, dst_t * __restrict__ yy) { + const block_q6_K * x = (const block_q6_K *) vx; + + const int64_t i = blockIdx.x; +#if QK_K == 256 + + // assume 64 threads - this is very slightly better than the one below + const int64_t tid = threadIdx.x; + const int64_t ip = tid/32; // ip is 0 or 1 + const int64_t il = tid - 32*ip; // 0...32 + const int64_t is = 8*ip + il/16; + + dst_t * y = yy + i*QK_K + 128*ip + il; + + const float d = x[i].d; + + const uint8_t * ql = x[i].ql + 64*ip + il; + const uint8_t qh = x[i].qh[32*ip + il]; + const int8_t * sc = x[i].scales + is; + + y[ 0] = d * sc[0] * ((int8_t)((ql[ 0] & 0xF) | (((qh >> 0) & 3) << 4)) - 32); + y[32] = d * sc[2] * ((int8_t)((ql[32] & 0xF) | (((qh >> 2) & 3) << 4)) - 32); + y[64] = d * sc[4] * ((int8_t)((ql[ 0] >> 4) | (((qh >> 4) & 3) << 4)) - 32); + y[96] = d * sc[6] * ((int8_t)((ql[32] >> 4) | (((qh >> 6) & 3) << 4)) - 32); +#else + + // assume 32 threads + const int64_t tid = threadIdx.x; + const int64_t ip = tid/16; // 0 or 1 + const int64_t il = tid - 16*ip; // 0...15 + + dst_t * y = yy + i*QK_K + 16*ip + il; + + const float d = x[i].d; + + const uint8_t ql = x[i].ql[16*ip + il]; + const uint8_t qh = x[i].qh[il] >> (2*ip); + const int8_t * sc = x[i].scales; + + y[ 0] = d * sc[ip+0] * ((int8_t)((ql & 0xF) | (((qh >> 0) & 3) << 4)) - 32); + y[32] = d * sc[ip+2] * ((int8_t)((ql >> 4) | (((qh >> 4) & 3) << 4)) - 32); +#endif +} + +template +static __device__ void dequantize_block_q8_0(const void * __restrict__ vx, dst_t * __restrict__ yy, int nb32) { + const int i = blockIdx.x; + + // assume 32 threads + const int tid = threadIdx.x; + const int il = tid/8; + const int ir = tid%8; + const int ib = 8*i + ir; + if (ib >= nb32) { + return; + } + + dst_t * y = yy + 256*i + 32*ir + 8*il; + + const block_q8_0 * x = (const block_q8_0 *)vx + ib; + const float d = __half2float(x->d); + + const int8_t * q = x->qs + 8*il; + + for (int l = 0; l < 8; ++l) { + y[l] = d * q[l]; + } +} + +template +static __device__ void dequantize_block_q8_K(const void * __restrict__ vx, dst_t * __restrict__ yy) { + const block_q8_K * x = (const block_q8_K *) vx; + + const int i = blockIdx.x; + +#if QK_K == 256 + // assume 32 threads + const int tid = threadIdx.x; + const int il = tid/8; + const int ir = tid%8; + const int n = 8; + + dst_t * y = yy + i*QK_K + 64*il + n*ir; + + const int8_t * q = x[i].qs + 64*il + n*ir; + + for (int l = 0; l < n; ++l) { + y[l] = q[l] * x[i].d; + } +#else + const int tid = threadIdx.x; + const uint8_t * q = x[i].qs; + float * y = yy + i*QK_K; + y[tid] = x[i].d * x[i].scales[0]; +#endif +} + +template +static __device__ void dequantize_block_q5_0(const void * __restrict__ vx, dst_t * __restrict__ yy, int nb32) { + return dequantize_block(vx, yy, nb32); +} + +template +static __device__ void dequantize_block_q5_1(const void * __restrict__ vx, dst_t * __restrict__ yy, int nb32) { + return dequantize_block(vx, yy, nb32); +} + +#define DEQUANTIZE_K(QNAME) \ +extern "C" __global__ void dequantize_block_##QNAME##_f32(const void * __restrict__ vx, float * __restrict__ y) { \ + dequantize_block_##QNAME(vx, y); \ +} \ +extern "C" __global__ void dequantize_block_##QNAME##_f16(const void * __restrict__ vx, half * __restrict__ y) { \ + dequantize_block_##QNAME(vx, y); \ +} \ + +#define DEQUANTIZE(QNAME) \ +extern "C" __global__ void dequantize_block_##QNAME##_f32(const void * __restrict__ vx, float * __restrict__ y, const int k) { \ + dequantize_block_##QNAME(vx, y, k); \ +} \ +extern "C" __global__ void dequantize_block_##QNAME##_f16(const void * __restrict__ vx, half * __restrict__ y, const int k) { \ + dequantize_block_##QNAME(vx, y, k); \ +} \ + +DEQUANTIZE_K(q2_K) +DEQUANTIZE_K(q3_K) +DEQUANTIZE_K(q4_K) +DEQUANTIZE_K(q5_K) +DEQUANTIZE_K(q6_K) +DEQUANTIZE_K(q8_K) +DEQUANTIZE(q4_0) +DEQUANTIZE(q4_1) +DEQUANTIZE(q5_0) +DEQUANTIZE(q5_1) +DEQUANTIZE(q8_0) + +template +static __device__ void get_rows_q( + const void * __restrict__ src0, const uint32_t * __restrict__ src1, float * __restrict__ dst, + const int64_t ne00, const size_t nb01) { + for (int64_t i00 = 2*(blockIdx.y*blockDim.x + threadIdx.x); i00 < ne00; i00 += gridDim.y*blockDim.x) { + const int i10 = blockIdx.x; + const int i01 = src1[i10]; + + float * dst_row = dst + i10*ne00; + const void * src0_row = (const char *) src0 + i01*nb01; + + const int ib = i00/qk; + const int iqs = (i00%qk)/qr; + const int iybs = i00 - i00%qk; + const int y_offset = qr == 1 ? 1 : qk/2; + + float2 v; + dequantize_kernel(src0_row, ib, iqs, v); + + dst_row[iybs + iqs + 0] = v.x; + dst_row[iybs + iqs + y_offset] = v.y; + } +} + +static __device__ __forceinline__ float get_rows_to_float(const float v) { + return v; +} + +static __device__ __forceinline__ float get_rows_to_float(const half v) { + return __half2float(v); +} + +static __device__ __forceinline__ float get_rows_to_float(const __nv_bfloat16 v) { + return __bfloat162float(v); +} + +template +static __device__ void get_rows_float( + const src0_t * src0, const uint32_t * src1, float * dst, const int64_t ne00) { + for (int64_t i00 = blockIdx.y*blockDim.x + threadIdx.x; i00 < ne00; i00 += gridDim.y*blockDim.x) { + const int i10 = blockIdx.x; + const int i01 = src1[i10]; + + float * dst_row = dst + i10*ne00; + const src0_t * src0_row = src0 + i01*ne00; + + dst_row[i00] = get_rows_to_float(src0_row[i00]); + } +} + +static __device__ void get_rows_q2_K_block(const block_q2_K * x, float * y) { + const int tid = threadIdx.x; +#if QK_K == 256 + const int n = tid/32; + const int l = tid - 32*n; + const int is = 8*n + l/16; + + const uint8_t q = x->qs[32*n + l]; + float * yy = y + 128*n; + + float dall = __low2half(x->dm); + float dmin = __high2half(x->dm); + yy[l+ 0] = dall * (x->scales[is+0] & 0xF) * ((q >> 0) & 3) - dmin * (x->scales[is+0] >> 4); + yy[l+32] = dall * (x->scales[is+2] & 0xF) * ((q >> 2) & 3) - dmin * (x->scales[is+2] >> 4); + yy[l+64] = dall * (x->scales[is+4] & 0xF) * ((q >> 4) & 3) - dmin * (x->scales[is+4] >> 4); + yy[l+96] = dall * (x->scales[is+6] & 0xF) * ((q >> 6) & 3) - dmin * (x->scales[is+6] >> 4); +#else + const int is = tid/16; + const int il = tid%16; + const uint8_t q = x->qs[il] >> (2*is); + float * yy = y + 16*is + il; + float dall = __low2half(x->dm); + float dmin = __high2half(x->dm); + yy[ 0] = dall * (x->scales[is+0] & 0xF) * ((q >> 0) & 3) - dmin * (x->scales[is+0] >> 4); + yy[32] = dall * (x->scales[is+2] & 0xF) * ((q >> 4) & 3) - dmin * (x->scales[is+2] >> 4); +#endif +} + +static __device__ void get_rows_q3_K_block(const block_q3_K * x, float * y) { +#if QK_K == 256 + const int r = threadIdx.x/4; + const int tid = r/2; + const int is0 = r%2; + const int l0 = 16*is0 + 4*(threadIdx.x%4); + const int n = tid / 4; + const int j = tid - 4*n; + + uint8_t m = 1 << (4*n + j); + int is = 8*n + 2*j + is0; + int shift = 2*j; + + int8_t us = is < 4 ? (x->scales[is-0] & 0xF) | (((x->scales[is+8] >> 0) & 3) << 4) : + is < 8 ? (x->scales[is-0] & 0xF) | (((x->scales[is+4] >> 2) & 3) << 4) : + is < 12 ? (x->scales[is-8] >> 4) | (((x->scales[is+0] >> 4) & 3) << 4) : + (x->scales[is-8] >> 4) | (((x->scales[is-4] >> 6) & 3) << 4); + float d_all = x->d; + float dl = d_all * (us - 32); + + float * yy = y + 128*n + 32*j; + const uint8_t * q = x->qs + 32*n; + const uint8_t * hm = x->hmask; + + for (int l = l0; l < l0+4; ++l) yy[l] = dl * ((int8_t)((q[l] >> shift) & 3) - ((hm[l] & m) ? 0 : 4)); +#else + const int tid = threadIdx.x; + const int is = tid/16; + const int il = tid%16; + const int im = il/8; + const int in = il%8; + + float * yy = y + 16*is + il; + + const uint8_t q = x->qs[il] >> (2*is); + const uint8_t h = x->hmask[in] >> (2*is + im); + const float d = (float)x->d; + + if (is == 0) { + yy[ 0] = d * ((x->scales[0] & 0xF) - 8) * ((int8_t)((q >> 0) & 3) - ((h >> 0) & 1 ? 0 : 4)); + yy[32] = d * ((x->scales[1] & 0xF) - 8) * ((int8_t)((q >> 4) & 3) - ((h >> 4) & 1 ? 0 : 4)); + } else { + yy[ 0] = d * ((x->scales[0] >> 4) - 8) * ((int8_t)((q >> 0) & 3) - ((h >> 0) & 1 ? 0 : 4)); + yy[32] = d * ((x->scales[1] >> 4) - 8) * ((int8_t)((q >> 4) & 3) - ((h >> 4) & 1 ? 0 : 4)); + } +#endif +} + +static __device__ void get_rows_q4_K_block(const block_q4_K * x, float * y) { +#if QK_K == 256 + const int tid = threadIdx.x; + const int il = tid/8; + const int ir = tid%8; + const int is = 2*il; + const int n = 4; + + float * yy = y + 64*il + n*ir; + + const float dall = __low2half(x->dm); + const float dmin = __high2half(x->dm); + + const uint8_t * q = x->qs + 32*il + n*ir; + + uint8_t sc, m; + get_scale_min_k4(is + 0, x->scales, sc, m); + const float d1 = dall * sc; const float m1 = dmin * m; + get_scale_min_k4(is + 1, x->scales, sc, m); + const float d2 = dall * sc; const float m2 = dmin * m; + for (int l = 0; l < n; ++l) { + yy[l + 0] = d1 * (q[l] & 0xF) - m1; + yy[l +32] = d2 * (q[l] >> 4) - m2; + } +#else + const int tid = threadIdx.x; + const uint8_t * q = x->qs; + const float d = (float)x->dm[0]; + const float m = (float)x->dm[1]; + y[tid+ 0] = d * (x->scales[0] & 0xF) * (q[tid] & 0xF) - m * (x->scales[0] >> 4); + y[tid+32] = d * (x->scales[1] & 0xF) * (q[tid] >> 4) - m * (x->scales[1] >> 4); +#endif +} + +static __device__ void get_rows_q5_K_block(const block_q5_K * x, float * y) { +#if QK_K == 256 + const int tid = threadIdx.x; + const int il = tid/16; + const int ir = tid%16; + const int is = 2*il; + + float * yy = y + 64*il + 2*ir; + + const float dall = __low2half(x->dm); + const float dmin = __high2half(x->dm); + + const uint8_t * ql = x->qs + 32*il + 2*ir; + const uint8_t * qh = x->qh + 2*ir; + + uint8_t sc, m; + get_scale_min_k4(is + 0, x->scales, sc, m); + const float d1 = dall * sc; const float m1 = dmin * m; + get_scale_min_k4(is + 1, x->scales, sc, m); + const float d2 = dall * sc; const float m2 = dmin * m; + + uint8_t hm = 1 << (2*il); + yy[ 0] = d1 * ((ql[ 0] & 0xF) + (qh[ 0] & hm ? 16 : 0)) - m1; + yy[ 1] = d1 * ((ql[ 1] & 0xF) + (qh[ 1] & hm ? 16 : 0)) - m1; + hm <<= 1; + yy[32] = d2 * ((ql[ 0] >> 4) + (qh[ 0] & hm ? 16 : 0)) - m2; + yy[33] = d2 * ((ql[ 1] >> 4) + (qh[ 1] & hm ? 16 : 0)) - m2; +#else + const int tid = threadIdx.x; + const uint8_t q = x->qs[tid]; + const int im = tid/8; + const int in = tid%8; + const int is = tid/16; + const uint8_t h = x->qh[in] >> im; + const float d = x->d; + y[tid + 0] = d * x->scales[is+0] * ((q & 0xF) - ((h >> 0) & 1 ? 0 : 16)); + y[tid+32] = d * x->scales[is+2] * ((q >> 4) - ((h >> 4) & 1 ? 0 : 16)); +#endif +} + +static __device__ void get_rows_q6_K_block(const block_q6_K * x, float * y) { +#if QK_K == 256 + const int64_t tid = threadIdx.x; + const int64_t ip = tid/32; + const int64_t il = tid - 32*ip; + const int64_t is = 8*ip + il/16; + + float * yy = y + 128*ip + il; + + const float d = x->d; + + const uint8_t * ql = x->ql + 64*ip + il; + const uint8_t qh = x->qh[32*ip + il]; + const int8_t * sc = x->scales + is; + + yy[ 0] = d * sc[0] * ((int8_t)((ql[ 0] & 0xF) | (((qh >> 0) & 3) << 4)) - 32); + yy[32] = d * sc[2] * ((int8_t)((ql[32] & 0xF) | (((qh >> 2) & 3) << 4)) - 32); + yy[64] = d * sc[4] * ((int8_t)((ql[ 0] >> 4) | (((qh >> 4) & 3) << 4)) - 32); + yy[96] = d * sc[6] * ((int8_t)((ql[32] >> 4) | (((qh >> 6) & 3) << 4)) - 32); +#else + const int64_t tid = threadIdx.x; + const int64_t ip = tid/16; + const int64_t il = tid - 16*ip; + + float * yy = y + 16*ip + il; + + const float d = x->d; + + const uint8_t ql = x->ql[16*ip + il]; + const uint8_t qh = x->qh[il] >> (2*ip); + const int8_t * sc = x->scales; + + yy[ 0] = d * sc[ip+0] * ((int8_t)((ql & 0xF) | (((qh >> 0) & 3) << 4)) - 32); + yy[32] = d * sc[ip+2] * ((int8_t)((ql >> 4) | (((qh >> 4) & 3) << 4)) - 32); +#endif +} + +#define GET_ROWS_K(QNAME, BLOCK_Q_T) \ +extern "C" __global__ void get_rows_##QNAME( \ + const void * __restrict__ src0, const uint32_t * __restrict__ src1, float * __restrict__ dst, \ + const int64_t ne00, const size_t nb01) { \ + const int i10 = blockIdx.x; \ + const int i01 = src1[i10]; \ + const int ib = blockIdx.y; \ + const int blocks_per_row = nb01 / sizeof(BLOCK_Q_T); \ + const BLOCK_Q_T * src0_row = (const BLOCK_Q_T *) src0 + i01*blocks_per_row; \ + float * dst_row = dst + i10*ne00; \ + get_rows_##QNAME##_block(src0_row + ib, dst_row + ib*QK_K); \ +} + +extern "C" __global__ void get_rows_q4_0( + const void * __restrict__ src0, const uint32_t * __restrict__ src1, float * __restrict__ dst, + const int64_t ne00, const size_t nb01) { + get_rows_q(src0, src1, dst, ne00, nb01); +} + +extern "C" __global__ void get_rows_q4_1( + const void * __restrict__ src0, const uint32_t * __restrict__ src1, float * __restrict__ dst, + const int64_t ne00, const size_t nb01) { + get_rows_q(src0, src1, dst, ne00, nb01); +} + +extern "C" __global__ void get_rows_q5_0( + const void * __restrict__ src0, const uint32_t * __restrict__ src1, float * __restrict__ dst, + const int64_t ne00, const size_t nb01) { + get_rows_q(src0, src1, dst, ne00, nb01); +} + +extern "C" __global__ void get_rows_q5_1( + const void * __restrict__ src0, const uint32_t * __restrict__ src1, float * __restrict__ dst, + const int64_t ne00, const size_t nb01) { + get_rows_q(src0, src1, dst, ne00, nb01); +} + +extern "C" __global__ void get_rows_q8_0( + const void * __restrict__ src0, const uint32_t * __restrict__ src1, float * __restrict__ dst, + const int64_t ne00, const size_t nb01) { + get_rows_q(src0, src1, dst, ne00, nb01); +} + +extern "C" __global__ void get_rows_f32( + const float * __restrict__ src0, const uint32_t * __restrict__ src1, float * __restrict__ dst, + const int64_t ne00, const size_t nb01) { + GGML_UNUSED(nb01); + get_rows_float(src0, src1, dst, ne00); +} + +extern "C" __global__ void get_rows_f16( + const half * __restrict__ src0, const uint32_t * __restrict__ src1, float * __restrict__ dst, + const int64_t ne00, const size_t nb01) { + GGML_UNUSED(nb01); + get_rows_float(src0, src1, dst, ne00); +} + +extern "C" __global__ void get_rows_bf16( + const __nv_bfloat16 * __restrict__ src0, const uint32_t * __restrict__ src1, float * __restrict__ dst, + const int64_t ne00, const size_t nb01) { + GGML_UNUSED(nb01); + get_rows_float(src0, src1, dst, ne00); +} + +GET_ROWS_K(q2_K, block_q2_K) +GET_ROWS_K(q3_K, block_q3_K) +GET_ROWS_K(q4_K, block_q4_K) +GET_ROWS_K(q5_K, block_q5_K) +GET_ROWS_K(q6_K, block_q6_K) + +template +static __device__ void dequantize_mul_mat_vec(const void * __restrict__ vx, const dfloat * __restrict__ y, float * __restrict__ dst, const int ncols, const int nrows) { + // qk = quantized weights per x block + // qr = number of quantized weights per data value in x block + const int row = blockIdx.x*blockDim.y + threadIdx.y; + + if (row >= nrows) { + return; + } + + const int tid = threadIdx.x; + + const int iter_stride = 2*GGML_CUDA_DMMV_X; + const int vals_per_iter = iter_stride / WARP_SIZE; // num quantized vals per thread and i iter + const int y_offset = qr == 1 ? 1 : qk/2; + +// partial sum for each thread +#ifdef GGML_CUDA_F16 + half2 tmp = {0.0f, 0.0f}; // two sums for f16 to take advantage of half2 intrinsics +#else + float tmp = 0.0f; +#endif // GGML_CUDA_F16 + + for (int i = 0; i < ncols; i += iter_stride) { + const int col = i + vals_per_iter*tid; + const int ib = (row*ncols + col)/qk; // x block index + const int iqs = (col%qk)/qr; // x quant index + const int iybs = col - col%qk; // y block start index + +// processing >2 values per i iter is faster for fast GPUs +#pragma unroll + for (int j = 0; j < vals_per_iter; j += 2) { + // process 2 vals per j iter + + // dequantize + // for qr = 2 the iqs needs to increase by 1 per j iter because 2 weights per data val + dfloat2 v; + dequantize_kernel(vx, ib, iqs + j/qr, v); + + // matrix multiplication + // for qr = 2 the y index needs to increase by 1 per j iter because of y_offset = qk/2 +#ifdef GGML_CUDA_F16 + tmp += __hmul2(v, { + y[iybs + iqs + j/qr + 0], + y[iybs + iqs + j/qr + y_offset] + }); +#else + tmp += v.x * y[iybs + iqs + j/qr + 0]; + tmp += v.y * y[iybs + iqs + j/qr + y_offset]; +#endif // GGML_CUDA_F16 + } + } + + // sum up partial sums and write back result +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + tmp += __shfl_xor_sync(0xffffffff, tmp, mask, 32); + } + + if (tid == 0) { +#ifdef GGML_CUDA_F16 + dst[row] = tmp.x + tmp.y; +#else + dst[row] = tmp; +#endif // GGML_CUDA_F16 + } +} + +extern "C" __global__ void dequantize_mul_mat_vec_q4_0_cuda(const void * vx, const dfloat * y, float * dst, const int ncols, const int nrows) { + dequantize_mul_mat_vec(vx, y, dst, ncols, nrows); +} + +extern "C" __global__ void dequantize_mul_mat_vec_q4_1_cuda(const void * vx, const dfloat * y, float * dst, const int ncols, const int nrows) { + dequantize_mul_mat_vec(vx, y, dst, ncols, nrows); +} + +extern "C" __global__ void dequantize_mul_mat_vec_q5_0_cuda(const void * vx, const dfloat * y, float * dst, const int ncols, const int nrows) { + dequantize_mul_mat_vec(vx, y, dst, ncols, nrows); +} + +extern "C" __global__ void dequantize_mul_mat_vec_q5_1_cuda(const void * vx, const dfloat * y, float * dst, const int ncols, const int nrows) { + dequantize_mul_mat_vec(vx, y, dst, ncols, nrows); +} +extern "C" __global__ void dequantize_mul_mat_vec_q8_0_cuda(const void * vx, const dfloat * y, float * dst, const int ncols, const int nrows) { + dequantize_mul_mat_vec(vx, y, dst, ncols, nrows); +} + +extern "C" __global__ void dequantize_mul_mat_vec_q2_k(const void * __restrict__ vx, const float * __restrict__ yy, float * __restrict__ dst, const int ncols, int nrows) { + + static_assert(16%K_QUANTS_PER_ITERATION == 0, "16 must be divisible by K_QUANTS_PER_ITERATION"); + + const int row = blockIdx.x*blockDim.y + threadIdx.y; + if (row > nrows) return; + + const int num_blocks_per_row = ncols / QK_K; + const int ib0 = row*num_blocks_per_row; + + const block_q2_K * x = (const block_q2_K *)vx + ib0; + + float tmp = 0; // partial sum for thread in warp + +#if QK_K == 256 + const int tid = threadIdx.x/K_QUANTS_PER_ITERATION; // 0...31 or 0...15 + const int ix = threadIdx.x%K_QUANTS_PER_ITERATION; // 0 or 0,1 + + const int step = 16/K_QUANTS_PER_ITERATION; + + const int im = tid/step; // 0 or 1. 0 computes 0..., 1 computes 128... + const int in = tid - step*im; // 0...15 or 0...7 + + const int l0 = K_QUANTS_PER_ITERATION*in; // 0...15 or 0...14 in steps of 2 + const int q_offset = 32*im + l0; + const int s_offset = 8*im; + const int y_offset = 128*im + l0; + + uint32_t aux[4]; + const uint8_t * d = (const uint8_t *)aux; + const uint8_t * m = (const uint8_t *)(aux + 2); + + for (int i = ix; i < num_blocks_per_row; i += K_QUANTS_PER_ITERATION) { + + const float * y = yy + i * QK_K + y_offset; + const uint8_t * q = x[i].qs + q_offset; + + const float dall = __low2half(x[i].dm); + const float dmin = __high2half(x[i].dm); + + const uint32_t * a = (const uint32_t *)(x[i].scales + s_offset); + aux[0] = a[0] & 0x0f0f0f0f; + aux[1] = a[1] & 0x0f0f0f0f; + aux[2] = (a[0] >> 4) & 0x0f0f0f0f; + aux[3] = (a[1] >> 4) & 0x0f0f0f0f; + + float sum1 = 0, sum2 = 0; + for (int l = 0; l < K_QUANTS_PER_ITERATION; ++l) { + sum1 += y[l+ 0] * d[0] * ((q[l+ 0] >> 0) & 3) + + y[l+32] * d[2] * ((q[l+ 0] >> 2) & 3) + + y[l+64] * d[4] * ((q[l+ 0] >> 4) & 3) + + y[l+96] * d[6] * ((q[l+ 0] >> 6) & 3) + + y[l+16] * d[1] * ((q[l+16] >> 0) & 3) + + y[l+48] * d[3] * ((q[l+16] >> 2) & 3) + + y[l+80] * d[5] * ((q[l+16] >> 4) & 3) + +y[l+112] * d[7] * ((q[l+16] >> 6) & 3); + sum2 += y[l+ 0] * m[0] + y[l+32] * m[2] + y[l+64] * m[4] + y[ l+96] * m[6] + + y[l+16] * m[1] + y[l+48] * m[3] + y[l+80] * m[5] + y[l+112] * m[7]; + + } + tmp += dall * sum1 - dmin * sum2; + + } +#else + const int tid = threadIdx.x/(2*K_QUANTS_PER_ITERATION); // 0...15 or 0...7 + const int ix = threadIdx.x%(2*K_QUANTS_PER_ITERATION); // 0....1 or 0...3 + const int offset = tid * K_QUANTS_PER_ITERATION; + + uint32_t uaux[2]; + const uint8_t * d = (const uint8_t *)uaux; + + for (int i = ix; i < num_blocks_per_row; i += 2*K_QUANTS_PER_ITERATION) { + + const float * y = yy + i * QK_K + offset; + const uint8_t * q = x[i].qs + offset; + const uint32_t * s = (const uint32_t *)x[i].scales; + + uaux[0] = s[0] & 0x0f0f0f0f; + uaux[1] = (s[0] >> 4) & 0x0f0f0f0f; + + const float2 dall = __half22float2(x[i].dm); + + float sum1 = 0, sum2 = 0; + for (int l = 0; l < K_QUANTS_PER_ITERATION; ++l) { + const uint8_t ql = q[l]; + sum1 += y[l+ 0] * d[0] * ((ql >> 0) & 3) + + y[l+16] * d[1] * ((ql >> 2) & 3) + + y[l+32] * d[2] * ((ql >> 4) & 3) + + y[l+48] * d[3] * ((ql >> 6) & 3); + sum2 += y[l+0] * d[4] + y[l+16] * d[5] + y[l+32] * d[6] + y[l+48] * d[7]; + } + tmp += dall.x * sum1 - dall.y * sum2; + } +#endif + + // sum up partial sums and write back result +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + tmp += __shfl_xor_sync(0xffffffff, tmp, mask, 32); + } + + if (threadIdx.x == 0) { + dst[row] = tmp; + } +} + +extern "C" __global__ void dequantize_mul_mat_vec_q3_k(const void * __restrict__ vx, const float * __restrict__ yy, float * __restrict__ dst, const int ncols, int nrows) { + + const int row = blockIdx.x*blockDim.y + threadIdx.y; + if (row > nrows) return; + + const int num_blocks_per_row = ncols / QK_K; + const int ib0 = row*num_blocks_per_row; + + const block_q3_K * x = (const block_q3_K *)vx + ib0; + + float tmp = 0; // partial sum for thread in warp + +#if QK_K == 256 + + const uint16_t kmask1 = 0x0303; + const uint16_t kmask2 = 0x0f0f; + + const int tid = threadIdx.x/K_QUANTS_PER_ITERATION; // 0...31 or 0...16 + const int ix = threadIdx.x%K_QUANTS_PER_ITERATION; // 0 or 0,1 + + const int n = K_QUANTS_PER_ITERATION; // iterations in the inner loop + const int step = 16/K_QUANTS_PER_ITERATION; + const int im = tid/step; // 0 or 1. 0 computes 0..., 1 computes 128... + const int in = tid - step*im; // 0....15 or 0...7 + + const uint8_t m = 1 << (4*im); + + const int l0 = n*in; // 0...15 or 0...14 in steps of 2 + const int q_offset = 32*im + l0; + const int y_offset = 128*im + l0; + + uint16_t utmp[4]; + const int8_t * s = (const int8_t *)utmp; + + const uint16_t s_shift = 4*im; + + for (int i = ix; i < num_blocks_per_row; i += K_QUANTS_PER_ITERATION) { + + const float * y = yy + i * QK_K + y_offset; + const uint8_t * q = x[i].qs + q_offset; + const uint8_t * h = x[i].hmask + l0; + + const uint16_t * a = (const uint16_t *)x[i].scales; + utmp[0] = ((a[0] >> s_shift) & kmask2) | (((a[4] >> (s_shift + 0)) & kmask1) << 4); + utmp[1] = ((a[1] >> s_shift) & kmask2) | (((a[5] >> (s_shift + 0)) & kmask1) << 4); + utmp[2] = ((a[2] >> s_shift) & kmask2) | (((a[4] >> (s_shift + 2)) & kmask1) << 4); + utmp[3] = ((a[3] >> s_shift) & kmask2) | (((a[5] >> (s_shift + 2)) & kmask1) << 4); + + const float d = x[i].d; + + float sum = 0; + for (int l = 0; l < n; ++l) { + sum += y[l+ 0] * (s[0] - 32) * (((q[l] >> 0) & 3) - (h[l] & (m << 0) ? 0 : 4)) + + y[l+32] * (s[2] - 32) * (((q[l] >> 2) & 3) - (h[l] & (m << 1) ? 0 : 4)) + + y[l+64] * (s[4] - 32) * (((q[l] >> 4) & 3) - (h[l] & (m << 2) ? 0 : 4)) + + y[l+96] * (s[6] - 32) * (((q[l] >> 6) & 3) - (h[l] & (m << 3) ? 0 : 4)); + sum += y[l+16] * (s[1] - 32) * (((q[l+16] >> 0) & 3) - (h[l+16] & (m << 0) ? 0 : 4)) + + y[l+48] * (s[3] - 32) * (((q[l+16] >> 2) & 3) - (h[l+16] & (m << 1) ? 0 : 4)) + + y[l+80] * (s[5] - 32) * (((q[l+16] >> 4) & 3) - (h[l+16] & (m << 2) ? 0 : 4)) + + y[l+112] * (s[7] - 32) * (((q[l+16] >> 6) & 3) - (h[l+16] & (m << 3) ? 0 : 4)); + } + tmp += d * sum; + + } +#else + + const int tid = threadIdx.x/(2*K_QUANTS_PER_ITERATION); // 0...15 or 0...7 + const int ix = threadIdx.x%(2*K_QUANTS_PER_ITERATION); // 0....1 or 0...3 + const int offset = tid * K_QUANTS_PER_ITERATION; // 0...15 or 0...14 + const int in = offset/8; // 0 or 1 + const int im = offset%8; // 0...7 + + for (int i = ix; i < num_blocks_per_row; i += 2*K_QUANTS_PER_ITERATION) { + + const float * y = yy + i * QK_K + offset; + const uint8_t * q = x[i].qs + offset; + const uint8_t * s = x[i].scales; + + const float dall = (float)x[i].d; + + float sum = 0; + for (int l = 0; l < K_QUANTS_PER_ITERATION; ++l) { + const uint8_t hl = x[i].hmask[im+l] >> in; + const uint8_t ql = q[l]; + sum += y[l+ 0] * dall * ((s[0] & 0xF) - 8) * ((int8_t)((ql >> 0) & 3) - ((hl >> 0) & 1 ? 0 : 4)) + + y[l+16] * dall * ((s[0] >> 4) - 8) * ((int8_t)((ql >> 2) & 3) - ((hl >> 2) & 1 ? 0 : 4)) + + y[l+32] * dall * ((s[1] & 0xF) - 8) * ((int8_t)((ql >> 4) & 3) - ((hl >> 4) & 1 ? 0 : 4)) + + y[l+48] * dall * ((s[1] >> 4) - 8) * ((int8_t)((ql >> 6) & 3) - ((hl >> 6) & 1 ? 0 : 4)); + } + tmp += sum; + } +#endif + + // sum up partial sums and write back result +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + tmp += __shfl_xor_sync(0xffffffff, tmp, mask, 32); + } + + if (threadIdx.x == 0) { + dst[row] = tmp; + } +} + +extern "C" __global__ void dequantize_mul_mat_vec_q4_k(const void * __restrict__ vx, const float * __restrict__ yy, float * __restrict__ dst, const int ncols, int nrows) { + + const int row = blockIdx.x*blockDim.y + threadIdx.y; + if (row > nrows) return; + const int num_blocks_per_row = ncols / QK_K; + const int ib0 = row*num_blocks_per_row; + + const block_q4_K * x = (const block_q4_K *)vx + ib0; + +#if QK_K == 256 + const uint16_t kmask1 = 0x3f3f; + const uint16_t kmask2 = 0x0f0f; + const uint16_t kmask3 = 0xc0c0; + + const int tid = threadIdx.x/K_QUANTS_PER_ITERATION; // 0...31 or 0...16 + const int ix = threadIdx.x%K_QUANTS_PER_ITERATION; // 0 or 0,1 + + const int step = 8/K_QUANTS_PER_ITERATION; // 8 or 4 + + const int il = tid/step; // 0...3 + const int ir = tid - step*il; // 0...7 or 0...3 + const int n = 2 * K_QUANTS_PER_ITERATION; // 2 or 4 + + const int im = il/2; // 0 or 1. 0 computes 0,32 + 128,160, 1 computes 64,96 + 192,224 + const int in = il%2; + + const int l0 = n*(2*ir + in); + const int q_offset = 32*im + l0; + const int y_offset = 64*im + l0; + + uint16_t aux[4]; + const uint8_t * sc = (const uint8_t *)aux; + +#if K_QUANTS_PER_ITERATION == 2 + uint32_t q32[4]; + const uint8_t * q4 = (const uint8_t *)q32; +#else + uint16_t q16[4]; + const uint8_t * q4 = (const uint8_t *)q16; +#endif + + float tmp = 0; // partial sum for thread in warp + + for (int i = ix; i < num_blocks_per_row; i += K_QUANTS_PER_ITERATION) { + + const float * y1 = yy + i*QK_K + y_offset; + const float * y2 = y1 + 128; + + const float dall = __low2half(x[i].dm); + const float dmin = __high2half(x[i].dm); + + const uint16_t * a = (const uint16_t *)x[i].scales; + aux[0] = a[im+0] & kmask1; + aux[1] = a[im+2] & kmask1; + aux[2] = ((a[im+4] >> 0) & kmask2) | ((a[im+0] & kmask3) >> 2); + aux[3] = ((a[im+4] >> 4) & kmask2) | ((a[im+2] & kmask3) >> 2); + +#if K_QUANTS_PER_ITERATION == 2 + const uint32_t * q1 = (const uint32_t *)(x[i].qs + q_offset); + const uint32_t * q2 = q1 + 16; + + q32[0] = q1[0] & 0x0f0f0f0f; + q32[1] = q1[0] & 0xf0f0f0f0; + q32[2] = q2[0] & 0x0f0f0f0f; + q32[3] = q2[0] & 0xf0f0f0f0; + + float4 s = {0.f, 0.f, 0.f, 0.f}; + float smin = 0; + for (int l = 0; l < 4; ++l) { + s.x += y1[l] * q4[l+0]; s.y += y1[l+32] * q4[l+ 4]; + s.z += y2[l] * q4[l+8]; s.w += y2[l+32] * q4[l+12]; + smin += y1[l] * sc[2] + y1[l+32] * sc[3] + y2[l] * sc[6] + y2[l+32] * sc[7]; + } + tmp += dall * (s.x * sc[0] + s.y * sc[1] * 1.f/16.f + s.z * sc[4] + s.w * sc[5] * 1.f/16.f) - dmin * smin; +#else + const uint16_t * q1 = (const uint16_t *)(x[i].qs + q_offset); + const uint16_t * q2 = q1 + 32; + + q16[0] = q1[0] & 0x0f0f; + q16[1] = q1[0] & 0xf0f0; + q16[2] = q2[0] & 0x0f0f; + q16[3] = q2[0] & 0xf0f0; + + float4 s = {0.f, 0.f, 0.f, 0.f}; + float smin = 0; + for (int l = 0; l < 2; ++l) { + s.x += y1[l] * q4[l+0]; s.y += y1[l+32] * q4[l+2]; + s.z += y2[l] * q4[l+4]; s.w += y2[l+32] * q4[l+6]; + smin += y1[l] * sc[2] + y1[l+32] * sc[3] + y2[l] * sc[6] + y2[l+32] * sc[7]; + } + tmp += dall * (s.x * sc[0] + s.y * sc[1] * 1.f/16.f + s.z * sc[4] + s.w * sc[5] * 1.f/16.f) - dmin * smin; +#endif + + } +#else + const int tid = threadIdx.x/(2*K_QUANTS_PER_ITERATION); // 0...15 + const int ix = threadIdx.x%(2*K_QUANTS_PER_ITERATION); + + const int step = tid * K_QUANTS_PER_ITERATION; + + uint16_t aux16[2]; + const uint8_t * s = (const uint8_t *)aux16; + + float tmp = 0; + + for (int i = ix; i < num_blocks_per_row; i += 2*K_QUANTS_PER_ITERATION) { + const uint8_t * q = x[i].qs + step; + const float * y = yy + i*QK_K + step; + const uint16_t * a = (const uint16_t *)x[i].scales; + aux16[0] = a[0] & 0x0f0f; + aux16[1] = (a[0] >> 4) & 0x0f0f; + const float d = (float)x[i].dm[0]; + const float m = (float)x[i].dm[1]; + float sum = 0.f; + for (int j = 0; j < K_QUANTS_PER_ITERATION; ++j) { + sum += y[j+ 0] * (d * s[0] * (q[j+ 0] & 0xF) - m * s[2]) + + y[j+16] * (d * s[0] * (q[j+16] & 0xF) - m * s[2]) + + y[j+32] * (d * s[1] * (q[j+ 0] >> 4) - m * s[3]) + + y[j+48] * (d * s[1] * (q[j+16] >> 4) - m * s[3]); + } + tmp += sum; + } + +#endif + + // sum up partial sums and write back result +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + tmp += __shfl_xor_sync(0xffffffff, tmp, mask, 32); + } + + if (tid == 0) { + dst[row] = tmp; + } +} + +extern "C" __global__ void dequantize_mul_mat_vec_q5_k(const void * __restrict__ vx, const float * __restrict__ yy, float * __restrict__ dst, const int ncols) { + + const int row = blockIdx.x; + const int num_blocks_per_row = ncols / QK_K; + const int ib0 = row*num_blocks_per_row; + + const block_q5_K * x = (const block_q5_K *)vx + ib0; + + float tmp = 0; // partial sum for thread in warp + +#if QK_K == 256 + const uint16_t kmask1 = 0x3f3f; + const uint16_t kmask2 = 0x0f0f; + const uint16_t kmask3 = 0xc0c0; + + const int tid = threadIdx.x/2; // 0...15 + const int ix = threadIdx.x%2; + + const int il = tid/4; // 0...3 + const int ir = tid - 4*il;// 0...3 + const int n = 2; + + const int im = il/2; // 0 or 1. 0 computes 0,32 + 128,160, 1 computes 64,96 + 192,224 + const int in = il%2; + + const int l0 = n*(2*ir + in); + const int q_offset = 32*im + l0; + const int y_offset = 64*im + l0; + + const uint8_t hm1 = 1 << (2*im); + const uint8_t hm2 = hm1 << 4; + + uint16_t aux[4]; + const uint8_t * sc = (const uint8_t *)aux; + + uint16_t q16[8]; + const uint8_t * q4 = (const uint8_t *)q16; + + for (int i = ix; i < num_blocks_per_row; i += 2) { + + const uint8_t * ql1 = x[i].qs + q_offset; + const uint8_t * qh = x[i].qh + l0; + const float * y1 = yy + i*QK_K + y_offset; + const float * y2 = y1 + 128; + + const float dall = __low2half(x[i].dm); + const float dmin = __high2half(x[i].dm); + + const uint16_t * a = (const uint16_t *)x[i].scales; + aux[0] = a[im+0] & kmask1; + aux[1] = a[im+2] & kmask1; + aux[2] = ((a[im+4] >> 0) & kmask2) | ((a[im+0] & kmask3) >> 2); + aux[3] = ((a[im+4] >> 4) & kmask2) | ((a[im+2] & kmask3) >> 2); + + float4 sum = {0.f, 0.f, 0.f, 0.f}; + float smin = 0; + const uint16_t * q1 = (const uint16_t *)ql1; + const uint16_t * q2 = q1 + 32; + q16[0] = q1[0] & 0x0f0f; + q16[1] = q1[8] & 0x0f0f; + q16[2] = (q1[0] >> 4) & 0x0f0f; + q16[3] = (q1[8] >> 4) & 0x0f0f; + q16[4] = q2[0] & 0x0f0f; + q16[5] = q2[8] & 0x0f0f; + q16[6] = (q2[0] >> 4) & 0x0f0f; + q16[7] = (q2[8] >> 4) & 0x0f0f; + for (int l = 0; l < n; ++l) { + sum.x += y1[l+ 0] * (q4[l +0] + (qh[l+ 0] & (hm1 << 0) ? 16 : 0)) + + y1[l+16] * (q4[l +2] + (qh[l+16] & (hm1 << 0) ? 16 : 0)); + sum.y += y1[l+32] * (q4[l +4] + (qh[l+ 0] & (hm1 << 1) ? 16 : 0)) + + y1[l+48] * (q4[l +6] + (qh[l+16] & (hm1 << 1) ? 16 : 0)); + sum.z += y2[l+ 0] * (q4[l +8] + (qh[l+ 0] & (hm2 << 0) ? 16 : 0)) + + y2[l+16] * (q4[l+10] + (qh[l+16] & (hm2 << 0) ? 16 : 0)); + sum.w += y2[l+32] * (q4[l+12] + (qh[l+ 0] & (hm2 << 1) ? 16 : 0)) + + y2[l+48] * (q4[l+14] + (qh[l+16] & (hm2 << 1) ? 16 : 0)); + smin += (y1[l] + y1[l+16]) * sc[2] + (y1[l+32] + y1[l+48]) * sc[3] + + (y2[l] + y2[l+16]) * sc[6] + (y2[l+32] + y2[l+48]) * sc[7]; + } + tmp += dall * (sum.x * sc[0] + sum.y * sc[1] + sum.z * sc[4] + sum.w * sc[5]) - dmin * smin; + } + +#else + const int tid = threadIdx.x/(2*K_QUANTS_PER_ITERATION); // 0...15 + const int ix = threadIdx.x%(2*K_QUANTS_PER_ITERATION); + const int step = tid * K_QUANTS_PER_ITERATION; + const int im = step/8; + const int in = step%8; + + for (int i = ix; i < num_blocks_per_row; i += 2*K_QUANTS_PER_ITERATION) { + const uint8_t * q = x[i].qs + step; + const int8_t * s = x[i].scales; + const float * y = yy + i*QK_K + step; + const float d = x[i].d; + float sum = 0.f; + for (int j = 0; j < K_QUANTS_PER_ITERATION; ++j) { + const uint8_t h = x[i].qh[in+j] >> im; + sum += y[j+ 0] * d * s[0] * ((q[j+ 0] & 0xF) - ((h >> 0) & 1 ? 0 : 16)) + + y[j+16] * d * s[1] * ((q[j+16] & 0xF) - ((h >> 2) & 1 ? 0 : 16)) + + y[j+32] * d * s[2] * ((q[j+ 0] >> 4) - ((h >> 4) & 1 ? 0 : 16)) + + y[j+48] * d * s[3] * ((q[j+16] >> 4) - ((h >> 6) & 1 ? 0 : 16)); + } + tmp += sum; + } +#endif + + // sum up partial sums and write back result +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + tmp += __shfl_xor_sync(0xffffffff, tmp, mask, 32); + } + + if (threadIdx.x == 0) { + dst[row] = tmp; + } +} + +extern "C" __global__ void dequantize_mul_mat_vec_q6_k(const void * __restrict__ vx, const float * __restrict__ yy, float * __restrict__ dst, const int ncols, int nrows) { + + static_assert(16%K_QUANTS_PER_ITERATION == 0, "16 must be divisible by K_QUANTS_PER_ITERATION"); + + const int row = blockIdx.x*blockDim.y + threadIdx.y; + if (row > nrows) return; + + const int num_blocks_per_row = ncols / QK_K; + const int ib0 = row*num_blocks_per_row; + + const block_q6_K * x = (const block_q6_K *)vx + ib0; + +#if QK_K == 256 + + const int tid = threadIdx.x/K_QUANTS_PER_ITERATION; // 0...31 or 0...16 + const int ix = threadIdx.x%K_QUANTS_PER_ITERATION; // 0 or 0, 1 + + const int step = 16/K_QUANTS_PER_ITERATION; // 16 or 8 + + const int im = tid/step; // 0 or 1. 0 computes 0..., 1 computes 128... + const int in = tid - step*im; // 0...15 or 0...7 + +#if K_QUANTS_PER_ITERATION == 1 + const int l0 = K_QUANTS_PER_ITERATION*in; // 0...15 + const int is = 0; +#else + const int l0 = 4 * in; // 0, 4, 8, ..., 28 + const int is = in / 4; +#endif + const int ql_offset = 64*im + l0; + const int qh_offset = 32*im + l0; + const int s_offset = 8*im + is; + const int y_offset = 128*im + l0; + + float tmp = 0; // partial sum for thread in warp + + for (int i = ix; i < num_blocks_per_row; i += K_QUANTS_PER_ITERATION) { + + const float * y = yy + i * QK_K + y_offset; + const uint8_t * ql = x[i].ql + ql_offset; + const uint8_t * qh = x[i].qh + qh_offset; + const int8_t * s = x[i].scales + s_offset; + + const float d = x[i].d; + +#if K_QUANTS_PER_ITERATION == 1 + float sum = y[ 0] * s[0] * d * ((int8_t)((ql[ 0] & 0xF) | ((qh[ 0] & 0x03) << 4)) - 32) + + y[16] * s[1] * d * ((int8_t)((ql[16] & 0xF) | ((qh[16] & 0x03) << 4)) - 32) + + y[32] * s[2] * d * ((int8_t)((ql[32] & 0xF) | ((qh[ 0] & 0x0c) << 2)) - 32) + + y[48] * s[3] * d * ((int8_t)((ql[48] & 0xF) | ((qh[16] & 0x0c) << 2)) - 32) + + y[64] * s[4] * d * ((int8_t)((ql[ 0] >> 4) | ((qh[ 0] & 0x30) >> 0)) - 32) + + y[80] * s[5] * d * ((int8_t)((ql[16] >> 4) | ((qh[16] & 0x30) >> 0)) - 32) + + y[96] * s[6] * d * ((int8_t)((ql[32] >> 4) | ((qh[ 0] & 0xc0) >> 2)) - 32) + +y[112] * s[7] * d * ((int8_t)((ql[48] >> 4) | ((qh[16] & 0xc0) >> 2)) - 32); + tmp += sum; +#else + float sum = 0; + for (int l = 0; l < 4; ++l) { + sum += y[l+ 0] * s[0] * d * ((int8_t)((ql[l+ 0] & 0xF) | (((qh[l] >> 0) & 3) << 4)) - 32) + + y[l+32] * s[2] * d * ((int8_t)((ql[l+32] & 0xF) | (((qh[l] >> 2) & 3) << 4)) - 32) + + y[l+64] * s[4] * d * ((int8_t)((ql[l+ 0] >> 4) | (((qh[l] >> 4) & 3) << 4)) - 32) + + y[l+96] * s[6] * d * ((int8_t)((ql[l+32] >> 4) | (((qh[l] >> 6) & 3) << 4)) - 32); + } + tmp += sum; +#endif + + } + +#else + + const int tid = threadIdx.x/(2*K_QUANTS_PER_ITERATION); // 0...7 + const int ix = threadIdx.x%(2*K_QUANTS_PER_ITERATION); // 0...3 + + const int step = tid * K_QUANTS_PER_ITERATION; + + float tmp = 0; // partial sum for thread in warp + + for (int i = ix; i < num_blocks_per_row; i += 2*K_QUANTS_PER_ITERATION) { + + const float * y = yy + i * QK_K + step; + const uint8_t * ql = x[i].ql + step; + const uint8_t * qh = x[i].qh + step; + const int8_t * s = x[i].scales; + + const float d = x[i+0].d; + + float sum = 0; + for (int j = 0; j < K_QUANTS_PER_ITERATION; ++j) { + sum += y[j+ 0] * s[0] * d * ((int8_t)((ql[j+ 0] & 0xF) | ((qh[j] & 0x03) << 4)) - 32) + + y[j+16] * s[1] * d * ((int8_t)((ql[j+16] & 0xF) | ((qh[j] & 0x0c) << 2)) - 32) + + y[j+32] * s[2] * d * ((int8_t)((ql[j+ 0] >> 4) | ((qh[j] & 0x30) >> 0)) - 32) + + y[j+48] * s[3] * d * ((int8_t)((ql[j+16] >> 4) | ((qh[j] & 0xc0) >> 2)) - 32); + } + tmp += sum; + + } + +#endif + + // sum up partial sums and write back result +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + tmp += __shfl_xor_sync(0xffffffff, tmp, mask, 32); + } + + if (tid == 0) { + dst[row] = tmp; + } +} + +// VDR = vec dot ratio, how many contiguous integers each thread processes when the vec dot kernel is called +// MMVQ = mul_mat_vec_q, MMQ = mul_mat_q + +#define VDR_Q4_0_Q8_1_MMVQ 2 +#define VDR_Q4_0_Q8_1_MMQ 4 + +template static __device__ __forceinline__ float vec_dot_q4_0_q8_1_impl( + const int * v, const int * u, const float & d4, const half2 & ds8) { + + int sumi = 0; + +#pragma unroll + for (int i = 0; i < vdr; ++i) { + const int vi0 = (v[i] >> 0) & 0x0F0F0F0F; + const int vi1 = (v[i] >> 4) & 0x0F0F0F0F; + + // SIMD dot product of quantized values + sumi = ggml_cuda_dp4a(vi0, u[2*i+0], sumi); + sumi = ggml_cuda_dp4a(vi1, u[2*i+1], sumi); + } + + const float2 ds8f = __half22float2(ds8); + + // second part effectively subtracts 8 from each quant value + return d4 * (sumi * ds8f.x - (8*vdr/QI4_0) * ds8f.y); +} + +#define VDR_Q4_1_Q8_1_MMVQ 2 +#define VDR_Q4_1_Q8_1_MMQ 4 + +template static __device__ __forceinline__ float vec_dot_q4_1_q8_1_impl( + const int * v, const int * u, const half2 & dm4, const half2 & ds8) { + int sumi = 0; + +#pragma unroll + for (int i = 0; i < vdr; ++i) { + const int vi0 = (v[i] >> 0) & 0x0F0F0F0F; + const int vi1 = (v[i] >> 4) & 0x0F0F0F0F; + + // SIMD dot product of quantized values + sumi = ggml_cuda_dp4a(vi0, u[2*i+0], sumi); + sumi = ggml_cuda_dp4a(vi1, u[2*i+1], sumi); + } + +#ifdef GGML_CUDA_F16 + const float2 tmp = __half22float2(__hmul2(dm4, ds8)); + const float d4d8 = tmp.x; + const float m4s8 = tmp.y; +#else + const float2 dm4f = __half22float2(dm4); + const float2 ds8f = __half22float2(ds8); + const float d4d8 = dm4f.x * ds8f.x; + const float m4s8 = dm4f.y * ds8f.y; +#endif // GGML_CUDA_F16 + + // scale second part of sum by QI8_1/(vdr * QR4_1) to compensate for multiple threads adding it + return sumi * d4d8 + m4s8 / (QI8_1 / (vdr * QR4_1)); +} + +#define VDR_Q5_0_Q8_1_MMVQ 2 +#define VDR_Q5_0_Q8_1_MMQ 4 + +template static __device__ __forceinline__ float vec_dot_q5_0_q8_1_impl( + const int * vl, const int * vh, const int * u, const float & d5, const half2 & ds8) { + + int sumi = 0; + +#pragma unroll + for (int i = 0; i < vdr; ++i) { + int vi0 = (vl[i] >> 0) & 0x0F0F0F0F; // lower 4 qs bits, still need qh as 5th bits + vi0 |= (vh[i] << 4) & 0x00000010; // 0 -> 4 + vi0 |= (vh[i] << 11) & 0x00001000; // 1 -> 12 + vi0 |= (vh[i] << 18) & 0x00100000; // 2 -> 20 + vi0 |= (vh[i] << 25) & 0x10000000; // 3 -> 28 + sumi = ggml_cuda_dp4a(vi0, u[2*i+0], sumi); // SIMD dot product of quantized values + + int vi1 = (vl[i] >> 4) & 0x0F0F0F0F; // upper 4 qs bits, still need qh as 5th bits + vi1 |= (vh[i] >> 12) & 0x00000010; // 16 -> 4 + vi1 |= (vh[i] >> 5) & 0x00001000; // 17 -> 12 + vi1 |= (vh[i] << 2) & 0x00100000; // 18 -> 20 + vi1 |= (vh[i] << 9) & 0x10000000; // 19 -> 28 + sumi = ggml_cuda_dp4a(vi1, u[2*i+1], sumi); // SIMD dot product of quantized values + } + + const float2 ds8f = __half22float2(ds8); + + // second part effectively subtracts 16 from each quant value + return d5 * (sumi * ds8f.x - (16*vdr/QI5_0) * ds8f.y); +} + +#define VDR_Q5_1_Q8_1_MMVQ 2 +#define VDR_Q5_1_Q8_1_MMQ 4 + +template static __device__ __forceinline__ float vec_dot_q5_1_q8_1_impl( + const int * vl, const int * vh, const int * u, const half2 & dm5, const half2 & ds8) { + + int sumi = 0; + +#pragma unroll + for (int i = 0; i < vdr; ++i) { + int vi0 = (vl[i] >> 0) & 0x0F0F0F0F; // lower 4 qs bits, still need qh as 5th bits + vi0 |= (vh[i] << 4) & 0x00000010; // 0 -> 4 + vi0 |= (vh[i] << 11) & 0x00001000; // 1 -> 12 + vi0 |= (vh[i] << 18) & 0x00100000; // 2 -> 20 + vi0 |= (vh[i] << 25) & 0x10000000; // 3 -> 28 + sumi = ggml_cuda_dp4a(vi0, u[2*i+0], sumi); // SIMD dot product of quantized values + + int vi1 = (vl[i] >> 4) & 0x0F0F0F0F; // upper 4 qs bits, still need qh as 5th bits + vi1 |= (vh[i] >> 12) & 0x00000010; // 16 -> 4 + vi1 |= (vh[i] >> 5) & 0x00001000; // 17 -> 12 + vi1 |= (vh[i] << 2) & 0x00100000; // 18 -> 20 + vi1 |= (vh[i] << 9) & 0x10000000; // 19 -> 28 + sumi = ggml_cuda_dp4a(vi1, u[2*i+1], sumi); // SIMD dot product of quantized values + } + +#ifdef GGML_CUDA_F16 + const float2 tmp = __half22float2(__hmul2(dm5, ds8)); + const float d5d8 = tmp.x; + const float m5s8 = tmp.y; +#else + const float2 dm5f = __half22float2(dm5); + const float2 ds8f = __half22float2(ds8); + const float d5d8 = dm5f.x * ds8f.x; + const float m5s8 = dm5f.y * ds8f.y; +#endif // GGML_CUDA_F16 + + // scale second part of sum by QI5_1 / vdr to compensate for multiple threads adding it + return sumi*d5d8 + m5s8 / (QI5_1 / vdr); +} + +#define VDR_Q8_0_Q8_1_MMVQ 2 +#define VDR_Q8_0_Q8_1_MMQ 8 + +template static __device__ __forceinline__ float vec_dot_q8_0_q8_1_impl( + const int * v, const int * u, const float & d8_0, const float & d8_1) { + + int sumi = 0; + +#pragma unroll + for (int i = 0; i < vdr; ++i) { + // SIMD dot product of quantized values + sumi = ggml_cuda_dp4a(v[i], u[i], sumi); + } + + return d8_0*d8_1 * sumi; +} + +template static __device__ __forceinline__ float vec_dot_q8_1_q8_1_impl( + const int * v, const int * u, const half2 & dm8, const half2 & ds8) { + + int sumi = 0; + +#pragma unroll + for (int i = 0; i < vdr; ++i) { + // SIMD dot product of quantized values + sumi = ggml_cuda_dp4a(v[i], u[i], sumi); + } + +#ifdef GGML_CUDA_F16 + const float2 tmp = __half22float2(__hmul2(dm8, ds8)); + const float d8d8 = tmp.x; + const float m8s8 = tmp.y; +#else + const float2 dm8f = __half22float2(dm8); + const float2 ds8f = __half22float2(ds8); + const float d8d8 = dm8f.x * ds8f.x; + const float m8s8 = dm8f.y * ds8f.y; +#endif // GGML_CUDA_F16 + + // scale second part of sum by QI8_1/ vdr to compensate for multiple threads adding it + return sumi*d8d8 + m8s8 / (QI8_1 / vdr); +} + +#define VDR_Q2_K_Q8_1_MMVQ 1 +#define VDR_Q2_K_Q8_1_MMQ 2 + +// contiguous v/x values +static __device__ __forceinline__ float vec_dot_q2_K_q8_1_impl_mmvq( + const int & v, const int * __restrict__ u, const uint8_t * __restrict__ scales, + const half2 & dm2, const float * __restrict__ d8) { + + float sumf_d = 0.0f; + float sumf_m = 0.0f; + +#pragma unroll + for (int i = 0; i < QR2_K; ++i) { + const int sc = scales[2*i]; + + const int vi = (v >> (2*i)) & 0x03030303; + + sumf_d += d8[i] * (ggml_cuda_dp4a(vi, u[i], 0) * (sc & 0xF)); // SIMD dot product + + // fill int with 4x m + int m = sc >> 4; + m |= m << 8; + m |= m << 16; + sumf_m += d8[i] * ggml_cuda_dp4a(m, u[i], 0); // multiply constant q2_K part with sum of q8_1 values + } + + const float2 dm2f = __half22float2(dm2); + + return dm2f.x*sumf_d - dm2f.y*sumf_m; +} + +// contiguous u/y values +static __device__ __forceinline__ float vec_dot_q2_K_q8_1_impl_mmq( + const int * __restrict__ v, const int * __restrict__ u, const uint8_t * __restrict__ scales, + const half2 & dm2, const float & d8) { + + int sumi_d = 0; + int sumi_m = 0; + +#pragma unroll + for (int i0 = 0; i0 < QI8_1; i0 += QI8_1/2) { + int sumi_d_sc = 0; + + const int sc = scales[i0 / (QI8_1/2)]; + + // fill int with 4x m + int m = sc >> 4; + m |= m << 8; + m |= m << 16; + +#pragma unroll + for (int i = i0; i < i0 + QI8_1/2; ++i) { + sumi_d_sc = ggml_cuda_dp4a(v[i], u[i], sumi_d_sc); // SIMD dot product + sumi_m = ggml_cuda_dp4a(m, u[i], sumi_m); // multiply sum of q8_1 values with m + } + + sumi_d += sumi_d_sc * (sc & 0xF); + } + + const float2 dm2f = __half22float2(dm2); + + return d8 * (dm2f.x*sumi_d - dm2f.y*sumi_m); +} + +#define VDR_Q3_K_Q8_1_MMVQ 1 +#define VDR_Q3_K_Q8_1_MMQ 2 + +// contiguous v/x values +static __device__ __forceinline__ float vec_dot_q3_K_q8_1_impl_mmvq( + const int & vl, const int & vh, const int * __restrict__ u, const uint8_t * __restrict__ scales, + const int & scale_offset, const float & d3, const float * __restrict__ d8) { + + float sumf = 0.0f; + +#pragma unroll + for (int i = 0; i < QR3_K; ++i) { + const int isc = scale_offset + 2*i; + + const int isc_low = isc % (QK_K/32); + const int sc_shift_low = 4 * (isc / (QK_K/32)); + const int sc_low = (scales[isc_low] >> sc_shift_low) & 0xF; + + const int isc_high = isc % (QK_K/64); + const int sc_shift_high = 2 * (isc / (QK_K/64)); + const int sc_high = ((scales[(QK_K/32) + isc_high] >> sc_shift_high) & 3) << 4; + + const int sc = (sc_low | sc_high) - 32; + + const int vil = (vl >> (2*i)) & 0x03030303; + + const int vih = ((vh >> i) << 2) & 0x04040404; + + const int vi = __vsubss4(vil, vih); + + sumf += d8[i] * (ggml_cuda_dp4a(vi, u[i], 0) * sc); // SIMD dot product + } + + return d3 * sumf; +} + +// contiguous u/y values +static __device__ __forceinline__ float vec_dot_q3_K_q8_1_impl_mmq( + const int * __restrict__ v, const int * __restrict__ u, const int8_t * __restrict__ scales, + const float & d3, const float & d8) { + + int sumi = 0; + +#pragma unroll + for (int i0 = 0; i0 < QR3_K*VDR_Q3_K_Q8_1_MMQ; i0 += QI8_1/2) { + int sumi_sc = 0; + + for (int i = i0; i < i0 + QI8_1/2; ++i) { + sumi_sc = ggml_cuda_dp4a(v[i], u[i], sumi_sc); // SIMD dot product + } + + sumi += sumi_sc * scales[i0 / (QI8_1/2)]; + } + + return d3*d8 * sumi; +} + +#define VDR_Q4_K_Q8_1_MMVQ 2 +#define VDR_Q4_K_Q8_1_MMQ 8 + +// contiguous v/x values +static __device__ __forceinline__ float vec_dot_q4_K_q8_1_impl_vmmq( + const int * __restrict__ v, const int * __restrict__ u, const uint8_t * __restrict__ sc, + const uint8_t * __restrict__ m, const half2 & dm4, const float * __restrict__ d8) { + + float sumf_d = 0.0f; + float sumf_m = 0.0f; + +#pragma unroll + for (int i = 0; i < QR4_K; ++i) { + const int v0i = (v[0] >> (4*i)) & 0x0F0F0F0F; + const int v1i = (v[1] >> (4*i)) & 0x0F0F0F0F; + + const int dot1 = ggml_cuda_dp4a(v1i, u[2*i+1], ggml_cuda_dp4a(v0i, u[2*i+0], 0)); // SIMD dot product + const int dot2 = ggml_cuda_dp4a(0x01010101, u[2*i+1], ggml_cuda_dp4a(0x01010101, u[2*i+0], 0)); // sum of u + + sumf_d += d8[i] * (dot1 * sc[i]); + sumf_m += d8[i] * (dot2 * m[i]); // multiply constant part of q4_K with sum of q8_1 values + } + + const float2 dm4f = __half22float2(dm4); + + return dm4f.x*sumf_d - dm4f.y*sumf_m; +} + +// contiguous u/y values +static __device__ __forceinline__ float vec_dot_q4_K_q8_1_impl_mmq( + const int * __restrict__ v, const int * __restrict__ u, const uint8_t * __restrict__ sc, + const uint8_t * __restrict__ m, const half2 & dm4, const half2 * __restrict__ ds8) { + + float sumf_d = 0.0f; + float sumf_m = 0.0f; + +#pragma unroll + for (int i = 0; i < QR4_K*VDR_Q4_K_Q8_1_MMQ/QI8_1; ++i) { + int sumi_d = 0; + +#pragma unroll + for (int j = 0; j < QI8_1; ++j) { + sumi_d = ggml_cuda_dp4a((v[j] >> (4*i)) & 0x0F0F0F0F, u[i*QI8_1 + j], sumi_d); // SIMD dot product + } + + const float2 ds8f = __half22float2(ds8[i]); + + sumf_d += ds8f.x * (sc[i] * sumi_d); + sumf_m += ds8f.y * m[i]; // sum of q8_1 block * q4_K min val + } + + const float2 dm4f = __half22float2(dm4); + + return dm4f.x*sumf_d - dm4f.y*sumf_m; +} + +#define VDR_Q5_K_Q8_1_MMVQ 2 +#define VDR_Q5_K_Q8_1_MMQ 8 + +// contiguous v/x values +static __device__ __forceinline__ float vec_dot_q5_K_q8_1_impl_vmmq( + const int * __restrict__ vl, const int * __restrict__ vh, const int * __restrict__ u, const uint8_t * __restrict__ sc, + const uint8_t * __restrict__ m, const half2 & dm5, const float * __restrict__ d8) { + + float sumf_d = 0.0f; + float sumf_m = 0.0f; + +#pragma unroll + for (int i = 0; i < QR5_K; ++i) { + const int vl0i = (vl[0] >> (4*i)) & 0x0F0F0F0F; + const int vl1i = (vl[1] >> (4*i)) & 0x0F0F0F0F; + + const int vh0i = ((vh[0] >> i) << 4) & 0x10101010; + const int vh1i = ((vh[1] >> i) << 4) & 0x10101010; + + const int v0i = vl0i | vh0i; + const int v1i = vl1i | vh1i; + + const int dot1 = ggml_cuda_dp4a(v0i, u[2*i+0], ggml_cuda_dp4a(v1i, u[2*i+1], 0)); // SIMD dot product + const int dot2 = ggml_cuda_dp4a(0x01010101, u[2*i+0], ggml_cuda_dp4a(0x01010101, u[2*i+1], 0)); // sum of u + + sumf_d += d8[i] * (dot1 * sc[i]); + sumf_m += d8[i] * (dot2 * m[i]); + + } + + const float2 dm5f = __half22float2(dm5); + + return dm5f.x*sumf_d - dm5f.y*sumf_m; +} + +// contiguous u/y values +static __device__ __forceinline__ float vec_dot_q5_K_q8_1_impl_mmq( + const int * __restrict__ v, const int * __restrict__ u, const uint8_t * __restrict__ sc, + const uint8_t * __restrict__ m, const half2 & dm4, const half2 * __restrict__ ds8) { + + float sumf_d = 0.0f; + float sumf_m = 0.0f; + +#pragma unroll + for (int i = 0; i < QR5_K*VDR_Q5_K_Q8_1_MMQ/QI8_1; ++i) { + int sumi_d = 0; + +#pragma unroll + for (int j = 0; j < QI8_1; ++j) { + sumi_d = ggml_cuda_dp4a(v[i*QI8_1 + j], u[i*QI8_1 + j], sumi_d); // SIMD dot product + } + + const float2 ds8f = __half22float2(ds8[i]); + + sumf_d += ds8f.x * (sc[i] * sumi_d); + sumf_m += ds8f.y * m[i]; // sum of q8_1 block * q4_K min val + } + + const float2 dm4f = __half22float2(dm4); + + return dm4f.x*sumf_d - dm4f.y*sumf_m; +} + +#define VDR_Q6_K_Q8_1_MMVQ 1 +#define VDR_Q6_K_Q8_1_MMQ 8 + +// contiguous v/x values +static __device__ __forceinline__ float vec_dot_q6_K_q8_1_impl_mmvq( + const int & vl, const int & vh, const int * __restrict__ u, const int8_t * __restrict__ scales, + const float & d, const float * __restrict__ d8) { + + float sumf = 0.0f; + +#pragma unroll + for (int i = 0; i < QR6_K; ++i) { + const int sc = scales[4*i]; + + const int vil = (vl >> (4*i)) & 0x0F0F0F0F; + + const int vih = ((vh >> (4*i)) << 4) & 0x30303030; + + const int vi = __vsubss4((vil | vih), 0x20202020); // vi = (vil | vih) - 32 + + sumf += d8[i] * (ggml_cuda_dp4a(vi, u[i], 0) * sc); // SIMD dot product + } + + return d*sumf; +} + +// contiguous u/y values +static __device__ __forceinline__ float vec_dot_q6_K_q8_1_impl_mmq( + const int * __restrict__ v, const int * __restrict__ u, const int8_t * __restrict__ sc, + const float & d6, const float * __restrict__ d8) { + + float sumf_d = 0.0f; + +#pragma unroll + for (int i0 = 0; i0 < VDR_Q6_K_Q8_1_MMQ; i0 += 4) { + int2 sumi_d = {0, 0}; // 2 q6_K scales per q8_1 scale + +#pragma unroll + for (int i = i0; i < i0 + 2; ++i) { + sumi_d.x = ggml_cuda_dp4a(v[2*i+0], u[2*i+0], sumi_d.x); // SIMD dot product + sumi_d.x = ggml_cuda_dp4a(v[2*i+1], u[2*i+1], sumi_d.x); // SIMD dot product + + sumi_d.y = ggml_cuda_dp4a(v[2*i+4], u[2*i+4], sumi_d.y); // SIMD dot product + sumi_d.y = ggml_cuda_dp4a(v[2*i+5], u[2*i+5], sumi_d.y); // SIMD dot product + } + + sumf_d += d8[i0/4] * (sc[i0/2+0]*sumi_d.x + sc[i0/2+1]*sumi_d.y); + } + + return d6 * sumf_d; +} + +static __device__ __forceinline__ float vec_dot_q4_0_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { + + const block_q4_0 * bq4_0 = (const block_q4_0 *) vbq; + + int v[VDR_Q4_0_Q8_1_MMVQ]; + int u[2*VDR_Q4_0_Q8_1_MMVQ]; + +#pragma unroll + for (int i = 0; i < VDR_Q4_0_Q8_1_MMVQ; ++i) { + v[i] = get_int_from_uint8(bq4_0->qs, iqs + i); + u[2*i+0] = get_int_from_int8_aligned(bq8_1->qs, iqs + i); + u[2*i+1] = get_int_from_int8_aligned(bq8_1->qs, iqs + i + QI4_0); + } + + return vec_dot_q4_0_q8_1_impl(v, u, bq4_0->d, bq8_1->ds); +} + + +static __device__ __forceinline__ float vec_dot_q4_1_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { + + const block_q4_1 * bq4_1 = (const block_q4_1 *) vbq; + + int v[VDR_Q4_1_Q8_1_MMVQ]; + int u[2*VDR_Q4_1_Q8_1_MMVQ]; + +#pragma unroll + for (int i = 0; i < VDR_Q4_1_Q8_1_MMVQ; ++i) { + v[i] = get_int_from_uint8_aligned(bq4_1->qs, iqs + i); + u[2*i+0] = get_int_from_int8_aligned(bq8_1->qs, iqs + i); + u[2*i+1] = get_int_from_int8_aligned(bq8_1->qs, iqs + i + QI4_1); + } + + return vec_dot_q4_1_q8_1_impl(v, u, bq4_1->dm, bq8_1->ds); +} + +static __device__ __forceinline__ float vec_dot_q5_0_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { + + const block_q5_0 * bq5_0 = (const block_q5_0 *) vbq; + + int vl[VDR_Q5_0_Q8_1_MMVQ]; + int vh[VDR_Q5_0_Q8_1_MMVQ]; + int u[2*VDR_Q5_0_Q8_1_MMVQ]; + +#pragma unroll + for (int i = 0; i < VDR_Q5_0_Q8_1_MMVQ; ++i) { + vl[i] = get_int_from_uint8(bq5_0->qs, iqs + i); + vh[i] = get_int_from_uint8(bq5_0->qh, 0) >> (4 * (iqs + i)); + u[2*i+0] = get_int_from_int8_aligned(bq8_1->qs, iqs + i); + u[2*i+1] = get_int_from_int8_aligned(bq8_1->qs, iqs + i + QI5_0); + } + + return vec_dot_q5_0_q8_1_impl(vl, vh, u, bq5_0->d, bq8_1->ds); +} + +static __device__ __forceinline__ float vec_dot_q5_1_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { + + const block_q5_1 * bq5_1 = (const block_q5_1 *) vbq; + + int vl[VDR_Q5_1_Q8_1_MMVQ]; + int vh[VDR_Q5_1_Q8_1_MMVQ]; + int u[2*VDR_Q5_1_Q8_1_MMVQ]; + +#pragma unroll + for (int i = 0; i < VDR_Q5_1_Q8_1_MMVQ; ++i) { + vl[i] = get_int_from_uint8_aligned(bq5_1->qs, iqs + i); + vh[i] = get_int_from_uint8_aligned(bq5_1->qh, 0) >> (4 * (iqs + i)); + u[2*i+0] = get_int_from_int8_aligned(bq8_1->qs, iqs + i); + u[2*i+1] = get_int_from_int8_aligned(bq8_1->qs, iqs + i + QI5_1); + } + + return vec_dot_q5_1_q8_1_impl(vl, vh, u, bq5_1->dm, bq8_1->ds); +} + +static __device__ __forceinline__ float vec_dot_q8_0_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { + + const block_q8_0 * bq8_0 = (const block_q8_0 *) vbq; + + int v[VDR_Q8_0_Q8_1_MMVQ]; + int u[VDR_Q8_0_Q8_1_MMVQ]; + +#pragma unroll + for (int i = 0; i < VDR_Q8_0_Q8_1_MMVQ; ++i) { + v[i] = get_int_from_int8(bq8_0->qs, iqs + i); + u[i] = get_int_from_int8_aligned(bq8_1->qs, iqs + i); + } + + return vec_dot_q8_0_q8_1_impl(v, u, bq8_0->d, __low2half(bq8_1->ds)); +} + +static __device__ __forceinline__ float vec_dot_q2_K_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { + + const block_q2_K * bq2_K = (const block_q2_K *) vbq; + + const int bq8_offset = QR2_K * (iqs / QI8_1); + const int scale_offset = iqs - iqs % QI8_1 + (iqs % QI8_1) / (QI8_1/2); + + const uint8_t * scales = bq2_K->scales + scale_offset; + + const int v = get_int_from_uint8_aligned(bq2_K->qs, iqs); + int u[QR2_K]; + float d8[QR2_K]; + +#pragma unroll + for (int i = 0; i < QR2_K; ++ i) { + u[i] = get_int_from_int8_aligned(bq8_1[bq8_offset + i].qs, iqs % QI8_1); + d8[i] = __low2float(bq8_1[bq8_offset + i].ds); + } + + return vec_dot_q2_K_q8_1_impl_mmvq(v, u, scales, bq2_K->dm, d8); +} + +static __device__ __forceinline__ float vec_dot_q3_K_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { + + const block_q3_K * bq3_K = (const block_q3_K *) vbq; + + const int bq8_offset = QR3_K * (iqs / (QI3_K/2)); + const int scale_offset = iqs - iqs % QI8_1 + (iqs % QI8_1) / (QI8_1/2); + + const float d = bq3_K->d; + + const int vl = get_int_from_uint8(bq3_K->qs, iqs); + + // invert the mask with ~ so that a 0/1 results in 4/0 being subtracted + const int vh = ~get_int_from_uint8(bq3_K->hmask, iqs % (QI3_K/2)) >> bq8_offset; + + int u[QR3_K]; + float d8[QR3_K]; + +#pragma unroll + for (int i = 0; i < QR3_K; ++i) { + u[i] = get_int_from_int8_aligned(bq8_1[bq8_offset + i].qs, iqs % QI8_1); + d8[i] = __low2float(bq8_1[bq8_offset + i].ds); + } + + return vec_dot_q3_K_q8_1_impl_mmvq(vl, vh, u, bq3_K->scales, scale_offset, d, d8); +} + +static __device__ __forceinline__ float vec_dot_q4_K_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { + +#ifndef GGML_QKK_64 + const block_q4_K * bq4_K = (const block_q4_K *) vbq; + + int v[2]; + int u[2*QR4_K]; + float d8[QR4_K]; + + // iqs is in 0,2..30. bq8_offset = iqs/4 -> bq8_offset = 0, 2, 4, 6 + const int bq8_offset = QR4_K * ((iqs/2) / (QI8_1/2)); + + // iqs = 0....3 -> bq8_offset = 0, want q4_offset = 0, 4, 8, 12 + // iqs = 4....7 -> bq8_offset = 2, want q4_offset = 32, 36, 40, 44 + // iqs = 8...11 -> bq8_offset = 4, want q4_offset = 64, 68, 72, 76 + // iqs = 12..15 -> bq8_offset = 6, want q4_offset = 96, 100, 104, 108 + + const int * q4 = (const int *)(bq4_K->qs + 16 * bq8_offset + 4 * ((iqs/2)%4)); + v[0] = q4[0]; + v[1] = q4[4]; + + const uint16_t * scales = (const uint16_t *)bq4_K->scales; + uint16_t aux[2]; + const int j = bq8_offset/2; + if (j < 2) { + aux[0] = scales[j+0] & 0x3f3f; + aux[1] = scales[j+2] & 0x3f3f; + } else { + aux[0] = ((scales[j+2] >> 0) & 0x0f0f) | ((scales[j-2] & 0xc0c0) >> 2); + aux[1] = ((scales[j+2] >> 4) & 0x0f0f) | ((scales[j-0] & 0xc0c0) >> 2); + } + const uint8_t * sc = (const uint8_t *)aux; + const uint8_t * m = sc + 2; + + for (int i = 0; i < QR4_K; ++i) { + const block_q8_1 * bq8i = bq8_1 + bq8_offset + i; + d8[i] = __low2float(bq8i->ds); + + const int * q8 = (const int *)bq8i->qs + ((iqs/2)%4); + u[2*i+0] = q8[0]; + u[2*i+1] = q8[4]; + } + + return vec_dot_q4_K_q8_1_impl_vmmq(v, u, sc, m, bq4_K->dm, d8); + +#else + + const block_q4_K * bq4_K = (const block_q4_K *) vbq; + + float sumf_d = 0.0f; + float sumf_m = 0.0f; + + uint16_t aux16[2]; + const uint8_t * s = (const uint8_t *)aux16; + + const uint16_t * a = (const uint16_t *)bq4_K->scales; + aux16[0] = a[0] & 0x0f0f; + aux16[1] = (a[0] >> 4) & 0x0f0f; + + const float dall = bq4_K->dm[0]; + const float dmin = bq4_K->dm[1]; + + const float d8_1 = __low2float(bq8_1[0].ds); + const float d8_2 = __low2float(bq8_1[1].ds); + + const int ui1 = *((const int *)bq8_1[0].qs + (iqs/2)); + const int ui2 = *((const int *)bq8_1[0].qs + (iqs/2) + 4); + const int ui3 = *((const int *)bq8_1[1].qs + (iqs/2)); + const int ui4 = *((const int *)bq8_1[1].qs + (iqs/2) + 4); + + const int * q4 = (const int *)bq4_K->qs + (iqs/2); + const int v1 = q4[0]; + const int v2 = q4[4]; + + const int dot1 = ggml_cuda_dp4a(ui2, v2 & 0x0f0f0f0f, ggml_cuda_dp4a(ui1, v1 & 0x0f0f0f0f, 0)); + const int dot2 = ggml_cuda_dp4a(ui4, (v2 >> 4) & 0x0f0f0f0f, ggml_cuda_dp4a(ui3, (v1 >> 4) & 0x0f0f0f0f, 0)); + const int dot3 = ggml_cuda_dp4a(0x01010101, ui2, ggml_cuda_dp4a(0x01010101, ui1, 0)); + const int dot4 = ggml_cuda_dp4a(0x01010101, ui4, ggml_cuda_dp4a(0x01010101, ui3, 0)); + + sumf_d += d8_1 * (dot1 * s[0]) + d8_2 * (dot2 * s[1]); + sumf_m += d8_1 * (dot3 * s[2]) + d8_2 * (dot4 * s[3]); + + return dall * sumf_d - dmin * sumf_m; +#endif +} + +static __device__ __forceinline__ float vec_dot_q5_K_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { + +#ifndef GGML_QKK_64 + const block_q5_K * bq5_K = (const block_q5_K *) vbq; + + int vl[2]; + int vh[2]; + int u[2*QR5_K]; + float d8[QR5_K]; + + const int bq8_offset = QR5_K * ((iqs/2) / (QI8_1/2)); + const int * ql = (const int *)(bq5_K->qs + 16 * bq8_offset + 4 * ((iqs/2)%4)); + const int * qh = (const int *)(bq5_K->qh + 4 * ((iqs/2)%4)); + + vl[0] = ql[0]; + vl[1] = ql[4]; + + vh[0] = qh[0] >> bq8_offset; + vh[1] = qh[4] >> bq8_offset; + + const uint16_t * scales = (const uint16_t *)bq5_K->scales; + uint16_t aux[2]; + const int j = bq8_offset/2; + if (j < 2) { + aux[0] = scales[j+0] & 0x3f3f; + aux[1] = scales[j+2] & 0x3f3f; + } else { + aux[0] = ((scales[j+2] >> 0) & 0x0f0f) | ((scales[j-2] & 0xc0c0) >> 2); + aux[1] = ((scales[j+2] >> 4) & 0x0f0f) | ((scales[j-0] & 0xc0c0) >> 2); + } + const uint8_t * sc = (const uint8_t *)aux; + const uint8_t * m = sc + 2; + +#pragma unroll + for (int i = 0; i < QR5_K; ++i) { + const block_q8_1 * bq8i = bq8_1 + bq8_offset + i; + d8[i] = __low2float(bq8i->ds); + + const int * q8 = (const int *)bq8i->qs + ((iqs/2)%4); + u[2*i+0] = q8[0]; + u[2*i+1] = q8[4]; + } + + return vec_dot_q5_K_q8_1_impl_vmmq(vl, vh, u, sc, m, bq5_K->dm, d8); + +#else + + const block_q5_K * bq5_K = (const block_q5_K *) vbq; + + const int8_t * s = bq5_K->scales; + + const float d = bq5_K->d; + + const float d8_1 = __low2half(bq8_1[0].ds); + const float d8_2 = __low2half(bq8_1[1].ds); + + const int ui1 = *((const int *)bq8_1[0].qs + (iqs/2)); + const int ui2 = *((const int *)bq8_1[0].qs + (iqs/2) + 4); + const int ui3 = *((const int *)bq8_1[1].qs + (iqs/2)); + const int ui4 = *((const int *)bq8_1[1].qs + (iqs/2) + 4); + + const int * ql = (const int *)bq5_K->qs + (iqs/2); + const int vl1 = ql[0]; + const int vl2 = ql[4]; + + const int step = 4 * (iqs/2); // 0, 4, 8, 12 + const int im = step/8; // = 0 for iqs = 0, 2, = 1 for iqs = 4, 6 + const int in = step%8; // 0, 4, 0, 4 + const int vh = (*((const int *)(bq5_K->qh + in))) >> im; + + const int v1 = (((vh << 4) & 0x10101010) ^ 0x10101010) | ((vl1 >> 0) & 0x0f0f0f0f); + const int v2 = (((vh << 2) & 0x10101010) ^ 0x10101010) | ((vl2 >> 0) & 0x0f0f0f0f); + const int v3 = (((vh >> 0) & 0x10101010) ^ 0x10101010) | ((vl1 >> 4) & 0x0f0f0f0f); + const int v4 = (((vh >> 2) & 0x10101010) ^ 0x10101010) | ((vl2 >> 4) & 0x0f0f0f0f); + + const float sumf_d = d8_1 * (ggml_cuda_dp4a(ui1, v1, 0) * s[0] + ggml_cuda_dp4a(ui2, v2, 0) * s[1]) + + d8_2 * (ggml_cuda_dp4a(ui3, v3, 0) * s[2] + ggml_cuda_dp4a(ui4, v4, 0) * s[3]); + + return d * sumf_d; +#endif +} + +static __device__ __forceinline__ float vec_dot_q6_K_q8_1( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs) { + + const block_q6_K * bq6_K = (const block_q6_K *) vbq; + + const int bq8_offset = 2 * QR6_K * (iqs / (QI6_K/2)) + (iqs % (QI6_K/2)) / (QI6_K/4); + const int scale_offset = (QI6_K/4) * (iqs / (QI6_K/2)) + (iqs % (QI6_K/2)) / (QI6_K/8); + const int vh_shift = 2 * ((iqs % (QI6_K/2)) / (QI6_K/4)); + + const int vl = get_int_from_uint8(bq6_K->ql, iqs); + const int vh = get_int_from_uint8(bq6_K->qh, (QI6_K/4) * (iqs / (QI6_K/2)) + iqs % (QI6_K/4)) >> vh_shift; + + const int8_t * scales = bq6_K->scales + scale_offset; + + int u[QR6_K]; + float d8[QR6_K]; + +#pragma unroll + for (int i = 0; i < QR6_K; ++i) { + u[i] = get_int_from_int8_aligned(bq8_1[bq8_offset + 2*i].qs, iqs % QI8_1); + d8[i] = __low2float(bq8_1[bq8_offset + 2*i].ds); + } + + return vec_dot_q6_K_q8_1_impl_mmvq(vl, vh, u, scales, bq6_K->d, d8); +} + +// https://github.com/ggerganov/llama.cpp/blob/c50a82ce0f71558cbb8e555146ba124251504b38/ggml-cuda/mmvq.cu#L4 +typedef float (*vec_dot_q_cuda_t)(const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & iqs); + +template +static __device__ void mul_mat_vec_q( + const void * __restrict__ vx, const void * __restrict__ vy, float * __restrict__ dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + +#if defined(GGML_USE_HIPBLAS) && defined(__HIP_PLATFORM_AMD__) && (defined(RDNA2) || defined(RDNA3)) + constexpr int nwarps = 1; + constexpr int rows_per_cuda_block = 1; +#else + constexpr int nwarps = ncols_y <= 4 ? 4 : 2; + constexpr int rows_per_cuda_block = ncols_y == 1 ? 1 : 2; +#endif // defined(GGML_USE_HIPBLAS) && defined(__HIP_PLATFORM_AMD__) && !defined(RDNA2) && !defined(RDNA3) + + const int tid = WARP_SIZE*threadIdx.y + threadIdx.x; + const int row0 = rows_per_cuda_block*blockIdx.x; + const int blocks_per_row_x = ncols_x / qk; + const int blocks_per_col_y = nrows_y / QK8_1; + constexpr int blocks_per_iter = vdr * nwarps*WARP_SIZE / qi; + +// partial sum for each thread + float tmp[ncols_y][rows_per_cuda_block] = {0.0f}; + + const block_q_t * x = (const block_q_t *) vx; + const block_q8_1 * y = (const block_q8_1 *) vy; + + for (int kbx = tid / (qi/vdr); kbx < blocks_per_row_x; kbx += blocks_per_iter) { + const int kby = kbx * (qk/QK8_1); // y block index that aligns with kbx + + // x block quant index when casting the quants to int + const int kqs = vdr * (tid % (qi/vdr)); + +#pragma unroll + for (int j = 0; j < ncols_y; ++j) { +#pragma unroll + for (int i = 0; i < rows_per_cuda_block; ++i) { + tmp[j][i] += vec_dot_q_cuda( + &x[kbx + (row0 + i)*blocks_per_row_x], &y[j*blocks_per_col_y + kby], kqs); + } + } + } + + __shared__ float tmp_shared[nwarps-1 > 0 ? nwarps-1 : 1][ncols_y][rows_per_cuda_block][WARP_SIZE]; + if (threadIdx.y > 0) { +#pragma unroll + for (int j = 0; j < ncols_y; ++j) { +#pragma unroll + for (int i = 0; i < rows_per_cuda_block; ++i) { + tmp_shared[threadIdx.y-1][j][i][threadIdx.x] = tmp[j][i]; + } + } + } + __syncthreads(); + if (threadIdx.y > 0) { + return; + } + + // sum up partial sums and write back result +#pragma unroll + for (int j = 0; j < ncols_y; ++j) { +#pragma unroll + for (int i = 0; i < rows_per_cuda_block; ++i) { +#pragma unroll + for (int l = 0; l < nwarps-1; ++l) { + tmp[j][i] += tmp_shared[l][j][i][threadIdx.x]; + } + tmp[j][i] = warp_reduce_sum(tmp[j][i]); + } + + if (threadIdx.x < rows_per_cuda_block) { + dst[j*nrows_dst + row0 + threadIdx.x] = tmp[j][threadIdx.x]; + } + } +} + +// batch size = 1 +extern "C" __global__ void mul_mat_vec_q4_0_q8_1_cuda1( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<1, QK4_0, QI4_0, block_q4_0, VDR_Q4_0_Q8_1_MMVQ, vec_dot_q4_0_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q4_1_q8_1_cuda1( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<1, QK4_1, QI4_1, block_q4_1, VDR_Q4_1_Q8_1_MMVQ, vec_dot_q4_1_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q5_0_q8_1_cuda1( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<1, QK5_0, QI5_0, block_q5_0, VDR_Q5_0_Q8_1_MMVQ, vec_dot_q5_0_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q5_1_q8_1_cuda1( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<1, QK5_1, QI5_1, block_q5_1, VDR_Q5_1_Q8_1_MMVQ, vec_dot_q5_1_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q8_0_q8_1_cuda1( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<1, QK8_0, QI8_0, block_q8_0, VDR_Q8_0_Q8_1_MMVQ, vec_dot_q8_0_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q2_K_q8_1_cuda1( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<1, QK_K, QI2_K, block_q2_K, VDR_Q2_K_Q8_1_MMVQ, vec_dot_q2_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q3_K_q8_1_cuda1( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<1, QK_K, QI3_K, block_q3_K, VDR_Q3_K_Q8_1_MMVQ, vec_dot_q3_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q4_K_q8_1_cuda1( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<1, QK_K, QI4_K, block_q4_K, VDR_Q4_K_Q8_1_MMVQ, vec_dot_q4_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q5_K_q8_1_cuda1( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<1, QK_K, QI5_K, block_q5_K, VDR_Q5_K_Q8_1_MMVQ, vec_dot_q5_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q6_K_q8_1_cuda1( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<1, QK_K, QI6_K, block_q6_K, VDR_Q6_K_Q8_1_MMVQ, vec_dot_q6_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +// batch size = 2 +extern "C" __global__ void mul_mat_vec_q4_0_q8_1_cuda2( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<2, QK4_0, QI4_0, block_q4_0, VDR_Q4_0_Q8_1_MMVQ, vec_dot_q4_0_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q4_1_q8_1_cuda2( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<2, QK4_1, QI4_1, block_q4_1, VDR_Q4_1_Q8_1_MMVQ, vec_dot_q4_1_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q5_0_q8_1_cuda2( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<2, QK5_0, QI5_0, block_q5_0, VDR_Q5_0_Q8_1_MMVQ, vec_dot_q5_0_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q5_1_q8_1_cuda2( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<2, QK5_1, QI5_1, block_q5_1, VDR_Q5_1_Q8_1_MMVQ, vec_dot_q5_1_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q8_0_q8_1_cuda2( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<2, QK8_0, QI8_0, block_q8_0, VDR_Q8_0_Q8_1_MMVQ, vec_dot_q8_0_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q2_K_q8_1_cuda2( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<2, QK_K, QI2_K, block_q2_K, VDR_Q2_K_Q8_1_MMVQ, vec_dot_q2_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q3_K_q8_1_cuda2( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<2, QK_K, QI3_K, block_q3_K, VDR_Q3_K_Q8_1_MMVQ, vec_dot_q3_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q4_K_q8_1_cuda2( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<2, QK_K, QI4_K, block_q4_K, VDR_Q4_K_Q8_1_MMVQ, vec_dot_q4_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q5_K_q8_1_cuda2( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<2, QK_K, QI5_K, block_q5_K, VDR_Q5_K_Q8_1_MMVQ, vec_dot_q5_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q6_K_q8_1_cuda2( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<2, QK_K, QI6_K, block_q6_K, VDR_Q6_K_Q8_1_MMVQ, vec_dot_q6_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +// batch size = 3 +extern "C" __global__ void mul_mat_vec_q4_0_q8_1_cuda3( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<3, QK4_0, QI4_0, block_q4_0, VDR_Q4_0_Q8_1_MMVQ, vec_dot_q4_0_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q4_1_q8_1_cuda3( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<3, QK4_1, QI4_1, block_q4_1, VDR_Q4_1_Q8_1_MMVQ, vec_dot_q4_1_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q5_0_q8_1_cuda3( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<3, QK5_0, QI5_0, block_q5_0, VDR_Q5_0_Q8_1_MMVQ, vec_dot_q5_0_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q5_1_q8_1_cuda3( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<3, QK5_1, QI5_1, block_q5_1, VDR_Q5_1_Q8_1_MMVQ, vec_dot_q5_1_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q8_0_q8_1_cuda3( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<3, QK8_0, QI8_0, block_q8_0, VDR_Q8_0_Q8_1_MMVQ, vec_dot_q8_0_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q2_K_q8_1_cuda3( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<3, QK_K, QI2_K, block_q2_K, VDR_Q2_K_Q8_1_MMVQ, vec_dot_q2_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q3_K_q8_1_cuda3( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<3, QK_K, QI3_K, block_q3_K, VDR_Q3_K_Q8_1_MMVQ, vec_dot_q3_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q4_K_q8_1_cuda3( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<3, QK_K, QI4_K, block_q4_K, VDR_Q4_K_Q8_1_MMVQ, vec_dot_q4_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q5_K_q8_1_cuda3( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<3, QK_K, QI5_K, block_q5_K, VDR_Q5_K_Q8_1_MMVQ, vec_dot_q5_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q6_K_q8_1_cuda3( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<3, QK_K, QI6_K, block_q6_K, VDR_Q6_K_Q8_1_MMVQ, vec_dot_q6_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +// batch size = 4 +extern "C" __global__ void mul_mat_vec_q4_0_q8_1_cuda4( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<4, QK4_0, QI4_0, block_q4_0, VDR_Q4_0_Q8_1_MMVQ, vec_dot_q4_0_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q4_1_q8_1_cuda4( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<4, QK4_1, QI4_1, block_q4_1, VDR_Q4_1_Q8_1_MMVQ, vec_dot_q4_1_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q5_0_q8_1_cuda4( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<4, QK5_0, QI5_0, block_q5_0, VDR_Q5_0_Q8_1_MMVQ, vec_dot_q5_0_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q5_1_q8_1_cuda4( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<4, QK5_1, QI5_1, block_q5_1, VDR_Q5_1_Q8_1_MMVQ, vec_dot_q5_1_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q8_0_q8_1_cuda4( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<4, QK8_0, QI8_0, block_q8_0, VDR_Q8_0_Q8_1_MMVQ, vec_dot_q8_0_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q2_K_q8_1_cuda4( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<4, QK_K, QI2_K, block_q2_K, VDR_Q2_K_Q8_1_MMVQ, vec_dot_q2_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q3_K_q8_1_cuda4( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<4, QK_K, QI3_K, block_q3_K, VDR_Q3_K_Q8_1_MMVQ, vec_dot_q3_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q4_K_q8_1_cuda4( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<4, QK_K, QI4_K, block_q4_K, VDR_Q4_K_Q8_1_MMVQ, vec_dot_q4_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q5_K_q8_1_cuda4( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<4, QK_K, QI5_K, block_q5_K, VDR_Q5_K_Q8_1_MMVQ, vec_dot_q5_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q6_K_q8_1_cuda4( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<4, QK_K, QI6_K, block_q6_K, VDR_Q6_K_Q8_1_MMVQ, vec_dot_q6_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +// batch size = 5 +extern "C" __global__ void mul_mat_vec_q4_0_q8_1_cuda5( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<5, QK4_0, QI4_0, block_q4_0, VDR_Q4_0_Q8_1_MMVQ, vec_dot_q4_0_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q4_1_q8_1_cuda5( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<5, QK4_1, QI4_1, block_q4_1, VDR_Q4_1_Q8_1_MMVQ, vec_dot_q4_1_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q5_0_q8_1_cuda5( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<5, QK5_0, QI5_0, block_q5_0, VDR_Q5_0_Q8_1_MMVQ, vec_dot_q5_0_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q5_1_q8_1_cuda5( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<5, QK5_1, QI5_1, block_q5_1, VDR_Q5_1_Q8_1_MMVQ, vec_dot_q5_1_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q8_0_q8_1_cuda5( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<5, QK8_0, QI8_0, block_q8_0, VDR_Q8_0_Q8_1_MMVQ, vec_dot_q8_0_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q2_K_q8_1_cuda5( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<5, QK_K, QI2_K, block_q2_K, VDR_Q2_K_Q8_1_MMVQ, vec_dot_q2_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q3_K_q8_1_cuda5( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<5, QK_K, QI3_K, block_q3_K, VDR_Q3_K_Q8_1_MMVQ, vec_dot_q3_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q4_K_q8_1_cuda5( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<5, QK_K, QI4_K, block_q4_K, VDR_Q4_K_Q8_1_MMVQ, vec_dot_q4_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q5_K_q8_1_cuda5( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<5, QK_K, QI5_K, block_q5_K, VDR_Q5_K_Q8_1_MMVQ, vec_dot_q5_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q6_K_q8_1_cuda5( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<5, QK_K, QI6_K, block_q6_K, VDR_Q6_K_Q8_1_MMVQ, vec_dot_q6_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +// batch size = 6 +extern "C" __global__ void mul_mat_vec_q4_0_q8_1_cuda6( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<6, QK4_0, QI4_0, block_q4_0, VDR_Q4_0_Q8_1_MMVQ, vec_dot_q4_0_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q4_1_q8_1_cuda6( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<6, QK4_1, QI4_1, block_q4_1, VDR_Q4_1_Q8_1_MMVQ, vec_dot_q4_1_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q5_0_q8_1_cuda6( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<6, QK5_0, QI5_0, block_q5_0, VDR_Q5_0_Q8_1_MMVQ, vec_dot_q5_0_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q5_1_q8_1_cuda6( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<6, QK5_1, QI5_1, block_q5_1, VDR_Q5_1_Q8_1_MMVQ, vec_dot_q5_1_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q8_0_q8_1_cuda6( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<6, QK8_0, QI8_0, block_q8_0, VDR_Q8_0_Q8_1_MMVQ, vec_dot_q8_0_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q2_K_q8_1_cuda6( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<6, QK_K, QI2_K, block_q2_K, VDR_Q2_K_Q8_1_MMVQ, vec_dot_q2_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q3_K_q8_1_cuda6( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<6, QK_K, QI3_K, block_q3_K, VDR_Q3_K_Q8_1_MMVQ, vec_dot_q3_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q4_K_q8_1_cuda6( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<6, QK_K, QI4_K, block_q4_K, VDR_Q4_K_Q8_1_MMVQ, vec_dot_q4_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q5_K_q8_1_cuda6( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<6, QK_K, QI5_K, block_q5_K, VDR_Q5_K_Q8_1_MMVQ, vec_dot_q5_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q6_K_q8_1_cuda6( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<6, QK_K, QI6_K, block_q6_K, VDR_Q6_K_Q8_1_MMVQ, vec_dot_q6_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +// batch size = 7 +extern "C" __global__ void mul_mat_vec_q4_0_q8_1_cuda7( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<7, QK4_0, QI4_0, block_q4_0, VDR_Q4_0_Q8_1_MMVQ, vec_dot_q4_0_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q4_1_q8_1_cuda7( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<7, QK4_1, QI4_1, block_q4_1, VDR_Q4_1_Q8_1_MMVQ, vec_dot_q4_1_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q5_0_q8_1_cuda7( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<7, QK5_0, QI5_0, block_q5_0, VDR_Q5_0_Q8_1_MMVQ, vec_dot_q5_0_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q5_1_q8_1_cuda7( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<7, QK5_1, QI5_1, block_q5_1, VDR_Q5_1_Q8_1_MMVQ, vec_dot_q5_1_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q8_0_q8_1_cuda7( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<7, QK8_0, QI8_0, block_q8_0, VDR_Q8_0_Q8_1_MMVQ, vec_dot_q8_0_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q2_K_q8_1_cuda7( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<7, QK_K, QI2_K, block_q2_K, VDR_Q2_K_Q8_1_MMVQ, vec_dot_q2_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q3_K_q8_1_cuda7( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<7, QK_K, QI3_K, block_q3_K, VDR_Q3_K_Q8_1_MMVQ, vec_dot_q3_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q4_K_q8_1_cuda7( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<7, QK_K, QI4_K, block_q4_K, VDR_Q4_K_Q8_1_MMVQ, vec_dot_q4_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q5_K_q8_1_cuda7( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<7, QK_K, QI5_K, block_q5_K, VDR_Q5_K_Q8_1_MMVQ, vec_dot_q5_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q6_K_q8_1_cuda7( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<7, QK_K, QI6_K, block_q6_K, VDR_Q6_K_Q8_1_MMVQ, vec_dot_q6_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +// batch size = 8 +extern "C" __global__ void mul_mat_vec_q4_0_q8_1_cuda8( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<8, QK4_0, QI4_0, block_q4_0, VDR_Q4_0_Q8_1_MMVQ, vec_dot_q4_0_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q4_1_q8_1_cuda8( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<8, QK4_1, QI4_1, block_q4_1, VDR_Q4_1_Q8_1_MMVQ, vec_dot_q4_1_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q5_0_q8_1_cuda8( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<8, QK5_0, QI5_0, block_q5_0, VDR_Q5_0_Q8_1_MMVQ, vec_dot_q5_0_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q5_1_q8_1_cuda8( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<8, QK5_1, QI5_1, block_q5_1, VDR_Q5_1_Q8_1_MMVQ, vec_dot_q5_1_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q8_0_q8_1_cuda8( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<8, QK8_0, QI8_0, block_q8_0, VDR_Q8_0_Q8_1_MMVQ, vec_dot_q8_0_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q2_K_q8_1_cuda8( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<8, QK_K, QI2_K, block_q2_K, VDR_Q2_K_Q8_1_MMVQ, vec_dot_q2_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q3_K_q8_1_cuda8( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<8, QK_K, QI3_K, block_q3_K, VDR_Q3_K_Q8_1_MMVQ, vec_dot_q3_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q4_K_q8_1_cuda8( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<8, QK_K, QI4_K, block_q4_K, VDR_Q4_K_Q8_1_MMVQ, vec_dot_q4_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q5_K_q8_1_cuda8( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<8, QK_K, QI5_K, block_q5_K, VDR_Q5_K_Q8_1_MMVQ, vec_dot_q5_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void mul_mat_vec_q6_K_q8_1_cuda8( + const void * vx, const void * vy, float * dst, + const int ncols_x, const int nrows_x, const int nrows_y, const int nrows_dst) { + + mul_mat_vec_q<8, QK_K, QI6_K, block_q6_K, VDR_Q6_K_Q8_1_MMVQ, vec_dot_q6_K_q8_1> + (vx, vy, dst, ncols_x, nrows_x, nrows_y, nrows_dst); +} + +extern "C" __global__ void quantize_q8_1(const float * __restrict__ x, void * __restrict__ vy, const int kx, const int kx_padded) { + const int ix = blockDim.x*blockIdx.x + threadIdx.x; + + if (ix >= kx_padded) { + return; + } + + const int iy = blockDim.y*blockIdx.y + threadIdx.y; + + const int i_padded = iy*kx_padded + ix; + + block_q8_1 * y = (block_q8_1 *) vy; + + const int ib = i_padded / QK8_1; // block index + const int iqs = i_padded % QK8_1; // quant index + + const float xi = ix < kx ? x[iy*kx + ix] : 0.0f; + float amax = fabsf(xi); + float sum = xi; + + amax = warp_reduce_max(amax); + sum = warp_reduce_sum(sum); + + const float d = amax / 127; + const int8_t q = amax == 0.0f ? 0 : roundf(xi / d); + + y[ib].qs[iqs] = q; + + if (iqs > 0) { + return; + } + + reinterpret_cast(y[ib].ds.x) = d; + reinterpret_cast(y[ib].ds.y) = sum; +} + +// Kernels from https://github.com/ggerganov/llama.cpp/blob/master/ggml-cuda/mmq.cu + +template static __device__ __forceinline__ void allocate_tiles_q5_0(int ** x_ql, half2 ** x_dm, int ** x_qh, int ** x_sc) { + GGML_UNUSED(x_qh); GGML_UNUSED(x_sc); + + __shared__ int tile_x_ql[mmq_y * (2*WARP_SIZE) + mmq_y]; + __shared__ float tile_x_d[mmq_y * (WARP_SIZE/QI5_0) + mmq_y/QI5_0]; + + *x_ql = tile_x_ql; + *x_dm = (half2 *) tile_x_d; +} + +template static __device__ __forceinline__ void load_tiles_q5_0( + const void * __restrict__ vx, int * __restrict__ x_ql, half2 * __restrict__ x_dm, int * __restrict__ x_qh, + int * __restrict__ x_sc, const int & i_offset, const int & i_max, const int & k, const int & blocks_per_row) { + GGML_UNUSED(x_qh); GGML_UNUSED(x_sc); + + GGML_CUDA_ASSUME(i_offset >= 0); + GGML_CUDA_ASSUME(i_offset < nwarps); + GGML_CUDA_ASSUME(k >= 0); + GGML_CUDA_ASSUME(k < WARP_SIZE); + + const int kbx = k / QI5_0; + const int kqsx = k % QI5_0; + + const block_q5_0 * bx0 = (const block_q5_0 *) vx; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { + int i = i0 + i_offset; + + if (need_check) { + i = min(i, i_max); + } + + const block_q5_0 * bxi = bx0 + i*blocks_per_row + kbx; + + const int ql = get_int_from_uint8(bxi->qs, kqsx); + const int qh = get_int_from_uint8(bxi->qh, 0) >> (4 * (k % QI5_0)); + + int qs0 = (ql >> 0) & 0x0F0F0F0F; + qs0 |= (qh << 4) & 0x00000010; // 0 -> 4 + qs0 |= (qh << 11) & 0x00001000; // 1 -> 12 + qs0 |= (qh << 18) & 0x00100000; // 2 -> 20 + qs0 |= (qh << 25) & 0x10000000; // 3 -> 28 + qs0 = __vsubss4(qs0, 0x10101010); // subtract 16 + + x_ql[i * (2*WARP_SIZE + 1) + 2*k+0] = qs0; + + int qs1 = (ql >> 4) & 0x0F0F0F0F; + qs1 |= (qh >> 12) & 0x00000010; // 16 -> 4 + qs1 |= (qh >> 5) & 0x00001000; // 17 -> 12 + qs1 |= (qh << 2) & 0x00100000; // 18 -> 20 + qs1 |= (qh << 9) & 0x10000000; // 19 -> 28 + qs1 = __vsubss4(qs1, 0x10101010); // subtract 16 + + x_ql[i * (2*WARP_SIZE + 1) + 2*k+1] = qs1; + } + + const int blocks_per_tile_x_row = WARP_SIZE / QI5_0; + const int kbxd = k % blocks_per_tile_x_row; + float * x_dmf = (float *) x_dm; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI5_0) { + int i = i0 + i_offset * QI5_0 + k / blocks_per_tile_x_row; + + if (need_check) { + i = min(i, i_max); + } + + const block_q5_0 * bxi = bx0 + i*blocks_per_row + kbxd; + + x_dmf[i * (WARP_SIZE/QI5_0) + i / QI5_0 + kbxd] = bxi->d; + } +} + +static __device__ __forceinline__ float vec_dot_q5_0_q8_1_mul_mat( + const int * __restrict__ x_ql, const half2 * __restrict__ x_dm, const int * __restrict__ x_qh, const int * __restrict__ x_sc, + const int * __restrict__ y_qs, const half2 * __restrict__ y_ds, const int & i, const int & j, const int & k) { + GGML_UNUSED(x_qh); GGML_UNUSED(x_sc); + + const int kyqs = k % (QI8_1/2) + QI8_1 * (k / (QI8_1/2)); + const int index_bx = i * (WARP_SIZE/QI5_0) + i/QI5_0 + k/QI5_0; + const float * x_dmf = (const float *) x_dm; + const float * y_df = (const float *) y_ds; + + int u[2*VDR_Q5_0_Q8_1_MMQ]; + +#pragma unroll + for (int l = 0; l < VDR_Q5_0_Q8_1_MMQ; ++l) { + u[2*l+0] = y_qs[j * WARP_SIZE + (kyqs + l) % WARP_SIZE]; + u[2*l+1] = y_qs[j * WARP_SIZE + (kyqs + l + QI5_0) % WARP_SIZE]; + } + + return vec_dot_q8_0_q8_1_impl + (&x_ql[i * (2*WARP_SIZE + 1) + 2 * k], u, x_dmf[index_bx], y_df[j * (WARP_SIZE/QI8_1) + (2*k/QI8_1) % (WARP_SIZE/QI8_1)]); +} + + +template static __device__ __forceinline__ void allocate_tiles_q5_1(int ** x_ql, half2 ** x_dm, int ** x_qh, int ** x_sc) { + GGML_UNUSED(x_qh); GGML_UNUSED(x_sc); + + __shared__ int tile_x_ql[mmq_y * (2*WARP_SIZE) + mmq_y]; + __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE/QI5_1) + mmq_y/QI5_1]; + + *x_ql = tile_x_ql; + *x_dm = tile_x_dm; +} + +template static __device__ __forceinline__ void load_tiles_q5_1( + const void * __restrict__ vx, int * __restrict__ x_ql, half2 * __restrict__ x_dm, int * __restrict__ x_qh, + int * __restrict__ x_sc, const int & i_offset, const int & i_max, const int & k, const int & blocks_per_row) { + GGML_UNUSED(x_qh); GGML_UNUSED(x_sc); + + GGML_CUDA_ASSUME(i_offset >= 0); + GGML_CUDA_ASSUME(i_offset < nwarps); + GGML_CUDA_ASSUME(k >= 0); + GGML_CUDA_ASSUME(k < WARP_SIZE); + + const int kbx = k / QI5_1; + const int kqsx = k % QI5_1; + + const block_q5_1 * bx0 = (const block_q5_1 *) vx; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { + int i = i0 + i_offset; + + if (need_check) { + i = min(i, i_max); + } + + const block_q5_1 * bxi = bx0 + i*blocks_per_row + kbx; + + const int ql = get_int_from_uint8_aligned(bxi->qs, kqsx); + const int qh = get_int_from_uint8_aligned(bxi->qh, 0) >> (4 * (k % QI5_1)); + + int qs0 = (ql >> 0) & 0x0F0F0F0F; + qs0 |= (qh << 4) & 0x00000010; // 0 -> 4 + qs0 |= (qh << 11) & 0x00001000; // 1 -> 12 + qs0 |= (qh << 18) & 0x00100000; // 2 -> 20 + qs0 |= (qh << 25) & 0x10000000; // 3 -> 28 + + x_ql[i * (2*WARP_SIZE + 1) + 2*k+0] = qs0; + + int qs1 = (ql >> 4) & 0x0F0F0F0F; + qs1 |= (qh >> 12) & 0x00000010; // 16 -> 4 + qs1 |= (qh >> 5) & 0x00001000; // 17 -> 12 + qs1 |= (qh << 2) & 0x00100000; // 18 -> 20 + qs1 |= (qh << 9) & 0x10000000; // 19 -> 28 + + x_ql[i * (2*WARP_SIZE + 1) + 2*k+1] = qs1; + } + + const int blocks_per_tile_x_row = WARP_SIZE / QI5_1; + const int kbxd = k % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI5_1) { + int i = i0 + i_offset * QI5_1 + k / blocks_per_tile_x_row; + + if (need_check) { + i = min(i, i_max); + } + + const block_q5_1 * bxi = bx0 + i*blocks_per_row + kbxd; + + x_dm[i * (WARP_SIZE/QI5_1) + i / QI5_1 + kbxd] = bxi->dm; + } +} + +static __device__ __forceinline__ float vec_dot_q5_1_q8_1_mul_mat( + const int * __restrict__ x_ql, const half2 * __restrict__ x_dm, const int * __restrict__ x_qh, const int * __restrict__ x_sc, + const int * __restrict__ y_qs, const half2 * __restrict__ y_ds, const int & i, const int & j, const int & k) { + GGML_UNUSED(x_qh); GGML_UNUSED(x_sc); + + const int kyqs = k % (QI8_1/2) + QI8_1 * (k / (QI8_1/2)); + const int index_bx = i * (WARP_SIZE/QI5_1) + + i/QI5_1 + k/QI5_1; + + int u[2*VDR_Q5_1_Q8_1_MMQ]; + +#pragma unroll + for (int l = 0; l < VDR_Q5_1_Q8_1_MMQ; ++l) { + u[2*l+0] = y_qs[j * WARP_SIZE + (kyqs + l) % WARP_SIZE]; + u[2*l+1] = y_qs[j * WARP_SIZE + (kyqs + l + QI5_1) % WARP_SIZE]; + } + + return vec_dot_q8_1_q8_1_impl + (&x_ql[i * (2*WARP_SIZE + 1) + 2 * k], u, x_dm[index_bx], y_ds[j * (WARP_SIZE/QI8_1) + (2*k/QI8_1) % (WARP_SIZE/QI8_1)]); +} + +template static __device__ __forceinline__ void allocate_tiles_q8_0(int ** x_ql, half2 ** x_dm, int ** x_qh, int ** x_sc) { + GGML_UNUSED(x_qh); GGML_UNUSED(x_sc); + + __shared__ int tile_x_qs[mmq_y * (WARP_SIZE) + mmq_y]; + __shared__ float tile_x_d[mmq_y * (WARP_SIZE/QI8_0) + mmq_y/QI8_0]; + + *x_ql = tile_x_qs; + *x_dm = (half2 *) tile_x_d; +} + +template static __device__ __forceinline__ void load_tiles_q8_0( + const void * __restrict__ vx, int * __restrict__ x_ql, half2 * __restrict__ x_dm, int * __restrict__ x_qh, + int * __restrict__ x_sc, const int & i_offset, const int & i_max, const int & k, const int & blocks_per_row) { + GGML_UNUSED(x_qh); GGML_UNUSED(x_sc); + + GGML_CUDA_ASSUME(i_offset >= 0); + GGML_CUDA_ASSUME(i_offset < nwarps); + GGML_CUDA_ASSUME(k >= 0); + GGML_CUDA_ASSUME(k < WARP_SIZE); + + const int kbx = k / QI8_0; + const int kqsx = k % QI8_0; + float * x_dmf = (float *) x_dm; + + const block_q8_0 * bx0 = (const block_q8_0 *) vx; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { + int i = i0 + i_offset; + + if (need_check) { + i = min(i, i_max); + } + + const block_q8_0 * bxi = bx0 + i*blocks_per_row + kbx; + + x_ql[i * (WARP_SIZE + 1) + k] = get_int_from_int8(bxi->qs, kqsx); + } + + const int blocks_per_tile_x_row = WARP_SIZE / QI8_0; + const int kbxd = k % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI8_0) { + int i = i0 + i_offset * QI8_0 + k / blocks_per_tile_x_row; + + if (need_check) { + i = min(i, i_max); + } + + const block_q8_0 * bxi = bx0 + i*blocks_per_row + kbxd; + + x_dmf[i * (WARP_SIZE/QI8_0) + i / QI8_0 + kbxd] = bxi->d; + } +} + +static __device__ __forceinline__ float vec_dot_q8_0_q8_1_mul_mat( + const int * __restrict__ x_ql, const half2 * __restrict__ x_dm, const int * __restrict__ x_qh, const int * __restrict__ x_sc, + const int * __restrict__ y_qs, const half2 * __restrict__ y_ds, const int & i, const int & j, const int & k) { + GGML_UNUSED(x_qh); GGML_UNUSED(x_sc); + + const float * x_dmf = (const float *) x_dm; + const float * y_df = (const float *) y_ds; + + return vec_dot_q8_0_q8_1_impl + (&x_ql[i * (WARP_SIZE + 1) + k], &y_qs[j * WARP_SIZE + k], x_dmf[i * (WARP_SIZE/QI8_0) + i/QI8_0 + k/QI8_0], + y_df[j * (WARP_SIZE/QI8_1) + k/QI8_1]); +} + +template static __device__ __forceinline__ void allocate_tiles_q2_K(int ** x_ql, half2 ** x_dm, int ** x_qh, int ** x_sc) { + GGML_UNUSED(x_qh); + + __shared__ int tile_x_ql[mmq_y * (WARP_SIZE) + mmq_y]; + __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE/QI2_K) + mmq_y/QI2_K]; + __shared__ int tile_x_sc[mmq_y * (WARP_SIZE/4) + mmq_y/4]; + + *x_ql = tile_x_ql; + *x_dm = tile_x_dm; + *x_sc = tile_x_sc; +} + +template static __device__ __forceinline__ void load_tiles_q2_K( + const void * __restrict__ vx, int * __restrict__ x_ql, half2 * __restrict__ x_dm, int * __restrict__ x_qh, + int * __restrict__ x_sc, const int & i_offset, const int & i_max, const int & k, const int & blocks_per_row) { + GGML_UNUSED(x_qh); + + GGML_CUDA_ASSUME(i_offset >= 0); + GGML_CUDA_ASSUME(i_offset < nwarps); + GGML_CUDA_ASSUME(k >= 0); + GGML_CUDA_ASSUME(k < WARP_SIZE); + + const int kbx = k / QI2_K; + const int kqsx = k % QI2_K; + + const block_q2_K * bx0 = (const block_q2_K *) vx; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { + int i = i0 + i_offset; + + if (need_check) { + i = min(i, i_max); + } + + const block_q2_K * bxi = bx0 + i*blocks_per_row + kbx; + + x_ql[i * (WARP_SIZE + 1) + k] = get_int_from_uint8_aligned(bxi->qs, kqsx); + } + + const int blocks_per_tile_x_row = WARP_SIZE / QI2_K; + const int kbxd = k % blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI2_K) { + int i = (i0 + i_offset * QI2_K + k / blocks_per_tile_x_row) % mmq_y; + + if (need_check) { + i = min(i, i_max); + } + + const block_q2_K * bxi = bx0 + i*blocks_per_row + kbxd; + + x_dm[i * (WARP_SIZE/QI2_K) + i / QI2_K + kbxd] = bxi->dm; + } + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 4) { + int i = i0 + i_offset * 4 + k / (WARP_SIZE/4); + + if (need_check) { + i = min(i, i_max); + } + + const block_q2_K * bxi = bx0 + i*blocks_per_row + (k % (WARP_SIZE/4)) / (QI2_K/4); + + x_sc[i * (WARP_SIZE/4) + i / 4 + k % (WARP_SIZE/4)] = get_int_from_uint8_aligned(bxi->scales, k % (QI2_K/4)); + } +} + +static __device__ __forceinline__ float vec_dot_q2_K_q8_1_mul_mat( + const int * __restrict__ x_ql, const half2 * __restrict__ x_dm, const int * __restrict__ x_qh, const int * __restrict__ x_sc, + const int * __restrict__ y_qs, const half2 * __restrict__ y_ds, const int & i, const int & j, const int & k) { + GGML_UNUSED(x_qh); + + const int kbx = k / QI2_K; + const int ky = (k % QI2_K) * QR2_K; + const float * y_df = (const float *) y_ds; + + int v[QR2_K*VDR_Q2_K_Q8_1_MMQ]; + + const int kqsx = i * (WARP_SIZE + 1) + kbx*QI2_K + (QI2_K/2) * (ky/(2*QI2_K)) + ky % (QI2_K/2); + const int shift = 2 * ((ky % (2*QI2_K)) / (QI2_K/2)); + +#pragma unroll + for (int l = 0; l < QR2_K*VDR_Q2_K_Q8_1_MMQ; ++l) { + v[l] = (x_ql[kqsx + l] >> shift) & 0x03030303; + } + + const uint8_t * scales = ((const uint8_t *) &x_sc[i * (WARP_SIZE/4) + i/4 + kbx*4]) + ky/4; + + const int index_y = j * WARP_SIZE + (QR2_K*k) % WARP_SIZE; + return vec_dot_q2_K_q8_1_impl_mmq(v, &y_qs[index_y], scales, x_dm[i * (WARP_SIZE/QI2_K) + i/QI2_K + kbx], y_df[index_y/QI8_1]); +} + +template static __device__ __forceinline__ void allocate_tiles_q3_K(int ** x_ql, half2 ** x_dm, int ** x_qh, int ** x_sc) { + + __shared__ int tile_x_ql[mmq_y * (WARP_SIZE) + mmq_y]; + __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE/QI3_K) + mmq_y/QI3_K]; + __shared__ int tile_x_qh[mmq_y * (WARP_SIZE/2) + mmq_y/2]; + __shared__ int tile_x_sc[mmq_y * (WARP_SIZE/4) + mmq_y/4]; + + *x_ql = tile_x_ql; + *x_dm = tile_x_dm; + *x_qh = tile_x_qh; + *x_sc = tile_x_sc; +} + +template static __device__ __forceinline__ void load_tiles_q3_K( + const void * __restrict__ vx, int * __restrict__ x_ql, half2 * __restrict__ x_dm, int * __restrict__ x_qh, + int * __restrict__ x_sc, const int & i_offset, const int & i_max, const int & k, const int & blocks_per_row) { + + GGML_CUDA_ASSUME(i_offset >= 0); + GGML_CUDA_ASSUME(i_offset < nwarps); + GGML_CUDA_ASSUME(k >= 0); + GGML_CUDA_ASSUME(k < WARP_SIZE); + + const int kbx = k / QI3_K; + const int kqsx = k % QI3_K; + + const block_q3_K * bx0 = (const block_q3_K *) vx; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { + int i = i0 + i_offset; + + if (need_check) { + i = min(i, i_max); + } + + const block_q3_K * bxi = bx0 + i*blocks_per_row + kbx; + + x_ql[i * (WARP_SIZE + 1) + k] = get_int_from_uint8(bxi->qs, kqsx); + } + + const int blocks_per_tile_x_row = WARP_SIZE / QI3_K; + const int kbxd = k % blocks_per_tile_x_row; + float * x_dmf = (float *) x_dm; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI3_K) { + int i = (i0 + i_offset * QI3_K + k / blocks_per_tile_x_row) % mmq_y; + + if (need_check) { + i = min(i, i_max); + } + + const block_q3_K * bxi = bx0 + i*blocks_per_row + kbxd; + + x_dmf[i * (WARP_SIZE/QI3_K) + i / QI3_K + kbxd] = bxi->d; + } + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 2) { + int i = i0 + i_offset * 2 + k / (WARP_SIZE/2); + + if (need_check) { + i = min(i, i_max); + } + + const block_q3_K * bxi = bx0 + i*blocks_per_row + (k % (WARP_SIZE/2)) / (QI3_K/2); + + // invert the mask with ~ so that a 0/1 results in 4/0 being subtracted + x_qh[i * (WARP_SIZE/2) + i / 2 + k % (WARP_SIZE/2)] = ~get_int_from_uint8(bxi->hmask, k % (QI3_K/2)); + } + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 4) { + int i = i0 + i_offset * 4 + k / (WARP_SIZE/4); + + if (need_check) { + i = min(i, i_max); + } + + const block_q3_K * bxi = bx0 + i*blocks_per_row + (k % (WARP_SIZE/4)) / (QI3_K/4); + + const int ksc = k % (QI3_K/4); + + const int ksc_low = ksc % (QI3_K/8); + const int shift_low = 4 * (ksc / (QI3_K/8)); + const int sc_low = (get_int_from_uint8(bxi->scales, ksc_low) >> shift_low) & 0x0F0F0F0F; + + const int ksc_high = QI3_K/8; + const int shift_high = 2 * ksc; + const int sc_high = ((get_int_from_uint8(bxi->scales, ksc_high) >> shift_high) << 4) & 0x30303030; + + const int sc = __vsubss4(sc_low | sc_high, 0x20202020); + + x_sc[i * (WARP_SIZE/4) + i / 4 + k % (WARP_SIZE/4)] = sc; + } +} + +static __device__ __forceinline__ float vec_dot_q3_K_q8_1_mul_mat( + const int * __restrict__ x_ql, const half2 * __restrict__ x_dm, const int * __restrict__ x_qh, const int * __restrict__ x_sc, + const int * __restrict__ y_qs, const half2 * __restrict__ y_ds, const int & i, const int & j, const int & k) { + + const int kbx = k / QI3_K; + const int ky = (k % QI3_K) * QR3_K; + const float * x_dmf = (const float *) x_dm; + const float * y_df = (const float *) y_ds; + + const int8_t * scales = ((const int8_t *) (x_sc + i * (WARP_SIZE/4) + i/4 + kbx*4)) + ky/4; + + int v[QR3_K*VDR_Q3_K_Q8_1_MMQ]; + +#pragma unroll + for (int l = 0; l < QR3_K*VDR_Q3_K_Q8_1_MMQ; ++l) { + const int kqsx = i * (WARP_SIZE + 1) + kbx*QI3_K + (QI3_K/2) * (ky/(2*QI3_K)) + ky % (QI3_K/2); + const int shift = 2 * ((ky % 32) / 8); + const int vll = (x_ql[kqsx + l] >> shift) & 0x03030303; + + const int vh = x_qh[i * (WARP_SIZE/2) + i/2 + kbx * (QI3_K/2) + (ky+l)%8] >> ((ky+l) / 8); + const int vlh = (vh << 2) & 0x04040404; + + v[l] = __vsubss4(vll, vlh); + } + + const int index_y = j * WARP_SIZE + (k*QR3_K) % WARP_SIZE; + return vec_dot_q3_K_q8_1_impl_mmq(v, &y_qs[index_y], scales, x_dmf[i * (WARP_SIZE/QI3_K) + i/QI3_K + kbx], y_df[index_y/QI8_1]); +} + +template static __device__ __forceinline__ void allocate_tiles_q4_K(int ** x_ql, half2 ** x_dm, int ** x_qh, int ** x_sc) { + GGML_UNUSED(x_qh); + + __shared__ int tile_x_ql[mmq_y * (WARP_SIZE) + mmq_y]; + __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE/QI4_K) + mmq_y/QI4_K]; + __shared__ int tile_x_sc[mmq_y * (WARP_SIZE/8) + mmq_y/8]; + + *x_ql = tile_x_ql; + *x_dm = tile_x_dm; + *x_sc = tile_x_sc; +} + +template static __device__ __forceinline__ void load_tiles_q4_K( + const void * __restrict__ vx, int * __restrict__ x_ql, half2 * __restrict__ x_dm, int * __restrict__ x_qh, + int * __restrict__ x_sc, const int & i_offset, const int & i_max, const int & k, const int & blocks_per_row) { + GGML_UNUSED(x_qh); + + GGML_CUDA_ASSUME(i_offset >= 0); + GGML_CUDA_ASSUME(i_offset < nwarps); + GGML_CUDA_ASSUME(k >= 0); + GGML_CUDA_ASSUME(k < WARP_SIZE); + + const int kbx = k / QI4_K; // == 0 if QK_K == 256 + const int kqsx = k % QI4_K; // == k if QK_K == 256 + + const block_q4_K * bx0 = (const block_q4_K *) vx; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { + int i = i0 + i_offset; + + if (need_check) { + i = min(i, i_max); + } + + const block_q4_K * bxi = bx0 + i*blocks_per_row + kbx; + + x_ql[i * (WARP_SIZE + 1) + k] = get_int_from_uint8_aligned(bxi->qs, kqsx); + } + + const int blocks_per_tile_x_row = WARP_SIZE / QI4_K; // == 1 if QK_K == 256 + const int kbxd = k % blocks_per_tile_x_row; // == 0 if QK_K == 256 + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI4_K) { + int i = (i0 + i_offset * QI4_K + k / blocks_per_tile_x_row) % mmq_y; + + if (need_check) { + i = min(i, i_max); + } + + const block_q4_K * bxi = bx0 + i*blocks_per_row + kbxd; + +#if QK_K == 256 + x_dm[i * (WARP_SIZE/QI4_K) + i / QI4_K + kbxd] = bxi->dm; +#else + x_dm[i * (WARP_SIZE/QI4_K) + i / QI4_K + kbxd] = {bxi->dm[0], bxi->dm[1]}; +#endif + } + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 8) { + int i = (i0 + i_offset * 8 + k / (WARP_SIZE/8)) % mmq_y; + + if (need_check) { + i = min(i, i_max); + } + + const block_q4_K * bxi = bx0 + i*blocks_per_row + (k % (WARP_SIZE/8)) / (QI4_K/8); + + const int * scales = (const int *) bxi->scales; + + const int ksc = k % (WARP_SIZE/8); + + // scale arrangement after the following two lines: sc0,...,sc3, sc4,...,sc7, m0,...,m3, m4,...,m8 + int scales8 = (scales[(ksc%2) + (ksc!=0)] >> (4 * (ksc & (ksc/2)))) & 0x0F0F0F0F; // lower 4 bits + scales8 |= (scales[ksc/2] >> (2 * (ksc % 2))) & 0x30303030; // upper 2 bits + + x_sc[i * (WARP_SIZE/8) + i / 8 + ksc] = scales8; + } +} + +static __device__ __forceinline__ float vec_dot_q4_K_q8_1_mul_mat( + const int * __restrict__ x_ql, const half2 * __restrict__ x_dm, const int * __restrict__ x_qh, const int * __restrict__ x_sc, + const int * __restrict__ y_qs, const half2 * __restrict__ y_ds, const int & i, const int & j, const int & k) { + GGML_UNUSED(x_qh); + + const uint8_t * sc = ((const uint8_t *) &x_sc[i * (WARP_SIZE/8) + i/8 + k/16]) + 2*((k % 16) / 8); + + const int index_y = j * WARP_SIZE + (QR4_K*k) % WARP_SIZE; + return vec_dot_q4_K_q8_1_impl_mmq(&x_ql[i * (WARP_SIZE + 1) + k], &y_qs[index_y], sc, sc+8, + x_dm[i * (WARP_SIZE/QI4_K) + i/QI4_K], &y_ds[index_y/QI8_1]); +} + +template static __device__ __forceinline__ void allocate_tiles_q5_K(int ** x_ql, half2 ** x_dm, int ** x_qh, int ** x_sc) { + GGML_UNUSED(x_qh); + + __shared__ int tile_x_ql[mmq_y * (2*WARP_SIZE) + mmq_y]; + __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE/QI5_K) + mmq_y/QI5_K]; + __shared__ int tile_x_sc[mmq_y * (WARP_SIZE/8) + mmq_y/8]; + + *x_ql = tile_x_ql; + *x_dm = tile_x_dm; + *x_sc = tile_x_sc; +} + +template static __device__ __forceinline__ void load_tiles_q5_K( + const void * __restrict__ vx, int * __restrict__ x_ql, half2 * __restrict__ x_dm, int * __restrict__ x_qh, + int * __restrict__ x_sc, const int & i_offset, const int & i_max, const int & k, const int & blocks_per_row) { + GGML_UNUSED(x_qh); + + GGML_CUDA_ASSUME(i_offset >= 0); + GGML_CUDA_ASSUME(i_offset < nwarps); + GGML_CUDA_ASSUME(k >= 0); + GGML_CUDA_ASSUME(k < WARP_SIZE); + + const int kbx = k / QI5_K; // == 0 if QK_K == 256 + const int kqsx = k % QI5_K; // == k if QK_K == 256 + + const block_q5_K * bx0 = (const block_q5_K *) vx; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { + int i = i0 + i_offset; + + if (need_check) { + i = min(i, i_max); + } + + const block_q5_K * bxi = bx0 + i*blocks_per_row + kbx; + const int ky = QR5_K*kqsx; + + const int ql = get_int_from_uint8_aligned(bxi->qs, kqsx); + const int ql0 = (ql >> 0) & 0x0F0F0F0F; + const int ql1 = (ql >> 4) & 0x0F0F0F0F; + + const int qh = get_int_from_uint8_aligned(bxi->qh, kqsx % (QI5_K/4)); + const int qh0 = ((qh >> (2 * (kqsx / (QI5_K/4)) + 0)) << 4) & 0x10101010; + const int qh1 = ((qh >> (2 * (kqsx / (QI5_K/4)) + 1)) << 4) & 0x10101010; + + const int kq0 = ky - ky % (QI5_K/2) + k % (QI5_K/4) + 0; + const int kq1 = ky - ky % (QI5_K/2) + k % (QI5_K/4) + (QI5_K/4); + + x_ql[i * (2*WARP_SIZE + 1) + kq0] = ql0 | qh0; + x_ql[i * (2*WARP_SIZE + 1) + kq1] = ql1 | qh1; + } + + const int blocks_per_tile_x_row = WARP_SIZE / QI5_K; // == 1 if QK_K == 256 + const int kbxd = k % blocks_per_tile_x_row; // == 0 if QK_K == 256 + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI5_K) { + int i = (i0 + i_offset * QI5_K + k / blocks_per_tile_x_row) % mmq_y; + + if (need_check) { + i = min(i, i_max); + } + + const block_q5_K * bxi = bx0 + i*blocks_per_row + kbxd; + +#if QK_K == 256 + x_dm[i * (WARP_SIZE/QI5_K) + i / QI5_K + kbxd] = bxi->dm; +#endif + } + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 8) { + int i = (i0 + i_offset * 8 + k / (WARP_SIZE/8)) % mmq_y; + + if (need_check) { + i = min(i, i_max); + } + + const block_q5_K * bxi = bx0 + i*blocks_per_row + (k % (WARP_SIZE/8)) / (QI5_K/8); + + const int * scales = (const int *) bxi->scales; + + const int ksc = k % (WARP_SIZE/8); + + // scale arrangement after the following two lines: sc0,...,sc3, sc4,...,sc7, m0,...,m3, m4,...,m8 + int scales8 = (scales[(ksc%2) + (ksc!=0)] >> (4 * (ksc & (ksc/2)))) & 0x0F0F0F0F; // lower 4 bits + scales8 |= (scales[ksc/2] >> (2 * (ksc % 2))) & 0x30303030; // upper 2 bits + + x_sc[i * (WARP_SIZE/8) + i / 8 + ksc] = scales8; + } +} + +static __device__ __forceinline__ float vec_dot_q5_K_q8_1_mul_mat( + const int * __restrict__ x_ql, const half2 * __restrict__ x_dm, const int * __restrict__ x_qh, const int * __restrict__ x_sc, + const int * __restrict__ y_qs, const half2 * __restrict__ y_ds, const int & i, const int & j, const int & k) { + GGML_UNUSED(x_qh); + + const uint8_t * sc = ((const uint8_t *) &x_sc[i * (WARP_SIZE/8) + i/8 + k/16]) + 2 * ((k % 16) / 8); + + const int index_x = i * (QR5_K*WARP_SIZE + 1) + QR5_K*k; + const int index_y = j * WARP_SIZE + (QR5_K*k) % WARP_SIZE; + return vec_dot_q5_K_q8_1_impl_mmq(&x_ql[index_x], &y_qs[index_y], sc, sc+8, + x_dm[i * (WARP_SIZE/QI5_K) + i/QI5_K], &y_ds[index_y/QI8_1]); +} + +template static __device__ __forceinline__ void allocate_tiles_q6_K(int ** x_ql, half2 ** x_dm, int ** x_qh, int ** x_sc) { + GGML_UNUSED(x_qh); + + __shared__ int tile_x_ql[mmq_y * (2*WARP_SIZE) + mmq_y]; + __shared__ half2 tile_x_dm[mmq_y * (WARP_SIZE/QI6_K) + mmq_y/QI6_K]; + __shared__ int tile_x_sc[mmq_y * (WARP_SIZE/8) + mmq_y/8]; + + *x_ql = tile_x_ql; + *x_dm = tile_x_dm; + *x_sc = tile_x_sc; +} + +template static __device__ __forceinline__ void load_tiles_q6_K( + const void * __restrict__ vx, int * __restrict__ x_ql, half2 * __restrict__ x_dm, int * __restrict__ x_qh, + int * __restrict__ x_sc, const int & i_offset, const int & i_max, const int & k, const int & blocks_per_row) { + GGML_UNUSED(x_qh); + + GGML_CUDA_ASSUME(i_offset >= 0); + GGML_CUDA_ASSUME(i_offset < nwarps); + GGML_CUDA_ASSUME(k >= 0); + GGML_CUDA_ASSUME(k < WARP_SIZE); + + const int kbx = k / QI6_K; // == 0 if QK_K == 256 + const int kqsx = k % QI6_K; // == k if QK_K == 256 + + const block_q6_K * bx0 = (const block_q6_K *) vx; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps) { + int i = i0 + i_offset; + + if (need_check) { + i = min(i, i_max); + } + + const block_q6_K * bxi = bx0 + i*blocks_per_row + kbx; + const int ky = QR6_K*kqsx; + + const int ql = get_int_from_uint8(bxi->ql, kqsx); + const int ql0 = (ql >> 0) & 0x0F0F0F0F; + const int ql1 = (ql >> 4) & 0x0F0F0F0F; + + const int qh = get_int_from_uint8(bxi->qh, (QI6_K/4) * (kqsx / (QI6_K/2)) + kqsx % (QI6_K/4)); + const int qh0 = ((qh >> (2 * ((kqsx % (QI6_K/2)) / (QI6_K/4)))) << 4) & 0x30303030; + const int qh1 = (qh >> (2 * ((kqsx % (QI6_K/2)) / (QI6_K/4)))) & 0x30303030; + + const int kq0 = ky - ky % QI6_K + k % (QI6_K/2) + 0; + const int kq1 = ky - ky % QI6_K + k % (QI6_K/2) + (QI6_K/2); + + x_ql[i * (2*WARP_SIZE + 1) + kq0] = __vsubss4(ql0 | qh0, 0x20202020); + x_ql[i * (2*WARP_SIZE + 1) + kq1] = __vsubss4(ql1 | qh1, 0x20202020); + } + + const int blocks_per_tile_x_row = WARP_SIZE / QI6_K; // == 1 if QK_K == 256 + const int kbxd = k % blocks_per_tile_x_row; // == 0 if QK_K == 256 + float * x_dmf = (float *) x_dm; + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * QI6_K) { + int i = (i0 + i_offset * QI6_K + k / blocks_per_tile_x_row) % mmq_y; + + if (need_check) { + i = min(i, i_max); + } + + const block_q6_K * bxi = bx0 + i*blocks_per_row + kbxd; + + x_dmf[i * (WARP_SIZE/QI6_K) + i / QI6_K + kbxd] = bxi->d; + } + +#pragma unroll + for (int i0 = 0; i0 < mmq_y; i0 += nwarps * 8) { + int i = (i0 + i_offset * 8 + k / (WARP_SIZE/8)) % mmq_y; + + if (need_check) { + i = min(i, i_max); + } + + const block_q6_K * bxi = bx0 + i*blocks_per_row + (k % (WARP_SIZE/8)) / 4; + + x_sc[i * (WARP_SIZE/8) + i / 8 + k % (WARP_SIZE/8)] = get_int_from_int8(bxi->scales, k % (QI6_K/8)); + } +} + +static __device__ __forceinline__ float vec_dot_q6_K_q8_1_mul_mat( + const int * __restrict__ x_ql, const half2 * __restrict__ x_dm, const int * __restrict__ x_qh, const int * __restrict__ x_sc, + const int * __restrict__ y_qs, const half2 * __restrict__ y_ds, const int & i, const int & j, const int & k) { + GGML_UNUSED(x_qh); + + const float * x_dmf = (const float *) x_dm; + const float * y_df = (const float *) y_ds; + + const int8_t * sc = ((const int8_t *) &x_sc[i * (WARP_SIZE/8) + i/8 + k/8]); + + const int index_x = i * (QR6_K*WARP_SIZE + 1) + QR6_K*k; + const int index_y = j * WARP_SIZE + (QR6_K*k) % WARP_SIZE; + return vec_dot_q6_K_q8_1_impl_mmq(&x_ql[index_x], &y_qs[index_y], sc, x_dmf[i * (WARP_SIZE/QI6_K) + i/QI6_K], &y_df[index_y/QI8_1]); +} + + +static __device__ __forceinline__ float vec_dot_q4_0_q8_1_mul_mat( + const int * __restrict__ x_ql, const half2 * __restrict__ x_dm, const int * __restrict__ x_qh, const int * __restrict__ x_sc, + const int * __restrict__ y_qs, const half2 * __restrict__ y_ds, const int & i, const int & j, const int & k) { + + const int kyqs = k % (QI8_1/2) + QI8_1 * (k / (QI8_1/2)); + const float * x_dmf = (const float *) x_dm; + + int u[2*VDR_Q4_0_Q8_1_MMQ]; + +#pragma unroll + for (int l = 0; l < VDR_Q4_0_Q8_1_MMQ; ++l) { + u[2*l+0] = y_qs[j * WARP_SIZE + (kyqs + l) % WARP_SIZE]; + u[2*l+1] = y_qs[j * WARP_SIZE + (kyqs + l + QI4_0) % WARP_SIZE]; + } + + return vec_dot_q4_0_q8_1_impl + (&x_ql[i * (WARP_SIZE + 1) + k], u, x_dmf[i * (WARP_SIZE/QI4_0) + i/QI4_0 + k/QI4_0], + y_ds[j * (WARP_SIZE/QI8_1) + (2*k/QI8_1) % (WARP_SIZE/QI8_1)]); +} + +static __device__ __forceinline__ float vec_dot_q4_1_q8_1_mul_mat( + const int * __restrict__ x_ql, const half2 * __restrict__ x_dm, const int * __restrict__ x_qh, const int * __restrict__ x_sc, + const int * __restrict__ y_qs, const half2 * __restrict__ y_ds, const int & i, const int & j, const int & k) { + GGML_UNUSED(x_qh); GGML_UNUSED(x_sc); + + const int kyqs = k % (QI8_1/2) + QI8_1 * (k / (QI8_1/2)); + + int u[2*VDR_Q4_1_Q8_1_MMQ]; + +#pragma unroll + for (int l = 0; l < VDR_Q4_1_Q8_1_MMQ; ++l) { + u[2*l+0] = y_qs[j * WARP_SIZE + (kyqs + l) % WARP_SIZE]; + u[2*l+1] = y_qs[j * WARP_SIZE + (kyqs + l + QI4_1) % WARP_SIZE]; + } + + return vec_dot_q4_1_q8_1_impl + (&x_ql[i * (WARP_SIZE + 1) + k], u, x_dm[i * (WARP_SIZE/QI4_1) + i/QI4_1 + k/QI4_1], + y_ds[j * (WARP_SIZE/QI8_1) + (2*k/QI8_1) % (WARP_SIZE/QI8_1)]); +} + + +extern "C" __global__ void + mul_mat_q4_0( + const void * __restrict__ vx, const void * __restrict__ vy, float * __restrict__ dst, + const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y, const int nrows_dst) { + const int mmq_x = MMQ_X_Q4_0_AMPERE; + const int mmq_y = MMQ_Y_Q4_0_AMPERE; + const int nwarps = NWARPS_Q4_0_AMPERE; + + mul_mat_q, + load_tiles_q4_0, VDR_Q4_0_Q8_1_MMQ, vec_dot_q4_0_q8_1_mul_mat> + (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); +} + +extern "C" __global__ void + mul_mat_q4_1( + const void * __restrict__ vx, const void * __restrict__ vy, float * __restrict__ dst, + const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y, const int nrows_dst) { + const int mmq_x = MMQ_X_Q4_1_AMPERE; + const int mmq_y = MMQ_Y_Q4_1_AMPERE; + const int nwarps = NWARPS_Q4_1_AMPERE; + + mul_mat_q, + load_tiles_q4_1, VDR_Q4_1_Q8_1_MMQ, vec_dot_q4_1_q8_1_mul_mat> + (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); +} + + +extern "C" __global__ void + mul_mat_q5_0( + const void * __restrict__ vx, const void * __restrict__ vy, float * __restrict__ dst, + const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y, const int nrows_dst) { + const int mmq_x = MMQ_X_Q5_0_AMPERE; + const int mmq_y = MMQ_Y_Q5_0_AMPERE; + const int nwarps = NWARPS_Q5_0_AMPERE; + + mul_mat_q, + load_tiles_q5_0, VDR_Q5_0_Q8_1_MMQ, vec_dot_q5_0_q8_1_mul_mat> + (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); +} + +extern "C" __global__ void +mul_mat_q5_1( + const void * __restrict__ vx, const void * __restrict__ vy, float * __restrict__ dst, + const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y, const int nrows_dst) { + const int mmq_x = MMQ_X_Q5_1_AMPERE; + const int mmq_y = MMQ_Y_Q5_1_AMPERE; + const int nwarps = NWARPS_Q5_1_AMPERE; + + mul_mat_q, + load_tiles_q5_1, VDR_Q5_1_Q8_1_MMQ, vec_dot_q5_1_q8_1_mul_mat> + (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); +} + +extern "C" __global__ void + mul_mat_q8_0( + const void * __restrict__ vx, const void * __restrict__ vy, float * __restrict__ dst, + const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y, const int nrows_dst) { + const int mmq_x = MMQ_X_Q8_0_AMPERE; + const int mmq_y = MMQ_Y_Q8_0_AMPERE; + const int nwarps = NWARPS_Q8_0_AMPERE; + + mul_mat_q, + load_tiles_q8_0, VDR_Q8_0_Q8_1_MMQ, vec_dot_q8_0_q8_1_mul_mat> + (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); +} + +extern "C" __global__ void +mul_mat_q2_K( + const void * __restrict__ vx, const void * __restrict__ vy, float * __restrict__ dst, + const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y, const int nrows_dst) { + const int mmq_x = MMQ_X_Q2_K_AMPERE; + const int mmq_y = MMQ_Y_Q2_K_AMPERE; + const int nwarps = NWARPS_Q2_K_AMPERE; + mul_mat_q, + load_tiles_q2_K, VDR_Q2_K_Q8_1_MMQ, vec_dot_q2_K_q8_1_mul_mat> + (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); +} + +extern "C" __global__ void + mul_mat_q3_K( + const void * __restrict__ vx, const void * __restrict__ vy, float * __restrict__ dst, + const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y, const int nrows_dst) { + const int mmq_x = MMQ_X_Q3_K_AMPERE; + const int mmq_y = MMQ_Y_Q3_K_AMPERE; + const int nwarps = NWARPS_Q3_K_AMPERE; + mul_mat_q, + load_tiles_q3_K, VDR_Q3_K_Q8_1_MMQ, vec_dot_q3_K_q8_1_mul_mat> + (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); +} + +extern "C" __global__ void + mul_mat_q4_K( + const void * __restrict__ vx, const void * __restrict__ vy, float * __restrict__ dst, + const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y, const int nrows_dst) { + const int mmq_x = MMQ_X_Q4_K_AMPERE; + const int mmq_y = MMQ_Y_Q4_K_AMPERE; + const int nwarps = NWARPS_Q4_K_AMPERE; + mul_mat_q, + load_tiles_q4_K, VDR_Q4_K_Q8_1_MMQ, vec_dot_q4_K_q8_1_mul_mat> + (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); +} + +extern "C" __global__ void +mul_mat_q5_K( + const void * __restrict__ vx, const void * __restrict__ vy, float * __restrict__ dst, + const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y, const int nrows_dst) { + const int mmq_x = MMQ_X_Q5_K_AMPERE; + const int mmq_y = MMQ_Y_Q5_K_AMPERE; + const int nwarps = NWARPS_Q5_K_AMPERE; + mul_mat_q, + load_tiles_q5_K, VDR_Q5_K_Q8_1_MMQ, vec_dot_q5_K_q8_1_mul_mat> + (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); +} + +extern "C" __global__ void + mul_mat_q6_K( + const void * __restrict__ vx, const void * __restrict__ vy, float * __restrict__ dst, + const int ncols_x, const int nrows_x, const int ncols_y, const int nrows_y, const int nrows_dst) { + const int mmq_x = MMQ_X_Q6_K_AMPERE; + const int mmq_y = MMQ_Y_Q6_K_AMPERE; + const int nwarps = NWARPS_Q6_K_AMPERE; + mul_mat_q, + load_tiles_q6_K, VDR_Q6_K_Q8_1_MMQ, vec_dot_q6_K_q8_1_mul_mat> + (vx, vy, dst, ncols_x, nrows_x, ncols_y, nrows_y, nrows_dst); +} + + +/** + * @brief Performs an indexed, batched matrix-vector multiplication for quantized tensors (for MoE models). + * + * This kernel handles a batch of `total_tasks` independent operations. Each task consists + * of multiplying a Q8_1 quantized input vector with a Q4_K quantized weight matrix selected + * by an index. + * + * Parallelization Strategy: + * - The grid is 2D: gridDim.y corresponds to the task index, and gridDim.x corresponds to the row blocks of the output matrix. + * - `blockIdx.y`: Identifies which task to perform from the batch (`0` to `total_tasks - 1`). + * - `blockIdx.x`: Used internally by `mul_mat_vec_q` to parallelize the dot products across the rows of the weight matrix. + * + * @author + * Guoqing Bao + * Part of the project: https://github.com/guoqingbao/vllm.rs/ + * @param all_weights Pointer to the beginning of the weight tensor [num_experts, n, k]. + * @param all_inputs Pointer to the beginning of the input tensor [batch * topk, k]. + * @param indices Pointer to the expert indices for each task [batch * topk]. + * @param all_outputs Pointer to the beginning of the output tensor [batch * topk, n]. + * @param n The number of output features (rows in the weight matrix). + * @param k The number of input features (columns in the weight matrix). + * @param total_tasks The total number of tasks to process, typically batch_size * topk. + * @param k_padded The value of k padded to a multiple of MATRIX_ROW_PADDING. + * @param weight_expert_stride_bytes The stride in bytes to get from one expert matrix to the next. + * @param input_task_stride_bytes The stride in bytes to get from one quantized input vector to the next. + * @param output_task_stride_elems The stride in elements (f32) to get from one output vector to the next. + */ +template +__device__ void indexed_moe_forward( + const void * __restrict__ all_weights, + const void * __restrict__ all_inputs, + const unsigned int * __restrict__ indices, + float * __restrict__ all_outputs, + const int n, + const int k, + const int batch, + const int topk, + const int k_padded, + const int input_dim1) { + + // `blockIdx.y` corresponds to the batch index (0 to batch_size-1) + const int current_batch = blockIdx.y; + // `blockIdx.z` corresponds to the topk index (0 to topk-1) + const int current_topk = blockIdx.z; + + // `gridDim.z` is the number of blocks in the z-dim, which is `topk`. + // This correctly flattens the (batch, topk) index into a single task ID. + const int task_id = current_batch * gridDim.z + current_topk; + if (task_id >= gridDim.y * gridDim.z) { + return; + } + // If input_dim1 is 1, all experts in a batch use the same input vector. + // Otherwise, each expert has a unique input vector. + const int input_idx = (input_dim1 == 1) ? current_batch : task_id; + + // The expert to use is found in the `indices` array at the flattened `task_id`. + const unsigned int expert_id = indices[task_id]; + + // Calculate strides + const size_t weight_block_size = sizeof(block_q_t); + const size_t input_block_size = sizeof(block_q8_1); + const size_t weight_expert_stride_bytes = (size_t)(n * k) / QK_K * weight_block_size; + const size_t input_task_stride_bytes = (size_t)k_padded / QK8_1 * input_block_size; + const size_t output_task_stride_elems = n; + + //data offsets of current task + const void * current_input_ptr = (const char *)all_inputs + input_idx * input_task_stride_bytes; + const void * current_weight_ptr = (const char *)all_weights + expert_id * weight_expert_stride_bytes; + float * current_output_ptr = all_outputs + task_id * output_task_stride_elems; + + //fixed for inner compute + constexpr int ncols_y = 1; + constexpr int nwarps = 4; + constexpr int rows_per_cuda_block = 1; + + const int tid = WARP_SIZE * threadIdx.y + threadIdx.x; + const int row0 = rows_per_cuda_block * blockIdx.x; // `blockIdx.x` is the row within the task + + if (row0 >= n) { + return; + } + + const int blocks_per_row_x = k / qk; + const int blocks_per_col_y = k_padded / QK8_1; + constexpr int blocks_per_iter = vdr * nwarps * WARP_SIZE / qi; + + float tmp = 0.0f; + + const block_q_t * w = (const block_q_t *) current_weight_ptr; + const block_q8_1 * x = (const block_q8_1 *) current_input_ptr; + + for (int kbx = tid / (qi / vdr); kbx < blocks_per_row_x; kbx += blocks_per_iter) { + const int kby = kbx * (qk / QK8_1); + const int kqs = vdr * (tid % (qi / vdr)); + tmp += vec_dot_q_cuda(&w[kbx + row0 * blocks_per_row_x], &x[kby], kqs); + } + + // --- Inter-warp reduction using shared memory --- + __shared__ float tmp_shared[nwarps - 1][WARP_SIZE]; + if (threadIdx.y > 0) { + tmp_shared[threadIdx.y - 1][threadIdx.x] = tmp; + } + __syncthreads(); + + if (threadIdx.y == 0) { + for (int l = 0; l < nwarps - 1; ++l) { + tmp += tmp_shared[l][threadIdx.x]; + } + tmp = warp_reduce_sum(tmp); + if (threadIdx.x == 0) { + current_output_ptr[row0] = tmp; + } + } +} + +extern "C" __global__ void indexed_moe_forward_q2k_q8_1( + const void * __restrict__ all_weights, + const void * __restrict__ all_inputs, + const unsigned int * __restrict__ indices, + float * __restrict__ all_outputs, + const int n, + const int k, + const int batch, + const int topk, + const int k_padded, + const int input_dim1) { + indexed_moe_forward + (all_weights, all_inputs, indices, all_outputs, n, k, batch, topk, k_padded, input_dim1); +} + +extern "C" __global__ void indexed_moe_forward_q3k_q8_1( + const void * __restrict__ all_weights, + const void * __restrict__ all_inputs, + const unsigned int * __restrict__ indices, + float * __restrict__ all_outputs, + const int n, + const int k, + const int batch, + const int topk, + const int k_padded, + const int input_dim1) { + indexed_moe_forward + (all_weights, all_inputs, indices, all_outputs, n, k, batch, topk, k_padded, input_dim1); +} + +extern "C" __global__ void indexed_moe_forward_q4k_q8_1( + const void * __restrict__ all_weights, + const void * __restrict__ all_inputs, + const unsigned int * __restrict__ indices, + float * __restrict__ all_outputs, + const int n, + const int k, + const int batch, + const int topk, + const int k_padded, + const int input_dim1) { + indexed_moe_forward + (all_weights, all_inputs, indices, all_outputs, n, k, batch, topk, k_padded, input_dim1); +} + +extern "C" __global__ void indexed_moe_forward_q5k_q8_1( + const void * __restrict__ all_weights, + const void * __restrict__ all_inputs, + const unsigned int * __restrict__ indices, + float * __restrict__ all_outputs, + const int n, + const int k, + const int batch, + const int topk, + const int k_padded, + const int input_dim1) { + indexed_moe_forward + (all_weights, all_inputs, indices, all_outputs, n, k, batch, topk, k_padded, input_dim1); +} + +extern "C" __global__ void indexed_moe_forward_q6k_q8_1( + const void * __restrict__ all_weights, + const void * __restrict__ all_inputs, + const unsigned int * __restrict__ indices, + float * __restrict__ all_outputs, + const int n, + const int k, + const int batch, + const int topk, + const int k_padded, + const int input_dim1) { + indexed_moe_forward + (all_weights, all_inputs, indices, all_outputs, n, k, batch, topk, k_padded, input_dim1); +} + +extern "C" __global__ void indexed_moe_forward_q8_0_q8_1( + const void * __restrict__ all_weights, + const void * __restrict__ all_inputs, + const unsigned int * __restrict__ indices, + float * __restrict__ all_outputs, + const int n, + const int k, + const int batch, + const int topk, + const int k_padded, + const int input_dim1) { + indexed_moe_forward + (all_weights, all_inputs, indices, all_outputs, n, k, batch, topk, k_padded, input_dim1); +} diff --git a/torchnative/rust/vendor/candle-kernels/src/reduce.cu b/torchnative/rust/vendor/candle-kernels/src/reduce.cu new file mode 100644 index 00000000..1dbb41c5 --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/reduce.cu @@ -0,0 +1,692 @@ +#include "cuda_utils.cuh" +#include +#include +#include + +#define WARP_SIZE 32 +const int BLOCK_SIZE = 1024; + +// Helpers to initialize reduction identities for both floating-point and +// integer types. For floats we keep using +/-INFINITY, while for integers +// we use well-defined numeric_limits values instead of relying on casting +// +/-INFINITY to an integer type (which is undefined behaviour and has been +// observed to break on newer GPU architectures such as Blackwell). +template +__device__ __forceinline__ T reduce_init_lowest() { + // Default implementation is used for floating-point types (__half, + // __nv_bfloat16, float, double). The conversion from -INFINITY (double) + // to these types is well-defined and produces -inf. + return -INFINITY; +} + +template +__device__ __forceinline__ T reduce_init_highest() { + // Default implementation is used for floating-point types (__half, + // __nv_bfloat16, float, double). The conversion from INFINITY (double) + // to these types is well-defined and produces +inf. + return INFINITY; +} + +// Integer specializations – use numeric_limits instead of +/-INFINITY. +template <> +__device__ __forceinline__ int64_t reduce_init_lowest() { + return ::cuda::std::numeric_limits::lowest(); +} + +template <> +__device__ __forceinline__ uint32_t reduce_init_lowest() { + return ::cuda::std::numeric_limits::lowest(); +} + +template <> +__device__ __forceinline__ uint8_t reduce_init_lowest() { + return ::cuda::std::numeric_limits::lowest(); +} + +template <> +__device__ __forceinline__ int64_t reduce_init_highest() { + return ::cuda::std::numeric_limits::max(); +} + +template <> +__device__ __forceinline__ uint32_t reduce_init_highest() { + return ::cuda::std::numeric_limits::max(); +} + +template <> +__device__ __forceinline__ uint8_t reduce_init_highest() { + return ::cuda::std::numeric_limits::max(); +} + +// TODO: Maybe add some fast_sum_f16_f32 variant that not only accumulate in f32 +// but also expect a f32 output so that this can be used for normalization e.g. +// in softmax. + +// Fast reduce sum kernel, this assumes that the dimensions to loop over are at +// the end, each block is responsible for populating one value in the output +// array. There are at most 1024 threads per block. +template +__device__ void +fast_sum(const size_t src_numel, const size_t el_to_sum_per_block, + const size_t num_dims, const size_t *info, const T *src, T *dst) { + const size_t *dims = info; + const size_t *strides = info + num_dims; + + __shared__ T shr[BLOCK_SIZE]; + size_t tid = threadIdx.x; + size_t dst_id = blockIdx.x; + + shr[tid] = 0; + // Elements summed in this block range from dst_id * el_to_sum_per_block + // to (dst_id + 1) * el_to_sum_per_block. + size_t start_idx = dst_id * el_to_sum_per_block; + size_t stop_idx = min(start_idx + el_to_sum_per_block, src_numel); + size_t idx = start_idx + tid; + + while (idx < stop_idx) { + // TODO: Fast version for the contiguous case. + size_t strided_i = get_strided_index(idx, num_dims, dims, strides); + shr[tid] += src[strided_i]; + idx += blockDim.x; + } + + // Parallel reduction, see the slides: + // https://www.olcf.ornl.gov/wp-content/uploads/2019/12/05_Atomics_Reductions_Warp_Shuffle.pdf + // https://stackoverflow.com/questions/66078814/is-cuda-atomicadd-operation-faster-than-launch-another-kernel-when-we-do-reduce + for (int s = blockDim.x / 2; s > 0; s >>= 1) { + __syncthreads(); + if (tid < s) + shr[tid] += shr[tid + s]; + } + + if (tid == 0) + dst[dst_id] = shr[0]; +} + +static __device__ __forceinline__ float2 warp_reduce_sum(float2 a) { +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + a.x += __shfl_xor_sync(0xffffffff, a.x, mask, 32); + a.y += __shfl_xor_sync(0xffffffff, a.y, mask, 32); + } + return a; +} + +static __device__ __forceinline__ float warp_reduce_sum(float x) { +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + x += __shfl_xor_sync(0xffffffff, x, mask, 32); + } + return x; +} + +// LayerNorm implementation adapted from ggml, accumulation is made using f32. +// https://github.com/ggerganov/llama.cpp/blob/d59bd97065cd7ded6c4ecab54b1d5e0b1b11e318/ggml-cuda.cu#L477 +template +__device__ void layernorm(const T * x, T * dst, const T * alpha, const T * beta, const int ncols, const int block_size, const float eps) { + const int row = blockIdx.x*blockDim.y + threadIdx.y; + const int tid = threadIdx.x; + + float2 mean_var = make_float2(0.f, 0.f); + + for (int col = tid; col < ncols; col += block_size) { + const float xi = x[row*ncols + col]; + mean_var.x += xi; + mean_var.y += xi * xi; + } + + // sum up partial sums + mean_var = warp_reduce_sum(mean_var); + if (block_size > WARP_SIZE) { + __shared__ float2 s_sum[32]; + int warp_id = threadIdx.x / WARP_SIZE; + int lane_id = threadIdx.x % WARP_SIZE; + if (lane_id == 0) { + s_sum[warp_id] = mean_var; + } + __syncthreads(); + mean_var = s_sum[lane_id]; + mean_var = warp_reduce_sum(mean_var); + } + + const float mean = mean_var.x / ncols; + const float var = mean_var.y / ncols - mean * mean; + const float inv_std = rsqrtf(var + eps); + + if (alpha == nullptr && beta == nullptr) { + for (int col = tid; col < ncols; col += block_size) { + float lhs = (static_cast(x[row*ncols + col]) - mean) * inv_std; + dst[row*ncols + col] = static_cast(lhs); + } + } + else if (alpha == nullptr && beta != nullptr) { + for (int col = tid; col < ncols; col += block_size) { + float b = static_cast(beta[col]); + float lhs = (static_cast(x[row*ncols + col]) - mean) * inv_std; + dst[row*ncols + col] = static_cast(lhs + b); + } + } + else if (alpha != nullptr && beta == nullptr) { + for (int col = tid; col < ncols; col += block_size) { + float a = static_cast(alpha[col]); + float lhs = (static_cast(x[row*ncols + col]) - mean) * inv_std; + dst[row*ncols + col] = static_cast(lhs * a); + } + } + else { + for (int col = tid; col < ncols; col += block_size) { + float a = static_cast(alpha[col]); + float b = static_cast(beta[col]); + float lhs = (static_cast(x[row*ncols + col]) - mean) * inv_std; + dst[row*ncols + col] = static_cast(lhs * a + b); + } + } +} + +// RmsNorm implementation adapted from ggml, accumulation is made using f32. +// https://github.com/ggerganov/llama.cpp/blob/d59bd97065cd7ded6c4ecab54b1d5e0b1b11e318/ggml-cuda.cu#L523 +template +__device__ void rmsnorm(const T * x, T * dst, const T * alpha, const int ncols, const int block_size, const float eps) { + const int row = blockIdx.x*blockDim.y + threadIdx.y; + const int tid = threadIdx.x; + + float tmp = 0.0f; // partial sum for thread in warp + + for (int col = tid; col < ncols; col += block_size) { + const float xi = static_cast(x[row*ncols + col]); + tmp += xi * xi; + } + + // sum up partial sums + tmp = warp_reduce_sum(tmp); + if (block_size > WARP_SIZE) { + __shared__ float s_sum[32]; + int warp_id = threadIdx.x / WARP_SIZE; + int lane_id = threadIdx.x % WARP_SIZE; + if (lane_id == 0) { + s_sum[warp_id] = tmp; + } + __syncthreads(); + tmp = s_sum[lane_id]; + tmp = warp_reduce_sum(tmp); + } + + const float mean = tmp / ncols; + const float scale = rsqrtf(mean + eps); + + if (alpha == nullptr) { + for (int col = tid; col < ncols; col += block_size) { + dst[row*ncols + col] = static_cast(scale * static_cast(x[row*ncols + col])); + } + } + else { + for (int col = tid; col < ncols; col += block_size) { + float a = static_cast(alpha[col]); + dst[row*ncols + col] = static_cast(scale * static_cast(x[row*ncols + col]) * a); + } + } +} + +// Softmax implementation adapted from ggml. +// https://github.com/ggerganov/llama.cpp/blob/d59bd97065cd7ded6c4ecab54b1d5e0b1b11e318/ggml-cuda.cu#L4159 +template +__device__ void softmax(const T * x, T * dst, const int ncols) { + const int row = blockDim.x*blockIdx.x + threadIdx.x; + const int block_size = blockDim.y; + const int tid = threadIdx.y; + + T max_val = -INFINITY; + + for (int col = tid; col < ncols; col += block_size) { + const int i = row*ncols + col; + max_val = maxg(max_val, x[i]); + } + + // find the max value in the block +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + max_val = maxg(max_val, __shfl_xor_sync(0xffffffff, max_val, mask, 32)); + } + + ACC tmp = 0.; + + for (int col = tid; col < ncols; col += block_size) { + const int i = row*ncols + col; + const T val = expg(x[i] - max_val); + tmp += static_cast(val); + dst[i] = val; + } + + // sum up partial sums +#pragma unroll + for (int mask = 16; mask > 0; mask >>= 1) { + tmp += __shfl_xor_sync(0xffffffff, tmp, mask, 32); + } + + const ACC inv_tmp = 1. / tmp; + + for (int col = tid; col < ncols; col += block_size) { + const int i = row*ncols + col; + dst[i] *= inv_tmp; + } +} + +template +__device__ void ropei(const T * src, const T * cos, const T * sin, T * dst, const uint32_t bh, const uint32_t td, const uint32_t stride_b) { + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (2 * idx >= bh * td) return; + + uint32_t rope_idx = idx % (td / 2); + if (stride_b > 0) { + uint32_t b_idx = (2 * idx) / stride_b; + rope_idx += b_idx * (td / 2); + } + T c = cos[rope_idx]; + T s = sin[rope_idx]; + + dst[2 * idx] = src[2 * idx] * c - src[2 * idx + 1] * s; + dst[2 * idx + 1] = src[2 * idx] * s + src[2 * idx + 1] * c; +} + +template +__device__ void rope(const T * src, const T * cos, const T * sin, T * dst, const uint32_t bh, const uint32_t td, const uint32_t d, const uint32_t stride_b) { + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (2 * idx >= bh * td) return; + + uint32_t i_bh = idx / (td / 2); + uint32_t i_td = idx - (td / 2) * i_bh; + uint32_t i_t = i_td / (d / 2); + uint32_t i_d = i_td - (d / 2) * i_t; + uint32_t i1 = i_bh * td + i_t * d + i_d; + uint32_t i2 = i1 + d / 2; + uint32_t i_cs = i_t * (d / 2) + i_d; + if (stride_b > 0) { + uint32_t b_idx = (2 * idx) / stride_b; + i_cs += b_idx * (td / 2); + } + T c = cos[i_cs]; + T s = sin[i_cs]; + + dst[i1] = src[i1] * c - src[i2] * s; + dst[i2] = src[i1] * s + src[i2] * c; +} + +template +__device__ void rope_thd( + const T * src, + const T * cos, + const T * sin, + T * dst, + const uint32_t b, + const uint32_t t, + const uint32_t h, + const uint32_t d, + const uint32_t stride_b +) { + const int idx = blockIdx.x * blockDim.x + threadIdx.x; + if (2 * idx >= b * t * h * d) return; + + uint32_t i_bth = idx / (d / 2); + uint32_t i_d = idx - (d / 2) * i_bth; + uint32_t i_t = (i_bth / h) % t; + uint32_t i1 = i_bth * d + i_d; + uint32_t i2 = i1 + d / 2; + uint32_t i_cs = i_t * (d / 2) + i_d; + if (stride_b > 0) { + uint32_t b_idx = (2 * idx) / stride_b; + i_cs += b_idx * ((t * d) / 2); + } + T c = cos[i_cs]; + T s = sin[i_cs]; + + dst[i1] = src[i1] * c - src[i2] * s; + dst[i2] = src[i1] * s + src[i2] * c; +} + +template +__device__ void +fast_max(const size_t src_numel, const size_t el_to_sum_per_block, + const size_t num_dims, const size_t *info, const T *src, T *dst) { + const size_t *dims = info; + const size_t *strides = info + num_dims; + + __shared__ T shr[BLOCK_SIZE]; + size_t tid = threadIdx.x; + size_t dst_id = blockIdx.x; + + // Initialize with the lowest representable value for T so that the first + // comparison in the reduction always picks a real element. + shr[tid] = reduce_init_lowest(); + // Elements summed in this block range from dst_id * el_to_sum_per_block + // to (dst_id + 1) * el_to_sum_per_block. + size_t start_idx = dst_id * el_to_sum_per_block; + size_t stop_idx = min(start_idx + el_to_sum_per_block, src_numel); + size_t idx = start_idx + tid; + + while (idx < stop_idx) { + // TODO: Fast version for the contiguous case. + size_t strided_i = get_strided_index(idx, num_dims, dims, strides); + shr[tid] = maxg(shr[tid], src[strided_i]); + idx += blockDim.x; + } + + // Parallel reduction, see the slides: + // https://www.olcf.ornl.gov/wp-content/uploads/2019/12/05_Atomics_Reductions_Warp_Shuffle.pdf + // https://stackoverflow.com/questions/66078814/is-cuda-atomicadd-operation-faster-than-launch-another-kernel-when-we-do-reduce + for (int s = blockDim.x / 2; s > 0; s >>= 1) { + __syncthreads(); + if (tid < s) + shr[tid] = maxg(shr[tid], shr[tid + s]); + } + + if (tid == 0) + dst[dst_id] = shr[0]; +} + +template +__device__ void +fast_min(const size_t src_numel, const size_t el_to_sum_per_block, + const size_t num_dims, const size_t *info, const T *src, T *dst) { + const size_t *dims = info; + const size_t *strides = info + num_dims; + + __shared__ T shr[BLOCK_SIZE]; + size_t tid = threadIdx.x; + size_t dst_id = blockIdx.x; + + // Initialize with the highest representable value for T so that the first + // comparison in the reduction always picks a real element. + shr[tid] = reduce_init_highest(); + // Elements summed in this block range from dst_id * el_to_sum_per_block + // to (dst_id + 1) * el_to_sum_per_block. + size_t start_idx = dst_id * el_to_sum_per_block; + size_t stop_idx = min(start_idx + el_to_sum_per_block, src_numel); + size_t idx = start_idx + tid; + + while (idx < stop_idx) { + // TODO: Fast version for the contiguous case. + size_t strided_i = get_strided_index(idx, num_dims, dims, strides); + shr[tid] = ming(shr[tid], src[strided_i]); + idx += blockDim.x; + } + + // Parallel reduction, see the slides: + // https://www.olcf.ornl.gov/wp-content/uploads/2019/12/05_Atomics_Reductions_Warp_Shuffle.pdf + // https://stackoverflow.com/questions/66078814/is-cuda-atomicadd-operation-faster-than-launch-another-kernel-when-we-do-reduce + for (int s = blockDim.x / 2; s > 0; s >>= 1) { + __syncthreads(); + if (tid < s) + shr[tid] = ming(shr[tid], shr[tid + s]); + } + + if (tid == 0) + dst[dst_id] = shr[0]; +} + +template +__device__ void +fast_argmin(const size_t src_numel, const size_t el_to_sum_per_block, + const size_t num_dims, const size_t *info, const T *src, uint32_t *dst) { + const size_t *dims = info; + const size_t *strides = info + num_dims; + + __shared__ T shr[BLOCK_SIZE]; + __shared__ uint32_t shr_index[BLOCK_SIZE]; + size_t tid = threadIdx.x; + size_t dst_id = blockIdx.x; + + // For floating types this uses +inf; for integer types we use the largest + // representable value instead of casting INFINITY to an integer. + shr[tid] = reduce_init_highest(); + shr_index[tid] = 0xFFFFFFFF; + bool not_set = true; + // Elements summed in this block range from dst_id * el_to_sum_per_block + // to (dst_id + 1) * el_to_sum_per_block. + size_t start_idx = dst_id * el_to_sum_per_block; + size_t stop_idx = min(start_idx + el_to_sum_per_block, src_numel); + size_t idx = start_idx + tid; + + while (idx < stop_idx) { + // TODO: Fast version for the contiguous case. + size_t strided_i = get_strided_index(idx, num_dims, dims, strides); + if (not_set || src[strided_i] < shr[tid]) { + shr[tid] = src[strided_i]; + // Assume that the reduction takes place over the last dimension which is contiguous. + shr_index[tid] = idx % dims[num_dims - 1]; + not_set = false; + } + idx += blockDim.x; + } + + // Parallel reduction, see the slides: + // https://www.olcf.ornl.gov/wp-content/uploads/2019/12/05_Atomics_Reductions_Warp_Shuffle.pdf + // https://stackoverflow.com/questions/66078814/is-cuda-atomicadd-operation-faster-than-launch-another-kernel-when-we-do-reduce + for (int s = blockDim.x / 2; s > 0; s >>= 1) { + __syncthreads(); + if (tid < s && shr[tid + s] < shr[tid]) { + shr[tid] = shr[tid + s]; + shr_index[tid] = shr_index[tid + s]; + } + } + + if (tid == 0) + dst[dst_id] = shr_index[0]; +} + +template +__device__ void +fast_argmax(const size_t src_numel, const size_t el_to_sum_per_block, + const size_t num_dims, const size_t *info, const T *src, uint32_t *dst) { + const size_t *dims = info; + const size_t *strides = info + num_dims; + + __shared__ T shr[BLOCK_SIZE]; + __shared__ uint32_t shr_index[BLOCK_SIZE]; + size_t tid = threadIdx.x; + size_t dst_id = blockIdx.x; + + // For floating types this uses -inf; for integer types we use the lowest + // representable value instead of casting -INFINITY to an integer. + shr[tid] = reduce_init_lowest(); + shr_index[tid] = 0xFFFFFFFF; + bool not_set = true; + // Elements summed in this block range from dst_id * el_to_sum_per_block + // to (dst_id + 1) * el_to_sum_per_block. + size_t start_idx = dst_id * el_to_sum_per_block; + size_t stop_idx = min(start_idx + el_to_sum_per_block, src_numel); + size_t idx = start_idx + tid; + + while (idx < stop_idx) { + // TODO: Fast version for the contiguous case. + size_t strided_i = get_strided_index(idx, num_dims, dims, strides); + if (not_set || src[strided_i] > shr[tid]) { + shr[tid] = src[strided_i]; + // Assume that the reduction takes place over the last dimension which is contiguous. + shr_index[tid] = idx % dims[num_dims - 1]; + not_set = false; + } + idx += blockDim.x; + } + + // Parallel reduction, see the slides: + // https://www.olcf.ornl.gov/wp-content/uploads/2019/12/05_Atomics_Reductions_Warp_Shuffle.pdf + // https://stackoverflow.com/questions/66078814/is-cuda-atomicadd-operation-faster-than-launch-another-kernel-when-we-do-reduce + for (int s = blockDim.x / 2; s > 0; s >>= 1) { + __syncthreads(); + if (tid < s && shr[tid + s] > shr[tid]) { + shr[tid] = shr[tid + s]; + shr_index[tid] = shr_index[tid + s]; + } + } + + if (tid == 0) + dst[dst_id] = shr_index[0]; +} + +#define FAST_OP(TYPENAME, MIN_NAME, MAX_NAME, ARGMIN_NAME, ARGMAX_NAME, SUM_NAME) \ + extern "C" __global__ void ARGMIN_NAME( \ + const size_t src_numel, const size_t el_to_sum_per_block, \ + const size_t num_dims, const size_t *info, const TYPENAME *src, \ + uint32_t *dst) { \ + fast_argmin(src_numel, el_to_sum_per_block, num_dims, info, src, dst); \ + } \ + extern "C" __global__ void ARGMAX_NAME( \ + const size_t src_numel, const size_t el_to_sum_per_block, \ + const size_t num_dims, const size_t *info, const TYPENAME *src, \ + uint32_t *dst) { \ + fast_argmax(src_numel, el_to_sum_per_block, num_dims, info, src, dst); \ + } \ + extern "C" __global__ void MIN_NAME( \ + const size_t src_numel, const size_t el_to_sum_per_block, \ + const size_t num_dims, const size_t *info, const TYPENAME *src, \ + TYPENAME *dst) { \ + fast_min(src_numel, el_to_sum_per_block, num_dims, info, src, dst); \ + } \ + extern "C" __global__ void MAX_NAME( \ + const size_t src_numel, const size_t el_to_sum_per_block, \ + const size_t num_dims, const size_t *info, const TYPENAME *src, \ + TYPENAME *dst) { \ + fast_max(src_numel, el_to_sum_per_block, num_dims, info, src, dst); \ + } \ + extern "C" __global__ void SUM_NAME( \ + const size_t src_numel, const size_t el_to_sum_per_block, \ + const size_t num_dims, const size_t *info, const TYPENAME *src, \ + TYPENAME *dst) { \ + fast_sum(src_numel, el_to_sum_per_block, num_dims, info, src, dst); \ + } + +#define SUM_OP(TYPENAME, FN_NAME) \ + extern "C" __global__ void FN_NAME( \ + const size_t numel, const size_t num_dims, const size_t num_sum_dims, \ + const size_t *info, const TYPENAME *inp, TYPENAME *out) { \ + const size_t *dims = info; \ + const size_t *strides = info + num_dims; \ + const size_t *sum_dims_l = info + 2 * num_dims; \ + const size_t *sum_dims_s = info + 2 * num_dims + num_sum_dims; \ + if (is_contiguous(num_dims, dims, strides)) { \ + for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; \ + i += blockDim.x * gridDim.x) { \ + size_t dst_index = i; \ + for (unsigned int nd = 0; nd < num_sum_dims; ++nd) { \ + size_t stride = sum_dims_s[nd]; \ + size_t pre = dst_index / stride; \ + size_t post = dst_index % stride; \ + dst_index = (pre / sum_dims_l[nd]) * stride + post; \ + } \ + atomicAdd(out + dst_index, inp[i]); \ + } \ + } else { \ + for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; \ + i += blockDim.x * gridDim.x) { \ + unsigned strided_i = get_strided_index(i, num_dims, dims, strides); \ + size_t dst_index = i; \ + for (unsigned int nd = 0; nd < num_sum_dims; ++nd) { \ + size_t stride = sum_dims_s[nd]; \ + size_t pre = dst_index / stride; \ + size_t post = dst_index % stride; \ + dst_index = (pre / sum_dims_l[nd]) * stride + post; \ + } \ + atomicAdd(out + dst_index, inp[strided_i]); \ + } \ + } \ + } + +#define SOFTMAX_OP(TYPENAME, ACC_TYPENAME, FN_NAME) \ + extern "C" __global__ void FN_NAME( \ + const TYPENAME *src, TYPENAME *dst, \ + const int n_cols) { \ + softmax(src, dst, n_cols); \ + } \ + +#define RMSNORM_OP(TYPENAME, FN_NAME) \ + extern "C" __global__ void FN_NAME( \ + const TYPENAME *src, TYPENAME *dst, const TYPENAME *alpha, \ + const int n_cols, const int block_size, const float eps) { \ + rmsnorm(src, dst, alpha, n_cols, block_size, eps); \ + } \ + +#define LAYERNORM_OP(TYPENAME, FN_NAME) \ + extern "C" __global__ void FN_NAME( \ + const TYPENAME *src, TYPENAME *dst, const TYPENAME *alpha, \ + const TYPENAME *beta, const int n_cols, const int block_size, const float eps) { \ + layernorm(src, dst, alpha, beta, n_cols, block_size, eps); \ + } \ + +#define ROPE_OP(TYPENAME, FN_NAME, FN_NAME_I, FN_NAME_THD) \ + extern "C" __global__ void FN_NAME_I( \ + const TYPENAME *src, \ + const TYPENAME *cos, \ + const TYPENAME *sin, \ + TYPENAME *dst, \ + const uint32_t bh, \ + const uint32_t td, \ + const uint32_t stride_b) { \ + ropei(src, cos, sin, dst, bh, td, stride_b); \ + } \ + extern "C" __global__ void FN_NAME( \ + const TYPENAME *src, \ + const TYPENAME *cos, \ + const TYPENAME *sin, \ + TYPENAME *dst, \ + const uint32_t bh, \ + const uint32_t td, \ + const uint32_t d, \ + const uint32_t stride_b) { \ + rope(src, cos, sin, dst, bh, td, d, stride_b); \ + } \ + extern "C" __global__ void FN_NAME_THD( \ + const TYPENAME *src, \ + const TYPENAME *cos, \ + const TYPENAME *sin, \ + TYPENAME *dst, \ + const uint32_t b, \ + const uint32_t t, \ + const uint32_t h, \ + const uint32_t d, \ + const uint32_t stride_b) { \ + rope_thd(src, cos, sin, dst, b, t, h, d, stride_b); \ + } \ + +#if __CUDA_ARCH__ >= 800 +SOFTMAX_OP(__nv_bfloat16, float, softmax_bf16) +RMSNORM_OP(__nv_bfloat16, rmsnorm_bf16) +LAYERNORM_OP(__nv_bfloat16, layernorm_bf16) +ROPE_OP(__nv_bfloat16, rope_bf16, rope_i_bf16, rope_thd_bf16) +SUM_OP(__nv_bfloat16, sum_bf16) +FAST_OP(__nv_bfloat16, fast_min_bf16, fast_max_bf16, fast_argmin_bf16, fast_argmax_bf16, fast_sum_bf16) + +// NOTE: No reduce ops for f8 +// SUM_OP(__nv_fp8_e4m3, sum_fp8_e4m3) +// SOFTMAX_OP(__nv_fp8_e4m3, float, softmax_fp8_e4m3) +// RMSNORM_OP(__nv_fp8_e4m3, rmsnorm_fp8_e4m3) +// LAYERNORM_OP(__nv_fp8_e4m3, layernorm_fp8_e4m3) +// ROPE_OP(__nv_fp8_e4m3, rope_fp8_e4m3, rope_i_fp8_e4m3, rope_thd_fp8_e4m3) +// FAST_OP(__nv_fp8_e4m3, fast_min_fp8_e4m3, fast_max_fp8_e4m3, fast_argmin_fp8_e4m3, fast_argmax_fp8_e4m3, fast_sum_fp8_e4m3) +#endif + +#if __CUDA_ARCH__ >= 530 +SOFTMAX_OP(__half, float, softmax_f16) +RMSNORM_OP(__half, rmsnorm_f16) +LAYERNORM_OP(__half, layernorm_f16) +ROPE_OP(__half, rope_f16, rope_i_f16, rope_thd_f16) +SUM_OP(__half, sum_f16) +FAST_OP(__half, fast_min_f16, fast_max_f16, fast_argmin_f16, fast_argmax_f16, fast_sum_f16) +#endif + +SUM_OP(float, sum_f32) +SUM_OP(double, sum_f64) +SUM_OP(uint32_t, sum_u32) +SOFTMAX_OP(float, float, softmax_f32) +SOFTMAX_OP(double, double, softmax_f64) +RMSNORM_OP(float, rmsnorm_f32) +RMSNORM_OP(double, rmsnorm_f64) +LAYERNORM_OP(float, layernorm_f32) +LAYERNORM_OP(double, layernorm_f64) +ROPE_OP(float, rope_f32, rope_i_f32, rope_thd_f32) +ROPE_OP(double, rope_f64, rope_i_f64, rope_thd_f64) + +FAST_OP(float, fast_min_f32, fast_max_f32, fast_argmin_f32, fast_argmax_f32, fast_sum_f32) +FAST_OP(double, fast_min_f64, fast_max_f64, fast_argmin_f64, fast_argmax_f64, fast_sum_f64) +FAST_OP(uint32_t, fast_min_u32, fast_max_u32, fast_argmin_u32, fast_argmax_u32, fast_sum_u32) +FAST_OP(int64_t, fast_min_i64, fast_max_i64, fast_argmin_i64, fast_argmax_i64, fast_sum_i64) +FAST_OP(uint8_t, fast_min_u8, fast_max_u8, fast_argmin_u8, fast_argmax_u8, fast_sum_u8) diff --git a/torchnative/rust/vendor/candle-kernels/src/sort.cu b/torchnative/rust/vendor/candle-kernels/src/sort.cu new file mode 100644 index 00000000..80ec69cd --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/sort.cu @@ -0,0 +1,90 @@ +// Adapted from https://github.com/ggerganov/llama.cpp/blob/master/ggml-cuda/argsort.cu +#define SORT_ORDER_ASC 1 +#define SORT_ORDER_DESC 0 +#include "cuda_utils.cuh" +#include + +template +static inline __device__ void ggml_cuda_swap(T & a, T & b) { + T tmp = a; + a = b; + b = tmp; +} + +template +static __device__ void k_argsort(const T * x, uint32_t * dst, const int ncols, int ncols_pad) { + // bitonic sort + int row = blockIdx.x; + + const T * x_row = x + row * ncols; + extern __shared__ int dst_row[]; + + // initialize indices - each thread handles multiple elements if ncols_pad > blockDim.x + for (int col = threadIdx.x; col < ncols_pad; col += blockDim.x) { + dst_row[col] = col; + } + + __syncthreads(); + + for (int k = 2; k <= ncols_pad; k *= 2) { + for (int j = k / 2; j > 0; j /= 2) { + for (int col = threadIdx.x; col < ncols_pad; col += blockDim.x) { + int ixj = col ^ j; + if (ixj > col) { + if ((col & k) == 0) { + if (dst_row[col] >= ncols || + (dst_row[ixj] < ncols && (order == SORT_ORDER_ASC ? + x_row[dst_row[col]] > x_row[dst_row[ixj]] : + x_row[dst_row[col]] < x_row[dst_row[ixj]])) + ) { + ggml_cuda_swap(dst_row[col], dst_row[ixj]); + } + } else { + if (dst_row[ixj] >= ncols || + (dst_row[col] < ncols && (order == SORT_ORDER_ASC ? + x_row[dst_row[col]] < x_row[dst_row[ixj]] : + x_row[dst_row[col]] > x_row[dst_row[ixj]])) + ) { + ggml_cuda_swap(dst_row[col], dst_row[ixj]); + } + } + } + } + __syncthreads(); + } + } + + // copy the result to dst without the padding + for (int col = threadIdx.x; col < ncols; col += blockDim.x) { + dst[row * ncols + col] = dst_row[col]; + } +} + +#define ASORT_OP(TYPENAME, RUST_NAME) \ +extern "C" __global__ void asort_asc_##RUST_NAME( \ + const TYPENAME * x, uint32_t * dst, const int ncols, int ncols_pad \ +) { \ + k_argsort(x, dst, ncols, ncols_pad); \ +} \ +extern "C" __global__ void asort_desc_##RUST_NAME( \ + const TYPENAME * x, uint32_t * dst, const int ncols, int ncols_pad \ +) { \ + k_argsort(x, dst, ncols, ncols_pad); \ +} \ + +#if __CUDA_ARCH__ >= 800 +ASORT_OP(__nv_bfloat16, bf16) + +// NOTE: No sort ops for f8 +// ASORT_OP(__nv_fp8_e4m3, fp8_e4m3) +#endif + +#if __CUDA_ARCH__ >= 530 +ASORT_OP(__half, f16) +#endif + +ASORT_OP(float, f32) +ASORT_OP(double, f64) +ASORT_OP(uint8_t, u8) +ASORT_OP(uint32_t, u32) +ASORT_OP(int64_t, i64) diff --git a/torchnative/rust/vendor/candle-kernels/src/ternary.cu b/torchnative/rust/vendor/candle-kernels/src/ternary.cu new file mode 100644 index 00000000..95a695ec --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/ternary.cu @@ -0,0 +1,71 @@ +#include "cuda_utils.cuh" +#include + +#define WHERE_OP(TYPENAME, ID_TYPENAME, FN_NAME) \ +extern "C" __global__ void FN_NAME( \ + const size_t numel, \ + const size_t num_dims, \ + const size_t *info, \ + const ID_TYPENAME *ids, \ + const TYPENAME *t, \ + const TYPENAME *f, \ + TYPENAME *out \ +) { \ + const size_t *dims = info; \ + const size_t *strides = info + num_dims; \ + const size_t *strides_t = info + 2*num_dims; \ + const size_t *strides_f = info + 3*num_dims; \ + if (is_contiguous(num_dims, dims, strides) \ + && is_contiguous(num_dims, dims, strides_f) \ + && is_contiguous(num_dims, dims, strides_t)) { \ + for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { \ + out[i] = ids[i] ? t[i] : f[i]; \ + } \ + } \ + else { \ + for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { \ + unsigned strided_i = get_strided_index(i, num_dims, dims, strides); \ + unsigned strided_i_t = get_strided_index(i, num_dims, dims, strides_t); \ + unsigned strided_i_f = get_strided_index(i, num_dims, dims, strides_f); \ + out[i] = ids[strided_i] ? t[strided_i_t] : f[strided_i_f]; \ + } \ + } \ +} \ + +#if __CUDA_ARCH__ >= 800 +WHERE_OP(__nv_bfloat16, int64_t, where_i64_bf16) +WHERE_OP(__nv_bfloat16, uint32_t, where_u32_bf16) +WHERE_OP(__nv_bfloat16, uint8_t, where_u8_bf16) +#endif + +#if __CUDA_ARCH__ >= 890 +WHERE_OP(__nv_fp8_e4m3, int16_t, where_i16_fp8_e4m3) +WHERE_OP(__nv_fp8_e4m3, int32_t, where_i32_fp8_e4m3) +WHERE_OP(__nv_fp8_e4m3, int64_t, where_i64_fp8_e4m3) +WHERE_OP(__nv_fp8_e4m3, uint32_t, where_u32_fp8_e4m3) +WHERE_OP(__nv_fp8_e4m3, uint8_t, where_u8_fp8_e4m3) +#endif + +#if __CUDA_ARCH__ >= 530 +WHERE_OP(__half, int64_t, where_i64_f16) +WHERE_OP(__half, uint32_t, where_u32_f16) +WHERE_OP(__half, uint8_t, where_u8_f16) +#endif + +WHERE_OP(float, int64_t, where_i64_f32) +WHERE_OP(double, int64_t, where_i64_f64) +WHERE_OP(uint8_t, int64_t, where_i64_u8) +WHERE_OP(uint32_t, int64_t, where_i64_u32) +WHERE_OP(int64_t, int64_t, where_i64_i64) + +WHERE_OP(float, uint32_t, where_u32_f32) +WHERE_OP(double, uint32_t, where_u32_f64) +WHERE_OP(uint8_t, uint32_t, where_u32_u8) +WHERE_OP(uint32_t, uint32_t, where_u32_u32) +WHERE_OP(int64_t, uint32_t, where_u32_i64) + +WHERE_OP(float, uint8_t, where_u8_f32) +WHERE_OP(double, uint8_t, where_u8_f64) +WHERE_OP(uint8_t, uint8_t, where_u8_u8) +WHERE_OP(uint32_t, uint8_t, where_u8_u32) +WHERE_OP(int64_t, uint8_t, where_u8_i64) diff --git a/torchnative/rust/vendor/candle-kernels/src/unary.cu b/torchnative/rust/vendor/candle-kernels/src/unary.cu new file mode 100644 index 00000000..3973b72b --- /dev/null +++ b/torchnative/rust/vendor/candle-kernels/src/unary.cu @@ -0,0 +1,233 @@ +#define _USE_MATH_DEFINES +#include +#include +#include "cuda_utils.cuh" + +#define UNARY_OP(TYPENAME, FN_NAME, FUNC) \ +extern "C" __global__ void FN_NAME( \ + const size_t numel, \ + const size_t num_dims, \ + const size_t *info, \ + const TYPENAME *inp, \ + TYPENAME *out \ +) { \ + const size_t *dims = info; \ + const size_t *strides = info + num_dims; \ + if (info == nullptr || is_contiguous(num_dims, dims, strides)) { \ + for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { \ + TYPENAME x = inp ? inp[i] : out[i]; \ + out[i] = FUNC; \ + } \ + } \ + else { \ + for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { \ + unsigned strided_i = get_strided_index(i, num_dims, dims, strides); \ + TYPENAME x = inp ? inp[strided_i] : out[i]; \ + out[i] = FUNC; \ + } \ + } \ +} \ + +template +__device__ __forceinline__ T gelu_erf_fwd(T x) { + return x * normcdfg(x); +} + +template +__device__ __forceinline__ T gelu_fwd(T x) { + T x_sq = x * x; + T x_cube = x_sq * x; + T alpha = x + static_cast(0.044715) * x_cube; + return static_cast(0.5) * x * (static_cast(1.0) + tanhg(static_cast(M_2_SQRTPI * M_SQRT1_2) * alpha)); +} + +template +__device__ __forceinline__ T elu_fwd(T x, T alpha) { + if (x > static_cast(0)) { + return x; + } + return alpha * (expg(x) - static_cast(1)); +} + +template +__device__ __forceinline__ T relu_fwd(T x) { + T zero = 0.; + return maxg(x, zero); +} + +template +__device__ __forceinline__ T silu_fwd(T x) { + return x / (static_cast(1) + expg(-x)); +} + +template +__device__ __forceinline__ T sigmoid_fwd(T x) { + return recipg(static_cast(1) + expg(-x)); +} + +#define UNARY_OP1(TYPENAME, FN_NAME, FUNC) \ +extern "C" __global__ void FN_NAME( \ + const size_t numel, \ + const size_t num_dims, \ + const size_t *info, \ + const TYPENAME param, \ + const TYPENAME *inp, \ + TYPENAME *out \ +) { \ + const size_t *dims = info; \ + const size_t *strides = info + num_dims; \ + if (info == nullptr || is_contiguous(num_dims, dims, strides)) { \ + for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { \ + TYPENAME x = inp ? inp[i] : out[i]; \ + out[i] = FUNC; \ + } \ + } \ + else { \ + for (unsigned int i = blockIdx.x * blockDim.x + threadIdx.x; i < numel; i += blockDim.x * gridDim.x) { \ + unsigned strided_i = get_strided_index(i, num_dims, dims, strides); \ + TYPENAME x = inp ? inp[strided_i] : out[i]; \ + out[i] = FUNC; \ + } \ + } \ +} \ + +template +__device__ T sign_(T t) { + return static_cast(t > static_cast(0)) - static_cast(t < static_cast(0)); +} + + +#if __CUDA_ARCH__ >= 800 +UNARY_OP(__nv_bfloat16, ucopy_bf16, x) +UNARY_OP(__nv_bfloat16, uneg_bf16, -x) +UNARY_OP(__nv_bfloat16, urecip_bf16, recipg(x)) +UNARY_OP(__nv_bfloat16, uexp_bf16, expg(x)) +UNARY_OP(__nv_bfloat16, ulog_bf16, logg(x)) +UNARY_OP(__nv_bfloat16, usin_bf16, sing(x)) +UNARY_OP(__nv_bfloat16, ucos_bf16, cosg(x)) +UNARY_OP(__nv_bfloat16, utanh_bf16, tanhg(x)) +UNARY_OP(__nv_bfloat16, uerf_bf16, erfg(x)) +UNARY_OP(__nv_bfloat16, uceil_bf16, ceilg(x)) +UNARY_OP(__nv_bfloat16, ufloor_bf16, floorg(x)) +UNARY_OP(__nv_bfloat16, uround_bf16, roundg(x)) +UNARY_OP(__nv_bfloat16, unormcdf_bf16, normcdfg(x)) +UNARY_OP(__nv_bfloat16, uabs_bf16, absg(x)) +UNARY_OP(__nv_bfloat16, usqr_bf16, x*x) +UNARY_OP(__nv_bfloat16, usqrt_bf16, sqrtg(x)) +UNARY_OP(__nv_bfloat16, ugelu_bf16, gelu_fwd(x)) +UNARY_OP(__nv_bfloat16, ugelu_erf_bf16, gelu_erf_fwd(x)) +UNARY_OP(__nv_bfloat16, urelu_bf16, relu_fwd(x)) +UNARY_OP1(__nv_bfloat16, uelu_bf16, elu_fwd(x, param)) +UNARY_OP(__nv_bfloat16, usilu_bf16, silu_fwd(x)) +UNARY_OP1(__nv_bfloat16, upowf_bf16, powg(x, param)) +UNARY_OP(__nv_bfloat16, usign_bf16, sign_(x)) +UNARY_OP(__nv_bfloat16, usigmoid_bf16, sigmoid_fwd(x)) +#endif + +#if __CUDA_ARCH__ >= 890 +#define F8E4M3_TO_FLOAT(x) __half2float(__nv_cvt_fp8_to_halfraw(x.__x, __NV_E4M3)) + +UNARY_OP(__nv_fp8_e4m3, ucopy_f8_e4m3, x) +UNARY_OP(__nv_fp8_e4m3, uneg_fp8_e4m3, __nv_fp8_e4m3(-F8E4M3_TO_FLOAT(x))) +UNARY_OP(__nv_fp8_e4m3, urecip_fp8_e4m3, recipg(x)) +UNARY_OP(__nv_fp8_e4m3, uexp_fp8_e4m3, expg(x)) +UNARY_OP(__nv_fp8_e4m3, ulog_fp8_e4m3, logg(x)) +UNARY_OP(__nv_fp8_e4m3, usin_fp8_e4m3, sing(x)) +UNARY_OP(__nv_fp8_e4m3, ucos_fp8_e4m3, cosg(x)) +UNARY_OP(__nv_fp8_e4m3, utanh_fp8_e4m3, tanhg(x)) +UNARY_OP(__nv_fp8_e4m3, uerf_fp8_e4m3, erfg(x)) +UNARY_OP(__nv_fp8_e4m3, uceil_fp8_e4m3, ceilg(x)) +UNARY_OP(__nv_fp8_e4m3, ufloor_fp8_e4m3, floorg(x)) +UNARY_OP(__nv_fp8_e4m3, uround_fp8_e4m3, roundg(x)) +UNARY_OP(__nv_fp8_e4m3, unormcdf_fp8_e4m3, normcdfg(x)) +UNARY_OP(__nv_fp8_e4m3, uabs_fp8_e4m3, absg(x)) +UNARY_OP(__nv_fp8_e4m3, usqr_fp8_e4m3, __nv_fp8_e4m3(F8E4M3_TO_FLOAT(x)*F8E4M3_TO_FLOAT(x))) +UNARY_OP(__nv_fp8_e4m3, usqrt_fp8_e4m3, sqrtg(x)) +UNARY_OP(__nv_fp8_e4m3, ugelu_fp8_e4m3, __nv_fp8_e4m3(gelu_fwd(F8E4M3_TO_FLOAT(x)))) +UNARY_OP(__nv_fp8_e4m3, ugelu_erf_fp8_e4m3, __nv_fp8_e4m3(gelu_erf_fwd(F8E4M3_TO_FLOAT(x)))) +UNARY_OP(__nv_fp8_e4m3, urelu_fp8_e4m3, __nv_fp8_e4m3(relu_fwd(F8E4M3_TO_FLOAT(x)))) +UNARY_OP1(__nv_fp8_e4m3, uelu_fp8_e4m3, __nv_fp8_e4m3(elu_fwd(F8E4M3_TO_FLOAT(x), F8E4M3_TO_FLOAT(param)))) +UNARY_OP(__nv_fp8_e4m3, usilu_fp8_e4m3, __nv_fp8_e4m3(silu_fwd(F8E4M3_TO_FLOAT(x)))) +UNARY_OP1(__nv_fp8_e4m3, upowf_fp8_e4m3, powg(x, param)) +UNARY_OP(__nv_fp8_e4m3, usign_fp8_e4m3, __nv_fp8_e4m3(sign_(F8E4M3_TO_FLOAT(x)))) +UNARY_OP(__nv_fp8_e4m3, usigmoid_fp8_e4m3, __nv_fp8_e4m3(sigmoid_fwd(F8E4M3_TO_FLOAT(x)))) +#endif + +#if __CUDA_ARCH__ >= 530 +UNARY_OP(__half, ucopy_f16, x) +UNARY_OP(__half, uneg_f16, -x) +UNARY_OP(__half, urecip_f16, recipg(x)) +UNARY_OP(__half, uexp_f16, expg(x)) +UNARY_OP(__half, ulog_f16, logg(x)) +UNARY_OP(__half, usin_f16, sing(x)) +UNARY_OP(__half, ucos_f16, cosg(x)) +UNARY_OP(__half, utanh_f16, tanhg(x)) +UNARY_OP(__half, uerf_f16, erfg(x)) +UNARY_OP(__half, uceil_f16, ceilg(x)) +UNARY_OP(__half, ufloor_f16, floorg(x)) +UNARY_OP(__half, uround_f16, roundg(x)) +UNARY_OP(__half, unormcdf_f16, normcdfg(x)) +UNARY_OP(__half, uabs_f16, absg(x)) +UNARY_OP(__half, usqr_f16, x*x) +UNARY_OP(__half, usqrt_f16, sqrtg(x)) +UNARY_OP(__half, ugelu_f16, gelu_fwd(x)) +UNARY_OP(__half, ugelu_erf_f16, gelu_erf_fwd(x)) +UNARY_OP(__half, urelu_f16, relu_fwd(x)) +UNARY_OP1(__half, uelu_f16, elu_fwd(x, param)) +UNARY_OP(__half, usilu_f16, silu_fwd(x)) +UNARY_OP1(__half, upowf_f16, powg(x, param)) +UNARY_OP(__half, usign_f16, sign_(x)) +UNARY_OP(__half, usigmoid_f16, sigmoid_fwd(x)) +#endif + +UNARY_OP(uint8_t, ucopy_u8, x) +UNARY_OP(uint32_t, ucopy_u32, x) +UNARY_OP(int64_t, ucopy_i64, x) +UNARY_OP(float, ucopy_f32, x) +UNARY_OP(double, ucopy_f64, x) +UNARY_OP(float, uneg_f32, -x) +UNARY_OP(double, uneg_f64, -x) +UNARY_OP(float, urecip_f32, recipg(x)) +UNARY_OP(double, urecip_f64, recipg(x)) +UNARY_OP(float, uexp_f32, expg(x)) +UNARY_OP(double, uexp_f64, expg(x)) +UNARY_OP(float, ulog_f32, logg(x)) +UNARY_OP(double, ulog_f64, logg(x)) +UNARY_OP(float, usin_f32, sing(x)) +UNARY_OP(double, usin_f64, sing(x)) +UNARY_OP(float, ucos_f32, cosg(x)) +UNARY_OP(double, ucos_f64, cosg(x)) +UNARY_OP(float, utanh_f32, tanhg(x)) +UNARY_OP(double, utanh_f64, tanhg(x)) +UNARY_OP(float, uerf_f32, erfg(x)) +UNARY_OP(double, uerf_f64, erfg(x)) +UNARY_OP(float, uceil_f32, ceilg(x)) +UNARY_OP(double, uceil_f64, ceilg(x)) +UNARY_OP(float, ufloor_f32, floorg(x)) +UNARY_OP(double, ufloor_f64, floorg(x)) +UNARY_OP(float, uround_f32, roundg(x)) +UNARY_OP(double, uround_f64, roundg(x)) +UNARY_OP(float, unormcdf_f32, normcdfg(x)) +UNARY_OP(double, unormcdf_f64, normcdfg(x)) +UNARY_OP(float, uabs_f32, absg(x)) +UNARY_OP(double, uabs_f64, absg(x)) +UNARY_OP(float, usqr_f32, x*x) +UNARY_OP(double, usqr_f64, x*x) +UNARY_OP(float, usqrt_f32, sqrtg(x)) +UNARY_OP(double, usqrt_f64, sqrtg(x)) +UNARY_OP(float, ugelu_f32, gelu_fwd(x)) +UNARY_OP(double, ugelu_f64, gelu_fwd(x)) +UNARY_OP(float, ugelu_erf_f32, gelu_erf_fwd(x)) +UNARY_OP(double, ugelu_erf_f64, gelu_erf_fwd(x)) +UNARY_OP(float, urelu_f32, relu_fwd(x)) +UNARY_OP(double, urelu_f64, relu_fwd(x)) +UNARY_OP1(float, uelu_f32, elu_fwd(x, param)) +UNARY_OP1(double, uelu_f64, elu_fwd(x, param)) +UNARY_OP(float, usilu_f32, silu_fwd(x)) +UNARY_OP(double, usilu_f64, silu_fwd(x)) +UNARY_OP1(float, upowf_f32, powg(x, param)) +UNARY_OP1(double, upowf_f64, powg(x, param)) +UNARY_OP(float, usign_f32, sign_(x)) +UNARY_OP(double, usign_f64, sign_(x)) +UNARY_OP(float, usigmoid_f32, sigmoid_fwd(x)) +UNARY_OP(double, usigmoid_f64, sigmoid_fwd(x)) diff --git a/torchnative/rust/vendor/cuda12-candle-kernels-0.11.0.patch b/torchnative/rust/vendor/cuda12-candle-kernels-0.11.0.patch new file mode 100644 index 00000000..87becdae --- /dev/null +++ b/torchnative/rust/vendor/cuda12-candle-kernels-0.11.0.patch @@ -0,0 +1,17 @@ +diff --git a/src/compatibility.cuh b/src/compatibility.cuh +index 8cb8f14..71a30dc 100644 +--- a/src/compatibility.cuh ++++ b/src/compatibility.cuh +@@ -7,7 +7,11 @@ + + // FIXME: the minimum compute capabilities are just guesses since the table is not specific enough + +-#if __CUDA_ARCH__ < 800 ++// torchnative (issue #86): CUDA 12.2 made cuda_fp16.hpp define __hmax_nan and ++// __hmin_nan for every architecture, so defining them again here is a ++// redefinition error on any toolkit from 12.2 on (measured on 12.6 with ++// CUDA_COMPUTE_CAP=75). Define them only for an older toolkit on an older arch. ++#if __CUDA_ARCH__ < 800 && (__CUDACC_VER_MAJOR__ < 12 || (__CUDACC_VER_MAJOR__ == 12 && __CUDACC_VER_MINOR__ < 2)) + __device__ __forceinline__ __half __hmax_nan(__half a, __half b) { + return __hisnan(a) ? a : (__hisnan(b) ? b : __hmax(a, b)); + } From 120101fb4b1aa5a1fc7654dd58a581b62b5fe61f Mon Sep 17 00:00:00 2001 From: BREW <81044362+b-re-w@users.noreply.github.com> Date: Fri, 16 Oct 2026 03:54:54 +0900 Subject: [PATCH 2/2] candle-core fork: cuda_backend handles DType::I8 (#86) The int8 fork added DType::I8 without touching cuda_backend, so the cuda feature failed with E0004 at ten matches once the kernels compiled. CudaStorageSlice has no I8 variant, so nine sites refuse I8 by name (CudaError::UnsupportedDtype with the op); Scalar::I8 passes through as a kernel argument. Patch regenerated; vendor_candle.sh --check passes. --- .../candle-core/src/cuda_backend/device.rs | 25 +++- .../candle-core/src/cuda_backend/mod.rs | 15 +++ .../rust/vendor/int8-candle-0.11.0-cpu.patch | 120 +++++++++++++++++- 3 files changed, 156 insertions(+), 4 deletions(-) diff --git a/torchnative/rust/vendor/candle-core/src/cuda_backend/device.rs b/torchnative/rust/vendor/candle-core/src/cuda_backend/device.rs index d14a2f96..3cbbc69e 100644 --- a/torchnative/rust/vendor/candle-core/src/cuda_backend/device.rs +++ b/torchnative/rust/vendor/candle-core/src/cuda_backend/device.rs @@ -487,6 +487,13 @@ impl BackendDevice for CudaDevice { let data = self.alloc_zeros::(elem_count)?; CudaStorageSlice::F8E4M3(data) } + DType::I8 => { + return Err(CudaError::UnsupportedDtype { + dtype: DType::I8, + op: "zeros", + } + .into()) + } DType::F6E2M3 | DType::F6E3M2 | DType::F4 | DType::F8E8M0 => { return Err( CudaError::InternalError("Dummy types not supported in CUDA backend").into(), @@ -506,6 +513,7 @@ impl BackendDevice for CudaDevice { // TODO: Add support for F16 and BF16 though this is likely to require some upstream // cudarc changes. DType::U8 + | DType::I8 | DType::U32 | DType::I16 | DType::I32 @@ -561,6 +569,7 @@ impl BackendDevice for CudaDevice { }; let slice = match dtype { DType::U8 + | DType::I8 | DType::U32 | DType::I16 | DType::I32 @@ -641,6 +650,13 @@ impl BackendDevice for CudaDevice { let data = self.alloc::(elem_count)?; CudaStorageSlice::F8E4M3(data) } + DType::I8 => { + return Err(CudaError::UnsupportedDtype { + dtype: DType::I8, + op: "alloc_uninit", + } + .into()) + } DType::F6E2M3 | DType::F6E3M2 | DType::F4 | DType::F8E8M0 => { return Err( CudaError::InternalError("Dummy types not supported in CUDA backend").into(), @@ -695,7 +711,8 @@ impl BackendDevice for CudaDevice { let data = self.clone_htod(storage)?; CudaStorageSlice::F8E4M3(data) } - CpuStorageRef::F4(_) + CpuStorageRef::I8(_) + | CpuStorageRef::F4(_) | CpuStorageRef::F6E2M3(_) | CpuStorageRef::F6E3M2(_) | CpuStorageRef::F8E8M0(_) => { @@ -754,7 +771,8 @@ impl BackendDevice for CudaDevice { let data = self.clone_htod(storage)?; CudaStorageSlice::F8E4M3(data) } - CpuStorage::F4(_) + CpuStorage::I8(_) + | CpuStorage::F4(_) | CpuStorage::F6E2M3(_) | CpuStorage::F6E3M2(_) | CpuStorage::F8E8M0(_) => { @@ -813,7 +831,8 @@ impl BackendDevice for CudaDevice { let data = self.clone_htod(&storage)?; CudaStorageSlice::F8E4M3(data) } - CpuStorage::F4(_) + CpuStorage::I8(_) + | CpuStorage::F4(_) | CpuStorage::F6E2M3(_) | CpuStorage::F6E3M2(_) | CpuStorage::F8E8M0(_) => { diff --git a/torchnative/rust/vendor/candle-core/src/cuda_backend/mod.rs b/torchnative/rust/vendor/candle-core/src/cuda_backend/mod.rs index 6986c6e2..c386165c 100644 --- a/torchnative/rust/vendor/candle-core/src/cuda_backend/mod.rs +++ b/torchnative/rust/vendor/candle-core/src/cuda_backend/mod.rs @@ -47,6 +47,7 @@ impl crate::scalar::Scalar { use crate::scalar::Scalar; match self { Scalar::U8(v) => builder.arg(v), + Scalar::I8(v) => builder.arg(v), Scalar::U32(v) => builder.arg(v), Scalar::I16(v) => builder.arg(v), Scalar::I32(v) => builder.arg(v), @@ -1367,6 +1368,13 @@ impl CudaStorage { let result = dst_stream.clone_dtod(cuda_slice).w()?; CudaStorageSlice::F8E8M0(result) } + DType::I8 => { + return Err(CudaError::UnsupportedDtype { + dtype: DType::I8, + op: "transfer_to_device", + } + .into()) + } }; Ok(Self { @@ -1669,6 +1677,13 @@ impl BackendStorage for CudaStorage { DType::I16 | DType::I32 => { return Err(CudaError::InternalError("i16,i32 dtypes are not supported").into()) } + DType::I8 => { + return Err(CudaError::UnsupportedDtype { + dtype: DType::I8, + op: "to_dtype", + } + .into()) + } DType::F6E2M3 | DType::F6E3M2 | DType::F4 | DType::F8E8M0 => { return Err( CudaError::InternalError("Dummy types not supported in CUDA backend").into(), diff --git a/torchnative/rust/vendor/int8-candle-0.11.0-cpu.patch b/torchnative/rust/vendor/int8-candle-0.11.0-cpu.patch index 809185ec..b7496328 100644 --- a/torchnative/rust/vendor/int8-candle-0.11.0-cpu.patch +++ b/torchnative/rust/vendor/int8-candle-0.11.0-cpu.patch @@ -618,6 +618,124 @@ index e909a6c..eac5657 100644 (C::I32(v1), C::I32(v2)) => Ok(C::U8(self.f(v1, l1, v2, l2)?)), (C::I64(v1), C::I64(v2)) => Ok(C::U8(self.f(v1, l1, v2, l2)?)), (C::BF16(v1), C::BF16(v2)) => Ok(C::U8(self.f(v1, l1, v2, l2)?)), +diff --git a/src/cuda_backend/device.rs b/src/cuda_backend/device.rs +index d14a2f9..3cbbc69 100644 +--- a/src/cuda_backend/device.rs ++++ b/src/cuda_backend/device.rs +@@ -487,6 +487,13 @@ impl BackendDevice for CudaDevice { + let data = self.alloc_zeros::(elem_count)?; + CudaStorageSlice::F8E4M3(data) + } ++ DType::I8 => { ++ return Err(CudaError::UnsupportedDtype { ++ dtype: DType::I8, ++ op: "zeros", ++ } ++ .into()) ++ } + DType::F6E2M3 | DType::F6E3M2 | DType::F4 | DType::F8E8M0 => { + return Err( + CudaError::InternalError("Dummy types not supported in CUDA backend").into(), +@@ -506,6 +513,7 @@ impl BackendDevice for CudaDevice { + // TODO: Add support for F16 and BF16 though this is likely to require some upstream + // cudarc changes. + DType::U8 ++ | DType::I8 + | DType::U32 + | DType::I16 + | DType::I32 +@@ -561,6 +569,7 @@ impl BackendDevice for CudaDevice { + }; + let slice = match dtype { + DType::U8 ++ | DType::I8 + | DType::U32 + | DType::I16 + | DType::I32 +@@ -641,6 +650,13 @@ impl BackendDevice for CudaDevice { + let data = self.alloc::(elem_count)?; + CudaStorageSlice::F8E4M3(data) + } ++ DType::I8 => { ++ return Err(CudaError::UnsupportedDtype { ++ dtype: DType::I8, ++ op: "alloc_uninit", ++ } ++ .into()) ++ } + DType::F6E2M3 | DType::F6E3M2 | DType::F4 | DType::F8E8M0 => { + return Err( + CudaError::InternalError("Dummy types not supported in CUDA backend").into(), +@@ -695,7 +711,8 @@ impl BackendDevice for CudaDevice { + let data = self.clone_htod(storage)?; + CudaStorageSlice::F8E4M3(data) + } +- CpuStorageRef::F4(_) ++ CpuStorageRef::I8(_) ++ | CpuStorageRef::F4(_) + | CpuStorageRef::F6E2M3(_) + | CpuStorageRef::F6E3M2(_) + | CpuStorageRef::F8E8M0(_) => { +@@ -754,7 +771,8 @@ impl BackendDevice for CudaDevice { + let data = self.clone_htod(storage)?; + CudaStorageSlice::F8E4M3(data) + } +- CpuStorage::F4(_) ++ CpuStorage::I8(_) ++ | CpuStorage::F4(_) + | CpuStorage::F6E2M3(_) + | CpuStorage::F6E3M2(_) + | CpuStorage::F8E8M0(_) => { +@@ -813,7 +831,8 @@ impl BackendDevice for CudaDevice { + let data = self.clone_htod(&storage)?; + CudaStorageSlice::F8E4M3(data) + } +- CpuStorage::F4(_) ++ CpuStorage::I8(_) ++ | CpuStorage::F4(_) + | CpuStorage::F6E2M3(_) + | CpuStorage::F6E3M2(_) + | CpuStorage::F8E8M0(_) => { +diff --git a/src/cuda_backend/mod.rs b/src/cuda_backend/mod.rs +index 6986c6e..c386165 100644 +--- a/src/cuda_backend/mod.rs ++++ b/src/cuda_backend/mod.rs +@@ -47,6 +47,7 @@ impl crate::scalar::Scalar { + use crate::scalar::Scalar; + match self { + Scalar::U8(v) => builder.arg(v), ++ Scalar::I8(v) => builder.arg(v), + Scalar::U32(v) => builder.arg(v), + Scalar::I16(v) => builder.arg(v), + Scalar::I32(v) => builder.arg(v), +@@ -1367,6 +1368,13 @@ impl CudaStorage { + let result = dst_stream.clone_dtod(cuda_slice).w()?; + CudaStorageSlice::F8E8M0(result) + } ++ DType::I8 => { ++ return Err(CudaError::UnsupportedDtype { ++ dtype: DType::I8, ++ op: "transfer_to_device", ++ } ++ .into()) ++ } + }; + + Ok(Self { +@@ -1669,6 +1677,13 @@ impl BackendStorage for CudaStorage { + DType::I16 | DType::I32 => { + return Err(CudaError::InternalError("i16,i32 dtypes are not supported").into()) + } ++ DType::I8 => { ++ return Err(CudaError::UnsupportedDtype { ++ dtype: DType::I8, ++ op: "to_dtype", ++ } ++ .into()) ++ } + DType::F6E2M3 | DType::F6E3M2 | DType::F4 | DType::F8E8M0 => { + return Err( + CudaError::InternalError("Dummy types not supported in CUDA backend").into(), diff --git a/src/display.rs b/src/display.rs index a9b5394..fd12b46 100644 --- a/src/display.rs @@ -811,7 +929,7 @@ index d628cfd..ee0298c 100644 pub use error::{Context, Error, Result}; pub use indexer::{IndexOp, TensorIndexer}; diff --git a/src/metal_backend/device.rs b/src/metal_backend/device.rs -index 9cae473..8a00bce 100644 +index 9cae473..a238e45 100644 --- a/src/metal_backend/device.rs +++ b/src/metal_backend/device.rs @@ -128,6 +128,36 @@ impl MetalDevice {