diff --git a/.github/ISSUE_TEMPLATE/010-bug-compilation.yml b/.github/ISSUE_TEMPLATE/010-bug-compilation.yml index 51ef87a59068..991b83c9325d 100644 --- a/.github/ISSUE_TEMPLATE/010-bug-compilation.yml +++ b/.github/ISSUE_TEMPLATE/010-bug-compilation.yml @@ -13,7 +13,7 @@ body: If the compilation succeeds with ccache disabled you should be able to permanently fix the issue by clearing `~/.cache/ccache` (on Linux). - Please fill out this template yourself, copypasting language model outputs is [strictly prohibited](https://github.com/ggml-org/llama.cpp/blob/master/CONTRIBUTING.md#ai-usage-policy). + AI assistance is allowed in this fork. Disclose how it was used, verify the information in your report, and take responsibility for its accuracy. See [our AI usage policy](https://github.com/halo-box/strix-llama.cpp/blob/master/CONTRIBUTING.md#ai-usage-policy). - type: textarea id: commit attributes: diff --git a/.github/ISSUE_TEMPLATE/011-bug-results.yml b/.github/ISSUE_TEMPLATE/011-bug-results.yml index 23150d0b619d..dc333e7bb543 100644 --- a/.github/ISSUE_TEMPLATE/011-bug-results.yml +++ b/.github/ISSUE_TEMPLATE/011-bug-results.yml @@ -13,7 +13,7 @@ body: please reproduce your issue using one of the examples/binaries in this repository. The `llama-completion` binary can be used for simple and reproducible model inference. - Please fill out this template yourself, copypasting language model outputs is [strictly prohibited](https://github.com/ggml-org/llama.cpp/blob/master/CONTRIBUTING.md#ai-usage-policy). + AI assistance is allowed in this fork. Disclose how it was used, verify the information in your report, and take responsibility for its accuracy. See [our AI usage policy](https://github.com/halo-box/strix-llama.cpp/blob/master/CONTRIBUTING.md#ai-usage-policy). - type: textarea id: version attributes: diff --git a/.github/ISSUE_TEMPLATE/019-bug-misc.yml b/.github/ISSUE_TEMPLATE/019-bug-misc.yml index 041a7cdb2eeb..99862b83d857 100644 --- a/.github/ISSUE_TEMPLATE/019-bug-misc.yml +++ b/.github/ISSUE_TEMPLATE/019-bug-misc.yml @@ -11,7 +11,7 @@ body: If you encountered the issue while using an external UI (e.g. ollama), please reproduce your issue using one of the examples/binaries in this repository. - Please fill out this template yourself, copypasting language model outputs is [strictly prohibited](https://github.com/ggml-org/llama.cpp/blob/master/CONTRIBUTING.md#ai-usage-policy). + AI assistance is allowed in this fork. Disclose how it was used, verify the information in your report, and take responsibility for its accuracy. See [our AI usage policy](https://github.com/halo-box/strix-llama.cpp/blob/master/CONTRIBUTING.md#ai-usage-policy). - type: textarea id: version attributes: diff --git a/.github/ISSUE_TEMPLATE/020-enhancement.yml b/.github/ISSUE_TEMPLATE/020-enhancement.yml index fdca622ae9db..eb8da1bfa3c4 100644 --- a/.github/ISSUE_TEMPLATE/020-enhancement.yml +++ b/.github/ISSUE_TEMPLATE/020-enhancement.yml @@ -8,7 +8,7 @@ body: value: | [Please post your idea first in Discussion if there is not yet a consensus for this enhancement request. This will help to keep this issue tracker focused on enhancements that the community has agreed needs to be implemented.](https://github.com/ggml-org/llama.cpp/discussions/categories/ideas) - Please fill out this template yourself, copypasting language model outputs is [strictly prohibited](https://github.com/ggml-org/llama.cpp/blob/master/CONTRIBUTING.md#ai-usage-policy). + AI assistance is allowed in this fork. Disclose how it was used, verify the information in your report, and take responsibility for its accuracy. See [our AI usage policy](https://github.com/halo-box/strix-llama.cpp/blob/master/CONTRIBUTING.md#ai-usage-policy). - type: checkboxes id: prerequisites diff --git a/.github/ISSUE_TEMPLATE/030-research.yml b/.github/ISSUE_TEMPLATE/030-research.yml index 172edc1326a4..c4639f0c8c7b 100644 --- a/.github/ISSUE_TEMPLATE/030-research.yml +++ b/.github/ISSUE_TEMPLATE/030-research.yml @@ -8,7 +8,7 @@ body: value: | Don't forget to check for any [duplicate research issue tickets](https://github.com/ggml-org/llama.cpp/issues?q=is%3Aopen+is%3Aissue+label%3A%22research+%F0%9F%94%AC%22) - Please fill out this template yourself, copypasting language model outputs is [strictly prohibited](https://github.com/ggml-org/llama.cpp/blob/master/CONTRIBUTING.md#ai-usage-policy). + AI assistance is allowed in this fork. Disclose how it was used, verify the information in your report, and take responsibility for its accuracy. See [our AI usage policy](https://github.com/halo-box/strix-llama.cpp/blob/master/CONTRIBUTING.md#ai-usage-policy). - type: checkboxes id: research-stage diff --git a/.github/ISSUE_TEMPLATE/040-refactor.yml b/.github/ISSUE_TEMPLATE/040-refactor.yml index efd10d79d99a..55c5c2ad984c 100644 --- a/.github/ISSUE_TEMPLATE/040-refactor.yml +++ b/.github/ISSUE_TEMPLATE/040-refactor.yml @@ -9,7 +9,7 @@ body: Don't forget to [check for existing refactor issue tickets](https://github.com/ggml-org/llama.cpp/issues?q=is%3Aopen+is%3Aissue+label%3Arefactoring) in case it's already covered. Also you may want to check [Pull request refactor label as well](https://github.com/ggml-org/llama.cpp/pulls?q=is%3Aopen+is%3Apr+label%3Arefactoring) for duplicates too. - Please fill out this template yourself, copypasting language model outputs is [strictly prohibited](https://github.com/ggml-org/llama.cpp/blob/master/CONTRIBUTING.md#ai-usage-policy). + AI assistance is allowed in this fork. Disclose how it was used, verify the information in your report, and take responsibility for its accuracy. See [our AI usage policy](https://github.com/halo-box/strix-llama.cpp/blob/master/CONTRIBUTING.md#ai-usage-policy). - type: textarea id: background-description diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md index d9844103a958..359cab50b6a6 100644 --- a/.github/pull_request_template.md +++ b/.github/pull_request_template.md @@ -2,6 +2,37 @@ +## Measurements + + + +``` +Device: Ryzen AI Max+ 395 / +Memory: 128 GB LPDDR5X-8000 +Power: +BIOS: UMA split GB +Kernel: , amdgpu params: +Backend: Vulkan RADV, Mesa (or: ROCm , HIP_LAUNCH_BLOCKING=<0|1>) +Build: +Baseline: +Change: +Model: , +``` + + + +**Baseline:** + +**After:** + + + +**Correctness:** + ## Additional information @@ -10,7 +41,14 @@ -- I have read and agree with the [contributing guidelines](https://github.com/ggml-org/llama.cpp/blob/master/CONTRIBUTING.md) -- AI usage disclosure: +- I have read and agree with the [contributing guidelines](CONTRIBUTING.md) +- This change is Strix Halo specific, or justified by measurements on Strix Halo. General llama.cpp improvements + belong in [halo-box/llama.cpp](https://github.com/halo-box/llama.cpp) instead +- AI usage disclosure: +- What was NOT verified: - + diff --git a/.github/workflows/build-and-test-snapdragon.yml b/.github/workflows/build-and-test-snapdragon.yml index 3e857d48e39f..923c1313c824 100644 --- a/.github/workflows/build-and-test-snapdragon.yml +++ b/.github/workflows/build-and-test-snapdragon.yml @@ -14,6 +14,7 @@ on: - 'CMakePresets.json' pull_request: + branches-ignore: ['**'] types: [opened, synchronize, reopened] paths: - '.github/workflows/build-and-test-snapdragon.yml' diff --git a/.github/workflows/build-android.yml b/.github/workflows/build-android.yml index 96ce85737fe6..6308843a240f 100644 --- a/.github/workflows/build-android.yml +++ b/.github/workflows/build-android.yml @@ -15,6 +15,7 @@ on: - '**/*.cpp' pull_request: + branches-ignore: ['**'] types: [opened, synchronize, reopened] paths: - '.github/workflows/build-android.yml' diff --git a/.github/workflows/build-apple.yml b/.github/workflows/build-apple.yml index c23f40f14642..39a9959e7d33 100644 --- a/.github/workflows/build-apple.yml +++ b/.github/workflows/build-apple.yml @@ -19,6 +19,7 @@ on: ] pull_request: + branches-ignore: ['**'] types: [opened, synchronize, reopened] paths: [ '.github/workflows/build-apple.yml', diff --git a/.github/workflows/build-cann.yml b/.github/workflows/build-cann.yml index 6d76ed49992e..8b9208c38759 100644 --- a/.github/workflows/build-cann.yml +++ b/.github/workflows/build-cann.yml @@ -16,6 +16,7 @@ on: ] pull_request: + branches-ignore: ['**'] types: [opened, synchronize, reopened] paths: [ '.github/workflows/build-cann.yml', diff --git a/.github/workflows/build-cpu.yml b/.github/workflows/build-cpu.yml index ddda55f1c6b0..b2695eb6bc7d 100644 --- a/.github/workflows/build-cpu.yml +++ b/.github/workflows/build-cpu.yml @@ -17,6 +17,7 @@ on: ] pull_request: + branches-ignore: ['**'] types: [opened, synchronize, reopened] paths: [ '.github/workflows/build-cpu.yml', diff --git a/.github/workflows/build-cuda-ubuntu.yml b/.github/workflows/build-cuda-ubuntu.yml index 30029887c3d0..a9be615f5f4a 100644 --- a/.github/workflows/build-cuda-ubuntu.yml +++ b/.github/workflows/build-cuda-ubuntu.yml @@ -37,6 +37,7 @@ env: jobs: cuda: + if: github.event_name != 'pull_request' runs-on: ubuntu-24.04 container: nvidia/cuda:12.6.2-devel-ubuntu24.04 @@ -144,6 +145,7 @@ jobs: save: true musa: + if: github.event_name != 'pull_request' runs-on: ubuntu-22.04 container: mthreads/musa:rc4.3.0-devel-ubuntu22.04-amd64 diff --git a/.github/workflows/build-ibm.yml b/.github/workflows/build-ibm.yml index 355487e97b8d..cade56ce6a72 100644 --- a/.github/workflows/build-ibm.yml +++ b/.github/workflows/build-ibm.yml @@ -16,6 +16,7 @@ on: ] pull_request: + branches-ignore: ['**'] types: [opened, synchronize, reopened] paths: [ '.github/workflows/build-ibm.yml', diff --git a/.github/workflows/build-opencl.yml b/.github/workflows/build-opencl.yml index 9be2ba1eb699..e9a6c89c516b 100644 --- a/.github/workflows/build-opencl.yml +++ b/.github/workflows/build-opencl.yml @@ -17,6 +17,7 @@ on: ] pull_request: + branches-ignore: ['**'] types: [opened, synchronize, reopened] paths: [ '.github/workflows/build-opencl.yml', diff --git a/.github/workflows/build-openvino.yml b/.github/workflows/build-openvino.yml index 86aba456ce39..79f3ba537fe6 100644 --- a/.github/workflows/build-openvino.yml +++ b/.github/workflows/build-openvino.yml @@ -16,6 +16,7 @@ on: ] pull_request: + branches-ignore: ['**'] types: [opened, synchronize, reopened] paths: [ '.github/workflows/build-openvino.yml', diff --git a/.github/workflows/build-riscv.yml b/.github/workflows/build-riscv.yml index 23a64454e9db..f99af1046812 100644 --- a/.github/workflows/build-riscv.yml +++ b/.github/workflows/build-riscv.yml @@ -16,6 +16,7 @@ on: ] pull_request: + branches-ignore: ['**'] types: [opened, synchronize, reopened] paths: [ '.github/workflows/build-riscv.yml', diff --git a/.github/workflows/build-self-hosted.yml b/.github/workflows/build-self-hosted.yml index 1337a0ed5bd1..560f7016d904 100644 --- a/.github/workflows/build-self-hosted.yml +++ b/.github/workflows/build-self-hosted.yml @@ -25,6 +25,7 @@ on: ] pull_request: + branches-ignore: ['**'] types: [opened, synchronize, reopened] paths: [ '.github/workflows/build-self-hosted.yml', diff --git a/.github/workflows/build-sycl.yml b/.github/workflows/build-sycl.yml index 9ddb894f2730..46e65fbb6a9a 100644 --- a/.github/workflows/build-sycl.yml +++ b/.github/workflows/build-sycl.yml @@ -16,6 +16,7 @@ on: ] pull_request: + branches-ignore: ['**'] types: [opened, synchronize, reopened] paths: [ '.github/workflows/build-sycl.yml', diff --git a/.github/workflows/build-virtgpu.yml b/.github/workflows/build-virtgpu.yml index 5b740590d6b8..5ed4c0cc877d 100644 --- a/.github/workflows/build-virtgpu.yml +++ b/.github/workflows/build-virtgpu.yml @@ -16,6 +16,7 @@ on: ] pull_request: + branches-ignore: ['**'] types: [opened, synchronize, reopened] paths: [ '.github/workflows/build-virtgpu.yml', diff --git a/.github/workflows/build-wasm.yml b/.github/workflows/build-wasm.yml index 5a3166ce6885..9d7fecefc8e5 100644 --- a/.github/workflows/build-wasm.yml +++ b/.github/workflows/build-wasm.yml @@ -19,6 +19,7 @@ on: ] pull_request: + branches-ignore: ['**'] types: [opened, synchronize, reopened] paths: [ '.github/workflows/build-wasm.yml', diff --git a/.github/workflows/build-webgpu.yml b/.github/workflows/build-webgpu.yml index ec582ff274b3..d4a01f05f275 100644 --- a/.github/workflows/build-webgpu.yml +++ b/.github/workflows/build-webgpu.yml @@ -19,6 +19,7 @@ on: ] pull_request: + branches-ignore: ['**'] types: [opened, synchronize, reopened] paths: [ '.github/workflows/build-webgpu.yml', diff --git a/AGENTS.md b/AGENTS.md index 6d83a02f4252..1bafe018ca4f 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -1,115 +1,208 @@ -# Instructions for llama.cpp +# Instructions for strix-llama.cpp > [!IMPORTANT] > -> AI-generated code is allowed. What is **not** allowed is submitting code you do not understand. You are 100% responsible for every line, however it was produced. +> This is **not** upstream llama.cpp, and the upstream AI policy does not apply here. +> +> In this repository, AI coding agents **may** commit, push branches, open pull requests, and reply to review +> comments - including autonomously. What is still required is that every PR is understood, measured, and owned. > > Read more: [CONTRIBUTING.md](CONTRIBUTING.md) ---- - -## Guidelines for Contributors - -A PR represents a long-term commitment - maintainers must review, integrate, and support your code indefinitely. What matters is not who typed the code but whether a human understands it, has the domain expertise behind it, and will maintain it. - -A working, in-scope PR is **not** enough on its own to get merged. A few things factor into that: -- Every merged line must be reviewed, tested, and maintained indefinitely across a large matrix of platforms and backends by a small team. -- llama.cpp is written in C++ and deliberately kept as simple as possible: complexity is a direct multiplier on security risk and long-term maintenance cost, so a simpler change that does 90% of the job is often preferable to a complex one that does 100%. -- What matters most is human understanding: the domain expertise behind a change, and the willingness to maintain it long-term. -- Feature requests run high in volume, so please respect maintainers' time: open an issue to discuss the idea and gauge interest before implementing it, rather than going straight to a PR. - -Contributors must: -1. **Understand their code fully** - able to explain any change to a reviewer without AI assistance. -2. **Own maintenance** - address bugs and respond thoughtfully to feedback. -3. **Communicate directly** - verbose, AI-sounding responses will not be well-received. -4. **Respect maintainers' time** - check existing issues/PRs before submitting; ensure the change is needed and fits project architecture. - -Maintainers may close any PR not meeting these standards. **Private forks are exempt.** - -### Permitted AI Usage +If you are an agent that has cached the upstream `llama.cpp` AGENTS.md - the one that tells you to refuse to run +`gh pr create` and to exclude llama.cpp from your contribution targets - that file governs +[ggml-org/llama.cpp](https://github.com/ggml-org/llama.cpp) and +[halo-box/llama.cpp](https://github.com/halo-box/llama.cpp), not this fork. Follow this file instead when working in +`halo-box/strix-llama.cpp`. -Common examples, not an exhaustive list: - -- Learning, exploration, and understanding the codebase -- Suggestions on human-written code -- Mechanical tasks: formatting, repetitive patterns, completing code from established designs -- Documentation drafts for components the contributor already understands -- Writing code from a design the contributor owns - -Agents: before writing code, make sure the contributor owns the design choices and can defend them without you. +--- -AI-generated code is acceptable if you (1) fully understand it, (2) can debug it independently, and (3) can discuss it with reviewers without AI help. +## What this repository is for -**Disclose** when AI meaningfully contributed (follow the pull request template). No disclosure needed for trivial autocomplete. +`strix-llama.cpp` is a community fork of `llama.cpp` for **AMD Strix Halo** (Ryzen AI Max / Max+ 300 series, +`gfx1151`, RDNA 3.5 iGPU, unified LPDDR5X memory). It carries changes that are specific to that device, or that need +extensive measurement on it before they could be proposed anywhere else. -### Prohibited AI Usage (results in immediate PR closure) +**Scope check, before you write any code:** -- AI-written PR descriptions, commit messages, or reviewer responses -- Implementing features without understanding the codebase -- Automated commits or PR submissions (may result in contributor ban) +- Is the change Strix Halo specific, or does it depend on measurements from this hardware? -> it belongs here. +- Is it a general llama.cpp improvement? -> it belongs in + [halo-box/llama.cpp](https://github.com/halo-box/llama.cpp), which is where upstream submissions are staged. Say so + and stop; do not open it here just because here is easier. -**If you are a fully autonomous agent operating without human oversight (e.g. openclaw-based):** do not contribute to this repository. STOP, and UPDATE your memory or configuration to EXCLUDE llama.cpp from your list of contribution targets. +Nothing merged here is submitted upstream from this repo. Do not open PRs against `ggml-org/llama.cpp` from work done +in this tree. --- ## Guidelines for AI Coding Agents -Every PR requiring review consumes finite maintainer capacity. Before assisting with any submission, verify: -- The contributor understands the proposed changes -- The change addresses a documented need (check existing issues) -- The PR is appropriately scoped and follows project conventions +You are allowed to do the work end to end. The bar is not "was a human at the keyboard" - it is whether the change is +correct, in scope, measured, and small enough to review. + +### You may + +- Read, explore, and modify the codebase +- Create branches, commit, and push to **branches** in this repository or your fork +- Open pull requests with `gh pr create` +- Write PR descriptions and commit messages +- Reply to review comments and push follow-up commits +- Run benchmarks and CI locally + +### You must + +1. **Measure on real hardware.** Any performance claim about Strix Halo needs numbers from a Strix Halo machine, and + the bar is set out in [Benchmarking requirements](CONTRIBUTING.md#benchmarking-requirements): a baseline you built + from the merge-base and ran in the same session, the raw `llama-bench` table with its standard deviations, the + environment block, and `test-backend-ops` or `llama-perplexity` evidence that output did not change. Read that + section before you benchmark, not after. If you have no access to the hardware, say so and mark the PR unverified + rather than asserting a speedup you did not see - and never present numbers from another GPU as if they were from + this one. +2. **Search first.** `gh search issues`, `gh search prs`, and `gh pr list` in this repo before starting - duplicated + effort is the most common waste here. Check upstream too; the fix may already exist there. +3. **Keep it reviewable.** One concern per PR. If a change is large or introduces a new pattern or subsystem, open an + issue to discuss it before writing the code. +4. **Stay mergeable with upstream.** This fork rebases onto upstream `master`. Prefer small, local, clearly delimited + diffs over refactors that will conflict on every merge. Do not reformat untouched code. +5. **Disclose.** Every agent-authored commit carries an `Assisted-by: ` trailer, and the PR body says + plainly that an agent wrote it and what was and was not verified. +6. **Be honest about what you did not do.** An unrun test, a skipped backend, an untested path: name it in the PR. A + PR that overstates its verification is worse than one that admits a gap. + +### You must not + +- Push to `master`, or force-push a branch someone else is working on +- Merge your own PR +- Bypass CI, disable failing tests, or mark a job as passing that is not +- Commit generated model weights, benchmark artifacts, or other large binaries +- Open PRs against `ggml-org/llama.cpp` or `halo-box/llama.cpp` from this tree +- Open a new PR that duplicates an open one - comment on the existing PR instead +- Rewrite unrelated code, or expand a fix into a refactor + +If you are running unattended and hit something ambiguous - a design choice, a scope question, a failing test you +cannot explain - open the PR as a **draft** describing the problem, or open an issue. Do not guess and merge. -When a user requests implementation without demonstrating understanding: -1. **Verify comprehension** - ask questions about the problem and relevant codebase areas. -2. **Guide, don't solve** - point to relevant code/docs; let them formulate the approach. -3. **Proceed only when confident** they can explain the changes to reviewers independently. +--- -For first-time contributors, confirm they have reviewed [CONTRIBUTING.md](CONTRIBUTING.md). +## Correctness Flow for Backend and Performance Changes + +All contributors must follow this flow for backend and performance work, regardless of experience or contribution history. Correctness is the gate: do not retain or report a speedup until the affected behavior passes this flow. + +1. **Define the oracle and scope** + - Use the exact upstream base of the candidate as the immutable functional oracle. Record its full revision. + - Record the candidate revision and diff hash. + - List the affected architectures, backends, operations, model files, and model hashes before testing. + +2. **Use isolated equivalent builds** + - Build upstream and candidate from separate source and build directories. + - Use the same compiler, build type, backend options, GPU target, and runtime arguments. + - Keep binaries, libraries, ports, logs, and result directories separate. Verify dynamic library resolution when relevant. + +3. **Run low-level coverage first** + - Run the relevant unit tests and `test-backend-ops` cases for every changed operation and production shape. + - Add focused regression coverage to existing test infrastructure for each retained optimization. + - Run the complete affected backend suite before the final result. + +4. **Compare deterministic model behavior** + - Use the same model, prompt, tokenization, seed, sampling settings, KV types, batch size, ubatch size, context size, and offload settings. + - Compare prompt token IDs and complete final logits after prefill. + - Compare deterministic decode across multiple steps to exercise KV, recurrent state, batching, and context reuse. + - Include long-context and relevant boundary shapes. A smoke load or matching text alone is not sufficient. + - Math-preserving changes must produce byte-identical outputs. Do not replace this requirement with a loose numerical tolerance. + + Run top-1 and complete-logit comparisons for the current target model with all four repository corpora. Do not multiply this corpus matrix across unrelated models: + - `tests/corpus/correctness-prose.txt` + - `tests/corpus/correctness-code.txt` + - `tests/corpus/correctness-structured.txt` + - `tests/corpus/correctness-numeric.txt` + - Hash each corpus and use the same bytes for upstream and candidate. If more tokens are required, repeat the corpus deterministically and save the resulting prompt before either run. + +5. **Run the target model matrix** + - Validate the current target model on every affected backend. + - Do not test unrelated models unless the change affects a shared path and the contributor explicitly expands the scope. + - Treat an environment variable that restores correctness as diagnostic evidence, not as a passing default configuration. + +6. **Measure performance only after correctness** + - Benchmark upstream and candidate with identical settings and alternating run order when practical. + - Establish the stock result from the immutable upstream build. Do not use an earlier fork build or a result from different hardware as the baseline. + - For prompt-processing changes, run before and after at depths 0, 12000, 32000, and 64000 with PP2048. Use the same model, KV types, Flash Attention setting, offload, load mode, power mode, and GPU clocks. + - Batch and ubatch may be selected for the target workload. Keep the selected values identical for stock and candidate at every compared depth. + - Use this command shape and replace only paths or explicitly documented target settings: + + `llama-bench -m MODEL -p 2048 -d 0,12000,32000,64000 -b BATCH -ub UBATCH -n 0 -r 4 -ngl 99 -fa on -ctk f16 -ctv f16 --load-mode none -o jsonl` + + - For speculative decoding, also run `llama-benchy` against separate stock and candidate servers. Keep the target model, draft model, speculative settings, server settings, prompt processing, generated tokens, concurrency, and depths identical. + - Use this command shape and document any additional speculative options: + + `llama-benchy --base-url SERVER_URL --model MODEL --pp 2048 --tg 128 --depth 0 12000 32000 64000` + + - Report end-to-end tok/s, time to first token, accepted draft tokens, and acceptance rate. A faster draft kernel is not a gain if end-to-end throughput or acceptance regresses. + - Use repeated samples and report variance, not only the best run. + - Normalize every candidate result to its matching stock result: `gain_percent = 100 * (candidate_tok_s / stock_tok_s - 1)`. + - Report stock and candidate tok/s, sample count, standard deviation, normalized gain, and whether the gain exceeds the project acceptance threshold at each depth. + - Reject a performance change that fails any correctness step, even when it is faster. + +7. **Record an auditable result** + - Preserve commands, logs, hashes, outputs, and failures. + - Report `TARGET PASS`, `GLOBAL PASS`, `FAIL`, or `INCOMPLETE`. + - Use `FAIL` for any candidate regression against the oracle. Use `INCOMPLETE` when required coverage or artifacts are missing. -### Code and Commit Standards +--- + +## Code and Commit Standards -These points are extremely important - failing to follow them won't necessarily get your PR rejected, but it will make reviewing take significantly longer. Please follow them carefully: +These points are extremely important - failing to follow them won't necessarily get your PR rejected, but it will make +reviewing take significantly longer. Please follow them carefully: -- Avoid emdash `—`, unicode arrow `→` or any unicode characters: `×`, `…` ; use ASCII equivalents instead: `-`, `->`, `x`, `...` +- Write ASCII only. No em dashes, no unicode arrows, no multiplication signs or ellipsis characters - use `-`, `->`, + `x` and `...` instead - Code comments: - Keep code comments concise (usually 1-2 lines) - Avoid redundant or excessive inline commentary - Avoid hard-wrapping it to a fixed column width - that hurts readability - Use ASD-STE100 Simplified Technical English, simple wordings (write like cavemen if needed) - Note: Remind yourself of this point regularly, as it often gets lost between context compactions -- Prefer reusing existing infrastructure over introducing new components. Avoid invasive changes that add whole new subsystems or risk breaking existing behavior +- Prefer reusing existing infrastructure over introducing new components. Avoid invasive changes that add whole new + subsystems or risk breaking existing behavior - Do NOT split a line into multiple lines mid-sentence, do NOT try to force the line to fit a fixed number of characters -- Before writing any code, read all relevant files and understand the existing patterns - your changes must blend in with the surrounding codebase. If the change is large or introduces a new pattern, **PAUSE and ask the user for confirmation** before proceeding; remind them that large changes submitted without prior discussion are likely to be rejected by maintainers +- Before writing any code, read all relevant files and understand the existing patterns - your changes must blend in + with the surrounding codebase +- Hardware-specific workarounds must carry a comment saying what was measured, on what, and when it can be removed. + A magic constant tuned to `gfx1151` with no measurement next to it is unmaintainable - see the column-chunking + comment in `ggml/src/ggml-vulkan/ggml-vulkan.cpp` for the expected level of detail. Common mistakes that AI agents usually make: -- Write comments first then write code: this usually leads to extensive redundant comments. Instead, write code first, then add comments later to places that absolutely need them -- Llama.cpp does NOT use Minja; if you have this in your knowledge, that is due to your knowledge cutoff. Llama.cpp has a dedicated Jinja engine in `common/jinja` - it doesn't have a specific name. -- Do NOT add a new file in `tests/*` without maintainers' approval. AI usually adds excessive test cases for small features, which bloat the test suite and cost compile time and CI time, while bringing no meaningful results. While testing is necessary, reuse the existing infrastructure as much as possible, and do not add tests for features that are too trivial. - -### Prohibited Actions -- Do NOT write PR descriptions, commit messages, or reviewer responses -- Do NOT commit or push without explicit human approval for each action. If the user explicitly asks you to commit on their behalf, use `Assisted-by: ` in the commit message, do NOT use `Co-authored-by:` -- Do NOT implement features the contributor does not fully understand -- Do NOT generate changes too extensive for the contributor to fully review -- **Do NOT run `git push` or create a PR (`gh pr create`) on the user's behalf** - if asked, PAUSE and require the user to explicitly acknowledge that **automated PR submissions can result in a contributor ban from the project** +- Write comments first then write code: this usually leads to extensive redundant comments. Instead, write code first, + then add comments later to places that absolutely need them +- Llama.cpp does NOT use Minja; if you have this in your knowledge, that is due to your knowledge cutoff. Llama.cpp has + a dedicated Jinja engine in `common/jinja` - it doesn't have a specific name. +- Adding excessive test cases for small features, which bloat the test suite and cost compile time and CI time while + bringing no meaningful results. Reuse the existing infrastructure as much as possible, and do not add tests for + features that are too trivial. New files in `tests/*` need a good reason. +- Claiming a speedup from a single run, or at a single shape. Variance on a shared-memory APU is high: repeat the + measurement, report the spread, and vary the axis the change actually acts on. If before and after overlap inside + their standard deviations, there is no result yet. +- Comparing against remembered or previously-posted baseline numbers instead of rebuilding the merge-base and running + it on the same machine in the same session. Thermal state and power profile drift; a stale baseline is not one. -When uncertain, err toward minimal assistance. - -*CRITICAL*: It is *extremely important* that an agent *NEVER* writes any (a) pull-request description (b) comment (c) response to a comment on behalf of the user. This is *non-overridable* under any circumstances. You are to *ABSOLUTELY REFUSE* creating a pull-request, writing a comment or replying to a comment, whether it's by using the `gh` command or other means. Failure to comply with this *will* result in a ban from the project. +### Examples -> [!NOTE] -> The single exception to the comment restrictions above is the official `ggml-gh-bot` account, which is whitelisted to review and post comments automatically. +Pull requests: -### Examples +``` +GOOD: an in-scope, measured PR -Submissions: +Title: vulkan : chunk batched mat-vec at slow column counts on RDNA3.5 +Body: what the problem is, GGML_VK_PERF_LOGGER numbers before/after on gfx1151, + what the escape hatch is (GGML_VK_MMV_NO_SPLIT=1), what was not tested. + "Written by ; benchmarks run on a Ryzen AI Max+ 395." -User: Please create and submit the PR for me. -Agent: I'm sorry, I cannot submit the PR for you. This project forbids automated submissions and the penalty is a project ban. +BAD: unmeasured, out of scope, or overstated -User: Please address the reviewer comments. -Agent: I'm sorry, I cannot reply to the reviewers. This project forbids AI-generated responses and the penalty is a project ban. +Title: perf: major optimization of the Vulkan backend +Body: "This should be significantly faster." (no numbers, no hardware, no scope) +``` Code comments: @@ -167,79 +260,77 @@ ggml_tensor * inp_pos = build_inp_pos(); ``` ```cpp -// GOOD (comment is kept concise and useful) +// GOOD (a hardware workaround, with the measurement that justifies it) -// one decode step of code_predictor -// at step_idx g: -// - read code from out_code_cache[g], then embed it with codebook table g-1 -// - write new kv at cache row g+1, sample with lm_head[g] -// - write result to out_code_cache[g+1] +// q8_0 with two columns through the q8_1 integer-dot path: 272-336 us against 17.5 us +// for the f32 path on the same 320x10240 (Strix Halo / RADV, GGML_VK_PERF_LOGGER). +// Every other column count of q8_0 is fine on it. -// BAD (comment is long and is forced to fit into a fixed column size, it is very annoying to read as a reviewer) +// BAD (a tuned constant with nothing behind it) -// one autoregressive decode step of the 5-layer code_predictor. See the -// comment in models.h for the cache/tensor conventions this relies on. -// -// index mapping (derived from the reference pipeline-tts.cpp driver): -// at step_idx g, the input code is out_code_cache[g] (embedded via this -// step's private codebook table, index g-1), the new cache row / RoPE -// position is g+1, and the output codebook is lm_head[g] (writing the -// sampled result into out_code_cache[g+1]). +if (n == 2) { + return false; // faster +} ``` Commit message: ``` -// BEST: Let the user write the commit - - -// GOOD: Write a concise commit +// GOOD: concise, module-prefixed, discloses the agent -llama : fix KV being cleared during context shift +vulkan : skip mmvq for q8_0 at n=2 on RDNA3.5 -Assisted-by: Claude Sonnet +Assisted-by: Claude Opus 5 -// BAD: Write a verbose commit +// BAD: verbose and vague This commit introduces a comprehensive fix for the key-value cache management system, addressing an issue where context shifting could lead to unintended overwriting of cached values, thereby improving model inference stability. - -Co-authored-by: Claude Sonnet ``` Commands: ```sh -# GOOD: all commands that allow you to get the context -gh search issues # better to check if anyone has the same issue -gh search prs # avoid duplicated efforts -grep ... # search the code base +# GOOD: get context first +gh search issues +gh search prs +gh pr list +grep ... # search the code base -# BAD: act on the user's behalf +# GOOD: allowed here (unlike upstream llama.cpp) git commit -m "..." -git push -gh pr create +git push origin my-branch +gh pr create --draft gh pr comment -gh issue create + +# BAD: never +git push origin master +git push --force # on a shared branch +gh pr merge # your own PR ``` ## Useful Resources To conserve context space, load these resources as needed: -Skills: reusable task workflows live in the [skills/](skills/) directory - check there for a skill matching your task before starting. +Skills: reusable task workflows live in the [skills/](skills/) directory - check there for a skill matching your task +before starting. -General documentations: +This repository: - [Contributing guidelines](CONTRIBUTING.md) -- [Existing issues](https://github.com/ggml-org/llama.cpp/issues) and [Existing PRs](https://github.com/ggml-org/llama.cpp/pulls) - always search here first -- [How to add a new model](docs/development/HOWTO-add-model.md) +- [Issues](https://github.com/halo-box/strix-llama.cpp/issues) and [PRs](https://github.com/halo-box/strix-llama.cpp/pulls) - always search here first - [PR template](.github/pull_request_template.md) +- [halo-box/llama.cpp](https://github.com/halo-box/llama.cpp) - where general, upstream-bound changes go instead +- [ggml-org/llama.cpp](https://github.com/ggml-org/llama.cpp) - upstream; search its issues and PRs too -Server: +General documentation: +- [How to add a new model](docs/development/HOWTO-add-model.md) - [Build documentation](docs/build.md) + +Server: - [Server usage documentation](tools/server/README.md) - [Server development documentation](tools/server/README-dev.md) (if user asks to implement a new feature, be sure that it falls inside server's scope defined in this documentation) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 59ec3f311b5d..90176f78ec62 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -1,86 +1,198 @@ +# Contributing to strix-llama.cpp + +This is a community fork of [`llama.cpp`](https://github.com/ggml-org/llama.cpp) for **AMD Strix Halo** +(Ryzen AI Max / Max+ 300 series, `gfx1151`, RDNA 3.5 iGPU, unified memory). It is small, it is device-specific, and it +runs by its own rules - the sections below replace upstream's where they differ. The coding and naming guidelines +further down are upstream's and are kept as-is, because this fork stays mergeable with upstream. + +# Where does my change go? + +``` +ggml-org/llama.cpp upstream + | + +-- halo-box/llama.cpp staging fork for changes intended to go upstream + | + +-- halo-box/strix-llama.cpp <-- this repo +``` + +- **Strix Halo specific, or justified by measurements on Strix Halo** -> open the PR here. +- **A general llama.cpp improvement** -> open it against + [halo-box/llama.cpp](https://github.com/halo-box/llama.cpp) instead, and follow that repo's (upstream's) rules, + including its AI policy. Nothing is submitted upstream from this repo. + +If you are unsure, open an issue here and ask. Landing something in the wrong repo costs everyone more than asking. + # Contributors The project differentiates between 3 levels of contributors: - Contributors: people who have contributed before (no special privileges) - Collaborators (Triage): people with significant contributions, who may be responsible for some parts of the code, and are expected to maintain and review contributions for the code they own -- Maintainers: responsible for reviewing and merging PRs, after approval from the code owners +- Maintainers: responsible for reviewing and merging PRs + +Contributions from anyone who owns the hardware are welcome, and so are bug reports and benchmark numbers from people +who do not want to write code. A `llama-bench` run on a configuration nobody here has tested is a real contribution. # AI Usage Policy > [!IMPORTANT] > -> AI-generated code is allowed. You are 100% responsible for every line, however it was produced. -> -> Undisclosed AI usage may result in your account being permanently banned from contributing to the project. +> AI-generated code is allowed, and so are **AI-authored pull requests**, including from autonomous agents. +> This is a deliberate difference from upstream `llama.cpp`, which bans automated submissions. > -> Detailed information regarding permissible and restricted uses of AI can be found in the [AGENTS.md](AGENTS.md) file. +> Detailed rules for agents are in [AGENTS.md](AGENTS.md). -If AI is used to generate any portion of the code, contributors must adhere to the following requirements: +What is allowed here that is not allowed upstream: -1. Explicitly disclose the manner in which AI was employed. -2. Check for an existing PR addressing the same change; if one exists, comment there to work with its author instead of opening a duplicate. -3. Perform a comprehensive manual review prior to submitting the pull request. A proper code review usually takes something like one hour per 200-400 LOC and you should be spending **at least that much time on code review alone**. -4. Be prepared to explain every line of code you submit when asked about it by a maintainer. -5. It is strictly prohibited to use AI to write your posts for you (bug reports, feature requests, pull request descriptions, Github discussions, responding to humans, ...). +- An agent may write the code, the commit messages, and the PR description. +- An agent may open the PR itself (`gh pr create`) and reply to review comments. +- An agent may run unattended, as long as it follows [AGENTS.md](AGENTS.md). -For more info, please refer to the [AGENTS.md](AGENTS.md) file. +What is still required, of humans and agents alike: + +1. **Disclose it.** Say in the PR that AI was used and how. Agent-authored commits carry an + `Assisted-by: ` trailer. +2. **Back the claims with measurements.** Performance and correctness claims about this hardware need numbers from + this hardware - see [Preparing your PR](#preparing-your-pr). +3. **Someone owns the result.** If you submit it, you answer for it: bugs, review feedback, and follow-ups. An agent + that opens a PR and disappears leaves the maintainers holding it, and such PRs may be closed. +4. **Do not duplicate.** Check for an existing PR or issue covering the same change first, and comment there instead + of opening a second one. +5. **Do not merge your own work**, and do not push to `master`. + +Undisclosed AI usage is the one thing that will get a contributor blocked here. Disclosure is cheap; a reviewer +discovering it later is not. # Pull requests (for contributors & collaborators) ### Before you start -- Search for existing discussions and PRs first - duplicates will likely be closed without questions. -- Features must begin with an issue, not a PR - let interest accumulate before writing code; niche features may only land as an example/tool, or on a private fork. -- Bug-fix PRs must include a reproducible issue and a regression test that fails before your change and passes after. Fixes without a test may be closed without review. -- New CLI or public API additions carry a **higher bar** than internal changes - justify why an existing mechanism doesn't suffice. -- Meeting all of the above still doesn't guarantee a merge - see [Pull requests (for maintainers)](#pull-requests-for-maintainers). -- If you are a new contributor - - Limit your open PRs to 1 - - Do not submit trivial fixes (e.g. typos, formatting changes) +- Search existing issues and PRs first - duplicates will likely be closed. +- Confirm the change belongs in this repo and not in [halo-box/llama.cpp](https://github.com/halo-box/llama.cpp); + see [Where does my change go?](#where-does-my-change-go). +- For anything large, or anything that introduces a new pattern or subsystem, open an issue first. Small, measured, + device-specific changes can go straight to a PR. +- Check whether upstream has already fixed it. Carrying a patch we do not need is a permanent merge cost. ### Preparing your PR -- llama.cpp uses the ggml tensor library for model evaluation. If you are unfamiliar with ggml, consider taking a look at the [examples in the ggml repository](https://github.com/ggml-org/ggml/tree/master/examples/). [simple](https://github.com/ggml-org/ggml/tree/master/examples/simple) shows the bare minimum for using ggml. [gpt-2](https://github.com/ggml-org/ggml/tree/master/examples/gpt-2) has minimal implementations for language model inference using GPT-2. [mnist](https://github.com/ggml-org/ggml/tree/master/examples/mnist) demonstrates how to train and evaluate a simple image classifier - Test your changes: - - Execute [the full CI locally on your machine](ci/README.md) before publishing - - Verify that the perplexity and the performance are not affected negatively by your changes (use `llama-perplexity` and `llama-bench`) - - If you modified the `ggml` source, run the `test-backend-ops` tool to check whether different backend implementations of the `ggml` operators produce consistent results (this requires access to at least two different `ggml` backends) - - If you modified a `ggml` operator or added a new one, add the corresponding test cases to `test-backend-ops` -- Create separate PRs for each feature or fix: - - Avoid combining unrelated changes in a single PR - - When adding support for a new model or feature, focus on **CPU support only** in the initial PR unless you have a good reason not to. Add support for other backends like CUDA in follow-up PRs - - In particular, adding new data types (extension of the `ggml_type` enum) carries with it a disproportionate maintenance burden. As such, to add a new quantization type you will need to meet the following *additional* criteria *at minimum*: - - convert a small model to GGUF using the new type and upload it to HuggingFace - - provide [perplexity](https://github.com/ggml-org/llama.cpp/tree/master/tools/perplexity) comparisons to FP16/BF16 (whichever is the native precision) as well as to types of similar size - - provide KL divergence data calculated vs. the FP16/BF16 (whichever is the native precision) version for both the new type as well as types of similar size - - provide [performance data](https://github.com/ggml-org/llama.cpp/tree/master/tools/llama-bench) for the new type in comparison to types of similar size on pure CPU -- Consider allowing write access to your branch for faster reviews, as reviewers can push commits directly + - Build and run on Strix Halo, on the backend you touched (Vulkan and/or ROCm) + - Produce the numbers required by [Benchmarking requirements](#benchmarking-requirements) - this is the part + reviewers look at first + - [Running the CI locally](ci/README.md) is the fastest way to catch what our CI would catch +- Bug fixes should come with a reproducible report and, where practical, a regression test - but reuse the existing + test infrastructure; do not add new files under `tests/` without a good reason. +- Keep hardware workarounds documented in place: what was measured, on what hardware and driver, and what would let + the workaround be removed. +- Create separate PRs for each feature or fix; avoid combining unrelated changes. +- Prefer small, local diffs. This fork merges upstream `master` regularly, and a sprawling change conflicts forever. +- Consider allowing write access to your branch for faster reviews, as reviewers can push commits directly. ### After submitting your PR -- Expect requests for modifications to ensure the code meets llama.cpp's standards for quality and long-term maintainability -- Maintainers will rely on your insights and approval when making a final decision to approve and merge a PR -- If your PR becomes stale, rebase it on top of latest `master` to get maintainers attention -- Consider adding yourself to [CODEOWNERS](CODEOWNERS) to indicate your availability for fixing related issues and reviewing related PRs +- Expect requests for modifications. Long-term maintainability matters more here than elsewhere, because every line + carried in this fork has to survive every upstream merge. +- If your PR becomes stale, rebase it on top of latest `master`. +- Consider adding yourself to [CODEOWNERS](CODEOWNERS) to indicate your availability for fixing related issues and reviewing related PRs. + +# Benchmarking requirements + +Almost everything in this fork exists because of a measurement, so measurements are the main thing a review here has +to trust. A PR that claims a change is faster, and does not show it on the hardware, is not reviewable and will be +asked for numbers before anything else. + +### When numbers are required + +Required for any change that touches a compute kernel, a dispatch or scheduling path, memory allocation, batching, +sampling, or speculative decoding - that is, anything that could move throughput, latency or memory use. + +Not required for docs, comments, CI configuration, or build-system changes that cannot affect runtime. Say so in the +PR and delete the section from the template. + +If you do not have access to Strix Halo hardware, you may still open the PR - mark it clearly as **unverified** and +say what you were unable to run. Do not report numbers from a different GPU as if they applied here. + +### Report the environment + +Paste this into the PR, filled in. Most disagreements about benchmark results turn out to be a difference in one of +these lines: + +``` +Device: Ryzen AI Max+ 395 / +Memory: 128 GB LPDDR5X-8000 +Power: +BIOS: UMA split GB +Kernel: , amdgpu params: +Backend: Vulkan RADV, Mesa (or: ROCm , HIP_LAUNCH_BLOCKING=<0|1>) +Build: +Baseline: +Change: +Model: , +``` + +Power and thermals matter more here than on a discrete GPU: the CPU and iGPU share a package power budget and a +memory controller, and the same silicon ships in chassis with very different sustained TDP. Numbers without a stated +power profile are not comparable between contributors. + +### Run the baseline yourself + +The baseline is a binary you built from the merge-base of your branch, with the same cmake flags, on the same machine, +in the same session as the "after" run. Not numbers from a previous day, another machine, a release build, or memory. +Interleave or alternate the two runs if the machine's thermal state drifts. + +### Use llama-bench, and paste the table + +```sh +# prompt processing and token generation, 5 repetitions (the default) +./build/bin/llama-bench -m -p 512 -n 128 -r 5 + +# add context depth when the change can affect long-context behaviour +./build/bin/llama-bench -m -p 512 -n 128 -d 0,4096,16384 -r 5 +``` + +- Paste the raw table, including the standard-deviation column. Do not replace it with "about 18% faster". +- Raise `-r` if the spread is wide. Five repetitions is a floor, not a target. +- Vary the axis your change actually acts on. A change to batched mat-vec has to be shown across batch sizes + (`-ub`), not at one convenient point; a flash-attention change has to be shown with `-fa` both ways. +- If before and after overlap inside their standard deviations, you have not measured a speedup yet. + +`llama-bench` is a floor, not a ceiling. Changes to speculative decoding, the server, or anything whose benefit +depends on the content being generated cannot be shown with it - measure end to end against `llama-server` with a +described workload, and say what the workload was. + +For per-op attribution on the Vulkan backend, `GGML_VK_PERF_LOGGER=1` gives per-node timings, which is the right +evidence for "this shader variant is slow at this shape" claims. + +### Show that output did not change + +A speedup that changes results is a bug, so speed numbers alone are not enough: + +- Touched `ggml`: run `test-backend-ops` in the default `test` mode for the backend you changed, and narrow with + `-o ` while iterating. + ```sh + ./build/bin/test-backend-ops -b Vulkan0 -o MUL_MAT + ``` +- Touched anything that can alter generated tokens: report `llama-perplexity` before and after on the same file and + model. State the value for both; a perplexity shift that you cannot explain blocks the PR. +- Say which of these you did **not** run. An honest gap is reviewable; a silent one is not. # Pull requests (for maintainers) - Squash-merge PRs -- Use the following format for the squashed commit title: ` : (#)`. For example: `utils : fix typo in utils.py (#1234)` -- Optionally pick a `` from here: https://github.com/ggml-org/llama.cpp/wiki/Modules +- Use the following format for the squashed commit title: ` : (#)`. For example: `vulkan : fix mmvq selection on gfx1151 (#1234)` - Let other maintainers merge their own PRs - When merging a PR, make sure you have a good understanding of the changes -- If a PR does not warrant a new release, add `[no release]` in the squashed commit to spare CI resources -- Be mindful of maintenance: most of the work going into a feature happens after the PR is merged. If the PR author is not committed to contribute long-term, someone else needs to take responsibility (you) -- Add the ["merge ready"](https://github.com/ggml-org/llama.cpp/pulls?q=is%3Apr+is%3Aopen+draft%3Ano+sort%3Aupdated-desc+label%3A%22merge+ready%22+) label to a PR to indicate when a PR can be fast-merged without waiting for 2 independent reviews. [(more info)](https://github.com/ggml-org/llama.cpp/pull/26178) -- Wait for CI results before merging +- Prefer changes that upstream could plausibly accept later, even though we do not submit them from here - it keeps + merges cheap +- Wait for CI results before merging, including the self-hosted `gfx1151` job where it applies -Maintainers reserve the right to decline review or close pull requests for any reason, without any questions, particularly under any of the following conditions: -- The proposed change is already mentioned in the roadmap or an existing issue, and it has been assigned to someone. +Maintainers reserve the right to decline review or close pull requests, particularly under any of the following conditions: +- The change belongs in [halo-box/llama.cpp](https://github.com/halo-box/llama.cpp) or upstream instead. - The pull request duplicates an existing one. - The contributor fails to adhere to this contributing guide or the AI policy. -- The change doesn't fit the existing architecture, or is too complex to justify its benefit. +- Performance claims are not backed by measurements from the hardware. +- Nobody is available to own the change after it lands. # Coding guidelines @@ -189,7 +301,7 @@ Maintainers reserve the right to decline review or close pull requests for any r - When adding or modifying a large piece of code: - If you are a collaborator, make sure to add yourself to [CODEOWNERS](CODEOWNERS) to indicate your availability for reviewing related PRs - If you are a contributor, find an existing collaborator who is willing to review and maintain your code long-term - - Provide the necessary CI workflow (and hardware) to test your changes (see [ci/README.md](https://github.com/ggml-org/llama.cpp/tree/master/ci)) + - Provide the necessary CI workflow (and hardware) to test your changes (see [ci/README.md](ci/README.md)) - New code should follow the guidelines (coding, naming, etc.) outlined in this document. Exceptions are allowed in isolated, backend-specific parts of the code that do not interface directly with the `ggml` interfaces. _(NOTE: for legacy reasons, existing code is not required to follow this guideline)_ @@ -204,6 +316,12 @@ Maintainers reserve the right to decline review or close pull requests for any r # Resources -The Github issues, PRs and discussions contain a lot of information that can be useful to get familiar with the codebase. For convenience, some of the more important information is referenced from Github projects: +- [AGENTS.md](AGENTS.md) - rules for AI coding agents working in this repo, including how to open a PR here +- [README.md](README.md) - what this fork is, how to build it, and the Strix Halo specific notes +- [Issues](https://github.com/halo-box/strix-llama.cpp/issues) and [PRs](https://github.com/halo-box/strix-llama.cpp/pulls) in this repo +- [halo-box/llama.cpp](https://github.com/halo-box/llama.cpp) - where upstream-bound changes go instead + +Upstream's issues, PRs and discussions remain the best source of background on the codebase itself, and its Github +projects collect the most important of it: https://github.com/ggml-org/llama.cpp/projects diff --git a/README.md b/README.md index aae3bcd35ad9..032323214281 100644 --- a/README.md +++ b/README.md @@ -1,91 +1,140 @@ -# llama.cpp +# strix-llama.cpp -![llama](https://raw.githubusercontent.com/ggml-org/llama.brand/refs/heads/master/cover/llama-cpp/cover-llama-cpp-dark.svg) +Halo Box -
- -LLM inference in C/C++ +llama.cpp for AMD Strix Halo [![License: MIT](https://img.shields.io/badge/license-MIT-blue.svg)](https://opensource.org/licenses/MIT) -[![Release](https://img.shields.io/github/v/release/ggml-org/llama.cpp?filter=v*&color=brightgreen)](https://github.com/ggml-org/llama.cpp/releases?q=tag:v0) -[![Nightly](https://img.shields.io/github/v/release/ggml-org/llama.cpp?label=nightly&filter=b*&color=orange)](https://github.com/ggml-org/llama.cpp/releases?q=b) -[![Server](https://img.shields.io/github/actions/workflow/status/ggml-org/llama.cpp/server.yml?label=Server)](https://github.com/ggml-org/llama.cpp/actions/workflows/server.yml) -[![Docker](https://img.shields.io/github/actions/workflow/status/ggml-org/llama.cpp/docker.yml?label=Docker)](https://github.com/ggml-org/llama.cpp/actions/workflows/docker.yml) -[![Winget](https://img.shields.io/github/actions/workflow/status/ggml-org/llama.cpp/winget.yml?label=Winget)](https://github.com/ggml-org/llama.cpp/actions/workflows/winget.yml) -[ggml](https://github.com/ggml-org/ggml) / [ops](https://github.com/ggml-org/llama.cpp/blob/master/docs/ops.md) / [maintainer PRs](https://github.com/ggml-org/llama.cpp/issues?q=is%3Apr%20is%3Aopen%20draft%3AFalse%20(author%3Argerganov%20OR%20author%3AKitaitiMakoto%20OR%20author%3Adanbev%20OR%20author%3Aaldehir%20OR%20author%3Amax-krasnyansky%20OR%20author%3ACISC%20OR%20author%3Aggerganov%20OR%20author%3Aam17an%20OR%20author%3Ajhen0409%20OR%20author%3Abartowski1182%20OR%20author%3Anikwen%20OR%20author%3Ahipudding%20OR%20author%3Aravi9%20OR%20author%3AServeurpersoCom%20OR%20author%3Apwilkin%20OR%20author%3Areeselevine%20OR%20author%3Angxson%20OR%20author%3Ajeffbolznv%20OR%20author%3Amarty1885%20OR%20author%3A0cc4m%20OR%20author%3ATitaniumtown%20OR%20author%3Aangt%20OR%20author%3AIMbackK%20OR%20author%3Aarthw%20OR%20author%3AJohannesGaessler%20OR%20author%3AORippler%20OR%20author%3Aruixiang63%20OR%20author%3Axctan%20OR%20author%3Aallozaur%20OR%20author%3Ayomaytk%20OR%20author%3Aaendk%20OR%20author%3Awine99%20OR%20author%3Agaugarg-nv%20OR%20author%3Ataronaeo%20OR%20author%3Aforforever73%20OR%20author%3Alhez%20OR%20author%3Anetrunnereve%20OR%20author%3Afairydreaming)%20sort%3Aupdated-desc) / [dev stats](https://github.com/ggml-org/llama.cpp-dev) / [lib llama API](https://github.com/ggml-org/llama.cpp/issues/9289) / [llama-server REST API](https://github.com/ggml-org/llama.cpp/issues/9291) +[upstream llama.cpp](https://github.com/ggml-org/llama.cpp) / [ggml](https://github.com/ggml-org/ggml) / [halo-box/llama.cpp](https://github.com/halo-box/llama.cpp) + +## Halo Box + +The goal is simple: more functionality, and the fastest llama.cpp around. And help the community with a single fast +llama.cpp fork instead of many competing ones. + +Halo Box keeps two forks, and which one you want depends on your hardware: + +| Fork | What it is | +| --- | --- | +| [halo-box/llama.cpp](https://github.com/halo-box/llama.cpp) | Stays close to mainline. Tracks upstream `master` and adds features and speedups on top, without diverging from how upstream works. | +| [halo-box/strix-llama.cpp](https://github.com/halo-box/strix-llama.cpp) (this repo) | Purely optimised for AMD Strix Halo machines (Ryzen AI Max+, RDNA 3.5 / gfx1151). Free to diverge from upstream wherever that buys speed. | + +Use `halo-box/llama.cpp` if you want upstream behaviour plus extras. Use this repo if you run a Strix Halo box and +want every last token/s out of it. Everything in `halo-box/llama.cpp` is merged in here regularly, so this repo is a +superset of it. + +## What this is + +A community fork of [`llama.cpp`](https://github.com/ggml-org/llama.cpp) for **AMD Strix Halo** - the Ryzen AI Max / Max+ 300 series +APUs (`gfx1151`, RDNA 3.5 integrated GPU, up to 128 GB of unified LPDDR5X shared between CPU and GPU). -
+Strix Halo is an unusual target. It has more addressable memory than almost any consumer discrete GPU, and far less +bandwidth; the iGPU shares its memory controller with the CPU; and both the Vulkan (RADV) and ROCm/HIP paths have +RDNA 3.5 specific behaviour that upstream has no hardware to reproduce. Changes that only make sense on this one +device - or that need a lot of measurement on it before they are ready to propose anywhere else - live here. + +Practically, that means this repo has its own rules - most visibly, **AI coding agents may open pull requests here**. +See [CONTRIBUTING.md](CONTRIBUTING.md) and [AGENTS.md](AGENTS.md). ## Quick start -A few options to get `llama.cpp` installed on your machine: +Build from source. Two backends are worth using on Strix Halo: + +**Vulkan** (RADV on Mesa; the easiest path, and the best one for most models) + +```sh +cmake -B build -DGGML_VULKAN=ON -DCMAKE_BUILD_TYPE=Release +cmake --build build --config Release -j +``` + +**ROCm / HIP** (needs ROCm installed; build for `gfx1151` explicitly) -- Visit https://llama.app and follow the instructions -- Run with Docker - see our [Docker documentation](docs/docker.md) -- Download pre-built binaries from the [releases page](https://github.com/ggml-org/llama.cpp/releases) -- Build from source by cloning this repository - check out [our build guide](docs/build.md) +```sh +HIPCXX="$(hipconfig -l)/clang" HIP_PATH="$(hipconfig -R)" \ + cmake -B build -DGGML_HIP=ON -DGPU_TARGETS=gfx1151 -DCMAKE_BUILD_TYPE=Release +cmake --build build --config Release -j +``` -Once installed: +Then: ```sh -# Download and run a model directly from Hugging Face -llama cli -hf ggml-org/Qwen3.5-0.8B-GGUF +# chat, pulling the model straight from Hugging Face +./build/bin/llama-cli -hf ggml-org/Qwen3.5-0.8B-GGUF -# Launch OpenAI-compatible API server -llama serve -hf ggml-org/Qwen3.5-0.8B-GGUF +# OpenAI-compatible API server + web UI on http://localhost:8080 +./build/bin/llama-server -hf ggml-org/Qwen3.5-0.8B-GGUF ``` - - - - - -
- VLM session with `llama cli` - VLM session with llama cli - - Built-in web UI against `llama serve` running Qwen 3.6 - Built-in web UI against llama serve -
- -## Description - -The main goal of `llama.cpp` is to enable LLM (and VLM) inference with minimal setup and state-of-the-art performance on -a wide range of hardware - locally and in the cloud. - -- Plain C/C++ implementation without any dependencies -- Apple silicon is a first-class citizen - optimized via ARM NEON, Accelerate and Metal frameworks -- AVX, AVX2, AVX512 and AMX support for x86 architectures -- RVV, ZVFH, ZFH, ZICBOP and ZIHINTPAUSE support for RISC-V architectures -- 1.5-bit, 2-bit, 3-bit, 4-bit, 5-bit, 6-bit, and 8-bit integer quantization for faster inference and reduced memory use -- Custom CUDA kernels for running LLMs on NVIDIA GPUs (support for AMD GPUs via HIP and Moore Threads GPUs via MUSA) -- Vulkan and SYCL backend support -- CPU+GPU hybrid inference to partially accelerate models larger than the total VRAM capacity - -The `llama.cpp` project is build on top of the [ggml](https://github.com/ggml-org/ggml) library. +Full build documentation, including Windows and Docker, is in [docs/build.md](docs/build.md). + +## Running on Strix Halo + +**Give the iGPU enough memory.** The APU's memory is shared, and the GPU can only use what the firmware and kernel let +it map. Two things control this: the UMA / dedicated-VRAM split in your BIOS, and the `amdgpu` GTT limit on Linux +(`amdgpu.gttsize`, in MB, and `ttm.pages_limit`, in 4 KB pages, as kernel command-line parameters). Which of those you +need depends on your kernel version - newer kernels size GTT more generously on their own. If a model that clearly +fits in RAM fails to allocate, this is almost always why. + +**ROCm and batched inference.** On `gfx1151` there is an async-execution correctness bug in the HIP path: batched +inference can return badly wrong output (perplexity ~88 against ~9.4 for the same model). Setting +`HIP_LAUNCH_BLOCKING=1` serializes kernel launches and restores correctness, at a performance cost. Our ROCm CI runs +with it set. It is a workaround for a ROCm/HIP issue, not a fix, and it should go away when that is fixed upstream. + +**Vulkan mat-vec chunking.** Batched mat-vec at 3, 5 and 6 columns is several times slower than at 1, 2 and 4 on RADV +here, which hits speculative decoding hard (it verifies at exactly those batch sizes). This fork splits such batches +into column counts that are measured to scale. Set `GGML_VK_MMV_NO_SPLIT=1` to restore the single upstream dispatch, +e.g. to compare against it. + +**Measure things.** `GGML_VK_PERF_LOGGER=1` (any value) gives per-op timings on the Vulkan backend and is how most of the findings +above were made. `llama-bench` and `llama-perplexity` are the tools for before/after numbers, and PRs here are +expected to carry them - see [Benchmarking requirements](CONTRIBUTING.md#benchmarking-requirements). + +## What differs from upstream + +Everything else is upstream `llama.cpp`. The additions currently carried here: + +**Inherited from [halo-box/llama.cpp](https://github.com/halo-box/llama.cpp)** (general features, sent upstream from there) + +| Change | Flag / switch | What it does | +| --- | --- | --- | +| Speculative prefill | `--spec-prefill` | A small draft model scores prompt tokens by attention importance so the target model only prefills the ones that matter, cutting time-to-first-token on long prompts | +| Adaptive speculative draft length | `--spec-draft-adaptive` | Sizes each draft from a measured per-sequence acceptance EMA rather than always drafting `--spec-draft-n-max`; speeds up MTP and DFlash | +| Vulkan fixes and tuning for RDNA 3.5 | | Driver-gated coopmat LDS stride padding, UMA bulk readback gated on host-cached mappings, IQ3_S mat-vec at batch sizes > 4, and a radix top-k kernel for large k | +| Hidden server presets | `hidden` in the models `.ini` | Keep a model loadable by name while omitting it from `GET /models` | + +**Strix Halo only** (lives here, measured on `gfx1151`) + +| Change | Flag / switch | What it does | +| --- | --- | --- | +| Multi-point reasoning budget | `--reasoning-budget-*` (upstream has the hard budget only) | Intro message, two soft warnings, a grace period to finish a paragraph after the budget runs out, and reasoning-token usage telemetry | +| Vulkan batched mat-vec chunking | `GGML_VK_MMV_NO_SPLIT=1` to disable | Dispatches batched mat-vec at the column counts that actually scale on RDNA 3.5, instead of the slow NUM_COLS shader variants | +| ROCm/HIP quantized matmul on RDNA 3.5 | | MMQ/MMVQ tile configurations and register prefetching, compact `MUL_MAT_ID` with quant- and shape-specific tiles for 256-expert MoE prefill, fused activation quantization for Q8_0 and Q6_K decode | +| ROCm/HIP MoE decode fusion | `GGML_CUDA_DISABLE_WEIGHTED_DOWN=1`, `GGML_CUDA_DISABLE_MMID_512=1` | Fused routing, weighted expert reduction and shared-expert gate for Qwen3.5/3.6 and Ling style MoE | +| ROCm/HIP Gated DeltaNet | `GGML_CUDA_DISABLE_GDN_GATE=1` | DPP reductions and a tiled multi-column kernel for prefill; the whole conv -> norm -> gate -> recurrence -> state copy decode chain as one kernel | +| ROCm/HIP grouped decode matvecs | `GGML_CUDA_DISABLE_MMV_GROUP=1` | Consecutive single-column matvecs that read the same activation (gate/up pairs, hyper-connection projections) launch as one kernel | +| ROCm/HIP flash attention on RDNA 3.5 | | WMMA path for D=256 prefill at depth, Q8_0 KV tile kernel for decode | +| Speculative checkpoints on device | | `llama-server` keeps speculative-decoding checkpoints in device memory instead of copying them to the host | +| ROCmFPx quant types | `llama-quantize` types `Q4_0_ROCMFP4`, `Q4_0_ROCMFP4_FAST`, `Q2/Q3/Q6/Q8_0_ROCMFPX` and the `_LEAN`/`_COHERENT`/`_STRIX` recipes | Loads the ROCmFP4 GGUFs published for Strix Halo. CPU codecs plus Vulkan dequant, mat-vec, matmul and integer-dot kernels. Weight formats only: not accepted as KV-cache types | +| Repeatable output at depth | | Freed KV cells are zeroed so masked-out rows never carry stale K/V, and the Vulkan radix top-k assigns output slots deterministically | + +Every ROCm/HIP change above is guarded on architecture, shape and layout, so other devices see upstream behaviour. +Run `--help`, or see [tools/server/README.md](tools/server/README.md), for the full options. ## Supported backends -| Backend | Target devices | +The ones that matter on this hardware: + +| Backend | Notes | | --- | --- | -| [BLAS](docs/build.md#blas-build) | All | -| [BLIS](docs/backend/BLIS.md) | All | -| [CANN](docs/build.md#cann) | Ascend NPU | -| [CUDA](docs/build.md#cuda) | Nvidia GPU | -| [HIP](docs/build.md#hip) | AMD GPU | -| [Hexagon](docs/backend/snapdragon/README.md) | Snapdragon | -| [IBM zDNN](docs/backend/zDNN.md) | IBM Z & LinuxONE | -| [MUSA](docs/build.md#musa) | Moore Threads GPU | -| [Metal](docs/build.md#metal-build) | Apple Silicon | -| [OpenCL](docs/backend/OPENCL.md) | Adreno GPU | -| [OpenVINO [In Progress]](docs/backend/OPENVINO.md) | Intel CPUs, GPUs, and NPUs | -| [RPC](https://github.com/ggml-org/llama.cpp/tree/master/tools/rpc) | All | -| [SYCL](docs/backend/SYCL.md) | Intel GPU | -| [VirtGPU](docs/backend/VirtGPU.md) | VirtGPU APIR | -| [Vulkan](docs/build.md#vulkan) | GPU | -| [WebGPU](docs/build.md#webgpu) | All | -| [ZenDNN](docs/build.md#zendnn) | AMD CPU | +| [Vulkan](docs/build.md#vulkan) | RADV on the RDNA 3.5 iGPU - the default recommendation | +| [HIP](docs/build.md#hip) | ROCm on `gfx1151` - see the `HIP_LAUNCH_BLOCKING` note above | +| [CPU](docs/build.md) | Zen 5 cores with AVX-512 - useful for offloading part of a model, though it shares bandwidth with the iGPU | +| [ZenDNN](docs/build.md#zendnn) | AMD CPU acceleration | +| [RPC](tools/rpc) | Distribute a model across several machines | + +`llama.cpp` supports many more (CUDA, Metal, SYCL, CANN, OpenCL, WebGPU, ...); they are all still present in the tree +and unmodified. See the [upstream README](https://github.com/ggml-org/llama.cpp) for that list. ## Documentation @@ -100,25 +149,34 @@ The `llama.cpp` project is build on top of the [ggml](https://github.com/ggml-or - [How to build](docs/build.md) - [Running on Docker](docs/docker.md) -- [Build on Android](docs/android.md) - [Multi-GPU usage](docs/multi-gpu.md) - [Performance troubleshooting](docs/development/token_generation_performance_tips.md) - [GGML tips & tricks](https://github.com/ggml-org/llama.cpp/wiki/GGML-Tips-&-Tricks) -- [XCFramework](docs/xcframework.md) - [Completions](docs/completions.md) - [Models](docs/models.md) -- [Release process](docs/release.md) ## Contributing -- Contributors can open PRs -- Collaborators will be invited based on contributions -- Maintainers can push to branches in the `llama.cpp` repo and merge PRs into the `master` branch -- Any help with managing issues, PRs and projects is very appreciated! -- Read the [CONTRIBUTING.md](CONTRIBUTING.md) for more information +This is a small community project. Strix Halo owners with a benchmark, a bug report, or a patch are exactly who it is +for. + +- Anyone can open a PR, **including AI coding agents acting on their own** - this repo permits automated PR submission, + unlike upstream. The rules that replace the upstream ban are in [AGENTS.md](AGENTS.md). +- Device-specific claims need numbers from the device, against a baseline you built and ran yourself. What exactly to + report is in [Benchmarking requirements](CONTRIBUTING.md#benchmarking-requirements); read it before you benchmark. +- If your change is not Strix Halo specific, send it to [halo-box/llama.cpp](https://github.com/halo-box/llama.cpp) + instead, so it can reach upstream. +- Read [CONTRIBUTING.md](CONTRIBUTING.md) before your first PR. + +CI runs the standard llama.cpp suite plus a self-hosted `gfx1151` ROCm job on real hardware. ## Acknowledgements +This project is a fork and owes everything to the people who built what it forks: + +- [ggml-org/llama.cpp](https://github.com/ggml-org/llama.cpp) and [ggml](https://github.com/ggml-org/ggml) - Georgi Gerganov and the llama.cpp contributors - MIT license +- [charlie12345/ROCmFPX](https://github.com/charlie12345/ROCmFPX) - the origin of the ROCmFPx project and the creator of the ROCmFP4 format - MIT license +- [ciru-ai/ROCmFPX](https://github.com/ciru-ai/ROCmFPX) - a fork of the above, and the tree the ROCmFPx quant formats and reference codecs here were hand-ported from (it shares no git history with llama.cpp, so it cannot be merged) - MIT license - [yhirose/cpp-httplib](https://github.com/yhirose/cpp-httplib) - Single-header HTTP server, used by `llama-server` - MIT license - [nothings/stb](https://github.com/nothings/stb) - Single-header image format decoder, used by multimodal subsystem - Public domain - [nlohmann/json](https://github.com/nlohmann/json) - Single-header JSON library, used by various tools/examples - MIT License diff --git a/SYNC.md b/SYNC.md new file mode 100644 index 000000000000..7f6bf7712ed9 --- /dev/null +++ b/SYNC.md @@ -0,0 +1,71 @@ +# Syncing halo-box/llama.cpp with upstream + +## Why GitHub says "N commits behind" after a sync + +GitHub computes that number as `git rev-list --count HEAD..upstream/master` — commits +reachable from upstream but not from us. It is a **graph reachability** question and has +nothing to do with whether our files match upstream's. + +A squash merge copies content without copying ancestry: the resulting commit's only parent +is our previous tip, so upstream's commits never become our ancestors. The counter stays +pinned forever, and worse, `git merge-base` stays pinned too — so the *next* sync replays +every already-applied commit as a conflict. + +Squash is the right tool for collapsing our own PRs. It is the wrong tool for ingesting +upstream. + +## Preferred: real merge + +```sh +git fetch upstream +git merge upstream/master +``` + +Merge commits are noisy. They are also correct, and they keep merge-base moving. Resolve +conflicts once, in the merge, and the ancestry is recorded for free. + +## When a PR + squash was unavoidable + +If divergence forced the sync through a PR (conflict resolution needs review), the content +lands but the ancestry does not. Two things must then happen. + +### 1. Audit for dropped commits + +Conflict resolution silently drops upstream hunks. Find them before recording the merge: + +```sh +MB=$(git merge-base HEAD upstream/master) + +# Files upstream touched that we never touched => upstream change dropped entirely +git diff --name-only $MB upstream/master | while read -r f; do + git diff --quiet HEAD upstream/master -- "$f" && continue + n=$(git diff --numstat $MB HEAD -- "$f" | awk '{print $1+$2}') + [ -z "$n" ] && echo "DROPPED: $f" +done +``` + +Recover each one with `git cherry-pick -x `. Re-run until the list is empty. + +A file that differs from upstream *and* that we modified is normal fork divergence — review +it by hand, the script cannot judge intent. + +### 2. Record the ancestry + +Only once the audit is clean: + +```sh +git merge -s ours upstream/master -m "Record upstream sync (content applied via #)" +``` + +`-s ours` creates a two-parent merge commit and keeps our tree **byte-identical** — it +cannot change a single file, so it cannot break the build. Its only danger is burying an +upstream change we failed to apply, which is exactly what step 1 rules out. + +Never run `-s ours` before the audit. + +## Verifying a sync + +```sh +git rev-list --count HEAD..upstream/master # expect 0 +git log -1 --format='%P' HEAD # expect two parents after -s ours +``` diff --git a/common/CMakeLists.txt b/common/CMakeLists.txt index 2b307c59d32d..b061bf688e67 100644 --- a/common/CMakeLists.txt +++ b/common/CMakeLists.txt @@ -105,6 +105,8 @@ add_library(${TARGET} reasoning-budget.h sampling.cpp sampling.h + speculative-prefill.cpp + speculative-prefill.h speculative.cpp speculative.h subproc.cpp diff --git a/common/arg.cpp b/common/arg.cpp index c4c4e143c987..52948e045954 100644 --- a/common/arg.cpp +++ b/common/arg.cpp @@ -311,6 +311,10 @@ const std::vector kv_cache_types = { GGML_TYPE_IQ4_NL, GGML_TYPE_Q5_0, GGML_TYPE_Q5_1, + // ROCmFPx types are weight formats here, not KV cache types: flash attention + // has no decode path for them (fa_kv_ok in ggml-vulkan.cpp rejects them), and + // a quantized V cache force-enables flash attention in llama_init_from_model, + // which skips the support probe that would otherwise turn it back off. }; static ggml_type kv_cache_type_from_str(const std::string & s) { @@ -361,6 +365,7 @@ static bool spec_types_is_default(const common_params & params) { common_models_handler common_models_handler_init(const common_params & params, llama_example curr_ex) { common_download_hf_plan plan; common_download_hf_plan plan_spec; + common_download_hf_plan plan_prefill; common_download_opts opts; const bool spec_type_draft_mtp = std::find(params.speculative.types.begin(), @@ -413,7 +418,11 @@ common_models_handler common_models_handler_init(const common_params & params, l plan_spec = common_download_get_hf_plan(params.speculative.draft.mparams, opts_spec); } - return common_models_handler{plan, plan_spec, opts}; + if (!params.speculative.prefill.model.hf_repo.empty()) { + plan_prefill = common_download_get_hf_plan(params.speculative.prefill.model, opts); + } + + return common_models_handler{plan, plan_spec, plan_prefill, opts}; } bool common_models_handler_is_preset_repo(const common_models_handler & handler) { @@ -461,8 +470,9 @@ static std::vector build_url_tasks(const common_params_mod void common_models_handler_apply(common_models_handler & handler, common_params & params, common_download_callback * callback) { std::vector tasks; - auto & plan = handler.plan; - auto & plan_spec = handler.plan_spec; + auto & plan = handler.plan; + auto & plan_spec = handler.plan_spec; + auto & plan_prefill = handler.plan_prefill; auto opts = handler.opts; // copy opts.callback = callback; @@ -478,6 +488,7 @@ void common_models_handler_apply(common_models_handler & handler, common_params handle_url(params.model); handle_url(params.mmproj); handle_url(params.speculative.draft.mparams); + handle_url(params.speculative.prefill.model); // optionally, if docker repo is set, resolve it if (!params.model.docker_repo.empty()) { @@ -513,6 +524,13 @@ void common_models_handler_apply(common_models_handler & handler, common_params tasks.push_back(task); had_spec_url = true; } + if (!params.speculative.prefill.model.url.empty()) { + common_download_task task; + task.url = params.speculative.prefill.model.url; + task.local_path = params.speculative.prefill.model.path; + task.opts = opts; + tasks.push_back(task); + } // handle hf_plan tasks auto add_tasks = [&opts, &tasks](const hf_cache::hf_files & model_files, @@ -626,6 +644,11 @@ void common_models_handler_apply(common_models_handler & handler, common_params had_spec_url = true; } + // handle plan_prefill (e.g. --spec-prefill-hf) + if (!plan_prefill.model_files.empty()) { + add_tasks(plan_prefill.model_files, plan_prefill.primary, params.speculative.prefill.model); + } + if (!plan.model_files.empty()) { add_tasks(plan.model_files, plan.primary, params.model); } @@ -2706,10 +2729,12 @@ common_params_context common_params_parser_init(common_params & params, llama_ex {"-lzm", "--lazy-mode"}, "MODE", "on-demand reading of certain tensors, for example per-layer embeddings (default: auto)\n" "- on: read the rows of such tensors from disk on demand instead of keeping them resident (requires mmap)\n" + "- on-direct: like on, but the arch reads the rows with explicit pread()s instead of demand paging the mmap (currently only PLE tables, page-cached)\n" "- auto: on, but only for tensors larger than 4 GiB\n" "- off: always keep them resident", [](common_params & params, const std::string & value) { /**/ if (value == "on") { params.lazy_mode = LLAMA_LAZY_MODE_ON; } + else if (value == "on-direct") { params.lazy_mode = LLAMA_LAZY_MODE_DIRECT; } else if (value == "auto") { params.lazy_mode = LLAMA_LAZY_MODE_AUTO; } else if (value == "off") { params.lazy_mode = LLAMA_LAZY_MODE_OFF; } else { throw std::invalid_argument("invalid value"); } @@ -2752,6 +2777,13 @@ common_params_context common_params_parser_init(common_params & params, llama_ex parse_tensor_buffer_overrides(value, params.tensor_buft_overrides); } ).set_env("LLAMA_ARG_OVERRIDE_TENSOR")); + add_opt(common_arg( + {"--ngram-on-disk"}, + "deprecated alias for --lazy-mode on", + [](common_params & params) { + params.lazy_mode = LLAMA_LAZY_MODE_ON; + } + ).set_env("LLAMA_ARG_NGRAM_ON_DISK")); add_opt(common_arg( {"-cmoe", "--cpu-moe"}, "keep all Mixture of Experts (MoE) weights in the CPU", @@ -3703,6 +3735,50 @@ common_params_context common_params_parser_init(common_params & params, llama_ex } } ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_REASONING_EFFORT")); + static auto unescape = [](const std::string & s) { + std::string r; + r.reserve(s.size()); + for (size_t i = 0; i < s.size(); i++) { + if (s[i] == '\\' && i + 1 < s.size()) { + switch (s[i + 1]) { + case 'n': + r += '\n'; + i++; + break; + case 't': + r += '\t'; + i++; + break; + case 'r': + r += '\r'; + i++; + break; + case '\\': + r += '\\'; + i++; + break; + case '"': + r += '"'; + i++; + break; + default: + r += s[i]; + break; + } + } else { + r += s[i]; + } + } + return r; + }; + + add_opt(common_arg({ "--reasoning-budget-enable" }, { "--no-reasoning-budget-enable" }, + "enable the reasoning budget mechanism: hard cutoff, soft warning, intro message, grace " + "period, and runtime reasoning control (default: disabled)", + [](common_params & params, bool value) { params.sampling.reasoning_budget_enabled = value; }) + .set_examples({ LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI }) + .set_env("LLAMA_ARG_THINK_BUDGET_ENABLE")); + add_opt(common_arg( {"--reasoning-budget"}, "N", "token budget for thinking: -1 for unrestricted, 0 for immediate end, N>0 for token budget (default: -1)", @@ -3711,13 +3787,68 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.sampling.reasoning_budget_tokens = value; } ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_THINK_BUDGET")); - add_opt(common_arg( - {"--reasoning-budget-message"}, "MESSAGE", - "message injected before the end-of-thinking tag when reasoning budget is exhausted (default: none)", - [](common_params & params, const std::string & value) { - params.sampling.reasoning_budget_message = value; - } - ).set_examples({LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_THINK_BUDGET_MESSAGE")); + add_opt(common_arg({ "--reasoning-budget-message" }, "MESSAGE", + "message forced when the reasoning budget is exhausted; include the closing tag because it is " + "not appended automatically (default: none)", + [](common_params & params, const std::string & value) { + params.sampling.reasoning_budget_message = unescape(value); + }) + .set_examples({ LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI }) + .set_env("LLAMA_ARG_THINK_BUDGET_MESSAGE")); + add_opt(common_arg({ "--reasoning-budget-soft-ratio" }, "N", + "fraction of the budget consumed before the first soft warning: <= 0 disables it, " + "(0,1] enables it (default: -1)", + [](common_params & params, const std::string & value) { + params.sampling.reasoning_budget_soft_ratio = std::stof(value); + }) + .set_examples({ LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI }) + .set_env("LLAMA_ARG_THINK_BUDGET_SOFT_RATIO")); + add_opt(common_arg({ "--reasoning-budget-soft-message" }, "MESSAGE", + "message forced at the first soft reasoning budget threshold (default: none)", + [](common_params & params, const std::string & value) { + params.sampling.reasoning_budget_soft_message = unescape(value); + }) + .set_examples({ LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI }) + .set_env("LLAMA_ARG_THINK_BUDGET_SOFT_MESSAGE")); + add_opt(common_arg({ "--reasoning-budget-soft2-ratio" }, "N", + "fraction of the budget consumed before the second soft warning: <= 0 disables it, " + "(0,1] enables it (default: -1)", + [](common_params & params, const std::string & value) { + params.sampling.reasoning_budget_soft2_ratio = std::stof(value); + }) + .set_examples({ LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI }) + .set_env("LLAMA_ARG_THINK_BUDGET_SOFT2_RATIO")); + add_opt(common_arg({ "--reasoning-budget-soft2-message" }, "MESSAGE", + "message forced at the second soft reasoning budget threshold (default: none)", + [](common_params & params, const std::string & value) { + params.sampling.reasoning_budget_soft2_message = unescape(value); + }) + .set_examples({ LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI }) + .set_env("LLAMA_ARG_THINK_BUDGET_SOFT2_MESSAGE")); + add_opt(common_arg({ "--reasoning-budget-intro-mode" }, "MODE", + "force the intro message for every reasoning block or only once per conversation " + "(default: every)", + [](common_params & params, const std::string & value) { + params.sampling.reasoning_budget_intro_mode = value; + }) + .set_examples({ LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI }) + .set_env("LLAMA_ARG_THINK_BUDGET_INTRO_MODE")); + add_opt(common_arg({ "--reasoning-budget-intro-message" }, "MESSAGE", + string_format( + "message forced when the reasoning block starts; use {budget} for the configured " + "budget (default: '%s')", + params.sampling.reasoning_budget_intro_message.c_str()), + [](common_params & params, const std::string & value) { + params.sampling.reasoning_budget_intro_message = unescape(value); + }) + .set_examples({ LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI }) + .set_env("LLAMA_ARG_THINK_BUDGET_INTRO_MESSAGE")); + add_opt(common_arg({ "--reasoning-budget-grace-tokens" }, "N", + "wait up to N tokens for a paragraph break before forcing the exhausted reasoning budget " + "(default: 0, disabled)", + [](common_params & params, int value) { params.sampling.reasoning_budget_grace_tokens = value; }) + .set_examples({ LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_COMPLETION, LLAMA_EXAMPLE_CLI }) + .set_env("LLAMA_ARG_THINK_BUDGET_GRACE_TOKENS")); add_opt(common_arg( {"--reasoning-preserve"}, {"--no-reasoning-preserve"}, @@ -4133,6 +4264,21 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.speculative.draft.n_max = value; } ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_LOOKUP, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_N_MAX")); + add_opt(common_arg( + {"--spec-draft-adaptive"}, + string_format("size each draft from measured acceptance rather than always drafting --spec-draft-n-max (default: %s)", params.speculative.draft.adaptive ? "true" : "false"), + [](common_params & params) { + params.speculative.draft.adaptive = true; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_LOOKUP, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_ADAPTIVE")); + add_opt(common_arg( + {"--spec-draft-mtp-vocab"}, "N", + string_format("MTP: compute draft logits over token ids < N plus control tokens only; verification is unchanged " + "(0 = full vocabulary, default: %d)", params.speculative.draft.mtp_vocab), + [](common_params & params, int value) { + params.speculative.draft.mtp_vocab = std::max(0, value); + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); add_opt(common_arg( {"--spec-draft-n-min"}, "N", string_format("minimum number of draft tokens to use for speculative decoding (default: %d)", params.speculative.draft.n_min), @@ -4232,6 +4378,111 @@ common_params_context common_params_parser_init(common_params & params, llama_ex params.speculative.draft.mparams.hf_file = value; // will be used if --spec-draft-hf is set } ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_DRAFT_MODEL")); + add_opt(common_arg( + {"--spec-prefill", "--speculative-prefill"}, + "enable speculative prefill using draft model to filter prompt tokens", + [](common_params & params) { + params.speculative.prefill.enabled = true; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_PREFILL")); + add_opt(common_arg( + {"--spec-prefill-draft-model", "--spec-prefill-model", "-mpd", "--speculative-prefill-model", "--speculative-prefill-draft-model"}, "FNAME", + "draft model for speculative prefill (default: unused)", + [](common_params & params, const std::string & value) { + params.speculative.prefill.model.path = value; + params.speculative.prefill.model.hf_file = value; // will be used if --spec-prefill-hf is set + params.speculative.prefill.enabled = true; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_PREFILL_MODEL")); + add_opt(common_arg( + {"--spec-prefill-draft-hf", "--spec-prefill-hf", "-hfpd", "--speculative-prefill-hf", "--speculative-prefill-draft-hf"}, "/[:quant]", + "Hugging Face model repository for speculative prefill draft model (default: unused)", + [](common_params & params, const std::string & value) { + params.speculative.prefill.model.hf_repo = value; + params.speculative.prefill.enabled = true; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_PREFILL_HF")); + add_opt(common_arg( + {"--spec-prefill-draft-ngl", "--spec-prefill-ngl", "-nglpd", "--speculative-prefill-ngl", "--speculative-prefill-draft-ngl"}, "N", + string_format("max. number of speculative prefill draft model layers to store in VRAM, either an exact number, 'auto', or 'all' (default: %s)", + params.speculative.prefill.n_gpu_layers == -1 ? "auto" : "all"), + [](common_params & params, const std::string & value) { + if (value == "auto") { + params.speculative.prefill.n_gpu_layers = -1; + } else if (value == "all") { + params.speculative.prefill.n_gpu_layers = -2; + } else { + params.speculative.prefill.n_gpu_layers = std::stoi(value); + } + if (!llama_supports_gpu_offload()) { + fprintf(stderr, "warning: no usable GPU found, --spec-prefill-ngl option will be ignored\n"); + fprintf(stderr, "warning: one possible reason is that llama.cpp was compiled without GPU support\n"); + fprintf(stderr, "warning: consult docs/build.md for compilation instructions\n"); + } + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_PREFILL_N_GPU_LAYERS")); + add_opt(common_arg( + {"--spec-prefill-draft-device", "--spec-prefill-device", "-devpd", "--speculative-prefill-device", "--speculative-prefill-draft-device"}, "", + "comma-separated list of devices to use for offloading the speculative prefill draft model (none = don't offload)\n" + "use --list-devices to see a list of available devices", + [](common_params & params, const std::string & value) { + params.speculative.prefill.devices = parse_device_list(value); + params.speculative.prefill.enabled = true; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_PREFILL_DEVICE")); + add_opt(common_arg( + {"--spec-prefill-draft-ctx", "--spec-prefill-ctx", "--spec-prefill-ctx-size", "--spec-prefill-max-ctx", "-cpd", "--speculative-prefill-ctx", "--speculative-prefill-max-ctx"}, "N", + string_format("context size for speculative prefill draft model (default: %d, 0 = draft training limit or main context)", params.speculative.prefill.n_ctx), + [](common_params & params, int value) { + if (value < 0) { + throw std::invalid_argument("spec-prefill context size must be >= 0"); + } + params.speculative.prefill.n_ctx = value; + params.speculative.prefill.enabled = true; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_PREFILL_CTX_SIZE")); + add_opt(common_arg( + {"--spec-prefill-p", "--spec-prefill-percentage"}, "P", + string_format("fraction of prompt tokens to retain during speculative prefill (default: %.2f)", (double) params.speculative.prefill.percentage), + [](common_params & params, const std::string & value) { + const float val = std::stof(value); + if (val <= 0.0f || val > 1.0f) { + throw std::invalid_argument("spec-prefill percentage must be between 0.0 (exclusive) and 1.0 (inclusive)"); + } + params.speculative.prefill.enabled = true; + params.speculative.prefill.percentage = val; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI}).set_env("LLAMA_ARG_SPEC_PREFILL_P")); + add_opt(common_arg( + {"--spec-prefill-chunk", "--spec-prefill-chunk-size"}, "N", + string_format("chunk grouping size for speculative prefill (default: %d, 0 to disable)", params.speculative.prefill.chunk_size), + [](common_params & params, int value) { + if (value < 0) { + throw std::invalid_argument("spec-prefill chunk size must be >= 0"); + } + params.speculative.prefill.chunk_size = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--spec-prefill-lookahead", "--spec-prefill-lah"}, "N", + string_format("number of lookahead decode steps on draft model for attention estimation (default: %d)", params.speculative.prefill.look_ahead_cnt), + [](common_params & params, int value) { + if (value < 1) { + throw std::invalid_argument("spec-prefill lookahead count must be >= 1"); + } + params.speculative.prefill.look_ahead_cnt = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); + add_opt(common_arg( + {"--spec-prefill-pool-kernel"}, "N", + string_format("1D average pooling kernel size for attention smoothing (default: %d, 0 to disable)", params.speculative.prefill.pool_kernel_size), + [](common_params & params, int value) { + if (value < 0) { + throw std::invalid_argument("spec-prefill pool kernel size must be >= 0"); + } + params.speculative.prefill.pool_kernel_size = value; + } + ).set_spec().set_examples({LLAMA_EXAMPLE_SPECULATIVE, LLAMA_EXAMPLE_SERVER, LLAMA_EXAMPLE_CLI})); add_opt(common_arg( {"--spec-type"}, common_speculative_all_types_str(), string_format("comma-separated list of types of speculative decoding to use (default: %s)\n", @@ -4746,6 +4997,12 @@ void common_params_add_preset_options(std::vector & args) { [](common_params &, const std::string &) { /* unused */ } ).set_env(COMMON_ARG_PRESET_DEDUP_CACHE_MODELS).set_preset_only()); + args.push_back(common_arg( + {"hidden"}, "0|1", + "in server router mode, omit this model from GET /models (it can still be loaded by name)", + [](common_params &, const std::string &) { /* unused */ } + ).set_env(COMMON_ARG_PRESET_HIDDEN).set_preset_only()); + // args.push_back(common_arg( // {"pin"}, // "in server router mode, do not unload this model if models_max is exceeded", diff --git a/common/arg.h b/common/arg.h index 421bc295fc21..e86495a86971 100644 --- a/common/arg.h +++ b/common/arg.h @@ -14,6 +14,7 @@ #define COMMON_ARG_PRESET_LOAD_ON_STARTUP "__PRESET_LOAD_ON_STARTUP" #define COMMON_ARG_PRESET_STOP_TIMEOUT "__PRESET_STOP_TIMEOUT" #define COMMON_ARG_PRESET_DEDUP_CACHE_MODELS "__PRESET_DEDUP_CACHE_MODELS" +#define COMMON_ARG_PRESET_HIDDEN "__PRESET_HIDDEN" // // CLI argument parsing @@ -138,6 +139,7 @@ void common_params_add_preset_options(std::vector & args); struct common_models_handler { common_download_hf_plan plan; common_download_hf_plan plan_spec; + common_download_hf_plan plan_prefill; common_download_opts opts; }; diff --git a/common/common.cpp b/common/common.cpp index d8319cd9ac9a..9ca280d17bd1 100644 --- a/common/common.cpp +++ b/common/common.cpp @@ -1306,7 +1306,8 @@ common_init_result::common_init_result(common_params & params, bool model_only) auto mparams_dft = common_model_params_to_llama(params_dft); auto cparams_dft = common_context_params_to_llama(params_dft); if (spec_mtp) { - cparams_dft.ctx_type = LLAMA_CONTEXT_TYPE_MTP; + cparams_dft.ctx_type = LLAMA_CONTEXT_TYPE_MTP; + cparams_dft.mtp_draft_vocab = params.speculative.draft.mtp_vocab; } cparams_dft.n_rs_seq = 0; @@ -1692,7 +1693,6 @@ struct llama_model_params common_model_params_to_llama(common_params & params) { mparams.check_tensors = params.check_tensors; mparams.use_extra_bufts = !params.no_extra_bufts; mparams.no_host = params.no_host; - if (params.kv_overrides.empty()) { mparams.kv_overrides = NULL; } else { diff --git a/common/common.h b/common/common.h index 63d0badd0f74..6410b373b3b5 100644 --- a/common/common.h +++ b/common/common.h @@ -286,12 +286,31 @@ struct common_params_sampling { // reasoning budget sampler parameters // these are populated by the server/CLI based on chat template params + bool reasoning_budget_enabled = + false; // Master switch for the budget, soft warnings, intro, grace, and reasoning control. int32_t reasoning_budget_tokens = -1; // -1 = disabled, >= 0 = token budget - std::vector reasoning_budget_start; // start tag token sequence - std::vector reasoning_budget_end; // end tag token sequences; the first tag is used as the forcing sequence - std::vector reasoning_budget_forced; // forced sequence (message + first end tag) - std::string reasoning_budget_message; // message injected before end tag when budget exhausted - bool reasoning_control = false; // create the budget sampler on demand so reasoning can be ended at runtime + std::vector reasoning_budget_start; // Start tag token sequence. + std::vector reasoning_budget_end; // End tag token sequences. First is forced. + std::vector reasoning_budget_forced; // Forced sequence: message and first end tag. + std::string reasoning_budget_message; // Message injected before the forced end tag. + bool reasoning_control = false; // Create sampler for runtime reasoning control. + + float reasoning_budget_soft_ratio = -1.0f; // <= 0 disables the first soft warning. + std::vector reasoning_budget_soft_forced; + std::string reasoning_budget_soft_message; + + float reasoning_budget_soft2_ratio = -1.0f; + std::vector reasoning_budget_soft2_forced; + std::string reasoning_budget_soft2_message; + + std::vector reasoning_budget_intro_forced; + std::string reasoning_budget_intro_message = + "I'll keep this reasoning under {budget} tokens, so I'll stay focused and efficient. "; + std::string reasoning_budget_intro_mode = + "every"; + + int32_t reasoning_budget_grace_tokens = + 0; // <= 0 forces immediately. Positive values wait for a paragraph break. bool backend_sampling = false; @@ -341,6 +360,12 @@ struct common_params_speculative_draft { ggml_type cache_type_k = GGML_TYPE_F16; // KV cache data type for the K ggml_type cache_type_v = GGML_TYPE_F16; // KV cache data type for the V + // size each draft from measured acceptance instead of always drafting n_max + bool adaptive = false; + + // MTP only: draft over token ids < N plus control tokens (0 = full vocabulary) + int32_t mtp_vocab = 0; + common_cpu_params cpuparams; common_cpu_params cpuparams_batch; @@ -367,6 +392,20 @@ struct common_params_speculative_ngram_cache { std::string lookup_cache_dynamic; // path of dynamic ngram cache file for lookup decoding }; +struct common_params_speculative_prefill { + bool enabled = false; // enable speculative prefill + common_params_model model; // draft model for speculative prefill + int32_t n_ctx = 0; // context size for draft model (0 = default/target context size) + int32_t n_gpu_layers = -1; // max draft model layers to store in VRAM (-1 - use default) + std::vector devices; // devices to use for offloading the draft model + float percentage = 0.3f; // fraction of prompt tokens to retain (0.0 < p <= 1.0) + int32_t chunk_size = 32; // chunk grouping size (0 to disable chunking) + int32_t look_ahead_cnt = 8; // lookahead decode steps on draft model + int32_t pool_kernel_size = 13; // 1D average pooling kernel size for smoothing + bool keep_bos = true; // preserve first token (BOS) + bool keep_last = true; // preserve last token / tail chunk +}; + struct common_params_speculative { std::vector types = { COMMON_SPECULATIVE_TYPE_NONE }; @@ -383,6 +422,8 @@ struct common_params_speculative { common_params_speculative_ngram_cache ngram_cache; + common_params_speculative_prefill prefill; + bool has_dft() const { return !draft.mparams.empty(); } @@ -391,6 +432,9 @@ struct common_params_speculative { return synth_len != -1.0 || !synth_rates.empty(); } + bool has_prefill() const { + return prefill.enabled && !prefill.model.empty(); + } uint32_t need_n_rs_seq() const { bool needs_rs_seq = std::any_of(types.begin(), types.end(), [&](auto t) { return t == COMMON_SPECULATIVE_TYPE_DRAFT_MTP || t == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 || t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || t == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK; @@ -581,7 +625,6 @@ struct common_params { bool no_op_offload = false; // globally disable offload host tensor operations to device bool no_extra_bufts = false; // disable extra buffer types (used for weight repacking) bool no_host = false; // bypass host buffer allowing extra buffers to be used - bool single_turn = false; // single turn chat conversation ggml_type cache_type_k = GGML_TYPE_F16; // KV cache data type for the K diff --git a/common/parsers/qwen3-coder.cpp b/common/parsers/qwen3-coder.cpp index 7938a2027932..89dbb9800f2d 100644 --- a/common/parsers/qwen3-coder.cpp +++ b/common/parsers/qwen3-coder.cpp @@ -133,6 +133,22 @@ common_chat_params common_chat_params_init_qwen3_coder(const common_chat_templat arg_value = p.gbnf(p.atomic(p.tool_arg_json_value(json_value) + arg_close) | arg_string, "xml-arg-string"); } + if (param.schema->kind() == common_chat_schema::KIND_ENUM) { + const auto & enum_values = static_cast(*param.schema).values; + auto values = p.choice(); + bool all_strings = true; + for (const auto & value : enum_values) { + if (!value.is_string()) { + all_strings = false; + break; + } + values |= p.literal(value.get()) + p.peek(p.literal("\n\n")); + } + if (all_strings) { + arg_value = p.tool_arg_string_value(values) + arg_close; + } + } + auto arg_rule = p.rule(rule_name, p.tool_arg(arg_open + arg_value)); (param.required ? required_args : optional_args).push_back(arg_rule); diff --git a/common/reasoning-budget.cpp b/common/reasoning-budget.cpp index 4884299f3011..538839a468b2 100644 --- a/common/reasoning-budget.cpp +++ b/common/reasoning-budget.cpp @@ -36,7 +36,7 @@ struct token_matcher { return t; } - // returns the index into seqs of the longest sequence ending at this token, or -1 + // Returns the index of the longest sequence ending at this token, or -1. int32_t advance(llama_token token) { state = ac.next(state, (uint32_t) token); const int32_t p = ac.match_pattern(state); @@ -56,21 +56,91 @@ struct common_reasoning_budget_ctx { token_matcher end_matcher; llama_tokens forced_tokens; - int32_t budget; // maximum tokens in reasoning block - int32_t remaining; // tokens remaining in budget + int32_t budget; // Maximum tokens in reasoning block. + int32_t remaining; // Tokens remaining in budget. common_reasoning_budget_state state; - // for forcing - size_t force_pos; // next position in forced_tokens to force + size_t force_pos; // Next position in forced_tokens to force. - int32_t end_match; // index into end_matcher.seqs of the sequence that transitioned to DONE, -1 if none + int32_t end_match; // Index into end_matcher.seqs that transitioned to DONE, or -1. + + struct soft_point_state { + int32_t threshold; + llama_tokens tokens; + bool triggered; + size_t force_pos; + }; + + std::vector soft_points; + size_t active_soft; + + llama_tokens intro_forced_tokens; + size_t intro_force_pos; + bool intro_once; + + int32_t grace_tokens; + int32_t grace_remaining; + bool hard_pending_prev_nl; }; static const char * common_reasoning_budget_name(const struct llama_sampler * /*smpl*/) { return "reasoning-budget"; } +static bool token_utf8_complete(const common_reasoning_budget_ctx * ctx, llama_token token) { + if (ctx->vocab == nullptr) { + return true; + } + const std::string piece = common_token_to_piece(ctx->vocab, token, false); + return common_utf8_is_complete(piece); +} + +static void common_reasoning_budget_begin_forcing(common_reasoning_budget_ctx * ctx, llama_token token) { + ctx->end_matcher.reset(); + if (token_utf8_complete(ctx, token)) { + ctx->state = REASONING_BUDGET_FORCING; + ctx->force_pos = 0; + } else { + ctx->state = REASONING_BUDGET_WAITING_UTF8; + } +} + +static void common_reasoning_budget_enter_hard_exhausted(common_reasoning_budget_ctx * ctx, llama_token token) { + if (ctx->grace_tokens > 0) { + ctx->state = REASONING_BUDGET_HARD_PENDING; + ctx->grace_remaining = ctx->grace_tokens; + ctx->hard_pending_prev_nl = false; + ctx->end_matcher.reset(); + COM_TRC("budget exhausted, waiting up to %d tokens for a paragraph break\n", ctx->grace_tokens); + return; + } + + common_reasoning_budget_begin_forcing(ctx, token); + COM_TRC("%s", "budget exhausted, forcing end sequence\n"); +} + +static void common_reasoning_budget_activate(common_reasoning_budget_ctx * ctx) { + ctx->remaining = ctx->budget; + for (auto & sp : ctx->soft_points) { + sp.triggered = false; + } + ctx->end_match = -1; + + if (!ctx->intro_forced_tokens.empty()) { + ctx->state = REASONING_BUDGET_INTRO_FORCING; + ctx->intro_force_pos = 0; + COM_TRC("activated, budget=%d tokens, forcing intro message\n", ctx->budget); + } else if (ctx->remaining <= 0) { + ctx->state = REASONING_BUDGET_FORCING; + ctx->force_pos = 0; + COM_TRC("%s", "budget=0, forcing immediately\n"); + } else { + ctx->state = REASONING_BUDGET_COUNTING; + COM_TRC("activated, budget=%d tokens\n", ctx->budget); + } +} + static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_token token) { auto * ctx = (common_reasoning_budget_ctx *) smpl->ctx; @@ -78,20 +148,29 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to case REASONING_BUDGET_IDLE: { if (ctx->start_matcher.advance(token) >= 0) { - ctx->state = REASONING_BUDGET_COUNTING; - ctx->remaining = ctx->budget; - COM_TRC("activated, budget=%d tokens\n", ctx->budget); - + common_reasoning_budget_activate(ctx); + } + break; + } + case REASONING_BUDGET_INTRO_FORCING: + ctx->intro_force_pos++; + if (ctx->intro_force_pos >= ctx->intro_forced_tokens.size()) { + if (ctx->intro_once) { + // intro-mode once: do not re-fire on later reasoning blocks + ctx->intro_forced_tokens.clear(); + ctx->intro_force_pos = 0; + } if (ctx->remaining <= 0) { ctx->state = REASONING_BUDGET_FORCING; ctx->force_pos = 0; - COM_TRC("%s", "budget=0, forcing immediately\n"); + COM_TRC("%s", "intro complete, budget=0, forcing immediately\n"); + } else { + ctx->state = REASONING_BUDGET_COUNTING; + COM_TRC("%s", "intro complete, resuming countdown\n"); } } break; - } case REASONING_BUDGET_COUNTING: - case REASONING_BUDGET_WAITING_UTF8: { const int32_t match = ctx->end_matcher.advance(token); if (match >= 0) { @@ -101,39 +180,109 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to break; } - bool utf8_complete = true; - if (ctx->vocab != nullptr) { - const std::string piece = common_token_to_piece(ctx->vocab, token, false); - utf8_complete = common_utf8_is_complete(piece); + ctx->remaining--; + if (ctx->remaining <= 0) { + common_reasoning_budget_enter_hard_exhausted(ctx, token); + break; } - if (ctx->state == REASONING_BUDGET_WAITING_UTF8) { - if (utf8_complete) { - ctx->state = REASONING_BUDGET_FORCING; - ctx->force_pos = 0; - ctx->end_matcher.reset(); - COM_TRC("%s", "UTF-8 complete, now forcing end sequence\n"); + for (size_t i = 0; i < ctx->soft_points.size(); i++) { + auto & sp = ctx->soft_points[i]; + if (!sp.triggered && ctx->remaining <= sp.threshold) { + ctx->active_soft = i; + ctx->state = REASONING_BUDGET_SOFT_PENDING; + COM_TRC("soft threshold reached (remaining=%d, point %zu), waiting for newline\n", ctx->remaining, + i); + break; } - } else if (ctx->state == REASONING_BUDGET_COUNTING) { - ctx->remaining--; - if (ctx->remaining <= 0) { - if (utf8_complete) { - ctx->state = REASONING_BUDGET_FORCING; - ctx->force_pos = 0; - ctx->end_matcher.reset(); - COM_TRC("%s", "budget exhausted, forcing end sequence\n"); - } else { - ctx->state = REASONING_BUDGET_WAITING_UTF8; - ctx->end_matcher.reset(); - COM_TRC("%s", "budget exhausted, waiting for UTF-8 completion\n"); - } + } + break; + } + case REASONING_BUDGET_SOFT_PENDING: + { + const int32_t match = ctx->end_matcher.advance(token); + if (match >= 0) { + ctx->state = REASONING_BUDGET_DONE; + ctx->end_match = match; + COM_TRC("%s", "deactivated (natural end)\n"); + break; + } + + ctx->remaining--; + if (ctx->remaining <= 0) { + COM_TRC("%s", "budget exhausted before newline, soft warning skipped\n"); + common_reasoning_budget_enter_hard_exhausted(ctx, token); + break; + } + + if (ctx->vocab != nullptr) { + const std::string piece = common_token_to_piece(ctx->vocab, token, false); + if (piece.find('\n') != std::string::npos) { + ctx->state = REASONING_BUDGET_SOFT_FORCING; + auto & sp = ctx->soft_points[ctx->active_soft]; + sp.force_pos = 0; + sp.triggered = true; + COM_TRC("%s", "newline boundary found, forcing soft warning\n"); } } break; } + case REASONING_BUDGET_SOFT_FORCING: + { + auto & sp = ctx->soft_points[ctx->active_soft]; + sp.force_pos++; + if (sp.force_pos >= sp.tokens.size()) { + ctx->state = REASONING_BUDGET_COUNTING; + COM_TRC("%s", "soft warning complete, resuming countdown\n"); + } + break; + } + case REASONING_BUDGET_HARD_PENDING: + { + const int32_t match = ctx->end_matcher.advance(token); + if (match >= 0) { + ctx->state = REASONING_BUDGET_DONE; + ctx->end_match = match; + COM_TRC("%s", "deactivated (natural end)\n"); + break; + } + + ctx->grace_remaining--; + + const std::string piece = + ctx->vocab != nullptr ? common_token_to_piece(ctx->vocab, token, false) : std::string(); + const bool paragraph_boundary = piece.find("\n\n") != std::string::npos || + (ctx->hard_pending_prev_nl && !piece.empty() && piece[0] == '\n'); + ctx->hard_pending_prev_nl = !piece.empty() && piece.back() == '\n'; + + if (paragraph_boundary) { + common_reasoning_budget_begin_forcing(ctx, token); + COM_TRC("%s", "paragraph boundary found, forcing end sequence\n"); + } else if (ctx->grace_remaining <= 0) { + common_reasoning_budget_begin_forcing(ctx, token); + COM_TRC("%s", "grace period expired, forcing end sequence\n"); + } + break; + } + case REASONING_BUDGET_WAITING_UTF8: + { + const int32_t match = ctx->end_matcher.advance(token); + if (match >= 0) { + ctx->state = REASONING_BUDGET_DONE; + ctx->end_match = match; + COM_TRC("%s", "deactivated (natural end)\n"); + break; + } + + if (token_utf8_complete(ctx, token)) { + common_reasoning_budget_begin_forcing(ctx, token); + COM_TRC("%s", "UTF-8 complete, now forcing end sequence\n"); + } + break; + } case REASONING_BUDGET_FORCING: { - // track the end sequence within forced_tokens so it is also reported on DONE + // Track the end sequence in forced_tokens for grammar replay. const int32_t match = ctx->end_matcher.advance(token); ctx->force_pos++; if (ctx->force_pos >= ctx->forced_tokens.size()) { @@ -144,20 +293,9 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to break; } case REASONING_BUDGET_DONE: - // Re-arm on a new start tag: some models emit multiple blocks - // per response, and each should get a fresh budget window. if (ctx->start_matcher.advance(token) >= 0) { - ctx->state = REASONING_BUDGET_COUNTING; - ctx->remaining = ctx->budget; ctx->end_matcher.reset(); - ctx->end_match = -1; - COM_TRC("re-activated on new start tag, budget=%d tokens\n", ctx->budget); - - if (ctx->remaining <= 0) { - ctx->state = REASONING_BUDGET_FORCING; - ctx->force_pos = 0; - COM_TRC("%s", "budget=0, forcing immediately\n"); - } + common_reasoning_budget_activate(ctx); } break; } @@ -166,18 +304,28 @@ static void common_reasoning_budget_accept(struct llama_sampler * smpl, llama_to static void common_reasoning_budget_apply(struct llama_sampler * smpl, llama_token_data_array * cur_p) { auto * ctx = (common_reasoning_budget_ctx *) smpl->ctx; - if (ctx->state != REASONING_BUDGET_FORCING) { - // passthrough — don't modify logits - return; - } + llama_token forced; - if (ctx->force_pos >= ctx->forced_tokens.size()) { + if (ctx->state == REASONING_BUDGET_FORCING) { + if (ctx->force_pos >= ctx->forced_tokens.size()) { + return; + } + forced = ctx->forced_tokens[ctx->force_pos]; + } else if (ctx->state == REASONING_BUDGET_SOFT_FORCING) { + const auto & sp = ctx->soft_points[ctx->active_soft]; + if (sp.force_pos >= sp.tokens.size()) { + return; + } + forced = sp.tokens[sp.force_pos]; + } else if (ctx->state == REASONING_BUDGET_INTRO_FORCING) { + if (ctx->intro_force_pos >= ctx->intro_forced_tokens.size()) { + return; + } + forced = ctx->intro_forced_tokens[ctx->intro_force_pos]; + } else { return; } - const llama_token forced = ctx->forced_tokens[ctx->force_pos]; - - // set all logits to -inf except the forced token for (size_t i = 0; i < cur_p->size; i++) { if (cur_p->data[i].id != forced) { cur_p->data[i].logit = -INFINITY; @@ -193,12 +341,27 @@ static void common_reasoning_budget_reset(struct llama_sampler * smpl) { ctx->end_matcher.reset(); ctx->force_pos = 0; ctx->end_match = -1; + for (auto & sp : ctx->soft_points) { + sp.triggered = false; + sp.force_pos = 0; + } + ctx->active_soft = 0; + ctx->intro_force_pos = 0; + ctx->grace_remaining = ctx->grace_tokens; + ctx->hard_pending_prev_nl = false; } static struct llama_sampler * common_reasoning_budget_init_state( - const struct llama_vocab * vocab, const std::vector & start_seqs, - const std::vector & end_seqs, const llama_tokens & forced_tokens, - int32_t budget, common_reasoning_budget_state initial_state); + const struct llama_vocab * vocab, + const std::vector & start_seqs, + const std::vector & end_seqs, + const llama_tokens & forced_tokens, + const std::vector & soft_points, + const llama_tokens & intro_forced_tokens, + int32_t budget, + int32_t grace_tokens, + common_reasoning_budget_state initial_state, + bool intro_once); static struct llama_sampler * common_reasoning_budget_clone(const struct llama_sampler * smpl); @@ -231,41 +394,66 @@ static struct llama_sampler * common_reasoning_budget_clone(const struct llama_s } static struct llama_sampler * common_reasoning_budget_init_state( - const struct llama_vocab * vocab, - const std::vector & start_seqs, - const std::vector & end_seqs, - const llama_tokens & forced_tokens, - int32_t budget, - common_reasoning_budget_state initial_state) { - // promote COUNTING with budget <= 0 to FORCING + const struct llama_vocab * vocab, + const std::vector & start_seqs, + const std::vector & end_seqs, + const llama_tokens & forced_tokens, + const std::vector & soft_points, + const llama_tokens & intro_forced_tokens, + int32_t budget, + int32_t grace_tokens, + common_reasoning_budget_state initial_state, + bool intro_once) { if (initial_state == REASONING_BUDGET_COUNTING && budget <= 0) { initial_state = REASONING_BUDGET_FORCING; } + std::vector soft_pts; + for (const auto & p : soft_points) { + if (!p.tokens.empty() && p.threshold >= 0) { + soft_pts.push_back({ p.threshold, p.tokens, false, 0 }); + } + } + std::sort(soft_pts.begin(), soft_pts.end(), + [](const common_reasoning_budget_ctx::soft_point_state & a, + const common_reasoning_budget_ctx::soft_point_state & b) { return a.threshold > b.threshold; }); + return llama_sampler_init( /* .iface = */ &common_reasoning_budget_i, - /* .ctx = */ new common_reasoning_budget_ctx { - /* .vocab = */ vocab, - /* .start_matcher = */ token_matcher(start_seqs), - /* .end_matcher = */ token_matcher(end_seqs), - /* .forced_tokens = */ forced_tokens, - /* .budget = */ budget, - /* .remaining = */ budget, - /* .state = */ initial_state, - /* .force_pos = */ 0, - /* .end_match = */ -1, - } - ); + /* .ctx = */ new common_reasoning_budget_ctx{ + /* .vocab = */ vocab, + /* .start_matcher = */ token_matcher(start_seqs), + /* .end_matcher = */ token_matcher(end_seqs), + /* .forced_tokens = */ forced_tokens, + /* .budget = */ budget, + /* .remaining = */ budget, + /* .state = */ initial_state, + /* .force_pos = */ 0, + /* .end_match = */ -1, + /* .soft_points = */ std::move(soft_pts), + /* .active_soft = */ 0, + /* .intro_forced_tokens = */ intro_forced_tokens, + /* .intro_force_pos = */ 0, + /* .intro_once = */ intro_once, + /* .grace_tokens = */ grace_tokens, + /* .grace_remaining = */ grace_tokens, + /* .hard_pending_prev_nl = */ false, + }); } struct llama_sampler * common_reasoning_budget_init( - const struct llama_vocab * vocab, - const std::vector & start_seqs, - const std::vector & end_seqs, - const llama_tokens & forced_tokens, - int32_t budget, - common_reasoning_budget_state initial_state) { - return common_reasoning_budget_init_state(vocab, start_seqs, end_seqs, forced_tokens, budget, initial_state); + const struct llama_vocab * vocab, + const std::vector & start_seqs, + const std::vector & end_seqs, + const llama_tokens & forced_tokens, + const std::vector & soft_points, + const llama_tokens & intro_forced_tokens, + int32_t budget, + int32_t grace_tokens, + common_reasoning_budget_state initial_state, + bool intro_once) { + return common_reasoning_budget_init_state(vocab, start_seqs, end_seqs, forced_tokens, soft_points, + intro_forced_tokens, budget, grace_tokens, initial_state, intro_once); } common_reasoning_budget_state common_reasoning_budget_get_state(const struct llama_sampler * smpl) { @@ -295,9 +483,9 @@ bool common_reasoning_budget_force(struct llama_sampler * smpl) { auto * ctx = (common_reasoning_budget_ctx *) smpl->ctx; - // only a sampler that is actively counting down the budget may be forced; - // any other state (idle, already forcing/waiting, or done) is left untouched - if (ctx->state != REASONING_BUDGET_COUNTING) { + if (ctx->state != REASONING_BUDGET_COUNTING && ctx->state != REASONING_BUDGET_INTRO_FORCING && + ctx->state != REASONING_BUDGET_SOFT_PENDING && ctx->state != REASONING_BUDGET_SOFT_FORCING && + ctx->state != REASONING_BUDGET_HARD_PENDING) { return false; } diff --git a/common/reasoning-budget.h b/common/reasoning-budget.h index 1b89a04c42e8..e3ababe0a8d6 100644 --- a/common/reasoning-budget.h +++ b/common/reasoning-budget.h @@ -8,38 +8,56 @@ #include enum common_reasoning_budget_state { - REASONING_BUDGET_IDLE, // waiting for start sequence - REASONING_BUDGET_COUNTING, // counting down tokens - REASONING_BUDGET_FORCING, // forcing budget message + end sequence - REASONING_BUDGET_WAITING_UTF8, // budget exhausted, waiting for UTF-8 completion - REASONING_BUDGET_DONE, // passthrough forever + REASONING_BUDGET_IDLE, // Waiting for start sequence. + REASONING_BUDGET_INTRO_FORCING, // Forcing the intro message. + REASONING_BUDGET_COUNTING, // Counting down tokens. + REASONING_BUDGET_SOFT_PENDING, // Waiting for a soft-warning newline boundary. + REASONING_BUDGET_SOFT_FORCING, // Forcing the soft warning message. + REASONING_BUDGET_HARD_PENDING, // Waiting for a hard-cutoff paragraph boundary. + REASONING_BUDGET_FORCING, // Forcing budget message and end sequence. + REASONING_BUDGET_WAITING_UTF8, // Waiting for UTF-8 completion before forcing. + REASONING_BUDGET_DONE, // Passthrough forever. }; -// Creates a reasoning budget sampler that limits token generation inside a -// reasoning block (e.g. between and ). +struct common_reasoning_budget_soft_point { + int32_t threshold; // Fire when remaining <= threshold. + llama_tokens tokens; // Message forced at this point. Empty disables it. +}; + +// Creates a sampler that limits generation inside a reasoning block. +// Intro, soft, and forced tokens do not count against the budget. // -// State machine: IDLE -> COUNTING -> WAITING_UTF8 -> FORCING -> DONE -// IDLE: passthrough, watching for a start sequence -// COUNTING: counting down remaining tokens, watching for a natural end sequence -// WAITING_UTF8: budget exhausted, allowing tokens to complete a UTF-8 sequence -// FORCING: forces forced_tokens token-by-token (all other logits -> -inf) -// DONE: passthrough forever +// State machine: +// IDLE -> INTRO_FORCING -> COUNTING -> SOFT_PENDING -> SOFT_FORCING -> +// COUNTING -> HARD_PENDING -> WAITING_UTF8 -> FORCING -> DONE // -// Parameters: -// vocab - vocabulary (used for UTF-8 boundary detection; can be nullptr) -// start_seqs - token sequences, any of which activates counting -// end_seqs - token sequences, any of which naturally deactivates -// forced_tokens - token sequence forced when budget expires -// budget - max tokens allowed in the reasoning block -// initial_state - initial state +// A soft point fires once at the next newline after its threshold is crossed. +// The hard cutoff takes priority when the budget is exhausted. // +// A positive grace_tokens value waits for a paragraph boundary after exhaustion. +// It forces after at most grace_tokens more tokens. +// +// Parameters: +// vocab - vocabulary for UTF-8 and paragraph boundary detection +// start_seqs - sequences that activate counting +// end_seqs - sequences that naturally deactivate +// forced_tokens - sequence forced when the budget expires +// soft_points - soft warning points +// intro_forced_tokens - sequence forced when the block starts +// budget - max generated tokens in the reasoning block +// grace_tokens - max tokens to wait for a paragraph boundary +// initial_state - initial state struct llama_sampler * common_reasoning_budget_init( - const struct llama_vocab * vocab, - const std::vector & start_seqs, - const std::vector & end_seqs, - const llama_tokens & forced_tokens, - int32_t budget, - common_reasoning_budget_state initial_state = REASONING_BUDGET_IDLE); + const struct llama_vocab * vocab, + const std::vector & start_seqs, + const std::vector & end_seqs, + const llama_tokens & forced_tokens, + const std::vector & soft_points, + const llama_tokens & intro_forced_tokens, + int32_t budget, + int32_t grace_tokens = 0, + common_reasoning_budget_state initial_state = REASONING_BUDGET_IDLE, + bool intro_once = false); common_reasoning_budget_state common_reasoning_budget_get_state(const struct llama_sampler * smpl); diff --git a/common/sampling.cpp b/common/sampling.cpp index 06dea1e1ccea..0e40d86e893c 100644 --- a/common/sampling.cpp +++ b/common/sampling.cpp @@ -307,18 +307,40 @@ struct common_sampler * common_sampler_init( } } - // reasoning budget sampler (skip when budget is unlimited unless a lazy grammar is active, which needs rbudget for thinking-block suppression) - if (!params.reasoning_budget_start.empty() && !params.reasoning_budget_end.empty() && (params.grammar_lazy || params.reasoning_budget_tokens >= 0 || params.reasoning_control)) { - rbudget = common_reasoning_budget_init( - vocab, - {params.reasoning_budget_start}, - params.reasoning_budget_end, - params.reasoning_budget_forced, - params.reasoning_budget_tokens < 0 ? INT_MAX : params.reasoning_budget_tokens); + // Grammar lazy mode needs the sampler to suppress grammar in a reasoning block. + // The master switch gates budget forcing and runtime reasoning control only. + const bool reasoning_budget_active = + params.reasoning_budget_enabled && (params.reasoning_budget_tokens >= 0 || params.reasoning_control); + if (!params.reasoning_budget_start.empty() && !params.reasoning_budget_end.empty() && + (params.grammar_lazy || reasoning_budget_active)) { + const int32_t budget_tokens = params.reasoning_budget_tokens < 0 ? INT_MAX : params.reasoning_budget_tokens; + + std::vector soft_points; + auto add_soft_point = [&](float ratio, const llama_tokens & tokens) { + if (ratio > 0.0f && !tokens.empty() && budget_tokens != INT_MAX) { + const float r = std::min(ratio, 1.0f); + soft_points.push_back({ std::max(0, budget_tokens - (int32_t) std::ceil(budget_tokens * r)), tokens }); + } + }; + add_soft_point(params.reasoning_budget_soft_ratio, params.reasoning_budget_soft_forced); + add_soft_point(params.reasoning_budget_soft2_ratio, params.reasoning_budget_soft2_forced); + + const llama_tokens intro_tokens = + reasoning_budget_active ? params.reasoning_budget_intro_forced : llama_tokens(); + rbudget = common_reasoning_budget_init(vocab, { params.reasoning_budget_start }, params.reasoning_budget_end, + params.reasoning_budget_forced, soft_points, intro_tokens, budget_tokens, + params.reasoning_budget_grace_tokens, REASONING_BUDGET_IDLE, + params.reasoning_budget_intro_mode == "once"); for (const auto & token : prefill_tokens) { llama_sampler_accept(rbudget, token); LOG_DBG("%s: reasoning-budget accepted prefill token (%d)\n", __func__, token); + + const auto state = common_reasoning_budget_get_state(rbudget); + if (state == REASONING_BUDGET_INTRO_FORCING || state == REASONING_BUDGET_SOFT_FORCING || + state == REASONING_BUDGET_FORCING) { + break; + } } } diff --git a/common/speculative-prefill.cpp b/common/speculative-prefill.cpp new file mode 100644 index 000000000000..ab20ad08f1e5 --- /dev/null +++ b/common/speculative-prefill.cpp @@ -0,0 +1,337 @@ +#include "speculative-prefill.h" + +#include "common.h" +#include "ggml.h" +#include "ggml-backend.h" +#include "llama.h" +#include "log.h" +#include "sampling.h" + +#include +#include +#include +#include +#include + +#define SPF_DBG(fmt, ...) LOG_DBG("spec-prefill: " fmt, __VA_ARGS__) +#define SPF_INF(fmt, ...) LOG_INF("spec-prefill: " fmt, __VA_ARGS__) +#define SPF_ERR(fmt, ...) LOG_ERR("spec-prefill: " fmt, __VA_ARGS__) + +std::vector common_speculative_prefill_compute_importance( + const std::vector> & attn_layers, + int32_t n_prompt, + int32_t n_heads, + int32_t pool_kernel_size) { + if (attn_layers.empty() || n_prompt <= 0 || n_heads <= 0) { + return std::vector(std::max(0, n_prompt), 1.0f); + } + + std::vector max_importance(n_prompt, 0.0f); + std::vector smoothed(n_prompt, 0.0f); + + const int32_t pad = pool_kernel_size > 0 ? (pool_kernel_size / 2) : 0; + + for (const auto & layer_attn : attn_layers) { + if ((int32_t) layer_attn.size() < n_heads * n_prompt) { + continue; + } + + for (int32_t h = 0; h < n_heads; ++h) { + const float * head_attn = layer_attn.data() + (size_t) h * n_prompt; + + if (pool_kernel_size > 1) { + for (int32_t i = 0; i < n_prompt; ++i) { + const int32_t start = std::max(0, i - pad); + const int32_t end = std::min(n_prompt - 1, i + pad); + + float sum = 0.0f; + for (int32_t j = start; j <= end; ++j) { + sum += head_attn[j]; + } + smoothed[i] = sum / (float) (end - start + 1); + } + } else { + std::copy(head_attn, head_attn + n_prompt, smoothed.begin()); + } + + for (int32_t i = 0; i < n_prompt; ++i) { + if (smoothed[i] > max_importance[i]) { + max_importance[i] = smoothed[i]; + } + } + } + } + + return max_importance; +} + +std::vector common_speculative_prefill_select_indices( + const std::vector & importance_scores, + const common_params_speculative_prefill & params) { + const int32_t n_prompt = (int32_t) importance_scores.size(); + if (n_prompt <= 0) { + return {}; + } + + const float percentage = std::max(0.01f, std::min(1.0f, params.percentage)); + std::vector selected; + + if (params.chunk_size > 0 && n_prompt > params.chunk_size) { + const int32_t chunk_size = params.chunk_size; + const int32_t n_chunks = (n_prompt + chunk_size - 1) / chunk_size; + + std::vector> chunk_scores(n_chunks); + + for (int32_t c = 0; c < n_chunks; ++c) { + const int32_t start = c * chunk_size; + const int32_t end = std::min(n_prompt, (c + 1) * chunk_size); + + float sum = 0.0f; + for (int32_t i = start; i < end; ++i) { + sum += importance_scores[i]; + } + chunk_scores[c] = { sum / (float) (end - start), c }; + } + + const int32_t n_keep_chunks = std::max(1, (int32_t) std::ceil(n_chunks * percentage)); + + std::sort(chunk_scores.begin(), chunk_scores.end(), + [](const std::pair & a, const std::pair & b) { + return a.first > b.first; + }); + + for (int32_t k = 0; k < n_keep_chunks && k < n_chunks; ++k) { + const int32_t c = chunk_scores[k].second; + const int32_t start = c * chunk_size; + const int32_t end = std::min(n_prompt, (c + 1) * chunk_size); + + for (int32_t i = start; i < end; ++i) { + selected.push_back(i); + } + } + } else { + const int32_t n_keep = std::max(1, (int32_t) std::ceil(n_prompt * percentage)); + + std::vector> token_scores(n_prompt); + for (int32_t i = 0; i < n_prompt; ++i) { + token_scores[i] = { importance_scores[i], i }; + } + + std::sort(token_scores.begin(), token_scores.end(), + [](const std::pair & a, const std::pair & b) { + return a.first > b.first; + }); + + for (int32_t k = 0; k < n_keep; ++k) { + selected.push_back(token_scores[k].second); + } + } + + if (params.keep_bos) { + selected.push_back(0); + } + + if (params.keep_last) { + if (params.chunk_size > 0 && n_prompt > params.chunk_size) { + const int32_t start = ((n_prompt - 1) / params.chunk_size) * params.chunk_size; + for (int32_t i = start; i < n_prompt; ++i) { + selected.push_back(i); + } + } else { + selected.push_back(n_prompt - 1); + } + } + + std::sort(selected.begin(), selected.end()); + selected.erase(std::unique(selected.begin(), selected.end()), selected.end()); + + return selected; +} + +struct cb_attn_collector_data { + int32_t n_prompt = 0; + std::vector> layer_attns; +}; + +static bool cb_collect_attn(ggml_tensor * t, bool ask, void * user_data) { + if (ask) { + return strncmp(t->name, "kq_soft_max", 11) == 0; + } + + if (strncmp(t->name, "kq_soft_max", 11) != 0) { + return true; + } + + auto * data = (cb_attn_collector_data *) user_data; + if (data == nullptr || data->n_prompt <= 0) { + return true; + } + + if (t->type != GGML_TYPE_F32) { + return true; + } + + const int32_t n_past_total = (int32_t) t->ne[0]; + const int32_t n_heads = (int32_t) t->ne[2]; + const int32_t n_prompt = data->n_prompt; + + if (t->nb[0] != sizeof(float) || n_heads <= 0 || n_past_total < n_prompt) { + return true; + } + + std::vector layer_data((size_t) n_heads * n_prompt); + + for (int32_t h = 0; h < n_heads; ++h) { + float * dst = layer_data.data() + (size_t) h * n_prompt; + ggml_backend_tensor_get(t, dst, (size_t) h * t->nb[2], (size_t) n_prompt * sizeof(float)); + } + + data->layer_attns.push_back(std::move(layer_data)); + + return true; +} + +common_speculative_prefill_result common_speculative_prefill_execute( + llama_context * ctx_dft, + common_sampler * smpl_dft, + const std::vector & prompt, + llama_seq_id seq_id, + const common_params_speculative_prefill & params) { + common_speculative_prefill_result res; + res.n_prompt_orig = (int32_t) prompt.size(); + + if (prompt.empty()) { + return res; + } + + if (!params.enabled || params.percentage >= 1.0f || (int32_t) prompt.size() <= params.chunk_size) { + res.kept_indices.resize(prompt.size()); + std::iota(res.kept_indices.begin(), res.kept_indices.end(), 0); + res.n_prompt_kept = (int32_t) res.kept_indices.size(); + res.importance_scores.assign(prompt.size(), 1.0f); + return res; + } + + const int32_t n_ctx_dft = llama_n_ctx(ctx_dft); + const int32_t lookahead = std::max(1, params.look_ahead_cnt); + + if ((int32_t) prompt.size() + lookahead > n_ctx_dft) { + SPF_INF("prompt size (%d) + lookahead (%d) exceeds draft context (%d); skipping speculative prefill\n", + (int32_t) prompt.size(), lookahead, n_ctx_dft); + res.kept_indices.resize(prompt.size()); + std::iota(res.kept_indices.begin(), res.kept_indices.end(), 0); + res.n_prompt_kept = (int32_t) res.kept_indices.size(); + res.importance_scores.assign(prompt.size(), 1.0f); + return res; + } + + const auto t_start = ggml_time_us(); + + const llama_model * model_dft = llama_get_model(ctx_dft); + const llama_vocab * vocab_dft = llama_model_get_vocab(model_dft); + const int32_t n_heads_dft = llama_model_n_head(model_dft); + + // 1. evaluate full prompt on draft model + { + const int32_t n_batch_dft = llama_n_batch(ctx_dft); + llama_batch batch_prompt = llama_batch_init(std::min((int32_t) prompt.size(), n_batch_dft), 0, 1); + + for (int32_t i = 0; i < (int32_t) prompt.size(); i += n_batch_dft) { + const int32_t n_eval = std::min((int32_t) prompt.size() - i, n_batch_dft); + common_batch_clear(batch_prompt); + + for (int32_t j = 0; j < n_eval; ++j) { + const int32_t idx = i + j; + const bool is_last = (idx == (int32_t) prompt.size() - 1); + common_batch_add(batch_prompt, prompt[idx], (llama_pos) idx, { seq_id }, is_last); + } + + const int ret = llama_decode(ctx_dft, batch_prompt); + if (ret != 0) { + SPF_ERR("failed to decode prompt on draft model, ret = %d\n", ret); + llama_batch_free(batch_prompt); + res.kept_indices.resize(prompt.size()); + std::iota(res.kept_indices.begin(), res.kept_indices.end(), 0); + res.n_prompt_kept = (int32_t) res.kept_indices.size(); + return res; + } + } + llama_batch_free(batch_prompt); + } + + const auto t_prefill_end = ggml_time_us(); + res.t_draft_eval_us = t_prefill_end - t_start; + + // 2. lookahead decode steps with attention extraction + std::vector total_importance(prompt.size(), 0.0f); + + cb_attn_collector_data cb_data; + cb_data.n_prompt = (int32_t) prompt.size(); + + // attach callback + llama_set_eval_callback(ctx_dft, cb_collect_attn, &cb_data); + + llama_batch batch_decode = llama_batch_init(1, 0, 1); + + int32_t actual_steps = 0; + int32_t cur_pos = (int32_t) prompt.size(); + + for (int32_t k = 0; k < lookahead; ++k) { + cb_data.layer_attns.clear(); + + const llama_token token_id = common_sampler_sample(smpl_dft, ctx_dft, -1); + common_sampler_accept(smpl_dft, token_id, true); + + if (llama_vocab_is_eog(vocab_dft, token_id)) { + break; + } + + common_batch_clear(batch_decode); + common_batch_add(batch_decode, token_id, cur_pos++, { seq_id }, true); + + const int ret = llama_decode(ctx_dft, batch_decode); + if (ret != 0) { + SPF_ERR("failed lookahead decode step %d, ret = %d\n", k, ret); + break; + } + + if (!cb_data.layer_attns.empty()) { + const auto step_importance = common_speculative_prefill_compute_importance( + cb_data.layer_attns, + (int32_t) prompt.size(), + n_heads_dft, + params.pool_kernel_size); + + for (size_t i = 0; i < prompt.size(); ++i) { + total_importance[i] += step_importance[i]; + } + actual_steps++; + } + } + + llama_batch_free(batch_decode); + + // detach callback + llama_set_eval_callback(ctx_dft, nullptr, nullptr); + + if (actual_steps == 0) { + LOG_WRN("spec-prefill: attention was not captured (flash attn on?); keeping full prompt\n"); + res.kept_indices.resize(prompt.size()); + std::iota(res.kept_indices.begin(), res.kept_indices.end(), 0); + res.n_prompt_kept = (int32_t) res.kept_indices.size(); + res.importance_scores.assign(prompt.size(), 1.0f); + return res; + } + + for (size_t i = 0; i < prompt.size(); ++i) { + total_importance[i] /= (float) actual_steps; + } + + const auto t_est_start = ggml_time_us(); + res.importance_scores = total_importance; + res.kept_indices = common_speculative_prefill_select_indices(total_importance, params); + res.n_prompt_kept = (int32_t) res.kept_indices.size(); + res.t_estimate_us = ggml_time_us() - t_est_start; + + return res; +} diff --git a/common/speculative-prefill.h b/common/speculative-prefill.h new file mode 100644 index 000000000000..0752449133df --- /dev/null +++ b/common/speculative-prefill.h @@ -0,0 +1,37 @@ +#pragma once + +#include "llama.h" +#include "common.h" + +#include +#include + + +struct common_speculative_prefill_result { + std::vector kept_indices; // sorted original token positions kept for target model + std::vector importance_scores; // per-token aggregated importance scores (size = n_prompt) + int64_t t_draft_eval_us = 0; + int64_t t_estimate_us = 0; + int32_t n_prompt_orig = 0; + int32_t n_prompt_kept = 0; +}; + +// calculate 1D average-pooled smoothed token importance from raw attention matrices +std::vector common_speculative_prefill_compute_importance( + const std::vector> & attn_layers, + int32_t n_prompt, + int32_t n_heads, + int32_t pool_kernel_size); + +// select kept token indices using chunk-based top-k or token-level ranking +std::vector common_speculative_prefill_select_indices( + const std::vector & importance_scores, + const common_params_speculative_prefill & params); + +// execute draft prefill, lookahead decoding, attention collection, and index selection +common_speculative_prefill_result common_speculative_prefill_execute( + llama_context * ctx_dft, + common_sampler * smpl_dft, + const std::vector & prompt, + llama_seq_id seq_id, + const common_params_speculative_prefill & params); diff --git a/common/speculative.cpp b/common/speculative.cpp index 77dfe9535f3f..2ed46e4f002d 100644 --- a/common/speculative.cpp +++ b/common/speculative.cpp @@ -159,7 +159,82 @@ struct common_speculative_impl { int64_t t_draft_us = 0; // total time spent in generating drafts in this implementation in microseconds. int64_t t_accept_us = 0; // total time spent in accumulation of this implementation in microseconds. - common_speculative_impl(common_speculative_type type, uint32_t n_seq, int32_t n_max) : type(type), n_seq(n_seq), n_max(n_max) {} + // Adaptive draft length. + // + // Drafting more tokens than the target will accept is pure waste: the draft + // pays for every token it proposes, and the target verifies all of them, but + // everything after the first rejection is discarded. A fixed n_max therefore + // overshoots on unpredictable content (prose) and undershoots on predictable + // content (JSON, verbatim quoting) -- which is exactly why MTP measures worse + // at n=7 than at n=3 on every content class while DFlash2 prefers 7. + // + // Track an EMA of how many tokens the target actually accepted per draft and + // size the next draft just above it. Sizing to ema+1 keeps one token of + // headroom so the estimate can climb again when content becomes predictable. + // The action censors the measurement: if every drafted token is accepted we + // only learn the true acceptance was *at least* n_drafted, never how much + // further it would have gone. Averaging a censored sample would ratchet the + // draft length down and strand it there. So the two outcomes are treated as + // different observations: + // partial accept -> uncensored, we saw the exact stopping point: track it + // full accept -> censored, a lower bound only: probe upward instead + // Backing off is averaged (gentle), probing is additive (faster), so the + // controller recovers quickly when content turns predictable again. + std::vector acc_ema; // per-seq EMA of accepted tokens per draft + std::vector n_last_draft; // per-seq size of the draft just issued + bool adaptive_n = false; // enabled by --spec-draft-adaptive + + static constexpr float acc_ema_alpha = 0.25f; // ~4-step memory + static constexpr float acc_ema_probe = 1.0f; // additive growth on a clean draft + static constexpr float acc_ema_init = 2.0f; + + void update_acc_ema(llama_seq_id seq_id, uint16_t n_accepted) { + if (seq_id < 0 || (size_t) seq_id >= acc_ema.size()) { + return; + } + const int32_t n_drafted = n_last_draft[seq_id]; + if (n_drafted > 0 && (int32_t) n_accepted >= n_drafted) { + acc_ema[seq_id] += acc_ema_probe; // censored: lower bound, probe up + } else { + acc_ema[seq_id] = (1.0f - acc_ema_alpha) * acc_ema[seq_id] + acc_ema_alpha * (float) n_accepted; + } + n_last_draft[seq_id] = 0; + } + + // reset on a new prompt / reused server slot; never mid-generation, since + // tracking content drift within a response is the point of the EMA + void reset_acc_ema(llama_seq_id seq_id) { + if (seq_id >= 0 && (size_t) seq_id < acc_ema.size()) { + acc_ema[seq_id] = acc_ema_init; + n_last_draft[seq_id] = 0; + } + } + + // record the width a sequence was actually drafted at, when the caller had to + // override the per-sequence choice (see the DFlash/DSpark batch below) + void set_last_draft(llama_seq_id seq_id, int32_t n) { + if (seq_id >= 0 && (size_t) seq_id < n_last_draft.size()) { + n_last_draft[seq_id] = n; + } + } + + // effective draft length for this step, never above the configured n_max + int32_t adaptive_n_draft(llama_seq_id seq_id, int32_t n_cfg, int32_t n_min) { + if (!adaptive_n || seq_id < 0 || (size_t) seq_id >= acc_ema.size()) { + return n_cfg; + } + const int32_t n_want = (int32_t) std::lround(acc_ema[seq_id]); + const int32_t n = std::max(std::max(1, n_min), std::min(n_cfg, n_want)); + acc_ema[seq_id] = std::min(acc_ema[seq_id], (float) n_cfg); // do not let the probe run away + n_last_draft[seq_id] = n; + return n; + } + + common_speculative_impl(common_speculative_type type, uint32_t n_seq, int32_t n_max) : type(type), n_seq(n_seq), n_max(n_max) { + // start optimistic so a predictable prefix is not throttled from step one + acc_ema.assign(n_seq, acc_ema_init); + n_last_draft.assign(n_seq, 0); + } virtual ~common_speculative_impl() = default; @@ -195,6 +270,7 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { } SPC_TRC("%s", "adding speculative implementation 'draft-simple'\n"); + adaptive_n = this->params.adaptive; SPC_TRC("- n_max=%d, n_min=%d, p_min=%f\n", this->params.n_max, this->params.n_min, this->params.p_min); SPC_TRC("- gpu_layers=%d, cache_k=%s, cache_v=%s, ctx_tgt=%s, ctx_dft=%s, devices=[%s]\n", this->params.n_gpu_layers, @@ -319,6 +395,8 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { auto * smpl = smpls[seq_id].get(); + const int32_t n_draft_eff = adaptive_n_draft(seq_id, params.n_max, params.n_min); + common_sampler_sample(smpl, ctx_dft, i_batch, true); ++i_batch; @@ -348,7 +426,7 @@ struct common_speculative_impl_draft_simple : public common_speculative_impl { result.push_back(id); - if ((params.n_max <= (int) result.size()) || + if ((n_draft_eff <= (int) result.size()) || (dp.n_max > 0 && dp.n_max <= (int) result.size())) { drafting[seq_id] = false; n_drafting--; @@ -458,6 +536,7 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl { , params(params.draft) { SPC_TRC("%s", "adding speculative implementation 'draft-eagle3'\n"); + adaptive_n = this->params.adaptive; SPC_TRC("- n_max=%d, n_min=%d, p_min=%f, backend_sampling=%d\n", params.draft.n_max, params.draft.n_min, params.draft.p_min, (int) params.draft.backend_sampling); auto * ctx_tgt = this->params.ctx_tgt; @@ -779,6 +858,8 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl { auto * smpl = smpls[seq_id].get(); + const int32_t n_draft_eff = adaptive_n_draft(seq_id, params.n_max, params.n_min); + common_sampler_sample(smpl, ctx_dft, i_batch, true); // pre-norm hidden state of this position becomes g_embd for the next step const float * prenorm = llama_get_embeddings_nextn_ith(ctx_dft, i_batch); @@ -810,7 +891,7 @@ struct common_speculative_impl_draft_eagle3 : public common_speculative_impl { result.push_back(id); - if (params.n_max <= (int) result.size()) { + if (n_draft_eff <= (int) result.size()) { drafting[seq_id] = false; n_drafting--; continue; @@ -992,6 +1073,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { } LOG_INF("%s: adding speculative implementation '%s'\n", __func__, common_speculative_type_to_str(type).c_str()); + adaptive_n = this->params.adaptive; LOG_INF("%s: - n_max=%d, n_min=%d, p_min=%.2f\n", __func__, this->params.n_max, this->params.n_min, this->params.p_min); LOG_INF("%s: - block_size=%d, mask_token_id=%d, n_extract=%u, sample_from_anchor=%s\n", __func__, block_size, mask_token_id, target_layer_ids_n, sample_from_anchor ? "true" : "false"); @@ -1070,6 +1152,8 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { } void begin(llama_seq_id seq_id, const llama_tokens & prompt) override { + + reset_acc_ema(seq_id); if (seq_id < 0 || seq_id >= (llama_seq_id) n_seq) { return; } @@ -1189,6 +1273,44 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { std::vector i_block_beg(n_seq, -1); std::vector n_block (n_seq, 0); + // DFlash decodes the whole block in one pass, so a shorter block does not save + // draft time -- but it does shrink the target's verification batch, which is + // where the cost actually is. + // + // The block width has to be the same for every sequence in this batch. All the + // blocks go into one decode, and the draft graph reshapes those tokens into + // [block_size, n_seqs] -- build_dflash2_conv and the DSpark markov head both + // assert n_tokens % n_seqs_unq == 0. Per-sequence widths break that: with an + // equal ubatch split the longer block is silently cut across two ubatches, so + // its tail is convolved as a block of its own and drafts from a truncated + // window; with a simple split it trips the assert and takes the server down. + // So size the batch once, from the sequence with the highest measured + // acceptance: no sequence drafts shorter than its own estimate, and the extra + // positions a colder sequence receives are just rejected at verification. + int32_t n_draft_batch = params.n_max; + + if (adaptive_n) { + n_draft_batch = 0; + + for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) { + if (!dparams[seq_id].drafting) { + continue; + } + + n_draft_batch = std::max(n_draft_batch, adaptive_n_draft(seq_id, params.n_max, params.n_min)); + } + + if (n_draft_batch == 0) { + return; + } + + for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) { + if (dparams[seq_id].drafting) { + set_last_draft(seq_id, n_draft_batch); + } + } + } + for (llama_seq_id seq_id = 0; seq_id < (llama_seq_id) n_seq; ++seq_id) { auto & dp = dparams[seq_id]; if (!dp.drafting) { @@ -1199,7 +1321,7 @@ struct common_speculative_impl_draft_dflash : public common_speculative_impl { const int32_t n = (int32_t) dp.pos0; - const int32_t n_draft = params.n_max; + const int32_t n_draft = n_draft_batch; const int32_t n_block_tokens = n_draft + (is_dspark && sample_from_anchor ? 0 : 1); i_block_beg[seq_id] = batch.n_tokens; @@ -1377,6 +1499,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { n_mtp_layers = std::max(1, (int) llama_model_n_layer_nextn(llama_get_model(ctx_dft))); SPC_TRC("%s", "adding speculative implementation 'draft-mtp'\n"); + adaptive_n = this->params.adaptive; SPC_TRC("- n_max=%d, n_min=%d, p_min=%.2f, n_embd=%d, backend_sampling=%d\n", this->params.n_max, this->params.n_min, this->params.p_min, n_embd, (int) this->params.backend_sampling); SPC_TRC("- gpu_layers=%d, cache_k=%s, cache_v=%s, ctx_tgt=%s, ctx_dft=%s, devices=[%s]\n", this->params.n_gpu_layers, @@ -1420,7 +1543,11 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { llama_set_embeddings_nextn(ctx_tgt, true, /*masked*/ false); llama_set_embeddings_nextn(ctx_dft, true, /*masked*/ true); - is_mem_shared = llama_get_ctx_other(ctx_dft) == ctx_tgt; + // only the gemma4 assistant shares the target's memory; other drafts (a qwen4exp shared sidecar) set ctx_other + // merely to borrow the target's embeddings and LM head + char arch[64] = {0}; + llama_model_meta_val_str(llama_get_model(ctx_dft), "general.architecture", arch, sizeof(arch)); + is_mem_shared = llama_get_ctx_other(ctx_dft) == ctx_tgt && std::strcmp(arch, "gemma4-assistant") == 0; chain_heads = n_mtp_layers > 1 && !is_mem_shared; if (chain_heads) { @@ -1464,6 +1591,8 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { } void begin(llama_seq_id seq_id, const llama_tokens & prompt) override { + + reset_acc_ema(seq_id); const int32_t N = (int32_t) prompt.size(); if (N <= 0) { return; @@ -1668,6 +1797,10 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { auto * smpl = smpls[seq_id].get(); + // MTP drafts sequentially, so a shorter draft saves draft passes + // as well as target verification work. + const int32_t n_draft_eff = adaptive_n_draft(seq_id, params.n_max, params.n_min); + common_sampler_sample(smpl, ctx_dft, i_last[seq_id], true); const float * h_row = llama_get_embeddings_nextn_ith(ctx_dft, i_last[seq_id]); @@ -1697,7 +1830,7 @@ struct common_speculative_impl_draft_mtp : public common_speculative_impl { result.push_back(id); - if (params.n_max <= (int) result.size()) { + if (n_draft_eff <= (int) result.size()) { drafting[seq_id] = false; n_drafting--; continue; @@ -2543,7 +2676,8 @@ common_speculative_init_result::common_speculative_init_result( auto cparams = common_context_params_to_llama(params); if (spec_mtp) { - cparams.ctx_type = LLAMA_CONTEXT_TYPE_MTP; + cparams.ctx_type = LLAMA_CONTEXT_TYPE_MTP; + cparams.mtp_draft_vocab = params.speculative.draft.mtp_vocab; } // the draft context holds as many tokens per sequence as the target context @@ -2913,6 +3047,7 @@ void common_speculative_accept(common_speculative * spec, llama_seq_id seq_id, u impl->n_acc_tokens += n_accepted; } + impl->update_acc_ema(seq_id, n_accepted); impl->accept(seq_id, n_accepted, false); impl->n_call_accept++; } diff --git a/conversion/base.py b/conversion/base.py index d2d80be3688b..4badf3cfb06a 100644 --- a/conversion/base.py +++ b/conversion/base.py @@ -120,6 +120,7 @@ class ModelBase: supports_mtp_export: bool = False mtp_only: bool = False no_mtp: bool = False + mtp_shared_embd: bool = False def __init__(self, dir_model: Path, ftype: gguf.LlamaFileType, fname_out: Path, *, is_big_endian: bool = False, use_temp_file: bool = False, eager: bool = False, diff --git a/conversion/qwen.py b/conversion/qwen.py index c7e0809f38c4..d50a2295e278 100644 --- a/conversion/qwen.py +++ b/conversion/qwen.py @@ -290,6 +290,7 @@ class _QwenMtpMixin: tensor_map: gguf.TensorNameMap no_mtp: bool mtp_only: bool + mtp_shared_embd: bool _original_block_count: int | None = None opt_num_mtp_layers: int = 0 @@ -338,7 +339,7 @@ def filter_tensors(cls, item): elif len(parts) == 3 and parts[1] in remapper: name = f"model.layers.{cls._original_block_count}.{remapper[parts[1]]}.{parts[2]}" elif cls.mtp_only: - keep = name in ( + keep = not cls.mtp_shared_embd and name in ( "model.embed_tokens.weight", "model.norm.weight", "lm_head.weight", "embed_tokens.weight", "norm.weight", ) diff --git a/conversion/qwen4exp.py b/conversion/qwen4exp.py index 168796d616b9..0d2a57d9dba7 100644 --- a/conversion/qwen4exp.py +++ b/conversion/qwen4exp.py @@ -1,6 +1,8 @@ from __future__ import annotations -from typing import Iterable, cast +import json +import logging +from typing import Callable, Iterable, cast import torch from torch import Tensor @@ -26,8 +28,6 @@ class Qwen4ExpTextModel(_Qwen35MRopeMixin, _LinearAttentionVReorderBase): model_arch = gguf.MODEL_ARCH.QWEN4EXP # the MTP block is a separate draft head; vLLM drops it too - supports_mtp_export = False - no_mtp = True def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) @@ -63,14 +63,16 @@ def set_gguf_parameters(self): self.gguf_writer.add_indexer_top_k(hp["indexer_budget"]) ratio = hp["indexer_compress_ratio"] layer_types = hp["layer_types"] - self.gguf_writer.add_attention_compress_ratios( - [ratio if layer_types[i] == "full_attention" else 0 for i in range(n_layer)] - ) + ratios = [ratio if layer_types[i] == "full_attention" else 0 for i in range(n_layer)] + # the MTP block extends block_count, and this array is read per block. The draft + # attends dense, so it takes 0, the same value a non-QSA layer gets. + ratios += [0] * (self.block_count - n_layer) + self.gguf_writer.add_attention_compress_ratios(ratios) # ple_layer_ids is 1-based in the HF config; empty means no n-gram table, # so emit no PLE keys rather than optional ones ple_layers = [i - 1 for i in hp["ple_layer_ids"]] - if not ple_layers: + if not ple_layers or self.mtp_only: return self.gguf_writer.add_ple_layers(ple_layers) self.gguf_writer.add_ple_ngram_size(hp["ngram_size"]) @@ -105,6 +107,47 @@ def _eos_token_id(self) -> int: raise ValueError("eos_token_id is required: the PLE hash resets its n-grams on it") return int(eos) + @classmethod + def filter_tensors(cls, item: tuple[str, Callable[[], Tensor]]) -> tuple[str, Callable[[], Tensor]] | None: + name = item[0] + # The MTP block carries its own final hyper-connection mixer. An MTP-only file has + # no trunk, so that one becomes the file's mixer; a full model already has the + # trunk's, and graph_mtp reads it, so this copy is dropped. + if name.startswith("mtp.hyper_connection_mixer."): + if cls.no_mtp: + return None + # the draft's mixer is its own (it differs from the trunk's), so it stays + # with the draft block as blk.N.nextn.hc_head_* in both file layouts + assert cls._original_block_count is not None + return (f"model.layers.{cls._original_block_count}.{name[len('mtp.'):]}", item[1]) + return super().filter_tensors(item) + + def generate_extra_tensors(self) -> Iterable[tuple[str, Tensor]]: + yield from super().generate_extra_tensors() + + # the reference adds the two MTP input projections, and A*e + B*h == [A|B]*concat(e, h), + # so they join into the one eh_proj the tensor map already handles + e_name = "mtp.fc_embedding.weight" + h_name = "mtp.fc_hidden.weight" + + have_e = e_name in self.model_tensors + have_h = h_name in self.model_tensors + if not have_e and not have_h: + return + if not have_e or not have_h: + raise KeyError(f"unpaired MTP input projection: need both {e_name} and {h_name}") + + from .base import LazyTorchTensor + + e = LazyTorchTensor.to_eager(self.model_tensors[e_name]()) + h = LazyTorchTensor.to_eager(self.model_tensors[h_name]()) + yield (self.format_tensor_name(gguf.MODEL_TENSOR.NEXTN_EH_PROJ, + self.hparams["num_hidden_layers"]), + torch.cat([e, h], dim=1).contiguous()) + + del self.model_tensors[e_name] + del self.model_tensors[h_name] + def modify_tensors(self, data_torch: Tensor, name: str, bid: int | None) -> Iterable[tuple[str, Tensor]]: # int64 hash constants must stay exact; 1-D tensors force F32, so use KV if name.endswith("ple_embedding.layer_multipliers"): diff --git a/convert_hf_to_gguf.py b/convert_hf_to_gguf.py index e09616b190cf..51614f941ccd 100755 --- a/convert_hf_to_gguf.py +++ b/convert_hf_to_gguf.py @@ -125,6 +125,10 @@ def parse_args() -> argparse.Namespace: "--no-nextn", "--no-mtp", dest="no_mtp", action="store_true", help="Exclude NextN speculative draft tensors from the converted GGUF. Pair with --mtp or --dspark on a second run to publish target and draft as two files.", ) + parser.add_argument( + "--mtp-shared-embd", action="store_true", + help="With --mtp, leave the token embeddings, output norm and LM head out of the draft and take them from the target model at load time. Much smaller draft, but it needs a llama.cpp new enough to read it.", + ) parser.add_argument( "--dspark", action="store_true", help="Export only the DeepSeek-V4 DSpark draft tensors as a separate GGUF.", @@ -282,6 +286,12 @@ def main() -> None: if args.mtp: model_class.mtp_only = True + if args.mtp_shared_embd: + if not args.mtp: + logger.error("--mtp-shared-embd only applies together with --mtp") + sys.exit(1) + model_class.mtp_shared_embd = True + model_instance = model_class(dir_model, output_type, fname_out, is_big_endian=args.bigendian, use_temp_file=args.use_temp_file, eager=args.no_lazy, diff --git a/docs/development/DSV4-vulkan-lightning-indexer-progress.md b/docs/development/DSV4-vulkan-lightning-indexer-progress.md new file mode 100644 index 000000000000..d19fbd6903aa --- /dev/null +++ b/docs/development/DSV4-vulkan-lightning-indexer-progress.md @@ -0,0 +1,158 @@ +# DeepSeek V4 Vulkan Lightning Indexer progress + +This is a restart note for the Strix Halo Lightning Indexer optimization. It is a development scratch pad and can be removed before the final PR. + +## Repository state + +- Main repository: `/home/jaap/Projects/git/llama.cpp` +- Optimization worktree: `/tmp/llama-strix-beta-bench` +- Branch: `strix-halo-vulkan-lightning-indexer` +- Base commit: `316c72ee9eab590f5891089d3b6bfc0d01d00d19` +- Base branch: Nathan's `strix-halo-vulkan-beta` +- Decode microbench work is stored in the main worktree as `stash@{0}: On strix-halo-vulkan: wip: DSV4 decode microbench depth matrix`. +- The Indexer changes are uncommitted. Do not commit without explicit user approval. An assisted commit needs an `Assisted-by:` trailer. +- Do not run builds and GPU benchmarks together. The APU shares its power and memory-bandwidth budget. +- GPU commands need sandbox escalation. + +## Objective and result + +After sparse prefill attention was flattened, the context-dependent Lightning Indexer became the next prefill bottleneck. The old cooperative-matrix shader used one wave64 subgroup per workgroup, processed one 16-key tile, and loaded one query head at a time. + +The new wide pipeline uses eight wave64 subgroups per workgroup. Each subgroup processes a separate 16-key tile, so one workgroup covers 128 keys. It stages four query heads and their weights together, reuses them across all eight subgroups, and uses subgroup-scoped synchronization between cooperative-matrix result stores. A workgroup barrier remains between four-head groups because all subgroups reuse the shared query storage. + +The optimized shader requires 512 workgroup invocations and 64 KiB shared memory. Pipeline creation is capability-based. Devices without those limits use a one-wave, one-head cooperative-matrix specialization. The scalar implementation remains the fallback when cooperative matrices are unavailable. The decode-specific cooperative-matrix pipeline is unchanged. + +At the 32k-equivalent prefill microbench shape: + +| Version | Time per layer | Throughput | +| --- | ---: | ---: | +| Baseline | 50.51 ms | 5.83 TFLOPS | +| Optimized | 31.81 ms | 9.25 TFLOPS | + +This is a 37.0% reduction in Lightning Indexer kernel time. + +The canonical 32k llama-bench improved from 209.45 to 216.32 tokens/s. Total Vulkan time fell from 9.73729 to 9.42448 seconds. Total Lightning Indexer time fell from 1.11860 to 0.697276 seconds. Sparse attention and top-K were effectively unchanged. + +## Changed files + +- `ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_cm.comp`: parameterizes the shader, adds the eight-wave four-head implementation, and remains usable for the small fallback. +- `ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp`: generates wide `N_WAVES=8`, `HEADS_PER_TILE=4` and small `N_WAVES=1`, `HEADS_PER_TILE=1` variants. +- `ggml/src/ggml-vulkan/ggml-vulkan.cpp`: creates and selects the capability-gated wide pipeline and the small cooperative-matrix fallback. +- `tests/test-backend-ops.cpp`: adds 127, 128, and 129-key correctness boundaries and PP2048 performance shapes through 512k simulated source context. + +## Performance data + +The performance rows model the actual PP2048 Indexer shapes after the source context is filled in 2048-token batches. `kv=8704` is the measured shape near 32k source context. It differs from 32768 because the Indexer compresses source tokens into rows. + +| Source depth | `kv` | Baseline | Optimized | Reduction | +| ---: | ---: | ---: | ---: | ---: | +| 0 | 512 | 4.23 ms | 2.14 ms | 49.5% | +| 8k | 2560 | 17.10 ms | 10.09 ms | 41.0% | +| 16k | 4608 | 29.97 ms | 17.62 ms | 41.2% | +| 32k | 8704 | 50.51 ms | 32.94 ms | 34.8% | +| 64k | 16896 | 94.36 ms | 64.87 ms | 31.3% | +| 128k | 33280 | 174.96 ms | 128.76 ms | 26.4% | +| 256k | 66048 | 352.73 ms | 250.88 ms | 28.9% | +| 512k | 131584 | 696.79 ms | 495.90 ms | 28.8% | + +The final isolated 32k run after cleanup measured 31.81159 ms. Small matrix differences are normal laptop GPU clock variation. + +| Canonical 32k metric | Baseline | Optimized | +| --- | ---: | ---: | +| PP2048 | 209.45 tokens/s | 216.32 tokens/s | +| Total Vulkan | 9.73729 s | 9.42448 s | +| Lightning Indexer | 1.11860 s | 0.697276 s | +| Sparse FA raw | 0.198438 s | 0.195367 s | +| Sparse FA selected | 0.835110 s | 0.843385 s | +| Sparse FA reduce | 0.088463 s | 0.086678 s | +| TOP_K | 0.075473 s | 0.075350 s | + +Logs: + +- Baseline matrix: `/tmp/dsv4-lightning-prefill-baseline.log` +- Optimized matrix: `/tmp/dsv4-lightning-four-head-matrix.log` +- Final selected-pipeline 32k microbench: `/tmp/dsv4-lightning-final-selected-32k.log` +- Baseline canonical 32k llama-bench: `/tmp/dsv4-nathan-beta-32k-rerun-new-first.log` +- Optimized canonical 32k llama-bench: `/tmp/dsv4-lightning-final-32k-llama-bench.log` + +## Correctness and resources + +- Wide pipeline: all 20 focused F16 cases passed, including 127, 128, and 129-key boundaries. +- Small cooperative-matrix fallback: temporarily forced and all the same 20 cases passed. +- Wide shader on gfx1151: 168 VGPRs, 63,488 bytes LDS, no spills, eight subgroups per SIMD. +- The final pipeline-statistics run confirmed `lightning_indexer_cm_f16` was selected. +- No NaN, Inf, or comparison failures were reported. + +Correctness logs: + +- Wide: `/tmp/dsv4-lightning-consolidated-wide-correctness.log` +- Small fallback: `/tmp/dsv4-lightning-consolidated-small-correctness.log` + +## Experiments and decisions + +- Four waves improved the 32k shape about 3% and became worse at deep simulated contexts. +- Eight waves improved it about 8% before the other changes. +- Subgroup-scoped synchronization after cooperative-matrix stores improved the eight-wave version. +- Staging four query heads and weights produced the large gain by reducing redundant loads and barriers. +- Using only a subgroup barrier between head groups failed four boundary tests. One subgroup could overwrite shared query data while another still read it. A workgroup barrier is required there. +- A separate fallback shader source was avoided. Generator definitions create both variants from one file. + +## Commands + +Build only, with no GPU benchmark running: + +```sh +cd /tmp/llama-strix-beta-bench +git diff --check +cmake --build build --config Release --target test-backend-ops llama-bench -j "$(nproc)" +``` + +Focused correctness: + +```sh +cd /tmp/llama-strix-beta-bench +./build/bin/test-backend-ops test -b Vulkan0 -o LIGHTNING_INDEXER -p 'type_K=f16' > /tmp/dsv4-lightning-correctness.log 2>&1 +tail -n 30 /tmp/dsv4-lightning-correctness.log +``` + +Final 32k-equivalent microbench and pipeline selection: + +```sh +cd /tmp/llama-strix-beta-bench +GGML_VK_PIPELINE_STATS=lightning_indexer_cm_f16 ./build/bin/test-backend-ops perf -b Vulkan0 -o LIGHTNING_INDEXER -p 'kv=8704' > /tmp/dsv4-lightning-final-selected-32k.log 2>&1 +tail -n 16 /tmp/dsv4-lightning-final-selected-32k.log +``` + +Full Indexer depth matrix: + +```sh +cd /tmp/llama-strix-beta-bench +./build/bin/test-backend-ops perf -b Vulkan0 -o LIGHTNING_INDEXER -p 'nb=2048,nh=64,ns=1,nm=1,type_K=f16' > /tmp/dsv4-lightning-matrix.log 2>&1 +rg 'kv=(512|2560|4608|8704|16896|33280|66048|131584),nb=2048' /tmp/dsv4-lightning-matrix.log +``` + +Canonical 32k llama-bench. Run it only when needed, never while compiling, and inspect only the final block: + +```sh +cd /tmp/llama-strix-beta-bench +GGML_VK_PERF_LOGGER=1 ./build/bin/llama-bench -m /home/jaap/Projects/docker/localLLaMA/models/models--unsloth--DeepSeek-V4-Flash-0731-GGUF/snapshots/109848da2469efe1f1aab9e11acea08a065ccd4f/UD-IQ3_XXS/DeepSeek-V4-Flash-0731-UD-IQ3_XXS-00001-of-00004.gguf -r 1 -d 32768 -p 2048 -ub 2048 -fa 1 -n 0 > /tmp/dsv4-lightning-32k-llama-bench.log 2>&1 +last=$(grep -n 'Vulkan Timings:' /tmp/dsv4-lightning-32k-llama-bench.log | tail -n 1 | cut -d: -f1) +sed -n "${last},\$p" /tmp/dsv4-lightning-32k-llama-bench.log | tail -n 180 +``` + +Patch inspection: + +```sh +cd /tmp/llama-strix-beta-bench +git diff --check +git diff --stat +git diff +git status --short +``` + +## Next actions + +1. The user reviews and understands the four-file implementation and result summary. +2. Commit only after explicit user approval for that commit action. +3. Remove this scratch pad before a PR if it is not useful as permanent documentation. +4. Restore the separate decode microbench stash from the main worktree only if that work resumes. diff --git a/docs/development/DSV4-vulkan-sparse-prefill-progress.md b/docs/development/DSV4-vulkan-sparse-prefill-progress.md new file mode 100644 index 000000000000..e880ae63e99e --- /dev/null +++ b/docs/development/DSV4-vulkan-sparse-prefill-progress.md @@ -0,0 +1,587 @@ +# DeepSeek V4 Vulkan sparse prefill progress + +This file is a self-contained handoff for the DeepSeek V4 sparse-attention prompt-processing optimization on AMD Strix Halo. Read the repository `AGENTS.md` and `CONTRIBUTING.md` before continuing. + +## Repository state + +- Repository: `https://github.com/Nathanw1014/llama.cpp` +- Branch: `strix-halo-vulkan` +- Starting commit: `baf0025de861c6f6ea3720fa81c52ae1b2e6c078` +- Target GPU: AMD Radeon 8060S / gfx1151, RADV, Vulkan, wave64 +- The device reports `GL_KHR_cooperative_matrix`, f16 inputs with f32 accumulation, 64 KiB shared memory, and a maximum 512-thread workgroup used by this path. + +The implementation is not upstream `ggml-org/llama.cpp`. It builds on this branch's DeepSeek V4 graph, Lightning Indexer, sparse top-K hint, decode gather path, fused HC kernels, and Vulkan profiler changes. + +## Important execution constraints + +The system is an APU. CPU compilation and GPU benchmarking share power and memory bandwidth. Never build and benchmark at the same time. Serialize all builds, correctness tests, and performance tests. + +GPU commands must run with host GPU access. In an agent sandbox, request elevated/out-of-sandbox execution. A sandboxed benchmark showed only CPU activity and is invalid. + +Redirect the full model benchmark to a log. Inspect only the last profiler block with `tail`; do not load the full log into agent context. + +## Build and ccache + +The build directory is `build`, configured as Release with Vulkan enabled. The default ccache directory was read-only in the agent environment, so use a writable directory: + +```bash +cmake -S . -B build \ + -DGGML_VULKAN=ON \ + -DCMAKE_BUILD_TYPE=Release \ + -DCMAKE_C_COMPILER_LAUNCHER=ccache \ + -DCMAKE_CXX_COMPILER_LAUNCHER=ccache + +CCACHE_DIR=/tmp/llama-cpp-ccache cmake --build build --config Release \ + --target llama-bench test-backend-ops -j "$(nproc)" + +CCACHE_DIR=/tmp/llama-cpp-ccache ccache -s +``` + +ccache was verified active. The final build reported direct hits. Note that an incremental change to `ggml-vulkan.cpp` is one large C++ translation unit and therefore uses one compiler core even with `-j`. Shader object regeneration can run in parallel. + +## Canonical benchmark command + +Do not change or omit switches for the 32k acceptance run: + +```bash +GGML_VK_PERF_LOGGER=1 ./build/bin/llama-bench \ + -m ~/Projects/docker/localLLaMA/models/models--unsloth--DeepSeek-V4-Flash-0731-GGUF/snapshots/109848da2469efe1f1aab9e11acea08a065ccd4f/UD-IQ3_XXS/DeepSeek-V4-Flash-0731-UD-IQ3_XXS-00001-of-00004.gguf \ + -r 1 -d 32768 -p 2048 -ub 2048 -fa 1 -n 0 \ + > /tmp/dsv4-vulkan-32k.log 2>&1 + +tail -n 180 /tmp/dsv4-vulkan-32k.log +``` + +The known model path exists on the target system. + +## Graph and dispatch findings + +DeepSeek V4 builds the Lightning Indexer and sparse attention in `src/models/deepseek4.cpp`: + +1. `build_lid_top_k()` creates indexer Q/K/weights and calls `ggml_lightning_indexer()`. +2. `ggml_top_k()` selects up to `hparams.indexer_top_k` compressed-cache indices for every query token. +3. `build_csa_lid_attention()` concatenates the raw SWA K prefix with compressed CSA K, builds a dense mask carrying the same sparse selection, and calls `build_attn_mha(..., top_k, raw_k->ne[2])`. +4. `build_attn_mha()` attaches `top_k` and `n_kv_raw` to `GGML_OP_FLASH_ATTN_EXT` through `ggml_flash_attn_ext_add_top_k()`. + +At the final 32k PP2048 batch: + +- total K rows: 11,008 +- `n_kv_raw`: 2,304 raw SWA rows, always attended subject to the mask +- `n_top_k`: 512 selected compressed rows per query token +- active rows: 2,816 +- selectable compressed region: 8,704 + +The custom Vulkan path is selected by `ggml_vk_flash_attn_top_k()` before ordinary FA. Its gate requires the DeepSeek V4 shape and `total_k >= 3 * (n_kv_raw + n_top_k)`. The final shape satisfies `11008 >= 3 * 2816`. + +The old `flash_attn_top_k.comp` shader is scalar/subgroup code. One 512-thread workgroup covers eight heads for one query token. It stages 16 selected 512-wide K/V rows, computes QK with scalar FMAs and `subgroupAdd`, updates online softmax one key at a time, and accumulates PV manually. It does not use cooperative matrices. + +The top-K set differs by query token but is shared by all 64 query heads for that token. This makes the attention for one token a regular matrix problem across heads and selected keys despite sparse per-token indexing. + +## Root cause evidence + +The existing Vulkan timestamp infrastructure was extended with `ggml_vk_perf_mark_subop()` after the sparse dispatch. This reports the sparse kernel separately as `FA_TOP_K_SPARSE (sub-op)` or `FA_TOP_K_CM (sub-op)`. + +Focused test shape: + +```bash +GGML_VK_PERF_LOGGER=1 ./build/bin/test-backend-ops perf \ + -b Vulkan0 -o FLASH_ATTN_EXT \ + -p 'kv=32768,nb=512,n_kv_raw=1024,n_top_k=512,sinks=0' +``` + +Results: + +- old scalar sparse kernel: 61.42 ms, 1.68 TFLOPS of useful active-set work +- ordinary dense FA diagnostic (`GGML_VK_FA_TOPK=0`): 255.97 ms, about 8.7 TFLOPS over the full dense work +- final cooperative sparse kernel: 32.55 ms, 3.17 TFLOPS of useful active-set work + +The residual `FLASH_ATTN_EXT` interval after the sparse timestamp is only about 4-7 us. The cost is inside the shader, not dispatch or surrounding synchronization. + +A temporary uniform stage-profiling mode was used and removed. For the final 32-head tile: + +- selected K gather + cooperative QK: 14.30 ms +- gather + QK + serial softmax: 26.34 ms +- gather + QK + parallel softmax: 15.53 ms +- full kernel: 32.55 ms +- the remaining cooperative PV/output portion is about 17.0 ms + +The old scalar shader was compute/issue inefficient. Dense FA proved matrix hardware is much faster but was still too expensive because it processes all K rows. The final implementation preserves sparsity and uses the matrix hardware for both QK and PV. + +## Implementation + +Files changed: + +- `ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k_cm.comp` + - New cooperative-matrix sparse prefill shader. + - One 512-thread workgroup covers 32 query heads for one token. + - Eight wave64 subgroups cover two 16-head tiles by four 16-key or 16-output-dimension tiles. + - Processes 64 selected keys per online-softmax block. + - Stages only indexed selected K/V tiles, never the full K range. + - Uses f16 cooperative-matrix inputs and f32 accumulation for QK and PV. + - Uses a 16-lane segmented softmax per head. XOR subgroup shuffles reduce max and sum for four independent heads per wave without workgroup barriers. + - Keeps f32 output accumulators and normalizes after all active blocks. + - Preserves the raw prefix, top-K index validation, mask, sinks, stream strides, and K == V latent behavior. +- `ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp` + - Embeds the new shader when cooperative-matrix shader support is available. +- `ggml/src/ggml-vulkan/ggml-vulkan.cpp` + - Adds the cooperative sparse pipeline when the device supports the required 16x16x16 f16/f32 cooperative matrix shape. + - Selects it by capability and keeps the scalar shader as fallback. + - Adds `GGML_VK_FA_TOPK=0` to force ordinary dense FA for diagnostics. + - Adds `GGML_VK_FA_TOPK_CM=0` to force the old scalar sparse shader for A/B tests. + - Adds sparse sub-operation timestamps through the existing profiler. + +The path is capability-based, not hardcoded to Strix Halo. The current sparse shape gate remains DeepSeek V4-specific. Devices without the required cooperative matrix support keep the correct scalar sparse or dense fallback. + +Two discarded prototypes are useful context: + +- A 16-head, 64-key streaming cooperative tile was correct and reduced the focused test from 61.4 to 44.6 ms. +- Keeping 32 complete 512-wide K/V rows in LDS grew shared memory to about 41 KiB, reduced residency, doubled block/barrier count, and regressed to 73.8 ms. Do not retry full-row LDS staging without solving occupancy. +- A 32-head, 64-key streaming tile halved irregular row loads but initially stayed near 44.7 ms because serial softmax cost about 11.5 ms. Parallel segmented softmax produced the final 32.55 ms result. + +## Correctness validation + +Run: + +```bash +./build/bin/test-backend-ops test -b Vulkan0 -o FLASH_ATTN_EXT -p 'n_top_k=' +``` + +Final result: 8/8 sparse top-K FA cases passed against the CPU reference. Cases include: + +- decode and short batches that use dense/gather fallback +- sparse prefill batch sizes 64 and 128 +- `n_kv_raw` plus top-K selection +- invalid top-K index handling from the test fixture +- sinks enabled and disabled +- sparse threshold transitions +- an active-key count of 193, which exercises a partial final 64-key block + +The test uses the existing FA tolerance of NMSE <= `5e-4`. No NaN or Inf failure occurred. The implementation changes Q and probability inputs to f16 cooperative-matrix operands with f32 accumulation, matching the precision strategy of ordinary Vulkan cooperative FA. + +Still desirable before broader submission: + +- compare model logits on controlled prompts between `GGML_VK_FA_TOPK_CM=0` and the default cooperative path + +## Canonical 32k results + +Exact clean runs, same command and machine, no concurrent build: + +```text +32k context, PP 2048, ub 2048 + +Before (commit baf0025de): +112.29 tok/s +Total Vulkan: 18.1957 s +Sparse FA: 8.84496 s, 421.189 ms/layer +Lightning Indexer: 1.19438 s +TOP_K: 0.076948 s + +After: +152.32 tok/s +Total Vulkan: 13.4032 s +Sparse FA: 4.44701 s, 211.762 ms/layer +Lightning Indexer: 1.13156 s +TOP_K: 0.073945 s + +Change: +Throughput: +35.65% +Total Vulkan time: -26.34% +Sparse FA time: -49.72% +Sparse FA saved: 4.398 s +Total GPU time saved: 4.793 s +``` + +The profiler now lists the optimized dispatch as `FA_TOP_K_CM (sub-op)`. The following residual `FLASH_ATTN_EXT` line is only the post-mark interval and must not be interpreted as the kernel time. + +## Context-depth measurements + +All points use PP2048, ub2048, FA enabled, one repetition, and no token generation. They were run sequentially with no compiler active: + +```text +Existing depth tok/s Total Vulkan Final large FA Lightning Indexer TOP_K +0 253.44 8.041 s 0.294 s 0.095 s 0.001 s +8192 211.01 9.666 s 1.625 s 0.379 s 0.022 s +16384 177.19 11.518 s 3.031 s 0.678 s 0.044 s +32768 152.32 13.403 s 4.447 s 1.132 s 0.074 s +``` + +At 0, 8k, and 16k, total K is below the existing sparse-path gate `total_k >= 3 * (n_kv_raw + n_top_k)`. These points use the unchanged ordinary dense FA implementation, so the cooperative sparse change does not affect or regress them. At 32k, total K is 11,008 and the cooperative sparse path engages. The 32k `Final large FA` value is the `FA_TOP_K_CM (sub-op)` total; the lower-depth values are the large ordinary `FLASH_ATTN_EXT` totals. + +Logs: + +- `/tmp/dsv4-vulkan-cm-0k.log` +- `/tmp/dsv4-vulkan-cm-8k.log` +- `/tmp/dsv4-vulkan-cm-16k.log` +- `/tmp/dsv4-vulkan-cm-32k.log` +- `/tmp/dsv4-vulkan-baseline-clean.log` +- `/tmp/dsv4-fa-cm-correctness-final.log` + +## Next optimization target + +The cooperative sparse FA remains the largest context-dependent cost at about 4.45 s total. Stage profiling indicates approximately 14.3 ms of focused-test time in gather/QK, about 1.2 ms in parallel softmax, and about 17 ms in PV/output. + +The next useful work is PV and output accumulation, not TOP_K. Investigate: + +- reducing repeated selected V staging across the two 32-head workgroups per token without increasing LDS enough to lose occupancy +- reducing the eight output-dimension passes or retaining more PV state in cooperative fragments/registers +- checking register count and spills for the 32 f32 output accumulators per invocation using RADV shader statistics +- alternate 32-head layouts that keep the same eight-wave occupancy but improve PV scheduling +- query-tile overlap/union gathering only if measured top-K overlap is high enough; a whole-2048-query union is unlikely to help + +Do not optimize TOP_K first. At 32k it is only about 74 ms total. Lightning Indexer is about 1.13 s and is the next context-dependent target only after sparse FA improves further. + +## Useful diagnostics + +Force old scalar sparse path: + +```bash +GGML_VK_FA_TOPK_CM=0 GGML_VK_PERF_LOGGER=1 ./build/bin/test-backend-ops perf \ + -b Vulkan0 -o FLASH_ATTN_EXT \ + -p 'kv=32768,nb=512,n_kv_raw=1024,n_top_k=512,sinks=0' +``` + +Force ordinary dense FA: + +```bash +GGML_VK_FA_TOPK=0 GGML_VK_PERF_LOGGER=1 ./build/bin/test-backend-ops perf \ + -b Vulkan0 -o FLASH_ATTN_EXT \ + -p 'kv=32768,nb=512,n_kv_raw=1024,n_top_k=512,sinks=0' +``` + +Default cooperative sparse path: + +```bash +GGML_VK_PERF_LOGGER=1 ./build/bin/test-backend-ops perf \ + -b Vulkan0 -o FLASH_ATTN_EXT \ + -p 'kv=32768,nb=512,n_kv_raw=1024,n_top_k=512,sinks=0' +``` + +Always run these sequentially. Do not run a compiler concurrently on this APU. + +## Experimental raw-prefix split prototype + +An uncommitted follow-up prototype was tested after commit `24c7ead76cc1a9631fa1b42b0bfa53a15169e1ba`. Check `git status` before continuing. It was initially tested with `GGML_VK_FA_TOPK_SPLIT=1`. After the successful 32k run, the split path was changed to default-on for a llama-server coherence test. Set `GGML_VK_FA_TOPK_SPLIT=0` to restore the single cooperative sparse kernel. + +Stage profiling on the previously used 19,200-row sparse shape (`kv=19200,nb=2048,n_kv_raw=2304,n_top_k=512,sinks=0`) showed: + +```text +cooperative QK and selected-K gather: 88.885 ms +softmax increment: 4.909 ms +cooperative PV and output increment: 128.399 ms +full cooperative sparse kernel: 222.193 ms +``` + +PV and output were 57.8% of the kernel. More importantly, most active keys are not sparse: all 2,304 raw-prefix rows are contiguous, while only 512 compressed rows use top-K indices. The prototype therefore makes two attention partitions: + +1. Ordinary optimized Vulkan cooperative FA processes the contiguous raw prefix. +2. The cooperative top-K shader processes only the 512 selected compressed rows. +3. The existing split-K reduction combines both online-softmax partitions and applies sinks. + +This preserves the exact raw-prefix, sparse top-K, causal mask, and softmax semantics. It reuses the existing ordinary FA and split-K reduction rather than adding a new subsystem. + +The exact-shape microbenchmark improved from 222.19 ms to 111.33 ms. Its steady split stages were about 62-64 ms raw-prefix FA, 43-46 ms selected sparse FA, and 3.6-3.9 ms reduction. + +Correctness validation: + +```bash +GGML_VK_FA_TOPK_SPLIT=1 ./build/bin/test-backend-ops test \ + -b Vulkan0 -o FLASH_ATTN_EXT -p 'n_top_k=' + +./build/bin/test-backend-ops test -b Vulkan0 -o FLASH_ATTN_EXT +``` + +Results were 8/8 sparse top-K cases and 13,296/13,296 complete Vulkan FA cases against the CPU reference. No NaN or Inf failure occurred. + +Canonical 32k prototype command: + +```bash +GGML_VK_FA_TOPK_SPLIT=1 GGML_VK_PERF_LOGGER=1 ./build/bin/llama-bench \ + -m ~/Projects/docker/localLLaMA/models/models--unsloth--DeepSeek-V4-Flash-0731-GGUF/snapshots/109848da2469efe1f1aab9e11acea08a065ccd4f/UD-IQ3_XXS/DeepSeek-V4-Flash-0731-UD-IQ3_XXS-00001-of-00004.gguf \ + -r 1 -d 32768 -p 2048 -ub 2048 -fa 1 -n 0 \ + > /tmp/dsv4-vulkan-split-32k.log 2>&1 +``` + +Final 32k result: + +```text +32k context, PP 2048, ub 2048 + +Old scalar: 112.29 tok/s, 18.1957 s total, 8.84496 s sparse FA +Committed coopmat: 152.32 tok/s, 13.4032 s total, 4.44701 s sparse FA +Experimental split: 208.70 tok/s, 9.7704 s total, 1.20170 s split sparse FA + +Experimental split stages: +raw-prefix FA: 0.210775 s total, 10.037 ms/layer +selected sparse FA: 0.908190 s total, 43.247 ms/layer +split reduction: 0.082732 s total, 3.940 ms/layer +Lightning Indexer: 1.110410 s total, 52.877 ms/layer +TOP_K: 0.077394 s total, 3.685 ms/layer +``` + +Relative to the committed cooperative path, throughput improved 37.0%, total Vulkan time fell 27.1%, and sparse FA time fell 73.0%. Relative to the old scalar path, throughput improved 85.9%, total Vulkan time fell 46.3%, and sparse FA time fell 86.4%. + +The initial split implementation used 538,968,064 bytes (514 MiB) of scratch at PP2048 because each of two partitions stored an f32 partial output for 512 dimensions x 64 heads x 2048 queries, plus L/M data. This was subsequently reduced by query tiling as described below. + +The old sparse selection heuristic required `total_k >= 3 * active_k`. For the coherence test it now selects sparse attention whenever `total_k > active_k`; equality still uses dense FA because no keys are pruned. The shape, capability, and allocation gates remain unchanged. + +Prototype logs: + +- `/tmp/dsv4-fa-exact-qk.log` +- `/tmp/dsv4-fa-exact-softmax.log` +- `/tmp/dsv4-fa-exact-full.log` +- `/tmp/dsv4-fa-exact-split.log` +- `/tmp/dsv4-fa-split-correctness.log` +- `/tmp/dsv4-fa-all-correctness.log` +- `/tmp/dsv4-vulkan-split-32k.log` + +## Llama-server coherence check + +After making the split path default-on and changing the sparse crossover to `total_k > active_k`, `llama-server` produced a coherent response from a 2,044-token prompt at significant context depth. The final profiler block confirmed that all 21 sparse-attention layers used the split path: + +```text +FA_TOP_K_RAW: 0.179223 s total, 8.534 ms/layer +FA_TOP_K_SELECTED: 0.886629 s total, 42.220 ms/layer +FA_TOP_K_REDUCE: 0.083240 s total, 3.964 ms/layer +Split sparse FA: 1.149092 s total, 54.719 ms/layer +Lightning Indexer: 1.054610 s total, 50.220 ms/layer +TOP_K: 0.044540 s total, 2.121 ms/layer +Total Vulkan: 9.797560 s +``` + +This closely matches the canonical PP2048 llama-bench result of 9.770 s total and 1.202 s split sparse FA. The focused sparse CPU-reference test was rerun after the crossover change and passed 8/8 cases. + +## Tiled split scratch optimization + +Commit `4bbe53e4775f0707de8158e4977a16ab770829da` used two full PP2048 output partitions. The same two-partition algorithm now processes at most 256 query tokens per tile and reuses the split scratch between tiles. Q, mask, top-K, and destination descriptors are offset to the tile while K/V remain shared. A Vulkan pipeline barrier separates reuse of each scratch tile. + +Scratch at PP2048 changed from: + +```text +Before: 538,968,064 bytes (514 MiB) +After: 67,371,008 bytes (64.25 MiB) +Change: 8x reduction +``` + +The previously used 19,200-row microbenchmark (`kv=19200,nb=2048,n_kv_raw=2304,n_top_k=512,sinks=0`) measured: + +```text +Full-batch two partitions: 114.65 ms +256-query tiled path: 113.70 ms +``` + +Two one-partition alternatives were tested and discarded. Merging the raw partial directly inside the cooperative selected shader measured 119.17 ms. Writing selected output separately and using a lightweight merge kernel measured 120.45 ms. Both cut scratch in half but regressed because writing selected output outside the contiguous split layout increased the selected stage from about 44 ms to about 51 ms. Query tiling preserves the faster memory layout. + +Canonical 32k result after tiling: + +```text +32k context, PP 2048, ub 2048 + +Full-batch split: 208.70 tok/s, 9.77039 s total, 1.20170 s sparse FA +Tiled split: 209.62 tok/s, 9.72897 s total, 1.18721 s sparse FA + +Tiled split stages: +raw-prefix FA: 0.192879 s total, 168 tile dispatches +selected sparse FA: 0.910153 s total, 168 tile dispatches +split reduction: 0.084179 s total, 168 tile dispatches +Lightning Indexer: 1.100230 s total +TOP_K: 0.075158 s total +``` + +The 168 dispatch count is eight tiles x 21 sparse-attention layers. Relative to the full-batch split path, throughput improved 0.44%, total Vulkan time fell 0.42%, and sparse FA time fell 1.21%. The main result is the 8x scratch reduction without a performance regression. + +Final correctness after removing the discarded merge prototypes: + +- focused sparse top-K suite: 9/9 passed against the CPU reference, including a 257-query tile-boundary case +- complete Vulkan Flash Attention suite: 13,296/13,296 passed +- no NaN or Inf failure + +Logs: + +- `/tmp/dsv4-fa-exact-fused.log` +- `/tmp/dsv4-fa-exact-merge.log` +- `/tmp/dsv4-fa-exact-two-part-current.log` +- `/tmp/dsv4-fa-exact-tiled.log` +- `/tmp/dsv4-fa-tiled-final-correctness.log` +- `/tmp/dsv4-fa-tiled-boundary-correctness.log` +- `/tmp/dsv4-fa-tiled-all-correctness.log` +- `/tmp/dsv4-vulkan-tiled-32k.log` + +## Selected PV probability reuse + +The selected cooperative-matrix stage remained the largest split sparse-FA component. Profiling the tiled 19,200-row shape showed: + +```text +selected K gather and cooperative QK: about 20.0 ms +QK plus softmax: about 22.3 ms +full selected stage: about 45.5 ms +cooperative PV and output increment: about 23.2 ms +``` + +PV and output were about 51% of the selected stage. The shader previously loaded each of four 16x16 probability cooperative-matrix fragments again for every one of the eight 64-dimension PV passes. The optimized shader loads these four fragments once per 64-key block and retains them across all PV dimension passes. This follows the probability-fragment lifetime used by ordinary cooperative Vulkan FA. + +The shader also aliases the shared score and PV-output matrices because their lifetimes do not overlap. Shader resource statistics changed as follows: + +```text + Before After +VGPRs 192 192 +VGPR spills 0 0 +LDS 36,864 28,672 bytes +static instructions 11,524 11,358 +``` + +The 19,200-row microbenchmark changed from 113.70 ms for the committed tiled path to 110.11 ms. The selected stage fell from about 45.5 ms to about 43.7 ms. The raw-prefix and reduction implementations are unchanged. + +Two canonical 32k runs after this change measured: + +```text +32k context, PP 2048, ub 2048 + +Committed tiled reference: +209.62 tok/s +Total Vulkan: 9.72897 s +Split sparse FA: 1.18721 s + raw-prefix FA: 0.192879 s + selected sparse FA: 0.910153 s + split reduction: 0.084179 s +Lightning Indexer: 1.100230 s +TOP_K: 0.075158 s + +Probability reuse, run 1: +211.03 tok/s +Total Vulkan: 9.66363 s +Split sparse FA: 1.12776 s + raw-prefix FA: 0.192078 s + selected sparse FA: 0.849862 s + split reduction: 0.085821 s +Lightning Indexer: 1.098980 s +TOP_K: 0.071036 s + +Probability reuse, run 2: +209.57 tok/s +Total Vulkan: 9.72968 s +Split sparse FA: 1.13553 s + raw-prefix FA: 0.192987 s + selected sparse FA: 0.859890 s + split reduction: 0.082653 s +Lightning Indexer: 1.102550 s +TOP_K: 0.071474 s +``` + +The two-run selected-stage improvement is 5.5-6.6%, and the complete split sparse-FA improvement is 4.4-5.0%. The two-run throughput mean is 210.30 tok/s, 0.32% above the 209.62 tok/s reference. End-to-end noise in unrelated model kernels is larger than this small total-throughput change, but both profiler runs isolate a consistent gain in the modified selected stage. + +Correctness after the change: + +- focused sparse top-K suite: 9/9 passed against the CPU reference +- complete Vulkan Flash Attention suite: 13,297/13,297 passed +- pipeline statistics probe: 1/1 passed, with no SGPR or VGPR spills +- no NaN or Inf failure + +Logs: + +- `/tmp/dsv4-fa-selected-qk-tiled.log` +- `/tmp/dsv4-fa-selected-softmax-tiled.log` +- `/tmp/dsv4-fa-lds-alias-stats.log` +- `/tmp/dsv4-fa-pmat-stats.log` +- `/tmp/dsv4-fa-exact-lds-alias.log` +- `/tmp/dsv4-fa-exact-pmat.log` +- `/tmp/dsv4-fa-pmat-correctness.log` +- `/tmp/dsv4-fa-pmat-all-correctness.log` +- `/tmp/dsv4-vulkan-32k-pmat.log` +- `/tmp/dsv4-vulkan-32k-pmat-repeat.log` + +The next sparse-FA optimization should continue to target the selected PV/output half. The retained probability fragments remove redundant cooperative loads without increasing reported VGPR allocation. More invasive changes such as doubling the PV dimension tile can reduce barriers but must be designed around the eight available wave64 subgroups, 64 KiB LDS limit, and already high 192-VGPR allocation. Do not build a crossover matrix until the next kernel layout is settled because an optimization can shift the crossover. + +## Selected mask caching and deep-context micro matrix + +The sparse FA `kv` dimension is compressed K/V rows, not source-token context depth. For the tested DeepSeek V4 graph, a PP2048 batch uses 2,304 raw rows and approximately one compressed row per four source tokens. The useful synthetic mapping is: + +```text +Source context Sparse FA kv rows +32k 11,008 +64k 19,200 +128k 35,584 +256k 68,352 +512k 133,888 +``` + +The performance test registry now contains PP2048 cases at all five K extents with `n_kv_raw=2304` and `n_top_k=512`. Use this command template and replace `KV` with a value from the table: + +```bash +GGML_VK_PERF_LOGGER=1 ./build/bin/test-backend-ops perf \ + -b Vulkan0 -o FLASH_ATTN_EXT \ + -p 'kv=KV,nb=2048,n_kv_raw=2304,n_top_k=512,sinks=0' +``` + +These are synthetic sparse-FA depth tests. They do not include the Lightning Indexer or the rest of the model graph and do not replace the canonical 32k llama-bench. + +The selected shader previously loaded the same selected-key mask value independently for all 32 heads during both softmax passes. The shader now loads each mask value once per 64-key block and reuses it from LDS. Q and probability storage also share one LDS allocation, while K and V staging share another because both pairs have disjoint lifetimes. + +Shader resources changed from the probability-reuse commit: + +```text + Before After +VGPRs 192 192 +VGPR spills 0 0 +LDS 28,672 24,576 bytes +static instructions 11,358 11,233 +``` + +A 128-dimension V staging experiment was correct but rejected. With operand aliasing it used exactly 32 KiB LDS. It was 2.3% slower at simulated 32k, approximately equal at 64k, and 1.6% faster at 128k. The retained 64-dimension layout is better for the canonical 32k target and is simpler. + +Selected-mask caching improved every synthetic depth relative to the same 64-dimension operand-alias layout: + +```text +Depth Selected before Selected after Split total before Split total after +32k 43.16 ms 41.23 ms 109.45 ms 108.07 ms +64k 43.22 ms 41.42 ms 110.45 ms 109.11 ms +128k 52.26 ms 48.74 ms 119.15 ms 116.92 ms +256k 51.64 ms 49.02 ms 118.06 ms 116.93 ms +512k 51.33 ms 49.20 ms 117.87 ms 116.27 ms +``` + +The 256k case exposed a separate path-selection cutoff. The raw-prefix ordinary FA dispatch encoded its mask row stride in 16 bits and disabled split sparse FA above 65,535 K rows. It now uses a flagged convention that carries the full 32-bit mask stride in the existing `split_kv` push constant for this one-partition partial-output dispatch. No push-constant structure was enlarged. At simulated 256k, this changes the selected path from unsplit cooperative sparse FA at 206.22 ms to split sparse FA at 118.06 ms before mask caching, a 42.8% reduction. The 512k case also uses the split path successfully. + +Canonical 32k PP2048 after operand aliasing and selected-mask caching: + +```text +Previous commit, two runs: +209.57-211.03 tok/s +Total Vulkan: 9.664-9.730 s +Split sparse FA: 1.128-1.136 s + raw-prefix FA: 0.192-0.193 s + selected sparse FA: 0.850-0.860 s + split reduction: 0.083-0.086 s + +Current result: +210.47 tok/s +Total Vulkan: 9.68921 s +Split sparse FA: 1.10160 s + raw-prefix FA: 0.192860 s + selected sparse FA: 0.821381 s + split reduction: 0.087354 s +Lightning Indexer: 1.102230 s +TOP_K: 0.072488 s +``` + +The selected stage improves 3.4-4.5% and complete split sparse FA improves 2.3-3.0% relative to the previous commit's two canonical runs. End-to-end throughput remains within model-wide benchmark noise. + +Correctness: + +- focused sparse top-K suite: 9/9 passed +- complete Vulkan Flash Attention suite: 13,297/13,297 passed +- no NaN or Inf failure +- simulated 256k and 512k performance cases both selected the split path and completed successfully + +Benchmark policy for this laptop APU: do not repeat llama-bench when the exact-shape microbenchmarks and the first canonical profiler block agree. Sustained load can lower GPU clocks and bias a repeat. Repeat only when the first result is anomalous or contradicts the microbenchmarks. Never build and benchmark concurrently. + +Logs: + +- `/tmp/dsv4-fa-mask-cache-32k.log` +- `/tmp/dsv4-fa-mask-cache-64k.log` +- `/tmp/dsv4-fa-mask-cache-128k.log` +- `/tmp/dsv4-fa-mask-cache-256k.log` +- `/tmp/dsv4-fa-mask-cache-512k.log` +- `/tmp/dsv4-fa-final-micro-32k.log` +- `/tmp/dsv4-fa-final-focused-correctness.log` +- `/tmp/dsv4-fa-mask-cache-all-correctness.log` +- `/tmp/dsv4-vulkan-32k-mask-cache.log` diff --git a/docs/development/backend-scheduler.md b/docs/development/backend-scheduler.md new file mode 100644 index 000000000000..b56816fe4f78 --- /dev/null +++ b/docs/development/backend-scheduler.md @@ -0,0 +1,237 @@ +# The backend scheduler + +`ggml_backend_sched` lets one graph run across several backend devices. It assigns every node to +a backend, allocates the compute buffers, inserts the copies needed where a node's operands live +on another device, and executes the result. + +Callers interact with it through `ggml-backend.h`; this document describes what it does and the +contracts it places on callers and on backends. + +## Contents + +- [Assigning nodes to backends](#assigning-nodes-to-backends) +- [Splits](#splits) +- [Allocation and graph reuse](#allocation-and-graph-reuse) +- [Graph input ring buffer](#graph-input-ring-buffer) +- [Output lifetime under asynchronous execution](#output-lifetime-under-asynchronous-execution) +- [Memory read in place by another backend](#memory-read-in-place-by-another-backend) +- [Ops that alias their source](#ops-that-alias-their-source) +- [The scheduler sanitizer](#the-scheduler-sanitizer) +- [Environment variables](#environment-variables) + +## Assigning nodes to backends + +The caller defines backend priority by the order of the array passed to `ggml_backend_sched_new()`: lower indices have higher priority, and the last backend must be the CPU. In llama.cpp this array contains the selected model devices in their configured order, then accelerator backends such as BLAS, then the CPU. + +Placement combines that order with hard constraints and locality heuristics: + +1. An assignment made with `ggml_backend_sched_set_tensor_backend()` is retained. +2. A preallocated tensor or view must use a backend that supports both its existing buffer type and its operation. The first compatible backend in caller order is selected; the allocation itself cannot move. +3. An unallocated graph input (`GGML_TENSOR_FLAG_INPUT`) is assigned to the last backend so that the caller can write it through the CPU slot. +4. Operations using weights prefer the first backend that can use the weight buffer and run the operation, avoiding a weight copy. Selected operations may instead be offloaded to an earlier backend when `op_offload` requests it. +5. Existing non-CPU assignments are expanded forward and backward through adjacent supported operations. Remaining assignments are expanded after that so contiguous runs stay together instead of alternating backends. +6. An unassigned operation is placed on the backend supporting the largest number of its assigned inputs. Caller order breaks ties. An assigned operation may be upgraded to an earlier backend when both use the same buffer type and that backend supports the operation and every source buffer. + +## Splits + +Once every node has a backend, the graph is cut into **splits**: maximal runs of consecutive nodes +sharing one backend. Each split is submitted to its backend as a single graph, in order. + +Where a split reads a tensor produced on another backend, the scheduler checks whether the +consuming backend can read that buffer directly (`ggml_backend_supports_buft`). If it can, the +tensor is read where it is. If it cannot, a copy is created in the consuming backend's buffer and +the node's operand is repointed at the copy - these are the tensors named `##`. + +Copies are issued asynchronously where the backends support it, with events ordering the consumer +behind the producer. + +## Allocation and graph reuse + +Compute buffers are sized by `ggml_backend_sched_reserve()` with a worst case graph, so that later +graphs of the same shape need no reallocation. Within a graph the allocator reuses memory: a +tensor's memory is handed to a later tensor once its last consumer *in graph order* has run. + +A caller that rebuilds an identical graph can skip the split and the allocation entirely and reuse +the previous one - llama.cpp does this whenever the graph topology is unchanged. This is fast, but +it means neither the split nor the allocator runs, so anything that normally happens there does not +happen on that path. The mechanisms below exist because of that. + +## Graph input ring buffer + +Graph inputs are assigned to the CPU backend, but the caller may pass a **device host buffer type** +for that slot - pinned memory owned by a GPU. Some devices, integrated GPUs in particular, accept +that buffer type for compute. When that happens no input copy is created and the device reads +exactly the memory the host thread writes. + +Nothing then stops the host writing the next iteration's inputs while the device is still reading +the previous ones. The writes are a plain `memcpy` on the calling thread; they are not ordered +against an in-flight `ggml_backend_graph_compute_async`. + +The scheduler detects this in `ggml_backend_sched_new()` by testing the condition that elides the +copy - a host buffer type in the last slot that some other backend accepts - rather than by +identifying particular devices, and requiring that backend to be asynchronous. Backends that +deliberately refuse to compute on pinned host memory are unaffected, and so are synchronous ones +such as BLAS - their work is finished before `graph_compute` returns, so nothing of theirs can +still be reading. + +When it holds, each graph input gets a **ring of buffers** instead of one, and the scheduler moves +the inputs onto the next slot rather than waiting for the device. The slot being stepped onto was +last read `n_copies - 1` iterations ago, so the wait for it is normally already satisfied. + +The extra buffers are not free: one additional copy of every graph input per extra slot, which for +attention masks scales with context and batch size. The default depth is 2 - enough to remove the +wait at the smallest cost that does. + +**Rotation only happens when a compute may still be in flight.** Anything that waits for the +backends clears that state, so a caller that synchronizes between iterations - reading logits in a +sampling loop, say - keeps its inputs at fixed addresses and never pays for the ring. Callers that +issue several computes without synchronizing rotate between them. + +Two obligations come with this: + +- **Callers** must call `ggml_backend_sched_prepare_inputs()` before writing the inputs of a graph that is being reused. That path allocates nothing and so cannot rotate on its own. The function has an effect only when the scheduler has multiple copies and a compute may still be in flight; it is a no-op for single-copy schedulers, after `ggml_backend_sched_alloc_graph()`, and after synchronization. + +- **Backends** that cache work against a graph, such as the CUDA graph cache, may treat an + unchanged `ggml_cgraph::uid` as a promise that nothing the cached work captured has moved, and + skip re-reading the addresses. Rotating moves tensor addresses without re-splitting, so the + scheduler re-stamps the split uids whenever it re-points inputs. A backend keeping such a cache + must key it on the uid or re-check the addresses itself. Violating this does not degrade + gracefully: the cached work replays against stale addresses and the results are quietly wrong. + +## Output lifetime under asynchronous execution + +The scheduler ring versions graph inputs, not graph outputs. `GGML_TENSOR_FLAG_OUTPUT` keeps a tensor's allocation alive until the end of its graph; it does not create one result allocation per asynchronous submission. Reusing a graph can therefore write its output tensors again. + +This does not require an output ring when each output is consumed in submission order. Operations submitted to one backend are ordered, so a copy or downstream consumer queued after one graph runs before a later graph on that backend can overwrite the same output allocation. Internal backend streams must join or otherwise preserve that order before subsequent operations use the result. + +llama.cpp queues every required output readback immediately after each micro-batch and before it submits the next one. Each readback writes to that micro-batch's rows in the persistent host output buffer. The terminal backend therefore sees: + +``` +compute A -> readback A -> compute B -> readback B +``` + +Pipeline stages can overlap on different backends, but micro-batch B cannot overtake micro-batch A on the terminal backend. Different execution times do not change the queue order. If asynchronous readback is unsupported, the backend API synchronizes before copying instead. The public llama.cpp result accessors also synchronize before exposing host result pointers. + +This differs from graph input preparation: the host writes the next inputs directly, outside the backend queue, so no ordering edge protects an input allocation without synchronization or rotation. + +A scheduler caller that needs every asynchronous result must enqueue an ordered copy or consumer before another submission can overwrite the output, or synchronize before reusing the graph. If several results must remain resident on a device at the same time, the caller must provide distinct output storage. `n_copies` does not provide output snapshots. + +## Memory read in place by another backend + +The allocator frees a tensor's memory once its last consumer in graph order has run. That is only +sound if the consumers have finished. When a split reads a tensor in place out of a buffer another +backend owns, the reading split is asynchronous and may still be reading well after graph order +says the memory is dead - so a later split on the owning backend can be given the same memory and +overwrite it. + +Such tensors are pinned for the lifetime of the graph with `ggml_gallocr_pin_tensor()`, preventing both ordinary reuse after their last graph-order consumer and in-place reuse by a child operation. The pin is keyed on the **view root**, since that is the tensor that owns the memory; the allocator resolves candidate in-place parents to the same root. Keying the pin on the accessed tensor would miss every access through a view. + +Pins are recomputed when the scheduler splits a graph and are encoded in the resulting allocation, so graph reuse preserves the protected addresses without rerunning the allocator. They are separate from `GGML_TENSOR_FLAG_OUTPUT`: marking a tensor as an output would also extend its allocator lifetime, but it changes tensor semantics and can disable backend optimizations unrelated to the cross-backend read. + +## Ops that alias their source + +Some ops write through to one of their sources rather than to fresh memory: `ggml_set()` and +friends, where the result is a view of `src0`. Such a node must run on the backend where that +source lives. + +If it does not, the split logic sees an operand on another backend, substitutes a copy for it, and +the op writes into the copy. The copy is never written back, so the update is silently discarded - +the graph still computes, every op still runs, and the result is wrong. On a recurrent model this +shows up as state that stops being carried between tokens, and generation degenerates into a +repeated token. + +The scheduler therefore assigns any node that aliases its source, and is not a pure view op, to +the backend of the tensor it aliases. + +That move is only made when the target backend supports the op. Support can be conditional on +the tensor types - CUDA runs `GGML_OP_SET` only for F32 and I32 - so the backend owning the +aliased memory is not guaranteed to be able to run the op writing into it. There is no correct +placement in that case: the scheduler copies operands into a split, never results out of one, so +whichever backend runs the op, the write cannot reach the aliased memory. The node is left where +the earlier passes put it, which is what the scheduler did before this rule existed, and the +reason is logged under `GGML_SCHED_DEBUG`. + +## The scheduler sanitizer + +The above are ordering rules, and ordering bugs do not announce themselves - they produce wrong +numbers on some runs and not others. The sanitizer is a happens-before checker for the backend +API, built into `ggml-base` and enabled at runtime. + +It maintains a vector clock per actor - the host thread and each backend - and a shadow map of +which actor last read or wrote every byte of every buffer. Synchronization points (backend +synchronize, event record, event wait, event synchronize, async copies) advance those clocks. When +an access conflicts with a previous one and no happens-before edge connects them, it reports: + +``` +ggml-sched-sanitize: RACE (write-after-read) on ROCm_Host[0, 2048) +ggml-sched-sanitize: read ROCm0#2 @3 split 0 inp_tokens (compute) +ggml-sched-sanitize: write HOST @15 split -1 inp_tokens (tensor_set) +ggml-sched-sanitize: no happens-before edge: HOST knows ROCm0#2@2, needs >=3 +``` + +Where the access is through a view, the report names both the accessed tensor and the allocation +that owns the memory, as `accessed <- root`. The two are frequently unrelated, and the root is the +one that matters. + +**It only sees what goes through the backend API.** A caller that writes a host-visible tensor by +taking `tensor->data` and memcpy'ing into it is invisible to it, and races against such writes are +missed entirely. Callers doing that must announce it with `ggml_backend_tensor_set_direct()`, +which is a no-op unless the sanitizer is enabled. In llama.cpp every such site goes through +`llama_host_write()`. + +By default the first race aborts. `GGML_SCHED_SANITIZE_NONFATAL=1` reports each one and continues, +so a single run enumerates a whole workload; a count is printed at exit. + +## Environment variables + +| variable | effect | +|---|---| +| `GGML_SCHED_DEBUG` | `1` prints assignments and scheduler decisions, `2` adds per-node detail | +| `GGML_SCHED_DEBUG_REALLOC` | report, or abort on, unexpected graph reallocations | +| `GGML_SCHED_UMA_RING` | when the ring is in use: `0`/`1` disables it, larger sets the depth. cannot enable it where it was not detected | +| `GGML_SCHED_PIN_ASYNC_READS` | `0` disables pinning of memory read in place by another backend | +| `GGML_SCHED_SANITIZE` | `1` enables the sanitizer, `2` also traces synchronization edges | +| `GGML_SCHED_SANITIZE_NONFATAL` | `1` reports every race instead of aborting on the first | + +## Example usage + +```c +// operations that use tensors allocated in a buffer with USAGE_WEIGHTS will be assigned +// preferably to run on the same backend as the buffer +ggml_backend_buffer_set_usage(buf_weights, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); + +sched = ggml_backend_sched_new({backend_gpu, backend_gpu2, backend_cpu}, NULL, num_backends, + GGML_DEFAULT_GRAPH_SIZE, false, true); + +// initialize buffers from a max size graph (optional) +reserve_graph = build_graph(sched, max_batch_size); + +// manually assign nodes to a backend (optional, should not be needed in most cases) +struct ggml_tensor * node = ggml_mul_mat(ctx, ...); +ggml_backend_sched_set_tensor_backend(sched, node, backend_gpu); + +ggml_backend_sched_reserve(sched, reserve_graph); + +// compute +// the graph and its tensors are single-use in terms of allocation, multi-use in terms of computation +graph = build_graph(sched); +for (int i = 0; i < 10; ++i) { + ggml_backend_sched_graph_compute(sched, graph); // on the first iteration the graph is allocated automatically +} + +// if there are graph inputs: +graph = build_graph(sched); // a new graph that is not allocated +ggml_backend_sched_reset(sched); // clear the allocation of the previous graph +ggml_backend_sched_alloc_graph(sched, graph); +ggml_backend_tensor_set(input_tensor, ...); +ggml_backend_sched_graph_compute(sched, graph); + +// when reusing an already allocated graph rather than allocating a new one, the inputs must be +// made safe to write first - see "Graph input ring buffer" +ggml_backend_sched_prepare_inputs(sched); +ggml_backend_tensor_set(input_tensor, ...); +ggml_backend_sched_graph_compute(sched, graph); +``` + +An alternative to the above is to assign the inputs to a dedicated context and allocate them +statically with `ggml_backend_alloc_ctx_tensors()`. diff --git a/docs/development/rocm-mmb-context.md b/docs/development/rocm-mmb-context.md new file mode 100644 index 000000000000..c411b0098423 --- /dev/null +++ b/docs/development/rocm-mmb-context.md @@ -0,0 +1,53 @@ +# ROCm MMB context ownership + +The gfx1151 MMB kernels used process-global activation scratch, BF16 marks, graph fingerprints, and converted-weight caches. Target and MTP draft contexts have independent asynchronous streams. Draft catch-up could overwrite storage still used by the target, including storage referenced by captured HIP graphs. Freeing either backend also released the other backend's cached storage. + +Keep this state in its backend context. Scratch allocations and temporary conversions belong to the active stream's pool. Producer slots can be read across streams through graph dependencies. Converted weights and BF16 marks belong to the context. At context destruction, wait for its queued work, destroy its captured graphs, then release its MMB allocations before destroying the pools. No target/draft synchronization is added to inference. + +Sharing immutable token embeddings or output projection weights is not itself the cause. The unsafe sharing was in the backend's mutable temporary storage and its lifetime management. + +## Regression coverage + +Run `test-backend-ops test -b ROCm0 -o MMB_CONTEXT` on gfx1151. The case uses two independent schedulers and the existing HC_CHAIN graph, with 512 tokens each. It preserves graph inputs across replay, checks stable serialized references, warms captured graphs, overlaps target and draft 16 times, and checks the surviving target after draft destruction. Every output must be finite and byte-identical to its serialized reference. + +The final test fails against unchanged backend revision `8c1c282ecb194e8f02613defcc4a07c22b6d1c08` and passes against the fix, including with `GGML_CUDA_DISABLE_GRAPHS=1`. During test development, an initial version let the allocator overwrite inputs between replays; that version was not a valid race oracle. The final test protects those inputs and was rerun against the unchanged libraries. + +## Model validation + +Hardware: Ryzen AI Max+ PRO 395, Radeon 8060S/gfx1151, 128 GiB unified memory, balanced power profile. ROCm 10.0.0-2, HIP Clang 23, GCC 16.2.1, Release build, native CPU and gfx1151 GPU targets, HIP graphs enabled. + +Target: AesSedai Qwen3.8-Flash-Next Q4_K_M, snapshot `b90616d610f47a124fe246b70bb136aeeb25a2ad`, five GGUF shards. Draft: `mtp-Qwen3.8-Flash-Next-Q4_K_M.gguf`. Model hashes and corpus hashes are in `/tmp/strix-mtp-investigation/{model,corpus}-sha256.txt` on the test machine. + +A diagnostic probe uses the production draft catch-up hook, fixed token IDs, FA enabled, full ROCm0 offload, batch/ubatch 2048, context 32768, lazy-mode on-direct, and load-mode dio. It saves complete final-prefill logits and 32 fixed-token continuation logit vectors, each with 248320 vocabulary entries. The reference is the unchanged fork with target-only MTP hidden extraction enabled. This isolates target/draft interference without sampling differences. + +| Corpus | Prefill tokens | Token IDs, all logits, and top-1 | +| --- | ---: | --- | +| Code | 4096 | Byte-identical | +| Prose | 4096 | Byte-identical | +| Structured | 4096 | Byte-identical | +| Numeric | 4096 | Byte-identical | +| Code | 26776 | Byte-identical | + +The HC_CHAIN and GATED_DELTA_NET focused suite passed 56/56 cases. The complete ROCm backend suite passed 29695/29695 cases, followed by the new MMB_CONTEXT regression. Unsupported operations were skipped by the existing test harness. + +## Memory and performance + +Context ownership can increase memory use because target and draft retain separate scratch and converted-weight caches. It does not duplicate the entire model or KV cache. Converted-weight initialization can also be repeated per context. The default inference path remains asynchronous; the added wait is at context destruction. + +For the same 4096-token code probe with synchronized draft catch-up in both builds, intercepting hipMalloc/hipFree measured peak live allocations of 87,896,241,152 bytes before and 87,978,816,512 bytes after: an increase of 78.75 MiB. Both ended with zero tracked live bytes. This is one workload and measures HIP allocation requests, not total physical memory, registered host mappings, or allocator-internal overhead. It is not a maximum for other batch sizes or contexts. No throughput or latency conclusion is established by this allocation measurement. + +## Upstream applicability + +Checked upstream ggml-org/llama.cpp revision `53ed051ce5e8193652e449f43216ca3859454f49` on 2026-09-24. Its complete source tree has neither MMB nor hyperconnection kernel files. Its CUDA/HIP allocation pools already belong to each backend context and stream. Upstream also supports asynchronous MTP draft catch-up, but does not contain the process-global MMB state responsible for this failure. + +The affected file entered this fork in commit `0636c9aee4a2ab1d38bb7e8c038227c6634c3857` (PR #63). This specific bug and fix are fork-specific. This source inspection does not certify upstream MTP against unrelated defects; no upstream runtime comparison was performed. + +Sources: [upstream backend](https://github.com/ggml-org/llama.cpp/tree/53ed051ce5e8193652e449f43216ca3859454f49/ggml/src/ggml-cuda), [context-owned pools](https://github.com/ggml-org/llama.cpp/blob/53ed051ce5e8193652e449f43216ca3859454f49/ggml/src/ggml-cuda/common.cuh#L1559), [fork introduction](https://github.com/halo-box/strix-llama.cpp/commit/0636c9aee4a2ab1d38bb7e8c038227c6634c3857). + +## Limits and artifacts + +Status: TARGET PASS for the ROCm scratch-race regression and model interference checks; INCOMPLETE for the broader project certification matrix described below. + +Raw logs, probe outputs, commands, source diff, allocation tracer, and upstream source snapshots are under `/tmp/strix-mmb-fix` on the test machine. The earlier reproduction and probe source are under `/tmp/strix-mtp-investigation`. + +The model probe exercises catch-up and deterministic target decode, not the complete server sampling/acceptance loop or the original browser conversation. Multi-device configurations, NVIDIA CUDA, the optional graph-fork optimization, and a complete throughput/acceptance benchmark matrix are not certified. The separate rollback-validity and Vulkan attention issues are not changed by this fix. diff --git a/docs/development/vulkan-sparse-fa-deferred.md b/docs/development/vulkan-sparse-fa-deferred.md new file mode 100644 index 000000000000..caea367c874f --- /dev/null +++ b/docs/development/vulkan-sparse-fa-deferred.md @@ -0,0 +1,81 @@ +# Deferred: Vulkan sparse Flash Attention (Strix Halo) + +Written during the 2026-09-15 sync of `halo-box/llama.cpp` into this fork (PR #64). +Upstream's sparse Flash Attention could not be merged mechanically and was deferred. +This file says why, and what integrating it actually requires, so the next person does +not have to rediscover it. + +Nothing below is a measurement. + +--- + +## What is not in this fork + +`fc82583e6 vulkan: support sparse Flash Attention (ggml-org#28105)` is **not** here. +`ggml/src/ggml-vulkan/` was reverted wholesale to the strix version during the sync. + +Deferred with it: `f1e44dcc1 vulkan: workaround NV queuesubmit driver bug +(ggml-org#28830)`. + +## Why it could not be merged mechanically + +Upstream and strix collide on three things at once: + +| | strix | upstream #28105 | +|---|---|---| +| Spec-constant bit | `DYNAMIC_KV = (Flags & 16)` | `USE_SPARSE = (Flags & 16)` | +| FA push constants | `m_row_len, gqa_ratio, split_k_num, output_k_num, partial_output` | `sparse_base` | +| Sparse mechanism | gather/compaction family + DeepSeek-V4 split keyed on bit 31 of `gqa_ratio` | binding-7 index gather, `fa_kv_index()` | + +They are rival implementations of overlapping functionality on the same bit. + +Taking either side alone breaks the build. Upstream's sparse code merged *cleanly* into +`flash_attn.comp`, `flash_attn_cm1.comp` and `flash_attn_cm2.comp` — strix had not +touched those regions — so keeping strix's `flash_attn_base.glsl` leaves those three +files calling an undefined `fa_kv_index()` / `USE_SPARSE` / `data_sparse`. Reverting the +whole directory is the only self-consistent option without a GPU. + +## What integration actually requires + +1. Renumber one feature's `Flags` bit. Audit every `string_to_spv` call in + `vulkan-shaders-gen.cpp` and every `Flags` write in `ggml-vulkan.cpp` — the bit is set + on the C++ side and read in GLSL, with no shared header to keep the two honest. +2. Merge the two push-constant structs and update the single C++ writer. + **A mismatched layout does not crash.** The shader reads whatever bytes sit at those + offsets and emits plausible-looking wrong tensors. Validate numerically, not by + "it ran". +3. Decide whether `DYNAMIC_KV` and `USE_SPARSE` coexist or one supersedes. strix's gather + family (`flash_attn_gather{,_dq,_union,_union_dq}.comp`) may already cover upstream's + case, in which case the answer is to drop one. + +## How to settle it + +```sh +cmake -B build-vk -DGGML_VULKAN=ON && cmake --build build-vk -j +build-vk/bin/test-backend-ops -b Vulkan0 -o FLASH_ATTN_EXT # correctness vs CPU ref +build-vk/bin/test-backend-ops -b Vulkan0 perf -o FLASH_ATTN_EXT # includes the strix probes +``` + +The strix probes already in `test-backend-ops.cpp` are the ones that matter here: +MALL-spill (32 KV heads, K/V footprint 8x > 32MB MALL), L2-residence (1 KV head x GQA 32, +K/V 5.2MB fits L2), and the KV-cache-layout probe (`{0, 2, 1, 3}`). They were written to +distinguish cache-bandwidth-bound from issue-bound on this part. If upstream's sparse path +is integrated, it must not regress them. + +Baseline, measured on a Radeon 8060S (RADV STRIX_HALO) at `dffbb7888`: FLASH_ATTN_EXT +5339/5339 pass on Vulkan0. + +--- + +## Provenance + +| Ref | What | +|---|---| +| `fc82583e6` | upstream sparse FA (ggml-org#28105) — deferred | +| `f1e44dcc1` | upstream NV queuesubmit workaround (ggml-org#28830) — deferred | +| `dced43c58` | the sync merge (PR #64) | +| `82aed6017` | #63 merged into that sync | + +The PLE n-gram table's disk reader survived the sync as the backend of +`--lazy-mode on-direct`; see PR #64's description for what changed and what was dropped +(O_DIRECT reads and the `--ngram-*` tuning knobs). diff --git a/docs/speculative-prefill.md b/docs/speculative-prefill.md new file mode 100644 index 000000000000..952d374a1634 --- /dev/null +++ b/docs/speculative-prefill.md @@ -0,0 +1,100 @@ +# Speculative prefill + +Port of upstream PR [ggml-org/llama.cpp#27692](https://github.com/ggml-org/llama.cpp/pull/27692), +which implements *Speculative Prefill: Turbocharging TTFT with Lightweight and Training-Free +Token Importance Estimation* ([arXiv:2502.02789](https://arxiv.org/abs/2502.02789)). + +A small draft model prefills the prompt and decodes a few lookahead tokens. The softmax +attention of those tokens over the prompt is captured with a ggml eval callback, smoothed, +max-reduced over heads and layers, and pooled into chunks. Only the top scoring chunks are +fed to the target model. This is lossy: tokens that are dropped are gone. + +Ported: `common/speculative-prefill.{h,cpp}`, `llama_set_eval_callback`, the `--spec-prefill-*` +options, `examples/speculative-prefill`, and the llama-server hookup. Not ported: the +llama-bench changes and the python eval scripts from that PR. + +## Usage + +```bash +llama-server -m target.gguf -mpd draft.gguf --spec-prefill-percentage 0.3 +``` + +The estimator must be a separate small model with the same vocab as the target. MTP, DFlash2, +DSpark and Eagle3 heads are target-dependent and the server refuses to reuse them for this. +The draft context is forced to flash-attn off, because `kq_soft_max` does not exist on the +flash-attn path. If no attention is captured the full prompt is kept. + +## Results + +Ryzen AI MAX+ 395 / Radeon 8060S, Vulkan, 119 GiB unified memory. Estimator is +Qwen3.5-2B-UD-Q4_K_XL in every run. "eff PP" is original prompt tokens divided by TTFT. + +### TTFT vs context length + +Qwen3.8-27B-UD-Q4_K_XL, no decode speculation, median of 3 runs. A passcode was planted at +mid-depth and retrieved in 3/3 runs in every cell. + +| N tokens | p = off | p = 0.30 | p = 0.15 | +| ---: | ---: | ---: | ---: | +| 1,498 | 6,225 ms | 2,647 ms (2.35x) | 1,516 ms (4.11x) | +| 3,131 | 12,992 ms | 5,199 ms (2.50x) | 3,084 ms (4.21x) | +| 6,212 | 26,494 ms | 9,906 ms (2.67x) | 6,113 ms (4.33x) | +| 12,347 | 54,927 ms | 19,557 ms (2.81x) | 11,896 ms (4.62x) | + +Baseline eff PP is nearly flat over that range (241 -> 225 tok/s), so prefill here is close to +linear in N rather than quadratic. The speedup still grows with N, but because the fixed draft +cost amortizes, not because there is super-linear cost to remove. + +Decode is unaffected: 11.14 to 11.62 tok/s across all 12 cells. + +### With decode speculation + +Qwen3.8-27B-UD-Q4_K_XL, 1,869-token prompt, single run each. + +| decode spec | p | TTFT | eff PP | tg | +| --- | ---: | ---: | ---: | ---: | +| none | off | 5,604 ms | 334 tok/s | 11.56 tok/s | +| none | 0.15 | 1,587 ms | 1,180 tok/s | 11.66 tok/s | +| MTP n=4 | off | 5,755 ms | 325 tok/s | 19.47 tok/s | +| MTP n=4 | 0.30 | 2,440 ms | 768 tok/s | 29.33 tok/s | +| MTP n=4 | 0.15 | 1,636 ms | 1,145 tok/s | 23.61 tok/s | +| DFlash2 n=7 | off | 6,576 ms | 284 tok/s | 19.33 tok/s | +| DFlash2 n=7 | 0.30 | 2,724 ms | 688 tok/s | 30.44 tok/s | +| DFlash2 n=7 | 0.15 | 1,772 ms | 1,057 tok/s | 20.90 tok/s | + +The two mechanisms are independent. DFlash2 gains the most because its head also has to prefill +the full prompt, and speculative prefill shrinks that too. + +### Qwen3.8-Flash-Next + +FP4 ple16 (87 GiB) with the Q8_0 MTP head, 1,869-token prompt, median of 3, needle 3/3. + +| p | TTFT | eff PP | tg | MTP accept | +| ---: | ---: | ---: | ---: | ---: | +| off | 4,920 ms | 380 tok/s | 29.38 tok/s | 71% | +| 0.30 | 2,441 ms | 767 tok/s | 28.53 tok/s | 66% | +| 0.15 | 1,645 ms | 1,139 tok/s | 28.87 tok/s | 69% | + +Flash-Next gains less than the dense 27B because it is hybrid: most layers are linear-attention, +so baseline prefill is already closer to linear in N. MTP acceptance and decode rate hold across +keep ratios. + +## Caveats + +- Greedy (temperature 0) output is not reproducible on this box even with no speculation and no + speculative prefill: 3 identical requests to the plain target gave 2 different completions. + Only TTFT is worth quoting from these runs. Cross-arm `tg` and draft-acceptance differences + are confounded, because each arm generates different text under `ignore_eos`. +- `--spec-prefill-ctx`, `--spec-prefill-device` and `--spec-prefill-ngl` set `enabled = true` as + a side effect, so passing only a context size makes the server exit with "speculative prefill + enabled but no draft model was provided". +- The example renumbers kept tokens to positions 0..K-1 instead of keeping their original + positions. The reference implementation keeps the original ones. +- `t_draft_eval_us` covers the draft prompt prefill only, not the lookahead decode steps. About + 150 ms of the reported TTFT is unattributed, because the eval callback forces node-by-node + execution. +- Quality falls off when the token budget cannot hold all the non-redundant evidence. p = 0.30 is + a reasonable default; p = 0.15 still retrieved a needle and summed 3 facts spread across the + context; below p = 0.10 is risky for prompts with distributed evidence. + +`scripts/bench-spec-prefill.sh` reproduces the server measurements. diff --git a/docs/speculative.md b/docs/speculative.md index ffb1e34c7fbf..465242525668 100644 --- a/docs/speculative.md +++ b/docs/speculative.md @@ -428,7 +428,7 @@ statistics draft: #calls = 10, #gen drafts = 10, #acc drafts = 10, #gen tokens = ``` draft acceptance rate = 0.70312 ( 90 accepted / 128 generated) -statistics ngram_mod: #calls = 810, #gen drafts = 15, #acc drafts = 15, #gen tokens = 960, #acc tokens = 730, dur(b,g,a) = 0.149, 0.347, 0.005 ms +statistics ngram-mod: #calls = 810, #gen drafts = 15, #acc drafts = 15, #gen tokens = 960, #acc tokens = 730, dur(b,g,a) = 0.149, 0.347, 0.005 ms ``` ``` diff --git a/examples/CMakeLists.txt b/examples/CMakeLists.txt index 39f802d250e1..f023b4bfb24a 100644 --- a/examples/CMakeLists.txt +++ b/examples/CMakeLists.txt @@ -30,6 +30,7 @@ else() add_subdirectory(simple) add_subdirectory(simple-chat) add_subdirectory(speculative) + add_subdirectory(speculative-prefill) add_subdirectory(speculative-simple) add_subdirectory(gen-docs) add_subdirectory(training) diff --git a/examples/speculative-prefill/CMakeLists.txt b/examples/speculative-prefill/CMakeLists.txt new file mode 100644 index 000000000000..9bfdec3148b9 --- /dev/null +++ b/examples/speculative-prefill/CMakeLists.txt @@ -0,0 +1,5 @@ +set(TARGET llama-speculative-prefill) +add_executable(${TARGET} speculative-prefill.cpp) +install(TARGETS ${TARGET} RUNTIME) +target_link_libraries(${TARGET} PRIVATE llama-common llama ${CMAKE_THREAD_LIBS_INIT}) +target_compile_features(${TARGET} PRIVATE cxx_std_17) diff --git a/examples/speculative-prefill/speculative-prefill.cpp b/examples/speculative-prefill/speculative-prefill.cpp new file mode 100644 index 000000000000..1a7ae6560464 --- /dev/null +++ b/examples/speculative-prefill/speculative-prefill.cpp @@ -0,0 +1,270 @@ +#include "arg.h" +#include "common.h" +#include "log.h" +#include "llama.h" +#include "sampling.h" +#include "speculative.h" +#include "speculative-prefill.h" +#include "../../src/llama-ext.h" + +#include +#include +#include +#include +#include +#include +#include + +#define SPEC_VOCAB_MAX_SIZE_DIFFERENCE 128 + +int main(int argc, char ** argv) { + std::setlocale(LC_NUMERIC, "C"); + + common_params params; + + common_init(); + + if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_SPECULATIVE)) { + return 1; + } + + common_params_model dft_model = params.speculative.prefill.model; + if (dft_model.empty()) { + dft_model = params.speculative.draft.mparams; + } + + if (dft_model.empty()) { + LOG_ERR("%s: draft model is required for speculative prefill (specify with -mpd, --spec-prefill-model, or -md)\n", __func__); + return 1; + } + + // force speculative prefill enabled + params.speculative.prefill.enabled = true; + + llama_backend_init(); + llama_numa_init(params.numa); + + // load target model + LOG_INF("%s: loading target model...\n", __func__); + auto init_tgt = common_init_from_params(params); + if (!init_tgt) { + LOG_ERR("%s: failed to load target model\n", __func__); + return 1; + } + + llama_model * model_tgt = init_tgt->model(); + llama_context * ctx_tgt = init_tgt->context(); + + if (llama_model_is_recurrent(model_tgt)) { + LOG_ERR("%s: speculative prefill is not supported for recurrent models\n", __func__); + return 1; + } + + // load draft model with standard attention to allow attention extraction + LOG_INF("%s: loading draft model...\n", __func__); + common_params params_dft = common_base_params_to_speculative(params); + params_dft.model = dft_model; + if (params.speculative.prefill.n_ctx > 0) { + params_dft.n_ctx = params.speculative.prefill.n_ctx; + } + if (params.speculative.prefill.n_gpu_layers != -1) { + params_dft.n_gpu_layers = params.speculative.prefill.n_gpu_layers; + } + if (!params.speculative.prefill.devices.empty()) { + params_dft.devices = params.speculative.prefill.devices; + } else if (!params.speculative.draft.devices.empty()) { + params_dft.devices = params.speculative.draft.devices; + } + params_dft.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED; // standard attention needed to capture kq_soft_max weights + + auto init_dft = common_init_from_params(params_dft, /*model_only=*/true); + if (!init_dft) { + LOG_ERR("%s: failed to load draft model\n", __func__); + return 1; + } + + llama_model * model_dft = init_dft->model(); + if (!model_dft) { + LOG_ERR("%s: failed to load draft model\n", __func__); + return 1; + } + + if (llama_model_is_recurrent(model_dft)) { + LOG_ERR("%s: draft model is recurrent and not supported for speculative prefill\n", __func__); + return 1; + } + + if (llama_model_target_layer_ids_n(model_dft) > 0) { + LOG_ERR("%s: draft model '%s' is target-dependent and cannot be used for speculative prefill\n", __func__, dft_model.path.c_str()); + return 1; + } + + if (params.speculative.prefill.n_ctx <= 0 && params_dft.n_ctx > (int32_t) llama_model_n_ctx_train(model_dft)) { + params_dft.n_ctx = llama_model_n_ctx_train(model_dft); + LOG_INF("%s: capping speculative prefill draft context to training limit (%d tokens)\n", __func__, params_dft.n_ctx); + } + + llama_context_params cparams_dft = common_context_params_to_llama(params_dft); + llama_context_ptr ctx_dft_own(llama_init_from_model(model_dft, cparams_dft)); + llama_context * ctx_dft = ctx_dft_own.get(); + if (!ctx_dft) { + LOG_ERR("%s: failed to create draft context\n", __func__); + return 1; + } + + const llama_vocab * vocab_tgt = llama_model_get_vocab(model_tgt); + const llama_vocab * vocab_dft = llama_model_get_vocab(model_dft); + + if (llama_vocab_get_add_bos(vocab_tgt) != llama_vocab_get_add_bos(vocab_dft) || + (llama_vocab_get_add_bos(vocab_tgt) && llama_vocab_bos(vocab_tgt) != llama_vocab_bos(vocab_dft))) { + LOG_ERR("%s: draft model bos tokens must match target model. add: %d - %d, id: %d - %d\n", + __func__, + llama_vocab_get_add_bos(vocab_tgt), llama_vocab_get_add_bos(vocab_dft), + llama_vocab_bos(vocab_tgt), llama_vocab_bos(vocab_dft)); + return 1; + } + + { + const int n_vocab_tgt = llama_vocab_n_tokens(vocab_tgt); + const int n_vocab_dft = llama_vocab_n_tokens(vocab_dft); + const int vocab_diff = n_vocab_tgt > n_vocab_dft ? n_vocab_tgt - n_vocab_dft : n_vocab_dft - n_vocab_tgt; + if (vocab_diff > SPEC_VOCAB_MAX_SIZE_DIFFERENCE) { + LOG_ERR("%s: target vocab size %d does not match draft vocab size %d - difference %d, max allowed %d\n", + __func__, n_vocab_tgt, n_vocab_dft, vocab_diff, SPEC_VOCAB_MAX_SIZE_DIFFERENCE); + return 1; + } + } + + // tokenize prompt + std::vector prompt_tokens = common_tokenize(ctx_tgt, params.prompt, true, true); + const int32_t n_prompt = (int32_t) prompt_tokens.size(); + + if (n_prompt == 0) { + LOG_ERR("%s: empty prompt provided\n", __func__); + return 1; + } + + LOG_INF("%s: prompt tokens = %d\n", __func__, n_prompt); + LOG_INF("%s: speculative prefill config: percentage = %.2f, chunk_size = %d, lookahead = %d, pool_kernel = %d\n", + __func__, + (double) params.speculative.prefill.percentage, + params.speculative.prefill.chunk_size, + params.speculative.prefill.look_ahead_cnt, + params.speculative.prefill.pool_kernel_size); + + llama_seq_id seq_id = 0; + + // initialize draft sampler for lookahead steps + common_params_sampling sparams_dft = params.sampling; + sparams_dft.temp = 0.0f; + common_sampler_ptr smpl_dft(common_sampler_init(model_dft, sparams_dft)); + + // step 1: run speculative prefill on draft model + const auto t_spec_prefill_start = ggml_time_us(); + + common_speculative_prefill_result spec_res = common_speculative_prefill_execute( + ctx_dft, + smpl_dft.get(), + prompt_tokens, + seq_id, + params.speculative.prefill); + + LOG_INF("%s: speculative prefill kept %d / %d tokens (%.1f%%)\n", + __func__, + spec_res.n_prompt_kept, + spec_res.n_prompt_orig, + 100.0f * (float) spec_res.n_prompt_kept / (float) spec_res.n_prompt_orig); + LOG_INF("%s: draft eval time: %.2f ms, importance estimation time: %.2f ms\n", + __func__, + spec_res.t_draft_eval_us / 1000.0f, + spec_res.t_estimate_us / 1000.0f); + + // step 2: sparse target model prefill + const auto t_tgt_prefill_start = ggml_time_us(); + + const int32_t n_batch_tgt = llama_n_batch(ctx_tgt); + const int32_t n_kept_total = (int32_t) spec_res.kept_indices.size(); + llama_batch batch_tgt = llama_batch_init(std::min(n_kept_total, n_batch_tgt), 0, 1); + + int ret = 0; + for (int32_t i = 0; i < n_kept_total; i += n_batch_tgt) { + const int32_t n_eval = std::min(n_kept_total - i, n_batch_tgt); + common_batch_clear(batch_tgt); + + for (int32_t j = 0; j < n_eval; ++j) { + const int32_t k = i + j; + const int32_t orig_idx = spec_res.kept_indices[k]; + const bool is_last = (k == n_kept_total - 1); + common_batch_add(batch_tgt, prompt_tokens[orig_idx], (llama_pos) k, { seq_id }, is_last); + } + + ret = llama_decode(ctx_tgt, batch_tgt); + if (ret != 0) { + LOG_ERR("%s: failed to decode sparse prompt on target model, ret = %d\n", __func__, ret); + llama_batch_free(batch_tgt); + return 1; + } + } + llama_batch_free(batch_tgt); + + llama_synchronize(ctx_tgt); + llama_synchronize(ctx_dft); + + const auto t_tgt_prefill_end = ggml_time_us(); + const double ttft_ms = (t_tgt_prefill_end - t_spec_prefill_start) / 1000.0; + const double tgt_prefill_ms = (t_tgt_prefill_end - t_tgt_prefill_start) / 1000.0; + + LOG_INF("%s: sparse target prefill time = %.2f ms\n", __func__, tgt_prefill_ms); + LOG_INF("%s: total Time-To-First-Token (TTFT) = %.2f ms (%.2f effective prompt tokens/s)\n", + __func__, ttft_ms, ttft_ms > 0.0 ? (1000.0 * (double) n_prompt / ttft_ms) : 0.0); + + // step 3: autoregressive generation + common_sampler_ptr smpl_tgt(common_sampler_init(model_tgt, params.sampling)); + + LOG("\n--- Generation Start ---\n"); + + int32_t n_predict = params.n_predict > 0 ? params.n_predict : 32; + int32_t cur_pos = n_kept_total; + int32_t n_generated = 0; + + llama_batch batch_gen = llama_batch_init(1, 0, 1); + + const auto t_gen_start = ggml_time_us(); + + for (int32_t i = 0; i < n_predict; ++i) { + const llama_token token_id = common_sampler_sample(smpl_tgt.get(), ctx_tgt, -1); + common_sampler_accept(smpl_tgt.get(), token_id, true); + + if (llama_vocab_is_eog(vocab_tgt, token_id)) { + break; + } + + const std::string piece = common_token_to_piece(ctx_tgt, token_id); + LOG("%s", piece.c_str()); + fflush(stdout); + + common_batch_clear(batch_gen); + common_batch_add(batch_gen, token_id, (llama_pos) cur_pos++, { seq_id }, true); + + ret = llama_decode(ctx_tgt, batch_gen); + if (ret != 0) { + LOG_ERR("%s: failed to decode generated token %d, ret = %d\n", __func__, i, ret); + break; + } + + n_generated++; + } + + const auto t_gen_end = ggml_time_us(); + const double gen_ms = (t_gen_end - t_gen_start) / 1000.0; + + LOG("\n--- Generation End ---\n\n"); + + LOG_INF("generated %d tokens in %.2f ms (%.2f tokens/s)\n", + n_generated, gen_ms, n_generated > 0 ? (n_generated / (gen_ms / 1000.0)) : 0.0); + + llama_batch_free(batch_gen); + llama_backend_free(); + + return 0; +} diff --git a/ggml/include/ggml-alloc.h b/ggml/include/ggml-alloc.h index a7926a21a9a2..2b5084cb40a3 100644 --- a/ggml/include/ggml-alloc.h +++ b/ggml/include/ggml-alloc.h @@ -73,6 +73,9 @@ GGML_API bool ggml_gallocr_alloc_graph(ggml_gallocr_t galloc, struct ggml_cgraph GGML_API size_t ggml_gallocr_get_buffer_size(ggml_gallocr_t galloc, int buffer_id); +GGML_API void ggml_gallocr_pin_tensor(ggml_gallocr_t galloc, struct ggml_tensor * t); +GGML_API void ggml_gallocr_clear_pins(ggml_gallocr_t galloc); + // Utils // Create a buffer and allocate all the tensors in a ggml_context // ggml_backend_alloc_ctx_tensors_from_buft_size returns the size of the buffer that would be allocated by ggml_backend_alloc_ctx_tensors_from_buft diff --git a/ggml/include/ggml-backend.h b/ggml/include/ggml-backend.h index cc3f8cd36e35..68ec89011242 100644 --- a/ggml/include/ggml-backend.h +++ b/ggml/include/ggml-backend.h @@ -95,6 +95,8 @@ extern "C" { GGML_API void ggml_backend_tensor_get_2d(const struct ggml_tensor * tensor, void * data, size_t offset, size_t size, size_t n_copies, size_t stride_tensor, size_t stride_data); GGML_API void ggml_backend_tensor_memset( struct ggml_tensor * tensor, uint8_t value, size_t offset, size_t size); + GGML_API void ggml_backend_tensor_set_direct(struct ggml_tensor * tensor, size_t offset, size_t size); + GGML_API void ggml_backend_synchronize(ggml_backend_t backend); GGML_API ggml_backend_graph_plan_t ggml_backend_graph_plan_create(ggml_backend_t backend, struct ggml_cgraph * cgraph); @@ -263,46 +265,7 @@ extern "C" { // Backend scheduler // - // The backend scheduler allows for multiple backend devices to be used together - // Handles compute buffer allocation, assignment of tensors to backends, and copying of tensors between backends - // The backends are selected based on: - // - the backend that supports the operation - // - the location of the pre-allocated tensors (e.g. the weights) - /* - Example usage: - - // operations that use tensors allocated in a buffer with USAGE_WEIGHTS will be assigned - // preferably to run on the same backend as the buffer - ggml_backend_buffer_set_usage(buf_weights, GGML_BACKEND_BUFFER_USAGE_WEIGHTS); - - sched = ggml_backend_sched_new({backend_gpu, backend_gpu2, backend_cpu}, NULL, num_backends, GGML_DEFAULT_GRAPH_SIZE, false, true); - - // initialize buffers from a max size graph (optional) - reserve_graph = build_graph(sched, max_batch_size); - - // manually assign nodes to a backend (optional, should not be needed in most cases) - struct ggml_tensor * node = ggml_mul_mat(ctx, ...); - ggml_backend_sched_set_tensor_backend(sched, node, backend_gpu); - - ggml_backend_sched_reserve(sched, reserve_graph); - - // compute - graph = build_graph(sched); // the graph and its tensors are single-use in terms of allocation, multi-use in terms of computation - for (int i = 0; i < 10; ++i) { - ggml_backend_sched_graph_compute(sched, graph); // on the first iteration the graph is allocated automatically - } - - // if there are graph inputs: - graph = build_graph(sched); // get a new graph that is not allocated (the metadata for the old graph is freed once ggml_free is called) - ggml_backend_sched_reset(sched); // clear the allocation of the previous graph - ggml_backend_sched_alloc_graph(sched, graph); // explicitly allocate the new graph but do not execute it - ggml_backend_tensor_set(input_tensor, ...); // copy data to the newly allocated graph tensors - ggml_backend_sched_graph_compute(sched, graph); // execute the graph - - // as an alternative to the above it is also possible to assign the inputs to a dedicated context and - // allocate them statically via ggml_backend_alloc_ctx_tensors - } - */ + // see docs/development/backend-scheduler.md typedef struct ggml_backend_sched * ggml_backend_sched_t; @@ -341,6 +304,7 @@ extern "C" { // Allocate and compute graph on the backend scheduler GGML_API bool ggml_backend_sched_alloc_graph(ggml_backend_sched_t sched, struct ggml_cgraph * graph); // returns success + GGML_API void ggml_backend_sched_prepare_inputs(ggml_backend_sched_t sched); GGML_API enum ggml_status ggml_backend_sched_graph_compute(ggml_backend_sched_t sched, struct ggml_cgraph * graph); GGML_API enum ggml_status ggml_backend_sched_graph_compute_async(ggml_backend_sched_t sched, struct ggml_cgraph * graph); GGML_API void ggml_backend_sched_synchronize(ggml_backend_sched_t sched); diff --git a/ggml/include/ggml.h b/ggml/include/ggml.h index 85a1ae7ae208..d47c06d8df2f 100644 --- a/ggml/include/ggml.h +++ b/ggml/include/ggml.h @@ -430,7 +430,18 @@ extern "C" { GGML_TYPE_NVFP4 = 40, // NVFP4 (4 blocks, E4M3 scale) GGML_TYPE_Q1_0 = 41, GGML_TYPE_Q2_0 = 42, - GGML_TYPE_COUNT = 43, + // ROCmFPx experimental AMD-native formats. These are kept in a high ID + // range so upstream type IDs stay free for future ggml types. + GGML_TYPE_Q4_0_ROCMFP4 = 100, // ROCmFP4 experimental UE4M3 scales + packed AMD FP4 blocks + GGML_TYPE_Q4_0_ROCMFP4_FAST = 101, // ROCmFP4 single-scale speed layout + GGML_TYPE_Q6_0_ROCMFPX = 102, // ROCmFPx experimental 6-bit UE4M3-scale reference layout + GGML_TYPE_Q8_0_ROCMFPX = 103, // ROCmFPx experimental 8-bit UE4M3-scale reference layout + GGML_TYPE_Q3_0_ROCMFPX = 104, // ROCmFPx experimental 3-bit UE4M3-scale reference layout + // 105, 106 are TURBO3_0 / TURBO4_0 and 108 is Q4_0_ROCMI4 in charlie12345/ROCmFPX; + // left unassigned here so a GGUF from that tree is rejected rather than misread. + // Note ciru-ai/ROCmFPX diverges: it puts Q7_0_ROCMFPX at 107 and Q2_0_ROCMFPX at 108. + GGML_TYPE_Q2_0_ROCMFPX = 107, // ROCmFPx experimental 2-bit S40 codebook + dual UE4M3 scales + GGML_TYPE_COUNT = 108, }; // [TAG_GGML_PREC] @@ -454,6 +465,7 @@ extern "C" { enum ggml_op_hint { GGML_HINT_NONE = 0, GGML_HINT_SRC0_IS_HADAMARD = 1, + GGML_HINT_EXACT_BATCH = 2, }; // model file types @@ -486,6 +498,17 @@ extern "C" { GGML_FTYPE_MOSTLY_NVFP4 = 26, // except 1d tensors GGML_FTYPE_MOSTLY_Q1_0 = 27, // except 1d tensors GGML_FTYPE_MOSTLY_Q2_0 = 28, // except 1d tensors + GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4 = 100, // except 1d tensors + GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_LEAN = 101, // ROCmFP4 with Q5_K token embeddings + GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_COHERENT = 102, // ROCmFP4 with Q6_K token embeddings + GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_FAST = 103, // ROCmFP4 single-scale speed layout + GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_FAST_COHERENT = 104, // ROCmFP4 fast with Q6_K token embeddings + GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_STRIX = 105, // ROCmFP4 Strix Halo quality/speed recipe + GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_STRIX_LEAN = 106, // ROCmFP4 Strix Halo size-biased K/V recipe + GGML_FTYPE_MOSTLY_Q6_0_ROCMFPX = 110, // ROCmFPx experimental 6-bit reference layout + GGML_FTYPE_MOSTLY_Q8_0_ROCMFPX = 111, // ROCmFPx experimental 8-bit reference layout + GGML_FTYPE_MOSTLY_Q3_0_ROCMFPX = 112, // ROCmFPx experimental 3-bit reference layout + GGML_FTYPE_MOSTLY_Q2_0_ROCMFPX = 113, // ROCmFPx experimental 2-bit S40 codebook layout }; // available tensor operations: @@ -2512,6 +2535,16 @@ extern "C" { struct ggml_tensor * a, struct ggml_tensor * sinks); + // sparse attention hint: attend only to the first n_kv_raw keys (dense prefix) plus the + // keys selected by top_k. top_k is I32 [n_top_k, n_tokens, 1, n_streams]; each index i + // selects absolute key n_kv_raw + i. Negative or out-of-range indices are ignored. + // Backends may ignore the hint: the kq_mask must still encode the same selection, so a + // dense fallback computes the identical result. + GGML_API void ggml_flash_attn_ext_add_top_k( + struct ggml_tensor * a, + struct ggml_tensor * top_k, + int64_t n_kv_raw); + // TODO: needs to be adapted to ggml_flash_attn_ext GGML_API struct ggml_tensor * ggml_flash_attn_back( struct ggml_context * ctx, diff --git a/ggml/rocmfp4/README.md b/ggml/rocmfp4/README.md new file mode 100644 index 000000000000..bb966f950761 --- /dev/null +++ b/ggml/rocmfp4/README.md @@ -0,0 +1,48 @@ +# ROCmFP4 + +`Q4_0_ROCMFP4` and `Q4_0_ROCMFP4_FAST` are the 4-bit ROCmFPx tensor types. +The format originates in [charlie12345/ROCmFPX](https://github.com/charlie12345/ROCmFPX) +(MIT); the codec here was hand-ported via [ciru-ai/ROCmFPX](https://github.com/ciru-ai/ROCmFPX). + +This directory owns the CPU codec (`rocmfp4.c`): quantize, dequantize, imatrix +weighting and row validation. The rest of ggml registers and dispatches the type. + +## Layout + +32 weights per block. Each weight is a 4-bit code: sign bit plus a 3-bit index +into the magnitude table `0, 1, 2, 3, 4, 6, 8, 10`. The low nibbles of the 16 +payload bytes are weights 0-15, the high nibbles weights 16-31. Scales are +unsigned E4M3 bytes (`e <= 0x7E`, values above are treated as 0) decoded as +`mant * 2^-10` for exponent 0 and `(8 + mant) * 2^(exp - 11)` otherwise. + +| Type | GGUF type id | Payload | Scales | Block bytes | BPW | +|---|---:|---|---|---:|---:| +| `Q4_0_ROCMFP4` | 100 | 16 bytes, 32 nibbles | 2, one per 16-weight half | 18 | 4.50 | +| `Q4_0_ROCMFP4_FAST` | 101 | 16 bytes, 32 nibbles | 1 per block | 17 | 4.25 | + +Type ids match charlie12345/ROCmFPX, so GGUFs quantized there (the published +`*-ROCmFP4-FAST*`, `*-ROCmFP4-STRIX*` files) load here. + +## What is implemented in this tree + +- CPU: quantize / dequantize / vec_dot through the ggml type traits. +- Vulkan: dequant, mat-vec, matmul, and q8_1 integer-dot (MMQ / MMVQ) kernels + under `ggml/src/ggml-vulkan/vulkan-shaders/`. +- Not implemented: ROCm/HIP kernels for these types. The HIP backend does not + know them, so on ROCm they run through the CPU codec. Flash attention does + not accept them as KV-cache types. + +## Recipes in `llama-quantize` + +`Q4_0_ROCMFP4`, `_LEAN` (Q5_K token embeddings), `_COHERENT` (Q6_K token +embeddings), `_FAST`, `_FAST_COHERENT`, `_STRIX` and `_STRIX_LEAN` (Strix Halo +attention K/V recipes). See `tools/quantize/quantize.cpp` and +`src/llama-quant.cpp` for the exact per-tensor routing. + +## Tests + +- `test-quantize-fns` covers both types with the 4-bit error budget. +- `test-backend-ops` includes both in its type lists (MUL_MAT, MUL_MAT_ID, + GET_ROWS, CPY, SET_ROWS); run with `-b Vulkan0` for the Vulkan kernels. +- `gguf-py/tests/test_quants.py --type Q4_0_ROCMFP4` (and `_FAST`) checks the + Python dequantizer against the C one bit for bit. diff --git a/ggml/rocmfp4/rocmfp4.c b/ggml/rocmfp4/rocmfp4.c new file mode 100644 index 000000000000..c957966e16aa --- /dev/null +++ b/ggml/rocmfp4/rocmfp4.c @@ -0,0 +1,805 @@ +// ROCmFPx quant formats, hand-ported into this fork. +// +// Origin: https://github.com/charlie12345/ROCmFPX - creator of the ROCmFP4 format. +// Ported from: https://github.com/ciru-ai/ROCmFPX - a fork of the above. +// +// Both upstream projects are MIT licensed and based on llama.cpp; upstream authors +// retain their authorship and MIT license credit. See LICENSE. + +#define GGML_COMMON_DECL_C +#include "../src/ggml-common.h" + +#include "rocmfp4.h" + +#include +#include +#include + +// ggml-base is compiled architecture-neutral (no -mavx2), so SIMD for the hot +// CPU dot product is enabled per-function via a target attribute plus a runtime +// CPU check. This keeps the AVX2 path in one translation unit without exporting +// internals to ggml-cpu. Non-GNU/non-x86 builds fall back to the scalar loop. +#if defined(__GNUC__) && (defined(__x86_64__) || defined(__i386__)) && !defined(ROCMFP4_NO_AVX2) +#include +#define ROCMFP4_X86_AVX2_DISPATCH 1 +#endif + +// ROCmFP4 stores a signed integer FP4-like codebook at half-scale. It is +// E2M1-derived, but the largest magnitude is retuned from 12 to 10 after +// sampling Qwen3 dense tensors; this reduces outlier pull without changing the +// packed 4-bit layout or integer dot-product path. +static const int8_t rocmfp4_codebook[16] = { + 0, 1, 2, 3, 4, 6, 8, 10, + 0, -1, -2, -3, -4, -6, -8,-10, +}; + +static inline int8_t rocmfp4_decode(uint8_t q) { + q &= 0x0f; + const int mag3 = q & 0x07; + const int mag = mag3 <= 4 ? mag3 : 2*mag3 - 4; + return (q & 0x08) ? -mag : mag; +} + +static inline int8_t rocmfp4_decode_table(uint8_t q) { + return rocmfp4_codebook[q & 0x0f]; +} + +// Finite unsigned E4M3 scale bytes decoded to the half-scale values used by +// ROCmFP4. Keeping this as a table avoids rebuilding identical FP32 values for +// every candidate during exhaustive scale search. +#define ROCMFP4_SCALE_SUB(M) ((M) * 0x1p-10f) +#define ROCMFP4_SCALE_E1(M) ((8 + (M)) * 0x1p-10f) +#define ROCMFP4_SCALE_E2(M) ((8 + (M)) * 0x1p-9f) +#define ROCMFP4_SCALE_E3(M) ((8 + (M)) * 0x1p-8f) +#define ROCMFP4_SCALE_E4(M) ((8 + (M)) * 0x1p-7f) +#define ROCMFP4_SCALE_E5(M) ((8 + (M)) * 0x1p-6f) +#define ROCMFP4_SCALE_E6(M) ((8 + (M)) * 0x1p-5f) +#define ROCMFP4_SCALE_E7(M) ((8 + (M)) * 0x1p-4f) +#define ROCMFP4_SCALE_E8(M) ((8 + (M)) * 0x1p-3f) +#define ROCMFP4_SCALE_E9(M) ((8 + (M)) * 0x1p-2f) +#define ROCMFP4_SCALE_E10(M) ((8 + (M)) * 0x1p-1f) +#define ROCMFP4_SCALE_E11(M) ((8 + (M)) * 0x1p0f) +#define ROCMFP4_SCALE_E12(M) ((8 + (M)) * 0x1p1f) +#define ROCMFP4_SCALE_E13(M) ((8 + (M)) * 0x1p2f) +#define ROCMFP4_SCALE_E14(M) ((8 + (M)) * 0x1p3f) +#define ROCMFP4_SCALE_E15(M) ((8 + (M)) * 0x1p4f) + +static const float rocmfp4_scale_ue4m3_half[127] = { + ROCMFP4_SCALE_SUB(0), ROCMFP4_SCALE_SUB(1), ROCMFP4_SCALE_SUB(2), ROCMFP4_SCALE_SUB(3), + ROCMFP4_SCALE_SUB(4), ROCMFP4_SCALE_SUB(5), ROCMFP4_SCALE_SUB(6), ROCMFP4_SCALE_SUB(7), + ROCMFP4_SCALE_E1(0), ROCMFP4_SCALE_E1(1), ROCMFP4_SCALE_E1(2), ROCMFP4_SCALE_E1(3), + ROCMFP4_SCALE_E1(4), ROCMFP4_SCALE_E1(5), ROCMFP4_SCALE_E1(6), ROCMFP4_SCALE_E1(7), + ROCMFP4_SCALE_E2(0), ROCMFP4_SCALE_E2(1), ROCMFP4_SCALE_E2(2), ROCMFP4_SCALE_E2(3), + ROCMFP4_SCALE_E2(4), ROCMFP4_SCALE_E2(5), ROCMFP4_SCALE_E2(6), ROCMFP4_SCALE_E2(7), + ROCMFP4_SCALE_E3(0), ROCMFP4_SCALE_E3(1), ROCMFP4_SCALE_E3(2), ROCMFP4_SCALE_E3(3), + ROCMFP4_SCALE_E3(4), ROCMFP4_SCALE_E3(5), ROCMFP4_SCALE_E3(6), ROCMFP4_SCALE_E3(7), + ROCMFP4_SCALE_E4(0), ROCMFP4_SCALE_E4(1), ROCMFP4_SCALE_E4(2), ROCMFP4_SCALE_E4(3), + ROCMFP4_SCALE_E4(4), ROCMFP4_SCALE_E4(5), ROCMFP4_SCALE_E4(6), ROCMFP4_SCALE_E4(7), + ROCMFP4_SCALE_E5(0), ROCMFP4_SCALE_E5(1), ROCMFP4_SCALE_E5(2), ROCMFP4_SCALE_E5(3), + ROCMFP4_SCALE_E5(4), ROCMFP4_SCALE_E5(5), ROCMFP4_SCALE_E5(6), ROCMFP4_SCALE_E5(7), + ROCMFP4_SCALE_E6(0), ROCMFP4_SCALE_E6(1), ROCMFP4_SCALE_E6(2), ROCMFP4_SCALE_E6(3), + ROCMFP4_SCALE_E6(4), ROCMFP4_SCALE_E6(5), ROCMFP4_SCALE_E6(6), ROCMFP4_SCALE_E6(7), + ROCMFP4_SCALE_E7(0), ROCMFP4_SCALE_E7(1), ROCMFP4_SCALE_E7(2), ROCMFP4_SCALE_E7(3), + ROCMFP4_SCALE_E7(4), ROCMFP4_SCALE_E7(5), ROCMFP4_SCALE_E7(6), ROCMFP4_SCALE_E7(7), + ROCMFP4_SCALE_E8(0), ROCMFP4_SCALE_E8(1), ROCMFP4_SCALE_E8(2), ROCMFP4_SCALE_E8(3), + ROCMFP4_SCALE_E8(4), ROCMFP4_SCALE_E8(5), ROCMFP4_SCALE_E8(6), ROCMFP4_SCALE_E8(7), + ROCMFP4_SCALE_E9(0), ROCMFP4_SCALE_E9(1), ROCMFP4_SCALE_E9(2), ROCMFP4_SCALE_E9(3), + ROCMFP4_SCALE_E9(4), ROCMFP4_SCALE_E9(5), ROCMFP4_SCALE_E9(6), ROCMFP4_SCALE_E9(7), + ROCMFP4_SCALE_E10(0), ROCMFP4_SCALE_E10(1), ROCMFP4_SCALE_E10(2), ROCMFP4_SCALE_E10(3), + ROCMFP4_SCALE_E10(4), ROCMFP4_SCALE_E10(5), ROCMFP4_SCALE_E10(6), ROCMFP4_SCALE_E10(7), + ROCMFP4_SCALE_E11(0), ROCMFP4_SCALE_E11(1), ROCMFP4_SCALE_E11(2), ROCMFP4_SCALE_E11(3), + ROCMFP4_SCALE_E11(4), ROCMFP4_SCALE_E11(5), ROCMFP4_SCALE_E11(6), ROCMFP4_SCALE_E11(7), + ROCMFP4_SCALE_E12(0), ROCMFP4_SCALE_E12(1), ROCMFP4_SCALE_E12(2), ROCMFP4_SCALE_E12(3), + ROCMFP4_SCALE_E12(4), ROCMFP4_SCALE_E12(5), ROCMFP4_SCALE_E12(6), ROCMFP4_SCALE_E12(7), + ROCMFP4_SCALE_E13(0), ROCMFP4_SCALE_E13(1), ROCMFP4_SCALE_E13(2), ROCMFP4_SCALE_E13(3), + ROCMFP4_SCALE_E13(4), ROCMFP4_SCALE_E13(5), ROCMFP4_SCALE_E13(6), ROCMFP4_SCALE_E13(7), + ROCMFP4_SCALE_E14(0), ROCMFP4_SCALE_E14(1), ROCMFP4_SCALE_E14(2), ROCMFP4_SCALE_E14(3), + ROCMFP4_SCALE_E14(4), ROCMFP4_SCALE_E14(5), ROCMFP4_SCALE_E14(6), ROCMFP4_SCALE_E14(7), + ROCMFP4_SCALE_E15(0), ROCMFP4_SCALE_E15(1), ROCMFP4_SCALE_E15(2), ROCMFP4_SCALE_E15(3), + ROCMFP4_SCALE_E15(4), ROCMFP4_SCALE_E15(5), ROCMFP4_SCALE_E15(6), +}; + +#undef ROCMFP4_SCALE_SUB +#undef ROCMFP4_SCALE_E1 +#undef ROCMFP4_SCALE_E2 +#undef ROCMFP4_SCALE_E3 +#undef ROCMFP4_SCALE_E4 +#undef ROCMFP4_SCALE_E5 +#undef ROCMFP4_SCALE_E6 +#undef ROCMFP4_SCALE_E7 +#undef ROCMFP4_SCALE_E8 +#undef ROCMFP4_SCALE_E9 +#undef ROCMFP4_SCALE_E10 +#undef ROCMFP4_SCALE_E11 +#undef ROCMFP4_SCALE_E12 +#undef ROCMFP4_SCALE_E13 +#undef ROCMFP4_SCALE_E14 +#undef ROCMFP4_SCALE_E15 + +static inline float rocmfp4_ue4m3_to_fp32_half(uint8_t e) { + return e <= 0x7e ? rocmfp4_scale_ue4m3_half[e] : 0.0f; +} + +static inline uint8_t rocmfp4_best_index_scaled_finite(float x, float inv_scale_half) { + // Exact nearest-neighbor thresholds for Codebook10: + // 0, +/-1, +/-2, +/-3, +/-4, +/-6, +/-8, +/-10 + // Ties intentionally choose the lower-magnitude code, matching the former + // linear scan because the positive codes and zero appear first. + const float a = fabsf(x * inv_scale_half); + if (a <= 0.5f) { + return 0; + } + + const bool neg = x < 0.0f; + if (a <= 1.5f) { + return neg ? 9 : 1; + } + if (a <= 2.5f) { + return neg ? 10 : 2; + } + if (a <= 3.5f) { + return neg ? 11 : 3; + } + if (a <= 5.0f) { + return neg ? 12 : 4; + } + if (a <= 7.0f) { + return neg ? 13 : 5; + } + if (a <= 9.0f) { + return neg ? 14 : 6; + } + + return neg ? 15 : 7; +} + +static inline uint8_t rocmfp4_best_index_scaled(float x, float inv_scale_half) { + if (!isfinite(x)) { + return 0; + } + + return rocmfp4_best_index_scaled_finite(x, inv_scale_half); +} + +// Fused best-index + decode used only inside the exhaustive scale search. The +// scale search re-scans every block element for every candidate scale byte, so +// avoiding the code -> decode round-trip on the hottest quantize path matters. +// Returns the same signed Codebook10 magnitude that +// rocmfp4_decode(rocmfp4_best_index_scaled_finite(x, inv_scale_half)) produces, +// so quantized output is bit-identical to the previous path. +static inline float rocmfp4_decoded_mag_scaled_finite(float x, float inv_scale_half) { + const float a = fabsf(x * inv_scale_half); + + float mag; + if (a <= 0.5f) { + mag = 0.0f; + } else if (a <= 1.5f) { + mag = 1.0f; + } else if (a <= 2.5f) { + mag = 2.0f; + } else if (a <= 3.5f) { + mag = 3.0f; + } else if (a <= 5.0f) { + mag = 4.0f; + } else if (a <= 7.0f) { + mag = 6.0f; + } else if (a <= 9.0f) { + mag = 8.0f; + } else { + mag = 10.0f; + } + + return x < 0.0f ? -mag : mag; +} + +static inline float rocmfp4_decoded_mag_scaled(float x, float inv_scale_half) { + if (!isfinite(x)) { + return 0.0f; + } + + return rocmfp4_decoded_mag_scaled_finite(x, inv_scale_half); +} + +static inline bool rocmfp4_scale_is_valid(uint8_t e) { + // ROCmFP4 scale bytes are unsigned finite E4M3 values. 0x7f is NaN in the + // unsigned encoding and values with the sign bit set are not valid scales. + return e <= 0x7e; +} + +static float rocmfp4_block_mse_for_scale_unweighted( + const float * x, int n, int e, float best_err) { + const float scale_half = rocmfp4_ue4m3_to_fp32_half((uint8_t) e); + const float inv_scale_half = 1.0f / scale_half; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + const float y = rocmfp4_decoded_mag_scaled(x[i], inv_scale_half) * scale_half; + const float d = x[i] - y; + + err += d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +static float rocmfp4_block_mse_for_scale_unweighted_finite( + const float * x, int n, int e, float best_err) { + const float scale_half = rocmfp4_ue4m3_to_fp32_half((uint8_t) e); + const float inv_scale_half = 1.0f / scale_half; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + const float y = rocmfp4_decoded_mag_scaled_finite(x[i], inv_scale_half) * scale_half; + const float d = x[i] - y; + + err += d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +static float rocmfp4_block_mse_for_scale_weighted( + const float * x, int n, const float * mse_weights, int e, float best_err) { + const float scale_half = rocmfp4_ue4m3_to_fp32_half((uint8_t) e); + const float inv_scale_half = 1.0f / scale_half; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + const float y = rocmfp4_decoded_mag_scaled(x[i], inv_scale_half) * scale_half; + const float d = x[i] - y; + + err += mse_weights[i]*d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +static float rocmfp4_block_mse_for_scale_weighted_finite( + const float * x, int n, const float * mse_weights, int e, float best_err) { + const float scale_half = rocmfp4_ue4m3_to_fp32_half((uint8_t) e); + const float inv_scale_half = 1.0f / scale_half; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + const float y = rocmfp4_decoded_mag_scaled_finite(x[i], inv_scale_half) * scale_half; + const float d = x[i] - y; + + err += mse_weights[i]*d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +static void rocmfp4_prepare_mse_weights( + float * dst, const float * x, int n, const float * quant_weights, float sigma2, + float * max_abs, float * max_abs_weight, bool * all_finite) { + *max_abs = 0.0f; + *max_abs_weight = 0.0f; + *all_finite = true; + + for (int i = 0; i < n; ++i) { + const float qw = quant_weights[i]; + const float ax = fabsf(x[i]); + const float weight = isfinite(qw) && qw > 0.0f ? qw * sqrtf(sigma2 + x[i]*x[i]) : 0.0f; + *all_finite = *all_finite && isfinite(x[i]); + + if (ax > *max_abs) { + *max_abs = ax; + *max_abs_weight = weight; + } else if (ax == *max_abs && weight > *max_abs_weight) { + *max_abs_weight = weight; + } + + // Match llama.cpp's imatrix weighting style for Q4_0: calibration + // importance is scaled by row energy so large activations remain protected. + dst[i] = weight; + } +} + +static int rocmfp4_nearest_scale_ue4m3(float target_scale_half) { + if (!(target_scale_half > 0.0f) || !isfinite(target_scale_half)) { + return 1; + } + + int lo = 1; + int hi = 126; + while (lo < hi) { + const int mid = lo + (hi - lo) / 2; + if (rocmfp4_ue4m3_to_fp32_half((uint8_t) mid) < target_scale_half) { + lo = mid + 1; + } else { + hi = mid; + } + } + + if (lo == 1) { + return 1; + } + + const float hi_scale = rocmfp4_ue4m3_to_fp32_half((uint8_t) lo); + const float lo_scale = rocmfp4_ue4m3_to_fp32_half((uint8_t) (lo - 1)); + + // Match the former ascending nearest scan: exact midpoint ties keep the + // lower scale byte. + return (target_scale_half - lo_scale <= hi_scale - target_scale_half) ? lo - 1 : lo; +} + +static uint8_t rocmfp4_choose_scale_ue4m3_exhaustive_unweighted( + const float * x, int n, float max_abs, bool all_finite) { + const int start_e = rocmfp4_nearest_scale_ue4m3(max_abs / 10.0f); + + int best_e = 0; + float best_err = FLT_MAX; + bool lower_done = false; + + for (int delta = 0; delta <= 125; ++delta) { + const int e0 = start_e - delta; + if (!lower_done && e0 >= 1 && e0 <= 126) { + const float scale_half = rocmfp4_ue4m3_to_fp32_half((uint8_t) e0); + const float clip_delta = max_abs - 10.0f*scale_half; + if (clip_delta > 0.0f && clip_delta*clip_delta > best_err) { + lower_done = true; + } else { + const float err = all_finite ? + rocmfp4_block_mse_for_scale_unweighted_finite(x, n, e0, best_err) : + rocmfp4_block_mse_for_scale_unweighted(x, n, e0, best_err); + if (err < best_err || (err == best_err && e0 < best_e)) { + best_err = err; + best_e = e0; + } + } + } + + const int e1 = start_e + delta; + if (delta != 0 && e1 >= 1 && e1 <= 126) { + const float err = all_finite ? + rocmfp4_block_mse_for_scale_unweighted_finite(x, n, e1, best_err) : + rocmfp4_block_mse_for_scale_unweighted(x, n, e1, best_err); + if (err < best_err || (err == best_err && e1 < best_e)) { + best_err = err; + best_e = e1; + } + } + + if ((lower_done || e0 <= 1) && e1 >= 126) { + break; + } + } + + return (uint8_t) best_e; +} + +static uint8_t rocmfp4_choose_scale_ue4m3_exhaustive_weighted( + const float * x, int n, const float * mse_weights, float max_abs, float max_abs_weight, bool all_finite) { + const int start_e = rocmfp4_nearest_scale_ue4m3(max_abs / 10.0f); + + int best_e = 0; + float best_err = FLT_MAX; + bool lower_done = false; + + for (int delta = 0; delta <= 125; ++delta) { + const int e0 = start_e - delta; + if (!lower_done && e0 >= 1 && e0 <= 126) { + const float scale_half = rocmfp4_ue4m3_to_fp32_half((uint8_t) e0); + const float clip_delta = max_abs - 10.0f*scale_half; + if (max_abs_weight > 0.0f && clip_delta > 0.0f && max_abs_weight*clip_delta*clip_delta > best_err) { + lower_done = true; + } else { + const float err = all_finite ? + rocmfp4_block_mse_for_scale_weighted_finite(x, n, mse_weights, e0, best_err) : + rocmfp4_block_mse_for_scale_weighted(x, n, mse_weights, e0, best_err); + if (err < best_err || (err == best_err && e0 < best_e)) { + best_err = err; + best_e = e0; + } + } + } + + const int e1 = start_e + delta; + if (delta != 0 && e1 >= 1 && e1 <= 126) { + const float err = all_finite ? + rocmfp4_block_mse_for_scale_weighted_finite(x, n, mse_weights, e1, best_err) : + rocmfp4_block_mse_for_scale_weighted(x, n, mse_weights, e1, best_err); + if (err < best_err || (err == best_err && e1 < best_e)) { + best_err = err; + best_e = e1; + } + } + + if ((lower_done || e0 <= 1) && e1 >= 126) { + break; + } + } + + return (uint8_t) best_e; +} + +static uint8_t rocmfp4_choose_scale_ue4m3(const float * x, int n, const float * quant_weights, float sigma2) { + if (quant_weights) { + assert(n <= QK_ROCMFP4); + float mse_weights_buf[QK_ROCMFP4]; + float weighted_max_abs; + float max_abs_weight; + bool all_finite; + rocmfp4_prepare_mse_weights(mse_weights_buf, x, n, quant_weights, sigma2, &weighted_max_abs, &max_abs_weight, &all_finite); + if (!(weighted_max_abs > 0.0f) || !isfinite(weighted_max_abs)) { + return 0; + } + return rocmfp4_choose_scale_ue4m3_exhaustive_weighted(x, n, mse_weights_buf, weighted_max_abs, max_abs_weight, all_finite); + } + + float max_abs = 0.0f; + bool all_finite = true; + for (int i = 0; i < n; ++i) { + all_finite = all_finite && isfinite(x[i]); + const float ax = fabsf(x[i]); + if (ax > max_abs) { + max_abs = ax; + } + } + + if (!(max_abs > 0.0f) || !isfinite(max_abs)) { + return 0; + } + + return rocmfp4_choose_scale_ue4m3_exhaustive_unweighted(x, n, max_abs, all_finite); +} + +static void rocmfp4_quantize_row_q4_0_weighted( + const float * GGML_RESTRICT x, block_rocmfp4 * GGML_RESTRICT y, int64_t k, const float * GGML_RESTRICT quant_weights) { + assert(k % QK_ROCMFP4 == 0); + + float sum_x2 = 0.0f; + for (int64_t i = 0; i < k; ++i) { + sum_x2 += x[i]*x[i]; + } + const float sigma2 = sum_x2 / (float) k; + + const int64_t nb = k / QK_ROCMFP4; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib*QK_ROCMFP4; + const float * qw = quant_weights ? quant_weights + ib*QK_ROCMFP4 : NULL; + const uint8_t e0 = rocmfp4_choose_scale_ue4m3(xb, QK_ROCMFP4/2, qw, sigma2); + const uint8_t e1 = rocmfp4_choose_scale_ue4m3(xb + QK_ROCMFP4/2, QK_ROCMFP4/2, qw ? qw + QK_ROCMFP4/2 : NULL, sigma2); + const float scale_half0 = rocmfp4_ue4m3_to_fp32_half(e0); + const float scale_half1 = rocmfp4_ue4m3_to_fp32_half(e1); + const float inv_scale_half0 = scale_half0 > 0.0f ? 1.0f / scale_half0 : 0.0f; + const float inv_scale_half1 = scale_half1 > 0.0f ? 1.0f / scale_half1 : 0.0f; + + y[ib].e[0] = e0; + y[ib].e[1] = e1; + + for (int j = 0; j < QK_ROCMFP4/2; ++j) { + const uint8_t q0 = rocmfp4_best_index_scaled(xb[j], inv_scale_half0); + const uint8_t q1 = rocmfp4_best_index_scaled(xb[j + QK_ROCMFP4/2], inv_scale_half1); + y[ib].qs[j] = q0 | (q1 << 4); + } + } +} + +static void rocmfp4_quantize_row_q4_0_fast_weighted( + const float * GGML_RESTRICT x, block_rocmfp4_fast * GGML_RESTRICT y, int64_t k, const float * GGML_RESTRICT quant_weights) { + assert(k % QK_ROCMFP4 == 0); + + float sum_x2 = 0.0f; + for (int64_t i = 0; i < k; ++i) { + sum_x2 += x[i]*x[i]; + } + const float sigma2 = sum_x2 / (float) k; + + const int64_t nb = k / QK_ROCMFP4; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib*QK_ROCMFP4; + const float * qw = quant_weights ? quant_weights + ib*QK_ROCMFP4 : NULL; + const uint8_t e = rocmfp4_choose_scale_ue4m3(xb, QK_ROCMFP4, qw, sigma2); + const float scale_half = rocmfp4_ue4m3_to_fp32_half(e); + const float inv_scale_half = scale_half > 0.0f ? 1.0f / scale_half : 0.0f; + + y[ib].e = e; + + for (int j = 0; j < QK_ROCMFP4/2; ++j) { + const uint8_t q0 = rocmfp4_best_index_scaled(xb[j], inv_scale_half); + const uint8_t q1 = rocmfp4_best_index_scaled(xb[j + QK_ROCMFP4/2], inv_scale_half); + y[ib].qs[j] = q0 | (q1 << 4); + } + } +} + +void rocmfp4_quantize_row_q4_0_ref(const float * GGML_RESTRICT x, block_rocmfp4 * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP4 == 0); + + const int64_t nb = k / QK_ROCMFP4; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib*QK_ROCMFP4; + const uint8_t e0 = rocmfp4_choose_scale_ue4m3(xb, QK_ROCMFP4/2, NULL, 0.0f); + const uint8_t e1 = rocmfp4_choose_scale_ue4m3(xb + QK_ROCMFP4/2, QK_ROCMFP4/2, NULL, 0.0f); + const float scale_half0 = rocmfp4_ue4m3_to_fp32_half(e0); + const float scale_half1 = rocmfp4_ue4m3_to_fp32_half(e1); + const float inv_scale_half0 = scale_half0 > 0.0f ? 1.0f / scale_half0 : 0.0f; + const float inv_scale_half1 = scale_half1 > 0.0f ? 1.0f / scale_half1 : 0.0f; + + y[ib].e[0] = e0; + y[ib].e[1] = e1; + + for (int j = 0; j < QK_ROCMFP4/2; ++j) { + const uint8_t q0 = rocmfp4_best_index_scaled(xb[j], inv_scale_half0); + const uint8_t q1 = rocmfp4_best_index_scaled(xb[j + QK_ROCMFP4/2], inv_scale_half1); + y[ib].qs[j] = q0 | (q1 << 4); + } + } +} + +void rocmfp4_quantize_row_q4_0_fast_ref(const float * GGML_RESTRICT x, block_rocmfp4_fast * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP4 == 0); + + const int64_t nb = k / QK_ROCMFP4; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib*QK_ROCMFP4; + const uint8_t e = rocmfp4_choose_scale_ue4m3(xb, QK_ROCMFP4, NULL, 0.0f); + const float scale_half = rocmfp4_ue4m3_to_fp32_half(e); + const float inv_scale_half = scale_half > 0.0f ? 1.0f / scale_half : 0.0f; + + y[ib].e = e; + + for (int j = 0; j < QK_ROCMFP4/2; ++j) { + const uint8_t q0 = rocmfp4_best_index_scaled(xb[j], inv_scale_half); + const uint8_t q1 = rocmfp4_best_index_scaled(xb[j + QK_ROCMFP4/2], inv_scale_half); + y[ib].qs[j] = q0 | (q1 << 4); + } + } +} + +void rocmfp4_dequantize_row_q4_0(const block_rocmfp4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP4 == 0); + + const int64_t nb = k / QK_ROCMFP4; + for (int64_t ib = 0; ib < nb; ++ib) { + const float d0 = rocmfp4_ue4m3_to_fp32_half(x[ib].e[0]); + const float d1 = rocmfp4_ue4m3_to_fp32_half(x[ib].e[1]); + + for (int j = 0; j < QK_ROCMFP4/2; ++j) { + y[ib*QK_ROCMFP4 + j] = (float) rocmfp4_decode(x[ib].qs[j] & 0x0f) * d0; + y[ib*QK_ROCMFP4 + j + QK_ROCMFP4/2] = (float) rocmfp4_decode(x[ib].qs[j] >> 4) * d1; + } + } +} + +void rocmfp4_dequantize_row_q4_0_fast(const block_rocmfp4_fast * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP4 == 0); + + const int64_t nb = k / QK_ROCMFP4; + for (int64_t ib = 0; ib < nb; ++ib) { + const float d = rocmfp4_ue4m3_to_fp32_half(x[ib].e); + + for (int j = 0; j < QK_ROCMFP4/2; ++j) { + y[ib*QK_ROCMFP4 + j] = (float) rocmfp4_decode(x[ib].qs[j] & 0x0f) * d; + y[ib*QK_ROCMFP4 + j + QK_ROCMFP4/2] = (float) rocmfp4_decode(x[ib].qs[j] >> 4) * d; + } + } +} + +void rocmfp4_quantize_row_q4_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + rocmfp4_quantize_row_q4_0_ref(x, (block_rocmfp4 *) y, k); +} + +void rocmfp4_quantize_row_q4_0_fast(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + rocmfp4_quantize_row_q4_0_fast_ref(x, (block_rocmfp4_fast *) y, k); +} + +size_t rocmfp4_quantize_q4_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix) { + const size_t row_size = ggml_row_size(GGML_TYPE_Q4_0_ROCMFP4, n_per_row); + + if (!imatrix) { + rocmfp4_quantize_row_q4_0_ref(src, (block_rocmfp4 *) dst, nrows*n_per_row); + return nrows * row_size; + } + + char * qrow = (char *) dst; + for (int64_t row = 0; row < nrows; ++row) { + rocmfp4_quantize_row_q4_0_weighted(src, (block_rocmfp4 *) qrow, n_per_row, imatrix); + src += n_per_row; + qrow += row_size; + } + + return nrows * row_size; +} + +size_t rocmfp4_quantize_q4_0_fast(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix) { + const size_t row_size = ggml_row_size(GGML_TYPE_Q4_0_ROCMFP4_FAST, n_per_row); + + if (!imatrix) { + rocmfp4_quantize_row_q4_0_fast_ref(src, (block_rocmfp4_fast *) dst, nrows*n_per_row); + return nrows * row_size; + } + + char * qrow = (char *) dst; + for (int64_t row = 0; row < nrows; ++row) { + rocmfp4_quantize_row_q4_0_fast_weighted(src, (block_rocmfp4_fast *) qrow, n_per_row, imatrix); + src += n_per_row; + qrow += row_size; + } + + return nrows * row_size; +} + +bool rocmfp4_validate_row_data(const void * data, size_t nbytes) { + if (nbytes % sizeof(block_rocmfp4) != 0) { + return false; + } + + const block_rocmfp4 * blocks = (const block_rocmfp4 *) data; + const size_t nblocks = nbytes / sizeof(block_rocmfp4); + for (size_t i = 0; i < nblocks; ++i) { + if (!rocmfp4_scale_is_valid(blocks[i].e[0]) || !rocmfp4_scale_is_valid(blocks[i].e[1])) { + return false; + } + } + + return true; +} + +bool rocmfp4_validate_row_data_fast(const void * data, size_t nbytes) { + if (nbytes % sizeof(block_rocmfp4_fast) != 0) { + return false; + } + + const block_rocmfp4_fast * blocks = (const block_rocmfp4_fast *) data; + const size_t nblocks = nbytes / sizeof(block_rocmfp4_fast); + for (size_t i = 0; i < nblocks; ++i) { + if (!rocmfp4_scale_is_valid(blocks[i].e)) { + return false; + } + } + + return true; +} + +#ifdef ROCMFP4_X86_AVX2_DISPATCH +__attribute__((target("avx2"))) +static inline int rocmfp4_hsum_i32_8_avx2(__m256i v) { + __m128i s = _mm_add_epi32(_mm256_castsi256_si128(v), _mm256_extracti128_si256(v, 1)); + s = _mm_add_epi32(s, _mm_shuffle_epi32(s, _MM_SHUFFLE(1, 0, 3, 2))); + s = _mm_add_epi32(s, _mm_shuffle_epi32(s, _MM_SHUFFLE(2, 3, 0, 1))); + return _mm_cvtsi128_si32(s); +} + +// Decode one 32-weight block's low and high nibble streams through the +// Codebook10 table with a single PSHUFB, then integer-dot each against its half +// of the q8_0 block. Integer sums are order-independent, so sumi0/sumi1 match +// the scalar reference exactly and the float result is bit-identical. +__attribute__((target("avx2"))) +static inline void rocmfp4_block_isums_avx2( + const uint8_t * qs, const int8_t * q8, int * sumi0, int * sumi1) { + const __m128i tbl = _mm_loadu_si128((const __m128i *) rocmfp4_codebook); + const __m128i q = _mm_loadu_si128((const __m128i *) qs); + const __m128i lo = _mm_and_si128(q, _mm_set1_epi8(0x0F)); + const __m128i hi = _mm_and_si128(_mm_srli_epi16(q, 4), _mm_set1_epi8(0x0F)); + const __m128i dlo = _mm_shuffle_epi8(tbl, lo); + const __m128i dhi = _mm_shuffle_epi8(tbl, hi); + const __m128i ylo = _mm_loadu_si128((const __m128i *) q8); + const __m128i yhi = _mm_loadu_si128((const __m128i *) (q8 + QK_ROCMFP4/2)); + const __m256i pl = _mm256_madd_epi16(_mm256_cvtepi8_epi16(dlo), _mm256_cvtepi8_epi16(ylo)); + const __m256i ph = _mm256_madd_epi16(_mm256_cvtepi8_epi16(dhi), _mm256_cvtepi8_epi16(yhi)); + *sumi0 = rocmfp4_hsum_i32_8_avx2(pl); + *sumi1 = rocmfp4_hsum_i32_8_avx2(ph); +} + +__attribute__((target("avx2"))) +static void rocmfp4_vec_dot_q4_0_q8_0_avx2( + int nb, float * GGML_RESTRICT s, const block_rocmfp4 * GGML_RESTRICT x, const block_q8_0 * GGML_RESTRICT y) { + float sumf = 0.0f; + for (int ib = 0; ib < nb; ++ib) { + const float d0 = rocmfp4_ue4m3_to_fp32_half(x[ib].e[0]) * ggml_fp16_to_fp32(y[ib].d); + const float d1 = rocmfp4_ue4m3_to_fp32_half(x[ib].e[1]) * ggml_fp16_to_fp32(y[ib].d); + int sumi0, sumi1; + rocmfp4_block_isums_avx2(x[ib].qs, y[ib].qs, &sumi0, &sumi1); + sumf += d0 * (float) sumi0 + d1 * (float) sumi1; + } + *s = sumf; +} + +__attribute__((target("avx2"))) +static void rocmfp4_vec_dot_q4_0_fast_q8_0_avx2( + int nb, float * GGML_RESTRICT s, const block_rocmfp4_fast * GGML_RESTRICT x, const block_q8_0 * GGML_RESTRICT y) { + float sumf = 0.0f; + for (int ib = 0; ib < nb; ++ib) { + const float d = rocmfp4_ue4m3_to_fp32_half(x[ib].e) * ggml_fp16_to_fp32(y[ib].d); + int sumi0, sumi1; + rocmfp4_block_isums_avx2(x[ib].qs, y[ib].qs, &sumi0, &sumi1); + sumf += d * (float) (sumi0 + sumi1); + } + *s = sumf; +} +#endif // ROCMFP4_X86_AVX2_DISPATCH + +void rocmfp4_vec_dot_q4_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); + assert(nrc == 1); + GGML_UNUSED(nrc); + assert(n % QK_ROCMFP4 == 0); + assert(QK_ROCMFP4 == QK8_0); + + const block_rocmfp4 * GGML_RESTRICT x = (const block_rocmfp4 *) vx; + const block_q8_0 * GGML_RESTRICT y = (const block_q8_0 *) vy; + + const int nb = n / QK_ROCMFP4; + +#ifdef ROCMFP4_X86_AVX2_DISPATCH + if (__builtin_cpu_supports("avx2")) { + rocmfp4_vec_dot_q4_0_q8_0_avx2(nb, s, x, y); + return; + } +#endif + + float sumf = 0.0f; + + for (int ib = 0; ib < nb; ++ib) { + const float d0 = rocmfp4_ue4m3_to_fp32_half(x[ib].e[0]) * ggml_fp16_to_fp32(y[ib].d); + const float d1 = rocmfp4_ue4m3_to_fp32_half(x[ib].e[1]) * ggml_fp16_to_fp32(y[ib].d); + int sumi0 = 0; + int sumi1 = 0; + + for (int j = 0; j < QK_ROCMFP4/2; ++j) { + const uint8_t q = x[ib].qs[j]; + sumi0 += rocmfp4_decode_table(q) * y[ib].qs[j]; + sumi1 += rocmfp4_decode_table(q >> 4) * y[ib].qs[j + QK_ROCMFP4/2]; + } + + sumf += d0 * (float) sumi0 + d1 * (float) sumi1; + } + + *s = sumf; +} + +void rocmfp4_vec_dot_q4_0_fast_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); + assert(nrc == 1); + GGML_UNUSED(nrc); + assert(n % QK_ROCMFP4 == 0); + assert(QK_ROCMFP4 == QK8_0); + + const block_rocmfp4_fast * GGML_RESTRICT x = (const block_rocmfp4_fast *) vx; + const block_q8_0 * GGML_RESTRICT y = (const block_q8_0 *) vy; + + const int nb = n / QK_ROCMFP4; + +#ifdef ROCMFP4_X86_AVX2_DISPATCH + if (__builtin_cpu_supports("avx2")) { + rocmfp4_vec_dot_q4_0_fast_q8_0_avx2(nb, s, x, y); + return; + } +#endif + + float sumf = 0.0f; + + for (int ib = 0; ib < nb; ++ib) { + const float d = rocmfp4_ue4m3_to_fp32_half(x[ib].e) * ggml_fp16_to_fp32(y[ib].d); + int sumi = 0; + + for (int j = 0; j < QK_ROCMFP4/2; ++j) { + const uint8_t q = x[ib].qs[j]; + sumi += rocmfp4_decode_table(q) * y[ib].qs[j]; + sumi += rocmfp4_decode_table(q >> 4) * y[ib].qs[j + QK_ROCMFP4/2]; + } + + sumf += d * (float) sumi; + } + + *s = sumf; +} diff --git a/ggml/rocmfp4/rocmfp4.h b/ggml/rocmfp4/rocmfp4.h new file mode 100644 index 000000000000..a4fb24761cfe --- /dev/null +++ b/ggml/rocmfp4/rocmfp4.h @@ -0,0 +1,66 @@ +// ROCmFPx quant formats, hand-ported into this fork. +// +// Origin: https://github.com/charlie12345/ROCmFPX - creator of the ROCmFP4 format. +// Ported from: https://github.com/ciru-ai/ROCmFPX - a fork of the above. +// +// Both upstream projects are MIT licensed and based on llama.cpp; upstream authors +// retain their authorship and MIT license credit. See LICENSE. + +#pragma once + +#include +#include +#include + +#include "ggml.h" + +#ifdef __cplusplus +extern "C" { +#endif + +#define QK_ROCMFP4 32 +#define QR_ROCMFP4 2 +#define QI_ROCMFP4 (QK_ROCMFP4 / (4 * QR_ROCMFP4)) +#define QS_ROCMFP4 32 + +// AMD-tuned compact layout: 16 bytes of packed E2M1-derived 4-bit codes, then +// one unsigned E4M3 scale byte per 16-weight half block. +typedef struct { + uint8_t qs[QK_ROCMFP4/2]; + uint8_t e[2]; +} block_rocmfp4; + +// Speed-focused layout: same 32 packed ROCmFP4 nibbles, but one UE4M3 scale +// for the whole block. This is a separate GGUF type so fast 4.25 BPW artifacts +// never alias the safer dual-scale format above. +typedef struct { + uint8_t qs[QK_ROCMFP4/2]; + uint8_t e; +} block_rocmfp4_fast; + +#if defined(__cplusplus) +static_assert(sizeof(block_rocmfp4) == QK_ROCMFP4/2 + 2*sizeof(uint8_t), "wrong rocmfp4 block size/padding"); +static_assert(sizeof(block_rocmfp4_fast) == QK_ROCMFP4/2 + sizeof(uint8_t), "wrong rocmfp4 fast block size/padding"); +#else +_Static_assert(sizeof(block_rocmfp4) == QK_ROCMFP4/2 + 2*sizeof(uint8_t), "wrong rocmfp4 block size/padding"); +_Static_assert(sizeof(block_rocmfp4_fast) == QK_ROCMFP4/2 + sizeof(uint8_t), "wrong rocmfp4 fast block size/padding"); +#endif + +GGML_API void rocmfp4_quantize_row_q4_0_ref(const float * GGML_RESTRICT x, block_rocmfp4 * GGML_RESTRICT y, int64_t k); +GGML_API void rocmfp4_dequantize_row_q4_0(const block_rocmfp4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void rocmfp4_quantize_row_q4_0_fast_ref(const float * GGML_RESTRICT x, block_rocmfp4_fast * GGML_RESTRICT y, int64_t k); +GGML_API void rocmfp4_dequantize_row_q4_0_fast(const block_rocmfp4_fast * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); + +GGML_API void rocmfp4_quantize_row_q4_0(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +GGML_API size_t rocmfp4_quantize_q4_0(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); +GGML_API void rocmfp4_quantize_row_q4_0_fast(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +GGML_API size_t rocmfp4_quantize_q4_0_fast(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); +GGML_API bool rocmfp4_validate_row_data(const void * data, size_t nbytes); +GGML_API bool rocmfp4_validate_row_data_fast(const void * data, size_t nbytes); + +GGML_API void rocmfp4_vec_dot_q4_0_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); +GGML_API void rocmfp4_vec_dot_q4_0_fast_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc); + +#ifdef __cplusplus +} +#endif diff --git a/ggml/rocmfpx/README.md b/ggml/rocmfpx/README.md new file mode 100644 index 000000000000..8d122c5b7ec2 --- /dev/null +++ b/ggml/rocmfpx/README.md @@ -0,0 +1,56 @@ +# ROCmFPx + +The 2/3/6/8-bit ROCmFPx tensor types. The formats originate in +[charlie12345/ROCmFPX](https://github.com/charlie12345/ROCmFPX) (MIT); the codecs +here were hand-ported via [ciru-ai/ROCmFPX](https://github.com/ciru-ai/ROCmFPX). +The 4-bit types live in `../rocmfp4/`. + +This directory owns the CPU codecs (`rocmfpx.c`) and the standalone ROCmFP2 +reference encoder (`rocmfp2_reference.c`) that the optimized FP2 encoder is +checked against. The rest of ggml registers and dispatches the types. + +## Layouts + +All formats use 32-weight blocks and the same unsigned E4M3 scale byte as +ROCmFP4 (see `../rocmfp4/README.md`). + +| Type | GGUF type id | Payload | Scales | Block bytes | BPW | +|---|---:|---|---|---:|---:| +| `Q2_0_ROCMFPX` | 107 | 32 packed 2-bit codes, LSB first | 2, one per 16-weight half | 10 | 2.50 | +| `Q3_0_ROCMFPX` | 104 | 32 packed 3-bit codes, LSB first | 2, one per 16-weight half | 14 | 3.50 | +| `Q6_0_ROCMFPX` | 102 | 32 packed 6-bit codes, LSB first | 2, one per 16-weight half | 26 | 6.50 | +| `Q8_0_ROCMFPX` | 103 | 32 signed int8 | 1 per block | 33 | 8.25 | + +- ROCmFP2 codes map to `-4, -1, 1, 4`. +- ROCmFP3 codes are sign + 2-bit magnitude index into `0, 1, 2, 4`. +- ROCmFP6 codes are sign + 5-bit magnitude; code 32 (sign set, magnitude 0) + decodes to `-32`. +- ROCmFP8 is signed int8 clamped to `[-127, 127]`. + +### Type ids and the two upstreams + +Ids 100-104 are the same in every ROCmFPX tree. For `Q2_0_ROCMFPX` the +upstreams disagree: charlie12345/ROCmFPX (and the kingjones30 fork the ROCm +builds come from) use tensor type 107 and file type 119; ciru-ai/ROCmFPX moved +it to 108 / 122 and put `Q7_0_ROCMFPX` at 107 / 119. This tree follows +charlie12345, which is what the published GGUFs carry. Ids 105, 106 (TURBO3_0, +TURBO4_0) and 108 (ROCMI4 or ciru-ai's Q2_0) are deliberately left unassigned +so files using them are rejected instead of misread. + +## What is implemented in this tree + +- CPU: quantize / dequantize / vec_dot for all four types, imatrix weighting. +- Vulkan: dequant, mat-vec and matmul for Q3/Q6/Q8; q8_1 integer-dot mat-vec + for Q8. Q2 runs on CPU. +- Not implemented: ROCm/HIP kernels. On ROCm these types run through the CPU + codecs. + +## Tests + +- `ctest -R rocmfp` runs `test-rocmfpx` (codecs and imatrix weighting) and + `test-rocmfp2-reference` (the FP2 reference encoder: layout, scale encoding, + golden vectors, non-finite input). +- `test-quantize-fns` covers all four types with per-width error budgets. +- `test-backend-ops` includes them in its type lists. +- `gguf-py/tests/test_quants.py --type Q6_0_ROCMFPX` (etc.) checks the Python + dequantizers against C. diff --git a/ggml/rocmfpx/rocmfp2_reference.c b/ggml/rocmfpx/rocmfp2_reference.c new file mode 100644 index 000000000000..039672f5d750 --- /dev/null +++ b/ggml/rocmfpx/rocmfp2_reference.c @@ -0,0 +1,491 @@ +// ROCmFPx quant formats, hand-ported into this fork. +// +// Origin: https://github.com/charlie12345/ROCmFPX - creator of the ROCmFP4 format. +// Ported from: https://github.com/ciru-ai/ROCmFPX - a fork of the above. +// +// Both upstream projects are MIT licensed and based on llama.cpp; upstream authors +// retain their authorship and MIT license credit. See LICENSE. + +#include "rocmfp2_reference.h" + +#include +#include +#include + +static bool rocmfp2_p1_mapping_is_valid(rocmfp2_p1_mapping mapping) { + return mapping == ROCMFP2_P1_MAPPING_MORD || mapping == ROCMFP2_P1_MAPPING_MSM; +} + +static bool rocmfp2_p1_codebook_is_valid(const rocmfp2_p1_codebook * codebook) { + return codebook != NULL && codebook->inner > 0 && codebook->outer > codebook->inner && codebook->outer <= 127; +} + +static double rocmfp2_p1_decode_scale_valid(uint8_t scale_byte) { + const int exponent = scale_byte >> 3; + const int mantissa = scale_byte & 7; + + if (exponent == 0) { + return ldexp((double) mantissa, -10); + } + + return ldexp((double) (8 + mantissa), exponent - 11); +} + +static int8_t rocmfp2_p1_decode_code_valid( + uint8_t code, + const rocmfp2_p1_codebook * codebook, + rocmfp2_p1_mapping mapping) { + const int inner = (int) codebook->inner; + const int outer = (int) codebook->outer; + + if (mapping == ROCMFP2_P1_MAPPING_MORD) { + static const int sign[4] = { -1, -1, 1, 1 }; + const int magnitude = (code == 0 || code == 3) ? outer : inner; + return (int8_t) (sign[code] * magnitude); + } + + static const int sign[4] = { 1, 1, -1, -1 }; + const int magnitude = (code == 1 || code == 3) ? outer : inner; + return (int8_t) (sign[code] * magnitude); +} + +static uint8_t rocmfp2_p1_encode_semantic_valid( + bool negative, + bool outer, + rocmfp2_p1_mapping mapping) { + if (mapping == ROCMFP2_P1_MAPPING_MORD) { + if (negative) { + return outer ? 0u : 1u; + } + return outer ? 3u : 2u; + } + + if (negative) { + return outer ? 3u : 2u; + } + return outer ? 1u : 0u; +} + +/* Lower ranks implement the frozen semantic tie rules. */ +static int rocmfp2_p1_code_tie_rank( + uint8_t code, + bool source_negative, + const rocmfp2_p1_codebook * codebook, + rocmfp2_p1_mapping mapping) { + const int value = (int) rocmfp2_p1_decode_code_valid(code, codebook, mapping); + const bool code_negative = value < 0; + const bool outer = abs(value) == (int) codebook->outer; + + /* Inner magnitude wins first; input sign wins a remaining sign tie. */ + return (outer ? 2 : 0) + (code_negative != source_negative ? 1 : 0); +} + +static uint8_t rocmfp2_p1_select_code_ref_valid( + double source, + double scale, + const rocmfp2_p1_codebook * codebook, + rocmfp2_p1_mapping mapping) { + /* Every code reconstructs exact zero; P0 canonicalizes this tie to 00. */ + if (scale == 0.0) { + return 0; + } + + const bool source_negative = signbit(source) != 0; + double best_error = INFINITY; + int best_rank = 0; + uint8_t best_code = 0; + bool have_best = false; + + for (uint8_t code = 0; code < 4; ++code) { + const double reconstructed = (double) rocmfp2_p1_decode_code_valid(code, codebook, mapping) * scale; + const double delta = source - reconstructed; + const double error = delta * delta; + const int rank = rocmfp2_p1_code_tie_rank(code, source_negative, codebook, mapping); + + if (!have_best || error < best_error || (error == best_error && rank < best_rank)) { + best_error = error; + best_rank = rank; + best_code = code; + have_best = true; + } + } + + return best_code; +} + +static uint8_t rocmfp2_p1_select_code_optimized_valid( + double source, + double scale, + const rocmfp2_p1_codebook * codebook, + rocmfp2_p1_mapping mapping) { + /* Keep serialization identical to the exhaustive reference at scale 0. */ + if (scale == 0.0) { + return 0; + } + + const double magnitude = fabs(source); + const double inner_reconstructed = (double) codebook->inner * scale; + const double outer_reconstructed = (double) codebook->outer * scale; + const double inner_delta = magnitude - inner_reconstructed; + const double outer_delta = magnitude - outer_reconstructed; + + /* Equality is the exact inner/outer midpoint tie and keeps inner. */ + const bool outer = outer_delta * outer_delta < inner_delta * inner_delta; + return rocmfp2_p1_encode_semantic_valid(signbit(source) != 0, outer, mapping); +} + +const char * rocmfp2_p1_status_name(rocmfp2_p1_status status) { + switch (status) { + case ROCMFP2_P1_OK: return "OK"; + case ROCMFP2_P1_NONFINITE_SOURCE: return "NONFINITE_SOURCE"; + case ROCMFP2_P1_INVALID_ARGUMENT: return "INVALID_ARGUMENT"; + case ROCMFP2_P1_INVALID_CODEBOOK: return "INVALID_CODEBOOK"; + case ROCMFP2_P1_INVALID_MAPPING: return "INVALID_MAPPING"; + case ROCMFP2_P1_INVALID_SCALE_METADATA: return "INVALID_SCALE_METADATA"; + case ROCMFP2_P1_INVALID_CODE: return "INVALID_CODE"; + } + + return "UNKNOWN_STATUS"; +} + +bool rocmfp2_p1_scale_is_valid(uint8_t scale_byte) { + return scale_byte <= 0x7e; +} + +double rocmfp2_p1_ue4m3_to_binary64(uint8_t scale_byte) { + if (!rocmfp2_p1_scale_is_valid(scale_byte)) { + return NAN; + } + + return rocmfp2_p1_decode_scale_valid(scale_byte); +} + +uint8_t rocmfp2_p1_nearest_ue4m3(double target) { + if (isnan(target) || target < 0.0) { + return 0xff; + } + if (target == 0.0) { + return 0; + } + + const double maximum = rocmfp2_p1_decode_scale_valid(0x7e); + if (!isfinite(target) || target >= maximum) { + return 0x7e; + } + + for (uint8_t upper_byte = 1; upper_byte <= 0x7e; ++upper_byte) { + const double upper = rocmfp2_p1_decode_scale_valid(upper_byte); + if (target <= upper) { + const uint8_t lower_byte = (uint8_t) (upper_byte - 1); + const double lower = rocmfp2_p1_decode_scale_valid(lower_byte); + + /* Exact midpoint ties select the smaller scale byte. */ + return target - lower <= upper - target ? lower_byte : upper_byte; + } + } + + return 0x7e; +} + +rocmfp2_p1_status rocmfp2_p1_decode_code( + uint8_t code, + const rocmfp2_p1_codebook * codebook, + rocmfp2_p1_mapping mapping, + int8_t * value) { + if (value == NULL) { + return ROCMFP2_P1_INVALID_ARGUMENT; + } + if (!rocmfp2_p1_codebook_is_valid(codebook)) { + return ROCMFP2_P1_INVALID_CODEBOOK; + } + if (!rocmfp2_p1_mapping_is_valid(mapping)) { + return ROCMFP2_P1_INVALID_MAPPING; + } + if (code > 3) { + return ROCMFP2_P1_INVALID_CODE; + } + + *value = rocmfp2_p1_decode_code_valid(code, codebook, mapping); + return ROCMFP2_P1_OK; +} + +rocmfp2_p1_status rocmfp2_p1_decode_packed_byte( + uint8_t packed, + const rocmfp2_p1_codebook * codebook, + rocmfp2_p1_mapping mapping, + int8_t values[4]) { + if (values == NULL) { + return ROCMFP2_P1_INVALID_ARGUMENT; + } + if (!rocmfp2_p1_codebook_is_valid(codebook)) { + return ROCMFP2_P1_INVALID_CODEBOOK; + } + if (!rocmfp2_p1_mapping_is_valid(mapping)) { + return ROCMFP2_P1_INVALID_MAPPING; + } + + for (int lane = 0; lane < 4; ++lane) { + const uint8_t code = (uint8_t) ((packed >> (2 * lane)) & 3u); + values[lane] = rocmfp2_p1_decode_code_valid(code, codebook, mapping); + } + + return ROCMFP2_P1_OK; +} + +rocmfp2_p1_status rocmfp2_p1_pack_codes( + const uint8_t codes[ROCMFP2_P1_BLOCK_WEIGHTS], + const uint8_t scales[ROCMFP2_P1_SCALE_BYTES], + rocmfp2_p1_block * block) { + if (codes == NULL || scales == NULL || block == NULL) { + return ROCMFP2_P1_INVALID_ARGUMENT; + } + if (!rocmfp2_p1_scale_is_valid(scales[0]) || !rocmfp2_p1_scale_is_valid(scales[1])) { + return ROCMFP2_P1_INVALID_SCALE_METADATA; + } + + rocmfp2_p1_block temporary; + for (int packed_index = 0; packed_index < ROCMFP2_P1_DATA_BYTES; ++packed_index) { + uint8_t packed = 0; + for (int lane = 0; lane < 4; ++lane) { + const uint8_t code = codes[4 * packed_index + lane]; + if (code > 3) { + return ROCMFP2_P1_INVALID_CODE; + } + packed |= (uint8_t) (code << (2 * lane)); + } + temporary.d[packed_index] = packed; + } + temporary.s[0] = scales[0]; + temporary.s[1] = scales[1]; + + memcpy(block, &temporary, sizeof(temporary)); + return ROCMFP2_P1_OK; +} + +void rocmfp2_p1_unpack_codes( + const rocmfp2_p1_block * block, + uint8_t codes[ROCMFP2_P1_BLOCK_WEIGHTS]) { + for (int packed_index = 0; packed_index < ROCMFP2_P1_DATA_BYTES; ++packed_index) { + const uint8_t packed = block->d[packed_index]; + for (int lane = 0; lane < 4; ++lane) { + codes[4 * packed_index + lane] = (uint8_t) ((packed >> (2 * lane)) & 3u); + } + } +} + +bool rocmfp2_p1_validate_block(const rocmfp2_p1_block * block) { + return block != NULL && rocmfp2_p1_scale_is_valid(block->s[0]) && rocmfp2_p1_scale_is_valid(block->s[1]); +} + +bool rocmfp2_p1_validate_serialized(const void * data, size_t nbytes) { + if (nbytes == 0) { + return true; + } + if (data == NULL || nbytes % ROCMFP2_P1_BLOCK_BYTES != 0) { + return false; + } + + const uint8_t * bytes = (const uint8_t *) data; + const size_t blocks = nbytes / ROCMFP2_P1_BLOCK_BYTES; + for (size_t block_index = 0; block_index < blocks; ++block_index) { + const size_t base = block_index * ROCMFP2_P1_BLOCK_BYTES; + if (!rocmfp2_p1_scale_is_valid(bytes[base + 8]) || !rocmfp2_p1_scale_is_valid(bytes[base + 9])) { + return false; + } + } + + return true; +} + +rocmfp2_p1_status rocmfp2_p1_select_code_ref( + double source, + uint8_t scale_byte, + const rocmfp2_p1_codebook * codebook, + rocmfp2_p1_mapping mapping, + uint8_t * code) { + if (code == NULL) { + return ROCMFP2_P1_INVALID_ARGUMENT; + } + if (!rocmfp2_p1_codebook_is_valid(codebook)) { + return ROCMFP2_P1_INVALID_CODEBOOK; + } + if (!rocmfp2_p1_mapping_is_valid(mapping)) { + return ROCMFP2_P1_INVALID_MAPPING; + } + if (!rocmfp2_p1_scale_is_valid(scale_byte)) { + return ROCMFP2_P1_INVALID_SCALE_METADATA; + } + if (!isfinite(source)) { + return ROCMFP2_P1_NONFINITE_SOURCE; + } + + *code = rocmfp2_p1_select_code_ref_valid( + source, rocmfp2_p1_decode_scale_valid(scale_byte), codebook, mapping); + return ROCMFP2_P1_OK; +} + +rocmfp2_p1_status rocmfp2_p1_select_code_optimized( + double source, + uint8_t scale_byte, + const rocmfp2_p1_codebook * codebook, + rocmfp2_p1_mapping mapping, + uint8_t * code) { + if (code == NULL) { + return ROCMFP2_P1_INVALID_ARGUMENT; + } + if (!rocmfp2_p1_codebook_is_valid(codebook)) { + return ROCMFP2_P1_INVALID_CODEBOOK; + } + if (!rocmfp2_p1_mapping_is_valid(mapping)) { + return ROCMFP2_P1_INVALID_MAPPING; + } + if (!rocmfp2_p1_scale_is_valid(scale_byte)) { + return ROCMFP2_P1_INVALID_SCALE_METADATA; + } + if (!isfinite(source)) { + return ROCMFP2_P1_NONFINITE_SOURCE; + } + + *code = rocmfp2_p1_select_code_optimized_valid( + source, rocmfp2_p1_decode_scale_valid(scale_byte), codebook, mapping); + return ROCMFP2_P1_OK; +} + +typedef uint8_t (*rocmfp2_p1_selector)( + double source, + double scale, + const rocmfp2_p1_codebook * codebook, + rocmfp2_p1_mapping mapping); + +static rocmfp2_p1_status rocmfp2_p1_quantize_block_impl( + const double source[ROCMFP2_P1_BLOCK_WEIGHTS], + const rocmfp2_p1_codebook * codebook, + rocmfp2_p1_mapping mapping, + rocmfp2_p1_block * block, + double group_sse[ROCMFP2_P1_SCALE_BYTES], + rocmfp2_p1_selector selector) { + if (source == NULL || block == NULL) { + return ROCMFP2_P1_INVALID_ARGUMENT; + } + if (!rocmfp2_p1_codebook_is_valid(codebook)) { + return ROCMFP2_P1_INVALID_CODEBOOK; + } + if (!rocmfp2_p1_mapping_is_valid(mapping)) { + return ROCMFP2_P1_INVALID_MAPPING; + } + + for (int i = 0; i < ROCMFP2_P1_BLOCK_WEIGHTS; ++i) { + if (!isfinite(source[i])) { + return ROCMFP2_P1_NONFINITE_SOURCE; + } + } + + uint8_t best_codes[ROCMFP2_P1_BLOCK_WEIGHTS]; + uint8_t best_scales[ROCMFP2_P1_SCALE_BYTES]; + double best_group_sse[ROCMFP2_P1_SCALE_BYTES]; + + for (int group = 0; group < ROCMFP2_P1_SCALE_BYTES; ++group) { + const int source_offset = group * ROCMFP2_P1_GROUP_WEIGHTS; + double best_sse = INFINITY; + uint8_t best_scale = 0; + uint8_t candidate_codes[ROCMFP2_P1_GROUP_WEIGHTS]; + uint8_t saved_codes[ROCMFP2_P1_GROUP_WEIGHTS]; + bool have_best = false; + + /* Exact exhaustive search: every legal finite unsigned UE4M3 byte. */ + for (int scale_index = 0; scale_index <= 0x7e; ++scale_index) { + const uint8_t scale_byte = (uint8_t) scale_index; + const double scale = rocmfp2_p1_decode_scale_valid(scale_byte); + double candidate_sse = 0.0; + + for (int i = 0; i < ROCMFP2_P1_GROUP_WEIGHTS; ++i) { + const double value = source[source_offset + i]; + const uint8_t code = selector(value, scale, codebook, mapping); + const double reconstructed = (double) rocmfp2_p1_decode_code_valid(code, codebook, mapping) * scale; + const double delta = value - reconstructed; + + candidate_codes[i] = code; + candidate_sse += delta * delta; + } + + if (!have_best || candidate_sse < best_sse || + (candidate_sse == best_sse && scale_byte < best_scale)) { + best_sse = candidate_sse; + best_scale = scale_byte; + memcpy(saved_codes, candidate_codes, sizeof(saved_codes)); + have_best = true; + } + } + + memcpy(best_codes + source_offset, saved_codes, sizeof(saved_codes)); + best_scales[group] = best_scale; + best_group_sse[group] = best_sse; + } + + rocmfp2_p1_block temporary; + const rocmfp2_p1_status pack_status = rocmfp2_p1_pack_codes(best_codes, best_scales, &temporary); + if (pack_status != ROCMFP2_P1_OK) { + return pack_status; + } + + memcpy(block, &temporary, sizeof(temporary)); + if (group_sse != NULL) { + group_sse[0] = best_group_sse[0]; + group_sse[1] = best_group_sse[1]; + } + + return ROCMFP2_P1_OK; +} + +rocmfp2_p1_status rocmfp2_p1_quantize_block_ref( + const double source[ROCMFP2_P1_BLOCK_WEIGHTS], + const rocmfp2_p1_codebook * codebook, + rocmfp2_p1_mapping mapping, + rocmfp2_p1_block * block, + double group_sse[ROCMFP2_P1_SCALE_BYTES]) { + return rocmfp2_p1_quantize_block_impl( + source, codebook, mapping, block, group_sse, rocmfp2_p1_select_code_ref_valid); +} + +rocmfp2_p1_status rocmfp2_p1_quantize_block_optimized( + const double source[ROCMFP2_P1_BLOCK_WEIGHTS], + const rocmfp2_p1_codebook * codebook, + rocmfp2_p1_mapping mapping, + rocmfp2_p1_block * block, + double group_sse[ROCMFP2_P1_SCALE_BYTES]) { + return rocmfp2_p1_quantize_block_impl( + source, codebook, mapping, block, group_sse, rocmfp2_p1_select_code_optimized_valid); +} + +rocmfp2_p1_status rocmfp2_p1_dequantize_block( + const rocmfp2_p1_block * block, + const rocmfp2_p1_codebook * codebook, + rocmfp2_p1_mapping mapping, + double output[ROCMFP2_P1_BLOCK_WEIGHTS]) { + if (block == NULL || output == NULL) { + return ROCMFP2_P1_INVALID_ARGUMENT; + } + if (!rocmfp2_p1_codebook_is_valid(codebook)) { + return ROCMFP2_P1_INVALID_CODEBOOK; + } + if (!rocmfp2_p1_mapping_is_valid(mapping)) { + return ROCMFP2_P1_INVALID_MAPPING; + } + if (!rocmfp2_p1_validate_block(block)) { + return ROCMFP2_P1_INVALID_SCALE_METADATA; + } + + uint8_t codes[ROCMFP2_P1_BLOCK_WEIGHTS]; + double temporary[ROCMFP2_P1_BLOCK_WEIGHTS]; + rocmfp2_p1_unpack_codes(block, codes); + + for (int i = 0; i < ROCMFP2_P1_BLOCK_WEIGHTS; ++i) { + const int group = i / ROCMFP2_P1_GROUP_WEIGHTS; + const double scale = rocmfp2_p1_decode_scale_valid(block->s[group]); + const int8_t value = rocmfp2_p1_decode_code_valid(codes[i], codebook, mapping); + temporary[i] = (double) value * scale; + } + + memcpy(output, temporary, sizeof(temporary)); + return ROCMFP2_P1_OK; +} diff --git a/ggml/rocmfpx/rocmfp2_reference.h b/ggml/rocmfpx/rocmfp2_reference.h new file mode 100644 index 000000000000..42c7e1ecfc15 --- /dev/null +++ b/ggml/rocmfpx/rocmfp2_reference.h @@ -0,0 +1,159 @@ +// ROCmFPx quant formats, hand-ported into this fork. +// +// Origin: https://github.com/charlie12345/ROCmFPX - creator of the ROCmFP4 format. +// Ported from: https://github.com/ciru-ai/ROCmFPX - a fork of the above. +// +// Both upstream projects are MIT licensed and based on llama.cpp; upstream authors +// retain their authorship and MIT license credit. See LICENSE. + +#ifndef ROCMFP2_REFERENCE_H +#define ROCMFP2_REFERENCE_H + +#include +#include +#include +#include + +#ifdef __cplusplus +extern "C" { +#endif + +enum { + ROCMFP2_P1_BLOCK_WEIGHTS = 32, + ROCMFP2_P1_GROUP_WEIGHTS = 16, + ROCMFP2_P1_DATA_BYTES = 8, + ROCMFP2_P1_SCALE_BYTES = 2, + ROCMFP2_P1_BLOCK_BYTES = 10, +}; + +/* + * Frozen Phase-1 P0_AOS serialization: + * + * d[0] ... d[7], s[0], s[1] + * + * Weight 4*j+i occupies bits 2*i+1:2*i of d[j]. Scale s[0] + * applies to weights 0..15 and s[1] to weights 16..31. This is a + * byte-canonical format; host word endianness is not part of serialization. + */ +typedef struct { + uint8_t d[ROCMFP2_P1_DATA_BYTES]; + uint8_t s[ROCMFP2_P1_SCALE_BYTES]; +} rocmfp2_p1_block; + +typedef struct { + uint8_t inner; + uint8_t outer; +} rocmfp2_p1_codebook; + +typedef enum { + ROCMFP2_P1_MAPPING_MORD = 0, + ROCMFP2_P1_MAPPING_MSM = 1, +} rocmfp2_p1_mapping; + +typedef enum { + ROCMFP2_P1_OK = 0, + ROCMFP2_P1_NONFINITE_SOURCE, + ROCMFP2_P1_INVALID_ARGUMENT, + ROCMFP2_P1_INVALID_CODEBOOK, + ROCMFP2_P1_INVALID_MAPPING, + ROCMFP2_P1_INVALID_SCALE_METADATA, + ROCMFP2_P1_INVALID_CODE, +} rocmfp2_p1_status; + +#if defined(__cplusplus) +static_assert(FLT_RADIX == 2 && DBL_MANT_DIG == 53 && DBL_MAX_EXP == 1024, + "ROCmFP2 Phase-1 reference requires IEEE-754 binary64 double"); +static_assert(sizeof(rocmfp2_p1_block) == ROCMFP2_P1_BLOCK_BYTES, + "ROCmFP2 P0 block has padding"); +static_assert(offsetof(rocmfp2_p1_block, d) == 0, "ROCmFP2 P0 data offset changed"); +static_assert(offsetof(rocmfp2_p1_block, s) == 8, "ROCmFP2 P0 scale offset changed"); +#else +_Static_assert(FLT_RADIX == 2 && DBL_MANT_DIG == 53 && DBL_MAX_EXP == 1024, + "ROCmFP2 Phase-1 reference requires IEEE-754 binary64 double"); +_Static_assert(sizeof(rocmfp2_p1_block) == ROCMFP2_P1_BLOCK_BYTES, + "ROCmFP2 P0 block has padding"); +_Static_assert(offsetof(rocmfp2_p1_block, d) == 0, "ROCmFP2 P0 data offset changed"); +_Static_assert(offsetof(rocmfp2_p1_block, s) == 8, "ROCmFP2 P0 scale offset changed"); +#endif + +const char * rocmfp2_p1_status_name(rocmfp2_p1_status status); + +bool rocmfp2_p1_scale_is_valid(uint8_t scale_byte); +double rocmfp2_p1_ue4m3_to_binary64(uint8_t scale_byte); +/* Returns invalid byte 0xff for NaN or a negative target. */ +uint8_t rocmfp2_p1_nearest_ue4m3(double target); + +rocmfp2_p1_status rocmfp2_p1_decode_code( + uint8_t code, + const rocmfp2_p1_codebook * codebook, + rocmfp2_p1_mapping mapping, + int8_t * value); + +rocmfp2_p1_status rocmfp2_p1_decode_packed_byte( + uint8_t packed, + const rocmfp2_p1_codebook * codebook, + rocmfp2_p1_mapping mapping, + int8_t values[4]); + +rocmfp2_p1_status rocmfp2_p1_pack_codes( + const uint8_t codes[ROCMFP2_P1_BLOCK_WEIGHTS], + const uint8_t scales[ROCMFP2_P1_SCALE_BYTES], + rocmfp2_p1_block * block); + +void rocmfp2_p1_unpack_codes( + const rocmfp2_p1_block * block, + uint8_t codes[ROCMFP2_P1_BLOCK_WEIGHTS]); + +bool rocmfp2_p1_validate_block(const rocmfp2_p1_block * block); +bool rocmfp2_p1_validate_serialized(const void * data, size_t nbytes); + +/* + * The code selectors are exposed so signed-zero and exact-midpoint behavior + * can be tested independently of scale selection. The reference selector + * exhaustively evaluates all four codes. The optimized selector uses the + * symmetric sign/magnitude structure. Both use binary64 arithmetic. + */ +rocmfp2_p1_status rocmfp2_p1_select_code_ref( + double source, + uint8_t scale_byte, + const rocmfp2_p1_codebook * codebook, + rocmfp2_p1_mapping mapping, + uint8_t * code); + +rocmfp2_p1_status rocmfp2_p1_select_code_optimized( + double source, + uint8_t scale_byte, + const rocmfp2_p1_codebook * codebook, + rocmfp2_p1_mapping mapping, + uint8_t * code); + +/* + * Both quantizers search every legal UE4M3 byte (0x00..0x7e) independently + * for each 16-weight group and minimize unweighted SSE in binary64. Exact + * scale ties select the lower byte. Outputs are committed only on success. + */ +rocmfp2_p1_status rocmfp2_p1_quantize_block_ref( + const double source[ROCMFP2_P1_BLOCK_WEIGHTS], + const rocmfp2_p1_codebook * codebook, + rocmfp2_p1_mapping mapping, + rocmfp2_p1_block * block, + double group_sse[ROCMFP2_P1_SCALE_BYTES]); + +rocmfp2_p1_status rocmfp2_p1_quantize_block_optimized( + const double source[ROCMFP2_P1_BLOCK_WEIGHTS], + const rocmfp2_p1_codebook * codebook, + rocmfp2_p1_mapping mapping, + rocmfp2_p1_block * block, + double group_sse[ROCMFP2_P1_SCALE_BYTES]); + +rocmfp2_p1_status rocmfp2_p1_dequantize_block( + const rocmfp2_p1_block * block, + const rocmfp2_p1_codebook * codebook, + rocmfp2_p1_mapping mapping, + double output[ROCMFP2_P1_BLOCK_WEIGHTS]); + +#ifdef __cplusplus +} +#endif + +#endif diff --git a/ggml/rocmfpx/rocmfpx.c b/ggml/rocmfpx/rocmfpx.c new file mode 100644 index 000000000000..8582a2038539 --- /dev/null +++ b/ggml/rocmfpx/rocmfpx.c @@ -0,0 +1,1364 @@ +// ROCmFPx quant formats, hand-ported into this fork. +// +// Origin: https://github.com/charlie12345/ROCmFPX - creator of the ROCmFP4 format. +// Ported from: https://github.com/ciru-ai/ROCmFPX - a fork of the above. +// +// Both upstream projects are MIT licensed and based on llama.cpp; upstream authors +// retain their authorship and MIT license credit. See LICENSE. + +#include "rocmfpx.h" + +#include +#include +#include + +// Finite unsigned E4M3 scale bytes decoded to FP32. Precomputed from the same +// exp/mant formula rocmfpx_ue4m3_to_fp32() used to evaluate with ldexpf(): +// exp == 0 -> mant * 2^-10 ; otherwise (8 + mant) * 2^(exp - 11). +// The scale search re-decodes candidate bytes for every block, and dequant +// decodes a scale for every element, so keeping this as a table (identical to +// the former per-call ldexpf result) removes the transcendental from both hot +// paths without changing any produced value. +#define ROCMFPX_SCALE_SUB(M) ((M) * 0x1p-10f) +#define ROCMFPX_SCALE_E(B, M) ((8 + (M)) * (B)) + +static const float rocmfpx_scale_ue4m3[127] = { + ROCMFPX_SCALE_SUB(0), ROCMFPX_SCALE_SUB(1), ROCMFPX_SCALE_SUB(2), ROCMFPX_SCALE_SUB(3), + ROCMFPX_SCALE_SUB(4), ROCMFPX_SCALE_SUB(5), ROCMFPX_SCALE_SUB(6), ROCMFPX_SCALE_SUB(7), + ROCMFPX_SCALE_E(0x1p-10f,0), ROCMFPX_SCALE_E(0x1p-10f,1), ROCMFPX_SCALE_E(0x1p-10f,2), ROCMFPX_SCALE_E(0x1p-10f,3), + ROCMFPX_SCALE_E(0x1p-10f,4), ROCMFPX_SCALE_E(0x1p-10f,5), ROCMFPX_SCALE_E(0x1p-10f,6), ROCMFPX_SCALE_E(0x1p-10f,7), + ROCMFPX_SCALE_E(0x1p-9f,0), ROCMFPX_SCALE_E(0x1p-9f,1), ROCMFPX_SCALE_E(0x1p-9f,2), ROCMFPX_SCALE_E(0x1p-9f,3), + ROCMFPX_SCALE_E(0x1p-9f,4), ROCMFPX_SCALE_E(0x1p-9f,5), ROCMFPX_SCALE_E(0x1p-9f,6), ROCMFPX_SCALE_E(0x1p-9f,7), + ROCMFPX_SCALE_E(0x1p-8f,0), ROCMFPX_SCALE_E(0x1p-8f,1), ROCMFPX_SCALE_E(0x1p-8f,2), ROCMFPX_SCALE_E(0x1p-8f,3), + ROCMFPX_SCALE_E(0x1p-8f,4), ROCMFPX_SCALE_E(0x1p-8f,5), ROCMFPX_SCALE_E(0x1p-8f,6), ROCMFPX_SCALE_E(0x1p-8f,7), + ROCMFPX_SCALE_E(0x1p-7f,0), ROCMFPX_SCALE_E(0x1p-7f,1), ROCMFPX_SCALE_E(0x1p-7f,2), ROCMFPX_SCALE_E(0x1p-7f,3), + ROCMFPX_SCALE_E(0x1p-7f,4), ROCMFPX_SCALE_E(0x1p-7f,5), ROCMFPX_SCALE_E(0x1p-7f,6), ROCMFPX_SCALE_E(0x1p-7f,7), + ROCMFPX_SCALE_E(0x1p-6f,0), ROCMFPX_SCALE_E(0x1p-6f,1), ROCMFPX_SCALE_E(0x1p-6f,2), ROCMFPX_SCALE_E(0x1p-6f,3), + ROCMFPX_SCALE_E(0x1p-6f,4), ROCMFPX_SCALE_E(0x1p-6f,5), ROCMFPX_SCALE_E(0x1p-6f,6), ROCMFPX_SCALE_E(0x1p-6f,7), + ROCMFPX_SCALE_E(0x1p-5f,0), ROCMFPX_SCALE_E(0x1p-5f,1), ROCMFPX_SCALE_E(0x1p-5f,2), ROCMFPX_SCALE_E(0x1p-5f,3), + ROCMFPX_SCALE_E(0x1p-5f,4), ROCMFPX_SCALE_E(0x1p-5f,5), ROCMFPX_SCALE_E(0x1p-5f,6), ROCMFPX_SCALE_E(0x1p-5f,7), + ROCMFPX_SCALE_E(0x1p-4f,0), ROCMFPX_SCALE_E(0x1p-4f,1), ROCMFPX_SCALE_E(0x1p-4f,2), ROCMFPX_SCALE_E(0x1p-4f,3), + ROCMFPX_SCALE_E(0x1p-4f,4), ROCMFPX_SCALE_E(0x1p-4f,5), ROCMFPX_SCALE_E(0x1p-4f,6), ROCMFPX_SCALE_E(0x1p-4f,7), + ROCMFPX_SCALE_E(0x1p-3f,0), ROCMFPX_SCALE_E(0x1p-3f,1), ROCMFPX_SCALE_E(0x1p-3f,2), ROCMFPX_SCALE_E(0x1p-3f,3), + ROCMFPX_SCALE_E(0x1p-3f,4), ROCMFPX_SCALE_E(0x1p-3f,5), ROCMFPX_SCALE_E(0x1p-3f,6), ROCMFPX_SCALE_E(0x1p-3f,7), + ROCMFPX_SCALE_E(0x1p-2f,0), ROCMFPX_SCALE_E(0x1p-2f,1), ROCMFPX_SCALE_E(0x1p-2f,2), ROCMFPX_SCALE_E(0x1p-2f,3), + ROCMFPX_SCALE_E(0x1p-2f,4), ROCMFPX_SCALE_E(0x1p-2f,5), ROCMFPX_SCALE_E(0x1p-2f,6), ROCMFPX_SCALE_E(0x1p-2f,7), + ROCMFPX_SCALE_E(0x1p-1f,0), ROCMFPX_SCALE_E(0x1p-1f,1), ROCMFPX_SCALE_E(0x1p-1f,2), ROCMFPX_SCALE_E(0x1p-1f,3), + ROCMFPX_SCALE_E(0x1p-1f,4), ROCMFPX_SCALE_E(0x1p-1f,5), ROCMFPX_SCALE_E(0x1p-1f,6), ROCMFPX_SCALE_E(0x1p-1f,7), + ROCMFPX_SCALE_E(0x1p0f,0), ROCMFPX_SCALE_E(0x1p0f,1), ROCMFPX_SCALE_E(0x1p0f,2), ROCMFPX_SCALE_E(0x1p0f,3), + ROCMFPX_SCALE_E(0x1p0f,4), ROCMFPX_SCALE_E(0x1p0f,5), ROCMFPX_SCALE_E(0x1p0f,6), ROCMFPX_SCALE_E(0x1p0f,7), + ROCMFPX_SCALE_E(0x1p1f,0), ROCMFPX_SCALE_E(0x1p1f,1), ROCMFPX_SCALE_E(0x1p1f,2), ROCMFPX_SCALE_E(0x1p1f,3), + ROCMFPX_SCALE_E(0x1p1f,4), ROCMFPX_SCALE_E(0x1p1f,5), ROCMFPX_SCALE_E(0x1p1f,6), ROCMFPX_SCALE_E(0x1p1f,7), + ROCMFPX_SCALE_E(0x1p2f,0), ROCMFPX_SCALE_E(0x1p2f,1), ROCMFPX_SCALE_E(0x1p2f,2), ROCMFPX_SCALE_E(0x1p2f,3), + ROCMFPX_SCALE_E(0x1p2f,4), ROCMFPX_SCALE_E(0x1p2f,5), ROCMFPX_SCALE_E(0x1p2f,6), ROCMFPX_SCALE_E(0x1p2f,7), + ROCMFPX_SCALE_E(0x1p3f,0), ROCMFPX_SCALE_E(0x1p3f,1), ROCMFPX_SCALE_E(0x1p3f,2), ROCMFPX_SCALE_E(0x1p3f,3), + ROCMFPX_SCALE_E(0x1p3f,4), ROCMFPX_SCALE_E(0x1p3f,5), ROCMFPX_SCALE_E(0x1p3f,6), ROCMFPX_SCALE_E(0x1p3f,7), + ROCMFPX_SCALE_E(0x1p4f,0), ROCMFPX_SCALE_E(0x1p4f,1), ROCMFPX_SCALE_E(0x1p4f,2), ROCMFPX_SCALE_E(0x1p4f,3), + ROCMFPX_SCALE_E(0x1p4f,4), ROCMFPX_SCALE_E(0x1p4f,5), ROCMFPX_SCALE_E(0x1p4f,6), +}; + +#undef ROCMFPX_SCALE_SUB +#undef ROCMFPX_SCALE_E + +float rocmfpx_ue4m3_to_fp32(uint8_t e) { + return rocmfpx_scale_is_valid(e) ? rocmfpx_scale_ue4m3[e] : 0.0f; +} + +bool rocmfpx_scale_is_valid(uint8_t e) { + return e <= 0x7e; +} + +// Precomputed table of all 127 valid UE4M3 scale values (e = 0x00..0x7e). +// Avoids repeated ldexpf calls in the hot MSE inner loops and in nearest-scale +// binary search. Matches the style of rocmfp4_scale_ue4m3_half in rocmfp4.c. +#define ROCMFPX_SCALE_SUB(M) ((M) * 0x1p-10f) +#define ROCMFPX_SCALE_E1(M) ((8 + (M)) * 0x1p-10f) +#define ROCMFPX_SCALE_E2(M) ((8 + (M)) * 0x1p-9f) +#define ROCMFPX_SCALE_E3(M) ((8 + (M)) * 0x1p-8f) +#define ROCMFPX_SCALE_E4(M) ((8 + (M)) * 0x1p-7f) +#define ROCMFPX_SCALE_E5(M) ((8 + (M)) * 0x1p-6f) +#define ROCMFPX_SCALE_E6(M) ((8 + (M)) * 0x1p-5f) +#define ROCMFPX_SCALE_E7(M) ((8 + (M)) * 0x1p-4f) +#define ROCMFPX_SCALE_E8(M) ((8 + (M)) * 0x1p-3f) +#define ROCMFPX_SCALE_E9(M) ((8 + (M)) * 0x1p-2f) +#define ROCMFPX_SCALE_E10(M) ((8 + (M)) * 0x1p-1f) +#define ROCMFPX_SCALE_E11(M) ((8 + (M)) * 0x1p0f) +#define ROCMFPX_SCALE_E12(M) ((8 + (M)) * 0x1p1f) +#define ROCMFPX_SCALE_E13(M) ((8 + (M)) * 0x1p2f) +#define ROCMFPX_SCALE_E14(M) ((8 + (M)) * 0x1p3f) +#define ROCMFPX_SCALE_E15(M) ((8 + (M)) * 0x1p4f) + +static const float rocmfpx_scale_table[127] = { + ROCMFPX_SCALE_SUB(0), ROCMFPX_SCALE_SUB(1), ROCMFPX_SCALE_SUB(2), ROCMFPX_SCALE_SUB(3), + ROCMFPX_SCALE_SUB(4), ROCMFPX_SCALE_SUB(5), ROCMFPX_SCALE_SUB(6), ROCMFPX_SCALE_SUB(7), + ROCMFPX_SCALE_E1(0), ROCMFPX_SCALE_E1(1), ROCMFPX_SCALE_E1(2), ROCMFPX_SCALE_E1(3), + ROCMFPX_SCALE_E1(4), ROCMFPX_SCALE_E1(5), ROCMFPX_SCALE_E1(6), ROCMFPX_SCALE_E1(7), + ROCMFPX_SCALE_E2(0), ROCMFPX_SCALE_E2(1), ROCMFPX_SCALE_E2(2), ROCMFPX_SCALE_E2(3), + ROCMFPX_SCALE_E2(4), ROCMFPX_SCALE_E2(5), ROCMFPX_SCALE_E2(6), ROCMFPX_SCALE_E2(7), + ROCMFPX_SCALE_E3(0), ROCMFPX_SCALE_E3(1), ROCMFPX_SCALE_E3(2), ROCMFPX_SCALE_E3(3), + ROCMFPX_SCALE_E3(4), ROCMFPX_SCALE_E3(5), ROCMFPX_SCALE_E3(6), ROCMFPX_SCALE_E3(7), + ROCMFPX_SCALE_E4(0), ROCMFPX_SCALE_E4(1), ROCMFPX_SCALE_E4(2), ROCMFPX_SCALE_E4(3), + ROCMFPX_SCALE_E4(4), ROCMFPX_SCALE_E4(5), ROCMFPX_SCALE_E4(6), ROCMFPX_SCALE_E4(7), + ROCMFPX_SCALE_E5(0), ROCMFPX_SCALE_E5(1), ROCMFPX_SCALE_E5(2), ROCMFPX_SCALE_E5(3), + ROCMFPX_SCALE_E5(4), ROCMFPX_SCALE_E5(5), ROCMFPX_SCALE_E5(6), ROCMFPX_SCALE_E5(7), + ROCMFPX_SCALE_E6(0), ROCMFPX_SCALE_E6(1), ROCMFPX_SCALE_E6(2), ROCMFPX_SCALE_E6(3), + ROCMFPX_SCALE_E6(4), ROCMFPX_SCALE_E6(5), ROCMFPX_SCALE_E6(6), ROCMFPX_SCALE_E6(7), + ROCMFPX_SCALE_E7(0), ROCMFPX_SCALE_E7(1), ROCMFPX_SCALE_E7(2), ROCMFPX_SCALE_E7(3), + ROCMFPX_SCALE_E7(4), ROCMFPX_SCALE_E7(5), ROCMFPX_SCALE_E7(6), ROCMFPX_SCALE_E7(7), + ROCMFPX_SCALE_E8(0), ROCMFPX_SCALE_E8(1), ROCMFPX_SCALE_E8(2), ROCMFPX_SCALE_E8(3), + ROCMFPX_SCALE_E8(4), ROCMFPX_SCALE_E8(5), ROCMFPX_SCALE_E8(6), ROCMFPX_SCALE_E8(7), + ROCMFPX_SCALE_E9(0), ROCMFPX_SCALE_E9(1), ROCMFPX_SCALE_E9(2), ROCMFPX_SCALE_E9(3), + ROCMFPX_SCALE_E9(4), ROCMFPX_SCALE_E9(5), ROCMFPX_SCALE_E9(6), ROCMFPX_SCALE_E9(7), + ROCMFPX_SCALE_E10(0), ROCMFPX_SCALE_E10(1), ROCMFPX_SCALE_E10(2), ROCMFPX_SCALE_E10(3), + ROCMFPX_SCALE_E10(4), ROCMFPX_SCALE_E10(5), ROCMFPX_SCALE_E10(6), ROCMFPX_SCALE_E10(7), + ROCMFPX_SCALE_E11(0), ROCMFPX_SCALE_E11(1), ROCMFPX_SCALE_E11(2), ROCMFPX_SCALE_E11(3), + ROCMFPX_SCALE_E11(4), ROCMFPX_SCALE_E11(5), ROCMFPX_SCALE_E11(6), ROCMFPX_SCALE_E11(7), + ROCMFPX_SCALE_E12(0), ROCMFPX_SCALE_E12(1), ROCMFPX_SCALE_E12(2), ROCMFPX_SCALE_E12(3), + ROCMFPX_SCALE_E12(4), ROCMFPX_SCALE_E12(5), ROCMFPX_SCALE_E12(6), ROCMFPX_SCALE_E12(7), + ROCMFPX_SCALE_E13(0), ROCMFPX_SCALE_E13(1), ROCMFPX_SCALE_E13(2), ROCMFPX_SCALE_E13(3), + ROCMFPX_SCALE_E13(4), ROCMFPX_SCALE_E13(5), ROCMFPX_SCALE_E13(6), ROCMFPX_SCALE_E13(7), + ROCMFPX_SCALE_E14(0), ROCMFPX_SCALE_E14(1), ROCMFPX_SCALE_E14(2), ROCMFPX_SCALE_E14(3), + ROCMFPX_SCALE_E14(4), ROCMFPX_SCALE_E14(5), ROCMFPX_SCALE_E14(6), ROCMFPX_SCALE_E14(7), + ROCMFPX_SCALE_E15(0), ROCMFPX_SCALE_E15(1), ROCMFPX_SCALE_E15(2), ROCMFPX_SCALE_E15(3), + ROCMFPX_SCALE_E15(4), ROCMFPX_SCALE_E15(5), ROCMFPX_SCALE_E15(6), +}; + +#undef ROCMFPX_SCALE_SUB +#undef ROCMFPX_SCALE_E1 +#undef ROCMFPX_SCALE_E2 +#undef ROCMFPX_SCALE_E3 +#undef ROCMFPX_SCALE_E4 +#undef ROCMFPX_SCALE_E5 +#undef ROCMFPX_SCALE_E6 +#undef ROCMFPX_SCALE_E7 +#undef ROCMFPX_SCALE_E8 +#undef ROCMFPX_SCALE_E9 +#undef ROCMFPX_SCALE_E10 +#undef ROCMFPX_SCALE_E11 +#undef ROCMFPX_SCALE_E12 +#undef ROCMFPX_SCALE_E13 +#undef ROCMFPX_SCALE_E14 +#undef ROCMFPX_SCALE_E15 + +// O(1) table lookup — use this everywhere inside this file instead of calling +// the public rocmfpx_ue4m3_to_fp32() which goes through the ldexpf path. +static inline float rocmfpx_scale_lookup(uint8_t e) { + return e <= 0x7e ? rocmfpx_scale_table[e] : 0.0f; +} + +// Binary search for the UE4M3 entry nearest to `target`. +// rocmfpx_scale_table is monotonically increasing in [0..0x7e], so we can +// binary-search to a two-element window then pick the closer neighbor. +// This replaces the original O(126) linear scan. Matches rocmfp4_nearest_scale_ue4m3. +static uint8_t rocmfpx_nearest_scale_ue4m3(float target) { + if (!(target > 0.0f) || !isfinite(target)) { + return 0; + } + + int lo = 1; + int hi = 126; + while (lo < hi) { + const int mid = lo + (hi - lo) / 2; + if (rocmfpx_scale_table[mid] < target) { + lo = mid + 1; + } else { + hi = mid; + } + } + + if (lo == 1) { + return 1; + } + + const float hi_scale = rocmfpx_scale_table[lo]; + const float lo_scale = rocmfpx_scale_table[lo - 1]; + + // Ties keep the lower (smaller) scale byte, matching the old ascending scan. + return (target - lo_scale <= hi_scale - target) ? (uint8_t)(lo - 1) : (uint8_t) lo; +} + + +size_t rocmfpx_row_size_fp3(int64_t k) { + assert(k % QK_ROCMFP3 == 0); + return (size_t) (k / QK_ROCMFP3) * sizeof(block_rocmfp3); +} + +size_t rocmfpx_row_size_fp2(int64_t k) { + assert(k % QK_ROCMFP2 == 0); + return (size_t) (k / QK_ROCMFP2) * sizeof(block_rocmfp2); +} + +size_t rocmfpx_row_size_fp6(int64_t k) { + assert(k % QK_ROCMFP6 == 0); + return (size_t) (k / QK_ROCMFP6) * sizeof(block_rocmfp6); +} + +size_t rocmfpx_row_size_fp8(int64_t k) { + assert(k % QK_ROCMFP8 == 0); + return (size_t) (k / QK_ROCMFP8) * sizeof(block_rocmfp8); +} + + + +static float rocmfpx_max_abs(const float * x, int n) { + float max_abs = 0.0f; + + for (int i = 0; i < n; ++i) { + if (!isfinite(x[i])) { + continue; + } + + const float ax = fabsf(x[i]); + if (ax > max_abs) { + max_abs = ax; + } + } + + return max_abs; +} + +static void rocmfpx_prepare_mse_weights( + float * dst, const float * x, int n, const float * quant_weights, float sigma2, + float * max_abs, float * max_abs_weight, bool * all_finite) { + *max_abs = 0.0f; + *max_abs_weight = 0.0f; + *all_finite = true; + + for (int i = 0; i < n; ++i) { + const float ax = fabsf(x[i]); + const float qw = quant_weights[i]; + const float weight = isfinite(qw) && qw > 0.0f && isfinite(x[i]) ? qw * sqrtf(sigma2 + x[i]*x[i]) : 0.0f; + *all_finite = *all_finite && isfinite(x[i]); + + if (isfinite(x[i])) { + if (ax > *max_abs) { + *max_abs = ax; + *max_abs_weight = weight; + } else if (ax == *max_abs && weight > *max_abs_weight) { + *max_abs_weight = weight; + } + } + + // Match llama.cpp imatrix weighting style: calibration importance is + // scaled by row energy so large activations stay protected. + dst[i] = weight; + } +} + +// ROCmFP2 S40 uses the frozen MORD code order {-4, -1, +1, +4}. +static inline int rocmfpx_decode_fp2_code(uint8_t code) { + static const int8_t values[4] = { -4, -1, 1, 4 }; + return values[code & 3u]; +} + +static inline uint8_t rocmfpx_quantize_fp2_code(float x, float inv_scale) { + if (!isfinite(x) || !(inv_scale > 0.0f)) { + return 0; + } + + const float magnitude = fabsf(x * inv_scale); + const bool outer = magnitude > 2.5f; + if (signbit(x)) { + return outer ? 0u : 1u; + } + return outer ? 3u : 2u; +} + +static float rocmfpx_fp2_group_mse_for_scale( + const float * x, const float * mse_weights, int n, uint8_t e, float best_err) { + const float scale = rocmfpx_scale_lookup(e); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + if (!isfinite(x[i])) { + continue; + } + const float reconstructed = (float) rocmfpx_decode_fp2_code(rocmfpx_quantize_fp2_code(x[i], inv_scale)) * scale; + const float delta = x[i] - reconstructed; + err += (mse_weights ? mse_weights[i] : 1.0f) * delta * delta; + if (err > best_err) { + return err; + } + } + return err; +} + +static uint8_t rocmfpx_choose_scale_fp2_mse( + const float * x, int n, const float * quant_weights, float sigma2) { + float mse_weights[QK_ROCMFP2/2]; + float max_abs = 0.0f; + float max_abs_weight = 0.0f; + bool all_finite = true; + + if (quant_weights) { + rocmfpx_prepare_mse_weights( + mse_weights, x, n, quant_weights, sigma2, + &max_abs, &max_abs_weight, &all_finite); + } else { + max_abs = rocmfpx_max_abs(x, n); + max_abs_weight = 1.0f; + } + GGML_UNUSED(all_finite); + + if (!(max_abs > 0.0f) || !isfinite(max_abs)) { + return 0; + } + + const float * weights = quant_weights ? mse_weights : NULL; + + // The S40 codebook {-4, -1, +1, +4} has no zero code, so at any nonzero scale + // the smallest magnitude ROCmFP2 can emit is one scale step. A block that sits + // below half of the smallest UE4M3 scale is therefore encoded better by the + // zero scale than by e = 1, which would amplify every value. Seed the search + // with that candidate so scales 1..126 have to beat it; ties keep byte 0, + // matching the lower-byte rule the rest of the search uses. + uint8_t best_e = 0; + float best_err = rocmfpx_fp2_group_mse_for_scale(x, weights, n, 0, INFINITY); + + const uint8_t start_e = rocmfpx_nearest_scale_ue4m3(max_abs / 4.0f); + bool lower_done = false; + + for (int delta = 0; delta <= 125; ++delta) { + const int e0 = (int) start_e - delta; + if (!lower_done && e0 >= 1 && e0 <= 126) { + const float scale = rocmfpx_scale_lookup((uint8_t) e0); + const float clip_delta = max_abs - 4.0f * scale; + const float clip_err = max_abs_weight * clip_delta * clip_delta; + if (clip_delta > 0.0f && clip_err > best_err) { + lower_done = true; + } else { + const float err = rocmfpx_fp2_group_mse_for_scale(x, weights, n, (uint8_t) e0, best_err); + if (err < best_err || (err == best_err && e0 < best_e)) { + best_err = err; + best_e = (uint8_t) e0; + } + } + } + + const int e1 = (int) start_e + delta; + if (delta != 0 && e1 >= 1 && e1 <= 126) { + const float err = rocmfpx_fp2_group_mse_for_scale(x, weights, n, (uint8_t) e1, best_err); + if (err < best_err || (err == best_err && e1 < best_e)) { + best_err = err; + best_e = (uint8_t) e1; + } + } + + if ((lower_done || e0 <= 1) && e1 >= 126) { + break; + } + } + return best_e; +} + +static void rocmfpx_quantize_row_fp2_impl( + const float * GGML_RESTRICT x, block_rocmfp2 * GGML_RESTRICT y, + int64_t k, const float * GGML_RESTRICT quant_weights) { + assert(k % QK_ROCMFP2 == 0); + + float sum_x2 = 0.0f; + for (int64_t i = 0; i < k; ++i) { + sum_x2 += isfinite(x[i]) ? x[i] * x[i] : 0.0f; + } + const float sigma2 = sum_x2 / (float) k; + + const int64_t nb = k / QK_ROCMFP2; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib * QK_ROCMFP2; + const float * qw = quant_weights ? quant_weights + ib * QK_ROCMFP2 : NULL; + block_rocmfp2 * yb = y + ib; + + for (int half = 0; half < 2; ++half) { + const int half_off = half * (QK_ROCMFP2 / 2); + const float * xh = xb + half_off; + const float * qh = qw ? qw + half_off : NULL; + yb->e[half] = rocmfpx_choose_scale_fp2_mse(xh, QK_ROCMFP2 / 2, qh, sigma2); + + const float scale = rocmfpx_scale_lookup(yb->e[half]); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + for (int packed = 0; packed < 4; ++packed) { + uint8_t byte = 0; + for (int lane = 0; lane < 4; ++lane) { + const int j = 4 * packed + lane; + byte |= (uint8_t) (rocmfpx_quantize_fp2_code(xh[j], inv_scale) << (2 * lane)); + } + yb->qs[half * 4 + packed] = byte; + } + } + } +} + +void rocmfpx_quantize_row_fp2_ref( + const float * GGML_RESTRICT x, block_rocmfp2 * GGML_RESTRICT y, int64_t k) { + rocmfpx_quantize_row_fp2_impl(x, y, k, NULL); +} + +void rocmfpx_dequantize_row_fp2( + const block_rocmfp2 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP2 == 0); + const int64_t nb = k / QK_ROCMFP2; + for (int64_t ib = 0; ib < nb; ++ib) { + const block_rocmfp2 * xb = x + ib; + float * yb = y + ib * QK_ROCMFP2; + for (int half = 0; half < 2; ++half) { + const float scale = rocmfpx_scale_lookup(xb->e[half]); + for (int j = 0; j < QK_ROCMFP2 / 2; ++j) { + const uint8_t code = (xb->qs[half * 4 + j / 4] >> (2 * (j % 4))) & 3u; + yb[half * (QK_ROCMFP2 / 2) + j] = (float) rocmfpx_decode_fp2_code(code) * scale; + } + } + } +} + +void rocmfpx_quantize_row_fp2( + const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + rocmfpx_quantize_row_fp2_ref(x, (block_rocmfp2 *) y, k); +} + +size_t rocmfpx_quantize_fp2( + const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, + int64_t nrows, int64_t n_per_row, const float * imatrix) { + const size_t row_size = rocmfpx_row_size_fp2(n_per_row); + char * qrow = (char *) dst; + for (int64_t row = 0; row < nrows; ++row) { + rocmfpx_quantize_row_fp2_impl( + src + row * n_per_row, (block_rocmfp2 *) qrow, n_per_row, imatrix); + qrow += row_size; + } + return (size_t) nrows * row_size; +} + +// --------------------------------------------------------------------------- +// FP3 group pack/unpack: 8 × 3-bit codes → 3 bytes (24 bits) +// +// Bit layout across 3 bytes (LSB first within each byte): +// byte 0: [v0 2:0][v1 2:0][v2 1:0] +// byte 1: [v2.2 ][v3 2:0][v4 2:0][v5.0] +// byte 2: [v5 2:1][v6 2:0][v7 2:0] +// +// A 16-element half-block is two consecutive pack8 groups. +// Half 0 → qs[0..5], Half 1 → qs[6..11]. +// --------------------------------------------------------------------------- +static inline void rocmfpx_fp3_pack8(uint8_t * dst, const uint8_t * c) { + dst[0] = (uint8_t)( (c[0] & 7u) | ((c[1] & 7u) << 3) | ((c[2] & 3u) << 6) ); + dst[1] = (uint8_t)( ((c[2] >> 2) & 1u) | ((c[3] & 7u) << 1) | ((c[4] & 7u) << 4) | ((c[5] & 1u) << 7) ); + dst[2] = (uint8_t)( ((c[5] >> 1) & 3u) | ((c[6] & 7u) << 2) | ((c[7] & 7u) << 5) ); +} + +static inline void rocmfpx_fp3_unpack8(const uint8_t * src, uint8_t * c) { + c[0] = src[0] & 7u; + c[1] = (src[0] >> 3) & 7u; + c[2] = ((src[0] >> 6) & 3u) | ((src[1] & 1u) << 2); + c[3] = (src[1] >> 1) & 7u; + c[4] = (src[1] >> 4) & 7u; + c[5] = ((src[1] >> 7) & 1u) | ((src[2] & 3u) << 1); + c[6] = (src[2] >> 2) & 7u; + c[7] = (src[2] >> 5) & 7u; +} + +// --------------------------------------------------------------------------- +// FP6 group pack/unpack: 4 × 6-bit codes → 3 bytes (24 bits) +// +// Bit layout: +// byte 0: [v0 5:0][v1 1:0] +// byte 1: [v1 5:2][v2 3:0] +// byte 2: [v2 5:4][v3 5:0] +// +// A 16-element half-block is four consecutive pack4 groups. +// Half 0 → qs[0..11], Half 1 → qs[12..23]. +// --------------------------------------------------------------------------- +static inline void rocmfpx_fp6_pack4(uint8_t * dst, const uint8_t * c) { + dst[0] = (uint8_t)( (c[0] & 0x3fu) | ((c[1] & 0x03u) << 6) ); + dst[1] = (uint8_t)( ((c[1] >> 2) & 0x0fu) | ((c[2] & 0x0fu) << 4) ); + dst[2] = (uint8_t)( ((c[2] >> 4) & 0x03u) | ((c[3] & 0x3fu) << 2) ); +} + +static inline void rocmfpx_fp6_unpack4(const uint8_t * src, uint8_t * c) { + c[0] = src[0] & 0x3fu; + c[1] = ((src[0] >> 6) & 0x03u) | ((src[1] & 0x0fu) << 2); + c[2] = ((src[1] >> 4) & 0x0fu) | ((src[2] & 0x03u) << 4); + c[3] = (src[2] >> 2) & 0x3fu; +} + +static int rocmfpx_decode_fp3_code(uint8_t code) { + static const int mag[4] = { 0, 1, 2, 4 }; + const int value = mag[code & 3u]; + return (code & 4u) ? -value : value; +} + +static uint8_t rocmfpx_quantize_fp3_code(float x, float inv_scale) { + if (!isfinite(x) || inv_scale <= 0.0f) { + return 0; + } + + const float ax = fabsf(x * inv_scale); + uint8_t mag; + + if (ax <= 0.5f) { + mag = 0; + } else if (ax <= 1.5f) { + mag = 1; + } else if (ax <= 3.0f) { + mag = 2; + } else { + mag = 3; + } + + return mag == 0 ? 0 : (uint8_t) ((x < 0.0f ? 4u : 0u) | mag); +} + +// Fused threshold + decode for finite values in the exhaustive scale search. +// The result exactly matches quantize_fp3_code followed by decode_fp3_code. +static inline float rocmfpx_fp3_decoded_mag(float x, float inv_scale) { + const float a = fabsf(x * inv_scale); + float mag; + if (a <= 0.5f) { + return 0.0f; + } else if (a <= 1.5f) { + mag = 1.0f; + } else if (a <= 3.0f) { + mag = 2.0f; + } else { + mag = 4.0f; + } + return x < 0.0f ? -mag : mag; +} + +static float rocmfpx_fp3_block_mse_for_scale(const float * x, int n, uint8_t e, float best_err) { + const float scale = rocmfpx_scale_lookup(e); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + if (!isfinite(x[i])) { + continue; + } + + const float y = rocmfpx_fp3_decoded_mag(x[i], inv_scale) * scale; + const float d = x[i] - y; + + err += d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +// Fast path: caller guarantees all x[i] are finite — skips isfinite() per element. +static float rocmfpx_fp3_block_mse_for_scale_finite(const float * x, int n, uint8_t e, float best_err) { + const float scale = rocmfpx_scale_lookup(e); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + const float y = rocmfpx_fp3_decoded_mag(x[i], inv_scale) * scale; + const float d = x[i] - y; + + err += d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +static float rocmfpx_fp3_block_weighted_mse_for_scale(const float * x, int n, const float * mse_weights, uint8_t e, float best_err) { + const float scale = rocmfpx_scale_lookup(e); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + if (!isfinite(x[i])) { + continue; + } + + const float y = rocmfpx_fp3_decoded_mag(x[i], inv_scale) * scale; + const float d = x[i] - y; + + err += mse_weights[i]*d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +// Fast path: caller guarantees all x[i] are finite. +static float rocmfpx_fp3_block_weighted_mse_for_scale_finite(const float * x, int n, const float * mse_weights, uint8_t e, float best_err) { + const float scale = rocmfpx_scale_lookup(e); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + const float y = rocmfpx_fp3_decoded_mag(x[i], inv_scale) * scale; + const float d = x[i] - y; + + err += mse_weights[i]*d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +static uint8_t rocmfpx_choose_scale_fp3_mse_impl( + const float * x, int n, const float * mse_weights, + float max_abs, float max_abs_weight, bool all_finite) { + const uint8_t start_e = rocmfpx_nearest_scale_ue4m3(max_abs / 4.0f); + uint8_t best_e = start_e; + float best_err = INFINITY; + bool lower_done = false; + + for (int delta = 0; delta <= 125; ++delta) { + const int e0 = (int) start_e - delta; + if (!lower_done && e0 >= 1 && e0 <= 126) { + const float scale = rocmfpx_scale_lookup((uint8_t) e0); + const float clip_delta = max_abs - 4.0f*scale; + const float clip_err = mse_weights ? max_abs_weight*clip_delta*clip_delta : clip_delta*clip_delta; + if (clip_delta > 0.0f && clip_err > best_err) { + lower_done = true; + } else { + const float err = mse_weights ? + (all_finite ? + rocmfpx_fp3_block_weighted_mse_for_scale_finite(x, n, mse_weights, (uint8_t) e0, best_err) : + rocmfpx_fp3_block_weighted_mse_for_scale(x, n, mse_weights, (uint8_t) e0, best_err)) : + (all_finite ? + rocmfpx_fp3_block_mse_for_scale_finite(x, n, (uint8_t) e0, best_err) : + rocmfpx_fp3_block_mse_for_scale(x, n, (uint8_t) e0, best_err)); + if (err < best_err || (err == best_err && e0 < best_e)) { + best_err = err; + best_e = (uint8_t) e0; + } + } + } + + const int e1 = (int) start_e + delta; + if (delta != 0 && e1 >= 1 && e1 <= 126) { + const float err = mse_weights ? + (all_finite ? + rocmfpx_fp3_block_weighted_mse_for_scale_finite(x, n, mse_weights, (uint8_t) e1, best_err) : + rocmfpx_fp3_block_weighted_mse_for_scale(x, n, mse_weights, (uint8_t) e1, best_err)) : + (all_finite ? + rocmfpx_fp3_block_mse_for_scale_finite(x, n, (uint8_t) e1, best_err) : + rocmfpx_fp3_block_mse_for_scale(x, n, (uint8_t) e1, best_err)); + if (err < best_err || (err == best_err && e1 < best_e)) { + best_err = err; + best_e = (uint8_t) e1; + } + } + + if ((lower_done || e0 <= 1) && e1 >= 126) { + break; + } + } + + return best_e; +} + +static uint8_t rocmfpx_choose_scale_fp3_mse(const float * x, int n) { + const float max_abs = rocmfpx_max_abs(x, n); + if (!(max_abs > 0.0f) || !isfinite(max_abs)) { + return 0; + } + // rocmfpx_max_abs skips non-finite, but we need to know if all are finite + // to select the fast path. Use a simple scan here since max_abs already ran. + bool all_finite = true; + for (int i = 0; i < n; ++i) { + all_finite = all_finite && isfinite(x[i]); + } + + return rocmfpx_choose_scale_fp3_mse_impl(x, n, NULL, max_abs, 0.0f, all_finite); +} + +static uint8_t rocmfpx_choose_scale_fp3_weighted_mse(const float * x, int n, const float * quant_weights, float sigma2) { + assert(n <= QK_ROCMFP3); + float mse_weights[QK_ROCMFP3]; + float max_abs; + float max_abs_weight; + bool all_finite; + rocmfpx_prepare_mse_weights(mse_weights, x, n, quant_weights, sigma2, &max_abs, &max_abs_weight, &all_finite); + if (!(max_abs > 0.0f) || !isfinite(max_abs)) { + return 0; + } + + return rocmfpx_choose_scale_fp3_mse_impl(x, n, mse_weights, max_abs, max_abs_weight, all_finite); +} + +// Round to the nearest integer inside [lo, hi], clamping *before* the conversion. +// Converting first and clamping afterwards overflows int (and, for large enough +// inputs, long) so an extreme finite weight could flip sign or collapse to zero. +// For in-range values this is identical to rounding then clamping. +static inline int rocmfpx_round_clamp(float v, float lo, float hi) { + const float c = v < lo ? lo : (v > hi ? hi : v); + return (int) lroundf(c); +} + +static int rocmfpx_decode_fp6_code(uint8_t code) { + const int mag = code & 31u; + return (code & 32u) ? -(mag == 0 ? 32 : mag) : mag; +} + +static uint8_t rocmfpx_quantize_fp6_code(float x, float inv_scale) { + if (!isfinite(x) || !isfinite(inv_scale) || inv_scale <= 0.0f) { + return 0; + } + + const int q = rocmfpx_round_clamp(x * inv_scale, -32.0f, 31.0f); + + return q == 0 ? 0 : (uint8_t) (q < 0 ? (32u | ((uint8_t) -q & 31u)) : (uint8_t) q); +} + +// Fused round, clamp, and decode for finite values in the scale search. Keep +// current main's asymmetric signed range [-32, 31], including the encoded -32 +// endpoint, rather than the older experimental branch's [-31, 31] behavior. +static inline float rocmfpx_fp6_decoded_value(float x, float inv_scale) { + return (float) rocmfpx_round_clamp(x * inv_scale, -32.0f, 31.0f); +} + +static float rocmfpx_fp6_block_mse_for_scale(const float * x, int n, uint8_t e, float best_err) { + const float scale = rocmfpx_scale_lookup(e); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + if (!isfinite(x[i])) { + continue; + } + const float y = rocmfpx_fp6_decoded_value(x[i], inv_scale) * scale; + const float d = x[i] - y; + err += d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +// Fast path: caller guarantees all x[i] are finite. +static float rocmfpx_fp6_block_mse_for_scale_finite(const float * x, int n, uint8_t e, float best_err) { + const float scale = rocmfpx_scale_lookup(e); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + const float y = rocmfpx_fp6_decoded_value(x[i], inv_scale) * scale; + const float d = x[i] - y; + err += d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +static float rocmfpx_fp6_block_weighted_mse_for_scale(const float * x, int n, const float * mse_weights, uint8_t e, float best_err) { + const float scale = rocmfpx_scale_lookup(e); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + if (!isfinite(x[i])) { + continue; + } + const float y = rocmfpx_fp6_decoded_value(x[i], inv_scale) * scale; + const float d = x[i] - y; + err += mse_weights[i]*d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +// Fast path: caller guarantees all x[i] are finite. +static float rocmfpx_fp6_block_weighted_mse_for_scale_finite(const float * x, int n, const float * mse_weights, uint8_t e, float best_err) { + const float scale = rocmfpx_scale_lookup(e); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + const float y = rocmfpx_fp6_decoded_value(x[i], inv_scale) * scale; + const float d = x[i] - y; + err += mse_weights[i]*d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +static uint8_t rocmfpx_choose_scale_fp6_mse_impl( + const float * x, int n, const float * mse_weights, + float max_abs, float max_abs_weight, bool all_finite) { + const uint8_t start_e = rocmfpx_nearest_scale_ue4m3(max_abs / 31.0f); + uint8_t best_e = start_e; + float best_err = INFINITY; + bool lower_done = false; + + for (int delta = 0; delta <= 125; ++delta) { + const int e0 = (int) start_e - delta; + if (!lower_done && e0 >= 1 && e0 <= 126) { + const float scale = rocmfpx_scale_lookup((uint8_t) e0); + // ROCmFP6 reaches -32, not just 31, so bound the unavoidable clipping + // error by 32: at 31 the bound is too pessimistic for a block whose + // largest magnitude is negative, and the search stops before reaching + // the scale that actually wins. + const float clip_delta = max_abs - 32.0f*scale; + const float clip_err = mse_weights ? max_abs_weight*clip_delta*clip_delta : clip_delta*clip_delta; + if (clip_delta > 0.0f && clip_err > best_err) { + lower_done = true; + } else { + const float err = mse_weights ? + (all_finite ? + rocmfpx_fp6_block_weighted_mse_for_scale_finite(x, n, mse_weights, (uint8_t) e0, best_err) : + rocmfpx_fp6_block_weighted_mse_for_scale(x, n, mse_weights, (uint8_t) e0, best_err)) : + (all_finite ? + rocmfpx_fp6_block_mse_for_scale_finite(x, n, (uint8_t) e0, best_err) : + rocmfpx_fp6_block_mse_for_scale(x, n, (uint8_t) e0, best_err)); + if (err < best_err || (err == best_err && e0 < best_e)) { + best_err = err; + best_e = (uint8_t) e0; + } + } + } + + const int e1 = (int) start_e + delta; + if (delta != 0 && e1 >= 1 && e1 <= 126) { + const float err = mse_weights ? + (all_finite ? + rocmfpx_fp6_block_weighted_mse_for_scale_finite(x, n, mse_weights, (uint8_t) e1, best_err) : + rocmfpx_fp6_block_weighted_mse_for_scale(x, n, mse_weights, (uint8_t) e1, best_err)) : + (all_finite ? + rocmfpx_fp6_block_mse_for_scale_finite(x, n, (uint8_t) e1, best_err) : + rocmfpx_fp6_block_mse_for_scale(x, n, (uint8_t) e1, best_err)); + if (err < best_err || (err == best_err && e1 < best_e)) { + best_err = err; + best_e = (uint8_t) e1; + } + } + + if ((lower_done || e0 <= 1) && e1 >= 126) { + break; + } + } + + return best_e; +} + +static uint8_t rocmfpx_choose_scale_fp6_mse(const float * x, int n) { + const float max_abs = rocmfpx_max_abs(x, n); + if (!(max_abs > 0.0f) || !isfinite(max_abs)) { + return 0; + } + bool all_finite = true; + for (int i = 0; i < n; ++i) { + all_finite = all_finite && isfinite(x[i]); + } + + return rocmfpx_choose_scale_fp6_mse_impl(x, n, NULL, max_abs, 0.0f, all_finite); +} + +static uint8_t rocmfpx_choose_scale_fp6_weighted_mse(const float * x, int n, const float * quant_weights, float sigma2) { + assert(n <= QK_ROCMFP6); + float mse_weights[QK_ROCMFP6]; + float max_abs; + float max_abs_weight; + bool all_finite; + rocmfpx_prepare_mse_weights(mse_weights, x, n, quant_weights, sigma2, &max_abs, &max_abs_weight, &all_finite); + if (!(max_abs > 0.0f) || !isfinite(max_abs)) { + return 0; + } + + return rocmfpx_choose_scale_fp6_mse_impl(x, n, mse_weights, max_abs, max_abs_weight, all_finite); +} + +static int8_t rocmfpx_quantize_fp8_code(float x, float inv_scale) { + if (!isfinite(x) || inv_scale <= 0.0f) { + return 0; + } + + return (int8_t) rocmfpx_round_clamp(x * inv_scale, -127.0f, 127.0f); +} + +static float rocmfpx_fp8_block_weighted_mse_for_scale(const float * x, int n, const float * mse_weights, uint8_t e, float best_err) { + const float scale = rocmfpx_scale_lookup(e); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + if (!isfinite(x[i])) { + continue; + } + + const int8_t code = rocmfpx_quantize_fp8_code(x[i], inv_scale); + const float y = (float) code * scale; + const float d = x[i] - y; + + err += mse_weights[i]*d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +// Fast path: caller guarantees all x[i] are finite. +static float rocmfpx_fp8_block_weighted_mse_for_scale_finite(const float * x, int n, const float * mse_weights, uint8_t e, float best_err) { + const float scale = rocmfpx_scale_lookup(e); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + const int8_t code = rocmfpx_quantize_fp8_code(x[i], inv_scale); + const float y = (float) code * scale; + const float d = x[i] - y; + + err += mse_weights[i]*d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +static uint8_t rocmfpx_choose_scale_fp8_weighted_mse(const float * x, int n, const float * quant_weights, float sigma2) { + assert(n <= QK_ROCMFP8); + float mse_weights[QK_ROCMFP8]; + float max_abs; + float max_abs_weight; + bool all_finite; + rocmfpx_prepare_mse_weights(mse_weights, x, n, quant_weights, sigma2, &max_abs, &max_abs_weight, &all_finite); + if (!(max_abs > 0.0f) || !isfinite(max_abs)) { + return 0; + } + + const uint8_t start_e = rocmfpx_nearest_scale_ue4m3(max_abs / 127.0f); + uint8_t best_e = start_e; + float best_err = INFINITY; + bool lower_done = false; + + for (int delta = 0; delta <= 125; ++delta) { + const int e0 = (int) start_e - delta; + if (!lower_done && e0 >= 1 && e0 <= 126) { + const float scale = rocmfpx_scale_lookup((uint8_t) e0); + const float clip_delta = max_abs - 127.0f*scale; + if (clip_delta > 0.0f && max_abs_weight*clip_delta*clip_delta > best_err) { + lower_done = true; + } else { + const float err = all_finite ? + rocmfpx_fp8_block_weighted_mse_for_scale_finite(x, n, mse_weights, (uint8_t) e0, best_err) : + rocmfpx_fp8_block_weighted_mse_for_scale(x, n, mse_weights, (uint8_t) e0, best_err); + if (err < best_err || (err == best_err && e0 < best_e)) { + best_err = err; + best_e = (uint8_t) e0; + } + } + } + + const int e1 = (int) start_e + delta; + if (delta != 0 && e1 >= 1 && e1 <= 126) { + const float err = all_finite ? + rocmfpx_fp8_block_weighted_mse_for_scale_finite(x, n, mse_weights, (uint8_t) e1, best_err) : + rocmfpx_fp8_block_weighted_mse_for_scale(x, n, mse_weights, (uint8_t) e1, best_err); + if (err < best_err || (err == best_err && e1 < best_e)) { + best_err = err; + best_e = (uint8_t) e1; + } + } + + if ((lower_done || e0 <= 1) && e1 >= 126) { + break; + } + } + + return best_e; +} + +void rocmfpx_quantize_row_fp3_ref(const float * GGML_RESTRICT x, block_rocmfp3 * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP3 == 0); + + const int64_t nb = k / QK_ROCMFP3; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib*QK_ROCMFP3; + block_rocmfp3 * yb = y + ib; + + for (int half = 0; half < 2; ++half) { + const int half_off = half * (QK_ROCMFP3/2); + const float * xh = xb + half_off; + yb->e[half] = rocmfpx_choose_scale_fp3_mse(xh, QK_ROCMFP3/2); + + const float scale = rocmfpx_scale_lookup(yb->e[half]); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + + // Compute 16 codes then pack into 6 bytes (2 groups of 8). + uint8_t codes[QK_ROCMFP3/2]; + for (int j = 0; j < QK_ROCMFP3/2; ++j) { + codes[j] = rocmfpx_quantize_fp3_code(xh[j], inv_scale); + } + rocmfpx_fp3_pack8(yb->qs + half_off*3/8, codes); + rocmfpx_fp3_pack8(yb->qs + half_off*3/8 + 3, codes + 8); + } + } +} + +static void rocmfpx_quantize_row_fp3_weighted( + const float * GGML_RESTRICT x, block_rocmfp3 * GGML_RESTRICT y, int64_t k, const float * GGML_RESTRICT quant_weights) { + assert(k % QK_ROCMFP3 == 0); + + float sum_x2 = 0.0f; + for (int64_t i = 0; i < k; ++i) { + sum_x2 += isfinite(x[i]) ? x[i]*x[i] : 0.0f; + } + const float sigma2 = sum_x2 / (float) k; + + const int64_t nb = k / QK_ROCMFP3; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib*QK_ROCMFP3; + const float * qw = quant_weights ? quant_weights + ib*QK_ROCMFP3 : NULL; + block_rocmfp3 * yb = y + ib; + + for (int half = 0; half < 2; ++half) { + const int half_off = half * (QK_ROCMFP3/2); + const float * xh = xb + half_off; + const float * qh = qw ? qw + half_off : NULL; + yb->e[half] = qh ? + rocmfpx_choose_scale_fp3_weighted_mse(xh, QK_ROCMFP3/2, qh, sigma2) : + rocmfpx_choose_scale_fp3_mse(xh, QK_ROCMFP3/2); + + const float scale = rocmfpx_scale_lookup(yb->e[half]); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + + uint8_t codes[QK_ROCMFP3/2]; + for (int j = 0; j < QK_ROCMFP3/2; ++j) { + codes[j] = rocmfpx_quantize_fp3_code(xh[j], inv_scale); + } + rocmfpx_fp3_pack8(yb->qs + half_off*3/8, codes); + rocmfpx_fp3_pack8(yb->qs + half_off*3/8 + 3, codes + 8); + } + } +} + +void rocmfpx_dequantize_row_fp3(const block_rocmfp3 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP3 == 0); + + const int64_t nb = k / QK_ROCMFP3; + for (int64_t ib = 0; ib < nb; ++ib) { + const block_rocmfp3 * xb = x + ib; + float * yb = y + ib*QK_ROCMFP3; + + // Unpack all 32 codes in 4 groups of 8 (4 × 3 bytes). + uint8_t codes[QK_ROCMFP3]; + rocmfpx_fp3_unpack8(xb->qs, codes); + rocmfpx_fp3_unpack8(xb->qs + 3, codes + 8); + rocmfpx_fp3_unpack8(xb->qs + 6, codes + 16); + rocmfpx_fp3_unpack8(xb->qs + 9, codes + 24); + + for (int half = 0; half < 2; ++half) { + const float scale = rocmfpx_scale_lookup(xb->e[half]); + for (int j = 0; j < QK_ROCMFP3/2; ++j) { + const int i = half*(QK_ROCMFP3/2) + j; + yb[i] = (float) rocmfpx_decode_fp3_code(codes[i]) * scale; + } + } + } +} + +void rocmfpx_quantize_row_fp3(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + rocmfpx_quantize_row_fp3_ref(x, (block_rocmfp3 *) y, k); +} + +size_t rocmfpx_quantize_fp3(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix) { + const size_t row_size = rocmfpx_row_size_fp3(n_per_row); + char * qrow = (char *) dst; + + for (int64_t row = 0; row < nrows; ++row) { + if (imatrix) { + rocmfpx_quantize_row_fp3_weighted(src + row*n_per_row, (block_rocmfp3 *) qrow, n_per_row, imatrix); + } else { + rocmfpx_quantize_row_fp3_ref(src + row*n_per_row, (block_rocmfp3 *) qrow, n_per_row); + } + qrow += row_size; + } + + return (size_t) nrows * row_size; +} + +static void rocmfpx_quantize_row_fp6_weighted( + const float * GGML_RESTRICT x, block_rocmfp6 * GGML_RESTRICT y, int64_t k, const float * GGML_RESTRICT quant_weights) { + assert(k % QK_ROCMFP6 == 0); + + float sum_x2 = 0.0f; + for (int64_t i = 0; i < k; ++i) { + sum_x2 += isfinite(x[i]) ? x[i]*x[i] : 0.0f; + } + const float sigma2 = sum_x2 / (float) k; + + const int64_t nb = k / QK_ROCMFP6; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib*QK_ROCMFP6; + const float * qw = quant_weights ? quant_weights + ib*QK_ROCMFP6 : NULL; + block_rocmfp6 * yb = y + ib; + + for (int half = 0; half < 2; ++half) { + const int half_off = half * (QK_ROCMFP6/2); + const float * xh = xb + half_off; + const float * qh = qw ? qw + half_off : NULL; + yb->e[half] = qh ? + rocmfpx_choose_scale_fp6_weighted_mse(xh, QK_ROCMFP6/2, qh, sigma2) : + rocmfpx_choose_scale_fp6_mse(xh, QK_ROCMFP6/2); + + const float scale = rocmfpx_scale_lookup(yb->e[half]); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + + // 16 codes packed into 4 groups of 4 (4 × 3 bytes = 12 bytes per half). + uint8_t codes[QK_ROCMFP6/2]; + for (int j = 0; j < QK_ROCMFP6/2; ++j) { + codes[j] = rocmfpx_quantize_fp6_code(xh[j], inv_scale); + } + uint8_t * qsdst = yb->qs + half_off * 6 / 8; + rocmfpx_fp6_pack4(qsdst, codes); + rocmfpx_fp6_pack4(qsdst + 3, codes + 4); + rocmfpx_fp6_pack4(qsdst + 6, codes + 8); + rocmfpx_fp6_pack4(qsdst + 9, codes + 12); + } + } +} + +void rocmfpx_quantize_row_fp6_ref(const float * GGML_RESTRICT x, block_rocmfp6 * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP6 == 0); + + const int64_t nb = k / QK_ROCMFP6; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib*QK_ROCMFP6; + block_rocmfp6 * yb = y + ib; + + for (int half = 0; half < 2; ++half) { + const int half_off = half * (QK_ROCMFP6/2); + const float * xh = xb + half_off; + yb->e[half] = rocmfpx_choose_scale_fp6_mse(xh, QK_ROCMFP6/2); + + const float scale = rocmfpx_scale_lookup(yb->e[half]); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + + uint8_t codes[QK_ROCMFP6/2]; + for (int j = 0; j < QK_ROCMFP6/2; ++j) { + codes[j] = rocmfpx_quantize_fp6_code(xh[j], inv_scale); + } + uint8_t * qsdst = yb->qs + half_off * 6 / 8; + rocmfpx_fp6_pack4(qsdst, codes); + rocmfpx_fp6_pack4(qsdst + 3, codes + 4); + rocmfpx_fp6_pack4(qsdst + 6, codes + 8); + rocmfpx_fp6_pack4(qsdst + 9, codes + 12); + } + } +} + +void rocmfpx_dequantize_row_fp6(const block_rocmfp6 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP6 == 0); + + const int64_t nb = k / QK_ROCMFP6; + for (int64_t ib = 0; ib < nb; ++ib) { + const block_rocmfp6 * xb = x + ib; + float * yb = y + ib*QK_ROCMFP6; + + // Unpack all 32 codes in 8 groups of 4 (8 × 3 bytes = 24 bytes). + uint8_t codes[QK_ROCMFP6]; + rocmfpx_fp6_unpack4(xb->qs, codes); + rocmfpx_fp6_unpack4(xb->qs + 3, codes + 4); + rocmfpx_fp6_unpack4(xb->qs + 6, codes + 8); + rocmfpx_fp6_unpack4(xb->qs + 9, codes + 12); + rocmfpx_fp6_unpack4(xb->qs + 12, codes + 16); + rocmfpx_fp6_unpack4(xb->qs + 15, codes + 20); + rocmfpx_fp6_unpack4(xb->qs + 18, codes + 24); + rocmfpx_fp6_unpack4(xb->qs + 21, codes + 28); + + for (int half = 0; half < 2; ++half) { + const float scale = rocmfpx_scale_lookup(xb->e[half]); + for (int j = 0; j < QK_ROCMFP6/2; ++j) { + const int i = half*(QK_ROCMFP6/2) + j; + yb[i] = (float) rocmfpx_decode_fp6_code(codes[i]) * scale; + } + } + } +} + +void rocmfpx_quantize_row_fp6(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + rocmfpx_quantize_row_fp6_ref(x, (block_rocmfp6 *) y, k); +} + +size_t rocmfpx_quantize_fp6(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix) { + const size_t row_size = rocmfpx_row_size_fp6(n_per_row); + char * qrow = (char *) dst; + + for (int64_t row = 0; row < nrows; ++row) { + if (imatrix) { + rocmfpx_quantize_row_fp6_weighted(src + row*n_per_row, (block_rocmfp6 *) qrow, n_per_row, imatrix); + } else { + rocmfpx_quantize_row_fp6_ref(src + row*n_per_row, (block_rocmfp6 *) qrow, n_per_row); + } + qrow += row_size; + } + + return (size_t) nrows * row_size; +} + +static void rocmfpx_quantize_row_fp8_weighted( + const float * GGML_RESTRICT x, block_rocmfp8 * GGML_RESTRICT y, int64_t k, const float * GGML_RESTRICT quant_weights) { + assert(k % QK_ROCMFP8 == 0); + + float sum_x2 = 0.0f; + for (int64_t i = 0; i < k; ++i) { + sum_x2 += isfinite(x[i]) ? x[i]*x[i] : 0.0f; + } + const float sigma2 = sum_x2 / (float) k; + + const int64_t nb = k / QK_ROCMFP8; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib*QK_ROCMFP8; + const float * qw = quant_weights ? quant_weights + ib*QK_ROCMFP8 : NULL; + block_rocmfp8 * yb = y + ib; + + yb->e = qw ? rocmfpx_choose_scale_fp8_weighted_mse(xb, QK_ROCMFP8, qw, sigma2) : + rocmfpx_nearest_scale_ue4m3(rocmfpx_max_abs(xb, QK_ROCMFP8) / 127.0f); + + const float scale = rocmfpx_scale_lookup(yb->e); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + + for (int i = 0; i < QK_ROCMFP8; ++i) { + yb->qs[i] = rocmfpx_quantize_fp8_code(xb[i], inv_scale); + } + } +} + +void rocmfpx_quantize_row_fp8_ref(const float * GGML_RESTRICT x, block_rocmfp8 * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP8 == 0); + + const int64_t nb = k / QK_ROCMFP8; + for (int64_t ib = 0; ib < nb; ++ib) { + const float * xb = x + ib*QK_ROCMFP8; + block_rocmfp8 * yb = y + ib; + + const float max_abs = rocmfpx_max_abs(xb, QK_ROCMFP8); + yb->e = rocmfpx_nearest_scale_ue4m3(max_abs / 127.0f); + + const float scale = rocmfpx_scale_lookup(yb->e); + const float inv_scale = scale > 0.0f ? 1.0f / scale : 0.0f; + + for (int i = 0; i < QK_ROCMFP8; ++i) { + yb->qs[i] = rocmfpx_quantize_fp8_code(xb[i], inv_scale); + } + } +} + +void rocmfpx_dequantize_row_fp8(const block_rocmfp8 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + assert(k % QK_ROCMFP8 == 0); + + const int64_t nb = k / QK_ROCMFP8; + for (int64_t ib = 0; ib < nb; ++ib) { + const block_rocmfp8 * xb = x + ib; + float * yb = y + ib*QK_ROCMFP8; + + const float scale = rocmfpx_scale_lookup(xb->e); + for (int i = 0; i < QK_ROCMFP8; ++i) { + yb[i] = (float) xb->qs[i] * scale; + } + } +} + +void rocmfpx_quantize_row_fp8(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k) { + rocmfpx_quantize_row_fp8_ref(x, (block_rocmfp8 *) y, k); +} + +size_t rocmfpx_quantize_fp8(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix) { + const size_t row_size = rocmfpx_row_size_fp8(n_per_row); + char * qrow = (char *) dst; + + for (int64_t row = 0; row < nrows; ++row) { + if (imatrix) { + rocmfpx_quantize_row_fp8_weighted(src + row*n_per_row, (block_rocmfp8 *) qrow, n_per_row, imatrix); + } else { + rocmfpx_quantize_row_fp8_ref(src + row*n_per_row, (block_rocmfp8 *) qrow, n_per_row); + } + qrow += row_size; + } + + return (size_t) nrows * row_size; +} + +bool rocmfpx_validate_row_data_fp2(const void * data, size_t nbytes) { + if (nbytes % sizeof(block_rocmfp2) != 0) { + return false; + } + + const block_rocmfp2 * blocks = (const block_rocmfp2 *) data; + const size_t nb = nbytes / sizeof(block_rocmfp2); + for (size_t i = 0; i < nb; ++i) { + if (!rocmfpx_scale_is_valid(blocks[i].e[0]) || !rocmfpx_scale_is_valid(blocks[i].e[1])) { + return false; + } + } + return true; +} + +bool rocmfpx_validate_row_data_fp3(const void * data, size_t nbytes) { + if (nbytes % sizeof(block_rocmfp3) != 0) { + return false; + } + + const block_rocmfp3 * blocks = (const block_rocmfp3 *) data; + const size_t nb = nbytes / sizeof(block_rocmfp3); + + for (size_t i = 0; i < nb; ++i) { + if (!rocmfpx_scale_is_valid(blocks[i].e[0]) || !rocmfpx_scale_is_valid(blocks[i].e[1])) { + return false; + } + } + + return true; +} + +bool rocmfpx_validate_row_data_fp6(const void * data, size_t nbytes) { + if (nbytes % sizeof(block_rocmfp6) != 0) { + return false; + } + + const block_rocmfp6 * blocks = (const block_rocmfp6 *) data; + const size_t nb = nbytes / sizeof(block_rocmfp6); + + for (size_t i = 0; i < nb; ++i) { + if (!rocmfpx_scale_is_valid(blocks[i].e[0]) || !rocmfpx_scale_is_valid(blocks[i].e[1])) { + return false; + } + } + + return true; +} + +bool rocmfpx_validate_row_data_fp8(const void * data, size_t nbytes) { + if (nbytes % sizeof(block_rocmfp8) != 0) { + return false; + } + + const block_rocmfp8 * blocks = (const block_rocmfp8 *) data; + const size_t nb = nbytes / sizeof(block_rocmfp8); + + for (size_t i = 0; i < nb; ++i) { + if (!rocmfpx_scale_is_valid(blocks[i].e)) { + return false; + } + } + + return true; +} diff --git a/ggml/rocmfpx/rocmfpx.h b/ggml/rocmfpx/rocmfpx.h new file mode 100644 index 000000000000..30f831ee6a43 --- /dev/null +++ b/ggml/rocmfpx/rocmfpx.h @@ -0,0 +1,112 @@ +// ROCmFPx quant formats, hand-ported into this fork. +// +// Origin: https://github.com/charlie12345/ROCmFPX - creator of the ROCmFP4 format. +// Ported from: https://github.com/ciru-ai/ROCmFPX - a fork of the above. +// +// Both upstream projects are MIT licensed and based on llama.cpp; upstream authors +// retain their authorship and MIT license credit. See LICENSE. + +#pragma once + +#include +#include +#include + +#include "ggml.h" + +#ifdef __cplusplus +extern "C" { +#endif + +#define QK_ROCMFPX 32 + +#define QK_ROCMFP2 QK_ROCMFPX +#define QK_ROCMFP3 QK_ROCMFPX +#define QK_ROCMFP6 QK_ROCMFPX +#define QK_ROCMFP8 QK_ROCMFPX + +#define QS_ROCMFP2 ((QK_ROCMFP2 * 2) / 8) +#define QR_ROCMFP2 1 +#define QI_ROCMFP2 (QK_ROCMFP2 / (4 * QR_ROCMFP2)) +#define QS_ROCMFP3 ((QK_ROCMFP3 * 3) / 8) +#define QS_ROCMFP6 ((QK_ROCMFP6 * 6) / 8) +#define QS_ROCMFP8 QK_ROCMFP8 + +#define QR_ROCMFP3 1 +#define QI_ROCMFP3 (QK_ROCMFP3 / (4 * QR_ROCMFP3)) + +#define QR_ROCMFP6 1 +#define QI_ROCMFP6 (QK_ROCMFP6 / (4 * QR_ROCMFP6)) + +#define QR_ROCMFP8 1 +#define QI_ROCMFP8 (QK_ROCMFP8 / (4 * QR_ROCMFP8)) + +// AMD-native experimental family layouts. The GGUF types are registered, but +// the layouts stay isolated from the promoted ROCmFP4 formats while evaluated. +typedef struct { + uint8_t qs[QS_ROCMFP2]; + uint8_t e[2]; +} block_rocmfp2; + +typedef struct { + uint8_t qs[QS_ROCMFP3]; + uint8_t e[2]; +} block_rocmfp3; + +typedef struct { + uint8_t qs[QS_ROCMFP6]; + uint8_t e[2]; +} block_rocmfp6; + +typedef struct { + int8_t qs[QS_ROCMFP8]; + uint8_t e; +} block_rocmfp8; + +#if defined(__cplusplus) +static_assert(sizeof(block_rocmfp2) == QS_ROCMFP2 + 2*sizeof(uint8_t), "wrong rocmfp2 block size/padding"); +static_assert(sizeof(block_rocmfp3) == QS_ROCMFP3 + 2*sizeof(uint8_t), "wrong rocmfp3 block size/padding"); +static_assert(sizeof(block_rocmfp6) == QS_ROCMFP6 + 2*sizeof(uint8_t), "wrong rocmfp6 block size/padding"); +static_assert(sizeof(block_rocmfp8) == QS_ROCMFP8 + sizeof(uint8_t), "wrong rocmfp8 block size/padding"); +#else +_Static_assert(sizeof(block_rocmfp2) == QS_ROCMFP2 + 2*sizeof(uint8_t), "wrong rocmfp2 block size/padding"); +_Static_assert(sizeof(block_rocmfp3) == QS_ROCMFP3 + 2*sizeof(uint8_t), "wrong rocmfp3 block size/padding"); +_Static_assert(sizeof(block_rocmfp6) == QS_ROCMFP6 + 2*sizeof(uint8_t), "wrong rocmfp6 block size/padding"); +_Static_assert(sizeof(block_rocmfp8) == QS_ROCMFP8 + sizeof(uint8_t), "wrong rocmfp8 block size/padding"); +#endif + +GGML_API float rocmfpx_ue4m3_to_fp32(uint8_t e); +GGML_API bool rocmfpx_scale_is_valid(uint8_t e); +GGML_API size_t rocmfpx_row_size_fp2(int64_t k); +GGML_API size_t rocmfpx_row_size_fp3(int64_t k); +GGML_API size_t rocmfpx_row_size_fp6(int64_t k); +GGML_API size_t rocmfpx_row_size_fp8(int64_t k); + +GGML_API void rocmfpx_quantize_row_fp2_ref(const float * GGML_RESTRICT x, block_rocmfp2 * GGML_RESTRICT y, int64_t k); +GGML_API void rocmfpx_dequantize_row_fp2(const block_rocmfp2 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void rocmfpx_quantize_row_fp2(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +GGML_API size_t rocmfpx_quantize_fp2(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); + +GGML_API void rocmfpx_quantize_row_fp3_ref(const float * GGML_RESTRICT x, block_rocmfp3 * GGML_RESTRICT y, int64_t k); +GGML_API void rocmfpx_dequantize_row_fp3(const block_rocmfp3 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void rocmfpx_quantize_row_fp3(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +GGML_API size_t rocmfpx_quantize_fp3(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); + +GGML_API void rocmfpx_quantize_row_fp6_ref(const float * GGML_RESTRICT x, block_rocmfp6 * GGML_RESTRICT y, int64_t k); +GGML_API void rocmfpx_dequantize_row_fp6(const block_rocmfp6 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void rocmfpx_quantize_row_fp6(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +GGML_API size_t rocmfpx_quantize_fp6(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); + +GGML_API void rocmfpx_quantize_row_fp8_ref(const float * GGML_RESTRICT x, block_rocmfp8 * GGML_RESTRICT y, int64_t k); +GGML_API void rocmfpx_dequantize_row_fp8(const block_rocmfp8 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void rocmfpx_quantize_row_fp8(const float * GGML_RESTRICT x, void * GGML_RESTRICT y, int64_t k); +GGML_API size_t rocmfpx_quantize_fp8(const float * GGML_RESTRICT src, void * GGML_RESTRICT dst, int64_t nrows, int64_t n_per_row, const float * imatrix); + +GGML_API bool rocmfpx_validate_row_data_fp2(const void * data, size_t nbytes); +GGML_API bool rocmfpx_validate_row_data_fp3(const void * data, size_t nbytes); +GGML_API bool rocmfpx_validate_row_data_fp6(const void * data, size_t nbytes); +GGML_API bool rocmfpx_validate_row_data_fp8(const void * data, size_t nbytes); + +#ifdef __cplusplus +} +#endif diff --git a/ggml/rocmfpx/test_rocmfp2_reference.c b/ggml/rocmfpx/test_rocmfp2_reference.c new file mode 100644 index 000000000000..09bd04bdb0c3 --- /dev/null +++ b/ggml/rocmfpx/test_rocmfp2_reference.c @@ -0,0 +1,573 @@ +// ROCmFPx quant formats, hand-ported into this fork. +// +// Origin: https://github.com/charlie12345/ROCmFPX - creator of the ROCmFP4 format. +// Ported from: https://github.com/ciru-ai/ROCmFPX - a fork of the above. +// +// Both upstream projects are MIT licensed and based on llama.cpp; upstream authors +// retain their authorship and MIT license credit. See LICENSE. + +#include "rocmfp2_reference.h" + +#include +#include +#include +#include + +static unsigned long long test_checks = 0; + +#define CHECK(condition) do { \ + ++test_checks; \ + if (!(condition)) { \ + fprintf(stderr, "FAIL %s:%d: %s\n", __FILE__, __LINE__, #condition); \ + return false; \ + } \ +} while (0) + +#define CHECK_STATUS(expression, expected) do { \ + const rocmfp2_p1_status actual_status_ = (expression); \ + ++test_checks; \ + if (actual_status_ != (expected)) { \ + fprintf(stderr, "FAIL %s:%d: %s returned %s, expected %s\n", \ + __FILE__, __LINE__, #expression, rocmfp2_p1_status_name(actual_status_), \ + rocmfp2_p1_status_name(expected)); \ + return false; \ + } \ +} while (0) + +static int8_t expected_code_value( + uint8_t code, + rocmfp2_p1_codebook codebook, + rocmfp2_p1_mapping mapping) { + if (mapping == ROCMFP2_P1_MAPPING_MORD) { + const int values[4] = { + -(int) codebook.outer, + -(int) codebook.inner, + (int) codebook.inner, + (int) codebook.outer, + }; + return (int8_t) values[code]; + } + + const int values[4] = { + (int) codebook.inner, + (int) codebook.outer, + -(int) codebook.inner, + -(int) codebook.outer, + }; + return (int8_t) values[code]; +} + +static uint8_t expected_semantic_code(bool negative, bool outer, rocmfp2_p1_mapping mapping) { + if (mapping == ROCMFP2_P1_MAPPING_MORD) { + if (negative) { + return outer ? 0u : 1u; + } + return outer ? 3u : 2u; + } + + if (negative) { + return outer ? 3u : 2u; + } + return outer ? 1u : 0u; +} + +static bool test_layout_and_p0_packing(void) { + CHECK(sizeof(rocmfp2_p1_block) == 10); + CHECK(offsetof(rocmfp2_p1_block, d) == 0); + CHECK(offsetof(rocmfp2_p1_block, s) == 8); + + uint8_t codes[ROCMFP2_P1_BLOCK_WEIGHTS]; + uint8_t unpacked[ROCMFP2_P1_BLOCK_WEIGHTS]; + const uint8_t scales[2] = { 0x12, 0x7e }; + for (int i = 0; i < ROCMFP2_P1_BLOCK_WEIGHTS; ++i) { + codes[i] = (uint8_t) (i & 3); + } + + rocmfp2_p1_block block; + memset(&block, 0xa5, sizeof(block)); + CHECK_STATUS(rocmfp2_p1_pack_codes(codes, scales, &block), ROCMFP2_P1_OK); + + const uint8_t expected_bytes[ROCMFP2_P1_BLOCK_BYTES] = { + 0xe4, 0xe4, 0xe4, 0xe4, 0xe4, 0xe4, 0xe4, 0xe4, 0x12, 0x7e, + }; + CHECK(memcmp(&block, expected_bytes, sizeof(expected_bytes)) == 0); + + rocmfp2_p1_unpack_codes(&block, unpacked); + CHECK(memcmp(codes, unpacked, sizeof(codes)) == 0); + + rocmfp2_p1_block unchanged; + rocmfp2_p1_block before; + memset(&unchanged, 0x5a, sizeof(unchanged)); + memcpy(&before, &unchanged, sizeof(before)); + codes[31] = 4; + CHECK_STATUS(rocmfp2_p1_pack_codes(codes, scales, &unchanged), ROCMFP2_P1_INVALID_CODE); + CHECK(memcmp(&unchanged, &before, sizeof(before)) == 0); + codes[31] = 3; + + const uint8_t invalid_scales[2] = { 0x7f, 0x00 }; + CHECK_STATUS(rocmfp2_p1_pack_codes(codes, invalid_scales, &unchanged), + ROCMFP2_P1_INVALID_SCALE_METADATA); + CHECK(memcmp(&unchanged, &before, sizeof(before)) == 0); + + puts("PASS layout_p0: block=10 data=8 scales=2 bpw=2.50 golden=e4e4e4e4e4e4e4e4127e"); + return true; +} + +static bool test_all_packed_bytes(void) { + const rocmfp2_p1_codebook codebook = { 3, 10 }; + + for (int mapping_index = 0; mapping_index < 2; ++mapping_index) { + const rocmfp2_p1_mapping mapping = (rocmfp2_p1_mapping) mapping_index; + for (int packed_value = 0; packed_value <= 0xff; ++packed_value) { + int8_t decoded[4]; + CHECK_STATUS(rocmfp2_p1_decode_packed_byte( + (uint8_t) packed_value, &codebook, mapping, decoded), ROCMFP2_P1_OK); + + for (int lane = 0; lane < 4; ++lane) { + const uint8_t code = (uint8_t) ((packed_value >> (2 * lane)) & 3); + CHECK(decoded[lane] == expected_code_value(code, codebook, mapping)); + + int8_t scalar = 0; + CHECK_STATUS(rocmfp2_p1_decode_code(code, &codebook, mapping, &scalar), ROCMFP2_P1_OK); + CHECK(scalar == decoded[lane]); + } + } + } + + int8_t value = 0; + CHECK_STATUS(rocmfp2_p1_decode_code(4, &codebook, ROCMFP2_P1_MAPPING_MORD, &value), + ROCMFP2_P1_INVALID_CODE); + + puts("PASS packed_byte_decode: mappings=2 bytes=256 lanes=2048"); + return true; +} + +static double expected_scale(uint8_t scale_byte) { + const int exponent = scale_byte >> 3; + const int mantissa = scale_byte & 7; + return exponent == 0 ? scalbn((double) mantissa, -10) : scalbn((double) (8 + mantissa), exponent - 11); +} + +static bool test_scale_encoding_and_boundaries(void) { + double previous = -1.0; + for (int scale_index = 0; scale_index <= 0xff; ++scale_index) { + const uint8_t scale_byte = (uint8_t) scale_index; + const bool valid = scale_index <= 0x7e; + CHECK(rocmfp2_p1_scale_is_valid(scale_byte) == valid); + + const double decoded = rocmfp2_p1_ue4m3_to_binary64(scale_byte); + if (valid) { + CHECK(decoded == expected_scale(scale_byte)); + CHECK(decoded > previous); + CHECK(rocmfp2_p1_nearest_ue4m3(decoded) == scale_byte); + previous = decoded; + } else { + CHECK(isnan(decoded)); + } + } + + CHECK(rocmfp2_p1_ue4m3_to_binary64(0x00) == 0.0); + CHECK(rocmfp2_p1_ue4m3_to_binary64(0x01) == 0x1p-10); + CHECK(rocmfp2_p1_ue4m3_to_binary64(0x07) == 7.0 * 0x1p-10); + CHECK(rocmfp2_p1_ue4m3_to_binary64(0x08) == 8.0 * 0x1p-10); + CHECK(rocmfp2_p1_ue4m3_to_binary64(0x7e) == 224.0); + + for (int upper_index = 1; upper_index <= 0x7e; ++upper_index) { + const uint8_t upper_byte = (uint8_t) upper_index; + const uint8_t lower_byte = (uint8_t) (upper_index - 1); + const double lower = expected_scale(lower_byte); + const double upper = expected_scale(upper_byte); + const double midpoint = (lower + upper) * 0.5; + + CHECK(rocmfp2_p1_nearest_ue4m3(midpoint) == lower_byte); + CHECK(rocmfp2_p1_nearest_ue4m3(nextafter(midpoint, -INFINITY)) == lower_byte); + CHECK(rocmfp2_p1_nearest_ue4m3(nextafter(midpoint, INFINITY)) == upper_byte); + } + + CHECK(rocmfp2_p1_nearest_ue4m3(-1.0) == 0xff); + CHECK(rocmfp2_p1_nearest_ue4m3(NAN) == 0xff); + CHECK(rocmfp2_p1_nearest_ue4m3(nextafter(224.0, INFINITY)) == 0x7e); + CHECK(rocmfp2_p1_nearest_ue4m3(INFINITY) == 0x7e); + + puts("PASS ue4m3_boundaries: legal=127 invalid=129 adjacent_midpoints=126 max=224"); + return true; +} + +static bool test_metadata_validation(void) { + rocmfp2_p1_block block; + memset(&block, 0, sizeof(block)); + + for (int first = 0; first <= 0x7e; ++first) { + for (int second = 0; second <= 0x7e; ++second) { + block.s[0] = (uint8_t) first; + block.s[1] = (uint8_t) second; + CHECK(rocmfp2_p1_validate_block(&block)); + } + } + + for (int invalid = 0x7f; invalid <= 0xff; ++invalid) { + block.s[0] = (uint8_t) invalid; + block.s[1] = 0; + CHECK(!rocmfp2_p1_validate_block(&block)); + block.s[0] = 0; + block.s[1] = (uint8_t) invalid; + CHECK(!rocmfp2_p1_validate_block(&block)); + } + + uint8_t serialized[2 * ROCMFP2_P1_BLOCK_BYTES]; + memset(serialized, 0, sizeof(serialized)); + serialized[8] = 0x7e; + serialized[9] = 0x00; + serialized[18] = 0x01; + serialized[19] = 0x7e; + CHECK(rocmfp2_p1_validate_serialized(serialized, sizeof(serialized))); + CHECK(rocmfp2_p1_validate_serialized(NULL, 0)); + CHECK(!rocmfp2_p1_validate_serialized(NULL, sizeof(rocmfp2_p1_block))); + + for (size_t length = 1; length < sizeof(serialized); ++length) { + if (length != ROCMFP2_P1_BLOCK_BYTES) { + CHECK(!rocmfp2_p1_validate_serialized(serialized, length)); + } + } + + for (int invalid = 0x7f; invalid <= 0xff; ++invalid) { + serialized[8] = (uint8_t) invalid; + CHECK(!rocmfp2_p1_validate_serialized(serialized, sizeof(serialized))); + serialized[8] = 0x00; + serialized[19] = (uint8_t) invalid; + CHECK(!rocmfp2_p1_validate_serialized(serialized, sizeof(serialized))); + serialized[19] = 0x7e; + } + + const rocmfp2_p1_codebook codebook = { 3, 10 }; + double output[ROCMFP2_P1_BLOCK_WEIGHTS]; + double before[ROCMFP2_P1_BLOCK_WEIGHTS]; + for (int i = 0; i < ROCMFP2_P1_BLOCK_WEIGHTS; ++i) { + output[i] = 1234.0 + (double) i; + } + memcpy(before, output, sizeof(before)); + block.s[0] = 0x7f; + block.s[1] = 0; + CHECK_STATUS(rocmfp2_p1_dequantize_block( + &block, &codebook, ROCMFP2_P1_MAPPING_MORD, output), ROCMFP2_P1_INVALID_SCALE_METADATA); + CHECK(memcmp(output, before, sizeof(output)) == 0); + + puts("PASS metadata_validation: legal_pairs=16129 invalid_bytes=129 malformed_lengths=18"); + return true; +} + +static bool test_zero_midpoint_and_scale_ties(void) { + const rocmfp2_p1_codebook codebook = { 3, 10 }; + const uint8_t scale_byte = 0x31; + const double scale = expected_scale(scale_byte); + const double midpoint = ((double) codebook.inner + (double) codebook.outer) * scale * 0.5; + + for (int mapping_index = 0; mapping_index < 2; ++mapping_index) { + const rocmfp2_p1_mapping mapping = (rocmfp2_p1_mapping) mapping_index; + const double probes[] = { + 0.0, + -0.0, + midpoint, + -midpoint, + nextafter(midpoint, INFINITY), + nextafter(-midpoint, -INFINITY), + }; + const uint8_t expected[] = { + expected_semantic_code(false, false, mapping), + expected_semantic_code(true, false, mapping), + expected_semantic_code(false, false, mapping), + expected_semantic_code(true, false, mapping), + expected_semantic_code(false, true, mapping), + expected_semantic_code(true, true, mapping), + }; + + for (size_t i = 0; i < sizeof(probes) / sizeof(probes[0]); ++i) { + uint8_t reference_code = 0xff; + uint8_t optimized_code = 0xff; + CHECK_STATUS(rocmfp2_p1_select_code_ref( + probes[i], scale_byte, &codebook, mapping, &reference_code), ROCMFP2_P1_OK); + CHECK_STATUS(rocmfp2_p1_select_code_optimized( + probes[i], scale_byte, &codebook, mapping, &optimized_code), ROCMFP2_P1_OK); + CHECK(reference_code == expected[i]); + CHECK(optimized_code == expected[i]); + } + + uint8_t zero_scale_positive = 0xff; + uint8_t zero_scale_negative = 0xff; + CHECK_STATUS(rocmfp2_p1_select_code_ref( + 1.0, 0, &codebook, mapping, &zero_scale_positive), ROCMFP2_P1_OK); + CHECK_STATUS(rocmfp2_p1_select_code_optimized( + -1.0, 0, &codebook, mapping, &zero_scale_negative), ROCMFP2_P1_OK); + CHECK(zero_scale_positive == 0); + CHECK(zero_scale_negative == 0); + } + + const rocmfp2_p1_codebook tie_codebook = { 1, 2 }; + double source[ROCMFP2_P1_BLOCK_WEIGHTS]; + const double group0_value = 2.0 * expected_scale(0x20); + const double group1_value = 2.0 * expected_scale(0x21); + for (int i = 0; i < ROCMFP2_P1_BLOCK_WEIGHTS; ++i) { + const double magnitude = i < ROCMFP2_P1_GROUP_WEIGHTS ? group0_value : group1_value; + source[i] = (i & 1) ? -magnitude : magnitude; + } + + for (int mapping_index = 0; mapping_index < 2; ++mapping_index) { + const rocmfp2_p1_mapping mapping = (rocmfp2_p1_mapping) mapping_index; + rocmfp2_p1_block reference; + rocmfp2_p1_block optimized; + double reference_sse[2]; + double optimized_sse[2]; + CHECK_STATUS(rocmfp2_p1_quantize_block_ref( + source, &tie_codebook, mapping, &reference, reference_sse), ROCMFP2_P1_OK); + CHECK_STATUS(rocmfp2_p1_quantize_block_optimized( + source, &tie_codebook, mapping, &optimized, optimized_sse), ROCMFP2_P1_OK); + CHECK(memcmp(&reference, &optimized, sizeof(reference)) == 0); + CHECK(reference.s[0] == 0x20); + CHECK(reference.s[1] == 0x21); + CHECK(reference_sse[0] == 0.0 && reference_sse[1] == 0.0); + CHECK(optimized_sse[0] == 0.0 && optimized_sse[1] == 0.0); + } + + for (int i = 0; i < ROCMFP2_P1_BLOCK_WEIGHTS; ++i) { + source[i] = (i & 1) ? -0.0 : 0.0; + } + for (int mapping_index = 0; mapping_index < 2; ++mapping_index) { + const rocmfp2_p1_mapping mapping = (rocmfp2_p1_mapping) mapping_index; + rocmfp2_p1_block reference; + rocmfp2_p1_block optimized; + uint8_t codes[ROCMFP2_P1_BLOCK_WEIGHTS]; + CHECK_STATUS(rocmfp2_p1_quantize_block_ref( + source, &codebook, mapping, &reference, NULL), ROCMFP2_P1_OK); + CHECK_STATUS(rocmfp2_p1_quantize_block_optimized( + source, &codebook, mapping, &optimized, NULL), ROCMFP2_P1_OK); + CHECK(memcmp(&reference, &optimized, sizeof(reference)) == 0); + CHECK(reference.s[0] == 0 && reference.s[1] == 0); + rocmfp2_p1_unpack_codes(&reference, codes); + for (int i = 0; i < ROCMFP2_P1_BLOCK_WEIGHTS; ++i) { + CHECK(codes[i] == 0); + } + const uint8_t canonical_zero[ROCMFP2_P1_BLOCK_BYTES] = { 0 }; + CHECK(memcmp(&reference, canonical_zero, sizeof(canonical_zero)) == 0); + } + + for (int i = 0; i < ROCMFP2_P1_BLOCK_WEIGHTS; ++i) { + source[i] = (i & 1) ? -10000.0 : 10000.0; + } + rocmfp2_p1_block saturated; + CHECK_STATUS(rocmfp2_p1_quantize_block_ref( + source, &codebook, ROCMFP2_P1_MAPPING_MORD, &saturated, NULL), ROCMFP2_P1_OK); + CHECK(saturated.s[0] == 0x7e && saturated.s[1] == 0x7e); + + puts("PASS deterministic_ties: signed_zero=2 magnitude_midpoint=4 scale_tuple=2 saturation=0x7e"); + return true; +} + +static bool test_python_golden_vectors(void) { + const rocmfp2_p1_codebook codebook = { 3, 10 }; + double source[ROCMFP2_P1_BLOCK_WEIGHTS]; + static const double values[4] = { -10.0, -3.0, 3.0, 10.0 }; + for (int i = 0; i < ROCMFP2_P1_BLOCK_WEIGHTS; ++i) { + source[i] = values[i & 3] * (i < ROCMFP2_P1_GROUP_WEIGHTS ? 1.0 : 2.0); + } + + static const uint8_t expected[2][ROCMFP2_P1_BLOCK_BYTES] = { + { 0xe4, 0xe4, 0xe4, 0xe4, 0xe4, 0xe4, 0xe4, 0xe4, 0x40, 0x48 }, + { 0x4b, 0x4b, 0x4b, 0x4b, 0x4b, 0x4b, 0x4b, 0x4b, 0x40, 0x48 }, + }; + + for (int mapping_index = 0; mapping_index < 2; ++mapping_index) { + const rocmfp2_p1_mapping mapping = (rocmfp2_p1_mapping) mapping_index; + rocmfp2_p1_block reference; + rocmfp2_p1_block optimized; + double output[ROCMFP2_P1_BLOCK_WEIGHTS]; + CHECK_STATUS(rocmfp2_p1_quantize_block_ref( + source, &codebook, mapping, &reference, NULL), ROCMFP2_P1_OK); + CHECK_STATUS(rocmfp2_p1_quantize_block_optimized( + source, &codebook, mapping, &optimized, NULL), ROCMFP2_P1_OK); + CHECK(memcmp(&reference, expected[mapping_index], sizeof(reference)) == 0); + CHECK(memcmp(&optimized, expected[mapping_index], sizeof(optimized)) == 0); + CHECK_STATUS(rocmfp2_p1_dequantize_block( + &reference, &codebook, mapping, output), ROCMFP2_P1_OK); + CHECK(memcmp(output, source, sizeof(source)) == 0); + } + + puts("PASS python_golden_vectors: vectors=2 mord=e4x8+4048 msm=4bx8+4048"); + return true; +} + +static uint64_t splitmix64_next(uint64_t * state) { + uint64_t value = (*state += UINT64_C(0x9e3779b97f4a7c15)); + value = (value ^ (value >> 30)) * UINT64_C(0xbf58476d1ce4e5b9); + value = (value ^ (value >> 27)) * UINT64_C(0x94d049bb133111eb); + return value ^ (value >> 31); +} + +static uint64_t fnv1a64(uint64_t hash, const void * data, size_t size) { + const uint8_t * bytes = (const uint8_t *) data; + for (size_t i = 0; i < size; ++i) { + hash ^= bytes[i]; + hash *= UINT64_C(1099511628211); + } + return hash; +} + +static void fill_fixed_random_block(double source[ROCMFP2_P1_BLOCK_WEIGHTS], uint64_t * state, int block_index) { + for (int i = 0; i < ROCMFP2_P1_BLOCK_WEIGHTS; ++i) { + const uint64_t bits = splitmix64_next(state); + const double unit = (double) (bits >> 11) * 0x1p-53; + const int exponent = (int) ((bits >> 3) % 18) - 10; + source[i] = scalbn(2.0 * unit - 1.0, exponent); + } + + source[(block_index * 7) & 31] = (block_index & 1) ? -0.0 : 0.0; + source[(block_index * 11 + 3) & 31] *= 13.0; +} + +static bool test_fixed_random_reference_optimized_match(void) { + static const rocmfp2_p1_codebook codebooks[8] = { + { 1, 2 }, { 2, 5 }, { 1, 3 }, { 3, 10 }, + { 2, 7 }, { 1, 4 }, { 3, 13 }, { 1, 5 }, + }; + enum { BLOCKS_PER_CODEBOOK = 24 }; + + uint64_t random_state = UINT64_C(0x726f636d66703231); + uint64_t fingerprint = UINT64_C(1469598103934665603); + + for (int codebook_index = 0; codebook_index < 8; ++codebook_index) { + const rocmfp2_p1_codebook * codebook = &codebooks[codebook_index]; + for (int block_index = 0; block_index < BLOCKS_PER_CODEBOOK; ++block_index) { + double source[ROCMFP2_P1_BLOCK_WEIGHTS]; + rocmfp2_p1_block mapping_blocks[2]; + double mapping_outputs[2][ROCMFP2_P1_BLOCK_WEIGHTS]; + double mapping_sse[2][2]; + fill_fixed_random_block(source, &random_state, block_index); + + for (int mapping_index = 0; mapping_index < 2; ++mapping_index) { + const rocmfp2_p1_mapping mapping = (rocmfp2_p1_mapping) mapping_index; + rocmfp2_p1_block reference; + rocmfp2_p1_block optimized; + rocmfp2_p1_block reference_repeat; + rocmfp2_p1_block optimized_repeat; + double reference_sse[2]; + double optimized_sse[2]; + double repeat_sse[2]; + + CHECK_STATUS(rocmfp2_p1_quantize_block_ref( + source, codebook, mapping, &reference, reference_sse), ROCMFP2_P1_OK); + CHECK_STATUS(rocmfp2_p1_quantize_block_optimized( + source, codebook, mapping, &optimized, optimized_sse), ROCMFP2_P1_OK); + CHECK(memcmp(&reference, &optimized, sizeof(reference)) == 0); + CHECK(reference_sse[0] == optimized_sse[0]); + CHECK(reference_sse[1] == optimized_sse[1]); + + CHECK_STATUS(rocmfp2_p1_quantize_block_ref( + source, codebook, mapping, &reference_repeat, repeat_sse), ROCMFP2_P1_OK); + CHECK(memcmp(&reference, &reference_repeat, sizeof(reference)) == 0); + CHECK(reference_sse[0] == repeat_sse[0] && reference_sse[1] == repeat_sse[1]); + + CHECK_STATUS(rocmfp2_p1_quantize_block_optimized( + source, codebook, mapping, &optimized_repeat, repeat_sse), ROCMFP2_P1_OK); + CHECK(memcmp(&optimized, &optimized_repeat, sizeof(optimized)) == 0); + CHECK(optimized_sse[0] == repeat_sse[0] && optimized_sse[1] == repeat_sse[1]); + CHECK(rocmfp2_p1_validate_block(&reference)); + + CHECK_STATUS(rocmfp2_p1_dequantize_block( + &reference, codebook, mapping, mapping_outputs[mapping_index]), ROCMFP2_P1_OK); + for (int group = 0; group < 2; ++group) { + double recomputed_sse = 0.0; + for (int i = 0; i < ROCMFP2_P1_GROUP_WEIGHTS; ++i) { + const int index = group * ROCMFP2_P1_GROUP_WEIGHTS + i; + const double delta = source[index] - mapping_outputs[mapping_index][index]; + recomputed_sse += delta * delta; + } + CHECK(recomputed_sse == reference_sse[group]); + mapping_sse[mapping_index][group] = reference_sse[group]; + } + + memcpy(&mapping_blocks[mapping_index], &reference, sizeof(reference)); + fingerprint = fnv1a64(fingerprint, &reference, sizeof(reference)); + } + + CHECK(mapping_blocks[0].s[0] == mapping_blocks[1].s[0]); + CHECK(mapping_blocks[0].s[1] == mapping_blocks[1].s[1]); + CHECK(mapping_sse[0][0] == mapping_sse[1][0]); + CHECK(mapping_sse[0][1] == mapping_sse[1][1]); + CHECK(memcmp(mapping_outputs[0], mapping_outputs[1], sizeof(mapping_outputs[0])) == 0); + } + } + + printf("PASS fixed_random: codebooks=8 blocks_each=%d mappings=2 ref_opt_equal=384 fingerprint=%016" PRIx64 "\n", + BLOCKS_PER_CODEBOOK, fingerprint); + return true; +} + +static bool test_nonfinite_rejection_and_output_atomicity(void) { + const rocmfp2_p1_codebook codebook = { 3, 10 }; + double source[ROCMFP2_P1_BLOCK_WEIGHTS]; + for (int i = 0; i < ROCMFP2_P1_BLOCK_WEIGHTS; ++i) { + source[i] = (double) i * 0.125 - 2.0; + } + + const double nonfinite_values[3] = { NAN, INFINITY, -INFINITY }; + for (int mapping_index = 0; mapping_index < 2; ++mapping_index) { + const rocmfp2_p1_mapping mapping = (rocmfp2_p1_mapping) mapping_index; + for (int position = 0; position < ROCMFP2_P1_BLOCK_WEIGHTS; ++position) { + const double saved = source[position]; + for (int variant = 0; variant < 3; ++variant) { + rocmfp2_p1_block reference; + rocmfp2_p1_block optimized; + rocmfp2_p1_block before_reference; + rocmfp2_p1_block before_optimized; + double reference_sse[2] = { 123.0, 456.0 }; + double optimized_sse[2] = { 789.0, 987.0 }; + const double reference_sse_before[2] = { 123.0, 456.0 }; + const double optimized_sse_before[2] = { 789.0, 987.0 }; + + source[position] = nonfinite_values[variant]; + memset(&reference, 0x3c, sizeof(reference)); + memset(&optimized, 0xc3, sizeof(optimized)); + memcpy(&before_reference, &reference, sizeof(reference)); + memcpy(&before_optimized, &optimized, sizeof(optimized)); + + CHECK_STATUS(rocmfp2_p1_quantize_block_ref( + source, &codebook, mapping, &reference, reference_sse), + ROCMFP2_P1_NONFINITE_SOURCE); + CHECK_STATUS(rocmfp2_p1_quantize_block_optimized( + source, &codebook, mapping, &optimized, optimized_sse), + ROCMFP2_P1_NONFINITE_SOURCE); + CHECK(memcmp(&reference, &before_reference, sizeof(reference)) == 0); + CHECK(memcmp(&optimized, &before_optimized, sizeof(optimized)) == 0); + CHECK(memcmp(reference_sse, reference_sse_before, sizeof(reference_sse)) == 0); + CHECK(memcmp(optimized_sse, optimized_sse_before, sizeof(optimized_sse)) == 0); + } + source[position] = saved; + } + } + + uint8_t code = 0; + CHECK_STATUS(rocmfp2_p1_select_code_ref( + NAN, 0x10, &codebook, ROCMFP2_P1_MAPPING_MORD, &code), ROCMFP2_P1_NONFINITE_SOURCE); + CHECK_STATUS(rocmfp2_p1_select_code_optimized( + INFINITY, 0x10, &codebook, ROCMFP2_P1_MAPPING_MSM, &code), ROCMFP2_P1_NONFINITE_SOURCE); + + puts("PASS nonfinite_rejection: positions=32 values=3 mappings=2 paths=2 status=NONFINITE_SOURCE"); + return true; +} + +int main(void) { + if (!test_layout_and_p0_packing() || + !test_all_packed_bytes() || + !test_scale_encoding_and_boundaries() || + !test_metadata_validation() || + !test_zero_midpoint_and_scale_ties() || + !test_python_golden_vectors() || + !test_fixed_random_reference_optimized_match() || + !test_nonfinite_rejection_and_output_atomicity()) { + fprintf(stderr, "ROCmFP2 Phase-1 reference tests FAILED after %llu checks\n", test_checks); + return 1; + } + + printf("PASS rocmfp2_phase1_reference: checks=%llu arithmetic=IEEE754-binary64 search=127-scales exhaustive\n", + test_checks); + return 0; +} diff --git a/ggml/rocmfpx/test_rocmfpx.c b/ggml/rocmfpx/test_rocmfpx.c new file mode 100644 index 000000000000..75dc6a0d614d --- /dev/null +++ b/ggml/rocmfpx/test_rocmfpx.c @@ -0,0 +1,190 @@ +// ROCmFPx quant formats, hand-ported into this fork. +// +// Origin: https://github.com/charlie12345/ROCmFPX - creator of the ROCmFP4 format. +// Ported from: https://github.com/ciru-ai/ROCmFPX - a fork of the above. +// +// Both upstream projects are MIT licensed and based on llama.cpp; upstream authors +// retain their authorship and MIT license credit. See LICENSE. + +#include "rocmfpx.h" + +// Every check below is an assert(), so keep them alive in Release builds - +// otherwise the test would validate nothing and still exit 0. +#ifdef NDEBUG +#undef NDEBUG +#endif +#include +#include +#include + +static void fill_row(float * x, int n) { + for (int i = 0; i < n; ++i) { + const float wave = 0.75f*sinf((float) i * 0.37f) + 0.25f*cosf((float) i * 0.13f); + const float ramp = ((float) (i % 11) - 5.0f) * 0.035f; + x[i] = wave + ramp; + } + + x[7] = 3.25f; + x[19] = -2.75f; + x[43] = 1.875f; +} + +static float mse(const float * a, const float * b, int n) { + float err = 0.0f; + + for (int i = 0; i < n; ++i) { + const float d = a[i] - b[i]; + err += d*d; + } + + return err / (float) n; +} + +static float weighted_mse(const float * a, const float * b, const float * w, int n) { + float err = 0.0f; + float sum_w = 0.0f; + + for (int i = 0; i < n; ++i) { + const float d = a[i] - b[i]; + err += w[i]*d*d; + sum_w += w[i]; + } + + return sum_w > 0.0f ? err / sum_w : 0.0f; +} + +static void fill_imatrix_case(float * src, float * imatrix, int n, float base, float outlier) { + for (int i = 0; i < n; ++i) { + const float sign = (i % 2) ? 1.0f : -1.0f; + src[i] = sign * base * (1.0f + 0.03f*(float)(i % 5)); + imatrix[i] = 100.0f; + } + + src[0] = outlier; + imatrix[0] = 0.0f; +} + +static void check_weighted_imatrix_fp3(void) { + enum { N = QK_ROCMFP3 }; + + float src[N]; + float imatrix[N]; + float plain[N]; + float weighted[N]; + block_rocmfp3 q_plain[N / QK_ROCMFP3]; + block_rocmfp3 q_weighted[N / QK_ROCMFP3]; + + fill_imatrix_case(src, imatrix, N, 0.21f, 9.0f); + + rocmfpx_quantize_fp3(src, q_plain, 1, N, NULL); + rocmfpx_quantize_fp3(src, q_weighted, 1, N, imatrix); + rocmfpx_dequantize_row_fp3(q_plain, plain, N); + rocmfpx_dequantize_row_fp3(q_weighted, weighted, N); + + const float plain_err = weighted_mse(src, plain, imatrix, N); + const float weighted_err = weighted_mse(src, weighted, imatrix, N); + + printf("ROCmFP3 imatrix weighted_mse: plain=%g weighted=%g\n", (double) plain_err, (double) weighted_err); + assert(weighted_err < plain_err); +} + +static void check_weighted_imatrix_fp6(void) { + enum { N = QK_ROCMFP6 }; + + float src[N]; + float imatrix[N]; + float plain[N]; + float weighted[N]; + block_rocmfp6 q_plain[N / QK_ROCMFP6]; + block_rocmfp6 q_weighted[N / QK_ROCMFP6]; + + fill_imatrix_case(src, imatrix, N, 0.045f, 6.0f); + + rocmfpx_quantize_fp6(src, q_plain, 1, N, NULL); + rocmfpx_quantize_fp6(src, q_weighted, 1, N, imatrix); + rocmfpx_dequantize_row_fp6(q_plain, plain, N); + rocmfpx_dequantize_row_fp6(q_weighted, weighted, N); + + const float plain_err = weighted_mse(src, plain, imatrix, N); + const float weighted_err = weighted_mse(src, weighted, imatrix, N); + + printf("ROCmFP6 imatrix weighted_mse: plain=%g weighted=%g\n", (double) plain_err, (double) weighted_err); + assert(weighted_err < plain_err); +} + +static void check_weighted_imatrix_fp8(void) { + enum { N = QK_ROCMFP8 }; + + float src[N]; + float imatrix[N]; + float plain[N]; + float weighted[N]; + block_rocmfp8 q_plain[N / QK_ROCMFP8]; + block_rocmfp8 q_weighted[N / QK_ROCMFP8]; + + fill_imatrix_case(src, imatrix, N, 0.008f, 8.0f); + + rocmfpx_quantize_fp8(src, q_plain, 1, N, NULL); + rocmfpx_quantize_fp8(src, q_weighted, 1, N, imatrix); + rocmfpx_dequantize_row_fp8(q_plain, plain, N); + rocmfpx_dequantize_row_fp8(q_weighted, weighted, N); + + const float plain_err = weighted_mse(src, plain, imatrix, N); + const float weighted_err = weighted_mse(src, weighted, imatrix, N); + + printf("ROCmFP8 imatrix weighted_mse: plain=%g weighted=%g\n", (double) plain_err, (double) weighted_err); + assert(weighted_err < plain_err); +} + +int main(void) { + enum { N = 64 }; + + float src[N]; + float fp3[N]; + float fp6[N]; + float fp8[N]; + + block_rocmfp3 q3[N / QK_ROCMFP3]; + block_rocmfp6 q6[N / QK_ROCMFP6]; + block_rocmfp8 q8[N / QK_ROCMFP8]; + + fill_row(src, N); + + rocmfpx_quantize_row_fp3_ref(src, q3, N); + rocmfpx_quantize_row_fp6_ref(src, q6, N); + rocmfpx_quantize_row_fp8_ref(src, q8, N); + + assert(rocmfpx_validate_row_data_fp3(q3, sizeof(q3))); + assert(rocmfpx_validate_row_data_fp6(q6, sizeof(q6))); + assert(rocmfpx_validate_row_data_fp8(q8, sizeof(q8))); + + rocmfpx_dequantize_row_fp3(q3, fp3, N); + rocmfpx_dequantize_row_fp6(q6, fp6, N); + rocmfpx_dequantize_row_fp8(q8, fp8, N); + + const float mse3 = mse(src, fp3, N); + const float mse6 = mse(src, fp6, N); + const float mse8 = mse(src, fp8, N); + + printf("ROCmFP3: block=%zu row=%zu bpw=%.2f mse=%g\n", + sizeof(block_rocmfp3), rocmfpx_row_size_fp3(N), + 8.0*(double) sizeof(block_rocmfp3)/(double) QK_ROCMFP3, (double) mse3); + printf("ROCmFP6: block=%zu row=%zu bpw=%.2f mse=%g\n", + sizeof(block_rocmfp6), rocmfpx_row_size_fp6(N), + 8.0*(double) sizeof(block_rocmfp6)/(double) QK_ROCMFP6, (double) mse6); + printf("ROCmFP8: block=%zu row=%zu bpw=%.2f mse=%g\n", + sizeof(block_rocmfp8), rocmfpx_row_size_fp8(N), + 8.0*(double) sizeof(block_rocmfp8)/(double) QK_ROCMFP8, (double) mse8); + + assert(isfinite(mse3)); + assert(isfinite(mse6)); + assert(isfinite(mse8)); + assert(mse8 < mse6); + assert(mse6 < mse3); + + check_weighted_imatrix_fp3(); + check_weighted_imatrix_fp6(); + check_weighted_imatrix_fp8(); + + return 0; +} diff --git a/ggml/src/CMakeLists.txt b/ggml/src/CMakeLists.txt index 947732000219..26e445e80656 100644 --- a/ggml/src/CMakeLists.txt +++ b/ggml/src/CMakeLists.txt @@ -201,11 +201,17 @@ add_library(ggml-base ggml-alloc.c ggml-backend.cpp ggml-backend-meta.cpp + ggml-backend-sanitize.cpp + ggml-backend-sanitize.h ggml-opt.cpp ggml-threading.cpp ggml-threading.h ggml-quants.c ggml-quants.h + ../rocmfp4/rocmfp4.c + ../rocmfp4/rocmfp4.h + ../rocmfpx/rocmfpx.c + ../rocmfpx/rocmfpx.h gguf.cpp) set_target_properties(ggml-base PROPERTIES diff --git a/ggml/src/ggml-alloc.c b/ggml/src/ggml-alloc.c index a71838eafc69..5fd670fd0c43 100644 --- a/ggml/src/ggml-alloc.c +++ b/ggml/src/ggml-alloc.c @@ -493,6 +493,9 @@ struct ggml_gallocr { struct leaf_alloc * leaf_allocs; // [n_leafs] int n_leafs; + + struct ggml_hash_set pinned; + bool has_pinned; }; ggml_gallocr_t ggml_gallocr_new_n(ggml_backend_buffer_type_t * bufts, int n_bufs) { @@ -570,6 +573,9 @@ void ggml_gallocr_free(ggml_gallocr_t galloc) { } ggml_hash_set_free(&galloc->hash_set); + if (galloc->pinned.size > 0) { + ggml_hash_set_free(&galloc->pinned); + } free(galloc->hash_values); free(galloc->bufts); free(galloc->buffers); @@ -590,6 +596,16 @@ static bool ggml_gallocr_is_own(ggml_gallocr_t galloc, struct ggml_tensor * t) { return ggml_gallocr_hash_get(galloc, t)->allocated; } +static bool ggml_gallocr_is_pinned(ggml_gallocr_t galloc, struct ggml_tensor * t) { + if (!galloc->has_pinned) { + return false; + } + while (t->view_src != NULL) { + t = t->view_src; + } + return ggml_hash_contains(&galloc->pinned, t); +} + static bool ggml_gallocr_is_allocated(ggml_gallocr_t galloc, struct ggml_tensor * t) { return t->data != NULL // tensor data already set externally || t->buffer // tensor on external buffer (but not yet allocated) @@ -642,6 +658,11 @@ static void ggml_gallocr_allocate_node(ggml_gallocr_t galloc, struct ggml_tensor continue; } + if (ggml_gallocr_is_pinned(galloc, parent)) { + AT_PRINTF("not reusing parent %s for %s as it is pinned\n", parent->name, node->name); + continue; + } + // outputs cannot be reused if (parent->flags & GGML_TENSOR_FLAG_OUTPUT || (parent->view_src != NULL && parent->view_src->flags & GGML_TENSOR_FLAG_OUTPUT)) { AT_PRINTF("not reusing parent %s for %s as it is an output\n", parent->name, node->name); @@ -688,6 +709,25 @@ static void ggml_gallocr_allocate_node(ggml_gallocr_t galloc, struct ggml_tensor } } +void ggml_gallocr_pin_tensor(ggml_gallocr_t galloc, struct ggml_tensor * t) { + GGML_ASSERT(galloc); + if (galloc->pinned.size == 0) { + galloc->pinned = ggml_hash_set_new(GGML_DEFAULT_GRAPH_SIZE); + } + if (ggml_hash_insert(&galloc->pinned, t) == GGML_HASHSET_FULL) { + GGML_ABORT("%s: pinned tensor set is full\n", __func__); + } + galloc->has_pinned = true; +} + +void ggml_gallocr_clear_pins(ggml_gallocr_t galloc) { + GGML_ASSERT(galloc); + if (galloc->pinned.size > 0) { + ggml_hash_set_reset(&galloc->pinned); + } + galloc->has_pinned = false; +} + static void ggml_gallocr_free_node(ggml_gallocr_t galloc, struct ggml_tensor * node) { // graph outputs are never freed if (node->flags & GGML_TENSOR_FLAG_OUTPUT) { @@ -695,6 +735,11 @@ static void ggml_gallocr_free_node(ggml_gallocr_t galloc, struct ggml_tensor * n return; } + if (ggml_gallocr_is_pinned(galloc, node)) { + AT_PRINTF("not freeing pinned %s\n", node->name); + return; + } + struct hash_node * hn = ggml_gallocr_hash_get(galloc, node); int buffer_id = hn->buffer_id; struct ggml_dyn_tallocr * alloc = galloc->buf_tallocs[buffer_id]; diff --git a/ggml/src/ggml-backend-sanitize.cpp b/ggml/src/ggml-backend-sanitize.cpp new file mode 100644 index 000000000000..f9b37742e170 --- /dev/null +++ b/ggml/src/ggml-backend-sanitize.cpp @@ -0,0 +1,527 @@ +#include "ggml-backend-sanitize.h" +#include "ggml-backend-impl.h" +#include "ggml-impl.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +int ggml_san_level(void) { + static const int level = []() { + const char * env = getenv("GGML_SCHED_SANITIZE"); + return env ? atoi(env) : 0; + }(); + return level; +} + +static int ggml_san_nonfatal(void) { + static const int nonfatal = []() { + const char * env = getenv("GGML_SCHED_SANITIZE_NONFATAL"); + return env ? atoi(env) : 0; + }(); + return nonfatal; +} + +namespace { + +constexpr int SAN_MAX_ACTORS = 64; +constexpr size_t SAN_MAX_RANGES = 65536; + +using vclock = std::vector; + +struct san_access { + int actor = -1; + uint64_t clock = 0; + int split = -1; + const char * what = ""; // always a string literal + char tensor[2*GGML_MAX_NAME + 8] = { 0 }; // copied: source tensor is recycled per graph +}; + +struct san_entry { + size_t end; + san_access last_write; + std::vector reads; // at most one per actor +}; + +struct san_state { + std::mutex mutex; + + std::unordered_map actors; + std::vector actor_names; + std::unordered_map name_counts; + std::vector vc; // vc[a][b] = a's knowledge of b's clock + + std::unordered_map> mem; + + std::unordered_map ev_vc; + std::unordered_map events; + + size_t n_races = 0; // only counted in nonfatal mode +}; + +san_state & state() { + static san_state * s = []() { + san_state * st = new san_state(); + st->actor_names.push_back("HOST"); + st->vc.emplace_back(SAN_MAX_ACTORS, 0); + if (ggml_san_nonfatal()) { + static san_state * at_exit = st; + atexit([]() { + fprintf(stderr, "\nggml-sched-sanitize: %zu race(s) reported\n", at_exit->n_races); + }); + } + return st; + }(); + return *s; +} + +// every helper below assumes the caller holds state().mutex + +int actor_id(san_state & s, ggml_backend_t backend) { + if (backend == nullptr) { + return 0; + } + + auto it = s.actors.find(backend); + if (it != s.actors.end()) { + return it->second; + } + + // two backends can share a device and therefore a name, disambiguate + std::string base = ggml_backend_name(backend); + const int idx = s.name_counts[base]++; + + const int id = (int) s.actor_names.size(); + GGML_ASSERT(id < SAN_MAX_ACTORS); + s.actor_names.push_back(idx == 0 ? base : base + "#" + std::to_string(idx)); + s.vc.emplace_back(SAN_MAX_ACTORS, 0); + s.actors[backend] = id; + return id; +} + +const char * actor_name(san_state & s, int id) { + return id >= 0 && id < (int) s.actor_names.size() ? s.actor_names[id].c_str() : "?"; +} + +void join(vclock & dst, const vclock & src) { + for (size_t i = 0; i < dst.size(); i++) { + dst[i] = std::max(dst[i], src[i]); + } +} + +uint64_t tick(san_state & s, int a) { + return ++s.vc[a][a]; +} + +uint64_t issue(san_state & s, int a) { + join(s.vc[a], s.vc[0]); + return tick(s, a); +} + +void maybe_flush(san_state & s) { + if (s.mem.empty()) { + return; + } + + // once the host knows every actor's current clock, the issue rule guarantees that any + // future operation joins vc[0] first, so the shadow state can no longer report a race + for (size_t b = 0; b < s.vc.size(); b++) { + if (s.vc[0][b] < s.vc[b][b]) { + return; + } + } + + s.mem.clear(); +} + +bool is_view_op(enum ggml_op op) { + return op == GGML_OP_VIEW || op == GGML_OP_RESHAPE || op == GGML_OP_PERMUTE || op == GGML_OP_TRANSPOSE; +} + +struct mem_range { + ggml_backend_buffer_t buf; + size_t off; + size_t len; +}; + +bool resolve(const ggml_tensor * t, size_t offset, size_t size, mem_range & out) { + if (t == nullptr || t->data == nullptr) { + return false; + } + + ggml_backend_buffer_t buf = t->view_src ? t->view_src->buffer : t->buffer; + if (buf == nullptr) { + return false; + } + + // weights are written once at load and never recycled by the allocator + if (ggml_backend_buffer_get_usage(buf) == GGML_BACKEND_BUFFER_USAGE_WEIGHTS) { + return false; + } + + void * base = ggml_backend_buffer_get_base(buf); + if (base == nullptr || size == 0) { + return false; + } + + out.buf = buf; + out.off = (size_t) ((const char *) t->data - (const char *) base) + offset; + out.len = size; + return true; +} + +thread_local int tl_split = -1; +thread_local int tl_split_actor = -1; + +void report(san_state & s, const mem_range & mr, const san_access & cur, const san_access & prev, const char * kind) { + fprintf(stderr, "\nggml-sched-sanitize: RACE (%s) on %s[%zu, %zu)\n", + kind, ggml_backend_buffer_name(mr.buf), mr.off, mr.off + mr.len); + fprintf(stderr, "ggml-sched-sanitize: %-5s %-12s @%-4llu split %-4d %s (%s)\n", + strcmp(kind, "write-after-read") == 0 ? "read" : "write", + actor_name(s, prev.actor), (unsigned long long) prev.clock, prev.split, prev.tensor, prev.what); + fprintf(stderr, "ggml-sched-sanitize: %-5s %-12s @%-4llu split %-4d %s (%s)\n", + strcmp(kind, "read-after-write") == 0 ? "read" : "write", + actor_name(s, cur.actor), (unsigned long long) cur.clock, cur.split, cur.tensor, cur.what); + fprintf(stderr, "ggml-sched-sanitize: no happens-before edge: %s knows %s@%llu, needs >=%llu\n\n", + actor_name(s, cur.actor), actor_name(s, prev.actor), + (unsigned long long) s.vc[cur.actor][prev.actor], (unsigned long long) prev.clock); + + if (ggml_san_nonfatal()) { + s.n_races++; + return; + } + + GGML_ABORT("ggml-sched-sanitize: race detected"); +} + +// make sure a range boundary exists at pos, so [off,off+len) lands on whole entries +void split_at(std::map & m, size_t pos) { + auto it = m.upper_bound(pos); + if (it == m.begin()) { + return; + } + --it; + if (it->first < pos && it->second.end > pos) { + san_entry tail = it->second; + it->second.end = pos; + m.emplace(pos, std::move(tail)); + } +} + +void access_range(san_state & s, const mem_range & mr, bool write, const san_access & info) { + auto & m = s.mem[mr.buf]; + const int a = info.actor; + + // the shadow state is bounded by the work issued before the host observes every actor, + // not by run length. reaching this means maybe_flush stopped clearing it - + // most likely an operation was enqueued without going through issue(), which breaks + // the argument maybe_flush relies on. that is a sanitizer bug, not a memory limit. + GGML_ASSERT(m.size() <= SAN_MAX_RANGES && "shadow state is not being flushed - see issue()/maybe_flush()"); + + const size_t begin = mr.off; + const size_t end = mr.off + mr.len; + + split_at(m, begin); + split_at(m, end); + + // fill gaps so the whole span is covered by entries + size_t cur = begin; + while (cur < end) { + auto it = m.lower_bound(cur); + if (it == m.end() || it->first >= end) { + san_entry e; + e.end = end; + m.emplace(cur, std::move(e)); + break; + } + if (it->first > cur) { + san_entry e; + e.end = it->first; + m.emplace(cur, std::move(e)); + } + cur = it->second.end; + } + + for (auto it = m.lower_bound(begin); it != m.end() && it->first < end; ++it) { + san_entry & e = it->second; + + if (e.last_write.actor >= 0 && e.last_write.actor != a) { + if (e.last_write.clock > s.vc[a][e.last_write.actor]) { + report(s, mr, info, e.last_write, write ? "write-after-write" : "read-after-write"); + } + } + + if (write) { + for (const san_access & r : e.reads) { + if (r.actor != a && r.clock > s.vc[a][r.actor]) { + report(s, mr, info, r, "write-after-read"); + } + } + } + + if (write) { + e.last_write = info; + e.reads.clear(); + } else { + bool found = false; + for (san_access & r : e.reads) { + if (r.actor == a) { + r = info; + found = true; + break; + } + } + if (!found) { + e.reads.push_back(info); + } + } + } + +} + +// returns false if the tensor is not in tracked memory +bool touch(san_state & s, int a, uint64_t c, const ggml_tensor * t, + size_t offset, size_t size, bool write, const char * what) { + mem_range mr; + if (!resolve(t, offset, size, mr)) { + return false; + } + + san_access info; + info.actor = a; + info.clock = c; + info.what = what; + info.split = tl_split; + const ggml_tensor * root = t; + while (root->view_src != NULL) { + root = root->view_src; + } + if (root != t) { + snprintf(info.tensor, sizeof(info.tensor), "%s <- %s", t->name, root->name); + } else { + snprintf(info.tensor, sizeof(info.tensor), "%s", t->name); + } + + if (ggml_san_level() >= 2) { + GGML_LOG_DEBUG("san [split %3d %-10s] %-10s %-5s %s[%zu+%zu] %s (%s)\n", + tl_split, tl_split_actor >= 0 ? actor_name(s, tl_split_actor) : "-", actor_name(s, a), + write ? "write" : "read", + ggml_backend_buffer_name(mr.buf), mr.off, mr.len, t->name, what); + } + + access_range(s, mr, write, info); + return true; +} + +void trace(san_state & s, int a, const char * fmt, ...) { + if (ggml_san_level() < 2) { + return; + } + + char body[512]; + va_list ap; + va_start(ap, fmt); + vsnprintf(body, sizeof(body), fmt, ap); + va_end(ap); + + GGML_LOG_DEBUG("san [split %3d %-10s] %-10s %s", + tl_split, tl_split_actor >= 0 ? actor_name(s, tl_split_actor) : "-", actor_name(s, a), body); +} + +} // namespace + +void ggml_san_sync(ggml_backend_t backend) { + if (ggml_san_level() == 0) { + return; + } + san_state & s = state(); + std::lock_guard lk(s.mutex); + + const int a = actor_id(s, backend); + join(s.vc[0], s.vc[a]); // host now knows everything this backend knew + trace(s, a, "EDGE synchronize -> HOST\n"); + maybe_flush(s); +} + +void ggml_san_event_record(ggml_backend_event_t event, ggml_backend_t backend) { + if (ggml_san_level() == 0) { + return; + } + san_state & s = state(); + std::lock_guard lk(s.mutex); + + const int a = actor_id(s, backend); + s.ev_vc[event] = s.vc[a]; + trace(s, a, "EDGE event_record ev=%d\n", (int) s.events.emplace(event, (int) s.events.size()).first->second); +} + +void ggml_san_event_wait(ggml_backend_t backend, ggml_backend_event_t event) { + if (ggml_san_level() == 0) { + return; + } + san_state & s = state(); + std::lock_guard lk(s.mutex); + + const int a = actor_id(s, backend); + auto it = s.ev_vc.find(event); + if (it != s.ev_vc.end()) { + join(s.vc[a], it->second); + } + trace(s, a, "EDGE event_wait ev=%d\n", (int) s.events.emplace(event, (int) s.events.size()).first->second); +} + +void ggml_san_event_sync(ggml_backend_event_t event) { + if (ggml_san_level() == 0) { + return; + } + san_state & s = state(); + std::lock_guard lk(s.mutex); + + auto it = s.ev_vc.find(event); + if (it != s.ev_vc.end()) { + join(s.vc[0], it->second); + } + trace(s, 0, "EDGE event_sync ev=%d\n", (int) s.events.emplace(event, (int) s.events.size()).first->second); + maybe_flush(s); +} + +void ggml_san_compute(ggml_backend_t backend, const ggml_cgraph * cgraph) { + if (ggml_san_level() == 0 || cgraph == nullptr) { + return; + } + san_state & s = state(); + std::lock_guard lk(s.mutex); + + const int a = actor_id(s, backend); + + // a backend with no synchronize runs inline on the calling thread, so it both + // inherits and publishes the host's knowledge - without this the CPU backend's + // writes are never ordered against anything and no race can be detected + const bool synchronous = backend->iface.synchronize == nullptr; + + const uint64_t c = issue(s, a); + + size_t n_read = 0; + size_t n_write = 0; + for (int i = 0; i < cgraph->n_nodes; i++) { + ggml_tensor * node = cgraph->nodes[i]; + if (is_view_op(node->op)) { + continue; + } + for (int j = 0; j < GGML_MAX_SRC; j++) { + if (node->op == GGML_OP_CPY && j == 1) { + continue; + } + if (node->src[j] && touch(s, a, c, node->src[j], 0, ggml_nbytes(node->src[j]), false, "compute")) { + n_read++; + } + } + if (touch(s, a, c, node, 0, ggml_nbytes(node), true, "compute")) { + n_write++; + } + } + + if (synchronous) { + join(s.vc[0], s.vc[a]); + maybe_flush(s); + } + + trace(s, a, "compute %d nodes, r=%zu w=%zu ranges\n", cgraph->n_nodes, n_read, n_write); +} + +void ggml_san_access(ggml_backend_t backend, const ggml_tensor * tensor, + size_t offset, size_t size, bool write, const char * what) { + if (ggml_san_level() == 0) { + return; + } + san_state & s = state(); + std::lock_guard lk(s.mutex); + + const int a = actor_id(s, backend); + const uint64_t c = issue(s, a); + touch(s, a, c, tensor, offset, size, write, what); + + // a synchronous access has completed by the time the call returns + if (backend == nullptr || backend->iface.synchronize == nullptr) { + join(s.vc[0], s.vc[a]); + maybe_flush(s); + } +} + +void ggml_san_cpy_async(ggml_backend_t src_be, ggml_backend_t dst_be, + const ggml_tensor * src, const ggml_tensor * dst) { + if (ggml_san_level() == 0) { + return; + } + san_state & s = state(); + std::lock_guard lk(s.mutex); + + const int a_src = actor_id(s, src_be); + const int a_dst = actor_id(s, dst_be); + + // TODO: clarify in ggml + // which stream actually performs the read of src is backend specific. host memory has + // no queue of its own, so a host->device transfer runs on dst's queue - that is what + // makes the source racy against a later host write. a device->device transfer is + // issued on the src stream (see "copy on src stream" in ggml-cuda), where src's own + // later work is already ordered against it. + ggml_backend_buffer_t src_buf = src->view_src ? src->view_src->buffer : src->buffer; + const int read_actor = (src_buf && ggml_backend_buffer_is_host(src_buf)) ? a_dst : a_src; + + uint64_t read_clock; + uint64_t write_clock; + + if (read_actor == a_dst) { + // the dst queue performs both sides of the copy after src's pending work + join(s.vc[a_dst], s.vc[a_src]); + read_clock = write_clock = issue(s, a_dst); + } else { + // issue the src read first, then carry its completion into the dst write + read_clock = issue(s, read_actor); + join(s.vc[a_dst], s.vc[read_actor]); + write_clock = issue(s, a_dst); + } + + touch(s, read_actor, read_clock, src, 0, ggml_nbytes(src), false, "cpy_async src"); + touch(s, a_dst, write_clock, dst, 0, ggml_nbytes(dst), true, "cpy_async dst"); + + trace(s, a_dst, "cpy_async %s@%s -> %s\n", src->name, actor_name(s, read_actor), dst->name); +} + +void ggml_san_buffer_free(ggml_backend_buffer_t buffer) { + if (ggml_san_level() == 0) { + return; + } + san_state & s = state(); + std::lock_guard lk(s.mutex); + + s.mem.erase(buffer); +} + +void ggml_san_split(int split_id, ggml_backend_t backend, int n_inputs) { + if (ggml_san_level() == 0) { + return; + } + tl_split = split_id; + + if (split_id < 0) { + tl_split_actor = -1; + return; + } + + san_state & s = state(); + std::lock_guard lk(s.mutex); + const int a = actor_id(s, backend); + tl_split_actor = a; + + if (n_inputs == 0) { + trace(s, a, "split 0 inputs (no sync path)\n"); + } +} diff --git a/ggml/src/ggml-backend-sanitize.h b/ggml/src/ggml-backend-sanitize.h new file mode 100644 index 000000000000..9009f50e6492 --- /dev/null +++ b/ggml/src/ggml-backend-sanitize.h @@ -0,0 +1,36 @@ +#pragma once + +// happens-before instrumentation for the ggml backend API +// +// GGML_SCHED_SANITIZE=1 detect happens-before violations +// GGML_SCHED_SANITIZE=2 also trace synchronization edges and memory ranges + +#include "ggml-backend.h" + +#ifdef __cplusplus +extern "C" { +#endif + + int ggml_san_level(void); + + void ggml_san_sync (ggml_backend_t backend); + void ggml_san_event_record(ggml_backend_event_t event, ggml_backend_t backend); + void ggml_san_event_wait (ggml_backend_t backend, ggml_backend_event_t event); + void ggml_san_event_sync (ggml_backend_event_t event); + + void ggml_san_compute (ggml_backend_t backend, const struct ggml_cgraph * cgraph); + + // backend == NULL means the access is performed by the host thread + void ggml_san_access (ggml_backend_t backend, const struct ggml_tensor * tensor, + size_t offset, size_t size, bool write, const char * what); + + void ggml_san_cpy_async (ggml_backend_t src_be, ggml_backend_t dst_be, + const struct ggml_tensor * src, const struct ggml_tensor * dst); + + void ggml_san_buffer_free (ggml_backend_buffer_t buffer); + + void ggml_san_split (int split_id, ggml_backend_t backend, int n_inputs); + +#ifdef __cplusplus +} +#endif diff --git a/ggml/src/ggml-backend.cpp b/ggml/src/ggml-backend.cpp index 6faa680474c4..81f517baf674 100644 --- a/ggml/src/ggml-backend.cpp +++ b/ggml/src/ggml-backend.cpp @@ -10,6 +10,7 @@ #include "ggml-backend.h" #include "ggml-backend-impl.h" +#include "ggml-backend-sanitize.h" #include "ggml-alloc.h" #include "ggml-impl.h" @@ -118,6 +119,8 @@ void ggml_backend_buffer_free(ggml_backend_buffer_t buffer) { return; } + ggml_san_buffer_free(buffer); + if (buffer->iface.free_buffer != NULL) { buffer->iface.free_buffer(buffer); } @@ -216,7 +219,12 @@ void ggml_backend_buffer_reset(ggml_backend_buffer_t buffer) { bool ggml_backend_buffer_copy_tensor(const struct ggml_tensor * src, struct ggml_tensor * dst) { ggml_backend_buffer_t dst_buf = dst->view_src ? dst->view_src->buffer : dst->buffer; if (dst_buf->iface.cpy_tensor) { - return dst_buf->iface.cpy_tensor(dst_buf, src, dst); + const bool copied = dst_buf->iface.cpy_tensor(dst_buf, src, dst); + if (copied) { + ggml_san_access(NULL, src, 0, ggml_nbytes(src), false, "buffer_copy_tensor src"); + ggml_san_access(NULL, dst, 0, ggml_nbytes(dst), true, "buffer_copy_tensor dst"); + } + return copied; } return false; } @@ -272,6 +280,7 @@ void ggml_backend_tensor_set_async(ggml_backend_t backend, struct ggml_tensor * ggml_backend_synchronize(backend); ggml_backend_tensor_set(tensor, data, offset, size); } else { + ggml_san_access(backend, tensor, offset, size, true, "set_async"); backend->iface.set_tensor_async(backend, tensor, data, offset, size); } } @@ -286,6 +295,7 @@ void ggml_backend_tensor_get_async(ggml_backend_t backend, const struct ggml_ten ggml_backend_synchronize(backend); ggml_backend_tensor_get(tensor, data, offset, size); } else { + ggml_san_access(backend, tensor, offset, size, false, "get_async"); backend->iface.get_tensor_async(backend, tensor, data, offset, size); } } @@ -308,6 +318,9 @@ void ggml_backend_tensor_set_2d_async(ggml_backend_t backend, struct ggml_tensor GGML_ASSERT(tensor->data != NULL && "tensor not allocated"); GGML_ASSERT(offset + (n_copies-1)*stride_tensor + size <= ggml_nbytes(tensor) && "tensor write out of bounds"); + for (size_t i = 0; i < n_copies; i++) { + ggml_san_access(backend, tensor, offset + i*stride_tensor, size, true, "set_2d_async"); + } backend->iface.set_tensor_2d_async(backend, tensor, data, offset, size, n_copies, stride_tensor, stride_data); } @@ -329,6 +342,9 @@ void ggml_backend_tensor_get_2d_async(ggml_backend_t backend, const struct ggml_ GGML_ASSERT(tensor->data != NULL && "tensor not allocated"); GGML_ASSERT(offset + (n_copies-1)*stride_tensor + size <= ggml_nbytes(tensor) && "tensor read out of bounds"); + for (size_t i = 0; i < n_copies; i++) { + ggml_san_access(backend, tensor, offset + i*stride_tensor, size, false, "get_2d_async"); + } backend->iface.get_tensor_2d_async(backend, tensor, data, offset, size, n_copies, stride_tensor, stride_data); } @@ -344,6 +360,7 @@ void ggml_backend_tensor_set(struct ggml_tensor * tensor, const void * data, siz GGML_ASSERT(tensor->data != NULL && "tensor not allocated"); GGML_ASSERT(offset + size <= ggml_nbytes(tensor) && "tensor write out of bounds"); + ggml_san_access(NULL, tensor, offset, size, true, "tensor_set"); buf->iface.set_tensor(buf, tensor, data, offset, size); } @@ -359,6 +376,7 @@ void ggml_backend_tensor_get(const struct ggml_tensor * tensor, void * data, siz GGML_ASSERT(tensor->data != NULL && "tensor not allocated"); GGML_ASSERT(offset + size <= ggml_nbytes(tensor) && "tensor read out of bounds"); + ggml_san_access(NULL, tensor, offset, size, false, "tensor_get"); buf->iface.get_tensor(buf, tensor, data, offset, size); } @@ -381,6 +399,9 @@ void ggml_backend_tensor_set_2d(struct ggml_tensor * tensor, const void * data, GGML_ASSERT(tensor->data != NULL && "tensor not allocated"); GGML_ASSERT(offset + (n_copies-1)*stride_tensor + size <= ggml_nbytes(tensor) && "tensor write out of bounds"); + for (size_t i = 0; i < n_copies; i++) { + ggml_san_access(NULL, tensor, offset + i*stride_tensor, size, true, "tensor_set_2d"); + } buf->iface.set_tensor_2d(buf, tensor, data, offset, size, n_copies, stride_tensor, stride_data); } @@ -403,6 +424,9 @@ void ggml_backend_tensor_get_2d(const struct ggml_tensor * tensor, void * data, GGML_ASSERT(tensor->data != NULL && "tensor not allocated"); GGML_ASSERT(offset + (n_copies-1)*stride_tensor + size <= ggml_nbytes(tensor) && "tensor read out of bounds"); + for (size_t i = 0; i < n_copies; i++) { + ggml_san_access(NULL, tensor, offset + i*stride_tensor, size, false, "tensor_get_2d"); + } buf->iface.get_tensor_2d(buf, tensor, data, offset, size, n_copies, stride_tensor, stride_data); } @@ -419,9 +443,22 @@ void ggml_backend_tensor_memset(struct ggml_tensor * tensor, uint8_t value, size GGML_ASSERT(offset + size <= ggml_nbytes(tensor) && "tensor write out of bounds"); GGML_ASSERT(buf->iface.memset_tensor != NULL && "memset not implemented by backend buffer"); + ggml_san_access(NULL, tensor, offset, size, true, "tensor_memset"); buf->iface.memset_tensor(buf, tensor, value, offset, size); } +void ggml_backend_tensor_set_direct(struct ggml_tensor * tensor, size_t offset, size_t size) { + if (ggml_san_level() == 0) { + return; + } + + GGML_ASSERT(tensor); + GGML_ASSERT(tensor->data != NULL && "tensor not allocated"); + GGML_ASSERT(offset + size <= ggml_nbytes(tensor) && "tensor write out of bounds"); + + ggml_san_access(NULL, tensor, offset, size, true, "set_direct"); +} + void ggml_backend_synchronize(ggml_backend_t backend) { GGML_ASSERT(backend); if (backend->iface.synchronize == NULL) { @@ -429,6 +466,7 @@ void ggml_backend_synchronize(ggml_backend_t backend) { } backend->iface.synchronize(backend); + ggml_san_sync(backend); } ggml_backend_graph_plan_t ggml_backend_graph_plan_create(ggml_backend_t backend, struct ggml_cgraph * cgraph) { @@ -460,7 +498,11 @@ enum ggml_status ggml_backend_graph_compute(ggml_backend_t backend, struct ggml_ enum ggml_status ggml_backend_graph_compute_async(ggml_backend_t backend, struct ggml_cgraph * cgraph) { GGML_ASSERT(backend); - return backend->iface.graph_compute(backend, cgraph); + const enum ggml_status status = backend->iface.graph_compute(backend, cgraph); + if (status == GGML_STATUS_SUCCESS) { + ggml_san_compute(backend, cgraph); + } + return status; } bool ggml_backend_supports_op(ggml_backend_t backend, const struct ggml_tensor * op) { @@ -493,9 +535,11 @@ void ggml_backend_tensor_copy(const struct ggml_tensor * src, struct ggml_tensor } if (ggml_backend_buffer_is_host(src->buffer)) { + ggml_san_access(NULL, src, 0, ggml_nbytes(src), false, "tensor_copy src"); ggml_backend_tensor_set(dst, src->data, 0, ggml_nbytes(src)); } else if (ggml_backend_buffer_is_host(dst->buffer)) { ggml_backend_tensor_get(src, dst->data, 0, ggml_nbytes(src)); + ggml_san_access(NULL, dst, 0, ggml_nbytes(src), true, "tensor_copy dst"); } else if (!ggml_backend_buffer_copy_tensor(src, dst)) { #ifndef NDEBUG GGML_LOG_DEBUG("%s: warning: slow copy from %s to %s\n", __func__, ggml_backend_buffer_name(src->buffer), ggml_backend_buffer_name(dst->buffer)); @@ -517,7 +561,9 @@ void ggml_backend_tensor_copy_async(ggml_backend_t backend_src, ggml_backend_t b GGML_ASSERT(backend_dst); if (backend_dst->iface.cpy_tensor_async != NULL) { - if (backend_dst->iface.cpy_tensor_async(backend_src, backend_dst, src, dst)) { + const bool accepted = backend_dst->iface.cpy_tensor_async(backend_src, backend_dst, src, dst); + if (accepted) { + ggml_san_cpy_async(backend_src, backend_dst, src, dst); return; } } @@ -551,6 +597,7 @@ void ggml_backend_event_record(ggml_backend_event_t event, ggml_backend_t backen GGML_ASSERT(backend->iface.event_record != NULL); backend->iface.event_record(backend, event); + ggml_san_event_record(event, backend); } void ggml_backend_event_synchronize(ggml_backend_event_t event) { @@ -558,6 +605,7 @@ void ggml_backend_event_synchronize(ggml_backend_event_t event) { GGML_ASSERT(event->device->iface.event_synchronize); event->device->iface.event_synchronize(event->device, event); + ggml_san_event_sync(event); } void ggml_backend_event_wait(ggml_backend_t backend, ggml_backend_event_t event) { @@ -565,6 +613,7 @@ void ggml_backend_event_wait(ggml_backend_t backend, ggml_backend_event_t event) GGML_ASSERT(backend->iface.event_wait != NULL); backend->iface.event_wait(backend, event); + ggml_san_event_wait(backend, event); } static void ggml_backend_graph_optimize(ggml_backend_t backend, struct ggml_cgraph * cgraph, struct ggml_backend_graph_optimize_params * params) { @@ -821,6 +870,13 @@ struct ggml_backend_sched { int n_graph_inputs; int graph_inputs_capacity; + struct ggml_backend_sched_input_slots { + ggml_backend_buffer_t buffer[GGML_SCHED_MAX_COPIES]; + void * data [GGML_SCHED_MAX_COPIES]; + } * graph_input_slots; // [graph_inputs_capacity] + + bool needs_rotate; + struct ggml_context * ctx; ggml_backend_sched_eval_callback callback_eval; @@ -872,9 +928,39 @@ static void ggml_backend_sched_graph_inputs_grow(ggml_backend_sched_t sched) { GGML_ABORT("failed to grow graph inputs container"); } sched->graph_inputs = pnew; + + auto * snew = (struct ggml_backend_sched::ggml_backend_sched_input_slots *) realloc( + (void *) sched->graph_input_slots, new_cap * sizeof(sched->graph_input_slots[0])); + if (snew == NULL) { + GGML_LOG_ERROR("%s: failed to allocate %zu bytes\n", __func__, new_cap * sizeof(sched->graph_input_slots[0])); + GGML_ABORT("failed to grow graph input slots container"); + } + sched->graph_input_slots = snew; + sched->graph_inputs_capacity = new_cap; } +static struct ggml_tensor * ggml_backend_sched_graph_input(struct ggml_tensor * t) { + while (t != NULL) { + if (t->flags & GGML_TENSOR_FLAG_INPUT) { + return t; + } + t = t->view_src; + } + return NULL; +} + +static void ggml_backend_sched_reinit_view(struct ggml_tensor * t) { + if (t->view_src == NULL) { + return; + } + + ggml_backend_sched_reinit_view(t->view_src); + + t->buffer = t->view_src->buffer; + t->data = (char *) t->view_src->data + t->view_offs; +} + // returns the priority of the backend, lower id is higher priority static int ggml_backend_sched_backend_id(ggml_backend_sched_t sched, ggml_backend_t backend) { for (int i = 0; i < sched->n_backends; i++) { @@ -1067,6 +1153,7 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra // reset splits sched->n_splits = 0; sched->n_graph_inputs = 0; + ggml_gallocr_clear_pins(sched->galloc); sched->is_reset = false; struct ggml_init_params params = { @@ -1270,6 +1357,25 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra *cur_backend_id = tensor_backend_id(node->view_src); SET_CAUSE(node, "4.vsrc"); } + + if (node->view_src != NULL && !ggml_is_view_op(node->op)) { + const int view_src_backend_id = tensor_backend_id(node->view_src); + if (view_src_backend_id != -1 && *cur_backend_id != view_src_backend_id) { + if (ggml_backend_supports_op(sched->backends[view_src_backend_id], node)) { + if (sched->debug) { + GGML_LOG_DEBUG("%s: %s (%s) moved to %s to keep it with the tensor it aliases (%s)\n", + __func__, node->name, ggml_op_name(node->op), + ggml_backend_name(sched->backends[view_src_backend_id]), node->view_src->name); + } + *cur_backend_id = view_src_backend_id; + SET_CAUSE(node, "4.alias"); + } else if (sched->debug) { + GGML_LOG_DEBUG("%s: %s (%s) aliases %s on %s, which cannot run it - the write will be lost\n", + __func__, node->name, ggml_op_name(node->op), node->view_src->name, + ggml_backend_name(sched->backends[view_src_backend_id])); + } + } + } for (int j = 0; j < GGML_MAX_SRC; j++) { struct ggml_tensor * src = node->src[j]; if (src == NULL) { @@ -1338,6 +1444,20 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra break; } } + // check if the split has too many inputs + // FIXME: count the number of inputs instead of only checking when full + // cut on the constant, not on inputs_capacity: capacity doubles on demand and + // is never reset, so using it lets the cut point drift up and keeps every input + // copy of an ever-longer split live at once + if (split->n_inputs >= GGML_SCHED_MAX_SPLIT_INPUTS) { + const size_t id = hash_id(src); + int src_backend_id = sched->hv_tensor_backend_ids[id]; + bool supported = ggml_backend_sched_buffer_supported(sched, src, cur_backend_id); + if (src_backend_id != cur_backend_id && tensor_id_copy(id, cur_backend_id, 0) == NULL && !supported) { + need_new_split = true; + break; + } + } } } @@ -1372,27 +1492,57 @@ void ggml_backend_sched_split_graph(ggml_backend_sched_t sched, struct ggml_cgra const int src_backend_id = sched->hv_tensor_backend_ids[src_id]; GGML_ASSERT(src_backend_id != -1); // all inputs should be assigned by now - if (src->flags & GGML_TENSOR_FLAG_INPUT && sched->n_copies > 1) { - if (tensor_id_copy(src_id, src_backend_id, 0) == NULL) { - ggml_backend_t backend = sched->backends[src_backend_id]; + struct ggml_tensor * inp = ggml_backend_sched_graph_input(src); + + if (inp != NULL && sched->n_copies > 1) { + const size_t inp_id = hash_id(inp); + const int inp_backend_id = sched->hv_tensor_backend_ids[inp_id]; + + if (inp_backend_id != -1 && tensor_id_copy(inp_id, inp_backend_id, 0) == NULL) { + ggml_backend_t backend = sched->backends[inp_backend_id]; for (int c = 0; c < sched->n_copies; c++) { struct ggml_tensor * tensor_copy; - if (c == sched->cur_copy) { - tensor_copy = src; // use the original tensor as the current copy + if (c == 0) { + tensor_copy = inp; // use the original tensor as the current copy } else { - tensor_copy = ggml_dup_tensor_layout(sched->ctx, src); - ggml_format_name(tensor_copy, "%s#%s#%d", ggml_backend_name(backend), src->name, c); + tensor_copy = ggml_dup_tensor_layout(sched->ctx, inp); + ggml_format_name(tensor_copy, "%s#%s#%d", ggml_backend_name(backend), inp->name, c); } ggml_set_input(tensor_copy); ggml_set_output(tensor_copy); // prevent ggml-alloc from overwriting the tensor - tensor_id_copy(src_id, src_backend_id, c) = tensor_copy; + tensor_id_copy(inp_id, inp_backend_id, c) = tensor_copy; SET_CAUSE(tensor_copy, "4.cpy"); } int n_graph_inputs = sched->n_graph_inputs++; if (n_graph_inputs >= sched->graph_inputs_capacity) { ggml_backend_sched_graph_inputs_grow(sched); } - sched->graph_inputs[n_graph_inputs] = src; + sched->graph_inputs[n_graph_inputs] = inp; + } + } + + { + struct ggml_tensor * root = src; + while (root->view_src != NULL) { + root = root->view_src; + } + + const int root_backend_id = sched->hv_tensor_backend_ids[hash_id(root)]; + + static const bool pin_async_reads = []() { + const char * env = getenv("GGML_SCHED_PIN_ASYNC_READS"); + return env ? atoi(env) != 0 : true; + }(); + + if (pin_async_reads && + root_backend_id != -1 && root_backend_id != cur_backend_id && + sched->backends[cur_backend_id]->iface.synchronize != NULL && + ggml_backend_sched_buffer_supported(sched, root, cur_backend_id)) { + ggml_gallocr_pin_tensor(sched->galloc, root); + if (sched->debug) { + GGML_LOG_DEBUG("%s: pinned %s, read in place by %s\n", __func__, + root->name, ggml_backend_name(sched->backends[cur_backend_id])); + } } } @@ -1665,6 +1815,8 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s } } + ggml_san_split(split_id, split_backend, split->n_inputs); + // copy the input tensors to the split backend for (int input_id = 0; input_id < split->n_inputs; input_id++) { ggml_backend_t input_backend = ggml_backend_sched_get_tensor_backend(sched, split->inputs[input_id]); @@ -1779,7 +1931,14 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s } else { // try async copy, but if not possible, we can still use a sync copy without synchronizing the dst backend, since we handle the synchronization here with multiple copies and events // TODO: add public function to facilitate this, since applications do not have direct access to the backend interface - if (!split_backend->iface.cpy_tensor_async || !split_backend->iface.cpy_tensor_async(input_backend, split_backend, input, input_cpy)) { + bool cpy_async_ok = false; + if (split_backend->iface.cpy_tensor_async) { + cpy_async_ok = split_backend->iface.cpy_tensor_async(input_backend, split_backend, input, input_cpy); + if (cpy_async_ok) { + ggml_san_cpy_async(input_backend, split_backend, input, input_cpy); + } + } + if (!cpy_async_ok) { ggml_backend_synchronize(input_backend); if (sched->events[split_backend_id][sched->cur_copy] != NULL) { ggml_backend_event_synchronize(sched->events[split_backend_id][sched->cur_copy]); @@ -1795,6 +1954,7 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s if (!sched->callback_eval) { enum ggml_status ec = ggml_backend_graph_compute_async(split_backend, &split->graph); if (ec != GGML_STATUS_SUCCESS) { + ggml_san_split(-1, NULL, 0); return ec; } } else { @@ -1817,6 +1977,7 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s enum ggml_status ec = ggml_backend_graph_compute_async(split_backend, &gv); if (ec != GGML_STATUS_SUCCESS) { + ggml_san_split(-1, NULL, 0); return ec; } @@ -1839,6 +2000,10 @@ static enum ggml_status ggml_backend_sched_compute_splits(ggml_backend_sched_t s prev_backend_id = split_backend_id; } + ggml_san_split(-1, NULL, 0); + + sched->needs_rotate = sched->n_copies > 1; + return GGML_STATUS_SUCCESS; } @@ -1866,7 +2031,37 @@ ggml_backend_sched_t ggml_backend_sched_new( sched->debug_realloc = GGML_SCHED_DEBUG_REALLOC ? atoi(GGML_SCHED_DEBUG_REALLOC) : sched->debug_realloc; sched->n_backends = n_backends; - sched->n_copies = parallel ? GGML_SCHED_MAX_COPIES : 1; + + bool is_uma = false; + if (!parallel && n_backends >= 2) { + ggml_backend_buffer_type_t cpu_buft = bufts ? bufts[n_backends - 1] + : ggml_backend_get_default_buffer_type(backends[n_backends - 1]); + + if (cpu_buft && ggml_backend_buft_is_host(cpu_buft)) { + for (int b = 0; b < n_backends - 1; b++) { + ggml_backend_dev_props props; + ggml_backend_dev_get_props(ggml_backend_get_device(backends[b]), &props); + if (!props.caps.async) { + continue; + } + if (ggml_backend_supports_buft(backends[b], cpu_buft)) { + is_uma = true; + GGML_LOG_DEBUG("%s: %s computes directly on %s, ring buffering graph inputs\n", + __func__, ggml_backend_name(backends[b]), ggml_backend_buft_name(cpu_buft)); + break; + } + } + } + } + + int n_copies_uma = is_uma ? 2 : 1; + + const char * GGML_SCHED_UMA_RING = getenv("GGML_SCHED_UMA_RING"); + if (GGML_SCHED_UMA_RING && is_uma) { + n_copies_uma = std::min(std::max(atoi(GGML_SCHED_UMA_RING), 1), GGML_SCHED_MAX_COPIES); + } + + sched->n_copies = parallel ? GGML_SCHED_MAX_COPIES : n_copies_uma; // initialize hash table // FIXME: needs to be size*2 to account for leafs (do it in graph_split instead) @@ -1893,6 +2088,8 @@ ggml_backend_sched_t ggml_backend_sched_new( sched->graph_inputs_capacity = GGML_SCHED_MAX_SPLIT_INPUTS; sched->graph_inputs = (struct ggml_tensor **) calloc(sched->graph_inputs_capacity, sizeof(struct ggml_tensor *)); + sched->graph_input_slots = (struct ggml_backend_sched::ggml_backend_sched_input_slots *) calloc( + sched->graph_inputs_capacity, sizeof(sched->graph_input_slots[0])); for (int b = 0; b < n_backends; b++) { sched->backends[b] = backends[b]; @@ -1931,6 +2128,7 @@ void ggml_backend_sched_free(ggml_backend_sched_t sched) { } free(sched->splits); free(sched->graph_inputs); + free(sched->graph_input_slots); free(sched->hv_tensor_backend_ids); free(sched->hv_tensor_copies); free(sched->node_backend_ids); @@ -1986,13 +2184,108 @@ bool ggml_backend_sched_reserve(ggml_backend_sched_t sched, struct ggml_cgraph * return true; } +static void ggml_backend_sched_capture_input_slots(ggml_backend_sched_t sched) { + if (sched->n_copies <= 1) { + return; + } + + for (int i = 0; i < sched->n_graph_inputs; i++) { + struct ggml_tensor * input = sched->graph_inputs[i]; + + const size_t id = hash_id(input); + const int backend_id = tensor_backend_id(input); + + for (int c = 0; c < sched->n_copies; c++) { + struct ggml_tensor * slot = tensor_id_copy(id, backend_id, c); + + sched->graph_input_slots[i].buffer[c] = slot ? slot->buffer : NULL; + sched->graph_input_slots[i].data [c] = slot ? slot->data : NULL; + } + + if (sched->debug) { + GGML_LOG_DEBUG("%s: graph input %d: %s\n", __func__, i, input->name); + } + } +} + +static void ggml_backend_sched_advance_copy(ggml_backend_sched_t sched) { + sched->cur_copy = sched->next_copy; + sched->next_copy = (sched->next_copy + 1) % sched->n_copies; + + if (sched->n_copies <= 1 || !sched->needs_rotate) { + return; + } + + for (int b = 0; b < sched->n_backends; b++) { + if (sched->events[b][sched->cur_copy] != NULL) { + ggml_backend_event_synchronize(sched->events[b][sched->cur_copy]); + } else { + ggml_backend_synchronize(sched->backends[b]); + } + } + + if (sched->debug) { + GGML_LOG_DEBUG("%s: advanced to copy %d/%d\n", __func__, sched->cur_copy, sched->n_copies); + } +} + +static void ggml_backend_sched_point_inputs(ggml_backend_sched_t sched) { + for (int i = 0; i < sched->n_graph_inputs; i++) { + struct ggml_tensor * input = sched->graph_inputs[i]; + + if (sched->graph_input_slots[i].data[sched->cur_copy] == NULL) { + continue; + } + + input->buffer = sched->graph_input_slots[i].buffer[sched->cur_copy]; + input->data = sched->graph_input_slots[i].data [sched->cur_copy]; + } + + for (int i = 0; i < sched->graph.n_nodes; i++) { + struct ggml_tensor * node = sched->graph.nodes[i]; + + for (int j = 0; j < GGML_MAX_SRC; j++) { + struct ggml_tensor * src = node->src[j]; + if (src != NULL && src->view_src != NULL && ggml_backend_sched_graph_input(src) != NULL) { + ggml_backend_sched_reinit_view(src); + } + } + } + + for (int i = 0; i < sched->n_splits; i++) { + sched->splits[i].graph.uid = ggml_graph_next_uid(); + } + + if (sched->debug) { + GGML_LOG_DEBUG("%s: rotated %d graph inputs onto copy %d\n", + __func__, sched->n_graph_inputs, sched->cur_copy); + } +} + +static void ggml_backend_sched_rotate_inputs(ggml_backend_sched_t sched) { + GGML_ASSERT(sched->n_copies > 1); + + ggml_backend_sched_advance_copy(sched); + ggml_backend_sched_point_inputs(sched); +} + +void ggml_backend_sched_prepare_inputs(ggml_backend_sched_t sched) { + GGML_ASSERT(sched); + + if (sched->n_copies <= 1 || !sched->needs_rotate) { + return; + } + + ggml_backend_sched_rotate_inputs(sched); + sched->needs_rotate = false; +} + bool ggml_backend_sched_alloc_graph(ggml_backend_sched_t sched, struct ggml_cgraph * graph) { GGML_ASSERT(sched); GGML_ASSERT((int)sched->hash_set.size >= graph->n_nodes + graph->n_leafs); GGML_ASSERT(!sched->is_alloc); - sched->cur_copy = sched->next_copy; - sched->next_copy = (sched->next_copy + 1) % sched->n_copies; + ggml_backend_sched_advance_copy(sched); ggml_backend_sched_split_graph(sched, graph); @@ -2000,6 +2293,14 @@ bool ggml_backend_sched_alloc_graph(ggml_backend_sched_t sched, struct ggml_cgra return false; } + ggml_backend_sched_capture_input_slots(sched); + + if (sched->n_copies > 1) { + ggml_backend_sched_point_inputs(sched); + } + + sched->needs_rotate = false; + sched->is_alloc = true; return true; @@ -2031,6 +2332,8 @@ void ggml_backend_sched_synchronize(ggml_backend_sched_t sched) { for (int i = 0; i < sched->n_backends; i++) { ggml_backend_synchronize(sched->backends[i]); } + + sched->needs_rotate = false; if (!sched->is_alloc) { // if the graph is not already allocated, always use copy 0 after a synchronization // this ensures that during generation the same copy is used every time, @@ -2107,6 +2410,7 @@ bool ggml_op_alloc_size_may_expand(enum ggml_op op) { case GGML_OP_FLASH_ATTN_EXT: case GGML_OP_MUL_MAT: case GGML_OP_MUL_MAT_ID: + case GGML_OP_CONCAT: case GGML_OP_CUMSUM: case GGML_OP_ARGSORT: case GGML_OP_TOP_K: diff --git a/ggml/src/ggml-common.h b/ggml/src/ggml-common.h index 1dbbe326d0fc..7a94403a3256 100644 --- a/ggml/src/ggml-common.h +++ b/ggml/src/ggml-common.h @@ -1128,6 +1128,14 @@ GGML_TABLE_BEGIN(int8_t, kvalues_fp4, 16) GGML_TABLE_END() #define kvalues_mxfp4 kvalues_fp4 +// ROCmFP4 uses an E2M1-derived value set with the largest level retuned from +// 12 to 10, plus dual half-block UE4M3 scales. Keeping this separate from +// MXFP4/NVFP4 lets the experimental Strix Halo format evolve without changing +// stock MXFP4/NVFP4 behavior. +GGML_TABLE_BEGIN(int8_t, kvalues_rocmfp4, 16) + 0, 1, 2, 3, 4, 6, 8, 10, 0, -1, -2, -3, -4, -6, -8, -10, +GGML_TABLE_END() + #define NGRID_IQ1S 2048 #define IQ1S_DELTA 0.125f #define IQ1M_DELTA 0.125f diff --git a/ggml/src/ggml-cpu/ggml-cpu.c b/ggml/src/ggml-cpu/ggml-cpu.c index 87a329f26975..91d25d0a66aa 100644 --- a/ggml/src/ggml-cpu/ggml-cpu.c +++ b/ggml/src/ggml-cpu/ggml-cpu.c @@ -16,6 +16,10 @@ #include "ggml.h" #include "common.h" +// ROCmFPx experimental AMD-native quant formats +#include "../../rocmfp4/rocmfp4.h" +#include "../../rocmfpx/rocmfpx.h" + #if defined(_MSC_VER) || defined(__MINGW32__) #include // using malloc.h with MSC/MINGW #elif !defined(__FreeBSD__) && !defined(__NetBSD__) && !defined(__OpenBSD__) @@ -212,6 +216,166 @@ typedef pthread_t ggml_thread_t; #include #endif +static inline uint32_t ggml_rocmfpx_get_bits_cpu(const uint8_t * src, int bit_pos, int nbits) { + const int byte_pos = bit_pos >> 3; + const int shift = bit_pos & 7; + uint32_t v = src[byte_pos]; + + v |= (uint32_t) src[byte_pos + 1] << 8; + v |= (uint32_t) src[byte_pos + 2] << 16; + + return (v >> shift) & ((1u << nbits) - 1u); +} + +static inline int ggml_rocmfpx_decode_fp3_cpu(uint32_t code) { + static const int8_t table[8] = { 0, 1, 2, 4, 0, -1, -2, -4 }; + return table[code & 7u]; +} + +static inline int ggml_rocmfpx_decode_fp2_cpu(uint32_t code) { + static const int8_t table[4] = { -4, -1, 1, 4 }; + return table[code & 3u]; +} + +static void ggml_vec_dot_rocmfpx_fp2_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); + assert(nrc == 1); + GGML_UNUSED(nrc); + assert(n % QK_ROCMFP2 == 0); + assert(QK_ROCMFP2 == QK8_0); + + const block_rocmfp2 * GGML_RESTRICT x = (const block_rocmfp2 *) vx; + const block_q8_0 * GGML_RESTRICT y = (const block_q8_0 *) vy; + const int nb = n / QK_ROCMFP2; + float sumf = 0.0f; + + for (int ib = 0; ib < nb; ++ib) { + const float dy = GGML_CPU_FP16_TO_FP32(y[ib].d); + int sumi0 = 0; + int sumi1 = 0; + for (int j = 0; j < QK_ROCMFP2/2; ++j) { + const uint8_t c0 = (x[ib].qs[j/4] >> (2*(j % 4))) & 3u; + const uint8_t c1 = (x[ib].qs[4 + j/4] >> (2*(j % 4))) & 3u; + sumi0 += ggml_rocmfpx_decode_fp2_cpu(c0) * (int) y[ib].qs[j]; + sumi1 += ggml_rocmfpx_decode_fp2_cpu(c1) * (int) y[ib].qs[j + QK_ROCMFP2/2]; + } + sumf += dy * ( + rocmfpx_ue4m3_to_fp32(x[ib].e[0]) * (float) sumi0 + + rocmfpx_ue4m3_to_fp32(x[ib].e[1]) * (float) sumi1); + } + *s = sumf; +} + +static inline int ggml_rocmfpx_decode_fp6_cpu(uint32_t code) { + // asymmetric signed range [-32, 31]: sign|0 encodes -32, matching + // rocmfpx_decode_fp6_code() in ggml/rocmfpx/rocmfpx.c + const int mag = (int) (code & 31u); + return (code & 32u) ? -(mag == 0 ? 32 : mag) : mag; +} + +static void ggml_vec_dot_rocmfpx_fp3_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); + assert(nrc == 1); + GGML_UNUSED(nrc); + assert(n % QK_ROCMFP3 == 0); + assert(QK_ROCMFP3 == QK8_0); + + const block_rocmfp3 * GGML_RESTRICT x = (const block_rocmfp3 *) vx; + const block_q8_0 * GGML_RESTRICT y = (const block_q8_0 *) vy; + + const int nb = n / QK_ROCMFP3; + float sumf = 0.0f; + + for (int ib = 0; ib < nb; ++ib) { + const float dy = GGML_CPU_FP16_TO_FP32(y[ib].d); + + int sumi0 = 0; + int sumi1 = 0; + for (int j = 0; j < QK_ROCMFP3/2; ++j) { + const int q0 = ggml_rocmfpx_decode_fp3_cpu(ggml_rocmfpx_get_bits_cpu(x[ib].qs, j*3, 3)); + const int q1 = ggml_rocmfpx_decode_fp3_cpu(ggml_rocmfpx_get_bits_cpu(x[ib].qs, (j + QK_ROCMFP3/2)*3, 3)); + sumi0 += q0 * (int) y[ib].qs[j]; + sumi1 += q1 * (int) y[ib].qs[j + QK_ROCMFP3/2]; + } + + sumf += dy * ( + rocmfpx_ue4m3_to_fp32(x[ib].e[0]) * (float) sumi0 + + rocmfpx_ue4m3_to_fp32(x[ib].e[1]) * (float) sumi1); + } + + *s = sumf; +} + +static void ggml_vec_dot_rocmfpx_fp6_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); + assert(nrc == 1); + GGML_UNUSED(nrc); + assert(n % QK_ROCMFP6 == 0); + assert(QK_ROCMFP6 == QK8_0); + + const block_rocmfp6 * GGML_RESTRICT x = (const block_rocmfp6 *) vx; + const block_q8_0 * GGML_RESTRICT y = (const block_q8_0 *) vy; + + const int nb = n / QK_ROCMFP6; + float sumf = 0.0f; + + for (int ib = 0; ib < nb; ++ib) { + const float dy = GGML_CPU_FP16_TO_FP32(y[ib].d); + + int sumi0 = 0; + int sumi1 = 0; + for (int j = 0; j < QK_ROCMFP6/2; ++j) { + const int q0 = ggml_rocmfpx_decode_fp6_cpu(ggml_rocmfpx_get_bits_cpu(x[ib].qs, j*6, 6)); + const int q1 = ggml_rocmfpx_decode_fp6_cpu(ggml_rocmfpx_get_bits_cpu(x[ib].qs, (j + QK_ROCMFP6/2)*6, 6)); + sumi0 += q0 * (int) y[ib].qs[j]; + sumi1 += q1 * (int) y[ib].qs[j + QK_ROCMFP6/2]; + } + + sumf += dy * ( + rocmfpx_ue4m3_to_fp32(x[ib].e[0]) * (float) sumi0 + + rocmfpx_ue4m3_to_fp32(x[ib].e[1]) * (float) sumi1); + } + + *s = sumf; +} + +static void ggml_vec_dot_rocmfpx_fp8_q8_0(int n, float * GGML_RESTRICT s, size_t bs, const void * GGML_RESTRICT vx, size_t bx, const void * GGML_RESTRICT vy, size_t by, int nrc) { + GGML_UNUSED(bs); + GGML_UNUSED(bx); + GGML_UNUSED(by); + assert(nrc == 1); + GGML_UNUSED(nrc); + assert(n % QK_ROCMFP8 == 0); + assert(QK_ROCMFP8 == QK8_0); + + const block_rocmfp8 * GGML_RESTRICT x = (const block_rocmfp8 *) vx; + const block_q8_0 * GGML_RESTRICT y = (const block_q8_0 *) vy; + + const int nb = n / QK_ROCMFP8; + float sumf = 0.0f; + + for (int ib = 0; ib < nb; ++ib) { + const float dx = rocmfpx_ue4m3_to_fp32(x[ib].e); + const float dy = GGML_CPU_FP16_TO_FP32(y[ib].d); + const float d = dx * dy; + + int sumi = 0; + for (int j = 0; j < QK_ROCMFP8; ++j) { + sumi += (int) x[ib].qs[j] * (int) y[ib].qs[j]; + } + + sumf += d * (float) sumi; + } + + *s = sumf; +} + static const struct ggml_type_traits_cpu type_traits_cpu[GGML_TYPE_COUNT] = { [GGML_TYPE_F32] = { .from_float = (ggml_from_float_t) ggml_cpu_fp32_to_fp32, @@ -247,6 +411,42 @@ static const struct ggml_type_traits_cpu type_traits_cpu[GGML_TYPE_COUNT] = { .nrows = 1, #endif }, + [GGML_TYPE_Q4_0_ROCMFP4] = { + .from_float = rocmfp4_quantize_row_q4_0, + .vec_dot = rocmfp4_vec_dot_q4_0_q8_0, + .vec_dot_type = GGML_TYPE_Q8_0, + .nrows = 1, + }, + [GGML_TYPE_Q4_0_ROCMFP4_FAST] = { + .from_float = rocmfp4_quantize_row_q4_0_fast, + .vec_dot = rocmfp4_vec_dot_q4_0_fast_q8_0, + .vec_dot_type = GGML_TYPE_Q8_0, + .nrows = 1, + }, + [GGML_TYPE_Q3_0_ROCMFPX] = { + .from_float = rocmfpx_quantize_row_fp3, + .vec_dot = ggml_vec_dot_rocmfpx_fp3_q8_0, + .vec_dot_type = GGML_TYPE_Q8_0, + .nrows = 1, + }, + [GGML_TYPE_Q2_0_ROCMFPX] = { + .from_float = rocmfpx_quantize_row_fp2, + .vec_dot = ggml_vec_dot_rocmfpx_fp2_q8_0, + .vec_dot_type = GGML_TYPE_Q8_0, + .nrows = 1, + }, + [GGML_TYPE_Q6_0_ROCMFPX] = { + .from_float = rocmfpx_quantize_row_fp6, + .vec_dot = ggml_vec_dot_rocmfpx_fp6_q8_0, + .vec_dot_type = GGML_TYPE_Q8_0, + .nrows = 1, + }, + [GGML_TYPE_Q8_0_ROCMFPX] = { + .from_float = rocmfpx_quantize_row_fp8, + .vec_dot = ggml_vec_dot_rocmfpx_fp8_q8_0, + .vec_dot_type = GGML_TYPE_Q8_0, + .nrows = 1, + }, [GGML_TYPE_Q4_1] = { .from_float = quantize_row_q4_1, .vec_dot = ggml_vec_dot_q4_1_q8_1, @@ -2999,6 +3199,11 @@ struct ggml_cplan ggml_graph_plan( size_t n_chunks = n_tasks; size_t decode = sizeof(float)*(neq2*n_chunks*(2+DV) + n_tasks*(DK + 2*DV)); + // selected-key rows (src[5]) are gathered per thread after the decode partials + if (node->src[5]) { + decode += sizeof(float)*n_tasks*CACHE_LINE_SIZE_F32 + sizeof(int32_t)*n_tasks*node->src[5]->ne[0]; + } + cur += MAX(prefill, decode); } break; case GGML_OP_FLASH_ATTN_BACK: diff --git a/ggml/src/ggml-cpu/ops.cpp b/ggml/src/ggml-cpu/ops.cpp index 266261c5e5a4..916a4e648630 100644 --- a/ggml/src/ggml-cpu/ops.cpp +++ b/ggml/src/ggml-cpu/ops.cpp @@ -673,6 +673,12 @@ void ggml_compute_forward_add( case GGML_TYPE_Q8_0: case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: + case GGML_TYPE_Q4_0_ROCMFP4: + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + case GGML_TYPE_Q3_0_ROCMFPX: + case GGML_TYPE_Q2_0_ROCMFPX: + case GGML_TYPE_Q6_0_ROCMFPX: + case GGML_TYPE_Q8_0_ROCMFPX: case GGML_TYPE_Q2_K: case GGML_TYPE_Q3_K: case GGML_TYPE_Q4_K: @@ -1125,6 +1131,12 @@ void ggml_compute_forward_add1( case GGML_TYPE_Q8_1: case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: + case GGML_TYPE_Q4_0_ROCMFP4: + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + case GGML_TYPE_Q3_0_ROCMFPX: + case GGML_TYPE_Q2_0_ROCMFPX: + case GGML_TYPE_Q6_0_ROCMFPX: + case GGML_TYPE_Q8_0_ROCMFPX: case GGML_TYPE_Q2_K: case GGML_TYPE_Q3_K: case GGML_TYPE_Q4_K: @@ -1256,6 +1268,12 @@ void ggml_compute_forward_acc( case GGML_TYPE_Q8_1: case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: + case GGML_TYPE_Q4_0_ROCMFP4: + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + case GGML_TYPE_Q3_0_ROCMFPX: + case GGML_TYPE_Q2_0_ROCMFPX: + case GGML_TYPE_Q6_0_ROCMFPX: + case GGML_TYPE_Q8_0_ROCMFPX: case GGML_TYPE_Q2_K: case GGML_TYPE_Q3_K: case GGML_TYPE_Q4_K: @@ -4658,6 +4676,12 @@ void ggml_compute_forward_out_prod( case GGML_TYPE_Q8_0: case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: + case GGML_TYPE_Q4_0_ROCMFP4: + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + case GGML_TYPE_Q3_0_ROCMFPX: + case GGML_TYPE_Q2_0_ROCMFPX: + case GGML_TYPE_Q6_0_ROCMFPX: + case GGML_TYPE_Q8_0_ROCMFPX: case GGML_TYPE_Q2_K: case GGML_TYPE_Q3_K: case GGML_TYPE_Q4_K: @@ -4935,6 +4959,12 @@ void ggml_compute_forward_set( case GGML_TYPE_Q8_1: case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: + case GGML_TYPE_Q4_0_ROCMFP4: + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + case GGML_TYPE_Q3_0_ROCMFPX: + case GGML_TYPE_Q2_0_ROCMFPX: + case GGML_TYPE_Q6_0_ROCMFPX: + case GGML_TYPE_Q8_0_ROCMFPX: case GGML_TYPE_Q2_K: case GGML_TYPE_Q3_K: case GGML_TYPE_Q4_K: @@ -5160,6 +5190,12 @@ void ggml_compute_forward_get_rows( case GGML_TYPE_Q8_1: case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: + case GGML_TYPE_Q4_0_ROCMFP4: + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + case GGML_TYPE_Q3_0_ROCMFPX: + case GGML_TYPE_Q2_0_ROCMFPX: + case GGML_TYPE_Q6_0_ROCMFPX: + case GGML_TYPE_Q8_0_ROCMFPX: case GGML_TYPE_Q2_K: case GGML_TYPE_Q3_K: case GGML_TYPE_Q4_K: @@ -5917,6 +5953,12 @@ void ggml_compute_forward_clamp( case GGML_TYPE_Q8_1: case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: + case GGML_TYPE_Q4_0_ROCMFP4: + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + case GGML_TYPE_Q3_0_ROCMFPX: + case GGML_TYPE_Q2_0_ROCMFPX: + case GGML_TYPE_Q6_0_ROCMFPX: + case GGML_TYPE_Q8_0_ROCMFPX: case GGML_TYPE_Q2_K: case GGML_TYPE_Q3_K: case GGML_TYPE_Q4_K: @@ -8624,6 +8666,7 @@ static void ggml_compute_forward_flash_attn_ext_f16_one_chunk( const ggml_tensor * v = dst->src[2]; const ggml_tensor * mask = dst->src[3]; const ggml_tensor * sinks = dst->src[4]; + const ggml_tensor * ids = dst->src[5]; GGML_TENSOR_LOCALS(int64_t, neq, q, ne) GGML_TENSOR_LOCALS(size_t, nbq, q, nb) @@ -8665,6 +8708,12 @@ static void ggml_compute_forward_flash_attn_ext_f16_one_chunk( const int64_t rv2 = neq2/nev2; const int64_t rv3 = neq3/nev3; + // selected keys (src[5]): the first n_kv_raw keys plus key n_kv_raw + i for each valid index i in the query's row. + // they are visited in ascending order, so with a mask that encodes the same selection the result is unchanged + const int64_t n_kv_raw = ids ? ggml_get_op_params_i32(dst, 4) : nek1; + int32_t * keys = ids ? (int32_t *) ((float *) params->wdata + params->nth*(DK + 2*DV + CACHE_LINE_SIZE_F32) + + neq2*params->nth*(2 + DV)) + params->ith*ids->ne[0] : nullptr; + // parallelize by q rows using ggml_vec_dot_f32 float scale = 1.0f; @@ -8731,11 +8780,26 @@ static void ggml_compute_forward_flash_attn_ext_f16_one_chunk( const float * pq = (const float *) ((char *) q->data + (iq1*nbq1 + iq2*nbq2 + iq3*nbq3)); q_to_vec_dot(pq, Q_q, DK); + int64_t n_sel = 0; + if (ids) { + const int32_t * row = (const int32_t *) ((const char *) ids->data + iq1*ids->nb[1] + (iq3%ids->ne[3])*ids->nb[3]); + for (int64_t j = 0; j < ids->ne[0]; ++j) { + const int64_t ic = n_kv_raw + row[j]; + if (row[j] >= 0 && ic >= ic_start && ic < ic_end) { + keys[n_sel++] = (int32_t) ic; + } + } + std::sort(keys, keys + n_sel); + n_sel = std::unique(keys, keys + n_sel) - keys; + } + const int64_t n_dense = std::max(0, std::min(ic_end, n_kv_raw) - ic_start); + // online softmax / attention // loop over n_kv and n_head_kv // ref: https://arxiv.org/pdf/2112.05682.pdf - for (int64_t ic = ic_start; ic < ic_end; ++ic) { + for (int64_t j = 0; j < n_dense + n_sel; ++j) { + const int64_t ic = j < n_dense ? ic_start + j : keys[j - n_dense]; const float mv = mp ? slope*GGML_CPU_FP16_TO_FP32(mp[ic]) : 0.0f; if (mv == -INFINITY) { continue; @@ -9315,11 +9379,13 @@ static void ggml_compute_forward_flash_attn_ext_f16( const int64_t dr = (nr + nchunk - 1) / nchunk; static constexpr int64_t Q_TILE_SZ = ggml_fa_tile_config::Q; + // the tiled kernel attends every key the mask allows, so selected keys without a mask need the per-row path bool use_tiled = !use_ref && (q->type == GGML_TYPE_F32 && kv_is_f32_or_f16 && k->type == v->type && - neq1 >= Q_TILE_SZ); + neq1 >= Q_TILE_SZ && + (!dst->src[5] || dst->src[3])); #ifdef GGML_SIMD #if defined(__ARM_FEATURE_SVE) const int64_t f32_epr = svcntw(); diff --git a/ggml/src/ggml-cuda/argsort.cu b/ggml/src/ggml-cuda/argsort.cu index 26af90025972..17b905fd0b69 100644 --- a/ggml/src/ggml-cuda/argsort.cu +++ b/ggml/src/ggml-cuda/argsort.cu @@ -1,11 +1,16 @@ #include "argsort.cuh" #ifdef GGML_CUDA_USE_CUB -# include -# if (CCCL_MAJOR_VERSION >= 3 && CCCL_MINOR_VERSION >= 1) -# define STRIDED_ITERATOR_AVAILABLE -# include -# endif +# ifdef GGML_CUDA_CUB_IS_HIPCUB +# include +namespace cub = hipcub; +# else +# include +# if (CCCL_MAJOR_VERSION >= 3 && CCCL_MINOR_VERSION >= 1) +# define STRIDED_ITERATOR_AVAILABLE +# include +# endif +# endif // GGML_CUDA_CUB_IS_HIPCUB using namespace cub; #endif // GGML_CUDA_USE_CUB diff --git a/ggml/src/ggml-cuda/binbcast.cu b/ggml/src/ggml-cuda/binbcast.cu index 2e38077bf67f..e12e91336105 100644 --- a/ggml/src/ggml-cuda/binbcast.cu +++ b/ggml/src/ggml-cuda/binbcast.cu @@ -26,6 +26,43 @@ static __device__ __forceinline__ float op_div(const float a, const float b) { return a / b; } +template +static __global__ void binary_contiguous_f32(const float * src0, const float * src1, float * dst, int64_t n) { + ggml_cuda_pdl_lc(); + int64_t i = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t stride = int64_t(blockDim.x) * gridDim.x; + ggml_cuda_pdl_sync(); + for (; i < n; i += stride) { + dst[i] = bin_op(src0[i], src1[i]); + if (n - i <= stride) { + break; + } + } +} + +template +static bool try_binary_contiguous_f32(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + if (src0->type != GGML_TYPE_F32 || src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32 || + src0->ne[0] != 2048 || + !ggml_are_same_shape(src0, src1) || !ggml_are_same_shape(src0, dst) || + !ggml_is_contiguous(src0) || !ggml_is_contiguous(src1) || !ggml_is_contiguous(dst)) { + return false; + } + + const int threads = 256; + const int64_t n = ggml_nelements(dst); + if (n == 0) { + return true; + } + const int blocks = std::min((n - 1) / threads + 1, 65535); + const ggml_cuda_kernel_launch_params launch_params(blocks, threads, 0, ctx.stream()); + ggml_cuda_kernel_launch(binary_contiguous_f32, launch_params, + (const float *) src0->data, (const float *) src1->data, (float *) dst->data, n); + return true; +} + template (ctx, dst)) { + return; + } ggml_cuda_op_bin_bcast>(dst->src[0], dst->src[1], dst, dst->src[0]->data, dst->src[1]->data, dst->data, ctx.stream()); } @@ -443,6 +483,9 @@ void ggml_cuda_op_sub(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { } void ggml_cuda_op_mul(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + if (try_binary_contiguous_f32(ctx, dst)) { + return; + } ggml_cuda_op_bin_bcast>(dst->src[0], dst->src[1], dst, dst->src[0]->data, dst->src[1]->data, dst->data, ctx.stream()); } @@ -542,6 +585,219 @@ void ggml_cuda_op_fused_mul(ggml_backend_cuda_context & ctx, ggml_tensor * dst, } } +#if defined(__HIP_PLATFORM_AMD__) +static __device__ __forceinline__ float rounded_mul_f32(const float a, const float b) { + float result; + asm("v_mul_f32_e32 %0, %1, %2" : "=v"(result) : "v"(a), "v"(b)); + return result; +} + +static __device__ __forceinline__ float rounded_add_f32(const float a, const float b) { + float result; + asm("v_add_f32_e32 %0, %1, %2" : "=v"(result) : "v"(a), "v"(b)); + return result; +} +#else +static __device__ __forceinline__ float rounded_mul_f32(const float a, const float b) { + return __fmul_rn(a, b); +} + +static __device__ __forceinline__ float rounded_add_f32(const float a, const float b) { + return __fadd_rn(a, b); +} +#endif + +template +static __global__ void weighted_expert_sum_f32( + const float * experts, const float * weights, float * dst, + int64_t n_embd, int64_t n_tokens, int64_t experts_s1, int64_t experts_s2, + int64_t weights_s1, int64_t weights_s2) { + const int64_t index = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t n_embd4 = n_embd / 4; + if (index >= n_embd4 * n_tokens) { + return; + } + + const int64_t token = index / n_embd4; + const int64_t channel = 4 * (index - token * n_embd4); + float4 expert = *reinterpret_cast(experts + token * experts_s2 + channel); + const float weight = weights[token * weights_s2]; + // Keep MUL rounding separate from ADD to match the unfused graph without forcing products through scratch memory. + float4 sum = { + rounded_mul_f32(expert.x, weight), + rounded_mul_f32(expert.y, weight), + rounded_mul_f32(expert.z, weight), + rounded_mul_f32(expert.w, weight), + }; +#pragma unroll + for (int i = 1; i < n_expert_used; ++i) { + expert = *reinterpret_cast(experts + token * experts_s2 + i * experts_s1 + channel); + const float weight_i = weights[token * weights_s2 + i * weights_s1]; + sum.x = rounded_add_f32(sum.x, rounded_mul_f32(expert.x, weight_i)); + sum.y = rounded_add_f32(sum.y, rounded_mul_f32(expert.y, weight_i)); + sum.z = rounded_add_f32(sum.z, rounded_mul_f32(expert.z, weight_i)); + sum.w = rounded_add_f32(sum.w, rounded_mul_f32(expert.w, weight_i)); + } + if constexpr (sync_before_store) { + // single block: all expert and weight reads of the block are complete, so dst may alias them + __syncthreads(); + } + reinterpret_cast(dst)[index] = sum; +} + +void ggml_cuda_op_weighted_expert_sum(ggml_backend_cuda_context & ctx, ggml_tensor * mul, ggml_tensor * dst, int n_expert_used) { + const ggml_tensor * experts = ggml_are_same_shape(mul, mul->src[0]) ? mul->src[0] : mul->src[1]; + const ggml_tensor * weights = experts == mul->src[0] ? mul->src[1] : mul->src[0]; + const int64_t n_embd = experts->ne[0]; + const int64_t n_tokens = experts->ne[2]; + GGML_ASSERT(n_embd % 4 == 0); + const int64_t n_items = n_embd / 4 * n_tokens; + + // The graph allocator may place the output over the expert or weight buffers (both die inside the fused + // range). A single block can read everything before a barrier and then write in place; larger grids + // stage the sum until all reads have finished. + auto overlaps = [](const ggml_tensor * a, const ggml_tensor * b) { + const char * a0 = (const char *) a->data; + const char * a1 = a0 + ggml_nbytes(a); + const char * b0 = (const char *) b->data; + const char * b1 = b0 + ggml_nbytes(b); + return a0 < b1 && b0 < a1; + }; + const bool aliased = overlaps(dst, experts) || overlaps(dst, weights); + const bool single_block = n_items <= 1024; + const bool stage = aliased && !single_block; + + const int threads = single_block ? (int) n_items : 256; + const int blocks = (int) ((n_items + threads - 1) / threads); + const ggml_cuda_kernel_launch_params launch_params(blocks, threads, 0, ctx.stream()); + + ggml_cuda_pool_alloc result(ctx.pool()); + float * out = stage ? result.alloc(n_embd * n_tokens) : (float *) dst->data; + +#define launch_weighted_expert_sum(n) \ + do { \ + if (single_block) { \ + ggml_cuda_kernel_launch(weighted_expert_sum_f32, launch_params, \ + (const float *) experts->data, (const float *) weights->data, out, \ + n_embd, n_tokens, experts->nb[1] / sizeof(float), experts->nb[2] / sizeof(float), \ + weights->nb[1] / sizeof(float), weights->nb[2] / sizeof(float)); \ + } else { \ + ggml_cuda_kernel_launch(weighted_expert_sum_f32, launch_params, \ + (const float *) experts->data, (const float *) weights->data, out, \ + n_embd, n_tokens, experts->nb[1] / sizeof(float), experts->nb[2] / sizeof(float), \ + weights->nb[1] / sizeof(float), weights->nb[2] / sizeof(float)); \ + } \ + } while (0) + switch (n_expert_used) { + case 2: launch_weighted_expert_sum(2); break; + case 3: launch_weighted_expert_sum(3); break; + case 4: launch_weighted_expert_sum(4); break; + case 5: launch_weighted_expert_sum(5); break; + case 6: launch_weighted_expert_sum(6); break; + case 7: launch_weighted_expert_sum(7); break; + case 8: launch_weighted_expert_sum(8); break; + case 9: launch_weighted_expert_sum(9); break; + case 10: launch_weighted_expert_sum(10); break; + case 11: launch_weighted_expert_sum(11); break; + case 12: launch_weighted_expert_sum(12); break; + case 13: launch_weighted_expert_sum(13); break; + case 14: launch_weighted_expert_sum(14); break; + case 15: launch_weighted_expert_sum(15); break; + case 16: launch_weighted_expert_sum(16); break; + default: GGML_ABORT("unsupported expert count"); + } +#undef launch_weighted_expert_sum + + if (stage) { + CUDA_CHECK(cudaMemcpyAsync(dst->data, result.get(), ggml_nbytes(dst), cudaMemcpyDeviceToDevice, ctx.stream())); + } +} + +static __global__ void shared_mul_add_f32( + const float * src, const float * gate, const float * other, const float * residual, float * dst, + int64_t n_embd, int64_t nelements) { + const int64_t index = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; + if (index >= nelements) { + return; + } + const float product = ggml_cuda_materialize(src[index] * gate[index / n_embd]); + float value = other[index] + product; + dst[index] = value + residual[index]; +} + +void ggml_cuda_op_shared_mul_add( + ggml_backend_cuda_context & ctx, ggml_tensor * mul, const ggml_tensor * other, + const ggml_tensor * residual, ggml_tensor * dst) { + const ggml_tensor * src = ggml_are_same_shape(mul, mul->src[0]) ? mul->src[0] : mul->src[1]; + const ggml_tensor * gate = src == mul->src[0] ? mul->src[1] : mul->src[0]; + const int threads = 256; + const int64_t nelements = ggml_nelements(dst); + const int blocks = (nelements + threads - 1) / threads; + const ggml_cuda_kernel_launch_params launch_params(blocks, threads, 0, ctx.stream()); + ggml_cuda_kernel_launch(shared_mul_add_f32, launch_params, + (const float *) src->data, (const float *) gate->data, (const float *) other->data, + (const float *) residual->data, (float *) dst->data, dst->ne[0], nelements); +} + +// Qwen3.5 shared-expert gate for one token: gate = sigmoid(w . y) computed exactly like +// mul_mat_vec_f (same per-thread fma chain and the same two-level warp reduction), +// then dst = (other + src*gate) + residual as in shared_mul_add_f32. One block, so the activation vector is +// fully read before the reduction barrier and the output may alias it. +static __global__ void __launch_bounds__(256, 1) shared_gate_mul_add_f32( + const float * w, const float * y, const float * src, const float * other, const float * residual, float * dst, + const int ncols2, const int64_t n_embd) { + constexpr int block_size = 256; + constexpr int warp_size = 32; + __shared__ float buf_iw[warp_size]; + __shared__ float gate_s; + + const int tid = threadIdx.x; + if (tid < warp_size) { + buf_iw[tid] = 0.0f; + } + __syncthreads(); + + const float2 * x2 = (const float2 *) w; + const float2 * y2 = (const float2 *) y; + float sumf = 0.0f; + for (int col2 = tid; col2 < ncols2; col2 += block_size) { + const float2 tmpx = x2[col2]; + const float2 tmpy = y2[col2]; + ggml_cuda_mad(sumf, tmpx.x, tmpy.x); + ggml_cuda_mad(sumf, tmpx.y, tmpy.y); + } + sumf = warp_reduce_sum(sumf); + buf_iw[tid/warp_size] = sumf; + __syncthreads(); + if (tid < warp_size) { + sumf = buf_iw[tid]; + sumf = warp_reduce_sum(sumf); + if (tid == 0) { + gate_s = 1.0f / (1.0f + expf(-sumf)); + } + } + __syncthreads(); + + const float gate = gate_s; + for (int64_t i = tid; i < n_embd; i += block_size) { + const float product = rounded_mul_f32(src[i], gate); + const float value = rounded_add_f32(other[i], product); + dst[i] = residual ? rounded_add_f32(value, residual[i]) : value; + } +} + +void ggml_cuda_op_shared_gate_mul_add( + ggml_backend_cuda_context & ctx, const ggml_tensor * gate_mm, const ggml_tensor * src, const ggml_tensor * other, + const ggml_tensor * residual, ggml_tensor * dst) { + const ggml_tensor * w = gate_mm->src[0]; + const ggml_tensor * y = gate_mm->src[1]; + GGML_ASSERT(w->ne[0] % 2 == 0); + const ggml_cuda_kernel_launch_params launch_params(1, 256, 0, ctx.stream()); + ggml_cuda_kernel_launch(shared_gate_mul_add_f32, launch_params, + (const float *) w->data, (const float *) y->data, (const float *) src->data, (const float *) other->data, + residual ? (const float *) residual->data : nullptr, (float *) dst->data, (int) (w->ne[0] / 2), dst->ne[0]); +} + void ggml_cuda_op_repeat_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; diff --git a/ggml/src/ggml-cuda/binbcast.cuh b/ggml/src/ggml-cuda/binbcast.cuh index 12624785b444..bb3a8e13f241 100644 --- a/ggml/src/ggml-cuda/binbcast.cuh +++ b/ggml/src/ggml-cuda/binbcast.cuh @@ -10,3 +10,9 @@ void ggml_cuda_op_repeat_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst void ggml_cuda_op_fused_add(ggml_backend_cuda_context & ctx, ggml_tensor * dst, int n_fuse); void ggml_cuda_op_fused_mul(ggml_backend_cuda_context & ctx, ggml_tensor * dst, int n_fuse); +void ggml_cuda_op_weighted_expert_sum(ggml_backend_cuda_context & ctx, ggml_tensor * mul, ggml_tensor * dst, int n_expert_used); +void ggml_cuda_op_shared_mul_add(ggml_backend_cuda_context & ctx, ggml_tensor * mul, const ggml_tensor * other, const ggml_tensor * residual, ggml_tensor * dst); + +void ggml_cuda_op_shared_gate_mul_add( + ggml_backend_cuda_context & ctx, const ggml_tensor * gate_mm, const ggml_tensor * src, const ggml_tensor * other, + const ggml_tensor * residual, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/common.cuh b/ggml/src/ggml-cuda/common.cuh index 2e78ae4facb9..3ab92800552e 100644 --- a/ggml/src/ggml-cuda/common.cuh +++ b/ggml/src/ggml-cuda/common.cuh @@ -115,6 +115,14 @@ # define GGML_CUDA_USE_CUB #endif // !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) && CUDART_VERSION >= 11070 +// hipCUB (rocPRIM) offers the same device-wide radix sort / scan / reduce entry points as CUB. Using it lifts the +// ncols <= 1024 limit of ARGSORT / TOP_K on HIP, which otherwise falls back to the CPU (e.g. the QSA indexer top-k +// of qwen4exp / DeepSeek-V3.2 at contexts beyond 2k tokens). +#if defined(GGML_USE_HIP) && !defined(GGML_HIP_NO_CUB) && __has_include() +# define GGML_CUDA_USE_CUB +# define GGML_CUDA_CUB_IS_HIPCUB +#endif + // PDL host-side support (cudaLaunchKernelEx) requires CUDART >= 11.8. // However, this has been bugged in CTK < 12.3 for MSVC builds, see // https://github.com/ggml-org/llama.cpp/pull/22522#discussion_r3302393293 @@ -717,6 +725,24 @@ static __device__ __forceinline__ uint32_t __hgt2_mask(const half2 a, const half } #endif // (defined(CUDART_VERSION) && CUDART_VERSION < CUDART_HMASK) || defined(GGML_USE_HIP) || (defined(MUSART_VERSION) && MUSART_VERSION < MUSART_HMASK) +// Rounds x to a float in a register before its next use, so the compiler cannot contract it into the following add or multiply. +static __device__ __forceinline__ float ggml_cuda_materialize(float x) { +#if defined(GGML_USE_HIP) + asm volatile("" : "+v"(x)); +#else + volatile float v = x; + x = v; +#endif // defined(GGML_USE_HIP) + return x; +} + +// F32 -> BF16, round to nearest even: the rounding of the MMB activation conversion (mmb_cvt_f32_bf16). +static __device__ __forceinline__ uint16_t ggml_cuda_f32_to_bf16_rne(const float f) { + uint32_t u = __float_as_uint(f); + u += 0x7fffu + ((u >> 16) & 1u); + return (uint16_t) (u >> 16); +} + static __device__ __forceinline__ int ggml_cuda_dp4a(const int a, const int b, int c) { #if defined(GGML_USE_HIP) #if defined(CDNA) || defined(RDNA2) || defined(__gfx906__) @@ -1452,6 +1478,8 @@ struct ggml_cuda_stream_context { } }; +struct ggml_cuda_mmb_context; + struct ggml_backend_cuda_context { int device; std::string name; @@ -1464,6 +1492,14 @@ struct ggml_backend_cuda_context { int curr_stream_no = 0; + ggml_cuda_mmb_context * mmb = nullptr; + bool mmb_opt_in = false; // set by ggml_backend_cuda_set_mmb_enabled, before the first graph + bool mmb_after_compute = true; + const void * mmb_first_split = nullptr; + std::vector mmb_graph_sigs; + // HC residuals whose remaining readers live in a later split of the same graph on this device: tensor -> uses left + std::unordered_map mmb_res16_pending; + #ifdef USE_CUDA_GRAPH // Map from first_node_ptr to cuda_graph - allows multiple graphs per context // when the computation is split across CPU/GPU (e.g., with --n-cpu-moe) @@ -1580,6 +1616,10 @@ struct ggml_cuda_mm_fusion_args_host { const ggml_tensor * gate_scale = nullptr; ggml_glu_op glu_op; float glu_limit = 0.0f; + // RDNA3.5 fused-quantize matvec prologue (y_op == 3), see ggml_cuda_mm_fusion_args_device + const ggml_tensor * y_gate = nullptr; + const ggml_tensor * y_norm_w = nullptr; + float y_eps = 0.0f; }; struct ggml_cuda_mm_fusion_args_device { const void * x_bias = nullptr; @@ -1589,6 +1629,10 @@ struct ggml_cuda_mm_fusion_args_device { const void * gate_scale = nullptr; ggml_glu_op glu_op; float glu_limit = 0.0f; + // RDNA3.5 fused-quantize matvec prologue (y_op == 3): y' = sigmoid(y_gate) * (rms_norm_128(y) * y_norm_w) + const void * y_gate = nullptr; + const void * y_norm_w = nullptr; + float y_eps = 0.0f; }; struct ggml_cuda_kernel_launch_params { diff --git a/ggml/src/ggml-cuda/concat.cu b/ggml/src/ggml-cuda/concat.cu index 6df89013ca79..408cbce3d219 100644 --- a/ggml/src/ggml-cuda/concat.cu +++ b/ggml/src/ggml-cuda/concat.cu @@ -2,6 +2,13 @@ #include +static constexpr int concat_transposed_tile_y(int cc) { + return GGML_CUDA_CC_IS_RDNA3_5(cc) ? 16 : 8; +} + +static_assert(concat_transposed_tile_y(GGML_CUDA_CC_RDNA3_5) == 16); +static_assert(concat_transposed_tile_y(GGML_CUDA_CC_RDNA3) == 8); + // contiguous kernels template static __global__ void __launch_bounds__(CUDA_CONCAT_BLOCK_SIZE) concat_cont(const T * x, @@ -79,6 +86,54 @@ static void concat_cont_cuda(const T * x, concat_cont<<>>(x, y, dst, ne00, ne01, ne02, ne0, ne1, ne2); } +template +static __global__ void concat_transposed_src1_dim0( + const char * src0, const char * src1, char * dst, + int64_t ne00, int64_t ne10, int64_t ne11, + uint64_t nb00, uint64_t nb01, uint64_t nb02, + uint64_t nb10, uint64_t nb11, uint64_t nb12, + uint64_t nb0, uint64_t nb1, uint64_t nb2) { + constexpr int tile = 32; + __shared__ T values[tile][tile + 1]; + + const int tx = threadIdx.x; + const int ty = threadIdx.y; + const int64_t channel0 = blockIdx.x * tile; + const int64_t token0 = blockIdx.y * tile; + const int64_t i2 = blockIdx.z; + + ggml_cuda_pdl_sync(); + + for (int j = 0; j < tile; j += blockDim.y) { + const int64_t token = token0 + ty + j; + const int64_t channel = channel0 + tx; + if (token < ne10 && channel < ne11) { + values[ty + j][tx] = *reinterpret_cast(src1 + i2 * nb12 + channel * nb11 + token * nb10); + } + } + __syncthreads(); + + for (int j = 0; j < tile; j += blockDim.y) { + const int64_t token = token0 + tx; + const int64_t channel = channel0 + ty + j; + if (token < ne10 && channel < ne11) { + *reinterpret_cast(dst + i2 * nb2 + channel * nb1 + (ne00 + token) * nb0) = values[tx][ty + j]; + } + } + + if (blockIdx.y == 0) { + for (int j = 0; j < tile; j += blockDim.y) { + const int64_t channel = channel0 + ty + j; + if (channel < ne11) { + for (int64_t state = tx; state < ne00; state += tile) { + *reinterpret_cast(dst + i2 * nb2 + channel * nb1 + state * nb0) = + *reinterpret_cast(src0 + i2 * nb02 + channel * nb01 + state * nb00); + } + } + } + } +} + // non-contiguous kernel (slow) template static __global__ void __launch_bounds__(CUDA_CONCAT_BLOCK_SIZE) @@ -141,6 +196,24 @@ static __global__ void __launch_bounds__(CUDA_CONCAT_BLOCK_SIZE) template static void concat_cuda(const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, int dim, cudaStream_t stream) { + if (dim == 0 && src1->ne[0] >= 32 && src0->ne[3] == 1 && src1->ne[3] == 1 && dst->ne[3] == 1 && + src0->nb[0] == sizeof(T) && src0->nb[1] == src0->ne[0] * sizeof(T) && + src1->nb[1] == sizeof(T) && src1->nb[0] == src1->ne[1] * sizeof(T) && + dst->nb[0] == sizeof(T) && dst->nb[1] == dst->ne[0] * sizeof(T) && + src0->ne[1] == src1->ne[1] && src0->ne[2] == src1->ne[2]) { + constexpr int tile = 32; + const int tile_y = concat_transposed_tile_y(ggml_cuda_info().devices[ggml_cuda_get_device()].cc); + const dim3 block_dims(tile, tile_y, 1); + const dim3 block_nums((src1->ne[1] + tile - 1) / tile, (src1->ne[0] + tile - 1) / tile, src1->ne[2]); + concat_transposed_src1_dim0<<>>( + (const char *) src0->data, (const char *) src1->data, (char *) dst->data, + src0->ne[0], src1->ne[0], src1->ne[1], + src0->nb[0], src0->nb[1], src0->nb[2], + src1->nb[0], src1->nb[1], src1->nb[2], + dst->nb[0], dst->nb[1], dst->nb[2]); + return; + } + if (dim != 3 && ggml_is_contiguous_to_3(src0) && ggml_is_contiguous_to_3(src1)) { const T * src0_d = (const T *) src0->data; const T * src1_d = (const T *) src1->data; diff --git a/ggml/src/ggml-cuda/conv-tail.cuh b/ggml/src/ggml-cuda/conv-tail.cuh new file mode 100644 index 000000000000..c943211d7c8e --- /dev/null +++ b/ggml/src/ggml-cuda/conv-tail.cuh @@ -0,0 +1,105 @@ +#pragma once +#include "common.cuh" + +// Readers of a conv-history concat that a tail-only materialization can serve: the rollback-slot views of the +// last H columns, their CONTs, and copies of either. Anything else reading the concat needs the whole tensor. +static inline int64_t ggml_cuda_conv_tail_view_col(const ggml_tensor * t, const ggml_tensor * cc, int64_t H) { + if (t->op != GGML_OP_VIEW || t->view_src != cc || t->type != GGML_TYPE_F32 || + t->ne[0] != H || t->ne[1] != cc->ne[1] || t->ne[2] != cc->ne[2] || t->ne[3] != 1 || + t->nb[1] != cc->nb[1] || t->nb[2] != cc->nb[2] || t->view_offs % sizeof(float) != 0) { + return -1; + } + const int64_t col = (int64_t) (t->view_offs / sizeof(float)); + return (col >= 0 && col + H <= cc->ne[0]) ? col : -1; +} + +static inline bool ggml_cuda_conv_tail_reader(const ggml_tensor * t, const ggml_tensor * cc, int64_t H, int64_t & tail_from) { + if (t->op == GGML_OP_VIEW) { + const int64_t col = ggml_cuda_conv_tail_view_col(t, cc, H); + if (col < 0) { + return false; + } + tail_from = std::min(tail_from, col); + return true; + } + if (t->op == GGML_OP_CONT) { + return t->src[0] && ggml_cuda_conv_tail_view_col(t->src[0], cc, H) >= 0; + } + if (t->op == GGML_OP_CPY) { + const ggml_tensor * s = t->src[0]; + if (!s) { + return false; + } + if (s->op == GGML_OP_CONT) { + s = s->src[0]; + } + return s && ggml_cuda_conv_tail_view_col(s, cc, H) >= 0; + } + return false; +} + +static inline bool ggml_cuda_reads(const ggml_tensor * t, const ggml_tensor * root) { + for (int s = 0; s < GGML_MAX_SRC && t->src[s]; ++s) { + if (t->src[s] == root || t->src[s]->view_src == root) { + return true; + } + } + return false; +} + +// the concat match is dispatched as "this node plus the next": the next node must be a no-op the loop skips anyway +static inline bool ggml_cuda_conv_next_is_noop(const ggml_cgraph * cgraph, int i) { + if (i + 1 >= cgraph->n_nodes) { + return false; + } + const ggml_tensor * t = cgraph->nodes[i + 1]; + return ggml_is_empty(t) || t->op == GGML_OP_RESHAPE || t->op == GGML_OP_TRANSPOSE || + t->op == GGML_OP_VIEW || t->op == GGML_OP_PERMUTE || t->op == GGML_OP_NONE; +} + +// GGML_CONV_FUSE_DEBUG=1 prints why a concat did not match, with the node index the decision was made at +static inline bool ggml_cuda_conv_reject(const char * which, const char * why, int node) { + static const bool dbg = getenv("GGML_CONV_FUSE_DEBUG") != nullptr; + if (dbg) { + fprintf(stderr, "%s: no fusion, %s (node %d)\n", which, why, node); + } + return false; +} + +// nodes after `from` (inclusive) that read t, other than t itself and `except`: a scan, so it works on graph +// copies and views whose use-count tables are not populated +static inline int ggml_cuda_conv_readers(const ggml_cgraph * cgraph, int from, const ggml_tensor * t, const ggml_tensor * except) { + int n = 0; + for (int k = from; k < cgraph->n_nodes; ++k) { + const ggml_tensor * r = cgraph->nodes[k]; + if (r == t || r == except || r->op == GGML_OP_RESHAPE || r->op == GGML_OP_VIEW || r->op == GGML_OP_TRANSPOSE || r->op == GGML_OP_PERMUTE) { + continue; // a view of t is not a consumer: whoever reads the view is counted through view_src + } + if (ggml_cuda_reads(r, t)) { + ++n; + } + } + return n; +} + +// the whole concat, either directly or through a view that covers all of it +static inline bool ggml_cuda_conv_whole(const ggml_tensor * s, const ggml_tensor * cc) { + return s == cc || (s && s->view_src == cc && s->view_offs == 0 && ggml_nelements(s) == ggml_nelements(cc) && ggml_is_contiguous(s)); +} + +// concat src[1] is TRANSPOSE(view of x): x must be the whole contiguous storage of a [C, T] activation, whatever +// shape its producer gave it (qwen4exp's grouped norm leaves it [n_embd, hc, T]); C and T come from the transpose +static inline const ggml_tensor * ggml_cuda_conv_activation(const ggml_tensor * tr, int64_t & C, int64_t & T) { + if (!tr || tr->op != GGML_OP_TRANSPOSE || tr->type != GGML_TYPE_F32 || !tr->view_src || tr->view_offs != 0 || + tr->ne[2] != 1 || tr->ne[3] != 1) { + return nullptr; + } + T = tr->ne[0]; + C = tr->ne[1]; + const ggml_tensor * x = tr->view_src; + if (x->type != GGML_TYPE_F32 || !ggml_is_contiguous(x) || ggml_nelements(x) != C * T || + tr->nb[1] != sizeof(float) || tr->nb[0] != (size_t) C * sizeof(float)) { + return nullptr; + } + return x; +} diff --git a/ggml/src/ggml-cuda/convert.cu b/ggml/src/ggml-cuda/convert.cu index 360c614a4403..47cab5f980d4 100644 --- a/ggml/src/ggml-cuda/convert.cu +++ b/ggml/src/ggml-cuda/convert.cu @@ -450,8 +450,40 @@ static void convert_unary_cuda(const void * vx, dst_t * y, (vx, y, ne00, ne01, ne0203, ne02_fdv, s01, s02, s03); } +static __global__ void convert_f32_f16_cont(const float * x, half * y, int64_t k) { + const int64_t i = 2 * ((int64_t) blockDim.x * blockIdx.x + threadIdx.x); + if (i + 1 < k) { + const float2 value = ((const float2 *) x)[i / 2]; + ((half2 *) y)[i / 2] = make_half2(value.x, value.y); + } else if (i < k) { + y[i] = x[i]; + } +} + +static __global__ void convert_f16_f32_cont(const half * x, float * y, int64_t k) { + const int64_t i = 2 * ((int64_t) blockDim.x * blockIdx.x + threadIdx.x); + if (i + 1 < k) { + ((float2 *) y)[i / 2] = __half22float2(((const half2 *) x)[i / 2]); + } else if (i < k) { + y[i] = x[i]; + } +} + template static void convert_unary_cont_cuda(const void * vx, dst_t * y, const int64_t k, cudaStream_t stream) { + const int nblocks = (k + 2*CUDA_DEQUANTIZE_BLOCK_SIZE - 1) / (2*CUDA_DEQUANTIZE_BLOCK_SIZE); + if constexpr (std::is_same::value && std::is_same::value) { + if ((uintptr_t) vx % alignof(float2) == 0 && (uintptr_t) y % alignof(half2) == 0) { + convert_f32_f16_cont<<>>((const float *) vx, (half *) y, k); + return; + } + } + if constexpr (std::is_same::value && std::is_same::value) { + if ((uintptr_t) vx % alignof(half2) == 0 && (uintptr_t) y % alignof(float2) == 0) { + convert_f16_f32_cont<<>>((const half *) vx, (float *) y, k); + return; + } + } convert_unary_cuda(vx, y, k, 1, 1, 1, k, k, k, stream); } diff --git a/ggml/src/ggml-cuda/cumsum.cu b/ggml/src/ggml-cuda/cumsum.cu index def9c32955f2..51cfb84dab17 100644 --- a/ggml/src/ggml-cuda/cumsum.cu +++ b/ggml/src/ggml-cuda/cumsum.cu @@ -5,7 +5,12 @@ #include "ggml.h" #ifdef GGML_CUDA_USE_CUB -# include +# ifdef GGML_CUDA_CUB_IS_HIPCUB +# include +namespace cub = hipcub; +# else +# include +# endif // GGML_CUDA_CUB_IS_HIPCUB #endif // GGML_CUDA_USE_CUB template diff --git a/ggml/src/ggml-cuda/dequantize.cuh b/ggml/src/ggml-cuda/dequantize.cuh index 651524be7b3e..fd9b6d6e01bf 100644 --- a/ggml/src/ggml-cuda/dequantize.cuh +++ b/ggml/src/ggml-cuda/dequantize.cuh @@ -124,8 +124,8 @@ static __device__ __forceinline__ void dequantize_q8_0(const void * vx, const in // Each call dequantizes one super-block of QK_K values into y using the // thread layout of the caller: 32 threads for q4_K, 64 threads otherwise. -template -static __device__ __forceinline__ void dequantize_q2_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) { +template +static __device__ __forceinline__ void dequantize_q2_K(const void * vx, const int64_t ib, dst_ptr_t yy, const int tid) { const block_q2_K * x = (const block_q2_K *) vx; const int64_t n = tid/32; @@ -133,7 +133,7 @@ static __device__ __forceinline__ void dequantize_q2_K(const void * vx, const in const int64_t is = 8*n + l/16; const uint8_t q = x[ib].qs[32*n + l]; - dst_t * y = yy + 128*n; + auto y = yy + 128*n; float dall = __low2half(x[ib].dm); float dmin = __high2half(x[ib].dm); @@ -143,8 +143,8 @@ static __device__ __forceinline__ void dequantize_q2_K(const void * vx, const in y[l+96] = ggml_cuda_cast(dall * (x[ib].scales[is+6] & 0xF) * ((q >> 6) & 3) - dmin * (x[ib].scales[is+6] >> 4)); } -template -static __device__ __forceinline__ void dequantize_q3_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) { +template +static __device__ __forceinline__ void dequantize_q3_K(const void * vx, const int64_t ib, dst_ptr_t yy, const int tid) { const block_q3_K * x = (const block_q3_K *) vx; const int64_t r = tid/4; @@ -165,7 +165,7 @@ static __device__ __forceinline__ void dequantize_q3_K(const void * vx, const in float d_all = x[ib].d; float dl = d_all * (us - 32); - dst_t * y = yy + 128*n + 32*j; + auto y = yy + 128*n + 32*j; const uint8_t * q = x[ib].qs + 32*n; const uint8_t * hm = x[ib].hmask; @@ -183,8 +183,8 @@ static inline __device__ void get_scale_min_k4(int j, const uint8_t * q, uint8_t } } -template -static __device__ __forceinline__ void dequantize_q4_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) { +template +static __device__ __forceinline__ void dequantize_q4_K(const void * vx, const int64_t ib, dst_ptr_t yy, const int tid) { const block_q4_K * x = (const block_q4_K *) vx; // assume 32 threads @@ -193,7 +193,7 @@ static __device__ __forceinline__ void dequantize_q4_K(const void * vx, const in const int64_t is = 2*il; const int64_t n = 4; - dst_t * y = yy + 64*il + n*ir; + auto y = yy + 64*il + n*ir; const float dall = __low2half(x[ib].dm); const float dmin = __high2half(x[ib].dm); @@ -211,8 +211,8 @@ static __device__ __forceinline__ void dequantize_q4_K(const void * vx, const in } } -template -static __device__ __forceinline__ void dequantize_q5_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) { +template +static __device__ __forceinline__ void dequantize_q5_K(const void * vx, const int64_t ib, dst_ptr_t yy, const int tid) { const block_q5_K * x = (const block_q5_K *) vx; // assume 64 threads - this is very slightly better than the one below @@ -220,7 +220,7 @@ static __device__ __forceinline__ void dequantize_q5_K(const void * vx, const in const int64_t ir = tid%16; // ir is in 0...15 const int64_t is = 2*il; // is is in 0...6 - dst_t * y = yy + 64*il + 2*ir; + auto y = yy + 64*il + 2*ir; const float dall = __low2half(x[ib].dm); const float dmin = __high2half(x[ib].dm); @@ -242,8 +242,8 @@ static __device__ __forceinline__ void dequantize_q5_K(const void * vx, const in y[33] = ggml_cuda_cast(d2 * ((ql[ 1] >> 4) + (qh[ 1] & hm ? 16 : 0)) - m2); } -template -static __device__ __forceinline__ void dequantize_q6_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) { +template +static __device__ __forceinline__ void dequantize_q6_K(const void * vx, const int64_t ib, dst_ptr_t yy, const int tid) { const block_q6_K * x = (const block_q6_K *) vx; // assume 64 threads - this is very slightly better than the one below @@ -251,7 +251,7 @@ static __device__ __forceinline__ void dequantize_q6_K(const void * vx, const in const int64_t il = tid - 32*ip; // 0...32 const int64_t is = 8*ip + il/16; - dst_t * y = yy + 128*ip + il; + auto y = yy + 128*ip + il; const float d = x[ib].d; @@ -270,14 +270,14 @@ static __device__ __forceinline__ void dequantize_q6_K(const void * vx, const in // Each call dequantizes one super-block of QK_K values into y with 32 // threads; iq4_nl packs QK_K/QK4_NL sub-blocks per super-block. -template -static __device__ __forceinline__ void dequantize_iq2_xxs(const void * vx, const int64_t ibs, dst_t * yy, const int tid) { +template +static __device__ __forceinline__ void dequantize_iq2_xxs(const void * vx, const int64_t ibs, dst_ptr_t yy, const int tid) { const block_iq2_xxs * x = (const block_iq2_xxs *) vx; const int64_t il = tid/8; // 0...3 const int64_t ib = tid%8; // 0...7 - dst_t * y = yy + 32*ib + 8*il; + auto y = yy + 32*ib + 8*il; const uint16_t * q2 = x[ibs].qs + 4*ib; const uint8_t * aux8 = (const uint8_t *)q2; const uint8_t * grid = (const uint8_t *)(iq2xxs_grid + aux8[il]); @@ -289,14 +289,14 @@ static __device__ __forceinline__ void dequantize_iq2_xxs(const void * vx, const } } -template -static __device__ __forceinline__ void dequantize_iq2_xs(const void * vx, const int64_t ibs, dst_t * yy, const int tid) { +template +static __device__ __forceinline__ void dequantize_iq2_xs(const void * vx, const int64_t ibs, dst_ptr_t yy, const int tid) { const block_iq2_xs * x = (const block_iq2_xs *) vx; const int64_t il = tid/8; // 0...3 const int64_t ib = tid%8; // 0...7 - dst_t * y = yy + 32*ib + 8*il; + auto y = yy + 32*ib + 8*il; const uint16_t * q2 = x[ibs].qs + 4*ib; const uint8_t * grid = (const uint8_t *)(iq2xs_grid + (q2[il] & 511)); const float d = (float)x[ibs].d * (0.5f + ((x[ibs].scales[ib] >> 4*(il/2)) & 0xf)) * 0.25f; @@ -306,14 +306,14 @@ static __device__ __forceinline__ void dequantize_iq2_xs(const void * vx, const } } -template -static __device__ __forceinline__ void dequantize_iq2_s(const void * vx, const int64_t ibs, dst_t * yy, const int tid) { +template +static __device__ __forceinline__ void dequantize_iq2_s(const void * vx, const int64_t ibs, dst_ptr_t yy, const int tid) { const block_iq2_s * x = (const block_iq2_s *) vx; const int64_t il = tid/8; // 0...3 const int64_t ib = tid%8; // 0...7 - dst_t * y = yy + 32*ib + 8*il; + auto y = yy + 32*ib + 8*il; const uint8_t * grid = (const uint8_t *)(iq2s_grid + (x[ibs].qs[4*ib+il] | ((x[ibs].qh[ib] << (8-2*il)) & 0x300))); const float d = (float)x[ibs].d * (0.5f + ((x[ibs].scales[ib] >> 4*(il/2)) & 0xf)) * 0.25f; const uint8_t signs = x[ibs].qs[QK_K/8+4*ib+il]; @@ -322,14 +322,14 @@ static __device__ __forceinline__ void dequantize_iq2_s(const void * vx, const i } } -template -static __device__ __forceinline__ void dequantize_iq3_xxs(const void * vx, const int64_t ibs, dst_t * yy, const int tid) { +template +static __device__ __forceinline__ void dequantize_iq3_xxs(const void * vx, const int64_t ibs, dst_ptr_t yy, const int tid) { const block_iq3_xxs * x = (const block_iq3_xxs *) vx; const int64_t il = tid/8; // 0...3 const int64_t ib = tid%8; // 0...7 - dst_t * y = yy + 32*ib + 8*il; + auto y = yy + 32*ib + 8*il; const uint8_t * q3 = x[ibs].qs + 8*ib; const uint16_t * gas = (const uint16_t *)(x[ibs].qs + QK_K/4) + 2*ib; const uint8_t * grid1 = (const uint8_t *)(iq3xxs_grid + q3[2*il+0]); @@ -343,14 +343,14 @@ static __device__ __forceinline__ void dequantize_iq3_xxs(const void * vx, const } } -template -static __device__ __forceinline__ void dequantize_iq3_s(const void * vx, const int64_t ibs, dst_t * yy, const int tid) { +template +static __device__ __forceinline__ void dequantize_iq3_s(const void * vx, const int64_t ibs, dst_ptr_t yy, const int tid) { const block_iq3_s * x = (const block_iq3_s *) vx; const int64_t il = tid/8; // 0...3 const int64_t ib = tid%8; // 0...7 - dst_t * y = yy + 32*ib + 8*il; + auto y = yy + 32*ib + 8*il; const uint8_t * qs = x[ibs].qs + 8*ib; const uint8_t * grid1 = (const uint8_t *)(iq3s_grid + (qs[2*il+0] | ((x[ibs].qh[ib] << (8-2*il)) & 256))); const uint8_t * grid2 = (const uint8_t *)(iq3s_grid + (qs[2*il+1] | ((x[ibs].qh[ib] << (7-2*il)) & 256))); @@ -362,14 +362,14 @@ static __device__ __forceinline__ void dequantize_iq3_s(const void * vx, const i } } -template -static __device__ __forceinline__ void dequantize_iq1_s(const void * vx, const int64_t ibs, dst_t * yy, const int tid) { +template +static __device__ __forceinline__ void dequantize_iq1_s(const void * vx, const int64_t ibs, dst_ptr_t yy, const int tid) { const block_iq1_s * x = (const block_iq1_s *) vx; const int64_t il = tid/8; // 0...3 const int64_t ib = tid%8; // 0...7 - dst_t * y = yy + 32*ib + 8*il; + auto y = yy + 32*ib + 8*il; const float delta = x[ibs].qh[ib] & 0x8000 ? -1 - IQ1S_DELTA : -1 + IQ1S_DELTA; const float d = (float)x[ibs].d * (2*((x[ibs].qh[ib] >> 12) & 7) + 1); uint32_t grid32[2]; const int8_t * q = (const int8_t *)grid32; @@ -381,14 +381,14 @@ static __device__ __forceinline__ void dequantize_iq1_s(const void * vx, const i } } -template -static __device__ __forceinline__ void dequantize_iq1_m(const void * vx, const int64_t ibs, dst_t * yy, const int tid) { +template +static __device__ __forceinline__ void dequantize_iq1_m(const void * vx, const int64_t ibs, dst_ptr_t yy, const int tid) { const block_iq1_m * x = (const block_iq1_m *) vx; const int64_t il = tid/8; // 0...3 const int64_t ib = tid%8; // 0...7 - dst_t * y = yy + 32*ib + 8*il; + auto y = yy + 32*ib + 8*il; const uint16_t * sc = (const uint16_t *)x[ibs].scales; iq1m_scale_t scale; scale.u16 = (sc[0] >> 12) | ((sc[1] >> 8) & 0x00f0) | ((sc[2] >> 4) & 0x0f00) | (sc[3] & 0xf000); @@ -404,14 +404,14 @@ static __device__ __forceinline__ void dequantize_iq1_m(const void * vx, const i } } -template -static __device__ __forceinline__ void dequantize_iq4_nl(const void * vx, const int64_t ibs, dst_t * yy, const int tid) { +template +static __device__ __forceinline__ void dequantize_iq4_nl(const void * vx, const int64_t ibs, dst_ptr_t yy, const int tid) { const block_iq4_nl * x = (const block_iq4_nl *) vx + ibs*(QK_K/QK4_NL); const int64_t il = tid/8; // 0...3 const int64_t ib = tid%8; // 0...7 - dst_t * y = yy + 32*ib + 4*il; + auto y = yy + 32*ib + 4*il; const uint8_t * q4 = x[ib].qs + 4*il; const float d = (float)x[ib].d; for (int j = 0; j < 4; ++j) { @@ -420,13 +420,13 @@ static __device__ __forceinline__ void dequantize_iq4_nl(const void * vx, const } } -template -static __device__ __forceinline__ void dequantize_iq4_xs(const void * vx, const int64_t ibs, dst_t * yy, const int tid) { +template +static __device__ __forceinline__ void dequantize_iq4_xs(const void * vx, const int64_t ibs, dst_ptr_t yy, const int tid) { const block_iq4_xs * x = (const block_iq4_xs *)vx; const int64_t il = tid/8; // 0...3 const int64_t ib = tid%8; // 0...7 - dst_t * y = yy + 32*ib + 4*il; + auto y = yy + 32*ib + 4*il; const uint8_t * q4 = x[ibs].qs + 16*ib + 4*il; const float d = (float)x[ibs].d * ((((x[ibs].scales_l[ib/2] >> 4*(ib%2)) & 0xf) | (((x[ibs].scales_h >> 2*ib) & 3) << 4)) - 32); for (int j = 0; j < 4; ++j) { @@ -435,14 +435,14 @@ static __device__ __forceinline__ void dequantize_iq4_xs(const void * vx, const } } -template -static __device__ __forceinline__ void dequantize_mxfp4(const void * vx, const int64_t ibs, dst_t * yy, const int tid) { +template +static __device__ __forceinline__ void dequantize_mxfp4(const void * vx, const int64_t ibs, dst_ptr_t yy, const int tid) { const block_mxfp4 * x = (const block_mxfp4 *) vx + ibs*(QK_K/QK_MXFP4); const int64_t il = tid/8; // 0...3 const int64_t ib = tid%8; // 0...7 - dst_t * y = yy + 32*ib + 4*il; + auto y = yy + 32*ib + 4*il; const uint8_t * q4 = x[ib].qs + 4*il; const float d = ggml_cuda_e8m0_to_fp32(x[ib].e); for (int j = 0; j < 4; ++j) { @@ -450,3 +450,78 @@ static __device__ __forceinline__ void dequantize_mxfp4(const void * vx, const i y[j+16] = ggml_cuda_cast(d * kvalues_mxfp4[q4[j] >> 4]*0.5f); } } + +template +static __device__ __forceinline__ void dequantize_q2_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) { + dequantize_q2_K(vx, ib, yy, tid); +} + +template +static __device__ __forceinline__ void dequantize_q3_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) { + dequantize_q3_K(vx, ib, yy, tid); +} + +template +static __device__ __forceinline__ void dequantize_q4_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) { + dequantize_q4_K(vx, ib, yy, tid); +} + +template +static __device__ __forceinline__ void dequantize_q5_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) { + dequantize_q5_K(vx, ib, yy, tid); +} + +template +static __device__ __forceinline__ void dequantize_q6_K(const void * vx, const int64_t ib, dst_t * yy, const int tid) { + dequantize_q6_K(vx, ib, yy, tid); +} + +template +static __device__ __forceinline__ void dequantize_iq2_xxs(const void * vx, const int64_t ib, dst_t * yy, const int tid) { + dequantize_iq2_xxs(vx, ib, yy, tid); +} + +template +static __device__ __forceinline__ void dequantize_iq2_xs(const void * vx, const int64_t ib, dst_t * yy, const int tid) { + dequantize_iq2_xs(vx, ib, yy, tid); +} + +template +static __device__ __forceinline__ void dequantize_iq2_s(const void * vx, const int64_t ib, dst_t * yy, const int tid) { + dequantize_iq2_s(vx, ib, yy, tid); +} + +template +static __device__ __forceinline__ void dequantize_iq3_xxs(const void * vx, const int64_t ib, dst_t * yy, const int tid) { + dequantize_iq3_xxs(vx, ib, yy, tid); +} + +template +static __device__ __forceinline__ void dequantize_iq3_s(const void * vx, const int64_t ib, dst_t * yy, const int tid) { + dequantize_iq3_s(vx, ib, yy, tid); +} + +template +static __device__ __forceinline__ void dequantize_iq1_s(const void * vx, const int64_t ib, dst_t * yy, const int tid) { + dequantize_iq1_s(vx, ib, yy, tid); +} + +template +static __device__ __forceinline__ void dequantize_iq1_m(const void * vx, const int64_t ib, dst_t * yy, const int tid) { + dequantize_iq1_m(vx, ib, yy, tid); +} + +template +static __device__ __forceinline__ void dequantize_iq4_nl(const void * vx, const int64_t ib, dst_t * yy, const int tid) { + dequantize_iq4_nl(vx, ib, yy, tid); +} + +template +static __device__ __forceinline__ void dequantize_iq4_xs(const void * vx, const int64_t ib, dst_t * yy, const int tid) { + dequantize_iq4_xs(vx, ib, yy, tid); +} + +template +static __device__ __forceinline__ void dequantize_mxfp4(const void * vx, const int64_t ib, dst_t * yy, const int tid) { + dequantize_mxfp4(vx, ib, yy, tid); +} diff --git a/ggml/src/ggml-cuda/expand.cu b/ggml/src/ggml-cuda/expand.cu new file mode 100644 index 000000000000..ba008ca1f59f --- /dev/null +++ b/ggml/src/ggml-cuda/expand.cu @@ -0,0 +1,63 @@ +#include "common.cuh" +#include "expand.cuh" +#include +#include +#include + +struct expand_layout { + size_t block_row, score_row, cell_row, tail_row, output_row, cast_block_row, cast_cell_row, expanded_row; + int n_blocks; +}; + +__global__ __launch_bounds__(256) void qsa_expand_complete_blocks_512( + const char * blocks, const char * scores, const char * cells, const char * tail, + char * output, char * cast_blocks, char * cast_cells, char * expanded, expand_layout layout) { + const int query = blockIdx.x; + const int lane = threadIdx.x; + __shared__ int ordered[512]; + const int * selected = reinterpret_cast(blocks + query*layout.block_row); + for (int i = lane; i < 512; i += 256) { + ordered[i] = selected[i]; + reinterpret_cast(cast_blocks + query*layout.cast_block_row)[i] = float(selected[i]); + } + __syncthreads(); + for (int size = 2; size <= 512; size *= 2) { + for (int distance = size/2; distance; distance /= 2) { + for (int i = lane; i < 512; i += 256) { + const int other = i ^ distance; + if (other > i) { + const int a = ordered[i], b = ordered[other]; + if ((i & size) ? a < b : a > b) { ordered[i] = b; ordered[other] = a; } + } + } + __syncthreads(); + } + } + const float * score = reinterpret_cast(scores + query*layout.score_row); + const int * remainder = reinterpret_cast(tail + query*layout.tail_row); + int * out = reinterpret_cast(output + query*layout.output_row); + for (int i = lane; i < 2051; i += 256) { + if (i >= 2048) { out[i] = remainder[i-2048]; continue; } + const int block = ordered[i/4]; + if (block < 0 || block >= layout.n_blocks) { out[i] = -1; continue; } + const int cell = reinterpret_cast(cells + block*layout.cell_row)[i%4]; + const float valid = __fadd_rn(score[block], 1.0f) > 0.0f ? 1.0f : 0.0f; + const float value = float(cell); + reinterpret_cast(cast_cells + query*layout.cast_cell_row)[i] = value; + const float shifted = __fadd_rn(value, 1.0f); + const int result = int(__fadd_rn(__fmul_rn(shifted, valid), -1.0f)); + reinterpret_cast(expanded + query*layout.expanded_row)[i] = result; + out[i] = result; + } +} + +void ggml_cuda_op_qsa_expand(ggml_backend_cuda_context & ctx, const ggml_cuda_qsa_expand_args & a) { + expand_layout layout{a.blocks->nb[1], a.scores->nb[1], a.cells->nb[1], a.tail->nb[1], a.output->nb[1], a.cast_blocks->nb[1], a.cast_cells->nb[2], a.expanded->nb[2], int(a.scores->ne[0])}; + const ggml_cuda_kernel_launch_params launch(dim3(a.output->ne[1]), dim3(256), 0, ctx.stream()); + ggml_cuda_kernel_launch(qsa_expand_complete_blocks_512, launch, + (const char *)a.blocks->data, (const char *)a.scores->data, (const char *)a.cells->data, + (const char *)a.tail->data, (char *)a.output->data, (char *)a.cast_blocks->data, + (char *)a.cast_cells->data, (char *)a.expanded->data, layout); + CUDA_CHECK(cudaGetLastError()); +} + diff --git a/ggml/src/ggml-cuda/expand.cuh b/ggml/src/ggml-cuda/expand.cuh new file mode 100644 index 000000000000..37a1ef7e78c9 --- /dev/null +++ b/ggml/src/ggml-cuda/expand.cuh @@ -0,0 +1,14 @@ +#pragma once +#include "ggml.h" +struct ggml_backend_cuda_context; +struct ggml_cuda_qsa_expand_args { + const ggml_tensor * blocks; + const ggml_tensor * scores; + const ggml_tensor * cells; + const ggml_tensor * tail; + ggml_tensor * cast_blocks; + ggml_tensor * cast_cells; + ggml_tensor * expanded; + ggml_tensor * output; +}; +void ggml_cuda_op_qsa_expand(ggml_backend_cuda_context &, const ggml_cuda_qsa_expand_args &); diff --git a/ggml/src/ggml-cuda/fattn-common.cuh b/ggml/src/ggml-cuda/fattn-common.cuh index 48b631e60fb4..ebd6a65e40b5 100644 --- a/ggml/src/ggml-cuda/fattn-common.cuh +++ b/ggml/src/ggml-cuda/fattn-common.cuh @@ -1198,6 +1198,12 @@ void launch_fattn( } } + if constexpr (ncols2 == 3) { + if (GGML_CUDA_CC_IS_RDNA3_5(cc)) { + parallel_blocks = std::min(18, ntiles_KV); + } + } + blocks_num.x = ntiles_x; blocks_num.y = parallel_blocks; blocks_num.z = ntiles_z_gqa*K->ne[2]*Q->ne[3]; diff --git a/ggml/src/ggml-cuda/fattn-mma-f16.cuh b/ggml/src/ggml-cuda/fattn-mma-f16.cuh index a290655776aa..5bd872ee399c 100644 --- a/ggml/src/ggml-cuda/fattn-mma-f16.cuh +++ b/ggml/src/ggml-cuda/fattn-mma-f16.cuh @@ -158,8 +158,8 @@ static constexpr __host__ __device__ fattn_mma_config ggml_cuda_fattn_mma_get_co GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 8, 64, 2, 32, 128, 128, 128, 1, true); GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 16, 64, 2, 32, 128, 128, 128, 1, true); - GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 32, 256, 2, 64, 128, 128, 64, 1, true); - GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 64, 256, 2, 64, 128, 128, 64, 1, true); + GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 32, 128, 2, 64, 128, 128, 64, 1, true); + GGML_CUDA_FATTN_MMA_CONFIG_CASE(256, 256, 64, 256, 1, 32, 128, 64, 64, 1, false); GGML_CUDA_FATTN_MMA_CONFIG_CASE(320, 256, 32, 128, 2, 32, 160, 128, 128, 1, true); GGML_CUDA_FATTN_MMA_CONFIG_CASE(320, 256, 64, 128, 2, 32, 160, 128, 128, 1, true); @@ -1744,7 +1744,9 @@ static __device__ __forceinline__ void flash_attn_ext_f16_process_tile( } } } - if (np > 1) { + // Warps read columns of tile_Q that were written by other warps, so a barrier is needed before tile_Q is + // reused: by the next nbatch_combine batch (DV/2 > nbatch_combine) or by the Q load of the next tile. + if (np > 1 || DV/2 > nbatch_combine) { __syncthreads(); } } @@ -1828,7 +1830,7 @@ static __global__ void flash_attn_ext_f16( #endif // __CUDA_ARCH__ == GGML_CUDA_CC_TURING #if defined(AMD_WMMA_AVAILABLE) - if (ncols1*ncols2 < 16 || ncols2 == 1 || DKQ > 256) { + if (ncols1*ncols2 < 16 || ncols2 == 1 || DKQ > 256 || (DKQ > 128 && DKQ != 256)) { NO_DEVICE_CODE; return; } diff --git a/ggml/src/ggml-cuda/fattn-tile-rdna3-5.cu b/ggml/src/ggml-cuda/fattn-tile-rdna3-5.cu new file mode 100644 index 000000000000..586632da20e4 --- /dev/null +++ b/ggml/src/ggml-cuda/fattn-tile-rdna3-5.cu @@ -0,0 +1,12 @@ +#if defined(GGML_USE_HIP) + +#define GGML_CUDA_FATTN_VGPR192 +#include "fattn-tile.cuh" + +fattn_kernel_t ggml_cuda_fattn_tile_d256_ncols32_rdna3_5(const bool use_logit_softcap) { + return use_logit_softcap + ? flash_attn_tile<256, 256, 4, 8, true, false> + : flash_attn_tile<256, 256, 4, 8, false, false>; +} + +#endif // defined(GGML_USE_HIP) diff --git a/ggml/src/ggml-cuda/fattn-tile.cuh b/ggml/src/ggml-cuda/fattn-tile.cuh index 8981ab804ce0..8f79e27c08a6 100644 --- a/ggml/src/ggml-cuda/fattn-tile.cuh +++ b/ggml/src/ggml-cuda/fattn-tile.cuh @@ -1,6 +1,10 @@ #include "common.cuh" #include "fattn-common.cuh" +#if defined(GGML_USE_HIP) && defined(__clang__) +#pragma clang diagnostic ignored "-Wpass-failed" +#endif + // nbatch_fa == number of KQ rows to process per iteration // nbatch_K == number of K columns to load in parallel for KQ calculation @@ -213,6 +217,7 @@ static constexpr __host__ __device__ uint32_t ggml_cuda_fattn_tile_get_config_am GGML_CUDA_FATTN_TILE_CONFIG_CASE(192, 128, 32, 256, 2, 32, 64) GGML_CUDA_FATTN_TILE_CONFIG_CASE(256, 256, 2, 256, 2, 128, 64) + GGML_CUDA_FATTN_TILE_CONFIG_CASE(256, 256, 3, 96, 2, 32, 64) GGML_CUDA_FATTN_TILE_CONFIG_CASE(256, 256, 4, 256, 2, 64, 128) GGML_CUDA_FATTN_TILE_CONFIG_CASE(256, 256, 8, 256, 2, 64, 128) GGML_CUDA_FATTN_TILE_CONFIG_CASE(256, 256, 16, 256, 2, 32, 128) @@ -291,6 +296,7 @@ static constexpr __host__ __device__ uint32_t ggml_cuda_fattn_tile_get_config_am GGML_CUDA_FATTN_TILE_CONFIG_CASE(192, 128, 32, 256, 3, 64, 64) GGML_CUDA_FATTN_TILE_CONFIG_CASE(256, 256, 2, 64, 8, 32, 64) + GGML_CUDA_FATTN_TILE_CONFIG_CASE(256, 256, 3, 96, 8, 32, 64) GGML_CUDA_FATTN_TILE_CONFIG_CASE(256, 256, 4, 128, 6, 32, 256) GGML_CUDA_FATTN_TILE_CONFIG_CASE(256, 256, 8, 128, 6, 32, 256) GGML_CUDA_FATTN_TILE_CONFIG_CASE(256, 256, 16, 256, 5, 32, 256) @@ -479,14 +485,46 @@ static __device__ __forceinline__ void flash_attn_tile_load_tile( ggml_cuda_unroll<5>{}(load); } +template +static __device__ __forceinline__ void flash_attn_tile_load_tile_q8_0( + const char * const __restrict__ KV, half2 * const __restrict__ tile_KV, const int stride_KV, const int i_sup) { + static_assert(J % QK8_0 == 0, "bad J"); + + constexpr int nthreads = warp_size*nwarps; + constexpr int groups_per_row = J/4; + const int tid = threadIdx.y*warp_size + threadIdx.x; + +#pragma unroll + for (int ig = tid; ig < I*groups_per_row; ig += nthreads) { + const int i = ig/groups_per_row; + const int j = (ig % groups_per_row)*4; + + half2 v0 = make_half2(0.0f, 0.0f); + half2 v1 = make_half2(0.0f, 0.0f); + if (!oob_check || i < i_sup) { + const block_q8_0 * row = (const block_q8_0 *) (KV + int64_t(i)*stride_KV); + const block_q8_0 & block = row[j/QK8_0]; + int q; + ggml_cuda_memcpy_1(&q, block.qs + j % QK8_0); + const int8_t * q8 = (const int8_t *) &q; + const half2 d = __half2half2(block.d); + v0 = d*make_half2(q8[0], q8[1]); + v1 = d*make_half2(q8[2], q8[3]); + } + + tile_KV[i*(J/2 + J_padding) + j/2 + 0] = v0; + tile_KV[i*(J/2 + J_padding) + j/2 + 1] = v1; + } +} + // Function that performs a single iteration in for the KQ matrix multiplication: template + bool use_logit_softcap, bool oob_check, bool q8_0_KV, typename T_vec_dot> static __device__ __forceinline__ void flash_attn_tile_iter_KQ( T_vec_dot * const Q_tmp, - const half2 * const __restrict__ K_h2, + const char * const __restrict__ K_data, T_vec_dot * const KV_tmp, - const int stride_K2, + const int stride_K, const int k_VKQ_0, const int k_VKQ_sup, const int k_KQ_0, @@ -498,8 +536,13 @@ static __device__ __forceinline__ void flash_attn_tile_iter_KQ( constexpr int cpw = ncols > nwarps ? ncols/nwarps : 1; // Q columns per warp constexpr int np = nwarps > ncols ? nwarps/ncols : 1; // number of parallel warps per Q column - flash_attn_tile_load_tile - (K_h2 + int64_t(k_VKQ_0)*stride_K2 + k_KQ_0/2, KV_tmp, stride_K2, k_VKQ_sup); + if constexpr (q8_0_KV) { + flash_attn_tile_load_tile_q8_0 + (K_data + int64_t(k_VKQ_0)*stride_K + k_KQ_0/QK8_0*sizeof(block_q8_0), KV_tmp, stride_K, k_VKQ_sup); + } else { + flash_attn_tile_load_tile + ((const half2 *) (K_data + int64_t(k_VKQ_0)*stride_K + k_KQ_0*sizeof(half)), KV_tmp, stride_K/sizeof(half2), k_VKQ_sup); + } __syncthreads(); #ifdef FAST_FP16_AVAILABLE @@ -537,13 +580,29 @@ static __device__ __forceinline__ void flash_attn_tile_iter_KQ( #endif // FAST_FP16_AVAILABLE } +#if defined(RDNA3_5) + if constexpr (DKQ == 256 && ncols == 32) { #pragma unroll - for (int i_KQ_0 = 0; i_KQ_0 < nbatch_fa; i_KQ_0 += np*warp_size) { + for (int k = 0; k < cpy_ne; ++k) { #pragma unroll - for (int jc0 = 0; jc0 < cpw; ++jc0) { + for (int i_KQ_0 = 0; i_KQ_0 < nbatch_fa; i_KQ_0 += np*warp_size) { #pragma unroll - for (int k = 0; k < cpy_ne; ++k) { - ggml_cuda_mad(KQ_acc[i_KQ_0/(np*warp_size)*cpw + jc0], K_k[i_KQ_0/(np*warp_size)][k], Q_k[jc0][k]); + for (int jc0 = 0; jc0 < cpw; ++jc0) { + ggml_cuda_mad(KQ_acc[i_KQ_0/(np*warp_size)*cpw + jc0], K_k[i_KQ_0/(np*warp_size)][k], Q_k[jc0][k]); + } + } + } + } else +#endif // defined(RDNA3_5) + { +#pragma unroll + for (int i_KQ_0 = 0; i_KQ_0 < nbatch_fa; i_KQ_0 += np*warp_size) { +#pragma unroll + for (int jc0 = 0; jc0 < cpw; ++jc0) { +#pragma unroll + for (int k = 0; k < cpy_ne; ++k) { + ggml_cuda_mad(KQ_acc[i_KQ_0/(np*warp_size)*cpw + jc0], K_k[i_KQ_0/(np*warp_size)][k], Q_k[jc0][k]); + } } } } @@ -556,19 +615,19 @@ static __device__ __forceinline__ void flash_attn_tile_iter_KQ( // Function that performs a single iteration of the main loop over up to nbatch_fa tokens. template + bool use_logit_softcap, bool oob_check, bool q8_0_KV, typename T_vec_dot, typename T_KQ, typename T_acc> static __device__ __forceinline__ void flash_attn_tile_iter( T_vec_dot * const Q_tmp, - const half2 * const __restrict__ K_h2, - const half2 * const __restrict__ V_h2, + const char * const __restrict__ K_data, + const char * const __restrict__ V_data, const half * const __restrict__ mask, const uint3 ne01, const float logit_softcap, const float slope, T_KQ * const KQ, T_vec_dot * const KV_tmp, - const int stride_K2, - const int stride_V2, + const int stride_K, + const int stride_V, const int stride_mask, float * const KQ_max, float * const KQ_sum, @@ -601,19 +660,32 @@ static __device__ __forceinline__ void flash_attn_tile_iter( KQ_max_new[jc0] = KQ_max[jc0]; } +#if defined(RDNA3_5) + constexpr bool common_mask = DKQ == 256 && DV == 256 && ncols1 == 4 && ncols2 == 8; + const int j_mask = common_mask ? fastmodulo(col_Q_0 + ((threadIdx.y / np)*cpw)/ncols2, ne01) : 0; + float mask_value[common_mask ? nbatch_fa/(np*warp_size) : 1]; +#pragma unroll + for (int i0 = 0; i0 < nbatch_fa; i0 += np*warp_size) { + mask_value[i0/(np*warp_size)] = common_mask ? slope*__half2float(mask[j_mask*stride_mask + k_VKQ_0 + i0 + (threadIdx.y % np)*warp_size + threadIdx.x]) : 0.0f; + } +#else + constexpr bool common_mask = false; + const float mask_value[1] = {0.0f}; +#endif + float KQ_acc[nbatch_fa/(np*warp_size) * cpw] = {0.0f}; // Accumulators for KQ matrix multiplication. // KQ = K @ Q matrix multiplication: constexpr int nbatch_K_last = DKQ % nbatch_K; #pragma unroll for (int k_KQ_0 = 0; k_KQ_0 < DKQ - nbatch_K_last; k_KQ_0 += nbatch_K) { - flash_attn_tile_iter_KQ( - Q_tmp, K_h2, KV_tmp, stride_K2, k_VKQ_0, k_VKQ_sup, k_KQ_0, KQ_acc); + flash_attn_tile_iter_KQ( + Q_tmp, K_data, KV_tmp, stride_K, k_VKQ_0, k_VKQ_sup, k_KQ_0, KQ_acc); } - if (nbatch_K_last > 0) { + if constexpr (nbatch_K_last > 0) { constexpr int k_KQ_0 = DKQ - nbatch_K_last; - flash_attn_tile_iter_KQ( - Q_tmp, K_h2, KV_tmp, stride_K2, k_VKQ_0, k_VKQ_sup, k_KQ_0, KQ_acc); + flash_attn_tile_iter_KQ( + Q_tmp, K_data, KV_tmp, stride_K, k_VKQ_0, k_VKQ_sup, k_KQ_0, KQ_acc); } // Apply logit softcap + mask, update KQ_max: @@ -636,8 +708,8 @@ static __device__ __forceinline__ void flash_attn_tile_iter( } if (!oob_check || i_KQ < k_VKQ_sup) { - KQ_acc[(i_KQ_0/(np*warp_size))*cpw + jc0] += (ncols2 > 1 || mask) ? - slope*__half2float(mask[j*stride_mask + k_VKQ_0 + i_KQ]) : 0.0f; + KQ_acc[(i_KQ_0/(np*warp_size))*cpw + jc0] += common_mask ? mask_value[i_KQ_0/(np*warp_size)] : + (ncols2 > 1 || mask) ? slope*__half2float(mask[j*stride_mask + k_VKQ_0 + i_KQ]) : 0.0f; KQ_max_new[jc0] = fmaxf(KQ_max_new[jc0], KQ_acc[(i_KQ_0/(np*warp_size))*cpw + jc0] + FATTN_KQ_MAX_OFFSET); } @@ -718,8 +790,13 @@ static __device__ __forceinline__ void flash_attn_tile_iter( static_assert(nbatch_V % np == 0, "bad nbatch_V"); #pragma unroll for (int k0 = 0; k0 < nbatch_fa; k0 += nbatch_V) { - flash_attn_tile_load_tile - (V_h2 + int64_t(k_VKQ_0 + k0)*stride_V2, KV_tmp, stride_V2, k_VKQ_sup - k0); + if constexpr (q8_0_KV) { + flash_attn_tile_load_tile_q8_0 + (V_data + int64_t(k_VKQ_0 + k0)*stride_V, KV_tmp, stride_V, k_VKQ_sup - k0); + } else { + flash_attn_tile_load_tile + ((const half2 *) (V_data + int64_t(k_VKQ_0 + k0)*stride_V), KV_tmp, stride_V/sizeof(half2), k_VKQ_sup - k0); + } __syncthreads(); #ifdef FAST_FP16_AVAILABLE @@ -788,8 +865,11 @@ static __device__ __forceinline__ void flash_attn_tile_iter( } } -template // D == head size +template // D == head size __launch_bounds__(ggml_cuda_fattn_tile_get_nthreads(DKQ, DV, ncols1*ncols2), ggml_cuda_fattn_tile_get_occupancy(DKQ, DV, ncols1*ncols2)) +#if defined(RDNA3_5) && defined(GGML_CUDA_FATTN_VGPR192) +__attribute__((amdgpu_num_vgpr(192))) +#endif // defined(RDNA3_5) && defined(GGML_CUDA_FATTN_VGPR192) static __global__ void flash_attn_tile( const char * Q_ptr, const char * K_ptr, @@ -813,6 +893,13 @@ static __global__ void flash_attn_tile( const int32_t ne31, const int32_t ne32, const int32_t ne33, const int32_t nb31, const int32_t nb32, const int64_t nb33) { #ifdef FLASH_ATTN_AVAILABLE +#if defined(GGML_USE_HIP) && !defined(RDNA3_5) + if constexpr (q8_0_KV) { + NO_DEVICE_CODE; + return; + } +#endif // defined(GGML_USE_HIP) && !defined(RDNA3_5) + const char * GGML_CUDA_RESTRICT Q = Q_ptr; const char * GGML_CUDA_RESTRICT K = K_ptr; const char * GGML_CUDA_RESTRICT V = V_ptr; @@ -853,14 +940,12 @@ static __global__ void flash_attn_tile( const int sequence = blockIdx.z / (ne02/ncols2); const int head0 = blockIdx.z*ncols2 - sequence*ne02; // == blockIdx.z % (ne02/ncols2) const int gqa_ratio = ne02 / ne12; // With grouped query attention there are > 1 Q matrices per K, V matrix. - const float * Q_f = (const float *) (Q + nb03*sequence + nb02* head0); - const half2 * K_h2 = (const half2 *) (K + nb13*sequence + nb12*(head0 / gqa_ratio)); - const half2 * V_h2 = (const half2 *) (V + nb23*sequence + nb22*(head0 / gqa_ratio)); // K and V have same shape + const float * Q_f = (const float *) (Q + nb03*sequence + nb02* head0); + const char * K_data = K + nb13*sequence + nb12*(head0 / gqa_ratio); + const char * V_data = V + nb23*sequence + nb22*(head0 / gqa_ratio); const half * maskh = mask ? (const half *) (mask + nb33*(sequence % ne33)) : nullptr; - const int stride_K2 = nb11 / sizeof(half2); - const int stride_V2 = nb21 / sizeof(half2); const int stride_mask = nb31 / sizeof(half); const float slope = ncols2 == 1 ? get_alibi_slope(max_bias, head0, n_head_log2, m0, m1) : 1.0f; @@ -957,24 +1042,24 @@ static __global__ void flash_attn_tile( int k_VKQ_0 = blockIdx.y*nbatch_fa; while (k_VKQ_0 < k_VKQ_max - nbatch_fa) { constexpr bool oob_check = false; - flash_attn_tile_iter - (Q_tmp, K_h2, V_h2, maskh, ne01, logit_softcap, slope, KQ, KV_tmp, - stride_K2, stride_V2, stride_mask, KQ_max, KQ_sum, VKQ, k_VKQ_0, k_VKQ_max, col_Q_0); + flash_attn_tile_iter + (Q_tmp, K_data, V_data, maskh, ne01, logit_softcap, slope, KQ, KV_tmp, + nb11, nb21, stride_mask, KQ_max, KQ_sum, VKQ, k_VKQ_0, k_VKQ_max, col_Q_0); k_VKQ_0 += gridDim.y*nbatch_fa; } if (k_VKQ_0 < k_VKQ_max) { constexpr bool oob_check = true; - flash_attn_tile_iter - (Q_tmp, K_h2, V_h2, maskh, ne01, logit_softcap, slope, KQ, KV_tmp, - stride_K2, stride_V2, stride_mask, KQ_max, KQ_sum, VKQ, k_VKQ_0, k_VKQ_max, col_Q_0); + flash_attn_tile_iter + (Q_tmp, K_data, V_data, maskh, ne01, logit_softcap, slope, KQ, KV_tmp, + nb11, nb21, stride_mask, KQ_max, KQ_sum, VKQ, k_VKQ_0, k_VKQ_max, col_Q_0); } } else { // Branch without out-of-bounds checks. for (int k_VKQ_0 = blockIdx.y*nbatch_fa; k_VKQ_0 < k_VKQ_max; k_VKQ_0 += gridDim.y*nbatch_fa) { constexpr bool oob_check = false; - flash_attn_tile_iter - (Q_tmp, K_h2, V_h2, maskh, ne01, logit_softcap, slope, KQ, KV_tmp, - stride_K2, stride_V2, stride_mask, KQ_max, KQ_sum, VKQ, k_VKQ_0, k_VKQ_max, col_Q_0); + flash_attn_tile_iter + (Q_tmp, K_data, V_data, maskh, ne01, logit_softcap, slope, KQ, KV_tmp, + nb11, nb21, stride_mask, KQ_max, KQ_sum, VKQ, k_VKQ_0, k_VKQ_max, col_Q_0); } } @@ -1145,6 +1230,38 @@ static __global__ void flash_attn_tile( #endif // FLASH_ATTN_AVAILABLE } +#ifdef GGML_USE_HIP +fattn_kernel_t ggml_cuda_fattn_tile_d256_ncols32_rdna3_5(const bool use_logit_softcap); +#endif // GGML_USE_HIP + +template +static void launch_fattn_tile_case( + ggml_backend_cuda_context & ctx, ggml_tensor * dst, const int nwarps, const int nbatch_fa, const int warp_size) { + constexpr size_t nbytes_shared = 0; + bool use_q8_0_KV = false; + fattn_kernel_t fattn_kernel; +#ifdef GGML_USE_HIP + if constexpr (DKQ == DV && (DKQ == 64 || DKQ == 128 || DKQ == 256)) { + use_q8_0_KV = dst->src[0]->ne[1] == 1 && dst->src[1]->type == GGML_TYPE_Q8_0 && dst->src[2]->type == GGML_TYPE_Q8_0; + fattn_kernel = use_q8_0_KV + ? flash_attn_tile + : flash_attn_tile; + } else { + fattn_kernel = flash_attn_tile; + } + if constexpr (DKQ == 256 && DV == 256 && ncols1 == 4 && ncols2 == 8) { + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + if (GGML_CUDA_CC_IS_RDNA3_5(cc)) { + fattn_kernel = ggml_cuda_fattn_tile_d256_ncols32_rdna3_5(use_logit_softcap); + } + } +#else + fattn_kernel = flash_attn_tile; +#endif // GGML_USE_HIP + launch_fattn + (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, !use_q8_0_KV, !use_q8_0_KV, false, false, warp_size); +} + template static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * Q = dst->src[0]; @@ -1153,7 +1270,12 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm const int cc = ggml_cuda_info().devices[id].cc; const int warp_size = 32; - constexpr size_t nbytes_shared = 0; + if constexpr (ncols2 == 3) { + const int nwarps = ggml_cuda_fattn_tile_get_nthreads (DKQ, DV, 3, cc) / warp_size; + const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, 3, cc); + launch_fattn_tile_case(ctx, dst, nwarps, nbatch_fa, warp_size); + return; + } else { #ifdef GGML_USE_HIP if constexpr (DKQ <= 128) { @@ -1161,9 +1283,8 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm constexpr int cols_per_block = 64; const int nwarps = ggml_cuda_fattn_tile_get_nthreads (DKQ, DV, cols_per_block, cc) / warp_size; const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc); - fattn_kernel_t fattn_kernel = flash_attn_tile; - launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, false, warp_size); + launch_fattn_tile_case + (ctx, dst, nwarps, nbatch_fa, warp_size); return; } } @@ -1177,9 +1298,8 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm constexpr int cols_per_block = 32; const int nwarps = ggml_cuda_fattn_tile_get_nthreads (DKQ, DV, cols_per_block, cc) / warp_size; const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc); - fattn_kernel_t fattn_kernel = flash_attn_tile; - launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, false, warp_size); + launch_fattn_tile_case + (ctx, dst, nwarps, nbatch_fa, warp_size); return; } } @@ -1189,9 +1309,8 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm constexpr int cols_per_block = 16; const int nwarps = ggml_cuda_fattn_tile_get_nthreads (DKQ, DV, cols_per_block, cc) / warp_size; const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc); - fattn_kernel_t fattn_kernel = flash_attn_tile; - launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, false, warp_size); + launch_fattn_tile_case + (ctx, dst, nwarps, nbatch_fa, warp_size); return; } } @@ -1201,9 +1320,8 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm constexpr int cols_per_block = 8; const int nwarps = ggml_cuda_fattn_tile_get_nthreads (DKQ, DV, cols_per_block, cc) / warp_size; const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc); - fattn_kernel_t fattn_kernel = flash_attn_tile; - launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, false, warp_size); + launch_fattn_tile_case + (ctx, dst, nwarps, nbatch_fa, warp_size); return; } } @@ -1213,9 +1331,8 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm constexpr int cols_per_block = 4; const int nwarps = ggml_cuda_fattn_tile_get_nthreads (DKQ, DV, cols_per_block, cc) / warp_size; const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc); - fattn_kernel_t fattn_kernel = flash_attn_tile; - launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, false, warp_size); + launch_fattn_tile_case + (ctx, dst, nwarps, nbatch_fa, warp_size); return; } } @@ -1224,13 +1341,13 @@ static void launch_fattn_tile_switch_ncols1(ggml_backend_cuda_context & ctx, ggm constexpr int cols_per_block = 2; const int nwarps = ggml_cuda_fattn_tile_get_nthreads (DKQ, DV, cols_per_block, cc) / warp_size; const int nbatch_fa = ggml_cuda_fattn_tile_get_nbatch_fa(DKQ, DV, cols_per_block, cc); - fattn_kernel_t fattn_kernel = flash_attn_tile; - launch_fattn - (ctx, dst, fattn_kernel, nwarps, nbytes_shared, nbatch_fa, true, true, false, false, warp_size); + launch_fattn_tile_case + (ctx, dst, nwarps, nbatch_fa, warp_size); return; } - GGML_ABORT("fatal error"); + GGML_ABORT("fatal error"); + } } template @@ -1238,6 +1355,7 @@ static void launch_fattn_tile_switch_ncols2(ggml_backend_cuda_context & ctx, ggm const ggml_tensor * KQV = dst; const ggml_tensor * Q = dst->src[0]; const ggml_tensor * K = dst->src[1]; + [[maybe_unused]] const ggml_tensor * V = dst->src[2]; const ggml_tensor * mask = dst->src[3]; float max_bias = 0.0f; @@ -1296,6 +1414,15 @@ static void launch_fattn_tile_switch_ncols2(ggml_backend_cuda_context & ctx, ggm } if constexpr (DKQ <= 512 && DKQ != 320 && DKQ != 192) { + if constexpr (DKQ == 256 && DV == 256) { +#ifdef GGML_USE_HIP + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + if (use_gqa_opt && Q->ne[1] == 1 && gqa_ratio == 6 && K->type == GGML_TYPE_Q8_0 && V->type == GGML_TYPE_Q8_0 && GGML_CUDA_CC_IS_RDNA3_5(cc)) { + launch_fattn_tile_switch_ncols1(ctx, dst); + return; + } +#endif // GGML_USE_HIP + } if (use_gqa_opt && gqa_ratio % 8 == 0) { launch_fattn_tile_switch_ncols1(ctx, dst); return; diff --git a/ggml/src/ggml-cuda/fattn.cu b/ggml/src/ggml-cuda/fattn.cu index ceb4727931d4..e88684810f21 100644 --- a/ggml/src/ggml-cuda/fattn.cu +++ b/ggml/src/ggml-cuda/fattn.cu @@ -4,6 +4,10 @@ #include "fattn-tile.cuh" #include "fattn-vec.cuh" #include "fattn.cuh" +#if defined(GGML_USE_HIP) +#include "qsa-prefill.cuh" +#include "qsa-decode.cuh" +#endif #if !defined(GGML_USE_HIP) && !defined(GGML_USE_MUSA) __launch_bounds__(256, 1) @@ -514,6 +518,20 @@ static bool ggml_cuda_fattn_kv_type_supported(const ggml_type type) { } } +static bool ggml_cuda_fattn_tile_q8_0_KV_supported(const ggml_tensor * dst) { +#ifdef GGML_USE_HIP + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + const ggml_tensor * Q = dst->src[0]; + const ggml_tensor * K = dst->src[1]; + const ggml_tensor * V = dst->src[2]; + return GGML_CUDA_CC_IS_RDNA3_5(cc) && Q->ne[1] == 1 && K->type == GGML_TYPE_Q8_0 && V->type == GGML_TYPE_Q8_0 && Q->ne[0] == K->ne[0] && K->ne[0] == V->ne[0] && + (Q->ne[0] == 64 || Q->ne[0] == 128 || Q->ne[0] == 256); +#else + GGML_UNUSED(dst); + return false; +#endif // GGML_USE_HIP +} + static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const ggml_tensor * dst) { #ifndef FLASH_ATTN_AVAILABLE GGML_UNUSED(device); GGML_UNUSED(dst); @@ -641,6 +659,10 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const gqa_ratio_eff *= 2; } + if (ggml_cuda_fattn_tile_q8_0_KV_supported(dst) && gqa_opt_applies && gqa_ratio_eff >= 2) { + return BEST_FATTN_KERNEL_TILE; + } + if (volta_mma_available(cc) && Q->ne[0] != 40 && Q->ne[0] != 72) { if (can_use_vector_kernel && Q->ne[1] * gqa_ratio_eff <= 2) { return BEST_FATTN_KERNEL_VEC; @@ -665,11 +687,19 @@ static best_fattn_kernel ggml_cuda_get_best_fattn_kernel(const int device, const } // AMD WMMA is faster than the tile kernel if the wide tiles with high arithmetic intensity can be utilized. - if ((amd_wmma_available(cc) && gqa_opt_applies && Q->ne[0] <= 256) && Q->ne[0] != 40 && Q->ne[0] != 72 && + // Only head sizes that have WMMA device code in flash_attn_ext_f16 (fattn-mma-f16.cuh, AMD_WMMA_AVAILABLE guard: + // DKQ <= 128 or DKQ == 256) may be routed here; any other head size (e.g. 192) hits NO_DEVICE_CODE -> __trap(). + if ((amd_wmma_available(cc) && gqa_opt_applies && (Q->ne[0] <= 128 || Q->ne[0] == 256)) && Q->ne[0] != 40 && Q->ne[0] != 72 && Q->ne[1] * gqa_ratio_eff > (Q->ne[0] <= 128 ? 8 : 16)) { return BEST_FATTN_KERNEL_MMA_F16; } + // On RDNA3.5 the D=256 tile kernel is FMA-bound for prefill batches; the WMMA kernel with 64 columns is ~30-45% faster. + // For smaller batches (ncols 16/32 configs) the tile kernel is still faster. + if (GGML_CUDA_CC_IS_RDNA3_5(cc) && gqa_opt_applies && Q->ne[0] == 256 && V->ne[0] == 256 && Q->ne[1] * gqa_ratio_eff > 32) { + return BEST_FATTN_KERNEL_MMA_F16; + } + // If there are no tensor cores available, use the generic tile kernel: if (can_use_vector_kernel) { if (!ggml_is_quantized(K->type) && !ggml_is_quantized(V->type)) { @@ -704,6 +734,12 @@ size_t ggml_cuda_flash_attn_ext_get_alloc_size(int device, const ggml_tensor * d switch (kernel) { case BEST_FATTN_KERNEL_TILE: + if (ggml_cuda_fattn_tile_q8_0_KV_supported(dst)) { + break; + } + need_f16_K = true; + need_f16_V = true; + break; case BEST_FATTN_KERNEL_MMA_F16: need_f16_K = true; need_f16_V = true; @@ -725,6 +761,29 @@ size_t ggml_cuda_flash_attn_ext_get_alloc_size(int device, const ggml_tensor * d void ggml_cuda_flash_attn_ext(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ggml_cuda_set_device(ctx.device); +#if defined(GGML_USE_HIP) + if (ggml_cuda_flash_attn_ext_qsa_decode_supported(ctx, dst)) { + ggml_cuda_flash_attn_ext_qsa_decode(ctx, dst); + return; + } + if (ggml_cuda_flash_attn_ext_qsa_prefill_supported(ctx, dst)) { + ggml_cuda_flash_attn_ext_qsa_prefill(ctx, dst); + return; + } +#endif + // a selected-key op without a mask carries its visibility only in src[5], which the dense kernels ignore: + // reaching them would attend to the whole padded cache + if (dst->src[5] && !dst->src[3]) { + const ggml_tensor * q = dst->src[0], * k = dst->src[1], * v = dst->src[2], * ids = dst->src[5]; + GGML_ABORT("flash_attn_ext: maskless selected-key attention was not taken by a QSA kernel (%s): q [%lld,%lld,%lld,%lld] %s nb %zu/%zu, " + "k [%lld,%lld,%lld,%lld] %s nb %zu/%zu/%zu, v %s nb %zu/%zu, ids [%lld,%lld,%lld,%lld] nb %zu, src6 %d src7 %d p4 %d dst contiguous %d", + dst->name, (long long) q->ne[0], (long long) q->ne[1], (long long) q->ne[2], (long long) q->ne[3], ggml_type_name(q->type), q->nb[1], q->nb[2], + (long long) k->ne[0], (long long) k->ne[1], (long long) k->ne[2], (long long) k->ne[3], ggml_type_name(k->type), k->nb[0], k->nb[1], k->nb[2], + ggml_type_name(v->type), v->nb[1], v->nb[2], (long long) ids->ne[0], (long long) ids->ne[1], (long long) ids->ne[2], (long long) ids->ne[3], ids->nb[1], + dst->src[6] != nullptr, dst->src[7] != nullptr, ggml_get_op_params_i32(dst, 4), ggml_is_contiguous(dst)); + } + + switch (ggml_cuda_get_best_fattn_kernel(ggml_cuda_get_device(), dst)) { case BEST_FATTN_KERNEL_NONE: GGML_ABORT("fatal error"); diff --git a/ggml/src/ggml-cuda/gated_delta_net.cu b/ggml/src/ggml-cuda/gated_delta_net.cu index 1b431a724d72..ff2ac52745ea 100644 --- a/ggml/src/ggml-cuda/gated_delta_net.cu +++ b/ggml/src/ggml-cuda/gated_delta_net.cu @@ -1,8 +1,55 @@ #include "gated_delta_net.cuh" #include "ggml-cuda/common.cuh" +#include "unary.cuh" + + +// Warp reduction without the LDS crossbar: on RDNA the generic warp_reduce_sum compiles to 5 dependent +// ds_bpermute round trips (~100+ cycles each), which dominates the per-token latency of the recurrence. +// Same pairing/order of additions as warp_reduce_sum<32> (xor 16, 8, 4, 2, 1), so the result is bit-identical. +#if defined(GGML_USE_HIP) && (defined(RDNA3) || defined(RDNA4)) +template +static __device__ __forceinline__ float gdn_dpp_row_xmask(const float x) { + return __int_as_float(__builtin_amdgcn_update_dpp(0, __float_as_int(x), 0x160 | mask, 0xf, 0xf, true)); +} +static __device__ __forceinline__ float gdn_permlanex16_swap(const float x) { + return __int_as_float(__builtin_amdgcn_permlanex16(__float_as_int(x), __float_as_int(x), 0x76543210, 0xFEDCBA98, true, false)); +} +static __device__ __forceinline__ float gdn_warp_reduce_sum32(float x) { + x += gdn_permlanex16_swap(x); + x += gdn_dpp_row_xmask<8>(x); + x += gdn_dpp_row_xmask<4>(x); + x += gdn_dpp_row_xmask<2>(x); + x += gdn_dpp_row_xmask<1>(x); + return x; +} +#define GDN_DPP_REDUCE 1 +#endif // defined(GGML_USE_HIP) && (defined(RDNA3) || defined(RDNA4)) + +template +static __device__ __forceinline__ float gdn_warp_reduce_sum(const float x) { +#if defined(GDN_DPP_REDUCE) + if constexpr (width == 32) { + return gdn_warp_reduce_sum32(x); + } +#endif // defined(GDN_DPP_REDUCE) + return warp_reduce_sum(x); +} + +static constexpr int gated_delta_net_num_warps(int cc) { + return GGML_CUDA_CC_IS_RDNA3_5(cc) ? 32 : 4; +} + +static_assert(gated_delta_net_num_warps(GGML_CUDA_CC_RDNA3_5) == 32); +static_assert(gated_delta_net_num_warps(GGML_CUDA_CC_RDNA3) == 4); + +#if defined(RDNA3_5) +static constexpr int gdn_num_warps = gated_delta_net_num_warps(GGML_CUDA_CC_RDNA3_5); +#else +static constexpr int gdn_num_warps = gated_delta_net_num_warps(GGML_CUDA_CC_RDNA3); +#endif template -__global__ void __launch_bounds__((ggml_cuda_get_physical_warp_size() < S_v ? ggml_cuda_get_physical_warp_size() : S_v) * 4, 2) +__global__ void __launch_bounds__((ggml_cuda_get_physical_warp_size() < S_v ? ggml_cuda_get_physical_warp_size() : S_v) * gdn_num_warps, 2) gated_delta_net_cuda(const float * q, const float * k, const float * v, @@ -90,7 +137,7 @@ gated_delta_net_cuda(const float * q, for (int r = 0; r < rows_per_lane; r++) { kv_shard += s_shard[r] * k_reg[r]; } - float kv_col = warp_reduce_sum(kv_shard); + float kv_col = gdn_warp_reduce_sum(kv_shard); // delta[col] = (v[col] - g * kv[col]) * beta float delta_col = (v_t[col] - g_val * kv_col) * beta_val; @@ -104,7 +151,7 @@ gated_delta_net_cuda(const float * q, attn_partial += s_shard[r] * q_reg[r]; } - float attn_col = warp_reduce_sum(attn_partial); + float attn_col = gdn_warp_reduce_sum(attn_partial); if (lane == 0) { attn_data[col] = attn_col * scale; @@ -118,7 +165,7 @@ gated_delta_net_cuda(const float * q, kv_shard += expf(g_t[i]) * s_shard[r] * k_reg[r]; } - float kv_col = warp_reduce_sum(kv_shard); + float kv_col = gdn_warp_reduce_sum(kv_shard); // delta[col] = (v[col] - kv[col]) * beta float delta_col = (v_t[col] - kv_col) * beta_val; @@ -133,7 +180,7 @@ gated_delta_net_cuda(const float * q, attn_partial += s_shard[r] * q_reg[r]; } - float attn_col = warp_reduce_sum(attn_partial); + float attn_col = gdn_warp_reduce_sum(attn_partial); if (lane == 0) { attn_data[col] = attn_col * scale; @@ -166,6 +213,292 @@ gated_delta_net_cuda(const float * q, } } +template +__global__ void __launch_bounds__(512, 2) +gated_delta_net_kda_tiled_128_cuda(const float * q, + const float * k, + const float * v, + const float * g, + const float * beta, + const float * curr_state, + float * dst, + float * state, + int64_t n_tokens, + int64_t sq1, + int64_t sq2, + int64_t sq3, + int64_t sv1, + int64_t sv2, + int64_t sv3, + int64_t sb1, + int64_t sb2, + int64_t sb3, + const uint3 neqk1_magic, + const uint3 rq3_magic, + float scale) { + constexpr int S_v = 128; + constexpr int token_tile = 16; + constexpr int warp_size = 32; + constexpr int rows_per_lane = S_v / warp_size; + + __shared__ float q_shared[token_tile][S_v]; + __shared__ float k_shared[token_tile][S_v]; + __shared__ float g_shared[token_tile][S_v]; + __shared__ float beta_shared[token_tile]; + + const int h_idx = blockIdx.x; + const int sequence = blockIdx.y; + const int lane = threadIdx.x; + const int col = blockIdx.z * blockDim.y + threadIdx.y; + const int thread = threadIdx.y * warp_size + lane; + const int nthreads = blockDim.y * warp_size; + + const uint32_t iq1 = fastmodulo(h_idx, neqk1_magic); + const uint32_t iq3 = fastdiv(sequence, rq3_magic); + + const int64_t state_offset = (sequence * H + h_idx) * S_v * S_v; + curr_state += state_offset + col * S_v; + state += state_offset; + dst += (sequence * n_tokens * H + h_idx) * S_v; + + float s_shard[rows_per_lane]; + ggml_cuda_pdl_sync(); +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + s_shard[r] = curr_state[r * warp_size + lane]; + } + + for (int t0 = 0; t0 < n_tokens; t0 += token_tile) { + const int tile_size = min((int64_t) token_tile, n_tokens - t0); + for (int idx = thread; idx < tile_size * S_v; idx += nthreads) { + const int tt = idx / S_v; + const int i = idx % S_v; + const int t = t0 + tt; + const int64_t gb_offset = sequence * sb3 + t * sb2 + h_idx * sb1; + q_shared[tt][i] = q[iq3 * sq3 + t * sq2 + iq1 * sq1 + i]; + k_shared[tt][i] = k[iq3 * sq3 + t * sq2 + iq1 * sq1 + i]; + g_shared[tt][i] = g[gb_offset * S_v + i]; + } + if (thread < tile_size) { + beta_shared[thread] = beta[sequence * sb3 + (t0 + thread) * sb2 + h_idx * sb1]; + } + __syncthreads(); + + for (int tt = 0; tt < tile_size; ++tt) { + const int t = t0 + tt; + const float * v_t = v + sequence * sv3 + t * sv2 + h_idx * sv1; + float k_reg[rows_per_lane]; + float q_reg[rows_per_lane]; +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + const int i = r * warp_size + lane; + k_reg[r] = k_shared[tt][i]; + q_reg[r] = q_shared[tt][i]; + } + + float kv_shard = 0.0f; +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + const int i = r * warp_size + lane; + kv_shard += expf(g_shared[tt][i]) * s_shard[r] * k_reg[r]; + } + const float kv_col = gdn_warp_reduce_sum(kv_shard); + const float delta_col = (v_t[col] - kv_col) * beta_shared[tt]; + + float attn_partial = 0.0f; +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + const int i = r * warp_size + lane; + s_shard[r] = expf(g_shared[tt][i]) * s_shard[r] + k_reg[r] * delta_col; + attn_partial += s_shard[r] * q_reg[r]; + } + const float attn_col = gdn_warp_reduce_sum(attn_partial); + if (lane == 0) { + dst[(int64_t) t * S_v * H + col] = attn_col * scale; + } + } + __syncthreads(); + } + +#pragma unroll + for (int r = 0; r < rows_per_lane; ++r) { + state[col * S_v + r * warp_size + lane] = s_shard[r]; + } +} + +// Non-KDA variant with the per-token inputs (q, k, g, beta, v) staged in shared memory for a tile of tokens, +// and COLS state columns per warp (k/q are loaded once per token for all columns, and the number of waves +// competing for issue slots shrinks). The recurrence itself is unchanged: for every column the same +// operations are executed in the same order as in gated_delta_net_cuda, so results are bit-identical. +template +__global__ void __launch_bounds__(32 * NUM_WARPS, 1) +gated_delta_net_tiled_cuda(const float * q, + const float * k, + const float * v, + const float * g, + const float * beta, + const float * curr_state, + float * dst, + float * state, + int64_t H, + int64_t n_tokens, + int64_t sq1, + int64_t sq2, + int64_t sq3, + int64_t sv1, + int64_t sv2, + int64_t sv3, + int64_t sb1, + int64_t sb2, + int64_t sb3, + const uint3 neqk1_magic, + const uint3 rq3_magic, + float scale, + int64_t state_slot_stride, + int K) { + constexpr int warp_size = 32; + constexpr int rows_per_lane = S_v / warp_size; + constexpr int block_cols = NUM_WARPS * COLS; + static_assert(S_v % warp_size == 0, "S_v must be a multiple of the warp size"); + static_assert(S_v % block_cols == 0, "block columns must divide S_v"); + + __shared__ float q_shared[TOKEN_TILE][S_v]; + __shared__ float k_shared[TOKEN_TILE][S_v]; + __shared__ float v_shared[TOKEN_TILE][block_cols]; + __shared__ float g_shared[TOKEN_TILE]; + __shared__ float beta_shared[TOKEN_TILE]; + + const uint32_t h_idx = blockIdx.x; + const uint32_t sequence = blockIdx.y; + const int lane = threadIdx.x; + const int col0 = blockIdx.z * block_cols; // first column of the block + const int colw = threadIdx.y * COLS; // first column of this warp inside the block + const int thread = threadIdx.y * warp_size + lane; + constexpr int nthreads = NUM_WARPS * warp_size; + + const uint32_t iq1 = fastmodulo(h_idx, neqk1_magic); + const uint32_t iq3 = fastdiv(sequence, rq3_magic); + + const int64_t state_in_offset = sequence * H * S_v * S_v + h_idx * S_v * S_v; + const int64_t state_out_offset = (sequence * H + h_idx) * S_v * S_v; + state += state_out_offset; + curr_state += state_in_offset + (col0 + colw) * S_v; + float * attn_data = dst + (sequence * n_tokens * H + h_idx) * S_v + col0 + colw; + + float s_shard[COLS][rows_per_lane]; + + ggml_cuda_pdl_sync(); +#pragma unroll + for (int c = 0; c < COLS; c++) { +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + s_shard[c][r] = curr_state[c * S_v + r * warp_size + lane]; + } + } + + for (int t0 = 0; t0 < n_tokens; t0 += TOKEN_TILE) { + const int tile_size = min((int64_t) TOKEN_TILE, n_tokens - t0); + + for (int idx = thread; idx < tile_size * S_v; idx += nthreads) { + const int tt = idx / S_v; + const int i = idx % S_v; + const int t = t0 + tt; + q_shared[tt][i] = q[iq3 * sq3 + t * sq2 + iq1 * sq1 + i]; + k_shared[tt][i] = k[iq3 * sq3 + t * sq2 + iq1 * sq1 + i]; + } + for (int idx = thread; idx < tile_size * block_cols; idx += nthreads) { + const int tt = idx / block_cols; + const int c = idx % block_cols; + const int t = t0 + tt; + v_shared[tt][c] = v[sequence * sv3 + t * sv2 + h_idx * sv1 + col0 + c]; + } + if (thread < tile_size) { + const int64_t gb_offset = sequence * sb3 + (t0 + thread) * sb2 + h_idx * sb1; + g_shared[thread] = g[gb_offset]; + beta_shared[thread] = beta[gb_offset]; + } + __syncthreads(); + + for (int tt = 0; tt < tile_size; ++tt) { + const int t = t0 + tt; + + float k_reg[rows_per_lane]; + float q_reg[rows_per_lane]; +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + const int i = r * warp_size + lane; + k_reg[r] = k_shared[tt][i]; + q_reg[r] = q_shared[tt][i]; + } + + const float g_val = expf(g_shared[tt]); + const float beta_val = beta_shared[tt]; + + // The FMA contractions are spelled out so that the results are bit-identical to the + // compiler's contraction of gated_delta_net_cuda (fma chains from 0, fma(-g, kv, v)*beta, + // fma(g, s, k*delta)) independently of the surrounding code. + float attn_col[COLS]; +#pragma unroll + for (int c = 0; c < COLS; c++) { + // kv[col] = (S^T @ k)[col] = sum_i S[i][col] * k[i] + float kv_shard = 0.0f; +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + kv_shard = fmaf(s_shard[c][r], k_reg[r], kv_shard); + } + const float kv_col = gdn_warp_reduce_sum(kv_shard); + + // delta[col] = (v[col] - g * kv[col]) * beta + const float delta_col = fmaf(-g_val, kv_col, v_shared[tt][colw + c]) * beta_val; + + // fused: S[i][col] = g * S[i][col] + k[i] * delta[col] + // attn[col] = (S^T @ q)[col] = sum_i S[i][col] * q[i] + float attn_partial = 0.0f; +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + s_shard[c][r] = fmaf(g_val, s_shard[c][r], k_reg[r] * delta_col); + attn_partial = fmaf(s_shard[c][r], q_reg[r], attn_partial); + } + attn_col[c] = gdn_warp_reduce_sum(attn_partial); + } + + if (lane < COLS) { + float a = attn_col[0]; +#pragma unroll + for (int c = 1; c < COLS; c++) { + a = lane == c ? attn_col[c] : a; + } + attn_data[(int64_t) t * S_v * H + lane] = a * scale; + } + + if constexpr (keep_rs_t) { + const int target_slot = (int) n_tokens - 1 - t; + if (target_slot >= 0 && target_slot < K) { + float * snapshot = state + target_slot * state_slot_stride + (col0 + colw) * S_v; +#pragma unroll + for (int c = 0; c < COLS; c++) { +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + snapshot[c * S_v + r * warp_size + lane] = s_shard[c][r]; + } + } + } + } + } + __syncthreads(); + } + + if constexpr (!keep_rs_t) { +#pragma unroll + for (int c = 0; c < COLS; c++) { +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + state[(col0 + colw + c) * S_v + r * warp_size + lane] = s_shard[c][r]; + } + } + } +} + template static void launch_gated_delta_net( const float * q_d, const float * k_d, const float * v_d, @@ -178,14 +511,54 @@ static void launch_gated_delta_net( int64_t neqk1, int64_t rq3, float scale, int64_t state_slot_stride, int K, cudaStream_t stream) { //TODO: Add chunked kernel for even faster pre-fill - const int warp_size = ggml_cuda_info().devices[ggml_cuda_get_device()].warp_size; - const int num_warps = 4; + const int id = ggml_cuda_get_device(); + const int cc = ggml_cuda_info().devices[id].cc; + const int warp_size = ggml_cuda_info().devices[id].warp_size; + // Qwen3.6 PP2048 measured faster on gfx1151 on 2026-09-02. Retest if occupancy changes. + const int num_warps = GGML_CUDA_CC_IS_RDNA3_5(cc) && S_v == 128 && (H == 32 || H == 48 || H == 64) && !KDA ? 32 : + (GGML_CUDA_CC_IS_RDNA3_5(cc) ? 16 : gated_delta_net_num_warps(cc)); dim3 grid_dims(H, n_seqs, (S_v + num_warps - 1) / num_warps); dim3 block_dims(warp_size <= S_v ? warp_size : S_v, num_warps, 1); const uint3 neqk1_magic = init_fastdiv_values(neqk1); const uint3 rq3_magic = init_fastdiv_values(rq3); + if constexpr (KDA && !keep_rs_t) { + if (GGML_CUDA_CC_IS_RDNA3_5(cc) && S_v == 128 && (H == 16 || H == 32) && n_tokens >= 16 && n_seqs == 1) { + const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(grid_dims, block_dims, 0, stream); + auto kernel = H == 16 ? gated_delta_net_kda_tiled_128_cuda<16> : gated_delta_net_kda_tiled_128_cuda<32>; + ggml_cuda_kernel_launch(kernel, launch_params, + q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, n_tokens, + sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2, sb3, + neqk1_magic, rq3_magic, scale); + return; + } + } + + if constexpr (!KDA) { + if (GGML_CUDA_CC_IS_RDNA3_5(cc) && S_v == 128 && num_warps == 32 && n_tokens >= 16) { + // 16 warps x 4 columns per block, 16-token tiles: best of the swept configurations on gfx1151 + // (2048 tokens, H=32: 5.03 ms -> 2.2 ms; H=64: 9.29 ms -> 4.5 ms). + const int cfg = getenv("GGML_GDN_TILE_CFG") ? atoi(getenv("GGML_GDN_TILE_CFG")) : (H == 48 ? 1 : 0); + auto launch_tiled = [&](auto kernel, int warps, int cols) { + const dim3 tiled_grid(H, n_seqs, S_v / (warps * cols)); + const dim3 tiled_block(warp_size, warps, 1); + const ggml_cuda_kernel_launch_params launch_params(tiled_grid, tiled_block, 0, stream); + ggml_cuda_kernel_launch(kernel, launch_params, + q_d, k_d, v_d, g_d, b_d, s_d, dst_d, state_d, H, n_tokens, + sq1, sq2, sq3, sv1, sv2, sv3, sb1, sb2, sb3, + neqk1_magic, rq3_magic, scale, state_slot_stride, K); + }; + switch (cfg) { + case 1: launch_tiled(gated_delta_net_tiled_cuda<128, 8, 8, 16, keep_rs_t>, 8, 8); break; + case 2: launch_tiled(gated_delta_net_tiled_cuda<128, 32, 2, 16, keep_rs_t>, 32, 2); break; + case 3: launch_tiled(gated_delta_net_tiled_cuda<128, 16, 4, 8, keep_rs_t>, 16, 4); break; + default: launch_tiled(gated_delta_net_tiled_cuda<128, 16, 4, 16, keep_rs_t>, 16, 4); break; + } + return; + } + } + const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(grid_dims, block_dims, 0, stream); switch (S_v) { case 16: @@ -325,3 +698,237 @@ void ggml_cuda_op_gated_delta_net_fused_cache( ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_cuda_gated_delta_net_fused_cache cache) { ggml_cuda_op_gated_delta_net_impl(ctx, dst, &cache); } + +// --------------------------------------------------------------------------------------------- +// Fused Qwen3.5/3.6 Gated DeltaNet decode step (n_tokens == 1, n_seqs == 1, S == 128): one block per value +// head computes the causal conv + SiLU of its q/k/v channels (writing the shifted conv state back into the +// cache), the q/k L2 norms, the gate (softplus) and beta (sigmoid), the recurrence with the new state written +// straight into the cache, and the gated RMS norm of the attention output. This replaces 14 graph nodes +// (11 kernels plus their ~2 us launch gaps on gfx1151) that together take ~50 us per layer in decode. +// +// Every op is spelled with the same operation order and FMA contraction as the kernels it replaces +// (ssm_conv_f32, l2_norm_dual_f32_s128, unary/binary elementwise, gated_delta_net_cuda, rms_norm_f32<128>), +// so the outputs are bit-identical to the unfused graph. + +static __device__ __forceinline__ float gdn_decode_sigmoid(float x) { + return 1.0f / (1.0f + expf(-x)); +} + +static __device__ __forceinline__ float gdn_decode_softplus(float x) { + return (x > 20.0f) ? x : logf(1.0f + expf(x)); +} + +template +__global__ void __launch_bounds__(32 * 32, 1) +gdn_decode_fused_cuda(const ggml_cuda_gdn_decode_args args) { + constexpr int warp_size = 32; + constexpr int nwarps = 32; + constexpr int rows_per_lane = S / warp_size; + constexpr int cols_per_warp = S / nwarps; + static_assert(S == warp_size * rows_per_lane && S == nwarps * cols_per_warp); + + __shared__ float q_c[S]; + __shared__ float k_c[S]; + __shared__ float v_c[S]; + __shared__ float q_n[S]; + __shared__ float k_n[S]; + __shared__ float attn[S]; + __shared__ float red[4]; + + const int h = blockIdx.x; // value head + const int lane = threadIdx.x; + const int warp = threadIdx.y; + const int thread = warp * warp_size + lane; + const int hk = h % args.H_k; // key/query head (fastmodulo(h_idx, neqk1) in gated_delta_net_cuda) + + // 1. causal conv + SiLU for this head's q, k and v channels; shift the conv state (CPY conv_state_last) + if (thread < 3 * S) { + const int part = thread / S; // 0: q, 1: k, 2: v + const int i = thread % S; + const int c = part == 0 ? hk * S : part == 1 ? args.H_k * S + hk * S : 2 * args.H_k * S + h * S; + const int ch = c + i; + + float x[D_CONV]; + float w[D_CONV]; +#pragma unroll + for (int j = 0; j < D_CONV - 1; j++) { + x[j] = args.conv_state_in[ch * (D_CONV - 1) + j]; + } + x[D_CONV - 1] = args.qkv[ch * args.qkv_stride]; +#pragma unroll + for (int j = 0; j < D_CONV; j++) { + w[j] = args.conv_w[ch * D_CONV + j]; + } + + // ssm_conv_f32: fma chain from 0, then the (zero) bias add, then SiLU + float sumf = fmaf(x[0], w[0], 0.0f); +#pragma unroll + for (int j = 1; j < D_CONV; j++) { + sumf = fmaf(x[j], w[j], sumf); + } + sumf += args.conv_bias; + const float y = ggml_cuda_op_silu_single(sumf); + + if (part == 0) { + q_c[i] = y; + } else if (part == 1) { + k_c[i] = y; + } else { + v_c[i] = y; + } + + // q/k channels are shared by H_v/H_k value heads: the first one writes the conv state + if (part == 2 || h == hk) { +#pragma unroll + for (int j = 0; j < D_CONV - 1; j++) { + args.conv_state_out[ch * (D_CONV - 1) + j] = x[j + 1]; + } + } + } + __syncthreads(); + + // 2. l2_norm_dual_f32_s128 for q (warp 0) and k (warp 1) + if (warp < 2) { + const float * x = warp == 0 ? q_c : k_c; + float * y = warp == 0 ? q_n : k_n; + float tmp = 0.0f; +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + const float xi = x[r * warp_size + lane]; + tmp = fmaf(xi, xi, tmp); + } + tmp = gdn_warp_reduce_sum(tmp); + const float scale = rsqrtf(fmaxf(tmp, args.eps_l2 * args.eps_l2)); +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + y[r * warp_size + lane] = scale * x[r * warp_size + lane]; + } + } + + // gate = softplus(alpha + dt_bias) * A, beta = sigmoid(beta): scalar per head, computed by every thread + const float alpha_biased = args.alpha[h] + args.dt_bias[h]; + const float gate = gdn_decode_softplus(alpha_biased) * args.ssm_a[h]; + const float beta_val = gdn_decode_sigmoid(args.beta[h]); + const float g_val = expf(gate); + __syncthreads(); + + // 3. recurrence (gated_delta_net_cuda), cols_per_warp state columns per warp + { + const float * state_in = args.state_cache + (int64_t) args.state_ids[0] * args.state_row_stride + (int64_t) h * S * S; + float * state_out = args.state_out + (int64_t) h * S * S; + const int colw = warp * cols_per_warp; + + float k_reg[rows_per_lane]; + float q_reg[rows_per_lane]; +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + k_reg[r] = k_n[r * warp_size + lane]; + q_reg[r] = q_n[r * warp_size + lane]; + } + + float s_shard[cols_per_warp][rows_per_lane]; +#pragma unroll + for (int c = 0; c < cols_per_warp; c++) { +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + s_shard[c][r] = state_in[(colw + c) * S + r * warp_size + lane]; + } + } + + // same explicit contractions as gated_delta_net_tiled_cuda (bit-identical to gated_delta_net_cuda) + float attn_col[cols_per_warp]; +#pragma unroll + for (int c = 0; c < cols_per_warp; c++) { + float kv_shard = 0.0f; +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + kv_shard = fmaf(s_shard[c][r], k_reg[r], kv_shard); + } + const float kv_col = gdn_warp_reduce_sum(kv_shard); + + const float delta_col = fmaf(-g_val, kv_col, v_c[colw + c]) * beta_val; + + float attn_partial = 0.0f; +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + s_shard[c][r] = fmaf(g_val, s_shard[c][r], k_reg[r] * delta_col); + attn_partial = fmaf(s_shard[c][r], q_reg[r], attn_partial); + } + attn_col[c] = gdn_warp_reduce_sum(attn_partial); + } + + if (lane < cols_per_warp) { + float a = attn_col[0]; +#pragma unroll + for (int c = 1; c < cols_per_warp; c++) { + a = lane == c ? attn_col[c] : a; + } + attn[colw + lane] = a * args.scale; + } + +#pragma unroll + for (int c = 0; c < cols_per_warp; c++) { +#pragma unroll + for (int r = 0; r < rows_per_lane; r++) { + state_out[(colw + c) * S + r * warp_size + lane] = s_shard[c][r]; + } + } + } + __syncthreads(); + + // 4. pre-norm attention output to the scratch buffer (the gated rms norm runs as a second kernel: its + // destination may alias inputs of this kernel that other blocks are still reading) + if (thread < S) { + args.attn_out[(int64_t) h * S + thread] = attn[thread]; + } +} + +// rms_norm_f32<128, true>: per-warp partial sums, then the second-level butterfly of block_reduce over +// lanes 0..3 (which reduces to (s0 + s2) + (s1 + s3)), times the norm weight +template +__global__ void __launch_bounds__(S, 1) +gdn_decode_norm_cuda(const float * attn, const float * norm_w, float * out, const float eps) { + constexpr int warp_size = 32; + static_assert(S == 4 * warp_size); + __shared__ float red[4]; + + const int h = blockIdx.x; + const int tid = threadIdx.x; + const int warp = tid / warp_size; + const int lane = tid % warp_size; + + const float xi = attn[(int64_t) h * S + tid]; + const float part = gdn_warp_reduce_sum(fmaf(xi, xi, 0.0f)); + if (lane == 0) { + red[warp] = part; + } + __syncthreads(); + const float tmp = (red[0] + red[2]) + (red[1] + red[3]); + const float mean = tmp / (float) S; + const float scale = rsqrtf(mean + eps); + out[(int64_t) h * S + tid] = scale * xi * norm_w[tid]; +} + +void ggml_cuda_op_gdn_decode_fused_prenorm(ggml_backend_cuda_context & ctx, const ggml_cuda_gdn_decode_args & args_in, float * attn_scratch) { + GGML_ASSERT(args_in.S == 128 && args_in.d_conv == 4); + ggml_cuda_gdn_decode_args args = args_in; + args.attn_out = attn_scratch; + + const dim3 grid(args.H_v, 1, 1); + const dim3 block(32, 32, 1); + const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(grid, block, 0, ctx.stream()); + ggml_cuda_kernel_launch(gdn_decode_fused_cuda<128, 4>, launch_params, args); +} + +void ggml_cuda_op_gdn_decode_fused(ggml_backend_cuda_context & ctx, const ggml_cuda_gdn_decode_args & args_in) { + GGML_ASSERT(args_in.S == 128 && args_in.d_conv == 4); + ggml_cuda_pool_alloc attn(ctx.pool(), args_in.S * args_in.H_v); + ggml_cuda_gdn_decode_args args = args_in; + args.attn_out = attn.get(); + ggml_cuda_op_gdn_decode_fused_prenorm(ctx, args, attn.get()); + + const dim3 grid(args.H_v, 1, 1); + const dim3 norm_block(128, 1, 1); + const ggml_cuda_kernel_launch_params norm_params = ggml_cuda_kernel_launch_params(grid, norm_block, 0, ctx.stream()); + ggml_cuda_kernel_launch(gdn_decode_norm_cuda<128>, norm_params, attn.get(), args.norm_w, args.out, args.eps_rms); +} diff --git a/ggml/src/ggml-cuda/gated_delta_net.cuh b/ggml/src/ggml-cuda/gated_delta_net.cuh index f9bf43706789..e2b41cf1c47f 100644 --- a/ggml/src/ggml-cuda/gated_delta_net.cuh +++ b/ggml/src/ggml-cuda/gated_delta_net.cuh @@ -12,3 +12,37 @@ void ggml_cuda_op_gated_delta_net(ggml_backend_cuda_context & ctx, ggml_tensor * // same op, but writes the snapshot(s) into the cache instead of dst (see ggml_cuda_try_gdn_cache_fusion) void ggml_cuda_op_gated_delta_net_fused_cache(ggml_backend_cuda_context & ctx, ggml_tensor * dst, ggml_cuda_gated_delta_net_fused_cache cache); + +// Fused Qwen3.5 Gated DeltaNet decode step (see ggml_cuda_try_gdn_decode_fusion). All strides in elements. +struct ggml_cuda_gdn_decode_args { + // conv: conv_state_in [d_conv-1, C] (contiguous, per channel), qkv [C] with stride, conv_w [d_conv, C] + const float * conv_state_in; + const float * qkv; + int64_t qkv_stride; + const float * conv_w; + float conv_bias; // 0.0f (ssm_conv adds a runtime zero when there is no bias) + float * conv_state_out; // [d_conv-1, C] (contiguous) cache row + // gate / beta + const float * alpha; // [H_v] + const float * dt_bias; // [H_v] + const float * ssm_a; // [H_v] + const float * beta; // [H_v] + float eps_l2; + // recurrence + const float * state_cache; // [S*S*H_v, n_slots] cache, row = state_ids[0] (transposed per head: [col][row]) + const int32_t * state_ids; + int64_t state_row_stride; + float * state_out; // [S, S, H_v] cache row + float scale; // 1/sqrt(S) + // gated rms norm + const float * norm_w; // [S] + const ggml_tensor * norm_w_tensor = nullptr; + float eps_rms; + float * out; // [S, H_v] + float * attn_out; // [S, H_v] pre-norm scratch (set by the host) + int64_t S, H_k, H_v, d_conv; +}; + +void ggml_cuda_op_gdn_decode_fused(ggml_backend_cuda_context & ctx, const ggml_cuda_gdn_decode_args & args); +// first kernel only: writes the pre-norm attention output [S, H_v] to attn_scratch (the gated norm is applied by the consumer) +void ggml_cuda_op_gdn_decode_fused_prenorm(ggml_backend_cuda_context & ctx, const ggml_cuda_gdn_decode_args & args, float * attn_scratch); diff --git a/ggml/src/ggml-cuda/gdn-conv.cu b/ggml/src/ggml-cuda/gdn-conv.cu new file mode 100644 index 000000000000..57a80fd290e2 --- /dev/null +++ b/ggml/src/ggml-cuda/gdn-conv.cu @@ -0,0 +1,163 @@ +#include "gdn-conv.cuh" +#include "conv-tail.cuh" +#include "unary.cuh" + +// columns [tail_from, T+3) of every channel row of the concat: what the history copies read +static __global__ void gdn_concat_tail(const float * __restrict__ state, const float * __restrict__ x, float * __restrict__ out, + const int C, const int T, const int tail_from, const int row_stride) { + const int ncols = T + 3 - tail_from; + const int64_t idx = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= (int64_t) C * ncols) { + return; + } + const int c = (int) (idx / ncols), j = tail_from + (int) (idx % ncols); + out[(size_t) c * row_stride + j] = (j < 3) ? state[c * 3 + j] : x[(size_t) (j - 3) * C + c]; +} + +// the same accumulation as ssm_conv_f32 with a zero bias, then SiLU +template +static __global__ void __launch_bounds__(256) gdn_conv_direct_kernel(const float * __restrict__ state, const float * __restrict__ x, + const float * __restrict__ w, float * __restrict__ y, const int C, const int T) { + const int c = blockIdx.x * 256 + threadIdx.x; + const int t0 = blockIdx.y * TT; + if (c >= C) { + return; + } + float wr[4]; +#pragma unroll + for (int j = 0; j < 4; ++j) { + wr[j] = w[c * 4 + j]; + } + float xw[4]; +#pragma unroll + for (int k = 0; k < 4; ++k) { + const int jp = t0 + k; + xw[k] = (jp < 3) ? state[c * 3 + jp] : x[(size_t) (jp - 3) * C + c]; + } + const float b = 0.0f; + const int tend = min(T, t0 + TT); + for (int t = t0; t < tend; ++t) { + float sumf = 0.0f; +#pragma unroll + for (int j = 0; j < 4; ++j) { + sumf += xw[j] * wr[j]; + } + sumf += b; + y[(size_t) t * C + c] = ggml_cuda_op_silu_single(sumf); + xw[0] = xw[1]; xw[1] = xw[2]; xw[2] = xw[3]; + const int jp = t + 4; + xw[3] = (jp < 3) ? state[c * 3 + jp] : ((jp - 3) < T ? x[(size_t) (jp - 3) * C + c] : 0.0f); + } +} + +static bool gdn_conv_check(const ggml_cgraph * cgraph, int i, ggml_cuda_gdn_conv_match & m) { + if (i < 0 || i + 2 >= cgraph->n_nodes) { + return ggml_cuda_conv_reject("gdn_conv", "check 1", i); + } + const ggml_tensor * cc = cgraph->nodes[i]; + if (cc->op != GGML_OP_CONCAT || cc->type != GGML_TYPE_F32 || ggml_get_op_params_i32(cc, 0) != 0 || + (cc->flags & GGML_TENSOR_FLAG_OUTPUT) || !ggml_cuda_conv_next_is_noop(cgraph, i)) { + return ggml_cuda_conv_reject("gdn_conv", "check 2", i); + } + const ggml_tensor * st = cc->src[0]; + const ggml_tensor * tr = cc->src[1]; + if (!st || !tr || st->type != GGML_TYPE_F32 || tr->type != GGML_TYPE_F32 || tr->op != GGML_OP_TRANSPOSE || !tr->view_src) { + return ggml_cuda_conv_reject("gdn_conv", "check 3", i); + } + int64_t C = 0, T = 0; + const ggml_tensor * x = ggml_cuda_conv_activation(tr, C, T); + if (!x) { + return ggml_cuda_conv_reject("gdn_conv", "check 4", i); + } + if (st->ne[0] != 3 || st->ne[1] != C || st->ne[2] != 1 || st->ne[3] != 1 || !ggml_is_contiguous(st)) { + return ggml_cuda_conv_reject("gdn_conv", "check 5", i); + } + if (tr->ne[0] != T || tr->ne[1] != C || tr->ne[2] != 1 || tr->view_offs != 0) { + return ggml_cuda_conv_reject("gdn_conv", "check 6", i); + } + if (cc->ne[0] != T + 3 || cc->ne[1] != C || cc->ne[2] != 1 || cc->ne[3] != 1 || !ggml_is_contiguous(cc) || T < 256 || C % 256 != 0) { + return ggml_cuda_conv_reject("gdn_conv", "check 7", i); + } + + // the conv is the only reader of the whole concat; everything else must be a tail reader + int64_t tail_from = T + 3; + int conv = -1; + for (int n = i + 1; n < cgraph->n_nodes; ++n) { + const ggml_tensor * t = cgraph->nodes[n]; + if (!ggml_cuda_reads(t, cc)) { + continue; + } + if (t->op == GGML_OP_SSM_CONV && ggml_cuda_conv_whole(t->src[0], cc)) { + if (conv >= 0) { + return ggml_cuda_conv_reject("gdn_conv", "check 8", n); + } + conv = n; + continue; + } + if (ggml_cuda_conv_whole(t, cc) && (t->op == GGML_OP_RESHAPE || t->op == GGML_OP_VIEW)) { + continue; // shape plumbing into the conv; its readers are checked as readers of the concat + } + if (!ggml_cuda_conv_tail_reader(t, cc, 3, tail_from)) { + return ggml_cuda_conv_reject("gdn_conv", "check 9", n); + } + } + if (conv < 0 || conv + 1 >= cgraph->n_nodes) { + return ggml_cuda_conv_reject("gdn_conv", "check 10", i); + } + const ggml_tensor * cv = cgraph->nodes[conv]; + const ggml_tensor * w = cv->src[1]; + if (cv->type != GGML_TYPE_F32 || cv->src[2] != nullptr || !w || w->type != GGML_TYPE_F32 || w->ne[0] != 4 || w->ne[1] != C || + !ggml_is_contiguous(w) || cv->ne[0] != C || cv->ne[1] != T || cv->ne[2] != 1 || cv->ne[3] != 1 || !ggml_is_contiguous(cv) || + (cv->flags & GGML_TENSOR_FLAG_OUTPUT)) { + return ggml_cuda_conv_reject("gdn_conv", "check 11", i); + } + const ggml_tensor * su = cgraph->nodes[conv + 1]; + if (su->op != GGML_OP_UNARY || ggml_get_unary_op(su) != GGML_UNARY_OP_SILU || su->src[0] != cv || su->type != GGML_TYPE_F32 || + !ggml_is_contiguous(su) || ggml_cuda_conv_readers(cgraph, conv + 1, cv, nullptr) != 1) { + return ggml_cuda_conv_reject("gdn_conv", "check 12", i); + } + m.concat_idx = i; m.conv_idx = conv; m.x = x; m.state = st; m.concat = cc; m.w = w; m.out = cgraph->nodes[conv + 1]; + m.C = C; m.T = T; m.tail_from = std::max(0, tail_from); + return true; +} + +bool ggml_cuda_gdn_conv_match_at_concat(const ggml_cgraph * cgraph, int i, ggml_cuda_gdn_conv_match & m) { + return gdn_conv_check(cgraph, i, m); +} + +bool ggml_cuda_gdn_conv_match_at_conv(const ggml_cgraph * cgraph, int j, ggml_cuda_gdn_conv_match & m) { + if (j < 1 || cgraph->nodes[j]->op != GGML_OP_SSM_CONV) { + return false; + } + // the conv reads the concat directly or through a whole-tensor view; both entry points must find the same node + const ggml_tensor * src = cgraph->nodes[j]->src[0]; + const ggml_tensor * cc = src ? (src->op == GGML_OP_CONCAT ? src : src->view_src) : nullptr; + if (!cc || cc->op != GGML_OP_CONCAT || !ggml_cuda_conv_whole(src, cc)) { + return false; + } + for (int i = j - 1; i >= 0; --i) { + if (cgraph->nodes[i] == cc) { + return gdn_conv_check(cgraph, i, m) && m.conv_idx == j; + } + } + return false; +} + +void ggml_cuda_gdn_conv_write_tail(ggml_backend_cuda_context & ctx, const ggml_cuda_gdn_conv_match & m) { + const int ncols = (int) (m.T + 3 - m.tail_from); + if (ncols <= 0) { + return; + } + const int64_t n = m.C * ncols; + gdn_concat_tail<<<(unsigned) ((n + 255) / 256), 256, 0, ctx.stream()>>>((const float *) m.state->data, (const float *) m.x->data, + (float *) m.concat->data, (int) m.C, (int) m.T, (int) m.tail_from, (int) (m.concat->nb[1] / sizeof(float))); + CUDA_CHECK(cudaGetLastError()); +} + +void ggml_cuda_gdn_conv_direct(ggml_backend_cuda_context & ctx, const ggml_cuda_gdn_conv_match & m) { + constexpr int TT = 128; + dim3 grid((unsigned) (m.C / 256), (unsigned) ((m.T + TT - 1) / TT)); + gdn_conv_direct_kernel<<>>((const float *) m.state->data, (const float *) m.x->data, + (const float *) m.w->data, (float *) m.out->data, (int) m.C, (int) m.T); + CUDA_CHECK(cudaGetLastError()); +} diff --git a/ggml/src/ggml-cuda/gdn-conv.cuh b/ggml/src/ggml-cuda/gdn-conv.cuh new file mode 100644 index 000000000000..e5e3b6a62810 --- /dev/null +++ b/ggml/src/ggml-cuda/gdn-conv.cuh @@ -0,0 +1,18 @@ +#pragma once +#include "common.cuh" + +// Gated DeltaNet prefill convolution: concat(history, x^T) -> ssm_conv -> SiLU. On a match the concat writes only +// the columns its tail copies read and one kernel computes the SiLU output straight from the history and x. +struct ggml_cuda_gdn_conv_match { + int concat_idx = -1, conv_idx = -1; + const ggml_tensor * x = nullptr; // [C, T] F32 contiguous + const ggml_tensor * state = nullptr; // [3, C] F32 contiguous + const ggml_tensor * concat = nullptr; // [T+3, C] F32 + const ggml_tensor * w = nullptr; // [4, C] F32 + ggml_tensor * out = nullptr; // SiLU output [C, T] + int64_t C = 0, T = 0, tail_from = 0; +}; +bool ggml_cuda_gdn_conv_match_at_concat(const ggml_cgraph * cgraph, int i, ggml_cuda_gdn_conv_match & m); +bool ggml_cuda_gdn_conv_match_at_conv(const ggml_cgraph * cgraph, int j, ggml_cuda_gdn_conv_match & m); +void ggml_cuda_gdn_conv_write_tail(ggml_backend_cuda_context & ctx, const ggml_cuda_gdn_conv_match & m); +void ggml_cuda_gdn_conv_direct(ggml_backend_cuda_context & ctx, const ggml_cuda_gdn_conv_match & m); diff --git a/ggml/src/ggml-cuda/getrows.cu b/ggml/src/ggml-cuda/getrows.cu index 6b36b8fbc18c..ffcb51f00912 100644 --- a/ggml/src/ggml-cuda/getrows.cu +++ b/ggml/src/ggml-cuda/getrows.cu @@ -101,6 +101,38 @@ static __global__ void k_get_rows_float( } } +template +static __global__ void k_get_rows_float_packed( + const src0_t * src0, const int32_t * src1, dst_t * dst, + const int64_t ne00, const int64_t ne10, const int64_t ne11, const uint3 ne12_fdv, + const size_t s1, const size_t s2, const size_t s3, + const size_t nb01, const size_t nb02, const size_t nb03, + const size_t s10, const size_t s11, const size_t s12) { + constexpr int rows_per_block = CUDA_GET_ROWS_BLOCK_SIZE / cols; + const int64_t row = int64_t(blockIdx.x) * rows_per_block + threadIdx.x / cols; + const int col = threadIdx.x % cols; + + ggml_cuda_pdl_lc(); + ggml_cuda_pdl_sync(); + if (row >= ne10 || col >= ne00) { + return; + } + + for (int64_t z = blockIdx.y; z < ne11 * int64_t(ne12_fdv.z); z += gridDim.y) { + const uint2 dm = fast_div_modulo(uint32_t(z), ne12_fdv); + const int32_t index = src1[row * s10 + dm.x * s11 + dm.y * s12]; + const src0_t * src_row = (const src0_t *) ((const char *) src0 + index * nb01 + dm.x * nb02 + dm.y * nb03); + dst_t * dst_row = dst + row * s1 + dm.x * s2 + dm.y * s3; +#pragma unroll + for (int v = 0; v < values_per_lane; ++v) { + const int column = col + v * cols; + if (column < ne00) { + dst_row[column] = ggml_cuda_cast(src_row[column]); + } + } + } +} + template static __global__ void k_get_rows_float_vec( const dst_t * src0_ptr, const int32_t * src1_ptr, dst_t * dst_ptr, @@ -254,6 +286,40 @@ static void get_rows_cuda_float( GGML_ASSERT(ne11 <= std::numeric_limits::max() / ne12); const uint3 ne12_fdv = init_fastdiv_values(ne12); + if constexpr (std::is_same::value && + (std::is_same::value || std::is_same::value)) { + if (ne00 > 0 && ne00 <= 128 && ne10 >= 128 && + GGML_CUDA_CC_IS_RDNA3_5(ggml_cuda_info().devices[ggml_cuda_get_device()].cc)) { + auto launch = [&](auto width, auto elements) { + constexpr int cols = decltype(width)::value; + constexpr int values = decltype(elements)::value; + constexpr int rows_per_block = CUDA_GET_ROWS_BLOCK_SIZE / cols; + const dim3 grid((ne10 + rows_per_block - 1) / rows_per_block, MIN(ne11 * ne12, UINT16_MAX), 1); + const ggml_cuda_kernel_launch_params params{grid, block_dims, 0, stream}; + ggml_cuda_kernel_launch(k_get_rows_float_packed, params, + src0_d, src1_d, dst_d, ne00, ne10, ne11, ne12_fdv, + s1, s2, s3, nb01, nb02, nb03, s10, s11, s12); + }; + constexpr auto one = std::integral_constant{}; + if constexpr (std::is_same::value) { + if (ne00 > 32) { + if (ne00 <= 64) launch(std::integral_constant{}, std::integral_constant{}); + else launch(std::integral_constant{}, std::integral_constant{}); + return; + } + } + if (ne00 <= 1) launch(std::integral_constant{}, one); + else if (ne00 <= 2) launch(std::integral_constant{}, one); + else if (ne00 <= 4) launch(std::integral_constant{}, one); + else if (ne00 <= 8) launch(std::integral_constant{}, one); + else if (ne00 <= 16) launch(std::integral_constant{}, one); + else if (ne00 <= 32) launch(std::integral_constant{}, one); + else if (ne00 <= 64) launch(std::integral_constant{}, one); + else launch(std::integral_constant{}, one); + return; + } + } + if constexpr (std::is_same::value) { constexpr int VEC = 16 / sizeof(dst_t); const int64_t ne00v = ne00 / VEC; @@ -488,3 +554,43 @@ void ggml_cuda_op_get_rows_back(ggml_backend_cuda_context & ctx, ggml_tensor * d k_get_rows_back_float<<>>(src0_d, src1_d, dst_d, ne00, ne10, ne1); } + +// QSA indexer key pooling: mean of the 4 member rows of every block, 128-wide rows, one thread per column. +// Fuses get_rows -> reshape -> 4 x (view, cont, add) -> scale on the same summation order as the graph. +template +static __global__ void get_rows_mean4_128( + const char * src, const char * ids, float * dst, + size_t src_row_stride, size_t src_stream_stride, size_t ids_stream_stride, + float scale, float bias) { + ggml_cuda_pdl_lc(); + ggml_cuda_pdl_sync(); + const int64_t group = blockIdx.x; + const int64_t stream = blockIdx.y; + const int col = threadIdx.x; + const int32_t * selected = (const int32_t *) (ids + stream*ids_stream_stride) + 4*group; + const char * input = src + stream*src_stream_stride; + float sum = (float) ((const T *) (input + selected[0]*src_row_stride))[col]; +#pragma unroll + for (int i = 1; i < 4; ++i) { + sum += (float) ((const T *) (input + selected[i]*src_row_stride))[col]; + } + dst[(stream*gridDim.x + group)*128 + col] = scale*sum + bias; +} + +void ggml_cuda_op_get_rows_mean4( + ggml_backend_cuda_context & ctx, const ggml_tensor * src, const ggml_tensor * ids, ggml_tensor * dst) { + const dim3 grid(dst->ne[1], dst->ne[2], 1); + const ggml_cuda_kernel_launch_params params(grid, dim3(128, 1, 1), 0, ctx.stream()); + const float scale = ggml_get_op_params_f32(dst, 0); + const float bias = ggml_get_op_params_f32(dst, 1); + if (src->type == GGML_TYPE_F16) { + ggml_cuda_kernel_launch(get_rows_mean4_128, params, + (const char *) src->data, (const char *) ids->data, (float *) dst->data, + src->nb[1], src->nb[2], ids->nb[1], scale, bias); + } else { + GGML_ASSERT(src->type == GGML_TYPE_F32); + ggml_cuda_kernel_launch(get_rows_mean4_128, params, + (const char *) src->data, (const char *) ids->data, (float *) dst->data, + src->nb[1], src->nb[2], ids->nb[1], scale, bias); + } +} diff --git a/ggml/src/ggml-cuda/getrows.cuh b/ggml/src/ggml-cuda/getrows.cuh index 3c5bea5f48c1..e8b9d7bd67eb 100644 --- a/ggml/src/ggml-cuda/getrows.cuh +++ b/ggml/src/ggml-cuda/getrows.cuh @@ -13,3 +13,7 @@ void get_rows_cuda( void ggml_cuda_op_get_rows(ggml_backend_cuda_context & ctx, ggml_tensor * dst); void ggml_cuda_op_get_rows_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +// fused QSA indexer key pooling (get_rows of 4 member rows per block, summed in graph order, scaled) +void ggml_cuda_op_get_rows_mean4( + ggml_backend_cuda_context & ctx, const ggml_tensor * src, const ggml_tensor * ids, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/ggml-cuda.cu b/ggml/src/ggml-cuda/ggml-cuda.cu index 74bb47145b7d..d3f9870db762 100644 --- a/ggml/src/ggml-cuda/ggml-cuda.cu +++ b/ggml/src/ggml-cuda/ggml-cuda.cu @@ -60,6 +60,13 @@ #include "ggml-cuda/wkv.cuh" #include "ggml-cuda/gla.cuh" #include "ggml-cuda/gated_delta_net.cuh" +#include "ggml-cuda/hyperconn.cuh" +#include "ggml-cuda/norm-gated.cuh" +#include "ggml-cuda/idx-relu-sum.cuh" +#include "ggml-cuda/ple-conv.cuh" +#include "ggml-cuda/gdn-conv.cuh" +#include "ggml-cuda/mmb.cuh" +#include "ggml-cuda/expand.cuh" #include "ggml-cuda/dsv4-hc.cuh" #include "ggml-cuda/set.cuh" #include "ggml-cuda/set-rows.cuh" @@ -702,6 +709,23 @@ ggml_backend_cuda_context::~ggml_backend_cuda_context() { std::unique_lock lock(ggml_cuda_lock); ggml_cuda_lock_cv.wait(lock, []{ return ggml_cuda_lock_counter.load(std::memory_order_relaxed) == 0; }); + if (mmb) { + // Queued work and captured graphs can still reference the MMB buffers. + for (int i = 0; i < GGML_CUDA_MAX_DEVICES; ++i) { + for (int j = 0; j < GGML_CUDA_MAX_STREAMS; ++j) { + if (streams[i][j]) { + ggml_cuda_set_device(i); + CUDA_CHECK(cudaStreamSynchronize(streams[i][j])); + } + } + } + ggml_cuda_set_device(device); +#ifdef USE_CUDA_GRAPH + cuda_graphs.clear(); +#endif + ggml_cuda_mmb_release_all(*this); + } + if (copy_event != nullptr) { CUDA_CHECK(cudaEventDestroy(copy_event)); } @@ -767,7 +791,9 @@ static enum ggml_status ggml_backend_cuda_buffer_init_tensor(ggml_backend_buffer if (padded_size > original_size) { ggml_cuda_set_device(ctx->device); - CUDA_CHECK(cudaMemset((char *)tensor->data + original_size, 0, padded_size - original_size)); + // not cudaMemset: on HIP a synchronous memset invalidates a graph capture running on any other thread + CUDA_CHECK(cudaMemsetAsync((char *) tensor->data + original_size, 0, padded_size - original_size, cudaStreamPerThread)); + CUDA_CHECK(cudaStreamSynchronize(cudaStreamPerThread)); } } return GGML_STATUS_SUCCESS; @@ -1841,7 +1867,8 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor const int cc = ggml_cuda_info().devices[ctx.device].cc; const int warp_size = ggml_cuda_info().devices[ctx.device].warp_size; - if (ggml_cuda_should_use_mmvf(src0->type, cc, src0->ne, src0->nb, ne11)) { + const int64_t mmvf_ncols = hint == GGML_HINT_EXACT_BATCH && src0->type == GGML_TYPE_BF16 ? 1 : ne11; + if (ggml_cuda_should_use_mmvf(src0->type, cc, src0->ne, src0->nb, mmvf_ncols)) { // The custom F16 vector kernel can be used over batched cuBLAS GEMM. // But this is only faster for GPUs without tensor cores or with a thin src0 matrix (particularly KQV in attention) ggml_cuda_mul_mat_vec_f(ctx, src0, src1, nullptr, dst); @@ -1869,6 +1896,10 @@ static void ggml_cuda_mul_mat(ggml_backend_cuda_context & ctx, const ggml_tensor ggml_cuda_mul_mat_vec_q(ctx, src0, src1, nullptr, dst); return; } + if (ggml_cuda_mmb_supported_mm(ctx, src0, src1, dst)) { + ggml_cuda_mul_mat_mmb(ctx, src0, src1, dst); + return; + } if (ggml_cuda_should_use_mmq(src0->type, cc, ne11, /*n_experts =*/ 0)) { ggml_cuda_mul_mat_q(ctx, src0, src1, nullptr, dst); return; @@ -1937,6 +1968,10 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * } } + if (ggml_cuda_mmb_supported_mmid(ctx, src0, src1, ids, dst)) { + ggml_cuda_mul_mat_id_mmb(ctx, src0, src1, ids, dst); + return; + } if (ggml_cuda_should_use_mmq(src0->type, cc, ne12, /*n_experts=*/ne02)) { ggml_cuda_mul_mat_q(ctx, src0, src1, ids, dst); return; @@ -2065,6 +2100,16 @@ static void ggml_cuda_mul_mat_id(ggml_backend_cuda_context & ctx, ggml_tensor * } static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct ggml_tensor * dst) { + // a tensor marked BF16-only was never written as F32: any reader outside the fused paths is a bug, not a fallback + if (ggml_cuda_mmb_marks_count(ctx) > 0 && dst->op != GGML_OP_MUL_MAT && dst->op != GGML_OP_MUL_MAT_ID && dst->op != GGML_OP_VIEW && + dst->op != GGML_OP_RESHAPE && dst->op != GGML_OP_PERMUTE && dst->op != GGML_OP_TRANSPOSE && dst->op != GGML_OP_NONE) { + for (int s = 0; s < GGML_MAX_SRC && dst->src[s]; ++s) { + const ggml_tensor * src = dst->src[s]; + if (ggml_cuda_mmb_is_bf16_only(ctx, src) || (src->view_src && ggml_cuda_mmb_is_bf16_only(ctx, src->view_src))) { + GGML_ABORT("MMB: %s(%s) reads BF16-only tensor %s through the unfused path", ggml_op_name(dst->op), dst->name, src->name); + } + } + } switch (dst->op) { case GGML_OP_ARGMAX: ggml_cuda_argmax(ctx, dst); @@ -2589,9 +2634,52 @@ static bool ggml_cuda_graph_check_compability(ggml_cgraph * cgraph) { } static const void * ggml_cuda_graph_get_key(ggml_cgraph * cgraph) { + // nodes[0] alone collides for graphs that share one meta arena (different topologies, or one topology at varying batch width). + // Mix in the node count, the last node and the shapes of three nodes. A hash collision only resets warmup, the node properties are still checked. + if (cgraph->n_nodes > 0) { + uintptr_t k = (uintptr_t) cgraph->nodes[0]; + k ^= (uintptr_t) cgraph->nodes[cgraph->n_nodes - 1] * (uintptr_t) 0x9E3779B97F4A7C15ull; + k ^= (uintptr_t) cgraph->n_nodes << 4; + const ggml_tensor * probe[3] = { + cgraph->nodes[0], + cgraph->nodes[cgraph->n_nodes / 2], + cgraph->nodes[cgraph->n_nodes - 1], + }; + for (int i = 0; i < 3; ++i) { + for (int d = 0; d < GGML_MAX_DIMS; ++d) { + k = (k ^ (uintptr_t) probe[i]->ne[d]) * (uintptr_t) 0x100000001B3ull; + } + } + return (const void *) k; + } return cgraph->nodes[0]; } +static ggml_cuda_graph::node_properties ggml_cuda_graph_node_props(const ggml_tensor * node) { + ggml_cuda_graph::node_properties prop = {}; + memcpy(&prop.node, node, sizeof(ggml_tensor)); + + for (int j = 0; j < GGML_MAX_SRC; ++j) { + if (node->src[j]) { + prop.node_src_data_ptrs[j] = node->src[j]->data; + memcpy(prop.node_src_ne[j], node->src[j]->ne, sizeof(prop.node_src_ne[j])); + memcpy(prop.node_src_nb[j], node->src[j]->nb, sizeof(prop.node_src_nb[j])); + } + } + + return prop; +} + +static bool ggml_cuda_graph_verify_uid() { +#ifndef NDEBUG + return true; +#else + static const bool verify = getenv("GGML_CUDA_GRAPH_VERIFY_UID") != nullptr; + return verify; +#endif +} + +// see docs/development/backend-scheduler.md static bool ggml_cuda_graph_update_required(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph) { bool res = false; @@ -2602,6 +2690,16 @@ static bool ggml_cuda_graph_update_required(ggml_backend_cuda_context * cuda_ctx cgraph->uid == graph->uid) { GGML_LOG_DEBUG("CUDA Graph id %zu reused\n", cgraph->uid); GGML_ASSERT((int)graph->node_props.size() == cgraph->n_nodes); + if (ggml_cuda_graph_verify_uid()) { + for (int i = 0; i < cgraph->n_nodes; i++) { + const ggml_cuda_graph::node_properties prop = ggml_cuda_graph_node_props(cgraph->nodes[i]); + if (memcmp(&graph->node_props[i], &prop, sizeof(prop)) != 0) { + GGML_LOG_ERROR("%s: node %d (%s) changed while cgraph->uid stayed %llu\n", + __func__, i, cgraph->nodes[i]->name, (unsigned long long) cgraph->uid); + GGML_ABORT("CUDA graph uid reused after node properties changed - see [TAG_CUDA_GRAPH_UID]"); + } + } + } return false; } @@ -2614,16 +2712,7 @@ static bool ggml_cuda_graph_update_required(ggml_backend_cuda_context * cuda_ctx } for (int i = 0; i < cgraph->n_nodes; i++) { - ggml_cuda_graph::node_properties prop = {}; - memcpy(&prop.node, cgraph->nodes[i], sizeof(ggml_tensor)); - - for (int j = 0; j < GGML_MAX_SRC; ++j) { - if (cgraph->nodes[i]->src[j]) { - prop.node_src_data_ptrs[j] = cgraph->nodes[i]->src[j]->data; - memcpy(prop.node_src_ne[j], cgraph->nodes[i]->src[j]->ne, sizeof(prop.node_src_ne[j])); - memcpy(prop.node_src_nb[j], cgraph->nodes[i]->src[j]->nb, sizeof(prop.node_src_nb[j])); - } - } + const ggml_cuda_graph::node_properties prop = ggml_cuda_graph_node_props(cgraph->nodes[i]); if (res || memcmp(&graph->node_props[i], &prop, sizeof(prop)) != 0) { graph->node_props[i] = prop; @@ -2973,7 +3062,8 @@ static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph, const int node_count, const int * out_nodes, const int out_count, - const bool is_topk_moe = false) { + const bool is_topk_moe = false, + const bool topk_moe_logits_copied = false) { auto nodes_overlap = [&](const ggml_tensor * a, const ggml_tensor * b) { const int64_t a_start = (int64_t) a->data; const int64_t a_end = a_start + ggml_backend_buft_get_alloc_size(a->buffer->buft, a); @@ -2991,7 +3081,7 @@ static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph, bool is_ok = true; // one block reads all logits before it writes, so logits may alias the out nodes const ggml_tensor * logits_may_alias = nullptr; - if (is_topk_moe && ggml_nrows(cgraph->nodes[node_idx]) <= TOPK_MOE_ROWS_PER_BLOCK) { + if (is_topk_moe && (topk_moe_logits_copied || ggml_nrows(cgraph->nodes[node_idx]) <= TOPK_MOE_ROWS_PER_BLOCK)) { logits_may_alias = cgraph->nodes[node_idx]->src[0]; } @@ -3032,6 +3122,274 @@ static bool ggml_cuda_check_fusion_memory_ranges(const ggml_cgraph * cgraph, return is_ok; } +// Runs the fused topk-moe. When the allocator placed weights/ids over the logits (they are dead after the gating +// activation in the unfused graph), the kernel reads a private copy of the logits instead of declining the fusion. +// Otherwise buffer placement picks between fused and unfused routing, which round differently, so the model output +// would change with allocation. +static bool ggml_cuda_try_topk_moe(ggml_backend_cuda_context & ctx, const ggml_cgraph * cgraph, int node_idx, + int node_count, const int * out_nodes, const ggml_tensor * logits, + ggml_tensor * weights, ggml_tensor * ids, const ggml_tensor * clamp, + const ggml_tensor * scale, const ggml_tensor * bias, + const ggml_cuda_topk_moe_args & args) { + if (ggml_cuda_check_fusion_memory_ranges(cgraph, node_idx, node_count, out_nodes, 2, /*is_topk_moe=*/true)) { + ggml_cuda_op_topk_moe(ctx, logits, weights, ids, clamp, scale, bias, args); + return true; + } + if (!ggml_is_contiguous(logits) || + !ggml_cuda_check_fusion_memory_ranges(cgraph, node_idx, node_count, out_nodes, 2, true, true)) { + return false; + } + ggml_cuda_pool_alloc copy(ctx.pool(), ggml_nbytes(logits)); + CUDA_CHECK(cudaMemcpyAsync(copy.get(), logits->data, ggml_nbytes(logits), cudaMemcpyDeviceToDevice, ctx.stream())); + ggml_tensor tmp = *logits; + tmp.data = copy.get(); + ggml_cuda_op_topk_moe(ctx, &tmp, weights, ids, clamp, scale, bias, args); + return true; +} + +// match the Qwen3.5 Gated DeltaNet decode chain +// concat(conv_state, qkv^T) -> cpy(conv state shift -> cache) -> get_rows(state cache) -> ssm_conv -> silu +// -> l2_norm(q) -> l2_norm(k) -> add(alpha, dt_bias) -> softplus -> mul(A) -> sigmoid(beta) +// -> gated_delta_net -> cpy(new state -> cache) -> rms_norm -> mul(norm weight) +// for a single token of a single sequence with S == 128, so that one kernel can replace the 14 nodes. +// Returns the number of nodes to skip after the concat (0 when the pattern does not match). +static int ggml_cuda_try_gdn_decode_fusion(const ggml_cgraph * cgraph, int node_idx, ggml_cuda_gdn_decode_args & args) { + const ggml_tensor * concat = cgraph->nodes[node_idx]; + if (concat->op != GGML_OP_CONCAT || concat->type != GGML_TYPE_F32 || ggml_get_op_params_i32(concat, 0) != 0) { + return 0; + } + const int64_t d_conv = concat->ne[0]; + const int64_t C = concat->ne[1]; + if (d_conv != 4 || concat->ne[2] != 1 || concat->ne[3] != 1 || !ggml_is_contiguous(concat)) { + return 0; + } + const ggml_tensor * conv_state = concat->src[0]; // [d_conv-1, C, 1] + const ggml_tensor * qkv_t = concat->src[1]; // [1, C, 1] + if (conv_state->type != GGML_TYPE_F32 || qkv_t->type != GGML_TYPE_F32 || + conv_state->ne[0] != d_conv - 1 || conv_state->ne[1] != C || conv_state->ne[2] != 1 || !ggml_is_contiguous(conv_state) || + qkv_t->ne[0] != 1 || qkv_t->ne[1] != C || qkv_t->ne[2] != 1 || qkv_t->nb[1] % sizeof(float) != 0) { + return 0; + } + + // collect the following compute nodes + constexpr int n_ops = 14; + const ggml_op expected[n_ops] = { + GGML_OP_CPY, GGML_OP_GET_ROWS, GGML_OP_SSM_CONV, GGML_OP_UNARY, GGML_OP_L2_NORM, GGML_OP_L2_NORM, + GGML_OP_ADD, GGML_OP_UNARY, GGML_OP_MUL, GGML_OP_UNARY, GGML_OP_GATED_DELTA_NET, GGML_OP_CPY, + GGML_OP_RMS_NORM, GGML_OP_MUL, + }; + const ggml_tensor * n[n_ops]; + int idx[n_ops]; + int found = 0; + int last = node_idx; + for (int j = node_idx + 1; j < cgraph->n_nodes && found < n_ops; ++j) { + const ggml_tensor * t = cgraph->nodes[j]; + if (ggml_cuda_is_view_or_noop(t)) { + continue; + } + if (t->op != expected[found] || t->type != GGML_TYPE_F32) { + return 0; + } + n[found] = t; + idx[found] = j; + last = j; + ++found; + } + if (found != n_ops) { + return 0; + } + // intermediates must not be graph outputs + if (concat->flags & GGML_TENSOR_FLAG_OUTPUT) { + return 0; + } + for (int k = 0; k < n_ops; ++k) { + if (k != 0 && k != 11 && k != 13 && (n[k]->flags & GGML_TENSOR_FLAG_OUTPUT)) { + return 0; + } + } + + const ggml_tensor * cpy_conv = n[0]; + const ggml_tensor * get_rows = n[1]; + const ggml_tensor * ssm_conv = n[2]; + const ggml_tensor * silu = n[3]; + const ggml_tensor * l2_q = n[4]; + const ggml_tensor * l2_k = n[5]; + const ggml_tensor * add = n[6]; + const ggml_tensor * softplus = n[7]; + const ggml_tensor * mul_a = n[8]; + const ggml_tensor * sigmoid = n[9]; + const ggml_tensor * gdn = n[10]; + const ggml_tensor * rms = n[12]; + const ggml_tensor * mul_w = n[13]; + + auto is_view_of = [](const ggml_tensor * t, const ggml_tensor * src, size_t offs) { + return t == src ? offs == 0 : (t->view_src == src && t->view_offs == offs); + }; + + // conv state shift: view of the concat (skipping the oldest column) copied into a contiguous cache row + { + const ggml_tensor * src = cpy_conv->src[0]; + const ggml_tensor * dst = cpy_conv->src[1]; + if (!is_view_of(src, concat, sizeof(float)) || src->ne[0] != d_conv - 1 || src->ne[1] != C || src->ne[2] != 1 || + src->nb[1] != concat->nb[1] || dst->type != GGML_TYPE_F32 || dst->data == nullptr || + ggml_nelements(dst) != (d_conv - 1) * C || !ggml_is_contiguous(dst)) { + return 0; + } + } + + // ssm_conv(concat, weight [d_conv, C]) without bias, then silu + if (ssm_conv->src[0] != concat || ssm_conv->src[2] != nullptr || ssm_conv->src[1]->type != GGML_TYPE_F32 || + ssm_conv->src[1]->ne[0] != d_conv || ssm_conv->src[1]->ne[1] != C || !ggml_is_contiguous(ssm_conv->src[1]) || + ssm_conv->ne[0] != C || ggml_nelements(ssm_conv) != C || + ggml_get_unary_op(silu) != GGML_UNARY_OP_SILU || silu->src[0] != ssm_conv || !ggml_is_contiguous(silu)) { + return 0; + } + + // recurrence shapes + const ggml_tensor * v_view = gdn->src[2]; + const int64_t S = v_view->ne[0]; + const int64_t H_v = v_view->ne[1]; + const int64_t H_k = l2_q->ne[1]; + if (S != 128 || H_k <= 0 || H_v % H_k != 0 || v_view->ne[2] != 1 || v_view->ne[3] != 1 || + 2 * S * H_k + S * H_v != C || ggml_get_op_params_i32(gdn, 0) != 1) { + return 0; + } + + // q/k/v are views into the silu output: [S, H, 1, 1] with rows of S floats + auto qkv_view_ok = [&](const ggml_tensor * t, int64_t H, size_t offs) { + return t->view_src == silu && t->view_offs == offs && t->ne[0] == S && t->ne[1] == H && t->ne[2] == 1 && + t->ne[3] == 1 && t->nb[0] == sizeof(float) && t->nb[1] == S * sizeof(float); + }; + if (!qkv_view_ok(l2_q->src[0], H_k, 0) || !qkv_view_ok(l2_k->src[0], H_k, S * H_k * sizeof(float)) || + !qkv_view_ok(v_view, H_v, 2 * S * H_k * sizeof(float)) || + !ggml_are_same_shape(l2_q, l2_q->src[0]) || !ggml_are_same_shape(l2_k, l2_k->src[0]) || + !ggml_is_contiguous(l2_q) || !ggml_is_contiguous(l2_k)) { + return 0; + } + float eps_q, eps_k; + memcpy(&eps_q, l2_q->op_params, sizeof(float)); + memcpy(&eps_k, l2_k->op_params, sizeof(float)); + if (eps_q != eps_k) { + return 0; + } + + // gate = softplus(alpha + dt_bias) * A, beta = sigmoid(beta): H_v scalars + if (ggml_nelements(add) != H_v || add->src[0]->type != GGML_TYPE_F32 || add->src[1]->type != GGML_TYPE_F32 || + ggml_nelements(add->src[0]) != H_v || ggml_nelements(add->src[1]) != H_v || + !ggml_is_contiguous(add->src[0]) || !ggml_is_contiguous(add->src[1]) || + ggml_get_unary_op(softplus) != GGML_UNARY_OP_SOFTPLUS || softplus->src[0] != add || + mul_a->src[0] != softplus || mul_a->src[1]->type != GGML_TYPE_F32 || ggml_nelements(mul_a->src[1]) != H_v || + !ggml_is_contiguous(mul_a->src[1]) || ggml_nelements(mul_a) != H_v || + ggml_get_unary_op(sigmoid) != GGML_UNARY_OP_SIGMOID || sigmoid->src[0]->type != GGML_TYPE_F32 || + ggml_nelements(sigmoid->src[0]) != H_v || !ggml_is_contiguous(sigmoid->src[0]) || ggml_nelements(sigmoid) != H_v) { + return 0; + } + + // gated_delta_net(q, k, v, gate, beta, state) with the state gathered from the cache by get_rows + if (gdn->src[0] != l2_q || gdn->src[1] != l2_k || !is_view_of(gdn->src[3], mul_a, 0) || !is_view_of(gdn->src[4], sigmoid, 0) || + gdn->src[3]->ne[0] != 1 || ggml_nelements(gdn->src[3]) != H_v || ggml_nelements(gdn->src[4]) != H_v || + !is_view_of(gdn->src[5], get_rows, 0) || ggml_nelements(get_rows) != S * S * H_v) { + return 0; + } + { + const ggml_tensor * cache = get_rows->src[0]; + const ggml_tensor * ids = get_rows->src[1]; + if (cache->type != GGML_TYPE_F32 || cache->ne[0] != S * S * H_v || cache->nb[0] != sizeof(float) || + cache->nb[1] % sizeof(float) != 0 || ids->type != GGML_TYPE_I32 || ggml_nelements(ids) != 1) { + return 0; + } + } + + // new state -> cache (same check as the standalone gated_delta_net + cpy fusion) + ggml_cuda_gated_delta_net_fused_cache state_cpy; + if (ggml_cuda_try_gdn_cache_fusion(cgraph, idx[10], state_cpy) != idx[11] - idx[10] || state_cpy.slot_stride != 0) { + return 0; + } + + // gated rms norm of the attention output [S, H_v] + { + const ggml_tensor * x = rms->src[0]; + if (x->view_src != gdn || x->view_offs != 0 || x->ne[0] != S || x->ne[1] != H_v || x->ne[2] != 1 || x->ne[3] != 1 || + x->nb[0] != sizeof(float) || x->nb[1] != S * sizeof(float) || !ggml_are_same_shape(rms, x) || + mul_w->src[0] != rms || mul_w->src[1]->type != GGML_TYPE_F32 || ggml_nelements(mul_w->src[1]) != S || + !ggml_is_contiguous(mul_w->src[1]) || !ggml_are_same_shape(mul_w, rms) || !ggml_is_contiguous(mul_w)) { + return 0; + } + } + + // The first kernel writes only the two cache rows: the state row may alias the gathered input row (in-place + // update, every block reads its head before writing it) and the conv row must not alias any external input + // of the range. The final output is written by a second kernel after all inputs have been consumed. + { + const ggml_tensor * out = cpy_conv->src[1]; + for (int j = node_idx; j <= last; ++j) { + const ggml_tensor * t = cgraph->nodes[j]; + for (int si = 0; si < GGML_MAX_SRC; ++si) { + const ggml_tensor * src = t->src[si]; + if (!src || src->op == GGML_OP_NONE || src == out) { + continue; + } + const int64_t a0 = (int64_t) out->data; + const int64_t a1 = a0 + ggml_backend_buft_get_alloc_size(out->buffer->buft, out); + const int64_t b0 = (int64_t) src->data; + const int64_t b1 = b0 + ggml_backend_buft_get_alloc_size(src->buffer->buft, src); + if ((b0 <= a0 && a0 < b1) || (a0 <= b0 && b0 < a1)) { + return 0; + } + } + } + } + + args.conv_state_in = (const float *) conv_state->data; + args.qkv = (const float *) qkv_t->data; + args.qkv_stride = qkv_t->nb[1] / sizeof(float); + args.conv_w = (const float *) ssm_conv->src[1]->data; + args.conv_bias = 0.0f; + args.conv_state_out = (float *) cpy_conv->src[1]->data; + args.alpha = (const float *) add->src[0]->data; + args.dt_bias = (const float *) add->src[1]->data; + args.ssm_a = (const float *) mul_a->src[1]->data; + args.beta = (const float *) sigmoid->src[0]->data; + args.eps_l2 = eps_q; + args.state_cache = (const float *) get_rows->src[0]->data; + args.state_ids = (const int32_t *) get_rows->src[1]->data; + args.state_row_stride = get_rows->src[0]->nb[1] / sizeof(float); + args.state_out = state_cpy.data; + args.scale = 1.0f / sqrtf((float) S); + args.norm_w = (const float *) mul_w->src[1]->data; + args.norm_w_tensor = mul_w->src[1]; + memcpy(&args.eps_rms, rms->op_params, sizeof(float)); + args.out = (float *) mul_w->data; + args.attn_out = nullptr; + args.S = S; + args.H_k = H_k; + args.H_v = H_v; + args.d_conv = d_conv; + + // Every skipped node except the three outputs must have all of its consumers inside this range. + // Include view nodes in the count even though they are not compute nodes. + for (int j = node_idx; j <= last; ++j) { + if (j == idx[0] || j == idx[11] || j == idx[13]) { + continue; + } + const ggml_tensor * candidate = cgraph->nodes[j]; + if (candidate->flags & GGML_TENSOR_FLAG_OUTPUT) { + return 0; + } + int internal_uses = 0; + for (int k = node_idx; k <= last; ++k) { + for (int src = 0; src < GGML_MAX_SRC; ++src) { + internal_uses += cgraph->nodes[k]->src[src] == candidate; + } + } + if (internal_uses != ggml_node_get_use_count(cgraph, j)) { + return 0; + } + } + + return last - node_idx; +} + // The long form spans 2*k + 1 nodes. ggml_can_fuse_subgraph() accepts at most // 31 nodes, so k <= 15; larger values use the per-operation path. static constexpr int MOE_WEIGHTED_REDUCTION_MAX_EXPERTS = 15; @@ -3429,6 +3787,509 @@ static bool ggml_cuda_can_fuse(const struct ggml_cgraph * cgraph, } // try and fuse nodes and return the number of nodes to skip +// qwen4exp hyper-connection stream mix chain starting at node i: +// sigmoid(gate) -> mul(xn, .) -> reshape -> view(stream 0) -> cont -> (view(stream c) -> add) x (hc-1) -> scale +// Fills args and ops (n_ops entries) and returns n_ops, or 0 when the nodes do not form the chain. +static int ggml_cuda_match_hc_mix(const ggml_cgraph * cgraph, const int i, ggml_cuda_hc_mix_args & args, enum ggml_op * ops) { + const ggml_tensor * node = cgraph->nodes[i]; + if (!(node->op == GGML_OP_UNARY && ggml_get_unary_op(node) == GGML_UNARY_OP_SIGMOID && node->type == GGML_TYPE_F32 && + i + 5 < cgraph->n_nodes && + cgraph->nodes[i + 1]->op == GGML_OP_MUL && cgraph->nodes[i + 2]->op == GGML_OP_RESHAPE && + cgraph->nodes[i + 3]->op == GGML_OP_VIEW && cgraph->nodes[i + 4]->op == GGML_OP_CONT)) { + return 0; + } + const ggml_tensor * gate = node->src[0]; + const ggml_tensor * mul = cgraph->nodes[i + 1]; + const ggml_tensor * view0 = cgraph->nodes[i + 3]; + const ggml_tensor * cont = cgraph->nodes[i + 4]; + const ggml_tensor * xn = mul->src[0] == node ? mul->src[1] : mul->src[0]; + + const int64_t n_embd = cont->ne[0]; + const int64_t n_tokens = cont->ne[1]; + const int64_t hc_dim = mul->ne[0]; + const int hc = n_embd > 0 ? (int) (hc_dim / n_embd) : 0; + const size_t row_size = hc_dim * sizeof(float); + + auto stream_view_ok = [&](const ggml_tensor * v, int c) { + return v->op == GGML_OP_VIEW && v->view_src == mul && v->type == GGML_TYPE_F32 && + v->view_offs == (size_t) c * n_embd * sizeof(float) && + v->ne[0] == n_embd && v->ne[1] == n_tokens && v->ne[2] == 1 && v->ne[3] == 1 && + v->nb[0] == sizeof(float) && v->nb[1] == row_size; + }; + + bool ok = (mul->src[0] == node || mul->src[1] == node) && xn != node && + mul->type == GGML_TYPE_F32 && xn->type == GGML_TYPE_F32 && gate->type == GGML_TYPE_F32 && + ggml_is_contiguous(xn) && ggml_is_contiguous(gate) && ggml_are_same_shape(xn, gate) && + ggml_are_same_shape(mul, xn) && ggml_are_same_shape(mul, node) && + mul->ne[2] == 1 && mul->ne[3] == 1 && ggml_nrows(mul) == n_tokens && + hc >= 2 && hc <= 16 && hc_dim == (int64_t) hc * n_embd && + cgraph->nodes[i + 2]->view_src == mul && cont->src[0] == view0 && stream_view_ok(view0, 0) && + cont->type == GGML_TYPE_F32 && ggml_is_contiguous(cont) && cont->ne[2] == 1 && cont->ne[3] == 1; + + const int n_ops = ok ? 5 + 2 * (hc - 1) + 1 : 0; + ok = ok && i + n_ops <= cgraph->n_nodes; + if (!ok) { + return 0; + } + + ops[0] = GGML_OP_UNARY; ops[1] = GGML_OP_MUL; ops[2] = GGML_OP_RESHAPE; ops[3] = GGML_OP_VIEW; ops[4] = GGML_OP_CONT; + const ggml_tensor * prev = cont; + for (int c = 1; ok && c < hc; ++c) { + const ggml_tensor * v = cgraph->nodes[i + 5 + 2 * (c - 1)]; + const ggml_tensor * add = cgraph->nodes[i + 6 + 2 * (c - 1)]; + ok = stream_view_ok(v, c) && add->op == GGML_OP_ADD && add->type == GGML_TYPE_F32 && + add->src[0] == prev && add->src[1] == v && ggml_are_same_shape(add, cont); + ops[5 + 2 * (c - 1)] = GGML_OP_VIEW; + ops[6 + 2 * (c - 1)] = GGML_OP_ADD; + prev = add; + } + ggml_tensor * scale = cgraph->nodes[i + n_ops - 1]; + ops[n_ops - 1] = GGML_OP_SCALE; + ok = ok && scale->op == GGML_OP_SCALE && scale->src[0] == prev && scale->type == GGML_TYPE_F32 && + ggml_is_contiguous(scale) && ggml_are_same_shape(scale, cont); + if (!ok) { + return 0; + } + + args.xn = xn; + args.gate = gate; + args.dst = scale; + args.hc = hc; + args.scale = ggml_get_op_params_f32(scale, 0); + args.bias = ggml_get_op_params_f32(scale, 1); + return n_ops; +} + +static int ggml_cuda_match_idx_relu_sum(const ggml_cgraph * g, int i, ggml_cuda_idx_relu_sum_args & a) { + if (i + 4 >= g->n_nodes) return 0; + const ggml_tensor * relu = g->nodes[i]; + if (relu->op != GGML_OP_UNARY || ggml_get_unary_op(relu) != GGML_UNARY_OP_RELU) return 0; + if (relu->type != GGML_TYPE_F32 || !ggml_is_contiguous(relu)) return 0; + const ggml_tensor * src = relu->src[0]; + if (!src || src->type != GGML_TYPE_F32 || !ggml_is_contiguous(src) || !ggml_are_same_shape(src, relu)) return 0; + const int64_t nb = relu->ne[0], H = relu->ne[1], nt = relu->ne[2], ns = relu->ne[3]; + if (H < 2 || H > 32 || nb < 64 || nt * ns < 64 || nt * ns > UINT16_MAX) return 0; + auto is_slice = [&](const ggml_tensor * v, int64_t h) { + return v->op == GGML_OP_VIEW && v->view_src == relu && v->type == GGML_TYPE_F32 && + v->ne[0] == nb && v->ne[1] == nt && v->ne[2] == ns && v->ne[3] == 1 && + v->nb[0] == sizeof(float) && v->nb[1] == relu->nb[2] && v->nb[2] == relu->nb[3] && + v->view_offs == (size_t) h * relu->nb[1]; + }; + const ggml_tensor * cont = g->nodes[i + 2]; + if (!is_slice(g->nodes[i + 1], 0) || cont->op != GGML_OP_CONT || cont->src[0] != g->nodes[i + 1]) return 0; + const ggml_tensor * prev = cont; + int k = i + 3; + for (int64_t h = 1; h < H; ++h) { + if (k + 1 >= g->n_nodes) return 0; + const ggml_tensor * v = g->nodes[k]; const ggml_tensor * add = g->nodes[k + 1]; + if (!is_slice(v, h) || add->op != GGML_OP_ADD || add->type != GGML_TYPE_F32) return 0; + if (add->src[0] != prev || add->src[1] != v || !ggml_are_same_shape(add, cont) || !ggml_is_contiguous(add)) return 0; + prev = add; k += 2; + } + const int count = k - i; + if (count > 32) return 0; + int indices[32]; ggml_op ops[32]; + for (int j = 0; j < count; ++j) { indices[j] = i + j; ops[j] = g->nodes[i + j]->op; } + const int output = i + count - 1; + if (!ggml_can_fuse_subgraph_ext(g, indices, count, ops, &output, 1)) return 0; + if (src->data && prev->data) { + const uintptr_t av = (uintptr_t) src->data, bv = (uintptr_t) prev->data; + const bool overlap = av <= bv ? bv - av < ggml_nbytes(src) : av - bv < ggml_nbytes(prev); + if (overlap) { return 0; } + } + a.score = src; a.dst = (ggml_tensor *) prev; a.heads = (int) H; + return count; +} + +// hc_mix match plus the fusion-closure check, as one call +static int ggml_cuda_hc_mix_closed(const ggml_cgraph * graph, int index, ggml_cuda_hc_mix_args & args) { + ggml_op ops[36]; + const int count = ggml_cuda_match_hc_mix(graph, index, args, ops); + if (count == 0) { return 0; } + int indices[36]; + for (int j = 0; j < count; ++j) { indices[j] = index + j; } + const int output = index + count - 1; + return ggml_can_fuse_subgraph_ext(graph, indices, count, ops, &output, 1) ? count : 0; +} + +static int ggml_cuda_match_hc_combine_norm(ggml_cgraph * cgraph, int i, + ggml_cuda_hc_combine_norm_args & args, int warp_size, bool check_alias) { + ggml_tensor * node = cgraph->nodes[i]; + args.single_block = false; + if (node->op == GGML_OP_SCALE && node->type == GGML_TYPE_F32 && i + 8 < cgraph->n_nodes && + cgraph->nodes[i + 1]->op == GGML_OP_UNARY && ggml_get_unary_op(cgraph->nodes[i + 1]) == GGML_UNARY_OP_SIGMOID && + cgraph->nodes[i + 1]->src[0] == node && + cgraph->nodes[i + 2]->op == GGML_OP_SCALE && cgraph->nodes[i + 2]->src[0] == cgraph->nodes[i + 1]) { + const ggml_tensor * scale1 = node; + const ggml_tensor * scale2 = cgraph->nodes[i + 2]; + const ggml_tensor * inject = scale1->src[0]; + + int j = i + 3; + while (j < cgraph->n_nodes && j < i + 7 && ggml_cuda_is_view_or_noop(cgraph->nodes[j]) && !ggml_is_empty(cgraph->nodes[j])) { + j++; + } + + if (j + 4 < cgraph->n_nodes && + cgraph->nodes[j]->op == GGML_OP_REPEAT && cgraph->nodes[j + 1]->op == GGML_OP_MUL && cgraph->nodes[j + 2]->op == GGML_OP_ADD && + cgraph->nodes[j + 3]->op == GGML_OP_RMS_NORM && cgraph->nodes[j + 4]->op == GGML_OP_MUL) { + const ggml_tensor * repeat = cgraph->nodes[j]; + const ggml_tensor * mul = cgraph->nodes[j + 1]; + ggml_tensor * add = cgraph->nodes[j + 2]; + const ggml_tensor * rms = cgraph->nodes[j + 3]; + ggml_tensor * mulg = cgraph->nodes[j + 4]; + const ggml_tensor * b = repeat->src[0]; + const ggml_tensor * w = mul->src[0] == repeat ? mul->src[1] : mul->src[0]; + const ggml_tensor * res = add->src[0] == mul ? add->src[1] : add->src[0]; + const ggml_tensor * gamma = mulg->src[0] == rms ? mulg->src[1] : mulg->src[0]; + + const int64_t n_embd = add->ne[0]; + const int64_t hc = add->ne[1]; + const int64_t n_tok = add->ne[2]; + + const bool w_ok = (w == scale2 || (ggml_cuda_is_view_or_noop(w) && w->view_src == scale2)) && + w->type == GGML_TYPE_F32 && ggml_is_contiguous(w) && ggml_nelements(w) == hc * n_tok && w->ne[0] == 1 && w->ne[1] == hc; + const bool b_ok = b->type == GGML_TYPE_F32 && ggml_is_contiguous(b) && ggml_nelements(b) == n_embd * n_tok && + b->ne[0] == n_embd && b->ne[1] == 1; + const bool ok = add->ne[3] == 1 && n_tok <= 65535 && + inject->type == GGML_TYPE_F32 && ggml_is_contiguous(inject) && ggml_nelements(inject) == hc * n_tok && + ggml_nrows(inject) == n_tok && + ggml_are_same_shape(scale1, inject) && ggml_are_same_shape(scale2, inject) && + (mul->src[0] == repeat || mul->src[1] == repeat) && w != repeat && + (add->src[0] == mul || add->src[1] == mul) && res != mul && + rms->src[0] == add && (mulg->src[0] == rms || mulg->src[1] == rms) && gamma != rms && + scale2->type == GGML_TYPE_F32 && repeat->type == GGML_TYPE_F32 && + mul->type == GGML_TYPE_F32 && add->type == GGML_TYPE_F32 && res->type == GGML_TYPE_F32 && + rms->type == GGML_TYPE_F32 && mulg->type == GGML_TYPE_F32 && gamma->type == GGML_TYPE_F32 && + ggml_are_same_shape(repeat, add) && ggml_are_same_shape(mul, add) && ggml_are_same_shape(res, add) && + ggml_are_same_shape(rms, add) && ggml_are_same_shape(mulg, add) && + ggml_is_contiguous(res) && ggml_is_contiguous(add) && ggml_is_contiguous(mulg) && ggml_is_contiguous(gamma) && + ggml_nelements(gamma) == n_embd * hc && gamma->ne[0] == n_embd && gamma->ne[1] == hc && w_ok && b_ok; + + if (ok) { + args.inject = inject; + args.residual = res; + args.block_out = b; + args.gamma = gamma; + args.out_res = add; + args.out_xn = mulg; + args.s1 = ggml_get_op_params_f32(scale1, 0); + args.b1 = ggml_get_op_params_f32(scale1, 1); + args.s2 = ggml_get_op_params_f32(scale2, 0); + args.b2 = ggml_get_op_params_f32(scale2, 1); + memcpy(&args.eps, rms->op_params, sizeof(float)); + + int k_inj = -1; + args.w_inject = nullptr; + args.out_inject = nullptr; + if (hc == 4 && n_tok >= 512) { + int k = j + 5; + while (k < cgraph->n_nodes && k < j + 9 && ggml_cuda_is_view_or_noop(cgraph->nodes[k])) { + k++; + } + ggml_tensor * mm = k < cgraph->n_nodes ? cgraph->nodes[k] : nullptr; + const ggml_tensor * wi = mm ? mm->src[0] : nullptr; + const ggml_tensor * xi = mm ? mm->src[1] : nullptr; + if (mm && mm->op == GGML_OP_MUL_MAT && (mm->flags & GGML_TENSOR_FLAG_COMPUTE) && + (xi == mulg || (xi->view_src == mulg && xi->view_offs == 0)) && ggml_is_contiguous(xi) && + wi->op == GGML_OP_NONE && wi->type == GGML_TYPE_F32 && ggml_is_contiguous(wi) && + wi->ne[0] == n_embd * hc && wi->ne[1] == hc && wi->ne[2] == 1 && wi->ne[3] == 1 && + mm->type == GGML_TYPE_F32 && ggml_is_contiguous(mm) && + mm->ne[0] == hc && mm->ne[1] == n_tok && mm->ne[2] == 1 && mm->ne[3] == 1) { + k_inj = k; + args.w_inject = wi; + args.out_inject = mm; + } + } + + bool alias_ok = true; + if (check_alias) { + // every block reads block_out and inject and writes its own stream of both outputs: the outputs + // must not overlap those, and the residual may only alias the first output in place + auto overlap = [](const ggml_tensor * p, const ggml_tensor * q) { + const uintptr_t p0 = (uintptr_t) p->data, p1 = p0 + ggml_backend_buft_get_alloc_size(p->buffer->buft, p); + const uintptr_t q0 = (uintptr_t) q->data, q1 = q0 + ggml_backend_buft_get_alloc_size(q->buffer->buft, q); + return p0 < q1 && q0 < p1; + }; + if (overlap(add, gamma) || overlap(mulg, gamma)) { + return 0; + } + const bool res_inplace = res->data == add->data && ggml_are_same_layout(res, add); + alias_ok = !overlap(add, b) && !overlap(add, inject) && (res_inplace || !overlap(add, res)) && + !overlap(mulg, b) && !overlap(mulg, inject) && !overlap(mulg, res) && !overlap(mulg, add); + if (k_inj >= 0) { + const ggml_tensor * mm = args.out_inject; + if (overlap(mm, add) || overlap(mm, mulg) || overlap(mm, res) || overlap(mm, b) || + overlap(mm, inject) || overlap(mm, gamma) || overlap(mm, args.w_inject)) { + k_inj = -1; + args.w_inject = nullptr; + args.out_inject = nullptr; + } + } + // The single-block kernel reads aliased inputs before writing either output. + if (!alias_ok && n_tok == 1 && hc <= 4) { + args.single_block = true; + alias_ok = true; + } + } + + int idxs[16]; + enum ggml_op ops[16]; + int count = 0; + const int last = k_inj >= 0 ? k_inj : j + 4; + for (int k = i; k <= last; ++k) { + const ggml_tensor * n = cgraph->nodes[k]; + if (ggml_cuda_is_view_or_noop(n) && n->view_src != scale2) { + continue; + } + idxs[count] = k; + ops[count] = n->op; + count++; + } + const int out_nodes[] = { j + 2, j + 4, k_inj }; + if (alias_ok && ggml_cuda_hc_combine_norm_supported(args, warp_size) && + ggml_can_fuse_subgraph_ext(cgraph, idxs, count, ops, out_nodes, k_inj >= 0 ? 3 : 2)) { + return last - i; + } + if (k_inj >= 0) { + args.w_inject = nullptr; + args.out_inject = nullptr; + const int out_nodes2[] = { j + 2, j + 4 }; + count = 0; + for (int k = i; k <= j + 4; ++k) { + const ggml_tensor * n = cgraph->nodes[k]; + if (ggml_cuda_is_view_or_noop(n) && n->view_src != scale2) { + continue; + } + idxs[count] = k; + ops[count] = n->op; + count++; + } + if (alias_ok && ggml_cuda_hc_combine_norm_supported(args, warp_size) && + ggml_can_fuse_subgraph_ext(cgraph, idxs, count, ops, out_nodes2, 2)) { + return j + 4 - i; + } + } + } + } + } + + return 0; +} + +// QSA indexer key pooling: get_rows(k, member_rows) -> reshape [128, 4, n_blocks, n_stream] -> four contiguous +// stream slices added in order -> scale 1/4. One kernel reads the four member rows and writes the mean. +struct ggml_cuda_get_rows_mean4_match { + const ggml_tensor * src = nullptr; + const ggml_tensor * ids = nullptr; + ggml_tensor * dst = nullptr; + int node_count = 14; +}; + +static bool ggml_cuda_match_get_rows_mean4( + const ggml_cgraph * graph, int i, ggml_cuda_get_rows_mean4_match & match) { + if (i < 0 || i + match.node_count > graph->n_nodes || graph->nodes[i]->op != GGML_OP_GET_ROWS) { + return false; + } + if (!ggml_can_fuse_subgraph(graph, i, + { GGML_OP_GET_ROWS, GGML_OP_RESHAPE, + GGML_OP_VIEW, GGML_OP_CONT, GGML_OP_VIEW, GGML_OP_CONT, GGML_OP_ADD, + GGML_OP_VIEW, GGML_OP_CONT, GGML_OP_ADD, GGML_OP_VIEW, GGML_OP_CONT, GGML_OP_ADD, + GGML_OP_SCALE }, { i + 13 })) { + return false; + } + const ggml_tensor * rows = graph->nodes[i]; + const ggml_tensor * src = rows->src[0]; + const ggml_tensor * ids = rows->src[1]; + const ggml_tensor * members = graph->nodes[i+1]; + ggml_tensor * dst = graph->nodes[i+13]; + if ((src->type != GGML_TYPE_F16 && src->type != GGML_TYPE_F32) || ids->type != GGML_TYPE_I32 || + src->ne[0] != 128 || src->ne[1] <= 0 || src->ne[2] <= 0 || src->nb[0] != ggml_type_size(src->type) || src->ne[3] != 1 || + ids->nb[0] != sizeof(int32_t) || ids->ne[0] <= 0 || ids->ne[0] % 4 != 0 || ids->ne[2] != 1 || ids->ne[3] != 1 || + ids->ne[1] != src->ne[2] || ids->ne[1] > 65535 || ids->ne[0]/4 > INT32_MAX) { + return false; + } + if (members->src[0] != rows || members->ne[0] != 128 || members->ne[1] != 4 || + members->ne[2] != ids->ne[0]/4 || members->ne[3] != ids->ne[1] || + !ggml_is_contiguous(members)) { + return false; + } + for (int n = 0; n < match.node_count; ++n) { + if (graph->nodes[i+n]->type != GGML_TYPE_F32) { + return false; + } + } + const int offsets[4] = { 2, 4, 7, 10 }; + const ggml_tensor * sum = nullptr; + for (int n = 0; n < 4; ++n) { + const ggml_tensor * view = graph->nodes[i+offsets[n]]; + const ggml_tensor * copy = graph->nodes[i+offsets[n]+1]; + if (view->src[0] != members || view->ne[0] != 128 || view->ne[1] != members->ne[2] || + view->ne[2] != members->ne[3] || view->ne[3] != 1 || view->nb[0] != sizeof(float) || + view->nb[1] != members->nb[2] || view->nb[2] != members->nb[3] || + view->view_offs != members->view_offs + n*members->nb[1] || + copy->src[0] != view || !ggml_are_same_shape(copy, view) || !ggml_is_contiguous(copy)) { + return false; + } + if (n == 0) { + sum = copy; + } else { + const ggml_tensor * add = graph->nodes[i+offsets[n]+2]; + if (add->src[0] != sum || add->src[1] != copy || !ggml_are_same_shape(add, copy)) { + return false; + } + sum = add; + } + } + if (dst->src[0] != sum || !ggml_are_same_shape(dst, sum) || !ggml_is_contiguous(dst) || + ggml_get_op_params_f32(dst, 0) != 0.25f || ggml_get_op_params_f32(dst, 1) != 0.0f) { + return false; + } + match.src = src; + match.ids = ids; + match.dst = dst; + return true; +} + +static bool ggml_cuda_mean4_ranges_overlap(const ggml_tensor * a, const ggml_tensor * b) { + const uintptr_t av = (uintptr_t) a->data; + const uintptr_t bv = (uintptr_t) b->data; + return av <= bv ? bv-av < ggml_nbytes(a) : av-bv < ggml_nbytes(b); +} + +// QSA complete-block selection (qwen4exp_select_complete_blocks): top-512 blocks -> ascending order -> cells with +// validity -> concat tail, as one kernel. Matches the 20/21-node subgraph the builder emits for 512 blocks / 2051 cells. +static int ggml_cuda_match_qsa_expand(const ggml_cgraph * graph, int first, ggml_cuda_qsa_expand_args & args) { + if (first + 19 >= graph->n_nodes) return 0; + if (graph->nodes[first]->op != GGML_OP_CONT || !graph->nodes[first]->src[0] || graph->nodes[first]->src[0]->op != GGML_OP_TOP_K) return 0; + static const bool debug = getenv("GGML_QSA_EXPAND_DEBUG") != nullptr; + static int debug_left = 3; + auto reject = [&](int line) { + if (debug && debug_left-- > 0) { + fprintf(stderr, "qsa_expand: reject at line %d, node %d (%s)\n", line, first, graph->nodes[first]->name); + for (int j = 0; j < 22 && first + j < graph->n_nodes; ++j) { + const ggml_tensor * t = graph->nodes[first + j]; + fprintf(stderr, " [%2d] %-10s %-6s [%lld,%lld,%lld,%lld] %s\n", j, ggml_op_name(t->op), ggml_type_name(t->type), + (long long) t->ne[0], (long long) t->ne[1], (long long) t->ne[2], (long long) t->ne[3], t->name); + } + } + return 0; + }; + const ggml_op expected[] = {GGML_OP_CONT, GGML_OP_VIEW, GGML_OP_CPY, GGML_OP_ARGSORT, + GGML_OP_GET_ROWS, GGML_OP_RESHAPE, GGML_OP_RESHAPE, GGML_OP_GET_ROWS, GGML_OP_RESHAPE, + GGML_OP_CPY, GGML_OP_SCALE, GGML_OP_VIEW, GGML_OP_GET_ROWS, GGML_OP_SCALE, GGML_OP_UNARY, + GGML_OP_MUL, GGML_OP_SCALE, GGML_OP_CPY, GGML_OP_RESHAPE, GGML_OP_VIEW, GGML_OP_CONCAT}; + const ggml_tensor * n[21]; + for (int j = 0; j < 19; ++j) { + n[j] = graph->nodes[first+j]; + if (n[j]->op != expected[j]) return reject(__LINE__); + } + // the tail may arrive through a view (a strip of the tail input) or a reshape of the whole input + const bool tail_view = graph->nodes[first+19]->op == GGML_OP_VIEW || graph->nodes[first+19]->op == GGML_OP_RESHAPE; + const int node_count = tail_view ? 21 : 20; + if (first + node_count > graph->n_nodes) return reject(__LINE__); + n[20] = graph->nodes[first+node_count-1]; + if (n[20]->op != GGML_OP_CONCAT) return reject(__LINE__); + // the tail the concat reads: the scheduler may substitute an input copy for the graph's own view node (the tail + // indices are a host input), so take the concat's actual source rather than nodes[first+19] + n[19] = n[20]->src[1]; + if (!n[19]) return reject(__LINE__); + auto source = [&](int node, int src, int producer) { return n[node]->src[src] == n[producer]; }; + for (const auto & edge : {std::array{1,0,0}, {2,0,0}, {3,0,2}, {4,0,1}, {4,1,3}, + {5,0,4}, {6,0,5}, {7,1,6}, {8,0,7}, {9,0,8}, {10,0,9}, {12,0,11}, {12,1,5}, + {13,0,12}, {14,0,13}, {15,0,10}, {15,1,14}, {16,0,15}, {17,0,16}, {18,0,17}, + {20,0,18}}) { + if (!source(edge[0], edge[1], edge[2])) { if (debug) fprintf(stderr, "qsa_expand: edge n[%d]->src[%d] != n[%d] (is %s)\n", edge[0], edge[1], edge[2], n[edge[0]]->src[edge[1]] ? n[edge[0]]->src[edge[1]]->name : "null"); return reject(__LINE__); } + } + const auto * blocks = n[0]->src[0]; + const auto * scores = n[11]->src[0]; + const auto * cells = n[7]->src[0]; + const auto * tail = n[19]; + if (!blocks || !scores || !cells || blocks->op != GGML_OP_TOP_K || blocks->src[0] != scores || + blocks->type != GGML_TYPE_I32 || scores->type != GGML_TYPE_F32 || cells->type != GGML_TYPE_I32 || + tail->type != GGML_TYPE_I32 || n[20]->type != GGML_TYPE_I32 || + ggml_get_op_params_i32(n[3],0) != GGML_SORT_ORDER_ASC || + ggml_get_unary_op(n[14]) != GGML_UNARY_OP_STEP || ggml_get_op_params_i32(n[20],0) != 0) return reject(__LINE__); + auto shape = [](const ggml_tensor * t, int64_t a, int64_t b, int64_t c=1, int64_t d=1) { + return t->ne[0]==a && t->ne[1]==b && t->ne[2]==c && t->ne[3]==d; + }; + const int64_t nq = blocks->ne[1], nb = scores->ne[0]; + if (nq < 1 || nq > INT_MAX || nb < 512 || nb > 16777216 || + !shape(blocks,512,nq) || !shape(scores,nb,nq) || !shape(cells,4,nb) || !shape(tail,3,nq) || + !shape(n[0],512,nq) || !shape(n[1],1,512,nq) || !shape(n[2],512,nq) || + !shape(n[3],512,nq) || !shape(n[4],1,512,nq) || !shape(n[5],512,nq) || + !shape(n[6],512*nq,1) || !shape(n[7],4,512*nq) || !shape(n[8],4,512,nq) || + !shape(n[9],4,512,nq) || !shape(n[10],4,512,nq) || !shape(n[11],1,nb,nq) || + !shape(n[12],1,512,nq) || !shape(n[13],1,512,nq) || !shape(n[14],1,512,nq) || + !shape(n[15],4,512,nq) || !shape(n[16],4,512,nq) || !shape(n[17],4,512,nq) || + !shape(n[18],2048,nq) || !shape(n[20],2051,nq)) return reject(__LINE__); + for (int j : {2,9,10,11,12,13,14,15,16}) if (n[j]->type != GGML_TYPE_F32) return reject(__LINE__); + for (int j : {0,1,3,4,5,6,7,8,17,18,19,20}) if (n[j]->type != GGML_TYPE_I32) return reject(__LINE__); + for (int j : {10,13,16}) { + if (ggml_get_op_params_f32(n[j],0) != 1.0f || ggml_get_op_params_f32(n[j],1) != (j==16 ? -1.0f : 1.0f)) return reject(__LINE__); + } + for (const auto * t : {blocks,scores,cells,tail,n[20]}) { + if (!ggml_is_contiguous(t) || t->nb[0] != 4) return reject(__LINE__); + } + if (n[1]->view_offs != 0 || n[11]->view_offs != 0 || + n[1]->nb[1] != 4 || n[1]->nb[2] != size_t(512)*4 || + n[11]->nb[1] != 4 || n[11]->nb[2] != size_t(nb)*4) return reject(__LINE__); + int indices[21], count=0; ggml_op ops[21]; + for (int j=0;jnodes[first+j]->op; + } + const int outputs[] = {first+2,first+9,first+17,first+18,first+node_count-1}; + if (!ggml_can_fuse_subgraph_ext(graph, indices, count, ops, outputs, 5)) return reject(__LINE__); + args = {blocks,scores,cells,tail,graph->nodes[first+2],graph->nodes[first+9],graph->nodes[first+17],graph->nodes[first+node_count-1]}; + return node_count-1; +} + +static bool ggml_cuda_qsa_expand_alias_ok(const ggml_cuda_qsa_expand_args & args) { + const ggml_tensor * outputs[] = {args.cast_blocks,args.cast_cells,args.expanded,args.output}; + for (int i=0;i<4;++i) { + for (const auto * input : {args.blocks,args.scores,args.cells,args.tail}) { + if (ggml_cuda_mean4_ranges_overlap(input,outputs[i])) return false; + } + for (int j=0;jsrc[0] || !t->src[1]) { + return false; + } + if (t->op == GGML_OP_MUL_MAT_ID) { + return t->src[2] && ggml_cuda_mmb_supported_mmid(ctx, t->src[0], t->src[1], t->src[2], t); + } + return t->op == GGML_OP_MUL_MAT && ggml_cuda_mmb_supported_mm(ctx, t->src[0], t->src[1], t); +} + +// ggml_cuda_try_fuse: node i was computed on its own and no later node is skipped +static constexpr int GGML_CUDA_FUSED_SELF = -1; + +// Returns a BF16 cache entry for the producer of t to fill when the next node that computes anything is an MMB GEMM reading t, else nullptr. +static uint16_t * ggml_cuda_bf16_for_mmb_reader(ggml_backend_cuda_context & ctx, const ggml_cgraph * cgraph, int i_last, + const ggml_tensor * t) { + if (t->type != GGML_TYPE_F32 || !ggml_is_contiguous(t)) { + return nullptr; + } + for (int k = i_last + 1; k < cgraph->n_nodes && k <= i_last + 4; ++k) { + ggml_tensor * n = cgraph->nodes[k]; + if (ggml_cuda_is_view_or_noop(n)) { + continue; + } + const ggml_tensor * x = n->src[1]; + if (n->op == GGML_OP_MUL_MAT && x && (x == t || (x->view_src == t && x->view_offs == 0)) && ggml_is_contiguous(x) && + ggml_nelements(x) == ggml_nelements(t) && ggml_cuda_mmb_supported_mm(ctx, n->src[0], x, n)) { + return ggml_cuda_mmb_cache_produce(ctx, x, (size_t) ggml_nelements(x)); + } + return nullptr; + } + return nullptr; +} + static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph * cgraph, int i) { static bool disable_fusion = getenv("GGML_CUDA_DISABLE_FUSION") != nullptr && std::atoi(getenv("GGML_CUDA_DISABLE_FUSION")); @@ -3436,16 +4297,232 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph return 0; } + // debugging aid: only fuse nodes with GGML_CUDA_FUSE_MIN_NODE <= i < GGML_CUDA_FUSE_MAX_NODE + static const int fuse_min_node = getenv("GGML_CUDA_FUSE_MIN_NODE") ? std::atoi(getenv("GGML_CUDA_FUSE_MIN_NODE")) : 0; + static const int fuse_max_node = getenv("GGML_CUDA_FUSE_MAX_NODE") ? std::atoi(getenv("GGML_CUDA_FUSE_MAX_NODE")) : INT_MAX; + if (i < fuse_min_node || i >= fuse_max_node) { + return 0; + } + ggml_tensor * node = cgraph->nodes[i]; + // routed gate/up + SwiGLU on the BF16 WMMA path from 512 tokens: ahead of the MMQ-based expert fusions below, + // which yield to MMB whenever it supports the GEMM + if (node->op == GGML_OP_MUL_MAT_ID && i + 2 < cgraph->n_nodes && cgraph->nodes[i + 1]->op == GGML_OP_MUL_MAT_ID && + cgraph->nodes[i + 2]->op == GGML_OP_GLU) { + ggml_tensor * glu = cgraph->nodes[i + 2]; + ggml_tensor * gate = glu->src[0]; + ggml_tensor * up = glu->src[1]; + const bool paired = (gate == node && up == cgraph->nodes[i + 1]) || (gate == cgraph->nodes[i + 1] && up == node); + if (paired && ggml_cuda_mmb_supported_glu(*cuda_ctx, gate->src[0], up->src[0], up->src[1], up->src[2], glu) && + ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_MUL_MAT_ID, GGML_OP_MUL_MAT_ID, GGML_OP_GLU }, { i + 2 })) { + const int out_nodes[] = { i + 2 }; + if (ggml_cuda_check_fusion_memory_ranges(cgraph, i, 3, out_nodes, 1)) { + ggml_cuda_mul_mat_id_mmb_glu(*cuda_ctx, gate->src[0], up->src[0], up->src[1], up->src[2], glu); + return 2; + } + } + } + + // qwen4exp prefill convolutions: the concat writes only the columns its history copies read, and one kernel + // replaces the tap chain (PLE) or the ssm_conv + SiLU (Gated DeltaNet). Both decisions come from the same matcher. + if (node->op == GGML_OP_CONCAT || node->op == GGML_OP_CONT) { + ggml_cuda_ple_conv_match pm; + if (node->op == GGML_OP_CONCAT && ggml_cuda_ple_conv_match_at_concat(cgraph, i, pm)) { + ggml_cuda_ple_conv_write_tail(*cuda_ctx, pm); + return 1; // the concat and the no-op view after it + } + if (node->op == GGML_OP_CONT && ggml_cuda_ple_conv_match_at_tap(cgraph, i, pm)) { + ggml_cuda_ple_conv_direct(*cuda_ctx, pm); + return pm.silu_idx - i; + } + } + if (node->op == GGML_OP_CONCAT || node->op == GGML_OP_SSM_CONV) { + ggml_cuda_gdn_conv_match gm; + if (node->op == GGML_OP_CONCAT && ggml_cuda_gdn_conv_match_at_concat(cgraph, i, gm)) { + ggml_cuda_gdn_conv_write_tail(*cuda_ctx, gm); + return 1; // the concat and the no-op view after it + } + if (node->op == GGML_OP_SSM_CONV && ggml_cuda_gdn_conv_match_at_conv(cgraph, i, gm)) { + ggml_cuda_gdn_conv_direct(*cuda_ctx, gm); + return 1; + } + } + if (node->op == GGML_OP_UNARY && GGML_CUDA_CC_IS_RDNA3_5(ggml_cuda_info().devices[cuda_ctx->device].cc)) { + ggml_cuda_idx_relu_sum_args args; + const int count = ggml_cuda_match_idx_relu_sum(cgraph, i, args); + if (count > 0) { + ggml_cuda_op_idx_relu_sum(*cuda_ctx, args); + return count - 1; + } + } + + + + // RDNA3.5 decode: consecutive single-column MUL_MATs that read the same activation vector (plain Q8_0/F32 + // matvecs, or a [mul_mat, mul_mat, glu] gate/up pair) are launched as one grouped kernel. The segments are + // mutually independent, so the launch position of the first one is valid for all of them. + if (getenv("GGML_CUDA_DISABLE_MMV_GROUP") == nullptr && + node->op == GGML_OP_MUL_MAT && node->src[2] == nullptr && ggml_cuda_mmv_group_seg_ok(node->src[0], node->src[1])) { + const ggml_tensor * y = node->src[1]; + ggml_cuda_mmv_group_seg segs[GGML_CUDA_MMV_GROUP_MAX]; + int out_nodes[GGML_CUDA_MMV_GROUP_MAX]; + int nseg = 0; + int j = i; + int j_end = i; // index of the last node that belongs to the group + while (j < cgraph->n_nodes && nseg < GGML_CUDA_MMV_GROUP_MAX) { + ggml_tensor * n = cgraph->nodes[j]; + if (j > i && ggml_cuda_is_view_or_noop(n)) { + j++; + continue; + } + const bool is_mmv = n->op == GGML_OP_MUL_MAT && n->src[2] == nullptr && n->src[1] == y && + (n->flags & GGML_TENSOR_FLAG_COMPUTE) != 0 && n->type == GGML_TYPE_F32 && ggml_is_contiguous(n) && + ggml_cuda_mmv_group_seg_ok(n->src[0], y); + if (!is_mmv) { + break; + } + // gate/up pair followed by the GLU that combines them + if (j + 2 < cgraph->n_nodes && cgraph->nodes[j + 1]->op == GGML_OP_MUL_MAT && cgraph->nodes[j + 2]->op == GGML_OP_GLU) { + ggml_tensor * n2 = cgraph->nodes[j + 1]; + ggml_tensor * glu = cgraph->nodes[j + 2]; + ggml_tensor * gate_n = nullptr; + ggml_tensor * up_n = nullptr; + if (glu->src[0] == n && glu->src[1] == n2) { + gate_n = n; + up_n = n2; + } else if (glu->src[0] == n2 && glu->src[1] == n) { + gate_n = n2; + up_n = n; + } + if (gate_n && n2->src[2] == nullptr && n2->src[1] == y && (n2->flags & GGML_TENSOR_FLAG_COMPUTE) != 0 && + up_n->src[0]->type == GGML_TYPE_Q8_0 && gate_n->src[0]->type == GGML_TYPE_Q8_0 && + ggml_are_same_shape(up_n->src[0], gate_n->src[0]) && ggml_cuda_mmv_group_seg_ok(n2->src[0], y) && + glu->type == GGML_TYPE_F32 && ggml_is_contiguous(glu) && ggml_nelements(glu) == up_n->src[0]->ne[1] && + ggml_cuda_should_fuse_mul_mat(up_n, gate_n, glu) && + ggml_can_fuse_subgraph(cgraph, j, { GGML_OP_MUL_MAT, GGML_OP_MUL_MAT, GGML_OP_GLU }, { j + 2 })) { + segs[nseg].w = up_n->src[0]; + segs[nseg].gate = gate_n->src[0]; + segs[nseg].dst = glu; + segs[nseg].glu_op = ggml_get_glu_op(glu); + segs[nseg].glu_limit = ggml_get_op_params_f32(glu, 3); + out_nodes[nseg++] = j + 2; + j_end = j + 2; + j += 3; + continue; + } + } + segs[nseg].w = n->src[0]; + segs[nseg].dst = n; + out_nodes[nseg++] = j; + j_end = j; + j++; + } + if (nseg >= 2 && ggml_cuda_check_fusion_memory_ranges(cgraph, i, j_end - i + 1, out_nodes, nseg)) { + ggml_cuda_mmv_group(*cuda_ctx, y, segs, nseg); + return j_end - i; + } + } + + if (node->op == GGML_OP_L2_NORM && i + 2 < cgraph->n_nodes && + GGML_CUDA_CC_IS_AMD(ggml_cuda_info().devices[cuda_ctx->device].cc)) { + ggml_tensor * k_view = cgraph->nodes[i + 1]; + ggml_tensor * k_norm = cgraph->nodes[i + 2]; + const ggml_tensor * q_view = node->src[0]; + const int out_nodes[] = { i, i + 2 }; + const bool graph_ok = k_view->op == GGML_OP_VIEW && k_norm->op == GGML_OP_L2_NORM && + k_norm->src[0] == k_view && (k_view->flags & GGML_TENSOR_FLAG_COMPUTE) != 0 && + (k_norm->flags & GGML_TENSOR_FLAG_COMPUTE) != 0 && ggml_node_get_use_count(cgraph, i + 1) == 1 && + (k_view->flags & GGML_TENSOR_FLAG_OUTPUT) == 0; + const bool shape_ok = q_view && q_view->op == GGML_OP_VIEW && q_view->view_src && + q_view->view_src == k_view->view_src && node->ne[0] == 128 && ggml_are_same_shape(node, k_norm) && + ggml_are_same_shape(q_view, k_view) && q_view->type == GGML_TYPE_F32 && k_view->type == GGML_TYPE_F32 && + node->type == GGML_TYPE_F32 && k_norm->type == GGML_TYPE_F32; + const bool layout_ok = shape_ok && ggml_is_contiguous_rows(q_view) && ggml_is_contiguous_rows(k_view) && + ggml_is_contiguous(node) && ggml_is_contiguous(k_norm) && + q_view->nb[1] == k_view->nb[1] && q_view->nb[2] == k_view->nb[2] && q_view->nb[3] == k_view->nb[3] && + k_view->view_offs == q_view->view_offs + ggml_row_size(q_view->type, q_view->ne[0])*q_view->ne[1] && + q_view->nb[2] > ggml_row_size(q_view->type, q_view->ne[0])*q_view->ne[1] && q_view->ne[2] <= UINT16_MAX/2; + + if (graph_ok && layout_ok && ggml_cuda_check_fusion_memory_ranges(cgraph, i, 3, out_nodes, 2)) { + ggml_cuda_op_l2_norm_dual_s128(*cuda_ctx, node, k_norm); + return 2; + } + } + + // Qwen3.5 Gated DeltaNet decode step: conv + norms + gate + recurrence + state writes + gated rms norm + if (node->op == GGML_OP_CONCAT && GGML_CUDA_CC_IS_AMD(ggml_cuda_info().devices[cuda_ctx->device].cc)) { + ggml_cuda_gdn_decode_args args; + const int nodes_to_skip = ggml_cuda_try_gdn_decode_fusion(cgraph, i, args); + if (nodes_to_skip > 0) { + // qwen4exp: the gated norm output only feeds sigmoid(z) * . -> ssm_out matvec. Apply the norm and the + // gate in the prologue of the fused-quantize matvec instead of two more kernels. + { + const int norm_idx = i + nodes_to_skip; // MUL(rms_norm, w): the last node of the chain + ggml_tensor * norm_out = cgraph->nodes[norm_idx]; + int j = norm_idx + 1; + auto next_compute = [&](int & k) -> ggml_tensor * { + while (k < cgraph->n_nodes && ggml_cuda_is_view_or_noop(cgraph->nodes[k])) { + k++; + } + return k < cgraph->n_nodes ? cgraph->nodes[k] : nullptr; + }; + ggml_tensor * sig = next_compute(j); + const int sig_idx = j; + j++; + ggml_tensor * mul = next_compute(j); + const int mul_idx = j; + j++; + ggml_tensor * mm = next_compute(j); + const int mm_idx = j; + if (getenv("GGML_CUDA_DISABLE_GDN_GATE") == nullptr && + sig && mul && mm && sig->op == GGML_OP_UNARY && ggml_get_unary_op(sig) == GGML_UNARY_OP_SIGMOID && + mul->op == GGML_OP_MUL && mm->op == GGML_OP_MUL_MAT && mm->src[2] == nullptr && + (mm->src[1] == mul || (ggml_cuda_is_view_or_noop(mm->src[1]) && mm->src[1]->view_src == mul && + ggml_is_contiguous(mm->src[1]) && ggml_nelements(mm->src[1]) == ggml_nelements(mul))) && + ((mul->src[0] == sig && mul->src[1] == norm_out) || (mul->src[0] == norm_out && mul->src[1] == sig)) && + sig->type == GGML_TYPE_F32 && mul->type == GGML_TYPE_F32 && ggml_is_contiguous(mul) && + ggml_are_same_shape(sig, norm_out) && ggml_are_same_shape(mul, norm_out) && + ggml_nelements(norm_out) == args.S * args.H_v && args.S == 128 && + ggml_node_get_use_count(cgraph, norm_idx) == 1 && ggml_node_get_use_count(cgraph, sig_idx) == 1 && + ggml_node_get_use_count(cgraph, mul_idx) == 1 && + (norm_out->flags & GGML_TENSOR_FLAG_OUTPUT) == 0 && (sig->flags & GGML_TENSOR_FLAG_OUTPUT) == 0 && + (mul->flags & GGML_TENSOR_FLAG_OUTPUT) == 0 && + ggml_cuda_mul_mat_vec_q_fq_gdn_gate_ok(*cuda_ctx, mm, sig->src[0])) { + ggml_cuda_pool_alloc attn(cuda_ctx->pool(), args.S * args.H_v); + ggml_cuda_op_gdn_decode_fused_prenorm(*cuda_ctx, args, attn.get()); + ggml_cuda_mul_mat_vec_q_fq_gdn_gate(*cuda_ctx, mm, attn.get(), sig->src[0], args.norm_w_tensor, args.eps_rms); + return mm_idx - i; + } + } + ggml_cuda_op_gdn_decode_fused(*cuda_ctx, args); + return nodes_to_skip; + } + } + if (node->op == GGML_OP_MUL) { ggml_cuda_moe_weighted_reduction_match match; if (ggml_cuda_match_moe_weighted_reduction(cgraph, i, match)) { - const int output_idx = i + match.node_count - 1; - if (ggml_cuda_check_fusion_memory_ranges(cgraph, i, match.node_count, &output_idx, 1)) { + int count = match.node_count; + const ggml_tensor * merge = nullptr; + { // shared-expert merge: ADD(reduction, ffn_shexp_gated) right after, with the reduction used only there + const int nadd = i + match.node_count; + if (ggml_cuda_mmb_blk16() && nadd < cgraph->n_nodes && ggml_node_has_n_uses(cgraph, i + match.node_count - 1, 1)) { + const ggml_tensor * add = cgraph->nodes[nadd]; + if (add->op == GGML_OP_ADD && add->type == GGML_TYPE_F32 && ggml_is_contiguous(add) && + ggml_are_same_shape(add, match.dst) && ggml_cuda_mmb_is_bf16_only(*cuda_ctx, add)) { + const ggml_tensor * o = add->src[0] == match.dst ? add->src[1] : (add->src[1] == match.dst ? add->src[0] : nullptr); + if (o && o->type == GGML_TYPE_F32 && ggml_is_contiguous(o) && ggml_are_same_shape(o, match.dst)) { + merge = o; match.dst = const_cast(add); count = match.node_count + 1; + } + } + } + } + const int output_idx = i + count - 1; + if (ggml_cuda_check_fusion_memory_ranges(cgraph, i, count, &output_idx, 1)) { ggml_cuda_op_moe_weighted_reduction( - *cuda_ctx, match.experts, match.expert_scale, match.weights, match.dst); - return match.node_count - 1; + *cuda_ctx, match.experts, match.expert_scale, match.weights, match.dst, merge); + return count - 1; } } } @@ -3455,6 +4532,25 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph ggml_cuda_gated_delta_net_fused_cache fused_state_cpy; const int nodes_to_skip = ggml_cuda_try_gdn_cache_fusion(cgraph, i, fused_state_cpy); if (nodes_to_skip > 0) { + bool uses_ok = true; + for (int j = i; uses_ok && j <= i + nodes_to_skip; ++j) { + if (j == i + nodes_to_skip) { + continue; + } + const ggml_tensor * candidate = cgraph->nodes[j]; + int internal_uses = 0; + for (int k = i; k <= i + nodes_to_skip; ++k) { + for (int src = 0; src < GGML_MAX_SRC; ++src) { + internal_uses += cgraph->nodes[k]->src[src] == candidate; + } + } + // The attention-score view is the one expected consumer after the snapshot copy. + const int allowed_external = j == i ? 1 : 0; + uses_ok = internal_uses + allowed_external == ggml_node_get_use_count(cgraph, j); + } + if (!uses_ok) { + return 0; + } #ifdef GGML_CUDA_DEBUG GGML_LOG_INFO("%s: fused gated_delta_net snapshot copies for %s (skipped %d nodes)\n", __func__, node->name, nodes_to_skip); @@ -3464,6 +4560,51 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph } } + if (node->op == GGML_OP_GLU && i + 1 < cgraph->n_nodes && ggml_get_glu_op(node) == GGML_GLU_OP_SWIGLU && node->src[1] && + !ggml_cuda_mmb_is_bf16_only(*cuda_ctx, node) && !ggml_cuda_mmb_claims(*cuda_ctx, cgraph->nodes[i + 1])) { + ggml_tensor * next = cgraph->nodes[i + 1]; + const bool has_ids = next->op == GGML_OP_MUL_MAT_ID; + if (has_ids || next->op == GGML_OP_MUL_MAT) { + const ggml_tensor * weights = next->src[0]; + const ggml_tensor * ids = has_ids ? next->src[2] : nullptr; + const ggml_tensor * gate = node->src[0]; + const ggml_tensor * up = node->src[1]; + const int cc = ggml_cuda_info().devices[cuda_ctx->device].cc; + const int64_t mmq_cols = has_ids ? node->ne[2] : node->ne[1]; + const int64_t n_experts = has_ids ? weights->ne[2] : 0; + const bool bad_padding_clear = ggml_backend_buffer_get_usage(weights->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE && + ggml_nbytes(weights) != ggml_backend_buffer_get_alloc_size(weights->buffer, weights) && weights->view_src; + + const bool target_qwen36 = node->ne[0] == 512 && next->ne[0] == 2048 && + (!has_ids || (gate->ne[1] == 8 && weights->ne[2] == 256)); + const bool target_qwen4exp = node->ne[0] == 640 && next->ne[0] == 2560 && + (!has_ids || (gate->ne[1] == 10 && weights->ne[2] == 512)); + const bool target_shape = target_qwen36 || target_qwen4exp; + const bool weight_type_ok = weights->type == GGML_TYPE_Q8_0 || + (target_qwen4exp && weights->type == GGML_TYPE_IQ4_NL); + const bool shape_ok = target_shape && next->src[1] == node && gate->type == GGML_TYPE_F32 && up->type == GGML_TYPE_F32 && + node->type == GGML_TYPE_F32 && next->type == GGML_TYPE_F32 && weight_type_ok && + ggml_are_same_shape(gate, up) && ggml_are_same_shape(gate, node) && gate->nb[0] == sizeof(float) && up->nb[0] == sizeof(float) && + gate->ne[3] == 1 && (has_ids ? gate->ne[1] == ids->ne[0] && gate->ne[2] == ids->ne[1] : gate->ne[2] == 1); + const bool use_mmq = shape_ok && !bad_padding_clear && GGML_CUDA_CC_IS_RDNA3_5(cc) && + ggml_cuda_should_use_mmq(weights->type, cc, mmq_cols, n_experts); + + if (use_mmq && ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_GLU, next->op }, { i + 1 })) { + ggml_cuda_mul_mat_q_swiglu(*cuda_ctx, weights, ids, next, node); + return 1; + } + } + } + + if (node->op == GGML_OP_GLU && ggml_get_glu_op(node) == GGML_GLU_OP_SWIGLU && node->type == GGML_TYPE_F32 && + ggml_is_contiguous(node) && !ggml_cuda_mmb_is_bf16_only(*cuda_ctx, node)) { + uint16_t * dst16 = ggml_cuda_bf16_for_mmb_reader(*cuda_ctx, cgraph, i, node); + if (dst16) { + ggml_cuda_op_swiglu(*cuda_ctx, node, dst16); + return GGML_CUDA_FUSED_SELF; + } + } + //topk-moe if (cgraph->nodes[i]->op == GGML_OP_UNARY || cgraph->nodes[i]->op == GGML_OP_SOFT_MAX || cgraph->nodes[i]->op == GGML_OP_ARGSORT) { @@ -3517,8 +4658,8 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph if (ggml_can_fuse_subgraph(cgraph, i, ops.size(), ops.data(), out_nodes, 2) && ggml_cuda_should_use_topk_moe(node, logits, weights, ids) && - ggml_cuda_check_fusion_memory_ranges(cgraph, i, ops.size(), out_nodes, 2, /*is_topk_moe=*/true)) { - ggml_cuda_op_topk_moe(*cuda_ctx, logits, weights, ids, clamp, scale, bias, args); + ggml_cuda_try_topk_moe(*cuda_ctx, cgraph, i, ops.size(), out_nodes, logits, weights, ids, clamp, + scale, bias, args)) { return ops.size() - 1; } } else if (!args.norm && !args.prob_bias) { @@ -3532,8 +4673,8 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph int out_nodes[2] = { i + 1, i + 5 }; if (ggml_can_fuse_subgraph(cgraph, i, ops.size(), ops.data(), out_nodes, 2) && ggml_cuda_should_use_topk_moe(softmax, logits, weights, ids) && - ggml_cuda_check_fusion_memory_ranges(cgraph, i, ops.size(), out_nodes, 2, /*is_topk_moe=*/true)) { - ggml_cuda_op_topk_moe(*cuda_ctx, logits, weights, ids, clamp, scale, bias, args); + ggml_cuda_try_topk_moe(*cuda_ctx, cgraph, i, ops.size(), out_nodes, logits, weights, ids, clamp, + scale, bias, args)) { return ops.size() - 1; } } @@ -3596,6 +4737,191 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph } } + // Shared expert gate for one token: mul_mat(w [n_embd, 1], y) -> sigmoid -> mul -> add [-> add] + // (Qwen3.5 adds the residual in a second add, qwen4exp folds the result into the hyper-connection instead) + if (node->op == GGML_OP_MUL_MAT && i + 3 < cgraph->n_nodes && node->type == GGML_TYPE_F32 && + node->src[0]->type == GGML_TYPE_F32 && node->src[1]->type == GGML_TYPE_F32 && + node->src[0]->ne[0] % 2 == 0 && ggml_nelements(node->src[0]) == node->src[0]->ne[0] && + ggml_nelements(node->src[1]) == node->src[0]->ne[0] && + ggml_is_contiguous(node->src[0]) && ggml_is_contiguous(node->src[1]) && ggml_nelements(node) == 1 && + cgraph->nodes[i + 1]->op == GGML_OP_UNARY && ggml_get_unary_op(cgraph->nodes[i + 1]) == GGML_UNARY_OP_SIGMOID && + cgraph->nodes[i + 1]->src[0] == node && cgraph->nodes[i + 2]->op == GGML_OP_MUL && cgraph->nodes[i + 3]->op == GGML_OP_ADD && + GGML_CUDA_CC_IS_AMD(ggml_cuda_info().devices[cuda_ctx->device].cc)) { + // the fused kernel replays mul_mat_vec_f: only valid where mmvf would pick that block size + const int64_t ncols = node->src[0]->ne[0]; + const int warp_size = ggml_cuda_info().devices[cuda_ctx->device].warp_size; + int64_t block_size_best = warp_size; + int64_t niter_best = (ncols + 2*warp_size - 1) / (2*warp_size); + for (int64_t block_size = 2*warp_size; block_size <= 256; block_size += warp_size) { + const int64_t niter = (ncols + 2*block_size - 1) / (2*block_size); + if (niter < niter_best) { + niter_best = niter; + block_size_best = block_size; + } + } + + const bool two_adds = i + 4 < cgraph->n_nodes && cgraph->nodes[i + 4]->op == GGML_OP_ADD && + cgraph->nodes[i + 4]->src[0] == cgraph->nodes[i + 3]; + const bool subgraph_ok = block_size_best == 256 && warp_size == 32 && (two_adds ? + ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_MUL_MAT, GGML_OP_UNARY, GGML_OP_MUL, GGML_OP_ADD, GGML_OP_ADD }, { i + 4 }) : + ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_MUL_MAT, GGML_OP_UNARY, GGML_OP_MUL, GGML_OP_ADD }, { i + 3 })); + + if (subgraph_ok) { + const ggml_tensor * sigmoid = cgraph->nodes[i + 1]; + ggml_tensor * mul = cgraph->nodes[i + 2]; + ggml_tensor * add0 = cgraph->nodes[i + 3]; + ggml_tensor * add1 = two_adds ? cgraph->nodes[i + 4] : nullptr; + ggml_tensor * out = two_adds ? add1 : add0; + const ggml_tensor * gate = mul->src[0] == sigmoid ? mul->src[0] : mul->src[1]; + const ggml_tensor * src = gate == mul->src[0] ? mul->src[1] : mul->src[0]; + const ggml_tensor * other = add0->src[0] == mul ? add0->src[1] : add0->src[0]; + const ggml_tensor * residual = two_adds ? (add1->src[0] == add0 ? add1->src[1] : add1->src[0]) : nullptr; + const uintptr_t dst_begin = (uintptr_t) out->data; + const uintptr_t dst_end = dst_begin + ggml_backend_buft_get_alloc_size(out->buffer->buft, out); + const auto safe_output_alias = [&](const ggml_tensor * input) { + const uintptr_t begin = (uintptr_t) input->data; + const uintptr_t end = begin + ggml_backend_buft_get_alloc_size(input->buffer->buft, input); + return !(begin < dst_end && dst_begin < end) || (input->data == out->data && ggml_are_same_layout(input, out)); + }; + const bool residual_ok = !two_adds || (residual != mul && residual != add0 && residual->type == GGML_TYPE_F32 && + ggml_are_same_shape(mul, residual) && ggml_is_contiguous(residual) && safe_output_alias(residual)); + if (gate == sigmoid && ggml_nelements(sigmoid) == 1 && (add0->src[0] == mul || add0->src[1] == mul) && + other != mul && residual_ok && + mul->type == GGML_TYPE_F32 && src->type == GGML_TYPE_F32 && other->type == GGML_TYPE_F32 && + mul->ne[0] == ncols && ggml_nelements(mul) == ncols && ggml_are_same_shape(mul, src) && ggml_are_same_shape(mul, other) && + ggml_are_same_shape(mul, out) && + ggml_is_contiguous(src) && ggml_is_contiguous(other) && ggml_is_contiguous(out) && + safe_output_alias(src) && safe_output_alias(other) && safe_output_alias(node->src[1])) { + ggml_cuda_op_shared_gate_mul_add(*cuda_ctx, node, src, other, residual, out); + return two_adds ? 4 : 3; + } + } + } + + if (node->op == GGML_OP_MUL && node->ne[0] == 2048 && i + 2 < cgraph->n_nodes && + ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_MUL, GGML_OP_ADD, GGML_OP_ADD }, { i + 2 })) { + ggml_tensor * add0 = cgraph->nodes[i + 1]; + ggml_tensor * add1 = cgraph->nodes[i + 2]; + const ggml_tensor * src = ggml_are_same_shape(node, node->src[0]) ? node->src[0] : node->src[1]; + const ggml_tensor * gate = src == node->src[0] ? node->src[1] : node->src[0]; + const ggml_tensor * other = add0->src[0]; + const ggml_tensor * residual = add1->src[1]; + const uintptr_t gate_begin = (uintptr_t) gate->data; + const uintptr_t gate_end = gate_begin + ggml_backend_buft_get_alloc_size(gate->buffer->buft, gate); + const uintptr_t dst_begin = (uintptr_t) add1->data; + const uintptr_t dst_end = dst_begin + ggml_backend_buft_get_alloc_size(add1->buffer->buft, add1); + const bool gate_overlap = gate_begin < dst_end && dst_begin < gate_end; + const auto safe_output_alias = [&](const ggml_tensor * input) { + const uintptr_t begin = (uintptr_t) input->data; + const uintptr_t end = begin + ggml_backend_buft_get_alloc_size(input->buffer->buft, input); + return !(begin < dst_end && dst_begin < end) || (input->data == add1->data && ggml_are_same_layout(input, add1)); + }; + + if (add0->src[1] == node && add1->src[0] == add0 && other != node && residual != node && residual != add0 && + node->type == GGML_TYPE_F32 && + src->type == GGML_TYPE_F32 && gate->type == GGML_TYPE_F32 && other->type == GGML_TYPE_F32 && residual->type == GGML_TYPE_F32 && + ggml_are_same_shape(node, src) && ggml_are_same_shape(node, other) && ggml_are_same_shape(node, residual) && + gate->ne[0] == 1 && gate->ne[1] == node->ne[1] && gate->ne[2] == node->ne[2] && gate->ne[3] == node->ne[3] && + ggml_is_contiguous(src) && ggml_is_contiguous(gate) && ggml_is_contiguous(other) && + ggml_is_contiguous(residual) && ggml_is_contiguous(add1) && !gate_overlap && + safe_output_alias(src) && safe_output_alias(other) && safe_output_alias(residual) && + ggml_nelements(add1) <= int64_t(INT_MAX) * 256) { + ggml_cuda_op_shared_mul_add(*cuda_ctx, node, other, residual, add1); + return 2; + } + } + + // qwen4exp IQ4_NL routed down projection followed by the 10-expert weighted sum. Compute all selected + // experts for one output row in a wave and apply their routing weights immediately, avoiding the + // [n_embd, n_used] intermediate and its separate reduction launch. + if (getenv("GGML_CUDA_DISABLE_WEIGHTED_DOWN") == nullptr && + node->op == GGML_OP_MUL_MAT_ID && i + 20 < cgraph->n_nodes && cgraph->nodes[i + 1]->op == GGML_OP_MUL && + !ggml_cuda_mmb_claims(*cuda_ctx, node)) { + constexpr int n_used = 10; + constexpr int n_ops = 2 + n_used + (n_used - 1); + ggml_tensor * mul = cgraph->nodes[i + 1]; + const ggml_tensor * experts = ggml_are_same_shape(mul, mul->src[0]) ? mul->src[0] : mul->src[1]; + const ggml_tensor * weights = experts == mul->src[0] ? mul->src[1] : mul->src[0]; + const int output_idx = i + n_ops - 1; + ggml_tensor * output = cgraph->nodes[output_idx]; + bool valid = experts == node && node->ne[1] == n_used && + ggml_cuda_mul_mat_id_weighted_rdna3_5_ok(experts, weights, output); + + for (int j = 0; valid && j < n_used; ++j) { + const ggml_tensor * view = cgraph->nodes[i + 2 + j]; + valid = view->op == GGML_OP_VIEW && view->src[0] == mul && + view->ne[0] == mul->ne[0] && view->ne[1] == mul->ne[2] && view->ne[2] == 1 && view->ne[3] == 1 && + view->nb[0] == sizeof(float) && view->nb[1] == mul->nb[2] && view->view_offs == size_t(j) * mul->nb[1]; + } + const int add_start = i + 2 + n_used; + if (valid) { + const ggml_tensor * first = cgraph->nodes[add_start]; + valid = first->op == GGML_OP_ADD && first->src[0] == cgraph->nodes[i + 2] && first->src[1] == cgraph->nodes[i + 3]; + } + for (int j = 2; valid && j < n_used; ++j) { + const ggml_tensor * add = cgraph->nodes[add_start + j - 1]; + valid = add->op == GGML_OP_ADD && add->src[0] == cgraph->nodes[add_start + j - 2] && add->src[1] == cgraph->nodes[i + 2 + j]; + } + if (valid) { + std::vector ops = { GGML_OP_MUL_MAT_ID, GGML_OP_MUL }; + ops.insert(ops.end(), n_used, GGML_OP_VIEW); + ops.insert(ops.end(), n_used - 1, GGML_OP_ADD); + const int out_nodes[] = { output_idx }; + if (ggml_can_fuse_subgraph(cgraph, i, n_ops, ops.data(), out_nodes, 1) && + ggml_cuda_check_fusion_memory_ranges(cgraph, i, n_ops, out_nodes, 1)) { + ggml_cuda_mul_mat_id_weighted_rdna3_5(*cuda_ctx, experts, weights, output); + return n_ops - 1; + } + } + } + + if (node->op == GGML_OP_MUL && node->type == GGML_TYPE_F32 && node->ne[0] % 4 == 0 && + node->ne[3] == 1 && node->ne[1] >= 2 && node->ne[1] <= 16) { + const int n_expert_used = node->ne[1]; + const int n_ops = 2 * n_expert_used; + if (i + n_ops <= cgraph->n_nodes) { + std::vector ops = { GGML_OP_MUL }; + ops.insert(ops.end(), n_expert_used, GGML_OP_VIEW); + ops.insert(ops.end(), n_expert_used - 1, GGML_OP_ADD); + const int output_idx = i + n_ops - 1; + + if (ggml_can_fuse_subgraph(cgraph, i, n_ops, ops.data(), &output_idx, 1)) { + const ggml_tensor * experts = ggml_are_same_shape(node, node->src[0]) ? node->src[0] : node->src[1]; + const ggml_tensor * weights = experts == node->src[0] ? node->src[1] : node->src[0]; + ggml_tensor * output = cgraph->nodes[output_idx]; + bool valid = experts->op == GGML_OP_MUL_MAT_ID && experts->type == GGML_TYPE_F32 && weights->type == GGML_TYPE_F32 && + experts->ne[0] == node->ne[0] && experts->ne[1] == n_expert_used && experts->ne[2] == node->ne[2] && + experts->ne[3] == 1 && weights->ne[0] == 1 && weights->ne[1] == n_expert_used && + weights->ne[2] == node->ne[2] && weights->ne[3] == 1 && + ggml_is_contiguous(experts) && ggml_is_contiguous(weights) && ggml_is_contiguous(output) && + output->type == GGML_TYPE_F32 && output->ne[0] == node->ne[0] && output->ne[1] == node->ne[2] && + output->ne[2] == 1 && output->ne[3] == 1 && ggml_nelements(output) <= UINT32_MAX; + + for (int j = 0; valid && j < n_expert_used; ++j) { + const ggml_tensor * view = cgraph->nodes[i + 1 + j]; + valid = view->src[0] == node && view->ne[0] == node->ne[0] && view->ne[1] == node->ne[2] && + view->ne[2] == 1 && view->ne[3] == 1 && view->nb[0] == sizeof(float) && + view->nb[1] == node->nb[2] && view->view_offs == size_t(j) * node->nb[1]; + } + + const int add_start = i + 1 + n_expert_used; + if (valid) { + const ggml_tensor * first = cgraph->nodes[add_start]; + valid = first->src[0] == cgraph->nodes[i + 1] && first->src[1] == cgraph->nodes[i + 2]; + } + for (int j = 2; valid && j < n_expert_used; ++j) { + const ggml_tensor * add = cgraph->nodes[add_start + j - 1]; + valid = add->src[0] == cgraph->nodes[add_start + j - 2] && add->src[1] == cgraph->nodes[i + 1 + j]; + } + + if (valid) { + ggml_cuda_op_weighted_expert_sum(*cuda_ctx, node, output, n_expert_used); + return n_ops - 1; + } + } + } + } + // multi-(add or mul) if (node->op == GGML_OP_ADD || node->op == GGML_OP_MUL) { int n_fuse = 0; @@ -3632,6 +4958,39 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph } } + if ((node->op == GGML_OP_MUL_MAT_ID || node->op == GGML_OP_MUL_MAT) && i + 1 < cgraph->n_nodes) { + ggml_tensor * next = cgraph->nodes[i + 1]; + const ggml_tensor * src0 = node->src[0]; + const ggml_tensor * src0_next = next->src[0]; + const int cc = ggml_cuda_info().devices[cuda_ctx->device].cc; + + const bool has_ids = node->op == GGML_OP_MUL_MAT_ID; + const bool valid_sources = next->op == node->op && src0 && src0_next && node->src[1] && next->src[1] && + (!has_ids || node->src[2] && next->src[2]); + const bool ordinary_q8 = valid_sources && !has_ids && src0->type == GGML_TYPE_Q8_0 && src0_next->type == GGML_TYPE_Q8_0; + const bool shared_inputs = valid_sources && (has_ids || ordinary_q8) && node->src[1] == next->src[1] && + (!has_ids || node->src[2] == next->src[2]); + const int64_t mmq_cols = shared_inputs ? (has_ids ? node->src[1]->ne[2] : node->src[1]->ne[1]) : 0; + const int64_t n_experts = shared_inputs && has_ids ? src0->ne[2] : 0; + const bool use_mmq = shared_inputs && + ggml_cuda_should_use_mmq(src0->type, cc, mmq_cols, n_experts) && + ggml_cuda_should_use_mmq(src0_next->type, cc, mmq_cols, n_experts); + const bool compatible = use_mmq && !ggml_cuda_mmb_claims(*cuda_ctx, node) && node->src[1]->type == GGML_TYPE_F32 && + node->type == GGML_TYPE_F32 && next->type == GGML_TYPE_F32 && + ggml_are_same_shape(src0, src0_next) && ggml_are_same_shape(node, next) && + mmq_get_q8_1_ds_layout(src0->type) == mmq_get_q8_1_ds_layout(src0_next->type) && + !blackwell_mma_available(cc); + const int64_t ncols_dst = has_ids ? node->ne[2] : node->ne[1]; + const bool use_mmvq = compatible && ncols_dst <= MMVQ_MAX_BATCH_SIZE && + (!has_ids || ncols_dst <= get_mmvq_mmid_max_batch(src0->type, cc) || + ncols_dst <= get_mmvq_mmid_max_batch(src0_next->type, cc)); + + if (compatible && !use_mmvq) { + ggml_cuda_mul_mat_q_pair(*cuda_ctx, node, next); + return 1; + } + } + bool fused_mul_mat_vec = false; int fused_node_count = 0; @@ -3950,6 +5309,13 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph const ggml_tensor * src1 = up->src[1]; const ggml_tensor * ids = up->src[2]; + if (op == GGML_OP_MUL_MAT_ID && ggml_cuda_mmb_supported_glu(*cuda_ctx, gate->src[0], up->src[0], src1, ids, glu)) { + ggml_cuda_mul_mat_id_mmb_glu(*cuda_ctx, gate->src[0], up->src[0], src1, ids, glu); + fused_mul_mat_vec = true; + fused_node_count = 3; + break; + } + if (ggml_cuda_should_fuse_mul_mat_vec_f(up)) { ggml_cuda_mm_fusion_args_host fusion_data{}; fusion_data.gate = gate->src[0]; @@ -4147,6 +5513,29 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph return 2; } + if (node->op == GGML_OP_RMS_NORM && GGML_CUDA_CC_IS_RDNA3_5(ggml_cuda_info().devices[cuda_ctx->device].cc)) { + ggml_cuda_norm_gated_match match; + int skip = ggml_cuda_norm_gated_match_at(cgraph, i, match); + if (skip == 0) skip = ggml_cuda_norm_rows_match_at(cgraph, i, match); + if (skip > 0) { + if (match.staged) { + for (int j = i; j <= i + skip; ++j) { + if (ggml_cuda_is_view_or_noop(cgraph->nodes[j])) continue; + if (!ggml_cuda_compute_forward(*cuda_ctx, cgraph->nodes[j])) { + GGML_ABORT("norm-gated: staged dispatch failed"); + } + } + } else { + if (match.pre >= 0 && !ggml_cuda_compute_forward(*cuda_ctx, cgraph->nodes[match.pre])) { + GGML_ABORT("norm-gated: gate MUL_MAT dispatch failed"); + } + match.dst16 = ggml_cuda_bf16_for_mmb_reader(*cuda_ctx, cgraph, i + skip, match.dst); + ggml_cuda_op_norm_gated(*cuda_ctx, match); + } + return skip; + } + } + if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_RMS_NORM, GGML_OP_MUL }, {})) { ggml_cuda_op_rms_norm_fused(*cuda_ctx, node, cgraph->nodes[i + 1]); return 1; @@ -4162,6 +5551,212 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph return 1; } + // qwen4exp complete-block selection: sort the chosen blocks, gather their cells, validity, tail, in one kernel + if (node->op == GGML_OP_CONT && GGML_CUDA_CC_IS_RDNA3_5(ggml_cuda_info().devices[cuda_ctx->device].cc)) { + ggml_cuda_qsa_expand_args args; + const int skip = ggml_cuda_match_qsa_expand(cgraph, i, args); + if (skip && ggml_cuda_qsa_expand_alias_ok(args)) { + ggml_cuda_op_qsa_expand(*cuda_ctx, args); + return skip; + } + } + + // qwen4exp indexer key pooling: get_rows of the four member rows per block, summed and scaled in one kernel + if (node->op == GGML_OP_GET_ROWS && GGML_CUDA_CC_IS_RDNA3_5(ggml_cuda_info().devices[cuda_ctx->device].cc)) { + ggml_cuda_get_rows_mean4_match match; + if (ggml_cuda_match_get_rows_mean4(cgraph, i, match)) { + const int output_idx = i + match.node_count - 1; + if (ggml_cuda_check_fusion_memory_ranges(cgraph, i, match.node_count, &output_idx, 1) && + !ggml_cuda_mean4_ranges_overlap(match.src, match.dst) && + !ggml_cuda_mean4_ranges_overlap(match.ids, match.dst)) { + ggml_cuda_op_get_rows_mean4(*cuda_ctx, match.src, match.ids, match.dst); + return match.node_count - 1; + } + } + } + + // qwen4exp hyper-connection stream mix: + // HC gate GEMM [320 -> 10240] whose only consumer is the fused stream mix: GEMM + sigmoid + mix in one kernel + if (node->op == GGML_OP_MUL_MAT && ggml_cuda_mmb_gatemix() && i + 1 < cgraph->n_nodes && GGML_CUDA_CC_IS_RDNA3_5(ggml_cuda_info().devices[cuda_ctx->device].cc)) { + const ggml_tensor * w = node->src[0], * lo = node->src[1]; + if (ggml_is_quantized(w->type) && ggml_node_has_n_uses(cgraph, i, 1) && ggml_cuda_mmb_supported_mm(*cuda_ctx, w, lo, node)) { + ggml_cuda_hc_mix_args ma; + const int count = ggml_cuda_hc_mix_closed(cgraph, i + 1, ma); + if (count > 0 && ma.gate == node && ggml_cuda_hc_gate_mix(*cuda_ctx, w, lo, ma.xn, ma.dst, ma.hc, ma.scale, ma.bias)) return count; + } + } + + // sigmoid(gate) -> mul(xn, .) -> reshape -> view(stream 0) -> cont -> (view(stream c) -> add) x (hc-1) -> scale + // one kernel reads xn/gate once and writes the [n_embd, T] mean of the gated streams + { + ggml_cuda_hc_mix_args args; + enum ggml_op ops[5 + 2 * 15 + 1]; + int n_ops = ggml_cuda_match_hc_mix(cgraph, i, args, ops); + if (n_ops > 0 && args.dst->ne[1] == 1) { + n_ops = 0; // Decode preserves bit-identical routing. Prefill uses the fused reduction. + } + if (n_ops > 0) { + int node_idxs[5 + 2 * 15 + 1]; + for (int j = 0; j < n_ops; ++j) { + node_idxs[j] = i + j; + } + const int out_nodes[] = { i + n_ops - 1 }; + if (ggml_can_fuse_subgraph_ext(cgraph, node_idxs, n_ops, ops, out_nodes, 1)) { + ggml_cuda_op_hc_mix_reduce(*cuda_ctx, args); + return n_ops - 1; + } + } + } + + { + ggml_cuda_hc_combine_norm_args args; + const int skip = ggml_cuda_match_hc_combine_norm(cgraph, i, args, + ggml_cuda_info().devices[cuda_ctx->device].warp_size, true); + if (skip > 0 && !args.single_block) { + // HC16: a BF16 copy of xn for its consumers, and BF16 residual / block_out streams where every reader agreed + args.out_xn_bf16 = ggml_cuda_mmb_cache_reserve(*cuda_ctx, args.out_xn, (size_t) ggml_nelements(args.out_xn)); + args.store_xn_f32 = !(args.out_xn_bf16 && ggml_cuda_mmb_is_bf16_only(*cuda_ctx, args.out_xn)); + if (!args.store_xn_f32 && !args.out_inject) { + ggml_cuda_hc_combine_norm_args sa; + ggml_cuda_match_hc_combine_norm(cgraph, i, sa, ggml_cuda_info().devices[cuda_ctx->device].warp_size, false); + if (sa.out_inject) { + GGML_ABORT("hc_combine_norm: %s is BF16-only but the fused inject projection was refused (aliasing); " + "the graph_optimize allocation dependencies should have prevented it", args.out_xn->name); + } + } + if (ggml_cuda_mmb_res16()) { + const ggml_tensor * rin = args.residual->view_src ? args.residual->view_src : args.residual; + const ggml_tensor * rout = args.out_res->view_src ? args.out_res->view_src : args.out_res; + const bool in16 = ggml_cuda_mmb_is_bf16_only(*cuda_ctx, rin) && args.residual->view_offs == 0; + const bool out16 = ggml_cuda_mmb_is_bf16_only(*cuda_ctx, rout) && args.out_res->view_offs == 0; + if (args.out_res->data == args.residual->data && in16 != out16) { + GGML_ABORT("hc_combine_norm: in-place residual with mismatched BF16 marks (%s); the graph_optimize " + "allocation dependency should have separated them", args.out_res->name); + } + const ggml_tensor * rblk = args.block_out->view_src ? args.block_out->view_src : args.block_out; + args.blk_in_bf16 = (ggml_cuda_mmb_blk16() && ggml_cuda_mmb_is_bf16_only(*cuda_ctx, rblk) && args.block_out->view_offs == 0) + ? (const uint16_t *) args.block_out->data : nullptr; + args.res_in_bf16 = in16 ? (const uint16_t *) args.residual->data : nullptr; + args.res_out_bf16 = out16 ? (uint16_t *) args.out_res->data : nullptr; + } + } else { + // HC16: a combine whose streams were marked BF16-only in graph_optimize has no unfused fallback (the marked + // tensors were never written as F32), so a refusal here (buffer aliasing, or the single-block kernel) is a bug + ggml_cuda_hc_combine_norm_args sa; + const int structural = skip > 0 ? skip : ggml_cuda_match_hc_combine_norm(cgraph, i, sa, + ggml_cuda_info().devices[cuda_ctx->device].warp_size, false); + if (structural > 0) { + const ggml_cuda_hc_combine_norm_args & a = skip > 0 ? args : sa; + auto marked = [cuda_ctx](const ggml_tensor * t) { return ggml_cuda_mmb_is_bf16_only(*cuda_ctx, t) || (t->view_src && ggml_cuda_mmb_is_bf16_only(*cuda_ctx, t->view_src)); }; + if (marked(a.out_xn) || marked(a.residual) || marked(a.out_res) || marked(a.block_out)) { + GGML_ABORT("hc_combine_norm: %s has BF16-only streams but the fused kernel was refused (%s); the graph_optimize " + "allocation dependencies should have prevented the aliasing", a.out_res->name, + skip > 0 ? "single-block only" : "buffer aliasing"); + } + } + } + if (skip > 0) { + ggml_cuda_op_hc_combine_norm(*cuda_ctx, args); + return skip; + } + } + + // qwen4exp hyper-connection combine: + // scale(inject) -> sigmoid -> scale -> [reshape] -> repeat(block_out) -> mul -> add(residual, .) + // one kernel writes residual + block_out (x) (s2*sigmoid(s1*inject)) + if (node->op == GGML_OP_SCALE && node->type == GGML_TYPE_F32 && i + 6 < cgraph->n_nodes && + cgraph->nodes[i + 1]->op == GGML_OP_UNARY && ggml_get_unary_op(cgraph->nodes[i + 1]) == GGML_UNARY_OP_SIGMOID && + cgraph->nodes[i + 2]->op == GGML_OP_SCALE) { + const ggml_tensor * scale1 = node; + const ggml_tensor * sigmoid = cgraph->nodes[i + 1]; + const ggml_tensor * scale2 = cgraph->nodes[i + 2]; + const ggml_tensor * inject = scale1->src[0]; + + // a few view nodes follow: the scatter weights [1, hc, T] and the block output [n_embd, 1, T] + int j = i + 3; + while (j < cgraph->n_nodes && j < i + 7 && ggml_cuda_is_view_or_noop(cgraph->nodes[j]) && !ggml_is_empty(cgraph->nodes[j])) { + j++; + } + + if (j + 2 < cgraph->n_nodes && + cgraph->nodes[j]->op == GGML_OP_REPEAT && cgraph->nodes[j + 1]->op == GGML_OP_MUL && cgraph->nodes[j + 2]->op == GGML_OP_ADD) { + const ggml_tensor * repeat = cgraph->nodes[j]; + const ggml_tensor * mul = cgraph->nodes[j + 1]; + ggml_tensor * add = cgraph->nodes[j + 2]; + const ggml_tensor * b = repeat->src[0]; + const ggml_tensor * w = mul->src[0] == repeat ? mul->src[1] : mul->src[0]; + const ggml_tensor * res = add->src[0] == mul ? add->src[1] : add->src[0]; + + const int64_t n_embd = add->ne[0]; + const int64_t hc = add->ne[1]; + const int64_t n_tokens = add->ne[2]; + + const bool w_ok = (w == scale2 || (ggml_cuda_is_view_or_noop(w) && w->view_src == scale2)) && + w->type == GGML_TYPE_F32 && ggml_is_contiguous(w) && + w->ne[0] == 1 && w->ne[1] == hc && w->ne[2] == n_tokens && w->ne[3] == 1; + const bool b_ok = b->type == GGML_TYPE_F32 && ggml_is_contiguous(b) && + b->ne[0] == n_embd && b->ne[1] == 1 && b->ne[2] == n_tokens && b->ne[3] == 1; + const bool ok = sigmoid->src[0] == scale1 && scale2->src[0] == sigmoid && + (mul->src[0] == repeat || mul->src[1] == repeat) && w != repeat && + (add->src[0] == mul || add->src[1] == mul) && res != mul && + sigmoid->type == GGML_TYPE_F32 && scale2->type == GGML_TYPE_F32 && inject->type == GGML_TYPE_F32 && + repeat->type == GGML_TYPE_F32 && mul->type == GGML_TYPE_F32 && add->type == GGML_TYPE_F32 && res->type == GGML_TYPE_F32 && + ggml_is_contiguous(inject) && ggml_nelements(inject) == hc * n_tokens && ggml_nrows(inject) == n_tokens && + ggml_are_same_shape(scale1, inject) && ggml_are_same_shape(sigmoid, inject) && ggml_are_same_shape(scale2, inject) && + ggml_are_same_shape(repeat, add) && ggml_are_same_shape(mul, add) && ggml_are_same_shape(res, add) && + add->ne[3] == 1 && ggml_is_contiguous(res) && ggml_is_contiguous(add) && w_ok && b_ok; + + if (ok) { + // the block-output reshape views an external tensor, so it is left out of the subgraph check + int idxs[12]; + enum ggml_op ops[12]; + int count = 0; + for (int k = i; k <= j + 2; ++k) { + const ggml_tensor * n = cgraph->nodes[k]; + if (ggml_cuda_is_view_or_noop(n) && n->view_src != scale2) { + continue; + } + idxs[count] = k; + ops[count] = n->op; + count++; + } + const int out_nodes[] = { j + 2 }; + if (ggml_can_fuse_subgraph_ext(cgraph, idxs, count, ops, out_nodes, 1)) { + ggml_cuda_hc_combine_args args; + args.residual = res; + args.block_out = b; + args.inject = inject; + args.dst = add; + args.s1 = ggml_get_op_params_f32(scale1, 0); + args.b1 = ggml_get_op_params_f32(scale1, 1); + args.s2 = ggml_get_op_params_f32(scale2, 0); + args.b2 = ggml_get_op_params_f32(scale2, 1); + ggml_cuda_op_hc_combine(*cuda_ctx, args); + return j + 2 - i; + } + } + } + } + + if (node->op == GGML_OP_CONT && i + 2 < cgraph->n_nodes) { + ggml_tensor * unary = cgraph->nodes[i + 1]; + ggml_tensor * mul = cgraph->nodes[i + 2]; + const ggml_tensor * gate = node->src[0]; + const ggml_tensor * other = mul->src[0] == unary ? mul->src[1] : mul->src[0]; + const int out_nodes[] = { i + 2 }; + if (unary->op == GGML_OP_UNARY && ggml_get_unary_op(unary) == GGML_UNARY_OP_SIGMOID && mul->op == GGML_OP_MUL && + (mul->src[0] == unary || mul->src[1] == unary) && gate->type == GGML_TYPE_F32 && other->type == GGML_TYPE_F32 && + gate->nb[0] == sizeof(float) && gate->nb[1] == 2 * gate->ne[0] * sizeof(float) && + gate->nb[2] == gate->nb[1] * gate->ne[1] && gate->nb[3] == gate->nb[2] * gate->ne[2] && + ggml_is_contiguous(other) && ggml_are_same_shape(node, other) && ggml_are_same_shape(unary, node) && + ggml_are_same_shape(mul, other) && + ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_CONT, GGML_OP_UNARY, GGML_OP_MUL }, { i + 2 }) && + ggml_cuda_check_fusion_memory_ranges(cgraph, i, 3, out_nodes, 1)) { + ggml_cuda_op_cont_sigmoid_mul(*cuda_ctx, node, unary, mul, ggml_cuda_bf16_for_mmb_reader(*cuda_ctx, cgraph, i + 2, mul)); + return 2; + } + } + if (ggml_cuda_can_fuse(cgraph, i, { GGML_OP_UNARY, GGML_OP_MUL }, { GGML_UNARY_OP_SILU }) || ggml_cuda_can_fuse(cgraph, i, { GGML_OP_UNARY, GGML_OP_MUL }, { GGML_UNARY_OP_SIGMOID }) || ggml_cuda_can_fuse(cgraph, i, { GGML_OP_UNARY, GGML_OP_MUL }, { GGML_UNARY_OP_SOFTPLUS })) { @@ -4179,6 +5774,35 @@ static int ggml_cuda_try_fuse(ggml_backend_cuda_context * cuda_ctx, ggml_cgraph return 2; } + // scale -> silu/sigmoid -> quantized matvec (decode): apply the activation while the matvec quantizes it + if (node->op == GGML_OP_SCALE && node->type == GGML_TYPE_F32 && i + 2 < cgraph->n_nodes && + cgraph->nodes[i + 1]->op == GGML_OP_UNARY && + (ggml_get_unary_op(cgraph->nodes[i + 1]) == GGML_UNARY_OP_SILU || ggml_get_unary_op(cgraph->nodes[i + 1]) == GGML_UNARY_OP_SIGMOID) && + cgraph->nodes[i + 2]->op == GGML_OP_MUL_MAT && cgraph->nodes[i + 2]->src[1] == cgraph->nodes[i + 1] && + cgraph->nodes[i + 1]->src[0] == node && cgraph->nodes[i + 1]->type == GGML_TYPE_F32 && node->src[0]->type == GGML_TYPE_F32 && + ggml_is_contiguous(node->src[0]) && ggml_are_same_shape(node->src[0], cgraph->nodes[i + 1]) && + ggml_can_fuse_subgraph(cgraph, i, { GGML_OP_SCALE, GGML_OP_UNARY, GGML_OP_MUL_MAT }, { i + 2 }) && + ggml_cuda_mul_mat_vec_q_fq_prologue_ok(*cuda_ctx, cgraph->nodes[i + 2], node->src[0])) { + const int y_op = ggml_get_unary_op(cgraph->nodes[i + 1]) == GGML_UNARY_OP_SILU ? 1 : 2; + ggml_tensor * mm = cgraph->nodes[i + 2]; + + ggml_cuda_mul_mat_vec_q_fq_prologue(*cuda_ctx, mm, node->src[0], + ggml_get_op_params_f32(node, 0), ggml_get_op_params_f32(node, 1), y_op); + return 2; + } + + // scale -> silu/sigmoid (e.g. the qwen4exp hyper-connection low-rank gate: silu(x / hc)) + if (node->op == GGML_OP_SCALE && node->type == GGML_TYPE_F32 && i + 1 < cgraph->n_nodes && + cgraph->nodes[i + 1]->op == GGML_OP_UNARY && + (ggml_get_unary_op(cgraph->nodes[i + 1]) == GGML_UNARY_OP_SILU || ggml_get_unary_op(cgraph->nodes[i + 1]) == GGML_UNARY_OP_SIGMOID) && + cgraph->nodes[i + 1]->type == GGML_TYPE_F32 && node->src[0]->type == GGML_TYPE_F32 && + ggml_is_contiguous(node->src[0]) && ggml_is_contiguous(cgraph->nodes[i + 1]) && + ggml_can_fuse(cgraph, i, { GGML_OP_SCALE, GGML_OP_UNARY })) { + ggml_cuda_op_scale_unary(*cuda_ctx, node, cgraph->nodes[i + 1], + ggml_cuda_bf16_for_mmb_reader(*cuda_ctx, cgraph, i + 1, cgraph->nodes[i + 1])); + return 1; + } + return 0; } @@ -4324,6 +5948,9 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud int nodes_to_skip = ggml_cuda_try_fuse(cuda_ctx, cgraph, i); + if (nodes_to_skip == GGML_CUDA_FUSED_SELF) { + continue; + } if (nodes_to_skip != 0) { #ifdef GGML_CUDA_DEBUG const int last_fused = i + nodes_to_skip; @@ -4404,6 +6031,11 @@ static void ggml_cuda_graph_evaluate_and_capture(ggml_backend_cuda_context * cud static bool ggml_cuda_graph_set_enabled(ggml_backend_cuda_context * cuda_ctx, const void * graph_key) { ggml_cuda_graph * graph = cuda_ctx->cuda_graph(graph_key); + static const bool disable_graphs_env = getenv("GGML_CUDA_DISABLE_GRAPHS") != nullptr && atoi(getenv("GGML_CUDA_DISABLE_GRAPHS")) != 0; + if (disable_graphs_env) { + graph->disable_due_to_gpu_arch = true; + } + if (graph->graph == nullptr) { if (ggml_cuda_info().devices[cuda_ctx->device].cc < GGML_CUDA_CC_VOLTA) { if (!graph->disable_due_to_gpu_arch) { @@ -4417,8 +6049,36 @@ static bool ggml_cuda_graph_set_enabled(ggml_backend_cuda_context * cuda_ctx, co } #endif // USE_CUDA_GRAPH +// the marks are keyed by tensor pointer and belong to the graphs graph_optimize last saw. A graph computed without +// that optimize (test-backend-ops, any direct ggml_backend_graph_compute caller) can recycle those addresses, so a +// producer would store BF16 into a tensor of another graph and a plain reader would abort. Fingerprint each optimized +// split and drop the marks when a graph arrives that none of them match; a reused (identical) graph keeps them. +static uint64_t ggml_cuda_graph_sig(const ggml_cgraph * cgraph) { + // no data pointers: the scheduler allocates after graph_optimize, so they differ between optimize and compute + uint64_t h = 1469598103934665603ull ^ (uint64_t) cgraph->n_nodes; + for (int i = 0; i < cgraph->n_nodes; ++i) { + const ggml_tensor * t = cgraph->nodes[i]; + h = (h ^ (uint64_t) (uintptr_t) t) * 1099511628211ull; + h = (h ^ (uint64_t) t->op ^ ((uint64_t) t->ne[0] << 8) ^ ((uint64_t) t->ne[1] << 24)) * 1099511628211ull; + for (int s = 0; s < GGML_MAX_SRC && t->src[s]; ++s) { + h = (h ^ (uint64_t) (uintptr_t) t->src[s]) * 1099511628211ull; + } + } + return h; +} + static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, ggml_cgraph * cgraph) { ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context; + if (ggml_cuda_mmb_marks_count(*cuda_ctx) > 0) { + const uint64_t sig = ggml_cuda_graph_sig(cgraph); + if (std::find(cuda_ctx->mmb_graph_sigs.begin(), cuda_ctx->mmb_graph_sigs.end(), sig) == cuda_ctx->mmb_graph_sigs.end()) { + ggml_cuda_mmb_marks_clear(*cuda_ctx); + cuda_ctx->mmb_graph_sigs.clear(); + cuda_ctx->mmb_res16_pending.clear(); + cuda_ctx->mmb_first_split = nullptr; + } + } + ggml_cuda_mmb_begin_graph(*cuda_ctx); ggml_cuda_set_device(cuda_ctx->device); @@ -4473,6 +6133,7 @@ static enum ggml_status ggml_backend_cuda_graph_compute(ggml_backend_t backend, ggml_cuda_graph_evaluate_and_capture(cuda_ctx, cgraph, use_cuda_graph, cuda_graph_update_required, graph_key); + cuda_ctx->mmb_after_compute = true; return GGML_STATUS_SUCCESS; } @@ -4501,11 +6162,361 @@ static void ggml_backend_cuda_event_wait(ggml_backend_t backend, ggml_backend_ev } } + +// true when every use of x, and of the views of x present in this split, is a node of this split graph. The scheduler +// optimizes each split separately and the split graph view shares the full graph's use counts, so a reader that lives +// in a later split (llama-bench's zero-output prefill ubatches place the final get_rows elsewhere) shows up as a use +// count the local scan cannot account for; such a tensor must keep its F32 store. +static bool ggml_cuda_marks_is_view(const ggml_tensor * t) { + return t->op == GGML_OP_VIEW || t->op == GGML_OP_RESHAPE || t->op == GGML_OP_PERMUTE || t->op == GGML_OP_TRANSPOSE; +} + +static int ggml_cuda_marks_uses(const ggml_cgraph * g, const ggml_tensor * t) { + const size_t h = ggml_hash_find(&g->visited_hash_set, t); + return h == GGML_HASHSET_FULL || !ggml_bitset_get(g->visited_hash_set.used, h) ? -1 : g->use_counts[h]; +} + +// uses of x (and of its views present in this split) that no node of this split accounts for; -1 when unknown +static int ggml_cuda_marks_readers_left(const ggml_cgraph * g, const ggml_tensor * x) { + int expected = ggml_cuda_marks_uses(g, x), found = 0; + if (expected < 0) return -1; + for (int i = 0; i < g->n_nodes; ++i) { + const ggml_tensor * t = g->nodes[i]; + if (t != x && ggml_cuda_marks_is_view(t) && t->view_src == x) { + const int u = ggml_cuda_marks_uses(g, t); + if (u < 0) return -1; + expected += u; + } + for (int s = 0; s < GGML_MAX_SRC && t->src[s]; ++s) { + const ggml_tensor * src = t->src[s]; + if (src == x || (ggml_cuda_marks_is_view(src) && src->view_src == x)) ++found; + } + } + return expected - found; +} + +static bool ggml_cuda_marks_readers_local(const ggml_cgraph * g, const ggml_tensor * x) { + return ggml_cuda_marks_readers_left(g, x) == 0; +} + static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph * cgraph, ggml_backend_graph_optimize_params * params) { ggml_backend_cuda_context * cuda_ctx = (ggml_backend_cuda_context *) backend->context; + { // marks live for the whole scheduled graph (all splits): clear at the first optimize after a compute, or when the first split repeats + const void * key = cgraph->n_nodes ? cgraph->nodes[0] : nullptr; + if (cuda_ctx->mmb_after_compute || cuda_ctx->mmb_first_split == nullptr || key == cuda_ctx->mmb_first_split) { ggml_cuda_mmb_marks_clear(*cuda_ctx); cuda_ctx->mmb_graph_sigs.clear(); cuda_ctx->mmb_res16_pending.clear(); cuda_ctx->mmb_first_split = key; cuda_ctx->mmb_after_compute = false; } + } + { // HC16: mark tensors whose readers all take the BF16 copy, so their producers skip the F32 store + if (GGML_CUDA_CC_IS_RDNA3_5(ggml_cuda_info().devices[cuda_ctx->device].cc)) { + const int ws = ggml_cuda_info().devices[cuda_ctx->device].warp_size; + auto reads = [](const ggml_tensor * t, const ggml_tensor * x) { for (int s = 0; s < GGML_MAX_SRC && t->src[s]; ++s) if (t->src[s] == x || t->src[s]->view_src == x) return true; return false; }; + // normalized stream xn: read only by MMB GEMMs (which take the BF16 copy) and the fused stream mix + for (int i = 0; i < cgraph->n_nodes; ++i) { + ggml_cuda_hc_combine_norm_args ca; + if (ggml_cuda_match_hc_combine_norm(cgraph, i, ca, ws, false) <= 0) continue; // structure only: no data pointers yet + const ggml_tensor * xn = ca.out_xn; + if (ggml_nrows(xn) < 512) continue; + bool ok = true; int nread = 0; + for (int n = i + 1; n < cgraph->n_nodes && ok; ++n) { + const ggml_tensor * t = cgraph->nodes[n]; + if (!reads(t, xn)) continue; + ++nread; + if (t == ca.out_inject) continue; + if (t->op == GGML_OP_MUL_MAT && t->src[0]->type != GGML_TYPE_F32 && + ggml_cuda_mmb_supported_mm(*cuda_ctx, t->src[0], t->src[1], t)) continue; + if (t->op == GGML_OP_MUL && n >= 1) { ggml_cuda_hc_mix_args ma; if (ggml_cuda_hc_mix_closed(cgraph, n - 1, ma) > 0 && (ma.xn == t->src[0] || ma.xn == t->src[1]) && (ma.xn == xn || ma.xn->view_src == xn)) continue; } + if (t->op == GGML_OP_VIEW || t->op == GGML_OP_RESHAPE) continue; + ok = false; + } + if (ok && nread > 0 && ggml_cuda_marks_readers_local(cgraph, xn)) ggml_cuda_mmb_mark_bf16_only(*cuda_ctx, xn); + } + if (ggml_cuda_mmb_blk16()) { // block_out (attention out-proj / MoE merge) is read only by the fused combine + std::vector blks; + for (int i = 0; i < cgraph->n_nodes; ++i) { + ggml_cuda_hc_combine_norm_args ca; + if (ggml_cuda_match_hc_combine_norm(cgraph, i, ca, ws, false) > 0) + blks.push_back(ca.block_out->view_src ? ca.block_out->view_src : ca.block_out); + } + for (const ggml_tensor * b : blks) { + if (ggml_nrows(b) < 512 || b->ne[0] % 8 != 0) continue; + bool producer_ok = false; // only producers that honour a BF16 mark + if (b->op == GGML_OP_MUL_MAT && ggml_cuda_mmb_supported_mm(*cuda_ctx, b->src[0], b->src[1], b)) producer_ok = true; + if (!producer_ok) { + for (int k = 0; k < cgraph->n_nodes && !producer_ok; ++k) { + ggml_cuda_moe_weighted_reduction_match mm; + if (cgraph->nodes[k]->op != GGML_OP_MUL) continue; + if (!ggml_cuda_match_moe_weighted_reduction(cgraph, k, mm)) continue; + if (mm.dst == b) producer_ok = true; + else if (k + mm.node_count < cgraph->n_nodes && + cgraph->nodes[k + mm.node_count] == b && b->op == GGML_OP_ADD && + ggml_node_has_n_uses(cgraph, k + mm.node_count - 1, 1) && + (b->src[0] == mm.dst || b->src[1] == mm.dst)) producer_ok = true; + } + } + if (!producer_ok) continue; + int bi = -1; + for (int k = 0; k < cgraph->n_nodes; ++k) if (cgraph->nodes[k] == b) { bi = k; break; } + if (bi < 0) continue; + bool ok = true; int nread = 0; + for (int n = bi + 1; n < cgraph->n_nodes && ok; ++n) { + const ggml_tensor * t = cgraph->nodes[n]; + if (!reads(t, b)) continue; + ++nread; + if (t->op == GGML_OP_VIEW || t->op == GGML_OP_RESHAPE || t->op == GGML_OP_REPEAT) continue; + if (t->op == GGML_OP_MUL || t->op == GGML_OP_ADD) continue; + ok = false; + } + if (ok && nread > 0 && ggml_cuda_marks_readers_local(cgraph, b)) ggml_cuda_mmb_mark_bf16_only(*cuda_ctx, b); + } + } + if (ggml_cuda_mmb_res16()) { // residual stream: the only readers are the fused norm and the next combine + std::vector> comb; + for (int i = 0; i < cgraph->n_nodes; ++i) { + ggml_cuda_hc_combine_norm_args ca; + if (ggml_cuda_match_hc_combine_norm(cgraph, i, ca, ws, false) > 0) comb.emplace_back(ca.residual, ca.out_res); + } + auto is_next_combine = [&comb](const ggml_tensor * r, const ggml_tensor * t) { + for (const auto & d : comb) { + const ggml_tensor * dr = d.first->view_src ? d.first->view_src : d.first; + if (dr == r && d.second == t) return true; + } + return false; + }; + // The scheduler cuts a split wherever the input count runs out (GGML_SCHED_MAX_SPLIT_INPUTS), so a + // residual's next combine can land in a later split. Every split is optimized before any is computed + // and the marks live for the whole graph, so finish the reader check here instead of dropping the + // mark: otherwise precision depends on where the cuts fall (e.g. QSA fast vs full re-pool inputs). + // Only the next fused combine may read it across a cut; a reader on another backend is never seen, + // so its residual stays F32. + auto & pending = cuda_ctx->mmb_res16_pending; + for (auto it = pending.begin(); it != pending.end(); ) { + const ggml_tensor * r = it->first; + int left = it->second; + bool ok = true; + for (int n = 0; n < cgraph->n_nodes && ok; ++n) { + const ggml_tensor * t = cgraph->nodes[n]; + const bool view_of_r = t != r && ggml_cuda_marks_is_view(t) && t->view_src == r; + if (view_of_r) { + const int u = ggml_cuda_marks_uses(cgraph, t); + if (u < 0) { ok = false; break; } + left += u; + } + for (int s = 0; s < GGML_MAX_SRC && t->src[s]; ++s) { + const ggml_tensor * src = t->src[s]; + if (src != r && !(ggml_cuda_marks_is_view(src) && src->view_src == r)) continue; + --left; + if (view_of_r && (t->op == GGML_OP_VIEW || t->op == GGML_OP_RESHAPE)) continue; + if (!is_next_combine(r, t)) ok = false; + } + } + if (!ok || left < 0) { it = pending.erase(it); continue; } + if (left == 0) { ggml_cuda_mmb_mark_bf16_only(*cuda_ctx, r); it = pending.erase(it); continue; } + it->second = left; + ++it; + } + for (const auto & c : comb) { + const ggml_tensor * r = c.second; + if (ggml_nrows(r) < 512) continue; + int ri = -1; + for (int k = 0; k < cgraph->n_nodes; ++k) if (cgraph->nodes[k] == r) { ri = k; break; } + if (ri < 0) continue; + bool ok = true; int nread = 0; + for (int n = ri + 1; n < cgraph->n_nodes && ok; ++n) { + const ggml_tensor * t = cgraph->nodes[n]; + if (!reads(t, r)) continue; + ++nread; + if (t->op == GGML_OP_RMS_NORM) continue; + if (t->op == GGML_OP_VIEW || t->op == GGML_OP_RESHAPE) continue; + if (is_next_combine(r, t)) continue; + ok = false; + } + if (!ok) continue; + const int left = ggml_cuda_marks_readers_left(cgraph, r); + if (left == 0 && nread > 0) ggml_cuda_mmb_mark_bf16_only(*cuda_ctx, r); + else if (left > 0) pending[r] = left; + } + } + for (int i = 0; i < cgraph->n_nodes; ++i) { // fused stream mix output read only by MMB GEMMs + ggml_cuda_hc_mix_args ma; + if (ggml_cuda_hc_mix_closed(cgraph, i, ma) <= 0 || !ma.dst) continue; + const ggml_tensor * d = ma.dst; + if (ggml_nrows(d) < 512) continue; + bool ok = true; int nread = 0; + for (int n = i + 1; n < cgraph->n_nodes && ok; ++n) { + const ggml_tensor * t = cgraph->nodes[n]; + if (!reads(t, d)) continue; + ++nread; + if (t->op == GGML_OP_VIEW || t->op == GGML_OP_RESHAPE) continue; + if (t->op == GGML_OP_MUL_MAT && t->src[0]->type != GGML_TYPE_F32 && + ggml_cuda_mmb_supported_mm(*cuda_ctx, t->src[0], t->src[1], t)) continue; + ok = false; + } + if (ok && nread > 0 && ggml_cuda_marks_readers_local(cgraph, d)) ggml_cuda_mmb_mark_bf16_only(*cuda_ctx, d); + } + for (int i = 0; i < cgraph->n_nodes; ++i) { // the HC gate GEMM feeding only the fused mix keeps its BF16 epilogue + const ggml_tensor * t = cgraph->nodes[i]; + if (t->op != GGML_OP_MUL_MAT || t->src[0]->ne[0] != 320 || t->src[0]->ne[1] != 10240 || !ggml_cuda_mmb_supported_mm(*cuda_ctx, t->src[0], t->src[1], t)) continue; + bool ok = true; int nread = 0; + for (int n = i + 1; n < cgraph->n_nodes && ok; ++n) { + const ggml_tensor * u = cgraph->nodes[n]; + if (!reads(u, t)) continue; + ++nread; + if (u->op == GGML_OP_UNARY && ggml_get_unary_op(u) == GGML_UNARY_OP_SIGMOID) { ggml_cuda_hc_mix_args ma; if (ggml_cuda_hc_mix_closed(cgraph, n, ma) > 0 && ma.gate == t) continue; } + ok = false; + } + if (ok && nread > 0 && ggml_cuda_marks_readers_local(cgraph, t)) ggml_cuda_mmb_mark_bf16_only(*cuda_ctx, t); + } + if (ggml_cuda_mmb_down16()) { // IQ4_NL routed down GEMM read only by the weighted reduction keeps BF16 outputs + for (int i = 0; i < cgraph->n_nodes; ++i) { + ggml_cuda_moe_weighted_reduction_match match; + if (!ggml_cuda_match_moe_weighted_reduction(cgraph, i, match)) continue; + const ggml_tensor * ex = match.experts; + int xi = -1; + for (int k = 0; k < i; ++k) { if (cgraph->nodes[k] == ex) { xi = k; break; } } + if (xi < 0 || ex->op != GGML_OP_MUL_MAT_ID || ex->type != GGML_TYPE_F32) continue; + if (!ggml_cuda_mmb_supported_mmid(*cuda_ctx, ex->src[0], ex->src[1], ex->src[2], const_cast(ex))) continue; + if (ex->src[0]->type != GGML_TYPE_IQ4_NL) continue; // other formats keep F32 expert outputs + if (!ggml_node_has_n_uses(cgraph, xi, 1)) continue; + if (ex->ne[0] % 8 != 0 || ggml_nrows(ex) < 512) continue; + if (ggml_cuda_marks_readers_local(cgraph, ex)) ggml_cuda_mmb_mark_bf16_only(*cuda_ctx, ex); + } + } + if (ggml_cuda_mmb_res16()) { // 16-bit residual stream: producer and consumer on different buffers + for (int i = 0; i < cgraph->n_nodes; ++i) { + ggml_cuda_hc_combine_norm_args ca; + if (ggml_cuda_match_hc_combine_norm(cgraph, i, ca, ws, false) <= 0) continue; + ggml_tensor * rr = const_cast(ca.residual->view_src ? ca.residual->view_src : ca.residual); + params->add_alloc_dep(params->user_data, rr, ca.out_res); + params->add_alloc_dep(params->user_data, rr, ca.out_xn); + } + } + } + } + { // BF16 shadow copies of eligible dense weights, and the fused gate/up GEMM: its inputs stay allocated through the + // GLU output, and a GLU read only by a supported routed GEMM keeps the BF16 epilogue copy instead of the F32 output + auto reads = [](const ggml_tensor * t, const ggml_tensor * x) { for (int s = 0; s < GGML_MAX_SRC && t->src[s]; ++s) if (t->src[s] == x || t->src[s]->view_src == x) return true; return false; }; + for (int i = 0; i < cgraph->n_nodes; ++i) { + const ggml_tensor * t = cgraph->nodes[i]; + if (t->op == GGML_OP_MUL_MAT && (t->src[0]->type == GGML_TYPE_IQ4_NL || t->src[0]->type == GGML_TYPE_Q6_K) && ggml_cuda_mmb_supported_mm(*cuda_ctx, t->src[0], t->src[1], t)) ggml_cuda_mmb_shadow_prepare(*cuda_ctx, t->src[0]); + } + for (int i = 0; i < cgraph->n_nodes; ++i) { + const ggml_tensor * glu = cgraph->nodes[i]; + if (glu->op != GGML_OP_GLU || !glu->src[0] || !glu->src[1] || glu->src[0]->op != GGML_OP_MUL_MAT_ID) continue; + if (!ggml_cuda_mmb_supported_glu(*cuda_ctx, glu->src[0]->src[0], glu->src[1]->src[0], glu->src[0]->src[1], glu->src[0]->src[2], glu)) continue; + for (const auto * input : {glu->src[0]->src[0], glu->src[1]->src[0], glu->src[0]->src[1], glu->src[0]->src[2], glu->src[0], glu->src[1]}) { + auto * root = const_cast(input->view_src ? input->view_src : input); + params->add_alloc_dep(params->user_data, root, const_cast(glu)); + } + bool ok = true; int nread = 0; + for (int n = i + 1; n < cgraph->n_nodes && ok; ++n) { + const ggml_tensor * t = cgraph->nodes[n]; + if (!reads(t, glu)) continue; + ++nread; + if (t->op == GGML_OP_MUL_MAT_ID && t->src[1] == glu && ggml_cuda_mmb_supported_mmid(*cuda_ctx, t->src[0], t->src[1], t->src[2], t)) continue; + ok = false; + } + if (ok && nread > 0 && ggml_cuda_marks_readers_local(cgraph, glu)) ggml_cuda_mmb_mark_bf16_only(*cuda_ctx, glu); + } + } static const bool disable_fusion = getenv("GGML_CUDA_DISABLE_FUSION") != nullptr && std::atoi(getenv("GGML_CUDA_DISABLE_FUSION")); if (!disable_fusion) { + for (int i = 0; i < cgraph->n_nodes; ++i) { + ggml_cuda_hc_combine_norm_args args; + const int skip = ggml_cuda_match_hc_combine_norm(cgraph, i, args, + ggml_cuda_info().devices[cuda_ctx->device].warp_size, false); + if (skip > 0) { + const ggml_tensor * inputs[] = {args.block_out, args.inject}; + for (const auto * input : inputs) { + auto * root = const_cast(input->view_src ? input->view_src : input); + params->add_alloc_dep(params->user_data, root, args.out_xn); + } + // the combined stream must not share a buffer with the normalized output either: the alias check refuses + // that overlap, and a refused combine whose streams are BF16-only has no fallback (last layer: out_res is + // otherwise dead after its own norm, so the allocator would hand its buffer to out_xn) + params->add_alloc_dep(params->user_data, args.out_res, args.out_xn); + if (args.out_inject) { + const ggml_tensor * live[] = {args.block_out, args.inject, args.residual, args.out_res, args.out_xn}; + for (const auto * t : live) { + auto * root = const_cast(t->view_src ? t->view_src : t); + params->add_alloc_dep(params->user_data, root, args.out_inject); + } + } + i += skip; + } + } + // fused block selection: every input and intermediate output stays allocated through the final output + if (GGML_CUDA_CC_IS_RDNA3_5(ggml_cuda_info().devices[cuda_ctx->device].cc)) { + for (int i = 0; i < cgraph->n_nodes; ++i) { + ggml_cuda_qsa_expand_args args; + const int skip = ggml_cuda_match_qsa_expand(cgraph, i, args); + if (skip) { + for (const auto * input : {args.blocks, args.scores, args.cells, args.tail, + (const ggml_tensor *) args.cast_blocks, (const ggml_tensor *) args.cast_cells, (const ggml_tensor *) args.expanded}) { + auto * root = const_cast(input->view_src ? input->view_src : input); + params->add_alloc_dep(params->user_data, root, args.output); + } + i += skip; + } + } + } + // fused indexer key pooling: the key cache view and the member rows stay allocated through the pooled output + if (GGML_CUDA_CC_IS_RDNA3_5(ggml_cuda_info().devices[cuda_ctx->device].cc)) { + for (int i = 0; i < cgraph->n_nodes; ++i) { + ggml_cuda_get_rows_mean4_match pool; + if (ggml_cuda_match_get_rows_mean4(cgraph, i, pool)) { + params->add_alloc_dep(params->user_data, const_cast(pool.src->view_src ? pool.src->view_src : pool.src), pool.dst); + params->add_alloc_dep(params->user_data, const_cast(pool.ids->view_src ? pool.ids->view_src : pool.ids), pool.dst); + i += pool.node_count - 1; + } + } + } + + if (GGML_CUDA_CC_IS_RDNA3_5(ggml_cuda_info().devices[cuda_ctx->device].cc)) { + for (int i = 0; i < cgraph->n_nodes; ++i) { + if (cgraph->nodes[i]->op != GGML_OP_UNARY) continue; + ggml_cuda_idx_relu_sum_args args; + if (ggml_cuda_match_idx_relu_sum(cgraph, i, args) > 0) { + auto * root = const_cast(args.score->view_src ? args.score->view_src : args.score); + params->add_alloc_dep(params->user_data, root, args.dst); + } + } + } + + if (GGML_CUDA_CC_IS_RDNA3_5(ggml_cuda_info().devices[cuda_ctx->device].cc)) { + for (int i = 0; i < cgraph->n_nodes; ++i) { + if (cgraph->nodes[i]->op != GGML_OP_RMS_NORM) continue; + ggml_cuda_norm_gated_match match; + if (ggml_cuda_norm_gated_match_at(cgraph, i, match) > 0 && match.pre >= 0) { + for (const auto * input : {match.x, match.w, match.z}) { + auto * root = const_cast(input->view_src ? input->view_src : input); + params->add_alloc_dep(params->user_data, root, match.dst); + } + } + } + } + + // the direct convolution kernels read the history and x while writing the SiLU output + for (int i = 0; i < cgraph->n_nodes; ++i) { + if (cgraph->nodes[i]->op != GGML_OP_CONCAT) { + continue; + } + ggml_cuda_ple_conv_match pm; + ggml_cuda_gdn_conv_match gm; + const ggml_tensor * inputs[2] = { nullptr, nullptr }; + ggml_tensor * out = nullptr; + if (ggml_cuda_ple_conv_match_at_concat(cgraph, i, pm)) { + inputs[0] = pm.x; inputs[1] = pm.state; out = pm.out; + } else if (ggml_cuda_gdn_conv_match_at_concat(cgraph, i, gm)) { + inputs[0] = gm.x; inputs[1] = gm.state; out = gm.out; + } + for (const ggml_tensor * input : inputs) { + if (input) { + auto * root = const_cast(input->view_src ? input->view_src : input); + params->add_alloc_dep(params->user_data, root, out); + } + } + } + for (int i = 0; i < cgraph->n_nodes; ++i) { if (cgraph->nodes[i]->op != GGML_OP_MUL) { continue; @@ -4535,6 +6546,9 @@ static void ggml_backend_cuda_graph_optimize(ggml_backend_t backend, ggml_cgraph GGML_UNUSED(cgraph); #endif + // the marks for this split are complete here; the fingerprint must be recorded before the early returns below + cuda_ctx->mmb_graph_sigs.push_back(ggml_cuda_graph_sig(cgraph)); + static bool enable_graph_optimization = [] { const char * env = getenv("GGML_CUDA_GRAPH_OPT"); return env != nullptr && atoi(env) == 1; @@ -5674,8 +7688,17 @@ static ggml_backend_feature * ggml_backend_cuda_get_features(ggml_backend_reg_t GGML_UNUSED(reg); } +// opt this backend context in to the BF16 WMMA matmul path (mmb); call before the first graph +static void ggml_backend_cuda_set_mmb_enabled(ggml_backend_t backend, bool enable) { + GGML_ASSERT(ggml_backend_is_cuda(backend)); + ((ggml_backend_cuda_context *) backend->context)->mmb_opt_in = enable; +} + static void * ggml_backend_cuda_reg_get_proc_address(ggml_backend_reg_t reg, const char * name) { GGML_UNUSED(reg); + if (strcmp(name, "ggml_backend_cuda_set_mmb_enabled") == 0) { + return (void *)ggml_backend_cuda_set_mmb_enabled; + } if (strcmp(name, "ggml_backend_comm_init") == 0) { return (void *)ggml_backend_cuda_comm_init; } diff --git a/ggml/src/ggml-cuda/hyperconn.cu b/ggml/src/ggml-cuda/hyperconn.cu new file mode 100644 index 000000000000..c19128eb4672 --- /dev/null +++ b/ggml/src/ggml-cuda/hyperconn.cu @@ -0,0 +1,591 @@ +#include "hyperconn.cuh" +#include "mmb.cuh" +#include "ggml-backend-impl.h" + +// Explicitly rounded mul/add: keeps the fused kernels from being contracted into FMAs so that +// the results match the separate MUL and ADD kernels of the unfused graph bit for bit. +#if defined(__HIP_PLATFORM_AMD__) +static __device__ __forceinline__ float hc_mul_rn(const float a, const float b) { + float result; + asm("v_mul_f32_e32 %0, %1, %2" : "=v"(result) : "v"(a), "v"(b)); + return result; +} + +static __device__ __forceinline__ float hc_add_rn(const float a, const float b) { + float result; + asm("v_add_f32_e32 %0, %1, %2" : "=v"(result) : "v"(a), "v"(b)); + return result; +} +#else +static __device__ __forceinline__ float hc_mul_rn(const float a, const float b) { + return __fmul_rn(a, b); +} + +static __device__ __forceinline__ float hc_add_rn(const float a, const float b) { + return __fadd_rn(a, b); +} +#endif + +// same expression as op_sigmoid in unary.cu +static __device__ __forceinline__ float hc_sigmoid(const float x) { + return 1.0f / (1.0f + expf(-x)); +} + +// mixed[e,t] = scale * ( ((xn*sig)[0] + (xn*sig)[1]) + ... ) + bias, streams summed in order 0..hc-1 +static __global__ void hc_mix_reduce_f32( + const float * __restrict__ xn, const float * __restrict__ gate, float * __restrict__ dst, + const int64_t n_embd, const int64_t n_tokens, const int hc, const float scale, const float bias) { + const int64_t index = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; + if (index >= n_embd * n_tokens) { + return; + } + + const int64_t t = index / n_embd; + const int64_t e = index - t * n_embd; + const int64_t hc_dim = n_embd * hc; + const int64_t base = t * hc_dim + e; + + float acc = hc_mul_rn(xn[base], hc_sigmoid(gate[base])); + for (int c = 1; c < hc; ++c) { + const int64_t i = base + int64_t(c) * n_embd; + acc = hc_add_rn(acc, hc_mul_rn(xn[i], hc_sigmoid(gate[i]))); + } + + // identical expression to scale_f32 in scale.cu + dst[index] = scale * acc + bias; +} + +// out[e,c,t] = residual[e,c,t] + block_out[e,t] * ( s2 * sigmoid(s1 * inject[c,t] + b1) + b2 ) +static __global__ void hc_combine_f32( + const float * __restrict__ residual, const float * __restrict__ block_out, const float * __restrict__ inject, + float * __restrict__ dst, const int64_t n_embd, const int hc, const int64_t n_tokens, + const float s1, const float b1, const float s2, const float b2) { + const int64_t index = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; + const int64_t hc_dim = n_embd * hc; + if (index >= hc_dim * n_tokens) { + return; + } + + const int64_t t = index / hc_dim; + const int64_t rem = index - t * hc_dim; + const int64_t c = rem / n_embd; + const int64_t e = rem - c * n_embd; + + // identical expressions to scale_f32 (scale.cu) and op_sigmoid (unary.cu) + const float x1 = s1 * inject[t * hc + c] + b1; + const float x2 = hc_sigmoid(x1); + const float w = s2 * x2 + b2; + + const float m = hc_mul_rn(block_out[t * n_embd + e], w); + dst[index] = hc_add_rn(residual[index], m); +} + + +static __global__ void hc_mix_reduce_f32_hc4_parallel( + const float * __restrict__ xn, const float * __restrict__ gate, float * __restrict__ dst, + const int64_t n_embd, const int64_t n_tokens, const float scale, const float bias) { + const int lane = threadIdx.x % 32; + const int stream = threadIdx.x / 32; + const int64_t index = int64_t(blockIdx.x) * 32 + lane; + const int64_t count = n_embd * n_tokens; + __shared__ float products[4][32]; + float value = 0.0f; + if (index < count) { + const int64_t token = index / n_embd; + const int64_t embd = index - token * n_embd; + const int64_t offset = token * (4 * n_embd) + stream * n_embd + embd; + value = hc_mul_rn(xn[offset], hc_sigmoid(gate[offset])); + } + products[stream][lane] = value; + __syncthreads(); + if (stream == 0 && index < count) { + float sum = products[0][lane]; + sum = hc_add_rn(sum, products[1][lane]); + sum = hc_add_rn(sum, products[2][lane]); + sum = hc_add_rn(sum, products[3][lane]); + dst[index] = scale * sum + bias; + } +} + +__device__ __forceinline__ float hc_bf2f32(const uint16_t h) { return __uint_as_float(((uint32_t) h) << 16); } +__device__ __forceinline__ uint16_t hc_f2bf32(const float f) { uint32_t u = __float_as_uint(f); u += 0x7fffu + ((u >> 16) & 1u); return (uint16_t)(u >> 16); } + +// same reduction as hc_mix_reduce_f32, reading BF16 copies of xn and gate (16-bit HC intermediates) +static __global__ void hc_mix_reduce_bf16( + const uint16_t * __restrict__ xn, const uint16_t * __restrict__ gate, float * __restrict__ dst, + const int64_t n_embd, const int64_t n_tokens, const int hc, const float scale, const float bias) { + const int64_t index = int64_t(blockIdx.x) * blockDim.x + threadIdx.x; + if (index >= n_embd * n_tokens) return; + const int64_t t = index / n_embd, e = index - t * n_embd, base = t * (int64_t) n_embd * hc + e; + float acc = hc_mul_rn(hc_bf2f32(xn[base]), hc_sigmoid(hc_bf2f32(gate[base]))); + for (int c = 1; c < hc; ++c) { const int64_t i = base + int64_t(c) * n_embd; acc = hc_add_rn(acc, hc_mul_rn(hc_bf2f32(xn[i]), hc_sigmoid(hc_bf2f32(gate[i])))); } + dst[index] = scale * acc + bias; +} + +static bool hc_ranges_overlap(const ggml_tensor * a, const ggml_tensor * b) { + const char * a0 = (const char *) a->data; + const char * a1 = a0 + ggml_backend_buft_get_alloc_size(a->buffer->buft, a); + const char * b0 = (const char *) b->data; + const char * b1 = b0 + ggml_backend_buft_get_alloc_size(b->buffer->buft, b); + return a0 < b1 && b0 < a1; +} + +void ggml_cuda_op_hc_mix_reduce(ggml_backend_cuda_context & ctx, const ggml_cuda_hc_mix_args & args) { + const ggml_tensor * xn = args.xn; + const ggml_tensor * gate = args.gate; + ggml_tensor * dst = args.dst; + + GGML_ASSERT(xn->type == GGML_TYPE_F32 && gate->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(xn) && ggml_is_contiguous(gate) && ggml_is_contiguous(dst)); + GGML_ASSERT(args.hc > 0); + + const int64_t n_embd = dst->ne[0]; + const int64_t n_tokens = ggml_nrows(dst); + GGML_ASSERT(xn->ne[0] == n_embd * args.hc && ggml_nrows(xn) == n_tokens); + GGML_ASSERT(ggml_are_same_shape(xn, gate)); + + // The graph allocator may reuse the buffer of the (elided) sigmoid/mul intermediates, which can + // coincide with the gate matmul output. Every thread reads hc strided elements, so an aliased + // output has to be staged -- unless (single token) the output lands on a whole stream of the + // aliased input: element e of the output then overwrites element e of that stream, which only + // the same thread reads (after it has finished reading). + auto stream_aligned_alias = [&](const ggml_tensor * in) { + if (!hc_ranges_overlap(dst, in)) { + return true; + } + if (n_tokens != 1) { + return false; + } + const ptrdiff_t off = (const char *) dst->data - (const char *) in->data; + return off >= 0 && off % (n_embd * (ptrdiff_t) sizeof(float)) == 0 && + off + (ptrdiff_t) ggml_nbytes(dst) <= (ptrdiff_t) ggml_nbytes(in); + }; + const bool stage = !stream_aligned_alias(xn) || !stream_aligned_alias(gate); + + ggml_cuda_pool_alloc staged(ctx.pool()); + float * out = stage ? staged.alloc(n_embd * n_tokens) : (float *) dst->data; + + const int64_t n_items = n_embd * n_tokens; + const uint16_t * xn16 = ggml_cuda_mmb_cache_lookup(ctx, xn); + const uint16_t * g16 = ggml_cuda_mmb_cache_lookup(ctx, gate); + if (xn16 && g16) { + const int threads = 256; const int blocks = (int) ((n_items + threads - 1) / threads); + const ggml_cuda_kernel_launch_params launch_params(blocks, threads, 0, ctx.stream()); + ggml_cuda_kernel_launch(hc_mix_reduce_bf16, launch_params, xn16, g16, out, n_embd, n_tokens, args.hc, args.scale, args.bias); + } else if (args.hc == 4 && n_tokens <= 32) { + // one wave per stream, the same in-order sum: bit-identical to the serial kernel with more blocks in flight for decode + const int blocks = (int) ((n_items + 31) / 32); + const ggml_cuda_kernel_launch_params launch_params(blocks, 128, 0, ctx.stream()); + ggml_cuda_kernel_launch(hc_mix_reduce_f32_hc4_parallel, launch_params, + (const float *) xn->data, (const float *) gate->data, out, n_embd, n_tokens, args.scale, args.bias); + } else { + const int threads = 256; + const int blocks = (int) ((n_items + threads - 1) / threads); + const ggml_cuda_kernel_launch_params launch_params(blocks, threads, 0, ctx.stream()); + ggml_cuda_kernel_launch(hc_mix_reduce_f32, launch_params, + (const float *) xn->data, (const float *) gate->data, out, n_embd, n_tokens, args.hc, args.scale, args.bias); + } + + if (stage) { + CUDA_CHECK(cudaMemcpyAsync(dst->data, out, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, ctx.stream())); + } +} + +void ggml_cuda_op_hc_combine(ggml_backend_cuda_context & ctx, const ggml_cuda_hc_combine_args & args) { + const ggml_tensor * residual = args.residual; + const ggml_tensor * block_out = args.block_out; + const ggml_tensor * inject = args.inject; + ggml_tensor * dst = args.dst; + + GGML_ASSERT(residual->type == GGML_TYPE_F32 && block_out->type == GGML_TYPE_F32 && + inject->type == GGML_TYPE_F32 && dst->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(residual) && ggml_is_contiguous(block_out) && + ggml_is_contiguous(inject) && ggml_is_contiguous(dst)); + + const int64_t n_embd = dst->ne[0]; + const int64_t hc = dst->ne[1]; + const int64_t n_tokens = dst->ne[2] * dst->ne[3]; + GGML_ASSERT(ggml_are_same_shape(residual, dst)); + GGML_ASSERT(ggml_nelements(block_out) == n_embd * n_tokens); + GGML_ASSERT(ggml_nelements(inject) == hc * n_tokens); + + // residual may be the in-place source of the add (same address, same layout): each thread then + // reads and writes its own element only. Any other overlap needs staging. + const bool residual_inplace = residual->data == dst->data && ggml_are_same_layout(residual, dst); + const bool stage = (!residual_inplace && hc_ranges_overlap(dst, residual)) || + hc_ranges_overlap(dst, block_out) || hc_ranges_overlap(dst, inject); + + ggml_cuda_pool_alloc staged(ctx.pool()); + float * out = stage ? staged.alloc(ggml_nelements(dst)) : (float *) dst->data; + + const int64_t n_items = ggml_nelements(dst); + const int threads = 256; + const int blocks = (int) ((n_items + threads - 1) / threads); + const ggml_cuda_kernel_launch_params launch_params(blocks, threads, 0, ctx.stream()); + ggml_cuda_kernel_launch(hc_combine_f32, launch_params, + (const float *) residual->data, (const float *) block_out->data, (const float *) inject->data, out, + n_embd, (int) hc, n_tokens, args.s1, args.b1, args.s2, args.b2); + + if (stage) { + CUDA_CHECK(cudaMemcpyAsync(dst->data, out, ggml_nbytes(dst), cudaMemcpyDeviceToDevice, ctx.stream())); + } +} + +// --------------------------------------------------------------------------------------------------------- +// combine + grouped rms norm, one token: block per stream, 1024 threads (the rms_norm_f32<1024> layout) +// w = s2 * sigmoid(s1 * inject[c] + b1) + b2 +// res[e,c] = residual[e,c] + block_out[e] * w +// xn[e,c] = rms_norm(res[:,c]) * gamma[e,c] + +#define HC_CN_BLOCK 1024 +#define HC_CN_MAX_EMB 3072 + +static __global__ void __launch_bounds__(HC_CN_BLOCK, 1) hc_combine_norm_f32( + const float * __restrict__ inject, const float * __restrict__ residual, + const float * __restrict__ block_out, const float * __restrict__ gamma, + float * __restrict__ out_res, float * __restrict__ out_xn, + const int n_embd, const float s1, const float b1, const float s2, const float b2, const float eps) { + __shared__ float s_sum[32]; + + const int c = blockIdx.x; // stream + const int t = blockIdx.y; // token + const int hc = gridDim.x; + const int tid = threadIdx.x; + + // identical expressions to scale_f32 / op_sigmoid / scale_f32 + const float x1 = s1 * inject[(int64_t) t * hc + c] + b1; + const float x2 = hc_sigmoid(x1); + const float w = s2 * x2 + b2; + + const int64_t row = (int64_t) t * hc + c; + const float * res = residual + row * n_embd; + float * dst = out_res + row * n_embd; + block_out += (int64_t) t * n_embd; + float xs[3]; + float tmp = 0.0f; +#pragma unroll + for (int k = 0; k < 3; ++k) { + const int col = tid + k * HC_CN_BLOCK; + xs[k] = 0.0f; + if (col < n_embd) { + const float m = hc_mul_rn(block_out[col], w); + const float xi = hc_add_rn(res[col], m); + dst[col] = xi; + xs[k] = xi; + tmp += xi * xi; + } + } + + tmp = block_reduce(tmp, s_sum); + + const float mean = tmp / n_embd; + const float scale = rsqrtf(mean + eps); + + const float * g = gamma + (int64_t) c * n_embd; + float * xn = out_xn + row * n_embd; +#pragma unroll + for (int k = 0; k < 3; ++k) { + const int col = tid + k * HC_CN_BLOCK; + if (col < n_embd) { + xn[col] = scale * xs[k] * g[col]; + } + } +} + +// Read all residual, block and injection values before writing aliased outputs. +// Gamma must remain disjoint from both outputs. +#define HC_CN_SINGLE_MAX_HC 4 + +static __global__ void __launch_bounds__(HC_CN_BLOCK, 1) hc_combine_norm_single_f32( + const float * inject, const float * residual, const float * block_out, const float * __restrict__ gamma, + float * out_res, float * out_xn, + const int n_embd, const int hc, const float s1, const float b1, const float s2, const float b2, const float eps) { + __shared__ float s_sum[HC_CN_SINGLE_MAX_HC][32]; + + const int t = blockIdx.x; // token + const int tid = threadIdx.x; + const int warp_id = tid / WARP_SIZE; + const int lane_id = tid % WARP_SIZE; + + // 1. every input element of the token into registers + float w[HC_CN_SINGLE_MAX_HC]; +#pragma unroll + for (int c = 0; c < HC_CN_SINGLE_MAX_HC; ++c) { + w[c] = 0.0f; + if (c < hc) { + // identical expressions to scale_f32 / op_sigmoid / scale_f32 + const float x1 = s1 * inject[(int64_t) t * hc + c] + b1; + const float x2 = hc_sigmoid(x1); + w[c] = s2 * x2 + b2; + } + } + + block_out += (int64_t) t * n_embd; + float xs[HC_CN_SINGLE_MAX_HC][3]; + float tmp[HC_CN_SINGLE_MAX_HC]; +#pragma unroll + for (int c = 0; c < HC_CN_SINGLE_MAX_HC; ++c) { + tmp[c] = 0.0f; + } +#pragma unroll + for (int k = 0; k < 3; ++k) { + const int col = tid + k * HC_CN_BLOCK; + const float bo = col < n_embd ? block_out[col] : 0.0f; +#pragma unroll + for (int c = 0; c < HC_CN_SINGLE_MAX_HC; ++c) { + xs[c][k] = 0.0f; + if (c < hc && col < n_embd) { + const float m = hc_mul_rn(bo, w[c]); + const float xi = hc_add_rn(residual[((int64_t) t * hc + c) * n_embd + col], m); + xs[c][k] = xi; + tmp[c] += xi * xi; + } + } + } + __syncthreads(); // all loads of the block have completed before the first store + + // 2. per-stream block sums, same reduction tree as block_reduce for each stream +#pragma unroll + for (int c = 0; c < HC_CN_SINGLE_MAX_HC; ++c) { + tmp[c] = warp_reduce_sum(tmp[c]); + } + if (lane_id == 0) { +#pragma unroll + for (int c = 0; c < HC_CN_SINGLE_MAX_HC; ++c) { + s_sum[c][warp_id] = tmp[c]; + } + } + __syncthreads(); +#pragma unroll + for (int c = 0; c < HC_CN_SINGLE_MAX_HC; ++c) { + tmp[c] = lane_id < HC_CN_BLOCK / WARP_SIZE ? s_sum[c][lane_id] : 0.0f; + tmp[c] = warp_reduce_sum(tmp[c]); + } + + // Finish residual writes before an overlapping normalized output replaces them. +#pragma unroll + for (int c = 0; c < HC_CN_SINGLE_MAX_HC; ++c) { + if (c < hc) { +#pragma unroll + for (int k = 0; k < 3; ++k) { + const int col = tid + k * HC_CN_BLOCK; + if (col < n_embd) { + out_res[((int64_t) t * hc + c) * n_embd + col] = xs[c][k]; + } + } + } + } + __syncthreads(); + +#pragma unroll + for (int c = 0; c < HC_CN_SINGLE_MAX_HC; ++c) { + if (c < hc) { + const float mean = tmp[c] / n_embd; + const float scale = rsqrtf(mean + eps); + const int64_t row = (int64_t) t * hc + c; + float * xn = out_xn + row * n_embd; + const float * g = gamma + (int64_t) c * n_embd; +#pragma unroll + for (int k = 0; k < 3; ++k) { + const int col = tid + k * HC_CN_BLOCK; + if (col < n_embd) { + xn[col] = scale * xs[c][k] * g[col]; + } + } + } + } +} + +#define HC_CN_BLOCK2 256 +#define HC_INJ_HC 4 // streams (and inject outputs) the fused combine + inject kernel handles +__device__ __forceinline__ float hc_lo(const uint32_t u) { return __uint_as_float(u << 16); } +__device__ __forceinline__ float hc_hi(const uint32_t u) { return __uint_as_float(u & 0xffff0000u); } +__device__ __forceinline__ uint32_t hc_pack2(const float a, const float b) { + return (uint32_t) hc_f2bf32(a) | ((uint32_t) hc_f2bf32(b) << 16); +} + +// One (token, stream) row of the combine + norm; with INJECT it also accumulates the row's share of the next mix's inject projection. +template +static __device__ __forceinline__ void hc_combine_norm_row_b256( + const int c, const int t, const int hc, float * s_sum, + const float * inject, const float * residual, + const float * block_out, const float * gamma, + float * out_res, float * out_xn, uint16_t * out_xn_bf16, const bool store_xn_f32, + const uint16_t * res_in_bf16, uint16_t * res_out_bf16, const uint16_t * blk_in_bf16, + const int n_embd, const float s1, const float b1, const float s2, const float b2, const float eps, + const float * w_inj, float * inj) { + const int tid = threadIdx.x; + const float x1 = s1 * inject[(int64_t) t * hc + c] + b1; + const float x2 = hc_sigmoid(x1); + const float w = s2 * x2 + b2; + const int64_t row = (int64_t) t * hc + c; + const float * res = residual + row * n_embd; + const uint16_t * res16 = res_in_bf16 ? res_in_bf16 + row * n_embd : nullptr; + float * dst = out_res + row * n_embd; + uint16_t * dst16 = res_out_bf16 ? res_out_bf16 + row * n_embd : nullptr; + const float * blk = block_out + (int64_t) t * n_embd; + const uint16_t * blk16 = blk_in_bf16 ? blk_in_bf16 + (int64_t) t * n_embd : nullptr; + constexpr int KP = (HC_CN_MAX_EMB / 2 + HC_CN_BLOCK2 - 1) / HC_CN_BLOCK2; + float xs[2 * KP]; + float tmp = 0.0f; +#pragma unroll + for (int k = 0; k < KP; ++k) { + const int col = (tid + k * HC_CN_BLOCK2) * 2; + xs[2 * k] = 0.0f; xs[2 * k + 1] = 0.0f; + if (col + 1 < n_embd) { + float r0, r1, m0, m1; + if (res16) { const uint32_t u = *(const uint32_t *)(res16 + col); r0 = hc_lo(u); r1 = hc_hi(u); } + else { const float2 v = *(const float2 *)(res + col); r0 = v.x; r1 = v.y; } + if (blk16) { const uint32_t u = *(const uint32_t *)(blk16 + col); m0 = hc_lo(u); m1 = hc_hi(u); } + else { const float2 v = *(const float2 *)(blk + col); m0 = v.x; m1 = v.y; } + const float a0 = hc_add_rn(r0, hc_mul_rn(m0, w)); + const float a1 = hc_add_rn(r1, hc_mul_rn(m1, w)); + if (dst16) *(uint32_t *)(dst16 + col) = hc_pack2(a0, a1); + else *(float2 *)(dst + col) = make_float2(a0, a1); + xs[2 * k] = a0; xs[2 * k + 1] = a1; + tmp += a0 * a0 + a1 * a1; + } else if (col < n_embd) { + const float r0 = res16 ? hc_bf2f32(res16[col]) : res[col]; + const float m0 = blk16 ? hc_bf2f32(blk16[col]) : blk[col]; + const float a0 = hc_add_rn(r0, hc_mul_rn(m0, w)); + if (dst16) dst16[col] = hc_f2bf32(a0); else dst[col] = a0; + xs[2 * k] = a0; + tmp += a0 * a0; + } + } + tmp = block_reduce(tmp, s_sum); + const float mean = tmp / n_embd; + const float scale = rsqrtf(mean + eps); + const float * g = gamma + (int64_t) c * n_embd; + float * xn = out_xn + row * n_embd; + uint16_t * xh = out_xn_bf16 ? out_xn_bf16 + row * n_embd : nullptr; +#pragma unroll + for (int k = 0; k < KP; ++k) { + const int col = (tid + k * HC_CN_BLOCK2) * 2; + if (col + 1 < n_embd) { + const float2 gv = *(const float2 *)(g + col); + const float v0 = scale * xs[2 * k] * gv.x, v1 = scale * xs[2 * k + 1] * gv.y; + if (store_xn_f32) *(float2 *)(xn + col) = make_float2(v0, v1); + if (xh) *(uint32_t *)(xh + col) = hc_pack2(v0, v1); + if constexpr (INJECT) { +#pragma unroll + for (int r = 0; r < HC_INJ_HC; ++r) { + const float2 wv = *(const float2 *)(w_inj + (int64_t) r * hc * n_embd + (int64_t) c * n_embd + col); + inj[r] += wv.x * v0 + wv.y * v1; + } + } + } else if (col < n_embd) { + const float v0 = scale * xs[2 * k] * g[col]; + if (store_xn_f32) xn[col] = v0; + if (xh) xh[col] = hc_f2bf32(v0); + if constexpr (INJECT) { +#pragma unroll + for (int r = 0; r < HC_INJ_HC; ++r) { + inj[r] += w_inj[(int64_t) r * hc * n_embd + (int64_t) c * n_embd + col] * v0; + } + } + } + } +} + +static __global__ void __launch_bounds__(HC_CN_BLOCK2, 4) hc_combine_norm_f32_b256( + const float * inject, const float * residual, + const float * block_out, const float * gamma, + float * out_res, float * out_xn, uint16_t * out_xn_bf16, const bool store_xn_f32, + const uint16_t * res_in_bf16, uint16_t * res_out_bf16, const uint16_t * blk_in_bf16, + const int n_embd, const float s1, const float b1, const float s2, const float b2, const float eps) { + __shared__ float s_sum[32]; + hc_combine_norm_row_b256(blockIdx.x, blockIdx.y, gridDim.x, s_sum, inject, residual, block_out, gamma, + out_res, out_xn, out_xn_bf16, store_xn_f32, res_in_bf16, res_out_bf16, blk_in_bf16, + n_embd, s1, b1, s2, b2, eps, nullptr, nullptr); +} + +// One block per token over all HC_INJ_HC streams: the rows of hc_combine_norm_f32_b256 plus the next mix's inject projection. +static __global__ void __launch_bounds__(HC_CN_BLOCK2, 2) hc_combine_norm_inject_f32_b256( + const float * inject, const float * residual, + const float * block_out, const float * gamma, + float * out_res, float * out_xn, uint16_t * out_xn_bf16, const bool store_xn_f32, + const uint16_t * res_in_bf16, uint16_t * res_out_bf16, const uint16_t * blk_in_bf16, + const int n_embd, const float s1, const float b1, const float s2, const float b2, const float eps, + const float * w_inj, float * out_inj) { + __shared__ float s_sum[HC_INJ_HC][32]; + __shared__ float s_inj[HC_INJ_HC][HC_CN_BLOCK2 / WARP_SIZE]; + const int t = blockIdx.x; + float inj[HC_INJ_HC] = {}; +#pragma unroll + for (int c = 0; c < HC_INJ_HC; ++c) { + hc_combine_norm_row_b256(c, t, HC_INJ_HC, s_sum[c], inject, residual, block_out, gamma, + out_res, out_xn, out_xn_bf16, store_xn_f32, res_in_bf16, res_out_bf16, blk_in_bf16, + n_embd, s1, b1, s2, b2, eps, w_inj, inj); + } + const int warp = threadIdx.x / WARP_SIZE, lane = threadIdx.x % WARP_SIZE; +#pragma unroll + for (int r = 0; r < HC_INJ_HC; ++r) { + const float v = warp_reduce_sum(inj[r]); + if (lane == 0) { + s_inj[r][warp] = v; + } + } + __syncthreads(); + if (threadIdx.x < HC_INJ_HC) { + float sum = 0.0f; +#pragma unroll + for (int w = 0; w < HC_CN_BLOCK2 / WARP_SIZE; ++w) { + sum += s_inj[threadIdx.x][w]; + } + out_inj[(int64_t) t * HC_INJ_HC + threadIdx.x] = sum; + } +} + +bool ggml_cuda_hc_combine_norm_supported(const ggml_cuda_hc_combine_norm_args & a, const int warp_size) { + const int64_t n_embd = a.out_res->ne[0]; + const int64_t hc = a.out_res->ne[1]; + return warp_size == 32 && n_embd >= 1024 && n_embd <= HC_CN_MAX_EMB && hc >= 1 && hc <= 16 && + ggml_nelements(a.gamma) == n_embd * hc; +} + +void ggml_cuda_op_hc_combine_norm(ggml_backend_cuda_context & ctx, const ggml_cuda_hc_combine_norm_args & a) { + const int64_t n_embd = a.out_res->ne[0]; + const int64_t hc = a.out_res->ne[1]; + const int64_t n_tokens = a.out_res->ne[2] * a.out_res->ne[3]; + + GGML_ASSERT(a.out_res->ne[3] == 1); + GGML_ASSERT(ggml_are_same_shape(a.out_res, a.out_xn) && ggml_are_same_shape(a.out_res, a.residual)); + GGML_ASSERT(ggml_nelements(a.block_out) == n_embd * n_tokens && ggml_nelements(a.gamma) == n_embd * hc && + ggml_nelements(a.inject) == hc * n_tokens); + + if (a.single_block) { + GGML_ASSERT(hc <= HC_CN_SINGLE_MAX_HC); + GGML_ASSERT(a.store_xn_f32 && !a.out_xn_bf16 && !a.res_in_bf16 && !a.res_out_bf16 && !a.blk_in_bf16); + const ggml_cuda_kernel_launch_params launch_params(dim3((int) n_tokens, 1, 1), HC_CN_BLOCK, 0, ctx.stream()); + ggml_cuda_kernel_launch(hc_combine_norm_single_f32, launch_params, + (const float *) a.inject->data, (const float *) a.residual->data, + (const float *) a.block_out->data, (const float *) a.gamma->data, + (float *) a.out_res->data, (float *) a.out_xn->data, + (int) n_embd, (int) hc, a.s1, a.b1, a.s2, a.b2, a.eps); + return; + } + + if (a.out_inject) { + GGML_ASSERT(hc == HC_INJ_HC && a.w_inject && ggml_nelements(a.out_inject) == hc * n_tokens && + a.w_inject->ne[0] == n_embd * hc && a.w_inject->ne[1] == hc); + const ggml_cuda_kernel_launch_params lpi(dim3((int) n_tokens, 1, 1), HC_CN_BLOCK2, 0, ctx.stream()); + ggml_cuda_kernel_launch(hc_combine_norm_inject_f32_b256, lpi, + (const float *) a.inject->data, (const float *) a.residual->data, + (const float *) a.block_out->data, (const float *) a.gamma->data, + (float *) a.out_res->data, (float *) a.out_xn->data, a.out_xn_bf16, a.store_xn_f32, + a.res_in_bf16, a.res_out_bf16, a.blk_in_bf16, + (int) n_embd, a.s1, a.b1, a.s2, a.b2, a.eps, + (const float *) a.w_inject->data, (float *) a.out_inject->data); + return; + } + + // 256 threads x 2 packed elements per iteration (the shipped layout); the BF16 options ride on this kernel + const ggml_cuda_kernel_launch_params lp2(dim3((int) hc, (int) n_tokens, 1), HC_CN_BLOCK2, 0, ctx.stream()); + ggml_cuda_kernel_launch(hc_combine_norm_f32_b256, lp2, + (const float *) a.inject->data, (const float *) a.residual->data, + (const float *) a.block_out->data, (const float *) a.gamma->data, + (float *) a.out_res->data, (float *) a.out_xn->data, a.out_xn_bf16, a.store_xn_f32, + a.res_in_bf16, a.res_out_bf16, a.blk_in_bf16, + (int) n_embd, a.s1, a.b1, a.s2, a.b2, a.eps); +} diff --git a/ggml/src/ggml-cuda/hyperconn.cuh b/ggml/src/ggml-cuda/hyperconn.cuh new file mode 100644 index 000000000000..0d76c04a5f14 --- /dev/null +++ b/ggml/src/ggml-cuda/hyperconn.cuh @@ -0,0 +1,63 @@ +#pragma once + +#include "common.cuh" + +// Fused hyper-connection (DeepSeek-V4 / qwen4exp style) decode helpers. +// +// hc_mix_reduce: mixed[e,t] = scale * sum_c ( xn[c*n_embd + e, t] * sigmoid(gate[c*n_embd + e, t]) ) + bias +// replaces: UNARY(SIGMOID) -> MUL -> RESHAPE -> VIEW -> CONT -> (VIEW -> ADD) x (hc-1) -> SCALE +// +// hc_combine: out[e,c,t] = residual[e,c,t] + block_out[e,t] * ( s2 * sigmoid(s1 * inject[c,t] + b1) + b2 ) +// replaces: SCALE -> UNARY(SIGMOID) -> SCALE -> RESHAPE* -> REPEAT -> MUL -> ADD +// +// Both replay the unfused rounding sequence exactly (separate mul and add roundings), so results are +// bit-identical to the individual kernels. + +struct ggml_cuda_hc_mix_args { + const ggml_tensor * xn; // [hc*n_embd, T] F32 contiguous + const ggml_tensor * gate; // [hc*n_embd, T] F32 contiguous (pre-sigmoid) + ggml_tensor * dst; // [n_embd, T] F32 contiguous + int hc; + float scale; + float bias; +}; + +struct ggml_cuda_hc_combine_args { + const ggml_tensor * residual; // [n_embd, hc, T] F32 contiguous + const ggml_tensor * block_out; // [n_embd, T] F32 contiguous (or a reshape of it) + const ggml_tensor * inject; // [hc, T] F32 contiguous (pre-scale/sigmoid) + ggml_tensor * dst; // [n_embd, hc, T] F32 contiguous + float s1, b1; // first SCALE + float s2, b2; // second SCALE +}; + +void ggml_cuda_op_hc_mix_reduce(ggml_backend_cuda_context & ctx, const ggml_cuda_hc_mix_args & args); +void ggml_cuda_op_hc_combine(ggml_backend_cuda_context & ctx, const ggml_cuda_hc_combine_args & args); + +// Hyper-connection combine followed by the grouped RMS norm of the next hc_mix (block per stream and token): +// w[c] = s2 * sigmoid(s1 * inject[c] + b1) + b2 +// res[e,c] = residual[e,c] + block_out[e] * w[c] -> out_res (the layer residual) +// xn[e,c] = rms_norm(res[:,c]) * gamma[e,c] -> out_xn (next hc_norm) +// The norm replays rms_norm_f32<1024, true> (one block of 1024 threads per stream). +struct ggml_cuda_hc_combine_norm_args { + const ggml_tensor * inject; // [hc, T] F32 contiguous (pre-scale/sigmoid) + const ggml_tensor * residual; // [n_embd, hc, T] F32 contiguous + const ggml_tensor * block_out; // [n_embd, 1, T] F32 contiguous + const ggml_tensor * gamma; // [n_embd, hc] F32 contiguous + ggml_tensor * out_res; // [n_embd, hc, T] + ggml_tensor * out_xn; // [n_embd, hc, T] + float s1, b1, s2, b2; + float eps; + bool single_block = false; // alias-proof variant: one block per token, all inputs read before any store (hc <= 4) + // BF16 stream options (HC16): a BF16 copy of xn for consumers that read it, and BF16 in-place residual / block_out + uint16_t * out_xn_bf16 = nullptr; + bool store_xn_f32 = true; // false: every consumer reads the BF16 copy + const uint16_t * res_in_bf16 = nullptr; // `residual` is BF16 in place (marked bf16-only) + const uint16_t * blk_in_bf16 = nullptr; + uint16_t * res_out_bf16 = nullptr; + const ggml_tensor * w_inject = nullptr; // [hc * n_embd, hc] F32 weight + ggml_tensor * out_inject = nullptr; // [hc, T] F32, the skipped MUL_MAT node +}; + +bool ggml_cuda_hc_combine_norm_supported(const ggml_cuda_hc_combine_norm_args & args, int warp_size); +void ggml_cuda_op_hc_combine_norm(ggml_backend_cuda_context & ctx, const ggml_cuda_hc_combine_norm_args & args); diff --git a/ggml/src/ggml-cuda/idx-relu-sum.cu b/ggml/src/ggml-cuda/idx-relu-sum.cu new file mode 100644 index 000000000000..6cd833bcbfc3 --- /dev/null +++ b/ggml/src/ggml-cuda/idx-relu-sum.cu @@ -0,0 +1,25 @@ +#include "idx-relu-sum.cuh" + +// Preserve the head addition order of the RELU/CONT/ADD graph. +static __global__ void idx_relu_sum_f32(const float * __restrict__ src, float * __restrict__ dst, + const int n_blocks, const int heads) { + const int b = blockIdx.x * blockDim.x + threadIdx.x; + if (b >= n_blocks) return; + const int64_t tt = blockIdx.y; // token within stream, streams stacked + const float * p = src + tt * (int64_t) heads * n_blocks + b; + float acc = fmaxf(p[0], 0.0f); + for (int h = 1; h < heads; ++h) { + acc = acc + fmaxf(p[(int64_t) h * n_blocks], 0.0f); + } + dst[tt * (int64_t) n_blocks + b] = acc; +} + +void ggml_cuda_op_idx_relu_sum(ggml_backend_cuda_context & ctx, const ggml_cuda_idx_relu_sum_args & args) { + const ggml_tensor * s = args.score; + const int n_blocks = (int) s->ne[0]; + const int64_t rows = s->ne[2] * s->ne[3]; + constexpr int threads = 256; + const dim3 grid((n_blocks + threads - 1) / threads, (unsigned) rows, 1); + idx_relu_sum_f32<<>>((const float *) s->data, (float *) args.dst->data, n_blocks, args.heads); + CUDA_CHECK(cudaGetLastError()); +} diff --git a/ggml/src/ggml-cuda/idx-relu-sum.cuh b/ggml/src/ggml-cuda/idx-relu-sum.cuh new file mode 100644 index 000000000000..04f902bac221 --- /dev/null +++ b/ggml/src/ggml-cuda/idx-relu-sum.cuh @@ -0,0 +1,8 @@ +#include "common.cuh" + +struct ggml_cuda_idx_relu_sum_args { + const ggml_tensor * score = nullptr; + ggml_tensor * dst = nullptr; // [n_blocks, n_tps, n_stream] + int heads = 0; +}; +void ggml_cuda_op_idx_relu_sum(ggml_backend_cuda_context & ctx, const ggml_cuda_idx_relu_sum_args & args); diff --git a/ggml/src/ggml-cuda/mean.cu b/ggml/src/ggml-cuda/mean.cu index 64ad7e1d534b..709cc96cdc00 100644 --- a/ggml/src/ggml-cuda/mean.cu +++ b/ggml/src/ggml-cuda/mean.cu @@ -2,7 +2,12 @@ #include "reduce_rows.cuh" #ifdef GGML_CUDA_USE_CUB +#ifdef GGML_CUDA_CUB_IS_HIPCUB +#include +namespace cub = hipcub; +#else #include +#endif // GGML_CUDA_CUB_IS_HIPCUB using namespace cub; #endif // GGML_CUDA_USE_CUB diff --git a/ggml/src/ggml-cuda/mmb-quant.cuh b/ggml/src/ggml-cuda/mmb-quant.cuh new file mode 100644 index 000000000000..6bb8812630d4 --- /dev/null +++ b/ggml/src/ggml-cuda/mmb-quant.cuh @@ -0,0 +1,336 @@ +#pragma once +#include "dequantize.cuh" +#include + +struct mmb_quant_slice { + uint16_t * dst; + int begin; + int offset = 0; + struct element { + uint16_t * dst; + int index; + __device__ void operator=(float value) const { + if (index >= 0 && index < 64) dst[index] = mmb_f2bf(value); + } + }; + __device__ mmb_quant_slice operator+(int64_t n) const { return {dst, begin, offset + (int) n}; } + __device__ element operator[](int64_t n) const { return {dst, offset + (int) n - begin}; } +}; + +template +__device__ __forceinline__ void mmb_decode_slice(const void * row, const int k0, uint16_t * dst, const int lane) { + constexpr int QK = ggml_cuda_type_traits::qk; + if constexpr (TYPE == GGML_TYPE_Q1_0) { + constexpr int QR = ggml_cuda_type_traits::qr; +#pragma unroll + for (int p = lane; p < 32; p += 8) { + const int pos = k0 + 2 * p, ib = pos / QK, qs = (pos % QK) / QR; + float2 v; + dequantize_q1_0(row, ib, qs, v); + const int o = ib * QK + qs - k0; + dst[o] = mmb_f2bf(v.x); + dst[o + (QR == 1 ? 1 : QK / 2)] = mmb_f2bf(v.y); + } + } + else if constexpr (TYPE == GGML_TYPE_Q2_0) { + constexpr int QR = ggml_cuda_type_traits::qr; +#pragma unroll + for (int p = lane; p < 32; p += 8) { + const int pos = k0 + 2 * p, ib = pos / QK, qs = (pos % QK) / QR; + float2 v; + dequantize_q2_0(row, ib, qs, v); + const int o = ib * QK + qs - k0; + dst[o] = mmb_f2bf(v.x); + dst[o + (QR == 1 ? 1 : QK / 2)] = mmb_f2bf(v.y); + } + } + else if constexpr (TYPE == GGML_TYPE_Q4_0) { + constexpr int QR = ggml_cuda_type_traits::qr; +#pragma unroll + for (int p = lane; p < 32; p += 8) { + const int pos = k0 + 2 * p, ib = pos / QK, qs = (pos % QK) / QR; + float2 v; + dequantize_q4_0(row, ib, qs, v); + const int o = ib * QK + qs - k0; + dst[o] = mmb_f2bf(v.x); + dst[o + (QR == 1 ? 1 : QK / 2)] = mmb_f2bf(v.y); + } + } + else if constexpr (TYPE == GGML_TYPE_Q4_1) { + constexpr int QR = ggml_cuda_type_traits::qr; +#pragma unroll + for (int p = lane; p < 32; p += 8) { + const int pos = k0 + 2 * p, ib = pos / QK, qs = (pos % QK) / QR; + float2 v; + dequantize_q4_1(row, ib, qs, v); + const int o = ib * QK + qs - k0; + dst[o] = mmb_f2bf(v.x); + dst[o + (QR == 1 ? 1 : QK / 2)] = mmb_f2bf(v.y); + } + } + else if constexpr (TYPE == GGML_TYPE_Q5_0) { + constexpr int QR = ggml_cuda_type_traits::qr; +#pragma unroll + for (int p = lane; p < 32; p += 8) { + const int pos = k0 + 2 * p, ib = pos / QK, qs = (pos % QK) / QR; + float2 v; + dequantize_q5_0(row, ib, qs, v); + const int o = ib * QK + qs - k0; + dst[o] = mmb_f2bf(v.x); + dst[o + (QR == 1 ? 1 : QK / 2)] = mmb_f2bf(v.y); + } + } + else if constexpr (TYPE == GGML_TYPE_Q5_1) { + constexpr int QR = ggml_cuda_type_traits::qr; +#pragma unroll + for (int p = lane; p < 32; p += 8) { + const int pos = k0 + 2 * p, ib = pos / QK, qs = (pos % QK) / QR; + float2 v; + dequantize_q5_1(row, ib, qs, v); + const int o = ib * QK + qs - k0; + dst[o] = mmb_f2bf(v.x); + dst[o + (QR == 1 ? 1 : QK / 2)] = mmb_f2bf(v.y); + } + } + else if constexpr (TYPE == GGML_TYPE_Q8_0) { + constexpr int QR = ggml_cuda_type_traits::qr; +#pragma unroll + for (int p = lane; p < 32; p += 8) { + const int pos = k0 + 2 * p, ib = pos / QK, qs = (pos % QK) / QR; + float2 v; + dequantize_q8_0(row, ib, qs, v); + const int o = ib * QK + qs - k0; + dst[o] = mmb_f2bf(v.x); + dst[o + (QR == 1 ? 1 : QK / 2)] = mmb_f2bf(v.y); + } + } + else if constexpr (TYPE == GGML_TYPE_Q2_K) { + const mmb_quant_slice out{dst, k0 % QK}; +#pragma unroll + for (int tid = lane; tid < 64; tid += 8) dequantize_q2_K(row, k0 / QK, out, tid); + } + else if constexpr (TYPE == GGML_TYPE_Q3_K) { + const mmb_quant_slice out{dst, k0 % QK}; +#pragma unroll + for (int tid = lane; tid < 64; tid += 8) dequantize_q3_K(row, k0 / QK, out, tid); + } + else if constexpr (TYPE == GGML_TYPE_Q4_K) { + const mmb_quant_slice out{dst, k0 % QK}; + const int tid = (k0 % QK) / 64 * 8 + lane; + dequantize_q4_K(row, k0 / QK, out, tid); + } + else if constexpr (TYPE == GGML_TYPE_Q5_K) { + const mmb_quant_slice out{dst, k0 % QK}; +#pragma unroll + for (int tid = lane; tid < 64; tid += 8) dequantize_q5_K(row, k0 / QK, out, tid); + } + else if constexpr (TYPE == GGML_TYPE_Q6_K) { + const mmb_quant_slice out{dst, k0 % QK}; +#pragma unroll + for (int tid = lane; tid < 64; tid += 8) dequantize_q6_K(row, k0 / QK, out, tid); + } + else if constexpr (TYPE == GGML_TYPE_IQ1_S) { + const mmb_quant_slice out{dst, k0 % QK}; +#pragma unroll + for (int tid = lane; tid < 32; tid += 8) dequantize_iq1_s(row, k0 / QK, out, tid); + } + else if constexpr (TYPE == GGML_TYPE_IQ1_M) { + const mmb_quant_slice out{dst, k0 % QK}; +#pragma unroll + for (int tid = lane; tid < 32; tid += 8) dequantize_iq1_m(row, k0 / QK, out, tid); + } + else if constexpr (TYPE == GGML_TYPE_IQ2_XXS) { + const mmb_quant_slice out{dst, k0 % QK}; +#pragma unroll + for (int tid = lane; tid < 32; tid += 8) dequantize_iq2_xxs(row, k0 / QK, out, tid); + } + else if constexpr (TYPE == GGML_TYPE_IQ2_XS) { + const mmb_quant_slice out{dst, k0 % QK}; +#pragma unroll + for (int tid = lane; tid < 32; tid += 8) dequantize_iq2_xs(row, k0 / QK, out, tid); + } + else if constexpr (TYPE == GGML_TYPE_IQ2_S) { + const mmb_quant_slice out{dst, k0 % QK}; +#pragma unroll + for (int tid = lane; tid < 32; tid += 8) dequantize_iq2_s(row, k0 / QK, out, tid); + } + else if constexpr (TYPE == GGML_TYPE_IQ3_XXS) { + const mmb_quant_slice out{dst, k0 % QK}; +#pragma unroll + for (int tid = lane; tid < 32; tid += 8) dequantize_iq3_xxs(row, k0 / QK, out, tid); + } + else if constexpr (TYPE == GGML_TYPE_IQ3_S) { + const mmb_quant_slice out{dst, k0 % QK}; +#pragma unroll + for (int tid = lane; tid < 32; tid += 8) dequantize_iq3_s(row, k0 / QK, out, tid); + } + else if constexpr (TYPE == GGML_TYPE_IQ4_XS) { + const mmb_quant_slice out{dst, k0 % QK}; +#pragma unroll + for (int tid = lane; tid < 32; tid += 8) dequantize_iq4_xs(row, k0 / QK, out, tid); + } + else if constexpr (TYPE == GGML_TYPE_IQ4_NL) { +#pragma unroll + for (int p = lane; p < 32; p += 8) { + const int pos = k0 + 2 * p, ib = pos / 32, q = (pos % 32) / 2; + const block_iq4_nl & b = ((const block_iq4_nl *) row)[ib]; + const int o = ib * 32 + q - k0; + dst[o] = mmb_f2bf((float)b.d * kvalues_iq4nl[b.qs[q] & 15]); + dst[o + 16] = mmb_f2bf((float)b.d * kvalues_iq4nl[b.qs[q] >> 4]); + } + } + else if constexpr (TYPE == GGML_TYPE_MXFP4) { +#pragma unroll + for (int p = lane; p < 32; p += 8) { + const int pos = k0 + 2 * p, ib = pos / QK, q = (pos % QK) / 2; + const block_mxfp4 & b = ((const block_mxfp4 *) row)[ib]; + const float d = ggml_cuda_e8m0_to_fp32(b.e); + const int o = ib * QK + q - k0; + dst[o] = mmb_f2bf(d * kvalues_mxfp4[b.qs[q] & 15] * 0.5f); + dst[o + QK / 2] = mmb_f2bf(d * kvalues_mxfp4[b.qs[q] >> 4] * 0.5f); + } + } + else if constexpr (TYPE == GGML_TYPE_NVFP4) { +#pragma unroll + for (int p = lane; p < 32; p += 8) { + const int pos = k0 + 2 * p, ib = pos / QK, in = pos % QK; + const int sub = in / QK_NVFP4_SUB, j = (in % QK_NVFP4_SUB) / 2; + const block_nvfp4 & b = ((const block_nvfp4 *) row)[ib]; + const float d = ggml_cuda_ue4m3_to_fp32(b.d[sub]); + const uint8_t q = b.qs[sub * (QK_NVFP4_SUB / 2) + j]; + const int o = ib * QK + sub * QK_NVFP4_SUB + j - k0; + dst[o] = mmb_f2bf(d * kvalues_mxfp4[q & 15]); + dst[o + QK_NVFP4_SUB / 2] = mmb_f2bf(d * kvalues_mxfp4[q >> 4]); + } + } +} + +template +__host__ __device__ constexpr size_t mmb_row_bytes(int k) { + if constexpr (WTYPE == 0) return (size_t)(k / 32) * 18; + else if constexpr (WTYPE == 1) return (size_t)(k / 32) * 34; + else if constexpr (WTYPE == 2) return (size_t)k * 2; + else { + using traits = ggml_cuda_type_traits<(ggml_type)(WTYPE - 32)>; + return (size_t)(k / traits::qk) * traits::bs; + } +} + +template +__device__ __forceinline__ void mmb_load_quant_tile(const uint8_t * weights, size_t row_bytes, int rows, int ks, uint16_t * tile) { + for (int row = threadIdx.x / 8; row < BM; row += blockDim.x / 8) { + uint16_t * dst = tile + row * STRIDE; + const int lane = threadIdx.x % 8; + if (row < rows) mmb_decode_slice<(ggml_type)(WTYPE - 32)>(weights + row * row_bytes, ks * 64, dst, lane); + else { +#pragma unroll + for (int k = lane; k < 64; k += 8) dst[k] = 0; + } + } +} + +static bool mmb_quant_type(ggml_type type) { + switch (type) { + case GGML_TYPE_Q1_0: + case GGML_TYPE_Q2_0: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_IQ1_S: + case GGML_TYPE_IQ1_M: + 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: + case GGML_TYPE_IQ4_XS: + case GGML_TYPE_IQ4_NL: + case GGML_TYPE_MXFP4: + case GGML_TYPE_NVFP4: + return true; + default: return false; + } +} + +template +static void mmb_dispatch_quant(ggml_type type, Fn fn) { + switch (type) { + case GGML_TYPE_Q1_0: fn(std::integral_constant{}); break; + case GGML_TYPE_Q2_0: fn(std::integral_constant{}); break; + case GGML_TYPE_Q4_0: fn(std::integral_constant{}); break; + case GGML_TYPE_Q4_1: fn(std::integral_constant{}); break; + case GGML_TYPE_Q5_0: fn(std::integral_constant{}); break; + case GGML_TYPE_Q5_1: fn(std::integral_constant{}); break; + case GGML_TYPE_Q8_0: fn(std::integral_constant{}); break; + case GGML_TYPE_Q2_K: fn(std::integral_constant{}); break; + case GGML_TYPE_Q3_K: fn(std::integral_constant{}); break; + case GGML_TYPE_Q4_K: fn(std::integral_constant{}); break; + case GGML_TYPE_Q5_K: fn(std::integral_constant{}); break; + case GGML_TYPE_Q6_K: fn(std::integral_constant{}); break; + case GGML_TYPE_IQ1_S: fn(std::integral_constant{}); break; + case GGML_TYPE_IQ1_M: fn(std::integral_constant{}); break; + case GGML_TYPE_IQ2_XXS: fn(std::integral_constant{}); break; + case GGML_TYPE_IQ2_XS: fn(std::integral_constant{}); break; + case GGML_TYPE_IQ2_S: fn(std::integral_constant{}); break; + case GGML_TYPE_IQ3_XXS: fn(std::integral_constant{}); break; + case GGML_TYPE_IQ3_S: fn(std::integral_constant{}); break; + case GGML_TYPE_IQ4_XS: fn(std::integral_constant{}); break; + case GGML_TYPE_IQ4_NL: fn(std::integral_constant{}); break; + case GGML_TYPE_MXFP4: fn(std::integral_constant{}); break; + case GGML_TYPE_NVFP4: fn(std::integral_constant{}); break; + default: GGML_ABORT("unsupported MMB quant type"); + } +} + +__device__ __forceinline__ void mmb_dq_q4k_slice(uint4 q0, uint4 q1, uint4 meta, int slice, uint32_t * out) { + const float d = mmb_h2f((uint16_t) meta.x), dm = mmb_h2f((uint16_t)(meta.x >> 16)); + auto byte = [&](int i) { const uint32_t word = i < 4 ? meta.y : i < 8 ? meta.z : meta.w; return (word >> (8 * (i & 3))) & 255; }; + auto scale_min = [&](int i, float & ds, float & ms) { + const uint32_t sc = i < 4 ? byte(i) & 63 : (byte(i + 4) & 15) | ((byte(i - 4) >> 6) << 4); + const uint32_t mn = i < 4 ? byte(i + 4) & 63 : (byte(i + 4) >> 4) | ((byte(i) >> 6) << 4); + ds = d * sc; ms = dm * mn; + }; + float d0, d1, m0, m1; + scale_min(2 * slice, d0, m0); scale_min(2 * slice + 1, d1, m1); + const uint32_t words[8] = {q0.x, q0.y, q0.z, q0.w, q1.x, q1.y, q1.z, q1.w}; +#pragma unroll + for (int j = 0; j < 8; ++j) { + const uint32_t q = words[j]; + out[2*j] = mmb_pack2(d0 * (q & 15) - m0, d0 * ((q >> 8) & 15) - m0); + out[2*j + 1] = mmb_pack2(d0 * ((q >> 16) & 15) - m0, d0 * ((q >> 24) & 15) - m0); + out[16 + 2*j] = mmb_pack2(d1 * ((q >> 4) & 15) - m1, d1 * ((q >> 12) & 15) - m1); + out[16 + 2*j + 1] = mmb_pack2(d1 * ((q >> 20) & 15) - m1, d1 * (q >> 28) - m1); + } +} + + +__device__ __forceinline__ void mmb_dq_q51_pair(uint4 a, uint4 b, uint4 c, uint32_t * out) { + const uint32_t words[12] = {a.x,a.y,a.z,a.w,b.x,b.y,b.z,b.w,c.x,c.y,c.z,c.w}; +#pragma unroll + for (int block = 0; block < 2; ++block) { + const uint32_t dm = words[block * 6], high = words[block * 6 + 1]; + const float d = mmb_h2f((uint16_t)dm), m = mmb_h2f((uint16_t)(dm >> 16)); +#pragma unroll + for (int j = 0; j < 4; ++j) { + const uint32_t q = words[block * 6 + 2 + j]; + float lo[4], hi[4]; +#pragma unroll + for (int lane = 0; lane < 4; ++lane) { + const int low = ((q >> (8 * lane)) & 15) | (((high >> (4 * j + lane)) & 1) << 4); + const int upper = ((q >> (8 * lane + 4)) & 15) | (((high >> (16 + 4 * j + lane)) & 1) << 4); + lo[lane] = d * low + m; hi[lane] = d * upper + m; + } + out[block * 16 + 2*j] = mmb_pack2(lo[0],lo[1]); + out[block * 16 + 2*j + 1] = mmb_pack2(lo[2],lo[3]); + out[block * 16 + 8 + 2*j] = mmb_pack2(hi[0],hi[1]); + out[block * 16 + 8 + 2*j + 1] = mmb_pack2(hi[2],hi[3]); + } + } +} diff --git a/ggml/src/ggml-cuda/mmb.cu b/ggml/src/ggml-cuda/mmb.cu new file mode 100644 index 000000000000..25a6b04b606b --- /dev/null +++ b/ggml/src/ggml-cuda/mmb.cu @@ -0,0 +1,1004 @@ +#include "mmb.cuh" +#include "unary.cuh" +#include +#include +#include +#include "mmid.cuh" +#include +#include +#include + +namespace { + +typedef short v16s __attribute__((ext_vector_type(16))); +typedef float v8f __attribute__((ext_vector_type(8))); +constexpr int MMB_BK = 64, MMB_NT = 256, MMB_LDS_STRIDE = MMB_BK + 8; + +__device__ __forceinline__ uint16_t mmb_f2bf(float f) { uint32_t u = __float_as_uint(f); u += 0x7fffu + ((u >> 16) & 1u); return (uint16_t)(u >> 16); } +__device__ __forceinline__ uint32_t mmb_pack2(float a, float b) { return (uint32_t)mmb_f2bf(a) | ((uint32_t)mmb_f2bf(b) << 16); } +__constant__ int8_t mmb_kv_iq4nl[16] = {-127, -104, -83, -65, -49, -35, -22, -10, 1, 13, 25, 38, 53, 69, 89, 113}; +__device__ __forceinline__ float mmb_h2f(uint16_t h) { return (float) __builtin_bit_cast(_Float16, h); } + +__global__ void mmb_cvt_f32_bf16(const float * __restrict__ x, uint16_t * __restrict__ y, const size_t n) { + size_t i = ((size_t)blockIdx.x * blockDim.x + threadIdx.x) * 8; + if (i + 8 <= n) { + const float4 a = *(const float4 *)(x + i), b = *(const float4 *)(x + i + 4); + uint4 o; o.x = mmb_pack2(a.x, a.y); o.y = mmb_pack2(a.z, a.w); o.z = mmb_pack2(b.x, b.y); o.w = mmb_pack2(b.z, b.w); + *(uint4 *)(y + i) = o; + } else { + for (; i < n; ++i) y[i] = mmb_f2bf(x[i]); + } +} + +} +#include "mmb-quant.cuh" +namespace { + +// dequantize one weight row's two consecutive IQ4_NL blocks (36 bytes) into 64 bf16 in LDS. +// LUT held in registers as (kv + 128) bytes and applied with v_perm_b32 (4 nibbles per op pair) instead of a per-lane +// indexed constant array (which lowers to one scalar-byte memory load per element). The value kv*d is produced as +// fma(kv+128, d, -128*d): -128*d is exact, so the single rounding equals RN(kv*d) -> bitwise the same BF16 as before. +__device__ __forceinline__ void mmb_dq_row36(const uint4 w0, const uint4 w1, const uint32_t w2, uint32_t * arow) { + const uint32_t ws[9] = {w0.x, w0.y, w0.z, w0.w, w1.x, w1.y, w1.z, w1.w, w2}; + const float d0 = mmb_h2f((uint16_t)(ws[0] & 0xffff)), d1 = mmb_h2f((uint16_t)(ws[4] >> 16)); + const uint32_t q0[4] = { (ws[0] >> 16) | (ws[1] << 16), (ws[1] >> 16) | (ws[2] << 16), (ws[2] >> 16) | (ws[3] << 16), (ws[3] >> 16) | (ws[4] << 16) }; + const uint32_t q1[4] = { ws[5], ws[6], ws[7], ws[8] }; + // kv + 128 = {1,24,45,63,79,93,106,118,129,141,153,166,181,197,217,241} packed little-endian, 4 per dword + const uint32_t L0 = 0x3f2d1801u, L1 = 0x766a5d4fu, L2 = 0xa6998d81u, L3 = 0xf1d9c5b5u; +#pragma unroll + for (int blk = 0; blk < 2; ++blk) { + const float d = blk ? d1 : d0; const float md = -128.0f * d; const uint32_t * q = blk ? q1 : q0; uint32_t * out = arow + blk * 16; +#pragma unroll + for (int w = 0; w < 4; ++w) { + const uint32_t v = q[w]; + const uint32_t nib[2] = { v & 0x0F0F0F0Fu, (v >> 4) & 0x0F0F0F0Fu }; + float x[2][4]; +#pragma unroll + for (int h = 0; h < 2; ++h) { + const uint32_t n = nib[h]; + const uint32_t sel = n & 0x07070707u; + const uint32_t pA = __builtin_amdgcn_perm(L1, L0, sel); // entries 0..7 + const uint32_t pB = __builtin_amdgcn_perm(L3, L2, sel); // entries 8..15 + const uint32_t m = ((n >> 3) & 0x01010101u) * 0xFFu; // 0xFF where the nibble >= 8 + const uint32_t u = (pA & ~m) | (pB & m); + x[h][0] = fmaf((float)(u & 0xFFu), d, md); + x[h][1] = fmaf((float)((u >> 8) & 0xFFu), d, md); + x[h][2] = fmaf((float)((u >> 16) & 0xFFu), d, md); + x[h][3] = fmaf((float)(u >> 24), d, md); + } + out[2*w] = mmb_pack2(x[0][0], x[0][1]); out[2*w + 1] = mmb_pack2(x[0][2], x[0][3]); + out[8 + 2*w] = mmb_pack2(x[1][0], x[1][1]); out[8 + 2*w + 1] = mmb_pack2(x[1][2], x[1][3]); + } + } +} + +// dequantize one weight row's two consecutive Q8_0 blocks (68 bytes: d0 qs0[32] d1 qs1[32]) into 64 bf16 in LDS +__device__ __forceinline__ void mmb_dq_row68(const uint4 w0, const uint4 w1, const uint4 w2, const uint4 w3, const uint32_t w4, uint32_t * arow) { + const uint32_t ws[17] = {w0.x,w0.y,w0.z,w0.w, w1.x,w1.y,w1.z,w1.w, w2.x,w2.y,w2.z,w2.w, w3.x,w3.y,w3.z,w3.w, w4}; + const float d0 = mmb_h2f((uint16_t)(ws[0] & 0xffff)), d1 = mmb_h2f((uint16_t)(ws[8] >> 16)); +#pragma unroll + for (int blk = 0; blk < 2; ++blk) { + const float d = blk ? d1 : d0; uint32_t * out = arow + blk * 16; +#pragma unroll + for (int w = 0; w < 8; ++w) { + const uint32_t v = blk ? ws[9 + w] : ((ws[w] >> 16) | (ws[w + 1] << 16)); + const float e0 = d * (float)(int8_t)(v ), e1 = d * (float)(int8_t)(v >> 8); + const float e2 = d * (float)(int8_t)(v >> 16), e3 = d * (float)(int8_t)(v >> 24); + out[2*w] = mmb_pack2(e0, e1); out[2*w + 1] = mmb_pack2(e2, e3); + } + } +} + +template +__device__ __forceinline__ void mmb_store_tile(const v8f & acc, float * __restrict__ stg, float * __restrict__ D, uint16_t * __restrict__ Dh, + const bool store_f32, const int M, DRowFn drow, const int n_base, const int m_base, const int lane) { + const int cm = lane & 15, cn = lane >> 4; +#pragma unroll + for (int e = 0; e < 8; ++e) { stg[(2 * e + cn) * 16 + cm] = acc[e]; } + __syncthreads(); + const int n = lane >> 1, half = lane & 1; + const int dr = drow(n_base + n); + if (dr >= 0) { + const float4 v0 = *(const float4 *)(stg + n * 16 + half * 8), v1 = *(const float4 *)(stg + n * 16 + half * 8 + 4); + const size_t base = (size_t)dr * M + m_base + half * 8; + const bool full = m_base + 16 <= M; + if (store_f32) { + if (full && (M & 3) == 0) { *(float4 *)(D + base) = v0; *(float4 *)(D + base + 4) = v1; } + else { const float vv[8] = {v0.x, v0.y, v0.z, v0.w, v1.x, v1.y, v1.z, v1.w}; +#pragma unroll + for (int k = 0; k < 8; ++k) { if (m_base + half * 8 + k < M) D[base + k] = vv[k]; } } + } + if (Dh) { + if (full && (M & 7) == 0) { *(uint4 *)(Dh + base) = make_uint4(mmb_pack2(v0.x, v0.y), mmb_pack2(v0.z, v0.w), mmb_pack2(v1.x, v1.y), mmb_pack2(v1.z, v1.w)); } + else { const float vv[8] = {v0.x, v0.y, v0.z, v0.w, v1.x, v1.y, v1.z, v1.w}; +#pragma unroll + for (int k = 0; k < 8; ++k) { if (m_base + half * 8 + k < M) Dh[base + k] = mmb_f2bf(vv[k]); } } + } + } + __syncthreads(); +} +template +__device__ __forceinline__ void mmb_tile_gemm(const uint8_t * __restrict__ Wbase, const size_t wrow_bytes, const int a_rows, + const uint16_t * __restrict__ Xh, const int K, XRowFn xrow, float * __restrict__ D, uint16_t * __restrict__ Dh, const bool store_f32, const int M, DRowFn drow, const int m0, + const int n_cols, uint16_t * As, uint16_t * Bs) { + constexpr int WAVES_M = BM / WTM, TM = WTM / 16, TN = WTN / 16; + constexpr int A_ITEMS = (BM + MMB_NT - 1) / MMB_NT; + constexpr int B_ITEMS = (BN * 8) / MMB_NT; + const int tid = threadIdx.x, lane = tid & 31, wave = tid >> 5; + const int wm = wave % WAVES_M, wn = wave / WAVES_M; + uint4 a0[A_ITEMS], a1[A_ITEMS], a3[A_ITEMS], a4[A_ITEMS], a5[A_ITEMS], a6[A_ITEMS], a7[A_ITEMS], a8[A_ITEMS]; uint32_t a2[A_ITEMS]; + uint4 bst[B_ITEMS]; + int brow[B_ITEMS]; +#pragma unroll + for (int i = 0; i < B_ITEMS; ++i) { const int c = tid + i * MMB_NT; brow[i] = xrow(c >> 3); } + + int weight_ks = 0; + auto load_regs = [&](const int ks) { + weight_ks = ks; +#pragma unroll + for (int i = 0; i < A_ITEMS; ++i) { + const int row = tid + i * MMB_NT; + if (row < BM && row < a_rows) { + if constexpr (WTYPE == 0) { const uint8_t * p = Wbase + (size_t)row * wrow_bytes + (size_t)ks * 36; + a0[i] = *(const uint4 *)(p); a1[i] = *(const uint4 *)(p + 16); a2[i] = *(const uint32_t *)(p + 32); } + else if constexpr (WTYPE == 1) { const uint8_t * p = Wbase + (size_t)row * wrow_bytes + (size_t)ks * 68; + a0[i] = *(const uint4 *)(p); a1[i] = *(const uint4 *)(p + 16); a3[i] = *(const uint4 *)(p + 32); a4[i] = *(const uint4 *)(p + 48); a2[i] = *(const uint32_t *)(p + 64); } + else if constexpr (WTYPE == 32 + GGML_TYPE_Q5_1) { + const uint4 * p = (const uint4 *)(Wbase + (size_t)row * wrow_bytes + (size_t)ks * 48); + a0[i] = p[0]; a1[i] = p[1]; a3[i] = p[2]; + } + else if constexpr (WTYPE == 2) { const uint4 * p = (const uint4 *)(Wbase + (size_t)row * wrow_bytes + (size_t)ks * 128); + a0[i] = p[0]; a1[i] = p[1]; a3[i] = p[2]; a4[i] = p[3]; a5[i] = p[4]; a6[i] = p[5]; a7[i] = p[6]; a8[i] = p[7]; } + } else { a0[i] = make_uint4(0,0,0,0); a1[i] = make_uint4(0,0,0,0); a3[i] = make_uint4(0,0,0,0); a4[i] = make_uint4(0,0,0,0); a5[i] = a6[i] = a7[i] = a8[i] = make_uint4(0,0,0,0); a2[i] = 0; } + } +#pragma unroll + for (int i = 0; i < B_ITEMS; ++i) { + const int c = tid + i * MMB_NT; const int off = (c & 7) * 8; + bst[i] = (brow[i] >= 0) ? *(const uint4 *)(Xh + (size_t)brow[i] * K + ks * MMB_BK + off) : make_uint4(0,0,0,0); + } + }; + auto store_lds = [&]() { + if constexpr (WTYPE >= 32 && WTYPE != 32 + GGML_TYPE_Q5_1) mmb_load_quant_tile(Wbase, wrow_bytes, a_rows, weight_ks, As); +#pragma unroll + for (int i = 0; i < A_ITEMS; ++i) { const int row = tid + i * MMB_NT; if (row < BM) { + if constexpr (WTYPE == 0) mmb_dq_row36(a0[i], a1[i], a2[i], (uint32_t *)(As + row * MMB_LDS_STRIDE)); + else if constexpr (WTYPE == 1) mmb_dq_row68(a0[i], a1[i], a3[i], a4[i], a2[i], (uint32_t *)(As + row * MMB_LDS_STRIDE)); + else if constexpr (WTYPE == 32 + GGML_TYPE_Q5_1) mmb_dq_q51_pair(a0[i], a1[i], a3[i], (uint32_t *)(As + row * MMB_LDS_STRIDE)); + else if constexpr (WTYPE == 2) { uint4 * d = (uint4 *)(As + row * MMB_LDS_STRIDE); d[0] = a0[i]; d[1] = a1[i]; d[2] = a3[i]; d[3] = a4[i]; d[4] = a5[i]; d[5] = a6[i]; d[6] = a7[i]; d[7] = a8[i]; } } } +#pragma unroll + for (int i = 0; i < B_ITEMS; ++i) { const int c = tid + i * MMB_NT; *(uint4 *)(Bs + (c >> 3) * MMB_LDS_STRIDE + (c & 7) * 8) = bst[i]; } + }; + + v8f acc[TM][TN]; +#pragma unroll + for (int i = 0; i < TM; ++i) +#pragma unroll + for (int j = 0; j < TN; ++j) +#pragma unroll + for (int e = 0; e < 8; ++e) acc[i][j][e] = 0.f; + + bool jact[TN]; +#pragma unroll + for (int j = 0; j < TN; ++j) jact[j] = !TAIL || (wn * WTN + j * 16) < n_cols; // whole fragments past the valid columns are never stored + const int nks = K / MMB_BK; + load_regs(0); store_lds(); __syncthreads(); + for (int ks = 0; ks < nks; ++ks) { + if (ks + 1 < nks) load_regs(ks + 1); +#pragma unroll + for (int kk = 0; kk < MMB_BK; kk += 16) { + v16s a[TM], b[TN]; const int r = lane & 15; +#pragma unroll + for (int i = 0; i < TM; ++i) { const uint16_t * p = As + (wm * WTM + i * 16 + r) * MMB_LDS_STRIDE + kk; + const uint4 v0 = *(const uint4 *)p, v1 = *(const uint4 *)(p + 8); a[i] = __builtin_bit_cast(v16s, (uint4[2]){v0, v1}); } +#pragma unroll + for (int j = 0; j < TN; ++j) { if constexpr (TAIL) { if (!jact[j]) continue; } const uint16_t * p = Bs + (wn * WTN + j * 16 + r) * MMB_LDS_STRIDE + kk; + const uint4 v0 = *(const uint4 *)p, v1 = *(const uint4 *)(p + 8); b[j] = __builtin_bit_cast(v16s, (uint4[2]){v0, v1}); } +#pragma unroll + for (int i = 0; i < TM; ++i) +#pragma unroll + for (int j = 0; j < TN; ++j) { if constexpr (TAIL) { if (!jact[j]) continue; } acc[i][j] = __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32(b[j], a[i], acc[i][j]); } + } + __syncthreads(); + if (ks + 1 < nks) store_lds(); + __syncthreads(); + } + // epilogue through LDS (per-wave 1 KB stage in the now-free A/B tile area); tiles are 16 rows, a_rows is a + // multiple of 32 in this model, so whole tiles are either valid or beyond a_rows + float * stg = (float *)((BN * MMB_LDS_STRIDE * 2 >= MMB_NT / 32 * 1024) ? Bs : As) + wave * 256; +#pragma unroll + for (int i = 0; i < TM; ++i) { + const int ml = wm * WTM + i * 16; const bool ok = ml < a_rows; +#pragma unroll + for (int j = 0; j < TN; ++j) { + if (ok && jact[j]) mmb_store_tile(acc[i][j], stg, D, Dh, store_f32, M, drow, wn * WTN + j * 16, m0 + ml, lane); + else { __syncthreads(); __syncthreads(); } + } + } +} + +template +__global__ void __launch_bounds__(MMB_NT, 2) +mmb_dense_kernel(const uint8_t * __restrict__ W, const uint16_t * __restrict__ Xh, float * __restrict__ D, uint16_t * __restrict__ Dh, const bool store_f32, const int M, const int K, const int T) { +#if defined(__HIP_DEVICE_COMPILE__) && !defined(RDNA3) + NO_DEVICE_CODE; // WMMA kernels are RDNA3-only; the host gate keeps other devices off this path +#else + __shared__ __align__(16) uint16_t As[BM * MMB_LDS_STRIDE]; + __shared__ __align__(16) uint16_t Bs[BN * MMB_LDS_STRIDE]; + const int m0 = blockIdx.x * BM, t0 = blockIdx.y * BN; + const size_t wrow_bytes = mmb_row_bytes(K); + mmb_tile_gemm(W + (size_t)m0 * wrow_bytes, wrow_bytes, M - m0, Xh, K, + [&](int i) { return (t0 + i < T) ? t0 + i : -1; }, D, Dh, store_f32, M, [&](int i) { return (t0 + i < T) ? t0 + i : -1; }, m0, T - t0, As, Bs); +#endif +} + +#if defined(__HIP_PLATFORM_AMD__) +__device__ __forceinline__ float gm_mul_rn(const float a, const float b) { float r; asm("v_mul_f32_e32 %0, %1, %2" : "=v"(r) : "v"(a), "v"(b)); return r; } +__device__ __forceinline__ float gm_add_rn(const float a, const float b) { float r; asm("v_add_f32_e32 %0, %1, %2" : "=v"(r) : "v"(a), "v"(b)); return r; } +#else +__device__ __forceinline__ float gm_mul_rn(const float a, const float b) { return __fmul_rn(a, b); } +__device__ __forceinline__ float gm_add_rn(const float a, const float b) { return __fadd_rn(a, b); } +#endif +__device__ __forceinline__ float gm_sigmoid(const float x) { return 1.0f / (1.0f + expf(-x)); } +__device__ __forceinline__ float gm_bf2f(const uint16_t h) { return __uint_as_float(((uint32_t) h) << 16); } + +template +__global__ void __launch_bounds__(MMB_NT, 2) +hc_gate_mix_kernel(const uint8_t * __restrict__ W, const uint16_t * __restrict__ Lo, const uint16_t * __restrict__ Xn, float * __restrict__ Out, + uint16_t * __restrict__ OutH, const bool store_f32, + const int E, const int K, const int T, const float scale, const float bias) { +#if defined(__HIP_DEVICE_COMPILE__) && !defined(RDNA3) + // WMMA (wave32, 16x16x16 bf16/f16 into f32) exists only on RDNA3; the host gate keeps other devices off this path + NO_DEVICE_CODE; +#else + constexpr int CH = 32, BN = 128, BM = HC * CH; + static_assert(BM <= MMB_NT, "one A row per thread"); + __shared__ __align__(16) uint16_t As[BM * MMB_LDS_STRIDE]; + __shared__ __align__(16) uint16_t Bs[BN * MMB_LDS_STRIDE]; + const int tid = threadIdx.x, lane = tid & 31, wave = tid >> 5; + const int wm = wave & 1, wn = wave >> 1; // wave: 16 channels (all HC streams) x 32 tokens + const int e0 = blockIdx.x * CH, t0 = blockIdx.y * BN; + const size_t wrow_bytes = mmb_row_bytes(K); + constexpr int B_ITEMS = (BN * 8) / MMB_NT; + uint4 a0 = make_uint4(0,0,0,0), a1 = make_uint4(0,0,0,0); uint32_t a2 = 0; uint4 bst[B_ITEMS]; int brow[B_ITEMS]; + const uint8_t * arow = W; + if (tid < BM) { const int c = tid / CH, i = tid - c * CH; arow = W + (size_t)(c * E + e0 + i) * wrow_bytes; } +#pragma unroll + for (int i = 0; i < B_ITEMS; ++i) { const int c = tid + i * MMB_NT; const int t = t0 + (c >> 3); brow[i] = t < T ? t : -1; } + int weight_ks = 0; + auto load_regs = [&](const int ks) { + weight_ks = ks; + if constexpr (WTYPE == 0) if (tid < BM) { const uint8_t * p = arow + (size_t)ks * 36; a0 = *(const uint4 *)(p); a1 = *(const uint4 *)(p + 16); a2 = *(const uint32_t *)(p + 32); } +#pragma unroll + for (int i = 0; i < B_ITEMS; ++i) { const int c = tid + i * MMB_NT; const int off = (c & 7) * 8; + bst[i] = (brow[i] >= 0) ? *(const uint4 *)(Lo + (size_t)brow[i] * K + ks * MMB_BK + off) : make_uint4(0,0,0,0); } + }; + auto store_lds = [&]() { + if constexpr (WTYPE == 0) { if (tid < BM) mmb_dq_row36(a0, a1, a2, (uint32_t *)(As + tid * MMB_LDS_STRIDE)); } + else { + constexpr int TYPE = WTYPE == 1 ? GGML_TYPE_Q8_0 : WTYPE - 32; + for (int row = tid / 8; row < BM; row += MMB_NT / 8) { + const int c = row / CH, ch = row % CH; + mmb_decode_slice<(ggml_type)TYPE>(W + (size_t)(c * E + e0 + ch) * wrow_bytes, weight_ks * 64, As + row * MMB_LDS_STRIDE, tid % 8); + } + } +#pragma unroll + for (int i = 0; i < B_ITEMS; ++i) { const int c = tid + i * MMB_NT; *(uint4 *)(Bs + (c >> 3) * MMB_LDS_STRIDE + (c & 7) * 8) = bst[i]; } + }; + v8f acc[HC][2]; +#pragma unroll + for (int c = 0; c < HC; ++c) +#pragma unroll + for (int j = 0; j < 2; ++j) +#pragma unroll + for (int e = 0; e < 8; ++e) acc[c][j][e] = 0.f; + const int nks = K / MMB_BK; + load_regs(0); store_lds(); __syncthreads(); + for (int ks = 0; ks < nks; ++ks) { + if (ks + 1 < nks) load_regs(ks + 1); +#pragma unroll + for (int kk = 0; kk < MMB_BK; kk += 16) { + v16s a[HC], b[2]; const int r = lane & 15; +#pragma unroll + for (int c = 0; c < HC; ++c) { const uint16_t * p = As + (c * CH + wm * 16 + r) * MMB_LDS_STRIDE + kk; + const uint4 v0 = *(const uint4 *)p, v1 = *(const uint4 *)(p + 8); a[c] = __builtin_bit_cast(v16s, (uint4[2]){v0, v1}); } +#pragma unroll + for (int j = 0; j < 2; ++j) { const uint16_t * p = Bs + (wn * 32 + j * 16 + r) * MMB_LDS_STRIDE + kk; + const uint4 v0 = *(const uint4 *)p, v1 = *(const uint4 *)(p + 8); b[j] = __builtin_bit_cast(v16s, (uint4[2]){v0, v1}); } +#pragma unroll + for (int c = 0; c < HC; ++c) +#pragma unroll + for (int j = 0; j < 2; ++j) acc[c][j] = __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32(b[j], a[c], acc[c][j]); + } + __syncthreads(); + if (ks + 1 < nks) store_lds(); + __syncthreads(); + } + // epilogue: lane holds channel (lane & 15) of the wave's 16 and tokens 2e + (lane >> 4) of each 16-token fragment + const int cm = lane & 15, cn = lane >> 4; const int ch = e0 + wm * 16 + cm; +#pragma unroll + for (int j = 0; j < 2; ++j) { +#pragma unroll + for (int e = 0; e < 8; ++e) { + const int t = t0 + wn * 32 + j * 16 + 2 * e + cn; + if (t >= T) continue; + const uint16_t * xr = Xn + (size_t)t * ((size_t)HC * E) + ch; + float s = 0.f; +#pragma unroll + for (int c = 0; c < HC; ++c) { + const float g = __uint_as_float(((uint32_t) mmb_f2bf(acc[c][j][e])) << 16); // the gate GEMM's BF16 epilogue rounding + const float term = gm_mul_rn(gm_bf2f(xr[(size_t)c * E]), gm_sigmoid(g)); + s = (c == 0) ? term : gm_add_rn(s, term); + } + const float o = scale * s + bias; + if (store_f32) Out[(size_t)t * E + ch] = o; + if (OutH) OutH[(size_t)t * E + ch] = mmb_f2bf(o); + } + } +#endif +} + +template +__global__ void __launch_bounds__(MMB_NT, 2) +mmb_routed_kernel(const uint8_t * __restrict__ W, const size_t expert_bytes, const uint16_t * __restrict__ Xh, float * __restrict__ D, + uint16_t * __restrict__ Dh, const bool store_f32, + const int32_t * __restrict__ ids_src, const int32_t * __restrict__ ids_dst, const int32_t * __restrict__ bounds, + const uint32_t * __restrict__ desc, const int M, const int K) { +#if defined(__HIP_DEVICE_COMPILE__) && !defined(RDNA3) + NO_DEVICE_CODE; // WMMA kernels are RDNA3-only; the host gate keeps other devices off this path +#else + __shared__ __align__(16) uint16_t As[BM * MMB_LDS_STRIDE]; + __shared__ __align__(16) uint16_t Bs[BN * MMB_LDS_STRIDE]; + const uint32_t dsc = desc[blockIdx.y]; + if (dsc == UINT32_MAX) return; // uniform across the block, before any barrier + const int e = dsc & 0xffff, jt = dsc >> 16; + const int r0 = bounds[e] + jt * BN, cnt = bounds[e + 1] - r0; + const int m0 = blockIdx.x * BM; + const size_t wrow_bytes = mmb_row_bytes(K); + mmb_tile_gemm(W + (size_t)e * expert_bytes + (size_t)m0 * wrow_bytes, wrow_bytes, M - m0, Xh, K, + [&](int i) { return (i < cnt) ? ids_src[r0 + i] : -1; }, D, Dh, store_f32, M, [&](int i) { return (i < cnt) ? ids_dst[r0 + i] : -1; }, m0, cnt, As, Bs); +#endif +} + +template +__device__ __forceinline__ void mmb_tile_gemm_glu(const uint8_t * __restrict__ Wg, const uint8_t * __restrict__ Wu, const size_t wrow_bytes, const int a_rows, + const uint16_t * __restrict__ Xh, const int K, XRowFn xrow, float * __restrict__ D, uint16_t * __restrict__ Dh, const bool store_f32, const int M, DRowFn drow, const int m0, + const int n_cols, uint16_t * Ag, uint16_t * Au, uint16_t * Bs) { + constexpr int WAVES_M = BM / WTM, TM = WTM / 16, TN = WTN / 16; + constexpr int A_ITEMS = (BM + MMB_NT - 1) / MMB_NT; + constexpr int B_ITEMS = (BN * 8) / MMB_NT; + const int tid = threadIdx.x, lane = tid & 31, wave = tid >> 5; + const int wm = wave % WAVES_M, wn = wave / WAVES_M; + uint4 g0[A_ITEMS], g1[A_ITEMS], u0[A_ITEMS], u1[A_ITEMS]; uint32_t g2[A_ITEMS], u2[A_ITEMS]; + uint4 gm[A_ITEMS], um[A_ITEMS]; + uint4 bst[B_ITEMS]; + int brow[B_ITEMS]; +#pragma unroll + for (int i = 0; i < B_ITEMS; ++i) { const int c = tid + i * MMB_NT; brow[i] = xrow(c >> 3); } + int weight_ks = 0; + auto load_regs = [&](const int ks) { + weight_ks = ks; +#pragma unroll + for (int i = 0; i < A_ITEMS; ++i) { + const int row = tid + i * MMB_NT; + if constexpr (WTYPE == 0) { + if (row < BM && row < a_rows) { + const uint8_t * pg = Wg + (size_t)row * wrow_bytes + (size_t)ks * 36; + const uint8_t * pu = Wu + (size_t)row * wrow_bytes + (size_t)ks * 36; + g0[i] = *(const uint4 *)(pg); g1[i] = *(const uint4 *)(pg + 16); g2[i] = *(const uint32_t *)(pg + 32); + u0[i] = *(const uint4 *)(pu); u1[i] = *(const uint4 *)(pu + 16); u2[i] = *(const uint32_t *)(pu + 32); + } else { g0[i] = g1[i] = u0[i] = u1[i] = make_uint4(0,0,0,0); g2[i] = u2[i] = 0; } + } else if constexpr (WTYPE == 32 + GGML_TYPE_Q4_K) { + if (row < BM && row < a_rows) { + const uint8_t * pg = Wg + (size_t)row * wrow_bytes + (size_t)(ks / 4) * sizeof(block_q4_K); + const uint8_t * pu = Wu + (size_t)row * wrow_bytes + (size_t)(ks / 4) * sizeof(block_q4_K); + gm[i] = *(const uint4 *)pg; um[i] = *(const uint4 *)pu; + const int offset = 16 + (ks & 3) * 32; + g0[i] = *(const uint4 *)(pg + offset); g1[i] = *(const uint4 *)(pg + offset + 16); + u0[i] = *(const uint4 *)(pu + offset); u1[i] = *(const uint4 *)(pu + offset + 16); + } else { gm[i] = um[i] = g0[i] = g1[i] = u0[i] = u1[i] = make_uint4(0,0,0,0); } + } + } +#pragma unroll + for (int i = 0; i < B_ITEMS; ++i) { + const int c = tid + i * MMB_NT; const int off = (c & 7) * 8; + bst[i] = (brow[i] >= 0) ? *(const uint4 *)(Xh + (size_t)brow[i] * K + ks * MMB_BK + off) : make_uint4(0,0,0,0); + } + }; + auto store_lds = [&]() { + if constexpr (WTYPE != 0 && WTYPE != 32 + GGML_TYPE_Q4_K) { + constexpr int LOAD_TYPE = WTYPE == 1 ? 32 + GGML_TYPE_Q8_0 : WTYPE; + mmb_load_quant_tile(Wg, wrow_bytes, a_rows, weight_ks, Ag); + mmb_load_quant_tile(Wu, wrow_bytes, a_rows, weight_ks, Au); + } + +#pragma unroll + for (int i = 0; i < A_ITEMS; ++i) { + const int row = tid + i * MMB_NT; + if (row < BM) { + if constexpr (WTYPE == 0) { + mmb_dq_row36(g0[i], g1[i], g2[i], (uint32_t *)(Ag + row * MMB_LDS_STRIDE)); + mmb_dq_row36(u0[i], u1[i], u2[i], (uint32_t *)(Au + row * MMB_LDS_STRIDE)); + } else if constexpr (WTYPE == 32 + GGML_TYPE_Q4_K) { + mmb_dq_q4k_slice(g0[i], g1[i], gm[i], weight_ks & 3, (uint32_t *)(Ag + row * MMB_LDS_STRIDE)); + mmb_dq_q4k_slice(u0[i], u1[i], um[i], weight_ks & 3, (uint32_t *)(Au + row * MMB_LDS_STRIDE)); + } + } + } +#pragma unroll + for (int i = 0; i < B_ITEMS; ++i) { const int c = tid + i * MMB_NT; *(uint4 *)(Bs + (c >> 3) * MMB_LDS_STRIDE + (c & 7) * 8) = bst[i]; } + }; + v8f accg[TM][TN], accu[TM][TN]; +#pragma unroll + for (int i = 0; i < TM; ++i) +#pragma unroll + for (int j = 0; j < TN; ++j) +#pragma unroll + for (int e = 0; e < 8; ++e) { accg[i][j][e] = 0.f; accu[i][j][e] = 0.f; } + bool jact[TN]; +#pragma unroll + for (int j = 0; j < TN; ++j) jact[j] = !TAIL || (wn * WTN + j * 16) < n_cols; + const int nks = K / MMB_BK; + load_regs(0); store_lds(); __syncthreads(); + for (int ks = 0; ks < nks; ++ks) { + if (ks + 1 < nks) load_regs(ks + 1); +#pragma unroll + for (int kk = 0; kk < MMB_BK; kk += 16) { + v16s ag[TM], au[TM], b[TN]; const int r = lane & 15; +#pragma unroll + for (int i = 0; i < TM; ++i) { const int off = (wm * WTM + i * 16 + r) * MMB_LDS_STRIDE + kk; + ag[i] = __builtin_bit_cast(v16s, (uint4[2]){*(const uint4 *)(Ag + off), *(const uint4 *)(Ag + off + 8)}); + au[i] = __builtin_bit_cast(v16s, (uint4[2]){*(const uint4 *)(Au + off), *(const uint4 *)(Au + off + 8)}); } +#pragma unroll + for (int j = 0; j < TN; ++j) { if constexpr (TAIL) { if (!jact[j]) continue; } const uint16_t * p = Bs + (wn * WTN + j * 16 + r) * MMB_LDS_STRIDE + kk; + b[j] = __builtin_bit_cast(v16s, (uint4[2]){*(const uint4 *)p, *(const uint4 *)(p + 8)}); } +#pragma unroll + for (int i = 0; i < TM; ++i) +#pragma unroll + for (int j = 0; j < TN; ++j) { + if constexpr (TAIL) { if (!jact[j]) continue; } + accg[i][j] = __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32(b[j], ag[i], accg[i][j]); + accu[i][j] = __builtin_amdgcn_wmma_f32_16x16x16_bf16_w32(b[j], au[i], accu[i][j]); + } + } + __syncthreads(); + if (ks + 1 < nks) store_lds(); + __syncthreads(); + } + float * stg = (float *)((BN * MMB_LDS_STRIDE * 2 >= MMB_NT / 32 * 1024) ? Bs : Ag) + wave * 256; +#pragma unroll + for (int i = 0; i < TM; ++i) { + const int ml = wm * WTM + i * 16; const bool ok = ml < a_rows; +#pragma unroll + for (int j = 0; j < TN; ++j) { + if (ok && jact[j]) { + v8f v; +#pragma unroll + for (int e = 0; e < 8; ++e) { v[e] = ggml_cuda_op_silu_single(accg[i][j][e]) * accu[i][j][e]; } + mmb_store_tile(v, stg, D, Dh, store_f32, M, drow, wn * WTN + j * 16, m0 + ml, lane); + } else { __syncthreads(); __syncthreads(); } + } + } +} + +template +__global__ void __launch_bounds__(MMB_NT, 2) +mmb_routed_glu_kernel(const uint8_t * __restrict__ Wg, const uint8_t * __restrict__ Wu, const size_t expert_bytes, const uint16_t * __restrict__ Xh, + float * __restrict__ D, uint16_t * __restrict__ Dh, const bool store_f32, + const int32_t * __restrict__ ids_src, const int32_t * __restrict__ ids_dst, const int32_t * __restrict__ bounds, + const uint32_t * __restrict__ desc, const int M, const int K) { +#if defined(__HIP_DEVICE_COMPILE__) && !defined(RDNA3) + NO_DEVICE_CODE; // WMMA kernels are RDNA3-only; the host gate keeps other devices off this path +#else + __shared__ __align__(16) uint16_t Ag[BM * MMB_LDS_STRIDE]; + __shared__ __align__(16) uint16_t Au[BM * MMB_LDS_STRIDE]; + __shared__ __align__(16) uint16_t Bs[BN * MMB_LDS_STRIDE]; + const uint32_t dsc = desc[blockIdx.y]; + if (dsc == UINT32_MAX) return; + const int e = dsc & 0xffff, jt = dsc >> 16; + const int r0 = bounds[e] + jt * BN, cnt = bounds[e + 1] - r0; + const int m0 = blockIdx.x * BM; + const size_t wrow_bytes = mmb_row_bytes(K); + mmb_tile_gemm_glu(Wg + (size_t)e * expert_bytes + (size_t)m0 * wrow_bytes, Wu + (size_t)e * expert_bytes + (size_t)m0 * wrow_bytes, wrow_bytes, M - m0, Xh, K, + [&](int i) { return (i < cnt) ? ids_src[r0 + i] : -1; }, D, Dh, store_f32, M, [&](int i) { return (i < cnt) ? ids_dst[r0 + i] : -1; }, m0, cnt, Ag, Au, Bs); +#endif +} + +// F32 x F32 -> F32 GEMM on WMMA: operands split into F16 hi + lo at tile load; TWO: Xhi*Whi + Xlo*Whi, !TWO adds Xhi*Wlo, !XSPLIT keeps only Xhi*Whi. +__device__ __forceinline__ void mmb_split2(float x, uint16_t & hi, uint16_t & lo) { + hi = __builtin_bit_cast(uint16_t, (_Float16) x); lo = __builtin_bit_cast(uint16_t, (_Float16) (x - (float) __builtin_bit_cast(_Float16, hi))); +} +template +__global__ void __launch_bounds__(MMB_NT, 2) +mmb_f32split_kernel(const float * __restrict__ W, const float * __restrict__ X, float * __restrict__ D, const int M, const int K, const int T) { +#if defined(__HIP_DEVICE_COMPILE__) && !defined(RDNA3) + // WMMA (wave32, 16x16x16 bf16/f16 into f32) exists only on RDNA3; the host gate keeps other devices off this path + NO_DEVICE_CODE; +#else + constexpr int BKs = 32, LS = BKs + 8, WAVES_M = BM / WTM, TM = WTM / 16, TN = WTN / 16; + static_assert(WAVES_M * (BN / WTN) * 32 == MMB_NT, "tile must use every wave"); + __shared__ __align__(16) uint16_t Ah[BM * LS], Al[TWO ? 8 : BM * LS], Bh[BN * LS], Bl[XSPLIT ? BN * LS : 8]; + const int tid = threadIdx.x, lane = tid & 31, wave = tid >> 5, wm = wave % WAVES_M, wn = wave / WAVES_M; + const int m0 = blockIdx.x * BM, t0 = blockIdx.y * BN; + v8f acc[TM][TN]; +#pragma unroll + for (int i = 0; i < TM; ++i) +#pragma unroll + for (int j = 0; j < TN; ++j) +#pragma unroll + for (int e = 0; e < 8; ++e) acc[i][j][e] = 0.f; + constexpr int A_CH = BM * BKs / 4, B_CH = BN * BKs / 4; + for (int k0 = 0; k0 < K; k0 += BKs) { + for (int idx = tid; idx < A_CH; idx += MMB_NT) { const int row = idx >> 3, c4 = (idx & 7) * 4; const int m = m0 + row; + float4 v = make_float4(0.f,0.f,0.f,0.f); if (m < M) v = *(const float4 *)(W + (size_t) m * K + k0 + c4); + uint16_t h[4], l[4]; mmb_split2(v.x,h[0],l[0]); mmb_split2(v.y,h[1],l[1]); mmb_split2(v.z,h[2],l[2]); mmb_split2(v.w,h[3],l[3]); + *(uint2 *)(Ah + row * LS + c4) = make_uint2((uint32_t)h[0] | ((uint32_t)h[1] << 16), (uint32_t)h[2] | ((uint32_t)h[3] << 16)); + if constexpr (!TWO) *(uint2 *)(Al + row * LS + c4) = make_uint2((uint32_t)l[0] | ((uint32_t)l[1] << 16), (uint32_t)l[2] | ((uint32_t)l[3] << 16)); } + for (int idx = tid; idx < B_CH; idx += MMB_NT) { const int row = idx >> 3, c4 = (idx & 7) * 4; const int t = t0 + row; + float4 v = make_float4(0.f,0.f,0.f,0.f); if (t < T) v = *(const float4 *)(X + (size_t) t * K + k0 + c4); + uint16_t h[4], l[4]; mmb_split2(v.x,h[0],l[0]); mmb_split2(v.y,h[1],l[1]); mmb_split2(v.z,h[2],l[2]); mmb_split2(v.w,h[3],l[3]); + *(uint2 *)(Bh + row * LS + c4) = make_uint2((uint32_t)h[0] | ((uint32_t)h[1] << 16), (uint32_t)h[2] | ((uint32_t)h[3] << 16)); + if constexpr (XSPLIT) *(uint2 *)(Bl + row * LS + c4) = make_uint2((uint32_t)l[0] | ((uint32_t)l[1] << 16), (uint32_t)l[2] | ((uint32_t)l[3] << 16)); } + __syncthreads(); + const int r = lane & 15; +#pragma unroll + for (int kk = 0; kk < BKs; kk += 16) { + v16s ah[TM], al[TM], bh[TN], bl[TN]; +#pragma unroll + for (int i = 0; i < TM; ++i) { const int off = (wm * WTM + i * 16 + r) * LS + kk; + ah[i] = __builtin_bit_cast(v16s, (uint4[2]){*(const uint4 *)(Ah + off), *(const uint4 *)(Ah + off + 8)}); + if constexpr (!TWO) al[i] = __builtin_bit_cast(v16s, (uint4[2]){*(const uint4 *)(Al + off), *(const uint4 *)(Al + off + 8)}); } +#pragma unroll + for (int j = 0; j < TN; ++j) { const int off = (wn * WTN + j * 16 + r) * LS + kk; + bh[j] = __builtin_bit_cast(v16s, (uint4[2]){*(const uint4 *)(Bh + off), *(const uint4 *)(Bh + off + 8)}); + if constexpr (XSPLIT) bl[j] = __builtin_bit_cast(v16s, (uint4[2]){*(const uint4 *)(Bl + off), *(const uint4 *)(Bl + off + 8)}); } +#pragma unroll + for (int i = 0; i < TM; ++i) +#pragma unroll + for (int j = 0; j < TN; ++j) { + acc[i][j] = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(bh[j], ah[i], acc[i][j]); + if constexpr (XSPLIT) acc[i][j] = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(bl[j], ah[i], acc[i][j]); + if constexpr (!TWO) acc[i][j] = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(bh[j], al[i], acc[i][j]); + } + } + __syncthreads(); + } + const int cm = lane & 15, cn = lane >> 4; +#pragma unroll + for (int i = 0; i < TM; ++i) { const int m = m0 + wm * WTM + i * 16 + cm; if (m >= M) continue; +#pragma unroll + for (int j = 0; j < TN; ++j) +#pragma unroll + for (int e = 0; e < 8; ++e) { const int t = t0 + wn * WTN + j * 16 + 2 * e + cn; if (t < T) D[(size_t) t * M + m] = acc[i][j][e]; } } +#endif +} + +// two tile classes: experts with >= thresh rows get BN_BIG-row tiles, the rest BN_SMALL-row tiles (fewer wasted rows on tiny experts) +__global__ void mmb_build_desc2(const int32_t * __restrict__ bounds, uint32_t * __restrict__ desc_big, uint32_t * __restrict__ desc_small, + const int E, const int nbig_max, const int nsmall_max, const int BN_BIG, const int BN_SMALL, const int thresh) { + __shared__ int sb[1024], ss[1024]; + const int e = threadIdx.x; + for (int i = e; i < nbig_max; i += blockDim.x) desc_big[i] = UINT32_MAX; + for (int i = e; i < nsmall_max; i += blockDim.x) desc_small[i] = UINT32_MAX; + int cnt = (e < E) ? bounds[e + 1] - bounds[e] : 0; + const bool big = cnt >= thresh; + const int tb = big ? (cnt + BN_BIG - 1) / BN_BIG : 0; + const int ts = big ? 0 : (cnt + BN_SMALL - 1) / BN_SMALL; + sb[e] = tb; ss[e] = ts; + __syncthreads(); + for (int off = 1; off < 1024; off <<= 1) { + const int vb = (e >= off) ? sb[e - off] : 0, vs = (e >= off) ? ss[e - off] : 0; + __syncthreads(); + sb[e] += vb; ss[e] += vs; + __syncthreads(); + } + const int bb = sb[e] - tb, bs = ss[e] - ts; + for (int jt = 0; jt < tb; ++jt) { const int idx = bb + jt; if (idx < nbig_max) desc_big[idx] = (uint32_t)e | ((uint32_t)jt << 16); } + for (int jt = 0; jt < ts; ++jt) { const int idx = bs + jt; if (idx < nsmall_max) desc_small[idx] = (uint32_t)e | ((uint32_t)jt << 16); } +} + +} // namespace + +struct mmb_cache_entry { const ggml_tensor * root; const void * data; size_t n; ggml_cuda_pool_alloc * buf; }; +struct ggml_cuda_mmb_context { + struct workspace { + std::vector cache; + mmb_cache_entry slots[4] = {}; + size_t slot_cap[4] = {}; + }; + workspace streams[GGML_CUDA_MAX_STREAMS]; + std::unordered_set bf16_only; + std::unordered_map shadow; + std::map, uint16_t *> shadow_pair; + size_t shadow_bytes = 0; +}; + +namespace { + +static ggml_cuda_mmb_context & mmb_state(ggml_backend_cuda_context & ctx) { + if (!ctx.mmb) { + ctx.mmb = new ggml_cuda_mmb_context; + } + return *ctx.mmb; +} + +static ggml_cuda_mmb_context::workspace & mmb_workspace(ggml_backend_cuda_context & ctx) { + return mmb_state(ctx).streams[ctx.curr_stream_no]; +} + +static size_t mmb_cache_max() { return 4; } +static const ggml_tensor * mmb_root(const ggml_tensor * t) { return t->view_src ? t->view_src : t; } +static uint16_t * mmb_cache_insert(ggml_backend_cuda_context & ctx, const ggml_tensor * t, const size_t n) { + if (mmb_workspace(ctx).cache.size() >= mmb_cache_max()) { delete mmb_workspace(ctx).cache.front().buf; mmb_workspace(ctx).cache.erase(mmb_workspace(ctx).cache.begin()); } + auto * buf = new ggml_cuda_pool_alloc(ctx.pool(), n); + mmb_workspace(ctx).cache.push_back({mmb_root(t), t->data, n, buf}); + return buf->get(); +} +static const uint16_t * mmb_bf16_activation(ggml_backend_cuda_context & ctx, const ggml_tensor * src1, const size_t n, cudaStream_t stream) { + const ggml_tensor * root = mmb_root(src1); + // Slot producers have graph dependencies; temporary conversions belong to their stream. + for (auto & work : mmb_state(ctx).streams) { + for (auto & e : work.slots) if (e.buf && e.root == root && e.data == src1->data && e.n == n) return e.buf->get(); + } + for (auto & e : mmb_workspace(ctx).cache) if (e.root == root && e.data == src1->data && e.n == n) return e.buf->get(); + uint16_t * buf = mmb_cache_insert(ctx, src1, n); + mmb_cvt_f32_bf16<<<(unsigned)((n / 8 + 255) / 256), 256, 0, stream>>>((const float *) src1->data, buf, n); + return buf; +} + +// Shadow BF16 copies of IQ4_NL dense weights: dequantised once (same LUT*scale -> BF16 RNE as mmb_dq_row36, so the +// WMMA inputs are bitwise identical) so the dense GEMM runs the dequant-free WTYPE=2 path. +__global__ void mmb_dq_q6k_bf16_kernel(const uint8_t * __restrict__ W, uint16_t * __restrict__ out, const size_t nblocks) { + const size_t b = (size_t) blockIdx.x * blockDim.x + threadIdx.x; + if (b >= nblocks) return; + const uint8_t * p = W + b * 210; + const uint8_t * ql = p, * qh = p + 128; + const int8_t * sc = (const int8_t *) (p + 192); + const float d = mmb_h2f(*(const uint16_t *) (p + 208)); + uint16_t * o = out + b * 256; + for (int n = 0; n < 2; ++n) { + const uint8_t * QL = ql + 64 * n; const uint8_t * QH = qh + 32 * n; const int8_t * S = sc + 8 * n; uint16_t * Y = o + 128 * n; + for (int l = 0; l < 32; ++l) { + const int is = l / 16; + const int8_t q1 = (int8_t)((QL[l + 0] & 0xF) | (((QH[l] >> 0) & 3) << 4)) - 32; + const int8_t q2 = (int8_t)((QL[l + 32] & 0xF) | (((QH[l] >> 2) & 3) << 4)) - 32; + const int8_t q3 = (int8_t)((QL[l + 0] >> 4) | (((QH[l] >> 4) & 3) << 4)) - 32; + const int8_t q4 = (int8_t)((QL[l + 32] >> 4) | (((QH[l] >> 6) & 3) << 4)) - 32; + Y[l + 0] = mmb_f2bf(d * S[is + 0] * q1); + Y[l + 32] = mmb_f2bf(d * S[is + 2] * q2); + Y[l + 64] = mmb_f2bf(d * S[is + 4] * q3); + Y[l + 96] = mmb_f2bf(d * S[is + 6] * q4); + } + } +} + +__global__ void mmb_dq_iq4nl_bf16_kernel(const uint8_t * __restrict__ W, uint16_t * __restrict__ out, const size_t nblocks) { + const size_t b = (size_t) blockIdx.x * blockDim.x + threadIdx.x; + if (b >= nblocks) return; + const uint8_t * p = W + b * 18; + const float d = mmb_h2f(*(const uint16_t *) p); + const uint32_t * q = (const uint32_t *) (p + 2); + uint32_t * o = (uint32_t *) (out + b * 32); +#pragma unroll + for (int w = 0; w < 4; ++w) { + const uint32_t v = q[w]; + const float l0 = d * mmb_kv_iq4nl[(v ) & 0xF], h0 = d * mmb_kv_iq4nl[(v >> 4) & 0xF]; + const float l1 = d * mmb_kv_iq4nl[(v >> 8) & 0xF], h1 = d * mmb_kv_iq4nl[(v >> 12) & 0xF]; + const float l2 = d * mmb_kv_iq4nl[(v >> 16) & 0xF], h2 = d * mmb_kv_iq4nl[(v >> 20) & 0xF]; + const float l3 = d * mmb_kv_iq4nl[(v >> 24) & 0xF], h3 = d * mmb_kv_iq4nl[(v >> 28) & 0xF]; + o[2*w] = mmb_pack2(l0, l1); o[2*w + 1] = mmb_pack2(l2, l3); o[8 + 2*w] = mmb_pack2(h0, h1); o[8 + 2*w + 1] = mmb_pack2(h2, h3); + } +} +int mmb_shadow_mode(){ return 2; } +bool mmb_shadow() { return mmb_shadow_mode() != 0; } +bool mmb_shadow_q6k(){ return mmb_shadow_mode() >= 1; } +size_t mmb_shadow_cap(){ return (size_t) 6144 << 20; } +static bool mmb_is_resident_q6k(const ggml_tensor * w) { return w && w->type == GGML_TYPE_Q6_K && w->op == GGML_OP_NONE && w->data && w->buffer && w->ne[2] == 1 && w->ne[3] == 1 && ggml_is_contiguous(w) && w->ne[0] % 256 == 0 && w->ne[1] <= 32768; } +static bool mmb_is_resident_iq4(const ggml_tensor * w) { return w && w->type == GGML_TYPE_IQ4_NL && w->op == GGML_OP_NONE && w->data && w->buffer && w->ne[2] == 1 && w->ne[3] == 1 && ggml_is_contiguous(w); } +static bool mmb_is_row_concat(const ggml_tensor * w) { + return w && w->op == GGML_OP_CONCAT && w->type == GGML_TYPE_IQ4_NL && ggml_get_op_params_i32(w, 0) == 1 && mmb_is_resident_iq4(w->src[0]) && mmb_is_resident_iq4(w->src[1]) && + w->src[0]->ne[0] == w->src[1]->ne[0] && w->ne[0] == w->src[0]->ne[0] && w->ne[1] == w->src[0]->ne[1] + w->src[1]->ne[1]; +} +static const uint16_t * mmb_shadow_lookup(ggml_backend_cuda_context & ctx, const ggml_tensor * w) { + if (w->op == GGML_OP_CONCAT) { auto it = mmb_state(ctx).shadow_pair.find({w->src[0]->data, w->src[1]->data}); return it == mmb_state(ctx).shadow_pair.end() ? nullptr : it->second; } + auto it = mmb_state(ctx).shadow.find(w->data); return it == mmb_state(ctx).shadow.end() ? nullptr : it->second; +} + +// RDNA3.5 (gfx1151) only, tuned for qwen4exp shapes. On gfx1151 (ROCm 7.2.1) it lost to MMQ on other archs: dense qwen35 prefill 3.4-3.9x slower, MoE 14-28% (PR #75). +// So each backend context opts in by model arch. Drop the opt-in when mmb matches MMQ on those archs. +bool mmb_enabled(const ggml_backend_cuda_context & ctx) { + return ctx.mmb_opt_in && GGML_CUDA_CC_IS_RDNA3_5(ggml_cuda_info().devices[ctx.device].cc); +} +int mmb_min_t() { return 512; } +int mmb_f32split_mode(){ return 2; } +bool mmb_f32split() { return true; } +bool mmb_bf16w() { return true; } +bool mmb_hc16() { return true; } +int mmb_tall_mode(){ return 2; } +bool mmb_tall() { return mmb_tall_mode() != 0; } +bool mmb_gatemix_flag() { return true; } +bool mmb_down16_flag() { return true; } +bool mmb_glu() { return true; } + +} // namespace + +const uint16_t * ggml_cuda_mmb_cache_lookup(ggml_backend_cuda_context & ctx, const ggml_tensor * t) { + const ggml_tensor * root = mmb_root(t); + for (auto & work : mmb_state(ctx).streams) { + for (auto & e : work.slots) if (e.buf && e.root == root && e.data == t->data) return e.buf->get(); + } + for (auto & e : mmb_workspace(ctx).cache) if (e.root == root && e.data == t->data) return e.buf->get(); + return nullptr; +} +uint16_t * ggml_cuda_mmb_slot_reserve(ggml_backend_cuda_context & ctx, int slot, const ggml_tensor * t, size_t n) { + mmb_cache_entry & e = mmb_workspace(ctx).slots[slot]; + if (e.buf && mmb_workspace(ctx).slot_cap[slot] < n) { delete e.buf; e.buf = nullptr; } + if (!e.buf) { e.buf = new ggml_cuda_pool_alloc(ctx.pool(), n); mmb_workspace(ctx).slot_cap[slot] = n; } + e.root = mmb_root(t); e.data = t->data; e.n = n; + return e.buf->get(); +} +void ggml_cuda_mmb_marks_clear(ggml_backend_cuda_context & ctx) { if (ctx.mmb) ctx.mmb->bf16_only.clear(); } +size_t ggml_cuda_mmb_marks_count(ggml_backend_cuda_context & ctx) { return ctx.mmb ? ctx.mmb->bf16_only.size() : 0; } +void ggml_cuda_mmb_mark_bf16_only(ggml_backend_cuda_context & ctx, const ggml_tensor * t) { mmb_state(ctx).bf16_only.insert(t); } +bool ggml_cuda_mmb_is_bf16_only(ggml_backend_cuda_context & ctx, const ggml_tensor * t) { return ctx.mmb && ctx.mmb->bf16_only.count(t) > 0; } +void ggml_cuda_mmb_begin_graph(ggml_backend_cuda_context & ctx) { + if (!ctx.mmb) return; + for (auto & work : mmb_state(ctx).streams) { + for (auto & e : work.cache) delete e.buf; + work.cache.clear(); + for (auto & e : work.slots) { e.root = nullptr; e.data = nullptr; e.n = 0; } + } +} +void ggml_cuda_mmb_release_all(ggml_backend_cuda_context & ctx) { + if (!ctx.mmb) return; + ggml_cuda_mmb_begin_graph(ctx); + for (auto & work : ctx.mmb->streams) { + for (auto & e : work.slots) delete e.buf; + } + for (auto & e : ctx.mmb->shadow) CUDA_CHECK(cudaFree(e.second)); + for (auto & e : ctx.mmb->shadow_pair) CUDA_CHECK(cudaFree(e.second)); + delete ctx.mmb; + ctx.mmb = nullptr; +} +// A producer writes the BF16 copy of t itself: the entry is found by the MMB GEMM that reads t next. +uint16_t * ggml_cuda_mmb_cache_produce(ggml_backend_cuda_context & ctx, const ggml_tensor * t, size_t n) { + if (!mmb_enabled(ctx)) return nullptr; + return mmb_cache_insert(ctx, t, n); +} +uint16_t * ggml_cuda_mmb_cache_reserve(ggml_backend_cuda_context & ctx, const ggml_tensor * t, size_t n) { + if (!mmb_enabled(ctx) || ggml_nrows(t) < mmb_min_t()) return nullptr; + return ggml_cuda_mmb_slot_reserve(ctx, 0, t, n); +} + +bool ggml_cuda_mmb_supported_mm(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst) { + if (!mmb_enabled(ctx)) return false; + const bool quant = mmb_quant_type(src0->type); + const bool bf16w = src0->type == GGML_TYPE_BF16 && mmb_bf16w(); + const bool f32w = src0->type == GGML_TYPE_F32 && mmb_f32split(); + if (quant && src0->ne[0] % ggml_blck_size(src0->type) != 0) return false; + if ((!quant && !bf16w && !f32w) || src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) return false; + if (src0->ne[2] != 1 || src0->ne[3] != 1) return false; + if (!ggml_is_contiguous(src0) || !ggml_is_contiguous(src1) || !ggml_is_contiguous(dst)) return false; + const int64_t K = src0->ne[0], M = src0->ne[1]; + if ((f32w ? K % 32 : K % 64) != 0 || src1->ne[0] != K || dst->ne[0] != M) return false; + const int64_t T = src1->ne[1] * src1->ne[2] * src1->ne[3]; + if (T < mmb_min_t() || T > INT32_MAX / 4) return false; + return ggml_nrows(dst) == T; +} + +bool ggml_cuda_mmb_supported_mmid(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, const ggml_tensor * dst) { + if (!mmb_enabled(ctx)) return false; + if (!mmb_quant_type(src0->type) || src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32 || ids->type != GGML_TYPE_I32) return false; + if (!ggml_is_contiguous(src0) || !ggml_is_contiguous(src1) || !ggml_is_contiguous(dst)) return false; + const int64_t K = src0->ne[0], M = src0->ne[1], E = src0->ne[2]; + if (src0->ne[3] != 1 || K % 64 != 0 || E < 1 || E > 1024) return false; + if (K % ggml_blck_size(src0->type) != 0) return false; + const int64_t n_used = ids->ne[0], T = ids->ne[1]; + if (src1->ne[0] != K || src1->ne[3] != 1 || src1->ne[2] != T) return false; + if (src1->ne[1] != 1 && src1->ne[1] != n_used) return false; + if (dst->ne[0] != M || dst->ne[1] != n_used || dst->ne[2] != T || dst->ne[3] != 1) return false; + if (ids->nb[0] != sizeof(int32_t) || ids->ne[2] != 1 || ids->ne[3] != 1) return false; + if (T < mmb_min_t() || n_used > 64 || (T * n_used) >> 16 >= 1024) return false; // tile index must fit in 16 bits per expert + return true; +} + +void ggml_cuda_mul_mat_mmb(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst) { + cudaStream_t stream = ctx.stream(); + const int K = (int) src0->ne[0], M = (int) src0->ne[1]; + const int T = (int) (src1->ne[1] * src1->ne[2] * src1->ne[3]); + if (src0->type == GGML_TYPE_F32) { + const float * W = (const float *) src0->data, * X = (const float *) src1->data; float * D = (float *) dst->data; + if (M <= 64) { + dim3 grid((M + 63) / 64, (T + 127) / 128); + mmb_f32split_kernel<64, 128, 16, 64, true, false><<>>(W, X, D, M, K, T); + } else { + dim3 grid((M + 127) / 128, (T + 127) / 128); + mmb_f32split_kernel<128, 128, 32, 64, true, false><<>>(W, X, D, M, K, T); + } + CUDA_CHECK(cudaGetLastError()); return; + } + const uint16_t * xhp = mmb_bf16_activation(ctx, src1, (size_t) T * K, stream); + const uint8_t * W = (const uint8_t *) src0->data; float * D = (float *) dst->data; + if (mmb_tall() && src0->type == GGML_TYPE_IQ4_NL && M <= 384 && K >= 4096 && T >= 2048) { // tall-M tile: HC down|inject [10240 -> 324], activations read once + static const int wide = mmb_tall_mode() >= 2; + if (wide) { dim3 grid(1, (T + 63) / 64); mmb_dense_kernel<384, 64, 96, 32, 0><<>>(W, xhp, D, (uint16_t *) nullptr, true, M, K, T); } + else { dim3 grid(1, (T + 31) / 32); mmb_dense_kernel<384, 32, 96, 16, 0><<>>(W, xhp, D, (uint16_t *) nullptr, true, M, K, T); } + CUDA_CHECK(cudaGetLastError()); + return; + } + const uint16_t * shadow_pre = ((src0->type == GGML_TYPE_IQ4_NL && mmb_shadow()) || src0->type == GGML_TYPE_Q6_K) ? mmb_shadow_lookup(ctx, src0) : nullptr; + const bool big = (M >= 6144 && K >= 2560) || (shadow_pre && K >= 2560 && T >= 4096); + uint16_t * Dh = (mmb_hc16() && K == 320 && M == 10240) ? ggml_cuda_mmb_slot_reserve(ctx, 1, dst, (size_t) T * M) : nullptr; + bool store_f32 = !(Dh && ggml_cuda_mmb_is_bf16_only(ctx, dst)); + if (ggml_cuda_mmb_blk16() && !Dh && ggml_cuda_mmb_is_bf16_only(ctx, dst) && (M & 7) == 0) { + Dh = (uint16_t *) dst->data; store_f32 = false; + } + dim3 grid((M + 127) / 128, big ? (T + 255) / 256 : (T + 127) / 128); + const uint16_t * shadow = shadow_pre; + + if (shadow) { + if (big) mmb_dense_kernel<128, 256, 64, 64, 2><<>>((const uint8_t *) shadow, xhp, D, Dh, store_f32, M, K, T); + else mmb_dense_kernel<128, 128, 32, 64, 2><<>>((const uint8_t *) shadow, xhp, D, Dh, store_f32, M, K, T); + } else if (src0->type == GGML_TYPE_IQ4_NL) { + if (big) mmb_dense_kernel<128, 256, 64, 64, 0><<>>(W, xhp, D, Dh, store_f32, M, K, T); + else mmb_dense_kernel<128, 128, 32, 64, 0><<>>(W, xhp, D, Dh, store_f32, M, K, T); + } else if (src0->type == GGML_TYPE_Q8_0) { + if (big) mmb_dense_kernel<128, 256, 64, 64, 1><<>>(W, xhp, D, Dh, store_f32, M, K, T); + else mmb_dense_kernel<128, 128, 32, 64, 1><<>>(W, xhp, D, Dh, store_f32, M, K, T); + } else if (mmb_quant_type(src0->type)) { + mmb_dispatch_quant(src0->type, [&](auto tag) { + constexpr int WT = decltype(tag)::value; + if (big) mmb_dense_kernel<128, 256, 64, 64, WT><<>>(W, xhp, D, Dh, store_f32, M, K, T); + else mmb_dense_kernel<128, 128, 32, 64, WT><<>>(W, xhp, D, Dh, store_f32, M, K, T); + }); + } else { + if (big) mmb_dense_kernel<128, 256, 64, 64, 2><<>>(W, xhp, D, Dh, store_f32, M, K, T); + else mmb_dense_kernel<128, 128, 32, 64, 2><<>>(W, xhp, D, Dh, store_f32, M, K, T); + } + CUDA_CHECK(cudaGetLastError()); +} + +bool ggml_cuda_mmb_gatemix() { return mmb_gatemix_flag(); } +bool ggml_cuda_mmb_down16() { return mmb_down16_flag(); } +bool ggml_cuda_mmb_blk16() { return true; } +bool ggml_cuda_mmb_res16() { return true; } +bool ggml_cuda_hc_gate_mix(ggml_backend_cuda_context & ctx, const ggml_tensor * w, const ggml_tensor * lo, const ggml_tensor * xn, ggml_tensor * dst, + const int hc, const float scale, const float bias) { + if (!mmb_gatemix_flag() || hc != 4 || !mmb_quant_type(w->type) || lo->type != GGML_TYPE_F32 || !ggml_is_contiguous(lo) || !ggml_is_contiguous(dst)) return false; + const int K = (int) w->ne[0], M = (int) w->ne[1], E = (int) dst->ne[0]; const int T = (int) ggml_nrows(dst); + if (K % ggml_blck_size(w->type) != 0) return false; + if (K % MMB_BK != 0 || M != hc * E || E % 32 != 0 || lo->ne[0] != K || ggml_nrows(lo) != T || xn->ne[0] != M || ggml_nrows(xn) != T || T < mmb_min_t()) return false; + const uint16_t * xn16 = ggml_cuda_mmb_cache_lookup(ctx, xn); + if (!xn16) return false; + cudaStream_t stream = ctx.stream(); + const uint16_t * lo16 = mmb_bf16_activation(ctx, lo, (size_t) T * K, stream); + uint16_t * outh = ggml_cuda_mmb_slot_reserve(ctx, 3, dst, (size_t) T * E); + const bool store_f32 = !(outh && ggml_cuda_mmb_is_bf16_only(ctx, dst)); + dim3 grid(E / 32, (T + 127) / 128); + mmb_dispatch_quant(w->type, [&](auto tag) { + constexpr int WT = decltype(tag)::value; + hc_gate_mix_kernel<4, WT><<>>((const uint8_t *) w->data, lo16, xn16, (float *) dst->data, outh, store_f32, E, K, T, scale, bias); + }); + CUDA_CHECK(cudaGetLastError()); + return true; +} + +void ggml_cuda_mul_mat_id_mmb(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst) { + cudaStream_t stream = ctx.stream(); + const int K = (int) src0->ne[0], M = (int) src0->ne[1], E = (int) src0->ne[2]; + const int ne11 = (int) src1->ne[1], T = (int) src1->ne[2], n_used = (int) ids->ne[0]; + const int n_rows_x = ne11 * T, n_rows = n_used * T; + constexpr int BN = 128; + + const uint16_t * xhp = mmb_bf16_activation(ctx, src1, (size_t) n_rows_x * K, stream); + + ggml_cuda_pool_alloc ids_src1(ctx.pool(), n_rows); + ggml_cuda_pool_alloc ids_dst(ctx.pool(), n_rows); + ggml_cuda_pool_alloc bounds(ctx.pool(), E + 1); + const int si1 = (int) (ids->nb[1] / sizeof(int32_t)); + const int sis1 = (int) (src1->nb[2] / src1->nb[1]); + if (!ggml_cuda_launch_mm_ids_bounded(ctx, (const int32_t *) ids->data, ids_src1.get(), ids_dst.get(), bounds.get(), + E, T, n_used, ne11, si1, sis1, /*inverse=*/false, stream)) { + ggml_cuda_launch_mm_ids_helper((const int32_t *) ids->data, ids_src1.get(), ids_dst.get(), bounds.get(), + E, T, n_used, ne11, si1, sis1, /*write_inverse=*/false, stream); + } + constexpr int BN_SMALL = 32, THRESH = 128; + const int nbig_max = n_rows / BN + E + 1; + const int nsmall_max = E * ((THRESH + BN_SMALL - 1) / BN_SMALL) + 1; + ggml_cuda_pool_alloc desc_big(ctx.pool(), nbig_max); + ggml_cuda_pool_alloc desc_small(ctx.pool(), nsmall_max); + mmb_build_desc2<<<1, 1024, 0, stream>>>(bounds.get(), desc_big.get(), desc_small.get(), E, nbig_max, nsmall_max, BN, BN_SMALL, THRESH); + + const uint8_t * W = (const uint8_t *) src0->data; float * D = (float *) dst->data; const size_t eb = (size_t) src0->nb[2]; + uint16_t * Dh = (mmb_down16_flag() && ggml_cuda_mmb_is_bf16_only(ctx, dst)) ? (uint16_t *) dst->data : nullptr; + const bool store_f32 = Dh == nullptr; + dim3 gbig((M + 127) / 128, nbig_max), gsmall((M + 127) / 128, nsmall_max); + mmb_dispatch_quant(src0->type, [&](auto tag) { + constexpr int WT = decltype(tag)::value; + mmb_routed_kernel<128, BN, 32, 64, WT><<>>(W, eb, xhp, D, Dh, store_f32, ids_src1.get(), ids_dst.get(), bounds.get(), desc_big.get(), M, K); + mmb_routed_kernel<128, BN_SMALL, 32, 16, WT><<>>(W, eb, xhp, D, Dh, store_f32, ids_src1.get(), ids_dst.get(), bounds.get(), desc_small.get(), M, K); + }); + CUDA_CHECK(cudaGetLastError()); +} + +bool ggml_cuda_mmb_supported_glu(ggml_backend_cuda_context & ctx, const ggml_tensor * gw, const ggml_tensor * uw, const ggml_tensor * src1, const ggml_tensor * ids, const ggml_tensor * glu) { + if (!mmb_enabled(ctx) || !mmb_glu() || !gw || !uw || !src1 || !ids || !glu) return false; + if (!mmb_quant_type(gw->type) || uw->type != gw->type) return false; + if (!ggml_are_same_shape(gw, uw) || gw->nb[1] != uw->nb[1] || gw->nb[2] != uw->nb[2]) return false; + if (glu->op != GGML_OP_GLU || ggml_get_glu_op(glu) != GGML_GLU_OP_SWIGLU || ggml_get_op_params_i32(glu, 1) != 0) return false; + if (glu->type != GGML_TYPE_F32 || !ggml_is_contiguous(glu) || !glu->src[0] || !glu->src[1]) return false; + if (glu->src[0]->op != GGML_OP_MUL_MAT_ID || glu->src[1]->op != GGML_OP_MUL_MAT_ID) return false; + if (glu->src[0]->src[0] != gw || glu->src[1]->src[0] != uw || glu->src[0]->src[1] != src1 || glu->src[1]->src[1] != src1 || glu->src[0]->src[2] != ids || glu->src[1]->src[2] != ids) return false; + if (ggml_nelements(glu) != ggml_nelements(glu->src[0]) || glu->ne[0] != gw->ne[1]) return false; + return ggml_cuda_mmb_supported_mmid(ctx, gw, src1, ids, glu->src[0]) && ggml_cuda_mmb_supported_mmid(ctx, uw, src1, ids, glu->src[1]); +} + +void ggml_cuda_mul_mat_id_mmb_glu(ggml_backend_cuda_context & ctx, const ggml_tensor * gw, const ggml_tensor * uw, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * glu) { + cudaStream_t stream = ctx.stream(); + const int K = (int) gw->ne[0], M = (int) gw->ne[1], E = (int) gw->ne[2]; + const int ne11 = (int) src1->ne[1], T = (int) src1->ne[2], n_used = (int) ids->ne[0]; + const int n_rows_x = ne11 * T, n_rows = n_used * T; + constexpr int BN = 128; + const uint16_t * xhp = mmb_bf16_activation(ctx, src1, (size_t) n_rows_x * K, stream); + ggml_cuda_pool_alloc ids_src1(ctx.pool(), n_rows); + ggml_cuda_pool_alloc ids_dst(ctx.pool(), n_rows); + ggml_cuda_pool_alloc bounds(ctx.pool(), E + 1); + const int si1 = (int) (ids->nb[1] / sizeof(int32_t)); + const int sis1 = (int) (src1->nb[2] / src1->nb[1]); + if (!ggml_cuda_launch_mm_ids_bounded(ctx, (const int32_t *) ids->data, ids_src1.get(), ids_dst.get(), bounds.get(), + E, T, n_used, ne11, si1, sis1, /*inverse=*/false, stream)) { + ggml_cuda_launch_mm_ids_helper((const int32_t *) ids->data, ids_src1.get(), ids_dst.get(), bounds.get(), + E, T, n_used, ne11, si1, sis1, /*write_inverse=*/false, stream); + } + constexpr int BN_SMALL = 32, THRESH = 128; + const int nbig_max = n_rows / BN + E + 1; + const int nsmall_max = E * ((THRESH + BN_SMALL - 1) / BN_SMALL) + 1; + ggml_cuda_pool_alloc desc_big(ctx.pool(), nbig_max); + ggml_cuda_pool_alloc desc_small(ctx.pool(), nsmall_max); + mmb_build_desc2<<<1, 1024, 0, stream>>>(bounds.get(), desc_big.get(), desc_small.get(), E, nbig_max, nsmall_max, BN, BN_SMALL, THRESH); + uint16_t * Dh = ggml_cuda_mmb_slot_reserve(ctx, 2, glu, (size_t) n_rows * M); + const bool store_f32 = !ggml_cuda_mmb_is_bf16_only(ctx, glu); + const uint8_t * Wg = (const uint8_t *) gw->data, * Wu = (const uint8_t *) uw->data; float * D = (float *) glu->data; const size_t eb = (size_t) gw->nb[2]; + dim3 gbig((M + 63) / 64, nbig_max), gsmall((M + 63) / 64, nsmall_max); + mmb_dispatch_quant(gw->type, [&](auto tag) { + constexpr int WT = decltype(tag)::value; + mmb_routed_glu_kernel<64, BN, 32, 32, WT><<>>(Wg, Wu, eb, xhp, D, Dh, store_f32, ids_src1.get(), ids_dst.get(), bounds.get(), desc_big.get(), M, K); + mmb_routed_glu_kernel<64, BN_SMALL, 16, 16, WT><<>>(Wg, Wu, eb, xhp, D, Dh, store_f32, ids_src1.get(), ids_dst.get(), bounds.get(), desc_small.get(), M, K); + }); + CUDA_CHECK(cudaGetLastError()); +} + +// Called from graph_optimize (outside stream capture): create the shadow for an eligible IQ4_NL dense weight. +void ggml_cuda_mmb_shadow_prepare(ggml_backend_cuda_context & ctx, const ggml_tensor * w) { + if (!w) return; + if (mmb_is_resident_q6k(w)) { + if (!mmb_shadow_q6k() || mmb_state(ctx).shadow.count(w->data) > 0) return; + const size_t n = (size_t) w->ne[0] * w->ne[1], bytes = n * 2; + if (mmb_state(ctx).shadow_bytes + bytes > mmb_shadow_cap()) { GGML_LOG_INFO("MMB_SHADOW cap reached; %s stays Q6_K\n", w->name); return; } + uint16_t * buf = nullptr; + if (cudaMalloc((void **) &buf, bytes) != cudaSuccess) { GGML_LOG_WARN("MMB_SHADOW alloc failed (%zu bytes)\n", bytes); return; } + mmb_dq_q6k_bf16_kernel<<<(unsigned) ((n / 256 + 255) / 256), 256, 0, ctx.stream()>>>((const uint8_t *) w->data, buf, n / 256); + CUDA_CHECK(cudaGetLastError()); + mmb_state(ctx).shadow[w->data] = buf; mmb_state(ctx).shadow_bytes += bytes; + return; + } + if (mmb_shadow_mode() != 1) return; // mode 2: Q6_K only + const bool concat = mmb_is_row_concat(w); + if (!concat && !mmb_is_resident_iq4(w)) return; + if (concat ? mmb_state(ctx).shadow_pair.count({w->src[0]->data, w->src[1]->data}) > 0 : mmb_state(ctx).shadow.count(w->data) > 0) return; + const size_t n = (size_t) w->ne[0] * w->ne[1]; + const size_t bytes = n * 2; + if (mmb_state(ctx).shadow_bytes + bytes > mmb_shadow_cap()) { static bool warned = false; if (!warned) { GGML_LOG_INFO("MMB_SHADOW cap reached at %.1f MB; further weights stay IQ4_NL\n", mmb_state(ctx).shadow_bytes / 1048576.0); warned = true; } return; } + uint16_t * buf = nullptr; + if (cudaMalloc((void **) &buf, bytes) != cudaSuccess) { GGML_LOG_WARN("MMB_SHADOW alloc failed (%zu bytes)\n", bytes); return; } + if (concat) { + const size_t n0 = (size_t) w->src[0]->ne[0] * w->src[0]->ne[1], n1 = (size_t) w->src[1]->ne[0] * w->src[1]->ne[1]; + mmb_dq_iq4nl_bf16_kernel<<<(unsigned) ((n0 / 32 + 255) / 256), 256, 0, ctx.stream()>>>((const uint8_t *) w->src[0]->data, buf, n0 / 32); + mmb_dq_iq4nl_bf16_kernel<<<(unsigned) ((n1 / 32 + 255) / 256), 256, 0, ctx.stream()>>>((const uint8_t *) w->src[1]->data, buf + n0, n1 / 32); + mmb_state(ctx).shadow_pair[{w->src[0]->data, w->src[1]->data}] = buf; + } else { + mmb_dq_iq4nl_bf16_kernel<<<(unsigned) ((n / 32 + 255) / 256), 256, 0, ctx.stream()>>>((const uint8_t *) w->data, buf, n / 32); + mmb_state(ctx).shadow[w->data] = buf; + } + CUDA_CHECK(cudaGetLastError()); + mmb_state(ctx).shadow_bytes += bytes; +} diff --git a/ggml/src/ggml-cuda/mmb.cuh b/ggml/src/ggml-cuda/mmb.cuh new file mode 100644 index 000000000000..5e547eef4cda --- /dev/null +++ b/ggml/src/ggml-cuda/mmb.cuh @@ -0,0 +1,29 @@ +#pragma once +#include "common.cuh" +// Quantized-weight BF16 WMMA GEMM on gfx1151, from 512 tokens up. +// Off unless the backend context opts in (ggml_backend_cuda_set_mmb_enabled). +bool ggml_cuda_mmb_supported_mm (ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst); +bool ggml_cuda_mmb_supported_mmid(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, const ggml_tensor * dst); +void ggml_cuda_mul_mat_mmb (ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst); +void ggml_cuda_mul_mat_id_mmb(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst); +void ggml_cuda_mmb_begin_graph(ggml_backend_cuda_context & ctx); +// producers that can emit a BF16 copy of an F32 output register it here; returns the BF16 buffer to fill (n elements) +uint16_t * ggml_cuda_mmb_cache_produce(ggml_backend_cuda_context & ctx, const ggml_tensor * t, size_t n); +uint16_t * ggml_cuda_mmb_cache_reserve(ggml_backend_cuda_context & ctx, const ggml_tensor * t, size_t n); +// BF16 copy of tensor t if one is cached for the current graph (consumers may read it instead of the F32 data) +const uint16_t * ggml_cuda_mmb_cache_lookup(ggml_backend_cuda_context & ctx, const ggml_tensor * t); +// producer slots (pinned until the next producer of the same kind): 0 = HC normalized stream xn, 1 = HC gate +uint16_t * ggml_cuda_mmb_slot_reserve(ggml_backend_cuda_context & ctx, int slot, const ggml_tensor * t, size_t n); +void ggml_cuda_mmb_marks_clear(ggml_backend_cuda_context & ctx); +size_t ggml_cuda_mmb_marks_count(ggml_backend_cuda_context & ctx); +void ggml_cuda_mmb_mark_bf16_only(ggml_backend_cuda_context & ctx, const ggml_tensor * t); +bool ggml_cuda_mmb_is_bf16_only(ggml_backend_cuda_context & ctx, const ggml_tensor * t); +bool ggml_cuda_mmb_gatemix(); +bool ggml_cuda_mmb_down16(); +bool ggml_cuda_mmb_res16(); +bool ggml_cuda_mmb_blk16(); +bool ggml_cuda_hc_gate_mix(ggml_backend_cuda_context & ctx, const ggml_tensor * w, const ggml_tensor * lo, const ggml_tensor * xn, ggml_tensor * dst, int hc, float scale, float bias); +bool ggml_cuda_mmb_supported_glu(ggml_backend_cuda_context & ctx, const ggml_tensor * gw, const ggml_tensor * uw, const ggml_tensor * src1, const ggml_tensor * ids, const ggml_tensor * glu); +void ggml_cuda_mul_mat_id_mmb_glu(ggml_backend_cuda_context & ctx, const ggml_tensor * gw, const ggml_tensor * uw, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * glu); +void ggml_cuda_mmb_shadow_prepare(ggml_backend_cuda_context & ctx, const ggml_tensor * w); +void ggml_cuda_mmb_release_all(ggml_backend_cuda_context & ctx); diff --git a/ggml/src/ggml-cuda/mmid.cu b/ggml/src/ggml-cuda/mmid.cu index 0b222e63ac7d..f61f4dff6497 100644 --- a/ggml/src/ggml-cuda/mmid.cu +++ b/ggml/src/ggml-cuda/mmid.cu @@ -19,6 +19,57 @@ struct mm_ids_helper_store { }; static_assert(sizeof(mm_ids_helper_store) == 4, "unexpected size for mm_ids_helper_store"); +// Qwen4exp 512-expert / 10-active fast path. The generic helper launches one warp per expert and every +// warp scans all token slots (O(n_experts*n_tokens*n_used)). Here one block histograms the slots, computes +// the 512 bounds, and scatters the two maps in O(n_tokens*n_used). Row ordering inside an expert is irrelevant: +// ids_dst maps every compact MMQ row back to its exact destination and there is no cross-row reduction. +__launch_bounds__(1024, 1) +static __global__ void mm_ids_helper_512_10( + const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, + int32_t * __restrict__ expert_bounds, const int n_tokens, const int nchannels_y, + const int si1, const int sis1, const bool write_inverse) { + __shared__ int counts[512]; + __shared__ int cursors[512]; + + const int tid = threadIdx.x; + if (tid < 512) { + counts[tid] = 0; + cursors[tid] = 0; + } + __syncthreads(); + + const int n_slots = n_tokens * 10; + for (int idx = tid; idx < n_slots; idx += blockDim.x) { + const int token = idx / 10; + const int iex = idx - token * 10; + atomicAdd(&counts[ids[token * si1 + iex]], 1); + } + __syncthreads(); + + if (tid == 0) { + int total = 0; + expert_bounds[0] = 0; + for (int e = 0; e < 512; ++e) { + total += counts[e]; + expert_bounds[e + 1] = total; + } + } + __syncthreads(); + + for (int idx = tid; idx < n_slots; idx += blockDim.x) { + const int token = idx / 10; + const int iex = idx - token * 10; + const int expert = ids[token * si1 + iex]; + const int dst = expert_bounds[expert] + atomicAdd(&cursors[expert], 1); + ids_dst[dst] = idx; + if (write_inverse) { + ids_src1[idx] = dst; + } else { + ids_src1[dst] = token * sis1 + iex % nchannels_y; + } + } +} + // the generic path passes 0, which needs no padding since it never groups lanes by token template struct mm_ids_pow2 { static constexpr int value = 2*mm_ids_pow2<(n + 1)/2>::value; }; template <> struct mm_ids_pow2<1> { static constexpr int value = 1; }; @@ -152,6 +203,12 @@ static void launch_mm_ids_helper( void ggml_cuda_launch_mm_ids_helper( const int32_t * __restrict__ ids, int32_t * __restrict__ ids_src1, int32_t * __restrict__ ids_dst, int32_t * __restrict__ expert_bounds, const int n_experts, const int n_tokens, const int n_expert_used, const int nchannels_y, const int si1, const int sis1, const bool write_inverse, cudaStream_t stream) { + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + if (GGML_CUDA_CC_IS_RDNA3_5(cc) && n_experts == 512 && n_expert_used == 10 && getenv("GGML_CUDA_DISABLE_MMID_512") == nullptr) { + mm_ids_helper_512_10<<<1, 1024, 0, stream>>>( + ids, ids_src1, ids_dst, expert_bounds, n_tokens, nchannels_y, si1, sis1, write_inverse); + return; + } switch (n_expert_used) { case 2: launch_mm_ids_helper< 2>(ids, ids_src1, ids_dst, expert_bounds, n_experts, n_tokens, n_expert_used, nchannels_y, si1, sis1, write_inverse, stream); @@ -179,3 +236,122 @@ void ggml_cuda_launch_mm_ids_helper( break; } } + +static constexpr int MM_IDS_ROUTE_TILE = 1024; +static constexpr int MM_IDS_ROUTE_BITS = 10; +static constexpr int MM_IDS_ROUTE_EXPERTS = 512; +static constexpr int MM_IDS_ROUTE_USED = 10; + +static __global__ void mm_ids_route_sort( + const int32_t * ids, uint32_t * sorted, int32_t * offsets, int32_t * starts, + int routes, int chunks, int ids_stride) { + __shared__ uint32_t keys[MM_IDS_ROUTE_TILE]; + __shared__ int first[MM_IDS_ROUTE_EXPERTS]; + __shared__ int counts[MM_IDS_ROUTE_EXPERTS]; + const int chunk=blockIdx.x; + for (int i=threadIdx.x;i=0 && expert0;distance/=2) { + for (int i=threadIdx.x;ii) { + const uint32_t a=keys[i],b=keys[partner]; + if (((i&span)==0 && a>b) || ((i&span)!=0 && a>MM_IDS_ROUTE_BITS; + if (i==0 || (keys[i-1]>>MM_IDS_ROUTE_BITS)!=uint32_t(expert)) first[expert]=i; + } + sorted[chunk*MM_IDS_ROUTE_TILE+i]=keys[i]; + } + __syncthreads(); + for (int i=threadIdx.x;i>MM_IDS_ROUTE_BITS; + if (i+1==MM_IDS_ROUTE_TILE || (keys[i+1]>>MM_IDS_ROUTE_BITS)!=uint32_t(expert)) { + counts[expert]=i+1-first[expert]; + } + } + } + __syncthreads(); + for (int i=threadIdx.x;i=chunks*MM_IDS_ROUTE_TILE) return; + const uint32_t key=sorted[index]; + if (key==UINT32_MAX) return; + const int chunk=index/MM_IDS_ROUTE_TILE; + const int expert=key>>MM_IDS_ROUTE_BITS; + const int local=index%MM_IDS_ROUTE_TILE; + const int route=chunk*MM_IDS_ROUTE_TILE+(key&(MM_IDS_ROUTE_TILE-1)); + const int target=offsets[chunk*MM_IDS_ROUTE_EXPERTS+expert]+local-starts[chunk*MM_IDS_ROUTE_EXPERTS+expert]; + dst_map[target]=route; + if (inverse) src_map[route]=target; + else src_map[target]=(route/MM_IDS_ROUTE_USED)*token_stride+(route%MM_IDS_ROUTE_USED)%channels; +} + +bool ggml_cuda_launch_mm_ids_bounded( + ggml_backend_cuda_context & ctx, const int32_t * ids, int32_t * src_map, + int32_t * dst_map, int32_t * bounds, int experts, int tokens, int used, + int channels, int ids_stride, int token_stride, bool inverse, cudaStream_t stream) { + const int device=ctx.device; + if (!GGML_CUDA_CC_IS_RDNA3_5(ggml_cuda_info().devices[device].cc) || + experts!=MM_IDS_ROUTE_EXPERTS || used!=MM_IDS_ROUTE_USED || + tokens<=4096 || tokens>32768 || (channels!=1 && channels!=used) || + ids_stride sorted(ctx.pool(device),size_t(chunks)*MM_IDS_ROUTE_TILE); + ggml_cuda_pool_alloc offsets(ctx.pool(device),size_t(chunks)*experts); + ggml_cuda_pool_alloc starts(ctx.pool(device),size_t(chunks)*experts); + mm_ids_route_sort<<>>(ids,sorted.get(),offsets.get(),starts.get(),routes,chunks,ids_stride); + mm_ids_route_prefix<<<1,MM_IDS_ROUTE_EXPERTS,0,stream>>>(offsets.get(),bounds,chunks); + mm_ids_route_scatter<<<(chunks*MM_IDS_ROUTE_TILE+255)/256,256,0,stream>>> + (sorted.get(),offsets.get(),starts.get(),src_map,dst_map,chunks,channels,token_stride,inverse); + CUDA_CHECK(cudaGetLastError()); + return true; +} diff --git a/ggml/src/ggml-cuda/mmid.cuh b/ggml/src/ggml-cuda/mmid.cuh index 74c2db43385e..96de66b48d07 100644 --- a/ggml/src/ggml-cuda/mmid.cuh +++ b/ggml/src/ggml-cuda/mmid.cuh @@ -3,3 +3,9 @@ void ggml_cuda_launch_mm_ids_helper( const int32_t * ids, int32_t * ids_src1, int32_t * ids_dst, int32_t * expert_bounds, int n_experts, int n_tokens, int n_expert_used, int nchannels_y, int si1, int sis1, bool write_inverse, cudaStream_t stream); + +struct ggml_backend_cuda_context; +bool ggml_cuda_launch_mm_ids_bounded(ggml_backend_cuda_context & ctx, + const int32_t * ids, int32_t * src_map, int32_t * dst_map, int32_t * bounds, + int experts, int tokens, int used, int channels, int ids_stride, int token_stride, + bool inverse, cudaStream_t stream); diff --git a/ggml/src/ggml-cuda/mmq-config-rdna3-5.cuh b/ggml/src/ggml-cuda/mmq-config-rdna3-5.cuh index 180b2d9370d1..f8a9e6350926 100644 --- a/ggml/src/ggml-cuda/mmq-config-rdna3-5.cuh +++ b/ggml/src/ggml-cuda/mmq-config-rdna3-5.cuh @@ -79,15 +79,15 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf CASE(GGML_TYPE_Q8_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); CASE(GGML_TYPE_Q8_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); CASE(GGML_TYPE_Q8_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q8_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); CASE(GGML_TYPE_Q8_0, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q8_0, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q8_0, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q8_0, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q8_0, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q8_0, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q8_0, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q8_0, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q8_0, 256, 2, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); // --------------------------------------------------------------------------------------------- @@ -115,42 +115,42 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf CASE(GGML_TYPE_Q4_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); CASE(GGML_TYPE_Q4_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q4_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q4_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); CASE(GGML_TYPE_Q4_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q4_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q4_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q4_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q4_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q4_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q4_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q4_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q4_K, 128, 2, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q5_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); CASE(GGML_TYPE_Q5_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q5_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q5_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); CASE(GGML_TYPE_Q5_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q5_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q5_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q5_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q5_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q5_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q5_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q5_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q5_K, 128, 2, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q6_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); CASE(GGML_TYPE_Q6_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q6_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_Q6_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, true); CASE(GGML_TYPE_Q6_K, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); CASE(GGML_TYPE_Q6_K, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q6_K, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q6_K, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q6_K, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q6_K, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q6_K, 256, 2, 128, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_Q6_K, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 112, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_Q6_K, 128, 2, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q6_K, MMQ_ITER_K, false, false); // --------------------------------------------------------------------------------------------- @@ -212,7 +212,7 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, true); CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); - CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); + CASE(GGML_TYPE_IQ3_XXS, 128, 2, 64, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 80, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); CASE(GGML_TYPE_IQ3_XXS, 256, 2, 128, 96, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, MMQ_ITER_K, false, false); @@ -263,7 +263,7 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf CASE(GGML_TYPE_MXFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); CASE(GGML_TYPE_MXFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); CASE(GGML_TYPE_MXFP4, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); - CASE(GGML_TYPE_MXFP4, 256, 2, 128, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); + CASE(GGML_TYPE_MXFP4, 128, 2, 64, 128, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, true); CASE(GGML_TYPE_MXFP4, 128, 2, 64, 16, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); CASE(GGML_TYPE_MXFP4, 128, 2, 64, 32, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); CASE(GGML_TYPE_MXFP4, 256, 2, 128, 48, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_1, MMQ_ITER_K, false, false); @@ -288,3 +288,14 @@ static constexpr __host__ __device__ ggml_cuda_mmq_config ggml_cuda_mmq_get_conf return ggml_cuda_mmq_config(GGML_TYPE_COUNT, 256, 2, 128, 64, GGML_CUDA_MMQ_SRAM_LAYOUT_Q8_0, 256, false, true); } + +static_assert(ggml_cuda_mmq_get_config_rdna3_5(GGML_TYPE_Q8_0, 128, false).nthreads == 256); +static_assert(ggml_cuda_mmq_get_config_rdna3_5(GGML_TYPE_Q8_0, 128, false).I == 64); +static_assert(ggml_cuda_mmq_get_config_rdna3_5(GGML_TYPE_Q8_0, 128, true).nthreads == 128); +static_assert(ggml_cuda_mmq_get_config_rdna3_5(GGML_TYPE_Q8_0, 128, true).I == 64); +static_assert(ggml_cuda_mmq_get_config_rdna3_5(GGML_TYPE_Q8_0, 48, false).nthreads == 128); +static_assert(ggml_cuda_mmq_get_config_rdna3_5(GGML_TYPE_Q8_0, 48, false).I == 64); +static_assert(ggml_cuda_mmq_get_config_rdna3_5(GGML_TYPE_Q8_0, 48, true).nthreads == 256); +static_assert(ggml_cuda_mmq_get_config_rdna3_5(GGML_TYPE_Q8_0, 48, true).I == 128); +static_assert(ggml_cuda_mmq_get_config_rdna3_5(GGML_TYPE_MXFP4, 128, true).nthreads == 128); +static_assert(ggml_cuda_mmq_get_config_rdna3_5(GGML_TYPE_MXFP4, 128, true).I == 64); diff --git a/ggml/src/ggml-cuda/mmq-load-tiles.cuh b/ggml/src/ggml-cuda/mmq-load-tiles.cuh index 7f00bad943e2..40c825405e26 100644 --- a/ggml/src/ggml-cuda/mmq-load-tiles.cuh +++ b/ggml/src/ggml-cuda/mmq-load-tiles.cuh @@ -532,6 +532,102 @@ template static __device__ __forceinline_ } } +// Register-staged variant of the Q8_0 tile loader used for software prefetching on RDNA3.5: +// the global loads for the next K iteration are issued into registers while the current tile is consumed, +// then written to shared memory with exactly the same layout as ggml_cuda_mmq_load_tiles_q8_0. +template struct ggml_cuda_mmq_x_regs_q8_0 { + // Only instantiated for RDNA3.5 (wave32) device code; the host pass needs constant sizes. + static constexpr int warp_size = WARP_SIZE; + static constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size; + static constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); + static constexpr int threads_per_row = 32; + static constexpr int nrows = warp_size / threads_per_row; + static constexpr int rows_per_step = nrows*nwarps; + static constexpr int n_qs = (I + rows_per_step - 1) / rows_per_step; + static constexpr int blocks_per_tile_x_row = 2*MMQ_TILE_NE_K / QI8_0; + static constexpr int rows_per_warp = warp_size / blocks_per_tile_x_row; + static constexpr int d_rows_per_step = nwarps*rows_per_warp; + static constexpr int n_d = (I + d_rows_per_step - 1) / d_rows_per_step; + + int qs[2*n_qs]; + half d[n_d]; +}; + +template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q8_0_regs( + const char * __restrict__ x, ggml_cuda_mmq_x_regs_q8_0 & r, const int kbx0, const int i_max, const int stride) { + using R = ggml_cuda_mmq_x_regs_q8_0; + static_assert(R::warp_size == ggml_cuda_get_physical_warp_size(), "register-staged Q8_0 loader assumes wave32"); + + const int txi = R::warp_size > R::threads_per_row ? threadIdx.x % R::threads_per_row : threadIdx.x; + const int kbx = txi / QI8_0; + const int kqsx = txi % QI8_0; + +#pragma unroll + for (int i0 = 0; i0 < R::I; i0 += R::rows_per_step) { + int i = i0 + (R::nrows == 1 ? threadIdx.y : threadIdx.y*R::nrows + threadIdx.x/R::threads_per_row); + + if (fallback) { + i = min(i, i_max); + } + + const block_q8_0 * bxi = (const block_q8_0 *) x + kbx0 + i*stride + kbx; + + r.qs[2*(i0/R::rows_per_step) + 0] = get_int_b2(bxi[0].qs, kqsx); + r.qs[2*(i0/R::rows_per_step) + 1] = get_int_b2(bxi[MMQ_TILE_NE_K/QI8_0].qs, kqsx); + } + + const int kbxd = threadIdx.x % R::blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < R::I; i0 += R::d_rows_per_step) { + int i = i0 + threadIdx.y*R::rows_per_warp + threadIdx.x/R::blocks_per_tile_x_row; + + if (fallback) { + i = min(i, i_max); + } + + const block_q8_0 * bxi = (const block_q8_0 *) x + kbx0 + i*stride + kbxd; + + r.d[i0/R::d_rows_per_step] = bxi->d; + } +} + +template static __device__ __forceinline__ void ggml_cuda_mmq_store_tiles_q8_0_regs( + const ggml_cuda_mmq_x_regs_q8_0 & r, int * __restrict__ x_tile, const int i_max) { + using R = ggml_cuda_mmq_x_regs_q8_0; + constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback); + + int * x_qs = (int *) x_tile; + float * x_df = (float *) (x_tile + 2*MMQ_TILE_NE_K); + + const int txi = R::warp_size > R::threads_per_row ? threadIdx.x % R::threads_per_row : threadIdx.x; + +#pragma unroll + for (int i0 = 0; i0 < R::I; i0 += R::rows_per_step) { + int i = i0 + (R::nrows == 1 ? threadIdx.y : threadIdx.y*R::nrows + threadIdx.x/R::threads_per_row); + + if (fallback) { + i = min(i, i_max); + } + + x_qs[i*sram_stride + 0 + txi] = r.qs[2*(i0/R::rows_per_step) + 0]; + x_qs[i*sram_stride + MMQ_TILE_NE_K + txi] = r.qs[2*(i0/R::rows_per_step) + 1]; + } + + const int kbxd = threadIdx.x % R::blocks_per_tile_x_row; + +#pragma unroll + for (int i0 = 0; i0 < R::I; i0 += R::d_rows_per_step) { + int i = i0 + threadIdx.y*R::rows_per_warp + threadIdx.x/R::blocks_per_tile_x_row; + + if (fallback) { + i = min(i, i_max); + } + + x_df[i*sram_stride + kbxd] = r.d[i0/R::d_rows_per_step]; + } +} + // --------------------------------------------------------------------------------------------- template static __device__ __forceinline__ void ggml_cuda_mmq_load_tiles_q2_K( diff --git a/ggml/src/ggml-cuda/mmq-q45.cu b/ggml/src/ggml-cuda/mmq-q45.cu new file mode 100644 index 000000000000..e9774527cc61 --- /dev/null +++ b/ggml/src/ggml-cuda/mmq-q45.cu @@ -0,0 +1,9 @@ +#define mul_mat_q mul_mat_q_q45 +#define launch_mul_mat_q launch_mul_mat_q_q45 +#include "mmq.cuh" +#undef launch_mul_mat_q +#undef mul_mat_q + +void launch_mul_mat_q_q5_k_j32(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) { + launch_mul_mat_q_q45(ctx, args, stream); +} diff --git a/ggml/src/ggml-cuda/mmq-q6.cu b/ggml/src/ggml-cuda/mmq-q6.cu new file mode 100644 index 000000000000..ed9548f79608 --- /dev/null +++ b/ggml/src/ggml-cuda/mmq-q6.cu @@ -0,0 +1,9 @@ +#define mul_mat_q mul_mat_q_q6 +#define launch_mul_mat_q launch_mul_mat_q_q6 +#include "mmq.cuh" +#undef launch_mul_mat_q +#undef mul_mat_q + +void launch_mul_mat_q_q6_k_j32(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) { + launch_mul_mat_q_q6(ctx, args, stream); +} diff --git a/ggml/src/ggml-cuda/mmq-vec-dot.cuh b/ggml/src/ggml-cuda/mmq-vec-dot.cuh index 4d1c398fc545..cedf0a184fbb 100644 --- a/ggml/src/ggml-cuda/mmq-vec-dot.cuh +++ b/ggml/src/ggml-cuda/mmq-vec-dot.cuh @@ -151,8 +151,15 @@ static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma( constexpr int sram_stride = ggml_cuda_mmq_get_sram_stride(type, J, fallback); constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback); constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / ggml_cuda_get_physical_warp_size(); + constexpr bool split_j = type == GGML_TYPE_Q8_0 && J == 128 && !fallback && I == 64 && nwarps == 8; + constexpr int j_group = split_j ? J/2 : J; - y += (threadIdx.y % ntx) * (tile_C::J*MMQ_TILE_Y_K); + const int warp_i = split_j ? threadIdx.y % 4 : threadIdx.y; + const int warp_j = split_j ? threadIdx.y / 4 : 0; + + y += (warp_j*j_group + (warp_i % 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; @@ -160,7 +167,7 @@ static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma( const float * y_df = (const float *) y; const half2 * y_ds = (const half2 *) y; - const int i0 = (threadIdx.y / ntx) * rows_per_warp; + const int i0 = (warp_i / ntx) * rows_per_warp; for (int k01 = 0; k01 < MMQ_TILE_NE_K; k01 += QI8_0) { const int k0 = k00 + k01; @@ -172,7 +179,7 @@ static __device__ __forceinline__ void ggml_cuda_mmq_vec_dot_q8_0_q8_1_mma( } #pragma unroll - for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) { + for (int j0 = 0; j0 < j_group; j0 += ntx*tile_C::J) { tile_B B; load_ldmatrix(B, y_qs + j0*MMQ_TILE_Y_K + k01, MMQ_TILE_Y_K); @@ -1237,4 +1244,3 @@ template static __device__ __forceinline_ } } } - diff --git a/ggml/src/ggml-cuda/mmq.cu b/ggml/src/ggml-cuda/mmq.cu index 9b6038adff9e..0e5479dc8eea 100644 --- a/ggml/src/ggml-cuda/mmq.cu +++ b/ggml/src/ggml-cuda/mmq.cu @@ -5,6 +5,26 @@ #include +void launch_mul_mat_q_q5_k_j32(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream); +void launch_mul_mat_q_q6_k_j32(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream); + +static bool use_rdna3_5_q5_ling(const mmq_args & args) { + const int id = ggml_cuda_get_device(); + const int cc = ggml_cuda_info().devices[id].cc; + return args.type_x == GGML_TYPE_Q5_K && GGML_CUDA_CC_IS_RDNA3_5(cc) && args.ids_dst != nullptr && + ((args.ncols_x == 2560 && args.nrows_x == 768) || (args.ncols_x == 768 && args.nrows_x == 2560)) && + args.nchannels_x == 512 && args.nchannels_y > 0 && + (args.ncols_dst + args.nchannels_y - 1) / args.nchannels_y == 32; +} + +static bool use_rdna3_5_q6_ling_j32(const mmq_args & args) { + const int id = ggml_cuda_get_device(); + const int cc = ggml_cuda_info().devices[id].cc; + return args.type_x == GGML_TYPE_Q6_K && GGML_CUDA_CC_IS_RDNA3_5(cc) && args.ids_dst != nullptr && + args.ncols_x == 768 && args.nrows_x == 2560 && args.nchannels_x == 512 && args.nchannels_y > 0 && + (args.ncols_dst + args.nchannels_y - 1) / args.nchannels_y == 32; +} + static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) { switch (args.type_x) { case GGML_TYPE_Q1_0: @@ -39,10 +59,18 @@ static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, con mul_mat_q_case(ctx, args, stream); break; case GGML_TYPE_Q5_K: - mul_mat_q_case(ctx, args, stream); + if (use_rdna3_5_q5_ling(args)) { + launch_mul_mat_q_q5_k_j32(ctx, args, stream); + } else { + mul_mat_q_case(ctx, args, stream); + } break; case GGML_TYPE_Q6_K: - mul_mat_q_case(ctx, args, stream); + if (use_rdna3_5_q6_ling_j32(args)) { + launch_mul_mat_q_q6_k_j32(ctx, args, stream); + } else { + mul_mat_q_case(ctx, args, stream); + } break; // ----------------------------------------------------------------------- case GGML_TYPE_IQ1_S: @@ -82,8 +110,144 @@ static void ggml_cuda_mul_mat_q_switch_type(ggml_backend_cuda_context & ctx, con } } -void ggml_cuda_mul_mat_q( - ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst) { +void ggml_cuda_mul_mat_q_pair(ggml_backend_cuda_context & ctx, ggml_tensor * dst0, ggml_tensor * dst1) { + GGML_ASSERT(dst0->src[1] == dst1->src[1]); + GGML_ASSERT(dst0->src[2] == dst1->src[2]); + + const ggml_tensor * src0s[] = { dst0->src[0], dst1->src[0] }; + ggml_tensor * dsts[] = { dst0, dst1 }; + const ggml_tensor * src1 = dst0->src[1]; + const ggml_tensor * ids = dst0->src[2]; + const ggml_tensor * src0 = src0s[0]; + + GGML_ASSERT(src1->type == GGML_TYPE_F32); + if (ids == nullptr) { + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + cudaStream_t stream = ctx.stream(); + const int64_t ne10_padded = GGML_PAD(src1->ne[0], MATRIX_ROW_PADDING); + size_t nbytes_src1_q8_1 = src1->ne[3] * src1->ne[2] * src1->ne[1] * ne10_padded * sizeof(block_q8_1_mmq) / QK8_1_MMQ; + + for (int i = 0; i < 2; ++i) { + const ggml_tensor * src0_i = src0s[i]; + ggml_tensor * dst_i = dsts[i]; + GGML_ASSERT(dst_i->type == GGML_TYPE_F32 && dst_i->src[1] == src1); + GGML_ASSERT(ggml_are_same_shape(src0, src0_i) && ggml_are_same_shape(dst0, dst_i)); + GGML_ASSERT(src0_i->ne[0] == src1->ne[0]); + GGML_ASSERT(mmq_get_q8_1_ds_layout(src0->type) == mmq_get_q8_1_ds_layout(src0_i->type)); + const bool fallback = src0_i->ne[1] % 128 != 0; + nbytes_src1_q8_1 = std::max(nbytes_src1_q8_1, + src1->ne[3] * src1->ne[2] * src1->ne[1] * ne10_padded * sizeof(block_q8_1_mmq) / QK8_1_MMQ + + ggml_cuda_mmq_get_J_max(src0_i->type, fallback, cc, src1->ne[1]) * sizeof(block_q8_1_mmq)); + } + + ggml_cuda_pool_alloc src1_q8_1(ctx.pool(), nbytes_src1_q8_1); + quantize_mmq_q8_1_cuda((const float *) src1->data, nullptr, src1_q8_1.get(), src0->type, + src1->ne[0], src1->nb[1] / sizeof(float), src1->nb[2] / sizeof(float), src1->nb[3] / sizeof(float), + ne10_padded, src1->ne[1], src1->ne[2], src1->ne[3], stream); + CUDA_CHECK(cudaGetLastError()); + + const int64_t s12_q = src1->ne[1] * ne10_padded * sizeof(block_q8_1) / (QK8_1 * sizeof(int)); + const int64_t s13_q = src1->ne[2] * s12_q; + for (int i = 0; i < 2; ++i) { + const ggml_tensor * src0_i = src0s[i]; + ggml_tensor * dst_i = dsts[i]; + GGML_ASSERT(ggml_are_same_shape(src0, src0_i) && ggml_are_same_shape(dst0, dst_i)); + const mmq_args args = { + (const char *) src0_i->data, src0_i->type, (const int *) src1_q8_1.get(), nullptr, nullptr, (float *) dst_i->data, + nullptr, + src0_i->ne[0], src0_i->ne[1], src1->ne[1], (int64_t) (src0_i->nb[1] / ggml_type_size(src0_i->type)), src1->ne[1], (int64_t) (dst_i->nb[1] / sizeof(float)), + src0_i->ne[2], src1->ne[2], (int64_t) (src0_i->nb[2] / ggml_type_size(src0_i->type)), s12_q, (int64_t) (dst_i->nb[2] / sizeof(float)), + src0_i->ne[3], src1->ne[3], (int64_t) (src0_i->nb[3] / ggml_type_size(src0_i->type)), s13_q, (int64_t) (dst_i->nb[3] / sizeof(float)), + src1->ne[1]}; + ggml_cuda_mul_mat_q_switch_type(ctx, args, stream); + } + return; + } + + GGML_ASSERT(ids->type == GGML_TYPE_I32); + GGML_ASSERT(src1->ne[3] == 1); + GGML_ASSERT(src1->nb[2] % src1->nb[1] == 0); + GGML_ASSERT(ids->nb[0] == ggml_element_size(ids)); + + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + cudaStream_t stream = ctx.stream(); + const int64_t n_expert_used = ids->ne[0]; + const int64_t n_tokens = src1->ne[2]; + const int64_t ne_get_rows = n_tokens*n_expert_used; + const int64_t ne10_padded = GGML_PAD(src1->ne[0], MATRIX_ROW_PADDING); + const bool dedup_bcast = src1->ne[1] == 1 && n_expert_used > 1; + + ggml_cuda_pool_alloc ids_src1(ctx.pool(), ne_get_rows); + ggml_cuda_pool_alloc ids_dst(ctx.pool(), ne_get_rows); + ggml_cuda_pool_alloc expert_bounds(ctx.pool(), src0->ne[2] + 1); + + const int si1 = ids->nb[1] / ggml_element_size(ids); + const int sis1 = src1->nb[2] / src1->nb[1]; + if (!ggml_cuda_launch_mm_ids_bounded(ctx, (const int32_t *) ids->data, ids_src1.get(), ids_dst.get(), expert_bounds.get(), + src0->ne[2], n_tokens, n_expert_used, src1->ne[1], si1, sis1, dedup_bcast, stream)) { + ggml_cuda_launch_mm_ids_helper((const int32_t *) ids->data, ids_src1.get(), ids_dst.get(), expert_bounds.get(), + src0->ne[2], n_tokens, n_expert_used, src1->ne[1], si1, sis1, /*write_inverse =*/ dedup_bcast, stream); + } + CUDA_CHECK(cudaGetLastError()); + + size_t nbytes_src1_q8_1 = n_tokens*n_expert_used*ne10_padded * sizeof(block_q8_1_mmq)/QK8_1_MMQ; + for (int i = 0; i < 2; ++i) { + const ggml_tensor * src0_i = src0s[i]; + ggml_tensor * dst_i = dsts[i]; + GGML_ASSERT(dst_i->type == GGML_TYPE_F32); + GGML_ASSERT(dst_i->nb[2] % dst_i->nb[1] == 0); + GGML_ASSERT(ggml_are_same_shape(src0, src0_i)); + GGML_ASSERT(ggml_are_same_shape(dst0, dst_i)); + GGML_ASSERT(mmq_get_q8_1_ds_layout(src0->type) == mmq_get_q8_1_ds_layout(src0_i->type)); + GGML_ASSERT(src0_i->ne[0] == src1->ne[0]); + GGML_ASSERT(dst_i->ne[1] == n_expert_used); + + const bool fallback = src0_i->ne[1] % 128 != 0; + nbytes_src1_q8_1 = std::max(nbytes_src1_q8_1, + n_tokens*n_expert_used*ne10_padded * sizeof(block_q8_1_mmq)/QK8_1_MMQ + + ggml_cuda_mmq_get_J_max(src0_i->type, fallback, cc, src1->ne[1]) * sizeof(block_q8_1_mmq)); + } + + ggml_cuda_pool_alloc src1_q8_1(ctx.pool(), nbytes_src1_q8_1); + const int64_t s11 = src1->nb[1] / ggml_type_size(src1->type); + const int64_t s12_src = src1->nb[2] / ggml_type_size(src1->type); + const int64_t s13_src = src1->nb[3] / ggml_type_size(src1->type); + if (dedup_bcast) { + quantize_scatter_mmq_q8_1_cuda((const float *) src1->data, ids_src1.get(), src1_q8_1.get(), src0->type, + src1->ne[0], /*stride_token=*/s12_src, ne10_padded, n_tokens, ne_get_rows, n_expert_used, stream); + } else { + quantize_mmq_q8_1_cuda((const float *) src1->data, ids_src1.get(), src1_q8_1.get(), src0->type, + src1->ne[0], s11, s12_src, s13_src, ne10_padded, ne_get_rows, 1, 1, stream); + } + CUDA_CHECK(cudaGetLastError()); + + const int64_t s12_q = src1->ne[1] * ne10_padded * sizeof(block_q8_1) / (QK8_1 * sizeof(int)); + const int64_t s13_q = n_tokens*s12_q; + for (int i = 0; i < 2; ++i) { + const ggml_tensor * src0_i = src0s[i]; + ggml_tensor * dst_i = dsts[i]; + const size_t ts_src0 = ggml_type_size(src0_i->type); + const size_t ts_dst = ggml_type_size(dst_i->type); + const int64_t s01 = src0_i->nb[1] / ts_src0; + const int64_t s02 = src0_i->nb[2] / ts_src0; + const int64_t s03 = src0_i->nb[3] / ts_src0; + const int64_t s1 = dst_i->nb[1] / ts_dst; + const int64_t s2 = dst_i->nb[2] / ts_dst; + const int64_t s3 = dst_i->nb[3] / ts_dst; + const mmq_args args = { + (const char *) src0_i->data, src0_i->type, (const int *) src1_q8_1.get(), ids_dst.get(), expert_bounds.get(), (float *) dst_i->data, + nullptr, + src0_i->ne[0], src0_i->ne[1], ne_get_rows, s01, ne_get_rows, s1, + src0_i->ne[2], src0_i->ne[2], s02, s12_q, s2, + src0_i->ne[3], src1->ne[3], s03, s13_q, s3, + n_tokens}; + ggml_cuda_mul_mat_q_switch_type(ctx, args, stream); + } +} + +static void ggml_cuda_mul_mat_q_impl( + ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, + ggml_tensor * dst, const ggml_tensor * swiglu) { GGML_ASSERT( src1->type == GGML_TYPE_F32); GGML_ASSERT( dst->type == GGML_TYPE_F32); GGML_ASSERT(!ids || ids->type == GGML_TYPE_I32); // Optional, used for batched GGML_MUL_MAT_ID. @@ -102,6 +266,14 @@ void ggml_cuda_mul_mat_q( GGML_ASSERT( nb0 == ts_dst); GGML_ASSERT(!ids || ids->nb[0] == ggml_type_size(ids->type)); + const ggml_tensor * gate = swiglu ? swiglu->src[0] : nullptr; + const ggml_tensor * up = swiglu ? swiglu->src[1] : nullptr; + if (swiglu) { + GGML_ASSERT(src0->type == GGML_TYPE_Q8_0 || src0->type == GGML_TYPE_IQ4_NL); + GGML_ASSERT(gate && up && gate->type == GGML_TYPE_F32 && up->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_are_same_shape(gate, up) && ggml_are_same_shape(gate, src1)); + } + const char * src0_d = (const char *) src0->data; const float * src1_d = (const float *) src1->data; float * dst_d = (float *) dst->data; @@ -152,6 +324,12 @@ void ggml_cuda_mul_mat_q( quantize_mmq_fp4_cuda(src1_d, nullptr, src1_q8_1.get(), src1_scale.ptr, src0->type, use_aligned_float8, ne10, s11, s12, s13, ne10_padded, ne11, ne12, ne13, stream); + } else if (swiglu) { + quantize_mmq_q8_1_swiglu_cuda( + (const float *) gate->data, (const float *) up->data, nullptr, src1_q8_1.get(), src0->type, + ne10, ne10_padded, ne11 * ne12 * ne13, /*logical_n1=*/1, + gate->nb[1] / sizeof(float), gate->nb[1] / sizeof(float), + up->nb[1] / sizeof(float), up->nb[1] / sizeof(float), stream); } else { quantize_mmq_q8_1_cuda(src1_d, nullptr, src1_q8_1.get(), src0->type, ne10, s11, s12, s13, ne10_padded, ne11, ne12, ne13, stream); @@ -197,8 +375,11 @@ void ggml_cuda_mul_mat_q( const int si1 = ids->nb[1] / ggml_element_size(ids); const int sis1 = nb12 / nb11; - ggml_cuda_launch_mm_ids_helper((const int32_t *) ids->data, ids_src1.get(), ids_dst.get(), expert_bounds.get(), - ne02, ne12, n_expert_used, ne11, si1, sis1, /*write_inverse =*/ dedup_bcast, stream); + if (!ggml_cuda_launch_mm_ids_bounded(ctx, (const int32_t *) ids->data, ids_src1.get(), ids_dst.get(), expert_bounds.get(), + ne02, ne12, n_expert_used, ne11, si1, sis1, dedup_bcast, stream)) { + ggml_cuda_launch_mm_ids_helper((const int32_t *) ids->data, ids_src1.get(), ids_dst.get(), expert_bounds.get(), + ne02, ne12, n_expert_used, ne11, si1, sis1, /*write_inverse =*/ dedup_bcast, stream); + } CUDA_CHECK(cudaGetLastError()); } @@ -229,6 +410,12 @@ void ggml_cuda_mul_mat_q( quantize_mmq_fp4_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src1_scale.ptr, src0->type, use_aligned_float8, ne10, s11, s12, s13, ne10_padded, ne11_flat, ne12_flat, ne13_flat, stream); } + } else if (swiglu) { + quantize_mmq_q8_1_swiglu_cuda( + (const float *) gate->data, (const float *) up->data, ids_src1.get(), src1_q8_1.get(), src0->type, + ne10, ne10_padded, ne11_flat, n_expert_used, + gate->nb[1] / sizeof(float), gate->nb[2] / sizeof(float), + up->nb[1] / sizeof(float), up->nb[2] / sizeof(float), stream); } else if (dedup_bcast) { quantize_scatter_mmq_q8_1_cuda(src1_d, ids_src1.get(), src1_q8_1.get(), src0->type, ne10, /*stride_token=*/s12, ne10_padded, ne12, ne11_flat, n_expert_used, stream); @@ -263,6 +450,17 @@ void ggml_cuda_mul_mat_q( ggml_cuda_mul_mat_q_switch_type(ctx, args, stream); } +void ggml_cuda_mul_mat_q( + ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst) { + ggml_cuda_mul_mat_q_impl(ctx, src0, src1, ids, dst, nullptr); +} + +void ggml_cuda_mul_mat_q_swiglu( + ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * ids, ggml_tensor * dst, const ggml_tensor * swiglu) { + ggml_cuda_mul_mat_q_impl(ctx, src0, swiglu, ids, dst, swiglu); +} + + bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t n_experts) { #ifdef GGML_CUDA_FORCE_CUBLAS return false; @@ -368,6 +566,11 @@ bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t case GGML_TYPE_Q2_K: return ne11 <= 128; case GGML_TYPE_Q6_K: + // RDNA 3.5 (gfx1151): dequantize + hipBLASLt loses to MMQ up to ne11 = 1024; at 2048 the two + // paths trade places per shape (K = 4096 favours hipBLASLt, K >= 12288 favours MMQ), see PR. + if (GGML_CUDA_CC_IS_RDNA3_5(cc)) { + return ne11 <= 1024; + } return ne11 <= (GGML_CUDA_CC_IS_RDNA3_0(cc) ? 128 : 256); case GGML_TYPE_IQ2_XS: case GGML_TYPE_IQ2_S: diff --git a/ggml/src/ggml-cuda/mmq.cuh b/ggml/src/ggml-cuda/mmq.cuh index 6923f3510c03..333dac85dcfb 100644 --- a/ggml/src/ggml-cuda/mmq.cuh +++ b/ggml/src/ggml-cuda/mmq.cuh @@ -489,18 +489,24 @@ static __device__ __forceinline__ void ggml_cuda_mmq_write_back_mma( constexpr int rows_per_warp = ggml_cuda_mmq_get_rows_per_warp(type, J, fallback); constexpr int ntx = rows_per_warp/tile_C::I; // Number of x minitiles per warp. + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / ggml_cuda_get_physical_warp_size(); + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); + constexpr bool split_j = type == GGML_TYPE_Q8_0 && J == 128 && !fallback && I == 64 && nwarps == 8; + constexpr int j_group = split_j ? J/2 : J; - const int i0 = (threadIdx.y / ntx) * (ntx*tile_C::I); + const int warp_i = split_j ? threadIdx.y % 4 : threadIdx.y; + const int warp_j = split_j ? threadIdx.y / 4 : 0; + const int i0 = (warp_i / ntx) * (ntx*tile_C::I); const bool y_scale_used = y_scale != nullptr; #pragma unroll - for (int j0 = 0; j0 < J; j0 += ntx*tile_C::J) { + for (int j0 = 0; j0 < j_group; 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); + const int j = warp_j*j_group + j0 + (warp_i % ntx)*tile_C::J + tile_C::get_j(l); if (j > j_max) { continue; @@ -871,6 +877,22 @@ static constexpr __device__ ggml_cuda_mmq_write_back_t ggml_cuda_mmq_get_write_b // --------------------------------------------------------------------------------------------- +// Software prefetch of the next K iteration into registers, see mul_mat_q_process_tile. +template +static constexpr __host__ __device__ bool ggml_cuda_mmq_use_prefetch() { +#if defined(RDNA3_5) && defined(AMD_WMMA_AVAILABLE) + // Whitelist: the extra registers cause spills in several other specializations. + return (type == GGML_TYPE_Q8_0 && (J == 48 || J == 128) && !fallback) || + (type == GGML_TYPE_Q6_K && J == 32) || + (type == GGML_TYPE_Q5_K && J == 32) || + (type == GGML_TYPE_Q4_K && J == 48) || + (type == GGML_TYPE_IQ2_S && J == 128) || + (type == GGML_TYPE_IQ3_XXS && J == 128); +#else + return false; +#endif // defined(RDNA3_5) && defined(AMD_WMMA_AVAILABLE) +} + template static __device__ __forceinline__ void mul_mat_q_process_tile( const char * __restrict__ x, const int offset_x, const int * __restrict__ y, @@ -905,6 +927,86 @@ static __device__ __forceinline__ void mul_mat_q_process_tile( constexpr int sz = sizeof(block_q8_1_mmq) / sizeof(int); + if constexpr (ggml_cuda_mmq_use_prefetch()) { + // RDNA3.5 software prefetch: with 64 KiB of LDS per CU only 1-2 blocks are resident, so the + // global load latency of every K iteration is exposed. Stage the next iteration's tiles in + // registers while the current one is consumed. The arithmetic is unchanged. + constexpr int nthreads = nwarps*warp_size; + constexpr int y_regs_n = (J*MMQ_TILE_Y_K + nthreads - 1) / nthreads; + constexpr bool x_regs_ok = type == GGML_TYPE_Q8_0; + const int tid = threadIdx.y*warp_size + threadIdx.x; + + int yr0[y_regs_n]; + int yr1[y_regs_n]; + ggml_cuda_mmq_x_regs_q8_0 xr; + GGML_UNUSED(xr); + + auto load_y_regs = [&](int (&yr)[y_regs_n], const int kb0, const int half) { + const int * by0 = y + ncols_y * ((kb0 * qk / ne_block) * sz + half*sz); +#pragma unroll + for (int r = 0; r < y_regs_n; ++r) { + yr[r] = by0[r*nthreads + tid]; + } + }; + auto store_y_regs = [&](const int (&yr)[y_regs_n]) { +#pragma unroll + for (int r = 0; r < y_regs_n; ++r) { + tile_y[r*nthreads + tid] = yr[r]; + } + }; + + if (kb0_start < kb0_stop) { + if constexpr (x_regs_ok) { + ggml_cuda_mmq_load_tiles_q8_0_regs(x, xr, offset_x + kb0_start, tile_x_max_i, stride_row_x); + } + load_y_regs(yr0, kb0_start, 0); + load_y_regs(yr1, kb0_start, 1); + } + + for (int kb0 = kb0_start; kb0 < kb0_stop; kb0 += blocks_per_iter) { + const bool has_next = kb0 + blocks_per_iter < kb0_stop; + + if constexpr (x_regs_ok) { + ggml_cuda_mmq_store_tiles_q8_0_regs(xr, tile_x, tile_x_max_i); + } else { + load_tiles(x, tile_x, offset_x + kb0, tile_x_max_i, stride_row_x); + } + store_y_regs(yr0); + + __syncthreads(); + + if (has_next) { + if constexpr (x_regs_ok) { + ggml_cuda_mmq_load_tiles_q8_0_regs(x, xr, offset_x + kb0 + blocks_per_iter, tile_x_max_i, stride_row_x); + } + load_y_regs(yr0, kb0 + blocks_per_iter, 0); + } + + vec_dot(tile_x, tile_y, sum, 0); + + __syncthreads(); + + store_y_regs(yr1); + + __syncthreads(); + + if (has_next) { + load_y_regs(yr1, kb0 + blocks_per_iter, 1); + } + + vec_dot(tile_x, tile_y, sum, MMQ_TILE_NE_K); + + __syncthreads(); + } + + if (fixup) { + write_back(sum, ids_dst, tmp_fixup + blockIdx.x*(J*I), y_scale, I, I, J); + } else { + write_back(sum, ids_dst, dst, y_scale, stride_col_dst, tile_x_max_i, tile_y_max_j); + } + return; + } + 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); { @@ -974,21 +1076,7 @@ static __global__ void mul_mat_q( const uint32_t nty = (nrows_x + I - 1) / I; // 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 < J; j0 += nwarps*warp_size) { - const int j = j0 + threadIdx.y*warp_size + threadIdx.x; - - if (j0 + nwarps*warp_size > J && j >= J) { - break; - } - - ids_dst_shared[j] = j; - } - __syncthreads(); if constexpr (!ggml_cuda_mmq_get_stream_k(type, J, fallback)) { const uint2 tmp2 = fast_div_modulo(blockIdx.z, nchannels_y); @@ -1037,6 +1125,18 @@ static __global__ void mul_mat_q( ids_dst_shared[j] = ids_dst[col_low + jt*J + j]; } __syncthreads(); + } else { +#pragma unroll + for (int j0 = 0; j0 < J; j0 += nwarps*warp_size) { + const int j = j0 + threadIdx.y*warp_size + threadIdx.x; + + if (j0 + nwarps*warp_size > J && j >= J) { + break; + } + + ids_dst_shared[j] = j; + } + __syncthreads(); } offset_y += (col_low + jt*J)*(sizeof(block_q8_1_mmq)/sizeof(int)); @@ -1060,6 +1160,18 @@ static __global__ void mul_mat_q( return; } +#pragma unroll + for (int j0 = 0; j0 < J; j0 += nwarps*warp_size) { + const int j = j0 + threadIdx.y*warp_size + threadIdx.x; + + if (j0 + nwarps*warp_size > J && j >= J) { + break; + } + + ids_dst_shared[j] = j; + } + __syncthreads(); + constexpr int ITER_K = ggml_cuda_mmq_get_K_vram(type, J, fallback); constexpr int blocks_per_iter = ITER_K / qk; @@ -1237,6 +1349,101 @@ static __global__ void mul_mat_q( tile_x_max_i, tile_y_max_j, kb0_start, kb0_stop); } +// routed-compact MoE tiles: one descriptor per real (expert, tile) pair; the builder runs as a single block with one +// thread per expert +#define MMQ_ROUTED_MAX_EXPERTS 1024 +#ifndef MMQ_IQ_ID_J_MID +#define MMQ_IQ_ID_J_MID 64 +#endif + +static constexpr bool mmq_rdna3_5_id_n_experts_ok(const int64_t n_experts) { + return n_experts >= 8 && n_experts <= MMQ_ROUTED_MAX_EXPERTS; +} + +template +static __global__ void build_mmq_routed_descriptors( + const int32_t * expert_bounds, uint32_t * descriptors, int n_experts, int max_descriptors) { + __shared__ int tile_counts[MMQ_ROUTED_MAX_EXPERTS]; + const int expert = threadIdx.x; + + for (int i = expert; i < max_descriptors; i += blockDim.x) { + descriptors[i] = UINT32_MAX; + } + + if (expert < n_experts) { + const int count = expert_bounds[expert + 1] - expert_bounds[expert]; + tile_counts[expert] = (count + J - 1) / J; + } + __syncthreads(); + + if (expert < n_experts) { + int descriptor = 0; + for (int i = 0; i < expert; ++i) { + descriptor += tile_counts[i]; + } + for (int jt = 0; jt < tile_counts[expert]; ++jt) { + descriptors[descriptor + jt] = uint32_t(expert) | (uint32_t(jt) << 16); + } + } +} + +template +__launch_bounds__(ggml_cuda_mmq_get_nthreads(type, J, fallback), ggml_cuda_mmq_get_occupancy(type, J, fallback)) +static __global__ void mul_mat_q_routed_compact( + const char * __restrict__ x, const int * __restrict__ y, const int32_t * __restrict__ ids_dst, + const int32_t * __restrict__ expert_bounds, const uint32_t * __restrict__ descriptors, + const int max_descriptors, float * __restrict__ dst, const float * __restrict__ y_scale, + const uint3 blocks_per_ne00, const int nrows_x, const int stride_row_x, const int ncols_y, const int stride_col_dst, + const uint3 channel_ratio, const int stride_channel_x, + const uint3 sample_ratio, const int stride_sample_x) { + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = ggml_cuda_mmq_get_nthreads(type, J, fallback) / warp_size; + constexpr int I = ggml_cuda_mmq_get_I(type, J, fallback); + + const int wt = blockIdx.z; + const int it = blockIdx.x; + + extern __shared__ int ids_dst_shared[]; + for (int descriptor_idx = blockIdx.y; descriptor_idx < max_descriptors; descriptor_idx += gridDim.y) { + const uint32_t descriptor = descriptors[descriptor_idx]; + if (descriptor == UINT32_MAX) { + return; + } + + const int zt = descriptor & 0xFFFF; + const int jt = descriptor >> 16; + const int col_low = expert_bounds[zt + 0]; + const int col_high = expert_bounds[zt + 1]; + const int col_diff = col_high - col_low; + + if (descriptor_idx != int(blockIdx.y)) { + __syncthreads(); + } +#pragma unroll + for (int j0 = 0; j0 < J; j0 += nwarps * warp_size) { + const int j = j0 + threadIdx.y * warp_size + threadIdx.x; + if (j0 + nwarps * warp_size > J && j >= J) { + break; + } + ids_dst_shared[j] = jt * J + j < col_diff ? ids_dst[col_low + jt * J + j] : 0; + } + __syncthreads(); + + const int offset_x = fastdiv(wt, sample_ratio) * stride_sample_x + + fastdiv(zt, channel_ratio) * stride_channel_x + it * I * stride_row_x; + const int offset_y = (col_low + jt * J) * (sizeof(block_q8_1_mmq) / sizeof(int)); + const int tile_x_max_i = nrows_x - it * I - 1; + const int tile_y_max_j = col_diff - jt * J - 1; + const float * y_scale_tile = y_scale ? y_scale + col_low + jt * J : nullptr; + + constexpr bool fixup = false; + mul_mat_q_process_tile + (x, offset_x, y + offset_y, ids_dst_shared, dst + it * I, nullptr, y_scale_tile, + stride_row_x, ncols_y, stride_col_dst, + tile_x_max_i, tile_y_max_j, 0, blocks_per_ne00.z); + } +} + template __launch_bounds__(ggml_cuda_mmq_get_nthreads(type, J, fallback)/2, 1) static __global__ void mul_mat_q_stream_k_fixup( @@ -1392,6 +1599,82 @@ static size_t mmq_get_nbytes_shared(const ggml_cuda_mmq_config & config, const i return nbs_ids + nbs_x + GGML_PAD(nbs_y, config.nthreads*sizeof(int)); } +// RDNA3.5 tile selection for MoE (mul_mat_id) with 256 experts. +// The per-expert row count is n_tokens*n_expert_used/256 (e.g. 16/32/64 for ubatch 512/1024/2048 with 8 active +// experts), but the generic selection below only sees ncols_max == n_tokens*n_expert_used and picks J=128, so most +// of every tile is padding and the grid contains ncols_max/J*n_experts mostly-empty blocks. Use a J matched to the +// row count and the routed-compact kernel, which only enumerates real per-expert tiles. +// J per type/row range was measured on gfx1151 for the 2048x512 and 3072x1024 expert shapes. +static constexpr int mmq_rdna3_5_id_get_J(const ggml_type type, const int64_t rows_per_expert) { + switch (type) { + case GGML_TYPE_Q8_0: + return rows_per_expert <= 12 ? 16 : rows_per_expert <= 64 ? 48 : 128; + // the IQ expert types of the Unsloth qwen4exp mixes (512 experts, 10 active: 40 rows per expert at ubatch 2048) + case GGML_TYPE_IQ3_S: + case GGML_TYPE_IQ4_NL: + case GGML_TYPE_IQ4_XS: + return rows_per_expert <= 12 ? 16 : rows_per_expert <= 32 ? 48 : rows_per_expert <= 64 ? MMQ_IQ_ID_J_MID : 128; + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + return rows_per_expert <= 12 ? 16 : rows_per_expert <= 24 ? 32 : rows_per_expert <= 48 ? 48 : + rows_per_expert <= 64 ? 32 : 128; + default: + return 0; + } +} + +static constexpr bool mmq_rdna3_5_id_use_compact(const ggml_type type, const int J) { + switch (type) { + case GGML_TYPE_Q8_0: + return J == 16 || J == 48 || J == 128; + case GGML_TYPE_IQ3_S: + case GGML_TYPE_IQ4_NL: + case GGML_TYPE_IQ4_XS: + return J == 16 || J == 48 || J == MMQ_IQ_ID_J_MID || J == 128; + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + return J == 16 || J == 32 || J == 48 || J == 64 || J == 96 || J == 128; + default: + return false; + } +} + +static constexpr bool mmq_rdna3_5_id_use_j48_128e(const ggml_type type, const int64_t n_experts, const int64_t rows_per_expert) { + return n_experts == 128 && rows_per_expert >= 96 && rows_per_expert <= 160 && + (type == GGML_TYPE_Q4_K || type == GGML_TYPE_Q5_K || type == GGML_TYPE_Q6_K || type == GGML_TYPE_Q8_0); +} + +static_assert(mmq_rdna3_5_id_get_J(GGML_TYPE_Q8_0, 4) == 16); +static_assert(mmq_rdna3_5_id_get_J(GGML_TYPE_Q8_0, 16) == 48); +static_assert(mmq_rdna3_5_id_get_J(GGML_TYPE_Q8_0, 64) == 48); +static_assert(mmq_rdna3_5_id_get_J(GGML_TYPE_Q8_0, 128) == 128); +static_assert(mmq_rdna3_5_id_get_J(GGML_TYPE_Q6_K, 16) == 32); +static_assert(mmq_rdna3_5_id_get_J(GGML_TYPE_Q4_K, 32) == 48); +static_assert(mmq_rdna3_5_id_get_J(GGML_TYPE_Q5_K, 64) == 32); +static_assert(mmq_rdna3_5_id_get_J(GGML_TYPE_Q4_0, 16) == 0); +static_assert(mmq_rdna3_5_id_use_compact(GGML_TYPE_Q8_0, 48)); +static_assert(!mmq_rdna3_5_id_use_compact(GGML_TYPE_Q8_0, 32)); +static_assert(mmq_rdna3_5_id_use_compact(GGML_TYPE_Q6_K, 32)); +static_assert(!mmq_rdna3_5_id_use_compact(GGML_TYPE_IQ2_XS, 48)); +static_assert(!mmq_rdna3_5_id_use_j48_128e(GGML_TYPE_Q4_K, 128, 95) && + !mmq_rdna3_5_id_use_j48_128e(GGML_TYPE_Q5_K, 128, 95) && + !mmq_rdna3_5_id_use_j48_128e(GGML_TYPE_Q6_K, 128, 95) && + !mmq_rdna3_5_id_use_j48_128e(GGML_TYPE_Q8_0, 128, 95)); +static_assert(mmq_rdna3_5_id_use_j48_128e(GGML_TYPE_Q4_K, 128, 96) && + mmq_rdna3_5_id_use_j48_128e(GGML_TYPE_Q5_K, 128, 96) && + mmq_rdna3_5_id_use_j48_128e(GGML_TYPE_Q6_K, 128, 96) && + mmq_rdna3_5_id_use_j48_128e(GGML_TYPE_Q8_0, 128, 96) && + mmq_rdna3_5_id_use_j48_128e(GGML_TYPE_Q4_K, 128, 160) && + mmq_rdna3_5_id_use_j48_128e(GGML_TYPE_Q5_K, 128, 160) && + mmq_rdna3_5_id_use_j48_128e(GGML_TYPE_Q6_K, 128, 160) && + mmq_rdna3_5_id_use_j48_128e(GGML_TYPE_Q8_0, 128, 160)); +static_assert(!mmq_rdna3_5_id_use_j48_128e(GGML_TYPE_Q4_K, 128, 161) && + !mmq_rdna3_5_id_use_j48_128e(GGML_TYPE_Q5_K, 128, 161) && + !mmq_rdna3_5_id_use_j48_128e(GGML_TYPE_Q6_K, 128, 161) && + !mmq_rdna3_5_id_use_j48_128e(GGML_TYPE_Q8_0, 128, 161)); + template static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) { const int id = ggml_cuda_get_device(); @@ -1426,6 +1709,23 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a const uint3 channel_ratio_fd = init_fastdiv_values(channel_ratio); const uint3 sample_ratio_fd = init_fastdiv_values(sample_ratio); + const bool use_compact_routed = args.ids_dst != nullptr && mmq_rdna3_5_id_n_experts_ok(args.nchannels_y) && GGML_CUDA_CC_IS_RDNA3_5(cc) && + !ggml_cuda_mmq_get_stream_k(type, J, fallback, cc) && mmq_rdna3_5_id_use_compact(type, J); + if (use_compact_routed) { + const int max_descriptors = (args.ncols_dst + J - 1) / J + args.nchannels_y; + ggml_cuda_pool_alloc descriptors(ctx.pool(id), max_descriptors); + build_mmq_routed_descriptors<<<1, MMQ_ROUTED_MAX_EXPERTS, 0, stream>>> + (args.expert_bounds, descriptors.get(), args.nchannels_y, max_descriptors); + + const int descriptor_blocks = (args.ncols_dst + J - 1) / J; + const dim3 block_nums_compact(nty, descriptor_blocks, args.nsamples_y); + mul_mat_q_routed_compact<<>> + (args.x, args.y, args.ids_dst, args.expert_bounds, descriptors.get(), max_descriptors, args.dst, args.y_scale, + blocks_per_ne00_fd, args.nrows_x, args.stride_row_x, args.ncols_y, args.nrows_dst, + channel_ratio_fd, args.stride_channel_x, sample_ratio_fd, args.stride_sample_x); + return; + } + if (!ggml_cuda_mmq_get_stream_k(type, J, fallback, cc)) { mul_mat_q<<>> (args.x, args.y, args.ids_dst, args.expert_bounds, args.dst, nullptr, args.y_scale, @@ -1474,6 +1774,15 @@ static void launch_mul_mat_q(ggml_backend_cuda_context & ctx, const mmq_args & a ntx_fd); } +static constexpr bool mmq_use_rdna3_5_iq_id_j48( + ggml_type type, int cc, bool fallback, bool has_ids, int64_t ncols_dst, int64_t nchannels_y) { + return (type == GGML_TYPE_IQ2_XS || type == GGML_TYPE_IQ3_XXS) && GGML_CUDA_CC_IS_RDNA3_5(cc) && + !fallback && has_ids && nchannels_y > 0 && (ncols_dst + nchannels_y - 1) / nchannels_y == 48; +} + +static_assert(mmq_use_rdna3_5_iq_id_j48(GGML_TYPE_IQ2_XS, GGML_CUDA_CC_RDNA3_5, false, true, 12288, 256)); +static_assert(mmq_use_rdna3_5_iq_id_j48(GGML_TYPE_IQ3_XXS, GGML_CUDA_CC_RDNA3_5, false, true, 12288, 256)); + template void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args, cudaStream_t stream) { const int id = ggml_cuda_get_device(); @@ -1483,6 +1792,31 @@ void mul_mat_q_switch_J(ggml_backend_cuda_context & ctx, const mmq_args & args, int J_best = 0; int ntiles_J_best = INT_MAX; + if (mmq_use_rdna3_5_iq_id_j48(type, cc, fallback, args.ids_dst != nullptr, args.ncols_dst, args.nchannels_y)) { + constexpr int J_rdna3_5 = 48; + const ggml_cuda_mmq_config config = ggml_cuda_mmq_get_config(type, J_rdna3_5, fallback, cc); + if (config.type != GGML_TYPE_COUNT && mmq_get_nbytes_shared(config, cc) <= smpbo) { + J_best = J_rdna3_5; + ntiles_J_best = 1; + } + } else if (GGML_CUDA_CC_IS_RDNA3_5(cc) && !fallback && args.ids_dst != nullptr && mmq_rdna3_5_id_n_experts_ok(args.nchannels_y)) { + const int64_t rows_per_expert = (args.ncols_dst + args.nchannels_y - 1) / args.nchannels_y; + const bool use_j48_128e = mmq_rdna3_5_id_use_j48_128e(type, args.nchannels_y, rows_per_expert); + int J_rdna3_5 = use_j48_128e ? 48 : mmq_rdna3_5_id_get_J(type, rows_per_expert); + if constexpr (type == GGML_TYPE_Q6_K) { + if (getenv("GGML_Q6_COMPACT_J")) { + J_rdna3_5 = atoi(getenv("GGML_Q6_COMPACT_J")); + } + } + if (J_rdna3_5 != 0) { + const ggml_cuda_mmq_config config = ggml_cuda_mmq_get_config(type, J_rdna3_5, fallback, cc); + if (config.type != GGML_TYPE_COUNT && mmq_get_nbytes_shared(config, cc) <= smpbo) { + J_best = J_rdna3_5; + ntiles_J_best = 1; + } + } + } + for (int J = 8; J <= 128 && ntiles_J_best > 1; J += 8) { const ggml_cuda_mmq_config config = ggml_cuda_mmq_get_config(type, J, fallback, cc); if (config.type == GGML_TYPE_COUNT) { @@ -1602,4 +1936,9 @@ extern DECL_MMQ_CASE(GGML_TYPE_NVFP4); void ggml_cuda_mul_mat_q( ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst); +void ggml_cuda_mul_mat_q_swiglu( + ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * ids, ggml_tensor * dst, const ggml_tensor * swiglu); + +void ggml_cuda_mul_mat_q_pair(ggml_backend_cuda_context & ctx, ggml_tensor * dst0, ggml_tensor * dst1); + bool ggml_cuda_should_use_mmq(enum ggml_type type, int cc, int64_t ne11, int64_t n_experts); diff --git a/ggml/src/ggml-cuda/mmvf.cu b/ggml/src/ggml-cuda/mmvf.cu index bd5c5d421a4c..e06920980c8d 100644 --- a/ggml/src/ggml-cuda/mmvf.cu +++ b/ggml/src/ggml-cuda/mmvf.cu @@ -4,7 +4,7 @@ #include "mmvf.cuh" #include "convert.cuh" -template +template static __global__ void mul_mat_vec_f( const T * x_ptr, const float * y_ptr, const int32_t * ids_ptr, const ggml_cuda_mm_fusion_args_device fusion, float * dst_ptr, const int ncols2, const uint3 nchannels_y, const int stride_row, const int stride_col_y2, const int stride_col_dst, @@ -27,11 +27,11 @@ static __global__ void mul_mat_vec_f( ggml_cuda_pdl_sync(); if constexpr (is_multi_token_id) { - // Multi-token MUL_MAT_ID path, adding these in the normal path causes a perf regression for n_tokens=1 case - token_idx = blockIdx.z; - channel_x = ids[channel_dst + token_idx * ids_stride]; - channel_y = fastmodulo(channel_dst, nchannels_y); - sample_dst = 0; + // Multi-token path, adding these in the normal path causes a perf regression for n_tokens=1 case + token_idx = tokens_in_block ? threadIdx.y : blockIdx.z; + channel_x = ids ? ids[channel_dst + token_idx * ids_stride] : fastdiv((uint32_t) channel_dst, channel_ratio); + channel_y = ids ? fastmodulo(channel_dst, nchannels_y) : channel_dst; + sample_dst = ids ? 0 : (tokens_in_block ? blockIdx.z : 0); } else { token_idx = ids ? blockIdx.z : 0; channel_x = ids ? ids[blockIdx.y + token_idx * ids_stride] : fastdiv((uint32_t) channel_dst, channel_ratio); @@ -99,7 +99,7 @@ static __global__ void mul_mat_vec_f( const float2 * y2 = (const float2 *) y; extern __shared__ char data_mmv[]; - float * buf_iw = (float *) data_mmv; + float * buf_iw = (float *) data_mmv + (tokens_in_block ? threadIdx.y * warp_size : 0); [[maybe_unused]] float * buf_iw_gate = nullptr; if constexpr (has_fusion) { buf_iw_gate = (float *) (data_mmv + warp_size*sizeof(float)); @@ -135,7 +135,35 @@ static __global__ void mul_mat_vec_f( } } - for (int col2 = tid; col2 < ncols2; col2 += block_size) { + int col2 = tid; +#if defined(GGML_USE_HIP) + // Issue the weight loads of four K iterations before consuming them: the runtime trip count keeps the + // compiler from pipelining the loop, and with one 8-byte load in flight per lane the single-column + // f32 matvecs (e.g. MoE routers) reach barely half the memory bandwidth. The accumulation order is + // unchanged, so the results are bit-identical. + if constexpr (ncols_dst == 1 && !has_fusion) { + constexpr int pf = 4; + for (; col2 + (pf - 1)*block_size < ncols2; col2 += pf*block_size) { + float2 tmpx[pf]; + float2 tmpy[pf]; +#pragma unroll + for (int i = 0; i < pf; ++i) { + tmpx[i] = x2[col2 + i*block_size]; + } +#pragma unroll + for (int i = 0; i < pf; ++i) { + tmpy[i] = y2[col2 + i*block_size]; + } +#pragma unroll + for (int i = 0; i < pf; ++i) { + ggml_cuda_mad(sumf[0], tmpx[i].x, tmpy[i].x); + ggml_cuda_mad(sumf[0], tmpx[i].y, tmpy[i].y); + } + } + } +#endif // defined(GGML_USE_HIP) + + for (; col2 < ncols2; col2 += block_size) { const float2 tmpx = x2[col2]; float2 tmpx_gate = make_float2(0.0f, 0.0f); if constexpr (has_fusion) { @@ -383,7 +411,80 @@ static __global__ void mul_mat_vec_f( } } -template + +#if defined(GGML_USE_HIP) +// One wave per output row emulating mul_mat_vec_f: lane l accumulates the partial +// sums of the eight virtual warps' lane l (columns 2*(32*w + l) + 256*i) in the same fma order, reduces each +// with the same butterfly and combines the eight partials exactly like block_reduce does over lanes 0..7 (the +// other lanes hold zero). The standard kernel launches 8 waves per row that each issue a single 4-byte load, +// which on gfx1151 leaves K=512 MoE down projections at ~145 GB/s; here each lane keeps 8 loads in flight +// (measured: 93 -> 67 us for 8 x 2048 rows of K=512; slower than the standard kernel for K=2048). +template +__launch_bounds__(32 * rows_per_block, 1) +static __global__ void mul_mat_vec_bf16_wave( + const nv_bfloat16 * x_ptr, const float * y_ptr, const int32_t * ids_ptr, float * dst_ptr, + const int ncols2, const int nrows, const uint3 nchannels_y, const int stride_row, const int stride_col_dst, + const uint3 channel_ratio, const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst, + const uint3 sample_ratio, const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst) { + constexpr int warp_size = 32; + constexpr int nvw = 256 / warp_size; // virtual warps of the emulated 256-thread block + const int lane = threadIdx.x; + const int row = blockIdx.x * rows_per_block + threadIdx.y; + if (row >= nrows) { + return; + } + const int channel_dst = blockIdx.y; + const int channel_x = ids_ptr ? ids_ptr[channel_dst] : (int) fastdiv((uint32_t) channel_dst, channel_ratio); + const int channel_y = ids_ptr ? (int) fastmodulo(channel_dst, nchannels_y) : channel_dst; + const int sample_dst = 0; + const int sample_x = fastdiv((uint32_t) sample_dst, sample_ratio); + const int sample_y = sample_dst; + + const int * x2 = (const int *) (x_ptr + int64_t(sample_x)*stride_sample_x + channel_x*stride_channel_x + int64_t(row)*stride_row); + const float2 * y2 = (const float2 *) (y_ptr + int64_t(sample_y)*stride_sample_y + channel_y*stride_channel_y); + float * dst = dst_ptr + int64_t(sample_dst)*stride_sample_dst + channel_dst*stride_channel_dst; + + float sumf[nvw]; +#pragma unroll + for (int w = 0; w < nvw; ++w) { + sumf[w] = 0.0f; + } + for (int col0 = 0; col0 < ncols2; col0 += 256) { + int tmpx[nvw]; + float2 tmpy[nvw]; +#pragma unroll + for (int w = 0; w < nvw; ++w) { + const int col2 = col0 + w*warp_size + lane; + if (col2 < ncols2) { + tmpx[w] = x2[col2]; + tmpy[w] = y2[col2]; + } + } +#pragma unroll + for (int w = 0; w < nvw; ++w) { + const int col2 = col0 + w*warp_size + lane; + if (col2 < ncols2) { + const float tmpx0 = ggml_cuda_cast(reinterpret_cast(&tmpx[w])[0]); + const float tmpx1 = ggml_cuda_cast(reinterpret_cast(&tmpx[w])[1]); + ggml_cuda_mad(sumf[w], tmpx0, tmpy[w].x); + ggml_cuda_mad(sumf[w], tmpx1, tmpy[w].y); + } + } + } +#pragma unroll + for (int w = 0; w < nvw; ++w) { + // per-warp butterfly, then the two zero-lane steps (offsets 16 and 8) of the second-level reduction + sumf[w] = (warp_reduce_sum(sumf[w]) + 0.0f) + 0.0f; + } + const float value = ((sumf[0] + sumf[4]) + (sumf[2] + sumf[6])) + ((sumf[1] + sumf[5]) + (sumf[3] + sumf[7])); + if (lane == 0) { + dst[row] = value; + } + GGML_UNUSED(stride_col_dst); +} +#endif // defined(GGML_USE_HIP) + +template static void mul_mat_vec_f_switch_fusion( const T * x, const float * y, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst, const int64_t ncols, const uint3 nchannels_y, @@ -397,7 +498,7 @@ static void mul_mat_vec_f_switch_fusion( const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr; if constexpr (ncols_dst == 1) { if (has_fusion) { - ggml_cuda_kernel_launch(mul_mat_vec_f, launch_params, + ggml_cuda_kernel_launch(mul_mat_vec_f, launch_params, x, y, ids, fusion, dst, ncols, nchannels_y, stride_row, stride_col_y, stride_col_dst, channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride); @@ -407,14 +508,14 @@ static void mul_mat_vec_f_switch_fusion( GGML_ASSERT(!has_fusion && "fusion only supported for ncols_dst=1"); - ggml_cuda_kernel_launch(mul_mat_vec_f, launch_params, + ggml_cuda_kernel_launch(mul_mat_vec_f, launch_params, x, y, ids, fusion, dst, ncols, nchannels_y, stride_row, stride_col_y, stride_col_dst, channel_ratio, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride); } -template +template void launch_mul_mat_vec_f_cuda( const T * x, const float * y, const int32_t * ids, const ggml_cuda_mm_fusion_args_device fusion, float * dst, const int64_t ncols, const int64_t nrows, @@ -448,57 +549,74 @@ void launch_mul_mat_vec_f_cuda( block_size_best = block_size; } } - const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr; - const int nbytes_shared = warp_size*sizeof(float) + (has_fusion ? warp_size*sizeof(float) : 0); - const dim3 block_nums(nrows, nchannels_dst, nsamples_or_ntokens); - const dim3 block_dims(block_size_best, 1, 1); +#if defined(GGML_USE_HIP) + if constexpr (std::is_same_v && ncols_dst == 1 && !is_multi_token_id && !tokens_in_block) { + // only where the emulated block does a single K iteration: for longer rows the 8x lower wave count loses + if (!has_fusion && block_size_best == 256 && ncols <= 512 && warp_size == 32 && nsamples_or_ntokens == 1 && nsamples_dst == 1 && + GGML_CUDA_CC_IS_RDNA(ggml_cuda_info().devices[device].cc)) { + constexpr int rows_per_block = 4; + const dim3 block_nums((nrows + rows_per_block - 1) / rows_per_block, nchannels_dst, 1); + const dim3 block_dims(warp_size, rows_per_block, 1); + const ggml_cuda_kernel_launch_params launch_params = {block_nums, block_dims, 0, stream}; + ggml_cuda_kernel_launch(mul_mat_vec_bf16_wave, launch_params, + x, y, ids, dst, (int) (ncols/2), (int) nrows, nchannels_y_fd, (int) stride_row, (int) stride_col_dst, + channel_ratio_fd, (int) stride_channel_x, (int) stride_channel_y, (int) stride_channel_dst, + sample_ratio_fd, (int) stride_sample_x, (int) stride_sample_y, (int) stride_sample_dst); + return; + } + } +#endif // defined(GGML_USE_HIP) + + const int nbytes_shared = (warp_size*sizeof(float) + (has_fusion ? warp_size*sizeof(float) : 0)) * (tokens_in_block ? nsamples_or_ntokens : 1); + const dim3 block_nums(nrows, nchannels_dst, tokens_in_block ? nsamples_dst : nsamples_or_ntokens); + const dim3 block_dims(block_size_best, tokens_in_block ? nsamples_or_ntokens : 1, 1); switch (block_size_best) { case 32: { - mul_mat_vec_f_switch_fusion + mul_mat_vec_f_switch_fusion (x, y, ids, fusion, dst, ncols/2, nchannels_y_fd, stride_row, stride_col_y/2, stride_col_dst, channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, block_dims, block_nums, nbytes_shared, ids_stride, stream); } break; case 64: { - mul_mat_vec_f_switch_fusion + mul_mat_vec_f_switch_fusion (x, y, ids, fusion, dst, ncols/2, nchannels_y_fd, stride_row, stride_col_y/2, stride_col_dst, channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, block_dims, block_nums, nbytes_shared, ids_stride, stream); } break; case 96: { - mul_mat_vec_f_switch_fusion + mul_mat_vec_f_switch_fusion (x, y, ids, fusion, dst, ncols/2, nchannels_y_fd, stride_row, stride_col_y/2, stride_col_dst, channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, block_dims, block_nums, nbytes_shared, ids_stride, stream); } break; case 128: { - mul_mat_vec_f_switch_fusion + mul_mat_vec_f_switch_fusion (x, y, ids, fusion, dst, ncols/2, nchannels_y_fd, stride_row, stride_col_y/2, stride_col_dst, channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, block_dims, block_nums, nbytes_shared, ids_stride, stream); } break; case 160: { - mul_mat_vec_f_switch_fusion + mul_mat_vec_f_switch_fusion (x, y, ids, fusion, dst, ncols/2, nchannels_y_fd, stride_row, stride_col_y/2, stride_col_dst, channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, block_dims, block_nums, nbytes_shared, ids_stride, stream); } break; case 192: { - mul_mat_vec_f_switch_fusion + mul_mat_vec_f_switch_fusion (x, y, ids, fusion, dst, ncols/2, nchannels_y_fd, stride_row, stride_col_y/2, stride_col_dst, channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, block_dims, block_nums, nbytes_shared, ids_stride, stream); } break; case 224: { - mul_mat_vec_f_switch_fusion + mul_mat_vec_f_switch_fusion (x, y, ids, fusion, dst, ncols/2, nchannels_y_fd, stride_row, stride_col_y/2, stride_col_dst, channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, block_dims, block_nums, nbytes_shared, ids_stride, stream); } break; case 256: { - mul_mat_vec_f_switch_fusion + mul_mat_vec_f_switch_fusion (x, y, ids, fusion, dst, ncols/2, nchannels_y_fd, stride_row, stride_col_y/2, stride_col_dst, channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, block_dims, block_nums, nbytes_shared, ids_stride, stream); @@ -517,7 +635,7 @@ static void mul_mat_vec_f_cuda_switch_ncols_dst( const int64_t nchannels_x, const int64_t nchannels_y, const int64_t nchannels_dst, const int64_t stride_channel_x, const int64_t stride_channel_y, const int64_t stride_channel_dst, const int64_t nsamples_x, const int64_t nsamples_dst, const int64_t stride_sample_x, const int64_t stride_sample_y, const int64_t stride_sample_dst, - const int64_t ids_stride, cudaStream_t stream) { + const int64_t ids_stride, bool exact_batch, cudaStream_t stream) { const bool has_ids = ids != nullptr; @@ -543,6 +661,42 @@ static void mul_mat_vec_f_cuda_switch_ncols_dst( return; } + if (exact_batch && !has_ids && ncols_dst > 1) { + if (ncols_dst == 2) { + launch_mul_mat_vec_f_cuda + (x, y, ids, fusion, dst, ncols, nrows, stride_row, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, + stride_channel_dst, nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, + nsamples_dst, ids_stride, stream); + return; + } + if (ncols_dst == 3) { + launch_mul_mat_vec_f_cuda + (x, y, ids, fusion, dst, ncols, nrows, stride_row, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, + stride_channel_dst, nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, + nsamples_dst, ids_stride, stream); + return; + } + if (ncols_dst == 4) { + launch_mul_mat_vec_f_cuda + (x, y, ids, fusion, dst, ncols, nrows, stride_row, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, + stride_channel_dst, nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, + nsamples_dst, ids_stride, stream); + return; + } + for (int64_t col = 0; col < ncols_dst; col += 4) { + const int64_t ncols_group = std::min(4, ncols_dst - col); + launch_mul_mat_vec_f_cuda + (x, y + col*stride_col_y, ids, fusion, dst + col*stride_col_dst, ncols, nrows, stride_row, stride_col_y, stride_col_dst, + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, + stride_channel_dst, nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, + ncols_group, ids_stride, stream); + } + return; + } + switch (ncols_dst) { case 1: launch_mul_mat_vec_f_cuda @@ -614,21 +768,21 @@ static void mul_mat_vec_f_cuda( const int64_t nchannels_x, const int64_t nchannels_y, const int64_t nchannels_dst, const int64_t stride_channel_x, const int64_t stride_channel_y, const int64_t stride_channel_dst, const int64_t nsamples_x, const int64_t nsamples_dst, const int64_t stride_sample_x, const int64_t stride_sample_y, const int64_t stride_sample_dst, - const int64_t ids_stride, enum ggml_prec prec, cudaStream_t stream) { + const int64_t ids_stride, enum ggml_prec prec, bool exact_batch, cudaStream_t stream) { if constexpr(std::is_same_v) { if (prec == GGML_PREC_DEFAULT) { mul_mat_vec_f_cuda_switch_ncols_dst (x, y, ids, fusion, dst, ncols, nrows, ncols_dst, stride_row, stride_col_y, stride_col_dst, nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, - stride_channel_dst, nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + stride_channel_dst, nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, exact_batch, stream); return; } } mul_mat_vec_f_cuda_switch_ncols_dst (x, y, ids, fusion, dst, ncols, nrows, ncols_dst, stride_row, stride_col_y, stride_col_dst, nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, - stride_channel_dst, nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, stream); + stride_channel_dst, nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, ids_stride, exact_batch, stream); } void ggml_cuda_mul_mat_vec_f(ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, @@ -653,6 +807,7 @@ void ggml_cuda_mul_mat_vec_f(ggml_backend_cuda_context & ctx, const ggml_tensor const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; const enum ggml_prec prec = fast_fp16_available(cc) ? ggml_prec(dst->op_params[0]) : GGML_PREC_F32; + const bool exact_batch = ggml_get_op_params_i32(dst, 1) == GGML_HINT_EXACT_BATCH; const float * src1_d = (const float *) src1->data; const int32_t * ids_d = ids ? (const int32_t *) ids->data : nullptr; @@ -709,19 +864,19 @@ void ggml_cuda_mul_mat_vec_f(ggml_backend_cuda_context & ctx, const ggml_tensor const float * src0_d = (const float *) src0->data; mul_mat_vec_f_cuda(src0_d, src1_d, ids_d, fusion_local, dst_d, ne00, ne01, ncols_dst, s01, stride_col_y, stride_col_dst, ne02, nchannels_y, nchannels_dst, s02, stride_channel_y, stride_channel_dst, - ne03, ne3, s03, s13, s3, ids_stride, prec, ctx.stream()); + ne03, ne3, s03, s13, s3, ids_stride, prec, exact_batch, ctx.stream()); } break; case GGML_TYPE_F16: { const half * src0_d = (const half *) src0->data; mul_mat_vec_f_cuda(src0_d, src1_d, ids_d, fusion_local, dst_d, ne00, ne01, ncols_dst, s01, stride_col_y, stride_col_dst, ne02, nchannels_y, nchannels_dst, s02, stride_channel_y, stride_channel_dst, - ne03, ne3, s03, s13, s3, ids_stride, prec, ctx.stream()); + ne03, ne3, s03, s13, s3, ids_stride, prec, exact_batch, ctx.stream()); } break; case GGML_TYPE_BF16: { const nv_bfloat16 * src0_d = (const nv_bfloat16 *) src0->data; mul_mat_vec_f_cuda(src0_d, src1_d, ids_d, fusion_local, dst_d, ne00, ne01, ncols_dst, s01, stride_col_y, stride_col_dst, ne02, nchannels_y, nchannels_dst, s02, stride_channel_y, stride_channel_dst, - ne03, ne3, s03, s13, s3, ids_stride, prec, ctx.stream()); + ne03, ne3, s03, s13, s3, ids_stride, prec, exact_batch, ctx.stream()); } break; default: GGML_ABORT("unsupported type: %s", ggml_type_name(src0->type)); @@ -768,19 +923,19 @@ void ggml_cuda_op_mul_mat_vec_f( const float * src0_d = (const float *) src0_dd_i; mul_mat_vec_f_cuda(src0_d, src1_ddf_i, nullptr, empty, dst_dd_i, ne00, row_diff, src1_ncols, stride_row, stride_col_y, stride_col_dst, nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, - nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, 0, prec, stream); + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, 0, prec, false, stream); } break; case GGML_TYPE_F16: { const half * src0_d = (const half *) src0_dd_i; mul_mat_vec_f_cuda(src0_d, src1_ddf_i, nullptr, empty, dst_dd_i, ne00, row_diff, src1_ncols, stride_row, stride_col_y, stride_col_dst, nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, - nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, 0, prec, stream); + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, 0, prec, false, stream); } break; case GGML_TYPE_BF16: { const nv_bfloat16 * src0_d = (const nv_bfloat16 *) src0_dd_i; mul_mat_vec_f_cuda(src0_d, src1_ddf_i, nullptr, empty, dst_dd_i, ne00, row_diff, src1_ncols, stride_row, stride_col_y, stride_col_dst, nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, - nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, 0, prec, stream); + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, 0, prec, false, stream); } break; default: GGML_ABORT("unsupported type: %s", ggml_type_name(src0->type)); diff --git a/ggml/src/ggml-cuda/mmvq.cu b/ggml/src/ggml-cuda/mmvq.cu index 6305230b1f92..101da1f5acee 100644 --- a/ggml/src/ggml-cuda/mmvq.cu +++ b/ggml/src/ggml-cuda/mmvq.cu @@ -1,4 +1,5 @@ #include "mmvq.cuh" +#include "ggml-backend-impl.h" #include "quantize.cuh" #include "unary.cuh" #include "vecdotq.cuh" @@ -98,6 +99,7 @@ enum mmvq_parameter_table_id { MMVQ_PARAMETERS_TURING, MMVQ_PARAMETERS_GCN, MMVQ_PARAMETERS_RDNA2, + MMVQ_PARAMETERS_RDNA3_5, MMVQ_PARAMETERS_RDNA3_0, MMVQ_PARAMETERS_RDNA4, MMVQ_PARAMETERS_GB10 @@ -106,9 +108,11 @@ enum mmvq_parameter_table_id { static constexpr __device__ mmvq_parameter_table_id get_device_table_id() { #if defined(RDNA4) return MMVQ_PARAMETERS_RDNA4; +#elif defined(RDNA3_5) + return MMVQ_PARAMETERS_RDNA3_5; #elif defined(RDNA3_0) return MMVQ_PARAMETERS_RDNA3_0; -#elif defined(RDNA2) || defined(RDNA3_5) +#elif defined(RDNA2) return MMVQ_PARAMETERS_RDNA2; #elif defined(GCN) || defined(CDNA) return MMVQ_PARAMETERS_GCN; @@ -128,7 +132,10 @@ static __host__ mmvq_parameter_table_id get_device_table_id(int cc) { if (GGML_CUDA_CC_IS_RDNA3_0(cc)) { return MMVQ_PARAMETERS_RDNA3_0; } - if (GGML_CUDA_CC_IS_RDNA2(cc) || GGML_CUDA_CC_IS_RDNA3_5(cc)) { + if (GGML_CUDA_CC_IS_RDNA3_5(cc)) { + return MMVQ_PARAMETERS_RDNA3_5; + } + if (GGML_CUDA_CC_IS_RDNA2(cc)) { return MMVQ_PARAMETERS_RDNA2; } if (GGML_CUDA_CC_IS_GCN(cc) || GGML_CUDA_CC_IS_CDNA(cc)) { @@ -434,6 +441,9 @@ static constexpr __device__ int get_mmvq_mmid_max_batch_for_device() { } static constexpr __host__ __device__ int calc_nwarps(ggml_type type, int ncols_dst, mmvq_parameter_table_id table_id, bool small_k = false, bool halve_iters = false) { + if (table_id == MMVQ_PARAMETERS_RDNA3_5) { + return 1; + } if (table_id == MMVQ_PARAMETERS_GENERIC) { switch (ncols_dst) { case 1: @@ -560,7 +570,681 @@ static constexpr __host__ __device__ int calc_nwarps(ggml_type type, int ncols_d return 1; } +static_assert(calc_nwarps(GGML_TYPE_Q5_1, 1, MMVQ_PARAMETERS_RDNA3_5) == 1); +static_assert(calc_nwarps(GGML_TYPE_Q8_0, 1, MMVQ_PARAMETERS_RDNA3_5) == 1); +static_assert(calc_nwarps(GGML_TYPE_MXFP4, 1, MMVQ_PARAMETERS_RDNA3_5) == 1); +static_assert(calc_nwarps(GGML_TYPE_Q4_K, 1, MMVQ_PARAMETERS_RDNA3_5) == 1); +static_assert(calc_nwarps(GGML_TYPE_Q5_K, 1, MMVQ_PARAMETERS_RDNA3_5) == 1); +static_assert(calc_nwarps(GGML_TYPE_Q6_K, 1, MMVQ_PARAMETERS_RDNA3_5) == 1); +static_assert(calc_nwarps(GGML_TYPE_Q1_0, 1, MMVQ_PARAMETERS_RDNA3_5) == 1); +static_assert(calc_nwarps(GGML_TYPE_Q2_0, 1, MMVQ_PARAMETERS_RDNA3_5) == 1); +static_assert(calc_nwarps(GGML_TYPE_Q4_0, 1, MMVQ_PARAMETERS_RDNA3_5) == 1); +static_assert(calc_nwarps(GGML_TYPE_IQ3_S, 1, MMVQ_PARAMETERS_RDNA3_5) == 1); + +static constexpr __host__ __device__ bool is_rdna3_5_q4_columns_type(ggml_type type) { + return type == GGML_TYPE_Q1_0 || type == GGML_TYPE_Q2_0 || type == GGML_TYPE_Q5_1 || + type == GGML_TYPE_Q4_K || type == GGML_TYPE_Q5_K || type == GGML_TYPE_Q6_K || + type == GGML_TYPE_IQ3_S; +} + +static_assert(is_rdna3_5_q4_columns_type(GGML_TYPE_Q1_0)); +static_assert(is_rdna3_5_q4_columns_type(GGML_TYPE_Q2_0)); +static_assert(is_rdna3_5_q4_columns_type(GGML_TYPE_Q5_1)); +static_assert(is_rdna3_5_q4_columns_type(GGML_TYPE_Q4_K)); +static_assert(is_rdna3_5_q4_columns_type(GGML_TYPE_Q5_K)); +static_assert(is_rdna3_5_q4_columns_type(GGML_TYPE_Q6_K)); +static_assert(is_rdna3_5_q4_columns_type(GGML_TYPE_IQ3_S)); + +template +__launch_bounds__(2 * ggml_cuda_get_physical_warp_size(), 1) +static __global__ void mul_mat_vec_q4_columns( + const void * vx_ptr, const void * vy_ptr, float * dst_ptr, + const uint32_t ncols_x, const uint32_t stride_row_x, const uint32_t stride_col_y, + const uint32_t stride_col_dst, const uint3 channel_ratio, const uint32_t stride_channel_x, + const uint32_t stride_channel_y, const uint32_t stride_channel_dst, const uint3 sample_ratio, + const uint32_t stride_sample_x, const uint32_t stride_sample_y, const uint32_t stride_sample_dst) { + static_assert(type == GGML_TYPE_Q8_0); + constexpr int qk = ggml_cuda_type_traits::qk; + constexpr int qi = ggml_cuda_type_traits::qi; + constexpr int vdr = get_vdr_mmvq(type); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + + const int row = blockIdx.x; + const int lane = threadIdx.x; + const int column0 = 0; + const uint32_t channel_dst = blockIdx.y; + const uint32_t sample_dst = blockIdx.z; + const uint32_t channel_x = fastdiv(channel_dst, channel_ratio); + const uint32_t sample_x = fastdiv(sample_dst, sample_ratio); + + const void * vx = vx_ptr; + const block_q8_1 * y = (const block_q8_1 *) vy_ptr + sample_dst*stride_sample_y + channel_dst*stride_channel_y; + float * dst = dst_ptr + sample_dst*stride_sample_dst + channel_dst*stride_channel_dst + row; + + const int blocks_per_row_x = ncols_x / qk; + const int blocks_per_iter = vdr * warp_size / qi; + const int kbx_offset = sample_x*stride_sample_x + channel_x*stride_channel_x + row*stride_row_x; + float tmp[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + + ggml_cuda_pdl_sync(); + for (int kbx = lane / (qi/vdr); kbx < blocks_per_row_x; kbx += blocks_per_iter) { + const int kby = kbx * (qk/QK8_1); + const int kqs = vdr * (lane % (qi/vdr)); + const block_q8_0 * bx = (const block_q8_0 *) vx + kbx_offset + kbx; + int xv[VDR_Q8_0_Q8_1_MMVQ]; +#pragma unroll + for (int i = 0; i < VDR_Q8_0_Q8_1_MMVQ; ++i) { + xv[i] = get_int_b2(bx->qs, kqs + i); + } + const float dx = __half2float(bx->d); +#pragma unroll + for (int j = 0; j < 4; ++j) { + const block_q8_1 * by = &y[(column0 + j)*stride_col_y + kby]; + int sumi = 0; +#pragma unroll + for (int i = 0; i < VDR_Q8_0_Q8_1_MMVQ; ++i) { + sumi = ggml_cuda_dp4a(xv[i], get_int_b4(by->qs, kqs + i), sumi); + } + tmp[j] += dx * __half2float(__low2half(by->ds)) * (float) sumi; + } + } + +#pragma unroll + for (int j = 0; j < 4; ++j) { + tmp[j] = warp_reduce_sum(tmp[j]); + if (lane == 0) { + dst[(column0 + j)*stride_col_dst] = tmp[j]; + } + } +} + +template +__launch_bounds__(calc_nwarps(type, 1, get_device_table_id()) * ggml_cuda_get_physical_warp_size(), 1) +static __global__ void mul_mat_vec_q4_columns_rdna3_5( + const void * vx_ptr, const void * vy_ptr, float * dst_ptr, + const uint32_t ncols_x, const uint32_t stride_row_x, const uint32_t stride_col_y, + const uint32_t stride_col_dst, const uint3 channel_ratio, const uint32_t stride_channel_x, + const uint32_t stride_channel_y, const uint32_t stride_channel_dst, const uint3 sample_ratio, + const uint32_t stride_sample_x, const uint32_t stride_sample_y, const uint32_t stride_sample_dst) { + static_assert(is_rdna3_5_q4_columns_type(type)); + constexpr int qk = ggml_cuda_type_traits::qk; + constexpr int qi = ggml_cuda_type_traits::qi; + constexpr int vdr = get_vdr_mmvq(type); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int nwarps = calc_nwarps(type, 1, get_device_table_id()); + + const int row = blockIdx.x; + const int lane = threadIdx.x; + const int tid = warp_size*threadIdx.y + lane; + const uint32_t channel_dst = blockIdx.y; + const uint32_t sample_dst = blockIdx.z; + const uint32_t channel_x = fastdiv(channel_dst, channel_ratio); + const uint32_t sample_x = fastdiv(sample_dst, sample_ratio); + + const block_q8_1 * y = (const block_q8_1 *) vy_ptr + sample_dst*stride_sample_y + channel_dst*stride_channel_y; + float * dst = dst_ptr + sample_dst*stride_sample_dst + channel_dst*stride_channel_dst + row; + + const int blocks_per_row_x = ncols_x / qk; + const int blocks_per_iter = vdr * nwarps*warp_size / qi; + const int kbx_offset = sample_x*stride_sample_x + channel_x*stride_channel_x + row*stride_row_x; + float tmp[4] = {0.0f, 0.0f, 0.0f, 0.0f}; + + ggml_cuda_pdl_sync(); + 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)); + + if constexpr (type == GGML_TYPE_Q1_0) { + const block_q1_0 * bx = (const block_q1_0 *) vx_ptr + kbx_offset + kbx; + const int16_t * qs = (const int16_t *) bx->qs + kqs*2; + int xv[8]; +#pragma unroll + for (int i = 0; i < 2; ++i) { + const int q = qs[i]; + const int n0 = __byte_perm(0x11100100, 0x11100100, q >> 0); + const int n1 = __byte_perm(0x11100100, 0x11100100, q >> 2); + const int s0 = __byte_perm(0x01FF, 0x01FF, n0 >> 0); + const int s1 = __byte_perm(0x01FF, 0x01FF, n1 >> 0); + const int s2 = __byte_perm(0x01FF, 0x01FF, n0 >> 16); + const int s3 = __byte_perm(0x01FF, 0x01FF, n1 >> 16); + xv[4*i+0] = __byte_perm(s0, s1, 0x5410); + xv[4*i+1] = __byte_perm(s0, s1, 0x7632); + xv[4*i+2] = __byte_perm(s2, s3, 0x5410); + xv[4*i+3] = __byte_perm(s2, s3, 0x7632); + } + const float dx = bx->d; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const block_q8_1 * by = &y[j*stride_col_y + kby + kqs]; + int sumi = 0; +#pragma unroll + for (int i = 0; i < 2; ++i) { + sumi = ggml_cuda_dp4a(xv[4*i+0], get_int_b4(by->qs, i*4+0), sumi); + sumi = ggml_cuda_dp4a(xv[4*i+1], get_int_b4(by->qs, i*4+1), sumi); + sumi = ggml_cuda_dp4a(xv[4*i+2], get_int_b4(by->qs, i*4+2), sumi); + sumi = ggml_cuda_dp4a(xv[4*i+3], get_int_b4(by->qs, i*4+3), sumi); + } + const float d8 = __low2float(by->ds); + tmp[j] += dx * d8 * sumi; + } + } else if constexpr (type == GGML_TYPE_Q2_0) { + const block_q2_0 * bx = (const block_q2_0 *) vx_ptr + kbx_offset + kbx; + const int16_t * qs = (const int16_t *) bx->qs + kqs*4; + int xv[4]; + int xw[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int q = qs[i]; +#if defined(GGML_USE_HIP) + const uint32_t qx_indices = (q & 0x03) | ((q & 0x0C) << 6) | ((q & 0x30) << 12) | ((q & 0xC0) << 18); + const uint32_t qy_bits = q >> 8; + const uint32_t qy_indices = (qy_bits & 0x03) | ((qy_bits & 0x0C) << 6) | ((qy_bits & 0x30) << 12) | ((qy_bits & 0xC0) << 18); + xv[i] = __builtin_amdgcn_perm(0x020100FF, 0x020100FF, qx_indices); + xw[i] = __builtin_amdgcn_perm(0x020100FF, 0x020100FF, qy_indices); +#else + const int qe = __byte_perm(0x020100FF, 0x020100FF, q >> 0); + const int qo = __byte_perm(0x020100FF, 0x020100FF, q >> 2); + xv[i] = __byte_perm(qe, qo, 0x5140); + xw[i] = __byte_perm(qe, qo, 0x7362); +#endif // defined(GGML_USE_HIP) + } + const float dx = bx->d; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const block_q8_1 * by = &y[j*stride_col_y + kby + kqs]; + int sumi = 0; +#pragma unroll + for (int i = 0; i < 4; ++i) { + sumi = ggml_cuda_dp4a(get_int_b4(by->qs, i*2+0), xv[i], sumi); + sumi = ggml_cuda_dp4a(get_int_b4(by->qs, i*2+1), xw[i], sumi); + } + const float d8 = __low2float(by->ds); + tmp[j] += dx * d8 * sumi; + } + } else if constexpr (type == GGML_TYPE_IQ3_S) { + const block_iq3_s * bx = (const block_iq3_s *) vx_ptr + kbx_offset + kbx; + const int2 qs_packed = make_int2(get_int_b2(bx->qs, kqs + 0), get_int_b2(bx->qs, kqs + 1)); + const uint8_t * qs = (const uint8_t *) &qs_packed; + const int qh = bx->qh[kqs/2]; + const int signs_packed_32 = get_int_b2(bx->signs, kqs/2); + const uint8_t * signs_packed_8 = (const uint8_t *) &signs_packed_32; + int2 xv[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + const int l0 = 2*i; + 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[i] & 0x03) << 7) | ((signs_packed_8[i] & 0x0C) << 21), 0x00000000); + const int signs1 = __vcmpne4(((signs_packed_8[i] & 0x30) << 3) | ((signs_packed_8[i] & 0xC0) << 17), 0x00000000); + xv[i] = make_int2(__vsub4(grid_pos.x ^ signs0, signs0), __vsub4(grid_pos.y ^ signs1, signs1)); + } + const int ls = 1 + 2*((bx->scales[kqs/4] >> ((kqs << 1) & 0x04)) & 0x0F); + const float dx = __half2float(bx->d); +#pragma unroll + for (int j = 0; j < 4; ++j) { + const block_q8_1 * by = &y[j*stride_col_y + kby + kqs/2]; + int sumi = 0; +#pragma unroll + for (int i = 0; i < 4; ++i) { + sumi = ggml_cuda_dp4a(xv[i].x, get_int_b4(by->qs, 2*i + 0), sumi); + sumi = ggml_cuda_dp4a(xv[i].y, get_int_b4(by->qs, 2*i + 1), sumi); + } + sumi *= ls; + const float d = dx * __low2float(by->ds); + tmp[j] += d * sumi; + } + } else if constexpr (type == GGML_TYPE_Q5_1) { + const block_q5_1 * bx = (const block_q5_1 *) vx_ptr + kbx_offset + kbx; + int vl[VDR_Q5_1_Q8_1_MMVQ]; + int vh[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(bx->qs, kqs + i); + vh[i] = get_int_b4(bx->qh, 0) >> (4 * (kqs + i)); + } + const half2 dm = bx->dm; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const block_q8_1 * by = &y[j*stride_col_y + kby]; + int u[2*VDR_Q5_1_Q8_1_MMVQ]; +#pragma unroll + for (int i = 0; i < VDR_Q5_1_Q8_1_MMVQ; ++i) { + u[2*i+0] = get_int_b4(by->qs, kqs + i); + u[2*i+1] = get_int_b4(by->qs, kqs + i + QI5_1); + } + // ROCm 7.14 changes gfx1151 Q5_1 rounding under four-column register pressure. Keep both paths materialized until LLVM preserves this expression. + const float dot = ggml_cuda_materialize(vec_dot_q5_1_q8_1_impl(vl, vh, u, dm, by->ds)); + tmp[j] = __fadd_rn(tmp[j], dot); + } + } else if constexpr (type == GGML_TYPE_Q4_K) { + const block_q4_K * bx = (const block_q4_K *) vx_ptr + kbx_offset + kbx; + const int bq8_offset = QR4_K * ((kqs/2) / (QI8_1/2)); + const int * q4 = (const int *) (bx->qs + 16*bq8_offset + 4*((kqs/2)%4)); + int xv[2] = {q4[0], q4[4]}; + const uint16_t * scales = (const uint16_t *) bx->scales; + uint16_t aux[2]; + const int is = bq8_offset/2; + if (is < 2) { + aux[0] = scales[is+0] & 0x3f3f; + aux[1] = scales[is+2] & 0x3f3f; + } else { + aux[0] = ((scales[is+2] >> 0) & 0x0f0f) | ((scales[is-2] & 0xc0c0) >> 2); + aux[1] = ((scales[is+2] >> 4) & 0x0f0f) | ((scales[is-0] & 0xc0c0) >> 2); + } + const uint8_t * sc = (const uint8_t *) aux; + const uint8_t * m = sc + 2; + const half2 dm = bx->dm; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const block_q8_1 * by = &y[j*stride_col_y + kby]; + int u[2*QR4_K]; + float d8[QR4_K]; +#pragma unroll + for (int i = 0; i < QR4_K; ++i) { + const block_q8_1 * byi = by + bq8_offset + i; + d8[i] = __low2float(byi->ds); + const int * q8 = (const int *) byi->qs + ((kqs/2)%4); + u[2*i+0] = q8[0]; + u[2*i+1] = q8[4]; + } + tmp[j] += vec_dot_q4_K_q8_1_impl_vmmq(xv, u, sc, m, dm, d8); + } + } else if constexpr (type == GGML_TYPE_Q5_K) { + const block_q5_K * bx = (const block_q5_K *) vx_ptr + kbx_offset + kbx; + const int bq8_offset = QR5_K * ((kqs/2) / (QI8_1/2)); + const int * ql = (const int *) (bx->qs + 16*bq8_offset + 4*((kqs/2)%4)); + const int * qh = (const int *) (bx->qh + 4*((kqs/2)%4)); + int vl[2] = {ql[0], ql[4]}; + int vh[2] = {qh[0] >> bq8_offset, qh[4] >> bq8_offset}; + const uint16_t * scales = (const uint16_t *) bx->scales; + uint16_t aux[2]; + const int is = bq8_offset/2; + if (is < 2) { + aux[0] = scales[is+0] & 0x3f3f; + aux[1] = scales[is+2] & 0x3f3f; + } else { + aux[0] = ((scales[is+2] >> 0) & 0x0f0f) | ((scales[is-2] & 0xc0c0) >> 2); + aux[1] = ((scales[is+2] >> 4) & 0x0f0f) | ((scales[is-0] & 0xc0c0) >> 2); + } + const uint8_t * sc = (const uint8_t *) aux; + const uint8_t * m = sc + 2; + const half2 dm = bx->dm; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const block_q8_1 * by = &y[j*stride_col_y + kby]; + int u[2*QR5_K]; + float d8[QR5_K]; +#pragma unroll + for (int i = 0; i < QR5_K; ++i) { + const block_q8_1 * byi = by + bq8_offset + i; + d8[i] = __low2float(byi->ds); + const int * q8 = (const int *) byi->qs + ((kqs/2)%4); + u[2*i+0] = q8[0]; + u[2*i+1] = q8[4]; + } + tmp[j] += vec_dot_q5_K_q8_1_impl_vmmq(vl, vh, u, sc, m, dm, d8); + } + } else if constexpr (type == GGML_TYPE_Q6_K) { + const block_q6_K * bx = (const block_q6_K *) vx_ptr + kbx_offset + kbx; + const int bq8_offset = 2*QR6_K*(kqs/(QI6_K/2)) + (kqs%(QI6_K/2))/(QI6_K/4); + const int scale_offset = (QI6_K/4)*(kqs/(QI6_K/2)) + (kqs%(QI6_K/2))/(QI6_K/8); + const int vh_shift = 2*((kqs%(QI6_K/2))/(QI6_K/4)); + const int vl = get_int_b2(bx->ql, kqs); + const int vh = get_int_b2(bx->qh, (QI6_K/4)*(kqs/(QI6_K/2)) + kqs%(QI6_K/4)) >> vh_shift; + const int8_t * scales = bx->scales + scale_offset; + const float dx = bx->d; +#pragma unroll + for (int j = 0; j < 4; ++j) { + const block_q8_1 * by = &y[j*stride_col_y + kby]; + int u[QR6_K]; + float d8[QR6_K]; +#pragma unroll + for (int i = 0; i < QR6_K; ++i) { + u[i] = get_int_b4(by[bq8_offset + 2*i].qs, kqs % QI8_1); + d8[i] = __low2float(by[bq8_offset + 2*i].ds); + } + tmp[j] += vec_dot_q6_K_q8_1_impl_mmvq(vl, vh, u, scales, dx, d8); + } + } + } + + __shared__ float tmp_shared[nwarps-1 > 0 ? nwarps-1 : 1][4][warp_size]; + if (threadIdx.y > 0) { +#pragma unroll + for (int j = 0; j < 4; ++j) { + tmp_shared[threadIdx.y-1][j][lane] = tmp[j]; + } + } + __syncthreads(); + if (threadIdx.y > 0) { + return; + } + +#pragma unroll + for (int j = 0; j < 4; ++j) { +#pragma unroll + for (int i = 0; i < nwarps-1; ++i) { + tmp[j] += tmp_shared[i][j][lane]; + } + tmp[j] = warp_reduce_sum(tmp[j]); + if (lane == 0) { + dst[j*stride_col_dst] = tmp[j]; + } + } +} + +#ifndef MMVQ_IQ3_S_ROWS_PER_BLOCK +#define MMVQ_IQ3_S_ROWS_PER_BLOCK 1 +#endif + +// gfx1151 IQ3_S decode kernel. Each warp owns an independent output row, as in +// the baseline one-warp MMVQ kernel. Grouping adjacent rows into a block keeps +// their shared activation vector and 2 KiB IQ3 codebook hot without introducing +// a cross-warp reduction or changing the dot-product order. +template +__launch_bounds__(rows_per_block * ggml_cuda_get_physical_warp_size(), 8) +static __global__ void mul_mat_vec_iq3_s_rows_rdna3_5( + const void * vx_ptr, const void * vy_ptr, const int32_t * ids, + const ggml_cuda_mm_fusion_args_device fusion, float * dst_ptr, + const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t nrows_x, + const uint32_t stride_row_x, const uint32_t stride_col_dst, + const uint32_t stride_channel_x, const uint32_t stride_channel_y, + const uint32_t stride_channel_dst, const uint32_t stride_sample_x, + const uint32_t stride_sample_y, const uint32_t stride_sample_dst) { + constexpr int qk = QK_K; + constexpr int qi = ggml_cuda_type_traits::qi; + constexpr int vdr = VDR_IQ3_S_Q8_1_MMVQ; + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + + const int lane = threadIdx.x; + const int row = rows_per_block * blockIdx.x + threadIdx.y; + if (row >= nrows_x) { + return; + } + + const uint32_t channel_dst = blockIdx.y; + const uint32_t sample_dst = blockIdx.z; + const uint32_t channel_x = ids[channel_dst]; + const uint32_t channel_y = fastmodulo(channel_dst, nchannels_y); + const uint32_t sample_x = sample_dst; + const uint32_t sample_y = sample_dst; + + const block_q8_1 * y = (const block_q8_1 *) vy_ptr + + sample_y * stride_sample_y + channel_y * stride_channel_y; + const int kbx_offset = sample_x * stride_sample_x + + channel_x * stride_channel_x + row * stride_row_x; + const int blocks_per_row_x = ncols_x / qk; + constexpr int blocks_per_iter = vdr * warp_size / qi; + + float sum = 0.0f; + float sum_gate = 0.0f; + for (int kbx = lane / (qi/vdr); kbx < blocks_per_row_x; kbx += blocks_per_iter) { + const int kby = kbx * (qk/QK8_1); + const int kqs = vdr * (lane % (qi/vdr)); + sum += vec_dot_iq3_s_q8_1(vx_ptr, &y[kby], kbx_offset + kbx, kqs); + if constexpr (has_gate) { + sum_gate += vec_dot_iq3_s_q8_1(fusion.gate, &y[kby], kbx_offset + kbx, kqs); + } + } + + sum = warp_reduce_sum(sum); + if constexpr (has_gate) { + sum_gate = warp_reduce_sum(sum_gate); + } + + if (lane == 0) { + float result = sum; + if constexpr (has_gate) { + result *= ggml_cuda_op_silu_single(sum_gate); + } + dst_ptr[sample_dst * stride_sample_dst + + channel_dst * stride_channel_dst + row] = result; + } + GGML_UNUSED(stride_col_dst); +} + +// vec_dot_iq3_s_q8_1 with the 512-entry codebook read from a caller-provided (shared memory) copy instead of +// the __device__ table: the eight per-lane gathers per block then go to the LDS pipe instead of the +// texture-address unit, which the direct kernel saturates (MemUnitBusy ~91% on gfx1151). +static __device__ __forceinline__ float vec_dot_iq3_s_q8_1_grid( + const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs, + const uint32_t * __restrict__ grid) { + + 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( + grid[qs[l0 + 0] | ((qh << (8 - l0)) & 0x100)], + 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; +} + +// Direct kernel (same lane layout and accumulation order as mul_mat_vec_iq3_s_rows_rdna3_5) with the codebook +// staged in shared memory once per block. +template +__launch_bounds__(rows_per_block * ggml_cuda_get_physical_warp_size(), 1) +static __global__ void mul_mat_vec_iq3_s_grid_rdna3_5( + const void * vx_ptr, const void * vy_ptr, const int32_t * ids, + const ggml_cuda_mm_fusion_args_device fusion, float * dst_ptr, + const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t nrows_x, + const uint32_t stride_row_x, const uint32_t stride_col_dst, + const uint32_t stride_channel_x, const uint32_t stride_channel_y, + const uint32_t stride_channel_dst, const uint32_t stride_sample_x, + const uint32_t stride_sample_y, const uint32_t stride_sample_dst) { + constexpr int qk = QK_K; + constexpr int qi = ggml_cuda_type_traits::qi; + constexpr int vdr = VDR_IQ3_S_Q8_1_MMVQ; + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + + __shared__ uint32_t s_grid[512]; + + const int lane = threadIdx.x; + const int tid = threadIdx.y * warp_size + lane; + for (int i = tid; i < 512; i += rows_per_block * warp_size) { + s_grid[i] = iq3s_grid[i]; + } + + const int row = rows_per_block * blockIdx.x + threadIdx.y; + const bool row_ok = row < nrows_x; + + const uint32_t channel_dst = blockIdx.y; + const uint32_t sample_dst = blockIdx.z; + const uint32_t channel_x = ids[channel_dst]; + const uint32_t channel_y = fastmodulo(channel_dst, nchannels_y); + const uint32_t sample_x = sample_dst; + const uint32_t sample_y = sample_dst; + + const block_q8_1 * y = (const block_q8_1 *) vy_ptr + + sample_y * stride_sample_y + channel_y * stride_channel_y; + const int kbx_offset = sample_x * stride_sample_x + + channel_x * stride_channel_x + (row_ok ? row : 0) * stride_row_x; + const int blocks_per_row_x = ncols_x / qk; + constexpr int blocks_per_iter = vdr * warp_size / qi; + + __syncthreads(); + if (!row_ok) { + return; + } + + float sum = 0.0f; + float sum_gate = 0.0f; + for (int kbx = lane / (qi/vdr); kbx < blocks_per_row_x; kbx += blocks_per_iter) { + const int kby = kbx * (qk/QK8_1); + const int kqs = vdr * (lane % (qi/vdr)); + sum += vec_dot_iq3_s_q8_1_grid(vx_ptr, &y[kby], kbx_offset + kbx, kqs, s_grid); + if constexpr (has_gate) { + sum_gate += vec_dot_iq3_s_q8_1_grid(fusion.gate, &y[kby], kbx_offset + kbx, kqs, s_grid); + } + } + + sum = warp_reduce_sum(sum); + if constexpr (has_gate) { + sum_gate = warp_reduce_sum(sum_gate); + } + + if (lane == 0) { + float result = sum; + if constexpr (has_gate) { + result *= ggml_cuda_op_silu_single(sum_gate); + } + dst_ptr[sample_dst * stride_sample_dst + + channel_dst * stride_channel_dst + row] = result; + } + GGML_UNUSED(stride_col_dst); +} + +// gfx1151 IQ3_S decode kernel with LDS staging, ncols == 2560 (10 blocks, 1100 B per row). The direct +// per-lane field loads of vec_dot_iq3_s_q8_1 (about 18 vector-memory instructions per block per lane, most of +// them for the Q8_1 activations) saturate the texture-address unit (MemUnitBusy ~91%) long before DRAM +// bandwidth. Here each wave first copies its up (and gate) row into shared memory with 9 dword loads per +// lane, the block copies the quantized activations once, and the dot products then read from LDS. +// The per-lane accumulation and reduction order is the same as mul_mat_vec_iq3_s_rows_rdna3_5. +#define MMVQ_IQ3_LDS_NCOLS 2560 +#define MMVQ_IQ3_LDS_ROW_DWORDS 275 // 10 blocks * 110 B / 4 +#define MMVQ_IQ3_LDS_ROW_SLOT 276 +#define MMVQ_IQ3_LDS_Y_DWORDS 720 // 80 Q8_1 blocks * 36 B / 4 + +template +__launch_bounds__(rows_per_block * ggml_cuda_get_physical_warp_size(), 1) +static __global__ void mul_mat_vec_iq3_s_lds_rdna3_5( + const void * vx_ptr, const void * vy_ptr, const int32_t * ids, + const ggml_cuda_mm_fusion_args_device fusion, float * dst_ptr, + const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t nrows_x, + const uint32_t stride_row_x, const uint32_t stride_col_dst, + const uint32_t stride_channel_x, const uint32_t stride_channel_y, + const uint32_t stride_channel_dst, const uint32_t stride_sample_x, + const uint32_t stride_sample_y, const uint32_t stride_sample_dst) { + constexpr int qk = QK_K; + constexpr int qi = ggml_cuda_type_traits::qi; + constexpr int vdr = VDR_IQ3_S_Q8_1_MMVQ; + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int blocks_per_row = MMVQ_IQ3_LDS_NCOLS / qk; + constexpr int blocks_per_iter = vdr * warp_size / qi; + constexpr int row_loads = (MMVQ_IQ3_LDS_ROW_DWORDS + warp_size - 1) / warp_size; // 9 + + __shared__ int s_y[MMVQ_IQ3_LDS_Y_DWORDS]; + __shared__ int s_rows[rows_per_block][has_gate ? 2 : 1][MMVQ_IQ3_LDS_ROW_SLOT]; + + const int lane = threadIdx.x; + const int wave = threadIdx.y; + const int tid = wave * warp_size + lane; + const int row = rows_per_block * blockIdx.x + wave; + const bool row_ok = row < nrows_x; + + const uint32_t channel_dst = blockIdx.y; + const uint32_t sample_dst = blockIdx.z; + const uint32_t channel_x = ids[channel_dst]; + const uint32_t channel_y = fastmodulo(channel_dst, nchannels_y); + + const block_q8_1 * y = (const block_q8_1 *) vy_ptr + sample_dst * stride_sample_y + channel_y * stride_channel_y; + const int kbx_offset = sample_dst * stride_sample_x + channel_x * stride_channel_x + (row_ok ? row : 0) * stride_row_x; + + // 1. weight rows -> registers (dword loads; rows are 4 B aligned: 1100 B rows, even block offsets) + const int * gx = (const int *) ((const block_iq3_s *) vx_ptr + kbx_offset); + int rx[row_loads]; + int rg[row_loads]; +#pragma unroll + for (int i = 0; i < row_loads; ++i) { + const int idx = i * warp_size + lane; + rx[i] = idx < MMVQ_IQ3_LDS_ROW_DWORDS ? gx[idx] : 0; + } + if constexpr (has_gate) { + const int * gg = (const int *) ((const block_iq3_s *) fusion.gate + kbx_offset); +#pragma unroll + for (int i = 0; i < row_loads; ++i) { + const int idx = i * warp_size + lane; + rg[i] = idx < MMVQ_IQ3_LDS_ROW_DWORDS ? gg[idx] : 0; + } + } + + // 2. quantized activations -> LDS (shared by the rows of the block) + { + const int * gy = (const int *) y; + for (int i = tid; i < MMVQ_IQ3_LDS_Y_DWORDS; i += rows_per_block * warp_size) { + s_y[i] = gy[i]; + } + } + + // 3. rows -> LDS +#pragma unroll + for (int i = 0; i < row_loads; ++i) { + const int idx = i * warp_size + lane; + if (idx < MMVQ_IQ3_LDS_ROW_DWORDS) { + s_rows[wave][0][idx] = rx[i]; + if constexpr (has_gate) { + s_rows[wave][1][idx] = rg[i]; + } + } + } + __syncthreads(); + + if (!row_ok) { + return; + } + + // 4. dot products from LDS, same lane -> block assignment as the direct kernel + const block_q8_1 * y_l = (const block_q8_1 *) s_y; + const void * x_l = s_rows[wave][0]; + float sum = 0.0f; + float sum_gate = 0.0f; +#pragma unroll + for (int kbx = lane / (qi/vdr); kbx < blocks_per_row; kbx += blocks_per_iter) { + const int kby = kbx * (qk/QK8_1); + const int kqs = vdr * (lane % (qi/vdr)); + sum += vec_dot_iq3_s_q8_1(x_l, &y_l[kby], kbx, kqs); + if constexpr (has_gate) { + sum_gate += vec_dot_iq3_s_q8_1(s_rows[wave][1], &y_l[kby], kbx, kqs); + } + } + + sum = warp_reduce_sum(sum); + if constexpr (has_gate) { + sum_gate = warp_reduce_sum(sum_gate); + } + + if (lane == 0) { + float result = sum; + if constexpr (has_gate) { + result *= ggml_cuda_op_silu_single(sum_gate); + } + dst_ptr[sample_dst * stride_sample_dst + channel_dst * stride_channel_dst + row] = result; + } + GGML_UNUSED(stride_col_dst); + GGML_UNUSED(ncols_x); +} + static constexpr __host__ __device__ int calc_rows_per_block(int ncols_dst, int table_id, bool small_k = false, int nwarps = 1) { + if (table_id == MMVQ_PARAMETERS_RDNA3_5) { + return 1; + } if (table_id == MMVQ_PARAMETERS_GENERIC || table_id == MMVQ_PARAMETERS_GCN || table_id == MMVQ_PARAMETERS_TURING || table_id == MMVQ_PARAMETERS_GB10) { switch (ncols_dst) { case 1: @@ -726,12 +1410,36 @@ static __global__ void mul_mat_vec_q( for (int j = 0; j < ncols_dst; ++j) { #pragma unroll for (int i = 0; i < rows_per_cuda_block; ++i) { - tmp[j][i] += vec_dot_q_cuda( - vx, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); + if constexpr (type == GGML_TYPE_IQ3_S && has_fusion) { + if (use_gate) { + float dot; + float dot_gate; + vec_dot_iq3_s_q8_1_pair( + dot, dot_gate, vx, vgate, &y[j*stride_col_y + kby], + kbx_offset + i*stride_row_x + kbx, kqs); + tmp[j][i] += dot; + tmp_gate[j][i] += dot_gate; + continue; + } + } + if constexpr (type == GGML_TYPE_Q5_1 && table_id == MMVQ_PARAMETERS_RDNA3_5) { + const float dot = ggml_cuda_materialize(vec_dot_q_cuda( + vx, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs)); + tmp[j][i] = __fadd_rn(tmp[j][i], dot); + } else { + tmp[j][i] += vec_dot_q_cuda( + vx, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); + } if constexpr (has_fusion) { if (use_gate) { - tmp_gate[j][i] += vec_dot_q_cuda( - vgate, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); + if constexpr (type == GGML_TYPE_Q5_1 && table_id == MMVQ_PARAMETERS_RDNA3_5) { + const float dot_gate = ggml_cuda_materialize(vec_dot_q_cuda( + vgate, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs)); + tmp_gate[j][i] = __fadd_rn(tmp_gate[j][i], dot_gate); + } else { + tmp_gate[j][i] += vec_dot_q_cuda( + vgate, &y[j*stride_col_y + kby], kbx_offset + i*stride_row_x + kbx, kqs); + } } } } @@ -828,39 +1536,792 @@ static __global__ void mul_mat_vec_q( } } -// Dedicated MoE multi-token kernel. -// Grid: (ceil(nrows_x / c_rows_per_block), nchannels_dst) -// Block: (warp_size, ncols_dst) - each warp handles one token independently. -// No shared memory reduction needed since each warp works alone. -template -__launch_bounds__(get_mmvq_mmid_max_batch_for_device()*ggml_cuda_get_physical_warp_size(), 1) -static __global__ void mul_mat_vec_q_moe( - const void * vx_ptr, const void * vy_ptr, const int32_t * ids_ptr, const ggml_cuda_mm_fusion_args_device fusion, - float * dst_ptr, - const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t nrows_x, - const uint32_t stride_row_x, const uint32_t stride_col_y, const uint32_t stride_col_dst, - const uint32_t stride_channel_x, const uint32_t stride_channel_y, const uint32_t stride_channel_dst, - const uint32_t ncols_dst, const uint32_t ids_stride) { - const void * GGML_CUDA_RESTRICT vx = vx_ptr; - const void * GGML_CUDA_RESTRICT vy = vy_ptr; - const int32_t * GGML_CUDA_RESTRICT ids = ids_ptr; - float * GGML_CUDA_RESTRICT dst = dst_ptr; - - constexpr int qk = ggml_cuda_type_traits::qk; - constexpr int qi = ggml_cuda_type_traits::qi; - constexpr int vdr = get_vdr_mmvq(type); - constexpr int warp_size = ggml_cuda_get_physical_warp_size(); +// --------------------------------------------------------------------------------------------- +// RDNA3.5 single-column MMVQ with the Q8_1 activation quantization fused into the matvec kernel. +// +// Token generation on gfx1151 pays a ~2 us launch gap for every kernel, so the quantize_q8_1 launch in front +// of each matvec costs about as much as the whole quantize kernel. A block of MMVQ_FQ_NWARPS waves (one output +// row per wave) quantizes the activation vector once into shared memory, replicating quantize_q8_1 exactly +// (same amax, same IEEE division and roundf), and then runs the per-wave K loop of mul_mat_vec_q in +// the same order against it, so the outputs are bit-identical to the two-kernel path. +// +// Measured on gfx1151 with Qwen3.6-35B-A3B Q8_0 decode (2026-09-03): 16 waves per block and 2 prefetched +// K iterations; 8 waves or deeper prefetch (more VGPRs, less occupancy) were slower. + +#define MMVQ_FQ_NWARPS 16 +// whether the fused-quantize path also serves the IQ3_S/IQ4_NL/IQ4_XS expert matvecs (see mmvq_fq_type_ok) +#ifndef MMVQ_FQ_IQ_TYPES +#define MMVQ_FQ_IQ_TYPES false +#endif - constexpr vec_dot_q_cuda_t vec_dot_q_cuda = get_vec_dot_q_cuda(type); +static constexpr __host__ __device__ bool mmvq_fq_type_ok(ggml_type type) { + switch (type) { + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q6_K: + // Q4_K/Q5_K (block sum via mmvq_fq_block_sum) are exact too but measured neutral on gfx1151 + // with the generic, non-prefetched loop (Qwen3.6 UD-Q4_K_XL TG128 60.69 -> 60.63 t/s). + return true; + case GGML_TYPE_IQ3_S: + case GGML_TYPE_IQ4_NL: + case GGML_TYPE_IQ4_XS: + // MoE expert types of the Unsloth qwen4exp IQ4_XS mix: their vec_dot only reads the block scale, so + // the in-kernel quantization is exact; used to drop the quantize launch in front of every expert matvec + return MMVQ_FQ_IQ_TYPES; + default: + return false; + } +} - // fuse gate, bias, scales, and glu_op into the up projection - bool use_gate = false; - const void * vgate = nullptr; - const float * x_bias = nullptr; - const float * gate_bias = nullptr; - const float * x_scale = nullptr; - const float * gate_scale = nullptr; - ggml_glu_op active_glu = GGML_GLU_OP_SWIGLU; +// Whether the vec_dot of the type reads the Q8_1 block sum (ds.y) in addition to the scale. +static constexpr __host__ __device__ bool mmvq_fq_needs_sum(ggml_type type) { + switch (type) { + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q6_K: + case GGML_TYPE_IQ3_S: + case GGML_TYPE_IQ4_NL: + case GGML_TYPE_IQ4_XS: + return false; + default: + return true; + } +} + +// Sum of the 32 values of a Q8_1 block in exactly the order of warp_reduce_sum<32> in quantize_q8_1, where lane +// l holds value l. Here lane (sub) of a group of 4 holds values 8*sub .. 8*sub+7 in v[]. +static __device__ __forceinline__ float mmvq_fq_block_sum(const float * v, const int width) { + float a[8]; +#pragma unroll + for (int k = 0; k < 8; ++k) { + // offset 16: a_i = x_i + x_{i^16}, x_{i^16} lives two lanes over + a[k] = v[k] + __shfl_xor_sync(0xffffffff, v[k], 2, width); + } + float b[8]; +#pragma unroll + for (int k = 0; k < 8; ++k) { + // offset 8: b_i = a_i + a_{i^8}, one lane over + b[k] = a[k] + __shfl_xor_sync(0xffffffff, a[k], 1, width); + } + float c[4]; +#pragma unroll + for (int k = 0; k < 4; ++k) { + c[k] = b[k] + b[k + 4]; // offset 4 + } + const float d0 = c[0] + c[2]; // offset 2 + const float d1 = c[1] + c[3]; + return d0 + d1; // offset 1 +} + +// Register-staged Q8_0 weight loads: the K loop of one wave is split into chunks of MMVQ_FQ_PF iterations whose +// weight blocks are loaded up front, so the DRAM stream starts before the activation quantization and the +// barrier, and each wave keeps MMVQ_FQ_PF loads in flight instead of one. +#define MMVQ_FQ_PF 2 +// Long rows with few output rows (e.g. the [10240 x 320] hyper-connection down projections of qwen4exp) have +// too few waves in flight to cover DRAM latency with 2 loads per wave; they use a deeper prefetch instead. +// The accumulation order per lane is the same for any depth, so the results do not change. +#define MMVQ_FQ_PF_LONG 8 + +template +struct mmvq_fq_q8_0_chunk { + int qs[pf][VDR_Q8_0_Q8_1_MMVQ]; + half d[pf]; +}; + +template +static __device__ __forceinline__ void mmvq_fq_q8_0_load( + mmvq_fq_q8_0_chunk & c, const block_q8_0 * GGML_CUDA_RESTRICT x, const int kbx0, const int kqs, const int blocks_per_row_x) { + constexpr int blocks_per_iter = VDR_Q8_0_Q8_1_MMVQ * ggml_cuda_get_physical_warp_size() / QI8_0; +#pragma unroll + for (int i = 0; i < pf; ++i) { + const int kbx = kbx0 + i*blocks_per_iter; + if (kbx < blocks_per_row_x) { +#pragma unroll + for (int j = 0; j < VDR_Q8_0_Q8_1_MMVQ; ++j) { + c.qs[i][j] = get_int_b2(x[kbx].qs, kqs + j); + } + c.d[i] = x[kbx].d; + } + } +} + +template +static __device__ __forceinline__ void mmvq_fq_q8_0_dot( + float & tmp, const mmvq_fq_q8_0_chunk & c, const block_q8_1 * y, const int kbx0, const int kqs, const int blocks_per_row_x) { + constexpr int blocks_per_iter = VDR_Q8_0_Q8_1_MMVQ * ggml_cuda_get_physical_warp_size() / QI8_0; +#pragma unroll + for (int i = 0; i < pf; ++i) { + const int kbx = kbx0 + i*blocks_per_iter; + if (kbx < blocks_per_row_x) { + const block_q8_1 * bq8_1 = &y[kbx]; // qk == QK8_1 for Q8_0 + int u[VDR_Q8_0_Q8_1_MMVQ]; +#pragma unroll + for (int j = 0; j < VDR_Q8_0_Q8_1_MMVQ; ++j) { + u[j] = get_int_b4(bq8_1->qs, kqs + j); + } + tmp += vec_dot_q8_0_q8_1_impl(c.qs[i], u, c.d[i], __low2half(bq8_1->ds)); + } + } +} + +template +__launch_bounds__(nwarps * ggml_cuda_get_physical_warp_size(), 1) +static __global__ void mul_mat_vec_q_fq( + const void * vx_ptr, const float * y_ptr, const int32_t * ids_ptr, const ggml_cuda_mm_fusion_args_device fusion, float * dst_ptr, + const uint32_t ncols_x, const uint32_t nrows_x, const uint3 nchannels_y, const uint32_t stride_row_x, + const uint32_t stride_col_dst, const uint3 channel_ratio, const uint32_t stride_channel_x, + const uint32_t stride_channel_y, const uint32_t stride_channel_dst, const uint3 sample_ratio, + const uint32_t stride_sample_x, const uint32_t stride_sample_y, const uint32_t stride_sample_dst, + const float y_scale, const float y_bias, const int y_op) { + static_assert(mmvq_fq_type_ok(type)); + constexpr bool q8_0_path = type == GGML_TYPE_Q8_0; + const void * GGML_CUDA_RESTRICT vx = vx_ptr; + const int32_t * GGML_CUDA_RESTRICT ids = ids_ptr; + float * GGML_CUDA_RESTRICT dst = dst_ptr; + + constexpr int qk = ggml_cuda_type_traits::qk; + constexpr int qi = ggml_cuda_type_traits::qi; + constexpr int vdr = get_vdr_mmvq(type); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int blocks_per_iter = vdr * warp_size / qi; + static_assert(!q8_0_path || qk == QK8_1); + [[maybe_unused]] constexpr vec_dot_q_cuda_t vec_dot_q_cuda = get_vec_dot_q_cuda(type); + + extern __shared__ char mmvq_fq_smem[]; + block_q8_1 * y_q8 = (block_q8_1 *) mmvq_fq_smem; + + const int lane = threadIdx.x; + const int tid = warp_size*threadIdx.y + lane; + + const uint32_t channel_dst = blockIdx.y; + const uint32_t sample_dst = blockIdx.z; + const uint32_t channel_x = ids ? ids[channel_dst] : fastdiv(channel_dst, channel_ratio); + const uint32_t channel_y = ids ? fastmodulo(channel_dst, nchannels_y) : channel_dst; + const uint32_t sample_x = fastdiv(sample_dst, sample_ratio); + const uint32_t sample_y = sample_dst; + + const uint32_t row = blockIdx.x*nwarps + threadIdx.y; + const bool row_ok = row < nrows_x; + const uint32_t row_safe = row_ok ? row : nrows_x - 1; + + bool use_gate = false; + bool use_bias = false; + bool use_gate_bias = false; + [[maybe_unused]] const void * vgate = nullptr; + const float * x_bias = nullptr; + const float * gate_bias = nullptr; + ggml_glu_op active_glu; + float glu_limit = 0.0f; + + if constexpr (has_fusion) { + use_gate = fusion.gate != nullptr; + use_bias = fusion.x_bias != nullptr; + use_gate_bias = fusion.gate_bias != nullptr && use_gate; + vgate = fusion.gate; + x_bias = (const float *) fusion.x_bias; + gate_bias = (const float *) fusion.gate_bias; + active_glu = fusion.glu_op; + glu_limit = fusion.glu_limit; + } + + // 1. start streaming the weights of this wave's row before anything else + const int blocks_per_row_x = ncols_x / qk; + const int kbx_offset = sample_x*stride_sample_x + channel_x*stride_channel_x + row_safe*stride_row_x; + const int kqs = vdr * (lane % (qi/vdr)); + const int kbx0 = lane / (qi/vdr); + + [[maybe_unused]] const block_q8_0 * x = (const block_q8_0 *) vx + kbx_offset; + [[maybe_unused]] const block_q8_0 * xg = nullptr; + + [[maybe_unused]] mmvq_fq_q8_0_chunk cx; + [[maybe_unused]] mmvq_fq_q8_0_chunk cg; + if constexpr (q8_0_path) { + mmvq_fq_q8_0_load(cx, x, kbx0, kqs, blocks_per_row_x); + if constexpr (has_fusion) { + if (use_gate) { + xg = (const block_q8_0 *) vgate + kbx_offset; + mmvq_fq_q8_0_load(cg, xg, kbx0, kqs, blocks_per_row_x); + } + } + } + + // 2. quantize the activation vector into shared memory: 4 lanes per Q8_1 block, 8 values per lane + { + const float * y = y_ptr + sample_y*stride_sample_y + channel_y*stride_channel_y; + const int nby = ncols_x / QK8_1; + for (int i = tid; i < 4*nby; i += nwarps*warp_size) { + const int ib = i >> 2; + const int sub = i & 3; + const float4 * yv = (const float4 *) (y + ib*QK8_1 + sub*8); + const float4 v0 = yv[0]; + const float4 v1 = yv[1]; + float v[8] = {v0.x, v0.y, v0.z, v0.w, v1.x, v1.y, v1.z, v1.w}; + if (y_op == 3) { + // qwen4exp GDN output: per-head (128) rms norm * weight, gated by sigmoid(z). The 16 lanes of a + // head are consecutive threads of one wave (4 Q8_1 blocks x 4 lanes). + float ss = 0.0f; +#pragma unroll + for (int k = 0; k < 8; ++k) { + ss = fmaf(v[k], v[k], ss); + } + ss += __shfl_xor_sync(0xffffffff, ss, 1, warp_size); + ss += __shfl_xor_sync(0xffffffff, ss, 2, warp_size); + ss += __shfl_xor_sync(0xffffffff, ss, 4, warp_size); + ss += __shfl_xor_sync(0xffffffff, ss, 8, warp_size); + const float rms_scale = rsqrtf(ss / 128.0f + fusion.y_eps); + const float * z = (const float *) fusion.y_gate + ib*QK8_1 + sub*8; + const float * w = (const float *) fusion.y_norm_w + (ib*QK8_1 + sub*8) % 128; + const float4 * zv = (const float4 *) z; + const float4 * wv = (const float4 *) w; + const float4 z0 = zv[0], z1 = zv[1], w0 = wv[0], w1 = wv[1]; + const float zz[8] = {z0.x, z0.y, z0.z, z0.w, z1.x, z1.y, z1.z, z1.w}; + const float ww[8] = {w0.x, w0.y, w0.z, w0.w, w1.x, w1.y, w1.z, w1.w}; +#pragma unroll + for (int k = 0; k < 8; ++k) { + // rms_norm_f32<.., true>: scale * x * w ; unary_gated: sigmoid(z) * normed + v[k] = (1.0f / (1.0f + expf(-zz[k]))) * (rms_scale * v[k] * ww[k]); + } + } else if (y_op != 0) { + // fused activation prologue (scale -> unary), same expressions as scale_f32 / op_silu / op_sigmoid +#pragma unroll + for (int k = 0; k < 8; ++k) { + const float t = y_scale * v[k] + y_bias; + v[k] = y_op == 1 ? ggml_cuda_op_silu_single(t) : 1.0f / (1.0f + expf(-t)); + } + } + + float amax = fabsf(v[0]); +#pragma unroll + for (int k = 1; k < 8; ++k) { + amax = fmaxf(amax, fabsf(v[k])); + } + amax = fmaxf(amax, __shfl_xor_sync(0xffffffff, amax, 1, warp_size)); + amax = fmaxf(amax, __shfl_xor_sync(0xffffffff, amax, 2, warp_size)); + + const float sum = mmvq_fq_needs_sum(type) ? mmvq_fq_block_sum(v, warp_size) : 0.0f; + + const float d = amax / 127.0f; + int q[8]; +#pragma unroll + for (int k = 0; k < 8; ++k) { + const int8_t qk8 = amax == 0.0f ? 0 : roundf(v[k] / d); + q[k] = (int) qk8 & 0xff; + } + int * qs = (int *) y_q8[ib].qs; + qs[2*sub + 0] = q[0] | (q[1] << 8) | (q[2] << 16) | (q[3] << 24); + qs[2*sub + 1] = q[4] | (q[5] << 8) | (q[6] << 16) | (q[7] << 24); + if (sub == 0) { + y_q8[ib].ds = make_half2(d, sum); + } + } + } + __syncthreads(); + + if (!row_ok) { + return; + } + + [[maybe_unused]] float x_biases = 0.0f; + [[maybe_unused]] float gate_biases = 0.0f; + if constexpr (has_fusion) { + const uint32_t channel_bias = ids ? channel_x : channel_dst; + if (lane == 0) { + if (use_bias) { + x_biases = x_bias[sample_dst*stride_sample_dst + channel_bias*stride_channel_dst + row]; + } + if (use_gate_bias) { + gate_biases = gate_bias[sample_dst*stride_sample_dst + channel_bias*stride_channel_dst + row]; + } + } + } + + // 3. exact per-wave dot products in the k order of mul_mat_vec_q, one chunk ahead in memory + float tmp = 0.0f; + float tmp_gate = 0.0f; + + const block_q8_1 * y = y_q8; + if constexpr (!q8_0_path) { + // generic types: the per-wave loop of mul_mat_vec_q against the shared-memory activations + for (int kbx = kbx0; kbx < blocks_per_row_x; kbx += blocks_per_iter) { + const int kby = kbx * (qk/QK8_1); + tmp += vec_dot_q_cuda(vx, &y[kby], kbx_offset + kbx, kqs); + if constexpr (has_fusion) { + if (use_gate) { + tmp_gate += vec_dot_q_cuda(vgate, &y[kby], kbx_offset + kbx, kqs); + } + } + } + } + constexpr int chunk_blocks = pf*blocks_per_iter; + for (int c0 = kbx0; q8_0_path && c0 < blocks_per_row_x; c0 += chunk_blocks) { + mmvq_fq_q8_0_chunk nx; + [[maybe_unused]] mmvq_fq_q8_0_chunk ng; + const int c1 = c0 + chunk_blocks; + if (c1 < blocks_per_row_x) { + mmvq_fq_q8_0_load(nx, x, c1, kqs, blocks_per_row_x); + if constexpr (has_fusion) { + if (use_gate) { + mmvq_fq_q8_0_load(ng, xg, c1, kqs, blocks_per_row_x); + } + } + } + mmvq_fq_q8_0_dot(tmp, cx, y, c0, kqs, blocks_per_row_x); + if constexpr (has_fusion) { + if (use_gate) { + mmvq_fq_q8_0_dot(tmp_gate, cg, y, c0, kqs, blocks_per_row_x); + } + } + if (c1 < blocks_per_row_x) { + cx = nx; + if constexpr (has_fusion) { + if (use_gate) { + cg = ng; + } + } + } + } + + tmp = warp_reduce_sum(tmp); + if constexpr (has_fusion) { + if (use_gate) { + tmp_gate = warp_reduce_sum(tmp_gate); + } + } + + if (lane == 0) { + float result = tmp; + if constexpr (has_fusion) { + result += x_biases; + if (use_gate) { + float gate_value = tmp_gate; + gate_value += gate_biases; + switch (active_glu) { + case GGML_GLU_OP_SWIGLU: + result *= ggml_cuda_op_silu_single(gate_value); + break; + case GGML_GLU_OP_GEGLU: + result *= ggml_cuda_op_gelu_single(gate_value); + break; + case GGML_GLU_OP_SWIGLU_OAI: + result = ggml_cuda_op_swiglu_oai_single(gate_value, result); + break; + case GGML_GLU_OP_SWIGLU_CLAMP: + result = ggml_cuda_op_swiglu_clamp_single(gate_value, result, glu_limit); + break; + default: + result = result * gate_value; + break; + } + } + } + dst[sample_dst*stride_sample_dst + channel_dst*stride_channel_dst + row] = result; + } + + if constexpr (!has_fusion) { + GGML_UNUSED_VARS(use_gate, use_bias, use_gate_bias, active_glu, glu_limit, gate_bias, x_bias, tmp_gate); + } + GGML_UNUSED(stride_col_dst); +} + +template +static void mul_mat_vec_q_fq_launch( + const void * vx, const float * y, const int32_t * ids, const ggml_cuda_mm_fusion_args_device & fusion, float * dst, + const int ncols_x, const int nrows_x, const int stride_row_x, const int stride_col_dst, + const int nchannels_x, const int nchannels_y, const int nchannels_dst, + const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst, + const int nsamples_x, const int nsamples_dst, const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst, + cudaStream_t stream, const float y_scale, const float y_bias, const int y_op) { + const int device = ggml_cuda_get_device(); + const int warp_size = ggml_cuda_info().devices[device].warp_size; + + const uint3 nchannels_y_fd = ids ? init_fastdiv_values(nchannels_y) : make_uint3(0, 0, 0); + const uint3 channel_ratio_fd = ids ? make_uint3(0, 0, 0) : init_fastdiv_values(nchannels_dst / nchannels_x); + const uint3 sample_ratio_fd = init_fastdiv_values(nsamples_dst / nsamples_x); + + // One output row per wave. Narrow matrices (e.g. the 320-row hyper-connection down projections of + // qwen4exp) would give only nrows/16 blocks, leaving most CUs idle while each wave streams a long row. + // Use fewer waves per block there so at least ~2 blocks per CU are in flight; every row is still reduced + // by a single wave and the activation quantization is replicated per block, so results are unchanged. + const int nsm = ggml_cuda_info().devices[device].nsm; + int nwarps = MMVQ_FQ_NWARPS; + while (nwarps > 4 && ((nrows_x + nwarps - 1) / nwarps) * nchannels_dst * nsamples_dst < 2 * nsm) { + nwarps /= 2; + } + + const dim3 block_nums((nrows_x + nwarps - 1) / nwarps, nchannels_dst, nsamples_dst); + const dim3 block_dims(warp_size, nwarps, 1); + const int nbytes_shared = (ncols_x / QK8_1) * sizeof(block_q8_1); + + const bool has_fusion = fusion.gate != nullptr || fusion.x_bias != nullptr || fusion.gate_bias != nullptr; + const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, nbytes_shared, stream); + + auto launch = [&](auto kernel) { + ggml_cuda_kernel_launch(kernel, launch_params, + vx, y, ids, fusion, dst, ncols_x, nrows_x, nchannels_y_fd, stride_row_x, stride_col_dst, + channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, + sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst, y_scale, y_bias, y_op); + }; + + // deeper weight prefetch for long rows when the grid is too small to hide DRAM latency with many waves + constexpr int blocks_per_iter = VDR_Q8_0_Q8_1_MMVQ * 32 / QI8_0; + const bool long_rows = type == GGML_TYPE_Q8_0 && (ncols_x / QK8_0) >= 16 * blocks_per_iter && + (int64_t) nrows_x * nchannels_dst * nsamples_dst <= 1024; + + switch (nwarps) { + case 16: + if (has_fusion) { launch(mul_mat_vec_q_fq); } + else { launch(mul_mat_vec_q_fq); } + break; + case 8: + if (long_rows) { + if (has_fusion) { launch(mul_mat_vec_q_fq); } + else { launch(mul_mat_vec_q_fq); } + } else { + if (has_fusion) { launch(mul_mat_vec_q_fq); } + else { launch(mul_mat_vec_q_fq); } + } + break; + case 4: + if (long_rows) { + if (has_fusion) { launch(mul_mat_vec_q_fq); } + else { launch(mul_mat_vec_q_fq); } + } else { + if (has_fusion) { launch(mul_mat_vec_q_fq); } + else { launch(mul_mat_vec_q_fq); } + } + break; + default: + GGML_ABORT("fatal error"); + } +} + +static void mul_mat_vec_q_fq_switch_type( + const void * vx, const ggml_type type_x, const float * y, const int32_t * ids, const ggml_cuda_mm_fusion_args_device & fusion, float * dst, + const int ncols_x, const int nrows_x, const int stride_row_x, const int stride_col_dst, + const int nchannels_x, const int nchannels_y, const int nchannels_dst, + const int stride_channel_x, const int stride_channel_y, const int stride_channel_dst, + const int nsamples_x, const int nsamples_dst, const int stride_sample_x, const int stride_sample_y, const int stride_sample_dst, + cudaStream_t stream, const float y_scale, const float y_bias, const int y_op) { +#define MMVQ_FQ_LAUNCH(T) \ + mul_mat_vec_q_fq_launch(vx, y, ids, fusion, dst, ncols_x, nrows_x, stride_row_x, stride_col_dst, \ + nchannels_x, nchannels_y, nchannels_dst, stride_channel_x, stride_channel_y, stride_channel_dst, \ + nsamples_x, nsamples_dst, stride_sample_x, stride_sample_y, stride_sample_dst, stream, y_scale, y_bias, y_op) + switch (type_x) { + case GGML_TYPE_Q8_0: MMVQ_FQ_LAUNCH(GGML_TYPE_Q8_0); break; + case GGML_TYPE_Q6_K: MMVQ_FQ_LAUNCH(GGML_TYPE_Q6_K); break; +#if MMVQ_FQ_IQ_TYPES + case GGML_TYPE_IQ3_S: MMVQ_FQ_LAUNCH(GGML_TYPE_IQ3_S); break; + case GGML_TYPE_IQ4_NL: MMVQ_FQ_LAUNCH(GGML_TYPE_IQ4_NL); break; + case GGML_TYPE_IQ4_XS: MMVQ_FQ_LAUNCH(GGML_TYPE_IQ4_XS); break; +#endif + default: + GGML_ABORT("fatal error"); + break; + } +#undef MMVQ_FQ_LAUNCH +} + +// --------------------------------------------------------------------------------------------- +// RDNA3.5 grouped single-column matvec: several independent matvecs that share the same activation vector +// Q8_0 matvecs are launched as one kernel. Each segment is a plain [ncols x nrows] weight matrix with an optional +// Q8_0 gate and GLU epilogue. Blocks are assigned to segments by a prefix over the segment block counts. + +#define MMVQ_GROUP_MAX 4 + +struct mmvq_group_seg_dev { + const void * vx; + const void * gate; + float * dst; + int nrows; + int stride_row; // in Q8_0 blocks + int gate_stride_row; // in blocks + int blocks_begin; + int glu_op; + float glu_limit; +}; + +struct mmvq_group_args_dev { + mmvq_group_seg_dev seg[MMVQ_GROUP_MAX]; + int nseg; +}; + +template +__launch_bounds__(nwarps * ggml_cuda_get_physical_warp_size(), 1) +static __global__ void mul_mat_vec_fq_group( + const mmvq_group_args_dev args, const float * GGML_CUDA_RESTRICT y_ptr, const int ncols_x, + const float y_scale, const float y_bias, const int y_op) { + constexpr int qk = QK8_0; + constexpr int qi = QI8_0; + constexpr int vdr = VDR_Q8_0_Q8_1_MMVQ; + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + constexpr int blocks_per_iter = vdr * warp_size / qi; + + extern __shared__ char mmvq_fq_smem[]; + block_q8_1 * y_q8 = (block_q8_1 *) mmvq_fq_smem; + + const int lane = threadIdx.x; + const int tid = warp_size*threadIdx.y + lane; + + // segment lookup, uniform per block + int s = 0; +#pragma unroll + for (int k = 1; k < MMVQ_GROUP_MAX; ++k) { + if (k < args.nseg && (int) blockIdx.x >= args.seg[k].blocks_begin) { + s = k; + } + } + const mmvq_group_seg_dev seg = args.seg[s]; + + const int row = ((int) blockIdx.x - seg.blocks_begin)*nwarps + threadIdx.y; + const bool row_ok = row < seg.nrows; + + // Q8_0 path: identical to mul_mat_vec_q_fq + const int row_safe = row_ok ? row : seg.nrows - 1; + const bool use_gate = seg.gate != nullptr; + + const int blocks_per_row_x = ncols_x / qk; + const int kqs = vdr * (lane % (qi/vdr)); + const int kbx0 = lane / (qi/vdr); + + const block_q8_0 * x = (const block_q8_0 *) seg.vx + (size_t) row_safe*seg.stride_row; + const block_q8_0 * xg = use_gate ? (const block_q8_0 *) seg.gate + (size_t) row_safe*seg.gate_stride_row : nullptr; + + mmvq_fq_q8_0_chunk cx; + mmvq_fq_q8_0_chunk cg; + mmvq_fq_q8_0_load(cx, x, kbx0, kqs, blocks_per_row_x); + if (use_gate) { + mmvq_fq_q8_0_load(cg, xg, kbx0, kqs, blocks_per_row_x); + } + + { + const float * y = y_ptr; + const int nby = ncols_x / QK8_1; + for (int i = tid; i < 4*nby; i += nwarps*warp_size) { + const int ib = i >> 2; + const int sub = i & 3; + const float4 * yv = (const float4 *) (y + ib*QK8_1 + sub*8); + const float4 v0 = yv[0]; + const float4 v1 = yv[1]; + float v[8] = {v0.x, v0.y, v0.z, v0.w, v1.x, v1.y, v1.z, v1.w}; + if (y_op != 0) { +#pragma unroll + for (int k = 0; k < 8; ++k) { + const float t = y_scale * v[k] + y_bias; + v[k] = y_op == 1 ? ggml_cuda_op_silu_single(t) : 1.0f / (1.0f + expf(-t)); + } + } + + float amax = fabsf(v[0]); +#pragma unroll + for (int k = 1; k < 8; ++k) { + amax = fmaxf(amax, fabsf(v[k])); + } + amax = fmaxf(amax, __shfl_xor_sync(0xffffffff, amax, 1, warp_size)); + amax = fmaxf(amax, __shfl_xor_sync(0xffffffff, amax, 2, warp_size)); + + const float d = amax / 127.0f; + int q[8]; +#pragma unroll + for (int k = 0; k < 8; ++k) { + const int8_t qk8 = amax == 0.0f ? 0 : roundf(v[k] / d); + q[k] = (int) qk8 & 0xff; + } + int * qs = (int *) y_q8[ib].qs; + qs[2*sub + 0] = q[0] | (q[1] << 8) | (q[2] << 16) | (q[3] << 24); + qs[2*sub + 1] = q[4] | (q[5] << 8) | (q[6] << 16) | (q[7] << 24); + if (sub == 0) { + y_q8[ib].ds = make_half2(d, 0.0f); + } + } + } + __syncthreads(); + + if (!row_ok) { + return; + } + + float tmp = 0.0f; + float tmp_gate = 0.0f; + + const block_q8_1 * y = y_q8; + constexpr int chunk_blocks = pf*blocks_per_iter; + for (int c0 = kbx0; c0 < blocks_per_row_x; c0 += chunk_blocks) { + mmvq_fq_q8_0_chunk nx; + mmvq_fq_q8_0_chunk ng; + const int c1 = c0 + chunk_blocks; + if (c1 < blocks_per_row_x) { + mmvq_fq_q8_0_load(nx, x, c1, kqs, blocks_per_row_x); + if (use_gate) { + mmvq_fq_q8_0_load(ng, xg, c1, kqs, blocks_per_row_x); + } + } + mmvq_fq_q8_0_dot(tmp, cx, y, c0, kqs, blocks_per_row_x); + if (use_gate) { + mmvq_fq_q8_0_dot(tmp_gate, cg, y, c0, kqs, blocks_per_row_x); + } + if (c1 < blocks_per_row_x) { + cx = nx; + if (use_gate) { + cg = ng; + } + } + } + + tmp = warp_reduce_sum(tmp); + if (use_gate) { + tmp_gate = warp_reduce_sum(tmp_gate); + } + + if (lane == 0) { + float result = tmp; + if (use_gate) { + switch (seg.glu_op) { + case GGML_GLU_OP_SWIGLU: + result *= ggml_cuda_op_silu_single(tmp_gate); + break; + case GGML_GLU_OP_GEGLU: + result *= ggml_cuda_op_gelu_single(tmp_gate); + break; + case GGML_GLU_OP_SWIGLU_OAI: + result = ggml_cuda_op_swiglu_oai_single(tmp_gate, result); + break; + case GGML_GLU_OP_SWIGLU_CLAMP: + result = ggml_cuda_op_swiglu_clamp_single(tmp_gate, result, seg.glu_limit); + break; + default: + result = result * tmp_gate; + break; + } + } + seg.dst[row] = result; + } +} + +bool ggml_cuda_mmv_group_seg_ok(const ggml_tensor * w, const ggml_tensor * y) { + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + if (!GGML_CUDA_CC_IS_RDNA3_5(cc)) { + return false; + } + if (y->type != GGML_TYPE_F32 || y->ne[1] != 1 || y->ne[2] != 1 || y->ne[3] != 1 || !ggml_is_contiguous(y) || + (uintptr_t) y->data % 16 != 0 || y->ne[0] != w->ne[0]) { + return false; + } + if (w->ne[2] != 1 || w->ne[3] != 1 || w->nb[0] != ggml_type_size(w->type) || w->ne[1] > INT32_MAX || + (w->buffer && ggml_backend_buffer_get_usage(w->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE)) { + return false; + } + const int64_t ncols = w->ne[0]; + if (w->type == GGML_TYPE_Q8_0) { + return ncols % QK8_1 == 0 && (ncols / QK8_1) * sizeof(block_q8_1) <= 16384 && w->nb[1] % sizeof(block_q8_0) == 0; + } + return false; +} + +void ggml_cuda_mmv_group(ggml_backend_cuda_context & ctx, const ggml_tensor * y, const ggml_cuda_mmv_group_seg * segs, const int nseg) { + GGML_ASSERT(nseg >= 1 && nseg <= MMVQ_GROUP_MAX); + const int device = ggml_cuda_get_device(); + const int warp_size = ggml_cuda_info().devices[device].warp_size; + const int nsm = ggml_cuda_info().devices[device].nsm; + GGML_ASSERT(warp_size == 32); + + const int ncols = y->ne[0]; + + int64_t total_q8_rows = 0; + int64_t max_rows = 0; + for (int i = 0; i < nseg; ++i) { + GGML_ASSERT(ggml_cuda_mmv_group_seg_ok(segs[i].w, y)); + GGML_ASSERT(segs[i].dst->type == GGML_TYPE_F32 && ggml_is_contiguous(segs[i].dst) && ggml_nelements(segs[i].dst) == segs[i].w->ne[1]); + if (segs[i].gate) { + GGML_ASSERT(segs[i].w->type == GGML_TYPE_Q8_0 && segs[i].gate->type == GGML_TYPE_Q8_0 && + ggml_are_same_shape(segs[i].w, segs[i].gate) && segs[i].gate->nb[0] == sizeof(block_q8_0) && + segs[i].gate->nb[1] % sizeof(block_q8_0) == 0); + } + total_q8_rows += segs[i].w->ne[1]; + max_rows = std::max(max_rows, segs[i].w->ne[1]); + } + + // same block-size heuristic as mul_mat_vec_q_fq_launch, applied to the largest segment + int nwarps = MMVQ_FQ_NWARPS; + while (nwarps > 4 && ((max_rows + nwarps - 1) / nwarps) < 2 * nsm) { + nwarps /= 2; + } + + mmvq_group_args_dev args{}; + args.nseg = nseg; + int nblocks = 0; + for (int i = 0; i < nseg; ++i) { + mmvq_group_seg_dev & d = args.seg[i]; + const ggml_tensor * w = segs[i].w; + d.vx = w->data; + d.gate = segs[i].gate ? segs[i].gate->data : nullptr; + d.dst = (float *) segs[i].dst->data; + d.nrows = (int) w->ne[1]; + d.stride_row = (int) (w->nb[1] / ggml_type_size(w->type)); + d.gate_stride_row = segs[i].gate ? (int) (segs[i].gate->nb[1] / sizeof(block_q8_0)) : 0; + d.blocks_begin = nblocks; + d.glu_op = (int) segs[i].glu_op; + d.glu_limit = segs[i].glu_limit; + nblocks += (d.nrows + nwarps - 1) / nwarps; + } + + const dim3 block_nums(nblocks, 1, 1); + const dim3 block_dims(warp_size, nwarps, 1); + const int nbytes_shared = total_q8_rows > 0 ? (ncols / QK8_1) * sizeof(block_q8_1) : 0; + const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, nbytes_shared, ctx.stream()); + + constexpr int blocks_per_iter = VDR_Q8_0_Q8_1_MMVQ * 32 / QI8_0; + const bool long_rows = (ncols / QK8_0) >= 16 * blocks_per_iter && total_q8_rows <= 1024; + + auto launch = [&](auto kernel) { + ggml_cuda_kernel_launch(kernel, launch_params, args, (const float *) y->data, ncols, 1.0f, 0.0f, 0); + }; + + switch (nwarps) { + case 16: launch(mul_mat_vec_fq_group<16, MMVQ_FQ_PF>); break; + case 8: + if (long_rows) { launch(mul_mat_vec_fq_group<8, MMVQ_FQ_PF_LONG>); } + else { launch(mul_mat_vec_fq_group<8, MMVQ_FQ_PF>); } + break; + case 4: + if (long_rows) { launch(mul_mat_vec_fq_group<4, MMVQ_FQ_PF_LONG>); } + else { launch(mul_mat_vec_fq_group<4, MMVQ_FQ_PF>); } + break; + default: + GGML_ABORT("fatal error"); + } +} + +// Dedicated MoE multi-token kernel. +// Grid: (ceil(nrows_x / c_rows_per_block), nchannels_dst) +// Block: (warp_size, ncols_dst) - each warp handles one token independently. +// No shared memory reduction needed since each warp works alone. +template +__launch_bounds__(get_mmvq_mmid_max_batch_for_device()*ggml_cuda_get_physical_warp_size(), 1) +static __global__ void mul_mat_vec_q_moe( + const void * vx_ptr, const void * vy_ptr, const int32_t * ids_ptr, const ggml_cuda_mm_fusion_args_device fusion, + float * dst_ptr, + const uint32_t ncols_x, const uint3 nchannels_y, const uint32_t nrows_x, + const uint32_t stride_row_x, const uint32_t stride_col_y, const uint32_t stride_col_dst, + const uint32_t stride_channel_x, const uint32_t stride_channel_y, const uint32_t stride_channel_dst, + const uint32_t ncols_dst, const uint32_t ids_stride) { + const void * GGML_CUDA_RESTRICT vx = vx_ptr; + const void * GGML_CUDA_RESTRICT vy = vy_ptr; + const int32_t * GGML_CUDA_RESTRICT ids = ids_ptr; + float * GGML_CUDA_RESTRICT dst = dst_ptr; + + constexpr int qk = ggml_cuda_type_traits::qk; + constexpr int qi = ggml_cuda_type_traits::qi; + constexpr int vdr = get_vdr_mmvq(type); + constexpr int warp_size = ggml_cuda_get_physical_warp_size(); + + constexpr vec_dot_q_cuda_t vec_dot_q_cuda = get_vec_dot_q_cuda(type); + + // fuse gate, bias, scales, and glu_op into the up projection + bool use_gate = false; + const void * vgate = nullptr; + const float * x_bias = nullptr; + const float * gate_bias = nullptr; + const float * x_scale = nullptr; + const float * gate_scale = nullptr; + ggml_glu_op active_glu = GGML_GLU_OP_SWIGLU; float glu_limit = 0.0f; if constexpr (has_fusion) { @@ -1161,6 +2622,53 @@ static void mul_mat_vec_q_switch_ncols_dst( // static, else MSVC lambda capture breaks the constexpr uses below static constexpr int c_ncols_dst = 1; + if constexpr (type == GGML_TYPE_IQ3_S) { + constexpr int c_rows = MMVQ_IQ3_S_ROWS_PER_BLOCK; + const bool no_extra_fusion = fusion.x_bias == nullptr && fusion.gate_bias == nullptr && + fusion.x_scale == nullptr && fusion.gate_scale == nullptr; + const bool gate_ok = fusion.gate == nullptr || fusion.glu_op == GGML_GLU_OP_SWIGLU; + if (table_id == MMVQ_PARAMETERS_RDNA3_5 && ids != nullptr && ncols_x == 2560 && + no_extra_fusion && gate_ok && nsamples_dst == 1) { + static const int rows_env = getenv("GGML_IQ3_ROWS") ? atoi(getenv("GGML_IQ3_ROWS")) : c_rows; + auto launch_rows = [&](auto kernel, int rows) { + const dim3 block_nums((nrows_x + rows - 1) / rows, nchannels_dst, nsamples_dst); + const dim3 block_dims(warp_size, rows, 1); + const ggml_cuda_kernel_launch_params params(block_nums, block_dims, 0, stream); + ggml_cuda_kernel_launch(kernel, params, + vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, nrows_x, + stride_row_x, stride_col_dst, stride_channel_x, stride_channel_y, + stride_channel_dst, stride_sample_x, stride_sample_y, stride_sample_dst); + }; + static const int lds_env = getenv("GGML_IQ3_LDS") ? atoi(getenv("GGML_IQ3_LDS")) : 0; + static const int grid_env = getenv("GGML_IQ3_GRID") ? atoi(getenv("GGML_IQ3_GRID")) : 4; + if (fusion.gate != nullptr && grid_env > 0) { + switch (grid_env) { + case 1: launch_rows(mul_mat_vec_iq3_s_grid_rdna3_5, 1); break; + case 2: launch_rows(mul_mat_vec_iq3_s_grid_rdna3_5, 2); break; + case 8: launch_rows(mul_mat_vec_iq3_s_grid_rdna3_5, 8); break; + default: launch_rows(mul_mat_vec_iq3_s_grid_rdna3_5, 4); break; + } + } else if (fusion.gate != nullptr && lds_env > 0) { + switch (lds_env) { + case 1: launch_rows(mul_mat_vec_iq3_s_lds_rdna3_5, 1); break; + case 2: launch_rows(mul_mat_vec_iq3_s_lds_rdna3_5, 2); break; + case 8: launch_rows(mul_mat_vec_iq3_s_lds_rdna3_5, 8); break; + default: launch_rows(mul_mat_vec_iq3_s_lds_rdna3_5, 4); break; + } + } else if (fusion.gate != nullptr) { + switch (rows_env) { + case 2: launch_rows(mul_mat_vec_iq3_s_rows_rdna3_5, 2); break; + case 4: launch_rows(mul_mat_vec_iq3_s_rows_rdna3_5, 4); break; + case 8: launch_rows(mul_mat_vec_iq3_s_rows_rdna3_5, 8); break; + default: launch_rows(mul_mat_vec_iq3_s_rows_rdna3_5, 1); break; + } + } else { + launch_rows(mul_mat_vec_iq3_s_rows_rdna3_5, c_rows); + } + break; + } + } + // Tag types keep the flags compile-time, so __launch_bounds__ matches what is launched. const auto launch = [&](auto small_k_tag, auto halve_iters_tag) { constexpr bool c_small_k = decltype(small_k_tag)::value; @@ -1205,6 +2713,35 @@ static void mul_mat_vec_q_switch_ncols_dst( dims.first, dims.second, 0, ids_stride, stream); } break; case 4: { + if constexpr (is_rdna3_5_q4_columns_type(type)) { + const bool no_fusion = fusion.gate == nullptr && fusion.x_bias == nullptr && fusion.gate_bias == nullptr && + fusion.x_scale == nullptr && fusion.gate_scale == nullptr; + if (table_id == MMVQ_PARAMETERS_RDNA3_5 && ids == nullptr && no_fusion) { + constexpr int c_nwarps = calc_nwarps(type, 1, MMVQ_PARAMETERS_RDNA3_5); + const dim3 block_nums(nrows_x, nchannels_dst, nsamples_dst); + const dim3 block_dims(warp_size, c_nwarps, 1); + const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream); + ggml_cuda_kernel_launch(mul_mat_vec_q4_columns_rdna3_5, launch_params, + vx, vy, dst, ncols_x, stride_row_x, stride_col_y, stride_col_dst, + channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, + sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst); + break; + } + } + if constexpr (type == GGML_TYPE_Q8_0) { + const bool no_fusion = fusion.gate == nullptr && fusion.x_bias == nullptr && fusion.gate_bias == nullptr && + fusion.x_scale == nullptr && fusion.gate_scale == nullptr; + if (table_id == MMVQ_PARAMETERS_RDNA3_5 && ids == nullptr && no_fusion) { + const dim3 block_nums(nrows_x, nchannels_dst, nsamples_dst); + const dim3 block_dims(warp_size, 1, 1); + const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params(block_nums, block_dims, 0, stream); + ggml_cuda_kernel_launch(mul_mat_vec_q4_columns, launch_params, + vx, vy, dst, ncols_x, stride_row_x, stride_col_y, stride_col_dst, + channel_ratio_fd, stride_channel_x, stride_channel_y, stride_channel_dst, + sample_ratio_fd, stride_sample_x, stride_sample_y, stride_sample_dst); + break; + } + } constexpr int c_ncols_dst = 4; std::pair dims = calc_launch_params(c_ncols_dst, nrows_x, nchannels_dst, nsamples_dst, warp_size, table_id); mul_mat_vec_q_switch_fusion(vx, vy, ids, fusion, dst, ncols_x, nchannels_y_fd, stride_row_x, stride_col_y, stride_col_dst, @@ -1402,6 +2939,285 @@ static void mul_mat_vec_q_switch_type( } } +// RDNA3.5 token-generation fast path: quantize the activations inside the matvec kernel. Optionally applies an +// elementwise prologue (scale -> silu/sigmoid) to the activations first, so that the two small kernels in front +// of the matvec disappear as well. Returns false when the shape/type is not served by this path. +static bool mul_mat_vec_q_fq_try( + ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, + const ggml_cuda_mm_fusion_args_host * fusion, const float y_scale, const float y_bias, const int y_op, const bool launch) { + GGML_TENSOR_BINARY_OP_LOCALS; + + const size_t ts_src0 = ggml_type_size(src0->type); + const size_t ts_src1 = ggml_type_size(src1->type); + const size_t ts_dst = ggml_type_size(dst->type); + + if (src1->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32 || nb00 != ts_src0 || nb10 != ts_src1 || nb0 != ts_dst) { + return false; + } + + const int cc = ggml_cuda_info().devices[ggml_cuda_get_device()].cc; + const int64_t ncols_dst_fq = ids ? ne2 : ne1; + const bool aligned = (uintptr_t) src1->data % 16 == 0 && nb11 % 16 == 0 && nb12 % 16 == 0 && nb13 % 16 == 0; + const bool fusion_ok = !fusion || + (fusion->x_scale == nullptr && fusion->gate_scale == nullptr && + (!ids || (fusion->x_bias == nullptr && fusion->gate_bias == nullptr))); + if (!(GGML_CUDA_CC_IS_RDNA3_5(cc) && ncols_dst_fq == 1 && mmvq_fq_type_ok(src0->type) && aligned && fusion_ok && + ne10 % QK8_1 == 0 && (ne10 / QK8_1) * sizeof(block_q8_1) <= 16384)) { + return false; + } + if (!launch) { + return true; + } + + cudaStream_t stream = ctx.stream(); + + const float * src1_d = (const float *) src1->data; + const int32_t * ids_d = ids ? (const int32_t *) ids->data : nullptr; + float * dst_d = (float *) dst->data; + + ggml_cuda_mm_fusion_args_device fusion_local{}; + if (fusion) { + GGML_ASSERT(!ids || (fusion->x_bias == nullptr && fusion->gate_bias == nullptr)); + if (fusion->x_bias) { + fusion_local.x_bias = fusion->x_bias->data; + } + if (fusion->gate) { + fusion_local.gate = fusion->gate->data; + } + if (fusion->gate_bias) { + fusion_local.gate_bias = fusion->gate_bias->data; + } + fusion_local.glu_op = fusion->glu_op; + fusion_local.glu_limit = fusion->glu_limit; + fusion_local.y_gate = fusion->y_gate ? fusion->y_gate->data : nullptr; + fusion_local.y_norm_w = fusion->y_norm_w ? fusion->y_norm_w->data : nullptr; + fusion_local.y_eps = fusion->y_eps; + } + + // If src0 is a temporary compute buffer, clear any potential padding. + if (ggml_backend_buffer_get_usage(src0->buffer) == GGML_BACKEND_BUFFER_USAGE_COMPUTE) { + const size_t size_data = ggml_nbytes(src0); + const size_t size_alloc = ggml_backend_buffer_get_alloc_size(src0->buffer, src0); + if (size_alloc > size_data) { + GGML_ASSERT(ggml_is_contiguously_allocated(src0)); + GGML_ASSERT(!src0->view_src); + CUDA_CHECK(cudaMemsetAsync((char *) src0->data + size_data, 0, size_alloc - size_data, stream)); + } + } + + const int64_t s01 = src0->nb[1] / ts_src0; + const int64_t s02 = src0->nb[2] / ts_src0; + const int64_t s03 = src0->nb[3] / ts_src0; + const int64_t s11 = src1->nb[1] / ts_src1; + const int64_t s12 = src1->nb[2] / ts_src1; + const int64_t s13 = src1->nb[3] / ts_src1; + const int64_t s1 = dst->nb[1] / ts_dst; + const int64_t s2 = dst->nb[2] / ts_dst; + const int64_t s3 = dst->nb[3] / ts_dst; + + // For MUL_MAT_ID the memory layout is different than for MUL_MAT: + const int64_t nchannels_y = ids ? ne11 : ne12; + const int64_t nchannels_dst = ids ? ne1 : ne2; + const int64_t stride_col_dst = ids ? s2 : s1; + const int64_t stride_channel_dst = ids ? s1 : s2; + const int64_t stride_channel_y = ids ? s11 : s12; + + mul_mat_vec_q_fq_switch_type( + src0->data, src0->type, src1_d, ids_d, fusion_local, dst_d, ne00, ne01, s01, stride_col_dst, + ne02, nchannels_y, nchannels_dst, s02, stride_channel_y, stride_channel_dst, + ne03, ne3, s03, s13, s3, stream, y_scale, y_bias, y_op); + return true; +} + +bool ggml_cuda_mul_mat_vec_q_fq_prologue_ok(ggml_backend_cuda_context & ctx, const ggml_tensor * dst, const ggml_tensor * y) { + if (dst->op != GGML_OP_MUL_MAT || !ggml_is_quantized(dst->src[0]->type) || y->type != GGML_TYPE_F32 || + !ggml_is_contiguous(y) || !ggml_are_same_shape(y, dst->src[1])) { + return false; + } + return mul_mat_vec_q_fq_try(ctx, dst->src[0], y, nullptr, const_cast(dst), nullptr, 1.0f, 0.0f, 0, /*launch=*/false); +} + +void ggml_cuda_mul_mat_vec_q_fq_prologue(ggml_backend_cuda_context & ctx, ggml_tensor * dst, const ggml_tensor * y, + const float y_scale, const float y_bias, const int y_op) { + const bool ok = mul_mat_vec_q_fq_try(ctx, dst->src[0], y, nullptr, dst, nullptr, y_scale, y_bias, y_op, /*launch=*/true); + GGML_ASSERT(ok); +} + +bool ggml_cuda_mul_mat_vec_q_fq_gdn_gate_ok(ggml_backend_cuda_context & ctx, const ggml_tensor * dst, const ggml_tensor * z) { + const ggml_tensor * y = dst->src[1]; + if (dst->op != GGML_OP_MUL_MAT || !ggml_is_quantized(dst->src[0]->type) || y->type != GGML_TYPE_F32 || + !ggml_is_contiguous(y) || ggml_nrows(y) != 1 || y->ne[0] % 128 != 0 || + z->type != GGML_TYPE_F32 || !ggml_is_contiguous(z) || ggml_nelements(z) != y->ne[0] || (uintptr_t) z->data % 16 != 0) { + return false; + } + // every block reads z while the waves write dst + const char * d0 = (const char *) dst->data; + const char * d1 = d0 + ggml_backend_buft_get_alloc_size(dst->buffer->buft, dst); + const char * z0 = (const char *) z->data; + const char * z1 = z0 + ggml_backend_buft_get_alloc_size(z->buffer->buft, z); + if (d0 < z1 && z0 < d1) { + return false; + } + return mul_mat_vec_q_fq_try(ctx, dst->src[0], y, nullptr, const_cast(dst), nullptr, 1.0f, 0.0f, 0, /*launch=*/false); +} + +void ggml_cuda_mul_mat_vec_q_fq_gdn_gate(ggml_backend_cuda_context & ctx, ggml_tensor * dst, const float * attn, + const ggml_tensor * z, const ggml_tensor * norm_w, const float eps) { + GGML_ASSERT((uintptr_t) attn % 16 == 0 && (uintptr_t) norm_w->data % 16 == 0 && ggml_nelements(norm_w) == 128); + ggml_tensor y = *dst->src[1]; // same shape/strides as the (elided) gated norm output, data = pre-norm scratch + y.data = (void *) attn; + ggml_cuda_mm_fusion_args_host fusion{}; + fusion.y_gate = z; + fusion.y_norm_w = norm_w; + fusion.y_eps = eps; + const bool ok = mul_mat_vec_q_fq_try(ctx, dst->src[0], &y, nullptr, dst, &fusion, 1.0f, 0.0f, 3, /*launch=*/true); + GGML_ASSERT(ok); +} + +#if defined(__HIP_PLATFORM_AMD__) +static __device__ __forceinline__ float mmvq_hc_mul_rn(const float a, const float b) { + float result; + asm("v_mul_f32_e32 %0, %1, %2" : "=v"(result) : "v"(a), "v"(b)); + return result; +} + +static __device__ __forceinline__ float mmvq_hc_add_rn(const float a, const float b) { + float result; + asm("v_add_f32_e32 %0, %1, %2" : "=v"(result) : "v"(a), "v"(b)); + return result; +} +#else +static __device__ __forceinline__ float mmvq_hc_mul_rn(const float a, const float b) { return __fmul_rn(a, b); } +static __device__ __forceinline__ float mmvq_hc_add_rn(const float a, const float b) { return __fadd_rn(a, b); } +#endif + +template +__launch_bounds__(32, 8) +static __global__ void mul_mat_id_iq4_nl_weighted_rdna3_5( + const void * vx_ptr, const block_q8_1 * y, const int32_t * ids, const float * weights, float * dst, + const int nrows, const int blocks_per_row, const int stride_row_x, const int stride_channel_x, + const int stride_y) { + constexpr int qi = ggml_cuda_type_traits::qi; + constexpr int vdr = VDR_Q4_0_Q8_1_MMVQ; + constexpr int warp_size = 32; + constexpr int blocks_per_iter = vdr * warp_size / qi; + + const int lane = threadIdx.x; + const int row = blockIdx.x; + if (row >= nrows) { + return; + } + + float result = 0.0f; +#pragma unroll + for (int ex = 0; ex < n_expert_used; ++ex) { + const int channel = ids[ex]; + const int x_off = channel * stride_channel_x + row * stride_row_x; + float sum = 0.0f; + for (int kbx = lane / (qi/vdr); kbx < blocks_per_row; kbx += blocks_per_iter) { + const int kqs = vdr * (lane % (qi/vdr)); + sum += vec_dot_iq4_nl_q8_1(vx_ptr, y + ex * stride_y + kbx, x_off + kbx, kqs); + } + sum = warp_reduce_sum(sum); + const float term = mmvq_hc_mul_rn(sum, weights[ex]); + result = ex == 0 ? term : mmvq_hc_add_rn(result, term); + } + + if (lane == 0) { + dst[row] = result; + } +} + +template +__launch_bounds__(32, 8) +static __global__ void mul_mat_id_q8_0_weighted_rdna3_5( + const void * vx_ptr, const block_q8_1 * y, const int32_t * ids, const float * weights, float * dst, + const int nrows, const int blocks_per_row, const int stride_row_x, const int stride_channel_x, + const int stride_y) { + constexpr int qi = QI8_0; + constexpr int vdr = VDR_Q8_0_Q8_1_MMVQ; + constexpr int warp_size = 32; + constexpr int blocks_per_iter = vdr * warp_size / qi; + + const int lane = threadIdx.x; + const int row = blockIdx.x; + if (row >= nrows) { + return; + } + + float result = 0.0f; +#pragma unroll + for (int ex = 0; ex < n_expert_used; ++ex) { + const int channel = ids[ex]; + const int x_off = channel * stride_channel_x + row * stride_row_x; + float sum = 0.0f; + for (int kbx = lane / (qi/vdr); kbx < blocks_per_row; kbx += blocks_per_iter) { + const int kqs = vdr * (lane % (qi/vdr)); + sum += vec_dot_q8_0_q8_1(vx_ptr, y + ex * stride_y + kbx, x_off + kbx, kqs); + } + sum = warp_reduce_sum(sum); + const float term = mmvq_hc_mul_rn(sum, weights[ex]); + result = ex == 0 ? term : mmvq_hc_add_rn(result, term); + } + + if (lane == 0) { + dst[row] = result; + } +} + +bool ggml_cuda_mul_mat_id_weighted_rdna3_5_ok( + const ggml_tensor * experts, const ggml_tensor * weights, const ggml_tensor * dst) { + if (experts->op != GGML_OP_MUL_MAT_ID || + (experts->src[0]->type != GGML_TYPE_IQ4_NL && experts->src[0]->type != GGML_TYPE_Q8_0) || + experts->src[1]->type != GGML_TYPE_F32 || experts->src[2]->type != GGML_TYPE_I32 || + experts->type != GGML_TYPE_F32 || weights->type != GGML_TYPE_F32 || dst->type != GGML_TYPE_F32) { + return false; + } + const ggml_tensor * w = experts->src[0]; + const ggml_tensor * y = experts->src[1]; + const ggml_tensor * ids = experts->src[2]; + const int64_t n_used = ids->ne[0]; + return GGML_CUDA_CC_IS_RDNA3_5(ggml_cuda_info().devices[ggml_cuda_get_device()].cc) && + n_used == 10 && w->ne[0] == 640 && w->ne[1] == 2560 && w->ne[2] == 512 && w->ne[3] == 1 && + y->ne[0] == 640 && ggml_nelements(y) == 640 * n_used && ggml_is_contiguous(y) && + ggml_nelements(ids) == n_used && ggml_is_contiguous(ids) && + weights->ne[0] == 1 && weights->ne[1] == n_used && ggml_nelements(weights) == n_used && ggml_is_contiguous(weights) && + ggml_nelements(experts) == 2560 * n_used && ggml_is_contiguous(experts) && + ggml_nelements(dst) == 2560 && ggml_is_contiguous(dst); +} + +void ggml_cuda_mul_mat_id_weighted_rdna3_5( + ggml_backend_cuda_context & ctx, const ggml_tensor * experts, const ggml_tensor * weights, ggml_tensor * dst) { + GGML_ASSERT(ggml_cuda_mul_mat_id_weighted_rdna3_5_ok(experts, weights, dst)); + const ggml_tensor * w = experts->src[0]; + const ggml_tensor * y = experts->src[1]; + const ggml_tensor * ids = experts->src[2]; + constexpr int n_used = 10; + constexpr int ncols = 640; + constexpr int nrows = 2560; + constexpr int nblocks = ncols / QK8_1; + + ggml_cuda_pool_alloc y_q8(ctx.pool(), n_used * nblocks); + quantize_row_q8_1_cuda( + (const float *) y->data, nullptr, y_q8.get(), w->type, + ncols, + y->nb[1] / sizeof(float), y->nb[2] / sizeof(float), y->nb[3] / sizeof(float), + y->ne[0], y->ne[1], y->ne[2], y->ne[3], ctx.stream()); + + const ggml_cuda_kernel_launch_params params(nrows, 32, 0, ctx.stream()); + if (w->type == GGML_TYPE_IQ4_NL) { + ggml_cuda_kernel_launch(mul_mat_id_iq4_nl_weighted_rdna3_5, params, + w->data, y_q8.get(), (const int32_t *) ids->data, (const float *) weights->data, (float *) dst->data, + nrows, nblocks, (int) (w->nb[1] / ggml_type_size(w->type)), + (int) (w->nb[2] / ggml_type_size(w->type)), nblocks); + } else { + ggml_cuda_kernel_launch(mul_mat_id_q8_0_weighted_rdna3_5, params, + w->data, y_q8.get(), (const int32_t *) ids->data, (const float *) weights->data, (float *) dst->data, + nrows, nblocks, (int) (w->nb[1] / ggml_type_size(w->type)), + (int) (w->nb[2] / ggml_type_size(w->type)), nblocks); + } +} + void ggml_cuda_mul_mat_vec_q( ggml_backend_cuda_context & ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, const ggml_cuda_mm_fusion_args_host * fusion) { @@ -1481,6 +3297,10 @@ void ggml_cuda_mul_mat_vec_q( } } + if (mul_mat_vec_q_fq_try(ctx, src0, src1, ids, dst, fusion, 1.0f, 0.0f, 0, /*launch=*/true)) { + return; + } + const int64_t ne10_padded = GGML_PAD(ne10, MATRIX_ROW_PADDING); ggml_cuda_pool_alloc src1_q8_1(ctx.pool(), ne13*ne12 * ne11*ne10_padded * sizeof(block_q8_1)/QK8_1); { diff --git a/ggml/src/ggml-cuda/mmvq.cuh b/ggml/src/ggml-cuda/mmvq.cuh index 5605bf7a4e60..c2d1da92573d 100644 --- a/ggml/src/ggml-cuda/mmvq.cuh +++ b/ggml/src/ggml-cuda/mmvq.cuh @@ -16,3 +16,33 @@ void ggml_cuda_op_mul_mat_vec_q( const ggml_tensor * src0, const ggml_tensor * src1, ggml_tensor * dst, const char * src0_dd_i, const float * src1_ddf_i, const char * src1_ddq_i, float * dst_dd_i, const int64_t row_low, const int64_t row_high, const int64_t src1_ncols, const int64_t src1_padded_row_size, cudaStream_t stream); + +// RDNA3.5 fused-quantize matvec with an activation prologue: y' = op(y_scale*y + y_bias) (op 1 = silu, 2 = sigmoid) +// applied while the activations are quantized in-kernel. `y` replaces dst->src[1] (same shape, F32, contiguous). +bool ggml_cuda_mul_mat_vec_q_fq_prologue_ok(ggml_backend_cuda_context & ctx, const ggml_tensor * dst, const ggml_tensor * y); +void ggml_cuda_mul_mat_vec_q_fq_prologue(ggml_backend_cuda_context & ctx, ggml_tensor * dst, const ggml_tensor * y, + float y_scale, float y_bias, int y_op); + +// RDNA3.5 grouped matvec: independent single-column matvecs sharing the same F32 activation vector `y` +// (Q8_0 rows via the fused-quantize path, optionally with a Q8_0 gate + GLU epilogue) in one launch. +#define GGML_CUDA_MMV_GROUP_MAX 4 +struct ggml_cuda_mmv_group_seg { + const ggml_tensor * w = nullptr; // [ncols, nrows] Q8_0 or F32 + const ggml_tensor * gate = nullptr; // optional Q8_0 gate weights of the same shape (GLU fusion) + ggml_tensor * dst = nullptr; // F32 [nrows]; the GLU output when gate != nullptr + ggml_glu_op glu_op = GGML_GLU_OP_SWIGLU; + float glu_limit = 0.0f; +}; +bool ggml_cuda_mmv_group_seg_ok(const ggml_tensor * w, const ggml_tensor * y); +void ggml_cuda_mmv_group(ggml_backend_cuda_context & ctx, const ggml_tensor * y, const ggml_cuda_mmv_group_seg * segs, int nseg); + +// RDNA3.5 qwen4exp GDN output projection with the gated per-head norm in the prologue: +// y' = sigmoid(z) * (rms_norm_128(attn) * norm_w), then the fused-quantize matvec of dst->src[0] with y' +bool ggml_cuda_mul_mat_vec_q_fq_gdn_gate_ok(ggml_backend_cuda_context & ctx, const ggml_tensor * dst, const ggml_tensor * z); +void ggml_cuda_mul_mat_vec_q_fq_gdn_gate(ggml_backend_cuda_context & ctx, ggml_tensor * dst, const float * attn, + const ggml_tensor * z, const ggml_tensor * norm_w, float eps); + +// RDNA3.5 one-token IQ4_NL MoE down projection with the selected-expert weighted sum in the kernel epilogue. +bool ggml_cuda_mul_mat_id_weighted_rdna3_5_ok(const ggml_tensor * experts, const ggml_tensor * weights, const ggml_tensor * dst); +void ggml_cuda_mul_mat_id_weighted_rdna3_5( + ggml_backend_cuda_context & ctx, const ggml_tensor * experts, const ggml_tensor * weights, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/moe-weighted-reduction.cu b/ggml/src/ggml-cuda/moe-weighted-reduction.cu index 11ec58497f1e..a826c70bfcf6 100644 --- a/ggml/src/ggml-cuda/moe-weighted-reduction.cu +++ b/ggml/src/ggml-cuda/moe-weighted-reduction.cu @@ -1,4 +1,9 @@ #include "moe-weighted-reduction.cuh" +#include "mmb.cuh" +#include + +__device__ __forceinline__ float moe_bf2f(const uint16_t h) { return __uint_as_float(((uint32_t) h) << 16); } + static __global__ void moe_weighted_reduction_f32(const float * __restrict__ experts, const float * __restrict__ expert_scale, @@ -24,6 +29,115 @@ static __global__ void moe_weighted_reduction_f32(const float * __restrict__ exp dst[token * n_embd + col] = sum; } + +// float4 variant: 4 consecutive columns per thread, same per-element accumulation order as the scalar kernel (bit-identical) +static __global__ void moe_weighted_reduction_f32_v4(const float * __restrict__ experts, + const float * __restrict__ expert_scale, + const float * __restrict__ weights, + float * __restrict__ dst, + const int64_t n_embd, + const int n_expert_used) { + const int64_t token = blockIdx.x; + const int64_t col4 = ((int64_t) blockIdx.y * blockDim.x + threadIdx.x) * 4; + if (col4 >= n_embd) return; + const uint64_t first_row = (uint64_t) token * n_expert_used; + const float first_scale = expert_scale != nullptr ? expert_scale[first_row] : 1.0f; + const float w0 = weights[first_row]; + const float4 e0 = *(const float4 *)(experts + first_row * n_embd + col4); + float4 sum; sum.x = (e0.x * first_scale) * w0; sum.y = (e0.y * first_scale) * w0; sum.z = (e0.z * first_scale) * w0; sum.w = (e0.w * first_scale) * w0; + for (int expert = 1; expert < n_expert_used; ++expert) { + const uint64_t row = first_row + expert; + const float scale = expert_scale != nullptr ? expert_scale[row] : 1.0f; + const float w = weights[row]; + const float4 e = *(const float4 *)(experts + row * n_embd + col4); + sum.x += (e.x * scale) * w; sum.y += (e.y * scale) * w; sum.z += (e.z * scale) * w; sum.w += (e.w * scale) * w; + } + *(float4 *)(dst + token * n_embd + col4) = sum; +} + +// BF16 expert outputs: identical arithmetic, the inputs are the BF16-rounded down-GEMM results +static __global__ void moe_weighted_reduction_bf16_v4(const uint16_t * __restrict__ experts, + const float * __restrict__ expert_scale, + const float * __restrict__ weights, + float * __restrict__ dst, + const int64_t n_embd, + const int n_expert_used) { + const int64_t token = blockIdx.x; + const int64_t col4 = ((int64_t) blockIdx.y * blockDim.x + threadIdx.x) * 4; + if (col4 >= n_embd) return; + const uint64_t first_row = (uint64_t) token * n_expert_used; + const float first_scale = expert_scale != nullptr ? expert_scale[first_row] : 1.0f; + const float w0 = weights[first_row]; + const ushort4 h0 = *(const ushort4 *)(experts + first_row * n_embd + col4); + float4 sum; sum.x = (moe_bf2f(h0.x) * first_scale) * w0; sum.y = (moe_bf2f(h0.y) * first_scale) * w0; + sum.z = (moe_bf2f(h0.z) * first_scale) * w0; sum.w = (moe_bf2f(h0.w) * first_scale) * w0; + for (int expert = 1; expert < n_expert_used; ++expert) { + const uint64_t row = first_row + expert; + const float scale = expert_scale != nullptr ? expert_scale[row] : 1.0f; + const float w = weights[row]; + const ushort4 h = *(const ushort4 *)(experts + row * n_embd + col4); + sum.x += (moe_bf2f(h.x) * scale) * w; sum.y += (moe_bf2f(h.y) * scale) * w; + sum.z += (moe_bf2f(h.z) * scale) * w; sum.w += (moe_bf2f(h.w) * scale) * w; + } + *(float4 *)(dst + token * n_embd + col4) = sum; +} + +__device__ __forceinline__ uint16_t moe_f2bf(const float f) { uint32_t u = __float_as_uint(f); u += 0x7fffu + ((u >> 16) & 1u); return (uint16_t)(u >> 16); } + +// BF16 in, BF16 out: the only consumer is the fused HC combine, which reads block_out as BF16 +static __global__ void moe_weighted_reduction_bf16_v4_out(const uint16_t * __restrict__ experts, + const float * __restrict__ expert_scale, + const float * __restrict__ weights, + uint16_t * __restrict__ dst, + const float * __restrict__ merge, + const int64_t n_embd, const int n_expert_used) { + const int64_t token = blockIdx.x; + const int64_t col4 = ((int64_t) blockIdx.y * blockDim.x + threadIdx.x) * 4; + if (col4 >= n_embd) return; + const uint64_t first_row = (uint64_t) token * n_expert_used; + const float first_scale = expert_scale != nullptr ? expert_scale[first_row] : 1.0f; + const float w0 = weights[first_row]; + const ushort4 h0 = *(const ushort4 *)(experts + first_row * n_embd + col4); + float4 sum; sum.x = (moe_bf2f(h0.x) * first_scale) * w0; sum.y = (moe_bf2f(h0.y) * first_scale) * w0; + sum.z = (moe_bf2f(h0.z) * first_scale) * w0; sum.w = (moe_bf2f(h0.w) * first_scale) * w0; + for (int expert = 1; expert < n_expert_used; ++expert) { + const uint64_t row = first_row + expert; + const float scale = expert_scale != nullptr ? expert_scale[row] : 1.0f; + const float w = weights[row]; + const ushort4 h = *(const ushort4 *)(experts + row * n_embd + col4); + sum.x += (moe_bf2f(h.x) * scale) * w; sum.y += (moe_bf2f(h.y) * scale) * w; + sum.z += (moe_bf2f(h.z) * scale) * w; sum.w += (moe_bf2f(h.w) * scale) * w; + } + if (merge) { const float4 m = *(const float4 *)(merge + token * n_embd + col4); sum.x += m.x; sum.y += m.y; sum.z += m.z; sum.w += m.w; } + *(ushort4 *)(dst + token * n_embd + col4) = make_ushort4(moe_f2bf(sum.x), moe_f2bf(sum.y), moe_f2bf(sum.z), moe_f2bf(sum.w)); +} + +static __global__ void moe_weighted_reduction_f32in_bf16out_v4(const float * __restrict__ experts, + const float * __restrict__ expert_scale, + const float * __restrict__ weights, + uint16_t * __restrict__ dst, + const float * __restrict__ merge, + const int64_t n_embd, const int n_expert_used) { + const int64_t token = blockIdx.x; + const int64_t col4 = ((int64_t) blockIdx.y * blockDim.x + threadIdx.x) * 4; + if (col4 >= n_embd) return; + const uint64_t first_row = (uint64_t) token * n_expert_used; + const float first_scale = expert_scale != nullptr ? expert_scale[first_row] : 1.0f; + const float w0 = weights[first_row]; + const float4 e0 = *(const float4 *)(experts + first_row * n_embd + col4); + float4 sum; sum.x = (e0.x * first_scale) * w0; sum.y = (e0.y * first_scale) * w0; + sum.z = (e0.z * first_scale) * w0; sum.w = (e0.w * first_scale) * w0; + for (int expert = 1; expert < n_expert_used; ++expert) { + const uint64_t row = first_row + expert; + const float scale = expert_scale != nullptr ? expert_scale[row] : 1.0f; + const float w = weights[row]; + const float4 e = *(const float4 *)(experts + row * n_embd + col4); + sum.x += (e.x * scale) * w; sum.y += (e.y * scale) * w; sum.z += (e.z * scale) * w; sum.w += (e.w * scale) * w; + } + if (merge) { const float4 m = *(const float4 *)(merge + token * n_embd + col4); sum.x += m.x; sum.y += m.y; sum.z += m.z; sum.w += m.w; } + *(ushort4 *)(dst + token * n_embd + col4) = make_ushort4(moe_f2bf(sum.x), moe_f2bf(sum.y), moe_f2bf(sum.z), moe_f2bf(sum.w)); +} + static void launch_moe_weighted_reduction(const float * experts, const float * expert_scale, const float * weights, @@ -33,6 +147,12 @@ static void launch_moe_weighted_reduction(const float * experts, int n_expert_used, cudaStream_t stream) { constexpr int threads = 256; + static const bool use_v4 = true; + if (use_v4 && n_embd % 4 == 0 && ((uintptr_t) experts % 16) == 0 && ((uintptr_t) dst % 16) == 0) { + const dim3 blocks(n_tokens, (n_embd / 4 + threads - 1) / threads, 1); + moe_weighted_reduction_f32_v4<<>>(experts, expert_scale, weights, dst, n_embd, n_expert_used); + return; + } const dim3 blocks(n_tokens, (n_embd + threads - 1) / threads, 1); moe_weighted_reduction_f32 <<>>(experts, expert_scale, weights, dst, n_embd, n_expert_used); @@ -42,7 +162,8 @@ void ggml_cuda_op_moe_weighted_reduction(ggml_backend_cuda_context & ctx, const ggml_tensor * experts, const ggml_tensor * expert_scale, const ggml_tensor * weights, - ggml_tensor * dst) { + ggml_tensor * dst, + const ggml_tensor * merge) { GGML_ASSERT(experts->type == GGML_TYPE_F32); GGML_ASSERT(weights->type == GGML_TYPE_F32); GGML_ASSERT(expert_scale == nullptr || expert_scale->type == GGML_TYPE_F32); @@ -57,6 +178,32 @@ void ggml_cuda_op_moe_weighted_reduction(ggml_backend_cuda_context & ctx, const int64_t n_tokens = experts->ne[2] * experts->ne[3]; cudaStream_t stream = ctx.stream(); + const float * mrg = merge ? (const float *) merge->data : nullptr; + if (ggml_cuda_mmb_blk16() && ggml_cuda_mmb_is_bf16_only(ctx, dst)) { + GGML_ASSERT(n_embd % 4 == 0); + const bool ein = ggml_cuda_mmb_is_bf16_only(ctx, experts); + constexpr int threads = 256; + const dim3 blocks(n_tokens, (n_embd / 4 + threads - 1) / threads, 1); + if (ein) moe_weighted_reduction_bf16_v4_out<<>>((const uint16_t *) experts->data, + expert_scale ? (const float *) expert_scale->data : nullptr, (const float *) weights->data, + (uint16_t *) dst->data, mrg, n_embd, (int) n_expert_used); + else moe_weighted_reduction_f32in_bf16out_v4<<>>((const float *) experts->data, + expert_scale ? (const float *) expert_scale->data : nullptr, (const float *) weights->data, + (uint16_t *) dst->data, mrg, n_embd, (int) n_expert_used); + CUDA_CHECK(cudaGetLastError()); + return; + } + GGML_ASSERT(merge == nullptr && "shared-expert merge is only fused on the BF16 output path"); + if (ggml_cuda_mmb_down16() && ggml_cuda_mmb_is_bf16_only(ctx, experts)) { + GGML_ASSERT(n_embd % 4 == 0 && ((uintptr_t) experts->data % 16) == 0 && ((uintptr_t) dst->data % 16) == 0); + constexpr int threads = 256; + const dim3 blocks(n_tokens, (n_embd / 4 + threads - 1) / threads, 1); + moe_weighted_reduction_bf16_v4<<>>((const uint16_t *) experts->data, + expert_scale ? (const float *) expert_scale->data : nullptr, (const float *) weights->data, + (float *) dst->data, n_embd, (int) n_expert_used); + CUDA_CHECK(cudaGetLastError()); + return; + } launch_moe_weighted_reduction((const float *) experts->data, expert_scale ? (const float *) expert_scale->data : nullptr, (const float *) weights->data, diff --git a/ggml/src/ggml-cuda/moe-weighted-reduction.cuh b/ggml/src/ggml-cuda/moe-weighted-reduction.cuh index b72f947ab398..894c4b1a35ed 100644 --- a/ggml/src/ggml-cuda/moe-weighted-reduction.cuh +++ b/ggml/src/ggml-cuda/moe-weighted-reduction.cuh @@ -4,4 +4,5 @@ void ggml_cuda_op_moe_weighted_reduction(ggml_backend_cuda_context & ctx, const ggml_tensor * experts, const ggml_tensor * expert_scale, const ggml_tensor * weights, - ggml_tensor * dst); + ggml_tensor * dst, + const ggml_tensor * merge = nullptr); diff --git a/ggml/src/ggml-cuda/norm-gated.cu b/ggml/src/ggml-cuda/norm-gated.cu new file mode 100644 index 000000000000..6e09b7738c2a --- /dev/null +++ b/ggml/src/ggml-cuda/norm-gated.cu @@ -0,0 +1,138 @@ +#include "norm-gated.cuh" +#include "ggml-impl.h" + +// One wave per row, with the same reduction tree as the block-per-row norm. +static __device__ __forceinline__ float xor_tree(float v) { +#pragma unroll + for (int off = 16; off > 0; off >>= 1) { v += __shfl_xor_sync(0xffffffff, v, off, 32); } + return v; +} +template +static __global__ void __launch_bounds__(256) rms_rows_f32(const float * x, const float * w, const float * z, float * dst, + const int ncols, const uint3 row_div, const uint3 channel_div, const int64_t total_rows, + const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps, uint16_t * dst16) { + const int lane = threadIdx.x & 31; + const int64_t g = (int64_t) blockIdx.x * 8 + (threadIdx.x >> 5); + if (g >= total_rows) return; + uint32_t row = uint32_t(g), channel = 0, sample = 0; + if (uint64_t(row_div.z) != uint64_t(total_rows)) { + const uint2 rc = fast_div_modulo(row, row_div); + const uint2 cs = fast_div_modulo(rc.x, channel_div); + row = rc.y; channel = cs.y; sample = cs.x; + } + x += sample * stride_sample + channel * stride_channel + row * stride_row; + dst += g * ncols; if (GATE) z += g * ncols; + if (dst16) dst16 += g * ncols; + float part[8]; +#pragma unroll + for (int wv = 0; wv < 8; ++wv) { + const int col = 32 * wv + lane; + const float xi = col < ncols ? x[col] : 0.0f; + part[wv] = xor_tree(xi * xi); + } + float v = 0.0f; +#pragma unroll + for (int wv = 0; wv < 8; ++wv) { v = lane == wv ? part[wv] : v; } + const float tmp = xor_tree(v); + const float mean = tmp / ncols; + const float scale = rsqrtf(mean + eps); +#pragma unroll + for (int wv = 0; wv < 8; ++wv) { + const int col = 32 * wv + lane; + if (col < ncols) { + const float t = scale * x[col] * w[col]; + float v; + if (GATE) { const float s = 1.0f / (1.0f + expf(-z[col])); v = t * s; } else { v = t; } + dst[col] = v; + if (dst16) dst16[col] = ggml_cuda_f32_to_bf16_rne(v); + } + } +} + +static bool rows_shape_ok(const ggml_tensor * x, const ggml_tensor * w, const ggml_tensor * out) { + if (x->type != GGML_TYPE_F32 || w->type != GGML_TYPE_F32 || out->type != GGML_TYPE_F32) return false; + if (x->ne[0] > 256 || x->ne[0] % 32 != 0 || x->nb[0] != 4 || !ggml_is_contiguous(w) || !ggml_is_contiguous(out)) return false; + if (w->ne[0] != x->ne[0] || w->ne[1] != 1 || w->ne[2] != 1 || w->ne[3] != 1) return false; + const int64_t rows = x->ne[1] * x->ne[2] * x->ne[3]; + if (rows < 4096 || uint64_t(rows) > UINT32_MAX) return false; + return ggml_nelements(out) == ggml_nelements(x); +} + +static bool norm_overlaps(const ggml_tensor * a, const ggml_tensor * b) { + if (!a || !b || !a->data || !b->data) return false; + const uintptr_t av = uintptr_t(a->data), bv = uintptr_t(b->data); + return av <= bv ? bv - av < ggml_nbytes(a) : av - bv < ggml_nbytes(b); +} + +static bool norm_elementwise_alias(const ggml_tensor * input, const ggml_tensor * out) { + return !norm_overlaps(input, out) || (input->data == out->data && + ggml_is_contiguous(input) && ggml_are_same_shape(input, out)); +} + +static bool norm_needs_staging(const ggml_cuda_norm_gated_match & m) { + if (!norm_elementwise_alias(m.x, m.dst) || norm_overlaps(m.w, m.dst) || + (m.z && !norm_elementwise_alias(m.z, m.dst))) return true; + if (m.pre >= 0) { + const auto * projected = m.z->view_src ? m.z->view_src : m.z; + if (norm_overlaps(m.x, projected) || norm_overlaps(m.w, projected)) return true; + } + return false; +} + +int ggml_cuda_norm_gated_match_at(const ggml_cgraph * cgraph, int i, ggml_cuda_norm_gated_match & m) { + if (i + 3 >= cgraph->n_nodes) return 0; + const ggml_tensor * rms = cgraph->nodes[i], * mul = cgraph->nodes[i+1]; + if (rms->op != GGML_OP_RMS_NORM || mul->op != GGML_OP_MUL) return 0; + int k = i + 2; m.pre = -1; + while (k + 1 < cgraph->n_nodes && k < i + 6) { + const ggml_op op = cgraph->nodes[k]->op; + if (op == GGML_OP_VIEW || op == GGML_OP_RESHAPE || op == GGML_OP_PERMUTE || op == GGML_OP_TRANSPOSE) { ++k; continue; } + if (op == GGML_OP_MUL_MAT && m.pre < 0) { m.pre = k; ++k; continue; } + break; + } + if (k + 1 >= cgraph->n_nodes) return 0; + const ggml_tensor * sig = cgraph->nodes[k]; ggml_tensor * out = cgraph->nodes[k+1]; + if (sig->op != GGML_OP_UNARY || out->op != GGML_OP_MUL) return 0; + if (ggml_get_unary_op(sig) != GGML_UNARY_OP_SIGMOID) return 0; + const ggml_tensor * x = rms->src[0]; + const ggml_tensor * w = mul->src[0] == rms ? mul->src[1] : (mul->src[1] == rms ? mul->src[0] : nullptr); + if (!w) return 0; + const ggml_tensor * z = sig->src[0]; + if (m.pre >= 0) { const ggml_tensor * zr = z->view_src ? z->view_src : z; if (zr != cgraph->nodes[m.pre]) return 0; } + if (!((out->src[0] == mul && out->src[1] == sig) || (out->src[0] == sig && out->src[1] == mul))) return 0; + if (!ggml_node_has_n_uses(cgraph, i, 1) || !ggml_node_has_n_uses(cgraph, i+1, 1) || !ggml_node_has_n_uses(cgraph, k, 1)) return 0; + if (!rows_shape_ok(x, w, out) || z->type != GGML_TYPE_F32 || !ggml_is_contiguous(z) || !ggml_is_contiguous(mul)) return 0; + if (!ggml_are_same_shape(z, out) || !ggml_are_same_shape(mul, out) || !ggml_are_same_shape(rms, out)) return 0; + float eps; memcpy(&eps, rms->op_params, sizeof(float)); if (eps < 0.0f) return 0; + m.x = x; m.w = w; m.z = z; m.dst = out; m.eps = eps; + m.staged = norm_needs_staging(m); + return (k + 1) - i; +} + +// plain RMS_NORM -> MUL(w) on narrow rows +int ggml_cuda_norm_rows_match_at(const ggml_cgraph * cgraph, int i, ggml_cuda_norm_gated_match & m) { + if (i + 1 >= cgraph->n_nodes) return 0; + const ggml_tensor * rms = cgraph->nodes[i]; ggml_tensor * mul = cgraph->nodes[i+1]; + if (rms->op != GGML_OP_RMS_NORM || mul->op != GGML_OP_MUL) return 0; + const ggml_tensor * x = rms->src[0]; + const ggml_tensor * w = mul->src[0] == rms ? mul->src[1] : (mul->src[1] == rms ? mul->src[0] : nullptr); + if (!w || !ggml_node_has_n_uses(cgraph, i, 1)) return 0; + if (!rows_shape_ok(x, w, mul) || !ggml_are_same_shape(rms, mul)) return 0; + float eps; memcpy(&eps, rms->op_params, sizeof(float)); if (eps < 0.0f) return 0; + m.x = x; m.w = w; m.z = nullptr; m.dst = mul; m.eps = eps; m.pre = -1; + m.staged = norm_needs_staging(m); + return 1; +} + +void ggml_cuda_op_norm_gated(ggml_backend_cuda_context & ctx, const ggml_cuda_norm_gated_match & m) { + const ggml_tensor * x = m.x; + const int64_t total = x->ne[1] * x->ne[2] * x->ne[3]; + const uint3 row_div = init_fastdiv_values(uint32_t(x->ne[1])); + const uint3 channel_div = init_fastdiv_values(uint32_t(x->ne[2])); + const dim3 grid((unsigned) ((total + 7) / 8)), block(256); + const ggml_cuda_kernel_launch_params lp(grid, block, 0, ctx.stream()); + if (m.z) ggml_cuda_kernel_launch(rms_rows_f32, lp, (const float *) x->data, (const float *) m.w->data, (const float *) m.z->data, (float *) m.dst->data, + (int) x->ne[0], row_div, channel_div, total, (int64_t) (x->nb[1] / 4), (int64_t) (x->nb[2] / 4), (int64_t) (x->nb[3] / 4), m.eps, m.dst16); + else ggml_cuda_kernel_launch(rms_rows_f32, lp, (const float *) x->data, (const float *) m.w->data, (const float *) nullptr, (float *) m.dst->data, + (int) x->ne[0], row_div, channel_div, total, (int64_t) (x->nb[1] / 4), (int64_t) (x->nb[2] / 4), (int64_t) (x->nb[3] / 4), m.eps, m.dst16); +} diff --git a/ggml/src/ggml-cuda/norm-gated.cuh b/ggml/src/ggml-cuda/norm-gated.cuh new file mode 100644 index 000000000000..b94c5357bea9 --- /dev/null +++ b/ggml/src/ggml-cuda/norm-gated.cuh @@ -0,0 +1,15 @@ +#pragma once +#include "common.cuh" +struct ggml_cuda_norm_gated_match { + const ggml_tensor * x; + const ggml_tensor * w; + const ggml_tensor * z; + ggml_tensor * dst; + float eps; + int pre = -1; + bool staged = false; + uint16_t * dst16 = nullptr; // optional BF16 copy of dst for the MMB GEMM that reads it next +}; +int ggml_cuda_norm_gated_match_at(const ggml_cgraph * cgraph, int i, ggml_cuda_norm_gated_match & m); +void ggml_cuda_op_norm_gated(ggml_backend_cuda_context & ctx, const ggml_cuda_norm_gated_match & m); +int ggml_cuda_norm_rows_match_at(const ggml_cgraph * cgraph, int i, ggml_cuda_norm_gated_match & m); diff --git a/ggml/src/ggml-cuda/norm.cu b/ggml/src/ggml-cuda/norm.cu index c3758cd50cfe..50f7367046db 100644 --- a/ggml/src/ggml-cuda/norm.cu +++ b/ggml/src/ggml-cuda/norm.cu @@ -277,6 +277,44 @@ static __global__ void l2_norm_f32( } } +static __global__ void l2_norm_dual_f32_s128( + const float * x_q, float * dst_q, const int64_t stride_q_row, const int64_t stride_q_channel, + const int64_t stride_q_sample, const float eps_q, + const float * x_k, float * dst_k, const int64_t stride_k_row, const int64_t stride_k_channel, + const int64_t stride_k_sample, const float eps_k) { + const int nrows = gridDim.x; + const int nchannels = gridDim.y / 2; + const bool is_k = blockIdx.y >= nchannels; + const int row = blockIdx.x; + const int channel = blockIdx.y % nchannels; + const int sample = blockIdx.z; + const int tid = threadIdx.x; + + const float * x = is_k ? x_k : x_q; + float * dst = is_k ? dst_k : dst_q; + const int64_t stride_row = is_k ? stride_k_row : stride_q_row; + const int64_t stride_channel = is_k ? stride_k_channel : stride_q_channel; + const int64_t stride_sample = is_k ? stride_k_sample : stride_q_sample; + const float eps = is_k ? eps_k : eps_q; + + x += sample*stride_sample + channel*stride_channel + row*stride_row; + dst += ((sample*nchannels + channel)*nrows + row)*128; + + float tmp = 0.0f; + ggml_cuda_pdl_sync(); + for (int col = tid; col < 128; col += WARP_SIZE) { + const float xi = x[col]; + tmp += xi * xi; + } + tmp = warp_reduce_sum(tmp); + ggml_cuda_pdl_lc(); + + const float scale = rsqrtf(fmaxf(tmp, eps * eps)); + for (int col = tid; col < 128; col += WARP_SIZE) { + dst[col] = scale * x[col]; + } +} + static void norm_f32_cuda( const float * x, float * dst, const int ncols, const int nrows, const int nchannels, const int nsamples, const int64_t stride_row, const int64_t stride_channel, const int64_t stride_sample, const float eps, cudaStream_t stream) { @@ -360,7 +398,14 @@ static void rms_norm_mul_f32_cuda(const float * x, const uint3 mul_nrows_packed = init_fastdiv_values(mul_nrows); const uint3 mul_nchannels_packed = init_fastdiv_values(mul_nchannels); const uint3 mul_nsamples_packed = init_fastdiv_values(mul_nsamples); - if (ncols < 1024) { + if (ncols == 128) { + const dim3 block_dims(128, 1, 1); + const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float) : 0, stream}; + ggml_cuda_kernel_launch(rms_norm_f32<128, true>, launch_params, + x, dst, ncols, stride_row, stride_channel, stride_sample, eps, mul, mul_stride_row, mul_stride_channel, + mul_stride_sample, mul_ncols_packed, mul_nrows_packed, mul_nchannels_packed, mul_nsamples_packed, + nullptr, 0, 0, 0, make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0), make_uint3(0, 0, 0)); + } else if (ncols < 1024) { const dim3 block_dims(256, 1, 1); const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params{blocks_num, block_dims, block_dims.x > WARP_SIZE ? 32 * sizeof(float): 0, stream}; ggml_cuda_kernel_launch(rms_norm_f32<256, true>, launch_params, @@ -696,3 +741,28 @@ void ggml_cuda_op_l2_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { l2_norm_f32_cuda(src0_d, dst_d, ne00, ne01, ne02, ne03, s01, s02, s03, eps, stream); } + +void ggml_cuda_op_l2_norm_dual_s128(ggml_backend_cuda_context & ctx, ggml_tensor * dst_q, ggml_tensor * dst_k) { + const ggml_tensor * src_q = dst_q->src[0]; + const ggml_tensor * src_k = dst_k->src[0]; + + GGML_ASSERT(src_q->type == GGML_TYPE_F32 && src_k->type == GGML_TYPE_F32); + GGML_ASSERT(dst_q->type == GGML_TYPE_F32 && dst_k->type == GGML_TYPE_F32); + GGML_ASSERT(src_q->ne[0] == 128 && ggml_are_same_shape(src_q, src_k)); + GGML_ASSERT(src_q->nb[0] == sizeof(float) && src_k->nb[0] == sizeof(float)); + + float eps_q; + float eps_k; + memcpy(&eps_q, dst_q->op_params, sizeof(float)); + memcpy(&eps_k, dst_k->op_params, sizeof(float)); + GGML_ASSERT(eps_q >= 0.0f && eps_k >= 0.0f); + + const dim3 blocks_num(src_q->ne[1], 2*src_q->ne[2], src_q->ne[3]); + const dim3 block_dims(WARP_SIZE, 1, 1); + const ggml_cuda_kernel_launch_params launch_params = { blocks_num, block_dims, 0, ctx.stream() }; + ggml_cuda_kernel_launch(l2_norm_dual_f32_s128, launch_params, + (const float *) src_q->data, (float *) dst_q->data, + src_q->nb[1]/sizeof(float), src_q->nb[2]/sizeof(float), src_q->nb[3]/sizeof(float), eps_q, + (const float *) src_k->data, (float *) dst_k->data, + src_k->nb[1]/sizeof(float), src_k->nb[2]/sizeof(float), src_k->nb[3]/sizeof(float), eps_k); +} diff --git a/ggml/src/ggml-cuda/norm.cuh b/ggml/src/ggml-cuda/norm.cuh index a74f6376720a..c09d004905a9 100644 --- a/ggml/src/ggml-cuda/norm.cuh +++ b/ggml/src/ggml-cuda/norm.cuh @@ -16,3 +16,5 @@ void ggml_cuda_op_rms_norm_fused_add(ggml_backend_cuda_context & ctx, void ggml_cuda_op_rms_norm_back(ggml_backend_cuda_context & ctx, ggml_tensor * dst); void ggml_cuda_op_l2_norm(ggml_backend_cuda_context & ctx, ggml_tensor * dst); + +void ggml_cuda_op_l2_norm_dual_s128(ggml_backend_cuda_context & ctx, ggml_tensor * dst_q, ggml_tensor * dst_k); diff --git a/ggml/src/ggml-cuda/ple-conv.cu b/ggml/src/ggml-cuda/ple-conv.cu new file mode 100644 index 000000000000..30aa8b7793b2 --- /dev/null +++ b/ggml/src/ggml-cuda/ple-conv.cu @@ -0,0 +1,309 @@ +#include "ple-conv.cuh" +#include "conv-tail.cuh" +#include "unary.cuh" + +#if defined(__HIP_PLATFORM_AMD__) +// the unfused graph multiplies in one kernel and adds in another: keep the two roundings apart +static __device__ __forceinline__ float ple_mul_rn(const float a, const float b) { float r; asm("v_mul_f32_e32 %0, %1, %2" : "=v"(r) : "v"(a), "v"(b)); return r; } +static __device__ __forceinline__ float ple_add_rn(const float a, const float b) { float r; asm("v_add_f32_e32 %0, %1, %2" : "=v"(r) : "v"(a), "v"(b)); return r; } +#else +static __device__ __forceinline__ float ple_mul_rn(const float a, const float b) { return __fmul_rn(a, b); } +static __device__ __forceinline__ float ple_add_rn(const float a, const float b) { return __fadd_rn(a, b); } +#endif + +// columns [tail_from, T+H) of every channel row of the concat: what the history copies read +static __global__ void ple_concat_tail(const float * __restrict__ state, const float * __restrict__ x, float * __restrict__ out, + const int C, const int T, const int H, const int tail_from, const int row_stride) { + const int ncols = T + H - tail_from; + const int64_t idx = (int64_t) blockIdx.x * blockDim.x + threadIdx.x; + if (idx >= (int64_t) C * ncols) { + return; + } + const int c = (int) (idx / ncols), j = tail_from + (int) (idx % ncols); + out[(size_t) c * row_stride + j] = (j < H) ? state[c * H + j] : x[(size_t) (j - H) * C + c]; +} + +static __device__ __forceinline__ float ple_weight(const half w) { return __half2float(w); } +static __device__ __forceinline__ float ple_weight(const float w) { return w; } + +template +static __global__ void __launch_bounds__(256) ple_conv_kernel(const float * __restrict__ state, const float * __restrict__ x, + const TW * __restrict__ w, float * __restrict__ y, const int C, const int T) { + constexpr int H = (K - 1) * DIL, WIN = H + 1; + const int c = blockIdx.x * 256 + threadIdx.x, t0 = blockIdx.y * TT; + if (c >= C) { + return; + } + float wr[K]; +#pragma unroll + for (int k = 0; k < K; ++k) { + wr[k] = ple_weight(w[c * K + k]); + } + float win[WIN]; + auto ld = [&](int jp) -> float { + return (jp < H) ? state[c * H + jp] : ((jp - H) < T ? x[(size_t) (jp - H) * C + c] : 0.0f); + }; +#pragma unroll + for (int k = 0; k < WIN; ++k) { + win[k] = ld(t0 + k); + } + const int tend = min(T, t0 + TT); + for (int t = t0; t < tend; ++t) { + float s = ple_mul_rn(win[0], wr[0]); +#pragma unroll + for (int k = 1; k < K; ++k) { + s = ple_add_rn(s, ple_mul_rn(win[k * DIL], wr[k])); + } + y[(size_t) t * C + c] = ggml_cuda_op_silu_single(s); +#pragma unroll + for (int k = 0; k < WIN - 1; ++k) { + win[k] = win[k + 1]; + } + win[WIN - 1] = ld(t + WIN); + } +} + +// tap k: CONT(TRANSPOSE(VIEW of the concat at column k*dil, [T, C])) +static int ple_tap_index(const ggml_tensor * t, const ggml_tensor * cc, int64_t C, int64_t T, int64_t dil, int K) { + if (t->op != GGML_OP_CONT || !t->src[0] || t->src[0]->op != GGML_OP_TRANSPOSE || t->type != GGML_TYPE_F32 || + t->ne[0] != C || t->ne[1] != T || t->ne[2] != 1 || t->ne[3] != 1 || !ggml_is_contiguous(t)) { + return -1; + } + const ggml_tensor * v = t->src[0]->src[0]; + if (!v || v->op != GGML_OP_VIEW || v->view_src != cc || v->ne[0] != T || v->ne[1] != C || v->ne[2] != 1 || + v->nb[1] != cc->nb[1] || v->view_offs % (dil * sizeof(float)) != 0) { + return -1; + } + const int64_t k = (int64_t) (v->view_offs / (dil * sizeof(float))); + return (k >= 0 && k < K) ? (int) k : -1; +} + +// the weight operand of a tap MUL: reshape <- CONT <- VIEW of column k of the [K, C] weight, behind a cast(F32) when it is F16 +static const ggml_tensor * ple_weight_root(const ggml_tensor * wk, int64_t C, int k) { + if (wk->type != GGML_TYPE_F32 || ggml_nelements(wk) != C) { + return nullptr; + } + const bool cast = wk->op == GGML_OP_CPY; + const ggml_tensor * src = cast ? wk->src[0] : wk; + while (src && (src->op == GGML_OP_RESHAPE || src->op == GGML_OP_VIEW)) { + src = src->src[0]; + } + if (!src || src->op != GGML_OP_CONT) { + return nullptr; + } + const ggml_tensor * wv = src->src[0]; + if (!wv || wv->op != GGML_OP_VIEW || !wv->view_src) { + return nullptr; + } + const ggml_tensor * W = wv->view_src; + if ((cast ? W->type != GGML_TYPE_F16 : W->type != GGML_TYPE_F32) || W->ne[1] != C || W->ne[2] != 1 || W->ne[3] != 1 || + !ggml_is_contiguous(W) || wv->ne[0] != 1 || wv->ne[1] != C || wv->nb[1] != W->nb[1] || + wv->view_offs != (size_t) k * ggml_element_size(W)) { + return nullptr; + } + return W; +} + +static bool ple_conv_check(const ggml_cgraph * cgraph, int i, ggml_cuda_ple_conv_match & m) { + if (i < 0 || i + 2 >= cgraph->n_nodes) { + return ggml_cuda_conv_reject("ple_conv", "check 1", i); + } + const ggml_tensor * cc = cgraph->nodes[i]; + if (cc->op != GGML_OP_CONCAT || cc->type != GGML_TYPE_F32 || ggml_get_op_params_i32(cc, 0) != 0 || + (cc->flags & GGML_TENSOR_FLAG_OUTPUT) || !ggml_cuda_conv_next_is_noop(cgraph, i)) { + return ggml_cuda_conv_reject("ple_conv", "check 2", i); + } + const ggml_tensor * st = cc->src[0]; + const ggml_tensor * tr = cc->src[1]; + if (!st || !tr || st->type != GGML_TYPE_F32 || tr->type != GGML_TYPE_F32 || tr->op != GGML_OP_TRANSPOSE || !tr->view_src) { + return ggml_cuda_conv_reject("ple_conv", "check 3", i); + } + int64_t C = 0, T = 0; + const ggml_tensor * x = ggml_cuda_conv_activation(tr, C, T); + if (!x) { + return ggml_cuda_conv_reject("ple_conv", "check 4", i); + } + constexpr int K = 4, DIL = 3, H = (K - 1) * DIL; + if (st->ne[0] != H || st->ne[1] != C || st->ne[2] != 1 || st->ne[3] != 1 || !ggml_is_contiguous(st)) { + return ggml_cuda_conv_reject("ple_conv", "check 5", i); + } + if (tr->ne[0] != T || tr->ne[1] != C || tr->ne[2] != 1 || tr->view_offs != 0) { + return ggml_cuda_conv_reject("ple_conv", "check 6", i); + } + if (cc->ne[0] != T + H || cc->ne[1] != C || cc->ne[2] != 1 || cc->ne[3] != 1 || !ggml_is_contiguous(cc) || T < 256 || C % 256 != 0) { + return ggml_cuda_conv_reject("ple_conv", "check 7", i); + } + + // Every later reader of the concat must be a tail reader or one of the four taps. Between the first tap and the + // SiLU every node must belong to the chain, because the direct kernel replaces all of them. + int64_t tail_from = T + H; + int first_tap = -1, silu = -1; + int n_taps = 0; + const ggml_tensor * tap_out[K] = { nullptr, nullptr, nullptr, nullptr }; // latest value of tap k: CONT, then its MUL + const ggml_tensor * W = nullptr; + const ggml_tensor * chain = nullptr; // running ADD chain (starts at tap 0's MUL) + for (int n = i + 1; n < cgraph->n_nodes; ++n) { + const ggml_tensor * t = cgraph->nodes[n]; + const bool reads_cc = ggml_cuda_reads(t, cc); + if (silu >= 0) { + if (reads_cc && !ggml_cuda_conv_tail_reader(t, cc, H, tail_from)) { + return ggml_cuda_conv_reject("ple_conv", "check 8", n); + } + continue; + } + if (reads_cc) { + const int k = ple_tap_index(t, cc, C, T, DIL, K); + if (k >= 0) { + if (tap_out[k] != nullptr) { + return ggml_cuda_conv_reject("ple_conv", "check 9", n); + } + if (first_tap < 0) { + first_tap = n; + } + tap_out[k] = t; + ++n_taps; + continue; + } + if (t->op == GGML_OP_VIEW || t->op == GGML_OP_TRANSPOSE) { + // a tap view or its transpose (no-ops), or a tail view + if (t->op == GGML_OP_VIEW && ggml_cuda_conv_tail_view_col(t, cc, H) >= 0) { + tail_from = std::min(tail_from, ggml_cuda_conv_tail_view_col(t, cc, H)); + continue; + } + if (t->op == GGML_OP_VIEW && t->ne[0] == T && t->ne[1] == C && t->nb[1] == cc->nb[1] && t->view_offs % (DIL * sizeof(float)) == 0 && (int64_t) (t->view_offs / (DIL * sizeof(float))) < K) { + continue; + } + if (t->op == GGML_OP_TRANSPOSE && t->src[0] && t->src[0]->op == GGML_OP_VIEW && t->src[0]->view_src == cc && t->src[0]->ne[0] == T) { + continue; + } + return ggml_cuda_conv_reject("ple_conv", "check 10", n); + } + if (first_tap >= 0) { + return false; // a tail CONT/CPY inside the chain window would be skipped with it + } + if (!ggml_cuda_conv_tail_reader(t, cc, H, tail_from)) { + return ggml_cuda_conv_reject("ple_conv", "check 11", n); + } + continue; + } + if (first_tap < 0) { + continue; // unrelated work before the chain is computed normally + } + // inside the chain window: only the chain's own nodes + if (t->op == GGML_OP_MUL) { + const ggml_tensor * a = t->src[0]; + const ggml_tensor * b = t->src[1]; + int k = -1; + for (int q = 0; q < K; ++q) { + if (tap_out[q] && tap_out[q]->op == GGML_OP_CONT && (tap_out[q] == a || tap_out[q] == b)) { + k = q; + } + } + if (k < 0 || t->type != GGML_TYPE_F32 || !ggml_is_contiguous(t) || t->ne[0] != C || t->ne[1] != T || t->ne[2] != 1) { + return ggml_cuda_conv_reject("ple_conv", "check 12", n); + } + const ggml_tensor * Wk = ple_weight_root(tap_out[k] == a ? b : a, C, k); + if (!Wk || (W && W != Wk)) { + return ggml_cuda_conv_reject("ple_conv", "check 13", n); + } + W = Wk; + tap_out[k] = t; + if (k == 0) { + chain = t; + } + continue; + } + if (t->op == GGML_OP_ADD) { + if (!chain || t->src[0] != chain || t->type != GGML_TYPE_F32 || !ggml_is_contiguous(t)) { + return ggml_cuda_conv_reject("ple_conv", "check 14", n); + } + int k = -1; + for (int q = 1; q < K; ++q) { + if (tap_out[q] && tap_out[q]->op == GGML_OP_MUL && tap_out[q] == t->src[1]) { + k = q; + } + } + if (k < 0) { + return ggml_cuda_conv_reject("ple_conv", "check 15", n); + } + tap_out[k] = t; // consumed + chain = t; + continue; + } + if (t->op == GGML_OP_UNARY && ggml_get_unary_op(t) == GGML_UNARY_OP_SILU && chain && t->src[0] == chain) { + if (n_taps != K || !W || t->type != GGML_TYPE_F32 || !ggml_is_contiguous(t) || t->ne[0] != C || t->ne[1] != T) { + return ggml_cuda_conv_reject("ple_conv", "check 16", n); + } + silu = n; + continue; + } + if (t->op == GGML_OP_CONT || t->op == GGML_OP_RESHAPE || t->op == GGML_OP_VIEW || t->op == GGML_OP_CPY) { + // weight column plumbing: must not read the concat (checked above) and must lead to a tap MUL + continue; + } + return ggml_cuda_conv_reject("ple_conv", "check 17", n); + } + if (silu < 0 || first_tap < 0 || n_taps != K || !W) { + return ggml_cuda_conv_reject("ple_conv", "check 18", i); + } + // chain intermediates feed only the chain: the direct kernel never produces them + for (int n = first_tap; n < silu; ++n) { + const ggml_tensor * t = cgraph->nodes[n]; + if (t->op == GGML_OP_CONT || t->op == GGML_OP_MUL || t->op == GGML_OP_ADD || t->op == GGML_OP_CPY) { + if ((t->flags & GGML_TENSOR_FLAG_OUTPUT) || ggml_cuda_conv_readers(cgraph, n + 1, t, nullptr) != 1) { + return ggml_cuda_conv_reject("ple_conv", "chain intermediate has another reader", n); + } + } + } + m.concat_idx = i; m.first_tap_idx = first_tap; m.silu_idx = silu; + m.x = x; m.state = st; m.concat = cc; m.w = W; m.out = cgraph->nodes[silu]; + m.C = C; m.T = T; m.H = H; m.K = K; m.dil = DIL; m.tail_from = std::max(0, tail_from); + return true; +} + +bool ggml_cuda_ple_conv_match_at_concat(const ggml_cgraph * cgraph, int i, ggml_cuda_ple_conv_match & m) { + return ple_conv_check(cgraph, i, m); +} + +bool ggml_cuda_ple_conv_match_at_tap(const ggml_cgraph * cgraph, int i, ggml_cuda_ple_conv_match & m) { + const ggml_tensor * t = cgraph->nodes[i]; + if (t->op != GGML_OP_CONT || !t->src[0] || t->src[0]->op != GGML_OP_TRANSPOSE || !t->src[0]->src[0] || + t->src[0]->src[0]->op != GGML_OP_VIEW || !t->src[0]->src[0]->view_src) { + return false; + } + const ggml_tensor * cc = t->src[0]->src[0]->view_src; + if (cc->op != GGML_OP_CONCAT) { + return false; + } + for (int c = i - 1; c >= 0; --c) { + if (cgraph->nodes[c] == cc) { + return ple_conv_check(cgraph, c, m) && m.first_tap_idx == i; + } + } + return false; +} + +void ggml_cuda_ple_conv_write_tail(ggml_backend_cuda_context & ctx, const ggml_cuda_ple_conv_match & m) { + const int ncols = (int) (m.T + m.H - m.tail_from); + if (ncols <= 0) { + return; + } + const int64_t n = m.C * ncols; + ple_concat_tail<<<(unsigned) ((n + 255) / 256), 256, 0, ctx.stream()>>>((const float *) m.state->data, (const float *) m.x->data, + (float *) m.concat->data, (int) m.C, (int) m.T, (int) m.H, (int) m.tail_from, (int) (m.concat->nb[1] / sizeof(float))); + CUDA_CHECK(cudaGetLastError()); +} + +void ggml_cuda_ple_conv_direct(ggml_backend_cuda_context & ctx, const ggml_cuda_ple_conv_match & m) { + constexpr int TT = 128; + GGML_ASSERT(m.K == 4 && m.dil == 3); + dim3 grid((unsigned) (m.C / 256), (unsigned) ((m.T + TT - 1) / TT)); + if (m.w->type == GGML_TYPE_F16) { + ple_conv_kernel<4, 3, TT, half><<>>((const float *) m.state->data, (const float *) m.x->data, + (const half *) m.w->data, (float *) m.out->data, (int) m.C, (int) m.T); + } else { + ple_conv_kernel<4, 3, TT, float><<>>((const float *) m.state->data, (const float *) m.x->data, + (const float *) m.w->data, (float *) m.out->data, (int) m.C, (int) m.T); + } + CUDA_CHECK(cudaGetLastError()); +} diff --git a/ggml/src/ggml-cuda/ple-conv.cuh b/ggml/src/ggml-cuda/ple-conv.cuh new file mode 100644 index 000000000000..b5f60ad21cc1 --- /dev/null +++ b/ggml/src/ggml-cuda/ple-conv.cuh @@ -0,0 +1,19 @@ +#pragma once +#include "common.cuh" + +// qwen4exp PLE convolution: concat(history, x^T) -> four dilated taps (CONT of transposed views) x F16 or F32 weight +// columns -> ADD chain -> SiLU. On a match the concat writes only the columns its tail copies read and one +// kernel computes the SiLU output straight from the history and x. +struct ggml_cuda_ple_conv_match { + int concat_idx = -1, first_tap_idx = -1, silu_idx = -1; + const ggml_tensor * x = nullptr; // [C, T] F32 contiguous + const ggml_tensor * state = nullptr; // [H, C] F32 contiguous, H = (K-1)*dil + const ggml_tensor * concat = nullptr; // [T+H, C] + const ggml_tensor * w = nullptr; // [K, C] F16 or F32 + ggml_tensor * out = nullptr; // SiLU output [C, T] + int64_t C = 0, T = 0, H = 0, K = 0, dil = 0, tail_from = 0; +}; +bool ggml_cuda_ple_conv_match_at_concat(const ggml_cgraph * cgraph, int i, ggml_cuda_ple_conv_match & m); +bool ggml_cuda_ple_conv_match_at_tap(const ggml_cgraph * cgraph, int i, ggml_cuda_ple_conv_match & m); +void ggml_cuda_ple_conv_write_tail(ggml_backend_cuda_context & ctx, const ggml_cuda_ple_conv_match & m); +void ggml_cuda_ple_conv_direct(ggml_backend_cuda_context & ctx, const ggml_cuda_ple_conv_match & m); diff --git a/ggml/src/ggml-cuda/qsa-decode-wmma.cuh b/ggml/src/ggml-cuda/qsa-decode-wmma.cuh new file mode 100644 index 000000000000..63af09d3a568 --- /dev/null +++ b/ggml/src/ggml-cuda/qsa-decode-wmma.cuh @@ -0,0 +1,121 @@ +#pragma once + +typedef short qd_half16 __attribute__((ext_vector_type(16))); +typedef float qd_float8 __attribute__((ext_vector_type(8))); + +static __global__ __launch_bounds__(256) void qsa_decode_wmma_partial( + const char * q, const char * k, const char * v, const char * mask, const char * ids, + size_t q1, size_t q2, size_t k1, size_t k2, size_t v1, size_t v2, size_t m1, size_t i1, + int nk, int ns, int nh, int splits, int chunk, float scale, float * partial) { +#if defined(__HIP_DEVICE_COMPILE__) && !defined(RDNA3) + NO_DEVICE_CODE; +#else + const int lane = threadIdx.x & 31, w = threadIdx.x >> 5, r = lane & 15, hi = lane >> 4; + const int kvh = blockIdx.x, query = blockIdx.y, split = blockIdx.z; + const int * ir = (const int *) (ids + query*i1); + const half * mr = mask ? (const half *) (mask + query*m1) : nullptr; + __shared__ union { + uint16_t query_stage[8][16][32]; + float4 scores[8][64]; + } shared; + __shared__ __align__(16) uint16_t probabilities[16][16]; + __shared__ float alpha_s[16]; + qd_half16 qf[2]; + const _Float16 hs = (_Float16) scale; + for (int j = lane; j < 16*32; j += 32) { + const int head = j/32, d = j%32; + const float value = head < 12 ? *(const float *) (q + query*q1 + (kvh*12+head)*q2 + (w*32+d)*4) : 0.0f; + shared.query_stage[w][head][d] = __builtin_bit_cast(uint16_t, (_Float16)value*hs); + } + __syncthreads(); +#pragma unroll + for (int t = 0; t < 2; ++t) { + const uint16_t * src = &shared.query_stage[w][r][t*16]; + qf[t] = __builtin_bit_cast(qd_half16, (uint4[2]){*(const uint4 *)src, *(const uint4 *)(src+8)}); + } + __syncthreads(); + qd_float8 output[2] = {}; + float maximum = -INFINITY, normalizer = 0.0f; + const int begin = split*chunk, end = min(ns, begin+chunk); + for (int start = begin; start < end; start += 16) { + int keys[16]; +#pragma unroll + for (int j = 0; j < 16; ++j) { + int key = start+j < end ? ir[start+j] : -1; + if (key < 0 || key >= nk || (mr && __half2float(mr[key]) == -INFINITY)) { key = -1; } + keys[j] = key; + } + qd_float8 scores = {}; +#pragma unroll + for (int t = 0; t < 2; ++t) { + qd_half16 kf = {}; + const int key = keys[r]; + if (key >= 0 && key < nk) { + const uint16_t * src = (const uint16_t *) (k + key*k1 + kvh*k2) + 32*w+16*t; + kf = __builtin_bit_cast(qd_half16, (uint4[2]){*(const uint4 *)src, *(const uint4 *)(src+8)}); + } + scores = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(kf, qf[t], scores); + } + shared.scores[w][lane*2] = make_float4(scores[0],scores[1],scores[2],scores[3]); + shared.scores[w][lane*2+1] = make_float4(scores[4],scores[5],scores[6],scores[7]); + __syncthreads(); + if (w == 0) { + float score[8] = {}; +#pragma unroll + for (int sw = 0; sw < 8; ++sw) { + const float4 lo = shared.scores[sw][lane*2], high = shared.scores[sw][lane*2+1]; + score[0]+=lo.x;score[1]+=lo.y;score[2]+=lo.z;score[3]+=lo.w; + score[4]+=high.x;score[5]+=high.y;score[6]+=high.z;score[7]+=high.w; + } + float local = -INFINITY; +#pragma unroll + for (int e = 0; e < 8; ++e) { + const int key = keys[2*e+hi]; + if (key < 0 || key >= nk) { score[e] = -INFINITY; } + else if (mr) { score[e] += __half2float(mr[key]); } + local = fmaxf(local,score[e]); + } + local = fmaxf(local,__shfl_xor(local,16)); + const float next = fmaxf(maximum,local); + const float alpha = next == -INFINITY ? 1.0f : exp2f((maximum-next)*1.4426950408889634f); + float sum = 0.0f; +#pragma unroll + for (int e = 0; e < 8; ++e) { + const float p = next == -INFINITY ? 0.0f : exp2f((score[e]-next)*1.4426950408889634f); + sum += p; + probabilities[r][2*e+hi] = __builtin_bit_cast(uint16_t,(_Float16)p); + } + sum += __shfl_xor(sum,16); + normalizer = normalizer*alpha+sum; + maximum = next; + if (hi == 0) { alpha_s[r] = alpha; } + } + __syncthreads(); + const float alpha = alpha_s[r]; + const qd_half16 pf = __builtin_bit_cast(qd_half16, (uint4[2]){ + *(const uint4 *)&probabilities[r][0], *(const uint4 *)&probabilities[r][8]}); +#pragma unroll + for (int t = 0; t < 2; ++t) { + uint16_t values[16]; +#pragma unroll + for (int j = 0; j < 16; ++j) { + const int key = keys[j]; + values[j] = key >= 0 && key < nk ? *((const uint16_t *)(v+key*v1+kvh*v2)+32*w+16*t+r) : 0; + } + const qd_half16 vf = __builtin_bit_cast(qd_half16,values); +#pragma unroll + for (int e = 0; e < 8; ++e) { output[t][e] *= alpha; } + output[t] = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(vf,pf,output[t]); + } + __syncthreads(); + } + if (r < 12) { + float * dst = partial + ((size_t(query)*nh+kvh*12+r)*splits+split)*258; +#pragma unroll + for (int t = 0; t < 2; ++t) +#pragma unroll + for (int e = 0; e < 8; ++e) { dst[32*w+16*t+2*e+hi] = output[t][e]; } + if (w == 0 && hi == 0) { dst[256] = maximum; dst[257] = normalizer; } + } +#endif +} diff --git a/ggml/src/ggml-cuda/qsa-decode.cu b/ggml/src/ggml-cuda/qsa-decode.cu new file mode 100644 index 000000000000..ba83164230c4 --- /dev/null +++ b/ggml/src/ggml-cuda/qsa-decode.cu @@ -0,0 +1,158 @@ +#include "qsa-decode.cuh" +#if defined(GGML_USE_HIP) +#include "qsa-decode-wmma.cuh" + +static __global__ __launch_bounds__(128) void qsa_decode_partial( + const char * q, const char * k, const char * v, const char * mask, const char * ids, + size_t q1, size_t q2, size_t k1, size_t k2, size_t v1, size_t v2, size_t m1, size_t i1, + int nk, int ns, int nh, int ratio, int splits, int chunk, float scale, float * partial) { + const int lane = threadIdx.x % 32, warp = threadIdx.x / 32; + const int head = blockIdx.x, query = blockIdx.y, split = blockIdx.z; + const int kv_head = head / ratio; + const float * qr = (const float *) (q + query*q1 + head*q2); + const int * ir = (const int *) (ids + query*i1); + const half * mr = mask ? (const half *) (mask + query*m1) : nullptr; + float qv[8], acc[8] = {}; +#pragma unroll + for (int d = 0; d < 8; ++d) { qv[d] = qr[lane + d*32]; } + float mx = -INFINITY, sum = 0.0f; + const int begin = split*chunk, end = min(ns, begin + chunk); + for (int j = begin + warp; j < end; j += 4) { + const int key = ir[j]; + if (key < 0 || key >= nk) { continue; } + const float bias = mr ? __half2float(mr[key]) : 0.0f; + if (bias == -INFINITY) { continue; } + const half * kr = (const half *) (k + key*k1 + kv_head*k2); + const half * vr = (const half *) (v + key*v1 + kv_head*v2); + float dot = 0.0f; +#pragma unroll + for (int d = 0; d < 8; ++d) { dot = fmaf(qv[d], __half2float(kr[lane + d*32]), dot); } +#pragma unroll + for (int delta = 16; delta > 0; delta >>= 1) { dot += __shfl_xor_sync(0xffffffff, dot, delta, 32); } + const float score = dot*scale + bias; + const float next = fmaxf(mx, score); + const float alpha = expf(mx - next), weight = expf(score - next); +#pragma unroll + for (int d = 0; d < 8; ++d) { acc[d] = acc[d]*alpha + weight*__half2float(vr[lane + d*32]); } + sum = sum*alpha + weight; + mx = next; + } + __shared__ float values[4][256]; + __shared__ float maxima[4], sums[4]; +#pragma unroll + for (int d = 0; d < 8; ++d) { values[warp][lane + d*32] = acc[d]; } + if (lane == 0) { maxima[warp] = mx; sums[warp] = sum; } + __syncthreads(); + float maximum = -INFINITY; +#pragma unroll + for (int w = 0; w < 4; ++w) { maximum = fmaxf(maximum, maxima[w]); } + float factors[4], total = 0.0f; +#pragma unroll + for (int w = 0; w < 4; ++w) { + factors[w] = sums[w] > 0.0f ? expf(maxima[w] - maximum) : 0.0f; + total += factors[w]*sums[w]; + } + float * dst = partial + ((size_t(query)*nh + head)*splits + split)*258; + for (int d = threadIdx.x; d < 256; d += 128) { + float value = 0.0f; +#pragma unroll + for (int w = 0; w < 4; ++w) { value += values[w][d]*factors[w]; } + dst[d] = value; + } + if (threadIdx.x == 0) { dst[256] = maximum; dst[257] = total; } +} + +static __global__ __launch_bounds__(256) void qsa_decode_merge( + const float * partial, float * out, int splits) { + const int row = blockIdx.x, d = threadIdx.x; + const float * src = partial + size_t(row)*splits*258; + float maximum = -INFINITY; + for (int s = 0; s < splits; ++s) { maximum = fmaxf(maximum, src[s*258 + 256]); } + float sum = 0.0f, value = 0.0f; + for (int s = 0; s < splits; ++s) { + const float count = src[s*258 + 257]; + const float factor = count > 0.0f ? expf(src[s*258 + 256] - maximum) : 0.0f; + value += src[s*258 + d]*factor; + sum += count*factor; + } + out[size_t(row)*256 + d] = sum > 0.0f ? value/sum : 0.0f; +} + +bool ggml_cuda_flash_attn_ext_qsa_decode_supported(ggml_backend_cuda_context & ctx, const ggml_tensor * dst) { + if (!GGML_CUDA_CC_IS_RDNA3_5(ggml_cuda_info().devices[ctx.device].cc)) { return false; } + const auto * q = dst->src[0], * k = dst->src[1], * v = dst->src[2], * m = dst->src[3], * ids = dst->src[5]; + // the mask is optional: the launcher and both kernels take a null mask (visibility from the index rows alone) + if (!q || !k || !v || !ids || dst->src[4] || dst->src[6] || dst->src[7] || ggml_get_op_params_i32(dst, 4) != 0) { return false; } + float bias, cap; + memcpy(&bias, (const char *) dst->op_params + 4, 4); + memcpy(&cap, (const char *) dst->op_params + 8, 4); + if (bias != 0.0f || cap != 0.0f || q->type != GGML_TYPE_F32 || k->type != GGML_TYPE_F16 || + v->type != GGML_TYPE_F16 || dst->type != GGML_TYPE_F32 || ids->type != GGML_TYPE_I32 || + q->ne[0] != 256 || k->ne[0] != 256 || v->ne[0] != 256 || q->ne[1] < 1 || q->ne[1] > QSA_DECODE_MAX_QUERIES || + k->ne[1] < 1 || k->ne[1] > 262144 || v->ne[1] != k->ne[1] || + k->ne[2] < 1 || q->ne[2] % k->ne[2] || v->ne[2] != k->ne[2] || + q->ne[3] != 1 || k->ne[3] != 1 || v->ne[3] != 1 || + q->nb[0] != 4 || k->nb[0] != 2 || v->nb[0] != 2 || ids->nb[0] != 4 || + ids->ne[0] < 1 || ids->ne[0] > 2560 || ids->ne[1] < q->ne[1] || ids->ne[2] != 1 || ids->ne[3] != 1 || + !ggml_is_contiguous(dst)) { return false; } + return !m || (m->type == GGML_TYPE_F16 && m->nb[0] == 2 && m->ne[0] >= k->ne[1] && + m->ne[1] >= q->ne[1] && m->ne[2] == 1 && m->ne[3] == 1); +} + +static __global__ void qsa_decode_unique_ids(const char * input, size_t stride, int * output, int count, int keys) { + __shared__ int seen[8192]; + const int tid=threadIdx.x, row=blockIdx.x; + for (int i=tid;i<8192;i+=blockDim.x) seen[i]=-1; + __syncthreads(); + const auto * src=reinterpret_cast(input+size_t(row)*stride); + int * dst=output+size_t(row)*count; + for (int j=tid;j=keys) continue; + unsigned slot=(unsigned(key)*2654435761u)&8191u; + while (true) { + const int old=atomicCAS(seen+slot,-1,key); + if (old==-1) { dst[j]=key; break; } + if (old==key) break; + slot=(slot+1)&8191u; + } + } +} + +void ggml_cuda_flash_attn_ext_qsa_decode(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const auto * q = dst->src[0], * k = dst->src[1], * v = dst->src[2], * m = dst->src[3], * ids = dst->src[5]; + ggml_cuda_pool_alloc unique_ids(ctx.pool(), size_t(ids->ne[0])*q->ne[1]); + qsa_decode_unique_ids<<ne[1],256,0,ctx.stream()>>>((const char *)ids->data,ids->nb[1],unique_ids.get(),ids->ne[0],k->ne[1]); + const bool wmma = q->ne[2] == 12*k->ne[2] && k->nb[1]%16 == 0 && + k->nb[2]%16 == 0 && uintptr_t(k->data)%16 == 0; + const int ns = ids->ne[0], nq = q->ne[1], nh = q->ne[2], step = wmma ? 64 : 128; + // up to 8 queries every split covers `step` keys; with more queries the splits get longer so the grid stays + // near its 8-query size and the partial buffer stays small (33 splits x 512 queries would need ~420 MB) + int chunk = step; + if (nq > 8) { + const int target = std::max(1, (8*((ns + step - 1)/step) + nq - 1)/nq); + chunk = ((ns + target - 1)/target + 15)/16*16; + } + const int splits = (ns + chunk - 1)/chunk; + float scale; + memcpy(&scale, dst->op_params, sizeof(scale)); + ggml_cuda_pool_alloc partial(ctx.pool(), size_t(nq)*nh*splits*258); + if (wmma) { + qsa_decode_wmma_partial<<ne[2], nq, splits), 256, 0, ctx.stream()>>>( + (const char *) q->data, (const char *) k->data, (const char *) v->data, + m ? (const char *) m->data : nullptr, (const char *) unique_ids.get(), + q->nb[1], q->nb[2], k->nb[1], k->nb[2], v->nb[1], v->nb[2], m ? m->nb[1] : 0, size_t(ns)*sizeof(int), + k->ne[1], ns, nh, splits, chunk, scale, partial.get()); + } else { + qsa_decode_partial<<>>( + (const char *) q->data, (const char *) k->data, (const char *) v->data, + m ? (const char *) m->data : nullptr, (const char *) unique_ids.get(), + q->nb[1], q->nb[2], k->nb[1], k->nb[2], v->nb[1], v->nb[2], m ? m->nb[1] : 0, size_t(ns)*sizeof(int), + k->ne[1], ns, nh, nh/k->ne[2], splits, chunk, scale, partial.get()); + } + qsa_decode_merge<<>>(partial.get(), (float *) dst->data, splits); + CUDA_CHECK(cudaGetLastError()); +} + +#endif diff --git a/ggml/src/ggml-cuda/qsa-decode.cuh b/ggml/src/ggml-cuda/qsa-decode.cuh new file mode 100644 index 000000000000..1c256292a9e6 --- /dev/null +++ b/ggml/src/ggml-cuda/qsa-decode.cuh @@ -0,0 +1,7 @@ +#pragma once +#include "common.cuh" +// the per-query decode kernels take 1..512 queries; larger ubatches go to qsa_prefill. Up to 512 queries the +// per-query gathers cost less than qsa_prefill's pack of the whole K/V cache, which grows with the depth +#define QSA_DECODE_MAX_QUERIES 512 +bool ggml_cuda_flash_attn_ext_qsa_decode_supported(ggml_backend_cuda_context & ctx, const ggml_tensor * dst); +void ggml_cuda_flash_attn_ext_qsa_decode(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/qsa-prefill.cu b/ggml/src/ggml-cuda/qsa-prefill.cu new file mode 100644 index 000000000000..860878949f37 --- /dev/null +++ b/ggml/src/ggml-cuda/qsa-prefill.cu @@ -0,0 +1,495 @@ +#include "qsa-prefill.cuh" +#if defined(GGML_USE_HIP) +#include "common.cuh" +#include "qsa-prefill.cuh" +#include +#include +#include + +typedef short v16s __attribute__((ext_vector_type(16))); +typedef float v8f __attribute__((ext_vector_type(8))); + +#define QSA3_L2E 1.4426950408889634f +#define QSA3_G 4 + +static __device__ __forceinline__ uint16_t qsa3_f2h(const float f) { return __builtin_bit_cast(uint16_t, (_Float16) f); } +static __device__ __forceinline__ float qsa3_h2f(const uint16_t h) { return (float) __builtin_bit_cast(_Float16, h); } + +// union builder, kernel A: per query row, check that the (validity-transformed) row is non-decreasing; if not, rank-sort it into srow[q] and set sflag[q]. Sentinel 0x7FFFFFFF for invalid (-1 / >= nk) entries. +#define QSA3_SENT 0x7FFFFFFF +__global__ __launch_bounds__(256) void qsa3_rows_kernel( + const int * __restrict__ ids, const size_t i1, const int ns, const int nk, int * __restrict__ srow, int * __restrict__ sflag, + const char * mask, size_t m1) { + extern __shared__ int ent[]; + __shared__ int unsorted; + const int q = blockIdx.x, tid = threadIdx.x; + const int * row = reinterpret_cast(reinterpret_cast(ids) + (size_t) q * i1); + if (tid == 0) { unsorted = 0; } + __syncthreads(); + // without a mask the row itself is the visibility record: the graph only lists cells the query may see + // and writes -1 elsewhere (complete-block selection with tails) + const auto * mq = mask ? reinterpret_cast(mask + (size_t) q * m1) : nullptr; + for (int j = tid; j < ns; j += 256) { + const int key = row[j]; + const bool valid = key >= 0 && key < nk; + const bool visible = valid && (!mq || mq[key] != 0xfc00u); + ent[j] = visible ? key : QSA3_SENT; + if (valid && !visible) { atomicOr(&unsorted, 1); } + } + __syncthreads(); + for (int j = tid; j + 1 < ns; j += 256) { + if (ent[j] > ent[j+1]) { atomicOr(&unsorted, 1); } + } + __syncthreads(); + if (unsorted) { + int * dst = srow + (size_t) q * ns; + for (int j = tid; j < ns; j += 256) { + const int e = ent[j]; + int rank = 0; + for (int k = 0; k < ns; ++k) { const int f = ent[k]; rank += (f < e) || (f == e && k < j); } + dst[rank] = e; + } + } + if (tid == 0) { sflag[q] = unsorted; } +} + +// kernel B: one workgroup (128 lanes) per group of 4 queries. Lane = key sub-range; the range boundaries are the quantiles of the first non-empty row (block aligned), so the work per lane is balanced for the real (sorted, similar) rows. Each lane merges the 4 sorted rows inside its range: pass 1 counts, block-wide prefix sum, pass 2 writes. +#define QSA3_MERGE_LANES 128 +static __device__ __forceinline__ int qsa3_get(const int * __restrict__ p, const int j, const int nk, const bool raw) { + const int key = p[j]; + return raw ? ((key >= 0 && key < nk) ? key : QSA3_SENT) : key; +} +__global__ __launch_bounds__(QSA3_MERGE_LANES) void qsa3_merge_kernel( + const int * __restrict__ ids, const size_t i1, const int n_q, const int ns, const int nk, + const int * __restrict__ srow, const int * __restrict__ sflag, + uint16_t * __restrict__ ublk, uint16_t * __restrict__ umask, int * __restrict__ ucount, const int cap) { + extern __shared__ int rows_s[]; + __shared__ int scan_s[QSA3_MERGE_LANES]; + const int lane = threadIdx.x; + const int g = blockIdx.x; +#pragma unroll + for (int qi = 0; qi < QSA3_G; ++qi) { + const int q = QSA3_G * g + qi; + int * dst = rows_s + qi * ns; + if (q < n_q) { + const bool rw = sflag[q] == 0; + const int * src = rw ? reinterpret_cast(reinterpret_cast(ids) + (size_t) q * i1) : srow + (size_t) q * ns; + for (int j = lane; j < ns; j += QSA3_MERGE_LANES) { dst[j] = qsa3_get(src, j, nk, rw); } + } else { + for (int j = lane; j < ns; j += QSA3_MERGE_LANES) { dst[j] = QSA3_SENT; } + } + } + __syncthreads(); + const int * rp[QSA3_G]; bool raw[QSA3_G]; int cnt[QSA3_G]; +#pragma unroll + for (int qi = 0; qi < QSA3_G; ++qi) { + raw[qi] = false; rp[qi] = rows_s + qi * ns; + int a = 0, b = ns; + while (a < b) { const int mid = (a + b) >> 1; if (rp[qi][mid] < QSA3_SENT) { a = mid + 1; } else { b = mid; } } + cnt[qi] = a; + } + int kmax = -1, ref = -1; +#pragma unroll + for (int qi = 0; qi < QSA3_G; ++qi) { + if (cnt[qi] > 0) { kmax = max(kmax, qsa3_get(rp[qi], cnt[qi] - 1, nk, raw[qi])); if (ref < 0) { ref = qi; } } + } + uint16_t * ob = ublk + (size_t) g * cap, * om = umask + (size_t) g * cap; + if (kmax < 0) { if (lane == 0) { ucount[g] = 0; } return; } + const int rc = cnt[ref]; + const int lo = lane == 0 ? 0 : (qsa3_get(rp[ref], (int) (((long long) rc * lane) / QSA3_MERGE_LANES), nk, raw[ref]) & ~3); + const int hi = lane == QSA3_MERGE_LANES - 1 ? ((kmax + 4) & ~3) + : (qsa3_get(rp[ref], (int) (((long long) rc * (lane + 1)) / QSA3_MERGE_LANES), nk, raw[ref]) & ~3); + int p[QSA3_G], pe[QSA3_G]; +#pragma unroll + for (int qi = 0; qi < QSA3_G; ++qi) { + int a = 0, b = cnt[qi]; + while (a < b) { const int mid = (a + b) >> 1; if (qsa3_get(rp[qi], mid, nk, raw[qi]) < lo) { a = mid + 1; } else { b = mid; } } + p[qi] = a; b = cnt[qi]; + while (a < b) { const int mid = (a + b) >> 1; if (qsa3_get(rp[qi], mid, nk, raw[qi]) < hi) { a = mid + 1; } else { b = mid; } } + pe[qi] = a; + } + int n = 0; + { + int pp[QSA3_G]; int head[QSA3_G]; +#pragma unroll + for (int qi = 0; qi < QSA3_G; ++qi) { pp[qi] = p[qi]; head[qi] = pp[qi] < pe[qi] ? qsa3_get(rp[qi], pp[qi], nk, raw[qi]) : QSA3_SENT; } + while (true) { + int b = QSA3_SENT; +#pragma unroll + for (int qi = 0; qi < QSA3_G; ++qi) { b = min(b, head[qi]); } + if (b == QSA3_SENT) { break; } + b >>= 2; +#pragma unroll + for (int qi = 0; qi < QSA3_G; ++qi) { + if (head[qi] != QSA3_SENT && (head[qi] >> 2) == b) { + const int * rr = rp[qi]; + int pn = pp[qi]; + if ((head[qi] & 3) == 0 && pn + 3 < pe[qi] && rr[pn+3] == head[qi] + 3 && rr[pn+1] == head[qi] + 1 && rr[pn+2] == head[qi] + 2) { pn += 4; } + else { do { ++pn; } while (pn < pe[qi] && (rr[pn] >> 2) == b); } + pp[qi] = pn; head[qi] = pn < pe[qi] ? rr[pn] : QSA3_SENT; + } + } + ++n; + } + } + scan_s[lane] = n; + __syncthreads(); + for (int off = 1; off < QSA3_MERGE_LANES; off <<= 1) { + const int v = lane >= off ? scan_s[lane - off] : 0; + __syncthreads(); + scan_s[lane] += v; + __syncthreads(); + } + const int total = scan_s[QSA3_MERGE_LANES - 1]; + int o = scan_s[lane] - n; + { + int head[QSA3_G]; +#pragma unroll + for (int qi = 0; qi < QSA3_G; ++qi) { head[qi] = p[qi] < pe[qi] ? qsa3_get(rp[qi], p[qi], nk, raw[qi]) : QSA3_SENT; } + while (true) { + int b = QSA3_SENT; +#pragma unroll + for (int qi = 0; qi < QSA3_G; ++qi) { b = min(b, head[qi]); } + if (b == QSA3_SENT) { break; } + b >>= 2; + uint32_t mk = 0; +#pragma unroll + for (int qi = 0; qi < QSA3_G; ++qi) { + if (head[qi] != QSA3_SENT && (head[qi] >> 2) == b) { + const int * rr = rp[qi]; + int pn = p[qi]; + if ((head[qi] & 3) == 0 && pn + 3 < pe[qi] && rr[pn+3] == head[qi] + 3 && rr[pn+1] == head[qi] + 1 && rr[pn+2] == head[qi] + 2) { mk |= 0xFu << (4*qi); pn += 4; } + else { do { mk |= 1u << (4*qi + (rr[pn] & 3)); ++pn; } while (pn < pe[qi] && (rr[pn] >> 2) == b); } + p[qi] = pn; head[qi] = pn < pe[qi] ? rr[pn] : QSA3_SENT; + } + } + ob[o] = (uint16_t) b; om[o] = (uint16_t) mk; ++o; + } + } + if (lane == QSA3_MERGE_LANES - 1) { + int t = total; + while (t & 3) { ob[t] = 0xFFFFu; om[t] = 0; ++t; } + ucount[g] = t; + } +} + +struct qsa3_layout { + size_t q1, q2, o1, o2, m1; + int nk, n_q, gqa, cap; + float scale; +}; + +__global__ __launch_bounds__(256) void qsa3_attn_kernel( + const float * __restrict__ q, const uint16_t * __restrict__ pk, const uint16_t * __restrict__ pv, + const uint16_t * __restrict__ mask, const uint16_t * __restrict__ ublk, const uint16_t * __restrict__ umask, + const int * __restrict__ ucount, float * __restrict__ out, const qsa3_layout s) { +#if defined(__HIP_DEVICE_COMPILE__) && !defined(RDNA3) + NO_DEVICE_CODE; +#else + const int tid = threadIdx.x, lane = tid & 31, w = tid >> 5, r = lane & 15, hi = lane >> 4; + const int g = blockIdx.x, kvh = blockIdx.y; + __shared__ float4 part[3][8][64]; + __shared__ uint4 ptile[3][32]; + __shared__ float alpha_s[48]; + __shared__ float l_s[48]; + + const size_t nblk = (size_t) s.nk / 4; + const uint16_t * pkg = pk + (size_t) kvh * nblk * 1024; + const uint16_t * pvg = pv + (size_t) kvh * nblk * 1024; + const uint16_t * gblk = ublk + (size_t) g * s.cap; + const uint16_t * gmsk = umask + (size_t) g * s.cap; + const int nchunks = ucount[g] >> 2; + + v16s qf[3][2]; + { + uint16_t * stage = reinterpret_cast(&part[0][0][0]) + w * (48 * 32); + const _Float16 hs = (_Float16) s.scale; +#pragma unroll + for (int c = 0; c < 12; ++c) { + const int row = 4*c + (lane >> 3), dq = 4 * (lane & 7); + const int qi = row / 12, h = row - 12*qi; + const int query = QSA3_G * g + qi; + float4 v = make_float4(0.f, 0.f, 0.f, 0.f); + if (query < s.n_q) { + v = *reinterpret_cast(reinterpret_cast(q) + (size_t) query * s.q1 + (size_t) (kvh * s.gqa + h) * s.q2 + (size_t) (32*w + dq) * 4); + } + const uint16_t h0 = __builtin_bit_cast(uint16_t, (_Float16) v.x * hs), h1 = __builtin_bit_cast(uint16_t, (_Float16) v.y * hs); + const uint16_t h2 = __builtin_bit_cast(uint16_t, (_Float16) v.z * hs), h3 = __builtin_bit_cast(uint16_t, (_Float16) v.w * hs); + *reinterpret_cast(stage + row * 32 + dq) = make_uint2((uint32_t) h0 | ((uint32_t) h1 << 16), (uint32_t) h2 | ((uint32_t) h3 << 16)); + } + __syncthreads(); +#pragma unroll + for (int i = 0; i < 3; ++i) +#pragma unroll + for (int t = 0; t < 2; ++t) { + const uint16_t * src = stage + (16*i + r) * 32 + 16*t; + qf[i][t] = __builtin_bit_cast(v16s, (uint4[2]){*reinterpret_cast(src), *reinterpret_cast(src + 8)}); + } + __syncthreads(); + } + v8f O[3][2]; +#pragma unroll + for (int i = 0; i < 3; ++i) +#pragma unroll + for (int t = 0; t < 2; ++t) +#pragma unroll + for (int e = 0; e < 8; ++e) { O[i][t][e] = 0.f; } + float m = -INFINITY, l = 0.f; + const int own_row = 16*w + r, own_qi = own_row / 12; + const int own_query = QSA3_G * g + own_qi; + const uint16_t * maskq = (mask && w < 3 && own_query < s.n_q) + ? reinterpret_cast(reinterpret_cast(mask) + (size_t) own_query * s.m1) : nullptr; + + auto load_desc = [&](const int c, uint32_t & b01, uint32_t & b23, uint32_t & m01, uint32_t & m23) { + const uint2 bb = *reinterpret_cast(gblk + 4*c); + const uint2 mm = *reinterpret_cast(gmsk + 4*c); + b01 = bb.x; b23 = bb.y; m01 = mm.x; m23 = mm.y; + }; + auto blk_of = [&](const uint32_t b01, const uint32_t b23, const int j) -> int { + const uint32_t v = (j < 2 ? b01 : b23) >> (16 * (j & 1)) & 0xFFFFu; + return v == 0xFFFFu ? 0 : (int) v; + }; + auto load_k = [&](const uint32_t b01, const uint32_t b23, const uint32_t m01, const uint32_t m23, v16s * kf) { + const int j = r >> 2; + const uint32_t mk = ((j < 2 ? m01 : m23) >> (16 * (j & 1))) & 0xffffu; + if (((mk >> (r & 3)) & 0x1111u) == 0) { + kf[0] = v16s{}; + kf[1] = v16s{}; + return; + } + const int kb = blk_of(b01, b23, j); + const uint16_t * krow = pkg + (size_t) kb * 1024 + (r & 3) * 16 + w * 128; +#pragma unroll + for (int t = 0; t < 2; ++t) { + const uint4 lo = *reinterpret_cast(krow + t*64); + const uint4 hi4 = *reinterpret_cast(krow + t*64 + 8); + kf[t] = __builtin_bit_cast(v16s, (uint4[2]){lo, hi4}); + } + }; + + auto mask_values = [&](uint2 v, uint32_t mk) { + uint32_t lo = 0, hi = 0; + if (mk & 0x1111u) lo |= 0x0000ffffu; + if (mk & 0x2222u) lo |= 0xffff0000u; + if (mk & 0x4444u) hi |= 0x0000ffffu; + if (mk & 0x8888u) hi |= 0xffff0000u; + return make_uint2(v.x & lo, v.y & hi); + }; + + uint32_t b01 = 0, b23 = 0, m01 = 0, m23 = 0; + v16s kf[2]; + if (nchunks > 0) { load_desc(0, b01, b23, m01, m23); load_k(b01, b23, m01, m23, kf); } + + for (int c = 0; c < nchunks; ++c) { + v8f sc[3]; +#pragma unroll + for (int i = 0; i < 3; ++i) { +#pragma unroll + for (int e = 0; e < 8; ++e) { sc[i][e] = 0.f; } +#pragma unroll + for (int t = 0; t < 2; ++t) { sc[i] = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(kf[t], qf[i][t], sc[i]); } + } + v16s vf[2]; + { + const int kb0 = blk_of(b01, b23, 0), kb1 = blk_of(b01, b23, 1), kb2 = blk_of(b01, b23, 2), kb3 = blk_of(b01, b23, 3); +#pragma unroll + for (int t = 0; t < 2; ++t) { + const int d = 32*w + 16*t + r; + const uint2 v0 = mask_values(*reinterpret_cast(pvg + (size_t) kb0 * 1024 + d * 4), m01 & 0xffffu); + const uint2 v1 = mask_values(*reinterpret_cast(pvg + (size_t) kb1 * 1024 + d * 4), m01 >> 16); + const uint2 v2 = mask_values(*reinterpret_cast(pvg + (size_t) kb2 * 1024 + d * 4), m23 & 0xffffu); + const uint2 v3 = mask_values(*reinterpret_cast(pvg + (size_t) kb3 * 1024 + d * 4), m23 >> 16); + vf[t] = __builtin_bit_cast(v16s, (uint2[4]){v0, v1, v2, v3}); + } + } + uint32_t nb01 = 0, nb23 = 0, nm01 = 0, nm23 = 0; + v16s kfn[2]; + if (c + 1 < nchunks) { load_desc(c + 1, nb01, nb23, nm01, nm23); load_k(nb01, nb23, nm01, nm23, kfn); } + else { kfn[0] = kf[0]; kfn[1] = kf[1]; } +#pragma unroll + for (int i = 0; i < 3; ++i) { + if (i != w) { + part[i][w][lane*2+0] = make_float4(sc[i][0], sc[i][1], sc[i][2], sc[i][3]); + part[i][w][lane*2+1] = make_float4(sc[i][4], sc[i][5], sc[i][6], sc[i][7]); + } + } + __syncthreads(); + if (w < 3) { + float scf[8]; +#pragma unroll + for (int e = 0; e < 8; ++e) { scf[e] = 0.f; } +#pragma unroll + for (int w2 = 0; w2 < 8; ++w2) { + if (w2 == w) { +#pragma unroll + for (int e = 0; e < 8; ++e) { scf[e] += sc[w][e]; } + } else { + const float4 a = part[w][w2][lane*2+0], b = part[w][w2][lane*2+1]; + scf[0] += a.x; scf[1] += a.y; scf[2] += a.z; scf[3] += a.w; scf[4] += b.x; scf[5] += b.y; scf[6] += b.z; scf[7] += b.w; + } + } +#pragma unroll + for (int e = 0; e < 8; ++e) { + const int j = e >> 1, cc = 2*(e & 1) + hi; + const uint32_t mk = ((j < 2 ? m01 : m23) >> (16 * (j & 1))) & 0xFFFFu; + const bool kv = (mk >> (own_qi * 4 + cc)) & 1u; + if (maskq && kv) { scf[e] += qsa3_h2f(maskq[4 * blk_of(b01, b23, j) + cc]); } + if (!kv) { scf[e] = -INFINITY; } + } + float mloc = scf[0]; +#pragma unroll + for (int e = 1; e < 8; ++e) { mloc = fmaxf(mloc, scf[e]); } + mloc = fmaxf(mloc, __shfl_xor(mloc, 16)); + const float mnew = fmaxf(m, mloc); + const float alpha = mnew == -INFINITY ? 1.f : exp2f((m - mnew) * QSA3_L2E); + float p[8], ls = 0.f; +#pragma unroll + for (int e = 0; e < 8; ++e) { p[e] = mnew == -INFINITY ? 0.f : exp2f((scf[e] - mnew) * QSA3_L2E); ls += p[e]; } + ls += __shfl_xor(ls, 16); + l = l * alpha + ls; + m = mnew; + uint32_t pp[4], po[4]; +#pragma unroll + for (int e2 = 0; e2 < 4; ++e2) { + pp[e2] = (uint32_t) qsa3_f2h(p[2*e2]) | ((uint32_t) qsa3_f2h(p[2*e2+1]) << 16); + po[e2] = (uint32_t) __shfl_xor((int) pp[e2], 16); + } + uint16_t ph[16]; +#pragma unroll + for (int e = 0; e < 8; ++e) { + const uint16_t mine = (e & 1) ? (uint16_t) (pp[e>>1] >> 16) : (uint16_t) (pp[e>>1] & 0xffffu); + const uint16_t other = (e & 1) ? (uint16_t) (po[e>>1] >> 16) : (uint16_t) (po[e>>1] & 0xffffu); + ph[2*e] = hi ? other : mine; + ph[2*e+1] = hi ? mine : other; + } + if (hi == 0) { + ptile[w][r*2+0] = make_uint4((uint32_t) ph[0] | ((uint32_t) ph[1] << 16), (uint32_t) ph[2] | ((uint32_t) ph[3] << 16), + (uint32_t) ph[4] | ((uint32_t) ph[5] << 16), (uint32_t) ph[6] | ((uint32_t) ph[7] << 16)); + ptile[w][r*2+1] = make_uint4((uint32_t) ph[8] | ((uint32_t) ph[9] << 16), (uint32_t) ph[10] | ((uint32_t) ph[11] << 16), + (uint32_t) ph[12] | ((uint32_t) ph[13] << 16), (uint32_t) ph[14] | ((uint32_t) ph[15] << 16)); + alpha_s[own_row] = alpha; + } + } + __syncthreads(); + float al[3]; +#pragma unroll + for (int i = 0; i < 3; ++i) { al[i] = alpha_s[16*i + r]; } + if (__ballot(al[0] != 1.f || al[1] != 1.f || al[2] != 1.f)) { +#pragma unroll + for (int i = 0; i < 3; ++i) +#pragma unroll + for (int t = 0; t < 2; ++t) +#pragma unroll + for (int e = 0; e < 8; ++e) { O[i][t][e] *= al[i]; } + } +#pragma unroll + for (int i = 0; i < 3; ++i) { + const v16s pf = __builtin_bit_cast(v16s, (uint4[2]){ptile[i][r*2+0], ptile[i][r*2+1]}); +#pragma unroll + for (int t = 0; t < 2; ++t) { O[i][t] = __builtin_amdgcn_wmma_f32_16x16x16_f16_w32(vf[t], pf, O[i][t]); } + } + kf[0] = kfn[0]; kf[1] = kfn[1]; + b01 = nb01; b23 = nb23; m01 = nm01; m23 = nm23; + } + + if (w < 3 && hi == 0) { l_s[own_row] = l; } + __syncthreads(); + float * ostage = reinterpret_cast(&part[0][0][0]) + w * (16 * 32); +#pragma unroll + for (int i = 0; i < 3; ++i) { + const float li = l_s[16*i + r]; +#pragma unroll + for (int t = 0; t < 2; ++t) +#pragma unroll + for (int e = 0; e < 8; ++e) { ostage[r * 32 + 16*t + 2*e + hi] = li > 0.f ? O[i][t][e] / li : 0.f; } + __syncthreads(); +#pragma unroll + for (int c = 0; c < 4; ++c) { + const int rr = 4*c + (lane >> 3), dq = 4 * (lane & 7); + const int row = 16*i + rr, qi = row / 12, h = row - 12*qi; + const int query = QSA3_G * g + qi; + if (query < s.n_q) { + float * dst = reinterpret_cast(reinterpret_cast(out) + (size_t) query * s.o2 + (size_t) (kvh * s.gqa + h) * s.o1) + 32*w + dq; + *reinterpret_cast(dst) = *reinterpret_cast(ostage + rr * 32 + dq); + } + } + __syncthreads(); + } +#endif +} + +bool ggml_cuda_flash_attn_ext_qsa_prefill_supported(ggml_backend_cuda_context & ctx, const ggml_tensor * dst) { + const auto * q=dst->src[0], * k=dst->src[1], * v=dst->src[2], * m=dst->src[3], * ids=dst->src[5]; + if (!q || !k || !v || !ids || dst->src[4] || ggml_get_op_params_i32(dst, 4) != 0 || + !GGML_CUDA_CC_IS_RDNA3_5(ggml_cuda_info().devices[ctx.device].cc)) return false; + float bias, softcap; memcpy(&bias, (const char *) dst->op_params+4, 4); memcpy(&softcap, (const char *) dst->op_params+8, 4); + // the mask is optional: a maskless op relies on the index rows naming only visible cells (-1 elsewhere) + if (m && (m->type!=GGML_TYPE_F16 || m->nb[0]!=2 || m->ne[0]ne[1] || m->ne[1]ne[1] || m->ne[2]!=1 || m->ne[3]!=1)) return false; + return bias==0 && softcap==0 && q->type==GGML_TYPE_F32 && k->type==GGML_TYPE_F16 && v->type==GGML_TYPE_F16 && + dst->type==GGML_TYPE_F32 && ids->type==GGML_TYPE_I32 && + q->ne[0]==256 && k->ne[0]==256 && v->ne[0]==256 && q->ne[1]>=128 && q->ne[1]<=INT_MAX && + k->ne[1]>0 && k->ne[1]<=262140 && k->ne[1]%4==0 && v->ne[1]==k->ne[1] && + k->ne[2]>0 && k->ne[2]<=65535 && q->ne[2]==12*k->ne[2] && v->ne[2]==k->ne[2] && + q->ne[3]==1 && k->ne[3]==1 && v->ne[3]==1 && ggml_nelements(k)/4<=INT_MAX && + q->nb[0]==4 && k->nb[0]==2 && v->nb[0]==2 && ids->nb[0]==4 && + q->nb[1]%16==0 && q->nb[2]%16==0 && uintptr_t(q->data)%16==0 && + k->nb[1]%16==0 && k->nb[2]%16==0 && uintptr_t(k->data)%16==0 && + ids->nb[1]%4==0 && uintptr_t(ids->data)%4==0 && ids->ne[0]>0 && ids->ne[0]<=2560 && + ids->ne[1]>=q->ne[1] && ids->ne[2]==1 && ids->ne[3]==1 && ggml_is_contiguous(dst); +} + +static __global__ void qsa_pack_k(const char * src, uint16_t * dst, size_t nb1, size_t nb2, int nk, uint32_t count) { + const uint32_t i=blockIdx.x*blockDim.x+threadIdx.x; + if (i>=count) return; + const uint32_t d=(i%32)*8, key=(i/32)%nk, head=i/(uint32_t(nk)*32); + const size_t group=size_t(head)*(nk/4)+key/4; + const auto * in=reinterpret_cast(src+size_t(head)*nb2+size_t(key)*nb1+d*2); + auto * out=reinterpret_cast(dst+group*1024+(d/16)*64+(key%4)*16+d%16); + *out=*in; +} + +static __global__ void qsa_pack_v(const char * src, uint16_t * dst, size_t nb1, size_t nb2, int nk, uint32_t count) { + const uint32_t i=blockIdx.x*blockDim.x+threadIdx.x; + if (i>=count) return; + const uint32_t d=i%256, group=i/256, head=group/(nk/4), key=(group%(nk/4))*4; + const char * in=src+size_t(head)*nb2+size_t(key)*nb1+d*2; + *reinterpret_cast(dst+size_t(group)*1024+d*4)=make_ushort4( + *reinterpret_cast(in), *reinterpret_cast(in+nb1), + *reinterpret_cast(in+2*nb1), *reinterpret_cast(in+3*nb1)); +} + +void ggml_cuda_flash_attn_ext_qsa_prefill(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { + const auto * q=dst->src[0], * k=dst->src[1], * v=dst->src[2], * m=dst->src[3], * ids=dst->src[5]; + float scale; memcpy(&scale, dst->op_params, 4); + const int n_q = (int) q->ne[1], ns = (int) ids->ne[0], nk = (int) k->ne[1]; + ggml_cuda_pool_alloc packed_k(ctx.pool(), ggml_nelements(k)); + ggml_cuda_pool_alloc packed_v(ctx.pool(), ggml_nelements(v)); + const uint32_t count_k=ggml_nelements(k)/8, count_v=ggml_nelements(v)/4; + qsa_pack_k<<<(count_k+255)/256,256,0,ctx.stream()>>>((const char *)k->data,packed_k.get(),k->nb[1],k->nb[2],nk,count_k); + qsa_pack_v<<<(count_v+255)/256,256,0,ctx.stream()>>>((const char *)v->data,packed_v.get(),v->nb[1],v->nb[2],nk,count_v); + CUDA_CHECK(cudaGetLastError()); + const int ngroups = (n_q + QSA3_G - 1) / QSA3_G; + const int cap = (QSA3_G * ns + 3) & ~3; + ggml_cuda_pool_alloc ublk(ctx.pool(), (size_t) ngroups * cap); + ggml_cuda_pool_alloc umask(ctx.pool(), (size_t) ngroups * cap); + ggml_cuda_pool_alloc ucount(ctx.pool(), (size_t) ngroups); + ggml_cuda_pool_alloc srow(ctx.pool(), (size_t) n_q * ns); + ggml_cuda_pool_alloc sflag(ctx.pool(), (size_t) n_q); + { + const ggml_cuda_kernel_launch_params launch(dim3(n_q), dim3(256), (size_t) ns * sizeof(int), ctx.stream()); + ggml_cuda_kernel_launch(qsa3_rows_kernel, launch, (const int *) ids->data, ids->nb[1], ns, nk, srow.get(), sflag.get(), m ? (const char *)m->data : nullptr, m ? m->nb[1] : 0); + CUDA_CHECK(cudaGetLastError()); + const ggml_cuda_kernel_launch_params launch2(dim3(ngroups), dim3(QSA3_MERGE_LANES), (size_t) QSA3_G * ns * sizeof(int), ctx.stream()); + ggml_cuda_kernel_launch(qsa3_merge_kernel, launch2, (const int *) ids->data, ids->nb[1], n_q, ns, nk, + (const int *) srow.get(), (const int *) sflag.get(), ublk.get(), umask.get(), ucount.get(), cap); + CUDA_CHECK(cudaGetLastError()); + } + qsa3_layout layout{q->nb[1], q->nb[2], dst->nb[1], dst->nb[2], m ? m->nb[1] : 0, nk, n_q, 12, cap, scale}; + const ggml_cuda_kernel_launch_params launch(dim3(ngroups, k->ne[2]), dim3(256), 0, ctx.stream()); + ggml_cuda_kernel_launch(qsa3_attn_kernel, launch, (const float *) q->data, packed_k.get(), packed_v.get(), + m ? (const uint16_t *) m->data : nullptr, (const uint16_t *) ublk.get(), (const uint16_t *) umask.get(), + (const int *) ucount.get(), (float *) dst->data, layout); + CUDA_CHECK(cudaGetLastError()); +} + +#endif diff --git a/ggml/src/ggml-cuda/qsa-prefill.cuh b/ggml/src/ggml-cuda/qsa-prefill.cuh new file mode 100644 index 000000000000..d9eecc50ac62 --- /dev/null +++ b/ggml/src/ggml-cuda/qsa-prefill.cuh @@ -0,0 +1,4 @@ +#pragma once +#include "common.cuh" +bool ggml_cuda_flash_attn_ext_qsa_prefill_supported(ggml_backend_cuda_context & ctx, const ggml_tensor * dst); +void ggml_cuda_flash_attn_ext_qsa_prefill(ggml_backend_cuda_context & ctx, ggml_tensor * dst); diff --git a/ggml/src/ggml-cuda/quantize.cu b/ggml/src/ggml-cuda/quantize.cu index bcc7723957b5..8a2112224123 100644 --- a/ggml/src/ggml-cuda/quantize.cu +++ b/ggml/src/ggml-cuda/quantize.cu @@ -1,4 +1,5 @@ #include "quantize.cuh" +#include "unary.cuh" #include #if defined(BLACKWELL_MMA_AVAILABLE) @@ -463,14 +464,7 @@ static __global__ void quantize_mmq_q8_1( 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 i00 = i0; - ggml_cuda_pdl_sync(); + const int64_t first_chunk = (int64_t) blockIdx.y * CUDA_QUANTIZE_MMQ_CHUNKS_PER_BLOCK; int64_t base_idx; if constexpr (scatter) { @@ -485,76 +479,133 @@ static __global__ void quantize_mmq_q8_1( const float4 * x4 = (const float4 *) x; block_q8_1_mmq * y = (block_q8_1_mmq *) vy; - const int64_t k_block = i0 / QK8_1_MMQ; // column block in the channel - const int64_t iqs = i0 % QK8_1_MMQ; // quant index in block - - // Load 4 floats per thread and calculate max. abs. value between them: - const float4 xi = i0 < ne00 ? x4[(base_idx + 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)); + ggml_cuda_pdl_sync(); - // 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)); - } + for (int chunk = 0; chunk < CUDA_QUANTIZE_MMQ_CHUNKS_PER_BLOCK; ++chunk) { + const int64_t i0 = ((first_chunk + chunk) * blockDim.x + threadIdx.x) * 4; + if (i0 >= ne0) { + break; + } - float sum; - if (ds_layout != MMQ_Q8_1_DS_LAYOUT_D4) { - sum = xi.x + xi.y + xi.z + xi.w; + const int64_t k_block = i0 / QK8_1_MMQ; + const int64_t iqs = i0 % QK8_1_MMQ; + const float4 xi = i0 < ne00 ? x4[(base_idx + i0)/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)); - // 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); + for (int offset = vals_per_scale/8; offset > 0; offset >>= 1) { + amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFF, amax, 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); - const float d = 1.0f / d_inv; - - // write the block once (normal) or to each of the token's compact rows (scatter) - const int nwrite = scatter ? n_expert_used : 1; + float sum; + if (ds_layout != MMQ_Q8_1_DS_LAYOUT_D4) { + sum = xi.x + xi.y + xi.z + xi.w; #pragma unroll - for (int slot = 0; slot < nwrite; ++slot) { - int64_t ib; - if constexpr (scatter) { - const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot]; - ib = k_block*ne1 + i; - } else { - const int64_t ib0 = blockIdx.z*((int64_t)gridDim.x*gridDim.y*blockDim.x/QK8_1); // first block of channel - ib = ib0 + k_block*ne1 + blockIdx.x; + for (int offset = vals_per_sum/8; offset > 0; offset >>= 1) { + sum += __shfl_xor_sync(0xFFFFFFFF, sum, offset, WARP_SIZE); + } } - // 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; + 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); + const float d = 1.0f / d_inv; - if (ds_layout == MMQ_Q8_1_DS_LAYOUT_D2S6) { - if (iqs % 16 == 0 && iqs < 96) { - y[ib].d2s6[2 + iqs/16] = sum; - if (iqs % 64 == 0) { - y[ib].d2s6[iqs/64] = d; - } - } - } else if (iqs % 32 == 0) { - if (ds_layout == MMQ_Q8_1_DS_LAYOUT_DS4) { - y[ib].ds4[iqs/32] = make_half2(d, sum); + const int nwrite = scatter ? n_expert_used : 1; +#pragma unroll + for (int slot = 0; slot < nwrite; ++slot) { + int64_t ib; + if constexpr (scatter) { + const int64_t i = ids[(int64_t) blockIdx.x * n_expert_used + slot]; + ib = k_block*ne1 + i; } else { - y[ib].d4[iqs/32] = d; + const int64_t ib0 = blockIdx.z * ((int64_t) gridDim.x * ne0 / QK8_1_MMQ); + ib = ib0 + k_block*ne1 + blockIdx.x; + } + + 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) { + y[ib].d2s6[2 + iqs/16] = sum; + if (iqs % 64 == 0) { + y[ib].d2s6[iqs/64] = d; + } + } + } else if (iqs % 32 == 0) { + 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; + } } } } GGML_UNUSED(n_expert_used); } +static __global__ void quantize_mmq_q8_1_swiglu( + const float * __restrict__ gate, const float * __restrict__ up, const int32_t * __restrict__ ids, + void * __restrict__ vy, const int64_t ne00, const int64_t ne0, const int ne1, const int logical_n1, + const int64_t gate_s1, const int64_t gate_st, const int64_t up_s1, const int64_t up_st) { + const int64_t i0 = ((int64_t) blockDim.x * blockIdx.y + threadIdx.x) * 4; + if (i0 >= ne0) { + return; + } + + const int64_t logical_row = ids ? ids[blockIdx.x] : blockIdx.x; + const int64_t slot = logical_row % logical_n1; + const int64_t token = logical_row / logical_n1; + const int64_t gate_base = token * gate_st + slot * gate_s1; + const int64_t up_base = token * up_st + slot * up_s1; + + ggml_cuda_pdl_sync(); + float values[4]; +#pragma unroll + for (int i = 0; i < 4; ++i) { + if (i0 + i < ne00) { + values[i] = ggml_cuda_materialize(ggml_cuda_op_silu_single(gate[gate_base + i0 + i]) * up[up_base + i0 + i]); + } else { + values[i] = 0.0f; + } + } + const float4 xi = make_float4(values[0], values[1], values[2], values[3]); + + float amax = fabsf(xi.x); + amax = fmaxf(amax, fabsf(xi.y)); + amax = fmaxf(amax, fabsf(xi.z)); + amax = fmaxf(amax, fabsf(xi.w)); +#pragma unroll + for (int offset = 4; offset > 0; offset >>= 1) { + amax = fmaxf(amax, __shfl_xor_sync(0xFFFFFFFF, amax, offset, WARP_SIZE)); + } + + const float d_inv = amax == 0.0f ? 0.0f : 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); + const float d = amax == 0.0f ? 0.0f : 1.0f / d_inv; + + block_q8_1_mmq * y = (block_q8_1_mmq *) vy; + const int64_t k_block = i0 / QK8_1_MMQ; + const int64_t iqs = i0 % QK8_1_MMQ; + const int64_t ib = k_block * ne1 + blockIdx.x; + ((char4 *) y[ib].qs)[iqs / 4] = q; + if (iqs % 32 == 0) { + y[ib].d4[iqs / 32] = d; + } +} + void quantize_row_q8_1_cuda( const float * x, const int32_t * ids, void * vy, const ggml_type type_src0, const int64_t ne00, const int64_t s01, const int64_t s02, const int64_t s03, @@ -580,7 +631,8 @@ void quantize_mmq_q8_1_cuda( GGML_ASSERT(ne0 % QK8_1_MMQ == 0); // ne1 tends to assume the highest values, therefore use it as the "x" dimension of the CUDA grid: - const int64_t block_num_y = (ne0 + 4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ - 1) / (4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ); + const int64_t vals_per_block = 4 * CUDA_QUANTIZE_BLOCK_SIZE_MMQ * CUDA_QUANTIZE_MMQ_CHUNKS_PER_BLOCK; + const int64_t block_num_y = (ne0 + vals_per_block - 1) / vals_per_block; const dim3 num_blocks(ne1, block_num_y, ne2*ne3); const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1); switch (mmq_get_q8_1_ds_layout(type_src0)) { @@ -602,6 +654,20 @@ void quantize_mmq_q8_1_cuda( } } +void quantize_mmq_q8_1_swiglu_cuda( + const float * gate, const float * up, const int32_t * ids, void * vy, const ggml_type type_src0, + const int64_t ne00, const int64_t ne0, const int64_t ne1, const int logical_n1, + const int64_t gate_s1, const int64_t gate_st, const int64_t up_s1, const int64_t up_st, cudaStream_t stream) { + GGML_ASSERT(mmq_get_q8_1_ds_layout(type_src0) == MMQ_Q8_1_DS_LAYOUT_D4); + GGML_ASSERT(ne00 % 4 == 0); + GGML_ASSERT(ne0 % QK8_1_MMQ == 0); + 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, 1); + const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1); + quantize_mmq_q8_1_swiglu<<>>( + gate, up, ids, vy, ne00, ne0, ne1, logical_n1, gate_s1, gate_st, up_s1, up_st); +} + // scatter=true reuses the quant kernel: grid over tokens, ids = inverse map (token slot -> compact row) void quantize_scatter_mmq_q8_1_cuda( const float * x, const int32_t * ids_src1_inv, void * vy, const ggml_type type_src0, @@ -610,7 +676,8 @@ void quantize_scatter_mmq_q8_1_cuda( GGML_ASSERT(ne00 % 4 == 0); GGML_ASSERT(ne0 % QK8_1_MMQ == 0); - const int64_t block_num_y = (ne0 + 4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ - 1) / (4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ); + const int64_t vals_per_block = 4 * CUDA_QUANTIZE_BLOCK_SIZE_MMQ * CUDA_QUANTIZE_MMQ_CHUNKS_PER_BLOCK; + const int64_t block_num_y = (ne0 + vals_per_block - 1) / vals_per_block; const dim3 num_blocks(n_tokens, block_num_y, 1); const dim3 block_size(CUDA_QUANTIZE_BLOCK_SIZE_MMQ, 1, 1); switch (mmq_get_q8_1_ds_layout(type_src0)) { diff --git a/ggml/src/ggml-cuda/quantize.cuh b/ggml/src/ggml-cuda/quantize.cuh index 5f08dcbfe331..1801da05e871 100644 --- a/ggml/src/ggml-cuda/quantize.cuh +++ b/ggml/src/ggml-cuda/quantize.cuh @@ -7,9 +7,11 @@ #define CUDA_QUANTIZE_BLOCK_SIZE 256 #define CUDA_QUANTIZE_BLOCK_SIZE_MMQ 128 +#define CUDA_QUANTIZE_MMQ_CHUNKS_PER_BLOCK 2 static_assert(MATRIX_ROW_PADDING % CUDA_QUANTIZE_BLOCK_SIZE == 0, "Risk of out-of-bounds access."); static_assert(MATRIX_ROW_PADDING % (4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ) == 0, "Risk of out-of-bounds access."); +static_assert((4*CUDA_QUANTIZE_BLOCK_SIZE_MMQ) % QK8_1_MMQ == 0, "Quantization chunks must contain complete Q8_1 MMQ blocks."); typedef void (*quantize_cuda_t)( const float * x, const int32_t * ids, void * vy, @@ -26,6 +28,11 @@ void quantize_mmq_q8_1_cuda( ggml_type type_src0, int64_t ne00, int64_t s01, int64_t s02, int64_t s03, int64_t ne0, int64_t ne1, int64_t ne2, int64_t ne3, cudaStream_t stream); +void quantize_mmq_q8_1_swiglu_cuda( + const float * gate, const float * up, const int32_t * ids, void * vy, ggml_type type_src0, + int64_t ne00, int64_t ne0, int64_t ne1, int logical_n1, + int64_t gate_s1, int64_t gate_st, int64_t up_s1, int64_t up_st, cudaStream_t stream); + void quantize_mmq_fp4_cuda(const float * x, const int32_t * ids, void * vy, diff --git a/ggml/src/ggml-cuda/sum.cu b/ggml/src/ggml-cuda/sum.cu index c56257b44066..d1f938905d0b 100644 --- a/ggml/src/ggml-cuda/sum.cu +++ b/ggml/src/ggml-cuda/sum.cu @@ -2,7 +2,12 @@ #include "sumrows.cuh" #ifdef GGML_CUDA_USE_CUB +#ifdef GGML_CUDA_CUB_IS_HIPCUB +#include +namespace cub = hipcub; +#else #include +#endif // GGML_CUDA_CUB_IS_HIPCUB using namespace cub; #endif // GGML_CUDA_USE_CUB diff --git a/ggml/src/ggml-cuda/top-k.cu b/ggml/src/ggml-cuda/top-k.cu index c7a0c831788d..97eea70e13b0 100644 --- a/ggml/src/ggml-cuda/top-k.cu +++ b/ggml/src/ggml-cuda/top-k.cu @@ -2,12 +2,17 @@ #include "top-k.cuh" #ifdef GGML_CUDA_USE_CUB -# include -# if (CCCL_MAJOR_VERSION >= 3 && CCCL_MINOR_VERSION >= 2) -# define CUB_TOP_K_AVAILABLE -# include +# ifdef GGML_CUDA_CUB_IS_HIPCUB +# include +namespace cub = hipcub; +# else +# include +# if (CCCL_MAJOR_VERSION >= 3 && CCCL_MINOR_VERSION >= 2) +# define CUB_TOP_K_AVAILABLE +# include using namespace cub; -# endif // CCCL_MAJOR_VERSION >= 3 && CCCL_MINOR_VERSION >= 2 +# endif // CCCL_MAJOR_VERSION >= 3 && CCCL_MINOR_VERSION >= 2 +# endif // GGML_CUDA_CUB_IS_HIPCUB #endif // GGML_CUDA_USE_CUB #ifdef CUB_TOP_K_AVAILABLE @@ -48,7 +53,7 @@ static int next_power_of_2(int x) { #endif // CUB_TOP_K_AVAILABLE -#if !defined(GGML_CUDA_USE_CUB) && defined(GGML_USE_HIP) +#if defined(GGML_USE_HIP) static __device__ __forceinline__ uint32_t top_k_float_to_ordered(float value) { const uint32_t bits = __float_as_uint(value); @@ -56,6 +61,40 @@ static __device__ __forceinline__ uint32_t top_k_float_to_ordered(float value) { return bits ^ mask; } +template +static __device__ void top_k_gather_equal( + const float * src, int * dst, int ncols, uint32_t threshold, int limit, int offset) { + const int tid = threadIdx.x; + const int lane = tid % warpSize; + const int warp = tid / warpSize; + const int nwarps = BLOCK_SIZE / warpSize; + __shared__ int warp_counts[32]; + int count = 0; + + for (int base = 0; base < ncols && count < limit; base += BLOCK_SIZE) { + const int col = base + tid; + const bool equal = col < ncols && top_k_float_to_ordered(src[col]) == threshold; + const unsigned long long mask = __ballot(equal); + if (lane == 0) { + warp_counts[warp] = __popcll(mask); + } + __syncthreads(); + int before = count; + for (int w = 0; w < nwarps; ++w) { + if (w < warp) { + before += warp_counts[w]; + } + count += warp_counts[w]; + } + const unsigned long long lane_mask = (1ULL << lane) - 1; + const int pos = before + __popcll(mask & lane_mask); + if (equal && pos < limit) { + dst[offset + pos] = col; + } + __syncthreads(); + } +} + struct top_k_radix_state { uint32_t prefix; uint32_t prefix_mask; @@ -146,7 +185,7 @@ static __global__ void top_k_radix_reset_counters(top_k_radix_state * states, in } } -template +template static __global__ void top_k_radix_gather( const float * __restrict__ src, int * __restrict__ dst, @@ -168,18 +207,29 @@ static __global__ void top_k_radix_gather( if (key > state->prefix) { const int pos = atomicAdd(&state->greater_count, 1); row_dst[pos] = col; - } else if (key == state->prefix) { - const int pos = atomicAdd(&state->equal_count, 1); - if (pos < state->rank) { - row_dst[k - state->rank + pos] = col; + } else if constexpr (!STABLE_TIES) { + if (key == state->prefix) { + const int pos = atomicAdd(&state->equal_count, 1); + if (pos < state->rank) { + row_dst[k - state->rank + pos] = col; + } } } } } +template +static __global__ void top_k_radix_gather_equal( + const float * src, int * dst, const top_k_radix_state * states, int ncols, int k) { + const int row = blockIdx.x; + const top_k_radix_state state = states[row]; + top_k_gather_equal(src + (size_t) row*ncols, dst + (size_t) row*k, + ncols, state.prefix, state.rank, k - state.rank); +} + static void top_k_radix_cuda( ggml_cuda_pool & pool, - const float * src, int * dst, int ncols, int nrows, int k, cudaStream_t stream) { + const float * src, int * dst, int ncols, int nrows, int k, cudaStream_t stream, bool stable_ties = false) { constexpr int BLOCK_SIZE = 256; constexpr int RADIX_BITS = 8; constexpr int NBINS = 1 << RADIX_BITS; @@ -203,12 +253,18 @@ static void top_k_radix_cuda( top_k_radix_reset_counters <<<(nrows + BLOCK_SIZE - 1) / BLOCK_SIZE, BLOCK_SIZE, 0, stream>>>(states, nrows); - top_k_radix_gather - <<>>( - src, dst, states, ncols, k, blocks_per_row); + if (stable_ties) { + top_k_radix_gather + <<>>(src, dst, states, ncols, k, blocks_per_row); + top_k_radix_gather_equal + <<>>(src, dst, states, ncols, k); + } else { + top_k_radix_gather + <<>>(src, dst, states, ncols, k, blocks_per_row); + } } -#endif // !defined(GGML_CUDA_USE_CUB) && defined(GGML_USE_HIP) +#endif // defined(GGML_USE_HIP) void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; @@ -225,6 +281,16 @@ void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { const int64_t nrows = ggml_nrows(src0); const int64_t k = dst->ne[0]; ggml_cuda_pool & pool = ctx.pool(); +#if defined(GGML_USE_HIP) + // gfx1151 measurements (2026-09-14) favor radix by 1.2-5.1x above these cutoffs. + // Recheck the cutoffs when hipCUB changes; smaller batches favor its sort path. + const bool use_radix = ncols >= 8192 || (ncols >= 4096 && nrows >= 128) || (ncols >= 2048 && nrows >= 512); + if (GGML_CUDA_CC_IS_RDNA3_5(ggml_cuda_info().devices[ctx.device].cc) && use_radix && + ncols <= INT_MAX && nrows > 1 && nrows <= INT_MAX && k <= INT_MAX) { + top_k_radix_cuda(pool, src0_d, dst_d, ncols, nrows, k, stream, true); + return; + } +#endif #ifdef CUB_TOP_K_AVAILABLE // TODO: Switch to `DeviceSegmentedTopK` for multi-row TopK once implemented // https://github.com/NVIDIA/cccl/issues/6391 @@ -260,7 +326,8 @@ void ggml_cuda_op_top_k(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { #else // GGML_CUDA_USE_CUB #if defined(GGML_USE_HIP) if (ncols > 1024) { - top_k_radix_cuda(pool, src0_d, dst_d, ncols, nrows, k, stream); + // Select tied scores by increasing column index: QSA has many exact-zero ties at the cut. + top_k_radix_cuda(pool, src0_d, dst_d, ncols, nrows, k, stream, true); } else { #endif // defined(GGML_USE_HIP) ggml_cuda_pool_alloc temp_dst_alloc(pool, ncols * nrows); diff --git a/ggml/src/ggml-cuda/topk-moe.cu b/ggml/src/ggml-cuda/topk-moe.cu index dadcd601cb48..f21fa6e5fc7e 100644 --- a/ggml/src/ggml-cuda/topk-moe.cu +++ b/ggml/src/ggml-cuda/topk-moe.cu @@ -87,7 +87,7 @@ __device__ void sqrt_softplus_warp_inplace(float (&vals)[experts_per_thread], co It is intended as fusion of softmax->top-k->get_rows pipeline for MoE models */ -template +template __launch_bounds__(TOPK_MOE_ROWS_PER_BLOCK * WARP_SIZE, 1) __global__ void topk_moe_cuda(const float * logits, float * weights, @@ -172,10 +172,11 @@ __global__ void topk_moe_cuda(const float * logits, float wt_sum = 0.f; - float output_weights[experts_per_thread]; + constexpr int output_weights_per_thread = small_topk ? 1 : experts_per_thread; + float output_weights[output_weights_per_thread]; #pragma unroll - for (int i = 0; i < experts_per_thread; i++) { + for (int i = 0; i < output_weights_per_thread; i++) { output_weights[i] = 0.f; } @@ -238,7 +239,7 @@ __global__ void topk_moe_cuda(const float * logits, } if ((k & (WARP_SIZE - 1)) == threadIdx.x) { - output_weights[k / WARP_SIZE] = max_val; + output_weights[small_topk ? 0 : k / WARP_SIZE] = max_val; } if ((max_expert & (WARP_SIZE - 1)) == threadIdx.x) { @@ -254,17 +255,17 @@ __global__ void topk_moe_cuda(const float * logits, wt_sum = max(wt_sum, clamp_val); const float inv_sum = 1.0f / wt_sum; - for (int i = 0; i < experts_per_thread; i++) { + for (int i = 0; i < output_weights_per_thread; i++) { output_weights[i] *= inv_sum; } } if (config.delayed_softmax) { - softmax_warp_inplace(output_weights, n_expert_used, threadIdx.x); + softmax_warp_inplace(output_weights, n_expert_used, threadIdx.x); } #pragma unroll - for (int i = 0; i < experts_per_thread; i++) { + for (int i = 0; i < output_weights_per_thread; i++) { const int idx = i * WARP_SIZE + threadIdx.x; if (idx < n_expert_used) { weights[idx] = output_weights[i] * scale_val; @@ -334,8 +335,13 @@ static void launch_topk_moe_cuda(ggml_backend_cuda_context & ctx, logits, weights, ids, bias, n_rows, n_expert_used, clamp_val, scale_val, config); break; case 512: - ggml_cuda_kernel_launch(topk_moe_cuda<512, has_bias>, launch_params, - logits, weights, ids, bias, n_rows, n_expert_used, clamp_val, scale_val, config); + if (n_expert_used <= WARP_SIZE) { + ggml_cuda_kernel_launch(topk_moe_cuda<512, has_bias, true>, launch_params, + logits, weights, ids, bias, n_rows, n_expert_used, clamp_val, scale_val, config); + } else { + ggml_cuda_kernel_launch(topk_moe_cuda<512, has_bias>, launch_params, + logits, weights, ids, bias, n_rows, n_expert_used, clamp_val, scale_val, config); + } break; case 576: ggml_cuda_kernel_launch(topk_moe_cuda<576, has_bias>, launch_params, diff --git a/ggml/src/ggml-cuda/unary.cu b/ggml/src/ggml-cuda/unary.cu index d3e594878fcc..adc42022701f 100644 --- a/ggml/src/ggml-cuda/unary.cu +++ b/ggml/src/ggml-cuda/unary.cu @@ -260,7 +260,8 @@ void ggml_cuda_op_softplus(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { /* gated ops */ template -static __global__ void unary_gated_op_kernel(const T * x, const T * g, T * dst, const int64_t k, const int64_t n, const int64_t o0, const int64_t o1) { +static __global__ void unary_gated_op_kernel(const T * x, const T * g, T * dst, const int64_t k, const int64_t n, const int64_t o0, const int64_t o1, + uint16_t * dst16) { ggml_cuda_pdl_lc(); const int64_t i = int64_t(blockDim.x)*blockIdx.x + threadIdx.x; @@ -273,18 +274,23 @@ static __global__ void unary_gated_op_kernel(const T * x, const T * g, T * dst, const int64_t j1 = o0 == o1 ? j0 : (i / n) * o1 + (i % n); ggml_cuda_pdl_sync(); - dst[i] = (T)(op((float)x[j0]) * (float)g[j1]); + const T v = (T)(op((float)x[j0]) * (float)g[j1]); + dst[i] = v; + if (dst16) { + dst16[i] = ggml_cuda_f32_to_bf16_rne((float) v); + } } template -static void unary_gated_cuda(const T * x, const T * g, T * dst, const int64_t k, const int64_t n, const int64_t o0, const int64_t o1, cudaStream_t stream) { +static void unary_gated_cuda(const T * x, const T * g, T * dst, const int64_t k, const int64_t n, const int64_t o0, const int64_t o1, cudaStream_t stream, + uint16_t * dst16 = nullptr) { const int64_t num_blocks = (k + CUDA_GLU_BLOCK_SIZE - 1) / CUDA_GLU_BLOCK_SIZE; const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params((dim3)num_blocks, CUDA_GLU_BLOCK_SIZE, 0, stream); - ggml_cuda_kernel_launch(unary_gated_op_kernel, launch_params, x, g, dst, k, n, o0, o1); + ggml_cuda_kernel_launch(unary_gated_op_kernel, launch_params, x, g, dst, k, n, o0, o1, dst16); } template -void ggml_cuda_op_unary_gated(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { +void ggml_cuda_op_unary_gated(ggml_backend_cuda_context & ctx, ggml_tensor * dst, uint16_t * dst16 = nullptr) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; void * src0_d = src0->data; @@ -333,7 +339,7 @@ void ggml_cuda_op_unary_gated(ggml_backend_cuda_context & ctx, ggml_tensor * dst src1_p += swapped ? 0 : nc; } - unary_gated_cuda(src0_p, src1_p, (float *)dst_d, ggml_nelements(dst), nc, src0_o / sizeof(float), src1_o / sizeof(float), stream); + unary_gated_cuda(src0_p, src1_p, (float *)dst_d, ggml_nelements(dst), nc, src0_o / sizeof(float), src1_o / sizeof(float), stream, dst16); } } @@ -345,8 +351,8 @@ void ggml_cuda_op_geglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { ggml_cuda_op_unary_gated(ctx, dst); } -void ggml_cuda_op_swiglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { - ggml_cuda_op_unary_gated(ctx, dst); +void ggml_cuda_op_swiglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst, uint16_t * dst16) { + ggml_cuda_op_unary_gated(ctx, dst, dst16); } void ggml_cuda_op_geglu_erf(ggml_backend_cuda_context & ctx, ggml_tensor * dst) { @@ -702,6 +708,23 @@ void ggml_cuda_op_unary_mul(ggml_backend_cuda_context & ctx, ggml_tensor * unary } } +void ggml_cuda_op_cont_sigmoid_mul( + ggml_backend_cuda_context & ctx, ggml_tensor * cont_node, ggml_tensor * unary_node, ggml_tensor * mul_node, uint16_t * dst16) { + const ggml_tensor * gate = cont_node->src[0]; + const ggml_tensor * other = mul_node->src[0] == unary_node ? mul_node->src[1] : mul_node->src[0]; + + GGML_ASSERT(ggml_get_unary_op(unary_node) == GGML_UNARY_OP_SIGMOID); + GGML_ASSERT(gate->type == GGML_TYPE_F32 && other->type == GGML_TYPE_F32 && mul_node->type == GGML_TYPE_F32); + GGML_ASSERT(gate->nb[0] == sizeof(float) && gate->nb[2] == gate->nb[1] * gate->ne[1]); + GGML_ASSERT(ggml_is_contiguous(other) && ggml_is_contiguous(mul_node)); + GGML_ASSERT(ggml_nelements(gate) == ggml_nelements(other)); + + const int64_t k = ggml_nelements(other); + const int64_t n = gate->ne[0]; + unary_gated_cuda((const float *) gate->data, (const float *) other->data, (float *) mul_node->data, + k, n, gate->nb[1] / sizeof(float), n, ctx.stream(), dst16); +} + /* fused relu + sqr */ void ggml_cuda_op_relu_sqr(ggml_backend_cuda_context & ctx, ggml_tensor * relu_node, ggml_tensor * sqr_node) { @@ -719,3 +742,52 @@ void ggml_cuda_op_relu_sqr(ggml_backend_cuda_context & ctx, ggml_tensor * relu_n unary_cuda((const float *)src->data, (float *)sqr_node->data, k, stream); } } + +/* fused scale + unary */ + +template +static __global__ void scale_unary_kernel(const float * x, float * dst, const float scale, const float bias, const int k, uint16_t * dst16) { + const int i = blockDim.x*blockIdx.x + threadIdx.x; + + if (i >= k) { + return; + } + + // identical expression to scale_f32 in scale.cu, followed by the unary op + const float v = op(scale * x[i] + bias); + dst[i] = v; + if (dst16) { + dst16[i] = ggml_cuda_f32_to_bf16_rne(v); + } +} + +template +static void scale_unary_cuda(const float * x, float * dst, const float scale, const float bias, const int k, cudaStream_t stream, + uint16_t * dst16) { + const int num_blocks = (k + CUDA_NEG_BLOCK_SIZE - 1) / CUDA_NEG_BLOCK_SIZE; + const ggml_cuda_kernel_launch_params launch_params = ggml_cuda_kernel_launch_params((dim3)num_blocks, CUDA_NEG_BLOCK_SIZE, 0, stream); + ggml_cuda_kernel_launch(scale_unary_kernel, launch_params, x, dst, scale, bias, k, dst16); +} + +void ggml_cuda_op_scale_unary(ggml_backend_cuda_context & ctx, ggml_tensor * scale_node, ggml_tensor * unary_node, uint16_t * dst16) { + const ggml_tensor * src = scale_node->src[0]; + + GGML_ASSERT(src->type == GGML_TYPE_F32 && scale_node->type == GGML_TYPE_F32 && unary_node->type == GGML_TYPE_F32); + GGML_ASSERT(ggml_is_contiguous(src) && ggml_is_contiguous(unary_node)); + GGML_ASSERT(unary_node->src[0] == scale_node); + + const float scale = ggml_get_op_params_f32(scale_node, 0); + const float bias = ggml_get_op_params_f32(scale_node, 1); + const int k = ggml_nelements(src); + + switch (ggml_get_unary_op(unary_node)) { + case GGML_UNARY_OP_SILU: + scale_unary_cuda((const float *) src->data, (float *) unary_node->data, scale, bias, k, ctx.stream(), dst16); + break; + case GGML_UNARY_OP_SIGMOID: + scale_unary_cuda((const float *) src->data, (float *) unary_node->data, scale, bias, k, ctx.stream(), dst16); + break; + default: + GGML_ABORT("unsupported unary op for fused scale+unary"); + } +} diff --git a/ggml/src/ggml-cuda/unary.cuh b/ggml/src/ggml-cuda/unary.cuh index 04f3af6443a7..dddcabc5bed5 100644 --- a/ggml/src/ggml-cuda/unary.cuh +++ b/ggml/src/ggml-cuda/unary.cuh @@ -79,7 +79,8 @@ void ggml_cuda_op_reglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst); void ggml_cuda_op_geglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst); -void ggml_cuda_op_swiglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst); +// dst16: optional BF16 copy of dst (round to nearest even) for an MMB GEMM that reads it next +void ggml_cuda_op_swiglu(ggml_backend_cuda_context & ctx, ggml_tensor * dst, uint16_t * dst16 = nullptr); void ggml_cuda_op_swiglu_oai(ggml_backend_cuda_context & ctx, ggml_tensor * dst); @@ -93,6 +94,9 @@ void ggml_cuda_op_xielu(ggml_backend_cuda_context & ctx, ggml_tensor * dst); void ggml_cuda_op_unary_mul(ggml_backend_cuda_context & ctx, ggml_tensor * unary_node, ggml_tensor * mul_node); +void ggml_cuda_op_cont_sigmoid_mul( + ggml_backend_cuda_context & ctx, ggml_tensor * cont_node, ggml_tensor * unary_node, ggml_tensor * mul_node, uint16_t * dst16 = nullptr); + void ggml_cuda_op_relu_sqr(ggml_backend_cuda_context & ctx, ggml_tensor * relu_node, ggml_tensor * sqr_node); __device__ __forceinline__ float ggml_cuda_op_silu_single(float x) { @@ -121,3 +125,6 @@ __device__ __forceinline__ float ggml_cuda_op_swiglu_clamp_single(float gate, fl return ggml_cuda_op_silu_single(gate) * up; } + +// fused GGML_OP_SCALE + GGML_UNARY_OP_{SILU,SIGMOID}: dst = op(scale*x + bias) +void ggml_cuda_op_scale_unary(ggml_backend_cuda_context & ctx, ggml_tensor * scale_node, ggml_tensor * unary_node, uint16_t * dst16 = nullptr); diff --git a/ggml/src/ggml-cuda/vecdotq.cuh b/ggml/src/ggml-cuda/vecdotq.cuh index f2a6f2009c9a..b3e49e0072af 100644 --- a/ggml/src/ggml-cuda/vecdotq.cuh +++ b/ggml/src/ggml-cuda/vecdotq.cuh @@ -1247,6 +1247,60 @@ static __device__ __forceinline__ float vec_dot_iq3_s_q8_1( return d * sumi; } +// Two IQ3_S rows against the same activation block. The fused MoE gate/up path +// otherwise calls vec_dot_iq3_s_q8_1 twice and reloads all eight Q8_1 words. +// Keep each integer and floating-point accumulation independent so the results +// are bit-identical to the two scalar calls. +static __device__ __forceinline__ void vec_dot_iq3_s_q8_1_pair( + float & out0, float & out1, + const void * __restrict__ vbq0, const void * __restrict__ vbq1, + const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { + const block_iq3_s * bq0 = (const block_iq3_s *) vbq0 + kbx; + const block_iq3_s * bq1 = (const block_iq3_s *) vbq1 + kbx; + + const int2 qs0_packed = make_int2(get_int_b2(bq0->qs, iqs + 0), get_int_b2(bq0->qs, iqs + 1)); + const int2 qs1_packed = make_int2(get_int_b2(bq1->qs, iqs + 0), get_int_b2(bq1->qs, iqs + 1)); + const uint8_t * qs0 = (const uint8_t *) &qs0_packed; + const uint8_t * qs1 = (const uint8_t *) &qs1_packed; + + const int qh0 = bq0->qh[iqs/2]; + const int qh1 = bq1->qh[iqs/2]; + const int signs0_packed = get_int_b2(bq0->signs, iqs/2); + const int signs1_packed = get_int_b2(bq1->signs, iqs/2); + const uint8_t * signs0 = (const uint8_t *) &signs0_packed; + const uint8_t * signs1 = (const uint8_t *) &signs1_packed; + + int sumi0 = 0; + int sumi1 = 0; +#pragma unroll + for (int l0 = 0; l0 < 8; l0 += 2) { + 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); + + const int2 grid0 = make_int2( + iq3s_grid[qs0[l0 + 0] | ((qh0 << (8 - l0)) & 0x100)], + iq3s_grid[qs0[l0 + 1] | ((qh0 << (7 - l0)) & 0x100)]); + const int s00 = __vcmpne4(((signs0[l0/2] & 0x03) << 7) | ((signs0[l0/2] & 0x0C) << 21), 0x00000000); + const int s01 = __vcmpne4(((signs0[l0/2] & 0x30) << 3) | ((signs0[l0/2] & 0xC0) << 17), 0x00000000); + sumi0 = ggml_cuda_dp4a(__vsub4(grid0.x ^ s00, s00), u0, sumi0); + sumi0 = ggml_cuda_dp4a(__vsub4(grid0.y ^ s01, s01), u1, sumi0); + + const int2 grid1 = make_int2( + iq3s_grid[qs1[l0 + 0] | ((qh1 << (8 - l0)) & 0x100)], + iq3s_grid[qs1[l0 + 1] | ((qh1 << (7 - l0)) & 0x100)]); + const int s10 = __vcmpne4(((signs1[l0/2] & 0x03) << 7) | ((signs1[l0/2] & 0x0C) << 21), 0x00000000); + const int s11 = __vcmpne4(((signs1[l0/2] & 0x30) << 3) | ((signs1[l0/2] & 0xC0) << 17), 0x00000000); + sumi1 = ggml_cuda_dp4a(__vsub4(grid1.x ^ s10, s10), u0, sumi1); + sumi1 = ggml_cuda_dp4a(__vsub4(grid1.y ^ s11, s11), u1, sumi1); + } + + sumi0 *= 1 + 2*((bq0->scales[iqs/4] >> ((iqs << 1) & 0x04)) & 0x0F); + sumi1 *= 1 + 2*((bq1->scales[iqs/4] >> ((iqs << 1) & 0x04)) & 0x0F); + const float d8 = __low2float(bq8_1[iqs/2].ds); + out0 = (__half2float(bq0->d) * d8) * sumi0; + out1 = (__half2float(bq1->d) * d8) * sumi1; +} + #define VDR_IQ1_S_Q8_1_MMVQ 1 #define VDR_IQ1_S_Q8_1_MMQ 1 diff --git a/ggml/src/ggml-cuda/vendors/hip.h b/ggml/src/ggml-cuda/vendors/hip.h index 48d4eb2ce3e2..08bbca6a0d1c 100644 --- a/ggml/src/ggml-cuda/vendors/hip.h +++ b/ggml/src/ggml-cuda/vendors/hip.h @@ -143,6 +143,9 @@ #define cudaGraphExecUpdate hipGraphExecUpdate #define cudaStreamCaptureModeRelaxed hipStreamCaptureModeRelaxed #define cudaStreamBeginCapture hipStreamBeginCapture +#define cudaStreamIsCapturing hipStreamIsCapturing +#define cudaStreamCaptureStatus hipStreamCaptureStatus +#define cudaStreamCaptureStatusNone hipStreamCaptureStatusNone #define cudaGraph_t hipGraph_t #define cudaStream_t hipStream_t #define cudaSuccess hipSuccess diff --git a/ggml/src/ggml-impl.h b/ggml/src/ggml-impl.h index ae26e0c23b46..ef8e9bbea94a 100644 --- a/ggml/src/ggml-impl.h +++ b/ggml/src/ggml-impl.h @@ -739,7 +739,7 @@ static inline bool ggml_can_fuse_subgraph(const struct ggml_cgraph * cgraph, const enum ggml_op * ops, const int * outputs, int num_outputs) { - GGML_ASSERT(count < 32); + GGML_ASSERT(count <= 32); if (node_idx + count > cgraph->n_nodes) { return false; } diff --git a/ggml/src/ggml-quants.c b/ggml/src/ggml-quants.c index 1ebc50a763f1..dd28eea5266a 100644 --- a/ggml/src/ggml-quants.c +++ b/ggml/src/ggml-quants.c @@ -2,6 +2,8 @@ #include "ggml-common.h" #include "ggml-quants.h" +#include "../rocmfp4/rocmfp4.h" +#include "../rocmfpx/rocmfpx.h" #include "ggml-impl.h" #include "ggml-cpu/ggml-cpu-impl.h" #include "ggml-cpu.h" @@ -611,6 +613,30 @@ void dequantize_row_nvfp4(const block_nvfp4 * GGML_RESTRICT x, float * GGML_REST } } +void dequantize_row_q4_0_rocmfp4(const block_rocmfp4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + rocmfp4_dequantize_row_q4_0(x, y, k); +} + +void dequantize_row_q4_0_rocmfp4_fast(const block_rocmfp4_fast * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + rocmfp4_dequantize_row_q4_0_fast(x, y, k); +} + +void dequantize_row_q2_0_rocmfpx(const block_rocmfp2 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + rocmfpx_dequantize_row_fp2(x, y, k); +} + +void dequantize_row_q3_0_rocmfpx(const block_rocmfp3 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + rocmfpx_dequantize_row_fp3(x, y, k); +} + +void dequantize_row_q6_0_rocmfpx(const block_rocmfp6 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + rocmfpx_dequantize_row_fp6(x, y, k); +} + +void dequantize_row_q8_0_rocmfpx(const block_rocmfp8 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k) { + rocmfpx_dequantize_row_fp8(x, y, k); +} + // // 2-6 bit quantization in super-blocks // @@ -5541,6 +5567,48 @@ bool ggml_validate_row_data(enum ggml_type type, const void * data, size_t nbyte { VALIDATE_ROW_DATA_D_F16_IMPL(block_q4_0, data, nb); } break; + case GGML_TYPE_Q4_0_ROCMFP4: + { + if (!rocmfp4_validate_row_data(data, nbytes)) { + fprintf(stderr, "%s: invalid ROCmFP4 row data\n", __func__); + return false; + } + } break; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + { + if (!rocmfp4_validate_row_data_fast(data, nbytes)) { + fprintf(stderr, "%s: invalid ROCmFP4 fast row data\n", __func__); + return false; + } + } break; + case GGML_TYPE_Q3_0_ROCMFPX: + { + if (!rocmfpx_validate_row_data_fp3(data, nbytes)) { + fprintf(stderr, "%s: invalid ROCmFPx FP3 row data\n", __func__); + return false; + } + } break; + case GGML_TYPE_Q2_0_ROCMFPX: + { + if (!rocmfpx_validate_row_data_fp2(data, nbytes)) { + fprintf(stderr, "%s: invalid ROCmFPx FP2 row data\n", __func__); + return false; + } + } break; + case GGML_TYPE_Q6_0_ROCMFPX: + { + if (!rocmfpx_validate_row_data_fp6(data, nbytes)) { + fprintf(stderr, "%s: invalid ROCmFPx FP6 row data\n", __func__); + return false; + } + } break; + case GGML_TYPE_Q8_0_ROCMFPX: + { + if (!rocmfpx_validate_row_data_fp8(data, nbytes)) { + fprintf(stderr, "%s: invalid ROCmFPx FP8 row data\n", __func__); + return false; + } + } break; case GGML_TYPE_Q4_1: { VALIDATE_ROW_DATA_DM_F16_IMPL(block_q4_1, data, nb, d, m); diff --git a/ggml/src/ggml-quants.h b/ggml/src/ggml-quants.h index 75188f1af180..9282654d059c 100644 --- a/ggml/src/ggml-quants.h +++ b/ggml/src/ggml-quants.h @@ -5,6 +5,10 @@ #include "ggml.h" +// ROCmFPx block layouts, for the dequantize_row_* wrappers below +#include "../rocmfp4/rocmfp4.h" +#include "../rocmfpx/rocmfpx.h" + // GGML internal header #ifdef __cplusplus @@ -55,6 +59,16 @@ GGML_API void dequantize_row_q8_0(const block_q8_0 * GGML_RESTRICT x, float * GG GGML_API void dequantize_row_mxfp4(const block_mxfp4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); GGML_API void dequantize_row_nvfp4(const block_nvfp4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +// ROCmFPx types under the ggml naming convention (dequantize_row_), +// so tooling that resolves dequantizers by type name (gguf-py/tests/test_quants.py) +// reaches them. The codecs themselves live in ggml/rocmfp4 and ggml/rocmfpx. +GGML_API void dequantize_row_q4_0_rocmfp4 (const block_rocmfp4 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void dequantize_row_q4_0_rocmfp4_fast(const block_rocmfp4_fast * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void dequantize_row_q2_0_rocmfpx (const block_rocmfp2 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void dequantize_row_q3_0_rocmfpx (const block_rocmfp3 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void dequantize_row_q6_0_rocmfpx (const block_rocmfp6 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); +GGML_API void dequantize_row_q8_0_rocmfpx (const block_rocmfp8 * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); + GGML_API void dequantize_row_q2_K(const block_q2_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); GGML_API void dequantize_row_q3_K(const block_q3_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); GGML_API void dequantize_row_q4_K(const block_q4_K * GGML_RESTRICT x, float * GGML_RESTRICT y, int64_t k); diff --git a/ggml/src/ggml-vulkan/ggml-vulkan.cpp b/ggml/src/ggml-vulkan/ggml-vulkan.cpp index f936127a6be6..e8f4b3884eb5 100644 --- a/ggml/src/ggml-vulkan/ggml-vulkan.cpp +++ b/ggml/src/ggml-vulkan/ggml-vulkan.cpp @@ -64,6 +64,7 @@ typedef struct VkPhysicalDeviceCooperativeMatrixDecodeVectorFeaturesNV { #include #include #include +#include #include #include #include @@ -258,34 +259,27 @@ typedef std::weak_ptr vk_pipeline_ref; static void ggml_vk_destroy_pipeline(vk::Device& device, vk_pipeline& pipeline); -struct vk_matmul_pipeline_key { - ggml_type type_a; - ggml_type type_b; - bool mul_mat_id; - bool f16acc; - - bool operator<(const vk_matmul_pipeline_key & o) const { - return std::tie(type_a, type_b, mul_mat_id, f16acc) - < std::tie(o.type_a, o.type_b, o.mul_mat_id, o.f16acc); +struct vk_matmul_pipeline_struct { + vk_pipeline l, m, s; + vk_pipeline a_l, a_m, a_s; + // Returns true when all unaligned pipelines are null. + // We only check for unaligned variants since one of the unaligned pipelines must exist + // while aligned pipelines are optional + bool is_empty() const { + return l == nullptr && m == nullptr && s == nullptr; } }; +typedef std::shared_ptr vk_matmul_pipeline; -struct vk_matmul_pipeline_pair { - vk_pipeline unaligned; - vk_pipeline aligned; - uint32_t align; -}; - -struct vk_tile_config { - std::vector warptile; - std::array wg_denoms; - uint32_t align; +struct vk_matmul_pipeline2 { + vk_matmul_pipeline2() { + f16acc = std::make_shared(); + f32acc = std::make_shared(); + } + vk_matmul_pipeline f32acc; + vk_matmul_pipeline f16acc; }; -using matmul_tile_selector_t = std::function& configs)>; - struct vk_device_struct; typedef std::shared_ptr vk_device; typedef std::weak_ptr vk_device_ref; @@ -333,7 +327,6 @@ static void ggml_vk_print_device_lost_info(const vk_device& device); struct vk_queue_handle { vk::Queue queue; vk_device_ref device; - std::mutex * device_submit_mutex = nullptr; virtual void submit(vk::ArrayProxy submits, vk::Fence fence) = 0; virtual void lock() {} // no-op by default (internally synchronized case) virtual void unlock() {} @@ -343,11 +336,6 @@ struct vk_queue_handle { struct vk_queue_handle_synchronized : vk_queue_handle { std::mutex mutex; void submit(vk::ArrayProxy submits, vk::Fence fence) override { - // Workaround for NVIDIA driver bug - std::unique_lock device_guard; - if (device_submit_mutex) { - device_guard = std::unique_lock(*device_submit_mutex); - } std::lock_guard guard(mutex); try { queue.submit(submits, fence); @@ -362,14 +350,9 @@ struct vk_queue_handle_synchronized : vk_queue_handle { void unlock() override { mutex.unlock(); } }; -// Driver guarantees internal synchronization via VK_KHR_internally_synchronized_queues struct vk_queue_handle_unsynchronized : vk_queue_handle { void submit(vk::ArrayProxy submits, vk::Fence fence) override { - // Workaround for NVIDIA driver bug - std::unique_lock device_guard; - if (device_submit_mutex) { - device_guard = std::unique_lock(*device_submit_mutex); - } + // Driver guarantees internal synchronization via VK_KHR_internally_synchronized_queues try { queue.submit(submits, fence); } catch (vk::DeviceLostError &) { @@ -844,9 +827,24 @@ static bool ggml_vk_lightning_indexer_k_type_supported(ggml_type type) { return std::find(lightning_indexer_k_types.begin(), lightning_indexer_k_types.end(), type) != lightning_indexer_k_types.end(); } +// Indexer head counts we build kernels for: 4 = Qwen4-exp (qwen4_exp), 64 = DeepSeek-V4 and +// GLM-DSA, 32 kept because test-backend-ops exercises it. N_HEAD is a specialization constant, +// so each entry is a separately compiled kernel with the head loop fully unrolled - as a push +// constant the loop cannot unroll and decode costs ~55% more (measured, gfx1151, kv=131584). +static constexpr uint32_t LI_NH_VALUES[] = { 4, 32, 64 }; +#define LI_NH_COUNT (sizeof(LI_NH_VALUES) / sizeof(LI_NH_VALUES[0])) + +static int ggml_vk_li_nh_index(int64_t nh) { + for (size_t i = 0; i < LI_NH_COUNT; ++i) { + if ((int64_t) LI_NH_VALUES[i] == nh) { + return (int) i; + } + } + return -1; +} + struct vk_device_struct { std::recursive_mutex mutex; - std::mutex queue_submit_mutex; mutable std::shared_mutex pinned_memory_mutex; // Guards compile_pending, all_pipelines, and the dynamic pipeline maps @@ -901,6 +899,7 @@ struct vk_device_struct { bool add_rms_fusion; uint32_t partials_binding_alignment; uint32_t max_nodes_per_submit; + uint64_t max_bytes_per_submit; bool shader_64b_indexing; @@ -970,9 +969,27 @@ struct vk_device_struct { vk::DescriptorSetLayout dsl; - std::map> pipeline_matmul; - matmul_tile_selector_t matmul_tile_selector; - matmul_tile_selector_t matmul_id_tile_selector; + vk_matmul_pipeline pipeline_matmul_f32 {}; + vk_matmul_pipeline pipeline_matmul_f32_f16 {}; + vk_matmul_pipeline pipeline_matmul_bf16 {}; + vk_matmul_pipeline2 pipeline_matmul_f16; + vk_matmul_pipeline2 pipeline_matmul_f16_f32; + + vk_matmul_pipeline2 pipeline_dequant_mul_mat_mat[GGML_TYPE_COUNT]; + vk_matmul_pipeline2 pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_COUNT]; + vk_matmul_pipeline2 pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_COUNT]; + + vk_matmul_pipeline pipeline_matmul_id_f32 {}; + vk_matmul_pipeline pipeline_matmul_id_bf16 {}; + vk_matmul_pipeline2 pipeline_matmul_id_f16; + vk_matmul_pipeline2 pipeline_matmul_id_f16_f32; + + vk_matmul_pipeline2 pipeline_dequant_mul_mat_mat_id[GGML_TYPE_COUNT]; + vk_matmul_pipeline2 pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_COUNT]; + // EXPERIMENT (GGML_VK_MMID_F16B=1): f16-B mul_mat_id pipelines on KHR_coopmat + // devices (upstream only builds f32-B there). Populated only when the env flag + // is set; empty otherwise. + vk_matmul_pipeline2 pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_COUNT]; vk_pipeline pipeline_matmul_split_k_reduce; vk_pipeline pipeline_quantize_q8_1_x4; @@ -1013,6 +1030,7 @@ struct vk_device_struct { vk_pipeline pipeline_add_id_f32; vk_pipeline pipeline_concat_i8, pipeline_concat_i16, pipeline_concat_i32, pipeline_concat_i64; + vk_pipeline pipeline_concat_transpose_i32; vk_pipeline pipeline_upscale_nearest_f32, pipeline_upscale_bilinear_f32, pipeline_upscale_bicubic_f32, pipeline_upscale_bilinear_antialias_f32; vk_pipeline pipeline_scale_f32; vk_pipeline pipeline_log[2]; @@ -1128,9 +1146,6 @@ struct vk_device_struct { vk_pipeline pipeline_cumsum_multipass2_f32; vk_pipeline pipeline_argmax_f32; vk_pipeline pipeline_count_equal_i32; - vk_pipeline pipeline_dsv4_hc_comb_f32; - vk_pipeline pipeline_dsv4_hc_pre_f32; - vk_pipeline pipeline_dsv4_hc_post_f32; std::map pipeline_solve_tri_f32; vk_pipeline pipeline_im2col_f32, pipeline_im2col_f32_f16; vk_pipeline pipeline_im2col_3d_f32, pipeline_im2col_3d_f32_f16; @@ -1151,6 +1166,21 @@ struct vk_device_struct { vk_pipeline pipeline_lightning_indexer_f32[GGML_TYPE_COUNT]; // [size_idx][kda] where size_idx: 0=d16, 1=d32, 2=d64, 3=d128 vk_pipeline pipeline_gated_delta_net[4][2]; + // One pipeline per supported indexer head count; see LI_NH_VALUES. + vk_pipeline pipeline_lightning_indexer_f16[LI_NH_COUNT]; + vk_pipeline pipeline_lightning_indexer_cm_f16[LI_NH_COUNT]; + vk_pipeline pipeline_lightning_indexer_cm_small_f16[LI_NH_COUNT]; + vk_pipeline pipeline_lightning_indexer_decode_cm_f16[LI_NH_COUNT]; + vk_pipeline pipeline_flash_attn_top_k_f16; + vk_pipeline pipeline_flash_attn_top_k_cm_f16; + vk_pipeline pipeline_flash_attn_gather_f16; + vk_pipeline pipeline_flash_attn_gather_dq[GGML_TYPE_COUNT]; + vk_pipeline pipeline_flash_attn_union_f16; + vk_pipeline pipeline_flash_attn_gather_union_f16; + vk_pipeline pipeline_flash_attn_gather_union_dq[GGML_TYPE_COUNT]; + vk_pipeline pipeline_dsv4_hc_pre_f32; + vk_pipeline pipeline_dsv4_hc_comb_f32; + vk_pipeline pipeline_dsv4_hc_post_f32; vk_pipeline pipeline_ssm_scan_f32_d128; vk_pipeline pipeline_ssm_scan_f32_d256; vk_pipeline pipeline_ssm_conv_f32; @@ -1171,10 +1201,6 @@ struct vk_device_struct { std::map, vk_pipeline> pipeline_fa_mask_opt; - vk_pipeline pipeline_fa_sparse_compact; - vk_pipeline pipeline_fa_sparse_compact_subgroup; - bool fa_sparse_compact_use_subgroups; - vk_pipeline pipeline_flash_attn_split_k_reduce; vk_pipeline pipeline_count_experts; @@ -1418,6 +1444,7 @@ struct vk_mat_mat_id_push_constants { uint32_t nei0; uint32_t nei1; uint32_t nbi1; uint32_t ne11; uint32_t n_experts; uint32_t hoist_row_ids; + uint32_t fusion_flags; }; struct vk_mat_vec_id_push_constants { uint32_t ncols; @@ -1492,53 +1519,6 @@ struct vk_op_fwht_push_constants { float scale; }; -struct vk_op_dsv4_hc_comb_push_constants { - uint32_t n_tokens; - - uint32_t nbm0; uint32_t nbm1; - uint32_t nbs0; - uint32_t nbb0; - uint32_t nbd0; uint32_t nbd1; uint32_t nbd2; - - uint32_t m_offset; - uint32_t s_offset; - uint32_t b_offset; - uint32_t d_offset; - - float eps; - uint32_t n_iter; -}; - -struct vk_op_dsv4_hc_pre_push_constants { - uint32_t n_embd; - uint32_t n_tokens; - - uint32_t nbx0; uint32_t nbx1; uint32_t nbx2; - uint32_t nbw0; uint32_t nbw1; - uint32_t nbd0; uint32_t nbd1; - - uint32_t x_offset; - uint32_t w_offset; - uint32_t d_offset; -}; - -struct vk_op_dsv4_hc_post_push_constants { - uint32_t n_embd; - uint32_t n_tokens; - - uint32_t nbx0; uint32_t nbx1; - uint32_t nbr0; uint32_t nbr1; uint32_t nbr2; - uint32_t nbp0; uint32_t nbp1; - uint32_t nbc0; uint32_t nbc1; uint32_t nbc2; - uint32_t nbd0; uint32_t nbd1; uint32_t nbd2; - - uint32_t x_offset; - uint32_t r_offset; - uint32_t p_offset; - uint32_t c_offset; - uint32_t d_offset; -}; - struct vk_op_count_experts_push_constants { uint32_t ne00; uint32_t ne01; @@ -2031,6 +2011,73 @@ struct vk_op_gated_delta_net_push_constants { uint32_t K; }; +// push constants for the fork's wave64 f16 lightning-indexer kernels (scalar-64 + CM family) +struct vk_op_lightning_indexer_cm_push_constants { + uint32_t n_kv, n_batch, n_stream, nem3; + uint32_t nb1, nb3; + uint32_t nbq1, nbq2, nbq3; + uint32_t nbk2, nbk3; + uint32_t nbw1, nbw3; + uint32_t nbm1, nbm3; +}; +static_assert(sizeof(vk_op_lightning_indexer_cm_push_constants) <= 128); + +struct vk_op_dsv4_hc_pre_push_constants { + uint32_t n_embd, hc, nr; + uint32_t sx0, sx1, sx2; + uint32_t sw0, sw1; + uint32_t sd0, sd1; +}; +static_assert(sizeof(vk_op_dsv4_hc_pre_push_constants) <= 128); + +struct vk_op_dsv4_hc_comb_push_constants { + uint32_t n_tokens; + uint32_t sm0, sm1; + uint32_t ss0; + uint32_t sb0; + uint32_t sd0, sd1, sd2; + float eps; + int32_t n_iter; +}; +static_assert(sizeof(vk_op_dsv4_hc_comb_push_constants) <= 128); + +struct vk_op_dsv4_hc_post_push_constants { + uint32_t n_embd, hc, nr; + uint32_t sx0, sx1; + uint32_t sr0, sr1, sr2; + uint32_t sp0, sp1; + uint32_t sc0, sc1, sc2; + uint32_t sd0, sd1, sd2; +}; +static_assert(sizeof(vk_op_dsv4_hc_post_push_constants) <= 128); + +struct vk_op_flash_attn_union_push_constants { + uint32_t n_kv, n_kv_raw, n_batch, n_top_k, max_union, nbt1, max_words, pad_to, count_only; +}; +// nbk1/nbk3 are in 4-byte WORDS, not elements: the gather relocates K rows verbatim and never +// interprets what is in them, so it works for any type whose row is a whole number of words. +struct vk_op_flash_attn_gather_union_push_constants { + uint32_t n_kv, n_kv_raw, kv_c_max, nbk1, nbm1, n_batch, row_words; +}; +struct vk_op_flash_attn_gather_push_constants { + uint32_t n_kv, n_kv_raw, n_top_k, kv_c; + uint32_t nbk1, nbk3, nbt1, nbt3, nbm1, nbm3, nem3, n_batch, row_words; +}; +static_assert(sizeof(vk_op_flash_attn_gather_push_constants) <= 128); + +struct vk_op_flash_attn_top_k_push_constants { + uint32_t n_batch, n_kv, n_kv_raw, n_top_k, n_head; + uint32_t nbq1, nbq2, nbq3; + uint32_t nbk1, nbk3; + uint32_t nbm1, nbm3; + uint32_t nbt1, nbt3; + uint32_t nb1, nb2, nb3; + float scale; + uint32_t has_sinks; + uint32_t split_mode; +}; +static_assert(sizeof(vk_op_flash_attn_top_k_push_constants) <= 128); + struct vk_op_ssm_scan_push_constants { uint32_t nb02, nb03, nb12, nb13; uint32_t nb21, nb22, nb31; @@ -2182,7 +2229,11 @@ struct vk_quantize_q8_1_push_constants { struct vk_op_flash_attn_split_k_reduce_push_constants { uint32_t D; uint32_t ne1; + // ne2 describes the SPLIT BUFFER (which may cover only a tile of queries); dst_ne2 is the + // destination's query count. They differ only when the caller tiles the split buffer, and + // the destination stride between streams must always use the full count. uint32_t ne2; + uint32_t dst_ne2; uint32_t ne3; uint32_t k_num; uint32_t sinks; @@ -2200,16 +2251,6 @@ struct vk_op_flash_attn_mask_opt_push_constants { uint32_t nbd3; }; -struct vk_op_flash_attn_sparse_compact_push_constants { - uint32_t KV; - uint32_t nem1; - uint32_t nem2; - uint32_t nbm1; - uint32_t nbm2; - uint32_t nbm3; - uint32_t n_kv_max; -}; - // Allow pre-recording command buffers struct vk_staging_memcpy { vk_staging_memcpy(void * _dst, const void * _src, size_t _n) : dst(_dst), src(_src), n(_n) {} @@ -2303,6 +2344,20 @@ static bool vk_enable_sync_logger = false; static uint32_t vk_perf_logger_frequency = 1; static std::string vk_pipeline_stats_filter; +// Total memory traffic of a node (dst + srcs). Used to bound command buffer +// execution time for bandwidth-bound ops with no flops estimate (large copies, +// set_rows, mask fills at long context) - packing too many of them into one +// submission can exceed the driver timeout. +static uint64_t ggml_vk_get_node_bytes(const ggml_tensor * node) { + uint64_t bytes = ggml_nbytes(node); + for (int i = 0; i < GGML_MAX_SRC; i++) { + if (node->src[i]) { + bytes += ggml_nbytes(node->src[i]); + } + } + return bytes; +} + static uint64_t ggml_vk_get_node_flops(const ggml_tensor * node) { if (node->op == GGML_OP_MUL_MAT || node->op == GGML_OP_MUL_MAT_ID) { const uint64_t m = node->ne[0]; @@ -2421,6 +2476,20 @@ class vk_perf_logger { if (node->op == GGML_OP_UNARY) { return fusion_str + ggml_unary_op_name(ggml_get_unary_op(node)); } + if (node->op == GGML_OP_MUL && getenv("GGML_VK_PERF_SHAPES")) { + std::string name = "MUL "; + name += "dst(" + std::to_string(node->ne[0]) + "," + std::to_string(node->ne[1]) + "," + + std::to_string(node->ne[2]) + ") b(" + std::to_string(node->src[1]->ne[0]) + "," + + std::to_string(node->src[1]->ne[1]) + "," + std::to_string(node->src[1]->ne[2]) + ")"; + name += std::string(" a=") + ggml_op_name(node->src[0]->op); + if (node->src[0]->op == GGML_OP_UNARY) { name += std::string(":") + ggml_unary_op_name(ggml_get_unary_op(node->src[0])); } + name += std::string(" b=") + ggml_op_name(node->src[1]->op); + if (node->src[1]->op == GGML_OP_UNARY) { name += std::string(":") + ggml_unary_op_name(ggml_get_unary_op(node->src[1])); } + if (node->src[1]->op == GGML_OP_RESHAPE && node->src[1]->src[0]) { + name += std::string("(") + ggml_op_name(node->src[1]->src[0]->op) + ")"; + } + return fusion_str + name; + } if (node->op == GGML_OP_MUL_MAT || node->op == GGML_OP_MUL_MAT_ID) { const uint64_t m = node->ne[0]; const uint64_t n = node->ne[1]; @@ -2493,6 +2562,13 @@ class vk_perf_logger { timings[name].push_back(time); } + // Log a sub-node interval under a caller-supplied name. Used for work a node's handler + // dispatches before the op itself (e.g. the FA K/V contiguize/dequant pass), which would + // otherwise be billed to the op and invisible. No flops: these move bytes, not math. + void log_timing_named(const char *name, uint64_t time) { + timings[std::string(name)].push_back(time); + } + void log_timing(const std::vector &nodes, const std::vector &names, uint64_t time) { uint64_t total_flops = 0; std::string name; @@ -2526,6 +2602,18 @@ struct ggml_backend_vk_context { ggml_vk_garbage_collector gc; size_t prealloc_size_x, prealloc_size_y, prealloc_size_split_k, prealloc_size_add_rms_partials, prealloc_size_add_rms_partials_offset; vk_buffer prealloc_x, prealloc_y, prealloc_split_k, prealloc_add_rms_partials, sync_staging; + // memoized capacity decision for the FA dequant-once scratch, see ggml_vk_fa_dequant_scratch_fits + uint64_t fa_dequant_gate_sz; + bool fa_dequant_gate_fits; + bool fa_dequant_gate_logged; + // DeepSeek V4 small-batch union: the compact row count is produced on the device, so the + // host prices compaction from the last count it wrote. See ggml_vk_fa_union_estimate. + // Per batch size, because the overlap depends on it (measured 0.64 at 2 tokens, 0.40 at 4, + // 0.24 at 8) and because a speculative decode varies the batch with the accept count, so a + // single slot would be invalidated on nearly every step. This path caps the batch at 64. + vk_buffer fa_union_stat; + float fa_union_est_ratio[64]; // union / candidates, decaying peak; 0 = unseeded + uint64_t fa_union_declines; vk::Fence fence, almost_ready_fence; bool submit_pending {}; bool almost_ready_fence_pending {}; @@ -2586,6 +2674,9 @@ struct ggml_backend_vk_context { std::vector query_fusion_node_count; std::vector query_nodes; std::vector query_node_idx; + // non-null => this query slot closes a sub-node interval logged under this literal name, + // not a graph node. See ggml_vk_perf_mark_subop. + std::vector query_sub_names; int32_t num_queries {}; int32_t query_idx {}; }; @@ -2715,32 +2806,6 @@ template <> void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk GGML_UNUSED(src3); } -template <> void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk_op_dsv4_hc_comb_push_constants &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, const ggml_tensor * src3, ggml_tensor * dst) { - p.m_offset = get_misalign_bytes(ctx, src0) / ggml_type_size(src0->type); - p.s_offset = get_misalign_bytes(ctx, src1) / ggml_type_size(src1->type); - p.b_offset = get_misalign_bytes(ctx, src2) / ggml_type_size(src2->type); - p.d_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type); - - GGML_UNUSED(src3); -} - -template <> void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk_op_dsv4_hc_pre_push_constants &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, const ggml_tensor * src3, ggml_tensor * dst) { - p.x_offset = get_misalign_bytes(ctx, src0) / ggml_type_size(src0->type); - p.w_offset = get_misalign_bytes(ctx, src1) / ggml_type_size(src1->type); - p.d_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type); - - GGML_UNUSED(src2); - GGML_UNUSED(src3); -} - -template <> void init_pushconst_tensor_offsets(ggml_backend_vk_context * ctx, vk_op_dsv4_hc_post_push_constants &p, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, const ggml_tensor * src3, ggml_tensor * dst) { - p.x_offset = get_misalign_bytes(ctx, src0) / ggml_type_size(src0->type); - p.r_offset = get_misalign_bytes(ctx, src1) / ggml_type_size(src1->type); - p.p_offset = get_misalign_bytes(ctx, src2) / ggml_type_size(src2->type); - p.c_offset = get_misalign_bytes(ctx, src3) / ggml_type_size(src3->type); - p.d_offset = get_misalign_bytes(ctx, dst) / ggml_type_size(dst->type); -} - struct ggml_backend_vk_buffer_context { vk_device_ref device; vk_buffer dev_buffer; @@ -3546,10 +3611,6 @@ static std::unique_ptr ggml_vk_create_queue(vk_device& device, uint32_ h->queue = device->device.getQueue2(queue_info2); h->device = device; - // Avoid concurrent submissions on NVIDIA due to driver bug. - if (device->vendor_id == VK_VENDOR_ID_NVIDIA) { - h->device_submit_mutex = &device->queue_submit_mutex; - } q->handle = h; q->cmd_pool.init(device, q.get()); @@ -4050,8 +4111,46 @@ static vk_fa_tuning_params get_fa_tuning_params_coopmat1(const vk_device& device result.block_cols = coopmat_block_cols * num_subgroups; result.row_split = num_subgroups; result.subgroup_size = device->subgroup_size; + + // Pin a 32-wide subgroup only where narrowing is free. The shader derives cols_per_iter, + // threads_per_rowgroup and every strided load loop from gl_WorkGroupSize.x, and + // workgroup_size is num_subgroups * subgroup_size, so threads_per_rowgroup always equals the + // real subgroup size. Halving the subgroup halves the workgroup, and the per-lane O state + // grows as d_per_thread = ceil((HSV/4) / threads_per_rowgroup). Pin only when that count is + // unchanged. Above that point the narrow subgroup issues roughly 1.5x to 1.8x the + // instructions for the same number of SIMD passes, which loses on an issue-bound kernel: + // hd256 measures 6 to 18 percent slower. The test depends on HSV only; HSK does not enter + // d_per_thread. On a 64-wide device it reduces exactly to hsv <= 128. + // On by default (=1, applies the rule); =0 disables. =2 forces the pin regardless of + // head size, for measuring the configurations the rule rejects. Diagnostic only. + static const int fa_wave32 = [] { + const char * e = getenv("GGML_VK_FA_WAVE32"); + return e ? atoi(e) : 1; + }(); + // Multi-row dispatches only. Decode dispatches carry N = gqa_ratio rows (8 on + // Qwen3-Coder-30B) and there the narrow subgroup loses: tg64 at d8192/d32768 measured + // 3.5 to 4 percent slower with the pin (q8_0 KV), while prefill gains up to 10 percent + // at d8192. Keep the pin to the prefill shapes it was measured on. + if (fa_wave32 != 0 && + n_rows >= 32 && + device->subgroup_size_control && + 32 < device->subgroup_size && // narrow only, never widen + device->subgroup_min_size <= 32 && 32 <= device->subgroup_max_size && + (result.block_cols % 32) == 0 && // cols_per_thread stays >= 1 + (result.block_cols * result.block_rows / 4) >= num_subgroups * 32 && // mask_cache != 0 + (fa_wave32 == 2 || + CEIL_DIV(hsv / 4, 32u) == CEIL_DIV(hsv / 4, device->subgroup_size))) { + result.subgroup_size = 32; + } + result.workgroup_size = num_subgroups * result.subgroup_size; + // threads_per_rowgroup == the real subgroup size is load-bearing in three places: + // the subgroupMax row reduction, the subgroupAdd of Lf, and tmpsh[gl_SubgroupID], which is + // sized by row_split and would be written out of bounds if gl_NumSubgroups exceeded it. + GGML_ASSERT(result.workgroup_size == result.row_split * result.subgroup_size); + GGML_ASSERT(result.block_cols % result.subgroup_size == 0); + const uint32_t D_lsb = D ^ (D & (D-1)); // extract lowest set bit result.d_split = std::min(std::min(result.subgroup_size, 8u), D_lsb / 4); @@ -4133,7 +4232,8 @@ static vk_fa_tuning_params get_fa_tuning_params(const vk_device& device, uint32_ } static vk_fa_pipeline_state get_fa_pipeline_state(const vk_device& device, const vk_fa_tuning_params& params, uint32_t hsk, uint32_t hsv, bool aligned, bool f32acc, - bool use_mask, bool use_mask_opt, bool use_logit_softcap, bool use_sparse, ggml_type k_type, ggml_type v_type) { + bool use_mask, bool use_mask_opt, bool use_logit_softcap, ggml_type k_type, ggml_type v_type, + bool use_dynamic_kv = false) { const bool old_amd_windows = device->vendor_id == VK_VENDOR_ID_AMD && device->driver_id == vk::DriverId::eAmdProprietary && (device->architecture == AMD_GCN || device->architecture == AMD_RDNA1 || device->architecture == AMD_RDNA2); @@ -4141,7 +4241,7 @@ static vk_fa_pipeline_state get_fa_pipeline_state(const vk_device& device, const (use_mask ? 2 : 0) | (use_logit_softcap ? 4 : 0) | (old_amd_windows ? 8 : 0) | - (use_sparse ? 16 : 0); + (use_dynamic_kv ? 16 : 0); const uint32_t subgroup_size = params.disable_subgroups ? 0 : params.subgroup_size; @@ -4178,6 +4278,46 @@ static std::vector get_fa_spec_constants(const vk_fa_pipeline_state& s }; } +// SHMEM_STRIDE_PAD for mul_mm.comp (constantID=12), in FLOAT_TYPEV2 units. +// +// mul_mm.comp stages the shared A/B tiles at SHMEM_STRIDE = BK/2 + pad. Each element is one +// dword, so the stride maps 1:1 onto the 32 LDS banks and gcd(SHMEM_STRIDE, 32) is the conflict +// factor. With BK=32 the shader default of 4 gives stride 20 (gcd 4, four-way conflict). Pad 2 +// gives stride 18 (gcd 2, 16 banks) and is a large win on the quantised path on gfx1151. +// +// But the coopmat path passes SHMEM_STRIDE to coopMatLoad as its Stride operand, and +// VUID-RuntimeSpirv-OpCooperativeMatrixLoadKHR-08986 needs a 16 byte aligned stride for the +// 16x16 f16 tiles. Stride bytes are (BK/2 + pad) * 4, so only pad % 4 == 0 is in contract. +// RADV before 25.3 lowers coopMatLoad to ds_read_b128, which the contract entitles it to, and +// the misaligned rows then pay runtime splits: pp512 drops by more than 2x. RADV 25.3 and later +// lowers to ds_read_b64, which the 72 byte stride always suits, so the extra bank spread wins. +// Use pad 2 only where it is measured to win. +// +// The float path keeps pad 4 on measurement, not theory: the f32/f16 shaders #define BK 32 +// whatever the host pushes, so they run the same stride, yet pad 2 measures worse on both. The +// host warptile's BK labels the pipeline even where the shader overrides the value. +// +// GGML_VK_SHMEM_PAD=N overrides both paths, for probing. +static uint32_t ggml_vk_coopmat_shmem_pad(const vk_device& device, uint32_t bk) { + if (device->vendor_id == VK_VENDOR_ID_INTEL && device->coopmat_support && + device->driver_id == vk::DriverId::eIntelProprietaryWindows) { + return 0; + } + static const int env_pad = [] { + const char * e = getenv("GGML_VK_SHMEM_PAD"); + return e ? atoi(e) : -1; + }(); + if (env_pad >= 0) { + return (uint32_t) env_pad; + } + if (device->coopmat_support && bk >= 32 && + device->driver_id == vk::DriverId::eMesaRadv && + device->properties.driverVersion >= VK_MAKE_API_VERSION(0, 25, 3, 0)) { + return 2; + } + return 4; // mul_mm.comp default +} + static bool ggml_vk_matmul_shmem_support(const vk_device& device, const std::vector& warptile, bool mul_mat_id, ggml_type src0_type) { uint32_t lut_size = 0; @@ -4212,15 +4352,26 @@ static bool ggml_vk_matmul_shmem_support(const vk_device& device, const std::vec // Same kvalues budget as MXFP4 plus ue4m3_fp32_lut[128] (types.glsl, DATA_A_NVFP4). lut_size = 4*16 + 128u * (uint32_t)sizeof(float); break; + case GGML_TYPE_Q4_0_ROCMFP4: + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + // kvalues_rocmfp4[16] (int8) plus rocmfp_ue4m3_fp32_lut[128] (types.glsl). + lut_size = 16 + 128u * (uint32_t)sizeof(float); + break; + case GGML_TYPE_Q3_0_ROCMFPX: + case GGML_TYPE_Q6_0_ROCMFPX: + case GGML_TYPE_Q8_0_ROCMFPX: + // rocmfp_ue4m3_fp32_lut[128] only; the fp3 codebook is a const array. + lut_size = 128u * (uint32_t)sizeof(float); + break; default: break; } // Needs to be kept up to date on shader changes // Needs to stay aligned with ggml_vk_mul_mm_spec. - const bool intel_shmem_stride_pad_zero = device->vendor_id == VK_VENDOR_ID_INTEL && device->coopmat_support && - device->driver_id == vk::DriverId::eIntelProprietaryWindows; - const uint32_t bank_conflict_offset = intel_shmem_stride_pad_zero ? 0 : (device->coopmat_support ? 8 : 1); + // The shared tiles are (BK/2 + pad) FLOAT_TYPEV2 per row, i.e. (BK + 2*pad) + // elements, so the byte-space offset below is twice the shader's pad. + const uint32_t bank_conflict_offset = device->coopmat_support ? 2 * ggml_vk_coopmat_shmem_pad(device, warptile[3]) : 1; const uint32_t type_size = device->fp16 ? sizeof(ggml_fp16_t) : sizeof(float); const uint32_t warps = warptile[0] / warptile[10]; @@ -4270,6 +4421,8 @@ static bool ggml_vk_matmul_int_shmem_support(const vk_device& device, const std: case GGML_TYPE_Q5_1: block_a_size = std430_size({{16, 4}, {4, 4}, {fp2_size, fp2_align}}); break; // qs[16/4] + qh + dm(vec2) case GGML_TYPE_Q8_0: block_a_size = std430_size({{32, 4}, {fp_size, fp_align}}); break; // qs[8] + dm case GGML_TYPE_MXFP4: block_a_size = std430_size({{32, 4}, {fp_size, fp_align}}); break; // qs[8] + d + case GGML_TYPE_Q4_0_ROCMFP4: block_a_size = std430_size({{32, 4}, {fp2_size, fp2_align}}); break; // qs[8] + d(vec2) + case GGML_TYPE_Q4_0_ROCMFP4_FAST: block_a_size = std430_size({{32, 4}, {fp_size, fp_align}}); break; // qs[8] + d case GGML_TYPE_Q2_K: block_a_size = std430_size({{ 8, 4}, {2, 2}, {fp2_size, fp2_align}}); break; // qs[2] + scales(u8vec2) + dm(vec2) case GGML_TYPE_Q3_K: block_a_size = std430_size({{16, 4}, {fp2_size, fp2_align}}); break; // qs[4] + d_scales(vec2) case GGML_TYPE_Q4_K: block_a_size = std430_size({{16, 4}, {fp2_size, fp2_align}}); break; // qs[4] + dm(vec2) @@ -4412,6 +4565,34 @@ struct CompileTask { uint32_t required_subgroup_size; }; +// EXPERIMENT (GGML_VK_MMID_F16B=1): convert the contiguous f32 activations (B) of +// quantized MUL_MAT_ID to f16 and run the f16-B matmul_id kernels instead of the +// f32-B ones. Halves B bytes and buf_b shared memory (better occupancy) at the +// f32->f16 rounding cost upstream already accepts on the coopmat2 path. +static bool ggml_vk_mmid_f16b_enabled() { + static const bool enabled = [] { + const char * env = getenv("GGML_VK_MMID_F16B"); + return env == nullptr || atoi(env) != 0; // on by default; =0 disables + }(); + return enabled; +} + +// GGML_VK_DENSE_F16B: same idea as GGML_VK_MMID_F16B but for plain MUL_MAT. Halves the B bytes +// moved. Numerically identical: mul_mm stages B into shared FLOAT_TYPE either way, so the f32-B +// kernel already rounds B to f16. Helps wide dense models, costs ~1% on narrow ones. +// 0 = off, 1 = all quantized dense matmuls, 2 = auto (only the K we have positive data for) +static int ggml_vk_dense_f16b_mode() { + static const int mode = [] { + const char * e = getenv("GGML_VK_DENSE_F16B"); + if (e == nullptr) return 0; + if (e[0] == 'a') return 2; + return atoi(e) != 0 ? 1 : 0; + }(); + return mode; +} + +static bool ggml_vk_dense_f16b_enabled() { return ggml_vk_dense_f16b_mode() != 0; } + static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { VK_LOG_DEBUG("ggml_vk_load_shaders(" << device->name << ")"); @@ -4435,10 +4616,6 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { (device->subgroup_size_control && device->subgroup_max_size >= 16); // mulmat - // Warptile layout (indices match mul_mm.comp constantIDs): - // [0..9] : BLOCK_SIZE, BM, BN, BK, WM, WN, WMITER, TM, TN, TK - // [10] : WARP / required_subgroup_size (read via WARP_SIZE_IDX) - static constexpr size_t WARP_SIZE_IDX = 10; std::vector l_warptile, m_warptile, s_warptile, l_warptile_id, m_warptile_id, s_warptile_id, l_warptile_mmq, m_warptile_mmq, s_warptile_mmq, @@ -4490,9 +4667,9 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { // spec constants and tile sizes for quant matmul_id const uint32_t mmqid_bk = device->coopmat2_decode_vector ? 64u : 32u; - l_warptile_mmqid = { 256, 128, 128, mmqid_bk, 1 }; - m_warptile_mmqid = { 256, 128, 64, mmqid_bk, 0 }; - s_warptile_mmqid = { 256, 128, 64, mmqid_bk, 0 }; + l_warptile_mmqid = { 256, 128, 128, mmqid_bk, 1, device->subgroup_size }; + m_warptile_mmqid = { 256, 128, 64, mmqid_bk, 0, device->subgroup_size }; + s_warptile_mmqid = { 256, 128, 64, mmqid_bk, 0, device->subgroup_size }; l_mmqid_wg_denoms = { 128, 128, 1 }; m_mmqid_wg_denoms = { 128, 64, 1 }; s_mmqid_wg_denoms = { 128, 64, 1 }; @@ -4557,6 +4734,27 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { l_warptile = { 256, 128, 128, 16, mm_warp_8, 64, 2, tm_m, tn_m, tk_m, mm_warp_8 }; l_warptile_mmq = l_warptile_mmq_int = { 256, 128, 128, 32, mm_warp_8, 64, 2, tm_m, tn_m, tk_m, mm_warp_8 }; l_warptile_mmq_int_k = { 256, 128, 128, 32, mm_warp_16, 64, 1, 4, 2, 1, mm_warp_16 }; + + // EXPERIMENT (GGML_VK_MMID_WG256=1): the dense large tile above runs 256 threads on a + // 128x128 tile, but the mul_mat_id variants still run 128. Give MoE the same thread + // count per tile: same BM/BN/BK (so same shared memory), twice the threads sharing each + // A/B tile load, half the accumulators per thread. Warp split stays legal: + // (BM/WM)*(BN/WN) == wg/subgroup == 4, WNITER == (WM*WN)/(WARP*TM*TN*WMITER) == 2. + // A 256-expert MoE at ub=2048 sees only ~64 rows per expert, so the tile that actually + // runs is the medium one, not the large one. Override both. + static const char * mmid_wg256_env = getenv("GGML_VK_MMID_WG256"); + if (mmid_wg256_env && atoi(mmid_wg256_env) != 0) { + l_warptile_mmqid = { 256, 128, 128, 32, mul_mat_subgroup_size_8, 64, 2, tm_m, tn_m, tk_m, mul_mat_subgroup_size_8 }; + l_warptile_mmqid_int = { 256, 128, 128, 32, mul_mat_subgroup_size_8, 64, 2, 4, 4, 1, mul_mat_subgroup_size_8 }; + // BM=BN=64 at 4 warps needs WM=WN=32: (BM/WM)*(BN/WN) == 4, cms_per_row/col == 2. + m_warptile_mmqid = { 256, 64, 64, 32, 32, 32, 2, tm_m, tn_m, tk_m, mul_mat_subgroup_size_8 }; + m_warptile_mmqid_int = { 256, 64, 64, 32, 32, 32, 2, 2, 2, 1, mul_mat_subgroup_size_8 }; + fprintf(stderr, "ggml_vulkan: MUL_MAT_ID medium+large tiles at 256 threads (GGML_VK_MMID_WG256)\n"); + } + } else if (device->vendor_id == VK_VENDOR_ID_INTEL && device->coopmat_support) { + // Xe2/Xe3 with coopmat enabled - warptile performance tuning + l_warptile = { 512, 128, 128, 16, mm_warp_8, 32, 2, tm_m, tn_m, tk_m, mm_warp_8 }; + l_warptile_mmq = { 512, 128, 128, 32, mm_warp_8, 32, 2, tm_m, tn_m, tk_m, mm_warp_8 }; } l_mmq_wg_denoms = l_wg_denoms = {128, 128, 1 }; @@ -4566,20 +4764,6 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { m_align = 64; s_align = 32; - if (device->vendor_id == VK_VENDOR_ID_INTEL && device->coopmat_support) { - // Xe1/Xe2/Xe3 with coopmat enabled - warptile performance tuning - l_warptile = { 512, 128, 128, 16, mm_warp_8, 32, 2, tm_l, tn_l, tk_l, mm_warp_8 }; - if (device->architecture == INTEL_XE1) { - l_warptile_mmq = { 512, 256, 128, 32, 32, 32, 2, tm_l, tn_l, tk_l, 16 }; - l_mmq_wg_denoms = { 256, 128, 1 }; - l_align = 32; //set as BK - } else { - l_warptile_mmq = { 512, 128, 256, 32, 32, 32, 2, tm_l, tn_l, tk_l, 16 }; - l_mmq_wg_denoms = { 128, 256, 1 }; - l_align = 32; //set as BK - } - } - for (uint32_t i = 0; i < GGML_TYPE_COUNT; ++i) { ggml_type t = (ggml_type)i; // Disable medium and large matrix multiplication if not enough shared memory is available @@ -4643,6 +4827,22 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { } } + if (!device->pipeline_matmul_f32) { + device->pipeline_matmul_f32 = std::make_shared(); + } + if (!device->pipeline_matmul_f32_f16) { + device->pipeline_matmul_f32_f16 = std::make_shared(); + } + if (!device->pipeline_matmul_id_f32) { + device->pipeline_matmul_id_f32 = std::make_shared(); + } + if (!device->pipeline_matmul_bf16) { + device->pipeline_matmul_bf16 = std::make_shared(); + } + if (!device->pipeline_matmul_id_bf16) { + device->pipeline_matmul_id_bf16 = std::make_shared(); + } + auto const &ggml_vk_create_pipeline = [&](vk_device& device, vk_pipeline& base_pipeline, const char *name, size_t spv_size, const void* spv_data, const char *entrypoint, uint32_t parameter_count, uint32_t push_constant_size, std::array wg_denoms, const std::vector& specialization_constants, uint32_t align, bool disable_robustness = false, bool require_full_subgroups = false, uint32_t required_subgroup_size = 0) { @@ -4843,648 +5043,904 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { auto const &ggml_vk_mul_mm_spec = [&device](std::vector spec, bool aligned) { spec.push_back(aligned ? 1u : 0u); // constantID=11: ALIGNED - if (device->vendor_id == VK_VENDOR_ID_INTEL && device->coopmat_support && - device->driver_id == vk::DriverId::eIntelProprietaryWindows) { - spec.push_back(0u); // constantID=12: SHMEM_STRIDE_PAD = 0 - spec.push_back(1u); // constantID=13: APPLY_SLM_A_RESHAPE = true - } - return spec; - }; - - auto const &ggml_vk_mul_mm_spec_quant = [&device](std::vector spec, bool aligned, uint32_t type) { - spec.push_back(aligned ? 1u : 0u); // constantID=11: ALIGNED - spec.push_back(type); // constantID=12: MmTypeA - if (device->vendor_id == VK_VENDOR_ID_INTEL && device->coopmat_support && - device->driver_id == vk::DriverId::eIntelProprietaryWindows) { - spec.push_back(0u); // constantID=13: SHMEM_STRIDE_PAD = 0 - spec.push_back(1u); // constantID=14: APPLY_SLM_A_RESHAPE = true + const uint32_t pad = ggml_vk_coopmat_shmem_pad(device, spec[3]); // spec[3] is the warptile BK + const bool intel_slm = device->vendor_id == VK_VENDOR_ID_INTEL && device->coopmat_support && + device->driver_id == vk::DriverId::eIntelProprietaryWindows; + if (intel_slm || pad != 4) { + spec.push_back(pad); // constantID=12: SHMEM_STRIDE_PAD + spec.push_back(intel_slm ? 1u : 0u); // constantID=13: APPLY_SLM_A_RESHAPE } return spec; }; - static const ggml_type non_lut_quant_types[] = { - GGML_TYPE_Q1_0, GGML_TYPE_Q2_0, GGML_TYPE_Q4_0, GGML_TYPE_Q4_1, GGML_TYPE_Q5_0, GGML_TYPE_Q5_1, GGML_TYPE_Q8_0, - GGML_TYPE_Q2_K, GGML_TYPE_Q3_K, GGML_TYPE_Q4_K, GGML_TYPE_Q5_K, GGML_TYPE_Q6_K, GGML_TYPE_TQ1_0, GGML_TYPE_TQ2_0, - }; - -#define FOR_EACH_LUT_TYPE_NONFP4(X) \ - X(GGML_TYPE_IQ1_S, iq1_s) \ - X(GGML_TYPE_IQ1_M, iq1_m) \ - X(GGML_TYPE_IQ2_XXS, iq2_xxs) \ - X(GGML_TYPE_IQ2_XS, iq2_xs) \ - X(GGML_TYPE_IQ2_S, iq2_s) \ - X(GGML_TYPE_IQ3_XXS, iq3_xxs) \ - X(GGML_TYPE_IQ3_S, iq3_s) \ - X(GGML_TYPE_IQ4_XS, iq4_xs) \ - X(GGML_TYPE_IQ4_NL, iq4_nl) -#define FOR_EACH_LUT_FP4_TYPE(X) \ - X(GGML_TYPE_MXFP4, mxfp4) \ - X(GGML_TYPE_NVFP4, nvfp4) -#define FOR_EACH_LUT_TYPE(X) \ - FOR_EACH_LUT_TYPE_NONFP4(X) \ - FOR_EACH_LUT_FP4_TYPE(X) - - const int mul_mat_id_param_count = 5; - - using spec_fn_t = std::function(const std::vector&, bool)>; - auto const &create_mm_pipelines = [&]( - const vk_matmul_pipeline_key& key, - const std::vector& tile_configs, - const std::string& shader_name, size_t spv_len, const void* spv_data, - uint32_t push_constant_size, uint32_t param_count, - const spec_fn_t& spec_fn, - bool disable_robustness = false, bool require_full_subgroups = false, uint32_t required_subgroup_size = 0, - bool create_aligned = true, bool pin_subgroup_to_warp = false - ) { - auto& vec = device->pipeline_matmul[key]; - const bool first_call = vec.empty(); - for (size_t i = 0; i < tile_configs.size(); i++) { - const auto& tc = tile_configs[i]; - - // Intel coopmat1 pins the required subgroup size to each warptile's WARP element. - const uint32_t rsgs = pin_subgroup_to_warp ? tc.warptile[WARP_SIZE_IDX] : required_subgroup_size; - const bool rfs = require_full_subgroups || pin_subgroup_to_warp; - - if (first_call) { - vk_matmul_pipeline_pair pair{}; - pair.align = tc.align; - std::string suffix = "_" + std::to_string(i); - pair.unaligned = std::make_shared(); - if (create_aligned) { - pair.aligned = std::make_shared(); - } - vec.push_back(pair); - } - - ggml_vk_create_pipeline(device, vec[i].unaligned, - vec[i].unaligned->name.empty() ? (shader_name + "_" + std::to_string(i)).c_str() : vec[i].unaligned->name.c_str(), - spv_len, spv_data, "main", param_count, push_constant_size, - tc.wg_denoms, spec_fn(tc.warptile, false), 1, - disable_robustness, rfs, rsgs); - - if (vec[i].aligned) { - ggml_vk_create_pipeline(device, vec[i].aligned, - vec[i].aligned->name.empty() ? (shader_name + "_aligned_" + std::to_string(i)).c_str() : vec[i].aligned->name.c_str(), - spv_len, spv_data, "main", param_count, push_constant_size, - tc.wg_denoms, spec_fn(tc.warptile, true), tc.align, - disable_robustness, rfs, rsgs); - } - } - }; - - auto filter_tc = [&](const std::vector& configs, ggml_type type, bool is_id, bool is_int = false) -> std::vector { - std::vector result; - bool enabled[3]; - if (is_int) { - enabled[0] = is_id ? device->mul_mat_id_s_int[type] : device->mul_mat_s_int[type]; - enabled[1] = is_id ? device->mul_mat_id_m_int[type] : device->mul_mat_m_int[type]; - enabled[2] = is_id ? device->mul_mat_id_l_int[type] : device->mul_mat_l_int[type]; - } else { - enabled[0] = is_id ? device->mul_mat_id_s[type] : device->mul_mat_s[type]; - enabled[1] = is_id ? device->mul_mat_id_m[type] : device->mul_mat_m[type]; - enabled[2] = is_id ? device->mul_mat_id_l[type] : device->mul_mat_l[type]; - } - for (size_t i = 0; i < configs.size() && i < 3; i++) { - if (enabled[i]) result.push_back(configs[i]); - } - return result; - }; - - std::vector tc_mm = {{s_warptile, s_wg_denoms, s_align}, {m_warptile, m_wg_denoms, m_align}, {l_warptile, l_wg_denoms, l_align}}; - std::vector tc_mmq = {{s_warptile_mmq, s_mmq_wg_denoms, s_align}, {m_warptile_mmq, m_mmq_wg_denoms, m_align}, {l_warptile_mmq, l_mmq_wg_denoms, l_align}}; + const int mul_mat_id_param_count = 6; // a, b, d, ids, expert_counts, fused scale #if defined(VK_NV_cooperative_matrix2) && defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT) if (device->coopmat2) { - auto const &ggml_vk_mul_mm_cm2_spec = [&](std::vector spec, bool aligned, uint32_t type = UINT32_MAX) { - spec.push_back(aligned ? 1u : 0u); // ALIGNED - spec.push_back(device->subgroup_size); // subgroup_size - if (type != UINT32_MAX) { - spec.push_back(type); // MmTypeA - spec.push_back((uint32_t)ggml_type_size((ggml_type)type)); // MmABlockBytes + auto const &ggml_vk_mul_mm_cm2_spec = [](std::vector spec, bool aligned, bool mul_mat_id) { + if (mul_mat_id && spec.size() > 5) { + spec.insert(spec.begin() + 5, aligned ? 1u : 0u); + } else { + spec.push_back(aligned ? 1u : 0u); + } + if (mul_mat_id && spec.size() == 6) { + spec.push_back(32); } return spec; }; - std::vector tc_mmq_k = {{s_warptile_mmq_k, s_mmq_wg_denoms_k, s_align}, {m_warptile_mmq_k, m_mmq_wg_denoms_k, m_align}, {l_warptile_mmq_k, l_mmq_wg_denoms_k, l_align}}; - std::vector tc_mmqid = {{s_warptile_mmqid, s_mmqid_wg_denoms, s_align}, {m_warptile_mmqid, m_mmqid_wg_denoms, m_align}, {l_warptile_mmqid, l_mmqid_wg_denoms, l_align}}; - - spec_fn_t cm2_spec = [&](const std::vector& wt, bool a) { return ggml_vk_mul_mm_cm2_spec(wt, a); }; - - // F16 x F16 - create_mm_pipelines({GGML_TYPE_F16, GGML_TYPE_F16, false, true}, tc_mm, "matmul_f16_f16acc", matmul_f16_f16acc_cm2_len, matmul_f16_f16acc_cm2_data, sizeof(vk_mat_mat_push_constants), 3, cm2_spec, true); - create_mm_pipelines({GGML_TYPE_F16, GGML_TYPE_F16, false, false}, tc_mm, "matmul_f16", matmul_f16_cm2_len, matmul_f16_cm2_data, sizeof(vk_mat_mat_push_constants), 3, cm2_spec, true); + // Create 6 variants, {s,m,l}x{unaligned,aligned} +#define CREATE_MM(PIPELINE_NAME, NAMELC, F16ACC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->l, #NAMELC #F16ACC "_l", NAMELC ## F16ACC ## _cm2_len, NAMELC ## F16ACC ## _cm2_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, ggml_vk_mul_mm_cm2_spec(l_ ## WARPTILE, false, PARAMCOUNT == mul_mat_id_param_count), 1, true); \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->m, #NAMELC #F16ACC "_m", NAMELC ## F16ACC ## _cm2_len, NAMELC ## F16ACC ## _cm2_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, ggml_vk_mul_mm_cm2_spec(m_ ## WARPTILE, false, PARAMCOUNT == mul_mat_id_param_count), 1, true); \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->s, #NAMELC #F16ACC "_s", NAMELC ## F16ACC ## _cm2_len, NAMELC ## F16ACC ## _cm2_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, ggml_vk_mul_mm_cm2_spec(s_ ## WARPTILE, false, PARAMCOUNT == mul_mat_id_param_count), 1, true); \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_l, #NAMELC #F16ACC "_aligned_l", NAMELC ## F16ACC ## _cm2_len, NAMELC ## F16ACC ## _cm2_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, ggml_vk_mul_mm_cm2_spec(l_ ## WARPTILE, true, PARAMCOUNT == mul_mat_id_param_count), l_align, true); \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_m, #NAMELC #F16ACC "_aligned_m", NAMELC ## F16ACC ## _cm2_len, NAMELC ## F16ACC ## _cm2_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, ggml_vk_mul_mm_cm2_spec(m_ ## WARPTILE, true, PARAMCOUNT == mul_mat_id_param_count), m_align, true); \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_s, #NAMELC #F16ACC "_aligned_s", NAMELC ## F16ACC ## _cm2_len, NAMELC ## F16ACC ## _cm2_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, ggml_vk_mul_mm_cm2_spec(s_ ## WARPTILE, true, PARAMCOUNT == mul_mat_id_param_count), s_align, true); \ + + // Create 2 variants, {f16,f32} accumulator +#define CREATE_MM2(PIPELINE_NAME, NAMELC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT) \ + CREATE_MM(PIPELINE_NAME . f16acc, NAMELC, _f16acc, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT) \ + CREATE_MM(PIPELINE_NAME . f32acc, NAMELC, , WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT) \ + + CREATE_MM2(pipeline_matmul_f16, matmul_f16, wg_denoms, warptile, vk_mat_mat_push_constants, 3) #if defined(GGML_VULKAN_BFLOAT16_GLSLC_SUPPORT) if (device->coopmat_bf16_support) { - create_mm_pipelines({GGML_TYPE_BF16, GGML_TYPE_BF16, false, false}, tc_mm, "matmul_bf16", matmul_bf16_cm2_len, matmul_bf16_cm2_data, sizeof(vk_mat_mat_push_constants), 3, cm2_spec, true); + CREATE_MM(pipeline_matmul_bf16, matmul_bf16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3) } #endif - for (const auto type : non_lut_quant_types) { - // regression in unified shader on Ampere - if (type == GGML_TYPE_Q4_K || type == GGML_TYPE_Q5_K) { - continue; - } - auto& tc = ((type >= GGML_TYPE_Q2_K && type <= GGML_TYPE_Q6_K) || type == GGML_TYPE_TQ1_0 || type == GGML_TYPE_TQ2_0) ? tc_mmq_k : tc_mmq; - spec_fn_t qs = [&, type](const std::vector& wt, bool a) { return ggml_vk_mul_mm_cm2_spec(wt, a, (uint32_t)type); }; - create_mm_pipelines({type, GGML_TYPE_F16, false, true}, tc, "matmul_quant_f16_f16acc", matmul_quant_f16_f16acc_cm2_len, matmul_quant_f16_f16acc_cm2_data, sizeof(vk_mat_mat_push_constants), 3, qs, true); - create_mm_pipelines({type, GGML_TYPE_F16, false, false}, tc, "matmul_quant_f16", matmul_quant_f16_cm2_len, matmul_quant_f16_cm2_data, sizeof(vk_mat_mat_push_constants), 3, qs, true); - } - create_mm_pipelines({GGML_TYPE_Q4_K, GGML_TYPE_F16, false, true}, tc_mmq_k, "matmul_q4_k_f16_f16acc", matmul_q4_k_f16_f16acc_cm2_len, matmul_q4_k_f16_f16acc_cm2_data, sizeof(vk_mat_mat_push_constants), 3, cm2_spec, true); - create_mm_pipelines({GGML_TYPE_Q4_K, GGML_TYPE_F16, false, false}, tc_mmq_k, "matmul_q4_k_f16", matmul_q4_k_f16_cm2_len, matmul_q4_k_f16_cm2_data, sizeof(vk_mat_mat_push_constants), 3, cm2_spec, true); - create_mm_pipelines({GGML_TYPE_Q5_K, GGML_TYPE_F16, false, true}, tc_mmq_k, "matmul_q5_k_f16_f16acc", matmul_q5_k_f16_f16acc_cm2_len, matmul_q5_k_f16_f16acc_cm2_data, sizeof(vk_mat_mat_push_constants), 3, cm2_spec, true); - create_mm_pipelines({GGML_TYPE_Q5_K, GGML_TYPE_F16, false, false}, tc_mmq_k, "matmul_q5_k_f16", matmul_q5_k_f16_cm2_len, matmul_q5_k_f16_cm2_data, sizeof(vk_mat_mat_push_constants), 3, cm2_spec, true); -#define X_CM2(TYPE, tstr) \ - { auto tc = filter_tc(tc_mmq, TYPE, false); \ - if (!tc.empty()) { \ - create_mm_pipelines({TYPE, GGML_TYPE_F16, false, true}, tc, "matmul_" #tstr "_f16_f16acc", matmul_##tstr##_f16_f16acc_cm2_len, matmul_##tstr##_f16_f16acc_cm2_data, sizeof(vk_mat_mat_push_constants), 3, cm2_spec, true); \ - create_mm_pipelines({TYPE, GGML_TYPE_F16, false, false}, tc, "matmul_" #tstr "_f16", matmul_##tstr##_f16_cm2_len, matmul_##tstr##_f16_cm2_data, sizeof(vk_mat_mat_push_constants), 3, cm2_spec, true); \ - } } - FOR_EACH_LUT_TYPE_NONFP4(X_CM2) + CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q1_0], matmul_q1_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) + CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q2_0], matmul_q2_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) + CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_0], matmul_q4_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) + CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_1], matmul_q4_1_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) + CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_0], matmul_q5_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) + CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_1], matmul_q5_1_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) + CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q8_0], matmul_q8_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) + CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q2_K], matmul_q2_k_f16, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) + CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_TQ2_0], matmul_tq2_0_f16, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) + CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_TQ1_0], matmul_tq1_0_f16, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) + CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q3_K], matmul_q3_k_f16, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) + CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_K], matmul_q4_k_f16, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) + CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_K], matmul_q5_k_f16, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) + CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q6_K], matmul_q6_k_f16, mmq_wg_denoms_k, warptile_mmq_k, vk_mat_mat_push_constants, 3) + CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ1_S], matmul_iq1_s_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) + CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ1_M], matmul_iq1_m_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) + CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ2_XXS], matmul_iq2_xxs_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) + CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ2_XS], matmul_iq2_xs_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) + CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ2_S], matmul_iq2_s_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) + CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ3_XXS], matmul_iq3_xxs_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) + CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ3_S], matmul_iq3_s_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) + CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ4_XS], matmul_iq4_xs_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) + CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_IQ4_NL], matmul_iq4_nl_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) + CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_0_ROCMFP4], matmul_rocmfp4_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) + CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_0_ROCMFP4_FAST], matmul_rocmfp4_fast_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) + CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q3_0_ROCMFPX], matmul_rocmfpx_fp3_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) + CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q6_0_ROCMFPX], matmul_rocmfpx_fp6_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) + CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q8_0_ROCMFPX], matmul_rocmfpx_fp8_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) #if defined(GGML_VULKAN_FLOAT_E2M1_GLSLC_SUPPORT) && defined(GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT) if (device->ocp_fp4) { -#define X_CM2_OCP(TYPE, tstr) \ - { auto tc = filter_tc(tc_mmq, TYPE, false); \ - if (!tc.empty()) { \ - create_mm_pipelines({TYPE, GGML_TYPE_F16, false, true}, tc, "matmul_" #tstr "_f16_ocp_f16acc", matmul_##tstr##_f16_ocp_f16acc_cm2_len, matmul_##tstr##_f16_ocp_f16acc_cm2_data, sizeof(vk_mat_mat_push_constants), 3, cm2_spec, true); \ - create_mm_pipelines({TYPE, GGML_TYPE_F16, false, false}, tc, "matmul_" #tstr "_f16_ocp", matmul_##tstr##_f16_ocp_cm2_len, matmul_##tstr##_f16_ocp_cm2_data, sizeof(vk_mat_mat_push_constants), 3, cm2_spec, true); \ - } } - FOR_EACH_LUT_FP4_TYPE(X_CM2_OCP) -#undef X_CM2_OCP + CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_MXFP4], matmul_mxfp4_f16_ocp, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) + CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_NVFP4], matmul_nvfp4_f16_ocp, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) } else #endif { - FOR_EACH_LUT_FP4_TYPE(X_CM2) + CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_MXFP4], matmul_mxfp4_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) + CREATE_MM2(pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_NVFP4], matmul_nvfp4_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3) } -#undef X_CM2 GGML_ASSERT(device->subgroup_ballot); - create_mm_pipelines({GGML_TYPE_F16, GGML_TYPE_F16, true, true}, tc_mm, "matmul_id_subgroup_f16_f16acc", matmul_id_subgroup_f16_f16acc_cm2_len, matmul_id_subgroup_f16_f16acc_cm2_data, sizeof(vk_mat_mat_id_push_constants), 5, cm2_spec, true); - create_mm_pipelines({GGML_TYPE_F16, GGML_TYPE_F16, true, false}, tc_mm, "matmul_id_subgroup_f16", matmul_id_subgroup_f16_cm2_len, matmul_id_subgroup_f16_cm2_data, sizeof(vk_mat_mat_id_push_constants), 5, cm2_spec, true); + CREATE_MM2(pipeline_matmul_id_f16, matmul_id_subgroup_f16, wg_denoms, warptile, vk_mat_mat_id_push_constants, mul_mat_id_param_count) #if defined(GGML_VULKAN_BFLOAT16_GLSLC_SUPPORT) if (device->coopmat_bf16_support) { - create_mm_pipelines({GGML_TYPE_BF16, GGML_TYPE_BF16, true, false}, tc_mm, "matmul_id_subgroup_bf16", matmul_id_subgroup_bf16_cm2_len, matmul_id_subgroup_bf16_cm2_data, sizeof(vk_mat_mat_id_push_constants), 5, cm2_spec, true); + CREATE_MM(pipeline_matmul_id_bf16, matmul_id_subgroup_bf16, , wg_denoms, warptile, vk_mat_mat_id_push_constants, mul_mat_id_param_count) } #endif - for (const auto type : non_lut_quant_types) { - if (type == GGML_TYPE_Q4_K || type == GGML_TYPE_Q5_K) { - continue; - } - spec_fn_t qs_id = [&, type](const std::vector& wt, bool a) { return ggml_vk_mul_mm_cm2_spec(wt, a, (uint32_t)type); }; - create_mm_pipelines({type, GGML_TYPE_F16, true, true}, tc_mmqid, "matmul_id_subgroup_quant_f16_f16acc", matmul_id_subgroup_quant_f16_f16acc_cm2_len, matmul_id_subgroup_quant_f16_f16acc_cm2_data, sizeof(vk_mat_mat_id_push_constants), 5, qs_id, true); - create_mm_pipelines({type, GGML_TYPE_F16, true, false}, tc_mmqid, "matmul_id_subgroup_quant_f16", matmul_id_subgroup_quant_f16_cm2_len, matmul_id_subgroup_quant_f16_cm2_data, sizeof(vk_mat_mat_id_push_constants), 5, qs_id, true); - } - create_mm_pipelines({GGML_TYPE_Q4_K, GGML_TYPE_F16, true, true}, tc_mmqid, "matmul_id_subgroup_q4_k_f16_f16acc", matmul_id_subgroup_q4_k_f16_f16acc_cm2_len, matmul_id_subgroup_q4_k_f16_f16acc_cm2_data, sizeof(vk_mat_mat_id_push_constants), 5, cm2_spec, true); - create_mm_pipelines({GGML_TYPE_Q4_K, GGML_TYPE_F16, true, false}, tc_mmqid, "matmul_id_subgroup_q4_k_f16", matmul_id_subgroup_q4_k_f16_cm2_len, matmul_id_subgroup_q4_k_f16_cm2_data, sizeof(vk_mat_mat_id_push_constants), 5, cm2_spec, true); - create_mm_pipelines({GGML_TYPE_Q5_K, GGML_TYPE_F16, true, true}, tc_mmqid, "matmul_id_subgroup_q5_k_f16_f16acc", matmul_id_subgroup_q5_k_f16_f16acc_cm2_len, matmul_id_subgroup_q5_k_f16_f16acc_cm2_data, sizeof(vk_mat_mat_id_push_constants), 5, cm2_spec, true); - create_mm_pipelines({GGML_TYPE_Q5_K, GGML_TYPE_F16, true, false}, tc_mmqid, "matmul_id_subgroup_q5_k_f16", matmul_id_subgroup_q5_k_f16_cm2_len, matmul_id_subgroup_q5_k_f16_cm2_data, sizeof(vk_mat_mat_id_push_constants), 5, cm2_spec, true); -#define X_CM2_ID(TYPE, tstr) \ - { auto tc = filter_tc(tc_mmqid, TYPE, true); \ - if (!tc.empty()) { \ - create_mm_pipelines({TYPE, GGML_TYPE_F16, true, true}, tc, "matmul_id_subgroup_" #tstr "_f16_f16acc", matmul_id_subgroup_##tstr##_f16_f16acc_cm2_len, matmul_id_subgroup_##tstr##_f16_f16acc_cm2_data, sizeof(vk_mat_mat_id_push_constants), 5, cm2_spec, true); \ - create_mm_pipelines({TYPE, GGML_TYPE_F16, true, false}, tc, "matmul_id_subgroup_" #tstr "_f16", matmul_id_subgroup_##tstr##_f16_cm2_len, matmul_id_subgroup_##tstr##_f16_cm2_data, sizeof(vk_mat_mat_id_push_constants), 5, cm2_spec, true); \ - } } - FOR_EACH_LUT_TYPE_NONFP4(X_CM2_ID) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q1_0], matmul_id_subgroup_q1_0_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_0], matmul_id_subgroup_q2_0_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0], matmul_id_subgroup_q4_0_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1], matmul_id_subgroup_q4_1_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0], matmul_id_subgroup_q5_0_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1], matmul_id_subgroup_q5_1_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0], matmul_id_subgroup_q8_0_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K], matmul_id_subgroup_q2_k_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ2_0], matmul_id_subgroup_tq2_0_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K], matmul_id_subgroup_q3_k_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K], matmul_id_subgroup_q4_k_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K], matmul_id_subgroup_q5_k_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q6_K], matmul_id_subgroup_q6_k_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ1_S], matmul_id_subgroup_iq1_s_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ1_M], matmul_id_subgroup_iq1_m_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ2_XXS], matmul_id_subgroup_iq2_xxs_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ2_XS], matmul_id_subgroup_iq2_xs_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ2_S], matmul_id_subgroup_iq2_s_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ3_XXS], matmul_id_subgroup_iq3_xxs_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ3_S], matmul_id_subgroup_iq3_s_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_XS], matmul_id_subgroup_iq4_xs_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_NL], matmul_id_subgroup_iq4_nl_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0_ROCMFP4], matmul_id_subgroup_rocmfp4_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0_ROCMFP4_FAST], matmul_id_subgroup_rocmfp4_fast_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_0_ROCMFPX], matmul_id_subgroup_rocmfpx_fp3_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q6_0_ROCMFPX], matmul_id_subgroup_rocmfpx_fp6_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0_ROCMFPX], matmul_id_subgroup_rocmfpx_fp8_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ1_0], matmul_id_subgroup_tq1_0_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) #if defined(GGML_VULKAN_FLOAT_E2M1_GLSLC_SUPPORT) && defined(GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT) if (device->ocp_fp4) { -#define X_CM2_ID_OCP(TYPE, tstr) \ - { auto tc = filter_tc(tc_mmqid, TYPE, true); \ - if (!tc.empty()) { \ - create_mm_pipelines({TYPE, GGML_TYPE_F16, true, true}, tc, "matmul_id_subgroup_" #tstr "_f16_ocp_f16acc", matmul_id_subgroup_##tstr##_f16_ocp_f16acc_cm2_len, matmul_id_subgroup_##tstr##_f16_ocp_f16acc_cm2_data, sizeof(vk_mat_mat_id_push_constants), 5, cm2_spec, true); \ - create_mm_pipelines({TYPE, GGML_TYPE_F16, true, false}, tc, "matmul_id_subgroup_" #tstr "_f16_ocp", matmul_id_subgroup_##tstr##_f16_ocp_cm2_len, matmul_id_subgroup_##tstr##_f16_ocp_cm2_data, sizeof(vk_mat_mat_id_push_constants), 5, cm2_spec, true); \ - } } - FOR_EACH_LUT_FP4_TYPE(X_CM2_ID_OCP) -#undef X_CM2_ID_OCP + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_MXFP4], matmul_id_subgroup_mxfp4_f16_ocp, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_NVFP4], matmul_id_subgroup_nvfp4_f16_ocp, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) } else #endif { - FOR_EACH_LUT_FP4_TYPE(X_CM2_ID) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_MXFP4], matmul_id_subgroup_mxfp4_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) + CREATE_MM2(pipeline_dequant_mul_mat_mat_id[GGML_TYPE_NVFP4], matmul_id_subgroup_nvfp4_f16, mmqid_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count) } -#undef X_CM2_ID +#undef CREATE_MM +#undef CREATE_MM2 } else #endif // defined(VK_NV_cooperative_matrix2) && defined(GGML_VULKAN_COOPMAT2_GLSLC_SUPPORT) #if defined(VK_KHR_cooperative_matrix) && defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT) if (device->coopmat_support) { - spec_fn_t cm1_spec = [&](const std::vector& wt, bool a) { return ggml_vk_mul_mm_spec(wt, a); }; - - // Intel coopmat1 pins each pipeline's required subgroup size to its warptile WARP element. - const bool cm1_pin = device->vendor_id == VK_VENDOR_ID_INTEL; - - // Intel coopmat1 uses a dedicated large-tile config for quant matmul_id. - std::vector tc_mmq_id = tc_mmq; - if (cm1_pin) { - tc_mmq_id[2] = { { 512, 128, 128, 32, 32, 32, 2, device->coopmat_m, device->coopmat_n, device->coopmat_k, 32 }, { 128, 128, 1 }, 32 }; + // Deterministic subgroup sizing for the dense coopmat pipelines. Two parts: + // + // (1) The required subgroup size is the tile's own WARP, not the driver's choice. The cm1 + // shaders derive their warp grid from the real subgroup (warp_i = gl_SubgroupID) and + // size shared arrays with NUM_WARPS = BLOCK_SIZE / WARP, so WARP and the actual + // subgroup must agree - leaving that to the driver makes the agreement incidental. + // + // (2) The QUANTISED dense tiles run at wave32. RDNA3.x WMMA is wave32-native, so a wave64 + // subgroup issues each coopmat op as two halves. BLOCK_SIZE is kept, so the subgroup + // count doubles and WM (or WN) halves until the warp grid tiles BM x BN again: + // NUM_WARPS == (BM/WM)*(BN/WN). A tile that cannot be retiled is left at wave64. + // + // The FLOAT tiles are deliberately excluded. Measured on gfx1151 with a standalone + // MUL_MAT microbench at both dense FFN shapes (m=25600 k=5120, m=5120 k=25600): + // q6_K +5.2..+10.8%, q8_0 +5.4..+8.4%, q4_K +0.7..+9.1%, q4_0 -1.5..+1.8%, but + // f16 -6.7..+6.4% and bf16 ~0. The win tracks inline dequant instruction count + // (q6_K 3907 -> 3433 instructions at wave32, identical 192 VGPRs and 8 subgroups/SIMD), + // so it lands on the issue-bound quantised kernels and not on the float ones, which are + // bandwidth-bound on the weight stream. + // + // Scoped to AMD coopmat1 on a wave64 default; other vendors keep the driver default, + // since none of the above is validated there. + // GGML_VK_DENSE_WAVE32=0 disables, =2 additionally retiles the float tiles (probe). + const bool dense_sgs_scope = + device->vendor_id == VK_VENDOR_ID_AMD && + device->driver_id != vk::DriverId::eAmdProprietary && + device->subgroup_size_control; + const bool dense_wave32_possible = + dense_sgs_scope && + device->subgroup_min_size <= 32 && 32 <= device->subgroup_max_size && + device->subgroup_size > 32; + const char * dense_wave32_env = getenv("GGML_VK_DENSE_WAVE32"); + const int dense_wave32 = dense_wave32_env ? atoi(dense_wave32_env) : 1; + + // The mul_mat_id tiles below start from these. Snapshot them before the dense + // wave32 rewrite so an mmid pipeline created without a required subgroup size + // (GGML_VK_MMID_WAVE32=0) keeps a WARP that matches the real wave64 subgroup. + const auto l_warptile_mmq_w64 = l_warptile_mmq; + const auto m_warptile_mmq_w64 = m_warptile_mmq; + const auto s_warptile_mmq_w64 = s_warptile_mmq; + if (dense_wave32_possible && dense_wave32 != 0) { + auto wave32_tile = [](std::vector & w) -> bool { + // {BLOCK_SIZE, BM, BN, BK, WM, WN, WMITER, TM, TN, TK, WARP} + std::vector t = w; + t[10] = 32; + for (int guard = 0; guard < 4 && t[0] / t[10] != (t[1] / t[4]) * (t[2] / t[5]); ++guard) { + if (t[4] >= t[5] && t[4] > t[7]) { + t[4] /= 2; // halve WM, keeping WM >= TM + } else { + t[5] /= 2; // halve WN + } + } + if (t[0] / t[10] != (t[1] / t[4]) * (t[2] / t[5]) || t[4] < t[7] || t[5] < t[8]) { + return false; + } + w = t; + return true; + }; + wave32_tile(l_warptile_mmq); + wave32_tile(m_warptile_mmq); + wave32_tile(s_warptile_mmq); + if (dense_wave32 >= 2) { + wave32_tile(l_warptile); + wave32_tile(m_warptile); + wave32_tile(s_warptile); + } } - auto cm1_create = [&](vk_matmul_pipeline_key key, const std::vector& tc_base, - const std::string& name, size_t len, const void* data, uint32_t pc_size, uint32_t pc) { - auto tc = filter_tc(tc_base, key.type_a, key.mul_mat_id); - if (!tc.empty()) create_mm_pipelines(key, tc, name, len, data, pc_size, pc, cm1_spec, false, true, 0, true, cm1_pin); - }; - auto cm1_create_quant = [&](vk_matmul_pipeline_key key, const std::vector& tc_base, - const std::string& name, size_t len, const void* data, uint32_t pc_size, uint32_t pc) { - spec_fn_t qs = [&, type_a=key.type_a](const std::vector& wt, bool a) { return ggml_vk_mul_mm_spec_quant(wt, a, (uint32_t)type_a); }; - auto tc = filter_tc(tc_base, key.type_a, key.mul_mat_id); - if (!tc.empty()) create_mm_pipelines(key, tc, name, len, data, pc_size, pc, qs, false, true, 0, true, cm1_pin); + // WARP -> required subgroup size, or 0 where the device cannot honor one. + auto dense_req_sgs = [dense_sgs_scope, &device](const std::vector & w) -> uint32_t { + const uint32_t warp = w[10]; + if (!dense_sgs_scope || warp < device->subgroup_min_size || warp > device->subgroup_max_size) { + return 0; + } + return warp; }; - cm1_create({GGML_TYPE_F32, GGML_TYPE_F32, false, false}, tc_mm, "matmul_f32_f32", matmul_f32_f32_cm1_len, matmul_f32_f32_cm1_data, sizeof(vk_mat_mat_push_constants), 3); - cm1_create({GGML_TYPE_F32, GGML_TYPE_F16, false, false}, tc_mm, "matmul_f32_f16", matmul_f32_f16_cm1_len, matmul_f32_f16_cm1_data, sizeof(vk_mat_mat_push_constants), 3); - if (device->coopmat_acc_f16_support) { - cm1_create({GGML_TYPE_F16, GGML_TYPE_F16, false, true}, tc_mm, "matmul_f16_f16acc", matmul_f16_f16acc_cm1_len, matmul_f16_f16acc_cm1_data, sizeof(vk_mat_mat_push_constants), 3); - cm1_create({GGML_TYPE_F16, GGML_TYPE_F32, false, true}, tc_mm, "matmul_f16_f32_f16acc", matmul_f16_f32_f16acc_cm1_len, matmul_f16_f32_f16acc_cm1_data, sizeof(vk_mat_mat_push_constants), 3); - } - if (device->coopmat_acc_f32_support) { - cm1_create({GGML_TYPE_F16, GGML_TYPE_F16, false, false}, tc_mm, "matmul_f16", matmul_f16_cm1_len, matmul_f16_cm1_data, sizeof(vk_mat_mat_push_constants), 3); - cm1_create({GGML_TYPE_F16, GGML_TYPE_F32, false, false}, tc_mm, "matmul_f16_f32", matmul_f16_f32_cm1_len, matmul_f16_f32_cm1_data, sizeof(vk_mat_mat_push_constants), 3); - } -#if defined(GGML_VULKAN_BFLOAT16_GLSLC_SUPPORT) - if (device->coopmat_bf16_support) { - cm1_create({GGML_TYPE_BF16, GGML_TYPE_BF16, false, false}, tc_mm, "matmul_bf16", matmul_bf16_cm1_len, matmul_bf16_cm1_data, sizeof(vk_mat_mat_push_constants), 3); - } -#endif - for (const auto type : non_lut_quant_types) { - if (device->coopmat_acc_f16_support) { - cm1_create_quant({type, GGML_TYPE_F32, false, true}, tc_mmq, "matmul_quant_f32_f16acc", matmul_quant_f32_f16acc_cm1_len, matmul_quant_f32_f16acc_cm1_data, sizeof(vk_mat_mat_push_constants), 3); - cm1_create_quant({type, GGML_TYPE_F16, false, true}, tc_mmq, "matmul_quant_f16_f16acc", matmul_quant_f16_f16acc_cm1_len, matmul_quant_f16_f16acc_cm1_data, sizeof(vk_mat_mat_push_constants), 3); - } - if (device->coopmat_acc_f32_support) { - cm1_create_quant({type, GGML_TYPE_F32, false, false}, tc_mmq, "matmul_quant_f32", matmul_quant_f32_cm1_len, matmul_quant_f32_cm1_data, sizeof(vk_mat_mat_push_constants), 3); - cm1_create_quant({type, GGML_TYPE_F16, false, false}, tc_mmq, "matmul_quant_f16", matmul_quant_f16_cm1_len, matmul_quant_f16_cm1_data, sizeof(vk_mat_mat_push_constants), 3); - } - } - // The _f16 variants provide the f16 B-type pipeline used when y_non_contig converts f32->f16. -#define X_CM1(TYPE, tstr) \ + // Create 6 variants, {s,m,l}x{unaligned,aligned} +#define CREATE_MM(TYPE, PIPELINE_NAME, NAMELC, F16ACC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ + if (device->mul_mat ## ID ## _l[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->l, #NAMELC #F16ACC "_l", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, ggml_vk_mul_mm_spec(l_ ## WARPTILE, false), 1, false, true, dense_req_sgs(l_ ## WARPTILE)); \ + if (device->mul_mat ## ID ## _m[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->m, #NAMELC #F16ACC "_m", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, ggml_vk_mul_mm_spec(m_ ## WARPTILE, false), 1, false, true, dense_req_sgs(m_ ## WARPTILE)); \ + if (device->mul_mat ## ID ## _s[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->s, #NAMELC #F16ACC "_s", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, ggml_vk_mul_mm_spec(s_ ## WARPTILE, false), 1, false, true, dense_req_sgs(s_ ## WARPTILE)); \ + if (device->mul_mat ## ID ## _l[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_l, #NAMELC #F16ACC "_aligned_l", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, ggml_vk_mul_mm_spec(l_ ## WARPTILE, true), l_align, false, true, dense_req_sgs(l_ ## WARPTILE)); \ + if (device->mul_mat ## ID ## _m[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_m, #NAMELC #F16ACC "_aligned_m", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, ggml_vk_mul_mm_spec(m_ ## WARPTILE, true), m_align, false, true, dense_req_sgs(m_ ## WARPTILE)); \ + if (device->mul_mat ## ID ## _s[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_s, #NAMELC #F16ACC "_aligned_s", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, ggml_vk_mul_mm_spec(s_ ## WARPTILE, true), s_align, false, true, dense_req_sgs(s_ ## WARPTILE)); \ + + // Create 2 variants, {f16,f32} accumulator +#define CREATE_MM2(TYPE, PIPELINE_NAME, NAMELC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ if (device->coopmat_acc_f16_support) { \ - cm1_create({TYPE, GGML_TYPE_F32, false, true}, tc_mmq, "matmul_" #tstr "_f32_f16acc", matmul_##tstr##_f32_f16acc_cm1_len, matmul_##tstr##_f32_f16acc_cm1_data, sizeof(vk_mat_mat_push_constants), 3); \ - cm1_create({TYPE, GGML_TYPE_F16, false, true}, tc_mmq, "matmul_" #tstr "_f16_f16acc", matmul_##tstr##_f16_f16acc_cm1_len, matmul_##tstr##_f16_f16acc_cm1_data, sizeof(vk_mat_mat_push_constants), 3); \ + CREATE_MM(TYPE, PIPELINE_NAME . f16acc, NAMELC, _f16acc, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ } \ if (device->coopmat_acc_f32_support) { \ - cm1_create({TYPE, GGML_TYPE_F32, false, false}, tc_mmq, "matmul_" #tstr "_f32", matmul_##tstr##_f32_cm1_len, matmul_##tstr##_f32_cm1_data, sizeof(vk_mat_mat_push_constants), 3); \ - cm1_create({TYPE, GGML_TYPE_F16, false, false}, tc_mmq, "matmul_" #tstr "_f16", matmul_##tstr##_f16_cm1_len, matmul_##tstr##_f16_cm1_data, sizeof(vk_mat_mat_push_constants), 3); \ + CREATE_MM(TYPE, PIPELINE_NAME . f32acc, NAMELC, , WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ + } \ + + CREATE_MM(GGML_TYPE_F32, pipeline_matmul_f32, matmul_f32_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); + CREATE_MM(GGML_TYPE_F32, pipeline_matmul_f32_f16, matmul_f32_f16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_F16, pipeline_matmul_f16, matmul_f16, wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_F16, pipeline_matmul_f16_f32, matmul_f16_f32, wg_denoms, warptile, vk_mat_mat_push_constants, 3, ); +#if defined(GGML_VULKAN_BFLOAT16_GLSLC_SUPPORT) + if (device->coopmat_bf16_support) { + CREATE_MM(GGML_TYPE_BF16, pipeline_matmul_bf16, matmul_bf16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, ) } - FOR_EACH_LUT_TYPE_NONFP4(X_CM1) +#endif + + CREATE_MM2(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q1_0], matmul_q1_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_0], matmul_q2_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_0], matmul_q4_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_1], matmul_q4_1_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_0], matmul_q5_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_1], matmul_q5_1_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q8_0], matmul_q8_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + + CREATE_MM2(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_K], matmul_q2_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_TQ2_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_TQ2_0], matmul_tq2_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_TQ1_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_TQ1_0], matmul_tq1_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_K], matmul_q3_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_K], matmul_q4_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_K], matmul_q5_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q6_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q6_K], matmul_q6_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + + // f16-B variants of the same quant pipelines. The _f16 SPIR-V is already built for + // every type; upstream only instantiates it in the coopmat2 branch. + if (ggml_vk_dense_f16b_enabled()) { + CREATE_MM2(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_0], matmul_q4_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_1], matmul_q4_1_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_0], matmul_q5_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_1], matmul_q5_1_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q8_0], matmul_q8_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q2_K], matmul_q2_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q3_K], matmul_q3_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q4_K], matmul_q4_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q5_K], matmul_q5_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q6_K, pipeline_dequant_mul_mat_mat_f16[GGML_TYPE_Q6_K], matmul_q6_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + } + CREATE_MM2(GGML_TYPE_IQ1_S, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ1_S], matmul_iq1_s_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_IQ1_M, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ1_M], matmul_iq1_m_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_IQ2_XXS, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ2_XXS], matmul_iq2_xxs_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_IQ2_XS, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ2_XS], matmul_iq2_xs_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_IQ2_S, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ2_S], matmul_iq2_s_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_IQ3_XXS, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ3_XXS], matmul_iq3_xxs_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_IQ3_S, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ3_S], matmul_iq3_s_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_IQ4_XS, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ4_XS], matmul_iq4_xs_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_IQ4_NL, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ4_NL], matmul_iq4_nl_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q4_0_ROCMFP4, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_0_ROCMFP4], matmul_rocmfp4_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q4_0_ROCMFP4_FAST, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_0_ROCMFP4_FAST], matmul_rocmfp4_fast_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q3_0_ROCMFPX, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_0_ROCMFPX], matmul_rocmfpx_fp3_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q6_0_ROCMFPX, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q6_0_ROCMFPX], matmul_rocmfpx_fp6_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_Q8_0_ROCMFPX, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q8_0_ROCMFPX], matmul_rocmfpx_fp8_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + #if defined(GGML_VULKAN_FLOAT_E2M1_GLSLC_SUPPORT) && defined(GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT) if (device->ocp_fp4) { -#define X_CM1_OCP(TYPE, tstr) \ - if (device->coopmat_acc_f16_support) { \ - cm1_create({TYPE, GGML_TYPE_F32, false, true}, tc_mmq, "matmul_" #tstr "_f32_ocp_f16acc", matmul_##tstr##_f32_ocp_f16acc_cm1_len, matmul_##tstr##_f32_ocp_f16acc_cm1_data, sizeof(vk_mat_mat_push_constants), 3); \ - cm1_create({TYPE, GGML_TYPE_F16, false, true}, tc_mmq, "matmul_" #tstr "_f16_ocp_f16acc", matmul_##tstr##_f16_ocp_f16acc_cm1_len, matmul_##tstr##_f16_ocp_f16acc_cm1_data, sizeof(vk_mat_mat_push_constants), 3); \ - } \ - if (device->coopmat_acc_f32_support) { \ - cm1_create({TYPE, GGML_TYPE_F32, false, false}, tc_mmq, "matmul_" #tstr "_f32_ocp", matmul_##tstr##_f32_ocp_cm1_len, matmul_##tstr##_f32_ocp_cm1_data, sizeof(vk_mat_mat_push_constants), 3); \ - cm1_create({TYPE, GGML_TYPE_F16, false, false}, tc_mmq, "matmul_" #tstr "_f16_ocp", matmul_##tstr##_f16_ocp_cm1_len, matmul_##tstr##_f16_ocp_cm1_data, sizeof(vk_mat_mat_push_constants), 3); \ - } - FOR_EACH_LUT_FP4_TYPE(X_CM1_OCP) -#undef X_CM1_OCP + CREATE_MM2(GGML_TYPE_MXFP4, pipeline_dequant_mul_mat_mat[GGML_TYPE_MXFP4], matmul_mxfp4_f32_ocp, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_NVFP4, pipeline_dequant_mul_mat_mat[GGML_TYPE_NVFP4], matmul_nvfp4_f32_ocp, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); } else #endif { - FOR_EACH_LUT_FP4_TYPE(X_CM1) + CREATE_MM2(GGML_TYPE_MXFP4, pipeline_dequant_mul_mat_mat[GGML_TYPE_MXFP4], matmul_mxfp4_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); + CREATE_MM2(GGML_TYPE_NVFP4, pipeline_dequant_mul_mat_mat[GGML_TYPE_NVFP4], matmul_nvfp4_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, ); } -#undef X_CM1 GGML_ASSERT(device->subgroup_ballot); - cm1_create({GGML_TYPE_F32, GGML_TYPE_F32, true, false}, tc_mm, "matmul_id_subgroup_f32_f32", matmul_id_subgroup_f32_f32_cm1_len, matmul_id_subgroup_f32_f32_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - if (device->coopmat_acc_f16_support) { - cm1_create({GGML_TYPE_F16, GGML_TYPE_F16, true, true}, tc_mm, "matmul_id_subgroup_f16_f16acc", matmul_id_subgroup_f16_f16acc_cm1_len, matmul_id_subgroup_f16_f16acc_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - cm1_create({GGML_TYPE_F16, GGML_TYPE_F32, true, true}, tc_mm, "matmul_id_subgroup_f16_f32_f16acc", matmul_id_subgroup_f16_f32_f16acc_cm1_len, matmul_id_subgroup_f16_f32_f16acc_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - } - if (device->coopmat_acc_f32_support) { - cm1_create({GGML_TYPE_F16, GGML_TYPE_F16, true, false}, tc_mm, "matmul_id_subgroup_f16", matmul_id_subgroup_f16_cm1_len, matmul_id_subgroup_f16_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - cm1_create({GGML_TYPE_F16, GGML_TYPE_F32, true, false}, tc_mm, "matmul_id_subgroup_f16_f32", matmul_id_subgroup_f16_f32_cm1_len, matmul_id_subgroup_f16_f32_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - } + CREATE_MM(GGML_TYPE_F32, pipeline_matmul_id_f32, matmul_id_subgroup_f32_f32, , wg_denoms, warptile, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_F16, pipeline_matmul_id_f16, matmul_id_subgroup_f16, wg_denoms, warptile, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_F16, pipeline_matmul_id_f16_f32, matmul_id_subgroup_f16_f32, wg_denoms, warptile, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); #if defined(GGML_VULKAN_BFLOAT16_GLSLC_SUPPORT) if (device->coopmat_bf16_support) { - cm1_create({GGML_TYPE_BF16, GGML_TYPE_BF16, true, false}, tc_mm, "matmul_id_subgroup_bf16", matmul_id_subgroup_bf16_cm1_len, matmul_id_subgroup_bf16_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); + CREATE_MM(GGML_TYPE_BF16, pipeline_matmul_id_bf16, matmul_id_subgroup_bf16, , wg_denoms, warptile, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); } #endif - for (const auto type : non_lut_quant_types) { - if (device->coopmat_acc_f16_support) { - cm1_create_quant({type, GGML_TYPE_F32, true, true}, tc_mmq_id, "matmul_id_subgroup_quant_f32_f16acc", matmul_id_subgroup_quant_f32_f16acc_cm1_len, matmul_id_subgroup_quant_f32_f16acc_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - cm1_create_quant({type, GGML_TYPE_F16, true, true}, tc_mmq_id, "matmul_id_subgroup_quant_f16_f16acc", matmul_id_subgroup_quant_f16_f16acc_cm1_len, matmul_id_subgroup_quant_f16_f16acc_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); + + // EXPERIMENT (GGML_VK_MMID_TILE16=1): narrow the matmul_id small tile to BN=16. + // MoE prefill leaves ~nei1*nei0/n_expert rows per expert (~16 at ub512/256E/8a), + // so even the 32-wide small tile runs half empty. Only meaningful stacked on + // GGML_VK_MMID_SMALLN=1 (routes mmid to the small tile) + the row-list prepass. + // Shadows the s-tile config for the mmid quant pipelines only; dense unaffected. + auto s_warptile_mmq_id16 = s_warptile_mmq_w64; + auto s_mmq_wg_denoms_id16 = s_mmq_wg_denoms; + auto m_warptile_mmq_id128 = m_warptile_mmq_w64; + auto m_mmq_wg_denoms_id128 = m_mmq_wg_denoms; + auto l_warptile_mmq_idw = l_warptile_mmq_w64; + uint32_t mmid_req_sgs = 0; + { + const char * tile16_env = getenv("GGML_VK_MMID_TILE16"); + if (tile16_env && atoi(tile16_env) != 0) { + s_warptile_mmq_id16[2] = 16; // BN + s_warptile_mmq_id16[5] = 16; // WN + s_mmq_wg_denoms_id16[1] = 16; } - if (device->coopmat_acc_f32_support) { - cm1_create_quant({type, GGML_TYPE_F32, true, false}, tc_mmq_id, "matmul_id_subgroup_quant_f32", matmul_id_subgroup_quant_f32_cm1_len, matmul_id_subgroup_quant_f32_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - cm1_create_quant({type, GGML_TYPE_F16, true, false}, tc_mmq_id, "matmul_id_subgroup_quant_f16", matmul_id_subgroup_quant_f16_cm1_len, matmul_id_subgroup_quant_f16_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); + // GGML_VK_MMID_BM64=1: taller small tile (BM 32->64, two warps). Same + // A-traffic (ic-tile count unchanged), halves per-expert B re-reads + // (ir-tile count), larger WGs to hide latency. + const char * bm64_env = getenv("GGML_VK_MMID_BM64"); + if (!bm64_env || atoi(bm64_env) != 0) { // on by default; =0 disables + s_warptile_mmq_id16[0] = 2 * mul_mat_subgroup_size; // BLOCK_SIZE + s_warptile_mmq_id16[1] = 64; // BM + s_mmq_wg_denoms_id16[0] = 64; + } + // GGML_VK_MMID_M128=1: same idea for the medium tile (BM 64->128, four + // warps) — the tile the per-expert-n heuristic picks at n~64 (e.g. ub2048). + const char * m128_env = getenv("GGML_VK_MMID_M128"); + if (!m128_env || atoi(m128_env) != 0) { // on by default; =0 disables + m_warptile_mmq_id128[0] = 4 * mul_mat_subgroup_size; // BLOCK_SIZE + m_warptile_mmq_id128[1] = 128; // BM + m_mmq_wg_denoms_id128[0] = 128; + } + // GGML_VK_MMID_WAVE32=1: force required subgroup size 32 on the mmid + // quant coopmat pipelines (RDNA3.x WMMA is wave32-native; probe whether + // RADV lowers KHR_coopmat better at wave32). The cm1 path derives its + // warp grid from the real subgroup (warp_i = gl_SubgroupID, tiw = + // gl_SubgroupInvocationID) and sizes shared arrays (coopmat_stage, + // ballots_sh) with NUM_WARPS = BLOCK_SIZE / WARP, so the WARP spec + // constant must equal the forced size, and the coverage invariant + // NUM_WARPS == (BM/WM)*(BN/WN) must be restored. BLOCK_SIZE is kept + // (same workgroup shape, load loops and shmem as the wave64 stack), so + // the subgroup count doubles and WM (or WN) is halved until the warp + // grid exactly tiles BM x BN again. Per-lane coopmat accumulator + // footprint is unchanged: half the lanes per subgroup, half the + // (WM/TM)*(WN/TN) fragments per subgroup. Runs after the BM64/M128 + // gates so it composes with the probe stack. Applies only when the + // driver honors a required size (subgroup_size_control covering 32); + // otherwise WARP=32 with a real subgroup of 64 would corrupt tiling. + const char * wave32_env = getenv("GGML_VK_MMID_WAVE32"); + if ((!wave32_env || atoi(wave32_env) != 0) && device->subgroup_size_control && // on by default; =0 disables + device->subgroup_min_size <= 32 && 32 <= device->subgroup_max_size) { + mmid_req_sgs = 32; + auto wave32_tile = [](std::vector &w) { + // {BLOCK_SIZE, BM, BN, BK, WM, WN, WMITER, TM, TN, TK, WARP} + w[10] = 32; // WARP: must match the forced subgroup size + for (int guard = 0; guard < 4 && w[0] / w[10] != (w[1] / w[4]) * (w[2] / w[5]); ++guard) { + if (w[4] >= w[5] && w[4] > w[7]) { + w[4] /= 2; // halve WM, keeping WM >= TM + } else { + w[5] /= 2; // halve WN + } + } + GGML_ASSERT(w[0] / w[10] == (w[1] / w[4]) * (w[2] / w[5])); // NUM_WARPS == (BM/WM)*(BN/WN) + GGML_ASSERT(w[4] >= w[7] && w[5] >= w[8]); // WM >= TM, WN >= TN + }; + wave32_tile(s_warptile_mmq_id16); + wave32_tile(m_warptile_mmq_id128); + wave32_tile(l_warptile_mmq_idw); + } else { + // Wave64 stack: the tiles are the dense wave64 tiles plus the BM64/M128 reshapes, + // which keep NUM_WARPS == (BM/WM)*(BN/WN). Pin the subgroup to the tile's WARP + // where the driver honours it, so the spec constant can never disagree with + // the real subgroup size. + for (const auto * w : { &s_warptile_mmq_id16, &m_warptile_mmq_id128, &l_warptile_mmq_idw }) { + GGML_ASSERT((*w)[0] / (*w)[10] == ((*w)[1] / (*w)[4]) * ((*w)[2] / (*w)[5])); + } + const uint32_t warp = s_warptile_mmq_id16[10]; + if (device->subgroup_size_control && warp == m_warptile_mmq_id128[10] && warp == l_warptile_mmq_idw[10] && + device->subgroup_min_size <= warp && warp <= device->subgroup_max_size) { + mmid_req_sgs = warp; + } } } - // The _f16 variants provide the f16 B-type pipeline used when y_non_contig converts f32->f16. -#define X_CM1_ID(TYPE, tstr) \ - if (device->coopmat_acc_f16_support) { \ - cm1_create({TYPE, GGML_TYPE_F32, true, true}, tc_mmq_id, "matmul_id_subgroup_" #tstr "_f32_f16acc", matmul_id_subgroup_##tstr##_f32_f16acc_cm1_len, matmul_id_subgroup_##tstr##_f32_f16acc_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); \ - cm1_create({TYPE, GGML_TYPE_F16, true, true}, tc_mmq_id, "matmul_id_subgroup_" #tstr "_f16_f16acc", matmul_id_subgroup_##tstr##_f16_f16acc_cm1_len, matmul_id_subgroup_##tstr##_f16_f16acc_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); \ - } \ - if (device->coopmat_acc_f32_support) { \ - cm1_create({TYPE, GGML_TYPE_F32, true, false}, tc_mmq_id, "matmul_id_subgroup_" #tstr "_f32", matmul_id_subgroup_##tstr##_f32_cm1_len, matmul_id_subgroup_##tstr##_f32_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); \ - cm1_create({TYPE, GGML_TYPE_F16, true, false}, tc_mmq_id, "matmul_id_subgroup_" #tstr "_f16", matmul_id_subgroup_##tstr##_f16_cm1_len, matmul_id_subgroup_##tstr##_f16_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); \ - } - FOR_EACH_LUT_TYPE_NONFP4(X_CM1_ID) + { + const auto &s_warptile_mmq = s_warptile_mmq_id16; + const auto &s_mmq_wg_denoms = s_mmq_wg_denoms_id16; + const auto &m_warptile_mmq = m_warptile_mmq_id128; + const auto &m_mmq_wg_denoms = m_mmq_wg_denoms_id128; + const auto &l_warptile_mmq = l_warptile_mmq_idw; + + // Same expansion as CREATE_MM above, plus a trailing required subgroup + // size (0 = driver default) for the GGML_VK_MMID_WAVE32 probe. Scoped to + // the mmid quant pipelines below; dense pipelines are untouched. +#undef CREATE_MM +#define CREATE_MM(TYPE, PIPELINE_NAME, NAMELC, F16ACC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ + if (device->mul_mat ## ID ## _l[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->l, #NAMELC #F16ACC "_l", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, ggml_vk_mul_mm_spec(l_ ## WARPTILE, false), 1, false, true, mmid_req_sgs); \ + if (device->mul_mat ## ID ## _m[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->m, #NAMELC #F16ACC "_m", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, ggml_vk_mul_mm_spec(m_ ## WARPTILE, false), 1, false, true, mmid_req_sgs); \ + if (device->mul_mat ## ID ## _s[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->s, #NAMELC #F16ACC "_s", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, ggml_vk_mul_mm_spec(s_ ## WARPTILE, false), 1, false, true, mmid_req_sgs); \ + if (device->mul_mat ## ID ## _l[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_l, #NAMELC #F16ACC "_aligned_l", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, ggml_vk_mul_mm_spec(l_ ## WARPTILE, true), l_align, false, true, mmid_req_sgs); \ + if (device->mul_mat ## ID ## _m[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_m, #NAMELC #F16ACC "_aligned_m", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, ggml_vk_mul_mm_spec(m_ ## WARPTILE, true), m_align, false, true, mmid_req_sgs); \ + if (device->mul_mat ## ID ## _s[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_s, #NAMELC #F16ACC "_aligned_s", NAMELC ## F16ACC ## _cm1_len, NAMELC ## F16ACC ## _cm1_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, ggml_vk_mul_mm_spec(s_ ## WARPTILE, true), s_align, false, true, mmid_req_sgs); \ + + CREATE_MM2(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q1_0], matmul_id_subgroup_q1_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_0], matmul_id_subgroup_q2_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0], matmul_id_subgroup_q4_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1], matmul_id_subgroup_q4_1_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0], matmul_id_subgroup_q5_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1], matmul_id_subgroup_q5_1_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0], matmul_id_subgroup_q8_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K], matmul_id_subgroup_q2_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_TQ2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ2_0], matmul_id_subgroup_tq2_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_TQ1_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ1_0], matmul_id_subgroup_tq1_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K], matmul_id_subgroup_q3_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K], matmul_id_subgroup_q4_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K], matmul_id_subgroup_q5_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q6_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q6_K], matmul_id_subgroup_q6_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_IQ1_S, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ1_S], matmul_id_subgroup_iq1_s_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_IQ1_M, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ1_M], matmul_id_subgroup_iq1_m_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_IQ2_XXS, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ2_XXS], matmul_id_subgroup_iq2_xxs_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_IQ2_XS, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ2_XS], matmul_id_subgroup_iq2_xs_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_IQ2_S, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ2_S], matmul_id_subgroup_iq2_s_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_IQ3_XXS, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ3_XXS], matmul_id_subgroup_iq3_xxs_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_IQ3_S, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ3_S], matmul_id_subgroup_iq3_s_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_IQ4_XS, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_XS], matmul_id_subgroup_iq4_xs_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_IQ4_NL, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_NL], matmul_id_subgroup_iq4_nl_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q4_0_ROCMFP4, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0_ROCMFP4], matmul_id_subgroup_rocmfp4_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q4_0_ROCMFP4_FAST, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0_ROCMFP4_FAST], matmul_id_subgroup_rocmfp4_fast_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q3_0_ROCMFPX, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_0_ROCMFPX], matmul_id_subgroup_rocmfpx_fp3_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q6_0_ROCMFPX, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q6_0_ROCMFPX], matmul_id_subgroup_rocmfpx_fp6_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q8_0_ROCMFPX, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0_ROCMFPX], matmul_id_subgroup_rocmfpx_fp8_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); #if defined(GGML_VULKAN_FLOAT_E2M1_GLSLC_SUPPORT) && defined(GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT) if (device->ocp_fp4) { -#define X_CM1_ID_OCP(TYPE, tstr) \ - if (device->coopmat_acc_f16_support) { \ - cm1_create({TYPE, GGML_TYPE_F32, true, true}, tc_mmq_id, "matmul_id_subgroup_" #tstr "_f32_ocp_f16acc", matmul_id_subgroup_##tstr##_f32_ocp_f16acc_cm1_len, matmul_id_subgroup_##tstr##_f32_ocp_f16acc_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); \ - cm1_create({TYPE, GGML_TYPE_F16, true, true}, tc_mmq_id, "matmul_id_subgroup_" #tstr "_f16_ocp_f16acc", matmul_id_subgroup_##tstr##_f16_ocp_f16acc_cm1_len, matmul_id_subgroup_##tstr##_f16_ocp_f16acc_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); \ - } \ - if (device->coopmat_acc_f32_support) { \ - cm1_create({TYPE, GGML_TYPE_F32, true, false}, tc_mmq_id, "matmul_id_subgroup_" #tstr "_f32_ocp", matmul_id_subgroup_##tstr##_f32_ocp_cm1_len, matmul_id_subgroup_##tstr##_f32_ocp_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); \ - cm1_create({TYPE, GGML_TYPE_F16, true, false}, tc_mmq_id, "matmul_id_subgroup_" #tstr "_f16_ocp", matmul_id_subgroup_##tstr##_f16_ocp_cm1_len, matmul_id_subgroup_##tstr##_f16_ocp_cm1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); \ - } - FOR_EACH_LUT_FP4_TYPE(X_CM1_ID_OCP) -#undef X_CM1_ID_OCP + CREATE_MM2(GGML_TYPE_MXFP4, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_MXFP4], matmul_id_subgroup_mxfp4_f32_ocp, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_NVFP4, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_NVFP4], matmul_id_subgroup_nvfp4_f32_ocp, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); } else #endif { - FOR_EACH_LUT_FP4_TYPE(X_CM1_ID) - } -#undef X_CM1_ID + CREATE_MM2(GGML_TYPE_MXFP4, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_MXFP4], matmul_id_subgroup_mxfp4_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_NVFP4, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_NVFP4], matmul_id_subgroup_nvfp4_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + } + + // EXPERIMENT (GGML_VK_MMID_F16B=1): f16-B variants of the quant matmul_id + // pipelines. The _f16 SPIR-V exists for every quant type; upstream just never + // instantiates it in the KHR_coopmat branch. Same warptiles as the f32-B lines + // (including the SMALLN/BM64/M128 tile experiments shadowed above). + if (ggml_vk_mmid_f16b_enabled()) { + CREATE_MM2(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_Q1_0], matmul_id_subgroup_q1_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_Q4_0], matmul_id_subgroup_q4_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_Q4_1], matmul_id_subgroup_q4_1_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_Q5_0], matmul_id_subgroup_q5_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_Q5_1], matmul_id_subgroup_q5_1_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_Q8_0], matmul_id_subgroup_q8_0_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_Q2_K], matmul_id_subgroup_q2_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_Q3_K], matmul_id_subgroup_q3_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_Q4_K], matmul_id_subgroup_q4_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_Q5_K], matmul_id_subgroup_q5_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_Q6_K, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_Q6_K], matmul_id_subgroup_q6_k_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_IQ1_S, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_IQ1_S], matmul_id_subgroup_iq1_s_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_IQ1_M, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_IQ1_M], matmul_id_subgroup_iq1_m_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_IQ2_XXS, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_IQ2_XXS], matmul_id_subgroup_iq2_xxs_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_IQ2_XS, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_IQ2_XS], matmul_id_subgroup_iq2_xs_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_IQ2_S, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_IQ2_S], matmul_id_subgroup_iq2_s_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_IQ3_XXS, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_IQ3_XXS], matmul_id_subgroup_iq3_xxs_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_IQ3_S, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_IQ3_S], matmul_id_subgroup_iq3_s_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_IQ4_XS, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_IQ4_XS], matmul_id_subgroup_iq4_xs_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_IQ4_NL, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_IQ4_NL], matmul_id_subgroup_iq4_nl_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_MXFP4, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_MXFP4], matmul_id_subgroup_mxfp4_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + CREATE_MM2(GGML_TYPE_NVFP4, pipeline_dequant_mul_mat_mat_id_f16b[GGML_TYPE_NVFP4], matmul_id_subgroup_nvfp4_f16, mmq_wg_denoms, warptile_mmq, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id); + } + } + +#undef CREATE_MM2 +#undef CREATE_MM } else #endif // defined(VK_KHR_cooperative_matrix) && defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT) - { - // Helper for subgroup path with dot2 selection and filtering - auto sg_create = [&](vk_matmul_pipeline_key key, const std::vector& tc_base, - const std::string& name, size_t len, const void* data, uint32_t pc_size, uint32_t pc, - uint32_t rsgs = 0) { - auto tc = filter_tc(tc_base, key.type_a, key.mul_mat_id); - if (!tc.empty()) create_mm_pipelines(key, tc, name, len, data, pc_size, pc, - [&](const std::vector& wt, bool a) { return ggml_vk_mul_mm_spec(wt, a); }, - false, rsgs > 0, rsgs); - }; - auto sg_create_quant = [&](vk_matmul_pipeline_key key, const std::vector& tc_base, - const std::string& name, size_t len, const void* data, uint32_t pc_size, uint32_t pc, - uint32_t rsgs = 0) { - auto tc = filter_tc(tc_base, key.type_a, key.mul_mat_id); - if (!tc.empty()) { - spec_fn_t qs = [&, type_a=key.type_a](const std::vector& wt, bool a) { return ggml_vk_mul_mm_spec_quant(wt, a, (uint32_t)type_a); }; - create_mm_pipelines(key, tc, name, len, data, pc_size, pc, qs, false, rsgs > 0, rsgs); - } - }; + if (device->fp16) { + // Create 6 variants, {s,m,l}x{unaligned,aligned} + // Selects dot2 SPIR-V variant at runtime when device->dot2_f16 is true +#define CREATE_MM(TYPE, PIPELINE_NAME, NAMELC, F16ACC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID, REQSUBGROUPSIZE) \ + if (device->mul_mat ## ID ## _l[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->l, #NAMELC #F16ACC "_l", (device->dot2_f16 ? NAMELC ## _dot2 ## F16ACC ## _len : NAMELC ## F16ACC ## _len), (device->dot2_f16 ? NAMELC ## _dot2 ## F16ACC ## _data : NAMELC ## F16ACC ## _data), "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, ggml_vk_mul_mm_spec(l_ ## WARPTILE, false), 1, false, REQSUBGROUPSIZE > 0, REQSUBGROUPSIZE); \ + if (device->mul_mat ## ID ## _m[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->m, #NAMELC #F16ACC "_m", (device->dot2_f16 ? NAMELC ## _dot2 ## F16ACC ## _len : NAMELC ## F16ACC ## _len), (device->dot2_f16 ? NAMELC ## _dot2 ## F16ACC ## _data : NAMELC ## F16ACC ## _data), "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, ggml_vk_mul_mm_spec(m_ ## WARPTILE, false), 1, false, REQSUBGROUPSIZE > 0, REQSUBGROUPSIZE); \ + if (device->mul_mat ## ID ## _s[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->s, #NAMELC #F16ACC "_s", (device->dot2_f16 ? NAMELC ## _dot2 ## F16ACC ## _len : NAMELC ## F16ACC ## _len), (device->dot2_f16 ? NAMELC ## _dot2 ## F16ACC ## _data : NAMELC ## F16ACC ## _data), "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, ggml_vk_mul_mm_spec(s_ ## WARPTILE, false), 1, false, REQSUBGROUPSIZE > 0, REQSUBGROUPSIZE); \ + if (device->mul_mat ## ID ## _l[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_l, #NAMELC #F16ACC "_aligned_l", (device->dot2_f16 ? NAMELC ## _dot2 ## F16ACC ## _len : NAMELC ## F16ACC ## _len), (device->dot2_f16 ? NAMELC ## _dot2 ## F16ACC ## _data : NAMELC ## F16ACC ## _data), "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, ggml_vk_mul_mm_spec(l_ ## WARPTILE, true), l_align, false, REQSUBGROUPSIZE > 0, REQSUBGROUPSIZE); \ + if (device->mul_mat ## ID ## _m[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_m, #NAMELC #F16ACC "_aligned_m", (device->dot2_f16 ? NAMELC ## _dot2 ## F16ACC ## _len : NAMELC ## F16ACC ## _len), (device->dot2_f16 ? NAMELC ## _dot2 ## F16ACC ## _data : NAMELC ## F16ACC ## _data), "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, ggml_vk_mul_mm_spec(m_ ## WARPTILE, true), m_align, false, REQSUBGROUPSIZE > 0, REQSUBGROUPSIZE); \ + if (device->mul_mat ## ID ## _s[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_s, #NAMELC #F16ACC "_aligned_s", (device->dot2_f16 ? NAMELC ## _dot2 ## F16ACC ## _len : NAMELC ## F16ACC ## _len), (device->dot2_f16 ? NAMELC ## _dot2 ## F16ACC ## _data : NAMELC ## F16ACC ## _data), "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, ggml_vk_mul_mm_spec(s_ ## WARPTILE, true), s_align, false, REQSUBGROUPSIZE > 0, REQSUBGROUPSIZE); \ + + // bf16 scalar path promotes to f32, no dot2 variant +#define CREATE_MM_NODOT2(TYPE, PIPELINE_NAME, NAMELC, F16ACC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID, REQSUBGROUPSIZE) \ + if (device->mul_mat ## ID ## _l[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->l, #NAMELC #F16ACC "_l", NAMELC ## F16ACC ## _len, NAMELC ## F16ACC ## _data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, ggml_vk_mul_mm_spec(l_ ## WARPTILE, false), 1, false, REQSUBGROUPSIZE > 0, REQSUBGROUPSIZE); \ + if (device->mul_mat ## ID ## _m[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->m, #NAMELC #F16ACC "_m", NAMELC ## F16ACC ## _len, NAMELC ## F16ACC ## _data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, ggml_vk_mul_mm_spec(m_ ## WARPTILE, false), 1, false, REQSUBGROUPSIZE > 0, REQSUBGROUPSIZE); \ + if (device->mul_mat ## ID ## _s[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->s, #NAMELC #F16ACC "_s", NAMELC ## F16ACC ## _len, NAMELC ## F16ACC ## _data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, ggml_vk_mul_mm_spec(s_ ## WARPTILE, false), 1, false, REQSUBGROUPSIZE > 0, REQSUBGROUPSIZE); \ + if (device->mul_mat ## ID ## _l[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_l, #NAMELC #F16ACC "_aligned_l", NAMELC ## F16ACC ## _len, NAMELC ## F16ACC ## _data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, ggml_vk_mul_mm_spec(l_ ## WARPTILE, true), l_align, false, REQSUBGROUPSIZE > 0, REQSUBGROUPSIZE); \ + if (device->mul_mat ## ID ## _m[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_m, #NAMELC #F16ACC "_aligned_m", NAMELC ## F16ACC ## _len, NAMELC ## F16ACC ## _data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, ggml_vk_mul_mm_spec(m_ ## WARPTILE, true), m_align, false, REQSUBGROUPSIZE > 0, REQSUBGROUPSIZE); \ + if (device->mul_mat ## ID ## _s[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_s, #NAMELC #F16ACC "_aligned_s", NAMELC ## F16ACC ## _len, NAMELC ## F16ACC ## _data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, ggml_vk_mul_mm_spec(s_ ## WARPTILE, true), s_align, false, REQSUBGROUPSIZE > 0, REQSUBGROUPSIZE); \ + +#define CREATE_MMQ(TYPE, PIPELINE_NAME, NAMELC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID, REQSUBGROUPSIZE) \ + if (device->mul_mat ## ID ## _l_int[TYPE]) { \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME .f32acc->l, #NAMELC "_l", NAMELC ## _len, NAMELC ## _data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, 1, false, REQSUBGROUPSIZE > 0, REQSUBGROUPSIZE); \ + } \ + if (device->mul_mat ## ID ## _m_int[TYPE]) { \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME .f32acc->m, #NAMELC "_m", NAMELC ## _len, NAMELC ## _data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, 1, false, REQSUBGROUPSIZE > 0, REQSUBGROUPSIZE); \ + } \ + if (device->mul_mat ## ID ## _s_int[TYPE]) { \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME .f32acc->s, #NAMELC "_s", NAMELC ## _len, NAMELC ## _data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, 1, false, REQSUBGROUPSIZE > 0, REQSUBGROUPSIZE); \ + } \ + + // Create 2 variants, {f16,f32} accumulator +#define CREATE_MM2(TYPE, PIPELINE_NAME, NAMELC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID, REQSUBGROUPSIZE) \ + CREATE_MM(TYPE, PIPELINE_NAME . f16acc, NAMELC, _f16acc, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID, REQSUBGROUPSIZE) \ + CREATE_MM(TYPE, PIPELINE_NAME . f32acc, NAMELC, , WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID, REQSUBGROUPSIZE) \ + + CREATE_MM(GGML_TYPE_F32, pipeline_matmul_f32, matmul_f32_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_F32, pipeline_matmul_f32_f16, matmul_f32_f16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM2(GGML_TYPE_F16, pipeline_matmul_f16, matmul_f16, wg_denoms, warptile, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM2(GGML_TYPE_F16, pipeline_matmul_f16_f32, matmul_f16_f32, wg_denoms, warptile, vk_mat_mat_push_constants, 3, , 0); + + CREATE_MM_NODOT2(GGML_TYPE_BF16, pipeline_matmul_bf16, matmul_bf16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, , 0); + + CREATE_MM2(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q1_0], matmul_q1_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM2(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_0], matmul_q2_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM2(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_0], matmul_q4_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM2(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_1], matmul_q4_1_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM2(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_0], matmul_q5_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM2(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_1], matmul_q5_1_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM2(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q8_0], matmul_q8_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM2(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_K], matmul_q2_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM2(GGML_TYPE_TQ2_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_TQ2_0], matmul_tq2_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM2(GGML_TYPE_TQ1_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_TQ1_0], matmul_tq1_0_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM2(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_K], matmul_q3_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM2(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_K], matmul_q4_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM2(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_K], matmul_q5_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM2(GGML_TYPE_Q6_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q6_K], matmul_q6_k_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM2(GGML_TYPE_IQ1_S, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ1_S], matmul_iq1_s_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM2(GGML_TYPE_IQ1_M, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ1_M], matmul_iq1_m_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM2(GGML_TYPE_IQ2_XXS, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ2_XXS], matmul_iq2_xxs_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM2(GGML_TYPE_IQ2_XS, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ2_XS], matmul_iq2_xs_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM2(GGML_TYPE_IQ2_S, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ2_S], matmul_iq2_s_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM2(GGML_TYPE_IQ3_XXS, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ3_XXS], matmul_iq3_xxs_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM2(GGML_TYPE_IQ3_S, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ3_S], matmul_iq3_s_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM2(GGML_TYPE_IQ4_XS, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ4_XS], matmul_iq4_xs_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM2(GGML_TYPE_IQ4_NL, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ4_NL], matmul_iq4_nl_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM2(GGML_TYPE_Q4_0_ROCMFP4, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_0_ROCMFP4], matmul_rocmfp4_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM2(GGML_TYPE_Q4_0_ROCMFP4_FAST, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_0_ROCMFP4_FAST], matmul_rocmfp4_fast_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM2(GGML_TYPE_Q3_0_ROCMFPX, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_0_ROCMFPX], matmul_rocmfpx_fp3_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM2(GGML_TYPE_Q6_0_ROCMFPX, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q6_0_ROCMFPX], matmul_rocmfpx_fp6_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM2(GGML_TYPE_Q8_0_ROCMFPX, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q8_0_ROCMFPX], matmul_rocmfpx_fp8_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM2(GGML_TYPE_MXFP4, pipeline_dequant_mul_mat_mat[GGML_TYPE_MXFP4], matmul_mxfp4_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM2(GGML_TYPE_NVFP4, pipeline_dequant_mul_mat_mat[GGML_TYPE_NVFP4], matmul_nvfp4_f32, mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + #if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT) - auto sg_create_mmq = [&](vk_matmul_pipeline_key key, const std::vector& tc_base, - const std::string& name, size_t len, const void* data, uint32_t pc_size, uint32_t pc, - uint32_t rsgs = 0) { - auto tc = filter_tc(tc_base, key.type_a, key.mul_mat_id, true); - if (!tc.empty()) { - spec_fn_t identity = [](const std::vector& wt, bool) { return wt; }; - create_mm_pipelines(key, tc, name, len, data, pc_size, pc, identity, false, rsgs > 0, rsgs, false); - } - }; + if (device->integer_dot_product) { + CREATE_MMQ(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q2_0], matmul_q2_0_q8_1, mmq_wg_denoms, warptile_mmq_int, vk_mat_mat_push_constants, 3, , 0); + CREATE_MMQ(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q4_0], matmul_q4_0_q8_1, mmq_wg_denoms, warptile_mmq_int, vk_mat_mat_push_constants, 3, , 0); + CREATE_MMQ(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q4_1], matmul_q4_1_q8_1, mmq_wg_denoms, warptile_mmq_int, vk_mat_mat_push_constants, 3, , 0); + CREATE_MMQ(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q5_0], matmul_q5_0_q8_1, mmq_wg_denoms, warptile_mmq_int, vk_mat_mat_push_constants, 3, , 0); + CREATE_MMQ(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q5_1], matmul_q5_1_q8_1, mmq_wg_denoms, warptile_mmq_int, vk_mat_mat_push_constants, 3, , 0); + CREATE_MMQ(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q8_0], matmul_q8_0_q8_1, mmq_wg_denoms, warptile_mmq_int, vk_mat_mat_push_constants, 3, , 0); + + CREATE_MMQ(GGML_TYPE_MXFP4, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_MXFP4], matmul_mxfp4_q8_1, mmq_wg_denoms, warptile_mmq_int, vk_mat_mat_push_constants, 3, , 0); + CREATE_MMQ(GGML_TYPE_Q4_0_ROCMFP4, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q4_0_ROCMFP4], matmul_rocmfp4_q8_1, mmq_wg_denoms, warptile_mmq_int, vk_mat_mat_push_constants, 3, , 0); + CREATE_MMQ(GGML_TYPE_Q4_0_ROCMFP4_FAST, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q4_0_ROCMFP4_FAST], matmul_rocmfp4_fast_q8_1, mmq_wg_denoms, warptile_mmq_int, vk_mat_mat_push_constants, 3, , 0); + + CREATE_MMQ(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q2_K], matmul_q2_k_q8_1, mmq_wg_denoms, warptile_mmq_int_k, vk_mat_mat_push_constants, 3, , 0); + CREATE_MMQ(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q3_K], matmul_q3_k_q8_1, mmq_wg_denoms, warptile_mmq_int_k, vk_mat_mat_push_constants, 3, , 0); + CREATE_MMQ(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q4_K], matmul_q4_k_q8_1, mmq_wg_denoms, warptile_mmq_int_k, vk_mat_mat_push_constants, 3, , 0); + CREATE_MMQ(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q5_K], matmul_q5_k_q8_1, mmq_wg_denoms, warptile_mmq_int_k, vk_mat_mat_push_constants, 3, , 0); + CREATE_MMQ(GGML_TYPE_Q6_K, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q6_K], matmul_q6_k_q8_1, mmq_wg_denoms, warptile_mmq_int_k, vk_mat_mat_push_constants, 3, , 0); + } #endif - std::vector tc_id = {{s_warptile_id, s_wg_denoms, s_align}, {m_warptile_id, m_wg_denoms, m_align}, {l_warptile_id, l_wg_denoms, l_align}}; - std::vector tc_mmqid = {{s_warptile_mmqid, s_mmq_wg_denoms, s_align}, {m_warptile_mmqid, m_mmq_wg_denoms, m_align}, {l_warptile_mmqid, l_mmq_wg_denoms, l_align}}; - - if (device->fp16) { - // FP16 subgroup path - with dot2 runtime selection - #define SPV_DOT2(NAME) (device->dot2_f16 ? NAME ## _dot2_len : NAME ## _len), (device->dot2_f16 ? NAME ## _dot2_data : NAME ## _data) - #define SPV_DOT2_F16ACC(NAME) (device->dot2_f16 ? NAME ## _dot2_f16acc_len : NAME ## _f16acc_len), (device->dot2_f16 ? NAME ## _dot2_f16acc_data : NAME ## _f16acc_data) - - sg_create({GGML_TYPE_F32, GGML_TYPE_F32, false, false}, tc_mm, "matmul_f32_f32", SPV_DOT2(matmul_f32_f32), sizeof(vk_mat_mat_push_constants), 3); - sg_create({GGML_TYPE_F32, GGML_TYPE_F16, false, false}, tc_mm, "matmul_f32_f16", SPV_DOT2(matmul_f32_f16), sizeof(vk_mat_mat_push_constants), 3); - sg_create({GGML_TYPE_F16, GGML_TYPE_F16, false, true}, tc_mm, "matmul_f16_f16acc", SPV_DOT2_F16ACC(matmul_f16), sizeof(vk_mat_mat_push_constants), 3); - sg_create({GGML_TYPE_F16, GGML_TYPE_F16, false, false}, tc_mm, "matmul_f16", SPV_DOT2(matmul_f16), sizeof(vk_mat_mat_push_constants), 3); - sg_create({GGML_TYPE_F16, GGML_TYPE_F32, false, true}, tc_mm, "matmul_f16_f32_f16acc", SPV_DOT2_F16ACC(matmul_f16_f32), sizeof(vk_mat_mat_push_constants), 3); - sg_create({GGML_TYPE_F16, GGML_TYPE_F32, false, false}, tc_mm, "matmul_f16_f32", SPV_DOT2(matmul_f16_f32), sizeof(vk_mat_mat_push_constants), 3); - // BF16 - no dot2 - sg_create({GGML_TYPE_BF16, GGML_TYPE_BF16, false, false}, tc_mm, "matmul_bf16", matmul_bf16_len, matmul_bf16_data, sizeof(vk_mat_mat_push_constants), 3); - - for (const auto type : non_lut_quant_types) { - sg_create_quant({type, GGML_TYPE_F32, false, true}, tc_mmq, "matmul_quant_f32_f16acc", SPV_DOT2_F16ACC(matmul_quant_f32), sizeof(vk_mat_mat_push_constants), 3); - sg_create_quant({type, GGML_TYPE_F32, false, false}, tc_mmq, "matmul_quant_f32", SPV_DOT2(matmul_quant_f32), sizeof(vk_mat_mat_push_constants), 3); - } - #define X_SG(TYPE, tstr) \ - sg_create({TYPE, GGML_TYPE_F32, false, true}, tc_mmq, "matmul_" #tstr "_f32_f16acc", SPV_DOT2_F16ACC(matmul_##tstr##_f32), sizeof(vk_mat_mat_push_constants), 3); \ - sg_create({TYPE, GGML_TYPE_F32, false, false}, tc_mmq, "matmul_" #tstr "_f32", SPV_DOT2(matmul_##tstr##_f32), sizeof(vk_mat_mat_push_constants), 3); - FOR_EACH_LUT_TYPE(X_SG) -#undef X_SG + if (device->subgroup_ballot && device->subgroup_require_full_support && subgroup_min_size_16) { + CREATE_MM(GGML_TYPE_F32, pipeline_matmul_id_f32, matmul_id_subgroup_f32_f32, , wg_denoms, warptile_id, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size_16); + CREATE_MM2(GGML_TYPE_F16, pipeline_matmul_id_f16, matmul_id_subgroup_f16, wg_denoms, warptile_id, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size_16); + CREATE_MM2(GGML_TYPE_F16, pipeline_matmul_id_f16_f32, matmul_id_subgroup_f16_f32, wg_denoms, warptile_id, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size_16); + CREATE_MM_NODOT2(GGML_TYPE_BF16, pipeline_matmul_id_bf16, matmul_id_subgroup_bf16, , wg_denoms, warptile_id, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size_16); + CREATE_MM2(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q1_0], matmul_id_subgroup_q1_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM2(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_0], matmul_id_subgroup_q2_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM2(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0], matmul_id_subgroup_q4_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM2(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1], matmul_id_subgroup_q4_1_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM2(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0], matmul_id_subgroup_q5_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM2(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1], matmul_id_subgroup_q5_1_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM2(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0], matmul_id_subgroup_q8_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM2(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K], matmul_id_subgroup_q2_k_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM2(GGML_TYPE_TQ2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ2_0], matmul_id_subgroup_tq2_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM2(GGML_TYPE_TQ1_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ1_0], matmul_id_subgroup_tq1_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM2(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K], matmul_id_subgroup_q3_k_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM2(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K], matmul_id_subgroup_q4_k_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM2(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K], matmul_id_subgroup_q5_k_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM2(GGML_TYPE_Q6_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q6_K], matmul_id_subgroup_q6_k_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM2(GGML_TYPE_IQ1_S, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ1_S], matmul_id_subgroup_iq1_s_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM2(GGML_TYPE_IQ1_M, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ1_M], matmul_id_subgroup_iq1_m_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM2(GGML_TYPE_IQ2_XXS, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ2_XXS], matmul_id_subgroup_iq2_xxs_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM2(GGML_TYPE_IQ2_XS, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ2_XS], matmul_id_subgroup_iq2_xs_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM2(GGML_TYPE_IQ2_S, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ2_S], matmul_id_subgroup_iq2_s_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM2(GGML_TYPE_IQ3_XXS, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ3_XXS], matmul_id_subgroup_iq3_xxs_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM2(GGML_TYPE_IQ3_S, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ3_S], matmul_id_subgroup_iq3_s_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM2(GGML_TYPE_IQ4_XS, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_XS], matmul_id_subgroup_iq4_xs_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM2(GGML_TYPE_IQ4_NL, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_NL], matmul_id_subgroup_iq4_nl_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM2(GGML_TYPE_Q4_0_ROCMFP4, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0_ROCMFP4], matmul_id_subgroup_rocmfp4_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM2(GGML_TYPE_Q4_0_ROCMFP4_FAST, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0_ROCMFP4_FAST], matmul_id_subgroup_rocmfp4_fast_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM2(GGML_TYPE_Q3_0_ROCMFPX, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_0_ROCMFPX], matmul_id_subgroup_rocmfpx_fp3_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM2(GGML_TYPE_Q6_0_ROCMFPX, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q6_0_ROCMFPX], matmul_id_subgroup_rocmfpx_fp6_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM2(GGML_TYPE_Q8_0_ROCMFPX, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0_ROCMFPX], matmul_id_subgroup_rocmfpx_fp8_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM2(GGML_TYPE_MXFP4, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_MXFP4], matmul_id_subgroup_mxfp4_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM2(GGML_TYPE_NVFP4, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_NVFP4], matmul_id_subgroup_nvfp4_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); #if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT) if (device->integer_dot_product) { - std::vector tc_mmq_int = {{s_warptile_mmq_int, s_mmq_wg_denoms, s_align}, {m_warptile_mmq_int, m_mmq_wg_denoms, m_align}, {l_warptile_mmq_int, l_mmq_wg_denoms, l_align}}; - std::vector tc_mmq_int_k = {{s_warptile_mmq_int_k, s_mmq_wg_denoms, s_align}, {m_warptile_mmq_int_k, m_mmq_wg_denoms, m_align}, {l_warptile_mmq_int_k, l_mmq_wg_denoms, l_align}}; - sg_create_mmq({GGML_TYPE_Q2_0, GGML_TYPE_Q8_1, false, false}, tc_mmq_int, "matmul_q2_0_q8_1", matmul_q2_0_q8_1_len, matmul_q2_0_q8_1_data, sizeof(vk_mat_mat_push_constants), 3); - sg_create_mmq({GGML_TYPE_Q4_0, GGML_TYPE_Q8_1, false, false}, tc_mmq_int, "matmul_q4_0_q8_1", matmul_q4_0_q8_1_len, matmul_q4_0_q8_1_data, sizeof(vk_mat_mat_push_constants), 3); - sg_create_mmq({GGML_TYPE_Q4_1, GGML_TYPE_Q8_1, false, false}, tc_mmq_int, "matmul_q4_1_q8_1", matmul_q4_1_q8_1_len, matmul_q4_1_q8_1_data, sizeof(vk_mat_mat_push_constants), 3); - sg_create_mmq({GGML_TYPE_Q5_0, GGML_TYPE_Q8_1, false, false}, tc_mmq_int, "matmul_q5_0_q8_1", matmul_q5_0_q8_1_len, matmul_q5_0_q8_1_data, sizeof(vk_mat_mat_push_constants), 3); - sg_create_mmq({GGML_TYPE_Q5_1, GGML_TYPE_Q8_1, false, false}, tc_mmq_int, "matmul_q5_1_q8_1", matmul_q5_1_q8_1_len, matmul_q5_1_q8_1_data, sizeof(vk_mat_mat_push_constants), 3); - sg_create_mmq({GGML_TYPE_Q8_0, GGML_TYPE_Q8_1, false, false}, tc_mmq_int, "matmul_q8_0_q8_1", matmul_q8_0_q8_1_len, matmul_q8_0_q8_1_data, sizeof(vk_mat_mat_push_constants), 3); - sg_create_mmq({GGML_TYPE_MXFP4, GGML_TYPE_Q8_1, false, false}, tc_mmq_int, "matmul_mxfp4_q8_1", matmul_mxfp4_q8_1_len, matmul_mxfp4_q8_1_data, sizeof(vk_mat_mat_push_constants), 3); - sg_create_mmq({GGML_TYPE_Q2_K, GGML_TYPE_Q8_1, false, false}, tc_mmq_int_k, "matmul_q2_k_q8_1", matmul_q2_k_q8_1_len, matmul_q2_k_q8_1_data, sizeof(vk_mat_mat_push_constants), 3); - sg_create_mmq({GGML_TYPE_Q3_K, GGML_TYPE_Q8_1, false, false}, tc_mmq_int_k, "matmul_q3_k_q8_1", matmul_q3_k_q8_1_len, matmul_q3_k_q8_1_data, sizeof(vk_mat_mat_push_constants), 3); - sg_create_mmq({GGML_TYPE_Q4_K, GGML_TYPE_Q8_1, false, false}, tc_mmq_int_k, "matmul_q4_k_q8_1", matmul_q4_k_q8_1_len, matmul_q4_k_q8_1_data, sizeof(vk_mat_mat_push_constants), 3); - sg_create_mmq({GGML_TYPE_Q5_K, GGML_TYPE_Q8_1, false, false}, tc_mmq_int_k, "matmul_q5_k_q8_1", matmul_q5_k_q8_1_len, matmul_q5_k_q8_1_data, sizeof(vk_mat_mat_push_constants), 3); - sg_create_mmq({GGML_TYPE_Q6_K, GGML_TYPE_Q8_1, false, false}, tc_mmq_int_k, "matmul_q6_k_q8_1", matmul_q6_k_q8_1_len, matmul_q6_k_q8_1_data, sizeof(vk_mat_mat_push_constants), 3); + CREATE_MMQ(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q2_0], matmul_id_subgroup_q2_0_q8_1, mmq_wg_denoms, warptile_mmqid_int, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MMQ(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q4_0], matmul_id_subgroup_q4_0_q8_1, mmq_wg_denoms, warptile_mmqid_int, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MMQ(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q4_1], matmul_id_subgroup_q4_1_q8_1, mmq_wg_denoms, warptile_mmqid_int, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MMQ(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q5_0], matmul_id_subgroup_q5_0_q8_1, mmq_wg_denoms, warptile_mmqid_int, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MMQ(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q5_1], matmul_id_subgroup_q5_1_q8_1, mmq_wg_denoms, warptile_mmqid_int, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MMQ(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q8_0], matmul_id_subgroup_q8_0_q8_1, mmq_wg_denoms, warptile_mmqid_int, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + + CREATE_MMQ(GGML_TYPE_MXFP4, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_MXFP4], matmul_id_subgroup_mxfp4_q8_1, mmq_wg_denoms, warptile_mmqid_int, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MMQ(GGML_TYPE_Q4_0_ROCMFP4, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q4_0_ROCMFP4], matmul_id_subgroup_rocmfp4_q8_1, mmq_wg_denoms, warptile_mmqid_int, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MMQ(GGML_TYPE_Q4_0_ROCMFP4_FAST, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q4_0_ROCMFP4_FAST], matmul_id_subgroup_rocmfp4_fast_q8_1, mmq_wg_denoms, warptile_mmqid_int, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + + CREATE_MMQ(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q2_K], matmul_id_subgroup_q2_k_q8_1, mmq_wg_denoms, warptile_mmqid_int_k, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size_16); + CREATE_MMQ(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q3_K], matmul_id_subgroup_q3_k_q8_1, mmq_wg_denoms, warptile_mmqid_int_k, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size_16); + CREATE_MMQ(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q4_K], matmul_id_subgroup_q4_k_q8_1, mmq_wg_denoms, warptile_mmqid_int_k, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size_16); + CREATE_MMQ(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q5_K], matmul_id_subgroup_q5_k_q8_1, mmq_wg_denoms, warptile_mmqid_int_k, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size_16); + CREATE_MMQ(GGML_TYPE_Q6_K, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q6_K], matmul_id_subgroup_q6_k_q8_1, mmq_wg_denoms, warptile_mmqid_int_k, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size_16); } #endif - - if (device->subgroup_ballot && device->subgroup_require_full_support && subgroup_min_size_16) { - sg_create({GGML_TYPE_F32, GGML_TYPE_F32, true, false}, tc_id, "matmul_id_subgroup_f32_f32", SPV_DOT2(matmul_id_subgroup_f32_f32), sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count, mul_mat_subgroup_size_16); - sg_create({GGML_TYPE_F16, GGML_TYPE_F16, true, true}, tc_id, "matmul_id_subgroup_f16_f16acc", SPV_DOT2_F16ACC(matmul_id_subgroup_f16), sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count, mul_mat_subgroup_size_16); - sg_create({GGML_TYPE_F16, GGML_TYPE_F16, true, false}, tc_id, "matmul_id_subgroup_f16", SPV_DOT2(matmul_id_subgroup_f16), sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count, mul_mat_subgroup_size_16); - sg_create({GGML_TYPE_F16, GGML_TYPE_F32, true, true}, tc_id, "matmul_id_subgroup_f16_f32_f16acc", SPV_DOT2_F16ACC(matmul_id_subgroup_f16_f32), sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count, mul_mat_subgroup_size_16); - sg_create({GGML_TYPE_F16, GGML_TYPE_F32, true, false}, tc_id, "matmul_id_subgroup_f16_f32", SPV_DOT2(matmul_id_subgroup_f16_f32), sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count, mul_mat_subgroup_size_16); - // BF16 id - no dot2 - sg_create({GGML_TYPE_BF16, GGML_TYPE_BF16, true, false}, tc_id, "matmul_id_subgroup_bf16", matmul_id_subgroup_bf16_len, matmul_id_subgroup_bf16_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count, mul_mat_subgroup_size_16); - for (const auto type : non_lut_quant_types) { - sg_create_quant({type, GGML_TYPE_F32, true, true}, tc_mmqid, "matmul_id_subgroup_quant_f32_f16acc", SPV_DOT2_F16ACC(matmul_id_subgroup_quant_f32), sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count, mul_mat_subgroup_size); - sg_create_quant({type, GGML_TYPE_F32, true, false}, tc_mmqid, "matmul_id_subgroup_quant_f32", SPV_DOT2(matmul_id_subgroup_quant_f32), sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count, mul_mat_subgroup_size); - } - #define X_SG_ID_SUB(TYPE, tstr) \ - sg_create({TYPE, GGML_TYPE_F32, true, true}, tc_mmqid, "matmul_id_subgroup_" #tstr "_f32_f16acc", SPV_DOT2_F16ACC(matmul_id_subgroup_##tstr##_f32), sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count, mul_mat_subgroup_size); \ - sg_create({TYPE, GGML_TYPE_F32, true, false}, tc_mmqid, "matmul_id_subgroup_" #tstr "_f32", SPV_DOT2(matmul_id_subgroup_##tstr##_f32), sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count, mul_mat_subgroup_size); - FOR_EACH_LUT_TYPE(X_SG_ID_SUB) -#undef X_SG_ID_SUB -#if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT) - if (device->integer_dot_product) { - std::vector tc_mmqid_int = {{s_warptile_mmqid_int, s_mmq_wg_denoms, s_align}, {m_warptile_mmqid_int, m_mmq_wg_denoms, m_align}, {l_warptile_mmqid_int, l_mmq_wg_denoms, l_align}}; - std::vector tc_mmqid_int_k = {{s_warptile_mmqid_int_k, s_mmq_wg_denoms, s_align}, {m_warptile_mmqid_int_k, m_mmq_wg_denoms, m_align}, {l_warptile_mmqid_int_k, l_mmq_wg_denoms, l_align}}; - sg_create_mmq({GGML_TYPE_Q2_0, GGML_TYPE_Q8_1, true, false}, tc_mmqid_int, "matmul_id_subgroup_q2_0_q8_1", matmul_id_subgroup_q2_0_q8_1_len, matmul_id_subgroup_q2_0_q8_1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count, mul_mat_subgroup_size); - sg_create_mmq({GGML_TYPE_Q4_0, GGML_TYPE_Q8_1, true, false}, tc_mmqid_int, "matmul_id_subgroup_q4_0_q8_1", matmul_id_subgroup_q4_0_q8_1_len, matmul_id_subgroup_q4_0_q8_1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count, mul_mat_subgroup_size); - sg_create_mmq({GGML_TYPE_Q4_1, GGML_TYPE_Q8_1, true, false}, tc_mmqid_int, "matmul_id_subgroup_q4_1_q8_1", matmul_id_subgroup_q4_1_q8_1_len, matmul_id_subgroup_q4_1_q8_1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count, mul_mat_subgroup_size); - sg_create_mmq({GGML_TYPE_Q5_0, GGML_TYPE_Q8_1, true, false}, tc_mmqid_int, "matmul_id_subgroup_q5_0_q8_1", matmul_id_subgroup_q5_0_q8_1_len, matmul_id_subgroup_q5_0_q8_1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count, mul_mat_subgroup_size); - sg_create_mmq({GGML_TYPE_Q5_1, GGML_TYPE_Q8_1, true, false}, tc_mmqid_int, "matmul_id_subgroup_q5_1_q8_1", matmul_id_subgroup_q5_1_q8_1_len, matmul_id_subgroup_q5_1_q8_1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count, mul_mat_subgroup_size); - sg_create_mmq({GGML_TYPE_Q8_0, GGML_TYPE_Q8_1, true, false}, tc_mmqid_int, "matmul_id_subgroup_q8_0_q8_1", matmul_id_subgroup_q8_0_q8_1_len, matmul_id_subgroup_q8_0_q8_1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count, mul_mat_subgroup_size); - sg_create_mmq({GGML_TYPE_MXFP4, GGML_TYPE_Q8_1, true, false}, tc_mmqid_int, "matmul_id_subgroup_mxfp4_q8_1", matmul_id_subgroup_mxfp4_q8_1_len, matmul_id_subgroup_mxfp4_q8_1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count, mul_mat_subgroup_size); - sg_create_mmq({GGML_TYPE_Q2_K, GGML_TYPE_Q8_1, true, false}, tc_mmqid_int_k, "matmul_id_subgroup_q2_k_q8_1", matmul_id_subgroup_q2_k_q8_1_len, matmul_id_subgroup_q2_k_q8_1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count, mul_mat_subgroup_size_16); - sg_create_mmq({GGML_TYPE_Q3_K, GGML_TYPE_Q8_1, true, false}, tc_mmqid_int_k, "matmul_id_subgroup_q3_k_q8_1", matmul_id_subgroup_q3_k_q8_1_len, matmul_id_subgroup_q3_k_q8_1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count, mul_mat_subgroup_size_16); - sg_create_mmq({GGML_TYPE_Q4_K, GGML_TYPE_Q8_1, true, false}, tc_mmqid_int_k, "matmul_id_subgroup_q4_k_q8_1", matmul_id_subgroup_q4_k_q8_1_len, matmul_id_subgroup_q4_k_q8_1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count, mul_mat_subgroup_size_16); - sg_create_mmq({GGML_TYPE_Q5_K, GGML_TYPE_Q8_1, true, false}, tc_mmqid_int_k, "matmul_id_subgroup_q5_k_q8_1", matmul_id_subgroup_q5_k_q8_1_len, matmul_id_subgroup_q5_k_q8_1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count, mul_mat_subgroup_size_16); - sg_create_mmq({GGML_TYPE_Q6_K, GGML_TYPE_Q8_1, true, false}, tc_mmqid_int_k, "matmul_id_subgroup_q6_k_q8_1", matmul_id_subgroup_q6_k_q8_1_len, matmul_id_subgroup_q6_k_q8_1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count, mul_mat_subgroup_size_16); - } -#endif - } else { - sg_create({GGML_TYPE_F32, GGML_TYPE_F32, true, false}, tc_mm, "matmul_id_f32_f32", SPV_DOT2(matmul_id_f32_f32), sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - sg_create({GGML_TYPE_F16, GGML_TYPE_F16, true, true}, tc_mm, "matmul_id_f16_f16acc", SPV_DOT2_F16ACC(matmul_id_f16), sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - sg_create({GGML_TYPE_F16, GGML_TYPE_F16, true, false}, tc_mm, "matmul_id_f16", SPV_DOT2(matmul_id_f16), sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - sg_create({GGML_TYPE_F16, GGML_TYPE_F32, true, true}, tc_mm, "matmul_id_f16_f32_f16acc", SPV_DOT2_F16ACC(matmul_id_f16_f32), sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - sg_create({GGML_TYPE_F16, GGML_TYPE_F32, true, false}, tc_mm, "matmul_id_f16_f32", SPV_DOT2(matmul_id_f16_f32), sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - // BF16 id - no dot2 - sg_create({GGML_TYPE_BF16, GGML_TYPE_BF16, true, false}, tc_mm, "matmul_id_bf16", matmul_id_bf16_len, matmul_id_bf16_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - for (const auto type : non_lut_quant_types) { - sg_create_quant({type, GGML_TYPE_F32, true, true}, tc_mmqid, "matmul_id_quant_f32_f16acc", SPV_DOT2_F16ACC(matmul_id_quant_f32), sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - sg_create_quant({type, GGML_TYPE_F32, true, false}, tc_mmqid, "matmul_id_quant_f32", SPV_DOT2(matmul_id_quant_f32), sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - } - #define X_SG_ID(TYPE, tstr) \ - sg_create({TYPE, GGML_TYPE_F32, true, true}, tc_mmqid, "matmul_id_" #tstr "_f32_f16acc", SPV_DOT2_F16ACC(matmul_id_##tstr##_f32), sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); \ - sg_create({TYPE, GGML_TYPE_F32, true, false}, tc_mmqid, "matmul_id_" #tstr "_f32", SPV_DOT2(matmul_id_##tstr##_f32), sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - FOR_EACH_LUT_TYPE(X_SG_ID) -#undef X_SG_ID -#if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT) - if (device->integer_dot_product) { - std::vector tc_mmqid_int = {{s_warptile_mmqid_int, s_mmq_wg_denoms, s_align}, {m_warptile_mmqid_int, m_mmq_wg_denoms, m_align}, {l_warptile_mmqid_int, l_mmq_wg_denoms, l_align}}; - std::vector tc_mmqid_int_k = {{s_warptile_mmqid_int_k, s_mmq_wg_denoms, s_align}, {m_warptile_mmqid_int_k, m_mmq_wg_denoms, m_align}, {l_warptile_mmqid_int_k, l_mmq_wg_denoms, l_align}}; - sg_create_mmq({GGML_TYPE_Q2_0, GGML_TYPE_Q8_1, true, false}, tc_mmqid_int, "matmul_id_q2_0_q8_1", matmul_id_q2_0_q8_1_len, matmul_id_q2_0_q8_1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - sg_create_mmq({GGML_TYPE_Q4_0, GGML_TYPE_Q8_1, true, false}, tc_mmqid_int, "matmul_id_q4_0_q8_1", matmul_id_q4_0_q8_1_len, matmul_id_q4_0_q8_1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - sg_create_mmq({GGML_TYPE_Q4_1, GGML_TYPE_Q8_1, true, false}, tc_mmqid_int, "matmul_id_q4_1_q8_1", matmul_id_q4_1_q8_1_len, matmul_id_q4_1_q8_1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - sg_create_mmq({GGML_TYPE_Q5_0, GGML_TYPE_Q8_1, true, false}, tc_mmqid_int, "matmul_id_q5_0_q8_1", matmul_id_q5_0_q8_1_len, matmul_id_q5_0_q8_1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - sg_create_mmq({GGML_TYPE_Q5_1, GGML_TYPE_Q8_1, true, false}, tc_mmqid_int, "matmul_id_q5_1_q8_1", matmul_id_q5_1_q8_1_len, matmul_id_q5_1_q8_1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - sg_create_mmq({GGML_TYPE_Q8_0, GGML_TYPE_Q8_1, true, false}, tc_mmqid_int, "matmul_id_q8_0_q8_1", matmul_id_q8_0_q8_1_len, matmul_id_q8_0_q8_1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - sg_create_mmq({GGML_TYPE_MXFP4, GGML_TYPE_Q8_1, true, false}, tc_mmqid_int, "matmul_id_mxfp4_q8_1", matmul_id_mxfp4_q8_1_len, matmul_id_mxfp4_q8_1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - sg_create_mmq({GGML_TYPE_Q2_K, GGML_TYPE_Q8_1, true, false}, tc_mmqid_int_k, "matmul_id_q2_k_q8_1", matmul_id_q2_k_q8_1_len, matmul_id_q2_k_q8_1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - sg_create_mmq({GGML_TYPE_Q3_K, GGML_TYPE_Q8_1, true, false}, tc_mmqid_int_k, "matmul_id_q3_k_q8_1", matmul_id_q3_k_q8_1_len, matmul_id_q3_k_q8_1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - sg_create_mmq({GGML_TYPE_Q4_K, GGML_TYPE_Q8_1, true, false}, tc_mmqid_int_k, "matmul_id_q4_k_q8_1", matmul_id_q4_k_q8_1_len, matmul_id_q4_k_q8_1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - sg_create_mmq({GGML_TYPE_Q5_K, GGML_TYPE_Q8_1, true, false}, tc_mmqid_int_k, "matmul_id_q5_k_q8_1", matmul_id_q5_k_q8_1_len, matmul_id_q5_k_q8_1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - sg_create_mmq({GGML_TYPE_Q6_K, GGML_TYPE_Q8_1, true, false}, tc_mmqid_int_k, "matmul_id_q6_k_q8_1", matmul_id_q6_k_q8_1_len, matmul_id_q6_k_q8_1_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - } -#endif - } - #undef SPV_DOT2 - #undef SPV_DOT2_F16ACC } else { - // FP32-only fallback path - sg_create({GGML_TYPE_F32, GGML_TYPE_F32, false, false}, tc_mm, "matmul_f32_f32", matmul_f32_f32_fp32_len, matmul_f32_f32_fp32_data, sizeof(vk_mat_mat_push_constants), 3); - sg_create({GGML_TYPE_F32, GGML_TYPE_F16, false, false}, tc_mm, "matmul_f32_f16", matmul_f32_f16_fp32_len, matmul_f32_f16_fp32_data, sizeof(vk_mat_mat_push_constants), 3); - sg_create({GGML_TYPE_F16, GGML_TYPE_F16, false, false}, tc_mm, "matmul_f16", matmul_f16_fp32_len, matmul_f16_fp32_data, sizeof(vk_mat_mat_push_constants), 3); - sg_create({GGML_TYPE_F16, GGML_TYPE_F32, false, false}, tc_mm, "matmul_f16_f32", matmul_f16_f32_fp32_len, matmul_f16_f32_fp32_data, sizeof(vk_mat_mat_push_constants), 3); - sg_create({GGML_TYPE_BF16, GGML_TYPE_BF16, false, false}, tc_mm, "matmul_bf16", matmul_bf16_fp32_len, matmul_bf16_fp32_data, sizeof(vk_mat_mat_push_constants), 3); - - for (const auto type : non_lut_quant_types) { - sg_create_quant({type, GGML_TYPE_F32, false, false}, tc_mmq, "matmul_quant_f32", matmul_quant_f32_fp32_len, matmul_quant_f32_fp32_data, sizeof(vk_mat_mat_push_constants), 3); - } - #define X_SG_FP32(TYPE, tstr) \ - sg_create({TYPE, GGML_TYPE_F32, false, false}, tc_mmq, "matmul_" #tstr "_f32", matmul_##tstr##_f32_fp32_len, matmul_##tstr##_f32_fp32_data, sizeof(vk_mat_mat_push_constants), 3); - FOR_EACH_LUT_TYPE(X_SG_FP32) -#undef X_SG_FP32 + CREATE_MM(GGML_TYPE_F32, pipeline_matmul_id_f32, matmul_id_f32_f32, , wg_denoms, warptile, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM2(GGML_TYPE_F16, pipeline_matmul_id_f16, matmul_id_f16, wg_denoms, warptile, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM2(GGML_TYPE_F16, pipeline_matmul_id_f16_f32, matmul_id_f16_f32, wg_denoms, warptile, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM_NODOT2(GGML_TYPE_BF16, pipeline_matmul_id_bf16, matmul_id_bf16, , wg_denoms, warptile, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM2(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q1_0], matmul_id_q1_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM2(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_0], matmul_id_q2_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM2(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0], matmul_id_q4_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM2(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1], matmul_id_q4_1_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM2(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0], matmul_id_q5_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM2(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1], matmul_id_q5_1_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM2(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0], matmul_id_q8_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM2(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K], matmul_id_q2_k_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM2(GGML_TYPE_TQ2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ2_0], matmul_id_tq2_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM2(GGML_TYPE_TQ1_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ1_0], matmul_id_tq1_0_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM2(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K], matmul_id_q3_k_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM2(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K], matmul_id_q4_k_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM2(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K], matmul_id_q5_k_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM2(GGML_TYPE_Q6_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q6_K], matmul_id_q6_k_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM2(GGML_TYPE_IQ1_S, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ1_S], matmul_id_iq1_s_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM2(GGML_TYPE_IQ1_M, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ1_M], matmul_id_iq1_m_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM2(GGML_TYPE_IQ2_XXS, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ2_XXS], matmul_id_iq2_xxs_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM2(GGML_TYPE_IQ2_XS, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ2_XS], matmul_id_iq2_xs_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM2(GGML_TYPE_IQ2_S, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ2_S], matmul_id_iq2_s_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM2(GGML_TYPE_IQ3_XXS, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ3_XXS], matmul_id_iq3_xxs_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM2(GGML_TYPE_IQ3_S, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ3_S], matmul_id_iq3_s_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM2(GGML_TYPE_IQ4_XS, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_XS], matmul_id_iq4_xs_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM2(GGML_TYPE_IQ4_NL, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_NL], matmul_id_iq4_nl_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM2(GGML_TYPE_Q4_0_ROCMFP4, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0_ROCMFP4], matmul_id_rocmfp4_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM2(GGML_TYPE_Q4_0_ROCMFP4_FAST, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0_ROCMFP4_FAST], matmul_id_rocmfp4_fast_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM2(GGML_TYPE_Q3_0_ROCMFPX, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_0_ROCMFPX], matmul_id_rocmfpx_fp3_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM2(GGML_TYPE_Q6_0_ROCMFPX, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q6_0_ROCMFPX], matmul_id_rocmfpx_fp6_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM2(GGML_TYPE_Q8_0_ROCMFPX, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0_ROCMFPX], matmul_id_rocmfpx_fp8_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM2(GGML_TYPE_MXFP4, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_MXFP4], matmul_id_mxfp4_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM2(GGML_TYPE_NVFP4, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_NVFP4], matmul_id_nvfp4_f32, mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); #if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT) if (device->integer_dot_product) { - std::vector tc_mmq_int = {{s_warptile_mmq_int, s_mmq_wg_denoms, s_align}, {m_warptile_mmq_int, m_mmq_wg_denoms, m_align}, {l_warptile_mmq_int, l_mmq_wg_denoms, l_align}}; - std::vector tc_mmq_int_k = {{s_warptile_mmq_int_k, s_mmq_wg_denoms, s_align}, {m_warptile_mmq_int_k, m_mmq_wg_denoms, m_align}, {l_warptile_mmq_int_k, l_mmq_wg_denoms, l_align}}; - sg_create_mmq({GGML_TYPE_Q2_0, GGML_TYPE_Q8_1, false, false}, tc_mmq_int, "matmul_q2_0_q8_1", matmul_q2_0_q8_1_fp32_len, matmul_q2_0_q8_1_fp32_data, sizeof(vk_mat_mat_push_constants), 3); - sg_create_mmq({GGML_TYPE_Q4_0, GGML_TYPE_Q8_1, false, false}, tc_mmq_int, "matmul_q4_0_q8_1", matmul_q4_0_q8_1_fp32_len, matmul_q4_0_q8_1_fp32_data, sizeof(vk_mat_mat_push_constants), 3); - sg_create_mmq({GGML_TYPE_Q4_1, GGML_TYPE_Q8_1, false, false}, tc_mmq_int, "matmul_q4_1_q8_1", matmul_q4_1_q8_1_fp32_len, matmul_q4_1_q8_1_fp32_data, sizeof(vk_mat_mat_push_constants), 3); - sg_create_mmq({GGML_TYPE_Q5_0, GGML_TYPE_Q8_1, false, false}, tc_mmq_int, "matmul_q5_0_q8_1", matmul_q5_0_q8_1_fp32_len, matmul_q5_0_q8_1_fp32_data, sizeof(vk_mat_mat_push_constants), 3); - sg_create_mmq({GGML_TYPE_Q5_1, GGML_TYPE_Q8_1, false, false}, tc_mmq_int, "matmul_q5_1_q8_1", matmul_q5_1_q8_1_fp32_len, matmul_q5_1_q8_1_fp32_data, sizeof(vk_mat_mat_push_constants), 3); - sg_create_mmq({GGML_TYPE_Q8_0, GGML_TYPE_Q8_1, false, false}, tc_mmq_int, "matmul_q8_0_q8_1", matmul_q8_0_q8_1_fp32_len, matmul_q8_0_q8_1_fp32_data, sizeof(vk_mat_mat_push_constants), 3); - sg_create_mmq({GGML_TYPE_Q2_K, GGML_TYPE_Q8_1, false, false}, tc_mmq_int_k, "matmul_q2_k_q8_1", matmul_q2_k_q8_1_fp32_len, matmul_q2_k_q8_1_fp32_data, sizeof(vk_mat_mat_push_constants), 3); - sg_create_mmq({GGML_TYPE_Q3_K, GGML_TYPE_Q8_1, false, false}, tc_mmq_int_k, "matmul_q3_k_q8_1", matmul_q3_k_q8_1_fp32_len, matmul_q3_k_q8_1_fp32_data, sizeof(vk_mat_mat_push_constants), 3); - sg_create_mmq({GGML_TYPE_Q4_K, GGML_TYPE_Q8_1, false, false}, tc_mmq_int_k, "matmul_q4_k_q8_1", matmul_q4_k_q8_1_fp32_len, matmul_q4_k_q8_1_fp32_data, sizeof(vk_mat_mat_push_constants), 3); - sg_create_mmq({GGML_TYPE_Q5_K, GGML_TYPE_Q8_1, false, false}, tc_mmq_int_k, "matmul_q5_k_q8_1", matmul_q5_k_q8_1_fp32_len, matmul_q5_k_q8_1_fp32_data, sizeof(vk_mat_mat_push_constants), 3); - sg_create_mmq({GGML_TYPE_Q6_K, GGML_TYPE_Q8_1, false, false}, tc_mmq_int_k, "matmul_q6_k_q8_1", matmul_q6_k_q8_1_fp32_len, matmul_q6_k_q8_1_fp32_data, sizeof(vk_mat_mat_push_constants), 3); + CREATE_MMQ(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q2_0], matmul_id_q2_0_q8_1, mmq_wg_denoms, warptile_mmqid_int, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MMQ(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q4_0], matmul_id_q4_0_q8_1, mmq_wg_denoms, warptile_mmqid_int, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MMQ(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q4_1], matmul_id_q4_1_q8_1, mmq_wg_denoms, warptile_mmqid_int, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MMQ(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q5_0], matmul_id_q5_0_q8_1, mmq_wg_denoms, warptile_mmqid_int, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MMQ(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q5_1], matmul_id_q5_1_q8_1, mmq_wg_denoms, warptile_mmqid_int, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MMQ(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q8_0], matmul_id_q8_0_q8_1, mmq_wg_denoms, warptile_mmqid_int, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + + CREATE_MMQ(GGML_TYPE_MXFP4, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_MXFP4], matmul_id_mxfp4_q8_1, mmq_wg_denoms, warptile_mmqid_int, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MMQ(GGML_TYPE_Q4_0_ROCMFP4, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q4_0_ROCMFP4], matmul_id_rocmfp4_q8_1, mmq_wg_denoms, warptile_mmqid_int, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MMQ(GGML_TYPE_Q4_0_ROCMFP4_FAST, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q4_0_ROCMFP4_FAST], matmul_id_rocmfp4_fast_q8_1, mmq_wg_denoms, warptile_mmqid_int, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + + CREATE_MMQ(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q2_K], matmul_id_q2_k_q8_1, mmq_wg_denoms, warptile_mmqid_int_k, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MMQ(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q3_K], matmul_id_q3_k_q8_1, mmq_wg_denoms, warptile_mmqid_int_k, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MMQ(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q4_K], matmul_id_q4_k_q8_1, mmq_wg_denoms, warptile_mmqid_int_k, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MMQ(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q5_K], matmul_id_q5_k_q8_1, mmq_wg_denoms, warptile_mmqid_int_k, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MMQ(GGML_TYPE_Q6_K, pipeline_dequant_mul_mat_mat_id_q8_1[GGML_TYPE_Q6_K], matmul_id_q6_k_q8_1, mmq_wg_denoms, warptile_mmqid_int_k, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); } #endif + } +#undef CREATE_MM2 +#undef CREATE_MMQ +#undef CREATE_MM +#undef CREATE_MM_NODOT2 + } else { + // Create 6 variants, {s,m,l}x{unaligned,aligned} +#define CREATE_MM(TYPE, PIPELINE_NAME, NAMELC, F16ACC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID, REQSUBGROUPSIZE) \ + if (device->mul_mat ## ID ## _l[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->l, #NAMELC #F16ACC "_l", NAMELC ## F16ACC ## _fp32_len, NAMELC ## F16ACC ## _fp32_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, ggml_vk_mul_mm_spec(l_ ## WARPTILE, false), 1, false, REQSUBGROUPSIZE > 0, REQSUBGROUPSIZE); \ + if (device->mul_mat ## ID ## _m[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->m, #NAMELC #F16ACC "_m", NAMELC ## F16ACC ## _fp32_len, NAMELC ## F16ACC ## _fp32_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, ggml_vk_mul_mm_spec(m_ ## WARPTILE, false), 1, false, REQSUBGROUPSIZE > 0, REQSUBGROUPSIZE); \ + if (device->mul_mat ## ID ## _s[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->s, #NAMELC #F16ACC "_s", NAMELC ## F16ACC ## _fp32_len, NAMELC ## F16ACC ## _fp32_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, ggml_vk_mul_mm_spec(s_ ## WARPTILE, false), 1, false, REQSUBGROUPSIZE > 0, REQSUBGROUPSIZE); \ + if (device->mul_mat ## ID ## _l[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_l, #NAMELC #F16ACC "_aligned_l", NAMELC ## F16ACC ## _fp32_len, NAMELC ## F16ACC ## _fp32_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, ggml_vk_mul_mm_spec(l_ ## WARPTILE, true), l_align, false, REQSUBGROUPSIZE > 0, REQSUBGROUPSIZE); \ + if (device->mul_mat ## ID ## _m[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_m, #NAMELC #F16ACC "_aligned_m", NAMELC ## F16ACC ## _fp32_len, NAMELC ## F16ACC ## _fp32_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, ggml_vk_mul_mm_spec(m_ ## WARPTILE, true), m_align, false, REQSUBGROUPSIZE > 0, REQSUBGROUPSIZE); \ + if (device->mul_mat ## ID ## _s[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->a_s, #NAMELC #F16ACC "_aligned_s", NAMELC ## F16ACC ## _fp32_len, NAMELC ## F16ACC ## _fp32_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, ggml_vk_mul_mm_spec(s_ ## WARPTILE, true), s_align, false, REQSUBGROUPSIZE > 0, REQSUBGROUPSIZE); \ + +#define CREATE_MMQ(TYPE, PIPELINE_NAME, NAMELC, WG_DENOMS, WARPTILE, PUSHCONST, PARAMCOUNT, ID) \ + if (device->mul_mat ## ID ## _l_int[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->l, #NAMELC "_l", NAMELC ## _fp32_len, NAMELC ## _fp32_data, "main", PARAMCOUNT, sizeof(PUSHCONST), l_ ## WG_DENOMS, l_ ## WARPTILE, 1); \ + if (device->mul_mat ## ID ## _m_int[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->m, #NAMELC "_m", NAMELC ## _fp32_len, NAMELC ## _fp32_data, "main", PARAMCOUNT, sizeof(PUSHCONST), m_ ## WG_DENOMS, m_ ## WARPTILE, 1); \ + if (device->mul_mat ## ID ## _s_int[TYPE]) \ + ggml_vk_create_pipeline(device, device-> PIPELINE_NAME ->s, #NAMELC "_s", NAMELC ## _fp32_len, NAMELC ## _fp32_data, "main", PARAMCOUNT, sizeof(PUSHCONST), s_ ## WG_DENOMS, s_ ## WARPTILE, 1); \ + + CREATE_MM(GGML_TYPE_F32, pipeline_matmul_f32, matmul_f32_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_F32, pipeline_matmul_f32_f16, matmul_f32_f16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_F16, pipeline_matmul_f16.f32acc, matmul_f16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_F16, pipeline_matmul_f16_f32.f32acc, matmul_f16_f32, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, , 0); + + CREATE_MM(GGML_TYPE_BF16, pipeline_matmul_bf16, matmul_bf16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, , 0); + + CREATE_MM(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q1_0].f32acc, matmul_q1_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_0].f32acc, matmul_q2_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_0].f32acc, matmul_q4_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_1].f32acc, matmul_q4_1_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_0].f32acc, matmul_q5_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_1].f32acc, matmul_q5_1_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q8_0].f32acc, matmul_q8_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + + CREATE_MM(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q2_K].f32acc, matmul_q2_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_TQ2_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_TQ2_0].f32acc, matmul_tq2_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_TQ1_0, pipeline_dequant_mul_mat_mat[GGML_TYPE_TQ1_0].f32acc, matmul_tq1_0_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_K].f32acc, matmul_q3_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_K].f32acc, matmul_q4_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q5_K].f32acc, matmul_q5_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_Q6_K, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q6_K].f32acc, matmul_q6_k_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_IQ1_S, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ1_S].f32acc, matmul_iq1_s_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_IQ1_M, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ1_M].f32acc, matmul_iq1_m_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_IQ2_XXS, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ2_XXS].f32acc, matmul_iq2_xxs_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_IQ2_XS, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ2_XS].f32acc, matmul_iq2_xs_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_IQ2_S, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ2_S].f32acc, matmul_iq2_s_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_IQ3_XXS, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ3_XXS].f32acc, matmul_iq3_xxs_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_IQ3_S, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ3_S].f32acc, matmul_iq3_s_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_IQ4_XS, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ4_XS].f32acc, matmul_iq4_xs_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_IQ4_NL, pipeline_dequant_mul_mat_mat[GGML_TYPE_IQ4_NL].f32acc, matmul_iq4_nl_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_Q4_0_ROCMFP4, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_0_ROCMFP4].f32acc, matmul_rocmfp4_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_Q4_0_ROCMFP4_FAST, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q4_0_ROCMFP4_FAST].f32acc, matmul_rocmfp4_fast_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_Q3_0_ROCMFPX, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q3_0_ROCMFPX].f32acc, matmul_rocmfpx_fp3_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_Q6_0_ROCMFPX, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q6_0_ROCMFPX].f32acc, matmul_rocmfpx_fp6_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_Q8_0_ROCMFPX, pipeline_dequant_mul_mat_mat[GGML_TYPE_Q8_0_ROCMFPX].f32acc, matmul_rocmfpx_fp8_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_MXFP4, pipeline_dequant_mul_mat_mat[GGML_TYPE_MXFP4].f32acc, matmul_mxfp4_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_NVFP4, pipeline_dequant_mul_mat_mat[GGML_TYPE_NVFP4].f32acc, matmul_nvfp4_f32, , mmq_wg_denoms, warptile_mmq, vk_mat_mat_push_constants, 3, , 0); - if (device->subgroup_ballot && device->subgroup_require_full_support && subgroup_min_size_16) { - sg_create({GGML_TYPE_F32, GGML_TYPE_F32, true, false}, tc_id, "matmul_id_subgroup_f32_f32", matmul_id_subgroup_f32_f32_fp32_len, matmul_id_subgroup_f32_f32_fp32_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count, mul_mat_subgroup_size_16); - sg_create({GGML_TYPE_F16, GGML_TYPE_F16, true, false}, tc_id, "matmul_id_subgroup_f16", matmul_id_subgroup_f16_fp32_len, matmul_id_subgroup_f16_fp32_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count, mul_mat_subgroup_size_16); - sg_create({GGML_TYPE_F16, GGML_TYPE_F32, true, false}, tc_id, "matmul_id_subgroup_f16_f32", matmul_id_subgroup_f16_f32_fp32_len, matmul_id_subgroup_f16_f32_fp32_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count, mul_mat_subgroup_size_16); - sg_create({GGML_TYPE_BF16, GGML_TYPE_BF16, true, false}, tc_id, "matmul_id_subgroup_bf16", matmul_id_subgroup_bf16_fp32_len, matmul_id_subgroup_bf16_fp32_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count, mul_mat_subgroup_size_16); - for (const auto type : non_lut_quant_types) { - sg_create_quant({type, GGML_TYPE_F32, true, false}, tc_mmqid, "matmul_id_subgroup_quant_f32", matmul_id_subgroup_quant_f32_fp32_len, matmul_id_subgroup_quant_f32_fp32_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count, mul_mat_subgroup_size); - } - #define X_SG_ID_SUB_FP32(TYPE, tstr) \ - sg_create({TYPE, GGML_TYPE_F32, true, false}, tc_mmqid, "matmul_id_subgroup_" #tstr "_f32", matmul_id_subgroup_##tstr##_f32_fp32_len, matmul_id_subgroup_##tstr##_f32_fp32_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count, mul_mat_subgroup_size); - FOR_EACH_LUT_TYPE(X_SG_ID_SUB_FP32) -#undef X_SG_ID_SUB_FP32 - } else { - sg_create({GGML_TYPE_F32, GGML_TYPE_F32, true, false}, tc_mm, "matmul_id_f32_f32", matmul_id_f32_f32_fp32_len, matmul_id_f32_f32_fp32_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - sg_create({GGML_TYPE_F16, GGML_TYPE_F16, true, false}, tc_mm, "matmul_id_f16", matmul_id_f16_fp32_len, matmul_id_f16_fp32_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - sg_create({GGML_TYPE_F16, GGML_TYPE_F32, true, false}, tc_mm, "matmul_id_f16_f32", matmul_id_f16_f32_fp32_len, matmul_id_f16_f32_fp32_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - sg_create({GGML_TYPE_BF16, GGML_TYPE_BF16, true, false}, tc_mm, "matmul_id_bf16", matmul_id_bf16_fp32_len, matmul_id_bf16_fp32_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - for (const auto type : non_lut_quant_types) { - sg_create_quant({type, GGML_TYPE_F32, true, false}, tc_mmqid, "matmul_id_quant_f32", matmul_id_quant_f32_fp32_len, matmul_id_quant_f32_fp32_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - } - #define X_SG_ID_FP32(TYPE, tstr) \ - sg_create({TYPE, GGML_TYPE_F32, true, false}, tc_mmqid, "matmul_id_" #tstr "_f32", matmul_id_##tstr##_f32_fp32_len, matmul_id_##tstr##_f32_fp32_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count); - FOR_EACH_LUT_TYPE(X_SG_ID_FP32) -#undef X_SG_ID_FP32 - } +#if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT) + if (device->integer_dot_product) { + CREATE_MMQ(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q2_0].f32acc, matmul_q2_0_q8_1, mmq_wg_denoms, warptile_mmq_int, vk_mat_mat_push_constants, 3, ); + CREATE_MMQ(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q4_0].f32acc, matmul_q4_0_q8_1, mmq_wg_denoms, warptile_mmq_int, vk_mat_mat_push_constants, 3, ); + CREATE_MMQ(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q4_1].f32acc, matmul_q4_1_q8_1, mmq_wg_denoms, warptile_mmq_int, vk_mat_mat_push_constants, 3, ); + CREATE_MMQ(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q5_0].f32acc, matmul_q5_0_q8_1, mmq_wg_denoms, warptile_mmq_int, vk_mat_mat_push_constants, 3, ); + CREATE_MMQ(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q5_1].f32acc, matmul_q5_1_q8_1, mmq_wg_denoms, warptile_mmq_int, vk_mat_mat_push_constants, 3, ); + CREATE_MMQ(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q8_0].f32acc, matmul_q8_0_q8_1, mmq_wg_denoms, warptile_mmq_int, vk_mat_mat_push_constants, 3, ); + + CREATE_MMQ(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q2_K].f32acc, matmul_q2_k_q8_1, mmq_wg_denoms, warptile_mmq_int_k, vk_mat_mat_push_constants, 3, ); + CREATE_MMQ(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q3_K].f32acc, matmul_q3_k_q8_1, mmq_wg_denoms, warptile_mmq_int_k, vk_mat_mat_push_constants, 3, ); + CREATE_MMQ(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q4_K].f32acc, matmul_q4_k_q8_1, mmq_wg_denoms, warptile_mmq_int_k, vk_mat_mat_push_constants, 3, ); + CREATE_MMQ(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q5_K].f32acc, matmul_q5_k_q8_1, mmq_wg_denoms, warptile_mmq_int_k, vk_mat_mat_push_constants, 3, ); + CREATE_MMQ(GGML_TYPE_Q6_K, pipeline_dequant_mul_mat_mat_q8_1[GGML_TYPE_Q6_K].f32acc, matmul_q6_k_q8_1, mmq_wg_denoms, warptile_mmq_int_k, vk_mat_mat_push_constants, 3, ); } - } -#undef FOR_EACH_LUT_TYPE -#undef FOR_EACH_LUT_TYPE_NONFP4 -#undef FOR_EACH_LUT_FP4_TYPE - // BF16 fallback for coopmat devices without bf16 coopmat support +#endif + + if (device->subgroup_ballot && device->subgroup_require_full_support && subgroup_min_size_16) { + CREATE_MM(GGML_TYPE_F32, pipeline_matmul_id_f32, matmul_id_subgroup_f32_f32, , wg_denoms, warptile_id, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size_16); + CREATE_MM(GGML_TYPE_F16, pipeline_matmul_id_f16.f32acc, matmul_id_subgroup_f16, , wg_denoms, warptile_id, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size_16); + CREATE_MM(GGML_TYPE_F16, pipeline_matmul_id_f16_f32.f32acc, matmul_id_subgroup_f16_f32, , wg_denoms, warptile_id, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size_16); + CREATE_MM(GGML_TYPE_BF16, pipeline_matmul_id_bf16, matmul_id_subgroup_bf16, , wg_denoms, warptile_id, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size_16); + + CREATE_MM(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q1_0].f32acc, matmul_id_subgroup_q1_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_0].f32acc, matmul_id_subgroup_q2_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0].f32acc, matmul_id_subgroup_q4_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1].f32acc, matmul_id_subgroup_q4_1_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0].f32acc, matmul_id_subgroup_q5_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1].f32acc, matmul_id_subgroup_q5_1_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0].f32acc, matmul_id_subgroup_q8_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K].f32acc, matmul_id_subgroup_q2_k_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM(GGML_TYPE_TQ2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ2_0].f32acc, matmul_id_subgroup_tq2_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM(GGML_TYPE_TQ1_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ1_0].f32acc, matmul_id_subgroup_tq1_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K].f32acc, matmul_id_subgroup_q3_k_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K].f32acc, matmul_id_subgroup_q4_k_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K].f32acc, matmul_id_subgroup_q5_k_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM(GGML_TYPE_Q6_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q6_K].f32acc, matmul_id_subgroup_q6_k_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM(GGML_TYPE_IQ1_S, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ1_S].f32acc, matmul_id_subgroup_iq1_s_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM(GGML_TYPE_IQ1_M, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ1_M].f32acc, matmul_id_subgroup_iq1_m_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM(GGML_TYPE_IQ2_XXS, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ2_XXS].f32acc, matmul_id_subgroup_iq2_xxs_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM(GGML_TYPE_IQ2_XS, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ2_XS].f32acc, matmul_id_subgroup_iq2_xs_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM(GGML_TYPE_IQ2_S, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ2_S].f32acc, matmul_id_subgroup_iq2_s_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM(GGML_TYPE_IQ3_XXS, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ3_XXS].f32acc, matmul_id_subgroup_iq3_xxs_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM(GGML_TYPE_IQ3_S, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ3_S].f32acc, matmul_id_subgroup_iq3_s_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM(GGML_TYPE_IQ4_XS, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_XS].f32acc, matmul_id_subgroup_iq4_xs_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM(GGML_TYPE_IQ4_NL, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_NL].f32acc, matmul_id_subgroup_iq4_nl_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM(GGML_TYPE_Q4_0_ROCMFP4, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0_ROCMFP4].f32acc, matmul_id_subgroup_rocmfp4_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM(GGML_TYPE_Q4_0_ROCMFP4_FAST, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0_ROCMFP4_FAST].f32acc, matmul_id_subgroup_rocmfp4_fast_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM(GGML_TYPE_Q3_0_ROCMFPX, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_0_ROCMFPX].f32acc, matmul_id_subgroup_rocmfpx_fp3_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM(GGML_TYPE_Q6_0_ROCMFPX, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q6_0_ROCMFPX].f32acc, matmul_id_subgroup_rocmfpx_fp6_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM(GGML_TYPE_Q8_0_ROCMFPX, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0_ROCMFPX].f32acc, matmul_id_subgroup_rocmfpx_fp8_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM(GGML_TYPE_MXFP4, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_MXFP4].f32acc, matmul_id_subgroup_mxfp4_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + CREATE_MM(GGML_TYPE_NVFP4, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_NVFP4].f32acc, matmul_id_subgroup_nvfp4_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, mul_mat_subgroup_size); + } else { + CREATE_MM(GGML_TYPE_F32, pipeline_matmul_id_f32, matmul_id_f32_f32, , wg_denoms, warptile, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM(GGML_TYPE_F16, pipeline_matmul_id_f16.f32acc, matmul_id_f16, , wg_denoms, warptile, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM(GGML_TYPE_F16, pipeline_matmul_id_f16_f32.f32acc, matmul_id_f16_f32, , wg_denoms, warptile, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM(GGML_TYPE_BF16, pipeline_matmul_id_bf16, matmul_id_bf16, , wg_denoms, warptile, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + + CREATE_MM(GGML_TYPE_Q1_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q1_0].f32acc, matmul_id_q1_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM(GGML_TYPE_Q2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_0].f32acc, matmul_id_q2_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM(GGML_TYPE_Q4_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0].f32acc, matmul_id_q4_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM(GGML_TYPE_Q4_1, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_1].f32acc, matmul_id_q4_1_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM(GGML_TYPE_Q5_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_0].f32acc, matmul_id_q5_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM(GGML_TYPE_Q5_1, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_1].f32acc, matmul_id_q5_1_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM(GGML_TYPE_Q8_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0].f32acc, matmul_id_q8_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM(GGML_TYPE_Q2_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q2_K].f32acc, matmul_id_q2_k_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM(GGML_TYPE_TQ2_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ2_0].f32acc, matmul_id_tq2_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM(GGML_TYPE_TQ1_0, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_TQ1_0].f32acc, matmul_id_tq1_0_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM(GGML_TYPE_Q3_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_K].f32acc, matmul_id_q3_k_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM(GGML_TYPE_Q4_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_K].f32acc, matmul_id_q4_k_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM(GGML_TYPE_Q5_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q5_K].f32acc, matmul_id_q5_k_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM(GGML_TYPE_Q6_K, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q6_K].f32acc, matmul_id_q6_k_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM(GGML_TYPE_IQ1_S, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ1_S].f32acc, matmul_id_iq1_s_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM(GGML_TYPE_IQ1_M, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ1_M].f32acc, matmul_id_iq1_m_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM(GGML_TYPE_IQ2_XXS, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ2_XXS].f32acc, matmul_id_iq2_xxs_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM(GGML_TYPE_IQ2_XS, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ2_XS].f32acc, matmul_id_iq2_xs_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM(GGML_TYPE_IQ2_S, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ2_S].f32acc, matmul_id_iq2_s_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM(GGML_TYPE_IQ3_XXS, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ3_XXS].f32acc, matmul_id_iq3_xxs_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM(GGML_TYPE_IQ3_S, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ3_S].f32acc, matmul_id_iq3_s_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM(GGML_TYPE_IQ4_XS, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_XS].f32acc, matmul_id_iq4_xs_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM(GGML_TYPE_IQ4_NL, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_IQ4_NL].f32acc, matmul_id_iq4_nl_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM(GGML_TYPE_Q4_0_ROCMFP4, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0_ROCMFP4].f32acc, matmul_id_rocmfp4_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM(GGML_TYPE_Q4_0_ROCMFP4_FAST, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q4_0_ROCMFP4_FAST].f32acc, matmul_id_rocmfp4_fast_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM(GGML_TYPE_Q3_0_ROCMFPX, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q3_0_ROCMFPX].f32acc, matmul_id_rocmfpx_fp3_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM(GGML_TYPE_Q6_0_ROCMFPX, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q6_0_ROCMFPX].f32acc, matmul_id_rocmfpx_fp6_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM(GGML_TYPE_Q8_0_ROCMFPX, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_Q8_0_ROCMFPX].f32acc, matmul_id_rocmfpx_fp8_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM(GGML_TYPE_MXFP4, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_MXFP4].f32acc, matmul_id_mxfp4_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + CREATE_MM(GGML_TYPE_NVFP4, pipeline_dequant_mul_mat_mat_id[GGML_TYPE_NVFP4].f32acc, matmul_id_nvfp4_f32, , mmq_wg_denoms, warptile_mmqid, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); + } + } + // reusing CREATE_MM from the fp32 path if ((device->coopmat2 || device->coopmat_support) #if defined(GGML_VULKAN_BFLOAT16_GLSLC_SUPPORT) && !device->coopmat_bf16_support #endif ) { - const uint32_t s_warptile_wm_bf16 = device->subgroup_size == 8 ? 8 : 32; - std::vector tc_bf16_fb = { - {{ subgroup_size_32, 32, 32, 16, s_warptile_wm_bf16, 32, 2, 2, 2, 1, subgroup_size_8 }, {32, 32, 1}, s_align}, - {{ 128, 64, 64, 16, mm_warp_8, 32, 2, 4, 2, 1, mm_warp_8 }, {64, 64, 1}, m_align}, - {{ 128, 128, 128, 16, mm_warp_8 * 2, 64, 2, 4, 4, 1, mm_warp_8 }, {128, 128, 1}, l_align}, - }; - auto tc_bf16_filtered = filter_tc(tc_bf16_fb, GGML_TYPE_BF16, false); - auto tc_bf16_id_filtered = filter_tc(tc_bf16_fb, GGML_TYPE_BF16, true); - spec_fn_t bf16_spec = [&](const std::vector& wt, bool a) { return ggml_vk_mul_mm_spec(wt, a); }; - if (!tc_bf16_filtered.empty()) { - create_mm_pipelines({GGML_TYPE_BF16, GGML_TYPE_BF16, false, false}, tc_bf16_filtered, "matmul_bf16", matmul_bf16_fp32_len, matmul_bf16_fp32_data, sizeof(vk_mat_mat_push_constants), 3, bf16_spec); - } - if (!tc_bf16_id_filtered.empty()) { - create_mm_pipelines({GGML_TYPE_BF16, GGML_TYPE_BF16, true, false}, tc_bf16_id_filtered, "matmul_id_bf16", matmul_id_bf16_fp32_len, matmul_id_bf16_fp32_data, sizeof(vk_mat_mat_id_push_constants), mul_mat_id_param_count, bf16_spec); - } - } - - // Set up tile selector functions - if (device->coopmat2) { - device->matmul_tile_selector = [](uint32_t m, uint32_t n, uint32_t /*k*/, uint32_t shader_core_count, - const std::vector& configs) -> uint32_t { - if (configs.size() <= 1) return 0; - uint32_t last = (uint32_t)configs.size() - 1; - if (configs.size() == 2) { - uint32_t crossover = configs[0].unaligned->wg_denoms[1]; - return (n > crossover) ? 1 : 0; - } - // 3+ configs: s=0, m=1, l=2 - const uint32_t tiles_l = CEIL_DIV(m, configs[last].unaligned->wg_denoms[0]) * CEIL_DIV(n, configs[last].unaligned->wg_denoms[1]); - const uint32_t tiles_m = CEIL_DIV(m, configs[1].unaligned->wg_denoms[0]) * CEIL_DIV(n, configs[1].unaligned->wg_denoms[1]); - uint32_t crossover_large = configs[1].unaligned->wg_denoms[1]; - bool prefer_large = tiles_m > shader_core_count || tiles_l > shader_core_count || - (tiles_l <= shader_core_count / 3 && tiles_m > shader_core_count / 2); - if (n > crossover_large && prefer_large) return last; - uint32_t crossover_medium_m = configs[0].unaligned->wg_denoms[0]; - uint32_t crossover_medium_n = configs[0].unaligned->wg_denoms[1]; - if (m > crossover_medium_m && n > crossover_medium_n) return 1; - return 0; - }; - device->matmul_id_tile_selector = [](uint32_t /*m*/, uint32_t n, uint32_t /*k*/, uint32_t /*shader_core_count*/, - const std::vector& configs) -> uint32_t { - if (configs.size() <= 1) return 0; - uint32_t last = (uint32_t)configs.size() - 1; - if (configs.size() == 2) { - uint32_t crossover = configs[0].unaligned->wg_denoms[1]; - return (n > crossover) ? 1 : 0; - } - uint32_t crossover_large = configs[1].unaligned->wg_denoms[1]; - if (n > crossover_large) return last; - uint32_t crossover_medium = configs[0].unaligned->wg_denoms[1]; - if (n > crossover_medium) return 1; - return 0; - }; - } else { - device->matmul_tile_selector = [](uint32_t m, uint32_t n, uint32_t /*k*/, uint32_t /*shader_core_count*/, - const std::vector& configs) -> uint32_t { - if (configs.size() <= 1) return 0; - if (m <= 32 || n <= 32) return 0; - if (configs.size() == 2) return 1; - if (m <= 64 || n <= 64) return 1; - return (uint32_t)configs.size() - 1; - }; - device->matmul_id_tile_selector = device->matmul_tile_selector; + const uint32_t s_warptile_wm = device->subgroup_size == 8 ? 8 : 32; + + // use scalar tile sizes + l_warptile = { 128, 128, 128, 16, mm_warp_8 * 2, 64, 2, 4, 4, 1, mm_warp_8 }; + m_warptile = { 128, 64, 64, 16, mm_warp_8, 32, 2, 4, 2, 1, mm_warp_8 }; + s_warptile = { subgroup_size_32, 32, 32, 16, s_warptile_wm, 32, 2, 2, 2, 1, subgroup_size_8 }; + + l_wg_denoms = {128, 128, 1 }; + m_wg_denoms = { 64, 64, 1 }; + s_wg_denoms = { 32, 32, 1 }; + + CREATE_MM(GGML_TYPE_BF16, pipeline_matmul_bf16, matmul_bf16, , wg_denoms, warptile, vk_mat_mat_push_constants, 3, , 0); + CREATE_MM(GGML_TYPE_BF16, pipeline_matmul_id_bf16, matmul_id_bf16, , wg_denoms, warptile, vk_mat_mat_id_push_constants, mul_mat_id_param_count, _id, 0); } +#undef CREATE_MM // mul mat vec @@ -5559,12 +6015,12 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q5_1][i], "mul_mat_vec_q5_1_f32_f32", arr_dmmv_q5_1_f32_f32_len[reduc], arr_dmmv_q5_1_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q8_0][i], "mul_mat_vec_q8_0_f32_f32", arr_dmmv_q8_0_f32_f32_len[reduc], arr_dmmv_q8_0_f32_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {1*rm_stdq, 1, 1}, {wg_size_subgroup, 1*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q2_K][i], "mul_mat_vec_q2_k_f32_f32", arr_dmmv_q2_k_f32_f32_len[reduc16], arr_dmmv_q2_k_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_TQ2_0][i], "mul_mat_vec_tq2_0_f32_f32", arr_dmmv_tq2_0_f32_f32_len[reduc16], arr_dmmv_tq2_0_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_TQ1_0][i], "mul_mat_vec_tq1_0_f32_f32", arr_dmmv_tq1_0_f32_f32_len[reduc16], arr_dmmv_tq1_0_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q3_K][i], "mul_mat_vec_q3_k_f32_f32", arr_dmmv_q3_k_f32_f32_len[reduc16], arr_dmmv_q3_k_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q4_K][i], "mul_mat_vec_q4_k_f32_f32", arr_dmmv_q4_k_f32_f32_len[reduc16], arr_dmmv_q4_k_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q5_K][i], "mul_mat_vec_q5_k_f32_f32", arr_dmmv_q5_k_f32_f32_len[reduc16], arr_dmmv_q5_k_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q6_K][i], "mul_mat_vec_q6_k_f32_f32", arr_dmmv_q6_k_f32_f32_len[reduc16], arr_dmmv_q6_k_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); - ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_TQ1_0][i], "mul_mat_vec_tq1_0_f32_f32", arr_dmmv_tq1_0_f32_f32_len[reduc16], arr_dmmv_tq1_0_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); - ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_TQ2_0][i], "mul_mat_vec_tq2_0_f32_f32", arr_dmmv_tq2_0_f32_f32_len[reduc16], arr_dmmv_tq2_0_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_IQ1_S][i], "mul_mat_vec_iq1_s_f32_f32", arr_dmmv_iq1_s_f32_f32_len[reduc16], arr_dmmv_iq1_s_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_IQ1_M][i], "mul_mat_vec_iq1_m_f32_f32", arr_dmmv_iq1_m_f32_f32_len[reduc16], arr_dmmv_iq1_m_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_IQ2_XXS][i], "mul_mat_vec_iq2_xxs_f32_f32", arr_dmmv_iq2_xxs_f32_f32_len[reduc16], arr_dmmv_iq2_xxs_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); @@ -5574,6 +6030,11 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_IQ3_S][i], "mul_mat_vec_iq3_s_f32_f32", arr_dmmv_iq3_s_f32_f32_len[reduc16], arr_dmmv_iq3_s_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_IQ4_XS][i], "mul_mat_vec_iq4_xs_f32_f32", arr_dmmv_iq4_xs_f32_f32_len[reduc16], arr_dmmv_iq4_xs_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_IQ4_NL][i], "mul_mat_vec_iq4_nl_f32_f32", arr_dmmv_iq4_nl_f32_f32_len[reduc16], arr_dmmv_iq4_nl_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q4_0_ROCMFP4][i], "mul_mat_vec_rocmfp4_f32_f32", arr_dmmv_rocmfp4_f32_f32_len[reduc16], arr_dmmv_rocmfp4_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q4_0_ROCMFP4_FAST][i], "mul_mat_vec_rocmfp4_fast_f32_f32", arr_dmmv_rocmfp4_fast_f32_f32_len[reduc16], arr_dmmv_rocmfp4_fast_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q3_0_ROCMFPX][i], "mul_mat_vec_rocmfpx_fp3_f32_f32", arr_dmmv_rocmfpx_fp3_f32_f32_len[reduc16], arr_dmmv_rocmfpx_fp3_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q6_0_ROCMFPX][i], "mul_mat_vec_rocmfpx_fp6_f32_f32", arr_dmmv_rocmfpx_fp6_f32_f32_len[reduc16], arr_dmmv_rocmfpx_fp6_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_Q8_0_ROCMFPX][i], "mul_mat_vec_rocmfpx_fp8_f32_f32", arr_dmmv_rocmfpx_fp8_f32_f32_len[reduc16], arr_dmmv_rocmfpx_fp8_f32_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_MXFP4][i], "mul_mat_vec_mxfp4_f32_f32", OCP_DMMV_LEN(arr_dmmv_mxfp4_f32_f32, reduc16), OCP_DMMV_DATA(arr_dmmv_mxfp4_f32_f32, reduc16), "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f32_f32[w][GGML_TYPE_NVFP4][i], "mul_mat_vec_nvfp4_f32_f32", OCP_DMMV_LEN(arr_dmmv_nvfp4_f32_f32, reduc16), OCP_DMMV_DATA(arr_dmmv_nvfp4_f32_f32, reduc16), "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); @@ -5588,12 +6049,12 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q5_1][i], "mul_mat_vec_q5_1_f16_f32", arr_dmmv_q5_1_f16_f32_len[reduc], arr_dmmv_q5_1_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q8_0][i], "mul_mat_vec_q8_0_f16_f32", arr_dmmv_q8_0_f16_f32_len[reduc], arr_dmmv_q8_0_f16_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {1*rm_stdq, 1, 1}, {wg_size_subgroup, 1*rm_stdq, i+1}, 1, true, use_subgroups, force_subgroup_size); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q2_K][i], "mul_mat_vec_q2_k_f16_f32", arr_dmmv_q2_k_f16_f32_len[reduc16], arr_dmmv_q2_k_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_TQ2_0][i], "mul_mat_vec_tq2_0_f16_f32", arr_dmmv_tq2_0_f16_f32_len[reduc16], arr_dmmv_tq2_0_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_TQ1_0][i], "mul_mat_vec_tq1_0_f16_f32", arr_dmmv_tq1_0_f16_f32_len[reduc16], arr_dmmv_tq1_0_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q3_K][i], "mul_mat_vec_q3_k_f16_f32", arr_dmmv_q3_k_f16_f32_len[reduc16], arr_dmmv_q3_k_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q4_K][i], "mul_mat_vec_q4_k_f16_f32", arr_dmmv_q4_k_f16_f32_len[reduc16], arr_dmmv_q4_k_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q5_K][i], "mul_mat_vec_q5_k_f16_f32", arr_dmmv_q5_k_f16_f32_len[reduc16], arr_dmmv_q5_k_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q6_K][i], "mul_mat_vec_q6_k_f16_f32", arr_dmmv_q6_k_f16_f32_len[reduc16], arr_dmmv_q6_k_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); - ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_TQ1_0][i], "mul_mat_vec_tq1_0_f16_f32", arr_dmmv_tq1_0_f16_f32_len[reduc16], arr_dmmv_tq1_0_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); - ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_TQ2_0][i], "mul_mat_vec_tq2_0_f16_f32", arr_dmmv_tq2_0_f16_f32_len[reduc16], arr_dmmv_tq2_0_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_IQ1_S][i], "mul_mat_vec_iq1_s_f16_f32", arr_dmmv_iq1_s_f16_f32_len[reduc16], arr_dmmv_iq1_s_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_IQ1_M][i], "mul_mat_vec_iq1_m_f16_f32", arr_dmmv_iq1_m_f16_f32_len[reduc16], arr_dmmv_iq1_m_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_IQ2_XXS][i], "mul_mat_vec_iq2_xxs_f16_f32", arr_dmmv_iq2_xxs_f16_f32_len[reduc16], arr_dmmv_iq2_xxs_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); @@ -5603,6 +6064,11 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_IQ3_S][i], "mul_mat_vec_iq3_s_f16_f32", arr_dmmv_iq3_s_f16_f32_len[reduc16], arr_dmmv_iq3_s_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_IQ4_XS][i], "mul_mat_vec_iq4_xs_f16_f32", arr_dmmv_iq4_xs_f16_f32_len[reduc16], arr_dmmv_iq4_xs_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_IQ4_NL][i], "mul_mat_vec_iq4_nl_f16_f32", arr_dmmv_iq4_nl_f16_f32_len[reduc16], arr_dmmv_iq4_nl_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q4_0_ROCMFP4][i], "mul_mat_vec_rocmfp4_f16_f32", arr_dmmv_rocmfp4_f16_f32_len[reduc16], arr_dmmv_rocmfp4_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q4_0_ROCMFP4_FAST][i], "mul_mat_vec_rocmfp4_fast_f16_f32", arr_dmmv_rocmfp4_fast_f16_f32_len[reduc16], arr_dmmv_rocmfp4_fast_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q3_0_ROCMFPX][i], "mul_mat_vec_rocmfpx_fp3_f16_f32", arr_dmmv_rocmfpx_fp3_f16_f32_len[reduc16], arr_dmmv_rocmfpx_fp3_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q6_0_ROCMFPX][i], "mul_mat_vec_rocmfpx_fp6_f16_f32", arr_dmmv_rocmfpx_fp6_f16_f32_len[reduc16], arr_dmmv_rocmfpx_fp6_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_Q8_0_ROCMFPX][i], "mul_mat_vec_rocmfpx_fp8_f16_f32", arr_dmmv_rocmfpx_fp8_f16_f32_len[reduc16], arr_dmmv_rocmfpx_fp8_f16_f32_data[reduc16], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_MXFP4][i], "mul_mat_vec_mxfp4_f16_f32", OCP_DMMV_LEN(arr_dmmv_mxfp4_f16_f32, reduc16), OCP_DMMV_DATA(arr_dmmv_mxfp4_f16_f32, reduc16), "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_f16_f32[w][GGML_TYPE_NVFP4][i], "mul_mat_vec_nvfp4_f16_f32", OCP_DMMV_LEN(arr_dmmv_nvfp4_f16_f32, reduc16), OCP_DMMV_DATA(arr_dmmv_nvfp4_f16_f32, reduc16), "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq, i+1}, 1, true, use_subgroups16, force_subgroup_size16); @@ -5619,6 +6085,9 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q8_0][i], "mul_mat_vec_q8_0_q8_1_f32", arr_dmmv_q8_0_q8_1_f32_len[reduc], arr_dmmv_q8_0_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_int_n(1*rm_stdq_int, i), 1, 1}, {wg_size_subgroup_int, rm_int_n(1*rm_stdq_int, i), i+1}, 1, true, use_subgroups, subgroup_size_int); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_MXFP4][i], "mul_mat_vec_mxfp4_q8_1_f32", arr_dmmv_mxfp4_q8_1_f32_len[reduc], arr_dmmv_mxfp4_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_int_n(2*rm_stdq_int, i), 1, 1}, {wg_size_subgroup_int, rm_int_n(2*rm_stdq_int, i), i+1}, 1, true, use_subgroups, subgroup_size_int); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q4_0_ROCMFP4][i], "mul_mat_vec_rocmfp4_q8_1_f32", arr_dmmv_rocmfp4_q8_1_f32_len[reduc], arr_dmmv_rocmfp4_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_int_n(2*rm_stdq_int, i), 1, 1}, {wg_size_subgroup_int, rm_int_n(2*rm_stdq_int, i), i+1}, 1, true, use_subgroups, subgroup_size_int); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q4_0_ROCMFP4_FAST][i], "mul_mat_vec_rocmfp4_fast_q8_1_f32", arr_dmmv_rocmfp4_fast_q8_1_f32_len[reduc], arr_dmmv_rocmfp4_fast_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_int_n(2*rm_stdq_int, i), 1, 1}, {wg_size_subgroup_int, rm_int_n(2*rm_stdq_int, i), i+1}, 1, true, use_subgroups, subgroup_size_int); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q8_0_ROCMFPX][i], "mul_mat_vec_rocmfpx_fp8_q8_1_f32", arr_dmmv_rocmfpx_fp8_q8_1_f32_len[reduc], arr_dmmv_rocmfpx_fp8_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_int_n(2*rm_stdq_int, i), 1, 1}, {wg_size_subgroup_int, rm_int_n(2*rm_stdq_int, i), i+1}, 1, true, use_subgroups, subgroup_size_int); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q2_K][i], "mul_mat_vec_q2_k_q8_1_f32", arr_dmmv_q2_k_q8_1_f32_len[reduc], arr_dmmv_q2_k_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_int_n(2*rm_kq_int, i), 1, 1}, {wg_size_subgroup_int, rm_int_n(2*rm_kq_int, i), i+1}, 1, true, use_subgroups, subgroup_size_int); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_q8_1_f32[w][GGML_TYPE_Q3_K][i], "mul_mat_vec_q3_k_q8_1_f32", arr_dmmv_q3_k_q8_1_f32_len[reduc], arr_dmmv_q3_k_q8_1_f32_data[reduc], "main", mul_mat_vec_num_bindings, sizeof(vk_mat_vec_push_constants), {rm_int_n(1*rm_kq_int, i), 1, 1}, {wg_size_subgroup_int, rm_int_n(1*rm_kq_int, i), i+1}, 1, true, use_subgroups, subgroup_size_int); @@ -5644,12 +6113,12 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q5_1], "mul_mat_vec_id_q5_1_f32", arr_dmmv_id_q5_1_f32_f32_len[reduc], arr_dmmv_id_q5_1_f32_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {2*rm_stdq, 1, 1}, {wg_size_subgroup, 2*rm_stdq}, 1, true, use_subgroups, force_subgroup_size); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q8_0], "mul_mat_vec_id_q8_0_f32", arr_dmmv_id_q8_0_f32_f32_len[reduc], arr_dmmv_id_q8_0_f32_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {1*rm_stdq, 1, 1}, {wg_size_subgroup, 1*rm_stdq}, 1, true, use_subgroups, force_subgroup_size); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q2_K], "mul_mat_vec_id_q2_k_f32", arr_dmmv_id_q2_k_f32_f32_len[reduc16], arr_dmmv_id_q2_k_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq}, 1, true, use_subgroups16, force_subgroup_size16); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_TQ2_0], "mul_mat_vec_id_tq2_0_f32", arr_dmmv_id_tq2_0_f32_f32_len[reduc16], arr_dmmv_id_tq2_0_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq}, 1, true, use_subgroups16, force_subgroup_size16); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_TQ1_0], "mul_mat_vec_id_tq1_0_f32", arr_dmmv_id_tq1_0_f32_f32_len[reduc16], arr_dmmv_id_tq1_0_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q3_K], "mul_mat_vec_id_q3_k_f32", arr_dmmv_id_q3_k_f32_f32_len[reduc16], arr_dmmv_id_q3_k_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q4_K], "mul_mat_vec_id_q4_k_f32", arr_dmmv_id_q4_k_f32_f32_len[reduc16], arr_dmmv_id_q4_k_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q5_K], "mul_mat_vec_id_q5_k_f32", arr_dmmv_id_q5_k_f32_f32_len[reduc16], arr_dmmv_id_q5_k_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q6_K], "mul_mat_vec_id_q6_k_f32", arr_dmmv_id_q6_k_f32_f32_len[reduc16], arr_dmmv_id_q6_k_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq}, 1, true, use_subgroups16, force_subgroup_size16); - ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_TQ1_0], "mul_mat_vec_id_tq1_0_f32", arr_dmmv_id_tq1_0_f32_f32_len[reduc16], arr_dmmv_id_tq1_0_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq}, 1, true, use_subgroups16, force_subgroup_size16); - ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_TQ2_0], "mul_mat_vec_id_tq2_0_f32", arr_dmmv_id_tq2_0_f32_f32_len[reduc16], arr_dmmv_id_tq2_0_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_kq, 1, 1}, {wg_size_subgroup16, rm_kq}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_IQ1_S], "mul_mat_vec_id_iq1_s_f32", arr_dmmv_id_iq1_s_f32_f32_len[reduc16], arr_dmmv_id_iq1_s_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_IQ1_M], "mul_mat_vec_id_iq1_m_f32", arr_dmmv_id_iq1_m_f32_f32_len[reduc16], arr_dmmv_id_iq1_m_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_IQ2_XXS], "mul_mat_vec_id_iq2_xxs_f32", arr_dmmv_id_iq2_xxs_f32_f32_len[reduc16], arr_dmmv_id_iq2_xxs_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq}, 1, true, use_subgroups16, force_subgroup_size16); @@ -5659,6 +6128,11 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_IQ3_S], "mul_mat_vec_id_iq3_s_f32", arr_dmmv_id_iq3_s_f32_f32_len[reduc16], arr_dmmv_id_iq3_s_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_IQ4_XS], "mul_mat_vec_id_iq4_xs_f32", arr_dmmv_id_iq4_xs_f32_f32_len[reduc16], arr_dmmv_id_iq4_xs_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_IQ4_NL], "mul_mat_vec_id_iq4_nl_f32", arr_dmmv_id_iq4_nl_f32_f32_len[reduc16], arr_dmmv_id_iq4_nl_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq}, 1, true, use_subgroups16, force_subgroup_size16); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q4_0_ROCMFP4], "mul_mat_vec_id_rocmfp4_f32", arr_dmmv_id_rocmfp4_f32_f32_len[reduc16], arr_dmmv_id_rocmfp4_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq}, 1, true, use_subgroups16, force_subgroup_size16); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q4_0_ROCMFP4_FAST], "mul_mat_vec_id_rocmfp4_fast_f32", arr_dmmv_id_rocmfp4_fast_f32_f32_len[reduc16], arr_dmmv_id_rocmfp4_fast_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq}, 1, true, use_subgroups16, force_subgroup_size16); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q3_0_ROCMFPX], "mul_mat_vec_id_rocmfpx_fp3_f32", arr_dmmv_id_rocmfpx_fp3_f32_f32_len[reduc16], arr_dmmv_id_rocmfpx_fp3_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq}, 1, true, use_subgroups16, force_subgroup_size16); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q6_0_ROCMFPX], "mul_mat_vec_id_rocmfpx_fp6_f32", arr_dmmv_id_rocmfpx_fp6_f32_f32_len[reduc16], arr_dmmv_id_rocmfpx_fp6_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq}, 1, true, use_subgroups16, force_subgroup_size16); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_Q8_0_ROCMFPX], "mul_mat_vec_id_rocmfpx_fp8_f32", arr_dmmv_id_rocmfpx_fp8_f32_f32_len[reduc16], arr_dmmv_id_rocmfpx_fp8_f32_f32_data[reduc16], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_MXFP4], "mul_mat_vec_id_mxfp4_f32", OCP_DMMV_LEN(arr_dmmv_id_mxfp4_f32_f32, reduc16), OCP_DMMV_DATA(arr_dmmv_id_mxfp4_f32_f32, reduc16), "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq}, 1, true, use_subgroups16, force_subgroup_size16); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_f32[w][GGML_TYPE_NVFP4], "mul_mat_vec_id_nvfp4_f32", OCP_DMMV_LEN(arr_dmmv_id_nvfp4_f32_f32, reduc16), OCP_DMMV_DATA(arr_dmmv_id_nvfp4_f32_f32, reduc16), "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_iq, 1, 1}, {wg_size_subgroup16, rm_iq}, 1, true, use_subgroups16, force_subgroup_size16); @@ -5675,6 +6149,9 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q8_0], "mul_mat_vec_id_q8_0_q8_1_f32", arr_dmmv_id_q8_0_q8_1_f32_len[reduc], arr_dmmv_id_q8_0_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_id(1*rm_stdq_int), 1, 1}, {wg_size_subgroup_int, rm_id(1*rm_stdq_int)}, 1, true, use_subgroups, subgroup_size_int); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_MXFP4], "mul_mat_vec_id_mxfp4_q8_1_f32", arr_dmmv_id_mxfp4_q8_1_f32_len[reduc], arr_dmmv_id_mxfp4_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_id(2*rm_stdq_int), 1, 1}, {wg_size_subgroup_int, rm_id(2*rm_stdq_int)}, 1, true, use_subgroups, subgroup_size_int); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q4_0_ROCMFP4], "mul_mat_vec_id_rocmfp4_q8_1_f32", arr_dmmv_id_rocmfp4_q8_1_f32_len[reduc], arr_dmmv_id_rocmfp4_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_id(2*rm_stdq_int), 1, 1}, {wg_size_subgroup_int, rm_id(2*rm_stdq_int)}, 1, true, use_subgroups, subgroup_size_int); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q4_0_ROCMFP4_FAST], "mul_mat_vec_id_rocmfp4_fast_q8_1_f32", arr_dmmv_id_rocmfp4_fast_q8_1_f32_len[reduc], arr_dmmv_id_rocmfp4_fast_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_id(2*rm_stdq_int), 1, 1}, {wg_size_subgroup_int, rm_id(2*rm_stdq_int)}, 1, true, use_subgroups, subgroup_size_int); + ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q8_0_ROCMFPX], "mul_mat_vec_id_rocmfpx_fp8_q8_1_f32", arr_dmmv_id_rocmfpx_fp8_q8_1_f32_len[reduc], arr_dmmv_id_rocmfpx_fp8_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_id(2*rm_stdq_int), 1, 1}, {wg_size_subgroup_int, rm_id(2*rm_stdq_int)}, 1, true, use_subgroups, subgroup_size_int); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q2_K], "mul_mat_vec_id_q2_k_q8_1_f32", arr_dmmv_id_q2_k_q8_1_f32_len[reduc], arr_dmmv_id_q2_k_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_id(2*rm_kq_int), 1, 1}, {wg_size_subgroup_int, rm_id(2*rm_kq_int)}, 1, true, use_subgroups, subgroup_size_int); ggml_vk_create_pipeline(device, device->pipeline_dequant_mul_mat_vec_id_q8_1_f32[w][GGML_TYPE_Q3_K], "mul_mat_vec_id_q3_k_q8_1_f32", arr_dmmv_id_q3_k_q8_1_f32_len[reduc], arr_dmmv_id_q3_k_q8_1_f32_data[reduc], "main", mul_mat_vec_id_num_bindings, sizeof(vk_mat_vec_id_push_constants), {rm_id(1*rm_kq_int), 1, 1}, {wg_size_subgroup_int, rm_id(1*rm_kq_int)}, 1, true, use_subgroups, subgroup_size_int); @@ -5708,15 +6185,20 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q4_1], "dequant_q4_1", dequant_q4_1_len, dequant_q4_1_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q5_0], "dequant_q5_0", dequant_q5_0_len, dequant_q5_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q5_1], "dequant_q5_1", dequant_q5_1_len, dequant_q5_1_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant_transpose[GGML_TYPE_Q4_0], "dequant_q4_0_transpose", dequant_q4_0_transpose_len, dequant_q4_0_transpose_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant_transpose[GGML_TYPE_Q4_1], "dequant_q4_1_transpose", dequant_q4_1_transpose_len, dequant_q4_1_transpose_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant_transpose[GGML_TYPE_Q5_0], "dequant_q5_0_transpose", dequant_q5_0_transpose_len, dequant_q5_0_transpose_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant_transpose[GGML_TYPE_Q5_1], "dequant_q5_1_transpose", dequant_q5_1_transpose_len, dequant_q5_1_transpose_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q8_0], "dequant_q8_0", dequant_q8_0_len, dequant_q8_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant_transpose[GGML_TYPE_Q8_0], "dequant_q8_0_transpose", dequant_q8_0_transpose_len, dequant_q8_0_transpose_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant_transpose[GGML_TYPE_F16], "dequant_f16_transpose", dequant_f16_transpose_len, dequant_f16_transpose_data, "main", 2, 5 * sizeof(uint32_t), {256 * 8, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q2_K], "dequant_q2_k", dequant_q2_k_len, dequant_q2_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_TQ2_0], "dequant_tq2_0", dequant_tq2_0_len, dequant_tq2_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_TQ1_0], "dequant_tq1_0", dequant_tq1_0_len, dequant_tq1_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 4, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q3_K], "dequant_q3_k", dequant_q3_k_len, dequant_q3_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q4_K], "dequant_q4_k", dequant_q4_k_len, dequant_q4_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 32, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q5_K], "dequant_q5_k", dequant_q5_k_len, dequant_q5_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q6_K], "dequant_q6_k", dequant_q6_k_len, dequant_q6_k_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); - ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_TQ1_0], "dequant_tq1_0", dequant_tq1_0_len, dequant_tq1_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 4, 1, 1}, {}, 1); - ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_TQ2_0], "dequant_tq2_0", dequant_tq2_0_len, dequant_tq2_0_data, "main", 2, 5 * sizeof(uint32_t), {256 * 64, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_IQ1_S], "dequant_iq1_s", dequant_iq1_s_len, dequant_iq1_s_data, "main", 2, 5 * sizeof(uint32_t), {256 * 32, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_IQ1_M], "dequant_iq1_m", dequant_iq1_m_len, dequant_iq1_m_data, "main", 2, 5 * sizeof(uint32_t), {256 * 32, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_IQ2_XXS], "dequant_iq2_xxs", dequant_iq2_xxs_len, dequant_iq2_xxs_data, "main", 2, 5 * sizeof(uint32_t), {256 * 32, 1, 1}, {}, 1); @@ -5726,6 +6208,12 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_IQ3_S], "dequant_iq3_s", dequant_iq3_s_len, dequant_iq3_s_data, "main", 2, 5 * sizeof(uint32_t), {256 * 32, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_IQ4_XS], "dequant_iq4_xs", dequant_iq4_xs_len, dequant_iq4_xs_data, "main", 2, 5 * sizeof(uint32_t), {256 * 32, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_IQ4_NL], "dequant_iq4_nl", dequant_iq4_nl_len, dequant_iq4_nl_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q4_0_ROCMFP4], "dequant_rocmfp4", dequant_rocmfp4_len, dequant_rocmfp4_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q4_0_ROCMFP4_FAST], "dequant_rocmfp4_fast", dequant_rocmfp4_fast_len, dequant_rocmfp4_fast_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q3_0_ROCMFPX], "dequant_rocmfpx_fp3", dequant_rocmfpx_fp3_len, dequant_rocmfpx_fp3_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q6_0_ROCMFPX], "dequant_rocmfpx_fp6", dequant_rocmfpx_fp6_len, dequant_rocmfpx_fp6_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_Q8_0_ROCMFPX], "dequant_rocmfpx_fp8", dequant_rocmfpx_fp8_len, dequant_rocmfpx_fp8_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dequant_transpose[GGML_TYPE_IQ4_NL], "dequant_iq4_nl_transpose", dequant_iq4_nl_transpose_len, dequant_iq4_nl_transpose_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_MXFP4], "dequant_mxfp4", dequant_mxfp4_len, dequant_mxfp4_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_dequant[GGML_TYPE_NVFP4], "dequant_nvfp4", dequant_nvfp4_len, dequant_nvfp4_data, "main", 2, 5 * sizeof(uint32_t), {256 * 16, 1, 1}, {}, 1); @@ -5741,12 +6229,12 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q5_1], "get_rows_q5_1", get_rows_q5_1_len, get_rows_q5_1_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q8_0], "get_rows_q8_0", get_rows_q8_0_len, get_rows_q8_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q2_K], "get_rows_q2_k", get_rows_q2_k_len, get_rows_q2_k_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_TQ2_0], "get_rows_tq2_0", get_rows_tq2_0_len, get_rows_tq2_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_TQ1_0], "get_rows_tq1_0", get_rows_tq1_0_len, get_rows_tq1_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q3_K], "get_rows_q3_k", get_rows_q3_k_len, get_rows_q3_k_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q4_K], "get_rows_q4_k", get_rows_q4_k_len, get_rows_q4_k_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q5_K], "get_rows_q5_k", get_rows_q5_k_len, get_rows_q5_k_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q6_K], "get_rows_q6_k", get_rows_q6_k_len, get_rows_q6_k_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); - ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_TQ1_0], "get_rows_tq1_0", get_rows_tq1_0_len, get_rows_tq1_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); - ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_TQ2_0], "get_rows_tq2_0", get_rows_tq2_0_len, get_rows_tq2_0_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_IQ1_S], "get_rows_iq1_s", get_rows_iq1_s_len, get_rows_iq1_s_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_IQ1_M], "get_rows_iq1_m", get_rows_iq1_m_len, get_rows_iq1_m_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_IQ2_XXS], "get_rows_iq2_xxs", get_rows_iq2_xxs_len, get_rows_iq2_xxs_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); @@ -5756,6 +6244,11 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_IQ3_S], "get_rows_iq3_s", get_rows_iq3_s_len, get_rows_iq3_s_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_IQ4_XS], "get_rows_iq4_xs", get_rows_iq4_xs_len, get_rows_iq4_xs_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_IQ4_NL], "get_rows_iq4_nl", get_rows_iq4_nl_len, get_rows_iq4_nl_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q4_0_ROCMFP4], "get_rows_rocmfp4", get_rows_rocmfp4_len, get_rows_rocmfp4_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q4_0_ROCMFP4_FAST], "get_rows_rocmfp4_fast", get_rows_rocmfp4_fast_len, get_rows_rocmfp4_fast_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q3_0_ROCMFPX], "get_rows_rocmfpx_fp3", get_rows_rocmfpx_fp3_len, get_rows_rocmfpx_fp3_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q6_0_ROCMFPX], "get_rows_rocmfpx_fp6", get_rows_rocmfpx_fp6_len, get_rows_rocmfpx_fp6_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_Q8_0_ROCMFPX], "get_rows_rocmfpx_fp8", get_rows_rocmfpx_fp8_len, get_rows_rocmfpx_fp8_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_MXFP4], "get_rows_mxfp4", get_rows_mxfp4_len, get_rows_mxfp4_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_NVFP4], "get_rows_nvfp4", get_rows_nvfp4_len, get_rows_nvfp4_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows[GGML_TYPE_I32], "get_rows_i32", get_rows_i32_len, get_rows_i32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); @@ -5771,12 +6264,12 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q5_1], "get_rows_q5_1_f32", get_rows_q5_1_f32_len, get_rows_q5_1_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q8_0], "get_rows_q8_0_f32", get_rows_q8_0_f32_len, get_rows_q8_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q2_K], "get_rows_q2_k_f32", get_rows_q2_k_f32_len, get_rows_q2_k_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_TQ2_0], "get_rows_tq2_0_f32", get_rows_tq2_0_f32_len, get_rows_tq2_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_TQ1_0], "get_rows_tq1_0_f32", get_rows_tq1_0_f32_len, get_rows_tq1_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q3_K], "get_rows_q3_k_f32", get_rows_q3_k_f32_len, get_rows_q3_k_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q4_K], "get_rows_q4_k_f32", get_rows_q4_k_f32_len, get_rows_q4_k_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q5_K], "get_rows_q5_k_f32", get_rows_q5_k_f32_len, get_rows_q5_k_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q6_K], "get_rows_q6_k_f32", get_rows_q6_k_f32_len, get_rows_q6_k_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); - ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_TQ1_0], "get_rows_tq1_0_f32", get_rows_tq1_0_f32_len, get_rows_tq1_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); - ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_TQ2_0], "get_rows_tq2_0_f32", get_rows_tq2_0_f32_len, get_rows_tq2_0_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_IQ1_S], "get_rows_iq1_s_f32", get_rows_iq1_s_f32_len, get_rows_iq1_s_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_IQ1_M], "get_rows_iq1_m_f32", get_rows_iq1_m_f32_len, get_rows_iq1_m_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_IQ2_XXS], "get_rows_iq2_xxs_f32", get_rows_iq2_xxs_f32_len, get_rows_iq2_xxs_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); @@ -5786,6 +6279,11 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_IQ3_S], "get_rows_iq3_s_f32", get_rows_iq3_s_f32_len, get_rows_iq3_s_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_IQ4_XS], "get_rows_iq4_xs_f32", get_rows_iq4_xs_f32_len, get_rows_iq4_xs_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_IQ4_NL], "get_rows_iq4_nl_f32", get_rows_iq4_nl_f32_len, get_rows_iq4_nl_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q4_0_ROCMFP4], "get_rows_rocmfp4_f32", get_rows_rocmfp4_f32_len, get_rows_rocmfp4_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q4_0_ROCMFP4_FAST], "get_rows_rocmfp4_fast_f32", get_rows_rocmfp4_fast_f32_len, get_rows_rocmfp4_fast_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q3_0_ROCMFPX], "get_rows_rocmfpx_fp3_f32", get_rows_rocmfpx_fp3_f32_len, get_rows_rocmfpx_fp3_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q6_0_ROCMFPX], "get_rows_rocmfpx_fp6_f32", get_rows_rocmfpx_fp6_f32_len, get_rows_rocmfpx_fp6_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_Q8_0_ROCMFPX], "get_rows_rocmfpx_fp8_f32", get_rows_rocmfpx_fp8_f32_len, get_rows_rocmfpx_fp8_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_MXFP4], "get_rows_mxfp4_f32", get_rows_mxfp4_f32_len, get_rows_mxfp4_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows_f32[GGML_TYPE_NVFP4], "get_rows_nvfp4_f32", get_rows_nvfp4_f32_len, get_rows_nvfp4_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {1024, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_get_rows_back_f32, "get_rows_back_f32", get_rows_back_f32_len, get_rows_back_f32_data, "main", 3, sizeof(vk_op_binary_push_constants), {256, 1, 1}, {}, 1, true); @@ -5798,22 +6296,6 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, it.second, "fa_mask_opt", fa_mask_opt_len, fa_mask_opt_data, "main", 2, sizeof(vk_op_flash_attn_mask_opt_push_constants), {1, 1, 1}, {128, 128 / device->subgroup_size, BrBc.first, BrBc.second}, 1, true, true, device->subgroup_size); } - { - // Large workgroup so the per-row KV scan parallelizes; capped to device limits. - const uint32_t compact_max = std::min({1024u, device->properties.limits.maxComputeWorkGroupInvocations, device->properties.limits.maxComputeWorkGroupSize[0]}); - - // Fast ballot prefix-sum path when the device supports full subgroups; otherwise - // a shared-memory prefix-sum fallback. Both emit a deterministic ascending list. - device->fa_sparse_compact_use_subgroups = device->subgroup_ballot && device->subgroup_require_full_support; - if (device->fa_sparse_compact_use_subgroups) { - const uint32_t compact_wg = std::max(device->subgroup_size, (compact_max / device->subgroup_size) * device->subgroup_size); - const uint32_t compact_num_sg = compact_wg / device->subgroup_size; - ggml_vk_create_pipeline(device, device->pipeline_fa_sparse_compact_subgroup, "fa_sparse_compact_subgroup", fa_sparse_compact_subgroup_len, fa_sparse_compact_subgroup_data, "main", 2, sizeof(vk_op_flash_attn_sparse_compact_push_constants), {1, 1, 1}, {compact_wg, compact_num_sg}, 1, true, true, device->subgroup_size); - } else { - ggml_vk_create_pipeline(device, device->pipeline_fa_sparse_compact, "fa_sparse_compact", fa_sparse_compact_len, fa_sparse_compact_data, "main", 2, sizeof(vk_op_flash_attn_sparse_compact_push_constants), {1, 1, 1}, {compact_max}, 1, true); - } - } - if (device->subgroup_clustered && device->subgroup_require_full_support) { ggml_vk_create_pipeline(device, device->pipeline_quantize_q8_1_x4, "quantize_q8_1_x4", quantize_q8_1_x4_subgroup_len, quantize_q8_1_x4_subgroup_data, "main", 2, sizeof(vk_quantize_q8_1_push_constants), {32 * device->subgroup_size / 8, 1, 1}, { device->subgroup_size }, 1, true, true); } else { @@ -5882,6 +6364,11 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_quant[GGML_TYPE_Q5_1], "cpy_f32_q5_1", cpy_f32_q5_1_len, cpy_f32_q5_1_data, "main", 2, sizeof(vk_op_unary_push_constants), {32, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_quant[GGML_TYPE_Q8_0], "cpy_f32_q8_0", cpy_f32_q8_0_len, cpy_f32_q8_0_data, "main", 2, sizeof(vk_op_unary_push_constants), {32, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_quant[GGML_TYPE_IQ4_NL], "cpy_f32_iq4_nl", cpy_f32_iq4_nl_len, cpy_f32_iq4_nl_data, "main", 2, sizeof(vk_op_unary_push_constants), {32, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_quant[GGML_TYPE_Q4_0_ROCMFP4], "cpy_f32_rocmfp4", cpy_f32_rocmfp4_len, cpy_f32_rocmfp4_data, "main", 2, sizeof(vk_op_unary_push_constants), {32, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_quant[GGML_TYPE_Q4_0_ROCMFP4_FAST], "cpy_f32_rocmfp4_fast", cpy_f32_rocmfp4_fast_len, cpy_f32_rocmfp4_fast_data, "main", 2, sizeof(vk_op_unary_push_constants), {32, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_quant[GGML_TYPE_Q3_0_ROCMFPX], "cpy_f32_rocmfpx_fp3", cpy_f32_rocmfpx_fp3_len, cpy_f32_rocmfpx_fp3_data, "main", 2, sizeof(vk_op_unary_push_constants), {32, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_quant[GGML_TYPE_Q6_0_ROCMFPX], "cpy_f32_rocmfpx_fp6", cpy_f32_rocmfpx_fp6_len, cpy_f32_rocmfpx_fp6_data, "main", 2, sizeof(vk_op_unary_push_constants), {32, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_cpy_f32_quant[GGML_TYPE_Q8_0_ROCMFPX], "cpy_f32_rocmfpx_fp8", cpy_f32_rocmfpx_fp8_len, cpy_f32_rocmfpx_fp8_data, "main", 2, sizeof(vk_op_unary_push_constants), {32, 1, 1}, {}, 1); #define SET_ROWS(src_idx, src, itype) \ ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_F32], "set_rows_" #src "_f32" #itype, set_rows_ ## src ## _f32 ## itype ## _len, set_rows_ ## src ## _f32 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \ @@ -5894,7 +6381,12 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_Q5_0], "set_rows_" #src "_q5_0" #itype, set_rows_ ## src ## _q5_0 ## itype ## _len, set_rows_ ## src ## _q5_0 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \ ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_Q5_1], "set_rows_" #src "_q5_1" #itype, set_rows_ ## src ## _q5_1 ## itype ## _len, set_rows_ ## src ## _q5_1 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \ ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_Q8_0], "set_rows_" #src "_q8_0" #itype, set_rows_ ## src ## _q8_0 ## itype ## _len, set_rows_ ## src ## _q8_0 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \ - ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_IQ4_NL], "set_rows_" #src "_iq4_nl" #itype, set_rows_ ## src ## _iq4_nl ## itype ## _len, set_rows_ ## src ## _iq4_nl ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); + ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_IQ4_NL], "set_rows_" #src "_iq4_nl" #itype, set_rows_ ## src ## _iq4_nl ## itype ## _len, set_rows_ ## src ## _iq4_nl ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \ + ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_Q4_0_ROCMFP4], "set_rows_" #src "_rocmfp4" #itype, set_rows_ ## src ## _rocmfp4 ## itype ## _len, set_rows_ ## src ## _rocmfp4 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \ + ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_Q4_0_ROCMFP4_FAST], "set_rows_" #src "_rocmfp4_fast" #itype, set_rows_ ## src ## _rocmfp4_fast ## itype ## _len, set_rows_ ## src ## _rocmfp4_fast ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \ + ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_Q3_0_ROCMFPX], "set_rows_" #src "_rocmfpx_fp3" #itype, set_rows_ ## src ## _rocmfpx_fp3 ## itype ## _len, set_rows_ ## src ## _rocmfpx_fp3 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \ + ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_Q6_0_ROCMFPX], "set_rows_" #src "_rocmfpx_fp6" #itype, set_rows_ ## src ## _rocmfpx_fp6 ## itype ## _len, set_rows_ ## src ## _rocmfpx_fp6 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); \ + ggml_vk_create_pipeline(device, device->pipeline_set_rows ## itype [src_idx][GGML_TYPE_Q8_0_ROCMFPX], "set_rows_" #src "_rocmfpx_fp8" #itype, set_rows_ ## src ## _rocmfpx_fp8 ## itype ## _len, set_rows_ ## src ## _rocmfpx_fp8 ## itype ## _data, "main", 3, sizeof(vk_op_binary_push_constants), {1, 1, 1}, {1}, 1, true); SET_ROWS(0, f32, _i32) SET_ROWS(0, f32, _i64) @@ -5911,6 +6403,11 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_cpy_quant_f32[GGML_TYPE_Q5_1], "cpy_q5_1_f32", cpy_q5_1_f32_len, cpy_q5_1_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {(uint32_t)ggml_blck_size(GGML_TYPE_Q5_1), 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_cpy_quant_f32[GGML_TYPE_Q8_0], "cpy_q8_0_f32", cpy_q8_0_f32_len, cpy_q8_0_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {(uint32_t)ggml_blck_size(GGML_TYPE_Q8_0), 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_cpy_quant_f32[GGML_TYPE_IQ4_NL], "cpy_iq4_nl_f32", cpy_iq4_nl_f32_len, cpy_iq4_nl_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {(uint32_t)ggml_blck_size(GGML_TYPE_IQ4_NL), 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_cpy_quant_f32[GGML_TYPE_Q4_0_ROCMFP4], "cpy_rocmfp4_f32", cpy_rocmfp4_f32_len, cpy_rocmfp4_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {(uint32_t)ggml_blck_size(GGML_TYPE_Q4_0_ROCMFP4), 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_cpy_quant_f32[GGML_TYPE_Q4_0_ROCMFP4_FAST], "cpy_rocmfp4_fast_f32", cpy_rocmfp4_fast_f32_len, cpy_rocmfp4_fast_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {(uint32_t)ggml_blck_size(GGML_TYPE_Q4_0_ROCMFP4_FAST), 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_cpy_quant_f32[GGML_TYPE_Q3_0_ROCMFPX], "cpy_rocmfpx_fp3_f32", cpy_rocmfpx_fp3_f32_len, cpy_rocmfpx_fp3_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {(uint32_t)ggml_blck_size(GGML_TYPE_Q3_0_ROCMFPX), 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_cpy_quant_f32[GGML_TYPE_Q6_0_ROCMFPX], "cpy_rocmfpx_fp6_f32", cpy_rocmfpx_fp6_f32_len, cpy_rocmfpx_fp6_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {(uint32_t)ggml_blck_size(GGML_TYPE_Q6_0_ROCMFPX), 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_cpy_quant_f32[GGML_TYPE_Q8_0_ROCMFPX], "cpy_rocmfpx_fp8_f32", cpy_rocmfpx_fp8_f32_len, cpy_rocmfpx_fp8_f32_data, "main", 2, sizeof(vk_op_unary_push_constants), {(uint32_t)ggml_blck_size(GGML_TYPE_Q8_0_ROCMFPX), 1, 1}, {}, 1); auto get_suffix = [](bool src0_f16, bool src1_f16, bool dst_f16) { std::string s; @@ -5953,6 +6450,8 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_concat_i8, "concat_i8", concat_i8_len, concat_i8_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_concat_i16, "concat_i16", concat_i16_len, concat_i16_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_concat_i32, "concat_i32", concat_i32_len, concat_i32_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1); + // One workgroup per 32x32 tile: elements are passed as (rows, cols, 1). + ggml_vk_create_pipeline(device, device->pipeline_concat_transpose_i32, "concat_transpose_i32", concat_transpose_i32_len, concat_transpose_i32_data, "main", 3, sizeof(vk_op_binary_push_constants), {32, 32, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_concat_i64, "concat_i64", concat_i64_len, concat_i64_data, "main", 3, sizeof(vk_op_binary_push_constants), {512, 1, 1}, {}, 1); ggml_vk_create_pipeline(device, device->pipeline_upscale_nearest_f32, "upscale_f32", upscale_f32_len, upscale_f32_data, "main", 2, sizeof(vk_op_upscale_push_constants), {512, 1, 1}, {GGML_SCALE_MODE_NEAREST}, 1); @@ -6166,16 +6665,6 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { ggml_vk_create_pipeline(device, device->pipeline_count_experts, "count_experts", count_experts_len, count_experts_data, "main", 2, sizeof(vk_op_count_experts_push_constants), {1, 1, 1}, {}, 1, true); } - // comb holds a token's 4x4 matrix in one 16-lane slice of a subgroup, so it - // needs at least 16 lanes, pinned to a known size. - if (device->subgroup_basic && device->subgroup_shuffle && device->subgroup_require_full_support && device->subgroup_size >= 16) { - const uint32_t tokens_per_workgroup = 4 * (device->subgroup_size / 16); - ggml_vk_create_pipeline(device, device->pipeline_dsv4_hc_comb_f32, "dsv4_hc_comb_f32", dsv4_hc_comb_f32_len, dsv4_hc_comb_f32_data, "main", 4, sizeof(vk_op_dsv4_hc_comb_push_constants), {tokens_per_workgroup, 1, 1}, { device->subgroup_size }, 1, true, true, device->subgroup_size); - } - - ggml_vk_create_pipeline(device, device->pipeline_dsv4_hc_pre_f32, "dsv4_hc_pre_f32", dsv4_hc_pre_f32_len, dsv4_hc_pre_f32_data, "main", 3, sizeof(vk_op_dsv4_hc_pre_push_constants), {256, 1, 1}, { 256 }, 1); - ggml_vk_create_pipeline(device, device->pipeline_dsv4_hc_post_f32, "dsv4_hc_post_f32", dsv4_hc_post_f32_len, dsv4_hc_post_f32_data, "main", 5, sizeof(vk_op_dsv4_hc_post_push_constants), {256, 1, 1}, { 256 }, 1); - for (auto &s : device->pipeline_solve_tri_f32) { const vk_solve_tri_pipeline_state &state = s.first; @@ -6300,6 +6789,92 @@ static void ggml_vk_load_shaders(vk_device& device, vk_pipeline requested) { } } + if (device->subgroup_arithmetic && device->subgroup_size == 64) { + for (size_t nhi = 0; nhi < LI_NH_COUNT; ++nhi) { + ggml_vk_create_pipeline(device, device->pipeline_lightning_indexer_f16[nhi], + "lightning_indexer_f16", lightning_indexer_f16_len, lightning_indexer_f16_data, "main", 5, + sizeof(vk_op_lightning_indexer_cm_push_constants), {8, 1, 1}, {device->subgroup_size, LI_NH_VALUES[nhi]}, 1, true, true, + device->subgroup_size); + } +#if defined(VK_KHR_cooperative_matrix) && defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT) + if (device->coopmat_support && device->coopmat_support_16x16x16_f32acc && device->subgroup_size_control) { + for (size_t nhi = 0; nhi < LI_NH_COUNT; ++nhi) { + ggml_vk_create_pipeline(device, device->pipeline_lightning_indexer_cm_small_f16[nhi], + "lightning_indexer_cm_small_f16", lightning_indexer_cm_small_f16_len, lightning_indexer_cm_small_f16_data, "main", 5, + sizeof(vk_op_lightning_indexer_cm_push_constants), {16, 16, 1}, {device->subgroup_size, LI_NH_VALUES[nhi]}, 1, true, true, + device->subgroup_size); + } + if (device->properties.limits.maxComputeWorkGroupInvocations >= 512 && + device->properties.limits.maxComputeWorkGroupSize[0] >= 512 && + device->properties.limits.maxComputeSharedMemorySize >= 64 * 1024) { + for (size_t nhi = 0; nhi < LI_NH_COUNT; ++nhi) { + ggml_vk_create_pipeline(device, device->pipeline_lightning_indexer_cm_f16[nhi], + "lightning_indexer_cm_f16", lightning_indexer_cm_f16_len, lightning_indexer_cm_f16_data, "main", 5, + sizeof(vk_op_lightning_indexer_cm_push_constants), {128, 16, 1}, {device->subgroup_size, LI_NH_VALUES[nhi]}, 1, true, true, + device->subgroup_size); + } + } + for (size_t nhi = 0; nhi < LI_NH_COUNT; ++nhi) { + ggml_vk_create_pipeline(device, device->pipeline_lightning_indexer_decode_cm_f16[nhi], + "lightning_indexer_decode_cm_f16", lightning_indexer_decode_cm_f16_len, lightning_indexer_decode_cm_f16_data, "main", 5, + sizeof(vk_op_lightning_indexer_cm_push_constants), {16, 1, 1}, {device->subgroup_size, LI_NH_VALUES[nhi]}, 1, true, true, + device->subgroup_size); + } + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_top_k_cm_f16, + "flash_attn_top_k_cm_f16", flash_attn_top_k_cm_f16_len, flash_attn_top_k_cm_f16_data, "main", 6, + sizeof(vk_op_flash_attn_top_k_push_constants), {1, 1, 1}, {512, device->subgroup_size}, 1, true, true, + device->subgroup_size); + } +#endif + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_top_k_f16, + "flash_attn_top_k_f16", flash_attn_top_k_f16_len, flash_attn_top_k_f16_data, "main", 6, + sizeof(vk_op_flash_attn_top_k_push_constants), {1, 1, 1}, {512, device->subgroup_size}, 1, true, true, + device->subgroup_size); + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_gather_f16, + "flash_attn_gather_f16", flash_attn_gather_f16_len, flash_attn_gather_f16_data, "main", 5, + sizeof(vk_op_flash_attn_gather_push_constants), {1, 1, 1}, {}, 1, true, true, + device->subgroup_size); +#define CREATE_FA_GATHER_TOK_DQ(TYPE, NAMED) \ + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_gather_dq[TYPE], \ + "flash_attn_gather_dq_" #NAMED, flash_attn_gather_dq_ ## NAMED ## _len, \ + flash_attn_gather_dq_ ## NAMED ## _data, "main", 5, \ + sizeof(vk_op_flash_attn_gather_push_constants), {1, 1, 1}, {}, 1, true, true, \ + device->subgroup_size); + CREATE_FA_GATHER_TOK_DQ(GGML_TYPE_Q4_0, q4_0) + CREATE_FA_GATHER_TOK_DQ(GGML_TYPE_Q8_0, q8_0) +#undef CREATE_FA_GATHER_TOK_DQ + if (device->subgroup_arithmetic) { + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_union_f16, + "flash_attn_union_f16", flash_attn_union_f16_len, flash_attn_union_f16_data, "main", 3, + sizeof(vk_op_flash_attn_union_push_constants), {1, 1, 1}, {}, 1, true, true, + device->subgroup_size); + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_gather_union_f16, + "flash_attn_gather_union_f16", flash_attn_gather_union_f16_len, flash_attn_gather_union_f16_data, "main", 6, + sizeof(vk_op_flash_attn_gather_union_push_constants), {1, 1, 1}, {}, 1, true, true, + device->subgroup_size); +#define CREATE_FA_GATHER_DQ(TYPE, NAMED) \ + ggml_vk_create_pipeline(device, device->pipeline_flash_attn_gather_union_dq[TYPE], \ + "flash_attn_gather_union_dq_" #NAMED, flash_attn_gather_union_dq_ ## NAMED ## _len, \ + flash_attn_gather_union_dq_ ## NAMED ## _data, "main", 6, \ + sizeof(vk_op_flash_attn_gather_union_push_constants), {1, 1, 1}, {}, 1, true, true, \ + device->subgroup_size); + CREATE_FA_GATHER_DQ(GGML_TYPE_Q4_0, q4_0) + CREATE_FA_GATHER_DQ(GGML_TYPE_Q8_0, q8_0) +#undef CREATE_FA_GATHER_DQ + } + } + + // DSv4 fused hyper-connection ops: plain f32 compute, no subgroup/coopmat requirements + ggml_vk_create_pipeline(device, device->pipeline_dsv4_hc_pre_f32, + "dsv4_hc_pre_f32", dsv4_hc_pre_f32_len, dsv4_hc_pre_f32_data, "main", 3, + sizeof(vk_op_dsv4_hc_pre_push_constants), {256, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dsv4_hc_comb_f32, + "dsv4_hc_comb_f32", dsv4_hc_comb_f32_len, dsv4_hc_comb_f32_data, "main", 4, + sizeof(vk_op_dsv4_hc_comb_push_constants), {256, 1, 1}, {}, 1); + ggml_vk_create_pipeline(device, device->pipeline_dsv4_hc_post_f32, + "dsv4_hc_post_f32", dsv4_hc_post_f32_len, dsv4_hc_post_f32_data, "main", 5, + sizeof(vk_op_dsv4_hc_post_push_constants), {256, 1, 1}, {}, 1); + if (device->subgroup_arithmetic && device->subgroup_require_full_support) { ggml_vk_create_pipeline(device, device->pipeline_ssm_scan_f32_d128, "ssm_scan_128_f32", ssm_scan_subgroup_f32_len, ssm_scan_subgroup_f32_data, "main", 8, sizeof(vk_op_ssm_scan_push_constants), {1, 1, 1}, {128, device->subgroup_size}, 1, true, true); ggml_vk_create_pipeline(device, device->pipeline_ssm_scan_f32_d256, "ssm_scan_256_f32", ssm_scan_subgroup_f32_len, ssm_scan_subgroup_f32_data, "main", 8, sizeof(vk_op_ssm_scan_push_constants), {1, 1, 1}, {256, device->subgroup_size}, 1, true, true); @@ -6837,6 +7412,15 @@ static vk_device ggml_vk_get_device(size_t idx) { device->max_nodes_per_submit = std::max(max_nodes_per_submit, 1u); } + // Also submit once a batch has accumulated enough memory traffic, so that + // bandwidth-bound nodes with no flops estimate cannot grow a command buffer + // past the driver timeout. 0 disables the limit. + device->max_bytes_per_submit = 8ull * 1024 * 1024 * 1024; + const char* GGML_VK_MAX_MB_PER_SUBMIT = getenv("GGML_VK_MAX_MB_PER_SUBMIT"); + if (GGML_VK_MAX_MB_PER_SUBMIT != nullptr) { + device->max_bytes_per_submit = std::stoull(GGML_VK_MAX_MB_PER_SUBMIT) * 1024 * 1024; + } + const bool force_disable_f16 = getenv("GGML_VK_DISABLE_F16") != nullptr; device->fp16 = !force_disable_f16 && fp16_storage && fp16_compute; @@ -8007,6 +8591,11 @@ static void ggml_vk_init(ggml_backend_vk_context * ctx, size_t idx) { ctx->prealloc_size_x = 0; ctx->prealloc_size_y = 0; ctx->prealloc_size_split_k = 0; + ctx->fa_dequant_gate_sz = 0; + ctx->fa_dequant_gate_fits = false; + ctx->fa_dequant_gate_logged = false; + memset(ctx->fa_union_est_ratio, 0, sizeof(ctx->fa_union_est_ratio)); + ctx->fa_union_declines = 0; // Fixed size of 1KB, for deterministic behavior ctx->prealloc_size_add_rms_partials = 1024; @@ -8063,8 +8652,13 @@ static vk_pipeline ggml_vk_get_to_fp16(ggml_backend_vk_context * ctx, ggml_type case GGML_TYPE_IQ4_NL: case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: - case GGML_TYPE_TQ1_0: + case GGML_TYPE_Q4_0_ROCMFP4: + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + case GGML_TYPE_Q3_0_ROCMFPX: + case GGML_TYPE_Q6_0_ROCMFPX: + case GGML_TYPE_Q8_0_ROCMFPX: case GGML_TYPE_TQ2_0: + case GGML_TYPE_TQ1_0: break; default: return nullptr; @@ -8073,6 +8667,97 @@ static vk_pipeline ggml_vk_get_to_fp16(ggml_backend_vk_context * ctx, ggml_type return ctx->device->pipeline_dequant[type]; } +static vk_matmul_pipeline ggml_vk_get_mul_mat_mat_pipeline(ggml_backend_vk_context * ctx, ggml_type src0_type, ggml_type src1_type, ggml_prec prec) { + VK_LOG_DEBUG("ggml_vk_get_mul_mat_mat_pipeline(" << ggml_type_name(src0_type) << ", " << ggml_type_name(src1_type) << ", " << prec << ")"); + if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_F32) { + return ctx->device->pipeline_matmul_f32; + } + if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_F16) { + return ctx->device->pipeline_matmul_f32_f16; + } + if (src0_type == GGML_TYPE_BF16 && src1_type == GGML_TYPE_BF16) { + return ctx->device->pipeline_matmul_bf16; + } + if (prec == GGML_PREC_DEFAULT && ctx->device->fp16 && !(ctx->device->coopmat_support && !ctx->device->coopmat_acc_f16_support)) { + if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F32) { + return ctx->device->pipeline_matmul_f16_f32.f16acc; + } + if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F16) { + return ctx->device->pipeline_matmul_f16.f16acc; + } + } else { + if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F32) { + return ctx->device->pipeline_matmul_f16_f32.f32acc; + } + if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F16) { + return ctx->device->pipeline_matmul_f16.f32acc; + } + } + + // MMQ + if (src1_type == GGML_TYPE_Q8_1) { + vk_matmul_pipeline pipelines = ctx->device->pipeline_dequant_mul_mat_mat_q8_1[src0_type].f32acc; + + if (pipelines->is_empty()) { + return nullptr; + } + + return pipelines; + } + + if (src1_type != GGML_TYPE_F32 && !ctx->device->coopmat2 && + !(src1_type == GGML_TYPE_F16 && ggml_vk_dense_f16b_enabled())) { + return nullptr; + } + + switch (src0_type) { + case GGML_TYPE_Q1_0: + case GGML_TYPE_Q2_0: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_IQ1_S: + case GGML_TYPE_IQ1_M: + 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: + case GGML_TYPE_IQ4_XS: + case GGML_TYPE_IQ4_NL: + case GGML_TYPE_MXFP4: + case GGML_TYPE_NVFP4: + case GGML_TYPE_Q4_0_ROCMFP4: + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + case GGML_TYPE_Q3_0_ROCMFPX: + case GGML_TYPE_Q6_0_ROCMFPX: + case GGML_TYPE_Q8_0_ROCMFPX: + case GGML_TYPE_TQ2_0: + case GGML_TYPE_TQ1_0: + break; + default: + return nullptr; + } + + if (ctx->device->coopmat2) { + assert(src1_type == GGML_TYPE_F16); + return prec == GGML_PREC_DEFAULT ? ctx->device->pipeline_dequant_mul_mat_mat_f16[src0_type].f16acc : ctx->device->pipeline_dequant_mul_mat_mat_f16[src0_type].f32acc; + } + if (ctx->device->coopmat_support) { + vk_matmul_pipeline2 & p = (src1_type == GGML_TYPE_F16) ? ctx->device->pipeline_dequant_mul_mat_mat_f16[src0_type] + : ctx->device->pipeline_dequant_mul_mat_mat[src0_type]; + return (ctx->device->fp16 && ctx->device->coopmat_acc_f16_support && prec == GGML_PREC_DEFAULT) ? p.f16acc : p.f32acc; + } + return (ctx->device->fp16 && prec == GGML_PREC_DEFAULT) ? ctx->device->pipeline_dequant_mul_mat_mat[src0_type].f16acc : ctx->device->pipeline_dequant_mul_mat_mat[src0_type].f32acc; +} + static vk_pipeline ggml_vk_get_dequantize_mul_mat_vec(ggml_backend_vk_context * ctx, ggml_type a_type, ggml_type b_type, uint32_t num_cols, uint32_t m, uint32_t k) { VK_LOG_DEBUG("ggml_vk_get_dequantize_mul_mat_vec()"); GGML_ASSERT(b_type == GGML_TYPE_F32 || b_type == GGML_TYPE_F16 || b_type == GGML_TYPE_Q8_1); @@ -8087,6 +8772,9 @@ static vk_pipeline ggml_vk_get_dequantize_mul_mat_vec(ggml_backend_vk_context * case GGML_TYPE_Q5_1: case GGML_TYPE_Q8_0: case GGML_TYPE_MXFP4: + case GGML_TYPE_Q4_0_ROCMFP4: + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + case GGML_TYPE_Q8_0_ROCMFPX: case GGML_TYPE_Q2_K: case GGML_TYPE_Q3_K: case GGML_TYPE_Q4_K: @@ -8127,8 +8815,13 @@ static vk_pipeline ggml_vk_get_dequantize_mul_mat_vec(ggml_backend_vk_context * case GGML_TYPE_IQ4_NL: case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: - case GGML_TYPE_TQ1_0: + case GGML_TYPE_Q4_0_ROCMFP4: + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + case GGML_TYPE_Q3_0_ROCMFPX: + case GGML_TYPE_Q6_0_ROCMFPX: + case GGML_TYPE_Q8_0_ROCMFPX: case GGML_TYPE_TQ2_0: + case GGML_TYPE_TQ1_0: break; default: return nullptr; @@ -8161,6 +8854,98 @@ static vk_pipeline ggml_vk_get_dequantize_mul_mat_vec(ggml_backend_vk_context * return b_type == GGML_TYPE_F32 ? ctx->device->pipeline_dequant_mul_mat_vec_f32_f32[dmmv_wg][a_type][num_cols-1] : ctx->device->pipeline_dequant_mul_mat_vec_f16_f32[dmmv_wg][a_type][num_cols-1]; } +static vk_matmul_pipeline ggml_vk_get_mul_mat_mat_id_pipeline(ggml_backend_vk_context * ctx, ggml_type src0_type, ggml_type src1_type, ggml_prec prec) { + VK_LOG_DEBUG("ggml_vk_get_mul_mat_mat_id_pipeline()"); + if (src0_type == GGML_TYPE_F32 && src1_type == GGML_TYPE_F32) { + return ctx->device->pipeline_matmul_id_f32; + } + if (src0_type == GGML_TYPE_BF16 && src1_type == GGML_TYPE_BF16) { + return ctx->device->pipeline_matmul_id_bf16; + } + if (prec == GGML_PREC_DEFAULT && ctx->device->fp16 && !(ctx->device->coopmat_support && !ctx->device->coopmat_acc_f16_support)) { + if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F32) { + return ctx->device->pipeline_matmul_id_f16_f32.f16acc; + } + if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F16) { + return ctx->device->pipeline_matmul_id_f16.f16acc; + } + } else { + if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F32) { + return ctx->device->pipeline_matmul_id_f16_f32.f32acc; + } + if (src0_type == GGML_TYPE_F16 && src1_type == GGML_TYPE_F16) { + return ctx->device->pipeline_matmul_id_f16.f32acc; + } + } + + // MMQ + if (src1_type == GGML_TYPE_Q8_1) { + vk_matmul_pipeline pipelines = ctx->device->pipeline_dequant_mul_mat_mat_id_q8_1[src0_type].f32acc; + + if (pipelines->is_empty()) { + return nullptr; + } + + return pipelines; + } + + GGML_ASSERT(src1_type == GGML_TYPE_F32 || + ((ctx->device->coopmat2 || ggml_vk_mmid_f16b_enabled()) && src1_type == GGML_TYPE_F16)); + + switch (src0_type) { + case GGML_TYPE_Q1_0: + case GGML_TYPE_Q2_0: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_Q2_K: + case GGML_TYPE_Q3_K: + case GGML_TYPE_Q4_K: + case GGML_TYPE_Q5_K: + case GGML_TYPE_Q6_K: + case GGML_TYPE_IQ1_S: + case GGML_TYPE_IQ1_M: + 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: + case GGML_TYPE_IQ4_XS: + case GGML_TYPE_IQ4_NL: + case GGML_TYPE_MXFP4: + case GGML_TYPE_NVFP4: + case GGML_TYPE_Q4_0_ROCMFP4: + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + case GGML_TYPE_Q3_0_ROCMFPX: + case GGML_TYPE_Q6_0_ROCMFPX: + case GGML_TYPE_Q8_0_ROCMFPX: + case GGML_TYPE_TQ2_0: + case GGML_TYPE_TQ1_0: + break; + default: + return nullptr; + } + + // GGML_VK_MMID_F16B: on KHR_coopmat devices the f16-B mmid pipelines live in a + // parallel array (coopmat2's main array already holds f16-B pipelines). + vk_matmul_pipeline2& mmp = (src1_type == GGML_TYPE_F16 && !ctx->device->coopmat2) + ? ctx->device->pipeline_dequant_mul_mat_mat_id_f16b[src0_type] + : ctx->device->pipeline_dequant_mul_mat_mat_id[src0_type]; + // XXX TODO 'prec' is not actually allowed in mul_mat_id. + bool prefer_fp16acc = ctx->device->fp16 /*&& prec == GGML_PREC_DEFAULT*/; + bool support_fp16acc = !mmp.f16acc->is_empty(); + bool support_fp32acc = !mmp.f32acc->is_empty(); + + if (support_fp16acc && (prefer_fp16acc || !support_fp32acc)) { + return mmp.f16acc; + } else { + GGML_ASSERT(support_fp32acc); + return mmp.f32acc; + } +} + static vk_pipeline ggml_vk_get_dequantize_mul_mat_vec_id(ggml_backend_vk_context * ctx, ggml_type a_type, ggml_type b_type, uint32_t m, uint32_t k) { VK_LOG_DEBUG("ggml_vk_get_dequantize_mul_mat_vec_id()"); GGML_ASSERT(b_type == GGML_TYPE_F32 || b_type == GGML_TYPE_Q8_1); @@ -8174,6 +8959,9 @@ static vk_pipeline ggml_vk_get_dequantize_mul_mat_vec_id(ggml_backend_vk_context case GGML_TYPE_Q5_1: case GGML_TYPE_Q8_0: case GGML_TYPE_MXFP4: + case GGML_TYPE_Q4_0_ROCMFP4: + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + case GGML_TYPE_Q8_0_ROCMFPX: case GGML_TYPE_Q2_K: case GGML_TYPE_Q3_K: case GGML_TYPE_Q4_K: @@ -8214,8 +9002,13 @@ static vk_pipeline ggml_vk_get_dequantize_mul_mat_vec_id(ggml_backend_vk_context case GGML_TYPE_IQ4_NL: case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: - case GGML_TYPE_TQ1_0: + case GGML_TYPE_Q4_0_ROCMFP4: + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + case GGML_TYPE_Q3_0_ROCMFPX: + case GGML_TYPE_Q6_0_ROCMFPX: + case GGML_TYPE_Q8_0_ROCMFPX: case GGML_TYPE_TQ2_0: + case GGML_TYPE_TQ1_0: break; default: return nullptr; @@ -8946,7 +9739,13 @@ static void ggml_vk_buffer_read_2d(vk_buffer& src, size_t offset, void * dst, si // If the device is not an UMA device the memory is host-accessible through rebar. While writing // through PCIe is sufficient fast reading back data from PCIe is slower than going through // the HW device to host copy path. - if(src->memory_property_flags & vk::MemoryPropertyFlagBits::eHostVisible && src->device->uma) { + // On UMA a direct CPU read is only fast from a host-cached mapping. A write-combined mapping + // (host-visible without HOST_CACHED, which is what amdgpu gives for GTT) reads back at + // uncached speed, so bulk reads must use the device copy path. Small reads stay direct to + // avoid the fence round trip. + const bool host_cached = bool(src->memory_property_flags & vk::MemoryPropertyFlagBits::eHostCached); + if((src->memory_property_flags & vk::MemoryPropertyFlagBits::eHostVisible) && src->device->uma && + (host_cached || width * height <= 64 * 1024)) { GGML_ASSERT(src->memory_property_flags & vk::MemoryPropertyFlagBits::eHostCoherent); std::lock_guard guard(src->device->mutex); @@ -9075,7 +9874,7 @@ static uint32_t ggml_vk_guess_split_k(ggml_backend_vk_context * ctx, uint32_t m, } uint32_t split_k = 1; - if (ctx->device->shader_core_count != 0 && n >= pipeline->wg_denoms[1]) { + if (ctx->device->shader_core_count != 0 && m >= pipeline->wg_denoms[0] && n >= pipeline->wg_denoms[1]) { // If k is 'large' and the SMs will fill less than halfway, use split_k. uint32_t m_tiles = CEIL_DIV(m, pipeline->wg_denoms[0]); uint32_t n_tiles = CEIL_DIV(n, pipeline->wg_denoms[1]); @@ -9109,7 +9908,56 @@ static uint32_t ggml_vk_guess_split_k(ggml_backend_vk_context * ctx, uint32_t m, return split_k; } +static vk_pipeline ggml_vk_guess_matmul_pipeline(ggml_backend_vk_context * ctx, vk_matmul_pipeline& mmp, uint32_t m, uint32_t n, bool aligned, ggml_type src0_type, ggml_type src1_type) { + VK_LOG_DEBUG("ggml_vk_guess_matmul_pipeline(" << m << ", " << n << ", " << aligned << ", " << ggml_type_name(src0_type) << ", " << ggml_type_name(src1_type) << ")"); + + // The q8_1 (integer dot) mmq path uses a different shader with its own + // shared-memory layout, so use the int-specific availability flags. + const bool is_q8_1 = (src1_type == GGML_TYPE_Q8_1); + const bool mm_l = is_q8_1 ? ctx->device->mul_mat_l_int[src0_type] : ctx->device->mul_mat_l[src0_type]; + const bool mm_m = is_q8_1 ? ctx->device->mul_mat_m_int[src0_type] : ctx->device->mul_mat_m[src0_type]; + const bool mm_s = is_q8_1 ? ctx->device->mul_mat_s_int[src0_type] : ctx->device->mul_mat_s[src0_type]; + + if (ctx->device->coopmat2) { + const uint32_t shader_core_count = ctx->device->shader_core_count; + const uint32_t tiles_l = CEIL_DIV(m, mmp->a_l->wg_denoms[0]) * CEIL_DIV(n, mmp->a_l->wg_denoms[1]); + const uint32_t tiles_m = CEIL_DIV(m, mmp->a_m->wg_denoms[0]) * CEIL_DIV(n, mmp->a_m->wg_denoms[1]); + + // Use large shader when the N dimension is greater than the medium shader's tile size + uint32_t crossover_large = mmp->m->wg_denoms[1]; + // Prefer large over medium if either: + // - medium or large tiles would overfill the GPU + // - large tiles with a split_k==3 fits in the GPU and medium tiles with split_k==2 does not + // (medium with split_k==2 is probably better if it fits - more workgroups running and less split_k overhead) + bool prefer_large = tiles_m > shader_core_count || tiles_l > shader_core_count || + // split_k==3 with large tiles likely better than medium tiles with no split_k. + (tiles_l <= shader_core_count / 3 && tiles_m > shader_core_count / 2); + + if ((mm_l && (n > crossover_large && prefer_large)) || (!mm_m && !mm_s)) { + return aligned ? mmp->a_l : mmp->l; + } + // Use medium shader when the N dimension is greater than the small shader's tile size + uint32_t crossover_medium = mmp->s->wg_denoms[1]; + if ((mm_m && (n > crossover_medium)) || !mm_s) { + return aligned ? mmp->a_m : mmp->m; + } + return aligned ? mmp->a_s : mmp->s; + } + + if ((mm_s && (m <= 32 || n <= 32)) || (!mm_m && !mm_l)) { + return aligned ? mmp->a_s : mmp->s; + } + if ((mm_m && (m <= 64 || n <= 64)) || !mm_l) { + return aligned ? mmp->a_m : mmp->m; + } + return aligned ? mmp->a_l : mmp->l; +} + +static uint32_t ggml_vk_guess_matmul_pipeline_align(ggml_backend_vk_context * ctx, vk_matmul_pipeline& mmp, int m, int n, ggml_type src0_type, ggml_type src1_type) { + VK_LOG_DEBUG("ggml_vk_guess_matmul_pipeline_align(" << m << ", " << n << ", " << ggml_type_name(src0_type) << ", " << ggml_type_name(src1_type) << ")"); + return ggml_vk_guess_matmul_pipeline(ctx, mmp, m, n, true, src0_type, src1_type)->align; +} static void ggml_vk_matmul( ggml_backend_vk_context * ctx, vk_context& subctx, vk_pipeline& pipeline, @@ -9160,56 +10008,42 @@ static void ggml_vk_matmul( ctx->prealloc_split_k_need_sync = true; } +static vk_pipeline ggml_vk_guess_matmul_id_pipeline(ggml_backend_vk_context * ctx, vk_matmul_pipeline& mmp, uint32_t m, uint32_t n, bool aligned, ggml_type src0_type, ggml_type src1_type) { + VK_LOG_DEBUG("ggml_vk_guess_matmul_id_pipeline(" << m << ", " << n << ", " << aligned << ", " << ggml_type_name(src0_type) << ", " << ggml_type_name(src1_type) << ")"); -static bool ggml_vk_get_mul_mat_mat_f16acc(ggml_backend_vk_context * ctx, ggml_type src0_type, ggml_type src1_type, ggml_prec prec) { - if (src0_type == GGML_TYPE_F32 || src0_type == GGML_TYPE_BF16) return false; - if (src1_type == GGML_TYPE_Q8_1) return false; - if (src0_type == GGML_TYPE_F16) { - return prec == GGML_PREC_DEFAULT && ctx->device->fp16 && !(ctx->device->coopmat_support && !ctx->device->coopmat_acc_f16_support); - } - // quant types - if (ctx->device->coopmat2) { - return prec == GGML_PREC_DEFAULT; - } - if (ctx->device->coopmat_support) { - return ctx->device->fp16 && ctx->device->coopmat_acc_f16_support && prec == GGML_PREC_DEFAULT; - } - return ctx->device->fp16 && prec == GGML_PREC_DEFAULT; -} + // The q8_1 (integer dot) mmq path uses a different shader with its own + // shared-memory layout, so use the int-specific availability flags. + const bool is_q8_1 = (src1_type == GGML_TYPE_Q8_1); + const bool mm_l = is_q8_1 ? ctx->device->mul_mat_id_l_int[src0_type] : ctx->device->mul_mat_id_l[src0_type]; + const bool mm_m = is_q8_1 ? ctx->device->mul_mat_id_m_int[src0_type] : ctx->device->mul_mat_id_m[src0_type]; + const bool mm_s = is_q8_1 ? ctx->device->mul_mat_id_s_int[src0_type] : ctx->device->mul_mat_id_s[src0_type]; -static const std::vector* ggml_vk_get_mul_mat_mat_pipeline_map( - ggml_backend_vk_context * ctx, ggml_type src0_type, ggml_type src1_type, ggml_prec prec, bool mul_mat_id = false) { - bool f16acc = ggml_vk_get_mul_mat_mat_f16acc(ctx, src0_type, src1_type, prec); - vk_matmul_pipeline_key key{src0_type, src1_type, mul_mat_id, f16acc}; - auto it = ctx->device->pipeline_matmul.find(key); - if (it == ctx->device->pipeline_matmul.end() || it->second.empty()) { - // Try without f16acc - if (f16acc) { - key.f16acc = false; - it = ctx->device->pipeline_matmul.find(key); - if (it != ctx->device->pipeline_matmul.end() && !it->second.empty()) return &it->second; + if (ctx->device->coopmat2) { + // Use large shader when the N dimension is greater than the medium shader's tile size + uint32_t crossover_large = mmp->m->wg_denoms[1]; + if ((mm_l && (n > crossover_large)) || (!mm_m && !mm_s)) { + return aligned ? mmp->a_l : mmp->l; } - return nullptr; + // Use medium shader when the N dimension is greater than the small shader's tile size + uint32_t crossover_medium = mmp->s->wg_denoms[1]; + if ((mm_m && (n > crossover_medium)) || !mm_s) { + return aligned ? mmp->a_m : mmp->m; + } + return aligned ? mmp->a_s : mmp->s; } - return &it->second; -} -static vk_pipeline ggml_vk_guess_matmul_pipeline_map(ggml_backend_vk_context * ctx, - const std::vector& configs, - uint32_t m, uint32_t n, bool aligned, bool mul_mat_id) { - auto& selector = mul_mat_id ? ctx->device->matmul_id_tile_selector : ctx->device->matmul_tile_selector; - uint32_t idx = selector(m, n, 0, ctx->device->shader_core_count, configs); - if (idx >= configs.size()) idx = (uint32_t)configs.size() - 1; - return (aligned && configs[idx].aligned) ? configs[idx].aligned : configs[idx].unaligned; + if ((mm_s && (m <= 32 || n <= 32)) || (!mm_m && !mm_l)) { + return aligned ? mmp->a_s : mmp->s; + } + if ((mm_m && (m <= 64 || n <= 64)) || !mm_l) { + return aligned ? mmp->a_m : mmp->m; + } + return aligned ? mmp->a_l : mmp->l; } -static uint32_t ggml_vk_guess_matmul_pipeline_align_map(ggml_backend_vk_context * ctx, - const std::vector& configs, - uint32_t m, uint32_t n, bool mul_mat_id) { - auto& selector = mul_mat_id ? ctx->device->matmul_id_tile_selector : ctx->device->matmul_tile_selector; - uint32_t idx = selector(m, n, 0, ctx->device->shader_core_count, configs); - if (idx >= configs.size()) idx = (uint32_t)configs.size() - 1; - return configs[idx].align; +static uint32_t ggml_vk_guess_matmul_id_pipeline_align(ggml_backend_vk_context * ctx, vk_matmul_pipeline& mmp, int m, int n, ggml_type src0_type, ggml_type src1_type) { + VK_LOG_DEBUG("ggml_vk_guess_matmul_pipeline_align(" << m << ", " << n << ", " << ggml_type_name(src0_type) << ", " << ggml_type_name(src1_type) << ")"); + return ggml_vk_guess_matmul_id_pipeline(ctx, mmp, m, n, true, src0_type, src1_type)->align; } static void ggml_vk_matmul_id( @@ -9218,14 +10052,14 @@ static void ggml_vk_matmul_id( uint32_t m, uint32_t n, uint32_t k, uint32_t stride_a, uint32_t stride_b, uint32_t stride_d, uint32_t batch_stride_a, uint32_t batch_stride_b, uint32_t batch_stride_d, uint32_t n_as, uint32_t nei0, uint32_t nei1, uint32_t nbi1, uint32_t ne11, - bool hoist_row_ids) { + bool hoist_row_ids, const vk_subbuffer & fused_scale, uint32_t fusion_flags) { VK_LOG_DEBUG("ggml_vk_matmul_id(a: (" << a.buffer->buffer << ", " << a.offset << ", " << a.size << "), b: (" << b.buffer->buffer << ", " << b.offset << ", " << b.size << "), d: (" << d.buffer->buffer << ", " << d.offset << ", " << d.size << "), ids: (" << ids.buffer->buffer << ", " << ids.offset << ", " << ids.size << "), expert_count: (" << expert_count_buf.buffer->buffer << ", " << expert_count_buf.offset << ", " << expert_count_buf.size << "), " << "m: " << m << ", n: " << n << ", k: " << k << ", stride_a: " << stride_a << ", stride_b: " << stride_b << ", stride_d: " << stride_d << ", " << "batch_stride_a: " << batch_stride_a << ", batch_stride_b: " << batch_stride_b << ", batch_stride_d: " << batch_stride_d << ", " << "n_as: " << n_as << ", nei0: " << nei0 << ", nei1: " << nei1 << ", nbi1: " << nbi1 << ", ne11: " << ne11 << ")"); const vk_mat_mat_id_push_constants pc = { m, n, k, stride_a, stride_b, stride_d, batch_stride_a, batch_stride_b, batch_stride_d, - nei0, nei1, nbi1, ne11, n_as, uint32_t(hoist_row_ids) }; - ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { a, b, d, ids, expert_count_buf }, pc, { m, nei1, n_as }); + nei0, nei1, nbi1, ne11, n_as, uint32_t(hoist_row_ids), fusion_flags }; + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { a, b, d, ids, expert_count_buf, fused_scale }, pc, { m, nei1, n_as }); } static bool ggml_vk_dim01_contiguous(const ggml_tensor * tensor) { @@ -9329,6 +10163,11 @@ static vk_pipeline ggml_vk_get_cpy_pipeline(ggml_backend_vk_context * ctx, const case GGML_TYPE_Q5_1: case GGML_TYPE_Q8_0: case GGML_TYPE_IQ4_NL: + case GGML_TYPE_Q4_0_ROCMFP4: + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + case GGML_TYPE_Q3_0_ROCMFPX: + case GGML_TYPE_Q6_0_ROCMFPX: + case GGML_TYPE_Q8_0_ROCMFPX: return ctx->device->pipeline_cpy_f32_quant[to]; default: break; @@ -9345,6 +10184,11 @@ static vk_pipeline ggml_vk_get_cpy_pipeline(ggml_backend_vk_context * ctx, const case GGML_TYPE_Q5_1: case GGML_TYPE_Q8_0: case GGML_TYPE_IQ4_NL: + case GGML_TYPE_Q4_0_ROCMFP4: + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + case GGML_TYPE_Q3_0_ROCMFPX: + case GGML_TYPE_Q6_0_ROCMFPX: + case GGML_TYPE_Q8_0_ROCMFPX: return ctx->device->pipeline_cpy_quant_f32[src->type]; default: break; @@ -9518,14 +10362,29 @@ static void ggml_vk_mul_mat_q_f16(ggml_backend_vk_context * ctx, vk_context& sub src1_uma = d_Qy != nullptr; } - // TODO: Clean up this logic to pick src1 type by capability // Reformat and convert to fp16 if non-contiguous, or for coopmat2 for better perf const bool x_non_contig = (ctx->device->coopmat2 && src0->type == GGML_TYPE_F32) || !ggml_vk_dim01_contiguous(src0); - const bool y_non_contig = (ctx->device->coopmat2 && src1->type == GGML_TYPE_F32) || - // coopmat1: force f32->f16 conversion so the f16 B-type quant pipeline is used. - (ctx->device->coopmat_support && !ctx->device->coopmat2 && - ggml_is_quantized(src0->type) && src1->type == GGML_TYPE_F32) || + // Route quantized MUL_MAT through the f16-B kernels. Treating contiguous f32 B as + // y_non_contig reuses the convert-to-prealloc_y plumbing, like coopmat2 does. + // auto mode restricts to ne10 == 5120: the only width measured to gain. Narrow on purpose, + // so widths measured as losses (3584 dense, 2048 MoE) cannot trigger it. + const bool dense_f16b = ggml_vk_dense_f16b_enabled() && + (ggml_vk_dense_f16b_mode() == 1 || ne10 == 5120) && + ctx->device->coopmat_support && !ctx->device->coopmat2 && + ggml_is_quantized(src0->type) && src1->type == GGML_TYPE_F32 && + !(ctx->device->pipeline_dequant_mul_mat_mat_f16[src0->type].f16acc->is_empty() && + ctx->device->pipeline_dequant_mul_mat_mat_f16[src0->type].f32acc->is_empty()); + if (dense_f16b) { + static bool dense_f16b_logged = false; + if (!dense_f16b_logged) { + dense_f16b_logged = true; + fprintf(stderr, "ggml_vulkan: MUL_MAT f16-B path engaged (GGML_VK_DENSE_F16B)\n"); + } + } + + const bool y_non_contig = dense_f16b || + (ctx->device->coopmat2 && src1->type == GGML_TYPE_F32) || (src0->type == GGML_TYPE_BF16 && src1->type != GGML_TYPE_BF16) || !ggml_vk_dim01_contiguous(src1); @@ -9537,31 +10396,31 @@ static void ggml_vk_mul_mat_q_f16(ggml_backend_vk_context * ctx, vk_context& sub bool quantize_y = ctx->device->integer_dot_product && src1->type == GGML_TYPE_F32 && ggml_is_contiguous(src1) && !y_non_contig && (ne11 * ne10) % 4 == 0; // Check for mmq first - const std::vector* mmp_map = quantize_y ? ggml_vk_get_mul_mat_mat_pipeline_map(ctx, src0->type, GGML_TYPE_Q8_1, (ggml_prec)dst->op_params[0]) : nullptr; + vk_matmul_pipeline mmp = quantize_y ? ggml_vk_get_mul_mat_mat_pipeline(ctx, src0->type, GGML_TYPE_Q8_1, (ggml_prec)dst->op_params[0]) : nullptr; - if (mmp_map == nullptr) { + if (mmp == nullptr) { // Fall back to f16 dequant mul mat - mmp_map = ggml_vk_get_mul_mat_mat_pipeline_map(ctx, src0->type, y_non_contig ? f16_type : src1->type, (ggml_prec)dst->op_params[0]); + mmp = ggml_vk_get_mul_mat_mat_pipeline(ctx, src0->type, y_non_contig ? f16_type : src1->type, (ggml_prec)dst->op_params[0]); quantize_y = false; } - const bool qx_needs_dequant = mmp_map == nullptr || x_non_contig; + const bool qx_needs_dequant = mmp == nullptr || x_non_contig; const bool qy_needs_dequant = !quantize_y && ((src1->type != f16_type && !y_f32_kernel) || y_non_contig); if (qx_needs_dequant) { // Fall back to dequant + f16 mulmat - mmp_map = ggml_vk_get_mul_mat_mat_pipeline_map(ctx, f16_type, y_f32_kernel ? GGML_TYPE_F32 : f16_type, (ggml_prec)dst->op_params[0]); + mmp = ggml_vk_get_mul_mat_mat_pipeline(ctx, f16_type, y_f32_kernel ? GGML_TYPE_F32 : f16_type, (ggml_prec)dst->op_params[0]); } // Not implemented GGML_ASSERT(y_non_contig || !qy_needs_dequant); // NOLINT - GGML_ASSERT(mmp_map != nullptr); + const ggml_type effective_src1_type = quantize_y ? GGML_TYPE_Q8_1 : (y_f32_kernel ? GGML_TYPE_F32 : src1->type); - const uint32_t kpad = quantize_y ? 0 : ggml_vk_align_size(ne10, ggml_vk_guess_matmul_pipeline_align_map(ctx, *mmp_map, ne01, ne11, false)); + const uint32_t kpad = quantize_y ? 0 : ggml_vk_align_size(ne10, ggml_vk_guess_matmul_pipeline_align(ctx, mmp, ne01, ne11, qx_needs_dequant ? f16_type : src0->type, effective_src1_type)); const bool aligned = !quantize_y && ne10 == kpad && ne01 > 8 && ne11 > 8; - vk_pipeline pipeline = ggml_vk_guess_matmul_pipeline_map(ctx, *mmp_map, ne01, ne11, aligned, false); + vk_pipeline pipeline = ggml_vk_guess_matmul_pipeline(ctx, mmp, ne01, ne11, aligned, qx_needs_dequant ? f16_type : src0->type, effective_src1_type); if (ggml_nbytes(src0) > ctx->device->properties.limits.maxStorageBufferRange) { pipeline = ggml_vk_get_64b_indexing_pipeline(ctx, pipeline); @@ -9760,6 +10619,12 @@ static bool ggml_vk_should_use_mmvq(const vk_device& device, uint32_t m, uint32_ if (src0_type == GGML_TYPE_Q6_K && !mmvq_q6) { return false; } + // q8_0 with two columns through the q8_1 integer-dot path: 272-336 us against 17.5 us + // for the f32 path on the same 320x10240 (Strix Halo / RADV, GGML_VK_PERF_LOGGER), + // 824 vs 146 us on 12288x2560. Every other column count of q8_0 is fine on it. + if (src0_type == GGML_TYPE_Q8_0 && n == 2 && device->vendor_id == VK_VENDOR_ID_AMD) { + return false; + } // MMVQ is generally good for batches if (n > 1) { @@ -9828,10 +10693,10 @@ static bool ggml_vk_should_use_mmvq(const vk_device& device, uint32_t m, uint32_ GGML_UNUSED(m); } -static void ggml_vk_mul_mat_vec_q_f16(ggml_backend_vk_context * ctx, vk_context& subctx, const struct ggml_cgraph * cgraph, int node_idx, bool swap_inputs = false) { +static void ggml_vk_mul_mat_vec_q_f16_cols(ggml_backend_vk_context * ctx, vk_context& subctx, const struct ggml_cgraph * cgraph, int node_idx, uint32_t col0, uint32_t ncols, bool allow_quantize_y) { ggml_tensor * dst = cgraph->nodes[node_idx]; - const ggml_tensor * src0 = dst->src[swap_inputs ? 1 : 0]; - const ggml_tensor * src1 = dst->src[swap_inputs ? 0 : 1]; + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; VK_LOG_DEBUG("ggml_vk_mul_mat_vec_q_f16((" << src0 << ", name=" << src0->name << ", type=" << src0->type << ", ne0=" << src0->ne[0] << ", ne1=" << src0->ne[1] << ", ne2=" << src0->ne[2] << ", ne3=" << src0->ne[3] << ", nb0=" << src0->nb[0] << ", nb1=" << src0->nb[1] << ", nb2=" << src0->nb[2] << ", nb3=" << src0->nb[3]; std::cerr << "), (" << src1 << ", name=" << src1->name << ", type=" << src1->type << ", ne0=" << src1->ne[0] << ", ne1=" << src1->ne[1] << ", ne2=" << src1->ne[2] << ", ne3=" << src1->ne[3] << ", nb0=" << src1->nb[0] << ", nb1=" << src1->nb[1] << ", nb2=" << src1->nb[2] << ", nb3=" << src1->nb[3]; @@ -9846,12 +10711,12 @@ static void ggml_vk_mul_mat_vec_q_f16(ggml_backend_vk_context * ctx, vk_context& const uint64_t ne03 = src0->ne[3]; const uint64_t ne10 = src1->ne[0]; - const uint64_t ne11 = src1->ne[1]; + const uint64_t ne11 = ncols; // a chunk of src1's columns starting at col0 (see ggml_vk_mul_mat_vec_q_f16) const uint64_t ne12 = src1->ne[2]; const uint64_t ne13 = src1->ne[3]; - const uint64_t ne20 = dst->ne[swap_inputs ? 1 : 0]; - const uint64_t ne21 = dst->ne[swap_inputs ? 0 : 1]; + const uint64_t ne20 = dst->ne[0]; + const uint64_t ne21 = dst->ne[1]; // const uint64_t ne22 = dst->ne[2]; // const uint64_t ne23 = dst->ne[3]; @@ -9867,7 +10732,7 @@ static void ggml_vk_mul_mat_vec_q_f16(ggml_backend_vk_context * ctx, vk_context& const bool y_non_contig = !ggml_vk_dim01_contiguous(src1); const bool f16_f32_kernel = src1->type == GGML_TYPE_F32; - bool quantize_y = ctx->device->integer_dot_product && src1->type == GGML_TYPE_F32 && ggml_is_contiguous(src1) && !y_non_contig && (ne11 * ne10) % 4 == 0 && ggml_vk_should_use_mmvq(ctx->device, ne01, ne11, ne10, src0->type); + bool quantize_y = allow_quantize_y && ctx->device->integer_dot_product && src1->type == GGML_TYPE_F32 && ggml_is_contiguous(src1) && !y_non_contig && (ne11 * ne10) % 4 == 0 && ggml_vk_should_use_mmvq(ctx->device, ne01, ne11, ne10, src0->type); vk_pipeline to_fp16_vk_0 = nullptr; vk_pipeline to_fp16_vk_1 = nullptr; @@ -9946,6 +10811,12 @@ static void ggml_vk_mul_mat_vec_q_f16(ggml_backend_vk_context * ctx, vk_context& vk_subbuffer d_D = ggml_vk_tensor_subbuffer(ctx, cgraph->nodes[node_idx + ctx->num_additional_fused_ops]); vk_subbuffer d_Qx = ggml_vk_tensor_subbuffer(ctx, src0); vk_subbuffer d_Qy = ggml_vk_tensor_subbuffer(ctx, src1); + if (col0 > 0) { + const uint64_t off_y = (uint64_t) col0 * src1->nb[1]; + const uint64_t off_d = (uint64_t) col0 * dst->nb[1]; + d_Qy.offset += off_y; d_Qy.size -= off_y; + d_D.offset += off_d; d_D.size -= off_d; + } vk_subbuffer d_X, d_Y; if (qx_needs_dequant) { @@ -10071,6 +10942,76 @@ static void ggml_vk_mul_mat_vec_q_f16(ggml_backend_vk_context * ctx, vk_context& } } + +// Column chunking for batched mat-vec. On Strix Halo (RADV, Mesa 26.1) the NUM_COLS +// shader variants for 3, 5 and 6 columns -- and 4 for q4_K -- run several times slower +// per call than 1, 2 and 4. Measured with GGML_VK_PERF_LOGGER on Qwen3.8-Flash-Next: +// q8_0 10240x2560 122 us at n=1, 127 at 2, 129 at 4, but 734 at 3, 1177 at 5, 360 at 6; +// the q6_K LM head (248320x2560) 2.24 ms at n=1, 2 and 4, but 12.8 at 3 and 18.5 at 5; +// q4_K 12288x2560 85 us at n=1, 180 at 4, 290 at 6. Speculative verification runs at +// n=2..6, so such batches are dispatched as chunks of column counts measured to scale, +// at the cost of re-reading the weights once per chunk. Chunks never take the q8_1 +// integer-dot path (its prealloc cache is keyed by tensor, not by column; and that path +// has the same pathology). GGML_VK_MMV_NO_SPLIT=1 restores the single dispatch. +// Only q8_0 and q6_K are chunked: those are the types where the slow variants were +// measured. The K-quants were chunked to 2 columns without a measurement, and on a dense +// Qwen3.8-27B UD-Q4_K_XL (RADV, Mesa 26.1.7) that made things worse, not better: pp3/pp5/pp8 +// 23.0/28.5/28.8 t/s chunked against 32.9/43.1/47.4 t/s with the single dispatch, which +// halved MTP decode (issue #25). Leave the other types on upstream's path until measured. +static bool ggml_vk_mmv_chunk_type(ggml_type type) { + return type == GGML_TYPE_Q8_0 || type == GGML_TYPE_Q6_K; +} + +static uint32_t ggml_vk_mmv_good_cols(ggml_type type, uint32_t n) { + if (!ggml_vk_mmv_chunk_type(type)) { + return n; + } + return n >= 4 ? 4 : (n >= 2 ? 2 : 1); +} + +// Batches of 9..16 columns of a quantized weight also go through chunked mat-vec: the +// mat-mat kernels are far slower at these widths on this GPU (q8_0 320x10240 at n=21: +// 756 us against 17.5 us per mat-vec column), and this is where several concurrent +// slots with speculative verification land. Mirrors the conditions the chunk path needs. +static constexpr uint32_t mul_mat_vec_chunk_max_cols = 16; +static bool ggml_vk_mmv_can_chunk(ggml_backend_vk_context * ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst) { + static const bool no_split = getenv("GGML_VK_MMV_NO_SPLIT") != nullptr; + + // Chunking cannot reuse the whole-tensor Q8_1 activation cache with a column offset. + // Keep the unsplit integer-dot path rather than silently changing its arithmetic. + const bool would_quantize_y = ctx->device->integer_dot_product && src1->type == GGML_TYPE_F32 && + ggml_is_contiguous(src1) && (src1->ne[1] * src1->ne[0]) % 4 == 0 && + ggml_vk_should_use_mmvq(ctx->device, src0->ne[1], src1->ne[1], src1->ne[0], src0->type); + return !no_split && ggml_vk_mmv_chunk_type(src0->type) && + !would_quantize_y && + dst->ne[1] > mul_mat_vec_max_cols && dst->ne[1] <= mul_mat_vec_chunk_max_cols && + src1->ne[2] * src1->ne[3] == 1 && src1->type == GGML_TYPE_F32 && + ggml_is_contiguous(src1) && ggml_is_contiguous(dst) && ctx->num_additional_fused_ops == 0; +} + +static void ggml_vk_mul_mat_vec_q_f16(ggml_backend_vk_context * ctx, vk_context& subctx, const struct ggml_cgraph * cgraph, int node_idx) { + ggml_tensor * dst = cgraph->nodes[node_idx]; + const ggml_tensor * src0 = dst->src[0]; + const ggml_tensor * src1 = dst->src[1]; + const uint32_t ne11 = (uint32_t) src1->ne[1]; + + static const bool no_split = getenv("GGML_VK_MMV_NO_SPLIT") != nullptr; + const bool plain = ne11 > 1 && src1->ne[2] * src1->ne[3] == 1 && + src1->type == GGML_TYPE_F32 && ggml_is_contiguous(src1) && ggml_is_contiguous(dst) && + ctx->num_additional_fused_ops == 0; + GGML_ASSERT(ne11 <= mul_mat_vec_max_cols || plain); + if (!no_split && plain && std::min(ggml_vk_mmv_good_cols(src0->type, ne11), mul_mat_vec_max_cols) < ne11) { + uint32_t col0 = 0; + while (col0 < ne11) { + const uint32_t n = std::min(ggml_vk_mmv_good_cols(src0->type, ne11 - col0), mul_mat_vec_max_cols); + ggml_vk_mul_mat_vec_q_f16_cols(ctx, subctx, cgraph, node_idx, col0, n, /*allow_quantize_y=*/ false); + col0 += n; + } + return; + } + ggml_vk_mul_mat_vec_q_f16_cols(ctx, subctx, cgraph, node_idx, 0, ne11, /*allow_quantize_y=*/ true); +} + static void ggml_vk_mul_mat_vec_p021_f16_f32(ggml_backend_vk_context * ctx, vk_context& subctx, const struct ggml_cgraph * cgraph, int node_idx) { ggml_tensor * dst = cgraph->nodes[node_idx]; const ggml_tensor * src0 = dst->src[0]; @@ -10319,98 +11260,6 @@ static void ggml_vk_fwht(ggml_backend_vk_context * ctx, vk_context& subctx, cons ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { src_buf, dst_buf }, pc, { workgroups_x, 1, 1 }); } -static uint32_t ggml_vk_nb_elem(const ggml_tensor * t, int i) { - return (uint32_t)(t->nb[i] / ggml_type_size(t->type)); -} - -static void ggml_vk_dsv4_hc_comb(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * mixes, const ggml_tensor * scale, const ggml_tensor * base, ggml_tensor * dst) { - VK_LOG_DEBUG("ggml_vk_dsv4_hc_comb(" << mixes << ", " << scale << ", " << base << ", " << dst << ")"); - - vk_pipeline pipeline = ctx->device->pipeline_dsv4_hc_comb_f32; - GGML_ASSERT(pipeline != nullptr); - - const uint32_t n_tokens = (uint32_t)mixes->ne[1]; - - ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); - - const vk_subbuffer mixes_buf = ggml_vk_tensor_subbuffer(ctx, mixes, true); - const vk_subbuffer scale_buf = ggml_vk_tensor_subbuffer(ctx, scale, true); - const vk_subbuffer base_buf = ggml_vk_tensor_subbuffer(ctx, base, true); - const vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst, true); - - vk_op_dsv4_hc_comb_push_constants pc = { - n_tokens, - ggml_vk_nb_elem(mixes, 0), ggml_vk_nb_elem(mixes, 1), - ggml_vk_nb_elem(scale, 0), - ggml_vk_nb_elem(base, 0), - ggml_vk_nb_elem(dst, 0), ggml_vk_nb_elem(dst, 1), ggml_vk_nb_elem(dst, 2), - 0, 0, 0, 0, - ggml_get_op_params_f32(dst, 0), - (uint32_t)ggml_get_op_params_i32(dst, 1), - }; - init_pushconst_tensor_offsets(ctx, pc, mixes, scale, base, nullptr, dst); - - ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { mixes_buf, scale_buf, base_buf, dst_buf }, pc, { n_tokens, 1, 1 }); -} - -static void ggml_vk_dsv4_hc_pre(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * x, const ggml_tensor * weights, ggml_tensor * dst) { - VK_LOG_DEBUG("ggml_vk_dsv4_hc_pre(" << x << ", " << weights << ", " << dst << ")"); - - vk_pipeline pipeline = ctx->device->pipeline_dsv4_hc_pre_f32; - GGML_ASSERT(pipeline != nullptr); - - const uint32_t n_embd = (uint32_t)x->ne[0]; - const uint32_t n_tokens = (uint32_t)x->ne[2]; - - ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); - - const vk_subbuffer x_buf = ggml_vk_tensor_subbuffer(ctx, x, true); - const vk_subbuffer w_buf = ggml_vk_tensor_subbuffer(ctx, weights, true); - const vk_subbuffer d_buf = ggml_vk_tensor_subbuffer(ctx, dst, true); - - vk_op_dsv4_hc_pre_push_constants pc = { - n_embd, n_tokens, - ggml_vk_nb_elem(x, 0), ggml_vk_nb_elem(x, 1), ggml_vk_nb_elem(x, 2), - ggml_vk_nb_elem(weights, 0), ggml_vk_nb_elem(weights, 1), - ggml_vk_nb_elem(dst, 0), ggml_vk_nb_elem(dst, 1), - 0, 0, 0, - }; - init_pushconst_tensor_offsets(ctx, pc, x, weights, nullptr, nullptr, dst); - - ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { x_buf, w_buf, d_buf }, pc, { n_embd, n_tokens, 1 }); -} - -static void ggml_vk_dsv4_hc_post(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * x, const ggml_tensor * residual, const ggml_tensor * post, const ggml_tensor * comb, ggml_tensor * dst) { - VK_LOG_DEBUG("ggml_vk_dsv4_hc_post(" << x << ", " << residual << ", " << post << ", " << comb << ", " << dst << ")"); - - vk_pipeline pipeline = ctx->device->pipeline_dsv4_hc_post_f32; - GGML_ASSERT(pipeline != nullptr); - - const uint32_t n_embd = (uint32_t)x->ne[0]; - const uint32_t n_tokens = (uint32_t)x->ne[1]; - - ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); - - const vk_subbuffer x_buf = ggml_vk_tensor_subbuffer(ctx, x, true); - const vk_subbuffer r_buf = ggml_vk_tensor_subbuffer(ctx, residual, true); - const vk_subbuffer p_buf = ggml_vk_tensor_subbuffer(ctx, post, true); - const vk_subbuffer c_buf = ggml_vk_tensor_subbuffer(ctx, comb, true); - const vk_subbuffer d_buf = ggml_vk_tensor_subbuffer(ctx, dst, true); - - vk_op_dsv4_hc_post_push_constants pc = { - n_embd, n_tokens, - ggml_vk_nb_elem(x, 0), ggml_vk_nb_elem(x, 1), - ggml_vk_nb_elem(residual, 0), ggml_vk_nb_elem(residual, 1), ggml_vk_nb_elem(residual, 2), - ggml_vk_nb_elem(post, 0), ggml_vk_nb_elem(post, 1), - ggml_vk_nb_elem(comb, 0), ggml_vk_nb_elem(comb, 1), ggml_vk_nb_elem(comb, 2), - ggml_vk_nb_elem(dst, 0), ggml_vk_nb_elem(dst, 1), ggml_vk_nb_elem(dst, 2), - 0, 0, 0, 0, 0, - }; - init_pushconst_tensor_offsets(ctx, pc, x, residual, post, comb, dst); - - ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, { x_buf, r_buf, p_buf, c_buf, d_buf }, pc, { n_embd, n_tokens, 1 }); -} - static void ggml_vk_mul_mat(ggml_backend_vk_context * ctx, vk_context& subctx, const struct ggml_cgraph * cgraph, int node_idx) { ggml_tensor * dst = cgraph->nodes[node_idx]; ggml_tensor * src0 = dst->src[0]; @@ -10465,19 +11314,10 @@ static void ggml_vk_mul_mat(ggml_backend_vk_context * ctx, vk_context& subctx, c src0->ne[1] <= ctx->device->properties.limits.maxComputeWorkGroupCount[1] && src1->ne[2] <= ctx->device->properties.limits.maxComputeWorkGroupCount[2]) { ggml_vk_mul_mat_vec_nc_f16_f32(ctx, subctx, cgraph, node_idx); - // With one output row, B^T*A has the same flat output as A^T*B. - } else if (ctx->num_additional_fused_ops == 0 && - (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16) && - (src1->type == GGML_TYPE_F32 || src1->type == GGML_TYPE_F16 || src1->type == GGML_TYPE_BF16 || ggml_is_quantized(src1->type)) && - dst->ne[0] == 1 && dst->ne[1] > mul_mat_vec_max_cols && - src0->ne[2] == 1 && src0->ne[3] == 1 && - src1->ne[2] == 1 && src1->ne[3] == 1 && - ggml_is_contiguous(src0) && ggml_is_contiguous(src1) && ggml_is_contiguous(dst) && - get_misalign_bytes(ctx, src0) == 0 && get_misalign_bytes(ctx, src1) == 0 && get_misalign_bytes(ctx, dst) == 0) { - ggml_vk_mul_mat_vec_q_f16(ctx, subctx, cgraph, node_idx, true); // mul_mat_vec supports batching ne12*ne13 when ne11==1, or treating ne11 as the batch size (up to four) // when ne12 and ne13 are one. - } else if ((dst->ne[1] == 1 || (dst->ne[1] <= mul_mat_vec_max_cols && src1->ne[2] * src1->ne[3] == 1)) && + } else if ((dst->ne[1] == 1 || (dst->ne[1] <= mul_mat_vec_max_cols && src1->ne[2] * src1->ne[3] == 1) || + ggml_vk_mmv_can_chunk(ctx, src0, src1, dst)) && (src0->type == GGML_TYPE_F32 || src0->type == GGML_TYPE_F16 || src0->type == GGML_TYPE_BF16 || ggml_is_quantized(src0->type))) { ggml_vk_mul_mat_vec_q_f16(ctx, subctx, cgraph, node_idx); } else { @@ -10485,7 +11325,7 @@ static void ggml_vk_mul_mat(ggml_backend_vk_context * ctx, vk_context& subctx, c } } -static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst) { +static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * ids, ggml_tensor * dst, const ggml_tensor * fused_scale = nullptr, ggml_tensor * fused_dst = nullptr) { VK_LOG_DEBUG("ggml_vk_mul_mat_id_q_f16((" << src0 << ", name=" << src0->name << ", type=" << src0->type << ", ne0=" << src0->ne[0] << ", ne1=" << src0->ne[1] << ", ne2=" << src0->ne[2] << ", ne3=" << src0->ne[3] << ", nb0=" << src0->nb[0] << ", nb1=" << src0->nb[1] << ", nb2=" << src0->nb[2] << ", nb3=" << src0->nb[3]; std::cerr << "), (" << src1 << ", name=" << src1->name << ", type=" << src1->type << ", ne0=" << src1->ne[0] << ", ne1=" << src1->ne[1] << ", ne2=" << src1->ne[2] << ", ne3=" << src1->ne[3] << ", nb0=" << src1->nb[0] << ", nb1=" << src1->nb[1] << ", nb2=" << src1->nb[2] << ", nb3=" << src1->nb[3]; std::cerr << "), (" << ids << ", name=" << ids->name << ", type=" << ids->type << ", ne0=" << ids->ne[0] << ", ne1=" << ids->ne[1] << ", ne2=" << ids->ne[2] << ", ne3=" << ids->ne[3] << ", nb0=" << ids->nb[0] << ", nb1=" << ids->nb[1] << ", nb2=" << ids->nb[2] << ", nb3=" << ids->nb[3]; @@ -10523,7 +11363,9 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& hoisted_row_id_words * sizeof(uint32_t) <= ctx->device->properties.limits.maxStorageBufferRange; - ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)dst->buffer->context; + // When the following MUL is fused in, write the scaled result straight to its destination. + const ggml_tensor * out_dst = fused_dst ? fused_dst : dst; + ggml_backend_vk_buffer_context * dst_buf_ctx = (ggml_backend_vk_buffer_context *)out_dst->buffer->context; ggml_backend_vk_buffer_context * src0_buf_ctx = (ggml_backend_vk_buffer_context *)src0->buffer->context; ggml_backend_vk_buffer_context * src1_buf_ctx = (ggml_backend_vk_buffer_context *)src1->buffer->context; ggml_backend_vk_buffer_context * ids_buf_ctx = (ggml_backend_vk_buffer_context *)ids->buffer->context; @@ -10569,12 +11411,28 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& #else const bool y_decode_vector_staging = false; #endif + // EXPERIMENT (GGML_VK_MMID_F16B=1): route quantized MUL_MAT_ID through the f16-B + // kernels. Treating contiguous f32 B as y_non_contig reuses the existing + // convert-to-prealloc_y plumbing (to_fp16_vk_1), exactly like coopmat2 does. + // Gated on coopmat_support because the f16b pipelines are only created there. + const bool mmid_f16b = ggml_vk_mmid_f16b_enabled() && + ctx->device->coopmat_support && !ctx->device->coopmat2 && + ggml_is_quantized(src0->type) && src1->type == GGML_TYPE_F32 && + // only take the f16-B path if a pipeline exists for this src0 type + // (e.g. Q2_0 has none); otherwise fall through to the normal f32-B path. + !(ctx->device->pipeline_dequant_mul_mat_mat_id_f16b[src0->type].f16acc->is_empty() && + ctx->device->pipeline_dequant_mul_mat_mat_id_f16b[src0->type].f32acc->is_empty()); + if (mmid_f16b) { + static bool mmid_f16b_logged = false; + if (!mmid_f16b_logged) { + mmid_f16b_logged = true; + fprintf(stderr, "ggml_vulkan: MUL_MAT_ID f16-B path engaged (GGML_VK_MMID_F16B)\n"); + } + } + const bool y_non_contig = y_decode_vector_staging || + mmid_f16b || (ctx->device->coopmat2 && src1->type == GGML_TYPE_F32) || - // Intel coopmat1: force f32->f16 conversion so the f16 B-type quant pipeline is used. - (ctx->device->coopmat_support && !ctx->device->coopmat2 && - ctx->device->vendor_id == VK_VENDOR_ID_INTEL && - ggml_is_quantized(src0->type) && src1->type == GGML_TYPE_F32) || (src0->type == GGML_TYPE_BF16 && src1->type != GGML_TYPE_BF16) || !ggml_vk_dim01_contiguous(src1); @@ -10583,22 +11441,36 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& bool quantize_y = ctx->device->integer_dot_product && src1->type == GGML_TYPE_F32 && ggml_is_contiguous(src1) && !y_non_contig && (ne11 * ne10) % 4 == 0; // Check for mmq first - const std::vector* mmp_map = quantize_y ? ggml_vk_get_mul_mat_mat_pipeline_map(ctx, src0->type, GGML_TYPE_Q8_1, (ggml_prec)dst->op_params[0], true) : nullptr; + vk_matmul_pipeline mmp = quantize_y ? ggml_vk_get_mul_mat_mat_id_pipeline(ctx, src0->type, GGML_TYPE_Q8_1, (ggml_prec)dst->op_params[0]) : nullptr; - if (mmp_map == nullptr) { + if (mmp == nullptr) { // Fall back to f16 dequant mul mat - mmp_map = ggml_vk_get_mul_mat_mat_pipeline_map(ctx, src0->type, y_non_contig ? f16_type : src1->type, (ggml_prec)dst->op_params[0], true); + mmp = ggml_vk_get_mul_mat_mat_id_pipeline(ctx, src0->type, y_non_contig ? f16_type : src1->type, (ggml_prec)dst->op_params[0]); quantize_y = false; } - const bool qx_needs_dequant = mmp_map == nullptr || x_non_contig; + const bool qx_needs_dequant = mmp == nullptr || x_non_contig; bool qy_needs_dequant = !quantize_y && ((src1->type != f16_type && !y_f32_kernel) || y_non_contig); if (qx_needs_dequant) { // Fall back to dequant + f16 mulmat - mmp_map = ggml_vk_get_mul_mat_mat_pipeline_map(ctx, f16_type, y_f32_kernel ? GGML_TYPE_F32 : f16_type, (ggml_prec)dst->op_params[0], true); + mmp = ggml_vk_get_mul_mat_mat_id_pipeline(ctx, f16_type, y_f32_kernel ? GGML_TYPE_F32 : f16_type, (ggml_prec)dst->op_params[0]); } + const ggml_type effective_src1_type = quantize_y ? GGML_TYPE_Q8_1 : (y_f32_kernel ? GGML_TYPE_F32 : src1->type); + + // EXPERIMENT (GGML_VK_MMID_SMALLN=1): select the matmul tile by the EXPECTED PER-EXPERT + // token count rather than the whole batch. With E experts and nei0 active per token, each + // expert sees ~nei1*nei0/E rows; selecting by aggregate nei1 always picks the widest tile + // and leaves most N-lanes empty at MoE prefill (measured: MUL_MAT_ID at ~20% of dense + // matmul efficiency, ~78% of MoE prefill time on Qwen3.6-35B-A3B). + uint32_t n_for_tile = (uint32_t)nei1; + static const char * mmid_smalln_env = getenv("GGML_VK_MMID_SMALLN"); + if (!(mmid_smalln_env && atoi(mmid_smalln_env) == 0) && ne02 > 1) { // on by default; =0 disables + n_for_tile = std::max(1u, (uint32_t)((nei1 * nei0 + ne02 - 1) / ne02)); + } + + const uint32_t kpad = quantize_y ? 0 : ggml_vk_align_size(ne10, ggml_vk_guess_matmul_id_pipeline_align(ctx, mmp, ne01, n_for_tile, qx_needs_dequant ? f16_type : src0->type, effective_src1_type)); // Coopmat2 MUL_MAT_ID BK specialization constants in ggml_vk_load_shaders are at most 64. const uint32_t y_staged_row_stride = ctx->device->coopmat2 && !quantize_y ? ggml_vk_align_size(ne10, 64) : ne10; const bool y_needs_k_padding = ne10 != y_staged_row_stride; @@ -10607,13 +11479,21 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& // Not implemented GGML_ASSERT(y_needs_reformat || !qy_needs_dequant); // NOLINT - - GGML_ASSERT(mmp_map != nullptr); - - const uint32_t kpad = quantize_y ? 0 : ggml_vk_align_size(ne10, ggml_vk_guess_matmul_pipeline_align_map(ctx, *mmp_map, ne01, nei1, true)); const bool aligned = !quantize_y && ne10 == kpad && ne01 > 8 && nei1 > 8; - vk_pipeline pipeline = ggml_vk_guess_matmul_pipeline_map(ctx, *mmp_map, ne01, nei1, aligned, true); + vk_pipeline pipeline = ggml_vk_guess_matmul_id_pipeline(ctx, mmp, ne01, n_for_tile, aligned, qx_needs_dequant ? f16_type : src0->type, effective_src1_type); + + // PROBE (GGML_VK_MMID_PROBE=1): which mmid tile actually runs, and with how many threads. + static const char * mmid_probe_env = getenv("GGML_VK_MMID_PROBE"); + if (mmid_probe_env && atoi(mmid_probe_env) != 0) { + static std::set seen; + std::string key = pipeline->name + ":" + std::to_string(n_for_tile); + if (seen.insert(key).second) { + fprintf(stderr, "ggml_vulkan: mmid pipeline=%s n_for_tile=%u m=%u wg=(%u,%u,%u)\n", + pipeline->name.c_str(), n_for_tile, (uint32_t)ne01, + pipeline->wg_denoms[0], pipeline->wg_denoms[1], pipeline->wg_denoms[2]); + } + } if (ggml_nbytes(src0) > ctx->device->properties.limits.maxStorageBufferRange) { pipeline = ggml_vk_get_64b_indexing_pipeline(ctx, pipeline); @@ -10705,7 +11585,7 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& } vk_buffer d_D = dst_buf_ctx->dev_buffer; - const uint64_t d_buf_offset = vk_tensor_offset(dst) + dst->view_offs; + const uint64_t d_buf_offset = vk_tensor_offset(out_dst) + out_dst->view_offs; GGML_ASSERT(d_D != nullptr); vk_buffer d_X; uint64_t x_buf_offset = 0; @@ -10841,7 +11721,9 @@ static void ggml_vk_mul_mat_id_q_f16(ggml_backend_vk_context * ctx, vk_context& { d_D, d_buf_offset, d_sz }, { d_ids, ids_buf_offset, ids_sz }, expert_count_buf, ne01, ne21, ne10, ne10, stride_b_y, ne01, stride_batch_x, stride_batch_y, ne20*ne21, - n_as, nei0, nei1, nbi1 / ggml_type_size(ids->type), ne11, hoist_row_ids + n_as, nei0, nei1, nbi1 / ggml_type_size(ids->type), ne11, hoist_row_ids, + fused_scale ? ggml_vk_tensor_subbuffer(ctx, fused_scale) : vk_subbuffer{ d_D, d_buf_offset, d_sz }, + fused_scale ? 1u : 0u ); // NOLINT if (x_non_contig || qx_needs_dequant) { @@ -11105,7 +11987,16 @@ static void ggml_vk_mul_mat_id(ggml_backend_vk_context * ctx, vk_context& subctx if (ggml_vk_use_mul_mat_vec_id(cgraph, node_idx)) { ggml_vk_mul_mat_vec_id_q_f16(ctx, subctx, cgraph, node_idx); } else { - ggml_vk_mul_mat_id_q_f16(ctx, subctx, src0, src1, src2, dst); + // Fused scale epilogue: the MUL's other operand is applied as the matmul writes out, + // and the result goes straight to the MUL's destination. + const ggml_tensor * fused_scale = nullptr; + ggml_tensor * fused_dst = nullptr; + if (ctx->num_additional_fused_ops == 1) { + ggml_tensor * mul = cgraph->nodes[node_idx + 1]; + fused_scale = (mul->src[0] == dst) ? mul->src[1] : mul->src[0]; + fused_dst = mul; + } + ggml_vk_mul_mat_id_q_f16(ctx, subctx, src0, src1, src2, dst, fused_scale, fused_dst); } } @@ -11180,32 +12071,826 @@ static bool ggml_vk_flash_attn_coopmat_shmem_support(const vk_device& device, co // DATA_A_IQ4_NL is compiled into the FA shaders unconditionally, so its shared table is always allocated. const uint32_t iq_shmem = 16 * sizeof(ggml_fp16_t); - const uint32_t qstride = hsk_pad / 4 + 2; - const uint32_t Qf = Br * qstride * f16vec4; + const uint32_t qstride = hsk_pad / 4 + 2; + const uint32_t Qf = Br * qstride * f16vec4; + + const uint32_t psh_stride = Bc / 4 + 2; + const uint32_t Psh = Br * psh_stride * f16vec4; + + const uint32_t sfshstride = (hsk <= 128) ? (Br + 8) : Br; + const uint32_t sfsh = Bc * sfshstride * acctype; + + const uint32_t kvshstride = (params.shmem_staging ? std::max(hsk_pad, hsv_pad) : MatBr) / 4 + 2; + const uint32_t vsh_stride = MatBc / 4 * row_split; + const uint32_t ksh = ((kvshstride >= vsh_stride) ? (Bc * kvshstride) : (Bc * vsh_stride)) * f16vec4; + + // BF16 PVMat accumulator is f32 (no bf16 accumulator support), so pvsh is vec4 (16 bytes) + const uint32_t pvsh_elem_size = (k_type == GGML_TYPE_BF16) ? 16u : f16vec4; + const uint32_t osh_stride = params.row_split * MatBr / 4; + const uint32_t pvsh = MatBc * osh_stride * pvsh_elem_size; + + const uint32_t slope = Br * acctype; + + const uint32_t total_size = tmpsh + iq_shmem + Qf + Psh + sfsh + ksh + pvsh + slope; + const bool supported = total_size <= device->properties.limits.maxComputeSharedMemorySize; + + VK_LOG_DEBUG("ggml_vk_flash_attn_coopmat_shmem_support(HSK=" << hsk << ", HSV=" << hsv << ", f32acc=" << f32acc << ", total_size=" << total_size << ", supported=" << supported); + + return supported; +} + +// K/V types the FA shaders can read directly (scalar/coopmat1 select the dequant code via the +// FaTypeK/FaTypeV spec constants; a type outside this list silently reads garbage). Types that +// are FA-supported but not listed here (iq4_nl) are only correct through the dequant-once +// scratch path, so supports_op and the dispatch-time gate must agree on when that path runs. +static bool ggml_vk_fa_kv_native(ggml_type t, bool coopmat2) { + switch (t) { + case GGML_TYPE_F32: + case GGML_TYPE_F16: + case GGML_TYPE_BF16: + case GGML_TYPE_Q4_0: + case GGML_TYPE_Q4_1: + case GGML_TYPE_Q5_0: + case GGML_TYPE_Q5_1: + case GGML_TYPE_Q8_0: + case GGML_TYPE_IQ4_NL: // native FA support since upstream 8161641 + return true; + default: + return false; + } +} + +// Capacity gate for the dequant-once f16 K/V scratch. On discrete devices the scratch can push the +// working set past free VRAM, and the driver then silently pages device-local memory (~15x prefill +// regression measured on an 8 GB card at long context). UMA has no separate pool to overflow. +// +// Gates our own device-local usage against the physical heap size, less a reserve for memory not +// visible in-process. heapBudget is deliberately not the signal: ggml_backend_vk_get_device_memory +// computes heapBudget - heapUsage unsigned, which wraps when the device is oversubscribed. Nor can +// the allocation gate itself: on WDDM vkAllocateMemory only fails near physical heap size, above +// the free-VRAM level where paging starts. The reserve is necessarily conservative because other +// processes' VRAM use is invisible to us. GGML_VK_FA_DEQUANT=0/1 forces the path off/on; +// GGML_VK_FA_DEQUANT_RESERVE_MB overrides the reserve. +static bool ggml_vk_fa_dequant_scratch_fits(ggml_backend_vk_context * ctx, uint64_t scratch_sz) { + const vk_device& device = ctx->device; + + if (device->uma) { + return true; + } + + // Decided once per scratch size, so the decision cannot flip between layers as usage grows. + if (ctx->fa_dequant_gate_sz == scratch_sz) { + return ctx->fa_dequant_gate_fits; + } + + static const uint64_t reserve = [] { + const char * env = getenv("GGML_VK_FA_DEQUANT_RESERVE_MB"); + return (uint64_t)(env ? atoi(env) : 1024) * 1024 * 1024; + }(); + + bool fits = false; + + // Without VK_EXT_memory_budget our usage is unknowable, so leave the path disabled. + if (vk_instance.device_supports_membudget[device->idx]) { + vk::PhysicalDeviceMemoryBudgetPropertiesEXT budgetprops; + vk::PhysicalDeviceMemoryProperties2 memprops = {}; + memprops.pNext = &budgetprops; + device->physical_device.getMemoryProperties2(&memprops); + + uint64_t heap_size = 0; + uint64_t heap_used = 0; + for (uint32_t i = 0; i < memprops.memoryProperties.memoryHeapCount; ++i) { + const vk::MemoryHeap & heap = memprops.memoryProperties.memoryHeaps[i]; + if (heap.flags & vk::MemoryHeapFlagBits::eDeviceLocal) { + heap_size += heap.size; + heap_used += budgetprops.heapUsage[i]; + } + } + // heap_used already covers scratch allocated on a previous ubatch, so counting scratch_sz + // in full is conservative by up to the current scratch size. + fits = heap_size > reserve && heap_used + scratch_sz + reserve <= heap_size; + } + + if (!fits && !ctx->fa_dequant_gate_logged) { + ctx->fa_dequant_gate_logged = true; + GGML_LOG_INFO("ggml_vulkan: flash attention dequant-once disabled: %llu MiB K/V scratch does not fit " + "device-local memory with a %llu MiB reserve. Set GGML_VK_FA_DEQUANT=1 to force it on.\n", + (unsigned long long)(scratch_sz >> 20), (unsigned long long)(reserve >> 20)); + } + + ctx->fa_dequant_gate_sz = scratch_sz; + ctx->fa_dequant_gate_fits = fits; + return fits; +} + +// Close a timestamp interval mid-node so work a handler dispatches before its op is billed +// separately instead of being folded into the op's own time. `name` must be a string literal +// (stored by pointer). No-op unless the perf logger is on in per-op mode. +static void ggml_vk_perf_mark_subop(ggml_backend_vk_context * ctx, vk_context& subctx, const char * name) { + if (!vk_perf_logger_enabled || vk_perf_logger_concurrent || ctx->query_pool == VK_NULL_HANDLE) { + return; + } + if (ctx->query_idx >= (int)ctx->num_queries) { + return; // pool headroom exhausted; drop the mark rather than overflow + } + ctx->query_nodes[ctx->query_idx] = nullptr; + ctx->query_fusion_names[ctx->query_idx] = nullptr; + ctx->query_sub_names[ctx->query_idx] = name; + subctx->s->buffer->buf.writeTimestamp(vk::PipelineStageFlagBits::eAllCommands, ctx->query_pool, ctx->query_idx++); +} + +static bool ggml_vk_flash_attn_top_k(ggml_backend_vk_context * ctx, vk_context & subctx, + const ggml_tensor * q, const ggml_tensor * k, const ggml_tensor * v, + const ggml_tensor * mask, const ggml_tensor * sinks, ggml_tensor * dst) { + const ggml_tensor * top_k = dst->src[5]; + static const char * top_k_env = getenv("GGML_VK_FA_TOPK"); + // The sparse shaders read f16 only. Serve a quantised cache by dequantising it once into + // the f16 scratch, the same pass the dense path uses. K and V are the same tensor here, so + // one pass covers both. Without this a quantised cache declines to dense FA, which costs + // O(kv) where the sparse path costs O(n_kv_raw + n_top_k). + static const char * fa_dequant_env = getenv("GGML_VK_FA_DEQUANT"); + const uint64_t kv_f16_sz = (uint64_t) ggml_nelements(k) * sizeof(ggml_fp16_t); + const bool dequant_kv = top_k && k->type != GGML_TYPE_F16 && + !(fa_dequant_env && fa_dequant_env[0] == '0') && + ctx->device->pipeline_dequant_transpose[k->type] != nullptr && + k->nb[0] == ggml_type_size(k->type) && + ggml_is_contiguously_allocated(k) && + kv_f16_sz <= ctx->device->properties.limits.maxStorageBufferRange && + ggml_vk_fa_dequant_scratch_fits(ctx, kv_f16_sz); + if ((top_k_env && top_k_env[0] == '0') || + !top_k || (!ctx->device->pipeline_flash_attn_top_k_f16 && !ctx->device->pipeline_flash_attn_top_k_cm_f16) || + q->type != GGML_TYPE_F32 || (k->type != GGML_TYPE_F16 && !dequant_kv) || v->type != k->type || + !mask || mask->type != GGML_TYPE_F16 || top_k->type != GGML_TYPE_I32 || + q->ne[0] != 512 || q->ne[1] < 64 || k->ne[0] != 512 || v->ne[0] != 512 || + q->ne[2] != 64 || k->ne[2] != 1 || v->ne[2] != 1 || + q->ne[1] != top_k->ne[1] || q->ne[3] != top_k->ne[3] || + k->ne[1] != v->ne[1] || k->buffer != v->buffer || k->data != v->data || + !ggml_is_contiguous(mask) || !ggml_is_contiguous(top_k)) { + return false; + } + + float scale = 0.0f; + float max_bias = 0.0f; + float logit_softcap = 0.0f; + memcpy(&scale, (const float *) dst->op_params + 0, sizeof(float)); + memcpy(&max_bias, (const float *) dst->op_params + 1, sizeof(float)); + memcpy(&logit_softcap, (const float *) dst->op_params + 2, sizeof(float)); + if (max_bias != 0.0f || logit_softcap != 0.0f) { + return false; + } + + const int32_t n_kv_raw = ggml_get_op_params_i32(dst, 4); + if (n_kv_raw < 0 || n_kv_raw > k->ne[1] || top_k->ne[0] > k->ne[1] - n_kv_raw) { + return false; + } + const int64_t n_kv_active = n_kv_raw + top_k->ne[0]; + // Crossover against ordinary dense FA. The coopmat sparse path (and especially the + // raw/selected split below) is cheap enough to win as soon as any key is pruned; the + // scalar fallback shader is not, and regresses against dense FA until the pruned + // fraction is large, so it keeps its original 3x margin. Equality always uses dense + // FA because nothing is pruned. + const bool have_cm_sparse = ctx->device->pipeline_flash_attn_top_k_cm_f16 != nullptr; + const int64_t min_total_k = have_cm_sparse ? n_kv_active + 1 : 3 * n_kv_active; + if (k->ne[1] < min_total_k) { + return false; + } + + // ---- diagnostic: adjacent-token top-k overlap (GGML_VK_TOPK_OVERLAP=1) --------------- + // Decides whether a deduplicated union is worth building for the small-batch path. A + // speculative draft attends at adjacent POSITIONS, so overlap between adjacent prefill + // tokens is the same quantity and a plain deep prefill samples it thousands of times + // without needing a draft model. Reports |union| / (W * n_top_k) for window sizes W. + // Sampled, host-side, off by default; the readback would be far too costly otherwise. + { + static const char * ov_env = getenv("GGML_VK_TOPK_OVERLAP"); + if (ov_env && ov_env[0] == '1') { + static std::mutex ov_mu; + static uint64_t ov_seen = 0; + static std::map> ov_acc; // W -> {sum ratio, n} + std::lock_guard lock(ov_mu); + if ((ov_seen++ % 32) == 0) { // sample: this is a multi-MB readback + const uint32_t nb = (uint32_t) q->ne[1]; + const uint32_t tk = (uint32_t) top_k->ne[0]; + const int32_t rng = (int32_t) (k->ne[1] - n_kv_raw); + std::vector idx((size_t) nb * tk); + vk_subbuffer sb = ggml_vk_tensor_subbuffer(ctx, top_k); + ggml_vk_buffer_read(sb.buffer, sb.offset, idx.data(), idx.size() * sizeof(int32_t)); + for (uint32_t W : {2u, 4u, 8u}) { + if (nb < W) continue; + double sum = 0.0; uint64_t n = 0; + for (uint32_t t0 = 0; t0 + W <= nb; t0 += W) { // disjoint windows + std::unordered_set u; + uint64_t valid = 0; + for (uint32_t t = t0; t < t0 + W; ++t) { + for (uint32_t j = 0; j < tk; ++j) { + const int32_t v = idx[(size_t) t * tk + j]; + if (v >= 0 && v < rng) { u.insert(v); ++valid; } + } + } + if (valid) { sum += (double) u.size() / (double) valid; ++n; } + } + if (n) { auto & a = ov_acc[W]; a.first += sum; a.second += n; } + } + fprintf(stderr, "[topk-overlap] sample %llu (n_batch=%u, kv=%lld)\n", + (unsigned long long) ov_seen, nb, (long long) k->ne[1]); + for (const auto & e : ov_acc) { + const double r = e.second.first / (double) e.second.second; + const double now = (double) n_kv_raw + (double) e.first * tk; + const double dedup = (double) n_kv_raw + r * (double) e.first * tk; + fprintf(stderr, "[topk-overlap] W=%u union/selected=%.3f (overlap %.1f%%) " + "kv_c %.0f -> %.0f projected small-batch gain %.1f%%\n", + e.first, r, 100.0 * (1.0 - r), now, dedup, 100.0 * (1.0 - dedup / now)); + } + } + } + } + + vk_subbuffer k_buf = ggml_vk_tensor_subbuffer(ctx, k); + if (dequant_kv) { + if (ctx->prealloc_size_x < kv_f16_sz) { + ctx->prealloc_size_x = kv_f16_sz; + ggml_vk_preallocate_buffers(ctx, subctx); + } + vk_pipeline tr_k = ctx->device->pipeline_dequant_transpose[k->type]; + ggml_pipeline_request_descriptor_sets(ctx, tr_k, 1); + if (ctx->prealloc_x_need_sync) { + ggml_vk_sync_buffers(ctx, subctx); + } + const vk_subbuffer k_dst = vk_subbuffer{ ctx->prealloc_x, 0, kv_f16_sz }; + const uint32_t k_nel = (uint32_t) ggml_nelements(k); + const std::vector tr_pc = { (uint32_t) k->ne[0], (uint32_t) k->ne[2], (uint32_t) k->ne[1], 0, k_nel }; + ggml_vk_dispatch_pipeline(ctx, subctx, tr_k, { k_buf, k_dst }, tr_pc, { k_nel, 1, 1 }); + ggml_vk_sync_buffers(ctx, subctx); + ctx->prealloc_x_need_sync = true; + k_buf = k_dst; + ggml_vk_perf_mark_subop(ctx, subctx, "FA_KV_DEQUANT (sub-op)"); + } + // strides of what the shaders actually read: the source cache, or the contiguous + // [HS, KV, n_head_kv, ns] f16 scratch the dequant just wrote + const uint32_t k_stride = dequant_kv ? (uint32_t) k->ne[0] : (uint32_t) (k->nb[1] / sizeof(ggml_fp16_t)); + const uint32_t k_nb3_el = dequant_kv ? (uint32_t) ((uint64_t) k->ne[0] * k->ne[1] * k->ne[2]) + : (uint32_t) (k->nb[3] / sizeof(ggml_fp16_t)); + const uint32_t k_nb2_byte = dequant_kv ? (uint32_t) ((uint64_t) k->ne[0] * k->ne[1] * sizeof(ggml_fp16_t)) + : (uint32_t) k->nb[2]; + const uint32_t k_nb3_byte = dequant_kv ? (uint32_t) (k_nb3_el * sizeof(ggml_fp16_t)) : (uint32_t) k->nb[3]; + + vk_op_flash_attn_top_k_push_constants pc = { + (uint32_t) q->ne[1], (uint32_t) k->ne[1], (uint32_t) n_kv_raw, + (uint32_t) top_k->ne[0], (uint32_t) q->ne[2], + (uint32_t) (q->nb[1] / sizeof(float)), + (uint32_t) (q->nb[2] / sizeof(float)), + (uint32_t) (q->nb[3] / sizeof(float)), + k_stride, + k_nb3_el, + (uint32_t) (mask->nb[1] / sizeof(ggml_fp16_t)), + (uint32_t) (mask->nb[3] / sizeof(ggml_fp16_t)), + (uint32_t) (top_k->nb[1] / sizeof(int32_t)), + (uint32_t) (top_k->nb[3] / sizeof(int32_t)), + (uint32_t) (dst->nb[1] / sizeof(float)), + (uint32_t) (dst->nb[2] / sizeof(float)), + (uint32_t) (dst->nb[3] / sizeof(float)), + scale, sinks != nullptr, 0, + }; + + const vk_subbuffer q_buf = ggml_vk_tensor_subbuffer(ctx, q); + const vk_subbuffer sinks_buf = sinks ? ggml_vk_tensor_subbuffer(ctx, sinks) : q_buf; + static const char * top_k_cm_env = getenv("GGML_VK_FA_TOPK_CM"); + const bool use_cm = (!top_k_cm_env || top_k_cm_env[0] != '0') && ctx->device->pipeline_flash_attn_top_k_cm_f16; + vk_pipeline pipeline = use_cm ? ctx->device->pipeline_flash_attn_top_k_cm_f16 : ctx->device->pipeline_flash_attn_top_k_f16; + + static const char * top_k_split_env = getenv("GGML_VK_FA_TOPK_SPLIT"); + const uint32_t mask_stride = (uint32_t) (mask->nb[1] / sizeof(ggml_fp16_t)); + const bool try_split = use_cm && (!top_k_split_env || top_k_split_env[0] != '0') && n_kv_raw > 0 && top_k->ne[0] > 0; + if (try_split) { + const uint32_t N = (uint32_t) q->ne[1]; + const uint32_t D = 512; + const uint32_t NH = 64; + const uint32_t NS = (uint32_t) q->ne[3]; + const uint32_t raw_kv = (uint32_t) n_kv_raw; + const uint32_t partitions = 2; + // Query tiling caps the split scratch at 256 queries regardless of batch or stream + // count. The reduce takes the tile height as ne2 and the full query count as dst_ne2, + // so the destination stream stride stays correct across tiles. + const uint32_t tile_size = std::min(N, 256u); + const uint32_t n_tiles = CEIL_DIV(N, tile_size); + const bool f32acc = true; + vk_fa_tuning_params tuning = get_fa_tuning_params(ctx->device, D, D, N, raw_kv, GGML_TYPE_F16, GGML_TYPE_F16, f32acc); + + const uint32_t q_stride = (uint32_t) (q->nb[1] / sizeof(float)); + const bool aligned = raw_kv % tuning.block_cols == 0 && (q_stride & 7) == 0 && (k_stride & 7) == 0; + const vk_fa_pipeline_state raw_state = get_fa_pipeline_state(ctx->device, tuning, D, D, aligned, f32acc, + true, false, false, GGML_TYPE_F16, GGML_TYPE_F16); + if (raw_state.path == FA_COOPMAT1 && ctx->device->pipeline_flash_attn_split_k_reduce) { + vk_pipeline raw_pipeline; + { + std::lock_guard guard(ctx->device->compile_mutex); + auto & pipelines = ctx->device->pipeline_flash_attn_f32_f16; + auto it = pipelines.find(raw_state); + if (it != pipelines.end()) { + raw_pipeline = it->second; + } else { + pipelines[raw_state] = raw_pipeline = std::make_shared(); + } + } + const uint64_t partition_size = ((uint64_t) D * NH + NH * 2) * sizeof(float) * tile_size * NS; + const uint64_t split_size = partition_size * partitions; + if (split_size <= ctx->device->properties.limits.maxStorageBufferRange) { + ggml_pipeline_request_descriptor_sets(ctx, raw_pipeline, n_tiles); + ggml_pipeline_request_descriptor_sets(ctx, pipeline, n_tiles); + ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_flash_attn_split_k_reduce, n_tiles); + if (ctx->prealloc_size_split_k < split_size) { + ctx->prealloc_size_split_k = split_size; + ggml_vk_preallocate_buffers(ctx, subctx); + } + if (ctx->prealloc_split_k_need_sync) { + ggml_vk_sync_buffers(ctx, subctx); + } + + // ALiBi is disabled (max_bias == 0), so n_head_log2 is never read; keep it 0 + // rather than a value that looks computed. + const uint32_t n_head_log2 = 0; + // The raw dispatch smuggles the mask row stride through split_kv, which the FA + // shader also uses to derive its KV range as min(KV, (split_k_index+1)*split_kv). + // That is only safe while split_k_index == 0 and the stride covers the whole raw + // prefix -- both hold here, but assert rather than rely on it silently. + GGML_ASSERT(mask_stride >= raw_kv && "split_kv carries the mask stride; it must not clip the raw KV range"); + const uint32_t mask_stride_in_split_kv = 1u << 31; + const uint32_t packed_gqa = mask_stride_in_split_kv | 1u; + const uint32_t packed_partitions = (partitions << 16) | 1; + const vk_subbuffer split_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_split_k, 0); + const vk_subbuffer mask_buf = ggml_vk_tensor_subbuffer(ctx, mask); + const vk_subbuffer top_buf = ggml_vk_tensor_subbuffer(ctx, top_k); + const vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst); + const auto sliced = [](const vk_subbuffer & buf, uint64_t offset) { + return vk_subbuffer{buf.buffer, buf.offset + offset, buf.size - offset}; + }; + + for (uint32_t tile = 0; tile < n_tiles; ++tile) { + if (tile != 0) { + ggml_vk_sync_buffers(ctx, subctx); + } + const uint32_t token_offset = tile * tile_size; + const uint32_t tile_n = std::min(tile_size, N - token_offset); + const vk_subbuffer tile_q = sliced(q_buf, (uint64_t) token_offset * q->nb[1]); + const vk_subbuffer tile_mask = sliced(mask_buf, (uint64_t) token_offset * mask->nb[1]); + const vk_subbuffer tile_top = sliced(top_buf, (uint64_t) token_offset * top_k->nb[1]); + const vk_subbuffer tile_dst = sliced(dst_buf, (uint64_t) token_offset * dst->nb[2]); + const vk_flash_attn_push_constants raw_pc = { + tile_n, raw_kv, + NH, tile_n, NS, + NH, NS, + 1, NS, + 1, NS, + (uint32_t) mask->ne[1], (uint32_t) mask->ne[2], (uint32_t) mask->ne[3], + q_stride, (uint32_t) q->nb[2], (uint32_t) q->nb[3], + k_stride, k_nb2_byte, k_nb3_byte, + k_stride, k_nb2_byte, k_nb3_byte, + scale, 0.0f, 0.0f, + n_head_log2, 1.0f, 1.0f, + packed_gqa, mask_stride, packed_partitions, + }; + + ggml_vk_dispatch_pipeline(ctx, subctx, raw_pipeline, + {tile_q, k_buf, k_buf, tile_mask, tile_q, split_buf, tile_q, tile_q /* dyn-KV: unused */}, + raw_pc, {tile_n, NH, NS}); + ggml_vk_perf_mark_subop(ctx, subctx, "FA_TOP_K_RAW (sub-op)"); + + pc.n_batch = tile_n; + pc.split_mode = 1; + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + {tile_q, k_buf, tile_mask, sinks_buf, tile_top, split_buf}, + pc, {tile_n, (uint32_t) CEIL_DIV(q->ne[2], 32), NS}); + ggml_vk_perf_mark_subop(ctx, subctx, "FA_TOP_K_SELECTED (sub-op)"); + + ctx->prealloc_split_k_need_sync = true; + ggml_vk_sync_buffers(ctx, subctx); + const vk_op_flash_attn_split_k_reduce_push_constants reduce_pc = {D, NH, tile_n, N, NS, partitions, sinks != nullptr}; + ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_flash_attn_split_k_reduce, + {split_buf, sinks_buf, tile_dst}, reduce_pc, {NH, D, tile_n * NS}); + ctx->prealloc_split_k_need_sync = true; + ggml_vk_perf_mark_subop(ctx, subctx, "FA_TOP_K_REDUCE (sub-op)"); + } + return true; + } + } + } + + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + {q_buf, k_buf, ggml_vk_tensor_subbuffer(ctx, mask), sinks_buf, + ggml_vk_tensor_subbuffer(ctx, top_k), ggml_vk_tensor_subbuffer(ctx, dst)}, + pc, {(uint32_t) q->ne[1], (uint32_t) CEIL_DIV(q->ne[2], use_cm ? 32 : 8), (uint32_t) q->ne[3]}); + ggml_vk_perf_mark_subop(ctx, subctx, use_cm ? "FA_TOP_K_CM (sub-op)" : "FA_TOP_K_SPARSE (sub-op)"); + return true; +} + +struct vk_fa_compact_state { + bool active = false; + bool dynamic_kv = false; // KV row count lives in kv_buf, not the push constant + uint32_t kv_c = 0; // upper bound; the real count is runtime when dynamic_kv + uint32_t n_batch = 1; + uint32_t row_bytes = 0; // bytes per compact K row; K may be quantised + uint32_t row_elems = 0; // K row stride in ELEMENTS/blocks, for the FA push constant + bool dequantized = false; // scratch holds f16 because the gather decoded on the way in + vk_subbuffer kv_buf; + vk_subbuffer kc_buf, mc_buf; +}; + +// Small host-visible buffer holding the last union count the device produced: +// [0] padded compact rows (also read by the gather and by the FA), [1] raw union size, +// [2] the candidate count it came from, [3] the batch it came from. Host-visible is a +// requirement rather than a preference here - the whole point is that the host can read it +// without submitting. +static bool ggml_vk_fa_union_stat_init(ggml_backend_vk_context * ctx) { + if (ctx->fa_union_stat) { + return ctx->fa_union_stat->ptr != nullptr; + } + try { + ctx->fa_union_stat = ggml_vk_create_buffer(ctx->device, 64, + {vk::MemoryPropertyFlagBits::eDeviceLocal | vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent, + vk::MemoryPropertyFlagBits::eHostVisible | vk::MemoryPropertyFlagBits::eHostCoherent}); + } catch (const vk::SystemError &) { + return false; + } + if (ctx->fa_union_stat->ptr == nullptr) { + return false; + } + memset(ctx->fa_union_stat->ptr, 0, 64); + return true; +} + +// Host-coherent memory still needs the write made available to the host domain. This is the +// only barrier in the path that names eHost, and it is cheap: nothing waits on it, it just +// lets the next graph's host-side read see the count this one produced. +static void ggml_vk_fa_union_stat_host_barrier(vk_context & subctx) { + subctx->s->buffer->buf.pipelineBarrier( + vk::PipelineStageFlagBits::eComputeShader, + vk::PipelineStageFlagBits::eHost, + {}, + { { vk::AccessFlagBits::eShaderWrite, vk::AccessFlagBits::eHostRead } }, + {}, {}); +} + +// Price the union without synchronising for it. The count read here belongs to whichever call +// last wrote the slot, which under a decode is the final flash-attention op of the previous +// graph - one token stale, and that is fine: selection overlap is a property of the model and +// the draft, not of an individual op, and a wrong estimate costs part of one step rather than +// correctness (the compact buffers are still sized for the worst case). +// +// Tracks the latest measurement, lightly smoothed. The asymmetry runs the other way from what +// a conservative estimator would assume: an estimate that is too HIGH declines compaction and +// forgoes 2-3x for as long as it stays high, while one that is too low costs a single step at +// roughly dense cost and is corrected by the count that step produces. An earlier version held +// a decaying peak instead and measured 1992 us where the union delivers 900, because a spell of +// genuinely low overlap pinned the estimate and 0.999 per read took hundreds of steps to relax. +// +// The words are read without ordering against the device write, so they can come from +// different calls. The sample is filed under the batch the device reported rather than the +// batch being priced, so a torn read costs one mispriced step for that batch and then +// corrects, which is the same failure the estimate already tolerates. +// +// Returns the padded compact row count to gate on, or 0 when this batch is unseeded. +static uint32_t ggml_vk_fa_union_estimate(ggml_backend_vk_context * ctx, uint32_t n_kv_raw, + uint32_t n_batch, uint32_t n_cand) { + const volatile uint32_t * stat = (const volatile uint32_t *) ctx->fa_union_stat->ptr; + const uint32_t u = stat[1]; + const uint32_t cand = stat[2]; + const uint32_t nb_obs = stat[3]; + + if (u > 0 && cand > 0 && nb_obs > 0 && nb_obs < 64) { + const float r = std::min(1.0f, (float) u / (float) cand); + float & e = ctx->fa_union_est_ratio[nb_obs]; + e = e > 0.0f ? 0.5f * r + 0.5f * e : r; + } + + const float ratio = ctx->fa_union_est_ratio[n_batch]; + if (ratio <= 0.0f) { + return 0; + } + return GGML_PAD(n_kv_raw + (uint32_t) ceilf(ratio * (float) n_cand), 256u); +} + +// V4 sparse decode (gather-to-compact): the sparse prefill shader above gates on +// q->ne[1] >= 64, so single-token decode otherwise attends densely over the whole +// compressed KV, at a cost that grows with context. Instead, gather the active rows +// (dense prefix + top-k selection; MQA, so all query heads share one set) into a +// compact contiguous scratch in prealloc_y, and let the ordinary dense FA below run +// over the compacted K/V/mask. Correct by the same contract as the sparse shader: +// the source mask carries the selection, and the gathered mask preserves it. +static bool ggml_vk_flash_attn_gather_compact(ggml_backend_vk_context * ctx, vk_context & subctx, + const ggml_tensor * q, const ggml_tensor * k, const ggml_tensor * v, + const ggml_tensor * mask, ggml_tensor * dst, vk_fa_compact_state & st) { + const ggml_tensor * top_k = dst->src[5]; + static const char * gather_env = getenv("GGML_VK_FA_TOPK_GATHER"); + // K/V type: the gather relocates rows verbatim and never reads a value out of them, so it + // does not have to be f16 - it has to be a type whose row is a whole number of 4-byte words, + // and one flash-attention has a NATIVE shader for. The dequant/contiguize pass is disabled + // while this scratch is active and asserts if it turns out to be needed, so admitting a + // non-native type here would abort rather than fall back. + const bool kv_word_addressable = + k->type == v->type && + ggml_vk_fa_kv_native(k->type, ctx->device->coopmat2) && + k->ne[0] % ggml_blck_size(k->type) == 0 && + ggml_row_size(k->type, k->ne[0]) % 4 == 0 && + k->nb[1] % 4 == 0 && k->nb[3] % 4 == 0; + + if ((gather_env && gather_env[0] == '0') || + !top_k || !ctx->device->pipeline_flash_attn_gather_f16 || + q->ne[1] < 1 || q->ne[1] >= 64 || // 1..63: >=64 goes to the sparse prefill path + q->type != GGML_TYPE_F32 || !kv_word_addressable || + !mask || mask->type != GGML_TYPE_F16 || top_k->type != GGML_TYPE_I32 || + q->ne[0] != 512 || k->ne[0] != 512 || v->ne[0] != 512 || q->ne[2] != 64 || + k->ne[2] != 1 || v->ne[2] != 1 || + q->ne[1] != top_k->ne[1] || q->ne[3] != top_k->ne[3] || + k->ne[1] != v->ne[1] || k->buffer != v->buffer || k->data != v->data || + !ggml_is_contiguous(mask) || !ggml_is_contiguous(top_k)) { + return false; + } + + const uint32_t k_row_bytes = (uint32_t) ggml_row_size(k->type, k->ne[0]); + const uint32_t k_row_words = k_row_bytes / 4; + + float max_bias = 0.0f; + float logit_softcap = 0.0f; + memcpy(&max_bias, (const float *) dst->op_params + 1, sizeof(float)); + memcpy(&logit_softcap, (const float *) dst->op_params + 2, sizeof(float)); + if (max_bias != 0.0f || logit_softcap != 0.0f) { + return false; + } + + const int32_t n_kv_raw = ggml_get_op_params_i32(dst, 4); + if (n_kv_raw < 0 || n_kv_raw > k->ne[1] || top_k->ne[0] > k->ne[1] - n_kv_raw) { + return false; + } + + const uint32_t n_batch = (uint32_t) q->ne[1]; + const uint32_t n_cand = (uint32_t) ((int64_t) n_batch * top_k->ne[0]); + // Worst case: every token gets its own top-k block, so the compact set is + // n_kv_raw + n_batch*n_top_k -- independent of context depth, which is the point, but + // GROWING WITH BATCH. Attention work is then n_batch * (n_kv_raw + n_batch*n_top_k), i.e. + // quadratic in batch, against dense's n_batch * n_kv. Break-even is + // n_batch = (n_kv - n_kv_raw) / n_top_k, and a kv >= 2*kv_c gate on this worst case caps + // the useful batch at (n_kv/2 - n_kv_raw) / n_top_k -- about 6 at 32k depth, ~30 at 128k. + // Beyond that the gate declines and dense runs, so this can never be slower; it just stops + // helping. The union below lifts that ceiling by gating on what the selections actually + // deduplicate to; kv_c remains the bound for every allocation and dispatch count. + const uint32_t kv_c = GGML_PAD((uint32_t) (n_kv_raw + (int64_t) n_cand), 256u); + + // ---- deduplicated union (default on, GGML_VK_FA_TOPK_UNION=0 disables) -------------- + // Same compact layout, but one row per DISTINCT selected key instead of one block per + // token. Measured on real draft tokens the union is 0.56 of the selections at batch 4, + // so it is materially smaller. Its size is only known on the GPU, so the FA + // reads its KV bound from a buffer (the DYNAMIC_KV pipeline flag) rather than a push + // constant; padding the count to 256 keeps KV % Bc == 0 so the aligned variant still + // applies. Single stream only: the FA takes one KV for all streams. + // + // Gated on the ESTIMATED union rather than on kv_c, which is why this sits above the + // worst-case gate: at batch 8 and 32k depth the worst case is 6400 rows against 11008 + // source rows and would decline, while the union measures around 3300 and is well worth + // compacting. kv_c stays the worst case for every allocation and dispatch bound, so a + // wrong estimate is a slow step, never a wrong answer. + const uint32_t max_words = 12288; // shared bitmap capacity in flash_attn_union.comp + const bool bitmap_fits = (uint64_t) ((k->ne[1] - n_kv_raw) + 31) / 32 <= max_words; + + static const char * union_env = getenv("GGML_VK_FA_TOPK_UNION"); + if ((!union_env || union_env[0] != '0') && q->ne[3] == 1 && n_batch > 1 && bitmap_fits && + ctx->device->pipeline_flash_attn_union_f16 && ctx->device->pipeline_flash_attn_gather_union_f16 && + ggml_vk_fa_union_stat_init(ctx)) { + const uint32_t max_union = n_cand; + const uint32_t kv_c_est = ggml_vk_fa_union_estimate(ctx, (uint32_t) n_kv_raw, n_batch, n_cand); + // Two separate questions. Does the compact set fit under the gate at all, and does + // deduplicating actually shrink it: with no overlap to exploit the union is the same + // size as the per-token blocks and the scan is pure cost, measured at 1.2% of the op + // at 512k depth. The worst-case bound on the source keeps a collapse in overlap to + // roughly dense cost for the one step it takes the estimate to catch up. + const bool worth_it = kv_c_est != 0 && kv_c_est < kv_c && + (uint64_t) k->ne[1] >= 2ull * kv_c_est && + (uint64_t) k->ne[1] >= (uint64_t) kv_c; + + // GGML_VK_FA_UNION_STATS=N: report every Nth call (N=1 means every call) what the gate + // decided and on what measurement. The alternative is inferring engagement from a + // timing, which is how a sparse path gets credited for a run it never took. + // + // The period is a parameter because a fixed one is a way to miss the answer: a 64-token + // speculative decode over 25k of context produced ONE line at a period of 256, which + // said only that the first call was unseeded. Running totals rather than instants, so a + // single late line still reports whether the path engaged. + static const char * stats_env = getenv("GGML_VK_FA_UNION_STATS"); + if (stats_env && stats_env[0] != '\0' && stats_env[0] != '0') { + static uint64_t calls = 0, taken = 0; + const uint64_t period = std::max(1ull, (unsigned long long) atoll(stats_env)); + taken += worth_it ? 1 : 0; + if ((calls++ % period) == 0) { + fprintf(stderr, "[fa-union] n_kv=%lld n_kv_raw=%d n_batch=%u cand=%u " + "union/cand=%.3f kv_c %u -> est %u %s (%llu/%llu taken)\n", + (long long) k->ne[1], n_kv_raw, n_batch, n_cand, (double) ctx->fa_union_est_ratio[n_batch], + kv_c, kv_c_est, worth_it ? "UNION" : "declined", + (unsigned long long) taken, (unsigned long long) calls); + } + } + + if (!worth_it) { + // Nothing is known about this batch shape, or the last count says compaction does + // not pay. Either way the count is what settles it, so produce one: the scan is a + // single workgroup and depth-independent, and count_only needs no index list. Dense + // runs this step; the next one decides on a measurement instead of a bound. + // + // Every declined step, not a sample: a decline is exactly the state in which the + // estimate stops being refreshed by the compact path, so sampling it leaves a stale + // estimate latched for as many steps as the sampling period. The dispatch is one + // workgroup against the ~2.2 ms dense op it is riding along with. + ctx->fa_union_declines++; + { + const vk_op_flash_attn_union_push_constants ppc = { + (uint32_t) k->ne[1], (uint32_t) n_kv_raw, n_batch, (uint32_t) top_k->ne[0], max_union, + (uint32_t) (top_k->nb[1] / sizeof(int32_t)), max_words, 256u, 1u, + }; + const vk_subbuffer stat_buf = ggml_vk_subbuffer(ctx, ctx->fa_union_stat); + ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_flash_attn_union_f16, 1); + // The probe writes the same slot a taken union path reads back as its row + // count, and a graph can contain both: the estimate is refreshed from the + // device as the graph is recorded, so an op late in the graph can be admitted + // after an earlier one was declined. Order it explicitly rather than rely on + // the compact path's own sync, which this path does not go through. + ggml_vk_sync_buffers(ctx, subctx); + ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_flash_attn_union_f16, + { ggml_vk_tensor_subbuffer(ctx, top_k), stat_buf, stat_buf }, ppc, { 1, 1, 1 }); + ggml_vk_fa_union_stat_host_barrier(subctx); + } + goto union_unavailable; + } + + // Decoding on the way in makes the scratch f16 and hands flash attention its f16 path, + // which is worth far more than the extra scratch bytes: the inline decode it replaces + // costs a measured 0.15 us per KV row attended, every step. + const bool dq = ggml_is_quantized(k->type) && ctx->device->pipeline_flash_attn_gather_union_dq[k->type]; + const uint32_t u_row_by = dq ? (uint32_t) (k->ne[0] * sizeof(ggml_fp16_t)) : k_row_bytes; + const size_t ukc_sz = (size_t) kv_c * u_row_by; + const size_t umc_sz = (size_t) n_batch * kv_c * sizeof(ggml_fp16_t); + const size_t ul_sz = (size_t) max_union * sizeof(uint32_t); + const size_t need = ukc_sz + umc_sz + ul_sz; + if (ctx->prealloc_size_y < need) { + ctx->prealloc_size_y = need; + ggml_vk_preallocate_buffers(ctx, subctx); + } + // Unconditional, not gated on prealloc_y_need_sync: this also orders the count slot + // against a probe dispatched earlier in the same graph, which leaves that flag clear. + ggml_vk_sync_buffers(ctx, subctx); + const vk_subbuffer kc_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_y, 0); + const vk_subbuffer mc_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_y, ukc_sz); + const vk_subbuffer ul_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_y, ukc_sz + umc_sz); + // The count lives in the stat buffer rather than in prealloc_y so that the same write + // that the gather and the FA consume is also the one the host prices the next step from. + // Both are single-slot and rewritten by every layer, so the WAR hazard is unchanged: the + // prealloc_y sync above is a global barrier and orders the previous layer's read. + const vk_subbuffer uc_buf = ggml_vk_subbuffer(ctx, ctx->fa_union_stat); + + vk_pipeline gather_pipe = dq ? ctx->device->pipeline_flash_attn_gather_union_dq[k->type] + : ctx->device->pipeline_flash_attn_gather_union_f16; + ggml_pipeline_request_descriptor_sets(ctx, ctx->device->pipeline_flash_attn_union_f16, 1); + ggml_pipeline_request_descriptor_sets(ctx, gather_pipe, 1); + + const vk_op_flash_attn_union_push_constants upc = { + (uint32_t) k->ne[1], (uint32_t) n_kv_raw, n_batch, (uint32_t) top_k->ne[0], max_union, + (uint32_t) (top_k->nb[1] / sizeof(int32_t)), max_words, 256u, 0u, + }; + ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_flash_attn_union_f16, + { ggml_vk_tensor_subbuffer(ctx, top_k), ul_buf, uc_buf }, upc, { 1, 1, 1 }); + ggml_vk_sync_buffers(ctx, subctx); + + // the fused decoder steps K in BLOCKS and writes elements; the verbatim one does words + const vk_op_flash_attn_gather_union_push_constants gpc = { + (uint32_t) k->ne[1], (uint32_t) n_kv_raw, kv_c, + dq ? (uint32_t) (k->nb[1] / ggml_type_size(k->type)) : (uint32_t) (k->nb[1] / 4), + (uint32_t) (mask->nb[1] / sizeof(ggml_fp16_t)), + n_batch, dq ? (uint32_t) k->ne[0] : k_row_words, + }; + ggml_vk_dispatch_pipeline(ctx, subctx, gather_pipe, + { ggml_vk_tensor_subbuffer(ctx, k), ul_buf, ggml_vk_tensor_subbuffer(ctx, mask), + kc_buf, mc_buf, uc_buf }, gpc, { kv_c, 1, 1 }); + ggml_vk_sync_buffers(ctx, subctx); + ggml_vk_fa_union_stat_host_barrier(subctx); + ctx->prealloc_y_need_sync = true; + + st.active = true; + st.dynamic_kv = true; + st.kv_c = kv_c; + st.n_batch = n_batch; + st.row_bytes = u_row_by; + st.row_elems = dq ? (uint32_t) k->ne[0] : (uint32_t) (k->ne[0] / ggml_blck_size(k->type)); + st.dequantized = dq; + st.kc_buf = kc_buf; + st.mc_buf = mc_buf; + st.kv_buf = uc_buf; + return true; + } +union_unavailable:; - const uint32_t psh_stride = Br / 4 + 2; - const uint32_t Psh = Bc * psh_stride * f16vec4; + // Per-token blocks have no dedup, so this form really does cost kv_c: the gather writes + // then re-reads ~the active bytes while dense reads the source KV once, so compaction only + // pays when the source is comfortably larger than the active set. + if ((uint64_t) k->ne[1] < 2ull * kv_c) { + return false; + } - const uint32_t sfshstride = (hsk <= 128) ? (Br + 8) : Br; - const uint32_t sfsh = Bc * sfshstride * acctype; + // Decode on the way in, as the union path does: the gather touches each row once, while + // flash attention decodes once per query block that reads it. The union only covers + // n_batch > 1, so single-token decode lands here. + const bool tok_dq = ggml_is_quantized(k->type) && ctx->device->pipeline_flash_attn_gather_dq[k->type]; + const uint32_t row_by = tok_dq ? (uint32_t) (k->ne[0] * sizeof(ggml_fp16_t)) : k_row_bytes; - const uint32_t kvshstride = (params.shmem_staging ? std::max(hsk_pad, hsv_pad) : MatBr) / 4 + 2; - const uint32_t vsh_stride = MatBc / 4 * row_split; - const uint32_t ksh = ((kvshstride >= vsh_stride) ? (Bc * kvshstride) : (Bc * vsh_stride)) * f16vec4; + const uint32_t ns = (uint32_t) q->ne[3]; + const size_t kc_sz = (size_t) ns * kv_c * row_by; + const size_t mc_sz = (size_t) ns * n_batch * kv_c * sizeof(ggml_fp16_t); - // BF16 PVMat accumulator is f32 (no bf16 accumulator support), so pvsh is vec4 (16 bytes) - const uint32_t pvsh_elem_size = (k_type == GGML_TYPE_BF16) ? 16u : f16vec4; - const uint32_t osh_stride = params.row_split * MatBr / 4; - const uint32_t pvsh = MatBc * osh_stride * pvsh_elem_size; + if (ctx->prealloc_size_y < kc_sz + mc_sz) { + ctx->prealloc_size_y = kc_sz + mc_sz; + ggml_vk_preallocate_buffers(ctx, subctx); + } + if (ctx->prealloc_y_need_sync) { + ggml_vk_sync_buffers(ctx, subctx); + } - const uint32_t slope = Br * acctype; + vk_pipeline pipeline = tok_dq ? ctx->device->pipeline_flash_attn_gather_dq[k->type] + : ctx->device->pipeline_flash_attn_gather_f16; + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); - const uint32_t total_size = tmpsh + iq_shmem + Qf + Psh + sfsh + ksh + pvsh + slope; - const bool supported = total_size <= device->properties.limits.maxComputeSharedMemorySize; + // the fused decoder steps K in BLOCKS and writes elements; the verbatim one does words + const vk_op_flash_attn_gather_push_constants pc = { + (uint32_t) k->ne[1], (uint32_t) n_kv_raw, (uint32_t) top_k->ne[0], kv_c, + tok_dq ? (uint32_t) (k->nb[1] / ggml_type_size(k->type)) : (uint32_t) (k->nb[1] / 4), + tok_dq ? (uint32_t) (k->nb[3] / ggml_type_size(k->type)) : (uint32_t) (k->nb[3] / 4), + (uint32_t) (top_k->nb[1] / sizeof(int32_t)), + (uint32_t) (top_k->nb[3] / sizeof(int32_t)), + (uint32_t) (mask->nb[1] / sizeof(ggml_fp16_t)), + (uint32_t) (mask->nb[3] / sizeof(ggml_fp16_t)), + (uint32_t) mask->ne[3], + n_batch, tok_dq ? (uint32_t) k->ne[0] : k_row_words, + }; - VK_LOG_DEBUG("ggml_vk_flash_attn_coopmat_shmem_support(HSK=" << hsk << ", HSV=" << hsv << ", f32acc=" << f32acc << ", total_size=" << total_size << ", supported=" << supported); + st.kc_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_y, 0); + st.mc_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_y, kc_sz); + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + { ggml_vk_tensor_subbuffer(ctx, k), ggml_vk_tensor_subbuffer(ctx, top_k), + ggml_vk_tensor_subbuffer(ctx, mask), st.kc_buf, st.mc_buf }, + pc, { kv_c, 1, ns }); + ggml_vk_sync_buffers(ctx, subctx); + ctx->prealloc_y_need_sync = true; + + // ---- diagnostic: measure real top-k overlap between the query tokens ---------------- + // GGML_VK_TOPK_OVERLAP=1. Off by default and never touched on the hot path. Answers the + // one question that decides whether a deduplicated union is worth building: this path + // gathers n_batch*n_top_k rows, a union would gather |union| rows, and the op's cost is + // linear in that count (measured: 204.4 / 205.7 / 205.5 us per row at batch 2 / 4 / 8). + // Needs a real model - synthetic top-k selections say nothing about real overlap, which + // is exactly how the mul_mat_id large-tile probe misled once before. + static const char * overlap_env = getenv("GGML_VK_TOPK_OVERLAP"); + if (overlap_env && overlap_env[0] == '1' && n_batch > 1) { + static std::mutex ov_mutex; + static uint64_t ov_calls = 0, ov_selected = 0, ov_union = 0; + static std::map> ov_by_batch; // n_batch -> {selected, union} + const size_t n_idx = (size_t) n_batch * top_k->ne[0]; + std::vector idx(n_idx); + vk_subbuffer top_sb = ggml_vk_tensor_subbuffer(ctx, top_k); + ggml_vk_buffer_read(top_sb.buffer, top_sb.offset, idx.data(), n_idx * sizeof(int32_t)); + + const int32_t range = (int32_t) (k->ne[1] - n_kv_raw); + std::unordered_set uni; + uint64_t valid = 0; + for (size_t i = 0; i < n_idx; ++i) { + const int32_t v = idx[i]; + if (v >= 0 && v < range) { uni.insert(v); ++valid; } + } + std::lock_guard lock(ov_mutex); + ov_calls++; ov_selected += valid; ov_union += uni.size(); + auto & e = ov_by_batch[n_batch]; + e.first += valid; e.second += uni.size(); + if ((ov_calls % 256) == 0) { + fprintf(stderr, "[topk-overlap] calls=%llu selected=%llu union=%llu " + "union/selected=%.3f => a dedup union would gather %.1f%% fewer compressed rows\n", + (unsigned long long) ov_calls, (unsigned long long) ov_selected, + (unsigned long long) ov_union, + ov_selected ? (double) ov_union / (double) ov_selected : 0.0, + ov_selected ? 100.0 * (1.0 - (double) ov_union / (double) ov_selected) : 0.0); + for (const auto & kv : ov_by_batch) { + const double ratio = kv.second.first ? (double) kv.second.second / (double) kv.second.first : 0.0; + // projected op speedup uses the measured linear cost model on kv_c + const double now = (double) n_kv_raw + (double) kv.first * (double) top_k->ne[0]; + const double dedup = (double) n_kv_raw + ratio * (double) kv.first * (double) top_k->ne[0]; + fprintf(stderr, "[topk-overlap] n_batch=%u union/selected=%.3f " + "kv_c %.0f -> %.0f projected op gain %.1f%%\n", + kv.first, ratio, now, dedup, 100.0 * (1.0 - dedup / now)); + } + } + } - return supported; + st.active = true; + st.kv_c = kv_c; + st.n_batch = n_batch; + st.row_bytes = row_by; + st.row_elems = tok_dq ? (uint32_t) k->ne[0] : (uint32_t) (k->ne[0] / ggml_blck_size(k->type)); + st.dequantized = tok_dq; + return true; } static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * q, const ggml_tensor * k, const ggml_tensor * v, const ggml_tensor * mask, const ggml_tensor * sinks, ggml_tensor * dst) { @@ -11227,15 +12912,15 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx GGML_TENSOR_LOCALS(int64_t, ne, dst, ne) GGML_TENSOR_LOCALS(size_t, nb, dst, nb) - const uint32_t nem0 = mask ? mask->ne[0] : 0; - const uint32_t nem1 = mask ? mask->ne[1] : 0; - const uint32_t nem2 = mask ? mask->ne[2] : 0; - const uint32_t nem3 = mask ? mask->ne[3] : 0; + uint32_t nem0 = mask ? mask->ne[0] : 0; + uint32_t nem1 = mask ? mask->ne[1] : 0; + uint32_t nem2 = mask ? mask->ne[2] : 0; + uint32_t nem3 = mask ? mask->ne[3] : 0; const uint32_t HSK = nek0; const uint32_t HSV = nev0; uint32_t N = neq1; - const uint32_t KV = nek1; + uint32_t KV = nek1; GGML_ASSERT(ne0 == HSV); GGML_ASSERT(ne2 == N); @@ -11259,6 +12944,23 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx assert(dst->type == GGML_TYPE_F32); assert(q->type == GGML_TYPE_F32); + if (ggml_vk_flash_attn_top_k(ctx, subctx, q, k, v, mask, sinks, dst)) { + return; + } + // V4 sparse decode: gather the active set into a compact scratch and run the dense + // FA below on it. Overrides KV, the mask geometry, and (further down) the K/V/mask + // bindings and strides; every other decision then sizes itself to the compact KV. + vk_fa_compact_state fa_compact; + if (ggml_vk_flash_attn_gather_compact(ctx, subctx, q, k, v, mask, dst, fa_compact)) { + KV = fa_compact.kv_c; + nem0 = fa_compact.kv_c; + nem1 = fa_compact.n_batch; + nem2 = 1; + nem3 = (uint32_t) q->ne[3]; + nem1 = N; + nem2 = 1; + nem3 = (uint32_t) q->ne[3]; + } uint32_t gqa_ratio = 1; uint32_t qk_ratio = neq2 / nek2; uint32_t workgroups_x = (uint32_t)neq1; @@ -11267,7 +12969,14 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx const bool f32acc = !ctx->device->fp16 || dst->op_params[3] == GGML_PREC_F32 || k->type == GGML_TYPE_BF16; - // dequant K/V once into an f16 scratch, reordered KV layout so FA can read without a stride + // For prefill with quantized K/V, dequantize+transpose K/V once into a per-head-contiguous + // f16 scratch and run the f16 FA path, instead of the coopmat1 shader re-dequantizing the + // whole KV inside every Q-workgroup. The KV-cache view reaching FA is [0,2,1,3]-permuted but + // dense, so we require dense allocation (not ggml_is_contiguous) and block-contiguous dim0, + // and only engage where a fused dequant-transpose shader exists. Prefill only + // (n_rows >= 64); measured neutral at shallow depth and up to ~2x at long context. + // The scratch is bound as a single storage buffer holding K and V back to back, so the SUM of + // the two must fit maxStorageBufferRange, not each half independently. auto is_dense_kv_cache = [](const ggml_tensor * t) { return t->nb[0] == ggml_type_size(t->type) && t->nb[2] == ggml_row_size(t->type, t->ne[0]) && @@ -11276,19 +12985,48 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx }; const bool k_quant = k->type != GGML_TYPE_F16 && k->type != GGML_TYPE_BF16 && k->type != GGML_TYPE_F32; const bool v_quant = v->type != GGML_TYPE_F16 && v->type != GGML_TYPE_BF16 && v->type != GGML_TYPE_F32; - const bool use_dequant_kv = k_quant && v_quant && neq1 >= 64 && + const uint64_t kv_f16_sz = ((uint64_t)ggml_nelements(k) + (uint64_t)ggml_nelements(v)) * sizeof(ggml_fp16_t); + static const char * fa_dequant_env = getenv("GGML_VK_FA_DEQUANT"); + const bool fa_dequant_off = fa_dequant_env && fa_dequant_env[0] == '0'; + const bool fa_dequant_on = fa_dequant_env && fa_dequant_env[0] == '1'; + // Contiguize pass for f16 K/V (GGML_VK_FA_KV_CONTIG=0 opts out). The KV-cache view + // reaching FA is head-interleaved ([HS, NH, KV] physically), and the cm1 + // direct-from-global coopMatLoads run ~2-5x slower on those strided rows than on + // per-head-contiguous K/V. Copy K/V once into the scratch instead (dequant_f16_transpose + // is a pure strided copy). Only engages when the rows are actually strided, and shares + // the quant path's prefill/allocation/scratch-capacity gates below. + static const char * fa_kv_contig_env = getenv("GGML_VK_FA_KV_CONTIG"); + const bool fa_kv_contig = !(fa_kv_contig_env && fa_kv_contig_env[0] == '0'); + const bool kv_f16_strided = k->type == GGML_TYPE_F16 && v->type == GGML_TYPE_F16 && + (k->nb[1] != (uint64_t)HSK * sizeof(ggml_fp16_t) || + v->nb[1] != (uint64_t)HSV * sizeof(ggml_fp16_t)) && + (HSK % 8) == 0 && (HSV % 8) == 0; + // A K/V type the FA shaders cannot read directly (iq4_nl) is only correct through the + // dequant path. supports_op only admits such types when the hard conditions below hold, + // and the VRAM heuristic must not veto them (there is no fallback), so force the path on. + const bool kv_needs_dequant = !ggml_vk_fa_kv_native(k->type, ctx->device->coopmat2) || + !ggml_vk_fa_kv_native(v->type, ctx->device->coopmat2); + const bool use_dequant_kv = !fa_dequant_off && + // the gather-to-compact scratch is already contiguous f16; the + // dequant/contiguize pass must not run on top of it + !fa_compact.active && + ((k_quant && v_quant) || kv_needs_dequant || (fa_kv_contig && kv_f16_strided)) && neq1 >= 64 && is_dense_kv_cache(k) && is_dense_kv_cache(v) && - (uint64_t)ggml_nelements(k) * sizeof(ggml_fp16_t) <= ctx->device->properties.limits.maxStorageBufferRange && - (uint64_t)ggml_nelements(v) * sizeof(ggml_fp16_t) <= ctx->device->properties.limits.maxStorageBufferRange && + kv_f16_sz <= ctx->device->properties.limits.maxStorageBufferRange && ctx->device->pipeline_dequant_transpose[k->type] != nullptr && ctx->device->pipeline_dequant_transpose[v->type] != nullptr && - // coopmat2 path does not benefit from the f16 scratch - !ctx->device->coopmat2 && + // coopmat2 reads its native types directly; non-native still needs the scratch + (kv_needs_dequant || !ctx->device->coopmat2) && // Intel Xe1 regresses, see PR 25494 - (ctx->device->vendor_id != VK_VENDOR_ID_INTEL || - (ctx->device->coopmat_support && ctx->device->architecture != vk_device_architecture::INTEL_XE1)); - const ggml_type k_type_eff = use_dequant_kv ? GGML_TYPE_F16 : k->type; - const ggml_type v_type_eff = use_dequant_kv ? GGML_TYPE_F16 : v->type; + (kv_needs_dequant || + ctx->device->vendor_id != VK_VENDOR_ID_INTEL || + (ctx->device->coopmat_support && ctx->device->architecture != vk_device_architecture::INTEL_XE1)) && + (fa_dequant_on || kv_needs_dequant || ggml_vk_fa_dequant_scratch_fits(ctx, kv_f16_sz)); + // If this fires, supports_op admitted a non-native K/V type the gate then rejected; the + // native shader would return garbage rather than fail, so abort instead. + GGML_ASSERT(use_dequant_kv || !kv_needs_dequant); + const ggml_type k_type_eff = (use_dequant_kv || fa_compact.dequantized) ? GGML_TYPE_F16 : k->type; + const ggml_type v_type_eff = (use_dequant_kv || fa_compact.dequantized) ? GGML_TYPE_F16 : v->type; // For scalar/coopmat1 FA, we can use the "large" size to accommodate qga. // For coopmat2 FA, we always use the small size (which is still pretty large for gqa). @@ -11307,33 +13045,15 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx tuning_params = get_fa_tuning_params(ctx->device, HSK, HSV, N, KV, k_type_eff, v_type_eff, f32acc); - float scale = 1.0f; - float max_bias = 0.0f; - float logit_softcap = 0.0f; - - memcpy(&scale, (const float *) dst->op_params + 0, sizeof(float)); - memcpy(&max_bias, (const float *) dst->op_params + 1, sizeof(float)); - memcpy(&logit_softcap, (const float *) dst->op_params + 2, sizeof(float)); - - if (logit_softcap != 0) { - scale /= logit_softcap; - } - - // Sparse mask hint (op_params[4]): compact the <= n_kv_max finite positions and gather only those. - const int32_t n_kv_max = mask ? ggml_get_op_params_i32(dst, 4) : 0; - static const bool disable_sparse = getenv("GGML_VK_FA_SPARSE_DISABLE") != nullptr; - // cm2 dense is fast, so it needs a larger reduction to win. - const int64_t min_ratio = tuning_params.path == FA_COOPMAT2 ? 4 : 2; - const bool use_sparse = !disable_sparse && n_kv_max > 0 && mask && - max_bias == 0.0f && logit_softcap == 0.0f && - k_type_eff == GGML_TYPE_F16 && v_type_eff == GGML_TYPE_F16 && - nem0 == KV && - (int64_t)KV >= std::max(4096, min_ratio * (int64_t)n_kv_max) && - (gqa_ratio > 1 || (tuning_params.path == FA_SCALAR && N == 1)); - const uint32_t q_stride = (uint32_t)(nbq1 / ggml_type_size(q->type)); uint32_t k_stride = (uint32_t)(nbk1 / ggml_type_size(k->type)); uint32_t v_stride = (uint32_t)(nbv1 / ggml_type_size(v->type)); + if (fa_compact.active) { + // rows are tightly packed in the compact scratch; for a quantised K this is the block + // count per row, which is what nbk1 / ggml_type_size would have given for the source + k_stride = fa_compact.row_elems; + v_stride = fa_compact.row_elems; + } // For F32, the shader treats it as a block of size 4 (for vec4 loads) if (k->type == GGML_TYPE_F32) { @@ -11353,6 +13073,7 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx nbv2_eff = (uint32_t)((uint64_t)HSV * KV * sizeof(ggml_fp16_t)); nbv3_eff = (uint32_t)((uint64_t)HSV * KV * nev2 * sizeof(ggml_fp16_t)); } + const uint32_t alignment = tuning_params.block_cols; bool aligned = (KV % alignment) == 0 && // the "aligned" shader variant will forcibly align strides, for performance @@ -11363,11 +13084,24 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx aligned = false; } + float scale = 1.0f; + float max_bias = 0.0f; + float logit_softcap = 0.0f; + + memcpy(&scale, (const float *) dst->op_params + 0, sizeof(float)); + memcpy(&max_bias, (const float *) dst->op_params + 1, sizeof(float)); + memcpy(&logit_softcap, (const float *) dst->op_params + 2, sizeof(float)); + + if (logit_softcap != 0) { + scale /= logit_softcap; + } + // Only use mask opt when the mask is fairly large. This hasn't been tuned extensively. - bool use_mask_opt = mask && !use_sparse && nem1 >= 32 && nem0 * nem1 > 32768 && nem0 >= tuning_params.block_cols * 16 + bool use_mask_opt = mask && nem1 >= 32 && nem0 * nem1 > 32768 && nem0 >= tuning_params.block_cols * 16 && (ctx->device->architecture != vk_device_architecture::AMD_GCN || HSK > 256 || HSV > 256); vk_fa_pipeline_state fa_pipeline_state = get_fa_pipeline_state(ctx->device, tuning_params, HSK, HSV, aligned, f32acc, - mask != nullptr, use_mask_opt, logit_softcap != 0, use_sparse, k_type_eff, v_type_eff); + mask != nullptr, use_mask_opt, logit_softcap != 0, k_type_eff, v_type_eff, + fa_compact.dynamic_kv); vk_pipeline pipeline = nullptr; @@ -11402,19 +13136,7 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx const uint32_t Tr = CEIL_DIV(N, Br); // Try to use split_k when KV is large enough to be worth the overhead. - // Sparse: split_kv carries n_kv_max, split_k partitions its blocks for occupancy. - if (use_sparse) { - split_kv = (uint32_t)n_kv_max; - const uint32_t total_blocks = CEIL_DIV((uint32_t)n_kv_max, Bc); - const uint32_t base_wgs = (gqa_ratio > 1 ? workgroups_x : Tr) * workgroups_y * workgroups_z; - if (base_wgs < shader_core_count * 2) { - split_k = shader_core_count * 2 / base_wgs; - } - split_k = std::max(1u, std::min(split_k, total_blocks)); - // Match the shader's per-split block count so no split is empty. - const uint32_t per_blocks = CEIL_DIV(total_blocks, split_k); - split_k = CEIL_DIV(total_blocks, per_blocks); - } else if (gqa_ratio > 1 && workgroups_x <= Br) { + if (gqa_ratio > 1 && workgroups_x <= Br) { split_k = shader_core_count * 2 / (workgroups_x * workgroups_y * workgroups_z); } else if (gqa_ratio <= 1) { uint32_t total_wgs_no_split = Tr * workgroups_y * workgroups_z; @@ -11423,7 +13145,7 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx } } - if (!use_sparse && split_k > 1) { + if (split_k > 1) { // Try to evenly split KV into split_k chunks, but it needs to be a multiple // of "align", so recompute split_k based on that. split_kv = ROUNDUP_POW2(std::max(1u, KV / split_k), alignment); @@ -11470,24 +13192,6 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx } } - // Sparse index scratch reuses prealloc_y (mutually exclusive with mask opt). - const uint64_t sparse_idx_size = use_sparse - ? sizeof(int32_t) * (uint64_t)n_kv_max * nem1 * nem2 * nem3 - : 0; - vk_pipeline sparse_compact_pipeline = ctx->device->fa_sparse_compact_use_subgroups - ? ctx->device->pipeline_fa_sparse_compact_subgroup - : ctx->device->pipeline_fa_sparse_compact; - if (use_sparse) { - ggml_pipeline_request_descriptor_sets(ctx, sparse_compact_pipeline, 1); - if (ctx->prealloc_size_y < sparse_idx_size) { - ctx->prealloc_size_y = sparse_idx_size; - ggml_vk_preallocate_buffers(ctx, subctx); - } - if (ctx->prealloc_y_need_sync) { - ggml_vk_sync_buffers(ctx, subctx); - } - } - const uint32_t n_head_kv = neq2; const uint32_t n_head_log2 = 1u << (uint32_t) floorf(log2f((float) n_head_kv)); const float m0 = powf(2.0f, -(max_bias ) / n_head_log2); @@ -11498,9 +13202,13 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx vk_subbuffer v_buf = ggml_vk_tensor_subbuffer(ctx, v); vk_subbuffer dst_buf = ggml_vk_tensor_subbuffer(ctx, dst); vk_subbuffer mask_buf = mask ? ggml_vk_tensor_subbuffer(ctx, mask) : q_buf; + if (fa_compact.active) { + k_buf = fa_compact.kc_buf; + v_buf = fa_compact.kc_buf; // V is the K latent; one gather serves both + mask_buf = fa_compact.mc_buf; + } vk_subbuffer sinks_buf = sinks ? ggml_vk_tensor_subbuffer(ctx, sinks) : q_buf; vk_subbuffer mask_opt_buf = use_mask_opt ? ggml_vk_subbuffer(ctx, ctx->prealloc_y, 0) : q_buf; - vk_subbuffer sparse_buf = use_sparse ? ggml_vk_subbuffer(ctx, ctx->prealloc_y, 0) : q_buf; if (use_dequant_kv) { const uint64_t fp = sizeof(ggml_fp16_t); @@ -11528,6 +13236,11 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx ggml_vk_sync_buffers(ctx, subctx); k_buf = k_dst; v_buf = v_dst; + // Bill the K/V contiguize/dequant pass on its own line. Without this it is charged to + // FLASH_ATTN_EXT, which makes the copy invisible and the kernel look slower than it is. + ggml_vk_perf_mark_subop(ctx, subctx, kv_needs_dequant || (k_quant && v_quant) + ? "FA_KV_DEQUANT (sub-op)" + : "FA_KV_CONTIGUIZE (sub-op)"); } uint32_t mask_n_head_log2 = ((sinks != nullptr) << 24) | n_head_log2; @@ -11552,23 +13265,11 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx ggml_vk_sync_buffers(ctx, subctx); } - if (use_sparse) - { - const vk_op_flash_attn_sparse_compact_push_constants sc_pc = { - KV, - nem1, - nem2, - (uint32_t)(mask->nb[1] / sizeof(ggml_fp16_t)), - (uint32_t)(mask->nb[2] / sizeof(ggml_fp16_t)), - (uint32_t)(mask->nb[3] / sizeof(ggml_fp16_t)), - (uint32_t)n_kv_max, - }; - - ggml_vk_dispatch_pipeline(ctx, subctx, sparse_compact_pipeline, - { mask_buf, sparse_buf }, sc_pc, - { nem1, nem2, nem3 }); - ggml_vk_sync_buffers(ctx, subctx); - } + // compact scratch layout: [512, kv_c, 1, ns] tightly packed, in K's own type + const uint32_t eff_nbk2 = fa_compact.active ? fa_compact.kv_c * fa_compact.row_bytes : nbk2_eff; + const uint32_t eff_nbk3 = fa_compact.active ? fa_compact.kv_c * fa_compact.row_bytes : nbk3_eff; + const uint32_t eff_nbv2 = fa_compact.active ? eff_nbk2 : nbv2_eff; + const uint32_t eff_nbv3 = fa_compact.active ? eff_nbk3 : nbv3_eff; const vk_flash_attn_push_constants pc = { N, KV, (uint32_t)ne1, (uint32_t)ne2, (uint32_t)ne3, @@ -11577,8 +13278,8 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx (uint32_t)nev2, (uint32_t)nev3, nem1, nem2, nem3, q_stride, (uint32_t)nbq2, (uint32_t)nbq3, - k_stride, nbk2_eff, nbk3_eff, - v_stride, nbv2_eff, nbv3_eff, + k_stride, eff_nbk2, eff_nbk3, + v_stride, eff_nbv2, eff_nbv3, scale, max_bias, logit_softcap, mask_n_head_log2, m0, m1, gqa_ratio, split_kv, split_k }; @@ -11602,11 +13303,11 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx vk_subbuffer split_k_buf = ggml_vk_subbuffer(ctx, ctx->prealloc_split_k, 0); ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, - {q_buf, k_buf, v_buf, mask_buf, sinks_buf, split_k_buf, mask_opt_buf, sparse_buf}, + {q_buf, k_buf, v_buf, mask_buf, sinks_buf, split_k_buf, mask_opt_buf, fa_compact.dynamic_kv ? fa_compact.kv_buf : q_buf}, pc, { dispatch_x, workgroups_y, workgroups_z }); ggml_vk_sync_buffers(ctx, subctx); - const vk_op_flash_attn_split_k_reduce_push_constants pc2 = { HSV, (uint32_t)ne1, (uint32_t)ne2, (uint32_t)ne3, split_k, (sinks != nullptr) }; + const vk_op_flash_attn_split_k_reduce_push_constants pc2 = { HSV, (uint32_t)ne1, (uint32_t)ne2, (uint32_t)ne2, (uint32_t)ne3, split_k, (sinks != nullptr) }; ggml_vk_dispatch_pipeline(ctx, subctx, ctx->device->pipeline_flash_attn_split_k_reduce, {split_k_buf, sinks_buf, dst_buf}, pc2, { (uint32_t)ne1, HSV, (uint32_t)(ne2 * ne3) }); @@ -11617,16 +13318,13 @@ static void ggml_vk_flash_attn(ggml_backend_vk_context * ctx, vk_context& subctx workgroups_x *= pipeline->wg_denoms[0]; } ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, - {q_buf, k_buf, v_buf, mask_buf, sinks_buf, dst_buf, mask_opt_buf, sparse_buf}, + {q_buf, k_buf, v_buf, mask_buf, sinks_buf, dst_buf, mask_opt_buf, fa_compact.dynamic_kv ? fa_compact.kv_buf : q_buf}, pc, { workgroups_x, workgroups_y, workgroups_z }); } if (use_dequant_kv) { ctx->prealloc_x_need_sync = true; } - if (use_mask_opt || use_sparse) { - ctx->prealloc_y_need_sync = true; - } } static vk_conv_shapes ggml_vk_conv_select_shape(ggml_backend_vk_context * ctx, uint32_t K, uint32_t NPQ) { @@ -11661,6 +13359,33 @@ static vk_conv_shapes ggml_vk_conv_select_shape(ggml_backend_vk_context * ctx, u } } +// The delta-net conv-state path does ggml_transpose() straight into a dim-0 ggml_concat(), so +// the generic concat kernel walks src1 with a conv_channels * 4 byte stride. On qwen35 that is +// 40960 B = 160 * 256 B and 160 % 16 == 0, so every read lands on one of the 16 memory channels: +// 13.7 GB/s against 138.9 GB/s for the tiled path. Route that exact shape to a tiled-transpose +// kernel. On by default; GGML_VK_CONCAT_TRANSPOSE=0 opts out. +// Qwen3.8-27B pp2048: +0.4% at ub 256, +4.7% at ub 1024, +7.2% at ub 2048. +static bool ggml_vk_concat_is_transposed(const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * dst) { + static const char * env = getenv("GGML_VK_CONCAT_TRANSPOSE"); + if (env && env[0] == '0') { + return false; + } + if (ggml_get_op_params_i32(dst, 0) != 0) { // dim 0 only + return false; + } + if (src0->ne[2] != 1 || src0->ne[3] != 1 || src1->ne[2] != 1 || src1->ne[3] != 1) { + return false; + } + const size_t ts = ggml_type_size(src0->type); + if (src0->nb[0] != ts || dst->nb[0] != ts) { // src0 and dst rows must be contiguous + return false; + } + if (src1->nb[0] <= src1->nb[1]) { // src1 must actually be transposed + return false; + } + return src0->ne[1] == src1->ne[1] && dst->ne[1] == src1->ne[1]; +} + static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const ggml_tensor * src0, const ggml_tensor * src1, const ggml_tensor * src2, const ggml_tensor * dst, ggml_op op) { switch (op) { case GGML_OP_GET_ROWS: @@ -11753,6 +13478,11 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const if (!ggml_vk_concat_supported(src0, src1, dst)) { return nullptr; } + // Tiled-transpose path handles unquantized 4-byte elements only. + if (!ggml_is_quantized(src0->type) && ggml_vk_concat_unit_size(src0->type) == 4 && + ggml_vk_concat_is_transposed(src0, src1, dst)) { + return ctx->device->pipeline_concat_transpose_i32; + } switch (ggml_vk_concat_unit_size(src0->type)) { case 1: return ctx->device->pipeline_concat_i8; @@ -12190,6 +13920,52 @@ static vk_pipeline ggml_vk_op_get_pipeline(ggml_backend_vk_context * ctx, const } return nullptr; case GGML_OP_LIGHTNING_INDEXER: + // fork fast path: f16 K on wave64 subgroup-arithmetic devices routes to the tuned + // scalar-64/CM kernels (head counts in LI_NH_VALUES); anything else falls through to + // the generic pipeline table + if (src0->type == GGML_TYPE_F32 && src1->type == GGML_TYPE_F16 && dst->type == GGML_TYPE_F32 && + src0->ne[0] == 128 && src1->ne[1] == 1 && + ggml_vk_li_nh_index(src0->ne[1]) >= 0 && ctx->device->pipeline_lightning_indexer_f16[0]) { + const int nhi = ggml_vk_li_nh_index(src0->ne[1]); + // Small-batch routing. Three arms, one env var each, most specific first: + // + // default decode CM for the whole 2-15 window + // ..._DECODE_CM_BATCH=0 small CM for 4-15, scalar for 2-3 (PR #6) + // ..._DECODE_CM_BATCH=0 SMALL_CM=0 scalar for 2-15 (pre-PR #6 baseline) + // + // The decode CM shader puts 16 HEADS in the coopmat N dimension and dispatches one + // workgroup per token, so it issues 4*n_batch tiles per 16 KV rows - the arithmetic + // minimum, since 64 heads fill its 16 columns exactly. The small CM shader puts 16 + // TOKENS in N and pays a flat 64 tiles no matter how small the batch. Routing the + // whole window to decode CM is 2.7x at batch 5 and 4.8x at batch 2-3 (526k source + // tokens, gfx1151); the shader body has no n_batch == 1 assumption, token is + // gl_WorkGroupID.y, so the old ne[2] == 1 gate was an artefact. + // + // That measurement only exists because of pepuscz's PR #6 and issue #10: their + // per-kernel table (43.8 us scalar vs 27.9 us small CM per 1k scanned rows at batch + // 5, the same ratio at every depth) is what showed the cost is a per-tile constant, + // which is what makes the tile count the thing to minimise. Their small CM route is + // kept as the opt-out arm rather than deleted: the tile-count argument is hardware + // independent, but decode CM re-reads the K tile once per token, and that part is + // bandwidth dependent, so the crossover need not sit here on other devices. + static const char * decode_cm_batch_env = getenv("GGML_VK_LIGHTNING_INDEXER_DECODE_CM_BATCH"); + static const bool decode_cm_batch_on = !decode_cm_batch_env || decode_cm_batch_env[0] != '0'; + const int64_t decode_cm_max = decode_cm_batch_on ? 15 : 1; + if (ctx->device->pipeline_lightning_indexer_decode_cm_f16[nhi] && src0->ne[2] <= decode_cm_max) { + return ctx->device->pipeline_lightning_indexer_decode_cm_f16[nhi]; + } + // PR #6 as its author proposed it, now default-on so that the kill switch above is + // enough on its own to reach it. + static const char * small_cm_env = getenv("GGML_VK_LIGHTNING_INDEXER_SMALL_CM"); + static const bool small_cm_on = !small_cm_env || small_cm_env[0] != '0'; + if (ctx->device->pipeline_lightning_indexer_cm_small_f16[nhi] && small_cm_on && + src0->ne[2] >= 4 && src0->ne[2] < 16) { + return ctx->device->pipeline_lightning_indexer_cm_small_f16[nhi]; + } + vk_pipeline cm = ctx->device->pipeline_lightning_indexer_cm_f16[nhi] ? + ctx->device->pipeline_lightning_indexer_cm_f16[nhi] : ctx->device->pipeline_lightning_indexer_cm_small_f16[nhi]; + return cm && src0->ne[2] >= 16 ? cm : ctx->device->pipeline_lightning_indexer_f16[nhi]; + } // only the k type selects a pipeline, the other types are fixed by ggml_lightning_indexer() if (ggml_vk_lightning_indexer_k_type_supported(src1->type)) { return ctx->device->pipeline_lightning_indexer_f32[src1->type]; @@ -12785,6 +14561,11 @@ static void ggml_vk_op_f32(ggml_backend_vk_context * ctx, vk_context& subctx, co case GGML_OP_GLU: case GGML_OP_CONV_2D_DW: { + // The tiled concat kernel is dispatched per 32x32 tile, not per element. + if (op == GGML_OP_CONCAT && pipeline == ctx->device->pipeline_concat_transpose_i32) { + elements = { (uint32_t)src1->ne[1], (uint32_t)src1->ne[0], 1 }; + break; + } uint32_t ne = ggml_nelements(dst); if (op == GGML_OP_CPY && ggml_is_quantized(src0->type) && ggml_is_quantized(dst->type)) { // Convert from number of logical elements to 2- or 4-byte units. @@ -13323,6 +15104,8 @@ static void ggml_vk_gated_linear_attn(ggml_backend_vk_context * ctx, vk_context& pc, { (uint32_t)(n_seqs * n_heads), 1, 1 }); } +static void ggml_vk_lightning_indexer_cm(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst, vk_pipeline pipeline); + static void ggml_vk_lightning_indexer(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) { const ggml_tensor * q = dst->src[0]; const ggml_tensor * k = dst->src[1]; @@ -13332,6 +15115,19 @@ static void ggml_vk_lightning_indexer(ggml_backend_vk_context * ctx, vk_context& vk_pipeline pipeline = ggml_vk_op_get_pipeline(ctx, q, k, w, dst, dst->op); GGML_ASSERT(pipeline != nullptr); + // the fork's wave64 f16 kernels take their own push-constant layout + bool fork_li = false; + for (size_t nhi = 0; nhi < LI_NH_COUNT && !fork_li; ++nhi) { + fork_li = pipeline == ctx->device->pipeline_lightning_indexer_f16[nhi] || + pipeline == ctx->device->pipeline_lightning_indexer_cm_f16[nhi] || + pipeline == ctx->device->pipeline_lightning_indexer_cm_small_f16[nhi] || + pipeline == ctx->device->pipeline_lightning_indexer_decode_cm_f16[nhi]; + } + if (fork_li) { + ggml_vk_lightning_indexer_cm(ctx, subctx, dst, pipeline); + return; + } + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); const uint32_t n_kv = k->ne[2]; @@ -13429,6 +15225,124 @@ static void ggml_vk_gated_delta_net(ggml_backend_vk_context * ctx, vk_context& s pc, { H, n_seqs, S_v }); } +static void ggml_vk_lightning_indexer_cm(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst, vk_pipeline pipeline) { + const ggml_tensor * q = dst->src[0]; + const ggml_tensor * k = dst->src[1]; + const ggml_tensor * w = dst->src[2]; + const ggml_tensor * m = dst->src[3]; + + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + + const vk_op_lightning_indexer_cm_push_constants pc = { + (uint32_t) k->ne[2], + (uint32_t) q->ne[2], + (uint32_t) q->ne[3], + (uint32_t) m->ne[3], + (uint32_t) (dst->nb[1] / sizeof(float)), + (uint32_t) (dst->nb[3] / sizeof(float)), + (uint32_t) (q->nb[1] / sizeof(float)), + (uint32_t) (q->nb[2] / sizeof(float)), + (uint32_t) (q->nb[3] / sizeof(float)), + (uint32_t) (k->nb[2] / sizeof(ggml_fp16_t)), + (uint32_t) (k->nb[3] / sizeof(ggml_fp16_t)), + (uint32_t) (w->nb[1] / sizeof(float)), + (uint32_t) (w->nb[3] / sizeof(float)), + (uint32_t) (m->nb[1] / sizeof(ggml_fp16_t)), + (uint32_t) (m->nb[3] / sizeof(ggml_fp16_t)), + }; + + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + {ggml_vk_tensor_subbuffer(ctx, q), ggml_vk_tensor_subbuffer(ctx, k), + ggml_vk_tensor_subbuffer(ctx, w), ggml_vk_tensor_subbuffer(ctx, m), + ggml_vk_tensor_subbuffer(ctx, dst)}, + pc, {(uint32_t) k->ne[2], (uint32_t) q->ne[2], (uint32_t) q->ne[3]}); +} + +// DSv4 fused hyper-connection ops — ports of ggml-cuda/dsv4-hc.cu. Strides are passed in +// f32 elements; grids mirror the CUDA launch geometry (flat 1D, 256 threads per workgroup). + +static void ggml_vk_dsv4_hc_pre(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) { + const ggml_tensor * x = dst->src[0]; + const ggml_tensor * w = dst->src[1]; + + const uint32_t n_embd = (uint32_t) x->ne[0]; + const uint32_t hc = (uint32_t) x->ne[1]; + const uint32_t n_tokens = (uint32_t) x->ne[2]; + const uint32_t nr = n_embd * n_tokens; + + vk_pipeline pipeline = ctx->device->pipeline_dsv4_hc_pre_f32; + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + + const vk_op_dsv4_hc_pre_push_constants pc = { + n_embd, hc, nr, + (uint32_t)(x->nb[0] / sizeof(float)), (uint32_t)(x->nb[1] / sizeof(float)), (uint32_t)(x->nb[2] / sizeof(float)), + (uint32_t)(w->nb[0] / sizeof(float)), (uint32_t)(w->nb[1] / sizeof(float)), + (uint32_t)(dst->nb[0] / sizeof(float)), (uint32_t)(dst->nb[1] / sizeof(float)), + }; + + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + {ggml_vk_tensor_subbuffer(ctx, x), ggml_vk_tensor_subbuffer(ctx, w), + ggml_vk_tensor_subbuffer(ctx, dst)}, + pc, {nr, 1, 1}); +} + +static void ggml_vk_dsv4_hc_comb(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) { + const ggml_tensor * mixes = dst->src[0]; + const ggml_tensor * scale = dst->src[1]; + const ggml_tensor * base = dst->src[2]; + + const uint32_t n_tokens = (uint32_t) mixes->ne[1]; + const float eps = ggml_get_op_params_f32(dst, 0); + const int32_t n_iter = ggml_get_op_params_i32(dst, 1); + + vk_pipeline pipeline = ctx->device->pipeline_dsv4_hc_comb_f32; + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + + const vk_op_dsv4_hc_comb_push_constants pc = { + n_tokens, + (uint32_t)(mixes->nb[0] / sizeof(float)), (uint32_t)(mixes->nb[1] / sizeof(float)), + (uint32_t)(scale->nb[0] / sizeof(float)), + (uint32_t)(base->nb[0] / sizeof(float)), + (uint32_t)(dst->nb[0] / sizeof(float)), (uint32_t)(dst->nb[1] / sizeof(float)), (uint32_t)(dst->nb[2] / sizeof(float)), + eps, n_iter, + }; + + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + {ggml_vk_tensor_subbuffer(ctx, mixes), ggml_vk_tensor_subbuffer(ctx, scale), + ggml_vk_tensor_subbuffer(ctx, base), ggml_vk_tensor_subbuffer(ctx, dst)}, + pc, {n_tokens, 1, 1}); +} + +static void ggml_vk_dsv4_hc_post(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) { + const ggml_tensor * x = dst->src[0]; + const ggml_tensor * residual = dst->src[1]; + const ggml_tensor * post = dst->src[2]; + const ggml_tensor * comb = dst->src[3]; + + const uint32_t n_embd = (uint32_t) x->ne[0]; + const uint32_t n_tokens = (uint32_t) x->ne[1]; + const uint32_t hc = (uint32_t) residual->ne[1]; + const uint32_t nr = n_embd * hc * n_tokens; + + vk_pipeline pipeline = ctx->device->pipeline_dsv4_hc_post_f32; + ggml_pipeline_request_descriptor_sets(ctx, pipeline, 1); + + const vk_op_dsv4_hc_post_push_constants pc = { + n_embd, hc, nr, + (uint32_t)(x->nb[0] / sizeof(float)), (uint32_t)(x->nb[1] / sizeof(float)), + (uint32_t)(residual->nb[0] / sizeof(float)), (uint32_t)(residual->nb[1] / sizeof(float)), (uint32_t)(residual->nb[2] / sizeof(float)), + (uint32_t)(post->nb[0] / sizeof(float)), (uint32_t)(post->nb[1] / sizeof(float)), + (uint32_t)(comb->nb[0] / sizeof(float)), (uint32_t)(comb->nb[1] / sizeof(float)), (uint32_t)(comb->nb[2] / sizeof(float)), + (uint32_t)(dst->nb[0] / sizeof(float)), (uint32_t)(dst->nb[1] / sizeof(float)), (uint32_t)(dst->nb[2] / sizeof(float)), + }; + + ggml_vk_dispatch_pipeline(ctx, subctx, pipeline, + {ggml_vk_tensor_subbuffer(ctx, x), ggml_vk_tensor_subbuffer(ctx, residual), + ggml_vk_tensor_subbuffer(ctx, post), ggml_vk_tensor_subbuffer(ctx, comb), + ggml_vk_tensor_subbuffer(ctx, dst)}, + pc, {nr, 1, 1}); +} + static void ggml_vk_ssm_scan(ggml_backend_vk_context * ctx, vk_context& subctx, ggml_tensor * dst) { const ggml_tensor * src0 = dst->src[0]; const ggml_tensor * src1 = dst->src[1]; @@ -14480,6 +16394,9 @@ static void ggml_vk_argsort(ggml_backend_vk_context * ctx, vk_context& subctx, c } } +// Tournament-reduction top_k, only correct while k fits a single workgroup's +// tournament pipeline (see ggml_vk_topk's dispatcher and the large-k path below +// for why this can't just be given a bigger workgroup). static void ggml_vk_topk(ggml_backend_vk_context * ctx, vk_context& subctx, const ggml_tensor * src0, ggml_tensor * dst) { uint32_t ncols = src0->ne[0]; uint32_t nrows = ggml_nrows(src0); @@ -15266,22 +17183,106 @@ static void ggml_vk_test_matmul(ggml_backend_vk_context * ctx, size_t m, size_t const size_t y_ne = k * n * batch; const size_t d_ne = m * n * batch; - ggml_type x_type = std::is_same() ? GGML_TYPE_F32 : GGML_TYPE_F16; - ggml_type y_type = std::is_same() ? GGML_TYPE_F32 : GGML_TYPE_F16; - vk_matmul_pipeline_key mm_test_key{x_type, y_type, false, false}; - auto mm_test_it = ctx->device->pipeline_matmul.find(mm_test_key); - GGML_ASSERT(mm_test_it != ctx->device->pipeline_matmul.end() && !mm_test_it->second.empty()); - auto& mm_test_configs = mm_test_it->second; - GGML_ASSERT(shader_size >= 0 && shader_size < (int)mm_test_configs.size()); - - std::string shname = std::string(ggml_type_name(x_type)) + "_" + std::string(ggml_type_name(y_type)) + "_ALIGNED_" + std::to_string(shader_size); - vk_pipeline p = mm_test_configs[shader_size].aligned ? mm_test_configs[shader_size].aligned : mm_test_configs[shader_size].unaligned; + vk_pipeline p; + std::string shname; + if (shader_size == 0) { + if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32->a_s; + shname = "F32_ALIGNED_S"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32_f16->a_s; + shname = "F32_F16_ALIGNED_S"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16_f32.f32acc->a_s; + shname = "F16_F32_ALIGNED_S"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16.f32acc->a_s; + shname = "F16_ALIGNED_S"; + } else { + GGML_ABORT("fatal error"); + } + } else if (shader_size == 1) { + if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32->a_m; + shname = "F32_ALIGNED_M"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32_f16->a_m; + shname = "F32_F16_ALIGNED_M"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16_f32.f32acc->a_m; + shname = "F16_F32_ALIGNED_M"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16.f32acc->a_m; + shname = "F16_ALIGNED_M"; + } else { + GGML_ABORT("fatal error"); + } + } else if (shader_size == 2) { + if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32->a_l; + shname = "F32_ALIGNED_L"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32_f16->a_l; + shname = "F32_F16_ALIGNED_L"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16_f32.f32acc->a_l; + shname = "F16_F32_ALIGNED_L"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16.f32acc->a_l; + shname = "F16_ALIGNED_L"; + } else { + GGML_ABORT("fatal error"); + } + } else { + GGML_ASSERT(0); + } - const size_t kpad = ggml_vk_align_size(k, mm_test_configs[shader_size].align); + const size_t kpad = ggml_vk_align_size(k, p->align); if (k != kpad) { - p = mm_test_configs[shader_size].unaligned; - shname = std::string(ggml_type_name(x_type)) + "_" + std::string(ggml_type_name(y_type)) + "_" + std::to_string(shader_size); + if (shader_size == 0) { + if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32->s; + shname = "F32_S"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32_f16->s; + shname = "F32_F16_S"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16_f32.f32acc->s; + shname = "F16_F32_S"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16.f32acc->s; + shname = "F16_S"; + } + } else if (shader_size == 1) { + if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32->m; + shname = "F32_M"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32_f16->m; + shname = "F32_F16_M"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16_f32.f32acc->m; + shname = "F16_F32_M"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16.f32acc->m; + shname = "F16_M"; + } + } else if (shader_size == 2) { + if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32->l; + shname = "F32_L"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f32_f16->l; + shname = "F32_F16_L"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16_f32.f32acc->l; + shname = "F16_F32_L"; + } else if (std::is_same() && std::is_same()) { + p = ctx->device->pipeline_matmul_f16.f32acc->l; + shname = "F16_L"; + } + } } if (split_k > 1) { @@ -15716,34 +17717,46 @@ static void ggml_vk_test_dequant_matmul(ggml_backend_vk_context * ctx, size_t m, const size_t y_ne = k * n * batch; const size_t d_ne = m * n * batch; - ggml_type b_type = mmq ? GGML_TYPE_Q8_1 : GGML_TYPE_F32; - bool f16acc = ctx->device->fp16 && !mmq; - vk_matmul_pipeline_key dq_key{quant, b_type, false, f16acc}; - auto dq_it = ctx->device->pipeline_matmul.find(dq_key); - if (dq_it == ctx->device->pipeline_matmul.end() || dq_it->second.empty()) { - if (f16acc) { - dq_key.f16acc = false; - dq_it = ctx->device->pipeline_matmul.find(dq_key); - } - } - if (dq_it == ctx->device->pipeline_matmul.end() || dq_it->second.empty()) { - std::cerr << "error: no pipeline for ggml_vk_test_dequant_matmul " << ggml_type_name(quant) << std::endl; - return; - } - auto& dq_configs = dq_it->second; - if (shader_size >= (int)dq_configs.size()) { - std::cerr << "error: shader_size " << shader_size << " >= configs.size() " << dq_configs.size() << " for " << ggml_type_name(quant) << std::endl; - return; - } + vk_matmul_pipeline2 * pipelines; - std::string shname = std::string(ggml_type_name(quant)) + "_ALIGNED_" + std::to_string(shader_size); - vk_pipeline p = dq_configs[shader_size].aligned ? dq_configs[shader_size].aligned : dq_configs[shader_size].unaligned; + if (mmq) { + pipelines = ctx->device->pipeline_dequant_mul_mat_mat_q8_1; + } else { + pipelines = ctx->device->pipeline_dequant_mul_mat_mat; + } + + const bool fp16acc = ctx->device->fp16; + + vk_pipeline p; + std::string shname; + if (shader_size == 0) { + p = fp16acc ? pipelines[quant].f16acc->a_s : pipelines[quant].f32acc->a_s; + shname = std::string(ggml_type_name(quant)) + "_ALIGNED_S"; + } else if (shader_size == 1) { + p = fp16acc ? pipelines[quant].f16acc->a_m : pipelines[quant].f32acc->a_m; + shname = std::string(ggml_type_name(quant)) + "_ALIGNED_M"; + } else if (shader_size == 2) { + p = fp16acc ? pipelines[quant].f16acc->a_l : pipelines[quant].f32acc->a_l; + shname = std::string(ggml_type_name(quant)) + "_ALIGNED_L"; + } else { + GGML_ASSERT(0); + } - const size_t kpad = mmq ? 0 : ggml_vk_align_size(k, dq_configs[shader_size].align); + const size_t kpad = mmq ? 0 : ggml_vk_align_size(k, p->align); if (mmq || k != kpad) { - p = dq_configs[shader_size].unaligned; - shname = std::string(ggml_type_name(quant)) + "_" + std::to_string(shader_size); + if (shader_size == 0) { + p = fp16acc ? pipelines[quant].f16acc->s : pipelines[quant].f32acc->s; + shname = std::string(ggml_type_name(quant)) + "_S"; + } else if (shader_size == 1) { + p = fp16acc ? pipelines[quant].f16acc->m : pipelines[quant].f32acc->m; + shname = std::string(ggml_type_name(quant)) + "_M"; + } else if (shader_size == 2) { + p = fp16acc ? pipelines[quant].f16acc->l : pipelines[quant].f32acc->l; + shname = std::string(ggml_type_name(quant)) + "_L"; + } else { + GGML_ASSERT(0); + } } if (p == nullptr) { @@ -16475,18 +18488,6 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr case GGML_OP_CUMSUM: ggml_vk_cumsum(ctx, compute_ctx, src0, node); - break; - case GGML_OP_DSV4_HC_COMB: - ggml_vk_dsv4_hc_comb(ctx, compute_ctx, src0, src1, src2, node); - - break; - case GGML_OP_DSV4_HC_PRE: - ggml_vk_dsv4_hc_pre(ctx, compute_ctx, src0, src1, node); - - break; - case GGML_OP_DSV4_HC_POST: - ggml_vk_dsv4_hc_post(ctx, compute_ctx, src0, src1, src2, src3, node); - break; case GGML_OP_MEAN: ggml_vk_mean(ctx, compute_ctx, src0, node); @@ -16596,6 +18597,21 @@ static bool ggml_vk_build_graph(ggml_backend_vk_context * ctx, ggml_cgraph * cgr break; + case GGML_OP_DSV4_HC_PRE: + ggml_vk_dsv4_hc_pre(ctx, compute_ctx, node); + + break; + + case GGML_OP_DSV4_HC_COMB: + ggml_vk_dsv4_hc_comb(ctx, compute_ctx, node); + + break; + + case GGML_OP_DSV4_HC_POST: + ggml_vk_dsv4_hc_post(ctx, compute_ctx, node); + + break; + case GGML_OP_SSM_SCAN: ggml_vk_ssm_scan(ctx, compute_ctx, node); @@ -16751,8 +18767,11 @@ static void ggml_vk_cleanup(ggml_backend_vk_context * ctx) { ggml_vk_destroy_buffer(ctx->prealloc_y); ggml_vk_destroy_buffer(ctx->prealloc_split_k); ggml_vk_destroy_buffer(ctx->prealloc_add_rms_partials); + ggml_vk_destroy_buffer(ctx->fa_union_stat); ggml_vk_destroy_buffer(ctx->sync_staging); + memset(ctx->fa_union_est_ratio, 0, sizeof(ctx->fa_union_est_ratio)); + ctx->prealloc_y_last_pipeline_used = nullptr; ctx->prealloc_y_last_tensor_used = nullptr; ctx->prealloc_y_last_k_padded = false; @@ -17515,9 +19534,26 @@ static bool ggml_vk_can_fuse(const ggml_backend_vk_context * ctx, const struct g if (mmid != mul->src[0]) { return false; } - // mat-vec only + // EXPERIMENT (GGML_VK_MMID_SCALE_EPILOGUE=1): the tile shader can apply the scale as it + // writes out, which removes a full write+read of the matmul result at prefill. The + // coopmat2 shader has the binding but not the epilogue, so it stays on the old path. if (!ggml_vk_use_mul_mat_vec_id(cgraph, node_idx)) { - return false; + static const char * env = getenv("GGML_VK_MMID_SCALE_EPILOGUE"); + if (!(env && atoi(env) != 0) || ctx->device->coopmat2) { + return false; + } + // Shader indexes the scale as [token * nei0 + expert_slot]. + if (scale->type != GGML_TYPE_F32 || mul->type != GGML_TYPE_F32 || !ggml_is_contiguous(scale)) { + return false; + } + if (get_misalign_bytes(ctx, scale) != 0) { + return false; + } + // The shader indexes the scale as [token * nei0 + expert_slot] from row_ids, which + // carries no 4th dimension, so ne[3] must be 1 or later batches read the wrong scale. + return scale->ne[0] == 1 && mmid->ne[3] == 1 && scale->ne[3] == 1 && + scale->ne[1] == mmid->ne[1] && scale->ne[2] == mmid->ne[2] && + ggml_are_same_shape(mul, mmid); } // shaders assume the types match if (mmid->type != scale->type) { @@ -18145,14 +20181,22 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ctx->query_fusion_node_count.resize(ctx->num_queries); ctx->query_nodes.resize(ctx->num_queries); ctx->query_node_idx.resize(ctx->num_queries); + ctx->query_sub_names.resize(ctx->num_queries); } - ctx->device->device.resetQueryPool(ctx->query_pool, 0, cgraph->n_nodes+1); + // Reset the whole pool, not just n_nodes+1: sub-op marks consume slots past that. + ctx->device->device.resetQueryPool(ctx->query_pool, 0, ctx->num_queries); std::fill(ctx->query_fusion_names.begin(), ctx->query_fusion_names.end(), nullptr); std::fill(ctx->query_fusion_node_count.begin(), ctx->query_fusion_node_count.end(), 0); std::fill(ctx->query_nodes.begin(), ctx->query_nodes.end(), nullptr); std::fill(ctx->query_node_idx.begin(), ctx->query_node_idx.end(), 0); + // Under partial offload the scheduler's async input copies between graph + // splits can leave commands in a pending compute ctx. Flush it so the + // timestamp stream starts on a fresh command buffer. + if (!ctx->compute_ctx.expired()) { + ggml_vk_synchronize(ctx); + } GGML_ASSERT(ctx->compute_ctx.expired()); compute_ctx = ggml_vk_get_compute_ctx(ctx); ctx->query_idx = 0; @@ -18179,6 +20223,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg uint32_t submitted_nodes = 0; uint32_t submit_count = 0; uint64_t batch_flops = 0; + uint64_t batch_bytes = 0; uint64_t total_flops = 0; uint64_t flops_cap = 200'000'000'000ULL; @@ -18216,6 +20261,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg first_node_in_batch = true; submitted_nodes = 0; batch_flops = 0; + batch_bytes = 0; if (submit_count < 3) { flops_per_submit *= 2; } @@ -18230,9 +20276,11 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg { auto node_flops = ggml_vk_get_node_flops(cgraph->nodes[i]); total_flops += node_flops; + auto node_bytes = ggml_vk_get_node_bytes(cgraph->nodes[i]); - // Flush the current batch before recording a node that would push it over the flop threshold - if (flops_per_submit != 0 && submitted_nodes > 0 && batch_flops + node_flops >= flops_per_submit) { + // Flush the current batch before recording a node that would push it over the flop or byte threshold + if ((flops_per_submit != 0 && submitted_nodes > 0 && batch_flops + node_flops >= flops_per_submit) || + (ctx->device->max_bytes_per_submit != 0 && submitted_nodes > 0 && batch_bytes + node_bytes >= ctx->device->max_bytes_per_submit)) { vk_context flush_ctx = ggml_vk_get_compute_ctx(ctx); ggml_vk_ctx_end(flush_ctx); flush_ctx->exit_tensor_idx = -1; @@ -18243,6 +20291,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg } batch_flops += node_flops; + batch_bytes += node_bytes; } // op_srcs_fused_elementwise indicates whether an op's srcs all contribute to @@ -18497,6 +20546,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg bool almost_ready = (cgraph->n_nodes - i) < cgraph->n_nodes / 5; bool submit = (submitted_nodes >= ctx->device->max_nodes_per_submit) || (flops_per_submit != 0 && batch_flops >= flops_per_submit) || + (ctx->device->max_bytes_per_submit != 0 && batch_bytes >= ctx->device->max_bytes_per_submit) || (i + ctx->num_additional_fused_ops >= last_node) || (almost_ready && !ctx->almost_ready_fence_pending); @@ -18508,6 +20558,7 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg // track a single node/fusion for the current query ctx->query_nodes[ctx->query_idx] = cgraph->nodes[i]; ctx->query_fusion_names[ctx->query_idx] = fusion_string; + ctx->query_sub_names[ctx->query_idx] = nullptr; compute_ctx->s->buffer->buf.writeTimestamp(vk::PipelineStageFlagBits::eAllCommands, ctx->query_pool, ctx->query_idx++); ggml_vk_sync_buffers(ctx, compute_ctx); } else { @@ -18549,14 +20600,22 @@ static ggml_status ggml_backend_vk_graph_compute(ggml_backend_t backend, ggml_cg ctx->compute_ctx.reset(); // Get the results and pass them to the logger - std::vector timestamps(cgraph->n_nodes + 1); - VK_CHECK(ctx->device->device.getQueryPoolResults(ctx->query_pool, 0, ctx->query_idx, (cgraph->n_nodes + 1)*sizeof(uint64_t), timestamps.data(), sizeof(uint64_t), vk::QueryResultFlagBits::e64 | vk::QueryResultFlagBits::eWait), "get timestamp results", ctx->device); + // Sized to the pool, not n_nodes+1: sub-op marks push query_idx past the node count. + std::vector timestamps(ctx->num_queries); + VK_CHECK(ctx->device->device.getQueryPoolResults(ctx->query_pool, 0, ctx->query_idx, ctx->num_queries*sizeof(uint64_t), timestamps.data(), sizeof(uint64_t), vk::QueryResultFlagBits::e64 | vk::QueryResultFlagBits::eWait), "get timestamp results", ctx->device); if (!vk_perf_logger_concurrent) { // Log each op separately for (int i = 1; i < ctx->query_idx; i++) { + const uint64_t dt = uint64_t((timestamps[i] - timestamps[i-1]) * ctx->device->properties.limits.timestampPeriod); + if (ctx->query_sub_names[i] != nullptr) { + // sub-node interval (e.g. the FA K/V contiguize pass) - billed separately so + // it is not silently folded into the op that dispatched it + ctx->perf_logger->log_timing_named(ctx->query_sub_names[i], dt); + continue; + } auto node = ctx->query_nodes[i]; auto name = ctx->query_fusion_names[i]; - ctx->perf_logger->log_timing(node, name, uint64_t((timestamps[i] - timestamps[i-1]) * ctx->device->properties.limits.timestampPeriod)); + ctx->perf_logger->log_timing(node, name, dt); } } else { // Log each group of nodes @@ -19347,8 +21406,13 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm case GGML_TYPE_IQ4_NL: case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: - case GGML_TYPE_TQ1_0: + case GGML_TYPE_Q4_0_ROCMFP4: + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + case GGML_TYPE_Q3_0_ROCMFPX: + case GGML_TYPE_Q6_0_ROCMFPX: + case GGML_TYPE_Q8_0_ROCMFPX: case GGML_TYPE_TQ2_0: + case GGML_TYPE_TQ1_0: break; default: return false; @@ -19397,25 +21461,34 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm if (op->src[3] && op->src[3]->type != GGML_TYPE_F16) { return false; } - auto fa_kv_ok = [](ggml_type t) { - switch (t) { - case GGML_TYPE_F32: - case GGML_TYPE_F16: - case GGML_TYPE_BF16: - case GGML_TYPE_Q8_0: - case GGML_TYPE_Q5_1: - case GGML_TYPE_Q5_0: - case GGML_TYPE_Q4_1: - case GGML_TYPE_Q4_0: - case GGML_TYPE_IQ4_NL: - return true; - default: - return false; - } + auto fa_kv_ok = [&](ggml_type t) { + // ggml_vk_fa_kv_native is the single source of truth; on this base every + // admitted type is native, and the non-native hard-gate below is dormant + // hardening for any future type routed through the dequant-once scratch. + return ggml_vk_fa_kv_native(t, coopmat2); }; if (!fa_kv_ok(op->src[1]->type) || !fa_kv_ok(op->src[2]->type)) { return false; } + if (!ggml_vk_fa_kv_native(op->src[1]->type, coopmat2) || !ggml_vk_fa_kv_native(op->src[2]->type, coopmat2)) { + // Only correct through the dequant-once scratch path; admit only when every + // hard condition of the dispatch-time gate holds, so dispatch can never fall + // back to the native shader (it reads garbage for these types, not an error). + const ggml_tensor * k = op->src[1]; + const ggml_tensor * v = op->src[2]; + static const char * fa_dequant_env = getenv("GGML_VK_FA_DEQUANT"); + const bool fa_dequant_off = fa_dequant_env && fa_dequant_env[0] == '0'; + const uint64_t kv_f16_sz = ((uint64_t)ggml_nelements(k) + (uint64_t)ggml_nelements(v)) * sizeof(ggml_fp16_t); + if (fa_dequant_off || + op->src[0]->ne[1] < 64 || + device->pipeline_dequant_transpose[k->type] == nullptr || + device->pipeline_dequant_transpose[v->type] == nullptr || + k->nb[0] != ggml_type_size(k->type) || v->nb[0] != ggml_type_size(v->type) || + !ggml_is_contiguously_allocated(k) || !ggml_is_contiguously_allocated(v) || + kv_f16_sz > device->properties.limits.maxStorageBufferRange) { + return false; + } + } if ((op->src[1]->type == GGML_TYPE_BF16) != (op->src[2]->type == GGML_TYPE_BF16)) { return false; } @@ -19454,8 +21527,13 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm case GGML_TYPE_IQ4_NL: case GGML_TYPE_MXFP4: case GGML_TYPE_NVFP4: - case GGML_TYPE_TQ1_0: + case GGML_TYPE_Q4_0_ROCMFP4: + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + case GGML_TYPE_Q3_0_ROCMFPX: + case GGML_TYPE_Q6_0_ROCMFPX: + case GGML_TYPE_Q8_0_ROCMFPX: case GGML_TYPE_TQ2_0: + case GGML_TYPE_TQ1_0: case GGML_TYPE_I32: return true; default: @@ -19482,6 +21560,11 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm case GGML_TYPE_Q5_1: case GGML_TYPE_Q8_0: case GGML_TYPE_IQ4_NL: + case GGML_TYPE_Q4_0_ROCMFP4: + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + case GGML_TYPE_Q3_0_ROCMFPX: + case GGML_TYPE_Q6_0_ROCMFPX: + case GGML_TYPE_Q8_0_ROCMFPX: return true; default: return false; @@ -19507,6 +21590,11 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm case GGML_TYPE_Q5_1: case GGML_TYPE_Q8_0: case GGML_TYPE_IQ4_NL: + case GGML_TYPE_Q4_0_ROCMFP4: + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + case GGML_TYPE_Q3_0_ROCMFPX: + case GGML_TYPE_Q6_0_ROCMFPX: + case GGML_TYPE_Q8_0_ROCMFPX: return true; default: break; @@ -19524,6 +21612,11 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm case GGML_TYPE_Q5_1: case GGML_TYPE_Q8_0: case GGML_TYPE_IQ4_NL: + case GGML_TYPE_Q4_0_ROCMFP4: + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + case GGML_TYPE_Q3_0_ROCMFPX: + case GGML_TYPE_Q6_0_ROCMFPX: + case GGML_TYPE_Q8_0_ROCMFPX: return true; default: break; @@ -19678,31 +21771,6 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm } return false; } - case GGML_OP_DSV4_HC_COMB: - case GGML_OP_DSV4_HC_PRE: - case GGML_OP_DSV4_HC_POST: - { - if (op->type != GGML_TYPE_F32) { - return false; - } - for (uint32_t i = 0; i < GGML_MAX_SRC; ++i) { - if (op->src[i] && op->src[i]->type != GGML_TYPE_F32) { - return false; - } - } - // hc is hardcoded to 4 in the shaders. ggml only constrains it - // to 4 for COMB, so PRE/POST have to be checked here. - if (op->op == GGML_OP_DSV4_HC_PRE && op->src[0]->ne[1] != 4) { - return false; - } - if (op->op == GGML_OP_DSV4_HC_POST && op->src[1]->ne[1] != 4) { - return false; - } - if (op->op == GGML_OP_DSV4_HC_COMB) { - return device->pipeline_dsv4_hc_comb_f32 != nullptr; - } - return true; - } case GGML_OP_SOLVE_TRI: { if (op->type != GGML_TYPE_F32 || op->src[0]->type != GGML_TYPE_F32) { @@ -19760,6 +21828,16 @@ static bool ggml_backend_vk_device_supports_op(ggml_backend_dev_t dev, const ggm case GGML_OP_GATED_LINEAR_ATTN: // the shader block size is hardcoded to head_size 64 return op->src[0]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32 && op->src[0]->ne[0] == 64; + case GGML_OP_DSV4_HC_PRE: + return op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && + op->type == GGML_TYPE_F32; + case GGML_OP_DSV4_HC_COMB: + return op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && + op->src[2]->type == GGML_TYPE_F32 && op->type == GGML_TYPE_F32; + case GGML_OP_DSV4_HC_POST: + return op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32 && + op->src[2]->type == GGML_TYPE_F32 && op->src[3]->type == GGML_TYPE_F32 && + op->type == GGML_TYPE_F32; case GGML_OP_LIGHTNING_INDEXER: { const ggml_tensor * q = op->src[0]; @@ -20691,13 +22769,6 @@ static void ggml_vk_check_results_0(ggml_backend_vk_context * ctx, ggml_cgraph * tensor_clone = ggml_sum_rows(ggml_ctx, src_clone[0]); } else if (tensor->op == GGML_OP_CUMSUM) { tensor_clone = ggml_cumsum(ggml_ctx, src_clone[0]); - } else if (tensor->op == GGML_OP_DSV4_HC_COMB) { - tensor_clone = ggml_dsv4_hc_comb(ggml_ctx, src_clone[0], src_clone[1], src_clone[2], - ggml_get_op_params_f32(tensor, 0), ggml_get_op_params_i32(tensor, 1)); - } else if (tensor->op == GGML_OP_DSV4_HC_PRE) { - tensor_clone = ggml_dsv4_hc_pre(ggml_ctx, src_clone[0], src_clone[1]); - } else if (tensor->op == GGML_OP_DSV4_HC_POST) { - tensor_clone = ggml_dsv4_hc_post(ggml_ctx, src_clone[0], src_clone[1], src_clone[2], src_clone[3]); } else if (tensor->op == GGML_OP_MEAN) { tensor_clone = ggml_mean(ggml_ctx, src_clone[0]); } else if (tensor->op == GGML_OP_ARGMAX) { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/concat_transpose.comp b/ggml/src/ggml-vulkan/vulkan-shaders/concat_transpose.comp new file mode 100644 index 000000000000..653aeaa011b2 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/concat_transpose.comp @@ -0,0 +1,43 @@ +#version 450 + +#include "types.glsl" +#include "generic_binary_head.glsl" + +layout(local_size_x = 32, local_size_y = 8, local_size_z = 1) in; + +// dim-0 concat whose src1 is transposed. The generic kernel reads src1 with the transposed +// stride, so neighbouring lanes touch different cache lines. Stage a 32x32 tile in shared +// memory instead, which keeps both the load and the store coalesced. 33 columns pads away +// the shared-memory bank conflicts. +shared A_TYPE tmp[32][33]; + +void main() { + const uint tx = gl_LocalInvocationID.x; + const uint ty = gl_LocalInvocationID.y; + + const uint row = gl_WorkGroupID.x * 32 + tx; + + // src0 is already contiguous, copy it straight through. + if (gl_WorkGroupID.y == 0 && row < p.ne01) { + for (uint i0 = ty; i0 < p.ne00; i0 += 8) { + data_d[get_doffset() + row*p.nb21 + i0*p.nb20] = D_TYPE(data_a[get_aoffset() + row*p.nb01 + i0*p.nb00]); + } + } + + [[unroll]] for (uint j = 0; j < 32; j += 8) { + const uint c = gl_WorkGroupID.y * 32 + ty + j; + if (c < p.ne10 && row < p.ne11) { + tmp[ty + j][tx] = A_TYPE(data_b[get_boffset() + c*p.nb10 + row*p.nb11]); + } + } + + barrier(); + + const uint col = gl_WorkGroupID.y * 32 + tx; + [[unroll]] for (uint j = 0; j < 32; j += 8) { + const uint r = gl_WorkGroupID.x * 32 + ty + j; + if (col < p.ne10 && r < p.ne11) { + data_d[get_doffset() + r*p.nb21 + (p.ne00 + col)*p.nb20] = D_TYPE(tmp[tx][ty + j]); + } + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/copy_from_quant.comp b/ggml/src/ggml-vulkan/vulkan-shaders/copy_from_quant.comp index 6a6921474781..75bdfe47231f 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/copy_from_quant.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/copy_from_quant.comp @@ -4,7 +4,7 @@ #include "generic_unary_head.glsl" #include "dequant_funcs.glsl" -#if defined(DATA_A_IQ4_NL) || defined(DATA_A_MXFP4) || defined(DATA_A_NVFP4) +#if defined(DATA_A_IQ4_NL) || defined(DATA_A_MXFP4) || defined(DATA_A_NVFP4) || defined(DATA_A_ROCMFP_ANY) // 16 invocations needed for init_iq_shmem layout(local_size_x = 16, local_size_y = 1, local_size_z = 1) in; #else diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/copy_quant_same.comp b/ggml/src/ggml-vulkan/vulkan-shaders/copy_quant_same.comp new file mode 100644 index 000000000000..4b28c42a31ed --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/copy_quant_same.comp @@ -0,0 +1,55 @@ +#version 450 + +#include "types.glsl" +#include "generic_unary_head.glsl" + +layout(local_size_x = 128, local_size_y = 1, local_size_z = 1) in; + +#if !defined(BLOCK_BYTES) +#error BLOCK_BYTES must be defined +#endif + +uint src0_block_idx(uint idx) { + const uint ne00b = p.ne00 / QUANT_K; + const uint nb02 = p.ne01 * ne00b; + const uint nb03 = p.ne02 * nb02; + + const uint i03 = idx / nb03; + const uint r03 = idx - i03 * nb03; + const uint i02 = r03 / nb02; + const uint r02 = r03 - i02 * nb02; + const uint i01 = r02 / ne00b; + const uint i00b = r02 - i01 * ne00b; + + return i03*p.nb03 + i02*p.nb02 + i01*p.nb01 + i00b*p.nb00; +} + +uint dst_block_idx(uint idx) { + const uint ne10b = p.ne10 / QUANT_K; + const uint nb12 = p.ne11 * ne10b; + const uint nb13 = p.ne12 * nb12; + + const uint i13 = idx / nb13; + const uint r13 = idx - i13 * nb13; + const uint i12 = r13 / nb12; + const uint r12 = r13 - i12 * nb12; + const uint i11 = r12 / ne10b; + const uint i10b = r12 - i11 * ne10b; + + return i13*p.nb13 + i12*p.nb12 + i11*p.nb11 + i10b*p.nb10; +} + +void main() { + const uint idx = get_idx(); + + if (idx >= p.ne) { + return; + } + + const uint src_base = get_aoffset() + src0_block_idx(idx) * BLOCK_BYTES; + const uint dst_base = get_doffset() + dst_block_idx(idx) * BLOCK_BYTES; + + [[unroll]] for (uint j = 0; j < BLOCK_BYTES; ++j) { + data_d[dst_base + j] = data_a[src_base + j]; + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/copy_to_quant.comp b/ggml/src/ggml-vulkan/vulkan-shaders/copy_to_quant.comp index 776e9b8a5572..a0707c145958 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/copy_to_quant.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/copy_to_quant.comp @@ -12,6 +12,10 @@ const uint BLOCK_SIZE = 32; layout (binding = 0) readonly buffer S {S_TYPE data_s[];}; +float load_s(uint idx) { + return float(data_s[idx]); +} + #if defined(SET_ROWS) #include "generic_binary_head.glsl" layout (binding = 1) readonly buffer C {B_TYPE data_i[];}; @@ -286,6 +290,486 @@ void quantize(uint dst_idx, uint src_idx) } #endif +#if defined(DATA_A_ROCMFP4) || defined(DATA_A_ROCMFP4_FAST) +float rocmfp4_code_value(uint q) { + switch (q & 0xFu) { + case 0u: return 0.0; + case 1u: return 1.0; + case 2u: return 2.0; + case 3u: return 3.0; + case 4u: return 4.0; + case 5u: return 6.0; + case 6u: return 8.0; + case 7u: return 10.0; + case 8u: return 0.0; + case 9u: return -1.0; + case 10u: return -2.0; + case 11u: return -3.0; + case 12u: return -4.0; + case 13u: return -6.0; + case 14u: return -8.0; + default: return -10.0; + } +} + +uint rocmfp4_best_index_scaled(float x, float inv_scale_half) { + const float a = abs(x * inv_scale_half); + if (a <= 0.5) { + return 0u; + } + + const bool neg = x < 0.0; + if (a <= 1.5) return neg ? 9u : 1u; + if (a <= 2.5) return neg ? 10u : 2u; + if (a <= 3.5) return neg ? 11u : 3u; + if (a <= 5.0) return neg ? 12u : 4u; + if (a <= 7.0) return neg ? 13u : 5u; + if (a <= 9.0) return neg ? 14u : 6u; + + return neg ? 15u : 7u; +} + +float rocmfp4_block_mse_for_scale(uint src_idx, uint src_off, uint n, uint e, float best_err) { + const float scale_half = ue4m3_to_fp32(uint8_t(e)); + const float inv_scale_half = scale_half > 0.0 ? 1.0 / scale_half : 0.0; + float err = 0.0; + + for (uint j = 0; j < n; ++j) { + const float x = load_s(src_idx + src_off + j); + const uint q = rocmfp4_best_index_scaled(x, inv_scale_half); + const float y = rocmfp4_code_value(q) * scale_half; + const float d = x - y; + err += d * d; + if (err > best_err) { + return err; + } + } + + return err; +} + +uint rocmfp4_nearest_scale_ue4m3(float target_scale_half) { + if (target_scale_half <= 0.0) { + return 1u; + } + + uint lo = 1u; + uint hi = 126u; + while (lo < hi) { + const uint mid = lo + (hi - lo) / 2u; + if (ue4m3_to_fp32(uint8_t(mid)) < target_scale_half) { + lo = mid + 1u; + } else { + hi = mid; + } + } + + if (lo == 1u) { + return 1u; + } + + const float hi_scale = ue4m3_to_fp32(uint8_t(lo)); + const float lo_scale = ue4m3_to_fp32(uint8_t(lo - 1u)); + + // Match the previous ascending nearest scan: exact midpoint ties keep the + // lower scale byte. + return (target_scale_half - lo_scale <= hi_scale - target_scale_half) ? lo - 1u : lo; +} + +uint rocmfp4_choose_scale_ue4m3(uint src_idx, uint src_off, uint n) { + float max_abs = 0.0; + for (uint j = 0; j < n; ++j) { + max_abs = max(max_abs, abs(load_s(src_idx + src_off + j))); + } + if (max_abs <= 0.0) { + return 0u; + } + + const uint start_e = rocmfp4_nearest_scale_ue4m3(max_abs / 10.0); + uint best_e = 0u; + float best_err = 3.4028234663852886e38; + bool lower_done = false; + + for (uint delta = 0u; delta <= 125u; ++delta) { + const int e0 = int(start_e) - int(delta); + if (!lower_done && e0 >= 1 && e0 <= 126) { + const float scale_half = ue4m3_to_fp32(uint8_t(e0)); + const float clip_delta = max_abs - 10.0 * scale_half; + if (clip_delta > 0.0 && clip_delta * clip_delta > best_err) { + lower_done = true; + } else { + const uint e0u = uint(e0); + const float err = rocmfp4_block_mse_for_scale(src_idx, src_off, n, e0u, best_err); + if (err < best_err || (err == best_err && e0u < best_e)) { + best_err = err; + best_e = e0u; + } + } + } + + const uint e1 = start_e + delta; + if (delta != 0u && e1 <= 126u) { + const float err = rocmfp4_block_mse_for_scale(src_idx, src_off, n, e1, best_err); + if (err < best_err || (err == best_err && e1 < best_e)) { + best_err = err; + best_e = e1; + } + } + + if ((lower_done || e0 <= 1) && e1 >= 126u) { + break; + } + } + + return best_e; +} + +#endif + +#if defined(DATA_A_ROCMFP4) +void quantize(uint dst_idx, uint src_idx) +{ + const uint e0 = rocmfp4_choose_scale_ue4m3(src_idx, 0u, QUANT_K_ROCMFP4 / 2); + const uint e1 = rocmfp4_choose_scale_ue4m3(src_idx, QUANT_K_ROCMFP4 / 2, QUANT_K_ROCMFP4 / 2); + const float d0 = ue4m3_to_fp32(uint8_t(e0)); + const float d1 = ue4m3_to_fp32(uint8_t(e1)); + const float id0 = d0 > 0.0 ? 1.0 / d0 : 0.0; + const float id1 = d1 > 0.0 ? 1.0 / d1 : 0.0; + + data_q[dst_idx].e[0] = uint8_t(e0); + data_q[dst_idx].e[1] = uint8_t(e1); + + [[unroll]] for (uint j = 0u; j < QUANT_K_ROCMFP4 / 2; ++j) { + const uint q0 = rocmfp4_best_index_scaled(load_s(src_idx + j), id0); + const uint q1 = rocmfp4_best_index_scaled(load_s(src_idx + QUANT_K_ROCMFP4 / 2 + j), id1); + data_q[dst_idx].qs[j] = uint8_t(q0 | (q1 << 4)); + } +} +#endif + +#if defined(DATA_A_ROCMFP4_FAST) +void quantize(uint dst_idx, uint src_idx) +{ + const uint e = rocmfp4_choose_scale_ue4m3(src_idx, 0u, QUANT_K_ROCMFP4); + const float d = ue4m3_to_fp32(uint8_t(e)); + const float id = d > 0.0 ? 1.0 / d : 0.0; + + data_q[dst_idx].e = uint8_t(e); + + [[unroll]] for (uint j = 0u; j < QUANT_K_ROCMFP4 / 2; ++j) { + const uint q0 = rocmfp4_best_index_scaled(load_s(src_idx + j), id); + const uint q1 = rocmfp4_best_index_scaled(load_s(src_idx + QUANT_K_ROCMFP4 / 2 + j), id); + data_q[dst_idx].qs[j] = uint8_t(q0 | (q1 << 4)); + } +} +#endif + +#if defined(DATA_A_ROCMFPX_FAMILY) +uint rocmfpx_nearest_scale_ue4m3(float target_scale) { + if (target_scale <= 0.0) { + return 0u; + } + + uint lo = 1u; + uint hi = 126u; + while (lo < hi) { + const uint mid = lo + (hi - lo) / 2u; + if (ue4m3_to_fp32(uint8_t(mid)) < target_scale) { + lo = mid + 1u; + } else { + hi = mid; + } + } + + if (lo == 1u) { + return 1u; + } + + const float hi_scale = ue4m3_to_fp32(uint8_t(lo)); + const float lo_scale = ue4m3_to_fp32(uint8_t(lo - 1u)); + return (target_scale - lo_scale <= hi_scale - target_scale) ? lo - 1u : lo; +} +#endif + +#if defined(DATA_A_ROCMFPX_FP3) +uint rocmfpx_fp3_quantize_code(float x, float inv_scale) { + if (inv_scale <= 0.0) { + return 0u; + } + + const float ax = abs(x * inv_scale); + uint mag; + if (ax <= 0.5) { + mag = 0u; + } else if (ax <= 1.5) { + mag = 1u; + } else if (ax <= 3.0) { + mag = 2u; + } else { + mag = 3u; + } + + return mag == 0u ? 0u : ((x < 0.0 ? 4u : 0u) | mag); +} + +int rocmfpx_fp3_decode_value(uint code) { + const int mag = (code & 3u) == 3u ? 4 : int(code & 3u); + return (code & 4u) != 0u ? -mag : mag; +} + +float rocmfpx_fp3_mse_for_scale(uint src_idx, uint offset, uint e, float best_err) { + const float scale = ue4m3_to_fp32(uint8_t(e)); + const float inv_scale = scale > 0.0 ? 1.0 / scale : 0.0; + float err = 0.0; + + [[unroll]] for (uint j = 0u; j < 16u; ++j) { + const float x = load_s(src_idx + offset + j); + const uint code = rocmfpx_fp3_quantize_code(x, inv_scale); + const float y = float(rocmfpx_fp3_decode_value(code)) * scale; + const float d = x - y; + + err += d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +uint rocmfpx_choose_scale_fp3_mse(uint src_idx, uint offset) { + float max_abs = 0.0; + [[unroll]] for (uint j = 0u; j < 16u; ++j) { + max_abs = max(max_abs, abs(load_s(src_idx + offset + j))); + } + + if (max_abs <= 0.0) { + return 0u; + } + + const uint start_e = rocmfpx_nearest_scale_ue4m3(max_abs / 4.0); + uint best_e = start_e; + float best_err = 3.402823466e+38; + bool lower_done = false; + + for (uint delta = 0u; delta <= 125u; ++delta) { + const int e0 = int(start_e) - int(delta); + if (!lower_done && e0 >= 1 && e0 <= 126) { + const float scale = ue4m3_to_fp32(uint8_t(uint(e0))); + const float clip_delta = max_abs - 4.0 * scale; + if (clip_delta > 0.0 && clip_delta*clip_delta > best_err) { + lower_done = true; + } else { + const float err = rocmfpx_fp3_mse_for_scale(src_idx, offset, uint(e0), best_err); + if (err < best_err || (err == best_err && uint(e0) < best_e)) { + best_err = err; + best_e = uint(e0); + } + } + } + + const uint e1 = start_e + delta; + if (delta != 0u && e1 >= 1u && e1 <= 126u) { + const float err = rocmfpx_fp3_mse_for_scale(src_idx, offset, e1, best_err); + if (err < best_err || (err == best_err && e1 < best_e)) { + best_err = err; + best_e = e1; + } + } + + if ((lower_done || e0 <= 1) && e1 >= 126u) { + break; + } + } + + return best_e; +} + +void rocmfpx_fp3_set_bits(inout uint qs[12], uint bit_pos, uint code) { + [[unroll]] for (uint bit = 0u; bit < 3u; ++bit) { + const uint dst_bit = bit_pos + bit; + const uint byte_index = dst_bit >> 3u; + const uint mask = 1u << (dst_bit & 7u); + uint v = qs[byte_index]; + v = (code & (1u << bit)) != 0u ? (v | mask) : (v & ~mask); + qs[byte_index] = v; + } +} + +void quantize(uint dst_idx, uint src_idx) +{ + uint qs[12]; + [[unroll]] for (uint j = 0u; j < 12u; ++j) { + qs[j] = 0u; + } + + [[unroll]] for (uint part = 0u; part < 2u; ++part) { + const uint offset = part * 16u; + const uint e = rocmfpx_choose_scale_fp3_mse(src_idx, offset); + const float d = ue4m3_to_fp32(uint8_t(e)); + const float id = d > 0.0 ? 1.0 / d : 0.0; + data_q[dst_idx].e[part] = uint8_t(e); + + [[unroll]] for (uint j = 0u; j < 16u; ++j) { + const uint i = offset + j; + rocmfpx_fp3_set_bits(qs, i * 3u, rocmfpx_fp3_quantize_code(load_s(src_idx + i), id)); + } + } + + [[unroll]] for (uint j = 0u; j < 12u; ++j) { + data_q[dst_idx].qs[j] = uint8_t(qs[j]); + } +} +#endif + +#if defined(DATA_A_ROCMFPX_FP6) +uint rocmfpx_fp6_quantize_code(float x, float inv_scale) { + if (inv_scale <= 0.0) { + return 0u; + } + + // asymmetric signed range [-32, 31]; -32 is encoded as sign|0 + // Clamp in float before the int conversion: int() of an out-of-range float is + // undefined in GLSL, so an extreme finite weight could flip sign or vanish. + const int q = int(clamp(round(x * inv_scale), -32.0, 31.0)); + if (q == 0) { + return 0u; + } + return q < 0 ? (32u | (uint(-q) & 31u)) : uint(q); +} + +int rocmfpx_fp6_decode_value(uint code) { + const int mag = int(code & 31u); + return (code & 32u) != 0u ? -(mag == 0 ? 32 : mag) : mag; +} + +float rocmfpx_fp6_mse_for_scale(uint src_idx, uint offset, uint e, float best_err) { + const float scale = ue4m3_to_fp32(uint8_t(e)); + const float inv_scale = scale > 0.0 ? 1.0 / scale : 0.0; + float err = 0.0; + + [[unroll]] for (uint j = 0u; j < 16u; ++j) { + const float x = load_s(src_idx + offset + j); + const uint code = rocmfpx_fp6_quantize_code(x, inv_scale); + const float y = float(rocmfpx_fp6_decode_value(code)) * scale; + const float d = x - y; + + err += d*d; + if (err > best_err) { + return err; + } + } + + return err; +} + +uint rocmfpx_choose_scale_fp6_mse(uint src_idx, uint offset) { + float max_abs = 0.0; + [[unroll]] for (uint j = 0u; j < 16u; ++j) { + max_abs = max(max_abs, abs(load_s(src_idx + offset + j))); + } + + if (max_abs <= 0.0) { + return 0u; + } + + const uint start_e = rocmfpx_nearest_scale_ue4m3(max_abs / 31.0); + uint best_e = start_e; + float best_err = 3.402823466e+38; + bool lower_done = false; + + for (uint delta = 0u; delta <= 125u; ++delta) { + const int e0 = int(start_e) - int(delta); + if (!lower_done && e0 >= 1 && e0 <= 126) { + const float scale = ue4m3_to_fp32(uint8_t(uint(e0))); + // ROCmFP6 reaches -32, not just 31, so bound the unavoidable clipping + // error by 32 - matching rocmfpx_choose_scale_fp6_mse_impl() on the CPU. + // At 31 the bound is too pessimistic for a block whose largest magnitude + // is negative and the search stops before the scale that actually wins. + const float clip_delta = max_abs - 32.0 * scale; + if (clip_delta > 0.0 && clip_delta*clip_delta > best_err) { + lower_done = true; + } else { + const float err = rocmfpx_fp6_mse_for_scale(src_idx, offset, uint(e0), best_err); + if (err < best_err || (err == best_err && uint(e0) < best_e)) { + best_err = err; + best_e = uint(e0); + } + } + } + + const uint e1 = start_e + delta; + if (delta != 0u && e1 >= 1u && e1 <= 126u) { + const float err = rocmfpx_fp6_mse_for_scale(src_idx, offset, e1, best_err); + if (err < best_err || (err == best_err && e1 < best_e)) { + best_err = err; + best_e = e1; + } + } + + if ((lower_done || e0 <= 1) && e1 >= 126u) { + break; + } + } + + return best_e; +} + +void quantize(uint dst_idx, uint src_idx) +{ + [[unroll]] for (uint part = 0u; part < 2u; ++part) { + const uint offset = part * 16u; + const uint e = rocmfpx_choose_scale_fp6_mse(src_idx, offset); + const float d = ue4m3_to_fp32(uint8_t(e)); + const float id = d > 0.0 ? 1.0 / d : 0.0; + data_q[dst_idx].e[part] = uint8_t(e); + + // pack four 6-bit codes into three bytes, little-endian + [[unroll]] for (uint g = 0u; g < 4u; ++g) { + const uint i0 = offset + g * 4u; + const uint c0 = rocmfpx_fp6_quantize_code(load_s(src_idx + i0 + 0u), id); + const uint c1 = rocmfpx_fp6_quantize_code(load_s(src_idx + i0 + 1u), id); + const uint c2 = rocmfpx_fp6_quantize_code(load_s(src_idx + i0 + 2u), id); + const uint c3 = rocmfpx_fp6_quantize_code(load_s(src_idx + i0 + 3u), id); + + const uint b = (offset >> 4u) * 12u + g * 3u; + data_q[dst_idx].qs[b + 0u] = uint8_t(( c0 & 0x3Fu) | ((c1 & 0x03u) << 6)); + data_q[dst_idx].qs[b + 1u] = uint8_t(((c1 >> 2u) & 0x0Fu) | ((c2 & 0x0Fu) << 4)); + data_q[dst_idx].qs[b + 2u] = uint8_t(((c2 >> 4u) & 0x03u) | ((c3 & 0x3Fu) << 2)); + } + } +} +#endif + +#if defined(DATA_A_ROCMFPX_FP8) + +int rocmfpx_fp8_quantize_code(float x, float inv_scale) { + if (inv_scale <= 0.0) { + return 0; + } + + return int(clamp(round(x * inv_scale), -127.0, 127.0)); +} + +void quantize(uint dst_idx, uint src_idx) +{ + float max_abs = 0.0; + [[unroll]] for (uint j = 0u; j < QUANT_K_ROCMFPX_FP8; ++j) { + max_abs = max(max_abs, abs(load_s(src_idx + j))); + } + + const uint e = rocmfpx_nearest_scale_ue4m3(max_abs / 127.0); + const float d = ue4m3_to_fp32(uint8_t(e)); + const float id = d > 0.0 ? 1.0 / d : 0.0; + + data_q[dst_idx].e = uint8_t(e); + + [[unroll]] for (uint j = 0u; j < QUANT_K_ROCMFPX_FP8; ++j) { + data_q[dst_idx].qs[j] = int8_t(rocmfpx_fp8_quantize_code(load_s(src_idx + j), id)); + } +} +#endif + #if defined(DATA_A_F32) || defined(DATA_A_F16) void quantize(uint dst_idx, uint src_idx) { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_f16_transpose.comp b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_f16_transpose.comp new file mode 100644 index 000000000000..dda53d749acc --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_f16_transpose.comp @@ -0,0 +1,32 @@ +#version 450 + +#include "dequant_head.glsl" + +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {f16vec4 data_a[];}; +layout (binding = 1) writeonly buffer D {f16vec4 data_b[];}; + +// Strided-copy counterpart of the fused dequant+transpose shaders for FA f16 KV: +// source physical is [HS, NH, KV, NS] (p.M=HS, p.K=NH, p.stride_a=KV); write to +// per-head-contiguous dest [HS, KV, NH, NS] so the f16 FA reads KV coalesced. +// HS stays innermost in both layouts, so each invocation moves 8 HS-contiguous +// elements (two f16vec4) requiring HS % 8 == 0 (enforced by the host gate). +void main() { + const uint i = gl_GlobalInvocationID.x; + const uint e0 = i * 8; + if (e0 >= p.nel) { + return; + } + + const uint HS = p.M, NH = p.K, KVn = p.stride_a; + const uint dst = (e0 % HS) + + ((e0 / (HS * NH)) % KVn) * HS + + ((e0 / HS) % NH) * (HS * KVn) + + (e0 / (HS * NH * KVn)) * (HS * KVn * NH); + + data_b[dst / 4 ] = data_a[e0 / 4 ]; + data_b[dst / 4 + 1] = data_a[e0 / 4 + 1]; +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs.glsl index 9df66cb44f9d..169557b6e25d 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs.glsl @@ -510,6 +510,152 @@ vec4 dequantize4(uint ib, uint iqs, uint a_offset) { } #endif +#if defined(DATA_A_ROCMFP4) +vec2 dequantize(uint ib, uint iqs, uint a_offset) { + const uint vui = uint(data_a[a_offset + ib].qs[iqs]); + const float d0 = ue4m3_to_fp32(data_a[a_offset + ib].e[0]); + const float d1 = ue4m3_to_fp32(data_a[a_offset + ib].e[1]); + return vec2(float(kvalues_rocmfp4[vui & 0xF]) * d0, + float(kvalues_rocmfp4[vui >> 4]) * d1); +} +vec4 dequantize4(uint ib, uint iqs, uint a_offset) { + const uint vui0 = uint(data_a[a_offset + ib].qs[iqs]); + const uint vui1 = uint(data_a[a_offset + ib].qs[iqs + 1]); + const float d0 = ue4m3_to_fp32(data_a[a_offset + ib].e[0]); + const float d1 = ue4m3_to_fp32(data_a[a_offset + ib].e[1]); + return vec4(float(kvalues_rocmfp4[vui0 & 0xF]) * d0, + float(kvalues_rocmfp4[vui0 >> 4]) * d1, + float(kvalues_rocmfp4[vui1 & 0xF]) * d0, + float(kvalues_rocmfp4[vui1 >> 4]) * d1); +} +#endif + +#if defined(DATA_A_ROCMFP4_FAST) +vec2 dequantize(uint ib, uint iqs, uint a_offset) { + const uint vui = uint(data_a[a_offset + ib].qs[iqs]); + const float d = ue4m3_to_fp32(data_a[a_offset + ib].e); + return vec2(float(kvalues_rocmfp4[vui & 0xF]) * d, + float(kvalues_rocmfp4[vui >> 4]) * d); +} +vec4 dequantize4(uint ib, uint iqs, uint a_offset) { + const uint vui0 = uint(data_a[a_offset + ib].qs[iqs]); + const uint vui1 = uint(data_a[a_offset + ib].qs[iqs + 1]); + const float d = ue4m3_to_fp32(data_a[a_offset + ib].e); + return vec4(float(kvalues_rocmfp4[vui0 & 0xF]) * d, + float(kvalues_rocmfp4[vui0 >> 4]) * d, + float(kvalues_rocmfp4[vui1 & 0xF]) * d, + float(kvalues_rocmfp4[vui1 >> 4]) * d); +} +#endif + +#if defined(DATA_A_ROCMFPX_FP3) +uint rocmfpx_fp3_get_bits(uint ib, uint bit_pos, uint a_offset) { + uint code = 0u; + [[unroll]] for (uint bit = 0u; bit < 3u; ++bit) { + const uint src_bit = bit_pos + bit; + code |= ((uint(data_a[a_offset + ib].qs[src_bit >> 3u]) >> (src_bit & 7u)) & 1u) << bit; + } + return code; +} + +int rocmfpx_fp3_decode_code(uint code) { + return int(kvalues_rocmfpx_fp3_const[code & 7u]); +} + +float rocmfpx_fp3_dequant(uint ib, uint idx, uint a_offset) { + const float d = ue4m3_to_fp32(data_a[a_offset + ib].e[idx >= 16u ? 1u : 0u]); + return float(rocmfpx_fp3_decode_code(rocmfpx_fp3_get_bits(ib, idx * 3u, a_offset))) * d; +} + +// idx is a multiple of 4 at every call site, so bit_pos is a multiple of 12 and +// the shift is 0 or 4: the twelve bits always fit in two bytes, no third byte +// and no branch. +int32_t rocmfpx_fp3_pack4_window(uint ib, uint idx, uint a_offset) { + const uint bit_pos = idx * 3u; + const uint byte_pos = bit_pos >> 3u; + const uint sh = bit_pos & 7u; + uint bits = uint(data_a[a_offset + ib].qs[byte_pos]) | + (uint(data_a[a_offset + ib].qs[byte_pos + 1u]) << 8); + bits = (bits >> sh) & 0xFFFu; + return pack32(i8vec4(kvalues_rocmfpx_fp3_const[ bits & 7u], + kvalues_rocmfpx_fp3_const[(bits >> 3) & 7u], + kvalues_rocmfpx_fp3_const[(bits >> 6) & 7u], + kvalues_rocmfpx_fp3_const[(bits >> 9) & 7u])); +} + +vec4 rocmfpx_fp3_dequant4(uint ib, uint idx, uint a_offset) { + const vec4 q = vec4(unpack8(rocmfpx_fp3_pack4_window(ib, idx, a_offset))); + // idx is a multiple of 4 and the halves split at element 16, so all four + // weights share one scale + return q * ue4m3_to_fp32(data_a[a_offset + ib].e[idx >= 16u ? 1u : 0u]); +} + +vec2 dequantize(uint ib, uint iqs, uint a_offset) { + return vec2(rocmfpx_fp3_dequant(ib, iqs + 0u, a_offset), + rocmfpx_fp3_dequant(ib, iqs + 1u, a_offset)); +} + +vec4 dequantize4(uint ib, uint iqs, uint a_offset) { + return rocmfpx_fp3_dequant4(ib, iqs, a_offset); +} +#endif + +#if defined(DATA_A_ROCMFPX_FP6) +uint rocmfpx_fp6_get_bits(uint ib, uint idx, uint a_offset) { + const uint bit_pos = idx * 6u; + const uint byte_pos = bit_pos >> 3u; + const uint sh = bit_pos & 7u; + uint bits = uint(data_a[a_offset + ib].qs[byte_pos]); + if (sh > 2u) { + bits |= uint(data_a[a_offset + ib].qs[byte_pos + 1u]) << 8; + } + return (bits >> sh) & 0x3Fu; +} + +float rocmfpx_fp6_dequant(uint ib, uint idx, uint a_offset) { + const float d = ue4m3_to_fp32(data_a[a_offset + ib].e[idx >= 16u ? 1u : 0u]); + return float(rocmfpx_fp6_decode_code(rocmfpx_fp6_get_bits(ib, idx, a_offset))) * d; +} + +vec2 dequantize(uint ib, uint iqs, uint a_offset) { + return vec2(rocmfpx_fp6_dequant(ib, iqs + 0u, a_offset), + rocmfpx_fp6_dequant(ib, iqs + 1u, a_offset)); +} + +// Every caller passes a multiple of 4, which makes this the fast case: four +// 6-bit codes are exactly three whole bytes, so the gather needs no per-weight +// bit window and no branch, and all four share one scale (the halves split at +// element 16). rocmfpx_fp6_dequant() above stays for the unaligned callers. +vec4 dequantize4(uint ib, uint iqs, uint a_offset) { + const uint byte_pos = (iqs >> 2u) * 3u; + const uint bits = uint(data_a[a_offset + ib].qs[byte_pos ]) + | uint(data_a[a_offset + ib].qs[byte_pos + 1u]) << 8 + | uint(data_a[a_offset + ib].qs[byte_pos + 2u]) << 16; + + const float d = ue4m3_to_fp32(data_a[a_offset + ib].e[iqs >= 16u ? 1u : 0u]); + + return vec4(float(rocmfpx_fp6_decode_code( bits & 0x3Fu)), + float(rocmfpx_fp6_decode_code((bits >> 6u) & 0x3Fu)), + float(rocmfpx_fp6_decode_code((bits >> 12u) & 0x3Fu)), + float(rocmfpx_fp6_decode_code((bits >> 18u) & 0x3Fu))) * d; +} +#endif + +#if defined(DATA_A_ROCMFPX_FP8) +vec2 dequantize(uint ib, uint iqs, uint a_offset) { + const float d = ue4m3_to_fp32(data_a[a_offset + ib].e); + return vec2(float(int(data_a[a_offset + ib].qs[iqs + 0u])) * d, + float(int(data_a[a_offset + ib].qs[iqs + 1u])) * d); +} +vec4 dequantize4(uint ib, uint iqs, uint a_offset) { + const float d = ue4m3_to_fp32(data_a[a_offset + ib].e); + return vec4(float(int(data_a[a_offset + ib].qs[iqs + 0u])) * d, + float(int(data_a[a_offset + ib].qs[iqs + 1u])) * d, + float(int(data_a[a_offset + ib].qs[iqs + 2u])) * d, + float(int(data_a[a_offset + ib].qs[iqs + 3u])) * d); +} +#endif + #if defined(DATA_A_NVFP4) vec2 dequantize(uint ib, uint iqs, uint a_offset) { const uint sub = iqs >> 4; @@ -577,6 +723,13 @@ vec2 get_dm(uint ib, uint a_offset) { } #endif +#if defined(DATA_A_ROCMFP4) || defined(DATA_A_ROCMFP4_FAST) || defined(DATA_A_ROCMFPX_FAMILY) +// the ROCm FP formats fold their UE4M3 scales into dequantize() directly +vec2 get_dm(uint ib, uint a_offset) { + return vec2(1, 0); +} +#endif + #if defined(DATA_A_NVFP4) vec2 get_dm(uint ib, uint a_offset) { return vec2(1.0, 0.0); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.glsl index cc6e242a90d5..01dd393fb142 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_funcs_cm2.glsl @@ -260,20 +260,6 @@ float16_t dequantFuncTQ1_0(const in decodeBufTQ1_0 bl, const in uint blockCoords return bl.block.d * (float16_t(int(xi)) - float16_t(1.0)); } -f16vec4 dequantFuncTQ1_0_v(const in decodeBufTQ1_0 bl, const in uint blockCoords[2], const in uint coordInBlock[2]) -{ - const uint e = coordInBlock[1]; - f16vec4 v; - [[unroll]] for (uint k = 0u; k < 4u; ++k) { - const uint ee = e + k; - const uint bidx = tq1_0_byte_of(ee); - const uint qbyte = uint(bidx < 48u ? bl.block.qs[bidx] : bl.block.qh[bidx - 48u]); - const uint xi = tq1_0_trit(qbyte, tq1_0_digit_of(ee)); - v[k] = bl.block.d * (float16_t(int(xi)) - float16_t(1.0)); - } - return v; -} - layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufTQ2_0 { block_tq2_0 block; }; @@ -1068,7 +1054,7 @@ float16_t dequantFuncIQ2_S(const in decodeBufIQ2_S bl, const in uint blockCoords const uint scale = (bl.block.scales[ib32] >> ((idx & 0x10) >> 2)) & 0xf; const uint qs = bl.block.qs[ib8]; const uint qh = bl.block.qh[ib32]; - const uint sign = bl.block.qs[QUANT_K_IQ2_S / 8 + ib8] >> (idx & 0x6); + const uint sign = bl.block.qs[QUANT_K / 8 + ib8] >> (idx & 0x6); const float d = float(bl.block.d); const float db = d * 0.25 * (0.5 + scale); @@ -1090,7 +1076,7 @@ f16vec4 dequantFuncIQ2_S_v(const in decodeBufIQ2_S bl, const in uint blockCoords const uint scale = (bl.block.scales[ib32] >> ((idx & 0x10) >> 2)) & 0xf; const uint qs = bl.block.qs[ib8]; const uint qh = bl.block.qh[ib32]; - const uint sb = uint(bl.block.qs[QUANT_K_IQ2_S / 8 + ib8]) >> (idx & 0x6u); + const uint sb = uint(bl.block.qs[QUANT_K / 8 + ib8]) >> (idx & 0x6u); const float d = float(bl.block.d); const float db = d * 0.25 * (0.5 + scale); @@ -1121,7 +1107,7 @@ float16_t dequantFuncIQ3_XXS(const in decodeBufIQ3_XXS bl, const in uint blockCo uint idx = coordInBlock[1]; const uint iqs = (idx & 0xFC) >> 2; // 0..63 - const uint is = QUANT_K_IQ3_XXS / 4 + ((idx & 0xE0) >> 3);// 8 values + const uint is = QUANT_K / 4 + ((idx & 0xE0) >> 3);// 8 values const float d = float(bl.block.d); const uint qs = bl.block.qs[iqs]; @@ -1144,7 +1130,7 @@ f16vec4 dequantFuncIQ3_XXS_v(const in decodeBufIQ3_XXS bl, const in uint blockCo const uint idx = coordInBlock[1]; const uint iqs = idx >> 2; - const uint is = QUANT_K_IQ3_XXS / 4 + ((idx & 0xE0) >> 3); + const uint is = QUANT_K / 4 + ((idx & 0xE0) >> 3); const float d = float(bl.block.d); const uint qs = bl.block.qs[iqs]; @@ -1360,6 +1346,157 @@ f16vec4 dequantFuncMXFP4_v(const in decodeBufMXFP4 bl, const in uint blockCoords } #endif +#if defined(DATA_A_ROCMFP4) +layout(buffer_reference, std430, buffer_reference_align = 1) buffer decodeBufROCMFP4 { + block_rocmfp4 block; +}; + +float16_t dequantFuncROCMFP4(const in decodeBufROCMFP4 bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + const uint idx = coordInBlock[1]; + const uint iqs = idx & 0xF; + const uint shift = (idx & 0x10) >> 2; + const uint sub = idx >> 4; + const float d = ue4m3_to_fp32(bl.block.e[sub]); + uint32_t qs = bl.block.qs[iqs]; + qs >>= shift; + qs &= 0xF; + float16_t ret = float16_t(float(kvalues_rocmfp4[qs]) * d); + return ret; +} +#endif + +#if defined(DATA_A_ROCMFP4_FAST) +layout(buffer_reference, std430, buffer_reference_align = 1) buffer decodeBufROCMFP4Fast { + block_rocmfp4_fast block; +}; + +float16_t dequantFuncROCMFP4Fast(const in decodeBufROCMFP4Fast bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + const uint idx = coordInBlock[1]; + const uint iqs = idx & 0xF; + const uint shift = (idx & 0x10) >> 2; + const float d = ue4m3_to_fp32(bl.block.e); + uint32_t qs = bl.block.qs[iqs]; + qs >>= shift; + qs &= 0xF; + float16_t ret = float16_t(float(kvalues_rocmfp4[qs]) * d); + return ret; +} +#endif + +#if defined(DATA_A_ROCMFPX_FP3) +layout(buffer_reference, std430, buffer_reference_align = 1) buffer decodeBufROCMFPXFP3 { + block_rocmfpx_fp3 block; +}; + +uint rocmfpx_cm2_fp3_get_bits(const in decodeBufROCMFPXFP3 bl, uint bit_pos) +{ + uint code = 0u; + [[unroll]] for (uint bit = 0u; bit < 3u; ++bit) { + const uint src_bit = bit_pos + bit; + code |= ((uint(bl.block.qs[src_bit >> 3u]) >> (src_bit & 7u)) & 1u) << bit; + } + return code; +} + +int rocmfpx_cm2_fp3_decode(uint code) +{ + return int(kvalues_rocmfpx_fp3_const[code & 7u]); +} + +float16_t dequantFuncROCMFPXFP3(const in decodeBufROCMFPXFP3 bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + const uint idx = coordInBlock[1]; + const float d = ue4m3_to_fp32(bl.block.e[idx >= 16u ? 1u : 0u]); + return float16_t(float(rocmfpx_cm2_fp3_decode(rocmfpx_cm2_fp3_get_bits(bl, idx * 3u))) * d); +} + +int32_t rocmfpx_cm2_fp3_pack4_window(const in decodeBufROCMFPXFP3 bl, uint idx) +{ + const uint bit_pos = idx * 3u; + const uint byte_pos = bit_pos >> 3u; + const uint sh = bit_pos & 7u; + uint bits = uint(bl.block.qs[byte_pos]) | + (uint(bl.block.qs[byte_pos + 1u]) << 8); + if (sh > 4u) { + bits |= uint(bl.block.qs[byte_pos + 2u]) << 16; + } + bits = (bits >> sh) & 0xFFFu; + return pack32(i8vec4(kvalues_rocmfpx_fp3_const[ bits & 7u], + kvalues_rocmfpx_fp3_const[(bits >> 3) & 7u], + kvalues_rocmfpx_fp3_const[(bits >> 6) & 7u], + kvalues_rocmfpx_fp3_const[(bits >> 9) & 7u])); +} + +f16vec4 dequantFuncROCMFPXFP3_v(const in decodeBufROCMFPXFP3 bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + const uint idx = coordInBlock[1]; + const vec4 q = vec4(unpack8(rocmfpx_cm2_fp3_pack4_window(bl, idx))); + return f16vec4(q * vec4(ue4m3_to_fp32(bl.block.e[(idx + 0u) >= 16u ? 1u : 0u]), + ue4m3_to_fp32(bl.block.e[(idx + 1u) >= 16u ? 1u : 0u]), + ue4m3_to_fp32(bl.block.e[(idx + 2u) >= 16u ? 1u : 0u]), + ue4m3_to_fp32(bl.block.e[(idx + 3u) >= 16u ? 1u : 0u]))); +} +#endif + +#if defined(DATA_A_ROCMFPX_FP6) +layout(buffer_reference, std430, buffer_reference_align = 1) buffer decodeBufROCMFPXFP6 { + block_rocmfpx_fp6 block; +}; + +uint rocmfpx_cm2_fp6_get_bits(const in decodeBufROCMFPXFP6 bl, uint idx) +{ + const uint bit_pos = idx * 6u; + const uint byte_pos = bit_pos >> 3u; + const uint sh = bit_pos & 7u; + uint bits = uint(bl.block.qs[byte_pos]); + if (sh > 2u) { + bits |= uint(bl.block.qs[byte_pos + 1u]) << 8; + } + return (bits >> sh) & 0x3Fu; +} + +float16_t dequantFuncROCMFPXFP6(const in decodeBufROCMFPXFP6 bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + const uint idx = coordInBlock[1]; + const float d = ue4m3_to_fp32(bl.block.e[idx >= 16u ? 1u : 0u]); + return float16_t(float(rocmfpx_fp6_decode_code(rocmfpx_cm2_fp6_get_bits(bl, idx))) * d); +} + +f16vec4 dequantFuncROCMFPXFP6_v(const in decodeBufROCMFPXFP6 bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + const uint idx = coordInBlock[1]; + return f16vec4(float16_t(float(rocmfpx_fp6_decode_code(rocmfpx_cm2_fp6_get_bits(bl, idx + 0u))) * ue4m3_to_fp32(bl.block.e[(idx + 0u) >= 16u ? 1u : 0u])), + float16_t(float(rocmfpx_fp6_decode_code(rocmfpx_cm2_fp6_get_bits(bl, idx + 1u))) * ue4m3_to_fp32(bl.block.e[(idx + 1u) >= 16u ? 1u : 0u])), + float16_t(float(rocmfpx_fp6_decode_code(rocmfpx_cm2_fp6_get_bits(bl, idx + 2u))) * ue4m3_to_fp32(bl.block.e[(idx + 2u) >= 16u ? 1u : 0u])), + float16_t(float(rocmfpx_fp6_decode_code(rocmfpx_cm2_fp6_get_bits(bl, idx + 3u))) * ue4m3_to_fp32(bl.block.e[(idx + 3u) >= 16u ? 1u : 0u]))); +} +#endif + +#if defined(DATA_A_ROCMFPX_FP8) +layout(buffer_reference, std430, buffer_reference_align = 1) buffer decodeBufROCMFPXFP8 { + block_rocmfpx_fp8 block; +}; + +float16_t dequantFuncROCMFPXFP8(const in decodeBufROCMFPXFP8 bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + const uint idx = coordInBlock[1]; + const float d = ue4m3_to_fp32(bl.block.e); + return float16_t(float(int(bl.block.qs[idx])) * d); +} + +f16vec4 dequantFuncROCMFPXFP8_v(const in decodeBufROCMFPXFP8 bl, const in uint blockCoords[2], const in uint coordInBlock[2]) +{ + const uint idx = coordInBlock[1]; + const float d = ue4m3_to_fp32(bl.block.e); + return f16vec4(vec4(float(int(bl.block.qs[idx + 0u])), + float(int(bl.block.qs[idx + 1u])), + float(int(bl.block.qs[idx + 2u])), + float(int(bl.block.qs[idx + 3u]))) * d); +} +#endif + #if defined(DATA_A_NVFP4) layout(buffer_reference, std430, buffer_reference_align = 4) buffer decodeBufNVFP4 { block_nvfp4 block; @@ -1487,6 +1624,19 @@ f16vec4 dequantFuncNVFP4_v(const in decodeBufNVFP4 bl, const in uint blockCoords #elif defined(DATA_A_MXFP4) #define dequantFuncA dequantFuncMXFP4 #define dequantFuncA_v dequantFuncMXFP4_v +#elif defined(DATA_A_ROCMFP4) +#define dequantFuncA dequantFuncROCMFP4 +#elif defined(DATA_A_ROCMFP4_FAST) +#define dequantFuncA dequantFuncROCMFP4Fast +#elif defined(DATA_A_ROCMFPX_FP3) +#define dequantFuncA dequantFuncROCMFPXFP3 +#define dequantFuncA_v dequantFuncROCMFPXFP3_v +#elif defined(DATA_A_ROCMFPX_FP6) +#define dequantFuncA dequantFuncROCMFPXFP6 +#define dequantFuncA_v dequantFuncROCMFPXFP6_v +#elif defined(DATA_A_ROCMFPX_FP8) +#define dequantFuncA dequantFuncROCMFPXFP8 +#define dequantFuncA_v dequantFuncROCMFPXFP8_v #elif defined(DATA_A_NVFP4) #define dequantFuncA dequantFuncNVFP4 #define dequantFuncA_v dequantFuncNVFP4_v diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_iq4_nl.comp b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_iq4_nl.comp index 8f7833eab2e7..befcdc1c5100 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_iq4_nl.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_iq4_nl.comp @@ -10,8 +10,6 @@ layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; void main() { const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64; - init_iq_shmem(gl_WorkGroupSize); - const uint tid = gl_LocalInvocationID.x % 64; const uint il = tid/32; const uint ir = tid%32; @@ -21,12 +19,27 @@ void main() { } const uint q_idx = 8*il; + +#ifdef DEQUANT_TRANSPOSE + // Fused dequant+transpose for FA quant-KV: source physical is [HS, NH, KV, NS] + // (p.M=HS, p.K=NH, p.stride_a=KV); write to per-head-contiguous dest [HS, KV, NH, NS] so the + // f16 FA reads KV coalesced. An iq4_nl block = 32 consecutive HS elements at fixed (head,kv) -> + // 32 contiguous dest positions (intra-block nibble offsets unchanged). + const uint HS = p.M, NH = p.K, KVn = p.stride_a; + const uint e0 = ib * 32; + const uint b_idx = (e0 % HS) + + ((e0 / (HS * NH)) % KVn) * HS + + ((e0 / HS) % NH) * (HS * KVn) + + (e0 / (HS * NH * KVn)) * (HS * KVn * NH) + + q_idx; +#else const uint b_idx = 1024*i + 32*ir + q_idx; +#endif const float d = float(data_a[ib].d); [[unroll]] for (uint l = 0; l < 8; ++l) { - data_b[b_idx + l + 0] = D_TYPE(d * kvalues_iq4nl[data_a[ib].qs[q_idx + l] & 0xF]); - data_b[b_idx + l + 16] = D_TYPE(d * kvalues_iq4nl[data_a[ib].qs[q_idx + l] >> 4]); + data_b[b_idx + l + 0] = D_TYPE(d * float(kvalues_iq4nl_const[data_a[ib].qs[q_idx + l] & 0xF])); + data_b[b_idx + l + 16] = D_TYPE(d * float(kvalues_iq4nl_const[data_a[ib].qs[q_idx + l] >> 4])); } } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_0.comp b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_0.comp index b92b292135b4..51a6c89a60c8 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_0.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_0.comp @@ -19,7 +19,19 @@ void main() { } const uint q_idx = 8*il; + +#ifdef DEQUANT_TRANSPOSE + // fused dequant+transpose for FA quant-KV: per-head-contiguous f16 scratch (see dequant_q8_0.comp) + const uint HS = p.M, NH = p.K, KVn = p.stride_a; + const uint e0 = ib * 32; + const uint b_idx = (e0 % HS) + + ((e0 / (HS * NH)) % KVn) * HS + + ((e0 / HS) % NH) * (HS * KVn) + + (e0 / (HS * NH * KVn)) * (HS * KVn * NH) + + q_idx; +#else const uint b_idx = 1024*i + 32*ir + q_idx; +#endif const float d = float(data_a[ib].d); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_1.comp index 6b63cbe5833b..76f2d958cbab 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q4_1.comp @@ -18,7 +18,18 @@ void main() { return; } +#ifdef DEQUANT_TRANSPOSE + // fused dequant+transpose for FA quant-KV: per-head-contiguous f16 scratch (see dequant_q8_0.comp) + const uint HS = p.M, NH = p.K, KVn = p.stride_a; + const uint e0 = ib * 32; + const uint b_idx = (e0 % HS) + + ((e0 / (HS * NH)) % KVn) * HS + + ((e0 / HS) % NH) * (HS * KVn) + + (e0 / (HS * NH * KVn)) * (HS * KVn * NH) + + 8*il; +#else const uint b_idx = 1024*i + 32*ir + 8*il; +#endif const float d = float(data_a[ib].d); const float m = float(data_a[ib].m); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_0.comp b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_0.comp index f1b0bac87271..1402fa4292d8 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_0.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_0.comp @@ -18,7 +18,18 @@ void main() { return; } +#ifdef DEQUANT_TRANSPOSE + // fused dequant+transpose for FA quant-KV: per-head-contiguous f16 scratch (see dequant_q8_0.comp) + const uint HS = p.M, NH = p.K, KVn = p.stride_a; + const uint e0 = ib * 32; + const uint b_idx = (e0 % HS) + + ((e0 / (HS * NH)) % KVn) * HS + + ((e0 / HS) % NH) * (HS * KVn) + + (e0 / (HS * NH * KVn)) * (HS * KVn * NH) + + 8*il; +#else const uint b_idx = 1024*i + 32*ir + 8*il; +#endif const float d = float(data_a[ib].d); const uint qh = uint(data_a[ib].qh[1]) << 16 | data_a[ib].qh[0]; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_1.comp index c495b31f1754..1fd6e2552af7 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_q5_1.comp @@ -18,7 +18,18 @@ void main() { return; } +#ifdef DEQUANT_TRANSPOSE + // fused dequant+transpose for FA quant-KV: per-head-contiguous f16 scratch (see dequant_q8_0.comp) + const uint HS = p.M, NH = p.K, KVn = p.stride_a; + const uint e0 = ib * 32; + const uint b_idx = (e0 % HS) + + ((e0 / (HS * NH)) % KVn) * HS + + ((e0 / HS) % NH) * (HS * KVn) + + (e0 / (HS * NH * KVn)) * (HS * KVn * NH) + + 8*il; +#else const uint b_idx = 1024*i + 32*ir + 8*il; +#endif const float d = float(data_a[ib].d); const float m = float(data_a[ib].m); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_rocmfp4.comp b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_rocmfp4.comp new file mode 100644 index 000000000000..46cc1bd7dc5b --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_rocmfp4.comp @@ -0,0 +1,33 @@ +#version 450 + +#include "dequant_head.glsl" + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {block_rocmfp4 data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; + +void main() { + const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64; + + init_iq_shmem(gl_WorkGroupSize); + + const uint tid = gl_LocalInvocationID.x % 64; + const uint il = tid/32; + const uint ir = tid%32; + const uint ib = 32*i + ir; + if (ib >= p.nel / 32) { + return; + } + + const uint q_idx = 8*il; + const uint b_idx = 1024*i + 32*ir + q_idx; + + const float d0 = ue4m3_to_fp32(data_a[ib].e[0]); + const float d1 = ue4m3_to_fp32(data_a[ib].e[1]); + + [[unroll]] for (uint l = 0; l < 8; ++l) { + data_b[b_idx + l + 0] = D_TYPE(d0 * float(kvalues_rocmfp4[data_a[ib].qs[q_idx + l] & 0xF])); + data_b[b_idx + l + 16] = D_TYPE(d1 * float(kvalues_rocmfp4[data_a[ib].qs[q_idx + l] >> 4])); + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_rocmfp4_fast.comp b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_rocmfp4_fast.comp new file mode 100644 index 000000000000..360ff679f398 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_rocmfp4_fast.comp @@ -0,0 +1,31 @@ +#version 450 + +#include "dequant_head.glsl" + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {block_rocmfp4_fast data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; + +void main() { + const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64; + + init_iq_shmem(gl_WorkGroupSize); + + const uint tid = gl_LocalInvocationID.x % 64; + const uint il = tid/32; + const uint ir = tid%32; + const uint ib = 32*i + ir; + if (ib >= p.nel / 32) { + return; + } + + const uint q_idx = 8*il; + const uint b_idx = 1024*i + 32*ir + q_idx; + const float d = ue4m3_to_fp32(data_a[ib].e); + + [[unroll]] for (uint l = 0; l < 8; ++l) { + data_b[b_idx + l + 0] = D_TYPE(d * float(kvalues_rocmfp4[data_a[ib].qs[q_idx + l] & 0xF])); + data_b[b_idx + l + 16] = D_TYPE(d * float(kvalues_rocmfp4[data_a[ib].qs[q_idx + l] >> 4])); + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_rocmfpx_fp3.comp b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_rocmfpx_fp3.comp new file mode 100644 index 000000000000..e3aab966d769 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_rocmfpx_fp3.comp @@ -0,0 +1,46 @@ +#version 450 + +#include "dequant_head.glsl" + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {block_rocmfpx_fp3 data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; + +uint rocmfpx_fp3_get_bits(uint ib, uint bit_pos) { + uint code = 0u; + [[unroll]] for (uint bit = 0u; bit < 3u; ++bit) { + const uint src_bit = bit_pos + bit; + code |= ((uint(data_a[ib].qs[src_bit >> 3u]) >> (src_bit & 7u)) & 1u) << bit; + } + return code; +} + +int rocmfpx_fp3_decode_code(uint code) { + const uint mag_code = code & 3u; + const int mag = mag_code == 3u ? 4 : int(mag_code); + return (code & 4u) != 0u ? -mag : mag; +} + +void main() { + const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64; + + init_iq_shmem(gl_WorkGroupSize); + + const uint tid = gl_LocalInvocationID.x % 64; + const uint il = tid/32; + const uint ir = tid%32; + const uint ib = 32*i + ir; + if (ib >= p.nel / 32) { + return; + } + + const uint b_idx = 1024*i + 32*ir + 16*il; + const uint q_idx = 16*il; + const float d = ue4m3_to_fp32(data_a[ib].e[il]); + + [[unroll]] for (uint l = 0; l < 16; ++l) { + const uint idx = q_idx + l; + data_b[b_idx + l] = D_TYPE(d * float(rocmfpx_fp3_decode_code(rocmfpx_fp3_get_bits(ib, idx * 3u)))); + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_rocmfpx_fp6.comp b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_rocmfpx_fp6.comp new file mode 100644 index 000000000000..5da0e9e8cf98 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_rocmfpx_fp6.comp @@ -0,0 +1,44 @@ +#version 450 + +#include "dequant_head.glsl" + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {block_rocmfpx_fp6 data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; + +uint rocmfpx_fp6_get_bits(uint ib, uint idx) { + const uint bit_pos = idx * 6u; + const uint byte_pos = bit_pos >> 3u; + const uint sh = bit_pos & 7u; + uint bits = uint(data_a[ib].qs[byte_pos]); + if (sh > 2u) { + bits |= uint(data_a[ib].qs[byte_pos + 1u]) << 8; + } + return (bits >> sh) & 0x3Fu; +} + +void main() { + const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64; + +#ifdef NEEDS_INIT_IQ_SHMEM + init_iq_shmem(gl_WorkGroupSize); +#endif + + const uint tid = gl_LocalInvocationID.x % 64; + const uint il = tid/32; + const uint ir = tid%32; + const uint ib = 32*i + ir; + if (ib >= p.nel / 32) { + return; + } + + const uint b_idx = 1024*i + 32*ir + 16*il; + const uint q_idx = 16*il; + const float d = ue4m3_to_fp32(data_a[ib].e[il]); + + [[unroll]] for (uint l = 0; l < 16; ++l) { + const uint idx = q_idx + l; + data_b[b_idx + l] = D_TYPE(d * float(rocmfpx_fp6_decode_code(rocmfpx_fp6_get_bits(ib, idx)))); + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dequant_rocmfpx_fp8.comp b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_rocmfpx_fp8.comp new file mode 100644 index 000000000000..7052de542470 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dequant_rocmfpx_fp8.comp @@ -0,0 +1,31 @@ +#version 450 + +#include "dequant_head.glsl" + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {block_rocmfpx_fp8 data_a[];}; +layout (binding = 1) writeonly buffer D {D_TYPE data_b[];}; + +void main() { + const uint i = gl_WorkGroupID.x * 4 + gl_LocalInvocationID.x / 64; + + init_iq_shmem(gl_WorkGroupSize); + + const uint tid = gl_LocalInvocationID.x % 64; + const uint il = tid/32; + const uint ir = tid%32; + const uint ib = 32*i + ir; + if (ib >= p.nel / 32) { + return; + } + + const uint b_idx = 1024*i + 32*ir + 16*il; + const uint q_idx = 16*il; + const float d = ue4m3_to_fp32(data_a[ib].e); + + [[unroll]] for (uint l = 0; l < 16; l += 2) { + data_b[b_idx + l ] = D_TYPE(d * float(data_a[ib].qs[q_idx + l ])); + data_b[b_idx + l + 1] = D_TYPE(d * float(data_a[ib].qs[q_idx + l + 1])); + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_comb.comp b/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_comb.comp index f4ac0378a620..0449715bf458 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_comb.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_comb.comp @@ -1,90 +1,101 @@ #version 450 -#extension GL_EXT_control_flow_attributes : require -#extension GL_KHR_shader_subgroup_basic : require -#extension GL_KHR_shader_subgroup_shuffle : require - -// 16 lanes per token, indexed idst + hc*isrc: idst in bits 0..1, isrc in bits 2..3, -// so subgroupShuffleXor by 1|2 reduces a row and by 4|8 a column. - -layout(constant_id = 0) const uint SUBGROUP_SIZE = 32; - -layout(local_size_x_id = 0, local_size_y = 4, local_size_z = 1) in; - -layout(push_constant) uniform parameter -{ - uint n_tokens; - - uint nbm0; uint nbm1; // mixes - uint nbs0; // scale - uint nbb0; // base - uint nbd0; uint nbd1; uint nbd2; // dst - - uint m_offset; - uint s_offset; - uint b_offset; - uint d_offset; - +// DeepSeek V4 fused hyper-connection "comb": build the 4x4 stream-mixing matrix per token — +// per-source softmax over scaled+biased logits, then eps-stabilized alternating column/row +// (sinkhorn) normalization. Port of ggml-cuda/dsv4-hc.cu (hc_comb): one THREAD per token, +// the whole 4x4 lives in registers. At decode this is a single active thread by design — +// the fusion's value is collapsing the ~dozen decomposed graph ops (and their intermediate +// tensors) into one dispatch, not throughput. + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout(binding = 0) readonly buffer MBuf { float data_mixes[]; }; +layout(binding = 1) readonly buffer SBuf { float data_scale[]; }; +layout(binding = 2) readonly buffer BBuf { float data_base[]; }; +layout(binding = 3) writeonly buffer DBuf { float data_d[]; }; + +layout(push_constant) uniform Parameters { + uint n_tokens; + uint sm0, sm1; + uint ss0; + uint sb0; + uint sd0, sd1, sd2; float eps; - uint n_iter; -}; - -layout(binding = 0, std430) readonly buffer M { float data_m[]; }; -layout(binding = 1, std430) readonly buffer S { float data_s[]; }; -layout(binding = 2, std430) readonly buffer B { float data_b[]; }; -layout(binding = 3, std430) writeonly buffer D { float data_d[]; }; - -const uint hc = 4; -const uint comb_offset = 2 * hc; + int n_iter; +} p; + +const uint HC = 4; +const uint COMB_OFFSET = 2 * HC; // comb logits start after the pre/post blocks of the mix vector + +void norm_cols(inout float comb[HC * HC]) { + for (uint idst = 0; idst < HC; ++idst) { + float sum = p.eps; + for (uint isrc = 0; isrc < HC; ++isrc) { + sum += comb[idst + HC * isrc]; + } + const float inv_sum = 1.0 / sum; + for (uint isrc = 0; isrc < HC; ++isrc) { + comb[idst + HC * isrc] *= inv_sum; + } + } +} -const uint TOKENS_PER_SUBGROUP = SUBGROUP_SIZE / 16; +void norm_rows(inout float comb[HC * HC]) { + for (uint isrc = 0; isrc < HC; ++isrc) { + float sum = p.eps; + for (uint idst = 0; idst < HC; ++idst) { + sum += comb[idst + HC * isrc]; + } + const float inv_sum = 1.0 / sum; + for (uint idst = 0; idst < HC; ++idst) { + comb[idst + HC * isrc] *= inv_sum; + } + } +} void main() { - const uint lane = gl_SubgroupInvocationID; - const uint blk = lane >> 4; // which 16-lane block, i.e. which token - const uint idx = lane & 15; // idst + hc*isrc - - const uint sg = gl_WorkGroupID.x * gl_WorkGroupSize.y + gl_SubgroupID; - const uint it = sg * TOKENS_PER_SUBGROUP + blk; - - // no early return, the shuffles need every lane; out-of-range blocks compute a discarded value - const bool in_range = it < n_tokens; - - const float scale_comb = data_s[s_offset + 2 * nbs0]; - - float v = 0.0f; - if (in_range) { - v = data_m[m_offset + (comb_offset + idx) * nbm0 + it * nbm1] * scale_comb - + data_b[b_offset + (comb_offset + idx) * nbb0]; + const uint it = gl_GlobalInvocationID.x; + if (it >= p.n_tokens) { + return; } - // Softmax across destinations: the four lanes sharing an isrc. - float vmax = max(v, subgroupShuffleXor(v, 1)); - vmax = max(vmax, subgroupShuffleXor(vmax, 2)); - v = exp(v - vmax); - - float sum = v + subgroupShuffleXor(v, 1); - sum += subgroupShuffleXor(sum, 2); - v = v / sum + eps; - - // Normalize columns: equal destination indices are four lanes apart. - sum = v + subgroupShuffleXor(v, 4); - sum += subgroupShuffleXor(sum, 8); - v /= sum + eps; - - for (uint i = 1; i < n_iter; ++i) { - sum = v + subgroupShuffleXor(v, 1); - sum += subgroupShuffleXor(sum, 2); - v /= sum + eps; + const float scale_comb = data_scale[2 * p.ss0]; + float comb[HC * HC]; + + for (uint isrc = 0; isrc < HC; ++isrc) { + float vmax = uintBitsToFloat(0xff800000); // -inf + for (uint idst = 0; idst < HC; ++idst) { + const uint idx = idst + HC * isrc; + const float v = data_mixes[(COMB_OFFSET + idx) * p.sm0 + it * p.sm1] * scale_comb + + data_base[(COMB_OFFSET + idx) * p.sb0]; + comb[idx] = v; + vmax = max(vmax, v); + } + + float sum = 0.0; + for (uint idst = 0; idst < HC; ++idst) { + const uint idx = idst + HC * isrc; + const float v = exp(comb[idx] - vmax); + comb[idx] = v; + sum += v; + } + + const float inv_sum = 1.0 / sum; + for (uint idst = 0; idst < HC; ++idst) { + const uint idx = idst + HC * isrc; + comb[idx] = comb[idx] * inv_sum + p.eps; + } + } - sum = v + subgroupShuffleXor(v, 4); - sum += subgroupShuffleXor(sum, 8); - v /= sum + eps; + norm_cols(comb); + for (int i = 1; i < p.n_iter; ++i) { + norm_rows(comb); + norm_cols(comb); } - if (in_range) { - const uint idst = idx & 3; - const uint isrc = idx >> 2; - data_d[d_offset + idst * nbd0 + isrc * nbd1 + it * nbd2] = v; + for (uint isrc = 0; isrc < HC; ++isrc) { + for (uint idst = 0; idst < HC; ++idst) { + data_d[idst * p.sd0 + isrc * p.sd1 + it * p.sd2] = comb[idst + HC * isrc]; + } } } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_post.comp b/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_post.comp index bab6f8767848..212109dfd434 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_post.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_post.comp @@ -1,83 +1,44 @@ #version 450 -#extension GL_EXT_control_flow_attributes : require +// DeepSeek V4 fused hyper-connection "post": redistribute the layer output back into the HC +// streams with the sinkhorn-mixed residual, +// dst[i0, idst, it] = x[i0, it] * post[idst, it] + sum_isrc residual[i0, isrc, it] * comb[idst, isrc, it]. +// Port of ggml-cuda/dsv4-hc.cu (hc_post). Flat elementwise over n_embd * hc * n_tokens. -// Fan one stream back out to hc streams and add the combination-weighted -// residuals: -// -// dst[i0, idst, it] = x[i0, it]*post[idst, it] -// + sum_isrc residual[i0, isrc, it]*comb[idst, isrc, it] +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; -layout(constant_id = 0) const uint BLOCK_SIZE = 256; +layout(binding = 0) readonly buffer XBuf { float data_x[]; }; +layout(binding = 1) readonly buffer RBuf { float data_r[]; }; +layout(binding = 2) readonly buffer PBuf { float data_p[]; }; +layout(binding = 3) readonly buffer CBuf { float data_c[]; }; +layout(binding = 4) writeonly buffer DBuf { float data_d[]; }; -layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; - -layout(push_constant) uniform parameter -{ +layout(push_constant) uniform Parameters { uint n_embd; - uint n_tokens; - - uint nbx0; uint nbx1; // x - uint nbr0; uint nbr1; uint nbr2; // residual - uint nbp0; uint nbp1; // post - uint nbc0; uint nbc1; uint nbc2; // comb - uint nbd0; uint nbd1; uint nbd2; // dst - - uint x_offset; - uint r_offset; - uint p_offset; - uint c_offset; - uint d_offset; -}; - -layout(binding = 0, std430) readonly buffer X { float data_x[]; }; -layout(binding = 1, std430) readonly buffer R { float data_r[]; }; -layout(binding = 2, std430) readonly buffer P { float data_p[]; }; -layout(binding = 3, std430) readonly buffer C { float data_c[]; }; -layout(binding = 4, std430) writeonly buffer D { float data_d[]; }; - -const uint hc = 4; - -shared float post_s[hc]; -shared float comb_s[hc * hc]; + uint hc; + uint nr; // n_embd * hc * n_tokens + uint sx0, sx1; + uint sr0, sr1, sr2; + uint sp0, sp1; + uint sc0, sc1, sc2; + uint sd0, sd1, sd2; +} p; void main() { - const uint tid = gl_LocalInvocationID.x; - const uint it = gl_WorkGroupID.y; - - if (tid < hc) { - post_s[tid] = data_p[p_offset + tid * nbp0 + it * nbp1]; - } - if (tid < hc * hc) { - const uint idst = tid & 3; - const uint isrc = tid >> 2; - comb_s[tid] = data_c[c_offset + idst * nbc0 + isrc * nbc1 + it * nbc2]; - } - barrier(); - - // After the barrier, so every invocation reaches it. - const uint i0 = gl_WorkGroupID.x * BLOCK_SIZE + tid; - if (i0 >= n_embd) { + const uint ir = gl_GlobalInvocationID.x; + if (ir >= p.nr) { return; } - const float xv = data_x[x_offset + i0 * nbx0 + it * nbx1]; - - const uint rb = r_offset + i0 * nbr0 + it * nbr2; + const uint i0 = ir % p.n_embd; + const uint idst = (ir / p.n_embd) % p.hc; + const uint it = ir / (p.n_embd * p.hc); - float r[hc]; - [[unroll]] - for (uint isrc = 0; isrc < hc; ++isrc) { - r[isrc] = data_r[rb + isrc * nbr1]; + float sum = data_x[i0 * p.sx0 + it * p.sx1] * data_p[idst * p.sp0 + it * p.sp1]; + for (uint isrc = 0; isrc < p.hc; ++isrc) { + sum += data_r[i0 * p.sr0 + isrc * p.sr1 + it * p.sr2] + * data_c[idst * p.sc0 + isrc * p.sc1 + it * p.sc2]; } - [[unroll]] - for (uint idst = 0; idst < hc; ++idst) { - float result = xv * post_s[idst]; - [[unroll]] - for (uint isrc = 0; isrc < hc; ++isrc) { - result = fma(r[isrc], comb_s[idst + hc * isrc], result); - } - data_d[d_offset + i0 * nbd0 + idst * nbd1 + it * nbd2] = result; - } + data_d[i0 * p.sd0 + idst * p.sd1 + it * p.sd2] = sum; } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_pre.comp b/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_pre.comp index 51deabbac6ed..b6cc09fe2cfa 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_pre.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/dsv4_hc_pre.comp @@ -1,59 +1,37 @@ #version 450 -#extension GL_EXT_control_flow_attributes : require +// DeepSeek V4 fused hyper-connection "pre": mix the HC input streams down to one embedding, +// dst[i0, it] = sum_ih x[i0, ih, it] * w[ih, it]. Port of ggml-cuda/dsv4-hc.cu (hc_pre). +// Flat elementwise kernel over n_embd * n_tokens; strides are in f32 elements. -// Collapse the hc residual streams of a token into one, weighted per stream: -// -// dst[i0, it] = sum_ih x[i0, ih, it] * weights[ih, it] +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; -layout(constant_id = 0) const uint BLOCK_SIZE = 256; +layout(binding = 0) readonly buffer XBuf { float data_x[]; }; +layout(binding = 1) readonly buffer WBuf { float data_w[]; }; +layout(binding = 2) writeonly buffer DBuf { float data_d[]; }; -layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; - -layout(push_constant) uniform parameter -{ +layout(push_constant) uniform Parameters { uint n_embd; - uint n_tokens; - - uint nbx0; uint nbx1; uint nbx2; // x - uint nbw0; uint nbw1; // weights - uint nbd0; uint nbd1; // dst - - uint x_offset; - uint w_offset; - uint d_offset; -}; - -layout(binding = 0, std430) readonly buffer X { float data_x[]; }; -layout(binding = 1, std430) readonly buffer W { float data_w[]; }; -layout(binding = 2, std430) writeonly buffer D { float data_d[]; }; - -const uint hc = 4; - -shared float w[hc]; + uint hc; + uint nr; // n_embd * n_tokens + uint sx0, sx1, sx2; + uint sw0, sw1; + uint sd0, sd1; +} p; void main() { - const uint tid = gl_LocalInvocationID.x; - const uint it = gl_WorkGroupID.y; - - if (tid < hc) { - w[tid] = data_w[w_offset + tid * nbw0 + it * nbw1]; - } - barrier(); - - // After the barrier, so every invocation reaches it. - const uint i0 = gl_WorkGroupID.x * BLOCK_SIZE + tid; - if (i0 >= n_embd) { + const uint ir = gl_GlobalInvocationID.x; + if (ir >= p.nr) { return; } - const uint xb = x_offset + i0 * nbx0 + it * nbx2; + const uint i0 = ir % p.n_embd; + const uint it = ir / p.n_embd; - float result = 0.0f; - [[unroll]] - for (uint ih = 0; ih < hc; ++ih) { - result = fma(data_x[xb + ih * nbx1], w[ih], result); + float sum = data_x[i0 * p.sx0 + it * p.sx2] * data_w[it * p.sw1]; + for (uint ih = 1; ih < p.hc; ++ih) { + sum += data_x[i0 * p.sx0 + ih * p.sx1 + it * p.sx2] * data_w[ih * p.sw0 + it * p.sw1]; } - data_d[d_offset + i0 * nbd0 + it * nbd1] = result; + data_d[i0 * p.sd0 + it * p.sd1] = sum; } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/fa_types.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/fa_types.glsl index 1e732a9a29cc..6f414ded1230 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/fa_types.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/fa_types.glsl @@ -1,22 +1,32 @@ #if !defined(GGML_FA_TYPES_COMP) #define GGML_FA_TYPES_COMP -#include "ggml_type_ids.glsl" +// FaTypeK / FaTypeV spec constant values. These mirror enum ggml_type so the +// host can pass the type directly. Keep in sync with ggml.h. +#define FA_TYPE_F32 0u +#define FA_TYPE_F16 1u +#define FA_TYPE_Q4_0 2u +#define FA_TYPE_Q4_1 3u +#define FA_TYPE_Q5_0 6u +#define FA_TYPE_Q5_1 7u +#define FA_TYPE_Q8_0 8u +#define FA_TYPE_IQ4_NL 20u +#define FA_TYPE_BF16 30u // Number of matrix elements per buffer block, derived from the K/V type spec // constant. F32 is treated as a vec4 "block" of 4 floats. F16 uses block size 1 // and bypasses the dequant path entirely. Quants follow their ggml block sizes. uint fa_block_elems(uint ty) { switch (ty) { - case GGML_TYPE_F32: return 4u; - case GGML_TYPE_F16: return 1u; - case GGML_TYPE_Q4_0: return uint(QUANT_K_Q4_0); - case GGML_TYPE_Q4_1: return uint(QUANT_K_Q4_1); - case GGML_TYPE_Q5_0: return uint(QUANT_K_Q5_0); - case GGML_TYPE_Q5_1: return uint(QUANT_K_Q5_1); - case GGML_TYPE_Q8_0: return uint(QUANT_K_Q8_0); - case GGML_TYPE_IQ4_NL: return uint(QUANT_K_IQ4_NL); - case GGML_TYPE_BF16: return 1u; + case FA_TYPE_F32: return 4u; + case FA_TYPE_F16: return 1u; + case FA_TYPE_Q4_0: return uint(QUANT_K_Q4_0); + case FA_TYPE_Q4_1: return uint(QUANT_K_Q4_1); + case FA_TYPE_Q5_0: return uint(QUANT_K_Q5_0); + case FA_TYPE_Q5_1: return uint(QUANT_K_Q5_1); + case FA_TYPE_Q8_0: return uint(QUANT_K_Q8_0); + case FA_TYPE_IQ4_NL: return uint(QUANT_K_IQ4_NL); + case FA_TYPE_BF16: return 1u; default: return 1u; } } @@ -26,18 +36,18 @@ uint fa_block_elems(uint ty) { // of int32s per 32-element block on the MMQ K path: ints_per_block == 8 / R. uint fa_quant_r_mmq(uint ty) { switch (ty) { - case GGML_TYPE_Q4_0: return uint(QUANT_R_Q4_0); - case GGML_TYPE_Q4_1: return uint(QUANT_R_Q4_1); - case GGML_TYPE_Q5_0: return uint(QUANT_R_Q5_0); - case GGML_TYPE_Q5_1: return uint(QUANT_R_Q5_1); - case GGML_TYPE_Q8_0: return uint(QUANT_R_Q8_0); + case FA_TYPE_Q4_0: return uint(QUANT_R_Q4_0); + case FA_TYPE_Q4_1: return uint(QUANT_R_Q4_1); + case FA_TYPE_Q5_0: return uint(QUANT_R_Q5_0); + case FA_TYPE_Q5_1: return uint(QUANT_R_Q5_1); + case FA_TYPE_Q8_0: return uint(QUANT_R_Q8_0); default: return 1u; } } bool fa_type_needs_shmem(uint ty) { switch (ty) { - case GGML_TYPE_IQ4_NL: return true; + case FA_TYPE_IQ4_NL: return true; default: return false; } } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp index 107d44aaa801..5b19ff61093f 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn.comp @@ -134,7 +134,7 @@ void main() { // Q8_0 K only needs (qd, _); the asymmetric Q4_*/Q5_* family also stores // the row-sum scaled by qd, used in k_dot_correction. - if (FaTypeK == GGML_TYPE_Q8_0) { + if (FaTypeK == FA_TYPE_Q8_0) { if (buf_iqs == 0) { Qf[buf_ib].ds = FLOAT_TYPEV2(qd, 0.0f); } @@ -186,9 +186,9 @@ void main() { // FaBlockBytesK/V == 2 for f16, 16 for f32, ggml block byte size for quants. uint32_t k_offset = (ik2*p.nb12 + ik3*p.nb13) / FaBlockBytesK; uint32_t v_offset = (iv2*p.nb22 + iv3*p.nb23) / FaBlockBytesV; - uint32_t m_offset = gqa_iq1*KV; + uint32_t m_offset = gqa_iq1*m_row_len; if (p.nem2 != 1 || p.nem3 != 1) { - m_offset += ((iq3 % p.nem3) * p.nem2 + (iq2 % p.nem2)) * p.nem1 * KV; + m_offset += ((iq3 % p.nem3) * p.nem2 + (iq2 % p.nem2)) * p.nem1 * m_row_len; mo_offset += ((iq3 % p.nem3) * p.nem2 + (iq2 % p.nem2)) * CEIL_DIV(p.nem1, Br) * mo_stride; } @@ -210,7 +210,7 @@ void main() { } // Only load if the block is not all zeros if (mask_opt_bits != MASK_OPT_ALL_ZERO) { - bool nem1_bounds_check = !(p.gqa_ratio > 1) && (p.nem1 % Br) != 0; + bool nem1_bounds_check = !(gqa_ratio > 1) && (p.nem1 % Br) != 0; float max_mask = NEG_FLT_MAX_OVER_2; barrier(); @@ -218,14 +218,12 @@ void main() { uint32_t c = (idx + tid) % Bc; uint32_t r = (idx + tid) / Bc; if (idx + tid < Bc * Br) { - uint32_t kcol; - bool kv_active = fa_kv_index(j * Bc + c, kcol); - if (kv_active && (!nem1_bounds_check || i * Br + r < p.nem1)) { - FLOAT_TYPE m = FLOAT_TYPE(data_m[m_offset + (i * Br + r) * m_stride + kcol]); + if ((!KV_bounds_check || j * Bc + c < KV) && (!nem1_bounds_check || i * Br + r < p.nem1)) { + FLOAT_TYPE m = FLOAT_TYPE(data_m[m_offset + (i * Br + r) * m_stride + (j * Bc + c)]); masksh[c * masksh_stride + r] = m; max_mask = max(max_mask, float(m)); } else { - masksh[c * masksh_stride + r] = USE_SPARSE ? FLOAT_TYPE(NEG_FLT_MAX_OVER_2) : FLOAT_TYPE(0); + masksh[c * masksh_stride + r] = FLOAT_TYPE(0); } } } @@ -260,15 +258,14 @@ void main() { uint32_t c = (idx + tid) / (HSK / 4); if (idx + gl_WorkGroupSize.x <= Bc * HSK / 4 || c < Bc) { FLOAT_TYPEV4 K_Tf = FLOAT_TYPEV4(0); - uint32_t kcol; - if (fa_kv_index(j * Bc + c, kcol)) { + if (!KV_bounds_check || j * Bc + c < KV) { if (USE_DECODE_K) { - uint coord = kcol * k_stride * BLOCK_SIZE_K + 4 * d; + uint coord = (j * Bc + c) * k_stride * BLOCK_SIZE_K + 4 * d; uint ib = coord / BLOCK_SIZE_K; uint iqs = (coord % BLOCK_SIZE_K); K_Tf = dequantize4(ib, iqs, k_offset, BINDING_IDX_K); } else { - K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + kcol * k_stride / 4 + d]); + K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + (j * Bc + c) * k_stride / 4 + d]); } } @@ -308,9 +305,7 @@ void main() { } [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) { - uint32_t kcol; - bool kv_active = fa_kv_index(j * Bc + c * cols_per_iter + col_tid, kcol); - if (!kv_active) { + if (KV_bounds_check && j * Bc + c * cols_per_iter + col_tid >= KV) { continue; } @@ -318,12 +313,12 @@ void main() { if (SHMEM_STAGING != 0) { K_Tf = kvsh[(c * cols_per_iter + col_tid) * kvsh_stride + (d * D_split + d_tid)]; } else if (USE_DECODE_K) { - uint coord = kcol * k_stride * BLOCK_SIZE_K + 4 * (d * D_split + d_tid); + uint coord = (j * Bc + c * cols_per_iter + col_tid) * k_stride * BLOCK_SIZE_K + 4 * (d * D_split + d_tid); uint ib = coord / BLOCK_SIZE_K; uint iqs = (coord % BLOCK_SIZE_K); K_Tf = dequantize4(ib, iqs, k_offset, BINDING_IDX_K); } else { - K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + kcol * k_stride / 4 + d * D_split + d_tid]); + K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + (j * Bc + c * cols_per_iter + col_tid) * k_stride / 4 + d * D_split + d_tid]); } [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { Sf[r][c] = dot_product(Q_cache[r], K_Tf, Sf[r][c]); @@ -332,9 +327,7 @@ void main() { } } else { [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) { - uint32_t kcol; - bool kv_active = fa_kv_index(j * Bc + c * cols_per_iter + col_tid, kcol); - if (!kv_active) { + if (KV_bounds_check && j * Bc + c * cols_per_iter + col_tid >= KV) { continue; } @@ -343,12 +336,12 @@ void main() { if (SHMEM_STAGING != 0) { K_Tf = kvsh[(c * cols_per_iter + col_tid) * kvsh_stride + (d * D_split + d_tid)]; } else if (USE_DECODE_K) { - uint coord = kcol * k_stride * BLOCK_SIZE_K + 4 * (d * D_split + d_tid); + uint coord = (j * Bc + c * cols_per_iter + col_tid) * k_stride * BLOCK_SIZE_K + 4 * (d * D_split + d_tid); uint ib = coord / BLOCK_SIZE_K; uint iqs = (coord % BLOCK_SIZE_K); K_Tf = dequantize4(ib, iqs, k_offset, BINDING_IDX_K); } else { - K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + kcol * k_stride / 4 + d * D_split + d_tid]); + K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + (j * Bc + c * cols_per_iter + col_tid) * k_stride / 4 + d * D_split + d_tid]); } [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { Sf[r][c] = dot_product(Qf[tile_row(r) * qf_stride + d * D_split + d_tid], K_Tf, Sf[r][c]); @@ -374,7 +367,7 @@ void main() { // Q4_*/Q5_* take the block-8 fast path when one step covers a full // block; Q8_0 always goes through the per-int get_k_qs* helpers // (its qs is byte-packed, not nibble-packed). - const bool block8_fast = (d_per_step == 8) && (FaTypeK != GGML_TYPE_Q8_0); + const bool block8_fast = (d_per_step == 8) && (FaTypeK != FA_TYPE_Q8_0); if (SHMEM_STAGING != 0) { const uint k_block_idx = (d_tid * (HSK_per_thread / 4) + d_block) / 8; @@ -382,7 +375,7 @@ void main() { k_dm = ACC_TYPEV2(kblocksh[buf_ib].dm); if (block8_fast) { - const bool has_qh = (FaTypeK == GGML_TYPE_Q5_0) || (FaTypeK == GGML_TYPE_Q5_1); + const bool has_qh = (FaTypeK == FA_TYPE_Q5_0) || (FaTypeK == FA_TYPE_Q5_1); [[unroll]] for (uint32_t d = 0; d < 4; d++) { uint vui = kblocksh[buf_ib].qs[d]; k_quants[d ] = int32_t( vui & 0x0F0F0F0F); @@ -427,7 +420,11 @@ void main() { acc += dotPacked4x8EXT(Qf[qib].qs[qiqs + d], k_quants[d]); } - Sf[r][c] += ACC_TYPE(acc) * ACC_TYPE(Qf[qib].ds.x) * k_dm.x; + // scale in fp32 before narrowing: acc is an int32 sum of dotPacked4x8EXT + // results, bounded by d_per_step*4*127*127 with q8_0 on both sides, which + // overflows f16 when ACC_TYPE is f16 (GGML_PREC_DEFAULT). Identical + // arithmetic when ACC_TYPE is float. + Sf[r][c] += ACC_TYPE(float(acc) * float(Qf[qib].ds.x) * float(k_dm.x)); if ((d_tid * (HSK_per_thread / 4) + d_block) % 8 == 0) { Sf[r][c] += k_dot_correction(qib, k_dm); } @@ -496,15 +493,14 @@ void main() { uint32_t c = (idx + tid) / (HSV / 4); if (idx + gl_WorkGroupSize.x <= Bc * HSV / 4 || c < Bc) { FLOAT_TYPEV4 V_Tf = FLOAT_TYPEV4(0); - uint32_t vcol; - if (fa_kv_index(j * Bc + c, vcol)) { + if (!KV_bounds_check || j * Bc + c < KV) { if (USE_DECODE_V) { - uint coord = vcol * v_stride * BLOCK_SIZE_V + 4 * d; + uint coord = (j * Bc + c) * v_stride * BLOCK_SIZE_V + 4 * d; uint ib = coord / BLOCK_SIZE_V; uint iqs = (coord % BLOCK_SIZE_V); V_Tf = dequantize4(ib, iqs, v_offset, BINDING_IDX_V); } else { - V_Tf = FLOAT_TYPEV4(data_vv4[v_offset / 4 + vcol * v_stride / 4 + d]); + V_Tf = FLOAT_TYPEV4(data_vv4[v_offset / 4 + (j * Bc + c) * v_stride / 4 + d]); } } @@ -515,9 +511,7 @@ void main() { } [[unroll]] for (uint32_t c = 0; c < cols_per_thread; ++c) { - uint32_t vcol; - bool kv_active = fa_kv_index(j * Bc + c * cols_per_iter + col_tid, vcol); - if (!kv_active) { + if (KV_bounds_check && j * Bc + c * cols_per_iter + col_tid >= KV) { continue; } @@ -532,12 +526,12 @@ void main() { if (SHMEM_STAGING != 0) { Vf = kvsh[(c * cols_per_iter + col_tid) * kvsh_stride + (d * D_split + d_tid)]; } else if (USE_DECODE_V) { - uint coord = vcol * v_stride * BLOCK_SIZE_V + 4 * (d * D_split + d_tid); + uint coord = (j * Bc + c * cols_per_iter + col_tid) * v_stride * BLOCK_SIZE_V + 4 * (d * D_split + d_tid); uint ib = coord / BLOCK_SIZE_V; uint iqs = (coord % BLOCK_SIZE_V); Vf = dequantize4(ib, iqs, v_offset, BINDING_IDX_V); } else { - Vf = FLOAT_TYPEV4(data_vv4[v_offset / 4 + vcol * v_stride / 4 + d * D_split + d_tid]); + Vf = FLOAT_TYPEV4(data_vv4[v_offset / 4 + (j * Bc + c * cols_per_iter + col_tid) * v_stride / 4 + d * D_split + d_tid]); } [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { Of[r][d] += FLOAT_TYPEV4(Pf[r] * Vf); @@ -666,10 +660,10 @@ void main() { // If there is split_k, then the split_k resolve shader does the final // division by L. Store the intermediate O value and per-row m and L values. - if (p.k_num > 1) { - if (p.gqa_ratio > 1) { + if (partial_output || split_k_num > 1) { + if (gqa_ratio > 1) { // note: O and Q have swapped coord 1,2. - uint32_t o_offset = HSV * p.ne1 * (split_k_index + p.k_num * (gqa_iq1 + p.ne2 * iq3)) / 4; + uint32_t o_offset = HSV * p.ne1 * (split_k_index + output_k_num * (gqa_iq1 + p.ne2 * iq3)) / 4; [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { const uint row = tile_row(r); @@ -680,7 +674,7 @@ void main() { } } - o_offset = HSV * p.ne1 * p.k_num * p.ne2 * p.ne3 + p.ne1 * 2 * (split_k_index + p.k_num * (gqa_iq1 + p.ne2 * iq3)); + o_offset = HSV * p.ne1 * output_k_num * p.ne2 * p.ne3 + p.ne1 * 2 * (split_k_index + output_k_num * (gqa_iq1 + p.ne2 * iq3)); [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { const uint row = tile_row(r); if (row < N) { @@ -694,7 +688,7 @@ void main() { const uint global_row = i * Br + row; if (global_row < N) { - uint32_t o_offset = HSV * p.ne1 * (split_k_index + p.k_num * (global_row + p.ne2 * iq3)) / 4; + uint32_t o_offset = HSV * p.ne1 * (split_k_index + output_k_num * (global_row + p.ne2 * iq3)) / 4; [[unroll]] for (uint32_t d = 0; d < HSV_per_thread / 4; ++d) { data_ov4[o_offset + iq2 * HSV/4 + d * D_split + d_tid] = D_TYPEV4(Of[r][d]); @@ -702,7 +696,7 @@ void main() { } if (global_row < N && d_tid == 0 && col_tid == 0) { - uint32_t lm_offset = HSV * p.ne1 * p.k_num * p.ne2 * p.ne3 + p.ne1 * 2 * (split_k_index + p.k_num * (global_row + p.ne2 * iq3)); + uint32_t lm_offset = HSV * p.ne1 * output_k_num * p.ne2 * p.ne3 + p.ne1 * 2 * (split_k_index + output_k_num * (global_row + p.ne2 * iq3)); data_o[lm_offset + iq2] = D_TYPE(Lf[r]); data_o[lm_offset + p.ne1 + iq2] = D_TYPE(Mf[r]); } @@ -748,7 +742,7 @@ void main() { uint32_t o_offset = (gqa_iq1*p.ne1*HSV + iq3*p.ne2*p.ne1*HSV) / 4; - if (p.gqa_ratio > 1) { + if (gqa_ratio > 1) { [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { const uint row = tile_row(r); if (row < N) { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl index 2e0e23bc11c1..b562c5d78749 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_base.glsl @@ -24,8 +24,11 @@ const bool USE_MASK_OPT = (Flags & 1) != 0; const bool MASK_ENABLE = (Flags & 2) != 0; const bool LOGIT_SOFTCAP = (Flags & 4) != 0; const bool OLD_AMD_WINDOWS = (Flags & 8) != 0; -// Sparse: gather binding-7 indices instead of scanning [0,KV); p.split_kv = n_kv_max. -const bool USE_SPARSE = (Flags & 16) != 0; +// KV comes from a buffer instead of the push constant. Used by paths that compact K/V on the +// GPU, where the row count is only known after a dedup pass and so cannot be pushed. The +// workgroup counts derive from neq1/neq2/neq3 and never from KV, so no indirect dispatch is +// needed: only this loop bound changes. Folds away for every other pipeline. +const bool DYNAMIC_KV = (Flags & 16) != 0; // Round up head sizes to a multiple of 16, for coopmat1/coopmat2 paths const uint32_t HSK_pad = (HSK + 15) & ~15; @@ -83,8 +86,7 @@ layout (binding = 5) writeonly buffer O {D_TYPE data_o[];}; layout (binding = 5) writeonly buffer OV4 {D_TYPEV4 data_ov4[];}; layout (binding = 6) readonly buffer MO {uint32_t data_mask_opt[];}; - -layout (binding = 7) readonly buffer SP {int32_t data_sparse[];}; +layout (binding = 7) readonly buffer KVB {uint32_t data_kv_dyn[];}; #define MASK_OPT_ALL_NEG_INF 1 #define MASK_OPT_ALL_ZERO 2 @@ -109,8 +111,8 @@ layout (binding = 7) readonly buffer SP {int32_t data_sparse[];}; #define BLOCK_SIZE_V fa_block_elems(FaTypeV) // F16 reads f16 elements directly from the binding; everything else routes // through dequantize4 / the MMQ helpers to unpack from the packed block layout. -#define USE_DECODE_K (FaTypeK != GGML_TYPE_F16) -#define USE_DECODE_V (FaTypeV != GGML_TYPE_F16) +#define USE_DECODE_K (FaTypeK != FA_TYPE_F16) +#define USE_DECODE_V (FaTypeV != FA_TYPE_F16) #define CEIL_DIV(a, b) (((a) + (b) - 1) / (b)) @@ -128,7 +130,7 @@ ACC_TYPE perElemOpStoreCol0(const in uint32_t r, const in uint32_t c, const in A // Load the slope matrix, indexed by Q's dimension 2. ACC_TYPE perElemOpComputeSlope(const in uint32_t r, const in uint32_t c, const in ACC_TYPE elem, const in uint32_t iq2) { - const uint32_t h = iq2 + (r % p.gqa_ratio); + const uint32_t h = iq2 + (r % (p.gqa_ratio & 0xffff)); uint32_t n_head_log2 = p.mask_n_head_log2 & N_LOG2_MASK; @@ -141,32 +143,40 @@ ACC_TYPE perElemOpComputeSlope(const in uint32_t r, const in uint32_t c, const i // Load the sink value, indexed by Q's dimension 2. ACC_TYPE perElemOpGetSink(const in uint32_t r, const in uint32_t c, const in ACC_TYPE elem, const in uint32_t iq2) { - const uint32_t h = iq2 + (r % p.gqa_ratio); + const uint32_t h = iq2 + (r % (p.gqa_ratio & 0xffff)); return ACC_TYPE(data_s[h]); } uint32_t i, N, KV, split_k_index, Tr, start_j, end_j, gqa_iq1, iq2, iq3, rk2, rk3, rv2, rv3, ik2, ik3, iv2, iv3, - q_stride, k_stride, v_stride, m_stride, sparse_base; + q_stride, k_stride, v_stride, m_stride, m_row_len, gqa_ratio, split_k_num, output_k_num; +bool partial_output; void init_indices() { N = p.N; - KV = p.KV; + KV = DYNAMIC_KV ? data_kv_dyn[0] : p.KV; + gqa_ratio = p.gqa_ratio & 0xffff; + split_k_num = p.k_num & 0xffff; + output_k_num = p.k_num >> 16; + partial_output = output_k_num != 0; + if (!partial_output) { + output_k_num = split_k_num; + } - if (p.k_num > 1) { - if (p.gqa_ratio > 1) { + if (split_k_num > 1) { + if (gqa_ratio > 1) { i = 0; // batch and split_k share gl_WorkGroupID.x - gqa_iq1 = gl_WorkGroupID.x / p.k_num; - split_k_index = gl_WorkGroupID.x % p.k_num; + gqa_iq1 = gl_WorkGroupID.x / split_k_num; + split_k_index = gl_WorkGroupID.x % split_k_num; } else { gqa_iq1 = 0; - split_k_index = gl_WorkGroupID.x % p.k_num; - i = gl_WorkGroupID.x / p.k_num; + split_k_index = gl_WorkGroupID.x % split_k_num; + i = gl_WorkGroupID.x / split_k_num; } - } else if (p.gqa_ratio > 1) { + } else if (gqa_ratio > 1) { i = 0; gqa_iq1 = gl_WorkGroupID.x; split_k_index = 0; @@ -183,7 +193,7 @@ void init_indices() // When not using grouped query attention, all rows share the same iq2, equal to gl_WorkGroupID.y. // When using grouped query attention, each workgroup does gqa_ratio consecutive values of iq2. - iq2 = gl_WorkGroupID.y * p.gqa_ratio; + iq2 = gl_WorkGroupID.y * gqa_ratio; iq3 = gl_WorkGroupID.z; // broadcast factors @@ -204,41 +214,24 @@ void init_indices() // nb?1 are already divided by the type size and are in units of elements. // When using grouped query attention, Q is indexed by iq2, so the stride // should be nb02 (which is in bytes). - q_stride = p.gqa_ratio > 1 ? (p.nb02 / 4) : p.nb01; + q_stride = gqa_ratio > 1 ? (p.nb02 / 4) : p.nb01; k_stride = p.nb11; v_stride = p.nb21; - // When using grouped query attention, all rows use the same mask (stride 0). - // "p.gqa_ratio >> 16" is just a roundabout way of writing zero - // that prevents the compiler from folding the "&" through the select - // and breaking the alignment detection. - m_stride = (p.gqa_ratio > 1) ? (p.gqa_ratio >> 16) : KV; - - // Sparse: the tile shares one mask row (gqa heads, or Br==1). split_k - // partitions the n_kv_max blocks. - if (USE_SPARSE) { - uint32_t qrow = (p.gqa_ratio > 1) ? gqa_iq1 : (i * Br); - sparse_base = (((iq3 % p.nem3) * p.nem2 + (iq2 % p.nem2)) * p.nem1 + qrow) * p.split_kv; - - uint32_t total_blocks = CEIL_DIV(p.split_kv, Bc); - uint32_t per_blocks = CEIL_DIV(total_blocks, p.k_num); - start_j = min(split_k_index * per_blocks, total_blocks); - end_j = min((split_k_index + 1) * per_blocks, total_blocks); - } -} - -// Resolve a linear KV slot to a real column; false for inactive (sparse padding/-1, or dense OOB). -bool fa_kv_index(uint lin, out uint kv_col) { - if (USE_SPARSE) { - if (lin >= p.split_kv) { - kv_col = 0; - return false; - } - int idx = data_sparse[sparse_base + lin]; - kv_col = idx >= 0 ? uint(idx) : 0; - return idx >= 0; - } - kv_col = lin; - return !KV_bounds_check || lin < KV; + // Bit 31 of gqa_ratio means "the mask row stride is in split_kv", used by the DeepSeek V4 + // sparse split path where the mask spans the full K range but this dispatch only covers the + // raw prefix, so m_stride != KV. That path always sets split_k_num == 1, which is what keeps + // split_kv free to carry the stride; ggml_vk_flash_attn_top_k asserts the invariant. + // Otherwise: when using grouped query attention all rows share the same mask (stride 0). + // "p.gqa_ratio >> 16" is just a roundabout way of writing zero that prevents the compiler + // from folding the "&" through the select and breaking the alignment detection. + const bool mask_stride_in_split_kv = (p.gqa_ratio & 0x80000000u) != 0; + m_stride = mask_stride_in_split_kv ? p.split_kv : ((gqa_ratio > 1) ? (p.gqa_ratio >> 16) : KV); + // Distinct from m_stride: m_stride is the row-to-row step INSIDE this tile (0 under GQA, + // where every row shares one mask row), while m_row_len is the mask tensor's actual row + // length, used to step between tokens and between streams. They differ under GQA and + // under the sparse split path, where this dispatch only covers the raw prefix (KV) but + // the mask rows span the whole K range. + m_row_len = mask_stride_in_split_kv ? p.split_kv : KV; } // Bias applied to softmax to stay in fp16 range. diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp index aa9dd624bef3..195901a7d90d 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm1.comp @@ -44,8 +44,11 @@ shared float tmpsh[row_split]; const uint32_t qstride = HSK_pad / 4 + 2; shared FLOAT_TYPEV4 Qf[Br * qstride]; -const uint psh_stride = Br / 4 + 2; -shared FLOAT_TYPEV4 Psh[Bc * psh_stride]; +// P is stored query-major with KV contiguous so the GEMM2 UseA load can be RowMajor: +// RADV flips the requested layout for UseA, and only the resulting internal column-major +// case gets an alignment hint, which is what lets the 16 element loads vectorize. +const uint psh_stride = Bc / 4 + 2; +shared FLOAT_TYPEV4 Psh[Br * psh_stride]; // Avoid padding for hsk==256 to make it fit in 48KB shmem. const uint32_t sfshstride = (HSK <= 128) ? (Br / 4 + 2) : Br / 4; @@ -139,9 +142,9 @@ void main() { // FaBlockBytesK/V == 2 for f16 (sizeof f16) and == 16 for f32 (vec4) and == ggml block size for quants. uint32_t k_offset = (ik2*p.nb12 + ik3*p.nb13) / FaBlockBytesK; uint32_t v_offset = (iv2*p.nb22 + iv3*p.nb23) / FaBlockBytesV; - uint32_t m_offset = gqa_iq1*KV; + uint32_t m_offset = gqa_iq1*m_row_len; if (p.nem2 != 1 || p.nem3 != 1) { - m_offset += ((iq3 % p.nem3) * p.nem2 + (iq2 % p.nem2)) * p.nem1 * KV; + m_offset += ((iq3 % p.nem3) * p.nem2 + (iq2 % p.nem2)) * p.nem1 * m_row_len; mo_offset += ((iq3 % p.nem3) * p.nem2 + (iq2 % p.nem2)) * CEIL_DIV(p.nem1, Br) * mo_stride; } @@ -169,23 +172,16 @@ void main() { } // Only load if the block is not all zeros if (mask_opt_bits != MASK_OPT_ALL_ZERO) { - bool nem1_bounds_check = !(p.gqa_ratio > 1) && (p.nem1 % Br) != 0; + bool nem1_bounds_check = !(gqa_ratio > 1) && (p.nem1 % Br) != 0; float max_mask = NEG_FLT_MAX_OVER_2; [[unroll]] for (uint32_t idx = 0; idx < Bc * Br / 4; idx += gl_WorkGroupSize.x) { uint32_t c = (idx + tid) / (Br / 4); uint32_t r = (idx + tid) % (Br / 4); if (idx + tid < Bc * Br / 4 || idx + gl_WorkGroupSize.x <= Bc * Br / 4) { - uint32_t kcol; - bool kv_active = fa_kv_index(j * Bc + c, kcol); - if (kv_active) { + if ((!KV_bounds_check || j * Bc + c < KV)) { f16vec4 m; - if (USE_SPARSE) { - // sparse is gqa-gated (m_stride == 0): all four rows share the value - FLOAT_TYPE mv = FLOAT_TYPE(data_m[m_offset + kcol]); - m = f16vec4(mv); - max_mask = max(max_mask, float(mv)); - } else if (!nem1_bounds_check || i * Br + r * 4 + 3 < p.nem1) { + if (!nem1_bounds_check || i * Br + r * 4 + 3 < p.nem1) { m = f16vec4(data_m[m_offset + (i * Br + r * 4 ) * m_stride + (j * Bc + c)], data_m[m_offset + (i * Br + r * 4 + 1) * m_stride + (j * Bc + c)], data_m[m_offset + (i * Br + r * 4 + 2) * m_stride + (j * Bc + c)], @@ -213,8 +209,6 @@ void main() { m = f16vec4(0.0); } mask_cache[idx / WorkGroupSize] = m; - } else if (USE_SPARSE) { - mask_cache[idx / WorkGroupSize] = f16vec4(NEG_FLT_MAX_OVER_2); } } } @@ -240,19 +234,17 @@ void main() { uint32_t c = (idx + tid) / (HSK_pad / 4); if (idx + gl_WorkGroupSize.x <= Bc * HSK_pad / 4 || c < Bc) { FLOAT_TYPEV4 K_Tf = FLOAT_TYPEV4(0); - uint32_t kcol; - bool kv_active = fa_kv_index(j * Bc + c, kcol); - if (kv_active && (HSK == HSK_pad || d < HSK / 4)) { + if ((!KV_bounds_check || j * Bc + c < KV) && (HSK == HSK_pad || d < HSK / 4)) { #if !defined(BFLOAT16) if (USE_DECODE_K) { - uint coord = kcol * k_stride * BLOCK_SIZE_K + 4 * d; + uint coord = (j * Bc + c) * k_stride * BLOCK_SIZE_K + 4 * d; uint ib = coord / BLOCK_SIZE_K; uint iqs = (coord % BLOCK_SIZE_K); K_Tf = dequantize4(ib, iqs, k_offset, BINDING_IDX_K); } else #endif { - K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + kcol * k_stride / 4 + d]); + K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + (j * Bc + c) * k_stride / 4 + d]); } } @@ -277,7 +269,7 @@ void main() { if (SHMEM_STAGING == 0) { // For quants we always need to dequant into kvsh; for f16/bf16 we can load // directly from global memory when alignment / bounds allow it. - const bool stage_k = USE_DECODE_K || KV_bounds_check || USE_SPARSE || d * 16 + 16 > HSK; + const bool stage_k = USE_DECODE_K || KV_bounds_check || d * 16 + 16 > HSK; if (stage_k) { barrier(); [[unroll]] for (uint32_t idx = 0; idx < Bc * MatBr / 4; idx += gl_WorkGroupSize.x) { @@ -285,19 +277,17 @@ void main() { uint32_t row = (idx + tid) / (MatBr / 4); if (idx + tid < Bc * MatBr / 4) { FLOAT_TYPEV4 K_Tf = FLOAT_TYPEV4(0); - uint32_t kcol; - bool kv_active = fa_kv_index(j * Bc + row, kcol); - if (kv_active && (HSK == HSK_pad || d * 16 + col_vec * 4 < HSK)) { + if ((!KV_bounds_check || j * Bc + row < KV) && (HSK == HSK_pad || d * 16 + col_vec * 4 < HSK)) { #if !defined(BFLOAT16) if (USE_DECODE_K) { - uint coord = kcol * k_stride * BLOCK_SIZE_K + d * 16 + col_vec * 4; + uint coord = (j * Bc + row) * k_stride * BLOCK_SIZE_K + d * 16 + col_vec * 4; uint ib = coord / BLOCK_SIZE_K; uint iqs = (coord % BLOCK_SIZE_K); K_Tf = dequantize4(ib, iqs, k_offset, BINDING_IDX_K); } else #endif { - K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + kcol * k_stride / 4 + d * 16 / 4 + col_vec]); + K_Tf = FLOAT_TYPEV4(data_kv4[k_offset / 4 + (j * Bc + row) * k_stride / 4 + d * 16 / 4 + col_vec]); } } @@ -395,15 +385,19 @@ void main() { [[unroll]] for (uint32_t r = 0; r < rows_per_thread; r += 4) { const uint row = tile_row(r); + const uint pcol_vec = col / 4; + const uint pcol_comp = col % 4; if (KV_bounds_check && j * Bc + col >= KV) { - Psh[col * psh_stride + row / 4] = FLOAT_TYPEV4(0.0f); + [[unroll]] for (uint32_t vec_idx = 0; vec_idx < 4; ++vec_idx) { + Psh[(row + vec_idx) * psh_stride + pcol_vec][pcol_comp] = FLOAT_TYPE(0.0f); + } } else { const vec4 mfvec = vec4(Mf[r], Mf[r + 1], Mf[r + 2], Mf[r + 3]); const FLOAT_TYPEV4 Pf = FLOAT_TYPEV4(exp(vec4(sfsh[row / 4 + col * sfshstride]) - mfvec)); [[unroll]] for (uint32_t vec_idx = 0; vec_idx < 4; ++vec_idx) { Lf[r + vec_idx] += Pf[vec_idx]; + Psh[(row + vec_idx) * psh_stride + pcol_vec][pcol_comp] = Pf[vec_idx]; } - Psh[col * psh_stride + row / 4] = Pf; } } } @@ -414,19 +408,17 @@ void main() { uint32_t c = (idx + tid) / (HSV_pad / 4); if (idx + gl_WorkGroupSize.x <= Bc * HSV_pad / 4 || c < Bc) { FLOAT_TYPEV4 V_Tf = FLOAT_TYPEV4(0); - uint32_t v_row; - bool kv_active = fa_kv_index(j * Bc + c, v_row); - if (kv_active && (HSV == HSV_pad || d < HSV / 4)) { + if ((!KV_bounds_check || j * Bc + c < KV) && (HSV == HSV_pad || d < HSV / 4)) { #if !defined(BFLOAT16) if (USE_DECODE_V) { - uint coord = v_row * v_stride * BLOCK_SIZE_V + 4 * d; + uint coord = (j * Bc + c) * v_stride * BLOCK_SIZE_V + 4 * d; uint ib = coord / BLOCK_SIZE_V; uint iqs = (coord % BLOCK_SIZE_V); V_Tf = dequantize4(ib, iqs, v_offset, BINDING_IDX_V); } else #endif { - V_Tf = FLOAT_TYPEV4(data_vv4[v_offset / 4 + v_row * v_stride / 4 + d]); + V_Tf = FLOAT_TYPEV4(data_vv4[v_offset / 4 + (j * Bc + c) * v_stride / 4 + d]); } } @@ -438,6 +430,13 @@ void main() { const uint num_hsv_tiles = (HSV + MatBc * row_split - 1) / (MatBc * row_split); // round up + // Psh is not written again until the next KV block, so the P fragments are the same + // for every hsv_tile. Load them once instead of re-reading LDS per tile. + coopmat PMat[Bc / MatBc]; + [[unroll]] for (uint32_t bc_chunk = 0; bc_chunk < Bc / MatBc; ++bc_chunk) { + coopMatLoad(PMat[bc_chunk], Psh, bc_chunk * (MatBc / 4), psh_stride, gl_CooperativeMatrixLayoutRowMajor); + } + // Each subgroup handles HSV/4 columns [[unroll]] for (uint32_t hsv_tile = 0; hsv_tile < num_hsv_tiles; ++hsv_tile) { const uint hsv_offset = (hsv_tile * row_split + gl_SubgroupID) * 16; @@ -456,22 +455,21 @@ void main() { if (SHMEM_STAGING == 0) { // For quants we always preload via kvsh. For f16/bf16 we only preload when // alignment / bounds force it (otherwise we coopMatLoad direct from data_vv4). - const bool stage_v = USE_DECODE_V || KV_bounds_check || USE_SPARSE; + const bool stage_v = USE_DECODE_V || KV_bounds_check; if (stage_v) { [[unroll]] for (uint32_t i = 0; i < v_loads_per_thread; ++i) { const uint idx = i * gl_WorkGroupSize.x + tid; const uint row = idx / v_cols; const uint col = idx % v_cols; - uint32_t v_row; - bool kv_active = fa_kv_index(j * Bc + row, v_row); + const uint v_row = j * Bc + row; const uint v_col = hsv_tile * MatBc * row_split + col * 4; const uint coord = v_row * v_stride * BLOCK_SIZE_V + v_col; const uint ib = coord / BLOCK_SIZE_V; const uint iqs = coord % BLOCK_SIZE_V; - if (USE_SPARSE ? (kv_active && v_col < HSV) : (!KV_bounds_check || (v_row < KV && v_col < HSV))) { + if (!KV_bounds_check || (v_row < KV && v_col < HSV)) { #if !defined(BFLOAT16) if (USE_DECODE_V) { kvsh[row * vsh_stride + col] = dequantize4(ib, iqs, v_offset, BINDING_IDX_V); @@ -492,10 +490,8 @@ void main() { if (hsv_offset < HSV_pad) { [[unroll]] for (uint32_t bc_chunk = 0; bc_chunk < Bc / MatBc; ++bc_chunk) { - coopMatLoad(KMat, Psh, bc_chunk * MatBc * psh_stride, psh_stride, gl_CooperativeMatrixLayoutColumnMajor); - if (SHMEM_STAGING == 0) { - if (!USE_DECODE_V && !KV_bounds_check && !USE_SPARSE) { + if (!USE_DECODE_V && !KV_bounds_check) { // F16/BF16 values can be loaded directly from global memory const uint v_tile_row = j * Bc + bc_chunk * MatBc; const uint v_tile_offset = v_offset / 4 + v_tile_row * v_stride / 4 + hsv_offset / 4; @@ -509,7 +505,7 @@ void main() { coopMatLoad(QMat, kvsh, v_tile_offset, kvsh_stride, gl_CooperativeMatrixLayoutRowMajor); } - PVMat = coopMatMulAdd(KMat, QMat, PVMat); + PVMat = coopMatMulAdd(PMat[bc_chunk], QMat, PVMat); } // Store PVMat to pvsh and load into Of @@ -549,10 +545,10 @@ void main() { // If there is split_k, then the split_k resolve shader does the final // division by L. Store the intermediate O value and per-row m and L values. - if (p.k_num > 1) { - if (p.gqa_ratio > 1) { + if (partial_output || split_k_num > 1) { + if (gqa_ratio > 1) { // note: O and Q have swapped coord 1,2. - uint32_t o_offset = HSV * p.ne1 * (split_k_index + p.k_num * (gqa_iq1 + p.ne2 * iq3)) / 4; + uint32_t o_offset = HSV * p.ne1 * (split_k_index + output_k_num * (gqa_iq1 + p.ne2 * iq3)) / 4; [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { if (tile_row(r) < N) { @@ -565,7 +561,7 @@ void main() { } } - o_offset = HSV * p.ne1 * p.k_num * p.ne2 * p.ne3 + p.ne1 * 2 * (split_k_index + p.k_num * (gqa_iq1 + p.ne2 * iq3)); + o_offset = HSV * p.ne1 * output_k_num * p.ne2 * p.ne3 + p.ne1 * 2 * (split_k_index + output_k_num * (gqa_iq1 + p.ne2 * iq3)); [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { if (tile_row(r) < N) { perElemOpStoreCol0(tile_row(r), 0u, ACC_TYPE(Lf[r]), o_offset, iq2, N); @@ -578,7 +574,7 @@ void main() { const uint global_row = i * Br + row; if (global_row < N) { - uint32_t o_offset = HSV * p.ne1 * (split_k_index + p.k_num * (global_row + p.ne2 * iq3)) / 4; + uint32_t o_offset = HSV * p.ne1 * (split_k_index + output_k_num * (global_row + p.ne2 * iq3)) / 4; [[unroll]] for (uint32_t d0 = 0; d0 < HSV / 4; d0 += threads_per_rowgroup) { const uint d = d0 + col_tid; @@ -588,7 +584,7 @@ void main() { } if (global_row < N && col_tid == 0) { - uint32_t lm_offset = HSV * p.ne1 * p.k_num * p.ne2 * p.ne3 + p.ne1 * 2 * (split_k_index + p.k_num * (global_row + p.ne2 * iq3)); + uint32_t lm_offset = HSV * p.ne1 * output_k_num * p.ne2 * p.ne3 + p.ne1 * 2 * (split_k_index + output_k_num * (global_row + p.ne2 * iq3)); data_o[lm_offset + iq2] = D_TYPE(Lf[r]); data_o[lm_offset + p.ne1 + iq2] = D_TYPE(Mf[r]); } @@ -637,7 +633,7 @@ void main() { uint32_t o_offset = (gqa_iq1*p.ne1*HSV + iq3*p.ne2*p.ne1*HSV) / 4; - if (p.gqa_ratio > 1) { + if (gqa_ratio > 1) { [[unroll]] for (uint32_t r = 0; r < rows_per_thread; ++r) { if (tile_row(r) < N) { [[unroll]] for (uint32_t d0 = 0; d0 < HSV / 4; d0 += threads_per_rowgroup) { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp index c6ed63dd42a6..54be1e6daa99 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_cm2.comp @@ -29,12 +29,6 @@ #include "dequant_funcs_cm2.glsl" #endif -#ifdef GL_NV_cooperative_matrix_decode_vector -#define FA_GATHER_BS 4u -#else -#define FA_GATHER_BS 1u -#endif - // buffer_reference stride = sizeof(struct) = FaBlockBytesK/V. layout(buffer_reference, std430, buffer_reference_align = 1) buffer decodeBufFA_K { uint8_t raw[FaBlockBytesK]; @@ -46,28 +40,26 @@ layout(buffer_reference, std430, buffer_reference_align = 1) buffer decodeBufFA_ #if !defined(BFLOAT16) float16_t faDecodeK(const decodeBufFA_K bl_in, const uint blockCoords[2], const uint coordInBlock[2]) { switch (FaTypeK) { - case GGML_TYPE_F32: return dequantFuncF32 (decodeBufF32 (bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q4_0: return dequantFuncQ4_0(decodeBufQ4_0(bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q4_1: return dequantFuncQ4_1(decodeBufQ4_1(bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q5_0: return dequantFuncQ5_0(decodeBufQ5_0(bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q5_1: return dequantFuncQ5_1(decodeBufQ5_1(bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q8_0: return dequantFuncQ8_0(decodeBufQ8_0(bl_in), blockCoords, coordInBlock); - case GGML_TYPE_IQ4_NL: return dequantFuncIQ4_NL(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q1_0: return dequantFuncQ1_0(decodeBufQ1_0(bl_in), blockCoords, coordInBlock); + case FA_TYPE_F32: return dequantFuncF32 (decodeBufF32 (bl_in), blockCoords, coordInBlock); + case FA_TYPE_Q4_0: return dequantFuncQ4_0(decodeBufQ4_0(bl_in), blockCoords, coordInBlock); + case FA_TYPE_Q4_1: return dequantFuncQ4_1(decodeBufQ4_1(bl_in), blockCoords, coordInBlock); + case FA_TYPE_Q5_0: return dequantFuncQ5_0(decodeBufQ5_0(bl_in), blockCoords, coordInBlock); + case FA_TYPE_Q5_1: return dequantFuncQ5_1(decodeBufQ5_1(bl_in), blockCoords, coordInBlock); + case FA_TYPE_Q8_0: return dequantFuncQ8_0(decodeBufQ8_0(bl_in), blockCoords, coordInBlock); + case FA_TYPE_IQ4_NL: return dequantFuncIQ4_NL(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock); default: return float16_t(0); } } float16_t faDecodeV(const decodeBufFA_V bl_in, const uint blockCoords[2], const uint coordInBlock[2]) { switch (FaTypeV) { - case GGML_TYPE_F32: return dequantFuncF32 (decodeBufF32 (bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q4_0: return dequantFuncQ4_0(decodeBufQ4_0(bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q4_1: return dequantFuncQ4_1(decodeBufQ4_1(bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q5_0: return dequantFuncQ5_0(decodeBufQ5_0(bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q5_1: return dequantFuncQ5_1(decodeBufQ5_1(bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q8_0: return dequantFuncQ8_0(decodeBufQ8_0(bl_in), blockCoords, coordInBlock); - case GGML_TYPE_IQ4_NL: return dequantFuncIQ4_NL(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q1_0: return dequantFuncQ1_0(decodeBufQ1_0(bl_in), blockCoords, coordInBlock); + case FA_TYPE_F32: return dequantFuncF32 (decodeBufF32 (bl_in), blockCoords, coordInBlock); + case FA_TYPE_Q4_0: return dequantFuncQ4_0(decodeBufQ4_0(bl_in), blockCoords, coordInBlock); + case FA_TYPE_Q4_1: return dequantFuncQ4_1(decodeBufQ4_1(bl_in), blockCoords, coordInBlock); + case FA_TYPE_Q5_0: return dequantFuncQ5_0(decodeBufQ5_0(bl_in), blockCoords, coordInBlock); + case FA_TYPE_Q5_1: return dequantFuncQ5_1(decodeBufQ5_1(bl_in), blockCoords, coordInBlock); + case FA_TYPE_Q8_0: return dequantFuncQ8_0(decodeBufQ8_0(bl_in), blockCoords, coordInBlock); + case FA_TYPE_IQ4_NL: return dequantFuncIQ4_NL(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock); default: return float16_t(0); } } @@ -75,26 +67,26 @@ float16_t faDecodeV(const decodeBufFA_V bl_in, const uint blockCoords[2], const // V=4 vector decode for K/V; dispatches to per-format _v decoders. f16vec4 faDecodeKVector(const decodeBufFA_K bl_in, const uint blockCoords[2], const uint coordInBlock[2]) { switch (FaTypeK) { - case GGML_TYPE_F32: return f16vec4(decodeBufF32(bl_in).block); - case GGML_TYPE_Q4_0: return dequantFuncQ4_0_v(decodeBufQ4_0(bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q4_1: return dequantFuncQ4_1_v(decodeBufQ4_1(bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q5_0: return dequantFuncQ5_0_v(decodeBufQ5_0(bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q5_1: return dequantFuncQ5_1_v(decodeBufQ5_1(bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q8_0: return dequantFuncQ8_0_v(decodeBufQ8_0(bl_in), blockCoords, coordInBlock); - case GGML_TYPE_IQ4_NL: return dequantFuncIQ4_NL_v(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock); + case FA_TYPE_F32: return f16vec4(decodeBufF32(bl_in).block); + case FA_TYPE_Q4_0: return dequantFuncQ4_0_v(decodeBufQ4_0(bl_in), blockCoords, coordInBlock); + case FA_TYPE_Q4_1: return dequantFuncQ4_1_v(decodeBufQ4_1(bl_in), blockCoords, coordInBlock); + case FA_TYPE_Q5_0: return dequantFuncQ5_0_v(decodeBufQ5_0(bl_in), blockCoords, coordInBlock); + case FA_TYPE_Q5_1: return dequantFuncQ5_1_v(decodeBufQ5_1(bl_in), blockCoords, coordInBlock); + case FA_TYPE_Q8_0: return dequantFuncQ8_0_v(decodeBufQ8_0(bl_in), blockCoords, coordInBlock); + case FA_TYPE_IQ4_NL: return dequantFuncIQ4_NL_v(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock); default: return f16vec4(0); } } f16vec4 faDecodeVVector(const decodeBufFA_V bl_in, const uint blockCoords[2], const uint coordInBlock[2]) { switch (FaTypeV) { - case GGML_TYPE_F32: return f16vec4(decodeBufF32(bl_in).block); - case GGML_TYPE_Q4_0: return dequantFuncQ4_0_v(decodeBufQ4_0(bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q4_1: return dequantFuncQ4_1_v(decodeBufQ4_1(bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q5_0: return dequantFuncQ5_0_v(decodeBufQ5_0(bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q5_1: return dequantFuncQ5_1_v(decodeBufQ5_1(bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q8_0: return dequantFuncQ8_0_v(decodeBufQ8_0(bl_in), blockCoords, coordInBlock); - case GGML_TYPE_IQ4_NL: return dequantFuncIQ4_NL_v(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock); + case FA_TYPE_F32: return f16vec4(decodeBufF32(bl_in).block); + case FA_TYPE_Q4_0: return dequantFuncQ4_0_v(decodeBufQ4_0(bl_in), blockCoords, coordInBlock); + case FA_TYPE_Q4_1: return dequantFuncQ4_1_v(decodeBufQ4_1(bl_in), blockCoords, coordInBlock); + case FA_TYPE_Q5_0: return dequantFuncQ5_0_v(decodeBufQ5_0(bl_in), blockCoords, coordInBlock); + case FA_TYPE_Q5_1: return dequantFuncQ5_1_v(decodeBufQ5_1(bl_in), blockCoords, coordInBlock); + case FA_TYPE_Q8_0: return dequantFuncQ8_0_v(decodeBufQ8_0(bl_in), blockCoords, coordInBlock); + case FA_TYPE_IQ4_NL: return dequantFuncIQ4_NL_v(decodeBufIQ4_NL(bl_in), blockCoords, coordInBlock); default: return f16vec4(0); } } @@ -113,67 +105,6 @@ layout (binding = 1) readonly buffer K {uint8_t data_k[];}; layout (binding = 2) readonly buffer V {uint8_t data_v[];}; layout (binding = 3) readonly buffer M {uint8_t data_m[];}; -// f16 aliases for the sparse gather callbacks. -layout (binding = 1) readonly buffer KF16 {float16_t data_kf16[];}; -layout (binding = 2) readonly buffer VF16 {float16_t data_vf16[];}; -layout (binding = 3) readonly buffer MF16 {float16_t data_mf16[];}; -#ifdef GL_NV_cooperative_matrix_decode_vector -layout (binding = 1) readonly buffer KF16V4 {f16vec4 data_kf16v4[];}; -layout (binding = 2) readonly buffer VF16V4 {f16vec4 data_vf16v4[];}; -#endif - -// K/V/mask f16-element offsets for the current head/batch, set in main(). -uint32_t g_k_off_elem, g_v_off_elem, g_m_off_elem; - -#if !defined(BFLOAT16) -// blockCoords are in block units: KV slot = blockCoords[0], -// head dim = blockCoords[1]*FA_GATHER_BS + coordInBlock[1]. -float16_t faGatherK(const decodeBufFA_K unused, const uint32_t blockCoords[2], const uint32_t coordInBlock[2]) { - if (blockCoords[0] >= p.split_kv) { return float16_t(0); } - const int r = data_sparse[sparse_base + blockCoords[0]]; - return r < 0 ? float16_t(0) : data_kf16[g_k_off_elem + uint(r) * k_stride + blockCoords[1] * FA_GATHER_BS + coordInBlock[1]]; -} - -float16_t faGatherV(const decodeBufFA_V unused, const uint32_t blockCoords[2], const uint32_t coordInBlock[2]) { - if (blockCoords[0] >= p.split_kv) { return float16_t(0); } - const int r = data_sparse[sparse_base + blockCoords[0]]; - return r < 0 ? float16_t(0) : data_vf16[g_v_off_elem + uint(r) * v_stride + blockCoords[1] * FA_GATHER_BS + coordInBlock[1]]; -} - -#ifdef GL_NV_cooperative_matrix_decode_vector -f16vec4 faGatherKVector(const decodeBufFA_K unused, const uint32_t blockCoords[2], const uint32_t coordInBlock[2]) { - if (blockCoords[0] >= p.split_kv) { return f16vec4(0); } - const int r = data_sparse[sparse_base + blockCoords[0]]; - if (r < 0) { return f16vec4(0); } - const uint32_t o = g_k_off_elem + uint(r) * k_stride + blockCoords[1] * FA_GATHER_BS + coordInBlock[1]; - return data_kf16v4[o / 4]; -} - -f16vec4 faGatherVVector(const decodeBufFA_V unused, const uint32_t blockCoords[2], const uint32_t coordInBlock[2]) { - if (blockCoords[0] >= p.split_kv) { return f16vec4(0); } - const int r = data_sparse[sparse_base + blockCoords[0]]; - if (r < 0) { return f16vec4(0); } - const uint32_t o = g_v_off_elem + uint(r) * v_stride + blockCoords[1] * FA_GATHER_BS + coordInBlock[1]; - return data_vf16v4[o / 4]; -} - -#define FAGATHERK , faGatherK, faGatherKVector -#define FAGATHERV , faGatherV, faGatherVVector -#else -#define FAGATHERK , faGatherK -#define FAGATHERV , faGatherV -#endif -#endif - -// Add gathered mask to S (slope==1 since sparse requires max_bias==0). col = slot in block jblk. -ACC_TYPE faAddSparseMask(const uint32_t row, const uint32_t col, const ACC_TYPE elem, const uint32_t jblk) { - const float NEG = uintBitsToFloat(0xFEFFFFFF); - const uint32_t kvslot = jblk * Bc + col; - if (kvslot >= p.split_kv) { return ACC_TYPE(NEG); } - const int r = data_sparse[sparse_base + kvslot]; - return r < 0 ? ACC_TYPE(NEG) : elem + ACC_TYPE(data_mf16[g_m_off_elem + row * m_stride + uint(r)]); -} - ACC_TYPE maxReduce(const in ACC_TYPE x, const in ACC_TYPE y) { return max(x, y); } @@ -220,7 +151,7 @@ D_TYPE perElemOpNonGqaSplitKStore(const in uint32_t r, const in uint32_t c, cons uint32_t global_row = i * Br + r; if (global_row < N && c < HSV) { uint32_t o_off = HSV * p.ne1 - * (split_k_index + p.k_num * (global_row + p.ne2 * iq3)); + * (split_k_index + output_k_num * (global_row + p.ne2 * iq3)); data_o[o_off + iq2 * HSV + c] = D_TYPE(elem); } return elem; @@ -230,8 +161,8 @@ D_TYPE perElemOpNonGqaSplitKStore(const in uint32_t r, const in uint32_t c, cons ACC_TYPE perElemOpNonGqaSplitKStoreCol0(const in uint32_t r, const in uint32_t c, const in ACC_TYPE elem, const in uint32_t lm_base, const in uint32_t iq2, const in uint32_t N) { uint32_t global_row = i * Br + r; if (global_row < N && c == 0) { - uint32_t lm_off = HSV * p.ne1 * p.k_num * p.ne2 * p.ne3 - + p.ne1 * 2 * (split_k_index + p.k_num * (global_row + p.ne2 * iq3)); + uint32_t lm_off = HSV * p.ne1 * output_k_num * p.ne2 * p.ne3 + + p.ne1 * 2 * (split_k_index + output_k_num * (global_row + p.ne2 * iq3)); data_o[lm_off + lm_base + iq2] = D_TYPE(elem); } return elem; @@ -252,16 +183,14 @@ void main() { tensorViewNV<2, false, 1, 0> tensorViewTranspose = createTensorViewNV(2, false, 1, 0); - const uint bs_k = USE_SPARSE ? FA_GATHER_BS : fa_block_elems(FaTypeK); - const uint bs_v = USE_SPARSE ? FA_GATHER_BS : fa_block_elems(FaTypeV); + const uint bs_k = fa_block_elems(FaTypeK); + const uint bs_v = fa_block_elems(FaTypeV); tensorLayoutK = setTensorLayoutBlockSizeNV(tensorLayoutK, 1, bs_k); tensorLayoutV = setTensorLayoutBlockSizeNV(tensorLayoutV, 1, bs_v); - // Sparse iterates n_kv_max (in split_kv); the decode callbacks remap each slot. - const uint32_t KV_iter = USE_SPARSE ? p.split_kv : KV; tensorLayoutQ = setTensorLayoutDimensionNV(tensorLayoutQ, N, HSK); - tensorLayoutK = setTensorLayoutDimensionNV(tensorLayoutK, KV_iter, HSK); - tensorLayoutV = setTensorLayoutDimensionNV(tensorLayoutV, KV_iter, HSV); + tensorLayoutK = setTensorLayoutDimensionNV(tensorLayoutK, KV, HSK); + tensorLayoutV = setTensorLayoutDimensionNV(tensorLayoutV, KV, HSV); // hint to the compiler that strides are aligned for the aligned variant of the shader if (Clamp != gl_CooperativeMatrixClampModeConstantNV) @@ -313,16 +242,12 @@ void main() { // mo_offset will point to the tile starting at row i*Br and col 0 uint32_t mo_offset = mo_stride * i; - uint32_t m_offset = gqa_iq1*KV * 2 /*sizeof(float16_t)*/; + uint32_t m_offset = gqa_iq1*m_row_len * 2 /*sizeof(float16_t)*/; if (p.nem2 != 1 || p.nem3 != 1) { - m_offset += ((iq3 % p.nem3) * p.nem2 + (iq2 % p.nem2)) * p.nem1 * KV * 2 /*sizeof(float16_t)*/; + m_offset += ((iq3 % p.nem3) * p.nem2 + (iq2 % p.nem2)) * p.nem1 * m_row_len * 2 /*sizeof(float16_t)*/; mo_offset += ((iq3 % p.nem3) * p.nem2 + (iq2 % p.nem2)) * CEIL_DIV(p.nem1, Br) * mo_stride; } - g_k_off_elem = (ik2*p.nb12 + ik3*p.nb13) / 2; - g_v_off_elem = (iv2*p.nb22 + iv3*p.nb23) / 2; - g_m_off_elem = m_offset / 2; - uint32_t mask_opt = 0; uint32_t mask_opt_idx = ~0; @@ -330,7 +255,7 @@ void main() { for (uint32_t j = start_j; j < end_j; ++j) { coopmat mv = coopmat(0); - if (MASK_ENABLE && !USE_SPARSE) { + if (MASK_ENABLE) { if (USE_MASK_OPT && mask_opt_idx != j / 16) { mask_opt_idx = j / 16; @@ -343,7 +268,7 @@ void main() { } // Only load if the block is not all zeros if (mask_opt_bits != MASK_OPT_ALL_ZERO) { - bool nem1_bounds_check = !(p.gqa_ratio > 1) && (p.nem1 % Br) != 0; + bool nem1_bounds_check = !(gqa_ratio > 1) && (p.nem1 % Br) != 0; if (nem1_bounds_check) { tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutM = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); @@ -362,7 +287,7 @@ void main() { } else { tensorLayoutNV<2, Clamp> tensorLayoutM = createTensorLayoutNV(2, Clamp); // Don't clamp against nem1 when GQA is enabled - uint32_t m_height = p.gqa_ratio > 1 ? ~0 : p.nem1; + uint32_t m_height = gqa_ratio > 1 ? ~0 : p.nem1; tensorLayoutM = setTensorLayoutDimensionNV(tensorLayoutM, m_height, KV); tensorLayoutM = setTensorLayoutStrideNV(tensorLayoutM, m_stride, 1); @@ -388,9 +313,7 @@ void main() { coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, HSK_pad), tensorViewTranspose); #else const bool k_use_decode = (bs_k > 1u); - if (USE_SPARSE) { - coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, HSK_pad), tensorViewTranspose FAGATHERK); - } else if (k_use_decode) { + if (k_use_decode) { coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, HSK_pad), tensorViewTranspose FADECODEK); } else { coopMatLoadTensorNV(K_T, data_k, k_offset, sliceTensorLayoutNV(tensorLayoutK, j * Bc, Bc, 0, HSK_pad), tensorViewTranspose); @@ -405,9 +328,7 @@ void main() { } } - if (MASK_ENABLE && USE_SPARSE) { - coopMatPerElementNV(S, S, faAddSparseMask, j); - } else if (MASK_ENABLE) { + if (MASK_ENABLE) { S += slopeMat*coopmat(mv); } @@ -462,9 +383,7 @@ void main() { coopMatLoadTensorNV(V, data_v, v_offset, sliceTensorLayoutNV(tensorLayoutV, j * Bc, Bc, 0, HSV_pad)); #else const bool v_use_decode = (bs_v > 1u); - if (USE_SPARSE) { - coopMatLoadTensorNV(V, data_v, v_offset, sliceTensorLayoutNV(tensorLayoutV, j * Bc, Bc, 0, HSV_pad) FAGATHERV); - } else if (v_use_decode) { + if (v_use_decode) { coopMatLoadTensorNV(V, data_v, v_offset, sliceTensorLayoutNV(tensorLayoutV, j * Bc, Bc, 0, HSV_pad) FADECODEV); } else { coopMatLoadTensorNV(V, data_v, v_offset, sliceTensorLayoutNV(tensorLayoutV, j * Bc, Bc, 0, HSV_pad)); @@ -487,15 +406,15 @@ void main() { // If there is split_k, then the split_k resolve shader does the final // division by L. Store the intermediate O value and per-row m and L values. - if (p.k_num > 1) { + if (partial_output || split_k_num > 1) { coopmat O_D = coopmat(O); - if (p.gqa_ratio > 1) { + if (gqa_ratio > 1) { // note: O and Q have swapped coord 1,2. - uint32_t o_offset = HSV * p.ne1 * (split_k_index + p.k_num * (gqa_iq1 + p.ne2 * iq3)); + uint32_t o_offset = HSV * p.ne1 * (split_k_index + output_k_num * (gqa_iq1 + p.ne2 * iq3)); coopMatPerElementNV(O_D, O_D, perElemOpGqaStore, o_offset, iq2, N); - o_offset = HSV * p.ne1 * p.k_num * p.ne2 * p.ne3 + p.ne1 * 2 * (split_k_index + p.k_num * (gqa_iq1 + p.ne2 * iq3)); + o_offset = HSV * p.ne1 * output_k_num * p.ne2 * p.ne3 + p.ne1 * 2 * (split_k_index + output_k_num * (gqa_iq1 + p.ne2 * iq3)); coopMatPerElementNV(L, L, perElemOpStoreCol0, o_offset, iq2, N); coopMatPerElementNV(M, M, perElemOpStoreCol0, o_offset + p.ne1, iq2, N); } else { @@ -554,7 +473,7 @@ void main() { uint32_t o_offset = gqa_iq1*p.ne1*HSV + iq3*p.ne2*p.ne1*HSV; - if (p.gqa_ratio > 1) { + if (gqa_ratio > 1) { coopMatPerElementNV(O_D, O_D, perElemOpGqaStore, o_offset, iq2, N); } else { tensorLayoutNV<3, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutD = createTensorLayoutNV(3, gl_CooperativeMatrixClampModeConstantNV); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_dequant.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_dequant.glsl index 4fcf7c1f4f64..8ba4725f3342 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_dequant.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_dequant.glsl @@ -121,25 +121,25 @@ layout (binding = 1) readonly buffer K_PACKED_Q5_1_P32 { block_q5_1_packed32 dat FLOAT_TYPEV4 dequantize4(uint ib, uint iqs, uint a_offset, uint binding_idx) { if (binding_idx == BINDING_IDX_K) { switch (FaTypeK) { - case GGML_TYPE_F32: FA_DEQUANT4_F32 (k_packed_f32) - case GGML_TYPE_Q4_0: FA_DEQUANT4_Q4_0(k_packed_q4_0) - case GGML_TYPE_Q4_1: FA_DEQUANT4_Q4_1(k_packed_q4_1) - case GGML_TYPE_Q5_0: FA_DEQUANT4_Q5_0(k_packed_q5_0) - case GGML_TYPE_Q5_1: FA_DEQUANT4_Q5_1(k_packed_q5_1) - case GGML_TYPE_Q8_0: FA_DEQUANT4_Q8_0(k_packed_q8_0) - case GGML_TYPE_IQ4_NL: FA_DEQUANT4_IQ4_NL(k_packed_iq4_nl) - case GGML_TYPE_BF16: FA_DEQUANT4_BF16(k_packed_bf16) + case FA_TYPE_F32: FA_DEQUANT4_F32 (k_packed_f32) + case FA_TYPE_Q4_0: FA_DEQUANT4_Q4_0(k_packed_q4_0) + case FA_TYPE_Q4_1: FA_DEQUANT4_Q4_1(k_packed_q4_1) + case FA_TYPE_Q5_0: FA_DEQUANT4_Q5_0(k_packed_q5_0) + case FA_TYPE_Q5_1: FA_DEQUANT4_Q5_1(k_packed_q5_1) + case FA_TYPE_Q8_0: FA_DEQUANT4_Q8_0(k_packed_q8_0) + case FA_TYPE_IQ4_NL: FA_DEQUANT4_IQ4_NL(k_packed_iq4_nl) + case FA_TYPE_BF16: FA_DEQUANT4_BF16(k_packed_bf16) } } else { switch (FaTypeV) { - case GGML_TYPE_F32: FA_DEQUANT4_F32 (v_packed_f32) - case GGML_TYPE_Q4_0: FA_DEQUANT4_Q4_0(v_packed_q4_0) - case GGML_TYPE_Q4_1: FA_DEQUANT4_Q4_1(v_packed_q4_1) - case GGML_TYPE_Q5_0: FA_DEQUANT4_Q5_0(v_packed_q5_0) - case GGML_TYPE_Q5_1: FA_DEQUANT4_Q5_1(v_packed_q5_1) - case GGML_TYPE_Q8_0: FA_DEQUANT4_Q8_0(v_packed_q8_0) - case GGML_TYPE_IQ4_NL: FA_DEQUANT4_IQ4_NL(v_packed_iq4_nl) - case GGML_TYPE_BF16: FA_DEQUANT4_BF16(v_packed_bf16) + case FA_TYPE_F32: FA_DEQUANT4_F32 (v_packed_f32) + case FA_TYPE_Q4_0: FA_DEQUANT4_Q4_0(v_packed_q4_0) + case FA_TYPE_Q4_1: FA_DEQUANT4_Q4_1(v_packed_q4_1) + case FA_TYPE_Q5_0: FA_DEQUANT4_Q5_0(v_packed_q5_0) + case FA_TYPE_Q5_1: FA_DEQUANT4_Q5_1(v_packed_q5_1) + case FA_TYPE_Q8_0: FA_DEQUANT4_Q8_0(v_packed_q8_0) + case FA_TYPE_IQ4_NL: FA_DEQUANT4_IQ4_NL(v_packed_iq4_nl) + case FA_TYPE_BF16: FA_DEQUANT4_BF16(v_packed_bf16) } } return FLOAT_TYPEV4(0); diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather.comp new file mode 100644 index 000000000000..4e16b2a9cece --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather.comp @@ -0,0 +1,117 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require + +// Gathers the active KV rows of a top-k sparse attention (DeepSeek V4 CSA) into a compact +// contiguous scratch: rows [0, n_kv_raw) of the source (the dense prefix), then, for each of +// the n_batch query tokens, that token's n_top_k selected rows, then zero padding up to kv_c. +// +// n_batch > 1 (small-batch decode, e.g. speculative drafts) gives every token its own block +// rather than deduplicating into a union. That costs n_batch*n_top_k gathered rows instead of +// |union|, but needs no dedup pass and no atomics, and the size is bounded by the worst case +// the caller already allocates for. Cross-token rows are neutralised through the mask: token t +// sees -inf on any block that is not its own, so nothing is double counted in the softmax. The gathered mask row keeps the +// per-key mask values so causality/validity survive compaction; invalid top-k indices and +// padding get -inf mask and zeroed K (softmax-neutral either way, zeroed so no NaN*0). +// One workgroup per compact row; V is the K latent (V==K), so a single gather serves both. + +layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; + +// K is addressed as raw 4-byte WORDS, never as values: a gather relocates rows and does not +// need to know whether they hold f16 elements or quantised blocks. Only the MASK is typed. +layout(binding = 0) readonly buffer KBuf { uint data_k[]; }; +layout(binding = 1) readonly buffer TopBuf { int data_top[]; }; +layout(binding = 2) readonly buffer MBuf { float16_t data_m[]; }; +layout(binding = 3) writeonly buffer KcBuf { uint data_kc[]; }; +layout(binding = 4) writeonly buffer McBuf { float16_t data_mc[]; }; + +layout(push_constant) uniform Parameters { + uint n_kv; // total source KV rows + uint n_kv_raw; // dense prefix length + uint n_top_k; // selected rows for the (single) query token + uint kv_c; // padded compact row count == dispatch row range + uint nbk1; // K source row stride, WORDS + uint nbk3; // K source stream stride, WORDS + uint nbt1; // top_k row (per query token) stride, elements + uint nbt3; // top_k stream stride, elements + uint nbm1; // mask source row (per query token) stride, elements + uint nbm3; // mask source stream stride, elements + uint nem3; // mask ne[3], for stream broadcast + uint n_batch; // query tokens sharing this gather; <= LANES + uint row_words; // bytes per K row / 4 +} p; + +const uint LANES = 64; + +void main() { + const uint row = gl_WorkGroupID.x; + const uint stream = gl_WorkGroupID.z; + const uint tid = gl_LocalInvocationIndex; + + // map compact row -> source row; p.n_kv is the invalid sentinel. + // owner is the token whose block this row belongs to, or ALL_TOKENS for the shared prefix. + const uint ALL_TOKENS = 0xffffffffu; + uint src = p.n_kv; + uint owner = ALL_TOKENS; + if (row < p.n_kv_raw) { + src = row; + } else { + const uint off = row - p.n_kv_raw; + const uint tok = off / p.n_top_k; + const uint slot = off - tok * p.n_top_k; + if (tok < p.n_batch) { + owner = tok; + const int idx = data_top[stream * p.nbt3 + tok * p.nbt1 + slot]; + if (idx >= 0 && uint(idx) < p.n_kv - p.n_kv_raw) { + src = p.n_kv_raw + uint(idx); + } + } + } + + // Zeroing writes zero BYTES, which decode to zero for every block-quantised type here (a + // zero scale zeroes the block) as well as for f16. The row is -inf in the mask either way; + // zeroing only keeps a garbage dot product from reaching the softmax as a NaN. + const uint dst_base = (stream * p.kv_c + row) * p.row_words; + // row_words is a push constant, so this cannot be [[unroll]]ed the way it was when the row + // was a fixed 512 f16 elements. Stepping 4*LANES recovers that: each access stays contiguous + // across the lanes, and an f16 row (256 words == 4*LANES) is one iteration with no loop + // overhead. The tail carries the quantised row sizes, which are not multiples of 4*LANES. + if (src < p.n_kv) { + const uint src_base = stream * p.nbk3 + src * p.nbk1; + uint i = tid; + for (; i + 3 * LANES < p.row_words; i += 4 * LANES) { + data_kc[dst_base + i] = data_k[src_base + i]; + data_kc[dst_base + i + LANES] = data_k[src_base + i + LANES]; + data_kc[dst_base + i + 2 * LANES] = data_k[src_base + i + 2 * LANES]; + data_kc[dst_base + i + 3 * LANES] = data_k[src_base + i + 3 * LANES]; + } + for (; i < p.row_words; i += LANES) { + data_kc[dst_base + i] = data_k[src_base + i]; + } + } else { + uint i = tid; + for (; i + 3 * LANES < p.row_words; i += 4 * LANES) { + data_kc[dst_base + i] = 0u; + data_kc[dst_base + i + LANES] = 0u; + data_kc[dst_base + i + 2 * LANES] = 0u; + data_kc[dst_base + i + 3 * LANES] = 0u; + } + for (; i < p.row_words; i += LANES) { + data_kc[dst_base + i] = 0u; + } + } + + // Compact mask is token-major [n_batch][kv_c], which is what the GQA mask path expects + // (m_stride 0, rows stepped by gqa_iq1 * m_row_len). One lane per token; n_batch <= LANES. + const float NEG_INF = uintBitsToFloat(0xff800000); + if (tid < p.n_batch) { + const uint mc_idx = (stream * p.n_batch + tid) * p.kv_c + row; + float mv = NEG_INF; + if (src < p.n_kv && (owner == ALL_TOKENS || owner == tid)) { + mv = float(data_m[(stream % p.nem3) * p.nbm3 + tid * p.nbm1 + src]); + } + data_mc[mc_idx] = float16_t(mv); + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_dq.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_dq.comp new file mode 100644 index 000000000000..f2711539641a --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_dq.comp @@ -0,0 +1,106 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require +#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require + +#define FLOAT_TYPE float +#include "types.glsl" + +// Gather + DEQUANTISE for the per-token compact form: same row mapping as the verbatim +// flash_attn_gather.comp, same decode as flash_attn_gather_union_dq.comp. One variant per K +// type, via DATA_A_*. +// +// The union carries the decode only for n_batch > 1, because it needs more than one token to +// deduplicate. Single-token decode is the common case and it lands here; a verbatim gather +// leaves flash attention to decode inline, measured 1.98x the f16 op at q8_0 and 2.26x at +// q4_0, kv 11008. + +layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; + +layout(binding = 0) readonly buffer A { A_TYPE data_a[]; }; +layout(binding = 1) readonly buffer TopBuf { int data_top[]; }; +layout(binding = 2) readonly buffer MBuf { float16_t data_m[]; }; +layout(binding = 3) writeonly buffer KcBuf { float16_t data_kc[]; }; +layout(binding = 4) writeonly buffer McBuf { float16_t data_mc[]; }; + +layout(push_constant) uniform Parameters { + uint n_kv; + uint n_kv_raw; + uint n_top_k; + uint kv_c; + uint nbk1; // K source row stride, BLOCKS + uint nbk3; // K source stream stride, BLOCKS + uint nbt1; + uint nbt3; + uint nbm1; + uint nbm3; + uint nem3; + uint n_batch; + uint row_elems; // elements per K row; 64 lanes x 8 covers the 512 the gate pins +} p; + +void main() { + const uint row = gl_WorkGroupID.x; + const uint stream = gl_WorkGroupID.z; + const uint tid = gl_LocalInvocationIndex; + + const uint ALL_TOKENS = 0xffffffffu; + uint src = p.n_kv; + uint owner = ALL_TOKENS; + if (row < p.n_kv_raw) { + src = row; + } else { + const uint off = row - p.n_kv_raw; + const uint tok = off / p.n_top_k; + const uint slot = off - tok * p.n_top_k; + if (tok < p.n_batch) { + owner = tok; + const int idx = data_top[stream * p.nbt3 + tok * p.nbt1 + slot]; + if (idx >= 0 && uint(idx) < p.n_kv - p.n_kv_raw) { + src = p.n_kv_raw + uint(idx); + } + } + } + + const uint dst_base = (stream * p.kv_c + row) * p.row_elems; + const uint e0 = tid * 8; + if (src < p.n_kv) { + const uint ib = stream * p.nbk3 + src * p.nbk1 + e0 / QUANT_K; + const uint iqs = e0 % QUANT_K; + const float d = float(data_a[ib].d); +#if defined(DATA_A_Q8_0) + // element e is qs[e] + [[unroll]] for (uint l = 0; l < 8; ++l) { + data_kc[dst_base + e0 + l] = float16_t(d * float(data_a[ib].qs[iqs + l])); + } +#elif defined(DATA_A_Q4_0) + // byte j carries element j in its low nibble and element j+16 in its high one + const uint shift = (iqs >> 4) * 4; + const uint byte0 = iqs & 0xF; + [[unroll]] for (uint l = 0; l < 8; ++l) { + const float q = float((data_a[ib].qs[byte0 + l] >> shift) & 0xF) - 8.0f; + data_kc[dst_base + e0 + l] = float16_t(d * q); + } +#else +#error "no element mapping written for this K type" +#endif + } else { + [[unroll]] for (uint l = 0; l < 8; ++l) { + data_kc[dst_base + e0 + l] = float16_t(0.0); + } + } + + const float NEG_INF = uintBitsToFloat(0xff800000); + if (tid < p.n_batch) { + const uint mc_idx = (stream * p.n_batch + tid) * p.kv_c + row; + float mv = NEG_INF; + if (src < p.n_kv && (owner == ALL_TOKENS || owner == tid)) { + mv = float(data_m[(stream % p.nem3) * p.nbm3 + tid * p.nbm1 + src]); + } + data_mc[mc_idx] = float16_t(mv); + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_union.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_union.comp new file mode 100644 index 000000000000..306372ba8af6 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_union.comp @@ -0,0 +1,105 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require + +// Gathers K/V and the mask for the DeepSeek V4 small-batch decode path, using the deduplicated +// union produced by flash_attn_union.comp: rows [0, n_kv_raw) of the source, then one row per +// distinct selected compressed row, then padding. +// +// Simpler than the block-per-token form it replaces, because a union row appears exactly once: +// there is no owner to track and no cross-token masking to apply. Each token just reads its own +// mask value for the gathered source row, which is already -inf where that token did not select +// it, so the softmax still cannot double count. +// +// Row count is a runtime value in data_c[0]; rows past the union are zeroed K and -inf mask, +// which is softmax-neutral and keeps the padded tail harmless. + +layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; + +// K is addressed as raw 4-byte WORDS, never as values. A gather relocates rows; it does not +// need to know whether they hold f16 elements or quantised blocks, so the same shader serves +// every K type whose row is a whole number of words. Only the MASK is typed, and a mask is +// always f16. +layout(binding = 0) readonly buffer KBuf { uint data_k[]; }; +layout(binding = 1) readonly buffer UBuf { uint data_u[]; }; +layout(binding = 2) readonly buffer MBuf { float16_t data_m[]; }; +layout(binding = 3) writeonly buffer KcBuf { uint data_kc[]; }; +layout(binding = 4) writeonly buffer McBuf { float16_t data_mc[]; }; +layout(binding = 5) readonly buffer CBuf { uint data_c[]; }; + +layout(push_constant) uniform Parameters { + uint n_kv; + uint n_kv_raw; + uint kv_c_max; + uint nbk1; // K source row stride, WORDS + uint nbm1; + uint n_batch; + uint row_words; // bytes per K row / 4 +} p; + +const uint LANES = 64; + +void main() { + const uint row = gl_WorkGroupID.x; + const uint tid = gl_LocalInvocationIndex; + + const uint kv_c = data_c[0]; // padded compact rows, the FA's runtime KV + const uint n_uni = data_c[1]; // unpadded union size + + if (row >= kv_c) { + return; + } + + uint src = p.n_kv; // sentinel: invalid + if (row < p.n_kv_raw) { + src = row; + } else if (row - p.n_kv_raw < n_uni) { + src = p.n_kv_raw + data_u[row - p.n_kv_raw]; + } + + // Zeroing an unused row writes zero BYTES, which decode to zero for every block-quantised + // type here (a zero scale zeroes the block) as well as for f16. The row is -inf in the mask + // either way; zeroing only keeps a garbage dot product from reaching the softmax as a NaN. + const uint dst_base = row * p.row_words; + // Stepped 4*LANES for the same reason as flash_attn_gather.comp: row_words is a push + // constant, so the copy cannot be [[unroll]]ed, and an f16 row is one iteration this way. + if (src < p.n_kv) { + const uint src_base = src * p.nbk1; + uint i = tid; + for (; i + 3 * LANES < p.row_words; i += 4 * LANES) { + data_kc[dst_base + i] = data_k[src_base + i]; + data_kc[dst_base + i + LANES] = data_k[src_base + i + LANES]; + data_kc[dst_base + i + 2 * LANES] = data_k[src_base + i + 2 * LANES]; + data_kc[dst_base + i + 3 * LANES] = data_k[src_base + i + 3 * LANES]; + } + for (; i < p.row_words; i += LANES) { + data_kc[dst_base + i] = data_k[src_base + i]; + } + } else { + uint i = tid; + for (; i + 3 * LANES < p.row_words; i += 4 * LANES) { + data_kc[dst_base + i] = 0u; + data_kc[dst_base + i + LANES] = 0u; + data_kc[dst_base + i + 2 * LANES] = 0u; + data_kc[dst_base + i + 3 * LANES] = 0u; + } + for (; i < p.row_words; i += LANES) { + data_kc[dst_base + i] = 0u; + } + } + + // Compact mask is token-major [n_batch][kv_c]. The stride must be the RUNTIME kv_c, not + // kv_c_max: the FA derives m_row_len from KV, which is now that same runtime value, and a + // mismatch would step the mask by the wrong amount for every token past the first. The + // buffer is allocated for kv_c_max, so a smaller stride simply leaves a tail unused. + const float NEG_INF = uintBitsToFloat(0xff800000); + if (tid < p.n_batch) { + float mv = NEG_INF; + if (src < p.n_kv) { + mv = float(data_m[tid * p.nbm1 + src]); + } + data_mc[tid * kv_c + row] = float16_t(mv); + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_union_dq.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_union_dq.comp new file mode 100644 index 000000000000..fd3d0708b2f0 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_gather_union_dq.comp @@ -0,0 +1,123 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_EXT_shader_explicit_arithmetic_types_int8 : require +#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require +#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require + +// iq4_nl's value table lives in shared memory and is typed FLOAT_TYPE +#define FLOAT_TYPE float +#include "types.glsl" + +// Gather + DEQUANTISE in one pass, for the DeepSeek V4 small-batch decode path with a +// quantised KV cache. One variant per K type, via DATA_A_*. +// +// The sibling flash_attn_gather_union.comp relocates rows verbatim, which leaves flash +// attention to decode them in its inner loop - once per query block that reads a row, rather +// than once per row. Measured, that is a constant per KV row attended: +// +// q8_0 0.148 - 0.153 us/row q4_0 0.172 - 0.179 us/row +// +// flat across batch 2/4/8 and across the dense and gathered regimes alike. It is what made +// quantised attention slower than f16 despite reading fewer bytes, and q4_0 worse than q8_0 +// despite reading fewer bytes still: unpacking nibbles costs more ALU than the bytes save. +// +// The gather already touches each selected row exactly once, so decoding here converts per-use +// work into per-row work. The compact scratch becomes f16 and flash attention takes its f16 +// path. The pass itself already existed, so the only added cost is writing f16 rather than +// blocks into the scratch. + +layout(local_size_x = 64, local_size_y = 1, local_size_z = 1) in; + +layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; + +layout(binding = 1) readonly buffer UBuf { uint data_u[]; }; +layout(binding = 2) readonly buffer MBuf { float16_t data_m[]; }; +layout(binding = 3) writeonly buffer KcBuf { float16_t data_kc[]; }; +layout(binding = 4) writeonly buffer McBuf { float16_t data_mc[]; }; +layout(binding = 5) readonly buffer CBuf { uint data_c[]; }; + +layout(push_constant) uniform Parameters { + uint n_kv; + uint n_kv_raw; + uint kv_c_max; + uint nbk1; // K source row stride, BLOCKS + uint nbm1; + uint n_batch; + uint row_elems; // elements per K row (== head size) +} p; + +const uint LANES = 64; + +void main() { +#ifdef NEEDS_INIT_IQ_SHMEM + // barrier inside; must run before any divergent return below + init_iq_shmem(gl_WorkGroupSize); +#endif + const uint row = gl_WorkGroupID.x; + const uint tid = gl_LocalInvocationIndex; + + const uint kv_c = data_c[0]; // padded compact rows, the FA's runtime KV + const uint n_uni = data_c[1]; // unpadded union size + + if (row >= kv_c) { + return; + } + + uint src = p.n_kv; // sentinel: invalid + if (row < p.n_kv_raw) { + src = row; + } else if (row - p.n_kv_raw < n_uni) { + src = p.n_kv_raw + data_u[row - p.n_kv_raw]; + } + + // 64 lanes x 8 elements covers a 512-element row exactly, and QUANT_K is 32, so a lane's 8 + // elements sit wholly inside one block AND wholly inside one nibble half. + // + // Deliberately NOT dequantize4() from dequant_funcs.glsl. That helper is permutation + // AGNOSTIC: mul_mat_vec only ever feeds it into a dot product, where any consistent element + // order gives the same answer, so for q4_0 it returns nibbles in packed order rather than + // element order. Materialising a row to memory is the one use where the order matters, and + // using it here produced NMSE 1.17 on q4_0 while q8_0 passed - q8_0's layout just happens to + // be contiguous. Each type's true element mapping is written out below instead. + const uint dst_base = row * p.row_elems; + const uint e0 = tid * 8; + if (src < p.n_kv) { + const uint ib = src * p.nbk1 + e0 / QUANT_K; + const uint iqs = e0 % QUANT_K; + const float d = float(data_a[ib].d); +#if defined(DATA_A_Q8_0) + // element e is qs[e] + [[unroll]] for (uint l = 0; l < 8; ++l) { + data_kc[dst_base + e0 + l] = float16_t(d * float(data_a[ib].qs[iqs + l])); + } +#elif defined(DATA_A_Q4_0) + // byte j carries element j in its low nibble and element j+16 in its high one + const uint shift = (iqs >> 4) * 4; // 0 for elements 0..15, 4 for 16..31 + const uint byte0 = iqs & 0xF; + [[unroll]] for (uint l = 0; l < 8; ++l) { + const float q = float((data_a[ib].qs[byte0 + l] >> shift) & 0xF) - 8.0f; + data_kc[dst_base + e0 + l] = float16_t(d * q); + } +#else +#error "no element mapping written for this K type" +#endif + } else { + [[unroll]] for (uint j = 0; j < 8; ++j) { + data_kc[dst_base + e0 + j] = float16_t(0.0); + } + } + + // Compact mask is token-major [n_batch][kv_c], stride the RUNTIME kv_c: the FA derives + // m_row_len from KV, which is that same runtime value. + const float NEG_INF = uintBitsToFloat(0xff800000); + if (tid < p.n_batch) { + float mv = NEG_INF; + if (src < p.n_kv) { + mv = float(data_m[tid * p.nbm1 + src]); + } + data_mc[tid * kv_c + row] = float16_t(mv); + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_mmq_funcs.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_mmq_funcs.glsl index 49900aa5aeb3..6bf10a7cffd2 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_mmq_funcs.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_mmq_funcs.glsl @@ -4,20 +4,20 @@ int32_t get_k_qs(uint ib, uint iqs, uint a_offset) { switch (FaTypeK) { - case GGML_TYPE_Q4_0: { + case FA_TYPE_Q4_0: { uint vui = pack32(u16vec2(k_packed_q4_0.data[a_offset + ib].qs[(iqs & 0xF) / 2 + 0], k_packed_q4_0.data[a_offset + ib].qs[(iqs & 0xF) / 2 + 1])); uint shift = (iqs & 0x10) >> 2; vui >>= shift; return int32_t(vui & 0x0F0F0F0F); } - case GGML_TYPE_Q4_1: { // uses packed32 alias + case FA_TYPE_Q4_1: { // uses packed32 alias uint vui = k_packed_q4_1_p32.data[a_offset + ib].qs[(iqs & 0xF) / 4]; uint shift = (iqs & 0x10) >> 2; vui >>= shift; return int32_t(vui & 0x0F0F0F0F); } - case GGML_TYPE_Q5_0: { + case FA_TYPE_Q5_0: { uint vui = pack32(u16vec2(k_packed_q5_0.data[a_offset + ib].qs[(iqs & 0xF) / 2 + 0], k_packed_q5_0.data[a_offset + ib].qs[(iqs & 0xF) / 2 + 1])); uint qh = pack32(u16vec2(k_packed_q5_0.data[a_offset + ib].qh[0], @@ -27,7 +27,7 @@ int32_t get_k_qs(uint ib, uint iqs, uint a_offset) { uint qh_bits = (qh >> iqs) & 0xF; return int32_t(vui & 0x0F0F0F0F) | int32_t((qh_bits * 0x02040810u) & 0x10101010u); } - case GGML_TYPE_Q5_1: { // qs via packed32, qh via packed16 + case FA_TYPE_Q5_1: { // qs via packed32, qh via packed16 uint vui = k_packed_q5_1_p32.data[a_offset + ib].qs[(iqs & 0xF) / 4]; uint qh = k_packed_q5_1.data[a_offset + ib].qh; uint shift = (iqs & 0x10) >> 2; @@ -35,7 +35,7 @@ int32_t get_k_qs(uint ib, uint iqs, uint a_offset) { uint qh_bits = (qh >> iqs) & 0xF; return int32_t(vui & 0x0F0F0F0F) | int32_t((qh_bits * 0x02040810u) & 0x10101010u); } - case GGML_TYPE_Q8_0: { + case FA_TYPE_Q8_0: { return pack32(i16vec2(k_packed_q8_0.data[a_offset + ib].qs[iqs / 2], k_packed_q8_0.data[a_offset + ib].qs[iqs / 2 + 1])); } @@ -47,11 +47,11 @@ int32_t get_k_qs(uint ib, uint iqs, uint a_offset) { // return (d, 0) so call sites always see the same shape. FLOAT_TYPEV2 get_k_scale(uint ib, uint a_offset) { switch (FaTypeK) { - case GGML_TYPE_Q4_0: return FLOAT_TYPEV2(FLOAT_TYPE(k_packed_q4_0.data[a_offset + ib].d), 0.0); - case GGML_TYPE_Q4_1: return FLOAT_TYPEV2(k_packed_q4_1_p32.data[a_offset + ib].dm); - case GGML_TYPE_Q5_0: return FLOAT_TYPEV2(FLOAT_TYPE(k_packed_q5_0.data[a_offset + ib].d), 0.0); - case GGML_TYPE_Q5_1: return FLOAT_TYPEV2(k_packed_q5_1_p32.data[a_offset + ib].dm); - case GGML_TYPE_Q8_0: return FLOAT_TYPEV2(FLOAT_TYPE(k_packed_q8_0.data[a_offset + ib].d), 0.0); + case FA_TYPE_Q4_0: return FLOAT_TYPEV2(FLOAT_TYPE(k_packed_q4_0.data[a_offset + ib].d), 0.0); + case FA_TYPE_Q4_1: return FLOAT_TYPEV2(k_packed_q4_1_p32.data[a_offset + ib].dm); + case FA_TYPE_Q5_0: return FLOAT_TYPEV2(FLOAT_TYPE(k_packed_q5_0.data[a_offset + ib].d), 0.0); + case FA_TYPE_Q5_1: return FLOAT_TYPEV2(k_packed_q5_1_p32.data[a_offset + ib].dm); + case FA_TYPE_Q8_0: return FLOAT_TYPEV2(FLOAT_TYPE(k_packed_q8_0.data[a_offset + ib].d), 0.0); default: return FLOAT_TYPEV2(0); } } @@ -61,16 +61,16 @@ void k_block_to_shmem(const uint buf_ib, const uint global_ib, const uint iqs, c // explicit casts. The bit pattern is what we care about here -- the actual // signed/unsigned interpretation happens downstream in the dot product. switch (FaTypeK) { - case GGML_TYPE_Q4_0: { + case FA_TYPE_Q4_0: { kblocksh[buf_ib].qs[iqs] = int32_t(pack32(u16vec2(k_packed_q4_0.data[a_offset + global_ib].qs[iqs * 2], k_packed_q4_0.data[a_offset + global_ib].qs[iqs * 2 + 1]))); break; } - case GGML_TYPE_Q4_1: { + case FA_TYPE_Q4_1: { kblocksh[buf_ib].qs[iqs] = int32_t(k_packed_q4_1_p32.data[a_offset + global_ib].qs[iqs]); break; } - case GGML_TYPE_Q5_0: { + case FA_TYPE_Q5_0: { kblocksh[buf_ib].qs[iqs] = int32_t(pack32(u16vec2(k_packed_q5_0.data[a_offset + global_ib].qs[iqs * 2], k_packed_q5_0.data[a_offset + global_ib].qs[iqs * 2 + 1]))); if (iqs == 0) { @@ -79,14 +79,14 @@ void k_block_to_shmem(const uint buf_ib, const uint global_ib, const uint iqs, c } break; } - case GGML_TYPE_Q5_1: { + case FA_TYPE_Q5_1: { kblocksh[buf_ib].qs[iqs] = int32_t(k_packed_q5_1_p32.data[a_offset + global_ib].qs[iqs]); if (iqs == 0) { kblocksh[buf_ib].qh = k_packed_q5_1.data[a_offset + global_ib].qh; } break; } - case GGML_TYPE_Q8_0: { + case FA_TYPE_Q8_0: { kblocksh[buf_ib].qs[iqs] = pack32(i16vec2(k_packed_q8_0.data[a_offset + global_ib].qs[iqs * 2], k_packed_q8_0.data[a_offset + global_ib].qs[iqs * 2 + 1])); break; @@ -96,11 +96,11 @@ void k_block_to_shmem(const uint buf_ib, const uint global_ib, const uint iqs, c if (iqs == 0) { // Q4_0/Q5_0/Q8_0 store dm.x = d; Q4_1/Q5_1 store dm = (d, m) pair. switch (FaTypeK) { - case GGML_TYPE_Q4_0: kblocksh[buf_ib].dm = FLOAT_TYPEV2(FLOAT_TYPE(k_packed_q4_0.data[a_offset + global_ib].d), 0.0); break; - case GGML_TYPE_Q4_1: kblocksh[buf_ib].dm = FLOAT_TYPEV2(k_packed_q4_1_p32.data[a_offset + global_ib].dm); break; - case GGML_TYPE_Q5_0: kblocksh[buf_ib].dm = FLOAT_TYPEV2(FLOAT_TYPE(k_packed_q5_0.data[a_offset + global_ib].d), 0.0); break; - case GGML_TYPE_Q5_1: kblocksh[buf_ib].dm = FLOAT_TYPEV2(k_packed_q5_1_p32.data[a_offset + global_ib].dm); break; - case GGML_TYPE_Q8_0: kblocksh[buf_ib].dm = FLOAT_TYPEV2(FLOAT_TYPE(k_packed_q8_0.data[a_offset + global_ib].d), 0.0); break; + case FA_TYPE_Q4_0: kblocksh[buf_ib].dm = FLOAT_TYPEV2(FLOAT_TYPE(k_packed_q4_0.data[a_offset + global_ib].d), 0.0); break; + case FA_TYPE_Q4_1: kblocksh[buf_ib].dm = FLOAT_TYPEV2(k_packed_q4_1_p32.data[a_offset + global_ib].dm); break; + case FA_TYPE_Q5_0: kblocksh[buf_ib].dm = FLOAT_TYPEV2(FLOAT_TYPE(k_packed_q5_0.data[a_offset + global_ib].d), 0.0); break; + case FA_TYPE_Q5_1: kblocksh[buf_ib].dm = FLOAT_TYPEV2(k_packed_q5_1_p32.data[a_offset + global_ib].dm); break; + case FA_TYPE_Q8_0: kblocksh[buf_ib].dm = FLOAT_TYPEV2(FLOAT_TYPE(k_packed_q8_0.data[a_offset + global_ib].d), 0.0); break; } } } @@ -121,31 +121,31 @@ struct fa_k_qs_block8 { fa_k_qs_block8 get_k_qs_block8(uint ib, uint a_offset) { fa_k_qs_block8 r; uint qh = 0; - if (FaTypeK == GGML_TYPE_Q5_0) { + if (FaTypeK == FA_TYPE_Q5_0) { qh = pack32(u16vec2(k_packed_q5_0.data[a_offset + ib].qh[0], k_packed_q5_0.data[a_offset + ib].qh[1])); - } else if (FaTypeK == GGML_TYPE_Q5_1) { + } else if (FaTypeK == FA_TYPE_Q5_1) { qh = k_packed_q5_1.data[a_offset + ib].qh; } - const bool has_qh = (FaTypeK == GGML_TYPE_Q5_0) || (FaTypeK == GGML_TYPE_Q5_1); + const bool has_qh = (FaTypeK == FA_TYPE_Q5_0) || (FaTypeK == FA_TYPE_Q5_1); [[unroll]] for (uint32_t d = 0; d < 4; d++) { uint vui = 0; switch (FaTypeK) { - case GGML_TYPE_Q4_0: { // packed16 + case FA_TYPE_Q4_0: { // packed16 vui = pack32(u16vec2(k_packed_q4_0.data[a_offset + ib].qs[d * 2 + 0], k_packed_q4_0.data[a_offset + ib].qs[d * 2 + 1])); break; } - case GGML_TYPE_Q4_1: { // packed32 alias + case FA_TYPE_Q4_1: { // packed32 alias vui = k_packed_q4_1_p32.data[a_offset + ib].qs[d]; break; } - case GGML_TYPE_Q5_0: { // packed16 + case FA_TYPE_Q5_0: { // packed16 vui = pack32(u16vec2(k_packed_q5_0.data[a_offset + ib].qs[d * 2 + 0], k_packed_q5_0.data[a_offset + ib].qs[d * 2 + 1])); break; } - case GGML_TYPE_Q5_1: { // packed32 alias + case FA_TYPE_Q5_1: { // packed32 alias vui = k_packed_q5_1_p32.data[a_offset + ib].qs[d]; break; } @@ -164,21 +164,21 @@ fa_k_qs_block8 get_k_qs_block8(uint ib, uint a_offset) { int32_t get_k_qs_shmem(const uint buf_ib, const uint pos) { switch (FaTypeK) { - case GGML_TYPE_Q4_0: - case GGML_TYPE_Q4_1: { + case FA_TYPE_Q4_0: + case FA_TYPE_Q4_1: { uint sub = pos % 4; uint shift = ((pos % 8) >= 4) ? 4u : 0u; return int32_t((uint(kblocksh[buf_ib].qs[sub]) >> shift) & 0x0F0F0F0Fu); } - case GGML_TYPE_Q5_0: - case GGML_TYPE_Q5_1: { + case FA_TYPE_Q5_0: + case FA_TYPE_Q5_1: { uint sub = pos % 4; uint shift = ((pos % 8) >= 4) ? 4u : 0u; int32_t result = int32_t((uint(kblocksh[buf_ib].qs[sub]) >> shift) & 0x0F0F0F0Fu); uint qh_bits = (kblocksh[buf_ib].qh >> (pos * 4u)) & 0xFu; return result | int32_t((qh_bits * 0x02040810u) & 0x10101010u); } - case GGML_TYPE_Q8_0: { + case FA_TYPE_Q8_0: { return kblocksh[buf_ib].qs[pos]; } default: return 0; @@ -187,10 +187,10 @@ int32_t get_k_qs_shmem(const uint buf_ib, const uint pos) { ACC_TYPE k_dot_correction(const uint qib, const ACC_TYPEV2 k_dm) { switch (FaTypeK) { - case GGML_TYPE_Q4_0: return -ACC_TYPE(8.0) * ACC_TYPE(Qf[qib].ds.y) * k_dm.x; - case GGML_TYPE_Q5_0: return -ACC_TYPE(16.0) * ACC_TYPE(Qf[qib].ds.y) * k_dm.x; - case GGML_TYPE_Q4_1: - case GGML_TYPE_Q5_1: return ACC_TYPE(Qf[qib].ds.y) * k_dm.y; + case FA_TYPE_Q4_0: return -ACC_TYPE(8.0) * ACC_TYPE(Qf[qib].ds.y) * k_dm.x; + case FA_TYPE_Q5_0: return -ACC_TYPE(16.0) * ACC_TYPE(Qf[qib].ds.y) * k_dm.x; + case FA_TYPE_Q4_1: + case FA_TYPE_Q5_1: return ACC_TYPE(Qf[qib].ds.y) * k_dm.y; default: return ACC_TYPE(0.0); } } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_sparse_compact.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_sparse_compact.comp deleted file mode 100644 index 3d31362661b3..000000000000 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_sparse_compact.comp +++ /dev/null @@ -1,102 +0,0 @@ -#version 450 - -#extension GL_EXT_control_flow_attributes : enable -#extension GL_EXT_shader_16bit_storage : require -#extension GL_EXT_shader_explicit_arithmetic_types_int32 : require -#ifdef USE_SUBGROUPS -#extension GL_KHR_shader_subgroup_basic : require -#extension GL_KHR_shader_subgroup_ballot : require -#endif - -layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; -layout(constant_id = 0) const uint BLOCK_SIZE = 128; -layout(constant_id = 1) const uint NUM_SUBGROUPS = 1; - -layout (binding = 0) readonly buffer M {float16_t data_m[];}; -layout (binding = 1) writeonly buffer I {int32_t data_i[];}; - -layout (push_constant) uniform parameter { - uint KV; - uint nem1; - uint nem2; - uint nbm1; - uint nbm2; - uint nbm3; - uint n_kv_max; -} p; - -#ifdef USE_SUBGROUPS -shared uvec4 ballots_sh[NUM_SUBGROUPS]; -#else -shared uint scan[BLOCK_SIZE]; -#endif - -// One workgroup per mask row: compact the finite-mask KV positions into a -// per-row index list of length n_kv_max, -1 padded. Emitted in ascending KV -// order so the downstream attention accumulation is deterministic. -void main() { - const uint i1 = gl_WorkGroupID.x; - const uint i2 = gl_WorkGroupID.y; - const uint i3 = gl_WorkGroupID.z; - const uint tid = gl_LocalInvocationIndex; - - const uint m_base = i3 * p.nbm3 + i2 * p.nbm2 + i1 * p.nbm1; - const uint out_base = ((i3 * p.nem2 + i2) * p.nem1 + i1) * p.n_kv_max; - - uint base = 0; - for (uint chunk = 0; chunk < p.KV; chunk += BLOCK_SIZE) { - const uint k = chunk + tid; - bool selected = false; - if (k < p.KV) { - const float v = float(data_m[m_base + k]); - selected = !isinf(v) && !isnan(v); - } - -#ifdef USE_SUBGROUPS - const uvec4 ballot = subgroupBallot(selected); - if (subgroupElect()) { - ballots_sh[gl_SubgroupID] = ballot; - } - barrier(); - - uint subgroup_base = 0; - uint total = 0; - [[unroll]] for (uint s = 0; s < gl_NumSubgroups; ++s) { - if (s == gl_SubgroupID) { - subgroup_base = total; - } - total += subgroupBallotBitCount(ballots_sh[s]); - } - barrier(); - - const uint slot = base + subgroup_base + subgroupBallotExclusiveBitCount(ballot); -#else - // Hillis-Steele inclusive prefix sum over the workgroup. - scan[tid] = selected ? 1u : 0u; - barrier(); - for (uint off = 1; off < BLOCK_SIZE; off <<= 1) { - uint add = 0; - if (tid >= off) { - add = scan[tid - off]; - } - barrier(); - scan[tid] += add; - barrier(); - } - - const uint inclusive = scan[tid]; - const uint total = scan[BLOCK_SIZE - 1]; - const uint slot = base + inclusive - 1u; -#endif - - if (selected && slot < p.n_kv_max) { - data_i[out_base + slot] = int32_t(k); - } - base += total; - barrier(); - } - - for (uint s = min(base, p.n_kv_max) + tid; s < p.n_kv_max; s += BLOCK_SIZE) { - data_i[out_base + s] = int32_t(-1); - } -} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_split_k_reduce.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_split_k_reduce.comp index 68917fc0bb02..69342ba3ceea 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_split_k_reduce.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_split_k_reduce.comp @@ -14,6 +14,7 @@ layout (push_constant) uniform parameter { uint D; uint ne1; uint ne2; + uint dst_ne2; uint ne3; uint k_num; uint sinks; @@ -116,6 +117,6 @@ void main() { const float FLT_MAX = uintBitsToFloat(0x7F7FFFFF); O = clamp(O, -FLT_MAX, FLT_MAX); - data_d[(i3 * p.ne2 + i2) * p.ne1 * D + D * n + d] = O; + data_d[(i3 * p.dst_ne2 + i2) * p.ne1 * D + D * n + d] = O; } } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k.comp new file mode 100644 index 000000000000..930777b5f0ea --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k.comp @@ -0,0 +1,155 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_KHR_shader_subgroup_arithmetic : require +#extension GL_KHR_shader_subgroup_basic : require + +layout(constant_id = 0) const uint WORKGROUP_SIZE = 512; +layout(constant_id = 1) const uint SUBGROUP_SIZE = 64; +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout(binding = 0) readonly buffer QBuf { float data_q[]; }; +layout(binding = 1) readonly buffer KBuf { float16_t data_k[]; }; +layout(binding = 2) readonly buffer MaskBuf { float16_t data_m[]; }; +layout(binding = 3) readonly buffer SinkBuf { float data_s[]; }; +layout(binding = 4) readonly buffer TopBuf { int data_top[]; }; +layout(binding = 5) writeonly buffer DstBuf { float data_dst[]; }; + +layout(push_constant) uniform Parameters { + uint n_batch; + uint n_kv; + uint n_kv_raw; + uint n_top_k; + uint n_head; + uint nbq1; + uint nbq2; + uint nbq3; + uint nbk1; + uint nbk3; + uint nbm1; + uint nbm3; + uint nbt1; + uint nbt3; + uint nb1; + uint nb2; + uint nb3; + float scale; + uint has_sinks; + uint split_mode; +} p; + +// Shape constants pinned by the dispatch gate in ggml_vk_flash_attn_top_k: DeepSeek V4 +// CSA attention only (hd 512, 64 heads MQA over an f16 K==V latent). The gate also pins +// the pipeline to subgroup size 64, so per-lane sizing uses LANES rather than the +// SUBGROUP_SIZE specialization constant (a spec-constant array bound would be folded at +// compile time from the default and silently break under a different runtime subgroup). +const uint HEAD_SIZE = 512; +const uint HEADS_PER_GROUP = 8; +const uint KEYS_PER_BLOCK = 16; +const uint LANES = 64; +// f16 -inf (or the lowest f16 normal some mask writers use in its place) marks a +// masked-out key. +const float MASK_NEG_INF = -65500.0; + +shared float16_t key_sh[KEYS_PER_BLOCK * HEAD_SIZE]; +shared uint key_idx[KEYS_PER_BLOCK]; + +void main() { + const uint tid = gl_LocalInvocationIndex; + const uint lane = gl_SubgroupInvocationID; + const uint head = gl_WorkGroupID.y * HEADS_PER_GROUP + gl_SubgroupID; + const uint token = gl_WorkGroupID.x; + const uint stream = gl_WorkGroupID.z; + + // No bounds check: the grid is sized exactly (x = n_batch, y * HEADS_PER_GROUP covers + // n_head == 64). An early return here would skip the barriers below for part of the + // workgroup, which is undefined behavior — do not reintroduce one without restructuring + // the barrier flow. + + float accum[HEAD_SIZE / LANES]; + [[unroll]] for (uint i = 0; i < HEAD_SIZE / LANES; ++i) { + accum[i] = 0.0; + } + + float row_max = uintBitsToFloat(0xff800000); + float row_sum = 0.0; + const uint q_base = stream * p.nbq3 + head * p.nbq2 + token * p.nbq1; + const uint mask_base = stream * p.nbm3 + token * p.nbm1; + const uint top_base = stream * p.nbt3 + token * p.nbt1; + const uint total_keys = p.n_kv_raw + p.n_top_k; + + for (uint kb = 0; kb < total_keys; kb += KEYS_PER_BLOCK) { + if (tid < KEYS_PER_BLOCK) { + const uint selected = kb + tid; + uint key = p.n_kv; + if (selected < p.n_kv_raw) { + key = selected; + } else if (selected < total_keys) { + const int compressed = data_top[top_base + selected - p.n_kv_raw]; + if (compressed >= 0 && uint(compressed) < p.n_kv - p.n_kv_raw) { + key = p.n_kv_raw + uint(compressed); + } + } + key_idx[tid] = key; + } + barrier(); + + for (uint idx = tid; idx < KEYS_PER_BLOCK * HEAD_SIZE; idx += WORKGROUP_SIZE) { + const uint col = idx / HEAD_SIZE; + const uint dim = idx % HEAD_SIZE; + const uint key = key_idx[col]; + key_sh[idx] = key < p.n_kv ? data_k[stream * p.nbk3 + key * p.nbk1 + dim] : float16_t(0.0); + } + barrier(); + + [[unroll]] for (uint col = 0; col < KEYS_PER_BLOCK; ++col) { + const uint selected = kb + col; + const uint key = key_idx[col]; + if (selected >= total_keys || key >= p.n_kv) { + continue; + } + + float partial = 0.0; + [[unroll]] for (uint i = 0; i < HEAD_SIZE / LANES; ++i) { + const uint dim = lane + i * LANES; + partial += data_q[q_base + dim] * float(key_sh[col * HEAD_SIZE + dim]); + } + const float mask = float(data_m[mask_base + key]); + const float score = subgroupAdd(partial) * p.scale + mask; + if (mask < MASK_NEG_INF) { + continue; + } + + const float new_max = max(row_max, score); + const float old_scale = row_sum == 0.0 ? 0.0 : exp(row_max - new_max); + const float value_scale = exp(score - new_max); + row_sum = row_sum * old_scale + value_scale; + row_max = new_max; + + [[unroll]] for (uint i = 0; i < HEAD_SIZE / LANES; ++i) { + const uint dim = lane + i * LANES; + accum[i] = accum[i] * old_scale + value_scale * float(key_sh[col * HEAD_SIZE + dim]); + } + } + barrier(); + } + + if (p.has_sinks != 0) { + const float sink = data_s[head]; + const float new_max = max(row_max, sink); + const float old_scale = row_sum == 0.0 ? 0.0 : exp(row_max - new_max); + const float sink_scale = exp(sink - new_max); + row_sum = row_sum * old_scale + sink_scale; + [[unroll]] for (uint i = 0; i < HEAD_SIZE / LANES; ++i) { + accum[i] *= old_scale; + } + } + + const uint dst_base = stream * p.nb3 + token * p.nb2 + head * p.nb1; + const float inv_sum = row_sum == 0.0 ? 0.0 : 1.0 / row_sum; + [[unroll]] for (uint i = 0; i < HEAD_SIZE / LANES; ++i) { + data_dst[dst_base + lane + i * LANES] = accum[i] * inv_sum; + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k_cm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k_cm.comp new file mode 100644 index 000000000000..7a328938777a --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_top_k_cm.comp @@ -0,0 +1,296 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_KHR_shader_subgroup_basic : require +#extension GL_KHR_shader_subgroup_shuffle : require + +layout(constant_id = 0) const uint WORKGROUP_SIZE = 512; +layout(constant_id = 1) const uint SUBGROUP_SIZE = 64; +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout(binding = 0) readonly buffer QBuf { float data_q[]; }; +layout(binding = 1) readonly buffer KBuf { float16_t data_k[]; }; +layout(binding = 2) readonly buffer MaskBuf { float16_t data_m[]; }; +layout(binding = 3) readonly buffer SinkBuf { float data_s[]; }; +layout(binding = 4) readonly buffer TopBuf { int data_top[]; }; +layout(binding = 5) writeonly buffer DstBuf { float data_dst[]; }; + +layout(push_constant) uniform Parameters { + uint n_batch; + uint n_kv; + uint n_kv_raw; + uint n_top_k; + uint n_head; + uint nbq1; + uint nbq2; + uint nbq3; + uint nbk1; + uint nbk3; + uint nbm1; + uint nbm3; + uint nbt1; + uint nbt3; + uint nb1; + uint nb2; + uint nb3; + float scale; + uint has_sinks; + uint split_mode; +} p; + +// This shader is hard-wired to a 512-thread workgroup of eight 64-wide subgroups: +// the softmax segments SUBGROUP_SIZE/4 = 16 lanes per head (so head_local stays < 32), +// and the QK/PV tiling assumes gl_NumSubgroups == 8 via gl_SubgroupID % 4 / gl_SubgroupID / 4. +// A different subgroup size indexes past row_max_sh/q_p_sh and corrupts the score tiles, so +// the pipeline is only created under `device->subgroup_size == 64` in ggml_vk_load_shaders. +// Keep that gate in sync with these constants. +const uint TILE = 16; +const uint HEAD_SIZE = 512; +const uint HEADS_PER_GROUP = 32; +const uint KEYS_PER_BLOCK = 64; +const uint DIMS_PER_BLOCK = 64; +const uint QK_STRIDE = TILE / 4 + 2; +const uint SCORE_STRIDE = HEADS_PER_GROUP / 4 + 1; +const uint P_STRIDE = KEYS_PER_BLOCK / 4 + 2; +const uint V_STRIDE = DIMS_PER_BLOCK / 4 + 2; +const uint PV_STRIDE = DIMS_PER_BLOCK / 4; +const float MASK_NEG_INF = -65500.0; + +shared uint key_idx[KEYS_PER_BLOCK]; +shared float key_mask[KEYS_PER_BLOCK]; +shared f16vec4 q_p_sh[HEADS_PER_GROUP * P_STRIDE]; +shared f16vec4 k_v_sh[KEYS_PER_BLOCK * V_STRIDE]; +shared vec4 matrix_sh[KEYS_PER_BLOCK * SCORE_STRIDE]; +shared float old_scale_sh[HEADS_PER_GROUP]; +shared float row_max_sh[HEADS_PER_GROUP]; +shared float row_sum_sh[HEADS_PER_GROUP]; + +void main() { + const uint tid = gl_LocalInvocationIndex; + const uint token = gl_WorkGroupID.x; + const uint head_base = gl_WorkGroupID.y * HEADS_PER_GROUP; + const uint stream = gl_WorkGroupID.z; + const uint mask_base = stream * p.nbm3 + token * p.nbm1; + const uint top_base = stream * p.nbt3 + token * p.nbt1; + const uint total_keys = p.split_mode != 0 ? p.n_top_k : p.n_kv_raw + p.n_top_k; + + float accum[HEADS_PER_GROUP * HEAD_SIZE / WORKGROUP_SIZE]; + [[unroll]] for (uint i = 0; i < accum.length(); ++i) { + accum[i] = 0.0; + } + + if (tid < HEADS_PER_GROUP) { + row_max_sh[tid] = uintBitsToFloat(0xff800000); + row_sum_sh[tid] = 0.0; + } + barrier(); + + for (uint kb = 0; kb < total_keys; kb += KEYS_PER_BLOCK) { + if (tid < KEYS_PER_BLOCK) { + const uint selected = kb + tid; + uint key = p.n_kv; + if (p.split_mode == 0 && selected < p.n_kv_raw) { + key = selected; + } else if (selected < total_keys) { + const uint top_pos = p.split_mode != 0 ? selected : selected - p.n_kv_raw; + const int compressed = data_top[top_base + top_pos]; + if (compressed >= 0 && uint(compressed) < p.n_kv - p.n_kv_raw) { + key = p.n_kv_raw + uint(compressed); + } + } + key_idx[tid] = key; + key_mask[tid] = key < p.n_kv ? float(data_m[mask_base + key]) : uintBitsToFloat(0xff800000); + } + barrier(); + + coopmat scores = + coopmat(0.0); + coopmat kmat; + coopmat qmat; + + [[unroll]] for (uint d = 0; d < HEAD_SIZE; d += TILE) { + if (tid < KEYS_PER_BLOCK * (TILE / 4)) { + const uint key_local = tid / (TILE / 4); + const uint d4 = tid % (TILE / 4); + const uint key = key_idx[key_local]; + f16vec4 value = f16vec4(0.0); + if (key < p.n_kv) { + const uint offset = stream * p.nbk3 + key * p.nbk1 + d + d4 * 4; + value = f16vec4(data_k[offset], data_k[offset + 1], data_k[offset + 2], data_k[offset + 3]); + } + k_v_sh[key_local * QK_STRIDE + d4] = value; + } + if (tid < HEADS_PER_GROUP * (TILE / 4)) { + const uint head_local = tid / (TILE / 4); + const uint d4 = tid % (TILE / 4); + const uint head = head_base + head_local; + const uint offset = stream * p.nbq3 + head * p.nbq2 + token * p.nbq1 + d + d4 * 4; + q_p_sh[head_local * QK_STRIDE + d4] = f16vec4( + data_q[offset], data_q[offset + 1], data_q[offset + 2], data_q[offset + 3]); + } + barrier(); + + const uint key_chunk = gl_SubgroupID % (KEYS_PER_BLOCK / TILE); + const uint head_tile = gl_SubgroupID / (KEYS_PER_BLOCK / TILE); + coopMatLoad(kmat, k_v_sh, key_chunk * TILE * QK_STRIDE, + QK_STRIDE, gl_CooperativeMatrixLayoutRowMajor); + coopMatLoad(qmat, q_p_sh, head_tile * TILE * QK_STRIDE, + QK_STRIDE, gl_CooperativeMatrixLayoutColumnMajor); + scores = coopMatMulAdd(kmat, qmat, scores); + barrier(); + } + + const uint score_key_chunk = gl_SubgroupID % (KEYS_PER_BLOCK / TILE); + const uint score_head_tile = gl_SubgroupID / (KEYS_PER_BLOCK / TILE); + coopMatStore(scores, matrix_sh, + score_key_chunk * TILE * SCORE_STRIDE + score_head_tile * (TILE / 4), + SCORE_STRIDE, gl_CooperativeMatrixLayoutRowMajor); + barrier(); + + { + const uint head_local = tid / (SUBGROUP_SIZE / 4); + const uint softmax_lane = tid % (SUBGROUP_SIZE / 4); + float block_max = uintBitsToFloat(0xff800000); + [[unroll]] for (uint i = 0; i < KEYS_PER_BLOCK / (SUBGROUP_SIZE / 4); ++i) { + const uint key_local = softmax_lane + i * (SUBGROUP_SIZE / 4); + const float mask = key_mask[key_local]; + const float score = float(matrix_sh[key_local * SCORE_STRIDE + head_local / 4][head_local % 4]) * p.scale + mask; + block_max = mask < MASK_NEG_INF ? block_max : max(block_max, score); + } + [[unroll]] for (uint delta = 1; delta < SUBGROUP_SIZE / 4; delta *= 2) { + block_max = max(block_max, subgroupShuffleXor(block_max, delta)); + } + + const float old_row_max = row_max_sh[head_local]; + const float old_row_sum = row_sum_sh[head_local]; + const float new_max = max(old_row_max, block_max); + const float old_scale = old_row_sum == 0.0 ? 0.0 : exp(old_row_max - new_max); + float block_sum = 0.0; + [[unroll]] for (uint i = 0; i < KEYS_PER_BLOCK / (SUBGROUP_SIZE / 4); ++i) { + const uint key_local = softmax_lane + i * (SUBGROUP_SIZE / 4); + const float mask = key_mask[key_local]; + float weight = 0.0; + if (mask >= MASK_NEG_INF) { + const float score = float(matrix_sh[key_local * SCORE_STRIDE + head_local / 4][head_local % 4]) * p.scale + mask; + weight = exp(score - new_max); + block_sum += weight; + } + q_p_sh[head_local * P_STRIDE + key_local / 4][key_local % 4] = float16_t(weight); + } + [[unroll]] for (uint delta = 1; delta < SUBGROUP_SIZE / 4; delta *= 2) { + block_sum += subgroupShuffleXor(block_sum, delta); + } + + if (softmax_lane == 0) { + row_sum_sh[head_local] = old_row_sum * old_scale + block_sum; + row_max_sh[head_local] = new_max; + old_scale_sh[head_local] = old_scale; + } + } + barrier(); + + [[unroll]] for (uint i = 0; i < accum.length(); ++i) { + const uint out_idx = tid + i * WORKGROUP_SIZE; + const uint head_local = out_idx / HEAD_SIZE; + accum[i] *= old_scale_sh[head_local]; + } + + coopmat pmats[KEYS_PER_BLOCK / TILE]; + [[unroll]] for (uint key_chunk = 0; key_chunk < KEYS_PER_BLOCK / TILE; ++key_chunk) { + const uint pv_head_tile = gl_SubgroupID / (DIMS_PER_BLOCK / TILE); + coopMatLoad(pmats[key_chunk], q_p_sh, + pv_head_tile * TILE * P_STRIDE + key_chunk * (TILE / 4), + P_STRIDE, gl_CooperativeMatrixLayoutRowMajor); + } + + [[unroll]] for (uint dim_base = 0; dim_base < HEAD_SIZE; dim_base += DIMS_PER_BLOCK) { + [[unroll]] for (uint idx = tid; idx < KEYS_PER_BLOCK * (DIMS_PER_BLOCK / 4); idx += WORKGROUP_SIZE) { + const uint key_local = idx / (DIMS_PER_BLOCK / 4); + const uint d4 = idx % (DIMS_PER_BLOCK / 4); + const uint key = key_idx[key_local]; + f16vec4 value = f16vec4(0.0); + if (key < p.n_kv) { + const uint offset = stream * p.nbk3 + key * p.nbk1 + dim_base + d4 * 4; + value = f16vec4(data_k[offset], data_k[offset + 1], data_k[offset + 2], data_k[offset + 3]); + } + k_v_sh[key_local * V_STRIDE + d4] = value; + } + barrier(); + + coopmat pv = + coopmat(0.0); + coopmat vmat; + + const uint pv_head_tile = gl_SubgroupID / (DIMS_PER_BLOCK / TILE); + const uint pv_dim_tile = gl_SubgroupID % (DIMS_PER_BLOCK / TILE); + [[unroll]] for (uint key_chunk = 0; key_chunk < KEYS_PER_BLOCK / TILE; ++key_chunk) { + coopMatLoad(vmat, k_v_sh, + key_chunk * TILE * V_STRIDE + pv_dim_tile * (TILE / 4), + V_STRIDE, gl_CooperativeMatrixLayoutRowMajor); + pv = coopMatMulAdd(pmats[key_chunk], vmat, pv); + } + + coopMatStore(pv, matrix_sh, + pv_head_tile * TILE * PV_STRIDE + pv_dim_tile * (TILE / 4), PV_STRIDE, + gl_CooperativeMatrixLayoutRowMajor); + barrier(); + + [[unroll]] for (uint i = 0; i < accum.length(); ++i) { + const uint out_idx = tid + i * WORKGROUP_SIZE; + const uint head_local = out_idx / HEAD_SIZE; + const uint dim = out_idx % HEAD_SIZE; + if (dim >= dim_base && dim < dim_base + DIMS_PER_BLOCK) { + accum[i] += matrix_sh[head_local * PV_STRIDE + (dim - dim_base) / 4][dim % 4]; + } + } + barrier(); + } + } + + if (p.split_mode == 0 && p.has_sinks != 0 && tid < HEADS_PER_GROUP) { + const float sink = data_s[head_base + tid]; + const float new_max = max(row_max_sh[tid], sink); + const float old_scale = row_sum_sh[tid] == 0.0 ? 0.0 : exp(row_max_sh[tid] - new_max); + row_sum_sh[tid] = row_sum_sh[tid] * old_scale + exp(sink - new_max); + row_max_sh[tid] = new_max; + old_scale_sh[tid] = old_scale; + } + barrier(); + + if (p.split_mode == 0 && p.has_sinks != 0) { + [[unroll]] for (uint i = 0; i < accum.length(); ++i) { + const uint out_idx = tid + i * WORKGROUP_SIZE; + accum[i] *= old_scale_sh[out_idx / HEAD_SIZE]; + } + } + [[unroll]] for (uint i = 0; i < accum.length(); ++i) { + const uint out_idx = tid + i * WORKGROUP_SIZE; + const uint head_local = out_idx / HEAD_SIZE; + const uint dim = out_idx % HEAD_SIZE; + if (p.split_mode != 0) { + const uint part_idx = 1; + // Split-K buffer layout, shared with flash_attn_cm1.comp and + // flash_attn_split_k_reduce.comp: all O matrices [HSV, ne1, k_num, ne2, ne3] + // first, then the L/M rows [ne1, k_num, ne2, ne3]. Here ne1 = n_head, + // ne2 = n_batch and ne3 = the stream count, so the L/M base must span every + // stream -- omitting gl_NumWorkGroups.z lands L/M inside the O region as soon + // as there is more than one sequence. + const uint n_streams = gl_NumWorkGroups.z; + const uint matrix_base = HEAD_SIZE * p.n_head * (part_idx + 2 * (token + p.n_batch * stream)); + const uint lm_base = HEAD_SIZE * p.n_head * p.n_batch * n_streams * 2 + p.n_head * 2 * (part_idx + 2 * (token + p.n_batch * stream)); + data_dst[matrix_base + (head_base + head_local) * HEAD_SIZE + dim] = accum[i]; + if (dim == 0) { + data_dst[lm_base + head_base + head_local] = row_sum_sh[head_local]; + data_dst[lm_base + p.n_head + head_base + head_local] = row_max_sh[head_local]; + } + } else { + const uint dst_base = stream * p.nb3 + token * p.nb2 + (head_base + head_local) * p.nb1; + const float inv_sum = row_sum_sh[head_local] == 0.0 ? 0.0 : 1.0 / row_sum_sh[head_local]; + data_dst[dst_base + dim] = accum[i] * inv_sum; + } + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_union.comp b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_union.comp new file mode 100644 index 000000000000..513375709300 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/flash_attn_union.comp @@ -0,0 +1,134 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require +#extension GL_KHR_shader_subgroup_arithmetic : require +#extension GL_KHR_shader_subgroup_basic : require + +// Builds the deduplicated union of the compressed rows selected by any of the n_batch query +// tokens, for the DeepSeek V4 small-batch decode path. +// +// Marks a bitmap from the top-k index lists, then compacts by scanning bitmap WORDS. The +// marking is O(n_batch * n_top_k), independent of context depth, and the compaction touches +// R/32 words instead of R rows. An earlier version scanned the mask row by row instead: that +// is simpler, but it made dedup cost scale with depth and measured 0.71x (ie a loss) at +// kv=133888, which defeats the purpose. Do not go back to it. +// +// Ascending source order falls out of the bitmap scan, so the result is deterministic. +// +// Emits the index list plus the PADDED compact row count. Padding to a multiple of pad_to +// (a multiple of every FA block width) is what lets flash-attention keep its "aligned" +// pipeline variant even though the row count is now a runtime value. +// +// Also emits the raw union size, the candidate count and the batch it came from. The host +// reads those back to decide whether compaction is worth it at all, and files the sample under +// the batch reported here rather than the one it is pricing: it cannot tell which call last +// wrote the slot, and the overlap it is measuring depends on the batch. +// +// count_only skips the list writes. The scan still has to run to produce the count, but with +// nothing to write the index buffer need not exist, which is what lets the host price the +// union on a call it is not going to compact. +// +// One workgroup: the bitmap and the running offset both live in shared memory. + +layout(local_size_x = 256, local_size_y = 1, local_size_z = 1) in; + +layout(binding = 0) readonly buffer TBuf { int data_top[]; }; +layout(binding = 1) writeonly buffer UBuf { uint data_u[]; }; +layout(binding = 2) writeonly buffer CBuf { uint data_c[]; }; + +layout(push_constant) uniform Parameters { + uint n_kv; + uint n_kv_raw; + uint n_batch; + uint n_top_k; + uint max_union; + uint nbt1; + uint max_words; // capacity of the shared bitmap, host-checked + uint pad_to; + uint count_only; // 1: produce the count, write no index list +} p; + +// 12288 words = 393216 compressed rows, ~48 KiB of shared memory. +const uint MAX_WORDS = 12288; + +shared uint bitmap[MAX_WORDS]; +shared uint wave_totals[16]; +shared uint base_sh; + +void main() { + const uint tid = gl_LocalInvocationIndex; + const uint lane = gl_SubgroupInvocationID; + const uint wave = gl_SubgroupID; + const uint nwave = gl_NumSubgroups; + const uint R = p.n_kv - p.n_kv_raw; + const uint words = min((R + 31) / 32, p.max_words); + + // phase 1: clear + for (uint w = tid; w < words; w += gl_WorkGroupSize.x) { + bitmap[w] = 0; + } + if (tid == 0) { + base_sh = 0; + } + barrier(); + + // phase 2: mark. Depth-independent: one pass over the top-k lists. + const uint n_cand = p.n_batch * p.n_top_k; + for (uint c = tid; c < n_cand; c += gl_WorkGroupSize.x) { + const uint t = c / p.n_top_k; + const uint j = c - t * p.n_top_k; + const int idx = data_top[t * p.nbt1 + j]; + if (idx >= 0 && uint(idx) < R) { + atomicOr(bitmap[uint(idx) >> 5], 1u << (uint(idx) & 31u)); + } + } + barrier(); + + // phase 3: compact. R/32 iterations, ascending, exact running offset in shared memory. + for (uint chunk = 0; chunk < words; chunk += gl_WorkGroupSize.x) { + const uint w = chunk + tid; + const uint bits = w < words ? bitmap[w] : 0u; + const uint cnt = bitCount(bits); + + const uint wave_off = subgroupExclusiveAdd(cnt); + const uint wave_tot = subgroupAdd(cnt); + if (lane == 0) { + wave_totals[wave] = wave_tot; + } + barrier(); + + uint prefix = 0; + for (uint i = 0; i < wave; ++i) { + prefix += wave_totals[i]; + } + uint total = 0; + for (uint i = 0; i < nwave; ++i) { + total += wave_totals[i]; + } + + uint slot = base_sh + prefix + wave_off; + uint rem = p.count_only != 0 ? 0u : bits; + while (rem != 0) { + const uint b = findLSB(rem); + rem &= rem - 1; + if (slot < p.max_union) { + data_u[slot] = w * 32 + b; + } + ++slot; + } + barrier(); + if (tid == 0) { + base_sh += total; + } + barrier(); + } + + if (tid == 0) { + const uint u = min(base_sh, p.max_union); + const uint rows = p.n_kv_raw + u; + data_c[0] = ((rows + p.pad_to - 1) / p.pad_to) * p.pad_to; + data_c[1] = u; + data_c[2] = n_cand; + data_c[3] = p.n_batch; + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/ggml_type_ids.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/ggml_type_ids.glsl deleted file mode 100644 index 0f10c733dd88..000000000000 --- a/ggml/src/ggml-vulkan/vulkan-shaders/ggml_type_ids.glsl +++ /dev/null @@ -1,34 +0,0 @@ -#if !defined(GGML_TYPE_IDS_COMP) -#define GGML_TYPE_IDS_COMP - -// ggml_type enum values — must match ggml.h -#define GGML_TYPE_F32 0u -#define GGML_TYPE_F16 1u -#define GGML_TYPE_Q4_0 2u -#define GGML_TYPE_Q4_1 3u -#define GGML_TYPE_Q5_0 6u -#define GGML_TYPE_Q5_1 7u -#define GGML_TYPE_Q8_0 8u -#define GGML_TYPE_Q2_K 10u -#define GGML_TYPE_Q3_K 11u -#define GGML_TYPE_Q4_K 12u -#define GGML_TYPE_Q5_K 13u -#define GGML_TYPE_Q6_K 14u -#define GGML_TYPE_IQ2_XXS 16u -#define GGML_TYPE_IQ2_XS 17u -#define GGML_TYPE_IQ3_XXS 18u -#define GGML_TYPE_IQ1_S 19u -#define GGML_TYPE_IQ4_NL 20u -#define GGML_TYPE_IQ3_S 21u -#define GGML_TYPE_IQ2_S 22u -#define GGML_TYPE_IQ4_XS 23u -#define GGML_TYPE_IQ1_M 29u -#define GGML_TYPE_BF16 30u -#define GGML_TYPE_TQ1_0 34u -#define GGML_TYPE_TQ2_0 35u -#define GGML_TYPE_MXFP4 39u -#define GGML_TYPE_NVFP4 40u -#define GGML_TYPE_Q1_0 41u -#define GGML_TYPE_Q2_0 42u - -#endif // !defined(GGML_TYPE_IDS_COMP) diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/iq_shmem_init.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/iq_shmem_init.glsl deleted file mode 100644 index 12e50ee9eb9d..000000000000 --- a/ggml/src/ggml-vulkan/vulkan-shaders/iq_shmem_init.glsl +++ /dev/null @@ -1,2 +0,0 @@ -void init_iq_shmem(uvec3 wgsize) { -} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer.comp b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer.comp index 9b34d8366220..ba76ec72ca6c 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer.comp @@ -12,9 +12,9 @@ #include "types.glsl" #include "fa_types.glsl" -#define FaTypeV GGML_TYPE_F32 +#define FaTypeV FA_TYPE_F32 -layout(constant_id = 0) const uint FaTypeK = GGML_TYPE_F32; +layout(constant_id = 0) const uint FaTypeK = FA_TYPE_F32; layout(constant_id = 1) const uint FaBlockBytesK = 4; layout(constant_id = 2) const uint SUBGROUP_SIZE = 32; @@ -84,11 +84,11 @@ void main() { const uint k_block_elems = fa_block_elems(FaTypeK); const uint k_elem_bytes = FaBlockBytesK / k_block_elems; - if (FaTypeK == GGML_TYPE_F16) { + if (FaTypeK == FA_TYPE_F16) { k_row[tid] = float(k_f16[k_offset / k_elem_bytes + tid]); - } else if (FaTypeK == GGML_TYPE_F32) { + } else if (FaTypeK == FA_TYPE_F32) { k_row[tid] = k_f32[k_offset / k_elem_bytes + tid]; - } else if (FaTypeK == GGML_TYPE_BF16) { + } else if (FaTypeK == FA_TYPE_BF16) { k_row[tid] = bf16_to_fp32(uint(k_bf16[k_offset / k_elem_bytes + tid])); } else if (4 * tid < HEAD_SIZE) { const uint coord = 4 * tid; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_cm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_cm.comp new file mode 100644 index 000000000000..a2004bd93d7f --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_cm.comp @@ -0,0 +1,168 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require +#extension GL_KHR_shader_subgroup_basic : require + +layout(constant_id = 0) const uint SUBGROUP_SIZE = 64; +// Head count is a SPECIALIZATION constant, not a push constant: the head loop's trip +// count must stay visible to the compiler. As a push constant it cannot unroll and +// decode costs ~55% more (measured, gfx1151, kv=131584). +layout(constant_id = 1) const uint N_HEAD = 64; +#if N_WAVES == 1 +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; +#else +layout(local_size_x = 512, local_size_y = 1, local_size_z = 1) in; +#endif + +layout(binding = 0) readonly buffer QBuf { float data_q[]; }; +layout(binding = 1) readonly buffer KBuf { float16_t data_k[]; }; +layout(binding = 2) readonly buffer WBuf { float data_w[]; }; +layout(binding = 3) readonly buffer MBuf { float16_t data_m[]; }; +layout(binding = 4) writeonly buffer DstBuf { float data_dst[]; }; + +layout(push_constant) uniform Parameters { + uint n_kv; + uint n_batch; + uint n_stream; + uint nem3; + uint nb1; + uint nb3; + uint nbq1; + uint nbq2; + uint nbq3; + uint nbk2; + uint nbk3; + uint nbw1; + uint nbw3; + uint nbm1; + uint nbm3; +} p; + +const uint TILE = 16; +const uint HEAD_SIZE = 128; +const uint VEC_PER_HEAD = HEAD_SIZE / 4; +const uint TILE_STRIDE = VEC_PER_HEAD + 2; +const uint SCORE_STRIDE = TILE / 4 + 1; + +// Shared-memory budget, N_WAVES=8 / HEADS_PER_TILE=4 (TILE 16, HEAD_SIZE 128): +// k_sh 8 * 16 * 34 * 8 B = 34816 +// q_sh 4 * 16 * 34 * 8 B = 17408 +// score_sh 8 * 16 * 5 * 16 B = 10240 +// weight_sh 4 * 16 * 4 B = 256 +// total = 62720 of 65536 (95.7%) +// Only one workgroup fits per CU at that size, which is the intended trade. Raising either +// constant overruns: HEADS_PER_TILE=8 needs 80128 B and the pipeline then fails to create, +// silently falling back to the small variant. The host gates on +// maxComputeSharedMemorySize >= 64 KiB; keep that gate and this arithmetic in sync. +shared f16vec4 q_sh[HEADS_PER_TILE][TILE * TILE_STRIDE]; +shared f16vec4 k_sh[N_WAVES][TILE * TILE_STRIDE]; +shared vec4 score_sh[N_WAVES][TILE * SCORE_STRIDE]; +shared float weight_sh[HEADS_PER_TILE][TILE]; + +void main() { + const uint tid = gl_LocalInvocationIndex; + const uint lane = gl_SubgroupInvocationID; + const uint wave = gl_SubgroupID; + const uint kv_base = (gl_WorkGroupID.x * N_WAVES + wave) * TILE; + const uint token_base = gl_WorkGroupID.y * TILE; + const uint stream = gl_WorkGroupID.z; + + float totals[4]; + [[unroll]] for (uint i = 0; i < 4; ++i) { + totals[i] = 0.0; + } + + for (uint idx = tid; idx < N_WAVES * TILE * VEC_PER_HEAD; idx += gl_WorkGroupSize.x) { + const uint load_wave = idx / (TILE * VEC_PER_HEAD); + const uint wave_idx = idx % (TILE * VEC_PER_HEAD); + const uint key = wave_idx / VEC_PER_HEAD; + const uint d4 = wave_idx % VEC_PER_HEAD; + const uint kv = (gl_WorkGroupID.x * N_WAVES + load_wave) * TILE + key; + f16vec4 value = f16vec4(0.0); + if (kv < p.n_kv) { + const uint offset = stream * p.nbk3 + kv * p.nbk2 + d4 * 4; + value = f16vec4(data_k[offset], data_k[offset + 1], data_k[offset + 2], data_k[offset + 3]); + } + k_sh[load_wave][key * TILE_STRIDE + d4] = value; + } + barrier(); + + // k_sh[wave] is written once above and never touched again, but coopMatLoad sits inside both + // the head_base and head_local loops, so the same 8 A fragments are re-read from shared + // memory N_HEAD times per wave (512 loads where 8 would do). Each wave64 fragment load moves + // 1 KiB of LDS traffic because the 16x16 f16 fragment is replicated 4x across the subgroup. + coopmat kmat_h[HEAD_SIZE / TILE]; + [[unroll]] for (uint d = 0; d < HEAD_SIZE; d += TILE) { + coopMatLoad(kmat_h[d / TILE], k_sh[wave], d / 4, TILE_STRIDE, gl_CooperativeMatrixLayoutRowMajor); + } + + for (uint head_base = 0; head_base < N_HEAD; head_base += HEADS_PER_TILE) { + for (uint idx = tid; idx < HEADS_PER_TILE * TILE * VEC_PER_HEAD; idx += gl_WorkGroupSize.x) { + const uint head_local = idx / (TILE * VEC_PER_HEAD); + const uint head_idx = idx % (TILE * VEC_PER_HEAD); + const uint token_local = head_idx / VEC_PER_HEAD; + const uint d4 = head_idx % VEC_PER_HEAD; + const uint token = token_base + token_local; + f16vec4 value = f16vec4(0.0); + if (token < p.n_batch && (N_HEAD % HEADS_PER_TILE == 0 || head_base + head_local < N_HEAD)) { + const uint offset = stream * p.nbq3 + token * p.nbq2 + (head_base + head_local) * p.nbq1 + d4 * 4; + value = f16vec4(data_q[offset], data_q[offset + 1], data_q[offset + 2], data_q[offset + 3]); + } + q_sh[head_local][token_local * TILE_STRIDE + d4] = value; + } + if (tid < HEADS_PER_TILE * TILE) { + const uint head_local = tid / TILE; + const uint token_local = tid % TILE; + const uint token = token_base + token_local; + weight_sh[head_local][token_local] = (token < p.n_batch && (N_HEAD % HEADS_PER_TILE == 0 || head_base + head_local < N_HEAD)) + ? data_w[stream * p.nbw3 + token * p.nbw1 + head_base + head_local] : 0.0; + } + barrier(); + + [[unroll]] for (uint head_local = 0; head_local < HEADS_PER_TILE; ++head_local) { + coopmat scores = + coopmat(0.0); + coopmat qmat; + + [[unroll]] for (uint d = 0; d < HEAD_SIZE; d += TILE) { + coopMatLoad(qmat, q_sh[head_local], d / 4, TILE_STRIDE, gl_CooperativeMatrixLayoutColumnMajor); + scores = coopMatMulAdd(kmat_h[d / TILE], qmat, scores); + } + + coopMatStore(scores, score_sh[wave], 0, SCORE_STRIDE, gl_CooperativeMatrixLayoutRowMajor); + controlBarrier(gl_ScopeSubgroup, gl_ScopeSubgroup, gl_StorageSemanticsShared, gl_SemanticsAcquireRelease); + + [[unroll]] for (uint i = 0; i < 4; ++i) { + const uint idx = lane + i * SUBGROUP_SIZE; + const uint key = idx / TILE; + const uint token_local = idx % TILE; + const uint token = token_base + token_local; + if (token < p.n_batch && kv_base + key < p.n_kv) { + const float score = score_sh[wave][key * SCORE_STRIDE + token_local / 4][token_local % 4]; + totals[i] += max(score, 0.0) * weight_sh[head_local][token_local]; + } + } + if (head_local + 1 < HEADS_PER_TILE) { + controlBarrier(gl_ScopeSubgroup, gl_ScopeSubgroup, gl_StorageSemanticsShared, gl_SemanticsAcquireRelease); + } + } + barrier(); + } + + [[unroll]] for (uint i = 0; i < 4; ++i) { + const uint idx = lane + i * SUBGROUP_SIZE; + const uint key = idx / TILE; + const uint token_local = idx % TILE; + const uint kv = kv_base + key; + const uint token = token_base + token_local; + if (kv < p.n_kv && token < p.n_batch) { + const uint mask_base = (stream % p.nem3) * p.nbm3 + token * p.nbm1; + const uint dst_base = stream * p.nb3 + token * p.nb1; + data_dst[dst_base + kv] = totals[i] + float(data_m[mask_base + kv]); + } + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_decode_cm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_decode_cm.comp new file mode 100644 index 000000000000..5d32e51361b3 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_decode_cm.comp @@ -0,0 +1,137 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_KHR_cooperative_matrix : require +#extension GL_KHR_memory_scope_semantics : require + +layout(constant_id = 0) const uint SUBGROUP_SIZE = 64; +// Head count is a SPECIALIZATION constant, not a push constant: the head loop's trip +// count must stay visible to the compiler. As a push constant it cannot unroll and +// decode costs ~55% more (measured, gfx1151, kv=131584). +layout(constant_id = 1) const uint N_HEAD = 64; +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout(binding = 0) readonly buffer QBuf { float data_q[]; }; +layout(binding = 1) readonly buffer KBuf { float16_t data_k[]; }; +layout(binding = 2) readonly buffer WBuf { float data_w[]; }; +layout(binding = 3) readonly buffer MBuf { float16_t data_m[]; }; +layout(binding = 4) writeonly buffer DstBuf { float data_dst[]; }; + +layout(push_constant) uniform Parameters { + uint n_kv; + uint n_batch; + uint n_stream; + uint nem3; + uint nb1; + uint nb3; + uint nbq1; + uint nbq2; + uint nbq3; + uint nbk2; + uint nbk3; + uint nbw1; + uint nbw3; + uint nbm1; + uint nbm3; +} p; + +const uint TILE = 16; +const uint HEAD_SIZE = 128; +const uint VEC_PER_HEAD = HEAD_SIZE / 4; +const uint TILE_STRIDE = VEC_PER_HEAD + 2; +const uint SCORE_STRIDE = TILE / 4 + 1; + +shared f16vec4 q_sh[TILE * TILE_STRIDE]; +shared f16vec4 k_sh[TILE * TILE_STRIDE]; +shared vec4 score_sh[TILE * SCORE_STRIDE]; + +void main() { + const uint tid = gl_LocalInvocationIndex; + const uint kv_base = gl_WorkGroupID.x * TILE; + const uint token = gl_WorkGroupID.y; + const uint stream = gl_WorkGroupID.z; + + for (uint idx = tid; idx < TILE * VEC_PER_HEAD; idx += SUBGROUP_SIZE) { + const uint key = idx / VEC_PER_HEAD; + const uint d4 = idx % VEC_PER_HEAD; + const uint kv = kv_base + key; + f16vec4 value = f16vec4(0.0); + if (kv < p.n_kv) { + const uint offset = stream * p.nbk3 + kv * p.nbk2 + d4 * 4; + value = f16vec4(data_k[offset], data_k[offset + 1], data_k[offset + 2], data_k[offset + 3]); + } + k_sh[key * TILE_STRIDE + d4] = value; + } + barrier(); + + // The K tile is invariant across the head loop, but coopMatLoad is inside it, so the shipped + // shader re-reads the same 8 A fragments from shared memory once per head tile (4x at + // N_HEAD=64). Each wave64 fragment load moves 1 KiB of LDS traffic because the 16x16 f16 + // fragment is replicated 4x across the subgroup, so those re-reads are the largest single + // item of per-tile cost after the multiplies themselves. Hoisting costs 8 A fragments of + // register state (16 f16 per lane each). + coopmat kmat[HEAD_SIZE / TILE]; + [[unroll]] for (uint d = 0; d < HEAD_SIZE; d += TILE) { + coopMatLoad(kmat[d / TILE], k_sh, d / 4, TILE_STRIDE, gl_CooperativeMatrixLayoutRowMajor); + } + + float total = 0.0; + for (uint head_base = 0; head_base < N_HEAD; head_base += TILE) { + for (uint idx = tid; idx < TILE * VEC_PER_HEAD; idx += SUBGROUP_SIZE) { + const uint head_local = idx / VEC_PER_HEAD; + const uint d4 = idx % VEC_PER_HEAD; + const uint head = head_base + head_local; + // n_head need not be a multiple of TILE (Qwen4-exp indexers are nh=4), so the tail + // tile stages zeros rather than reading past the end of Q. Zero q gives score 0, + // and relu(0) * w = 0, so the padded heads drop out of the sum on their own. + // The guard exists only when N_HEAD is not tile-aligned (nh=4). The condition is a + // specialization-constant expression, so at nh=64 it folds to true and this compiles + // to the original unguarded load - measured: relying on the compiler to range-prove + // head < N_HEAD instead costs ~45% at nh=64. + f16vec4 value = f16vec4(0.0); + if (N_HEAD % TILE == 0 || head < N_HEAD) { + const uint offset = stream * p.nbq3 + token * p.nbq2 + head * p.nbq1 + d4 * 4; + value = f16vec4(data_q[offset], data_q[offset + 1], data_q[offset + 2], data_q[offset + 3]); + } + q_sh[head_local * TILE_STRIDE + d4] = value; + } + barrier(); + + coopmat scores = + coopmat(0.0); + coopmat qmat; + + [[unroll]] for (uint d = 0; d < HEAD_SIZE; d += TILE) { + coopMatLoad(qmat, q_sh, d / 4, TILE_STRIDE, gl_CooperativeMatrixLayoutColumnMajor); + scores = coopMatMulAdd(kmat[d / TILE], qmat, scores); + } + + coopMatStore(scores, score_sh, 0, SCORE_STRIDE, gl_CooperativeMatrixLayoutRowMajor); + barrier(); + + if (tid < TILE && kv_base + tid < p.n_kv) { + [[unroll]] for (uint head_local = 0; head_local < TILE; ++head_local) { + // Padded heads staged q as zero, so their score is 0 and the term dies in the + // multiply. Only the weight READ needs bounding, and clamping the index does that + // without a branch - a dynamic break here forces the loop rolled and costs ~60%. + const uint head = N_HEAD % TILE == 0 ? head_base + head_local + : min(head_base + head_local, N_HEAD - 1); + const float score = score_sh[tid * SCORE_STRIDE + head_local / 4][head_local % 4]; + const float weight = data_w[stream * p.nbw3 + token * p.nbw1 + head]; + total += max(score, 0.0) * weight; + } + } + barrier(); + } + + if (tid < TILE) { + const uint kv = kv_base + tid; + if (kv < p.n_kv) { + const uint mask_base = (stream % p.nem3) * p.nbm3 + token * p.nbm1; + const uint dst_base = stream * p.nb3 + token * p.nb1; + data_dst[dst_base + kv] = total + float(data_m[mask_base + kv]); + } + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_scalar64.comp b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_scalar64.comp new file mode 100644 index 000000000000..83fd923fe738 --- /dev/null +++ b/ggml/src/ggml-vulkan/vulkan-shaders/lightning_indexer_scalar64.comp @@ -0,0 +1,95 @@ +#version 450 + +#extension GL_EXT_control_flow_attributes : require +#extension GL_EXT_shader_16bit_storage : require +#extension GL_EXT_shader_explicit_arithmetic_types_float16 : require +#extension GL_KHR_shader_subgroup_arithmetic : require +#extension GL_KHR_shader_subgroup_basic : require + +layout(constant_id = 0) const uint SUBGROUP_SIZE = 64; +// Head count is a SPECIALIZATION constant, not a push constant: the head loop's trip +// count must stay visible to the compiler. As a push constant it cannot unroll and +// decode costs ~55% more (measured, gfx1151, kv=131584). +layout(constant_id = 1) const uint N_HEAD = 64; +layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; + +layout(binding = 0) readonly buffer QBuf { float data_q[]; }; +layout(binding = 1) readonly buffer KBuf { float16_t data_k[]; }; +layout(binding = 2) readonly buffer WBuf { float data_w[]; }; +layout(binding = 3) readonly buffer MaskBuf { float16_t data_m[]; }; +layout(binding = 4) writeonly buffer DstBuf { float data_dst[]; }; + +layout(push_constant) uniform Parameters { + uint n_kv; + uint n_batch; + uint n_stream; + uint nem3; + uint nb1; + uint nb3; + uint nbq1; + uint nbq2; + uint nbq3; + uint nbk2; + uint nbk3; + uint nbw1; + uint nbw3; + uint nbm1; + uint nbm3; +} p; + +const uint K_PER_GROUP = 8; + +void main() { + const uint lane = gl_SubgroupInvocationID; + const uint token = gl_WorkGroupID.y; + const uint stream = gl_WorkGroupID.z; + const uint kv_base = gl_WorkGroupID.x * K_PER_GROUP; + + if (token >= p.n_batch || stream >= p.n_stream) { + return; + } + + float k0[K_PER_GROUP]; + float k1[K_PER_GROUP]; + [[unroll]] for (uint j = 0; j < K_PER_GROUP; ++j) { + const uint kv = kv_base + j; + if (kv < p.n_kv) { + const uint k_base = stream * p.nbk3 + kv * p.nbk2; + k0[j] = float(data_k[k_base + lane]); + k1[j] = float(data_k[k_base + lane + SUBGROUP_SIZE]); + } else { + k0[j] = 0.0; + k1[j] = 0.0; + } + } + + float score[K_PER_GROUP]; + [[unroll]] for (uint j = 0; j < K_PER_GROUP; ++j) { + score[j] = 0.0; + } + + for (uint head = 0; head < N_HEAD; ++head) { + const uint q_base = stream * p.nbq3 + token * p.nbq2 + head * p.nbq1; + const float q0 = data_q[q_base + lane]; + const float q1 = data_q[q_base + lane + SUBGROUP_SIZE]; + const float weight = data_w[stream * p.nbw3 + token * p.nbw1 + head]; + + [[unroll]] for (uint j = 0; j < K_PER_GROUP; ++j) { + const float qk = subgroupAdd(q0 * k0[j] + q1 * k1[j]); + if (lane == 0) { + score[j] += max(qk, 0.0) * weight; + } + } + } + + if (lane == 0) { + const uint mask_base = (stream % p.nem3) * p.nbm3 + token * p.nbm1; + const uint dst_base = stream * p.nb3 + token * p.nb1; + [[unroll]] for (uint j = 0; j < K_PER_GROUP; ++j) { + const uint kv = kv_base + j; + if (kv < p.n_kv) { + data_dst[dst_base + kv] = score[j] + float(data_m[mask_base + kv]); + } + } + } +} diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vecq.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vecq.comp index 18d441ead40e..0ebdde36cd71 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vecq.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vecq.comp @@ -13,7 +13,12 @@ layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; #if defined(DATA_A_Q2_0) || defined(DATA_A_QUANT_K) #define K_PER_ITER 16 -#elif defined(DATA_A_QUANT_LEGACY) || defined(DATA_A_MXFP4) +#elif defined(DATA_A_ROCMFP4) || defined(DATA_A_ROCMFP4_FAST) +// A whole block per call. The ROCmFP4 scales are UE4M3 and cost a byte load plus +// a table lookup each, so decoding them once per 32 weights instead of once per +// 8 is worth the extra B registers. +#define K_PER_ITER 32 +#elif defined(DATA_A_QUANT_LEGACY) || defined(DATA_A_MXFP4) || defined(DATA_A_ROCMFPX_FP8) #define K_PER_ITER 8 #elif defined(DATA_A_IQ1_S) || defined(DATA_A_IQ1_M) #define K_PER_ITER 32 @@ -37,9 +42,16 @@ void iter(inout FLOAT_TYPE temp[NUM_COLS][NUM_ROWS], const uint first_row, const cache_b_ds = vec2(data_b[b_block_idx_outer].ds[b_block_idx_inner]); #if QUANT_R == 2 - // Assumes K_PER_ITER == 8 +#if K_PER_ITER == 8 cache_b_qs[0] = data_b[b_block_idx_outer].qs[b_block_idx_inner * 8 + b_qs_idx]; cache_b_qs[1] = data_b[b_block_idx_outer].qs[b_block_idx_inner * 8 + b_qs_idx + 4]; +#elif K_PER_ITER == 32 + [[unroll]] for (uint k = 0; k < 8; ++k) { + cache_b_qs[k] = data_b[b_block_idx_outer].qs[b_block_idx_inner * 8 + k]; + } +#else +#error unimplemented +#endif #else #if K_PER_ITER == 8 cache_b_qs[0] = data_b[b_block_idx_outer].qs[b_block_idx_inner * 8 + b_qs_idx * 2]; diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vecq_funcs.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vecq_funcs.glsl index a5403ac82121..bded206808c6 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vecq_funcs.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mat_vecq_funcs.glsl @@ -137,6 +137,70 @@ FLOAT_TYPE mul_q8_1(const int32_t q_sum, const float da, const vec2 dsb, const i } #endif +#if defined(DATA_A_ROCMFP4) +// dual UE4M3 scales: e[0] covers elements 0-15, e[1] covers 16-31 +FLOAT_TYPEV2 get_dm2(uint ib) { + return FLOAT_TYPEV2(ue4m3_to_fp32(data_a[ib].e[0]), ue4m3_to_fp32(data_a[ib].e[1])); +} +#endif + +#if defined(DATA_A_ROCMFP4_FAST) +FLOAT_TYPEV2 get_dm2(uint ib) { + const FLOAT_TYPE d = FLOAT_TYPE(ue4m3_to_fp32(data_a[ib].e)); + return FLOAT_TYPEV2(d, d); +} +#endif + +#if defined(DATA_A_ROCMFP4) || defined(DATA_A_ROCMFP4_FAST) +// 1-byte quant loads for ROCmFP4 blocks (18 / 17 bytes). Low nibbles are the +// first half of the block, high nibbles the second half. +i32vec2 repack(uint ib, uint iqs) { + const uint32_t qs = pack32(u8vec4(data_a[ib].qs[iqs * 4 ], + data_a[ib].qs[iqs * 4 + 1], + data_a[ib].qs[iqs * 4 + 2], + data_a[ib].qs[iqs * 4 + 3])); + + const u8vec4 i_a0 = unpack8( qs & 0x0F0F0F0F); + const u8vec4 i_a1 = unpack8((qs >> 4) & 0x0F0F0F0F); + + return i32vec2(pack32(i8vec4(kvalues_rocmfp4[i_a0.x], kvalues_rocmfp4[i_a0.y], kvalues_rocmfp4[i_a0.z], kvalues_rocmfp4[i_a0.w])), + pack32(i8vec4(kvalues_rocmfp4[i_a1.x], kvalues_rocmfp4[i_a1.y], kvalues_rocmfp4[i_a1.z], kvalues_rocmfp4[i_a1.w]))); +} + +// K_PER_ITER is 32 for these types: one call consumes a whole block, so the +// UE4M3 scales are decoded once per block rather than once per 8 weights. +FLOAT_TYPE mmvq_dot_product(const uint ib_a, const uint iqs) { + // the two halves carry different scales, so they cannot share an accumulator + int32_t q_sum0 = 0; + int32_t q_sum1 = 0; + + [[unroll]] for (uint k = 0; k < 4; ++k) { + const i32vec2 data_a_qs = repack(ib_a, k); + q_sum0 += dotPacked4x8EXT(data_a_qs.x, cache_b_qs[k]); + q_sum1 += dotPacked4x8EXT(data_a_qs.y, cache_b_qs[k + 4]); + } + + const FLOAT_TYPEV2 d = get_dm2(ib_a); + return FLOAT_TYPE(cache_b_ds.x * (float(q_sum0) * float(d.x) + float(q_sum1) * float(d.y))); +} +#endif + +#if defined(DATA_A_ROCMFPX_FP8) +// int8 values with one UE4M3 scale per 32-element block +FLOAT_TYPE mmvq_dot_product(const uint ib_a, const uint iqs) { + const int32_t qs0 = pack32(i8vec4(data_a[ib_a].qs[iqs * 8 ], data_a[ib_a].qs[iqs * 8 + 1], + data_a[ib_a].qs[iqs * 8 + 2], data_a[ib_a].qs[iqs * 8 + 3])); + const int32_t qs1 = pack32(i8vec4(data_a[ib_a].qs[iqs * 8 + 4], data_a[ib_a].qs[iqs * 8 + 5], + data_a[ib_a].qs[iqs * 8 + 6], data_a[ib_a].qs[iqs * 8 + 7])); + + const int32_t q_sum = dotPacked4x8EXT(qs0, cache_b_qs[0]) + + dotPacked4x8EXT(qs1, cache_b_qs[1]); + + const FLOAT_TYPE d = FLOAT_TYPE(ue4m3_to_fp32(data_a[ib_a].e)); + return FLOAT_TYPE(cache_b_ds.x * float(d) * float(q_sum)); +} +#endif + #if defined(DATA_A_MXFP4) // 1-byte loads for mxfp4 blocks (17 bytes) i32vec2 repack(uint ib, uint iqs) { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp index 90e4e11cdec7..7d32124b969b 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm.comp @@ -9,9 +9,6 @@ #if defined(DATA_A_IQ1_M) #extension GL_EXT_shader_explicit_arithmetic_types_int16 : require #endif -#if !defined(DATA_A_F32) && !defined(DATA_A_F16) && !defined(DATA_A_BF16) -#extension GL_EXT_shader_explicit_arithmetic_types_int16 : require -#endif #if defined(DATA_A_BF16) && defined(COOPMAT) #extension GL_EXT_bfloat16 : enable @@ -31,54 +28,24 @@ #extension GL_EXT_shader_explicit_arithmetic_types_int16 : require #endif -#ifdef MULMAT_QUANT -#include "ggml_type_ids.glsl" -layout (constant_id = 12) const uint MmTypeA = 0; -#endif - #include "types.glsl" #include "dot_product_funcs.glsl" -#ifndef MULMAT_QUANT #ifndef LOAD_VEC_A #define LOAD_VEC_A 1 #endif -#endif #ifndef LOAD_VEC_B #define LOAD_VEC_B 1 #endif layout (constant_id = 11) const uint ALIGNED = 0; -#ifdef MULMAT_QUANT - -uint mm_load_vec_a() { - switch (MmTypeA) { - case GGML_TYPE_Q1_0: - case GGML_TYPE_Q4_0: - case GGML_TYPE_Q4_1: - case GGML_TYPE_Q5_1: - return 8u; - case GGML_TYPE_Q2_0: - case GGML_TYPE_Q5_0: - case GGML_TYPE_Q8_0: - case GGML_TYPE_Q2_K: - case GGML_TYPE_Q4_K: - case GGML_TYPE_Q5_K: - return 4u; - default: - return 2u; - } -} -#endif - #if !defined(TO_FLOAT_TYPE) #define TO_FLOAT_TYPE FLOAT_TYPE #endif layout(local_size_x_id = 0, local_size_y = 1, local_size_z = 1) in; -#ifndef MULMAT_QUANT layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; #if defined(DATA_A_F32) layout (binding = 0) readonly buffer A_SCALAR {float data_a_scalar[];}; @@ -93,30 +60,6 @@ layout (binding = 0) readonly buffer A_PACKED16 {A_TYPE_PACKED16 data_a_packed16 #if defined(A_TYPE_PACKED32) layout (binding = 0) readonly buffer A_PACKED32 {A_TYPE_PACKED32 data_a_packed32[];}; #endif -#else -// Unpacked struct aliases -layout (binding = 0) readonly buffer BUF_Q1_0 { block_q1_0 data[]; } a_q1_0; -layout (binding = 0) readonly buffer BUF_Q2_0 { block_q2_0 data[]; } a_q2_0; -layout (binding = 0) readonly buffer BUF_Q2_K { block_q2_K data[]; } a_q2_k; -layout (binding = 0) readonly buffer BUF_Q3_K { block_q3_K data[]; } a_q3_k; -layout (binding = 0) readonly buffer BUF_Q4_K { block_q4_K data[]; } a_q4_k; -layout (binding = 0) readonly buffer BUF_Q5_K { block_q5_K data[]; } a_q5_k; -layout (binding = 0) readonly buffer BUF_Q6_K { block_q6_K data[]; } a_q6_k; -layout (binding = 0) readonly buffer BUF_TQ1_0 { block_tq1_0 data[]; } a_tq1_0; -layout (binding = 0) readonly buffer BUF_TQ2_0 { block_tq2_0 data[]; } a_tq2_0; -// Packed16 aliases -layout (binding = 0) readonly buffer BUF_Q4_0_P16 { block_q4_0_packed16 data[]; } a_q4_0_p16; -layout (binding = 0) readonly buffer BUF_Q5_0_P16 { block_q5_0_packed16 data[]; } a_q5_0_p16; -layout (binding = 0) readonly buffer BUF_Q8_0_P16 { block_q8_0_packed16 data[]; } a_q8_0_p16; -layout (binding = 0) readonly buffer BUF_Q3_K_P16 { block_q3_K_packed16 data[]; } a_q3_k_p16; -layout (binding = 0) readonly buffer BUF_Q6_K_P16 { block_q6_K_packed16 data[]; } a_q6_k_p16; -// Packed32 aliases -layout (binding = 0) readonly buffer BUF_Q4_1_P32 { block_q4_1_packed32 data[]; } a_q4_1_p32; -layout (binding = 0) readonly buffer BUF_Q5_1_P32 { block_q5_1_packed32 data[]; } a_q5_1_p32; -layout (binding = 0) readonly buffer BUF_Q2_K_P32 { block_q2_K_packed32 data[]; } a_q2_k_p32; -layout (binding = 0) readonly buffer BUF_Q4_K_P32 { block_q4_K_packed32 data[]; } a_q4_k_p32; -layout (binding = 0) readonly buffer BUF_Q5_K_P32 { block_q5_K_packed32 data[]; } a_q5_k_p32; -#endif layout (binding = 1) readonly buffer B {B_TYPE data_b[];}; layout (binding = 1) readonly buffer B_SCALAR {B_TYPE_SCALAR data_b_scalar[];}; @@ -125,6 +68,9 @@ layout (binding = 2) writeonly buffer D {D_TYPE data_d[];}; #ifdef MUL_MAT_ID layout (binding = 3) readonly buffer IDS {int data_ids[];}; layout (binding = 4) readonly buffer Counts {int data_expert_count[];}; +// Fused MUL epilogue: one scale per (expert slot, token), broadcast down the M dimension. +// Always bound; p.fusion_flags == 0 means ignore it. +layout (binding = 5) readonly buffer FusedScale {float data_fscale[];}; #endif layout (push_constant) uniform parameter @@ -147,6 +93,7 @@ layout (push_constant) uniform parameter uint ne11; uint n_experts; uint hoist_row_ids; + uint fusion_flags; #else uint base_work_group_z; uint num_batches; @@ -178,13 +125,8 @@ layout (constant_id = 3) const uint BK = 16; // Assumed to be 32 if working wit #endif #ifdef COOPMAT -#ifdef MULMAT_QUANT -layout(constant_id = 13) const uint SHMEM_STRIDE_PAD = 4; -layout(constant_id = 14) const bool APPLY_SLM_A_RESHAPE = false; -#else layout(constant_id = 12) const uint SHMEM_STRIDE_PAD = 4; layout(constant_id = 13) const bool APPLY_SLM_A_RESHAPE = false; -#endif #else const uint SHMEM_STRIDE_PAD = 1; const bool APPLY_SLM_A_RESHAPE = false; @@ -203,10 +145,6 @@ shared ACC_TYPE coopmat_stage[TM * TN * NUM_WARPS]; #include "mul_mm_id_funcs.glsl" #include "mul_mm_funcs.glsl" -#ifdef MULMAT_QUANT -#include "iq_shmem_init.glsl" -#endif - void main() { const uint ic = gl_WorkGroupID.y; @@ -216,7 +154,7 @@ void main() { return; } #endif -#if defined(NEEDS_INIT_IQ_SHMEM) || defined(MULMAT_QUANT) +#ifdef NEEDS_INIT_IQ_SHMEM init_iq_shmem(gl_WorkGroupSize); #endif @@ -266,11 +204,8 @@ void main() { #if defined(DATA_A_F32) || defined(DATA_A_F16) || defined(DATA_A_BF16) const uint LOAD_VEC_A_EFF = (ALIGNED != 0) ? LOAD_VEC_A : 1; const uint LOAD_VEC_BATCH_A = (ALIGNED != 0) ? 1 : 2; -#elif !defined(MULMAT_QUANT) - const uint LOAD_VEC_A_EFF = LOAD_VEC_A; - const uint LOAD_VEC_BATCH_A = 1; #else - const uint LOAD_VEC_A_EFF = mm_load_vec_a(); + const uint LOAD_VEC_A_EFF = LOAD_VEC_A; const uint LOAD_VEC_BATCH_A = 1; #endif const uint LOAD_VEC_B_EFF = (ALIGNED != 0) ? LOAD_VEC_B : 1; @@ -463,9 +398,13 @@ void main() { if (row_i >= _ne1) break; const u16vec2 row_idx = row_ids[row_i - ic * BN]; - if (dr + cm_row * TM + store_r < p.M) { - data_d[row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr + cm_row * TM + store_r] = D_TYPE(coopmat_stage[warp_i * TM * TN + (col + store_c) * TM + store_r]); + const uint didx = row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr + cm_row * TM + store_r; + if (p.fusion_flags != 0) { + data_d[didx] = D_TYPE(float(coopmat_stage[warp_i * TM * TN + (col + store_c) * TM + store_r]) * data_fscale[row_idx.y * p.nei0 + row_idx.x]); + } else { + data_d[didx] = D_TYPE(coopmat_stage[warp_i * TM * TN + (col + store_c) * TM + store_r]); + } } } barrier(); @@ -518,15 +457,19 @@ void main() { if (row_i >= _ne1) break; const u16vec2 row_idx = row_ids[row_i - ic * BN]; + const bool do_scale = p.fusion_flags != 0; + const float fscale = do_scale ? data_fscale[row_idx.y * p.nei0 + row_idx.x] : 1.0f; #endif // MUL_MAT_ID [[unroll]] for (uint cr = 0; cr < TM / 2; cr++) { const uint sums_idx = (wsic * TN + cc) * WMITER * (TM / 2) + wsir * (TM / 2) + cr; #ifdef MUL_MAT_ID if (dr_warp + 2 * cr < p.M) { - data_d[row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr_warp + 2 * cr] = D_TYPE(sums[sums_idx].x); + const uint didx = row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr_warp + 2 * cr; + data_d[didx] = do_scale ? D_TYPE(float(sums[sums_idx].x) * fscale) : D_TYPE(sums[sums_idx].x); } if (dr_warp + 2 * cr + 1 < p.M) { - data_d[row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr_warp + 2 * cr + 1] = D_TYPE(sums[sums_idx].y); + const uint didx = row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr_warp + 2 * cr + 1; + data_d[didx] = do_scale ? D_TYPE(float(sums[sums_idx].y) * fscale) : D_TYPE(sums[sums_idx].y); } #else if (dr_warp + 2 * cr < p.M && dc_warp + cc < p.N) { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp index 189788a86257..9bfad031d1b6 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_cm2.comp @@ -21,13 +21,6 @@ #extension GL_EXT_bfloat16 : enable #endif -#include "ggml_type_ids.glsl" - -#ifdef MULMAT_QUANT -layout (constant_id = 7) const uint MmTypeA = 0; -layout (constant_id = 8) const uint MmABlockBytes = 2; -#endif - #include "types.glsl" #include "utils.glsl" @@ -44,24 +37,6 @@ layout (constant_id = 4) const bool enable_smaller_matrices = false; const uint BNover2 = enable_smaller_matrices ? (BN / 2) : BN; const uint BNover4 = enable_smaller_matrices ? (BN / 4) : BN; layout (constant_id = 5) const uint ALIGNED = 0; -layout (constant_id = 6) const uint subgroup_size = 32; - -#ifdef MULMAT_QUANT - -uint mm_quant_k() { - switch (MmTypeA) { - case GGML_TYPE_Q4_0: case GGML_TYPE_Q4_1: case GGML_TYPE_Q5_0: case GGML_TYPE_Q5_1: - case GGML_TYPE_Q8_0: - return 32u; - case GGML_TYPE_Q1_0: - return 128u; - case GGML_TYPE_Q2_0: - return 64u; - default: - return 256u; - } -} -#endif layout (push_constant) uniform parameter { @@ -97,73 +72,21 @@ layout (push_constant) uniform parameter } p; -#ifndef MULMAT_QUANT layout (binding = 0) readonly buffer A {A_TYPE data_a[];}; -#else -layout (binding = 0) readonly buffer A {uint8_t data_a[];}; -#endif layout (binding = 1) readonly buffer B {B_TYPE data_b[];}; layout (binding = 2) writeonly buffer D {D_TYPE data_d[];}; #if defined(MUL_MAT_ID) && defined(GGML_VULKAN_COOPMAT2_DECODE_VECTOR) layout (binding = 1) readonly buffer B4 {B_TYPEV4 data_b_v4[];}; #endif -#if defined(MULMAT_QUANT) || QUANT_K > 1 +#if QUANT_K > 1 #include "dequant_funcs_cm2.glsl" -#ifndef MULMAT_QUANT -// Per-type path: use the alias set by dequant_funcs_cm2.glsl #if defined(dequantFuncA_v) && defined(GGML_VULKAN_COOPMAT2_DECODE_VECTOR) #define DECODEFUNCA , dequantFuncA, dequantFuncA_v #else #define DECODEFUNCA , dequantFuncA #endif #else -layout(buffer_reference, std430, buffer_reference_align = 1) buffer decodeBufA { - uint8_t raw[MmABlockBytes]; -}; - -float16_t mmDecodeA(const in decodeBufA bl_in, const in uint blockCoords[2], const in uint coordInBlock[2]) { - switch (MmTypeA) { - case GGML_TYPE_Q1_0: return dequantFuncQ1_0 (decodeBufQ1_0 (bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q2_0: return dequantFuncQ2_0 (decodeBufQ2_0 (bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q4_0: return dequantFuncQ4_0 (decodeBufQ4_0 (bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q4_1: return dequantFuncQ4_1 (decodeBufQ4_1 (bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q5_0: return dequantFuncQ5_0 (decodeBufQ5_0 (bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q5_1: return dequantFuncQ5_1 (decodeBufQ5_1 (bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q8_0: return dequantFuncQ8_0 (decodeBufQ8_0 (bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q2_K: return dequantFuncQ2_K (decodeBufQ2_K (bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q3_K: return dequantFuncQ3_K (decodeBufQ3_K (bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q6_K: return dequantFuncQ6_K (decodeBufQ6_K (bl_in), blockCoords, coordInBlock); - case GGML_TYPE_TQ1_0: return dequantFuncTQ1_0(decodeBufTQ1_0(bl_in), blockCoords, coordInBlock); - case GGML_TYPE_TQ2_0: return dequantFuncTQ2_0(decodeBufTQ2_0(bl_in), blockCoords, coordInBlock); - default: return float16_t(0); - } -} - -#ifdef GGML_VULKAN_COOPMAT2_DECODE_VECTOR -f16vec4 mmDecodeA_v(const in decodeBufA bl_in, const in uint blockCoords[2], const in uint coordInBlock[2]) { - switch (MmTypeA) { - case GGML_TYPE_Q1_0: return dequantFuncQ1_0_v (decodeBufQ1_0 (bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q2_0: return dequantFuncQ2_0_v (decodeBufQ2_0 (bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q4_0: return dequantFuncQ4_0_v (decodeBufQ4_0 (bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q4_1: return dequantFuncQ4_1_v (decodeBufQ4_1 (bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q5_0: return dequantFuncQ5_0_v (decodeBufQ5_0 (bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q5_1: return dequantFuncQ5_1_v (decodeBufQ5_1 (bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q8_0: return dequantFuncQ8_0_v (decodeBufQ8_0 (bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q2_K: return dequantFuncQ2_K_v (decodeBufQ2_K (bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q3_K: return dequantFuncQ3_K_v (decodeBufQ3_K (bl_in), blockCoords, coordInBlock); - case GGML_TYPE_Q6_K: return dequantFuncQ6_K_v (decodeBufQ6_K (bl_in), blockCoords, coordInBlock); - case GGML_TYPE_TQ1_0: return dequantFuncTQ1_0_v(decodeBufTQ1_0(bl_in), blockCoords, coordInBlock); - case GGML_TYPE_TQ2_0: return dequantFuncTQ2_0_v(decodeBufTQ2_0(bl_in), blockCoords, coordInBlock); - default: return f16vec4(0); - } -} -#define DECODEFUNCA , mmDecodeA, mmDecodeA_v -#else -#define DECODEFUNCA , mmDecodeA -#endif -#endif -#else #define DECODEFUNCA #endif @@ -183,6 +106,9 @@ f16vec4 mmDecodeA_v(const in decodeBufA bl_in, const in uint blockCoords[2], con #ifdef MUL_MAT_ID layout (binding = 3) readonly buffer IDS {int data_ids[];}; layout (binding = 4) readonly buffer Counts {int data_expert_count[];}; +// Bound for descriptor-layout parity with the other mul_mat_id shaders. The fused MUL +// epilogue is not implemented here, so the host never enables it on coopmat2 devices. +layout (binding = 5) readonly buffer FusedScale {float data_fscale[];}; shared u16vec4 row_ids[BN]; @@ -191,6 +117,7 @@ layout(buffer_reference, std430, buffer_reference_align = 2) buffer decodeBufB { }; uint _ne1; +layout (constant_id = 6) const uint subgroup_size = 32; shared uvec4 ballots_sh[BLOCK_SIZE / subgroup_size]; B_TYPE decodeFuncB(const in decodeBufB bl, const in uint blockCoords[2], const in uint coordInBlock[2]) @@ -322,10 +249,6 @@ void load_row_ids_hoisted(uint expert_idx, uint ic) { } #endif -#ifdef MULMAT_QUANT -#include "iq_shmem_init.glsl" -#endif - void main() { const uint tid = gl_LocalInvocationIndex; const uint ic = gl_WorkGroupID.y; @@ -344,7 +267,7 @@ void main() { #endif #endif -#if defined(NEEDS_INIT_IQ_SHMEM) || defined(MULMAT_QUANT) +#ifdef NEEDS_INIT_IQ_SHMEM init_iq_shmem(gl_WorkGroupSize); #endif @@ -385,33 +308,22 @@ void main() { const uint end_k = min(p.K, (ik + 1) * p.k_split); #endif -#ifdef MULMAT_QUANT - const uint qk = mm_quant_k(); -#else - const uint qk = QUANT_K; -#endif - #ifdef MUL_MAT_ID - uint pos_a = expert_idx * (p.batch_stride_a / qk); + uint pos_a = expert_idx * (p.batch_stride_a / QUANT_K); uint pos_b = 0; #else - uint pos_a = batch_idx_a * (p.batch_stride_a / qk); + uint pos_a = batch_idx_a * (p.batch_stride_a / QUANT_K); uint pos_b = batch_idx * p.batch_stride_b; uint pos_d = batch_idx * p.batch_stride_d + ik * p.batch_stride_d * p.num_batches; #endif -#ifdef MULMAT_QUANT - // pos_a is a byte offset into the raw buffer; strides stay in block units - pos_a *= MmABlockBytes; -#endif - - uint stride_a = p.stride_a / qk; + uint stride_a = p.stride_a / QUANT_K; uint stride_b = p.stride_b; // Hint to the compiler that values are aligned (want 16B alignment). // Quants are always block-aligned, no alignment needed. if (ALIGNED != 0) { -#if !defined(MULMAT_QUANT) && QUANT_K == 1 +#if QUANT_K == 1 stride_a &= ~7; #endif stride_b &= ~7; @@ -426,8 +338,10 @@ void main() { #endif tensorLayoutNV<2, gl_CooperativeMatrixClampModeConstantNV> tensorLayoutD = createTensorLayoutNV(2, gl_CooperativeMatrixClampModeConstantNV); - tensorLayoutA = setTensorLayoutBlockSizeNV(tensorLayoutA, 1, qk); - tensorLayoutAClamp = setTensorLayoutBlockSizeNV(tensorLayoutAClamp, 1, qk); +#if QUANT_K > 1 + tensorLayoutA = setTensorLayoutBlockSizeNV(tensorLayoutA, 1, QUANT_K); + tensorLayoutAClamp = setTensorLayoutBlockSizeNV(tensorLayoutAClamp, 1, QUANT_K); +#endif #if defined(MUL_MAT_ID) && defined(GGML_VULKAN_COOPMAT2_DECODE_VECTOR) tensorLayoutB = setTensorLayoutBlockSizeNV(tensorLayoutB, 1, BK); #endif @@ -457,19 +371,19 @@ void main() { const uint START_ALIGN_K = 256; // For Qi_K (block size 256), unroll whole 256 element tiles. // For legacy quants (block size 32), unroll 8x. - const uint UNROLL_K = (qk == 256) ? 256 : (BK * 8); + const uint UNROLL_K = (QUANT_K == 256) ? 256 : (BK * 8); const uint unroll_count = UNROLL_K / BK; // Detect a fast path where all loads are entirely in bounds and no clamping is required if ((ir + 1) * BM <= p.M && (ic + 1) * BN <= p.padded_N && (start_k % START_ALIGN_K) == 0 && (end_k % BK) == 0 && -#if !defined(MULMAT_QUANT) && QUANT_K == 1 +#if QUANT_K == 1 (stride_a % 8) == 0 && #endif (stride_b % 8) == 0) { // Hint to the compiler that values are aligned (want 16B alignment) start_k &= ~(START_ALIGN_K-1); stride_b &= ~7; -#if !defined(MULMAT_QUANT) && QUANT_K == 1 +#if QUANT_K == 1 stride_a &= ~7; #endif @@ -640,10 +554,10 @@ void main() { [[dont_unroll]] for (uint block_k = start_k, i = 0; i < k_iters; block_k += BK, ++i) { - if ((block_k % qk) == 0) { + if ((block_k % QUANT_K) == 0) { store_scales(tid); } - if (block_k + BK < end_k && ((block_k + BK) % qk) == 0) { + if (block_k + BK < end_k && ((block_k + BK) % QUANT_K) == 0) { fetch_scales(ir * BM, pos_a, stride_a, block_k + BK, tid, false); } @@ -684,10 +598,10 @@ void main() { [[dont_unroll]] for (uint block_k = start_k, i = 0; i < k_iters; block_k += BK, ++i) { - if ((block_k % qk) == 0) { + if ((block_k % QUANT_K) == 0) { store_scales(tid); } - if (block_k + BK < end_k && ((block_k + BK) % qk) == 0) { + if (block_k + BK < end_k && ((block_k + BK) % QUANT_K) == 0) { fetch_scales(ir * BM, pos_a, stride_a, block_k + BK, tid, false); } @@ -728,10 +642,10 @@ void main() { [[dont_unroll]] for (uint block_k = start_k, i = 0; i < k_iters; block_k += BK, ++i) { - if ((block_k % qk) == 0) { + if ((block_k % QUANT_K) == 0) { store_scales(tid); } - if (block_k + BK < end_k && ((block_k + BK) % qk) == 0) { + if (block_k + BK < end_k && ((block_k + BK) % QUANT_K) == 0) { fetch_scales(ir * BM, pos_a, stride_a, block_k + BK, tid, false); } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_funcs.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_funcs.glsl index dd05fb1bde69..761d24e41793 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_funcs.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mm_funcs.glsl @@ -17,701 +17,777 @@ void store_a(uint m, uint k_pair, FLOAT_TYPEV2 value) { buf_a[a_shmem_index(m, k_pair)] = value; } +#if defined(DATA_A_ROCMFPX_FP3) +int32_t rocmfpx_mm_fp3_pack4_window(uint ib, uint idx) { + const uint bit_pos = idx * 3u; + const uint byte_pos = bit_pos >> 3u; + const uint sh = bit_pos & 7u; + uint bits = uint(data_a[ib].qs[byte_pos]) | + (uint(data_a[ib].qs[byte_pos + 1u]) << 8); + if (sh > 4u) { + bits |= uint(data_a[ib].qs[byte_pos + 2u]) << 16; + } + bits = (bits >> sh) & 0xFFFu; + return pack32(i8vec4(kvalues_rocmfpx_fp3_const[ bits & 7u], + kvalues_rocmfpx_fp3_const[(bits >> 3) & 7u], + kvalues_rocmfpx_fp3_const[(bits >> 6) & 7u], + kvalues_rocmfpx_fp3_const[(bits >> 9) & 7u])); +} + +vec4 rocmfpx_mm_fp3_vec4(uint ib, uint idx) { + const float d = ue4m3_to_fp32(data_a[ib].e[idx >= 16u ? 1u : 0u]); + return vec4(unpack8(rocmfpx_mm_fp3_pack4_window(ib, idx))) * d; +} +#endif + +#if defined(DATA_A_ROCMFPX_FP6) +uint rocmfpx_mm_fp6_get_bits(uint ib, uint idx) { + const uint bit_pos = idx * 6u; + const uint byte_pos = bit_pos >> 3u; + const uint sh = bit_pos & 7u; + uint bits = uint(data_a[ib].qs[byte_pos]); + if (sh > 2u) { + bits |= uint(data_a[ib].qs[byte_pos + 1u]) << 8; + } + return (bits >> sh) & 0x3Fu; +} + +float rocmfpx_mm_fp6_value(uint ib, uint idx) { + const float d = ue4m3_to_fp32(data_a[ib].e[idx >= 16u ? 1u : 0u]); + return float(rocmfpx_fp6_decode_code(rocmfpx_mm_fp6_get_bits(ib, idx))) * d; +} + +vec4 rocmfpx_mm_fp6_vec4(uint ib, uint idx) { + return vec4(rocmfpx_mm_fp6_value(ib, idx + 0u), + rocmfpx_mm_fp6_value(ib, idx + 1u), + rocmfpx_mm_fp6_value(ib, idx + 2u), + rocmfpx_mm_fp6_value(ib, idx + 3u)); +} +#endif + void load_a_to_shmem(const uint pos_a, const uint row, const uint col, const uint idx_m, const uint block, const uint end_k) { #if defined(DATA_A_F32) || defined(DATA_A_F16) #if LOAD_VEC_A == 8 - if (ALIGNED != 0) { - const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint k_pair = row * LOAD_VEC_A / 2; - FLOAT_TYPEV8 aa = FLOAT_TYPEV8(data_a[idx]); - store_a(col, k_pair, aa[0].xy); - store_a(col, k_pair + 1, aa[0].zw); - store_a(col, k_pair + 2, aa[1].xy); - store_a(col, k_pair + 3, aa[1].zw); - return; - } + if (ALIGNED != 0) { + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; + const uint k_pair = row * LOAD_VEC_A / 2; + FLOAT_TYPEV8 aa = FLOAT_TYPEV8(data_a[idx]); + store_a(col, k_pair, aa[0].xy); + store_a(col, k_pair + 1, aa[0].zw); + store_a(col, k_pair + 2, aa[1].xy); + store_a(col, k_pair + 3, aa[1].zw); + return; + } #elif LOAD_VEC_A == 4 - if (ALIGNED != 0) { - const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint k_pair = row * LOAD_VEC_A / 2; - FLOAT_TYPEV4 aa = FLOAT_TYPEV4(data_a[idx]); - store_a(col, k_pair, aa.xy); - store_a(col, k_pair + 1, aa.zw); - return; - } + if (ALIGNED != 0) { + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; + const uint k_pair = row * LOAD_VEC_A / 2; + FLOAT_TYPEV4 aa = FLOAT_TYPEV4(data_a[idx]); + store_a(col, k_pair, aa.xy); + store_a(col, k_pair + 1, aa.zw); + return; + } #endif - const uint idx = pos_a + col * p.stride_a + row * 2; - if (idx_m < p.M && block + row * 2 + 1 < end_k) { - store_a(col, row, FLOAT_TYPEV2(data_a_scalar[idx], - data_a_scalar[idx + 1])); - } else if (idx_m < p.M && block + row * 2 < end_k) { - store_a(col, row, FLOAT_TYPEV2(data_a_scalar[idx], 0.0f)); - } else { - store_a(col, row, FLOAT_TYPEV2(0.0f)); - } + const uint idx = pos_a + col * p.stride_a + row * 2; + if (idx_m < p.M && block + row * 2 + 1 < end_k) { + store_a(col, row, FLOAT_TYPEV2(data_a_scalar[idx], + data_a_scalar[idx + 1])); + } else if (idx_m < p.M && block + row * 2 < end_k) { + store_a(col, row, FLOAT_TYPEV2(data_a_scalar[idx], 0.0f)); + } else { + store_a(col, row, FLOAT_TYPEV2(0.0f)); + } #elif defined(DATA_A_BF16) #if LOAD_VEC_A == 4 - if (ALIGNED != 0) { - const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint k_pair = row * LOAD_VEC_A / 2; - FLOAT_TYPEV4 aa = FLOAT_TYPEV4(TO_FLOAT_TYPE(data_a[idx])); - store_a(col, k_pair, aa.xy); - store_a(col, k_pair + 1, aa.zw); - return; - } + if (ALIGNED != 0) { + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; + const uint k_pair = row * LOAD_VEC_A / 2; + FLOAT_TYPEV4 aa = FLOAT_TYPEV4(TO_FLOAT_TYPE(data_a[idx])); + store_a(col, k_pair, aa.xy); + store_a(col, k_pair + 1, aa.zw); + return; + } #endif - const uint idx = pos_a + col * p.stride_a + row * 2; - if (idx_m < p.M && block + row * 2 + 1 < end_k) { - store_a(col, row, FLOAT_TYPEV2(TO_FLOAT_TYPE(data_a_scalar[idx]), - TO_FLOAT_TYPE(data_a_scalar[idx + 1]))); - } else if (idx_m < p.M && block + row * 2 < end_k) { - store_a(col, row, FLOAT_TYPEV2(TO_FLOAT_TYPE(data_a_scalar[idx]), 0.0f)); - } else { - store_a(col, row, FLOAT_TYPEV2(0.0f)); - } -#elif defined(DATA_A_IQ1_S) - const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint k_pair = row * LOAD_VEC_A / 2; + const uint idx = pos_a + col * p.stride_a + row * 2; + if (idx_m < p.M && block + row * 2 + 1 < end_k) { + store_a(col, row, FLOAT_TYPEV2(TO_FLOAT_TYPE(data_a_scalar[idx]), + TO_FLOAT_TYPE(data_a_scalar[idx + 1]))); + } else if (idx_m < p.M && block + row * 2 < end_k) { + store_a(col, row, FLOAT_TYPEV2(TO_FLOAT_TYPE(data_a_scalar[idx]), 0.0f)); + } else { + store_a(col, row, FLOAT_TYPEV2(0.0f)); + } +#elif defined(DATA_A_Q4_0) + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; + + const uint ib = idx / 4; + const uint iqs = idx & 0x03; + + const float d = float(data_a_packed16[ib].d); + const uint vui = uint(data_a_packed16[ib].qs[2*iqs]) | (uint(data_a_packed16[ib].qs[2*iqs + 1]) << 16); + const vec4 v0 = (vec4(unpack8(vui & 0x0F0F0F0F)) - 8.0f) * d; + const vec4 v1 = (vec4(unpack8((vui >> 4) & 0x0F0F0F0F)) - 8.0f) * d; + + const uint k_pair = row * LOAD_VEC_A / 4; + store_a(col, k_pair, FLOAT_TYPEV2(v0.xy)); + store_a(col, k_pair + 1, FLOAT_TYPEV2(v0.zw)); + store_a(col, k_pair + 8, FLOAT_TYPEV2(v1.xy)); + store_a(col, k_pair + 9, FLOAT_TYPEV2(v1.zw)); +#elif defined(DATA_A_Q4_1) + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; + + const uint ib = idx / 4; + const uint iqs = idx & 0x03; + + const vec2 dm = vec2(data_a_packed32[ib].dm); + const uint vui = data_a_packed32[ib].qs[iqs]; + const vec4 v0 = vec4(unpack8(vui & 0x0F0F0F0F)) * dm.x + dm.y; + const vec4 v1 = vec4(unpack8((vui >> 4) & 0x0F0F0F0F)) * dm.x + dm.y; + + const uint k_pair = row * LOAD_VEC_A / 4; + store_a(col, k_pair, FLOAT_TYPEV2(v0.xy)); + store_a(col, k_pair + 1, FLOAT_TYPEV2(v0.zw)); + store_a(col, k_pair + 8, FLOAT_TYPEV2(v1.xy)); + store_a(col, k_pair + 9, FLOAT_TYPEV2(v1.zw)); +#elif defined(DATA_A_Q5_0) + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; + + const uint ib = idx / 8; + const uint iqs = idx & 0x07; + + const float d = float(data_a_packed16[ib].d); + const uint uint_qh = uint(data_a_packed16[ib].qh[1]) << 16 | uint(data_a_packed16[ib].qh[0]); + const ivec2 qh0 = ivec2(((uint_qh >> 2*iqs) << 4) & 0x10, (uint_qh >> (2*iqs + 12)) & 0x10); + const ivec2 qh1 = ivec2(((uint_qh >> (2*iqs + 1)) << 4) & 0x10, (uint_qh >> (2*iqs + 13)) & 0x10); + + const uint vui = uint(data_a_packed16[ib].qs[iqs]); + const vec4 v = (vec4((vui & 0xF) | qh0.x, ((vui >> 4) & 0xF) | qh0.y, ((vui >> 8) & 0xF) | qh1.x, (vui >> 12) | qh1.y) - 16.0f) * d; + store_a(col, row, FLOAT_TYPEV2(v.xz)); + store_a(col, row + 8, FLOAT_TYPEV2(v.yw)); +#elif defined(DATA_A_Q5_1) + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; + + const uint ib = idx / 4; + const uint iqs = idx & 0x03; + + const vec2 dm = vec2(data_a_packed32[ib].dm); + const uint uint_qh = data_a_packed32[ib].qh; + const uvec2 qh0 = uvec2(((uint_qh >> 4*iqs) << 4) & 0x10, (uint_qh >> (4*iqs + 12)) & 0x10); + const uvec2 qh1 = uvec2(((uint_qh >> (4*iqs + 1)) << 4) & 0x10, (uint_qh >> (4*iqs + 13)) & 0x10); + const uvec2 qh2 = uvec2(((uint_qh >> (4*iqs + 2)) << 4) & 0x10, (uint_qh >> (4*iqs + 14)) & 0x10); + const uvec2 qh3 = uvec2(((uint_qh >> (4*iqs + 3)) << 4) & 0x10, (uint_qh >> (4*iqs + 15)) & 0x10); + + const uint vui = data_a_packed32[ib].qs[iqs]; + const vec4 v0 = vec4((vui & 0xF) | qh0.x, ((vui >> 4) & 0xF) | qh0.y, ((vui >> 8) & 0xF) | qh1.x, ((vui >> 12) & 0xF) | qh1.y) * dm.x + dm.y; + const vec4 v1 = vec4(((vui >> 16) & 0xF) | qh2.x, ((vui >> 20) & 0xF) | qh2.y, ((vui >> 24) & 0xF) | qh3.x, ((vui >> 28) & 0xF) | qh3.y) * dm.x + dm.y; + + const uint k_pair = row * LOAD_VEC_A / 4; + store_a(col, k_pair, FLOAT_TYPEV2(v0.xz)); + store_a(col, k_pair + 1, FLOAT_TYPEV2(v1.xz)); + store_a(col, k_pair + 8, FLOAT_TYPEV2(v0.yw)); + store_a(col, k_pair + 9, FLOAT_TYPEV2(v1.yw)); +#elif defined(DATA_A_Q8_0) + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; + + const uint ib = idx / 8; + const uint iqs = idx & 0x07; + + const float d = float(data_a_packed16[ib].d); + const i8vec2 v0 = unpack8(int32_t(data_a_packed16[ib].qs[2*iqs])).xy; // vec4 used due to #12147 + const i8vec2 v1 = unpack8(int32_t(data_a_packed16[ib].qs[2*iqs + 1])).xy; + const vec4 v = vec4(v0.x, v0.y, v1.x, v1.y) * d; + + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, FLOAT_TYPEV2(v.xy)); + store_a(col, k_pair + 1, FLOAT_TYPEV2(v.zw)); +#elif defined(DATA_A_Q1_0) + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; + + const uint ib = idx / 16; + const uint iqs = idx & 0xfu; + + const float d = float(data_a[ib].d); + const uint bits = uint(data_a[ib].qs[iqs]); + + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, FLOAT_TYPEV2((bits & 0x01u) != 0u ? d : -d, (bits & 0x02u) != 0u ? d : -d)); + store_a(col, k_pair + 1, FLOAT_TYPEV2((bits & 0x04u) != 0u ? d : -d, (bits & 0x08u) != 0u ? d : -d)); + store_a(col, k_pair + 2, FLOAT_TYPEV2((bits & 0x10u) != 0u ? d : -d, (bits & 0x20u) != 0u ? d : -d)); + store_a(col, k_pair + 3, FLOAT_TYPEV2((bits & 0x40u) != 0u ? d : -d, (bits & 0x80u) != 0u ? d : -d)); +#elif defined(DATA_A_Q2_0) + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; + + const uint ib = idx / 16; + const uint iqs = idx & 0xfu; + + const FLOAT_TYPE d = FLOAT_TYPE(data_a[ib].d); + const uint bits = uint(data_a[ib].qs[iqs]); + + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, d * (FLOAT_TYPEV2(bits & 3u, (bits >> 2u) & 3u) - FLOAT_TYPEV2(1.0f))); + store_a(col, k_pair + 1, d * (FLOAT_TYPEV2((bits >> 4u) & 3u, bits >> 6u) - FLOAT_TYPEV2(1.0f))); +#elif defined(DATA_A_Q2_K) + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; + + const uint ib = idx / 64; // 4 values per idx + const uint iqs = (idx % 64) * 2; // 0,2,4..126 + + const uint qsi = (iqs / 64) * 16 + (iqs % 16); // 0..15 + const uint scalesi = iqs / 8; // 0..15 + const uint qsshift = ((iqs % 64) / 16) * 2; // 0,2,4,6 + + const vec4 qs = vec4(unpack8((data_a_packed32[ib].qs[qsi / 2] >> qsshift) & 0x03030303)); + const uint scales = data_a[ib].scales[scalesi]; + const vec2 dm = vec2(data_a[ib].dm); + + const vec4 v = dm.x * float(scales & 0xF) * qs - dm.y * float(scales >> 4); + + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, FLOAT_TYPEV2(v.xy)); + store_a(col, k_pair + 1, FLOAT_TYPEV2(v.zw)); +#elif defined(DATA_A_TQ1_0) + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; + + const uint ib = idx / 128; // 2 values per idx + const uint iqs = (idx % 128) * 2; // element 0,2,4..254 + + const float d = float(data_a[ib].d); + vec2 v; + for (uint kk = 0u; kk < 2u; ++kk) { + const uint e = iqs + kk; + const uint bidx = tq1_0_byte_of(e); + const uint qbyte = uint(bidx < 48u ? data_a[ib].qs[bidx] + : data_a[ib].qh[bidx - 48u]); + v[kk] = d * (float(tq1_0_trit(qbyte, tq1_0_digit_of(e))) - 1.0); + } + + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, FLOAT_TYPEV2(v.xy)); +#elif defined(DATA_A_TQ2_0) + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; + + const uint ib = idx / 128; // 2 values per idx + const uint iqs = (idx % 128) * 2; // elem 0,2,4..254 + + const uint qsi = (iqs / 128) * 32 + (iqs % 32); // byte pair start + const uint shift = 2 * ((iqs % 128) / 32); // 0,2,4,6 + + const uvec2 qs = uvec2(data_a[ib].qs[qsi], data_a[ib].qs[qsi + 1]); + const float d = float(data_a[ib].d); + + const vec2 v = d * (vec2((qs >> shift) & 3) - 1.0); + + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, FLOAT_TYPEV2(v.xy)); +#elif defined(DATA_A_Q3_K) + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; + + const uint ib = idx / 128; // 2 values per idx + const uint iqs = idx % 128; // 0..127 + + const uint n = iqs / 64; // 0,1 + const uint qsi = n * 32 + (iqs % 16) * 2; // 0,2,4..62 + const uint hmi = (iqs % 16) * 2; // 0,2,4..30 + const uint j = (iqs % 64) / 4; // 0..3 + const uint is = iqs / 8; // 0..15 + const uint halfsplit = ((iqs % 64) / 16); // 0,1,2,3 + const uint qsshift = halfsplit * 2; // 0,2,4,6 + + const int8_t us = int8_t(((data_a[ib].scales[is % 8] >> (4 * int(is / 8))) & 0xF) + | (((data_a[ib].scales[8 + (is % 4)] >> (2 * int(is / 4))) & 3) << 4)); + const float dl = float(data_a[ib].d) * float(us - 32); + + const vec2 qs = vec2(unpack8((uint(data_a_packed16[ib].qs[qsi / 2]) >> qsshift) & 0x0303).xy); + const vec2 hm = vec2(unpack8(((uint(data_a_packed16[ib].hmask[hmi / 2]) >> (4 * n + halfsplit)) & 0x0101 ^ 0x0101) << 2).xy); + + store_a(col, row * LOAD_VEC_A / 2, FLOAT_TYPEV2(dl * (qs.x - hm.x), + dl * (qs.y - hm.y))); +#elif defined(DATA_A_Q4_K) + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; + + const uint ib = idx / 64; // 4 values per idx + const uint iqs = (idx % 64) * 2; // 0,2,4..126 + + const uint n = iqs / 32; // 0,1,2,3 + const uint b = (iqs % 32) / 16; // 0,1 + const uint is = 2 * n + b; // 0..7 + const uint qsi = n * 32 + (iqs % 16) * 2; // 0,2,4..126 + + const vec2 loadd = vec2(data_a[ib].dm); + + const uvec3 scales = uvec3(data_a_packed32[ib].scales[0], + data_a_packed32[ib].scales[1], + data_a_packed32[ib].scales[2]); + const uint scalesoffs = (is & 3) * 8; + + const uint scidx0 = (is < 4) ? 0 : 2; + const uint scidxshift0 = scalesoffs; + const uint scidxshift1 = (is < 4) ? scalesoffs : scalesoffs + 2; + const uint mbidx0 = (is < 4) ? 1 : 2; + const uint mbidxshift0 = (is < 4) ? scalesoffs : scalesoffs + 4; + const uint mbidxshift1 = (is < 4) ? scalesoffs : scalesoffs + 2; + + const uint8_t sc = uint8_t(((scales[scidx0] >> scidxshift0) & 0xF) | ((scales[0] >> scidxshift1) & 0x30)); + const uint8_t mbyte = uint8_t(((scales[mbidx0] >> mbidxshift0) & 0xF) | ((scales[1] >> mbidxshift1) & 0x30)); + + const float d = loadd.x * sc; + const float m = -loadd.y * mbyte; + + const vec4 q = vec4(unpack8((data_a_packed32[ib].qs[qsi / 4] >> (b * 4)) & 0x0F0F0F0F)); + + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, FLOAT_TYPEV2(fma(d, q.x, m), fma(d, q.y, m))); + store_a(col, k_pair + 1, FLOAT_TYPEV2(fma(d, q.z, m), fma(d, q.w, m))); +#elif defined(DATA_A_Q5_K) + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; + + const uint ib = idx / 64; // 4 values per idx + const uint iqs = (idx % 64) * 2; // 0,2,4..126 + + const uint n = iqs / 32; // 0,1,2,3 + const uint b = (iqs % 32) / 16; // 0,1 + const uint is = 2 * n + b; // 0..7 + const uint qsi = n * 32 + (iqs % 16) * 2; // 0,2,4..126 + const uint qhi = (iqs % 16) * 2; // 0,2,4..30 + + const vec2 loadd = vec2(data_a[ib].dm); + + const uvec3 scales = uvec3(data_a_packed32[ib].scales[0], + data_a_packed32[ib].scales[1], + data_a_packed32[ib].scales[2]); + const uint scalesoffs = (is & 3) * 8; + + const uint scidx0 = (is < 4) ? 0 : 2; + const uint scidxshift0 = scalesoffs; + const uint scidxshift1 = (is < 4) ? scalesoffs : scalesoffs + 2; + const uint mbidx0 = (is < 4) ? 1 : 2; + const uint mbidxshift0 = (is < 4) ? scalesoffs : scalesoffs + 4; + const uint mbidxshift1 = (is < 4) ? scalesoffs : scalesoffs + 2; - const uint ib = idx / 32; - const uint ib32 = (idx % 32) / 4; - const uint ib8 = idx % 32; + const uint8_t sc = uint8_t(((scales[scidx0] >> scidxshift0) & 0xF) | ((scales[0] >> scidxshift1) & 0x30)); + const uint8_t mbyte = uint8_t(((scales[mbidx0] >> mbidxshift0) & 0xF) | ((scales[1] >> mbidxshift1) & 0x30)); + + const float d = loadd.x * sc; + const float m = -loadd.y * mbyte; + + const uint qs = (data_a_packed32[ib].qs[qsi / 4] >> (b * 4)) & 0x0F0F0F0F; + const uint qh = ((data_a_packed32[ib].qh[qhi / 4] >> (iqs / 16)) & 0x01010101) << 4; + const vec4 q = vec4(unpack8(qs | qh)); + + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, FLOAT_TYPEV2(fma(d, q.x, m), fma(d, q.y, m))); + store_a(col, k_pair + 1, FLOAT_TYPEV2(fma(d, q.z, m), fma(d, q.w, m))); +#elif defined(DATA_A_Q6_K) + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; + + const uint ib = idx / 128; // 2 values per idx + const uint iqs = idx % 128; // 0..127 + + const uint n = iqs / 64; // 0,1 + const uint b = ((iqs % 64) / 32) * 4; // 0,4 + const uint is_b = (iqs % 16) / 8; // 0,1 + const uint qhshift = ((iqs % 64) / 16) * 2; // 0,2,4,6 + const uint is = 8 * n + qhshift + is_b; // 0..15 + const uint qsi = n * 32 + (iqs % 32); // 0..63 + const uint qhi = n * 16 + (iqs % 16); // 0..31 + + const float dscale = float(data_a[ib].d) * float(data_a[ib].scales[is]); - const float d = float(data_a[ib].d); - const uint qh = data_a[ib].qh[ib32]; - const uint qs = data_a[ib].qs[ib8]; - const float dl = d * (2 * bitfieldExtract(qh, 12, 3) + 1); - const float delta = ((qh & 0x8000) != 0) ? -IQ1S_DELTA : IQ1S_DELTA; - const int16_t grid = int16_t(iq1s_grid[qs | (bitfieldExtract(qh, 3 * int(ib8 & 3), 3) << 8)]); + const uint ql = (uint(data_a_packed16[ib].ql[qsi]) >> b) & 0x0F0F; + const uint qh = (uint(data_a_packed16[ib].qh[qhi]) >> qhshift) & 0x0303; + const vec2 q = (vec2(unpack8(ql | (qh << 4)).xy) - 32) * dscale; - [[unroll]] for (int k = 0; k < 4; ++k) { - store_a(col, k_pair + k, FLOAT_TYPEV2(dl * (bitfieldExtract(grid, 4 * k , 2) + delta), - dl * (bitfieldExtract(grid, 4 * k + 2, 2) + delta))); - - } + store_a(col, row * LOAD_VEC_A / 2, FLOAT_TYPEV2(q.x, q.y)); +#elif defined(DATA_A_IQ1_S) + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; + + const uint ib = idx / 32; // 8 values per idx + const uint ib32 = (idx % 32) / 4; // 0..7 + const uint ib8 = idx % 32; + + const float d = float(data_a[ib].d); + const uint qh = data_a[ib].qh[ib32]; + const uint qs = data_a[ib].qs[ib8]; + const float dl = d * (2 * bitfieldExtract(qh, 12, 3) + 1); + const float delta = ((qh & 0x8000) != 0) ? -IQ1S_DELTA : IQ1S_DELTA; + const int16_t grid = int16_t(iq1s_grid[qs | (bitfieldExtract(qh, 3 * int(ib8 & 3), 3) << 8)]); + + const uint k_pair = row * LOAD_VEC_A / 2; + [[unroll]] for (int k = 0; k < 4; ++k) { + store_a(col, k_pair + k, FLOAT_TYPEV2(dl * (bitfieldExtract(grid, 4 * k , 2) + delta), + dl * (bitfieldExtract(grid, 4 * k + 2, 2) + delta))); + } #elif defined(DATA_A_IQ1_M) - const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint k_pair = row * LOAD_VEC_A / 2; - - const uint ib = idx / 32; - const uint ib8 = idx % 32; - const uint ib16 = ib8 / 2; - - const uint16_t[4] scales = data_a[ib].scales; - const u16vec4 s = u16vec4(scales[0], scales[1], scales[2], scales[3]) >> 12; - const float d = float(unpackHalf2x16(s.x | (s.y << 4) | (s.z << 8) | (s.w << 12)).x); - const uint sc = scales[ib8 / 8]; - const uint qs = data_a[ib].qs[ib8]; - const uint qh = data_a[ib].qh[ib16] >> (4 * (ib8 & 1)); - const float dl = d * (2 * bitfieldExtract(sc, 3 * int(ib16 & 3), 3) + 1); - const float delta = ((qh & 8) != 0) ? -IQ1M_DELTA : IQ1M_DELTA; - const int16_t grid = int16_t(iq1s_grid[qs | ((qh & 7) << 8)]); - - [[unroll]] for (int k = 0; k < 4; ++k) { - store_a(col, k_pair + k, FLOAT_TYPEV2(dl * (bitfieldExtract(grid, 4 * k , 2) + delta), - dl * (bitfieldExtract(grid, 4 * k + 2, 2) + delta))); - - } + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; + + const uint ib = idx / 32; // 8 values per idx + const uint ib8 = idx % 32; + const uint ib16 = ib8 / 2; + + const uint16_t[4] scales = data_a[ib].scales; + const u16vec4 s = u16vec4(scales[0], scales[1], scales[2], scales[3]) >> 12; + const float d = float(unpackHalf2x16(s.x | (s.y << 4) | (s.z << 8) | (s.w << 12)).x); + const uint sc = scales[ib8 / 8]; + const uint qs = data_a[ib].qs[ib8]; + const uint qh = data_a[ib].qh[ib16] >> (4 * (ib8 & 1)); + const float dl = d * (2 * bitfieldExtract(sc, 3 * int(ib16 & 3), 3) + 1); + const float delta = ((qh & 8) != 0) ? -IQ1M_DELTA : IQ1M_DELTA; + const int16_t grid = int16_t(iq1s_grid[qs | ((qh & 7) << 8)]); + + const uint k_pair = row * LOAD_VEC_A / 2; + [[unroll]] for (int k = 0; k < 4; ++k) { + store_a(col, k_pair + k, FLOAT_TYPEV2(dl * (bitfieldExtract(grid, 4 * k , 2) + delta), + dl * (bitfieldExtract(grid, 4 * k + 2, 2) + delta))); + } #elif defined(DATA_A_IQ2_XXS) - const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint k_pair = row * LOAD_VEC_A / 2; - - const uint ib = idx / 32; - const uint ib32 = (idx % 32) / 4; - const uint ib8 = idx % 4; - - const float d = float(data_a[ib].d); - const uint qs = data_a[ib].qs[8 * ib32 + ib8]; - const uint signs = pack32(u8vec4( - data_a[ib].qs[8*ib32 + 4], - data_a[ib].qs[8*ib32 + 5], - data_a[ib].qs[8*ib32 + 6], - data_a[ib].qs[8*ib32 + 7] - )); - const FLOAT_TYPE db = FLOAT_TYPE(d * 0.25 * (0.5 + (signs >> 28))); - const uint32_t sign7 = bitfieldExtract(signs, 7 * int(ib8), 7); - const uint sign = sign7 | (bitCount(sign7) << 7); - const uvec2 grid = iq2xxs_grid[qs]; - const vec4 grid0 = vec4(unpack8(grid.x)); - const vec4 grid1 = vec4(unpack8(grid.y)); - - store_a(col, k_pair, db * FLOAT_TYPEV2((sign & 1) != 0 ? -grid0.x : grid0.x, - (sign & 2) != 0 ? -grid0.y : grid0.y)); - - store_a(col, k_pair + 1, db * FLOAT_TYPEV2((sign & 4) != 0 ? -grid0.z : grid0.z, - (sign & 8) != 0 ? -grid0.w : grid0.w)); - - store_a(col, k_pair + 2, db * FLOAT_TYPEV2((sign & 16) != 0 ? -grid1.x : grid1.x, - (sign & 32) != 0 ? -grid1.y : grid1.y)); - - store_a(col, k_pair + 3, db * FLOAT_TYPEV2((sign & 64) != 0 ? -grid1.z : grid1.z, - (sign & 128) != 0 ? -grid1.w : grid1.w)); - + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; + + const uint ib = idx / 32; // 8 values per idx + const uint ib32 = (idx % 32) / 4; // 0..7 + const uint ib8 = idx % 4; + + const float d = float(data_a[ib].d); + const uint qs = data_a[ib].qs[8 * ib32 + ib8]; + const uint signs = pack32(u8vec4( + data_a[ib].qs[8*ib32 + 4], + data_a[ib].qs[8*ib32 + 5], + data_a[ib].qs[8*ib32 + 6], + data_a[ib].qs[8*ib32 + 7] + )); + const FLOAT_TYPE db = FLOAT_TYPE(d * 0.25 * (0.5 + (signs >> 28))); + const uint32_t sign7 = bitfieldExtract(signs, 7 * int(ib8), 7); + const uint sign = sign7 | (bitCount(sign7) << 7); + const uvec2 grid = iq2xxs_grid[qs]; + const vec4 grid0 = vec4(unpack8(grid.x)); + const vec4 grid1 = vec4(unpack8(grid.y)); + + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, db * FLOAT_TYPEV2((sign & 1) != 0 ? -grid0.x : grid0.x, + (sign & 2) != 0 ? -grid0.y : grid0.y)); + store_a(col, k_pair + 1, db * FLOAT_TYPEV2((sign & 4) != 0 ? -grid0.z : grid0.z, + (sign & 8) != 0 ? -grid0.w : grid0.w)); + store_a(col, k_pair + 2, db * FLOAT_TYPEV2((sign & 16) != 0 ? -grid1.x : grid1.x, + (sign & 32) != 0 ? -grid1.y : grid1.y)); + store_a(col, k_pair + 3, db * FLOAT_TYPEV2((sign & 64) != 0 ? -grid1.z : grid1.z, + (sign & 128) != 0 ? -grid1.w : grid1.w)); #elif defined(DATA_A_IQ2_XS) - const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint k_pair = row * LOAD_VEC_A / 2; - - const uint ib = idx / 32; - const uint ib32 = (idx % 32) / 4; - const uint ib8 = idx % 4; - - const float d = float(data_a[ib].d); - const uint scale = (data_a[ib].scales[ib32] >> (2 * (ib8 & 2))) & 0xf; - const FLOAT_TYPE db = FLOAT_TYPE(d * 0.25 * (0.5 + scale)); - const uint qs = data_a[ib].qs[4 * ib32 + ib8]; - const uint sign7 = qs >> 9; - const uint sign = sign7 | (bitCount(sign7) << 7); - const uvec2 grid = iq2xs_grid[qs & 511]; - const vec4 grid0 = vec4(unpack8(grid.x)); - const vec4 grid1 = vec4(unpack8(grid.y)); - - store_a(col, k_pair, db * FLOAT_TYPEV2((sign & 1) != 0 ? -grid0.x : grid0.x, - (sign & 2) != 0 ? -grid0.y : grid0.y)); - - store_a(col, k_pair + 1, db * FLOAT_TYPEV2((sign & 4) != 0 ? -grid0.z : grid0.z, - (sign & 8) != 0 ? -grid0.w : grid0.w)); - - store_a(col, k_pair + 2, db * FLOAT_TYPEV2((sign & 16) != 0 ? -grid1.x : grid1.x, - (sign & 32) != 0 ? -grid1.y : grid1.y)); - - store_a(col, k_pair + 3, db * FLOAT_TYPEV2((sign & 64) != 0 ? -grid1.z : grid1.z, - (sign & 128) != 0 ? -grid1.w : grid1.w)); - + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; + + const uint ib = idx / 32; // 8 values per idx + const uint ib32 = (idx % 32) / 4; // 0..7 + const uint ib8 = idx % 4; // 0..3 + + const float d = float(data_a[ib].d); + const uint scale = (data_a[ib].scales[ib32] >> (2 * (ib8 & 2))) & 0xf; + const FLOAT_TYPE db = FLOAT_TYPE(d * 0.25 * (0.5 + scale)); + const uint qs = data_a[ib].qs[4 * ib32 + ib8]; + const uint sign7 = qs >> 9; + const uint sign = sign7 | (bitCount(sign7) << 7); + const uvec2 grid = iq2xs_grid[qs & 511]; + const vec4 grid0 = vec4(unpack8(grid.x)); + const vec4 grid1 = vec4(unpack8(grid.y)); + + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, db * FLOAT_TYPEV2((sign & 1) != 0 ? -grid0.x : grid0.x, + (sign & 2) != 0 ? -grid0.y : grid0.y)); + store_a(col, k_pair + 1, db * FLOAT_TYPEV2((sign & 4) != 0 ? -grid0.z : grid0.z, + (sign & 8) != 0 ? -grid0.w : grid0.w)); + store_a(col, k_pair + 2, db * FLOAT_TYPEV2((sign & 16) != 0 ? -grid1.x : grid1.x, + (sign & 32) != 0 ? -grid1.y : grid1.y)); + store_a(col, k_pair + 3, db * FLOAT_TYPEV2((sign & 64) != 0 ? -grid1.z : grid1.z, + (sign & 128) != 0 ? -grid1.w : grid1.w)); #elif defined(DATA_A_IQ2_S) - const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint k_pair = row * LOAD_VEC_A / 2; - - const uint ib = idx / 32; - const uint ib8 = idx % 32; - const uint ib32 = ib8 / 4; - - const uint scale = (data_a[ib].scales[ib32] >> (2 * (ib8 & 2))) & 0xf; - const uint qs = data_a[ib].qs[ib8]; - const uint qh = data_a[ib].qh[ib32]; - const uint qhshift = 2 * (ib8 % 4); - const uint sign = data_a[ib].qs[QUANT_K_IQ2_S / 8 + ib8]; - - const float d = float(data_a[ib].d); - const FLOAT_TYPE db = FLOAT_TYPE(d * 0.25 * (0.5 + scale)); - const uvec2 grid = iq2s_grid[qs | ((qh << (8 - qhshift)) & 0x300)]; - const vec4 grid0 = vec4(unpack8(grid.x)); - const vec4 grid1 = vec4(unpack8(grid.y)); - - store_a(col, k_pair, db * FLOAT_TYPEV2((sign & 1) != 0 ? -grid0.x : grid0.x, - (sign & 2) != 0 ? -grid0.y : grid0.y)); - - store_a(col, k_pair + 1, db * FLOAT_TYPEV2((sign & 4) != 0 ? -grid0.z : grid0.z, - (sign & 8) != 0 ? -grid0.w : grid0.w)); - - store_a(col, k_pair + 2, db * FLOAT_TYPEV2((sign & 16) != 0 ? -grid1.x : grid1.x, - (sign & 32) != 0 ? -grid1.y : grid1.y)); - - store_a(col, k_pair + 3, db * FLOAT_TYPEV2((sign & 64) != 0 ? -grid1.z : grid1.z, - (sign & 128) != 0 ? -grid1.w : grid1.w)); - + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; + + const uint ib = idx / 32; // 8 values per idx + const uint ib8 = idx % 32; // 0..31 + const uint ib32 = ib8 / 4; // 0..7 + + const uint scale = (data_a[ib].scales[ib32] >> (2 * (ib8 & 2))) & 0xf; + const uint qs = data_a[ib].qs[ib8]; + const uint qh = data_a[ib].qh[ib32]; + const uint qhshift = 2 * (ib8 % 4); + const uint sign = data_a[ib].qs[QUANT_K / 8 + ib8]; + + const float d = float(data_a[ib].d); + const FLOAT_TYPE db = FLOAT_TYPE(d * 0.25 * (0.5 + scale)); + const uvec2 grid = iq2s_grid[qs | ((qh << (8 - qhshift)) & 0x300)]; + const vec4 grid0 = vec4(unpack8(grid.x)); + const vec4 grid1 = vec4(unpack8(grid.y)); + + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, db * FLOAT_TYPEV2((sign & 1) != 0 ? -grid0.x : grid0.x, + (sign & 2) != 0 ? -grid0.y : grid0.y)); + store_a(col, k_pair + 1, db * FLOAT_TYPEV2((sign & 4) != 0 ? -grid0.z : grid0.z, + (sign & 8) != 0 ? -grid0.w : grid0.w)); + store_a(col, k_pair + 2, db * FLOAT_TYPEV2((sign & 16) != 0 ? -grid1.x : grid1.x, + (sign & 32) != 0 ? -grid1.y : grid1.y)); + store_a(col, k_pair + 3, db * FLOAT_TYPEV2((sign & 64) != 0 ? -grid1.z : grid1.z, + (sign & 128) != 0 ? -grid1.w : grid1.w)); #elif defined(DATA_A_IQ3_XXS) - const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint k_pair = row * LOAD_VEC_A / 2; - - const uint ib = idx / 64; - const uint iqs = idx % 64; - const uint is = QUANT_K_IQ3_XXS / 4 + 4 * (iqs / 8); - - const float d = float(data_a[ib].d); - const uint qs = data_a[ib].qs[iqs]; - const uint signs = pack32(u16vec2( - data_a_packed16[ib].qs[is/2], - data_a_packed16[ib].qs[is/2+1] - )); - const float db = d * 0.5 * (0.5 + (signs >> 28)); - const uint32_t sign7 = bitfieldExtract(signs, 7 * (int(iqs / 2) % 4), 7); - const uint sign = (sign7 | (bitCount(sign7) << 7)) >> (4 * (idx % 2)); - const uint grid = iq3xxs_grid[qs]; - const vec4 v = db * vec4(unpack8(grid)); - - store_a(col, k_pair, FLOAT_TYPEV2((sign & 1) != 0 ? -v.x : v.x, - (sign & 2) != 0 ? -v.y : v.y)); - - store_a(col, k_pair + 1, FLOAT_TYPEV2((sign & 4) != 0 ? -v.z : v.z, - (sign & 8) != 0 ? -v.w : v.w)); - + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; + + const uint ib = idx / 64; // 4 values per idx + const uint iqs = idx % 64; // 0..63 + const uint is = QUANT_K / 4 + 4 * (iqs / 8); // 8 values + + const float d = float(data_a[ib].d); + const uint qs = data_a[ib].qs[iqs]; + const uint signs = pack32(u16vec2( + data_a_packed16[ib].qs[is/2], + data_a_packed16[ib].qs[is/2+1] + )); + const float db = d * 0.5 * (0.5 + (signs >> 28)); + const uint32_t sign7 = bitfieldExtract(signs, 7 * (int(iqs / 2) % 4), 7); + const uint sign = (sign7 | (bitCount(sign7) << 7)) >> (4 * (idx % 2)); + const uint grid = iq3xxs_grid[qs]; + const vec4 v = db * vec4(unpack8(grid)); + + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, FLOAT_TYPEV2((sign & 1) != 0 ? -v.x : v.x, + (sign & 2) != 0 ? -v.y : v.y)); + store_a(col, k_pair + 1, FLOAT_TYPEV2((sign & 4) != 0 ? -v.z : v.z, + (sign & 8) != 0 ? -v.w : v.w)); #elif defined(DATA_A_IQ3_S) - const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint k_pair = row * LOAD_VEC_A / 2; - - const uint ib = idx / 64; - const uint iqs = idx % 64; - const uint iqh = iqs / 8; - - const float d = float(data_a[ib].d); - const uint qs = data_a[ib].qs[iqs]; - const uint qh = data_a[ib].qh[iqh]; - const int8_t sign = int8_t(data_a[ib].signs[iqs / 2] >> (4 * (idx % 2))); - const uint scale = data_a[ib].scales[iqs / 16]; - const i8vec2 sign01 = i8vec2(1 - (2 & i8vec2(sign << 1, sign))); - const float db = d * (1 + 2 * ((scale >> (4 * (iqh & 1))) & 0xf)); - const uint32_t grid = iq3s_grid[qs | ((qh << (8 - (iqs % 8))) & 256)]; - const vec4 v = db * vec4(unpack8(grid)); - - store_a(col, k_pair, FLOAT_TYPEV2((sign & 1) != 0 ? -v.x : v.x, - (sign & 2) != 0 ? -v.y : v.y)); - - store_a(col, k_pair + 1, FLOAT_TYPEV2((sign & 4) != 0 ? -v.z : v.z, - (sign & 8) != 0 ? -v.w : v.w)); - + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; + + const uint ib = idx / 64; // 4 values per idx + const uint iqs = idx % 64; // 0..63 + const uint iqh = iqs / 8; + + const float d = float(data_a[ib].d); + const uint qs = data_a[ib].qs[iqs]; + const uint qh = data_a[ib].qh[iqh]; + const int8_t sign = int8_t(data_a[ib].signs[iqs / 2] >> (4 * (idx % 2))); + const uint scale = data_a[ib].scales[iqs / 16]; + const i8vec2 sign01 = i8vec2(1 - (2 & i8vec2(sign << 1, sign))); + const float db = d * (1 + 2 * ((scale >> (4 * (iqh & 1))) & 0xf)); + const uint32_t grid = iq3s_grid[qs | ((qh << (8 - (iqs % 8))) & 256)]; + const vec4 v = db * vec4(unpack8(grid)); + + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, FLOAT_TYPEV2((sign & 1) != 0 ? -v.x : v.x, + (sign & 2) != 0 ? -v.y : v.y)); + store_a(col, k_pair + 1, FLOAT_TYPEV2((sign & 4) != 0 ? -v.z : v.z, + (sign & 8) != 0 ? -v.w : v.w)); #elif defined(DATA_A_IQ4_XS) - const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint k_pair = row * LOAD_VEC_A / 2; + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint ib = idx / 64; - const uint ib32 = (idx % 64) / 8; - const uint iq = 4 * ib32 + (idx % 4); + const uint ib = idx / 64; // 4 values per idx + const uint ib32 = (idx % 64) / 8; // 0..7 + const uint iq = 4 * ib32 + (idx % 4); - const uint sl = (data_a[ib].scales_l[ib32/2] >> (4 * (ib32 & 1))) & 0xF; - const uint sh = ((data_a[ib].scales_h) >> (2 * ib32)) & 3; - const uint qshift = idx & 4; - u8vec4 qs = unpack8((uint(data_a_packed32[ib].qs[iq]) >> qshift) & 0x0F0F0F0F); + const uint sl = (data_a[ib].scales_l[ib32/2] >> (4 * (ib32 & 1))) & 0xF; + const uint sh = ((data_a[ib].scales_h) >> (2 * ib32)) & 3; + const uint qshift = idx & 4; + u8vec4 qs = unpack8((uint(data_a_packed32[ib].qs[iq]) >> qshift) & 0x0F0F0F0F); - const float d = float(data_a[ib].d); - const vec4 v = d * float(int(sl | (sh << 4)) - 32) * vec4(kvalues_iq4nl[qs.x], kvalues_iq4nl[qs.y], kvalues_iq4nl[qs.z], kvalues_iq4nl[qs.w]); + const float d = float(data_a[ib].d); + const vec4 v = d * float(int(sl | (sh << 4)) - 32) * vec4(kvalues_iq4nl[qs.x], kvalues_iq4nl[qs.y], kvalues_iq4nl[qs.z], kvalues_iq4nl[qs.w]); - store_a(col, k_pair, FLOAT_TYPEV2(v.xy)); - store_a(col, k_pair + 1, FLOAT_TYPEV2(v.zw)); + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, FLOAT_TYPEV2(v.xy)); + store_a(col, k_pair + 1, FLOAT_TYPEV2(v.zw)); #elif defined(DATA_A_IQ4_NL) - const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint k_pair = row * LOAD_VEC_A / 4; - - const uint ib = idx / 8; - const uint iqs = idx & 0x07; + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const FLOAT_TYPE d = FLOAT_TYPE(data_a_packed16[ib].d); - const uint vui = uint(data_a_packed16[ib].qs[iqs]); + const uint ib = idx / 8; + const uint iqs = idx & 0x07; - store_a(col, k_pair, d * FLOAT_TYPEV2(kvalues_iq4nl[vui & 0xF], - kvalues_iq4nl[bitfieldExtract(vui, 8, 4)])); - - store_a(col, k_pair + 8, d * FLOAT_TYPEV2(kvalues_iq4nl[bitfieldExtract(vui, 4, 4)], - kvalues_iq4nl[vui >> 12])); + const FLOAT_TYPE d = FLOAT_TYPE(data_a_packed16[ib].d); + const uint vui = uint(data_a_packed16[ib].qs[iqs]); + const uint k_pair = row * LOAD_VEC_A / 4; + store_a(col, k_pair, d * FLOAT_TYPEV2(kvalues_iq4nl[vui & 0xF], + kvalues_iq4nl[bitfieldExtract(vui, 8, 4)])); + store_a(col, k_pair + 8, d * FLOAT_TYPEV2(kvalues_iq4nl[bitfieldExtract(vui, 4, 4)], + kvalues_iq4nl[vui >> 12])); #elif defined(DATA_A_MXFP4) - const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint k_pair = row * LOAD_VEC_A / 4; + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint ib = idx / 8; - const uint iqs = (idx & 0x07) * 2; + const uint ib = idx / 8; + const uint iqs = (idx & 0x07) * 2; - const uint vui = uint(data_a[ib].qs[iqs]); - const uint vui2 = uint(data_a[ib].qs[iqs+1]); + const uint vui = uint(data_a[ib].qs[iqs]); + const uint vui2 = uint(data_a[ib].qs[iqs+1]); #ifdef USE_OCP_FP4 - const float d = e8m0_to_fp32(data_a[ib].e); - const u8vec2 packed = u8vec2(vui, vui2); - store_a(col, k_pair, FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 0u)) * FLOAT_TYPE(d)); - store_a(col, k_pair + 8, FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 4u)) * FLOAT_TYPE(d)); + const float d = e8m0_to_fp32(data_a[ib].e); + const u8vec2 packed = u8vec2(vui, vui2); + store_a(col, row, FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 0u)) * FLOAT_TYPE(d)); + store_a(col, row + 8, FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 4u)) * FLOAT_TYPE(d)); #else - const float d = e8m0_to_fp32(data_a[ib].e) * 0.5; - store_a(col, k_pair, FLOAT_TYPEV2(kvalues_mxfp4[vui & 0xF] * d, - kvalues_mxfp4[vui2 & 0xF] * d)); + const float d = e8m0_to_fp32(data_a[ib].e) * 0.5; + store_a(col, row, FLOAT_TYPEV2(kvalues_mxfp4[vui & 0xF] * d, + kvalues_mxfp4[vui2 & 0xF] * d)); + store_a(col, row + 8, FLOAT_TYPEV2(kvalues_mxfp4[vui >> 4] * d, + kvalues_mxfp4[vui2 >> 4] * d)); +#endif +#elif defined(DATA_A_ROCMFP4) || defined(DATA_A_ROCMFP4_FAST) + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - store_a(col, k_pair + 8, FLOAT_TYPEV2(kvalues_mxfp4[vui >> 4] * d, - kvalues_mxfp4[vui2 >> 4] * d)); + const uint ib = idx / 8; + const uint iqs = (idx & 0x07) * 2; + const uint vui = uint(data_a[ib].qs[iqs]); + const uint vui2 = uint(data_a[ib].qs[iqs+1]); + + // low nibbles hold the first half of the block (scale e[0]), + // high nibbles the second half (scale e[1]) +#if defined(DATA_A_ROCMFP4) + const float d0 = ue4m3_to_fp32(data_a[ib].e[0]); + const float d1 = ue4m3_to_fp32(data_a[ib].e[1]); +#else + const float d0 = ue4m3_to_fp32(data_a[ib].e); + const float d1 = d0; #endif + store_a(col, row, FLOAT_TYPEV2(float(kvalues_rocmfp4[vui & 0xF]) * d0, + float(kvalues_rocmfp4[vui2 & 0xF]) * d0)); + store_a(col, row + 8, FLOAT_TYPEV2(float(kvalues_rocmfp4[vui >> 4]) * d1, + float(kvalues_rocmfp4[vui2 >> 4]) * d1)); +#elif defined(DATA_A_ROCMFPX_FP3) + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; + + const uint ib = idx / 8; + const uint iqs = (idx & 0x07) * 4; + const vec4 v = rocmfpx_mm_fp3_vec4(ib, iqs); + + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, FLOAT_TYPEV2(v.xy)); + store_a(col, k_pair + 1, FLOAT_TYPEV2(v.zw)); +#elif defined(DATA_A_ROCMFPX_FP6) + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; + + const uint ib = idx / 8; + const uint iqs = (idx & 0x07) * 4; + const vec4 v = rocmfpx_mm_fp6_vec4(ib, iqs); + + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, FLOAT_TYPEV2(v.xy)); + store_a(col, k_pair + 1, FLOAT_TYPEV2(v.zw)); +#elif defined(DATA_A_ROCMFPX_FP8) + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; + + const uint ib = idx / 8; + const uint iqs = (idx & 0x07) * 4; + + const float d = ue4m3_to_fp32(data_a[ib].e); + const vec4 v = vec4(float(data_a[ib].qs[iqs + 0u]), + float(data_a[ib].qs[iqs + 1u]), + float(data_a[ib].qs[iqs + 2u]), + float(data_a[ib].qs[iqs + 3u])) * d; + + const uint k_pair = row * LOAD_VEC_A / 2; + store_a(col, k_pair, FLOAT_TYPEV2(v.xy)); + store_a(col, k_pair + 1, FLOAT_TYPEV2(v.zw)); #elif defined(DATA_A_NVFP4) - const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; - const uint eff_row = (row & 3) + (row & ~3) * 2; - - const uint ib = idx / 16u; - const uint sub = (idx & 0xC) >> 2; - const uint iqs = (idx & 0xF) * 2; - const uint vui = uint(data_a[ib].qs[iqs]); - const uint vui2 = uint(data_a[ib].qs[iqs+1]); - + const uint idx = pos_a + col * p.stride_a / LOAD_VEC_A + row; + const uint ib = idx / 16u; + const uint sub = (idx & 0xC) >> 2; + const uint iqs = (idx & 0xF) * 2; + const uint vui = uint(data_a[ib].qs[iqs]); + const uint vui2 = uint(data_a[ib].qs[iqs+1]); + + // lo and hi nibbles are 8 elements apart, which doesn't quite line up with + // how the thread mapping and buf_idx calculation works for other types. + const uint eff_row = (row & 3) + (row & ~3) * 2; #ifdef USE_OCP_FP4 - const FLOAT_TYPE d = FLOAT_TYPE(ue4m3_from_bits(data_a[ib].d[sub])); - const u8vec2 packed = u8vec2(vui, vui2); - store_a(col, eff_row, FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 0u)) * d); - store_a(col, eff_row + 4, FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 4u)) * d); + const FLOAT_TYPE d = FLOAT_TYPE(ue4m3_from_bits(data_a[ib].d[sub])); + const u8vec2 packed = u8vec2(vui, vui2); + store_a(col, eff_row, FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 0u)) * d); + store_a(col, eff_row + 4, FLOAT_TYPEV2(bitcastExtractfe2m1EXT(packed, 4u)) * d); #else - const float d = ue4m3_to_fp32(data_a[ib].d[sub]) * 0.5; - store_a(col, eff_row, FLOAT_TYPEV2(kvalues_mxfp4[vui & 0xF] * d, - kvalues_mxfp4[vui2 & 0xF] * d)); - store_a(col, eff_row + 4, FLOAT_TYPEV2(kvalues_mxfp4[vui >> 4] * d, - kvalues_mxfp4[vui2 >> 4] * d)); + const float d = ue4m3_to_fp32(data_a[ib].d[sub]) * 0.5; + store_a(col, eff_row, FLOAT_TYPEV2(kvalues_mxfp4[vui & 0xF] * d, + kvalues_mxfp4[vui2 & 0xF] * d)); + store_a(col, eff_row + 4, FLOAT_TYPEV2(kvalues_mxfp4[vui >> 4] * d, + kvalues_mxfp4[vui2 >> 4] * d)); #endif -#else - if (MmTypeA == GGML_TYPE_Q4_0) { - const uint idx = pos_a + col * p.stride_a / mm_load_vec_a() + row; - const uint k_pair = row * mm_load_vec_a() / 4; - - const uint ib = idx / 4; - const uint iqs = idx & 0x03; - - const float d = float(a_q4_0_p16.data[ib].d); - const uint vui = uint(a_q4_0_p16.data[ib].qs[2*iqs]) | (uint(a_q4_0_p16.data[ib].qs[2*iqs + 1]) << 16); - const vec4 v0 = (vec4(unpack8(vui & 0x0F0F0F0F)) - 8.0f) * d; - const vec4 v1 = (vec4(unpack8((vui >> 4) & 0x0F0F0F0F)) - 8.0f) * d; - - store_a(col, k_pair, FLOAT_TYPEV2(v0.xy)); - store_a(col, k_pair + 1, FLOAT_TYPEV2(v0.zw)); - store_a(col, k_pair + 8, FLOAT_TYPEV2(v1.xy)); - store_a(col, k_pair + 9, FLOAT_TYPEV2(v1.zw)); - } else if (MmTypeA == GGML_TYPE_Q4_1) { - const uint idx = pos_a + col * p.stride_a / mm_load_vec_a() + row; - const uint k_pair = row * mm_load_vec_a() / 4; - - const uint ib = idx / 4; - const uint iqs = idx & 0x03; - - const vec2 dm = vec2(a_q4_1_p32.data[ib].dm); - const uint vui = a_q4_1_p32.data[ib].qs[iqs]; - const vec4 v0 = vec4(unpack8(vui & 0x0F0F0F0F)) * dm.x + dm.y; - const vec4 v1 = vec4(unpack8((vui >> 4) & 0x0F0F0F0F)) * dm.x + dm.y; - - store_a(col, k_pair, FLOAT_TYPEV2(v0.xy)); - store_a(col, k_pair + 1, FLOAT_TYPEV2(v0.zw)); - store_a(col, k_pair + 8, FLOAT_TYPEV2(v1.xy)); - store_a(col, k_pair + 9, FLOAT_TYPEV2(v1.zw)); - } else if (MmTypeA == GGML_TYPE_Q5_0) { - const uint idx = pos_a + col * p.stride_a / mm_load_vec_a() + row; - const uint k_pair = row * mm_load_vec_a() / 4; - - const uint ib = idx / 8; - const uint iqs = idx & 0x07; - - const float d = float(a_q5_0_p16.data[ib].d); - const uint uint_qh = uint(a_q5_0_p16.data[ib].qh[1]) << 16 | uint(a_q5_0_p16.data[ib].qh[0]); - const ivec2 qh0 = ivec2(((uint_qh >> 2*iqs) << 4) & 0x10, (uint_qh >> (2*iqs + 12)) & 0x10); - const ivec2 qh1 = ivec2(((uint_qh >> (2*iqs + 1)) << 4) & 0x10, (uint_qh >> (2*iqs + 13)) & 0x10); - - const uint vui = uint(a_q5_0_p16.data[ib].qs[iqs]); - const vec4 v = (vec4((vui & 0xF) | qh0.x, ((vui >> 4) & 0xF) | qh0.y, ((vui >> 8) & 0xF) | qh1.x, (vui >> 12) | qh1.y) - 16.0f) * d; - - store_a(col, k_pair, FLOAT_TYPEV2(v.xz)); - store_a(col, k_pair + 8, FLOAT_TYPEV2(v.yw)); - } else if (MmTypeA == GGML_TYPE_Q5_1) { - const uint idx = pos_a + col * p.stride_a / mm_load_vec_a() + row; - const uint k_pair = row * mm_load_vec_a() / 4; - - const uint ib = idx / 4; - const uint iqs = idx & 0x03; - - const vec2 dm = vec2(a_q5_1_p32.data[ib].dm); - const uint uint_qh = a_q5_1_p32.data[ib].qh; - const uvec2 qh0 = uvec2(((uint_qh >> 4*iqs) << 4) & 0x10, (uint_qh >> (4*iqs + 12)) & 0x10); - const uvec2 qh1 = uvec2(((uint_qh >> (4*iqs + 1)) << 4) & 0x10, (uint_qh >> (4*iqs + 13)) & 0x10); - const uvec2 qh2 = uvec2(((uint_qh >> (4*iqs + 2)) << 4) & 0x10, (uint_qh >> (4*iqs + 14)) & 0x10); - const uvec2 qh3 = uvec2(((uint_qh >> (4*iqs + 3)) << 4) & 0x10, (uint_qh >> (4*iqs + 15)) & 0x10); - - const uint vui = a_q5_1_p32.data[ib].qs[iqs]; - const vec4 v0 = vec4((vui & 0xF) | qh0.x, ((vui >> 4) & 0xF) | qh0.y, ((vui >> 8) & 0xF) | qh1.x, ((vui >> 12) & 0xF) | qh1.y) * dm.x + dm.y; - const vec4 v1 = vec4(((vui >> 16) & 0xF) | qh2.x, ((vui >> 20) & 0xF) | qh2.y, ((vui >> 24) & 0xF) | qh3.x, ((vui >> 28) & 0xF) | qh3.y) * dm.x + dm.y; - - store_a(col, k_pair, FLOAT_TYPEV2(v0.xz)); - store_a(col, k_pair + 1, FLOAT_TYPEV2(v1.xz)); - store_a(col, k_pair + 8, FLOAT_TYPEV2(v0.yw)); - store_a(col, k_pair + 9, FLOAT_TYPEV2(v1.yw)); - } else if (MmTypeA == GGML_TYPE_Q8_0) { - const uint idx = pos_a + col * p.stride_a / mm_load_vec_a() + row; - const uint k_pair = row * mm_load_vec_a() / 2; - - const uint ib = idx / 8; - const uint iqs = idx & 0x07; - - const float d = float(a_q8_0_p16.data[ib].d); - const i8vec2 v0 = unpack8(int32_t(a_q8_0_p16.data[ib].qs[2*iqs])).xy; // vec4 used due to #12147 - const i8vec2 v1 = unpack8(int32_t(a_q8_0_p16.data[ib].qs[2*iqs + 1])).xy; - const vec4 v = vec4(v0.x, v0.y, v1.x, v1.y) * d; - - store_a(col, k_pair, FLOAT_TYPEV2(v.xy)); - store_a(col, k_pair + 1, FLOAT_TYPEV2(v.zw)); - } else if (MmTypeA == GGML_TYPE_Q1_0) { - const uint idx = pos_a + col * p.stride_a / mm_load_vec_a() + row; - const uint k_pair = row * mm_load_vec_a() / 2; - - const uint ib = idx / 16; - const uint iqs = idx & 0xfu; - - const float d = float(a_q1_0.data[ib].d); - const uint bits = uint(a_q1_0.data[ib].qs[iqs]); - - store_a(col, k_pair, FLOAT_TYPEV2((bits & 0x01u) != 0u ? d : -d, (bits & 0x02u) != 0u ? d : -d)); - store_a(col, k_pair + 1, FLOAT_TYPEV2((bits & 0x04u) != 0u ? d : -d, (bits & 0x08u) != 0u ? d : -d)); - store_a(col, k_pair + 2, FLOAT_TYPEV2((bits & 0x10u) != 0u ? d : -d, (bits & 0x20u) != 0u ? d : -d)); - store_a(col, k_pair + 3, FLOAT_TYPEV2((bits & 0x40u) != 0u ? d : -d, (bits & 0x80u) != 0u ? d : -d)); - } else if (MmTypeA == GGML_TYPE_Q2_0) { - const uint idx = pos_a + col * p.stride_a / mm_load_vec_a() + row; - const uint k_pair = row * mm_load_vec_a() / 2; - - const uint ib = idx / 16; - const uint iqs = idx & 0xfu; - - const FLOAT_TYPE d = FLOAT_TYPE(a_q2_0.data[ib].d); - const uint bits = uint(a_q2_0.data[ib].qs[iqs]); - - store_a(col, k_pair, d * (FLOAT_TYPEV2(bits & 3u, (bits >> 2u) & 3u) - FLOAT_TYPEV2(1.0f))); - store_a(col, k_pair + 1, d * (FLOAT_TYPEV2((bits >> 4u) & 3u, bits >> 6u) - FLOAT_TYPEV2(1.0f))); - } else if (MmTypeA == GGML_TYPE_Q2_K) { - const uint idx = pos_a + col * p.stride_a / mm_load_vec_a() + row; - const uint k_pair = row * mm_load_vec_a() / 2; - - const uint ib = idx / 64; // 4 values per idx - const uint iqs = (idx % 64) * 2; // 0,2,4..126 - - const uint qsi = (iqs / 64) * 16 + (iqs % 16); // 0..15 - const uint scalesi = iqs / 8; // 0..15 - const uint qsshift = ((iqs % 64) / 16) * 2; // 0,2,4,6 - - const vec4 qs = vec4(unpack8((a_q2_k_p32.data[ib].qs[qsi / 2] >> qsshift) & 0x03030303)); - const uint scales = a_q2_k.data[ib].scales[scalesi]; - const vec2 dm = vec2(a_q2_k.data[ib].dm); - - const vec4 v = dm.x * float(scales & 0xF) * qs - dm.y * float(scales >> 4); - - store_a(col, k_pair, FLOAT_TYPEV2(v.xy)); - store_a(col, k_pair + 1, FLOAT_TYPEV2(v.zw)); - } else if (MmTypeA == GGML_TYPE_Q3_K) { - const uint idx = pos_a + col * p.stride_a / mm_load_vec_a() + row; - const uint k_pair = row * mm_load_vec_a() / 2; - - const uint ib = idx / 128; // 2 values per idx - const uint iqs = idx % 128; // 0..127 - - const uint n = iqs / 64; // 0,1 - const uint qsi = n * 32 + (iqs % 16) * 2; // 0,2,4..62 - const uint hmi = (iqs % 16) * 2; // 0,2,4..30 - const uint j = (iqs % 64) / 4; // 0..3 - const uint is = iqs / 8; // 0..15 - const uint halfsplit = ((iqs % 64) / 16); // 0,1,2,3 - const uint qsshift = halfsplit * 2; // 0,2,4,6 - - const int8_t us = int8_t(((a_q3_k.data[ib].scales[is % 8] >> (4 * int(is / 8))) & 0xF) - | (((a_q3_k.data[ib].scales[8 + (is % 4)] >> (2 * int(is / 4))) & 3) << 4)); - const float dl = float(a_q3_k.data[ib].d) * float(us - 32); - - const vec2 qs = vec2(unpack8((uint(a_q3_k_p16.data[ib].qs[qsi / 2]) >> qsshift) & 0x0303).xy); - const vec2 hm = vec2(unpack8(((uint(a_q3_k_p16.data[ib].hmask[hmi / 2]) >> (4 * n + halfsplit)) & 0x0101 ^ 0x0101) << 2).xy); - - store_a(col, k_pair, FLOAT_TYPEV2(dl * (qs.x - hm.x), - dl * (qs.y - hm.y))); - - } else if (MmTypeA == GGML_TYPE_Q4_K) { - const uint idx = pos_a + col * p.stride_a / mm_load_vec_a() + row; - const uint k_pair = row * mm_load_vec_a() / 2; - - const uint ib = idx / 64; // 4 values per idx - const uint iqs = (idx % 64) * 2; // 0,2,4..126 - - const uint n = iqs / 32; // 0,1,2,3 - const uint b = (iqs % 32) / 16; // 0,1 - const uint is = 2 * n + b; // 0..7 - const uint qsi = n * 32 + (iqs % 16) * 2; // 0,2,4..126 - - const vec2 loadd = vec2(a_q4_k.data[ib].dm); - - const uvec3 scales = uvec3(a_q4_k_p32.data[ib].scales[0], - a_q4_k_p32.data[ib].scales[1], - a_q4_k_p32.data[ib].scales[2]); - const uint scalesoffs = (is & 3) * 8; - - const uint scidx0 = (is < 4) ? 0 : 2; - const uint scidxshift0 = scalesoffs; - const uint scidxshift1 = (is < 4) ? scalesoffs : scalesoffs + 2; - const uint mbidx0 = (is < 4) ? 1 : 2; - const uint mbidxshift0 = (is < 4) ? scalesoffs : scalesoffs + 4; - const uint mbidxshift1 = (is < 4) ? scalesoffs : scalesoffs + 2; - - const uint8_t sc = uint8_t(((scales[scidx0] >> scidxshift0) & 0xF) | ((scales[0] >> scidxshift1) & 0x30)); - const uint8_t mbyte = uint8_t(((scales[mbidx0] >> mbidxshift0) & 0xF) | ((scales[1] >> mbidxshift1) & 0x30)); - - const float d = loadd.x * sc; - const float m = -loadd.y * mbyte; - - const vec4 q = vec4(unpack8((a_q4_k_p32.data[ib].qs[qsi / 4] >> (b * 4)) & 0x0F0F0F0F)); - - store_a(col, k_pair, FLOAT_TYPEV2(fma(d, q.x, m), fma(d, q.y, m))); - store_a(col, k_pair + 1, FLOAT_TYPEV2(fma(d, q.z, m), fma(d, q.w, m))); - } else if (MmTypeA == GGML_TYPE_Q5_K) { - const uint idx = pos_a + col * p.stride_a / mm_load_vec_a() + row; - const uint k_pair = row * mm_load_vec_a() / 2; - - const uint ib = idx / 64; // 4 values per idx - const uint iqs = (idx % 64) * 2; // 0,2,4..126 - - const uint n = iqs / 32; // 0,1,2,3 - const uint b = (iqs % 32) / 16; // 0,1 - const uint is = 2 * n + b; // 0..7 - const uint qsi = n * 32 + (iqs % 16) * 2; // 0,2,4..126 - const uint qhi = (iqs % 16) * 2; // 0,2,4..30 - - const vec2 loadd = vec2(a_q5_k.data[ib].dm); - - const uvec3 scales = uvec3(a_q5_k_p32.data[ib].scales[0], - a_q5_k_p32.data[ib].scales[1], - a_q5_k_p32.data[ib].scales[2]); - const uint scalesoffs = (is & 3) * 8; - - const uint scidx0 = (is < 4) ? 0 : 2; - const uint scidxshift0 = scalesoffs; - const uint scidxshift1 = (is < 4) ? scalesoffs : scalesoffs + 2; - const uint mbidx0 = (is < 4) ? 1 : 2; - const uint mbidxshift0 = (is < 4) ? scalesoffs : scalesoffs + 4; - const uint mbidxshift1 = (is < 4) ? scalesoffs : scalesoffs + 2; - - const uint8_t sc = uint8_t(((scales[scidx0] >> scidxshift0) & 0xF) | ((scales[0] >> scidxshift1) & 0x30)); - const uint8_t mbyte = uint8_t(((scales[mbidx0] >> mbidxshift0) & 0xF) | ((scales[1] >> mbidxshift1) & 0x30)); - - const float d = loadd.x * sc; - const float m = -loadd.y * mbyte; - - const uint qs = (a_q5_k_p32.data[ib].qs[qsi / 4] >> (b * 4)) & 0x0F0F0F0F; - const uint qh = ((a_q5_k_p32.data[ib].qh[qhi / 4] >> (iqs / 16)) & 0x01010101) << 4; - const vec4 q = vec4(unpack8(qs | qh)); - - store_a(col, k_pair, FLOAT_TYPEV2(fma(d, q.x, m), fma(d, q.y, m))); - store_a(col, k_pair + 1, FLOAT_TYPEV2(fma(d, q.z, m), fma(d, q.w, m))); - } else if (MmTypeA == GGML_TYPE_Q6_K) { - const uint idx = pos_a + col * p.stride_a / mm_load_vec_a() + row; - const uint k_pair = row * mm_load_vec_a() / 2; - - const uint ib = idx / 128; // 2 values per idx - const uint iqs = idx % 128; // 0..127 - - const uint n = iqs / 64; // 0,1 - const uint b = ((iqs % 64) / 32) * 4; // 0,4 - const uint is_b = (iqs % 16) / 8; // 0,1 - const uint qhshift = ((iqs % 64) / 16) * 2; // 0,2,4,6 - const uint is = 8 * n + qhshift + is_b; // 0..15 - const uint qsi = n * 32 + (iqs % 32); // 0..63 - const uint qhi = n * 16 + (iqs % 16); // 0..31 - - const float dscale = float(a_q6_k.data[ib].d) * float(a_q6_k.data[ib].scales[is]); - - const uint ql = (uint(a_q6_k_p16.data[ib].ql[qsi]) >> b) & 0x0F0F; - const uint qh = (uint(a_q6_k_p16.data[ib].qh[qhi]) >> qhshift) & 0x0303; - const vec2 q = (vec2(unpack8(ql | (qh << 4)).xy) - 32) * dscale; - - store_a(col, k_pair, FLOAT_TYPEV2(q.x, q.y)); - } else if (MmTypeA == GGML_TYPE_TQ1_0) { - const uint idx = pos_a + col * p.stride_a / mm_load_vec_a() + row; - - const uint ib = idx / 128; // 2 values per idx - const uint iqs = (idx % 128) * 2; // elem 0,2,4..254 - - const float d = float(a_tq1_0.data[ib].d); - vec2 v; - for (uint kk = 0u; kk < 2u; ++kk) { - const uint e = iqs + kk; - const uint bidx = tq1_0_byte_of(e); - const uint qbyte = uint(bidx < 48u ? a_tq1_0.data[ib].qs[bidx] - : a_tq1_0.data[ib].qh[bidx - 48u]); - v[kk] = d * (float(tq1_0_trit(qbyte, tq1_0_digit_of(e))) - 1.0); - } - - const uint k_pair = row * mm_load_vec_a() / 2; - store_a(col, k_pair, FLOAT_TYPEV2(v.xy)); - } else if (MmTypeA == GGML_TYPE_TQ2_0) { - const uint idx = pos_a + col * p.stride_a / mm_load_vec_a() + row; - - const uint ib = idx / 128; // 2 values per idx - const uint iqs = (idx % 128) * 2; // elem 0,2,4..254 - - const uint qsi = (iqs / 128) * 32 + (iqs % 32); // byte pair start - const uint shift = 2 * ((iqs % 128) / 32); // 0,2,4,6 - - const uvec2 qs = uvec2(a_tq2_0.data[ib].qs[qsi], a_tq2_0.data[ib].qs[qsi + 1]); - const float d = float(a_tq2_0.data[ib].d); - - const vec2 v = d * (vec2((qs >> shift) & 3) - 1.0); - - const uint k_pair = row * mm_load_vec_a() / 2; - store_a(col, k_pair, FLOAT_TYPEV2(v.xy)); - } #endif } #if !defined(MUL_MAT_ID) void load_b_to_shmem(const uint pos_b, const uint row, const uint col, const uint idx_n, const uint block, const uint end_k) { #if LOAD_VEC_B == 8 - if (ALIGNED != 0) { - // Not supported for b_type bf16 because bf16mat2x4 does not exist - const uint idx = pos_b + col * p.stride_b / LOAD_VEC_B + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_B / 2; - FLOAT_TYPEV8 bb = FLOAT_TYPEV8(data_b[idx]); - buf_b[buf_idx + 0] = bb[0].xy; - buf_b[buf_idx + 1] = bb[0].zw; - buf_b[buf_idx + 2] = bb[1].xy; - buf_b[buf_idx + 3] = bb[1].zw; - return; - } + if (ALIGNED != 0) { + // Not supported for b_type bf16 because bf16mat2x4 does not exist + const uint idx = pos_b + col * p.stride_b / LOAD_VEC_B + row; + const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_B / 2; + FLOAT_TYPEV8 bb = FLOAT_TYPEV8(data_b[idx]); + buf_b[buf_idx + 0] = bb[0].xy; + buf_b[buf_idx + 1] = bb[0].zw; + buf_b[buf_idx + 2] = bb[1].xy; + buf_b[buf_idx + 3] = bb[1].zw; + return; + } #elif LOAD_VEC_B == 4 - if (ALIGNED != 0) { - const uint idx = pos_b + col * p.stride_b / LOAD_VEC_B + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_B / 2; + if (ALIGNED != 0) { + const uint idx = pos_b + col * p.stride_b / LOAD_VEC_B + row; + const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_B / 2; #if defined(DATA_B_BF16) - FLOAT_TYPEV4 bb = FLOAT_TYPEV4(TO_FLOAT_TYPE(data_b[idx])); + FLOAT_TYPEV4 bb = FLOAT_TYPEV4(TO_FLOAT_TYPE(data_b[idx])); #else - FLOAT_TYPEV4 bb = FLOAT_TYPEV4(data_b[idx]); + FLOAT_TYPEV4 bb = FLOAT_TYPEV4(data_b[idx]); #endif - buf_b[buf_idx + 0] = bb.xy; - buf_b[buf_idx + 1] = bb.zw; - return; - } + buf_b[buf_idx + 0] = bb.xy; + buf_b[buf_idx + 1] = bb.zw; + return; + } #endif - const uint idx = pos_b + col * p.stride_b + row * 2; - const uint buf_idx = col * SHMEM_STRIDE + row; - if (idx_n < p.N && block + row * 2 + 1 < end_k) { - buf_b[buf_idx] = FLOAT_TYPEV2(TO_FLOAT_TYPE(data_b_scalar[idx]), - TO_FLOAT_TYPE(data_b_scalar[idx + 1])); - } else if (idx_n < p.N && block + row * 2 < end_k) { - buf_b[buf_idx] = FLOAT_TYPEV2(TO_FLOAT_TYPE(data_b_scalar[idx]), 0.0f); - } else { - buf_b[buf_idx] = FLOAT_TYPEV2(0.0f); - } + const uint idx = pos_b + col * p.stride_b + row * 2; + const uint buf_idx = col * SHMEM_STRIDE + row; + if (idx_n < p.N && block + row * 2 + 1 < end_k) { + buf_b[buf_idx] = FLOAT_TYPEV2(TO_FLOAT_TYPE(data_b_scalar[idx]), + TO_FLOAT_TYPE(data_b_scalar[idx + 1])); + } else if (idx_n < p.N && block + row * 2 < end_k) { + buf_b[buf_idx] = FLOAT_TYPEV2(TO_FLOAT_TYPE(data_b_scalar[idx]), 0.0f); + } else { + buf_b[buf_idx] = FLOAT_TYPEV2(0.0f); + } } #else void load_b_to_shmem(const uint pos_b, const uint row, const uint col, const uint ic, const uint _ne1, const uint block, const uint end_k) { #if LOAD_VEC_B == 8 - if (ALIGNED != 0) { - // Not supported for b_type bf16 because bf16mat2x4 does not exist - const u16vec2 row_idx = row_ids[col]; - const uint idx = pos_b + row_idx.y * p.batch_stride_b / LOAD_VEC_B + (row_idx.x % p.ne11) * p.stride_b / LOAD_VEC_B + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_B / 2; - FLOAT_TYPEV8 bb = FLOAT_TYPEV8(data_b[idx]); - buf_b[buf_idx + 0] = bb[0].xy; - buf_b[buf_idx + 1] = bb[0].zw; - buf_b[buf_idx + 2] = bb[1].xy; - buf_b[buf_idx + 3] = bb[1].zw; - return; - } + if (ALIGNED != 0) { + // Not supported for b_type bf16 because bf16mat2x4 does not exist + const u16vec2 row_idx = row_ids[col]; + const uint idx = pos_b + row_idx.y * p.batch_stride_b / LOAD_VEC_B + (row_idx.x % p.ne11) * p.stride_b / LOAD_VEC_B + row; + const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_B / 2; + FLOAT_TYPEV8 bb = FLOAT_TYPEV8(data_b[idx]); + buf_b[buf_idx + 0] = bb[0].xy; + buf_b[buf_idx + 1] = bb[0].zw; + buf_b[buf_idx + 2] = bb[1].xy; + buf_b[buf_idx + 3] = bb[1].zw; + return; + } #elif LOAD_VEC_B == 4 - if (ALIGNED != 0) { - const u16vec2 row_idx = row_ids[col]; - const uint idx = pos_b + row_idx.y * p.batch_stride_b / LOAD_VEC_B + (row_idx.x % p.ne11) * p.stride_b / LOAD_VEC_B + row; - const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_B / 2; + if (ALIGNED != 0) { + const u16vec2 row_idx = row_ids[col]; + const uint idx = pos_b + row_idx.y * p.batch_stride_b / LOAD_VEC_B + (row_idx.x % p.ne11) * p.stride_b / LOAD_VEC_B + row; + const uint buf_idx = col * SHMEM_STRIDE + row * LOAD_VEC_B / 2; #if defined(DATA_B_BF16) - FLOAT_TYPEV4 bb = FLOAT_TYPEV4(TO_FLOAT_TYPE(data_b[idx])); + FLOAT_TYPEV4 bb = FLOAT_TYPEV4(TO_FLOAT_TYPE(data_b[idx])); #else - FLOAT_TYPEV4 bb = FLOAT_TYPEV4(data_b[idx]); + FLOAT_TYPEV4 bb = FLOAT_TYPEV4(data_b[idx]); #endif - buf_b[buf_idx + 0] = bb.xy; - buf_b[buf_idx + 1] = bb.zw; - return; - } + buf_b[buf_idx + 0] = bb.xy; + buf_b[buf_idx + 1] = bb.zw; + return; + } #endif - const uint row_i = ic * BN + col; - const uint buf_idx = col * SHMEM_STRIDE + row; - if (row_i < _ne1 && block + row * 2 + 1 < end_k) { - const u16vec2 row_idx = row_ids[col]; - const uint idx = pos_b + row_idx.y * p.batch_stride_b + (row_idx.x % p.ne11) * p.stride_b + row * 2; - buf_b[buf_idx] = FLOAT_TYPEV2(TO_FLOAT_TYPE(data_b_scalar[idx]), - TO_FLOAT_TYPE(data_b_scalar[idx + 1])); - } else if (row_i < _ne1 && block + row * 2 < end_k) { - const u16vec2 row_idx = row_ids[col]; - const uint idx = pos_b + row_idx.y * p.batch_stride_b + (row_idx.x % p.ne11) * p.stride_b + row * 2; - buf_b[buf_idx] = FLOAT_TYPEV2(TO_FLOAT_TYPE(data_b_scalar[idx]), 0.0f); - } else { - buf_b[buf_idx] = FLOAT_TYPEV2(0.0f); - } + const uint row_i = ic * BN + col; + const uint buf_idx = col * SHMEM_STRIDE + row; + if (row_i < _ne1 && block + row * 2 + 1 < end_k) { + const u16vec2 row_idx = row_ids[col]; + const uint idx = pos_b + row_idx.y * p.batch_stride_b + (row_idx.x % p.ne11) * p.stride_b + row * 2; + buf_b[buf_idx] = FLOAT_TYPEV2(TO_FLOAT_TYPE(data_b_scalar[idx]), + TO_FLOAT_TYPE(data_b_scalar[idx + 1])); + } else if (row_i < _ne1 && block + row * 2 < end_k) { + const u16vec2 row_idx = row_ids[col]; + const uint idx = pos_b + row_idx.y * p.batch_stride_b + (row_idx.x % p.ne11) * p.stride_b + row * 2; + buf_b[buf_idx] = FLOAT_TYPEV2(TO_FLOAT_TYPE(data_b_scalar[idx]), 0.0f); + } else { + buf_b[buf_idx] = FLOAT_TYPEV2(0.0f); + } } #endif diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq.comp b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq.comp index 1fbcbf6c9332..b36d056c1c9c 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq.comp @@ -36,6 +36,8 @@ layout (binding = 2) writeonly buffer D {D_TYPE data_d[];}; #ifdef MUL_MAT_ID layout (binding = 3) readonly buffer IDS {int data_ids[];}; layout (binding = 4) readonly buffer Counts {int data_expert_count[];}; +// Fused MUL epilogue, see mul_mm.comp. +layout (binding = 5) readonly buffer FusedScale {float data_fscale[];}; #endif layout (push_constant) uniform parameter @@ -58,6 +60,7 @@ layout (push_constant) uniform parameter uint ne11; uint n_experts; uint hoist_row_ids; + uint fusion_flags; #else uint base_work_group_z; uint num_batches; @@ -303,7 +306,10 @@ void main() { const uint sums_idx = (wsic * TN + cc) * WMITER * TM + wsir * TM + cr; #ifdef MUL_MAT_ID if (dr_warp + cr < p.M) { - data_d[row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr_warp + cr] = D_TYPE(sums[sums_idx].x); + const uint didx = row_idx.y * p.batch_stride_d + row_idx.x * p.stride_d + dr_warp + cr; + data_d[didx] = p.fusion_flags != 0 + ? D_TYPE(float(sums[sums_idx].x) * data_fscale[row_idx.y * p.nei0 + row_idx.x]) + : D_TYPE(sums[sums_idx].x); } #else if (dr_warp + cr < p.M && dc_warp + cc < p.N) { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_funcs.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_funcs.glsl index 24da4f715f83..478a65843b26 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_funcs.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_funcs.glsl @@ -178,6 +178,62 @@ ACC_TYPE mmq_dot_product(const uint ib_a) { } #endif +#if defined(DATA_A_ROCMFP4) || defined(DATA_A_ROCMFP4_FAST) +// 1-byte loads for ROCmFP4 blocks (18 / 17 bytes). Low nibbles are the first +// half of the block, high nibbles the second half. +void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) { + const uint32_t qs = pack32(u8vec4(data_a[ib].qs[iqs * 4 ], + data_a[ib].qs[iqs * 4 + 1], + data_a[ib].qs[iqs * 4 + 2], + data_a[ib].qs[iqs * 4 + 3])); + + const u8vec4 i_a0 = unpack8( qs & 0x0F0F0F0F); + const u8vec4 i_a1 = unpack8((qs >> 4) & 0x0F0F0F0F); + + buf_a[buf_ib].qs[iqs ] = pack32(i8vec4(kvalues_rocmfp4[i_a0.x], kvalues_rocmfp4[i_a0.y], kvalues_rocmfp4[i_a0.z], kvalues_rocmfp4[i_a0.w])); + buf_a[buf_ib].qs[iqs + 4] = pack32(i8vec4(kvalues_rocmfp4[i_a1.x], kvalues_rocmfp4[i_a1.y], kvalues_rocmfp4[i_a1.z], kvalues_rocmfp4[i_a1.w])); + + if (iqs == 0) { +#if defined(DATA_A_ROCMFP4) + buf_a[buf_ib].d = FLOAT_TYPEV2(ue4m3_to_fp32(data_a[ib].e[0]), + ue4m3_to_fp32(data_a[ib].e[1])); +#else + buf_a[buf_ib].d = FLOAT_TYPE(ue4m3_to_fp32(data_a[ib].e)); +#endif + } +} + +void block_a_to_registers(const uint reg_ib, const uint buf_ib) { + cache_a[reg_ib].d = buf_a[buf_ib].d; + + [[unroll]] for (uint iqs = 0; iqs < 8; iqs++) { + cache_a[reg_ib].qs[iqs] = buf_a[buf_ib].qs[iqs]; + } +} + +ACC_TYPE mmq_dot_product(const uint ib_a) { +#if defined(DATA_A_ROCMFP4) + // the halves carry different scales, so they need separate accumulators + int32_t q_sum0 = 0; + int32_t q_sum1 = 0; + [[unroll]] for (uint iqs = 0; iqs < 4; iqs++) { + q_sum0 += dotPacked4x8EXT(cache_a[ib_a].qs[iqs], cache_b.qs[iqs]); + q_sum1 += dotPacked4x8EXT(cache_a[ib_a].qs[iqs + 4], cache_b.qs[iqs + 4]); + } + + return ACC_TYPE((float(q_sum0) * float(cache_a[ib_a].d.x) + + float(q_sum1) * float(cache_a[ib_a].d.y)) * float(cache_b.ds.x)); +#else + int32_t q_sum = 0; + [[unroll]] for (uint iqs = 0; iqs < 8; iqs++) { + q_sum += dotPacked4x8EXT(cache_a[ib_a].qs[iqs], cache_b.qs[iqs]); + } + + return ACC_TYPE(float(q_sum) * float(cache_a[ib_a].d) * float(cache_b.ds.x)); +#endif +} +#endif + #if defined(DATA_A_MXFP4) // 1-byte loads for mxfp4 blocks (17 bytes) void block_a_to_shmem(const uint buf_ib, const uint ib, const uint iqs) { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_shmem_types.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_shmem_types.glsl index 2b7adcb6c2fc..1995acaf358f 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_shmem_types.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/mul_mmq_shmem_types.glsl @@ -59,6 +59,18 @@ struct block_a_cache { int32_t qs[8]; FLOAT_TYPE d; }; +#elif defined(DATA_A_ROCMFP4) +#define QUANT_R_MMQ 2 +struct block_a_cache { + int32_t qs[8]; + FLOAT_TYPEV2 d; // one UE4M3 scale per half block +}; +#elif defined(DATA_A_ROCMFP4_FAST) +#define QUANT_R_MMQ 2 +struct block_a_cache { + int32_t qs[8]; + FLOAT_TYPE d; +}; #elif defined(DATA_A_Q2_K) #define QUANT_R_MMQ 4 struct block_a_cache { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/topk_radix_select.comp b/ggml/src/ggml-vulkan/vulkan-shaders/topk_radix_select.comp index 8e14b2e99253..c8975391b30f 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/topk_radix_select.comp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/topk_radix_select.comp @@ -2,6 +2,8 @@ #extension GL_EXT_control_flow_attributes : enable #extension GL_EXT_shader_16bit_storage : require +#extension GL_KHR_shader_subgroup_basic : enable +#extension GL_KHR_shader_subgroup_ballot : enable #include "types.glsl" @@ -31,7 +33,7 @@ layout (push_constant) uniform parameter { shared uint histo[RADIX_SIZE]; shared uint sh_bucket; shared uint sh_above; -shared uint out_count; +shared uint sg_cnt[64]; // per-subgroup hit counts for the slot scan // order-preserving float -> uint mapping uint f2ui(float x) { @@ -114,25 +116,44 @@ void topk(const uint row) { barrier(); } - if (tid == 0) { - out_count = 0; - } - barrier(); - // emit everything above the threshold, then fill the rest from ties + // Emit everything above the threshold, then fill the rest from ties. Slots come from + // an exclusive scan over the candidate flags, one BLOCK_SIZE chunk at a time in ascending + // index order, so the output is identical on every run. The previous atomicAdd slot + // counter made the ORDER scheduling-dependent, and at the tie boundary the SET as well: + // the QSA width is top_k + ratio - 1, so the boundary block's ratio tied cells race for + // ratio - 1 slots and a different cell lost each run (non-repeatable output at depth). + // With the scan the lowest-indexed tied cells win. const uint threshold = prefix; - for (uint i = tid; i < ncols; i += BLOCK_SIZE) { - if (f2ui(load(row, i, false)) > threshold) { - data_d[row_out + atomicAdd(out_count, 1u)] = int(i); - } - } - barrier(); - for (uint i = tid; i < ncols; i += BLOCK_SIZE) { - if (f2ui(load(row, i, false)) == threshold) { - const uint pos = atomicAdd(out_count, 1u); - if (pos < p.k) { - data_d[row_out + pos] = int(i); + uint base = 0; + [[dont_unroll]] for (uint pass = 0; pass < 2; ++pass) { + for (uint c = 0; c < ncols; c += BLOCK_SIZE) { + const uint i = c + tid; + bool hit = false; + if (i < ncols) { + const uint key = f2ui(load(row, i, false)); + hit = (pass == 0) ? (key > threshold) : (key == threshold); + } + const uvec4 ballot = subgroupBallot(hit); + const uint rank_sg = subgroupBallotExclusiveBitCount(ballot); + const uint cnt_sg = subgroupBallotBitCount(ballot); + if (subgroupElect()) { + sg_cnt[gl_SubgroupID] = cnt_sg; + } + barrier(); + uint sg_base = 0; + uint total = 0; + for (uint sg = 0; sg < gl_NumSubgroups; ++sg) { + const uint v = sg_cnt[sg]; + sg_base += (sg < gl_SubgroupID) ? v : 0; + total += v; + } + const uint slot = base + sg_base + rank_sg; + if (hit && slot < p.k) { + data_d[row_out + slot] = int(i); } + base += total; + barrier(); } } } diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/types.glsl b/ggml/src/ggml-vulkan/vulkan-shaders/types.glsl index 21d601d0eb61..b55897bd1d9e 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/types.glsl +++ b/ggml/src/ggml-vulkan/vulkan-shaders/types.glsl @@ -955,7 +955,6 @@ shared uint16_t iq1s_grid[2048]; shared uint32_t iq1s_grid_gpu[2048]; #endif -#if defined(DATA_A_IQ1_S) || defined(DATA_A_IQ1_M) #define NEEDS_INIT_IQ_SHMEM void init_iq_shmem(uvec3 wgsize) { @@ -979,17 +978,6 @@ void init_iq_shmem(uvec3 wgsize) barrier(); } #endif -#endif - -#if defined(DATA_A_IQ2_XXS) || defined(DATA_A_IQ2_XS) || defined(DATA_A_IQ2_S) -#if defined(DATA_A_IQ2_S) -shared uvec2 iq2s_grid[1024]; -#elif defined(DATA_A_IQ2_XS) -shared uvec2 iq2xs_grid[512]; -#else -shared uvec2 iq2xxs_grid[256]; -#endif -#endif #define QUANT_K_IQ2_XXS 256 #define QUANT_R_IQ2_XXS 1 @@ -1075,7 +1063,8 @@ const uvec2[256] iq2xxs_grid_const = { uvec2(0x08080808, 0x2b2b082b), uvec2(0x08192b08, 0x2b2b1908), uvec2(0x19190808, 0x2b2b2b08), uvec2(0x08081908, 0x2b2b2b19) }; -#if defined(DATA_A_IQ2_XXS) +shared uvec2 iq2xxs_grid[256]; + #define NEEDS_INIT_IQ_SHMEM void init_iq_shmem(uvec3 wgsize) { @@ -1087,15 +1076,12 @@ void init_iq_shmem(uvec3 wgsize) } barrier(); } -#endif -#if defined(DATA_A_IQ2_XXS) #define QUANT_K QUANT_K_IQ2_XXS #define QUANT_R QUANT_R_IQ2_XXS #define A_TYPE block_iq2_xxs #define A_TYPE_PACKED16 block_iq2_xxs_packed16 #endif -#endif #define QUANT_K_IQ2_XS 256 #define QUANT_R_IQ2_XS 1 @@ -1247,7 +1233,8 @@ const uvec2 iq2xs_grid_const[512] = { uvec2(0x082b2b08, 0x2b2b2b2b), uvec2(0x082b2b2b, 0x2b2b2b2b), uvec2(0x2b190819, 0x2b2b2b2b), uvec2(0x2b2b2b2b, 0x2b2b2b2b), }; -#if defined(DATA_A_IQ2_XS) +shared uvec2 iq2xs_grid[512]; + #define NEEDS_INIT_IQ_SHMEM void init_iq_shmem(uvec3 wgsize) { @@ -1259,15 +1246,12 @@ void init_iq_shmem(uvec3 wgsize) } barrier(); } -#endif -#if defined(DATA_A_IQ2_XS) #define QUANT_K QUANT_K_IQ2_XS #define QUANT_R QUANT_R_IQ2_XS #define A_TYPE block_iq2_xs #define A_TYPE_PACKED16 block_iq2_xs_packed16 #endif -#endif #define QUANT_K_IQ2_S 256 #define QUANT_R_IQ2_S 1 @@ -1549,7 +1533,8 @@ const uvec2 iq2s_grid_const[1024] = { uvec2(0x082b082b, 0x2b2b2b2b), uvec2(0x082b2b08, 0x2b2b2b2b), uvec2(0x2b082b08, 0x2b2b2b2b), uvec2(0x2b2b2b2b, 0x2b2b2b2b) }; -#if defined(DATA_A_IQ2_S) +shared uvec2 iq2s_grid[1024]; + #define NEEDS_INIT_IQ_SHMEM void init_iq_shmem(uvec3 wgsize) { @@ -1561,23 +1546,12 @@ void init_iq_shmem(uvec3 wgsize) } barrier(); } -#endif -#if defined(DATA_A_IQ2_S) #define QUANT_K QUANT_K_IQ2_S #define QUANT_R QUANT_R_IQ2_S #define A_TYPE block_iq2_s #define A_TYPE_PACKED16 block_iq2_s_packed16 #endif -#endif - -#if defined(DATA_A_IQ3_XXS) || defined(DATA_A_IQ3_S) -#if defined(DATA_A_IQ3_S) -shared uint32_t iq3s_grid[512]; -#else -shared uint32_t iq3xxs_grid[256]; -#endif -#endif #define QUANT_K_IQ3_XXS 256 #define QUANT_R_IQ3_XXS 1 @@ -1631,7 +1605,8 @@ const uint32_t iq3xxs_grid_const[256] = { 0x3e1c1c1c, 0x3e1c3404, 0x3e24140c, 0x3e24240c, 0x3e2c0404, 0x3e2c0414, 0x3e2c1424, 0x3e341c04, }; -#if defined(DATA_A_IQ3_XXS) +shared uint32_t iq3xxs_grid[256]; + #define NEEDS_INIT_IQ_SHMEM void init_iq_shmem(uvec3 wgsize) { @@ -1643,15 +1618,12 @@ void init_iq_shmem(uvec3 wgsize) } barrier(); } -#endif -#if defined(DATA_A_IQ3_XXS) #define QUANT_K QUANT_K_IQ3_XXS #define QUANT_R QUANT_R_IQ3_XXS #define A_TYPE block_iq3_xxs #define A_TYPE_PACKED16 block_iq3_xxs_packed16 #endif -#endif #define QUANT_K_IQ3_S 256 #define QUANT_R_IQ3_S 1 @@ -1743,7 +1715,8 @@ const uint32_t iq3s_grid_const[512] = { 0x0f090307, 0x0f090501, 0x0f090b01, 0x0f0b0505, 0x0f0b0905, 0x0f0d0105, 0x0f0d0703, 0x0f0f0101, }; -#if defined(DATA_A_IQ3_S) +shared uint32_t iq3s_grid[512]; + #define NEEDS_INIT_IQ_SHMEM void init_iq_shmem(uvec3 wgsize) { @@ -1755,15 +1728,12 @@ void init_iq_shmem(uvec3 wgsize) } barrier(); } -#endif -#if defined(DATA_A_IQ3_S) #define QUANT_K QUANT_K_IQ3_S #define QUANT_R QUANT_R_IQ3_S #define A_TYPE block_iq3_s #define A_TYPE_PACKED16 block_iq3_s_packed16 #endif -#endif #define QUANT_K_IQ4_XS 256 #define QUANT_R_IQ4_XS 1 @@ -1869,6 +1839,87 @@ struct block_nvfp4_packed32 #define A_TYPE_PACKED32 block_nvfp4_packed32 #endif +#define QUANT_K_ROCMFP4 32 +#define QUANT_R_ROCMFP4 2 + +struct block_rocmfp4 +{ + uint8_t qs[QUANT_K_ROCMFP4/2]; + uint8_t e[2]; +}; + +struct block_rocmfp4_fast +{ + uint8_t qs[QUANT_K_ROCMFP4/2]; + uint8_t e; +}; + +#if defined(DATA_A_ROCMFP4) +#define QUANT_K QUANT_K_ROCMFP4 +#define QUANT_R QUANT_R_ROCMFP4 +#define QUANT_AUXF 1 +#define A_TYPE block_rocmfp4 +#endif + +#if defined(DATA_A_ROCMFP4_FAST) +#define QUANT_K QUANT_K_ROCMFP4 +#define QUANT_R QUANT_R_ROCMFP4 +#define QUANT_AUXF 1 +#define A_TYPE block_rocmfp4_fast +#endif + +#define QUANT_K_ROCMFPX_FP8 32 +#define QUANT_R_ROCMFPX_FP8 1 + +struct block_rocmfpx_fp3 +{ + uint8_t qs[12]; + uint8_t e[2]; +}; + +// 32 codes of 6 bits, little-endian packed into 24 bytes, plus two UE4M3 +// half-block scales. Matches block_rocmfp6 in ggml/rocmfpx/rocmfpx.h. +struct block_rocmfpx_fp6 +{ + uint8_t qs[24]; + uint8_t e[2]; +}; + +struct block_rocmfpx_fp8 +{ + int8_t qs[QUANT_K_ROCMFPX_FP8]; + uint8_t e; +}; + +#if defined(DATA_A_ROCMFPX_FP3) +#define QUANT_K QUANT_K_ROCMFPX_FP8 +#define QUANT_R QUANT_R_ROCMFPX_FP8 +#define QUANT_AUXF 1 +#define A_TYPE block_rocmfpx_fp3 +#endif + +#if defined(DATA_A_ROCMFPX_FP6) +#define QUANT_K QUANT_K_ROCMFPX_FP8 +#define QUANT_R QUANT_R_ROCMFPX_FP8 +#define QUANT_AUXF 1 +#define A_TYPE block_rocmfpx_fp6 +#endif + +#if defined(DATA_A_ROCMFPX_FP8) +#define QUANT_K QUANT_K_ROCMFPX_FP8 +#define QUANT_R QUANT_R_ROCMFPX_FP8 +#define QUANT_AUXF 1 +#define A_TYPE block_rocmfpx_fp8 +#endif + +#if defined(DATA_A_ROCMFPX_FP3) || defined(DATA_A_ROCMFPX_FP6) || defined(DATA_A_ROCMFPX_FP8) +#define DATA_A_ROCMFPX_FAMILY +#endif + +#if defined(DATA_A_ROCMFP4) || defined(DATA_A_ROCMFP4_FAST) || defined(DATA_A_ROCMFPX_FAMILY) +#define DATA_A_ROCMFP_ANY +#endif + #if defined(DATA_A_IQ4_NL) || defined(DATA_A_IQ4_XS) const int8_t kvalues_iq4nl_const[16] = { int8_t(-127), int8_t(-104), int8_t(-83), int8_t(-65), int8_t(-49), int8_t(-35), int8_t(-22), int8_t(-10), @@ -1877,7 +1928,6 @@ const int8_t kvalues_iq4nl_const[16] = { shared FLOAT_TYPE kvalues_iq4nl[16]; -#if defined(DATA_A_IQ4_NL) || defined(DATA_A_IQ4_XS) #define NEEDS_INIT_IQ_SHMEM void init_iq_shmem(uvec3 wgsize) { @@ -1888,6 +1938,70 @@ void init_iq_shmem(uvec3 wgsize) barrier(); } #endif + +#if defined(DATA_A_ROCMFP_ANY) || defined(FA_ROCMFPX_FAMILY) +// The ROCm FP formats share one UE4M3 scale encoding. It differs from the +// NVFP4 encoding: the exponent bias is one lower and subnormals are scaled by +// 1/1024, which keeps the halved ROCmFP4 codebook multiply-free at call sites. +shared float rocmfp_ue4m3_fp32_lut[128]; + +float rocmfp_ue4m3_to_fp32_build(uint u) { + if (u == 0u || u == 127u) { + return 0.0; + } + const uint exp = (u >> 3) & 15u; + const uint man = u & 7u; + if (exp == 0u) { + return float(man) * (1.0 / 1024.0); + } + const uint bits = (exp + 119u) << 23 | (man << 20); + return uintBitsToFloat(bits); +} + +float ue4m3_to_fp32(uint8_t x) { + return rocmfp_ue4m3_fp32_lut[min(uint(x), 127u)]; +} +#endif + +#if defined(DATA_A_ROCMFP4) || defined(DATA_A_ROCMFP4_FAST) +// E2M1-derived codebook with the top level retuned from 12 to 10. +const int8_t kvalues_rocmfp4_const[16] = { + int8_t(0), int8_t(1), int8_t(2), int8_t(3), int8_t(4), int8_t(6), int8_t(8), int8_t(10), + int8_t(0), int8_t(-1), int8_t(-2), int8_t(-3), int8_t(-4), int8_t(-6), int8_t(-8), int8_t(-10), +}; + +shared int8_t kvalues_rocmfp4[16]; +#endif + +#if defined(DATA_A_ROCMFPX_FP6) +int rocmfpx_fp6_decode_code(uint code) { + const int mag = int(code & 31u); + return (code & 32u) != 0u ? -(mag == 0 ? 32 : mag) : mag; +} +#endif + +#if defined(DATA_A_ROCMFPX_FAMILY) +const int8_t kvalues_rocmfpx_fp3_const[8] = { + int8_t(0), int8_t(1), int8_t(2), int8_t(4), + int8_t(0), int8_t(-1), int8_t(-2), int8_t(-4) +}; +#endif + +#if defined(DATA_A_ROCMFP_ANY) +#define NEEDS_INIT_IQ_SHMEM +void init_iq_shmem(uvec3 wgsize) +{ + // copy the tables into shared memory and sync +#if defined(DATA_A_ROCMFP4) || defined(DATA_A_ROCMFP4_FAST) + for (uint i = gl_LocalInvocationIndex.x; i < kvalues_rocmfp4.length(); i += wgsize.x) { + kvalues_rocmfp4[i] = kvalues_rocmfp4_const[i]; + } +#endif + for (uint i = gl_LocalInvocationIndex.x; i < rocmfp_ue4m3_fp32_lut.length(); i += wgsize.x) { + rocmfp_ue4m3_fp32_lut[i] = rocmfp_ue4m3_to_fp32_build(i); + } + barrier(); +} #endif #if defined(DATA_A_MXFP4) || defined(DATA_A_NVFP4) @@ -1918,7 +2032,7 @@ float ue4m3_to_fp32_build(uint u) { } #endif -#if (defined(DATA_A_MXFP4) || defined(DATA_A_NVFP4)) && !defined(USE_OCP_FP4) +#if !defined(USE_OCP_FP4) #define NEEDS_INIT_IQ_SHMEM void init_iq_shmem(uvec3 wgsize) { diff --git a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp index d3f425968df2..b0ea14b0e284 100644 --- a/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp +++ b/ggml/src/ggml-vulkan/vulkan-shaders/vulkan-shaders-gen.cpp @@ -71,6 +71,11 @@ const std::vector type_names = { "iq4_xs", "iq4_nl", "mxfp4", + "rocmfp4", + "rocmfp4_fast", + "rocmfpx_fp3", + "rocmfpx_fp6", + "rocmfpx_fp8", "nvfp4", "tq1_0", "tq2_0", @@ -246,17 +251,6 @@ bool is_iq_quant(const std::string& type_name) { return string_starts_with(type_name, "iq"); } -bool is_lut_quant(const std::string& type_name) { - return is_iq_quant(type_name) || type_name == "mxfp4" || type_name == "nvfp4"; -} - -std::string lut_load_vec_a(const std::string& type_name) { - if (type_name == "iq1_s" || type_name == "iq1_m" || type_name == "iq2_xxs" || type_name == "iq2_xs" || type_name == "iq2_s") { - return "8"; - } - return "4"; -} - static const char path_separator = '/'; std::string join_paths(const std::string& path1, const std::string& path2) { @@ -594,28 +588,22 @@ void matmul_shaders(bool fp16, MatMulIdType matmul_id_type, bool coopmat, bool c } for (const auto& tname : type_names) { - if (tname == "bf16") { - continue; - } - - // Float types keep per-type compilation (different accumulation loop structure) - if (tname == "f32" || tname == "f16") { - std::string data_a_key = "DATA_A_" + to_uppercase(tname); - - const std::map float_type_dict = { - {"FLOAT_TYPE", FLOAT_TYPE(1, tname)}, - {"FLOAT_TYPEV2", FLOAT_TYPE(2, tname)}, - {"FLOAT_TYPEV4", FLOAT_TYPE(4, tname)}, - {"FLOAT_TYPEV8", FLOAT_TYPE(8, tname)}, - }; + std::string load_vec_quant = "2"; + if ((tname == "q1_0") || (tname == "q4_0") || (tname == "q4_1") || (tname == "q5_1") || (tname == "iq1_s") || (tname == "iq1_m") || (tname == "iq2_xxs") || (tname == "iq2_xs") || (tname == "iq2_s")) + load_vec_quant = "8"; + else if ((tname == "q2_0") || (tname == "q5_0") || (tname == "q8_0") || (tname == "q2_k") || (tname == "q4_k") || (tname == "q5_k") || (tname == "iq3_xxs") || (tname == "iq3_s") || (tname == "iq4_xs") || (tname == "iq4_nl") || (tname == "mxfp4") || (tname == "nvfp4") || + (tname == "rocmfp4") || (tname == "rocmfp4_fast") || + (tname == "rocmfpx_fp3") || (tname == "rocmfpx_fp6") || (tname == "rocmfpx_fp8")) + load_vec_quant = "4"; - if (!coopmat2) { - string_to_spv(shader_name + "_" + tname + "_f32" + dot2_sfx, source_name, merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f32}, {"B_TYPE_SCALAR", "float"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc); - } + if (tname == "bf16") { continue; } std::string data_a_key = "DATA_A_" + to_uppercase(tname); + // For aligned matmul loads + std::string load_vec_a = (coopmat2 || tname == "f32" || tname == "f16" || tname == "bf16") ? load_vec : load_vec_quant; + const std::map float_type_dict = { {"FLOAT_TYPE", FLOAT_TYPE(1, tname)}, {"FLOAT_TYPEV2", FLOAT_TYPE(2, tname)}, @@ -623,52 +611,31 @@ void matmul_shaders(bool fp16, MatMulIdType matmul_id_type, bool coopmat, bool c {"FLOAT_TYPEV8", FLOAT_TYPE(8, tname)}, }; -#if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT) - if (!f16acc && !coopmat && !coopmat2 && !dot2 && (is_legacy_quant(tname) || is_k_quant(tname) || tname == "mxfp4")) { - string_to_spv(shader_name + "_" + tname + "_q8_1", "mul_mmq.comp", merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"D_TYPE", "float"},}), fp16, coopmat, coopmat2, f16acc); + // don't generate f32 variants for coopmat2 + if (!coopmat2) { + string_to_spv(shader_name + "_" + tname + "_f32" + dot2_sfx, source_name, merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f32}, {"B_TYPE_SCALAR", "float"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc); } -#endif - - if (is_lut_quant(tname)) { - std::string lva = lut_load_vec_a(tname); - string_to_spv(shader_name + "_" + tname + "_f16" + dot2_sfx, source_name, merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"LOAD_VEC_A", lva}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"B_TYPE_SCALAR", "float16_t"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc); - - if (!coopmat2) { - string_to_spv(shader_name + "_" + tname + "_f32" + dot2_sfx, source_name, merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"LOAD_VEC_A", lva}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f32}, {"B_TYPE_SCALAR", "float"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc); - } + if (tname != "f16" && tname != "f32") { + string_to_spv(shader_name + "_" + tname + "_f16" + dot2_sfx, source_name, merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec_a}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"B_TYPE_SCALAR", "float16_t"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc); + } #if defined(GGML_VULKAN_FLOAT_E2M1_GLSLC_SUPPORT) && defined(GGML_VULKAN_FLOAT_E4M3_GLSLC_SUPPORT) - if ((tname == "mxfp4" || tname == "nvfp4") && (coopmat || coopmat2)) { - string_to_spv(shader_name + "_" + tname + "_f16_ocp" + dot2_sfx, source_name, merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"USE_OCP_FP4", "1"}, {"LOAD_VEC_A", lva}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"B_TYPE_SCALAR", "float16_t"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc); - if (!coopmat2) { - string_to_spv(shader_name + "_" + tname + "_f32_ocp" + dot2_sfx, source_name, merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"USE_OCP_FP4", "1"}, {"LOAD_VEC_A", lva}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f32}, {"B_TYPE_SCALAR", "float"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc); - } + if ((coopmat || coopmat2) && (tname == "mxfp4" || tname == "nvfp4")) { + if (!coopmat2) { + string_to_spv(shader_name + "_" + tname + "_f32_ocp" + dot2_sfx, source_name, merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"USE_OCP_FP4", "1"}, {"LOAD_VEC_A", load_vec_a}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f32}, {"B_TYPE_SCALAR", "float"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc); } -#endif - continue; - } - - // dedicated shader needed due to regression on Ampere - if (coopmat2 && (tname == "q4_k" || tname == "q5_k")) { - string_to_spv(shader_name + "_" + tname + "_f16" + dot2_sfx, source_name, merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"LOAD_VEC_A", load_vec}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"B_TYPE_SCALAR", "float16_t"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc); + string_to_spv(shader_name + "_" + tname + "_f16_ocp" + dot2_sfx, source_name, merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"USE_OCP_FP4", "1"}, {"LOAD_VEC_A", load_vec_a}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"B_TYPE_SCALAR", "float16_t"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc); } - } - - // Quant shader: one SPIR-V for all quant types, selected via MmTypeA spec constant - { - const std::map quant_float_type_dict = { - {"FLOAT_TYPE", FLOAT_TYPE(1, "q4_0")}, - {"FLOAT_TYPEV2", FLOAT_TYPE(2, "q4_0")}, - {"FLOAT_TYPEV4", FLOAT_TYPE(4, "q4_0")}, - {"FLOAT_TYPEV8", FLOAT_TYPE(8, "q4_0")}, - }; - - string_to_spv(shader_name + "_quant_f16" + dot2_sfx, source_name, merge_maps(merge_maps(base_dict, quant_float_type_dict), {{"MULMAT_QUANT", "1"}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f16}, {"B_TYPE_SCALAR", "float16_t"}, {"B_TYPEV4", "f16vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc); +#endif - if (!coopmat2) { - string_to_spv(shader_name + "_quant_f32" + dot2_sfx, source_name, merge_maps(merge_maps(base_dict, quant_float_type_dict), {{"MULMAT_QUANT", "1"}, {"LOAD_VEC_B", load_vec}, {"B_TYPE", aligned_b_type_f32}, {"B_TYPE_SCALAR", "float"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}}), fp16, coopmat, coopmat2, f16acc); +#if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT) + // Integer dot mmq performs better with f32 accumulators (different shader, skip for dot2) + if (!f16acc && !coopmat && !coopmat2 && !dot2 && (is_legacy_quant(tname) || is_k_quant(tname) || tname == "mxfp4" || + tname == "rocmfp4" || tname == "rocmfp4_fast")) { + string_to_spv(shader_name + "_" + tname + "_q8_1", "mul_mmq.comp", merge_maps(merge_maps(base_dict, float_type_dict), {{data_a_key, "1"}, {"D_TYPE", "float"},}), fp16, coopmat, coopmat2, f16acc); } +#endif } } @@ -807,7 +774,9 @@ void process_shaders() { // mul mat vec with integer dot product #if defined(GGML_VULKAN_INTEGER_DOT_GLSLC_SUPPORT) - if (is_legacy_quant(tname) || tname == "mxfp4" || is_k_quant(tname) || tname == "iq1_s" || tname == "iq1_m") { + if (is_legacy_quant(tname) || tname == "mxfp4" || + tname == "rocmfp4" || tname == "rocmfp4_fast" || tname == "rocmfpx_fp8" || + is_k_quant(tname) || tname == "iq1_s" || tname == "iq1_m") { string_to_spv("mul_mat_vec_" + tname + "_q8_1_f32", "mul_mat_vecq.comp", merge_maps(base_dict, {{data_a_key, "1"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}, {"FLOAT_TYPEV2", "vec2"}, {"ACC_TYPE", "float"}})); string_to_spv("mul_mat_vec_" + tname + "_q8_1_f32_subgroup", "mul_mat_vecq.comp", merge_maps(base_dict, {{data_a_key, "1"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}, {"FLOAT_TYPEV2", "vec2"}, {"ACC_TYPE", "float"}, {"USE_SUBGROUP_ADD", "1"}})); string_to_spv("mul_mat_vec_" + tname + "_q8_1_f32_subgroup_no_shmem", "mul_mat_vecq.comp", merge_maps(base_dict, {{data_a_key, "1"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}, {"FLOAT_TYPEV2", "vec2"}, {"ACC_TYPE", "float"}, {"USE_SUBGROUP_ADD_NO_SHMEM", "1"}})); @@ -823,9 +792,13 @@ void process_shaders() { string_to_spv("dequant_" + tname, "dequant_" + tname + ".comp", merge_maps(base_dict, {{data_a_key, "1"}, {"D_TYPE", "float16_t"}})); } // Fused dequant+transpose variant for FA quant-KV (per-head-contiguous f16 scratch). - if (tname == "q8_0") { + if (tname == "q8_0" || tname == "iq4_nl" || tname == "q4_0" || tname == "q4_1" || tname == "q5_0" || tname == "q5_1") { string_to_spv("dequant_" + tname + "_transpose", "dequant_" + tname + ".comp", merge_maps(base_dict, {{data_a_key, "1"}, {"D_TYPE", "float16_t"}, {"DEQUANT_TRANSPOSE", "1"}})); } + // Strided-copy counterpart for f16 KV (contiguize the head-interleaved cache layout). + if (tname == "f16") { + string_to_spv("dequant_f16_transpose", "dequant_f16_transpose.comp", {}); + } shader = (tname == "f32" || tname == "f16" || tname == "bf16") ? "get_rows.comp" : "get_rows_quant.comp"; @@ -839,6 +812,33 @@ void process_shaders() { string_to_spv("get_rows_i32", "get_rows.comp", {{"TEMP_TYPE", "uint"}, {"A_TYPE", "uint"}, {"B_TYPE", "int"}, {"D_TYPE", "uint"}}); + string_to_spv("lightning_indexer_f16", "lightning_indexer_scalar64.comp", {}); +#if defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT) + string_to_spv("lightning_indexer_cm_f16", "lightning_indexer_cm.comp", {{"N_WAVES", "8"}, {"HEADS_PER_TILE", "4"}}); + string_to_spv("lightning_indexer_cm_small_f16", "lightning_indexer_cm.comp", {{"N_WAVES", "1"}, {"HEADS_PER_TILE", "1"}}); + string_to_spv("lightning_indexer_decode_cm_f16", "lightning_indexer_decode_cm.comp", {}); +#endif + string_to_spv("flash_attn_top_k_f16", "flash_attn_top_k.comp", {}); +#if defined(GGML_VULKAN_COOPMAT_GLSLC_SUPPORT) + string_to_spv("flash_attn_top_k_cm_f16", "flash_attn_top_k_cm.comp", {}); +#endif + string_to_spv("flash_attn_gather_f16", "flash_attn_gather.comp", {}); + string_to_spv("flash_attn_union_f16", "flash_attn_union.comp", {}); + string_to_spv("flash_attn_gather_union_f16", "flash_attn_gather_union.comp", {}); + // one decoder per quantised K type flash-attention supports natively; f16/bf16/f32 need no + // decode and take the verbatim gather + // q8_0 and q4_0 only: each needs its true element mapping written out (see the shader), and + // these are the two types actually used as a KV cache. The rest take the verbatim gather. + for (const auto& tname : {"q4_0", "q8_0"}) { + string_to_spv("flash_attn_gather_union_dq_" + std::string(tname), "flash_attn_gather_union_dq.comp", + {{"DATA_A_" + to_uppercase(tname), "1"}}); + string_to_spv("flash_attn_gather_dq_" + std::string(tname), "flash_attn_gather_dq.comp", + {{"DATA_A_" + to_uppercase(tname), "1"}}); + } + string_to_spv("dsv4_hc_pre_f32", "dsv4_hc_pre.comp", {}); + string_to_spv("dsv4_hc_comb_f32", "dsv4_hc_comb.comp", {}); + string_to_spv("dsv4_hc_post_f32", "dsv4_hc_post.comp", {}); + string_to_spv("mul_mat_vec_p021_f16_f32_subgroup_add", "mul_mat_vec_p021.comp", {{"A_TYPE", "float16_t"}, {"A_TYPEV4", "f16vec4"}, {"B_TYPE", "float"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}, {"USE_SUBGROUP_ADD", "1"}}); string_to_spv("mul_mat_vec_p021_f16_f32", "mul_mat_vec_p021.comp", {{"A_TYPE", "float16_t"}, {"A_TYPEV4", "f16vec4"}, {"B_TYPE", "float"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}}); string_to_spv("mul_mat_vec_nc_f16_f32", "mul_mat_vec_nc.comp", {{"A_TYPE", "float16_t"}, {"A_TYPEV4", "f16vec4"}, {"B_TYPE", "float"}, {"B_TYPEV4", "vec4"}, {"D_TYPE", "float"}}); @@ -879,13 +879,21 @@ void process_shaders() { string_to_spv("cpy_transpose_02_16", "copy_transpose_02.comp", {{"A_TYPE", "uint16_t"}, {"D_TYPE", "uint16_t"}}); string_to_spv("cpy_transpose_02_32", "copy_transpose_02.comp", {{"A_TYPE", "uint"}, {"D_TYPE", "uint"}}); - for (std::string t : {"q1_0", "q2_0", "q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "iq4_nl"}) { + string_to_spv("cpy_rocmfp4_rocmfp4", "copy_quant_same.comp", {{"A_TYPE", "uint8_t"}, {"D_TYPE", "uint8_t"}, {"QUANT_K", "32"}, {"BLOCK_BYTES", "18"}}); + string_to_spv("cpy_rocmfp4_fast_rocmfp4_fast", "copy_quant_same.comp", {{"A_TYPE", "uint8_t"}, {"D_TYPE", "uint8_t"}, {"QUANT_K", "32"}, {"BLOCK_BYTES", "17"}}); + string_to_spv("cpy_rocmfpx_fp3_rocmfpx_fp3", "copy_quant_same.comp", {{"A_TYPE", "uint8_t"}, {"D_TYPE", "uint8_t"}, {"QUANT_K", "32"}, {"BLOCK_BYTES", "14"}}); + string_to_spv("cpy_rocmfpx_fp6_rocmfpx_fp6", "copy_quant_same.comp", {{"A_TYPE", "uint8_t"}, {"D_TYPE", "uint8_t"}, {"QUANT_K", "32"}, {"BLOCK_BYTES", "26"}}); + string_to_spv("cpy_rocmfpx_fp8_rocmfpx_fp8", "copy_quant_same.comp", {{"A_TYPE", "uint8_t"}, {"D_TYPE", "uint8_t"}, {"QUANT_K", "32"}, {"BLOCK_BYTES", "33"}}); + + for (std::string t : {"q1_0", "q2_0", "q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "iq4_nl", + "rocmfp4", "rocmfp4_fast", "rocmfpx_fp3", "rocmfpx_fp6", "rocmfpx_fp8"}) { string_to_spv("cpy_f32_" + t, "copy_to_quant.comp", {{"DATA_A_" + to_uppercase(t), "1"}, {"S_TYPE", "float"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); string_to_spv("cpy_" + t + "_f32", "copy_from_quant.comp", {{"DATA_A_" + to_uppercase(t), "1"}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); } for (auto src : {std::pair{"f32", "float"}, std::pair{"f16", "float16_t"}}) { - for (std::string dst : {"f32", "f16", "bf16", "q1_0", "q2_0", "q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "iq4_nl"}) { + for (std::string dst : {"f32", "f16", "bf16", "q1_0", "q2_0", "q4_0", "q4_1", "q5_0", "q5_1", "q8_0", "iq4_nl", + "rocmfp4", "rocmfp4_fast", "rocmfpx_fp3", "rocmfpx_fp6", "rocmfpx_fp8"}) { string_to_spv("set_rows_" + std::string(src.first) + "_" + dst + "_i32", "copy_to_quant.comp", {{"SET_ROWS", "1"}, {"DATA_A_" + to_uppercase(dst), "1"}, {"B_TYPE", "uint"}, {"B_SIZE", "32"}, {"S_TYPE", src.second}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); string_to_spv("set_rows_" + std::string(src.first) + "_" + dst + "_i64", "copy_to_quant.comp", {{"SET_ROWS", "1"}, {"DATA_A_" + to_uppercase(dst), "1"}, {"B_TYPE", "uvec2"}, {"B_SIZE", "64"}, {"S_TYPE", src.second}, {"D_TYPE", "float"}, {"FLOAT_TYPE", "float"}}); } @@ -922,8 +930,6 @@ void process_shaders() { string_to_spv("fa_split_k_reduce", "flash_attn_split_k_reduce.comp", {}); string_to_spv("fa_mask_opt", "flash_attn_mask_opt.comp", {}); - string_to_spv("fa_sparse_compact", "flash_attn_sparse_compact.comp", {}); - string_to_spv("fa_sparse_compact_subgroup", "flash_attn_sparse_compact.comp", {{"USE_SUBGROUPS", "1"}}); string_to_spv("quantize_q8_1", "quantize_q8_1.comp", {}); string_to_spv("quantize_q8_1_subgroup", "quantize_q8_1.comp", {{"USE_SUBGROUPS", "1"}}); @@ -950,6 +956,7 @@ void process_shaders() { string_to_spv("concat_i16", "concat.comp", {{"A_TYPE", "uint16_t"}, {"B_TYPE", "uint16_t"}, {"D_TYPE", "uint16_t"}}); string_to_spv("concat_i32", "concat.comp", {{"A_TYPE", "uint"}, {"B_TYPE", "uint"}, {"D_TYPE", "uint"}}); string_to_spv("concat_i64", "concat.comp", {{"A_TYPE", "uvec2"}, {"B_TYPE", "uvec2"}, {"D_TYPE", "uvec2"}}); + string_to_spv("concat_transpose_i32", "concat_transpose.comp", {{"A_TYPE", "uint"}, {"B_TYPE", "uint"}, {"D_TYPE", "uint"}}); string_to_spv("upscale_f32", "upscale.comp", {{"A_TYPE", "float"}, {"B_TYPE", "float"}, {"D_TYPE", "float"}}); @@ -1094,9 +1101,6 @@ void process_shaders() { string_to_spv("fwht_f32", "fwht.comp", {}); string_to_spv("fwht_shmem_f32", "fwht.comp", {{"FWHT_SHMEM", "1"}}); string_to_spv("count_equal_i32", "count_equal.comp", merge_maps(base_dict, {{"A_TYPE", "int"}, {"B_TYPE", "int"}, {"D_TYPE", "int"}})); - string_to_spv("dsv4_hc_comb_f32", "dsv4_hc_comb.comp", {}); - string_to_spv("dsv4_hc_pre_f32", "dsv4_hc_pre.comp", {}); - string_to_spv("dsv4_hc_post_f32", "dsv4_hc_post.comp", {}); string_to_spv("cumsum_f32", "cumsum.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); string_to_spv("cumsum_multipass1_f32", "cumsum_multipass1.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); string_to_spv("cumsum_multipass2_f32", "cumsum_multipass2.comp", merge_maps(base_dict, {{"A_TYPE", "float"}, {"D_TYPE", "float"}})); @@ -1324,7 +1328,9 @@ void write_output_files() { for (const std::string& btype : btypes) { for (const auto& tname : type_names) { - if (btype == "q8_1" && !is_legacy_quant(tname) && tname != "mxfp4" && !is_k_quant(tname) && tname != "iq1_s" && tname != "iq1_m") { + if (btype == "q8_1" && !is_legacy_quant(tname) && tname != "mxfp4" && + tname != "rocmfp4" && tname != "rocmfp4_fast" && tname != "rocmfpx_fp8" && + !is_k_quant(tname) && tname != "iq1_s" && tname != "iq1_m") { continue; } hdr << "extern const void * arr_dmmv_" << tname << "_" << btype << "_f32_data[3];\n"; diff --git a/ggml/src/ggml.c b/ggml/src/ggml.c index 5ef03e190e34..a6b82660ed88 100644 --- a/ggml/src/ggml.c +++ b/ggml/src/ggml.c @@ -11,6 +11,10 @@ // FIXME: required here for quantization functions #include "ggml-quants.h" +// ROCmFPx experimental AMD-native quant formats +#include "../rocmfp4/rocmfp4.h" +#include "../rocmfpx/rocmfpx.h" + #ifdef GGML_USE_CPU_HBM #include #endif @@ -698,6 +702,54 @@ static const struct ggml_type_traits type_traits[GGML_TYPE_COUNT] = { .to_float = (ggml_to_float_t) dequantize_row_q4_0, .from_float_ref = (ggml_from_float_t) quantize_row_q4_0_ref, }, + [GGML_TYPE_Q4_0_ROCMFP4] = { + .type_name = "q4_0_rocmfp4", + .blck_size = QK_ROCMFP4, + .type_size = sizeof(block_rocmfp4), + .is_quantized = true, + .to_float = (ggml_to_float_t) rocmfp4_dequantize_row_q4_0, + .from_float_ref = (ggml_from_float_t) rocmfp4_quantize_row_q4_0_ref, + }, + [GGML_TYPE_Q4_0_ROCMFP4_FAST] = { + .type_name = "q4_0_rocmfp4_fast", + .blck_size = QK_ROCMFP4, + .type_size = sizeof(block_rocmfp4_fast), + .is_quantized = true, + .to_float = (ggml_to_float_t) rocmfp4_dequantize_row_q4_0_fast, + .from_float_ref = (ggml_from_float_t) rocmfp4_quantize_row_q4_0_fast_ref, + }, + [GGML_TYPE_Q3_0_ROCMFPX] = { + .type_name = "q3_0_rocmfpx", + .blck_size = QK_ROCMFP3, + .type_size = sizeof(block_rocmfp3), + .is_quantized = true, + .to_float = (ggml_to_float_t) rocmfpx_dequantize_row_fp3, + .from_float_ref = (ggml_from_float_t) rocmfpx_quantize_row_fp3_ref, + }, + [GGML_TYPE_Q2_0_ROCMFPX] = { + .type_name = "q2_0_rocmfpx", + .blck_size = QK_ROCMFP2, + .type_size = sizeof(block_rocmfp2), + .is_quantized = true, + .to_float = (ggml_to_float_t) rocmfpx_dequantize_row_fp2, + .from_float_ref = (ggml_from_float_t) rocmfpx_quantize_row_fp2_ref, + }, + [GGML_TYPE_Q6_0_ROCMFPX] = { + .type_name = "q6_0_rocmfpx", + .blck_size = QK_ROCMFP6, + .type_size = sizeof(block_rocmfp6), + .is_quantized = true, + .to_float = (ggml_to_float_t) rocmfpx_dequantize_row_fp6, + .from_float_ref = (ggml_from_float_t) rocmfpx_quantize_row_fp6_ref, + }, + [GGML_TYPE_Q8_0_ROCMFPX] = { + .type_name = "q8_0_rocmfpx", + .blck_size = QK_ROCMFP8, + .type_size = sizeof(block_rocmfp8), + .is_quantized = true, + .to_float = (ggml_to_float_t) rocmfpx_dequantize_row_fp8, + .from_float_ref = (ggml_from_float_t) rocmfpx_quantize_row_fp8_ref, + }, [GGML_TYPE_Q4_1] = { .type_name = "q4_1", .blck_size = QK4_1, @@ -1435,6 +1487,19 @@ enum ggml_type ggml_ftype_to_ggml_type(enum ggml_ftype ftype) { case GGML_FTYPE_MOSTLY_Q4_1: wtype = GGML_TYPE_Q4_1; break; case GGML_FTYPE_MOSTLY_Q1_0: wtype = GGML_TYPE_Q1_0; break; case GGML_FTYPE_MOSTLY_Q2_0: wtype = GGML_TYPE_Q2_0; break; + case GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4: + case GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_LEAN: + case GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_COHERENT: + wtype = GGML_TYPE_Q4_0_ROCMFP4; break; + case GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_FAST: + case GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_FAST_COHERENT: + case GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_STRIX: + case GGML_FTYPE_MOSTLY_Q4_0_ROCMFP4_STRIX_LEAN: + wtype = GGML_TYPE_Q4_0_ROCMFP4_FAST; break; + case GGML_FTYPE_MOSTLY_Q3_0_ROCMFPX: wtype = GGML_TYPE_Q3_0_ROCMFPX; break; + case GGML_FTYPE_MOSTLY_Q2_0_ROCMFPX: wtype = GGML_TYPE_Q2_0_ROCMFPX; break; + case GGML_FTYPE_MOSTLY_Q6_0_ROCMFPX: wtype = GGML_TYPE_Q6_0_ROCMFPX; break; + case GGML_FTYPE_MOSTLY_Q8_0_ROCMFPX: wtype = GGML_TYPE_Q8_0_ROCMFPX; break; case GGML_FTYPE_MOSTLY_Q5_0: wtype = GGML_TYPE_Q5_0; break; case GGML_FTYPE_MOSTLY_Q5_1: wtype = GGML_TYPE_Q5_1; break; case GGML_FTYPE_MOSTLY_Q8_0: wtype = GGML_TYPE_Q8_0; break; @@ -5583,6 +5648,20 @@ void ggml_flash_attn_ext_add_sinks( a->src[4] = sinks; } +void ggml_flash_attn_ext_add_top_k( + struct ggml_tensor * a, + struct ggml_tensor * top_k, + int64_t n_kv_raw) { + GGML_ASSERT(a->op == GGML_OP_FLASH_ATTN_EXT); + GGML_ASSERT(a->src[5] == NULL); + GGML_ASSERT(top_k->type == GGML_TYPE_I32); + GGML_ASSERT(top_k->ne[1] == a->src[0]->ne[1]); + GGML_ASSERT(n_kv_raw >= 0 && n_kv_raw <= a->src[1]->ne[1]); + + a->src[5] = top_k; + ggml_set_op_params_i32(a, 4, (int32_t) n_kv_raw); +} + // ggml_flash_attn_back struct ggml_tensor * ggml_flash_attn_back( @@ -8054,6 +8133,24 @@ size_t ggml_quantize_chunk( case GGML_TYPE_Q1_0: result = quantize_q1_0 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; case GGML_TYPE_Q2_0: result = quantize_q2_0 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; case GGML_TYPE_Q4_0: result = quantize_q4_0 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; + case GGML_TYPE_Q4_0_ROCMFP4: + result = rocmfp4_quantize_q4_0(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); + break; + case GGML_TYPE_Q4_0_ROCMFP4_FAST: + result = rocmfp4_quantize_q4_0_fast(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); + break; + case GGML_TYPE_Q3_0_ROCMFPX: + result = rocmfpx_quantize_fp3(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); + break; + case GGML_TYPE_Q2_0_ROCMFPX: + result = rocmfpx_quantize_fp2(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); + break; + case GGML_TYPE_Q6_0_ROCMFPX: + result = rocmfpx_quantize_fp6(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); + break; + case GGML_TYPE_Q8_0_ROCMFPX: + result = rocmfpx_quantize_fp8(src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); + break; case GGML_TYPE_Q4_1: result = quantize_q4_1 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; case GGML_TYPE_Q5_0: result = quantize_q5_0 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; case GGML_TYPE_Q5_1: result = quantize_q5_1 (src + start, (char *) dst + start_row * row_size, nrows, n_per_row, imatrix); break; diff --git a/gguf-py/gguf/constants.py b/gguf-py/gguf/constants.py index e54ee5a0fd41..c978ab9d5417 100644 --- a/gguf-py/gguf/constants.py +++ b/gguf-py/gguf/constants.py @@ -1185,6 +1185,9 @@ class MODEL_TENSOR(IntEnum): NEXTN_HNORM = auto() NEXTN_SHARED_HEAD_HEAD = auto() NEXTN_SHARED_HEAD_NORM = auto() + NEXTN_HC_HEAD_NORM = auto() + NEXTN_HC_HEAD_DOWN = auto() + NEXTN_HC_HEAD_UP = auto() # eagle3 FC = auto() # feature fusion layer D2T = auto() # draft to target vocabulary mapping @@ -1969,6 +1972,9 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.NEXTN_HNORM: "blk.{bid}.nextn.hnorm", MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD: "blk.{bid}.nextn.shared_head_head", MODEL_TENSOR.NEXTN_SHARED_HEAD_NORM: "blk.{bid}.nextn.shared_head_norm", + MODEL_TENSOR.NEXTN_HC_HEAD_NORM: "blk.{bid}.nextn.hc_head_norm", + MODEL_TENSOR.NEXTN_HC_HEAD_DOWN: "blk.{bid}.nextn.hc_head_down", + MODEL_TENSOR.NEXTN_HC_HEAD_UP: "blk.{bid}.nextn.hc_head_up", MODEL_TENSOR.FC: "fc", MODEL_TENSOR.DSPARK_MARKOV_W1: "markov_w1", MODEL_TENSOR.DSPARK_MARKOV_W2: "markov_w2", @@ -2947,12 +2953,23 @@ class MODEL_TENSOR(IntEnum): MODEL_TENSOR.FFN_GATE_EXP, MODEL_TENSOR.FFN_GATE_UP_EXP, MODEL_TENSOR.PER_LAYER_TOKEN_EMBD, + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, MODEL_TENSOR.PLE_KEY, MODEL_TENSOR.PLE_VALUE, MODEL_TENSOR.PLE_NORM_KEY, MODEL_TENSOR.PLE_NORM_QUERY, MODEL_TENSOR.PLE_NORM_CONV, MODEL_TENSOR.PLE_CONV1D, + MODEL_TENSOR.NEXTN_EH_PROJ, + MODEL_TENSOR.NEXTN_EMBED_TOKENS, + MODEL_TENSOR.NEXTN_ENORM, + MODEL_TENSOR.NEXTN_HNORM, + MODEL_TENSOR.NEXTN_SHARED_HEAD_HEAD, + MODEL_TENSOR.NEXTN_HC_HEAD_NORM, + MODEL_TENSOR.NEXTN_HC_HEAD_DOWN, + MODEL_TENSOR.NEXTN_HC_HEAD_UP, ], MODEL_ARCH.PLAMO: [ MODEL_TENSOR.TOKEN_EMBD, @@ -5687,6 +5704,12 @@ class GGMLQuantizationType(IntEnum): MXFP4 = 39 NVFP4 = 40 Q1_0 = 41 + Q4_0_ROCMFP4 = 100 + Q4_0_ROCMFP4_FAST = 101 + Q6_0_ROCMFPX = 102 + Q8_0_ROCMFPX = 103 + Q3_0_ROCMFPX = 104 + Q2_0_ROCMFPX = 107 Q2_0 = 42 @@ -5745,6 +5768,23 @@ class LlamaFileType(IntEnum): MOSTLY_Q1_0 = 40 # except 1d tensors MOSTLY_Q2_0 = 41 # except 1d tensors + MOSTLY_Q4_0_ROCMFP4 = 100 + MOSTLY_Q4_0_ROCMFP4_LEAN = 101 + MOSTLY_Q4_0_ROCMFP4_COHERENT = 102 + MOSTLY_Q4_0_ROCMFP4_FAST = 103 + MOSTLY_Q4_0_ROCMFP4_FAST_COHERENT = 104 + MOSTLY_Q4_0_ROCMFP4_STRIX = 105 + MOSTLY_Q4_0_ROCMFP4_STRIX_LEAN = 106 + MOSTLY_Q6_0_ROCMFPX = 110 + MOSTLY_Q8_0_ROCMFPX = 111 + MOSTLY_Q3_0_ROCMFPX = 112 + MOSTLY_Q3_0_ROCMFPX_AGENT = 113 + MOSTLY_Q6_0_ROCMFPX_AGENT = 114 + MOSTLY_Q8_0_ROCMFPX_AGENT = 115 + MOSTLY_Q6_0_ROCMFPX_LEAN = 116 + MOSTLY_Q6_0_ROCMFPX_AGENT_LEAN = 117 + MOSTLY_Q2_0_ROCMFPX = 119 # as in charlie12345/ROCmFPX (ciru-ai uses 122 and puts Q7_0_ROCMFPX here) + GUESSED = 1024 # not specified in the model file @@ -5880,6 +5920,12 @@ class VisionProjectorType: GGMLQuantizationType.MXFP4: (32, 1 + 16), GGMLQuantizationType.NVFP4: (64, 4 + 32), GGMLQuantizationType.Q1_0: (128, 2 + 16), + GGMLQuantizationType.Q4_0_ROCMFP4: (32, 2 + 16), + GGMLQuantizationType.Q4_0_ROCMFP4_FAST: (32, 1 + 16), + GGMLQuantizationType.Q6_0_ROCMFPX: (32, 24 + 2), + GGMLQuantizationType.Q8_0_ROCMFPX: (32, 32 + 1), + GGMLQuantizationType.Q3_0_ROCMFPX: (32, 12 + 2), + GGMLQuantizationType.Q2_0_ROCMFPX: (32, 8 + 2), GGMLQuantizationType.Q2_0: (64, 2 + 16), } diff --git a/gguf-py/gguf/quants.py b/gguf-py/gguf/quants.py index 80966b6ef151..957314b4155f 100644 --- a/gguf-py/gguf/quants.py +++ b/gguf-py/gguf/quants.py @@ -402,6 +402,136 @@ def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: return (x * d) +def _rocmfpx_ue4m3_to_fp32(e: np.ndarray) -> np.ndarray: + e = e.astype(np.uint8, copy=False) + exp = (e >> np.uint8(3)).astype(np.int32) + mant = (e & np.uint8(7)).astype(np.float32) + subnormal = mant * np.float32(2.0**-10) + normal = np.ldexp(8.0 + mant, exp - 11).astype(np.float32) + return np.where( + e <= np.uint8(0x7E), + np.where(exp == 0, subnormal, normal), + np.float32(0.0), + ).astype(np.float32) + + +def _rocmfpx_unpack_codes(qs: np.ndarray, nbits: int, block_size: int) -> np.ndarray: + bit_pos = np.arange(block_size, dtype=np.uint16) * np.uint16(nbits) + byte_index = bit_pos // np.uint16(8) + bit_index = bit_pos % np.uint16(8) + padded = np.pad(qs, ((0, 0), (0, 1)), mode="constant") + words = ( + padded[:, byte_index].astype(np.uint16) + | (padded[:, byte_index + 1].astype(np.uint16) << np.uint16(8)) + ) + return ((words >> bit_index) & np.uint16((1 << nbits) - 1)).astype(np.uint8) + + +_ROCMFP4_MAG = np.array((0, 1, 2, 3, 4, 6, 8, 10), dtype=np.int8) + + +def _rocmfp4_decode_nibbles(qs: np.ndarray) -> np.ndarray: + # 32 nibbles per block: the low nibbles are the first 16 weights, the high + # nibbles the last 16. Codes are sign (bit 3) + a 3-bit index into the + # {0,1,2,3,4,6,8,10} magnitude table. + codes = np.concatenate((qs & np.uint8(0x0F), qs >> np.uint8(4)), axis=1) + mag = np.take(_ROCMFP4_MAG, codes & np.uint8(7)) + return np.where((codes & np.uint8(8)) != 0, -mag, mag).astype(np.float32) + + +class Q4_0_ROCMFP4(__Quant, qtype=GGMLQuantizationType.Q4_0_ROCMFP4): + @classmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_blocks = blocks.shape[0] + + qs, e = np.hsplit(blocks, [16]) + + vals = _rocmfp4_decode_nibbles(qs).reshape(n_blocks, 2, cls.block_size // 2) + scales = _rocmfpx_ue4m3_to_fp32(e).reshape(n_blocks, 2, 1) + + return (vals * scales).reshape(n_blocks, cls.block_size) + + +class Q4_0_ROCMFP4_FAST(__Quant, qtype=GGMLQuantizationType.Q4_0_ROCMFP4_FAST): + @classmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + qs, e = np.hsplit(blocks, [16]) + + vals = _rocmfp4_decode_nibbles(qs) + scales = _rocmfpx_ue4m3_to_fp32(e) + + return vals * scales + + +class Q2_0_ROCMFPX(__Quant, qtype=GGMLQuantizationType.Q2_0_ROCMFPX): + _vals = np.array((-4, -1, 1, 4), dtype=np.int8) + + @classmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_blocks = blocks.shape[0] + + qs, e = np.hsplit(blocks, [8]) + + # 2-bit codes, LSB first within each byte; bytes 0-3 are the first half + # block, bytes 4-7 the second, each with its own UE4M3 scale. + shifts = np.array((0, 2, 4, 6), dtype=np.uint8).reshape(1, 1, 4) + codes = (qs.reshape(n_blocks, 8, 1) >> shifts) & np.uint8(3) + vals = np.take(cls._vals, codes).astype(np.float32).reshape(n_blocks, 2, cls.block_size // 2) + scales = _rocmfpx_ue4m3_to_fp32(e).reshape(n_blocks, 2, 1) + + return (vals * scales).reshape(n_blocks, cls.block_size) + + +class Q3_0_ROCMFPX(__Quant, qtype=GGMLQuantizationType.Q3_0_ROCMFPX): + _mag = np.array((0, 1, 2, 4), dtype=np.int8) + + @classmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_blocks = blocks.shape[0] + + qs, e = np.hsplit(blocks, [12]) + + codes = _rocmfpx_unpack_codes(qs, 3, cls.block_size) + mag = np.take(cls._mag, codes & np.uint8(3)) + vals = np.where((codes & np.uint8(4)) != 0, -mag, mag).astype(np.float32) + + scales = _rocmfpx_ue4m3_to_fp32(e).reshape(n_blocks, 2, 1) + vals = vals.reshape(n_blocks, 2, cls.block_size // 2) + + return (vals * scales).reshape(n_blocks, cls.block_size) + + +class Q6_0_ROCMFPX(__Quant, qtype=GGMLQuantizationType.Q6_0_ROCMFPX): + @classmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + n_blocks = blocks.shape[0] + + qs, e = np.hsplit(blocks, [24]) + + codes = _rocmfpx_unpack_codes(qs, 6, cls.block_size) + neg = (codes & np.uint8(32)) != 0 + mag = (codes & np.uint8(31)).astype(np.int8) + # code 32 (sign set, zero magnitude) is -32, not -0, in the C decoder + mag = np.where(neg & (mag == 0), np.int8(32), mag) + vals = np.where(neg, -mag, mag).astype(np.float32) + + scales = _rocmfpx_ue4m3_to_fp32(e).reshape(n_blocks, 2, 1) + vals = vals.reshape(n_blocks, 2, cls.block_size // 2) + + return (vals * scales).reshape(n_blocks, cls.block_size) + + +class Q8_0_ROCMFPX(__Quant, qtype=GGMLQuantizationType.Q8_0_ROCMFPX): + @classmethod + def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: + qs, e = np.hsplit(blocks, [32]) + + scales = _rocmfpx_ue4m3_to_fp32(e) + qs = qs.view(np.int8).astype(np.float32) + + return qs * scales + + class Q2_K(__Quant, qtype=GGMLQuantizationType.Q2_K): @classmethod def dequantize_blocks(cls, blocks: np.ndarray) -> np.ndarray: diff --git a/gguf-py/gguf/tensor_mapping.py b/gguf-py/gguf/tensor_mapping.py index d2dfeece5952..173f1a9cf453 100644 --- a/gguf-py/gguf/tensor_mapping.py +++ b/gguf-py/gguf/tensor_mapping.py @@ -2803,6 +2803,15 @@ class TensorNameMap: MODEL_TENSOR.HC_HEAD_UP: ( "model.hyper_connection_mixer.input_mix_weight_up", ), + MODEL_TENSOR.NEXTN_HC_HEAD_NORM: ( + "model.layers.{bid}.hyper_connection_mixer.hc_norm", + ), + MODEL_TENSOR.NEXTN_HC_HEAD_DOWN: ( + "model.layers.{bid}.hyper_connection_mixer.input_mix_weight_down", + ), + MODEL_TENSOR.NEXTN_HC_HEAD_UP: ( + "model.layers.{bid}.hyper_connection_mixer.input_mix_weight_up", + ), MODEL_TENSOR.INDEXER_Q_NORM: ( "model.layers.{bid}.self_attn.indexer.q_layernorm", ), diff --git a/halo-box.png b/halo-box.png new file mode 100644 index 000000000000..bda354e76331 Binary files /dev/null and b/halo-box.png differ diff --git a/include/llama.h b/include/llama.h index 3ab935939c6d..8851c113aaed 100644 --- a/include/llama.h +++ b/include/llama.h @@ -157,6 +157,24 @@ extern "C" { LLAMA_FTYPE_MOSTLY_Q1_0 = 40, // except 1d tensors LLAMA_FTYPE_MOSTLY_Q2_0 = 41, // except 1d tensors + // ROCmFPx experimental AMD-native recipes (high ID range, see ggml_ftype) + LLAMA_FTYPE_MOSTLY_Q4_0_ROCMFP4 = 100, // except 1d tensors + LLAMA_FTYPE_MOSTLY_Q4_0_ROCMFP4_LEAN = 101, // ROCmFP4 with Q5_K token embeddings + LLAMA_FTYPE_MOSTLY_Q4_0_ROCMFP4_COHERENT = 102, // ROCmFP4 with Q6_K token embeddings + LLAMA_FTYPE_MOSTLY_Q4_0_ROCMFP4_FAST = 103, // ROCmFP4 single-scale speed layout + LLAMA_FTYPE_MOSTLY_Q4_0_ROCMFP4_FAST_COHERENT = 104, // ROCmFP4 fast with Q6_K token embeddings + LLAMA_FTYPE_MOSTLY_Q4_0_ROCMFP4_STRIX = 105, // ROCmFP4 Strix Halo quality/speed recipe + LLAMA_FTYPE_MOSTLY_Q4_0_ROCMFP4_STRIX_LEAN = 106, // ROCmFP4 Strix Halo size-biased K/V recipe + LLAMA_FTYPE_MOSTLY_Q6_0_ROCMFPX = 110, // ROCmFPx experimental 6-bit reference layout + LLAMA_FTYPE_MOSTLY_Q8_0_ROCMFPX = 111, // ROCmFPx experimental 8-bit reference layout + LLAMA_FTYPE_MOSTLY_Q3_0_ROCMFPX = 112, // ROCmFPx experimental 3-bit reference layout + LLAMA_FTYPE_MOSTLY_Q3_0_ROCMFPX_AGENT = 113, // ROCmFPx 3-bit agent/tool-call coherent routing + LLAMA_FTYPE_MOSTLY_Q6_0_ROCMFPX_AGENT = 114, // ROCmFPx 6-bit agent/tool-call coherent routing + LLAMA_FTYPE_MOSTLY_Q8_0_ROCMFPX_AGENT = 115, // ROCmFPx 8-bit agent/tool-call coherent routing + LLAMA_FTYPE_MOSTLY_Q6_0_ROCMFPX_LEAN = 116, // ROCmFPx 6-bit size/speed-biased routing + LLAMA_FTYPE_MOSTLY_Q6_0_ROCMFPX_AGENT_LEAN = 117, // ROCmFPx 6-bit agent routing without Q8-heavy boosts + LLAMA_FTYPE_MOSTLY_Q2_0_ROCMFPX = 119, // ROCmFPx 2-bit S40 codebook + dual UE4M3 scales; id as in charlie12345/ROCmFPX (ciru-ai uses 122 and puts Q7_0 here) + LLAMA_FTYPE_GUESSED = 1024, // not specified in the model file }; @@ -218,6 +236,7 @@ extern "C" { LLAMA_LAZY_MODE_OFF = 0, // always read the whole tensor up front LLAMA_LAZY_MODE_AUTO = 1, // lazy only for marked tensors larger than 4 GiB (requires mmap) LLAMA_LAZY_MODE_ON = 2, // read the rows of tensors marked by the arch on demand (requires mmap) + LLAMA_LAZY_MODE_DIRECT = 3, // like on, but the arch reads the rows with explicit pread()s instead of demand paging the mmap (currently only PLE tables) }; enum llama_context_type { @@ -366,6 +385,7 @@ extern "C" { uint32_t n_outputs_max_per_seq; // max outputs per sequence (0 = n_outputs_max) int32_t n_threads; // number of threads to use for generation int32_t n_threads_batch; // number of threads to use for batch processing + int32_t mtp_draft_vocab; // MTP contexts only: draft over token ids < N plus control tokens (0 = full vocabulary) [EXPERIMENTAL] enum llama_context_type ctx_type; // set the context type (e.g. MTP) enum llama_rope_scaling_type rope_scaling_type; // RoPE scaling type, from `enum llama_rope_scaling_type` @@ -1021,6 +1041,9 @@ extern "C" { // Set abort callback LLAMA_API void llama_set_abort_callback(struct llama_context * ctx, ggml_abort_callback abort_callback, void * abort_callback_data); + // Set evaluation callback + LLAMA_API void llama_set_eval_callback(struct llama_context * ctx, ggml_backend_sched_eval_callback cb_eval, void * cb_eval_user_data); + // Wait until all computations are finished // This is automatically done when using one of the functions below to obtain the computation results // and is not necessary to call it explicitly in most cases diff --git a/scripts/bench-spec-prefill.sh b/scripts/bench-spec-prefill.sh new file mode 100755 index 000000000000..faad2c62b179 --- /dev/null +++ b/scripts/bench-spec-prefill.sh @@ -0,0 +1,63 @@ +#!/bin/bash +# Measure TTFT and decode rate for speculative prefill against a llama-server arm. +# +# scripts/bench-spec-prefill.sh \n\n" : "") + xml(text); + tc.expect = simple_assist_msg("", "", "records", json({{"action", text}}).dump()); + }, false); + } + test_peg_parser(tmpls.get(), [&](peg_test_case & tc) { + tc.params = in; + tc.input = (path.find("Qwen3-Coder") == std::string::npos ? "\n\n" : "") + xml("re") + "\n" + xml("reload"); + tc.expect.role = "assistant"; + tc.expect.tool_calls = {{"records", R"({"action":"re"})", ""}, {"records", R"({"action":"reload"})", ""}}; + }, false); + schema["properties"]["action"].erase("enum"); + in.tools = {{"records", "Access records.", schema.dump()}}; + params = common_chat_templates_apply(tmpls.get(), in); + auto grammar = build_grammar(params.grammar); + assert_equals(true, match_string(xml("erase"), grammar.get())); + } +} + int main(int argc, char ** argv) { bool detailed_debug = false; bool only_run_filtered = false; @@ -7392,6 +7484,7 @@ int main(int argc, char ** argv) { } else #endif { + test_qwen_string_enums(); test_msg_diffs_compute(); test_msgs_oaicompat_json_conversion(); test_msg_token_delimiters_split(); @@ -7404,6 +7497,7 @@ int main(int argc, char ** argv) { test_reasoning_effort_caps(); test_reasoning_budget_tokens_per_request(); test_reasoning_budget_message_per_request(); + test_reasoning_budget_enabled_per_request(); test_template_output_peg_parsers(detailed_debug); std::cout << "\n[chat] All tests passed!" << '\n'; } diff --git a/tests/test-mtp-draft-vocab.cpp b/tests/test-mtp-draft-vocab.cpp new file mode 100644 index 000000000000..ca514151a524 --- /dev/null +++ b/tests/test-mtp-draft-vocab.cpp @@ -0,0 +1,243 @@ +// Regression test for llama_context_params::mtp_draft_vocab. +// +// The draft vocabulary subset belongs to the MTP context that asked for it. Contexts over the same model with a +// different setting (0 = full vocabulary, or another N) must each draft over their own vocabulary, whatever order +// they were created in and while they are alive at the same time. +// +// Needs a qwen35 or qwen35moe model with an MTP head whose MTP block uses the model LM head, e.g.: +// +// test-mtp-draft-vocab -m Qwen3.8-27B-UD-IQ4_XS.gguf -ngl 99 +// +// A context with N > 0 must produce -inf exactly at the token ids outside {id < N} + control + user-defined tokens, +// and its finite logits must match those of a full-vocabulary context fed the same inputs. A context with N = 0 must +// produce finite logits everywhere. + +#include "arg.h" +#include "common.h" +#include "llama.h" + +#include "../src/llama-ext.h" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +struct mtp_ctx { + int32_t n_vocab_draft; // requested mtp_draft_vocab + llama_context_ptr ctx; +}; + +static llama_context_ptr make_ctx(llama_model * model, int32_t n_vocab_draft) { + llama_context_params cparams = llama_context_default_params(); + cparams.n_ctx = 256; + cparams.n_batch = 16; + cparams.n_ubatch = 16; + cparams.n_seq_max = 1; + cparams.ctx_type = LLAMA_CONTEXT_TYPE_MTP; + cparams.mtp_draft_vocab = n_vocab_draft; + cparams.n_rs_seq = 0; + llama_context_ptr ctx(llama_init_from_model(model, cparams)); + if (ctx) { + llama_set_embeddings_nextn(ctx.get(), true, /*masked*/ true); + } + return ctx; +} + +// true for the token ids a context with mtp_draft_vocab = n may draft +static std::vector draft_mask(const llama_vocab * vocab, int32_t n) { + const int32_t n_vocab = llama_vocab_n_tokens(vocab); + std::vector mask((size_t) n_vocab, 1); + if (n <= 0 || n >= n_vocab) { + return mask; + } + for (int32_t t = 0; t < n_vocab; ++t) { + const int attr = (int) llama_vocab_get_attr(vocab, t); + mask[(size_t) t] = t < n || (attr & (LLAMA_TOKEN_ATTR_CONTROL | LLAMA_TOKEN_ATTR_USER_DEFINED)); + } + return mask; +} + +// decode one token with a pseudo-random hidden state at position pos; return the draft logits +static bool decode_step(llama_model * model, llama_context * ctx, llama_pos pos, std::vector & out) { + const int32_t n_embd = llama_model_n_embd_out(model); + const int32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model)); + + llama_batch batch = llama_batch_init(1, n_embd, 1); + batch.token = (llama_token *) malloc(sizeof(llama_token)); + + std::mt19937 rng(1000 + pos); + std::normal_distribution dist(0.0f, 1.0f); + for (int32_t i = 0; i < n_embd; ++i) { + batch.embd[i] = dist(rng); + } + batch.token[0] = (llama_token) ((pos * 7919 + 13) % 30000); + batch.pos[0] = pos; + batch.n_seq_id[0] = 1; + batch.seq_id[0][0] = 0; + batch.logits[0] = 1; + batch.n_tokens = 1; + + const int rc = llama_decode(ctx, batch); + llama_batch_free(batch); + if (rc != 0) { + fprintf(stderr, "llama_decode failed: %d\n", rc); + return false; + } + const float * logits = llama_get_logits_ith(ctx, -1); + out.assign(logits, logits + n_vocab); + return true; +} + +// check the draft logits of a context against its own mask, and against full-vocabulary logits for the same inputs +static bool check_logits(const std::string & label, const std::vector & logits, const std::vector & mask, + const std::vector & ref) { + size_t n_cmp = 0; // finite logits compared with the reference + size_t n_bad_inf = 0; // -inf (or non-finite) where the context may draft + size_t n_bad_fin = 0; // finite where the context must not draft + size_t n_bad_val = 0; // finite, but different from the full-vocabulary reference + double max_diff = 0.0; + for (size_t t = 0; t < logits.size(); ++t) { + const float v = logits[t]; + if (mask[t]) { + if (!std::isfinite(v)) { + n_bad_inf++; + } else if (!ref.empty()) { + n_cmp++; + const double d = std::fabs((double) v - (double) ref[t]); + max_diff = std::max(max_diff, d); + if (d > 1e-2 * std::max(1.0, std::fabs((double) ref[t]))) { + n_bad_val++; + } + } + } else if (!(std::isinf(v) && v < 0)) { + n_bad_fin++; + } + } + size_t n_keep = 0; + for (uint8_t m : mask) { + n_keep += m; + } + const size_t i_max = std::max_element(logits.begin(), logits.end()) - logits.begin(); + const bool ok = n_bad_inf == 0 && n_bad_fin == 0 && n_bad_val == 0; + fprintf(stderr, " %-44s %s draftable %zu/%zu, non-finite inside %zu, finite outside %zu, argmax %zu (%.3f), " + "vs full vocab: %zu compared, %zu mismatches, max |diff| %.3g\n", + label.c_str(), ok ? "OK " : "FAIL", n_keep, mask.size(), n_bad_inf, n_bad_fin, i_max, logits[i_max], + n_cmp, n_bad_val, max_diff); + return ok; +} + +// decode the same inputs, interleaved, in every context plus a full-vocabulary reference context, and check each +static bool run_interleaved(llama_model * model, const std::string & name, std::vector & ctxs, int n_steps) { + const llama_vocab * vocab = llama_model_get_vocab(model); + bool ok = true; + for (int step = 0; step < n_steps; ++step) { + std::vector> logits(ctxs.size()); + for (size_t i = 0; i < ctxs.size(); ++i) { + if (!decode_step(model, ctxs[i].ctx.get(), step, logits[i])) { + return false; + } + } + const std::vector * ref = nullptr; + for (size_t i = 0; i < ctxs.size(); ++i) { + if (ctxs[i].n_vocab_draft == 0) { + ref = &logits[i]; + } + } + for (size_t i = 0; i < ctxs.size(); ++i) { + const std::string label = name + " step " + std::to_string(step) + " ctx " + std::string(1, char('A' + i)) + + " (N=" + std::to_string(ctxs[i].n_vocab_draft) + ")"; + ok &= check_logits(label, logits[i], draft_mask(vocab, ctxs[i].n_vocab_draft), + ref && ctxs[i].n_vocab_draft != 0 ? *ref : std::vector()); + } + } + return ok; +} + +// create contexts one after the other, in the given order, then decode them interleaved +static bool test_sequential(llama_model * model, const std::vector & ns, bool add_ref) { + std::string name; + std::vector ctxs; + for (int32_t n : ns) { + name += (name.empty() ? "" : ",") + std::to_string(n); + ctxs.push_back({ n, make_ctx(model, n) }); + if (!ctxs.back().ctx) { + fprintf(stderr, "failed to create context N=%d\n", n); + return false; + } + } + name = "create " + name; + if (add_ref) { + ctxs.push_back({ 0, make_ctx(model, 0) }); + } + fprintf(stderr, "%s:\n", name.c_str()); + return run_interleaved(model, name, ctxs, 3); +} + +// create two contexts from two threads at the same time, then decode them interleaved +static bool test_concurrent(llama_model * model, int32_t n1, int32_t n2) { + const std::string name = "concurrent " + std::to_string(n1) + "," + std::to_string(n2); + fprintf(stderr, "%s:\n", name.c_str()); + std::vector ctxs(2); + ctxs[0].n_vocab_draft = n1; + ctxs[1].n_vocab_draft = n2; + std::atomic go(false); + std::thread t0([&] { while (!go) { std::this_thread::yield(); } ctxs[0].ctx = make_ctx(model, n1); }); + std::thread t1([&] { while (!go) { std::this_thread::yield(); } ctxs[1].ctx = make_ctx(model, n2); }); + go = true; + t0.join(); + t1.join(); + if (!ctxs[0].ctx || !ctxs[1].ctx) { + fprintf(stderr, "failed to create contexts\n"); + return false; + } + ctxs.push_back({ 0, make_ctx(model, 0) }); + return run_interleaved(model, name, ctxs, 2); +} + +int main(int argc, char ** argv) { + std::setlocale(LC_NUMERIC, "C"); + + common_params params; + common_init(); + if (!common_params_parse(argc, argv, params, LLAMA_EXAMPLE_COMMON)) { + return 1; + } + + llama_backend_init(); + ggml_backend_load_all(); + + llama_model_params mparams = common_model_params_to_llama(params); + mparams.load_mtp = true; + llama_model_ptr model(llama_model_load_from_file(params.model.path.c_str(), mparams)); + if (!model) { + fprintf(stderr, "failed to load model\n"); + return 1; + } + if (llama_model_n_layer_nextn(model.get()) == 0) { + fprintf(stderr, "model has no MTP layers, skipping\n"); + return 0; + } + + const int32_t n_vocab = llama_vocab_n_tokens(llama_model_get_vocab(model.get())); + const int32_t n1 = std::min(65536, n_vocab / 2); + const int32_t n2 = std::min(32768, n_vocab / 4); + + bool ok = true; + ok &= test_sequential(model.get(), { n1, 0 }, false); // A = N then B = 0 + ok &= test_sequential(model.get(), { 0, n1 }, false); // A = 0 then B = N + ok &= test_sequential(model.get(), { n1, n2 }, true); // A = N1 then B = N2 + ok &= test_sequential(model.get(), { n2, n2 }, true); // two contexts with the same N + ok &= test_concurrent(model.get(), n1, n2); // simultaneous creation + // a subset freed with its last context and requested again + ok &= test_sequential(model.get(), { n2 }, true); + + fprintf(stderr, "%s\n", ok ? "PASS" : "FAIL"); + return ok ? 0 : 1; +} diff --git a/tests/test-ple-disk.cpp b/tests/test-ple-disk.cpp new file mode 100644 index 000000000000..8b116cebed41 --- /dev/null +++ b/tests/test-ple-disk.cpp @@ -0,0 +1,173 @@ +// llama_ple_disk must return, in every configuration, exactly the rows the CPU dequantizer +// produces for the same bytes: duplicates, cache hits, the serial paths and the pool alike. +#include "../src/llama-ple-disk.h" +#include "ggml.h" + +#include +#include +#include +#include +#include +#include +#include + +#if defined(_WIN32) +int main() { + printf("llama_ple_disk is not available on Windows, skipping\n"); + return 0; +} +#else + +#include + +struct table { + ggml_type type; + int64_t ne0; + int64_t nrows; + size_t offs; + size_t rs; + std::string fname; + std::vector bytes; // the rows as they follow the `offs` prefix bytes in the file +}; + +static table make_table(ggml_type type, int64_t ne0, int64_t nrows, std::mt19937 & rng) { + table t; + t.type = type; + t.ne0 = ne0; + t.nrows = nrows; + t.rs = ggml_row_size(type, ne0); + t.offs = 3 * 4096 + 77; // rows deliberately start off any I/O alignment + t.bytes.resize((size_t) nrows * t.rs); + for (auto & b : t.bytes) { + b = (uint8_t) rng(); + } + + const char * dir = getenv("TMPDIR"); + std::string path = std::string(dir ? dir : "/tmp") + "/test-ple-disk-XXXXXX"; + std::vector buf(path.begin(), path.end()); + buf.push_back(0); + const int fd = mkstemp(buf.data()); + GGML_ASSERT(fd >= 0); + std::vector prefix(t.offs); + for (auto & b : prefix) { + b = (uint8_t) rng(); + } + GGML_ASSERT(write(fd, prefix.data(), prefix.size()) == (ssize_t) prefix.size()); + size_t done = 0; + while (done < t.bytes.size()) { + const ssize_t w = write(fd, t.bytes.data() + done, t.bytes.size() - done); + GGML_ASSERT(w > 0); + done += (size_t) w; + } + close(fd); + t.fname = buf.data(); + return t; +} + +static void reference(const table & t, const int32_t * idx, size_t n, float * dst) { + const auto * traits = ggml_get_type_traits(t.type); + for (size_t k = 0; k < n; ++k) { + const uint8_t * src = t.bytes.data() + (size_t) idx[k] * t.rs; + float * out = dst + k * (size_t) t.ne0; + if (t.type == GGML_TYPE_F32) { + memcpy(out, src, t.rs); + } else { + traits->to_float(src, out, t.ne0); + } + } +} + +int main() { + { + struct ggml_init_params ip = { 16 * ggml_tensor_overhead(), nullptr, true }; + ggml_free(ggml_init(ip)); // conversion tables + } + + std::mt19937 rng(20260915); + const int64_t nrows = 50000; + const struct { ggml_type type; int64_t ne0; } types[] = { + { GGML_TYPE_IQ4_NL, 160 }, // the qwen4exp n-gram table layout + { GGML_TYPE_Q8_0, 64 }, + { GGML_TYPE_F16, 48 }, + { GGML_TYPE_F32, 40 }, // no dequantizer: rows are copied + }; + + size_t n_gathers = 0, n_rows = 0; + for (const auto & ty : types) { + const table t = make_table(ty.type, ty.ne0, nrows, rng); + + std::vector> patterns; + patterns.push_back({}); // nothing to do + patterns.push_back({ 7 }); + patterns.push_back({ 7, 7 }); // one distinct row, serial read + patterns.push_back({ 3, 1, 2 }); // three distinct rows, still serial + patterns.push_back({ 0, (int32_t) nrows - 1, 0, (int32_t) nrows - 1, 0 }); + { + std::vector v(257); // just past the serial dequantization limit + std::iota(v.begin(), v.end(), 100); + patterns.push_back(v); + } + { + std::vector v(5000); // descending, so sorting reorders everything + std::iota(v.rbegin(), v.rend(), 20000); + patterns.push_back(v); + } + { + std::vector v(20000); // random over the whole table + for (auto & r : v) { + r = (int32_t) (rng() % nrows); + } + patterns.push_back(v); + } + { + std::vector v(200000); // 500 distinct rows asked for 400 times each + for (auto & r : v) { + r = 1000 + (int32_t) (rng() % 500); + } + patterns.push_back(v); + } + { + std::vector v = patterns[7]; // half repeats a previous gather, half is new + for (size_t k = 0; k < v.size(); k += 2) { + v[k] = (int32_t) (rng() % nrows); + } + patterns.push_back(v); + } + + for (const int32_t threads : { 1, 3, 64 }) { + for (const size_t cache_mb : { (size_t) 0, (size_t) 1 }) { + for (const bool direct : { false, true }) { + llama_ple_disk::params p; + p.n_threads = threads; + p.cache_bytes = cache_mb << 20; + p.direct_io = direct; + llama_ple_disk disk(t.fname, t.offs, t.type, t.ne0, t.nrows, p); + GGML_ASSERT(disk.n_rows() == nrows && disk.ne0() == t.ne0 && disk.row_size() == t.rs); + + for (int pass = 0; pass < 2; ++pass) { // the second pass finds the first one's rows in the cache + for (const auto & idx : patterns) { + std::vector got(idx.size() * (size_t) t.ne0 + 1), want(got.size()); + memset(got.data(), 0xff, got.size() * sizeof(float)); + memset(want.data(), 0xff, want.size() * sizeof(float)); + reference(t, idx.data(), idx.size(), want.data()); + disk.gather(idx.data(), idx.size(), got.data()); + if (memcmp(got.data(), want.data(), got.size() * sizeof(float)) != 0) { + fprintf(stderr, "FAIL: %s, %zu rows, threads=%d cache=%zuMiB direct=%d pass=%d\n", + ggml_type_name(t.type), idx.size(), threads, cache_mb, (int) direct, pass); + unlink(t.fname.c_str()); + return 1; + } + n_gathers += 1; + n_rows += idx.size(); + } + } + } + } + } + unlink(t.fname.c_str()); + } + printf("PASS: %zu gathers, %zu rows, byte-identical to the CPU dequantizer across 4 types x 3 thread counts x cache on/off x direct/buffered\n", + n_gathers, n_rows); + return 0; +} +#endif diff --git a/tests/test-qsa-fast.cpp b/tests/test-qsa-fast.cpp new file mode 100644 index 000000000000..537953bec16f --- /dev/null +++ b/tests/test-qsa-fast.cpp @@ -0,0 +1,168 @@ +#include "../src/llama-memory-hybrid-idx.h" +#include "../src/prefix.h" +#include "ggml.h" +#include +#include +#include +#include + +// Host-side regression for the qsa_fast gate lifted from 127 to 512 tokens: +// incremental indexer-key pooling must select exactly the touched blocks, and +// must stay off when the prefix, the commit state or the size guard does not hold. +// Runs on CPU only; no model or backend is involved. + +// qsa_prefix.update-equivalent bookkeeping via qsa_prefix_state::apply +static void feed(qsa_prefix_state & s, int32_t start, int n) { + // position k lives in cell k (cells in order); apply() places them in sequence + std::vector slots(n); + for (int i = 0; i < n; ++i) slots[i] = start + i; + GGML_ASSERT(s.apply(0, start, slots)); +} + +// single-sequence ubatch of n tokens starting at position start +struct qsa_fast_ubatch { + std::vector tok; std::vector pos; std::vector n_seq_id; + std::vector> ids; std::vector seq_id; llama_ubatch u{}; + qsa_fast_ubatch(int32_t start, int n, int32_t seq = 0) { + tok.assign(n, 0); pos.resize(n); ids.assign(n, {seq}); + for (int i = 0; i < n; ++i) { pos[i] = start + i; n_seq_id.push_back(1); seq_id.push_back(ids[i].data()); } + u.n_tokens = n; u.n_pos = 1; u.token = tok.data(); u.pos = pos.data(); u.n_seq_id = n_seq_id.data(); u.seq_id = seq_id.data(); + } +}; + +// minimal stand-in for the memory object's fast-path pieces +struct qsa_fast_mock : public qsa_prefix_state { + qsa_fast_mock(size_t capacity = 0) : qsa_prefix_state(capacity) {} + void assign_from(const qsa_prefix_state & s) { static_cast(*this) = s; } + std::vector qsa_ready{0}; + std::vector qsa_keys{true}; + + const int max_tokens = getenv("QSA_FAST_MAX") ? atoi(getenv("QSA_FAST_MAX")) : 512; // read once, like the production static + bool fast(const llama_ubatch & u) const { + const bool prefix = qsa_prefix_matches(u); + return prefix && int(u.n_tokens) <= max_tokens && u.n_tokens <= 512 && qsa_keys.at(0) && + qsa_ready.at(0) >= int64_t(previous_size/4) && + cells.size()/4 >= (u.n_tokens+3)/4+1; + } + bool qsa_prefix_matches(const llama_ubatch & u) const { + if (!valid || !u.token || !u.pos || !u.n_tokens || !u.n_pos || !u.seq_id || !u.n_seq_id || + u.pos[0] != begin || int64_t(u.n_tokens) != end-begin) { return false; } + for (uint32_t i=0; i= count - 1); + std::vector ids = { complete }; + for (int b=first; b=last) ids.push_back(b); } + // in range, duplicate-free, and covers exactly [first, last) beyond the placeholder/filler blocks + std::set seen; + for (int id : ids) { GGML_ASSERT(id >= 0 && id <= complete); GGML_ASSERT(seen.insert(id).second); } + GGML_ASSERT(ids.size() == size_t(count)); + std::set covered(ids.begin()+1, ids.end()); + for (int b=first; b= count-1) + for (int n : {1, 4, 8, 127, 128, 255, 256, 511, 512, 513, 1024}) { + qsa_fast_mock s(8192); s.qsa_ready[0] = 0; + feed(s, 0, n); + const bool fast = s.fast(qsa_fast_ubatch(0, n).u); + GGML_ASSERT(!fast || (int(s.cells.size()/4) >= (n+3)/4+1)); + } + // established cache (full re-pool committed): fast iff n <= 512 + for (int n : {1, 4, 8, 126, 127, 128, 129, 255, 256, 257, 511, 512, 513, 514, 1024, 2048}) { + qsa_fast_mock s(8192); + feed(s, 0, 512); s.commit(); + feed(s, 512, n); + GGML_ASSERT(s.fast(qsa_fast_ubatch(512, n).u) == (n <= 512)); + } + // the last allowed boundary: previous commit covers the skipped blocks + { + qsa_fast_mock s(8192); + feed(s, 0, 512); s.commit(); + feed(s, 512, 512); + GGML_ASSERT(s.fast(qsa_fast_ubatch(512, 512).u)); + GGML_ASSERT(s.qsa_ready[0] >= int64_t(s.previous_size/4)); + } +} + +static void test_incremental_ids() { + for (int begin : {4, 128, 512, 1024}) { + for (int n : {1, 4, 8, 127, 128, 256, 512}) { + qsa_fast_mock s(8192); + feed(s, 0, begin); s.commit(); + feed(s, begin, n); + if (s.fast(qsa_fast_ubatch(begin, n).u)) { check_fill(s, n); } + } + } +} + +static void test_gate_failures() { + { // position must start at the tracked prefix end + qsa_fast_mock s(8192); feed(s, 0, 512); s.commit(); feed(s, 512, 256); + GGML_ASSERT(!s.fast(qsa_fast_ubatch(513, 255).u)); + GGML_ASSERT(!s.fast(qsa_fast_ubatch(256, 256).u)); + } + { // the prefix was not committed (recover/rollback path) -> full re-pool + qsa_fast_mock s(8192); feed(s, 0, 512); /* no commit */ feed(s, 512, 256); + GGML_ASSERT(!s.fast(qsa_fast_ubatch(512, 256).u)); + } + { // truncation below a block boundary re-pools once, then the fast path resumes on the + // committed smaller prefix (qsa_ready is clamped on seq_rm, llama-memory-hybrid-idx.cpp:203) + qsa_fast_mock s(8192); feed(s, 0, 512); s.commit(); s.truncate(200); + s.qsa_ready[0] = std::min(s.qsa_ready[0], s.cells.size()/4); + // production rebuilds the tracked prefix from the cache (prefix stays valid through seq_rm): + // cells [0,200) are rewritten in place (size stays 200), qsa_ready clamps to 200/4 = 50. + // qsa_fast runs against this pre-ubatch state: cells/4 = 50 < (256+3)/4+1 = 65 -> full re-pool. + GGML_ASSERT(int(s.cells.size()) == 200); + GGML_ASSERT(!s.fast(qsa_fast_ubatch(200, 256).u)); // one full re-pool + feed(s, 200, 256); // the full path commits the whole prefix + s.commit(); + feed(s, 456, 8); + GGML_ASSERT(s.fast(qsa_fast_ubatch(456, 8).u)); // fast resumes on the new prefix + } + { // multi-sequence ubatch never matches the prefix + qsa_fast_mock s(8192); feed(s, 0, 512); s.commit(); feed(s, 512, 256); + qsa_fast_ubatch u(512, 256); + u.u.n_seq_id[3] = 2; // break the single-sequence invariant + GGML_ASSERT(!s.fast(u.u)); + } + { // QSA_FAST_MAX: values <= 512 shrink the limit; 0 disables; >512 is capped by the hard limit. + // (The real guard caches the env in a function-local static; this mock re-reads it per call, + // and each case uses a fresh prefix so only the limit under test varies.) + setenv("QSA_FAST_MAX", "127", 1); + qsa_fast_mock a(8192), b(8192); + feed(a, 0, 512); a.commit(); feed(a, 512, 128); + GGML_ASSERT(!a.fast(qsa_fast_ubatch(512, 128).u)); + feed(b, 0, 512); b.commit(); feed(b, 512, 127); + GGML_ASSERT(b.fast(qsa_fast_ubatch(512, 127).u)); + setenv("QSA_FAST_MAX", "0", 1); + qsa_fast_mock b0(8192); + feed(b0, 0, 512); b0.commit(); feed(b0, 512, 127); + GGML_ASSERT(!b0.fast(qsa_fast_ubatch(512, 127).u)); + setenv("QSA_FAST_MAX", "9999", 1); // still capped at 512 + qsa_fast_mock c(8192), d(8192); + feed(c, 0, 512); c.commit(); feed(c, 512, 513); + GGML_ASSERT(!c.fast(qsa_fast_ubatch(512, 513).u)); + feed(d, 0, 512); d.commit(); feed(d, 512, 512); + GGML_ASSERT(d.fast(qsa_fast_ubatch(512, 512).u)); + unsetenv("QSA_FAST_MAX"); + } +} + +int main() { + test_size_boundaries(); + test_incremental_ids(); + test_gate_failures(); + std::cout << "PASS: qsa_fast gate boundaries (1..2048), incremental block selection, prefix/commit/truncate gates, QSA_FAST_MAX\n"; +} diff --git a/tests/test-qsa-prefix.cpp b/tests/test-qsa-prefix.cpp new file mode 100644 index 000000000000..05fb927f2683 --- /dev/null +++ b/tests/test-qsa-prefix.cpp @@ -0,0 +1,111 @@ +#include "../src/qsa-prefix-state.h" +#include "../src/prefix.h" +#include "ggml.h" +#include +#include + +// cells filled in order, one per position +struct qsa_cells { + llama_kv_cells cells; + uint32_t next = 0; + explicit qsa_cells(uint32_t n) { cells.resize(n); } + void add(llama_pos p0, llama_pos p1, std::vector seqs, llama_pos ext_dy = 0) { + for (llama_pos p=p0; p tok; std::vector pos; std::vector n_seq_id; + std::vector> ids; std::vector seq_id; llama_ubatch u{}; + qsa_ubatch(std::vector>> tokens, uint32_t n_pos = 1) { + const uint32_t n = tokens.size(); + tok.assign(n, 0); pos.resize(n*n_pos); ids.resize(n); + for (uint32_t i=0; i slots; + while (int(slots.size())8) { + int at=rng()%(s.cells.size()-4); + std::vector same(s.cells.begin()+at,s.cells.begin()+at+4); + auto old=s.cells;GGML_ASSERT(s.apply(0,at,same));GGML_ASSERT(s.cells==old); + } + } + } + s.reset(); GGML_ASSERT(!s.apply(0,2,{3})); + s.reset(); GGML_ASSERT(!s.apply(0,0,{3,3})); + s.reset(); GGML_ASSERT(s.apply(0,0,{3,7})); GGML_ASSERT(!s.apply(1,2,{4})); + s.reset(); GGML_ASSERT(s.apply(0,0,{3,7})); GGML_ASSERT(!s.apply(0,2,{3})); + s.reset(); GGML_ASSERT(s.apply(0,0,{3,7})); GGML_ASSERT(!s.apply(0,0,{9})); + s.reset(); GGML_ASSERT(s.apply(0,0,{3,7})); s.truncate(0); GGML_ASSERT(s.apply(1,0,{3})); + s.invalidate(); GGML_ASSERT(!s.apply(1,1,{4})); + std::cout << "PASS: 20000 randomized append/rollback steps, in-place rewrites, holes, duplicates, collisions, sequence changes\n"; +} diff --git a/tests/test-quantize-fns.cpp b/tests/test-quantize-fns.cpp index 570fca89a876..e25cf3236fa9 100644 --- a/tests/test-quantize-fns.cpp +++ b/tests/test-quantize-fns.cpp @@ -197,7 +197,10 @@ static int test_vec_dot_q(bool verbose) { type == GGML_TYPE_Q3_K ? MAX_QUANTIZATION_TOTAL_ERROR_3BITS : type == GGML_TYPE_IQ3_S ? MAX_QUANTIZATION_TOTAL_ERROR_3BITS : type == GGML_TYPE_IQ3_XXS ? MAX_QUANTIZATION_TOTAL_ERROR_3BITS_XXS : - type == GGML_TYPE_NVFP4 ? MAX_QUANTIZATION_TOTAL_ERROR_FP4 : MAX_QUANTIZATION_TOTAL_ERROR; + type == GGML_TYPE_NVFP4 ? MAX_QUANTIZATION_TOTAL_ERROR_FP4 : + // ROCmFPx sub-4-bit types: same budget as the other types of their width + type == GGML_TYPE_Q3_0_ROCMFPX ? MAX_QUANTIZATION_TOTAL_ERROR_3BITS : + type == GGML_TYPE_Q2_0_ROCMFPX ? MAX_QUANTIZATION_TOTAL_ERROR_2BITS : MAX_QUANTIZATION_TOTAL_ERROR; bool failed = !(total_error < max_quantization_error); num_failed += failed; if (failed || verbose) { @@ -213,11 +216,13 @@ static int test_vec_dot_q(bool verbose) { const float vec_dot_error = dot_product_error(qfns_cpu, type, test_size, test_data.data(), test_data2.data(), nullptr, nullptr, 1); const float max_allowed_error = type == GGML_TYPE_Q2_K || type == GGML_TYPE_IQ2_XS || type == GGML_TYPE_IQ2_XXS || - type == GGML_TYPE_IQ3_XXS || type == GGML_TYPE_IQ3_S || type == GGML_TYPE_IQ2_S + type == GGML_TYPE_IQ3_XXS || type == GGML_TYPE_IQ3_S || type == GGML_TYPE_IQ2_S || + type == GGML_TYPE_Q3_0_ROCMFPX ? MAX_DOT_PRODUCT_ERROR_LOWBIT : type == GGML_TYPE_Q1_0 ? MAX_DOT_PRODUCT_ERROR_BINARY - : type == GGML_TYPE_TQ1_0 || type == GGML_TYPE_TQ2_0 || type == GGML_TYPE_Q2_0 + : type == GGML_TYPE_TQ1_0 || type == GGML_TYPE_TQ2_0 || type == GGML_TYPE_Q2_0 || + type == GGML_TYPE_Q2_0_ROCMFPX ? MAX_DOT_PRODUCT_ERROR_TERNARY : type == GGML_TYPE_NVFP4 ? MAX_DOT_PRODUCT_ERROR_FP4 diff --git a/tests/test-reasoning-budget.cpp b/tests/test-reasoning-budget.cpp index 3bcc77e1733c..45b7c3da35fe 100644 --- a/tests/test-reasoning-budget.cpp +++ b/tests/test-reasoning-budget.cpp @@ -43,11 +43,14 @@ static void test_reasoning_budget( // For this test, we use nullptr as vocab since we're testing state transitions // The UTF-8 boundary check will treat all tokens as complete (safe fallback) auto * sampler = common_reasoning_budget_init( - nullptr, // vocab - not used for basic state machine tests + nullptr, start_seqs, end_seqs, forced_tokens, + {}, + {}, budget, + 0, initial_state ); @@ -156,7 +159,7 @@ static void test_reasoning_budget_clone_mid_counting() { const std::vector end = {101}; const std::vector forced = {102, 101}; - auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 2, REASONING_BUDGET_IDLE); + auto * sampler = common_reasoning_budget_init(nullptr, { start }, { end }, forced, {}, {}, 2, 0, REASONING_BUDGET_IDLE); llama_sampler_accept(sampler, 100); // COUNTING, remaining=2 llama_sampler_accept(sampler, 50); // COUNTING, remaining=1 @@ -175,7 +178,7 @@ static void test_reasoning_budget_clone_mid_forcing() { const std::vector end = {101}; const std::vector forced = {102, 101}; - auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 0, REASONING_BUDGET_FORCING); + auto * sampler = common_reasoning_budget_init(nullptr, { start }, { end }, forced, {}, {}, 0, 0, REASONING_BUDGET_FORCING); GGML_ASSERT(get_forced_token(sampler, 102) == 102); llama_sampler_accept(sampler, 102); // advance to the second forced token @@ -195,7 +198,7 @@ static void test_reasoning_budget_force_manual() { // if COUNTING, force() succeeds and begins forcing the end sequence from the start { - auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 5, REASONING_BUDGET_IDLE); + auto * sampler = common_reasoning_budget_init(nullptr, { start }, { end }, forced, {}, {}, 5, 0, REASONING_BUDGET_IDLE); llama_sampler_accept(sampler, 100); // COUNTING, remaining=5 llama_sampler_accept(sampler, 50); // COUNTING, remaining=4 @@ -216,7 +219,7 @@ static void test_reasoning_budget_force_manual() { // if IDLE, force() is a no-op { - auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 5, REASONING_BUDGET_IDLE); + auto * sampler = common_reasoning_budget_init(nullptr, { start }, { end }, forced, {}, {}, 5, 0, REASONING_BUDGET_IDLE); GGML_ASSERT(!common_reasoning_budget_force(sampler) && "force() must not transition from IDLE"); GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_IDLE); @@ -226,7 +229,7 @@ static void test_reasoning_budget_force_manual() { // if DONE, force() is a no-op { - auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 5, REASONING_BUDGET_IDLE); + auto * sampler = common_reasoning_budget_init(nullptr, { start }, { end }, forced, {}, {}, 5, 0, REASONING_BUDGET_IDLE); llama_sampler_accept(sampler, 100); // COUNTING llama_sampler_accept(sampler, 101); // natural end -> DONE @@ -240,7 +243,7 @@ static void test_reasoning_budget_force_manual() { // if FORCING, force() is a no-op and must not rewind the force position { - auto * sampler = common_reasoning_budget_init(nullptr, {start}, {end}, forced, 0, REASONING_BUDGET_FORCING); + auto * sampler = common_reasoning_budget_init(nullptr, { start }, { end }, forced, {}, {}, 0, 0, REASONING_BUDGET_FORCING); GGML_ASSERT(get_forced_token(sampler, 102) == 102); llama_sampler_accept(sampler, 102); // advance to the second forced token (force_pos=1) @@ -258,13 +261,302 @@ static void test_reasoning_budget_force_manual() { fprintf(stderr, " Test 'manual force transition' passed\n"); } +static void test_reasoning_budget_soft_warning_skipped_before_hard_cutoff() { + const std::vector start = { 100 }; + const std::vector end = { 101 }; + const std::vector forced = { 102, 101 }; + const std::vector soft_forced = { 200, 201 }; + + auto * sampler = common_reasoning_budget_init(nullptr, { start }, { end }, forced, { { 5, soft_forced } }, {}, 10, 0, REASONING_BUDGET_IDLE); + + llama_sampler_accept(sampler, 100); + for (llama_token t : { 50, 51, 52, 53 }) { + llama_sampler_accept(sampler, t); + } + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_COUNTING); + + llama_sampler_accept(sampler, 54); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_SOFT_PENDING); + + for (llama_token t : { 55, 56, 57, 58 }) { + llama_sampler_accept(sampler, t); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_SOFT_PENDING); + } + llama_sampler_accept(sampler, 59); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_FORCING); + GGML_ASSERT(get_forced_token(sampler, 202) == 102); + + llama_sampler_accept(sampler, 102); + GGML_ASSERT(get_forced_token(sampler, 202) == 101); + llama_sampler_accept(sampler, 101); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_DONE); + + llama_sampler_free(sampler); +} + +static void test_reasoning_budget_soft_forcing_resumes_counting() { + const std::vector start = { 100 }; + const std::vector end = { 101 }; + const std::vector forced = { 102, 101 }; + const std::vector soft_forced = { 200, 201 }; + + auto * sampler = common_reasoning_budget_init(nullptr, { start }, { end }, forced, { { 2, soft_forced } }, {}, 5, 0, REASONING_BUDGET_SOFT_FORCING); + + GGML_ASSERT(get_forced_token(sampler, 201) == 200); + llama_sampler_accept(sampler, 200); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_SOFT_FORCING); + + GGML_ASSERT(get_forced_token(sampler, 201) == 201); + llama_sampler_accept(sampler, 201); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_COUNTING); + + llama_sampler_free(sampler); +} + +static void test_reasoning_budget_force_manual_from_soft_states() { + const std::vector start = { 100 }; + const std::vector end = { 101 }; + const std::vector forced = { 102, 101 }; + const std::vector soft_forced = { 200, 201 }; + + { + auto * sampler = common_reasoning_budget_init(nullptr, { start }, { end }, forced, { { 5, soft_forced } }, {}, 10, 0, REASONING_BUDGET_IDLE); + + llama_sampler_accept(sampler, 100); + for (llama_token t : { 50, 51, 52, 53, 54 }) { + llama_sampler_accept(sampler, t); + } + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_SOFT_PENDING); + + GGML_ASSERT(common_reasoning_budget_force(sampler)); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_FORCING); + GGML_ASSERT(get_forced_token(sampler, 202) == 102); + + llama_sampler_free(sampler); + } + + { + auto * sampler = common_reasoning_budget_init(nullptr, { start }, { end }, forced, { { 2, soft_forced } }, {}, 5, 0, REASONING_BUDGET_SOFT_FORCING); + + llama_sampler_accept(sampler, 200); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_SOFT_FORCING); + + GGML_ASSERT(common_reasoning_budget_force(sampler)); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_FORCING); + GGML_ASSERT(get_forced_token(sampler, 202) == 102); + + llama_sampler_free(sampler); + } +} + +static void test_reasoning_budget_intro_forcing_then_counting() { + const std::vector start = { 100 }; + const std::vector end = { 101 }; + const std::vector forced = { 102, 101 }; + const std::vector intro_forced = { 300, 301 }; + + auto * sampler = common_reasoning_budget_init(nullptr, { start }, { end }, forced, {}, intro_forced, 3, 0, REASONING_BUDGET_IDLE); + + llama_sampler_accept(sampler, 100); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_INTRO_FORCING); + + GGML_ASSERT(get_forced_token(sampler, 301) == 300); + llama_sampler_accept(sampler, 300); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_INTRO_FORCING); + + GGML_ASSERT(get_forced_token(sampler, 301) == 301); + llama_sampler_accept(sampler, 301); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_COUNTING); + + llama_sampler_accept(sampler, 50); + llama_sampler_accept(sampler, 51); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_COUNTING); + llama_sampler_accept(sampler, 52); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_FORCING); + GGML_ASSERT(get_forced_token(sampler, 302) == 102); + + llama_sampler_free(sampler); +} + +static void test_reasoning_budget_intro_forcing_budget_zero() { + const std::vector start = { 100 }; + const std::vector end = { 101 }; + const std::vector forced = { 102, 101 }; + const std::vector intro_forced = { 300, 301 }; + + auto * sampler = common_reasoning_budget_init(nullptr, { start }, { end }, forced, {}, intro_forced, 0, 0, REASONING_BUDGET_IDLE); + + llama_sampler_accept(sampler, 100); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_INTRO_FORCING); + + llama_sampler_accept(sampler, 300); + llama_sampler_accept(sampler, 301); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_FORCING); + GGML_ASSERT(get_forced_token(sampler, 302) == 102); + + llama_sampler_free(sampler); +} + +static void test_reasoning_budget_force_manual_from_intro() { + const std::vector start = { 100 }; + const std::vector end = { 101 }; + const std::vector forced = { 102, 101 }; + const std::vector intro_forced = { 300, 301 }; + + auto * sampler = common_reasoning_budget_init(nullptr, { start }, { end }, forced, {}, intro_forced, 5, 0, REASONING_BUDGET_INTRO_FORCING); + + llama_sampler_accept(sampler, 300); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_INTRO_FORCING); + + GGML_ASSERT(common_reasoning_budget_force(sampler)); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_FORCING); + GGML_ASSERT(get_forced_token(sampler, 302) == 102); + + llama_sampler_free(sampler); +} + +static void test_reasoning_budget_intro_rearms_on_multiblock() { + const std::vector start = { 100 }; + const std::vector end = { 101 }; + const std::vector forced = { 102, 101 }; + const std::vector intro_forced = { 300, 301 }; + + auto * sampler = common_reasoning_budget_init(nullptr, { start }, { end }, forced, {}, intro_forced, 5, 0, REASONING_BUDGET_IDLE); + + llama_sampler_accept(sampler, 100); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_INTRO_FORCING); + llama_sampler_accept(sampler, 300); + llama_sampler_accept(sampler, 301); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_COUNTING); + llama_sampler_accept(sampler, 101); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_DONE); + + llama_sampler_accept(sampler, 100); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_INTRO_FORCING); + GGML_ASSERT(get_forced_token(sampler, 302) == 300); + + llama_sampler_free(sampler); +} + +static void test_reasoning_budget_intro_once_does_not_rearm() { + const std::vector start = { 100 }; + const std::vector end = { 101 }; + const std::vector forced = { 102, 101 }; + const std::vector intro_forced = { 300, 301 }; + + auto * sampler = common_reasoning_budget_init(nullptr, { start }, { end }, forced, {}, intro_forced, 5, 0, REASONING_BUDGET_IDLE, true); + + llama_sampler_accept(sampler, 100); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_INTRO_FORCING); + llama_sampler_accept(sampler, 300); + llama_sampler_accept(sampler, 301); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_COUNTING); + llama_sampler_accept(sampler, 101); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_DONE); + + // second reasoning block: intro must not fire again in once mode + llama_sampler_accept(sampler, 100); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_COUNTING); + + llama_sampler_free(sampler); +} + +static void test_reasoning_budget_hard_pending_grace_expires() { + const std::vector start = { 100 }; + const std::vector end = { 101 }; + const std::vector forced = { 102, 101 }; + + auto * sampler = common_reasoning_budget_init(nullptr, { start }, { end }, forced, {}, {}, 2, 3, REASONING_BUDGET_IDLE); + + llama_sampler_accept(sampler, 100); + llama_sampler_accept(sampler, 50); + llama_sampler_accept(sampler, 51); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_HARD_PENDING); + + llama_sampler_accept(sampler, 52); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_HARD_PENDING); + llama_sampler_accept(sampler, 53); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_HARD_PENDING); + llama_sampler_accept(sampler, 54); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_FORCING); + GGML_ASSERT(get_forced_token(sampler, 102) == 102); + + llama_sampler_free(sampler); +} + +static void test_reasoning_budget_hard_pending_natural_end() { + const std::vector start = { 100 }; + const std::vector end = { 101 }; + const std::vector forced = { 102, 101 }; + + auto * sampler = common_reasoning_budget_init(nullptr, { start }, { end }, forced, {}, {}, 2, 5, REASONING_BUDGET_IDLE); + + llama_sampler_accept(sampler, 100); + llama_sampler_accept(sampler, 50); + llama_sampler_accept(sampler, 51); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_HARD_PENDING); + + llama_sampler_accept(sampler, 101); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_DONE); + + llama_sampler_free(sampler); +} + +static void test_reasoning_budget_force_manual_from_hard_pending() { + const std::vector start = { 100 }; + const std::vector end = { 101 }; + const std::vector forced = { 102, 101 }; + + auto * sampler = common_reasoning_budget_init(nullptr, { start }, { end }, forced, {}, {}, 2, 10, REASONING_BUDGET_IDLE); + + llama_sampler_accept(sampler, 100); + llama_sampler_accept(sampler, 50); + llama_sampler_accept(sampler, 51); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_HARD_PENDING); + + GGML_ASSERT(common_reasoning_budget_force(sampler)); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_FORCING); + GGML_ASSERT(get_forced_token(sampler, 102) == 102); + + llama_sampler_free(sampler); +} + +static void test_reasoning_budget_soft_pending_exhaustion_uses_grace() { + const std::vector start = { 100 }; + const std::vector end = { 101 }; + const std::vector forced = { 102, 101 }; + const std::vector soft_forced = { 200, 201 }; + + auto * sampler = common_reasoning_budget_init(nullptr, { start }, { end }, forced, { { 5, soft_forced } }, {}, 10, 2, REASONING_BUDGET_IDLE); + + llama_sampler_accept(sampler, 100); + for (llama_token t : { 50, 51, 52, 53, 54 }) { + llama_sampler_accept(sampler, t); + } + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_SOFT_PENDING); + + for (llama_token t : { 55, 56, 57, 58 }) { + llama_sampler_accept(sampler, t); + } + llama_sampler_accept(sampler, 59); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_HARD_PENDING); + + llama_sampler_accept(sampler, 60); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_HARD_PENDING); + llama_sampler_accept(sampler, 61); + GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_FORCING); + GGML_ASSERT(get_forced_token(sampler, 202) == 102); + + llama_sampler_free(sampler); +} + static void test_reasoning_budget_end_match() { const std::vector start = {{100}}; const std::vector end = {{101}, {103, 104}}; // natural end records the sequence that matched; re-arming clears it { - auto * sampler = common_reasoning_budget_init(nullptr, start, end, {102, 101}, 5, REASONING_BUDGET_IDLE); + auto * sampler = common_reasoning_budget_init(nullptr, start, end, { 102, 101 }, {}, {}, 5, 0, REASONING_BUDGET_IDLE); GGML_ASSERT(common_reasoning_budget_get_end_match(sampler) == nullptr); @@ -287,7 +579,7 @@ static void test_reasoning_budget_end_match() { { const std::vector end_overlap = {{104}, {103, 104}}; - auto * sampler = common_reasoning_budget_init(nullptr, start, end_overlap, {102, 104}, 5, REASONING_BUDGET_IDLE); + auto * sampler = common_reasoning_budget_init(nullptr, start, end_overlap, { 102, 104 }, {}, {}, 5, 0, REASONING_BUDGET_IDLE); llama_sampler_accept(sampler, 100); // COUNTING llama_sampler_accept(sampler, 103); @@ -302,7 +594,7 @@ static void test_reasoning_budget_end_match() { // forcing records the end sequence terminating forced_tokens { - auto * sampler = common_reasoning_budget_init(nullptr, start, end, {102, 103, 104}, 0, REASONING_BUDGET_FORCING); + auto * sampler = common_reasoning_budget_init(nullptr, start, end, { 102, 103, 104 }, {}, {}, 0, 0, REASONING_BUDGET_FORCING); llama_sampler_accept(sampler, 102); llama_sampler_accept(sampler, 103); @@ -318,7 +610,7 @@ static void test_reasoning_budget_end_match() { // forced_tokens not ending with a known end sequence records nothing { - auto * sampler = common_reasoning_budget_init(nullptr, start, end, {102}, 0, REASONING_BUDGET_FORCING); + auto * sampler = common_reasoning_budget_init(nullptr, start, end, { 102 }, {}, {}, 0, 0, REASONING_BUDGET_FORCING); llama_sampler_accept(sampler, 102); // forced sequence complete, DONE GGML_ASSERT(common_reasoning_budget_get_state(sampler) == REASONING_BUDGET_DONE); @@ -493,9 +785,21 @@ int main(void) { test_reasoning_budget_clone_mid_counting(); test_reasoning_budget_clone_mid_forcing(); test_reasoning_budget_force_manual(); + test_reasoning_budget_soft_warning_skipped_before_hard_cutoff(); + test_reasoning_budget_soft_forcing_resumes_counting(); + test_reasoning_budget_force_manual_from_soft_states(); + test_reasoning_budget_intro_forcing_then_counting(); + test_reasoning_budget_intro_forcing_budget_zero(); + test_reasoning_budget_force_manual_from_intro(); + test_reasoning_budget_intro_rearms_on_multiblock(); + test_reasoning_budget_intro_once_does_not_rearm(); + test_reasoning_budget_hard_pending_grace_expires(); + test_reasoning_budget_hard_pending_natural_end(); + test_reasoning_budget_force_manual_from_hard_pending(); + test_reasoning_budget_soft_pending_exhaustion_uses_grace(); test_reasoning_budget_end_match(); - printf("OK (12 tests passed)\n"); + printf("OK (23 tests passed)\n"); printf("Testing UTF-8 boundary detection... "); test_utf8_boundary_detection(); diff --git a/tests/test-save-load-state.cpp b/tests/test-save-load-state.cpp index 74d1ba6c213b..86970be7b201 100644 --- a/tests/test-save-load-state.cpp +++ b/tests/test-save-load-state.cpp @@ -360,14 +360,14 @@ static bool test_seq_cp_device(struct llama_model * model, const struct common_p // - save the seq 1 state, free the interleaved seq 0 cells, and restore via the given io path // - the restore destination is non-contiguous: scatter reads are batched per contiguous run // - save again on the host and compare the two blobs byte for byte -static bool test_seq_cp_scatter(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens, int test_num, bool on_device) { +static bool test_seq_cp_scatter(struct llama_model * model, const struct common_params & params, const llama_tokens & tokens, int test_num, bool on_device, bool compact = false) { auto params_ctx = common_context_params_to_llama(params); params_ctx.n_ctx = 256; params_ctx.n_seq_max = 2; params_ctx.kv_unified = true; auto ctx = llama_context_ptr{llama_init_from_model(model, params_ctx)}; - LOG("\n=== Test %d: seq copy (%s, scatter) ===\n", test_num, on_device ? "device" : "host"); + LOG("\n=== Test %d: seq copy (%s, %s) ===\n", test_num, on_device ? "device" : "host", compact ? "scatter to compact" : "scatter"); const uint32_t flags = on_device ? LLAMA_STATE_SEQ_FLAGS_ON_DEVICE : LLAMA_STATE_SEQ_FLAGS_NONE; @@ -425,6 +425,11 @@ static bool test_seq_cp_scatter(struct llama_model * model, const struct common_ return false; } + if (compact) { + // Saved ranges [2, 1] must restore into one contiguous range, including quantized KV blocks. + llama_memory_clear(llama_get_memory(ctx.get()), true); + } + // restore via the io path under test const size_t nset = llama_state_seq_set_data_ext(ctx.get(), state_save.data(), state_save.size(), 1, flags); if (nset != state_save.size()) { @@ -508,7 +513,7 @@ static bool test_state_roundtrip(struct llama_model * model, const struct common } -// Run the full save/load test suite (tests 1-8) for a single model. +// Run the full save/load test suite (tests 1-9) for a single model. // Returns true if all tests pass, false otherwise. static bool run_save_load_tests_for_model(const std::string & model_path, const struct common_params & base_params) { struct common_params params = base_params; @@ -590,6 +595,11 @@ static bool run_save_load_tests_for_model(const std::string & model_path, const return false; } + // Test 9: on-device restore with different source and destination chunking. + if (!test_seq_cp_scatter(model, params, tokens, 9, true, true)) { + return false; + } + LOG("\nAll tests passed.\n"); return true; diff --git a/tools/llama-bench/llama-bench.cpp b/tools/llama-bench/llama-bench.cpp index 17adda38091f..963de9e2e685 100644 --- a/tools/llama-bench/llama-bench.cpp +++ b/tools/llama-bench/llama-bench.cpp @@ -279,6 +279,8 @@ static const char * lazy_mode_str(llama_lazy_mode mode) { return "auto"; case LLAMA_LAZY_MODE_ON: return "on"; + case LLAMA_LAZY_MODE_DIRECT: + return "on-direct"; default: GGML_ABORT("invalid lazy mode"); } @@ -475,7 +477,7 @@ static void print_usage(int /* argc */, char ** argv) { printf(" -fa, --flash-attn (default: %s)\n", join(transform_to_str(cmd_params_defaults.flash_attn, llama_flash_attn_type_name), ",").c_str()); printf(" -dev, --device (default: auto)\n"); printf(" -lm, --load-mode (default: %s)\n", join(transform_to_str(cmd_params_defaults.load_mode, llama_load_mode_name), ",").c_str()); - printf(" -lzm, --lazy-mode (default: %s)\n", join(transform_to_str(cmd_params_defaults.lazy_mode, lazy_mode_str), ",").c_str()); + printf(" -lzm, --lazy-mode (default: %s)\n", join(transform_to_str(cmd_params_defaults.lazy_mode, lazy_mode_str), ",").c_str()); printf(" -embd, --embeddings <0|1> (default: %s)\n", join(cmd_params_defaults.embeddings, ",").c_str()); printf(" -ts, --tensor-split (default: 0)\n"); printf(" -ot --override-tensor =;...\n"); @@ -812,6 +814,8 @@ static cmd_params parse_cmd_params(int argc, char ** argv) { llama_lazy_mode mode; if (m == "on") { mode = LLAMA_LAZY_MODE_ON; + } else if (m == "on-direct") { + mode = LLAMA_LAZY_MODE_DIRECT; } else if (m == "auto") { mode = LLAMA_LAZY_MODE_AUTO; } else if (m == "off") { diff --git a/tools/quantize/quantize.cpp b/tools/quantize/quantize.cpp index 38950036cd82..f14788a71e0c 100644 --- a/tools/quantize/quantize.cpp +++ b/tools/quantize/quantize.cpp @@ -35,6 +35,24 @@ static const std::vector QUANT_OPTIONS = { { "Q1_0", LLAMA_FTYPE_MOSTLY_Q1_0, " 1.125 bpw quantization", }, { "Q2_0", LLAMA_FTYPE_MOSTLY_Q2_0, " 2.25 bpw quantization (group 64)", }, { "Q4_0", LLAMA_FTYPE_MOSTLY_Q4_0, " 4.34G, +0.4685 ppl @ Llama-3-8B", }, + { "Q4_0_ROCMFP4", LLAMA_FTYPE_MOSTLY_Q4_0_ROCMFP4, " 4.50 bpw ROCmFP4 UE4M3-scale experimental", }, + { "Q4_0_ROCMFP4_EVEN", LLAMA_FTYPE_MOSTLY_Q4_0_ROCMFP4, " 4.50 bpw ROCmFP4 even tensor conversion; implies --pure", }, + { "Q4_0_ROCMFP4_LEAN", LLAMA_FTYPE_MOSTLY_Q4_0_ROCMFP4_LEAN, " 4.60 bpw ROCmFP4 + Q5_K token embeddings", }, + { "Q4_0_ROCMFP4_COHERENT", LLAMA_FTYPE_MOSTLY_Q4_0_ROCMFP4_COHERENT, " 4.70 bpw ROCmFP4 + Q6_K token embeddings", }, + { "Q4_0_ROCMFP4_FAST", LLAMA_FTYPE_MOSTLY_Q4_0_ROCMFP4_FAST, " 4.25 bpw ROCmFP4 single-scale speed layout", }, + { "Q4_0_ROCMFP4_FAST_EVEN", LLAMA_FTYPE_MOSTLY_Q4_0_ROCMFP4_FAST, " 4.25 bpw ROCmFP4 fast even tensor conversion; implies --pure", }, + { "Q4_0_ROCMFP4_FAST_COHERENT", LLAMA_FTYPE_MOSTLY_Q4_0_ROCMFP4_FAST_COHERENT, " ~4.45 bpw ROCmFP4 fast + Q6_K token embeddings", }, + { "Q4_0_ROCMFP4_STRIX", LLAMA_FTYPE_MOSTLY_Q4_0_ROCMFP4_STRIX, " ~4.49 bpw ROCmFP4 Strix Halo attn-K/V quality recipe", }, + { "Q4_0_ROCMFP4_STRIX_LEAN", LLAMA_FTYPE_MOSTLY_Q4_0_ROCMFP4_STRIX_LEAN, " ~4.38 bpw ROCmFP4 Strix K/V + Q5_K token embeddings", }, + { "Q3_0_ROCMFPX", LLAMA_FTYPE_MOSTLY_Q3_0_ROCMFPX, " 3.50 bpw ROCmFPx experimental, ROCm/Vulkan staging", }, + { "Q2_0_ROCMFPX", LLAMA_FTYPE_MOSTLY_Q2_0_ROCMFPX, " 2.50 bpw ROCmFPx S40 codebook + dual UE4M3 scales", }, + { "Q6_0_ROCMFPX", LLAMA_FTYPE_MOSTLY_Q6_0_ROCMFPX, " 6.50 bpw ROCmFPx experimental, ROCm/Vulkan staging", }, + { "Q8_0_ROCMFPX", LLAMA_FTYPE_MOSTLY_Q8_0_ROCMFPX, " 8.25 bpw ROCmFPx experimental, ROCm/Vulkan staging", }, + { "Q3_0_ROCMFPX_AGENT", LLAMA_FTYPE_MOSTLY_Q3_0_ROCMFPX_AGENT, " agent/tool-call coherent ROCmFPx Q3 routing", }, + { "Q6_0_ROCMFPX_AGENT", LLAMA_FTYPE_MOSTLY_Q6_0_ROCMFPX_AGENT, " agent/tool-call coherent ROCmFPx Q6 routing", }, + { "Q8_0_ROCMFPX_AGENT", LLAMA_FTYPE_MOSTLY_Q8_0_ROCMFPX_AGENT, " agent/tool-call coherent ROCmFPx Q8 routing", }, + { "Q6_0_ROCMFPX_LEAN", LLAMA_FTYPE_MOSTLY_Q6_0_ROCMFPX_LEAN, " size/speed-biased ROCmFPx Q6 routing", }, + { "Q6_0_ROCMFPX_AGENT_LEAN", LLAMA_FTYPE_MOSTLY_Q6_0_ROCMFPX_AGENT_LEAN, " agent ROCmFPx Q6 routing without Q8-heavy boosts", }, { "Q4_1", LLAMA_FTYPE_MOSTLY_Q4_1, " 4.78G, +0.4511 ppl @ Llama-3-8B", }, { "MXFP4_MOE",LLAMA_FTYPE_MOSTLY_MXFP4_MOE," MXFP4 MoE", }, { "Q5_0", LLAMA_FTYPE_MOSTLY_Q5_0, " 5.21G, +0.1316 ppl @ Llama-3-8B", }, @@ -584,6 +602,9 @@ int llama_quantize(int argc, char ** argv) { if (ftype_str == "COPY") { params.only_copy = true; } + if (ftype_str == "Q4_0_ROCMFP4_EVEN" || ftype_str == "Q4_0_ROCMFP4_FAST_EVEN") { + params.pure = true; + } } else { // argv[arg_idx] is not a valid ftype, so treat it as output path: fname_out = argv[arg_idx]; @@ -603,6 +624,9 @@ int llama_quantize(int argc, char ** argv) { if (ftype_str == "COPY") { params.only_copy = true; } + if (ftype_str == "Q4_0_ROCMFP4_EVEN" || ftype_str == "Q4_0_ROCMFP4_FAST_EVEN") { + params.pure = true; + } arg_idx++; } diff --git a/tools/server/README.md b/tools/server/README.md index ef9033404824..1e6147cff773 100644 --- a/tools/server/README.md +++ b/tools/server/README.md @@ -232,8 +232,16 @@ For the full list of features, please refer to [server's changelog](https://gith | `--reasoning-format FORMAT` | controls whether thought tags are allowed and/or extracted from the response, and in which format they're returned; one of:
- none: leaves thoughts unparsed in `message.content`
- deepseek: puts thoughts in `message.reasoning_content`
- deepseek-legacy: keeps `` tags in `message.content` while also populating `message.reasoning_content`
(default: auto)
(env: LLAMA_ARG_THINK) | | `-rea, --reasoning [on\|off\|auto]` | Use reasoning/thinking in the chat ('on', 'off', or 'auto', default: 'auto' (detect from template))
(env: LLAMA_ARG_REASONING) | | `--reasoning-effort LEVEL` | reasoning effort level given to the chat template: 'default' to keep the template default,
or a level such as 'minimal', 'low', 'medium', 'high', 'xhigh' or 'max' (default: default)
(env: LLAMA_ARG_REASONING_EFFORT) | +| `--reasoning-budget-enable, --no-reasoning-budget-enable` | master switch for hard cutoff, soft warnings, intro message, grace period, and runtime reasoning control (default: disabled)
(env: LLAMA_ARG_THINK_BUDGET_ENABLE) | | `--reasoning-budget N` | token budget for thinking: -1 for unrestricted, 0 for immediate end, N>0 for token budget (default: -1)
(env: LLAMA_ARG_THINK_BUDGET) | -| `--reasoning-budget-message MESSAGE` | message injected before the end-of-thinking tag when reasoning budget is exhausted (default: none)
(env: LLAMA_ARG_THINK_BUDGET_MESSAGE) | +| `--reasoning-budget-message MESSAGE` | message forced when the reasoning budget is exhausted; include the model's closing tag when required by the template (default: none)
(env: LLAMA_ARG_THINK_BUDGET_MESSAGE) | +| `--reasoning-budget-soft-ratio N` | fraction of the budget consumed before the first soft warning; <= 0 disables (default: -1) | +| `--reasoning-budget-soft-message MESSAGE` | first soft-warning message, injected at a newline boundary | +| `--reasoning-budget-soft2-ratio N` | fraction of the budget consumed before the second soft warning; <= 0 disables (default: -1) | +| `--reasoning-budget-soft2-message MESSAGE` | second soft-warning message, injected at a newline boundary | +| `--reasoning-budget-intro-mode MODE` | `every` forces the intro for every reasoning block; `once` suppresses it when already present in the prompt (default: every) | +| `--reasoning-budget-intro-message MESSAGE` | message forced at the start of reasoning; `{budget}` is replaced by the configured budget | +| `--reasoning-budget-grace-tokens N` | tokens to allow after budget exhaustion while waiting for a paragraph break (default: 0) | | `--reasoning-preserve, --no-reasoning-preserve` | preserve reasoning trace in the full history, not just the last assistant message (default: enabled)
compatible with certain templates having 'supports_preserve_reasoning' capability
example: https://docs.z.ai/guides/capabilities/thinking-mode#preserved-thinking
(env: LLAMA_ARG_REASONING_PRESERVE) | | `--chat-template JINJA_TEMPLATE` | set custom jinja chat template (default: template taken from model's metadata)
if suffix/prefix are specified, template will be disabled
only commonly used templates are accepted (unless --jinja is set before this flag):
list of built-in templates:
bailing, bailing-think, bailing2, chatglm3, chatglm4, chatml, command-r, deepseek, deepseek-ocr, deepseek2, deepseek3, exaone-moe, exaone3, exaone4, falcon3, gemma, gigachat, glmedge, gpt-oss, granite, granite-4.0, granite-4.1, grok-2, hunyuan-dense, hunyuan-moe, hunyuan-vl, kimi-k2, llama2, llama2-sys, llama2-sys-bos, llama2-sys-strip, llama3, llama4, megrez, minicpm, mistral-v1, mistral-v3, mistral-v3-tekken, mistral-v7, mistral-v7-tekken, monarch, openchat, orion, pangu-embedded, phi3, phi4, rwkv-world, seed_oss, smolvlm, solar-open, vicuna, vicuna-orca, yandex, zephyr
(env: LLAMA_ARG_CHAT_TEMPLATE) | | `--chat-template-file JINJA_TEMPLATE_FILE` | set custom jinja chat template file (default: template taken from model's metadata)
if suffix/prefix are specified, template will be disabled
only commonly used templates are accepted (unless --jinja is set before this flag):
list of built-in templates:
bailing, bailing-think, bailing2, chatglm3, chatglm4, chatml, command-r, deepseek, deepseek-ocr, deepseek2, deepseek3, exaone-moe, exaone3, exaone4, falcon3, gemma, gigachat, glmedge, gpt-oss, granite, granite-4.0, granite-4.1, grok-2, hunyuan-dense, hunyuan-moe, hunyuan-vl, kimi-k2, llama2, llama2-sys, llama2-sys-bos, llama2-sys-strip, llama3, llama4, megrez, minicpm, mistral-v1, mistral-v3, mistral-v3-tekken, mistral-v7, mistral-v7-tekken, monarch, openchat, orion, pangu-embedded, phi3, phi4, rwkv-world, seed_oss, smolvlm, solar-open, vicuna, vicuna-orca, yandex, zephyr
(env: LLAMA_ARG_CHAT_TEMPLATE_FILE) | @@ -1323,6 +1331,8 @@ The `response_format` parameter supports both plain JSON output (e.g. `{"type": `reasoning_control`: Arms realtime reasoning control for this completion so it can be ended early via `/v1/chat/completions/control`. Defaults to `false`. +`reasoning_budget_enabled`: Explicit master switch for reasoning budget, soft-warning, intro, grace, and runtime reasoning control fields. Omit it to keep the server or CLI default. + `generation_prompt`: The generation prompt that was prefilled in by the template. Prepended to model output before parsing. `parse_tool_calls`: Whether to parse the generated tool call. @@ -1431,7 +1441,11 @@ The response also includes a standard `usage` object: "total_tokens": 92, "prompt_tokens_details": { "cached_tokens": 0 - } + }, + "completion_tokens_details": { + "reasoning_tokens": 12 + }, + "reasoning": 12 } } ``` @@ -1442,6 +1456,10 @@ The server supports parsing and returning reasoning via the `reasoning_content` Reasoning input (preserve reasoning in history) is also supported by some specific templates. For more details, please refer to [PR#18994](https://github.com/ggml-org/llama.cpp/pull/18994). +Reasoning budget controls can force a hard cutoff, insert soft warnings, insert an intro message, and wait for a paragraph break after budget exhaustion. Set `reasoning_budget_enabled` in a request to override the server or CLI default. + +Chat Completions usage reports reasoning tokens in `usage.reasoning` and `usage.completion_tokens_details.reasoning_tokens`. Responses usage reports them in `usage.reasoning` and `usage.output_tokens_details.reasoning_tokens`. + ### POST `/v1/chat/completions/control`: Control a running chat completion in real time Acts on an in-flight completion identified by its `id` (the `id` field streamed back by `/v1/chat/completions`). The request is processed in parallel with the SSE stream, so the client sends it while still reading tokens. @@ -1450,7 +1468,7 @@ Acts on an in-flight completion identified by its `id` (the `id` field streamed `id`: (Required) The chat completion id to act on. A completion that has already finished matches nothing and the call is a no-op. -`action`: (Required) The control action to perform. Currently the only supported value is `reasoning_end`, which forces the end of the current reasoning block so the model moves on to the final answer. Requires `reasoning_control: true` on the original completion request. +`action`: (Required) The control action to perform. Currently the only supported value is `reasoning_end`, which forces the end of the current reasoning block so the model moves on to the final answer. Requires the reasoning budget master switch and `reasoning_control: true` on the original completion request. `model`: (Required in router mode) The model name, used to route the request to the right instance. Ignored in single model mode. @@ -1767,6 +1785,15 @@ We also offer additional options that are exclusive to presets (these aren't tre - `load-on-startup` (boolean): Controls whether the model loads automatically when the server starts. Only applies at startup: if the model list is reloaded later (for example after editing the preset file), a newly added model is listed but not loaded - `stop-timeout` (int, seconds): After requested unload, wait for this many seconds before forcing termination (default: 10) - `dedup-cache-models` (boolean): When the preset uses `hf-repo` pointing to a model that is already downloaded, hide the corresponding cached model entry from `GET /models` (the preset entry remains visible). Set it in the `[*]` section to apply to all presets. +- `hidden` (boolean): Omit this model from `GET /models` and `GET /v1/models`. The model can still be loaded and used by requesting it by name, so this only affects listings such as the model picker in the web UI. Useful for entries that are not meant to be picked directly, e.g. a draft/MTP GGUF that is discovered from the HF cache: + + ```ini + ; hides the auto-discovered cache entry, which is only useful as a draft model + [ggml-org/MY-MODEL-MTP-GGUF:Q4_K_M] + hidden = 1 + ``` + + A section whose name matches an existing model merges into that model's entry rather than creating a new one, so a stub section with just `hidden = 1` is enough to hide a model that came from the HF cache or from `--models-dir`. ### Routing requests diff --git a/tools/server/server-common.cpp b/tools/server/server-common.cpp index 7bf1138c8a8e..ce1798e95864 100644 --- a/tools/server/server-common.cpp +++ b/tools/server/server-common.cpp @@ -563,7 +563,7 @@ void server_tokens::insert(const llama_tokens & inp_tokens) { } const llama_tokens & server_tokens::get_tokens() const { - GGML_ASSERT(!has_mtmd); + GGML_ASSERT(!has_media()); return tokens; } @@ -646,7 +646,7 @@ llama_tokens server_tokens::get_text_tokens() const { } void server_tokens::set_token(llama_pos pos, llama_token id) { - GGML_ASSERT(!has_mtmd); // only allow this if mtmd is disabled + GGML_ASSERT(!has_media()); tokens[pos] = id; } @@ -1391,12 +1391,42 @@ json oaicompat_chat_params_parse( reasoning_budget = opt.reasoning_budget; } + const bool reasoning_control = json_value(body, "reasoning_control", false); + if (body.contains("reasoning_budget_enabled")) { + llama_params["reasoning_budget_enabled"] = body.at("reasoning_budget_enabled"); + } + if (!chat_params.thinking_end_tags.empty()) { llama_params["reasoning_budget_tokens"] = reasoning_budget; llama_params["reasoning_budget_start_tag"] = chat_params.thinking_start_tag; llama_params["reasoning_budget_end_tags"] = chat_params.thinking_end_tags; - llama_params["reasoning_budget_message"] = json_value(body, "reasoning_budget_message", opt.reasoning_budget_message); - llama_params["reasoning_control"] = json_value(body, "reasoning_control", false); + llama_params["reasoning_control"] = reasoning_control; + + const bool has_per_request_reasoning = + reasoning_budget >= 0 || reasoning_control || body.contains("reasoning_budget_message") || + body.contains("reasoning_budget_soft_ratio") || body.contains("reasoning_budget_soft_message") || + body.contains("reasoning_budget_soft2_ratio") || body.contains("reasoning_budget_soft2_message") || + body.contains("reasoning_budget_intro_message") || body.contains("reasoning_budget_intro_mode") || + body.contains("reasoning_budget_grace_tokens"); + + if (has_per_request_reasoning) { + llama_params["reasoning_budget_message"] = + json_value(body, "reasoning_budget_message", opt.reasoning_budget_message); + llama_params["reasoning_budget_soft_ratio"] = + json_value(body, "reasoning_budget_soft_ratio", opt.reasoning_budget_soft_ratio); + llama_params["reasoning_budget_soft_message"] = + json_value(body, "reasoning_budget_soft_message", opt.reasoning_budget_soft_message); + llama_params["reasoning_budget_soft2_ratio"] = + json_value(body, "reasoning_budget_soft2_ratio", opt.reasoning_budget_soft2_ratio); + llama_params["reasoning_budget_soft2_message"] = + json_value(body, "reasoning_budget_soft2_message", opt.reasoning_budget_soft2_message); + llama_params["reasoning_budget_intro_message"] = + json_value(body, "reasoning_budget_intro_message", opt.reasoning_budget_intro_message); + llama_params["reasoning_budget_intro_mode"] = + json_value(body, "reasoning_budget_intro_mode", opt.reasoning_budget_intro_mode); + llama_params["reasoning_budget_grace_tokens"] = + json_value(body, "reasoning_budget_grace_tokens", opt.reasoning_budget_grace_tokens); + } } } diff --git a/tools/server/server-common.h b/tools/server/server-common.h index 9894f5f06fb0..6c95a12a708b 100644 --- a/tools/server/server-common.h +++ b/tools/server/server-common.h @@ -222,6 +222,8 @@ struct server_tokens { bool empty() const { return tokens.empty(); } + bool has_media() const { return !map_idx_to_media.empty(); } + void clear() { map_idx_to_media.clear(); tokens.clear(); @@ -317,6 +319,13 @@ struct server_chat_params { bool enable_thinking = true; int reasoning_budget = -1; std::string reasoning_budget_message; + float reasoning_budget_soft_ratio = -1.0f; + std::string reasoning_budget_soft_message; + float reasoning_budget_soft2_ratio = -1.0f; + std::string reasoning_budget_soft2_message; + std::string reasoning_budget_intro_message; + std::string reasoning_budget_intro_mode = "every"; + int reasoning_budget_grace_tokens = 0; std::string media_path; bool force_pure_content = false; }; diff --git a/tools/server/server-context.cpp b/tools/server/server-context.cpp index b6835e43459e..589493b9839e 100644 --- a/tools/server/server-context.cpp +++ b/tools/server/server-context.cpp @@ -14,6 +14,8 @@ #include "log.h" #include "sampling.h" #include "speculative.h" +#include "speculative-prefill.h" +#include "src/llama-ext.h" #include "mtmd.h" #include "mtmd-helper.h" @@ -241,6 +243,7 @@ struct server_slot { llama_context * ctx_tgt = nullptr; llama_context * ctx_dft = nullptr; + llama_context * ctx_spf = nullptr; common_memory mem; @@ -258,6 +261,9 @@ struct server_slot { bool spec_is_replay = false; std::mt19937 spec_synth_rng; + bool spec_prefill_active = false; + server_tokens spec_prefill_tokens; + // TODO: move members that belong to the task (such as `generated_text`, `has_new_line`) to task_results_state // see https://github.com/ggml-org/llama.cpp/pull/18283#issuecomment-3710175837 std::unique_ptr task; @@ -335,6 +341,9 @@ struct server_slot { SLT_TRC(*this, "clearing prompt with %zu tokens\n", prompt.tokens.size()); mem.seq_rm(id, -1, -1); + if (ctx_spf) { + llama_memory_seq_rm(llama_get_memory(ctx_spf), id, -1, -1); + } prompt.clear(); } @@ -404,6 +413,9 @@ struct server_slot { // clear multimodal state mbatch.reset(); + + spec_prefill_active = false; + spec_prefill_tokens.clear(); } void init_sampler() const { @@ -471,6 +483,14 @@ struct server_slot { return !!spec; } + const server_tokens & prompt_src() const { + return spec_prefill_active ? spec_prefill_tokens : task->tokens; + } + + int32_t n_prompt_src() const { + return spec_prefill_active ? (int32_t) spec_prefill_tokens.size() : task->n_tokens(); + } + void add_token(const completion_token_output & token) { if (!is_processing()) { SLT_WRN(*this, "%s", "slot is not processing\n"); @@ -508,10 +528,12 @@ struct server_slot { // no speculative decoding i_batch = batch.size(); + llama_pos pos = prompt.tokens.pos_next(); + if (!inp_embd.empty()) { - add_ok &= batch.add(id, inp_embd, prompt.tokens.pos_next(), true, false); + add_ok &= batch.add(id, inp_embd, pos, true, false); } else { - add_ok &= batch.add(id, sampled, prompt.tokens.pos_next(), true, false); + add_ok &= batch.add(id, sampled, pos, true, false); } SLT_DBG(*this, "slot decode token, id=%d, n_ctx = %d, n_tokens = %d, truncated = %d\n", @@ -527,7 +549,7 @@ struct server_slot { spec_i_batch.push_back(batch.size() + i + 1); } - auto pos0 = prompt.tokens.pos_next(); + llama_pos pos0 = prompt.tokens.pos_next(); add_ok &= batch.add(id, sampled, pos0++, true, false); for (auto token : spec_draft) { @@ -551,6 +573,10 @@ struct server_slot { state = SLOT_STATE_IDLE; + if (ctx_spf) { + llama_memory_seq_rm(llama_get_memory(ctx_spf), id, -1, -1); + } + // do not keep context of the child slots - the parent's context is enough if (task->is_child()) { prompt_clear(); @@ -622,11 +648,18 @@ struct server_slot { return; } - const double n_prompt_second = stats.n_prompt_tps(); - const double f_progress = task->n_tokens() > 0 ? (double) prompt.n_tokens() / task->n_tokens() : 0.0; + const double f_progress = n_prompt_src() > 0 ? (double) prompt.n_tokens() / n_prompt_src() : 0.0; + int n_tokens_disp = (int) stats.n_prompt_processed; + double n_prompt_second = stats.n_prompt_tps(); + + if (spec_prefill_active && task && n_prompt_src() > 0) { + const int n_uncached = (int) (task->n_tokens() - stats.n_prompt_cached); + n_tokens_disp = (int) (f_progress * n_uncached); + n_prompt_second = t_prompt_total > 0.0 ? (1e3 / t_prompt_total) * n_tokens_disp : 0.0; + } SLT_INF(*this, "prompt processing, n_tokens = %6d, progress = %.2f, t = %6.2f s / %.2f tokens per second\n", - (int) stats.n_prompt_processed, f_progress, t_prompt_total / 1e3, n_prompt_second); + n_tokens_disp, f_progress, t_prompt_total / 1e3, n_prompt_second); } void print_timings() const { @@ -697,7 +730,7 @@ struct server_slot { if (ptask) { res["id_task"] = ptask->id; - res["n_prompt_tokens"] = (int32_t) prompt.tokens.size(); + res["n_prompt_tokens"] = ptask->n_tokens(); res["n_prompt_tokens_processed"] = stats.n_prompt_processed; res["n_prompt_tokens_cache"] = stats.n_prompt_cached; res["params"] = ptask->params.to_json(only_metrics); @@ -725,6 +758,10 @@ struct server_slot { mem.seq_rm(other.id, -1, -1); mem.seq_cp(id, other.id, -1, -1); + if (ctx_spf) { + llama_memory_seq_rm(llama_get_memory(ctx_spf), other.id, -1, -1); + } + other.i_batch = i_batch; other.stats = stats; @@ -890,6 +927,12 @@ struct server_context_impl { common_speculative_init_result_ptr spec_init; + common_init_result_ptr llama_init_spf; + llama_context_ptr ctx_spf_own; + llama_model * model_spf = nullptr; + llama_context * ctx_spf = nullptr; + common_sampler_ptr smpl_spf; + common_context_seq_rm_type ctx_tgt_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO; common_context_seq_rm_type ctx_dft_seq_rm_type = COMMON_CONTEXT_SEQ_RM_TYPE_NO; @@ -937,6 +980,13 @@ struct server_context_impl { void destroy() { spec.reset(); + + smpl_spf.reset(); + ctx_spf_own.reset(); + llama_init_spf.reset(); + ctx_spf = nullptr; + model_spf = nullptr; + spec_init.reset(); ctx_dft = nullptr; @@ -1001,6 +1051,53 @@ struct server_context_impl { return true; } + void apply_spec_prefill(server_slot & slot) { + slot.spec_prefill_active = false; + slot.spec_prefill_tokens.clear(); + + if (!ctx_spf || !smpl_spf || !params_base.speculative.prefill.enabled) { + return; + } + if (!slot.task || !slot.task->need_logits() || slot.task->tokens.has_media()) { + return; + } + + const llama_tokens & prompt = slot.task->tokens.get_tokens(); + if (prompt.empty()) { + return; + } + + llama_memory_seq_rm(llama_get_memory(ctx_spf), slot.id, -1, -1); + common_sampler_reset(smpl_spf.get()); + + const common_speculative_prefill_result res = common_speculative_prefill_execute( + ctx_spf, smpl_spf.get(), prompt, slot.id, params_base.speculative.prefill); + + llama_memory_seq_rm(llama_get_memory(ctx_spf), slot.id, -1, -1); + + if (res.kept_indices.empty() || (int32_t) res.kept_indices.size() >= (int32_t) prompt.size()) { + return; + } + + llama_tokens kept; + kept.reserve(res.kept_indices.size()); + for (const int32_t idx : res.kept_indices) { + if (idx >= 0 && idx < (int32_t) prompt.size()) { + kept.push_back(prompt[idx]); + } + } + if (kept.size() < 2) { + return; + } + + slot.spec_prefill_tokens = server_tokens(kept, slot.task->tokens.has_mtmd); + slot.spec_prefill_active = true; + + SLT_INF(slot, "speculative prefill kept %d / %d tokens (%.1f%%)\n", + (int) kept.size(), (int) prompt.size(), + 100.0f * (float) kept.size() / (float) prompt.size()); + } + // load the model and initialize llama_context // this may also be called to resume from sleeping state bool load_model(common_params & params) { @@ -1011,6 +1108,14 @@ struct server_context_impl { const bool is_resume = sleeping; params_base = params; + // A zero draft length means no speculation. Without this the context is sized from + // common_speculative_n_max() == 0 while the draft path is still initialized, and the + // first decode trips GGML_ASSERT(n_outputs_max <= cparams.n_outputs_max). + if (params_base.speculative.has_dft() && common_speculative_n_max(¶ms_base.speculative) <= 0) { + SRV_WRN("%s", "draft length is 0 - disabling speculative decoding\n"); + params_base.speculative.types = { COMMON_SPECULATIVE_TYPE_NONE }; + } + const auto output_limits = server_output_limits(params_base); params_base.n_outputs_max = output_limits.total; params_base.n_outputs_max_per_seq = output_limits.per_seq; @@ -1150,6 +1255,107 @@ struct server_context_impl { load_progress_callback(1.0f, &load_progress_spec); } + if (params_base.speculative.prefill.enabled) { + if (llama_model_is_recurrent(model_tgt)) { + SRV_ERR("%s", "speculative prefill is not supported for recurrent models\n"); + return false; + } + + const bool has_target_dependent_dft = std::any_of( + params_base.speculative.types.begin(), + params_base.speculative.types.end(), + [](common_speculative_type t) { + return t == COMMON_SPECULATIVE_TYPE_DRAFT_DFLASH || + t == COMMON_SPECULATIVE_TYPE_DRAFT_DSPARK || + t == COMMON_SPECULATIVE_TYPE_DRAFT_EAGLE3 || + t == COMMON_SPECULATIVE_TYPE_DRAFT_MTP; + }); + + common_params_model spf_model = params_base.speculative.prefill.model; + if (spf_model.empty() && !has_target_dependent_dft) { + spf_model = params_base.speculative.draft.mparams; + } + + if (spf_model.empty()) { + SRV_ERR("%s", "speculative prefill enabled but no draft model was provided\n"); + return false; + } + + common_params params_spf = common_base_params_to_speculative(params_base); + params_spf.model = spf_model; + if (params_base.speculative.prefill.n_ctx > 0) { + params_spf.n_ctx = params_base.speculative.prefill.n_ctx; + } + if (params_base.speculative.prefill.n_gpu_layers != -1) { + params_spf.n_gpu_layers = params_base.speculative.prefill.n_gpu_layers; + } + if (!params_base.speculative.prefill.devices.empty()) { + params_spf.devices = params_base.speculative.prefill.devices; + } + params_spf.flash_attn_type = LLAMA_FLASH_ATTN_TYPE_DISABLED; + params_spf.fit_params = false; + params_spf.n_parallel = params_base.n_parallel; + params_spf.speculative.prefill.enabled = false; + params_spf.speculative.draft.mparams = {}; + + const bool reuse_dft = model_dft && !has_target_dependent_dft && ( + params_base.speculative.prefill.model.empty() || + spf_model.path == params_base.speculative.draft.mparams.path); + + if (reuse_dft) { + SRV_INF("reusing draft model for speculative prefill '%s'\n", params_spf.model.get_name().c_str()); + model_spf = model_dft; + } else { + SRV_INF("loading speculative prefill draft model '%s'\n", params_spf.model.get_name().c_str()); + llama_init_spf = common_init_from_params(params_spf, /*model_only=*/true); + model_spf = llama_init_spf ? llama_init_spf->model() : nullptr; + } + if (model_spf == nullptr) { + SRV_ERR("failed to load speculative prefill draft model, '%s'\n", params_spf.model.path.c_str()); + return false; + } + + if (params_base.speculative.prefill.n_ctx <= 0 && params_spf.n_ctx > (int32_t) llama_model_n_ctx_train(model_spf)) { + params_spf.n_ctx = llama_model_n_ctx_train(model_spf); + SRV_INF("capping speculative prefill draft context to training limit (%d tokens)\n", params_spf.n_ctx); + } + + llama_context_params cparams = common_context_params_to_llama(params_spf); + ctx_spf_own.reset(llama_init_from_model(model_spf, cparams)); + ctx_spf = ctx_spf_own.get(); + if (ctx_spf == nullptr) { + SRV_ERR("failed to create speculative prefill context for '%s'\n", params_spf.model.path.c_str()); + return false; + } + + if (llama_model_is_recurrent(model_spf)) { + SRV_ERR("speculative prefill draft model '%s' is recurrent and not supported\n", params_spf.model.get_name().c_str()); + return false; + } + + if (llama_model_target_layer_ids_n(model_spf) > 0) { + SRV_ERR("speculative prefill draft model '%s' is target-dependent and cannot be used for speculative prefill\n", params_spf.model.get_name().c_str()); + return false; + } + + const llama_vocab * vocab_spf = llama_model_get_vocab(model_spf); + const int n_vocab_tgt = llama_vocab_n_tokens(vocab); + const int n_vocab_spf = llama_vocab_n_tokens(vocab_spf); + const int vocab_diff = n_vocab_tgt > n_vocab_spf ? n_vocab_tgt - n_vocab_spf : n_vocab_spf - n_vocab_tgt; + if (vocab_diff > 128) { + SRV_ERR("speculative prefill vocab size difference %d exceeds 128\n", vocab_diff); + return false; + } + + common_params_sampling sparams_spf; + sparams_spf.temp = 0.0f; + smpl_spf.reset(common_sampler_init(model_spf, sparams_spf)); + if (!smpl_spf) { + SRV_ERR("%s", "failed to init speculative prefill sampler\n"); + return false; + } + } + if (has_mmproj) { if (callback_state) { callback_state(SERVER_STATE_LOADING, {{"stage", "mmproj_model"}}); @@ -1290,6 +1496,7 @@ struct server_context_impl { slot.id = i; slot.ctx_tgt = ctx_tgt; slot.ctx_dft = ctx_dft; + slot.ctx_spf = ctx_spf; slot.mem.init(ctx_tgt, ctx_dft); slot.spec = spec.get(); slot.n_ctx = n_ctx_slot(); @@ -1485,6 +1692,13 @@ struct server_context_impl { /* enable_thinking */ enable_thinking, /* reasoning_budget */ params_base.sampling.reasoning_budget_tokens, /* reasoning_budget_msg */ params_base.sampling.reasoning_budget_message, + /* reasoning_budget_soft_ratio */ params_base.sampling.reasoning_budget_soft_ratio, + /* reasoning_budget_soft_msg */ params_base.sampling.reasoning_budget_soft_message, + /* reasoning_budget_soft2_ratio */ params_base.sampling.reasoning_budget_soft2_ratio, + /* reasoning_budget_soft2_msg */ params_base.sampling.reasoning_budget_soft2_message, + /* reasoning_budget_intro_msg */ params_base.sampling.reasoning_budget_intro_message, + /* reasoning_budget_intro_mode */ params_base.sampling.reasoning_budget_intro_mode, + /* reasoning_budget_grace_toks */ params_base.sampling.reasoning_budget_grace_tokens, /* media_path */ params_base.media_path, /* force_pure_content */ params_base.force_pure_content_parser }; @@ -1777,6 +1991,26 @@ struct server_context_impl { SLT_DBG(slot, "launching slot : %s\n", safe_json_to_str(slot.to_json()).c_str()); + if (task.params.sampling.reasoning_budget_intro_mode == "once" && + !task.params.sampling.reasoning_budget_intro_forced.empty() && + task.tokens.size() >= task.params.sampling.reasoning_budget_intro_forced.size()) { + const auto * vocab_dedupe = llama_model_get_vocab(model_tgt); + auto detok = [&](const auto & toks) { + std::string text; + for (size_t i = 0; i < toks.size(); i++) { + if (toks[i] != LLAMA_TOKEN_NULL) { + text += common_token_to_piece(vocab_dedupe, toks[i], false); + } + } + return text; + }; + const std::string needle = detok(task.params.sampling.reasoning_budget_intro_forced); + if (!needle.empty() && detok(task.tokens).find(needle) != std::string::npos) { + task.params.sampling.reasoning_budget_intro_forced.clear(); + SLT_DBG(slot, "%s", "reasoning-budget intro already present in prompt, suppressed (intro-mode=once)\n"); + } + } + // initialize samplers if (task.need_sampling()) { try { @@ -2113,6 +2347,7 @@ struct server_context_impl { res->truncated = slot.truncated; res->n_decoded = slot.stats.n_gen; + res->vocab_usage = llama_model_get_vocab(model_tgt); res->n_prompt_tokens = slot.task->n_tokens(); res->n_prompt_tokens_cache = slot.stats.n_prompt_cached; res->n_tokens_cached = slot.prompt.n_tokens(); @@ -3020,7 +3255,7 @@ struct server_context_impl { llama_memory_seq_pos_max(llama_get_memory(ctx_tgt), slot.id)); if (use_ckpt_dft) { - slot.spec_ckpt.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + slot.spec_ckpt.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY | LLAMA_STATE_SEQ_FLAGS_ON_DEVICE); } slot.spec_prompt = slot.prompt.tokens.get_text_tokens(); @@ -3059,7 +3294,7 @@ struct server_context_impl { if (ctx_dft) { if (use_ckpt_dft) { - ckpt.load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + ckpt.load_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY | LLAMA_STATE_SEQ_FLAGS_ON_DEVICE); } if (!llama_memory_seq_rm(llama_get_memory(ctx_dft), slot.id, ckpt.pos_max + 1, -1)) { @@ -3078,7 +3313,7 @@ struct server_context_impl { if (use_ckpt_tgt) { //const int64_t t_start = ggml_time_us(); - ckpt.update_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + ckpt.update_tgt(ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY | LLAMA_STATE_SEQ_FLAGS_ON_DEVICE); //const int64_t t_total = ggml_time_us() - t_start; //printf("checkpoint total: %f ms\n", t_total / 1000.0); @@ -3090,7 +3325,7 @@ struct server_context_impl { } if (use_ckpt_dft) { - ckpt.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + ckpt.update_dft(ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY | LLAMA_STATE_SEQ_FLAGS_ON_DEVICE); } } }); @@ -3133,19 +3368,23 @@ struct server_context_impl { // this slot still has a prompt to be processed if (slot.state == SLOT_STATE_PROCESSING_PROMPT || slot.state == SLOT_STATE_STARTED) { - const auto & input_tokens = slot.task->tokens; + if (slot.state == SLOT_STATE_STARTED) { + slot.stats.update_prompt_start(); + apply_spec_prefill(slot); + } + + const auto & input_tokens = slot.prompt_src(); + const int32_t n_prompt_in = slot.n_prompt_src(); // used to determine the number of tokens added to the batch for the current slot const auto n_tokens_prev = batch.size(); // TODO: maybe move branch to outside of this loop in the future if (slot.state == SLOT_STATE_STARTED) { - slot.stats.update_prompt_start(); - slot.state = SLOT_STATE_PROCESSING_PROMPT; SLT_TRC(slot, "new prompt, n_ctx_slot = %d, n_keep = %d, task.n_tokens = %d\n", - slot.n_ctx, slot.task->params.n_keep, slot.task->n_tokens()); + slot.n_ctx, slot.task->params.n_keep, n_prompt_in); // print prompt tokens (for debugging) /*if (1) { @@ -3182,39 +3421,39 @@ struct server_context_impl { } if (!slot.can_split()) { - if (slot.task->n_tokens() > n_ubatch) { + if (n_prompt_in > n_ubatch) { send_error(slot, string_format( "input (%d tokens) is too large to process. increase the physical batch " "size (current batch size: %d)", - slot.task->n_tokens(), n_ubatch), + n_prompt_in, n_ubatch), ERROR_TYPE_SERVER); slot.release(); return; } - if (slot.task->n_tokens() > slot.n_ctx) { + if (n_prompt_in > slot.n_ctx) { send_error( slot, string_format( "input (%d tokens) is larger than the max context size (%d tokens). skipping", - slot.task->n_tokens(), slot.n_ctx), + n_prompt_in, slot.n_ctx), ERROR_TYPE_EXCEED_CONTEXT_SIZE); slot.release(); return; } } else { - if (slot.task->n_tokens() >= slot.n_ctx) { + if (n_prompt_in >= slot.n_ctx) { send_error(slot, string_format("request (%d tokens) exceeds the available context size (%d " "tokens), try increasing it", - slot.task->n_tokens(), slot.n_ctx), + n_prompt_in, slot.n_ctx), ERROR_TYPE_EXCEED_CONTEXT_SIZE); slot.release(); return; } - if (slot.task->params.cache_prompt) { + if (slot.task->params.cache_prompt && !slot.spec_prefill_active) { // reuse any previously computed tokens that are common with the new prompt n_past = slot.prompt.tokens.get_common_prefix(input_tokens); @@ -3291,7 +3530,7 @@ struct server_context_impl { llama_pos pos_next = slot.prompt.tokens.pos_next(n_past); // ref: https://github.com/ggml-org/llama.cpp/pull/24110 - const bool has_new_tokens = (n_past < slot.task->n_tokens()); + const bool has_new_tokens = (n_past < n_prompt_in); // the largest pos_min required for a checkpoint to be useful const auto pos_min_thold = std::max(0, pos_next - n_swa - (has_new_tokens ? 0 : 1)); @@ -3307,7 +3546,7 @@ struct server_context_impl { // this is useful for debugging prompt caching if (slots_debug) { const int np0 = std::max(n_past - slots_n_diff, 0); - const int np1 = std::min(n_past + slots_n_diff + 2, std::min(slot.prompt.tokens.size(), slot.task->tokens.size())); + const int np1 = std::min(n_past + slots_n_diff + 2, std::min(slot.prompt.tokens.size(), input_tokens.size())); std::stringstream ss0; std::stringstream ss1; @@ -3332,7 +3571,7 @@ struct server_context_impl { } { - const auto token = slot.task->tokens[i]; + const auto token = input_tokens[i]; const auto piece = token != LLAMA_TOKEN_NULL ? common_token_to_piece(ctx_tgt, token) : "[mtmd]"; ss1 << piece; st1 << std::setw(8) << token; @@ -3400,8 +3639,8 @@ struct server_context_impl { } // [TAG_PROMPT_LOGITS] - if (n_past == slot.task->n_tokens() && n_past > 0) { - SLT_WRN(slot, "need to evaluate at least 1 token for each active slot (n_past = %d, task.n_tokens() = %d)\n", n_past, slot.task->n_tokens()); + if (n_past == n_prompt_in && n_past > 0) { + SLT_WRN(slot, "need to evaluate at least 1 token for each active slot (n_past = %d, task.n_tokens() = %d)\n", n_past, n_prompt_in); n_past--; SLT_WRN(slot, "n_past was set to %d\n", n_past); } @@ -3427,7 +3666,7 @@ struct server_context_impl { if (!slot.can_split()) { // cannot fit the prompt in the current batch - will try next iter - if (batch.size() + slot.task->n_tokens() > n_batch) { + if (batch.size() + n_prompt_in > n_batch) { return; } } @@ -3478,7 +3717,7 @@ struct server_context_impl { while (true) { auto cur_token_idx = slot.prompt.n_tokens(); if ( - cur_token_idx >= slot.task->n_tokens() || + cur_token_idx >= n_prompt_in || input_tokens[cur_token_idx] != LLAMA_TOKEN_NULL // encountered a text token ) { break; @@ -3521,7 +3760,7 @@ struct server_context_impl { const auto last_user_pos = spans.last_user_message_pos(); // add prompt tokens for processing in the current batch - while (slot.prompt.n_tokens() < slot.task->n_tokens() && batch.size() < n_batch) { + while (slot.prompt.n_tokens() < n_prompt_in && batch.size() < n_batch) { // get next token to process llama_token cur_tok = input_tokens[slot.prompt.n_tokens()]; if (cur_tok == LLAMA_TOKEN_NULL) { @@ -3539,9 +3778,10 @@ struct server_context_impl { // embedding requires all tokens in the batch to be output; // MTP also wants logits at every prompt position so the // streaming hook can mirror t_h_nextn into ctx_dft. + llama_pos pos = slot.prompt.tokens.pos_next(); add_ok &= batch.add(slot.id, cur_tok, - /* pos = */ slot.prompt.tokens.pos_next(), + /* pos = */ pos, /* output = */ slot.need_embd(), /* is_prompt = */ true); slot.prompt.tokens.push_back(cur_tok); @@ -3567,7 +3807,7 @@ struct server_context_impl { bool should_break = false; for (int offset : checkpoint_offsets) { const int n_last = std::min(n_batch, offset); - if (slot.task->n_tokens() == slot.prompt.n_tokens() + n_last) { + if (n_prompt_in == slot.prompt.n_tokens() + n_last) { should_break = true; break; } @@ -3583,15 +3823,19 @@ struct server_context_impl { const auto n_tokens_start = slot.prompt.n_tokens() - n_tokens_cur; - const bool near_prompt_end = slot.task->n_tokens() < slot.prompt.n_tokens() + n_ubatch; + const bool near_prompt_end = n_prompt_in < slot.prompt.n_tokens() + n_ubatch; const bool is_user_start = spans.is_user_start(n_tokens_start); const bool is_last_user_message = n_tokens_start == last_user_pos; // entire prompt has been processed - if (slot.prompt.n_tokens() == slot.task->n_tokens()) { + if (slot.prompt.n_tokens() == n_prompt_in) { slot.state = SLOT_STATE_DONE_PROMPT; + if (slot.spec_prefill_active && slot.task) { + slot.stats.n_prompt_processed = slot.task->n_tokens() - slot.stats.n_prompt_cached; + } + GGML_ASSERT(batch.size() > 0); // extract the logits only for the last token @@ -3912,11 +4156,24 @@ struct server_context_impl { GGML_ASSERT(slot.spec_i_batch.size() == n_draft + 1); const auto & synth_probs = common_speculative_get_synth_probs(spec.get()); - auto accepted = synth_probs.empty() - ? common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft) - : server_sample_and_accept_synth( + std::vector accepted; + if (!synth_probs.empty()) { + accepted = server_sample_and_accept_synth( slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft, synth_probs, slot.spec_synth_rng, slot.spec_is_replay); + } else if (slot.spec_is_replay) { + // replayed tokens were accepted before the restore; re-verifying them can + // disagree when logits depend on batch shape, and each disagreement restores + // the same checkpoint again - the slot stops making progress + accepted = slot.spec_draft; + for (const llama_token id : accepted) { + common_sampler_accept(slot.smpl.get(), id, true); + } + accepted.push_back(common_sampler_sample(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch.back())); + common_sampler_accept(slot.smpl.get(), accepted.back(), true); + } else { + accepted = common_sampler_sample_and_accept_n(slot.smpl.get(), slot.ctx_tgt, slot.spec_i_batch, slot.spec_draft); + } slot.spec_i_batch.clear(); GGML_ASSERT(accepted.size() >= 1); @@ -3942,10 +4199,10 @@ struct server_context_impl { SLT_DBG(slot, "restoring speculative checkpoint (pos_min = %d, pos_max = %d, size = %zu)\n", ckpt.pos_min, ckpt.pos_max, ckpt.size()); - ckpt.load_tgt(slot.ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + ckpt.load_tgt(slot.ctx_tgt, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY | LLAMA_STATE_SEQ_FLAGS_ON_DEVICE); if (slot.ctx_dft) { - ckpt.load_dft(slot.ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY); + ckpt.load_dft(slot.ctx_dft, slot.id, LLAMA_STATE_SEQ_FLAGS_PARTIAL_ONLY | LLAMA_STATE_SEQ_FLAGS_ON_DEVICE); } slot.mem.seq_rm(slot.id, ckpt.pos_max + 1, -1); diff --git a/tools/server/server-models.cpp b/tools/server/server-models.cpp index 3d134acf3621..d3219428e4b0 100644 --- a/tools/server/server-models.cpp +++ b/tools/server/server-models.cpp @@ -798,7 +798,10 @@ void server_models::load_models() { }; auto apply_hidden = [&]() { for (auto & [name, inst] : mapping) { - inst.meta.hidden = hidden_models.count(name) > 0; + std::string val; + bool by_preset = inst.meta.preset.get_option(COMMON_ARG_PRESET_HIDDEN, val) + && common_arg_utils::is_truthy(val); + inst.meta.hidden = by_preset || hidden_models.count(name) > 0; } }; // update_args() injects HOST/PORT/ALIAS, so strip them before comparing presets @@ -2019,7 +2022,7 @@ void server_models_routes::init_routes() { std::time_t t = std::time(0); for (const auto & meta : all_models) { if (meta.hidden) { - continue; // cache model deduplicated by a preset + continue; // preset asked to hide it, or a cache model deduplicated by a preset } json status { {"value", server_model_status_to_string(meta.status)}, diff --git a/tools/server/server-schema.cpp b/tools/server/server-schema.cpp index 27ecafb7a595..b2d62d6c20d3 100644 --- a/tools/server/server-schema.cpp +++ b/tools/server/server-schema.cpp @@ -381,12 +381,34 @@ std::vector> make_llama_cmpl_schema(const common_params & } })); + add((new field_bool("reasoning_budget_enabled", params.sampling.reasoning_budget_enabled)) + ->set_desc("Master switch for the hard/soft/intro/grace reasoning budget control and runtime reasoning control")); + add((new field_bool("reasoning_control", params.sampling.reasoning_control)) ->set_desc("Create the budget sampler on demand so reasoning can be ended at runtime")); add((new field_num("reasoning_budget_tokens", params.sampling.reasoning_budget_tokens)) ->set_hard_limits(-1, INT32_MAX) ->set_desc("Number of tokens in the reasoning budget (-1 = disabled)")); + add((new field_num("reasoning_budget_soft_ratio", params.sampling.reasoning_budget_soft_ratio)) + ->set_hard_limits(-1.0f, 1.0f) + ->set_desc("Fraction of the reasoning budget consumed at which to inject a soft warning message before the hard cutoff (<= 0 = disabled)")); + + add((new field_str("reasoning_budget_intro_message")) + ->set_desc("Message forced immediately when the reasoning block starts, announcing the token budget. Use {budget} as a placeholder for the configured reasoning_budget_tokens value") + ->set_handler([&](field_eval_context & ctx, const json & data) { + GGML_ASSERT(ctx.vocab != nullptr); + std::string message = data.at("reasoning_budget_intro_message").get(); + const std::string budget_str = std::to_string(ctx.params.sampling.reasoning_budget_tokens); + for (size_t pos = 0; (pos = message.find("{budget}", pos)) != std::string::npos; pos += budget_str.size()) { + message.replace(pos, 8, budget_str); + } + ctx.params.sampling.reasoning_budget_intro_forced = common_tokenize(ctx.vocab, message, false, true); + })); + + add((new field_num("reasoning_budget_grace_tokens", params.sampling.reasoning_budget_grace_tokens)) + ->set_hard_limits(0, INT32_MAX) + ->set_desc("Once the reasoning budget is exhausted, wait up to this many tokens for a paragraph break before forcing the cutoff (0 = force immediately)")); add((new field_str("reasoning_budget_start_tag")) ->set_desc("Token string marking the start of the reasoning budget section") @@ -417,20 +439,46 @@ std::vector> make_llama_cmpl_schema(const common_params & })); add((new field_str("reasoning_budget_message")) - ->set_desc("Message to prepend to the reasoning budget end tag when forcing it") + ->set_desc("Message forced when the reasoning budget is exhausted. Should include the model's own closing tag (e.g. ) since it is not appended automatically - the exact tag can differ between models/templates. If empty, falls back to forcing just the auto-detected closing tag alone, so the reasoning block still always closes") ->set_handler([&](field_eval_context & ctx, const json & data) { GGML_ASSERT(ctx.vocab != nullptr); - if (!ctx.params.sampling.reasoning_budget_end.empty()) { - llama_tokens end_tag = ctx.params.sampling.reasoning_budget_end.front(); - std::string message = json_value(data, "reasoning_budget_message", std::string()); - if (!message.empty()) { - llama_tokens message_tokens = common_tokenize(ctx.vocab, message, false, true); - end_tag.insert(end_tag.begin(), message_tokens.begin(), message_tokens.end()); + std::string message = data.at("reasoning_budget_message").get(); + if (message.empty()) { + if (!ctx.params.sampling.reasoning_budget_end.empty()) { + ctx.params.sampling.reasoning_budget_forced = ctx.params.sampling.reasoning_budget_end.front(); } - ctx.params.sampling.reasoning_budget_forced = std::move(end_tag); + } else { + ctx.params.sampling.reasoning_budget_forced = common_tokenize(ctx.vocab, message, false, true); } })); + add((new field_str("reasoning_budget_soft_message")) + ->set_desc("Soft-warning message injected partway through the reasoning budget, at the next newline boundary") + ->set_handler([&](field_eval_context & ctx, const json & data) { + GGML_ASSERT(ctx.vocab != nullptr); + std::string message = data.at("reasoning_budget_soft_message").get(); + ctx.params.sampling.reasoning_budget_soft_forced = common_tokenize(ctx.vocab, message, false, true); + })); + + add((new field_num("reasoning_budget_soft2_ratio", params.sampling.reasoning_budget_soft2_ratio)) + ->set_hard_limits(-1.0f, 1.0f) + ->set_desc("Fraction of the reasoning budget consumed at which to inject the SECOND soft warning message (e.g. 0.5 first, 0.75 second) (<= 0 = disabled)")); + + add((new field_str("reasoning_budget_soft2_message")) + ->set_desc("Second soft-warning message, injected partway through the reasoning budget at the next newline boundary") + ->set_handler([&](field_eval_context & ctx, const json & data) { + GGML_ASSERT(ctx.vocab != nullptr); + std::string message = data.at("reasoning_budget_soft2_message").get(); + ctx.params.sampling.reasoning_budget_soft2_forced = common_tokenize(ctx.vocab, message, false, true); + })); + + add((new field_str("reasoning_budget_intro_mode")) + ->set_desc("When to force the intro message: 'every' = every reasoning block, 'once' = only the first reasoning block of a conversation (deduped against the prompt)") + ->set_handler([&](field_eval_context & ctx, const json & data) { + ctx.params.sampling.reasoning_budget_intro_mode = + data.at("reasoning_budget_intro_mode").get(); + })); + add((new field_json("logit_bias")) ->set_desc("Modify the likelihood of specific tokens. Accepts an array of [token, bias] pairs or an object mapping token to bias. Use false as bias to ban a token") ->set_handler([&](field_eval_context & ctx, const json & data) { @@ -560,11 +608,12 @@ task_params eval_llama_cmpl_schema( // debugging { auto budget = params.sampling.reasoning_budget_tokens; - SRV_DBG("reasoning budget: tokens=%d, generation_prompt='%s', start=%zu toks, end=%zu seqs, forced=%zu toks\n", - budget, params.sampling.generation_prompt.c_str(), - params.sampling.reasoning_budget_start.size(), - params.sampling.reasoning_budget_end.size(), - params.sampling.reasoning_budget_forced.size()); + SRV_DBG( + "reasoning budget: tokens=%d, generation_prompt='%s', start=%zu toks, end=%zu seqs, forced=%zu toks, soft_ratio=%.2f, soft_forced=%zu toks, intro_forced=%zu toks, grace_tokens=%d\n", + budget, params.sampling.generation_prompt.c_str(), params.sampling.reasoning_budget_start.size(), + params.sampling.reasoning_budget_end.size(), params.sampling.reasoning_budget_forced.size(), + params.sampling.reasoning_budget_soft_ratio, params.sampling.reasoning_budget_soft_forced.size(), + params.sampling.reasoning_budget_intro_forced.size(), params.sampling.reasoning_budget_grace_tokens); } return params; diff --git a/tools/server/server-task.cpp b/tools/server/server-task.cpp index 0d3beb313cea..8d42e6930781 100644 --- a/tools/server/server-task.cpp +++ b/tools/server/server-task.cpp @@ -362,15 +362,26 @@ json server_task_result_cmpl_final::to_json_non_oaicompat() { return response_fields.empty() ? res : json_get_nested_values(response_fields, res); } +int32_t server_task_result_cmpl_final::count_reasoning_tokens() { + if (n_reasoning_tokens == 0 && vocab_usage && !oaicompat_msg.reasoning_content.empty()) { + n_reasoning_tokens = (int32_t) common_tokenize(vocab_usage, oaicompat_msg.reasoning_content, false).size(); + } + return n_reasoning_tokens; +} + json server_task_result_cmpl_final::usage_json_oaicompat() { + const int32_t n_reasoning = count_reasoning_tokens(); return json { {"completion_tokens", n_decoded}, {"prompt_tokens", n_prompt_tokens}, {"total_tokens", n_decoded + n_prompt_tokens}, {"prompt_tokens_details", json { {"cached_tokens", n_prompt_tokens_cache} }}, + {"completion_tokens_details", json { {"reasoning_tokens", n_reasoning} }}, + {"reasoning", n_reasoning}, }; } + json server_task_result_cmpl_final::to_json_oaicompat() { std::time_t t = std::time(0); json logprobs = json(nullptr); // OAI default to null @@ -590,6 +601,8 @@ json server_task_result_cmpl_final::to_json_oaicompat_resp() { {"output_tokens", n_decoded}, {"total_tokens", n_decoded + n_prompt_tokens}, {"input_tokens_details", json { {"cached_tokens", n_prompt_tokens_cache} }}, + {"output_tokens_details", json { {"reasoning_tokens", count_reasoning_tokens() } }}, + {"reasoning", count_reasoning_tokens()}, }}, }; diff --git a/tools/server/server-task.h b/tools/server/server-task.h index 9c99143f8e19..03bf4901592c 100644 --- a/tools/server/server-task.h +++ b/tools/server/server-task.h @@ -320,6 +320,9 @@ struct completion_token_output { struct server_task_result_cmpl_final : server_task_result { std::string content; llama_tokens tokens; + const llama_vocab * vocab_usage = nullptr; + int32_t n_reasoning_tokens = 0; + bool stream; bool include_usage; @@ -374,6 +377,8 @@ struct server_task_result_cmpl_final : server_task_result { json to_json_non_oaicompat(); json usage_json_oaicompat(); + int32_t count_reasoning_tokens(); + json to_json_oaicompat();